diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..1336b80 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +vendor/parquet/src/bin/parquet-fromcsv-help.txt whitespace=-blank-at-eol diff --git a/.github/workflows/quality.yml b/.github/workflows/quality.yml index 561a569..882bd3d 100644 --- a/.github/workflows/quality.yml +++ b/.github/workflows/quality.yml @@ -64,7 +64,7 @@ jobs: with: tool: cargo-deny@0.19.2 - uses: Swatinem/rust-cache@d1527c8f67ef9af357290475d4af4b0fab8f12e0 - - run: timeout -k 5 60 cargo +1.95 fmt --all -- --check + - run: timeout -k 5 60 cargo +1.95 fmt --package allstack -- --check - run: timeout -k 5 60 cargo +1.95 fmt --manifest-path examples/spendcheck/Cargo.toml -- --check - run: timeout -k 5 90 cargo +1.95 deny check - run: timeout -k 5 270 cargo +1.95 clippy --all-targets --locked -- -D warnings diff --git a/Cargo.lock b/Cargo.lock index 5354eff..86da78e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1468,8 +1468,6 @@ dependencies = [ [[package]] name = "parquet" version = "59.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff322f54b1a0f9288e614ed1f2d329b380af5476420db19f46ffb865e1163d73" dependencies = [ "ahash", "arrow-array", @@ -1923,7 +1921,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.4.3", + "getrandom 0.3.4", "once_cell", "rustix", "windows-sys", @@ -2479,27 +2477,27 @@ checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" [[package]] name = "zstd" -version = "0.13.3" +version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +checksum = "db3ae9bf90af6e40197bef767eeefc46ad6397420fe873830bd85e5fbbfb3c70" dependencies = [ "zstd-safe", ] [[package]] name = "zstd-safe" -version = "7.3.0" +version = "8.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "64d80649ab6db9d9f6f9c80a40becd948eda4714a0a5ac8c4d157a32231c7882" +checksum = "ded520f16b29e1a9b771c722221b2b0a0398c2316d9bb3c42b458a5414d7434b" dependencies = [ "zstd-sys", ] [[package]] name = "zstd-sys" -version = "2.1.0+zstd.1.5.7" +version = "2.1.2+zstd.1.5.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ef0a8027ec3ee71300ab3bcbcd0393f434aa72b91ca6d635a39941deae8eea0" +checksum = "481ada9c543b2323cbb7bbdbb5249177fbf0188ed959ed2c15e4b22421c84ed3" dependencies = [ "cc", "pkg-config", diff --git a/Cargo.toml b/Cargo.toml index 054b8a2..b10aa44 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,6 +22,9 @@ uuid = { version = "1", features = ["v4"] } [dev-dependencies] tempfile = "3" +[patch.crates-io] +parquet = { path = "vendor/parquet" } + [[test]] name = "workspace" harness = false diff --git a/docs/README.md b/docs/README.md index 8ad49eb..45d08fa 100644 --- a/docs/README.md +++ b/docs/README.md @@ -2,6 +2,8 @@ ## Design +- [2026-10-27: Native evidence checkpoints and Chronis qualification](plans/2026-10-27-native-evidence-and-chronis-design.md) — hosted-proof lessons in existing skills, preserved receipt streams, and real CLI integrity acceptance across Codex and Claude. + - [2026-10-25: Native runtime checks and AllSource Prime validation](plans/2026-10-25-native-runtime-checks-design.md) — lessons from real native delivery, bounded command receipts, and authenticated AllSource application qualification. - [2026-10-25: Spendcheck category plan](plans/2026-10-25-spendcheck-category-plan.md) — real Claude native shaping/planning output for the expense-report application change. @@ -16,6 +18,12 @@ ## Reference +- [Native evidence checkpoints and real Chronis delivery](reports/2026-10-27-native-evidence-chronis-verification.md) — actual Codex/Claude calls, baseline failures, independent runtime QA, durable handoff and exact-commit publication. + +- [Fixed compression dependencies with Arrow 59 compatibility](decisions/2026-10-27-parquet-compression-patch.md) — provenance and removal condition for the manifest-only Parquet patch required by strict dependency auditing. + +- [Native evidence request](reference/native-evidence-request.md) — accepted method changes and real Chronis application qualification. + - [Native AllSource workflow verification](reports/2026-10-26-native-allsource-verification.md) — installed runtime receipts, Chronis export compatibility, real hosts, and Prime application acceptance. - [Native skills and complete application verification](reports/2026-10-25-native-skills-application-verification.md) — installed interface, real Codex/Claude use, frozen expense acceptance, and durable publication evidence. @@ -54,3 +62,5 @@ The repository is an implementation preview. Source-level contracts and fixture - [Allstack delivery plan 9dc4c81f-73ca-4b92-89a2-be3c2ca44deb](plans/allstack-9dc4c81f-73ca-4b92-89a2-be3c2ca44deb.md) - [Allstack delivery plan 26340e4d-9eba-4ad5-bbe0-0832086978b5](plans/allstack-26340e4d-9eba-4ad5-bbe0-0832086978b5.md) + +- [Allstack delivery plan f870b303-2b8f-4385-9f6a-c0ed60bb6b61](plans/allstack-f870b303-2b8f-4385-9f6a-c0ed60bb6b61.md) diff --git a/docs/decisions/2026-10-27-parquet-compression-patch.md b/docs/decisions/2026-10-27-parquet-compression-patch.md new file mode 100644 index 0000000..3e501f2 --- /dev/null +++ b/docs/decisions/2026-10-27-parquet-compression-patch.md @@ -0,0 +1,13 @@ +# Fixed compression dependencies with Arrow 59 compatibility + +Allstack's strict dependency audit rejects zstd 0.13.3, zstd-safe 7.3.0 and zstd-sys 2.1.0. The affected packages arrive through published `allsource-core 0.26.1` → `parquet 59.3.0`. The required fixed versions are zstd 0.14.3, zstd-safe 8.2.0 and zstd-sys 2.1.2. + +Published compatible Parquet 59.3.1 and 59.3.2 still require zstd 0.13. The Arrow 59 maintenance branch has the same dependency constraint. Moving Parquet to a newer Arrow major would change the RecordBatch type boundary used by Core. Disabling compression would remove support for existing encoded data. Neither is necessary for this repair. + +The root manifest patches Parquet to `vendor/parquet`, copied from the exact 59.3.0 crate archive whose checksum is recorded in the lockfile. Only its normalized and original manifests change: both runtime and development zstd requirements become 0.14.3. Library source, bundled Apache license and notice, codec features and version stay byte-identical to the published package. The lockfile selects the fixed compression stack. No advisory is ignored and no audit severity is reduced. + +The patch keeps the existing Arrow/Parquet 59 API and compression formats. Qualification must compare the vendored files with the verified archive, execute strict dependency auditing, rebuild the installed toolkit, and exercise persistence reopening plus native delivery. Remove the local patch when a compatible upstream Parquet release carries the fixed dependency requirement; verify the actual manifest rather than assuming a patch-version bump includes it. + +`tests/vendor_integrity.rs` verifies the complete file set and every upstream digest, reversing only the two manifest requirement edits before comparison. First-party formatting targets the Allstack package explicitly; external source stays byte-identical to its archive and is checked by the provenance test. This preserves first-party formatting coverage without rewriting upstream code. + +The published `parquet-fromcsv-help.txt` fixture contains eight space-only help-output lines. Its file-specific Git whitespace attribute permits those trailing spaces; the integrity test still rejects any byte change to the fixture. Whitespace checks remain active for first-party source and all other files. diff --git a/docs/plans/2026-10-27-native-evidence-and-chronis-design.md b/docs/plans/2026-10-27-native-evidence-and-chronis-design.md new file mode 100644 index 0000000..4b74940 --- /dev/null +++ b/docs/plans/2026-10-27-native-evidence-and-chronis-design.md @@ -0,0 +1,64 @@ +# Native evidence checkpoints and Chronis qualification + +## Complete user job + +A developer asks for a real application repair through native skills. Claude reads the installed planning method and current source, Codex implements the accepted change, Claude independently reviews and exercises the application, and Allstack publishes the verified commit through one PR per repository. AllSource retains delivery history, Chronis owns claims and completion, and project Prime returns source-linked plan and outcome records. Evidence remains readable after owned worktrees are removed. + +## Observed constraints + +Merged [AllSource PR #364](https://github.com/all-source-os/all-source/pull/364) records existing-consumer compatibility, deployed-source provenance, actual host calls, and a failed HTTP observation whose cause could not be established because its body was missing. Its [verification report](https://github.com/all-source-os/all-source/blob/fabca9d0e420578774f0d80eecaa8480f15e8706/docs/reports/2026-10-26-managed-key-lifecycle-verification.md) separates synthetic hosted sync from the later existing-key admission proof. These are bounded observations, not evidence that every tenant or deployment succeeded. + +The current native interface already has seven methods, bounded `allstack check` receipts, source-linked memory, native adoption, and exact-commit shipping. The missing guidance concerns evidence decisions around those existing capabilities. No new engine behavior or catalog entry is needed for this correction. + +## Method changes + +| Method | Decision added | +| --- | --- | +| `allstack-plan` | Read producer revision/count semantics and representative retained history; include an existing consumer's next operation. Compare frozen gates with changed components and name supplemental coverage. | +| `allstack-review` | Trace legacy compatibility and the complete existing-consumer path. Successful authentication alone does not establish application success. | +| `allstack-qa` | Preserve safely redacted failed HTTP responses and later successes separately. Bind host execution to actual calls; distinguish running-image proof from final-commit checks. | +| `allstack-handoff` | Ask the receiving native host to read its installed method and current source; retain calls and results instead of accepting its summary as execution evidence. | +| `allstack-ship` | Archive epoch receipts and raw streams before shipping removes accepted worktrees; verify lengths/digests, read back tracker/memory through their engines, and preserve state before owned cleanup. | + +The installed native guide carries the packaging implications. Skill bodies remain identical across Codex and Claude apart from existing host overlays. The seven names, invocation modes, permission boundaries, CLI contract and delivery state machine remain stable. + +## Alternatives + +1. **Update existing method data, then qualify a real change.** This uses the established packaging and runtime mechanisms. Recommended because the observed gaps concern what an agent checks and reports. +2. Add an orchestration layer that forces a new receipt schema or deployment state machine. This would expand the product before the method-level repair is exercised; current receipts already retain the necessary streams and identities. +3. Add a retrospective only. This preserves evidence but leaves incoming native hosts without the operational guidance at the relevant decision point. + +## Real application target: Chronis TOON integrity + +Current AllSource main `e1416e58eb485bce243f1cc728160f1302289e1c` contains Chronis 0.9.1. `apps/chronis/src/presentation/toon.rs` escapes pipes and newlines in titles but leaves backslashes ambiguous and prints `claimed_by` without escaping. An isolated real CLI reproduction with installed `cn 0.9.1` demonstrated two failures: + +- A title containing literal backslash plus `n` and a title containing an actual newline produce identical rendered title values. +- A custom `CN_AGENT_ID` containing a pipe and newline inserts an extra table column and a `status:done`-looking detail line, although actual task state remains in progress. + +Private raw commands and outputs are retained under the founder-mode ignored `.allstack/proof/native-lessons-2026-10-27/` directory. Installed-binary reproduction is supporting evidence; qualification must also rebuild and exercise the current baseline source. + +Repair the TOON presentation boundary. Preserve exact stored task/claim data and the JSON representation; do not sanitize writes or change task lifecycle rules. Define unambiguous escaping for backslashes, pipes and line controls, apply it consistently to free-text values, and verify ordinary records remain stable. Inspect downstream readers before fixing the contract. No new dependency, hosted credential, production deployment or store migration is required. + +Authenticated Claude session `49403bff-e4f2-4447-94ef-9045edaccfd1` read the installed native planning skill, actual baseline receipt and current formatter/callers. Its 12 observed file-tool calls support the proposed repair site; it performed no application execution. Native receipt `1f00b611-0ee9-4c99-b662-f6572fbce943` records a normal zero exit, and all archived raw streams matched their recorded lengths and digests. Its source inspection does not replace the required current-baseline build. + +The acceptance parser must scan escapes once and count only unescaped field delimiters. A plain `split('|')` would reject valid escaped pipes, and chained reverse substitutions can collapse literal escape sequences. Assert round trips through a single-pass decoder used by the acceptance harness. Preserve the existing comma-separated dependency representation for generated IDs; changing that subformat is outside this defect's observed path. + +Expected application files are the existing Chronis TOON encoder, focused sibling/integration tests, and an indexed application plan/report. Final file selection follows source-grounded native planning. Do not edit unrelated primary-checkout changes or another session's open deployment PR. + +## Qualification + +1. Validate changed skills with the bundled skill validator. Run existing native packaging/conflict checks and applicable Allstack gates; do not add tests that merely match instruction wording. +2. Install the tested Allstack executable and read back its digest, bundled methods, and generated methods for both hosts. Replace only resources verified against their previous generated form; preserve local customizations. +3. Freeze Chronis acceptance before implementation. Build baseline and candidate from their actual source, with one machine compiler at a time and commands bounded below 300 seconds. Warm caches do not establish artifact provenance. +4. Through the real `cn` CLI in an isolated workspace, create distinct literal/control-character titles, claim using a delimiter-bearing agent identity, reopen with separate processes, list/show, complete, and read back persisted state through Chronis. Assert unambiguous output, stable field/row counts using an escape-aware parser, no injected status field, preserved JSON data, and unchanged ordinary output. Include invalid lifecycle transitions with explicit expected exits. +5. Use authenticated Codex and Claude with their normal configured models and permission modes. Retain the receiving host's actual skill-read calls, source reads, test commands, output and executed case counts. A source review alone cannot satisfy runtime QA. +6. Adopt the native application candidate with a frozen policy covering Chronis, verify the single commit, and ship its one PR. Read back remote branch/PR head, final CI, Chronis completion, and project Prime plan/review/QA outcomes. +7. Preserve receipt/raw-stream paths, byte lengths and digests outside disposable epoch worktrees before shipment. After publication and store readback, archive isolated coordinator state and remove only owned idle worktrees/builds. + +The Allstack report must distinguish baseline failure, candidate behavior, actual work performed by each host, installed toolkit provenance, policy coverage, and material limits. Qualification covers sequential real hosts on this machine; it does not imply a second physical machine, hosted sync on a new revision, or deployment. + +## Branch and publication prerequisites + +The user authorized `feat/native-evidence-checkpoints` in Allstack and `fix/chronis-toon-escaping` in AllSource, with commits and one PR each. Both branches start from their current main. The isolated AllSource coordinator preserves unrelated changes in the primary checkout and remains separate from the deployment-gate PR owned elsewhere. + +Supplied machine policy requires explicit new-branch authorization and prohibits trunk pushes. Commit and PR publication proceed only on those authorized live branches. Merge and deployment remain separate actions. diff --git a/docs/plans/allstack-f870b303-2b8f-4385-9f6a-c0ed60bb6b61.md b/docs/plans/allstack-f870b303-2b8f-4385-9f6a-c0ed60bb6b61.md new file mode 100644 index 0000000..4b74940 --- /dev/null +++ b/docs/plans/allstack-f870b303-2b8f-4385-9f6a-c0ed60bb6b61.md @@ -0,0 +1,64 @@ +# Native evidence checkpoints and Chronis qualification + +## Complete user job + +A developer asks for a real application repair through native skills. Claude reads the installed planning method and current source, Codex implements the accepted change, Claude independently reviews and exercises the application, and Allstack publishes the verified commit through one PR per repository. AllSource retains delivery history, Chronis owns claims and completion, and project Prime returns source-linked plan and outcome records. Evidence remains readable after owned worktrees are removed. + +## Observed constraints + +Merged [AllSource PR #364](https://github.com/all-source-os/all-source/pull/364) records existing-consumer compatibility, deployed-source provenance, actual host calls, and a failed HTTP observation whose cause could not be established because its body was missing. Its [verification report](https://github.com/all-source-os/all-source/blob/fabca9d0e420578774f0d80eecaa8480f15e8706/docs/reports/2026-10-26-managed-key-lifecycle-verification.md) separates synthetic hosted sync from the later existing-key admission proof. These are bounded observations, not evidence that every tenant or deployment succeeded. + +The current native interface already has seven methods, bounded `allstack check` receipts, source-linked memory, native adoption, and exact-commit shipping. The missing guidance concerns evidence decisions around those existing capabilities. No new engine behavior or catalog entry is needed for this correction. + +## Method changes + +| Method | Decision added | +| --- | --- | +| `allstack-plan` | Read producer revision/count semantics and representative retained history; include an existing consumer's next operation. Compare frozen gates with changed components and name supplemental coverage. | +| `allstack-review` | Trace legacy compatibility and the complete existing-consumer path. Successful authentication alone does not establish application success. | +| `allstack-qa` | Preserve safely redacted failed HTTP responses and later successes separately. Bind host execution to actual calls; distinguish running-image proof from final-commit checks. | +| `allstack-handoff` | Ask the receiving native host to read its installed method and current source; retain calls and results instead of accepting its summary as execution evidence. | +| `allstack-ship` | Archive epoch receipts and raw streams before shipping removes accepted worktrees; verify lengths/digests, read back tracker/memory through their engines, and preserve state before owned cleanup. | + +The installed native guide carries the packaging implications. Skill bodies remain identical across Codex and Claude apart from existing host overlays. The seven names, invocation modes, permission boundaries, CLI contract and delivery state machine remain stable. + +## Alternatives + +1. **Update existing method data, then qualify a real change.** This uses the established packaging and runtime mechanisms. Recommended because the observed gaps concern what an agent checks and reports. +2. Add an orchestration layer that forces a new receipt schema or deployment state machine. This would expand the product before the method-level repair is exercised; current receipts already retain the necessary streams and identities. +3. Add a retrospective only. This preserves evidence but leaves incoming native hosts without the operational guidance at the relevant decision point. + +## Real application target: Chronis TOON integrity + +Current AllSource main `e1416e58eb485bce243f1cc728160f1302289e1c` contains Chronis 0.9.1. `apps/chronis/src/presentation/toon.rs` escapes pipes and newlines in titles but leaves backslashes ambiguous and prints `claimed_by` without escaping. An isolated real CLI reproduction with installed `cn 0.9.1` demonstrated two failures: + +- A title containing literal backslash plus `n` and a title containing an actual newline produce identical rendered title values. +- A custom `CN_AGENT_ID` containing a pipe and newline inserts an extra table column and a `status:done`-looking detail line, although actual task state remains in progress. + +Private raw commands and outputs are retained under the founder-mode ignored `.allstack/proof/native-lessons-2026-10-27/` directory. Installed-binary reproduction is supporting evidence; qualification must also rebuild and exercise the current baseline source. + +Repair the TOON presentation boundary. Preserve exact stored task/claim data and the JSON representation; do not sanitize writes or change task lifecycle rules. Define unambiguous escaping for backslashes, pipes and line controls, apply it consistently to free-text values, and verify ordinary records remain stable. Inspect downstream readers before fixing the contract. No new dependency, hosted credential, production deployment or store migration is required. + +Authenticated Claude session `49403bff-e4f2-4447-94ef-9045edaccfd1` read the installed native planning skill, actual baseline receipt and current formatter/callers. Its 12 observed file-tool calls support the proposed repair site; it performed no application execution. Native receipt `1f00b611-0ee9-4c99-b662-f6572fbce943` records a normal zero exit, and all archived raw streams matched their recorded lengths and digests. Its source inspection does not replace the required current-baseline build. + +The acceptance parser must scan escapes once and count only unescaped field delimiters. A plain `split('|')` would reject valid escaped pipes, and chained reverse substitutions can collapse literal escape sequences. Assert round trips through a single-pass decoder used by the acceptance harness. Preserve the existing comma-separated dependency representation for generated IDs; changing that subformat is outside this defect's observed path. + +Expected application files are the existing Chronis TOON encoder, focused sibling/integration tests, and an indexed application plan/report. Final file selection follows source-grounded native planning. Do not edit unrelated primary-checkout changes or another session's open deployment PR. + +## Qualification + +1. Validate changed skills with the bundled skill validator. Run existing native packaging/conflict checks and applicable Allstack gates; do not add tests that merely match instruction wording. +2. Install the tested Allstack executable and read back its digest, bundled methods, and generated methods for both hosts. Replace only resources verified against their previous generated form; preserve local customizations. +3. Freeze Chronis acceptance before implementation. Build baseline and candidate from their actual source, with one machine compiler at a time and commands bounded below 300 seconds. Warm caches do not establish artifact provenance. +4. Through the real `cn` CLI in an isolated workspace, create distinct literal/control-character titles, claim using a delimiter-bearing agent identity, reopen with separate processes, list/show, complete, and read back persisted state through Chronis. Assert unambiguous output, stable field/row counts using an escape-aware parser, no injected status field, preserved JSON data, and unchanged ordinary output. Include invalid lifecycle transitions with explicit expected exits. +5. Use authenticated Codex and Claude with their normal configured models and permission modes. Retain the receiving host's actual skill-read calls, source reads, test commands, output and executed case counts. A source review alone cannot satisfy runtime QA. +6. Adopt the native application candidate with a frozen policy covering Chronis, verify the single commit, and ship its one PR. Read back remote branch/PR head, final CI, Chronis completion, and project Prime plan/review/QA outcomes. +7. Preserve receipt/raw-stream paths, byte lengths and digests outside disposable epoch worktrees before shipment. After publication and store readback, archive isolated coordinator state and remove only owned idle worktrees/builds. + +The Allstack report must distinguish baseline failure, candidate behavior, actual work performed by each host, installed toolkit provenance, policy coverage, and material limits. Qualification covers sequential real hosts on this machine; it does not imply a second physical machine, hosted sync on a new revision, or deployment. + +## Branch and publication prerequisites + +The user authorized `feat/native-evidence-checkpoints` in Allstack and `fix/chronis-toon-escaping` in AllSource, with commits and one PR each. Both branches start from their current main. The isolated AllSource coordinator preserves unrelated changes in the primary checkout and remains separate from the deployment-gate PR owned elsewhere. + +Supplied machine policy requires explicit new-branch authorization and prohibits trunk pushes. Commit and PR publication proceed only on those authorized live branches. Merge and deployment remain separate actions. diff --git a/docs/reference/native-evidence-request.md b/docs/reference/native-evidence-request.md new file mode 100644 index 0000000..13c082f --- /dev/null +++ b/docs/reference/native-evidence-request.md @@ -0,0 +1,7 @@ +# Native evidence checkpoints + +Update existing native planning, review, QA, handoff and shipping methods from the recorded hosted-proof failures. Retain failed HTTP observations, distinguish actual host calls from summaries, bind runtime and final-source evidence separately, cover changed components, and archive receipt streams before disposable worktrees disappear. + +Use the [accepted design](../plans/2026-10-27-native-evidence-and-chronis-design.md). Preserve the seven native skill names and existing delivery mechanics. Validate installed packaging for both hosts, then qualify a real Chronis formatter repair through authenticated Codex and Claude and publish one verified PR per repository. AllSource, Chronis and project Prime remain the underlying history, ownership and memory systems. + +The user authorized `feat/native-evidence-checkpoints`, commits and one PR. Merge, deployment and hosted credential changes are outside this request. diff --git a/docs/reports/2026-10-27-native-evidence-chronis-verification.md b/docs/reports/2026-10-27-native-evidence-chronis-verification.md new file mode 100644 index 0000000..d4823a4 --- /dev/null +++ b/docs/reports/2026-10-27-native-evidence-chronis-verification.md @@ -0,0 +1,53 @@ +# Native evidence checkpoints and real Chronis delivery + +## Delivered behavior + +The seven native skills remain the user interface. Updated planning and review methods require producer/history compatibility and an existing consumer's next operation. QA retains failed responses and distinguishes actual host calls, running artifacts and final-commit gates. Handoff requires the receiving host to read its installed method and current source. Shipping archives epoch-local streams before cleanup and reads tracker/memory state through their engines. + +The [accepted design](../plans/2026-10-27-native-evidence-and-chronis-design.md) and [request](../reference/native-evidence-request.md) define this change. Existing AllSource history, Chronis ownership and project Prime remain underneath; no additional orchestrator or workflow catalog entry was introduced. + +## Real application delivered + +[AllSource PR #368](https://github.com/all-source-os/all-source/pull/368) publishes exact verified commit `d6f1b0673785cd53e1c1faa5e90ad677f66ecdd0`. Its [versioned application report](https://github.com/all-source-os/all-source/blob/d6f1b0673785cd53e1c1faa5e90ad677f66ecdd0/docs/reports/2026-10-27-chronis-native-verification.md) records source, executable and receipt identities. + +Chronis now escapes backslashes, pipes, line controls and non-final action-field colons at the presentation boundary. Stored titles and claimant identities remain unchanged. Delimiter-bearing identities cannot insert an extra column or status field. Ordinary output, legacy IDs and silent repeated completion retain their contracts. + +| Evidence | Observed result | +| --- | --- | +| Frozen four-case acceptance against freshly built baseline | 1 passed, 3 failed; expected exit 101 | +| Same acceptance against rebuilt candidate | 4 passed, 0 failed, 0 ignored | +| Independent authenticated Claude execution | Same 4 cases passed through installed Allstack | +| Final exact-commit Chronis policy | Formatting, 102 executed tests, strict all-target Clippy and whitespace passed | +| Native destination readback | Remote branch and PR head equal the verified commit; one commit | +| Actual Chronis readback | Task `t-691ec2` done; reason: remote commit and PR head verified | +| Actual project Prime readback | Source-linked plan/review/QA records and published outcome retrieved through Allstack | + +Acceptance SHA-256 is `46eeaa13966fc153bbc65fb10f36ba9c082474e4dfb0ba0f77cb4e9e77df534a` in both builds. The test executes separate real CLI processes for create, claim, list/show, completion, rejected transitions and JSON state readback. No direct storage-file parsing substitutes for Chronis. + +The baseline CLI independently reproduced title collision and an injected status-looking line while the actual task remained in progress. Baseline and candidate use the same one-line corrected standalone lockfile. Two inherited Clippy warnings and existing claim-test formatting were repaired without suppressions or reduced checks. + +## Actual hosts and ownership + +Authenticated Claude session `49403bff-e4f2-4447-94ef-9045edaccfd1` read the installed planning method and current source: 12 observed file calls. Real Codex session `01a17d58-bec1-7a03-a656-ef5d5a3e9836` wrote the application formatter and acceptance test. An initial bounded inspection timed out after edits; retained source and receipts supported continuation without replaying implementation. + +Independent Claude source review identified colon ambiguity for legacy action IDs. Codex corrected that boundary and the completion expectation after actual dispatch inspection. Claude QA session `f9da996b-5cb5-4dbc-a154-4fb612330dfc` read its installed QA method, checked executable hashes and executed the acceptance binary through Allstack. Inner receipt `34f692ef-8dd8-4dc6-b278-0a6392df64f6` records four passing tests. Parent readback confirmed the original CLI hash after QA. + +Native run `b67f2a72-0d88-40b5-bf4d-efdc492f4f99` moved Codex epoch 1 → Claude epoch 2 → Codex epoch 3. Both receiving hosts read installed handoff instructions and successfully resumed. Actual Chronis claimant readbacks matched each host and epoch. Codex's initial read-only resume failed permission checks despite a successful host-process exit; the subsequent workspace-permitted receipt establishes successful continuation. Final independent Claude inspection also verified the behavior-preserving Kanban match guards. + +## Installed artifact and package integrity + +The workspace-local installed Allstack executable has SHA-256 `c29cc5f1d143878e38e726e4a1fc45d6cfb40b0b5e49b674067407fb77724d78`. It is the macOS artifact from [green CI run 40510117757](https://github.com/all-source-os/allstack/actions/runs/40510117757). CI source was synthetic merge `7efd61db76e2cb980145e7074f637d7cf10affaf`; its Git tree `9742a7b59db3dcdb4a9886c265ba3c428bb9740c` equals feature head `0f7f764b492ac9ff9c5b4906201adc516c79b7aa`. Later changes add verification documentation, policy environment binding and a byte-preserving Git whitespace attribute; bundled methods and binary inputs remain the same. + +Packaging qualification compared all 14 generated Codex/Claude skill files, created no stores and preserved a deliberately conflicting customization. The installed fixed-dependency executable independently lists the same seven skills without opening stores. Application installation replaced only resources verified against their previous generated contents. Explicit host directories in PATH resolved Claude discovery; actual application doctor then reported setup and executables ready. + +Allstack's complete native policy executed 74 toolkit tests with two configured skips, eight expense-example tests and four handoff-fixture tests. Toolkit/example Clippy, first-party formatting and source whitespace checks passed. The byte-verified upstream help fixture has a file-specific trailing-space attribute; other whitespace checks remain active. Actual authenticated Chronis work above supplies real host evidence independently of fixture execution. + +Strict dependency auditing exposed inherited compression advisories. The [indexed compression decision](../decisions/2026-10-27-parquet-compression-patch.md) records a manifest-only patch to the exact published Parquet archive. All 227 upstream files are checked by an integrity test; only two zstd requirements change. Fixed versions are zstd 0.14.3, zstd-safe 8.2.0 and zstd-sys 2.1.2. Unique package count remains 258. Local cargo-deny and final-current Linux/macOS CI passed; no advisory exception was added. + +## Evidence retention and limits + +Before application shipping removed accepted epoch worktrees, 258 evidence files were copied to ignored proof storage. All 19 command receipts and 57 raw streams were checked against recorded lengths and SHA-256 digests. Exported native history and final gate receipts preserve AllSource engine readback. Published-outcome Prime node `bc3cf8f2-3f3f-4876-a271-979a3e698cd4` points to PR #368. Private transcripts and source-linked local state remain outside tracked deliverables. + +Qualification covers sequential real authenticated hosts on this machine and isolated local task data. It does not establish a second physical machine, a new hosted sync revision, browser acceptance, merge or deployment. Dependency-ID comma subformat remains unchanged. A failed nested-epoch Cargo preflight and initial host timeouts remain distinguishable from executed tests. + +Neither repository currently has applicable branch protection/rulesets; the configured hook directory has no commit/push hooks. CI is visible but does not enforce merge. Publication uses explicit native exact-commit verification and remote destination readback. diff --git a/skills/allstack-handoff/SKILL.md b/skills/allstack-handoff/SKILL.md index 2dc4cda..b200b6a 100644 --- a/skills/allstack-handoff/SKILL.md +++ b/skills/allstack-handoff/SKILL.md @@ -13,4 +13,6 @@ Use `allstack handoff --epoch --to codex|claude-code`, then `alls If invoking another real host CLI, use existing supervised Allstack plan/build operations with normal authentication and permission modes. Native tools used outside those operations have the host's own authority and are not fenced by the epoch. Do not claim global protection. +For independent native planning, review, or QA, ask the receiving host to read its installed skill and current source. Retain actual calls and command outputs with the target revision. Do not report a fresh skill read or executed check from the receiving host's final text alone. + Report which job moved, the receiving host, preserved work, and the next user-relevant action. Recovery must inspect receipts and actual edits; never replay an uncertain build merely because its final response is missing. Retain ownership and process cleanup until the operation is complete. diff --git a/skills/allstack-plan/SKILL.md b/skills/allstack-plan/SKILL.md index 1e72861..ad2584e 100644 --- a/skills/allstack-plan/SKILL.md +++ b/skills/allstack-plan/SKILL.md @@ -16,9 +16,12 @@ Read the request, repository instructions, relevant code, and actual gates. Trac - **Architecture:** choose the owner of each invariant; list the smallest coherent files and boundaries that must change. Compare an existing seam with a new abstraction and explain the choice. - **Failure behavior:** follow invalid input, empty results, partial writes, retries, interruption, and duplicate execution where relevant. Name which component detects each failure and what the user sees. - **Compatibility:** inspect current callers, serialized data, command output, and persisted state that the change touches. Specify what remains stable and how compatibility will be exercised. +- For persisted contracts, inspect the producer's actual revision/count semantics and representative retained history, including legacy records. Include an existing consumer's complete job; a fresh synthetic fixture can miss a supported caller's next operation. - **Data correctness:** identify units, rounding, ordering, identifiers, and trust boundaries. For money, use exact minor units and checked arithmetic; never hide malformed input behind a favorable result. - **Verification:** map each user-visible requirement to a real command or runtime action. Include a negative case that would catch the likely defect, with its expected exit and output channels. State when a build only compiles and when a test actually executes. Name how the executable will be rebuilt from the intended source; a digest alone cannot establish source provenance, and cloned warm artifacts can retain stale timestamps. +Compare the frozen delivery policy with the changed components. Name supplemental checks and CI where the policy covers another component; its green result proves only its executed commands. Bind each acceptance claim to its source revision, target, host, and observable result. + Do not expand the assignment into speculative platforms, services, or catalog entries. Escalate a concrete decision only when choosing without it risks the wrong outcome. Do not translate unknowns into time estimates. ## Hand off an executable plan diff --git a/skills/allstack-qa/SKILL.md b/skills/allstack-qa/SKILL.md index 7b32a34..a0ff28d 100644 --- a/skills/allstack-qa/SKILL.md +++ b/skills/allstack-qa/SKILL.md @@ -23,6 +23,12 @@ When installed, `allstack check --expect-exit 4 --seconds 120 -- /absolute/progr Record the target version, command or interaction, input, observed result, expected result, and executed test count where available. A zero-case run, timeout, skipped browser, or unavailable credential remains unverified. Inspect the underlying result rather than trusting a model summary or green wrapper. +For HTTP failures, retain the status, a safely redacted response body, timestamp, and request identifier when supplied. Preserve the failed observation alongside later successes. Read the relevant runtime conditions before assigning a cause; missing response evidence leaves the cause unresolved. Keep credentials and identity payloads out of shared reports. + +For independent host proof, inspect actual tool calls and retained outputs: which skill file was read, which program ran, and which assertions executed. Separate source review, proposed commands, denied calls, and completed runtime checks. Another host's summary cannot fill a missing execution. + +When qualifying a deployment, bind observations to the running image and its source/configuration. Keep earlier deployment evidence separate. If the final publication adds only tests or docs, compare the actual runtime build inputs before claiming the deployed application still matches; final-commit CI and runtime proof establish different facts. + ## Deliver the outcome Report reproducible defects with steps and consequences, then map the completed user flow to actual evidence. Fix only within the user's requested scope; repeat affected checks after a repair. Save required reports under indexed repo docs, and use `allstack remember --kind qa --source --summary ` for durable context when available. Stop when the requested flow is proved or a concrete external prerequisite prevents completion. diff --git a/skills/allstack-review/SKILL.md b/skills/allstack-review/SKILL.md index 40e4a20..8807b43 100644 --- a/skills/allstack-review/SKILL.md +++ b/skills/allstack-review/SKILL.md @@ -14,6 +14,7 @@ Start with the changed user behavior and the invariants it relies on. Follow val - Check error paths, empty inputs, ordering, boundary values, and type/unit conversions. Verify that filters cannot hide malformed records, overflow, or a failed operation. - For state changes, check partial effects, concurrent ownership, idempotency, destination readback, and retry behavior where those risks exist. - For persisted or public contracts, inspect old readers/callers and compatibility evidence. A schema declaration does not prove actual output is valid. +- Check the producer's actual ordering and revision semantics against retained legacy records, and follow an existing consumer through its next required operation. Authentication admission alone does not prove that consumer's complete job. - Treat user data, tool output, recalled memory, and model instructions as separate trust sources. Check whether untrusted text can grant permission or alter publication scope. - Read the relevant tests. Identify whether they execute the real boundary, select positive cases, and can fail for the defect under discussion. Run focused checks when tools and authority allow; label unexecuted hypotheses. diff --git a/skills/allstack-ship/SKILL.md b/skills/allstack-ship/SKILL.md index 7c8544a..f66a2f3 100644 --- a/skills/allstack-ship/SKILL.md +++ b/skills/allstack-ship/SKILL.md @@ -21,6 +21,10 @@ Existing guarded configuration files can remain in the tracked base: snapshots b Use the returned run identity and epoch internally. `allstack verify` prepares the owned single commit and runs the frozen policy against that exact SHA. Confirm positive test execution and unchanged candidate/environment. Any edit or amended commit requires fresh affected verification; do not ship stale receipts. A pending host call must be reconciled, never bypassed. +Before `allstack ship`, preserve required epoch-local receipts and separate raw streams outside disposable worktrees, with paths, lengths, and digests verified. Shipping can remove accepted epoch worktrees. Put the public evidence summary under indexed repo docs and keep private transcripts in ignored storage; preserving only a receipt whose stream paths will disappear is insufficient. + Use `allstack ship` to publish the exact verified commit and read back the branch and PR head. Confirm Chronis task completion and the Prime source-linked outcome. If publication is uncertain, read the destination before repeating a side effect. Never bypass hooks, push trunk, weaken gates, or rewrite another author's history. +Read back required project memory through its engine before archiving an isolated coordinator. Use queries that match each recorded summary; an unrelated lexical search can omit a node that still exists. Preserve tracker/memory state and verified receipts before removing owned checkouts and idle build directories. Recheck PR liveness before any later commit; a merged branch requires newly authorized continuation. + Report the application behavior delivered, PR, exact commit, actual executed checks, and material limits. Required documentation stays indexed under repo docs. Keep private process receipts out of source. The user should not need to manage run IDs or epochs to ask for shipping. diff --git a/templates/native-skills.md b/templates/native-skills.md index 8f209d8..510d7aa 100644 --- a/templates/native-skills.md +++ b/templates/native-skills.md @@ -27,3 +27,5 @@ For a complete durable delivery, initialize with a project policy using the CLI/ AllSource owns managed run history and receipts, Chronis owns delivery claims and completion, and Prime supplies source-linked context. The coordinator retains exact-commit verification, immutable handoffs, bounded process ownership, and destination reconciliation. Its plan/build host calls remain more restricted than independent native execution. Shipping ends at a verified remote commit and open PR; merge and deployment remain separate actions. Keep the coordinator at its original path through verification and publication. Native adoption verifies the candidate in that checkout; an epoch worktree may retain the base. Preserve existing project tracker/memory state and archive path-bound history before removing an isolated coordinator. + +Archive required epoch-local receipts and raw streams before `allstack ship`, which can remove accepted epoch worktrees. Verify copied byte lengths and digests. Report host execution from actual calls, preserve failed HTTP observations with safely redacted responses, and distinguish running-image proof from final-commit checks. Existing consumer compatibility and changed-component checks belong in acceptance even when a frozen policy covers another component. diff --git a/tests/vendor_integrity.rs b/tests/vendor_integrity.rs new file mode 100644 index 0000000..3a5e3fc --- /dev/null +++ b/tests/vendor_integrity.rs @@ -0,0 +1,51 @@ +use std::{collections::BTreeSet, fs, path::Path}; + +use allstack::infrastructure::digest; + +fn collect(root: &Path, path: &Path, files: &mut BTreeSet) { + for entry in fs::read_dir(path).unwrap() { + let path = entry.unwrap().path(); + let metadata = fs::symlink_metadata(&path).unwrap(); + assert!(!metadata.file_type().is_symlink()); + if metadata.is_dir() { + collect(root, &path, files); + } else { + assert!(metadata.is_file()); + files.insert(path.strip_prefix(root).unwrap().to_str().unwrap().into()); + } + } +} + +#[test] +fn parquet_patch_preserves_published_source_except_compression_requirements() { + let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("vendor/parquet"); + let metadata: serde_json::Value = + serde_json::from_slice(&fs::read(root.join("upstream.json")).unwrap()).unwrap(); + assert_eq!(metadata["crate"], "parquet"); + assert_eq!(metadata["version"], "59.3.0"); + let expected = metadata["files"].as_object().unwrap(); + let mut actual = BTreeSet::new(); + collect(&root, &root, &mut actual); + actual.remove("upstream.json"); + assert_eq!(actual, expected.keys().cloned().collect()); + for (name, checksum) in expected { + let bytes = fs::read(root.join(name)).unwrap(); + let original = match name.as_str() { + "Cargo.toml" => { + let text = String::from_utf8(bytes).unwrap(); + assert_eq!(text.matches("version = \"0.14.3\"").count(), 2); + text.replace("version = \"0.14.3\"", "version = \"0.13\"") + .into_bytes() + } + "Cargo.toml.orig" => { + let text = String::from_utf8(bytes).unwrap(); + let updated = "zstd = { version = \"0.14.3\""; + assert_eq!(text.matches(updated).count(), 2); + text.replace(updated, "zstd = { version = \"0.13\"") + .into_bytes() + } + _ => bytes, + }; + assert_eq!(digest(&original), checksum.as_str().unwrap(), "{name}"); + } +} diff --git a/vendor/parquet/CONTRIBUTING.md b/vendor/parquet/CONTRIBUTING.md new file mode 100644 index 0000000..febe57e --- /dev/null +++ b/vendor/parquet/CONTRIBUTING.md @@ -0,0 +1,63 @@ + + +## Build + +Run `cargo build` or `cargo build --release` to build in release mode. +Some features take advantage of SSE4.2 instructions, which can be +enabled by adding `RUSTFLAGS="-C target-feature=+sse4.2"` before the +`cargo build` command. + +## Test + +Run `cargo test` for unit tests. To also run tests related to the binaries, use `cargo test --features cli`. + +## Binaries + +The following binaries are provided (use `cargo install --features cli` to install them): + +- **parquet-schema** for printing Parquet file schema and metadata. + `Usage: parquet-schema `, where `file-path` is the path to a Parquet file. Use `-v/--verbose` flag + to print full metadata or schema only (when not specified only schema will be printed). + +- **parquet-read** for reading records from a Parquet file. + `Usage: parquet-read [num-records]`, where `file-path` is the path to a Parquet file, + and `num-records` is the number of records to read from a file (when not specified all records will + be printed). Use `-j/--json` to print records in JSON lines format. + +- **parquet-rowcount** for reporting the number of records in one or more Parquet files. + `Usage: parquet-rowcount ...`, where `...` is a space separated list of one or more + files to read. + +If you see `Library not loaded` error, please make sure `LD_LIBRARY_PATH` is set properly: + +``` +export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$(rustc --print sysroot)/lib +``` + +## Benchmarks + +Run `cargo bench` for benchmarks. + +## Docs + +To build documentation, run `cargo doc --no-deps --all-features`. +To compile and view in the browser, run `cargo doc --no-deps --all-features --open`. + +Before submitting a pull request, run `cargo fmt --all` to format the change. diff --git a/vendor/parquet/Cargo.lock b/vendor/parquet/Cargo.lock new file mode 100644 index 0000000..ab18653 --- /dev/null +++ b/vendor/parquet/Cargo.lock @@ -0,0 +1,3373 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "const-random", + "getrandom 0.3.4", + "once_cell", + "version_check", + "zerocopy", +] + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "alloc-no-stdlib" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" + +[[package]] +name = "alloc-stdlib" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94fb8275041c72129eb51b7d0322c29b8387a0386127718b096429201a5d6ece" +dependencies = [ + "alloc-no-stdlib", +] + +[[package]] +name = "alloca" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4" +dependencies = [ + "cc", +] + +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[package]] +name = "approx" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cab112f0a86d568ea0e627cc1d6be74a1e9cd55214684db5561995f6dad897c6" +dependencies = [ + "num-traits", +] + +[[package]] +name = "arrow" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c14b3d39f306bc28fd639d59f06e17a0f377d0021e1b7e9054e4d6fedc98774" +dependencies = [ + "arrow-arith", + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-data", + "arrow-ipc", + "arrow-json", + "arrow-ord", + "arrow-row", + "arrow-schema", + "arrow-select", + "arrow-string", + "half", + "rand 0.9.4", +] + +[[package]] +name = "arrow-arith" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce2961626677665b2195eb59242af4c7befe7b8737ca2050295389362380104e" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "num-traits", +] + +[[package]] +name = "arrow-array" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e5f6adeffdf587d7a31db5d2266189624b526730cd3627f9ff9fedae97ad584" +dependencies = [ + "ahash", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "half", + "hashbrown 0.17.1", + "num-complex", + "num-integer", + "num-traits", +] + +[[package]] +name = "arrow-buffer" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "097d193003ce7995d5d087089069ec2a6e0187faf5a6f8c9f38af2645d987182" +dependencies = [ + "bytes", + "half", + "num-bigint", + "num-traits", +] + +[[package]] +name = "arrow-cast" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "635c9c635668ad26adf76cce8fb276c4be7cf06e63bd516de7da514f9680ee53" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-ord", + "arrow-schema", + "arrow-select", + "atoi", + "base64 0.23.0", + "chrono", + "comfy-table", + "half", + "lexical-core", + "num-traits", + "ryu", +] + +[[package]] +name = "arrow-csv" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c2ebf8d631e79b02c16cf5ae860561272c26024ec88fce389a56aaddd558e86" +dependencies = [ + "arrow-array", + "arrow-cast", + "arrow-schema", + "chrono", + "csv", + "csv-core", + "regex", +] + +[[package]] +name = "arrow-data" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ba2f832eaeca24b8f26143dba750e42ee4ab51cf7d65e701ca9607cfda9f358" +dependencies = [ + "arrow-buffer", + "arrow-schema", + "half", + "num-integer", + "num-traits", +] + +[[package]] +name = "arrow-ipc" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dcc41681ea80f521df14c36725b74d4c60702c47f0793af2be469c04527e2599" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "flatbuffers", +] + +[[package]] +name = "arrow-json" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2f57d7a81969f24ccf80809587b76c09897e6f829d2d65a5976bfb3218851f1" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-ord", + "arrow-schema", + "arrow-select", + "chrono", + "half", + "indexmap", + "itoa", + "lexical-core", + "memchr", + "num-traits", + "ryu", + "serde_core", + "serde_json", + "simdutf8", +] + +[[package]] +name = "arrow-ord" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c900759f3bd8354fd4196bc4403eee846894dc2adf66b4225472006a0bf18c5" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", +] + +[[package]] +name = "arrow-row" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4c6425032e28266e3fc4ff680805e57e670d6ea92473043f3e65b7ed6ac79f2" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "half", +] + +[[package]] +name = "arrow-schema" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10fab8d4563491417ba801fab29d205104d20d4bdf37bda6cd1cf425cff598cd" +dependencies = [ + "serde_core", + "serde_json", +] + +[[package]] +name = "arrow-select" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc58569193c2525915f3cc6310edba3792f1200f65d6e9ed330aa33e691493b8" +dependencies = [ + "ahash", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "num-traits", +] + +[[package]] +name = "arrow-string" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e0813f3c35c1cfea65e14c20a953440f7783c088b7ad2d0db162ccdeefcec14" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "memchr", + "num-traits", + "regex", + "regex-syntax", +] + +[[package]] +name = "async-trait" +version = "0.1.89" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "atoi" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f28d99ec8bfea296261ca1af174f24225171fea9664ba9003cbebee704810528" +dependencies = [ + "num-traits", +] + +[[package]] +name = "atomic-waker" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + +[[package]] +name = "autocfg" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + +[[package]] +name = "base64" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b25655df2c3cdd83c5e5b293b88acd880332b2ddadd7c30ac43144fdc0033da9" + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[package]] +name = "brotli" +version = "8.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cc91aac060a7a1e25823bdccbfb6af1875b88f17c6daac97894eed8207166b3" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", + "brotli-decompressor", +] + +[[package]] +name = "brotli-decompressor" +version = "5.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "874bb8112abecc98cbd6d81ea4fa7e94fb9449648c93cc89aa40c81c24d7de03" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", +] + +[[package]] +name = "bumpalo" +version = "3.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + +[[package]] +name = "bytes" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" + +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + +[[package]] +name = "cc" +version = "1.2.61" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d16d90359e986641506914ba71350897565610e87ce0ad9e6f28569db3dd5c6d" +dependencies = [ + "find-msvc-tools", + "jobserver", + "libc", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + +[[package]] +name = "chacha20" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" +dependencies = [ + "cfg-if", + "cpufeatures", + "rand_core 0.10.1", +] + +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "iana-time-zone", + "num-traits", + "serde", + "windows-link", +] + +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + +[[package]] +name = "clap" +version = "4.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d91e0c145792ef73a6ad36d27c75ac09f1832222a3c209689d90f534685ee5b7" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f09628afdcc538b57f3c6341e9c8e9970f18e4a481690a64974d7023bd33548b" +dependencies = [ + "anstyle", + "clap_lex", +] + +[[package]] +name = "clap_derive" +version = "4.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d012d2b9d65aca7f18f4d9878a045bc17899bba951561ba5ec3c2ba1eed9a061" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "comfy-table" +version = "7.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "958c5d6ecf1f214b4c2bbbbf6ab9523a864bd136dcf71a7e8904799acfe1ad47" +dependencies = [ + "unicode-segmentation", + "unicode-width", +] + +[[package]] +name = "console" +version = "0.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d64e8af5551369d19cf50138de61f1c42074ab970f74e99be916646777f8fc87" +dependencies = [ + "encode_unicode", + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "const-random" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359" +dependencies = [ + "const-random-macro", +] + +[[package]] +name = "const-random-macro" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e" +dependencies = [ + "getrandom 0.2.17", + "once_cell", + "tiny-keccak", +] + +[[package]] +name = "core-foundation" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "criterion" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "950046b2aa2492f9a536f5f4f9a3de7b9e2476e575e05bd6c333371add4d98f3" +dependencies = [ + "alloca", + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "futures", + "itertools 0.13.0", + "num-traits", + "oorandom", + "page_size", + "regex", + "serde", + "serde_json", + "tinytemplate", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea" +dependencies = [ + "cast", + "itertools 0.13.0", +] + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "csv" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52cd9d68cf7efc6ddfaaee42e7288d3a99d613d4b50f76ce9827ae0c6e14f938" +dependencies = [ + "csv-core", + "itoa", + "ryu", + "serde_core", +] + +[[package]] +name = "csv-core" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "704a3c26996a80471189265814dbc2c257598b96b8a7feae2d31ace646bb9782" +dependencies = [ + "memchr", +] + +[[package]] +name = "displaydoc" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "either" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" + +[[package]] +name = "encode_unicode" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "fastrand" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "flatbuffers" +version = "25.12.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35f6839d7b3b98adde531effaf34f0c2badc6f4735d26fe74709d8e513a96ef3" +dependencies = [ + "bitflags", + "rustc_version", +] + +[[package]] +name = "flate2" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +dependencies = [ + "crc32fast", + "miniz_oxide", + "zlib-rs", +] + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "form_urlencoded" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" +dependencies = [ + "percent-encoding", +] + +[[package]] +name = "futures" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a88cf1f829d945f548cf8fec32c61b1f202b6d93b45848602fc02af4b12ad218" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "262590f4fe6afeb0bc83be1daa64e52657fe185690a958af7f3ad0e92085c5ae" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" + +[[package]] +name = "futures-executor" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6754879cc9f2c66f88c6e5c35344bb0bdb0708b0352b1201815667c7eabc7458" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4577ecaa3c4f96589d473f679a71b596316f6641bc350038b962a5daf0085d7a" + +[[package]] +name = "futures-macro" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6d3cde68c518367be28956066ddfef33813991b77a55005a69dae04bf3b10b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "futures-sink" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e34418ac499d6305c2fb5ad0ed2f6ac998c5f8ca209b4510f7f94242c647e307" + +[[package]] +name = "futures-task" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109" + +[[package]] +name = "futures-util" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "geo-traits" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e7c353d12a704ccfab1ba8bfb1a7fe6cb18b665bf89d37f4f7890edcd260206" +dependencies = [ + "geo-types", +] + +[[package]] +name = "geo-types" +version = "0.7.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94776032c45f950d30a13af6113c2ad5625316c9abfbccee4dd5a6695f8fe0f5" +dependencies = [ + "approx", + "num-traits", + "serde", +] + +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "wasi", + "wasm-bindgen", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "r-efi 5.3.0", + "wasip2", + "wasm-bindgen", +] + +[[package]] +name = "getrandom" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0de51e6874e94e7bf76d726fc5d13ba782deca734ff60d5bb2fb2607c7406555" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", + "rand_core 0.10.1", + "wasip2", + "wasip3", +] + +[[package]] +name = "h2" +version = "0.4.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9f37a958b41b3b19ee2707c06439c0e9e547e847223eb791ecb0cb821c65e27" +dependencies = [ + "atomic-waker", + "bytes", + "fnv", + "futures-core", + "futures-sink", + "http", + "indexmap", + "slab", + "tokio", + "tokio-util", + "tracing", +] + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "num-traits", + "rand 0.9.4", + "rand_distr", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "http" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "918d3568bebf352712bc2ef3d46a8bcf1a75b373be6539de198e9105cbbf9ce0" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "http-body" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca2a8f2913ee65f60facd6a5905613afaa448497a0230cc41ce022d93290bc2c" +dependencies = [ + "bytes", + "http", +] + +[[package]] +name = "http-body-util" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "pin-project-lite", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + +[[package]] +name = "humantime" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "135b12329e5e3ce057a9f972339ea52bc954fe1e9358ef27f95e89716fbc5424" + +[[package]] +name = "hyper" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6299f016b246a94207e63da54dbe807655bf9e00044f73ded42c3ac5305fbcca" +dependencies = [ + "atomic-waker", + "bytes", + "futures-channel", + "futures-core", + "h2", + "http", + "http-body", + "httparse", + "itoa", + "pin-project-lite", + "smallvec", + "tokio", + "want", +] + +[[package]] +name = "hyper-rustls" +version = "0.27.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" +dependencies = [ + "http", + "hyper", + "hyper-util", + "rustls", + "rustls-native-certs", + "tokio", + "tokio-rustls", + "tower-service", +] + +[[package]] +name = "hyper-util" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-channel", + "futures-util", + "http", + "http-body", + "hyper", + "ipnet", + "libc", + "percent-encoding", + "pin-project-lite", + "socket2", + "tokio", + "tower-service", + "tracing", +] + +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + +[[package]] +name = "icu_collections" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +dependencies = [ + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" + +[[package]] +name = "icu_properties" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" + +[[package]] +name = "icu_provider" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + +[[package]] +name = "id-arena" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" + +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", + "serde", + "serde_core", +] + +[[package]] +name = "insta" +version = "1.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "86f0f8fee8c926415c58d6ae43a08523a26faccb2323f5e6b644fe7dd4ef6b82" +dependencies = [ + "console", + "once_cell", + "similar", + "tempfile", +] + +[[package]] +name = "ipnet" +version = "2.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" + +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "jobserver" +version = "0.1.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33" +dependencies = [ + "getrandom 0.3.4", + "libc", +] + +[[package]] +name = "js-sys" +version = "0.3.98" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67df7112613f8bfd9150013a0314e196f4800d3201ae742489d999db2f979f08" +dependencies = [ + "cfg-if", + "futures-util", + "once_cell", + "wasm-bindgen", +] + +[[package]] +name = "leb128fmt" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" + +[[package]] +name = "lexical-core" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d8d125a277f807e55a77304455eb7b1cb52f2b18c143b60e766c120bd64a594" +dependencies = [ + "lexical-parse-float", + "lexical-parse-integer", + "lexical-util", + "lexical-write-float", + "lexical-write-integer", +] + +[[package]] +name = "lexical-parse-float" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52a9f232fbd6f550bc0137dcb5f99ab674071ac2d690ac69704593cb4abbea56" +dependencies = [ + "lexical-parse-integer", + "lexical-util", +] + +[[package]] +name = "lexical-parse-integer" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a7a039f8fb9c19c996cd7b2fcce303c1b2874fe1aca544edc85c4a5f8489b34" +dependencies = [ + "lexical-util", +] + +[[package]] +name = "lexical-util" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2604dd126bb14f13fb5d1bd6a66155079cb9fa655b37f875b3a742c705dbed17" + +[[package]] +name = "lexical-write-float" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50c438c87c013188d415fbabbb1dceb44249ab81664efbd31b14ae55dabb6361" +dependencies = [ + "lexical-util", + "lexical-write-integer", +] + +[[package]] +name = "lexical-write-integer" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "409851a618475d2d5796377cad353802345cba92c867d9fbcde9cf4eac4e14df" +dependencies = [ + "lexical-util", +] + +[[package]] +name = "libc" +version = "0.2.189" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[package]] +name = "libm" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + +[[package]] +name = "litemap" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" + +[[package]] +name = "lru-slab" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" + +[[package]] +name = "lz4_flex" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ecbdfe44b1bd960b68170b417450a628c43f7cf56bb3c5317e61cb230ee7f226" +dependencies = [ + "twox-hash", +] + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + +[[package]] +name = "mio" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50b7e5b27aa02a74bac8c3f23f448f8d87ff11f92d3aac1a6ed369ee08cc56c1" +dependencies = [ + "libc", + "wasi", + "windows-sys 0.61.2", +] + +[[package]] +name = "ntapi" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3b335231dfd352ffb0f8017f3b6027a4917f7df785ea2143d8af2adc66980ae" +dependencies = [ + "winapi", +] + +[[package]] +name = "num-bigint" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", + "libm", +] + +[[package]] +name = "num_enum" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d0bca838442ec211fa11de3a8b0e0e8f3a4522575b5c4c06ed722e005036f26" +dependencies = [ + "num_enum_derive", + "rustversion", +] + +[[package]] +name = "num_enum_derive" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "680998035259dcfcafe653688bf2aa6d3e2dc05e98be6ab46afb089dc84f1df8" +dependencies = [ + "proc-macro-crate", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "objc2-core-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags", +] + +[[package]] +name = "objc2-io-kit" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33fafba39597d6dc1fb709123dfa8289d39406734be322956a69f0931c73bb15" +dependencies = [ + "libc", + "objc2-core-foundation", +] + +[[package]] +name = "object_store" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "622acbc9100d3c10e2ee15804b0caa40e55c933d5aa53814cd520805b7958a49" +dependencies = [ + "async-trait", + "base64 0.22.1", + "bytes", + "chrono", + "form_urlencoded", + "futures-channel", + "futures-core", + "futures-util", + "http", + "http-body-util", + "httparse", + "humantime", + "hyper", + "itertools 0.14.0", + "parking_lot", + "percent-encoding", + "quick-xml", + "rand 0.10.1", + "reqwest", + "ring", + "serde", + "serde_json", + "serde_urlencoded", + "thiserror 2.0.18", + "tokio", + "tracing", + "url", + "walkdir", + "wasm-bindgen-futures", + "web-time", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + +[[package]] +name = "openssl-probe" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" + +[[package]] +name = "page_size" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da" +dependencies = [ + "libc", + "winapi", +] + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "parquet" +version = "59.3.0" +dependencies = [ + "ahash", + "arrow", + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-csv", + "arrow-data", + "arrow-ipc", + "arrow-schema", + "arrow-select", + "base64 0.23.0", + "brotli", + "bytes", + "chrono", + "clap", + "crc32fast", + "criterion", + "flate2", + "futures", + "half", + "hashbrown 0.17.1", + "insta", + "lz4_flex", + "num-bigint", + "num-integer", + "num-traits", + "object_store", + "parquet-geospatial", + "parquet-variant", + "parquet-variant-compute", + "parquet-variant-json", + "rand 0.9.4", + "ring", + "seq-macro", + "serde", + "serde_json", + "simdutf8", + "snap", + "sysinfo", + "tempfile", + "tokio", + "twox-hash", + "zstd", +] + +[[package]] +name = "parquet-geospatial" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "634827365851180ca26bf54c76e970dc1c14c075684c0580ee97d8d0c953faf1" +dependencies = [ + "arrow-schema", + "geo-traits", + "serde", + "serde_json", + "wkb", +] + +[[package]] +name = "parquet-variant" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3950502c5d383bc98ac83424c50ad06214cfb8a8554ef6f19413e71894079288" +dependencies = [ + "arrow", + "arrow-schema", + "chrono", + "half", + "indexmap", + "num-traits", + "simdutf8", + "uuid", +] + +[[package]] +name = "parquet-variant-compute" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28c4fc004d60a76d727fe4a383b48e8c4d629384b9d9f0cc359164ec2cbd3326" +dependencies = [ + "arrow", + "arrow-schema", + "chrono", + "half", + "indexmap", + "parquet-variant", + "parquet-variant-json", + "serde_json", + "uuid", +] + +[[package]] +name = "parquet-variant-json" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "216b1aecab3ee2e8f429ade4dfef24e233d008ee91ce745fb273c39fc5936b85" +dependencies = [ + "arrow-schema", + "base64 0.23.0", + "chrono", + "parquet-variant", + "serde_json", + "uuid", +] + +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + +[[package]] +name = "potential_utf" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + +[[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quick-xml" +version = "0.39.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdcc8dd4e2f670d309a5f0e83fe36dfdc05af317008fea29144da1a2ac858e5e" +dependencies = [ + "memchr", + "serde", +] + +[[package]] +name = "quinn" +version = "0.11.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c1a41e437b6bbd489372cd4971de128e85c855f56c57f283d20ff016cf7c0a8" +dependencies = [ + "bytes", + "cfg_aliases", + "pin-project-lite", + "quinn-proto", + "quinn-udp", + "rustc-hash", + "rustls", + "socket2", + "thiserror 2.0.18", + "tokio", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-proto" +version = "0.11.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fcb935c5bec503c2f0e306bdd3e58bb9029dcb14fa8d9ac76e3a5256ac0763e" +dependencies = [ + "bytes", + "getrandom 0.3.4", + "lru-slab", + "rand 0.9.4", + "ring", + "rustc-hash", + "rustls", + "rustls-pki-types", + "slab", + "thiserror 2.0.18", + "tinyvec", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-udp" +version = "0.5.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd" +dependencies = [ + "cfg_aliases", + "libc", + "once_cell", + "socket2", + "tracing", + "windows-sys 0.60.2", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rand" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" +dependencies = [ + "rand_chacha", + "rand_core 0.9.5", +] + +[[package]] +name = "rand" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2e8e8bcc7961af1fdac401278c6a831614941f6164ee3bf4ce61b7edb162207" +dependencies = [ + "chacha20", + "getrandom 0.4.2", + "rand_core 0.10.1", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + +[[package]] +name = "rand_distr" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463" +dependencies = [ + "num-traits", + "rand 0.9.4", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fcfdb36bda0c880c5931cdc7a2bcdc8ba4556847b9d912bca70bc94708711ad" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reqwest" +version = "0.12.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-core", + "futures-util", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls", + "rustls-native-certs", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tokio-util", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams", + "web-sys", +] + +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + +[[package]] +name = "rustc-hash" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" + +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustls" +version = "0.23.40" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef86cd5876211988985292b91c96a8f2d298df24e75989a43a3c73f2d4d8168b" +dependencies = [ + "once_cell", + "ring", + "rustls-pki-types", + "rustls-webpki", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-native-certs" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "612460d5f7bea540c490b2b6395d8e34a953e52b491accd6c86c8164c5932a63" +dependencies = [ + "openssl-probe", + "rustls-pki-types", + "schannel", + "security-framework", +] + +[[package]] +name = "rustls-pki-types" +version = "1.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" +dependencies = [ + "web-time", + "zeroize", +] + +[[package]] +name = "rustls-webpki" +version = "0.103.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + +[[package]] +name = "rustversion" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" + +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "schannel" +version = "0.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "security-framework" +version = "3.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" +dependencies = [ + "bitflags", + "core-foundation", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "semver" +version = "1.0.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" + +[[package]] +name = "seq-macro" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "serde_json" +version = "1.0.151" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_urlencoded" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "simd-adler32" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" + +[[package]] +name = "simdutf8" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" + +[[package]] +name = "similar" +version = "2.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbbb5d9659141646ae647b42fe094daf6c6192d1620870b449d9557f748b2daa" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" + +[[package]] +name = "snap" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "199905e6153d6405f9728fe44daace35f8f837bbf830bb6e85fbd5828709a886" + +[[package]] +name = "socket2" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "sync_wrapper" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +dependencies = [ + "futures-core", +] + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "sysinfo" +version = "0.39.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2071df9448915b71c4fe6d25deaf1c22f12bd234f01540b77312bb8e41361e6" +dependencies = [ + "libc", + "memchr", + "ntapi", + "objc2-core-foundation", + "objc2-io-kit", + "windows", +] + +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.2", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl 2.0.18", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "tiny-keccak" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" +dependencies = [ + "crunchy", +] + +[[package]] +name = "tinystr" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +dependencies = [ + "displaydoc", + "zerovec", +] + +[[package]] +name = "tinytemplate" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] + +[[package]] +name = "tinyvec" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + +[[package]] +name = "tokio" +version = "1.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" +dependencies = [ + "bytes", + "libc", + "mio", + "pin-project-lite", + "socket2", + "tokio-macros", + "windows-sys 0.61.2", +] + +[[package]] +name = "tokio-macros" +version = "2.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "tokio-rustls" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" +dependencies = [ + "rustls", + "tokio", +] + +[[package]] +name = "tokio-util" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.25.11+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b59c4d22ed448339746c59b905d24568fcbb3ab65a500494f7b8c3e97739f2b" +dependencies = [ + "indexmap", + "toml_datetime", + "toml_parser", + "winnow", +] + +[[package]] +name = "toml_parser" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" +dependencies = [ + "winnow", +] + +[[package]] +name = "tower" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" +dependencies = [ + "futures-core", + "futures-util", + "pin-project-lite", + "sync_wrapper", + "tokio", + "tower-layer", + "tower-service", +] + +[[package]] +name = "tower-http" +version = "0.6.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68d6fdd9f81c2819c9a8b0e0cd91660e7746a8e6ea2ba7c6b2b057985f6bcb51" +dependencies = [ + "bitflags", + "bytes", + "futures-util", + "http", + "http-body", + "pin-project-lite", + "tower", + "tower-layer", + "tower-service", + "url", +] + +[[package]] +name = "tower-layer" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" + +[[package]] +name = "tower-service" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", +] + +[[package]] +name = "try-lock" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" + +[[package]] +name = "twox-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8464ec13c3691491391d9fce00f6416c9a48e46972f72d7865688be2080192c9" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unicode-segmentation" +version = "1.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9629274872b2bfaf8d66f5f15725007f635594914870f65218920345aa11aa8c" + +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + +[[package]] +name = "unicode-xid" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" + +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + +[[package]] +name = "url" +version = "2.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", + "serde", +] + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "uuid" +version = "1.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf3923a6f5c4c6382e0b653c4117f48d631ea17f38ed86e2a828e6f7412f5239" +dependencies = [ + "getrandom 0.4.2", + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + +[[package]] +name = "want" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e" +dependencies = [ + "try-lock", +] + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasip2" +version = "1.0.3+wasi-0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20064672db26d7cdc89c7798c48a0fdfac8213434a1186e5ef29fd560ae223d6" +dependencies = [ + "wit-bindgen 0.57.1", +] + +[[package]] +name = "wasip3" +version = "0.4.0+wasi-0.3.0-rc-2026-01-06" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" +dependencies = [ + "wit-bindgen 0.51.0", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.121" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49ace1d07c165b0864824eee619580c4689389afa9dc9ed3a4c75040d82e6790" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-futures" +version = "0.4.71" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96492d0d3ffba25305a7dc88720d250b1401d7edca02cc3bcd50633b424673b8" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.121" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e68e6f4afd367a562002c05637acb8578ff2dea1943df76afb9e83d177c8578" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.121" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d95a9ec35c64b2a7cb35d3fead40c4238d0940c86d107136999567a4703259f2" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn 2.0.119", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.121" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4e0100b01e9f0d03189a92b96772a1fb998639d981193d7dbab487302513441" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "wasm-encoder" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "990065f2fe63003fe337b932cfb5e3b80e0b4d0f5ff650e6985b1048f62c8319" +dependencies = [ + "leb128fmt", + "wasmparser", +] + +[[package]] +name = "wasm-metadata" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb0e353e6a2fbdc176932bbaab493762eb1255a7900fe0fea1a2f96c296cc909" +dependencies = [ + "anyhow", + "indexmap", + "wasm-encoder", + "wasmparser", +] + +[[package]] +name = "wasm-streams" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + +[[package]] +name = "wasmparser" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" +dependencies = [ + "bitflags", + "hashbrown 0.15.5", + "indexmap", + "semver", +] + +[[package]] +name = "web-sys" +version = "0.3.98" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b572dff8bcf38bad0fa19729c89bb5748b2b9b1d8be70cf90df697e3a8f32aa" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + +[[package]] +name = "windows" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "527fadee13e0c05939a6a05d5bd6eec6cd2e3dbd648b9f8e447c6518133d8580" +dependencies = [ + "windows-collections", + "windows-core", + "windows-future", + "windows-numerics", +] + +[[package]] +name = "windows-collections" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b2d95af1a8a14a3c7367e1ed4fc9c20e0a26e79551b1454d72583c97cc6610" +dependencies = [ + "windows-core", +] + +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-future" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1d6f90251fe18a279739e78025bd6ddc52a7e22f921070ccdc67dde84c605cb" +dependencies = [ + "windows-core", + "windows-link", + "windows-threading", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-numerics" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e2e40844ac143cdb44aead537bbf727de9b044e107a0f1220392177d15b0f26" +dependencies = [ + "windows-core", + "windows-link", +] + +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-sys" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +dependencies = [ + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-sys" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" +dependencies = [ + "windows-targets 0.53.5", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm 0.52.6", + "windows_aarch64_msvc 0.52.6", + "windows_i686_gnu 0.52.6", + "windows_i686_gnullvm 0.52.6", + "windows_i686_msvc 0.52.6", + "windows_x86_64_gnu 0.52.6", + "windows_x86_64_gnullvm 0.52.6", + "windows_x86_64_msvc 0.52.6", +] + +[[package]] +name = "windows-targets" +version = "0.53.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" +dependencies = [ + "windows-link", + "windows_aarch64_gnullvm 0.53.1", + "windows_aarch64_msvc 0.53.1", + "windows_i686_gnu 0.53.1", + "windows_i686_gnullvm 0.53.1", + "windows_i686_msvc 0.53.1", + "windows_x86_64_gnu 0.53.1", + "windows_x86_64_gnullvm 0.53.1", + "windows_x86_64_msvc 0.53.1", +] + +[[package]] +name = "windows-threading" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3949bd5b99cafdf1c7ca86b43ca564028dfe27d66958f2470940f73d86d75b37" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_i686_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" + +[[package]] +name = "winnow" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ee1708bef14716a11bae175f579062d4554d95be2c6829f518df847b7b3fdd0" +dependencies = [ + "memchr", +] + +[[package]] +name = "wit-bindgen" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7249219f66ced02969388cf2bb044a09756a083d0fab1e566056b04d9fbcaa5" +dependencies = [ + "wit-bindgen-rust-macro", +] + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "wit-bindgen-core" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea61de684c3ea68cb082b7a88508a8b27fcc8b797d738bfc99a82facf1d752dc" +dependencies = [ + "anyhow", + "heck", + "wit-parser", +] + +[[package]] +name = "wit-bindgen-rust" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7c566e0f4b284dd6561c786d9cb0142da491f46a9fbed79ea69cdad5db17f21" +dependencies = [ + "anyhow", + "heck", + "indexmap", + "prettyplease", + "syn 2.0.119", + "wasm-metadata", + "wit-bindgen-core", + "wit-component", +] + +[[package]] +name = "wit-bindgen-rust-macro" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c0f9bfd77e6a48eccf51359e3ae77140a7f50b1e2ebfe62422d8afdaffab17a" +dependencies = [ + "anyhow", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", + "wit-bindgen-core", + "wit-bindgen-rust", +] + +[[package]] +name = "wit-component" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" +dependencies = [ + "anyhow", + "bitflags", + "indexmap", + "log", + "serde", + "serde_derive", + "serde_json", + "wasm-encoder", + "wasm-metadata", + "wasmparser", + "wit-parser", +] + +[[package]] +name = "wit-parser" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ecc8ac4bc1dc3381b7f59c34f00b67e18f910c2c0f50015669dde7def656a736" +dependencies = [ + "anyhow", + "id-arena", + "indexmap", + "log", + "semver", + "serde", + "serde_derive", + "serde_json", + "unicode-xid", + "wasmparser", +] + +[[package]] +name = "wkb" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a120b336c7ad17749026d50427c23d838ecb50cd64aaea6254b5030152f890a9" +dependencies = [ + "byteorder", + "geo-traits", + "num_enum", + "thiserror 1.0.69", +] + +[[package]] +name = "writeable" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" + +[[package]] +name = "yoke" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "abe8c5fda708d9ca3df187cae8bfb9ceda00dd96231bed36e445a1a48e66f9ca" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", + "synstructure", +] + +[[package]] +name = "zerocopy" +version = "0.8.48" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eed437bf9d6692032087e337407a86f04cd8d6a16a37199ed57949d415bd68e9" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.48" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70e3cd084b1788766f53af483dd21f93881ff30d7320490ec3ef7526d203bad4" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zerofrom" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69faa1f2a1ea75661980b013019ed6687ed0e83d069bc1114e2cc74c6c04c4df" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", + "synstructure", +] + +[[package]] +name = "zeroize" +version = "1.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" + +[[package]] +name = "zerotrie" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zlib-rs" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3be3d40e40a133f9c916ee3f9f4fa2d9d63435b5fbe1bfc6d9dae0aa0ada1513" + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" + +[[package]] +name = "zstd" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "7.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.0.16+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748" +dependencies = [ + "cc", + "pkg-config", +] diff --git a/vendor/parquet/Cargo.toml b/vendor/parquet/Cargo.toml new file mode 100644 index 0000000..d03b681 --- /dev/null +++ b/vendor/parquet/Cargo.toml @@ -0,0 +1,651 @@ +# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO +# +# When uploading crates to the registry Cargo will automatically +# "normalize" Cargo.toml files for maximal compatibility +# with all versions of Cargo and also rewrite `path` dependencies +# to registry (e.g., crates.io) dependencies. +# +# If you are reading this file be aware that the original Cargo.toml +# will likely look very different (and much more reasonable). +# See Cargo.toml.orig for the original contents. + +[package] +edition = "2024" +rust-version = "1.85" +name = "parquet" +version = "59.3.0" +authors = ["Apache Arrow "] +build = false +autolib = false +autobins = false +autoexamples = false +autotests = false +autobenches = false +description = "Apache Parquet implementation in Rust" +homepage = "https://github.com/apache/arrow-rs" +readme = "README.md" +keywords = [ + "arrow", + "parquet", + "hadoop", +] +license = "Apache-2.0" +repository = "https://github.com/apache/arrow-rs" +resolver = "2" + +[package.metadata.docs.rs] +all-features = true + +[features] +arrow = [ + "base64", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "arrow-ipc", +] +arrow_canonical_extension_types = ["arrow-schema?/canonical_extension_types"] +async = [ + "futures", + "tokio", +] +cli = [ + "json", + "base64", + "clap", + "arrow-csv", + "serde", +] +crc = ["dep:crc32fast"] +default = [ + "arrow", + "snap", + "brotli", + "flate2-zlib-rs", + "lz4", + "zstd", + "base64", + "simdutf8", +] +encryption = ["dep:ring"] +experimental = ["variant_experimental"] +flate2-rust_backend = ["flate2/rust_backend"] +flate2-rust_backened = ["flate2-rust_backend"] +flate2-zlib-rs = ["flate2/zlib-rs"] +geospatial = ["parquet-geospatial"] +json = [ + "serde_json", + "base64", +] +lz4 = ["lz4_flex"] +object_store = [ + "dep:object_store", + "async", +] +simdutf8 = ["dep:simdutf8"] +test_common = ["arrow/test_utils"] +variant_experimental = [ + "arrow", + "parquet-variant", + "parquet-variant-json", + "parquet-variant-compute", +] +zstd = ["dep:zstd"] + +[lib] +name = "parquet" +path = "src/lib.rs" +bench = false + +[[bin]] +name = "parquet-concat" +path = "src/bin/parquet-concat.rs" +required-features = ["cli"] + +[[bin]] +name = "parquet-fromcsv" +path = "src/bin/parquet-fromcsv.rs" +required-features = [ + "arrow", + "cli", + "snap", + "brotli", + "flate2", + "lz4", + "zstd", +] + +[[bin]] +name = "parquet-index" +path = "src/bin/parquet-index.rs" +required-features = ["cli"] + +[[bin]] +name = "parquet-layout" +path = "src/bin/parquet-layout.rs" +required-features = ["cli"] + +[[bin]] +name = "parquet-read" +path = "src/bin/parquet-read.rs" +required-features = ["cli"] + +[[bin]] +name = "parquet-rewrite" +path = "src/bin/parquet-rewrite.rs" +required-features = [ + "arrow", + "cli", +] + +[[bin]] +name = "parquet-rowcount" +path = "src/bin/parquet-rowcount.rs" +required-features = ["cli"] + +[[bin]] +name = "parquet-schema" +path = "src/bin/parquet-schema.rs" +required-features = ["cli"] + +[[bin]] +name = "parquet-show-bloom-filter" +path = "src/bin/parquet-show-bloom-filter.rs" +required-features = ["cli"] + +[[example]] +name = "external_metadata" +path = "examples/external_metadata.rs" +required-features = [ + "arrow", + "async", +] + +[[example]] +name = "object_store" +path = "examples/object_store.rs" +required-features = [ + "arrow", + "async", +] + +[[example]] +name = "read_parquet" +path = "examples/read_parquet.rs" +required-features = ["arrow"] + +[[example]] +name = "read_with_rowgroup" +path = "examples/read_with_rowgroup.rs" +required-features = [ + "arrow", + "async", +] + +[[example]] +name = "write_parquet" +path = "examples/write_parquet.rs" +required-features = ["cli"] + +[[test]] +name = "arrow_reader" +path = "tests/arrow_reader/mod.rs" +required-features = ["arrow"] + +[[test]] +name = "arrow_writer" +path = "tests/arrow_writer.rs" +required-features = ["arrow"] + +[[test]] +name = "arrow_writer_layout" +path = "tests/arrow_writer_layout.rs" +required-features = ["arrow"] + +[[test]] +name = "encryption" +path = "tests/encryption/mod.rs" +required-features = ["arrow"] + +[[test]] +name = "geospatial" +path = "tests/geospatial.rs" + +[[test]] +name = "variant_integration" +path = "tests/variant_integration.rs" +required-features = [ + "arrow", + "variant_experimental", + "serde", +] + +[[bench]] +name = "arrow_reader" +path = "benches/arrow_reader.rs" +harness = false +required-features = [ + "arrow", + "test_common", + "experimental", +] + +[[bench]] +name = "arrow_reader_clickbench" +path = "benches/arrow_reader_clickbench.rs" +harness = false +required-features = [ + "arrow", + "async", +] + +[[bench]] +name = "arrow_reader_peak_memory" +path = "benches/arrow_reader_peak_memory.rs" +harness = false +required-features = [ + "arrow", + "test_common", + "experimental", +] + +[[bench]] +name = "arrow_reader_row_filter" +path = "benches/arrow_reader_row_filter.rs" +harness = false +required-features = [ + "arrow", + "async", +] + +[[bench]] +name = "arrow_reader_row_selection_policy" +path = "benches/arrow_reader_row_selection_policy.rs" +harness = false +required-features = [ + "arrow", + "async", +] + +[[bench]] +name = "arrow_statistics" +path = "benches/arrow_statistics.rs" +harness = false +required-features = ["arrow"] + +[[bench]] +name = "arrow_writer" +path = "benches/arrow_writer.rs" +harness = false +required-features = ["arrow"] + +[[bench]] +name = "bloom_filter" +path = "benches/bloom_filter.rs" +harness = false + +[[bench]] +name = "compression" +path = "benches/compression.rs" +harness = false +required-features = [ + "experimental", + "default", +] + +[[bench]] +name = "encoding" +path = "benches/encoding.rs" +harness = false +required-features = [ + "experimental", + "default", +] + +[[bench]] +name = "metadata" +path = "benches/metadata.rs" +harness = false + +[[bench]] +name = "parquet_round_trip" +path = "benches/parquet_round_trip.rs" +harness = false +required-features = ["arrow"] + +[[bench]] +name = "push_decoder" +path = "benches/push_decoder.rs" +harness = false +required-features = ["arrow"] + +[[bench]] +name = "row_group_index_reader" +path = "benches/row_group_index_reader.rs" +harness = false +required-features = ["arrow"] + +[[bench]] +name = "row_selection_cursor" +path = "benches/row_selection_cursor.rs" +harness = false +required-features = ["arrow"] + +[[bench]] +name = "row_selector" +path = "benches/row_selector.rs" +harness = false +required-features = ["arrow"] + +[[bench]] +name = "row_selector_boolean_buffer" +path = "benches/row_selector_boolean_buffer.rs" +harness = false +required-features = ["arrow"] + +[[bench]] +name = "writer_overhead" +path = "benches/writer_overhead.rs" +harness = false +required-features = ["arrow"] + +[dependencies.arrow-array] +version = "59.3.0" +optional = true + +[dependencies.arrow-buffer] +version = "59.3.0" +optional = true + +[dependencies.arrow-csv] +version = "59.3.0" +optional = true + +[dependencies.arrow-data] +version = "59.3.0" +optional = true + +[dependencies.arrow-ipc] +version = "59.3.0" +optional = true + +[dependencies.arrow-schema] +version = "59.3.0" +optional = true + +[dependencies.arrow-select] +version = "59.3.0" +optional = true + +[dependencies.base64] +version = "0.23" +features = ["std"] +optional = true +default-features = false + +[dependencies.brotli] +version = "8.0" +features = ["std"] +optional = true +default-features = false + +[dependencies.bytes] +version = "1.1" +features = ["std"] +default-features = false + +[dependencies.chrono] +version = "0.4.40" +features = ["clock"] +default-features = false + +[dependencies.clap] +version = "4.1" +features = [ + "std", + "derive", + "env", + "help", + "error-context", + "usage", +] +optional = true +default-features = false + +[dependencies.crc32fast] +version = "1.4.2" +optional = true +default-features = false + +[dependencies.flate2] +version = "1.1" +optional = true +default-features = false + +[dependencies.futures] +version = "0.3" +features = ["std"] +optional = true +default-features = false + +[dependencies.half] +version = "2.1" +features = ["num-traits"] +default-features = false + +[dependencies.hashbrown] +version = "0.17" +default-features = false + +[dependencies.lz4_flex] +version = "0.14" +features = [ + "std", + "frame", +] +optional = true +default-features = false + +[dependencies.num-bigint] +version = "0.5" +default-features = false + +[dependencies.num-integer] +version = "0.1.46" +features = ["std"] +default-features = false + +[dependencies.num-traits] +version = "0.2.19" +features = ["std"] +default-features = false + +[dependencies.object_store] +version = "0.13.2" +features = ["tokio"] +optional = true +default-features = false + +[dependencies.parquet-geospatial] +version = "59.3.0" +optional = true + +[dependencies.parquet-variant] +version = "59.3.0" +optional = true + +[dependencies.parquet-variant-compute] +version = "59.3.0" +optional = true + +[dependencies.parquet-variant-json] +version = "59.3.0" +optional = true + +[dependencies.ring] +version = "0.17" +features = ["std"] +optional = true +default-features = false + +[dependencies.seq-macro] +version = "0.3" +default-features = false + +[dependencies.serde] +version = "1.0" +features = ["derive"] +optional = true +default-features = false + +[dependencies.serde_json] +version = "1.0" +features = ["std"] +optional = true +default-features = false + +[dependencies.simdutf8] +version = "0.1.5" +optional = true +default-features = false + +[dependencies.snap] +version = "1.0" +optional = true +default-features = false + +[dependencies.tokio] +version = "1.0" +features = [ + "macros", + "rt", + "io-util", +] +optional = true +default-features = false + +[dependencies.twox-hash] +version = "2.0" +features = ["xxhash64"] +default-features = false + +[dependencies.zstd] +version = "0.14.3" +optional = true +default-features = false + +[dev-dependencies.arrow] +version = "59.3.0" +features = [ + "ipc", + "test_utils", + "prettyprint", + "json", +] +default-features = false + +[dev-dependencies.arrow-cast] +version = "59.3.0" + +[dev-dependencies.base64] +version = "0.23" +features = ["std"] +default-features = false + +[dev-dependencies.brotli] +version = "8.0" +features = ["std"] +default-features = false + +[dev-dependencies.criterion] +version = "0.8.0" +features = ["async_futures"] +default-features = false + +[dev-dependencies.flate2] +version = "1.0" +features = ["rust_backend"] +default-features = false + +[dev-dependencies.insta] +version = "1.46.3" +default-features = true + +[dev-dependencies.lz4_flex] +version = "0.14" +features = [ + "std", + "frame", +] +default-features = false + +[dev-dependencies.object_store] +version = "0.13.2" +features = [ + "azure", + "fs", +] +default-features = false + +[dev-dependencies.rand] +version = "0.9" +features = [ + "std", + "std_rng", + "thread_rng", +] +default-features = false + +[dev-dependencies.serde_json] +version = "1.0" +features = ["std"] +default-features = false + +[dev-dependencies.snap] +version = "1.0" +default-features = false + +[dev-dependencies.sysinfo] +version = "0.39.6" +features = ["system"] +default-features = false + +[dev-dependencies.tempfile] +version = "3.0" +default-features = false + +[dev-dependencies.tokio] +version = "1.0" +features = [ + "macros", + "rt-multi-thread", + "io-util", + "fs", + "sync", +] +default-features = false + +[dev-dependencies.zstd] +version = "0.14.3" +default-features = false + +[target.'cfg(not(target_arch = "wasm32"))'.dependencies.ahash] +version = "0.8" +features = ["runtime-rng"] +default-features = false + +[target.'cfg(target_arch = "wasm32")'.dependencies.ahash] +version = "0.8" +features = ["compile-time-rng"] +default-features = false + +[target.'cfg(target_arch = "wasm32")'.dependencies.ring] +version = "0.17" +features = [ + "wasm32_unknown_unknown_js", + "std", +] +optional = true +default-features = false diff --git a/vendor/parquet/Cargo.toml.orig b/vendor/parquet/Cargo.toml.orig new file mode 100644 index 0000000..1fa1f66 --- /dev/null +++ b/vendor/parquet/Cargo.toml.orig @@ -0,0 +1,316 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[package] +name = "parquet" +version = { workspace = true } +license = { workspace = true } +description = "Apache Parquet implementation in Rust" +homepage = { workspace = true } +repository = { workspace = true } +authors = { workspace = true } +keywords = ["arrow", "parquet", "hadoop"] +readme = "README.md" +edition = { workspace = true } +rust-version = { workspace = true } + +[target.'cfg(target_arch = "wasm32")'.dependencies] +ahash = { version = "0.8", default-features = false, features = ["compile-time-rng"] } +# See https://github.com/briansmith/ring/issues/918#issuecomment-2077788925 +ring = { version = "0.17", default-features = false, features = ["wasm32_unknown_unknown_js", "std"], optional = true } + +[target.'cfg(not(target_arch = "wasm32"))'.dependencies] +ahash = { version = "0.8", default-features = false, features = ["runtime-rng"] } + +[dependencies] +arrow-array = { workspace = true, optional = true } +arrow-buffer = { workspace = true, optional = true } +arrow-csv = { workspace = true, optional = true } +arrow-data = { workspace = true, optional = true } +arrow-schema = { workspace = true, optional = true } +arrow-select = { workspace = true, optional = true } +arrow-ipc = { workspace = true, optional = true } +parquet-geospatial = { workspace = true, optional = true } +parquet-variant = { workspace = true, optional = true } +parquet-variant-json = { workspace = true, optional = true } +parquet-variant-compute = { workspace = true, optional = true } + +object_store = { workspace = true, optional = true, features = ["tokio"] } + +bytes = { version = "1.1", default-features = false, features = ["std"] } +snap = { version = "1.0", default-features = false, optional = true } +brotli = { version = "8.0", default-features = false, features = ["std"], optional = true } +# To use `flate2` you must enable either the `flate2-zlib-rs` or `flate2-rust_backend` backends +flate2 = { version = "1.1", default-features = false, optional = true } +lz4_flex = { version = "0.14", default-features = false, features = ["std", "frame"], optional = true } +zstd = { version = "0.14.3", optional = true, default-features = false } +chrono = { workspace = true } +num-bigint = { version = "0.5", default-features = false } +num-integer = { version = "0.1.46", default-features = false, features = ["std"] } +num-traits = { version = "0.2.19", default-features = false, features = ["std"] } +base64 = { version = "0.23", default-features = false, features = ["std", ], optional = true } +clap = { version = "4.1", default-features = false, features = ["std", "derive", "env", "help", "error-context", "usage"], optional = true } +serde = { version = "1.0", default-features = false, features = ["derive"], optional = true } +serde_json = { version = "1.0", default-features = false, features = ["std"], optional = true } +seq-macro = { version = "0.3", default-features = false } +futures = { version = "0.3", default-features = false, features = ["std"], optional = true } +tokio = { version = "1.0", optional = true, default-features = false, features = ["macros", "rt", "io-util"] } +hashbrown = { version = "0.17", default-features = false } +twox-hash = { version = "2.0", default-features = false, features = ["xxhash64"] } +half = { version = "2.1", default-features = false, features = ["num-traits"] } +crc32fast = { version = "1.4.2", optional = true, default-features = false } +simdutf8 = { workspace = true , optional = true } +ring = { version = "0.17", default-features = false, features = ["std"], optional = true } + +[dev-dependencies] +base64 = { version = "0.23", default-features = false, features = ["std"] } +criterion = { workspace = true, default-features = false, features = ["async_futures"] } +snap = { version = "1.0", default-features = false } +tempfile = { version = "3.0", default-features = false } +insta = { workspace = true, default-features = true } +brotli = { version = "8.0", default-features = false, features = ["std"] } +flate2 = { version = "1.0", default-features = false, features = ["rust_backend"] } +lz4_flex = { version = "0.14", default-features = false, features = ["std", "frame"] } +zstd = { version = "0.14.3", default-features = false } +serde_json = { version = "1.0", features = ["std"], default-features = false } +arrow = { workspace = true, features = ["ipc", "test_utils", "prettyprint", "json"] } +arrow-cast = { workspace = true } +tokio = { version = "1.0", default-features = false, features = ["macros", "rt-multi-thread", "io-util", "fs", "sync"] } +rand = { version = "0.9", default-features = false, features = ["std", "std_rng", "thread_rng"] } +object_store = { workspace = true, features = ["azure", "fs"] } +sysinfo = { version = "0.39.6", default-features = false, features = ["system"] } + +[package.metadata.docs.rs] +all-features = true + +[features] +default = ["arrow", "snap", "brotli", "flate2-zlib-rs", "lz4", "zstd", "base64", "simdutf8"] +# Enable lz4 +lz4 = ["lz4_flex"] +# Enable arrow reader/writer APIs +arrow = ["base64", "arrow-array", "arrow-buffer", "arrow-data", "arrow-schema", "arrow-select", "arrow-ipc"] +# Enable support for arrow canonical extension types +arrow_canonical_extension_types = ["arrow-schema?/canonical_extension_types"] +# Enable CLI tools +cli = ["json", "base64", "clap", "arrow-csv", "serde"] +# Enable JSON APIs +json = ["serde_json", "base64"] +# Enable internal testing APIs +test_common = ["arrow/test_utils"] +# Experimental, unstable functionality primarily used for testing +experimental = ["variant_experimental"] +# Enable async APIs +async = ["futures", "tokio"] +# Enable object_store integration +# This feature will be removed in a future release. +# See https://github.com/apache/arrow-rs/issues/10308 +object_store = ["dep:object_store", "async"] +# Group Zstd dependencies +zstd = ["dep:zstd"] +# Verify 32-bit CRC checksum when decoding parquet pages +crc = ["dep:crc32fast"] +# Enable SIMD UTF-8 validation +simdutf8 = ["dep:simdutf8"] +# Enable Parquet modular encryption support +encryption = ["dep:ring"] +# Enable flate2's default Rust backend +flate2-rust_backend = ["flate2/rust_backend"] +# Enable flate2's zlib-rs backend +flate2-zlib-rs = ["flate2/zlib-rs"] +# Enable parquet variant support +variant_experimental = ["arrow", "parquet-variant", "parquet-variant-json", "parquet-variant-compute"] +# Enable geospatial support +geospatial = ["parquet-geospatial"] +# Alias for typo. TODO: Remove in future major version bump. +flate2-rust_backened = ["flate2-rust_backend"] + + +[[example]] +name = "external_metadata" +required-features = ["arrow", "async"] +path = "./examples/external_metadata.rs" + +[[example]] +name = "read_parquet" +required-features = ["arrow"] +path = "./examples/read_parquet.rs" + +[[example]] +name = "write_parquet" +required-features = ["cli"] +path = "./examples/write_parquet.rs" + +[[example]] +name = "read_with_rowgroup" +required-features = ["arrow", "async"] +path = "./examples/read_with_rowgroup.rs" + +[[example]] +name = "object_store" +required-features = ["arrow", "async"] +path = "./examples/object_store.rs" + +[[test]] +name = "arrow_writer_layout" +required-features = ["arrow"] + +[[test]] +name = "arrow_reader" +required-features = ["arrow"] +path = "./tests/arrow_reader/mod.rs" + +[[test]] +name = "arrow_writer" +required-features = ["arrow"] +path = "./tests/arrow_writer.rs" + +[[test]] +name = "encryption" +required-features = ["arrow"] +path = "./tests/encryption/mod.rs" + +[[test]] +name = "variant_integration" +required-features = ["arrow", "variant_experimental", "serde"] +path = "./tests/variant_integration.rs" + +[[bin]] +name = "parquet-read" +required-features = ["cli"] + +[[bin]] +name = "parquet-rewrite" +required-features = ["arrow", "cli"] + +[[bin]] +name = "parquet-schema" +required-features = ["cli"] + +[[bin]] +name = "parquet-rowcount" +required-features = ["cli"] + +[[bin]] +name = "parquet-concat" +required-features = ["cli"] + +[[bin]] +name = "parquet-fromcsv" +required-features = ["arrow", "cli", "snap", "brotli", "flate2", "lz4", "zstd"] + +[[bin]] +name = "parquet-show-bloom-filter" +required-features = ["cli"] + +[[bin]] +name = "parquet-layout" +required-features = ["cli"] + +[[bin]] +name = "parquet-index" +required-features = ["cli"] + +[[bench]] +name = "arrow_writer" +required-features = ["arrow"] +harness = false + +[[bench]] +name = "push_decoder" +required-features = ["arrow"] +harness = false + +[[bench]] +name = "writer_overhead" +required-features = ["arrow"] +harness = false + +[[bench]] +name = "arrow_reader" +required-features = ["arrow", "test_common", "experimental"] +harness = false + +[[bench]] +name = "arrow_statistics" +required-features = ["arrow"] +harness = false + +[[bench]] +name = "arrow_reader_row_filter" +required-features = ["arrow", "async"] +harness = false + +[[bench]] +name = "arrow_reader_row_selection_policy" +required-features = ["arrow", "async"] +harness = false + +[[bench]] +name = "arrow_reader_clickbench" +required-features = ["arrow", "async"] +harness = false + +[[bench]] +name = "compression" +required-features = ["experimental", "default"] +harness = false + +[[bench]] +name = "encoding" +required-features = ["experimental", "default"] +harness = false + +[[bench]] +name = "metadata" +harness = false + +[[bench]] +name = "parquet_round_trip" +required-features = ["arrow"] +harness = false + +[[bench]] +name = "row_selector" +harness = false +required-features = ["arrow"] + +[[bench]] +name = "row_selector_boolean_buffer" +harness = false +required-features = ["arrow"] + +[[bench]] +name = "row_group_index_reader" +required-features = ["arrow"] +harness = false + +[[bench]] +name = "row_selection_cursor" +harness = false +required-features = ["arrow"] + +[[bench]] +name = "bloom_filter" +harness = false + +[[bench]] +name = "arrow_reader_peak_memory" +required-features = ["arrow", "test_common", "experimental"] +harness = false + +[lib] +bench = false diff --git a/vendor/parquet/LICENSE.txt b/vendor/parquet/LICENSE.txt new file mode 100644 index 0000000..d645695 --- /dev/null +++ b/vendor/parquet/LICENSE.txt @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/vendor/parquet/NOTICE.txt b/vendor/parquet/NOTICE.txt new file mode 100644 index 0000000..3e997a0 --- /dev/null +++ b/vendor/parquet/NOTICE.txt @@ -0,0 +1,13 @@ +Apache Arrow +Copyright 2016-2026 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + +This product includes software from the chronoutil crate (MIT) + * Copyright (c) 2020-2022 Oliver Margetts + * https://github.com/olliemath/chronoutil + +This product includes software from the compact-thrift project (Apache 2.0) + * Copyright Jörn Horstmann + * https://github.com/jhorstmann/compact-thrift diff --git a/vendor/parquet/README.md b/vendor/parquet/README.md new file mode 100644 index 0000000..8fb4885 --- /dev/null +++ b/vendor/parquet/README.md @@ -0,0 +1,90 @@ + + +# Apache Parquet Official Native Rust Implementation + +[![crates.io](https://img.shields.io/crates/v/parquet.svg)](https://crates.io/crates/parquet) +[![docs.rs](https://img.shields.io/docsrs/parquet.svg)](https://docs.rs/parquet/latest/parquet/) + +This crate contains the official Native Rust implementation of [Apache Parquet](https://parquet.apache.org/), which is part of the [Apache Arrow](https://arrow.apache.org/) project. + +See the [API documentation](https://docs.rs/parquet/latest) for examples and the full API. + +The API documentation for most recent, unreleased code is available [here](https://arrow.apache.org/rust/parquet/index.html). + +## Rust Version Compatibility + +This crate is tested with the latest stable version of Rust. We do not currently test against other, older versions of the Rust compiler. + +## Versioning / Releases + +The `parquet` crate follows the [SemVer standard] defined by Cargo and works well +within the Rust crate ecosystem. See the [repository README] for more details on +the release schedule, version and deprecation policy. + +[semver standard]: https://doc.rust-lang.org/cargo/reference/semver.html +[repository readme]: https://github.com/apache/arrow-rs + +Note that for historical reasons, this crate uses versions with major numbers +greater than `0.x` (e.g. `19.0.0`), unlike many other crates in the Rust +ecosystem which spend extended time releasing versions `0.x` to signal planned +ongoing API changes. Minor arrow releases contain only compatible changes, while +major releases may contain breaking API changes. + +## Feature Flags + +The `parquet` crate provides the following features which may be enabled in your `Cargo.toml`: + +- `arrow` (default) - support for reading / writing [`arrow`] arrays to / from Parquet +- `async` - support `async` APIs for reading Parquet +- `json` - support for reading / writing `json` data to / from Parquet +- `brotli` (default) - support for Parquet using `brotli` compression +- `flate2` (default) - support for Parquet using `gzip` compression +- `lz4` (default) - support for Parquet using `lz4` compression +- `zstd` (default) - support for Parquet using `zstd` compression +- `snap` (default) - support for Parquet using `snappy` compression +- `cli` - parquet [CLI tools](https://github.com/apache/arrow-rs/tree/main/parquet/src/bin) +- `crc` - enables functionality to automatically verify checksums of each page (if present) when decoding +- `experimental` - Experimental APIs which may change, even between minor releases +- `simdutf8` (default) - Use the [`simdutf8`] crate for SIMD-accelerated UTF-8 validation +- `encryption` - support for reading / writing encrypted Parquet files +- `variant_experimental` - ⚠️ Experimental [Parquet Variant] support, which may change, even between minor releases. +- `geospatial` - ⚠️ Experimental geospatial support, which may change, even between minor releases. + +[`arrow`]: https://crates.io/crates/arrow +[`simdutf8`]: https://crates.io/crates/simdutf8 +[parquet variant]: https://github.com/apache/parquet-format/blob/master/VariantEncoding.md + +## Parquet Feature Status + +Please see the [Implementation Status Page] on the [Apache Parquet] website for +information on the status of this implementation. + +[implementation status page]: https://parquet.apache.org/docs/file-format/implementationstatus/ +[apache parquet]: https://parquet.apache.org/ + +## Security + +See the [Security Policy] for information on the security model and how to report vulnerabilities. + +[security policy]: https://github.com/apache/arrow-rs/blob/main/SECURITY.md + +## License + +Licensed under the Apache License, Version 2.0: . diff --git a/vendor/parquet/THRIFT.md b/vendor/parquet/THRIFT.md new file mode 100644 index 0000000..599b33f --- /dev/null +++ b/vendor/parquet/THRIFT.md @@ -0,0 +1,478 @@ + + +# Thrift serialization in the parquet crate + +For both performance and flexibility reasons, this crate uses custom Thrift parsers and +serialization mechanisms. For many of the objects defined by the Parquet specification macros +are used to generate the objects as well as the code to serialize them. But in certain instances +(performance bottlenecks, additions to the spec, etc.), it becomes necessary to implement the +serialization code manually. This document serves to document both the standard usage of the +Thrift macros, as well as how to implement custom encoders and decoders. + +## Thrift macros + +The Parquet specification utilizes Thrift enums, unions, and structs, defined by an Interface +Description Language (IDL). This IDL is usually parsed by a Thrift code generator to produce +language specific structures and serialization/deserialization code. This crate, however, uses +Rust macros to perform the same function. In addition to skipping creation of additional duplicate +structures, doing so allows for customizations that produce more performant code, as well as the +ability to pick and choose which fields to process. + +### Enums + +Thrift enums are the simplest structure, and are logically identical to Rust enums with unit +variants. The IDL description will look like + +``` +enum Type { + BOOLEAN = 0; + INT32 = 1; + INT64 = 2; + INT96 = 3; + FLOAT = 4; + DOUBLE = 5; + BYTE_ARRAY = 6; + FIXED_LEN_BYTE_ARRAY = 7; +} +``` + +The `thrift_enum` macro can be used in this instance. + +```rust +thrift_enum!( +enum Type { + BOOLEAN = 0; + INT32 = 1; + INT64 = 2; + INT96 = 3; + FLOAT = 4; + DOUBLE = 5; + BYTE_ARRAY = 6; + FIXED_LEN_BYTE_ARRAY = 7; +} +); +``` + +which will produce a public Rust enum + +```rust +pub enum Type { + BOOLEAN, + INT32, + INT64, + INT96, + FLOAT, + DOUBLE, + BYTE_ARRAY, + FIXED_LEN_BYTE_ARRAY, +} +``` + +All Rust `enum`s produced with this macro will have `pub` visibility. + +### Unions + +Thrift unions are a special kind of struct in which only a single field is populated. In this +regard they are much like Rust enums which can have a mix of unit and tuple variants. Because of +this flexibility, specifying unions is a little bit trickier. + +Often times a union will be defined for which all the variants are typed with empty structs. For +example the `TimeUnit` union used for `LogicalType`s. + +``` +struct MilliSeconds {} +struct MicroSeconds {} +struct NanoSeconds {} +union TimeUnit { + 1: MilliSeconds MILLIS + 2: MicroSeconds MICROS + 3: NanoSeconds NANOS +} +``` + +When serialized, these empty structs become a single `0` (to mark the end of the struct). As an +optimization, and to allow for a simpler interface, the `thrift_union_all_empty` macro can be used. + +```rust +thrift_union_all_empty!( +union TimeUnit { + 1: MilliSeconds MILLIS + 2: MicroSeconds MICROS + 3: NanoSeconds NANOS +} +); +``` + +This macro will ignore the types specified for each variant, and will produce the following Rust +`enum`: + +```rust +pub enum TimeUnit { + MILLIS, + MICROS, + NANOS, +} +``` + +For unions with mixed variant types, some modifications to the IDL are necessary. Take the +definition of `ColumnCryptoMetadata`. + +``` +struct EncryptionWithFooterKey { +} + +struct EncryptionWithColumnKey { + /** Column path in schema **/ + 1: required list path_in_schema + + /** Retrieval metadata of column encryption key **/ + 2: optional binary key_metadata +} + +union ColumnCryptoMetaData { + 1: EncryptionWithFooterKey ENCRYPTION_WITH_FOOTER_KEY + 2: EncryptionWithColumnKey ENCRYPTION_WITH_COLUMN_KEY +} +``` + +The `ENCRYPTION_WITH_FOOTER_KEY` variant is typed with an empty struct, while +`ENCRYPTION_WITH_COLUMN_KEY` has the type of a struct with fields. In this case, the `thrift_union` +macro is used. + +```rust +thrift_union!( +union ColumnCryptoMetaData { + 1: ENCRYPTION_WITH_FOOTER_KEY + 2: (EncryptionWithColumnKey) ENCRYPTION_WITH_COLUMN_KEY +} +); +``` + +Here, the type has been omitted for `ENCRYPTION_WITH_FOOTER_KEY` to indicate it should be a unit +variant, while the type for `ENCRYPTION_WITH_COLUMN_KEY` is enclosed in parens. The parens are +necessary to provide a semantic clue to the macro that the identifier is a type. The above will +produce the Rust enum + +```rust +pub enum ColumnCryptoMetaData { + ENCRYPTION_WITH_FOOTER_KEY, + ENCRYPTION_WITH_COLUMN_KEY(EncryptionWithColumnKey), +} +``` + +All Rust `enum`s produced with either macro will have `pub` visibility. `thrift_union` also allows +for lifetime annotations, but this capability is not currently utilized. + +### Structs + +The `thrift_struct` macro is used for structs. This macro is a little more flexible than the others +because it allows for the visibility to be specified, and also allows for lifetimes to be specified +for the defined structs as well as their fields. An example of this is the `SchemaElement` struct. +This is defined in this crate as + +```rust +thrift_struct!( +pub(crate) struct SchemaElement<'a> { + 1: optional Type r#type; + 2: optional i32 type_length; + 3: optional Repetition repetition_type; + 4: required string<'a> name; + 5: optional i32 num_children; + 6: optional ConvertedType converted_type; + 7: optional i32 scale + 8: optional i32 precision + 9: optional i32 field_id; + 10: optional LogicalType logical_type +} +); +``` + +Here the `string` field `name` is given a lifetime annotation, which is then propagated to the +struct definition. Without this annotation, the resultant field would be a `String` type, rather +than a string slice. The visibility of this struct (and all fields) will be `pub(crate)`. The +resultant Rust struct will be + +```rust +pub(crate) struct SchemaElement<'a> { + pub(crate) r#type: Type, // here we've changed the name `type` to `r#type` to avoid reserved words + pub(crate) type_length: i32, + pub(crate) repetition_type: Repetition, + pub(crate) name: &'a str, + ... +} +``` + +The lifetime annotations can also be added to list elements, as in + +```rust +thrift_struct!( +struct FileMetaData<'a> { + /** Version of this file **/ + 1: required i32 version + 2: required list<'a> schema; + 3: required i64 num_rows + 4: required list<'a> row_groups + 5: optional list key_value_metadata + 6: optional string created_by + 7: optional list column_orders; + 8: optional EncryptionAlgorithm encryption_algorithm + 9: optional binary footer_signing_key_metadata +} +); +``` + +Note that the lifetime annotation precedes the element type specification. + +## Serialization traits + +Serialization is performed via several Rust traits. On the deserialization, objects implement +the `ReadThrift` trait. This defines a `read_thrift` function that takes a +`ThriftCompactInputProtocol` I/O object as an argument. The `read_thrift` function performs +all steps necessary to deserialize the object from the input stream, and is usually produced by +one of the macros mentioned above. + +On the serialization side, the `WriteThrift` and `WriteThriftField` traits are used in conjunction +with a `ThriftCompactOutputProtocol` struct. As above, the Thrift macros produce the necessary +implementations needed to perform serialization. + +While the macros can be used in most circumstances, sometimes more control is needed. The following +sections provide information on how to provide custom implementations for the serialization +traits. + +### ReadThrift Customization + +Thrift enums are serialized as a single `i32` value. The process of reading an enum is straightforward: +read the enum discriminant, and then match on the possible values. For instance, reading the +`ConvertedType` enum becomes: + +```rust +impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for ConvertedType { + fn read_thrift(prot: &mut R) -> Result { + let val = prot.read_i32()?; + Ok(match val { + 0 => Self::UTF8, + 1 => Self::MAP, + 2 => Self::MAP_KEY_VALUE, + ... + 21 => Self::INTERVAL, + _ => return Err(general_err!("Unexpected ConvertedType {}", val)), + }) + } +} +``` + +The default behavior is to return an error when an unexpected field is encountered. One could, +however, provide an `Unknown` variant if forward compatibility is neeeded in the case of an +evolving enum. + +Deserializing structs is more involved, but still fairly easy. A thrift struct is serialized as +repeated `(field_id,field_type,field)` tuples. The `field_id` and `field_type` usually occupy a +single byte, followed by the Thrift encoded field. Because only 4 bits are available for the id, +encoders usually will instead use deltas from the preceding field. If the delta will exceed 15, +then the `field_id` nibble will be set to `0`, and the `field_id` will instead be encoded as a +varint, following the `field_type`. Fields will generally be read in a loop, with the `field_id` +and `field_type` read first, and then the `field_id` used to determine which field to read. +When a `field_id` of `0` is encountered, this marks the end of the struct and processing ceases. +Here is an example of the processing loop: + +```rust + let mut last_field_id = 0i16; + loop { + // read the field id and field type. break if we encounter `Stop` + let field_ident = prot.read_field_begin(last_field_id)?; + if field_ident.field_type == FieldType::Stop { + break; + } + // match on the field id + match field_ident.id { + 1 => { + let val = i32::read_thrift(&mut *prot)?; + num_values = Some(val); + } + 2 => { + let val = Encoding::read_thrift(&mut *prot)?; + encoding = Some(val); + } + 3 => { + let val = Encoding::read_thrift(&mut *prot)?; + definition_level_encoding = Some(val); + } + 4 => { + let val = Encoding::read_thrift(&mut *prot)?; + repetition_level_encoding = Some(val); + } + // Thrift structs are meant to be forward compatible, so do not error + // here. Instead, simply skip unknown fields. + _ => { + prot.skip(field_ident.field_type)?; + } + }; + // set the last seen field id to calculate the next field_id + last_field_id = field_ident.id; + } +``` + +Thrift unions are encoded as structs, but only a single field will be encoded. The loop above +can be eliminated, and only the `match` on the id performed. A subsequent call to +`read_field_begin` must return `Stop`, or an error should be returned. Here's an example from +the decoding of the `LogicalType` union: + +```rust + // read the discriminant, error if it is `0` + let field_ident = prot.read_field_begin(0)?; + if field_ident.field_type == FieldType::Stop { + return Err(general_err!("received empty union from remote LogicalType")); + } + let ret = match field_ident.id { + 1 => { + prot.skip_empty_struct()?; + Self::String + } + ... + _ => { + // LogicalType needs to be forward compatible, so we have defined an `_Unknown` + // variant for it. This can return an error if forward compatibility is not desired. + prot.skip(field_ident.field_type)?; + Self::_Unknown { + field_id: field_ident.id, + } + } + }; + // test to ensure there is only one field present + let field_ident = prot.read_field_begin(field_ident.id)?; + if field_ident.field_type != FieldType::Stop { + return Err(general_err!( + "Received multiple fields for union from remote LogicalType" + )); + } +``` + +### WriteThrift Customization + +On the serialization side, there are two traits to implement. The first, `WriteThrift`, is used +for actually serializing the object. The other, `WriteThriftField`, handles serializing objects +as struct fields. + +Serializing enums is as simple as writing the discriminant as an `i32`. For example, here is the +custom serialization code for `ConvertedType`: + +```rust +impl WriteThrift for ConvertedType { + const ELEMENT_TYPE: ElementType = ElementType::I32; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + // because we've added NONE, the variant values are off by 1, so correct that here + writer.write_i32(*self as i32 - 1) + } +} +``` + +Structs and unions are serialized by field. When performing the serialization, one needs to keep +track of the last field that has been written, as this is needed to calculate the delta in the +Thrift field header. For required fields this is not strictly necessary, but when writing +optional fields it is. A typical `write_thrift` implementation will look like: + +```rust + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + // required field f1 + self.f1.write_thrift_field(writer, 1, 0)?; // field_id == 1, last_field_id == 0 + // required field f2 + self.f2.write_thrift_field(writer, 2, 1)?; // field_id == 2, last_field_id == 1 + // final required field f3, we now save the last_field_id, which is returned by write_thrift_field + let mut last_field_id = self.f3.write_thrift_field(writer, 3, 2)?; // field_id == 3, last_field_id == 2 + + // optional field f4 + if let Some(val) = self.f4.as_ref() { + last_field_id = val.write_thrift_field(writer, 4, last_field_id)?; + } + // optional field f5 + if let Some(val) = self.f5.as_ref() { + last_field_id = val.write_thrift_field(writer, 5, last_field_id)?; + } + // write end of struct + writer.write_struct_end() + } +``` + +In most instances, the `WriteThriftField` implementation can be handled by the `write_thrift_field` +macro. The first argument is the unqualified name of an object that implements `WriteThrift`, and +the second is the field type (which will be `FieldType::Struct` for Thrift structs and unions, +and `FieldType::I32` for Thrift enums). + +```rust +write_thrift_field!(MyNewStruct, FieldType::Struct); +``` + +which expands to: + +```rust +impl WriteThriftField for MyNewStruct { + fn write_thrift_field( + &self, + writer: &mut ThriftCompactOutputProtocol, + field_id: i16, + last_field_id: i16, + ) -> Result { + writer.write_field_begin(FieldType::Struct, field_id, last_field_id)?; + self.write_thrift(writer)?; + Ok(field_id) + } +} +``` + +### Handling for lists + +Lists of serialized objects can usually be read using `parquet_thrift::read_thrift_vec` and written +using the `WriteThrift::write_thrift` implementation for vectors of objects that implement +`WriteThrift`. + +When reading a list, one first reads the list header which will provide the number of elements +that have been encoded, and then read elements one at a time. + +```rust + // read the list header + let list_ident = prot.read_list_begin()?; + // allocate vector with enough capacity + let mut page_locations = Vec::with_capacity(list_ident.size as usize); + // read elements + for _ in 0..list_ident.size { + page_locations.push(read_page_location(prot)?); + } +``` + +Writing is simply the reverse: write the list header, and then serialize the elements: + +```rust + // write the list header + writer.write_list_begin(ElementType::Struct, page_locations.len)?; + // write the elements + for i in 0..len { + page_locations[i].write_thrift(writer)?; + } +``` + +## More examples + +For more examples, the easiest thing to do is to [expand](https://github.com/dtolnay/cargo-expand) +the thrift macros. For instance, to see the implementations generated in the `basic` module, type: + +```sh +% cargo expand -p parquet --lib --all-features basic +``` diff --git a/vendor/parquet/benches/arrow_reader.rs b/vendor/parquet/benches/arrow_reader.rs new file mode 100644 index 0000000..6e91a04 --- /dev/null +++ b/vendor/parquet/benches/arrow_reader.rs @@ -0,0 +1,2618 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::Array; +use arrow::datatypes::DataType; +use arrow_schema::Field; +use criterion::measurement::WallTime; +use criterion::{BenchmarkGroup, Criterion, criterion_group, criterion_main}; +use half::f16; +use num_bigint::BigInt; +use num_traits::FromPrimitive; +use parquet::arrow::array_reader::{ + ListArrayReader, make_byte_array_reader, make_byte_view_array_reader, + make_fixed_len_byte_array_reader, +}; +use parquet::arrow::arrow_reader::DEFAULT_BATCH_SIZE; +use parquet::basic::Type; +use parquet::data_type::{ByteArray, FixedLenByteArrayType}; +use parquet::util::{DataPageBuilder, DataPageBuilderImpl, InMemoryPageIterator}; +use parquet::{ + arrow::array_reader::ArrayReader, + basic::Encoding, + column::page::PageIterator, + data_type::{BoolType, ByteArrayType, Int32Type, Int64Type}, + schema::types::{ColumnDescPtr, SchemaDescPtr}, +}; +use rand::distr::uniform::SampleUniform; +use rand::{Rng, SeedableRng, rngs::StdRng}; +use std::{collections::VecDeque, sync::Arc}; + +fn build_test_schema() -> SchemaDescPtr { + use parquet::schema::{parser::parse_message_type, types::SchemaDescriptor}; + let message_type = " + message test_schema { + REQUIRED INT32 mandatory_int32_leaf; + OPTIONAL INT32 optional_int32_leaf; + REQUIRED BYTE_ARRAY mandatory_string_leaf (UTF8); + OPTIONAL BYTE_ARRAY optional_string_leaf (UTF8); + REQUIRED INT64 mandatory_int64_leaf; + OPTIONAL INT64 optional_int64_leaf; + REQUIRED INT32 mandatory_decimal1_leaf (DECIMAL(8,2)); + OPTIONAL INT32 optional_decimal1_leaf (DECIMAL(8,2)); + REQUIRED INT64 mandatory_decimal2_leaf (DECIMAL(16,2)); + OPTIONAL INT64 optional_decimal2_leaf (DECIMAL(16,2)); + REQUIRED BYTE_ARRAY mandatory_decimal3_leaf (DECIMAL(16,2)); + OPTIONAL BYTE_ARRAY optional_decimal3_leaf (DECIMAL(16,2)); + REQUIRED FIXED_LEN_BYTE_ARRAY (16) mandatory_decimal4_leaf (DECIMAL(16,2)); + OPTIONAL FIXED_LEN_BYTE_ARRAY (16) optional_decimal4_leaf (DECIMAL(16,2)); + OPTIONAL GROUP string_list (LIST) { + repeated group list { + optional BYTE_ARRAY element (UTF8); + } + } + REQUIRED BYTE_ARRAY mandatory_binary_leaf; + OPTIONAL BYTE_ARRAY optional_binary_leaf; + REQUIRED FIXED_LEN_BYTE_ARRAY (2) mandatory_f16_leaf (Float16); + OPTIONAL FIXED_LEN_BYTE_ARRAY (2) optional_f16_leaf (Float16); + REQUIRED FIXED_LEN_BYTE_ARRAY (2) mandatory_flba2_leaf; + OPTIONAL FIXED_LEN_BYTE_ARRAY (2) optional_flba2_leaf; + REQUIRED FIXED_LEN_BYTE_ARRAY (4) mandatory_flba4_leaf; + OPTIONAL FIXED_LEN_BYTE_ARRAY (4) optional_flba4_leaf; + REQUIRED FIXED_LEN_BYTE_ARRAY (8) mandatory_flba8_leaf; + OPTIONAL FIXED_LEN_BYTE_ARRAY (8) optional_flba8_leaf; + REQUIRED FIXED_LEN_BYTE_ARRAY (16) mandatory_flba16_leaf; + OPTIONAL FIXED_LEN_BYTE_ARRAY (16) optional_flba16_leaf; + REQUIRED INT32 mandatory_uint8_leaf (INTEGER(8, false)); + OPTIONAL INT32 optional_uint8_leaf (INTEGER(8, false)); + REQUIRED INT32 mandatory_uint16_leaf (INTEGER(16, false)); + OPTIONAL INT32 optional_uint16_leaf (INTEGER(16, false)); + REQUIRED INT32 mandatory_uint32_leaf (INTEGER(32, false)); + OPTIONAL INT32 optional_uint32_leaf (INTEGER(32, false)); + REQUIRED INT32 mandatory_int8_leaf (INTEGER(8, true)); + OPTIONAL INT32 optional_int8_leaf (INTEGER(8, true)); + REQUIRED INT32 mandatory_int16_leaf (INTEGER(16, true)); + OPTIONAL INT32 optional_int16_leaf (INTEGER(16, true)); + REQUIRED INT64 mandatory_uint64_leaf (INTEGER(64, false)); + OPTIONAL INT64 optional_uint64_leaf (INTEGER(64, false)); + REQUIRED GROUP mandatory_struct_optional_int32_leaf { + OPTIONAL INT32 element; + } + OPTIONAL GROUP optional_struct_optional_int32_leaf { + OPTIONAL INT32 element; + } + OPTIONAL GROUP int32_list (LIST) { + repeated group list { + optional INT32 element; + } + } + OPTIONAL GROUP fixed32_list (LIST) { + repeated group list { + optional FIXED_LEN_BYTE_ARRAY(32) element; + } + } + REQUIRED BOOLEAN mandatory_bool_leaf; + OPTIONAL BOOLEAN optional_bool_leaf; + } + "; + parse_message_type(message_type) + .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t)))) + .unwrap() +} + +// test data params +const NUM_ROW_GROUPS: usize = 2; +const PAGES_PER_GROUP: usize = 4; +const VALUES_PER_PAGE: usize = 10_000; +const BATCH_SIZE: usize = 8192; +const MAX_LIST_LEN: usize = 10; +const EXPECTED_VALUE_COUNT: usize = NUM_ROW_GROUPS * PAGES_PER_GROUP * VALUES_PER_PAGE; + +pub fn seedable_rng() -> StdRng { + StdRng::seed_from_u64(42) +} + +// support byte array for float16 +fn build_encoded_f16_bytes_page_iterator( + column_desc: ColumnDescPtr, + null_density: f32, + encoding: Encoding, + min: f32, + max: f32, +) -> impl PageIterator + Clone +where + T: parquet::data_type::DataType, + T::T: From>, +{ + let max_def_level = column_desc.max_def_level(); + let max_rep_level = column_desc.max_rep_level(); + let rep_levels = vec![0; VALUES_PER_PAGE]; + let mut rng = seedable_rng(); + let mut pages: Vec> = Vec::new(); + for _i in 0..NUM_ROW_GROUPS { + let mut column_chunk_pages = Vec::new(); + for _j in 0..PAGES_PER_GROUP { + // generate page + let mut values = Vec::with_capacity(VALUES_PER_PAGE); + let mut def_levels = Vec::with_capacity(VALUES_PER_PAGE); + for _k in 0..VALUES_PER_PAGE { + let def_level = if rng.random::() < null_density { + max_def_level - 1 + } else { + max_def_level + }; + if def_level == max_def_level { + // create the Float16 value + let value = f16::from_f32(rng.random_range(min..max)); + // Float16 in parquet is stored little-endian + let bytes = match column_desc.physical_type() { + Type::FIXED_LEN_BYTE_ARRAY => { + // Float16 annotates FIXED_LEN_BYTE_ARRAY(2) + assert_eq!(column_desc.type_length(), 2); + value.to_le_bytes().to_vec() + } + _ => unimplemented!(), + }; + let value = T::T::from(bytes); + values.push(value); + } + def_levels.push(def_level); + } + let mut page_builder = + DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true); + page_builder.add_rep_levels(max_rep_level, &rep_levels); + page_builder.add_def_levels(max_def_level, &def_levels); + page_builder.add_values::(encoding, &values); + column_chunk_pages.push(page_builder.consume()); + } + pages.push(column_chunk_pages); + } + InMemoryPageIterator::new(pages) +} + +// support byte array for decimal +fn build_encoded_decimal_bytes_page_iterator( + column_desc: ColumnDescPtr, + null_density: f32, + encoding: Encoding, + min: i128, + max: i128, +) -> impl PageIterator + Clone +where + T: parquet::data_type::DataType, + T::T: From>, +{ + let max_def_level = column_desc.max_def_level(); + let max_rep_level = column_desc.max_rep_level(); + let rep_levels = vec![0; VALUES_PER_PAGE]; + let mut rng = seedable_rng(); + let mut pages: Vec> = Vec::new(); + for _i in 0..NUM_ROW_GROUPS { + let mut column_chunk_pages = Vec::new(); + for _j in 0..PAGES_PER_GROUP { + // generate page + let mut values = Vec::with_capacity(VALUES_PER_PAGE); + let mut def_levels = Vec::with_capacity(VALUES_PER_PAGE); + for _k in 0..VALUES_PER_PAGE { + let def_level = if rng.random::() < null_density { + max_def_level - 1 + } else { + max_def_level + }; + if def_level == max_def_level { + // create the decimal value + let value = rng.random_range(min..max); + // decimal of parquet use the big-endian to store + let bytes = match column_desc.physical_type() { + Type::BYTE_ARRAY => { + // byte array use the unfixed size + let big_int = BigInt::from(value); + big_int.to_signed_bytes_be() + } + Type::FIXED_LEN_BYTE_ARRAY => { + assert_eq!(column_desc.type_length(), 16); + // fixed length byte array use the fixed size + // the size is 16 + value.to_be_bytes().to_vec() + } + _ => unimplemented!(), + }; + let value = T::T::from(bytes); + values.push(value); + } + def_levels.push(def_level); + } + let mut page_builder = + DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true); + page_builder.add_rep_levels(max_rep_level, &rep_levels); + page_builder.add_def_levels(max_def_level, &def_levels); + page_builder.add_values::(encoding, &values); + column_chunk_pages.push(page_builder.consume()); + } + pages.push(column_chunk_pages); + } + InMemoryPageIterator::new(pages) +} + +// support for fixed_len_byte_arrays +fn build_encoded_flba_bytes_page_iterator( + column_desc: ColumnDescPtr, + null_density: f32, + encoding: Encoding, +) -> impl PageIterator + Clone { + let max_def_level = column_desc.max_def_level(); + let max_rep_level = column_desc.max_rep_level(); + let rep_levels = vec![0; VALUES_PER_PAGE]; + let mut rng = seedable_rng(); + let mut pages: Vec> = Vec::new(); + for _i in 0..NUM_ROW_GROUPS { + let mut column_chunk_pages = Vec::new(); + for _j in 0..PAGES_PER_GROUP { + // generate page + let mut values = Vec::with_capacity(VALUES_PER_PAGE); + let mut def_levels = Vec::with_capacity(VALUES_PER_PAGE); + for _k in 0..VALUES_PER_PAGE { + let def_level = if rng.random::() < null_density { + max_def_level - 1 + } else { + max_def_level + }; + if def_level == max_def_level { + // create the FLBA(BYTE_LENGTH) value + let value = (0..BYTE_LENGTH).map(|_| rng.random()).collect::>(); + let value = + ::T::from(value); + values.push(value); + } + def_levels.push(def_level); + } + let mut page_builder = + DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true); + page_builder.add_rep_levels(max_rep_level, &rep_levels); + page_builder.add_def_levels(max_def_level, &def_levels); + page_builder.add_values::(encoding, &values); + column_chunk_pages.push(page_builder.consume()); + } + pages.push(column_chunk_pages); + } + InMemoryPageIterator::new(pages) +} + +fn build_encoded_primitive_page_iterator( + column_desc: ColumnDescPtr, + null_density: f32, + encoding: Encoding, + min: usize, + max: usize, +) -> impl PageIterator + Clone +where + T: parquet::data_type::DataType, + T::T: SampleUniform + FromPrimitive, +{ + let max_def_level = column_desc.max_def_level(); + let max_rep_level = column_desc.max_rep_level(); + let rep_levels = vec![0; VALUES_PER_PAGE]; + let mut rng = seedable_rng(); + let mut pages: Vec> = Vec::new(); + for _i in 0..NUM_ROW_GROUPS { + let mut column_chunk_pages = Vec::new(); + for _j in 0..PAGES_PER_GROUP { + // generate page + let mut values = Vec::with_capacity(VALUES_PER_PAGE); + let mut def_levels = Vec::with_capacity(VALUES_PER_PAGE); + for _k in 0..VALUES_PER_PAGE { + let def_level = if rng.random::() < null_density { + max_def_level - 1 + } else { + max_def_level + }; + if def_level == max_def_level { + let value = FromPrimitive::from_usize(rng.random_range(min..max)).unwrap(); + values.push(value); + } + def_levels.push(def_level); + } + let mut page_builder = + DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true); + page_builder.add_rep_levels(max_rep_level, &rep_levels); + page_builder.add_def_levels(max_def_level, &def_levels); + page_builder.add_values::(encoding, &values); + column_chunk_pages.push(page_builder.consume()); + } + pages.push(column_chunk_pages); + } + + InMemoryPageIterator::new(pages) +} + +fn build_encoded_bool_page_iterator( + column_desc: ColumnDescPtr, + null_density: f32, + encoding: Encoding, +) -> impl PageIterator + Clone { + let max_def_level = column_desc.max_def_level(); + let max_rep_level = column_desc.max_rep_level(); + let rep_levels = vec![0; VALUES_PER_PAGE]; + let mut rng = seedable_rng(); + let mut pages: Vec> = Vec::new(); + for _i in 0..NUM_ROW_GROUPS { + let mut column_chunk_pages = Vec::new(); + for _j in 0..PAGES_PER_GROUP { + // generate page + let mut values = Vec::with_capacity(VALUES_PER_PAGE); + let mut def_levels = Vec::with_capacity(VALUES_PER_PAGE); + for _k in 0..VALUES_PER_PAGE { + let def_level = if rng.random::() < null_density { + max_def_level - 1 + } else { + max_def_level + }; + if def_level == max_def_level { + let value = rng.random_bool(0.5); + values.push(value); + } + def_levels.push(def_level); + } + let mut page_builder = + DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true); + page_builder.add_rep_levels(max_rep_level, &rep_levels); + page_builder.add_def_levels(max_def_level, &def_levels); + page_builder.add_values::(encoding, &values); + column_chunk_pages.push(page_builder.consume()); + } + pages.push(column_chunk_pages); + } + + InMemoryPageIterator::new(pages) +} + +fn build_delta_encoded_incr_primitive_page_iterator( + column_desc: ColumnDescPtr, + null_density: f32, + increment: usize, + stepped: bool, +) -> impl PageIterator + Clone +where + T: parquet::data_type::DataType, + T::T: SampleUniform + FromPrimitive, +{ + let max_def_level = column_desc.max_def_level(); + let max_rep_level = column_desc.max_rep_level(); + let rep_levels = vec![0; VALUES_PER_PAGE]; + let mut rng = seedable_rng(); + let mut pages: Vec> = Vec::new(); + let mut running_val: usize = 1; + for _i in 0..NUM_ROW_GROUPS { + let mut column_chunk_pages = Vec::new(); + for _j in 0..PAGES_PER_GROUP { + // generate page + let mut values = Vec::with_capacity(VALUES_PER_PAGE); + let mut def_levels = Vec::with_capacity(VALUES_PER_PAGE); + for k in 0..VALUES_PER_PAGE { + let def_level = if rng.random::() < null_density { + max_def_level - 1 + } else { + max_def_level + }; + if def_level == max_def_level { + let value = FromPrimitive::from_usize(running_val).unwrap(); + running_val = if !stepped || k % 2 == 1 { + running_val + increment + } else { + running_val + }; + values.push(value); + } + def_levels.push(def_level); + } + let mut page_builder = + DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true); + page_builder.add_rep_levels(max_rep_level, &rep_levels); + page_builder.add_def_levels(max_def_level, &def_levels); + page_builder.add_values::(Encoding::DELTA_BINARY_PACKED, &values); + column_chunk_pages.push(page_builder.consume()); + } + pages.push(column_chunk_pages); + } + + InMemoryPageIterator::new(pages) +} + +fn build_dictionary_encoded_primitive_page_iterator( + column_desc: ColumnDescPtr, + null_density: f32, +) -> impl PageIterator + Clone +where + T: parquet::data_type::DataType, + T::T: SampleUniform + FromPrimitive + Copy, +{ + use parquet::encoding::{DictEncoder, Encoder}; + let max_def_level = column_desc.max_def_level(); + let max_rep_level = column_desc.max_rep_level(); + let rep_levels = vec![0; VALUES_PER_PAGE]; + // generate 1% unique values + const NUM_UNIQUE_VALUES: usize = VALUES_PER_PAGE / 100; + let unique_values: Vec = (0..NUM_UNIQUE_VALUES) + .map(|x| FromPrimitive::from_usize(x + 1).unwrap()) + .collect::>(); + let mut rng = seedable_rng(); + let mut pages: Vec> = Vec::new(); + for _i in 0..NUM_ROW_GROUPS { + let mut column_chunk_pages = VecDeque::new(); + let mut dict_encoder = DictEncoder::::new(column_desc.clone()); + // add data pages + for _j in 0..PAGES_PER_GROUP { + // generate page + let mut values = Vec::with_capacity(VALUES_PER_PAGE); + let mut def_levels = Vec::with_capacity(VALUES_PER_PAGE); + for _k in 0..VALUES_PER_PAGE { + let def_level = if rng.random::() < null_density { + max_def_level - 1 + } else { + max_def_level + }; + if def_level == max_def_level { + // select random value from list of unique values + let value = unique_values[rng.random_range(0..NUM_UNIQUE_VALUES)]; + values.push(value); + } + def_levels.push(def_level); + } + let mut page_builder = + DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true); + page_builder.add_rep_levels(max_rep_level, &rep_levels); + page_builder.add_def_levels(max_def_level, &def_levels); + let _ = dict_encoder.put(&values); + let indices = dict_encoder + .write_indices() + .expect("write_indices() should be OK"); + page_builder.add_indices(indices); + column_chunk_pages.push_back(page_builder.consume()); + } + // add dictionary page + let dict = dict_encoder + .write_dict() + .expect("write_dict() should be OK"); + let dict_page = parquet::column::page::Page::DictionaryPage { + buf: dict, + num_values: dict_encoder.num_entries() as u32, + encoding: Encoding::RLE_DICTIONARY, + is_sorted: false, + }; + column_chunk_pages.push_front(dict_page); + pages.push(column_chunk_pages.into()); + } + + InMemoryPageIterator::new(pages) +} + +fn build_plain_encoded_byte_array_page_iterator_inner( + column_desc: ColumnDescPtr, + null_density: f32, + short_string: bool, +) -> impl PageIterator + Clone { + let max_def_level = column_desc.max_def_level(); + let max_rep_level = column_desc.max_rep_level(); + let rep_levels = vec![0; VALUES_PER_PAGE]; + let mut rng = seedable_rng(); + let mut pages: Vec> = Vec::new(); + for i in 0..NUM_ROW_GROUPS { + let mut column_chunk_pages = Vec::new(); + for j in 0..PAGES_PER_GROUP { + // generate page + let mut values = Vec::with_capacity(VALUES_PER_PAGE); + let mut def_levels = Vec::with_capacity(VALUES_PER_PAGE); + for k in 0..VALUES_PER_PAGE { + let def_level = if rng.random::() < null_density { + max_def_level - 1 + } else { + max_def_level + }; + if def_level == max_def_level { + let string_value = if short_string { + format!("{k}{i}{j}") + } else { + format!("Test value {k}, row group: {i}, page: {j}") + }; + values.push(parquet::data_type::ByteArray::from(string_value.as_str())); + } + def_levels.push(def_level); + } + let mut page_builder = + DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true); + page_builder.add_rep_levels(max_rep_level, &rep_levels); + page_builder.add_def_levels(max_def_level, &def_levels); + page_builder.add_values::(Encoding::PLAIN, &values); + column_chunk_pages.push(page_builder.consume()); + } + pages.push(column_chunk_pages); + } + + InMemoryPageIterator::new(pages) +} + +fn build_constant_prefix_byte_array_page_iterator( + column_desc: ColumnDescPtr, + null_density: f32, + encoding: Encoding, + const_string: bool, +) -> impl PageIterator + Clone { + let max_def_level = column_desc.max_def_level(); + let max_rep_level = column_desc.max_rep_level(); + let rep_levels = vec![0; VALUES_PER_PAGE]; + let mut rng = seedable_rng(); + let mut pages: Vec> = Vec::new(); + for i in 0..NUM_ROW_GROUPS { + let mut column_chunk_pages = Vec::new(); + for j in 0..PAGES_PER_GROUP { + // generate page + let mut values = Vec::with_capacity(VALUES_PER_PAGE); + let mut def_levels = Vec::with_capacity(VALUES_PER_PAGE); + for k in 0..VALUES_PER_PAGE { + let def_level = if rng.random::() < null_density { + max_def_level - 1 + } else { + max_def_level + }; + if def_level == max_def_level { + let string_value = if const_string { + "01234567890123456789012345678901".to_string() + } else { + format!("01234567890123456789012345678901:{:x}{j}{i}", (k % 16)) + }; + values.push(parquet::data_type::ByteArray::from(string_value.as_str())); + } + def_levels.push(def_level); + } + let mut page_builder = + DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true); + page_builder.add_rep_levels(max_rep_level, &rep_levels); + page_builder.add_def_levels(max_def_level, &def_levels); + page_builder.add_values::(encoding, &values); + column_chunk_pages.push(page_builder.consume()); + } + pages.push(column_chunk_pages); + } + + InMemoryPageIterator::new(pages) +} + +fn build_plain_encoded_byte_array_page_iterator( + column_desc: ColumnDescPtr, + null_density: f32, +) -> impl PageIterator + Clone { + build_plain_encoded_byte_array_page_iterator_inner(column_desc, null_density, false) +} + +fn build_dictionary_encoded_string_page_iterator( + column_desc: ColumnDescPtr, + null_density: f32, +) -> impl PageIterator + Clone { + use parquet::encoding::{DictEncoder, Encoder}; + let max_def_level = column_desc.max_def_level(); + let max_rep_level = column_desc.max_rep_level(); + let rep_levels = vec![0; VALUES_PER_PAGE]; + // generate 1% unique values + const NUM_UNIQUE_VALUES: usize = VALUES_PER_PAGE / 100; + let unique_values = (0..NUM_UNIQUE_VALUES) + .map(|x| format!("Dictionary value {x}")) + .collect::>(); + let mut rng = seedable_rng(); + let mut pages: Vec> = Vec::new(); + for _i in 0..NUM_ROW_GROUPS { + let mut column_chunk_pages = VecDeque::new(); + let mut dict_encoder = DictEncoder::::new(column_desc.clone()); + // add data pages + for _j in 0..PAGES_PER_GROUP { + // generate page + let mut values = Vec::with_capacity(VALUES_PER_PAGE); + let mut def_levels = Vec::with_capacity(VALUES_PER_PAGE); + for _k in 0..VALUES_PER_PAGE { + let def_level = if rng.random::() < null_density { + max_def_level - 1 + } else { + max_def_level + }; + if def_level == max_def_level { + // select random value from list of unique values + let string_value = + unique_values[rng.random_range(0..NUM_UNIQUE_VALUES)].as_str(); + values.push(parquet::data_type::ByteArray::from(string_value)); + } + def_levels.push(def_level); + } + let mut page_builder = + DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true); + page_builder.add_rep_levels(max_rep_level, &rep_levels); + page_builder.add_def_levels(max_def_level, &def_levels); + let _ = dict_encoder.put(&values); + let indices = dict_encoder + .write_indices() + .expect("write_indices() should be OK"); + page_builder.add_indices(indices); + column_chunk_pages.push_back(page_builder.consume()); + } + // add dictionary page + let dict = dict_encoder + .write_dict() + .expect("write_dict() should be OK"); + let dict_page = parquet::column::page::Page::DictionaryPage { + buf: dict, + num_values: dict_encoder.num_entries() as u32, + encoding: Encoding::RLE_DICTIONARY, + is_sorted: false, + }; + column_chunk_pages.push_front(dict_page); + pages.push(column_chunk_pages.into()); + } + + InMemoryPageIterator::new(pages) +} + +fn build_string_list_page_iterator( + column_desc: ColumnDescPtr, + null_density: f32, +) -> impl PageIterator + Clone { + let max_def_level = column_desc.max_def_level(); + let max_rep_level = column_desc.max_rep_level(); + assert_eq!(max_def_level, 3); + assert_eq!(max_rep_level, 1); + + let mut rng = seedable_rng(); + let mut pages: Vec> = Vec::new(); + for i in 0..NUM_ROW_GROUPS { + let mut column_chunk_pages = Vec::new(); + for j in 0..PAGES_PER_GROUP { + // generate page + let mut values: Vec = Vec::with_capacity(VALUES_PER_PAGE * MAX_LIST_LEN); + let mut def_levels = Vec::with_capacity(VALUES_PER_PAGE * MAX_LIST_LEN); + let mut rep_levels = Vec::with_capacity(VALUES_PER_PAGE * MAX_LIST_LEN); + for k in 0..VALUES_PER_PAGE { + rep_levels.push(0); + if rng.random::() < null_density { + // Null list + def_levels.push(0); + continue; + } + let len = rng.random_range(0..MAX_LIST_LEN); + if len == 0 { + // Empty list + def_levels.push(1); + continue; + } + + (1..len).for_each(|_| rep_levels.push(1)); + + for l in 0..len { + if rng.random::() < null_density { + // Null element + def_levels.push(2); + } else { + def_levels.push(3); + let value = format!("Test value {k}[{l}], row group: {i}, page: {j}"); + values.push(value.as_str().into()); + } + } + } + let mut page_builder = + DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true); + page_builder.add_rep_levels(max_rep_level, &rep_levels); + page_builder.add_def_levels(max_def_level, &def_levels); + page_builder.add_values::(Encoding::PLAIN, &values); + column_chunk_pages.push(page_builder.consume()); + } + pages.push(column_chunk_pages); + } + + InMemoryPageIterator::new(pages) +} + +fn build_int32_list_page_iterator( + column_desc: ColumnDescPtr, + null_density: f32, +) -> impl PageIterator + Clone { + let max_def_level = column_desc.max_def_level(); + let max_rep_level = column_desc.max_rep_level(); + assert_eq!(max_def_level, 3); + assert_eq!(max_rep_level, 1); + + let mut rng = seedable_rng(); + let mut pages: Vec> = Vec::new(); + for _i in 0..NUM_ROW_GROUPS { + let mut column_chunk_pages = Vec::new(); + for _j in 0..PAGES_PER_GROUP { + let mut values: Vec = Vec::with_capacity(VALUES_PER_PAGE * MAX_LIST_LEN); + let mut def_levels = Vec::with_capacity(VALUES_PER_PAGE * MAX_LIST_LEN); + let mut rep_levels = Vec::with_capacity(VALUES_PER_PAGE * MAX_LIST_LEN); + for _k in 0..VALUES_PER_PAGE { + rep_levels.push(0); + if rng.random::() < null_density { + def_levels.push(0); + continue; + } + let len = rng.random_range(0..MAX_LIST_LEN); + if len == 0 { + def_levels.push(1); + continue; + } + + (1..len).for_each(|_| rep_levels.push(1)); + + for _l in 0..len { + if rng.random::() < null_density { + def_levels.push(2); + } else { + def_levels.push(3); + values.push(rng.random()); + } + } + } + let mut page_builder = + DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true); + page_builder.add_rep_levels(max_rep_level, &rep_levels); + page_builder.add_def_levels(max_def_level, &def_levels); + page_builder.add_values::(Encoding::PLAIN, &values); + column_chunk_pages.push(page_builder.consume()); + } + pages.push(column_chunk_pages); + } + + InMemoryPageIterator::new(pages) +} + +fn create_int32_list_reader( + page_iterator: impl PageIterator + 'static, + column_desc: ColumnDescPtr, +) -> Box { + use parquet::arrow::array_reader::PrimitiveArrayReader; + let items = Box::new( + PrimitiveArrayReader::::new( + Box::new(page_iterator), + column_desc, + None, + DEFAULT_BATCH_SIZE, + Some(2), + ) + .unwrap(), + ) as Box; + let field = Field::new_list_field(DataType::Int32, true); + let data_type = DataType::List(Arc::new(field)); + Box::new(ListArrayReader::::new( + items, data_type, 2, 1, true, None, + )) +} + +const FIXED_BYTE_LEN: usize = 32; + +fn build_fixed32_list_page_iterator( + column_desc: ColumnDescPtr, + null_density: f32, +) -> impl PageIterator + Clone { + let max_def_level = column_desc.max_def_level(); + let max_rep_level = column_desc.max_rep_level(); + assert_eq!(max_def_level, 3); + assert_eq!(max_rep_level, 1); + + let mut rng = seedable_rng(); + let mut pages: Vec> = Vec::new(); + for _i in 0..NUM_ROW_GROUPS { + let mut column_chunk_pages = Vec::new(); + for _j in 0..PAGES_PER_GROUP { + let mut values: Vec = + Vec::with_capacity(VALUES_PER_PAGE * MAX_LIST_LEN); + let mut def_levels = Vec::with_capacity(VALUES_PER_PAGE * MAX_LIST_LEN); + let mut rep_levels = Vec::with_capacity(VALUES_PER_PAGE * MAX_LIST_LEN); + for _k in 0..VALUES_PER_PAGE { + rep_levels.push(0); + if rng.random::() < null_density { + def_levels.push(0); + continue; + } + let len = rng.random_range(0..MAX_LIST_LEN); + if len == 0 { + def_levels.push(1); + continue; + } + (1..len).for_each(|_| rep_levels.push(1)); + for _l in 0..len { + if rng.random::() < null_density { + def_levels.push(2); + } else { + def_levels.push(3); + let mut buf = vec![0u8; FIXED_BYTE_LEN]; + rng.fill(&mut buf[..]); + values.push(buf.into()); + } + } + } + let mut page_builder = + DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true); + page_builder.add_rep_levels(max_rep_level, &rep_levels); + page_builder.add_def_levels(max_def_level, &def_levels); + page_builder.add_values::(Encoding::PLAIN, &values); + column_chunk_pages.push(page_builder.consume()); + } + pages.push(column_chunk_pages); + } + + InMemoryPageIterator::new(pages) +} + +fn create_fixed32_list_reader( + page_iterator: impl PageIterator + 'static, + column_desc: ColumnDescPtr, +) -> Box { + let items = make_fixed_len_byte_array_reader( + Box::new(page_iterator), + column_desc, + None, + DEFAULT_BATCH_SIZE, + Some(2), + ) + .unwrap(); + let field = Field::new_list_field(DataType::FixedSizeBinary(FIXED_BYTE_LEN as i32), true); + let data_type = DataType::List(Arc::new(field)); + Box::new(ListArrayReader::::new( + items, data_type, 2, 1, true, None, + )) +} + +fn bench_array_reader(mut array_reader: Box) -> usize { + // test procedure: read data in batches of 8192 until no more data + let mut total_count = 0; + loop { + let array = array_reader.next_batch(BATCH_SIZE); + let array_len = array.unwrap().len(); + total_count += array_len; + if array_len < BATCH_SIZE { + break; + } + } + total_count +} + +fn bench_array_reader_skip(mut array_reader: Box) -> usize { + // test procedure: read data in batches of 8192 until no more data + let mut total_count = 0; + let mut skip = false; + let mut array_len; + loop { + if skip { + array_len = array_reader.skip_records(BATCH_SIZE).unwrap(); + } else { + let array = array_reader.next_batch(BATCH_SIZE); + array_len = array.unwrap().len(); + } + total_count += array_len; + skip = !skip; + if array_len < BATCH_SIZE { + break; + } + } + total_count +} + +fn create_primitive_array_reader( + page_iterator: impl PageIterator + 'static, + column_desc: ColumnDescPtr, +) -> Box { + use parquet::arrow::array_reader::PrimitiveArrayReader; + match column_desc.physical_type() { + Type::INT32 => { + let reader = PrimitiveArrayReader::::new( + Box::new(page_iterator), + column_desc, + None, + DEFAULT_BATCH_SIZE, + None, + ) + .unwrap(); + Box::new(reader) + } + Type::INT64 => { + let reader = PrimitiveArrayReader::::new( + Box::new(page_iterator), + column_desc, + None, + DEFAULT_BATCH_SIZE, + None, + ) + .unwrap(); + Box::new(reader) + } + Type::BOOLEAN => { + let reader = PrimitiveArrayReader::::new( + Box::new(page_iterator), + column_desc, + None, + DEFAULT_BATCH_SIZE, + None, + ) + .unwrap(); + Box::new(reader) + } + _ => unreachable!(), + } +} + +fn create_f16_by_bytes_reader( + page_iterator: impl PageIterator + 'static, + column_desc: ColumnDescPtr, +) -> Box { + let physical_type = column_desc.physical_type(); + match physical_type { + Type::FIXED_LEN_BYTE_ARRAY => make_fixed_len_byte_array_reader( + Box::new(page_iterator), + column_desc, + None, + DEFAULT_BATCH_SIZE, + None, + ) + .unwrap(), + _ => unimplemented!(), + } +} + +fn create_decimal_by_bytes_reader( + page_iterator: impl PageIterator + 'static, + column_desc: ColumnDescPtr, +) -> Box { + let physical_type = column_desc.physical_type(); + match physical_type { + Type::BYTE_ARRAY => make_byte_array_reader( + Box::new(page_iterator), + column_desc, + None, + DEFAULT_BATCH_SIZE, + None, + ) + .unwrap(), + Type::FIXED_LEN_BYTE_ARRAY => make_fixed_len_byte_array_reader( + Box::new(page_iterator), + column_desc, + None, + DEFAULT_BATCH_SIZE, + None, + ) + .unwrap(), + _ => unimplemented!(), + } +} + +fn create_fixed_len_byte_array_reader( + page_iterator: impl PageIterator + 'static, + column_desc: ColumnDescPtr, +) -> Box { + make_fixed_len_byte_array_reader( + Box::new(page_iterator), + column_desc, + None, + DEFAULT_BATCH_SIZE, + None, + ) + .unwrap() +} + +fn create_byte_array_reader( + page_iterator: impl PageIterator + 'static, + column_desc: ColumnDescPtr, +) -> Box { + make_byte_array_reader( + Box::new(page_iterator), + column_desc, + None, + DEFAULT_BATCH_SIZE, + None, + ) + .unwrap() +} + +fn create_byte_view_array_reader( + page_iterator: impl PageIterator + 'static, + column_desc: ColumnDescPtr, +) -> Box { + make_byte_view_array_reader( + Box::new(page_iterator), + column_desc, + None, + DEFAULT_BATCH_SIZE, + None, + ) + .unwrap() +} + +fn create_string_view_byte_array_reader( + page_iterator: impl PageIterator + 'static, + column_desc: ColumnDescPtr, +) -> Box { + make_byte_view_array_reader( + Box::new(page_iterator), + column_desc, + None, + DEFAULT_BATCH_SIZE, + None, + ) + .unwrap() +} + +fn create_string_byte_array_dictionary_reader( + page_iterator: impl PageIterator + 'static, + column_desc: ColumnDescPtr, +) -> Box { + use parquet::arrow::array_reader::make_byte_array_dictionary_reader; + let arrow_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + + make_byte_array_dictionary_reader( + Box::new(page_iterator), + column_desc, + Some(arrow_type), + DEFAULT_BATCH_SIZE, + None, + ) + .unwrap() +} + +fn create_string_list_reader( + page_iterator: impl PageIterator + 'static, + column_desc: ColumnDescPtr, +) -> Box { + let items = make_byte_array_reader( + Box::new(page_iterator), + column_desc, + None, + DEFAULT_BATCH_SIZE, + Some(2), + ) + .unwrap(); + let field = Field::new_list_field(DataType::Utf8, true); + let data_type = DataType::List(Arc::new(field)); + Box::new(ListArrayReader::::new( + items, data_type, 2, 1, true, None, + )) +} + +fn bench_byte_decimal( + group: &mut BenchmarkGroup, + mandatory_column_desc: &ColumnDescPtr, + optional_column_desc: &ColumnDescPtr, + encoding: Encoding, + min: i128, + max: i128, +) where + T: parquet::data_type::DataType, + T::T: From>, +{ + // all are plain encoding + let mut count: usize = 0; + + // no NULLs + let data = build_encoded_decimal_bytes_page_iterator::( + mandatory_column_desc.clone(), + 0.0, + encoding, + min, + max, + ); + group.bench_function( + encoding.to_string().to_lowercase() + " encoded, mandatory, no NULLs", + |b| { + b.iter(|| { + let array_reader = + create_decimal_by_bytes_reader(data.clone(), mandatory_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }, + ); + + let data = build_encoded_decimal_bytes_page_iterator::( + optional_column_desc.clone(), + 0.0, + encoding, + min, + max, + ); + group.bench_function( + encoding.to_string().to_lowercase() + " encoded, optional, no NULLs", + |b| { + b.iter(|| { + let array_reader = + create_decimal_by_bytes_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }, + ); + + // half null + let data = build_encoded_decimal_bytes_page_iterator::( + optional_column_desc.clone(), + 0.5, + encoding, + min, + max, + ); + group.bench_function( + encoding.to_string().to_lowercase() + " encoded, optional, half NULLs", + |b| { + b.iter(|| { + let array_reader = + create_decimal_by_bytes_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }, + ); +} + +fn bench_f16( + group: &mut BenchmarkGroup, + mandatory_column_desc: &ColumnDescPtr, + optional_column_desc: &ColumnDescPtr, + encoding: Encoding, + min: f32, + max: f32, +) where + T: parquet::data_type::DataType, + T::T: From>, +{ + let mut count: usize = 0; + + // byte_stream_split encoded, no NULLs + let data = build_encoded_f16_bytes_page_iterator::( + mandatory_column_desc.clone(), + 0.0, + encoding, + min, + max, + ); + group.bench_function( + encoding.to_string().to_lowercase() + " encoded, mandatory, no NULLs", + |b| { + b.iter(|| { + let array_reader = + create_f16_by_bytes_reader(data.clone(), mandatory_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }, + ); + + let data = build_encoded_f16_bytes_page_iterator::( + optional_column_desc.clone(), + 0.0, + encoding, + min, + max, + ); + group.bench_function( + encoding.to_string().to_lowercase() + " encoded, optional, no NULLs", + |b| { + b.iter(|| { + let array_reader = + create_f16_by_bytes_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }, + ); + + let data = build_encoded_f16_bytes_page_iterator::( + optional_column_desc.clone(), + 0.5, + encoding, + min, + max, + ); + group.bench_function( + encoding.to_string().to_lowercase() + " encoded, optional, half NULLs", + |b| { + b.iter(|| { + let array_reader = + create_f16_by_bytes_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }, + ); +} + +fn bench_flba( + group: &mut BenchmarkGroup, + mandatory_column_desc: &ColumnDescPtr, + optional_column_desc: &ColumnDescPtr, + encoding: Encoding, +) { + let mut count: usize = 0; + + encoding.to_string(); + // no NULLs + let data = build_encoded_flba_bytes_page_iterator::( + mandatory_column_desc.clone(), + 0.0, + encoding, + ); + group.bench_function( + encoding.to_string().to_lowercase() + " encoded, mandatory, no NULLs", + |b| { + b.iter(|| { + let array_reader = + create_fixed_len_byte_array_reader(data.clone(), mandatory_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }, + ); + + let data = build_encoded_flba_bytes_page_iterator::( + optional_column_desc.clone(), + 0.0, + encoding, + ); + group.bench_function( + encoding.to_string().to_lowercase() + " encoded, optional, no NULLs", + |b| { + b.iter(|| { + let array_reader = + create_fixed_len_byte_array_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }, + ); + + // half null + let data = build_encoded_flba_bytes_page_iterator::( + optional_column_desc.clone(), + 0.5, + encoding, + ); + group.bench_function( + encoding.to_string().to_lowercase() + " encoded, optional, half NULLs", + |b| { + b.iter(|| { + let array_reader = + create_fixed_len_byte_array_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }, + ); +} + +fn bench_fixed_len_byte_array( + group: &mut BenchmarkGroup, + mandatory_column_desc: &ColumnDescPtr, + optional_column_desc: &ColumnDescPtr, +) { + bench_flba::( + group, + mandatory_column_desc, + optional_column_desc, + Encoding::PLAIN, + ); + bench_flba::( + group, + mandatory_column_desc, + optional_column_desc, + Encoding::BYTE_STREAM_SPLIT, + ); +} + +fn bench_primitive( + group: &mut BenchmarkGroup, + mandatory_column_desc: &ColumnDescPtr, + optional_column_desc: &ColumnDescPtr, + min: usize, + max: usize, +) where + T: parquet::data_type::DataType, + T::T: SampleUniform + FromPrimitive + Copy, +{ + let mut count: usize = 0; + + // plain encoded, no NULLs + let data = build_encoded_primitive_page_iterator::( + mandatory_column_desc.clone(), + 0.0, + Encoding::PLAIN, + min, + max, + ); + group.bench_function("plain encoded, mandatory, no NULLs", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), mandatory_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + let data = build_encoded_primitive_page_iterator::( + optional_column_desc.clone(), + 0.0, + Encoding::PLAIN, + min, + max, + ); + group.bench_function("plain encoded, optional, no NULLs", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // plain encoded, half NULLs + let data = build_encoded_primitive_page_iterator::( + optional_column_desc.clone(), + 0.5, + Encoding::PLAIN, + min, + max, + ); + group.bench_function("plain encoded, optional, half NULLs", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // binary packed, no NULLs + let data = build_encoded_primitive_page_iterator::( + mandatory_column_desc.clone(), + 0.0, + Encoding::DELTA_BINARY_PACKED, + min, + max, + ); + group.bench_function("binary packed, mandatory, no NULLs", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), mandatory_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + let data = build_encoded_primitive_page_iterator::( + optional_column_desc.clone(), + 0.0, + Encoding::DELTA_BINARY_PACKED, + min, + max, + ); + group.bench_function("binary packed, optional, no NULLs", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // binary packed, half NULLs + let data = build_encoded_primitive_page_iterator::( + optional_column_desc.clone(), + 0.5, + Encoding::DELTA_BINARY_PACKED, + min, + max, + ); + group.bench_function("binary packed, optional, half NULLs", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // binary packed skip , no NULLs + let data = build_encoded_primitive_page_iterator::( + mandatory_column_desc.clone(), + 0.0, + Encoding::DELTA_BINARY_PACKED, + min, + max, + ); + group.bench_function("binary packed skip, mandatory, no NULLs", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), mandatory_column_desc.clone()); + count = bench_array_reader_skip(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + let data = build_encoded_primitive_page_iterator::( + optional_column_desc.clone(), + 0.0, + Encoding::DELTA_BINARY_PACKED, + min, + max, + ); + group.bench_function("binary packed skip, optional, no NULLs", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader_skip(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // binary packed skip, half NULLs + let data = build_encoded_primitive_page_iterator::( + optional_column_desc.clone(), + 0.5, + Encoding::DELTA_BINARY_PACKED, + min, + max, + ); + group.bench_function("binary packed skip, optional, half NULLs", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader_skip(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // binary packed same value + let data = build_delta_encoded_incr_primitive_page_iterator::( + mandatory_column_desc.clone(), + 0.0, + 0, + false, + ); + group.bench_function("binary packed single value", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), mandatory_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + let data = build_delta_encoded_incr_primitive_page_iterator::( + mandatory_column_desc.clone(), + 0.0, + 0, + false, + ); + group.bench_function("binary packed skip single value", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), mandatory_column_desc.clone()); + count = bench_array_reader_skip(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // binary packed monotonically increasing + let data = build_delta_encoded_incr_primitive_page_iterator::( + mandatory_column_desc.clone(), + 0.0, + 1, + false, + ); + group.bench_function("binary packed increasing value", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), mandatory_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + let data = build_delta_encoded_incr_primitive_page_iterator::( + mandatory_column_desc.clone(), + 0.0, + 1, + false, + ); + group.bench_function("binary packed skip increasing value", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), mandatory_column_desc.clone()); + count = bench_array_reader_skip(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // binary packed increasing stepped + let data = build_delta_encoded_incr_primitive_page_iterator::( + mandatory_column_desc.clone(), + 0.0, + 1, + true, + ); + group.bench_function("binary packed stepped increasing value", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), mandatory_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + let data = build_delta_encoded_incr_primitive_page_iterator::( + mandatory_column_desc.clone(), + 0.0, + 1, + true, + ); + group.bench_function("binary packed skip stepped increasing value", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), mandatory_column_desc.clone()); + count = bench_array_reader_skip(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // dictionary encoded, no NULLs + let data = + build_dictionary_encoded_primitive_page_iterator::(mandatory_column_desc.clone(), 0.0); + group.bench_function("dictionary encoded, mandatory, no NULLs", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), mandatory_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + let data = + build_dictionary_encoded_primitive_page_iterator::(optional_column_desc.clone(), 0.0); + group.bench_function("dictionary encoded, optional, no NULLs", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // dictionary encoded, half NULLs + let data = + build_dictionary_encoded_primitive_page_iterator::(optional_column_desc.clone(), 0.5); + group.bench_function("dictionary encoded, optional, half NULLs", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // byte_stream_split encoded, no NULLs + let data = build_encoded_primitive_page_iterator::( + mandatory_column_desc.clone(), + 0.0, + Encoding::BYTE_STREAM_SPLIT, + min, + max, + ); + group.bench_function("byte_stream_split encoded, mandatory, no NULLs", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), mandatory_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + let data = build_encoded_primitive_page_iterator::( + optional_column_desc.clone(), + 0.0, + Encoding::BYTE_STREAM_SPLIT, + min, + max, + ); + group.bench_function("byte_stream_split encoded, optional, no NULLs", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // plain encoded, half NULLs + let data = build_encoded_primitive_page_iterator::( + optional_column_desc.clone(), + 0.5, + Encoding::BYTE_STREAM_SPLIT, + min, + max, + ); + group.bench_function("byte_stream_split encoded, optional, half NULLs", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); +} + +fn bench_boolean( + group: &mut BenchmarkGroup, + mandatory_column_desc: &ColumnDescPtr, + optional_column_desc: &ColumnDescPtr, +) { + let mut count: usize = 0; + + // plain encoded, no NULLs + let data = + build_encoded_bool_page_iterator(mandatory_column_desc.clone(), 0.0, Encoding::PLAIN); + group.bench_function("plain encoded, mandatory, no NULLs", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), mandatory_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + let data = build_encoded_bool_page_iterator(optional_column_desc.clone(), 0.0, Encoding::PLAIN); + group.bench_function("plain encoded, optional, no NULLs", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // plain encoded, half NULLs + let data = build_encoded_bool_page_iterator(optional_column_desc.clone(), 0.5, Encoding::PLAIN); + group.bench_function("plain encoded, optional, half NULLs", |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); +} + +// Benchmark reading a struct with a single primitive field. +// No need to bench all encodings for the data, as that should already be covered by `bench_primitive`. +// The only performance difference should be caused by the additional definition level. +fn bench_struct_primitive( + group: &mut BenchmarkGroup, + mandatory_column_desc: &ColumnDescPtr, + optional_column_desc: &ColumnDescPtr, + min: usize, + max: usize, +) where + T: parquet::data_type::DataType, + T::T: SampleUniform + FromPrimitive + Copy, +{ + let mut count: usize = 0; + + let data = build_encoded_primitive_page_iterator::( + mandatory_column_desc.clone(), + 0.0, + Encoding::PLAIN, + min, + max, + ); + group.bench_function( + "plain encoded, mandatory struct, optional data, no NULLs", + |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), mandatory_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }, + ); + + let data = build_encoded_primitive_page_iterator::( + optional_column_desc.clone(), + 0.0, + Encoding::PLAIN, + min, + max, + ); + group.bench_function( + "plain encoded, optional struct, optional data, no NULLs", + |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }, + ); + + let data = build_encoded_primitive_page_iterator::( + mandatory_column_desc.clone(), + 0.5, + Encoding::PLAIN, + min, + max, + ); + group.bench_function( + "plain encoded, mandatory struct, optional data, half NULLs", + |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), mandatory_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }, + ); + + let data = build_encoded_primitive_page_iterator::( + optional_column_desc.clone(), + 0.5, + Encoding::PLAIN, + min, + max, + ); + group.bench_function( + "plain encoded, optional struct, optional data, half NULLs", + |b| { + b.iter(|| { + let array_reader = + create_primitive_array_reader(data.clone(), optional_column_desc.clone()); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }, + ); +} + +fn float16_benches(c: &mut Criterion) { + let schema = build_test_schema(); + + let mut group = c.benchmark_group("arrow_array_reader/FIXED_LEN_BYTE_ARRAY/Float16Array"); + let mandatory_f16_leaf_desc = schema.column(17); + let optional_f16_leaf_desc = schema.column(18); + bench_f16::( + &mut group, + &mandatory_f16_leaf_desc, + &optional_f16_leaf_desc, + Encoding::PLAIN, + -1.0, + 1.0, + ); + group.finish(); + + let mut group = c.benchmark_group("arrow_array_reader/FIXED_LEN_BYTE_ARRAY/Float16Array"); + let mandatory_f16_leaf_desc = schema.column(17); + let optional_f16_leaf_desc = schema.column(18); + bench_f16::( + &mut group, + &mandatory_f16_leaf_desc, + &optional_f16_leaf_desc, + Encoding::BYTE_STREAM_SPLIT, + -1.0, + 1.0, + ); + group.finish(); +} + +fn decimal_benches(c: &mut Criterion) { + let schema = build_test_schema(); + // parquet int32, logical type decimal(8,2) + let mandatory_decimal1_leaf_desc = schema.column(6); + let optional_decimal1_leaf_desc = schema.column(7); + let mut group = c.benchmark_group("arrow_array_reader/INT32/Decimal128Array"); + bench_primitive::( + &mut group, + &mandatory_decimal1_leaf_desc, + &optional_decimal1_leaf_desc, + // precision is 8: the max is 99999999 + 9999000, + 9999999, + ); + group.finish(); + + // parquet int64, logical type decimal(16,2) + let mut group = c.benchmark_group("arrow_array_reader/INT64/Decimal128Array"); + let mandatory_decimal2_leaf_desc = schema.column(8); + let optional_decimal2_leaf_desc = schema.column(9); + bench_primitive::( + &mut group, + &mandatory_decimal2_leaf_desc, + &optional_decimal2_leaf_desc, + // precision is 16: the max is 9999999999999999 + 9999999999999000, + 9999999999999999, + ); + group.finish(); + + // parquet BYTE_ARRAY, logical type decimal(16,2) + let mut group = c.benchmark_group("arrow_array_reader/BYTE_ARRAY/Decimal128Array"); + let mandatory_decimal3_leaf_desc = schema.column(10); + let optional_decimal3_leaf_desc = schema.column(11); + bench_byte_decimal::( + &mut group, + &mandatory_decimal3_leaf_desc, + &optional_decimal3_leaf_desc, + Encoding::PLAIN, + // precision is 16: the max is 9999999999999999 + 9999999999999000, + 9999999999999999, + ); + group.finish(); + + // parquet FIXED_LEN_BYTE_ARRAY, logical type decimal(16,2) + let mut group = c.benchmark_group("arrow_array_reader/FIXED_LEN_BYTE_ARRAY/Decimal128Array"); + let mandatory_decimal4_leaf_desc = schema.column(12); + let optional_decimal4_leaf_desc = schema.column(13); + bench_byte_decimal::( + &mut group, + &mandatory_decimal4_leaf_desc, + &optional_decimal4_leaf_desc, + Encoding::PLAIN, + // precision is 16: the max is 9999999999999999 + 9999999999999000, + 9999999999999999, + ); + group.finish(); + + let mut group = c.benchmark_group("arrow_array_reader/FIXED_LEN_BYTE_ARRAY/Decimal128Array"); + let mandatory_decimal4_leaf_desc = schema.column(12); + let optional_decimal4_leaf_desc = schema.column(13); + bench_byte_decimal::( + &mut group, + &mandatory_decimal4_leaf_desc, + &optional_decimal4_leaf_desc, + Encoding::BYTE_STREAM_SPLIT, + // precision is 16: the max is 9999999999999999 + 9999999999999000, + 9999999999999999, + ); + group.finish(); +} + +fn add_benches(c: &mut Criterion) { + let mut count: usize = 0; + + let schema = build_test_schema(); + let mandatory_int32_column_desc = schema.column(0); + let optional_int32_column_desc = schema.column(1); + let mandatory_string_column_desc = schema.column(2); + let optional_string_column_desc = schema.column(3); + let mandatory_int64_column_desc = schema.column(4); + let optional_int64_column_desc = schema.column(5); + let string_list_desc = schema.column(14); + let mandatory_binary_column_desc = schema.column(15); + let optional_binary_column_desc = schema.column(16); + let mandatory_uint8_column_desc = schema.column(27); + let optional_uint8_column_desc = schema.column(28); + let mandatory_uint16_column_desc = schema.column(29); + let optional_uint16_column_desc = schema.column(30); + let mandatory_uint32_column_desc = schema.column(31); + let optional_uint32_column_desc = schema.column(32); + let mandatory_int8_column_desc = schema.column(33); + let optional_int8_column_desc = schema.column(34); + let mandatory_int16_column_desc = schema.column(35); + let optional_int16_column_desc = schema.column(36); + let mandatory_uint64_column_desc = schema.column(37); + let optional_uint64_column_desc = schema.column(38); + let mandatory_struct_optional_in32_column_desc = schema.column(39); + let optional_struct_optional_in32_column_desc = schema.column(40); + let int32_list_desc = schema.column(41); + let fixed32_list_desc = schema.column(42); + let mandatory_bool_column_desc = schema.column(43); + let optional_bool_column_desc = schema.column(44); + + // primitive / int32 benchmarks + // ============================= + + let mut group = c.benchmark_group("arrow_array_reader/Int32Array"); + bench_primitive::( + &mut group, + &mandatory_int32_column_desc, + &optional_int32_column_desc, + 0, + 1000, + ); + group.finish(); + + // primitive int32 / logical uint8 benchmarks + let mut group = c.benchmark_group("arrow_array_reader/UInt8Array"); + bench_primitive::( + &mut group, + &mandatory_uint8_column_desc, + &optional_uint8_column_desc, + 0, + 256, + ); + group.finish(); + + // primitive int32 / logical int8 benchmarks + let mut group = c.benchmark_group("arrow_array_reader/Int8Array"); + bench_primitive::( + &mut group, + &mandatory_int8_column_desc, + &optional_int8_column_desc, + 0, + 128, + ); + group.finish(); + + // primitive int32 / logical uint16 benchmarks + let mut group = c.benchmark_group("arrow_array_reader/UInt16Array"); + bench_primitive::( + &mut group, + &mandatory_uint16_column_desc, + &optional_uint16_column_desc, + 0, + 65536, + ); + group.finish(); + + // primitive int32 / logical int16 benchmarks + let mut group = c.benchmark_group("arrow_array_reader/Int16Array"); + bench_primitive::( + &mut group, + &mandatory_int16_column_desc, + &optional_int16_column_desc, + 0, + 32768, + ); + group.finish(); + + // primitive int32 / logical uint32 benchmarks + let mut group = c.benchmark_group("arrow_array_reader/UInt32Array"); + bench_primitive::( + &mut group, + &mandatory_uint32_column_desc, + &optional_uint32_column_desc, + 0, + 1000, + ); + group.finish(); + + // primitive / int64 benchmarks + // ============================= + + let mut group = c.benchmark_group("arrow_array_reader/Int64Array"); + bench_primitive::( + &mut group, + &mandatory_int64_column_desc, + &optional_int64_column_desc, + 0, + 1000, + ); + group.finish(); + + // primitive int64 / logical uint64 benchmarks + let mut group = c.benchmark_group("arrow_array_reader/UInt64Array"); + bench_primitive::( + &mut group, + &mandatory_uint64_column_desc, + &optional_uint64_column_desc, + 0, + 1000, + ); + group.finish(); + + // boolean benchmarks + let mut group = c.benchmark_group("arrow_array_reader/BooleanArray"); + bench_boolean( + &mut group, + &mandatory_bool_column_desc, + &optional_bool_column_desc, + ); + group.finish(); + + let mut group = c.benchmark_group("arrow_array_reader/struct/Int32Array"); + bench_struct_primitive::( + &mut group, + &mandatory_struct_optional_in32_column_desc, + &optional_struct_optional_in32_column_desc, + 0, + 1000, + ); + group.finish(); + + // string benchmarks + //============================== + + let mut group = c.benchmark_group("arrow_array_reader/StringArray"); + + // string, plain encoded, no NULLs + let plain_string_no_null_data = + build_plain_encoded_byte_array_page_iterator(mandatory_string_column_desc.clone(), 0.0); + group.bench_function("plain encoded, mandatory, no NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_array_reader( + plain_string_no_null_data.clone(), + mandatory_string_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + let plain_string_no_null_data = + build_plain_encoded_byte_array_page_iterator(optional_string_column_desc.clone(), 0.0); + group.bench_function("plain encoded, optional, no NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_array_reader( + plain_string_no_null_data.clone(), + optional_string_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // string, plain encoded, half NULLs + let plain_string_half_null_data = + build_plain_encoded_byte_array_page_iterator(optional_string_column_desc.clone(), 0.5); + group.bench_function("plain encoded, optional, half NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_array_reader( + plain_string_half_null_data.clone(), + optional_string_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // string, dictionary encoded, no NULLs + let dictionary_string_no_null_data = + build_dictionary_encoded_string_page_iterator(mandatory_string_column_desc.clone(), 0.0); + group.bench_function("dictionary encoded, mandatory, no NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_array_reader( + dictionary_string_no_null_data.clone(), + mandatory_string_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + let dictionary_string_no_null_data = + build_dictionary_encoded_string_page_iterator(optional_string_column_desc.clone(), 0.0); + group.bench_function("dictionary encoded, optional, no NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_array_reader( + dictionary_string_no_null_data.clone(), + optional_string_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // string, dictionary encoded, half NULLs + let dictionary_string_half_null_data = + build_dictionary_encoded_string_page_iterator(optional_string_column_desc.clone(), 0.5); + group.bench_function("dictionary encoded, optional, half NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_array_reader( + dictionary_string_half_null_data.clone(), + optional_string_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // delta byte array with constant prefix and suffix lengths + let delta_string_const_prefix_no_null_data = build_constant_prefix_byte_array_page_iterator( + mandatory_string_column_desc.clone(), + 0.0, + Encoding::DELTA_BYTE_ARRAY, + false, + ); + group.bench_function( + "const prefix delta byte array encoded, mandatory, no NULLs", + |b| { + b.iter(|| { + let array_reader = create_byte_array_reader( + delta_string_const_prefix_no_null_data.clone(), + mandatory_string_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }, + ); + + // delta byte array with constant prefix and no suffix + let delta_string_const_no_null_data = build_constant_prefix_byte_array_page_iterator( + mandatory_string_column_desc.clone(), + 0.0, + Encoding::DELTA_BYTE_ARRAY, + true, + ); + group.bench_function("const delta byte array encoded, mandatory, no NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_array_reader( + delta_string_const_no_null_data.clone(), + mandatory_string_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // delta length byte array with constant lengths + let delta_string_const_no_null_data = build_constant_prefix_byte_array_page_iterator( + mandatory_string_column_desc.clone(), + 0.0, + Encoding::DELTA_LENGTH_BYTE_ARRAY, + true, + ); + group.bench_function( + "const delta length byte array encoded, mandatory, no NULLs", + |b| { + b.iter(|| { + let array_reader = create_byte_array_reader( + delta_string_const_no_null_data.clone(), + mandatory_string_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }, + ); + + group.finish(); + + // binary benchmarks + //============================== + + let mut group = c.benchmark_group("arrow_array_reader/BinaryArray"); + + // byte array, plain encoded, no NULLs + let plain_byte_array_no_null_data = + build_plain_encoded_byte_array_page_iterator(mandatory_binary_column_desc.clone(), 0.0); + group.bench_function("plain encoded, mandatory, no NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_array_reader( + plain_byte_array_no_null_data.clone(), + mandatory_binary_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + let plain_byte_array_no_null_data = + build_plain_encoded_byte_array_page_iterator(optional_binary_column_desc.clone(), 0.0); + group.bench_function("plain encoded, optional, no NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_array_reader( + plain_byte_array_no_null_data.clone(), + optional_binary_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // byte array, plain encoded, half NULLs + let plain_byte_array_half_null_data = + build_plain_encoded_byte_array_page_iterator(optional_binary_column_desc.clone(), 0.5); + group.bench_function("plain encoded, optional, half NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_array_reader( + plain_byte_array_half_null_data.clone(), + optional_binary_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // byte array, dictionary encoded, no NULLs + let dictionary_byte_array_no_null_data = + build_dictionary_encoded_string_page_iterator(mandatory_binary_column_desc.clone(), 0.0); + group.bench_function("dictionary encoded, mandatory, no NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_array_reader( + dictionary_byte_array_no_null_data.clone(), + mandatory_binary_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + let dictionary_byte_array_no_null_data = + build_dictionary_encoded_string_page_iterator(optional_binary_column_desc.clone(), 0.0); + group.bench_function("dictionary encoded, optional, no NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_array_reader( + dictionary_byte_array_no_null_data.clone(), + optional_binary_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // string, dictionary encoded, half NULLs + let dictionary_byte_array_half_null_data = + build_dictionary_encoded_string_page_iterator(optional_binary_column_desc.clone(), 0.5); + group.bench_function("dictionary encoded, optional, half NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_array_reader( + dictionary_byte_array_half_null_data.clone(), + optional_binary_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + group.finish(); + + // binary view benchmarks + //============================== + + let mut group = c.benchmark_group("arrow_array_reader/BinaryViewArray"); + + // binary view, plain encoded, no NULLs, short string + let plain_byte_array_no_null_data = build_plain_encoded_byte_array_page_iterator_inner( + mandatory_binary_column_desc.clone(), + 0.0, + true, + ); + + // Short strings should not be slower than long strings, however, as discussed in https://github.com/apache/arrow-rs/issues/6034, + // the current implementation is more than 2x slower. + // This benchmark tracks the performance of short strings so that we can optimize it. + group.bench_function("plain encoded, mandatory, no NULLs, short string", |b| { + b.iter(|| { + let array_reader = create_byte_view_array_reader( + plain_byte_array_no_null_data.clone(), + mandatory_binary_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // binary view, plain encoded, no NULLs + let plain_byte_array_no_null_data = + build_plain_encoded_byte_array_page_iterator(mandatory_binary_column_desc.clone(), 0.0); + group.bench_function("plain encoded, mandatory, no NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_view_array_reader( + plain_byte_array_no_null_data.clone(), + mandatory_binary_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + let plain_byte_array_no_null_data = + build_plain_encoded_byte_array_page_iterator(optional_binary_column_desc.clone(), 0.0); + group.bench_function("plain encoded, optional, no NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_view_array_reader( + plain_byte_array_no_null_data.clone(), + optional_binary_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // binary view, plain encoded, half NULLs + let plain_byte_array_half_null_data = + build_plain_encoded_byte_array_page_iterator(optional_binary_column_desc.clone(), 0.5); + group.bench_function("plain encoded, optional, half NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_view_array_reader( + plain_byte_array_half_null_data.clone(), + optional_binary_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // binary view, dictionary encoded, no NULLs + let dictionary_byte_array_no_null_data = + build_dictionary_encoded_string_page_iterator(mandatory_binary_column_desc.clone(), 0.0); + group.bench_function("dictionary encoded, mandatory, no NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_view_array_reader( + dictionary_byte_array_no_null_data.clone(), + mandatory_binary_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + let dictionary_byte_array_no_null_data = + build_dictionary_encoded_string_page_iterator(optional_binary_column_desc.clone(), 0.0); + group.bench_function("dictionary encoded, optional, no NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_view_array_reader( + dictionary_byte_array_no_null_data.clone(), + optional_binary_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // binary view, dictionary encoded, half NULLs + let dictionary_byte_array_half_null_data = + build_dictionary_encoded_string_page_iterator(optional_binary_column_desc.clone(), 0.5); + group.bench_function("dictionary encoded, optional, half NULLs", |b| { + b.iter(|| { + let array_reader = create_byte_view_array_reader( + dictionary_byte_array_half_null_data.clone(), + optional_binary_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + group.finish(); + + // string dictionary benchmarks + //============================== + + let mut group = c.benchmark_group("arrow_array_reader/StringDictionary"); + + group.bench_function("dictionary encoded, mandatory, no NULLs", |b| { + b.iter(|| { + let array_reader = create_string_byte_array_dictionary_reader( + dictionary_byte_array_no_null_data.clone(), + mandatory_string_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + group.bench_function("dictionary encoded, optional, no NULLs", |b| { + b.iter(|| { + let array_reader = create_string_byte_array_dictionary_reader( + dictionary_byte_array_no_null_data.clone(), + optional_string_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + group.bench_function("dictionary encoded, optional, half NULLs", |b| { + b.iter(|| { + let array_reader = create_string_byte_array_dictionary_reader( + dictionary_byte_array_half_null_data.clone(), + optional_string_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + group.finish(); + + // string view benchmarks + //============================== + + let mut group = c.benchmark_group("arrow_array_reader/StringViewArray"); + + // string, plain encoded, no NULLs + let plain_string_no_null_data = + build_plain_encoded_byte_array_page_iterator(mandatory_string_column_desc.clone(), 0.0); + group.bench_function("plain encoded, mandatory, no NULLs", |b| { + b.iter(|| { + let array_reader = create_string_view_byte_array_reader( + plain_string_no_null_data.clone(), + mandatory_string_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + let plain_string_no_null_data = + build_plain_encoded_byte_array_page_iterator(optional_string_column_desc.clone(), 0.0); + group.bench_function("plain encoded, optional, no NULLs", |b| { + b.iter(|| { + let array_reader = create_string_view_byte_array_reader( + plain_string_no_null_data.clone(), + optional_string_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // string, plain encoded, half NULLs + let plain_string_half_null_data = + build_plain_encoded_byte_array_page_iterator(optional_string_column_desc.clone(), 0.5); + group.bench_function("plain encoded, optional, half NULLs", |b| { + b.iter(|| { + let array_reader = create_string_view_byte_array_reader( + plain_string_half_null_data.clone(), + optional_string_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // string, dictionary encoded, no NULLs + let dictionary_string_no_null_data = + build_dictionary_encoded_string_page_iterator(mandatory_string_column_desc.clone(), 0.0); + group.bench_function("dictionary encoded, mandatory, no NULLs", |b| { + b.iter(|| { + let array_reader = create_string_view_byte_array_reader( + dictionary_string_no_null_data.clone(), + mandatory_string_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + let dictionary_string_no_null_data = + build_dictionary_encoded_string_page_iterator(optional_string_column_desc.clone(), 0.0); + group.bench_function("dictionary encoded, optional, no NULLs", |b| { + b.iter(|| { + let array_reader = create_string_view_byte_array_reader( + dictionary_string_no_null_data.clone(), + optional_string_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + // string, dictionary encoded, half NULLs + let dictionary_string_half_null_data = + build_dictionary_encoded_string_page_iterator(optional_string_column_desc.clone(), 0.5); + group.bench_function("dictionary encoded, optional, half NULLs", |b| { + b.iter(|| { + let array_reader = create_string_view_byte_array_reader( + dictionary_string_half_null_data.clone(), + optional_string_column_desc.clone(), + ); + count = bench_array_reader(array_reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + + group.finish(); + + // list benchmarks + //============================== + + let mut group = c.benchmark_group("arrow_array_reader/ListArray/StringList"); + for (label, null_density) in [ + ("no NULLs", 0.0), + ("half NULLs", 0.5), + ("90pct NULLs", 0.9), + ("99pct NULLs", 0.99), + ] { + let list_data = build_string_list_page_iterator(string_list_desc.clone(), null_density); + group.bench_function(label, |b| { + b.iter(|| { + let reader = create_string_list_reader(list_data.clone(), string_list_desc.clone()); + count = bench_array_reader(reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + } + group.finish(); + + let mut group = c.benchmark_group("arrow_array_reader/ListArray/Int32List"); + for (label, null_density) in [ + ("no NULLs", 0.0), + ("half NULLs", 0.5), + ("90pct NULLs", 0.9), + ("99pct NULLs", 0.99), + ] { + let list_data = build_int32_list_page_iterator(int32_list_desc.clone(), null_density); + group.bench_function(label, |b| { + b.iter(|| { + let reader = create_int32_list_reader(list_data.clone(), int32_list_desc.clone()); + count = bench_array_reader(reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + } + group.finish(); + + let mut group = c.benchmark_group("arrow_array_reader/ListArray/Fixed32List"); + for (label, null_density) in [ + ("no NULLs", 0.0), + ("half NULLs", 0.5), + ("90pct NULLs", 0.9), + ("99pct NULLs", 0.99), + ] { + let list_data = build_fixed32_list_page_iterator(fixed32_list_desc.clone(), null_density); + group.bench_function(label, |b| { + b.iter(|| { + let reader = + create_fixed32_list_reader(list_data.clone(), fixed32_list_desc.clone()); + count = bench_array_reader(reader); + }); + assert_eq!(count, EXPECTED_VALUE_COUNT); + }); + } + group.finish(); + + // fixed_len_byte_array benchmarks + //============================== + + let mut group = c.benchmark_group("arrow_array_reader/FixedLenByteArray(2)"); + let mandatory_flba2_leaf_desc = schema.column(19); + let optional_flba2_leaf_desc = schema.column(20); + bench_fixed_len_byte_array::<2>( + &mut group, + &mandatory_flba2_leaf_desc, + &optional_flba2_leaf_desc, + ); + group.finish(); + + let mut group = c.benchmark_group("arrow_array_reader/FixedLenByteArray(4)"); + let mandatory_flba4_leaf_desc = schema.column(21); + let optional_flba4_leaf_desc = schema.column(22); + bench_fixed_len_byte_array::<4>( + &mut group, + &mandatory_flba4_leaf_desc, + &optional_flba4_leaf_desc, + ); + group.finish(); + + let mut group = c.benchmark_group("arrow_array_reader/FixedLenByteArray(8)"); + let mandatory_flba8_leaf_desc = schema.column(23); + let optional_flba8_leaf_desc = schema.column(24); + bench_fixed_len_byte_array::<8>( + &mut group, + &mandatory_flba8_leaf_desc, + &optional_flba8_leaf_desc, + ); + group.finish(); + + let mut group = c.benchmark_group("arrow_array_reader/FixedLenByteArray(16)"); + let mandatory_flba16_leaf_desc = schema.column(25); + let optional_flba16_leaf_desc = schema.column(26); + bench_fixed_len_byte_array::<16>( + &mut group, + &mandatory_flba16_leaf_desc, + &optional_flba16_leaf_desc, + ); + group.finish(); +} + +criterion_group!(benches, add_benches, decimal_benches, float16_benches,); +criterion_main!(benches); diff --git a/vendor/parquet/benches/arrow_reader_clickbench.rs b/vendor/parquet/benches/arrow_reader_clickbench.rs new file mode 100644 index 0000000..f411b96 --- /dev/null +++ b/vendor/parquet/benches/arrow_reader_clickbench.rs @@ -0,0 +1,961 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmark for evaluating row filters and projections using the [ClickBench] queries and data. +//! +//! While the actual ClickBench queries often also include some sort of aggregation +//! or limit, this benchmark measures the raw speed of applying filtering +//! and projections, and optimize the performance of the `parquet` filter +//! evaluation in real world scenarios. +//! +//! This benchmark uses the hits_1.parquet file, is a 100,000 row samples of +//! the entire 100M row dataset. It is reasonable in size and speed to run +//! and seems to be a good representative of the entire dataset. +//! +//! See also `arrow_reader_row_filter` for more focused filter evaluation microbenchmarks +//! +//! [ClickBench]: https://benchmark.clickhouse.com/ + +use arrow::compute::kernels::cmp::{eq, neq}; +use arrow::compute::{like, nlike, or}; +use arrow_array::types::{Int16Type, Int32Type, Int64Type}; +use arrow_array::{ArrayRef, ArrowPrimitiveType, BooleanArray, PrimitiveArray, StringViewArray}; +use arrow_schema::{ArrowError, DataType, Schema}; +use bytes::Bytes; +use criterion::{Criterion, criterion_group, criterion_main}; +use futures::future::BoxFuture; +use futures::{FutureExt, StreamExt, TryFutureExt}; +use object_store::local::LocalFileSystem; +use object_store::{GetOptions, GetRange, ObjectStore, ObjectStoreExt}; +use parquet::arrow::arrow_reader::{ + ArrowPredicate, ArrowPredicateFn, ArrowReaderMetadata, ArrowReaderOptions, + ParquetRecordBatchReaderBuilder, RowFilter, +}; +use parquet::arrow::async_reader::{AsyncFileReader, MetadataSuffixFetch}; +use parquet::arrow::{ParquetRecordBatchStreamBuilder, ProjectionMask}; +use parquet::errors::ParquetError; +use parquet::file::metadata::PageIndexPolicy; +use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader}; +use parquet::schema::types::SchemaDescriptor; +use std::fmt::{Display, Formatter}; +use std::path::{Path, PathBuf}; +use std::sync::{Arc, OnceLock}; + +fn async_reader(c: &mut Criterion) { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + let mut async_group = c.benchmark_group("arrow_reader_clickbench/async"); + let handle = rt.handle(); + for query in all_queries() { + let query_name = query.to_string(); + let read_test = ReadTest::new(query); + async_group.bench_function(query_name, |b| { + b.iter(|| handle.block_on(async { read_test.run_async().await })) + }); + } +} + +fn async_reader_object_store(c: &mut Criterion) { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + let mut async_group = c.benchmark_group("arrow_reader_clickbench/async_object_store"); + let handle = rt.handle(); + for query in all_queries() { + let query_name = query.to_string(); + let read_test = ReadTest::new(query); + async_group.bench_function(query_name, |b| { + b.iter(|| handle.block_on(async { read_test.run_async_object_store().await })) + }); + } +} + +fn sync_reader(c: &mut Criterion) { + let mut sync_group = c.benchmark_group("arrow_reader_clickbench/sync"); + for query in all_queries() { + let query_name = query.to_string(); + let read_test = ReadTest::new(query); + sync_group.bench_function(query_name, |b| b.iter(|| read_test.run_sync())); + } +} + +criterion_group!( + benches, + sync_reader, + async_reader, + async_reader_object_store +); +criterion_main!(benches); + +fn to_parquet_err(e: object_store::Error) -> ParquetError { + ParquetError::External(Box::new(e)) +} + +/// An [`AsyncFileReader`] reading via an [`ObjectStore`], mirroring the +/// example on the [`AsyncFileReader`] trait documentation +#[derive(Clone)] +struct ObjectStoreReader { + store: Arc, + path: object_store::path::Path, +} + +impl AsyncFileReader for ObjectStoreReader { + fn get_bytes( + &mut self, + range: std::ops::Range, + ) -> BoxFuture<'_, Result> { + self.store + .get_range(&self.path, range) + .map_err(to_parquet_err) + .boxed() + } + + fn get_byte_ranges( + &mut self, + ranges: Vec>, + ) -> BoxFuture<'_, Result, ParquetError>> { + async move { + self.store + .get_ranges(&self.path, &ranges) + .await + .map_err(to_parquet_err) + } + .boxed() + } + + fn get_metadata<'a>( + &'a mut self, + options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, Result, ParquetError>> { + async move { + let metadata = ParquetMetaDataReader::new() + .with_arrow_reader_options(options) + .load_via_suffix_and_finish(self) + .await?; + Ok(Arc::new(metadata)) + } + .boxed() + } +} + +impl MetadataSuffixFetch for &mut ObjectStoreReader { + fn fetch_suffix(&mut self, suffix: usize) -> BoxFuture<'_, Result> { + let options = GetOptions { + range: Some(GetRange::Suffix(suffix as u64)), + ..Default::default() + }; + async move { + let resp = self + .store + .get_opts(&self.path, options) + .await + .map_err(to_parquet_err)?; + resp.bytes().await.map_err(to_parquet_err) + } + .boxed() + } +} + +/// Predicate Function. +/// +/// Functions are invoked with the requested array and return a [`BooleanArray`] +/// as described in [`ArrowPredicate::evaluate`]. +type ColumnPredicateFn = + dyn FnMut(&ArrayRef) -> Result + Send + Sync + 'static; + +/// ClickBench query pattern: a particular set of filter and projections used in the +/// [ClickBench queries] when run in [Apache DataFusion]. +/// +/// [ClickBench queries]: https://github.com/apache/datafusion/blob/main/benchmarks/queries/clickbench/queries.sql +/// [Apache DataFusion]: https://datafusion.apache.org/ +struct Query { + /// Human identifiable name + name: &'static str, + /// Which columns will be passed to the predicate functions? + /// + /// Must be in the same order as the columns in the schema + filter_columns: Vec<&'static str>, + /// Which columns will be projected (decoded after applying filters) + projection_columns: Vec<&'static str>, + /// Predicates to apply + predicates: Vec, + /// How many rows are expected to pass the predicate? This serves + /// as a sanity check that the benchmark is working correctly. + expected_row_count: usize, +} + +/// Table that describes each relevant query pattern in the ClickBench queries +fn all_queries() -> Vec { + vec![ + // Q0: SELECT COUNT(*) FROM hits; + // (no filters) + // Q1: SELECT COUNT(*) FROM hits WHERE "AdvEngineID" <> 0; + Query { + name: "Q1", + filter_columns: vec!["AdvEngineID"], + projection_columns: vec!["AdvEngineID"], + predicates: vec![ClickBenchPredicate::neq_literal::(0, 0)], + expected_row_count: 3312, + }, + // no filters in Q2-Q9, Q7 is same filter and projection as Q1 + // Q2: SELECT SUM("AdvEngineID"), COUNT(*), AVG("ResolutionWidth") FROM hits; + // Q3: SELECT AVG("UserID") FROM hits; + // Q4: SELECT COUNT(DISTINCT "UserID") FROM hits; + // Q5: SELECT COUNT(DISTINCT "SearchPhrase") FROM hits; + // Q6: SELECT MIN("EventDate"), MAX("EventDate") FROM hits; + // Q7: SELECT "AdvEngineID", COUNT(*) FROM hits WHERE "AdvEngineID" <> 0 GROUP BY "AdvEngineID" ORDER BY COUNT(*) DESC; + // Q8: SELECT "RegionID", COUNT(DISTINCT "UserID") AS u FROM hits GROUP BY "RegionID" ORDER BY u DESC LIMIT 10; + // Q9: SELECT "RegionID", SUM("AdvEngineID"), COUNT(*) AS c, AVG("ResolutionWidth"), COUNT(DISTINCT "UserID") FROM hits GROUP BY "RegionID" ORDER BY c DESC LIMIT 10; + // Q10: SELECT "MobilePhoneModel", COUNT(DISTINCT "UserID") AS u FROM hits WHERE "MobilePhoneModel" <> '' GROUP BY "MobilePhoneModel" ORDER BY u DESC LIMIT 10; + Query { + name: "Q10", + filter_columns: vec!["MobilePhoneModel"], + projection_columns: vec!["MobilePhoneModel", "UserID"], + predicates: vec![ClickBenchPredicate::not_empty(0)], + expected_row_count: 34276, + }, + // Q11: SELECT "MobilePhone", "MobilePhoneModel", COUNT(DISTINCT "UserID") AS u FROM hits WHERE "MobilePhoneModel" <> '' GROUP BY "MobilePhone", "MobilePhoneModel" ORDER BY u DESC LIMIT 10; + Query { + name: "Q11", + filter_columns: vec!["MobilePhoneModel"], + projection_columns: vec!["MobilePhone", "MobilePhoneModel", "UserID"], + predicates: vec![ClickBenchPredicate::not_empty(0)], + expected_row_count: 34276, + }, + // Q12: SELECT "SearchPhrase", COUNT(*) AS c FROM hits WHERE "SearchPhrase" <> '' GROUP BY "SearchPhrase" ORDER BY c DESC LIMIT 10; + Query { + name: "Q12", + filter_columns: vec!["SearchPhrase"], + projection_columns: vec!["SearchPhrase"], + predicates: vec![ClickBenchPredicate::not_empty(0)], + expected_row_count: 131559, + }, + // Q13: SELECT "SearchPhrase", COUNT(DISTINCT "UserID") AS u FROM hits WHERE "SearchPhrase" <> '' GROUP BY "SearchPhrase" ORDER BY u DESC LIMIT 10; + Query { + name: "Q13", + filter_columns: vec!["SearchPhrase"], + projection_columns: vec!["SearchPhrase", "UserID"], + predicates: vec![ClickBenchPredicate::not_empty(0)], + expected_row_count: 131559, + }, + // Q14: SELECT "SearchEngineID", "SearchPhrase", COUNT(*) AS c FROM hits WHERE "SearchPhrase" <> '' GROUP BY "SearchEngineID", "SearchPhrase" ORDER BY c DESC LIMIT 10; + Query { + name: "Q14", + filter_columns: vec!["SearchPhrase"], + projection_columns: vec!["SearchEngineID", "SearchPhrase"], + predicates: vec![ClickBenchPredicate::not_empty(0)], + expected_row_count: 131559, + }, + // No predicates in Q15-Q18 + // Q15: SELECT "UserID", COUNT(*) FROM hits GROUP BY "UserID" ORDER BY COUNT(*) DESC LIMIT 10; + // Q16: SELECT "UserID", "SearchPhrase", COUNT(*) FROM hits GROUP BY "UserID", "SearchPhrase" ORDER BY COUNT(*) DESC LIMIT 10; + // Q17: SELECT "UserID", "SearchPhrase", COUNT(*) FROM hits GROUP BY "UserID", "SearchPhrase" LIMIT 10; + // Q18: SELECT "UserID", extract(minute FROM to_timestamp_seconds("EventTime")) AS m, "SearchPhrase", COUNT(*) FROM hits GROUP BY "UserID", m, "SearchPhrase" ORDER BY COUNT(*) DESC LIMIT 10; + // Q19: SELECT "UserID" FROM hits WHERE "UserID" = 435090932899640449; + Query { + name: "Q19", + filter_columns: vec!["UserID"], + projection_columns: vec!["UserID"], + predicates: vec![ + // Original predicate is `UserID = 435090932899640449` + // However, 435090932899640449 does not exist in hits_0.parquet. There are 4 rows in the total dataset: + // ```sql + // > select count(*) from 'hits.parquet' where "UserID" = 435090932899640449; + // +----------+ + // | count(*) | + // +----------+ + // | 4 | + // +----------+ + // 1 row(s) fetched. + // ``` + // + // So use a user ID that actually exists in hits_0 and has 4 rows as well: + // + // `sql + // > select "UserID", count(*) as c from 'hits_1.parquet' GROUP BY "UserID" HAVING c = 4; + // +---------------------+---+ + // | UserID | c | + // +---------------------+---+ + // | 3233473875476175636 | 4 | + // | ... | . | + // +---------------------+---+ + // ``` + ClickBenchPredicate::eq_literal::(0, 3233473875476175636), + ], + expected_row_count: 4, + }, + // Q20: SELECT COUNT(*) FROM hits WHERE "URL" LIKE '%google%'; + Query { + name: "Q20", + filter_columns: vec!["URL"], + projection_columns: vec!["URL"], + predicates: vec![ClickBenchPredicate::like_google(0)], + expected_row_count: 137, + }, + // Q21: SELECT "SearchPhrase", MIN("URL"), COUNT(*) AS c FROM hits WHERE "URL" LIKE '%google%' AND "SearchPhrase" <> '' GROUP BY "SearchPhrase" ORDER BY c DESC LIMIT 10; + Query { + name: "Q21", + filter_columns: vec!["URL", "SearchPhrase"], + projection_columns: vec!["SearchPhrase", "URL"], + predicates: vec![ + ClickBenchPredicate::like_google(0), + ClickBenchPredicate::not_empty(1), + ], + expected_row_count: 16, + }, + // Q22: SELECT "SearchPhrase", MIN("URL"), MIN("Title"), COUNT(*) AS c, COUNT(DISTINCT "UserID") FROM hits WHERE "Title" LIKE '%Google%' AND "URL" NOT LIKE '%.google.%' AND "SearchPhrase" <> '' GROUP BY "SearchPhrase" ORDER BY c DESC LIMIT 10; + Query { + name: "Q22", + filter_columns: vec!["Title", "URL", "SearchPhrase"], + projection_columns: vec!["SearchPhrase", "URL", "Title", "UserID"], + predicates: vec![ + ClickBenchPredicate::like_Google(0), + ClickBenchPredicate::nlike_google(1), + ClickBenchPredicate::not_empty(2), + ], + expected_row_count: 46, + }, + // Q23: SELECT * FROM hits WHERE "URL" LIKE '%google%'; + Query { + name: "Q23", + filter_columns: vec!["URL"], + projection_columns: vec!["*"], // all columns + predicates: vec![ClickBenchPredicate::like_google(0)], + expected_row_count: 137, + }, + // Q24: SELECT "SearchPhrase" FROM hits WHERE "SearchPhrase" <> '' ORDER BY "EventTime" LIMIT 10; + Query { + name: "Q24", + filter_columns: vec!["SearchPhrase"], + projection_columns: vec!["SearchPhrase", "EventTime"], + predicates: vec![ClickBenchPredicate::not_empty(0)], + expected_row_count: 131559, + }, + // Same filters and projection as Q12 + // Q25: SELECT "SearchPhrase" FROM hits WHERE "SearchPhrase" <> '' ORDER BY "SearchPhrase" LIMIT 10; + // Same filters and projection as Q24 + // Q26: SELECT "SearchPhrase" FROM hits WHERE "SearchPhrase" <> '' ORDER BY "EventTime", "SearchPhrase" LIMIT 10; + // Q27: SELECT "CounterID", AVG(length("URL")) AS l, COUNT(*) AS c FROM hits WHERE "URL" <> '' GROUP BY "CounterID" HAVING COUNT(*) > 100000 ORDER BY l DESC LIMIT 25; + Query { + name: "Q27", + filter_columns: vec!["URL"], + projection_columns: vec!["CounterID", "URL"], + predicates: vec![ClickBenchPredicate::not_empty(0)], + expected_row_count: 999978, + }, + // Q28: SELECT REGEXP_REPLACE("Referer", '^https?://(?:www\.)?([^/]+)/.*$', '\1') AS k, AVG(length("Referer")) AS l, COUNT(*) AS c, MIN("Referer") FROM hits WHERE "Referer" <> '' GROUP BY k HAVING COUNT(*) > 100000 ORDER BY l DESC LIMIT 25; + Query { + name: "Q28", + filter_columns: vec!["Referer"], + projection_columns: vec!["Referer"], + predicates: vec![ClickBenchPredicate::not_empty(0)], + expected_row_count: 925813, + }, + // No predicates in Q29 + // Q29: SELECT SUM("ResolutionWidth"), SUM("ResolutionWidth" + 1), SUM("ResolutionWidth" + 2), SUM("ResolutionWidth" + 3), SUM("ResolutionWidth" + 4), SUM("ResolutionWidth" + 5), SUM("ResolutionWidth" + 6), SUM("ResolutionWidth" + 7), SUM("ResolutionWidth" + 8), SUM("ResolutionWidth" + 9), SUM("ResolutionWidth" + 10), SUM("ResolutionWidth" + 11), SUM("ResolutionWidth" + 12), SUM("ResolutionWidth" + 13), SUM("ResolutionWidth" + 14), SUM("ResolutionWidth" + 15), SUM("ResolutionWidth" + 16), SUM("ResolutionWidth" + 17), SUM("ResolutionWidth" + 18), SUM("ResolutionWidth" + 19), SUM("ResolutionWidth" + 20), SUM("ResolutionWidth" + 21), SUM("ResolutionWidth" + 22), SUM("ResolutionWidth" + 23), SUM("ResolutionWidth" + 24), SUM("ResolutionWidth" + 25), SUM("ResolutionWidth" + 26), SUM("ResolutionWidth" + 27), SUM("ResolutionWidth" + 28), SUM("ResolutionWidth" + 29), SUM("ResolutionWidth" + 30), SUM("ResolutionWidth" + 31), SUM("ResolutionWidth" + 32), SUM("ResolutionWidth" + 33), SUM("ResolutionWidth" + 34), SUM("ResolutionWidth" + 35), SUM("ResolutionWidth" + 36), SUM("ResolutionWidth" + 37), SUM("ResolutionWidth" + 38), SUM("ResolutionWidth" + 39), SUM("ResolutionWidth" + 40), SUM("ResolutionWidth" + 41), SUM("ResolutionWidth" + 42), SUM("ResolutionWidth" + 43), SUM("ResolutionWidth" + 44), SUM("ResolutionWidth" + 45), SUM("ResolutionWidth" + 46), SUM("ResolutionWidth" + 47), SUM("ResolutionWidth" + 48), SUM("ResolutionWidth" + 49), SUM("ResolutionWidth" + 50), SUM("ResolutionWidth" + 51), SUM("ResolutionWidth" + 52), SUM("ResolutionWidth" + 53), SUM("ResolutionWidth" + 54), SUM("ResolutionWidth" + 55), SUM("ResolutionWidth" + 56), SUM("ResolutionWidth" + 57), SUM("ResolutionWidth" + 58), SUM("ResolutionWidth" + 59), SUM("ResolutionWidth" + 60), SUM("ResolutionWidth" + 61), SUM("ResolutionWidth" + 62), SUM("ResolutionWidth" + 63), SUM("ResolutionWidth" + 64), SUM("ResolutionWidth" + 65), SUM("ResolutionWidth" + 66), SUM("ResolutionWidth" + 67), SUM("ResolutionWidth" + 68), SUM("ResolutionWidth" + 69), SUM("ResolutionWidth" + 70), SUM("ResolutionWidth" + 71), SUM("ResolutionWidth" + 72), SUM("ResolutionWidth" + 73), SUM("ResolutionWidth" + 74), SUM("ResolutionWidth" + 75), SUM("ResolutionWidth" + 76), SUM("ResolutionWidth" + 77), SUM("ResolutionWidth" + 78), SUM("ResolutionWidth" + 79), SUM("ResolutionWidth" + 80), SUM("ResolutionWidth" + 81), SUM("ResolutionWidth" + 82), SUM("ResolutionWidth" + 83), SUM("ResolutionWidth" + 84), SUM("ResolutionWidth" + 85), SUM("ResolutionWidth" + 86), SUM("ResolutionWidth" + 87), SUM("ResolutionWidth" + 88), SUM("ResolutionWidth" + 89) FROM hits; + // Q30: SELECT "SearchEngineID", "ClientIP", COUNT(*) AS c, SUM("IsRefresh"), AVG("ResolutionWidth") FROM hits WHERE "SearchPhrase" <> '' GROUP BY "SearchEngineID", "ClientIP" ORDER BY c DESC LIMIT 10; + Query { + name: "Q30", + filter_columns: vec!["SearchPhrase"], + projection_columns: vec!["SearchEngineID", "ClientIP", "IsRefresh", "ResolutionWidth"], + predicates: vec![ClickBenchPredicate::not_empty(0)], + expected_row_count: 131559, + }, + // Same filters and projection as Q30 + // Q31: SELECT "WatchID", "ClientIP", COUNT(*) AS c, SUM("IsRefresh"), AVG("ResolutionWidth") FROM hits WHERE "SearchPhrase" <> '' GROUP BY "WatchID", "ClientIP" ORDER BY c DESC LIMIT 10; + // No predicates in Q32-Q35 + // Q32: SELECT "WatchID", "ClientIP", COUNT(*) AS c, SUM("IsRefresh"), AVG("ResolutionWidth") FROM hits GROUP BY "WatchID", "ClientIP" ORDER BY c DESC LIMIT 10; + // Q33: SELECT "URL", COUNT(*) AS c FROM hits GROUP BY "URL" ORDER BY c DESC LIMIT 10; + // Q34: SELECT 1, "URL", COUNT(*) AS c FROM hits GROUP BY 1, "URL" ORDER BY c DESC LIMIT 10; + // Q35: SELECT "ClientIP", "ClientIP" - 1, "ClientIP" - 2, "ClientIP" - 3, COUNT(*) AS c FROM hits GROUP BY "ClientIP", "ClientIP" - 1, "ClientIP" - 2, "ClientIP" - 3 ORDER BY c DESC LIMIT 10; + // Q36: SELECT "URL", COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate" >= '2013-07-01' AND "EventDate" <= '2013-07-31' AND "DontCountHits" = 0 AND "IsRefresh" = 0 AND "URL" <> '' GROUP BY "URL" ORDER BY PageViews DESC LIMIT 10; + Query { + name: "Q36", + filter_columns: vec![ + "CounterID", + "EventDate", + "DontCountHits", + "IsRefresh", + "URL", + ], + projection_columns: vec!["URL"], + predicates: vec![ + ClickBenchPredicate::eq_literal::(0, 62), + // For now, elide the EventDate predicates + // The predicates don't contribute to the filter evaluation + // + // Note the predicates on EventDate are not pushed down into the Parquet scan in DataFusion + // This is because the EventDate column is a Int16 and DataFusion doesn't automatically + // coerce from Int16 to Date + // + // DataFusion does `CAST(EventDate AS VARCHAR)` which is not + // obviously correct in this case. + // + // You can cast the Int16 to Date like this: + // ```sql + // > select "EventDate"::int::date from 'hits.parquet' limit 10; + // +------------------------+ + // | hits.parquet.EventDate | + // +------------------------+ + // | 2013-07-28 | + // | .... | + // +------------------------+ + // ``` + // + // ClickBenchPredicate::gt_eq_literal::(1, str_to_i16_date("2013-07-01")), + // ClickBenchPredicate::lt_eq_literal::(1, str_to_i16_date("2013-07-31")), + ClickBenchPredicate::eq_literal::(2, 0), + ClickBenchPredicate::eq_literal::(3, 0), + ClickBenchPredicate::not_empty(4), + ], + expected_row_count: 181198, + }, + // Q37: SELECT "Title", COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate" >= '2013-07-01' AND "EventDate" <= '2013-07-31' AND "DontCountHits" = 0 AND "IsRefresh" = 0 AND "Title" <> '' GROUP BY "Title" ORDER BY PageViews DESC LIMIT 10; + Query { + name: "Q37", + filter_columns: vec![ + "CounterID", + "EventDate", + "DontCountHits", + "IsRefresh", + "Title", + ], + projection_columns: vec!["Title"], + predicates: vec![ + ClickBenchPredicate::eq_literal::(0, 62), + // see Q36 for EventDate predicates + //ClickBenchPredicate::gt_eq_literal::(1, str_to_i16_date("2013-07-01")), + //ClickBenchPredicate::lt_eq_literal::(1, str_to_i16_date("2013-07-31")), + ClickBenchPredicate::eq_literal::(2, 0), + ClickBenchPredicate::eq_literal::(3, 0), + ClickBenchPredicate::not_empty(4), + ], + expected_row_count: 178323, + }, + // Q38: SELECT "URL", COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate" >= '2013-07-01' AND "EventDate" <= '2013-07-31' AND "IsRefresh" = 0 AND "IsLink" <> 0 AND "IsDownload" = 0 GROUP BY "URL" ORDER BY PageViews DESC LIMIT 10 OFFSET 1000; + Query { + name: "Q38", + filter_columns: vec![ + "CounterID", + "EventDate", + "IsRefresh", + "IsLink", + "IsDownload", + ], + projection_columns: vec!["URL"], + predicates: vec![ + ClickBenchPredicate::eq_literal::(0, 62), + // see Q36 for EventDate predicates + // ClickBenchPredicate::gt_eq_literal::(1, str_to_i16_date("2013-07-01")), + // ClickBenchPredicate::lt_eq_literal::(1, str_to_i16_date("2013-07-31")), + ClickBenchPredicate::eq_literal::(2, 0), + ClickBenchPredicate::neq_literal::(3, 0), + ClickBenchPredicate::neq_literal::(4, 0), + ], + expected_row_count: 419, + }, + // Q39: SELECT "TraficSourceID", "SearchEngineID", "AdvEngineID", CASE WHEN ("SearchEngineID" = 0 AND "AdvEngineID" = 0) THEN "Referer" ELSE '' END AS Src, "URL" AS Dst, COUNT(*) AS PageViews FROM hits + // WHERE "CounterID" = 62 AND "EventDate" >= '2013-07-01' AND "EventDate" <= '2013-07-31' AND "IsRefresh" = 0 GROUP BY "TraficSourceID", "SearchEngineID", "AdvEngineID", Src, Dst ORDER BY PageViews DESC LIMIT 10 OFFSET 1000; + Query { + name: "Q39", + filter_columns: vec!["CounterID", "EventDate", "IsRefresh"], + projection_columns: vec![ + "TraficSourceID", + "SearchEngineID", + "AdvEngineID", + "Referer", + "URL", + ], + predicates: vec![ + ClickBenchPredicate::eq_literal::(0, 62), + // see Q36 for EventDate predicates + // ClickBenchPredicate::gt_eq_literal::(1, str_to_i16_date("2013-07-01")), + // ClickBenchPredicate::lt_eq_literal::(1, str_to_i16_date("2013-07-31")), + ClickBenchPredicate::eq_literal::(2, 0), + ], + expected_row_count: 194225, + }, + // Q40: SELECT "URLHash", "EventDate", COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate" >= '2013-07-01' AND "EventDate" <= '2013-07-31' AND "IsRefresh" = 0 AND "TraficSourceID" IN (-1, 6) AND "RefererHash" = 3594120000172545465 GROUP BY "URLHash", "EventDate" ORDER BY PageViews DESC LIMIT 10 OFFSET 100; + Query { + name: "Q40", + filter_columns: vec![ + "CounterID", + "EventDate", + "IsRefresh", + "TraficSourceID", + "RefererHash", + ], + projection_columns: vec!["URLHash", "EventDate"], + predicates: vec![ + ClickBenchPredicate::eq_literal::(0, 62), + // see Q36 for EventDate predicates + // ClickBenchPredicate::gt_eq_literal::(1, str_to_i16_date("2013-07-01")), + // ClickBenchPredicate::lt_eq_literal::(1, str_to_i16_date("2013-07-31")), + ClickBenchPredicate::eq_literal::(2, 0), + ClickBenchPredicate::in_list::(3, (-1, 6)), // IN -1, 6 + ClickBenchPredicate::eq_literal::(4, 3594120000172545465), + ], + expected_row_count: 24793, + }, + // Q41: SELECT "WindowClientWidth", "WindowClientHeight", COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate" >= '2013-07-01' AND "EventDate" <= '2013-07-31' AND "IsRefresh" = 0 AND "DontCountHits" = 0 AND "URLHash" = 2868770270353813622 GROUP BY "WindowClientWidth", "WindowClientHeight" ORDER BY PageViews DESC LIMIT 10 OFFSET 10000; + Query { + name: "Q41", + filter_columns: vec![ + "CounterID", + "EventDate", + "IsRefresh", + "DontCountHits", + "URLHash", + ], + projection_columns: vec!["WindowClientWidth", "WindowClientHeight"], + predicates: vec![ + ClickBenchPredicate::eq_literal::(0, 62), + // see Q36 for EventDate predicates + // ClickBenchPredicate::gt_eq_literal::(1, str_to_i16_date("2013-07-01")), + // ClickBenchPredicate::lt_eq_literal::(1, str_to_i16_date("2013-07-31")), + ClickBenchPredicate::eq_literal::(2, 0), + ClickBenchPredicate::eq_literal::(3, 0), + ClickBenchPredicate::eq_literal::(4, 2868770270353813622), + ], + expected_row_count: 29201, + }, + // Q42: SELECT DATE_TRUNC('minute', to_timestamp_seconds("EventTime")) AS M, COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate" >= '2013-07-14' AND "EventDate" <= '2013-07-15' AND "IsRefresh" = 0 AND "DontCountHits" = 0 GROUP BY DATE_TRUNC('minute', to_timestamp_seconds("EventTime")) ORDER BY DATE_TRUNC('minute', M) LIMIT 10 OFFSET 1000; + Query { + name: "Q42", + filter_columns: vec!["CounterID", "EventDate", "IsRefresh", "DontCountHits"], + projection_columns: vec!["EventTime"], + predicates: vec![ + ClickBenchPredicate::eq_literal::(0, 62), + // see Q36 for EventDate predicates + // ClickBenchPredicate::gt_eq_literal::(1, str_to_i16_date("2013-07-14")), + // ClickBenchPredicate::lt_eq_literal::(1, str_to_i16_date("2013-07-15")), + ClickBenchPredicate::eq_literal::(2, 0), + ClickBenchPredicate::eq_literal::(3, 0), + ], + expected_row_count: 181200, + }, + ] +} + +/// Evaluate a predicate the input column specified offset relative to +/// the provided filter column +/// +/// This provides a level of indirection needed because the columns passed to +/// ArrowPredicateFn is in the order of the schema rather than the order of the +/// specified predicate columns +struct ClickBenchPredicate { + column_index: usize, + /// Function that makes the actual [`ColumnPredicateFn`] + /// + /// This is necessary (and awkward) because `ArrowPredicateFn` does not + /// implement `Clone`, so it must be created for each reader instance. + predicate_factory: Box Box>, +} + +impl ClickBenchPredicate { + /// Create a new `ClickBenchPredicate` + /// + /// Parameters: + /// * `column_index`: the index of the column in the `filter_columns` list + /// * `f`: creates the actual predicate function + fn new Box + 'static>( + column_index: usize, + predicate_factory: F, + ) -> ClickBenchPredicate { + Self { + column_index, + predicate_factory: Box::new(predicate_factory), + } + } + + fn column_index(&self) -> usize { + self.column_index + } + + /// return a predicate function + fn predicate_fn(&self) -> Box { + (self.predicate_factory)() + } + + /// Create Predicate: col = literal + fn eq_literal(column_index: usize, literal_value: T::Native) -> Self { + Self::new(column_index, move || { + let literal = PrimitiveArray::::new_scalar(literal_value); + Box::new(move |col| eq(col, &literal)) + }) + } + + /// Create Predicate: col IN (lit1, lit2) + fn in_list( + column_index: usize, + literal_values: (T::Native, T::Native), + ) -> Self { + Self::new(column_index, move || { + let literal_1 = PrimitiveArray::::new_scalar(literal_values.0); + let literal_2 = PrimitiveArray::::new_scalar(literal_values.1); + Box::new(move |col| { + // evaluate like (col = lit1) OR (col = lit2) + let match1 = eq(&col, &literal_1)?; + let match2 = eq(&col, &literal_2)?; + or(&match1, &match2) + }) + }) + } + + /// Create predicate: col != '' + fn neq_literal(column_index: usize, literal_value: T::Native) -> Self { + Self::new(column_index, move || { + let literal = PrimitiveArray::::new_scalar(literal_value); + Box::new(move |col| neq(col, &literal)) + }) + } + + /// Create Predicate: col <> '' + fn not_empty(column_index: usize) -> Self { + Self::new(column_index, move || { + let empty_string = StringViewArray::new_scalar(""); + Box::new(move |col| neq(col, &empty_string)) + }) + } + + /// Create Predicate: col LIKE '%google%' + fn like_google(column_index: usize) -> Self { + Self::new(column_index, move || { + let google_url = StringViewArray::new_scalar("%google%"); + Box::new(move |col| like(col, &google_url)) + }) + } + + /// Create Predicate: col NOT LIKE '%google%' + fn nlike_google(column_index: usize) -> Self { + Self::new(column_index, move || { + let google_url = StringViewArray::new_scalar("%google%"); + Box::new(move |col| nlike(col, &google_url)) + }) + } + + /// Create Predicate: col LIKE '%Google%' + #[allow(non_snake_case)] + fn like_Google(column_index: usize) -> Self { + Self::new(column_index, move || { + let google_url = StringViewArray::new_scalar("%Google%"); + Box::new(move |col| like(col, &google_url)) + }) + } +} + +impl Display for Query { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.name) + } +} + +/// FULL path to the ClickBench hits_1.parquet file +static HITS_1_PATH: OnceLock = OnceLock::new(); + +/// Finds the paths to the ClickBench file, downloading it if not found +fn hits_1() -> &'static Path { + HITS_1_PATH.get_or_init(|| { + let current_dir = std::env::current_dir().expect("Failed to get current directory"); + println!( + "Looking for ClickBench files starting in current_dir and all parent directories: {current_dir:?}" + ); + + if let Some(hits_1_path) = find_file_if_exists(current_dir.clone(), "hits_1.parquet") { + return hits_1_path; + } + + // File not found, download it + let download_path = current_dir.join("hits_1.parquet"); + let url = "https://datasets.clickhouse.com/hits_compatible/athena_partitioned/hits_1.parquet"; + println!("hits_1.parquet not found, downloading from {url}..."); + + let status = std::process::Command::new("wget") + .args(["--continue", "-O"]) + .arg(&download_path) + .arg(url) + .status() + .expect("Failed to execute wget. Please install wget or download manually."); + + assert!( + status.success(), + "Failed to download hits_1.parquet. You can download it manually via:\n\ + wget --continue {url}" + ); + + println!("Downloaded hits_1.parquet to {download_path:?}"); + download_path + }) +} + +/// Searches for `file_name` in `current_dir` and all parent directories +fn find_file_if_exists(mut current_dir: PathBuf, file_name: &str) -> Option { + loop { + let potential_file_path = current_dir.join(file_name); + if potential_file_path.exists() { + return Some(potential_file_path); + } + // didn't find in current path, try parent + let Some(parent) = current_dir.parent() else { + break; + }; + current_dir = parent.to_path_buf(); + } + None +} + +/// Encapsulates the test parameters for a single benchmark +struct ReadTest { + /// Human identifiable name + name: &'static str, + /// Metadata from Parquet file + arrow_reader_metadata: ArrowReaderMetadata, + /// Which columns in the file should be projected (decoded after filter)? + projection_mask: ProjectionMask, + /// Schema indices for each filter column (in filter_columns order) + filter_schema_indices: Vec, + /// Predicates to apply + predicates: Vec, + /// How many rows are expected to pass the predicate? + /// + /// This value is a sanity check that the benchmark is working correctly. + expected_row_count: usize, +} + +impl ReadTest { + fn new(query: Query) -> Self { + let Query { + name, + filter_columns, + projection_columns, + predicates, + expected_row_count, + } = query; + + let arrow_reader_metadata = load_metadata(hits_1()); + let schema_descr = arrow_reader_metadata + .metadata() + .file_metadata() + .schema_descr(); + + // Determine the correct selection ("ProjectionMask") + let projection_mask = if projection_columns.contains(&"*") { + // * means all columns + ProjectionMask::all() + } else { + let projection_schema_indices = column_indices(schema_descr, &projection_columns); + ProjectionMask::leaves(schema_descr, projection_schema_indices) + }; + + let filter_schema_indices = column_indices(schema_descr, &filter_columns); + + Self { + name, + arrow_reader_metadata, + projection_mask, + filter_schema_indices, + predicates, + expected_row_count, + } + } + + /// Run the filter and projection using the async `ParquetRecordBatchStream` + /// reader: + /// + /// 1. open the file + /// 2. apply the filter + /// 3. read the projection_columns + async fn run_async(&self) { + let Ok(parquet_file) = tokio::fs::File::open(hits_1()).await else { + panic!("Failed to open {:?}", hits_1()); + }; + + // setup the reader + let mut stream = ParquetRecordBatchStreamBuilder::new_with_metadata( + parquet_file, + self.arrow_reader_metadata.clone(), + ) + .with_batch_size(8192) + .with_projection(self.projection_mask.clone()) + .with_row_filter(self.row_filter()) + .build() + .unwrap(); + + // run the stream to its end + let mut row_count = 0; + while let Some(b) = stream.next().await { + let b = b.unwrap(); + let num_rows = b.num_rows(); + row_count += num_rows; + } + self.check_row_count(row_count); + } + + /// Like [`Self::run_async`] but reading via an [`ObjectStore`] based reader + async fn run_async_object_store(&self) { + let hits_path = hits_1(); + let parent = hits_path.parent().unwrap(); + let file_name = hits_path.file_name().unwrap().to_str().unwrap(); + let store = Arc::new(LocalFileSystem::new_with_prefix(parent).unwrap()); + let location = object_store::path::Path::from(file_name); + + let reader = ObjectStoreReader { + store, + path: location, + }; + + // setup the reader + let mut stream = ParquetRecordBatchStreamBuilder::new_with_metadata( + reader, + self.arrow_reader_metadata.clone(), + ) + .with_batch_size(8192) + .with_projection(self.projection_mask.clone()) + .with_row_filter(self.row_filter()) + .build() + .unwrap(); + + // run the stream to its end + let mut row_count = 0; + while let Some(b) = stream.next().await { + let b = b.unwrap(); + let num_rows = b.num_rows(); + row_count += num_rows; + } + self.check_row_count(row_count); + } + + /// Like [`Self::run_async`] but for the sync parquet reader + fn run_sync(&self) { + let Ok(parquet_file) = std::fs::File::open(hits_1()) else { + panic!("Failed to open {:?}", hits_1()); + }; + + // setup the reader + let reader = ParquetRecordBatchReaderBuilder::new_with_metadata( + parquet_file, + self.arrow_reader_metadata.clone(), + ) + .with_batch_size(8192) + .with_projection(self.projection_mask.clone()) + .with_row_filter(self.row_filter()) + .build() + .unwrap(); + + // run the stream to its end + let mut row_count = 0; + for b in reader { + let b = b.unwrap(); + let num_rows = b.num_rows(); + row_count += num_rows; + } + self.check_row_count(row_count); + } + + /// Return a `RowFilter` to apply to the reader. + /// + /// Each predicate gets a ProjectionMask containing only the single column + /// it needs, rather than all filter columns. This avoids decoding expensive + /// columns (e.g. strings) when evaluating cheap predicates (e.g. integer equality). + fn row_filter(&self) -> RowFilter { + let schema_descr = self + .arrow_reader_metadata + .metadata() + .file_metadata() + .schema_descr(); + + let arrow_predicates: Vec<_> = self + .predicates + .iter() + .map(|pred| { + let schema_index = self.filter_schema_indices[pred.column_index()]; + let predicate_mask = ProjectionMask::leaves(schema_descr, [schema_index]); + let mut predicate_fn = pred.predicate_fn(); + Box::new(ArrowPredicateFn::new(predicate_mask, move |batch| { + (predicate_fn)(batch.column(0)) + })) as Box + }) + .collect(); + + RowFilter::new(arrow_predicates) + } + + fn check_row_count(&self, row_count: usize) { + let expected_row_count = self.expected_row_count; + assert_eq!( + row_count, expected_row_count, + "Expected {} rows, but got {} in {}", + expected_row_count, row_count, self.name, + ); + } +} + +/// Return a map from `column_names` in `filter_columns` to the index in the schema +fn column_indices(schema: &SchemaDescriptor, column_names: &Vec<&str>) -> Vec { + let fields = schema.root_schema().get_fields(); + let mut indices = vec![]; + for &name in column_names { + for (idx, field) in fields.iter().enumerate().take(schema.num_columns()) { + if name == field.name() { + indices.push(idx) + } + } + } + indices +} + +/// Loads Parquet metadata from the given path, including page indexes +fn load_metadata(path: &Path) -> ArrowReaderMetadata { + let file = std::fs::File::open(path).unwrap(); + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::from(true)); + let orig_metadata = + ArrowReaderMetadata::load(&file, options.clone()).expect("parquet-metadata loading failed"); + + // Update the arrow schema so that it reads View types for binary and utf8 columns + let new_fields = orig_metadata + .schema() + .fields() + .iter() + .map(|f| { + //println!("Converting field: {:?}", f); + // Read UTF8 fields as Utf8View + // + // The clickbench_partitioned dataset has textual fields listed as + // binary for some historical reason so translate Binary to Utf8View + if matches!( + f.data_type(), + DataType::Utf8 | DataType::Binary | DataType::BinaryView + ) { + let new_field = f.as_ref().clone().with_data_type(DataType::Utf8View); + Arc::new(new_field) + } else { + // otherwise use the inferred field type + Arc::clone(f) + } + }) + .collect::>(); + + let new_arrow_schema = Arc::new(Schema::new(new_fields)); + + let new_options = options.with_schema(new_arrow_schema); + ArrowReaderMetadata::try_new(Arc::clone(orig_metadata.metadata()), new_options).unwrap() +} diff --git a/vendor/parquet/benches/arrow_reader_peak_memory.rs b/vendor/parquet/benches/arrow_reader_peak_memory.rs new file mode 100644 index 0000000..2a717ec --- /dev/null +++ b/vendor/parquet/benches/arrow_reader_peak_memory.rs @@ -0,0 +1,619 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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 memory usage of ListArrayReader for sparse +//! list columns. +//! +//! Uses a global allocator with thread-local memory tracking so that +//! concurrent criterion threads don't interfere with measurements. + +use std::alloc::Layout; +use std::cell::Cell; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field}; +use criterion::measurement::{Measurement, ValueFormatter}; +use criterion::{BenchmarkGroup, Criterion, Throughput, criterion_group, criterion_main}; +use parquet::arrow::array_reader::make_byte_array_reader; +use parquet::arrow::array_reader::{ + ArrayReader, ListArrayReader, PrimitiveArrayReader, make_fixed_len_byte_array_reader, +}; +use parquet::arrow::arrow_reader::DEFAULT_BATCH_SIZE; +use parquet::basic::Encoding; +use parquet::column::page::PageIterator; +use parquet::data_type::{ByteArrayType, DoubleType, FixedLenByteArrayType, Int32Type}; +use parquet::schema::parser::parse_message_type; +use parquet::schema::types::{ColumnDescPtr, SchemaDescriptor}; +use parquet::util::{DataPageBuilder, DataPageBuilderImpl, InMemoryPageIterator}; +use rand::{Rng, SeedableRng, rngs::StdRng}; + +// --------------------------------------------------------------------------- +// 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)); + }); +} + +#[allow(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 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: &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 +// --------------------------------------------------------------------------- + +const NUM_ROW_GROUPS: usize = 2; +const PAGES_PER_GROUP: usize = 4; +const VALUES_PER_PAGE: usize = 10_000; +const MAX_LIST_LEN: usize = 10; +const BATCH_SIZE: usize = 8192; +const EXPECTED_VALUE_COUNT: usize = NUM_ROW_GROUPS * PAGES_PER_GROUP * VALUES_PER_PAGE; + +fn seedable_rng() -> StdRng { + StdRng::seed_from_u64(42) +} + +fn build_int32_list_schema() -> (SchemaDescriptor, ColumnDescPtr) { + let message_type = " + message schema { + OPTIONAL GROUP int32_list (LIST) { + repeated group list { + optional INT32 element; + } + } + } + "; + let schema = parse_message_type(message_type) + .map(|t| SchemaDescriptor::new(Arc::new(t))) + .unwrap(); + let col_desc = schema.column(0); + (schema, col_desc) +} + +fn build_string_list_schema() -> (SchemaDescriptor, ColumnDescPtr) { + let message_type = " + message schema { + OPTIONAL GROUP string_list (LIST) { + repeated group list { + optional BYTE_ARRAY element (UTF8); + } + } + } + "; + let schema = parse_message_type(message_type) + .map(|t| SchemaDescriptor::new(Arc::new(t))) + .unwrap(); + let col_desc = schema.column(0); + (schema, col_desc) +} + +const FIXED_BYTE_LEN: usize = 32; + +fn build_fixed32_list_schema() -> (SchemaDescriptor, ColumnDescPtr) { + let message_type = " + message schema { + OPTIONAL GROUP fixed32_list (LIST) { + repeated group list { + optional FIXED_LEN_BYTE_ARRAY(32) element; + } + } + } + "; + let schema = parse_message_type(message_type) + .map(|t| SchemaDescriptor::new(Arc::new(t))) + .unwrap(); + let col_desc = schema.column(0); + (schema, col_desc) +} + +fn build_double_list_schema() -> (SchemaDescriptor, ColumnDescPtr) { + let message_type = " + message schema { + OPTIONAL GROUP double_list (LIST) { + repeated group list { + optional DOUBLE element; + } + } + } + "; + let schema = parse_message_type(message_type) + .map(|t| SchemaDescriptor::new(Arc::new(t))) + .unwrap(); + let col_desc = schema.column(0); + (schema, col_desc) +} + +fn build_list_pages( + column_desc: ColumnDescPtr, + null_density: f32, + mut gen_value: impl FnMut(&mut StdRng) -> T::T, +) -> impl PageIterator + Clone { + let max_def_level = column_desc.max_def_level(); + let max_rep_level = column_desc.max_rep_level(); + let mut rng = seedable_rng(); + let mut pages: Vec> = Vec::new(); + + for _ in 0..NUM_ROW_GROUPS { + let mut column_chunk_pages = Vec::new(); + for _ in 0..PAGES_PER_GROUP { + let mut values = Vec::new(); + let mut def_levels = Vec::new(); + let mut rep_levels = Vec::new(); + + for _ in 0..VALUES_PER_PAGE { + rep_levels.push(0); + if rng.random::() < null_density { + def_levels.push(0); + continue; + } + let len = rng.random_range(0..MAX_LIST_LEN); + if len == 0 { + def_levels.push(1); + continue; + } + (1..len).for_each(|_| rep_levels.push(1)); + for _ in 0..len { + if rng.random::() < null_density { + def_levels.push(2); + } else { + def_levels.push(3); + values.push(gen_value(&mut rng)); + } + } + } + + let mut pb = DataPageBuilderImpl::new(column_desc.clone(), values.len() as u32, true); + pb.add_rep_levels(max_rep_level, &rep_levels); + pb.add_def_levels(max_def_level, &def_levels); + pb.add_values::(Encoding::PLAIN, &values); + column_chunk_pages.push(pb.consume()); + } + pages.push(column_chunk_pages); + } + InMemoryPageIterator::new(pages) +} + +fn drain_reader(mut reader: Box) -> usize { + let mut total_count = 0; + loop { + let batch = reader.next_batch(BATCH_SIZE).unwrap(); + let batch_len = batch.len(); + total_count += batch_len; + if batch_len < BATCH_SIZE { + break; + } + } + total_count +} + +// --------------------------------------------------------------------------- +// Benchmarks +// --------------------------------------------------------------------------- + +fn null_density_label(null_density: f32) -> &'static str { + match (null_density * 100.0) as u32 { + 0 => "no NULLs", + 50 => "half NULLs", + 90 => "90pct NULLs", + 99 => "99pct NULLs", + _ => unreachable!(), + } +} + +fn bench_list_memory( + group: &mut BenchmarkGroup, + null_density: f32, + pages: impl PageIterator + Clone + 'static, + make_reader: fn(Box, ColumnDescPtr) -> Box, + column_desc: &ColumnDescPtr, +) { + group.bench_function(null_density_label(null_density), |b| { + b.iter_batched( + || make_reader(Box::new(pages.clone()), column_desc.clone()), + |reader| { + let count = drain_reader(reader); + assert_eq!(count, EXPECTED_VALUE_COUNT); + count + }, + criterion::BatchSize::PerIteration, + ); + }); +} + +fn int32_list_wrapper( + pages: Box, + column_desc: ColumnDescPtr, +) -> Box { + let child: Box = Box::new( + PrimitiveArrayReader::::new( + pages, + column_desc, + None, + DEFAULT_BATCH_SIZE, + Some(2), + ) + .unwrap(), + ); + let field = Field::new_list_field(DataType::Int32, true); + let data_type = DataType::List(Arc::new(field)); + Box::new(ListArrayReader::::new( + child, data_type, 2, 1, true, None, + )) +} + +fn double_list_wrapper( + pages: Box, + column_desc: ColumnDescPtr, +) -> Box { + let child: Box = Box::new( + PrimitiveArrayReader::::new( + pages, + column_desc, + None, + DEFAULT_BATCH_SIZE, + Some(2), + ) + .unwrap(), + ); + let field = Field::new_list_field(DataType::Float64, true); + let data_type = DataType::List(Arc::new(field)); + Box::new(ListArrayReader::::new( + child, data_type, 2, 1, true, None, + )) +} + +fn fixed32_list_wrapper( + pages: Box, + column_desc: ColumnDescPtr, +) -> Box { + let child = + make_fixed_len_byte_array_reader(pages, column_desc, None, DEFAULT_BATCH_SIZE, Some(2)) + .unwrap(); + let field = Field::new_list_field(DataType::FixedSizeBinary(FIXED_BYTE_LEN as i32), true); + let data_type = DataType::List(Arc::new(field)); + Box::new(ListArrayReader::::new( + child, data_type, 2, 1, true, None, + )) +} + +fn string_list_wrapper( + pages: Box, + column_desc: ColumnDescPtr, +) -> Box { + let child = + make_byte_array_reader(pages, column_desc, None, DEFAULT_BATCH_SIZE, Some(2)).unwrap(); + let field = Field::new_list_field(DataType::Utf8, true); + let data_type = DataType::List(Arc::new(field)); + Box::new(ListArrayReader::::new( + child, data_type, 2, 1, true, None, + )) +} + +fn add_benches(c: &mut Criterion, measurement_name: &str) { + let (_schema, int32_desc) = build_int32_list_schema(); + let (_schema, double_desc) = build_double_list_schema(); + let (_schema, fixed32_desc) = build_fixed32_list_schema(); + let (_schema, string_desc) = build_string_list_schema(); + + let mut group = c.benchmark_group(format!( + "arrow_array_reader/ListArray_{measurement_name}/Int32List" + )); + for null_density in [0.0, 0.5, 0.9, 0.99] { + let pages = + build_list_pages::(int32_desc.clone(), null_density, |rng| rng.random()); + bench_list_memory( + &mut group, + null_density, + pages, + int32_list_wrapper, + &int32_desc, + ); + } + group.finish(); + + let mut group = c.benchmark_group(format!( + "arrow_array_reader/ListArray_{measurement_name}/DoubleList" + )); + for null_density in [0.0, 0.5, 0.9, 0.99] { + let pages = + build_list_pages::(double_desc.clone(), null_density, |rng| rng.random()); + bench_list_memory( + &mut group, + null_density, + pages, + double_list_wrapper, + &double_desc, + ); + } + group.finish(); + + let mut group = c.benchmark_group(format!( + "arrow_array_reader/ListArray_{measurement_name}/Fixed32List" + )); + for null_density in [0.0, 0.5, 0.9, 0.99] { + let pages = + build_list_pages::(fixed32_desc.clone(), null_density, |rng| { + let mut buf = vec![0u8; FIXED_BYTE_LEN]; + rng.fill(&mut buf[..]); + buf.into() + }); + bench_list_memory( + &mut group, + null_density, + pages, + fixed32_list_wrapper, + &fixed32_desc, + ); + } + group.finish(); + + let mut group = c.benchmark_group(format!( + "arrow_array_reader/ListArray_{measurement_name}/StringList" + )); + for null_density in [0.0, 0.5, 0.9, 0.99] { + let pages = build_list_pages::(string_desc.clone(), null_density, |rng| { + let len = rng.random_range(5..50); + let bytes: Vec = (0..len).map(|_| rng.random_range(b'a'..=b'z')).collect(); + bytes.into() + }); + bench_list_memory( + &mut group, + null_density, + pages, + string_list_wrapper, + &string_desc, + ); + } + 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/vendor/parquet/benches/arrow_reader_row_filter.rs b/vendor/parquet/benches/arrow_reader_row_filter.rs new file mode 100644 index 0000000..2b5a09e --- /dev/null +++ b/vendor/parquet/benches/arrow_reader_row_filter.rs @@ -0,0 +1,701 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmark for evaluating row filters and projections on a Parquet file. +//! +//! # Background: +//! +//! As described in [Efficient Filter Pushdown in Parquet], evaluating +//! pushdown filters is a two-step process: +//! +//! 1. Build a filter mask by decoding and evaluating filter functions on +//! the filter column(s). +//! +//! 2. Decode the rows that match the filter mask from the projected columns. +//! +//! The performance depends on factors such as the number of rows selected, +//! the clustering of results (which affects the efficiency of the filter mask), +//! and whether the same column is used for both filtering and projection. +//! +//! This benchmark helps measure the performance of these operations. +//! +//! [Efficient Filter Pushdown in Parquet]: https://datafusion.apache.org/blog/2025/03/21/parquet-pushdown/ +//! +//! The benchmark creates an in-memory Parquet file with 100K rows and ten columns. +//! The first four columns are: +//! - int64: random integers (range: 0..100) generated with a fixed seed. +//! - float64: random floating-point values (range: 0.0..100.0) generated with a fixed seed. +//! - utf8View: random strings with some empty values and occasional constant "const" values. +//! - ts: sequential timestamps in milliseconds. +//! +//! The following six columns (for filtering) are generated to mimic different +//! filter selectivity and clustering patterns: +//! - pt: for Point Lookup – exactly one row is set to "unique_point", all others are random strings. +//! - sel: for Selective Unclustered – exactly 1% of rows (those with i % 100 == 0) are "selected". +//! - mod_clustered: for Moderately Selective Clustered – in each 10K-row block, the first 10 rows are "mod_clustered". +//! - mod_unclustered: for Moderately Selective Unclustered – exactly 10% of rows (those with i % 10 == 1) are "mod_unclustered". +//! - unsel_unclustered: for Unselective Unclustered – exactly 99% of rows (those with i % 100 != 0) are "unsel_unclustered". +//! - unsel_clustered: for Unselective Clustered – in each 10K-row block, rows with an offset >= 1000 are "unsel_clustered". +//! + +use arrow::array::{ArrayRef, BooleanArray, Float64Array, Int64Array, TimestampMillisecondArray}; +use arrow::compute::and; +use arrow::compute::kernels::cmp::{eq, gt, lt, neq}; +use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; +use arrow::record_batch::RecordBatch; +use arrow_array::StringViewArray; +use arrow_array::builder::{ArrayBuilder, StringViewBuilder}; +use arrow_cast::pretty::pretty_format_batches; +use bytes::Bytes; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use futures::future::BoxFuture; +use futures::{FutureExt, StreamExt}; +use parquet::arrow::arrow_reader::{ + ArrowPredicateFn, ArrowReaderOptions, ParquetRecordBatchReaderBuilder, RowFilter, +}; +use parquet::arrow::async_reader::AsyncFileReader; +use parquet::arrow::{ArrowWriter, ParquetRecordBatchStreamBuilder, ProjectionMask}; +use parquet::basic::Compression; +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader}; +use parquet::file::properties::WriterProperties; +use rand::{Rng, SeedableRng, rngs::StdRng}; +use std::ops::Range; +use std::sync::Arc; + +/// Generates a random string. Has a 50% chance to generate a short string (3–11 characters) +/// or a long string (13–20 characters). +fn random_string(rng: &mut StdRng) -> String { + let charset = b"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"; + let is_long = rng.random_bool(0.5); + let len = if is_long { + rng.random_range(13..21) + } else { + rng.random_range(3..12) + }; + (0..len) + .map(|_| charset[rng.random_range(0..charset.len())] as char) + .collect() +} + +/// Creates an int64 array of a given size with random integers in [0, 100). +/// Then, it overwrites a single random index with 9999 to serve as the unique value for point lookup. +fn create_int64_array(size: usize) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(42); + let mut values: Vec = (0..size).map(|_| rng.random_range(0..100)).collect(); + let unique_index = rng.random_range(0..size); + values[unique_index] = 9999; // Unique value for point lookup + Arc::new(Int64Array::from(values)) as ArrayRef +} + +/// Creates a float64 array of a given size with random floats in [0.0, 100.0). +fn create_float64_array(size: usize) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(43); + let values: Vec = (0..size).map(|_| rng.random_range(0.0..100.0)).collect(); + Arc::new(Float64Array::from(values)) as ArrayRef +} + +/// Creates a utf8View array of a given size with random strings. +/// +/// This is modeled after the "SearchPhrase" column in the ClickBench benchmark. +/// +/// See for calculations. +/// +/// The important ClickBench data properties are: +/// * Selectivity is: 13172392 / 99997497 = 0.132 +/// * Number of RowSelections = 14054784 +/// * Average run length of each RowSelection: 99997497 / 14054784 = 7.114 +/// +/// The properties of this array are: +/// * Selectivity is: 15144 / 100000 = 0.15144 +/// * Number of RowSelections = 12904 +/// * Average run length of each RowSelection: 100000 / 12904 = 7.75 +fn create_utf8_view_array(size: usize) -> ArrayRef { + const AVG_RUN_LENGTH: usize = 4; // average number of empty/non-empty strings in a row + const EMPTY_DENSITY: u32 = 85; // percent chance that each run is an empty string + + let mut builder = StringViewBuilder::with_capacity(size); + let mut rng = StdRng::seed_from_u64(44); + while builder.len() < size { + let mut run_length = rng.random_range(1..AVG_RUN_LENGTH); + if builder.len() + run_length > size { + // cap to size rows + run_length = size - builder.len(); + } + + let choice = rng.random_range(0..100); + if choice < EMPTY_DENSITY { + for _ in 0..run_length { + builder.append_value(""); + } + } else { + for _ in 0..run_length { + builder.append_value(random_string(&mut rng)); + } + } + } + Arc::new(builder.finish()) as ArrayRef +} + +/// Creates a ts (timestamp) array of a given size. Each value is computed as i % 10_000, +/// which simulates repeating blocks (each block of 10,000) to model clustered patterns. +fn create_ts_array(size: usize) -> ArrayRef { + let values: Vec = (0..size).map(|i| (i % 10_000) as i64).collect(); + Arc::new(TimestampMillisecondArray::from(values)) as ArrayRef +} + +/// Creates a RecordBatch with 100K rows and 4 columns: int64, float64, utf8View, and ts. +fn create_record_batch(size: usize) -> RecordBatch { + let fields = vec![ + Field::new("int64", DataType::Int64, false), + Field::new("float64", DataType::Float64, false), + Field::new("utf8View", DataType::Utf8View, true), + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ]; + let schema = Arc::new(Schema::new(fields)); + + let int64_array = create_int64_array(size); + let float64_array = create_float64_array(size); + let utf8_array = create_utf8_view_array(size); + let ts_array = create_ts_array(size); + + let arrays: Vec = vec![int64_array, float64_array, utf8_array, ts_array]; + RecordBatch::try_new(schema, arrays).unwrap() +} + +/// Total number of rows. +const TOTAL_ROWS: usize = 500_000; + +/// Maximum rows per row group. +const ROW_GROUP_SIZE: usize = 100_000; + +/// Writes the RecordBatch to an in memory buffer, returning the buffer +fn write_parquet_file() -> Vec { + let batch = create_record_batch(TOTAL_ROWS); + println!("Batch created with {TOTAL_ROWS} rows, row group size = {ROW_GROUP_SIZE}"); + println!( + "First 100 rows:\n{}", + pretty_format_batches(&[batch.clone().slice(0, 100)]).unwrap() + ); + let schema = batch.schema(); + let props = WriterProperties::builder() + .set_compression(Compression::SNAPPY) + .set_max_row_group_row_count(Some(ROW_GROUP_SIZE)) + .build(); + let mut buffer = vec![]; + { + let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } + buffer +} + +/// ProjectionCase defines the projection mode for the benchmark: +/// either projecting all columns or excluding the column that is used for filtering. +#[derive(Clone)] +enum ProjectionCase { + AllColumns, + ExcludeFilterColumn, +} + +impl std::fmt::Display for ProjectionCase { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ProjectionCase::AllColumns => write!(f, "all_columns"), + ProjectionCase::ExcludeFilterColumn => write!(f, "exclude_filter_column"), + } + } +} + +/// FilterType encapsulates the different filter comparisons. +/// The variants correspond to the different filter patterns. +#[derive(Clone, Copy, Debug)] +enum FilterType { + /// "Point Lookup": selects a single row + /// ```text + /// ┌───────────────┐ ┌───────────────┐ + /// │ │ │ │ + /// │ │ │ ... │ + /// │ │ │ │ + /// │ │ │ │ + /// │ ... │ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// │ │ │ │ + /// │ │ │ ... │ + /// │ │ │ │ + /// │ │ │ │ + /// └───────────────┘ └───────────────┘ + /// ``` + /// (1 RowSelection of 1 row) + PointLookup, + /// selective (1%) unclustered filter + /// ```text + /// ┌───────────────┐ ┌───────────────┐ + /// │ ... │ │ │ + /// │ │ │ │ + /// │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ │ │ + /// │ │ │ ... │ + /// │ │ │ │ + /// │ │ │ │ + /// │ ... │ │ │ + /// │ │ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// │ │ │ │ + /// └───────────────┘ └───────────────┘ + /// ``` + /// (1000 RowSelection of 10 rows each) + SelectiveUnclustered, + /// moderately selective (10%) clustered filter + /// ```text + /// ┌───────────────┐ ┌───────────────┐ + /// │ │ │ │ + /// │ │ │ │ + /// │ │ │ ... │ + /// │ │ │ │ + /// │ ... │ │ │ + /// │ │ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// │ │ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// │ │ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// │ │ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// └───────────────┘ └───────────────┘ + /// ``` + /// (10 RowSelections of 10,000 rows each) + ModeratelySelectiveClustered, + /// moderately selective (10%) clustered filter + /// ```text + /// ┌───────────────┐ ┌───────────────┐ + /// │ ... │ │ │ + /// │ │ │ │ + /// │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ │ │ + /// │ │ │ │ + /// │ │ │ ... │ + /// │ ... │ │ │ + /// │ │ │ │ + /// │ │ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// └───────────────┘ └───────────────┘ + /// ``` + /// (10 RowSelections of 10,000 rows each) + ModeratelySelectiveUnclustered, + /// unselective (99%) unclustered filter + /// ```text + /// ┌───────────────┐ ┌───────────────┐ + /// │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ │ │ + /// │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// └───────────────┘ └───────────────┘ + /// ``` + /// (99,000 RowSelections of 10 rows each) + UnselectiveUnclustered, + /// unselective (90%) clustered filter + /// ```text + /// ┌───────────────┐ ┌───────────────┐ + /// │ │ │ │ + /// │ │ │ │ + /// │ │ │ ... │ + /// │ │ │ │ + /// │ ... │ │ │ + /// │ │ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// │ │ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// │ │ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// │ │ │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + /// └───────────────┘ └───────────────┘ + /// ``` + /// (99 RowSelection of 10,000 rows each) + UnselectiveClustered, + /// [`Self::SelectivelUnclusered`] `AND` + /// [`Self::ModeratelySelectiveClustered`] + Composite, + /// `utf8View <> ''` modeling [ClickBench] [Q21-Q27] + /// + /// [ClickBench]: https://github.com/ClickHouse/ClickBench + /// [Q21-Q27]: https://github.com/apache/datafusion/blob/b7177234e65cbbb2dcc04c252f6acd80bb026362/benchmarks/queries/clickbench/queries.sql#L22-L28 + Utf8ViewNonEmpty, +} + +impl std::fmt::Display for FilterType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let s = match self { + FilterType::PointLookup => "int64 == 9999", + FilterType::SelectiveUnclustered => "float64 > 99.0", + FilterType::ModeratelySelectiveClustered => "ts >= 9000", + FilterType::ModeratelySelectiveUnclustered => "int64 > 90", + FilterType::UnselectiveUnclustered => "float64 <= 99.0", + FilterType::UnselectiveClustered => "ts < 9000", + FilterType::Composite => "float64 > 99.0 AND ts >= 9000", + FilterType::Utf8ViewNonEmpty => "utf8View <> ''", + }; + write!(f, "{s}") + } +} + +impl FilterType { + /// Applies the specified filter on the given RecordBatch and returns a BooleanArray mask. + fn filter_batch(&self, batch: &RecordBatch) -> arrow::error::Result { + match self { + // Point Lookup on int64 column + FilterType::PointLookup => { + let array = batch.column(batch.schema().index_of("int64")?); + let scalar = Int64Array::new_scalar(9999); + eq(array, &scalar) + } + // Selective Unclustered on float64 column: float64 > 99.0 + FilterType::SelectiveUnclustered => { + let array = batch.column(batch.schema().index_of("float64")?); + let scalar = Float64Array::new_scalar(99.0); + gt(array, &scalar) + } + // Moderately Selective Clustered on ts column: ts >= 9000 (implemented as > 8999) + FilterType::ModeratelySelectiveClustered => { + let array = batch.column(batch.schema().index_of("ts")?); + gt(array, &TimestampMillisecondArray::new_scalar(8999)) + } + // Moderately Selective Unclustered on int64 column: int64 > 90 + FilterType::ModeratelySelectiveUnclustered => { + let array = batch.column(batch.schema().index_of("int64")?); + let scalar = Int64Array::new_scalar(90); + gt(array, &scalar) + } + // Unselective Unclustered on float64 column: NOT (float64 > 99.0) + FilterType::UnselectiveUnclustered => { + let array = batch.column(batch.schema().index_of("float64")?); + gt(array, &Float64Array::new_scalar(99.0)) + } + // Unselective Clustered on ts column: ts < 9000 + FilterType::UnselectiveClustered => { + let array = batch.column(batch.schema().index_of("ts")?); + lt(array, &TimestampMillisecondArray::new_scalar(9000)) + } + // Composite filter: logical AND of (float64 > 99.0) and (ts >= 9000) + FilterType::Composite => { + let mask1 = FilterType::SelectiveUnclustered.filter_batch(batch)?; + let mask2 = FilterType::ModeratelySelectiveClustered.filter_batch(batch)?; + and(&mask1, &mask2) + } + // Utf8ViewNonEmpty: selects rows where the utf8View column is not an empty string. + FilterType::Utf8ViewNonEmpty => { + let array = batch.column(batch.schema().index_of("utf8View")?); + let scalar = StringViewArray::new_scalar(""); + neq(array, &scalar) + } + } + } + + /// Return the indexes in the batch's schema that are used for filtering. + fn filter_projection(&self) -> &'static [usize] { + match self { + FilterType::PointLookup => &[0], + FilterType::SelectiveUnclustered => &[1], + FilterType::ModeratelySelectiveClustered => &[3], + FilterType::ModeratelySelectiveUnclustered => &[0], + FilterType::UnselectiveUnclustered => &[1], + FilterType::UnselectiveClustered => &[3], + FilterType::Composite => &[1, 3], // Use float64 column and ts column as representative for composite + FilterType::Utf8ViewNonEmpty => &[2], + } + } +} + +/// Benchmark filters and projections by reading the Parquet file. +/// This benchmark iterates over all individual filter types and two projection cases. +/// It measures the time to read and filter the Parquet file according to each scenario. +fn benchmark_filters_and_projections(c: &mut Criterion) { + // make the parquet file in memory that can be shared + let parquet_file = Bytes::from(write_parquet_file()); + let filter_types = vec![ + FilterType::PointLookup, + FilterType::SelectiveUnclustered, + FilterType::ModeratelySelectiveClustered, + FilterType::ModeratelySelectiveUnclustered, + FilterType::UnselectiveUnclustered, + FilterType::UnselectiveClustered, + FilterType::Utf8ViewNonEmpty, + FilterType::Composite, + ]; + let projection_cases = vec![ + ProjectionCase::AllColumns, + ProjectionCase::ExcludeFilterColumn, + ]; + + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap(); + + let mut group = c.benchmark_group("arrow_reader_row_filter"); + + for filter_type in filter_types { + for proj_case in &projection_cases { + // All indices corresponding to the 10 columns. + let all_indices = vec![0, 1, 2, 3]; + let filter_col = filter_type.filter_projection().to_vec(); + // For the projection, either select all columns or exclude the filter column(s). + let output_projection: Vec = match proj_case { + ProjectionCase::AllColumns => all_indices.clone(), + ProjectionCase::ExcludeFilterColumn => all_indices + .into_iter() + .filter(|i| !filter_col.contains(i)) + .collect(), + }; + + let reader = InMemoryReader::try_new(&parquet_file).unwrap(); + let metadata = Arc::clone(reader.metadata()); + + let schema_descr = metadata.file_metadata().schema_descr(); + let projection_mask = ProjectionMask::roots(schema_descr, output_projection.clone()); + let pred_mask = ProjectionMask::roots(schema_descr, filter_col.clone()); + + let benchmark_name = format!("{filter_type}/{proj_case}",); + + // run the benchmark for the async reader + let bench_id = BenchmarkId::new(benchmark_name.clone(), "async"); + let rt_captured = rt.handle().clone(); + group.bench_function(bench_id, |b| { + b.iter(|| { + let reader = reader.clone(); + let pred_mask = pred_mask.clone(); + let projection_mask = projection_mask.clone(); + // row filters are not clone, so must make it each iter + let filter = ArrowPredicateFn::new(pred_mask, move |batch: RecordBatch| { + Ok(filter_type.filter_batch(&batch).unwrap()) + }); + let row_filter = RowFilter::new(vec![Box::new(filter)]); + + rt_captured.block_on(async { + benchmark_async_reader(reader, projection_mask, row_filter).await; + }) + }); + }); + + // run the benchmark for the sync reader + let bench_id = BenchmarkId::new(benchmark_name, "sync"); + group.bench_function(bench_id, |b| { + b.iter(|| { + let reader = reader.clone(); + let pred_mask = pred_mask.clone(); + let projection_mask = projection_mask.clone(); + // row filters are not clone, so must make it each iter + let filter = ArrowPredicateFn::new(pred_mask, move |batch: RecordBatch| { + Ok(filter_type.filter_batch(&batch).unwrap()) + }); + let row_filter = RowFilter::new(vec![Box::new(filter)]); + + benchmark_sync_reader(reader, projection_mask, row_filter) + }); + }); + } + } +} + +/// Use async API +async fn benchmark_async_reader( + reader: InMemoryReader, + projection_mask: ProjectionMask, + row_filter: RowFilter, +) { + let mut stream = ParquetRecordBatchStreamBuilder::new(reader) + .await + .unwrap() + .with_batch_size(8192) + .with_projection(projection_mask) + .with_row_filter(row_filter) + .build() + .unwrap(); + while let Some(b) = stream.next().await { + b.unwrap(); // consume the batches, no buffering + } +} + +/// Like [`benchmark_async_reader`] but also threads `with_limit(limit)` into +/// the stream builder. Used by the `LIMIT` benchmark below. +async fn benchmark_async_reader_with_limit( + reader: InMemoryReader, + projection_mask: ProjectionMask, + row_filter: RowFilter, + limit: usize, +) { + let mut stream = ParquetRecordBatchStreamBuilder::new(reader) + .await + .unwrap() + .with_batch_size(8192) + .with_projection(projection_mask) + .with_row_filter(row_filter) + .with_limit(limit) + .build() + .unwrap(); + while let Some(b) = stream.next().await { + b.unwrap(); // consume the batches, no buffering + } +} + +/// Use sync API +fn benchmark_sync_reader( + reader: InMemoryReader, + projection_mask: ProjectionMask, + row_filter: RowFilter, +) { + let stream = ParquetRecordBatchReaderBuilder::try_new(reader.into_inner()) + .unwrap() + .with_batch_size(8192) + .with_projection(projection_mask) + .with_row_filter(row_filter) + .build() + .unwrap(); + for b in stream { + b.unwrap(); // consume the batches, no buffering + } +} + +/// Adapter to read asynchronously from in memory bytes and always loads the +/// metadata with page indexes. +#[derive(Debug, Clone)] +struct InMemoryReader { + inner: Bytes, + metadata: Arc, +} + +impl InMemoryReader { + fn try_new(inner: &Bytes) -> parquet::errors::Result { + let mut metadata_reader = + ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required); + metadata_reader.try_parse(inner)?; + let metadata = metadata_reader.finish().map(Arc::new)?; + + Ok(Self { + // clone of bytes is cheap -- increments a refcount + inner: inner.clone(), + metadata, + }) + } + + fn metadata(&self) -> &Arc { + &self.metadata + } + + fn into_inner(self) -> Bytes { + self.inner + } +} + +impl AsyncFileReader for InMemoryReader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, parquet::errors::Result> { + let data = self.inner.slice(range.start as usize..range.end as usize); + async move { Ok(data) }.boxed() + } + + fn get_metadata<'a>( + &'a mut self, + _options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, parquet::errors::Result>> { + let metadata = Arc::clone(&self.metadata); + async move { Ok(metadata) }.boxed() + } +} + +/// Benchmark filters with `LIMIT` short-circuit (`with_limit(N)`) +/// +/// `PointLookup` is excluded because the filter has only 1 match in the +/// whole file; `LIMIT 10` is not binding. +fn benchmark_filters_with_limit(c: &mut Criterion) { + const LIMIT: usize = 10; + + let parquet_file = Bytes::from(write_parquet_file()); + let filter_types = vec![ + FilterType::SelectiveUnclustered, + FilterType::ModeratelySelectiveClustered, + FilterType::ModeratelySelectiveUnclustered, + FilterType::UnselectiveUnclustered, + FilterType::UnselectiveClustered, + FilterType::Utf8ViewNonEmpty, + FilterType::Composite, + ]; + let projection_cases = vec![ + ProjectionCase::AllColumns, + ProjectionCase::ExcludeFilterColumn, + ]; + let all_indices = vec![0, 1, 2, 3]; + + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap(); + + let mut group = c.benchmark_group("arrow_reader_row_filter_limit"); + + for filter_type in filter_types { + for proj_case in &projection_cases { + let filter_col = filter_type.filter_projection().to_vec(); + let output_projection: Vec = match proj_case { + ProjectionCase::AllColumns => all_indices.clone(), + ProjectionCase::ExcludeFilterColumn => all_indices + .iter() + .copied() + .filter(|i| !filter_col.contains(i)) + .collect(), + }; + + let reader = InMemoryReader::try_new(&parquet_file).unwrap(); + let metadata = Arc::clone(reader.metadata()); + let schema_descr = metadata.file_metadata().schema_descr(); + let projection_mask = ProjectionMask::roots(schema_descr, output_projection); + let pred_mask = ProjectionMask::roots(schema_descr, filter_col); + + let benchmark_name = format!("{filter_type}/{proj_case}/limit{LIMIT}"); + + // async variant + let bench_id = BenchmarkId::new(benchmark_name.clone(), "async"); + let rt_handle = rt.handle().clone(); + let pred_mask_async = pred_mask.clone(); + let projection_mask_async = projection_mask.clone(); + let reader_async = reader.clone(); + group.bench_function(bench_id, |b| { + b.iter(|| { + let reader = reader_async.clone(); + let pred_mask = pred_mask_async.clone(); + let projection_mask = projection_mask_async.clone(); + // RowFilter and ArrowPredicateFn are not Clone — fresh each iter. + let predicate = ArrowPredicateFn::new(pred_mask, move |batch: RecordBatch| { + Ok(filter_type.filter_batch(&batch).unwrap()) + }); + let row_filter = RowFilter::new(vec![Box::new(predicate)]); + rt_handle.block_on(benchmark_async_reader_with_limit( + reader, + projection_mask, + row_filter, + LIMIT, + )); + }); + }); + } + } +} + +criterion_group!( + benches, + benchmark_filters_and_projections, + benchmark_filters_with_limit, +); +criterion_main!(benches); diff --git a/vendor/parquet/benches/arrow_reader_row_selection_policy.rs b/vendor/parquet/benches/arrow_reader_row_selection_policy.rs new file mode 100644 index 0000000..da3f4fc --- /dev/null +++ b/vendor/parquet/benches/arrow_reader_row_selection_policy.rs @@ -0,0 +1,57 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Auto-policy benchmark for async Parquet `RowSelection` execution. +//! +//! This benchmark is intended for direct comparisons between `main` and +//! candidate branches. Forced row-selection policies and decode-then-filter +//! diagnostics intentionally belong in separate oracle benchmarks. +//! Page indexes are intentionally disabled so this benchmark isolates +//! row-selection execution from page-level I/O pruning. + +mod row_selection_policy_common; + +use criterion::{Criterion, criterion_group, criterion_main}; +use row_selection_policy_common::cases::{ + ORDER_CASES, SCALE_CASES, SHAPE_CASES, assert_order_cases_are_reverses, +}; +use row_selection_policy_common::register::register_auto_group; +use row_selection_policy_common::shapes::assert_shape_contracts; + +fn benchmark_auto(c: &mut Criterion) { + assert_shape_contracts(); + assert_order_cases_are_reverses(); + + register_auto_group( + c, + "arrow_reader_row_selection_policy/auto/shape", + SHAPE_CASES, + ); + register_auto_group( + c, + "arrow_reader_row_selection_policy/auto/order", + ORDER_CASES, + ); + register_auto_group( + c, + "arrow_reader_row_selection_policy/auto/scale", + SCALE_CASES, + ); +} + +criterion_group!(benches, benchmark_auto); +criterion_main!(benches); diff --git a/vendor/parquet/benches/arrow_statistics.rs b/vendor/parquet/benches/arrow_statistics.rs new file mode 100644 index 0000000..6da816b --- /dev/null +++ b/vendor/parquet/benches/arrow_statistics.rs @@ -0,0 +1,273 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! 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_schema::{ + DataType::{self, *}, + Field, Schema, +}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use parquet::{ + arrow::arrow_reader::ArrowReaderOptions, + file::{metadata::PageIndexPolicy, properties::WriterProperties}, +}; +use parquet::{ + arrow::{ArrowWriter, arrow_reader::ArrowReaderBuilder}, + file::properties::EnabledStatistics, +}; +use std::sync::Arc; +use tempfile::NamedTempFile; +#[derive(Debug, Clone)] +enum TestTypes { + UInt64, + Int64, + F64, + String, + Dictionary, +} + +use parquet::arrow::arrow_reader::statistics::StatisticsConverter; +use std::fmt; + +impl fmt::Display for TestTypes { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + TestTypes::UInt64 => write!(f, "UInt64"), + TestTypes::Int64 => write!(f, "Int64"), + TestTypes::F64 => write!(f, "F64"), + TestTypes::String => write!(f, "String"), + TestTypes::Dictionary => write!(f, "Dictionary(Int32, String)"), + } + } +} + +fn create_parquet_file( + dtype: TestTypes, + row_groups: usize, + data_page_row_count_limit: &Option, +) -> NamedTempFile { + let schema = match dtype { + TestTypes::UInt64 => Arc::new(Schema::new(vec![Field::new("col", DataType::UInt64, true)])), + TestTypes::Int64 => Arc::new(Schema::new(vec![Field::new("col", DataType::Int64, true)])), + TestTypes::F64 => Arc::new(Schema::new(vec![Field::new( + "col", + DataType::Float64, + true, + )])), + TestTypes::String => Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, true)])), + TestTypes::Dictionary => Arc::new(Schema::new(vec![Field::new( + "col", + DataType::Dictionary(Box::new(Int32), Box::new(Utf8)), + true, + )])), + }; + + let mut props = WriterProperties::builder().set_max_row_group_row_count(Some(row_groups)); + if let Some(limit) = data_page_row_count_limit { + props = props + .set_data_page_row_count_limit(*limit) + .set_statistics_enabled(EnabledStatistics::Page); + }; + let props = props.build(); + + let file = tempfile::Builder::new() + .suffix(".parquet") + .tempfile() + .unwrap(); + let mut writer = + ArrowWriter::try_new(file.reopen().unwrap(), schema.clone(), Some(props)).unwrap(); + + for _ in 0..row_groups { + let batch = match dtype { + TestTypes::UInt64 => make_uint64_batch(), + TestTypes::Int64 => make_int64_batch(), + TestTypes::F64 => make_f64_batch(), + TestTypes::String => make_string_batch(), + TestTypes::Dictionary => make_dict_batch(), + }; + if data_page_row_count_limit.is_some() { + // Send batches one at a time. This allows the + // writer to apply the page limit, that is only + // checked on RecordBatch boundaries. + for i in 0..batch.num_rows() { + writer.write(&batch.slice(i, 1)).unwrap(); + } + } else { + writer.write(&batch).unwrap(); + } + } + writer.close().unwrap(); + file +} + +fn make_uint64_batch() -> RecordBatch { + let array: ArrayRef = Arc::new(UInt64Array::from(vec![ + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + ])); + RecordBatch::try_new( + Arc::new(arrow::datatypes::Schema::new(vec![ + arrow::datatypes::Field::new("col", UInt64, false), + ])), + vec![array], + ) + .unwrap() +} + +fn make_int64_batch() -> RecordBatch { + let array: ArrayRef = Arc::new(Int64Array::from(vec![ + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + ])); + RecordBatch::try_new( + Arc::new(arrow::datatypes::Schema::new(vec![ + arrow::datatypes::Field::new("col", Int64, false), + ])), + vec![array], + ) + .unwrap() +} + +fn make_f64_batch() -> RecordBatch { + let array: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0, 5.0])); + RecordBatch::try_new( + Arc::new(arrow::datatypes::Schema::new(vec![ + arrow::datatypes::Field::new("col", Float64, false), + ])), + vec![array], + ) + .unwrap() +} + +fn make_string_batch() -> RecordBatch { + let array: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"])); + RecordBatch::try_new( + Arc::new(arrow::datatypes::Schema::new(vec![ + arrow::datatypes::Field::new("col", Utf8, false), + ])), + vec![array], + ) + .unwrap() +} + +fn make_dict_batch() -> RecordBatch { + let keys = Int32Array::from(vec![0, 1, 2, 3, 4]); + let values = StringArray::from(vec!["a", "b", "c", "d", "e"]); + let array: ArrayRef = Arc::new(DictionaryArray::try_new(keys, Arc::new(values)).unwrap()); + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "col", + Dictionary(Box::new(Int32), Box::new(Utf8)), + false, + )])), + vec![array], + ) + .unwrap() +} + +fn criterion_benchmark(c: &mut Criterion) { + let row_groups = 100; + use TestTypes::*; + let types = vec![Int64, UInt64, F64, String, Dictionary]; + let data_page_row_count_limits = vec![None, Some(1)]; + + for dtype in types { + for data_page_row_count_limit in &data_page_row_count_limits { + let file = create_parquet_file(dtype.clone(), row_groups, data_page_row_count_limit); + let file = file.reopen().unwrap(); + let options = + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::from(true)); + let reader = ArrowReaderBuilder::try_new_with_options(file, options).unwrap(); + let metadata = reader.metadata(); + let row_groups = metadata.row_groups(); + let row_group_indices: Vec<_> = (0..row_groups.len()).collect(); + + let statistic_type = if data_page_row_count_limit.is_some() { + "data page" + } else { + "row group" + }; + + let mut group = c.benchmark_group(format!( + "Extract {} statistics for {}", + statistic_type, + dtype.clone() + )); + group.bench_function(BenchmarkId::new("extract_statistics", dtype.clone()), |b| { + b.iter(|| { + let converter = StatisticsConverter::try_new( + "col", + reader.schema(), + reader.parquet_schema(), + ) + .unwrap(); + + if data_page_row_count_limit.is_some() { + let column_page_index = reader + .metadata() + .column_index() + .expect("File should have column page indices"); + + let column_offset_index = reader + .metadata() + .offset_index() + .expect("File should have column offset indices"); + + let _ = converter.data_page_mins( + column_page_index, + column_offset_index, + &row_group_indices, + ); + let _ = converter.data_page_maxes( + column_page_index, + column_offset_index, + &row_group_indices, + ); + let _ = converter.data_page_null_counts( + column_page_index, + column_offset_index, + &row_group_indices, + ); + let _ = converter.data_page_row_counts( + column_offset_index, + row_groups, + &row_group_indices, + ); + } else { + let _ = converter.row_group_mins(row_groups.iter()).unwrap(); + let _ = converter.row_group_maxes(row_groups.iter()).unwrap(); + let _ = converter.row_group_null_counts(row_groups.iter()).unwrap(); + let _ = converter.row_group_row_counts(row_groups.iter()).unwrap(); + } + }) + }); + group.finish(); + } + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/vendor/parquet/benches/arrow_writer.rs b/vendor/parquet/benches/arrow_writer.rs new file mode 100644 index 0000000..adc4dcc --- /dev/null +++ b/vendor/parquet/benches/arrow_writer.rs @@ -0,0 +1,655 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#[macro_use] +extern crate criterion; + +use criterion::{Bencher, Criterion, Throughput}; +use parquet::arrow::ArrowWriter; +use parquet::basic::{Compression, ZstdLevel}; + +extern crate arrow; +extern crate parquet; + +use std::hint::black_box; +use std::io::Empty; +use std::sync::Arc; + +use arrow::datatypes::*; +use arrow::util::bench_util::{create_f16_array, create_f32_array, create_f64_array}; +use arrow::{record_batch::RecordBatch, util::data_gen::*}; +use arrow_array::{RecordBatchOptions, StringArray}; +use parquet::errors::Result; +use parquet::file::properties::{CdcOptions, WriterProperties, WriterVersion}; + +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_string_bench_batch( + size: usize, + null_density: f32, + true_density: f32, +) -> Result { + let fields = vec![ + Field::new("_1", DataType::Utf8, true), + Field::new("_2", DataType::LargeUtf8, true), + ]; + let schema = Schema::new(fields); + Ok(create_random_batch( + Arc::new(schema), + size, + null_density, + true_density, + )?) +} + +/// 1 M short, fixed-width 8-byte strings. Exercises the BYTE_ARRAY hot path +/// for the case where individual values are small enough that the byte-budget +/// based sub-batch sizing in `write_batch_internal` should always resolve to +/// the full chunk (no granular splitting, no regression vs. current behavior). +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)])?) +} + +/// `size` rows of `value_size`-byte strings. Exercises the BYTE_ARRAY path +/// where individual values are large enough that batching the default +/// `write_batch_size` of them would blow the page byte limit by orders of +/// magnitude — the case the page-size fix targets. +fn create_large_string_bench_batch(size: usize, value_size: usize) -> Result { + let value = "x".repeat(value_size); + let array = Arc::new(StringArray::from_iter_values( + (0..size).map(|_| value.as_str()), + )) as _; + Ok(RecordBatch::try_from_iter([("col", array)])?) +} + +fn create_string_and_binary_view_bench_batch( + size: usize, + null_density: f32, + true_density: f32, +) -> Result { + let fields = vec![ + Field::new("_1", DataType::Utf8View, true), + Field::new("_2", DataType::BinaryView, true), + ]; + let schema = Schema::new(fields); + Ok(create_random_batch( + Arc::new(schema), + size, + null_density, + true_density, + )?) +} + +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_ree_bench_batch( + value_dt: DataType, + size: usize, + null_pct: Option, + true_density: f32, +) -> Result { + const DEFAULT_NULL_PCT: u8 = 10; + let null_density = null_pct.unwrap_or(DEFAULT_NULL_PCT) as f32 / 100.0; + let fields = vec![Field::new( + "_1", + DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(Field::new("values", value_dt, true)), + ), + true, + )]; + let schema = Schema::new(fields); + Ok(create_random_batch( + Arc::new(schema), + size, + null_density, + true_density, + )?) +} + +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, + )?) +} + +fn create_bool_bench_batch( + size: usize, + null_density: f32, + true_density: f32, +) -> Result { + let fields = vec![Field::new("_1", DataType::Boolean, true)]; + 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_float_bench_batch_with_nans(size: usize, nan_density: f32) -> Result { + let fields = vec![ + Field::new("_1", DataType::Float16, false), + Field::new("_2", DataType::Float32, false), + Field::new("_3", DataType::Float64, false), + ]; + let schema = Schema::new(fields); + let columns: Vec = vec![ + Arc::new(create_f16_array(size, nan_density)), + Arc::new(create_f32_array(size, nan_density)), + Arc::new(create_f64_array(size, nan_density)), + ]; + Ok(RecordBatch::try_new_with_options( + Arc::new(schema), + columns, + &RecordBatchOptions::new().with_match_field_names(false), + )?) +} + +fn create_decimal_bench_batch(size: usize, null_density: f32) -> Result { + let fields = vec![ + Field::new("_1", Decimal32Type::DEFAULT_TYPE, false), + Field::new("_2", Decimal64Type::DEFAULT_TYPE, false), + Field::new("_3", Decimal128Type::DEFAULT_TYPE, false), + Field::new("_4", Decimal256Type::DEFAULT_TYPE, false), + ]; + let schema = Schema::new(fields); + Ok(create_random_batch( + Arc::new(schema), + size, + null_density, + 0.75, + )?) +} + +fn create_list_primitive_bench_batch( + size: usize, + null_density: f32, + true_density: f32, +) -> Result { + let fields = vec![ + Field::new( + "_1", + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))), + true, + ), + Field::new( + "_2", + DataType::List(Arc::new(Field::new_list_field(DataType::Boolean, true))), + true, + ), + Field::new( + "_3", + DataType::LargeList(Arc::new(Field::new_list_field(DataType::Utf8, true))), + true, + ), + ]; + let schema = Schema::new(fields); + Ok(create_random_batch( + Arc::new(schema), + size, + null_density, + true_density, + )?) +} + +fn create_list_primitive_bench_batch_non_null( + size: usize, + null_density: f32, + true_density: f32, +) -> Result { + let fields = vec![ + Field::new( + "_1", + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))), + false, + ), + Field::new( + "_2", + DataType::List(Arc::new(Field::new_list_field(DataType::Boolean, false))), + false, + ), + Field::new( + "_3", + DataType::LargeList(Arc::new(Field::new_list_field(DataType::Utf8, false))), + false, + ), + ]; + let schema = Schema::new(fields); + Ok(create_random_batch( + Arc::new(schema), + size, + null_density, + true_density, + )?) +} + +fn create_struct_bench_batch(size: usize, null_density: f32) -> Result { + let fields = vec![Field::new( + "_1", + DataType::Struct(Fields::from(vec![ + Field::new("_1", DataType::Int32, false), + Field::new("_2", DataType::Int64, false), + Field::new("_3", DataType::Float32, false), + ])), + true, + )]; + let schema = Schema::new(fields); + Ok(create_random_batch( + Arc::new(schema), + size, + null_density, + 0.75, + )?) +} + +fn create_fsb_bench_batch(size: usize, null_density: f32, len: i32) -> Result { + let fields = vec![Field::new("_1", DataType::FixedSizeBinary(len), true)]; + let schema = Schema::new(fields); + Ok(create_random_batch( + Arc::new(schema), + size, + null_density, + 0.75, + )?) +} + +fn create_nested_list_bench_batch(size: usize, null_density: f32) -> Result { + // List> — exercises the nested repetition (non-batched) path + let fields = vec![Field::new( + "_1", + DataType::List(Arc::new(Field::new_list_field( + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))), + true, + ))), + true, + )]; + let schema = Schema::new(fields); + Ok(create_random_batch( + Arc::new(schema), + size, + null_density, + 0.75, + )?) +} + +fn create_list_struct_with_list_batch(size: usize, null_density: f32) -> Result { + // List>> + // The struct child contains a nested list, so child_has_no_nested_rep() = false. + // This exercises the per-slot (non-batched) write path in level computation. + let fields = vec![Field::new( + "_1", + DataType::List(Arc::new(Field::new_list_field( + DataType::Struct(Fields::from(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Float32, true), + Field::new( + "c", + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))), + true, + ), + ])), + true, + ))), + true, + )]; + let schema = Schema::new(fields); + Ok(create_random_batch( + Arc::new(schema), + size, + null_density, + 0.75, + )?) +} + +fn _create_nested_bench_batch( + size: usize, + null_density: f32, + true_density: f32, +) -> Result { + let fields = vec![ + Field::new( + "_1", + DataType::Struct(Fields::from(vec![ + Field::new("_1", DataType::Int8, true), + Field::new( + "_2", + DataType::Struct(Fields::from(vec![ + Field::new("_1", DataType::Int8, true), + Field::new( + "_1", + DataType::Struct(Fields::from(vec![ + Field::new("_1", DataType::Int8, true), + Field::new("_2", DataType::Utf8, true), + ])), + true, + ), + Field::new("_2", DataType::UInt8, true), + ])), + true, + ), + ])), + true, + ), + Field::new( + "_2", + DataType::LargeList(Arc::new(Field::new_list_field( + DataType::List(Arc::new(Field::new_list_field( + DataType::Struct(Fields::from(vec![ + Field::new( + "_1", + DataType::Struct(Fields::from(vec![ + Field::new("_1", DataType::Int8, true), + Field::new("_2", DataType::Int16, true), + Field::new("_3", DataType::Int32, true), + ])), + true, + ), + Field::new( + "_2", + DataType::List(Arc::new(Field::new( + "", + DataType::FixedSizeBinary(2), + true, + ))), + true, + ), + ])), + true, + ))), + true, + ))), + true, + ), + ]; + let schema = Schema::new(fields); + Ok(create_random_batch( + Arc::new(schema), + size, + null_density, + true_density, + )?) +} + +fn write_batch_with_option( + bench: &mut Bencher, + batch: &RecordBatch, + props: Option, +) -> Result<()> { + let props = props.unwrap_or_default(); + + bench.iter(|| { + let mut file = Empty::default(); + let mut writer = + ArrowWriter::try_new(&mut file, batch.schema(), Some(props.clone())).unwrap(); + writer.write(black_box(batch)).unwrap(); + black_box(writer.close()).unwrap(); + }); + + Ok(()) +} + +fn create_batches() -> Vec<(&'static str, RecordBatch)> { + const BATCH_SIZE: usize = 1024 * 1024; + + let mut batches = vec![]; + + let batch = create_primitive_bench_batch(BATCH_SIZE, 0.25, 0.75).unwrap(); + batches.push(("primitive", batch)); + + let batch = create_primitive_bench_batch_non_null(BATCH_SIZE, 0.25, 0.75).unwrap(); + batches.push(("primitive_non_null", batch)); + + let batch = create_bool_bench_batch(BATCH_SIZE, 0.25, 0.75).unwrap(); + batches.push(("bool", batch)); + + let batch = create_bool_bench_batch_non_null(BATCH_SIZE, 0.25, 0.75).unwrap(); + batches.push(("bool_non_null", batch)); + + let batch = create_string_bench_batch(BATCH_SIZE, 0.25, 0.75).unwrap(); + batches.push(("string", batch)); + + let batch = create_short_string_bench_batch(BATCH_SIZE).unwrap(); + batches.push(("short_string_non_null", batch)); + + // 1024 rows × 256 KiB = 256 MiB total. With the default 1 MiB page byte + // limit, this is the case where the page-size fix kicks in: each value + // needs its own page, and `write_batch_size = 1024` would otherwise + // buffer all 256 MiB before the post-write check runs. + let batch = create_large_string_bench_batch(1024, 256 * 1024).unwrap(); + batches.push(("large_string_non_null", batch)); + + let batch = create_string_and_binary_view_bench_batch(BATCH_SIZE, 0.25, 0.75).unwrap(); + batches.push(("string_and_binary_view", batch)); + + let batch = create_string_dictionary_bench_batch(BATCH_SIZE, 0.25, 0.75).unwrap(); + batches.push(("string_dictionary", batch)); + + let batch = create_string_bench_batch_non_null(BATCH_SIZE, 0.25, 0.75).unwrap(); + batches.push(("string_non_null", batch)); + + let batch = create_ree_bench_batch(DataType::Utf8, BATCH_SIZE, None, 0.75).unwrap(); + batches.push(("string_ree", batch)); + + let batch = create_ree_bench_batch(DataType::Int32, BATCH_SIZE, None, 0.75).unwrap(); + batches.push(("int32_ree", batch)); + + let batch = create_ree_bench_batch(DataType::Boolean, BATCH_SIZE, None, 0.75).unwrap(); + batches.push(("bool_ree", batch)); + + let batch = + create_ree_bench_batch(DataType::FixedSizeBinary(16), BATCH_SIZE, None, 0.75).unwrap(); + batches.push(("fixed_size_binary_ree", batch)); + + let batch = create_ree_bench_batch(DataType::Utf8, BATCH_SIZE, Some(95), 0.75).unwrap(); + batches.push(("string_ree_95pct_null", batch)); + + let batch = create_ree_bench_batch(DataType::Int32, BATCH_SIZE, Some(95), 0.75).unwrap(); + batches.push(("int32_ree_95pct_null", batch)); + + let batch = create_float_bench_batch_with_nans(BATCH_SIZE, 0.5).unwrap(); + batches.push(("float_with_nans", batch)); + + let batch = create_decimal_bench_batch(BATCH_SIZE, 0.75).unwrap(); + batches.push(("decimal", batch)); + + let batch = create_list_primitive_bench_batch(BATCH_SIZE, 0.25, 0.75).unwrap(); + batches.push(("list_primitive", batch)); + + let batch = create_list_primitive_bench_batch_non_null(BATCH_SIZE, 0.25, 0.75).unwrap(); + batches.push(("list_primitive_non_null", batch)); + + let batch = create_primitive_bench_batch(BATCH_SIZE, 0.99, 0.75).unwrap(); + batches.push(("primitive_sparse_99pct_null", batch)); + + let batch = create_list_primitive_bench_batch(BATCH_SIZE, 0.99, 0.75).unwrap(); + batches.push(("list_primitive_sparse_99pct_null", batch)); + + let batch = create_primitive_bench_batch(BATCH_SIZE, 1.0, 0.75).unwrap(); + batches.push(("primitive_all_null", batch)); + + let batch = create_struct_bench_batch(BATCH_SIZE, 0.0).unwrap(); + batches.push(("struct_non_null", batch)); + + let batch = create_struct_bench_batch(BATCH_SIZE, 0.99).unwrap(); + batches.push(("struct_sparse_99pct_null", batch)); + + let batch = create_struct_bench_batch(BATCH_SIZE, 1.0).unwrap(); + batches.push(("struct_all_null", batch)); + + let batch = create_fsb_bench_batch(BATCH_SIZE, 0.9, 16).unwrap(); + batches.push(("fsb", batch)); + + let batch = create_nested_list_bench_batch(BATCH_SIZE, 0.25).unwrap(); + batches.push(("list_nested", batch)); + + let batch = create_list_struct_with_list_batch(BATCH_SIZE, 0.25).unwrap(); + batches.push(("list_struct_with_list", batch)); + + batches +} + +fn create_writer_props() -> Vec<(&'static str, WriterProperties)> { + let mut props = vec![]; + + props.push(("default", Default::default())); + + let prop = WriterProperties::builder() + .set_bloom_filter_enabled(true) + .build(); + props.push(("bloom_filter", prop)); + + let prop = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_2_0) + .build(); + props.push(("parquet_2", prop)); + + let prop = WriterProperties::builder() + .set_compression(Compression::ZSTD(ZstdLevel::default())) + .build(); + props.push(("zstd", prop)); + + let prop = WriterProperties::builder() + .set_compression(Compression::ZSTD(ZstdLevel::default())) + .set_writer_version(WriterVersion::PARQUET_2_0) + .build(); + props.push(("zstd_parquet_2", prop)); + + let prop = WriterProperties::builder() + .set_content_defined_chunking(Some(CdcOptions::default())) + .build(); + props.push(("cdc", prop)); + + props +} + +fn bench_all_writers(c: &mut Criterion) { + let batches = create_batches(); + let props = create_writer_props(); + + for (batch_name, batch) in &batches { + let mut group = c.benchmark_group(*batch_name); + group.throughput(Throughput::Bytes( + batch + .columns() + .iter() + .map(|f| f.get_array_memory_size() as u64) + .sum(), + )); + + for (prop_name, prop) in &props { + group.bench_function(*prop_name, |b| { + write_batch_with_option(b, batch, Some(prop.clone())).unwrap() + }); + } + group.finish(); + } +} + +criterion_group!(benches, bench_all_writers); +criterion_main!(benches); diff --git a/vendor/parquet/benches/bloom_filter.rs b/vendor/parquet/benches/bloom_filter.rs new file mode 100644 index 0000000..cb2bfb6 --- /dev/null +++ b/vendor/parquet/benches/bloom_filter.rs @@ -0,0 +1,230 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::hint; + +use criterion::{BatchSize, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use parquet::bloom_filter::Sbbf; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; + +/// Build a bloom filter sized for `initial_ndv` at `fpp`, insert `num_values` distinct values, +/// and return it ready for folding. +fn build_filter(initial_ndv: u64, fpp: f64, num_values: u64) -> Sbbf { + let mut sbbf = Sbbf::new_with_ndv_fpp(initial_ndv, fpp).unwrap(); + for i in 0..num_values { + sbbf.insert(&i); + } + sbbf +} + +fn bench_fold_to_target_fpp(c: &mut Criterion) { + let mut group = c.benchmark_group("fold_to_target_fpp"); + + // Realistic scenario: filter sized for 1M NDV, varying actual distinct values + let initial_ndv = 1_000_000u64; + let fpp = 0.05; + + for num_values in [1_000u64, 10_000, 100_000] { + let filter = build_filter(initial_ndv, fpp, num_values); + let num_blocks = filter.num_blocks(); + group.throughput(Throughput::Elements(num_blocks as u64)); + group.bench_with_input(BenchmarkId::new("ndv", num_values), &filter, |b, filter| { + b.iter_batched( + || filter.clone(), + |mut f| { + f.fold_to_target_fpp(fpp); + f + }, + BatchSize::SmallInput, + ); + }); + } + group.finish(); +} + +fn bench_insert_and_fold(c: &mut Criterion) { + let mut group = c.benchmark_group("insert_and_fold"); + + let initial_ndv = 1_000_000u64; + let fpp = 0.05; + + for num_values in [1_000u64, 10_000, 100_000] { + group.throughput(Throughput::Elements(num_values)); + group.bench_with_input( + BenchmarkId::new("values", num_values), + &num_values, + |b, &num_values| { + b.iter(|| { + let mut sbbf = Sbbf::new_with_ndv_fpp(initial_ndv, fpp).unwrap(); + for i in 0..num_values { + sbbf.insert(&i); + } + sbbf.fold_to_target_fpp(fpp); + sbbf + }); + }, + ); + } + group.finish(); +} + +fn bench_insert_only(c: &mut Criterion) { + let mut group = c.benchmark_group("insert_only"); + + let initial_ndv = 1_000_000u64; + let fpp = 0.05; + + for num_values in [1_000u64, 10_000, 100_000] { + group.throughput(Throughput::Elements(num_values)); + group.bench_with_input( + BenchmarkId::new("values", num_values), + &num_values, + |b, &num_values| { + b.iter(|| { + let mut sbbf = Sbbf::new_with_ndv_fpp(initial_ndv, fpp).unwrap(); + for i in 0..num_values { + sbbf.insert(&i); + } + sbbf + }); + }, + ); + } + group.finish(); +} + +/// Benchmark `Sbbf::insert` across the same three cache regimes as +/// `bench_check`. The filter is allocated once per regime and reused +/// across iterations — bloom inserts are idempotent on identical +/// input, so this measures the pure insert kernel cost (hash + mask + +/// load + OR + store) without per-iteration allocation noise. +fn bench_insert(c: &mut Criterion) { + let regimes: [(&str, usize); 3] = [ + ("s_128KiB", 128 * 1024), + ("m_2MiB", 2 * 1024 * 1024), + ("l_32MiB", 32 * 1024 * 1024), + ]; + const NUM_INSERTS: usize = 50_000; + + let mut group = c.benchmark_group("insert"); + + for (label, num_bytes) in regimes { + let mut rng = StdRng::seed_from_u64(0xC0FFEE); + let keys: Vec<[u8; 16]> = (0..NUM_INSERTS) + .map(|_| { + let mut k = [0u8; 16]; + rng.fill(&mut k); + k + }) + .collect(); + + let mut filter = Sbbf::new_with_num_of_bytes(num_bytes); + + group.throughput(Throughput::Elements(NUM_INSERTS as u64)); + group.bench_with_input(BenchmarkId::new("ins", label), &keys, |b, keys| { + b.iter(|| { + for k in keys { + filter.insert(hint::black_box(k.as_slice())); + } + }); + }); + } + group.finish(); +} + +/// Benchmark `Sbbf::check` across three cache regimes and both miss-heavy +/// (the common case for Parquet row-group skipping) and hit-heavy workloads. +/// +/// The three filter sizes span the cache hierarchy on a typical server: +/// 128 KB fits in L2, 2 MB lives in L3, 32 MB spills out of L3. +fn bench_check(c: &mut Criterion) { + let regimes: [(&str, usize, usize); 3] = [ + ("s_128KiB", 128 * 1024, 10_000), + ("m_2MiB", 2 * 1024 * 1024, 200_000), + ("l_32MiB", 32 * 1024 * 1024, 3_000_000), + ]; + const NUM_QUERIES: usize = 50_000; + + let mut group = c.benchmark_group("check"); + + for (label, num_bytes, ndv) in regimes { + let mut rng = StdRng::seed_from_u64(0xC0FFEE); + + // Clear the high bit of byte 0 on inserted keys so the miss set below + // — which forces the same bit to 1 — is genuinely disjoint by + // construction (not just disjoint at birthday-paradox probability). + let mut filter = Sbbf::new_with_num_of_bytes(num_bytes); + let mut inserted: Vec<[u8; 16]> = Vec::with_capacity(ndv); + for _ in 0..ndv { + let mut k = [0u8; 16]; + rng.fill(&mut k); + k[0] &= 0x7f; + filter.insert(k.as_slice()); + inserted.push(k); + } + + // Disjoint miss set: high bit of byte 0 is 1 here and 0 in `inserted`, + // so no miss key can equal an inserted key. + let miss_keys: Vec<[u8; 16]> = (0..NUM_QUERIES) + .map(|_| { + let mut k = [0u8; 16]; + rng.fill(&mut k); + k[0] |= 0x80; + k + }) + .collect(); + let hit_keys: Vec<[u8; 16]> = (0..NUM_QUERIES) + .map(|_| inserted[rng.random_range(0..inserted.len())]) + .collect(); + + group.throughput(Throughput::Elements(NUM_QUERIES as u64)); + group.bench_with_input(BenchmarkId::new("miss", label), &miss_keys, |b, keys| { + b.iter(|| { + let mut found = 0u64; + for k in keys { + if hint::black_box(filter.check(k.as_slice())) { + found += 1; + } + } + hint::black_box(found) + }); + }); + group.bench_with_input(BenchmarkId::new("hit", label), &hit_keys, |b, keys| { + b.iter(|| { + let mut found = 0u64; + for k in keys { + if hint::black_box(filter.check(k.as_slice())) { + found += 1; + } + } + hint::black_box(found) + }); + }); + } + group.finish(); +} + +criterion_group!( + benches, + bench_fold_to_target_fpp, + bench_insert_and_fold, + bench_insert_only, + bench_insert, + bench_check, +); +criterion_main!(benches); diff --git a/vendor/parquet/benches/compression.rs b/vendor/parquet/benches/compression.rs new file mode 100644 index 0000000..b90c048 --- /dev/null +++ b/vendor/parquet/benches/compression.rs @@ -0,0 +1,101 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use criterion::*; +use parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel}; +use parquet::compression::create_codec; +use rand::distr::Alphanumeric; +use rand::prelude::*; +use std::hint; + +fn do_bench(c: &mut Criterion, name: &str, uncompressed: &[u8]) { + let codecs = [ + Compression::BROTLI(BrotliLevel::default()), + Compression::GZIP(GzipLevel::default()), + Compression::LZ4, + Compression::LZ4_RAW, + Compression::SNAPPY, + Compression::GZIP(GzipLevel::default()), + Compression::ZSTD(ZstdLevel::default()), + ]; + + for compression in codecs { + let mut codec = create_codec(compression, &Default::default()) + .unwrap() + .unwrap(); + + c.bench_function(&format!("compress {compression} - {name}"), |b| { + b.iter(|| { + let mut out = Vec::new(); + codec.compress(uncompressed, &mut out).unwrap(); + out + }); + }); + + let mut compressed = Vec::new(); + codec.compress(uncompressed, &mut compressed).unwrap(); + println!( + "{compression} compressed {} bytes of {name} to {} bytes", + uncompressed.len(), + compressed.len() + ); + + c.bench_function(&format!("decompress {compression} - {name}"), |b| { + b.iter(|| { + let mut out = Vec::new(); + codec + .decompress( + hint::black_box(&compressed), + &mut out, + Some(uncompressed.len()), + ) + .unwrap(); + out + }); + }); + } +} + +fn criterion_benchmark(c: &mut Criterion) { + let mut rng = StdRng::seed_from_u64(42); + let rng = &mut rng; + const DATA_SIZE: usize = 1024 * 1024; + + let uncompressed: Vec<_> = rng.sample_iter(&Alphanumeric).take(DATA_SIZE).collect(); + do_bench(c, "alphanumeric", &uncompressed); + + // Create a collection of 64 words + let words: Vec> = (0..64) + .map(|_| { + let len = rng.random_range(1..12); + rng.sample_iter(&Alphanumeric).take(len).collect() + }) + .collect(); + + // Build data by concatenating these words randomly together + let mut uncompressed = Vec::with_capacity(DATA_SIZE); + while uncompressed.len() < DATA_SIZE { + let word = &words[rng.random_range(0..words.len())]; + uncompressed.extend_from_slice(&word[..word.len().min(DATA_SIZE - uncompressed.len())]) + } + assert_eq!(uncompressed.len(), DATA_SIZE); + + do_bench(c, "words", &uncompressed); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/vendor/parquet/benches/encoding.rs b/vendor/parquet/benches/encoding.rs new file mode 100644 index 0000000..65c2ec3 --- /dev/null +++ b/vendor/parquet/benches/encoding.rs @@ -0,0 +1,105 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use criterion::*; +use half::f16; +use parquet::basic::{Encoding, Type as ParquetType}; +use parquet::data_type::{ + DataType, DoubleType, FixedLenByteArray, FixedLenByteArrayType, FloatType, +}; +use parquet::decoding::{Decoder, get_decoder}; +use parquet::encoding::get_encoder; +use parquet::schema::types::{ColumnDescPtr, ColumnDescriptor, ColumnPath, Type}; +use rand::prelude::*; +use std::sync::Arc; + +fn bench_typed( + c: &mut Criterion, + values: &[T::T], + encoding: Encoding, + type_length: i32, +) { + let name = format!( + "dtype={}, encoding={:?}", + match T::get_physical_type() { + ParquetType::FIXED_LEN_BYTE_ARRAY => format!("FixedLenByteArray({type_length})"), + _ => std::any::type_name::().to_string(), + }, + encoding + ); + let column_desc_ptr = ColumnDescPtr::new(ColumnDescriptor::new( + Arc::new( + Type::primitive_type_builder("", T::get_physical_type()) + .with_length(type_length) + .build() + .unwrap(), + ), + 0, + 0, + ColumnPath::new(vec![]), + )); + c.bench_function(&format!("encoding: {name}"), |b| { + b.iter(|| { + let mut encoder = get_encoder::(encoding, &column_desc_ptr).unwrap(); + encoder.put(values).unwrap(); + encoder.flush_buffer().unwrap(); + }); + }); + + let mut encoder = get_encoder::(encoding, &column_desc_ptr).unwrap(); + encoder.put(values).unwrap(); + let encoded = encoder.flush_buffer().unwrap(); + println!("{} encoded as {} bytes", name, encoded.len(),); + + let mut buffer = vec![T::T::default(); values.len()]; + c.bench_function(&format!("decoding: {name}"), |b| { + b.iter(|| { + let mut decoder: Box> = + get_decoder(column_desc_ptr.clone(), encoding).unwrap(); + decoder.set_data(encoded.clone(), values.len()).unwrap(); + decoder.get(&mut buffer).unwrap(); + }); + }); +} + +fn criterion_benchmark(c: &mut Criterion) { + let mut rng = StdRng::seed_from_u64(0); + let n = 16 * 1024; + + let mut f16s = Vec::new(); + let mut f32s = Vec::new(); + let mut f64s = Vec::new(); + let mut d128s = Vec::new(); + for _ in 0..n { + f16s.push(FixedLenByteArray::from( + f16::from_f32(rng.random::()).to_le_bytes().to_vec(), + )); + f32s.push(rng.random::()); + f64s.push(rng.random::()); + d128s.push(FixedLenByteArray::from( + rng.random::().to_be_bytes().to_vec(), + )); + } + + bench_typed::(c, &f32s, Encoding::BYTE_STREAM_SPLIT, 0); + bench_typed::(c, &f64s, Encoding::BYTE_STREAM_SPLIT, 0); + bench_typed::(c, &f16s, Encoding::BYTE_STREAM_SPLIT, 2); + bench_typed::(c, &d128s, Encoding::BYTE_STREAM_SPLIT, 16); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/vendor/parquet/benches/metadata.rs b/vendor/parquet/benches/metadata.rs new file mode 100644 index 0000000..b6e0458 --- /dev/null +++ b/vendor/parquet/benches/metadata.rs @@ -0,0 +1,317 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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 std::sync::Arc; + +use parquet::basic::{Encoding, PageType, Type as PhysicalType}; +use parquet::file::metadata::{ + ColumnChunkMetaData, FileMetaData, LevelHistogram, PageEncodingStats, ParquetMetaData, + ParquetMetaDataOptions, ParquetMetaDataReader, ParquetMetaDataWriter, ParquetStatisticsPolicy, + RowGroupMetaData, +}; +use parquet::file::statistics::Statistics; +use parquet::file::writer::TrackedWrite; +use parquet::schema::parser::parse_message_type; +use parquet::schema::types::{ + ColumnDescPtr, ColumnDescriptor, ColumnPath, SchemaDescriptor, Type as SchemaType, +}; +use rand::Rng; + +use arrow::util::test_util::seedable_rng; +use bytes::Bytes; +use criterion::{Criterion, criterion_group, criterion_main}; +use parquet::file::reader::SerializedFileReader; +use parquet::file::serialized_reader::ReadOptionsBuilder; + +const NUM_COLUMNS: usize = 10_000; +const NUM_ROW_GROUPS: usize = 10; + +fn encoded_meta(is_nullable: bool, has_lists: bool, write_path_in_schema: bool) -> Vec { + let mut rng = seedable_rng(); + + let mut column_desc_ptrs: Vec = Vec::with_capacity(NUM_COLUMNS); + let mut message_type = "message test_schema {".to_string(); + for i in 0..NUM_COLUMNS { + message_type.push_str(&format!("REQUIRED FLOAT {};", i)); + column_desc_ptrs.push(ColumnDescPtr::new(ColumnDescriptor::new( + Arc::new( + SchemaType::primitive_type_builder(&i.to_string(), PhysicalType::FLOAT) + .build() + .unwrap(), + ), + 0, + 0, + ColumnPath::new(vec![]), + ))); + } + message_type.push('}'); + + let schema_descr = parse_message_type(&message_type) + .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t)))) + .unwrap(); + + let stats = Statistics::float(Some(rng.random()), Some(rng.random()), None, Some(0), false); + + let (var_size, rep_hist, def_hist) = match (is_nullable, has_lists) { + (true, true) => { + let rep_hist = LevelHistogram::from(vec![1500i64; 2]); + let def_hist = LevelHistogram::from(vec![1000i64; 3]); + ( + Some(rng.random_range(0..1000000000)), + Some(rep_hist), + Some(def_hist), + ) + } + (true, false) => { + let def_hist = LevelHistogram::from(vec![1500i64; 2]); + (Some(rng.random_range(0..1000000000)), None, Some(def_hist)) + } + (_, _) => (None, None, None), + }; + + let row_groups = (0..NUM_ROW_GROUPS) + .map(|i| { + let columns = (0..NUM_COLUMNS) + .map(|j| { + ColumnChunkMetaData::builder(column_desc_ptrs[j].clone()) + .set_encodings(vec![Encoding::PLAIN, Encoding::RLE_DICTIONARY]) + .set_compression_codec(parquet::basic::CompressionCodec::UNCOMPRESSED) + .set_num_values(rng.random_range(1..1000000)) + .set_total_compressed_size(rng.random_range(50000..5000000)) + .set_data_page_offset(rng.random_range(4..2000000000)) + .set_dictionary_page_offset(Some(rng.random_range(4..2000000000))) + .set_statistics(stats.clone()) + .set_page_encoding_stats(vec![ + PageEncodingStats { + page_type: PageType::DICTIONARY_PAGE, + encoding: Encoding::PLAIN, + count: 1, + }, + PageEncodingStats { + page_type: PageType::DATA_PAGE, + encoding: Encoding::RLE_DICTIONARY, + count: 10, + }, + ]) + .set_offset_index_offset(Some(rng.random_range(0..2000000000))) + .set_offset_index_length(Some(rng.random_range(1..100000))) + .set_column_index_offset(Some(rng.random_range(0..2000000000))) + .set_column_index_length(Some(rng.random_range(1..100000))) + .set_unencoded_byte_array_data_bytes(var_size) + .set_repetition_level_histogram(rep_hist.clone()) + .set_definition_level_histogram(def_hist.clone()) + .build() + .unwrap() + }) + .collect(); + + RowGroupMetaData::builder(schema_descr.clone()) + .set_column_metadata(columns) + .set_total_byte_size(rng.random_range(1..2000000000)) + .set_num_rows(rng.random_range(1..10000000000)) + .set_ordinal(i as i16) + .build() + .unwrap() + }) + .collect(); + + let file_metadata = FileMetaData::new( + 1, + rng.random_range(1..2000000000), + Some("parquet-rs".into()), + None, + schema_descr, + None, + ); + + let metadata = ParquetMetaData::new(file_metadata, row_groups); + let mut buffer = Vec::with_capacity(1024); + { + let buf = TrackedWrite::new(&mut buffer); + let mut writer = ParquetMetaDataWriter::new_with_tracked(buf, &metadata); + // use defaults unless `write_path_in_schema` is false + if !write_path_in_schema { + writer = writer.with_write_path_in_schema(write_path_in_schema); + } + writer.finish().unwrap(); + } + + buffer +} + +fn get_footer_bytes(data: Bytes) -> Bytes { + let footer_bytes = data.slice(data.len() - 8..); + let footer_len = footer_bytes[0] as u32 + | (footer_bytes[1] as u32) << 8 + | (footer_bytes[2] as u32) << 16 + | (footer_bytes[3] as u32) << 24; + let meta_start = data.len() - footer_len as usize - 8; + let meta_end = data.len() - 8; + data.slice(meta_start..meta_end) +} + +fn criterion_benchmark(c: &mut Criterion) { + // Read file into memory to isolate filesystem performance + let file = "../parquet-testing/data/alltypes_tiny_pages.parquet"; + let data = std::fs::read(file).unwrap(); + let data = Bytes::from(data); + + c.bench_function("open(default)", |b| { + b.iter(|| { + let options = ReadOptionsBuilder::new() + .with_encoding_stats_as_mask(false) + .build(); + SerializedFileReader::new_with_options(data.clone(), options).unwrap() + }) + }); + + c.bench_function("open(page index)", |b| { + b.iter(|| { + let options = ReadOptionsBuilder::new() + .with_page_index() + .with_encoding_stats_as_mask(false) + .build(); + SerializedFileReader::new_with_options(data.clone(), options).unwrap() + }) + }); + + let meta_data = get_footer_bytes(data.clone()); + let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false); + c.bench_function("decode parquet metadata", |b| { + b.iter(|| { + ParquetMetaDataReader::decode_metadata_with_options(&meta_data, Some(&options)) + .unwrap(); + }) + }); + + let schema = ParquetMetaDataReader::decode_schema(&meta_data).unwrap(); + let options = ParquetMetaDataOptions::new() + .with_schema(schema) + .with_encoding_stats_as_mask(false); + c.bench_function("decode metadata with schema", |b| { + b.iter(|| { + ParquetMetaDataReader::decode_metadata_with_options(&meta_data, Some(&options)) + .unwrap(); + }) + }); + + let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(true); + c.bench_function("decode metadata with stats mask", |b| { + b.iter(|| { + ParquetMetaDataReader::decode_metadata_with_options(&meta_data, Some(&options)) + .unwrap(); + }) + }); + + let options = + ParquetMetaDataOptions::new().with_encoding_stats_policy(ParquetStatisticsPolicy::SkipAll); + c.bench_function("decode metadata with skip PES", |b| { + b.iter(|| { + ParquetMetaDataReader::decode_metadata_with_options(&meta_data, Some(&options)) + .unwrap(); + }) + }); + + let options = ParquetMetaDataOptions::new() + .with_column_stats_policy(ParquetStatisticsPolicy::SkipAll) + .with_encoding_stats_as_mask(false); + c.bench_function("decode metadata with skip column stats", |b| { + b.iter(|| { + ParquetMetaDataReader::decode_metadata_with_options(&meta_data, Some(&options)) + .unwrap(); + }) + }); + + let buf: Bytes = black_box(encoded_meta(false, false, true)).into(); + let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false); + c.bench_function("decode parquet metadata (wide)", |b| { + b.iter(|| { + ParquetMetaDataReader::decode_metadata_with_options(&buf, Some(&options)).unwrap(); + }) + }); + + let schema = ParquetMetaDataReader::decode_schema(&buf).unwrap(); + let options = ParquetMetaDataOptions::new() + .with_schema(schema) + .with_encoding_stats_as_mask(false); + c.bench_function("decode metadata (wide) with schema", |b| { + b.iter(|| { + ParquetMetaDataReader::decode_metadata_with_options(&buf, Some(&options)).unwrap(); + }) + }); + + let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(true); + c.bench_function("decode metadata (wide) with stats mask", |b| { + b.iter(|| { + ParquetMetaDataReader::decode_metadata_with_options(&buf, Some(&options)).unwrap(); + }) + }); + + let options = + ParquetMetaDataOptions::new().with_encoding_stats_policy(ParquetStatisticsPolicy::SkipAll); + c.bench_function("decode metadata (wide) with skip PES", |b| { + b.iter(|| { + ParquetMetaDataReader::decode_metadata_with_options(&buf, Some(&options)).unwrap(); + }) + }); + + let options = ParquetMetaDataOptions::new() + .with_column_stats_policy(ParquetStatisticsPolicy::SkipAll) + .with_encoding_stats_as_mask(false); + c.bench_function("decode metadata (wide) with skip column stats", |b| { + b.iter(|| { + ParquetMetaDataReader::decode_metadata_with_options(&buf, Some(&options)).unwrap(); + }) + }); + + let buf: Bytes = black_box(encoded_meta(false, false, false)).into(); + let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false); + c.bench_function("decode parquet metadata no path_in_schema (wide)", |b| { + b.iter(|| { + ParquetMetaDataReader::decode_metadata_with_options(&buf, Some(&options)).unwrap(); + }) + }); + + let buf: Bytes = black_box(encoded_meta(true, true, true)).into(); + c.bench_function("decode parquet metadata w/ size stats (wide)", |b| { + b.iter(|| { + ParquetMetaDataReader::decode_metadata(&buf).unwrap(); + }) + }); + + let options = + ParquetMetaDataOptions::new().with_size_stats_policy(ParquetStatisticsPolicy::SkipAll); + c.bench_function("decode metadata (wide) with skip size stats", |b| { + b.iter(|| { + ParquetMetaDataReader::decode_metadata_with_options(&buf, Some(&options)).unwrap(); + }) + }); + + let options = ParquetMetaDataOptions::new() + .with_column_stats_policy(ParquetStatisticsPolicy::SkipAll) + .with_encoding_stats_policy(ParquetStatisticsPolicy::SkipAll) + .with_size_stats_policy(ParquetStatisticsPolicy::SkipAll); + c.bench_function("decode metadata (wide) with skip all stats", |b| { + b.iter(|| { + ParquetMetaDataReader::decode_metadata_with_options(&buf, Some(&options)).unwrap(); + }) + }); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/vendor/parquet/benches/parquet_round_trip.rs b/vendor/parquet/benches/parquet_round_trip.rs new file mode 100644 index 0000000..7f46b02 --- /dev/null +++ b/vendor/parquet/benches/parquet_round_trip.rs @@ -0,0 +1,473 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::{ArrayRef, RecordBatch}; +use arrow::datatypes::{DataType, Field, Float32Type, Float64Type, Int32Type, Int64Type, Schema}; +use arrow::util::bench_util::{ + create_binary_array_with_len_range_and_prefix_and_seed, create_primitive_array_with_seed, + create_string_array_with_len_range_and_prefix_and_seed, +}; +use arrow_array::{FixedSizeBinaryArray, StringViewArray}; +use bytes::Bytes; +use criterion::{Criterion, criterion_group, criterion_main}; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use parquet::basic::Encoding; +use parquet::file::properties::WriterProperties; +use rand::{ + Rng, SeedableRng, + distr::{Alphanumeric, StandardUniform}, + prelude::StdRng, +}; +use std::sync::Arc; + +#[derive(Copy, Clone)] +pub enum ColumnType { + String(usize), + StringView(usize), + Binary(usize), + FixedLen(i32), + Int32, + Int64, + Float, + Double, +} + +// arrow::util::bench_util::create_fsb_array with a seed + +/// Creates a random (but fixed-seeded) array of fixed size with a given null density and length +fn create_fsb_array_with_seed( + size: usize, + null_density: f32, + fixed_len: i32, + seed: u64, +) -> FixedSizeBinaryArray { + let mut rng = StdRng::seed_from_u64(seed); + + let rng = &mut rng; + FixedSizeBinaryArray::try_from_sparse_iter_with_size( + (0..size).map(|_| { + if rng.random::() < null_density { + None + } else { + let value = rng + .sample_iter::(StandardUniform) + .take(fixed_len as usize) + .collect::>(); + Some(value) + } + }), + fixed_len, + ) + .unwrap() +} + +// arrow::util::bench_util::create_string_view_array_with_max_len with a seed + +/// Creates a random (but fixed-seeded) array of rand size with a given max size, null density and length +pub fn create_string_view_array_with_seed( + size: usize, + null_density: f32, + max_str_len: usize, + seed: u64, +) -> StringViewArray { + let mut rng = StdRng::seed_from_u64(seed); + + let rng = &mut rng; + (0..size) + .map(|_| { + if rng.random::() < null_density { + None + } else { + let str_len = rng.random_range(max_str_len / 2..max_str_len); + let value = rng.sample_iter(&Alphanumeric).take(str_len).collect(); + let value = String::from_utf8(value).unwrap(); + Some(value) + } + }) + .collect() +} + +fn schema(column_type: ColumnType, num_columns: usize) -> Arc { + let field_type = match column_type { + ColumnType::Binary(_) => DataType::Binary, + ColumnType::String(_) => DataType::Utf8, + ColumnType::StringView(_) => DataType::Utf8View, + ColumnType::FixedLen(size) => DataType::FixedSizeBinary(size), + ColumnType::Int32 => DataType::Int32, + ColumnType::Int64 => DataType::Int64, + ColumnType::Float => DataType::Float32, + ColumnType::Double => DataType::Float64, + }; + + let fields: Vec = (0..num_columns) + .map(|i| Field::new(format!("col_{i}"), field_type.clone(), true)) + .collect(); + Arc::new(Schema::new(fields)) +} + +fn create_batch( + schema: &Arc, + column_type: ColumnType, + seed: usize, + num_columns: usize, + num_rows: usize, +) -> RecordBatch { + let null_density = 0.0001; + let mut arrays: Vec = vec![]; + match column_type { + ColumnType::Binary(max_len) => { + for i in 0..num_columns { + let array_seed = seed * num_columns + i; + let array = create_binary_array_with_len_range_and_prefix_and_seed::( + num_rows, + null_density, + max_len / 2, + max_len, + &[], + array_seed as u64, + ); + arrays.push(Arc::new(array)); + } + } + ColumnType::String(max_str_len) => { + for i in 0..num_columns { + let array_seed = seed * num_columns + i; + let array = create_string_array_with_len_range_and_prefix_and_seed::( + num_rows, + null_density, + max_str_len / 2, + max_str_len, + "", + array_seed as u64, + ); + arrays.push(Arc::new(array)); + } + } + ColumnType::StringView(max_str_len) => { + for i in 0..num_columns { + let array_seed = seed * num_columns + i; + let array = create_string_view_array_with_seed( + num_rows, + null_density, + max_str_len, + array_seed as u64, + ); + arrays.push(Arc::new(array)); + } + } + ColumnType::FixedLen(size) => { + for i in 0..num_columns { + let array_seed = seed * num_columns + i; + let array = + create_fsb_array_with_seed(num_rows, null_density, size, array_seed as u64); + arrays.push(Arc::new(array)); + } + } + ColumnType::Int32 => { + for i in 0..num_columns { + let array_seed = seed * num_columns + i; + let array = create_primitive_array_with_seed::( + num_rows, + null_density, + array_seed as u64, + ); + arrays.push(Arc::new(array)); + } + } + ColumnType::Int64 => { + for i in 0..num_columns { + let array_seed = seed * num_columns + i; + let array = create_primitive_array_with_seed::( + num_rows, + null_density, + array_seed as u64, + ); + arrays.push(Arc::new(array)); + } + } + ColumnType::Float => { + for i in 0..num_columns { + let array_seed = seed * num_columns + i; + let array = create_primitive_array_with_seed::( + num_rows, + null_density, + array_seed as u64, + ); + arrays.push(Arc::new(array)); + } + } + ColumnType::Double => { + for i in 0..num_columns { + let array_seed = seed * num_columns + i; + let array = create_primitive_array_with_seed::( + num_rows, + null_density, + array_seed as u64, + ); + arrays.push(Arc::new(array)); + } + } + } + RecordBatch::try_new(schema.clone(), arrays).unwrap() +} + +#[derive(Copy, Clone)] +pub struct ParquetFileSpec { + column_type: ColumnType, + num_columns: usize, + num_row_groups: usize, + rows_per_row_group: usize, + rows_per_page: usize, + encoding: Encoding, + use_dict: bool, +} + +const DEFAULT_NUM_COLUMNS: usize = 10; +const DEFAULT_NUM_ROWGROUPS: usize = 10; +const DEFAULT_ROWS_PER_PAGE: usize = 2_000; +const DEFAULT_ROWS_PER_ROWGROUP: usize = 10_000; + +impl ParquetFileSpec { + pub fn new(column_type: ColumnType) -> Self { + Self { + column_type, + num_columns: DEFAULT_NUM_COLUMNS, + num_row_groups: DEFAULT_NUM_ROWGROUPS, + rows_per_row_group: DEFAULT_ROWS_PER_ROWGROUP, + rows_per_page: DEFAULT_ROWS_PER_PAGE, + encoding: Encoding::PLAIN, + use_dict: true, + } + } + + pub fn with_num_columns(self, num_columns: usize) -> Self { + Self { + num_columns, + ..self + } + } + + pub fn with_num_row_groups(self, num_row_groups: usize) -> Self { + Self { + num_row_groups, + ..self + } + } + + pub fn with_rows_per_row_group(self, rows_per_row_group: usize) -> Self { + Self { + rows_per_row_group, + ..self + } + } + + pub fn with_rows_per_page(self, rows_per_page: usize) -> Self { + Self { + rows_per_page, + ..self + } + } + + pub fn with_encoding(self, encoding: Encoding) -> Self { + Self { encoding, ..self } + } + + pub fn with_use_dict(self, use_dict: bool) -> Self { + Self { use_dict, ..self } + } +} + +fn file_from_spec(spec: ParquetFileSpec, buffer: &mut Vec) { + const SEED: usize = 31; + let num_rows = spec.rows_per_row_group.min(100); + let rows_to_write = spec.num_row_groups * spec.rows_per_row_group; + + let schema = schema(spec.column_type, spec.num_columns); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(spec.rows_per_row_group)) + .set_data_page_row_count_limit(spec.rows_per_page) + .set_encoding(spec.encoding) + .set_dictionary_enabled(spec.use_dict) + .set_compression(parquet::basic::Compression::UNCOMPRESSED) + .build(); + + let mut writer = ArrowWriter::try_new(buffer, schema.clone(), Some(props)).unwrap(); + + // use the same batch repeatedly otherwise the data generation will dominate the time + let batch = create_batch(&schema, spec.column_type, SEED, spec.num_columns, num_rows); + + let mut rows_written = 0; + while rows_written < rows_to_write { + writer.write(&batch).unwrap(); + rows_written += num_rows; + } + + let parquet_metadata = writer.close().unwrap(); + assert_eq!(parquet_metadata.num_row_groups(), spec.num_row_groups); + assert_eq!( + parquet_metadata.file_metadata().num_rows() as usize, + rows_to_write + ); +} + +fn read_write(c: &mut Criterion, spec: ParquetFileSpec, msg: &str) { + let mut buffer = Vec::with_capacity(1_000_000); + + // read once to size the buffer + file_from_spec(spec, &mut buffer); + + c.bench_function(&format!("write {msg}"), |b| { + buffer.clear(); + b.iter(|| file_from_spec(spec, &mut buffer)) + }); + + let file_bytes = Bytes::from(buffer); + c.bench_function(&format!("read {msg}"), |b| { + b.iter(|| { + let record_reader = ParquetRecordBatchReaderBuilder::try_new(file_bytes.clone()) + .unwrap() + .build() + .unwrap(); + let mut num_rows = 0; + for maybe_batch in record_reader { + let batch = maybe_batch.unwrap(); + num_rows += batch.num_rows(); + } + assert_eq!(num_rows, spec.num_row_groups * spec.rows_per_row_group); + }) + }); +} + +fn int_benches(c: &mut Criterion, column_type: ColumnType) { + let ctype = match column_type { + ColumnType::Int32 => "int32", + ColumnType::Int64 => "int64", + _ => unreachable!(), + }; + + let spec = ParquetFileSpec::new(column_type).with_use_dict(true); + read_write(c, spec, &format!("{ctype} dict")); + + let spec = spec.with_use_dict(false).with_encoding(Encoding::PLAIN); + read_write(c, spec, &format!("{ctype} plain")); + + let spec = spec.with_encoding(Encoding::DELTA_BINARY_PACKED); + read_write(c, spec, &format!("{ctype} delta_binary")); + + let spec = spec.with_encoding(Encoding::BYTE_STREAM_SPLIT); + read_write(c, spec, &format!("{ctype} byte_stream_split")); +} + +fn float_benches(c: &mut Criterion, column_type: ColumnType) { + let ctype = match column_type { + ColumnType::Float => "f32", + ColumnType::Double => "f64", + _ => unreachable!(), + }; + + let spec = ParquetFileSpec::new(column_type).with_use_dict(true); + read_write(c, spec, &format!("{ctype} dict")); + + let spec = spec.with_use_dict(false).with_encoding(Encoding::PLAIN); + read_write(c, spec, &format!("{ctype} plain")); + + let spec = spec.with_encoding(Encoding::BYTE_STREAM_SPLIT); + read_write(c, spec, &format!("{ctype} byte_stream_split")); +} + +fn string_benches(c: &mut Criterion, max_str_len: usize) { + let spec = ParquetFileSpec::new(ColumnType::String(max_str_len)) + .with_num_columns(5) + .with_use_dict(true); + read_write(c, spec, &format!("String({max_str_len}) dict")); + + let spec = spec.with_use_dict(false).with_encoding(Encoding::PLAIN); + read_write(c, spec, &format!("String({max_str_len}) plain")); + + let spec = spec.with_encoding(Encoding::DELTA_LENGTH_BYTE_ARRAY); + read_write(c, spec, &format!("String({max_str_len}) delta_length")); + + let spec = spec.with_encoding(Encoding::DELTA_BYTE_ARRAY); + read_write(c, spec, &format!("String({max_str_len}) delta_byte_array")); + + let spec = ParquetFileSpec::new(ColumnType::StringView(max_str_len)) + .with_num_columns(5) + .with_use_dict(true); + read_write(c, spec, &format!("StringView({max_str_len}) dict")); + + let spec = spec.with_use_dict(false).with_encoding(Encoding::PLAIN); + read_write(c, spec, &format!("StringView({max_str_len}) plain")); + + let spec = spec.with_encoding(Encoding::DELTA_LENGTH_BYTE_ARRAY); + read_write(c, spec, &format!("StringView({max_str_len}) delta_length")); + + let spec = spec.with_encoding(Encoding::DELTA_BYTE_ARRAY); + read_write( + c, + spec, + &format!("StringView({max_str_len}) delta_byte_array"), + ); +} + +fn binary_benches(c: &mut Criterion, max_len: usize) { + let spec = ParquetFileSpec::new(ColumnType::Binary(max_len)) + .with_num_columns(5) + .with_use_dict(true); + read_write(c, spec, &format!("Binary({max_len}) dict")); + + let spec = spec.with_use_dict(false).with_encoding(Encoding::PLAIN); + read_write(c, spec, &format!("Binary({max_len}) plain")); + + let spec = spec.with_encoding(Encoding::DELTA_LENGTH_BYTE_ARRAY); + read_write(c, spec, &format!("Binary({max_len}) delta_length")); + + let spec = spec.with_encoding(Encoding::DELTA_BYTE_ARRAY); + read_write(c, spec, &format!("Binary({max_len}) delta_byte_array")); +} + +fn flba_benches(c: &mut Criterion, len: i32) { + let spec = ParquetFileSpec::new(ColumnType::FixedLen(len)) + .with_num_columns(5) + .with_use_dict(true); + read_write(c, spec, &format!("Fixed({len}) dict")); + + let spec = spec.with_use_dict(false).with_encoding(Encoding::PLAIN); + read_write(c, spec, &format!("Fixed({len}) plain")); + + let spec = spec.with_encoding(Encoding::BYTE_STREAM_SPLIT); + read_write(c, spec, &format!("Fixed({len}) byte_stream_split")); + + let spec = spec.with_encoding(Encoding::DELTA_BYTE_ARRAY); + read_write(c, spec, &format!("Fixed({len}) delta_byte_array")); +} + +fn criterion_benchmark(c: &mut Criterion) { + int_benches(c, ColumnType::Int32); + int_benches(c, ColumnType::Int64); + float_benches(c, ColumnType::Float); + float_benches(c, ColumnType::Double); + string_benches(c, 20); + string_benches(c, 100); + binary_benches(c, 20); + binary_benches(c, 100); + flba_benches(c, 2); + flba_benches(c, 16); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/vendor/parquet/benches/push_decoder.rs b/vendor/parquet/benches/push_decoder.rs new file mode 100644 index 0000000..190dd71 --- /dev/null +++ b/vendor/parquet/benches/push_decoder.rs @@ -0,0 +1,170 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks for the push-based decoder measuring PushBuffers overhead. +//! +//! Uses `try_next_reader` to build row group readers without decoding any +//! pages, isolating PushBuffers operations (has_range, get_bytes, clearing). + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow_array::{Float32Array, RecordBatch}; +use bytes::Bytes; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use parquet::DecodeResult; +use parquet::arrow::ArrowWriter; +use parquet::arrow::push_decoder::ParquetPushDecoderBuilder; +use parquet::file::metadata::ParquetMetaDataPushDecoder; +use parquet::file::properties::WriterProperties; + +fn make_wide_schema(num_columns: usize) -> SchemaRef { + let fields: Vec = (0..num_columns) + .map(|i| Field::new(format!("c{i}"), DataType::Float32, false)) + .collect(); + Arc::new(Schema::new(fields)) +} + +/// Write a Parquet file with `num_columns` columns, 10 row groups of 100 rows. +fn make_test_file(num_columns: usize) -> Bytes { + let num_rows = 1_000; + let rows_per_rg = 100; + let schema = make_wide_schema(num_columns); + let columns: Vec> = (0..num_columns) + .map(|_| Arc::new(Float32Array::from(vec![0.0f32; num_rows])) as _) + .collect(); + let batch = RecordBatch::try_new(schema.clone(), columns).unwrap(); + + let mut buf = Vec::new(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(rows_per_rg)) + .build(); + let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + Bytes::from(buf) +} + +fn decode_metadata(file_data: &Bytes) -> Arc { + let file_len = file_data.len() as u64; + let mut dec = ParquetMetaDataPushDecoder::try_new(file_len).unwrap(); + dec.push_range(0..file_len, file_data.clone()).unwrap(); + match dec.try_decode().unwrap() { + DecodeResult::Data(m) => Arc::new(m), + other => panic!("expected metadata, got {other:?}"), + } +} + +/// Push the entire file as one buffer, then build all row group readers. +fn build_readers_single_buffer( + file_data: &Bytes, + metadata: &Arc, +) { + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(metadata.clone()) + .unwrap() + .build() + .unwrap(); + + decoder + .push_range(0..file_data.len() as u64, file_data.clone()) + .unwrap(); + + loop { + match decoder.try_next_reader().unwrap() { + DecodeResult::Data(reader) => { + black_box(reader); + } + DecodeResult::Finished => break, + DecodeResult::NeedsData(r) => panic!("unexpected NeedsData: {r:?}"), + } + } +} + +/// Push one buffer per requested range, then build all row group readers. +fn build_readers_exact_ranges( + file_data: &Bytes, + metadata: &Arc, +) { + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(metadata.clone()) + .unwrap() + .build() + .unwrap(); + + loop { + match decoder.try_next_reader().unwrap() { + DecodeResult::Data(reader) => { + black_box(reader); + } + DecodeResult::Finished => break, + DecodeResult::NeedsData(ranges) => { + let buffers: Vec = ranges + .iter() + .map(|r| file_data.slice(r.start as usize..r.end as usize)) + .collect(); + decoder.push_ranges(ranges, buffers).unwrap(); + } + } + } +} + +fn bench_1buf(c: &mut Criterion) { + let mut group = c.benchmark_group("push_decoder/1buf"); + + for num_cols in [100, 1_000, 10_000, 50_000] { + let file_data = make_test_file(num_cols); + let metadata = decode_metadata(&file_data); + let num_ranges: usize = metadata + .row_groups() + .iter() + .map(|rg| rg.columns().len()) + .sum(); + + group.bench_with_input( + BenchmarkId::from_parameter(format!("{num_ranges}ranges")), + &(&file_data, &metadata), + |b, &(data, meta)| b.iter(|| build_readers_single_buffer(data, meta)), + ); + } + + group.finish(); +} + +fn bench_nbuf(c: &mut Criterion) { + let mut group = c.benchmark_group("push_decoder/Nbuf"); + + for num_cols in [100, 1_000, 10_000] { + let file_data = make_test_file(num_cols); + let metadata = decode_metadata(&file_data); + let num_ranges: usize = metadata + .row_groups() + .iter() + .map(|rg| rg.columns().len()) + .sum(); + + group.bench_with_input( + BenchmarkId::from_parameter(format!("{num_ranges}ranges")), + &(&file_data, &metadata), + |b, &(data, meta)| b.iter(|| build_readers_exact_ranges(data, meta)), + ); + } + + group.finish(); +} + +criterion_group!(benches, bench_1buf, bench_nbuf); +criterion_main!(benches); diff --git a/vendor/parquet/benches/row_group_index_reader.rs b/vendor/parquet/benches/row_group_index_reader.rs new file mode 100644 index 0000000..1fadfb5 --- /dev/null +++ b/vendor/parquet/benches/row_group_index_reader.rs @@ -0,0 +1,213 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks for RowGroupIndexReader performance +//! +//! This benchmark tests the performance of reading row group indices, +//! comparing single row group reads vs multiple row group reads. + +use arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; +use arrow::record_batch::RecordBatch; +use arrow_schema::extension::ExtensionType; +use bytes::Bytes; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::{ArrowReaderOptions, ParquetRecordBatchReaderBuilder}; +use parquet::basic::Compression; +use parquet::file::properties::WriterProperties; +use std::sync::Arc; + +#[derive(Debug, Default, Clone, Copy, PartialEq)] +struct RowGroupIndex; + +impl ExtensionType for RowGroupIndex { + const NAME: &'static str = "parquet.virtual.row_group_index"; + type Metadata = &'static str; + + fn metadata(&self) -> &Self::Metadata { + &"" + } + + fn serialize_metadata(&self) -> Option { + Some(String::default()) + } + + fn deserialize_metadata( + metadata: Option<&str>, + ) -> Result { + if metadata.is_some_and(str::is_empty) { + Ok("") + } else { + Err(arrow_schema::ArrowError::InvalidArgumentError( + "Virtual column extension type expects an empty string as metadata".to_owned(), + )) + } + } + + fn supports_data_type( + &self, + data_type: &ArrowDataType, + ) -> Result<(), arrow_schema::ArrowError> { + match data_type { + ArrowDataType::Int64 => Ok(()), + data_type => Err(arrow_schema::ArrowError::InvalidArgumentError(format!( + "Virtual column data type mismatch, expected Int64, found {data_type}" + ))), + } + } + + fn try_new( + data_type: &ArrowDataType, + _metadata: Self::Metadata, + ) -> Result { + RowGroupIndex.supports_data_type(data_type).map(|_| Self) + } +} + +fn create_test_file(num_row_groups: usize, rows_per_group: usize) -> Bytes { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", ArrowDataType::Int64, false), + Field::new("value", ArrowDataType::Int64, false), + ])); + + let mut buffer = Vec::new(); + let props = WriterProperties::builder() + .set_compression(Compression::SNAPPY) + .set_max_row_group_row_count(Some(rows_per_group)) + .build(); + + let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(), Some(props)).unwrap(); + + for group_idx in 0..num_row_groups { + let mut id_values = Vec::with_capacity(rows_per_group); + let mut value_values = Vec::with_capacity(rows_per_group); + + for i in 0..rows_per_group { + id_values.push((group_idx * rows_per_group + i) as i64); + value_values.push((i * 2) as i64); + } + + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(arrow::array::Int64Array::from(id_values)), + Arc::new(arrow::array::Int64Array::from(value_values)), + ], + ) + .unwrap(); + + writer.write(&batch).unwrap(); + } + + writer.close().unwrap(); + Bytes::from(buffer) +} + +fn bench_single_row_group_read(c: &mut Criterion) { + let mut group = c.benchmark_group("row_group_index_single"); + group.sample_size(50); + + // test different row counts to see how the optimization scales + // these numbers were tuned to ensure ms-range measurements + for rows_per_group in [100_000, 500_000, 1_000_000, 5_000_000] { + let file = create_test_file(3, rows_per_group); + + let row_group_index_field = Arc::new( + Field::new("row_group_index", ArrowDataType::Int64, false) + .with_extension_type(RowGroupIndex), + ); + + group.bench_with_input( + BenchmarkId::new("single_row_group", rows_per_group), + &rows_per_group, + |b, &_rows| { + b.iter(|| { + let options = ArrowReaderOptions::new() + .with_virtual_columns(vec![row_group_index_field.clone()]) + .unwrap(); + + let reader = ParquetRecordBatchReaderBuilder::try_new_with_options( + file.clone(), + options, + ) + .unwrap() + .with_row_groups(vec![1]) + .build() + .unwrap(); + + for batch in reader { + let batch = batch.unwrap(); + let _ = batch.column_by_name("row_group_index").unwrap(); + } + }); + }, + ); + } + + group.finish(); +} + +fn bench_multiple_row_group_read(c: &mut Criterion) { + let mut group = c.benchmark_group("row_group_index_multiple"); + group.sample_size(50); + + for num_row_groups in [2, 5, 10, 20] { + let rows_per_group = 100_000; + let file = create_test_file(num_row_groups, rows_per_group); + + let row_group_index_field = Arc::new( + Field::new("row_group_index", ArrowDataType::Int64, false) + .with_extension_type(RowGroupIndex), + ); + + group.bench_with_input( + BenchmarkId::new("num_row_groups", num_row_groups), + &num_row_groups, + |b, &num_rg| { + b.iter(|| { + let options = ArrowReaderOptions::new() + .with_virtual_columns(vec![row_group_index_field.clone()]) + .unwrap(); + + let row_groups: Vec = (0..num_rg).collect(); + let reader = ParquetRecordBatchReaderBuilder::try_new_with_options( + file.clone(), + options, + ) + .unwrap() + .with_row_groups(row_groups) + .build() + .unwrap(); + + for batch in reader { + let batch = batch.unwrap(); + let _ = batch.column_by_name("row_group_index").unwrap(); + } + }); + }, + ); + } + + group.finish(); +} + +criterion_group!( + benches, + bench_single_row_group_read, + bench_multiple_row_group_read +); +criterion_main!(benches); diff --git a/vendor/parquet/benches/row_selection_cursor.rs b/vendor/parquet/benches/row_selection_cursor.rs new file mode 100644 index 0000000..a957f0b --- /dev/null +++ b/vendor/parquet/benches/row_selection_cursor.rs @@ -0,0 +1,541 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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; +use std::sync::Arc; + +use arrow_array::builder::StringViewBuilder; +use arrow_array::{ArrayRef, Float64Array, Int32Array, RecordBatch, StringViewArray}; +use arrow_buffer::BooleanBufferBuilder; +use arrow_schema::{DataType, Field, Schema}; +use bytes::Bytes; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::{ + ParquetRecordBatchReaderBuilder, RowSelection, RowSelectionPolicy, RowSelector, +}; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; + +const TOTAL_ROWS: usize = 1 << 20; +const BATCH_SIZE: usize = 1 << 10; +const BASE_SEED: u64 = 0xA55AA55A; +const AVG_SELECTOR_LENGTHS: &[usize] = &[4, 8, 12, 16, 20, 24, 28, 32, 36, 40]; +const COLUMN_WIDTHS: &[usize] = &[2, 4, 8, 16, 32]; +const UTF8VIEW_LENS: &[usize] = &[4, 8, 16, 32, 64, 128, 256]; +const BENCH_MODES: &[BenchMode] = &[BenchMode::Selector, BenchMode::Mask, BenchMode::Auto]; +const BACKINGS: &[Backing] = &[Backing::Selectors, Backing::Mask]; + +struct DataProfile { + name: &'static str, + build_batch: fn(usize) -> RecordBatch, +} + +const DATA_PROFILES: &[DataProfile] = &[ + DataProfile { + name: "int32", + build_batch: build_int32_batch, + }, + DataProfile { + name: "float64", + build_batch: build_float64_batch, + }, + DataProfile { + name: "utf8view", + build_batch: build_utf8view_batch, + }, +]; + +fn criterion_benchmark(c: &mut Criterion) { + let scenarios = [ + /* uniform50 (50% selected, constant run lengths, starts with skip) + ```text + ┌───────────────┐ + │ │ skip + │ │ + │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ select + │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + │ │ skip + │ │ + │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ select + │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + │ ... │ + └───────────────┘ + ``` */ + Scenario { + name: "uniform50", + select_ratio: 0.5, + start_with_select: false, + distribution: RunDistribution::Constant, + }, + /* spread50 (50% selected, large jitter in run lengths, starts with skip) + ```text + ┌───────────────┐ + │ │ skip (long) + │ │ + │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ select (short) + │ │ skip (short) + │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ select (long) + │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + │ │ skip (medium) + │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ select (medium) + │ ... │ + └───────────────┘ + ``` */ + Scenario { + name: "spread50", + select_ratio: 0.5, + start_with_select: false, + distribution: RunDistribution::Uniform { spread: 0.9 }, + }, + /* sparse20 (20% selected, bimodal: occasional long runs, starts with skip) + ```text + ┌───────────────┐ + │ │ skip (long) + │ │ + │ │ + │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ select (short) + │ │ skip (long) + │ │ + │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ select (occasional long) + │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + │ ... │ + └───────────────┘ + ``` */ + Scenario { + name: "sparse20", + select_ratio: 0.2, + start_with_select: false, + distribution: RunDistribution::Bimodal { + long_factor: 6.0, + long_prob: 0.1, + }, + }, + /* dense80 (80% selected, bimodal: occasional long runs, starts with select) + ```text + ┌───────────────┐ + │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ select (long) + │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + │ │ skip (short) + │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ select (long) + │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + │ │ skip (very short) + │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ select (long) + │▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒│ + │ ... │ + └───────────────┘ + ``` */ + Scenario { + name: "dense80", + select_ratio: 0.8, + start_with_select: true, + distribution: RunDistribution::Bimodal { + long_factor: 4.0, + long_prob: 0.05, + }, + }, + ]; + + let base_parquet = build_parquet_data(TOTAL_ROWS, build_int32_batch); + let base_scenario = &scenarios[0]; + + for (idx, scenario) in scenarios.iter().enumerate() { + // The first scenario is a special case for backwards compatibility with + // existing benchmark result formats. + let suite = if idx == 0 { "len" } else { "scenario" }; + bench_over_lengths( + c, + suite, + scenario.name, + &base_parquet, + scenario, + BASE_SEED ^ ((idx as u64) << 16), + ); + } + + for (profile_idx, profile) in DATA_PROFILES.iter().enumerate() { + let parquet_data = build_parquet_data(TOTAL_ROWS, profile.build_batch); + bench_over_lengths( + c, + "dtype", + profile.name, + &parquet_data, + base_scenario, + BASE_SEED ^ ((profile_idx as u64) << 24), + ); + } + + for (offset, &column_count) in COLUMN_WIDTHS.iter().enumerate() { + let parquet_data = write_parquet_batch(build_int32_columns_batch(TOTAL_ROWS, column_count)); + let variant_label = format!("C{:02}", column_count); + bench_over_lengths( + c, + "columns", + &variant_label, + &parquet_data, + base_scenario, + BASE_SEED ^ ((offset as u64) << 32), + ); + } + + for (offset, &len) in UTF8VIEW_LENS.iter().enumerate() { + let batch = build_utf8view_batch_with_len(TOTAL_ROWS, len); + let parquet_data = write_parquet_batch(batch); + let variant_label = format!("utf8view-L{:03}", len); + bench_over_lengths( + c, + "utf8view-len", + &variant_label, + &parquet_data, + base_scenario, + BASE_SEED ^ ((offset as u64) << 40), + ); + } +} + +fn bench_over_lengths( + c: &mut Criterion, + suite: &str, + variant: &str, + parquet_data: &Bytes, + scenario: &Scenario, + seed_base: u64, +) { + for (offset, &avg_len) in AVG_SELECTOR_LENGTHS.iter().enumerate() { + let selectors = + generate_selectors(avg_len, TOTAL_ROWS, scenario, seed_base + offset as u64); + let stats = SelectorStats::new(&selectors); + let selection = RowSelection::from(selectors); + let suffix = format!( + "{}-{}-{}-L{:02}-avg{:.1}-sel{:02}", + suite, + scenario.name, + variant, + avg_len, + stats.average_selector_len, + (stats.select_ratio * 100.0).round() as u32 + ); + + for &backing in BACKINGS { + let bench_input = BenchInput { + parquet_data: parquet_data.clone(), + selection: backing.convert(&selection), + }; + + for &mode in BENCH_MODES { + c.bench_with_input( + BenchmarkId::new(backing.function_name(mode), &suffix), + &bench_input, + |b, input| { + b.iter(|| { + let total = + run_read(&input.parquet_data, &input.selection, mode.policy()); + hint::black_box(total); + }); + }, + ); + } + } + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); + +struct BenchInput { + parquet_data: Bytes, + selection: RowSelection, +} + +fn run_read(parquet_data: &Bytes, selection: &RowSelection, policy: RowSelectionPolicy) -> usize { + let reader = ParquetRecordBatchReaderBuilder::try_new(parquet_data.clone()) + .unwrap() + .with_batch_size(BATCH_SIZE) + .with_row_selection(selection.clone()) + .with_row_selection_policy(policy) + .build() + .unwrap(); + + let mut total_rows = 0usize; + for batch in reader { + let batch = batch.unwrap(); + total_rows += batch.num_rows(); + } + total_rows +} + +fn build_parquet_data(total_rows: usize, build_batch: fn(usize) -> RecordBatch) -> Bytes { + let batch = build_batch(total_rows); + write_parquet_batch(batch) +} + +fn build_single_column_batch(data_type: DataType, array: ArrayRef) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("value", data_type, false)])); + RecordBatch::try_new(schema, vec![array]).unwrap() +} + +fn build_int32_batch(total_rows: usize) -> RecordBatch { + let values = Int32Array::from_iter_values((0..total_rows).map(|v| v as i32)); + build_single_column_batch(DataType::Int32, Arc::new(values) as ArrayRef) +} + +fn build_float64_batch(total_rows: usize) -> RecordBatch { + let values = Float64Array::from_iter_values((0..total_rows).map(|v| v as f64)); + build_single_column_batch(DataType::Float64, Arc::new(values) as ArrayRef) +} + +fn build_utf8view_batch(total_rows: usize) -> RecordBatch { + let mut builder = StringViewBuilder::new(); + // Mix short and long values. + for i in 0..total_rows { + match i % 5 { + 0 => builder.append_value("alpha"), + 1 => builder.append_value("beta"), + 2 => builder.append_value("gamma"), + 3 => builder.append_value("delta"), + _ => builder.append_value("a longer utf8 string payload to test view storage"), + } + } + let values: StringViewArray = builder.finish(); + build_single_column_batch(DataType::Utf8View, Arc::new(values) as ArrayRef) +} + +fn build_utf8view_batch_with_len(total_rows: usize, len: usize) -> RecordBatch { + let mut builder = StringViewBuilder::new(); + let value: String = "a".repeat(len); + for _ in 0..total_rows { + builder.append_value(&value); + } + let values: StringViewArray = builder.finish(); + build_single_column_batch(DataType::Utf8View, Arc::new(values) as ArrayRef) +} + +fn build_int32_columns_batch(total_rows: usize, num_columns: usize) -> RecordBatch { + let base_values: ArrayRef = Arc::new(Int32Array::from_iter_values( + (0..total_rows).map(|v| v as i32), + )); + let mut fields = Vec::with_capacity(num_columns); + let mut columns = Vec::with_capacity(num_columns); + for idx in 0..num_columns { + fields.push(Field::new(format!("value{}", idx), DataType::Int32, false)); + columns.push(base_values.clone()); + } + let schema = Arc::new(Schema::new(fields)); + RecordBatch::try_new(schema, columns).unwrap() +} + +fn write_parquet_batch(batch: RecordBatch) -> Bytes { + let schema = batch.schema(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema.clone(), None).unwrap(); + writer.write(&batch).unwrap(); + let buffer = writer.into_inner().unwrap(); + Bytes::from(buffer) +} + +#[derive(Clone)] +struct Scenario { + name: &'static str, + select_ratio: f64, + start_with_select: bool, + distribution: RunDistribution, +} + +#[derive(Clone)] +enum RunDistribution { + Constant, + Uniform { spread: f64 }, + Bimodal { long_factor: f64, long_prob: f64 }, +} + +fn generate_selectors( + avg_selector_len: usize, + total_rows: usize, + scenario: &Scenario, + seed: u64, +) -> Vec { + assert!( + (0.0..=1.0).contains(&scenario.select_ratio), + "select_ratio must be in [0, 1]" + ); + + let mut select_mean = scenario.select_ratio * 2.0 * avg_selector_len as f64; + let mut skip_mean = (1.0 - scenario.select_ratio) * 2.0 * avg_selector_len as f64; + + select_mean = select_mean.max(1.0); + skip_mean = skip_mean.max(1.0); + + let sum = select_mean + skip_mean; + // Rebalance the sampled select/skip run lengths so their sum matches the requested + // average selector length while respecting the configured selectivity ratio. + let scale = if sum == 0.0 { + 1.0 + } else { + (2.0 * avg_selector_len as f64) / sum + }; + select_mean *= scale; + skip_mean *= scale; + + let mut rng = StdRng::seed_from_u64(seed ^ (avg_selector_len as u64).wrapping_mul(0x9E3779B1)); + let mut selectors = Vec::with_capacity(total_rows / avg_selector_len.max(1)); + let mut remaining = total_rows; + let mut is_select = scenario.start_with_select; + + while remaining > 0 { + let mean = if is_select { select_mean } else { skip_mean }; + let len = sample_length(mean, &scenario.distribution, &mut rng).max(1); + let len = len.min(remaining); + selectors.push(if is_select { + RowSelector::select(len) + } else { + RowSelector::skip(len) + }); + remaining -= len; + if remaining == 0 { + break; + } + is_select = !is_select; + } + + let selection: RowSelection = selectors.into(); + selection.into() +} + +fn sample_length(mean: f64, distribution: &RunDistribution, rng: &mut StdRng) -> usize { + match distribution { + RunDistribution::Constant => mean.round().max(1.0) as usize, + RunDistribution::Uniform { spread } => { + let spread = spread.clamp(0.0, 0.99); + let lower = (mean * (1.0 - spread)).max(1.0); + let upper = (mean * (1.0 + spread)).max(lower + f64::EPSILON); + if (upper - lower) < 1.0 { + lower.round().max(1.0) as usize + } else { + let low = lower.floor() as usize; + let high = upper.ceil() as usize; + rng.random_range(low..=high).max(1) + } + } + RunDistribution::Bimodal { + long_factor, + long_prob, + } => { + let long_prob = long_prob.clamp(0.0, 0.5); + let short_prob = 1.0 - long_prob; + let short_factor = if short_prob == 0.0 { + 1.0 / long_factor.max(f64::EPSILON) + } else { + (1.0 - long_prob * long_factor).max(0.0) / short_prob + }; + let use_long = rng.random_bool(long_prob); + let factor = if use_long { + *long_factor + } else { + short_factor.max(0.1) + }; + (mean * factor).round().max(1.0) as usize + } + } +} + +#[derive(Clone, Copy)] +enum BenchMode { + Selector, + Mask, + Auto, +} + +impl BenchMode { + fn label(self) -> &'static str { + match self { + BenchMode::Selector => "read_selector", + BenchMode::Mask => "read_mask", + BenchMode::Auto => "read_auto", + } + } + + fn policy(self) -> RowSelectionPolicy { + match self { + BenchMode::Selector => RowSelectionPolicy::Selectors, + BenchMode::Mask => RowSelectionPolicy::Mask, + BenchMode::Auto => RowSelectionPolicy::default(), + } + } +} + +/// Which representation backs the input [`RowSelection`] handed to the reader. +#[derive(Clone, Copy)] +enum Backing { + Selectors, + Mask, +} + +impl Backing { + /// Benchmark function name for this backing/mode pair. Selector-backed + /// inputs keep the historical names so results stay comparable with runs + /// from before mask-backed selections existed. + fn function_name(self, mode: BenchMode) -> String { + match self { + Backing::Selectors => mode.label().to_string(), + Backing::Mask => format!("{}_mask_backed", mode.label()), + } + } + + fn convert(self, selection: &RowSelection) -> RowSelection { + match self { + Backing::Selectors => selection.clone(), + Backing::Mask => { + let mut builder = BooleanBufferBuilder::new(TOTAL_ROWS); + for selector in selection.iter() { + builder.append_n(selector.row_count, !selector.skip); + } + RowSelection::from_boolean_buffer(builder.finish()) + } + } + } +} + +struct SelectorStats { + average_selector_len: f64, + select_ratio: f64, +} + +impl SelectorStats { + fn new(selectors: &[RowSelector]) -> Self { + if selectors.is_empty() { + return Self { + average_selector_len: 0.0, + select_ratio: 0.0, + }; + } + + let total_rows: usize = selectors.iter().map(|s| s.row_count).sum(); + let selected_rows: usize = selectors + .iter() + .filter(|s| !s.skip) + .map(|s| s.row_count) + .sum(); + + Self { + average_selector_len: total_rows as f64 / selectors.len() as f64, + select_ratio: if total_rows == 0 { + 0.0 + } else { + selected_rows as f64 / total_rows as f64 + }, + } + } +} diff --git a/vendor/parquet/benches/row_selection_policy_common/assertions.rs b/vendor/parquet/benches/row_selection_policy_common/assertions.rs new file mode 100644 index 0000000..057ce35 --- /dev/null +++ b/vendor/parquet/benches/row_selection_policy_common/assertions.rs @@ -0,0 +1,65 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use parquet::arrow::arrow_reader::RowSelectionPolicy; + +use super::fixture::CaseFixture; +use super::model::{CaseSpec, PAYLOAD_VALUE_MODULUS, ROWS_PER_GROUP}; +use super::runner::run_collect_payload0; +use super::shapes::expand_pattern; + +pub(crate) async fn preflight_auto(case: &CaseSpec, fixture: &CaseFixture) { + let actual = run_collect_payload0(fixture, RowSelectionPolicy::default()).await; + assert_eq!( + actual.row_count, fixture.expected_rows, + "{} returned an unexpected number of rows", + case.name + ); + + let expected = expected_selected_global_rows(case); + assert_eq!(actual.payload0.len(), expected.len(), "{}", case.name); + if let Some((output_row, (actual, expected))) = actual + .payload0 + .iter() + .zip(&expected) + .enumerate() + .find(|(_, (actual, expected))| actual != expected) + { + panic!( + "{} returned the wrong source row at output {output_row}: expected {expected}, got {actual}", + case.name + ); + } +} + +fn expected_selected_global_rows(case: &CaseSpec) -> Vec { + case.row_groups + .iter() + .copied() + .enumerate() + .flat_map(|(row_group_idx, pattern)| { + expand_pattern(pattern, ROWS_PER_GROUP) + .into_iter() + .enumerate() + .filter(|(_, selected)| *selected == 1) + .map(move |(row_idx, _)| { + let global_row = row_group_idx * ROWS_PER_GROUP + row_idx; + global_row.wrapping_rem(PAYLOAD_VALUE_MODULUS) as i32 + }) + }) + .collect() +} diff --git a/vendor/parquet/benches/row_selection_policy_common/cases.rs b/vendor/parquet/benches/row_selection_policy_common/cases.rs new file mode 100644 index 0000000..99e9068 --- /dev/null +++ b/vendor/parquet/benches/row_selection_policy_common/cases.rs @@ -0,0 +1,121 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use super::model::{CaseSpec, RowGroupPattern}; +use super::shapes::{ + BURSTY_50_SAME_SUMMARY, CLUSTERED_50_RUN128, DENSE_98_44_SKIP1_SELECT63, FRAGMENTED_50_RUN1, + MODERATE_12_5_RUN32, REGULAR_50_RUN32, SPARSE_1_56_RUN32, +}; + +const FOUR_SPARSE: &[RowGroupPattern] = &[RowGroupPattern::Cycle(SPARSE_1_56_RUN32); 4]; + +const FOUR_MODERATE: &[RowGroupPattern] = &[RowGroupPattern::Cycle(MODERATE_12_5_RUN32); 4]; + +const FOUR_FRAGMENTED: &[RowGroupPattern] = &[RowGroupPattern::Cycle(FRAGMENTED_50_RUN1); 4]; + +const FOUR_CLUSTERED: &[RowGroupPattern] = &[RowGroupPattern::Cycle(CLUSTERED_50_RUN128); 4]; + +const FOUR_REGULAR: &[RowGroupPattern] = &[RowGroupPattern::Cycle(REGULAR_50_RUN32); 4]; + +const FOUR_BURSTY: &[RowGroupPattern] = &[RowGroupPattern::Cycle(BURSTY_50_SAME_SUMMARY); 4]; + +const FOUR_DENSE: &[RowGroupPattern] = &[RowGroupPattern::Cycle(DENSE_98_44_SKIP1_SELECT63); 4]; + +const FOUR_ALL_SELECTED: &[RowGroupPattern] = &[RowGroupPattern::AllSelected; 4]; + +const SPARSE_TO_FRAGMENTED: &[RowGroupPattern] = &[ + RowGroupPattern::Cycle(SPARSE_1_56_RUN32), + RowGroupPattern::Cycle(SPARSE_1_56_RUN32), + RowGroupPattern::Cycle(FRAGMENTED_50_RUN1), + RowGroupPattern::Cycle(FRAGMENTED_50_RUN1), +]; + +const FRAGMENTED_TO_SPARSE: &[RowGroupPattern] = &[ + RowGroupPattern::Cycle(FRAGMENTED_50_RUN1), + RowGroupPattern::Cycle(FRAGMENTED_50_RUN1), + RowGroupPattern::Cycle(SPARSE_1_56_RUN32), + RowGroupPattern::Cycle(SPARSE_1_56_RUN32), +]; + +const ONE_FRAGMENTED: &[RowGroupPattern] = &[RowGroupPattern::Cycle(FRAGMENTED_50_RUN1)]; + +const EIGHT_FRAGMENTED: &[RowGroupPattern] = &[RowGroupPattern::Cycle(FRAGMENTED_50_RUN1); 8]; + +pub(crate) fn assert_order_cases_are_reverses() { + assert!( + SPARSE_TO_FRAGMENTED + .iter() + .eq(FRAGMENTED_TO_SPARSE.iter().rev()) + ); +} + +pub(crate) const SHAPE_CASES: &[CaseSpec] = &[ + CaseSpec { + name: "sparse_1_56_run32", + row_groups: FOUR_SPARSE, + }, + CaseSpec { + name: "moderate_12_5_run32", + row_groups: FOUR_MODERATE, + }, + CaseSpec { + name: "fragmented_50_run1", + row_groups: FOUR_FRAGMENTED, + }, + CaseSpec { + name: "clustered_50_run128", + row_groups: FOUR_CLUSTERED, + }, + CaseSpec { + name: "regular_50_run32", + row_groups: FOUR_REGULAR, + }, + CaseSpec { + name: "bursty_50_same_summary", + row_groups: FOUR_BURSTY, + }, + CaseSpec { + name: "dense_98_44_skip1_select63", + row_groups: FOUR_DENSE, + }, + CaseSpec { + name: "all_selected", + row_groups: FOUR_ALL_SELECTED, + }, +]; + +pub(crate) const ORDER_CASES: &[CaseSpec] = &[ + CaseSpec { + name: "sparse2_then_fragmented2", + row_groups: SPARSE_TO_FRAGMENTED, + }, + CaseSpec { + name: "fragmented2_then_sparse2", + row_groups: FRAGMENTED_TO_SPARSE, + }, +]; + +pub(crate) const SCALE_CASES: &[CaseSpec] = &[ + CaseSpec { + name: "fragmented_50_run1/rg01", + row_groups: ONE_FRAGMENTED, + }, + CaseSpec { + name: "fragmented_50_run1/rg08", + row_groups: EIGHT_FRAGMENTED, + }, +]; diff --git a/vendor/parquet/benches/row_selection_policy_common/fixture.rs b/vendor/parquet/benches/row_selection_policy_common/fixture.rs new file mode 100644 index 0000000..1470f5b --- /dev/null +++ b/vendor/parquet/benches/row_selection_policy_common/fixture.rs @@ -0,0 +1,163 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::ops::Range; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int32Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use bytes::Bytes; +use futures::FutureExt; +use futures::future::BoxFuture; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::ArrowReaderOptions; +use parquet::arrow::async_reader::AsyncFileReader; +use parquet::basic::Compression; +use parquet::errors::Result; +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader}; +use parquet::file::properties::WriterProperties; + +use super::model::{CaseSpec, PAYLOAD_COLUMNS, PAYLOAD_VALUE_MODULUS, ROWS_PER_GROUP}; +use super::shapes::{expand_pattern, selected_rows}; + +#[derive(Debug)] +pub(crate) struct CaseFixture { + bytes: Bytes, + metadata: Arc, + pub(crate) expected_rows: usize, +} + +impl CaseFixture { + pub(crate) fn reader(&self) -> InMemoryAsyncReader { + InMemoryAsyncReader { + bytes: self.bytes.clone(), + metadata: Arc::clone(&self.metadata), + } + } + + pub(crate) fn schema_descr(&self) -> &parquet::schema::types::SchemaDescriptor { + self.metadata.file_metadata().schema_descr() + } +} + +#[derive(Debug, Clone)] +pub(crate) struct InMemoryAsyncReader { + bytes: Bytes, + metadata: Arc, +} + +impl AsyncFileReader for InMemoryAsyncReader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + let bytes = self.bytes.slice(range.start as usize..range.end as usize); + async move { Ok(bytes) }.boxed() + } + + fn get_metadata<'a>( + &'a mut self, + _options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, Result>> { + let metadata = Arc::clone(&self.metadata); + async move { Ok(metadata) }.boxed() + } +} + +pub(crate) fn build_fixture(case: &CaseSpec) -> Result { + assert!( + !case.row_groups.is_empty(), + "benchmark case must contain at least one row group" + ); + + let schema = build_schema(); + let properties = WriterProperties::builder() + .set_compression(Compression::UNCOMPRESSED) + .set_dictionary_enabled(false) + .set_max_row_group_row_count(Some(ROWS_PER_GROUP)) + .build(); + + let mut encoded = Vec::new(); + { + let mut writer = ArrowWriter::try_new(&mut encoded, Arc::clone(&schema), Some(properties))?; + for (row_group_idx, pattern) in case.row_groups.iter().copied().enumerate() { + writer.write(&build_row_group_batch( + Arc::clone(&schema), + pattern, + row_group_idx, + )?)?; + } + writer.close()?; + } + + let bytes = Bytes::from(encoded); + let mut metadata_reader = + ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Skip); + metadata_reader.try_parse(&bytes)?; + let metadata = Arc::new(metadata_reader.finish()?); + + assert_eq!( + metadata.num_row_groups(), + case.row_groups.len(), + "writer did not preserve the requested row-group layout" + ); + for row_group in metadata.row_groups() { + assert_eq!(row_group.num_rows() as usize, ROWS_PER_GROUP); + } + + let expected_rows = case + .row_groups + .iter() + .copied() + .map(|pattern| selected_rows(pattern, ROWS_PER_GROUP)) + .sum(); + + Ok(CaseFixture { + bytes, + metadata, + expected_rows, + }) +} + +fn build_schema() -> SchemaRef { + let mut fields = Vec::with_capacity(PAYLOAD_COLUMNS + 1); + fields.push(Field::new("predicate", DataType::Int32, false)); + fields.extend( + (0..PAYLOAD_COLUMNS) + .map(|column_idx| Field::new(format!("payload_{column_idx}"), DataType::Int32, false)), + ); + Arc::new(Schema::new(fields)) +} + +fn build_row_group_batch( + schema: SchemaRef, + pattern: super::model::RowGroupPattern, + row_group_idx: usize, +) -> Result { + let predicate = expand_pattern(pattern, ROWS_PER_GROUP); + let mut columns = Vec::with_capacity(PAYLOAD_COLUMNS + 1); + columns.push(Arc::new(Int32Array::from(predicate)) as ArrayRef); + + for column_idx in 0..PAYLOAD_COLUMNS { + let values = Int32Array::from_iter_values((0..ROWS_PER_GROUP).map(|row_idx| { + let global_row = row_group_idx * ROWS_PER_GROUP + row_idx; + global_row + .wrapping_add(column_idx * 17) + .wrapping_rem(PAYLOAD_VALUE_MODULUS) as i32 + })); + columns.push(Arc::new(values) as ArrayRef); + } + + Ok(RecordBatch::try_new(schema, columns)?) +} diff --git a/vendor/parquet/benches/row_selection_policy_common/mod.rs b/vendor/parquet/benches/row_selection_policy_common/mod.rs new file mode 100644 index 0000000..a4053e8 --- /dev/null +++ b/vendor/parquet/benches/row_selection_policy_common/mod.rs @@ -0,0 +1,24 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +pub(crate) mod assertions; +pub(crate) mod cases; +pub(crate) mod fixture; +pub(crate) mod model; +pub(crate) mod register; +pub(crate) mod runner; +pub(crate) mod shapes; diff --git a/vendor/parquet/benches/row_selection_policy_common/model.rs b/vendor/parquet/benches/row_selection_policy_common/model.rs new file mode 100644 index 0000000..e44462e --- /dev/null +++ b/vendor/parquet/benches/row_selection_policy_common/model.rs @@ -0,0 +1,61 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +pub(crate) const ROWS_PER_GROUP: usize = 65_536; +pub(crate) const BATCH_SIZE: usize = 8_192; +pub(crate) const PAYLOAD_COLUMNS: usize = 8; +pub(crate) const PAYLOAD_VALUE_MODULUS: usize = 1_000_003; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct SelectionRun { + pub(crate) selected: bool, + pub(crate) len: usize, +} + +impl SelectionRun { + pub(crate) const fn select(len: usize) -> Self { + Self { + selected: true, + len, + } + } + + pub(crate) const fn skip(len: usize) -> Self { + Self { + selected: false, + len, + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum RowGroupPattern { + Cycle(&'static [SelectionRun]), + AllSelected, +} + +#[derive(Clone, Copy, Debug)] +pub(crate) struct CaseSpec { + pub(crate) name: &'static str, + pub(crate) row_groups: &'static [RowGroupPattern], +} + +impl CaseSpec { + pub(crate) const fn total_rows(self) -> usize { + self.row_groups.len() * ROWS_PER_GROUP + } +} diff --git a/vendor/parquet/benches/row_selection_policy_common/register.rs b/vendor/parquet/benches/row_selection_policy_common/register.rs new file mode 100644 index 0000000..7c2e8dc --- /dev/null +++ b/vendor/parquet/benches/row_selection_policy_common/register.rs @@ -0,0 +1,59 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::hint; +use std::sync::{Arc, OnceLock}; + +use criterion::{Criterion, Throughput}; + +use super::assertions::preflight_auto; +use super::fixture::{CaseFixture, build_fixture}; +use super::model::CaseSpec; +use super::runner::run_auto; + +pub(crate) fn register_auto_group(c: &mut Criterion, group_name: &str, cases: &'static [CaseSpec]) { + let runtime = Arc::new( + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(), + ); + let mut group = c.benchmark_group(group_name); + + for case in cases { + group.throughput(Throughput::Elements(case.total_rows() as u64)); + + let fixture: Arc>> = Arc::new(OnceLock::new()); + let preflight = Arc::new(OnceLock::new()); + let runtime = Arc::clone(&runtime); + + group.bench_function(case.name, move |b| { + let fixture = + Arc::clone(fixture.get_or_init(|| Arc::new(build_fixture(case).unwrap()))); + preflight.get_or_init(|| { + runtime.block_on(preflight_auto(case, &fixture)); + }); + + b.iter(|| { + let rows = runtime.block_on(run_auto(&fixture)); + hint::black_box(rows); + }); + }); + } + + group.finish(); +} diff --git a/vendor/parquet/benches/row_selection_policy_common/runner.rs b/vendor/parquet/benches/row_selection_policy_common/runner.rs new file mode 100644 index 0000000..bf89fd8 --- /dev/null +++ b/vendor/parquet/benches/row_selection_policy_common/runner.rs @@ -0,0 +1,92 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::{Int32Array, RecordBatch}; +use arrow::compute::kernels::cmp::eq; +use arrow::datatypes::Int32Type; +use arrow_array::cast::AsArray; +use futures::StreamExt; +use parquet::arrow::arrow_reader::{ArrowPredicateFn, RowFilter, RowSelectionPolicy}; +use parquet::arrow::{ParquetRecordBatchStreamBuilder, ProjectionMask}; + +use super::fixture::CaseFixture; +use super::model::{BATCH_SIZE, PAYLOAD_COLUMNS}; + +#[derive(Debug, Eq, PartialEq)] +pub(crate) struct RunResult { + pub(crate) row_count: usize, + pub(crate) payload0: Vec, +} + +async fn run_with_consumer( + fixture: &CaseFixture, + policy: RowSelectionPolicy, + mut consume: F, +) -> usize +where + F: FnMut(&RecordBatch), +{ + let predicate_projection = ProjectionMask::roots(fixture.schema_descr(), [0]); + let output_projection = ProjectionMask::roots(fixture.schema_descr(), 1..=PAYLOAD_COLUMNS); + let predicate = ArrowPredicateFn::new(predicate_projection, |batch: RecordBatch| { + eq(batch.column(0), &Int32Array::new_scalar(1)) + }); + let row_filter = RowFilter::new(vec![Box::new(predicate)]); + + let mut stream = ParquetRecordBatchStreamBuilder::new(fixture.reader()) + .await + .unwrap() + .with_batch_size(BATCH_SIZE) + .with_projection(output_projection) + .with_row_filter(row_filter) + .with_row_selection_policy(policy) + .build() + .unwrap(); + + let mut rows = 0; + while let Some(batch) = stream.next().await { + let batch = batch.unwrap(); + rows += batch.num_rows(); + consume(&batch); + } + rows +} + +pub(crate) async fn run(fixture: &CaseFixture, policy: RowSelectionPolicy) -> usize { + run_with_consumer(fixture, policy, |_| {}).await +} + +pub(crate) async fn run_auto(fixture: &CaseFixture) -> usize { + run(fixture, RowSelectionPolicy::default()).await +} + +pub(crate) async fn run_collect_payload0( + fixture: &CaseFixture, + policy: RowSelectionPolicy, +) -> RunResult { + let mut payload0 = Vec::with_capacity(fixture.expected_rows); + let row_count = run_with_consumer(fixture, policy, |batch| { + let values = batch.column(0).as_primitive::(); + payload0.extend(values.values().iter().copied()); + }) + .await; + + RunResult { + row_count, + payload0, + } +} diff --git a/vendor/parquet/benches/row_selection_policy_common/shapes.rs b/vendor/parquet/benches/row_selection_policy_common/shapes.rs new file mode 100644 index 0000000..e83667d --- /dev/null +++ b/vendor/parquet/benches/row_selection_policy_common/shapes.rs @@ -0,0 +1,150 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use super::model::{RowGroupPattern, SelectionRun}; + +pub(crate) const SPARSE_1_56_RUN32: &[SelectionRun] = + &[SelectionRun::skip(2_016), SelectionRun::select(32)]; + +pub(crate) const MODERATE_12_5_RUN32: &[SelectionRun] = + &[SelectionRun::skip(224), SelectionRun::select(32)]; + +pub(crate) const FRAGMENTED_50_RUN1: &[SelectionRun] = + &[SelectionRun::skip(1), SelectionRun::select(1)]; + +pub(crate) const CLUSTERED_50_RUN128: &[SelectionRun] = + &[SelectionRun::skip(128), SelectionRun::select(128)]; + +pub(crate) const REGULAR_50_RUN32: &[SelectionRun] = + &[SelectionRun::skip(32), SelectionRun::select(32)]; + +pub(crate) const DENSE_98_44_SKIP1_SELECT63: &[SelectionRun] = + &[SelectionRun::skip(1), SelectionRun::select(63)]; + +/// Has the same selectivity, run density, and mean selected/skipped run lengths +/// as [`REGULAR_50_RUN32`], but a different run variance and ordering. +pub(crate) const BURSTY_50_SAME_SUMMARY: &[SelectionRun] = &[ + SelectionRun::skip(1), + SelectionRun::select(1), + SelectionRun::skip(1), + SelectionRun::select(1), + SelectionRun::skip(1), + SelectionRun::select(1), + SelectionRun::skip(125), + SelectionRun::select(125), +]; + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +struct ShapeSummary { + selected_rows: usize, + skipped_rows: usize, + selector_count: usize, + selected_run_count: usize, + skipped_run_count: usize, +} + +impl ShapeSummary { + fn from_cycle(cycle: &[SelectionRun]) -> Self { + cycle.iter().fold(Self::default(), |mut summary, run| { + assert!(run.len > 0, "selection run must be non-empty"); + summary.selector_count += 1; + if run.selected { + summary.selected_rows += run.len; + summary.selected_run_count += 1; + } else { + summary.skipped_rows += run.len; + summary.skipped_run_count += 1; + } + summary + }) + } + + fn total_rows(self) -> usize { + self.selected_rows + self.skipped_rows + } + + fn selected_ratio(self) -> f64 { + self.selected_rows as f64 / self.total_rows() as f64 + } + + fn run_density(self) -> f64 { + self.selector_count as f64 / self.total_rows() as f64 + } + + fn average_selected_run(self) -> f64 { + self.selected_rows as f64 / self.selected_run_count as f64 + } + + fn average_skipped_run(self) -> f64 { + self.skipped_rows as f64 / self.skipped_run_count as f64 + } +} + +pub(crate) fn assert_shape_contracts() { + let regular = ShapeSummary::from_cycle(REGULAR_50_RUN32); + let bursty = ShapeSummary::from_cycle(BURSTY_50_SAME_SUMMARY); + + assert_eq!(regular.selected_ratio(), bursty.selected_ratio()); + assert_eq!(regular.run_density(), bursty.run_density()); + assert_eq!( + regular.average_selected_run(), + bursty.average_selected_run() + ); + assert_eq!(regular.average_skipped_run(), bursty.average_skipped_run()); + + let dense = ShapeSummary::from_cycle(DENSE_98_44_SKIP1_SELECT63); + assert_eq!(dense.selected_rows, 63); + assert_eq!(dense.skipped_rows, 1); + assert_eq!(dense.selector_count, 2); + assert_eq!(dense.average_selected_run(), 63.0); + assert_eq!(dense.average_skipped_run(), 1.0); +} + +pub(crate) fn expand_pattern(pattern: RowGroupPattern, row_count: usize) -> Vec { + match pattern { + RowGroupPattern::AllSelected => vec![1; row_count], + RowGroupPattern::Cycle(cycle) => { + assert!(!cycle.is_empty(), "selection cycle must not be empty"); + let summary = ShapeSummary::from_cycle(cycle); + assert_eq!( + row_count % summary.total_rows(), + 0, + "row group size must be divisible by cycle size" + ); + + let mut values = Vec::with_capacity(row_count); + while values.len() < row_count { + for run in cycle { + values.extend(std::iter::repeat_n(i32::from(run.selected), run.len)); + } + } + assert_eq!(values.len(), row_count); + values + } + } +} + +pub(crate) fn selected_rows(pattern: RowGroupPattern, row_count: usize) -> usize { + match pattern { + RowGroupPattern::AllSelected => row_count, + RowGroupPattern::Cycle(cycle) => { + let summary = ShapeSummary::from_cycle(cycle); + assert_eq!(row_count % summary.total_rows(), 0); + summary.selected_rows * (row_count / summary.total_rows()) + } + } +} diff --git a/vendor/parquet/benches/row_selector.rs b/vendor/parquet/benches/row_selector.rs new file mode 100644 index 0000000..38fb712 --- /dev/null +++ b/vendor/parquet/benches/row_selector.rs @@ -0,0 +1,88 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_array::BooleanArray; +use criterion::*; +use parquet::arrow::arrow_reader::RowSelection; +use rand::Rng; +use std::hint; + +/// Generates a random RowSelection with a specified selection ratio. +/// +/// # Arguments +/// +/// * `total_rows` - The total number of rows in the selection. +/// * `selection_ratio` - The ratio of rows to select (e.g., 1/3 for ~33% selection). +/// +/// # Returns +/// +/// * A `BooleanArray` instance with randomly selected rows based on the provided ratio. +fn generate_random_row_selection(total_rows: usize, selection_ratio: f64) -> BooleanArray { + let mut rng = rand::rng(); + let bools: Vec = (0..total_rows) + .map(|_| rng.random_bool(selection_ratio)) + .collect(); + BooleanArray::from(bools) +} + +fn criterion_benchmark(c: &mut Criterion) { + let total_rows = 300_000; + let selection_ratio = 1.0 / 3.0; + + // Generate two random RowSelections with approximately 1/3 of the rows selected. + let row_selection_a = + RowSelection::from_filters(&[generate_random_row_selection(total_rows, selection_ratio)]); + let row_selection_b = + RowSelection::from_filters(&[generate_random_row_selection(total_rows, selection_ratio)]); + + // Benchmark the intersection of the two RowSelections. + c.bench_function("intersection", |b| { + b.iter(|| { + let intersection = row_selection_a.intersection(&row_selection_b); + hint::black_box(intersection); + }) + }); + + c.bench_function("union", |b| { + b.iter(|| { + let union = row_selection_a.union(&row_selection_b); + hint::black_box(union); + }) + }); + + c.bench_function("from_filters", |b| { + let boolean_array = generate_random_row_selection(total_rows, selection_ratio); + b.iter(|| { + let array = boolean_array.clone(); + let selection = RowSelection::from_filters(&[array]); + hint::black_box(selection); + }) + }); + + c.bench_function("and_then", |b| { + let selected = row_selection_a.row_count(); + let sub_selection = + RowSelection::from_filters(&[generate_random_row_selection(selected, selection_ratio)]); + b.iter(|| { + let result = row_selection_a.and_then(&sub_selection); + hint::black_box(result); + }) + }); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/vendor/parquet/benches/row_selector_boolean_buffer.rs b/vendor/parquet/benches/row_selector_boolean_buffer.rs new file mode 100644 index 0000000..c5e1177 --- /dev/null +++ b/vendor/parquet/benches/row_selector_boolean_buffer.rs @@ -0,0 +1,233 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_array::{ArrayRef, BooleanArray, Int32Array, RecordBatch}; +use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder}; +use arrow_schema::{DataType, Field, Schema}; +use bytes::Bytes; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::{ParquetRecordBatchReaderBuilder, RowSelection}; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; +use std::hint; +use std::sync::Arc; + +const TOTAL_ROWS: usize = 3_000_000; +const SELECTIVITY_CASES: &[Selectivity] = &[ + Selectivity::new("select01", 1, 100), + Selectivity::new("select10", 1, 10), + Selectivity::new("select33", 1, 3), + Selectivity::new("select80", 4, 5), +]; + +/// Generates a deterministic random row mask with the specified selectivity. +fn generate_random_row_selection(total_rows: usize, selectivity: Selectivity) -> BooleanBuffer { + let mut rng = StdRng::seed_from_u64(0x5E1EC7_u64 ^ selectivity.seed()); + let bools: Vec = (0..total_rows) + .map(|_| rng.random_bool(selectivity.ratio())) + .collect(); + BooleanBuffer::from(bools) +} + +fn generate_fragmented_selection(total_rows: usize, selectivity: Selectivity) -> BooleanBuffer { + let mut builder = BooleanBufferBuilder::new(total_rows); + for row in 0..total_rows { + builder.append(row % selectivity.denominator < selectivity.numerator); + } + builder.finish() +} + +fn generate_clustered_selection(total_rows: usize, selectivity: Selectivity) -> BooleanBuffer { + const RUN: usize = 8 * 1024; + const JITTER: usize = RUN / 2; + let mut rng = StdRng::seed_from_u64(0xC1057E_u64 ^ selectivity.seed()); + let mut builder = BooleanBufferBuilder::new(total_rows); + let mut rows_remaining = total_rows; + let mut run_idx = 0usize; + + while rows_remaining != 0 { + let run_len = rng + .random_range((RUN - JITTER)..=(RUN + JITTER)) + .min(rows_remaining); + let selected = run_idx % selectivity.denominator < selectivity.numerator; + builder.append_n(run_len, selected); + rows_remaining -= run_len; + run_idx += 1; + } + builder.finish() +} + +fn boolean_array(mask: &BooleanBuffer) -> BooleanArray { + BooleanArray::new(mask.clone(), None) +} + +fn criterion_benchmark(c: &mut Criterion) { + let patterns = build_patterns(); + + let mut construction = c.benchmark_group("row_selector_boolean_buffer/construction"); + for case in &patterns { + construction.bench_with_input( + BenchmarkId::new("from_filters", case.label()), + &case.mask, + |b, mask| { + b.iter(|| { + let array = boolean_array(mask); + let selection = RowSelection::from_filters(&[array]); + hint::black_box(selection); + }) + }, + ); + + construction.bench_with_input( + BenchmarkId::new("from_boolean_buffer", case.label()), + &case.mask, + |b, mask| { + b.iter(|| { + let selection = RowSelection::from_boolean_buffer(mask.clone()); + hint::black_box(selection); + }) + }, + ); + } + construction.finish(); + + let parquet_data = write_parquet_file(TOTAL_ROWS); + let mut reader = c.benchmark_group("row_selector_boolean_buffer/reader"); + for case in &patterns { + reader.bench_with_input( + BenchmarkId::new("from_filters", case.label()), + &case.mask, + |b, mask| { + b.iter(|| { + let selection = RowSelection::from_filters(&[boolean_array(mask)]); + let rows = read_rows(&parquet_data, selection); + hint::black_box(rows); + }) + }, + ); + + reader.bench_with_input( + BenchmarkId::new("from_boolean_buffer", case.label()), + &case.mask, + |b, mask| { + b.iter(|| { + let selection = RowSelection::from_boolean_buffer(mask.clone()); + let rows = read_rows(&parquet_data, selection); + hint::black_box(rows); + }) + }, + ); + } + reader.finish(); +} + +fn build_patterns() -> Vec { + let mut patterns = Vec::new(); + for selectivity in SELECTIVITY_CASES { + patterns.push(BenchCase::new( + "fragmented", + *selectivity, + generate_fragmented_selection(TOTAL_ROWS, *selectivity), + )); + patterns.push(BenchCase::new( + "clustered", + *selectivity, + generate_clustered_selection(TOTAL_ROWS, *selectivity), + )); + patterns.push(BenchCase::new( + "random", + *selectivity, + generate_random_row_selection(TOTAL_ROWS, *selectivity), + )); + } + patterns +} + +fn write_parquet_file(total_rows: usize) -> Bytes { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + let values: ArrayRef = Arc::new(Int32Array::from_iter_values( + (0..total_rows).map(|row| row as i32), + )); + let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap(); + + let mut writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap(); + writer.write(&batch).unwrap(); + Bytes::from(writer.into_inner().unwrap()) +} + +fn read_rows(parquet_data: &Bytes, selection: RowSelection) -> usize { + let reader = ParquetRecordBatchReaderBuilder::try_new(parquet_data.clone()) + .unwrap() + .with_row_selection(selection) + .build() + .unwrap(); + + reader.map(|batch| batch.unwrap().num_rows()).sum::() +} + +struct BenchCase { + pattern: &'static str, + selectivity: Selectivity, + mask: BooleanBuffer, +} + +impl BenchCase { + fn new(pattern: &'static str, selectivity: Selectivity, mask: BooleanBuffer) -> Self { + Self { + pattern, + selectivity, + mask, + } + } + + fn label(&self) -> String { + format!("{}/{}", self.pattern, self.selectivity.label) + } +} + +#[derive(Clone, Copy)] +struct Selectivity { + label: &'static str, + numerator: usize, + denominator: usize, +} + +impl Selectivity { + const fn new(label: &'static str, numerator: usize, denominator: usize) -> Self { + Self { + label, + numerator, + denominator, + } + } + + fn ratio(self) -> f64 { + self.numerator as f64 / self.denominator as f64 + } + + fn seed(self) -> u64 { + ((self.numerator as u64) << 32) ^ self.denominator as u64 + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/vendor/parquet/benches/writer_overhead.rs b/vendor/parquet/benches/writer_overhead.rs new file mode 100644 index 0000000..fc4f616 --- /dev/null +++ b/vendor/parquet/benches/writer_overhead.rs @@ -0,0 +1,86 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks for writer per-column overhead at high column cardinality. +//! +//! These benchmarks measure the structural cost of creating, writing, and +//! closing a parquet file with many columns while keeping actual data +//! encoding negligible (1 row per column). This isolates overhead such as +//! `WriterProperties` per-column lookups, `GenericColumnWriter` allocation, +//! and metadata assembly. + +use criterion::{Criterion, criterion_group, criterion_main}; +use std::hint::black_box; +use std::io::Empty; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow_array::{Float32Array, RecordBatch}; +use parquet::arrow::ArrowWriter; +use parquet::basic::Compression; +use parquet::file::properties::WriterProperties; +use parquet::schema::types::ColumnPath; + +const COLUMN_COUNTS: &[usize] = &[1_000, 5_000, 10_000]; + +fn make_wide_schema(num_columns: usize) -> SchemaRef { + let fields: Vec = (0..num_columns) + .map(|i| Field::new(format!("c{i}"), DataType::Float32, false)) + .collect(); + Arc::new(Schema::new(fields)) +} + +fn make_single_row_batch(schema: &SchemaRef) -> RecordBatch { + let columns: Vec> = (0..schema.fields().len()) + .map(|_| Arc::new(Float32Array::from(vec![0.0f32])) as _) + .collect(); + RecordBatch::try_new(schema.clone(), columns).unwrap() +} + +/// Build WriterProperties with a per-column property set for every column, +/// populating the internal HashMap so that per-column lookups are exercised. +fn make_per_column_props(schema: &SchemaRef) -> WriterProperties { + let mut builder = WriterProperties::builder().set_dictionary_enabled(false); + for field in schema.fields() { + builder = builder.set_column_compression( + ColumnPath::from(field.name().as_str()), + Compression::UNCOMPRESSED, + ); + } + builder.build() +} + +fn bench_writer_overhead(c: &mut Criterion) { + for &num_cols in COLUMN_COUNTS { + let schema = make_wide_schema(num_cols); + let batch = make_single_row_batch(&schema); + let props = make_per_column_props(&schema); + + c.bench_function(&format!("writer_overhead/{num_cols}_cols"), |b| { + b.iter(|| { + let mut writer = + ArrowWriter::try_new(Empty::default(), schema.clone(), Some(props.clone())) + .unwrap(); + writer.write(black_box(&batch)).unwrap(); + black_box(writer.close()).unwrap(); + }); + }); + } +} + +criterion_group!(benches, bench_writer_overhead); +criterion_main!(benches); diff --git a/vendor/parquet/examples/external_metadata.rs b/vendor/parquet/examples/external_metadata.rs new file mode 100644 index 0000000..eeb1d90 --- /dev/null +++ b/vendor/parquet/examples/external_metadata.rs @@ -0,0 +1,238 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray}; +use arrow_cast::pretty::pretty_format_batches; +use futures::TryStreamExt; +use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; +use parquet::arrow::{ArrowWriter, ParquetRecordBatchStreamBuilder}; +use parquet::file::metadata::{ + PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader, ParquetMetaDataWriter, +}; +use parquet::file::properties::{EnabledStatistics, WriterProperties}; +use std::fs::File; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use tempfile::TempDir; + +/// This example demonstrates advanced usage of the Parquet metadata APIs. +/// +/// It is sometimes desired to copy the metadata for parquet files stored on +/// remote object storage (e.g. S3) to a local file or an in-memory cache, use a +/// query engine like DataFusion to analyze the metadata to determine which file +/// to read, and then read any matching files with a single subsequent object +/// store request. +/// +/// # Example Overview +/// +/// 1. Reads the metadata of a Parquet file using [`ParquetMetaDataReader`] +/// +/// 2. Removes column statistics from the metadata (to make it smaller) +/// +/// 3. Stores the metadata in a separate file using [`ParquetMetaDataWriter`] +/// +/// 4. Reads the metadata from the separate file and uses that to read the +/// Parquet file, thus avoiding a second IO to read metadata or reparsing +/// the footer. +/// +#[tokio::main(flavor = "current_thread")] +async fn main() -> parquet::errors::Result<()> { + let tempdir = TempDir::new().unwrap(); + let parquet_path = create_parquet_file(&tempdir); + let metadata_path = tempdir.path().join("thrift_metadata.dat"); + + // In this example, we use a tokio file to mimic an async remote data source + let mut remote_parquet_file = tokio::fs::File::open(&parquet_path).await?; + + let metadata = get_metadata_from_remote_parquet_file(&mut remote_parquet_file).await; + println!( + "Metadata from 'remote' Parquet file into memory: {} bytes", + metadata.memory_size() + ); + + // now slim down the metadata and write it to a "local" file + let metadata = prepare_metadata(metadata); + write_metadata_to_local_file(metadata, &metadata_path); + + // now read the metadata from the local file and use it to read the "remote" Parquet file + let metadata = read_metadata_from_local_file(&metadata_path); + println!("Read metadata from file"); + + let batches = read_remote_parquet_file_with_metadata(remote_parquet_file, metadata).await; + + // display the results + let batches_string = pretty_format_batches(&batches).unwrap().to_string(); + let batches_lines: Vec<_> = batches_string.split('\n').collect(); + + assert_eq!( + batches_lines, + [ + "+-----+-------------+", + "| id | description |", + "+-----+-------------+", + "| 100 | oranges |", + "| 200 | apples |", + "| 201 | grapefruit |", + "| 300 | bannanas |", + "| 102 | grapes |", + "| 33 | pears |", + "+-----+-------------+", + ], + "actual output:\n\n{batches_lines:#?}" + ); + + Ok(()) +} + +/// Reads the metadata from a "remote" parquet file +/// +/// Note that this function models reading from a remote file source using a +/// tokio file. In a real application, you would implement [`MetadataFetch`] for +/// your own remote source. +/// +/// [`MetadataFetch`]: parquet::arrow::async_reader::MetadataFetch +async fn get_metadata_from_remote_parquet_file( + remote_file: &mut tokio::fs::File, +) -> ParquetMetaData { + // the remote source must know the total file size (e.g. from an object store LIST operation) + let file_size = remote_file.metadata().await.unwrap().len(); + + // tell the reader to read the page index + ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .load_and_finish(remote_file, file_size) + .await + .unwrap() +} + +/// modifies the metadata to reduce its size +fn prepare_metadata(metadata: ParquetMetaData) -> ParquetMetaData { + let orig_size = metadata.memory_size(); + + let mut builder = metadata.into_builder(); + + // remove column statistics to reduce the size of the metadata by converting + // the various structures into their respective builders and modifying them + // as needed. + for row_group in builder.take_row_groups() { + let mut row_group_builder = row_group.into_builder(); + for column in row_group_builder.take_columns() { + let column = column.into_builder().clear_statistics().build().unwrap(); + row_group_builder = row_group_builder.add_column_metadata(column); + } + let row_group = row_group_builder.build().unwrap(); + builder = builder.add_row_group(row_group); + } + let metadata = builder.build(); + + // verifiy that the size has indeed been reduced + let new_size = metadata.memory_size(); + assert!(new_size < orig_size, "metadata size did not decrease"); + println!("Reduced metadata size from {orig_size} to {new_size}"); + metadata +} + +/// writes the metadata to a file +/// +/// The data is stored using the same thrift format as the Parquet file metadata +fn write_metadata_to_local_file(metadata: ParquetMetaData, file: impl AsRef) { + let file = File::create(file).unwrap(); + ParquetMetaDataWriter::new(file, &metadata) + .finish() + .unwrap() +} + +/// Reads the metadata from a file +/// +/// This function reads the format written by `write_metadata_to_file` +fn read_metadata_from_local_file(file: impl AsRef) -> ParquetMetaData { + let file = File::open(file).unwrap(); + ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .parse_and_finish(&file) + .unwrap() +} + +/// Reads the "remote" Parquet file using the metadata +/// +/// This shows how to read the Parquet file using previously read metadata +/// instead of the metadata in the Parquet file itself. This avoids an IO / +/// having to fetch and decode the metadata from the Parquet file before +/// beginning to read it. +/// +/// Note that this function models reading from a remote file source using a +/// tokio file. In a real application, you would implement [`AsyncFileReader`] +/// for your own remote source. +/// +/// In this example, we simply buffer the results in memory as Arrow record +/// batches but a real application would likely process the batches as they are +/// read. +/// +/// [`AsyncFileReader`]: parquet::arrow::async_reader::AsyncFileReader +async fn read_remote_parquet_file_with_metadata( + remote_file: tokio::fs::File, + metadata: ParquetMetaData, +) -> Vec { + let options = ArrowReaderOptions::new() + // tell the reader to read the page index + .with_page_index_policy(PageIndexPolicy::from(true)); + // create a reader with pre-existing metadata + let arrow_reader_metadata = ArrowReaderMetadata::try_new(metadata.into(), options).unwrap(); + let reader = + ParquetRecordBatchStreamBuilder::new_with_metadata(remote_file, arrow_reader_metadata) + .build() + .unwrap(); + + reader.try_collect::>().await.unwrap() +} + +/// Make a new parquet file in the temporary directory, and returns the path +fn create_parquet_file(tmpdir: &TempDir) -> PathBuf { + let path = tmpdir.path().join("example.parquet"); + let new_file = File::create(&path).unwrap(); + + let batch = RecordBatch::try_from_iter(vec![ + ( + "id", + Arc::new(Int32Array::from(vec![100, 200, 201, 300, 102, 33])) as ArrayRef, + ), + ( + "description", + Arc::new(StringArray::from(vec![ + "oranges", + "apples", + "grapefruit", + "bannanas", + "grapes", + "pears", + ])), + ), + ]) + .unwrap(); + + let props = WriterProperties::builder() + // ensure we write the page index level statistics + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + + let mut writer = ArrowWriter::try_new(new_file, batch.schema(), Some(props)).unwrap(); + + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + + path +} diff --git a/vendor/parquet/examples/object_store.rs b/vendor/parquet/examples/object_store.rs new file mode 100644 index 0000000..fa52931 --- /dev/null +++ b/vendor/parquet/examples/object_store.rs @@ -0,0 +1,167 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_array::{ArrayRef, Int64Array, RecordBatch}; +use bytes::Bytes; +use futures::future::BoxFuture; +use futures::{FutureExt, TryFutureExt, TryStreamExt}; +use object_store::buffered::BufWriter; +use object_store::memory::InMemory; +use object_store::path::Path; +use object_store::{GetOptions, GetRange, ObjectStore, ObjectStoreExt}; +use parquet::arrow::arrow_reader::ArrowReaderOptions; +use parquet::arrow::async_reader::{AsyncFileReader, MetadataSuffixFetch, SpawnedReader}; +use parquet::arrow::{AsyncArrowWriter, ParquetRecordBatchStreamBuilder}; +use parquet::errors::{ParquetError, Result}; +use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader}; +use std::ops::Range; +use std::sync::Arc; + +/// This example demonstrates reading and writing Parquet files on object +/// storage via the [`object_store`] crate. +/// +/// # Example Overview +/// +/// 1. Writes a Parquet file to an [`ObjectStore`] by passing an +/// [`object_store::buffered::BufWriter`] directly to [`AsyncArrowWriter`], +/// via the blanket [`AsyncFileWriter`] implementation for types +/// implementing [`AsyncWrite`] (equivalent of `ParquetObjectWriter`). +/// +/// 2. Reads it back with [`ObjectStoreReader`], a minimal [`AsyncFileReader`] +/// implementation on top of an [`ObjectStore`] (equivalent of +/// `ParquetObjectReader`). +/// +/// 3. Reads it again with the reader wrapped in a [`SpawnedReader`], which +/// performs all I/O on a separate tokio runtime so that the runtime +/// decoding Parquet is not also driving the I/O (equivalent of +/// `ParquetObjectReader::with_runtime`). +/// +/// [`AsyncFileWriter`]: parquet::arrow::async_writer::AsyncFileWriter +/// [`AsyncWrite`]: tokio::io::AsyncWrite +#[tokio::main] +async fn main() -> Result<()> { + let store: Arc = Arc::new(InMemory::new()); + let path = Path::from("example.parquet"); + + // 1. Write a Parquet file: a `BufWriter` implements `AsyncWrite` and can + // therefore be passed to `AsyncArrowWriter` directly. + let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef; + let batch = RecordBatch::try_from_iter([("col", col)]).unwrap(); + + let writer = BufWriter::new(Arc::clone(&store), path.clone()); + let mut writer = AsyncArrowWriter::try_new(writer, batch.schema(), None)?; + writer.write(&batch).await?; + writer.close().await?; + + // 2. Read it back with an `AsyncFileReader` implemented on `ObjectStore`. + let reader = ObjectStoreReader::new(Arc::clone(&store), path.clone()); + let builder = ParquetRecordBatchStreamBuilder::new(reader).await?; + let read: Vec = builder.build()?.try_collect().await?; + assert_eq!(read, vec![batch.clone()]); + println!("read {} rows", read[0].num_rows()); + + // 3. Read again, performing the I/O on a dedicated runtime. + let io_runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .enable_all() + .build() + .expect("failed to build I/O runtime"); + + let reader = ObjectStoreReader::new(Arc::clone(&store), path); + let reader = SpawnedReader::new(reader, io_runtime.handle().clone()); + let builder = ParquetRecordBatchStreamBuilder::new(reader).await?; + let read: Vec = builder.build()?.try_collect().await?; + assert_eq!(read, vec![batch]); + println!("read {} rows via dedicated I/O runtime", read[0].num_rows()); + + io_runtime.shutdown_background(); + Ok(()) +} + +fn to_parquet_err(e: object_store::Error) -> ParquetError { + ParquetError::External(Box::new(e)) +} + +/// An [`AsyncFileReader`] for a location in an [`ObjectStore`] +/// +/// This mirrors the example on the [`AsyncFileReader`] trait documentation. +#[derive(Clone, Debug)] +struct ObjectStoreReader { + store: Arc, + path: Path, +} + +impl ObjectStoreReader { + fn new(store: Arc, path: Path) -> Self { + Self { store, path } + } +} + +impl AsyncFileReader for ObjectStoreReader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + self.store + .get_range(&self.path, range) + .map_err(to_parquet_err) + .boxed() + } + + fn get_byte_ranges(&mut self, ranges: Vec>) -> BoxFuture<'_, Result>> { + async move { + self.store + .get_ranges(&self.path, &ranges) + .await + .map_err(to_parquet_err) + } + .boxed() + } + + /// Loads the metadata, respecting the provided [`ArrowReaderOptions`] via + /// [`ParquetMetaDataReader::with_arrow_reader_options`] + fn get_metadata<'a>( + &'a mut self, + options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, Result>> { + async move { + let metadata = ParquetMetaDataReader::new() + .with_arrow_reader_options(options) + .load_via_suffix_and_finish(self) + .await?; + Ok(Arc::new(metadata)) + } + .boxed() + } +} + +/// Supports fetching the Parquet footer without knowing the file size upfront, +/// via suffix range requests +impl MetadataSuffixFetch for &mut ObjectStoreReader { + fn fetch_suffix(&mut self, suffix: usize) -> BoxFuture<'_, Result> { + let options = GetOptions { + range: Some(GetRange::Suffix(suffix as u64)), + ..Default::default() + }; + async move { + let resp = self + .store + .get_opts(&self.path, options) + .await + .map_err(to_parquet_err)?; + resp.bytes().await.map_err(to_parquet_err) + } + .boxed() + } +} diff --git a/vendor/parquet/examples/read_parquet.rs b/vendor/parquet/examples/read_parquet.rs new file mode 100644 index 0000000..f374fcd --- /dev/null +++ b/vendor/parquet/examples/read_parquet.rs @@ -0,0 +1,43 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::util::pretty::print_batches; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use parquet::errors::Result; +use std::fs::File; + +fn main() -> Result<()> { + // Create parquet file that will be read. + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_plain.parquet"); + let file = File::open(path).unwrap(); + + // Create a sync parquet reader with batch_size. + // batch_size is the number of rows to read up to buffer once from pages, defaults to 1024 + let parquet_reader = ParquetRecordBatchReaderBuilder::try_new(file)? + .with_batch_size(8192) + .build()?; + + let mut batches = Vec::new(); + + for batch in parquet_reader { + batches.push(batch?); + } + + print_batches(&batches).unwrap(); + Ok(()) +} diff --git a/vendor/parquet/examples/read_with_rowgroup.rs b/vendor/parquet/examples/read_with_rowgroup.rs new file mode 100644 index 0000000..e3c714c --- /dev/null +++ b/vendor/parquet/examples/read_with_rowgroup.rs @@ -0,0 +1,200 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::util::pretty::print_batches; +use bytes::{Buf, Bytes}; +use parquet::arrow::arrow_reader::{ParquetRecordBatchReader, RowGroups, RowSelection}; +use parquet::arrow::async_reader::AsyncFileReader; +use parquet::arrow::{ProjectionMask, parquet_to_arrow_field_levels}; +use parquet::column::page::{PageIterator, PageReader}; +use parquet::errors::{ParquetError, Result}; +use parquet::file::metadata::{ParquetMetaData, RowGroupMetaData}; +use parquet::file::reader::{ChunkReader, Length}; +use parquet::file::serialized_reader::SerializedPageReader; +use std::sync::Arc; +use tokio::fs::File; + +#[tokio::main(flavor = "current_thread")] +async fn main() -> Result<()> { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_plain.parquet"); + let mut file = File::open(&path).await.unwrap(); + + // The metadata could be cached in other places, this example only shows how to read + let metadata = Arc::try_unwrap(file.get_metadata(None).await?).unwrap(); + + for row_group_idx in 0..metadata.row_groups().len() { + let mut rowgroup = + InMemoryRowGroup::create(metadata.clone(), row_group_idx, ProjectionMask::all()); + rowgroup.async_fetch_data(&mut file, None).await?; + let reader = rowgroup.build_reader(1024, None)?; + + for batch in reader { + let batch = batch?; + print_batches(&[batch])?; + } + } + + Ok(()) +} + +/// Implements [`PageIterator`] for a single column chunk, yielding a single [`PageReader`] +struct ColumnChunkIterator { + reader: Option>>, +} + +impl Iterator for ColumnChunkIterator { + type Item = Result>; + + fn next(&mut self) -> Option { + self.reader.take() + } +} + +impl PageIterator for ColumnChunkIterator {} + +/// An in-memory column chunk +#[derive(Clone)] +pub struct ColumnChunkData { + offset: usize, + data: Bytes, +} + +impl ColumnChunkData { + fn get(&self, start: u64) -> Result { + let start = start as usize - self.offset; + Ok(self.data.slice(start..)) + } +} + +impl Length for ColumnChunkData { + fn len(&self) -> u64 { + self.data.len() as u64 + } +} + +impl ChunkReader for ColumnChunkData { + type T = bytes::buf::Reader; + + fn get_read(&self, start: u64) -> Result { + Ok(self.get(start)?.reader()) + } + + fn get_bytes(&self, start: u64, length: usize) -> Result { + Ok(self.get(start)?.slice(..length)) + } +} + +#[derive(Clone)] +pub struct InMemoryRowGroup { + metadata: ParquetMetaData, + row_group_idx: usize, + mask: ProjectionMask, + column_chunks: Vec>>, +} + +impl RowGroups for InMemoryRowGroup { + fn num_rows(&self) -> usize { + self.row_group_metadata().num_rows() as usize + } + + fn column_chunks(&self, i: usize) -> Result> { + match &self.column_chunks[i] { + None => Err(ParquetError::General(format!( + "Invalid column index {i}, column was not fetched" + ))), + Some(data) => { + let page_reader: Box = Box::new(SerializedPageReader::new( + data.clone(), + self.row_group_metadata().column(i), + self.num_rows(), + None, + )?); + + Ok(Box::new(ColumnChunkIterator { + reader: Some(Ok(page_reader)), + })) + } + } + } + + fn row_groups(&self) -> Box + '_> { + Box::new(std::iter::once(self.row_group_metadata())) + } + + fn metadata(&self) -> &ParquetMetaData { + &self.metadata + } +} + +impl InMemoryRowGroup { + pub fn create(metadata: ParquetMetaData, row_group_idx: usize, mask: ProjectionMask) -> Self { + let column_chunks = metadata + .row_group(row_group_idx) + .columns() + .iter() + .map(|_| None) + .collect::>(); + + Self { + metadata, + row_group_idx, + mask, + column_chunks, + } + } + + pub fn row_group_metadata(&self) -> &RowGroupMetaData { + self.metadata.row_group(self.row_group_idx) + } + + pub fn build_reader( + &self, + batch_size: usize, + selection: Option, + ) -> Result { + let levels = parquet_to_arrow_field_levels( + &self.row_group_metadata().schema_descr_ptr(), + self.mask.clone(), + None, + )?; + + ParquetRecordBatchReader::try_new_with_row_groups(&levels, self, batch_size, selection) + } + + /// fetch data from a reader in sync mode + pub async fn async_fetch_data( + &mut self, + reader: &mut R, + _selection: Option<&RowSelection>, + ) -> Result<()> { + let mut vs = std::mem::take(&mut self.column_chunks); + for (leaf_idx, meta) in self.row_group_metadata().columns().iter().enumerate() { + if self.mask.leaf_included(leaf_idx) { + let (start, len) = meta.byte_range(); + let data = reader.get_bytes(start..(start + len)).await?; + + vs[leaf_idx] = Some(Arc::new(ColumnChunkData { + offset: start as usize, + data, + })); + } + } + self.column_chunks = std::mem::take(&mut vs); + Ok(()) + } +} diff --git a/vendor/parquet/examples/write_parquet.rs b/vendor/parquet/examples/write_parquet.rs new file mode 100644 index 0000000..ebdd952 --- /dev/null +++ b/vendor/parquet/examples/write_parquet.rs @@ -0,0 +1,140 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fs::File; +use std::path::PathBuf; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use arrow::array::{StructArray, UInt64Builder}; +use arrow::datatypes::DataType::UInt64; +use arrow::datatypes::{Field, Schema}; +use clap::{Parser, ValueEnum}; +use parquet::arrow::ArrowWriter as ParquetWriter; +use parquet::basic::Encoding; +use parquet::errors::Result; +use parquet::file::properties::{BloomFilterPosition, WriterProperties}; +use sysinfo::{ProcessRefreshKind, ProcessesToUpdate, RefreshKind, System}; + +#[derive(ValueEnum, Clone)] +enum BloomFilterPositionArg { + End, + AfterRowGroup, +} + +#[derive(Parser)] +#[command(version)] +/// Writes sequences of integers, with a Bloom Filter, while logging timing and memory usage. +struct Args { + #[arg(long, default_value_t = 1000)] + /// Number of batches to write + iterations: u64, + + #[arg(long, default_value_t = 1000000)] + /// Number of rows in each batch + batch: u64, + + #[arg(long, value_enum, default_value_t=BloomFilterPositionArg::AfterRowGroup)] + /// Where to write Bloom Filters + bloom_filter_position: BloomFilterPositionArg, + + /// Path to the file to write + path: PathBuf, +} + +fn now() -> String { + chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string() +} + +fn mem(system: &mut System) -> String { + let pid = match sysinfo::get_current_pid() { + Ok(pid) => pid, + Err(e) => return format!("Can't get process PID: {e}"), + }; + + let remove_dead_processes = true; + system.refresh_processes_specifics( + ProcessesToUpdate::Some(&[pid]), + remove_dead_processes, + ProcessRefreshKind::everything(), + ); + + system + .process(pid) + .map(|proc| format!("{}MB", proc.memory() / 1_000_000)) + .unwrap_or("N/A".to_string()) +} + +fn main() -> Result<()> { + let args = Args::parse(); + + let bloom_filter_position = match args.bloom_filter_position { + BloomFilterPositionArg::End => BloomFilterPosition::End, + BloomFilterPositionArg::AfterRowGroup => BloomFilterPosition::AfterRowGroup, + }; + + let properties = WriterProperties::builder() + .set_column_bloom_filter_enabled("id".into(), true) + .set_column_encoding("id".into(), Encoding::DELTA_BINARY_PACKED) + .set_bloom_filter_position(bloom_filter_position) + .build(); + let schema = Arc::new(Schema::new(vec![Field::new("id", UInt64, false)])); + // Create parquet file that will be read. + let file = File::create(args.path).unwrap(); + let mut writer = ParquetWriter::try_new(file, schema.clone(), Some(properties))?; + + let mut system = System::new_with_specifics(RefreshKind::everything()); + eprintln!( + "{} Writing {} batches of {} rows. RSS = {}", + now(), + args.iterations, + args.batch, + mem(&mut system) + ); + + let mut array_builder = UInt64Builder::new(); + let mut last_log = Instant::now(); + for i in 0..args.iterations { + if Instant::now() - last_log > Duration::new(10, 0) { + last_log = Instant::now(); + eprintln!( + "{} Iteration {}/{}. RSS = {}", + now(), + i + 1, + args.iterations, + mem(&mut system) + ); + } + for j in 0..args.batch { + array_builder.append_value(i + j); + } + writer.write( + &StructArray::new( + schema.fields().clone(), + vec![Arc::new(array_builder.finish())], + None, + ) + .into(), + )?; + } + writer.flush()?; + writer.close()?; + + eprintln!("{} Done. RSS = {}", now(), mem(&mut system)); + + Ok(()) +} diff --git a/vendor/parquet/pytest/requirements.in b/vendor/parquet/pytest/requirements.in new file mode 100644 index 0000000..575fb83 --- /dev/null +++ b/vendor/parquet/pytest/requirements.in @@ -0,0 +1,20 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +pytest +pyspark +black +pandas diff --git a/vendor/parquet/pytest/requirements.txt b/vendor/parquet/pytest/requirements.txt new file mode 100644 index 0000000..cf0adc7 --- /dev/null +++ b/vendor/parquet/pytest/requirements.txt @@ -0,0 +1,225 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# This file is autogenerated by pip-compile with python 3.10 +# To update, run: +# +# pip-compile --generate-hashes --resolver=backtracking +# +black==26.3.1 \ + --hash=sha256:0126ae5b7c09957da2bdbd91a9ba1207453feada9e9fe51992848658c6c8e01c \ + --hash=sha256:0f76ff19ec5297dd8e66eb64deda23631e642c9393ab592826fd4bdc97a4bce7 \ + --hash=sha256:28ef38aee69e4b12fda8dba75e21f9b4f979b490c8ac0baa7cb505369ac9e1ff \ + --hash=sha256:2bd5aa94fc267d38bb21a70d7410a89f1a1d318841855f698746f8e7f51acd1b \ + --hash=sha256:2c50f5063a9641c7eed7795014ba37b0f5fa227f3d408b968936e24bc0566b07 \ + --hash=sha256:2d6bfaf7fd0993b420bed691f20f9492d53ce9a2bcccea4b797d34e947318a78 \ + --hash=sha256:41cd2012d35b47d589cb8a16faf8a32ef7a336f56356babd9fcf70939ad1897f \ + --hash=sha256:474c27574d6d7037c1bc875a81d9be0a9a4f9ee95e62800dab3cfaadbf75acd5 \ + --hash=sha256:5602bdb96d52d2d0672f24f6ffe5218795736dd34807fd0fd55ccd6bf206168b \ + --hash=sha256:5e9d0d86df21f2e1677cc4bd090cd0e446278bcbbe49bf3659c308c3e402843e \ + --hash=sha256:5ed0ca58586c8d9a487352a96b15272b7fa55d139fc8496b519e78023a8dab0a \ + --hash=sha256:6c54a4a82e291a1fee5137371ab488866b7c86a3305af4026bdd4dc78642e1ac \ + --hash=sha256:6e131579c243c98f35bce64a7e08e87fb2d610544754675d4a0e73a070a5aa3a \ + --hash=sha256:855822d90f884905362f602880ed8b5df1b7e3ee7d0db2502d4388a954cc8c54 \ + --hash=sha256:86a8b5035fce64f5dcd1b794cf8ec4d31fe458cf6ce3986a30deb434df82a1d2 \ + --hash=sha256:8a33d657f3276328ce00e4d37fe70361e1ec7614da5d7b6e78de5426cb56332f \ + --hash=sha256:92c0ec1f2cc149551a2b7b47efc32c866406b6891b0ee4625e95967c8f4acfb1 \ + --hash=sha256:9a5e9f45e5d5e1c5b5c29b3bd4265dcc90e8b92cf4534520896ed77f791f4da5 \ + --hash=sha256:afc622538b430aa4c8c853f7f63bc582b3b8030fd8c80b70fb5fa5b834e575c2 \ + --hash=sha256:b07fc0dab849d24a80a29cfab8d8a19187d1c4685d8a5e6385a5ce323c1f015f \ + --hash=sha256:b5e6f89631eb88a7302d416594a32faeee9fb8fb848290da9d0a5f2903519fc1 \ + --hash=sha256:bf9bf162ed91a26f1adba8efda0b573bc6924ec1408a52cc6f82cb73ec2b142c \ + --hash=sha256:c7e72339f841b5a237ff14f7d3880ddd0fc7f98a1199e8c4327f9a4f478c1839 \ + --hash=sha256:ddb113db38838eb9f043623ba274cfaf7d51d5b0c22ecb30afe58b1bb8322983 \ + --hash=sha256:dfdd51fc3e64ea4f35873d1b3fb25326773d55d2329ff8449139ebaad7357efb \ + --hash=sha256:f1cd08e99d2f9317292a311dfe578fd2a24b15dbce97792f9c4d752275c1fa56 \ + --hash=sha256:f89f2ab047c76a9c03f78d0d66ca519e389519902fa27e7a91117ef7611c0568 + # via -r requirements.in +click==8.1.3 \ + --hash=sha256:7682dc8afb30297001674575ea00d1814d808d6a36af415a82bd481d37ba7b8e \ + --hash=sha256:bb4d8133cb15a609f44e8213d9b391b0809795062913b383c62be0ee95b1db48 + # via black +exceptiongroup==1.0.4 \ + --hash=sha256:542adf9dea4055530d6e1279602fa5cb11dab2395fa650b8674eaec35fc4a828 \ + --hash=sha256:bd14967b79cd9bdb54d97323216f8fdf533e278df937aa2a90089e7d6e06e5ec + # via pytest +iniconfig==1.1.1 \ + --hash=sha256:011e24c64b7f47f6ebd835bb12a743f2fbe9a26d4cecaa7f53bc4f35ee9da8b3 \ + --hash=sha256:bc3af051d7d14b2ee5ef9969666def0cd1a000e121eaea580d4a313df4b37f32 + # via pytest +mypy-extensions==0.4.3 \ + --hash=sha256:090fedd75945a69ae91ce1303b5824f428daf5a028d2f6ab8a299250a846f15d \ + --hash=sha256:2d82818f5bb3e369420cb3c4060a7970edba416647068eb4c5343488a6c604a8 + # via black +numpy==1.23.5 \ + --hash=sha256:01dd17cbb340bf0fc23981e52e1d18a9d4050792e8fb8363cecbf066a84b827d \ + --hash=sha256:06005a2ef6014e9956c09ba07654f9837d9e26696a0470e42beedadb78c11b07 \ + --hash=sha256:09b7847f7e83ca37c6e627682f145856de331049013853f344f37b0c9690e3df \ + --hash=sha256:0aaee12d8883552fadfc41e96b4c82ee7d794949e2a7c3b3a7201e968c7ecab9 \ + --hash=sha256:0cbe9848fad08baf71de1a39e12d1b6310f1d5b2d0ea4de051058e6e1076852d \ + --hash=sha256:1b1766d6f397c18153d40015ddfc79ddb715cabadc04d2d228d4e5a8bc4ded1a \ + --hash=sha256:33161613d2269025873025b33e879825ec7b1d831317e68f4f2f0f84ed14c719 \ + --hash=sha256:5039f55555e1eab31124a5768898c9e22c25a65c1e0037f4d7c495a45778c9f2 \ + --hash=sha256:522e26bbf6377e4d76403826ed689c295b0b238f46c28a7251ab94716da0b280 \ + --hash=sha256:56e454c7833e94ec9769fa0f86e6ff8e42ee38ce0ce1fa4cbb747ea7e06d56aa \ + --hash=sha256:58f545efd1108e647604a1b5aa809591ccd2540f468a880bedb97247e72db387 \ + --hash=sha256:5e05b1c973a9f858c74367553e236f287e749465f773328c8ef31abe18f691e1 \ + --hash=sha256:7903ba8ab592b82014713c491f6c5d3a1cde5b4a3bf116404e08f5b52f6daf43 \ + --hash=sha256:8969bfd28e85c81f3f94eb4a66bc2cf1dbdc5c18efc320af34bffc54d6b1e38f \ + --hash=sha256:92c8c1e89a1f5028a4c6d9e3ccbe311b6ba53694811269b992c0b224269e2398 \ + --hash=sha256:9c88793f78fca17da0145455f0d7826bcb9f37da4764af27ac945488116efe63 \ + --hash=sha256:a7ac231a08bb37f852849bbb387a20a57574a97cfc7b6cabb488a4fc8be176de \ + --hash=sha256:abdde9f795cf292fb9651ed48185503a2ff29be87770c3b8e2a14b0cd7aa16f8 \ + --hash=sha256:af1da88f6bc3d2338ebbf0e22fe487821ea4d8e89053e25fa59d1d79786e7481 \ + --hash=sha256:b2a9ab7c279c91974f756c84c365a669a887efa287365a8e2c418f8b3ba73fb0 \ + --hash=sha256:bf837dc63ba5c06dc8797c398db1e223a466c7ece27a1f7b5232ba3466aafe3d \ + --hash=sha256:ca51fcfcc5f9354c45f400059e88bc09215fb71a48d3768fb80e357f3b457e1e \ + --hash=sha256:ce571367b6dfe60af04e04a1834ca2dc5f46004ac1cc756fb95319f64c095a96 \ + --hash=sha256:d208a0f8729f3fb790ed18a003f3a57895b989b40ea4dce4717e9cf4af62c6bb \ + --hash=sha256:dbee87b469018961d1ad79b1a5d50c0ae850000b639bcb1b694e9981083243b6 \ + --hash=sha256:e9f4c4e51567b616be64e05d517c79a8a22f3606499941d97bb76f2ca59f982d \ + --hash=sha256:f063b69b090c9d918f9df0a12116029e274daf0181df392839661c4c7ec9018a \ + --hash=sha256:f9a909a8bae284d46bbfdefbdd4a262ba19d3bc9921b1e76126b1d21c3c34135 + # via pandas +packaging==24.0 \ + --hash=sha256:2ddfb553fdf02fb784c234c7ba6ccc288296ceabec964ad2eae3777778130bc5 \ + --hash=sha256:eb82c5e3e56209074766e6885bb04b8c38a0c015d0a30036ebe7ece34c9989e9 + # via + # black + # pytest +pandas==1.5.2 \ + --hash=sha256:0183cb04a057cc38fde5244909fca9826d5d57c4a5b7390c0cc3fa7acd9fa883 \ + --hash=sha256:1fc87eac0541a7d24648a001d553406f4256e744d92df1df8ebe41829a915028 \ + --hash=sha256:220b98d15cee0b2cd839a6358bd1f273d0356bf964c1a1aeb32d47db0215488b \ + --hash=sha256:2552bffc808641c6eb471e55aa6899fa002ac94e4eebfa9ec058649122db5824 \ + --hash=sha256:315e19a3e5c2ab47a67467fc0362cb36c7c60a93b6457f675d7d9615edad2ebe \ + --hash=sha256:344021ed3e639e017b452aa8f5f6bf38a8806f5852e217a7594417fb9bbfa00e \ + --hash=sha256:375262829c8c700c3e7cbb336810b94367b9c4889818bbd910d0ecb4e45dc261 \ + --hash=sha256:457d8c3d42314ff47cc2d6c54f8fc0d23954b47977b2caed09cd9635cb75388b \ + --hash=sha256:4aed257c7484d01c9a194d9a94758b37d3d751849c05a0050c087a358c41ad1f \ + --hash=sha256:530948945e7b6c95e6fa7aa4be2be25764af53fba93fe76d912e35d1c9ee46f5 \ + --hash=sha256:5ae7e989f12628f41e804847a8cc2943d362440132919a69429d4dea1f164da0 \ + --hash=sha256:71f510b0efe1629bf2f7c0eadb1ff0b9cf611e87b73cd017e6b7d6adb40e2b3a \ + --hash=sha256:73f219fdc1777cf3c45fde7f0708732ec6950dfc598afc50588d0d285fddaefc \ + --hash=sha256:8092a368d3eb7116e270525329a3e5c15ae796ccdf7ccb17839a73b4f5084a39 \ + --hash=sha256:82ae615826da838a8e5d4d630eb70c993ab8636f0eff13cb28aafc4291b632b5 \ + --hash=sha256:9608000a5a45f663be6af5c70c3cbe634fa19243e720eb380c0d378666bc7702 \ + --hash=sha256:a40dd1e9f22e01e66ed534d6a965eb99546b41d4d52dbdb66565608fde48203f \ + --hash=sha256:b4f5a82afa4f1ff482ab8ded2ae8a453a2cdfde2001567b3ca24a4c5c5ca0db3 \ + --hash=sha256:c009a92e81ce836212ce7aa98b219db7961a8b95999b97af566b8dc8c33e9519 \ + --hash=sha256:c218796d59d5abd8780170c937b812c9637e84c32f8271bbf9845970f8c1351f \ + --hash=sha256:cc3cd122bea268998b79adebbb8343b735a5511ec14efb70a39e7acbc11ccbdc \ + --hash=sha256:d0d8fd58df5d17ddb8c72a5075d87cd80d71b542571b5f78178fb067fa4e9c72 \ + --hash=sha256:e18bc3764cbb5e118be139b3b611bc3fbc5d3be42a7e827d1096f46087b395eb \ + --hash=sha256:e2b83abd292194f350bb04e188f9379d36b8dfac24dd445d5c87575f3beaf789 \ + --hash=sha256:e7469271497960b6a781eaa930cba8af400dd59b62ec9ca2f4d31a19f2f91090 \ + --hash=sha256:e9dbacd22555c2d47f262ef96bb4e30880e5956169741400af8b306bbb24a273 \ + --hash=sha256:f6257b314fc14958f8122779e5a1557517b0f8e500cfb2bd53fa1f75a8ad0af2 + # via -r requirements.in +pathspec==1.0.4 \ + --hash=sha256:0210e2ae8a21a9137c0d470578cb0e595af87edaa6ebf12ff176f14a02e0e645 \ + --hash=sha256:fb6ae2fd4e7c921a165808a552060e722767cfa526f99ca5156ed2ce45a5c723 + # via black +platformdirs==2.5.4 \ + --hash=sha256:1006647646d80f16130f052404c6b901e80ee4ed6bef6792e1f238a8969106f7 \ + --hash=sha256:af0276409f9a02373d540bf8480021a048711d572745aef4b7842dad245eba10 + # via black +pluggy==1.6.0 \ + --hash=sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3 \ + --hash=sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746 + # via pytest +py4j==0.10.9.7 \ + --hash=sha256:0b6e5315bb3ada5cf62ac651d107bb2ebc02def3dee9d9548e3baac644ea8dbb \ + --hash=sha256:85defdfd2b2376eb3abf5ca6474b51ab7e0de341c75a02f46dc9b5976f5a5c1b + # via pyspark +pygments==2.20.0 \ + --hash=sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f \ + --hash=sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176 + # via pytest +pyspark==3.4.4 \ + --hash=sha256:b831a99c1aeb058bc4a2498665de10006b899b2438ba15470cc88f577755ebf8 + # via -r requirements.in +pytest==9.0.3 \ + --hash=sha256:2c5efc453d45394fdd706ade797c0a81091eccd1d6e4bccfcd476e2b8e0ab5d9 \ + --hash=sha256:b86ada508af81d19edeb213c681b1d48246c1a91d304c6c81a427674c17eb91c + # via -r requirements.in +python-dateutil==2.8.2 \ + --hash=sha256:0123cacc1627ae19ddf3c27a5de5bd67ee4586fbdd6440d9748f8abb483d3e86 \ + --hash=sha256:961d03dc3453ebbc59dbdea9e4e11c5651520a876d0f4db161e8674aae935da9 + # via pandas +pytokens==0.4.1 \ + --hash=sha256:0fc71786e629cef478cbf29d7ea1923299181d0699dbe7c3c0f4a583811d9fc1 \ + --hash=sha256:11edda0942da80ff58c4408407616a310adecae1ddd22eef8c692fe266fa5009 \ + --hash=sha256:140709331e846b728475786df8aeb27d24f48cbcf7bcd449f8de75cae7a45083 \ + --hash=sha256:24afde1f53d95348b5a0eb19488661147285ca4dd7ed752bbc3e1c6242a304d1 \ + --hash=sha256:26cef14744a8385f35d0e095dc8b3a7583f6c953c2e3d269c7f82484bf5ad2de \ + --hash=sha256:27b83ad28825978742beef057bfe406ad6ed524b2d28c252c5de7b4a6dd48fa2 \ + --hash=sha256:292052fe80923aae2260c073f822ceba21f3872ced9a68bb7953b348e561179a \ + --hash=sha256:29d1d8fb1030af4d231789959f21821ab6325e463f0503a61d204343c9b355d1 \ + --hash=sha256:2a44ed93ea23415c54f3face3b65ef2b844d96aeb3455b8a69b3df6beab6acc5 \ + --hash=sha256:30f51edd9bb7f85c748979384165601d028b84f7bd13fe14d3e065304093916a \ + --hash=sha256:34bcc734bd2f2d5fe3b34e7b3c0116bfb2397f2d9666139988e7a3eb5f7400e3 \ + --hash=sha256:3ad72b851e781478366288743198101e5eb34a414f1d5627cdd585ca3b25f1db \ + --hash=sha256:3f901fe783e06e48e8cbdc82d631fca8f118333798193e026a50ce1b3757ea68 \ + --hash=sha256:42f144f3aafa5d92bad964d471a581651e28b24434d184871bd02e3a0d956037 \ + --hash=sha256:4a14d5f5fc78ce85e426aa159489e2d5961acf0e47575e08f35584009178e321 \ + --hash=sha256:4a58d057208cb9075c144950d789511220b07636dd2e4708d5645d24de666bdc \ + --hash=sha256:4e691d7f5186bd2842c14813f79f8884bb03f5995f0575272009982c5ac6c0f7 \ + --hash=sha256:5502408cab1cb18e128570f8d598981c68a50d0cbd7c61312a90507cd3a1276f \ + --hash=sha256:584c80c24b078eec1e227079d56dc22ff755e0ba8654d8383b2c549107528918 \ + --hash=sha256:5ad948d085ed6c16413eb5fec6b3e02fa00dc29a2534f088d3302c47eb59adf9 \ + --hash=sha256:670d286910b531c7b7e3c0b453fd8156f250adb140146d234a82219459b9640c \ + --hash=sha256:682fa37ff4d8e95f7df6fe6fe6a431e8ed8e788023c6bcc0f0880a12eab80ad1 \ + --hash=sha256:6d6c4268598f762bc8e91f5dbf2ab2f61f7b95bdc07953b602db879b3c8c18e1 \ + --hash=sha256:79fc6b8699564e1f9b521582c35435f1bd32dd06822322ec44afdeba666d8cb3 \ + --hash=sha256:8bdb9d0ce90cbf99c525e75a2fa415144fd570a1ba987380190e8b786bc6ef9b \ + --hash=sha256:8fcb9ba3709ff77e77f1c7022ff11d13553f3c30299a9fe246a166903e9091eb \ + --hash=sha256:941d4343bf27b605e9213b26bfa1c4bf197c9c599a9627eb7305b0defcfe40c1 \ + --hash=sha256:967cf6e3fd4adf7de8fc73cd3043754ae79c36475c1c11d514fc72cf5490094a \ + --hash=sha256:970b08dd6b86058b6dc07efe9e98414f5102974716232d10f32ff39701e841c4 \ + --hash=sha256:97f50fd18543be72da51dd505e2ed20d2228c74e0464e4262e4899797803d7fa \ + --hash=sha256:9bd7d7f544d362576be74f9d5901a22f317efc20046efe2034dced238cbbfe78 \ + --hash=sha256:add8bf86b71a5d9fb5b89f023a80b791e04fba57960aa790cc6125f7f1d39dfe \ + --hash=sha256:b35d7e5ad269804f6697727702da3c517bb8a5228afa450ab0fa787732055fc9 \ + --hash=sha256:b49750419d300e2b5a3813cf229d4e5a4c728dae470bcc89867a9ad6f25a722d \ + --hash=sha256:d31b97b3de0f61571a124a00ffe9a81fb9939146c122c11060725bd5aea79975 \ + --hash=sha256:d70e77c55ae8380c91c0c18dea05951482e263982911fc7410b1ffd1dadd3440 \ + --hash=sha256:d9907d61f15bf7261d7e775bd5d7ee4d2930e04424bab1972591918497623a16 \ + --hash=sha256:da5baeaf7116dced9c6bb76dc31ba04a2dc3695f3d9f74741d7910122b456edc \ + --hash=sha256:dc74c035f9bfca0255c1af77ddd2d6ae8419012805453e4b0e7513e17904545d \ + --hash=sha256:dcafc12c30dbaf1e2af0490978352e0c4041a7cde31f4f81435c2a5e8b9cabb6 \ + --hash=sha256:ee44d0f85b803321710f9239f335aafe16553b39106384cef8e6de40cb4ef2f6 \ + --hash=sha256:f66a6bbe741bd431f6d741e617e0f39ec7257ca1f89089593479347cc4d13324 + # via black +pytz==2022.6 \ + --hash=sha256:222439474e9c98fced559f1709d89e6c9cbf8d79c794ff3eb9f8800064291427 \ + --hash=sha256:e89512406b793ca39f5971bc999cc538ce125c0e51c27941bef4568b460095e2 + # via pandas +six==1.16.0 \ + --hash=sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926 \ + --hash=sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254 + # via python-dateutil +tomli==2.0.1 \ + --hash=sha256:939de3e7a6161af0c887ef91b7d41a53e7c5a1ca976325f429cb46ea9bc30ecc \ + --hash=sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f + # via + # black + # pytest +typing-extensions==4.10.0 \ + --hash=sha256:69b1a937c3a517342112fb4c6df7e72fc39a38e7891a5730ed4985b5214b5475 \ + --hash=sha256:b0abd7c89e8fb96f98db18d86106ff1d90ab692004eb746cf6eda2682f91b3cb + # via black diff --git a/vendor/parquet/pytest/test_parquet_integration.py b/vendor/parquet/pytest/test_parquet_integration.py new file mode 100755 index 0000000..4ff0163 --- /dev/null +++ b/vendor/parquet/pytest/test_parquet_integration.py @@ -0,0 +1,225 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +import pyspark.sql +from pyspark.sql.types import ( + ArrayType, + StructType, + StructField, + StringType, + IntegerType, + LongType, +) +import pandas as pd +import re +from tempfile import NamedTemporaryFile, TemporaryDirectory +import subprocess +import pathlib +import pytest + + +def create_data_and_spark_df(n): + spark = pyspark.sql.SparkSession.builder.getOrCreate() + spark.conf.set("parquet.bloom.filter.enabled", True) + spark.conf.set("parquet.bloom.filter.expected.ndv", 10) + spark.conf.set("parquet.bloom.filter.max.bytes", 32) + data = [(f"id-{i % 10}", f"name-{i%10}", i * 2, i * 2 + 1) for i in range(n)] + schema = StructType( + [ + StructField("id", StringType(), True), + StructField("name", StringType(), True), + StructField("int32", IntegerType(), True), + StructField("int64", LongType(), True), + ] + ) + df = spark.createDataFrame(data, schema).repartition(1) + return data, df + + +def create_spark_df_with_arrays(n): + def make_list(x): + return [x for _ in range(8)] + + spark = pyspark.sql.SparkSession.builder.getOrCreate() + data = [ + (f"id-{i % 10}", f"name-{i%10}", make_list(i + 1), make_list(i + 2)) + for i in range(n) + ] + schema = StructType( + [ + StructField("id", StringType(), True), + StructField("name", StringType(), True), + StructField("int32_list", ArrayType(IntegerType()), True), + StructField("int64_list", ArrayType(LongType()), True), + ] + ) + df = spark.createDataFrame(data, schema).repartition(1) + return data, df + + +def create_data_and_pandas_df(n): + data = [(f"id-{i % 10}", f"name-{i%10}", i * 2, i * 2 + 1) for i in range(n)] + df = pd.DataFrame(data, columns=["id", "name", "int32", "int64"]) + return data, df + + +def get_expected_output(expected_results): + expected = ["Row group #0", "=" * 80] + for value, result in expected_results: + result_str = "present" if result else "absent" + expected.append(f"Value {value} is {result_str} in bloom filter") + expected = "\n".join(expected) + "\n" + return expected.encode("utf-8") + + +def get_from_csv_cli_output(schema_file, output_file, csv_file): + args = [ + "parquet-fromcsv", + "--schema", + schema_file, + "--has-header", + "--enable-bloom-filter", + "true", + "--input-file", + csv_file, + "--output-file", + output_file, + ] + return subprocess.check_output(args) + + +def get_show_filter_cli_output(output_dir, col_name, test_values): + # take the first (and only) parquet file + (parquet_file,) = sorted(pathlib.Path(output_dir).glob("*.parquet")) + args = [ + "parquet-show-bloom-filter", + parquet_file, + col_name, + ] + for v in test_values: + args.append(str(v)) + return subprocess.check_output(args) + + +def get_index_cli_output(output_dir, col_name): + (parquet_file,) = sorted(pathlib.Path(output_dir).glob("*.parquet")) + args = ["parquet-index", parquet_file, col_name] + return subprocess.check_output(args) + + +SCHEMA = b"""message schema { + required binary id (UTF8); + required binary name (UTF8); + required int32 int32; + required int64 int64; +}""" + + +@pytest.mark.parametrize("n", [1, 10]) +class TestParquetIntegration: + def test_pyspark_bloom_filter(self, n): + data, df = create_data_and_spark_df(n) + with TemporaryDirectory() as output_dir: + df.write.parquet(output_dir, mode="overwrite") + self._test_column_filters(output_dir, data) + + def test_bloom_filter_round_trip(self, n): + data, df = create_data_and_pandas_df(n) + with NamedTemporaryFile(suffix=".csv") as csv_file, NamedTemporaryFile( + suffix=".schema" + ) as schema_file, TemporaryDirectory() as output_dir: + schema_file.write(SCHEMA) + schema_file.flush() + df.to_csv(csv_file.name, index=False, header=True) + parquet_file = pathlib.Path(output_dir) / "output.parquet" + cli_output = get_from_csv_cli_output( + schema_file.name, parquet_file, csv_file.name + ) + assert cli_output == b"" + self._test_column_filters(output_dir, data) + + def _test_column_filters(self, output_dir, data): + self._test_string_col_filter(output_dir, data) + self._test_int32_col_filter(output_dir, data) + self._test_int64_col_filter(output_dir, data) + + def _test_string_col_filter(self, output_dir, data): + test_values = [] + expected_results = [] + for v in data: + test_values.append(v[0]) + expected_results.append((v[0], True)) + for v in data: + test_values.append(v[1]) + expected_results.append((v[1], False)) + cli_output = get_show_filter_cli_output(output_dir, "id", test_values) + assert cli_output == get_expected_output(expected_results) + + def _test_int32_col_filter(self, output_dir, data): + test_values = [] + expected_results = [] + for v in data: + test_values.append(v[2]) + expected_results.append((v[2], True)) + for v in data: + test_values.append(v[3]) + expected_results.append((v[3], False)) + cli_output = get_show_filter_cli_output(output_dir, "int32", test_values) + assert cli_output == get_expected_output(expected_results) + + def _test_int64_col_filter(self, output_dir, data): + test_values = [] + expected_results = [] + for v in data: + test_values.append(v[3]) + expected_results.append((v[3], True)) + for v in data: + test_values.append(v[2]) + expected_results.append((v[2], False)) + cli_output = get_show_filter_cli_output(output_dir, "int64", test_values) + assert cli_output == get_expected_output(expected_results) + + +class TestParquetPageIndexIntegration: + @staticmethod + def _assert_cli_output(out, pattern): + normalized = re.sub(r"\s+", " ", out.decode("utf-8")).strip() + if re.search(pattern, normalized): + return + raise AssertionError("Failed to find expected pattern in cli output") + + def test_index(self): + _, df = create_spark_df_with_arrays(10) + with TemporaryDirectory() as output_dir: + df.write.parquet(output_dir, mode="overwrite") + + name_index_expected_pattern = r""" + Page 0 at offset \d+ with length \d+ and row count 10, min id-0, max id-9 + """.strip() + cli_output = get_index_cli_output(output_dir, "id") + self._assert_cli_output(cli_output, name_index_expected_pattern) + + int32_list_index_expected_pattern = r""" + Page 0 at offset \d+ with length \d+ and row count 10, min 1, max 10 + """.strip() + cli_output = get_index_cli_output(output_dir, "int32_list.list.element") + self._assert_cli_output(cli_output, int32_list_index_expected_pattern) + + cli_output = get_index_cli_output(output_dir, "int64_list.list.element") + int64_list_index_expected_pattern = r""" + Page 0 at offset \d+ with length \d+ and row count 10, min 2, max 11 + """.strip() + self._assert_cli_output(cli_output, int64_list_index_expected_pattern) diff --git a/vendor/parquet/src/arrow/array_reader/builder.rs b/vendor/parquet/src/arrow/array_reader/builder.rs new file mode 100644 index 0000000..f89a789 --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/builder.rs @@ -0,0 +1,697 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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::sync::{Arc, RwLock}; + +use arrow_schema::{DataType, Fields, SchemaBuilder}; + +use crate::arrow::ProjectionMask; +use crate::arrow::array_reader::byte_view_array::make_byte_view_array_reader; +use crate::arrow::array_reader::cached_array_reader::CacheRole; +use crate::arrow::array_reader::cached_array_reader::CachedArrayReader; +use crate::arrow::array_reader::empty_array::make_empty_array_reader; +use crate::arrow::array_reader::fixed_len_byte_array::make_fixed_len_byte_array_reader; +use crate::arrow::array_reader::row_group_cache::RowGroupCache; +use crate::arrow::array_reader::row_group_index::RowGroupIndexReader; +use crate::arrow::array_reader::row_number::RowNumberReader; +use crate::arrow::array_reader::{ + ArrayReader, FixedSizeListArrayReader, ListArrayReader, ListViewArrayReader, MapArrayReader, + NullArrayReader, PrimitiveArrayReader, RowGroups, StructArrayReader, + make_byte_array_dictionary_reader, make_byte_array_reader, +}; +use crate::arrow::arrow_reader::DEFAULT_BATCH_SIZE; +use crate::arrow::arrow_reader::metrics::ArrowReaderMetrics; +use crate::arrow::schema::{ParquetField, ParquetFieldType, VirtualColumnType}; +use crate::basic::Type as PhysicalType; +use crate::data_type::{BoolType, DoubleType, FloatType, Int32Type, Int64Type, Int96Type}; +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::ParquetMetaData; +use crate::schema::types::{ColumnDescriptor, ColumnPath, Type}; + +/// Builder for [`CacheOptions`] +#[derive(Debug, Clone)] +pub struct CacheOptionsBuilder<'a> { + /// Projection mask to apply to the cache + pub projection_mask: &'a ProjectionMask, + /// Cache to use for storing row groups + pub cache: &'a Arc>, +} + +impl<'a> CacheOptionsBuilder<'a> { + /// create a new cache options builder + pub fn new(projection_mask: &'a ProjectionMask, cache: &'a Arc>) -> Self { + Self { + projection_mask, + cache, + } + } + + /// Return a new [`CacheOptions`] for producing (populating) the cache + pub fn producer(self) -> CacheOptions<'a> { + CacheOptions { + projection_mask: self.projection_mask, + cache: self.cache, + role: CacheRole::Producer, + } + } + + /// return a new [`CacheOptions`] for consuming (reading) the cache + pub fn consumer(self) -> CacheOptions<'a> { + CacheOptions { + projection_mask: self.projection_mask, + cache: self.cache, + role: CacheRole::Consumer, + } + } +} + +/// Cache options containing projection mask, cache, and role +#[derive(Clone)] +pub struct CacheOptions<'a> { + pub projection_mask: &'a ProjectionMask, + pub cache: &'a Arc>, + pub role: CacheRole, +} + +/// Builds [`ArrayReader`]s from parquet schema, projection mask, and RowGroups reader +pub struct ArrayReaderBuilder<'a> { + /// Source of row group data + row_groups: &'a dyn RowGroups, + /// Optional cache options for the array reader + cache_options: Option<&'a CacheOptions<'a>>, + /// Parquet metadata for computing virtual column values + parquet_metadata: Option<&'a ParquetMetaData>, + /// metrics + metrics: &'a ArrowReaderMetrics, + /// Batch size hint for pre-allocating internal buffers (see [`Self::with_batch_size`]) + batch_size: usize, +} + +/// Arguments threaded through the recursive `build_*_reader` calls. +/// +/// Bundling the per-field arguments into a single struct keeps the recursive +/// builder signatures small and provides one documented place to add new +/// per-reader options in the future. +#[derive(Clone, Copy)] +struct ReaderArgs<'a> { + /// The parquet field the output array corresponds to. + field: &'a ParquetField, + /// Which leaf columns are being read. + mask: &'a ProjectionMask, + /// Definition-level threshold at or above which the nearest enclosing + /// list/map keeps a leaf value. Used to push selective null padding down + /// into leaf readers so they emit compact child arrays. `None` at the top + /// level (no enclosing list/map), which disables selective padding. + padding_threshold: Option, +} + +impl<'a> ReaderArgs<'a> { + /// Returns a copy of these arguments pointing at `field`. + /// + /// Used when recursing from a field into one of its children. + fn with_field(self, field: &'a ParquetField) -> Self { + Self { field, ..self } + } + + /// Returns a copy of these arguments with a different `padding_threshold`. + /// + /// Used by list/map readers to set the threshold their children should + /// use, without disturbing the threshold applied to the list/map itself. + fn with_padding_threshold(self, padding_threshold: Option) -> Self { + Self { + padding_threshold, + ..self + } + } +} + +impl<'a> ArrayReaderBuilder<'a> { + /// Create a new `ArrayReaderBuilder` + pub fn new(row_groups: &'a dyn RowGroups, metrics: &'a ArrowReaderMetrics) -> Self { + Self { + row_groups, + cache_options: None, + parquet_metadata: None, + metrics, + batch_size: DEFAULT_BATCH_SIZE, + } + } + + /// Set the batch size hint used to pre-allocate internal buffers. + /// + /// This is a hint, not a hard capacity contract. Readers without + /// selective null padding reserve their value buffers up front from this + /// size to avoid reallocations when reading the first batch. Readers that + /// emit compact (selectively padded) child arrays instead size their + /// buffers from the number of values actually decoded, and may allocate + /// less than the batch size. + pub fn with_batch_size(mut self, batch_size: usize) -> Self { + self.batch_size = batch_size; + self + } + + /// Add cache options to the builder + pub fn with_cache_options(mut self, cache_options: Option<&'a CacheOptions<'a>>) -> Self { + self.cache_options = cache_options; + self + } + + /// Add parquet metadata to the builder for computing virtual column values + pub fn with_parquet_metadata(mut self, parquet_metadata: &'a ParquetMetaData) -> Self { + self.parquet_metadata = Some(parquet_metadata); + self + } + + /// Create [`ArrayReader`] from parquet schema, projection mask, and parquet file reader. + pub fn build_array_reader( + &self, + field: Option<&ParquetField>, + mask: &ProjectionMask, + ) -> Result> { + let reader = field + .and_then(|field| { + self.build_reader(ReaderArgs { + field, + mask, + padding_threshold: None, + }) + .transpose() + }) + .transpose()? + .unwrap_or_else(|| make_empty_array_reader(self.num_rows())); + + Ok(reader) + } + + /// Return the total number of rows + fn num_rows(&self) -> usize { + self.row_groups.num_rows() + } + + fn build_reader(&self, args: ReaderArgs<'_>) -> Result>> { + match args.field.field_type { + ParquetFieldType::Primitive { col_idx, .. } => { + let Some(reader) = self.build_primitive_reader(args)? else { + return Ok(None); + }; + let Some(cache_options) = self.cache_options.as_ref() else { + return Ok(Some(reader)); + }; + + if cache_options.projection_mask.leaf_included(col_idx) { + Ok(Some(Box::new(CachedArrayReader::new( + reader, + Arc::clone(cache_options.cache), + col_idx, + cache_options.role, + self.metrics.clone(), // cheap clone + )))) + } else { + Ok(Some(reader)) + } + } + ParquetFieldType::Virtual(virtual_type) => { + // Virtual columns don't have data in the parquet file + // They need to be built by specialized readers + match virtual_type { + VirtualColumnType::RowNumber => Ok(Some(self.build_row_number_reader()?)), + VirtualColumnType::RowGroupIndex => { + Ok(Some(self.build_row_group_index_reader()?)) + } + } + } + ParquetFieldType::Group { .. } => match &args.field.arrow_type { + DataType::Map(_, _) => self.build_map_reader(args), + DataType::Struct(_) => self.build_struct_reader(args), + DataType::List(_) + | DataType::LargeList(_) + | DataType::ListView(_) + | DataType::LargeListView(_) => self.build_list_reader(args), + DataType::FixedSizeList(_, _) => self.build_fixed_size_list_reader(args), + d => unimplemented!("reading group type {} not implemented", d), + }, + } + } + + fn build_row_number_reader(&self) -> Result> { + let parquet_metadata = self.parquet_metadata.ok_or_else(|| { + ParquetError::General( + "ParquetMetaData is required to read virtual row number columns.".to_string(), + ) + })?; + Ok(Box::new(RowNumberReader::try_new( + parquet_metadata, + self.row_groups.row_groups(), + )?)) + } + + fn build_row_group_index_reader(&self) -> Result> { + let parquet_metadata = self.parquet_metadata.ok_or_else(|| { + ParquetError::General( + "ParquetMetaData is required to read virtual row group index columns.".to_string(), + ) + })?; + Ok(Box::new(RowGroupIndexReader::try_new( + parquet_metadata, + self.row_groups.row_groups(), + )?)) + } + + /// Build array reader for map type. + fn build_map_reader(&self, args: ReaderArgs<'_>) -> Result>> { + let field = args.field; + let padding_threshold = args.padding_threshold; + let children = field.children().unwrap(); + assert_eq!(children.len(), 2); + + // The map is internally List(Struct(key, value)). The key/value + // readers are children of the struct inside the list, so their + // padding threshold is the map's (= list's) def_level. + let child_args = args.with_padding_threshold(Some(field.def_level)); + let key_reader = self.build_reader(child_args.with_field(&children[0]))?; + let value_reader = self.build_reader(child_args.with_field(&children[1]))?; + + match (key_reader, value_reader) { + (Some(key_reader), Some(value_reader)) => { + // Need to retrieve underlying data type to handle projection + let key_type = key_reader.get_data_type().clone(); + let value_type = value_reader.get_data_type().clone(); + + let data_type = match &field.arrow_type { + DataType::Map(map_field, is_sorted) => match map_field.data_type() { + DataType::Struct(fields) => { + assert_eq!(fields.len(), 2); + let struct_field = map_field.as_ref().clone().with_data_type( + DataType::Struct(Fields::from(vec![ + fields[0].as_ref().clone().with_data_type(key_type), + fields[1].as_ref().clone().with_data_type(value_type), + ])), + ); + DataType::Map(Arc::new(struct_field), *is_sorted) + } + _ => unreachable!(), + }, + _ => unreachable!(), + }; + + Ok(Some(Box::new(MapArrayReader::new( + key_reader, + value_reader, + data_type, + field.def_level, + field.rep_level, + field.nullable, + padding_threshold, + )))) + } + (None, None) => Ok(None), + _ => Err(general_err!( + "partial projection of MapArray is not supported" + )), + } + } + + /// Build array reader for list type. + fn build_list_reader(&self, args: ReaderArgs<'_>) -> Result>> { + let field = args.field; + let padding_threshold = args.padding_threshold; + let children = field.children().unwrap(); + assert_eq!(children.len(), 1); + + // Build the child with this list's def_level as its padding threshold. + // Pass the received padding_threshold as this list's parent_threshold. + let child_args = args.with_padding_threshold(Some(field.def_level)); + let reader = match self.build_reader(child_args.with_field(&children[0]))? { + Some(item_reader) => { + // Need to retrieve underlying data type to handle projection + let item_type = item_reader.get_data_type().clone(); + let reader: Box = match &field.arrow_type { + DataType::List(f) => { + let data_type = + DataType::List(Arc::new(f.as_ref().clone().with_data_type(item_type))); + Box::new(ListArrayReader::::new( + item_reader, + data_type, + field.def_level, + field.rep_level, + field.nullable, + padding_threshold, + )) + } + DataType::LargeList(f) => { + let data_type = DataType::LargeList(Arc::new( + f.as_ref().clone().with_data_type(item_type), + )); + Box::new(ListArrayReader::::new( + item_reader, + data_type, + field.def_level, + field.rep_level, + field.nullable, + padding_threshold, + )) + } + DataType::ListView(f) => { + let data_type = DataType::ListView(Arc::new( + f.as_ref().clone().with_data_type(item_type), + )); + Box::new(ListViewArrayReader::::new( + item_reader, + data_type, + field.def_level, + field.rep_level, + field.nullable, + padding_threshold, + )) + } + DataType::LargeListView(f) => { + let data_type = DataType::LargeListView(Arc::new( + f.as_ref().clone().with_data_type(item_type), + )); + Box::new(ListViewArrayReader::::new( + item_reader, + data_type, + field.def_level, + field.rep_level, + field.nullable, + padding_threshold, + )) + } + _ => unreachable!(), + }; + Some(reader) + } + None => None, + }; + Ok(reader) + } + + /// Build array reader for fixed-size list type. + fn build_fixed_size_list_reader( + &self, + args: ReaderArgs<'_>, + ) -> Result>> { + let field = args.field; + let padding_threshold = args.padding_threshold; + let children = field.children().unwrap(); + assert_eq!(children.len(), 1); + + // Build the child with this list's def_level as its padding threshold. + let child_args = args.with_padding_threshold(Some(field.def_level)); + let reader = match self.build_reader(child_args.with_field(&children[0]))? { + Some(item_reader) => { + let item_type = item_reader.get_data_type().clone(); + let reader = match &field.arrow_type { + &DataType::FixedSizeList(ref f, size) => { + let data_type = DataType::FixedSizeList( + Arc::new(f.as_ref().clone().with_data_type(item_type)), + size, + ); + + Box::new(FixedSizeListArrayReader::new( + item_reader, + size as usize, + data_type, + field.def_level, + field.rep_level, + field.nullable, + padding_threshold, + )) as _ + } + _ => unimplemented!(), + }; + Some(reader) + } + None => None, + }; + Ok(reader) + } + + /// Creates primitive array reader for each primitive type. + fn build_primitive_reader(&self, args: ReaderArgs<'_>) -> Result>> { + let field = args.field; + let padding_threshold = args.padding_threshold; + let (col_idx, primitive_type) = match &field.field_type { + ParquetFieldType::Primitive { + col_idx, + primitive_type, + } => match primitive_type.as_ref() { + Type::PrimitiveType { .. } => (*col_idx, primitive_type.clone()), + Type::GroupType { .. } => unreachable!(), + }, + _ => unreachable!(), + }; + + if !args.mask.leaf_included(col_idx) { + return Ok(None); + } + + let physical_type = primitive_type.get_physical_type(); + + // We don't track the column path in ParquetField as it adds a potential source + // of bugs when the arrow mapping converts more than one level in the parquet + // schema into a single arrow field. + // + // None of the readers actually use this field, but it is required for this type, + // so just stick a placeholder in + let column_desc = Arc::new(ColumnDescriptor::new( + primitive_type, + field.def_level, + field.rep_level, + ColumnPath::new(vec![]), + )); + + let page_iterator = self.row_groups.column_chunks(col_idx)?; + let arrow_type = Some(field.arrow_type.clone()); + + // LogicalType::Unknown maps to DataType::Null. In the past it has been assumed + // that only INT32 can have this annotation, but this is not required by the Parquet + // specification. Since this can only annotate an entirely null column, the data type + // used for the NullArrayReader should be irrelevant. It's just needed to read the + // repetition and definition level data. + if matches!(arrow_type, Some(DataType::Null)) { + let reader = Box::new(NullArrayReader::::new( + page_iterator, + column_desc, + self.batch_size, + padding_threshold, + )?) as _; + return Ok(Some(reader)); + } + + let reader = match physical_type { + PhysicalType::BOOLEAN => Box::new(PrimitiveArrayReader::::new( + page_iterator, + column_desc, + arrow_type, + self.batch_size, + padding_threshold, + )?) as _, + PhysicalType::INT32 => Box::new(PrimitiveArrayReader::::new( + page_iterator, + column_desc, + arrow_type, + self.batch_size, + padding_threshold, + )?) as _, + PhysicalType::INT64 => Box::new(PrimitiveArrayReader::::new( + page_iterator, + column_desc, + arrow_type, + self.batch_size, + padding_threshold, + )?) as _, + PhysicalType::INT96 => Box::new(PrimitiveArrayReader::::new( + page_iterator, + column_desc, + arrow_type, + self.batch_size, + padding_threshold, + )?) as _, + PhysicalType::FLOAT => Box::new(PrimitiveArrayReader::::new( + page_iterator, + column_desc, + arrow_type, + self.batch_size, + padding_threshold, + )?) as _, + PhysicalType::DOUBLE => Box::new(PrimitiveArrayReader::::new( + page_iterator, + column_desc, + arrow_type, + self.batch_size, + padding_threshold, + )?) as _, + PhysicalType::BYTE_ARRAY => match arrow_type { + Some(DataType::Dictionary(_, _)) => make_byte_array_dictionary_reader( + page_iterator, + column_desc, + arrow_type, + self.batch_size, + padding_threshold, + )?, + Some(DataType::Utf8View | DataType::BinaryView) => make_byte_view_array_reader( + page_iterator, + column_desc, + arrow_type, + self.batch_size, + padding_threshold, + )?, + _ => make_byte_array_reader( + page_iterator, + column_desc, + arrow_type, + self.batch_size, + padding_threshold, + )?, + }, + PhysicalType::FIXED_LEN_BYTE_ARRAY => match arrow_type { + Some(DataType::Dictionary(_, _)) => make_byte_array_dictionary_reader( + page_iterator, + column_desc, + arrow_type, + self.batch_size, + padding_threshold, + )?, + _ => make_fixed_len_byte_array_reader( + page_iterator, + column_desc, + arrow_type, + self.batch_size, + padding_threshold, + )?, + }, + }; + Ok(Some(reader)) + } + + fn build_struct_reader(&self, args: ReaderArgs<'_>) -> Result>> { + let field = args.field; + let padding_threshold = args.padding_threshold; + let arrow_fields = match &field.arrow_type { + DataType::Struct(children) => children, + _ => unreachable!(), + }; + let children = field.children().unwrap(); + assert_eq!(arrow_fields.len(), children.len()); + + let mut readers = Vec::with_capacity(children.len()); + let mut builder = SchemaBuilder::with_capacity(children.len()); + + for (arrow, parquet) in arrow_fields.iter().zip(children) { + if let Some(reader) = self.build_reader(args.with_field(parquet))? { + // Need to retrieve underlying data type to handle projection + let child_type = reader.get_data_type().clone(); + builder.push(arrow.as_ref().clone().with_data_type(child_type)); + readers.push(reader); + } + } + + if readers.is_empty() { + return Ok(None); + } + + Ok(Some(Box::new(StructArrayReader::new( + DataType::Struct(builder.finish().fields), + readers, + field.def_level, + field.rep_level, + field.nullable, + padding_threshold, + )))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::schema::parquet_to_arrow_schema_and_fields; + use crate::arrow::schema::virtual_type::RowNumber; + use crate::file::reader::{FileReader, SerializedFileReader}; + use crate::util::test_common::file_util::get_test_file; + use arrow::datatypes::Field; + use std::sync::Arc; + + #[test] + fn test_create_array_reader() { + let file = get_test_file("nulls.snappy.parquet"); + let file_reader: Arc = Arc::new(SerializedFileReader::new(file).unwrap()); + + let file_metadata = file_reader.metadata().file_metadata(); + let mask = ProjectionMask::leaves(file_metadata.schema_descr(), [0]); + let (_, fields) = parquet_to_arrow_schema_and_fields( + file_metadata.schema_descr(), + ProjectionMask::all(), + file_metadata.key_value_metadata(), + &[], + ) + .unwrap(); + + let metrics = ArrowReaderMetrics::disabled(); + let array_reader = ArrayReaderBuilder::new(&file_reader, &metrics) + .with_batch_size(DEFAULT_BATCH_SIZE) + .build_array_reader(fields.as_ref(), &mask) + .unwrap(); + + // Create arrow types + let arrow_type = DataType::Struct(Fields::from(vec![Field::new( + "b_struct", + DataType::Struct(vec![Field::new("b_c_int", DataType::Int32, true)].into()), + true, + )])); + + assert_eq!(array_reader.get_data_type(), &arrow_type); + } + + #[test] + fn test_create_array_reader_with_row_numbers() { + let file = get_test_file("nulls.snappy.parquet"); + let file_reader: Arc = Arc::new(SerializedFileReader::new(file).unwrap()); + + let file_metadata = file_reader.metadata().file_metadata(); + let mask = ProjectionMask::leaves(file_metadata.schema_descr(), [0]); + let row_number_field = Arc::new( + Field::new("row_number", DataType::Int64, false).with_extension_type(RowNumber), + ); + let (_, fields) = parquet_to_arrow_schema_and_fields( + file_metadata.schema_descr(), + ProjectionMask::all(), + file_metadata.key_value_metadata(), + std::slice::from_ref(&row_number_field), + ) + .unwrap(); + + let metrics = ArrowReaderMetrics::disabled(); + let array_reader = ArrayReaderBuilder::new(&file_reader, &metrics) + .with_batch_size(DEFAULT_BATCH_SIZE) + .with_parquet_metadata(file_reader.metadata()) + .build_array_reader(fields.as_ref(), &mask) + .unwrap(); + + // Create arrow types + let arrow_type = DataType::Struct(Fields::from(vec![ + Field::new( + "b_struct", + DataType::Struct(vec![Field::new("b_c_int", DataType::Int32, true)].into()), + true, + ), + (*row_number_field).clone(), + ])); + + assert_eq!(array_reader.get_data_type(), &arrow_type); + } +} diff --git a/vendor/parquet/src/arrow/array_reader/byte_array.rs b/vendor/parquet/src/arrow/array_reader/byte_array.rs new file mode 100644 index 0000000..4fc32dd --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/byte_array.rs @@ -0,0 +1,756 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::array_reader::{ArrayReader, read_records, skip_records}; +use crate::arrow::buffer::bit_util::sign_extend_be; +use crate::arrow::buffer::offset_buffer::OffsetBuffer; +use crate::arrow::decoder::{DeltaByteArrayDecoder, DictIndexDecoder}; +use crate::arrow::record_reader::GenericRecordReader; +use crate::arrow::schema::parquet_to_arrow_field; +use crate::basic::{ConvertedType, Encoding}; +use crate::column::page::PageIterator; +use crate::column::reader::decoder::ColumnValueDecoder; +use crate::data_type::Int32Type; +use crate::encodings::decoding::{Decoder, DeltaBitPackDecoder}; +use crate::errors::{ParquetError, Result}; +use crate::schema::types::ColumnDescPtr; +use arrow_array::{ + Array, ArrayRef, BinaryArray, Decimal128Array, Decimal256Array, OffsetSizeTrait, +}; +use arrow_buffer::i256; +use arrow_schema::DataType as ArrowType; +use bytes::Bytes; +use std::any::Any; +use std::sync::Arc; + +/// Returns an [`ArrayReader`] that decodes the provided byte array column +pub fn make_byte_array_reader( + pages: Box, + column_desc: ColumnDescPtr, + arrow_type: Option, + batch_size: usize, + padding_threshold: Option, +) -> Result> { + // Check if Arrow type is specified, else create it from Parquet type + let data_type = match arrow_type { + Some(t) => t, + None => parquet_to_arrow_field(column_desc.as_ref())? + .data_type() + .clone(), + }; + + match data_type { + ArrowType::Binary + | ArrowType::Utf8 + | ArrowType::Decimal128(_, _) + | ArrowType::Decimal256(_, _) => { + let mut reader = GenericRecordReader::new(column_desc, batch_size); + if let Some(threshold) = padding_threshold { + reader.set_padding_threshold(threshold); + } + Ok(Box::new(ByteArrayReader::::new( + pages, data_type, reader, + ))) + } + ArrowType::LargeUtf8 | ArrowType::LargeBinary => { + let mut reader = GenericRecordReader::new(column_desc, batch_size); + if let Some(threshold) = padding_threshold { + reader.set_padding_threshold(threshold); + } + Ok(Box::new(ByteArrayReader::::new( + pages, data_type, reader, + ))) + } + _ => Err(general_err!( + "invalid data type for byte array reader - {}", + data_type + )), + } +} + +/// An [`ArrayReader`] for variable length byte arrays +struct ByteArrayReader { + data_type: ArrowType, + pages: Box, + def_levels_buffer: Option>, + rep_levels_buffer: Option>, + record_reader: GenericRecordReader, ByteArrayColumnValueDecoder>, +} + +impl ByteArrayReader { + fn new( + pages: Box, + data_type: ArrowType, + record_reader: GenericRecordReader, ByteArrayColumnValueDecoder>, + ) -> Self { + Self { + data_type, + pages, + def_levels_buffer: None, + rep_levels_buffer: None, + record_reader, + } + } +} + +impl ArrayReader for ByteArrayReader { + fn as_any(&self) -> &dyn Any { + self + } + + fn get_data_type(&self) -> &ArrowType { + &self.data_type + } + + fn read_records(&mut self, batch_size: usize) -> Result { + read_records(&mut self.record_reader, self.pages.as_mut(), batch_size) + } + + fn consume_batch(&mut self) -> Result { + let buffer = self.record_reader.consume_record_data(); + let null_buffer = self.record_reader.consume_compact_bitmap(); + self.def_levels_buffer = self.record_reader.consume_def_levels(); + self.rep_levels_buffer = self.record_reader.consume_rep_levels(); + self.record_reader.reset(); + + let array: ArrayRef = match self.data_type { + // Apply conversion to all elements regardless of null slots as the conversions + // are infallible. This improves performance by avoiding a branch in the inner + // loop (see docs for `PrimitiveArray::from_unary`). + ArrowType::Decimal128(p, s) => { + let array = buffer.into_array(null_buffer, ArrowType::Binary); + let binary = array.as_any().downcast_ref::().unwrap(); + // Null slots will have 0 length, so we need to check for that in the lambda + // or sign_extend_be will panic. + let decimal = Decimal128Array::from_unary(binary, |x| match x.len() { + 0 => i128::default(), + _ => i128::from_be_bytes(sign_extend_be(x)), + }) + .with_precision_and_scale(p, s)?; + Arc::new(decimal) + } + ArrowType::Decimal256(p, s) => { + let array = buffer.into_array(null_buffer, ArrowType::Binary); + let binary = array.as_any().downcast_ref::().unwrap(); + // Null slots will have 0 length, so we need to check for that in the lambda + // or sign_extend_be will panic. + let decimal = Decimal256Array::from_unary(binary, |x| match x.len() { + 0 => i256::default(), + _ => i256::from_be_bytes(sign_extend_be(x)), + }) + .with_precision_and_scale(p, s)?; + Arc::new(decimal) + } + _ => buffer.into_array(null_buffer, self.data_type.clone()), + }; + + Ok(array) + } + + fn skip_records(&mut self, num_records: usize) -> Result { + skip_records(&mut self.record_reader, self.pages.as_mut(), num_records) + } + + fn get_def_levels(&self) -> Option<&[i16]> { + self.def_levels_buffer.as_deref() + } + + fn get_rep_levels(&self) -> Option<&[i16]> { + self.rep_levels_buffer.as_deref() + } + + fn max_def_level(&self) -> i16 { + self.record_reader.max_def_level() + } +} + +/// A [`ColumnValueDecoder`] for variable length byte arrays +struct ByteArrayColumnValueDecoder { + dict: Option>, + decoder: Option, + validate_utf8: bool, +} + +impl ColumnValueDecoder for ByteArrayColumnValueDecoder { + type Buffer = OffsetBuffer; + + fn new(desc: &ColumnDescPtr) -> Self { + let validate_utf8 = desc.converted_type() == ConvertedType::UTF8; + Self { + dict: None, + decoder: None, + validate_utf8, + } + } + + fn set_dict( + &mut self, + buf: Bytes, + num_values: u32, + encoding: Encoding, + _is_sorted: bool, + ) -> Result<()> { + if !matches!( + encoding, + Encoding::PLAIN | Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY + ) { + return Err(nyi_err!( + "Invalid/Unsupported encoding type for dictionary: {}", + encoding + )); + } + + let mut buffer = OffsetBuffer::with_capacity(0); + let mut decoder = ByteArrayDecoderPlain::new( + buf, + num_values as usize, + Some(num_values as usize), + self.validate_utf8, + ); + decoder.read(&mut buffer, usize::MAX)?; + self.dict = Some(buffer); + Ok(()) + } + + fn set_data( + &mut self, + encoding: Encoding, + data: Bytes, + num_levels: usize, + num_values: Option, + ) -> Result<()> { + self.decoder = Some(ByteArrayDecoder::new( + encoding, + data, + num_levels, + num_values, + self.validate_utf8, + )?); + Ok(()) + } + + fn read(&mut self, out: &mut Self::Buffer, num_values: usize) -> Result { + let decoder = self + .decoder + .as_mut() + .ok_or_else(|| general_err!("no decoder set"))?; + + decoder.read(out, num_values, self.dict.as_ref()) + } + + fn skip_values(&mut self, num_values: usize) -> Result { + let decoder = self + .decoder + .as_mut() + .ok_or_else(|| general_err!("no decoder set"))?; + + decoder.skip(num_values, self.dict.as_ref()) + } +} + +/// A generic decoder from uncompressed parquet value data to [`OffsetBuffer`] +pub enum ByteArrayDecoder { + Plain(ByteArrayDecoderPlain), + Dictionary(ByteArrayDecoderDictionary), + DeltaLength(ByteArrayDecoderDeltaLength), + DeltaByteArray(ByteArrayDecoderDelta), +} + +impl ByteArrayDecoder { + pub fn new( + encoding: Encoding, + data: Bytes, + num_levels: usize, + num_values: Option, + validate_utf8: bool, + ) -> Result { + let decoder = match encoding { + Encoding::PLAIN => ByteArrayDecoder::Plain(ByteArrayDecoderPlain::new( + data, + num_levels, + num_values, + validate_utf8, + )), + Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY => ByteArrayDecoder::Dictionary( + ByteArrayDecoderDictionary::new(data, num_levels, num_values)?, + ), + Encoding::DELTA_LENGTH_BYTE_ARRAY => ByteArrayDecoder::DeltaLength( + ByteArrayDecoderDeltaLength::new(data, validate_utf8)?, + ), + Encoding::DELTA_BYTE_ARRAY => { + ByteArrayDecoder::DeltaByteArray(ByteArrayDecoderDelta::new(data, validate_utf8)?) + } + _ => { + return Err(general_err!( + "unsupported encoding for byte array: {}", + encoding + )); + } + }; + + Ok(decoder) + } + + /// Read up to `len` values to `out` with the optional dictionary + pub fn read( + &mut self, + out: &mut OffsetBuffer, + len: usize, + dict: Option<&OffsetBuffer>, + ) -> Result { + match self { + ByteArrayDecoder::Plain(d) => d.read(out, len), + ByteArrayDecoder::Dictionary(d) => { + let dict = + dict.ok_or_else(|| general_err!("missing dictionary page for column"))?; + + d.read(out, dict, len) + } + ByteArrayDecoder::DeltaLength(d) => d.read(out, len), + ByteArrayDecoder::DeltaByteArray(d) => d.read(out, len), + } + } + + /// Skip `len` values + pub fn skip( + &mut self, + len: usize, + dict: Option<&OffsetBuffer>, + ) -> Result { + match self { + ByteArrayDecoder::Plain(d) => d.skip(len), + ByteArrayDecoder::Dictionary(d) => { + let dict = + dict.ok_or_else(|| general_err!("missing dictionary page for column"))?; + + d.skip(dict, len) + } + ByteArrayDecoder::DeltaLength(d) => d.skip(len), + ByteArrayDecoder::DeltaByteArray(d) => d.skip(len), + } + } +} + +/// Decoder from [`Encoding::PLAIN`] data to [`OffsetBuffer`] +pub struct ByteArrayDecoderPlain { + buf: Bytes, + offset: usize, + validate_utf8: bool, + + /// This is a maximum as the null count is not always known, e.g. value data from + /// a v1 data page + max_remaining_values: usize, +} + +impl ByteArrayDecoderPlain { + pub fn new( + buf: Bytes, + num_levels: usize, + num_values: Option, + validate_utf8: bool, + ) -> Self { + Self { + buf, + validate_utf8, + offset: 0, + max_remaining_values: num_values.unwrap_or(num_levels), + } + } + + pub fn read( + &mut self, + output: &mut OffsetBuffer, + len: usize, + ) -> Result { + let initial_values_length = output.values.len(); + + let to_read = len.min(self.max_remaining_values); + output.offsets.reserve(to_read); + + let remaining_bytes = self.buf.len() - self.offset; + if remaining_bytes == 0 { + return Ok(0); + } + + let estimated_bytes = remaining_bytes + .checked_mul(to_read) + .map(|x| x / self.max_remaining_values) + .unwrap_or_default(); + + output.values.reserve(estimated_bytes); + + let mut read = 0; + + let buf = self.buf.as_ref(); + while self.offset < self.buf.len() && read != to_read { + if self.offset + 4 > buf.len() { + return Err(ParquetError::EOF("eof decoding byte array".into())); + } + let len_bytes: [u8; 4] = buf[self.offset..self.offset + 4].try_into().unwrap(); + let len = u32::from_le_bytes(len_bytes); + + let start_offset = self.offset + 4; + let end_offset = start_offset + len as usize; + if end_offset > buf.len() { + return Err(ParquetError::EOF("eof decoding byte array".into())); + } + + output.try_push(&buf[start_offset..end_offset], self.validate_utf8)?; + + self.offset = end_offset; + read += 1; + } + self.max_remaining_values -= to_read; + + if self.validate_utf8 { + output.check_valid_utf8(initial_values_length)?; + } + Ok(to_read) + } + + pub fn skip(&mut self, to_skip: usize) -> Result { + let to_skip = to_skip.min(self.max_remaining_values); + let mut skip = 0; + let buf = self.buf.as_ref(); + + while self.offset < self.buf.len() && skip != to_skip { + if self.offset + 4 > buf.len() { + return Err(ParquetError::EOF("eof decoding byte array".into())); + } + let len_bytes: [u8; 4] = buf[self.offset..self.offset + 4].try_into().unwrap(); + let len = u32::from_le_bytes(len_bytes) as usize; + skip += 1; + self.offset = self.offset + 4 + len; + } + self.max_remaining_values -= skip; + Ok(skip) + } +} + +/// Decoder from [`Encoding::DELTA_LENGTH_BYTE_ARRAY`] data to [`OffsetBuffer`] +pub struct ByteArrayDecoderDeltaLength { + lengths: Vec, + data: Bytes, + length_offset: usize, + data_offset: usize, + validate_utf8: bool, +} + +impl ByteArrayDecoderDeltaLength { + fn new(data: Bytes, validate_utf8: bool) -> Result { + let mut len_decoder = DeltaBitPackDecoder::::new(); + len_decoder.set_data(data.clone(), 0)?; + let values = len_decoder.values_left(); + + let mut lengths = vec![0; values]; + len_decoder.get(&mut lengths)?; + + let mut total_bytes = 0; + + for l in lengths.iter() { + if *l < 0 { + return Err(ParquetError::General( + "negative delta length byte array length".to_string(), + )); + } + total_bytes += *l as usize; + } + + if total_bytes + len_decoder.get_offset() > data.len() { + return Err(ParquetError::General( + "Insufficient delta length byte array bytes".to_string(), + )); + } + + Ok(Self { + lengths, + data, + validate_utf8, + length_offset: 0, + data_offset: len_decoder.get_offset(), + }) + } + + fn read( + &mut self, + output: &mut OffsetBuffer, + len: usize, + ) -> Result { + let initial_values_length = output.values.len(); + + let to_read = len.min(self.lengths.len() - self.length_offset); + let src_lengths = &self.lengths[self.length_offset..self.length_offset + to_read]; + let total_bytes: usize = src_lengths.iter().map(|x| *x as usize).sum(); + + // Reserve capacity for both offsets and values upfront + output.offsets.reserve(to_read); + output.values.reserve(total_bytes); + + // Delta length data is contiguous — copy all value bytes at once + let data_end = self.data_offset + total_bytes; + output + .values + .extend_from_slice(&self.data.as_ref()[self.data_offset..data_end]); + + // Compute and extend offsets in batch using extend + let base_offset = initial_values_length; + let mut running = base_offset; + output.offsets.extend(src_lengths.iter().map(|length| { + running += *length as usize; + I::from_usize(running).expect("index overflow decoding byte array") + })); + + self.data_offset = data_end; + self.length_offset += to_read; + + if self.validate_utf8 { + output.check_valid_utf8(initial_values_length)?; + } + Ok(to_read) + } + + fn skip(&mut self, to_skip: usize) -> Result { + let remain_values = self.lengths.len() - self.length_offset; + let to_skip = remain_values.min(to_skip); + + let src_lengths = &self.lengths[self.length_offset..self.length_offset + to_skip]; + let total_bytes: usize = src_lengths.iter().map(|x| *x as usize).sum(); + + self.data_offset += total_bytes; + self.length_offset += to_skip; + Ok(to_skip) + } +} + +/// Decoder from [`Encoding::DELTA_BYTE_ARRAY`] to [`OffsetBuffer`] +pub struct ByteArrayDecoderDelta { + decoder: DeltaByteArrayDecoder, + validate_utf8: bool, +} + +impl ByteArrayDecoderDelta { + fn new(data: Bytes, validate_utf8: bool) -> Result { + Ok(Self { + decoder: DeltaByteArrayDecoder::new(data)?, + validate_utf8, + }) + } + + fn read( + &mut self, + output: &mut OffsetBuffer, + len: usize, + ) -> Result { + let initial_values_length = output.values.len(); + output.offsets.reserve(len.min(self.decoder.remaining())); + + let read = self + .decoder + .read(len, |bytes| output.try_push(bytes, self.validate_utf8))?; + + if self.validate_utf8 { + output.check_valid_utf8(initial_values_length)?; + } + Ok(read) + } + + fn skip(&mut self, to_skip: usize) -> Result { + self.decoder.skip(to_skip) + } +} + +/// Decoder from [`Encoding::RLE_DICTIONARY`] to [`OffsetBuffer`] +pub struct ByteArrayDecoderDictionary { + decoder: DictIndexDecoder, +} + +impl ByteArrayDecoderDictionary { + fn new(data: Bytes, num_levels: usize, num_values: Option) -> Result { + Ok(Self { + decoder: DictIndexDecoder::new(data, num_levels, num_values)?, + }) + } + + fn read( + &mut self, + output: &mut OffsetBuffer, + dict: &OffsetBuffer, + len: usize, + ) -> Result { + // All data must be NULL + if dict.is_empty() { + return Ok(0); + } + + // Pre-reserve offsets capacity to avoid per-chunk reallocation + output.offsets.reserve(len); + + self.decoder.read(len, |keys| { + output.extend_from_dictionary(keys, dict.offsets.as_slice(), dict.values.as_slice()) + }) + } + + fn skip( + &mut self, + dict: &OffsetBuffer, + to_skip: usize, + ) -> Result { + // All data must be NULL + if dict.is_empty() { + return Ok(0); + } + + self.decoder.skip(to_skip) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::array_reader::test_util::{byte_array_all_encodings, utf8_column}; + use crate::arrow::record_reader::buffer::ValuesBuffer; + use arrow_array::{Array, StringArray}; + use arrow_buffer::Buffer; + + #[test] + fn test_byte_array_decoder() { + let (pages, encoded_dictionary) = + byte_array_all_encodings(vec!["hello", "world", "a", "b"]); + + let column_desc = utf8_column(); + let mut decoder = ByteArrayColumnValueDecoder::new(&column_desc); + + decoder + .set_dict(encoded_dictionary, 4, Encoding::RLE_DICTIONARY, false) + .unwrap(); + + for (encoding, page) in pages { + let mut output = OffsetBuffer::::with_capacity(0); + decoder.set_data(encoding, page, 4, Some(4)).unwrap(); + + assert_eq!(decoder.read(&mut output, 1).unwrap(), 1); + + assert_eq!(output.values.as_slice(), "hello".as_bytes()); + assert_eq!(output.offsets.as_slice(), &[0, 5]); + + assert_eq!(decoder.read(&mut output, 1).unwrap(), 1); + assert_eq!(output.values.as_slice(), "helloworld".as_bytes()); + assert_eq!(output.offsets.as_slice(), &[0, 5, 10]); + + assert_eq!(decoder.read(&mut output, 2).unwrap(), 2); + assert_eq!(output.values.as_slice(), "helloworldab".as_bytes()); + assert_eq!(output.offsets.as_slice(), &[0, 5, 10, 11, 12]); + + assert_eq!(decoder.read(&mut output, 4).unwrap(), 0); + + let valid = [false, false, true, true, false, true, true, false, false]; + let valid_buffer = Buffer::from_iter(valid.iter().cloned()); + + output + .pad_nulls(0, 4, valid.len(), valid_buffer.as_slice()) + .unwrap(); + let array = output.into_array(Some(valid_buffer), ArrowType::Utf8); + let strings = array.as_any().downcast_ref::().unwrap(); + + assert_eq!( + strings.iter().collect::>(), + vec![ + None, + None, + Some("hello"), + Some("world"), + None, + Some("a"), + Some("b"), + None, + None, + ] + ); + } + } + + #[test] + fn test_byte_array_decoder_skip() { + let (pages, encoded_dictionary) = + byte_array_all_encodings(vec!["hello", "world", "a", "b"]); + + let column_desc = utf8_column(); + let mut decoder = ByteArrayColumnValueDecoder::new(&column_desc); + + decoder + .set_dict(encoded_dictionary, 4, Encoding::RLE_DICTIONARY, false) + .unwrap(); + + for (encoding, page) in pages { + let mut output = OffsetBuffer::::with_capacity(0); + decoder.set_data(encoding, page, 4, Some(4)).unwrap(); + + assert_eq!(decoder.read(&mut output, 1).unwrap(), 1); + + assert_eq!(output.values.as_slice(), "hello".as_bytes()); + assert_eq!(output.offsets.as_slice(), &[0, 5]); + + assert_eq!(decoder.skip_values(1).unwrap(), 1); + assert_eq!(decoder.skip_values(1).unwrap(), 1); + + assert_eq!(decoder.read(&mut output, 1).unwrap(), 1); + assert_eq!(output.values.as_slice(), "hellob".as_bytes()); + assert_eq!(output.offsets.as_slice(), &[0, 5, 6]); + + assert_eq!(decoder.read(&mut output, 4).unwrap(), 0); + + let valid = [false, false, true, true, false, false]; + let valid_buffer = Buffer::from_iter(valid.iter().cloned()); + + output + .pad_nulls(0, 2, valid.len(), valid_buffer.as_slice()) + .unwrap(); + let array = output.into_array(Some(valid_buffer), ArrowType::Utf8); + let strings = array.as_any().downcast_ref::().unwrap(); + + assert_eq!( + strings.iter().collect::>(), + vec![None, None, Some("hello"), Some("b"), None, None,] + ); + } + } + + #[test] + fn test_byte_array_decoder_nulls() { + let (pages, encoded_dictionary) = byte_array_all_encodings(Vec::<&str>::new()); + + let column_desc = utf8_column(); + let mut decoder = ByteArrayColumnValueDecoder::new(&column_desc); + + decoder + .set_dict(encoded_dictionary, 4, Encoding::RLE_DICTIONARY, false) + .unwrap(); + + // test nulls read + for (encoding, page) in pages.clone() { + let mut output = OffsetBuffer::::with_capacity(0); + decoder.set_data(encoding, page, 4, None).unwrap(); + assert_eq!(decoder.read(&mut output, 1024).unwrap(), 0); + } + + // test nulls skip + for (encoding, page) in pages { + decoder.set_data(encoding, page, 4, None).unwrap(); + assert_eq!(decoder.skip_values(1024).unwrap(), 0); + } + } +} diff --git a/vendor/parquet/src/arrow/array_reader/byte_array_dictionary.rs b/vendor/parquet/src/arrow/array_reader/byte_array_dictionary.rs new file mode 100644 index 0000000..65f2d9c --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/byte_array_dictionary.rs @@ -0,0 +1,691 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::any::Any; +use std::marker::PhantomData; + +use arrow_array::{Array, ArrayRef, OffsetSizeTrait, new_empty_array}; +use arrow_buffer::ArrowNativeType; +use arrow_schema::DataType as ArrowType; +use bytes::Bytes; + +use crate::arrow::array_reader::byte_array::{ByteArrayDecoder, ByteArrayDecoderPlain}; +use crate::arrow::array_reader::{ArrayReader, read_records, skip_records}; +use crate::arrow::buffer::{dictionary_buffer::DictionaryBuffer, offset_buffer::OffsetBuffer}; +use crate::arrow::record_reader::GenericRecordReader; +use crate::arrow::schema::parquet_to_arrow_field; +use crate::basic::{ConvertedType, Encoding}; +use crate::column::page::PageIterator; +use crate::column::reader::decoder::ColumnValueDecoder; +use crate::encodings::rle::RleDecoder; +use crate::errors::{ParquetError, Result}; +use crate::schema::types::ColumnDescPtr; +use crate::util::bit_util::FromBitpacked; + +/// A macro to reduce verbosity of [`make_byte_array_dictionary_reader`] +macro_rules! make_reader { + ( + ($pages:expr, $column_desc:expr, $data_type:expr, $batch_size:expr, $padding_threshold:expr) => match ($k:expr, $v:expr) { + $(($key_arrow:pat, $value_arrow:pat) => ($key_type:ty, $value_type:ty),)+ + } + ) => { + match ($k, $v) { + $( + ($key_arrow, $value_arrow) => { + let mut reader = GenericRecordReader::new($column_desc, $batch_size); + if let Some(threshold) = $padding_threshold { + reader.set_padding_threshold(threshold); + } + Ok(Box::new(ByteArrayDictionaryReader::<$key_type, $value_type>::new( + $pages, $data_type, reader, + ))) + } + )+ + _ => Err(general_err!( + "unsupported data type for byte array dictionary reader - {}", + $data_type + )), + } + } +} + +/// Returns an [`ArrayReader`] that decodes the provided byte array column +/// +/// This will attempt to preserve any dictionary encoding present in the parquet data +/// +/// It will be unable to preserve the dictionary encoding if: +/// +/// * A single read spans across multiple column chunks +/// * A column chunk contains non-dictionary encoded pages +/// +/// It is therefore recommended that if `pages` contains data from multiple column chunks, +/// that the read batch size used is a divisor of the row group size +pub fn make_byte_array_dictionary_reader( + pages: Box, + column_desc: ColumnDescPtr, + arrow_type: Option, + batch_size: usize, + padding_threshold: Option, +) -> Result> { + // Check if Arrow type is specified, else create it from Parquet type + let data_type = match arrow_type { + Some(t) => t, + None => parquet_to_arrow_field(column_desc.as_ref())? + .data_type() + .clone(), + }; + + match &data_type { + ArrowType::Dictionary(key_type, value_type) => { + make_reader! { + (pages, column_desc, data_type, batch_size, padding_threshold) => match (key_type.as_ref(), value_type.as_ref()) { + (ArrowType::UInt8, ArrowType::Binary | ArrowType::Utf8 | ArrowType::FixedSizeBinary(_)) => (u8, i32), + (ArrowType::UInt8, ArrowType::LargeBinary | ArrowType::LargeUtf8) => (u8, i64), + (ArrowType::Int8, ArrowType::Binary | ArrowType::Utf8 | ArrowType::FixedSizeBinary(_)) => (i8, i32), + (ArrowType::Int8, ArrowType::LargeBinary | ArrowType::LargeUtf8) => (i8, i64), + (ArrowType::UInt16, ArrowType::Binary | ArrowType::Utf8 | ArrowType::FixedSizeBinary(_)) => (u16, i32), + (ArrowType::UInt16, ArrowType::LargeBinary | ArrowType::LargeUtf8) => (u16, i64), + (ArrowType::Int16, ArrowType::Binary | ArrowType::Utf8 | ArrowType::FixedSizeBinary(_)) => (i16, i32), + (ArrowType::Int16, ArrowType::LargeBinary | ArrowType::LargeUtf8) => (i16, i64), + (ArrowType::UInt32, ArrowType::Binary | ArrowType::Utf8 | ArrowType::FixedSizeBinary(_)) => (u32, i32), + (ArrowType::UInt32, ArrowType::LargeBinary | ArrowType::LargeUtf8) => (u32, i64), + (ArrowType::Int32, ArrowType::Binary | ArrowType::Utf8 | ArrowType::FixedSizeBinary(_)) => (i32, i32), + (ArrowType::Int32, ArrowType::LargeBinary | ArrowType::LargeUtf8) => (i32, i64), + (ArrowType::UInt64, ArrowType::Binary | ArrowType::Utf8 | ArrowType::FixedSizeBinary(_)) => (u64, i32), + (ArrowType::UInt64, ArrowType::LargeBinary | ArrowType::LargeUtf8) => (u64, i64), + (ArrowType::Int64, ArrowType::Binary | ArrowType::Utf8 | ArrowType::FixedSizeBinary(_)) => (i64, i32), + (ArrowType::Int64, ArrowType::LargeBinary | ArrowType::LargeUtf8) => (i64, i64), + } + } + } + _ => Err(general_err!( + "invalid non-dictionary data type for byte array dictionary reader - {}", + data_type + )), + } +} + +/// An [`ArrayReader`] for dictionary encoded variable length byte arrays +/// +/// Will attempt to preserve any dictionary encoding present in the parquet data +struct ByteArrayDictionaryReader { + data_type: ArrowType, + pages: Box, + def_levels_buffer: Option>, + rep_levels_buffer: Option>, + record_reader: GenericRecordReader, DictionaryDecoder>, +} + +impl ByteArrayDictionaryReader +where + K: FromBitpacked + Ord + ArrowNativeType, + V: OffsetSizeTrait, +{ + fn new( + pages: Box, + data_type: ArrowType, + record_reader: GenericRecordReader, DictionaryDecoder>, + ) -> Self { + Self { + data_type, + pages, + def_levels_buffer: None, + rep_levels_buffer: None, + record_reader, + } + } +} + +impl ArrayReader for ByteArrayDictionaryReader +where + K: FromBitpacked + Ord + ArrowNativeType, + V: OffsetSizeTrait, +{ + fn as_any(&self) -> &dyn Any { + self + } + + fn get_data_type(&self) -> &ArrowType { + &self.data_type + } + + fn read_records(&mut self, batch_size: usize) -> Result { + read_records(&mut self.record_reader, self.pages.as_mut(), batch_size) + } + + fn consume_batch(&mut self) -> Result { + // advance the def & rep level buffers + self.def_levels_buffer = self.record_reader.consume_def_levels(); + self.rep_levels_buffer = self.record_reader.consume_rep_levels(); + + if self.record_reader.num_values() == 0 { + // once the record_reader has been consumed, we've replaced its values with the default + // variant of DictionaryBuffer (Offset). If `consume_batch` then gets called again, we + // avoid using the wrong variant of the buffer by returning empty array. + return Ok(new_empty_array(&self.data_type)); + } + + let buffer = self.record_reader.consume_record_data(); + let null_buffer = self.record_reader.consume_compact_bitmap(); + let array = buffer.into_array(null_buffer, &self.data_type)?; + self.record_reader.reset(); + + Ok(array) + } + + fn skip_records(&mut self, num_records: usize) -> Result { + skip_records(&mut self.record_reader, self.pages.as_mut(), num_records) + } + + fn get_def_levels(&self) -> Option<&[i16]> { + self.def_levels_buffer.as_deref() + } + + fn get_rep_levels(&self) -> Option<&[i16]> { + self.rep_levels_buffer.as_deref() + } + + fn max_def_level(&self) -> i16 { + self.record_reader.max_def_level() + } +} + +/// If the data is dictionary encoded decode the key data directly, so that the dictionary +/// encoding can be preserved. Otherwise fallback to decoding using [`ByteArrayDecoder`] +/// and compute a fresh dictionary in [`ByteArrayDictionaryReader::next_batch`] +enum MaybeDictionaryDecoder { + Dict { + decoder: RleDecoder, + /// This is a maximum as the null count is not always known, e.g. value data from + /// a v1 data page + max_remaining_values: usize, + }, + Fallback(ByteArrayDecoder), +} + +/// A [`ColumnValueDecoder`] for dictionary encoded variable length byte arrays +struct DictionaryDecoder { + /// The current dictionary + dict: Option, + + /// Dictionary decoder + decoder: Option, + + validate_utf8: bool, + + value_type: ArrowType, + + phantom: PhantomData<(K, V)>, +} + +impl ColumnValueDecoder for DictionaryDecoder +where + K: FromBitpacked + Ord + ArrowNativeType, + V: OffsetSizeTrait, +{ + type Buffer = DictionaryBuffer; + + fn new(col: &ColumnDescPtr) -> Self { + let validate_utf8 = col.converted_type() == ConvertedType::UTF8; + + let value_type = match (V::IS_LARGE, col.converted_type() == ConvertedType::UTF8) { + (true, true) => ArrowType::LargeUtf8, + (true, false) => ArrowType::LargeBinary, + (false, true) => ArrowType::Utf8, + (false, false) => ArrowType::Binary, + }; + + Self { + dict: None, + decoder: None, + validate_utf8, + value_type, + phantom: Default::default(), + } + } + + fn set_dict( + &mut self, + buf: Bytes, + num_values: u32, + encoding: Encoding, + _is_sorted: bool, + ) -> Result<()> { + if !matches!( + encoding, + Encoding::PLAIN | Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY + ) { + return Err(nyi_err!( + "Invalid/Unsupported encoding type for dictionary: {}", + encoding + )); + } + + if K::from_usize(num_values as usize).is_none() { + return Err(general_err!("dictionary too large for index type")); + } + + let len = num_values as usize; + let mut buffer = OffsetBuffer::::with_capacity(0); + let mut decoder = ByteArrayDecoderPlain::new(buf, len, Some(len), self.validate_utf8); + decoder.read(&mut buffer, usize::MAX)?; + + let array = buffer.into_array(None, self.value_type.clone()); + self.dict = Some(array); + Ok(()) + } + + fn set_data( + &mut self, + encoding: Encoding, + data: Bytes, + num_levels: usize, + num_values: Option, + ) -> Result<()> { + let decoder = match encoding { + Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY => { + let bit_width = data[0]; + let mut decoder = RleDecoder::new(bit_width); + decoder.set_data(data.slice(1..))?; + MaybeDictionaryDecoder::Dict { + decoder, + max_remaining_values: num_values.unwrap_or(num_levels), + } + } + _ => MaybeDictionaryDecoder::Fallback(ByteArrayDecoder::new( + encoding, + data, + num_levels, + num_values, + self.validate_utf8, + )?), + }; + + self.decoder = Some(decoder); + Ok(()) + } + + fn read(&mut self, out: &mut Self::Buffer, num_values: usize) -> Result { + match self.decoder.as_mut().expect("decoder set") { + MaybeDictionaryDecoder::Fallback(decoder) => { + decoder.read(out.spill_values()?, num_values, None) + } + MaybeDictionaryDecoder::Dict { + decoder, + max_remaining_values, + } => { + let len = num_values.min(*max_remaining_values); + + let dict = self + .dict + .as_ref() + .ok_or_else(|| general_err!("missing dictionary page for column"))?; + + assert_eq!(dict.data_type(), &self.value_type); + + if dict.is_empty() { + return Ok(0); // All data must be NULL + } + + match out.as_keys(dict) { + Some(keys) => { + // Happy path - can just copy keys + // Keys will be validated on conversion to arrow + + // TODO: Push vec into decoder (#5177) + let start = keys.len(); + keys.resize(start + len, K::default()); + let len = decoder.get_batch(&mut keys[start..])?; + keys.truncate(start + len); + *max_remaining_values -= len; + Ok(len) + } + None => { + // Sad path - need to recompute dictionary + // + // This either means we crossed into a new column chunk whilst + // reading this batch, or encountered non-dictionary encoded data + let values = out.spill_values()?; + let mut keys = vec![K::default(); len]; + let len = decoder.get_batch(&mut keys)?; + + assert_eq!(dict.data_type(), &self.value_type); + + let data = dict.to_data(); + let dict_buffers = data.buffers(); + let dict_offsets = dict_buffers[0].typed_data::(); + let dict_values = dict_buffers[1].as_slice(); + + values.extend_from_dictionary(&keys[..len], dict_offsets, dict_values)?; + *max_remaining_values -= len; + Ok(len) + } + } + } + } + } + + fn skip_values(&mut self, num_values: usize) -> Result { + match self.decoder.as_mut().expect("decoder set") { + MaybeDictionaryDecoder::Fallback(decoder) => decoder.skip::(num_values, None), + MaybeDictionaryDecoder::Dict { + decoder, + max_remaining_values, + } => { + let num_values = num_values.min(*max_remaining_values); + *max_remaining_values -= num_values; + decoder.skip(num_values) + } + } + } +} + +#[cfg(test)] +mod tests { + use arrow::compute::cast; + use arrow_array::{Array, StringArray}; + use arrow_buffer::Buffer; + + use crate::arrow::array_reader::test_util::{ + byte_array_all_encodings, encode_dictionary, utf8_column, + }; + use crate::arrow::record_reader::buffer::ValuesBuffer; + use crate::data_type::ByteArray; + + use super::*; + + fn utf8_dictionary() -> ArrowType { + ArrowType::Dictionary(Box::new(ArrowType::Int32), Box::new(ArrowType::Utf8)) + } + + #[test] + fn test_dictionary_preservation() { + let data_type = utf8_dictionary(); + + let data: Vec<_> = vec!["0", "1", "0", "1", "2", "1", "2"] + .into_iter() + .map(ByteArray::from) + .collect(); + let (dict, encoded) = encode_dictionary(&data); + + let column_desc = utf8_column(); + let mut decoder = DictionaryDecoder::::new(&column_desc); + + decoder + .set_dict(dict, 3, Encoding::RLE_DICTIONARY, false) + .unwrap(); + + decoder + .set_data(Encoding::RLE_DICTIONARY, encoded, 14, Some(data.len())) + .unwrap(); + + let mut output = DictionaryBuffer::::with_capacity(0); + assert_eq!(decoder.read(&mut output, 3).unwrap(), 3); + + let mut valid = vec![false, false, true, true, false, true]; + let valid_buffer = Buffer::from_iter(valid.iter().cloned()); + output + .pad_nulls(0, 3, valid.len(), valid_buffer.as_slice()) + .unwrap(); + + assert!(matches!(output, DictionaryBuffer::Dict { .. })); + + assert_eq!(decoder.read(&mut output, 4).unwrap(), 4); + + valid.extend_from_slice(&[false, false, true, true, false, true, true, false]); + let valid_buffer = Buffer::from_iter(valid.iter().cloned()); + output.pad_nulls(6, 4, 8, valid_buffer.as_slice()).unwrap(); + + assert!(matches!(output, DictionaryBuffer::Dict { .. })); + + let array = output.into_array(Some(valid_buffer), &data_type).unwrap(); + assert_eq!(array.data_type(), &data_type); + + let array = cast(&array, &ArrowType::Utf8).unwrap(); + let strings = array.as_any().downcast_ref::().unwrap(); + assert_eq!(strings.len(), 14); + + assert_eq!( + strings.iter().collect::>(), + vec![ + None, + None, + Some("0"), + Some("1"), + None, + Some("0"), + None, + None, + Some("1"), + Some("2"), + None, + Some("1"), + Some("2"), + None + ] + ) + } + + #[test] + fn test_dictionary_preservation_skip() { + let data_type = utf8_dictionary(); + + let data: Vec<_> = vec!["0", "1", "0", "1", "2", "1", "2"] + .into_iter() + .map(ByteArray::from) + .collect(); + let (dict, encoded) = encode_dictionary(&data); + + let column_desc = utf8_column(); + let mut decoder = DictionaryDecoder::::new(&column_desc); + + decoder + .set_dict(dict, 3, Encoding::RLE_DICTIONARY, false) + .unwrap(); + + decoder + .set_data(Encoding::RLE_DICTIONARY, encoded, 7, Some(data.len())) + .unwrap(); + + let mut output = DictionaryBuffer::::with_capacity(0); + + // read two skip one + assert_eq!(decoder.read(&mut output, 2).unwrap(), 2); + assert_eq!(decoder.skip_values(1).unwrap(), 1); + + assert!(matches!(output, DictionaryBuffer::Dict { .. })); + + // read two skip one + assert_eq!(decoder.read(&mut output, 2).unwrap(), 2); + assert_eq!(decoder.skip_values(1).unwrap(), 1); + + // read one and test on skip at the end + assert_eq!(decoder.read(&mut output, 1).unwrap(), 1); + assert_eq!(decoder.skip_values(4).unwrap(), 0); + + let valid = [true, true, true, true, true]; + let valid_buffer = Buffer::from_iter(valid.iter().cloned()); + output.pad_nulls(0, 5, 5, valid_buffer.as_slice()).unwrap(); + + assert!(matches!(output, DictionaryBuffer::Dict { .. })); + + let array = output.into_array(Some(valid_buffer), &data_type).unwrap(); + assert_eq!(array.data_type(), &data_type); + + let array = cast(&array, &ArrowType::Utf8).unwrap(); + let strings = array.as_any().downcast_ref::().unwrap(); + assert_eq!(strings.len(), 5); + + assert_eq!( + strings.iter().collect::>(), + vec![Some("0"), Some("1"), Some("1"), Some("2"), Some("2"),] + ) + } + + #[test] + fn test_dictionary_fallback() { + let data_type = utf8_dictionary(); + let data = vec!["hello", "world", "a", "b"]; + + let (pages, encoded_dictionary) = byte_array_all_encodings(data.clone()); + let num_encodings = pages.len(); + + let column_desc = utf8_column(); + let mut decoder = DictionaryDecoder::::new(&column_desc); + + decoder + .set_dict(encoded_dictionary, 4, Encoding::RLE_DICTIONARY, false) + .unwrap(); + + // Read all pages into single buffer + let mut output = DictionaryBuffer::::with_capacity(0); + + for (encoding, page) in pages { + decoder.set_data(encoding, page, 4, Some(4)).unwrap(); + assert_eq!(decoder.read(&mut output, 1024).unwrap(), 4); + } + let array = output.into_array(None, &data_type).unwrap(); + assert_eq!(array.data_type(), &data_type); + + let array = cast(&array, &ArrowType::Utf8).unwrap(); + let strings = array.as_any().downcast_ref::().unwrap(); + assert_eq!(strings.len(), data.len() * num_encodings); + + // Should have a copy of `data` for each encoding + for i in 0..num_encodings { + assert_eq!( + strings + .iter() + .skip(i * data.len()) + .take(data.len()) + .map(|x| x.unwrap()) + .collect::>(), + data + ) + } + } + + #[test] + fn test_dictionary_skip_fallback() { + let data_type = utf8_dictionary(); + let data = vec!["hello", "world", "a", "b"]; + + let (pages, encoded_dictionary) = byte_array_all_encodings(data.clone()); + let num_encodings = pages.len(); + + let column_desc = utf8_column(); + let mut decoder = DictionaryDecoder::::new(&column_desc); + + decoder + .set_dict(encoded_dictionary, 4, Encoding::RLE_DICTIONARY, false) + .unwrap(); + + // Read all pages into single buffer + let mut output = DictionaryBuffer::::with_capacity(0); + + for (encoding, page) in pages { + decoder.set_data(encoding, page, 4, Some(4)).unwrap(); + decoder.skip_values(2).expect("skipping two values"); + assert_eq!(decoder.read(&mut output, 1024).unwrap(), 2); + } + let array = output.into_array(None, &data_type).unwrap(); + assert_eq!(array.data_type(), &data_type); + + let array = cast(&array, &ArrowType::Utf8).unwrap(); + let strings = array.as_any().downcast_ref::().unwrap(); + assert_eq!(strings.len(), (data.len() - 2) * num_encodings); + + // Should have a copy of `data` for each encoding + for i in 0..num_encodings { + assert_eq!( + &strings + .iter() + .skip(i * (data.len() - 2)) + .take(data.len() - 2) + .map(|x| x.unwrap()) + .collect::>(), + &data[2..] + ) + } + } + + #[test] + fn test_too_large_dictionary() { + let data: Vec<_> = (0..128) + .map(|x| ByteArray::from(x.to_string().as_str())) + .collect(); + let (dictionary, _) = encode_dictionary(&data); + + let column_desc = utf8_column(); + + let mut decoder = DictionaryDecoder::::new(&column_desc); + let err = decoder + .set_dict(dictionary.clone(), 128, Encoding::RLE_DICTIONARY, false) + .unwrap_err() + .to_string(); + + assert!(err.contains("dictionary too large for index type")); + + let mut decoder = DictionaryDecoder::::new(&column_desc); + decoder + .set_dict(dictionary, 128, Encoding::RLE_DICTIONARY, false) + .unwrap(); + } + + #[test] + fn test_nulls() { + let data_type = utf8_dictionary(); + let (pages, encoded_dictionary) = byte_array_all_encodings(Vec::<&str>::new()); + + let column_desc = utf8_column(); + let mut decoder = DictionaryDecoder::new(&column_desc); + + decoder + .set_dict(encoded_dictionary, 4, Encoding::PLAIN_DICTIONARY, false) + .unwrap(); + + for (encoding, page) in pages.clone() { + let mut output = DictionaryBuffer::::with_capacity(0); + decoder.set_data(encoding, page, 8, None).unwrap(); + assert_eq!(decoder.read(&mut output, 1024).unwrap(), 0); + + output.pad_nulls(0, 0, 8, &[0]).unwrap(); + let array = output + .into_array(Some(Buffer::from(&[0])), &data_type) + .unwrap(); + + assert_eq!(array.len(), 8); + assert_eq!(array.null_count(), 8); + assert_eq!(array.logical_null_count(), 8); + } + + for (encoding, page) in pages { + let mut output = DictionaryBuffer::::with_capacity(0); + decoder.set_data(encoding, page, 8, None).unwrap(); + assert_eq!(decoder.skip_values(1024).unwrap(), 0); + + output.pad_nulls(0, 0, 8, &[0]).unwrap(); + let array = output + .into_array(Some(Buffer::from(&[0])), &data_type) + .unwrap(); + + assert_eq!(array.len(), 8); + assert_eq!(array.null_count(), 8); + assert_eq!(array.logical_null_count(), 8); + } + } +} diff --git a/vendor/parquet/src/arrow/array_reader/byte_view_array.rs b/vendor/parquet/src/arrow/array_reader/byte_view_array.rs new file mode 100644 index 0000000..cc1abf7 --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/byte_view_array.rs @@ -0,0 +1,848 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::array_reader::{ArrayReader, read_records, skip_records}; +use crate::arrow::buffer::view_buffer::ViewBuffer; +use crate::arrow::decoder::{DeltaByteArrayDecoder, DictIndexDecoder}; +use crate::arrow::record_reader::GenericRecordReader; +use crate::arrow::schema::parquet_to_arrow_field; +use crate::basic::{ConvertedType, Encoding}; +use crate::column::page::PageIterator; +use crate::column::reader::decoder::ColumnValueDecoder; +use crate::data_type::Int32Type; +use crate::encodings::decoding::{Decoder, DeltaBitPackDecoder}; +use crate::errors::{ParquetError, Result}; +use crate::schema::types::ColumnDescPtr; +use crate::util::utf8::check_valid_utf8; +use arrow_array::{ArrayRef, builder::make_view}; +use arrow_buffer::Buffer; +use arrow_data::ByteView; +use arrow_schema::DataType as ArrowType; +use bytes::Bytes; +use std::any::Any; + +/// Returns an [`ArrayReader`] that decodes the provided byte array column to view types. +/// +/// `batch_size` is a hint used to pre-allocate internal buffers and avoid +/// reallocations when reading the first batch. It is not a hard capacity +/// contract: with selective null padding (`padding_threshold`), value buffers +/// are sized from the number of values actually decoded. +pub fn make_byte_view_array_reader( + pages: Box, + column_desc: ColumnDescPtr, + arrow_type: Option, + batch_size: usize, + padding_threshold: Option, +) -> Result> { + // Check if Arrow type is specified, else create it from Parquet type + let data_type = match arrow_type { + Some(t) => t, + None => match parquet_to_arrow_field(column_desc.as_ref())?.data_type() { + ArrowType::Utf8 | ArrowType::Utf8View => ArrowType::Utf8View, + _ => ArrowType::BinaryView, + }, + }; + + match data_type { + ArrowType::BinaryView | ArrowType::Utf8View => { + let mut reader = GenericRecordReader::new(column_desc, batch_size); + if let Some(threshold) = padding_threshold { + reader.set_padding_threshold(threshold); + } + Ok(Box::new(ByteViewArrayReader::new(pages, data_type, reader))) + } + + _ => Err(general_err!( + "invalid data type for byte array reader read to view type - {}", + data_type + )), + } +} + +/// An [`ArrayReader`] for variable length byte arrays +struct ByteViewArrayReader { + data_type: ArrowType, + pages: Box, + def_levels_buffer: Option>, + rep_levels_buffer: Option>, + record_reader: GenericRecordReader, +} + +impl ByteViewArrayReader { + fn new( + pages: Box, + data_type: ArrowType, + record_reader: GenericRecordReader, + ) -> Self { + Self { + data_type, + pages, + def_levels_buffer: None, + rep_levels_buffer: None, + record_reader, + } + } +} + +impl ArrayReader for ByteViewArrayReader { + fn as_any(&self) -> &dyn Any { + self + } + + fn get_data_type(&self) -> &ArrowType { + &self.data_type + } + + fn read_records(&mut self, batch_size: usize) -> Result { + read_records(&mut self.record_reader, self.pages.as_mut(), batch_size) + } + + fn consume_batch(&mut self) -> Result { + let buffer = self.record_reader.consume_record_data(); + let null_buffer = self.record_reader.consume_compact_bitmap(); + self.def_levels_buffer = self.record_reader.consume_def_levels(); + self.rep_levels_buffer = self.record_reader.consume_rep_levels(); + self.record_reader.reset(); + + let array = buffer.into_array(null_buffer, &self.data_type); + + Ok(array) + } + + fn skip_records(&mut self, num_records: usize) -> Result { + skip_records(&mut self.record_reader, self.pages.as_mut(), num_records) + } + + fn get_def_levels(&self) -> Option<&[i16]> { + self.def_levels_buffer.as_deref() + } + + fn get_rep_levels(&self) -> Option<&[i16]> { + self.rep_levels_buffer.as_deref() + } + + fn max_def_level(&self) -> i16 { + self.record_reader.max_def_level() + } +} + +/// A [`ColumnValueDecoder`] for variable length byte arrays +struct ByteViewArrayColumnValueDecoder { + dict: Option, + decoder: Option, + validate_utf8: bool, +} + +impl ColumnValueDecoder for ByteViewArrayColumnValueDecoder { + type Buffer = ViewBuffer; + + fn new(desc: &ColumnDescPtr) -> Self { + let validate_utf8 = desc.converted_type() == ConvertedType::UTF8; + Self { + dict: None, + decoder: None, + validate_utf8, + } + } + + fn set_dict( + &mut self, + buf: Bytes, + num_values: u32, + encoding: Encoding, + _is_sorted: bool, + ) -> Result<()> { + if !matches!( + encoding, + Encoding::PLAIN | Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY + ) { + return Err(nyi_err!( + "Invalid/Unsupported encoding type for dictionary: {}", + encoding + )); + } + + let num_values = num_values as usize; + let mut buffer = ViewBuffer::with_capacity(num_values); + let mut decoder = + ByteViewArrayDecoderPlain::new(buf, num_values, Some(num_values), self.validate_utf8); + decoder.read(&mut buffer, usize::MAX)?; + self.dict = Some(buffer); + Ok(()) + } + + fn set_data( + &mut self, + encoding: Encoding, + data: Bytes, + num_levels: usize, + num_values: Option, + ) -> Result<()> { + self.decoder = Some(ByteViewArrayDecoder::new( + encoding, + data, + num_levels, + num_values, + self.validate_utf8, + )?); + Ok(()) + } + + fn read(&mut self, out: &mut Self::Buffer, num_values: usize) -> Result { + let decoder = self + .decoder + .as_mut() + .ok_or_else(|| general_err!("no decoder set"))?; + + decoder.read(out, num_values, self.dict.as_ref()) + } + + fn skip_values(&mut self, num_values: usize) -> Result { + let decoder = self + .decoder + .as_mut() + .ok_or_else(|| general_err!("no decoder set"))?; + + decoder.skip(num_values, self.dict.as_ref()) + } +} + +/// A generic decoder from uncompressed parquet value data to [`ViewBuffer`] +pub enum ByteViewArrayDecoder { + Plain(ByteViewArrayDecoderPlain), + Dictionary(ByteViewArrayDecoderDictionary), + DeltaLength(ByteViewArrayDecoderDeltaLength), + DeltaByteArray(ByteViewArrayDecoderDelta), +} + +impl ByteViewArrayDecoder { + pub fn new( + encoding: Encoding, + data: Bytes, + num_levels: usize, + num_values: Option, + validate_utf8: bool, + ) -> Result { + let decoder = match encoding { + Encoding::PLAIN => ByteViewArrayDecoder::Plain(ByteViewArrayDecoderPlain::new( + data, + num_levels, + num_values, + validate_utf8, + )), + Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY => { + ByteViewArrayDecoder::Dictionary(ByteViewArrayDecoderDictionary::new( + data, num_levels, num_values, + )?) + } + Encoding::DELTA_LENGTH_BYTE_ARRAY => ByteViewArrayDecoder::DeltaLength( + ByteViewArrayDecoderDeltaLength::new(data, validate_utf8)?, + ), + Encoding::DELTA_BYTE_ARRAY => ByteViewArrayDecoder::DeltaByteArray( + ByteViewArrayDecoderDelta::new(data, validate_utf8)?, + ), + _ => { + return Err(general_err!( + "unsupported encoding for byte array: {}", + encoding + )); + } + }; + + Ok(decoder) + } + + /// Read up to `len` values to `out` with the optional dictionary + pub fn read( + &mut self, + out: &mut ViewBuffer, + len: usize, + dict: Option<&ViewBuffer>, + ) -> Result { + match self { + ByteViewArrayDecoder::Plain(d) => d.read(out, len), + ByteViewArrayDecoder::Dictionary(d) => { + let dict = dict + .ok_or_else(|| general_err!("dictionary required for dictionary encoding"))?; + d.read(out, dict, len) + } + ByteViewArrayDecoder::DeltaLength(d) => d.read(out, len), + ByteViewArrayDecoder::DeltaByteArray(d) => d.read(out, len), + } + } + + /// Skip `len` values + pub fn skip(&mut self, len: usize, dict: Option<&ViewBuffer>) -> Result { + match self { + ByteViewArrayDecoder::Plain(d) => d.skip(len), + ByteViewArrayDecoder::Dictionary(d) => { + let dict = dict + .ok_or_else(|| general_err!("dictionary required for dictionary encoding"))?; + d.skip(dict, len) + } + ByteViewArrayDecoder::DeltaLength(d) => d.skip(len), + ByteViewArrayDecoder::DeltaByteArray(d) => d.skip(len), + } + } +} + +/// Decoder from [`Encoding::PLAIN`] data to [`ViewBuffer`] +pub struct ByteViewArrayDecoderPlain { + buf: Buffer, + offset: usize, + + validate_utf8: bool, + + /// This is a maximum as the null count is not always known, e.g. value data from + /// a v1 data page + max_remaining_values: usize, +} + +impl ByteViewArrayDecoderPlain { + pub fn new( + buf: Bytes, + num_levels: usize, + num_values: Option, + validate_utf8: bool, + ) -> Self { + Self { + buf: Buffer::from(buf), + offset: 0, + max_remaining_values: num_values.unwrap_or(num_levels), + validate_utf8, + } + } + + pub fn read(&mut self, output: &mut ViewBuffer, len: usize) -> Result { + if self.validate_utf8 { + self.read_impl::(output, len) + } else { + self.read_impl::(output, len) + } + } + + fn read_impl( + &mut self, + output: &mut ViewBuffer, + len: usize, + ) -> Result { + // avoid creating a new buffer if the last buffer is the same as the current buffer + // This is especially useful when row-level filtering is applied, where we call lots of small `read` over the same buffer. + let block_id = { + if output.buffers.last().is_some_and(|x| x.ptr_eq(&self.buf)) { + output.buffers.len() as u32 - 1 + } else { + output.append_block(self.buf.clone()) + } + }; + + let to_read = len.min(self.max_remaining_values); + + let buf: &[u8] = self.buf.as_ref(); + let buf_len = buf.len(); + let mut end_offset = self.offset; + let mut utf8_validation_begin = end_offset; + + output.views.reserve(to_read); + + // Safety: we reserved enough space in output.views + // and we will only write up to to_read views / track how many views we wrote. + // Ideally, we would use `Vec::extend` here, but this generates sub-optimal code. + let views_ptr = output.views.as_mut_ptr().wrapping_add(output.views.len()); + for i in 0..to_read { + let start_offset = end_offset + 4; + + if start_offset > buf_len { + return Err(ParquetError::EOF("eof decoding byte array".into())); + } + + // Safety: we have checked that start_offset <= buf_len + let len = u32::from_le_bytes( + unsafe { buf.get_unchecked(end_offset..start_offset) } + .try_into() + .unwrap(), + ); + + end_offset = start_offset + len as usize; + + if end_offset > buf_len { + return Err(ParquetError::EOF("eof decoding byte array".into())); + } + + if VALIDATE_UTF8 { + // It seems you are trying to understand what's going on here, take a breath and be patient. + // Utf-8 validation is a non-trivial task, here are some background facts: + // (1) Validating one 2048-byte string is much faster than validating 128 of 16-byte string. + // As shown in https://github.com/apache/arrow-rs/pull/6009#issuecomment-2211174229 + // Potentially because the SIMD operations favor longer strings. + // (2) Practical strings are short, 99% of strings are smaller than 100 bytes, as shown in paper: + // https://www.vldb.org/pvldb/vol17/p148-zeng.pdf, Figure 5f. + // (3) Parquet plain encoding makes utf-8 validation harder, + // because it stores the length of each string right before the string. + // This means naive utf-8 validation will be slow, because the validation need to skip the length bytes. + // I.e., the validation cannot validate the buffer in one pass, but instead, validate strings chunk by chunk. + // + // Given the above observations, the goal is to do batch validation as much as possible. + // The key idea is that if the length is smaller than 128 (99% of the case), then the length bytes are valid utf-8, as reasoned below: + // If the length is smaller than 128, its 4-byte encoding are [0, 0, 0, len]. + // Each of the byte is a valid ASCII character, so they are valid utf-8. + // Since they are all smaller than 128, the won't break a utf-8 code point (won't mess with later bytes). + // + // The implementation keeps a water mark `utf8_validation_begin` to track the beginning of the buffer that is not validated. + // If the length is smaller than 128, then we continue to next string. + // If the length is larger than 128, then we validate the buffer before the length bytes, and move the water mark to the beginning of next string. + if len >= 128 { + // unfortunately, the len bytes may not be valid utf8, we need to wrap up and validate everything before it. + check_valid_utf8(unsafe { + buf.get_unchecked(utf8_validation_begin..start_offset - 4) + })?; + // move the cursor to skip the len bytes. + utf8_validation_begin = start_offset; + } + } + + let view = make_view( + unsafe { buf.get_unchecked(start_offset..end_offset) }, + block_id, + start_offset as u32, + ); + // Safety: views_ptr is valid for writes, and we have reserved enough space. + unsafe { + views_ptr.add(i).write(view); + } + } + + // Safety: we have written `to_read` views to `views_ptr` + unsafe { + output.views.set_len(output.views.len() + to_read); + } + if VALIDATE_UTF8 { + // validate values from the previously validated location up to (but not including) + // the length of this string + check_valid_utf8(unsafe { buf.get_unchecked(utf8_validation_begin..end_offset) })?; + } + + self.offset = end_offset; + self.max_remaining_values -= to_read; + + Ok(to_read) + } + + pub fn skip(&mut self, to_skip: usize) -> Result { + let to_skip = to_skip.min(self.max_remaining_values); + let mut skip = 0; + let buf: &[u8] = self.buf.as_ref(); + + while self.offset < self.buf.len() && skip != to_skip { + if self.offset + 4 > buf.len() { + return Err(ParquetError::EOF("eof decoding byte array".into())); + } + let len_bytes: [u8; 4] = buf[self.offset..self.offset + 4].try_into().unwrap(); + let len = u32::from_le_bytes(len_bytes) as usize; + skip += 1; + self.offset = self.offset + 4 + len; + } + self.max_remaining_values -= skip; + Ok(skip) + } +} + +pub struct ByteViewArrayDecoderDictionary { + decoder: DictIndexDecoder, +} + +impl ByteViewArrayDecoderDictionary { + fn new(data: Bytes, num_levels: usize, num_values: Option) -> Result { + Ok(Self { + decoder: DictIndexDecoder::new(data, num_levels, num_values)?, + }) + } + + /// Reads the next indexes from self.decoder + /// the indexes are assumed to be indexes into `dict` + /// the output values are written to output + /// + /// Assumptions / Optimization + /// This function checks if dict.buffers() are the last buffers in `output`, and if so + /// reuses the dictionary page buffers directly without copying data + /// + /// If the dictionary is empty, the buffer contains empty view. + fn read(&mut self, output: &mut ViewBuffer, dict: &ViewBuffer, len: usize) -> Result { + if dict.is_empty() || len == 0 { + return Ok(0); + } + + // Check if the last few buffer of `output`` are the same as the `dict` buffer + // This is to avoid creating a new buffers if the same dictionary is used for multiple `read` + let need_to_create_new_buffer = { + if output.buffers.len() >= dict.buffers.len() { + let offset = output.buffers.len() - dict.buffers.len(); + output.buffers[offset..] + .iter() + .zip(dict.buffers.iter()) + .any(|(a, b)| !a.ptr_eq(b)) + } else { + true + } + }; + + if need_to_create_new_buffer { + for b in dict.buffers.iter() { + output.buffers.push(b.clone()); + } + } + + // Calculate the offset of the dictionary buffers in the output buffers + // For example if the 2nd buffer in the dictionary is the 5th buffer in the output buffers, + // then the base_buffer_idx is 5 - 2 = 3 + let base_buffer_idx = output.buffers.len() as u32 - dict.buffers.len() as u32; + + // Pre-reserve output capacity to avoid per-chunk reallocation in extend + output.views.reserve(len); + + let mut error = None; + let read = self.decoder.read(len, |keys| { + if base_buffer_idx == 0 { + // the dictionary buffers are the last buffers in output, we can directly use the views + output + .views + .extend(keys.iter().map(|k| match dict.views.get(*k as usize) { + Some(&view) => view, + None => { + if error.is_none() { + error = Some(general_err!("invalid key={} for dictionary", *k)); + } + 0 + } + })); + Ok(()) + } else { + output + .views + .extend(keys.iter().map(|k| match dict.views.get(*k as usize) { + Some(&view) => { + let len = view as u32; + if len <= 12 { + view + } else { + let mut view = ByteView::from(view); + view.buffer_index += base_buffer_idx; + view.into() + } + } + None => { + if error.is_none() { + error = Some(general_err!("invalid key={} for dictionary", *k)); + } + 0 + } + })); + Ok(()) + } + })?; + if let Some(e) = error { + return Err(e); + } + Ok(read) + } + + fn skip(&mut self, dict: &ViewBuffer, to_skip: usize) -> Result { + if dict.is_empty() { + return Ok(0); + } + self.decoder.skip(to_skip) + } +} + +/// Decoder from [`Encoding::DELTA_LENGTH_BYTE_ARRAY`] data to [`ViewBuffer`] +pub struct ByteViewArrayDecoderDeltaLength { + lengths: Vec, + data: Bytes, + length_offset: usize, + data_offset: usize, + validate_utf8: bool, +} + +impl ByteViewArrayDecoderDeltaLength { + fn new(data: Bytes, validate_utf8: bool) -> Result { + let mut len_decoder = DeltaBitPackDecoder::::new(); + len_decoder.set_data(data.clone(), 0)?; + let values = len_decoder.values_left(); + + let mut lengths = vec![0; values]; + len_decoder.get(&mut lengths)?; + + let mut total_bytes = 0; + + for l in lengths.iter() { + if *l < 0 { + return Err(ParquetError::General( + "negative delta length byte array length".to_string(), + )); + } + total_bytes += *l as usize; + } + + if total_bytes + len_decoder.get_offset() > data.len() { + return Err(ParquetError::General( + "Insufficient delta length byte array bytes".to_string(), + )); + } + + Ok(Self { + lengths, + data, + validate_utf8, + length_offset: 0, + data_offset: len_decoder.get_offset(), + }) + } + + fn read(&mut self, output: &mut ViewBuffer, len: usize) -> Result { + let to_read = len.min(self.lengths.len() - self.length_offset); + output.views.reserve(to_read); + + let src_lengths = &self.lengths[self.length_offset..self.length_offset + to_read]; + + // Zero copy convert `bytes::Bytes` into `arrow_buffer::Buffer` + let bytes = Buffer::from(self.data.clone()); + let block_id = output.append_block(bytes); + + let mut current_offset = self.data_offset; + let initial_offset = current_offset; + + output.views.extend(src_lengths.iter().map(|length| { + let len = *length as u32; + let start_offset = current_offset; + current_offset += len as usize; + // # Safety + // The length and offset are guaranteed valid by the entry check in `new` + make_view( + &self.data[start_offset..start_offset + len as usize], + block_id, + start_offset as u32, + ) + })); + + // Delta length encoding has continuous strings, we can validate utf8 in one go + if self.validate_utf8 { + check_valid_utf8(&self.data[initial_offset..current_offset])?; + } + + self.data_offset = current_offset; + self.length_offset += to_read; + + Ok(to_read) + } + + fn skip(&mut self, to_skip: usize) -> Result { + let remain_values = self.lengths.len() - self.length_offset; + let to_skip = remain_values.min(to_skip); + + let src_lengths = &self.lengths[self.length_offset..self.length_offset + to_skip]; + let total_bytes: usize = src_lengths.iter().map(|x| *x as usize).sum(); + + self.data_offset += total_bytes; + self.length_offset += to_skip; + Ok(to_skip) + } +} + +/// Decoder from [`Encoding::DELTA_BYTE_ARRAY`] to [`ViewBuffer`] +pub struct ByteViewArrayDecoderDelta { + decoder: DeltaByteArrayDecoder, + validate_utf8: bool, +} + +impl ByteViewArrayDecoderDelta { + fn new(data: Bytes, validate_utf8: bool) -> Result { + Ok(Self { + decoder: DeltaByteArrayDecoder::new(data)?, + validate_utf8, + }) + } + + // Unlike other encodings, we need to copy the data. + // + // DeltaByteArray data is stored using shared prefixes/suffixes, + // which results in potentially non-contiguous + // strings, while Arrow encodings require contiguous strings + // + // + + fn read(&mut self, output: &mut ViewBuffer, len: usize) -> Result { + let to_reserve = len.min(self.decoder.remaining()); + output.views.reserve(to_reserve); + + // array buffer only have long strings + let mut array_buffer: Vec = Vec::with_capacity(4096); + + let buffer_id = output.buffers.len() as u32; + + // Use unsafe ptr writes instead of per-element push to avoid + // repeated length checks. Safety: we reserved enough space above. + let views_ptr = output.views.as_mut_ptr(); + let initial_len = output.views.len(); + let mut write_count = 0; + + let read = if !self.validate_utf8 { + self.decoder.read(len, |bytes| { + let offset = array_buffer.len(); + let view = make_view(bytes, buffer_id, offset as u32); + if bytes.len() > 12 { + // only copy the data to buffer if the string can not be inlined. + array_buffer.extend_from_slice(bytes); + } + + // Safety: views_ptr is valid for writes, we reserved enough space, + // and write_count < to_reserve. + unsafe { + views_ptr.add(initial_len + write_count).write(view); + } + write_count += 1; + Ok(()) + })? + } else { + // utf8 validation buffer has only short strings. These short + // strings are inlined into the views but we copy them into a + // contiguous buffer to accelerate validation. + let mut utf8_validation_buffer = Vec::with_capacity(4096); + + let v = self.decoder.read(len, |bytes| { + let offset = array_buffer.len(); + let view = make_view(bytes, buffer_id, offset as u32); + if bytes.len() > 12 { + // only copy the data to buffer if the string can not be inlined. + array_buffer.extend_from_slice(bytes); + } else { + utf8_validation_buffer.extend_from_slice(bytes); + } + + // Safety: views_ptr is valid for writes, we reserved enough space, + // and write_count < to_reserve. Utf-8 validation is done later. + unsafe { + views_ptr.add(initial_len + write_count).write(view); + } + write_count += 1; + Ok(()) + })?; + check_valid_utf8(&array_buffer)?; + check_valid_utf8(&utf8_validation_buffer)?; + v + }; + + // Safety: we wrote exactly `read` views via ptr writes above + unsafe { + output.views.set_len(initial_len + read); + } + + let actual_block_id = output.append_block(Buffer::from_vec(array_buffer)); + assert_eq!(actual_block_id, buffer_id); + Ok(read) + } + + fn skip(&mut self, to_skip: usize) -> Result { + self.decoder.skip(to_skip) + } +} + +#[cfg(test)] +mod tests { + use arrow_array::StringViewArray; + use arrow_buffer::Buffer; + + use crate::{ + arrow::{ + array_reader::test_util::{byte_array_all_encodings, encode_byte_array, utf8_column}, + buffer::view_buffer::ViewBuffer, + record_reader::buffer::ValuesBuffer, + }, + basic::Encoding, + column::reader::decoder::ColumnValueDecoder, + data_type::ByteArray, + }; + + use super::*; + + #[test] + fn test_byte_array_string_view_decoder() { + let (pages, encoded_dictionary) = + byte_array_all_encodings(vec!["hello", "world", "large payload over 12 bytes", "b"]); + + let column_desc = utf8_column(); + let mut decoder = ByteViewArrayColumnValueDecoder::new(&column_desc); + + decoder + .set_dict(encoded_dictionary, 4, Encoding::RLE_DICTIONARY, false) + .unwrap(); + + for (encoding, page) in pages { + let mut output = ViewBuffer::with_capacity(0); + decoder.set_data(encoding, page, 4, Some(4)).unwrap(); + + assert_eq!(decoder.read(&mut output, 1).unwrap(), 1); + assert_eq!(decoder.read(&mut output, 1).unwrap(), 1); + assert_eq!(decoder.read(&mut output, 2).unwrap(), 2); + assert_eq!(decoder.read(&mut output, 4).unwrap(), 0); + + assert_eq!(output.views.len(), 4); + + let valid = [false, false, true, true, false, true, true, false, false]; + let valid_buffer = Buffer::from_iter(valid.iter().cloned()); + + output + .pad_nulls(0, 4, valid.len(), valid_buffer.as_slice()) + .unwrap(); + let array = output.into_array(Some(valid_buffer), &ArrowType::Utf8View); + let strings = array.as_any().downcast_ref::().unwrap(); + + assert_eq!( + strings.iter().collect::>(), + vec![ + None, + None, + Some("hello"), + Some("world"), + None, + Some("large payload over 12 bytes"), + Some("b"), + None, + None, + ] + ); + } + } + + #[test] + fn test_byte_view_array_plain_decoder_reuse_buffer() { + let byte_array = vec!["hello", "world", "large payload over 12 bytes", "b"]; + let byte_array: Vec = byte_array.into_iter().map(|x| x.into()).collect(); + let pages = encode_byte_array(Encoding::PLAIN, &byte_array); + + let column_desc = utf8_column(); + let mut decoder = ByteViewArrayColumnValueDecoder::new(&column_desc); + + let mut view_buffer = ViewBuffer::with_capacity(0); + decoder.set_data(Encoding::PLAIN, pages, 4, None).unwrap(); + decoder.read(&mut view_buffer, 1).unwrap(); + decoder.read(&mut view_buffer, 1).unwrap(); + assert_eq!(view_buffer.buffers.len(), 1); + + decoder.read(&mut view_buffer, 1).unwrap(); + assert_eq!(view_buffer.buffers.len(), 1); + } +} diff --git a/vendor/parquet/src/arrow/array_reader/cached_array_reader.rs b/vendor/parquet/src/arrow/array_reader/cached_array_reader.rs new file mode 100644 index 0000000..7e31ff7 --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/cached_array_reader.rs @@ -0,0 +1,769 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`CachedArrayReader`] wrapper around [`ArrayReader`] + +use crate::arrow::array_reader::row_group_cache::BatchID; +use crate::arrow::array_reader::{ArrayReader, row_group_cache::RowGroupCache}; +use crate::arrow::arrow_reader::metrics::ArrowReaderMetrics; +use crate::errors::Result; +use arrow_array::{ArrayRef, BooleanArray, new_empty_array}; +use arrow_buffer::BooleanBufferBuilder; +use arrow_schema::DataType as ArrowType; +use std::any::Any; +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; + +/// Role of the cached array reader +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CacheRole { + /// Producer role: inserts data into the cache during filter phase + Producer, + /// Consumer role: removes consumed data from the cache during output building phase + Consumer, +} + +/// A cached wrapper around an ArrayReader that avoids duplicate decoding +/// when the same column appears in both filter predicates and output projection. +/// +/// This reader acts as a transparent layer over the inner reader, using a cache +/// to avoid redundant work when the same data is needed multiple times. +/// +/// The reader can operate in two roles: +/// - Producer: During filter phase, inserts decoded data into the cache +/// - Consumer: During output building, consumes and removes data from the cache +/// +/// This means the memory consumption of the cache has two stages: +/// 1. During the filter phase, the memory increases as the cache is populated +/// 2. It peaks when filters are built. +/// 3. It decreases as the cached data is consumed. +/// +/// ```text +/// ▲ +/// │ ╭─╮ +/// │ ╱ ╲ +/// │ ╱ ╲ +/// │ ╱ ╲ +/// │ ╱ ╲ +/// │╱ ╲ +/// └─────────────╲──────► Time +/// │ │ │ +/// Filter Peak Consume +/// Phase (Built) (Decrease) +/// ``` +pub struct CachedArrayReader { + /// The underlying array reader + inner: Box, + /// Shared cache for this row group + shared_cache: Arc>, + /// Column index for cache key generation + column_idx: usize, + /// Current logical position in the data stream for this reader (for cache key generation) + outer_position: usize, + /// Current position in `inner` + inner_position: usize, + /// Batch size for the cache + batch_size: usize, + /// Boolean buffer builder to track selections for the next consume_batch() + selections: BooleanBufferBuilder, + /// Role of this reader (Producer or Consumer) + role: CacheRole, + /// Local cache to store batches between read_records and consume_batch calls + /// This ensures data is available even if the shared cache evicts items + local_cache: HashMap, + /// Statistics to report on the Cache behavior + metrics: ArrowReaderMetrics, +} + +impl CachedArrayReader { + /// Creates a new cached array reader with the specified role + pub fn new( + inner: Box, + cache: Arc>, + column_idx: usize, + role: CacheRole, + metrics: ArrowReaderMetrics, + ) -> Self { + let batch_size = cache.read().unwrap().batch_size(); + + Self { + inner, + shared_cache: cache, + column_idx, + outer_position: 0, + inner_position: 0, + batch_size, + selections: BooleanBufferBuilder::new(0), + role, + local_cache: HashMap::new(), + metrics, + } + } + + fn get_batch_id_from_position(&self, row_id: usize) -> BatchID { + BatchID { + val: row_id / self.batch_size, + } + } + + /// Loads the batch with the given ID (first row offset) from the inner + /// reader + /// + /// After this call the required batch will be available in + /// `self.local_cache` and may also be stored in `self.shared_cache`. + /// + fn fetch_batch(&mut self, batch_id: BatchID) -> Result { + let first_row_offset = batch_id.val * self.batch_size; + if self.inner_position < first_row_offset { + let to_skip = first_row_offset - self.inner_position; + let skipped = self.inner.skip_records(to_skip)?; + assert_eq!(skipped, to_skip); + self.inner_position += skipped; + } + + // For sparse mask reads, this full-batch fallback relies on `MaskCursor` + // ending every chunk at a selected row. Predicate fetch expands cached + // columns to batch boundaries, so the batch containing that row is loaded. + let read = self.inner.read_records(self.batch_size)?; + + // If there are no remaining records (EOF), return immediately without + // attempting to cache an empty batch. This prevents inserting zero-length + // arrays into the cache which can later cause panics when slicing. + if read == 0 { + return Ok(0); + } + + let array = self.inner.consume_batch()?; + + // Store in both shared cache and local cache + // The shared cache is used to reuse results between readers + // The local cache ensures data is available for our consume_batch call + let _cached = + self.shared_cache + .write() + .unwrap() + .insert(self.column_idx, batch_id, array.clone()); + // Note: if the shared cache is full (_cached == false), we continue without caching + // The local cache will still store the data for this reader's use + + self.local_cache.insert(batch_id, array); + + self.inner_position += read; + Ok(read) + } + + /// Remove batches from cache that have been completely consumed + /// This is only called for Consumer role readers + fn cleanup_consumed_batches(&mut self) { + let current_batch_id = self.get_batch_id_from_position(self.outer_position); + + // 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, + }, + ); + } + } + } +} + +impl ArrayReader for CachedArrayReader { + fn as_any(&self) -> &dyn Any { + self + } + + fn get_data_type(&self) -> &ArrowType { + self.inner.get_data_type() + } + + fn read_records(&mut self, num_records: usize) -> Result { + let mut read = 0; + while read < num_records { + let batch_id = self.get_batch_id_from_position(self.outer_position); + + // Check local cache first + let cached = if let Some(array) = self.local_cache.get(&batch_id) { + Some(Arc::clone(array)) + } else { + // If not in local cache, i.e., we are consumer, check shared cache + let cache_content = self + .shared_cache + .read() + .unwrap() + .get(self.column_idx, batch_id); + if let Some(array) = cache_content.as_ref() { + // Store in local cache for later use in consume_batch + self.local_cache.insert(batch_id, Arc::clone(array)); + } + cache_content + }; + + match cached { + Some(array) => { + let array_len = array.len(); + if array_len + batch_id.val * self.batch_size > self.outer_position { + // the cache batch has some records that we can select + let v = array_len + batch_id.val * self.batch_size - self.outer_position; + let select_cnt = std::cmp::min(num_records - read, v); + read += select_cnt; + self.metrics.increment_cache_reads(select_cnt); + self.outer_position += select_cnt; + self.selections.append_n(select_cnt, true); + } else { + // this is last batch and we have used all records from it + break; + } + } + None => { + let read_from_inner = self.fetch_batch(batch_id)?; + // Reached end-of-file, no more records to read + if read_from_inner == 0 { + break; + } + self.metrics.increment_inner_reads(read_from_inner); + let select_from_this_batch = std::cmp::min( + num_records - read, + self.inner_position - self.outer_position, + ); + read += select_from_this_batch; + self.outer_position += select_from_this_batch; + self.selections.append_n(select_from_this_batch, true); + if read_from_inner < self.batch_size { + // this is last batch from inner reader + break; + } + } + } + } + Ok(read) + } + + fn skip_records(&mut self, num_records: usize) -> Result { + let mut skipped = 0; + while skipped < num_records { + let size = std::cmp::min(num_records - skipped, self.batch_size); + skipped += size; + self.selections.append_n(size, false); + self.outer_position += size; + } + Ok(num_records) + } + + fn consume_batch(&mut self) -> Result { + let row_count = self.selections.len(); + if row_count == 0 { + return Ok(new_empty_array(self.inner.get_data_type())); + } + + let start_position = self.outer_position - row_count; + + let selection_buffer = self.selections.finish(); + + let start_batch = start_position / self.batch_size; + let end_batch = (start_position + row_count - 1) / self.batch_size; + + let mut selected_arrays = Vec::new(); + for batch_id in start_batch..=end_batch { + let batch_start = batch_id * self.batch_size; + let batch_end = batch_start + self.batch_size - 1; + let batch_id = self.get_batch_id_from_position(batch_start); + + // Calculate the overlap between the start_position and the batch + let overlap_start = start_position.max(batch_start); + let overlap_end = (start_position + row_count - 1).min(batch_end); + + if overlap_start > overlap_end { + continue; + } + + let selection_start = overlap_start - start_position; + let selection_length = overlap_end - overlap_start + 1; + let mask = selection_buffer.slice(selection_start, selection_length); + + if mask.count_set_bits() == 0 { + continue; + } + + let mask_array = BooleanArray::from(mask); + // Read from local cache instead of shared cache to avoid cache eviction issues + let cached = self + .local_cache + .get(&batch_id) + .expect("data must be already cached in the read_records call, this is a bug"); + let cached = cached.slice(overlap_start - batch_start, selection_length); + let filtered = arrow_select::filter::filter(&cached, &mask_array)?; + selected_arrays.push(filtered); + } + + self.selections = BooleanBufferBuilder::new(0); + + // Only remove batches from local buffer that are completely behind current position + // Keep the current batch and any future batches as they might still be needed + let current_batch_id = self.get_batch_id_from_position(self.outer_position); + self.local_cache + .retain(|batch_id, _| batch_id.val >= current_batch_id.val); + + // For consumers, cleanup batches that have been completely consumed + // This reduces the memory usage of the shared cache + if self.role == CacheRole::Consumer { + self.cleanup_consumed_batches(); + } + + match selected_arrays.len() { + 0 => Ok(new_empty_array(self.inner.get_data_type())), + 1 => Ok(selected_arrays.into_iter().next().unwrap()), + _ => Ok(arrow_select::concat::concat( + &selected_arrays + .iter() + .map(|a| a.as_ref()) + .collect::>(), + )?), + } + } + + fn get_def_levels(&self) -> Option<&[i16]> { + None // we don't allow nullable parent for now. + } + + fn get_rep_levels(&self) -> Option<&[i16]> { + None + } + + fn max_def_level(&self) -> i16 { + self.inner.max_def_level() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::array_reader::ArrayReader; + use crate::arrow::array_reader::row_group_cache::RowGroupCache; + use arrow_array::{ArrayRef, Int32Array}; + use std::sync::{Arc, RwLock}; + + // Mock ArrayReader for testing + struct MockArrayReader { + data: Vec, + position: usize, + records_to_consume: usize, + data_type: ArrowType, + } + + impl MockArrayReader { + fn new(data: Vec) -> Self { + Self { + data, + position: 0, + records_to_consume: 0, + data_type: ArrowType::Int32, + } + } + } + + impl ArrayReader for MockArrayReader { + fn as_any(&self) -> &dyn Any { + self + } + + fn get_data_type(&self) -> &ArrowType { + &self.data_type + } + + fn read_records(&mut self, batch_size: usize) -> Result { + let remaining = self.data.len() - self.position; + let to_read = std::cmp::min(batch_size, remaining); + self.records_to_consume += to_read; + Ok(to_read) + } + + fn consume_batch(&mut self) -> Result { + let start = self.position; + let end = start + self.records_to_consume; + let slice = &self.data[start..end]; + self.position = end; + self.records_to_consume = 0; + Ok(Arc::new(Int32Array::from(slice.to_vec()))) + } + + fn skip_records(&mut self, num_records: usize) -> Result { + let remaining = self.data.len() - self.position; + let to_skip = std::cmp::min(num_records, remaining); + self.position += to_skip; + Ok(to_skip) + } + + fn get_def_levels(&self) -> Option<&[i16]> { + None + } + + fn get_rep_levels(&self) -> Option<&[i16]> { + None + } + } + + #[test] + fn test_cached_reader_basic() { + let metrics = ArrowReaderMetrics::disabled(); + let mock_reader = MockArrayReader::new(vec![1, 2, 3, 4, 5]); + let cache = Arc::new(RwLock::new(RowGroupCache::new(3, usize::MAX))); // Batch size 3 + let mut cached_reader = CachedArrayReader::new( + Box::new(mock_reader), + cache, + 0, + CacheRole::Producer, + metrics, + ); + + // Read 3 records + let records_read = cached_reader.read_records(3).unwrap(); + assert_eq!(records_read, 3); + + let array = cached_reader.consume_batch().unwrap(); + assert_eq!(array.len(), 3); + + let int32_array = array.as_any().downcast_ref::().unwrap(); + assert_eq!(int32_array.values(), &[1, 2, 3]); + + // Read 3 more records + let records_read = cached_reader.read_records(3).unwrap(); + assert_eq!(records_read, 2); + } + + #[test] + fn test_read_skip_pattern() { + let metrics = ArrowReaderMetrics::disabled(); + let mock_reader = MockArrayReader::new(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); + let cache = Arc::new(RwLock::new(RowGroupCache::new(5, usize::MAX))); // Batch size 5 + let mut cached_reader = CachedArrayReader::new( + Box::new(mock_reader), + cache, + 0, + CacheRole::Consumer, + metrics, + ); + + let read1 = cached_reader.read_records(2).unwrap(); + assert_eq!(read1, 2); + + let array1 = cached_reader.consume_batch().unwrap(); + assert_eq!(array1.len(), 2); + let int32_array = array1.as_any().downcast_ref::().unwrap(); + assert_eq!(int32_array.values(), &[1, 2]); + + let skipped = cached_reader.skip_records(2).unwrap(); + assert_eq!(skipped, 2); + + let read2 = cached_reader.read_records(1).unwrap(); + assert_eq!(read2, 1); + + // Consume it (should be the 5th element after skipping 3,4) + let array2 = cached_reader.consume_batch().unwrap(); + assert_eq!(array2.len(), 1); + let int32_array = array2.as_any().downcast_ref::().unwrap(); + assert_eq!(int32_array.values(), &[5]); + } + + #[test] + fn test_multiple_reads_before_consume() { + let metrics = ArrowReaderMetrics::disabled(); + let mock_reader = MockArrayReader::new(vec![1, 2, 3, 4, 5, 6]); + let cache = Arc::new(RwLock::new(RowGroupCache::new(3, usize::MAX))); // Batch size 3 + let mut cached_reader = CachedArrayReader::new( + Box::new(mock_reader), + cache, + 0, + CacheRole::Consumer, + metrics, + ); + + // Multiple reads should accumulate + let read1 = cached_reader.read_records(2).unwrap(); + assert_eq!(read1, 2); + + let read2 = cached_reader.read_records(1).unwrap(); + assert_eq!(read2, 1); + + // Consume should return all accumulated records + let array = cached_reader.consume_batch().unwrap(); + assert_eq!(array.len(), 3); + let int32_array = array.as_any().downcast_ref::().unwrap(); + assert_eq!(int32_array.values(), &[1, 2, 3]); + } + + #[test] + fn test_eof_behavior() { + let metrics = ArrowReaderMetrics::disabled(); + let mock_reader = MockArrayReader::new(vec![1, 2, 3]); + let cache = Arc::new(RwLock::new(RowGroupCache::new(5, usize::MAX))); // Batch size 5 + let mut cached_reader = CachedArrayReader::new( + Box::new(mock_reader), + cache, + 0, + CacheRole::Consumer, + metrics, + ); + + // Try to read more than available + let read1 = cached_reader.read_records(5).unwrap(); + assert_eq!(read1, 3); // Should only get 3 records (all available) + + let array1 = cached_reader.consume_batch().unwrap(); + assert_eq!(array1.len(), 3); + + // Further reads should return 0 + let read2 = cached_reader.read_records(1).unwrap(); + assert_eq!(read2, 0); + + let array2 = cached_reader.consume_batch().unwrap(); + assert_eq!(array2.len(), 0); + } + + #[test] + fn test_cache_sharing() { + let metrics = ArrowReaderMetrics::disabled(); + let cache = Arc::new(RwLock::new(RowGroupCache::new(5, usize::MAX))); // Batch size 5 + + // First reader - populate cache + let mock_reader1 = MockArrayReader::new(vec![1, 2, 3, 4, 5]); + let mut cached_reader1 = CachedArrayReader::new( + Box::new(mock_reader1), + cache.clone(), + 0, + CacheRole::Producer, + metrics.clone(), + ); + + cached_reader1.read_records(3).unwrap(); + let array1 = cached_reader1.consume_batch().unwrap(); + assert_eq!(array1.len(), 3); + + // Second reader with different column index should not interfere + let mock_reader2 = MockArrayReader::new(vec![10, 20, 30, 40, 50]); + let mut cached_reader2 = CachedArrayReader::new( + Box::new(mock_reader2), + cache.clone(), + 1, + CacheRole::Consumer, + metrics.clone(), + ); + + cached_reader2.read_records(2).unwrap(); + let array2 = cached_reader2.consume_batch().unwrap(); + assert_eq!(array2.len(), 2); + + // Verify the second reader got its own data, not from cache + let int32_array = array2.as_any().downcast_ref::().unwrap(); + assert_eq!(int32_array.values(), &[10, 20]); + } + + #[test] + fn test_consumer_removes_batches() { + let metrics = ArrowReaderMetrics::disabled(); + let mock_reader = MockArrayReader::new(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); + let cache = Arc::new(RwLock::new(RowGroupCache::new(3, usize::MAX))); // Batch size 3 + let mut consumer_reader = CachedArrayReader::new( + Box::new(mock_reader), + cache.clone(), + 0, + CacheRole::Consumer, + metrics, + ); + + // Read first batch (positions 0-2, batch 0) + let read1 = consumer_reader.read_records(3).unwrap(); + assert_eq!(read1, 3); + assert_eq!(consumer_reader.outer_position, 3); + // Check that batch 0 is in cache after read_records + assert!(cache.read().unwrap().get(0, BatchID { val: 0 }).is_some()); + + let array1 = consumer_reader.consume_batch().unwrap(); + assert_eq!(array1.len(), 3); + + // After first consume_batch, batch 0 should still be in cache + // (current_batch_id = 3/3 = 1, cleanup only happens if current_batch_id > 1) + assert!(cache.read().unwrap().get(0, BatchID { val: 0 }).is_some()); + + // Read second batch (positions 3-5, batch 1) + let read2 = consumer_reader.read_records(3).unwrap(); + assert_eq!(read2, 3); + assert_eq!(consumer_reader.outer_position, 6); + let array2 = consumer_reader.consume_batch().unwrap(); + assert_eq!(array2.len(), 3); + + // After second consume_batch, batch 0 should be removed + // (current_batch_id = 6/3 = 2, cleanup removes batches 0..(2-1) = 0..1, so removes batch 0) + assert!(cache.read().unwrap().get(0, BatchID { val: 0 }).is_none()); + assert!(cache.read().unwrap().get(0, BatchID { val: 1 }).is_some()); + + // Read third batch (positions 6-8, batch 2) + let read3 = consumer_reader.read_records(3).unwrap(); + assert_eq!(read3, 3); + assert_eq!(consumer_reader.outer_position, 9); + let array3 = consumer_reader.consume_batch().unwrap(); + assert_eq!(array3.len(), 3); + + // After third consume_batch, batches 0 and 1 should be removed + // (current_batch_id = 9/3 = 3, cleanup removes batches 0..(3-1) = 0..2, so removes batches 0 and 1) + assert!(cache.read().unwrap().get(0, BatchID { val: 0 }).is_none()); + assert!(cache.read().unwrap().get(0, BatchID { val: 1 }).is_none()); + assert!(cache.read().unwrap().get(0, BatchID { val: 2 }).is_some()); + } + + #[test] + fn test_producer_keeps_batches() { + let metrics = ArrowReaderMetrics::disabled(); + let mock_reader = MockArrayReader::new(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); + let cache = Arc::new(RwLock::new(RowGroupCache::new(3, usize::MAX))); // Batch size 3 + let mut producer_reader = CachedArrayReader::new( + Box::new(mock_reader), + cache.clone(), + 0, + CacheRole::Producer, + metrics, + ); + + // Read first batch (positions 0-2) + let read1 = producer_reader.read_records(3).unwrap(); + assert_eq!(read1, 3); + let array1 = producer_reader.consume_batch().unwrap(); + assert_eq!(array1.len(), 3); + + // Verify batch 0 is in cache + assert!(cache.read().unwrap().get(0, BatchID { val: 0 }).is_some()); + + // Read second batch (positions 3-5) - producer should NOT remove batch 0 + let read2 = producer_reader.read_records(3).unwrap(); + assert_eq!(read2, 3); + let array2 = producer_reader.consume_batch().unwrap(); + assert_eq!(array2.len(), 3); + + // Verify both batch 0 and batch 1 are still present (no removal for producer) + assert!(cache.read().unwrap().get(0, BatchID { val: 0 }).is_some()); + assert!(cache.read().unwrap().get(0, BatchID { val: 1 }).is_some()); + } + + #[test] + fn test_local_cache_protects_against_eviction() { + let metrics = ArrowReaderMetrics::disabled(); + let mock_reader = MockArrayReader::new(vec![1, 2, 3, 4, 5, 6]); + let cache = Arc::new(RwLock::new(RowGroupCache::new(3, usize::MAX))); // Batch size 3 + let mut cached_reader = CachedArrayReader::new( + Box::new(mock_reader), + cache.clone(), + 0, + CacheRole::Consumer, + metrics, + ); + + // Read records which should populate both shared and local cache + let records_read = cached_reader.read_records(3).unwrap(); + assert_eq!(records_read, 3); + + // Verify data is in both caches + assert!(cache.read().unwrap().get(0, BatchID { val: 0 }).is_some()); + assert!(cached_reader.local_cache.contains_key(&BatchID { val: 0 })); + + // Simulate cache eviction by manually removing from shared cache + cache.write().unwrap().remove(0, BatchID { val: 0 }); + assert!(cache.read().unwrap().get(0, BatchID { val: 0 }).is_none()); + + // Even though shared cache was evicted, consume_batch should still work + // because data is preserved in local cache + let array = cached_reader.consume_batch().unwrap(); + assert_eq!(array.len(), 3); + + let int32_array = array.as_any().downcast_ref::().unwrap(); + assert_eq!(int32_array.values(), &[1, 2, 3]); + + // Local cache should be cleared after consume_batch + assert!(cached_reader.local_cache.is_empty()); + } + + #[test] + fn test_local_cache_is_cleared_properly() { + let metrics = ArrowReaderMetrics::disabled(); + let mock_reader = MockArrayReader::new(vec![1, 2, 3, 4]); + let cache = Arc::new(RwLock::new(RowGroupCache::new(3, 0))); // Batch size 3, cache 0 + let mut cached_reader = CachedArrayReader::new( + Box::new(mock_reader), + cache.clone(), + 0, + CacheRole::Consumer, + metrics, + ); + + // Read records which should populate both shared and local cache + let records_read = cached_reader.read_records(1).unwrap(); + assert_eq!(records_read, 1); + let array = cached_reader.consume_batch().unwrap(); + assert_eq!(array.len(), 1); + + let records_read = cached_reader.read_records(3).unwrap(); + assert_eq!(records_read, 3); + let array = cached_reader.consume_batch().unwrap(); + assert_eq!(array.len(), 3); + } + + #[test] + fn test_batch_id_calculation_with_incremental_reads() { + let metrics = ArrowReaderMetrics::disabled(); + let mock_reader = MockArrayReader::new(vec![1, 2, 3, 4, 5, 6, 7, 8, 9]); + let cache = Arc::new(RwLock::new(RowGroupCache::new(3, usize::MAX))); // Batch size 3 + + // Create a producer to populate cache + let mut producer = CachedArrayReader::new( + Box::new(MockArrayReader::new(vec![1, 2, 3, 4, 5, 6, 7, 8, 9])), + cache.clone(), + 0, + CacheRole::Producer, + metrics.clone(), + ); + + // Populate cache with first batch (1, 2, 3) + producer.read_records(3).unwrap(); + producer.consume_batch().unwrap(); + + // Now create a consumer that will try to read from cache + let mut consumer = CachedArrayReader::new( + Box::new(mock_reader), + cache.clone(), + 0, + CacheRole::Consumer, + metrics, + ); + + // - We want to read 4 records starting from position 0 + // - First 3 records (positions 0-2) should come from cache (batch 0) + // - The 4th record (position 3) should come from the next batch + let records_read = consumer.read_records(4).unwrap(); + assert_eq!(records_read, 4); + + let array = consumer.consume_batch().unwrap(); + assert_eq!(array.len(), 4); + + let int32_array = array.as_any().downcast_ref::().unwrap(); + assert_eq!(int32_array.values(), &[1, 2, 3, 4]); + } +} diff --git a/vendor/parquet/src/arrow/array_reader/empty_array.rs b/vendor/parquet/src/arrow/array_reader/empty_array.rs new file mode 100644 index 0000000..9aba08b --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/empty_array.rs @@ -0,0 +1,88 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::array_reader::ArrayReader; +use crate::errors::Result; +use arrow_array::{ArrayRef, StructArray}; +use arrow_data::ArrayDataBuilder; +use arrow_schema::{DataType as ArrowType, Fields}; +use std::any::Any; +use std::sync::Arc; + +/// Returns an [`ArrayReader`] that yields [`StructArray`] with no columns +/// but with row counts that correspond to the amount of data in the file +/// +/// This is useful for when projection eliminates all columns within a collection +pub fn make_empty_array_reader(row_count: usize) -> Box { + Box::new(EmptyArrayReader::new(row_count)) +} + +struct EmptyArrayReader { + data_type: ArrowType, + remaining_rows: usize, + need_consume_records: usize, +} + +impl EmptyArrayReader { + pub fn new(row_count: usize) -> Self { + Self { + data_type: ArrowType::Struct(Fields::empty()), + remaining_rows: row_count, + need_consume_records: 0, + } + } +} + +impl ArrayReader for EmptyArrayReader { + fn as_any(&self) -> &dyn Any { + self + } + + fn get_data_type(&self) -> &ArrowType { + &self.data_type + } + + fn read_records(&mut self, batch_size: usize) -> Result { + let len = self.remaining_rows.min(batch_size); + self.remaining_rows -= len; + self.need_consume_records += len; + Ok(len) + } + + fn consume_batch(&mut self) -> Result { + let data = ArrayDataBuilder::new(self.data_type.clone()) + .len(self.need_consume_records) + .build() + .unwrap(); + self.need_consume_records = 0; + Ok(Arc::new(StructArray::from(data))) + } + + fn skip_records(&mut self, num_records: usize) -> Result { + let skipped = self.remaining_rows.min(num_records); + self.remaining_rows -= skipped; + Ok(skipped) + } + + fn get_def_levels(&self) -> Option<&[i16]> { + None + } + + fn get_rep_levels(&self) -> Option<&[i16]> { + None + } +} diff --git a/vendor/parquet/src/arrow/array_reader/fixed_len_byte_array.rs b/vendor/parquet/src/arrow/array_reader/fixed_len_byte_array.rs new file mode 100644 index 0000000..e895f79 --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/fixed_len_byte_array.rs @@ -0,0 +1,647 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::array_reader::{ArrayReader, read_records, skip_records}; +use crate::arrow::buffer::bit_util::{iter_set_bits_rev, sign_extend_be}; +use crate::arrow::decoder::{DeltaByteArrayDecoder, DictIndexDecoder}; +use crate::arrow::record_reader::GenericRecordReader; +use crate::arrow::record_reader::buffer::ValuesBuffer; +use crate::arrow::schema::parquet_to_arrow_field; +use crate::basic::{Encoding, Type}; +use crate::column::page::PageIterator; +use crate::column::reader::decoder::ColumnValueDecoder; +use crate::errors::{ParquetError, Result}; +use crate::schema::types::ColumnDescPtr; +use arrow_array::{ + ArrayRef, Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, + FixedSizeBinaryArray, Float16Array, IntervalDayTimeArray, IntervalYearMonthArray, +}; +use arrow_buffer::{Buffer, IntervalDayTime, i256}; +use arrow_data::ArrayDataBuilder; +use arrow_schema::{DataType as ArrowType, IntervalUnit}; +use bytes::Bytes; +use half::f16; +use std::any::Any; +use std::ops::Range; +use std::sync::Arc; + +/// Returns an [`ArrayReader`] that decodes the provided fixed length byte array column +pub fn make_fixed_len_byte_array_reader( + pages: Box, + column_desc: ColumnDescPtr, + arrow_type: Option, + batch_size: usize, + padding_threshold: Option, +) -> Result> { + // Check if Arrow type is specified, else create it from Parquet type + let data_type = match arrow_type { + Some(t) => t, + None => parquet_to_arrow_field(column_desc.as_ref())? + .data_type() + .clone(), + }; + + let byte_length = match column_desc.physical_type() { + Type::FIXED_LEN_BYTE_ARRAY => column_desc.type_length() as usize, + t => { + return Err(general_err!( + "invalid physical type for fixed length byte array reader - {}", + t + )); + } + }; + match &data_type { + ArrowType::FixedSizeBinary(_) => {} + ArrowType::Decimal32(_, _) => { + if byte_length > 4 { + return Err(general_err!( + "decimal 32 type too large, must be less then 4 bytes, got {}", + byte_length + )); + } + } + ArrowType::Decimal64(_, _) => { + if byte_length > 8 { + return Err(general_err!( + "decimal 64 type too large, must be less then 8 bytes, got {}", + byte_length + )); + } + } + ArrowType::Decimal128(_, _) => { + if byte_length > 16 { + return Err(general_err!( + "decimal 128 type too large, must be less than 16 bytes, got {}", + byte_length + )); + } + } + ArrowType::Decimal256(_, _) => { + if byte_length > 32 { + return Err(general_err!( + "decimal 256 type too large, must be less than 32 bytes, got {}", + byte_length + )); + } + } + ArrowType::Interval(_) => { + if byte_length != 12 { + // https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#interval + return Err(general_err!( + "interval type must consist of 12 bytes got {}", + byte_length + )); + } + } + ArrowType::Float16 => { + if byte_length != 2 { + return Err(general_err!( + "float 16 type must be 2 bytes, got {}", + byte_length + )); + } + } + _ => { + return Err(general_err!( + "invalid data type for fixed length byte array reader - {}", + data_type + )); + } + } + + Ok(Box::new(FixedLenByteArrayReader::new( + pages, + column_desc, + data_type, + byte_length, + batch_size, + padding_threshold, + ))) +} + +struct FixedLenByteArrayReader { + data_type: ArrowType, + byte_length: usize, + pages: Box, + def_levels_buffer: Option>, + rep_levels_buffer: Option>, + record_reader: GenericRecordReader, +} + +impl FixedLenByteArrayReader { + fn new( + pages: Box, + column_desc: ColumnDescPtr, + data_type: ArrowType, + byte_length: usize, + batch_size: usize, + padding_threshold: Option, + ) -> Self { + let mut record_reader = GenericRecordReader::new(column_desc, batch_size); + if let Some(threshold) = padding_threshold { + record_reader.set_padding_threshold(threshold); + } + Self { + data_type, + byte_length, + pages, + def_levels_buffer: None, + rep_levels_buffer: None, + record_reader, + } + } +} + +impl ArrayReader for FixedLenByteArrayReader { + fn as_any(&self) -> &dyn Any { + self + } + + fn get_data_type(&self) -> &ArrowType { + &self.data_type + } + + fn read_records(&mut self, batch_size: usize) -> Result { + read_records(&mut self.record_reader, self.pages.as_mut(), batch_size) + } + + fn consume_batch(&mut self) -> Result { + let len = self.record_reader.values_written(); + + let record_data = self.record_reader.consume_record_data(); + debug_assert_eq!( + record_data.buffer.len(), + len * self.byte_length, + "fixed-len byte array buffer size mismatch: {} bytes for {} elements of size {}", + record_data.buffer.len(), + len, + self.byte_length, + ); + let null_bit_buffer = self.record_reader.consume_compact_bitmap(); + + let array_data = ArrayDataBuilder::new(ArrowType::FixedSizeBinary(self.byte_length as i32)) + .len(len) + .add_buffer(Buffer::from_vec(record_data.buffer)) + .null_bit_buffer(null_bit_buffer); + + let binary = FixedSizeBinaryArray::from(unsafe { array_data.build_unchecked() }); + + // TODO: An improvement might be to do this conversion on read + // Note the conversions below apply to all elements regardless of null slots as the + // conversion lambdas are all infallible. This improves performance by avoiding a branch in + // the inner loop (see docs for `PrimitiveArray::from_unary`). + let array: ArrayRef = match &self.data_type { + ArrowType::Decimal32(p, s) => { + let f = |b: &[u8]| i32::from_be_bytes(sign_extend_be(b)); + Arc::new(Decimal32Array::from_unary(&binary, f).with_precision_and_scale(*p, *s)?) + as ArrayRef + } + ArrowType::Decimal64(p, s) => { + let f = |b: &[u8]| i64::from_be_bytes(sign_extend_be(b)); + Arc::new(Decimal64Array::from_unary(&binary, f).with_precision_and_scale(*p, *s)?) + as ArrayRef + } + ArrowType::Decimal128(p, s) => { + let f = |b: &[u8]| i128::from_be_bytes(sign_extend_be(b)); + Arc::new(Decimal128Array::from_unary(&binary, f).with_precision_and_scale(*p, *s)?) + as ArrayRef + } + ArrowType::Decimal256(p, s) => { + let f = |b: &[u8]| i256::from_be_bytes(sign_extend_be(b)); + Arc::new(Decimal256Array::from_unary(&binary, f).with_precision_and_scale(*p, *s)?) + as ArrayRef + } + ArrowType::Interval(unit) => { + // An interval is stored as 3x 32-bit unsigned integers storing months, days, + // and milliseconds + match unit { + IntervalUnit::YearMonth => { + let f = |b: &[u8]| i32::from_le_bytes(b[0..4].try_into().unwrap()); + Arc::new(IntervalYearMonthArray::from_unary(&binary, f)) as ArrayRef + } + IntervalUnit::DayTime => { + let f = |b: &[u8]| { + IntervalDayTime::new( + i32::from_le_bytes(b[4..8].try_into().unwrap()), + i32::from_le_bytes(b[8..12].try_into().unwrap()), + ) + }; + Arc::new(IntervalDayTimeArray::from_unary(&binary, f)) as ArrayRef + } + IntervalUnit::MonthDayNano => { + return Err(nyi_err!("MonthDayNano intervals not supported")); + } + } + } + ArrowType::Float16 => { + let f = |b: &[u8]| f16::from_le_bytes(b[..2].try_into().unwrap()); + Arc::new(Float16Array::from_unary(&binary, f)) as ArrayRef + } + _ => Arc::new(binary) as ArrayRef, + }; + + self.def_levels_buffer = self.record_reader.consume_def_levels(); + self.rep_levels_buffer = self.record_reader.consume_rep_levels(); + self.record_reader.reset(); + + Ok(array) + } + + fn skip_records(&mut self, num_records: usize) -> Result { + skip_records(&mut self.record_reader, self.pages.as_mut(), num_records) + } + + fn get_def_levels(&self) -> Option<&[i16]> { + self.def_levels_buffer.as_deref() + } + + fn get_rep_levels(&self) -> Option<&[i16]> { + self.rep_levels_buffer.as_deref() + } + + fn max_def_level(&self) -> i16 { + self.record_reader.max_def_level() + } +} + +#[derive(Default)] +struct FixedLenByteArrayBuffer { + buffer: Vec, + /// The length of each element in bytes + byte_length: Option, + /// Preserved value-count hint used to allocate `buffer` once `byte_length` + /// becomes known on the first decode. + values_capacity: Option, +} + +#[inline] +fn move_values( + buffer: &mut Vec, + byte_length: usize, + values_range: Range, + valid_mask: &[u8], + mut op: F, +) where + F: FnMut(&mut Vec, usize, usize, usize), +{ + for (value_pos, level_pos) in values_range.rev().zip(iter_set_bits_rev(valid_mask)) { + debug_assert!(level_pos >= value_pos); + if level_pos <= value_pos { + break; + } + + let level_pos_bytes = level_pos * byte_length; + let value_pos_bytes = value_pos * byte_length; + + op(buffer, level_pos_bytes, value_pos_bytes, byte_length) + } +} + +impl ValuesBuffer for FixedLenByteArrayBuffer { + fn with_capacity(capacity: usize) -> Self { + // `byte_length` is not known initially, so preserve the value-count + // hint so the first decode can allocate the initial byte capacity. + Self { + buffer: Vec::new(), + byte_length: None, + values_capacity: Some(capacity), + } + } + + fn reserve_exact(&mut self, additional: usize) { + match self.byte_length { + Some(byte_length) => self + .buffer + .reserve_exact(additional.saturating_mul(byte_length)), + None => { + let capacity = self.values_capacity.get_or_insert(0); + *capacity = (*capacity).max(additional); + } + } + } + + fn pad_nulls( + &mut self, + read_offset: usize, + values_read: usize, + levels_read: usize, + valid_mask: &[u8], + ) -> Result<()> { + let byte_length = self.byte_length.unwrap_or_default(); + + if self.buffer.len() != (read_offset + values_read) * byte_length { + return Err(general_err!( + "found inconsistent buffer length while padding nulls: expected {} bytes, got {}", + (read_offset + values_read) * byte_length, + self.buffer.len() + )); + } + self.buffer + .resize((read_offset + levels_read) * byte_length, 0); + + let values_range = read_offset..read_offset + values_read; + // Move the bytes from value_pos to level_pos. For values of `byte_length` <= 4, + // the simple loop is preferred as the compiler can eliminate the loop via unrolling. + // For `byte_length > 4`, we instead copy from non-overlapping slices. This allows + // the loop to be vectorized, yielding much better performance. + const VEC_CUTOFF: usize = 4; + if byte_length > VEC_CUTOFF { + let op = |buffer: &mut Vec, level_pos_bytes, value_pos_bytes, byte_length| { + let split = buffer.split_at_mut(level_pos_bytes); + let dst = &mut split.1[..byte_length]; + let src = &split.0[value_pos_bytes..value_pos_bytes + byte_length]; + dst.copy_from_slice(src); + }; + move_values(&mut self.buffer, byte_length, values_range, valid_mask, op); + } else { + let op = |buffer: &mut Vec, level_pos_bytes, value_pos_bytes, byte_length| { + for i in 0..byte_length { + buffer[level_pos_bytes + i] = buffer[value_pos_bytes + i] + } + }; + move_values(&mut self.buffer, byte_length, values_range, valid_mask, op); + } + Ok(()) + } +} + +struct ValueDecoder { + byte_length: usize, + dict_page: Option, + decoder: Option, +} + +impl ColumnValueDecoder for ValueDecoder { + type Buffer = FixedLenByteArrayBuffer; + + fn new(col: &ColumnDescPtr) -> Self { + Self { + byte_length: col.type_length() as usize, + dict_page: None, + decoder: None, + } + } + + fn set_dict( + &mut self, + buf: Bytes, + num_values: u32, + encoding: Encoding, + _is_sorted: bool, + ) -> Result<()> { + if !matches!( + encoding, + Encoding::PLAIN | Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY + ) { + return Err(nyi_err!( + "Invalid/Unsupported encoding type for dictionary: {}", + encoding + )); + } + let expected_len = num_values as usize * self.byte_length; + if expected_len > buf.len() { + return Err(general_err!( + "too few bytes in dictionary page, expected {} got {}", + expected_len, + buf.len() + )); + } + + self.dict_page = Some(buf); + Ok(()) + } + + fn set_data( + &mut self, + encoding: Encoding, + data: Bytes, + num_levels: usize, + num_values: Option, + ) -> Result<()> { + self.decoder = Some(match encoding { + Encoding::PLAIN => Decoder::Plain { + buf: data, + offset: 0, + }, + Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY => Decoder::Dict { + decoder: DictIndexDecoder::new(data, num_levels, num_values)?, + }, + Encoding::DELTA_BYTE_ARRAY => Decoder::Delta { + decoder: DeltaByteArrayDecoder::new(data)?, + }, + Encoding::BYTE_STREAM_SPLIT => Decoder::ByteStreamSplit { + buf: data, + offset: 0, + }, + _ => { + return Err(general_err!( + "unsupported encoding for fixed length byte array: {}", + encoding + )); + } + }); + Ok(()) + } + + fn read(&mut self, out: &mut Self::Buffer, num_values: usize) -> Result { + match out.byte_length { + Some(x) => assert_eq!(x, self.byte_length), + None => { + out.byte_length = Some(self.byte_length); + // TODO: collapse to a let-chain once MSRV ≥ 1.88 + // (`if out.buffer.is_empty() && let Some(cap) = out.values_capacity.take()`) + if out.buffer.is_empty() { + if let Some(values_capacity) = out.values_capacity.take() { + // now that the byte length per output element is known, + // allocate the actual needed space. + let byte_capacity = values_capacity.saturating_mul(self.byte_length); + out.buffer = Vec::with_capacity(byte_capacity); + } + } + } + } + + match self.decoder.as_mut().unwrap() { + Decoder::Plain { offset, buf } => { + let to_read = + (num_values * self.byte_length).min(buf.len() - *offset) / self.byte_length; + let end_offset = *offset + to_read * self.byte_length; + out.buffer + .extend_from_slice(&buf.as_ref()[*offset..end_offset]); + *offset = end_offset; + Ok(to_read) + } + Decoder::Dict { decoder } => { + let dict = self.dict_page.as_ref().unwrap(); + // All data must be NULL + if dict.is_empty() { + return Ok(0); + } + + decoder.read(num_values, |keys| { + out.buffer.reserve(keys.len() * self.byte_length); + for key in keys { + let offset = *key as usize * self.byte_length; + let val = &dict.as_ref()[offset..offset + self.byte_length]; + out.buffer.extend_from_slice(val); + } + Ok(()) + }) + } + Decoder::Delta { decoder } => { + let to_read = num_values.min(decoder.remaining()); + out.buffer.reserve(to_read * self.byte_length); + + decoder.read(to_read, |slice| { + if slice.len() != self.byte_length { + return Err(general_err!( + "encountered array with incorrect length, got {} expected {}", + slice.len(), + self.byte_length + )); + } + out.buffer.extend_from_slice(slice); + Ok(()) + }) + } + Decoder::ByteStreamSplit { buf, offset } => { + // we have n=`byte_length` streams of length `buf.len/byte_length` + // to read value i, we need the i'th byte from each of the streams + // so `offset` should be the value offset, not the byte offset + let total_values = buf.len() / self.byte_length; + let to_read = num_values.min(total_values - *offset); + + // now read the n streams and reassemble values into the output buffer + read_byte_stream_split(&mut out.buffer, buf, *offset, to_read, self.byte_length); + + *offset += to_read; + Ok(to_read) + } + } + } + + fn skip_values(&mut self, num_values: usize) -> Result { + match self.decoder.as_mut().unwrap() { + Decoder::Plain { offset, buf } => { + let to_read = num_values.min((buf.len() - *offset) / self.byte_length); + *offset += to_read * self.byte_length; + Ok(to_read) + } + Decoder::Dict { decoder } => decoder.skip(num_values), + Decoder::Delta { decoder } => decoder.skip(num_values), + Decoder::ByteStreamSplit { offset, buf } => { + let total_values = buf.len() / self.byte_length; + let to_read = num_values.min(total_values - *offset); + *offset += to_read; + Ok(to_read) + } + } + } +} + +// `src` is an array laid out like a NxM matrix where N == `data_width` and +// M == total_values_in_src. Each "row" of the matrix is a stream of bytes, with stream `i` +// containing the `ith` byte for each value. Each "column" is a single value. +// This will reassemble `num_values` values by reading columns of the matrix starting at +// `offset`. Values will be appended to `dst`. +fn read_byte_stream_split( + dst: &mut Vec, + src: &mut Bytes, + offset: usize, + num_values: usize, + data_width: usize, +) { + let stride = src.len() / data_width; + let idx = dst.len(); + dst.resize(idx + num_values * data_width, 0u8); + let dst_slc = &mut dst[idx..idx + num_values * data_width]; + for j in 0..data_width { + let src_slc = &src[offset + j * stride..offset + j * stride + num_values]; + for i in 0..num_values { + dst_slc[i * data_width + j] = src_slc[i]; + } + } +} + +enum Decoder { + Plain { buf: Bytes, offset: usize }, + Dict { decoder: DictIndexDecoder }, + Delta { decoder: DeltaByteArrayDecoder }, + ByteStreamSplit { buf: Bytes, offset: usize }, +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::ArrowWriter; + use crate::arrow::arrow_reader::ParquetRecordBatchReader; + use arrow::datatypes::Field; + use arrow::error::Result as ArrowResult; + use arrow_array::{Array, ListArray}; + use arrow_array::{Decimal256Array, RecordBatch}; + use bytes::Bytes; + use std::sync::Arc; + + #[test] + fn test_pending_reservation_tracks_max_capacity() { + let mut buffer = FixedLenByteArrayBuffer::with_capacity(0); + + buffer.reserve_exact(8); + buffer.reserve_exact(4); + assert_eq!(buffer.values_capacity, Some(8)); + + buffer.reserve_exact(12); + assert_eq!(buffer.values_capacity, Some(12)); + } + + #[test] + fn test_decimal_list() { + let decimals = Decimal256Array::from_iter_values( + [1, 2, 3, 4, 5, 6, 7, 8].into_iter().map(i256::from_i128), + ); + + // [[], [1], [2, 3], null, [4], null, [6, 7, 8]] + let data = ArrayDataBuilder::new(ArrowType::List(Arc::new(Field::new_list_field( + decimals.data_type().clone(), + false, + )))) + .len(7) + .add_buffer(Buffer::from_iter([0_i32, 0, 1, 3, 3, 4, 5, 8])) + .null_bit_buffer(Some(Buffer::from(&[0b01010111]))) + .child_data(vec![decimals.into_data()]) + .build() + .unwrap(); + + let written = + RecordBatch::try_from_iter([("list", Arc::new(ListArray::from(data)) as ArrayRef)]) + .unwrap(); + + let mut buffer = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buffer, written.schema(), None).unwrap(); + writer.write(&written).unwrap(); + writer.close().unwrap(); + + let read = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 3) + .unwrap() + .collect::>>() + .unwrap(); + + assert_eq!(&written.slice(0, 3), &read[0]); + assert_eq!(&written.slice(3, 3), &read[1]); + assert_eq!(&written.slice(6, 1), &read[2]); + } +} diff --git a/vendor/parquet/src/arrow/array_reader/fixed_size_list_array.rs b/vendor/parquet/src/arrow/array_reader/fixed_size_list_array.rs new file mode 100644 index 0000000..976aeb9 --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/fixed_size_list_array.rs @@ -0,0 +1,631 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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::sync::Arc; + +use crate::arrow::array_reader::ArrayReader; +use crate::arrow::record_reader::definition_levels::build_filtered_validity_bitmap; +use crate::errors::ParquetError; +use crate::errors::Result; +use arrow_array::FixedSizeListArray; +use arrow_array::{Array, ArrayRef, builder::BooleanBufferBuilder, new_empty_array}; +use arrow_data::{ArrayData, transform::MutableArrayData}; +use arrow_schema::DataType as ArrowType; + +/// Implementation of fixed-size list array reader. +pub struct FixedSizeListArrayReader { + item_reader: Box, + /// The number of child items in each row of the list array + fixed_size: usize, + data_type: ArrowType, + /// The definition level at which this list is not null + def_level: i16, + /// The repetition level that corresponds to a new value in this array + rep_level: i16, + /// If the list is nullable + nullable: bool, + /// When set, this reader is a child of another list and should + /// exclude entries where def < threshold from its output. + parent_threshold: Option, +} + +impl FixedSizeListArrayReader { + /// Construct fixed-size list array reader. + /// + /// The child reader must already be constructed with a padding + /// threshold of `def_level`. `parent_threshold` is set when this + /// reader is itself a child of another list. + pub fn new( + item_reader: Box, + fixed_size: usize, + data_type: ArrowType, + def_level: i16, + rep_level: i16, + nullable: bool, + parent_threshold: Option, + ) -> Self { + Self { + item_reader, + fixed_size, + data_type, + def_level, + rep_level, + nullable, + parent_threshold, + } + } +} + +impl ArrayReader for FixedSizeListArrayReader { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn get_data_type(&self) -> &ArrowType { + &self.data_type + } + + fn read_records(&mut self, batch_size: usize) -> Result { + let size = self.item_reader.read_records(batch_size)?; + Ok(size) + } + + /// Consume batch. The compact child array contains entries only for + /// items within non-null fixed-size lists. Null lists get `fixed_size` + /// null entries inserted via MutableArrayData. When all lists are + /// non-null the compact child is used directly without any copying. + fn consume_batch(&mut self) -> Result { + let child_array = self.item_reader.consume_batch()?; + + let def_levels = self + .get_def_levels() + .ok_or_else(|| general_err!("item_reader def levels are None"))?; + let rep_levels = self + .get_rep_levels() + .ok_or_else(|| general_err!("item_reader rep levels are None"))?; + + if def_levels.is_empty() { + return Ok(new_empty_array(&self.data_type)); + } + + if rep_levels[0] != 0 { + return Err(general_err!("first repetition level of batch must be 0")); + } + + // Single pass: collect the def level at each list boundary and + // validate that each non-null row contains exactly fixed_size items. + let mut boundary_defs: Vec = Vec::with_capacity(def_levels.len()); + let mut current_row_len = None; + for (&d, &r) in def_levels.iter().zip(rep_levels) { + if let Some(threshold) = self.parent_threshold { + if d < threshold { + continue; + } + } + if r < self.rep_level { + if let Some(row_len) = current_row_len.take().filter(|&len| len != self.fixed_size) + { + return Err(general_err!( + "Encountered misaligned row with length {} (expected length {})", + row_len, + self.fixed_size + )); + } + + boundary_defs.push(d); + current_row_len = (d >= self.def_level).then_some(1); + } else if r == self.rep_level { + if let Some(row_len) = current_row_len.as_mut() { + *row_len += 1; + } + } else if d < self.def_level { + return Err(general_err!( + "Encountered repetition level too large for definition level" + )); + } + } + + if let Some(row_len) = current_row_len.filter(|&len| len != self.fixed_size) { + return Err(general_err!( + "Encountered misaligned row with length {} (expected length {})", + row_len, + self.fixed_size + )); + } + + let list_len = boundary_defs.len(); + let items_count = boundary_defs + .iter() + .filter(|&&d| d >= self.def_level) + .count(); + + if child_array.len() != items_count * self.fixed_size { + return Err(general_err!( + "Fixed-size list child mismatch: expected {} entries but got {}", + items_count * self.fixed_size, + child_array.len() + )); + } + + // Build child data. If all lists are non-null, reuse the compact + // child directly. Otherwise insert fixed_size null blocks. + let child_data = if items_count == list_len { + child_array.to_data() + } else { + let data = child_array.to_data(); + let mut builder = MutableArrayData::new(vec![&data], true, list_len * self.fixed_size); + let mut src_idx = 0; + for &d in &boundary_defs { + if d >= self.def_level { + builder + .try_extend(0, src_idx, src_idx + self.fixed_size) + .map_err(|e| general_err!("{}", e))?; + src_idx += self.fixed_size; + } else { + builder + .try_extend_nulls(self.fixed_size) + .map_err(|e| general_err!("{}", e))?; + } + } + builder.freeze() + }; + + debug_assert_eq!(child_data.len(), list_len * self.fixed_size); + + let mut list_builder = ArrayData::builder(self.get_data_type().clone()) + .len(list_len) + .add_child_data(child_data); + + if self.nullable { + let mut validity = BooleanBufferBuilder::new(list_len); + build_filtered_validity_bitmap( + &boundary_defs, + None, + None, + self.def_level - 1, + &mut validity, + ); + list_builder = list_builder.null_bit_buffer(Some(validity.into())); + } + + let list_data = unsafe { list_builder.build_unchecked() }; + Ok(Arc::new(FixedSizeListArray::from(list_data))) + } + + fn skip_records(&mut self, num_records: usize) -> Result { + self.item_reader.skip_records(num_records) + } + + fn get_def_levels(&self) -> Option<&[i16]> { + self.item_reader.get_def_levels() + } + + fn get_rep_levels(&self) -> Option<&[i16]> { + self.item_reader.get_rep_levels() + } + + fn max_def_level(&self) -> i16 { + self.def_level + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::{ + ArrowWriter, + array_reader::{ListArrayReader, test_util::make_int32_page_reader}, + arrow_reader::{ArrowReaderBuilder, ArrowReaderOptions, ParquetRecordBatchReader}, + }; + use arrow::datatypes::{Field, Int32Type}; + use arrow_array::{ + FixedSizeListArray, ListArray, PrimitiveArray, RecordBatch, + builder::{FixedSizeListBuilder, Int32Builder, ListBuilder}, + cast::AsArray, + }; + use arrow_buffer::Buffer; + use arrow_data::ArrayDataBuilder; + use arrow_schema::Schema; + use bytes::Bytes; + + #[test] + fn test_nullable_list() { + // [null, [1, null, 2], null, [3, 4, 5], [null, null, null]] + let expected = FixedSizeListArray::from_iter_primitive::( + vec![ + None, + Some([Some(1), None, Some(2)]), + None, + Some([Some(3), Some(4), Some(5)]), + Some([None, None, None]), + ], + 3, + ); + + let item_array_reader = make_int32_page_reader( + &[1, 2, 3, 4, 5], + &[0, 3, 2, 3, 0, 3, 3, 3, 2, 2, 2], + &[0, 0, 1, 1, 0, 0, 1, 1, 0, 1, 1], + 3, + 1, + Some(2), + ); + + let mut list_array_reader = FixedSizeListArrayReader::new( + item_array_reader, + 3, + ArrowType::FixedSizeList(Arc::new(Field::new_list_field(ArrowType::Int32, true)), 3), + 2, + 1, + true, + None, + ); + let actual = list_array_reader.next_batch(1024).unwrap(); + let actual = actual + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(&expected, actual) + } + + #[test] + fn test_required_list() { + // [[1, null], [2, 3], [null, null], [4, 5]] + let expected = FixedSizeListArray::from_iter_primitive::( + vec![ + Some([Some(1), None]), + Some([Some(2), Some(3)]), + Some([None, None]), + Some([Some(4), Some(5)]), + ], + 2, + ); + + let item_array_reader = make_int32_page_reader( + &[1, 2, 3, 4, 5], + &[2, 1, 2, 2, 1, 1, 2, 2], + &[0, 1, 0, 1, 0, 1, 0, 1], + 2, + 1, + Some(1), + ); + + let mut list_array_reader = FixedSizeListArrayReader::new( + item_array_reader, + 2, + ArrowType::FixedSizeList(Arc::new(Field::new_list_field(ArrowType::Int32, true)), 2), + 1, + 1, + false, + None, + ); + let actual = list_array_reader.next_batch(1024).unwrap(); + let actual = actual + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(&expected, actual) + } + + #[test] + fn test_misaligned_rows_error() { + // Two rows with fixed_size == 2 are encoded with row lengths 1 and 3. + // The total child count is still 4, so aggregate count validation alone + // would silently repartition the values into two valid-looking rows. + let item_array_reader = + make_int32_page_reader(&[1, 2, 3, 4], &[2, 2, 2, 2], &[0, 0, 1, 1], 2, 1, Some(1)); + + let mut list_array_reader = FixedSizeListArrayReader::new( + item_array_reader, + 2, + ArrowType::FixedSizeList(Arc::new(Field::new_list_field(ArrowType::Int32, true)), 2), + 1, + 1, + false, + None, + ); + + let err = list_array_reader.next_batch(1024).unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: Encountered misaligned row with length 1 (expected length 2)" + ); + } + + #[test] + fn test_nested_list() { + // [ + // null, + // [[1, 2]], + // [[null, 3]], + // null, + // [[4, 5]], + // [[null, null]], + // ] + let l2_type = + ArrowType::FixedSizeList(Arc::new(Field::new_list_field(ArrowType::Int32, true)), 2); + let l1_type = + ArrowType::FixedSizeList(Arc::new(Field::new_list_field(l2_type.clone(), false)), 1); + + let array = PrimitiveArray::::from(vec![ + None, + None, + Some(1), + Some(2), + None, + Some(3), + None, + None, + Some(4), + Some(5), + None, + None, + ]); + + let l2 = ArrayDataBuilder::new(l2_type.clone()) + .len(6) + .add_child_data(array.into_data()) + .build() + .unwrap(); + + let l1 = ArrayDataBuilder::new(l1_type.clone()) + .len(6) + .add_child_data(l2) + .null_bit_buffer(Some(Buffer::from([0b110110]))) + .build() + .unwrap(); + + let expected = FixedSizeListArray::from(l1); + + let item_array_reader = make_int32_page_reader( + &[1, 2, 3, 4, 5], + &[0, 5, 5, 4, 5, 0, 5, 5, 4, 4], + &[0, 0, 2, 0, 2, 0, 0, 2, 0, 2], + 5, + 2, + Some(4), + ); + + let l2 = FixedSizeListArrayReader::new(item_array_reader, 2, l2_type, 4, 2, false, Some(3)); + let mut l1 = FixedSizeListArrayReader::new(Box::new(l2), 1, l1_type, 3, 1, true, None); + + let expected_1 = expected.slice(0, 2); + let expected_2 = expected.slice(2, 4); + + let actual = l1.next_batch(2).unwrap(); + assert_eq!(actual.as_ref(), &expected_1); + + let actual = l1.next_batch(1024).unwrap(); + assert_eq!(actual.as_ref(), &expected_2); + } + + #[test] + fn test_empty_list() { + // [null, [], null, []] + let expected = FixedSizeListArray::from_iter_primitive::( + vec![None, Some([]), None, Some([])], + 0, + ); + + let item_array_reader = + make_int32_page_reader(&[], &[0, 1, 0, 1], &[0, 0, 0, 0], 2, 1, Some(2)); + + let mut list_array_reader = FixedSizeListArrayReader::new( + item_array_reader, + 0, + ArrowType::FixedSizeList(Arc::new(Field::new_list_field(ArrowType::Int32, true)), 0), + 2, + 1, + true, + None, + ); + let actual = list_array_reader.next_batch(1024).unwrap(); + let actual = actual + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(&expected, actual) + } + + #[test] + fn test_nested_var_list() { + // [[[1, null, 3], null], [[4], []], [[5, 6], [null, null]], null] + let mut builder = FixedSizeListBuilder::new(ListBuilder::new(Int32Builder::new()), 2); + builder.values().append_value([Some(1), None, Some(3)]); + builder.values().append_null(); + builder.append(true); + builder.values().append_value([Some(4)]); + builder.values().append_value([]); + builder.append(true); + builder.values().append_value([Some(5), Some(6)]); + builder.values().append_value([None, None]); + builder.append(true); + builder.values().append_null(); + builder.values().append_null(); + builder.append(false); + let expected = builder.finish(); + + let inner_type = ArrowType::List(Arc::new(Field::new_list_field(ArrowType::Int32, true))); + let list_type = + ArrowType::FixedSizeList(Arc::new(Field::new_list_field(inner_type.clone(), true)), 2); + + let item_array_reader = make_int32_page_reader( + &[1, 3, 4, 5, 6], + &[5, 4, 5, 2, 5, 3, 5, 5, 4, 4, 0], + &[0, 2, 2, 1, 0, 1, 0, 2, 1, 2, 0], + 5, + 2, + Some(4), + ); + + let inner_array_reader = + ListArrayReader::::new(item_array_reader, inner_type, 4, 2, true, Some(2)); + + let mut list_array_reader = FixedSizeListArrayReader::new( + Box::new(inner_array_reader), + 2, + list_type, + 2, + 1, + true, + None, + ); + let actual = list_array_reader.next_batch(1024).unwrap(); + let actual = actual + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(&expected, actual) + } + + #[test] + fn test_read_list_column() { + // This test writes a Parquet file containing a fixed-length array column and a primitive column, + // then reads the columns back from the file. + + // [ + // [1, 2, 3, null], + // [5, 6, 7, 8], + // null, + // [9, null, 11, 12], + // ] + let list = FixedSizeListArray::from_iter_primitive::( + vec![ + Some(vec![Some(1), Some(2), Some(3), None]), + Some(vec![Some(5), Some(6), Some(7), Some(8)]), + None, + Some(vec![Some(9), None, Some(11), Some(12)]), + Some(vec![None, None, None, None]), + ], + 4, + ); + + // [null, 2, 3, null, 5] + let primitive = + PrimitiveArray::::from_iter(vec![None, Some(2), Some(3), None, Some(5)]); + + let schema = Arc::new(Schema::new(vec![ + Field::new( + "list", + ArrowType::FixedSizeList( + Arc::new(Field::new_list_field(ArrowType::Int32, true)), + 4, + ), + true, + ), + Field::new("primitive", ArrowType::Int32, true), + ])); + + // Create record batch with a fixed-length array column and a primitive column + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(list.clone()), Arc::new(primitive.clone())], + ) + .expect("unable to create record batch"); + + // Write record batch to Parquet + let mut buffer = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(), None) + .expect("unable to create parquet writer"); + writer.write(&batch).expect("unable to write record batch"); + writer.close().expect("unable to close parquet writer"); + + // Read record batch from Parquet + let reader = Bytes::from(buffer); + let mut batch_reader = ParquetRecordBatchReader::try_new(reader, 1024) + .expect("unable to create parquet reader"); + let actual = batch_reader + .next() + .expect("missing record batch") + .expect("unable to read record batch"); + + // Verify values of both read columns match + assert_eq!(schema, actual.schema()); + let actual_list = actual + .column(0) + .as_any() + .downcast_ref::() + .expect("unable to cast array to FixedSizeListArray"); + let actual_primitive = actual.column(1).as_primitive::(); + assert_eq!(actual_list, &list); + assert_eq!(actual_primitive, &primitive); + } + + #[test] + fn test_read_as_dyn_list() { + // This test verifies that fixed-size list arrays can be read from Parquet + // as variable-length list arrays. + + // [ + // [1, 2, 3, null], + // [5, 6, 7, 8], + // null, + // [9, null, 11, 12], + // ] + let list = FixedSizeListArray::from_iter_primitive::( + vec![ + Some(vec![Some(1), Some(2), Some(3), None]), + Some(vec![Some(5), Some(6), Some(7), Some(8)]), + None, + Some(vec![Some(9), None, Some(11), Some(12)]), + Some(vec![None, None, None, None]), + ], + 4, + ); + + let schema = Arc::new(Schema::new(vec![Field::new( + "list", + ArrowType::FixedSizeList(Arc::new(Field::new_list_field(ArrowType::Int32, true)), 4), + true, + )])); + + // Create record batch with a single fixed-length array column + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(list)]).unwrap(); + + // Write record batch to Parquet + let mut buffer = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buffer, schema, None) + .expect("unable to create parquet writer"); + writer.write(&batch).expect("unable to write record batch"); + writer.close().expect("unable to close parquet writer"); + + // Read record batch from Parquet - ignoring arrow metadata + let reader = Bytes::from(buffer); + let mut batch_reader = ArrowReaderBuilder::try_new_with_options( + reader, + ArrowReaderOptions::new().with_skip_arrow_metadata(true), + ) + .expect("unable to create reader builder") + .build() + .expect("unable to create parquet reader"); + let actual = batch_reader + .next() + .expect("missing record batch") + .expect("unable to read record batch"); + + // Verify the read column is a variable length list with values that match the input + let col = actual.column(0).as_list::(); + let expected = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2), Some(3), None]), + Some(vec![Some(5), Some(6), Some(7), Some(8)]), + None, + Some(vec![Some(9), None, Some(11), Some(12)]), + Some(vec![None, None, None, None]), + ]); + assert_eq!(col, &expected); + } +} diff --git a/vendor/parquet/src/arrow/array_reader/list_array.rs b/vendor/parquet/src/arrow/array_reader/list_array.rs new file mode 100644 index 0000000..66c93b7 --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/list_array.rs @@ -0,0 +1,696 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::array_reader::ArrayReader; +use crate::errors::ParquetError; +use crate::errors::Result; +use arrow_array::{ + Array, ArrayRef, GenericListArray, OffsetSizeTrait, builder::BooleanBufferBuilder, + new_empty_array, +}; +use arrow_buffer::Buffer; +use arrow_buffer::ToByteSlice; +use arrow_data::ArrayData; +use arrow_schema::DataType as ArrowType; +use std::any::Any; +use std::marker::PhantomData; +use std::sync::Arc; + +/// Implementation of list array reader. +pub struct ListArrayReader { + item_reader: Box, + data_type: ArrowType, + /// The definition level at which this list is not null + def_level: i16, + /// The repetition level that corresponds to a new value in this array + rep_level: i16, + /// If this list is nullable + nullable: bool, + /// When set, this list is itself a child of another list and should + /// exclude entries where def < threshold from its output. + parent_threshold: Option, + _marker: PhantomData, +} + +impl ListArrayReader { + /// Construct list array reader. + /// + /// The child reader must already be constructed with a padding + /// threshold of `def_level`, so it produces item-level padded arrays + /// (values + null items, no list-level padding). `consume_batch` + /// then computes offsets directly from def/rep levels without + /// MutableArrayData. + /// + /// `parent_threshold` is set when this list is itself a child of + /// another list — entries where `def < threshold` are excluded from + /// this list's output. + pub fn new( + item_reader: Box, + data_type: ArrowType, + def_level: i16, + rep_level: i16, + nullable: bool, + parent_threshold: Option, + ) -> Self { + Self { + item_reader, + data_type, + def_level, + rep_level, + nullable, + parent_threshold, + _marker: PhantomData, + } + } +} + +/// Implementation of ListArrayReader. +impl ArrayReader for ListArrayReader { + fn as_any(&self) -> &dyn Any { + self + } + + /// Returns data type. + /// This must be a List. + fn get_data_type(&self) -> &ArrowType { + &self.data_type + } + + fn read_records(&mut self, batch_size: usize) -> Result { + let size = self.item_reader.read_records(batch_size)?; + Ok(size) + } + + fn consume_batch(&mut self) -> Result { + let child_array = self.item_reader.consume_batch()?; + + let def_levels = self + .item_reader + .get_def_levels() + .ok_or_else(|| general_err!("item_reader def levels are None."))?; + + let rep_levels = self + .item_reader + .get_rep_levels() + .ok_or_else(|| general_err!("item_reader rep levels are None."))?; + + if def_levels.is_empty() { + return Ok(new_empty_array(&self.data_type)); + } + + if rep_levels[0] != 0 { + return Err(general_err!("first repetition level of batch must be 0")); + } + + if OffsetSize::from_usize(child_array.len()).is_none() { + return Err(general_err!( + "offset of {} would overflow list array", + child_array.len() + )); + } + + // Definition levels identify whether each list boundary contributes a + // child slot. For a nullable list, the states are: + // + // d >= self.def_level : list is present with a child item + // d == self.def_level - 1 : list is present but empty + // d <= self.def_level - 2 : list is null + // + // Required lists do not have the null state, but still use the same + // `d >= self.def_level` test to distinguish child items from empty + // lists. Repetition levels identify list boundaries and whether a + // child item belongs directly to this list or to a nested child list. + let levels_len = def_levels.len(); + let mut list_offsets: Vec = Vec::with_capacity(levels_len + 1); + let mut validity = self.nullable.then(|| BooleanBufferBuilder::new(levels_len)); + + // Count direct child values of this list. + // + // `d >= self.def_level` excludes null/empty parent lists: an empty + // list is encoded as `d == self.def_level - 1` (a null list lower + // still), so it carries no child value and must not advance the + // offset. Emptiness is handled by this definition-level guard, not + // by the repetition level. + // + // `r <= self.rep_level` (not `==`) keeps only direct children: the + // first element of each list has `r < self.rep_level` (it coincides + // with an outer/row boundary), so `==` would drop it; continuation + // entries of a nested child's own list have `r > self.rep_level` and + // are excluded. + // + // Split the scan by nullable/parent-threshold mode to avoid checking + // these options once per decoded level in the hot loop. + let mut cur_offset: usize = 0; + let def_level = self.def_level; + let rep_level = self.rep_level; + + match (self.parent_threshold, validity.as_mut()) { + (None, Some(validity)) => { + for (&d, &r) in def_levels.iter().zip(rep_levels) { + if r < rep_level { + list_offsets.push(OffsetSize::from_usize(cur_offset).unwrap()); + validity.append(d + 1 >= def_level); + } + + if d >= def_level && r <= rep_level { + cur_offset += 1; + } + } + } + (None, None) => { + for (&d, &r) in def_levels.iter().zip(rep_levels) { + if r < rep_level { + list_offsets.push(OffsetSize::from_usize(cur_offset).unwrap()); + } + + if d >= def_level && r <= rep_level { + cur_offset += 1; + } + } + } + (Some(threshold), Some(validity)) => { + for (&d, &r) in def_levels.iter().zip(rep_levels) { + // When this list is a child of another list, skip entries + // belonging to null/empty parent lists. + if d < threshold { + continue; + } + + if r < rep_level { + list_offsets.push(OffsetSize::from_usize(cur_offset).unwrap()); + validity.append(d + 1 >= def_level); + } + + if d >= def_level && r <= rep_level { + cur_offset += 1; + } + } + } + (Some(threshold), None) => { + for (&d, &r) in def_levels.iter().zip(rep_levels) { + // When this list is a child of another list, skip entries + // belonging to null/empty parent lists. + if d < threshold { + continue; + } + + if r < rep_level { + list_offsets.push(OffsetSize::from_usize(cur_offset).unwrap()); + } + + if d >= def_level && r <= rep_level { + cur_offset += 1; + } + } + } + } + + list_offsets.push(OffsetSize::from_usize(cur_offset).unwrap()); + + if cur_offset != child_array.len() { + return Err(general_err!( + "Failed to reconstruct list from level data: \ + expected {} child values but got {}", + cur_offset, + child_array.len() + )); + } + + debug_assert!(list_offsets.windows(2).all(|w| w[0] <= w[1])); + + let child_data = child_array.to_data(); + let value_offsets = Buffer::from(list_offsets.to_byte_slice()); + + let mut data_builder = ArrayData::builder(self.get_data_type().clone()) + .len(list_offsets.len() - 1) + .add_buffer(value_offsets) + .add_child_data(child_data); + + if let Some(builder) = validity { + if builder.len() != list_offsets.len() - 1 { + return Err(general_err!( + "Failed to decode level data for list array: \ + expected {} validity bits but got {}", + list_offsets.len() - 1, + builder.len() + )); + } + data_builder = data_builder.null_bit_buffer(Some(builder.into())); + } + + let list_data = unsafe { data_builder.build_unchecked() }; + Ok(Arc::new(GenericListArray::::from(list_data))) + } + + fn skip_records(&mut self, num_records: usize) -> Result { + self.item_reader.skip_records(num_records) + } + + fn get_def_levels(&self) -> Option<&[i16]> { + self.item_reader.get_def_levels() + } + + fn get_rep_levels(&self) -> Option<&[i16]> { + self.item_reader.get_rep_levels() + } + + fn max_def_level(&self) -> i16 { + self.def_level + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::array_reader::ArrayReaderBuilder; + use crate::arrow::array_reader::list_array::ListArrayReader; + use crate::arrow::array_reader::test_util::make_int32_page_reader; + use crate::arrow::arrow_reader::DEFAULT_BATCH_SIZE; + use crate::arrow::arrow_reader::metrics::ArrowReaderMetrics; + use crate::arrow::schema::parquet_to_arrow_schema_and_fields; + use crate::arrow::{ArrowWriter, ProjectionMask, parquet_to_arrow_schema}; + use crate::file::properties::WriterProperties; + use crate::file::reader::{FileReader, SerializedFileReader}; + use crate::schema::parser::parse_message_type; + use crate::schema::types::SchemaDescriptor; + use arrow::datatypes::{Field, Int32Type}; + use arrow_array::{Array, PrimitiveArray}; + use arrow_data::ArrayDataBuilder; + use arrow_schema::Fields; + use std::sync::Arc; + + fn list_type( + data_type: ArrowType, + item_nullable: bool, + ) -> ArrowType { + let field = Arc::new(Field::new_list_field(data_type, item_nullable)); + GenericListArray::::DATA_TYPE_CONSTRUCTOR(field) + } + + fn downcast(array: &ArrayRef) -> &'_ GenericListArray { + array + .as_any() + .downcast_ref::>() + .unwrap() + } + + fn to_offsets(values: Vec) -> Buffer { + Buffer::from_iter( + values + .into_iter() + .map(|x| OffsetSize::from_usize(x).unwrap()), + ) + } + + fn test_nested_list() { + // 3 lists, with first and third nullable + // [ + // [ + // [[1, null], null, [4], []], + // [], + // [[7]], + // [[]], + // [[1, 2, 3], [4, null, 6], null] + // ], + // null, + // [], + // [[[11]]] + // ] + + let l3_item_type = ArrowType::Int32; + let l3_type = list_type::(l3_item_type, true); + + let l2_item_type = l3_type.clone(); + let l2_type = list_type::(l2_item_type, true); + + let l1_item_type = l2_type.clone(); + let l1_type = list_type::(l1_item_type, false); + + let leaf = PrimitiveArray::::from_iter(vec![ + Some(1), + None, + Some(4), + Some(7), + Some(1), + Some(2), + Some(3), + Some(4), + None, + Some(6), + Some(11), + ]); + + // [[1, null], null, [4], [], [7], [], [1, 2, 3], [4, null, 6], null, [11]] + let offsets = to_offsets::(vec![0, 2, 2, 3, 3, 4, 4, 7, 10, 10, 11]); + let l3 = ArrayDataBuilder::new(l3_type.clone()) + .len(10) + .add_buffer(offsets) + .add_child_data(leaf.into_data()) + .null_bit_buffer(Some(Buffer::from([0b11111101, 0b00000010]))) + .build() + .unwrap(); + + // [[[1, null], null, [4], []], [], [[7]], [[]], [[1, 2, 3], [4, null, 6], null], [[11]]] + let offsets = to_offsets::(vec![0, 4, 4, 5, 6, 9, 10]); + let l2 = ArrayDataBuilder::new(l2_type.clone()) + .len(6) + .add_buffer(offsets) + .add_child_data(l3) + .build() + .unwrap(); + + let offsets = to_offsets::(vec![0, 5, 5, 5, 6]); + let l1 = ArrayDataBuilder::new(l1_type.clone()) + .len(4) + .add_buffer(offsets) + .add_child_data(l2) + .null_bit_buffer(Some(Buffer::from([0b00001101]))) + .build() + .unwrap(); + + let expected = GenericListArray::::from(l1); + + let item_array_reader = make_int32_page_reader( + &[1, 4, 7, 1, 2, 3, 4, 6, 11], + &[6, 5, 3, 6, 4, 2, 6, 4, 6, 6, 6, 6, 5, 6, 3, 0, 1, 6], + &[0, 3, 2, 2, 2, 1, 1, 1, 1, 3, 3, 2, 3, 3, 2, 0, 0, 0], + 6, + 3, + Some(5), + ); + + let l3 = + ListArrayReader::::new(item_array_reader, l3_type, 5, 3, true, Some(3)); + + let l2 = ListArrayReader::::new(Box::new(l3), l2_type, 3, 2, false, Some(2)); + + let mut l1 = ListArrayReader::::new(Box::new(l2), l1_type, 2, 1, true, None); + + let expected_1 = expected.slice(0, 2); + let expected_2 = expected.slice(2, 2); + + let actual = l1.next_batch(2).unwrap(); + assert_eq!(actual.as_ref(), &expected_1); + + let actual = l1.next_batch(1024).unwrap(); + assert_eq!(actual.as_ref(), &expected_2); + } + + fn test_required_list() { + // [[1, null, 2], [], [3, 4], [], [], [null, 1]] + let expected = + GenericListArray::::from_iter_primitive::(vec![ + Some(vec![Some(1), None, Some(2)]), + Some(vec![]), + Some(vec![Some(3), Some(4)]), + Some(vec![]), + Some(vec![]), + Some(vec![None, Some(1)]), + ]); + + let item_array_reader = make_int32_page_reader( + &[1, 2, 3, 4, 1], + &[2, 1, 2, 0, 2, 2, 0, 0, 1, 2], + &[0, 1, 1, 0, 0, 1, 0, 0, 0, 1], + 2, + 1, + Some(1), + ); + + let mut list_array_reader = ListArrayReader::::new( + item_array_reader, + list_type::(ArrowType::Int32, true), + 1, + 1, + false, + None, + ); + + let actual = list_array_reader.next_batch(1024).unwrap(); + let actual = downcast::(&actual); + + assert_eq!(&expected, actual) + } + + fn test_nullable_list() { + // [[1, null, 2], null, [], [3, 4], [], [], null, [], [null, 1]] + let expected = + GenericListArray::::from_iter_primitive::(vec![ + Some(vec![Some(1), None, Some(2)]), + None, + Some(vec![]), + Some(vec![Some(3), Some(4)]), + Some(vec![]), + Some(vec![]), + None, + Some(vec![]), + Some(vec![None, Some(1)]), + ]); + + let item_array_reader = make_int32_page_reader( + &[1, 2, 3, 4, 1], + &[3, 2, 3, 0, 1, 3, 3, 1, 1, 0, 1, 2, 3], + &[0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1], + 3, + 1, + Some(2), + ); + + let mut list_array_reader = ListArrayReader::::new( + item_array_reader, + list_type::(ArrowType::Int32, true), + 2, + 1, + true, + None, + ); + + let actual = list_array_reader.next_batch(1024).unwrap(); + let actual = downcast::(&actual); + + assert_eq!(&expected, actual) + } + + fn test_list_array() { + test_nullable_list::(); + test_required_list::(); + test_nested_list::(); + } + + #[test] + fn test_list_array_reader() { + test_list_array::(); + } + + #[test] + fn test_large_list_array_reader() { + test_list_array::() + } + + #[test] + fn test_list_multi_row_group_selective_padding() { + use crate::arrow::array_reader::PrimitiveArrayReader; + use crate::basic::{Encoding, Type as PhysicalType}; + use crate::data_type::Int32Type as ParquetInt32; + use crate::schema::types::{ColumnDescriptor, ColumnPath, Type}; + use crate::util::InMemoryPageIterator; + use crate::util::test_common::page_util::{DataPageBuilder, DataPageBuilderImpl}; + + // Schema: OPTIONAL GROUP my_list (LIST) { + // repeated group list { optional INT32 element; } + // } + // max_def_level = 3, max_rep_level = 1 + let leaf_type = Type::primitive_type_builder("element", PhysicalType::INT32) + .build() + .unwrap(); + + let desc = Arc::new(ColumnDescriptor::new( + Arc::new(leaf_type), + 3, // max_def_level + 1, // max_rep_level + ColumnPath::new(vec![]), + )); + + // Row group 0 — 3 records: + // Record 1: [10, null, 20] + // Record 2: null list + // Record 3: [null, 30] + let mut pb1 = DataPageBuilderImpl::new(desc.clone(), 6, true); + pb1.add_rep_levels(1, &[0, 1, 1, 0, 0, 1]); + pb1.add_def_levels(3, &[3, 2, 3, 0, 2, 3]); + pb1.add_values::(Encoding::PLAIN, &[10, 20, 30]); + + // Row group 1 — 2 records: + // Record 4: [40, null] + // Record 5: [null, null, 50] + let mut pb2 = DataPageBuilderImpl::new(desc.clone(), 5, true); + pb2.add_rep_levels(1, &[0, 1, 0, 1, 1]); + pb2.add_def_levels(3, &[3, 2, 2, 2, 3]); + pb2.add_values::(Encoding::PLAIN, &[40, 50]); + + // Two row groups, one page each → forces two read_one_batch calls + let pages = vec![vec![pb1.consume()], vec![pb2.consume()]]; + let page_iter = InMemoryPageIterator::new(pages); + + let item_reader = Box::new( + PrimitiveArrayReader::::new( + Box::new(page_iter), + desc, + None, + DEFAULT_BATCH_SIZE, + Some(2), // padding_threshold = list's def_level + ) + .unwrap(), + ); + + let list_type = ArrowType::List(Arc::new(Field::new_list_field(ArrowType::Int32, true))); + let mut list_reader = ListArrayReader::::new(item_reader, list_type, 2, 1, true, None); + + let result = list_reader.next_batch(5).unwrap(); + let list = result + .as_any() + .downcast_ref::>() + .unwrap(); + + assert_eq!(list.len(), 5); + + // Record 1: [10, null, 20] + assert!(list.is_valid(0)); + let r1 = list.value(0); + let r1 = r1 + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(r1.len(), 3); + assert_eq!(r1.value(0), 10); + assert!(r1.is_null(1)); + assert_eq!(r1.value(2), 20); + + // Record 2: null + assert!(list.is_null(1)); + + // Record 3: [null, 30] + assert!(list.is_valid(2)); + let r3 = list.value(2); + let r3 = r3 + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(r3.len(), 2); + assert!(r3.is_null(0)); + assert_eq!(r3.value(1), 30); + + // Record 4: [40, null] + assert!(list.is_valid(3)); + let r4 = list.value(3); + let r4 = r4 + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(r4.len(), 2); + assert_eq!(r4.value(0), 40); + assert!(r4.is_null(1)); + + // Record 5: [null, null, 50] + assert!(list.is_valid(4)); + let r5 = list.value(4); + let r5 = r5 + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(r5.len(), 3); + assert!(r5.is_null(0)); + assert!(r5.is_null(1)); + assert_eq!(r5.value(2), 50); + } + + #[test] + fn test_nested_lists() { + // Construct column schema + let message_type = " + message table { + REPEATED group table_info { + REQUIRED BYTE_ARRAY name; + REPEATED group cols { + REQUIRED BYTE_ARRAY name; + REQUIRED INT32 type; + OPTIONAL INT32 length; + } + REPEATED group tags { + REQUIRED BYTE_ARRAY name; + REQUIRED INT32 type; + OPTIONAL INT32 length; + } + } + } + "; + + let schema = parse_message_type(message_type) + .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t)))) + .unwrap(); + + let arrow_schema = parquet_to_arrow_schema(schema.as_ref(), None).unwrap(); + + let file = tempfile::tempfile().unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(200)) + .build(); + + let writer = ArrowWriter::try_new( + file.try_clone().unwrap(), + Arc::new(arrow_schema), + Some(props), + ) + .unwrap(); + writer.close().unwrap(); + + let file_reader: Arc = Arc::new(SerializedFileReader::new(file).unwrap()); + + let file_metadata = file_reader.metadata().file_metadata(); + let schema = file_metadata.schema_descr(); + let mask = ProjectionMask::leaves(schema, vec![0]); + let (_, fields) = parquet_to_arrow_schema_and_fields( + schema, + ProjectionMask::all(), + file_metadata.key_value_metadata(), + &[], + ) + .unwrap(); + + let metrics = ArrowReaderMetrics::disabled(); + let mut array_reader = ArrayReaderBuilder::new(&file_reader, &metrics) + .with_batch_size(DEFAULT_BATCH_SIZE) + .build_array_reader(fields.as_ref(), &mask) + .unwrap(); + + let batch = array_reader.next_batch(100).unwrap(); + assert_eq!(batch.data_type(), array_reader.get_data_type()); + assert_eq!( + batch.data_type(), + &ArrowType::Struct(Fields::from(vec![Field::new( + "table_info", + ArrowType::List(Arc::new(Field::new( + "table_info", + ArrowType::Struct(vec![Field::new("name", ArrowType::Binary, false)].into()), + false + ))), + false + )])) + ); + assert_eq!(batch.len(), 0); + } +} diff --git a/vendor/parquet/src/arrow/array_reader/list_view_array.rs b/vendor/parquet/src/arrow/array_reader/list_view_array.rs new file mode 100644 index 0000000..0a7f58c --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/list_view_array.rs @@ -0,0 +1,220 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::array_reader::{ArrayReader, ListArrayReader}; +use crate::errors::Result; +use arrow_array::cast::AsArray; +use arrow_array::{Array, ArrayRef, GenericListViewArray, OffsetSizeTrait, new_empty_array}; +use arrow_schema::DataType as ArrowType; +use std::any::Any; +use std::sync::Arc; + +/// Implementation of list view array reader. +/// This wraps a ListArrayReader and converts the result to ListViewArray. +pub struct ListViewArrayReader { + inner: ListArrayReader, + data_type: ArrowType, +} + +impl ListViewArrayReader { + /// Construct list view array reader. + pub fn new( + item_reader: Box, + data_type: ArrowType, + def_level: i16, + rep_level: i16, + nullable: bool, + parent_threshold: Option, + ) -> Self { + // Create the underlying ListArrayReader with the corresponding List type + let list_data_type = match &data_type { + ArrowType::ListView(f) => ArrowType::List(f.clone()), + ArrowType::LargeListView(f) => ArrowType::LargeList(f.clone()), + _ => unreachable!(), + }; + + let inner = ListArrayReader::new( + item_reader, + list_data_type, + def_level, + rep_level, + nullable, + parent_threshold, + ); + + Self { inner, data_type } + } +} + +impl ArrayReader for ListViewArrayReader { + fn as_any(&self) -> &dyn Any { + self + } + + /// Returns data type. + /// This must be a ListView. + fn get_data_type(&self) -> &ArrowType { + &self.data_type + } + + fn read_records(&mut self, batch_size: usize) -> Result { + self.inner.read_records(batch_size) + } + + fn consume_batch(&mut self) -> Result { + let array = self.inner.consume_batch()?; + if array.is_empty() { + return Ok(new_empty_array(&self.data_type)); + } + + // Convert ListArray to ListViewArray + let list_array = array.as_list::(); + + let list_view_array = + Arc::new(GenericListViewArray::::from(list_array.clone())); + + // Ensure the data type is correct + assert_eq!( + list_view_array.data_type(), + &self.data_type, + "Converted array type does not match expected type" + ); + + Ok(list_view_array) + } + + fn skip_records(&mut self, num_records: usize) -> Result { + self.inner.skip_records(num_records) + } + + fn get_def_levels(&self) -> Option<&[i16]> { + self.inner.get_def_levels() + } + + fn get_rep_levels(&self) -> Option<&[i16]> { + self.inner.get_rep_levels() + } + + fn max_def_level(&self) -> i16 { + self.inner.max_def_level() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::array_reader::test_util::make_int32_page_reader; + use arrow::datatypes::Int32Type as ArrowInt32; + + fn test_nullable_list_view() { + // [[1, null, 2], null, [], [3, 4], [], [], null, [], [null, 1]] + let expected = + GenericListViewArray::::from_iter_primitive::(vec![ + Some(vec![Some(1), None, Some(2)]), + None, + Some(vec![]), + Some(vec![Some(3), Some(4)]), + Some(vec![]), + Some(vec![]), + None, + Some(vec![]), + Some(vec![None, Some(1)]), + ]); + + let item_array_reader = make_int32_page_reader( + &[1, 2, 3, 4, 1], + &[3, 2, 3, 0, 1, 3, 3, 1, 1, 0, 1, 2, 3], + &[0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1], + 3, + 1, + Some(2), + ); + + let field = Arc::new(arrow_schema::Field::new_list_field(ArrowType::Int32, true)); + let data_type = if OffsetSize::IS_LARGE { + ArrowType::LargeListView(field) + } else { + ArrowType::ListView(field) + }; + + let mut list_view_array_reader = + ListViewArrayReader::::new(item_array_reader, data_type, 2, 1, true, None); + + let actual = list_view_array_reader.next_batch(1024).unwrap(); + let actual = actual + .as_any() + .downcast_ref::>() + .unwrap(); + + assert_eq!(&expected, actual) + } + + fn test_required_list_view() { + // [[1, null, 2], [], [3, 4], [], [], [null, 1]] + let expected = + GenericListViewArray::::from_iter_primitive::(vec![ + Some(vec![Some(1), None, Some(2)]), + Some(vec![]), + Some(vec![Some(3), Some(4)]), + Some(vec![]), + Some(vec![]), + Some(vec![None, Some(1)]), + ]); + + let item_array_reader = make_int32_page_reader( + &[1, 2, 3, 4, 1], + &[2, 1, 2, 0, 2, 2, 0, 0, 1, 2], + &[0, 1, 1, 0, 0, 1, 0, 0, 0, 1], + 2, + 1, + Some(1), + ); + + let field = Arc::new(arrow_schema::Field::new_list_field(ArrowType::Int32, true)); + let data_type = if OffsetSize::IS_LARGE { + ArrowType::LargeListView(field) + } else { + ArrowType::ListView(field) + }; + + let mut list_view_array_reader = + ListViewArrayReader::::new(item_array_reader, data_type, 1, 1, false, None); + + let actual = list_view_array_reader.next_batch(1024).unwrap(); + let actual = actual + .as_any() + .downcast_ref::>() + .unwrap(); + + assert_eq!(&expected, actual) + } + + fn test_list_view_array() { + test_nullable_list_view::(); + test_required_list_view::(); + } + + #[test] + fn test_list_view_array_reader() { + test_list_view_array::(); + } + + #[test] + fn test_large_list_view_array_reader() { + test_list_view_array::() + } +} diff --git a/vendor/parquet/src/arrow/array_reader/map_array.rs b/vendor/parquet/src/arrow/array_reader/map_array.rs new file mode 100644 index 0000000..5d88df6 --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/map_array.rs @@ -0,0 +1,223 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::array_reader::{ArrayReader, ListArrayReader, StructArrayReader}; +use crate::errors::Result; +use arrow_array::{Array, ArrayRef, MapArray}; +use arrow_schema::DataType as ArrowType; +use std::any::Any; +use std::sync::Arc; + +/// Implementation of a map array reader. +pub struct MapArrayReader { + data_type: ArrowType, + reader: ListArrayReader, +} + +impl MapArrayReader { + #[allow(rustdoc::private_intra_doc_links)] + /// Creates a new [`MapArrayReader`] with a `def_level`, `rep_level` and `nullable` + /// as defined on [`ParquetField`][crate::arrow::schema::ParquetField] + pub fn new( + key_reader: Box, + value_reader: Box, + data_type: ArrowType, + def_level: i16, + rep_level: i16, + nullable: bool, + parent_threshold: Option, + ) -> Self { + // The struct exists when the key (always required in maps) exists. + // Derive struct_def_level from the key's max def level rather than + // a fixed formula, so it matches the schema exactly. + let struct_def_level = key_reader.max_def_level(); + let struct_rep_level = rep_level + 1; + + let element = match &data_type { + ArrowType::Map(element, _) => match element.data_type() { + ArrowType::Struct(fields) if fields.len() == 2 => { + // Parquet cannot represent nullability at this level (#1697) + // and so encountering nullability here indicates some manner + // of schema inconsistency / inference bug + assert!(!element.is_nullable(), "map struct cannot be nullable"); + element + } + _ => unreachable!("expected struct with two fields"), + }, + _ => unreachable!("expected map type"), + }; + + let struct_reader = StructArrayReader::new( + element.data_type().clone(), + vec![key_reader, value_reader], + struct_def_level, + struct_rep_level, + false, + Some(def_level), + ); + + let reader = ListArrayReader::new( + Box::new(struct_reader), + ArrowType::List(element.clone()), + def_level, + rep_level, + nullable, + parent_threshold, + ); + + Self { data_type, reader } + } +} + +impl ArrayReader for MapArrayReader { + fn as_any(&self) -> &dyn Any { + self + } + + fn get_data_type(&self) -> &ArrowType { + &self.data_type + } + + fn read_records(&mut self, batch_size: usize) -> Result { + self.reader.read_records(batch_size) + } + + fn consume_batch(&mut self) -> Result { + // A MapArray is just a ListArray with a StructArray child + // we can therefore just alter the ArrayData + let array = self.reader.consume_batch().unwrap(); + let data = array.to_data(); + let builder = data.into_builder().data_type(self.data_type.clone()); + + // SAFETY - we can assume that ListArrayReader produces valid ListArray + // of the expected type, and as such its output can be reinterpreted as + // a MapArray without validation + Ok(Arc::new(MapArray::from(unsafe { + builder.build_unchecked() + }))) + } + + fn skip_records(&mut self, num_records: usize) -> Result { + self.reader.skip_records(num_records) + } + + fn get_def_levels(&self) -> Option<&[i16]> { + self.reader.get_def_levels() + } + + fn get_rep_levels(&self) -> Option<&[i16]> { + self.reader.get_rep_levels() + } + + fn max_def_level(&self) -> i16 { + self.reader.max_def_level() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::ArrowWriter; + use crate::arrow::arrow_reader::ParquetRecordBatchReader; + use arrow::datatypes::{Field, Int32Type, Schema}; + use arrow_array::RecordBatch; + use arrow_array::builder::{MapBuilder, PrimitiveBuilder, StringBuilder}; + use arrow_array::cast::*; + use arrow_schema::Fields; + use bytes::Bytes; + + #[test] + // This test writes a parquet file with the following data: + // +--------------------------------------------------------+ + // |map | + // +--------------------------------------------------------+ + // |null | + // |null | + // |{three -> 3, four -> 4, five -> 5, six -> 6, seven -> 7}| + // +--------------------------------------------------------+ + // + // It then attempts to read the data back and checks that the third record + // contains the expected values. + fn read_map_array_column() { + // Schema for single map of string to int32 + let schema = Schema::new(vec![Field::new( + "map", + ArrowType::Map( + Arc::new(Field::new( + "entries", + ArrowType::Struct(Fields::from(vec![ + Field::new("keys", ArrowType::Utf8, false), + Field::new("values", ArrowType::Int32, true), + ])), + false, + )), + false, // Map field not sorted + ), + true, + )]); + + // Create builders for map + let string_builder = StringBuilder::new(); + let ints_builder: PrimitiveBuilder = PrimitiveBuilder::new(); + let mut map_builder = MapBuilder::new(None, string_builder, ints_builder); + + // Add two null records and one record with five entries + map_builder.append(false).expect("adding null map entry"); + map_builder.append(false).expect("adding null map entry"); + map_builder.keys().append_value("three"); + map_builder.keys().append_value("four"); + map_builder.keys().append_value("five"); + map_builder.keys().append_value("six"); + map_builder.keys().append_value("seven"); + + map_builder.values().append_value(3); + map_builder.values().append_value(4); + map_builder.values().append_value(5); + map_builder.values().append_value(6); + map_builder.values().append_value(7); + map_builder.append(true).expect("adding map entry"); + + // Create record batch + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(map_builder.finish())]) + .expect("create record batch"); + + // Write record batch to file + let mut buffer = Vec::with_capacity(1024); + let mut writer = + ArrowWriter::try_new(&mut buffer, batch.schema(), None).expect("creat file writer"); + writer.write(&batch).expect("writing file"); + writer.close().expect("close writer"); + + // Read file + let reader = Bytes::from(buffer); + let record_batch_reader = ParquetRecordBatchReader::try_new(reader, 1024).unwrap(); + for maybe_record_batch in record_batch_reader { + let record_batch = maybe_record_batch.expect("Getting current batch"); + let col = record_batch.column(0); + assert!(col.is_null(0)); + assert!(col.is_null(1)); + let map_entry = as_map_array(col).value(2); + let struct_col = as_struct_array(&map_entry); + let key_col = as_string_array(struct_col.column(0)); // Key column + assert_eq!(key_col.value(0), "three"); + assert_eq!(key_col.value(1), "four"); + assert_eq!(key_col.value(2), "five"); + assert_eq!(key_col.value(3), "six"); + assert_eq!(key_col.value(4), "seven"); + } + } +} diff --git a/vendor/parquet/src/arrow/array_reader/mod.rs b/vendor/parquet/src/arrow/array_reader/mod.rs new file mode 100644 index 0000000..3470583 --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/mod.rs @@ -0,0 +1,258 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Logic for reading into arrow arrays: [`ArrayReader`] and [`RowGroups`] + +use crate::errors::Result; +use arrow_array::ArrayRef; +use arrow_schema::DataType as ArrowType; +use std::any::Any; +use std::sync::Arc; + +use crate::arrow::record_reader::GenericRecordReader; +use crate::arrow::record_reader::buffer::ValuesBuffer; +use crate::column::page::PageIterator; +use crate::column::reader::decoder::ColumnValueDecoder; +use crate::file::metadata::ParquetMetaData; +use crate::file::reader::{FilePageIterator, FileReader}; + +mod builder; +mod byte_array; +mod byte_array_dictionary; +mod byte_view_array; +mod cached_array_reader; +mod empty_array; +mod fixed_len_byte_array; +mod fixed_size_list_array; +mod list_array; +mod list_view_array; +mod map_array; +mod null_array; +mod primitive_array; +mod row_group_cache; +mod row_group_index; +mod row_number; +mod struct_array; + +#[cfg(test)] +pub(crate) mod test_util; + +// Note that this crate is public under the `experimental` feature flag. +use crate::file::metadata::RowGroupMetaData; +pub use builder::{ArrayReaderBuilder, CacheOptions, CacheOptionsBuilder}; +pub use byte_array::make_byte_array_reader; +pub use byte_array_dictionary::make_byte_array_dictionary_reader; +#[allow(unused_imports)] // Only used for benchmarks +pub use byte_view_array::make_byte_view_array_reader; +#[allow(unused_imports)] // Only used for benchmarks +pub use fixed_len_byte_array::make_fixed_len_byte_array_reader; +pub use fixed_size_list_array::FixedSizeListArrayReader; +pub use list_array::ListArrayReader; +pub use list_view_array::ListViewArrayReader; +pub use map_array::MapArrayReader; +pub use null_array::NullArrayReader; +pub use primitive_array::PrimitiveArrayReader; +pub use row_group_cache::RowGroupCache; +pub use struct_array::StructArrayReader; + +/// Reads Parquet data into Arrow Arrays. +/// +/// This is an internal implementation detail of the Parquet reader, and is not +/// intended for public use. +/// +/// This is the core trait for reading encoded Parquet data directly into Arrow +/// Arrays efficiently. There are various specializations of this trait for +/// different combinations of encodings and arrays, such as +/// [`PrimitiveArrayReader`], [`ListArrayReader`], etc. +/// +/// Each `ArrayReader` logically contains the following state +/// 1. A handle to the encoded Parquet data +/// 2. An in progress buffered Array +/// +/// Data can either be read in batches using [`ArrayReader::next_batch`] or +/// incrementally using [`ArrayReader::read_records`] and [`ArrayReader::skip_records`]. +pub trait ArrayReader: Send { + // TODO: this function is never used, and the trait is not public. Perhaps this should be + // removed. + #[allow(dead_code)] + fn as_any(&self) -> &dyn Any; + + /// Returns the arrow type of this array reader. + fn get_data_type(&self) -> &ArrowType; + + /// Reads at most `batch_size` records into an arrow array and return it. + #[cfg(any(feature = "experimental", test))] + fn next_batch(&mut self, batch_size: usize) -> Result { + self.read_records(batch_size)?; + self.consume_batch() + } + + /// Reads at most `batch_size` records' bytes into buffer + /// + /// Returns the number of records read, which can be less than `batch_size` if + /// pages is exhausted. + fn read_records(&mut self, batch_size: usize) -> Result; + + /// Consume all currently stored buffer data + /// into an arrow array and return it. + fn consume_batch(&mut self) -> Result; + + /// Skips over `num_records` records, returning the number of rows skipped + /// + /// Note that calling `skip_records` with large values of `num_records` is + /// efficient as it avoids decoding data into the the in-progress array. + /// However, there is overhead to calling this function, so for small values of + /// `num_records`, it can be more efficient to call read_records and apply + /// a filter to the resulting array. + fn skip_records(&mut self, num_records: usize) -> Result; + + /// If this array has a non-zero definition level, i.e. has a nullable parent + /// array, returns the definition levels of data from the last call of `next_batch` + /// + /// Otherwise returns None + /// + /// This is used by parent [`ArrayReader`] to compute their null bitmaps + fn get_def_levels(&self) -> Option<&[i16]>; + + /// If this array has a non-zero repetition level, i.e. has a repeated parent + /// array, returns the repetition levels of data from the last call of `next_batch` + /// + /// Otherwise returns None + /// + /// This is used by parent [`ArrayReader`] to compute their array offsets + fn get_rep_levels(&self) -> Option<&[i16]>; + + /// Returns the maximum definition level for this reader as defined by + /// the Parquet schema. For leaf readers this is the column's max def level; + /// for composite readers it is the level at which the composite itself is + /// fully defined. + /// + /// The default panics. Synthetic readers that are not backed by a Parquet + /// column (e.g. row-number or empty-struct readers) may rely on this + /// default because they are never used as children of schema-driven + /// composite readers (list, map, struct), which are the only callers. + fn max_def_level(&self) -> i16 { + panic!("max_def_level called on a reader that does not track definition levels") + } +} + +/// Interface for reading data pages from the columns of one or more RowGroups. +pub trait RowGroups { + /// Get the number of rows in this collection + fn num_rows(&self) -> usize; + + /// Returns a [`PageIterator`] for all pages in the specified column chunk + /// across all row groups in this collection. + fn column_chunks(&self, i: usize) -> Result>; + + /// Returns an iterator over the row groups in this collection + /// + /// Note this may not include all row groups in [`Self::metadata`]. + fn row_groups(&self) -> Box + '_>; + + /// Returns the parquet metadata + fn metadata(&self) -> &ParquetMetaData; +} + +impl RowGroups for Arc { + fn num_rows(&self) -> usize { + FileReader::metadata(self.as_ref()) + .file_metadata() + .num_rows() as usize + } + + fn column_chunks(&self, column_index: usize) -> Result> { + let iterator = FilePageIterator::new(column_index, Arc::clone(self))?; + Ok(Box::new(iterator)) + } + + fn row_groups(&self) -> Box + '_> { + Box::new(FileReader::metadata(self.as_ref()).row_groups().iter()) + } + + fn metadata(&self) -> &ParquetMetaData { + FileReader::metadata(self.as_ref()) + } +} + +/// Uses `record_reader` to read up to `batch_size` records from `pages` +/// +/// Returns the number of records read, which can be less than `batch_size` if +/// pages is exhausted. +fn read_records( + record_reader: &mut GenericRecordReader, + pages: &mut dyn PageIterator, + batch_size: usize, +) -> Result +where + V: ValuesBuffer, + CV: ColumnValueDecoder, +{ + let mut records_read = 0usize; + while records_read < batch_size { + let records_to_read = batch_size - records_read; + + let records_read_once = record_reader.read_records(records_to_read)?; + records_read += records_read_once; + + // Record reader exhausted + if records_read_once < records_to_read { + if let Some(page_reader) = pages.next() { + // Read from new page reader (i.e. column chunk) + record_reader.set_page_reader(page_reader?)?; + } else { + // Page reader also exhausted + break; + } + } + } + Ok(records_read) +} + +/// Uses `record_reader` to skip up to `batch_size` records from `pages` +/// +/// Returns the number of records skipped, which can be less than `batch_size` if +/// pages is exhausted +fn skip_records( + record_reader: &mut GenericRecordReader, + pages: &mut dyn PageIterator, + batch_size: usize, +) -> Result +where + V: ValuesBuffer, + CV: ColumnValueDecoder, +{ + let mut records_skipped = 0usize; + while records_skipped < batch_size { + let records_to_read = batch_size - records_skipped; + + let records_skipped_once = record_reader.skip_records(records_to_read)?; + records_skipped += records_skipped_once; + + // Record reader exhausted + if records_skipped_once < records_to_read { + if let Some(page_reader) = pages.next() { + // Read from new page reader (i.e. column chunk) + record_reader.set_page_reader(page_reader?)?; + } else { + // Page reader also exhausted + break; + } + } + } + Ok(records_skipped) +} diff --git a/vendor/parquet/src/arrow/array_reader/null_array.rs b/vendor/parquet/src/arrow/array_reader/null_array.rs new file mode 100644 index 0000000..d3740b5 --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/null_array.rs @@ -0,0 +1,120 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::array_reader::{ArrayReader, read_records, skip_records}; +use crate::arrow::record_reader::RecordReader; +use crate::column::page::PageIterator; +use crate::data_type::DataType; +use crate::errors::Result; +use crate::schema::types::ColumnDescPtr; +use arrow_array::ArrayRef; +use arrow_buffer::ArrowNativeType; +use arrow_schema::DataType as ArrowType; +use std::any::Any; +use std::sync::Arc; + +/// A NullArrayReader reads Parquet columns stored as null int32s with an Arrow +/// NullArray type. +pub struct NullArrayReader +where + T: DataType, + T::T: ArrowNativeType, +{ + data_type: ArrowType, + pages: Box, + def_levels_buffer: Option>, + rep_levels_buffer: Option>, + record_reader: RecordReader, +} + +impl NullArrayReader +where + T: DataType, + T::T: ArrowNativeType, +{ + /// Construct null array reader. + pub fn new( + pages: Box, + column_desc: ColumnDescPtr, + batch_size: usize, + padding_threshold: Option, + ) -> Result { + let mut record_reader = RecordReader::::new(column_desc, batch_size); + if let Some(threshold) = padding_threshold { + record_reader.set_padding_threshold(threshold); + } + + Ok(Self { + data_type: ArrowType::Null, + pages, + def_levels_buffer: None, + rep_levels_buffer: None, + record_reader, + }) + } +} + +/// Implementation of primitive array reader. +impl ArrayReader for NullArrayReader +where + T: DataType, + T::T: ArrowNativeType, +{ + fn as_any(&self) -> &dyn Any { + self + } + + /// Returns data type of primitive array. + fn get_data_type(&self) -> &ArrowType { + &self.data_type + } + + fn read_records(&mut self, batch_size: usize) -> Result { + read_records(&mut self.record_reader, self.pages.as_mut(), batch_size) + } + + fn consume_batch(&mut self) -> Result { + // convert to arrays + let array = arrow_array::NullArray::new(self.record_reader.values_written()); + + // save definition and repetition buffers + self.def_levels_buffer = self.record_reader.consume_def_levels(); + self.rep_levels_buffer = self.record_reader.consume_rep_levels(); + + // Must consume bitmap buffer + self.record_reader.consume_compact_bitmap(); + + self.record_reader.reset(); + Ok(Arc::new(array)) + } + + fn skip_records(&mut self, num_records: usize) -> Result { + skip_records(&mut self.record_reader, self.pages.as_mut(), num_records) + } + + fn get_def_levels(&self) -> Option<&[i16]> { + self.def_levels_buffer.as_deref() + } + + fn get_rep_levels(&self) -> Option<&[i16]> { + self.rep_levels_buffer.as_deref() + } + + fn max_def_level(&self) -> i16 { + self.record_reader.max_def_level() + } +} diff --git a/vendor/parquet/src/arrow/array_reader/primitive_array.rs b/vendor/parquet/src/arrow/array_reader/primitive_array.rs new file mode 100644 index 0000000..a6581ae --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/primitive_array.rs @@ -0,0 +1,1020 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::array_reader::{ArrayReader, read_records, skip_records}; +use crate::arrow::record_reader::RecordReader; +use crate::arrow::schema::parquet_to_arrow_field; +use crate::basic::Type as PhysicalType; +use crate::column::page::PageIterator; +use crate::data_type::{DataType, Int96}; +use crate::errors::Result; +use crate::schema::types::ColumnDescPtr; +use arrow_array::{ + Array, ArrayRef, BooleanArray, Date64Array, Decimal64Array, Decimal128Array, Decimal256Array, + Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, Int64Array, PrimitiveArray, + UInt8Array, UInt16Array, builder::PrimitiveDictionaryBuilder, cast::AsArray, downcast_integer, + types::*, +}; +use arrow_array::{ + TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, + TimestampSecondArray, UInt32Array, UInt64Array, +}; +use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer, ScalarBuffer, i256}; +use arrow_schema::{DataType as ArrowType, TimeUnit}; +use std::any::Any; +use std::sync::Arc; + +/// Provides conversion from `Vec` to `Buffer` +pub trait IntoBuffer { + fn into_buffer(self, target_type: &ArrowType) -> Buffer; +} + +macro_rules! native_buffer { + ($($t:ty),*) => { + $(impl IntoBuffer for Vec<$t> { + fn into_buffer(self, _target_type: &ArrowType) -> Buffer { + Buffer::from_vec(self) + } + })* + }; +} +native_buffer!(i8, i16, i32, i64, u8, u16, u32, u64, f32, f64); + +impl IntoBuffer for Vec { + fn into_buffer(self, _target_type: &ArrowType) -> Buffer { + BooleanBuffer::from(self.as_slice()).into_inner() + } +} + +impl IntoBuffer for Vec { + fn into_buffer(self, target_type: &ArrowType) -> Buffer { + let mut builder = Vec::with_capacity(self.len()); + match target_type { + ArrowType::Timestamp(TimeUnit::Second, _) => { + builder.extend(self.iter().map(|x| x.to_seconds())); + } + ArrowType::Timestamp(TimeUnit::Millisecond, _) => { + builder.extend(self.iter().map(|x| x.to_millis())); + } + ArrowType::Timestamp(TimeUnit::Microsecond, _) => { + builder.extend(self.iter().map(|x| x.to_micros())); + } + ArrowType::Timestamp(TimeUnit::Nanosecond, _) => { + builder.extend(self.iter().map(|x| x.to_nanos())); + } + _ => unreachable!("Invalid target_type for Int96."), + } + Buffer::from_vec(builder) + } +} + +/// Primitive array readers are leaves of array reader tree. They accept page iterator +/// and read them into primitive arrays. +pub struct PrimitiveArrayReader +where + T: DataType, + T::T: Copy + Default, + Vec: IntoBuffer, +{ + data_type: ArrowType, + pages: Box, + def_levels_buffer: Option>, + rep_levels_buffer: Option>, + record_reader: RecordReader, +} + +impl PrimitiveArrayReader +where + T: DataType, + T::T: Copy + Default, + Vec: IntoBuffer, +{ + /// Construct primitive array reader. + /// + /// `padding_threshold` controls how null padding is applied. When + /// `None`, the reader pads all null positions (full padding). When + /// `Some(threshold)`, entries with `def < threshold` are excluded + /// from the value buffer (selective padding for list children). + pub fn new( + pages: Box, + column_desc: ColumnDescPtr, + arrow_type: Option, + batch_size: usize, + padding_threshold: Option, + ) -> Result { + // Check if Arrow type is specified, else create it from Parquet type + let data_type = match arrow_type { + Some(t) => t, + None => parquet_to_arrow_field(column_desc.as_ref())? + .data_type() + .clone(), + }; + + let mut record_reader = RecordReader::::new(column_desc, batch_size); + if let Some(threshold) = padding_threshold { + record_reader.set_padding_threshold(threshold); + } + + Ok(Self { + data_type, + pages, + def_levels_buffer: None, + rep_levels_buffer: None, + record_reader, + }) + } +} + +/// Implementation of primitive array reader. +impl ArrayReader for PrimitiveArrayReader +where + T: DataType, + T::T: Copy + Default, + Vec: IntoBuffer, +{ + fn as_any(&self) -> &dyn Any { + self + } + + /// Returns data type of primitive array. + fn get_data_type(&self) -> &ArrowType { + &self.data_type + } + + fn read_records(&mut self, batch_size: usize) -> Result { + read_records(&mut self.record_reader, self.pages.as_mut(), batch_size) + } + + fn consume_batch(&mut self) -> Result { + let target_type = &self.data_type; + + // Convert physical data to equivalent arrow type, and then perform + // coercion as needed + let len = self.record_reader.values_written(); + + let record_data = self + .record_reader + .consume_record_data() + .into_buffer(target_type); + + let nulls = self + .record_reader + .consume_compact_bitmap() + .and_then(|b| NullBuffer::from_unsliced_buffer(b, len)); + + let array: ArrayRef = match T::get_physical_type() { + PhysicalType::BOOLEAN => Arc::new(BooleanArray::new( + BooleanBuffer::new(record_data, 0, len), + nulls, + )), + PhysicalType::INT32 => Arc::new(Int32Array::new( + ScalarBuffer::new(record_data, 0, len), + nulls, + )), + PhysicalType::INT64 => Arc::new(Int64Array::new( + ScalarBuffer::new(record_data, 0, len), + nulls, + )), + PhysicalType::FLOAT => Arc::new(Float32Array::new( + ScalarBuffer::new(record_data, 0, len), + nulls, + )), + PhysicalType::DOUBLE => Arc::new(Float64Array::new( + ScalarBuffer::new(record_data, 0, len), + nulls, + )), + PhysicalType::INT96 => Arc::new(Int64Array::new( + ScalarBuffer::new(record_data, 0, len), + nulls, + )), + PhysicalType::BYTE_ARRAY | PhysicalType::FIXED_LEN_BYTE_ARRAY => { + unreachable!("PrimitiveArrayReaders don't support complex physical types"); + } + }; + + // Coerce the arrow type to the desired array type + let array = coerce_array(array, target_type)?; + + // save definition and repetition buffers + self.def_levels_buffer = self.record_reader.consume_def_levels(); + self.rep_levels_buffer = self.record_reader.consume_rep_levels(); + self.record_reader.reset(); + Ok(array) + } + + fn skip_records(&mut self, num_records: usize) -> Result { + skip_records(&mut self.record_reader, self.pages.as_mut(), num_records) + } + + fn get_def_levels(&self) -> Option<&[i16]> { + self.def_levels_buffer.as_deref() + } + + fn get_rep_levels(&self) -> Option<&[i16]> { + self.rep_levels_buffer.as_deref() + } + + fn max_def_level(&self) -> i16 { + self.record_reader.max_def_level() + } +} + +/// Coerce the parquet physical type array to the target type +/// +/// This should match the logic in schema::primitive::apply_hint +fn coerce_array(array: ArrayRef, target_type: &ArrowType) -> Result { + if let ArrowType::Dictionary(key_type, value_type) = target_type { + let dictionary = pack_dictionary(key_type, array.as_ref())?; + let any_dictionary = dictionary.as_any_dictionary(); + + let coerced_values = + coerce_array(Arc::clone(any_dictionary.values()), value_type.as_ref())?; + + return Ok(any_dictionary.with_values(coerced_values)); + } + + match array.data_type() { + ArrowType::Int32 => coerce_i32(array.as_primitive(), target_type), + ArrowType::Int64 => coerce_i64(array.as_primitive(), target_type), + ArrowType::Boolean | ArrowType::Float32 | ArrowType::Float64 => Ok(array), + _ => unreachable!("Cannot coerce array of type {}", array.data_type()), + } +} + +fn coerce_i32(array: &Int32Array, target_type: &ArrowType) -> Result { + Ok(match target_type { + ArrowType::UInt8 => { + let array = array.unary(|i| i as u8) as UInt8Array; + Arc::new(array) as ArrayRef + } + ArrowType::Int8 => { + let array = array.unary(|i| i as i8) as Int8Array; + Arc::new(array) as ArrayRef + } + ArrowType::UInt16 => { + let array = array.unary(|i| i as u16) as UInt16Array; + Arc::new(array) as ArrayRef + } + ArrowType::Int16 => { + let array = array.unary(|i| i as i16) as Int16Array; + Arc::new(array) as ArrayRef + } + ArrowType::Int32 => Arc::new(array.clone()), + // follow C++ implementation and use overflow/reinterpret cast from i32 to u32 which will map + // `i32::MIN..0` to `(i32::MAX as u32)..u32::MAX` + ArrowType::UInt32 => Arc::new(UInt32Array::new( + array.values().inner().clone().into(), + array.nulls().cloned(), + )) as ArrayRef, + ArrowType::Date32 => Arc::new(array.reinterpret_cast::()) as _, + ArrowType::Date64 => { + let array: Date64Array = array.unary(|x| x as i64 * 86_400_000); + Arc::new(array) as ArrayRef + } + ArrowType::Time32(TimeUnit::Second) => { + Arc::new(array.reinterpret_cast::()) as ArrayRef + } + ArrowType::Time32(TimeUnit::Millisecond) => { + Arc::new(array.reinterpret_cast::()) as ArrayRef + } + ArrowType::Timestamp(time_unit, timezone) => match time_unit { + TimeUnit::Second => { + let array: TimestampSecondArray = array + .unary(|x| x as i64) + .with_timezone_opt(timezone.clone()); + Arc::new(array) as _ + } + TimeUnit::Millisecond => { + let array: TimestampMillisecondArray = array + .unary(|x| x as i64) + .with_timezone_opt(timezone.clone()); + Arc::new(array) as _ + } + TimeUnit::Microsecond => { + let array: TimestampMicrosecondArray = array + .unary(|x| x as i64) + .with_timezone_opt(timezone.clone()); + Arc::new(array) as _ + } + TimeUnit::Nanosecond => { + let array: TimestampNanosecondArray = array + .unary(|x| x as i64) + .with_timezone_opt(timezone.clone()); + Arc::new(array) as _ + } + }, + ArrowType::Decimal32(p, s) => { + let array = array + .reinterpret_cast::() + .with_precision_and_scale(*p, *s)?; + Arc::new(array) as ArrayRef + } + ArrowType::Decimal64(p, s) => { + let array: Decimal64Array = + array.unary(|i| i as i64).with_precision_and_scale(*p, *s)?; + Arc::new(array) as ArrayRef + } + ArrowType::Decimal128(p, s) => { + let array: Decimal128Array = array + .unary(|i| i as i128) + .with_precision_and_scale(*p, *s)?; + Arc::new(array) as ArrayRef + } + ArrowType::Decimal256(p, s) => { + let array: Decimal256Array = array + .unary(|i| i256::from_i128(i as i128)) + .with_precision_and_scale(*p, *s)?; + Arc::new(array) as ArrayRef + } + _ => unreachable!("Cannot coerce i32 to {target_type}"), + }) +} + +fn coerce_i64(array: &Int64Array, target_type: &ArrowType) -> Result { + Ok(match target_type { + ArrowType::Int64 => Arc::new(array.clone()) as _, + // follow C++ implementation and use overflow/reinterpret cast from i64 to u64 which will map + // `i64::MIN..0` to `(i64::MAX as u64)..u64::MAX` + ArrowType::UInt64 => Arc::new(UInt64Array::new( + array.values().inner().clone().into(), + array.nulls().cloned(), + )) as ArrayRef, + ArrowType::Date64 => Arc::new(array.reinterpret_cast::()) as _, + ArrowType::Time64(TimeUnit::Microsecond) => { + Arc::new(array.reinterpret_cast::()) as _ + } + ArrowType::Time64(TimeUnit::Nanosecond) => { + Arc::new(array.reinterpret_cast::()) as _ + } + ArrowType::Duration(unit) => match unit { + TimeUnit::Second => Arc::new(array.reinterpret_cast::()) as _, + TimeUnit::Millisecond => { + Arc::new(array.reinterpret_cast::()) as _ + } + TimeUnit::Microsecond => { + Arc::new(array.reinterpret_cast::()) as _ + } + TimeUnit::Nanosecond => { + Arc::new(array.reinterpret_cast::()) as _ + } + }, + ArrowType::Timestamp(time_unit, timezone) => match time_unit { + TimeUnit::Second => { + let array = array + .reinterpret_cast::() + .with_timezone_opt(timezone.clone()); + Arc::new(array) as _ + } + TimeUnit::Millisecond => { + let array = array + .reinterpret_cast::() + .with_timezone_opt(timezone.clone()); + Arc::new(array) as _ + } + TimeUnit::Microsecond => { + let array = array + .reinterpret_cast::() + .with_timezone_opt(timezone.clone()); + Arc::new(array) as _ + } + TimeUnit::Nanosecond => { + let array = array + .reinterpret_cast::() + .with_timezone_opt(timezone.clone()); + Arc::new(array) as _ + } + }, + ArrowType::Decimal64(p, s) => { + let array = array + .reinterpret_cast::() + .with_precision_and_scale(*p, *s)?; + Arc::new(array) as _ + } + ArrowType::Decimal128(p, s) => { + let array: Decimal128Array = array + .unary(|i| i as i128) + .with_precision_and_scale(*p, *s)?; + Arc::new(array) as _ + } + ArrowType::Decimal256(p, s) => { + let array: Decimal256Array = array + .unary(|i| i256::from_i128(i as i128)) + .with_precision_and_scale(*p, *s)?; + Arc::new(array) as _ + } + _ => unreachable!("Cannot coerce i64 to {target_type}"), + }) +} + +macro_rules! pack_dictionary_helper { + ($t:ty, $values:ident) => { + match $values.data_type() { + ArrowType::Int32 => pack_dictionary_impl::<$t, Int32Type>($values.as_primitive()), + ArrowType::Int64 => pack_dictionary_impl::<$t, Int64Type>($values.as_primitive()), + ArrowType::Float32 => pack_dictionary_impl::<$t, Float32Type>($values.as_primitive()), + ArrowType::Float64 => pack_dictionary_impl::<$t, Float64Type>($values.as_primitive()), + _ => unreachable!("Invalid physical type"), + } + }; +} + +fn pack_dictionary(key: &ArrowType, values: &dyn Array) -> Result { + downcast_integer! { + key => (pack_dictionary_helper, values), + _ => unreachable!("Invalid key type"), + } +} + +fn pack_dictionary_impl( + values: &PrimitiveArray, +) -> Result { + let mut builder = PrimitiveDictionaryBuilder::::with_capacity(1024, values.len()); + builder.extend(values); + Ok(Arc::new(builder.finish())) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::array_reader::test_util::EmptyPageIterator; + use crate::arrow::arrow_reader::DEFAULT_BATCH_SIZE; + use crate::basic::Encoding; + use crate::column::page::Page; + use crate::data_type::{Int32Type, Int64Type}; + use crate::schema::parser::parse_message_type; + use crate::schema::types::SchemaDescriptor; + use crate::util::InMemoryPageIterator; + use crate::util::test_common::rand_gen::make_pages; + use arrow::datatypes::ArrowPrimitiveType; + use arrow_array::{Array, Date32Array, PrimitiveArray}; + + use arrow::datatypes::DataType::{Date32, Decimal128}; + use rand::distr::uniform::SampleUniform; + use std::collections::VecDeque; + + #[allow(clippy::too_many_arguments)] + fn make_column_chunks( + column_desc: ColumnDescPtr, + encoding: Encoding, + num_levels: usize, + min_value: T::T, + max_value: T::T, + def_levels: &mut Vec, + rep_levels: &mut Vec, + values: &mut Vec, + page_lists: &mut Vec>, + use_v2: bool, + num_chunks: usize, + ) where + T::T: PartialOrd + SampleUniform + Copy, + { + for _i in 0..num_chunks { + let mut pages = VecDeque::new(); + let mut data = Vec::new(); + let mut page_def_levels = Vec::new(); + let mut page_rep_levels = Vec::new(); + + make_pages::( + column_desc.clone(), + encoding, + 1, + num_levels, + min_value, + max_value, + &mut page_def_levels, + &mut page_rep_levels, + &mut data, + &mut pages, + use_v2, + ); + + def_levels.append(&mut page_def_levels); + rep_levels.append(&mut page_rep_levels); + values.append(&mut data); + page_lists.push(Vec::from(pages)); + } + } + + #[test] + fn test_primitive_array_reader_empty_pages() { + // Construct column schema + let message_type = " + message test_schema { + REQUIRED INT32 leaf; + } + "; + + let schema = parse_message_type(message_type) + .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t)))) + .unwrap(); + + let mut array_reader = PrimitiveArrayReader::::new( + Box::::default(), + schema.column(0), + None, + DEFAULT_BATCH_SIZE, + None, + ) + .unwrap(); + + // expect no values to be read + let array = array_reader.next_batch(50).unwrap(); + assert!(array.is_empty()); + } + + #[test] + fn test_primitive_array_reader_data() { + // Construct column schema + let message_type = " + message test_schema { + REQUIRED INT32 leaf; + } + "; + + let schema = parse_message_type(message_type) + .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t)))) + .unwrap(); + + let column_desc = schema.column(0); + + // Construct page iterator + { + let mut data = Vec::new(); + let mut page_lists = Vec::new(); + make_column_chunks::( + column_desc.clone(), + Encoding::PLAIN, + 100, + 1, + 200, + &mut Vec::new(), + &mut Vec::new(), + &mut data, + &mut page_lists, + true, + 2, + ); + let page_iterator = InMemoryPageIterator::new(page_lists); + + let mut array_reader = PrimitiveArrayReader::::new( + Box::new(page_iterator), + column_desc, + None, + DEFAULT_BATCH_SIZE, + None, + ) + .unwrap(); + + // Read first 50 values, which are all from the first column chunk + let array = array_reader.next_batch(50).unwrap(); + let array = array.as_any().downcast_ref::().unwrap(); + + assert_eq!(&Int32Array::from(data[0..50].to_vec()), array); + + // Read next 100 values, the first 50 ones are from the first column chunk, + // and the last 50 ones are from the second column chunk + let array = array_reader.next_batch(100).unwrap(); + let array = array.as_any().downcast_ref::().unwrap(); + + assert_eq!(&Int32Array::from(data[50..150].to_vec()), array); + + // Try to read 100 values, however there are only 50 values + let array = array_reader.next_batch(100).unwrap(); + let array = array.as_any().downcast_ref::().unwrap(); + + assert_eq!(&Int32Array::from(data[150..200].to_vec()), array); + } + } + + macro_rules! test_primitive_array_reader_one_type { + ( + $arrow_parquet_type:ty, + $physical_type:expr, + $converted_type_str:expr, + $result_arrow_type:ty, + $result_arrow_cast_type:ty, + $result_primitive_type:ty + $(, $timezone:expr)? + ) => {{ + let message_type = format!( + " + message test_schema {{ + REQUIRED {:?} leaf ({}); + }} + ", + $physical_type, $converted_type_str + ); + let schema = parse_message_type(&message_type) + .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t)))) + .unwrap(); + + let column_desc = schema.column(0); + + // Construct page iterator + { + let mut data = Vec::new(); + let mut page_lists = Vec::new(); + make_column_chunks::<$arrow_parquet_type>( + column_desc.clone(), + Encoding::PLAIN, + 100, + 1, + 200, + &mut Vec::new(), + &mut Vec::new(), + &mut data, + &mut page_lists, + true, + 2, + ); + let page_iterator = InMemoryPageIterator::new(page_lists); + let mut array_reader = PrimitiveArrayReader::<$arrow_parquet_type>::new( + Box::new(page_iterator), + column_desc.clone(), + None, + DEFAULT_BATCH_SIZE, + None, + ) + .expect("Unable to get array reader"); + + let array = array_reader + .next_batch(50) + .expect("Unable to get batch from reader"); + + let result_data_type = <$result_arrow_type>::DATA_TYPE; + let array = array + .as_any() + .downcast_ref::>() + .expect( + format!( + "Unable to downcast {:?} to {:?}", + array.data_type(), + result_data_type + ) + .as_str(), + ) + $(.clone().with_timezone($timezone))? + ; + + // create expected array as primitive, and cast to result type + let expected = PrimitiveArray::<$result_arrow_cast_type>::from( + data[0..50] + .iter() + .map(|x| *x as $result_primitive_type) + .collect::>(), + ); + let expected = Arc::new(expected) as ArrayRef; + let expected = arrow::compute::cast(&expected, &result_data_type) + .expect("Unable to cast expected array"); + assert_eq!(expected.data_type(), &result_data_type); + let expected = expected + .as_any() + .downcast_ref::>() + .expect( + format!( + "Unable to downcast expected {:?} to {:?}", + expected.data_type(), + result_data_type + ) + .as_str(), + ) + $(.clone().with_timezone($timezone))? + ; + assert_eq!(expected, array); + } + }}; + } + + #[test] + fn test_primitive_array_reader_temporal_types() { + test_primitive_array_reader_one_type!( + crate::data_type::Int32Type, + PhysicalType::INT32, + "DATE", + arrow::datatypes::Date32Type, + arrow::datatypes::Int32Type, + i32 + ); + test_primitive_array_reader_one_type!( + crate::data_type::Int32Type, + PhysicalType::INT32, + "TIME_MILLIS", + arrow::datatypes::Time32MillisecondType, + arrow::datatypes::Int32Type, + i32 + ); + test_primitive_array_reader_one_type!( + crate::data_type::Int64Type, + PhysicalType::INT64, + "TIME_MICROS", + arrow::datatypes::Time64MicrosecondType, + arrow::datatypes::Int64Type, + i64 + ); + test_primitive_array_reader_one_type!( + crate::data_type::Int64Type, + PhysicalType::INT64, + "TIMESTAMP_MILLIS", + arrow::datatypes::TimestampMillisecondType, + arrow::datatypes::Int64Type, + i64, + "UTC" + ); + test_primitive_array_reader_one_type!( + crate::data_type::Int64Type, + PhysicalType::INT64, + "TIMESTAMP_MICROS", + arrow::datatypes::TimestampMicrosecondType, + arrow::datatypes::Int64Type, + i64, + "UTC" + ); + } + + #[test] + fn test_primitive_array_reader_def_and_rep_levels() { + // Construct column schema + let message_type = " + message test_schema { + REPEATED Group test_mid { + OPTIONAL INT32 leaf; + } + } + "; + + let schema = parse_message_type(message_type) + .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t)))) + .unwrap(); + + let column_desc = schema.column(0); + + // Construct page iterator + { + let mut def_levels = Vec::new(); + let mut rep_levels = Vec::new(); + let mut page_lists = Vec::new(); + make_column_chunks::( + column_desc.clone(), + Encoding::PLAIN, + 100, + 1, + 200, + &mut def_levels, + &mut rep_levels, + &mut Vec::new(), + &mut page_lists, + true, + 2, + ); + + let page_iterator = InMemoryPageIterator::new(page_lists); + + let mut array_reader = PrimitiveArrayReader::::new( + Box::new(page_iterator), + column_desc, + None, + DEFAULT_BATCH_SIZE, + None, + ) + .unwrap(); + + let mut accu_len: usize = 0; + + // Read first 50 values, which are all from the first column chunk + let array = array_reader.next_batch(50).unwrap(); + assert_eq!( + Some(&def_levels[accu_len..(accu_len + array.len())]), + array_reader.get_def_levels() + ); + assert_eq!( + Some(&rep_levels[accu_len..(accu_len + array.len())]), + array_reader.get_rep_levels() + ); + accu_len += array.len(); + + // Read next 100 values, the first 50 ones are from the first column chunk, + // and the last 50 ones are from the second column chunk + let array = array_reader.next_batch(100).unwrap(); + assert_eq!( + Some(&def_levels[accu_len..(accu_len + array.len())]), + array_reader.get_def_levels() + ); + assert_eq!( + Some(&rep_levels[accu_len..(accu_len + array.len())]), + array_reader.get_rep_levels() + ); + accu_len += array.len(); + + // Try to read 100 values, however there are only 50 values + let array = array_reader.next_batch(100).unwrap(); + assert_eq!( + Some(&def_levels[accu_len..(accu_len + array.len())]), + array_reader.get_def_levels() + ); + assert_eq!( + Some(&rep_levels[accu_len..(accu_len + array.len())]), + array_reader.get_rep_levels() + ); + } + } + + #[test] + fn test_primitive_array_reader_decimal_types() { + // parquet `INT32` to decimal + let message_type = " + message test_schema { + REQUIRED INT32 decimal1 (DECIMAL(8,2)); + } + "; + let schema = parse_message_type(message_type) + .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t)))) + .unwrap(); + let column_desc = schema.column(0); + + // create the array reader + { + let mut data = Vec::new(); + let mut page_lists = Vec::new(); + make_column_chunks::( + column_desc.clone(), + Encoding::PLAIN, + 100, + -99999999, + 99999999, + &mut Vec::new(), + &mut Vec::new(), + &mut data, + &mut page_lists, + true, + 2, + ); + let page_iterator = InMemoryPageIterator::new(page_lists); + + let mut array_reader = PrimitiveArrayReader::::new( + Box::new(page_iterator), + column_desc, + None, + DEFAULT_BATCH_SIZE, + None, + ) + .unwrap(); + + // read data from the reader + // the data type is decimal(8,2) + let array = array_reader.next_batch(50).unwrap(); + assert_eq!(array.data_type(), &Decimal128(8, 2)); + let array = array.as_any().downcast_ref::().unwrap(); + let data_decimal_array = data[0..50] + .iter() + .copied() + .map(|v| Some(v as i128)) + .collect::() + .with_precision_and_scale(8, 2) + .unwrap(); + assert_eq!(array, &data_decimal_array); + + // not equal with different data type(precision and scale) + let data_decimal_array = data[0..50] + .iter() + .copied() + .map(|v| Some(v as i128)) + .collect::() + .with_precision_and_scale(9, 0) + .unwrap(); + assert_ne!(array, &data_decimal_array) + } + + // parquet `INT64` to decimal + let message_type = " + message test_schema { + REQUIRED INT64 decimal1 (DECIMAL(18,4)); + } + "; + let schema = parse_message_type(message_type) + .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t)))) + .unwrap(); + let column_desc = schema.column(0); + + // create the array reader + { + let mut data = Vec::new(); + let mut page_lists = Vec::new(); + make_column_chunks::( + column_desc.clone(), + Encoding::PLAIN, + 100, + -999999999999999999, + 999999999999999999, + &mut Vec::new(), + &mut Vec::new(), + &mut data, + &mut page_lists, + true, + 2, + ); + let page_iterator = InMemoryPageIterator::new(page_lists); + + let mut array_reader = PrimitiveArrayReader::::new( + Box::new(page_iterator), + column_desc, + None, + DEFAULT_BATCH_SIZE, + None, + ) + .unwrap(); + + // read data from the reader + // the data type is decimal(18,4) + let array = array_reader.next_batch(50).unwrap(); + assert_eq!(array.data_type(), &Decimal128(18, 4)); + let array = array.as_any().downcast_ref::().unwrap(); + let data_decimal_array = data[0..50] + .iter() + .copied() + .map(|v| Some(v as i128)) + .collect::() + .with_precision_and_scale(18, 4) + .unwrap(); + assert_eq!(array, &data_decimal_array); + + // not equal with different data type(precision and scale) + let data_decimal_array = data[0..50] + .iter() + .copied() + .map(|v| Some(v as i128)) + .collect::() + .with_precision_and_scale(34, 0) + .unwrap(); + assert_ne!(array, &data_decimal_array) + } + } + + #[test] + fn test_primitive_array_reader_date32_type() { + // parquet `INT32` to date + let message_type = " + message test_schema { + REQUIRED INT32 date1 (DATE); + } + "; + let schema = parse_message_type(message_type) + .map(|t| Arc::new(SchemaDescriptor::new(Arc::new(t)))) + .unwrap(); + let column_desc = schema.column(0); + + // create the array reader + { + let mut data = Vec::new(); + let mut page_lists = Vec::new(); + make_column_chunks::( + column_desc.clone(), + Encoding::PLAIN, + 100, + -99999999, + 99999999, + &mut Vec::new(), + &mut Vec::new(), + &mut data, + &mut page_lists, + true, + 2, + ); + let page_iterator = InMemoryPageIterator::new(page_lists); + + let mut array_reader = PrimitiveArrayReader::::new( + Box::new(page_iterator), + column_desc, + None, + DEFAULT_BATCH_SIZE, + None, + ) + .unwrap(); + + // read data from the reader + // the data type is date + let array = array_reader.next_batch(50).unwrap(); + assert_eq!(array.data_type(), &Date32); + let array = array.as_any().downcast_ref::().unwrap(); + let data_date_array = data[0..50] + .iter() + .copied() + .map(Some) + .collect::(); + assert_eq!(array, &data_date_array); + } + } +} diff --git a/vendor/parquet/src/arrow/array_reader/row_group_cache.rs b/vendor/parquet/src/arrow/array_reader/row_group_cache.rs new file mode 100644 index 0000000..ef726e1 --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/row_group_cache.rs @@ -0,0 +1,206 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_array::{Array, ArrayRef}; +use arrow_schema::DataType; +use std::collections::HashMap; + +/// Starting row ID for this batch +/// +/// The `BatchID` is used to identify batches of rows within a row group. +/// +/// The row_index in the id are relative to the rows being read from the +/// underlying column reader (which might already have a RowSelection applied) +/// +/// The `BatchID` for any particular row is `row_index / batch_size`. The +/// integer division ensures that rows in the same batch share the same +/// the BatchID which can be calculated quickly from the row index +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] +pub struct BatchID { + pub val: usize, +} + +/// Cache key that uniquely identifies a batch within a row group +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct CacheKey { + /// Column index in the row group + pub column_idx: usize, + /// Starting row ID for this batch + pub batch_id: BatchID, +} + +fn get_array_memory_size_for_cache(array: &ArrayRef) -> usize { + match array.data_type() { + // TODO: this is temporary workaround. It's very difficult to measure the actual memory usage of one StringViewArray, + // because the underlying buffer is shared with multiple StringViewArrays. + DataType::Utf8View => { + use arrow_array::cast::AsArray; + let array = array.as_string_view(); + array.len() * 16 + array.total_buffer_bytes_used() + std::mem::size_of_val(array) + } + _ => array.get_array_memory_size(), + } +} + +/// Row group cache that stores decoded arrow arrays at batch granularity +/// +/// This cache is designed to avoid duplicate decoding when the same column +/// appears in both filter predicates and output projection. +#[derive(Debug)] +pub struct RowGroupCache { + /// Cache storage mapping (column_idx, row_id) -> ArrayRef + cache: HashMap, + /// Cache granularity + batch_size: usize, + /// Maximum cache size in bytes + max_cache_bytes: usize, + /// Current cache size in bytes + current_cache_size: usize, +} + +impl RowGroupCache { + /// Creates a new empty row group cache + pub fn new(batch_size: usize, max_cache_bytes: usize) -> Self { + Self { + cache: HashMap::new(), + batch_size, + max_cache_bytes, + current_cache_size: 0, + } + } + + /// Inserts an array into the cache for the given column and starting row ID + /// Returns true if the array was inserted, false if it would exceed the cache size limit + pub fn insert(&mut self, column_idx: usize, batch_id: BatchID, array: ArrayRef) -> bool { + let array_size = get_array_memory_size_for_cache(&array); + + // Check if adding this array would exceed the cache size limit + if self.current_cache_size + array_size > self.max_cache_bytes { + return false; // Cache is full, don't insert + } + + let key = CacheKey { + column_idx, + batch_id, + }; + + let existing = self.cache.insert(key, array); + assert!(existing.is_none()); + self.current_cache_size += array_size; + true + } + + /// Retrieves a cached array for the given column and row ID + /// Returns None if not found in cache + pub fn get(&self, column_idx: usize, batch_id: BatchID) -> Option { + let key = CacheKey { + column_idx, + batch_id, + }; + self.cache.get(&key).cloned() + } + + /// Gets the batch size for this cache + pub fn batch_size(&self) -> usize { + self.batch_size + } + + /// Removes a cached array for the given column and row ID + /// Returns true if the entry was found and removed, false otherwise + pub fn remove(&mut self, column_idx: usize, batch_id: BatchID) -> bool { + let key = CacheKey { + column_idx, + batch_id, + }; + if let Some(array) = self.cache.remove(&key) { + self.current_cache_size -= get_array_memory_size_for_cache(&array); + true + } else { + false + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::{ArrayRef, Int32Array}; + use std::sync::Arc; + + #[test] + fn test_cache_basic_operations() { + let mut cache = RowGroupCache::new(1000, usize::MAX); + + // Create test array + let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])); + + // Test insert and get + let batch_id = BatchID { val: 0 }; + assert!(cache.insert(0, batch_id, array.clone())); + let retrieved = cache.get(0, batch_id); + assert!(retrieved.is_some()); + assert_eq!(retrieved.unwrap().len(), 5); + + // Test miss + let miss = cache.get(1, batch_id); + assert!(miss.is_none()); + + // Test different row_id + let miss = cache.get(0, BatchID { val: 1000 }); + assert!(miss.is_none()); + } + + #[test] + fn test_cache_remove() { + let mut cache = RowGroupCache::new(1000, usize::MAX); + + // Create test arrays + let array1: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let array2: ArrayRef = Arc::new(Int32Array::from(vec![4, 5, 6])); + + // Insert arrays + assert!(cache.insert(0, BatchID { val: 0 }, array1.clone())); + assert!(cache.insert(0, BatchID { val: 1000 }, array2.clone())); + assert!(cache.insert(1, BatchID { val: 0 }, array1.clone())); + + // Verify they're there + assert!(cache.get(0, BatchID { val: 0 }).is_some()); + assert!(cache.get(0, BatchID { val: 1000 }).is_some()); + assert!(cache.get(1, BatchID { val: 0 }).is_some()); + + // Remove one entry + let removed = cache.remove(0, BatchID { val: 0 }); + assert!(removed); + assert!(cache.get(0, BatchID { val: 0 }).is_none()); + + // Other entries should still be there + assert!(cache.get(0, BatchID { val: 1000 }).is_some()); + assert!(cache.get(1, BatchID { val: 0 }).is_some()); + + // Try to remove non-existent entry + let not_removed = cache.remove(0, BatchID { val: 0 }); + assert!(!not_removed); + + // Remove remaining entries + assert!(cache.remove(0, BatchID { val: 1000 })); + assert!(cache.remove(1, BatchID { val: 0 })); + + // Cache should be empty + assert!(cache.get(0, BatchID { val: 1000 }).is_none()); + assert!(cache.get(1, BatchID { val: 0 }).is_none()); + } +} diff --git a/vendor/parquet/src/arrow/array_reader/row_group_index.rs b/vendor/parquet/src/arrow/array_reader/row_group_index.rs new file mode 100644 index 0000000..b4c296c --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/row_group_index.rs @@ -0,0 +1,383 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::array_reader::ArrayReader; +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::{ParquetMetaData, RowGroupMetaData}; +use arrow_array::{ArrayRef, Int64Array}; +use arrow_schema::DataType; +use std::any::Any; +use std::collections::HashMap; +use std::sync::Arc; + +pub(crate) struct RowGroupIndexReader { + buffered_indices: Vec, + state: ReaderState, +} + +enum ReaderState { + // fast path: single row group with constant index value + SingleRowGroup { + index: i64, + remaining_rows: usize, + }, + // general path: multiple row groups with iterator + MultipleRowGroups { + remaining_indices: std::iter::Flatten>>, + }, +} + +impl RowGroupIndexReader { + pub(crate) fn try_new<'a>( + parquet_metadata: &'a ParquetMetaData, + row_groups: impl Iterator, + ) -> Result { + let row_groups: Vec<_> = row_groups.collect(); + + // optimize for single row group case + if row_groups.len() == 1 { + let rg = row_groups[0]; + let ordinal = rg.ordinal().ok_or_else(|| { + ParquetError::General( + "Row group missing ordinal field, required to compute row group indices" + .to_string(), + ) + })?; + + // find the row group index by a linear scan through metadata + // this is O(n) but only done once, avoiding HashMap allocation + let index = parquet_metadata + .row_groups() + .iter() + .enumerate() + .find_map(|(idx, metadata_rg)| { + if metadata_rg.ordinal() == Some(ordinal) { + Some(idx as i64) + } else { + None + } + }) + .ok_or_else(|| { + ParquetError::General(format!( + "Row group with ordinal {} not found in metadata", + ordinal + )) + })?; + + return Ok(Self { + buffered_indices: Vec::new(), + state: ReaderState::SingleRowGroup { + index, + remaining_rows: rg.num_rows() as usize, + }, + }); + } + + // general path: many row groups + // builds a mapping from ordinal to row group index + // this is O(n) where n is the total number of row groups in the file + let ordinal_to_index: HashMap = + HashMap::from_iter(parquet_metadata.row_groups().iter().enumerate().filter_map( + |(row_group_index, rg)| { + rg.ordinal() + .map(|ordinal| (ordinal, row_group_index as i64)) + }, + )); + + // build repeating iterators in the order specified by the row_groups iterator + // this is O(m) where m is the number of selected row groups + let repeated_indices = row_groups + .iter() + .map(|rg| { + let ordinal = rg.ordinal().ok_or_else(|| { + ParquetError::General( + "Row group missing ordinal field, required to compute row group indices" + .to_string(), + ) + })?; + + let row_group_index = ordinal_to_index.get(&ordinal).ok_or_else(|| { + ParquetError::General(format!( + "Row group with ordinal {} not found in metadata", + ordinal + )) + })?; + + // repeat row group index for each row in this row group + Ok(std::iter::repeat_n( + *row_group_index, + rg.num_rows() as usize, + )) + }) + .collect::>>()?; + + Ok(Self { + buffered_indices: Vec::new(), + state: ReaderState::MultipleRowGroups { + remaining_indices: repeated_indices.into_iter().flatten(), + }, + }) + } +} + +impl ArrayReader for RowGroupIndexReader { + fn read_records(&mut self, batch_size: usize) -> Result { + match &mut self.state { + ReaderState::SingleRowGroup { + index, + remaining_rows, + } => { + let num_to_read = batch_size.min(*remaining_rows); + self.buffered_indices + .resize(self.buffered_indices.len() + num_to_read, *index); + *remaining_rows -= num_to_read; + Ok(num_to_read) + } + ReaderState::MultipleRowGroups { remaining_indices } => { + let starting_len = self.buffered_indices.len(); + self.buffered_indices + .extend((remaining_indices.by_ref()).take(batch_size)); + Ok(self.buffered_indices.len() - starting_len) + } + } + } + + fn skip_records(&mut self, num_records: usize) -> Result { + match &mut self.state { + ReaderState::SingleRowGroup { remaining_rows, .. } => { + let num_to_skip = num_records.min(*remaining_rows); + *remaining_rows -= num_to_skip; + Ok(num_to_skip) + } + ReaderState::MultipleRowGroups { remaining_indices } => { + // TODO: Use advance_by when it stabilizes to improve performance + Ok((remaining_indices.by_ref()).take(num_records).count()) + } + } + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn get_data_type(&self) -> &DataType { + &DataType::Int64 + } + + fn consume_batch(&mut self) -> Result { + Ok(Arc::new(Int64Array::from_iter( + self.buffered_indices.drain(..), + ))) + } + + fn get_def_levels(&self) -> Option<&[i16]> { + None + } + + fn get_rep_levels(&self) -> Option<&[i16]> { + None + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::basic::Type as PhysicalType; + use crate::file::metadata::{ + ColumnChunkMetaData, FileMetaData, ParquetMetaData, RowGroupMetaData, + }; + use crate::schema::types::{SchemaDescriptor, Type as SchemaType}; + use std::sync::Arc; + + fn create_test_schema() -> Arc { + let schema = SchemaType::group_type_builder("schema") + .with_fields(vec![Arc::new( + SchemaType::primitive_type_builder("test_col", PhysicalType::INT32) + .build() + .unwrap(), + )]) + .build() + .unwrap(); + Arc::new(SchemaDescriptor::new(Arc::new(schema))) + } + + fn create_test_parquet_metadata(row_groups: Vec<(i16, i64)>) -> ParquetMetaData { + let schema_descr = create_test_schema(); + + let mut row_group_metas = vec![]; + for (ordinal, num_rows) in row_groups { + let columns: Vec<_> = schema_descr + .columns() + .iter() + .map(|col| ColumnChunkMetaData::builder(col.clone()).build().unwrap()) + .collect(); + + let row_group = RowGroupMetaData::builder(schema_descr.clone()) + .set_num_rows(num_rows) + .set_ordinal(ordinal) + .set_total_byte_size(100) + .set_column_metadata(columns) + .build() + .unwrap(); + row_group_metas.push(row_group); + } + + let total_rows: i64 = row_group_metas.iter().map(|rg| rg.num_rows()).sum(); + let file_metadata = FileMetaData::new(1, total_rows, None, None, schema_descr, None); + + ParquetMetaData::new(file_metadata, row_group_metas) + } + + #[test] + fn test_row_group_index_reader_basic() { + // create metadata with 3 row groups, each with varying number of rows + let metadata = create_test_parquet_metadata(vec![ + (0, 2), // rg: 0, ordinal: 0, 2 rows + (1, 3), // rg: 1, ordinal: 1, 3 rows + (2, 1), // rg: 2, ordinal: 2, 1 row + ]); + + let selected_row_groups: Vec<_> = metadata.row_groups().iter().collect(); + + let mut reader = + RowGroupIndexReader::try_new(&metadata, selected_row_groups.into_iter()).unwrap(); + + // 2 rows + 3 rows + 1 row + let num_read = reader.read_records(6).unwrap(); + assert_eq!(num_read, 6); + + let array = reader.consume_batch().unwrap(); + let indices = array.as_any().downcast_ref::().unwrap(); + + let actual: Vec = indices.iter().map(|v| v.unwrap()).collect(); + assert_eq!(actual, [0, 0, 1, 1, 1, 2],); + } + + #[test] + fn test_row_group_index_reader_reverse_order() { + // create metadata with 3 row groups, each rg has 2 rows + let metadata = create_test_parquet_metadata(vec![(0, 2), (1, 2), (2, 2)]); + + // select only rgs with ordinals 2 and 0 (in that order) + // means select row group 2 first, then row group 0, skipping row group 1 + let selected_row_groups: Vec<_> = + vec![&metadata.row_groups()[2], &metadata.row_groups()[0]]; + + let mut reader = + RowGroupIndexReader::try_new(&metadata, selected_row_groups.into_iter()).unwrap(); + + let num_read = reader.read_records(6).unwrap(); + // 2 rgs * 2 rows each + assert_eq!(num_read, 4); + + let array = reader.consume_batch().unwrap(); + let indices = array.as_any().downcast_ref::().unwrap(); + + let actual: Vec = indices.iter().map(|v| v.unwrap()).collect(); + + assert_eq!(actual, [2, 2, 0, 0],); + } + + #[test] + fn test_row_group_index_reader_skip_records() { + // rg 0: 3 rows, rg 1: 4 rows, rg 2: 2 rows + // [0, 0, 0, 1, 1, 1, 1, 2, 2] + let metadata = create_test_parquet_metadata(vec![(0, 3), (1, 4), (2, 2)]); + + let selected_row_groups = metadata.row_groups().iter().collect::>(); + + let mut reader = + RowGroupIndexReader::try_new(&metadata, selected_row_groups.into_iter()).unwrap(); + + // skip first 5 rows + // [0, 0, 0, 1, 1, 1, 1, 2, 2] + // |---- skip ---| + let num_skipped = reader.skip_records(5).unwrap(); + assert_eq!(num_skipped, 5); + + let num_read = reader.read_records(10).unwrap(); + assert_eq!(num_read, 4); + + let array = reader.consume_batch().unwrap(); + let indices = array.as_any().downcast_ref::().unwrap(); + + let actual = indices.iter().map(|v| v.unwrap()).collect::>(); + assert_eq!(actual, [1, 1, 2, 2]); + } + + #[test] + fn test_row_group_index_reader_single_row_group() { + // 3 row groups + let metadata = create_test_parquet_metadata(vec![(0, 2), (1, 3), (2, 5)]); + + // select last row group [2, 2, 2, 2, 2] + let selected_row_groups = [&metadata.row_groups()[2]]; + + let mut reader = + RowGroupIndexReader::try_new(&metadata, selected_row_groups.into_iter()).unwrap(); + + assert!(matches!( + reader.state, + ReaderState::SingleRowGroup { index: 2, .. } + )); + + let num_read = reader.read_records(10).unwrap(); + assert_eq!(num_read, 5); + + let array = reader.consume_batch().unwrap(); + let indices = array.as_any().downcast_ref::().unwrap(); + + let actual: Vec = indices.iter().map(|v| v.unwrap()).collect(); + assert_eq!(actual, [2, 2, 2, 2, 2]); + } + + #[test] + fn test_row_group_index_reader_single_row_group_with_skip() { + let metadata = create_test_parquet_metadata(vec![(0, 10), (1, 10), (2, 10)]); + + // second row group (index 1, ordinal 1 with 10 rows) + let selected_row_groups = vec![&metadata.row_groups()[1]]; + + let mut reader = + RowGroupIndexReader::try_new(&metadata, selected_row_groups.into_iter()).unwrap(); + + assert!(matches!( + reader.state, + ReaderState::SingleRowGroup { index: 1, .. } + )); + + // skip first 3 rows + let num_skipped = reader.skip_records(3).unwrap(); + assert_eq!(num_skipped, 3); + + // read next 5 rows + let num_read = reader.read_records(5).unwrap(); + assert_eq!(num_read, 5); + + let array = reader.consume_batch().unwrap(); + let indices = array.as_any().downcast_ref::().unwrap(); + + let actual: Vec = indices.iter().map(|v| v.unwrap()).collect(); + assert_eq!(actual, [1, 1, 1, 1, 1]); + + if let ReaderState::SingleRowGroup { remaining_rows, .. } = reader.state { + assert_eq!(remaining_rows, 2); + } else { + panic!("Expected SingleRowGroup state"); + } + } +} diff --git a/vendor/parquet/src/arrow/array_reader/row_number.rs b/vendor/parquet/src/arrow/array_reader/row_number.rs new file mode 100644 index 0000000..def8864 --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/row_number.rs @@ -0,0 +1,305 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::array_reader::ArrayReader; +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::{ParquetMetaData, RowGroupMetaData}; +use arrow_array::{ArrayRef, Int64Array}; +use arrow_schema::DataType; +use std::any::Any; +use std::collections::{HashMap, VecDeque}; +use std::ops::Range; +use std::sync::Arc; + +/// Tracks row numbers within a Parquet file and emits them as an `Int64Array`. +pub(crate) struct RowNumberReader { + /// Pre-computed row ranges that are not read yet. + /// + /// This reader only keeps track of the ranges of row numbers for each row group. The range is + /// not materialized into a full array until it's needed. + remaining_row_ranges: VecDeque>, + /// Row ranges read but not emitted. + /// + /// These are either full or partial (split) row ranges taken from `remaining_row_ranges`. + buffered_row_ranges: Vec>, +} + +impl RowNumberReader { + pub(crate) fn try_new<'a>( + parquet_metadata: &'a ParquetMetaData, + row_groups: impl Iterator, + ) -> Result { + // Pass 1: Build a map from ordinal to first_row_index + // This is O(M) where M is the total number of row groups in the file + let mut ordinal_to_offset: HashMap = HashMap::new(); + let mut first_row_index: i64 = 0; + + for rg in parquet_metadata.row_groups() { + if let Some(ordinal) = rg.ordinal() { + ordinal_to_offset.insert(ordinal, first_row_index); + } + first_row_index += rg.num_rows(); + } + + // Pass 2: Build ranges in the order specified by the row_groups iterator + // This is O(N) where N is the number of selected row groups + // This preserves the user's requested order instead of sorting by ordinal + let ranges: VecDeque<_> = row_groups + .map(|rg| { + let ordinal = rg.ordinal().ok_or_else(|| { + ParquetError::General( + "Row group missing ordinal field, required to compute row numbers" + .to_string(), + ) + })?; + + let offset = ordinal_to_offset.get(&ordinal).ok_or_else(|| { + ParquetError::General(format!( + "Row group with ordinal {} not found in metadata", + ordinal + )) + })?; + + Ok(*offset..*offset + rg.num_rows()) + }) + .collect::>()?; + + Ok(Self { + buffered_row_ranges: Vec::new(), + remaining_row_ranges: ranges, + }) + } + + /// Take up to `count` rows from the first range, splitting it if needed. + /// + /// Returns `None` if no ranges remain. + fn take_range(&mut self, count: usize) -> Option> { + let first = self.remaining_row_ranges.front_mut()?; + if (first.end - first.start) <= count as i64 { + // take out the full range + self.remaining_row_ranges.pop_front() + } else { + // first range has more rows than we need. + // so we split the range and put the remaining back. + let split = first.start + count as i64; + let taken = first.start..split; + first.start = split; + Some(taken) + } + } +} + +impl ArrayReader for RowNumberReader { + fn read_records(&mut self, batch_size: usize) -> Result { + let mut remaining = batch_size; + while remaining > 0 { + let Some(range) = self.take_range(remaining) else { + break; + }; + remaining -= (range.end - range.start) as usize; + self.buffered_row_ranges.push(range); + } + Ok(batch_size - remaining) + } + + fn skip_records(&mut self, num_records: usize) -> Result { + let mut remaining = num_records; + while remaining > 0 { + let Some(range) = self.take_range(remaining) else { + break; + }; + remaining -= (range.end - range.start) as usize; + } + Ok(num_records - remaining) + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn get_data_type(&self) -> &DataType { + &DataType::Int64 + } + + fn consume_batch(&mut self) -> Result { + let total_rows: i64 = self + .buffered_row_ranges + .iter() + .map(|range| range.end - range.start) + .sum(); + let mut result = Vec::with_capacity(total_rows as usize); + + for range in self.buffered_row_ranges.drain(..) { + result.extend(range); + } + + Ok(Arc::new(Int64Array::from(result))) + } + + fn get_def_levels(&self) -> Option<&[i16]> { + None + } + + fn get_rep_levels(&self) -> Option<&[i16]> { + None + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::basic::Type as PhysicalType; + use crate::file::metadata::{ + ColumnChunkMetaData, FileMetaData, ParquetMetaData, RowGroupMetaData, + }; + use crate::schema::types::{SchemaDescriptor, Type as SchemaType}; + use std::sync::Arc; + + fn create_test_schema() -> Arc { + let schema = SchemaType::group_type_builder("schema") + .with_fields(vec![Arc::new( + SchemaType::primitive_type_builder("test_col", PhysicalType::INT32) + .build() + .unwrap(), + )]) + .build() + .unwrap(); + Arc::new(SchemaDescriptor::new(Arc::new(schema))) + } + + fn create_test_parquet_metadata(row_groups: Vec<(i16, i64)>) -> ParquetMetaData { + let schema_descr = create_test_schema(); + + let mut row_group_metas = vec![]; + for (ordinal, num_rows) in row_groups { + let columns: Vec<_> = schema_descr + .columns() + .iter() + .map(|col| ColumnChunkMetaData::builder(col.clone()).build().unwrap()) + .collect(); + + let row_group = RowGroupMetaData::builder(schema_descr.clone()) + .set_num_rows(num_rows) + .set_ordinal(ordinal) + .set_total_byte_size(100) + .set_column_metadata(columns) + .build() + .unwrap(); + row_group_metas.push(row_group); + } + + let total_rows: i64 = row_group_metas.iter().map(|rg| rg.num_rows()).sum(); + let file_metadata = FileMetaData::new( + 1, // version + total_rows, // num_rows + None, // created_by + None, // key_value_metadata + schema_descr, // schema_descr + None, // column_orders + ); + + ParquetMetaData::new(file_metadata, row_group_metas) + } + + fn consume_row_numbers(reader: &mut RowNumberReader) -> Vec { + let array = reader.consume_batch().unwrap(); + array + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + } + + fn reader_for_all(metadata: &ParquetMetaData) -> RowNumberReader { + let all_rgs: Vec<_> = metadata.row_groups().iter().collect(); + RowNumberReader::try_new(metadata, all_rgs.into_iter()).unwrap() + } + + #[test] + fn test_row_number_reader_reverse_order() { + // Create metadata with 3 row groups, each with 2 rows + let metadata = create_test_parquet_metadata(vec![ + (0, 2), // Row group 0: ordinal=0, rows 0-1 + (1, 2), // Row group 1: ordinal=1, rows 2-3 + (2, 2), // Row group 2: ordinal=2, rows 4-5 + ]); + + // Select only row groups with ordinals 2 and 0 (in that order) + // This means we want row group 2 first, then row group 0, skipping row group 1 + let selected_row_groups: Vec<_> = vec![ + &metadata.row_groups()[2], // ordinal 2 + &metadata.row_groups()[0], // ordinal 0 + ]; + + let mut reader = + RowNumberReader::try_new(&metadata, selected_row_groups.into_iter()).unwrap(); + + // Read all row numbers + let num_read = reader.read_records(6).unwrap(); + assert_eq!(num_read, 4); // Should read 4 rows total (2 from each selected group) + + // Expected: row group 2 first (rows 4-5), then row group 0 (rows 0-1) + assert_eq!(consume_row_numbers(&mut reader), vec![4, 5, 0, 1]); + } + + #[test] + fn test_range_splitting_across_batches() { + // One row group with 10 rows + let metadata = create_test_parquet_metadata(vec![(0, 10)]); + let mut reader = reader_for_all(&metadata); + + assert_eq!(reader.read_records(3).unwrap(), 3); + assert_eq!(consume_row_numbers(&mut reader), vec![0, 1, 2]); + + assert_eq!(reader.read_records(3).unwrap(), 3); + assert_eq!(consume_row_numbers(&mut reader), vec![3, 4, 5]); + + assert_eq!(reader.read_records(3).unwrap(), 3); + assert_eq!(consume_row_numbers(&mut reader), vec![6, 7, 8]); + + // Only 1 row left, requesting 3 + assert_eq!(reader.read_records(3).unwrap(), 1); + assert_eq!(consume_row_numbers(&mut reader), vec![9]); + } + + #[test] + fn test_interleaved_skip_and_read() { + // Row group 0: rows 0..5 + // Row group 1: rows 5..10 + let metadata = create_test_parquet_metadata(vec![(0, 5), (1, 5)]); + let mut reader = reader_for_all(&metadata); + + assert_eq!(reader.skip_records(2).unwrap(), 2); // skip [0,1] + assert_eq!(reader.read_records(2).unwrap(), 2); // read [2,3] + assert_eq!(reader.skip_records(3).unwrap(), 3); // skip [4] then [5,6] + assert_eq!(reader.read_records(3).unwrap(), 3); // read [7,8,9] + + assert_eq!(consume_row_numbers(&mut reader), vec![2, 3, 7, 8, 9]); + } + + #[test] + fn test_skip_then_read() { + // One row group with 10 rows + let metadata = create_test_parquet_metadata(vec![(0, 10)]); + let mut reader = reader_for_all(&metadata); + + assert_eq!(reader.skip_records(3).unwrap(), 3); + assert_eq!(reader.read_records(4).unwrap(), 4); + assert_eq!(consume_row_numbers(&mut reader), vec![3, 4, 5, 6]); + } +} diff --git a/vendor/parquet/src/arrow/array_reader/struct_array.rs b/vendor/parquet/src/arrow/array_reader/struct_array.rs new file mode 100644 index 0000000..2032768 --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/struct_array.rs @@ -0,0 +1,297 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::array_reader::ArrayReader; +use crate::arrow::record_reader::definition_levels::build_filtered_validity_bitmap; +use crate::errors::{ParquetError, Result}; +use arrow_array::{Array, ArrayRef, StructArray, builder::BooleanBufferBuilder}; +use arrow_buffer::NullBuffer; +use arrow_schema::{DataType as ArrowType, DataType}; +use std::any::Any; +use std::sync::Arc; + +/// Implementation of struct array reader. +pub struct StructArrayReader { + children: Vec>, + data_type: ArrowType, + struct_def_level: i16, + struct_rep_level: i16, + nullable: bool, + /// When set, entries with def < threshold are excluded from child + /// arrays. Set when this struct is inside a list (to the parent + /// list's def_level). + padding_threshold: Option, +} + +impl StructArrayReader { + /// Construct struct array reader. + pub fn new( + data_type: ArrowType, + children: Vec>, + def_level: i16, + rep_level: i16, + nullable: bool, + padding_threshold: Option, + ) -> Self { + Self { + data_type, + children, + struct_def_level: def_level, + struct_rep_level: rep_level, + nullable, + padding_threshold, + } + } +} + +impl ArrayReader for StructArrayReader { + fn as_any(&self) -> &dyn Any { + self + } + + /// Returns data type. + /// This must be a struct. + fn get_data_type(&self) -> &ArrowType { + &self.data_type + } + + fn read_records(&mut self, batch_size: usize) -> Result { + let mut read = None; + for child in self.children.iter_mut() { + let child_read = child.read_records(batch_size)?; + match read { + Some(expected) => { + if expected != child_read { + return Err(general_err!( + "StructArrayReader out of sync in read_records, expected {} read, got {}", + expected, + child_read + )); + } + } + None => read = Some(child_read), + } + } + Ok(read.unwrap_or(0)) + } + + fn consume_batch(&mut self) -> Result { + if self.children.is_empty() { + return Ok(Arc::new(StructArray::from(Vec::new()))); + } + + let children_arrays = self + .children + .iter_mut() + .map(|reader| reader.consume_batch()) + .collect::>>()?; + + let DataType::Struct(fields) = &self.data_type else { + return Err(general_err!( + "Internal: StructArrayReader must have struct data type, got {:?}", + self.data_type + )); + }; + let fields = fields.clone(); + + let item_count = children_arrays.first().map(|a| a.len()).unwrap_or(0); + + if !children_arrays.windows(2).all(|w| w[0].len() == w[1].len()) { + return Err(general_err!("Not all children array length are the same!")); + } + + // Build struct null bitmap if the struct is nullable. + // We iterate def/rep levels and select entries that correspond to + // struct rows: skip parent-level padding (d < threshold) and + // inner-list continuations (r > struct_rep_level). + let nulls = if self.nullable { + let def_levels = self.children[0] + .get_def_levels() + .ok_or_else(|| general_err!("child def levels are None"))?; + let rep_levels = self.children[0].get_rep_levels(); + + let mut bitmap = BooleanBufferBuilder::new(item_count); + build_filtered_validity_bitmap( + def_levels, + rep_levels.map(|r| (r, self.struct_rep_level)), + self.padding_threshold, + self.struct_def_level, + &mut bitmap, + ); + if bitmap.len() != item_count { + return Err(general_err!( + "Failed to decode level data for struct array: \ + expected {} validity bits but got {}", + item_count, + bitmap.len() + )); + } + Some(NullBuffer::from(bitmap)) + } else { + None + }; + + unsafe { + Ok(Arc::new(StructArray::new_unchecked_with_length( + fields, + children_arrays, + nulls, + item_count, + ))) + } + } + + fn skip_records(&mut self, num_records: usize) -> Result { + let mut skipped = None; + for child in self.children.iter_mut() { + let child_skipped = child.skip_records(num_records)?; + match skipped { + Some(expected) => { + if expected != child_skipped { + return Err(general_err!( + "StructArrayReader out of sync, expected {} skipped, got {}", + expected, + child_skipped + )); + } + } + None => skipped = Some(child_skipped), + } + } + Ok(skipped.unwrap_or(0)) + } + + fn get_def_levels(&self) -> Option<&[i16]> { + // Children definition levels should describe the same + // parent structure, so return first child's + self.children.first().and_then(|l| l.get_def_levels()) + } + + fn get_rep_levels(&self) -> Option<&[i16]> { + // Children definition levels should describe the same + // parent structure, so return first child's + self.children.first().and_then(|l| l.get_rep_levels()) + } + + fn max_def_level(&self) -> i16 { + self.struct_def_level + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::array_reader::ListArrayReader; + use crate::arrow::array_reader::test_util::make_int32_page_reader; + use arrow::buffer::Buffer; + use arrow::datatypes::Field; + use arrow_array::cast::AsArray; + use arrow_array::{Array, ListArray}; + use arrow_schema::Fields; + + #[test] + fn test_struct_array_reader() { + let array_reader_1 = + make_int32_page_reader(&[4], &[0, 1, 2, 3, 1], &[0, 1, 1, 1, 1], 3, 1, None); + + let array_reader_2 = + make_int32_page_reader(&[3], &[0, 1, 3, 1, 2], &[0, 1, 1, 1, 1], 3, 1, None); + + let struct_type = ArrowType::Struct(Fields::from(vec![ + Field::new("f1", ArrowType::Int32, true), + Field::new("f2", ArrowType::Int32, true), + ])); + + let mut struct_array_reader = StructArrayReader::new( + struct_type, + vec![array_reader_1, array_reader_2], + 1, + 1, + true, + None, + ); + + let struct_array = struct_array_reader.next_batch(5).unwrap(); + let struct_array = struct_array.as_struct(); + + assert_eq!(5, struct_array.len()); + assert_eq!( + vec![true, false, false, false, false], + (0..5) + .map(|idx| struct_array.is_null(idx)) + .collect::>() + ); + assert_eq!( + Some(vec![0, 1, 2, 3, 1].as_slice()), + struct_array_reader.get_def_levels() + ); + assert_eq!( + Some(vec![0, 1, 1, 1, 1].as_slice()), + struct_array_reader.get_rep_levels() + ); + } + + #[test] + fn test_struct_array_reader_list() { + use arrow::datatypes::Int32Type; + // [ + // {foo: [1, 2, null], + // {foo: []}, + // {foo: null}, + // null, + // ] + + let expected_l = Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2), None]), + Some(vec![]), + None, + None, + ])); + + let validity = Buffer::from([0b00000111]); + let struct_fields = vec![( + Arc::new(Field::new("foo", expected_l.data_type().clone(), true)), + expected_l.clone() as ArrayRef, + )]; + let expected = StructArray::from((struct_fields, validity)); + + let reader = make_int32_page_reader( + &[1, 2], + &[4, 4, 3, 2, 1, 0], + &[0, 1, 1, 0, 0, 0], + 4, + 1, + Some(3), + ); + + let list_reader = + ListArrayReader::::new(reader, expected_l.data_type().clone(), 3, 1, true, None); + + let mut struct_reader = StructArrayReader::new( + expected.data_type().clone(), + vec![Box::new(list_reader)], + 1, + 0, + true, + None, + ); + + let actual = struct_reader.next_batch(1024).unwrap(); + let actual = actual.as_struct(); + assert_eq!(actual, &expected) + } +} diff --git a/vendor/parquet/src/arrow/array_reader/test_util.rs b/vendor/parquet/src/arrow/array_reader/test_util.rs new file mode 100644 index 0000000..4cbf9d8 --- /dev/null +++ b/vendor/parquet/src/arrow/array_reader/test_util.rs @@ -0,0 +1,161 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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 bytes::Bytes; +use std::sync::Arc; + +use crate::arrow::array_reader::ArrayReader; +use crate::basic::{ConvertedType, Encoding, Type as PhysicalType}; +use crate::column::page::{PageIterator, PageReader}; +use crate::data_type::{ByteArray, ByteArrayType, Int32Type}; +use crate::encodings::encoding::{DictEncoder, Encoder, get_encoder}; +use crate::errors::Result; +use crate::schema::types::{ColumnDescPtr, ColumnDescriptor, ColumnPath, Type}; + +/// Returns a descriptor for a UTF-8 column +pub fn utf8_column() -> ColumnDescPtr { + let t = Type::primitive_type_builder("col", PhysicalType::BYTE_ARRAY) + .with_converted_type(ConvertedType::UTF8) + .build() + .unwrap(); + + Arc::new(ColumnDescriptor::new( + Arc::new(t), + 1, + 0, + ColumnPath::new(vec![]), + )) +} + +/// Encode `data` with the provided `encoding` +pub fn encode_byte_array(encoding: Encoding, data: &[ByteArray]) -> Bytes { + let desc = utf8_column(); + let mut encoder = get_encoder::(encoding, &desc).unwrap(); + + encoder.put(data).unwrap(); + encoder.flush_buffer().unwrap() +} + +/// Returns the encoded dictionary and value data +pub fn encode_dictionary(data: &[ByteArray]) -> (Bytes, Bytes) { + let mut dict_encoder = DictEncoder::::new(utf8_column()); + + dict_encoder.put(data).unwrap(); + let encoded_rle = dict_encoder.flush_buffer().unwrap(); + let encoded_dictionary = dict_encoder.write_dict().unwrap(); + + (encoded_dictionary, encoded_rle) +} + +/// Encodes `data` in all the possible encodings +/// +/// Returns an array of data with its associated encoding, along with an encoded dictionary +pub fn byte_array_all_encodings( + data: Vec>, +) -> (Vec<(Encoding, Bytes)>, Bytes) { + let data: Vec<_> = data.into_iter().map(Into::into).collect(); + let (encoded_dictionary, encoded_rle) = encode_dictionary(&data); + + // A column chunk with all the encodings! + let pages = vec![ + (Encoding::PLAIN, encode_byte_array(Encoding::PLAIN, &data)), + ( + Encoding::DELTA_BYTE_ARRAY, + encode_byte_array(Encoding::DELTA_BYTE_ARRAY, &data), + ), + ( + Encoding::DELTA_LENGTH_BYTE_ARRAY, + encode_byte_array(Encoding::DELTA_LENGTH_BYTE_ARRAY, &data), + ), + (Encoding::PLAIN_DICTIONARY, encoded_rle.clone()), + (Encoding::RLE_DICTIONARY, encoded_rle), + ]; + + (pages, encoded_dictionary) +} + +/// Build a real `PrimitiveArrayReader` from raw non-null values +/// and definition/repetition levels. This exercises the full production +/// `RecordReader` code path, including selective padding when +/// `padding_threshold` is set. +/// +/// `values` must contain only the non-null values (entries where +/// `def_levels[i] == max_def_level`), in order, as Parquet encodes them. +/// +/// `padding_threshold` controls null filtering: `None` for full padding +/// (top-level columns), `Some(threshold)` for selective padding (list +/// children — typically the parent list's def_level). +pub fn make_int32_page_reader( + values: &[i32], + def_levels: &[i16], + rep_levels: &[i16], + max_def_level: i16, + max_rep_level: i16, + padding_threshold: Option, +) -> Box { + use crate::arrow::array_reader::PrimitiveArrayReader; + use crate::arrow::arrow_reader::DEFAULT_BATCH_SIZE; + use crate::util::InMemoryPageIterator; + use crate::util::test_common::page_util::{DataPageBuilder, DataPageBuilderImpl}; + + let leaf_type = Type::primitive_type_builder("leaf", PhysicalType::INT32) + .build() + .unwrap(); + + let desc = Arc::new(ColumnDescriptor::new( + Arc::new(leaf_type), + max_def_level, + max_rep_level, + ColumnPath::new(vec![]), + )); + + let mut pb = DataPageBuilderImpl::new(desc.clone(), def_levels.len() as u32, true); + if max_rep_level > 0 { + pb.add_rep_levels(max_rep_level, rep_levels); + } + if max_def_level > 0 { + pb.add_def_levels(max_def_level, def_levels); + } + pb.add_values::(Encoding::PLAIN, values); + + let pages = vec![vec![pb.consume()]]; + let page_iter = InMemoryPageIterator::new(pages); + Box::new( + PrimitiveArrayReader::::new( + Box::new(page_iter), + desc, + None, + DEFAULT_BATCH_SIZE, + padding_threshold, + ) + .unwrap(), + ) +} + +/// Iterator for testing reading empty columns +#[derive(Default)] +pub struct EmptyPageIterator {} + +impl Iterator for EmptyPageIterator { + type Item = Result>; + + fn next(&mut self) -> Option { + None + } +} + +impl PageIterator for EmptyPageIterator {} diff --git a/vendor/parquet/src/arrow/arrow_reader/filter.rs b/vendor/parquet/src/arrow/arrow_reader/filter.rs new file mode 100644 index 0000000..3fd5e1d --- /dev/null +++ b/vendor/parquet/src/arrow/arrow_reader/filter.rs @@ -0,0 +1,201 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::ProjectionMask; +use arrow_array::{BooleanArray, RecordBatch}; +use arrow_schema::ArrowError; +use std::fmt::{Debug, Formatter}; + +/// A predicate operating on [`RecordBatch`] +/// +/// See also: +/// * [`RowFilter`] for more information on applying filters during the +/// Parquet decoding process. +/// * [`ArrowPredicateFn`] for a concrete implementation based on a function +pub trait ArrowPredicate: Send + 'static { + /// Returns the [`ProjectionMask`] that describes the columns required + /// to evaluate this predicate. + /// + /// All projected columns will be provided in the `batch` passed to + /// [`evaluate`](Self::evaluate). The projection mask should be as small as + /// possible because any columns needed for the overall projection mask are + /// decoded again after a predicate is applied. + fn projection(&self) -> &ProjectionMask; + + /// Evaluate this predicate for the given [`RecordBatch`] containing the columns + /// identified by [`Self::projection`] + /// + /// Must return a [`BooleanArray`] that has the same length as the input + /// `batch` where each row indicates whether the row should be returned: + /// * `true`:the row should be returned + /// * `false` or `null`: the row should not be returned + fn evaluate(&mut self, batch: RecordBatch) -> Result; +} + +/// An [`ArrowPredicate`] created from an [`FnMut`] and a [`ProjectionMask`] +/// +/// See [`RowFilter`] for more information on applying filters during the +/// Parquet decoding process. +/// +/// The function is passed `RecordBatch`es with only the columns specified in +/// the [`ProjectionMask`]. +/// +/// The function must return a [`BooleanArray`] that has the same length as the +/// input `batch` where each row indicates whether the row should be returned: +/// * `true`: the row should be returned +/// * `false` or `null`: the row should not be returned +/// +/// # Example: +/// +/// Given an input schema: `"a:int64", "b:int64"`, you can create a predicate that +/// evaluates `b > 0` like this: +/// +/// ``` +/// # use std::sync::Arc; +/// # use arrow::compute::kernels::cmp::gt; +/// # use arrow_array::{BooleanArray, Int64Array, RecordBatch}; +/// # use arrow_array::cast::AsArray; +/// # use arrow_array::types::Int64Type; +/// # use parquet::arrow::arrow_reader::ArrowPredicateFn; +/// # use parquet::arrow::ProjectionMask; +/// # use parquet::schema::types::{SchemaDescriptor, Type}; +/// # use parquet::basic; // note there are two `Type`s that are different +/// # // Schema for a table with one columns: "a" (int64) and "b" (int64) +/// # let descriptor = SchemaDescriptor::new( +/// # Arc::new( +/// # Type::group_type_builder("my_schema") +/// # .with_fields(vec![ +/// # Arc::new( +/// # Type::primitive_type_builder("a", basic::Type::INT64) +/// # .build().unwrap() +/// # ), +/// # Arc::new( +/// # Type::primitive_type_builder("b", basic::Type::INT64) +/// # .build().unwrap() +/// # ), +/// # ]) +/// # .build().unwrap() +/// # ) +/// # ); +/// // Create a mask for selecting only the second column "b" (index 1) +/// let projection_mask = ProjectionMask::leaves(&descriptor, [1]); +/// // Closure that evaluates "b > 0" +/// let predicate = |batch: RecordBatch| { +/// let scalar_0 = Int64Array::new_scalar(0); +/// let column = batch.column(0).as_primitive::(); +/// // call the gt kernel to compute `>` which returns a BooleanArray +/// gt(column, &scalar_0) +/// }; +/// // Create ArrowPredicateFn that can be passed to RowFilter +/// let arrow_predicate = ArrowPredicateFn::new(projection_mask, predicate); +/// ``` +pub struct ArrowPredicateFn { + f: F, + projection: ProjectionMask, +} + +impl ArrowPredicateFn +where + F: FnMut(RecordBatch) -> Result + Send + 'static, +{ + /// Create a new [`ArrowPredicateFn`] that invokes `f` on the columns + /// specified in `projection`. + pub fn new(projection: ProjectionMask, f: F) -> Self { + Self { f, projection } + } +} + +impl ArrowPredicate for ArrowPredicateFn +where + F: FnMut(RecordBatch) -> Result + Send + 'static, +{ + fn projection(&self) -> &ProjectionMask { + &self.projection + } + + fn evaluate(&mut self, batch: RecordBatch) -> Result { + (self.f)(batch) + } +} + +/// Filter applied *during* the parquet read process +/// +/// See example on [`ArrowReaderBuilder::with_row_filter`] +/// +/// [`RowFilter`] applies predicates in order, after decoding only the columns +/// required. As predicates eliminate rows, fewer rows from subsequent columns +/// may be required, thus potentially reducing IO and decode. This process is +/// also known as *push down* filtering and *late materialization*. +/// +/// A `RowFilter` consists of a list of [`ArrowPredicate`]s. Only the rows for which +/// all the predicates evaluate to `true` will be returned. +/// Any [`RowSelection`] provided to the reader will be applied prior +/// to the first predicate, and each predicate in turn will then be used to compute +/// a more refined [`RowSelection`] used when evaluating the subsequent predicates. +/// +/// Once all predicates have been evaluated, the final [`RowSelection`] is applied +/// to the top-level [`ProjectionMask`] to produce the final output [`RecordBatch`]. +/// +/// 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`] +/// * IO will be deferred until needed by a [`ProjectionMask`] +/// +/// As such there is a trade-off between a single large predicate, or multiple predicates, +/// that will depend on the shape of the data. Whilst multiple smaller predicates may +/// minimise the amount of data scanned/decoded, it may not be faster overall. +/// +/// For example, if a predicate that needs a single column of data filters out all but +/// 1% of the rows, applying it as one of the early `ArrowPredicateFn` will likely significantly +/// improve performance. +/// +/// As a counter example, if a predicate needs several columns of data to evaluate but +/// leaves 99% of the rows, it may be better to not filter the data from parquet and +/// apply the filter after the RecordBatch has been fully decoded. +/// +/// Additionally, even if a predicate eliminates a moderate number of rows, it may still be faster +/// to filter the data after the RecordBatch has been fully decoded, if the eliminated rows are +/// not contiguous. +/// +/// [`RowSelection`]: crate::arrow::arrow_reader::RowSelection +/// [`ArrowReaderBuilder::with_row_filter`]: crate::arrow::arrow_reader::ArrowReaderBuilder::with_row_filter +pub struct RowFilter { + /// A list of [`ArrowPredicate`] + pub(crate) predicates: Vec>, +} + +impl Debug for RowFilter { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "RowFilter {{ {} predicates: }}", self.predicates.len()) + } +} + +impl RowFilter { + /// Create a new [`RowFilter`] from an array of [`ArrowPredicate`] + pub fn new(predicates: Vec>) -> Self { + Self { predicates } + } + /// Returns the inner predicates + pub fn predicates(&self) -> &Vec> { + &self.predicates + } + /// Returns the inner predicates, consuming self + pub fn into_predicates(self) -> Vec> { + self.predicates + } +} diff --git a/vendor/parquet/src/arrow/arrow_reader/metrics.rs b/vendor/parquet/src/arrow/arrow_reader/metrics.rs new file mode 100644 index 0000000..b36d795 --- /dev/null +++ b/vendor/parquet/src/arrow/arrow_reader/metrics.rs @@ -0,0 +1,135 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [ArrowReaderMetrics] for collecting metrics about the Arrow reader + +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; + +/// This enum represents the state of Arrow reader metrics collection. +/// +/// The inner metrics are stored in an `Arc` +/// so cloning the `ArrowReaderMetrics` enum will not clone the inner metrics. +/// +/// To access metrics, create an `ArrowReaderMetrics` via [`ArrowReaderMetrics::enabled()`] +/// and configure the `ArrowReaderBuilder` with a clone. +#[derive(Debug, Clone)] +pub enum ArrowReaderMetrics { + /// Metrics are not collected (default) + Disabled, + /// Metrics are collected and stored in an `Arc`. + /// + /// Create this via [`ArrowReaderMetrics::enabled()`]. + Enabled(Arc), +} + +impl ArrowReaderMetrics { + /// Creates a new instance of [`ArrowReaderMetrics::Disabled`] + pub fn disabled() -> Self { + Self::Disabled + } + + /// Creates a new instance of [`ArrowReaderMetrics::Enabled`] + pub fn enabled() -> Self { + Self::Enabled(Arc::new(ArrowReaderMetricsInner::new())) + } + + /// Predicate Cache: number of records read directly from the inner reader + /// + /// This is the total number of records read from the inner reader (that is + /// actually decoding). It measures the amount of work that could not be + /// avoided with caching. + /// + /// It returns the number of records read across all columns, so if you read + /// 2 columns each with 100 records, this will return 200. + /// + /// + /// Returns None if metrics are disabled. + pub fn records_read_from_inner(&self) -> Option { + match self { + Self::Disabled => None, + Self::Enabled(inner) => Some( + inner + .records_read_from_inner + .load(std::sync::atomic::Ordering::Relaxed), + ), + } + } + + /// Predicate Cache: number of records read from the cache + /// + /// This is the total number of records read from the cache actually + /// decoding). It measures the amount of work that was avoided with caching. + /// + /// It returns the number of records read across all columns, so if you read + /// 2 columns each with 100 records from the cache, this will return 200. + /// + /// Returns None if metrics are disabled. + pub fn records_read_from_cache(&self) -> Option { + match self { + Self::Disabled => None, + Self::Enabled(inner) => Some( + inner + .records_read_from_cache + .load(std::sync::atomic::Ordering::Relaxed), + ), + } + } + + /// Increments the count of records read from the inner reader + pub(crate) fn increment_inner_reads(&self, count: usize) { + let Self::Enabled(inner) = self else { + return; + }; + inner + .records_read_from_inner + .fetch_add(count, std::sync::atomic::Ordering::Relaxed); + } + + /// Increments the count of records read from the cache + pub(crate) fn increment_cache_reads(&self, count: usize) { + let Self::Enabled(inner) = self else { + return; + }; + + inner + .records_read_from_cache + .fetch_add(count, std::sync::atomic::Ordering::Relaxed); + } +} + +/// Holds the actual metrics for the Arrow reader. +/// +/// Please see [`ArrowReaderMetrics`] for the public interface. +#[derive(Debug)] +pub struct ArrowReaderMetricsInner { + // Metrics for Predicate Cache + /// Total number of records read from the inner reader (uncached) + records_read_from_inner: AtomicUsize, + /// Total number of records read from previously cached pages + records_read_from_cache: AtomicUsize, +} + +impl ArrowReaderMetricsInner { + /// Creates a new instance of `ArrowReaderMetricsInner` + pub(crate) fn new() -> Self { + Self { + records_read_from_inner: AtomicUsize::new(0), + records_read_from_cache: AtomicUsize::new(0), + } + } +} diff --git a/vendor/parquet/src/arrow/arrow_reader/mod.rs b/vendor/parquet/src/arrow/arrow_reader/mod.rs new file mode 100644 index 0000000..5341739 --- /dev/null +++ b/vendor/parquet/src/arrow/arrow_reader/mod.rs @@ -0,0 +1,6040 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Contains reader which reads parquet data into arrow [`RecordBatch`] + +use arrow_array::cast::AsArray; +use arrow_array::{BooleanArray, RecordBatch, RecordBatchReader}; +use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder}; +use arrow_schema::{ArrowError, DataType as ArrowType, FieldRef, Schema, SchemaRef}; +use arrow_select::filter::filter_record_batch; +pub use filter::{ArrowPredicate, ArrowPredicateFn, RowFilter}; +use selection::MaskCursor; +pub use selection::{ + MaskRunIter, RowSelection, RowSelectionCursor, RowSelectionPolicy, RowSelector, +}; +use std::fmt::{Debug, Formatter}; +use std::sync::Arc; + +pub use crate::arrow::array_reader::RowGroups; +use crate::arrow::array_reader::{ArrayReader, ArrayReaderBuilder}; +use crate::arrow::schema::{ + ParquetField, parquet_to_arrow_schema_and_fields, virtual_type::is_virtual_column, +}; +use crate::arrow::{FieldLevels, ProjectionMask, parquet_to_arrow_field_levels_with_virtual}; +use crate::basic::{BloomFilterAlgorithm, BloomFilterCompression, BloomFilterHash}; +use crate::bloom_filter::{ + SBBF_HEADER_SIZE_ESTIMATE, Sbbf, chunk_read_bloom_filter_header_and_offset, +}; +use crate::column::page::{PageIterator, PageReader}; +#[cfg(feature = "encryption")] +use crate::encryption::decrypt::FileDecryptionProperties; +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::{ + PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions, ParquetMetaDataReader, + ParquetStatisticsPolicy, RowGroupMetaData, +}; +use crate::file::reader::{ChunkReader, SerializedPageReader}; +use crate::schema::types::SchemaDescriptor; + +use crate::arrow::arrow_reader::metrics::ArrowReaderMetrics; +// Exposed so integration tests and benchmarks can temporarily override the threshold. +pub use read_plan::{PredicateOptions, ReadPlan, ReadPlanBuilder}; + +mod filter; +pub mod metrics; +mod read_plan; +pub(crate) mod selection; +pub mod statistics; + +/// Default batch size for reading parquet files +pub const DEFAULT_BATCH_SIZE: usize = 1024; + +/// Builder for constructing Parquet readers that decode into [Apache Arrow] +/// arrays. +/// +/// Most users should use one of the following specializations: +/// +/// * synchronous API: [`ParquetRecordBatchReaderBuilder`] +/// * `async` API: [`ParquetRecordBatchStreamBuilder`] +/// * decoder API: [`ParquetPushDecoderBuilder`] +/// +/// # Features +/// * Projection pushdown: [`Self::with_projection`] +/// * Cached metadata: [`ArrowReaderMetadata::load`] +/// * Offset skipping: [`Self::with_offset`] and [`Self::with_limit`] +/// * Row group filtering: [`Self::with_row_groups`] +/// * Range filtering: [`Self::with_row_selection`] +/// * Row level filtering: [`Self::with_row_filter`] +/// +/// # Implementing Predicate Pushdown +/// +/// [`Self::with_row_filter`] permits filter evaluation *during* the decoding +/// process, which is efficient and allows the most low level optimizations. +/// +/// However, most Parquet based systems will apply filters at many steps prior +/// to decoding such as pruning files, row groups and data pages. This crate +/// provides the low level APIs needed to implement such filtering, but does not +/// include any logic to actually evaluate predicates. For example: +/// +/// * [`Self::with_row_groups`] for Row Group pruning +/// * [`Self::with_row_selection`] for data page pruning +/// * [`StatisticsConverter`] to convert Parquet statistics to Arrow arrays +/// +/// The rationale for this design is that implementing predicate pushdown is a +/// complex topic and varies significantly from system to system. For example +/// +/// 1. Predicates supported (do you support predicates like prefix matching, user defined functions, etc) +/// 2. Evaluating predicates on multiple files (with potentially different but compatible schemas) +/// 3. Evaluating predicates using information from an external metadata catalog (e.g. Apache Iceberg or similar) +/// 4. Interleaving fetching metadata, evaluating predicates, and decoding files +/// +/// You can read more about this design in the [Querying Parquet with +/// Millisecond Latency] Arrow blog post. +/// +/// [`ParquetRecordBatchStreamBuilder`]: crate::arrow::async_reader::ParquetRecordBatchStreamBuilder +/// [`ParquetPushDecoderBuilder`]: crate::arrow::push_decoder::ParquetPushDecoderBuilder +/// [Apache Arrow]: https://arrow.apache.org/ +/// [`StatisticsConverter`]: statistics::StatisticsConverter +/// [Querying Parquet with Millisecond Latency]: https://arrow.apache.org/blog/2022/12/26/querying-parquet-with-millisecond-latency/ +pub struct ArrowReaderBuilder { + /// The "input" to read parquet data from. + /// + /// Note in the case of the [`ParquetPushDecoderBuilder`] there is no + /// underlying reader; the input is instead [`PushDecoderInput`], the buffer that + /// caller-pushed bytes accumulate in. + /// + /// [`ParquetPushDecoderBuilder`]: crate::arrow::push_decoder::ParquetPushDecoderBuilder + /// [`PushDecoderInput`]: crate::arrow::push_decoder::PushDecoderInput + pub(crate) input: T, + + pub(crate) metadata: Arc, + + pub(crate) schema: SchemaRef, + + pub(crate) fields: Option>, + + pub(crate) batch_size: usize, + + pub(crate) row_groups: Option>, + + pub(crate) projection: ProjectionMask, + + pub(crate) filter: Option, + + pub(crate) selection: Option, + + pub(crate) row_selection_policy: RowSelectionPolicy, + + pub(crate) limit: Option, + + pub(crate) offset: Option, + + pub(crate) metrics: ArrowReaderMetrics, + + pub(crate) max_predicate_cache_size: usize, +} + +impl Debug for ArrowReaderBuilder { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ArrowReaderBuilder") + .field("input", &self.input) + .field("metadata", &self.metadata) + .field("schema", &self.schema) + .field("fields", &self.fields) + .field("batch_size", &self.batch_size) + .field("row_groups", &self.row_groups) + .field("projection", &self.projection) + .field("filter", &self.filter) + .field("selection", &self.selection) + .field("row_selection_policy", &self.row_selection_policy) + .field("limit", &self.limit) + .field("offset", &self.offset) + .field("metrics", &self.metrics) + .finish() + } +} + +impl ArrowReaderBuilder { + pub(crate) fn new_builder(input: T, metadata: ArrowReaderMetadata) -> Self { + Self { + input, + metadata: metadata.metadata, + schema: metadata.schema, + fields: metadata.fields, + batch_size: DEFAULT_BATCH_SIZE, + row_groups: None, + projection: ProjectionMask::all(), + filter: None, + selection: None, + row_selection_policy: RowSelectionPolicy::default(), + limit: None, + offset: None, + metrics: ArrowReaderMetrics::Disabled, + max_predicate_cache_size: 100 * 1024 * 1024, // 100MB default cache size + } + } + + /// Returns a reference to the [`ParquetMetaData`] for this parquet file + pub fn metadata(&self) -> &Arc { + &self.metadata + } + + /// Returns the parquet [`SchemaDescriptor`] for this parquet file + pub fn parquet_schema(&self) -> &SchemaDescriptor { + self.metadata.file_metadata().schema_descr() + } + + /// Returns the arrow [`SchemaRef`] for this parquet file + pub fn schema(&self) -> &SchemaRef { + &self.schema + } + + /// Set the size of [`RecordBatch`] to produce. Defaults to [`DEFAULT_BATCH_SIZE`]. + /// + /// This may be used as a hint for internal allocations, but does not + /// guarantee exact internal buffer capacities. + /// + /// If `batch_size` is more than the file row count, use the file row count. + pub fn with_batch_size(self, batch_size: usize) -> Self { + // Try to avoid allocate large buffer + let batch_size = batch_size.min(self.metadata.file_metadata().num_rows() as usize); + Self { batch_size, ..self } + } + + /// Only read data from the provided row group indexes + /// + /// This is also called row group filtering + pub fn with_row_groups(self, row_groups: Vec) -> Self { + Self { + row_groups: Some(row_groups), + ..self + } + } + + /// Only read data from the provided column indexes + pub fn with_projection(self, mask: ProjectionMask) -> Self { + Self { + projection: mask, + ..self + } + } + + /// Configure how row selections should be materialised during execution + /// + /// See [`RowSelectionPolicy`] for more details + pub fn with_row_selection_policy(self, policy: RowSelectionPolicy) -> Self { + Self { + row_selection_policy: policy, + ..self + } + } + + /// Provide a [`RowSelection`] to filter out rows, and avoid fetching their + /// data into memory. + /// + /// This feature is used to restrict which rows are decoded within row + /// groups, skipping ranges of rows that are not needed. Such selections + /// could be determined by evaluating predicates against the parquet page + /// [`Index`] or some other external information available to a query + /// engine. + /// + /// # Notes + /// + /// Row group filtering (see [`Self::with_row_groups`]) is applied prior to + /// applying the row selection, and therefore rows from skipped row groups + /// should not be included in the [`RowSelection`] (see example below) + /// + /// It is recommended to enable writing the page index if using this + /// functionality, to allow more efficient skipping over data pages. See + /// [`ArrowReaderOptions::with_page_index`]. + /// + /// # Example + /// + /// Given a parquet file with 4 row groups, and a row group filter of `[0, + /// 2, 3]`, in order to scan rows 50-100 in row group 2 and rows 200-300 in + /// row group 3: + /// + /// ```text + /// Row Group 0, 1000 rows (selected) + /// Row Group 1, 1000 rows (skipped) + /// Row Group 2, 1000 rows (selected, but want to only scan rows 50-100) + /// Row Group 3, 1000 rows (selected, but want to only scan rows 200-300) + /// ``` + /// + /// You could pass the following [`RowSelection`]: + /// + /// ```text + /// Select 1000 (scan all rows in row group 0) + /// Skip 50 (skip the first 50 rows in row group 2) + /// Select 50 (scan rows 50-100 in row group 2) + /// Skip 900 (skip the remaining rows in row group 2) + /// Skip 200 (skip the first 200 rows in row group 3) + /// Select 100 (scan rows 200-300 in row group 3) + /// Skip 700 (skip the remaining rows in row group 3) + /// ``` + /// Note there is no entry for the (entirely) skipped row group 1. + /// + /// Note you can represent the same selection with fewer entries. Instead of + /// + /// ```text + /// Skip 900 (skip the remaining rows in row group 2) + /// Skip 200 (skip the first 200 rows in row group 3) + /// ``` + /// + /// you could use + /// + /// ```text + /// Skip 1100 (skip the remaining 900 rows in row group 2 and the first 200 rows in row group 3) + /// ``` + /// + /// [`Index`]: crate::file::page_index::column_index::ColumnIndexMetaData + pub fn with_row_selection(self, selection: RowSelection) -> Self { + Self { + selection: Some(selection), + ..self + } + } + + /// Provide a [`RowFilter`] to skip decoding rows + /// + /// Row filters are applied after row group selection and row selection + /// + /// It is recommended to enable reading the page index if using this functionality, to allow + /// more efficient skipping over data pages. See [`ArrowReaderOptions::with_page_index`]. + /// + /// See the [blog post on late materialization] for a more technical explanation. + /// + /// [blog post on late materialization]: https://arrow.apache.org/blog/2025/12/11/parquet-late-materialization-deep-dive + /// + /// # Example + /// ```rust + /// # use std::fs::File; + /// # use arrow_array::Int32Array; + /// # use parquet::arrow::ProjectionMask; + /// # use parquet::arrow::arrow_reader::{ArrowPredicateFn, ParquetRecordBatchReaderBuilder, RowFilter}; + /// # fn main() -> Result<(), parquet::errors::ParquetError> { + /// # let testdata = arrow::util::test_util::parquet_test_data(); + /// # let path = format!("{testdata}/alltypes_plain.parquet"); + /// # let file = File::open(&path)?; + /// let builder = ParquetRecordBatchReaderBuilder::try_new(file)?; + /// let schema_desc = builder.metadata().file_metadata().schema_descr_ptr(); + /// // Create predicate that evaluates `int_col != 1`. + /// // `int_col` column has index 4 (zero based) in the schema + /// let projection = ProjectionMask::leaves(&schema_desc, [4]); + /// // Only the projection columns are passed to the predicate so + /// // int_col is column 0 in the predicate + /// let predicate = ArrowPredicateFn::new(projection, |batch| { + /// let int_col = batch.column(0); + /// arrow::compute::kernels::cmp::neq(int_col, &Int32Array::new_scalar(1)) + /// }); + /// let row_filter = RowFilter::new(vec![Box::new(predicate)]); + /// // The filter will be invoked during the reading process + /// let reader = builder.with_row_filter(row_filter).build()?; + /// # for b in reader { let _ = b?; } + /// # Ok(()) + /// # } + /// ``` + pub fn with_row_filter(self, filter: RowFilter) -> Self { + Self { + filter: Some(filter), + ..self + } + } + + /// Provide a limit to the number of rows to be read + /// + /// The limit will be applied after any [`Self::with_row_selection`] and [`Self::with_row_filter`] + /// allowing it to limit the final set of rows decoded after any pushed down predicates + /// + /// It is recommended to enable reading the page index if using this functionality, to allow + /// more efficient skipping over data pages. See [`ArrowReaderOptions::with_page_index`] + pub fn with_limit(self, limit: usize) -> Self { + Self { + limit: Some(limit), + ..self + } + } + + /// Provide an offset to skip over the given number of rows + /// + /// The offset will be applied after any [`Self::with_row_selection`] and [`Self::with_row_filter`] + /// allowing it to skip rows after any pushed down predicates + /// + /// It is recommended to enable reading the page index if using this functionality, to allow + /// more efficient skipping over data pages. See [`ArrowReaderOptions::with_page_index`] + pub fn with_offset(self, offset: usize) -> Self { + Self { + offset: Some(offset), + ..self + } + } + + /// Specify metrics collection during reading + /// + /// To access the metrics, create an [`ArrowReaderMetrics`] and pass a + /// clone of the provided metrics to the builder. + /// + /// For example: + /// + /// ```rust + /// # use std::sync::Arc; + /// # use bytes::Bytes; + /// # use arrow_array::{Int32Array, RecordBatch}; + /// # use arrow_schema::{DataType, Field, Schema}; + /// # use parquet::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder}; + /// use parquet::arrow::arrow_reader::metrics::ArrowReaderMetrics; + /// # use parquet::arrow::ArrowWriter; + /// # let mut file: Vec = Vec::with_capacity(1024); + /// # let schema = Arc::new(Schema::new(vec![Field::new("i32", DataType::Int32, false)])); + /// # let mut writer = ArrowWriter::try_new(&mut file, schema.clone(), None).unwrap(); + /// # let batch = RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1, 2, 3]))]).unwrap(); + /// # writer.write(&batch).unwrap(); + /// # writer.close().unwrap(); + /// # let file = Bytes::from(file); + /// // Create metrics object to pass into the reader + /// let metrics = ArrowReaderMetrics::enabled(); + /// let reader = ParquetRecordBatchReaderBuilder::try_new(file).unwrap() + /// // Configure the builder to use the metrics by passing a clone + /// .with_metrics(metrics.clone()) + /// // Build the reader + /// .build().unwrap(); + /// // .. read data from the reader .. + /// + /// // check the metrics + /// assert!(metrics.records_read_from_inner().is_some()); + /// ``` + pub fn with_metrics(self, metrics: ArrowReaderMetrics) -> Self { + Self { metrics, ..self } + } + + /// Set the maximum size (per row group) of the predicate cache in bytes for + /// the async decoder. + /// + /// Defaults to 100MB (across all columns). Set to `usize::MAX` to use + /// unlimited cache size. + /// + /// This cache is used to store decoded arrays that are used in + /// predicate evaluation ([`Self::with_row_filter`]). + /// + /// This cache is only used for the "async" decoder, [`ParquetRecordBatchStream`]. See + /// [this ticket] for more details and alternatives. + /// + /// [`ParquetRecordBatchStream`]: https://docs.rs/parquet/latest/parquet/arrow/async_reader/struct.ParquetRecordBatchStream.html + /// [this ticket]: https://github.com/apache/arrow-rs/issues/8000 + pub fn with_max_predicate_cache_size(self, max_predicate_cache_size: usize) -> Self { + Self { + max_predicate_cache_size, + ..self + } + } +} + +/// Options that control how [`ParquetMetaData`] is read when constructing +/// an Arrow reader. +/// +/// To use these options, pass them to one of the following methods: +/// * [`ParquetRecordBatchReaderBuilder::try_new_with_options`] +/// * [`ParquetRecordBatchStreamBuilder::new_with_options`] +/// +/// For fine-grained control over metadata loading, use +/// [`ArrowReaderMetadata::load`] to load metadata with these options, +/// +/// See [`ArrowReaderBuilder`] for how to configure how the column data +/// is then read from the file, including projection and filter pushdown +/// +/// [`ParquetRecordBatchStreamBuilder::new_with_options`]: crate::arrow::async_reader::ParquetRecordBatchStreamBuilder::new_with_options +#[derive(Debug, Clone, Default)] +pub struct ArrowReaderOptions { + /// Should the reader strip any user defined metadata from the Arrow schema + skip_arrow_metadata: bool, + /// If provided, used as the schema hint when determining the Arrow schema, + /// otherwise the schema hint is read from the [ARROW_SCHEMA_META_KEY] + /// + /// [ARROW_SCHEMA_META_KEY]: crate::arrow::ARROW_SCHEMA_META_KEY + supplied_schema: Option, + + pub(crate) column_index: PageIndexPolicy, + pub(crate) offset_index: PageIndexPolicy, + + /// Options to control reading of Parquet metadata + metadata_options: ParquetMetaDataOptions, + /// If encryption is enabled, the file decryption properties can be provided + #[cfg(feature = "encryption")] + pub(crate) file_decryption_properties: Option>, + + virtual_columns: Vec, +} + +impl ArrowReaderOptions { + /// Create a new [`ArrowReaderOptions`] with the default settings + pub fn new() -> Self { + Self::default() + } + + /// Skip decoding the embedded arrow metadata (defaults to `false`) + /// + /// Parquet files generated by some writers may contain embedded arrow + /// schema and metadata. + /// This may not be correct or compatible with your system, + /// for example, see [ARROW-16184](https://issues.apache.org/jira/browse/ARROW-16184) + pub fn with_skip_arrow_metadata(self, skip_arrow_metadata: bool) -> Self { + Self { + skip_arrow_metadata, + ..self + } + } + + /// Provide a schema hint to use when reading the Parquet file. + /// + /// If provided, this schema takes precedence over any arrow schema embedded + /// in the metadata (see the [`arrow`] documentation for more details). + /// + /// If the provided schema is not compatible with the data stored in the + /// parquet file schema, an error will be returned when constructing the + /// builder. + /// + /// This option is only required if you want to explicitly control the + /// conversion of Parquet types to Arrow types, such as casting a column to + /// a different type. For example, if you wanted to read an Int64 in + /// a Parquet file to a [`TimestampMicrosecondArray`] in the Arrow schema. + /// + /// [`arrow`]: crate::arrow + /// [`TimestampMicrosecondArray`]: arrow_array::TimestampMicrosecondArray + /// + /// # Notes + /// + /// The provided schema must have the same number of columns as the parquet schema and + /// the column names must be the same. + /// + /// # Example + /// ``` + /// # use std::sync::Arc; + /// # use bytes::Bytes; + /// # use arrow_array::{ArrayRef, Int32Array, RecordBatch}; + /// # use arrow_schema::{DataType, Field, Schema, TimeUnit}; + /// # use parquet::arrow::arrow_reader::{ArrowReaderOptions, ParquetRecordBatchReaderBuilder}; + /// # use parquet::arrow::ArrowWriter; + /// // Write data - schema is inferred from the data to be Int32 + /// let mut file = Vec::new(); + /// let batch = RecordBatch::try_from_iter(vec![ + /// ("col_1", Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef), + /// ]).unwrap(); + /// let mut writer = ArrowWriter::try_new(&mut file, batch.schema(), None).unwrap(); + /// writer.write(&batch).unwrap(); + /// writer.close().unwrap(); + /// let file = Bytes::from(file); + /// + /// // Read the file back. + /// // Supply a schema that interprets the Int32 column as a Timestamp. + /// let supplied_schema = Arc::new(Schema::new(vec![ + /// Field::new("col_1", DataType::Timestamp(TimeUnit::Nanosecond, None), false) + /// ])); + /// let options = ArrowReaderOptions::new().with_schema(supplied_schema.clone()); + /// let mut builder = ParquetRecordBatchReaderBuilder::try_new_with_options( + /// file.clone(), + /// options + /// ).expect("Error if the schema is not compatible with the parquet file schema."); + /// + /// // Create the reader and read the data using the supplied schema. + /// let mut reader = builder.build().unwrap(); + /// let _batch = reader.next().unwrap().unwrap(); + /// ``` + /// + /// # Example: Preserving Dictionary Encoding + /// + /// By default, Parquet string columns are read as `Utf8Array` (or `LargeUtf8Array`), + /// even if the underlying Parquet data uses dictionary encoding. You can preserve + /// the dictionary encoding by specifying a `Dictionary` type in the schema hint: + /// + /// ``` + /// # use std::sync::Arc; + /// # use bytes::Bytes; + /// # use arrow_array::{ArrayRef, RecordBatch, StringArray}; + /// # use arrow_schema::{DataType, Field, Schema}; + /// # use parquet::arrow::arrow_reader::{ArrowReaderOptions, ParquetRecordBatchReaderBuilder}; + /// # use parquet::arrow::ArrowWriter; + /// // Write a Parquet file with string data + /// let mut file = Vec::new(); + /// let schema = Arc::new(Schema::new(vec![ + /// Field::new("city", DataType::Utf8, false) + /// ])); + /// let cities = StringArray::from(vec!["Berlin", "Berlin", "Paris", "Berlin", "Paris"]); + /// let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(cities)]).unwrap(); + /// + /// let mut writer = ArrowWriter::try_new(&mut file, batch.schema(), None).unwrap(); + /// writer.write(&batch).unwrap(); + /// writer.close().unwrap(); + /// let file = Bytes::from(file); + /// + /// // Read the file back, requesting dictionary encoding preservation + /// let dict_schema = Arc::new(Schema::new(vec![ + /// Field::new("city", DataType::Dictionary( + /// Box::new(DataType::Int32), + /// Box::new(DataType::Utf8) + /// ), false) + /// ])); + /// let options = ArrowReaderOptions::new().with_schema(dict_schema); + /// let builder = ParquetRecordBatchReaderBuilder::try_new_with_options( + /// file.clone(), + /// options + /// ).unwrap(); + /// + /// let mut reader = builder.build().unwrap(); + /// let batch = reader.next().unwrap().unwrap(); + /// + /// // The column is now a DictionaryArray + /// assert!(matches!( + /// batch.column(0).data_type(), + /// DataType::Dictionary(_, _) + /// )); + /// ``` + /// + /// **Note**: Dictionary encoding preservation works best when: + /// 1. The original column was dictionary encoded (the default for string columns) + /// 2. There are a small number of distinct values + pub fn with_schema(self, schema: SchemaRef) -> Self { + Self { + supplied_schema: Some(schema), + skip_arrow_metadata: true, + ..self + } + } + + #[deprecated(since = "57.2.0", note = "Use `with_page_index_policy` instead")] + /// Enable reading the [`PageIndex`] from the metadata, if present (defaults to `false`) + /// + /// The `PageIndex` can be used to push down predicates to the parquet scan, + /// potentially eliminating unnecessary IO, by some query engines. + /// + /// If this is enabled, [`ParquetMetaData::column_index`] and + /// [`ParquetMetaData::offset_index`] will be populated if the corresponding + /// information is present in the file. + /// + /// [`PageIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + /// [`ParquetMetaData::column_index`]: crate::file::metadata::ParquetMetaData::column_index + /// [`ParquetMetaData::offset_index`]: crate::file::metadata::ParquetMetaData::offset_index + pub fn with_page_index(self, page_index: bool) -> Self { + self.with_page_index_policy(PageIndexPolicy::from(page_index)) + } + + /// Sets the [`PageIndexPolicy`] for both the column and offset indexes. + /// + /// The `PageIndex` consists of two structures: the `ColumnIndex` and `OffsetIndex`. + /// This method sets the same policy for both. For fine-grained control, use + /// [`Self::with_column_index_policy`] and [`Self::with_offset_index_policy`]. + /// + /// See [`Self::with_page_index`] for more details on page indexes. + pub fn with_page_index_policy(self, policy: PageIndexPolicy) -> Self { + self.with_column_index_policy(policy) + .with_offset_index_policy(policy) + } + + /// Sets the [`PageIndexPolicy`] for the Parquet [ColumnIndex] structure. + /// + /// The `ColumnIndex` contains min/max statistics for each page, which can be used + /// for predicate pushdown and page-level pruning. + /// + /// [ColumnIndex]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + pub fn with_column_index_policy(mut self, policy: PageIndexPolicy) -> Self { + self.column_index = policy; + self + } + + /// Sets the [`PageIndexPolicy`] for the Parquet [OffsetIndex] structure. + /// + /// The `OffsetIndex` contains the locations and sizes of each page, which enables + /// efficient page-level skipping and random access within column chunks. + /// + /// [OffsetIndex]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + pub fn with_offset_index_policy(mut self, policy: PageIndexPolicy) -> Self { + self.offset_index = policy; + self + } + + /// Provide a Parquet schema to use when decoding the metadata. The schema in the Parquet + /// footer will be skipped. + /// + /// This can be used to avoid reparsing the schema from the file when it is + /// already known. + pub fn with_parquet_schema(mut self, schema: Arc) -> Self { + self.metadata_options.set_schema(schema); + self + } + + /// Set whether to convert the [`encoding_stats`] in the Parquet `ColumnMetaData` to a bitmask + /// (defaults to `false`). + /// + /// See [`ColumnChunkMetaData::page_encoding_stats_mask`] for an explanation of why this + /// might be desirable. + /// + /// [`ColumnChunkMetaData::page_encoding_stats_mask`]: + /// crate::file::metadata::ColumnChunkMetaData::page_encoding_stats_mask + /// [`encoding_stats`]: + /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L917 + pub fn with_encoding_stats_as_mask(mut self, val: bool) -> Self { + self.metadata_options.set_encoding_stats_as_mask(val); + self + } + + /// Sets the decoding policy for [`encoding_stats`] in the Parquet `ColumnMetaData`. + /// + /// [`encoding_stats`]: + /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L917 + pub fn with_encoding_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self { + self.metadata_options.set_encoding_stats_policy(policy); + self + } + + /// Sets the decoding policy for [`statistics`] in the Parquet `ColumnMetaData`. + /// + /// [`statistics`]: + /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L912 + pub fn with_column_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self { + self.metadata_options.set_column_stats_policy(policy); + self + } + + /// Sets the decoding policy for [`size_statistics`] in the Parquet `ColumnMetaData`. + /// + /// [`size_statistics`]: + /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L936 + pub fn with_size_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self { + self.metadata_options.set_size_stats_policy(policy); + self + } + + /// Provide the file decryption properties to use when reading encrypted parquet files. + /// + /// If encryption is enabled and the file is encrypted, the `file_decryption_properties` must be provided. + #[cfg(feature = "encryption")] + pub fn with_file_decryption_properties( + self, + file_decryption_properties: Arc, + ) -> Self { + Self { + file_decryption_properties: Some(file_decryption_properties), + ..self + } + } + + /// Include virtual columns in the output. + /// + /// Virtual columns are columns that are not part of the Parquet schema, but are added to the output by the reader such as row numbers and row group indices. + /// + /// # Example + /// ``` + /// # use std::sync::Arc; + /// # use bytes::Bytes; + /// # use arrow_array::{ArrayRef, Int64Array, RecordBatch}; + /// # use arrow_schema::{DataType, Field, Schema}; + /// # use parquet::arrow::{ArrowWriter, RowNumber}; + /// # use parquet::arrow::arrow_reader::{ArrowReaderOptions, ParquetRecordBatchReaderBuilder}; + /// # + /// # fn main() -> Result<(), Box> { + /// // Create a simple record batch with some data + /// let values = Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef; + /// let batch = RecordBatch::try_from_iter(vec![("value", values)])?; + /// + /// // Write the batch to an in-memory buffer + /// let mut file = Vec::new(); + /// let mut writer = ArrowWriter::try_new( + /// &mut file, + /// batch.schema(), + /// None + /// )?; + /// writer.write(&batch)?; + /// writer.close()?; + /// let file = Bytes::from(file); + /// + /// // Create a virtual column for row numbers + /// let row_number_field = Arc::new(Field::new("row_number", DataType::Int64, false) + /// .with_extension_type(RowNumber)); + /// + /// // Configure options with virtual columns + /// let options = ArrowReaderOptions::new() + /// .with_virtual_columns(vec![row_number_field])?; + /// + /// // Create a reader with the options + /// let mut reader = ParquetRecordBatchReaderBuilder::try_new_with_options( + /// file, + /// options + /// )? + /// .build()?; + /// + /// // Read the batch - it will include both the original column and the virtual row_number column + /// let result_batch = reader.next().unwrap()?; + /// assert_eq!(result_batch.num_columns(), 2); // "value" + "row_number" + /// assert_eq!(result_batch.num_rows(), 3); + /// # + /// # Ok(()) + /// # } + /// ``` + pub fn with_virtual_columns(self, virtual_columns: Vec) -> Result { + // Validate that all fields are virtual columns + for field in &virtual_columns { + if !is_virtual_column(field) { + return Err(ParquetError::General(format!( + "Field '{}' is not a virtual column. Virtual columns must have extension type names starting with 'arrow.virtual.'", + field.name() + ))); + } + } + Ok(Self { + virtual_columns, + ..self + }) + } + + #[deprecated( + since = "57.2.0", + note = "Use `column_index_policy` or `offset_index_policy` instead" + )] + /// Returns whether page index reading is enabled. + /// + /// This returns `true` if both the column index and offset index policies are not [`PageIndexPolicy::Skip`]. + /// + /// This can be set via [`with_page_index`][Self::with_page_index] or + /// [`with_page_index_policy`][Self::with_page_index_policy]. + pub fn page_index(&self) -> bool { + self.offset_index != PageIndexPolicy::Skip && self.column_index != PageIndexPolicy::Skip + } + + /// Retrieve the currently set [`PageIndexPolicy`] for the offset index. + /// + /// This can be set via [`with_offset_index_policy`][Self::with_offset_index_policy] + /// or [`with_page_index_policy`][Self::with_page_index_policy]. + pub fn offset_index_policy(&self) -> PageIndexPolicy { + self.offset_index + } + + /// Retrieve the currently set [`PageIndexPolicy`] for the column index. + /// + /// This can be set via [`with_column_index_policy`][Self::with_column_index_policy] + /// or [`with_page_index_policy`][Self::with_page_index_policy]. + pub fn column_index_policy(&self) -> PageIndexPolicy { + self.column_index + } + + /// Retrieve the currently set metadata decoding options. + pub fn metadata_options(&self) -> &ParquetMetaDataOptions { + &self.metadata_options + } + + /// Retrieve the currently set file decryption properties. + /// + /// This can be set via + /// [`file_decryption_properties`][Self::with_file_decryption_properties]. + #[cfg(feature = "encryption")] + pub fn file_decryption_properties(&self) -> Option<&Arc> { + self.file_decryption_properties.as_ref() + } +} + +impl ParquetMetaDataReader { + /// Applies the metadata related settings from [`ArrowReaderOptions`], + /// such as the [`ParquetMetaDataOptions`], decryption properties, and + /// [`PageIndexPolicy`] to this reader. + /// + /// The page index policies are only applied if at least one of them is not + /// [`PageIndexPolicy::Skip`], so policies previously configured on this + /// reader (e.g. from a preload setting) are preserved when the options do + /// not request the page index. + /// + /// This encodes the canonical way to construct a `ParquetMetaDataReader` + /// inside `AsyncFileReader::get_metadata` (available with the `async` + /// feature), so implementations outside this crate do not need to + /// duplicate it. + pub fn with_arrow_reader_options(mut self, options: Option<&ArrowReaderOptions>) -> Self { + let Some(options) = options else { return self }; + + self = self.with_metadata_options(Some(options.metadata_options().clone())); + + #[cfg(feature = "encryption")] + { + self = self.with_decryption_properties( + options.file_decryption_properties.as_ref().map(Arc::clone), + ); + } + + if options.column_index_policy() != PageIndexPolicy::Skip + || options.offset_index_policy() != PageIndexPolicy::Skip + { + self = self + .with_column_index_policy(options.column_index_policy()) + .with_offset_index_policy(options.offset_index_policy()); + } + + self + } +} + +/// The metadata necessary to construct a [`ArrowReaderBuilder`] +/// +/// Note this structure is cheaply clone-able as it consists of several arcs. +/// +/// This structure allows +/// +/// 1. Loading metadata for a file once and then using that same metadata to +/// construct multiple separate readers, for example, to distribute readers +/// across multiple threads +/// +/// 2. Using a cached copy of the [`ParquetMetadata`] rather than reading it +/// from the file each time a reader is constructed. +/// +/// [`ParquetMetadata`]: crate::file::metadata::ParquetMetaData +#[derive(Debug, Clone)] +pub struct ArrowReaderMetadata { + /// The Parquet Metadata, if known aprior + pub(crate) metadata: Arc, + /// The Arrow Schema + pub(crate) schema: SchemaRef, + /// The Parquet schema (root field) + pub(crate) fields: Option>, +} + +impl ArrowReaderMetadata { + /// Create [`ArrowReaderMetadata`] from the provided [`ArrowReaderOptions`] + /// and [`ChunkReader`] + /// + /// See [`ParquetRecordBatchReaderBuilder::new_with_metadata`] for an + /// example of how this can be used + /// + /// # Notes + /// + /// If `options` has [`ArrowReaderOptions::with_page_index`] true, but + /// `Self::metadata` is missing the page index, this function will attempt + /// to load the page index by making an object store request. + pub fn load(reader: &T, options: ArrowReaderOptions) -> Result { + let metadata = ParquetMetaDataReader::new() + .with_column_index_policy(options.column_index) + .with_offset_index_policy(options.offset_index) + .with_metadata_options(Some(options.metadata_options.clone())); + #[cfg(feature = "encryption")] + let metadata = metadata.with_decryption_properties( + options.file_decryption_properties.as_ref().map(Arc::clone), + ); + let metadata = metadata.parse_and_finish(reader)?; + Self::try_new(Arc::new(metadata), options) + } + + /// Create a new [`ArrowReaderMetadata`] from a pre-existing + /// [`ParquetMetaData`] and [`ArrowReaderOptions`]. + /// + /// # Notes + /// + /// This function will not attempt to load the PageIndex if not present in the metadata, regardless + /// of the settings in `options`. See [`Self::load`] to load metadata including the page index if needed. + pub fn try_new(metadata: Arc, options: ArrowReaderOptions) -> Result { + match options.supplied_schema { + Some(supplied_schema) => Self::with_supplied_schema( + metadata, + supplied_schema.clone(), + &options.virtual_columns, + ), + None => { + let kv_metadata = match options.skip_arrow_metadata { + true => None, + false => metadata.file_metadata().key_value_metadata(), + }; + + let (schema, fields) = parquet_to_arrow_schema_and_fields( + metadata.file_metadata().schema_descr(), + ProjectionMask::all(), + kv_metadata, + &options.virtual_columns, + )?; + + Ok(Self { + metadata, + schema: Arc::new(schema), + fields: fields.map(Arc::new), + }) + } + } + } + + fn with_supplied_schema( + metadata: Arc, + supplied_schema: SchemaRef, + virtual_columns: &[FieldRef], + ) -> Result { + let parquet_schema = metadata.file_metadata().schema_descr(); + let field_levels = parquet_to_arrow_field_levels_with_virtual( + parquet_schema, + ProjectionMask::all(), + Some(supplied_schema.fields()), + virtual_columns, + )?; + let fields = field_levels.fields; + let inferred_len = fields.len(); + let supplied_len = supplied_schema.fields().len() + virtual_columns.len(); + // Ensure the supplied schema has the same number of columns as the parquet schema. + // parquet_to_arrow_field_levels is expected to throw an error if the schemas have + // different lengths, but we check here to be safe. + if inferred_len != supplied_len { + return Err(arrow_err!(format!( + "Incompatible supplied Arrow schema: expected {} columns received {}", + inferred_len, supplied_len + ))); + } + + let mut errors = Vec::new(); + + let field_iter = supplied_schema.fields().iter().zip(fields.iter()); + + for (field1, field2) in field_iter { + if field1.data_type() != field2.data_type() { + errors.push(format!( + "data type mismatch for field {}: requested {} but found {}", + field1.name(), + field1.data_type(), + field2.data_type() + )); + } + if field1.is_nullable() != field2.is_nullable() { + errors.push(format!( + "nullability mismatch for field {}: expected {:?} but found {:?}", + field1.name(), + field1.is_nullable(), + field2.is_nullable() + )); + } + if field1.metadata() != field2.metadata() { + errors.push(format!( + "metadata mismatch for field {}: expected {:?} but found {:?}", + field1.name(), + field1.metadata(), + field2.metadata() + )); + } + } + + if !errors.is_empty() { + let message = errors.join(", "); + return Err(ParquetError::ArrowError(format!( + "Incompatible supplied Arrow schema: {message}", + ))); + } + + Ok(Self { + metadata, + schema: supplied_schema, + fields: field_levels.levels.map(Arc::new), + }) + } + + /// Returns a reference to the [`ParquetMetaData`] for this parquet file + pub fn metadata(&self) -> &Arc { + &self.metadata + } + + /// Returns the parquet [`SchemaDescriptor`] for this parquet file + pub fn parquet_schema(&self) -> &SchemaDescriptor { + self.metadata.file_metadata().schema_descr() + } + + /// Returns the arrow [`SchemaRef`] for this parquet file + pub fn schema(&self) -> &SchemaRef { + &self.schema + } +} + +#[doc(hidden)] +// A newtype used within `ReaderOptionsBuilder` to distinguish sync readers from async +pub struct SyncReader(T); + +impl Debug for SyncReader { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_tuple("SyncReader").field(&self.0).finish() + } +} + +/// Creates [`ParquetRecordBatchReader`] for reading Parquet files into Arrow [`RecordBatch`]es +/// +/// # See Also +/// * [`crate::arrow::async_reader::ParquetRecordBatchStreamBuilder`] for an async API +/// * [`crate::arrow::push_decoder::ParquetPushDecoderBuilder`] for a SansIO decoder API +/// * [`ArrowReaderBuilder`] for additional member functions +pub type ParquetRecordBatchReaderBuilder = ArrowReaderBuilder>; + +impl ParquetRecordBatchReaderBuilder { + /// Create a new [`ParquetRecordBatchReaderBuilder`] + /// + /// ``` + /// # use std::sync::Arc; + /// # use bytes::Bytes; + /// # use arrow_array::{Int32Array, RecordBatch}; + /// # use arrow_schema::{DataType, Field, Schema}; + /// # use parquet::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder}; + /// # use parquet::arrow::ArrowWriter; + /// # let mut file: Vec = Vec::with_capacity(1024); + /// # let schema = Arc::new(Schema::new(vec![Field::new("i32", DataType::Int32, false)])); + /// # let mut writer = ArrowWriter::try_new(&mut file, schema.clone(), None).unwrap(); + /// # let batch = RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1, 2, 3]))]).unwrap(); + /// # writer.write(&batch).unwrap(); + /// # writer.close().unwrap(); + /// # let file = Bytes::from(file); + /// // Build the reader from anything that implements `ChunkReader` + /// // such as a `File`, or `Bytes` + /// let mut builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + /// // The builder has access to ParquetMetaData such + /// // as the number and layout of row groups + /// assert_eq!(builder.metadata().num_row_groups(), 1); + /// // Call build to create the reader + /// let mut reader: ParquetRecordBatchReader = builder.build().unwrap(); + /// // Read data + /// while let Some(batch) = reader.next().transpose()? { + /// println!("Read {} rows", batch.num_rows()); + /// } + /// # Ok::<(), parquet::errors::ParquetError>(()) + /// ``` + pub fn try_new(reader: T) -> Result { + Self::try_new_with_options(reader, Default::default()) + } + + /// Create a new [`ParquetRecordBatchReaderBuilder`] with [`ArrowReaderOptions`] + /// + /// Use this method if you want to control the options for reading the + /// [`ParquetMetaData`] + pub fn try_new_with_options(reader: T, options: ArrowReaderOptions) -> Result { + let metadata = ArrowReaderMetadata::load(&reader, options)?; + Ok(Self::new_with_metadata(reader, metadata)) + } + + /// Create a [`ParquetRecordBatchReaderBuilder`] from the provided [`ArrowReaderMetadata`] + /// + /// Use this method if you already have [`ParquetMetaData`] for a file. + /// This interface allows: + /// + /// 1. Loading metadata once and using it to create multiple builders with + /// potentially different settings or run on different threads + /// + /// 2. Using a cached copy of the metadata rather than re-reading it from the + /// file each time a reader is constructed. + /// + /// See the docs on [`ArrowReaderMetadata`] for more details + /// + /// # Example + /// ``` + /// # use std::fs::metadata; + /// # use std::sync::Arc; + /// # use bytes::Bytes; + /// # use arrow_array::{Int32Array, RecordBatch}; + /// # use arrow_schema::{DataType, Field, Schema}; + /// # use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder}; + /// # use parquet::arrow::ArrowWriter; + /// # + /// # let mut file: Vec = Vec::with_capacity(1024); + /// # let schema = Arc::new(Schema::new(vec![Field::new("i32", DataType::Int32, false)])); + /// # let mut writer = ArrowWriter::try_new(&mut file, schema.clone(), None).unwrap(); + /// # let batch = RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1, 2, 3]))]).unwrap(); + /// # writer.write(&batch).unwrap(); + /// # writer.close().unwrap(); + /// # let file = Bytes::from(file); + /// # + /// let metadata = ArrowReaderMetadata::load(&file, Default::default()).unwrap(); + /// let mut a = ParquetRecordBatchReaderBuilder::new_with_metadata(file.clone(), metadata.clone()).build().unwrap(); + /// let mut b = ParquetRecordBatchReaderBuilder::new_with_metadata(file, metadata).build().unwrap(); + /// + /// // Should be able to read from both in parallel + /// assert_eq!(a.next().unwrap().unwrap(), b.next().unwrap().unwrap()); + /// ``` + pub fn new_with_metadata(input: T, metadata: ArrowReaderMetadata) -> Self { + Self::new_builder(SyncReader(input), metadata) + } + + /// Read bloom filter for a column in a row group + /// + /// Returns `None` if the column does not have a bloom filter + /// + /// We should call this function after other forms pruning, such as projection and predicate pushdown. + pub fn get_row_group_column_bloom_filter( + &self, + row_group_idx: usize, + column_idx: usize, + ) -> Result> { + let metadata = self.metadata.row_group(row_group_idx); + let column_metadata = metadata.column(column_idx); + + let offset: u64 = if let Some(offset) = column_metadata.bloom_filter_offset() { + offset + .try_into() + .map_err(|_| ParquetError::General("Bloom filter offset is invalid".to_string()))? + } else { + return Ok(None); + }; + + let buffer = match column_metadata.bloom_filter_length() { + Some(length) => self.input.0.get_bytes(offset, length as usize), + None => self.input.0.get_bytes(offset, SBBF_HEADER_SIZE_ESTIMATE), + }?; + + let (header, bitset_offset) = + chunk_read_bloom_filter_header_and_offset(offset, buffer.clone())?; + + match header.algorithm { + BloomFilterAlgorithm::BLOCK => { + // this match exists to future proof the singleton algorithm enum + } + } + match header.compression { + BloomFilterCompression::UNCOMPRESSED => { + // this match exists to future proof the singleton compression enum + } + } + match header.hash { + BloomFilterHash::XXHASH => { + // this match exists to future proof the singleton hash enum + } + } + + let bitset = match column_metadata.bloom_filter_length() { + Some(_) => buffer.slice( + (TryInto::::try_into(bitset_offset).unwrap() + - TryInto::::try_into(offset).unwrap()).., + ), + None => { + let bitset_length: usize = header.num_bytes.try_into().map_err(|_| { + ParquetError::General("Bloom filter length is invalid".to_string()) + })?; + self.input.0.get_bytes(bitset_offset, bitset_length)? + } + }; + Ok(Some(Sbbf::new(&bitset))) + } + + /// Build a [`ParquetRecordBatchReader`] + /// + /// Note: this will eagerly evaluate any `RowFilter` before returning + pub fn build(self) -> Result { + let Self { + input, + metadata, + schema: _, + fields, + batch_size, + row_groups, + projection, + mut filter, + selection, + row_selection_policy, + limit, + offset, + metrics, + // Not used for the sync reader, see https://github.com/apache/arrow-rs/issues/8000 + max_predicate_cache_size: _, + } = self; + + // Try to avoid allocate large buffer + let batch_size = batch_size.min(metadata.file_metadata().num_rows() as usize); + + let row_groups = row_groups.unwrap_or_else(|| (0..metadata.num_row_groups()).collect()); + + let reader = ReaderRowGroups { + reader: Arc::new(input.0), + metadata, + row_groups, + }; + + let mut plan_builder = ReadPlanBuilder::new(batch_size) + .with_selection(selection) + .with_row_selection_policy(row_selection_policy); + + // Update selection based on any filters + if let Some(filter) = filter.as_mut() { + for predicate in filter.predicates.iter_mut() { + // break early if we have ruled out all rows + if !plan_builder.selects_any() { + break; + } + + let mut cache_projection = predicate.projection().clone(); + cache_projection.intersect(&projection); + + let array_reader = ArrayReaderBuilder::new(&reader, &metrics) + .with_batch_size(batch_size) + .with_parquet_metadata(&reader.metadata) + .build_array_reader(fields.as_deref(), predicate.projection())?; + + plan_builder = plan_builder.with_predicate(array_reader, predicate.as_mut())?; + } + } + + let array_reader = ArrayReaderBuilder::new(&reader, &metrics) + .with_batch_size(batch_size) + .with_parquet_metadata(&reader.metadata) + .build_array_reader(fields.as_deref(), &projection)?; + + let read_plan = plan_builder + .limited(reader.num_rows()) + .with_offset(offset) + .with_limit(limit) + .build_limited() + .build(); + + Ok(ParquetRecordBatchReader::new(array_reader, read_plan)) + } +} + +struct ReaderRowGroups { + reader: Arc, + + metadata: Arc, + /// Optional list of row group indices to scan + row_groups: Vec, +} + +impl RowGroups for ReaderRowGroups { + fn num_rows(&self) -> usize { + let meta = self.metadata.row_groups(); + self.row_groups + .iter() + .map(|x| meta[*x].num_rows() as usize) + .sum() + } + + fn column_chunks(&self, i: usize) -> Result> { + Ok(Box::new(ReaderPageIterator { + column_idx: i, + reader: self.reader.clone(), + metadata: self.metadata.clone(), + row_groups: self.row_groups.clone().into_iter(), + })) + } + + fn row_groups(&self) -> Box + '_> { + Box::new( + self.row_groups + .iter() + .map(move |i| self.metadata.row_group(*i)), + ) + } + + fn metadata(&self) -> &ParquetMetaData { + self.metadata.as_ref() + } +} + +struct ReaderPageIterator { + reader: Arc, + column_idx: usize, + row_groups: std::vec::IntoIter, + metadata: Arc, +} + +impl ReaderPageIterator { + /// Return the next SerializedPageReader + fn next_page_reader(&mut self, rg_idx: usize) -> Result> { + let rg = self.metadata.row_group(rg_idx); + let column_chunk_metadata = rg.column(self.column_idx); + let offset_index = self.metadata.offset_index(); + // `offset_index` may not exist and `i[rg_idx]` will be empty. + // To avoid `i[rg_idx][self.column_idx`] panic, we need to filter out empty `i[rg_idx]`. + let page_locations = offset_index + .filter(|i| !i[rg_idx].is_empty()) + .map(|i| i[rg_idx][self.column_idx].page_locations.clone()); + let total_rows = rg.num_rows() as usize; + let reader = self.reader.clone(); + + SerializedPageReader::new(reader, column_chunk_metadata, total_rows, page_locations)? + .add_crypto_context( + rg_idx, + self.column_idx, + self.metadata.as_ref(), + column_chunk_metadata, + ) + } +} + +impl Iterator for ReaderPageIterator { + type Item = Result>; + + fn next(&mut self) -> Option { + let rg_idx = self.row_groups.next()?; + let page_reader = self + .next_page_reader(rg_idx) + .map(|page_reader| Box::new(page_reader) as _); + Some(page_reader) + } +} + +impl PageIterator for ReaderPageIterator {} + +/// Reads Parquet data as Arrow [`RecordBatch`]es +/// +/// This struct implements the [`RecordBatchReader`] trait and is an +/// `Iterator>` that yields [`RecordBatch`]es. +/// +/// Typically, either reads from a file or an in memory buffer [`Bytes`] +/// +/// Created by [`ParquetRecordBatchReaderBuilder`] +/// +/// [`Bytes`]: bytes::Bytes +pub struct ParquetRecordBatchReader { + array_reader: Box, + schema: SchemaRef, + read_plan: ReadPlan, +} + +/// Accumulates filter masks for decoded chunks in one logical output batch. +/// +/// The first chunk keeps its [`BooleanBuffer`] without copying. A second chunk +/// promotes the accumulator to a [`BooleanBufferBuilder`], and later chunks are +/// appended to it. For example, chunks `1001` and `1` become `10011`: +/// +/// ```text +/// append(1001) append(1) +/// Empty ───────────────▶ Single ───────────────▶ Combined +/// 1001 10011 +/// (zero copy) (promoted to builder) +/// ``` +/// +/// The combined mask lines up with the rows that [`ArrayReader::read_records`] +/// buffered across the decoded chunks (see [`read_mask_batch`]). Consuming the +/// buffered batch and filtering it with the accumulated mask yields the output. +/// +/// ```text +/// decoded rows: 0 1 2 3 11 <-- buffered by the array reader +/// chunk masks: [1 0 0 1] [1] +/// finish(): 1 0 0 1 1 <-- filters the whole batch in one pass +/// ``` +#[derive(Default)] +enum FilterMaskAccumulator { + #[default] + Empty, + Single(BooleanBuffer), + Combined(BooleanBufferBuilder), +} + +impl FilterMaskAccumulator { + fn append(&mut self, mask: BooleanBuffer) { + *self = match std::mem::take(self) { + Self::Empty => Self::Single(mask), + Self::Single(first) => { + let mut combined = BooleanBufferBuilder::new(first.len() + mask.len()); + combined.append_buffer(&first); + combined.append_buffer(&mask); + Self::Combined(combined) + } + Self::Combined(mut combined) => { + combined.append_buffer(&mask); + Self::Combined(combined) + } + }; + } + + fn finish(self) -> Option { + match self { + Self::Empty => None, + Self::Single(mask) => Some(mask), + Self::Combined(combined) => Some(combined.build()), + } + } +} + +/// Converts the projection buffered by `array_reader` into a record batch. +fn consume_record_batch(array_reader: &mut dyn ArrayReader) -> Result { + let array = array_reader.consume_batch()?; + let struct_array = array.as_struct_opt().ok_or_else(|| { + ArrowError::ParquetError("Struct array reader should return struct array".to_string()) + })?; + Ok(RecordBatch::from(struct_array)) +} + +/// Reads one logical Mask batch, potentially spanning multiple loaded ranges. +/// +/// Each [`MaskCursor`] chunk is safe to decode because it stays within loaded +/// pages. Gaps are crossed with [`ArrayReader::skip_records`], while decoded +/// arrays and their mask fragments remain buffered. Once `batch_size` selected +/// rows have accumulated, this consumes the underlying batch and filters it +/// once with the combined mask. +fn read_mask_batch( + array_reader: &mut dyn ArrayReader, + mask_cursor: &mut MaskCursor, + batch_size: usize, +) -> Result> { + let mut selected_rows = 0; + let mut filter_mask = FilterMaskAccumulator::default(); + + while selected_rows < batch_size && !mask_cursor.is_empty() { + let mask_chunk = mask_cursor.next_chunk(batch_size - selected_rows)?; + + if mask_chunk.initial_skip > 0 { + let skipped = array_reader.skip_records(mask_chunk.initial_skip)?; + if skipped != mask_chunk.initial_skip { + return Err(general_err!( + "failed to skip rows, expected {}, got {}", + mask_chunk.initial_skip, + skipped + )); + } + } + + let mask = mask_cursor.mask_values_for(&mask_chunk)?; + let read = array_reader.read_records(mask_chunk.chunk_rows)?; + if read == 0 { + return Err(general_err!( + "reached end of column while expecting {} rows", + mask_chunk.chunk_rows + )); + } + if read != mask_chunk.chunk_rows { + return Err(general_err!( + "insufficient rows read from array reader - expected {}, got {}", + mask_chunk.chunk_rows, + read + )); + } + + filter_mask.append(mask.values().clone()); + selected_rows += mask_chunk.selected_rows; + } + + if selected_rows == 0 { + return Ok(None); + } + + let filter_mask = filter_mask + .finish() + .ok_or_else(|| general_err!("Internal Error: decoded Mask batch has no filter values"))?; + let batch = consume_record_batch(array_reader)?; + let filtered_batch = filter_record_batch(&batch, &BooleanArray::from(filter_mask))?; + if filtered_batch.num_rows() != selected_rows { + return Err(general_err!( + "filtered rows mismatch selection - expected {}, got {}", + selected_rows, + filtered_batch.num_rows() + )); + } + + Ok(Some(filtered_batch)) +} + +impl Debug for ParquetRecordBatchReader { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ParquetRecordBatchReader") + .field("array_reader", &"...") + .field("schema", &self.schema) + .field("read_plan", &self.read_plan) + .finish() + } +} + +impl Iterator for ParquetRecordBatchReader { + type Item = Result; + + fn next(&mut self) -> Option { + self.next_inner() + .map_err(|arrow_err| arrow_err.into()) + .transpose() + } +} + +impl ParquetRecordBatchReader { + /// Returns the next `RecordBatch` from the reader, or `None` if the reader + /// has reached the end of the file. + /// + /// Returns `Result>` rather than `Option>` to + /// simplify error handling with `?` + fn next_inner(&mut self) -> Result> { + let mut read_records = 0; + let batch_size = self.batch_size(); + if batch_size == 0 { + return Ok(None); + } + match self.read_plan.row_selection_cursor_mut() { + RowSelectionCursor::Mask(mask_cursor) => { + return read_mask_batch(self.array_reader.as_mut(), mask_cursor, batch_size); + } + RowSelectionCursor::Selectors(selectors_cursor) => { + while read_records < batch_size && !selectors_cursor.is_empty() { + let front = selectors_cursor.next_selector(); + if front.skip { + let skipped = self.array_reader.skip_records(front.row_count)?; + + if skipped != front.row_count { + return Err(general_err!( + "failed to skip rows, expected {}, got {}", + front.row_count, + skipped + )); + } + continue; + } + + //Currently, when RowSelectors with row_count = 0 are included then its interpreted as end of reader. + //Fix is to skip such entries. See https://github.com/apache/arrow-rs/issues/2669 + if front.row_count == 0 { + continue; + } + + // try to read record + let need_read = batch_size - read_records; + let to_read = match front.row_count.checked_sub(need_read) { + Some(remaining) if remaining != 0 => { + // if page row count less than batch_size we must set batch size to page row count. + // add check avoid dead loop + selectors_cursor.return_selector(RowSelector::select(remaining)); + need_read + } + _ => front.row_count, + }; + match self.array_reader.read_records(to_read)? { + 0 => break, + rec => read_records += rec, + }; + } + } + RowSelectionCursor::All => { + self.array_reader.read_records(batch_size)?; + } + } + + let batch = consume_record_batch(self.array_reader.as_mut())?; + Ok(if batch.num_rows() > 0 { + Some(batch) + } else { + None + }) + } +} + +impl RecordBatchReader for ParquetRecordBatchReader { + /// Returns the projected [`SchemaRef`] for reading the parquet file. + /// + /// Note that the schema metadata will be stripped here. See + /// [`ParquetRecordBatchReaderBuilder::schema`] if the metadata is desired. + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +impl ParquetRecordBatchReader { + /// Create a new [`ParquetRecordBatchReader`] from the provided chunk reader + /// + /// See [`ParquetRecordBatchReaderBuilder`] for more options + pub fn try_new(reader: T, batch_size: usize) -> Result { + ParquetRecordBatchReaderBuilder::try_new(reader)? + .with_batch_size(batch_size) + .build() + } + + /// Create a new [`ParquetRecordBatchReader`] from the provided [`RowGroups`] + /// + /// Note: this is a low-level interface see [`ParquetRecordBatchReader::try_new`] for a + /// higher-level interface for reading parquet data from a file + pub fn try_new_with_row_groups( + levels: &FieldLevels, + row_groups: &dyn RowGroups, + batch_size: usize, + selection: Option, + ) -> Result { + // note metrics are not supported in this API + let metrics = ArrowReaderMetrics::disabled(); + let array_reader = ArrayReaderBuilder::new(row_groups, &metrics) + .with_batch_size(batch_size) + .with_parquet_metadata(row_groups.metadata()) + .build_array_reader(levels.levels.as_ref(), &ProjectionMask::all())?; + + let read_plan = ReadPlanBuilder::new(batch_size) + .with_selection(selection) + .build(); + + Ok(Self { + array_reader, + schema: Arc::new(Schema::new(levels.fields.clone())), + read_plan, + }) + } + + /// Create a new [`ParquetRecordBatchReader`] that will read at most `batch_size` rows at + /// a time from [`ArrayReader`] based on the configured `selection`. If `selection` is `None` + /// all rows will be returned + pub(crate) fn new(array_reader: Box, read_plan: ReadPlan) -> Self { + let schema = match array_reader.get_data_type() { + ArrowType::Struct(fields) => Schema::new(fields.clone()), + _ => unreachable!("Struct array reader's data type is not struct!"), + }; + + Self { + array_reader, + schema: Arc::new(schema), + read_plan, + } + } + + #[inline(always)] + pub(crate) fn batch_size(&self) -> usize { + self.read_plan.batch_size() + } +} + +#[cfg(test)] +pub(crate) mod tests { + use std::cmp::min; + use std::collections::{HashMap, VecDeque}; + use std::fmt::Formatter; + use std::fs::File; + use std::io::Seek; + use std::path::PathBuf; + use std::sync::Arc; + + use rand::rngs::StdRng; + use rand::{Rng, RngCore, SeedableRng, random, rng}; + use tempfile::tempfile; + + use crate::arrow::arrow_reader::{ + ArrowPredicateFn, ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReader, + ParquetRecordBatchReaderBuilder, RowFilter, RowSelection, RowSelector, + }; + use crate::arrow::schema::{ + add_encoded_arrow_schema_to_metadata, + virtual_type::{RowGroupIndex, RowNumber}, + }; + use crate::arrow::{ArrowWriter, ProjectionMask}; + use crate::basic::{ConvertedType, Encoding, LogicalType, Repetition, Type as PhysicalType}; + use crate::column::reader::decoder::REPETITION_LEVELS_BATCH_SIZE; + use crate::data_type::{ + BoolType, ByteArray, ByteArrayType, DataType, DoubleType, FixedLenByteArray, + FixedLenByteArrayType, FloatType, Int32Type, Int64Type, Int96, Int96Type, + }; + use crate::errors::Result; + use crate::file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetStatisticsPolicy}; + use crate::file::properties::{EnabledStatistics, WriterProperties, WriterVersion}; + use crate::file::writer::{SerializedFileWriter, SerializedRowGroupWriter}; + use crate::schema::parser::parse_message_type; + use crate::schema::types::{Type, TypePtr}; + use crate::util::test_common::rand_gen::RandGen; + use arrow_array::builder::*; + use arrow_array::cast::AsArray; + use arrow_array::types::{ + Date32Type, Date64Type, Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, + DecimalType, Float16Type, Float32Type, Float64Type, Time32MillisecondType, + Time64MicrosecondType, + }; + use arrow_array::*; + use arrow_buffer::{ArrowNativeType, BooleanBuffer, Buffer, IntervalDayTime, NullBuffer, i256}; + use arrow_data::{ArrayData, ArrayDataBuilder}; + use arrow_schema::{DataType as ArrowDataType, Field, Fields, Schema, SchemaRef, TimeUnit}; + use arrow_select::concat::concat_batches; + use bytes::Bytes; + use half::f16; + use num_traits::PrimInt; + + #[test] + fn filter_mask_accumulator_handles_empty_single_and_multiple_chunks() { + let first = BooleanBuffer::from(vec![true, false, false, false]); + let second = BooleanBuffer::from(vec![true]); + let third = BooleanBuffer::from(vec![false, true]); + + assert!(super::FilterMaskAccumulator::default().finish().is_none()); + + let mut single = super::FilterMaskAccumulator::default(); + single.append(first.clone()); + assert_eq!(single.finish().unwrap(), first); + + let mut combined = super::FilterMaskAccumulator::default(); + combined.append(BooleanBuffer::from(vec![true, false, false, false])); + combined.append(second); + combined.append(third); + assert_eq!( + combined.finish().unwrap(), + BooleanBuffer::from(vec![true, false, false, false, true, false, true]) + ); + } + + #[test] + fn test_arrow_reader_all_columns() { + let file = get_test_file("parquet/generated_simple_numerics/blogs.parquet"); + + let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + let original_schema = Arc::clone(builder.schema()); + let reader = builder.build().unwrap(); + + // Verify that the schema was correctly parsed + assert_eq!(original_schema.fields(), reader.schema().fields()); + } + + #[test] + fn test_reuse_schema() { + let file = get_test_file("parquet/alltypes-java.parquet"); + + let builder = ParquetRecordBatchReaderBuilder::try_new(file.try_clone().unwrap()).unwrap(); + let expected = builder.metadata; + let schema = expected.file_metadata().schema_descr_ptr(); + + let arrow_options = ArrowReaderOptions::new().with_parquet_schema(schema.clone()); + let builder = + ParquetRecordBatchReaderBuilder::try_new_with_options(file, arrow_options).unwrap(); + + // Verify that the metadata matches + assert_eq!(expected.as_ref(), builder.metadata.as_ref()); + } + + #[test] + fn test_page_encoding_stats_mask() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_tiny_pages.parquet"); + let file = File::open(path).unwrap(); + + let arrow_options = ArrowReaderOptions::new().with_encoding_stats_as_mask(true); + let builder = + ParquetRecordBatchReaderBuilder::try_new_with_options(file, arrow_options).unwrap(); + + let row_group_metadata = builder.metadata.row_group(0); + + // test page encoding stats + let page_encoding_stats = row_group_metadata + .column(0) + .page_encoding_stats_mask() + .unwrap(); + assert!(page_encoding_stats.is_only(Encoding::PLAIN)); + let page_encoding_stats = row_group_metadata + .column(2) + .page_encoding_stats_mask() + .unwrap(); + assert!(page_encoding_stats.is_only(Encoding::PLAIN_DICTIONARY)); + } + + #[test] + fn test_stats_stats_skipped() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_tiny_pages.parquet"); + let file = File::open(path).unwrap(); + + // test skipping all + let arrow_options = ArrowReaderOptions::new() + .with_encoding_stats_policy(ParquetStatisticsPolicy::SkipAll) + .with_column_stats_policy(ParquetStatisticsPolicy::SkipAll); + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options( + file.try_clone().unwrap(), + arrow_options, + ) + .unwrap(); + + let row_group_metadata = builder.metadata.row_group(0); + for column in row_group_metadata.columns() { + assert!(column.page_encoding_stats().is_none()); + assert!(column.page_encoding_stats_mask().is_none()); + assert!(column.statistics().is_none()); + } + + // test skipping all but one column and converting to mask + let arrow_options = ArrowReaderOptions::new() + .with_encoding_stats_as_mask(true) + .with_encoding_stats_policy(ParquetStatisticsPolicy::skip_except(&[0])) + .with_column_stats_policy(ParquetStatisticsPolicy::skip_except(&[0])); + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options( + file.try_clone().unwrap(), + arrow_options, + ) + .unwrap(); + + let row_group_metadata = builder.metadata.row_group(0); + for (idx, column) in row_group_metadata.columns().iter().enumerate() { + assert!(column.page_encoding_stats().is_none()); + assert_eq!(column.page_encoding_stats_mask().is_some(), idx == 0); + assert_eq!(column.statistics().is_some(), idx == 0); + } + } + + #[test] + fn test_size_stats_stats_skipped() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/repeated_primitive_no_list.parquet"); + let file = File::open(path).unwrap(); + + // test skipping all + let arrow_options = + ArrowReaderOptions::new().with_size_stats_policy(ParquetStatisticsPolicy::SkipAll); + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options( + file.try_clone().unwrap(), + arrow_options, + ) + .unwrap(); + + let row_group_metadata = builder.metadata.row_group(0); + for column in row_group_metadata.columns() { + assert!(column.repetition_level_histogram().is_none()); + assert!(column.definition_level_histogram().is_none()); + assert!(column.unencoded_byte_array_data_bytes().is_none()); + } + + // test skipping all but one column and converting to mask + let arrow_options = ArrowReaderOptions::new() + .with_encoding_stats_as_mask(true) + .with_size_stats_policy(ParquetStatisticsPolicy::skip_except(&[1])); + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options( + file.try_clone().unwrap(), + arrow_options, + ) + .unwrap(); + + let row_group_metadata = builder.metadata.row_group(0); + for (idx, column) in row_group_metadata.columns().iter().enumerate() { + assert_eq!(column.repetition_level_histogram().is_some(), idx == 1); + assert_eq!(column.definition_level_histogram().is_some(), idx == 1); + assert_eq!(column.unencoded_byte_array_data_bytes().is_some(), idx == 1); + } + } + + #[test] + fn test_arrow_reader_single_column() { + let file = get_test_file("parquet/generated_simple_numerics/blogs.parquet"); + + let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + let original_schema = Arc::clone(builder.schema()); + + let mask = ProjectionMask::leaves(builder.parquet_schema(), [2]); + let reader = builder.with_projection(mask).build().unwrap(); + + // Verify that the schema was correctly parsed + assert_eq!(1, reader.schema().fields().len()); + assert_eq!(original_schema.fields()[1], reader.schema().fields()[0]); + } + + #[test] + fn test_arrow_reader_single_column_by_name() { + let file = get_test_file("parquet/generated_simple_numerics/blogs.parquet"); + + let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + let original_schema = Arc::clone(builder.schema()); + + let mask = ProjectionMask::columns(builder.parquet_schema(), ["blog_id"]); + let reader = builder.with_projection(mask).build().unwrap(); + + // Verify that the schema was correctly parsed + assert_eq!(1, reader.schema().fields().len()); + assert_eq!(original_schema.fields()[1], reader.schema().fields()[0]); + } + + #[test] + fn test_null_column_reader_test() { + let mut file = tempfile::tempfile().unwrap(); + + let schema = " + message message { + OPTIONAL INT32 int32; + } + "; + let schema = Arc::new(parse_message_type(schema).unwrap()); + + let def_levels = vec![vec![0, 0, 0], vec![0, 0, 0, 0]]; + generate_single_column_file_with_data::( + &[vec![], vec![]], + Some(&def_levels), + file.try_clone().unwrap(), // Cannot use &mut File (#1163) + schema, + Some(Field::new("int32", ArrowDataType::Null, true)), + &Default::default(), + ) + .unwrap(); + + file.rewind().unwrap(); + + let record_reader = ParquetRecordBatchReader::try_new(file, 2).unwrap(); + let batches = record_reader.collect::, _>>().unwrap(); + + assert_eq!(batches.len(), 4); + for batch in &batches[0..3] { + assert_eq!(batch.num_rows(), 2); + assert_eq!(batch.num_columns(), 1); + assert_eq!(batch.column(0).null_count(), 2); + } + + assert_eq!(batches[3].num_rows(), 1); + assert_eq!(batches[3].num_columns(), 1); + assert_eq!(batches[3].column(0).null_count(), 1); + } + + #[test] + fn test_primitive_single_column_reader_test() { + run_single_column_reader_tests::( + 2, + ConvertedType::NONE, + None, + |vals| Arc::new(BooleanArray::from_iter(vals.iter().cloned())), + &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY], + ); + run_single_column_reader_tests::( + 2, + ConvertedType::NONE, + None, + |vals| Arc::new(Int32Array::from_iter(vals.iter().cloned())), + &[ + Encoding::PLAIN, + Encoding::RLE_DICTIONARY, + Encoding::DELTA_BINARY_PACKED, + Encoding::BYTE_STREAM_SPLIT, + ], + ); + run_single_column_reader_tests::( + 2, + ConvertedType::NONE, + None, + |vals| Arc::new(Int64Array::from_iter(vals.iter().cloned())), + &[ + Encoding::PLAIN, + Encoding::RLE_DICTIONARY, + Encoding::DELTA_BINARY_PACKED, + Encoding::BYTE_STREAM_SPLIT, + ], + ); + run_single_column_reader_tests::( + 2, + ConvertedType::NONE, + None, + |vals| Arc::new(Float32Array::from_iter(vals.iter().cloned())), + &[Encoding::PLAIN, Encoding::BYTE_STREAM_SPLIT], + ); + } + + #[test] + fn test_unsigned_primitive_single_column_reader_test() { + run_single_column_reader_tests::( + 2, + ConvertedType::UINT_32, + Some(ArrowDataType::UInt32), + |vals| { + Arc::new(UInt32Array::from_iter( + vals.iter().map(|x| x.map(|x| x as u32)), + )) + }, + &[ + Encoding::PLAIN, + Encoding::RLE_DICTIONARY, + Encoding::DELTA_BINARY_PACKED, + ], + ); + run_single_column_reader_tests::( + 2, + ConvertedType::UINT_64, + Some(ArrowDataType::UInt64), + |vals| { + Arc::new(UInt64Array::from_iter( + vals.iter().map(|x| x.map(|x| x as u64)), + )) + }, + &[ + Encoding::PLAIN, + Encoding::RLE_DICTIONARY, + Encoding::DELTA_BINARY_PACKED, + ], + ); + } + + #[test] + fn test_unsigned_roundtrip() { + let schema = Arc::new(Schema::new(vec![ + Field::new("uint32", ArrowDataType::UInt32, true), + Field::new("uint64", ArrowDataType::UInt64, true), + ])); + + let mut buf = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None).unwrap(); + + let original = RecordBatch::try_new( + schema, + vec![ + Arc::new(UInt32Array::from_iter_values([ + 0, + i32::MAX as u32, + u32::MAX, + ])), + Arc::new(UInt64Array::from_iter_values([ + 0, + i64::MAX as u64, + u64::MAX, + ])), + ], + ) + .unwrap(); + + writer.write(&original).unwrap(); + writer.close().unwrap(); + + let mut reader = ParquetRecordBatchReader::try_new(Bytes::from(buf), 1024).unwrap(); + let ret = reader.next().unwrap().unwrap(); + assert_eq!(ret, original); + + // Check they can be downcast to the correct type + ret.column(0) + .as_any() + .downcast_ref::() + .unwrap(); + + ret.column(1) + .as_any() + .downcast_ref::() + .unwrap(); + } + + #[test] + fn test_float16_roundtrip() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("float16", ArrowDataType::Float16, false), + Field::new("float16-nullable", ArrowDataType::Float16, true), + ])); + + let mut buf = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None)?; + + let original = RecordBatch::try_new( + schema, + vec![ + Arc::new(Float16Array::from_iter_values([ + f16::EPSILON, + f16::MIN, + f16::MAX, + f16::NAN, + f16::INFINITY, + f16::NEG_INFINITY, + f16::ONE, + f16::NEG_ONE, + f16::ZERO, + f16::NEG_ZERO, + f16::E, + f16::PI, + f16::FRAC_1_PI, + ])), + Arc::new(Float16Array::from(vec![ + None, + None, + None, + Some(f16::NAN), + Some(f16::INFINITY), + Some(f16::NEG_INFINITY), + None, + None, + None, + None, + None, + None, + Some(f16::FRAC_1_PI), + ])), + ], + )?; + + writer.write(&original)?; + writer.close()?; + + let mut reader = ParquetRecordBatchReader::try_new(Bytes::from(buf), 1024)?; + let ret = reader.next().unwrap()?; + assert_eq!(ret, original); + + // Ensure can be downcast to the correct type + ret.column(0).as_primitive::(); + ret.column(1).as_primitive::(); + + Ok(()) + } + + #[test] + fn test_time_utc_roundtrip() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "time_millis", + ArrowDataType::Time32(TimeUnit::Millisecond), + true, + ) + .with_metadata(HashMap::from_iter(vec![( + "adjusted_to_utc".to_string(), + "".to_string(), + )])), + Field::new( + "time_micros", + ArrowDataType::Time64(TimeUnit::Microsecond), + true, + ) + .with_metadata(HashMap::from_iter(vec![( + "adjusted_to_utc".to_string(), + "".to_string(), + )])), + ])); + + let mut buf = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None)?; + + let original = RecordBatch::try_new( + schema, + vec![ + Arc::new(Time32MillisecondArray::from(vec![ + Some(-1), + Some(0), + Some(86_399_000), + Some(86_400_000), + Some(86_401_000), + None, + ])), + Arc::new(Time64MicrosecondArray::from(vec![ + Some(-1), + Some(0), + Some(86_399 * 1_000_000), + Some(86_400 * 1_000_000), + Some(86_401 * 1_000_000), + None, + ])), + ], + )?; + + writer.write(&original)?; + writer.close()?; + + let mut reader = ParquetRecordBatchReader::try_new(Bytes::from(buf), 1024)?; + let ret = reader.next().unwrap()?; + assert_eq!(ret, original); + + // Ensure can be downcast to the correct type + ret.column(0).as_primitive::(); + ret.column(1).as_primitive::(); + + Ok(()) + } + + #[test] + fn test_date32_roundtrip() -> Result<()> { + use arrow_array::Date32Array; + + let schema = Arc::new(Schema::new(vec![Field::new( + "date32", + ArrowDataType::Date32, + false, + )])); + + let mut buf = Vec::with_capacity(1024); + + let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None)?; + + let original = RecordBatch::try_new( + schema, + vec![Arc::new(Date32Array::from(vec![ + -1_000_000, -100_000, -10_000, -1_000, 0, 1_000, 10_000, 100_000, 1_000_000, + ]))], + )?; + + writer.write(&original)?; + writer.close()?; + + let mut reader = ParquetRecordBatchReader::try_new(Bytes::from(buf), 1024)?; + let ret = reader.next().unwrap()?; + assert_eq!(ret, original); + + // Ensure can be downcast to the correct type + ret.column(0).as_primitive::(); + + Ok(()) + } + + #[test] + fn test_date64_roundtrip() -> Result<()> { + use arrow_array::Date64Array; + + let schema = Arc::new(Schema::new(vec![ + Field::new("small-date64", ArrowDataType::Date64, false), + Field::new("big-date64", ArrowDataType::Date64, false), + Field::new("invalid-date64", ArrowDataType::Date64, false), + ])); + + let mut default_buf = Vec::with_capacity(1024); + let mut coerce_buf = Vec::with_capacity(1024); + + let coerce_props = WriterProperties::builder().set_coerce_types(true).build(); + + let mut default_writer = ArrowWriter::try_new(&mut default_buf, schema.clone(), None)?; + let mut coerce_writer = + ArrowWriter::try_new(&mut coerce_buf, schema.clone(), Some(coerce_props))?; + + static NUM_MILLISECONDS_IN_DAY: i64 = 1000 * 60 * 60 * 24; + + let original = RecordBatch::try_new( + schema, + vec![ + // small-date64 + Arc::new(Date64Array::from(vec![ + -1_000_000 * NUM_MILLISECONDS_IN_DAY, + -1_000 * NUM_MILLISECONDS_IN_DAY, + 0, + 1_000 * NUM_MILLISECONDS_IN_DAY, + 1_000_000 * NUM_MILLISECONDS_IN_DAY, + ])), + // big-date64 + Arc::new(Date64Array::from(vec![ + -10_000_000_000 * NUM_MILLISECONDS_IN_DAY, + -1_000_000_000 * NUM_MILLISECONDS_IN_DAY, + 0, + 1_000_000_000 * NUM_MILLISECONDS_IN_DAY, + 10_000_000_000 * NUM_MILLISECONDS_IN_DAY, + ])), + // invalid-date64 + Arc::new(Date64Array::from(vec![ + -1_000_000 * NUM_MILLISECONDS_IN_DAY + 1, + -1_000 * NUM_MILLISECONDS_IN_DAY + 1, + 1, + 1_000 * NUM_MILLISECONDS_IN_DAY + 1, + 1_000_000 * NUM_MILLISECONDS_IN_DAY + 1, + ])), + ], + )?; + + default_writer.write(&original)?; + coerce_writer.write(&original)?; + + default_writer.close()?; + coerce_writer.close()?; + + let mut default_reader = ParquetRecordBatchReader::try_new(Bytes::from(default_buf), 1024)?; + let mut coerce_reader = ParquetRecordBatchReader::try_new(Bytes::from(coerce_buf), 1024)?; + + let default_ret = default_reader.next().unwrap()?; + let coerce_ret = coerce_reader.next().unwrap()?; + + // Roundtrip should be successful when default writer used + assert_eq!(default_ret, original); + + // Only small-date64 should roundtrip successfully when coerce_types writer is used + assert_eq!(coerce_ret.column(0), original.column(0)); + assert_ne!(coerce_ret.column(1), original.column(1)); + assert_ne!(coerce_ret.column(2), original.column(2)); + + // Ensure both can be downcast to the correct type + default_ret.column(0).as_primitive::(); + coerce_ret.column(0).as_primitive::(); + + Ok(()) + } + struct RandFixedLenGen {} + + impl RandGen for RandFixedLenGen { + fn r#gen(len: i32) -> FixedLenByteArray { + let mut v = vec![0u8; len as usize]; + rng().fill_bytes(&mut v); + ByteArray::from(v).into() + } + } + + #[test] + fn test_fixed_length_binary_column_reader() { + run_single_column_reader_tests::( + 20, + ConvertedType::NONE, + None, + |vals| { + let mut builder = FixedSizeBinaryBuilder::with_capacity(vals.len(), 20); + for val in vals { + match val { + Some(b) => builder.append_value(b).unwrap(), + None => builder.append_null(), + } + } + Arc::new(builder.finish()) + }, + &[Encoding::PLAIN, Encoding::RLE_DICTIONARY], + ); + } + + #[test] + fn test_interval_day_time_column_reader() { + run_single_column_reader_tests::( + 12, + ConvertedType::INTERVAL, + None, + |vals| { + Arc::new( + vals.iter() + .map(|x| { + x.as_ref().map(|b| IntervalDayTime { + days: i32::from_le_bytes(b.as_ref()[4..8].try_into().unwrap()), + milliseconds: i32::from_le_bytes( + b.as_ref()[8..12].try_into().unwrap(), + ), + }) + }) + .collect::(), + ) + }, + &[Encoding::PLAIN, Encoding::RLE_DICTIONARY], + ); + } + + #[test] + fn test_int96_single_column_reader_test() { + let encodings = &[Encoding::PLAIN, Encoding::RLE_DICTIONARY]; + + type TypeHintAndConversionFunction = + (Option, fn(&[Option]) -> ArrayRef); + + let resolutions: Vec = vec![ + // Test without a specified ArrowType hint. + (None, |vals: &[Option]| { + Arc::new(TimestampNanosecondArray::from_iter( + vals.iter().map(|x| x.map(|x| x.to_nanos())), + )) as ArrayRef + }), + // Test other TimeUnits as ArrowType hints. + ( + Some(ArrowDataType::Timestamp(TimeUnit::Second, None)), + |vals: &[Option]| { + Arc::new(TimestampSecondArray::from_iter( + vals.iter().map(|x| x.map(|x| x.to_seconds())), + )) as ArrayRef + }, + ), + ( + Some(ArrowDataType::Timestamp(TimeUnit::Millisecond, None)), + |vals: &[Option]| { + Arc::new(TimestampMillisecondArray::from_iter( + vals.iter().map(|x| x.map(|x| x.to_millis())), + )) as ArrayRef + }, + ), + ( + Some(ArrowDataType::Timestamp(TimeUnit::Microsecond, None)), + |vals: &[Option]| { + Arc::new(TimestampMicrosecondArray::from_iter( + vals.iter().map(|x| x.map(|x| x.to_micros())), + )) as ArrayRef + }, + ), + ( + Some(ArrowDataType::Timestamp(TimeUnit::Nanosecond, None)), + |vals: &[Option]| { + Arc::new(TimestampNanosecondArray::from_iter( + vals.iter().map(|x| x.map(|x| x.to_nanos())), + )) as ArrayRef + }, + ), + // Test another timezone with TimeUnit as ArrowType hints. + ( + Some(ArrowDataType::Timestamp( + TimeUnit::Second, + Some(Arc::from("-05:00")), + )), + |vals: &[Option]| { + Arc::new( + TimestampSecondArray::from_iter( + vals.iter().map(|x| x.map(|x| x.to_seconds())), + ) + .with_timezone("-05:00"), + ) as ArrayRef + }, + ), + ]; + + resolutions.iter().for_each(|(arrow_type, converter)| { + run_single_column_reader_tests::( + 2, + ConvertedType::NONE, + arrow_type.clone(), + converter, + encodings, + ); + }) + } + + #[test] + fn test_int96_from_spark_file_with_provided_schema() { + // int96_from_spark.parquet was written based on Spark's microsecond timestamps which trade + // range for resolution compared to a nanosecond timestamp. We must provide a schema with + // microsecond resolution for the Parquet reader to interpret these values correctly. + use arrow_schema::DataType::Timestamp; + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/int96_from_spark.parquet"); + let file = File::open(path).unwrap(); + + let supplied_schema = Arc::new(Schema::new(vec![Field::new( + "a", + Timestamp(TimeUnit::Microsecond, None), + true, + )])); + let options = ArrowReaderOptions::new().with_schema(supplied_schema.clone()); + + let mut record_reader = + ParquetRecordBatchReaderBuilder::try_new_with_options(file, options) + .unwrap() + .build() + .unwrap(); + + let batch = record_reader.next().unwrap().unwrap(); + assert_eq!(batch.num_columns(), 1); + let column = batch.column(0); + assert_eq!(column.data_type(), &Timestamp(TimeUnit::Microsecond, None)); + + let expected = Arc::new(Int64Array::from(vec![ + Some(1704141296123456), + Some(1704070800000000), + Some(253402225200000000), + Some(1735599600000000), + None, + Some(9089380393200000000), + ])); + + // arrow-rs relies on the chrono library to convert between timestamps and strings, so + // instead compare as Int64. The underlying type should be a PrimitiveArray of Int64 + // anyway, so this should be a zero-copy non-modifying cast. + + let binding = arrow_cast::cast(batch.column(0), &arrow_schema::DataType::Int64).unwrap(); + let casted_timestamps = binding.as_primitive::(); + + assert_eq!(casted_timestamps.len(), expected.len()); + + casted_timestamps + .iter() + .zip(expected.iter()) + .for_each(|(lhs, rhs)| { + assert_eq!(lhs, rhs); + }); + } + + #[test] + fn test_int96_from_spark_file_without_provided_schema() { + // int96_from_spark.parquet was written based on Spark's microsecond timestamps which trade + // range for resolution compared to a nanosecond timestamp. Without a provided schema, some + // values when read as nanosecond resolution overflow and result in garbage values. + use arrow_schema::DataType::Timestamp; + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/int96_from_spark.parquet"); + let file = File::open(path).unwrap(); + + let mut record_reader = ParquetRecordBatchReaderBuilder::try_new(file) + .unwrap() + .build() + .unwrap(); + + let batch = record_reader.next().unwrap().unwrap(); + assert_eq!(batch.num_columns(), 1); + let column = batch.column(0); + assert_eq!(column.data_type(), &Timestamp(TimeUnit::Nanosecond, None)); + + let expected = Arc::new(Int64Array::from(vec![ + Some(1704141296123456000), // Reads as nanosecond fine (note 3 extra 0s) + Some(1704070800000000000), // Reads as nanosecond fine (note 3 extra 0s) + Some(-4852191831933722624), // Cannot be represented with nanos timestamp (year 9999) + Some(1735599600000000000), // Reads as nanosecond fine (note 3 extra 0s) + None, + Some(-4864435138808946688), // Cannot be represented with nanos timestamp (year 290000) + ])); + + // arrow-rs relies on the chrono library to convert between timestamps and strings, so + // instead compare as Int64. The underlying type should be a PrimitiveArray of Int64 + // anyway, so this should be a zero-copy non-modifying cast. + + let binding = arrow_cast::cast(batch.column(0), &arrow_schema::DataType::Int64).unwrap(); + let casted_timestamps = binding.as_primitive::(); + + assert_eq!(casted_timestamps.len(), expected.len()); + + casted_timestamps + .iter() + .zip(expected.iter()) + .for_each(|(lhs, rhs)| { + assert_eq!(lhs, rhs); + }); + } + + struct RandUtf8Gen {} + + impl RandGen for RandUtf8Gen { + fn r#gen(len: i32) -> ByteArray { + Int32Type::r#gen(len).to_string().as_str().into() + } + } + + #[test] + fn test_utf8_single_column_reader_test() { + fn string_converter(vals: &[Option]) -> ArrayRef { + Arc::new(GenericStringArray::::from_iter(vals.iter().map(|x| { + x.as_ref().map(|b| std::str::from_utf8(b.data()).unwrap()) + }))) + } + + let encodings = &[ + Encoding::PLAIN, + Encoding::RLE_DICTIONARY, + Encoding::DELTA_LENGTH_BYTE_ARRAY, + Encoding::DELTA_BYTE_ARRAY, + ]; + + run_single_column_reader_tests::( + 2, + ConvertedType::NONE, + None, + |vals| { + Arc::new(BinaryArray::from_iter( + vals.iter().map(|x| x.as_ref().map(|x| x.data())), + )) + }, + encodings, + ); + + run_single_column_reader_tests::( + 2, + ConvertedType::UTF8, + None, + string_converter::, + encodings, + ); + + run_single_column_reader_tests::( + 2, + ConvertedType::UTF8, + Some(ArrowDataType::Utf8), + string_converter::, + encodings, + ); + + run_single_column_reader_tests::( + 2, + ConvertedType::UTF8, + Some(ArrowDataType::LargeUtf8), + string_converter::, + encodings, + ); + + let small_key_types = [ArrowDataType::Int8, ArrowDataType::UInt8]; + for key in &small_key_types { + for encoding in encodings { + let mut opts = TestOptions::new(2, 20, 15).with_null_percent(50); + opts.encoding = *encoding; + + let data_type = + ArrowDataType::Dictionary(Box::new(key.clone()), Box::new(ArrowDataType::Utf8)); + + // Cannot run full test suite as keys overflow, run small test instead + single_column_reader_test::( + opts, + 2, + ConvertedType::UTF8, + Some(data_type.clone()), + move |vals| { + let vals = string_converter::(vals); + arrow::compute::cast(&vals, &data_type).unwrap() + }, + ); + } + } + + let key_types = [ + ArrowDataType::Int16, + ArrowDataType::UInt16, + ArrowDataType::Int32, + ArrowDataType::UInt32, + ArrowDataType::Int64, + ArrowDataType::UInt64, + ]; + + for key in &key_types { + let data_type = + ArrowDataType::Dictionary(Box::new(key.clone()), Box::new(ArrowDataType::Utf8)); + + run_single_column_reader_tests::( + 2, + ConvertedType::UTF8, + Some(data_type.clone()), + move |vals| { + let vals = string_converter::(vals); + arrow::compute::cast(&vals, &data_type).unwrap() + }, + encodings, + ); + + let data_type = ArrowDataType::Dictionary( + Box::new(key.clone()), + Box::new(ArrowDataType::LargeUtf8), + ); + + run_single_column_reader_tests::( + 2, + ConvertedType::UTF8, + Some(data_type.clone()), + move |vals| { + let vals = string_converter::(vals); + arrow::compute::cast(&vals, &data_type).unwrap() + }, + encodings, + ); + } + } + + #[test] + fn test_decimal_nullable_struct() { + let decimals = Decimal256Array::from_iter_values( + [1, 2, 3, 4, 5, 6, 7, 8].into_iter().map(i256::from_i128), + ); + + let data = ArrayDataBuilder::new(ArrowDataType::Struct(Fields::from(vec![Field::new( + "decimals", + decimals.data_type().clone(), + false, + )]))) + .len(8) + .null_bit_buffer(Some(Buffer::from(&[0b11101111]))) + .child_data(vec![decimals.into_data()]) + .build() + .unwrap(); + + let written = + RecordBatch::try_from_iter([("struct", Arc::new(StructArray::from(data)) as ArrayRef)]) + .unwrap(); + + let mut buffer = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buffer, written.schema(), None).unwrap(); + writer.write(&written).unwrap(); + writer.close().unwrap(); + + let read = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 3) + .unwrap() + .collect::, _>>() + .unwrap(); + + assert_eq!(&written.slice(0, 3), &read[0]); + assert_eq!(&written.slice(3, 3), &read[1]); + assert_eq!(&written.slice(6, 2), &read[2]); + } + + #[test] + fn test_int32_nullable_struct() { + let int32 = Int32Array::from_iter_values([1, 2, 3, 4, 5, 6, 7, 8]); + let data = ArrayDataBuilder::new(ArrowDataType::Struct(Fields::from(vec![Field::new( + "int32", + int32.data_type().clone(), + false, + )]))) + .len(8) + .null_bit_buffer(Some(Buffer::from(&[0b11101111]))) + .child_data(vec![int32.into_data()]) + .build() + .unwrap(); + + let written = + RecordBatch::try_from_iter([("struct", Arc::new(StructArray::from(data)) as ArrayRef)]) + .unwrap(); + + let mut buffer = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buffer, written.schema(), None).unwrap(); + writer.write(&written).unwrap(); + writer.close().unwrap(); + + let read = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 3) + .unwrap() + .collect::, _>>() + .unwrap(); + + assert_eq!(&written.slice(0, 3), &read[0]); + assert_eq!(&written.slice(3, 3), &read[1]); + assert_eq!(&written.slice(6, 2), &read[2]); + } + + #[test] + fn test_decimal_list() { + let decimals = Decimal128Array::from_iter_values([1, 2, 3, 4, 5, 6, 7, 8]); + + // [[], [1], [2, 3], null, [4], null, [6, 7, 8]] + let data = ArrayDataBuilder::new(ArrowDataType::List(Arc::new(Field::new_list_field( + decimals.data_type().clone(), + false, + )))) + .len(7) + .add_buffer(Buffer::from_iter([0_i32, 0, 1, 3, 3, 4, 5, 8])) + .null_bit_buffer(Some(Buffer::from(&[0b01010111]))) + .child_data(vec![decimals.into_data()]) + .build() + .unwrap(); + + let written = + RecordBatch::try_from_iter([("list", Arc::new(ListArray::from(data)) as ArrayRef)]) + .unwrap(); + + let mut buffer = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buffer, written.schema(), None).unwrap(); + writer.write(&written).unwrap(); + writer.close().unwrap(); + + let read = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 3) + .unwrap() + .collect::, _>>() + .unwrap(); + + assert_eq!(&written.slice(0, 3), &read[0]); + assert_eq!(&written.slice(3, 3), &read[1]); + assert_eq!(&written.slice(6, 1), &read[2]); + } + + #[test] + fn test_read_decimal_file() { + use arrow_array::Decimal128Array; + let testdata = arrow::util::test_util::parquet_test_data(); + let file_variants = vec![ + ("byte_array", 4), + ("fixed_length", 25), + ("int32", 4), + ("int64", 10), + ]; + for (prefix, target_precision) in file_variants { + let path = format!("{testdata}/{prefix}_decimal.parquet"); + let file = File::open(path).unwrap(); + let mut record_reader = ParquetRecordBatchReader::try_new(file, 32).unwrap(); + + let batch = record_reader.next().unwrap().unwrap(); + assert_eq!(batch.num_rows(), 24); + let col = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + + let expected = 1..25; + + assert_eq!(col.precision(), target_precision); + assert_eq!(col.scale(), 2); + + for (i, v) in expected.enumerate() { + assert_eq!(col.value(i), v * 100_i128); + } + } + } + + #[test] + fn test_read_float16_nonzeros_file() { + use arrow_array::Float16Array; + let testdata = arrow::util::test_util::parquet_test_data(); + // see https://github.com/apache/parquet-testing/pull/40 + let path = format!("{testdata}/float16_nonzeros_and_nans.parquet"); + let file = File::open(path).unwrap(); + let mut record_reader = ParquetRecordBatchReader::try_new(file, 32).unwrap(); + + let batch = record_reader.next().unwrap().unwrap(); + assert_eq!(batch.num_rows(), 8); + let col = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + + let f16_two = f16::ONE + f16::ONE; + + assert_eq!(col.null_count(), 1); + assert!(col.is_null(0)); + assert_eq!(col.value(1), f16::ONE); + assert_eq!(col.value(2), -f16_two); + assert!(col.value(3).is_nan()); + assert_eq!(col.value(4), f16::ZERO); + assert!(col.value(4).is_sign_positive()); + assert_eq!(col.value(5), f16::NEG_ONE); + assert_eq!(col.value(6), f16::NEG_ZERO); + assert!(col.value(6).is_sign_negative()); + assert_eq!(col.value(7), f16_two); + } + + #[test] + fn test_read_float16_zeros_file() { + use arrow_array::Float16Array; + let testdata = arrow::util::test_util::parquet_test_data(); + // see https://github.com/apache/parquet-testing/pull/40 + let path = format!("{testdata}/float16_zeros_and_nans.parquet"); + let file = File::open(path).unwrap(); + let mut record_reader = ParquetRecordBatchReader::try_new(file, 32).unwrap(); + + let batch = record_reader.next().unwrap().unwrap(); + assert_eq!(batch.num_rows(), 3); + let col = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(col.null_count(), 1); + assert!(col.is_null(0)); + assert_eq!(col.value(1), f16::ZERO); + assert!(col.value(1).is_sign_positive()); + assert!(col.value(2).is_nan()); + } + + #[test] + fn test_read_float32_float64_byte_stream_split() { + let path = format!( + "{}/byte_stream_split.zstd.parquet", + arrow::util::test_util::parquet_test_data(), + ); + let file = File::open(path).unwrap(); + let record_reader = ParquetRecordBatchReader::try_new(file, 128).unwrap(); + + let mut row_count = 0; + for batch in record_reader { + let batch = batch.unwrap(); + row_count += batch.num_rows(); + let f32_col = batch.column(0).as_primitive::(); + let f64_col = batch.column(1).as_primitive::(); + + // This file contains floats from a standard normal distribution + for &x in f32_col.values() { + assert!(x > -10.0); + assert!(x < 10.0); + } + for &x in f64_col.values() { + assert!(x > -10.0); + assert!(x < 10.0); + } + } + assert_eq!(row_count, 300); + } + + #[test] + fn test_read_extended_byte_stream_split() { + let path = format!( + "{}/byte_stream_split_extended.gzip.parquet", + arrow::util::test_util::parquet_test_data(), + ); + let file = File::open(path).unwrap(); + let record_reader = ParquetRecordBatchReader::try_new(file, 128).unwrap(); + + let mut row_count = 0; + for batch in record_reader { + let batch = batch.unwrap(); + row_count += batch.num_rows(); + + // 0,1 are f16 + let f16_col = batch.column(0).as_primitive::(); + let f16_bss = batch.column(1).as_primitive::(); + assert_eq!(f16_col.len(), f16_bss.len()); + f16_col + .iter() + .zip(f16_bss.iter()) + .for_each(|(l, r)| assert_eq!(l.unwrap(), r.unwrap())); + + // 2,3 are f32 + let f32_col = batch.column(2).as_primitive::(); + let f32_bss = batch.column(3).as_primitive::(); + assert_eq!(f32_col.len(), f32_bss.len()); + f32_col + .iter() + .zip(f32_bss.iter()) + .for_each(|(l, r)| assert_eq!(l.unwrap(), r.unwrap())); + + // 4,5 are f64 + let f64_col = batch.column(4).as_primitive::(); + let f64_bss = batch.column(5).as_primitive::(); + assert_eq!(f64_col.len(), f64_bss.len()); + f64_col + .iter() + .zip(f64_bss.iter()) + .for_each(|(l, r)| assert_eq!(l.unwrap(), r.unwrap())); + + // 6,7 are i32 + let i32_col = batch.column(6).as_primitive::(); + let i32_bss = batch.column(7).as_primitive::(); + assert_eq!(i32_col.len(), i32_bss.len()); + i32_col + .iter() + .zip(i32_bss.iter()) + .for_each(|(l, r)| assert_eq!(l.unwrap(), r.unwrap())); + + // 8,9 are i64 + let i64_col = batch.column(8).as_primitive::(); + let i64_bss = batch.column(9).as_primitive::(); + assert_eq!(i64_col.len(), i64_bss.len()); + i64_col + .iter() + .zip(i64_bss.iter()) + .for_each(|(l, r)| assert_eq!(l.unwrap(), r.unwrap())); + + // 10,11 are FLBA(5) + let flba_col = batch.column(10).as_fixed_size_binary(); + let flba_bss = batch.column(11).as_fixed_size_binary(); + assert_eq!(flba_col.len(), flba_bss.len()); + flba_col + .iter() + .zip(flba_bss.iter()) + .for_each(|(l, r)| assert_eq!(l.unwrap(), r.unwrap())); + + // 12,13 are FLBA(4) (decimal(7,3)) + let dec_col = batch.column(12).as_primitive::(); + let dec_bss = batch.column(13).as_primitive::(); + assert_eq!(dec_col.len(), dec_bss.len()); + dec_col + .iter() + .zip(dec_bss.iter()) + .for_each(|(l, r)| assert_eq!(l.unwrap(), r.unwrap())); + } + assert_eq!(row_count, 200); + } + + #[test] + fn test_read_incorrect_map_schema_file() { + let testdata = arrow::util::test_util::parquet_test_data(); + // see https://github.com/apache/parquet-testing/pull/47 + let path = format!("{testdata}/incorrect_map_schema.parquet"); + let file = File::open(path).unwrap(); + let mut record_reader = ParquetRecordBatchReader::try_new(file, 32).unwrap(); + + let batch = record_reader.next().unwrap().unwrap(); + assert_eq!(batch.num_rows(), 1); + + let expected_schema = Schema::new(vec![Field::new( + "my_map", + ArrowDataType::Map( + Arc::new(Field::new( + "key_value", + ArrowDataType::Struct(Fields::from(vec![ + Field::new("key", ArrowDataType::Utf8, false), + Field::new("value", ArrowDataType::Utf8, true), + ])), + false, + )), + false, + ), + true, + )]); + assert_eq!(batch.schema().as_ref(), &expected_schema); + + assert_eq!(batch.num_rows(), 1); + assert_eq!(batch.column(0).null_count(), 0); + assert_eq!( + batch.column(0).as_map().keys().as_ref(), + &StringArray::from(vec!["parent", "name"]) + ); + assert_eq!( + batch.column(0).as_map().values().as_ref(), + &StringArray::from(vec!["another", "report"]) + ); + } + + #[test] + fn test_read_dict_fixed_size_binary() { + let schema = Arc::new(Schema::new(vec![Field::new( + "a", + ArrowDataType::Dictionary( + Box::new(ArrowDataType::UInt8), + Box::new(ArrowDataType::FixedSizeBinary(8)), + ), + true, + )])); + let keys = UInt8Array::from_iter_values(vec![0, 0, 1]); + let values = FixedSizeBinaryArray::try_from_iter( + vec![ + (0u8..8u8).collect::>(), + (24u8..32u8).collect::>(), + ] + .into_iter(), + ) + .unwrap(); + let arr = UInt8DictionaryArray::new(keys, Arc::new(values)); + let batch = RecordBatch::try_new(schema, vec![Arc::new(arr)]).unwrap(); + + let mut buffer = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let read = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 3) + .unwrap() + .collect::, _>>() + .unwrap(); + + assert_eq!(read.len(), 1); + assert_eq!(&batch, &read[0]) + } + + #[test] + fn test_read_nullable_structs_with_binary_dict_as_first_child_column() { + // the `StructArrayReader` will check the definition and repetition levels of the first + // child column in the struct to determine nullability for the struct. If the first + // column's is being read by `ByteArrayDictionaryReader` we need to ensure that the + // nullability is interpreted correctly from the rep/def level buffers managed by the + // buffers managed by this array reader. + + let struct_fields = Fields::from(vec![ + Field::new( + "city", + ArrowDataType::Dictionary( + Box::new(ArrowDataType::UInt8), + Box::new(ArrowDataType::Utf8), + ), + true, + ), + Field::new("name", ArrowDataType::Utf8, true), + ]); + let schema = Arc::new(Schema::new(vec![Field::new( + "items", + ArrowDataType::Struct(struct_fields.clone()), + true, + )])); + + let items_arr = StructArray::new( + struct_fields, + vec![ + Arc::new(DictionaryArray::new( + UInt8Array::from_iter_values(vec![0, 1, 1, 0, 2]), + Arc::new(StringArray::from_iter_values(vec![ + "quebec", + "fredericton", + "halifax", + ])), + )), + Arc::new(StringArray::from_iter_values(vec![ + "albert", "terry", "lance", "", "tim", + ])), + ], + Some(NullBuffer::from_iter(vec![true, true, true, false, true])), + ); + + let batch = RecordBatch::try_new(schema, vec![Arc::new(items_arr)]).unwrap(); + let mut buffer = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let read = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 8) + .unwrap() + .collect::, _>>() + .unwrap(); + + assert_eq!(read.len(), 1); + assert_eq!(&batch, &read[0]) + } + + /// Parameters for single_column_reader_test + #[derive(Clone)] + struct TestOptions { + /// Number of row group to write to parquet (row group size = + /// num_row_groups / num_rows) + num_row_groups: usize, + /// Total number of rows per row group + num_rows: usize, + /// Size of batches to read back + record_batch_size: usize, + /// Percentage of nulls in column or None if required + null_percent: Option, + /// Set write batch size + /// + /// This is the number of rows that are written at once to a page and + /// therefore acts as a bound on the page granularity of a row group + write_batch_size: usize, + /// Maximum size of page in bytes + max_data_page_size: usize, + /// Maximum size of dictionary page in bytes + max_dict_page_size: usize, + /// Writer version + writer_version: WriterVersion, + /// Enabled statistics + enabled_statistics: EnabledStatistics, + /// Encoding + encoding: Encoding, + /// row selections and total selected row count + row_selections: Option<(RowSelection, usize)>, + /// row filter + row_filter: Option>, + /// limit + limit: Option, + /// offset + offset: Option, + } + + /// Manually implement this to avoid printing entire contents of row_selections and row_filter + impl std::fmt::Debug for TestOptions { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TestOptions") + .field("num_row_groups", &self.num_row_groups) + .field("num_rows", &self.num_rows) + .field("record_batch_size", &self.record_batch_size) + .field("null_percent", &self.null_percent) + .field("write_batch_size", &self.write_batch_size) + .field("max_data_page_size", &self.max_data_page_size) + .field("max_dict_page_size", &self.max_dict_page_size) + .field("writer_version", &self.writer_version) + .field("enabled_statistics", &self.enabled_statistics) + .field("encoding", &self.encoding) + .field("row_selections", &self.row_selections.is_some()) + .field("row_filter", &self.row_filter.is_some()) + .field("limit", &self.limit) + .field("offset", &self.offset) + .finish() + } + } + + impl Default for TestOptions { + fn default() -> Self { + Self { + num_row_groups: 2, + num_rows: 100, + record_batch_size: 15, + null_percent: None, + write_batch_size: 64, + max_data_page_size: 1024 * 1024, + max_dict_page_size: 1024 * 1024, + writer_version: WriterVersion::PARQUET_1_0, + enabled_statistics: EnabledStatistics::Page, + encoding: Encoding::PLAIN, + row_selections: None, + row_filter: None, + limit: None, + offset: None, + } + } + } + + impl TestOptions { + fn new(num_row_groups: usize, num_rows: usize, record_batch_size: usize) -> Self { + Self { + num_row_groups, + num_rows, + record_batch_size, + ..Default::default() + } + } + + fn with_null_percent(self, null_percent: usize) -> Self { + Self { + null_percent: Some(null_percent), + ..self + } + } + + fn with_max_data_page_size(self, max_data_page_size: usize) -> Self { + Self { + max_data_page_size, + ..self + } + } + + fn with_max_dict_page_size(self, max_dict_page_size: usize) -> Self { + Self { + max_dict_page_size, + ..self + } + } + + fn with_enabled_statistics(self, enabled_statistics: EnabledStatistics) -> Self { + Self { + enabled_statistics, + ..self + } + } + + fn with_row_selections(self) -> Self { + assert!(self.row_filter.is_none(), "Must set row selection first"); + + let mut rng = rng(); + let step = rng.random_range(self.record_batch_size..self.num_rows); + let row_selections = create_test_selection( + step, + self.num_row_groups * self.num_rows, + rng.random::(), + ); + Self { + row_selections: Some(row_selections), + ..self + } + } + + fn with_row_filter(self) -> Self { + let row_count = match &self.row_selections { + Some((_, count)) => *count, + None => self.num_row_groups * self.num_rows, + }; + + let mut rng = rng(); + Self { + row_filter: Some((0..row_count).map(|_| rng.random_bool(0.9)).collect()), + ..self + } + } + + fn with_limit(self, limit: usize) -> Self { + Self { + limit: Some(limit), + ..self + } + } + + fn with_offset(self, offset: usize) -> Self { + Self { + offset: Some(offset), + ..self + } + } + + fn writer_props(&self) -> WriterProperties { + let builder = WriterProperties::builder() + .set_data_page_size_limit(self.max_data_page_size) + .set_write_batch_size(self.write_batch_size) + .set_writer_version(self.writer_version) + .set_statistics_enabled(self.enabled_statistics); + + let builder = match self.encoding { + Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY => builder + .set_dictionary_enabled(true) + .set_dictionary_page_size_limit(self.max_dict_page_size), + _ => builder + .set_dictionary_enabled(false) + .set_encoding(self.encoding), + }; + + builder.build() + } + } + + /// Create a parquet file and then read it using + /// `ParquetFileArrowReader` using a standard set of parameters + /// `opts`. + /// + /// `rand_max` represents the maximum size of value to pass to to + /// value generator + fn run_single_column_reader_tests( + rand_max: i32, + converted_type: ConvertedType, + arrow_type: Option, + converter: F, + encodings: &[Encoding], + ) where + T: DataType, + G: RandGen, + F: Fn(&[Option]) -> ArrayRef, + { + let all_options = vec![ + // choose record_batch_batch (15) so batches cross row + // group boundaries (50 rows in 2 row groups) cases. + TestOptions::new(2, 100, 15), + // choose record_batch_batch (5) so batches sometime fall + // on row group boundaries and (25 rows in 3 row groups + // --> row groups of 10, 10, and 5). Tests buffer + // refilling edge cases. + TestOptions::new(3, 25, 5), + // Choose record_batch_size (25) so all batches fall + // exactly on row group boundary (25). Tests buffer + // refilling edge cases. + TestOptions::new(4, 100, 25), + // Set maximum page size so row groups have multiple pages + TestOptions::new(3, 256, 73).with_max_data_page_size(128), + // Set small dictionary page size to test dictionary fallback + TestOptions::new(3, 256, 57).with_max_dict_page_size(128), + // Test optional but with no nulls + TestOptions::new(2, 256, 127).with_null_percent(0), + // Test optional with nulls + TestOptions::new(2, 256, 93).with_null_percent(25), + // Test with limit of 0 + TestOptions::new(4, 100, 25).with_limit(0), + // Test with limit of 50 + TestOptions::new(4, 100, 25).with_limit(50), + // Test with limit equal to number of rows + TestOptions::new(4, 100, 25).with_limit(10), + // Test with limit larger than number of rows + TestOptions::new(4, 100, 25).with_limit(101), + // Test with limit + offset equal to number of rows + TestOptions::new(4, 100, 25).with_offset(30).with_limit(20), + // Test with limit + offset equal to number of rows + TestOptions::new(4, 100, 25).with_offset(20).with_limit(80), + // Test with limit + offset larger than number of rows + TestOptions::new(4, 100, 25).with_offset(20).with_limit(81), + // Test with no page-level statistics + TestOptions::new(2, 256, 91) + .with_null_percent(25) + .with_enabled_statistics(EnabledStatistics::Chunk), + // Test with no statistics + TestOptions::new(2, 256, 91) + .with_null_percent(25) + .with_enabled_statistics(EnabledStatistics::None), + // Test with all null + TestOptions::new(2, 128, 91) + .with_null_percent(100) + .with_enabled_statistics(EnabledStatistics::None), + // Test skip + + // choose record_batch_batch (15) so batches cross row + // group boundaries (50 rows in 2 row groups) cases. + TestOptions::new(2, 100, 15).with_row_selections(), + // choose record_batch_batch (5) so batches sometime fall + // on row group boundaries and (25 rows in 3 row groups + // --> row groups of 10, 10, and 5). Tests buffer + // refilling edge cases. + TestOptions::new(3, 25, 5).with_row_selections(), + // Choose record_batch_size (25) so all batches fall + // exactly on row group boundary (25). Tests buffer + // refilling edge cases. + TestOptions::new(4, 100, 25).with_row_selections(), + // Set maximum page size so row groups have multiple pages + TestOptions::new(3, 256, 73) + .with_max_data_page_size(128) + .with_row_selections(), + // Set small dictionary page size to test dictionary fallback + TestOptions::new(3, 256, 57) + .with_max_dict_page_size(128) + .with_row_selections(), + // Test optional but with no nulls + TestOptions::new(2, 256, 127) + .with_null_percent(0) + .with_row_selections(), + // Test optional with nulls + TestOptions::new(2, 256, 93) + .with_null_percent(25) + .with_row_selections(), + // Test optional with nulls + TestOptions::new(2, 256, 93) + .with_null_percent(25) + .with_row_selections() + .with_limit(10), + // Test optional with nulls + TestOptions::new(2, 256, 93) + .with_null_percent(25) + .with_row_selections() + .with_offset(20) + .with_limit(10), + // Test filter + + // Test with row filter + TestOptions::new(4, 100, 25).with_row_filter(), + // Test with row selection and row filter + TestOptions::new(4, 100, 25) + .with_row_selections() + .with_row_filter(), + // Test with nulls and row filter + TestOptions::new(2, 256, 93) + .with_null_percent(25) + .with_max_data_page_size(10) + .with_row_filter(), + // Test with nulls and row filter and small pages + TestOptions::new(2, 256, 93) + .with_null_percent(25) + .with_max_data_page_size(10) + .with_row_selections() + .with_row_filter(), + // Test with row selection and no offset index and small pages + TestOptions::new(2, 256, 93) + .with_enabled_statistics(EnabledStatistics::None) + .with_max_data_page_size(10) + .with_row_selections(), + ]; + + all_options.into_iter().for_each(|opts| { + for writer_version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] { + for encoding in encodings { + let opts = TestOptions { + writer_version, + encoding: *encoding, + ..opts.clone() + }; + + single_column_reader_test::( + opts, + rand_max, + converted_type, + arrow_type.clone(), + &converter, + ) + } + } + }); + } + + /// Create a parquet file and then read it using + /// `ParquetFileArrowReader` using the parameters described in + /// `opts`. + fn single_column_reader_test( + opts: TestOptions, + rand_max: i32, + converted_type: ConvertedType, + arrow_type: Option, + converter: F, + ) where + T: DataType, + G: RandGen, + F: Fn(&[Option]) -> ArrayRef, + { + // Print out options to facilitate debugging failures on CI + println!( + "Running type {:?} single_column_reader_test ConvertedType::{}/ArrowType::{:?} with Options: {:?}", + T::get_physical_type(), + converted_type, + arrow_type, + opts + ); + + //according to null_percent generate def_levels + let (repetition, def_levels) = match opts.null_percent.as_ref() { + Some(null_percent) => { + let mut rng = rng(); + + let def_levels: Vec> = (0..opts.num_row_groups) + .map(|_| { + std::iter::from_fn(|| { + Some((rng.next_u32() as usize % 100 >= *null_percent) as i16) + }) + .take(opts.num_rows) + .collect() + }) + .collect(); + (Repetition::OPTIONAL, Some(def_levels)) + } + None => (Repetition::REQUIRED, None), + }; + + //generate random table data + let values: Vec> = (0..opts.num_row_groups) + .map(|idx| { + let null_count = match def_levels.as_ref() { + Some(d) => d[idx].iter().filter(|x| **x == 0).count(), + None => 0, + }; + G::gen_vec(rand_max, opts.num_rows - null_count) + }) + .collect(); + + let len = match T::get_physical_type() { + crate::basic::Type::FIXED_LEN_BYTE_ARRAY => rand_max, + crate::basic::Type::INT96 => 12, + _ => -1, + }; + + let fields = vec![Arc::new( + Type::primitive_type_builder("leaf", T::get_physical_type()) + .with_repetition(repetition) + .with_converted_type(converted_type) + .with_length(len) + .build() + .unwrap(), + )]; + + let schema = Arc::new( + Type::group_type_builder("test_schema") + .with_fields(fields) + .build() + .unwrap(), + ); + + let arrow_field = arrow_type.map(|t| Field::new("leaf", t, false)); + + let mut file = tempfile::tempfile().unwrap(); + + generate_single_column_file_with_data::( + &values, + def_levels.as_ref(), + file.try_clone().unwrap(), // Cannot use &mut File (#1163) + schema, + arrow_field, + &opts, + ) + .unwrap(); + + file.rewind().unwrap(); + + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::from( + opts.enabled_statistics == EnabledStatistics::Page, + )); + + let mut builder = + ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap(); + + let expected_data = match opts.row_selections { + Some((selections, row_count)) => { + let mut without_skip_data = gen_expected_data::(def_levels.as_ref(), &values); + + let mut skip_data: Vec> = vec![]; + let dequeue: VecDeque = selections.clone().into(); + for select in dequeue { + if select.skip { + without_skip_data.drain(0..select.row_count); + } else { + skip_data.extend(without_skip_data.drain(0..select.row_count)); + } + } + builder = builder.with_row_selection(selections); + + assert_eq!(skip_data.len(), row_count); + skip_data + } + None => { + //get flatten table data + let expected_data = gen_expected_data::(def_levels.as_ref(), &values); + assert_eq!(expected_data.len(), opts.num_rows * opts.num_row_groups); + expected_data + } + }; + + let mut expected_data = match opts.row_filter { + Some(filter) => { + let expected_data = expected_data + .into_iter() + .zip(filter.iter()) + .filter_map(|(d, f)| f.then(|| d)) + .collect(); + + let mut filter_offset = 0; + let filter = RowFilter::new(vec![Box::new(ArrowPredicateFn::new( + ProjectionMask::all(), + move |b| { + let array = BooleanArray::from_iter( + filter + .iter() + .skip(filter_offset) + .take(b.num_rows()) + .map(|x| Some(*x)), + ); + filter_offset += b.num_rows(); + Ok(array) + }, + ))]); + + builder = builder.with_row_filter(filter); + expected_data + } + None => expected_data, + }; + + if let Some(offset) = opts.offset { + builder = builder.with_offset(offset); + expected_data = expected_data.into_iter().skip(offset).collect(); + } + + if let Some(limit) = opts.limit { + builder = builder.with_limit(limit); + expected_data = expected_data.into_iter().take(limit).collect(); + } + + let mut record_reader = builder + .with_batch_size(opts.record_batch_size) + .build() + .unwrap(); + + let mut total_read = 0; + loop { + let maybe_batch = record_reader.next(); + if total_read < expected_data.len() { + let end = min(total_read + opts.record_batch_size, expected_data.len()); + let batch = maybe_batch.unwrap().unwrap(); + assert_eq!(end - total_read, batch.num_rows()); + + let a = converter(&expected_data[total_read..end]); + let b = batch.column(0); + + assert_eq!(a.data_type(), b.data_type()); + assert_eq!(a.to_data(), b.to_data()); + assert_eq!( + a.as_any().type_id(), + b.as_any().type_id(), + "incorrect type ids" + ); + + total_read = end; + } else { + assert!(maybe_batch.is_none()); + break; + } + } + } + + fn gen_expected_data( + def_levels: Option<&Vec>>, + values: &[Vec], + ) -> Vec> { + let data: Vec> = match def_levels { + Some(levels) => { + let mut values_iter = values.iter().flatten(); + levels + .iter() + .flatten() + .map(|d| match d { + 1 => Some(values_iter.next().cloned().unwrap()), + 0 => None, + _ => unreachable!(), + }) + .collect() + } + None => values.iter().flatten().cloned().map(Some).collect(), + }; + data + } + + fn generate_single_column_file_with_data( + values: &[Vec], + def_levels: Option<&Vec>>, + file: File, + schema: TypePtr, + field: Option, + opts: &TestOptions, + ) -> Result { + let mut writer_props = opts.writer_props(); + if let Some(field) = field { + let arrow_schema = Schema::new(vec![field]); + add_encoded_arrow_schema_to_metadata(&arrow_schema, &mut writer_props); + } + + let mut writer = SerializedFileWriter::new(file, schema, Arc::new(writer_props))?; + + for (idx, v) in values.iter().enumerate() { + let def_levels = def_levels.map(|d| d[idx].as_slice()); + let mut row_group_writer = writer.next_row_group()?; + { + let mut column_writer = row_group_writer + .next_column()? + .expect("Column writer is none!"); + + column_writer + .typed::() + .write_batch(v, def_levels, None)?; + + column_writer.close()?; + } + row_group_writer.close()?; + } + + writer.close() + } + + fn get_test_file(file_name: &str) -> File { + let mut path = PathBuf::new(); + path.push(arrow::util::test_util::arrow_test_data()); + path.push(file_name); + + File::open(path.as_path()).expect("File not found!") + } + + #[test] + fn test_read_structs() { + // This particular test file has columns of struct types where there is + // a column that has the same name as one of the struct fields + // (see: ARROW-11452) + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/nested_structs.rust.parquet"); + let file = File::open(&path).unwrap(); + let record_batch_reader = ParquetRecordBatchReader::try_new(file, 60).unwrap(); + + for batch in record_batch_reader { + batch.unwrap(); + } + + let file = File::open(&path).unwrap(); + let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + + let mask = ProjectionMask::leaves(builder.parquet_schema(), [3, 8, 10]); + let projected_reader = builder + .with_projection(mask) + .with_batch_size(60) + .build() + .unwrap(); + + let expected_schema = Schema::new(vec![ + Field::new( + "roll_num", + ArrowDataType::Struct(Fields::from(vec![Field::new( + "count", + ArrowDataType::UInt64, + false, + )])), + false, + ), + Field::new( + "PC_CUR", + ArrowDataType::Struct(Fields::from(vec![ + Field::new("mean", ArrowDataType::Int64, false), + Field::new("sum", ArrowDataType::Int64, false), + ])), + false, + ), + ]); + + // Tests for #1652 and #1654 + assert_eq!(&expected_schema, projected_reader.schema().as_ref()); + + for batch in projected_reader { + let batch = batch.unwrap(); + assert_eq!(batch.schema().as_ref(), &expected_schema); + } + } + + #[test] + // same as test_read_structs but constructs projection mask via column names + fn test_read_structs_by_name() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/nested_structs.rust.parquet"); + let file = File::open(&path).unwrap(); + let record_batch_reader = ParquetRecordBatchReader::try_new(file, 60).unwrap(); + + for batch in record_batch_reader { + batch.unwrap(); + } + + let file = File::open(&path).unwrap(); + let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + + let mask = ProjectionMask::columns( + builder.parquet_schema(), + ["roll_num.count", "PC_CUR.mean", "PC_CUR.sum"], + ); + let projected_reader = builder + .with_projection(mask) + .with_batch_size(60) + .build() + .unwrap(); + + let expected_schema = Schema::new(vec![ + Field::new( + "roll_num", + ArrowDataType::Struct(Fields::from(vec![Field::new( + "count", + ArrowDataType::UInt64, + false, + )])), + false, + ), + Field::new( + "PC_CUR", + ArrowDataType::Struct(Fields::from(vec![ + Field::new("mean", ArrowDataType::Int64, false), + Field::new("sum", ArrowDataType::Int64, false), + ])), + false, + ), + ]); + + assert_eq!(&expected_schema, projected_reader.schema().as_ref()); + + for batch in projected_reader { + let batch = batch.unwrap(); + assert_eq!(batch.schema().as_ref(), &expected_schema); + } + } + + #[test] + fn test_read_maps() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/nested_maps.snappy.parquet"); + let file = File::open(path).unwrap(); + let record_batch_reader = ParquetRecordBatchReader::try_new(file, 60).unwrap(); + + for batch in record_batch_reader { + batch.unwrap(); + } + } + + // test that we can handle the UNKNOWN logical type annotation on any physical type + #[test] + fn test_unknown_logical_type() { + let message_type = "message uk { + OPTIONAL INT32 uki32 (UNKNOWN); + OPTIONAL INT64 uki64 (UNKNOWN); + OPTIONAL INT96 uki96 (UNKNOWN); + OPTIONAL BOOLEAN ukbool (UNKNOWN); + OPTIONAL FLOAT ukfloat (UNKNOWN); + OPTIONAL DOUBLE ukdbl (UNKNOWN); + OPTIONAL BYTE_ARRAY ukbytes (UNKNOWN); + OPTIONAL FIXED_LEN_BYTE_ARRAY(10) ukflba (UNKNOWN); + }"; + + let schema = Arc::new(parse_message_type(message_type).unwrap()); + let file = tempfile::tempfile().unwrap(); + + let mut writer = + SerializedFileWriter::new(file.try_clone().unwrap(), schema, Default::default()) + .unwrap(); + + let mut row_group_writer = writer.next_row_group().unwrap(); + + fn write_nulls(row_group_writer: &mut SerializedRowGroupWriter<'_, File>) { + let mut column_writer = row_group_writer.next_column().unwrap().unwrap(); + // write out a bunch of nulls + column_writer + .typed::() + .write_batch(&[], Some(&[0, 0, 0, 0]), None) + .unwrap(); + column_writer.close().unwrap(); + } + + // INT32 + write_nulls::(&mut row_group_writer); + + // INT64 + write_nulls::(&mut row_group_writer); + + // INT96 + write_nulls::(&mut row_group_writer); + + // BOOLEAN + write_nulls::(&mut row_group_writer); + + // FLOAT + write_nulls::(&mut row_group_writer); + + // DOUBLE + write_nulls::(&mut row_group_writer); + + // BYTE_ARRAY + write_nulls::(&mut row_group_writer); + + // FIXED_LEN_BYTE_ARRAY + write_nulls::(&mut row_group_writer); + + row_group_writer.close().unwrap(); + + writer.close().unwrap(); + + let mut reader = ParquetRecordBatchReader::try_new(file, 4).unwrap(); + let batch = reader.next().unwrap().unwrap(); + + for col in batch.columns() { + assert_eq!(col.len(), 4); + assert_eq!(col.logical_null_count(), 4); + assert_eq!(*col.data_type(), ArrowDataType::Null); + } + } + + #[test] + fn test_nested_nullability() { + let message_type = "message nested { + OPTIONAL Group group { + REQUIRED INT32 leaf; + } + }"; + + let file = tempfile::tempfile().unwrap(); + let schema = Arc::new(parse_message_type(message_type).unwrap()); + + { + // Write using low-level parquet API (#1167) + let mut writer = + SerializedFileWriter::new(file.try_clone().unwrap(), schema, Default::default()) + .unwrap(); + + { + let mut row_group_writer = writer.next_row_group().unwrap(); + let mut column_writer = row_group_writer.next_column().unwrap().unwrap(); + + column_writer + .typed::() + .write_batch(&[34, 76], Some(&[0, 1, 0, 1]), None) + .unwrap(); + + column_writer.close().unwrap(); + row_group_writer.close().unwrap(); + } + + writer.close().unwrap(); + } + + let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + let mask = ProjectionMask::leaves(builder.parquet_schema(), [0]); + + let reader = builder.with_projection(mask).build().unwrap(); + + let expected_schema = Schema::new(vec![Field::new( + "group", + ArrowDataType::Struct(vec![Field::new("leaf", ArrowDataType::Int32, false)].into()), + true, + )]); + + let batch = reader.into_iter().next().unwrap().unwrap(); + assert_eq!(batch.schema().as_ref(), &expected_schema); + assert_eq!(batch.num_rows(), 4); + assert_eq!(batch.column(0).null_count(), 2); + } + + #[test] + fn test_dictionary_preservation() { + let fields = vec![Arc::new( + Type::primitive_type_builder("leaf", PhysicalType::BYTE_ARRAY) + .with_repetition(Repetition::OPTIONAL) + .with_converted_type(ConvertedType::UTF8) + .build() + .unwrap(), + )]; + + let schema = Arc::new( + Type::group_type_builder("test_schema") + .with_fields(fields) + .build() + .unwrap(), + ); + + let dict_type = ArrowDataType::Dictionary( + Box::new(ArrowDataType::Int32), + Box::new(ArrowDataType::Utf8), + ); + + let arrow_field = Field::new("leaf", dict_type, true); + + let mut file = tempfile::tempfile().unwrap(); + + let values = vec![ + vec![ + ByteArray::from("hello"), + ByteArray::from("a"), + ByteArray::from("b"), + ByteArray::from("d"), + ], + vec![ + ByteArray::from("c"), + ByteArray::from("a"), + ByteArray::from("b"), + ], + ]; + + let def_levels = vec![ + vec![1, 0, 0, 1, 0, 0, 1, 1], + vec![0, 0, 1, 1, 0, 0, 1, 0, 0], + ]; + + let opts = TestOptions { + encoding: Encoding::RLE_DICTIONARY, + ..Default::default() + }; + + generate_single_column_file_with_data::( + &values, + Some(&def_levels), + file.try_clone().unwrap(), // Cannot use &mut File (#1163) + schema, + Some(arrow_field), + &opts, + ) + .unwrap(); + + file.rewind().unwrap(); + + let record_reader = ParquetRecordBatchReader::try_new(file, 3).unwrap(); + + let batches = record_reader + .collect::, _>>() + .unwrap(); + + assert_eq!(batches.len(), 6); + assert!(batches.iter().all(|x| x.num_columns() == 1)); + + let row_counts = batches + .iter() + .map(|x| (x.num_rows(), x.column(0).null_count())) + .collect::>(); + + assert_eq!( + row_counts, + vec![(3, 2), (3, 2), (3, 1), (3, 1), (3, 2), (2, 2)] + ); + + let get_dict = |batch: &RecordBatch| batch.column(0).to_data().child_data()[0].clone(); + + // First and second batch in same row group -> same dictionary + assert_eq!(get_dict(&batches[0]), get_dict(&batches[1])); + // Third batch spans row group -> computed dictionary + assert_ne!(get_dict(&batches[1]), get_dict(&batches[2])); + assert_ne!(get_dict(&batches[2]), get_dict(&batches[3])); + // Fourth, fifth and sixth from same row group -> same dictionary + assert_eq!(get_dict(&batches[3]), get_dict(&batches[4])); + assert_eq!(get_dict(&batches[4]), get_dict(&batches[5])); + } + + #[test] + fn test_read_null_list() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/null_list.parquet"); + let file = File::open(path).unwrap(); + let mut record_batch_reader = ParquetRecordBatchReader::try_new(file, 60).unwrap(); + + let batch = record_batch_reader.next().unwrap().unwrap(); + assert_eq!(batch.num_rows(), 1); + assert_eq!(batch.num_columns(), 1); + assert_eq!(batch.column(0).len(), 1); + + let list = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(list.len(), 1); + assert!(list.is_valid(0)); + + let val = list.value(0); + assert_eq!(val.len(), 0); + } + + #[test] + fn test_null_schema_inference() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/null_list.parquet"); + let file = File::open(path).unwrap(); + + let arrow_field = Field::new( + "emptylist", + ArrowDataType::List(Arc::new(Field::new_list_field(ArrowDataType::Null, true))), + true, + ); + + let options = ArrowReaderOptions::new().with_skip_arrow_metadata(true); + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap(); + let schema = builder.schema(); + assert_eq!(schema.fields().len(), 1); + assert_eq!(schema.field(0), &arrow_field); + } + + #[test] + fn test_skip_metadata() { + let col = Arc::new(TimestampNanosecondArray::from_iter_values(vec![0, 1, 2])); + let field = Field::new("col", col.data_type().clone(), true); + + let schema_without_metadata = Arc::new(Schema::new(vec![field.clone()])); + + let metadata = [("key".to_string(), "value".to_string())] + .into_iter() + .collect(); + + let schema_with_metadata = Arc::new(Schema::new(vec![field.with_metadata(metadata)])); + + assert_ne!(schema_with_metadata, schema_without_metadata); + + let batch = + RecordBatch::try_new(schema_with_metadata.clone(), vec![col as ArrayRef]).unwrap(); + + let file = |version: WriterVersion| { + let props = WriterProperties::builder() + .set_writer_version(version) + .build(); + + let file = tempfile().unwrap(); + let mut writer = + ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)) + .unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + file + }; + + let skip_options = ArrowReaderOptions::new().with_skip_arrow_metadata(true); + + let v1_reader = file(WriterVersion::PARQUET_1_0); + let v2_reader = file(WriterVersion::PARQUET_2_0); + + let arrow_reader = + ParquetRecordBatchReader::try_new(v1_reader.try_clone().unwrap(), 1024).unwrap(); + assert_eq!(arrow_reader.schema(), schema_with_metadata); + + let reader = + ParquetRecordBatchReaderBuilder::try_new_with_options(v1_reader, skip_options.clone()) + .unwrap() + .build() + .unwrap(); + assert_eq!(reader.schema(), schema_without_metadata); + + let arrow_reader = + ParquetRecordBatchReader::try_new(v2_reader.try_clone().unwrap(), 1024).unwrap(); + assert_eq!(arrow_reader.schema(), schema_with_metadata); + + let reader = ParquetRecordBatchReaderBuilder::try_new_with_options(v2_reader, skip_options) + .unwrap() + .build() + .unwrap(); + assert_eq!(reader.schema(), schema_without_metadata); + } + + fn write_parquet_from_iter(value: I) -> File + where + I: IntoIterator, + F: AsRef, + { + let batch = RecordBatch::try_from_iter(value).unwrap(); + let file = tempfile().unwrap(); + let mut writer = + ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema().clone(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + file + } + + fn run_schema_test_with_error(value: I, schema: SchemaRef, expected_error: &str) + where + I: IntoIterator, + F: AsRef, + { + let file = write_parquet_from_iter(value); + let options_with_schema = ArrowReaderOptions::new().with_schema(schema.clone()); + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options( + file.try_clone().unwrap(), + options_with_schema, + ); + assert_eq!(builder.err().unwrap().to_string(), expected_error); + } + + #[test] + fn test_schema_too_few_columns() { + run_schema_test_with_error( + vec![ + ("int64", Arc::new(Int64Array::from(vec![0])) as ArrayRef), + ("int32", Arc::new(Int32Array::from(vec![0])) as ArrayRef), + ], + Arc::new(Schema::new(vec![Field::new( + "int64", + ArrowDataType::Int64, + false, + )])), + "Arrow: incompatible arrow schema, expected 2 struct fields got 1", + ); + } + + #[test] + fn test_schema_too_many_columns() { + run_schema_test_with_error( + vec![("int64", Arc::new(Int64Array::from(vec![0])) as ArrayRef)], + Arc::new(Schema::new(vec![ + Field::new("int64", ArrowDataType::Int64, false), + Field::new("int32", ArrowDataType::Int32, false), + ])), + "Arrow: incompatible arrow schema, expected 1 struct fields got 2", + ); + } + + #[test] + fn test_schema_mismatched_column_names() { + run_schema_test_with_error( + vec![("int64", Arc::new(Int64Array::from(vec![0])) as ArrayRef)], + Arc::new(Schema::new(vec![Field::new( + "other", + ArrowDataType::Int64, + false, + )])), + "Arrow: incompatible arrow schema, expected field named int64 got other", + ); + } + + #[test] + fn test_schema_incompatible_columns() { + run_schema_test_with_error( + vec![ + ( + "col1_invalid", + Arc::new(Int64Array::from(vec![0])) as ArrayRef, + ), + ( + "col2_valid", + Arc::new(Int32Array::from(vec![0])) as ArrayRef, + ), + ( + "col3_invalid", + Arc::new(Date64Array::from(vec![0])) as ArrayRef, + ), + ], + Arc::new(Schema::new(vec![ + Field::new("col1_invalid", ArrowDataType::Int32, false), + Field::new("col2_valid", ArrowDataType::Int32, false), + Field::new("col3_invalid", ArrowDataType::Int32, false), + ])), + "Arrow: Incompatible supplied Arrow schema: data type mismatch for field col1_invalid: requested Int32 but found Int64, data type mismatch for field col3_invalid: requested Int32 but found Int64", + ); + } + + #[test] + fn test_one_incompatible_nested_column() { + let nested_fields = Fields::from(vec![ + Field::new("nested1_valid", ArrowDataType::Utf8, false), + Field::new("nested1_invalid", ArrowDataType::Int64, false), + ]); + let nested = StructArray::try_new( + nested_fields, + vec![ + Arc::new(StringArray::from(vec!["a"])) as ArrayRef, + Arc::new(Int64Array::from(vec![0])) as ArrayRef, + ], + None, + ) + .expect("struct array"); + let supplied_nested_fields = Fields::from(vec![ + Field::new("nested1_valid", ArrowDataType::Utf8, false), + Field::new("nested1_invalid", ArrowDataType::Int32, false), + ]); + run_schema_test_with_error( + vec![ + ("col1", Arc::new(Int64Array::from(vec![0])) as ArrayRef), + ("col2", Arc::new(Int32Array::from(vec![0])) as ArrayRef), + ("nested", Arc::new(nested) as ArrayRef), + ], + Arc::new(Schema::new(vec![ + Field::new("col1", ArrowDataType::Int64, false), + Field::new("col2", ArrowDataType::Int32, false), + Field::new( + "nested", + ArrowDataType::Struct(supplied_nested_fields), + false, + ), + ])), + "Arrow: Incompatible supplied Arrow schema: data type mismatch for field nested: \ + requested Struct(\"nested1_valid\": non-null Utf8, \"nested1_invalid\": non-null Int32) \ + but found Struct(\"nested1_valid\": non-null Utf8, \"nested1_invalid\": non-null Int64)", + ); + } + + /// Return parquet data with a single column of utf8 strings + fn utf8_parquet() -> Bytes { + let input = StringArray::from_iter_values(vec!["foo", "bar", "baz"]); + let batch = RecordBatch::try_from_iter(vec![("column1", Arc::new(input) as _)]).unwrap(); + let props = None; + // write parquet file with non nullable strings + let mut parquet_data = vec![]; + let mut writer = ArrowWriter::try_new(&mut parquet_data, batch.schema(), props).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + Bytes::from(parquet_data) + } + + #[test] + fn test_schema_error_bad_types() { + // verify incompatible schemas error on read + let parquet_data = utf8_parquet(); + + // Ask to read it back with an incompatible schema (int vs string) + let input_schema: SchemaRef = Arc::new(Schema::new(vec![Field::new( + "column1", + arrow::datatypes::DataType::Int32, + false, + )])); + + // read it back out + let reader_options = ArrowReaderOptions::new().with_schema(input_schema.clone()); + let err = + ParquetRecordBatchReaderBuilder::try_new_with_options(parquet_data, reader_options) + .unwrap_err(); + assert_eq!( + err.to_string(), + "Arrow: Incompatible supplied Arrow schema: data type mismatch for field column1: requested Int32 but found Utf8" + ) + } + + #[test] + fn test_schema_error_bad_nullability() { + // verify incompatible schemas error on read + let parquet_data = utf8_parquet(); + + // Ask to read it back with an incompatible schema (nullability mismatch) + let input_schema: SchemaRef = Arc::new(Schema::new(vec![Field::new( + "column1", + arrow::datatypes::DataType::Utf8, + true, + )])); + + // read it back out + let reader_options = ArrowReaderOptions::new().with_schema(input_schema.clone()); + let err = + ParquetRecordBatchReaderBuilder::try_new_with_options(parquet_data, reader_options) + .unwrap_err(); + assert_eq!( + err.to_string(), + "Arrow: Incompatible supplied Arrow schema: nullability mismatch for field column1: expected true but found false" + ) + } + + #[test] + fn test_read_binary_as_utf8() { + let file = write_parquet_from_iter(vec![ + ( + "binary_to_utf8", + Arc::new(BinaryArray::from(vec![ + b"one".as_ref(), + b"two".as_ref(), + b"three".as_ref(), + ])) as ArrayRef, + ), + ( + "large_binary_to_large_utf8", + Arc::new(LargeBinaryArray::from(vec![ + b"one".as_ref(), + b"two".as_ref(), + b"three".as_ref(), + ])) as ArrayRef, + ), + ( + "binary_view_to_utf8_view", + Arc::new(BinaryViewArray::from(vec![ + b"one".as_ref(), + b"two".as_ref(), + b"three".as_ref(), + ])) as ArrayRef, + ), + ]); + let supplied_fields = Fields::from(vec![ + Field::new("binary_to_utf8", ArrowDataType::Utf8, false), + Field::new( + "large_binary_to_large_utf8", + ArrowDataType::LargeUtf8, + false, + ), + Field::new("binary_view_to_utf8_view", ArrowDataType::Utf8View, false), + ]); + + let options = ArrowReaderOptions::new().with_schema(Arc::new(Schema::new(supplied_fields))); + let mut arrow_reader = ParquetRecordBatchReaderBuilder::try_new_with_options( + file.try_clone().unwrap(), + options, + ) + .expect("reader builder with schema") + .build() + .expect("reader with schema"); + + let batch = arrow_reader.next().unwrap().unwrap(); + assert_eq!(batch.num_columns(), 3); + assert_eq!(batch.num_rows(), 3); + assert_eq!( + batch + .column(0) + .as_string::() + .iter() + .collect::>(), + vec![Some("one"), Some("two"), Some("three")] + ); + + assert_eq!( + batch + .column(1) + .as_string::() + .iter() + .collect::>(), + vec![Some("one"), Some("two"), Some("three")] + ); + + assert_eq!( + batch.column(2).as_string_view().iter().collect::>(), + vec![Some("one"), Some("two"), Some("three")] + ); + } + + #[test] + #[should_panic(expected = "Invalid UTF8 sequence at")] + fn test_read_non_utf8_binary_as_utf8() { + let file = write_parquet_from_iter(vec![( + "non_utf8_binary", + Arc::new(BinaryArray::from(vec![ + b"\xDE\x00\xFF".as_ref(), + b"\xDE\x01\xAA".as_ref(), + b"\xDE\x02\xFF".as_ref(), + ])) as ArrayRef, + )]); + let supplied_fields = Fields::from(vec![Field::new( + "non_utf8_binary", + ArrowDataType::Utf8, + false, + )]); + + let options = ArrowReaderOptions::new().with_schema(Arc::new(Schema::new(supplied_fields))); + let mut arrow_reader = ParquetRecordBatchReaderBuilder::try_new_with_options( + file.try_clone().unwrap(), + options, + ) + .expect("reader builder with schema") + .build() + .expect("reader with schema"); + arrow_reader.next().unwrap().unwrap_err(); + } + + #[test] + fn test_with_schema() { + let nested_fields = Fields::from(vec![ + Field::new("utf8_to_dict", ArrowDataType::Utf8, false), + Field::new("int64_to_ts_nano", ArrowDataType::Int64, false), + ]); + + let nested_arrays: Vec = vec![ + Arc::new(StringArray::from(vec!["a", "a", "a", "b"])) as ArrayRef, + Arc::new(Int64Array::from(vec![1, 2, 3, 4])) as ArrayRef, + ]; + + let nested = StructArray::try_new(nested_fields, nested_arrays, None).unwrap(); + + let file = write_parquet_from_iter(vec![ + ( + "int32_to_ts_second", + Arc::new(Int32Array::from(vec![0, 1, 2, 3])) as ArrayRef, + ), + ( + "date32_to_date64", + Arc::new(Date32Array::from(vec![0, 1, 2, 3])) as ArrayRef, + ), + ("nested", Arc::new(nested) as ArrayRef), + ]); + + let supplied_nested_fields = Fields::from(vec![ + Field::new( + "utf8_to_dict", + ArrowDataType::Dictionary( + Box::new(ArrowDataType::Int32), + Box::new(ArrowDataType::Utf8), + ), + false, + ), + Field::new( + "int64_to_ts_nano", + ArrowDataType::Timestamp( + arrow::datatypes::TimeUnit::Nanosecond, + Some("+10:00".into()), + ), + false, + ), + ]); + + let supplied_schema = Arc::new(Schema::new(vec![ + Field::new( + "int32_to_ts_second", + ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Second, Some("+01:00".into())), + false, + ), + Field::new("date32_to_date64", ArrowDataType::Date64, false), + Field::new( + "nested", + ArrowDataType::Struct(supplied_nested_fields), + false, + ), + ])); + + let options = ArrowReaderOptions::new().with_schema(supplied_schema.clone()); + let mut arrow_reader = ParquetRecordBatchReaderBuilder::try_new_with_options( + file.try_clone().unwrap(), + options, + ) + .expect("reader builder with schema") + .build() + .expect("reader with schema"); + + assert_eq!(arrow_reader.schema(), supplied_schema); + let batch = arrow_reader.next().unwrap().unwrap(); + assert_eq!(batch.num_columns(), 3); + assert_eq!(batch.num_rows(), 4); + assert_eq!( + batch + .column(0) + .as_any() + .downcast_ref::() + .expect("downcast to timestamp second") + .value_as_datetime_with_tz(0, "+01:00".parse().unwrap()) + .map(|v| v.to_string()) + .expect("value as datetime"), + "1970-01-01 01:00:00 +01:00" + ); + assert_eq!( + batch + .column(1) + .as_any() + .downcast_ref::() + .expect("downcast to date64") + .value_as_date(0) + .map(|v| v.to_string()) + .expect("value as date"), + "1970-01-01" + ); + + let nested = batch + .column(2) + .as_any() + .downcast_ref::() + .expect("downcast to struct"); + + let nested_dict = nested + .column(0) + .as_any() + .downcast_ref::() + .expect("downcast to dictionary"); + + assert_eq!( + nested_dict + .values() + .as_any() + .downcast_ref::() + .expect("downcast to string") + .iter() + .collect::>(), + vec![Some("a"), Some("b")] + ); + + assert_eq!( + nested_dict.keys().iter().collect::>(), + vec![Some(0), Some(0), Some(0), Some(1)] + ); + + assert_eq!( + nested + .column(1) + .as_any() + .downcast_ref::() + .expect("downcast to timestamp nanosecond") + .value_as_datetime_with_tz(0, "+10:00".parse().unwrap()) + .map(|v| v.to_string()) + .expect("value as datetime"), + "1970-01-01 10:00:00.000000001 +10:00" + ); + } + + #[test] + fn test_empty_projection() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_plain.parquet"); + let file = File::open(path).unwrap(); + + let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + let file_metadata = builder.metadata().file_metadata(); + let expected_rows = file_metadata.num_rows() as usize; + + let mask = ProjectionMask::leaves(builder.parquet_schema(), []); + let batch_reader = builder + .with_projection(mask) + .with_batch_size(2) + .build() + .unwrap(); + + let mut total_rows = 0; + for maybe_batch in batch_reader { + let batch = maybe_batch.unwrap(); + total_rows += batch.num_rows(); + assert_eq!(batch.num_columns(), 0); + assert!(batch.num_rows() <= 2); + } + + assert_eq!(total_rows, expected_rows); + } + + fn test_row_group_batch(row_group_size: usize, batch_size: usize) { + let schema = Arc::new(Schema::new(vec![Field::new( + "list", + ArrowDataType::List(Arc::new(Field::new_list_field(ArrowDataType::Int32, true))), + true, + )])); + + let mut buf = Vec::with_capacity(1024); + + let mut writer = ArrowWriter::try_new( + &mut buf, + schema.clone(), + Some( + WriterProperties::builder() + .set_max_row_group_row_count(Some(row_group_size)) + .build(), + ), + ) + .unwrap(); + for _ in 0..2 { + let mut list_builder = ListBuilder::new(Int32Builder::with_capacity(batch_size)); + for _ in 0..(batch_size) { + list_builder.append(true); + } + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(list_builder.finish())]) + .unwrap(); + writer.write(&batch).unwrap(); + } + writer.close().unwrap(); + + let mut record_reader = + ParquetRecordBatchReader::try_new(Bytes::from(buf), batch_size).unwrap(); + assert_eq!( + batch_size, + record_reader.next().unwrap().unwrap().num_rows() + ); + assert_eq!( + batch_size, + record_reader.next().unwrap().unwrap().num_rows() + ); + } + + #[test] + fn test_row_group_exact_multiple() { + const BATCH_SIZE: usize = REPETITION_LEVELS_BATCH_SIZE; + test_row_group_batch(8, 8); + test_row_group_batch(10, 8); + test_row_group_batch(8, 10); + test_row_group_batch(BATCH_SIZE, BATCH_SIZE); + test_row_group_batch(BATCH_SIZE + 1, BATCH_SIZE); + test_row_group_batch(BATCH_SIZE, BATCH_SIZE + 1); + test_row_group_batch(BATCH_SIZE, BATCH_SIZE - 1); + test_row_group_batch(BATCH_SIZE - 1, BATCH_SIZE); + } + + /// Given a RecordBatch containing all the column data, return the expected batches given + /// a `batch_size` and `selection` + fn get_expected_batches( + column: &RecordBatch, + selection: &RowSelection, + batch_size: usize, + ) -> Vec { + let mut expected_batches = vec![]; + + let mut selection: VecDeque<_> = selection.clone().into(); + let mut row_offset = 0; + let mut last_start = None; + while row_offset < column.num_rows() && !selection.is_empty() { + let mut batch_remaining = batch_size.min(column.num_rows() - row_offset); + while batch_remaining > 0 && !selection.is_empty() { + let (to_read, skip) = match selection.front_mut() { + Some(selection) if selection.row_count > batch_remaining => { + selection.row_count -= batch_remaining; + (batch_remaining, selection.skip) + } + Some(_) => { + let select = selection.pop_front().unwrap(); + (select.row_count, select.skip) + } + None => break, + }; + + batch_remaining -= to_read; + + match skip { + true => { + 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 + } + } + } + } + + if let Some(last_start) = last_start.take() { + expected_batches.push(column.slice(last_start, row_offset - last_start)) + } + + // Sanity check, all batches except the final should be the batch size + for batch in &expected_batches[..expected_batches.len() - 1] { + assert_eq!(batch.num_rows(), batch_size); + } + + expected_batches + } + + fn create_test_selection( + step_len: usize, + total_len: usize, + skip_first: bool, + ) -> (RowSelection, usize) { + let mut remaining = total_len; + let mut skip = skip_first; + let mut vec = vec![]; + let mut selected_count = 0; + while remaining != 0 { + let step = if remaining > step_len { + step_len + } else { + remaining + }; + vec.push(RowSelector { + row_count: step, + skip, + }); + remaining -= step; + if !skip { + selected_count += step; + } + skip = !skip; + } + (vec.into(), selected_count) + } + + #[test] + fn test_scan_row_with_selection() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet"); + let test_file = File::open(&path).unwrap(); + + let mut serial_reader = + ParquetRecordBatchReader::try_new(File::open(&path).unwrap(), 7300).unwrap(); + let data = serial_reader.next().unwrap().unwrap(); + + let do_test = |batch_size: usize, selection_len: usize| { + for skip_first in [false, true] { + let selections = create_test_selection(batch_size, data.num_rows(), skip_first).0; + + let expected = get_expected_batches(&data, &selections, batch_size); + let skip_reader = create_skip_reader(&test_file, batch_size, selections); + assert_eq!( + skip_reader.collect::, _>>().unwrap(), + expected, + "batch_size: {batch_size}, selection_len: {selection_len}, skip_first: {skip_first}" + ); + } + }; + + // total row count 7300 + // 1. test selection len more than one page row count + do_test(1000, 1000); + + // 2. test selection len less than one page row count + do_test(20, 20); + + // 3. test selection_len less than batch_size + do_test(20, 5); + + // 4. test selection_len more than batch_size + // If batch_size < selection_len + do_test(20, 5); + + fn create_skip_reader( + test_file: &File, + batch_size: usize, + selections: RowSelection, + ) -> ParquetRecordBatchReader { + let options = + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let file = test_file.try_clone().unwrap(); + ParquetRecordBatchReaderBuilder::try_new_with_options(file, options) + .unwrap() + .with_batch_size(batch_size) + .with_row_selection(selections) + .build() + .unwrap() + } + } + + #[test] + fn test_batch_size_overallocate() { + let testdata = arrow::util::test_util::parquet_test_data(); + // `alltypes_plain.parquet` only have 8 rows + let path = format!("{testdata}/alltypes_plain.parquet"); + let test_file = File::open(path).unwrap(); + + let builder = ParquetRecordBatchReaderBuilder::try_new(test_file).unwrap(); + let num_rows = builder.metadata.file_metadata().num_rows(); + let reader = builder + .with_batch_size(1024) + .with_projection(ProjectionMask::all()) + .build() + .unwrap(); + assert_ne!(1024, num_rows); + assert_eq!(reader.read_plan.batch_size(), num_rows as usize); + } + + #[test] + fn test_read_with_page_index_enabled() { + let testdata = arrow::util::test_util::parquet_test_data(); + + { + // `alltypes_tiny_pages.parquet` has page index + let path = format!("{testdata}/alltypes_tiny_pages.parquet"); + let test_file = File::open(path).unwrap(); + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options( + test_file, + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required), + ) + .unwrap(); + assert!(!builder.metadata().offset_index().unwrap()[0].is_empty()); + let reader = builder.build().unwrap(); + let batches = reader.collect::, _>>().unwrap(); + assert_eq!(batches.len(), 8); + } + + { + // `alltypes_plain.parquet` doesn't have page index + let path = format!("{testdata}/alltypes_plain.parquet"); + let test_file = File::open(path).unwrap(); + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options( + test_file, + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required), + ) + .unwrap(); + // Although `Vec>` of each row group is empty, + // we should read the file successfully. + assert!(builder.metadata().offset_index().is_none()); + let reader = builder.build().unwrap(); + let batches = reader.collect::, _>>().unwrap(); + assert_eq!(batches.len(), 1); + } + } + + #[test] + fn test_raw_repetition() { + const MESSAGE_TYPE: &str = " + message Log { + OPTIONAL INT32 eventType; + REPEATED INT32 category; + REPEATED group filter { + OPTIONAL INT32 error; + } + } + "; + let schema = Arc::new(parse_message_type(MESSAGE_TYPE).unwrap()); + let props = Default::default(); + + let mut buf = Vec::with_capacity(1024); + let mut writer = SerializedFileWriter::new(&mut buf, schema, props).unwrap(); + let mut row_group_writer = writer.next_row_group().unwrap(); + + // column 0 + let mut col_writer = row_group_writer.next_column().unwrap().unwrap(); + col_writer + .typed::() + .write_batch(&[1], Some(&[1]), None) + .unwrap(); + col_writer.close().unwrap(); + // column 1 + let mut col_writer = row_group_writer.next_column().unwrap().unwrap(); + col_writer + .typed::() + .write_batch(&[1, 1], Some(&[1, 1]), Some(&[0, 1])) + .unwrap(); + col_writer.close().unwrap(); + // column 2 + let mut col_writer = row_group_writer.next_column().unwrap().unwrap(); + col_writer + .typed::() + .write_batch(&[1], Some(&[1]), Some(&[0])) + .unwrap(); + col_writer.close().unwrap(); + + let rg_md = row_group_writer.close().unwrap(); + assert_eq!(rg_md.num_rows(), 1); + writer.close().unwrap(); + + let bytes = Bytes::from(buf); + + let mut no_mask = ParquetRecordBatchReader::try_new(bytes.clone(), 1024).unwrap(); + let full = no_mask.next().unwrap().unwrap(); + + assert_eq!(full.num_columns(), 3); + + for idx in 0..3 { + let b = ParquetRecordBatchReaderBuilder::try_new(bytes.clone()).unwrap(); + let mask = ProjectionMask::leaves(b.parquet_schema(), [idx]); + let mut reader = b.with_projection(mask).build().unwrap(); + let projected = reader.next().unwrap().unwrap(); + + assert_eq!(projected.num_columns(), 1); + assert_eq!(full.column(idx), projected.column(0)); + } + } + + #[test] + fn test_read_lz4_raw() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/lz4_raw_compressed.parquet"); + let file = File::open(path).unwrap(); + + let batches = ParquetRecordBatchReader::try_new(file, 1024) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.len(), 1); + let batch = &batches[0]; + + assert_eq!(batch.num_columns(), 3); + assert_eq!(batch.num_rows(), 4); + + // https://github.com/apache/parquet-testing/pull/18 + let a: &Int64Array = batch.column(0).as_any().downcast_ref().unwrap(); + assert_eq!( + a.values(), + &[1593604800, 1593604800, 1593604801, 1593604801] + ); + + let a: &BinaryArray = batch.column(1).as_any().downcast_ref().unwrap(); + let a: Vec<_> = a.iter().flatten().collect(); + assert_eq!(a, &[b"abc", b"def", b"abc", b"def"]); + + let a: &Float64Array = batch.column(2).as_any().downcast_ref().unwrap(); + assert_eq!(a.values(), &[42.000000, 7.700000, 42.125000, 7.700000]); + } + + // This test is to ensure backward compatibility, it test 2 files containing the LZ4 CompressionCodec + // but different algorithms: LZ4_HADOOP and LZ4_RAW. + // 1. hadoop_lz4_compressed.parquet -> It is a file with LZ4 CompressionCodec which uses + // LZ4_HADOOP algorithm for compression. + // 2. non_hadoop_lz4_compressed.parquet -> It is a file with LZ4 CompressionCodec which uses + // LZ4_RAW algorithm for compression. This fallback is done to keep backward compatibility with + // older parquet-cpp versions. + // + // For more information, check: https://github.com/apache/arrow-rs/issues/2988 + #[test] + fn test_read_lz4_hadoop_fallback() { + for file in [ + "hadoop_lz4_compressed.parquet", + "non_hadoop_lz4_compressed.parquet", + ] { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/{file}"); + let file = File::open(path).unwrap(); + let expected_rows = 4; + + let batches = ParquetRecordBatchReader::try_new(file, expected_rows) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.len(), 1); + let batch = &batches[0]; + + assert_eq!(batch.num_columns(), 3); + assert_eq!(batch.num_rows(), expected_rows); + + let a: &Int64Array = batch.column(0).as_any().downcast_ref().unwrap(); + assert_eq!( + a.values(), + &[1593604800, 1593604800, 1593604801, 1593604801] + ); + + let b: &BinaryArray = batch.column(1).as_any().downcast_ref().unwrap(); + let b: Vec<_> = b.iter().flatten().collect(); + assert_eq!(b, &[b"abc", b"def", b"abc", b"def"]); + + let c: &Float64Array = batch.column(2).as_any().downcast_ref().unwrap(); + assert_eq!(c.values(), &[42.0, 7.7, 42.125, 7.7]); + } + } + + #[test] + fn test_read_lz4_hadoop_large() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/hadoop_lz4_compressed_larger.parquet"); + let file = File::open(path).unwrap(); + let expected_rows = 10000; + + let batches = ParquetRecordBatchReader::try_new(file, expected_rows) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.len(), 1); + let batch = &batches[0]; + + assert_eq!(batch.num_columns(), 1); + assert_eq!(batch.num_rows(), expected_rows); + + let a: &StringArray = batch.column(0).as_any().downcast_ref().unwrap(); + let a: Vec<_> = a.iter().flatten().collect(); + assert_eq!(a[0], "c7ce6bef-d5b0-4863-b199-8ea8c7fb117b"); + assert_eq!(a[1], "e8fb9197-cb9f-4118-b67f-fbfa65f61843"); + assert_eq!(a[expected_rows - 2], "ab52a0cc-c6bb-4d61-8a8f-166dc4b8b13c"); + assert_eq!(a[expected_rows - 1], "85440778-460a-41ac-aa2e-ac3ee41696bf"); + } + + #[test] + #[cfg(feature = "snap")] + fn test_read_nested_lists() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/nested_lists.snappy.parquet"); + let file = File::open(path).unwrap(); + + let f = file.try_clone().unwrap(); + let mut reader = ParquetRecordBatchReader::try_new(f, 60).unwrap(); + let expected = reader.next().unwrap().unwrap(); + assert_eq!(expected.num_rows(), 3); + + let selection = RowSelection::from(vec![ + RowSelector::skip(1), + RowSelector::select(1), + RowSelector::skip(1), + ]); + let mut reader = ParquetRecordBatchReaderBuilder::try_new(file) + .unwrap() + .with_row_selection(selection) + .build() + .unwrap(); + + let actual = reader.next().unwrap().unwrap(); + assert_eq!(actual.num_rows(), 1); + assert_eq!(actual.column(0), &expected.column(0).slice(1, 1)); + } + + #[test] + fn test_arbitrary_decimal() { + let values = [1, 2, 3, 4, 5, 6, 7, 8]; + let decimals_19_0 = Decimal128Array::from_iter_values(values) + .with_precision_and_scale(19, 0) + .unwrap(); + let decimals_12_0 = Decimal128Array::from_iter_values(values) + .with_precision_and_scale(12, 0) + .unwrap(); + let decimals_17_10 = Decimal128Array::from_iter_values(values) + .with_precision_and_scale(17, 10) + .unwrap(); + + let written = RecordBatch::try_from_iter([ + ("decimal_values_19_0", Arc::new(decimals_19_0) as ArrayRef), + ("decimal_values_12_0", Arc::new(decimals_12_0) as ArrayRef), + ("decimal_values_17_10", Arc::new(decimals_17_10) as ArrayRef), + ]) + .unwrap(); + + let mut buffer = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buffer, written.schema(), None).unwrap(); + writer.write(&written).unwrap(); + writer.close().unwrap(); + + let read = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 8) + .unwrap() + .collect::, _>>() + .unwrap(); + + assert_eq!(&written.slice(0, 8), &read[0]); + } + + #[test] + fn test_list_skip() { + let mut list = ListBuilder::new(Int32Builder::new()); + list.append_value([Some(1), Some(2)]); + list.append_value([Some(3)]); + list.append_value([Some(4)]); + let list = list.finish(); + let batch = RecordBatch::try_from_iter([("l", Arc::new(list) as _)]).unwrap(); + + // First page contains 2 values but only 1 row + let props = WriterProperties::builder() + .set_data_page_row_count_limit(1) + .set_write_batch_size(2) + .build(); + + let mut buffer = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let selection = vec![RowSelector::skip(2), RowSelector::select(1)]; + let mut reader = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(buffer)) + .unwrap() + .with_row_selection(selection.into()) + .build() + .unwrap(); + let out = reader.next().unwrap().unwrap(); + assert_eq!(out.num_rows(), 1); + assert_eq!(out, batch.slice(2, 1)); + } + + fn test_decimal32_roundtrip() { + let d = |values: Vec, p: u8| { + let iter = values.into_iter(); + PrimitiveArray::::from_iter_values(iter) + .with_precision_and_scale(p, 2) + .unwrap() + }; + + let d1 = d(vec![1, 2, 3, 4, 5], 9); + let batch = RecordBatch::try_from_iter([("d1", Arc::new(d1) as ArrayRef)]).unwrap(); + + let mut buffer = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(buffer)).unwrap(); + let t1 = builder.parquet_schema().columns()[0].physical_type(); + assert_eq!(t1, PhysicalType::INT32); + + let mut reader = builder.build().unwrap(); + assert_eq!(batch.schema(), reader.schema()); + + let out = reader.next().unwrap().unwrap(); + assert_eq!(batch, out); + } + + fn test_decimal64_roundtrip() { + // Precision <= 9 -> INT32 + // Precision <= 18 -> INT64 + + let d = |values: Vec, p: u8| { + let iter = values.into_iter(); + PrimitiveArray::::from_iter_values(iter) + .with_precision_and_scale(p, 2) + .unwrap() + }; + + let d1 = d(vec![1, 2, 3, 4, 5], 9); + let d2 = d(vec![1, 2, 3, 4, 10.pow(10) - 1], 10); + let d3 = d(vec![1, 2, 3, 4, 10.pow(18) - 1], 18); + + let batch = RecordBatch::try_from_iter([ + ("d1", Arc::new(d1) as ArrayRef), + ("d2", Arc::new(d2) as ArrayRef), + ("d3", Arc::new(d3) as ArrayRef), + ]) + .unwrap(); + + let mut buffer = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(buffer)).unwrap(); + let t1 = builder.parquet_schema().columns()[0].physical_type(); + assert_eq!(t1, PhysicalType::INT32); + let t2 = builder.parquet_schema().columns()[1].physical_type(); + assert_eq!(t2, PhysicalType::INT64); + let t3 = builder.parquet_schema().columns()[2].physical_type(); + assert_eq!(t3, PhysicalType::INT64); + + let mut reader = builder.build().unwrap(); + assert_eq!(batch.schema(), reader.schema()); + + let out = reader.next().unwrap().unwrap(); + assert_eq!(batch, out); + } + + fn test_decimal_roundtrip() { + // Precision <= 9 -> INT32 + // Precision <= 18 -> INT64 + // Precision > 18 -> FIXED_LEN_BYTE_ARRAY + + let d = |values: Vec, p: u8| { + let iter = values.into_iter().map(T::Native::usize_as); + PrimitiveArray::::from_iter_values(iter) + .with_precision_and_scale(p, 2) + .unwrap() + }; + + let d1 = d(vec![1, 2, 3, 4, 5], 9); + let d2 = d(vec![1, 2, 3, 4, 10.pow(10) - 1], 10); + let d3 = d(vec![1, 2, 3, 4, 10.pow(18) - 1], 18); + let d4 = d(vec![1, 2, 3, 4, 10.pow(19) - 1], 19); + + let batch = RecordBatch::try_from_iter([ + ("d1", Arc::new(d1) as ArrayRef), + ("d2", Arc::new(d2) as ArrayRef), + ("d3", Arc::new(d3) as ArrayRef), + ("d4", Arc::new(d4) as ArrayRef), + ]) + .unwrap(); + + let mut buffer = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(buffer)).unwrap(); + let t1 = builder.parquet_schema().columns()[0].physical_type(); + assert_eq!(t1, PhysicalType::INT32); + let t2 = builder.parquet_schema().columns()[1].physical_type(); + assert_eq!(t2, PhysicalType::INT64); + let t3 = builder.parquet_schema().columns()[2].physical_type(); + assert_eq!(t3, PhysicalType::INT64); + let t4 = builder.parquet_schema().columns()[3].physical_type(); + assert_eq!(t4, PhysicalType::FIXED_LEN_BYTE_ARRAY); + + let mut reader = builder.build().unwrap(); + assert_eq!(batch.schema(), reader.schema()); + + let out = reader.next().unwrap().unwrap(); + assert_eq!(batch, out); + } + + #[test] + fn test_decimal() { + test_decimal32_roundtrip(); + test_decimal64_roundtrip(); + test_decimal_roundtrip::(); + test_decimal_roundtrip::(); + } + + #[test] + fn test_list_selection() { + let schema = Arc::new(Schema::new(vec![Field::new_list( + "list", + Field::new_list_field(ArrowDataType::Utf8, true), + false, + )])); + let mut buf = Vec::with_capacity(1024); + + let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None).unwrap(); + + for i in 0..2 { + let mut list_a_builder = ListBuilder::new(StringBuilder::new()); + for j in 0..1024 { + list_a_builder.values().append_value(format!("{i} {j}")); + list_a_builder.append(true); + } + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(list_a_builder.finish())]) + .unwrap(); + writer.write(&batch).unwrap(); + } + let _metadata = writer.close().unwrap(); + + let buf = Bytes::from(buf); + let reader = ParquetRecordBatchReaderBuilder::try_new(buf) + .unwrap() + .with_row_selection(RowSelection::from(vec![ + RowSelector::skip(100), + RowSelector::select(924), + RowSelector::skip(100), + RowSelector::select(924), + ])) + .build() + .unwrap(); + + let batches = reader.collect::, _>>().unwrap(); + let batch = concat_batches(&schema, &batches).unwrap(); + + assert_eq!(batch.num_rows(), 924 * 2); + let list = batch.column(0).as_list::(); + + for w in list.value_offsets().windows(2) { + assert_eq!(w[0] + 1, w[1]) + } + let mut values = list.values().as_string::().iter(); + + for i in 0..2 { + for j in 100..1024 { + let expected = format!("{i} {j}"); + assert_eq!(values.next().unwrap().unwrap(), &expected); + } + } + } + + #[test] + fn test_list_selection_fuzz() { + let mut rng = rng(); + let schema = Arc::new(Schema::new(vec![Field::new_list( + "list", + Field::new_list( + Field::LIST_FIELD_DEFAULT_NAME, + Field::new_list_field(ArrowDataType::Int32, true), + true, + ), + true, + )])); + let mut buf = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None).unwrap(); + + let mut list_a_builder = ListBuilder::new(ListBuilder::new(Int32Builder::new())); + + for _ in 0..2048 { + if rng.random_bool(0.2) { + list_a_builder.append(false); + continue; + } + + let list_a_len = rng.random_range(0..10); + let list_b_builder = list_a_builder.values(); + + for _ in 0..list_a_len { + if rng.random_bool(0.2) { + list_b_builder.append(false); + continue; + } + + let list_b_len = rng.random_range(0..10); + let int_builder = list_b_builder.values(); + for _ in 0..list_b_len { + match rng.random_bool(0.2) { + true => int_builder.append_null(), + false => int_builder.append_value(rng.random()), + } + } + list_b_builder.append(true) + } + list_a_builder.append(true); + } + + let array = Arc::new(list_a_builder.finish()); + let batch = RecordBatch::try_new(schema, vec![array]).unwrap(); + + writer.write(&batch).unwrap(); + let _metadata = writer.close().unwrap(); + + let buf = Bytes::from(buf); + + let cases = [ + vec![ + RowSelector::skip(100), + RowSelector::select(924), + RowSelector::skip(100), + RowSelector::select(924), + ], + vec![ + RowSelector::select(924), + RowSelector::skip(100), + RowSelector::select(924), + RowSelector::skip(100), + ], + vec![ + RowSelector::skip(1023), + RowSelector::select(1), + RowSelector::skip(1023), + RowSelector::select(1), + ], + vec![ + RowSelector::select(1), + RowSelector::skip(1023), + RowSelector::select(1), + RowSelector::skip(1023), + ], + ]; + + for batch_size in [100, 1024, 2048] { + for selection in &cases { + let selection = RowSelection::from(selection.clone()); + let reader = ParquetRecordBatchReaderBuilder::try_new(buf.clone()) + .unwrap() + .with_row_selection(selection.clone()) + .with_batch_size(batch_size) + .build() + .unwrap(); + + let batches = reader.collect::, _>>().unwrap(); + let actual = concat_batches(batch.schema_ref(), &batches).unwrap(); + assert_eq!(actual.num_rows(), selection.row_count()); + + let mut batch_offset = 0; + let mut actual_offset = 0; + for selector in selection.iter() { + if selector.skip { + batch_offset += selector.row_count; + continue; + } + + assert_eq!( + batch.slice(batch_offset, selector.row_count), + actual.slice(actual_offset, selector.row_count) + ); + + batch_offset += selector.row_count; + actual_offset += selector.row_count; + } + } + } + } + + #[test] + fn test_read_old_nested_list() { + use arrow::datatypes::DataType; + use arrow::datatypes::ToByteSlice; + + let testdata = arrow::util::test_util::parquet_test_data(); + // message my_record { + // REQUIRED group a (LIST) { + // REPEATED group array (LIST) { + // REPEATED INT32 array; + // } + // } + // } + // should be read as list> + let path = format!("{testdata}/old_list_structure.parquet"); + let test_file = File::open(path).unwrap(); + + // create expected ListArray + let a_values = Int32Array::from(vec![1, 2, 3, 4]); + + // Construct a buffer for value offsets, for the nested array: [[1, 2], [3, 4]] + let a_value_offsets = arrow::buffer::Buffer::from([0, 2, 4].to_byte_slice()); + + // Construct a list array from the above two + let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new( + "array", + DataType::Int32, + false, + )))) + .len(2) + .add_buffer(a_value_offsets) + .add_child_data(a_values.into_data()) + .build() + .unwrap(); + let a = ListArray::from(a_list_data); + + let builder = ParquetRecordBatchReaderBuilder::try_new(test_file).unwrap(); + let mut reader = builder.build().unwrap(); + let out = reader.next().unwrap().unwrap(); + assert_eq!(out.num_rows(), 1); + assert_eq!(out.num_columns(), 1); + // grab first column + let c0 = out.column(0); + let c0arr = c0.as_any().downcast_ref::().unwrap(); + // get first row: [[1, 2], [3, 4]] + let r0 = c0arr.value(0); + let r0arr = r0.as_any().downcast_ref::().unwrap(); + assert_eq!(r0arr, &a); + } + + #[test] + fn test_map_no_value() { + // File schema: + // message schema { + // required group my_map (MAP) { + // repeated group key_value { + // required int32 key; + // optional int32 value; + // } + // } + // required group my_map_no_v (MAP) { + // repeated group key_value { + // required int32 key; + // } + // } + // required group my_list (LIST) { + // repeated group list { + // required int32 element; + // } + // } + // } + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/map_no_value.parquet"); + let file = File::open(path).unwrap(); + + let mut reader = ParquetRecordBatchReaderBuilder::try_new(file) + .unwrap() + .build() + .unwrap(); + let out = reader.next().unwrap().unwrap(); + assert_eq!(out.num_rows(), 3); + assert_eq!(out.num_columns(), 3); + // my_map_no_v and my_list columns should now be equivalent + let c0 = out.column(1).as_list::(); + let c1 = out.column(2).as_list::(); + assert_eq!(c0.len(), c1.len()); + c0.iter().zip(c1.iter()).for_each(|(l, r)| assert_eq!(l, r)); + } + + #[test] + fn test_read_unknown_logical_type() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/unknown-logical-type.parquet"); + let test_file = File::open(path).unwrap(); + + let builder = ParquetRecordBatchReaderBuilder::try_new(test_file) + .expect("Error creating reader builder"); + + let schema = builder.metadata().file_metadata().schema_descr(); + assert_eq!( + schema.column(0).logical_type_ref(), + Some(&LogicalType::String) + ); + assert_eq!( + schema.column(1).logical_type_ref(), + Some(&LogicalType::_Unknown { field_id: 2555 }) + ); + assert_eq!(schema.column(1).physical_type(), PhysicalType::BYTE_ARRAY); + + let mut reader = builder.build().unwrap(); + let out = reader.next().unwrap().unwrap(); + assert_eq!(out.num_rows(), 3); + assert_eq!(out.num_columns(), 2); + } + + #[test] + fn test_read_row_numbers() { + let file = write_parquet_from_iter(vec![( + "value", + Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef, + )]); + let supplied_fields = Fields::from(vec![Field::new("value", ArrowDataType::Int64, false)]); + + let row_number_field = Arc::new( + Field::new("row_number", ArrowDataType::Int64, false).with_extension_type(RowNumber), + ); + + let options = ArrowReaderOptions::new() + .with_schema(Arc::new(Schema::new(supplied_fields))) + .with_virtual_columns(vec![row_number_field.clone()]) + .unwrap(); + let mut arrow_reader = ParquetRecordBatchReaderBuilder::try_new_with_options( + file.try_clone().unwrap(), + options, + ) + .expect("reader builder with schema") + .build() + .expect("reader with schema"); + + let batch = arrow_reader.next().unwrap().unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("value", ArrowDataType::Int64, false), + (*row_number_field).clone(), + ])); + + assert_eq!(batch.schema(), schema); + assert_eq!(batch.num_columns(), 2); + assert_eq!(batch.num_rows(), 3); + assert_eq!( + batch + .column(0) + .as_primitive::() + .iter() + .collect::>(), + vec![Some(1), Some(2), Some(3)] + ); + assert_eq!( + batch + .column(1) + .as_primitive::() + .iter() + .collect::>(), + vec![Some(0), Some(1), Some(2)] + ); + } + + #[test] + fn test_read_only_row_numbers() { + let file = write_parquet_from_iter(vec![( + "value", + Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef, + )]); + let row_number_field = Arc::new( + Field::new("row_number", ArrowDataType::Int64, false).with_extension_type(RowNumber), + ); + let options = ArrowReaderOptions::new() + .with_virtual_columns(vec![row_number_field.clone()]) + .unwrap(); + let metadata = ArrowReaderMetadata::load(&file, options).unwrap(); + let num_columns = metadata + .metadata + .file_metadata() + .schema_descr() + .num_columns(); + + let mut arrow_reader = ParquetRecordBatchReaderBuilder::new_with_metadata(file, metadata) + .with_projection(ProjectionMask::none(num_columns)) + .build() + .expect("reader with schema"); + + let batch = arrow_reader.next().unwrap().unwrap(); + let schema = Arc::new(Schema::new(vec![row_number_field])); + + assert_eq!(batch.schema(), schema); + assert_eq!(batch.num_columns(), 1); + assert_eq!(batch.num_rows(), 3); + assert_eq!( + batch + .column(0) + .as_primitive::() + .iter() + .collect::>(), + vec![Some(0), Some(1), Some(2)] + ); + } + + #[test] + fn test_read_row_numbers_row_group_order() -> Result<()> { + // Make a parquet file with 100 rows split across 2 row groups + let array = Int64Array::from_iter_values(5000..5100); + let batch = RecordBatch::try_from_iter([("col", Arc::new(array) as ArrayRef)])?; + let mut buffer = Vec::new(); + let options = WriterProperties::builder() + .set_max_row_group_row_count(Some(50)) + .build(); + let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema().clone(), Some(options))?; + // write in 10 row batches as the size limits are enforced after each batch + for batch_chunk in (0..10).map(|i| batch.slice(i * 10, 10)) { + writer.write(&batch_chunk)?; + } + writer.close()?; + + let row_number_field = Arc::new( + Field::new("row_number", ArrowDataType::Int64, false).with_extension_type(RowNumber), + ); + + let buffer = Bytes::from(buffer); + + let options = + ArrowReaderOptions::new().with_virtual_columns(vec![row_number_field.clone()])?; + + // read out with normal options + let arrow_reader = + ParquetRecordBatchReaderBuilder::try_new_with_options(buffer.clone(), options.clone())? + .build()?; + + assert_eq!( + ValuesAndRowNumbers { + values: (5000..5100).collect(), + row_numbers: (0..100).collect() + }, + ValuesAndRowNumbers::new_from_reader(arrow_reader) + ); + + // Now read, out of order row groups + let arrow_reader = ParquetRecordBatchReaderBuilder::try_new_with_options(buffer, options)? + .with_row_groups(vec![1, 0]) + .build()?; + + assert_eq!( + ValuesAndRowNumbers { + values: (5050..5100).chain(5000..5050).collect(), + row_numbers: (50..100).chain(0..50).collect(), + }, + ValuesAndRowNumbers::new_from_reader(arrow_reader) + ); + + Ok(()) + } + + #[derive(Debug, PartialEq)] + struct ValuesAndRowNumbers { + values: Vec, + row_numbers: Vec, + } + impl ValuesAndRowNumbers { + fn new_from_reader(reader: ParquetRecordBatchReader) -> Self { + let mut values = vec![]; + let mut row_numbers = vec![]; + for batch in reader { + let batch = batch.expect("Could not read batch"); + values.extend( + batch + .column_by_name("col") + .expect("Could not get col column") + .as_primitive::() + .iter() + .map(|v| v.expect("Could not get value")), + ); + + row_numbers.extend( + batch + .column_by_name("row_number") + .expect("Could not get row_number column") + .as_primitive::() + .iter() + .map(|v| v.expect("Could not get row number")) + .collect::>(), + ); + } + Self { + values, + row_numbers, + } + } + } + + #[test] + fn test_with_virtual_columns_rejects_non_virtual_fields() { + // Try to pass a regular field (not a virtual column) to with_virtual_columns + let regular_field = Arc::new(Field::new("regular_column", ArrowDataType::Int64, false)); + assert_eq!( + ArrowReaderOptions::new() + .with_virtual_columns(vec![regular_field]) + .unwrap_err() + .to_string(), + "Parquet error: Field 'regular_column' is not a virtual column. Virtual columns must have extension type names starting with 'arrow.virtual.'" + ); + } + + #[test] + fn test_row_numbers_with_multiple_row_groups() { + test_row_numbers_with_multiple_row_groups_helper( + false, + |path, selection, _row_filter, batch_size| { + let file = File::open(path).unwrap(); + let row_number_field = Arc::new( + Field::new("row_number", ArrowDataType::Int64, false) + .with_extension_type(RowNumber), + ); + let options = ArrowReaderOptions::new() + .with_virtual_columns(vec![row_number_field]) + .unwrap(); + let reader = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options) + .unwrap() + .with_row_selection(selection) + .with_batch_size(batch_size) + .build() + .expect("Could not create reader"); + reader + .collect::, _>>() + .expect("Could not read") + }, + ); + } + + #[test] + fn test_row_numbers_with_multiple_row_groups_and_filter() { + test_row_numbers_with_multiple_row_groups_helper( + true, + |path, selection, row_filter, batch_size| { + let file = File::open(path).unwrap(); + let row_number_field = Arc::new( + Field::new("row_number", ArrowDataType::Int64, false) + .with_extension_type(RowNumber), + ); + let options = ArrowReaderOptions::new() + .with_virtual_columns(vec![row_number_field]) + .unwrap(); + let reader = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options) + .unwrap() + .with_row_selection(selection) + .with_batch_size(batch_size) + .with_row_filter(row_filter.expect("No filter")) + .build() + .expect("Could not create reader"); + reader + .collect::, _>>() + .expect("Could not read") + }, + ); + } + + #[test] + fn test_read_row_group_indices() { + // create a parquet file with 3 row groups, 2 rows each + let array1 = Int64Array::from(vec![1, 2]); + let array2 = Int64Array::from(vec![3, 4]); + let array3 = Int64Array::from(vec![5, 6]); + + let batch1 = + RecordBatch::try_from_iter(vec![("value", Arc::new(array1) as ArrayRef)]).unwrap(); + let batch2 = + RecordBatch::try_from_iter(vec![("value", Arc::new(array2) as ArrayRef)]).unwrap(); + let batch3 = + RecordBatch::try_from_iter(vec![("value", Arc::new(array3) as ArrayRef)]).unwrap(); + + let mut buffer = Vec::new(); + let options = WriterProperties::builder() + .set_max_row_group_row_count(Some(2)) + .build(); + let mut writer = ArrowWriter::try_new(&mut buffer, batch1.schema(), Some(options)).unwrap(); + writer.write(&batch1).unwrap(); + writer.write(&batch2).unwrap(); + writer.write(&batch3).unwrap(); + writer.close().unwrap(); + + let file = Bytes::from(buffer); + let row_group_index_field = Arc::new( + Field::new("row_group_index", ArrowDataType::Int64, false) + .with_extension_type(RowGroupIndex), + ); + + let options = ArrowReaderOptions::new() + .with_virtual_columns(vec![row_group_index_field.clone()]) + .unwrap(); + let mut arrow_reader = + ParquetRecordBatchReaderBuilder::try_new_with_options(file.clone(), options) + .expect("reader builder with virtual columns") + .build() + .expect("reader with virtual columns"); + + let batch = arrow_reader.next().unwrap().unwrap(); + + assert_eq!(batch.num_columns(), 2); + assert_eq!(batch.num_rows(), 6); + + assert_eq!( + batch + .column(0) + .as_primitive::() + .iter() + .collect::>(), + vec![Some(1), Some(2), Some(3), Some(4), Some(5), Some(6)] + ); + + assert_eq!( + batch + .column(1) + .as_primitive::() + .iter() + .collect::>(), + vec![Some(0), Some(0), Some(1), Some(1), Some(2), Some(2)] + ); + } + + #[test] + fn test_read_only_row_group_indices() { + let array1 = Int64Array::from(vec![1, 2, 3]); + let array2 = Int64Array::from(vec![4, 5]); + + let batch1 = + RecordBatch::try_from_iter(vec![("value", Arc::new(array1) as ArrayRef)]).unwrap(); + let batch2 = + RecordBatch::try_from_iter(vec![("value", Arc::new(array2) as ArrayRef)]).unwrap(); + + let mut buffer = Vec::new(); + let options = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .build(); + let mut writer = ArrowWriter::try_new(&mut buffer, batch1.schema(), Some(options)).unwrap(); + writer.write(&batch1).unwrap(); + writer.write(&batch2).unwrap(); + writer.close().unwrap(); + + let file = Bytes::from(buffer); + let row_group_index_field = Arc::new( + Field::new("row_group_index", ArrowDataType::Int64, false) + .with_extension_type(RowGroupIndex), + ); + + let options = ArrowReaderOptions::new() + .with_virtual_columns(vec![row_group_index_field.clone()]) + .unwrap(); + let metadata = ArrowReaderMetadata::load(&file, options).unwrap(); + let num_columns = metadata + .metadata + .file_metadata() + .schema_descr() + .num_columns(); + + let mut arrow_reader = ParquetRecordBatchReaderBuilder::new_with_metadata(file, metadata) + .with_projection(ProjectionMask::none(num_columns)) + .build() + .expect("reader with virtual columns only"); + + let batch = arrow_reader.next().unwrap().unwrap(); + let schema = Arc::new(Schema::new(vec![(*row_group_index_field).clone()])); + + assert_eq!(batch.schema(), schema); + assert_eq!(batch.num_columns(), 1); + assert_eq!(batch.num_rows(), 5); + + assert_eq!( + batch + .column(0) + .as_primitive::() + .iter() + .collect::>(), + vec![Some(0), Some(0), Some(0), Some(1), Some(1)] + ); + } + + #[test] + fn test_read_row_group_indices_with_selection() -> Result<()> { + let mut buffer = Vec::new(); + let options = WriterProperties::builder() + .set_max_row_group_row_count(Some(10)) + .build(); + + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + ArrowDataType::Int64, + false, + )])); + + let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(), Some(options))?; + + // write out 3 batches of 10 rows each + for i in 0..3 { + let start = i * 10; + let array = Int64Array::from_iter_values(start..start + 10); + let batch = RecordBatch::try_from_iter(vec![("value", Arc::new(array) as ArrayRef)])?; + writer.write(&batch)?; + } + writer.close()?; + + let file = Bytes::from(buffer); + let row_group_index_field = Arc::new( + Field::new("rg_idx", ArrowDataType::Int64, false).with_extension_type(RowGroupIndex), + ); + + let options = + ArrowReaderOptions::new().with_virtual_columns(vec![row_group_index_field])?; + + // test row groups are read in reverse order + let arrow_reader = + ParquetRecordBatchReaderBuilder::try_new_with_options(file.clone(), options.clone())? + .with_row_groups(vec![2, 1, 0]) + .build()?; + + let batches: Vec<_> = arrow_reader.collect::, _>>()?; + let combined = concat_batches(&batches[0].schema(), &batches)?; + + let values = combined.column(0).as_primitive::(); + let first_val = values.value(0); + let last_val = values.value(combined.num_rows() - 1); + // first row from rg 2 + assert_eq!(first_val, 20); + // the last row from rg 0 + assert_eq!(last_val, 9); + + let rg_indices = combined.column(1).as_primitive::(); + assert_eq!(rg_indices.value(0), 2); + assert_eq!(rg_indices.value(10), 1); + assert_eq!(rg_indices.value(20), 0); + + Ok(()) + } + + pub(crate) fn test_row_numbers_with_multiple_row_groups_helper( + use_filter: bool, + test_case: F, + ) where + F: FnOnce(PathBuf, RowSelection, Option, usize) -> Vec, + { + let seed: u64 = random(); + println!("test_row_numbers_with_multiple_row_groups seed: {}", seed); + let mut rng = StdRng::seed_from_u64(seed); + + use tempfile::TempDir; + let tempdir = TempDir::new().expect("Could not create temp dir"); + + let (bytes, metadata) = generate_file_with_row_numbers(&mut rng); + + let path = tempdir.path().join("test.parquet"); + std::fs::write(&path, bytes).expect("Could not write file"); + + let mut case = vec![]; + let mut remaining = metadata.file_metadata().num_rows(); + while remaining > 0 { + let row_count = rng.random_range(1..=remaining); + remaining -= row_count; + case.push(RowSelector { + row_count: row_count as usize, + skip: rng.random_bool(0.5), + }); + } + + let filter = use_filter.then(|| { + let filter = (0..metadata.file_metadata().num_rows()) + .map(|_| rng.random_bool(0.99)) + .collect::>(); + let mut filter_offset = 0; + RowFilter::new(vec![Box::new(ArrowPredicateFn::new( + ProjectionMask::all(), + move |b| { + let array = BooleanArray::from_iter( + filter + .iter() + .skip(filter_offset) + .take(b.num_rows()) + .map(|x| Some(*x)), + ); + filter_offset += b.num_rows(); + Ok(array) + }, + ))]) + }); + + let selection = RowSelection::from(case); + let batches = test_case(path, selection.clone(), filter, rng.random_range(1..4096)); + + if selection.skipped_row_count() == metadata.file_metadata().num_rows() as usize { + assert!(batches.into_iter().all(|batch| batch.num_rows() == 0)); + return; + } + let actual = concat_batches(batches.first().expect("No batches").schema_ref(), &batches) + .expect("Failed to concatenate"); + // assert_eq!(selection.row_count(), actual.num_rows()); + let values = actual + .column(0) + .as_primitive::() + .iter() + .collect::>(); + let row_numbers = actual + .column(1) + .as_primitive::() + .iter() + .collect::>(); + assert_eq!( + row_numbers + .into_iter() + .map(|number| number.map(|number| number + 1)) + .collect::>(), + values + ); + } + + fn generate_file_with_row_numbers(rng: &mut impl Rng) -> (Bytes, ParquetMetaData) { + let schema = Arc::new(Schema::new(Fields::from(vec![Field::new( + "value", + ArrowDataType::Int64, + false, + )]))); + + let mut buf = Vec::with_capacity(1024); + let mut writer = + ArrowWriter::try_new(&mut buf, schema.clone(), None).expect("Could not create writer"); + + let mut values = 1..=rng.random_range(1..4096); + while !values.is_empty() { + let batch_values = values + .by_ref() + .take(rng.random_range(1..4096)) + .collect::>(); + let array = Arc::new(Int64Array::from(batch_values)) as ArrayRef; + let batch = + RecordBatch::try_from_iter([("value", array)]).expect("Could not create batch"); + writer.write(&batch).expect("Could not write batch"); + writer.flush().expect("Could not flush"); + } + let metadata = writer.close().expect("Could not close writer"); + + (Bytes::from(buf), metadata) + } +} diff --git a/vendor/parquet/src/arrow/arrow_reader/read_plan.rs b/vendor/parquet/src/arrow/arrow_reader/read_plan.rs new file mode 100644 index 0000000..5b6cff4 --- /dev/null +++ b/vendor/parquet/src/arrow/arrow_reader/read_plan.rs @@ -0,0 +1,806 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`ReadPlan`] and [`ReadPlanBuilder`] for determining which rows to read +//! from a Parquet file + +use crate::arrow::array_reader::ArrayReader; +use crate::arrow::arrow_reader::selection::{ + LoadedRowRanges, RowSelectionInner, RowSelectionPolicy, RowSelectionStrategy, mask_to_selectors, +}; +use crate::arrow::arrow_reader::{ + ArrowPredicate, ParquetRecordBatchReader, RowSelection, RowSelectionCursor, RowSelector, +}; +use crate::errors::{ParquetError, Result}; +use arrow_array::{Array, BooleanArray}; +use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder}; +use arrow_select::filter::prep_null_mask_filter; +use std::collections::VecDeque; +use std::sync::Arc; + +/// Options for [`ReadPlanBuilder::with_predicate_options`]. +pub struct PredicateOptions<'a> { + array_reader: Box, + predicate: &'a mut dyn ArrowPredicate, + limit: Option, + total_rows: usize, +} + +impl<'a> PredicateOptions<'a> { + /// Create options for evaluating `predicate` against rows produced by + /// `array_reader`. + /// + /// By default there is no match-count limit; the predicate is evaluated + /// over every row the reader yields. Use [`Self::with_limit`] to enable + /// early termination. + pub fn new(array_reader: Box, predicate: &'a mut dyn ArrowPredicate) -> Self { + Self { + array_reader, + predicate, + limit: None, + total_rows: 0, + } + } + + /// Stop scanning `array_reader` once `limit` matches have accumulated. + /// + /// Performance optimization for `LIMIT` / TopK: when the cumulative + /// `true_count` reaches `limit`, the current filter batch is truncated + /// at the `limit`-th match and remaining batches are never decoded. + /// + /// `limit` counts predicate matches, not output rows — callers applying + /// an offset must pass `offset + limit`. + /// + /// `total_rows` is the row count `array_reader` would yield if iterated + /// to completion. It is used to pad un-evaluated trailing rows as "not + /// selected" so the returned [`RowSelection`] covers the full row group. + /// + /// Only valid for the *last* predicate in a filter chain: intermediate + /// predicates' match counts do not map 1:1 to output rows. + pub fn with_limit(mut self, limit: usize, total_rows: usize) -> Self { + self.limit = Some(limit); + self.total_rows = total_rows; + self + } +} + +/// A builder for [`ReadPlan`] +#[derive(Clone, Debug)] +pub struct ReadPlanBuilder { + batch_size: usize, + /// Which rows to select. Includes the result of all filters applied so far + selection: Option, + /// Policy to use when materializing the row selection + row_selection_policy: RowSelectionPolicy, + /// Row ranges with page data loaded for the current projection. + loaded_row_ranges: Option>, +} + +impl ReadPlanBuilder { + /// Create a `ReadPlanBuilder` with the given batch size + pub fn new(batch_size: usize) -> Self { + Self { + batch_size, + selection: None, + row_selection_policy: RowSelectionPolicy::default(), + loaded_row_ranges: None, + } + } + + /// Set the current selection to the given value + pub fn with_selection(mut self, selection: Option) -> Self { + self.selection = selection; + self + } + + /// Configure the policy to use when materialising the [`RowSelection`] + /// + /// Defaults to [`RowSelectionPolicy::Auto`] + pub fn with_row_selection_policy(mut self, policy: RowSelectionPolicy) -> Self { + self.row_selection_policy = policy; + self + } + + pub(crate) fn with_loaded_row_ranges(mut self, ranges: Option) -> Self { + self.loaded_row_ranges = ranges.map(Arc::new); + self + } + + /// Returns the current row selection policy + pub fn row_selection_policy(&self) -> &RowSelectionPolicy { + &self.row_selection_policy + } + + /// Returns the current selection, if any + pub fn selection(&self) -> Option<&RowSelection> { + self.selection.as_ref() + } + + /// Specifies the number of rows in the row group, before filtering is applied. + /// + /// Returns a [`LimitedReadPlanBuilder`] that can apply + /// offset and limit. + /// + /// Call [`LimitedReadPlanBuilder::build_limited`] to apply the limits to this + /// selection. + pub(crate) fn limited(self, row_count: usize) -> LimitedReadPlanBuilder { + LimitedReadPlanBuilder::new(self, row_count) + } + + /// Returns true if the current plan selects any rows + pub fn selects_any(&self) -> bool { + self.selection + .as_ref() + .map(|s| s.selects_any()) + .unwrap_or(true) + } + + /// Returns the number of rows selected, or `None` if all rows are selected. + pub fn num_rows_selected(&self) -> Option { + self.selection.as_ref().map(|s| s.row_count()) + } + + /// Returns the [`RowSelectionStrategy`] for this plan. + /// + /// 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 selection = match self.selection.as_ref() { + Some(selection) => selection, + None => return RowSelectionStrategy::Selectors, + }; + + selection.auto_selection_strategy(threshold) + } + } + } + + /// Evaluates an [`ArrowPredicate`], updating this plan's `selection` + /// + /// If the current `selection` is `Some`, the resulting [`RowSelection`] + /// will be the conjunction of the existing selection and the rows selected + /// by `predicate`. + /// + /// Note: pre-existing selections may come from evaluating a previous predicate + /// or if the [`ParquetRecordBatchReader`] specified an explicit + /// [`RowSelection`] in addition to one or more predicates. + pub fn with_predicate( + self, + array_reader: Box, + predicate: &mut dyn ArrowPredicate, + ) -> Result { + self.with_predicate_options(PredicateOptions::new(array_reader, predicate)) + } + + /// Evaluates an [`ArrowPredicate`] with the given [`PredicateOptions`], + /// updating this plan's `selection`. + /// + /// Like [`Self::with_predicate`], but allows additional options such as a + /// match-count limit for early termination (see + /// [`PredicateOptions::with_limit`]). + pub fn with_predicate_options(mut self, options: PredicateOptions<'_>) -> Result { + let PredicateOptions { + array_reader, + predicate, + limit, + total_rows, + } = options; + + // Target length for the concatenated filter output: + // - Prior selection ⇒ the reader yields that many rows; `and_then` + // below requires the filter output to match. + // - No prior selection ⇒ the reader yields `total_rows`. We only + // need to pad when `limit` may short-circuit the loop; otherwise + // iteration naturally exhausts. + let expected_rows = match self.selection.as_ref() { + Some(s) => Some(s.row_count()), + None => limit.map(|_| total_rows), + }; + + let reader = ParquetRecordBatchReader::new(array_reader, self.clone().build()); + let mut filters = vec![]; + let mut processed_rows: usize = 0; + let mut matched_rows: usize = 0; + for maybe_batch in reader { + let maybe_batch = maybe_batch?; + let input_rows = maybe_batch.num_rows(); + let filter = predicate.evaluate(maybe_batch)?; + // Since user supplied predicate, check error here to catch bugs quickly + if filter.len() != input_rows { + return Err(arrow_err!( + "ArrowPredicate predicate returned {} rows, expected {input_rows}", + filter.len() + )); + } + let filter = match filter.null_count() { + 0 => filter, + // RowSelection::from_filters expects non-null filters. Convert + // NULL predicate results to false so they are not selected. + _ => prep_null_mask_filter(&filter), + }; + + processed_rows += input_rows; + + match limit { + Some(limit) if limit - matched_rows <= filter.len() => { + let truncated = filter.take_n_true(limit - matched_rows); + matched_rows += truncated.true_count(); + filters.push(truncated); + if matched_rows >= limit { + break; + } + } + _ => { + matched_rows += filter.true_count(); + filters.push(filter); + } + } + } + + // Pad the tail so the filters cover `expected_rows` total. This keeps + // the invariant that the resulting `RowSelection` spans every row the + // reader would have produced — rows past the early break are marked + // "not selected". When no limit is set the loop always exhausts and + // no padding is needed. + if let Some(expected) = expected_rows { + if processed_rows < expected { + let pad_len = expected - processed_rows; + filters.push(BooleanArray::new(BooleanBuffer::new_unset(pad_len), None)); + } + } + + // If the predicate selected all rows, applying it is a no-op. With no + // prior selection this keeps selection as None, enabling coalesced page + // fetches; with a prior selection it avoids rebuilding the same + // selection. + let all_selected = filters.iter().all(|f| f.true_count() == f.len()); + if all_selected { + return Ok(self); + } + let raw = if self + .selection + .as_ref() + .is_some_and(|s| s.as_mask().is_some()) + { + RowSelection::from_boolean_buffer(filters_to_boolean_buffer(&filters)) + } else { + RowSelection::from_filters(&filters) + }; + self.selection = match self.selection.take() { + Some(selection) => Some(selection.and_then(&raw)), + None => Some(raw), + }; + Ok(self) + } + + /// Create a final `ReadPlan` the read plan for the scan + pub fn build(mut self) -> ReadPlan { + // If selection is empty, truncate + if !self.selects_any() { + self.selection = Some(RowSelection::from(vec![])); + } + + // Preferred strategy must not be Auto + let selection_strategy = self.resolve_selection_strategy(); + + let Self { + batch_size, + selection, + row_selection_policy: _, + loaded_row_ranges, + } = self; + + let row_selection_cursor = selection + .map(|s| build_cursor(s.trim(), selection_strategy, loaded_row_ranges)) + .unwrap_or(RowSelectionCursor::new_all()); + + ReadPlan { + batch_size, + row_selection_cursor, + } + } +} + +/// 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_to_selectors(mask.mask())) + } + } +} + +/// Builder for [`ReadPlan`] that applies a limit and offset to the read plan +/// +/// See [`ReadPlanBuilder::limited`] to create this builder. +pub(crate) struct LimitedReadPlanBuilder { + /// The underlying builder + inner: ReadPlanBuilder, + /// Total number of rows in the row group before the selection, limit or + /// offset are applied + row_count: usize, + /// The offset to apply, if any + offset: Option, + /// The limit to apply, if any + limit: Option, +} + +impl LimitedReadPlanBuilder { + /// Create a new `LimitedReadPlanBuilder` from the existing builder and number of rows + fn new(inner: ReadPlanBuilder, row_count: usize) -> Self { + Self { + inner, + row_count, + offset: None, + limit: None, + } + } + + /// Set the offset to apply to the read plan + pub(crate) fn with_offset(mut self, offset: Option) -> Self { + self.offset = offset; + self + } + + /// Set the limit to apply to the read plan + pub(crate) fn with_limit(mut self, limit: Option) -> Self { + self.limit = limit; + self + } + + /// Apply offset and limit, updating the selection on the underlying builder + /// and returning it. + pub(crate) fn build_limited(self) -> ReadPlanBuilder { + let Self { + mut inner, + row_count, + offset, + limit, + } = self; + + // If the selection is empty, truncate + if !inner.selects_any() { + inner.selection = Some(RowSelection::from(vec![])); + } + + // If an offset is defined, apply it to the `selection` + if let Some(offset) = offset { + inner.selection = Some(match row_count.checked_sub(offset) { + None => RowSelection::from(vec![]), + Some(remaining) => inner + .selection + .map(|selection| selection.offset(offset)) + .unwrap_or_else(|| { + RowSelection::from(vec![ + RowSelector::skip(offset), + RowSelector::select(remaining), + ]) + }), + }); + } + + // If a limit is defined, apply it to the final `selection` + if let Some(limit) = limit { + inner.selection = Some( + inner + .selection + .map(|selection| selection.limit(limit)) + .unwrap_or_else(|| { + RowSelection::from(vec![RowSelector::select(limit.min(row_count))]) + }), + ); + } + + inner + } +} + +fn filters_to_boolean_buffer(filters: &[BooleanArray]) -> BooleanBuffer { + let total_rows = filters.iter().map(|f| f.len()).sum(); + let mut builder = BooleanBufferBuilder::new(total_rows); + for filter in filters { + assert_eq!(filter.null_count(), 0); + builder.append_buffer(filter.values()); + } + builder.finish() +} + +/// A plan reading specific rows from a Parquet Row Group. +/// +/// See [`ReadPlanBuilder`] to create `ReadPlan`s +#[derive(Debug)] +pub struct ReadPlan { + /// The number of rows to read in each batch + batch_size: usize, + /// Row ranges to be selected from the data source + row_selection_cursor: RowSelectionCursor, +} + +impl ReadPlan { + /// Returns a mutable reference to the selection selectors, if any + #[deprecated(since = "57.1.0", note = "Use `row_selection_cursor_mut` instead")] + pub fn selection_mut(&mut self) -> Option<&mut VecDeque> { + if let RowSelectionCursor::Selectors(selectors_cursor) = &mut self.row_selection_cursor { + Some(selectors_cursor.selectors_mut()) + } else { + None + } + } + + /// Returns a mutable reference to the row selection cursor + pub fn row_selection_cursor_mut(&mut self) -> &mut RowSelectionCursor { + &mut self.row_selection_cursor + } + + /// Return the number of rows to read in each output batch + #[inline(always)] + pub fn batch_size(&self) -> usize { + self.batch_size + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn builder_with_selection(selection: RowSelection) -> ReadPlanBuilder { + ReadPlanBuilder::new(1024).with_selection(Some(selection)) + } + + #[test] + fn preferred_selection_strategy_prefers_mask_by_default() { + let selection = RowSelection::from(vec![RowSelector::select(8)]); + let builder = builder_with_selection(selection); + assert_eq!( + builder.resolve_selection_strategy(), + RowSelectionStrategy::Mask + ); + } + + #[test] + fn preferred_selection_strategy_prefers_selectors_when_threshold_small() { + let selection = RowSelection::from(vec![RowSelector::select(8)]); + let builder = builder_with_selection(selection) + .with_row_selection_policy(RowSelectionPolicy::Auto { threshold: 1 }); + assert_eq!( + builder.resolve_selection_strategy(), + RowSelectionStrategy::Selectors + ); + } + + #[test] + fn preferred_selection_strategy_handles_dense_mask_backing() { + let bits: Vec<_> = (0..16).map(|i| i % 2 == 0).collect(); + let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits)); + let builder = builder_with_selection(selection) + .with_row_selection_policy(RowSelectionPolicy::Auto { threshold: 4 }); + assert_eq!( + builder.resolve_selection_strategy(), + RowSelectionStrategy::Mask + ); + } + + #[test] + fn preferred_selection_strategy_handles_sparse_mask_backing() { + let bits: Vec<_> = (0..128).map(|i| i < 64).collect(); + let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits)); + let builder = builder_with_selection(selection) + .with_row_selection_policy(RowSelectionPolicy::Auto { threshold: 32 }); + assert_eq!( + builder.resolve_selection_strategy(), + RowSelectionStrategy::Selectors + ); + } + + #[test] + fn preferred_selection_strategy_preserves_mask_threshold_boundaries() { + let mask = BooleanBuffer::from(vec![true; 8]); + + let empty = builder_with_selection(RowSelection::from_boolean_buffer( + BooleanBuffer::new_unset(0), + )); + assert_eq!( + empty.resolve_selection_strategy(), + RowSelectionStrategy::Mask + ); + + let disabled = builder_with_selection(RowSelection::from_boolean_buffer(mask.clone())) + .with_row_selection_policy(RowSelectionPolicy::Auto { threshold: 0 }); + assert_eq!( + disabled.resolve_selection_strategy(), + RowSelectionStrategy::Selectors + ); + + // Equality does not satisfy the policy's strict inequality: 8 < 1 * 8. + let equal = builder_with_selection(RowSelection::from_boolean_buffer(mask.clone())) + .with_row_selection_policy(RowSelectionPolicy::Auto { threshold: 8 }); + assert_eq!( + equal.resolve_selection_strategy(), + RowSelectionStrategy::Selectors + ); + + let above = builder_with_selection(RowSelection::from_boolean_buffer(mask)) + .with_row_selection_policy(RowSelectionPolicy::Auto { threshold: 9 }); + assert_eq!( + above.resolve_selection_strategy(), + RowSelectionStrategy::Mask + ); + } + + #[test] + fn preferred_selection_strategy_mask_matches_selector_backing() { + use rand::{Rng, rng}; + + let mut rand = rng(); + for _ in 0..200 { + let len = rand.random_range(0..256); + let bits: Vec<_> = (0..len).map(|_| rand.random_bool(0.5)).collect(); + let mask_backed = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone())); + let selector_backed = RowSelection::from_filters(&[BooleanArray::from(bits)]); + + for threshold in [0, 1, 2, 8, 32, 64, usize::MAX] { + assert_eq!( + mask_backed.auto_selection_strategy(threshold), + selector_backed.auto_selection_strategy(threshold), + "strategy differs for len {len} and threshold {threshold}" + ); + } + } + } + + #[test] + fn mask_plan_trims_trailing_skips_before_chunking() { + let mut plan = ReadPlanBuilder::new(8) + .with_selection(Some(RowSelection::from(vec![ + RowSelector::select(1), + RowSelector::skip(7), + ]))) + .with_row_selection_policy(RowSelectionPolicy::Mask) + .build(); + let RowSelectionCursor::Mask(cursor) = plan.row_selection_cursor_mut() else { + panic!("expected a Mask cursor"); + }; + + let chunk = cursor.next_chunk(8).unwrap(); + assert_eq!(chunk.chunk_rows, 1); + assert_eq!(chunk.selected_rows, 1); + assert!(cursor.is_empty()); + } + + #[test] + fn selectors_policy_lowers_mask_backed_selection() { + let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ + true, false, false, true, true, + ])); + let mut plan = ReadPlanBuilder::new(4) + .with_selection(Some(selection)) + .with_row_selection_policy(RowSelectionPolicy::Selectors) + .build(); + let RowSelectionCursor::Selectors(cursor) = plan.row_selection_cursor_mut() else { + panic!("expected a Selectors cursor"); + }; + + assert_eq!(cursor.next_selector(), RowSelector::select(1)); + assert_eq!(cursor.next_selector(), RowSelector::skip(2)); + assert_eq!(cursor.next_selector(), RowSelector::select(2)); + assert!(cursor.is_empty()); + } + + #[test] + fn mask_backed_plan_respects_loaded_row_ranges() { + // Start from a non-byte-aligned BooleanBuffer::slice so the mask-backed + // cursor path (new_mask_from_buffer) is exercised with a bit offset. + // Rows 0 and 11 of 12 are selected; pages for rows [4, 10) are unloaded. + let mut bits = vec![false; 5]; + bits.push(true); // row 0 + bits.extend(std::iter::repeat_n(false, 10)); // rows 1..=10 + bits.push(true); // row 11 + bits.extend([true, false, true]); // trailing padding outside the slice + let mask = BooleanBuffer::from(bits).slice(5, 12); + let selection = RowSelection::from_boolean_buffer(mask); + assert!(selection.as_mask().is_some()); + + let loaded = LoadedRowRanges::from_selection(RowSelection::from(vec![ + RowSelector::select(4), + RowSelector::skip(6), + RowSelector::select(2), + ])); + + let mut plan = ReadPlanBuilder::new(12) + .with_selection(Some(selection)) + .with_row_selection_policy(RowSelectionPolicy::Mask) + .with_loaded_row_ranges(Some(loaded)) + .build(); + let RowSelectionCursor::Mask(cursor) = plan.row_selection_cursor_mut() else { + panic!("expected a Mask cursor"); + }; + + // The first chunk stops at its final selected row instead of carrying + // trailing skipped rows to the loaded range boundary. + let first = cursor.next_chunk(12).unwrap(); + assert_eq!(first.initial_skip, 0); + assert_eq!(first.chunk_rows, 1); + assert_eq!(first.selected_rows, 1); + + // The second chunk skips directly to the next selected row. + let second = cursor.next_chunk(12).unwrap(); + assert_eq!(second.initial_skip, 10); + assert_eq!(second.chunk_rows, 1); + assert_eq!(second.selected_rows, 1); + assert!(cursor.is_empty()); + } + + #[test] + fn with_predicate_options_limit_pads_tail_when_no_prior_selection() { + use crate::arrow::ProjectionMask; + use crate::arrow::array_reader::StructArrayReader; + use crate::arrow::array_reader::test_util::make_int32_page_reader; + use crate::arrow::arrow_reader::ArrowPredicateFn; + use arrow_schema::{DataType as ArrowType, Field, Fields}; + + // 100 rows, all match the predicate. Limit stops the loop after 10 + // matches — but the resulting RowSelection must still describe the + // full 100-row row group (90 trailing rows as "not selected"), not + // only the 10 rows we happened to evaluate before breaking. + const TOTAL_ROWS: usize = 100; + const LIMIT: usize = 10; + + let data: Vec = (0..TOTAL_ROWS as i32).collect(); + let levels = vec![0; TOTAL_ROWS]; + let leaf = make_int32_page_reader(&data, &levels, &levels, 0, 0, None); + let struct_type = ArrowType::Struct(Fields::from(vec![Field::new( + "c0", + ArrowType::Int32, + false, + )])); + let struct_reader = StructArrayReader::new(struct_type, vec![leaf], 0, 0, false, None); + + let mut predicate = ArrowPredicateFn::new(ProjectionMask::all(), |batch| { + Ok(BooleanArray::from(vec![true; batch.num_rows()])) + }); + + let builder = ReadPlanBuilder::new(16) + .with_predicate_options( + PredicateOptions::new(Box::new(struct_reader), &mut predicate) + .with_limit(LIMIT, TOTAL_ROWS), + ) + .unwrap(); + + let selection = builder + .selection() + .expect("limit-driven early break must produce a selection"); + + // `row_count` counts selected rows — must equal the limit. + assert_eq!(selection.row_count(), LIMIT); + + // Total rows covered (selects + skips) must equal the full row group + // so downstream offset/limit math stays in absolute-row space. + let total: usize = selection.iter().map(|s| s.row_count).sum(); + assert_eq!( + total, TOTAL_ROWS, + "selection must span the full row group, not only the prefix evaluated before the limit" + ); + } + + #[test] + fn with_predicate_options_preserves_mask_selection() { + use crate::arrow::ProjectionMask; + use crate::arrow::array_reader::StructArrayReader; + use crate::arrow::array_reader::test_util::make_int32_page_reader; + use crate::arrow::arrow_reader::ArrowPredicateFn; + use arrow_schema::{DataType as ArrowType, Field, Fields}; + + let data: Vec = (0..6).collect(); + let levels = vec![0; data.len()]; + let leaf = make_int32_page_reader(&data, &levels, &levels, 0, 0, None); + let struct_type = ArrowType::Struct(Fields::from(vec![Field::new( + "c0", + ArrowType::Int32, + false, + )])); + let struct_reader = StructArrayReader::new(struct_type, vec![leaf], 0, 0, false, None); + + let prior = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ + true, false, true, true, false, true, + ])); + let mut filters = vec![BooleanArray::from(vec![true, false, true, false])]; + let mut predicate = ArrowPredicateFn::new(ProjectionMask::all(), move |batch| { + assert_eq!(batch.num_rows(), 4); + Ok(filters.remove(0)) + }); + + let builder = ReadPlanBuilder::new(16) + .with_selection(Some(prior)) + .with_predicate_options(PredicateOptions::new( + Box::new(struct_reader), + &mut predicate, + )) + .unwrap(); + + let selection = builder.selection().unwrap(); + assert!(selection.as_mask().is_some()); + + let expected = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ + true, false, false, true, false, false, + ])); + assert_eq!(selection, &expected); + } + + #[test] + fn with_predicate_options_limit_handles_null_filters() { + use crate::arrow::ProjectionMask; + use crate::arrow::array_reader::StructArrayReader; + use crate::arrow::array_reader::test_util::make_int32_page_reader; + use crate::arrow::arrow_reader::ArrowPredicateFn; + use arrow_schema::{DataType as ArrowType, Field, Fields}; + + const TOTAL_ROWS: usize = 100; + const LIMIT: usize = 10; + + let data: Vec = (0..TOTAL_ROWS as i32).collect(); + let levels = vec![0; TOTAL_ROWS]; + let leaf = make_int32_page_reader(&data, &levels, &levels, 0, 0, None); + let struct_type = ArrowType::Struct(Fields::from(vec![Field::new( + "c0", + ArrowType::Int32, + false, + )])); + let struct_reader = StructArrayReader::new(struct_type, vec![leaf], 0, 0, false, None); + + let mut predicate = ArrowPredicateFn::new(ProjectionMask::all(), |batch| { + Ok((0..batch.num_rows()) + .map(|i| match i % 4 { + 0 | 2 => Some(true), + 1 => None, + _ => Some(false), + }) + .collect::()) + }); + + let builder = ReadPlanBuilder::new(16) + .with_predicate_options( + PredicateOptions::new(Box::new(struct_reader), &mut predicate) + .with_limit(LIMIT, TOTAL_ROWS), + ) + .unwrap(); + + let selection = builder + .selection() + .expect("limit-driven early break must produce a selection"); + + assert_eq!(selection.row_count(), LIMIT); + + let total: usize = selection.iter().map(|s| s.row_count).sum(); + assert_eq!(total, TOTAL_ROWS); + } +} diff --git a/vendor/parquet/src/arrow/arrow_reader/selection/algebra.rs b/vendor/parquet/src/arrow/arrow_reader/selection/algebra.rs new file mode 100644 index 0000000..9c25afb --- /dev/null +++ b/vendor/parquet/src/arrow/arrow_reader/selection/algebra.rs @@ -0,0 +1,828 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Set algebra backing [`RowSelection::and_then`], [`RowSelection::intersection`] +//! and [`RowSelection::union`] +//! +//! Each operation has two implementations, picked by the backing of its +//! operands: a merge of the [`RowSelector`] runs, and a bitwise variant over +//! [`BooleanBuffer`] masks. + +use super::{MaskRunIter, RowSelection, RowSelectionInner, RowSelector}; +use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder}; +use std::cmp::Ordering; +use std::iter::Peekable; + +/// Applies `second` to the rows selected by `first`, both selector-backed. +pub(super) fn and_then_row_selections( + first: &[RowSelector], + second: &[RowSelector], +) -> RowSelection { + let mut selectors = vec![]; + let mut first = first.iter().copied().peekable(); + let mut second = second.iter().copied().peekable(); + and_then_iter(&mut selectors, &mut first, &mut second); + RowSelection::from_selectors(selectors) +} + +/// Applies the mask `second` to the rows selected by the selector-backed `first`. +/// +/// The mask is streamed as [`RowSelector`] runs, so it is never materialized. +pub(super) fn and_then_selectors_with_mask( + first: &[RowSelector], + second: &BooleanBuffer, +) -> RowSelection { + let mut selectors = vec![]; + let mut first = first.iter().copied().peekable(); + let mut second = MaskRunIter::new(second).peekable(); + and_then_iter(&mut selectors, &mut first, &mut second); + RowSelection::from_selectors(selectors) +} + +fn and_then_iter( + selectors: &mut Vec, + first: &mut Peekable, + second: &mut Peekable, +) where + I: Iterator, + J: Iterator, +{ + let mut to_skip = 0; + while let Some(b) = second.peek_mut() { + let a = first + .peek_mut() + .expect("selection exceeds the number of selected rows"); + + if b.row_count == 0 { + second.next().unwrap(); + continue; + } + + if a.row_count == 0 { + first.next().unwrap(); + continue; + } + + if a.skip { + // Records were skipped when producing second + to_skip += a.row_count; + first.next().unwrap(); + continue; + } + + let skip = b.skip; + let to_process = a.row_count.min(b.row_count); + + a.row_count -= to_process; + b.row_count -= to_process; + + match skip { + true => to_skip += to_process, + false => { + if to_skip != 0 { + selectors.push(RowSelector::skip(to_skip)); + to_skip = 0; + } + selectors.push(RowSelector::select(to_process)) + } + } + } + + for v in first { + if v.row_count != 0 { + assert!( + v.skip, + "selection contains less than the number of selected rows" + ); + to_skip += v.row_count + } + } + + if to_skip != 0 { + selectors.push(RowSelector::skip(to_skip)); + } +} + +/// Combine two lists of `RowSelection` return the intersection of them +/// For example: +/// self: NNYYYYNNYYNYN +/// other: NYNNNNNNY +/// +/// returned: NNNNNNNNYYNYN +pub(super) fn intersect_row_selections( + left: &[RowSelector], + right: &[RowSelector], +) -> RowSelection { + let mut l_iter = left.iter().copied().peekable(); + let mut r_iter = right.iter().copied().peekable(); + + let iter = std::iter::from_fn(move || { + loop { + let l = l_iter.peek_mut(); + let r = r_iter.peek_mut(); + + match (l, r) { + (Some(a), _) if a.row_count == 0 => { + l_iter.next().unwrap(); + } + (_, Some(b)) if b.row_count == 0 => { + r_iter.next().unwrap(); + } + (Some(l), Some(r)) => { + return match (l.skip, r.skip) { + // Keep both ranges + (false, false) => { + if l.row_count < r.row_count { + r.row_count -= l.row_count; + l_iter.next() + } else { + l.row_count -= r.row_count; + r_iter.next() + } + } + // skip at least one + _ => { + if l.row_count < r.row_count { + let skip = l.row_count; + r.row_count -= l.row_count; + l_iter.next(); + Some(RowSelector::skip(skip)) + } else { + let skip = r.row_count; + l.row_count -= skip; + r_iter.next(); + Some(RowSelector::skip(skip)) + } + } + }; + } + (Some(_), None) => return l_iter.next(), + (None, Some(_)) => return r_iter.next(), + (None, None) => return None, + } + } + }); + + iter.collect() +} + +/// Combine two lists of `RowSelector` return the union of them +/// For example: +/// self: NNYYYYNNYYNYN +/// other: NYNNNNNNY +/// +/// returned: NYYYYYNNYYNYN +/// +/// This can be removed from here once RowSelection::union is in parquet::arrow +pub(super) fn union_row_selections(left: &[RowSelector], right: &[RowSelector]) -> RowSelection { + let mut l_iter = left.iter().copied().peekable(); + let mut r_iter = right.iter().copied().peekable(); + + let iter = std::iter::from_fn(move || { + loop { + let l = l_iter.peek_mut(); + let r = r_iter.peek_mut(); + + match (l, r) { + (Some(a), _) if a.row_count == 0 => { + l_iter.next().unwrap(); + } + (_, Some(b)) if b.row_count == 0 => { + r_iter.next().unwrap(); + } + (Some(l), Some(r)) => { + return match (l.skip, r.skip) { + // Skip both ranges + (true, true) => { + if l.row_count < r.row_count { + let skip = l.row_count; + r.row_count -= l.row_count; + l_iter.next(); + Some(RowSelector::skip(skip)) + } else { + let skip = r.row_count; + l.row_count -= skip; + r_iter.next(); + Some(RowSelector::skip(skip)) + } + } + // Keep rows from left + (false, true) => { + if l.row_count < r.row_count { + r.row_count -= l.row_count; + l_iter.next() + } else { + let r_row_count = r.row_count; + l.row_count -= r_row_count; + r_iter.next(); + Some(RowSelector::select(r_row_count)) + } + } + // Keep rows from right + (true, false) => { + if l.row_count < r.row_count { + let l_row_count = l.row_count; + r.row_count -= l_row_count; + l_iter.next(); + Some(RowSelector::select(l_row_count)) + } else { + l.row_count -= r.row_count; + r_iter.next() + } + } + // Keep at least one + _ => { + if l.row_count < r.row_count { + r.row_count -= l.row_count; + l_iter.next() + } else { + l.row_count -= r.row_count; + r_iter.next() + } + } + }; + } + (Some(_), None) => return l_iter.next(), + (None, Some(_)) => return r_iter.next(), + (None, None) => return None, + } + } + }); + + iter.collect() +} + +/// Bitwise AND of two mask-backed selections. Longer side's tail passes through. +pub(super) fn intersect_masks(l: &BooleanBuffer, r: &BooleanBuffer) -> BooleanBuffer { + if l.len() == r.len() { + return l & r; + } + let common = l.len().min(r.len()); + let head = &l.slice(0, common) & &r.slice(0, common); + let (longer, longer_len) = if l.len() > r.len() { + (l, l.len()) + } else { + (r, r.len()) + }; + let tail = longer.slice(common, longer_len - common); + let mut builder = BooleanBufferBuilder::new(longer_len); + builder.append_buffer(&head); + builder.append_buffer(&tail); + builder.finish() +} + +/// Bitwise OR of two mask-backed selections. Longer side's tail passes through. +pub(super) fn union_masks(l: &BooleanBuffer, r: &BooleanBuffer) -> BooleanBuffer { + if l.len() == r.len() { + return l | r; + } + let common = l.len().min(r.len()); + let head = &l.slice(0, common) | &r.slice(0, common); + let (longer, longer_len) = if l.len() > r.len() { + (l, l.len()) + } else { + (r, r.len()) + }; + let tail = longer.slice(common, longer_len - common); + let mut builder = BooleanBufferBuilder::new(longer_len); + builder.append_buffer(&head); + builder.append_buffer(&tail); + builder.finish() +} + +/// Applies `other` to the selected rows of `mask`, preserving the original row domain. +pub(super) fn and_then_mask(mask: &BooleanBuffer, other: &RowSelection) -> BooleanBuffer { + match &other.inner { + RowSelectionInner::Mask(other_mask) => and_then_masks(mask, other_mask.mask()), + RowSelectionInner::Selectors(selectors) => { + and_then_mask_from_selectors(mask, selectors.iter().copied()) + } + } +} + +fn and_then_mask_from_selectors(mask: &BooleanBuffer, other: I) -> BooleanBuffer +where + I: IntoIterator, +{ + let mut builder = BooleanBufferBuilder::new(mask.len()); + let mut other_iter = other.into_iter(); + let mut current = other_iter.next(); + let mut cursor = 0usize; + + // Iterate only over the set positions in `mask`; the gaps of unset bits + // are filled in bulk with `append_n` instead of bit-by-bit. + for set_idx in mask.set_indices() { + if set_idx > cursor { + builder.append_n(set_idx - cursor, false); + } + cursor = set_idx + 1; + + while current.as_ref().is_some_and(|s| s.row_count == 0) { + current = other_iter.next(); + } + let selector = current + .as_mut() + .expect("selection contains less than the number of selected rows"); + let selected = !selector.skip; + selector.row_count -= 1; + builder.append(selected); + } + if cursor < mask.len() { + builder.append_n(mask.len() - cursor, false); + } + + if current.is_some_and(|s| s.row_count != 0) || other_iter.any(|s| s.row_count != 0) { + panic!("selection exceeds the number of selected rows"); + } + + builder.finish() +} + +fn and_then_masks(mask: &BooleanBuffer, other: &BooleanBuffer) -> BooleanBuffer { + let selected_count = mask.count_set_bits(); + match other.len().cmp(&selected_count) { + Ordering::Less => panic!("selection contains less than the number of selected rows"), + Ordering::Greater => panic!("selection exceeds the number of selected rows"), + Ordering::Equal => {} + } + + let other_true_count = other.count_set_bits(); + if other_true_count == 0 { + return BooleanBuffer::new_unset(mask.len()); + } + if other_true_count == selected_count { + return mask.clone(); + } + + let mut builder = BooleanBufferBuilder::new(mask.len()); + let mut outer_set_indices = mask.set_indices(); + let mut next_selected_ordinal = 0usize; + let mut cursor = 0usize; + + for selected_ordinal in other.set_indices() { + let skip = selected_ordinal - next_selected_ordinal; + let set_idx = outer_set_indices + .nth(skip) + .expect("validated other length matches selected row count"); + if set_idx > cursor { + builder.append_n(set_idx - cursor, false); + } + builder.append(true); + cursor = set_idx + 1; + next_selected_ordinal = selected_ordinal + 1; + } + + if cursor < mask.len() { + builder.append_n(mask.len() - cursor, false); + } + + builder.finish() +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::BooleanArray; + use rand::{Rng, rng}; + + #[test] + fn test_and() { + let mut a = RowSelection::from(vec![ + RowSelector::skip(12), + RowSelector::select(23), + RowSelector::skip(3), + RowSelector::select(5), + ]); + + let b = RowSelection::from(vec![ + RowSelector::select(5), + RowSelector::skip(4), + RowSelector::select(15), + RowSelector::skip(4), + ]); + + let mut expected = RowSelection::from(vec![ + RowSelector::skip(12), + RowSelector::select(5), + RowSelector::skip(4), + RowSelector::select(14), + RowSelector::skip(3), + RowSelector::select(1), + RowSelector::skip(4), + ]); + + assert_eq!(a.and_then(&b), expected); + + a.split_off(7); + expected.split_off(7); + assert_eq!(a.and_then(&b), expected); + + let a = RowSelection::from(vec![RowSelector::select(5), RowSelector::skip(3)]); + + let b = RowSelection::from(vec![ + RowSelector::select(2), + RowSelector::skip(1), + RowSelector::select(1), + RowSelector::skip(1), + ]); + + assert_eq!( + a.and_then(&b).selectors(), + vec![ + RowSelector::select(2), + RowSelector::skip(1), + RowSelector::select(1), + RowSelector::skip(4) + ] + ); + } + + #[test] + #[should_panic(expected = "selection exceeds the number of selected rows")] + fn test_and_longer() { + let a = RowSelection::from(vec![ + RowSelector::select(3), + RowSelector::skip(33), + RowSelector::select(3), + RowSelector::skip(33), + ]); + let b = RowSelection::from(vec![RowSelector::select(36)]); + a.and_then(&b); + } + + #[test] + #[should_panic(expected = "selection contains less than the number of selected rows")] + fn test_and_shorter() { + let a = RowSelection::from(vec![ + RowSelector::select(3), + RowSelector::skip(33), + RowSelector::select(3), + RowSelector::skip(33), + ]); + let b = RowSelection::from(vec![RowSelector::select(3)]); + a.and_then(&b); + } + + #[test] + fn test_intersect_row_selection_and_combine() { + // a size equal b size + let a = vec![ + RowSelector::select(5), + RowSelector::skip(4), + RowSelector::select(1), + ]; + let b = vec![ + RowSelector::select(8), + RowSelector::skip(1), + RowSelector::select(1), + ]; + + let res = intersect_row_selections(&a, &b); + assert_eq!( + res.selectors(), + vec![ + RowSelector::select(5), + RowSelector::skip(4), + RowSelector::select(1), + ], + ); + + // a size larger than b size + let a = vec![ + RowSelector::select(3), + RowSelector::skip(33), + RowSelector::select(3), + RowSelector::skip(33), + ]; + let b = vec![RowSelector::select(36), RowSelector::skip(36)]; + let res = intersect_row_selections(&a, &b); + assert_eq!( + res.selectors(), + vec![RowSelector::select(3), RowSelector::skip(69)] + ); + + // a size less than b size + let a = vec![RowSelector::select(3), RowSelector::skip(7)]; + let b = vec![ + RowSelector::select(2), + RowSelector::skip(2), + RowSelector::select(2), + RowSelector::skip(2), + RowSelector::select(2), + ]; + let res = intersect_row_selections(&a, &b); + assert_eq!( + res.selectors(), + vec![RowSelector::select(2), RowSelector::skip(8)] + ); + + let a = vec![RowSelector::select(3), RowSelector::skip(7)]; + let b = vec![ + RowSelector::select(2), + RowSelector::skip(2), + RowSelector::select(2), + RowSelector::skip(2), + RowSelector::select(2), + ]; + let res = intersect_row_selections(&a, &b); + assert_eq!( + res.selectors(), + vec![RowSelector::select(2), RowSelector::skip(8)] + ); + } + + #[test] + fn test_and_fuzz() { + let mut rand = rng(); + for _ in 0..100 { + let a_len = rand.random_range(10..100); + let a_bools: Vec<_> = (0..a_len).map(|_| rand.random_bool(0.2)).collect(); + let a = RowSelection::from_filters(&[BooleanArray::from(a_bools.clone())]); + + let b_len: usize = a_bools.iter().map(|x| *x as usize).sum(); + let b_bools: Vec<_> = (0..b_len).map(|_| rand.random_bool(0.8)).collect(); + let b = RowSelection::from_filters(&[BooleanArray::from(b_bools.clone())]); + + let mut expected_bools = vec![false; a_len]; + + let mut iter_b = b_bools.iter(); + for (idx, b) in a_bools.iter().enumerate() { + if *b && *iter_b.next().unwrap() { + expected_bools[idx] = true; + } + } + + let expected = RowSelection::from_filters(&[BooleanArray::from(expected_bools)]); + + let total_rows: usize = expected.selectors().iter().map(|s| s.row_count).sum(); + assert_eq!(a_len, total_rows); + + assert_eq!(a.and_then(&b), expected); + } + } + + #[test] + fn test_intersection() { + let selection = RowSelection::from(vec![RowSelector::select(1048576)]); + let result = selection.intersection(&selection); + assert_eq!(result, selection); + + let a = RowSelection::from(vec![ + RowSelector::skip(10), + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(20), + ]); + + let b = RowSelection::from(vec![ + RowSelector::skip(20), + RowSelector::select(20), + RowSelector::skip(10), + ]); + + let result = a.intersection(&b); + assert_eq!( + result.selectors(), + vec![ + RowSelector::skip(30), + RowSelector::select(10), + RowSelector::skip(10) + ] + ); + } + + #[test] + fn test_union() { + let selection = RowSelection::from(vec![RowSelector::select(1048576)]); + let result = selection.union(&selection); + assert_eq!(result, selection); + + // NYNYY + let a = RowSelection::from(vec![ + RowSelector::skip(10), + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(20), + ]); + + // NNYYNYN + let b = RowSelection::from(vec![ + RowSelector::skip(20), + RowSelector::select(20), + RowSelector::skip(10), + RowSelector::select(10), + RowSelector::skip(10), + ]); + + let result = a.union(&b); + + // NYYYYYN + assert_eq!( + result.iter().copied().collect::>(), + vec![ + RowSelector::skip(10), + RowSelector::select(50), + RowSelector::skip(10), + ] + ); + } + + #[test] + fn test_mask_and_then_preserves_backing() { + let outer_bits = vec![false, true, true, false, true, false, true]; + let inner_bits = vec![true, false, true, false]; + let outer_mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(outer_bits.clone())); + let inner = RowSelection::from_filters(&[BooleanArray::from(inner_bits.clone())]); + + let result = outer_mask.and_then(&inner); + assert!(result.as_mask().is_some()); + + let outer_selectors = RowSelection::from_filters(&[BooleanArray::from(outer_bits)]); + let expected = outer_selectors.and_then(&inner); + assert_eq!(result, expected); + + let result_mask = result.as_mask().unwrap(); + let actual_bits: Vec<_> = (0..result_mask.len()) + .map(|i| result_mask.value(i)) + .collect(); + assert_eq!( + actual_bits, + vec![false, true, false, false, true, false, false] + ); + } + + #[test] + fn test_mask_and_then_mask_preserves_backing() { + let outer_bits = vec![false, true, true, false, true, false, true, true]; + let inner_bits = vec![false, true, false, true, false]; + let outer_mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(outer_bits.clone())); + let inner_mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(inner_bits)); + + let result = outer_mask.and_then(&inner_mask); + assert!(result.as_mask().is_some()); + + let outer_selectors = RowSelection::from_filters(&[BooleanArray::from(outer_bits)]); + let inner_selectors = RowSelection::from_filters(&[BooleanArray::from(vec![ + false, true, false, true, false, + ])]); + assert_eq!(result, outer_selectors.and_then(&inner_selectors)); + + let result_mask = result.as_mask().unwrap(); + let actual_bits: Vec<_> = (0..result_mask.len()) + .map(|i| result_mask.value(i)) + .collect(); + assert_eq!( + actual_bits, + vec![false, false, true, false, false, false, true, false] + ); + } + + #[test] + fn test_selector_and_then_mask() { + let outer = + RowSelection::from_filters(&[BooleanArray::from(vec![false, true, true, false, true])]); + let inner = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true])); + + let result = outer.and_then(&inner); + assert!(result.as_mask().is_none()); + assert_eq!( + result, + RowSelection::from_filters(&[BooleanArray::from(vec![ + false, true, false, false, true, + ])]) + ); + } + + #[test] + fn test_mask_and_then_none_selected_returns_all_unset() { + let outer = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ + false, true, true, false, true, + ])); + let inner = + RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![false, false, false])); + + let result = outer.and_then(&inner); + let mask = result.as_mask().unwrap(); + assert_eq!(mask.len(), 5); + assert_eq!(mask.count_set_bits(), 0); + } + + #[test] + fn test_mask_intersection_uses_bitwise() { + let a_bits = vec![true, true, false, true, false, true]; + let b_bits = vec![true, false, true, true, true, false]; + let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits.clone())); + let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits.clone())); + + let r = a.intersection(&b); + assert!(r.as_mask().is_some()); + + let expected: Vec = a_bits.iter().zip(&b_bits).map(|(x, y)| *x && *y).collect(); + let expected_sel = RowSelection::from_filters(&[BooleanArray::from(expected)]); + assert_eq!(r, expected_sel); + } + + #[test] + fn test_mask_union_uses_bitwise() { + let a_bits = vec![true, false, false, true, false, false]; + let b_bits = vec![false, true, false, false, true, false]; + let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits.clone())); + let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits.clone())); + + let r = a.union(&b); + assert!(r.as_mask().is_some()); + + let expected: Vec = a_bits.iter().zip(&b_bits).map(|(x, y)| *x || *y).collect(); + let expected_sel = RowSelection::from_filters(&[BooleanArray::from(expected)]); + assert_eq!(r, expected_sel); + } + + #[test] + fn test_mixed_mask_selector_intersection_and_union() { + let mask_bits = vec![true, false, true, false, true, false]; + let selector_bits = vec![false, true, true, false, false, true]; + let mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(mask_bits.clone())); + let selectors = RowSelection::from_filters(&[BooleanArray::from(selector_bits.clone())]); + + let intersection_bits: Vec<_> = mask_bits + .iter() + .zip(&selector_bits) + .map(|(x, y)| *x && *y) + .collect(); + let expected_intersection = + RowSelection::from_filters(&[BooleanArray::from(intersection_bits)]); + assert_eq!(mask.intersection(&selectors), expected_intersection); + assert_eq!(selectors.intersection(&mask), expected_intersection); + + let union_bits: Vec<_> = mask_bits + .iter() + .zip(&selector_bits) + .map(|(x, y)| *x || *y) + .collect(); + let expected_union = RowSelection::from_filters(&[BooleanArray::from(union_bits)]); + assert_eq!(mask.union(&selectors), expected_union); + assert_eq!(selectors.union(&mask), expected_union); + } + + #[test] + fn test_mask_intersection_uneven_passes_tail_through() { + let a_bits = vec![true, true, true, true, true]; + let b_bits = vec![true, false, true]; + let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits)); + let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits)); + + let r = a.intersection(&b); + let r_mask = r.as_mask().unwrap(); + assert_eq!(r_mask.len(), 5); + let bits: Vec = (0..5).map(|i| r_mask.value(i)).collect(); + assert_eq!(bits, vec![true, false, true, true, true]); + + // Swapped operands: the right side is longer and its tail passes through. + let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true])); + let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ + true, true, true, false, true, + ])); + let r = a.intersection(&b); + let r_mask = r.as_mask().unwrap(); + assert_eq!(r_mask.len(), 5); + let bits: Vec = (0..5).map(|i| r_mask.value(i)).collect(); + assert_eq!(bits, vec![true, false, true, false, true]); + } + + #[test] + fn test_mask_union_uneven_passes_tail_through() { + let a_bits = vec![true, false, true]; + let b_bits = vec![false, true, false, true, false]; + let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits)); + let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits)); + + let r = a.union(&b); + let r_mask = r.as_mask().unwrap(); + assert_eq!(r_mask.len(), 5); + let bits: Vec = (0..5).map(|i| r_mask.value(i)).collect(); + assert_eq!(bits, vec![true, true, true, true, false]); + + let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ + false, true, false, false, true, + ])); + let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, false])); + let r = a.union(&b); + let r_mask = r.as_mask().unwrap(); + let bits: Vec = (0..5).map(|i| r_mask.value(i)).collect(); + assert_eq!(bits, vec![true, true, false, false, true]); + } +} diff --git a/vendor/parquet/src/arrow/arrow_reader/selection/boolean.rs b/vendor/parquet/src/arrow/arrow_reader/selection/boolean.rs new file mode 100644 index 0000000..2ac879e --- /dev/null +++ b/vendor/parquet/src/arrow/arrow_reader/selection/boolean.rs @@ -0,0 +1,746 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The bitmap backed representation of a [`RowSelection`] and the primitives +//! operating on it: conversion to and from the run length ([`RowSelector`]) +//! form, and the transforms backing `split_off`, `trim`, `offset` and `limit`. +//! +//! The bitwise set algebra lives in the `algebra` module. +//! +//! [`RowSelection`]: crate::arrow::arrow_reader::RowSelection + +use super::RowSelector; +use arrow_buffer::bit_iterator::BitSliceIterator; +use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, Buffer}; +use std::sync::OnceLock; + +/// Mask-backed [`RowSelection`] storage. +/// +/// `selectors` is only populated if callers use the borrowed +/// [`RowSelection::iter`] compatibility API. Internal paths that can stream or +/// consume the bitmap avoid this cache. +/// +/// `count` caches the popcount; [`RowSelection::split_off`] propagates it to +/// both halves so repeated `row_count()` calls do not rescan the bitmap. +/// +/// [`RowSelection`]: crate::arrow::arrow_reader::RowSelection +/// [`RowSelection::iter`]: crate::arrow::arrow_reader::RowSelection::iter +/// [`RowSelection::split_off`]: crate::arrow::arrow_reader::RowSelection::split_off +#[derive(Debug)] +pub(crate) struct MaskSelection { + mask: BooleanBuffer, + selectors: OnceLock>, + count: OnceLock, +} + +impl MaskSelection { + pub(super) fn new(mask: BooleanBuffer) -> Self { + Self { + mask, + selectors: OnceLock::new(), + count: OnceLock::new(), + } + } + + /// Create a selection whose selected-row count is already known. + pub(super) fn with_count(mask: BooleanBuffer, count: usize) -> Self { + debug_assert!(count <= mask.len()); + let cell = OnceLock::new(); + let _ = cell.set(count); + Self { + mask, + selectors: OnceLock::new(), + count: cell, + } + } + + pub(crate) fn mask(&self) -> &BooleanBuffer { + &self.mask + } + + pub(crate) fn into_mask(self) -> BooleanBuffer { + let Self { mask, .. } = self; + mask + } + + /// Number of selected rows, computed once and cached. + pub(super) fn count(&self) -> usize { + *self.count.get_or_init(|| self.mask.count_set_bits()) + } + + /// The cached selected-row count, if it has been computed. + pub(super) fn cached_count(&self) -> Option { + self.count.get().copied() + } + + pub(super) fn selectors(&self) -> &[RowSelector] { + self.selectors + .get_or_init(|| mask_to_selectors(&self.mask)) + .as_slice() + } +} + +impl Clone for MaskSelection { + fn clone(&self) -> Self { + // Drop the selector cache but keep the cheap count cache. + Self { + mask: self.mask.clone(), + selectors: OnceLock::new(), + count: self.count.clone(), + } + } +} + +/// Streaming RLE view of a [`BooleanBuffer`], yielding owned [`RowSelector`]s +/// without allocation. +/// +/// Useful as a zero-cost alternative to [`RowSelection::iter`] for mask-backed +/// selections, via [`RowSelection::as_mask`]: +/// +/// ```ignore +/// if let Some(mask) = selection.as_mask() { +/// for run in MaskRunIter::new(mask) { ... } +/// } +/// ``` +/// +/// [`RowSelection::iter`]: crate::arrow::arrow_reader::RowSelection::iter +/// [`RowSelection::as_mask`]: crate::arrow::arrow_reader::RowSelection::as_mask +#[derive(Debug)] +pub struct MaskRunIter<'a> { + slices: BitSliceIterator<'a>, + cursor: usize, + total: usize, + pending: Option, + finished: bool, +} + +impl<'a> MaskRunIter<'a> { + /// Create a streaming RLE iterator over a [`BooleanBuffer`]. + pub fn new(mask: &'a BooleanBuffer) -> Self { + Self { + slices: mask.set_slices(), + cursor: 0, + total: mask.len(), + pending: None, + finished: false, + } + } +} + +impl Iterator for MaskRunIter<'_> { + type Item = RowSelector; + + fn next(&mut self) -> Option { + if let Some(p) = self.pending.take() { + return Some(p); + } + if self.finished { + return None; + } + match self.slices.next() { + Some((start, end)) => { + let select = RowSelector::select(end - start); + if start > self.cursor { + let skip = RowSelector::skip(start - self.cursor); + self.pending = Some(select); + self.cursor = end; + Some(skip) + } else { + self.cursor = end; + Some(select) + } + } + None => { + self.finished = true; + if self.cursor < self.total { + let skip = RowSelector::skip(self.total - self.cursor); + self.cursor = self.total; + Some(skip) + } else { + None + } + } + } + } +} + +/// Materialize a [`BooleanBuffer`] into its RLE form. +pub(crate) fn mask_to_selectors(mask: &BooleanBuffer) -> Vec { + let total_rows = mask.len(); + if total_rows == 0 { + return Vec::new(); + } + let mut selectors: Vec = Vec::new(); + let mut last_end = 0; + for (start, end) in mask.set_slices() { + if start > last_end { + selectors.push(RowSelector::skip(start - last_end)); + } + selectors.push(RowSelector::select(end - start)); + last_end = end; + } + if last_end != total_rows { + selectors.push(RowSelector::skip(total_rows - last_end)); + } + selectors +} + +/// Returns whether `mask` contains at least `min_runs` alternating set/unset runs. +/// +/// Stops as soon as the requested number of runs is found, avoiding a full scan +/// when callers only need to know whether a boundary has been crossed. +pub(super) fn mask_has_at_least_runs(mask: &BooleanBuffer, min_runs: usize) -> bool { + if min_runs == 0 { + return true; + } + + let total_rows = mask.len(); + if total_rows == 0 { + return false; + } + + let mut run_count = 0; + let mut last_end = 0; + for (start, end) in mask.set_slices() { + run_count += usize::from(start > last_end) + 1; + if run_count >= min_runs { + return true; + } + last_end = end; + } + + run_count + usize::from(last_end < total_rows) >= min_runs +} + +/// Split a mask into `(head, tail)` at `row_count`, preserving an empty mask tail +/// when the split point is past the end. +pub(super) fn split_off_mask( + mask: BooleanBuffer, + row_count: usize, +) -> (BooleanBuffer, BooleanBuffer) { + let total = mask.len(); + if row_count >= total { + return (mask, BooleanBuffer::new_unset(0)); + } + + let head = mask.slice(0, row_count); + let tail = mask.slice(row_count, total - row_count); + (head, tail) +} + +/// Position of the highest set bit in `mask`, scanning bytes from the end. +fn last_set_bit_position(mask: &BooleanBuffer) -> Option { + let values = mask.values(); + let offset = mask.offset(); + let end = offset + mask.len(); + for byte_idx in (offset / 8..end.div_ceil(8)).rev() { + let byte_start = byte_idx * 8; + let mut byte = values[byte_idx]; + if end - byte_start < 8 { + byte &= (1u8 << (end - byte_start)) - 1; + } + if byte_start < offset { + byte &= !((1u8 << (offset - byte_start)) - 1); + } + if byte != 0 { + return Some(byte_start + 7 - byte.leading_zeros() as usize - offset); + } + } + None +} + +/// Trims trailing unset bits from a mask-backed selection. +pub(super) fn trim_mask(mask: &BooleanBuffer) -> Option { + let len = mask.len(); + // Fast path: final bit set means there is nothing to trim. + if len == 0 || mask.value(len - 1) { + return None; + } + let new_len = last_set_bit_position(mask).map_or(0, |pos| pos + 1); + Some(mask.slice(0, new_len)) +} + +/// Skips the first `offset` selected rows of a mask-backed selection. +/// `popcount` is the caller's (possibly cached) set-bit count of `mask`. +pub(super) fn offset_mask(mask: BooleanBuffer, offset: usize, popcount: usize) -> BooleanBuffer { + if offset >= popcount { + return BooleanBuffer::new_unset(0); + } + // Position one past the `offset`-th set bit, i.e. the index of the first + // selected row to keep. + let pos = mask.find_nth_set_bit_position(0, offset); + let mut builder = BooleanBufferBuilder::new(mask.len()); + builder.append_n(pos, false); + builder.append_buffer(&mask.slice(pos, mask.len() - pos)); + builder.finish() +} + +/// Keeps only the first `limit` selected rows of a mask-backed selection. +pub(super) fn limit_mask(mask: BooleanBuffer, limit: usize) -> BooleanBuffer { + // `find_nth_set_bit_position` returns `mask.len()` when there are fewer + // than `limit` set bits, so the slice naturally degrades to the original + // mask in that case. + let cut = mask.find_nth_set_bit_position(0, limit); + mask.slice(0, cut) +} + +/// Set bits `[start, start + len)` in a zero-initialized little-endian bitmap. +fn set_bit_run(buf: &mut [u8], start: usize, len: usize) { + if len == 0 { + return; + } + let end = start + len; + let first_byte = start / 8; + let last_byte = (end - 1) / 8; + let start_mask = 0xFFu8 << (start % 8); + let end_mask = 0xFFu8 >> (8 - (end - last_byte * 8)); + if first_byte == last_byte { + buf[first_byte] |= start_mask & end_mask; + } else { + buf[first_byte] |= start_mask; + buf[first_byte + 1..last_byte].fill(0xFF); + buf[last_byte] |= end_mask; + } +} + +/// Build a bitmap from a selector sequence by filling bytes directly. +/// +/// This sits on the read hot path (`Mask` strategy over a selector-backed +/// selection) where per-selector `append_n` calls are too slow. +pub(super) fn boolean_mask_from_selectors(selectors: &[RowSelector]) -> BooleanBuffer { + let total_rows: usize = selectors.iter().map(|s| s.row_count).sum(); + let mut buf = vec![0u8; total_rows.div_ceil(8)]; + let mut position = 0usize; + for selector in selectors { + if !selector.skip { + set_bit_run(&mut buf, position, selector.row_count); + } + position += selector.row_count; + } + BooleanBuffer::new(Buffer::from(buf), 0, total_rows) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::arrow_reader::selection::{RowSelection, RowSelectionInner}; + use arrow_array::BooleanArray; + use rand::{Rng, rng}; + + #[test] + fn test_mask_iter_yields_borrowed_selectors() { + let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ + false, false, true, true, false, true, false, false, + ])); + + let borrowed: Vec<&RowSelector> = selection.iter().collect(); + assert_eq!( + borrowed, + vec![ + &RowSelector::skip(2), + &RowSelector::select(2), + &RowSelector::skip(1), + &RowSelector::select(1), + &RowSelector::skip(2), + ] + ); + } + + #[test] + fn test_mask_iter_clone_drops_cache() { + let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ + false, false, true, true, false, true, false, false, + ])); + + let _ = selection.iter().count(); + match &selection.inner { + RowSelectionInner::Mask(m) => assert!(m.selectors.get().is_some()), + _ => unreachable!(), + } + + let cloned = selection.clone(); + match &cloned.inner { + RowSelectionInner::Mask(m) => assert!(m.selectors.get().is_none()), + _ => unreachable!(), + } + + let round_tripped: Vec = cloned.iter().copied().collect(); + assert_eq!( + round_tripped, + vec![ + RowSelector::skip(2), + RowSelector::select(2), + RowSelector::skip(1), + RowSelector::select(1), + RowSelector::skip(2), + ] + ); + } + + #[test] + fn test_mask_run_iter_streams_without_cache() { + let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ + false, false, true, true, false, true, false, false, + ])); + let mut iter = MaskRunIter::new(selection.as_mask().unwrap()); + + assert_eq!(iter.next(), Some(RowSelector::skip(2))); + assert_eq!(iter.next(), Some(RowSelector::select(2))); + assert_eq!(iter.next(), Some(RowSelector::skip(1))); + assert_eq!(iter.next(), Some(RowSelector::select(1))); + assert_eq!(iter.next(), Some(RowSelector::skip(2))); + assert_eq!(iter.next(), None); + assert_eq!(iter.next(), None); + + let selection = + RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, true, false])); + let mut iter = MaskRunIter::new(selection.as_mask().unwrap()); + assert_eq!(iter.next(), Some(RowSelector::select(2))); + assert_eq!(iter.next(), Some(RowSelector::skip(1))); + assert_eq!(iter.next(), None); + } + + #[test] + fn test_from_boolean_buffer() { + let bits = vec![ + false, false, true, true, false, true, false, false, true, false, false, false, false, + false, false, true, + ]; + let buf = BooleanBuffer::from(bits.clone()); + let selection = RowSelection::from_boolean_buffer(buf.clone()); + + assert!(selection.as_mask().is_some()); + assert_eq!(selection.row_count(), 5); + assert_eq!(selection.skipped_row_count(), 11); + assert!(selection.selects_any()); + + let from_filters = RowSelection::from_filters(&[BooleanArray::from(bits)]); + assert_eq!(selection, from_filters); + + let bits_tail = vec![true, false, true, false, false, false]; + let trimmed = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits_tail)).trim(); + assert!(trimmed.as_mask().is_some()); + assert_eq!(trimmed.as_mask().unwrap().len(), 3); + } + + #[test] + fn test_from_boolean_buffer_empty() { + let empty = RowSelection::from_boolean_buffer(BooleanBuffer::from(Vec::::new())); + assert!(empty.as_mask().is_some()); + assert_eq!(empty.row_count(), 0); + assert_eq!(empty.skipped_row_count(), 0); + assert!(!empty.selects_any()); + assert!(empty.selectors().is_empty()); + } + + #[test] + fn test_from_boolean_buffer_all_unset_does_not_select() { + let all_zero = RowSelection::from_boolean_buffer(BooleanBuffer::new_unset(1024)); + assert!(all_zero.as_mask().is_some()); + assert!(!all_zero.selects_any()); + assert_eq!(all_zero.row_count(), 0); + assert_eq!(all_zero.skipped_row_count(), 1024); + } + + #[test] + fn test_from_boolean_buffer_via_from_impl() { + let buf = BooleanBuffer::from(vec![true, false, true, true]); + let a = RowSelection::from(buf.clone()); + let b = RowSelection::from_boolean_buffer(buf); + assert_eq!(a, b); + assert!(a.as_mask().is_some()); + } + + #[test] + fn test_mask_backing_clone_preserves_backing() { + let buf = BooleanBuffer::from(vec![true, false, true]); + let original = RowSelection::from_boolean_buffer(buf); + let cloned = original.clone(); + assert!(cloned.as_mask().is_some()); + assert_eq!(original, cloned); + } + + #[test] + fn test_mask_backing_mutation_equivalence() { + let bits = vec![true, true, false, false, true, false, true, true]; + + let from_mask = { + let mut s = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone())); + let split = s.split_off(3); + (split, s) + }; + let from_selectors = { + let mut s = RowSelection::from_filters(&[BooleanArray::from(bits.clone())]); + let split = s.split_off(3); + (split, s) + }; + assert_eq!(from_mask.0, from_selectors.0); + assert_eq!(from_mask.1, from_selectors.1); + assert!(from_mask.0.as_mask().is_some()); + assert!(from_mask.1.as_mask().is_some()); + + let limited_mask = + RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone())).limit(3); + let limited_sel = RowSelection::from_filters(&[BooleanArray::from(bits.clone())]).limit(3); + assert!(limited_mask.as_mask().is_some()); + assert_eq!(limited_mask, limited_sel); + + let offset_mask = + RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone())).offset(2); + let offset_sel = RowSelection::from_filters(&[BooleanArray::from(bits)]).offset(2); + assert!(offset_mask.as_mask().is_some()); + assert_eq!(offset_mask, offset_sel); + } + + #[test] + fn test_mask_backing_fuzz_equivalence() { + let mut rand = rng(); + for _ in 0..100 { + let len = rand.random_range(0..200); + let bits: Vec<_> = (0..len).map(|_| rand.random_bool(0.35)).collect(); + + let from_mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone())); + let from_filters = RowSelection::from_filters(&[BooleanArray::from(bits.clone())]); + + assert_eq!(from_mask, from_filters); + assert_eq!(from_mask.row_count(), from_filters.row_count()); + assert_eq!( + from_mask.skipped_row_count(), + from_filters.skipped_row_count() + ); + assert_eq!(from_mask.selects_any(), from_filters.selects_any()); + + let inner_len: usize = bits.iter().map(|b| *b as usize).sum(); + let inner_bits: Vec<_> = (0..inner_len).map(|_| rand.random_bool(0.7)).collect(); + let inner = RowSelection::from_filters(&[BooleanArray::from(inner_bits.clone())]); + let inner_mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(inner_bits)); + let and_then_mask = from_mask.and_then(&inner); + let and_then_both_masks = from_mask.and_then(&inner_mask); + assert!(and_then_mask.as_mask().is_some()); + assert!(and_then_both_masks.as_mask().is_some()); + assert_eq!(and_then_mask, from_filters.and_then(&inner)); + assert_eq!(and_then_both_masks, and_then_mask); + } + } + + #[test] + fn test_mask_offset_past_end_preserves_empty_mask_backing() { + let selection = + RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true])) + .offset(2); + + assert!(selection.as_mask().is_some()); + assert_eq!(selection.as_mask().unwrap().len(), 0); + assert_eq!(selection.row_count(), 0); + assert_eq!(selection.skipped_row_count(), 0); + } + + #[test] + fn test_mask_limit_truncates_at_nth_selected_row() { + let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ + false, true, false, true, false, true, false, + ])) + .limit(2); + + let mask = selection.as_mask().unwrap(); + assert_eq!(mask.len(), 4); + let actual_bits: Vec<_> = (0..mask.len()).map(|i| mask.value(i)).collect(); + assert_eq!(actual_bits, vec![false, true, false, true]); + } + + #[test] + fn test_mask_split_off_preserves_backing() { + let bits: Vec = (0..40).map(|i| i % 3 == 0).collect(); + let mut s = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone())); + let head = s.split_off(15); + + assert!(head.as_mask().is_some()); + assert!(s.as_mask().is_some()); + + let head_sel = RowSelection::from_filters(&[BooleanArray::from(bits[..15].to_vec())]); + let tail_sel = RowSelection::from_filters(&[BooleanArray::from(bits[15..].to_vec())]); + assert_eq!(head, head_sel); + assert_eq!(s, tail_sel); + } + + #[test] + fn test_mask_split_off_past_end_returns_whole() { + let bits = vec![true, false, true]; + let mut s = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone())); + let head = s.split_off(100); + + assert!(head.as_mask().is_some()); + assert_eq!(head.as_mask().unwrap().len(), 3); + // `self` keeps its mask backing and is left empty. + assert!(s.as_mask().is_some()); + assert_eq!(s.as_mask().unwrap().len(), 0); + assert_eq!(s.row_count(), 0); + assert_eq!(s.skipped_row_count(), 0); + } + + #[test] + fn test_mask_offset_exceeds_selected_returns_empty() { + let s = + RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, true, false, true])); + let r = s.offset(10); + assert_eq!(r.row_count(), 0); + assert_eq!(r.skipped_row_count(), 0); + + let from_selectors = + RowSelection::from_filters(&[BooleanArray::from(vec![true, true, false, true])]) + .offset(10); + assert_eq!(r, from_selectors); + } + + #[test] + fn test_mask_limit_exceeds_selected_returns_all() { + let bits = vec![true, true, false, true]; + let s = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone())); + let r = s.limit(10); + assert_eq!(r.row_count(), 3); + + let from_selectors = RowSelection::from_filters(&[BooleanArray::from(bits)]).limit(10); + assert_eq!(r, from_selectors); + } + + #[test] + fn test_mask_trim_all_zero_collapses_to_empty() { + let s = RowSelection::from_boolean_buffer(BooleanBuffer::new_unset(128)); + let trimmed = s.trim(); + assert!(trimmed.as_mask().is_some()); + assert_eq!(trimmed.as_mask().unwrap().len(), 0); + } + + #[test] + fn test_boolean_mask_from_selectors_fuzz_equivalence() { + let mut rand = rng(); + for _ in 0..200 { + let n_selectors = rand.random_range(0..30); + let mut selectors = Vec::with_capacity(n_selectors); + for _ in 0..n_selectors { + selectors.push(RowSelector { + row_count: rand.random_range(0..40), + skip: rand.random_bool(0.5), + }); + } + + let expected = { + let total_rows: usize = selectors.iter().map(|s| s.row_count).sum(); + let mut builder = BooleanBufferBuilder::new(total_rows); + for selector in &selectors { + builder.append_n(selector.row_count, !selector.skip); + } + builder.finish() + }; + + assert_eq!(boolean_mask_from_selectors(&selectors), expected); + } + } + + #[test] + fn test_mask_has_at_least_runs() { + fn assert_run_count(bits: Vec, expected_runs: usize) { + let mask = BooleanBuffer::from(bits); + for min_runs in 0..=expected_runs + 2 { + assert_eq!( + mask_has_at_least_runs(&mask, min_runs), + expected_runs >= min_runs, + "expected {expected_runs} runs with boundary {min_runs}" + ); + } + } + + assert_run_count(vec![], 0); + assert_run_count(vec![false; 8], 1); + assert_run_count(vec![true; 8], 1); + assert_run_count(vec![false, false, true, true, false], 3); + assert_run_count(vec![true, false, true, false, true, false], 6); + + // Exercise the unaligned iterator path as mask-backed selections can be slices. + let mask = BooleanBuffer::from(vec![true, false, false, true, true, false, true, true]) + .slice(1, 6); + for min_runs in 0..=6 { + assert_eq!(mask_has_at_least_runs(&mask, min_runs), 4 >= min_runs); + } + } + + #[test] + fn test_trim_mask_fuzz_equivalence() { + let mut rand = rng(); + for _ in 0..200 { + let len = rand.random_range(0..200); + let bits: Vec = (0..len).map(|_| rand.random_bool(0.3)).collect(); + let full = BooleanBuffer::from(bits.clone()); + // Exercise non-zero bit offsets via slicing + let start = rand.random_range(0..=len); + let slice_len = rand.random_range(0..=(len - start)); + let mask = full.slice(start, slice_len); + + let expected_len = bits[start..start + slice_len] + .iter() + .rposition(|&b| b) + .map_or(0, |pos| pos + 1); + + match trim_mask(&mask) { + Some(trimmed) => { + assert_ne!(expected_len, mask.len()); + assert_eq!(trimmed.len(), expected_len); + assert_eq!(trimmed, mask.slice(0, expected_len)); + } + None => assert_eq!(expected_len, mask.len()), + } + } + } + + #[test] + fn test_split_off_propagates_cached_count() { + let bits = vec![true, false, true, true, false, false, true, false]; + let mut selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits)); + // Populate the count cache, then verify both split halves. + assert_eq!(selection.row_count(), 4); + let head = selection.split_off(3); + assert_eq!(head.row_count(), 2); + assert_eq!(selection.row_count(), 2); + let tail_fresh = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ + true, false, false, true, false, + ])); + assert_eq!(selection, tail_fresh); + + // Splitting past the end keeps the whole selection as the head + let head = selection.split_off(100); + assert_eq!(head.row_count(), 2); + assert_eq!(selection.row_count(), 0); + } + + #[test] + fn test_trim_and_offset_and_limit_preserve_cached_count() { + let bits = vec![true, true, false, true, false, false]; + let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone())); + assert_eq!(selection.row_count(), 3); + + let trimmed = selection.trim(); + assert!(trimmed.as_mask().is_some()); + assert_eq!(trimmed.as_mask().unwrap().len(), 4); + assert_eq!(trimmed.row_count(), 3); + + let offset = trimmed.clone().offset(1); + assert_eq!(offset.row_count(), 2); + + let limited = trimmed.limit(2); + assert_eq!(limited.row_count(), 2); + } +} diff --git a/vendor/parquet/src/arrow/arrow_reader/selection/cursor.rs b/vendor/parquet/src/arrow/arrow_reader/selection/cursor.rs new file mode 100644 index 0000000..2b63c09 --- /dev/null +++ b/vendor/parquet/src/arrow/arrow_reader/selection/cursor.rs @@ -0,0 +1,428 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Execution time iteration over a [`RowSelection`]. +//! +//! A [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan) resolves a +//! [`RowSelectionPolicy`] into a [`RowSelectionStrategy`] and builds the +//! 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; +use arrow_buffer::BooleanBuffer; +use std::collections::VecDeque; +use std::ops::Range; +use std::sync::Arc; + +/// Policy for picking a strategy to materialize [`RowSelection`] during execution. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum RowSelectionPolicy { + /// Use a queue of [`RowSelector`] values + Selectors, + /// Use a boolean mask to materialize the selection + Mask, + /// Choose between [`Self::Mask`] and [`Self::Selectors`] based on selector density + Auto { + /// Average selector length below which masks are preferred + threshold: usize, + }, +} + +impl Default for RowSelectionPolicy { + fn default() -> Self { + Self::Auto { threshold: 32 } + } +} + +/// Fully resolved strategy for materializing [`RowSelection`] during execution. +/// +/// This is determined by [`RowSelectionPolicy`], including selector density for +/// [`RowSelectionPolicy::Auto`]. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum RowSelectionStrategy { + /// Use a queue of [`RowSelector`] values + Selectors, + /// Use a boolean mask to materialize the selection + Mask, +} + +/// Cursor for iterating a [`RowSelection`] during execution within a +/// [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan). +/// +/// This keeps per-reader state such as the current position and delegates the +/// actual storage strategy to the internal `RowSelectionInner`. +#[derive(Debug)] +pub enum RowSelectionCursor { + /// Reading all rows + All, + /// Use a bitmask to back the selection (dense selections) + Mask(MaskCursor), + /// Use a queue of selectors to back the selection (sparse selections) + Selectors(SelectorsCursor), +} + +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, + loaded_row_ranges: Option>, + ) -> Self { + debug_assert!( + mask.is_empty() || mask.value(mask.len() - 1), + "Mask selections must not end with a skip" + ); + Self::Mask(MaskCursor { + mask, + position: 0, + loaded_row_ranges, + }) + } + + /// Create a [`RowSelectionCursor::Selectors`] from the provided selectors + pub(crate) fn new_selectors(selectors: Vec) -> Self { + Self::Selectors(SelectorsCursor { + selectors: selectors.into(), + position: 0, + }) + } + + /// Create a cursor that selects all rows + pub(crate) fn new_all() -> Self { + Self::All + } +} + +/// Cursor for iterating a selector-backed [`RowSelection`] +/// +/// This is best for sparse selections where large contiguous +/// blocks of rows are selected or skipped. +#[derive(Debug)] +pub struct SelectorsCursor { + selectors: VecDeque, + /// Current absolute offset into the selection + position: usize, +} + +impl SelectorsCursor { + /// Returns `true` when no further rows remain + pub fn is_empty(&self) -> bool { + self.selectors.is_empty() + } + + pub(crate) fn selectors_mut(&mut self) -> &mut VecDeque { + &mut self.selectors + } + + /// Return the next [`RowSelector`] + pub(crate) fn next_selector(&mut self) -> RowSelector { + let selector = self.selectors.pop_front().unwrap(); + self.position += selector.row_count; + selector + } + + /// Return a selector to the front, rewinding the position + pub(crate) fn return_selector(&mut self, selector: RowSelector) { + self.position = self.position.saturating_sub(selector.row_count); + self.selectors.push_front(selector); + } +} + +/// Cursor for iterating a mask-backed [`RowSelection`] +/// +/// This is best for dense selections where there are many small skips +/// or selections. For example, selecting every other row. +/// +/// When page pruning produces sparse column data, `loaded_row_ranges` limits +/// each decoded chunk to rows whose pages are loaded for every projected leaf. +/// For example, two projected columns can have different page boundaries: +/// +/// ```text +/// Row ranges: [0, 4) [4, 6) [6, 8) [8, 10) [10, 12) +/// Selection mask: 1000 00 00 00 01 +/// Column A pages: loaded | missing [4, 8) | loaded [8, 12) +/// Column B pages: loaded [0, 6) | missing [6, 10) | loaded +/// LoadedRowRanges: [0, 4) [10, 12) +/// ``` +/// +/// The first chunk decodes `[0, 1)` with mask `1`. The next chunk skips to row +/// 11 and decodes `[11, 12)` with mask `1`. When loaded ranges are present, +/// every returned chunk ends at a selected row and never includes trailing +/// unselected rows. [`ParquetRecordBatchReader`] still accumulates both chunks +/// and applies the combined mask `11` once. +/// +/// [`ParquetRecordBatchReader`]: crate::arrow::arrow_reader::ParquetRecordBatchReader +#[derive(Debug)] +pub struct MaskCursor { + mask: BooleanBuffer, + /// Current absolute offset into the selection + position: usize, + /// Row ranges whose backing pages are loaded for every projected column. + loaded_row_ranges: Option>, +} + +impl MaskCursor { + /// Returns `true` when no further rows remain + pub fn is_empty(&self) -> bool { + self.position >= self.mask.len() + } + + /// Advance through the mask representation, producing the next chunk summary + pub fn next_mask_chunk(&mut self, batch_size: usize) -> Option { + if self.is_empty() { + return None; + } + + Some(self.next_mask_chunk_non_empty(batch_size)) + } + + /// Produces the next chunk for a non-empty, trailing-skip-free mask. + fn next_mask_chunk_non_empty(&mut self, batch_size: usize) -> MaskChunk { + debug_assert!(!self.is_empty()); + + let (initial_skip, chunk_rows, selected_rows, mask_start, end_position) = { + let mask = &self.mask; + let start_position = self.position; + let mut cursor = start_position; + let mut initial_skip = 0; + + while cursor < mask.len() && !mask.value(cursor) { + initial_skip += 1; + cursor += 1; + } + debug_assert!( + cursor < mask.len(), + "ReadPlan must remove trailing skips from Mask selections" + ); + + let mask_start = cursor; + let mut chunk_rows = 0; + let mut selected_rows = 0; + + // Advance until enough rows have been selected to satisfy the batch size, + // or until the mask is exhausted. This mirrors the behaviour of the legacy + // `RowSelector` queue-based iteration. + while cursor < mask.len() && selected_rows < batch_size { + chunk_rows += 1; + if mask.value(cursor) { + selected_rows += 1; + } + cursor += 1; + } + + (initial_skip, chunk_rows, selected_rows, mask_start, cursor) + }; + + self.position = end_position; + + MaskChunk { + initial_skip, + chunk_rows, + selected_rows, + mask_start, + } + } + + /// Returns the next non-empty mask chunk without crossing an unloaded row range. + /// When loaded ranges are present, every returned chunk ends immediately after a + /// selected row and therefore never contains trailing unselected rows. Those rows + /// remain for the next call's initial skip. + /// + /// The [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan) removes trailing + /// skips before constructing this cursor. Callers therefore only invoke + /// this method for a non-empty mask that has another selected row. + pub(crate) fn next_chunk(&mut self, batch_size: usize) -> Result { + debug_assert!(batch_size > 0); + debug_assert!(!self.is_empty()); + + if self.loaded_row_ranges.is_none() { + return Ok(self.next_mask_chunk_non_empty(batch_size)); + } + + let start_position = self.position; + let mut cursor = start_position; + while cursor < self.mask.len() && !self.mask.value(cursor) { + cursor += 1; + } + + if cursor == self.mask.len() { + return Err(ParquetError::General( + "Internal Error: Mask cursor reached the end without finding a selected row; \ + ReadPlan must remove trailing skips" + .to_string(), + )); + } + + let loaded_range_end = self + .loaded_row_ranges + .as_ref() + .and_then(|ranges| ranges.end_containing(cursor)) + .ok_or_else(|| { + ParquetError::General(format!( + "Internal Error: selected row {cursor} has no loaded page range" + )) + })?; + + let mask_start = cursor; + let mut selected_rows = 0; + let mut chunk_end = cursor; + while cursor < loaded_range_end && cursor < self.mask.len() && selected_rows < batch_size { + if self.mask.value(cursor) { + selected_rows += 1; + chunk_end = cursor + 1; + } + cursor += 1; + } + + self.position = chunk_end; + Ok(MaskChunk { + initial_skip: mask_start - start_position, + chunk_rows: chunk_end - mask_start, + selected_rows, + mask_start, + }) + } + + /// Materialise the boolean values for a mask-backed chunk + pub fn mask_values_for(&self, chunk: &MaskChunk) -> Result { + if chunk.mask_start.saturating_add(chunk.chunk_rows) > self.mask.len() { + return Err(ParquetError::General( + "Internal Error: MaskChunk exceeds mask length".to_string(), + )); + } + Ok(BooleanArray::from( + self.mask.slice(chunk.mask_start, chunk.chunk_rows), + )) + } +} + +/// Result of computing the next chunk to read when using a [`MaskCursor`] +#[derive(Debug)] +pub struct MaskChunk { + /// Number of leading rows to skip before reaching selected rows + pub initial_skip: usize, + /// Total rows covered by this chunk (selected + skipped) + pub chunk_rows: usize, + /// Rows actually selected within the chunk + pub selected_rows: usize, + /// Starting offset within the mask where the chunk begins + pub mask_start: usize, +} + +/// Row ranges whose backing pages are loaded for every projected column. +#[derive(Clone, Debug)] +pub(crate) struct LoadedRowRanges(Vec>); + +impl LoadedRowRanges { + pub(crate) fn from_selection(selection: RowSelection) -> Self { + let selectors: Vec = selection.into(); + let mut position = 0; + let ranges = selectors + .into_iter() + .filter_map(|selector| { + let start = position; + position += selector.row_count; + (!selector.skip).then_some(start..position) + }) + .collect(); + Self(ranges) + } + + fn end_containing(&self, row: usize) -> Option { + let idx = self.0.partition_point(|range| range.end <= row); + self.0 + .get(idx) + .filter(|range| range.start <= row) + .map(|range| range.end) + } + + #[cfg(test)] + pub(crate) fn ranges(&self) -> &[Range] { + &self.0 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[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)], + Some(loaded.into()), + ) else { + unreachable!() + }; + + let chunk = cursor.next_chunk(10).unwrap(); + assert_eq!(chunk.chunk_rows, 1); + assert!(cursor.is_empty()); + } + + #[test] + fn test_next_mask_chunk_until_cursor_is_empty() { + let RowSelectionCursor::Mask(mut cursor) = RowSelectionCursor::new_mask_from_selectors( + vec![ + RowSelector::skip(2), + RowSelector::select(2), + RowSelector::skip(1), + RowSelector::select(1), + ], + None, + ) else { + unreachable!() + }; + + let first = cursor.next_mask_chunk(2).unwrap(); + assert_eq!(first.initial_skip, 2); + assert_eq!(first.chunk_rows, 2); + assert_eq!(first.selected_rows, 2); + + let second = cursor.next_mask_chunk(2).unwrap(); + assert_eq!(second.initial_skip, 1); + assert_eq!(second.chunk_rows, 1); + assert_eq!(second.selected_rows, 1); + + assert!(cursor.next_mask_chunk(2).is_none()); + } +} diff --git a/vendor/parquet/src/arrow/arrow_reader/selection/mod.rs b/vendor/parquet/src/arrow/arrow_reader/selection/mod.rs new file mode 100644 index 0000000..ff8d5e9 --- /dev/null +++ b/vendor/parquet/src/arrow/arrow_reader/selection/mod.rs @@ -0,0 +1,954 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Logic for selecting which rows to read: [`RowSelection`] and [`RowSelector`] +//! +//! This module holds [`RowSelection`] and its public API, which dispatches to +//! one of the two backings depending on how the selection is stored: +//! +//! * `selector`: the run length backing, [`RowSelector`] and its primitives +//! * `boolean`: the bitmap backing, `MaskSelection` and its primitives +//! +//! The remaining modules hold the operations that are common to both: +//! +//! * `algebra`: `and_then`, `intersection` and `union` +//! * `ranges`: mapping a [`RowSelection`] onto page and batch ranges +//! * `cursor`: iterating a [`RowSelection`] while reading + +use crate::file::page_index::offset_index::PageLocation; +use arrow_array::{Array, BooleanArray}; +use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder}; +use arrow_select::filter::SlicesIterator; +use std::cmp::Ordering; +use std::collections::VecDeque; +use std::ops::Range; + +mod algebra; +mod boolean; +mod cursor; +mod ranges; +mod selector; + +use algebra::{ + and_then_mask, and_then_row_selections, and_then_selectors_with_mask, intersect_masks, + intersect_row_selections, union_masks, union_row_selections, +}; +pub use boolean::MaskRunIter; +pub(crate) use boolean::mask_to_selectors; +use boolean::{ + MaskSelection, 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}; +use ranges::{expand_to_batch_boundaries_from_selectors, scan_ranges_from_selectors}; +pub use selector::RowSelector; +use selector::{limit_selectors, offset_selectors, split_off_selectors}; + +/// [`RowSelection`] represents selecting a subset of rows +/// when scanning a parquet file. +/// +/// This is applied prior to reading column data, and can therefore +/// be used to skip IO to fetch data into memory +/// +/// A typical use-case would be using the [`PageIndex`] to filter out rows +/// that don't satisfy a predicate +/// +/// Depending on the pattern of rows to be selected, [`RowSelection`] has +/// either a bitmap or an RLE ([`RowSelector`]) based implementation. +/// +/// # Example +/// ``` +/// use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; +/// +/// let selectors = vec![ +/// RowSelector::skip(5), +/// RowSelector::select(5), +/// RowSelector::select(5), +/// RowSelector::skip(5), +/// ]; +/// +/// // Creating a selection will combine adjacent selectors +/// let selection: RowSelection = selectors.into(); +/// +/// let expected = vec![ +/// RowSelector::skip(5), +/// RowSelector::select(10), +/// RowSelector::skip(5), +/// ]; +/// +/// let actual: Vec = selection.into(); +/// assert_eq!(actual, expected); +/// +/// // you can also create a selection from consecutive ranges +/// let ranges = vec![5..10, 10..15]; +/// let selection = +/// RowSelection::from_consecutive_ranges(ranges.into_iter(), 20); +/// let actual: Vec = selection.into(); +/// assert_eq!(actual, expected); +/// +/// // or directly from a packed bitmap, when the upstream producer already +/// // has one. The bitmap is kept as-is rather than run-length-encoded. +/// use arrow_buffer::BooleanBuffer; +/// let mask = BooleanBuffer::from(vec![true, false, true, true]); +/// let selection = RowSelection::from_boolean_buffer(mask); +/// assert_eq!(selection.row_count(), 3); +/// ``` +/// +/// An RLE ([`RowSelector`]) backed [`RowSelection`] maintains the following +/// invariants (they do not apply to the bitmap backed implementation): +/// +/// * It contains no [`RowSelector`] of 0 rows +/// * Consecutive [`RowSelector`]s alternate skipping or selecting rows +/// +/// [`PageIndex`]: crate::file::page_index::column_index::ColumnIndexMetaData +#[derive(Default, Clone)] +pub struct RowSelection { + inner: RowSelectionInner, +} + +/// Internal storage for [`RowSelection`]. +#[derive(Debug, Clone)] +pub(crate) enum RowSelectionInner { + Selectors(Vec), + Mask(Box), +} + +impl Default for RowSelectionInner { + fn default() -> Self { + Self::Selectors(Vec::new()) + } +} + +impl std::fmt::Debug for RowSelection { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match &self.inner { + RowSelectionInner::Selectors(s) => f + .debug_struct("RowSelection") + .field("selectors", s) + .finish(), + RowSelectionInner::Mask(m) => f + .debug_struct("RowSelection") + .field("mask_len", &m.mask().len()) + .finish_non_exhaustive(), + } + } +} + +impl PartialEq for RowSelection { + fn eq(&self, other: &Self) -> bool { + match (&self.inner, &other.inner) { + (RowSelectionInner::Selectors(a), RowSelectionInner::Selectors(b)) => a == b, + (RowSelectionInner::Mask(a), RowSelectionInner::Mask(b)) => a.mask() == b.mask(), + (RowSelectionInner::Mask(mask), RowSelectionInner::Selectors(selectors)) + | (RowSelectionInner::Selectors(selectors), RowSelectionInner::Mask(mask)) => { + if selectors + .iter() + .try_fold(0usize, |acc, selector| acc.checked_add(selector.row_count)) + != Some(mask.mask().len()) + { + return false; + } + + let mut slices = mask.mask().set_slices().peekable(); + let mut cursor = 0usize; + + for selector in selectors { + let end = cursor + selector.row_count; + + if selector.skip { + if slices.peek().is_some_and(|(start, _)| *start < end) { + return false; + } + } else { + match slices.next() { + Some((start, slice_end)) if start == cursor && slice_end == end => {} + _ => return false, + } + } + + cursor = end; + } + + slices.next().is_none() + } + } + } +} + +impl Eq for RowSelection {} + +impl RowSelection { + /// Not `pub`: unlike `From>`, this performs no + /// validation/normalization of the selectors (e.g. combining adjacent + /// selectors), so callers must uphold the invariants themselves. + fn from_selectors(selectors: Vec) -> Self { + Self { + inner: RowSelectionInner::Selectors(selectors), + } + } + + /// Create a [`RowSelection`] from a packed [`BooleanBuffer`]. + /// + /// Each set bit selects a row, each unset bit skips one. Unlike + /// [`Self::from_filters`], the bitmap is kept as-is rather than + /// eagerly run-length-encoded. [`Self::iter`] materializes and caches the + /// RLE form on first use; use [`MaskRunIter`] to stream the RLE form + /// directly from the bitmap. + pub fn from_boolean_buffer(mask: BooleanBuffer) -> Self { + Self { + inner: RowSelectionInner::Mask(Box::new(MaskSelection::new(mask))), + } + } + + fn from_mask_selection(mask: MaskSelection) -> Self { + Self { + inner: RowSelectionInner::Mask(Box::new(mask)), + } + } + + /// Returns the underlying mask if this selection is mask-backed. + /// + /// Public so that engines composing selections (e.g. DataFusion's + /// [`ParquetAccessPlan::into_overall_row_selection`]) can concatenate + /// mask-backed selections without materialising the RLE form. + /// + /// [`ParquetAccessPlan::into_overall_row_selection`]: https://docs.rs/datafusion-datasource-parquet/latest/datafusion_datasource_parquet/access_plan/struct.ParquetAccessPlan.html#method.into_overall_row_selection + pub fn as_mask(&self) -> Option<&BooleanBuffer> { + match &self.inner { + RowSelectionInner::Mask(m) => Some(m.mask()), + _ => None, + } + } + + /// Consume the selection and return its internal storage. + pub(crate) fn into_inner(self) -> RowSelectionInner { + self.inner + } + + /// Choose the automatic materialisation strategy without converting between + /// selector and mask backing. + #[inline] + pub(crate) fn auto_selection_strategy(&self, threshold: usize) -> RowSelectionStrategy { + let (total_rows, effective_count) = match &self.inner { + RowSelectionInner::Selectors(selectors) => { + selectors.iter().fold((0usize, 0usize), |(rows, count), s| { + if s.row_count > 0 { + (rows + s.row_count, count + 1) + } else { + (rows, count) + } + }) + } + RowSelectionInner::Mask(mask) => { + let mask = mask.mask(); + let total_rows = mask.len(); + + if total_rows == 0 { + return RowSelectionStrategy::Mask; + } + + // A mask is preferred when: + // + // total_rows < run_count * threshold + // + // Therefore only scan until the first run count that can make + // the inequality true. Fragmented masks normally reach this + // boundary near the start instead of enumerating every run. + let min_mask_runs = total_rows + .checked_div(threshold) + .and_then(|max_selector_runs| max_selector_runs.checked_add(1)); + + return match min_mask_runs { + Some(min_runs) if mask_has_at_least_runs(mask, min_runs) => { + RowSelectionStrategy::Mask + } + _ => RowSelectionStrategy::Selectors, + }; + } + }; + + if effective_count == 0 { + return RowSelectionStrategy::Mask; + } + + if total_rows < effective_count.saturating_mul(threshold) { + RowSelectionStrategy::Mask + } else { + RowSelectionStrategy::Selectors + } + } + + #[cfg(test)] + fn selectors(&self) -> Vec { + self.iter().copied().collect() + } + + fn into_selectors_vec(self) -> Vec { + match self.inner { + RowSelectionInner::Selectors(s) => s, + RowSelectionInner::Mask(m) => mask_to_selectors(m.mask()), + } + } + + /// Creates a [`RowSelection`] from a slice of [`BooleanArray`] + /// + /// # Panic + /// + /// Panics if any of the [`BooleanArray`] contain nulls + pub fn from_filters(filters: &[BooleanArray]) -> Self { + let mut next_offset = 0; + let total_rows = filters.iter().map(|x| x.len()).sum(); + + let iter = filters.iter().flat_map(|filter| { + let offset = next_offset; + next_offset += filter.len(); + assert_eq!(filter.null_count(), 0); + SlicesIterator::new(filter).map(move |(start, end)| start + offset..end + offset) + }); + + Self::from_consecutive_ranges(iter, total_rows) + } + + /// Creates a [`RowSelection`] from an iterator of consecutive ranges to keep + pub fn from_consecutive_ranges>>( + ranges: I, + total_rows: usize, + ) -> Self { + let mut selectors: Vec = Vec::with_capacity(ranges.size_hint().0); + let mut last_end = 0; + for range in ranges { + let len = range.end - range.start; + if len == 0 { + continue; + } + + match range.start.cmp(&last_end) { + Ordering::Equal => match selectors.last_mut() { + Some(last) => last.row_count = last.row_count.checked_add(len).unwrap(), + None => selectors.push(RowSelector::select(len)), + }, + Ordering::Greater => { + selectors.push(RowSelector::skip(range.start - last_end)); + selectors.push(RowSelector::select(len)) + } + Ordering::Less => panic!("out of order"), + } + last_end = range.end; + } + + if last_end != total_rows { + selectors.push(RowSelector::skip(total_rows - last_end)) + } + + Self::from_selectors(selectors) + } + + /// Given an offset index, return the byte ranges for all data pages selected by `self` + /// + /// This is useful for determining what byte ranges to fetch from underlying storage + /// + /// Note: this method does not make any effort to combine consecutive ranges, nor coalesce + /// ranges that are close together. This is instead delegated to the IO subsystem to optimise, + /// e.g. `ObjectStore::get_ranges` in the [`object_store`] crate + /// + /// [`object_store`]: https://crates.io/crates/object_store + pub fn scan_ranges(&self, page_locations: &[PageLocation]) -> Vec> { + match &self.inner { + RowSelectionInner::Selectors(selectors) => { + scan_ranges_from_selectors(selectors.iter().copied(), page_locations) + } + RowSelectionInner::Mask(mask) => { + scan_ranges_from_selectors(MaskRunIter::new(mask.mask()), page_locations) + } + } + } + + /// Returns the complete row ranges of the pages selected by [`Self::scan_ranges`]. + pub(crate) fn row_ranges_for_selected_pages( + &self, + page_locations: &[PageLocation], + total_rows: usize, + ) -> Vec> { + let mut selected_pages = self.scan_ranges(page_locations).into_iter().peekable(); + let mut row_ranges = Vec::new(); + + for (idx, page) in page_locations.iter().enumerate() { + let Some(selected_page) = selected_pages.peek() else { + break; + }; + if selected_page.start != page.offset as u64 { + continue; + } + selected_pages.next(); + + let end = page_locations + .get(idx + 1) + .map(|next| next.first_row_index as usize) + .unwrap_or(total_rows); + row_ranges.push(page.first_row_index as usize..end); + } + + row_ranges + } + + /// Splits off the first `row_count` from this [`RowSelection`] + pub fn split_off(&mut self, row_count: usize) -> Self { + match std::mem::take(&mut self.inner) { + RowSelectionInner::Mask(mask) => { + let total = mask.cached_count(); + let (head, tail) = split_off_mask((*mask).into_mask(), row_count); + // Popcount only the head and derive the tail by subtraction, so + // repeated splits stay O(bitmap) overall. + let (head, tail) = match total { + Some(total) => { + let head_count = if tail.is_empty() { + total + } else { + head.count_set_bits() + }; + ( + MaskSelection::with_count(head, head_count), + MaskSelection::with_count(tail, total - head_count), + ) + } + None => (MaskSelection::new(head), MaskSelection::new(tail)), + }; + self.inner = RowSelectionInner::Mask(Box::new(tail)); + Self::from_mask_selection(head) + } + RowSelectionInner::Selectors(selectors) => { + let (head, tail) = split_off_selectors(selectors, row_count); + self.inner = RowSelectionInner::Selectors(tail); + Self::from_selectors(head) + } + } + } + + /// returns a [`RowSelection`] representing rows that are selected in both + /// input [`RowSelection`]s. + /// + /// This is equivalent to the logical `AND` / conjunction of the two + /// selections. + /// + /// # Example + /// If `N` means the row is not selected, and `Y` means it is + /// selected: + /// + /// ```text + /// self: NNNNNNNNNNNNYYYYYYYYYYYYYYYYYYYYYYNNNYYYYY + /// other: YYYYYNNNNYYYYYYYYYYYYY YYNNN + /// + /// returned: NNNNNNNNNNNNYYYYYNNNNYYYYYYYYYYYYYNNNYYNNN + /// ``` + /// + /// # Panics + /// + /// Panics if `other` does not have a length equal to the number of rows selected + /// by this RowSelection + /// + pub fn and_then(&self, other: &Self) -> Self { + match (&self.inner, &other.inner) { + (RowSelectionInner::Mask(mask), _) => { + Self::from_boolean_buffer(and_then_mask(mask.mask(), other)) + } + (RowSelectionInner::Selectors(first), RowSelectionInner::Selectors(second)) => { + and_then_row_selections(first, second) + } + (RowSelectionInner::Selectors(first), RowSelectionInner::Mask(second)) => { + and_then_selectors_with_mask(first, second.mask()) + } + } + } + + /// Compute the intersection of two [`RowSelection`] + /// For example: + /// self: NNYYYYNNYYNYN + /// other: NYNNNNNNY + /// + /// returned: NNNNNNNNYYNYN + pub fn intersection(&self, other: &Self) -> Self { + match (&self.inner, &other.inner) { + (RowSelectionInner::Mask(l), RowSelectionInner::Mask(r)) => { + Self::from_boolean_buffer(intersect_masks(l.mask(), r.mask())) + } + (RowSelectionInner::Selectors(l), RowSelectionInner::Selectors(r)) => { + intersect_row_selections(l, r) + } + (RowSelectionInner::Selectors(l), RowSelectionInner::Mask(r)) => { + let r = mask_to_selectors(r.mask()); + intersect_row_selections(l, &r) + } + (RowSelectionInner::Mask(l), RowSelectionInner::Selectors(r)) => { + let l = mask_to_selectors(l.mask()); + intersect_row_selections(&l, r) + } + } + } + + /// Compute the union of two [`RowSelection`] + /// For example: + /// self: NNYYYYNNYYNYN + /// other: NYNNNNNNN + /// + /// returned: NYYYYYNNYYNYN + pub fn union(&self, other: &Self) -> Self { + match &self.inner { + RowSelectionInner::Mask(l) => match &other.inner { + RowSelectionInner::Mask(r) => { + Self::from_boolean_buffer(union_masks(l.mask(), r.mask())) + } + RowSelectionInner::Selectors(r) => { + let l = mask_to_selectors(l.mask()); + union_row_selections(&l, r) + } + }, + RowSelectionInner::Selectors(l) => match &other.inner { + RowSelectionInner::Mask(r) => { + let r = mask_to_selectors(r.mask()); + union_row_selections(l, &r) + } + RowSelectionInner::Selectors(r) => union_row_selections(l, r), + }, + } + } + + /// Returns `true` if this [`RowSelection`] selects any rows + pub fn selects_any(&self) -> bool { + match &self.inner { + RowSelectionInner::Selectors(s) => s.iter().any(|x| !x.skip), + RowSelectionInner::Mask(m) => match m.cached_count() { + Some(count) => count > 0, + None => m.mask().set_indices().next().is_some(), + }, + } + } + + /// Trims this [`RowSelection`] removing any trailing skips + pub(crate) fn trim(self) -> Self { + match self.inner { + RowSelectionInner::Mask(m) => { + let trimmed = trim_mask(m.mask()); + let cached_count = m.cached_count(); + match trimmed { + // Trimming only drops trailing unset bits; the count is unchanged. + Some(mask) => match cached_count { + Some(count) => { + Self::from_mask_selection(MaskSelection::with_count(mask, count)) + } + None => Self::from_boolean_buffer(mask), + }, + // Nothing to trim, hand the existing box back untouched. + None => Self { + inner: RowSelectionInner::Mask(m), + }, + } + } + RowSelectionInner::Selectors(mut selectors) => { + while selectors.last().map(|x| x.skip).unwrap_or(false) { + selectors.pop(); + } + Self::from_selectors(selectors) + } + } + } + + /// Applies an offset to this [`RowSelection`], skipping the first `offset` selected rows + pub(crate) fn offset(self, offset: usize) -> Self { + if offset == 0 { + return self; + } + + match self.inner { + RowSelectionInner::Mask(mask) => { + let count = mask.count(); + let buffer = offset_mask((*mask).into_mask(), offset, count); + Self::from_mask_selection(MaskSelection::with_count( + buffer, + count.saturating_sub(offset), + )) + } + RowSelectionInner::Selectors(selectors) => { + Self::from_selectors(offset_selectors(selectors, offset)) + } + } + } + + /// Limit this [`RowSelection`] to only select `limit` rows + pub(crate) fn limit(self, limit: usize) -> Self { + match self.inner { + RowSelectionInner::Mask(mask) => { + let cached = mask.cached_count(); + let buffer = limit_mask((*mask).into_mask(), limit); + match cached { + Some(count) => Self::from_mask_selection(MaskSelection::with_count( + buffer, + count.min(limit), + )), + None => Self::from_boolean_buffer(buffer), + } + } + RowSelectionInner::Selectors(selectors) => { + Self::from_selectors(limit_selectors(selectors, limit)) + } + } + } + + /// Returns an iterator over the [`RowSelector`]s for this + /// [`RowSelection`]. + /// + /// Mask-backed selections materialize a `Vec` cache on first + /// call (one allocation, `O(set_slices)` work) so the iterator can hand out + /// `&RowSelector`; the cache is not copied on clone. For single-pass walks + /// over mask-backed selections, prefer streaming directly via + /// [`Self::as_mask`] + [`MaskRunIter::new`] — that path is allocation-free + /// and avoids populating the cache. + pub fn iter(&self) -> impl Iterator { + match &self.inner { + RowSelectionInner::Selectors(s) => s.iter(), + RowSelectionInner::Mask(m) => m.selectors().iter(), + } + } + + /// Returns the number of selected rows + pub fn row_count(&self) -> usize { + match &self.inner { + RowSelectionInner::Selectors(s) => { + s.iter().filter(|x| !x.skip).map(|x| x.row_count).sum() + } + RowSelectionInner::Mask(m) => m.count(), + } + } + + /// Returns the number of de-selected rows + pub fn skipped_row_count(&self) -> usize { + match &self.inner { + RowSelectionInner::Selectors(s) => { + s.iter().filter(|x| x.skip).map(|x| x.row_count).sum() + } + RowSelectionInner::Mask(m) => m.mask().len() - m.count(), + } + } + + /// Expands the selection to align with batch boundaries. + /// This is needed when using cached array readers to ensure that + /// the cached data covers full batches. + pub(crate) fn expand_to_batch_boundaries(&self, batch_size: usize, total_rows: usize) -> Self { + if batch_size == 0 { + return self.clone(); + } + + match &self.inner { + RowSelectionInner::Selectors(selectors) => expand_to_batch_boundaries_from_selectors( + selectors.iter().copied(), + batch_size, + total_rows, + ), + RowSelectionInner::Mask(mask) => expand_to_batch_boundaries_from_selectors( + MaskRunIter::new(mask.mask()), + batch_size, + total_rows, + ), + } + } +} + +impl From> for RowSelection { + fn from(selectors: Vec) -> Self { + selectors.into_iter().collect() + } +} + +impl From for RowSelection { + fn from(mask: BooleanBuffer) -> Self { + Self::from_boolean_buffer(mask) + } +} + +impl FromIterator for RowSelection { + fn from_iter>(iter: T) -> Self { + let iter = iter.into_iter(); + + // Capacity before filter + let mut selectors = Vec::with_capacity(iter.size_hint().0); + + let mut filtered = iter.filter(|x| x.row_count != 0); + if let Some(x) = filtered.next() { + selectors.push(x); + } + + for s in filtered { + if s.row_count == 0 { + continue; + } + + // Combine consecutive selectors + let last = selectors.last_mut().unwrap(); + if last.skip == s.skip { + last.row_count = last.row_count.checked_add(s.row_count).unwrap(); + } else { + selectors.push(s) + } + } + + Self::from_selectors(selectors) + } +} + +impl From for Vec { + fn from(r: RowSelection) -> Self { + r.into_selectors_vec() + } +} + +impl From for VecDeque { + fn from(r: RowSelection) -> Self { + r.into_selectors_vec().into() + } +} + +impl FromIterator for RowSelection { + /// Concatenate multiple [`RowSelection`]s in iterator order. + /// + /// When every input is mask-backed the result stays mask-backed + /// (`BooleanBuffer`s are appended); otherwise falls back to flattening + /// through the per-`RowSelector` form. + fn from_iter>(iter: T) -> Self { + let items: Vec = iter.into_iter().collect(); + + let all_mask = items + .iter() + .all(|s| matches!(&s.inner, RowSelectionInner::Mask(_))); + + if all_mask { + let total_len: usize = items + .iter() + .map(|s| match &s.inner { + RowSelectionInner::Mask(m) => m.mask().len(), + RowSelectionInner::Selectors(_) => unreachable!(), + }) + .sum(); + let mut builder = BooleanBufferBuilder::new(total_len); + for item in items { + match item.into_inner() { + RowSelectionInner::Mask(m) => builder.append_buffer(m.mask()), + RowSelectionInner::Selectors(_) => unreachable!(), + } + } + return Self::from_boolean_buffer(builder.finish()); + } + + items + .into_iter() + .flat_map(|s| s.into_selectors_vec()) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_offset_zero_and_zero_batch_expand_are_identity() { + let selection = + RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true])); + assert_eq!(selection.clone().offset(0), selection); + assert_eq!(selection.expand_to_batch_boundaries(0, 3), selection); + } + + #[test] + fn test_from_filters() { + let filters = vec![ + BooleanArray::from(vec![false, false, false, true, true, true, true]), + BooleanArray::from(vec![true, true, false, false, true, true, true]), + BooleanArray::from(vec![false, false, false, false]), + BooleanArray::from(Vec::::new()), + ]; + + let selection = RowSelection::from_filters(&filters[..1]); + assert!(selection.selects_any()); + assert_eq!( + selection.selectors(), + vec![RowSelector::skip(3), RowSelector::select(4)] + ); + + let selection = RowSelection::from_filters(&filters[..2]); + assert!(selection.selects_any()); + assert_eq!( + selection.selectors(), + vec![ + RowSelector::skip(3), + RowSelector::select(6), + RowSelector::skip(2), + RowSelector::select(3) + ] + ); + + let selection = RowSelection::from_filters(&filters); + assert!(selection.selects_any()); + assert_eq!( + selection.selectors(), + vec![ + RowSelector::skip(3), + RowSelector::select(6), + RowSelector::skip(2), + RowSelector::select(3), + RowSelector::skip(4) + ] + ); + + let selection = RowSelection::from_filters(&filters[2..3]); + assert!(!selection.selects_any()); + assert_eq!(selection.selectors(), vec![RowSelector::skip(4)]); + } + + #[test] + fn test_iter() { + // use the iter() API to show it does what is expected and + // avoid accidental deletion + let selectors = vec![ + RowSelector::select(3), + RowSelector::skip(33), + RowSelector::select(4), + ]; + + let round_tripped: Vec = RowSelection::from(selectors.clone()) + .iter() + .copied() + .collect(); + assert_eq!(selectors, round_tripped); + } + + #[test] + fn test_row_count() { + let selection = RowSelection::from(vec![ + RowSelector::skip(34), + RowSelector::select(12), + RowSelector::skip(3), + RowSelector::select(35), + ]); + + assert_eq!(selection.row_count(), 12 + 35); + assert_eq!(selection.skipped_row_count(), 34 + 3); + + let selection = RowSelection::from(vec![RowSelector::select(12), RowSelector::select(35)]); + + assert_eq!(selection.row_count(), 12 + 35); + assert_eq!(selection.skipped_row_count(), 0); + + let selection = RowSelection::from(vec![RowSelector::skip(34), RowSelector::skip(3)]); + + assert_eq!(selection.row_count(), 0); + assert_eq!(selection.skipped_row_count(), 34 + 3); + + let selection = RowSelection::from(vec![]); + + assert_eq!(selection.row_count(), 0); + assert_eq!(selection.skipped_row_count(), 0); + } + + #[test] + fn test_mixed_backing_equality_mismatches() { + let mask = + RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true, true])); + + // Total row counts differ + let longer = RowSelection::from(vec![ + RowSelector::select(1), + RowSelector::skip(1), + RowSelector::select(2), + RowSelector::skip(1), + ]); + assert_ne!(mask, longer); + assert_ne!(longer, mask); + + // A selected bit falls inside a skip run + let skip_overlap = RowSelection::from(vec![RowSelector::skip(2), RowSelector::select(2)]); + assert_ne!(mask, skip_overlap); + + // Select run boundaries do not line up + let misaligned = RowSelection::from(vec![ + RowSelector::select(2), + RowSelector::skip(1), + RowSelector::select(1), + ]); + assert_ne!(mask, misaligned); + + let equal = RowSelection::from(vec![ + RowSelector::select(1), + RowSelector::skip(1), + RowSelector::select(2), + ]); + assert_eq!(mask, equal); + assert_eq!(equal, mask); + } + + #[test] + fn test_from_iter_all_mask_preserves_mask_backing() { + let a_bits = vec![true, false, true, true]; + let b_bits = vec![false, true, false]; + let c_bits = vec![true, true, false, false, true]; + + let parts = vec![ + RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits.clone())), + RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits.clone())), + RowSelection::from_boolean_buffer(BooleanBuffer::from(c_bits.clone())), + ]; + let collected: RowSelection = parts.into_iter().collect(); + + let combined = a_bits + .iter() + .chain(b_bits.iter()) + .chain(c_bits.iter()) + .copied() + .collect::>(); + let expected = RowSelection::from_filters(&[BooleanArray::from(combined)]); + + assert!(collected.as_mask().is_some()); + assert_eq!(collected, expected); + } + + #[test] + fn test_from_iter_mixed_backing_falls_back_to_selectors() { + let a_bits = vec![true, false, true]; + let b_selectors = vec![RowSelector::skip(2), RowSelector::select(3)]; + let c_bits = vec![false, true]; + + let parts = vec![ + RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits.clone())), + RowSelection::from(b_selectors), + RowSelection::from_boolean_buffer(BooleanBuffer::from(c_bits.clone())), + ]; + let collected: RowSelection = parts.into_iter().collect(); + + assert!(collected.as_mask().is_none()); + + let combined_bits = vec![ + true, false, true, false, false, true, true, true, false, true, + ]; + let expected = RowSelection::from_filters(&[BooleanArray::from(combined_bits)]); + assert_eq!(collected, expected); + } + + #[test] + fn test_from_iter_empty_yields_empty_selection() { + let collected: RowSelection = std::iter::empty::().collect(); + assert_eq!(collected, RowSelection::default()); + assert!(collected.as_mask().is_some()); + assert_eq!(collected.as_mask().unwrap().len(), 0); + } +} diff --git a/vendor/parquet/src/arrow/arrow_reader/selection/ranges.rs b/vendor/parquet/src/arrow/arrow_reader/selection/ranges.rs new file mode 100644 index 0000000..4c7860d --- /dev/null +++ b/vendor/parquet/src/arrow/arrow_reader/selection/ranges.rs @@ -0,0 +1,283 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Mapping the [`RowSelector`] runs of a selection onto ranges: the byte ranges +//! of the data pages that must be fetched ([`RowSelection::scan_ranges`]) and +//! the expansion of a selection to batch boundaries. +//! +//! Both are shared by the selector and mask backings, which stream their runs +//! from a slice and a [`MaskRunIter`] respectively. +//! +//! [`MaskRunIter`]: crate::arrow::arrow_reader::MaskRunIter + +use super::{RowSelection, RowSelector}; +use crate::file::page_index::offset_index::PageLocation; +use std::ops::Range; + +/// Byte ranges of the data pages containing at least one selected row. +#[inline] +pub(super) fn scan_ranges_from_selectors( + selectors: I, + page_locations: &[PageLocation], +) -> Vec> +where + I: IntoIterator, +{ + let mut ranges: Vec> = vec![]; + let mut row_offset = 0; + + let mut pages = page_locations.iter().peekable(); + let mut selectors = selectors.into_iter(); + let mut current_selector = selectors.next(); + let mut current_page = pages.next(); + + let mut current_page_included = false; + + while let Some((selector, page)) = current_selector.as_mut().zip(current_page) { + if !(selector.skip || current_page_included) { + let start = page.offset as u64; + let end = start + page.compressed_page_size as u64; + ranges.push(start..end); + current_page_included = true; + } + + if let Some(next_page) = pages.peek() { + if row_offset + selector.row_count > next_page.first_row_index as usize { + let remaining_in_page = next_page.first_row_index as usize - row_offset; + selector.row_count -= remaining_in_page; + row_offset += remaining_in_page; + current_page = pages.next(); + current_page_included = false; + + continue; + } else { + if row_offset + selector.row_count == next_page.first_row_index as usize { + current_page = pages.next(); + current_page_included = false; + } + row_offset += selector.row_count; + current_selector = selectors.next(); + } + } else { + if !(selector.skip || current_page_included) { + let start = page.offset as u64; + let end = start + page.compressed_page_size as u64; + ranges.push(start..end); + } + current_selector = selectors.next() + } + } + + ranges +} + +/// Grows each selected run to the batch boundaries containing it, merging the +/// runs that overlap as a result. +#[inline] +pub(super) fn expand_to_batch_boundaries_from_selectors( + selectors: I, + batch_size: usize, + total_rows: usize, +) -> RowSelection +where + I: IntoIterator, +{ + let mut expanded_ranges = Vec::new(); + let mut row_offset = 0; + + for selector in selectors { + if selector.skip { + row_offset += selector.row_count; + } else { + let start = row_offset; + let end = row_offset + selector.row_count; + + // Expand start to batch boundary + let expanded_start = (start / batch_size) * batch_size; + // Expand end to batch boundary + let expanded_end = end.div_ceil(batch_size) * batch_size; + let expanded_end = expanded_end.min(total_rows); + + expanded_ranges.push(expanded_start..expanded_end); + row_offset += selector.row_count; + } + } + + // Sort ranges by start position + expanded_ranges.sort_by_key(|range| range.start); + + // Merge overlapping or consecutive ranges + let mut merged_ranges: Vec> = Vec::new(); + for range in expanded_ranges { + if let Some(last) = merged_ranges.last_mut() { + if range.start <= last.end { + // Overlapping or consecutive - merge them + last.end = last.end.max(range.end); + } else { + // No overlap - add new range + merged_ranges.push(range); + } + } else { + // First range + merged_ranges.push(range); + } + } + + RowSelection::from_consecutive_ranges(merged_ranges.into_iter(), total_rows) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_scan_ranges() { + let index = vec![ + PageLocation { + offset: 0, + compressed_page_size: 10, + first_row_index: 0, + }, + PageLocation { + offset: 10, + compressed_page_size: 10, + first_row_index: 10, + }, + PageLocation { + offset: 20, + compressed_page_size: 10, + first_row_index: 20, + }, + PageLocation { + offset: 30, + compressed_page_size: 10, + first_row_index: 30, + }, + PageLocation { + offset: 40, + compressed_page_size: 10, + first_row_index: 40, + }, + PageLocation { + offset: 50, + compressed_page_size: 10, + first_row_index: 50, + }, + PageLocation { + offset: 60, + compressed_page_size: 10, + first_row_index: 60, + }, + ]; + + let selection = RowSelection::from(vec![ + // Skip first page + RowSelector::skip(10), + // Multiple selects in same page + RowSelector::select(3), + RowSelector::skip(3), + RowSelector::select(4), + // Select to page boundary + RowSelector::skip(5), + RowSelector::select(5), + // Skip full page past page boundary + RowSelector::skip(12), + // Select across page boundaries + RowSelector::select(12), + // Skip final page + RowSelector::skip(12), + ]); + + let ranges = selection.scan_ranges(&index); + + // assert_eq!(mask, vec![false, true, true, false, true, true, false]); + assert_eq!(ranges, vec![10..20, 20..30, 40..50, 50..60]); + assert_eq!( + selection.row_ranges_for_selected_pages(&index, 70), + vec![10..20, 20..30, 40..50, 50..60] + ); + + let selection = RowSelection::from(vec![ + // Skip first page + RowSelector::skip(10), + // Multiple selects in same page + RowSelector::select(3), + RowSelector::skip(3), + RowSelector::select(4), + // Select to page boundary + RowSelector::skip(5), + RowSelector::select(5), + // Skip full page past page boundary + RowSelector::skip(12), + // Select across page boundaries + RowSelector::select(12), + RowSelector::skip(1), + // Select across page boundaries including final page + RowSelector::select(8), + ]); + + let ranges = selection.scan_ranges(&index); + + // assert_eq!(mask, vec![false, true, true, false, true, true, true]); + assert_eq!(ranges, vec![10..20, 20..30, 40..50, 50..60, 60..70]); + + let selection = RowSelection::from(vec![ + // Skip first page + RowSelector::skip(10), + // Multiple selects in same page + RowSelector::select(3), + RowSelector::skip(3), + RowSelector::select(4), + // Select to page boundary + RowSelector::skip(5), + RowSelector::select(5), + // Skip full page past page boundary + RowSelector::skip(12), + // Select to final page boundary + RowSelector::select(12), + RowSelector::skip(1), + // Skip across final page boundary + RowSelector::skip(8), + // Select from final page + RowSelector::select(4), + ]); + + let ranges = selection.scan_ranges(&index); + + // assert_eq!(mask, vec![false, true, true, false, true, true, true]); + assert_eq!(ranges, vec![10..20, 20..30, 40..50, 50..60, 60..70]); + + let selection = RowSelection::from(vec![ + // Skip first page + RowSelector::skip(10), + // Multiple selects in same page + RowSelector::select(3), + RowSelector::skip(3), + RowSelector::select(4), + // Select to remaining in page and first row of next page + RowSelector::skip(5), + RowSelector::select(6), + // Skip remaining + RowSelector::skip(50), + ]); + + let ranges = selection.scan_ranges(&index); + + // assert_eq!(mask, vec![false, true, true, false, true, true, true]); + assert_eq!(ranges, vec![10..20, 20..30, 30..40]); + } +} diff --git a/vendor/parquet/src/arrow/arrow_reader/selection/selector.rs b/vendor/parquet/src/arrow/arrow_reader/selection/selector.rs new file mode 100644 index 0000000..a7ba939 --- /dev/null +++ b/vendor/parquet/src/arrow/arrow_reader/selection/selector.rs @@ -0,0 +1,469 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The run length backed representation of a [`RowSelection`]: [`RowSelector`] +//! and the primitives operating on a sequence of them. +//! +//! This is the counterpart of the bitmap backing in the `boolean` module, and +//! provides the same set of transforms (`split_off`, `trim`, `offset`, +//! `limit`) over `Vec` instead of a `BooleanBuffer`. +//! +//! [`RowSelection`]: crate::arrow::arrow_reader::RowSelection + +/// [`RowSelection`] is a collection of [`RowSelector`] used to skip rows when +/// scanning a parquet file +/// +/// [`RowSelection`]: crate::arrow::arrow_reader::RowSelection +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +pub struct RowSelector { + /// The number of rows + pub row_count: usize, + + /// If true, skip `row_count` rows + pub skip: bool, +} + +impl RowSelector { + /// Select `row_count` rows + pub fn select(row_count: usize) -> Self { + Self { + row_count, + skip: false, + } + } + + /// Skip `row_count` rows + pub fn skip(row_count: usize) -> Self { + Self { + row_count, + skip: true, + } + } +} + +/// Splits `selectors` at the first `row_count` rows, returning `(head, tail)`. +pub(super) fn split_off_selectors( + mut selectors: Vec, + row_count: usize, +) -> (Vec, Vec) { + let mut total_count = 0; + + // Find the index where the selector exceeds the row count + let find = selectors.iter().position(|selector| { + total_count += selector.row_count; + total_count > row_count + }); + + let split_idx = match find { + Some(idx) => idx, + None => return (selectors, Vec::new()), + }; + + // `selectors` keeps the head, `tail` takes the rest. The selector straddling + // the boundary is split between the two. + let mut tail = selectors.split_off(split_idx); + + // Always present as `split_idx < selectors.len` + let next = tail.first_mut().unwrap(); + let overflow = total_count - row_count; + + if next.row_count != overflow { + selectors.push(RowSelector { + row_count: next.row_count - overflow, + skip: next.skip, + }) + } + next.row_count = overflow; + + (selectors, tail) +} + +/// Skips the first `offset` selected rows of `selectors`. +pub(super) fn offset_selectors(mut selectors: Vec, offset: usize) -> Vec { + let mut selected_count = 0; + let mut skipped_count = 0; + + // Find the index where the selector exceeds the row count + let find = selectors.iter().position(|selector| match selector.skip { + true => { + skipped_count += selector.row_count; + false + } + false => { + selected_count += selector.row_count; + selected_count > offset + } + }); + + let split_idx = match find { + Some(idx) => idx, + None => { + selectors.clear(); + return selectors; + } + }; + + let mut new_selectors = Vec::with_capacity(selectors.len() - split_idx + 1); + new_selectors.push(RowSelector::skip(skipped_count + offset)); + new_selectors.push(RowSelector::select(selected_count - offset)); + new_selectors.extend_from_slice(&selectors[split_idx + 1..]); + + new_selectors +} + +/// Keeps only the first `limit` selected rows of `selectors`. +pub(super) fn limit_selectors( + mut selectors: Vec, + mut limit: usize, +) -> Vec { + if limit == 0 { + selectors.clear(); + } + + for (idx, selection) in selectors.iter_mut().enumerate() { + if !selection.skip { + if selection.row_count >= limit { + selection.row_count = limit; + selectors.truncate(idx + 1); + break; + } else { + limit -= selection.row_count; + } + } + } + selectors +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::arrow_reader::selection::RowSelection; + + #[test] + fn test_from_selectors_skips_empty_selectors() { + let selection = RowSelection::from(vec![ + RowSelector::select(0), + RowSelector::skip(0), + RowSelector::select(2), + RowSelector::select(0), + RowSelector::skip(1), + ]); + assert_eq!( + selection.selectors(), + vec![RowSelector::select(2), RowSelector::skip(1)] + ); + } + + #[test] + fn test_split_off() { + let mut selection = RowSelection::from(vec![ + RowSelector::skip(34), + RowSelector::select(12), + RowSelector::skip(3), + RowSelector::select(35), + ]); + + let split = selection.split_off(34); + assert_eq!(split.selectors(), vec![RowSelector::skip(34)]); + assert_eq!( + selection.selectors(), + vec![ + RowSelector::select(12), + RowSelector::skip(3), + RowSelector::select(35) + ] + ); + + let split = selection.split_off(5); + assert_eq!(split.selectors(), vec![RowSelector::select(5)]); + assert_eq!( + selection.selectors(), + vec![ + RowSelector::select(7), + RowSelector::skip(3), + RowSelector::select(35) + ] + ); + + let split = selection.split_off(8); + assert_eq!( + split.selectors(), + vec![RowSelector::select(7), RowSelector::skip(1)] + ); + assert_eq!( + selection.selectors(), + vec![RowSelector::skip(2), RowSelector::select(35)] + ); + + let split = selection.split_off(200); + assert_eq!( + split.selectors(), + vec![RowSelector::skip(2), RowSelector::select(35)] + ); + assert!(selection.selectors().is_empty()); + } + + #[test] + fn test_offset() { + let selection = RowSelection::from(vec![ + RowSelector::select(5), + RowSelector::skip(23), + RowSelector::select(7), + RowSelector::skip(33), + RowSelector::select(6), + ]); + + let selection = selection.offset(2); + assert_eq!( + selection.selectors(), + vec![ + RowSelector::skip(2), + RowSelector::select(3), + RowSelector::skip(23), + RowSelector::select(7), + RowSelector::skip(33), + RowSelector::select(6), + ] + ); + + let selection = selection.offset(5); + assert_eq!( + selection.selectors(), + vec![ + RowSelector::skip(30), + RowSelector::select(5), + RowSelector::skip(33), + RowSelector::select(6), + ] + ); + + let selection = selection.offset(3); + assert_eq!( + selection.selectors(), + vec![ + RowSelector::skip(33), + RowSelector::select(2), + RowSelector::skip(33), + RowSelector::select(6), + ] + ); + + let selection = selection.offset(2); + assert_eq!( + selection.selectors(), + vec![RowSelector::skip(68), RowSelector::select(6),] + ); + + let selection = selection.offset(3); + assert_eq!( + selection.selectors(), + vec![RowSelector::skip(71), RowSelector::select(3),] + ); + } + + #[test] + fn test_combine() { + let a = vec![ + RowSelector::skip(3), + RowSelector::skip(3), + RowSelector::select(10), + RowSelector::skip(4), + ]; + + let b = vec![ + RowSelector::skip(3), + RowSelector::skip(3), + RowSelector::select(10), + RowSelector::skip(4), + RowSelector::skip(0), + ]; + + let c = vec![ + RowSelector::skip(2), + RowSelector::skip(4), + RowSelector::select(3), + RowSelector::select(3), + RowSelector::select(4), + RowSelector::skip(3), + RowSelector::skip(1), + RowSelector::skip(0), + ]; + + let expected = RowSelection::from(vec![ + RowSelector::skip(6), + RowSelector::select(10), + RowSelector::skip(4), + ]); + + assert_eq!(RowSelection::from_iter(a), expected); + assert_eq!(RowSelection::from_iter(b), expected); + assert_eq!(RowSelection::from_iter(c), expected); + } + + #[test] + fn test_combine_2elements() { + let a = vec![RowSelector::select(10), RowSelector::select(5)]; + let a_expect = vec![RowSelector::select(15)]; + assert_eq!(RowSelection::from_iter(a).selectors(), a_expect); + + let b = vec![RowSelector::select(10), RowSelector::skip(5)]; + let b_expect = vec![RowSelector::select(10), RowSelector::skip(5)]; + assert_eq!(RowSelection::from_iter(b).selectors(), b_expect); + + let c = vec![RowSelector::skip(10), RowSelector::select(5)]; + let c_expect = vec![RowSelector::skip(10), RowSelector::select(5)]; + assert_eq!(RowSelection::from_iter(c).selectors(), c_expect); + + let d = vec![RowSelector::skip(10), RowSelector::skip(5)]; + let d_expect = vec![RowSelector::skip(15)]; + assert_eq!(RowSelection::from_iter(d).selectors(), d_expect); + } + + #[test] + fn test_from_one_and_empty() { + let a = vec![RowSelector::select(10)]; + let selection1 = RowSelection::from(a.clone()); + assert_eq!(selection1.selectors(), a); + + let b = vec![]; + let selection1 = RowSelection::from(b.clone()); + assert_eq!(selection1.selectors(), b) + } + + #[test] + fn test_limit() { + // Limit to existing limit should no-op + let selection = RowSelection::from(vec![RowSelector::select(10), RowSelector::skip(90)]); + let limited = selection.limit(10); + assert_eq!(RowSelection::from(vec![RowSelector::select(10)]), limited); + + let selection = RowSelection::from(vec![ + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(10), + ]); + + let limited = selection.clone().limit(5); + let expected = vec![RowSelector::select(5)]; + assert_eq!(limited.selectors(), expected); + + let limited = selection.clone().limit(15); + let expected = vec![ + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(5), + ]; + assert_eq!(limited.selectors(), expected); + + let limited = selection.clone().limit(0); + let expected = vec![]; + assert_eq!(limited.selectors(), expected); + + let limited = selection.clone().limit(30); + let expected = vec![ + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(10), + ]; + assert_eq!(limited.selectors(), expected); + + let limited = selection.limit(100); + let expected = vec![ + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(10), + ]; + assert_eq!(limited.selectors(), expected); + } + + #[test] + fn test_from_ranges() { + let ranges = [1..3, 4..6, 6..6, 8..8, 9..10]; + let selection = RowSelection::from_consecutive_ranges(ranges.into_iter(), 10); + assert_eq!( + selection.selectors(), + vec![ + RowSelector::skip(1), + RowSelector::select(2), + RowSelector::skip(1), + RowSelector::select(2), + RowSelector::skip(3), + RowSelector::select(1) + ] + ); + + let out_of_order_ranges = [1..3, 8..10, 4..7]; + let result = std::panic::catch_unwind(|| { + RowSelection::from_consecutive_ranges(out_of_order_ranges.into_iter(), 10) + }); + assert!(result.is_err()); + } + + #[test] + fn test_empty_selector() { + let selection = RowSelection::from(vec![ + RowSelector::skip(0), + RowSelector::select(2), + RowSelector::skip(0), + RowSelector::select(2), + ]); + assert_eq!(selection.selectors(), vec![RowSelector::select(4)]); + + let selection = RowSelection::from(vec![ + RowSelector::select(0), + RowSelector::skip(2), + RowSelector::select(0), + RowSelector::skip(2), + ]); + assert_eq!(selection.selectors(), vec![RowSelector::skip(4)]); + } + + #[test] + fn test_trim() { + let selection = RowSelection::from(vec![ + RowSelector::skip(34), + RowSelector::select(12), + RowSelector::skip(3), + RowSelector::select(35), + ]); + + let expected = vec![ + RowSelector::skip(34), + RowSelector::select(12), + RowSelector::skip(3), + RowSelector::select(35), + ]; + + assert_eq!(selection.trim().selectors(), expected); + + let selection = RowSelection::from(vec![ + RowSelector::skip(34), + RowSelector::select(12), + RowSelector::skip(3), + ]); + + let expected = vec![RowSelector::skip(34), RowSelector::select(12)]; + + assert_eq!(selection.trim().selectors(), expected); + } +} diff --git a/vendor/parquet/src/arrow/arrow_reader/statistics.rs b/vendor/parquet/src/arrow/arrow_reader/statistics.rs new file mode 100644 index 0000000..c096448 --- /dev/null +++ b/vendor/parquet/src/arrow/arrow_reader/statistics.rs @@ -0,0 +1,1985 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`StatisticsConverter`] to convert statistics in parquet format to arrow [`ArrayRef`]. + +/// Notice that all the corresponding tests are in +/// `arrow-rs/parquet/tests/arrow_reader/statistics.rs`. +use crate::arrow::buffer::bit_util::sign_extend_be; +use crate::arrow::parquet_column; +use crate::basic::Type as PhysicalType; +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::{ParquetColumnIndex, ParquetOffsetIndex, RowGroupMetaData}; +use crate::file::page_index::column_index::ColumnIndexMetaData; +use crate::file::statistics::Statistics as ParquetStatistics; +use crate::schema::types::SchemaDescriptor; +use arrow_array::builder::{ + BinaryBuilder, BinaryViewBuilder, BooleanBuilder, Date32Builder, Date64Builder, + Decimal32Builder, Decimal64Builder, FixedSizeBinaryBuilder, Float16Builder, Float32Builder, + Float64Builder, Int8Builder, Int16Builder, Int32Builder, Int64Builder, LargeBinaryBuilder, + LargeStringBuilder, StringBuilder, StringViewBuilder, Time32MillisecondBuilder, + Time32SecondBuilder, Time64MicrosecondBuilder, Time64NanosecondBuilder, + TimestampMicrosecondBuilder, TimestampMillisecondBuilder, TimestampNanosecondBuilder, + TimestampSecondBuilder, UInt8Builder, UInt16Builder, UInt32Builder, UInt64Builder, +}; +use arrow_array::{ + ArrayRef, BinaryArray, BooleanArray, Date32Array, Date64Array, Decimal32Array, Decimal64Array, + Decimal128Array, Decimal256Array, Float16Array, Float32Array, Float64Array, Int8Array, + Int16Array, Int32Array, Int64Array, LargeBinaryArray, Time32MillisecondArray, + Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray, TimestampMicrosecondArray, + TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, UInt8Array, + UInt16Array, UInt32Array, UInt64Array, new_null_array, +}; +use arrow_buffer::{NullBufferBuilder, i256}; +use arrow_schema::{DataType, Field, Schema, TimeUnit}; +use half::f16; +use std::sync::Arc; + +// Convert the bytes array to i32. +// The endian of the input bytes array must be big-endian. +pub(crate) fn from_bytes_to_i32(b: &[u8]) -> i32 { + // The bytes array are from parquet file and must be the big-endian. + // The endian is defined by parquet format, and the reference document + // https://github.com/apache/parquet-format/blob/54e53e5d7794d383529dd30746378f19a12afd58/src/main/thrift/parquet.thrift#L66 + i32::from_be_bytes(sign_extend_be::<4>(b)) +} + +// Convert the bytes array to i64. +// The endian of the input bytes array must be big-endian. +pub(crate) fn from_bytes_to_i64(b: &[u8]) -> i64 { + i64::from_be_bytes(sign_extend_be::<8>(b)) +} + +// Convert the bytes array to i128. +// The endian of the input bytes array must be big-endian. +pub(crate) fn from_bytes_to_i128(b: &[u8]) -> i128 { + i128::from_be_bytes(sign_extend_be::<16>(b)) +} + +// Convert the bytes array to i256. +// The endian of the input bytes array must be big-endian. +pub(crate) fn from_bytes_to_i256(b: &[u8]) -> i256 { + i256::from_be_bytes(sign_extend_be::<32>(b)) +} + +// Convert the bytes array to f16 +pub(crate) fn from_bytes_to_f16(b: &[u8]) -> Option { + match b { + [low, high] => Some(f16::from_be_bytes([*high, *low])), + _ => None, + } +} + +/// Define an adapter iterator for extracting statistics from an iterator of +/// `ParquetStatistics` +/// +/// +/// Handles checking if the statistics are present and valid with the correct type. +/// +/// Parameters: +/// * `$iterator_type` is the name of the iterator type (e.g. `MinBooleanStatsIterator`) +/// * `$func` is the function to call to get the value (e.g. `min` or `max`) +/// * `$parquet_statistics_type` is the type of the statistics (e.g. `ParquetStatistics::Boolean`) +/// * `$stat_value_type` is the type of the statistics value (e.g. `bool`) +macro_rules! make_stats_iterator { + ($iterator_type:ident, $func:ident, $parquet_statistics_type:path, $stat_value_type:ty) => { + /// Maps an iterator of `ParquetStatistics` into an iterator of + /// `&$stat_value_type`` + /// + /// Yielded elements: + /// * Some(stats) if valid + /// * None if the statistics are not present, not valid, or not $stat_value_type + struct $iterator_type<'a, I> + where + I: Iterator>, + { + iter: I, + } + + impl<'a, I> $iterator_type<'a, I> + where + I: Iterator>, + { + /// Create a new iterator to extract the statistics + fn new(iter: I) -> Self { + Self { iter } + } + } + + /// Implement the Iterator trait for the iterator + impl<'a, I> Iterator for $iterator_type<'a, I> + where + I: Iterator>, + { + type Item = Option<&'a $stat_value_type>; + + /// return the next statistics value + fn next(&mut self) -> Option { + let next = self.iter.next(); + next.map(|x| { + x.and_then(|stats| match stats { + $parquet_statistics_type(s) => s.$func(), + _ => None, + }) + }) + } + + fn size_hint(&self) -> (usize, Option) { + self.iter.size_hint() + } + } + }; +} + +make_stats_iterator!( + MinBooleanStatsIterator, + min_opt, + ParquetStatistics::Boolean, + bool +); +make_stats_iterator!( + MaxBooleanStatsIterator, + max_opt, + ParquetStatistics::Boolean, + bool +); +make_stats_iterator!( + MinInt32StatsIterator, + min_opt, + ParquetStatistics::Int32, + i32 +); +make_stats_iterator!( + MaxInt32StatsIterator, + max_opt, + ParquetStatistics::Int32, + i32 +); +make_stats_iterator!( + MinInt64StatsIterator, + min_opt, + ParquetStatistics::Int64, + i64 +); +make_stats_iterator!( + MaxInt64StatsIterator, + max_opt, + ParquetStatistics::Int64, + i64 +); +make_stats_iterator!( + MinFloatStatsIterator, + min_opt, + ParquetStatistics::Float, + f32 +); +make_stats_iterator!( + MaxFloatStatsIterator, + max_opt, + ParquetStatistics::Float, + f32 +); +make_stats_iterator!( + MinDoubleStatsIterator, + min_opt, + ParquetStatistics::Double, + f64 +); +make_stats_iterator!( + MaxDoubleStatsIterator, + max_opt, + ParquetStatistics::Double, + f64 +); +make_stats_iterator!( + MinByteArrayStatsIterator, + min_bytes_opt, + ParquetStatistics::ByteArray, + [u8] +); +make_stats_iterator!( + MaxByteArrayStatsIterator, + max_bytes_opt, + ParquetStatistics::ByteArray, + [u8] +); +make_stats_iterator!( + MinFixedLenByteArrayStatsIterator, + min_bytes_opt, + ParquetStatistics::FixedLenByteArray, + [u8] +); +make_stats_iterator!( + MaxFixedLenByteArrayStatsIterator, + max_bytes_opt, + ParquetStatistics::FixedLenByteArray, + [u8] +); + +/// Special iterator adapter for extracting i128 values from from an iterator of +/// `ParquetStatistics` +/// +/// Handles checking if the statistics are present and valid with the correct type. +/// +/// Depending on the parquet file, the statistics for `Decimal128` can be stored as +/// `Int32`, `Int64` or `ByteArray` or `FixedSizeByteArray` :mindblown: +/// +/// This iterator handles all cases, extracting the values +/// and converting it to `stat_value_type`. +/// +/// Parameters: +/// * `$iterator_type` is the name of the iterator type (e.g. `MinBooleanStatsIterator`) +/// * `$func` is the function to call to get the value (e.g. `min` or `max`) +/// * `$bytes_func` is the function to call to get the value as bytes (e.g. `min_bytes` or `max_bytes`) +/// * `$stat_value_type` is the type of the statistics value (e.g. `i128`) +/// * `convert_func` is the function to convert the bytes to stats value (e.g. `from_bytes_to_i128`) +macro_rules! make_decimal_stats_iterator { + ($iterator_type:ident, $func:ident, $bytes_func:ident, $stat_value_type:ident, $convert_func: ident) => { + struct $iterator_type<'a, I> + where + I: Iterator>, + { + iter: I, + } + + impl<'a, I> $iterator_type<'a, I> + where + I: Iterator>, + { + fn new(iter: I) -> Self { + Self { iter } + } + } + + impl<'a, I> Iterator for $iterator_type<'a, I> + where + I: Iterator>, + { + type Item = Option<$stat_value_type>; + + fn next(&mut self) -> Option { + let next = self.iter.next(); + next.map(|x| { + x.and_then(|stats| match stats { + ParquetStatistics::Int32(s) => { + s.$func().map(|x| $stat_value_type::from(*x)) + } + ParquetStatistics::Int64(s) => s + .$func() + .map(|x| $stat_value_type::try_from(*x).ok()) + .flatten(), + ParquetStatistics::ByteArray(s) => s.$bytes_func().map($convert_func), + ParquetStatistics::FixedLenByteArray(s) => { + s.$bytes_func().map($convert_func) + } + _ => None, + }) + }) + } + + fn size_hint(&self) -> (usize, Option) { + self.iter.size_hint() + } + } + }; +} + +make_decimal_stats_iterator!( + MinDecimal32StatsIterator, + min_opt, + min_bytes_opt, + i32, + from_bytes_to_i32 +); +make_decimal_stats_iterator!( + MaxDecimal32StatsIterator, + max_opt, + max_bytes_opt, + i32, + from_bytes_to_i32 +); +make_decimal_stats_iterator!( + MinDecimal64StatsIterator, + min_opt, + min_bytes_opt, + i64, + from_bytes_to_i64 +); +make_decimal_stats_iterator!( + MaxDecimal64StatsIterator, + max_opt, + max_bytes_opt, + i64, + from_bytes_to_i64 +); +make_decimal_stats_iterator!( + MinDecimal128StatsIterator, + min_opt, + min_bytes_opt, + i128, + from_bytes_to_i128 +); +make_decimal_stats_iterator!( + MaxDecimal128StatsIterator, + max_opt, + max_bytes_opt, + i128, + from_bytes_to_i128 +); +make_decimal_stats_iterator!( + MinDecimal256StatsIterator, + min_opt, + min_bytes_opt, + i256, + from_bytes_to_i256 +); +make_decimal_stats_iterator!( + MaxDecimal256StatsIterator, + max_opt, + max_bytes_opt, + i256, + from_bytes_to_i256 +); + +/// Special macro to combine the statistics iterators for min and max. +/// This is used to avoid repeating the same code for min and max statistics extractions +/// +/// Parameters: +/// data_type: The data type of the statistics (e.g. `DataType::Int32`) +/// iterator: The iterator of [`ParquetStatistics`] to extract the statistics from. +macro_rules! get_statistics { + (Min, $data_type: ident, $iterator: ident, $physical_type: ident) => { + get_statistics!( + $data_type, + $iterator, + $physical_type, + MinBooleanStatsIterator, + MinInt32StatsIterator, + MinInt64StatsIterator, + MinFloatStatsIterator, + MinDoubleStatsIterator, + MinByteArrayStatsIterator, + MinFixedLenByteArrayStatsIterator, + MinDecimal32StatsIterator, + MinDecimal64StatsIterator, + MinDecimal128StatsIterator, + MinDecimal256StatsIterator, + min_statistics + ) + }; + (Max, $data_type: ident, $iterator: ident, $physical_type: ident) => { + get_statistics!( + $data_type, + $iterator, + $physical_type, + MaxBooleanStatsIterator, + MaxInt32StatsIterator, + MaxInt64StatsIterator, + MaxFloatStatsIterator, + MaxDoubleStatsIterator, + MaxByteArrayStatsIterator, + MaxFixedLenByteArrayStatsIterator, + MaxDecimal32StatsIterator, + MaxDecimal64StatsIterator, + MaxDecimal128StatsIterator, + MaxDecimal256StatsIterator, + max_statistics + ) + }; + ( + $data_type: ident, + $iterator: ident, + $physical_type: ident, + $boolean_iter: ident, + $int32_iter: ident, + $int64_iter: ident, + $float_iter: ident, + $double_iter: ident, + $byte_array_iter: ident, + $fixed_len_byte_array_iter: ident, + $decimal32_iter: ident, + $decimal64_iter: ident, + $decimal128_iter: ident, + $decimal256_iter: ident, + $dictionary_statistics: ident + ) => { + match $data_type { + DataType::Boolean => Ok(Arc::new(BooleanArray::from_iter( + $boolean_iter::new($iterator).map(|x| x.copied()), + ))), + DataType::Int8 => Ok(Arc::new(Int8Array::from_iter( + $int32_iter::new($iterator).map(|x| { + x.and_then(|x| i8::try_from(*x).ok()) + }), + ))), + DataType::Int16 => Ok(Arc::new(Int16Array::from_iter( + $int32_iter::new($iterator).map(|x| { + x.and_then(|x| i16::try_from(*x).ok()) + }), + ))), + DataType::Int32 => Ok(Arc::new(Int32Array::from_iter( + $int32_iter::new($iterator).map(|x| x.copied()), + ))), + DataType::Int64 => Ok(Arc::new(Int64Array::from_iter( + $int64_iter::new($iterator).map(|x| x.copied()), + ))), + DataType::UInt8 => Ok(Arc::new(UInt8Array::from_iter( + $int32_iter::new($iterator).map(|x| { + x.and_then(|x| u8::try_from(*x).ok()) + }), + ))), + DataType::UInt16 => Ok(Arc::new(UInt16Array::from_iter( + $int32_iter::new($iterator).map(|x| { + x.and_then(|x| u16::try_from(*x).ok()) + }), + ))), + DataType::UInt32 => Ok(Arc::new(UInt32Array::from_iter( + $int32_iter::new($iterator).map(|x| x.map(|x| *x as u32)), + ))), + DataType::UInt64 => Ok(Arc::new(UInt64Array::from_iter( + $int64_iter::new($iterator).map(|x| x.map(|x| *x as u64)), + ))), + DataType::Float16 => Ok(Arc::new(Float16Array::from_iter( + $fixed_len_byte_array_iter::new($iterator).map(|x| x.and_then(|x| { + from_bytes_to_f16(x) + })), + ))), + DataType::Float32 => Ok(Arc::new(Float32Array::from_iter( + $float_iter::new($iterator).map(|x| x.copied()), + ))), + DataType::Float64 => Ok(Arc::new(Float64Array::from_iter( + $double_iter::new($iterator).map(|x| x.copied()), + ))), + DataType::Date32 => Ok(Arc::new(Date32Array::from_iter( + $int32_iter::new($iterator).map(|x| x.copied()), + ))), + DataType::Date64 if $physical_type == Some(PhysicalType::INT32) => Ok(Arc::new(Date64Array::from_iter( + $int32_iter::new($iterator) + .map(|x| x.map(|x| i64::from(*x) * 24 * 60 * 60 * 1000))))), + DataType::Date64 if $physical_type == Some(PhysicalType::INT64) => Ok(Arc::new(Date64Array::from_iter( + $int64_iter::new($iterator).map(|x| x.copied()),))), + DataType::Timestamp(unit, timezone) =>{ + let iter = $int64_iter::new($iterator).map(|x| x.copied()); + Ok(match unit { + TimeUnit::Second => Arc::new(TimestampSecondArray::from_iter(iter).with_timezone_opt(timezone.clone())), + TimeUnit::Millisecond => Arc::new(TimestampMillisecondArray::from_iter(iter).with_timezone_opt(timezone.clone())), + TimeUnit::Microsecond => Arc::new(TimestampMicrosecondArray::from_iter(iter).with_timezone_opt(timezone.clone())), + TimeUnit::Nanosecond => Arc::new(TimestampNanosecondArray::from_iter(iter).with_timezone_opt(timezone.clone())), + }) + }, + DataType::Time32(unit) => { + Ok(match unit { + TimeUnit::Second => Arc::new(Time32SecondArray::from_iter( + $int32_iter::new($iterator).map(|x| x.copied()), + )), + TimeUnit::Millisecond => Arc::new(Time32MillisecondArray::from_iter( + $int32_iter::new($iterator).map(|x| x.copied()), + )), + _ => { + let len = $iterator.count(); + // don't know how to extract statistics, so return a null array + new_null_array($data_type, len) + } + }) + }, + DataType::Time64(unit) => { + Ok(match unit { + TimeUnit::Microsecond => Arc::new(Time64MicrosecondArray::from_iter( + $int64_iter::new($iterator).map(|x| x.copied()), + )), + TimeUnit::Nanosecond => Arc::new(Time64NanosecondArray::from_iter( + $int64_iter::new($iterator).map(|x| x.copied()), + )), + _ => { + let len = $iterator.count(); + // don't know how to extract statistics, so return a null array + new_null_array($data_type, len) + } + }) + }, + DataType::Binary => Ok(Arc::new(BinaryArray::from_iter( + $byte_array_iter::new($iterator) + ))), + DataType::LargeBinary => Ok(Arc::new(LargeBinaryArray::from_iter( + $byte_array_iter::new($iterator) + ))), + DataType::Utf8 => { + let iterator = $byte_array_iter::new($iterator); + let mut builder = StringBuilder::new(); + for x in iterator { + let Some(x) = x else { + builder.append_null(); // no statistics value + continue; + }; + + let Ok(x) = std::str::from_utf8(x) else { + builder.append_null(); + continue; + }; + + builder.append_value(x); + } + Ok(Arc::new(builder.finish())) + }, + DataType::LargeUtf8 => { + let iterator = $byte_array_iter::new($iterator); + let mut builder = LargeStringBuilder::new(); + for x in iterator { + let Some(x) = x else { + builder.append_null(); // no statistics value + continue; + }; + + let Ok(x) = std::str::from_utf8(x) else { + builder.append_null(); + continue; + }; + + builder.append_value(x); + } + Ok(Arc::new(builder.finish())) + }, + DataType::FixedSizeBinary(size) => { + let iterator = $fixed_len_byte_array_iter::new($iterator); + let mut builder = FixedSizeBinaryBuilder::new(*size); + for x in iterator { + let Some(x) = x else { + builder.append_null(); // no statistics value + continue; + }; + + // ignore invalid values + if x.len().try_into() != Ok(*size){ + builder.append_null(); + continue; + } + + builder.append_value(x).expect("ensure to append successfully here, because size have been checked before"); + } + Ok(Arc::new(builder.finish())) + }, + DataType::Decimal32(precision, scale) => { + let arr = Decimal32Array::from_iter( + $decimal32_iter::new($iterator) + ).with_precision_and_scale(*precision, *scale)?; + Ok(Arc::new(arr)) + }, + DataType::Decimal64(precision, scale) => { + let arr = Decimal64Array::from_iter( + $decimal64_iter::new($iterator) + ).with_precision_and_scale(*precision, *scale)?; + Ok(Arc::new(arr)) + }, + DataType::Decimal128(precision, scale) => { + let arr = Decimal128Array::from_iter( + $decimal128_iter::new($iterator) + ).with_precision_and_scale(*precision, *scale)?; + Ok(Arc::new(arr)) + }, + DataType::Decimal256(precision, scale) => { + let arr = Decimal256Array::from_iter( + $decimal256_iter::new($iterator) + ).with_precision_and_scale(*precision, *scale)?; + Ok(Arc::new(arr)) + }, + DataType::Dictionary(_, value_type) => { + $dictionary_statistics(value_type, $iterator, $physical_type) + }, + DataType::Utf8View => { + let iterator = $byte_array_iter::new($iterator); + let mut builder = StringViewBuilder::new(); + for x in iterator { + let Some(x) = x else { + builder.append_null(); // no statistics value + continue; + }; + + let Ok(x) = std::str::from_utf8(x) else { + builder.append_null(); + continue; + }; + + builder.append_value(x); + } + Ok(Arc::new(builder.finish())) + }, + DataType::BinaryView => { + let iterator = $byte_array_iter::new($iterator); + let mut builder = BinaryViewBuilder::new(); + for x in iterator { + let Some(x) = x else { + builder.append_null(); // no statistics value + continue; + }; + + builder.append_value(x); + } + Ok(Arc::new(builder.finish())) + } + + DataType::Map(_,_) | + DataType::Duration(_) | + DataType::Interval(_) | + DataType::Date64 | // required to cover $physical_type match guard + DataType::Null | + DataType::List(_) | + DataType::ListView(_) | + DataType::FixedSizeList(_, _) | + DataType::LargeList(_) | + DataType::LargeListView(_) | + DataType::Struct(_) | + DataType::Union(_, _) | + DataType::RunEndEncoded(_, _) => { + let len = $iterator.count(); + // don't know how to extract statistics, so return a null array + Ok(new_null_array($data_type, len)) + } + } + }; +} + +macro_rules! get_data_page_statistics { + (Min, $data_type: ident, $iterator: ident, $physical_type: ident) => { + get_data_page_statistics!( + $data_type, + $iterator, + $physical_type, + min_values_iter, + min_page_statistics + ) + }; + (Max, $data_type: ident, $iterator: ident, $physical_type: ident) => { + get_data_page_statistics!( + $data_type, + $iterator, + $physical_type, + max_values_iter, + max_page_statistics + ) + }; + ( + $data_type: ident, + $iterator: ident, + $physical_type: ident, + $values_iter: ident, + $page_statistics: ident + ) => {{ + let chunks: Vec<(usize, &ColumnIndexMetaData)> = $iterator.collect(); + let capacity: usize = chunks.iter().map(|c| c.0).sum(); + match $data_type { + DataType::Boolean => { + let mut b = BooleanBuilder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::BOOLEAN(index) => { + for val in index.$values_iter() { + b.append_option(val.copied()); + } + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::UInt8 => { + let mut b = UInt8Builder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT32(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.and_then(|&x| u8::try_from(x).ok())), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::UInt16 => { + let mut b = UInt16Builder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT32(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.and_then(|&x| u16::try_from(x).ok())), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::UInt32 => { + let mut b = UInt32Builder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT32(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.map(|&x| x as u32)), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::UInt64 => { + let mut b = UInt64Builder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT64(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.map(|&x| x as u64)), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::Int8 => { + let mut b = Int8Builder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT32(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.and_then(|&x| i8::try_from(x).ok())), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::Int16 => { + let mut b = Int16Builder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT32(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.and_then(|&x| i16::try_from(x).ok())), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::Int32 => { + let mut b = Int32Builder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT32(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.copied()), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::Int64 => { + let mut b = Int64Builder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT64(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.copied()), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::Float16 => { + let mut b = Float16Builder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.and_then(|x| from_bytes_to_f16(x))), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::Float32 => { + let mut b = Float32Builder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::FLOAT(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.copied()), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::Float64 => { + let mut b = Float64Builder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::DOUBLE(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.copied()), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::Binary => { + let mut b = BinaryBuilder::with_capacity(capacity, capacity * 10); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::BYTE_ARRAY(index) => { + for val in index.$values_iter() { + b.append_option(val.map(|x| x.as_ref())); + } + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::LargeBinary => { + let mut b = LargeBinaryBuilder::with_capacity(capacity, capacity * 10); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::BYTE_ARRAY(index) => { + for val in index.$values_iter() { + b.append_option(val.map(|x| x.as_ref())); + } + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::Utf8 => { + let mut b = StringBuilder::with_capacity(capacity, capacity * 10); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::BYTE_ARRAY(index) => { + for val in index.$values_iter() { + match val { + Some(x) => match std::str::from_utf8(x.as_ref()) { + Ok(s) => b.append_value(s), + _ => b.append_null(), + } + None => b.append_null(), + } + } + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::LargeUtf8 => { + let mut b = LargeStringBuilder::with_capacity(capacity, capacity * 10); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::BYTE_ARRAY(index) => { + for val in index.$values_iter() { + match val { + Some(x) => match std::str::from_utf8(x.as_ref()) { + Ok(s) => b.append_value(s), + _ => b.append_null(), + } + None => b.append_null(), + } + } + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::Dictionary(_, value_type) => { + $page_statistics(value_type, chunks.into_iter(), $physical_type) + }, + DataType::Timestamp(unit, timezone) => { + match unit { + TimeUnit::Second => { + let mut b = TimestampSecondBuilder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT64(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.copied()), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish().with_timezone_opt(timezone.clone()))) + } + TimeUnit::Millisecond => { + let mut b = TimestampMillisecondBuilder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT64(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.copied()), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish().with_timezone_opt(timezone.clone()))) + } + TimeUnit::Microsecond => { + let mut b = TimestampMicrosecondBuilder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT64(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.copied()), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish().with_timezone_opt(timezone.clone()))) + } + TimeUnit::Nanosecond => { + let mut b = TimestampNanosecondBuilder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT64(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.copied()), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish().with_timezone_opt(timezone.clone()))) + } + } + }, + DataType::Date32 => { + let mut b = Date32Builder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT32(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.copied()), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::Date64 if $physical_type == Some(PhysicalType::INT32)=> { + let mut b = Date64Builder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT32(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.map(|&x| (x as i64) * 24 * 60 * 60 * 1000)), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::Date64 if $physical_type == Some(PhysicalType::INT64) => { + let mut b = Date64Builder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT64(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.copied()), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::Decimal32(precision, scale) => { + let mut b = Decimal32Builder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT32(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.copied()), + ); + } + ColumnIndexMetaData::INT64(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.and_then(|&x| i32::try_from(x).ok())), + ); + } + ColumnIndexMetaData::BYTE_ARRAY(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.map(|x| from_bytes_to_i32(x.as_ref()))), + ); + } + ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.map(|x| from_bytes_to_i32(x.as_ref()))), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.with_precision_and_scale(*precision, *scale)?.finish())) + }, + DataType::Decimal64(precision, scale) => { + let mut b = Decimal64Builder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT32(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.map(|x| *x as i64)), + ); + } + ColumnIndexMetaData::INT64(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.copied()), + ); + } + ColumnIndexMetaData::BYTE_ARRAY(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.map(|x| from_bytes_to_i64(x.as_ref()))), + ); + } + ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.map(|x| from_bytes_to_i64(x.as_ref()))), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.with_precision_and_scale(*precision, *scale)?.finish())) + }, + DataType::Decimal128(precision, scale) => { + let mut b = Decimal128Array::builder(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT32(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.map(|x| *x as i128)), + ); + } + ColumnIndexMetaData::INT64(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.map(|x| *x as i128)), + ); + } + ColumnIndexMetaData::BYTE_ARRAY(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.map(|x| from_bytes_to_i128(x.as_ref()))), + ); + } + ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.map(|x| from_bytes_to_i128(x.as_ref()))), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.with_precision_and_scale(*precision, *scale)?.finish())) + }, + DataType::Decimal256(precision, scale) => { + let mut b = Decimal256Array::builder(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT32(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.map(|x| i256::from_i128(*x as i128))), + ); + } + ColumnIndexMetaData::INT64(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.map(|x| i256::from_i128(*x as i128))), + ); + } + ColumnIndexMetaData::BYTE_ARRAY(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.map(|x| from_bytes_to_i256(x.as_ref()))), + ); + } + ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.map(|x| from_bytes_to_i256(x.as_ref()))), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.with_precision_and_scale(*precision, *scale)?.finish())) + }, + DataType::Time32(unit) => { + match unit { + TimeUnit::Second => { + let mut b = Time32SecondBuilder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT32(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.copied()), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + } + TimeUnit::Millisecond => { + let mut b = Time32MillisecondBuilder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT32(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.copied()), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + } + _ => { + Ok(new_null_array($data_type, capacity)) + } + } + } + DataType::Time64(unit) => { + match unit { + TimeUnit::Microsecond => { + let mut b = Time64MicrosecondBuilder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT64(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.copied()), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + } + TimeUnit::Nanosecond => { + let mut b = Time64NanosecondBuilder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::INT64(index) => { + b.extend_from_iter_option( + index.$values_iter() + .map(|val| val.copied()), + ); + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + } + _ => { + Ok(new_null_array($data_type, capacity)) + } + } + }, + DataType::FixedSizeBinary(size) => { + let mut b = FixedSizeBinaryBuilder::with_capacity(capacity, *size); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(index) => { + for val in index.$values_iter() { + match val { + Some(v) => { + if v.len() == *size as usize { + let _ = b.append_value(v.as_ref())?; + } else { + b.append_null(); + } + } + None => b.append_null(), + } + } + } + _ => b.append_nulls(len), + } + } + Ok(Arc::new(b.finish())) + }, + DataType::Utf8View => { + let mut b = StringViewBuilder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::BYTE_ARRAY(index) => { + for val in index.$values_iter() { + match val { + Some(x) => match std::str::from_utf8(x.as_ref()) { + Ok(s) => b.append_value(s), + _ => b.append_null(), + } + None => b.append_null(), + } + } + } + _ => { + for _ in 0..len { b.append_null(); } + } + } + } + Ok(Arc::new(b.finish())) + }, + DataType::BinaryView => { + let mut b = BinaryViewBuilder::with_capacity(capacity); + for (len, index) in chunks { + match index { + ColumnIndexMetaData::BYTE_ARRAY(index) => { + for val in index.$values_iter() { + match val { + Some(v) => b.append_value(v.as_ref()), + None => b.append_null(), + } + } + } + _ => { + for _ in 0..len { b.append_null(); } + } + } + } + Ok(Arc::new(b.finish())) + }, + DataType::Date64 | // required to cover $physical_type match guard + DataType::Null | + DataType::Duration(_) | + DataType::Interval(_) | + DataType::List(_) | + DataType::ListView(_) | + DataType::FixedSizeList(_, _) | + DataType::LargeList(_) | + DataType::LargeListView(_) | + DataType::Struct(_) | + DataType::Union(_, _) | + DataType::Map(_, _) | + DataType::RunEndEncoded(_, _) => { + // don't know how to extract statistics, so return a null array + Ok(new_null_array($data_type, capacity)) + }, + } + } + }; +} +/// Extracts the min statistics from an iterator of [`ParquetStatistics`] to an +/// [`ArrayRef`] +/// +/// This is an internal helper -- see [`StatisticsConverter`] for public API +fn min_statistics<'a, I: Iterator>>( + data_type: &DataType, + iterator: I, + physical_type: Option, +) -> Result { + get_statistics!(Min, data_type, iterator, physical_type) +} + +/// Extracts the max statistics from an iterator of [`ParquetStatistics`] to an [`ArrayRef`] +/// +/// This is an internal helper -- see [`StatisticsConverter`] for public API +fn max_statistics<'a, I: Iterator>>( + data_type: &DataType, + iterator: I, + physical_type: Option, +) -> Result { + get_statistics!(Max, data_type, iterator, physical_type) +} + +/// Extracts the min statistics from an iterator +/// of parquet page [`ColumnIndexMetaData`]'s to an [`ArrayRef`] +pub(crate) fn min_page_statistics<'a, I>( + data_type: &DataType, + iterator: I, + physical_type: Option, +) -> Result +where + I: Iterator, +{ + get_data_page_statistics!(Min, data_type, iterator, physical_type) +} + +/// Extracts the max statistics from an iterator +/// of parquet page [`ColumnIndexMetaData`]'s to an [`ArrayRef`] +pub(crate) fn max_page_statistics<'a, I>( + data_type: &DataType, + iterator: I, + physical_type: Option, +) -> Result +where + I: Iterator, +{ + get_data_page_statistics!(Max, data_type, iterator, physical_type) +} + +/// Extracts the null count statistics from an iterator +/// of parquet page [`ColumnIndexMetaData`]'s to an [`ArrayRef`] +/// +/// The returned Array is an [`UInt64Array`] +pub(crate) fn null_counts_page_statistics<'a, I>(iterator: I) -> Result +where + I: Iterator, +{ + let chunks: Vec<_> = iterator.collect(); + let total_capacity: usize = chunks.iter().map(|(len, _)| *len).sum(); + let mut values = Vec::with_capacity(total_capacity); + let mut nulls = NullBufferBuilder::new(total_capacity); + for (len, index) in chunks { + match index.null_counts() { + Some(counts) => { + values.extend(counts.iter().map(|&x| x as u64)); + nulls.append_n_non_nulls(len); + } + None => { + values.resize(values.len() + len, 0); + nulls.append_n_nulls(len); + } + } + } + let null_buffer = nulls.build(); + let array = UInt64Array::new(values.into(), null_buffer); + Ok(array) +} + +/// Extracts Parquet statistics as Arrow arrays +/// +/// This is used to convert Parquet statistics to Arrow [`ArrayRef`], with +/// proper type conversions. This information can be used for pruning Parquet +/// files, row groups, and data pages based on the statistics embedded in +/// Parquet metadata. +/// +/// # Schemas +/// +/// The converter uses the schema of the Parquet file and the Arrow schema to +/// convert the underlying statistics value (stored as a parquet value) into the +/// corresponding Arrow value. For example, Decimals are stored as binary in +/// parquet files and this structure handles mapping them to the `i128` +/// representation used in Arrow. +/// +/// Note: The Parquet schema and Arrow schema do not have to be identical (for +/// example, the columns may be in different orders and one or the other schemas +/// may have additional columns). The function [`parquet_column`] is used to +/// match the column in the Parquet schema to the column in the Arrow schema +/// when using [`Self::try_new`]. For nested fields (e.g., struct fields), +/// where `parquet_column` does not support schema resolution, use +/// [`Self::from_column_index`] instead with a pre-resolved leaf column index. +#[derive(Debug)] +pub struct StatisticsConverter<'a> { + /// the index of the matched column in the Parquet schema + parquet_column_index: Option, + /// The field (with data type) of the column in the Arrow schema + arrow_field: &'a Field, + /// treat missing null_counts as 0 nulls + missing_null_counts_as_zero: bool, + /// The physical type of the matched column in the Parquet schema + physical_type: Option, +} + +impl<'a> StatisticsConverter<'a> { + /// Return the index of the column in the Parquet schema, if any + /// + /// Returns `None` if the column is was present in the Arrow schema, but not + /// present in the parquet file + pub fn parquet_column_index(&self) -> Option { + self.parquet_column_index + } + + /// Return the arrow schema's [`Field]` of the column in the Arrow schema + pub fn arrow_field(&self) -> &'a Field { + self.arrow_field + } + + /// Set the statistics converter to treat missing null counts as missing + /// + /// By default, the converter will treat missing null counts as though + /// the null count is known to be `0`. + /// + /// Note that parquet files written by parquet-rs currently do not store + /// null counts even when it is known there are zero nulls, and the reader + /// will return 0 for the null counts in that instance. This behavior may + /// change in a future release. + /// + /// Both parquet-java and parquet-cpp store null counts as 0 when there are + /// no nulls, and don't write unknown values to the null count field. + pub fn with_missing_null_counts_as_zero(mut self, missing_null_counts_as_zero: bool) -> Self { + self.missing_null_counts_as_zero = missing_null_counts_as_zero; + self + } + + /// Returns a [`UInt64Array`] with row counts for each row group + /// + /// # Return Value + /// + /// The returned array has no nulls, and has one value for each row group. + /// Each value is the number of rows in the row group. + /// + /// # Example + /// ```no_run + /// # use arrow::datatypes::Schema; + /// # use arrow_array::{ArrayRef, UInt64Array}; + /// # use parquet::arrow::arrow_reader::statistics::StatisticsConverter; + /// # use parquet::file::metadata::ParquetMetaData; + /// # fn get_parquet_metadata() -> ParquetMetaData { unimplemented!() } + /// # fn get_arrow_schema() -> Schema { unimplemented!() } + /// // Given the metadata for a parquet file and the arrow schema + /// let metadata: ParquetMetaData = get_parquet_metadata(); + /// let arrow_schema: Schema = get_arrow_schema(); + /// let parquet_schema = metadata.file_metadata().schema_descr(); + /// // create a converter + /// let converter = StatisticsConverter::try_new("foo", &arrow_schema, parquet_schema) + /// .unwrap(); + /// // get the row counts for each row group + /// let row_counts = converter.row_group_row_counts(metadata + /// .row_groups() + /// .iter() + /// ).unwrap(); + /// // file had 2 row groups, with 1024 and 23 rows respectively + /// assert_eq!(row_counts, Some(UInt64Array::from(vec![1024, 23]))); + /// ``` + pub fn row_group_row_counts(&self, metadatas: I) -> Result> + where + I: IntoIterator, + { + let Some(_) = self.parquet_column_index else { + return Ok(None); + }; + + let mut builder = UInt64Array::builder(10); + for metadata in metadatas.into_iter() { + let row_count = metadata.num_rows(); + let row_count: u64 = row_count.try_into().map_err(|e| { + arrow_err!(format!( + "Parquet row count {row_count} too large to convert to u64: {e}" + )) + })?; + builder.append_value(row_count); + } + Ok(Some(builder.finish())) + } + + /// Create a new `StatisticsConverter` to extract statistics for a column + /// + /// Note if there is no corresponding column in the parquet file, the returned + /// arrays will be null. This can happen if the column is in the arrow + /// schema but not in the parquet schema due to schema evolution. + /// + /// This constructor only supports top-level, non-nested columns. For nested + /// fields (e.g., fields within a struct), use [`Self::from_column_index`]. + /// + /// See example on [`Self::row_group_mins`] for usage + /// + /// # Errors + /// + /// * If the column is not found in the arrow schema + pub fn try_new<'b>( + column_name: &'b str, + arrow_schema: &'a Schema, + parquet_schema: &'a SchemaDescriptor, + ) -> Result { + // ensure the requested column is in the arrow schema + let Some((_idx, arrow_field)) = arrow_schema.column_with_name(column_name) else { + return Err(arrow_err!(format!( + "Column '{}' not found in schema for statistics conversion", + column_name + ))); + }; + + // find the column in the parquet schema, if not, return a null array + let parquet_index = match parquet_column(parquet_schema, arrow_schema, column_name) { + Some((parquet_idx, matched_field)) => { + // sanity check that matching field matches the arrow field + if matched_field.as_ref() != arrow_field { + return Err(arrow_err!(format!( + "Matched column '{:?}' does not match original matched column '{:?}'", + matched_field, arrow_field + ))); + } + Some(parquet_idx) + } + None => None, + }; + + Ok(Self { + parquet_column_index: parquet_index, + arrow_field, + missing_null_counts_as_zero: true, + physical_type: parquet_index.map(|idx| parquet_schema.column(idx).physical_type()), + }) + } + + /// Create a new `StatisticsConverter` from a Parquet leaf column index directly. + /// + /// Unlike [`Self::try_new`], this constructor bypasses schema resolution and + /// accepts a Parquet column index directly. This is useful for nested fields + /// (e.g., struct fields) where the caller has already resolved the mapping + /// from the Arrow field to the Parquet leaf column. + /// + /// # Arguments + /// + /// * `parquet_column_index` - The index of the leaf column in the Parquet schema + /// * `arrow_field` - The Arrow field describing the column's data type + /// * `parquet_schema` - The Parquet schema descriptor (used to look up the physical type) + /// + /// The caller must ensure that `arrow_field` describes the same leaf column as + /// `parquet_column_index`. This mapping is not validated by the converter; if + /// the Arrow type does not match the Parquet column statistics, extraction + /// returns null statistics values rather than an error. + /// + /// # Errors + /// + /// * If the `parquet_column_index` is out of bounds + pub fn from_column_index( + parquet_column_index: usize, + arrow_field: &'a Field, + parquet_schema: &'a SchemaDescriptor, + ) -> Result { + if parquet_column_index >= parquet_schema.columns().len() { + return Err(arrow_err!(format!( + "Parquet column index {} out of bounds, column count {}", + parquet_column_index, + parquet_schema.columns().len() + ))); + } + + let physical_type = parquet_schema.column(parquet_column_index).physical_type(); + + Ok(Self { + parquet_column_index: Some(parquet_column_index), + arrow_field, + missing_null_counts_as_zero: true, + physical_type: Some(physical_type), + }) + } + + /// Extract the minimum values from row group statistics in [`RowGroupMetaData`] + /// + /// # Return Value + /// + /// The returned array contains 1 value for each row group, in the same order as `metadatas` + /// + /// Each value is either + /// * the minimum value for the column + /// * a null value, if the statistics can not be extracted + /// + /// Note that a null value does NOT mean the min value was actually + /// `null` it means it the requested statistic is unknown + /// + /// # Errors + /// + /// Reasons for not being able to extract the statistics include: + /// * the column is not present in the parquet file + /// * statistics for the column are not present in the row group + /// * the stored statistic value can not be converted to the requested type + /// + /// # Example + /// ```no_run + /// # use std::sync::Arc; + /// # use arrow::datatypes::Schema; + /// # use arrow_array::{ArrayRef, Float64Array}; + /// # use parquet::arrow::arrow_reader::statistics::StatisticsConverter; + /// # use parquet::file::metadata::ParquetMetaData; + /// # fn get_parquet_metadata() -> ParquetMetaData { unimplemented!() } + /// # fn get_arrow_schema() -> Schema { unimplemented!() } + /// // Given the metadata for a parquet file and the arrow schema + /// let metadata: ParquetMetaData = get_parquet_metadata(); + /// let arrow_schema: Schema = get_arrow_schema(); + /// let parquet_schema = metadata.file_metadata().schema_descr(); + /// // create a converter + /// let converter = StatisticsConverter::try_new("foo", &arrow_schema, parquet_schema) + /// .unwrap(); + /// // get the minimum value for the column "foo" in the parquet file + /// let min_values: ArrayRef = converter + /// .row_group_mins(metadata.row_groups().iter()) + /// .unwrap(); + /// // if "foo" is a Float64 value, the returned array will contain Float64 values + /// assert_eq!(min_values, Arc::new(Float64Array::from(vec![Some(1.0), Some(2.0)])) as _); + /// ``` + pub fn row_group_mins(&self, metadatas: I) -> Result + where + I: IntoIterator, + { + let data_type = self.arrow_field.data_type(); + + let Some(parquet_index) = self.parquet_column_index else { + return Ok(self.make_null_array(data_type, metadatas)); + }; + + let iter = metadatas + .into_iter() + .map(|x| x.column(parquet_index).statistics()); + min_statistics(data_type, iter, self.physical_type) + } + + /// Extract the maximum values from row group statistics in [`RowGroupMetaData`] + /// + /// See docs on [`Self::row_group_mins`] for details + pub fn row_group_maxes(&self, metadatas: I) -> Result + where + I: IntoIterator, + { + let data_type = self.arrow_field.data_type(); + + let Some(parquet_index) = self.parquet_column_index else { + return Ok(self.make_null_array(data_type, metadatas)); + }; + + let iter = metadatas + .into_iter() + .map(|x| x.column(parquet_index).statistics()); + max_statistics(data_type, iter, self.physical_type) + } + + /// Extract the `is_max_value_exact` flags from row group statistics in [`RowGroupMetaData`] + /// + /// See docs on [`Self::row_group_maxes`] for details + pub fn row_group_is_max_value_exact(&self, metadatas: I) -> Result + where + I: IntoIterator, + { + let Some(parquet_index) = self.parquet_column_index else { + let num_row_groups = metadatas.into_iter().count(); + return Ok(BooleanArray::from_iter(std::iter::repeat_n( + None, + num_row_groups, + ))); + }; + + let is_max_value_exact = metadatas + .into_iter() + .map(|x| x.column(parquet_index).statistics()) + .map(|s| s.map(|s| s.max_is_exact())); + Ok(BooleanArray::from_iter(is_max_value_exact)) + } + + /// Extract the `is_min_value_exact` flags from row group statistics in [`RowGroupMetaData`] + /// + /// See docs on [`Self::row_group_mins`] for details + pub fn row_group_is_min_value_exact(&self, metadatas: I) -> Result + where + I: IntoIterator, + { + let Some(parquet_index) = self.parquet_column_index else { + let num_row_groups = metadatas.into_iter().count(); + return Ok(BooleanArray::from_iter(std::iter::repeat_n( + None, + num_row_groups, + ))); + }; + + let is_min_value_exact = metadatas + .into_iter() + .map(|x| x.column(parquet_index).statistics()) + .map(|s| s.map(|s| s.min_is_exact())); + Ok(BooleanArray::from_iter(is_min_value_exact)) + } + + /// Extract the null counts from row group statistics in [`RowGroupMetaData`] + /// + /// See docs on [`Self::row_group_mins`] for details + pub fn row_group_null_counts(&self, metadatas: I) -> Result + where + I: IntoIterator, + { + let Some(parquet_index) = self.parquet_column_index else { + let num_row_groups = metadatas.into_iter().count(); + return Ok(UInt64Array::from_iter(std::iter::repeat_n( + None, + num_row_groups, + ))); + }; + + let null_counts = metadatas + .into_iter() + .map(|x| x.column(parquet_index).statistics()) + .map(|s| { + s.and_then(|s| { + if self.missing_null_counts_as_zero { + Some(s.null_count_opt().unwrap_or(0)) + } else { + s.null_count_opt() + } + }) + }); + Ok(UInt64Array::from_iter(null_counts)) + } + + /// Extract the minimum values from Data Page statistics. + /// + /// In Parquet files, in addition to the Column Chunk level statistics + /// (stored for each column for each row group) there are also + /// optional statistics stored for each data page, as part of + /// the [`ParquetColumnIndex`]. + /// + /// Since a single Column Chunk is stored as one or more pages, + /// page level statistics can prune at a finer granularity. + /// + /// However since they are stored in a separate metadata + /// structure ([`ColumnIndexMetaData`]) there is different code to extract them as + /// compared to arrow statistics. + /// + /// # Parameters: + /// + /// * `column_page_index`: The parquet column page indices, read from + /// `ParquetMetaData` column_index + /// + /// * `column_offset_index`: The parquet column offset indices, read from + /// `ParquetMetaData` offset_index + /// + /// * `row_group_indices`: The indices of the row groups, that are used to + /// extract the column page index and offset index on a per row group + /// per column basis. + /// + /// # Return Value + /// + /// The returned array contains 1 value for each `NativeIndex` + /// in the underlying `Index`es, in the same order as they appear + /// in `metadatas`. + /// + /// For example, if there are two `Index`es in `metadatas`: + /// 1. the first having `3` `PageIndex` entries + /// 2. the second having `2` `PageIndex` entries + /// + /// The returned array would have 5 rows. + /// + /// Each value is either: + /// * the minimum value for the page + /// * a null value, if the statistics can not be extracted + /// + /// Note that a null value does NOT mean the min value was actually + /// `null` it means it the requested statistic is unknown + /// + /// # Errors + /// + /// Reasons for not being able to extract the statistics include: + /// * the column is not present in the parquet file + /// * statistics for the pages are not present in the row group + /// * the stored statistic value can not be converted to the requested type + pub fn data_page_mins( + &self, + column_page_index: &ParquetColumnIndex, + column_offset_index: &ParquetOffsetIndex, + row_group_indices: I, + ) -> Result + where + I: IntoIterator, + { + let data_type = self.arrow_field.data_type(); + + let Some(parquet_index) = self.parquet_column_index else { + return Ok(self.make_null_array(data_type, row_group_indices)); + }; + + let iter = row_group_indices.into_iter().map(|rg_index| { + let column_page_index_per_row_group_per_column = + &column_page_index[*rg_index][parquet_index]; + let num_data_pages = &column_offset_index[*rg_index][parquet_index] + .page_locations() + .len(); + + (*num_data_pages, column_page_index_per_row_group_per_column) + }); + + min_page_statistics(data_type, iter, self.physical_type) + } + + /// Extract the maximum values from Data Page statistics. + /// + /// See docs on [`Self::data_page_mins`] for details. + pub fn data_page_maxes( + &self, + column_page_index: &ParquetColumnIndex, + column_offset_index: &ParquetOffsetIndex, + row_group_indices: I, + ) -> Result + where + I: IntoIterator, + { + let data_type = self.arrow_field.data_type(); + + let Some(parquet_index) = self.parquet_column_index else { + return Ok(self.make_null_array(data_type, row_group_indices)); + }; + + let iter = row_group_indices.into_iter().map(|rg_index| { + let column_page_index_per_row_group_per_column = + &column_page_index[*rg_index][parquet_index]; + let num_data_pages = &column_offset_index[*rg_index][parquet_index] + .page_locations() + .len(); + + (*num_data_pages, column_page_index_per_row_group_per_column) + }); + + max_page_statistics(data_type, iter, self.physical_type) + } + + /// Returns a [`UInt64Array`] with null counts for each data page. + /// + /// See docs on [`Self::data_page_mins`] for details. + pub fn data_page_null_counts( + &self, + column_page_index: &ParquetColumnIndex, + column_offset_index: &ParquetOffsetIndex, + row_group_indices: I, + ) -> Result + where + I: IntoIterator, + { + let Some(parquet_index) = self.parquet_column_index else { + let num_row_groups = row_group_indices.into_iter().count(); + return Ok(UInt64Array::new_null(num_row_groups)); + }; + + let iter = row_group_indices.into_iter().map(|rg_index| { + let column_page_index_per_row_group_per_column = + &column_page_index[*rg_index][parquet_index]; + let num_data_pages = &column_offset_index[*rg_index][parquet_index] + .page_locations() + .len(); + + (*num_data_pages, column_page_index_per_row_group_per_column) + }); + null_counts_page_statistics(iter) + } + + /// Returns a [`UInt64Array`] with row counts for each data page. + /// + /// This function iterates over the given row group indexes and computes + /// the row count for each page in the specified column. + /// + /// # Parameters: + /// + /// * `column_offset_index`: The parquet column offset indices, read from + /// `ParquetMetaData` offset_index + /// + /// * `row_group_metadatas`: The metadata slice of the row groups, read + /// from `ParquetMetaData` row_groups + /// + /// * `row_group_indices`: The indices of the row groups, that are used to + /// extract the column offset index on a per row group per column basis. + /// + /// See docs on [`Self::data_page_mins`] for details. + pub fn data_page_row_counts( + &self, + column_offset_index: &ParquetOffsetIndex, + row_group_metadatas: &'a [RowGroupMetaData], + row_group_indices: I, + ) -> Result> + where + I: IntoIterator, + { + let Some(parquet_index) = self.parquet_column_index else { + // no matching column found in parquet_index; + // thus we cannot extract page_locations in order to determine + // the row count on a per DataPage basis. + return Ok(None); + }; + + let mut row_counts = Vec::new(); + let mut nulls = NullBufferBuilder::new(0); + for rg_idx in row_group_indices { + let page_locations = &column_offset_index[*rg_idx][parquet_index].page_locations(); + + let row_count_per_page = page_locations + .windows(2) + .map(|loc| Some(loc[1].first_row_index as u64 - loc[0].first_row_index as u64)); + + // append the last page row count + let num_rows_in_row_group = &row_group_metadatas[*rg_idx].num_rows(); + let row_count_per_page = row_count_per_page.chain(std::iter::once(Some( + *num_rows_in_row_group as u64 + - page_locations.last().unwrap().first_row_index as u64, + ))); + + row_counts.extend(row_count_per_page.clone().map(|x| x.unwrap_or(0))); + for val in row_count_per_page { + if val.is_some() { + nulls.append_non_null(); + } else { + nulls.append_null(); + } + } + } + + Ok(Some(UInt64Array::new(row_counts.into(), nulls.build()))) + } + + /// Returns a null array of data_type with one element per row group + fn make_null_array(&self, data_type: &DataType, metadatas: I) -> ArrayRef + where + I: IntoIterator, + { + // column was in the arrow schema but not in the parquet schema, so return a null array + let num_row_groups = metadatas.into_iter().count(); + new_null_array(data_type, num_row_groups) + } +} + +// See tests in parquet/tests/arrow_reader/statistics.rs diff --git a/vendor/parquet/src/arrow/arrow_writer/byte_array.rs b/vendor/parquet/src/arrow/arrow_writer/byte_array.rs new file mode 100644 index 0000000..6160c67 --- /dev/null +++ b/vendor/parquet/src/arrow/arrow_writer/byte_array.rs @@ -0,0 +1,822 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::basic::Encoding; +use crate::bloom_filter::Sbbf; +use crate::column::writer::encoder::{ + ColumnValueEncoder, DataPageValues, DictionaryPage, create_bloom_filter, +}; +use crate::data_type::{AsBytes, ByteArray, Int32Type}; +use crate::encodings::encoding::{DeltaBitPackEncoder, Encoder}; +use crate::encodings::rle::RleEncoder; +use crate::errors::{ParquetError, Result}; +use crate::file::properties::{EnabledStatistics, WriterProperties, WriterVersion}; +use crate::geospatial::accumulator::{GeoStatsAccumulator, try_new_geo_stats_accumulator}; +use crate::geospatial::statistics::GeospatialStatistics; +use crate::schema::types::ColumnDescPtr; +use crate::util::bit_util::num_required_bits; +use crate::util::interner::{Interner, Storage}; +use arrow_array::types::ByteArrayType; +use arrow_array::{ + Array, ArrayAccessor, BinaryArray, BinaryViewArray, DictionaryArray, FixedSizeBinaryArray, + GenericByteArray, LargeBinaryArray, LargeStringArray, StringArray, StringViewArray, +}; +use arrow_buffer::{ArrowNativeType, Buffer}; +use arrow_schema::DataType; + +macro_rules! downcast_dict_impl { + ($array:ident, $key:ident, $val:ident, $op:expr $(, $arg:expr)*) => {{ + $op($array + .as_any() + .downcast_ref::>() + .unwrap() + .downcast_dict::<$val>() + .unwrap()$(, $arg)*) + }}; +} + +macro_rules! downcast_dict_op { + ($key_type:expr, $val:ident, $array:ident, $op:expr $(, $arg:expr)*) => { + match $key_type.as_ref() { + DataType::UInt8 => downcast_dict_impl!($array, UInt8Type, $val, $op$(, $arg)*), + DataType::UInt16 => downcast_dict_impl!($array, UInt16Type, $val, $op$(, $arg)*), + DataType::UInt32 => downcast_dict_impl!($array, UInt32Type, $val, $op$(, $arg)*), + DataType::UInt64 => downcast_dict_impl!($array, UInt64Type, $val, $op$(, $arg)*), + DataType::Int8 => downcast_dict_impl!($array, Int8Type, $val, $op$(, $arg)*), + DataType::Int16 => downcast_dict_impl!($array, Int16Type, $val, $op$(, $arg)*), + DataType::Int32 => downcast_dict_impl!($array, Int32Type, $val, $op$(, $arg)*), + DataType::Int64 => downcast_dict_impl!($array, Int64Type, $val, $op$(, $arg)*), + _ => unreachable!(), + } + }; +} + +macro_rules! downcast_op { + ($data_type:expr, $array:ident, $op:expr $(, $arg:expr)*) => { + match $data_type { + DataType::Utf8 => $op($array.as_any().downcast_ref::().unwrap()$(, $arg)*), + DataType::LargeUtf8 => { + $op($array.as_any().downcast_ref::().unwrap()$(, $arg)*) + } + DataType::Utf8View => $op($array.as_any().downcast_ref::().unwrap()$(, $arg)*), + DataType::Binary => { + $op($array.as_any().downcast_ref::().unwrap()$(, $arg)*) + } + DataType::LargeBinary => { + $op($array.as_any().downcast_ref::().unwrap()$(, $arg)*) + } + DataType::BinaryView => { + $op($array.as_any().downcast_ref::().unwrap()$(, $arg)*) + } + DataType::Dictionary(key, value) => match value.as_ref() { + DataType::Utf8 => downcast_dict_op!(key, StringArray, $array, $op$(, $arg)*), + DataType::LargeUtf8 => { + downcast_dict_op!(key, LargeStringArray, $array, $op$(, $arg)*) + } + DataType::Binary => downcast_dict_op!(key, BinaryArray, $array, $op$(, $arg)*), + DataType::LargeBinary => { + downcast_dict_op!(key, LargeBinaryArray, $array, $op$(, $arg)*) + } + DataType::FixedSizeBinary(_) => { + downcast_dict_op!(key, FixedSizeBinaryArray, $array, $op$(, $arg)*) + } + d => unreachable!("cannot downcast {} dictionary value to byte array", d), + }, + d => unreachable!("cannot downcast {} to byte array", d), + } + }; +} + +/// A fallback encoder, i.e. non-dictionary, for [`ByteArray`] +struct FallbackEncoder { + encoder: FallbackEncoderImpl, + num_values: usize, + variable_length_bytes: i64, +} + +/// The fallback encoder in use +/// +/// Note: DeltaBitPackEncoder is boxed as it is rather large +enum FallbackEncoderImpl { + Plain { + buffer: Vec, + }, + DeltaLength { + buffer: Vec, + lengths: Box>, + }, + Delta { + buffer: Vec, + last_value: Vec, + prefix_lengths: Box>, + suffix_lengths: Box>, + }, +} + +impl FallbackEncoder { + /// Create the fallback encoder for the given [`ColumnDescPtr`] and [`WriterProperties`] + fn new(descr: &ColumnDescPtr, props: &WriterProperties) -> Result { + // Set either main encoder or fallback encoder. + let encoding = + props + .encoding(descr.path()) + .unwrap_or_else(|| match props.writer_version() { + WriterVersion::PARQUET_1_0 => Encoding::PLAIN, + WriterVersion::PARQUET_2_0 => Encoding::DELTA_BYTE_ARRAY, + }); + + let encoder = match encoding { + Encoding::PLAIN => FallbackEncoderImpl::Plain { buffer: vec![] }, + Encoding::DELTA_LENGTH_BYTE_ARRAY => FallbackEncoderImpl::DeltaLength { + buffer: vec![], + lengths: Box::new(DeltaBitPackEncoder::new()), + }, + Encoding::DELTA_BYTE_ARRAY => FallbackEncoderImpl::Delta { + buffer: vec![], + last_value: vec![], + prefix_lengths: Box::new(DeltaBitPackEncoder::new()), + suffix_lengths: Box::new(DeltaBitPackEncoder::new()), + }, + _ => { + return Err(general_err!( + "unsupported encoding {} for byte array", + encoding + )); + } + }; + + Ok(Self { + encoder, + num_values: 0, + variable_length_bytes: 0, + }) + } + + /// Encode `values` to the in-progress page + fn encode(&mut self, values: T, indices: impl ExactSizeIterator) + where + T: ArrayAccessor + Copy, + T::Item: AsRef<[u8]>, + { + self.num_values += indices.len(); + match &mut self.encoder { + FallbackEncoderImpl::Plain { buffer } => { + for idx in indices { + let value = values.value(idx); + let value = value.as_ref(); + buffer.extend_from_slice((value.len() as u32).as_bytes()); + buffer.extend_from_slice(value); + self.variable_length_bytes += value.len() as i64; + } + } + FallbackEncoderImpl::DeltaLength { buffer, lengths } => { + for idx in indices { + let value = values.value(idx); + let value = value.as_ref(); + lengths.put(&[value.len() as i32]).unwrap(); + buffer.extend_from_slice(value); + self.variable_length_bytes += value.len() as i64; + } + } + FallbackEncoderImpl::Delta { + buffer, + last_value, + prefix_lengths, + suffix_lengths, + } => { + for idx in indices { + let value = values.value(idx); + let value = value.as_ref(); + let mut prefix_length = 0; + + while prefix_length < last_value.len() + && prefix_length < value.len() + && last_value[prefix_length] == value[prefix_length] + { + prefix_length += 1; + } + + let suffix_length = value.len() - prefix_length; + + last_value.clear(); + last_value.extend_from_slice(value); + + buffer.extend_from_slice(&value[prefix_length..]); + prefix_lengths.put(&[prefix_length as i32]).unwrap(); + suffix_lengths.put(&[suffix_length as i32]).unwrap(); + self.variable_length_bytes += value.len() as i64; + } + } + } + } + + /// Returns an estimate of the data page size in bytes + /// + /// This includes: + /// + + fn estimated_data_page_size(&self) -> usize { + match &self.encoder { + FallbackEncoderImpl::Plain { buffer, .. } => buffer.len(), + FallbackEncoderImpl::DeltaLength { buffer, lengths } => { + buffer.len() + lengths.estimated_data_encoded_size() + } + FallbackEncoderImpl::Delta { + buffer, + prefix_lengths, + suffix_lengths, + .. + } => { + buffer.len() + + prefix_lengths.estimated_data_encoded_size() + + suffix_lengths.estimated_data_encoded_size() + } + } + } + + fn flush_data_page( + &mut self, + min_value: Option, + max_value: Option, + ) -> Result> { + let (buf, encoding) = match &mut self.encoder { + FallbackEncoderImpl::Plain { buffer } => (std::mem::take(buffer), Encoding::PLAIN), + FallbackEncoderImpl::DeltaLength { buffer, lengths } => { + let lengths = lengths.flush_buffer()?; + + let mut out = Vec::with_capacity(lengths.len() + buffer.len()); + out.extend_from_slice(&lengths); + out.extend_from_slice(buffer); + buffer.clear(); + (out, Encoding::DELTA_LENGTH_BYTE_ARRAY) + } + FallbackEncoderImpl::Delta { + buffer, + prefix_lengths, + suffix_lengths, + last_value, + } => { + let prefix_lengths = prefix_lengths.flush_buffer()?; + let suffix_lengths = suffix_lengths.flush_buffer()?; + + let mut out = + Vec::with_capacity(prefix_lengths.len() + suffix_lengths.len() + buffer.len()); + out.extend_from_slice(&prefix_lengths); + out.extend_from_slice(&suffix_lengths); + out.extend_from_slice(buffer); + buffer.clear(); + last_value.clear(); + (out, Encoding::DELTA_BYTE_ARRAY) + } + }; + + // Capture value of variable_length_bytes and reset for next page + let variable_length_bytes = Some(self.variable_length_bytes); + self.variable_length_bytes = 0; + + Ok(DataPageValues { + buf: buf.into(), + num_values: std::mem::take(&mut self.num_values), + encoding, + min_value, + max_value, + variable_length_bytes, + }) + } +} + +/// [`Storage`] for the [`Interner`] used by [`DictEncoder`] +#[derive(Debug, Default)] +struct ByteArrayStorage { + /// Encoded dictionary data + page: Vec, + + values: Vec>, +} + +impl Storage for ByteArrayStorage { + type Key = u64; + type Value = [u8]; + + fn get(&self, idx: Self::Key) -> &Self::Value { + &self.page[self.values[idx as usize].clone()] + } + + fn push(&mut self, value: &Self::Value) -> Self::Key { + let key = self.values.len(); + + self.page.reserve(4 + value.len()); + self.page.extend_from_slice((value.len() as u32).as_bytes()); + + let start = self.page.len(); + self.page.extend_from_slice(value); + self.values.push(start..self.page.len()); + + key as u64 + } + + #[allow(dead_code)] // not used in parquet_derive, so is dead there + fn estimated_memory_size(&self) -> usize { + self.page.capacity() * std::mem::size_of::() + + self.values.capacity() * std::mem::size_of::>() + } +} + +/// A dictionary encoder for byte array data +#[derive(Debug, Default)] +struct DictEncoder { + interner: Interner, + indices: Vec, + variable_length_bytes: i64, +} + +impl DictEncoder { + /// Encode `values` to the in-progress page + fn encode(&mut self, values: T, indices: impl ExactSizeIterator) + where + T: ArrayAccessor + Copy, + T::Item: AsRef<[u8]>, + { + self.indices.reserve(indices.len()); + + for idx in indices { + let value = values.value(idx); + let interned = self.interner.intern(value.as_ref()); + self.indices.push(interned); + self.variable_length_bytes += value.as_ref().len() as i64; + } + } + + fn bit_width(&self) -> u8 { + let length = self.interner.storage().values.len(); + num_required_bits(length.saturating_sub(1) as u64) + } + + fn estimated_memory_size(&self) -> usize { + self.interner.estimated_memory_size() + self.indices.capacity() * std::mem::size_of::() + } + + fn estimated_data_page_size(&self) -> usize { + let bit_width = self.bit_width(); + 1 + RleEncoder::max_buffer_size(bit_width, self.indices.len()) + } + + fn estimated_dict_page_size(&self) -> usize { + self.interner.storage().page.len() + } + + fn flush_dict_page(self) -> DictionaryPage { + let storage = self.interner.into_inner(); + + DictionaryPage { + buf: storage.page.into(), + num_values: storage.values.len(), + is_sorted: false, + } + } + + fn flush_data_page( + &mut self, + min_value: Option, + max_value: Option, + ) -> DataPageValues { + let num_values = self.indices.len(); + let buffer_len = self.estimated_data_page_size(); + let mut buffer = Vec::with_capacity(buffer_len); + buffer.push(self.bit_width()); + + let mut encoder = RleEncoder::new_from_buf(self.bit_width(), buffer); + for index in &self.indices { + encoder.put(*index) + } + + self.indices.clear(); + + // Capture value of variable_length_bytes and reset for next page + let variable_length_bytes = Some(self.variable_length_bytes); + self.variable_length_bytes = 0; + + DataPageValues { + buf: encoder.consume().into(), + num_values, + encoding: Encoding::RLE_DICTIONARY, + min_value, + max_value, + variable_length_bytes, + } + } +} + +pub struct ByteArrayEncoder { + fallback: FallbackEncoder, + dict_encoder: Option, + statistics_enabled: EnabledStatistics, + min_value: Option, + max_value: Option, + bloom_filter: Option, + bloom_filter_target_fpp: f64, + geo_stats_accumulator: Option>, +} + +impl ColumnValueEncoder for ByteArrayEncoder { + type T = ByteArray; + type Values = dyn Array; + fn flush_bloom_filter(&mut self) -> Option { + let mut sbbf = self.bloom_filter.take()?; + sbbf.fold_to_target_fpp(self.bloom_filter_target_fpp); + Some(sbbf) + } + + fn try_new(descr: &ColumnDescPtr, props: &WriterProperties) -> Result + where + Self: Sized, + { + let dictionary = props + .dictionary_enabled(descr.path()) + .then(DictEncoder::default); + + let fallback = FallbackEncoder::new(descr, props)?; + + let (bloom_filter, bloom_filter_target_fpp) = create_bloom_filter(props, descr)?; + + let statistics_enabled = props.statistics_enabled(descr.path()); + + let geo_stats_accumulator = try_new_geo_stats_accumulator(descr); + + Ok(Self { + fallback, + statistics_enabled, + bloom_filter, + bloom_filter_target_fpp, + dict_encoder: dictionary, + min_value: None, + max_value: None, + geo_stats_accumulator, + }) + } + + fn write(&mut self, _values: &Self::Values, _offset: usize, _len: usize) -> Result<()> { + unreachable!("should call write_gather instead") + } + + fn write_gather(&mut self, values: &Self::Values, indices: &[usize]) -> Result<()> { + downcast_op!( + values.data_type(), + values, + encode, + indices.iter().copied(), + self + ); + Ok(()) + } + + fn count_values_within_byte_budget_gather( + values: &Self::Values, + indices: &[usize], + byte_budget: usize, + ) -> Option { + // `ByteArrayEncoder` only ever writes via `write_gather`, so this + // is the relevant method. + // + // Two-stage walk for the simple offset-buffer byte array types: + // 1. If indices are contiguous, compute the total payload in + // O(1) via a single subtraction on the offsets buffer. + // When the total fits the budget — the overwhelmingly + // common "small values" case — return immediately. + // 2. Otherwise, walk per-value byte sizes from the offsets + // buffer (still cheap, no slice/UTF-8 construction) and + // exit at the first value that pushes the cumulative sum + // past the budget. This bounds skewed distributions: an + // outlier value is caught wherever it lands in the chunk. + let count = match values.data_type() { + DataType::Utf8 => count_within_budget_offsets( + values.as_any().downcast_ref::().unwrap(), + indices, + byte_budget, + ), + DataType::LargeUtf8 => count_within_budget_offsets( + values.as_any().downcast_ref::().unwrap(), + indices, + byte_budget, + ), + DataType::Binary => count_within_budget_offsets( + values.as_any().downcast_ref::().unwrap(), + indices, + byte_budget, + ), + DataType::LargeBinary => count_within_budget_offsets( + values.as_any().downcast_ref::().unwrap(), + indices, + byte_budget, + ), + // View arrays carry each value's length in the low 32 bits of + // its u128 view word, so lengths are scannable without touching + // any data buffer — and the common small-value case skips even + // that scan via an O(1) conservative bound. + DataType::Utf8View => { + let array = values.as_any().downcast_ref::().unwrap(); + count_within_budget_views( + array.views(), + indices, + byte_budget, + max_view_value_len(array.data_buffers()), + ) + } + DataType::BinaryView => { + let array = values.as_any().downcast_ref::().unwrap(); + count_within_budget_views( + array.views(), + indices, + byte_budget, + max_view_value_len(array.data_buffers()), + ) + } + // The values in an arrow dictionary are already small and + // deduplicated, so there is nothing to bound — treat every + // chunk as fitting and stay on the batched path. (A per-value + // walk through dict keys on every chunk also measured ~+30-80% + // slower than `main`.) + DataType::Dictionary(_, _) => indices.len(), + // Every byte-array type `ByteArrayEncoder` is constructed for + // has an explicit arm above. A `Dictionary(value = FixedSizeBinary)` + // column hits the `Dictionary(_, _)` arm (its `values.data_type()` + // is `Dictionary`), and a bare `FixedSizeBinary` column is routed + // to the generic column writer, never this encoder — so no other + // type can reach here. + data_type => unreachable!("ByteArrayEncoder cannot be constructed for {data_type:?}"), + }; + Some(count) + } + + fn num_values(&self) -> usize { + match &self.dict_encoder { + Some(encoder) => encoder.indices.len(), + None => self.fallback.num_values, + } + } + + fn has_dictionary(&self) -> bool { + self.dict_encoder.is_some() + } + + fn compresses_against_previous_value(&self) -> bool { + // While dictionary encoding is active the data page holds RLE + // indices, which carry no cross-value state; only the DELTA_BYTE_ARRAY + // fallback shares prefixes with the preceding value. + self.dict_encoder.is_none() + && matches!(self.fallback.encoder, FallbackEncoderImpl::Delta { .. }) + } + + fn estimated_memory_size(&self) -> usize { + let encoder_size = match &self.dict_encoder { + Some(encoder) => encoder.estimated_memory_size(), + // For the FallbackEncoder, these unflushed bytes are already encoded. + // Therefore, the size should be the same as estimated_data_page_size. + None => self.fallback.estimated_data_page_size(), + }; + + let bloom_filter_size = self + .bloom_filter + .as_ref() + .map(|bf| bf.estimated_memory_size()) + .unwrap_or_default(); + + let stats_size = self.min_value.as_ref().map(|v| v.len()).unwrap_or_default() + + self.max_value.as_ref().map(|v| v.len()).unwrap_or_default(); + + encoder_size + bloom_filter_size + stats_size + } + + fn estimated_dict_page_size(&self) -> Option { + Some(self.dict_encoder.as_ref()?.estimated_dict_page_size()) + } + + /// Returns an estimate of the data page size in bytes + /// + /// This includes: + /// + + fn estimated_data_page_size(&self) -> usize { + match &self.dict_encoder { + Some(encoder) => encoder.estimated_data_page_size(), + None => self.fallback.estimated_data_page_size(), + } + } + + fn flush_dict_page(&mut self) -> Result> { + match self.dict_encoder.take() { + Some(encoder) => { + if !encoder.indices.is_empty() { + return Err(general_err!( + "Must flush data pages before flushing dictionary" + )); + } + + Ok(Some(encoder.flush_dict_page())) + } + _ => Ok(None), + } + } + + fn flush_data_page(&mut self) -> Result> { + let min_value = self.min_value.take(); + let max_value = self.max_value.take(); + + match &mut self.dict_encoder { + Some(encoder) => Ok(encoder.flush_data_page(min_value, max_value)), + _ => self.fallback.flush_data_page(min_value, max_value), + } + } + + fn flush_geospatial_statistics(&mut self) -> Option> { + self.geo_stats_accumulator.as_mut().map(|a| a.finish())? + } +} + +/// Encodes the provided `values` and `indices` to `encoder` +/// +/// This is a free function so it can be used with `downcast_op!` +fn encode(values: T, indices: I, encoder: &mut ByteArrayEncoder) +where + T: ArrayAccessor + Copy, + T::Item: Copy + Ord + AsRef<[u8]>, + I: ExactSizeIterator + Clone, +{ + if encoder.statistics_enabled != EnabledStatistics::None { + if let Some(accumulator) = encoder.geo_stats_accumulator.as_mut() { + update_geo_stats_accumulator(accumulator.as_mut(), values, indices.clone()); + } else if let Some((min, max)) = compute_min_max(values, indices.clone()) { + if encoder.min_value.as_ref().is_none_or(|m| m > &min) { + encoder.min_value = Some(min); + } + + if encoder.max_value.as_ref().is_none_or(|m| m < &max) { + encoder.max_value = Some(max); + } + } + } + + // encode the values into bloom filter if enabled + if let Some(bloom_filter) = &mut encoder.bloom_filter { + for idx in indices.clone() { + bloom_filter.insert(values.value(idx).as_ref()); + } + } + + match &mut encoder.dict_encoder { + Some(dict_encoder) => dict_encoder.encode(values, indices), + None => encoder.fallback.encode(values, indices), + } +} + +/// Upper bound on any single value's byte length in a view array. +fn max_view_value_len(buffers: &[Buffer]) -> usize { + /// Bytes that fit inline in a u128 view word (the rest is len + prefix). + const MAX_INLINE_VIEW_LEN: usize = 12; + // An out-of-line view's data is a contiguous slice of exactly one data + // buffer, so it cannot exceed the largest buffer; inline views hold at + // most `MAX_INLINE_VIEW_LEN`. Loose (a value is usually far smaller than + // a whole buffer) but O(number of buffers) and always sound. + buffers + .iter() + .map(|b| b.len()) + .max() + .unwrap_or(0) + .max(MAX_INLINE_VIEW_LEN) +} + +/// Number of leading `indices` whose cumulative plain-encoded size fits +/// `byte_budget` (boundary value included), for view arrays (`Utf8View`, +/// `BinaryView`). +fn count_within_budget_views( + views: &[u128], + indices: &[usize], + byte_budget: usize, + max_value_len: usize, +) -> usize { + // Each plain-encoded BYTE_ARRAY value carries a 4-byte length prefix, so + // the budget is compared against `value_len + size_of::()` — the + // bytes actually written to the page, not just the payload. + // + // Stage 1: O(1) conservative bound. View arrays have no prefix-sum + // offsets buffer, so the exact span subtraction used by + // `count_within_budget_offsets` is unavailable; instead bound every + // value by `max_value_len`. Skips the walk for the common small-value + // case (what view arrays are built for, and where there is nothing to + // bound). + let per_value = max_value_len + std::mem::size_of::(); + if indices.len().saturating_mul(per_value) <= byte_budget { + return indices.len(); + } + // Stage 2: exact per-value scan, reading each length from the low 32 + // bits of its u128 view word (no data-buffer dereference). + let mut cum: usize = 0; + for (i, idx) in indices.iter().enumerate() { + let len = (views[*idx] as u32) as usize; + cum = cum.saturating_add(len + std::mem::size_of::()); + if cum > byte_budget { + return i + 1; + } + } + indices.len() +} + +/// Number of leading `indices` whose cumulative plain-encoded size fits +/// `byte_budget` (boundary value included), for offset-buffer byte arrays +/// (`Utf8`/`LargeUtf8`/`Binary`/`LargeBinary`). +/// +/// `indices` are assumed sorted ascending — they always are here, since +/// they come from `non_null_indices`, which is built in array order. +fn count_within_budget_offsets( + values: &GenericByteArray, + indices: &[usize], + byte_budget: usize, +) -> usize { + if indices.is_empty() { + return 0; + } + let n = indices.len(); + let first = indices[0]; + let last = indices[n - 1]; + let offsets = values.value_offsets(); + // Each plain-encoded value carries a 4-byte length prefix on the page. + let prefix_overhead = std::mem::size_of::(); + + // Stage 1: O(1) span upper bound. The span `offsets[last+1] - + // offsets[first]` covers every array position in `[first, last]`, a + // superset of `indices` — and the skipped positions in a nullable + // column are nulls with zero offset delta, so the span still equals the + // exact payload. If it fits the budget, every value fits. Covers the + // common small-value case for both non-null and (sparse) nullable + // columns. + if last >= first { + let payload = (offsets[last + 1] - offsets[first]).as_usize(); + if payload + n * prefix_overhead <= byte_budget { + return n; + } + } + + // Stage 2: scan per-index lengths from the offsets buffer. + let mut cum: usize = 0; + for (i, idx) in indices.iter().enumerate() { + let len = (offsets[idx + 1] - offsets[*idx]).as_usize() + prefix_overhead; + cum = cum.saturating_add(len); + if cum > byte_budget { + return i + 1; + } + } + n +} + +/// Computes the min and max for the provided array and indices +/// +/// This is a free function so it can be used with `downcast_op!` +fn compute_min_max( + array: T, + mut valid: impl Iterator, +) -> Option<(ByteArray, ByteArray)> +where + T: ArrayAccessor, + T::Item: Copy + Ord + AsRef<[u8]>, +{ + let first_idx = valid.next()?; + + let first_val = array.value(first_idx); + let mut min = first_val; + let mut max = first_val; + for idx in valid { + let val = array.value(idx); + min = min.min(val); + max = max.max(val); + } + Some((min.as_ref().to_vec().into(), max.as_ref().to_vec().into())) +} + +/// Updates geospatial statistics for the provided array and indices +fn update_geo_stats_accumulator( + bounder: &mut dyn GeoStatsAccumulator, + array: T, + valid: impl Iterator, +) where + T: ArrayAccessor, + T::Item: Copy + Ord + AsRef<[u8]>, +{ + if bounder.is_valid() { + for idx in valid { + let val = array.value(idx); + bounder.update_wkb(val.as_ref()); + } + } +} diff --git a/vendor/parquet/src/arrow/arrow_writer/levels.rs b/vendor/parquet/src/arrow/arrow_writer/levels.rs new file mode 100644 index 0000000..2781898 --- /dev/null +++ b/vendor/parquet/src/arrow/arrow_writer/levels.rs @@ -0,0 +1,2848 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Parquet definition and repetition levels +//! +//! Contains the algorithm for computing definition and repetition levels. +//! The algorithm works by tracking the slots of an array that should +//! ultimately be populated when writing to Parquet. +//! Parquet achieves nesting through definition levels and repetition levels \[1\]. +//! Definition levels specify how many optional fields in the part for the column +//! are defined. +//! Repetition levels specify at what repeated field (list) in the path a column +//! is defined. +//! +//! In a nested data structure such as `a.b.c`, one can see levels as defining +//! whether a record is defined at `a`, `a.b`, or `a.b.c`. +//! Optional fields are nullable fields, thus if all 3 fields +//! are nullable, the maximum definition could be = 3 if there are no lists. +//! +//! The algorithm in this module computes the necessary information to enable +//! the writer to keep track of which columns are at which levels, and to extract +//! the correct values at the correct slots from Arrow arrays. +//! +//! It works by walking a record batch's arrays, keeping track of what values +//! are non-null, their positions and computing what their levels are. +//! +//! \[1\] [parquet-format#nested-encoding](https://github.com/apache/parquet-format#nested-encoding) + +use crate::column::chunker::CdcChunk; +use crate::column::writer::LevelDataRef; +use crate::errors::{ParquetError, Result}; +use arrow_array::cast::AsArray; +use arrow_array::types::RunEndIndexType; +use arrow_array::{Array, ArrayRef, Int32Array, OffsetSizeTrait, RunArray, downcast_run_array}; +use arrow_buffer::bit_iterator::BitIndexIterator; +use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; +use arrow_schema::{DataType, Field}; +use std::ops::Range; +use std::sync::Arc; + +/// Expands a [`DataType::RunEndEncoded`] array into a flat (logical) array of its values type. +/// +/// use `arrow_select::take` to materialize the full-length flat array. +/// This is intentionally simple (O(n)); efficiency can/should be improved +fn expand_ree_array(array: &ArrayRef) -> Result { + downcast_run_array!( + array => expand_typed_ree(array), + _ => unreachable!("expand_ree_array called on non-REE array"), + ) +} + +fn expand_typed_ree(run_array: &RunArray) -> Result { + let run_ends = run_array.run_ends(); + let values = run_array.values(); + let len = run_array.len(); + let indices: Int32Array = (0..len) + .map(|i| run_ends.get_physical_index(i) as i32) + .collect(); + arrow_select::take::take(values.as_ref(), &indices, None) + .map_err(|e| arrow_err!("Failed to expand REE array: {}", e)) +} + +/// Performs a depth-first scan of the children of `array`, constructing [`ArrayLevels`] +/// for each leaf column encountered +pub(crate) fn calculate_array_levels(array: &ArrayRef, field: &Field) -> Result> { + let mut builder = LevelInfoBuilder::try_new(field, Default::default(), array)?; + builder.write(0..array.len()); + Ok(builder.finish()) +} + +/// Returns true if the DataType can be represented as a primitive parquet column, +/// i.e. a leaf array with no children +fn is_leaf(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Null + | DataType::Boolean + | DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + | DataType::Float16 + | DataType::Float32 + | DataType::Float64 + | DataType::Utf8 + | DataType::Utf8View + | DataType::LargeUtf8 + | DataType::Timestamp(_, _) + | DataType::Date32 + | DataType::Date64 + | DataType::Time32(_) + | DataType::Time64(_) + | DataType::Duration(_) + | DataType::Interval(_) + | DataType::Binary + | DataType::LargeBinary + | DataType::BinaryView + | DataType::Decimal32(_, _) + | DataType::Decimal64(_, _) + | DataType::Decimal128(_, _) + | DataType::Decimal256(_, _) + | DataType::FixedSizeBinary(_) + ) +} + +/// The definition and repetition level of an array within a potentially nested hierarchy +#[derive(Debug, Default, Clone, Copy)] +struct LevelContext { + /// The current repetition level + rep_level: i16, + /// The current definition level + def_level: i16, +} + +/// A helper to construct [`ArrayLevels`] from a potentially nested [`Field`] +#[derive(Debug)] +enum LevelInfoBuilder { + /// A primitive, leaf array + Primitive(ArrayLevels), + /// A list array + List( + Box, // Child Values + LevelContext, // Context + OffsetBuffer, // Offsets + Option, // Nulls + bool, // is_last_level (child has no nested rep) + ), + /// A large list array + LargeList( + Box, // Child Values + LevelContext, // Context + OffsetBuffer, // Offsets + Option, // Nulls + bool, // is_last_level (child has no nested rep) + ), + /// A fixed size list array + FixedSizeList( + Box, // Values + LevelContext, // Context + usize, // List Size + Option, // Nulls + ), + /// A list view array + ListView( + Box, // Child Values + LevelContext, // Context + ScalarBuffer, // Offsets + ScalarBuffer, // Sizes + Option, // Nulls + ), + /// A large list view array + LargeListView( + Box, // Child Values + LevelContext, // Context + ScalarBuffer, // Offsets + ScalarBuffer, // Sizes + Option, // Nulls + ), + /// A struct array + Struct(Vec, LevelContext, Option), +} + +/// Minimum sub-range length before the bulk-fill fast path in `write_leaf` +/// becomes profitable for null-heavy leaf columns. Below this, per-call +/// slice + popcount overhead regresses list/struct paths that call +/// `write_leaf` many times with tiny ranges. Picked via threshold sweep; +/// see for the rationale. +const BULK_FILL_MIN_LEN: usize = 64; + +impl LevelInfoBuilder { + /// Create a new [`LevelInfoBuilder`] for the given [`Field`] and parent [`LevelContext`] + fn try_new(field: &Field, parent_ctx: LevelContext, array: &ArrayRef) -> Result { + if !Self::types_compatible(field.data_type(), array.data_type()) { + return Err(arrow_err!(format!( + "Incompatible type. Field '{}' has type {}, array has type {}", + field.name(), + field.data_type(), + array.data_type(), + ))); + } + + let is_nullable = field.is_nullable(); + + match array.data_type() { + d if is_leaf(d) => { + let levels = ArrayLevels::new(parent_ctx, is_nullable, array.clone()); + Ok(Self::Primitive(levels)) + } + DataType::Dictionary(_, v) if is_leaf(v.as_ref()) => { + let levels = ArrayLevels::new(parent_ctx, is_nullable, array.clone()); + Ok(Self::Primitive(levels)) + } + DataType::RunEndEncoded(_, value_field) => { + let flat = expand_ree_array(array)?; + let flat_field = Field::new( + field.name(), + value_field.data_type().clone(), + field.is_nullable(), + ); + Self::try_new(&flat_field, parent_ctx, &flat) + } + DataType::Struct(children) => { + let array = array.as_struct(); + let def_level = match is_nullable { + true => parent_ctx.def_level + 1, + false => parent_ctx.def_level, + }; + + let ctx = LevelContext { + rep_level: parent_ctx.rep_level, + def_level, + }; + + let children = children + .iter() + .zip(array.columns()) + .map(|(f, a)| Self::try_new(f, ctx, a)) + .collect::>()?; + + Ok(Self::Struct(children, ctx, array.nulls().cloned())) + } + DataType::List(child) + | DataType::LargeList(child) + | DataType::Map(child, _) + | DataType::FixedSizeList(child, _) + | DataType::ListView(child) + | DataType::LargeListView(child) => { + let def_level = match is_nullable { + true => parent_ctx.def_level + 2, + false => parent_ctx.def_level + 1, + }; + + let ctx = LevelContext { + rep_level: parent_ctx.rep_level + 1, + def_level, + }; + + Ok(match field.data_type() { + DataType::List(_) => { + let list = array.as_list(); + let child = Self::try_new(child.as_ref(), ctx, list.values())?; + let is_last = child.child_has_no_nested_rep(); + let offsets = list.offsets().clone(); + Self::List( + Box::new(child), + ctx, + offsets, + list.nulls().cloned(), + is_last, + ) + } + DataType::LargeList(_) => { + let list = array.as_list(); + let child = Self::try_new(child.as_ref(), ctx, list.values())?; + let is_last = child.child_has_no_nested_rep(); + let offsets = list.offsets().clone(); + let nulls = list.nulls().cloned(); + Self::LargeList(Box::new(child), ctx, offsets, nulls, is_last) + } + DataType::Map(_, _) => { + let map = array.as_map(); + let entries = Arc::new(map.entries().clone()) as ArrayRef; + let child = Self::try_new(child.as_ref(), ctx, &entries)?; + let is_last = child.child_has_no_nested_rep(); + let offsets = map.offsets().clone(); + Self::List(Box::new(child), ctx, offsets, map.nulls().cloned(), is_last) + } + DataType::FixedSizeList(_, size) => { + let list = array.as_fixed_size_list(); + let child = Self::try_new(child.as_ref(), ctx, list.values())?; + let nulls = list.nulls().cloned(); + Self::FixedSizeList(Box::new(child), ctx, *size as _, nulls) + } + DataType::ListView(_) => { + let list = array.as_list_view(); + let child = Self::try_new(child.as_ref(), ctx, list.values())?; + let offsets = list.offsets().clone(); + let sizes = list.sizes().clone(); + let nulls = list.nulls().cloned(); + Self::ListView(Box::new(child), ctx, offsets, sizes, nulls) + } + DataType::LargeListView(_) => { + let list = array.as_list_view(); + let child = Self::try_new(child.as_ref(), ctx, list.values())?; + let offsets = list.offsets().clone(); + let sizes = list.sizes().clone(); + let nulls = list.nulls().cloned(); + Self::LargeListView(Box::new(child), ctx, offsets, sizes, nulls) + } + _ => unreachable!(), + }) + } + d => Err(nyi_err!("Datatype {} is not yet supported", d)), + } + } + + /// Finish this [`LevelInfoBuilder`] returning the [`ArrayLevels`] for the leaf columns + /// as enumerated by a depth-first search + fn finish(self) -> Vec { + match self { + LevelInfoBuilder::Primitive(v) => vec![v], + LevelInfoBuilder::List(v, _, _, _, _) + | LevelInfoBuilder::LargeList(v, _, _, _, _) + | LevelInfoBuilder::FixedSizeList(v, _, _, _) + | LevelInfoBuilder::ListView(v, _, _, _, _) + | LevelInfoBuilder::LargeListView(v, _, _, _, _) => v.finish(), + LevelInfoBuilder::Struct(v, _, _) => v.into_iter().flat_map(|l| l.finish()).collect(), + } + } + + /// Given an `array`, write the level data for the elements in `range` + fn write(&mut self, range: Range) { + match self { + LevelInfoBuilder::Primitive(info) => Self::write_leaf(info, range), + LevelInfoBuilder::List(child, ctx, offsets, nulls, is_last) => { + Self::write_list(child, ctx, offsets, nulls.as_ref(), range, *is_last) + } + LevelInfoBuilder::LargeList(child, ctx, offsets, nulls, is_last) => { + Self::write_list(child, ctx, offsets, nulls.as_ref(), range, *is_last) + } + LevelInfoBuilder::FixedSizeList(child, ctx, size, nulls) => { + Self::write_fixed_size_list(child, ctx, *size, nulls.as_ref(), range) + } + LevelInfoBuilder::ListView(child, ctx, offsets, sizes, nulls) => { + Self::write_list_view(child, ctx, offsets, sizes, nulls.as_ref(), range) + } + LevelInfoBuilder::LargeListView(child, ctx, offsets, sizes, nulls) => { + Self::write_list_view(child, ctx, offsets, sizes, nulls.as_ref(), range) + } + LevelInfoBuilder::Struct(children, ctx, nulls) => { + Self::write_struct(children, ctx, nulls.as_ref(), range) + } + } + } + + /// Returns `true` if the child contains no nested repetition levels, meaning + /// each child element produces exactly one rep_level entry in the leaf. + /// This is true for `Primitive` children and `Struct` trees with no list descendants. + fn child_has_no_nested_rep(&self) -> bool { + match self { + LevelInfoBuilder::Primitive(_) => true, + LevelInfoBuilder::Struct(children, _, _) => { + children.iter().all(|c| c.child_has_no_nested_rep()) + } + _ => false, + } + } + + /// Write `range` elements from ListArray `array` + /// + /// Note: MapArrays are `ListArray` under the hood and so are dispatched to this method + fn write_list( + child: &mut LevelInfoBuilder, + ctx: &LevelContext, + offsets: &[O], + nulls: Option<&NullBuffer>, + range: Range, + is_last_level: bool, + ) { + // Fast path: entire list array is null; emit bulk null rep/def levels + if nulls.is_some_and(|nulls| nulls.null_count() == nulls.len()) { + let count = range.end - range.start; + child.visit_leaves(|leaf| { + leaf.extend_uniform_levels(ctx.def_level - 2, ctx.rep_level - 1, count); + }); + return; + } + + // Dispatch to separate functions so the compiler can optimize each + // hot loop independently (function body size affects codegen quality). + if is_last_level { + Self::write_list_direct(child, ctx, offsets, nulls, range); + } else { + Self::write_list_scan(child, ctx, offsets, nulls, range); + } + } + + /// Batch write for lists whose child has no nested repetition. + /// + /// "direct" means writing the child rep levels using offsets without scanning. + fn write_list_direct( + child: &mut LevelInfoBuilder, + ctx: &LevelContext, + offsets: &[O], + nulls: Option<&NullBuffer>, + range: Range, + ) { + let list_start_rep = ctx.rep_level - 1; + + let emit_non_empty_run = |child: &mut LevelInfoBuilder, run_offsets: &[O]| { + debug_assert!(run_offsets.len() >= 2); + let values_start = run_offsets[0].as_usize(); + let values_end = run_offsets[run_offsets.len() - 1].as_usize(); + debug_assert!(values_end > values_start); + + child.write(values_start..values_end); + + // The first element of each list slot needs rep_level = + // list_start_rep to mark a new list boundary. Because there's a 1:1 + // mapping between child elements and rep_level entries, the position + // of each slot's first element is directly computable from offsets. + child.visit_leaves(|leaf| { + debug_assert!(leaf.max_rep_level == ctx.rep_level); + let rep_levels = leaf.rep_levels.materialize_mut().unwrap(); + let batch_len = values_end - values_start; + let batch_base = rep_levels.len() - batch_len; + for slot_offset in run_offsets.iter().take(run_offsets.len() - 1) { + let pos = batch_base + (slot_offset.as_usize() - values_start); + rep_levels[pos] = list_start_rep; + } + }); + }; + + Self::write_list_impl(child, ctx, offsets, nulls, range, emit_non_empty_run); + } + + /// Batch write for lists whose child has nested repetition. + /// + /// After batch-writing child elements, scans backward through rep_levels + /// counting child-element starts to find and stamp slot boundaries. + /// + /// Scan backward because we don't know start offset before writing. + fn write_list_scan( + child: &mut LevelInfoBuilder, + ctx: &LevelContext, + offsets: &[O], + nulls: Option<&NullBuffer>, + range: Range, + ) { + let list_start_rep = ctx.rep_level - 1; + + let emit_non_empty_run = |child: &mut LevelInfoBuilder, run_offsets: &[O]| { + debug_assert!(run_offsets.len() >= 2); + let values_start = run_offsets[0].as_usize(); + let values_end = run_offsets[run_offsets.len() - 1].as_usize(); + debug_assert!(values_end > values_start); + + child.write(values_start..values_end); + + child.visit_leaves(|leaf| { + let rep_levels = leaf.rep_levels.materialize_mut().unwrap(); + + if leaf.max_rep_level == ctx.rep_level { + // This algorithm is the same as write_list_direct. + // Use a separate function because the branch code size would affect codegen + // quality of the hot loop of write_list_direct. + let batch_len = values_end - values_start; + let batch_base = rep_levels.len() - batch_len; + for slot_offset in run_offsets.iter().take(run_offsets.len() - 1) { + let pos = batch_base + (slot_offset.as_usize() - values_start); + rep_levels[pos] = list_start_rep; + } + } else { + // Backward scan: count child-element starts (rep <= ctx.rep_level) + // and stamp list_start_rep at slot boundaries. + let mut slot_bounds = run_offsets[..run_offsets.len() - 1].iter().rev(); + let mut next_stamp_at = values_end - slot_bounds.next().unwrap().as_usize(); + let mut seen = 0usize; + + for rep in rep_levels.iter_mut().rev() { + // Asserting rep >= ctx.rep_level to ensure low level depth haven't + // being written. + debug_assert!(*rep >= ctx.rep_level); + // Count element starts by skipping nested reps (rep > ctx.rep_level). + // + // `==` would also work here: the child is written before the + // parent, so no entry within the batch has rep < ctx.rep_level. + // Benchmarks show no difference, so keep the more defensive `<=`. + if *rep <= ctx.rep_level { + seen += 1; + if seen == next_stamp_at { + *rep = list_start_rep; + match slot_bounds.next() { + Some(offset) => next_stamp_at = values_end - offset.as_usize(), + None => break, + } + } + } + } + } + }); + }; + + Self::write_list_impl(child, ctx, offsets, nulls, range, emit_non_empty_run); + } + + /// Shared run-classification loop for write_list_direct and write_list_scan. + /// Monomorphized per `emit_non_empty_run` closure type, giving the compiler + /// separate optimization contexts for each backfill strategy. + fn write_list_impl( + child: &mut LevelInfoBuilder, + ctx: &LevelContext, + offsets: &[O], + nulls: Option<&NullBuffer>, + range: Range, + mut emit_non_empty_run: impl FnMut(&mut LevelInfoBuilder, &[O]), + ) { + let null_offset = range.start; + let offsets = &offsets[range.start..range.end + 1]; + let list_start_rep = ctx.rep_level - 1; + + let emit_nulls = |child: &mut LevelInfoBuilder, count: usize| { + child.visit_leaves(|leaf| { + leaf.append_rep_level_run(list_start_rep, count); + leaf.append_def_level_run(ctx.def_level - 2, count); + }); + }; + + let emit_empties = |child: &mut LevelInfoBuilder, count: usize| { + child.visit_leaves(|leaf| { + leaf.append_rep_level_run(list_start_rep, count); + leaf.append_def_level_run(ctx.def_level - 1, count); + }); + }; + + // Classify each slot, detect run boundaries, flush on transition. + #[derive(Clone, Copy, PartialEq)] + enum SlotKind { + Null, + Empty, + NonEmpty, + } + + let num_slots = offsets.len() - 1; + if num_slots == 0 { + return; + } + + macro_rules! classify { + ($i:expr, $nulls:expr) => { + if !$nulls.is_valid($i + null_offset) { + SlotKind::Null + } else if offsets[$i] == offsets[$i + 1] { + SlotKind::Empty + } else { + SlotKind::NonEmpty + } + }; + } + + macro_rules! flush_run { + ($kind:expr, $start:expr, $end:expr) => { + match $kind { + SlotKind::Null => emit_nulls(child, $end - $start), + SlotKind::Empty => emit_empties(child, $end - $start), + SlotKind::NonEmpty => emit_non_empty_run(child, &offsets[$start..$end + 1]), + } + }; + } + + match nulls { + // A null buffer without any null can skip the per-slot validity + // checks and use the null-free classification loop below. + Some(nulls) if nulls.null_count() > 0 => { + let mut run_kind = classify!(0, nulls); + let mut run_start: usize = 0; + for i in 1..num_slots { + let kind = classify!(i, nulls); + if kind != run_kind { + flush_run!(run_kind, run_start, i); + run_kind = kind; + run_start = i; + } + } + flush_run!(run_kind, run_start, num_slots); + } + _ => { + let mut run_kind = if offsets[0] == offsets[1] { + SlotKind::Empty + } else { + SlotKind::NonEmpty + }; + let mut run_start: usize = 0; + for i in 1..num_slots { + let kind = if offsets[i] == offsets[i + 1] { + SlotKind::Empty + } else { + SlotKind::NonEmpty + }; + if kind != run_kind { + flush_run!(run_kind, run_start, i); + run_kind = kind; + run_start = i; + } + } + flush_run!(run_kind, run_start, num_slots); + } + } + } + + /// Write `range` elements from ListViewArray `array` + fn write_list_view( + child: &mut LevelInfoBuilder, + ctx: &LevelContext, + offsets: &[O], + sizes: &[O], + nulls: Option<&NullBuffer>, + range: Range, + ) { + let offsets = &offsets[range.start..range.end]; + let sizes = &sizes[range.start..range.end]; + + let write_non_null_slice = + |child: &mut LevelInfoBuilder, start_idx: usize, end_idx: usize| { + child.write(start_idx..end_idx); + child.visit_leaves(|leaf| { + let rep_levels = leaf.rep_levels.materialize_mut().unwrap(); + let mut rev = rep_levels.iter_mut().rev(); + let mut remaining = end_idx - start_idx; + + loop { + let next = rev.next().unwrap(); + if *next > ctx.rep_level { + // Nested element - ignore + continue; + } + + remaining -= 1; + if remaining == 0 { + *next = ctx.rep_level - 1; + break; + } + } + }) + }; + + let write_empty_slice = |child: &mut LevelInfoBuilder| { + child.visit_leaves(|leaf| { + leaf.append_rep_level_run(ctx.rep_level - 1, 1); + leaf.append_def_level_run(ctx.def_level - 1, 1); + }) + }; + + let write_null_slice = |child: &mut LevelInfoBuilder| { + child.visit_leaves(|leaf| { + leaf.append_rep_level_run(ctx.rep_level - 1, 1); + leaf.append_def_level_run(ctx.def_level - 2, 1); + }) + }; + + match nulls { + Some(nulls) => { + let null_offset = range.start; + // TODO: Faster bitmask iteration (#1757) + for (idx, (offset, size)) in offsets.iter().zip(sizes.iter()).enumerate() { + let is_valid = nulls.is_valid(idx + null_offset); + let start_idx = offset.as_usize(); + let size = size.as_usize(); + let end_idx = start_idx + size; + if !is_valid { + write_null_slice(child) + } else if size == 0 { + write_empty_slice(child) + } else { + write_non_null_slice(child, start_idx, end_idx) + } + } + } + None => { + for (offset, size) in offsets.iter().zip(sizes.iter()) { + let start_idx = offset.as_usize(); + let size = size.as_usize(); + let end_idx = start_idx + size; + if size == 0 { + write_empty_slice(child) + } else { + write_non_null_slice(child, start_idx, end_idx) + } + } + } + } + } + + /// Write `range` elements from StructArray `array` + fn write_struct( + children: &mut [LevelInfoBuilder], + ctx: &LevelContext, + nulls: Option<&NullBuffer>, + range: Range, + ) { + let write_null = |children: &mut [LevelInfoBuilder], range: Range| { + let len = range.end - range.start; + for child in children { + child.visit_leaves(|info| { + info.extend_uniform_levels(ctx.def_level - 1, ctx.rep_level, len); + }) + } + }; + + // Fast path: entire struct array is null; emit bulk null def/rep levels + if nulls.is_some_and(|nulls| nulls.null_count() == nulls.len()) { + write_null(children, range); + return; + } + + let write_non_null = |children: &mut [LevelInfoBuilder], range: Range| { + for child in children { + child.write(range.clone()) + } + }; + + match nulls { + Some(validity) => { + let mut last_non_null_idx = None; + let mut last_null_idx = None; + + // TODO: Faster bitmask iteration (#1757) + for i in range.clone() { + match validity.is_valid(i) { + true => { + if let Some(last_idx) = last_null_idx.take() { + write_null(children, last_idx..i) + } + last_non_null_idx.get_or_insert(i); + } + false => { + if let Some(last_idx) = last_non_null_idx.take() { + write_non_null(children, last_idx..i) + } + last_null_idx.get_or_insert(i); + } + } + } + + if let Some(last_idx) = last_null_idx.take() { + write_null(children, last_idx..range.end) + } + + if let Some(last_idx) = last_non_null_idx.take() { + write_non_null(children, last_idx..range.end) + } + } + None => write_non_null(children, range), + } + } + + /// Write `range` elements from FixedSizeListArray with child data `values` and null bitmap `nulls`. + fn write_fixed_size_list( + child: &mut LevelInfoBuilder, + ctx: &LevelContext, + fixed_size: usize, + nulls: Option<&NullBuffer>, + range: Range, + ) { + // Fast path: entire fixed-size list array is null + if nulls.is_some_and(|nulls| nulls.null_count() == nulls.len()) { + let count = range.end - range.start; + child.visit_leaves(|leaf| { + leaf.extend_uniform_levels(ctx.def_level - 2, ctx.rep_level - 1, count); + }); + return; + } + + let write_non_null = |child: &mut LevelInfoBuilder, start_idx: usize, end_idx: usize| { + let values_start = start_idx * fixed_size; + let values_end = end_idx * fixed_size; + child.write(values_start..values_end); + + child.visit_leaves(|leaf| { + let rep_levels = leaf.rep_levels.materialize_mut().unwrap(); + + let row_indices = (0..fixed_size) + .rev() + .cycle() + .take(values_end - values_start); + + // Step backward over the child rep levels and mark the start of each list + rep_levels + .iter_mut() + .rev() + // Filter out reps from nested children + .filter(|&&mut r| r == ctx.rep_level) + .zip(row_indices) + .for_each(|(r, idx)| { + if idx == 0 { + *r = ctx.rep_level - 1; + } + }); + }) + }; + + // If list size is 0, ignore values and just write rep/def levels. + let write_empty = |child: &mut LevelInfoBuilder, start_idx: usize, end_idx: usize| { + let len = end_idx - start_idx; + child.visit_leaves(|leaf| { + leaf.append_rep_level_run(ctx.rep_level - 1, len); + leaf.append_def_level_run(ctx.def_level - 1, len); + }) + }; + + let write_rows = |child: &mut LevelInfoBuilder, start_idx: usize, end_idx: usize| { + if fixed_size > 0 { + write_non_null(child, start_idx, end_idx) + } else { + write_empty(child, start_idx, end_idx) + } + }; + + match nulls { + Some(nulls) => { + let mut start_idx = None; + for idx in range.clone() { + if nulls.is_valid(idx) { + // Start a run of valid rows if not already inside of one + start_idx.get_or_insert(idx); + } else { + // Write out any pending valid rows + if let Some(start) = start_idx.take() { + write_rows(child, start, idx); + } + // Add null row + child.visit_leaves(|leaf| { + leaf.append_rep_level_run(ctx.rep_level - 1, 1); + leaf.append_def_level_run(ctx.def_level - 2, 1); + }) + } + } + // Write out any remaining valid rows + if let Some(start) = start_idx.take() { + write_rows(child, start, range.end); + } + } + // If all rows are valid then write the whole array + None => write_rows(child, range.start, range.end), + } + } + + /// Write a primitive array, as defined by [`is_leaf`] + fn write_leaf(info: &mut ArrayLevels, range: Range) { + let len = range.end - range.start; + + // Fast path: entire leaf array is null + if let Some(nulls) = &info.logical_nulls { + if !matches!(info.def_levels, LevelData::Absent) && nulls.null_count() == nulls.len() { + info.extend_uniform_levels(info.max_def_level - 1, info.max_rep_level, len); + return; + } + } + + if matches!(info.def_levels, LevelData::Absent) { + info.non_null_indices.extend(range.clone()); + } else { + let max_def_level = info.max_def_level; + match &info.logical_nulls { + Some(nulls) => { + assert!(range.end <= nulls.len()); + // Bulk-fill is profitable only on null-heavy ranges long enough to + // amortize the slice/popcount cost; see `BULK_FILL_MIN_LEN` and the + // PR description for the threshold sweep. The gate uses the cached + // buffer-wide `null_count` (O(1)) to stay cheap on the cold path. + if len >= BULK_FILL_MIN_LEN && nulls.null_count() * 2 >= nulls.len() { + let range_nulls = nulls.slice(range.start, len); + let valid_in_range = len - range_nulls.null_count(); + let null_def_level = max_def_level - 1; + let buf = info + .def_levels + .materialize_mut() + .expect("definition levels present"); + let base = buf.len(); + buf.resize(base + len, null_def_level); + for i in range_nulls.valid_indices() { + buf[base + i] = max_def_level; + } + info.non_null_indices.reserve(valid_in_range); + info.non_null_indices + .extend(range_nulls.valid_indices().map(|i| i + range.start)); + } else { + let bits = nulls.inner(); + info.def_levels.extend_from_iter(range.clone().map(|i| { + // Safety: range.end was asserted to be in bounds earlier + let valid = unsafe { bits.value_unchecked(i) }; + max_def_level - (!valid as i16) + })); + info.non_null_indices.reserve(len); + info.non_null_indices.extend( + BitIndexIterator::new(bits.inner(), bits.offset() + range.start, len) + .map(|i| i + range.start), + ); + } + } + None => { + info.append_def_level_run(max_def_level, len); + info.non_null_indices.reserve(len); + info.non_null_indices.extend(range.clone()); + } + } + } + + if !matches!(info.rep_levels, LevelData::Absent) { + info.append_rep_level_run(info.max_rep_level, len); + } + } + + /// Visits all children of this node in depth first order + fn visit_leaves(&mut self, visit: impl Fn(&mut ArrayLevels) + Copy) { + match self { + LevelInfoBuilder::Primitive(info) => visit(info), + LevelInfoBuilder::List(c, _, _, _, _) + | LevelInfoBuilder::LargeList(c, _, _, _, _) + | LevelInfoBuilder::FixedSizeList(c, _, _, _) + | LevelInfoBuilder::ListView(c, _, _, _, _) + | LevelInfoBuilder::LargeListView(c, _, _, _, _) => c.visit_leaves(visit), + LevelInfoBuilder::Struct(children, _, _) => { + for c in children { + c.visit_leaves(visit) + } + } + } + } + + /// Determine if the fields are compatible for purposes of constructing `LevelBuilderInfo`. + /// + /// Fields are compatible if they're the same type. Otherwise if one of them is a dictionary + /// and the other is a native array, the dictionary values must have the same type as the + /// native array + fn types_compatible(a: &DataType, b: &DataType) -> bool { + // if the Arrow data types are equal, the types are deemed compatible + if a.equals_datatype(b) { + return true; + } + + // get the values out of the dictionaries + let (a, b) = match (a, b) { + (DataType::Dictionary(_, va), DataType::Dictionary(_, vb)) => { + (va.as_ref(), vb.as_ref()) + } + (DataType::Dictionary(_, v), b) => (v.as_ref(), b), + (a, DataType::Dictionary(_, v)) => (a, v.as_ref()), + _ => (a, b), + }; + + // now that we've got the values from one/both dictionaries, if the values + // have the same Arrow data type, they're compatible + if a == b { + return true; + } + + // here we have different Arrow data types, but if the array contains the same type of data + // then we consider the type compatible + match a { + // String, StringView and LargeString are compatible + DataType::Utf8 => matches!(b, DataType::LargeUtf8 | DataType::Utf8View), + DataType::Utf8View => matches!(b, DataType::LargeUtf8 | DataType::Utf8), + DataType::LargeUtf8 => matches!(b, DataType::Utf8 | DataType::Utf8View), + + // Binary, BinaryView and LargeBinary are compatible + DataType::Binary => matches!(b, DataType::LargeBinary | DataType::BinaryView), + DataType::BinaryView => matches!(b, DataType::LargeBinary | DataType::Binary), + DataType::LargeBinary => matches!(b, DataType::Binary | DataType::BinaryView), + + // otherwise we have incompatible types + _ => false, + } + } +} + +/// The data necessary to write a primitive Arrow array to parquet, taking into account +/// any non-primitive parents it may have in the arrow representation +#[derive(Debug, Clone)] +pub(crate) enum LevelData { + Absent, + Materialized(Vec), + Uniform { value: i16, count: usize }, +} + +// Compare logical level contents rather than physical representation, so a +// uniform run compares equal to the equivalent materialized buffer. +impl PartialEq for LevelData { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (Self::Absent, Self::Absent) => true, + (Self::Materialized(a), Self::Materialized(b)) => a == b, + (Self::Uniform { value: v, count: n }, Self::Materialized(b)) + | (Self::Materialized(b), Self::Uniform { value: v, count: n }) => { + b.len() == *n && b.iter().all(|x| x == v) + } + ( + Self::Uniform { + value: v1, + count: n1, + }, + Self::Uniform { + value: v2, + count: n2, + }, + ) => v1 == v2 && n1 == n2, + _ => false, + } + } +} + +impl Eq for LevelData {} + +impl LevelData { + fn new(present: bool) -> Self { + match present { + true => Self::Materialized(Vec::new()), + false => Self::Absent, + } + } + + pub(crate) fn as_ref(&self) -> LevelDataRef<'_> { + match self { + Self::Absent => LevelDataRef::Absent, + Self::Materialized(values) => LevelDataRef::Materialized(values), + Self::Uniform { value, count } => LevelDataRef::Uniform { + value: *value, + count: *count, + }, + } + } + + pub(crate) fn slice(&self, offset: usize, len: usize) -> Self { + match self { + Self::Absent => Self::Absent, + Self::Materialized(values) => Self::Materialized(values[offset..offset + len].to_vec()), + Self::Uniform { value, .. } => Self::Uniform { + value: *value, + count: len, + }, + } + } + + fn append_run(&mut self, value: i16, count: usize) { + if count == 0 { + return; + } + + match self { + // No physical level stream exists for this schema. Higher-level + // traversal may still append implicit levels, so this remains a no-op. + Self::Absent => {} + // Start compact: the first appended run can be represented without + // allocating a level buffer. + Self::Materialized(values) if values.is_empty() => { + *self = Self::Uniform { value, count }; + } + // Already materialized, so preserve the buffer representation and append. + Self::Materialized(values) => values.extend(std::iter::repeat_n(value, count)), + // Preserve the compact representation while the appended run has + // the same value. + Self::Uniform { + value: uniform_value, + count: uniform_count, + } if *uniform_value == value => { + *uniform_count += count; + } + // A different value breaks the uniform representation. Materialize + // the existing run, then append the new run to the buffer. + Self::Uniform { .. } => { + let values = self.materialize_mut().unwrap(); + values.extend(std::iter::repeat_n(value, count)); + } + } + } + + fn extend_from_iter(&mut self, iter: I) + where + I: IntoIterator, + { + if let Some(values) = self.materialize_mut() { + values.extend(iter); + } + } + + /// Convert a uniform run into a materialized buffer if needed, then return + /// the mutable level buffer. Returns `None` when no physical level stream exists. + fn materialize_mut(&mut self) -> Option<&mut Vec> { + match self { + Self::Absent => None, + Self::Materialized(values) => Some(values), + Self::Uniform { value, count } => { + let values = vec![*value; *count]; + *self = Self::Materialized(values); + match self { + Self::Materialized(values) => Some(values), + _ => unreachable!(), + } + } + } + } +} + +#[derive(Debug, Clone)] +pub(crate) struct ArrayLevels { + /// Array's definition levels + /// + /// Present if `max_def_level != 0` + def_levels: LevelData, + + /// Array's optional repetition levels + /// + /// Present if `max_rep_level != 0` + rep_levels: LevelData, + + /// The corresponding array identifying non-null slices of data + /// from the primitive array + non_null_indices: Vec, + + /// The maximum definition level for this leaf column + max_def_level: i16, + + /// The maximum repetition for this leaf column + max_rep_level: i16, + + /// The arrow array + array: ArrayRef, + + /// cached logical nulls of the array. + logical_nulls: Option, +} + +impl PartialEq for ArrayLevels { + fn eq(&self, other: &Self) -> bool { + self.def_levels == other.def_levels + && self.rep_levels == other.rep_levels + && self.non_null_indices == other.non_null_indices + && self.max_def_level == other.max_def_level + && self.max_rep_level == other.max_rep_level + && self.array.as_ref() == other.array.as_ref() + && self.logical_nulls.as_ref() == other.logical_nulls.as_ref() + } +} +impl Eq for ArrayLevels {} + +impl ArrayLevels { + fn new(ctx: LevelContext, is_nullable: bool, array: ArrayRef) -> Self { + let max_rep_level = ctx.rep_level; + let max_def_level = match is_nullable { + true => ctx.def_level + 1, + false => ctx.def_level, + }; + + let logical_nulls = array.logical_nulls(); + + Self { + def_levels: LevelData::new(max_def_level != 0), + rep_levels: LevelData::new(max_rep_level != 0), + non_null_indices: vec![], + max_def_level, + max_rep_level, + array, + logical_nulls, + } + } + + pub fn array(&self) -> &ArrayRef { + &self.array + } + + pub(crate) fn def_level_data(&self) -> &LevelData { + &self.def_levels + } + + pub(crate) fn rep_level_data(&self) -> &LevelData { + &self.rep_levels + } + + pub fn non_null_indices(&self) -> &[usize] { + &self.non_null_indices + } + + /// Create a sliced view of this `ArrayLevels` for a CDC chunk. + /// + /// The chunk's `value_offset`/`num_values` select the relevant slice of + /// `non_null_indices`. The array is sliced to the range covered by + /// those indices, and they are shifted to be relative to the slice. + pub(crate) fn slice_for_chunk(&self, chunk: &CdcChunk) -> Self { + let def_levels = self.def_levels.slice(chunk.level_offset, chunk.num_levels); + let rep_levels = self.rep_levels.slice(chunk.level_offset, chunk.num_levels); + + // Select the non-null indices for this chunk. + let nni = &self.non_null_indices[chunk.value_offset..chunk.value_offset + chunk.num_values]; + // Compute the array range spanned by the non-null indices. + // When nni is empty (all-null chunk), start=0, end=0 → zero-length + // array slice; write_batch_internal will process only the def/rep + // levels and write no values. + let start = nni.first().copied().unwrap_or(0); + let end = nni.last().map_or(0, |&i| i + 1); + // Shift indices to be relative to the sliced array. + let non_null_indices = nni.iter().map(|&idx| idx - start).collect(); + // Slice the array to the computed range. + let array = self.array.slice(start, end - start); + let logical_nulls = array.logical_nulls(); + + Self { + def_levels, + rep_levels, + non_null_indices, + max_def_level: self.max_def_level, + max_rep_level: self.max_rep_level, + array, + logical_nulls, + } + } + + /// Bulk-emit `count` uniform def/rep levels. + fn extend_uniform_levels(&mut self, def_val: i16, rep_val: i16, count: usize) { + self.def_levels.append_run(def_val, count); + self.rep_levels.append_run(rep_val, count); + } + + fn append_def_level_run(&mut self, value: i16, count: usize) { + self.def_levels.append_run(value, count); + } + + fn append_rep_level_run(&mut self, value: i16, count: usize) { + self.rep_levels.append_run(value, count); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::column::chunker::CdcChunk; + + use arrow_array::builder::*; + use arrow_array::types::Int32Type; + use arrow_array::*; + use arrow_buffer::{Buffer, ToByteSlice}; + use arrow_cast::display::array_value_to_string; + use arrow_data::{ArrayData, ArrayDataBuilder}; + use arrow_schema::{Fields, Schema}; + + #[test] + fn test_calculate_array_levels_twitter_example() { + // based on the example at https://blog.twitter.com/engineering/en_us/a/2013/dremel-made-simple-with-parquet.html + // [[a, b, c], [d, e, f, g]], [[h], [i,j]] + + let leaf_type = Field::new_list_field(DataType::Int32, false); + let inner_type = DataType::List(Arc::new(leaf_type)); + let inner_field = Field::new("l2", inner_type.clone(), false); + let outer_type = DataType::List(Arc::new(inner_field)); + let outer_field = Field::new("l1", outer_type.clone(), false); + + let primitives = Int32Array::from_iter(0..10); + + // Cannot use from_iter_primitive as always infers nullable + let offsets = Buffer::from_iter([0_i32, 3, 7, 8, 10]); + let inner_list = ArrayDataBuilder::new(inner_type) + .len(4) + .add_buffer(offsets) + .add_child_data(primitives.to_data()) + .build() + .unwrap(); + + let offsets = Buffer::from_iter([0_i32, 2, 4]); + let outer_list = ArrayDataBuilder::new(outer_type) + .len(2) + .add_buffer(offsets) + .add_child_data(inner_list) + .build() + .unwrap(); + let outer_list = make_array(outer_list); + + let levels = calculate_array_levels(&outer_list, &outer_field).unwrap(); + assert_eq!(levels.len(), 1); + + let expected = ArrayLevels { + def_levels: LevelData::Materialized(vec![2; 10]), + rep_levels: LevelData::Materialized(vec![0, 2, 2, 1, 2, 2, 2, 0, 1, 2]), + non_null_indices: vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + max_def_level: 2, + max_rep_level: 2, + array: Arc::new(primitives), + logical_nulls: None, + }; + assert_eq!(&levels[0], &expected); + } + + #[test] + fn test_calculate_one_level_1() { + // This test calculates the levels for a non-null primitive array + let array = Arc::new(Int32Array::from_iter(0..10)) as ArrayRef; + let field = Field::new_list_field(DataType::Int32, false); + + let levels = calculate_array_levels(&array, &field).unwrap(); + assert_eq!(levels.len(), 1); + + let expected_levels = ArrayLevels { + def_levels: LevelData::Absent, + rep_levels: LevelData::Absent, + non_null_indices: (0..10).collect(), + max_def_level: 0, + max_rep_level: 0, + array, + logical_nulls: None, + }; + assert_eq!(&levels[0], &expected_levels); + } + + #[test] + fn test_calculate_one_level_2() { + // This test calculates the levels for a nullable primitive array + let array = Arc::new(Int32Array::from_iter([ + Some(0), + None, + Some(0), + Some(0), + None, + ])) as ArrayRef; + let field = Field::new_list_field(DataType::Int32, true); + + let levels = calculate_array_levels(&array, &field).unwrap(); + assert_eq!(levels.len(), 1); + + let logical_nulls = array.logical_nulls(); + let expected_levels = ArrayLevels { + def_levels: LevelData::Materialized(vec![1, 0, 1, 1, 0]), + rep_levels: LevelData::Absent, + non_null_indices: vec![0, 2, 3], + max_def_level: 1, + max_rep_level: 0, + array, + logical_nulls, + }; + assert_eq!(&levels[0], &expected_levels); + } + + #[test] + fn test_calculate_array_levels_1() { + let leaf_field = Field::new_list_field(DataType::Int32, false); + let list_type = DataType::List(Arc::new(leaf_field)); + + // if all array values are defined (e.g. batch>) + // [[0], [1], [2], [3], [4]] + + let leaf_array = Int32Array::from_iter(0..5); + // Cannot use from_iter_primitive as always infers nullable + let offsets = Buffer::from_iter(0_i32..6); + let list = ArrayDataBuilder::new(list_type.clone()) + .len(5) + .add_buffer(offsets) + .add_child_data(leaf_array.to_data()) + .build() + .unwrap(); + let list = make_array(list); + + let list_field = Field::new("list", list_type.clone(), false); + let levels = calculate_array_levels(&list, &list_field).unwrap(); + assert_eq!(levels.len(), 1); + + let expected_levels = ArrayLevels { + def_levels: LevelData::Materialized(vec![1; 5]), + rep_levels: LevelData::Materialized(vec![0; 5]), + non_null_indices: (0..5).collect(), + max_def_level: 1, + max_rep_level: 1, + array: Arc::new(leaf_array), + logical_nulls: None, + }; + assert_eq!(&levels[0], &expected_levels); + + // array: [[0, 0], NULL, [2, 2], [3, 3, 3, 3], [4, 4, 4]] + // all values are defined as we do not have nulls on the root (batch) + // repetition: + // 0: 0, 1 + // 1: 0 + // 2: 0, 1 + // 3: 0, 1, 1, 1 + // 4: 0, 1, 1 + let leaf_array = Int32Array::from_iter([0, 0, 2, 2, 3, 3, 3, 3, 4, 4, 4]); + let offsets = Buffer::from_iter([0_i32, 2, 2, 4, 8, 11]); + let list = ArrayDataBuilder::new(list_type.clone()) + .len(5) + .add_buffer(offsets) + .add_child_data(leaf_array.to_data()) + .null_bit_buffer(Some(Buffer::from([0b00011101]))) + .build() + .unwrap(); + let list = make_array(list); + + let list_field = Field::new("list", list_type, true); + let levels = calculate_array_levels(&list, &list_field).unwrap(); + assert_eq!(levels.len(), 1); + + let expected_levels = ArrayLevels { + def_levels: LevelData::Materialized(vec![2, 2, 0, 2, 2, 2, 2, 2, 2, 2, 2, 2]), + rep_levels: LevelData::Materialized(vec![0, 1, 0, 0, 1, 0, 1, 1, 1, 0, 1, 1]), + non_null_indices: (0..11).collect(), + max_def_level: 2, + max_rep_level: 1, + array: Arc::new(leaf_array), + logical_nulls: None, + }; + assert_eq!(&levels[0], &expected_levels); + } + + #[test] + fn test_calculate_array_levels_2() { + // If some values are null + // + // This emulates an array in the form: > + // with values: + // - 0: [0, 1], but is null because of the struct + // - 1: [] + // - 2: [2, 3], but is null because of the struct + // - 3: [4, 5, 6, 7] + // - 4: [8, 9, 10] + // + // If the first values of a list are null due to a parent, we have to still account for them + // while indexing, because they would affect the way the child is indexed + // i.e. in the above example, we have to know that [0, 1] has to be skipped + let leaf = Int32Array::from_iter(0..11); + let leaf_field = Field::new("leaf", DataType::Int32, false); + + let list_type = DataType::List(Arc::new(leaf_field)); + let list = ArrayData::builder(list_type.clone()) + .len(5) + .add_child_data(leaf.to_data()) + .add_buffer(Buffer::from_iter([0_i32, 2, 2, 4, 8, 11])) + .build() + .unwrap(); + + let list = make_array(list); + let list_field = Arc::new(Field::new("list", list_type, true)); + + let struct_array = + StructArray::from((vec![(list_field, list)], Buffer::from([0b00011010]))); + let array = Arc::new(struct_array) as ArrayRef; + + let struct_field = Field::new("struct", array.data_type().clone(), true); + + let levels = calculate_array_levels(&array, &struct_field).unwrap(); + assert_eq!(levels.len(), 1); + + let expected_levels = ArrayLevels { + def_levels: LevelData::Materialized(vec![0, 2, 0, 3, 3, 3, 3, 3, 3, 3]), + rep_levels: LevelData::Materialized(vec![0, 0, 0, 0, 1, 1, 1, 0, 1, 1]), + non_null_indices: (4..11).collect(), + max_def_level: 3, + max_rep_level: 1, + array: Arc::new(leaf), + logical_nulls: None, + }; + + assert_eq!(&levels[0], &expected_levels); + + // nested lists + + // 0: [[100, 101], [102, 103]] + // 1: [] + // 2: [[104, 105], [106, 107]] + // 3: [[108, 109], [110, 111], [112, 113], [114, 115]] + // 4: [[116, 117], [118, 119], [120, 121]] + + let leaf = Int32Array::from_iter(100..122); + let leaf_field = Field::new("leaf", DataType::Int32, true); + + let l1_type = DataType::List(Arc::new(leaf_field)); + let offsets = Buffer::from_iter([0_i32, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22]); + let l1 = ArrayData::builder(l1_type.clone()) + .len(11) + .add_child_data(leaf.to_data()) + .add_buffer(offsets) + .build() + .unwrap(); + + let l1_field = Field::new("l1", l1_type, true); + let l2_type = DataType::List(Arc::new(l1_field)); + let l2 = ArrayData::builder(l2_type) + .len(5) + .add_child_data(l1) + .add_buffer(Buffer::from_iter([0, 2, 2, 4, 8, 11])) + .build() + .unwrap(); + + let l2 = make_array(l2); + let l2_field = Field::new("l2", l2.data_type().clone(), true); + + let levels = calculate_array_levels(&l2, &l2_field).unwrap(); + assert_eq!(levels.len(), 1); + + let expected_levels = ArrayLevels { + def_levels: LevelData::Materialized(vec![ + 5, 5, 5, 5, 1, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, + ]), + rep_levels: LevelData::Materialized(vec![ + 0, 2, 1, 2, 0, 0, 2, 1, 2, 0, 2, 1, 2, 1, 2, 1, 2, 0, 2, 1, 2, 1, 2, + ]), + non_null_indices: (0..22).collect(), + max_def_level: 5, + max_rep_level: 2, + array: Arc::new(leaf), + logical_nulls: None, + }; + + assert_eq!(&levels[0], &expected_levels); + } + + #[test] + fn test_calculate_array_levels_nested_list() { + let leaf_field = Field::new("leaf", DataType::Int32, false); + let list_type = DataType::List(Arc::new(leaf_field)); + + // if all array values are defined (e.g. batch>) + // The array at this level looks like: + // 0: [a] + // 1: [a] + // 2: [a] + // 3: [a] + + let leaf = Int32Array::from_iter([0; 4]); + let list = ArrayData::builder(list_type.clone()) + .len(4) + .add_buffer(Buffer::from_iter(0_i32..5)) + .add_child_data(leaf.to_data()) + .build() + .unwrap(); + let list = make_array(list); + + let list_field = Field::new("list", list_type.clone(), false); + let levels = calculate_array_levels(&list, &list_field).unwrap(); + assert_eq!(levels.len(), 1); + + let expected_levels = ArrayLevels { + def_levels: LevelData::Materialized(vec![1; 4]), + rep_levels: LevelData::Materialized(vec![0; 4]), + non_null_indices: (0..4).collect(), + max_def_level: 1, + max_rep_level: 1, + array: Arc::new(leaf), + logical_nulls: None, + }; + assert_eq!(&levels[0], &expected_levels); + + // 0: null + // 1: [1, 2, 3] + // 2: [4, 5] + // 3: [6, 7] + let leaf = Int32Array::from_iter(0..8); + let list = ArrayData::builder(list_type.clone()) + .len(4) + .add_buffer(Buffer::from_iter([0_i32, 0, 3, 5, 7])) + .null_bit_buffer(Some(Buffer::from([0b00001110]))) + .add_child_data(leaf.to_data()) + .build() + .unwrap(); + let list = make_array(list); + let list_field = Arc::new(Field::new("list", list_type, true)); + + let struct_array = StructArray::from(vec![(list_field, list)]); + let array = Arc::new(struct_array) as ArrayRef; + + let struct_field = Field::new("struct", array.data_type().clone(), true); + let levels = calculate_array_levels(&array, &struct_field).unwrap(); + assert_eq!(levels.len(), 1); + + let expected_levels = ArrayLevels { + def_levels: LevelData::Materialized(vec![1, 3, 3, 3, 3, 3, 3, 3]), + rep_levels: LevelData::Materialized(vec![0, 0, 1, 1, 0, 1, 0, 1]), + non_null_indices: (0..7).collect(), + max_def_level: 3, + max_rep_level: 1, + array: Arc::new(leaf), + logical_nulls: None, + }; + assert_eq!(&levels[0], &expected_levels); + + // nested lists + // In a JSON syntax with the schema: >>>, this translates into: + // 0: {"struct": null } + // 1: {"struct": [ [201], [202, 203], [] ]} + // 2: {"struct": [ [204, 205, 206], [207, 208, 209, 210] ]} + // 3: {"struct": [ [], [211, 212, 213, 214, 215] ]} + + let leaf = Int32Array::from_iter(201..216); + let leaf_field = Field::new("leaf", DataType::Int32, false); + let list_1_type = DataType::List(Arc::new(leaf_field)); + let list_1 = ArrayData::builder(list_1_type.clone()) + .len(7) + .add_buffer(Buffer::from_iter([0_i32, 1, 3, 3, 6, 10, 10, 15])) + .add_child_data(leaf.to_data()) + .build() + .unwrap(); + + let list_1_field = Field::new("l1", list_1_type, true); + let list_2_type = DataType::List(Arc::new(list_1_field)); + let list_2 = ArrayData::builder(list_2_type.clone()) + .len(4) + .add_buffer(Buffer::from_iter([0_i32, 0, 3, 5, 7])) + .null_bit_buffer(Some(Buffer::from([0b00001110]))) + .add_child_data(list_1) + .build() + .unwrap(); + + let list_2 = make_array(list_2); + let list_2_field = Arc::new(Field::new("list_2", list_2_type, true)); + + let struct_array = + StructArray::from((vec![(list_2_field, list_2)], Buffer::from([0b00001111]))); + let struct_field = Field::new("struct", struct_array.data_type().clone(), true); + + let array = Arc::new(struct_array) as ArrayRef; + let levels = calculate_array_levels(&array, &struct_field).unwrap(); + assert_eq!(levels.len(), 1); + + let expected_levels = ArrayLevels { + def_levels: LevelData::Materialized(vec![ + 1, 5, 5, 5, 4, 5, 5, 5, 5, 5, 5, 5, 4, 5, 5, 5, 5, 5, + ]), + rep_levels: LevelData::Materialized(vec![ + 0, 0, 1, 2, 1, 0, 2, 2, 1, 2, 2, 2, 0, 1, 2, 2, 2, 2, + ]), + non_null_indices: (0..15).collect(), + max_def_level: 5, + max_rep_level: 2, + array: Arc::new(leaf), + logical_nulls: None, + }; + assert_eq!(&levels[0], &expected_levels); + } + + #[test] + fn test_calculate_nested_struct_levels() { + // tests a > + // array: + // - {a: {b: {c: 1}}} + // - {a: {b: {c: null}}} + // - {a: {b: {c: 3}}} + // - {a: {b: null}} + // - {a: null}} + // - {a: {b: {c: 6}}} + + let c = Int32Array::from_iter([Some(1), None, Some(3), None, Some(5), Some(6)]); + let leaf = Arc::new(c) as ArrayRef; + let c_field = Arc::new(Field::new("c", DataType::Int32, true)); + let b = StructArray::from(((vec![(c_field, leaf.clone())]), Buffer::from([0b00110111]))); + + let b_field = Arc::new(Field::new("b", b.data_type().clone(), true)); + let a = StructArray::from(( + (vec![(b_field, Arc::new(b) as ArrayRef)]), + Buffer::from([0b00101111]), + )); + + let a_field = Field::new("a", a.data_type().clone(), true); + let a_array = Arc::new(a) as ArrayRef; + + let levels = calculate_array_levels(&a_array, &a_field).unwrap(); + assert_eq!(levels.len(), 1); + + let logical_nulls = leaf.logical_nulls(); + let expected_levels = ArrayLevels { + def_levels: LevelData::Materialized(vec![3, 2, 3, 1, 0, 3]), + rep_levels: LevelData::Absent, + non_null_indices: vec![0, 2, 5], + max_def_level: 3, + max_rep_level: 0, + array: leaf, + logical_nulls, + }; + assert_eq!(&levels[0], &expected_levels); + } + + #[test] + fn list_single_column() { + // this tests the level generation from the arrow_writer equivalent test + + let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); + let a_value_offsets = arrow::buffer::Buffer::from_iter([0_i32, 1, 3, 3, 6, 10]); + let a_list_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))); + let a_list_data = ArrayData::builder(a_list_type.clone()) + .len(5) + .add_buffer(a_value_offsets) + .null_bit_buffer(Some(Buffer::from([0b00011011]))) + .add_child_data(a_values.to_data()) + .build() + .unwrap(); + + assert_eq!(a_list_data.null_count(), 1); + + let a = ListArray::from(a_list_data); + + let item_field = Field::new_list_field(a_list_type, true); + let mut builder = levels(&item_field, a); + builder.write(2..4); + let levels = builder.finish(); + + assert_eq!(levels.len(), 1); + + let list_level = &levels[0]; + + let expected_level = ArrayLevels { + def_levels: LevelData::Materialized(vec![0, 3, 3, 3]), + rep_levels: LevelData::Materialized(vec![0, 0, 1, 1]), + non_null_indices: vec![3, 4, 5], + max_def_level: 3, + max_rep_level: 1, + array: Arc::new(a_values), + logical_nulls: None, + }; + assert_eq!(list_level, &expected_level); + } + + #[test] + fn mixed_struct_list() { + // this tests the level generation from the equivalent arrow_writer_complex test + + // define schema + let struct_field_d = Arc::new(Field::new("d", DataType::Float64, true)); + let struct_field_f = Arc::new(Field::new("f", DataType::Float32, true)); + let struct_field_g = Arc::new(Field::new( + "g", + DataType::List(Arc::new(Field::new("items", DataType::Int16, false))), + false, + )); + let struct_field_e = Arc::new(Field::new( + "e", + DataType::Struct(vec![struct_field_f.clone(), struct_field_g.clone()].into()), + true, + )); + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, true), + Field::new( + "c", + DataType::Struct(vec![struct_field_d.clone(), struct_field_e.clone()].into()), + true, // https://github.com/apache/arrow-rs/issues/245 + ), + ]); + + // create some data + let a = Int32Array::from(vec![1, 2, 3, 4, 5]); + let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]); + let d = Float64Array::from(vec![None, None, None, Some(1.0), None]); + let f = Float32Array::from(vec![Some(0.0), None, Some(333.3), None, Some(5.25)]); + + let g_value = Int16Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); + + // Construct a buffer for value offsets, for the nested array: + // [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]] + let g_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice()); + + // Construct a list array from the above two + let g_list_data = ArrayData::builder(struct_field_g.data_type().clone()) + .len(5) + .add_buffer(g_value_offsets) + .add_child_data(g_value.into_data()) + .build() + .unwrap(); + let g = ListArray::from(g_list_data); + + let e = StructArray::from(vec![ + (struct_field_f, Arc::new(f.clone()) as ArrayRef), + (struct_field_g, Arc::new(g) as ArrayRef), + ]); + + let c = StructArray::from(vec![ + (struct_field_d, Arc::new(d.clone()) as ArrayRef), + (struct_field_e, Arc::new(e) as ArrayRef), + ]); + + // build a record batch + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![Arc::new(a.clone()), Arc::new(b.clone()), Arc::new(c)], + ) + .unwrap(); + + ////////////////////////////////////////////// + // calculate the list's level + let mut levels = vec![]; + batch + .columns() + .iter() + .zip(batch.schema().fields()) + .for_each(|(array, field)| { + let mut array_levels = calculate_array_levels(array, field).unwrap(); + levels.append(&mut array_levels); + }); + assert_eq!(levels.len(), 5); + + // test "a" levels + let list_level = &levels[0]; + + let expected_level = ArrayLevels { + def_levels: LevelData::Absent, + rep_levels: LevelData::Absent, + non_null_indices: vec![0, 1, 2, 3, 4], + max_def_level: 0, + max_rep_level: 0, + array: Arc::new(a), + logical_nulls: None, + }; + assert_eq!(list_level, &expected_level); + + // test "b" levels + let list_level = levels.get(1).unwrap(); + + let b_logical_nulls = b.logical_nulls(); + let expected_level = ArrayLevels { + def_levels: LevelData::Materialized(vec![1, 0, 0, 1, 1]), + rep_levels: LevelData::Absent, + non_null_indices: vec![0, 3, 4], + max_def_level: 1, + max_rep_level: 0, + array: Arc::new(b), + logical_nulls: b_logical_nulls, + }; + assert_eq!(list_level, &expected_level); + + // test "d" levels + let list_level = levels.get(2).unwrap(); + + let d_logical_nulls = d.logical_nulls(); + let expected_level = ArrayLevels { + def_levels: LevelData::Materialized(vec![1, 1, 1, 2, 1]), + rep_levels: LevelData::Absent, + non_null_indices: vec![3], + max_def_level: 2, + max_rep_level: 0, + array: Arc::new(d), + logical_nulls: d_logical_nulls, + }; + assert_eq!(list_level, &expected_level); + + // test "f" levels + let list_level = levels.get(3).unwrap(); + + let f_logical_nulls = f.logical_nulls(); + let expected_level = ArrayLevels { + def_levels: LevelData::Materialized(vec![3, 2, 3, 2, 3]), + rep_levels: LevelData::Absent, + non_null_indices: vec![0, 2, 4], + max_def_level: 3, + max_rep_level: 0, + array: Arc::new(f), + logical_nulls: f_logical_nulls, + }; + assert_eq!(list_level, &expected_level); + } + + #[test] + fn test_null_vs_nonnull_struct() { + // define schema + let offset_field = Arc::new(Field::new("offset", DataType::Int32, true)); + let schema = Schema::new(vec![Field::new( + "some_nested_object", + DataType::Struct(vec![offset_field.clone()].into()), + false, + )]); + + // create some data + let offset = Int32Array::from(vec![1, 2, 3, 4, 5]); + + let some_nested_object = + StructArray::from(vec![(offset_field, Arc::new(offset) as ArrayRef)]); + + // build a record batch + let batch = + RecordBatch::try_new(Arc::new(schema), vec![Arc::new(some_nested_object)]).unwrap(); + + let struct_null_level = + calculate_array_levels(batch.column(0), batch.schema().field(0)).unwrap(); + + // create second batch + // define schema + let offset_field = Arc::new(Field::new("offset", DataType::Int32, true)); + let schema = Schema::new(vec![Field::new( + "some_nested_object", + DataType::Struct(vec![offset_field.clone()].into()), + true, + )]); + + // create some data + let offset = Int32Array::from(vec![1, 2, 3, 4, 5]); + + let some_nested_object = + StructArray::from(vec![(offset_field, Arc::new(offset) as ArrayRef)]); + + // build a record batch + let batch = + RecordBatch::try_new(Arc::new(schema), vec![Arc::new(some_nested_object)]).unwrap(); + + let struct_non_null_level = + calculate_array_levels(batch.column(0), batch.schema().field(0)).unwrap(); + + // The 2 levels should not be the same + if struct_non_null_level == struct_null_level { + panic!("Levels should not be equal, to reflect the difference in struct nullness"); + } + } + + #[test] + fn test_map_array() { + // Note: we are using the JSON Arrow reader for brevity + let json_content = r#" + {"stocks":{"long": "$AAA", "short": "$BBB"}} + {"stocks":{"long": "$CCC", "short": null}} + {"stocks":{"hedged": "$YYY", "long": null, "short": "$D"}} + "#; + let entries_struct_type = DataType::Struct(Fields::from(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Utf8, true), + ])); + let stocks_field = Field::new( + "stocks", + DataType::Map( + Arc::new(Field::new("entries", entries_struct_type, false)), + false, + ), + // not nullable, so the keys have max level = 1 + false, + ); + let schema = Arc::new(Schema::new(vec![stocks_field])); + let builder = arrow::json::ReaderBuilder::new(schema).with_batch_size(64); + let mut reader = builder.build(std::io::Cursor::new(json_content)).unwrap(); + + let batch = reader.next().unwrap().unwrap(); + + // calculate the map's level + let mut levels = vec![]; + batch + .columns() + .iter() + .zip(batch.schema().fields()) + .for_each(|(array, field)| { + let mut array_levels = calculate_array_levels(array, field).unwrap(); + levels.append(&mut array_levels); + }); + assert_eq!(levels.len(), 2); + + let map = batch.column(0).as_map(); + let map_keys_logical_nulls = map.keys().logical_nulls(); + + // test key levels + let list_level = &levels[0]; + + let expected_level = ArrayLevels { + def_levels: LevelData::Materialized(vec![1; 7]), + rep_levels: LevelData::Materialized(vec![0, 1, 0, 1, 0, 1, 1]), + non_null_indices: vec![0, 1, 2, 3, 4, 5, 6], + max_def_level: 1, + max_rep_level: 1, + array: map.keys().clone(), + logical_nulls: map_keys_logical_nulls, + }; + assert_eq!(list_level, &expected_level); + + // test values levels + let list_level = levels.get(1).unwrap(); + let map_values_logical_nulls = map.values().logical_nulls(); + + let expected_level = ArrayLevels { + def_levels: LevelData::Materialized(vec![2, 2, 2, 1, 2, 1, 2]), + rep_levels: LevelData::Materialized(vec![0, 1, 0, 1, 0, 1, 1]), + non_null_indices: vec![0, 1, 2, 4, 6], + max_def_level: 2, + max_rep_level: 1, + array: map.values().clone(), + logical_nulls: map_values_logical_nulls, + }; + assert_eq!(list_level, &expected_level); + } + + #[test] + fn test_list_of_struct() { + // define schema + let int_field = Field::new("a", DataType::Int32, true); + let fields = Fields::from([Arc::new(int_field)]); + let item_field = Field::new_list_field(DataType::Struct(fields.clone()), true); + let list_field = Field::new("list", DataType::List(Arc::new(item_field)), true); + + let int_builder = Int32Builder::with_capacity(10); + let struct_builder = StructBuilder::new(fields, vec![Box::new(int_builder)]); + let mut list_builder = ListBuilder::new(struct_builder); + + // [{a: 1}], [], null, [null, null], [{a: null}], [{a: 2}] + // + // [{a: 1}] + let values = list_builder.values(); + values + .field_builder::(0) + .unwrap() + .append_value(1); + values.append(true); + list_builder.append(true); + + // [] + list_builder.append(true); + + // null + list_builder.append(false); + + // [null, null] + let values = list_builder.values(); + values + .field_builder::(0) + .unwrap() + .append_null(); + values.append(false); + values + .field_builder::(0) + .unwrap() + .append_null(); + values.append(false); + list_builder.append(true); + + // [{a: null}] + let values = list_builder.values(); + values + .field_builder::(0) + .unwrap() + .append_null(); + values.append(true); + list_builder.append(true); + + // [{a: 2}] + let values = list_builder.values(); + values + .field_builder::(0) + .unwrap() + .append_value(2); + values.append(true); + list_builder.append(true); + + let array = Arc::new(list_builder.finish()); + + let values = array.values().as_struct().column(0).clone(); + let values_len = values.len(); + assert_eq!(values_len, 5); + + let schema = Arc::new(Schema::new(vec![list_field])); + + let rb = RecordBatch::try_new(schema, vec![array]).unwrap(); + + let levels = calculate_array_levels(rb.column(0), rb.schema().field(0)).unwrap(); + let list_level = &levels[0]; + + let logical_nulls = values.logical_nulls(); + let expected_level = ArrayLevels { + def_levels: LevelData::Materialized(vec![4, 1, 0, 2, 2, 3, 4]), + rep_levels: LevelData::Materialized(vec![0, 0, 0, 0, 1, 0, 0]), + non_null_indices: vec![0, 4], + max_def_level: 4, + max_rep_level: 1, + array: values, + logical_nulls, + }; + + assert_eq!(list_level, &expected_level); + } + + #[test] + fn test_struct_mask_list() { + // Test the null mask of a struct array masking out non-empty slices of a child ListArray + let inner = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2)]), + Some(vec![None]), + Some(vec![]), + Some(vec![Some(3), None]), // Masked by struct array + Some(vec![Some(4), Some(5)]), + None, // Masked by struct array + None, + ]); + let values = inner.values().clone(); + + // This test assumes that nulls don't take up space + assert_eq!(inner.values().len(), 7); + + let field = Arc::new(Field::new("list", inner.data_type().clone(), true)); + let array = Arc::new(inner) as ArrayRef; + let nulls = Buffer::from([0b01010111]); + let struct_a = StructArray::from((vec![(field, array)], nulls)); + + let field = Field::new("struct", struct_a.data_type().clone(), true); + let array = Arc::new(struct_a) as ArrayRef; + let levels = calculate_array_levels(&array, &field).unwrap(); + + assert_eq!(levels.len(), 1); + + let logical_nulls = values.logical_nulls(); + let expected_level = ArrayLevels { + def_levels: LevelData::Materialized(vec![4, 4, 3, 2, 0, 4, 4, 0, 1]), + rep_levels: LevelData::Materialized(vec![0, 1, 0, 0, 0, 0, 1, 0, 0]), + non_null_indices: vec![0, 1, 5, 6], + max_def_level: 4, + max_rep_level: 1, + array: values, + logical_nulls, + }; + + assert_eq!(&levels[0], &expected_level); + } + + #[test] + fn test_list_mask_struct() { + // Test the null mask of a struct array and the null mask of a list array + // masking out non-null elements of their children + + let a1 = ListArray::from_iter_primitive::(vec![ + Some(vec![None]), // Masked by list array + Some(vec![]), // Masked by list array + Some(vec![Some(3), None]), + Some(vec![Some(4), Some(5), None, Some(6)]), // Masked by struct array + None, + None, + ]); + let a1_values = a1.values().clone(); + let a1 = Arc::new(a1) as ArrayRef; + + let a2 = Arc::new(Int32Array::from_iter(vec![ + Some(1), // Masked by list array + Some(2), // Masked by list array + None, + Some(4), // Masked by struct array + Some(5), + None, + ])) as ArrayRef; + let a2_values = a2.clone(); + + let field_a1 = Arc::new(Field::new("list", a1.data_type().clone(), true)); + let field_a2 = Arc::new(Field::new("integers", a2.data_type().clone(), true)); + + let nulls = Buffer::from([0b00110111]); + let struct_a = Arc::new(StructArray::from(( + vec![(field_a1, a1), (field_a2, a2)], + nulls, + ))) as ArrayRef; + + let offsets = Buffer::from_iter([0_i32, 0, 2, 2, 3, 5, 5]); + let nulls = Buffer::from([0b00111100]); + + let list_type = DataType::List(Arc::new(Field::new( + "struct", + struct_a.data_type().clone(), + true, + ))); + + let data = ArrayDataBuilder::new(list_type.clone()) + .len(6) + .null_bit_buffer(Some(nulls)) + .add_buffer(offsets) + .add_child_data(struct_a.into_data()) + .build() + .unwrap(); + + let list = make_array(data); + let list_field = Field::new("col", list_type, true); + + let expected = vec![ + r#""#.to_string(), + r#""#.to_string(), + r#"[]"#.to_string(), + r#"[{list: [3, ], integers: }]"#.to_string(), + r#"[, {list: , integers: 5}]"#.to_string(), + r#"[]"#.to_string(), + ]; + + let actual: Vec<_> = (0..6) + .map(|x| array_value_to_string(&list, x).unwrap()) + .collect(); + assert_eq!(actual, expected); + + let levels = calculate_array_levels(&list, &list_field).unwrap(); + + assert_eq!(levels.len(), 2); + + let a1_logical_nulls = a1_values.logical_nulls(); + let expected_level = ArrayLevels { + def_levels: LevelData::Materialized(vec![0, 0, 1, 6, 5, 2, 3, 1]), + rep_levels: LevelData::Materialized(vec![0, 0, 0, 0, 2, 0, 1, 0]), + non_null_indices: vec![1], + max_def_level: 6, + max_rep_level: 2, + array: a1_values, + logical_nulls: a1_logical_nulls, + }; + + assert_eq!(&levels[0], &expected_level); + + let a2_logical_nulls = a2_values.logical_nulls(); + let expected_level = ArrayLevels { + def_levels: LevelData::Materialized(vec![0, 0, 1, 3, 2, 4, 1]), + rep_levels: LevelData::Materialized(vec![0, 0, 0, 0, 0, 1, 0]), + non_null_indices: vec![4], + max_def_level: 4, + max_rep_level: 1, + array: a2_values, + logical_nulls: a2_logical_nulls, + }; + + assert_eq!(&levels[1], &expected_level); + } + + #[test] + fn test_fixed_size_list() { + // [[1, 2], null, null, [7, 8], null] + let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 2); + builder.values().append_slice(&[1, 2]); + builder.append(true); + builder.values().append_slice(&[3, 4]); + builder.append(false); + builder.values().append_slice(&[5, 6]); + builder.append(false); + builder.values().append_slice(&[7, 8]); + builder.append(true); + builder.values().append_slice(&[9, 10]); + builder.append(false); + let a = builder.finish(); + let values = a.values().clone(); + + let item_field = Field::new_list_field(a.data_type().clone(), true); + let mut builder = levels(&item_field, a); + builder.write(1..4); + let levels = builder.finish(); + + assert_eq!(levels.len(), 1); + + let list_level = &levels[0]; + + let logical_nulls = values.logical_nulls(); + let expected_level = ArrayLevels { + def_levels: LevelData::Materialized(vec![0, 0, 3, 3]), + rep_levels: LevelData::Materialized(vec![0, 0, 0, 1]), + non_null_indices: vec![6, 7], + max_def_level: 3, + max_rep_level: 1, + array: values, + logical_nulls, + }; + assert_eq!(list_level, &expected_level); + } + + #[test] + fn test_fixed_size_list_of_struct() { + // define schema + let field_a = Field::new("a", DataType::Int32, true); + let field_b = Field::new("b", DataType::Int64, false); + let fields = Fields::from([Arc::new(field_a), Arc::new(field_b)]); + let item_field = Field::new_list_field(DataType::Struct(fields.clone()), true); + let list_field = Field::new( + "list", + DataType::FixedSizeList(Arc::new(item_field), 2), + true, + ); + + let builder_a = Int32Builder::with_capacity(10); + let builder_b = Int64Builder::with_capacity(10); + let struct_builder = + StructBuilder::new(fields, vec![Box::new(builder_a), Box::new(builder_b)]); + let mut list_builder = FixedSizeListBuilder::new(struct_builder, 2); + + // [ + // [{a: 1, b: 2}, null], + // null, + // [null, null], + // [{a: null, b: 3}, {a: 2, b: 4}] + // ] + + // [{a: 1, b: 2}, null] + let values = list_builder.values(); + // {a: 1, b: 2} + values + .field_builder::(0) + .unwrap() + .append_value(1); + values + .field_builder::(1) + .unwrap() + .append_value(2); + values.append(true); + // null + values + .field_builder::(0) + .unwrap() + .append_null(); + values + .field_builder::(1) + .unwrap() + .append_value(0); + values.append(false); + list_builder.append(true); + + // null + let values = list_builder.values(); + // null + values + .field_builder::(0) + .unwrap() + .append_null(); + values + .field_builder::(1) + .unwrap() + .append_value(0); + values.append(false); + // null + values + .field_builder::(0) + .unwrap() + .append_null(); + values + .field_builder::(1) + .unwrap() + .append_value(0); + values.append(false); + list_builder.append(false); + + // [null, null] + let values = list_builder.values(); + // null + values + .field_builder::(0) + .unwrap() + .append_null(); + values + .field_builder::(1) + .unwrap() + .append_value(0); + values.append(false); + // null + values + .field_builder::(0) + .unwrap() + .append_null(); + values + .field_builder::(1) + .unwrap() + .append_value(0); + values.append(false); + list_builder.append(true); + + // [{a: null, b: 3}, {a: 2, b: 4}] + let values = list_builder.values(); + // {a: null, b: 3} + values + .field_builder::(0) + .unwrap() + .append_null(); + values + .field_builder::(1) + .unwrap() + .append_value(3); + values.append(true); + // {a: 2, b: 4} + values + .field_builder::(0) + .unwrap() + .append_value(2); + values + .field_builder::(1) + .unwrap() + .append_value(4); + values.append(true); + list_builder.append(true); + + let array = Arc::new(list_builder.finish()); + + assert_eq!(array.values().len(), 8); + assert_eq!(array.len(), 4); + + let struct_values = array.values().as_struct(); + let values_a = struct_values.column(0).clone(); + let values_b = struct_values.column(1).clone(); + + let schema = Arc::new(Schema::new(vec![list_field])); + let rb = RecordBatch::try_new(schema, vec![array]).unwrap(); + + let levels = calculate_array_levels(rb.column(0), rb.schema().field(0)).unwrap(); + let a_levels = &levels[0]; + let b_levels = &levels[1]; + + // [[{a: 1}, null], null, [null, null], [{a: null}, {a: 2}]] + let values_a_logical_nulls = values_a.logical_nulls(); + let expected_a = ArrayLevels { + def_levels: LevelData::Materialized(vec![4, 2, 0, 2, 2, 3, 4]), + rep_levels: LevelData::Materialized(vec![0, 1, 0, 0, 1, 0, 1]), + non_null_indices: vec![0, 7], + max_def_level: 4, + max_rep_level: 1, + array: values_a, + logical_nulls: values_a_logical_nulls, + }; + // [[{b: 2}, null], null, [null, null], [{b: 3}, {b: 4}]] + let values_b_logical_nulls = values_b.logical_nulls(); + let expected_b = ArrayLevels { + def_levels: LevelData::Materialized(vec![3, 2, 0, 2, 2, 3, 3]), + rep_levels: LevelData::Materialized(vec![0, 1, 0, 0, 1, 0, 1]), + non_null_indices: vec![0, 6, 7], + max_def_level: 3, + max_rep_level: 1, + array: values_b, + logical_nulls: values_b_logical_nulls, + }; + + assert_eq!(a_levels, &expected_a); + assert_eq!(b_levels, &expected_b); + } + + #[test] + fn test_fixed_size_list_empty() { + let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 0); + builder.append(true); + builder.append(false); + builder.append(true); + let array = builder.finish(); + let values = array.values().clone(); + + let item_field = Field::new_list_field(array.data_type().clone(), true); + let mut builder = levels(&item_field, array); + builder.write(0..3); + let levels = builder.finish(); + + assert_eq!(levels.len(), 1); + + let list_level = &levels[0]; + + let logical_nulls = values.logical_nulls(); + let expected_level = ArrayLevels { + def_levels: LevelData::Materialized(vec![1, 0, 1]), + rep_levels: LevelData::Materialized(vec![0, 0, 0]), + non_null_indices: vec![], + max_def_level: 3, + max_rep_level: 1, + array: values, + logical_nulls, + }; + assert_eq!(list_level, &expected_level); + } + + #[test] + fn test_fixed_size_list_of_var_lists() { + // [[[1, null, 3], null], [[4], []], [[5, 6], [null, null]], null] + let mut builder = FixedSizeListBuilder::new(ListBuilder::new(Int32Builder::new()), 2); + builder.values().append_value([Some(1), None, Some(3)]); + builder.values().append_null(); + builder.append(true); + builder.values().append_value([Some(4)]); + builder.values().append_value([]); + builder.append(true); + builder.values().append_value([Some(5), Some(6)]); + builder.values().append_value([None, None]); + builder.append(true); + builder.values().append_null(); + builder.values().append_null(); + builder.append(false); + let a = builder.finish(); + let values = a.values().as_list::().values().clone(); + + let item_field = Field::new_list_field(a.data_type().clone(), true); + let mut builder = levels(&item_field, a); + builder.write(0..4); + let levels = builder.finish(); + + let logical_nulls = values.logical_nulls(); + let expected_level = ArrayLevels { + def_levels: LevelData::Materialized(vec![5, 4, 5, 2, 5, 3, 5, 5, 4, 4, 0]), + rep_levels: LevelData::Materialized(vec![0, 2, 2, 1, 0, 1, 0, 2, 1, 2, 0]), + non_null_indices: vec![0, 2, 3, 4, 5], + max_def_level: 5, + max_rep_level: 2, + array: values, + logical_nulls, + }; + + assert_eq!(levels[0], expected_level); + } + + #[test] + fn test_null_dictionary_values() { + let values = Int32Array::new( + vec![1, 2, 3, 4].into(), + Some(NullBuffer::from(vec![true, false, true, true])), + ); + let keys = Int32Array::new( + vec![1, 54, 2, 0].into(), + Some(NullBuffer::from(vec![true, false, true, true])), + ); + // [NULL, NULL, 3, 0] + let dict = DictionaryArray::new(keys, Arc::new(values)); + + let item_field = Field::new_list_field(dict.data_type().clone(), true); + + let mut builder = levels(&item_field, dict.clone()); + builder.write(0..4); + let levels = builder.finish(); + + let logical_nulls = dict.logical_nulls(); + let expected_level = ArrayLevels { + def_levels: LevelData::Materialized(vec![0, 0, 1, 1]), + rep_levels: LevelData::Absent, + non_null_indices: vec![2, 3], + max_def_level: 1, + max_rep_level: 0, + array: Arc::new(dict), + logical_nulls, + }; + assert_eq!(levels[0], expected_level); + } + + #[test] + fn mismatched_types() { + let array = Arc::new(Int32Array::from_iter(0..10)) as ArrayRef; + let field = Field::new_list_field(DataType::Float64, false); + + let err = LevelInfoBuilder::try_new(&field, Default::default(), &array) + .unwrap_err() + .to_string(); + + assert_eq!( + err, + "Arrow: Incompatible type. Field 'item' has type Float64, array has type Int32", + ); + } + + fn levels(field: &Field, array: T) -> LevelInfoBuilder { + let v = Arc::new(array) as ArrayRef; + LevelInfoBuilder::try_new(field, Default::default(), &v).unwrap() + } + + #[test] + fn test_slice_for_chunk_flat() { + // Case 1: required field (max_def_level=0, no def/rep levels stored). + // Array has 6 values; all are non-null so non_null_indices covers every position. + // value_offset=2, num_values=3 → non_null_indices[2..5] = [2,3,4]. + // Array is sliced (no def_levels → write_batch_internal uses values.len()). + let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6])); + let logical_nulls = array.logical_nulls(); + let levels = ArrayLevels { + def_levels: LevelData::Absent, + rep_levels: LevelData::Absent, + non_null_indices: vec![0, 1, 2, 3, 4, 5], + max_def_level: 0, + max_rep_level: 0, + array, + logical_nulls, + }; + let sliced = levels.slice_for_chunk(&CdcChunk { + level_offset: 0, + num_levels: 0, + value_offset: 2, + num_values: 3, + }); + assert!(matches!(sliced.def_levels, LevelData::Absent)); + assert!(matches!(sliced.rep_levels, LevelData::Absent)); + assert_eq!(sliced.non_null_indices, vec![0, 1, 2]); + assert_eq!(sliced.array.len(), 3); + + // Case 2: optional field (max_def_level=1, def levels present, no rep levels). + // Array: [Some(1), None, Some(3), None, Some(5), Some(6)] + // non_null_indices: [0, 2, 4, 5] + // value_offset=1, num_values=1 → non_null_indices[1..2] = [2]. + // Array is not sliced (def_levels present → num_levels from def_levels.len()). + let array: ArrayRef = Arc::new(Int32Array::from(vec![ + Some(1), + None, + Some(3), + None, + Some(5), + Some(6), + ])); + let logical_nulls = array.logical_nulls(); + let levels = ArrayLevels { + def_levels: LevelData::Materialized(vec![1, 0, 1, 0, 1, 1]), + rep_levels: LevelData::Absent, + non_null_indices: vec![0, 2, 4, 5], + max_def_level: 1, + max_rep_level: 0, + array, + logical_nulls, + }; + let sliced = levels.slice_for_chunk(&CdcChunk { + level_offset: 1, + num_levels: 3, + value_offset: 1, + num_values: 1, + }); + assert_eq!(sliced.def_levels, LevelData::Materialized(vec![0, 1, 0])); + assert!(matches!(sliced.rep_levels, LevelData::Absent)); + assert_eq!(sliced.non_null_indices, vec![0]); // [2] shifted by -2 (nni[0]) + assert_eq!(sliced.array.len(), 1); + } + + #[test] + fn test_slice_for_chunk_nested_with_nulls() { + // Regression test for https://github.com/apache/arrow-rs/issues/9637 + // + // Simulates a List where null list entries have non-zero child + // ranges (valid per Arrow spec: "a null value may correspond to a + // non-empty segment in the child array"). This creates gaps in the + // leaf array that don't correspond to any levels. + // + // 5 rows with 2 null list entries owning non-empty child ranges: + // row 0: [1] → leaf[0] + // row 1: null list → owns leaf[1..3] (gap of 2) + // row 2: [2, null] → leaf[3], leaf[4]=null element + // row 3: null list → owns leaf[5..8] (gap of 3) + // row 4: [4, 5] → leaf[8], leaf[9] + // + // def_levels: [3, 0, 3, 2, 0, 3, 3] + // rep_levels: [0, 0, 0, 1, 0, 0, 1] + // non_null_indices: [0, 3, 8, 9] + // gaps in array: 0→3 (skip 1,2), 3→8 (skip 5,6,7) + let array: ArrayRef = Arc::new(Int32Array::from(vec![ + Some(1), // 0: row 0 + None, // 1: gap (null list row 1) + None, // 2: gap (null list row 1) + Some(2), // 3: row 2 + None, // 4: row 2, null element + None, // 5: gap (null list row 3) + None, // 6: gap (null list row 3) + None, // 7: gap (null list row 3) + Some(4), // 8: row 4 + Some(5), // 9: row 4 + ])); + let logical_nulls = array.logical_nulls(); + let levels = ArrayLevels { + def_levels: LevelData::Materialized(vec![3, 0, 3, 2, 0, 3, 3]), + rep_levels: LevelData::Materialized(vec![0, 0, 0, 1, 0, 0, 1]), + non_null_indices: vec![0, 3, 8, 9], + max_def_level: 3, + max_rep_level: 1, + array, + logical_nulls, + }; + + // Chunk 0: rows 0-1, nni=[0] → array sliced to [0..1] + let chunk0 = levels.slice_for_chunk(&CdcChunk { + level_offset: 0, + num_levels: 2, + value_offset: 0, + num_values: 1, + }); + assert_eq!(chunk0.non_null_indices, vec![0]); + assert_eq!(chunk0.array.len(), 1); + + // Chunk 1: rows 2-3, nni=[3] → array sliced to [3..4] + let chunk1 = levels.slice_for_chunk(&CdcChunk { + level_offset: 2, + num_levels: 3, + value_offset: 1, + num_values: 1, + }); + assert_eq!(chunk1.non_null_indices, vec![0]); + assert_eq!(chunk1.array.len(), 1); + + // Chunk 2: row 4, nni=[8, 9] → array sliced to [8..10] + let chunk2 = levels.slice_for_chunk(&CdcChunk { + level_offset: 5, + num_levels: 2, + value_offset: 2, + num_values: 2, + }); + assert_eq!(chunk2.non_null_indices, vec![0, 1]); + assert_eq!(chunk2.array.len(), 2); + } + + #[test] + fn test_slice_for_chunk_all_null() { + // All-null chunk: num_values=0 → empty nni slice → zero-length array. + let array: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, None, Some(4)])); + let logical_nulls = array.logical_nulls(); + let levels = ArrayLevels { + def_levels: LevelData::Materialized(vec![1, 0, 0, 1]), + rep_levels: LevelData::Absent, + non_null_indices: vec![0, 3], + max_def_level: 1, + max_rep_level: 0, + array, + logical_nulls, + }; + // Chunk covering only the two null rows (levels 1..3), zero non-null values. + let sliced = levels.slice_for_chunk(&CdcChunk { + level_offset: 1, + num_levels: 2, + value_offset: 1, + num_values: 0, + }); + assert_eq!(sliced.def_levels, LevelData::Materialized(vec![0, 0])); + assert_eq!(sliced.non_null_indices, Vec::::new()); + assert_eq!(sliced.array.len(), 0); + } + + #[test] + fn test_all_null_list() { + // List where every list slot is null. + // Schema: list (nullable) -> item (int32, nullable) + // Data: [null, null, null, null] + // + // Expected: max_def=3, max_rep=1, def/rep levels all 0. + let item_field = Arc::new(Field::new_list_field(DataType::Int32, true)); + let list = ListArray::new_null(item_field, 4); + let values = list.values().clone(); + let field = Field::new("list", list.data_type().clone(), true); + let array = Arc::new(list) as ArrayRef; + + let levels = calculate_array_levels(&array, &field).unwrap(); + assert_eq!(levels.len(), 1); + + let logical_nulls = values.logical_nulls(); + let expected = ArrayLevels { + def_levels: LevelData::Uniform { value: 0, count: 4 }, + rep_levels: LevelData::Uniform { value: 0, count: 4 }, + non_null_indices: vec![], + max_def_level: 3, + max_rep_level: 1, + array: values, + logical_nulls, + }; + assert_eq!(&levels[0], &expected); + } + + #[test] + fn test_all_null_fixed_size_list() { + // FixedSizeList where every list slot is null. + // Schema: list (nullable) -> item (int32, nullable) + // Data: [null, null, null] + // + // Expected: max_def=3, max_rep=1, def/rep levels all 0. + let item_field = Arc::new(Field::new_list_field(DataType::Int32, true)); + let list = FixedSizeListArray::new_null(item_field, 2, 3); + let values = list.values().clone(); + let field = Field::new("list", list.data_type().clone(), true); + let array = Arc::new(list) as ArrayRef; + + let levels = calculate_array_levels(&array, &field).unwrap(); + assert_eq!(levels.len(), 1); + + let logical_nulls = values.logical_nulls(); + let expected = ArrayLevels { + def_levels: LevelData::Uniform { value: 0, count: 3 }, + rep_levels: LevelData::Uniform { value: 0, count: 3 }, + non_null_indices: vec![], + max_def_level: 3, + max_rep_level: 1, + array: values, + logical_nulls, + }; + assert_eq!(&levels[0], &expected); + } + + #[test] + fn test_all_null_struct() { + // Struct where every struct slot is null. + // Schema: a (struct, nullable) -> c (int32, nullable) + // Data: [null, null, null, null] + // + // Expected: max_def=2, def_levels all 0 (struct is null → child never reached), + // leaf values are empty. + let c = Int32Array::from(vec![None::; 4]); + let leaf = Arc::new(c) as ArrayRef; + let c_field = Arc::new(Field::new("c", DataType::Int32, true)); + let a = StructArray::from((vec![(c_field, leaf.clone())], Buffer::from([0b00000000]))); + let a_field = Field::new("a", a.data_type().clone(), true); + let a_array = Arc::new(a) as ArrayRef; + + let levels = calculate_array_levels(&a_array, &a_field).unwrap(); + assert_eq!(levels.len(), 1); + + let expected = ArrayLevels { + def_levels: LevelData::Uniform { value: 0, count: 4 }, + rep_levels: LevelData::Absent, + non_null_indices: vec![], + max_def_level: 2, + max_rep_level: 0, + array: leaf, + logical_nulls: Some(NullBuffer::new_null(4)), + }; + assert_eq!(&levels[0], &expected); + } + + #[test] + fn test_all_null_nested_struct() { + // Struct> where the outer struct is entirely null. + // Schema: a (struct, nullable) -> b (struct, nullable) -> c (int32, nullable) + // Data: [null, null, null] + // + // Expected: max_def=3, def_levels all 0. + let c = Int32Array::from(vec![None::; 3]); + let leaf = Arc::new(c) as ArrayRef; + let c_field = Arc::new(Field::new("c", DataType::Int32, true)); + let b = StructArray::from((vec![(c_field, leaf.clone())], Buffer::from([0b00000000]))); + let b_field = Arc::new(Field::new("b", b.data_type().clone(), true)); + let a = StructArray::from(( + vec![(b_field, Arc::new(b) as ArrayRef)], + Buffer::from([0b00000000]), + )); + let a_field = Field::new("a", a.data_type().clone(), true); + let a_array = Arc::new(a) as ArrayRef; + + let levels = calculate_array_levels(&a_array, &a_field).unwrap(); + assert_eq!(levels.len(), 1); + + let expected = ArrayLevels { + def_levels: LevelData::Uniform { value: 0, count: 3 }, + rep_levels: LevelData::Absent, + non_null_indices: vec![], + max_def_level: 3, + max_rep_level: 0, + array: leaf, + logical_nulls: Some(NullBuffer::new_null(3)), + }; + assert_eq!(&levels[0], &expected); + } + + #[test] + fn test_all_null_struct_multiple_children() { + // Struct with two leaf children, entirely null. + // Schema: a (struct, nullable) -> { c1 (int32, nullable), c2 (int32, nullable) } + // Data: [null, null] + // + // Both leaf columns should get uniform def_levels=0. + let c1 = Arc::new(Int32Array::from(vec![None::; 2])) as ArrayRef; + let c2 = Arc::new(Int32Array::from(vec![None::; 2])) as ArrayRef; + let c1_field = Arc::new(Field::new("c1", DataType::Int32, true)); + let c2_field = Arc::new(Field::new("c2", DataType::Int32, true)); + let a = StructArray::from(( + vec![(c1_field, c1.clone()), (c2_field, c2.clone())], + Buffer::from([0b00000000]), + )); + let a_field = Field::new("a", a.data_type().clone(), true); + let a_array = Arc::new(a) as ArrayRef; + + let levels = calculate_array_levels(&a_array, &a_field).unwrap(); + assert_eq!(levels.len(), 2); + + for (i, leaf) in [c1, c2].into_iter().enumerate() { + let expected = ArrayLevels { + def_levels: LevelData::Uniform { value: 0, count: 2 }, + rep_levels: LevelData::Absent, + non_null_indices: vec![], + max_def_level: 2, + max_rep_level: 0, + array: leaf, + logical_nulls: Some(NullBuffer::new_null(2)), + }; + assert_eq!(&levels[i], &expected, "leaf {i} mismatch"); + } + } +} diff --git a/vendor/parquet/src/arrow/arrow_writer/mod.rs b/vendor/parquet/src/arrow/arrow_writer/mod.rs new file mode 100644 index 0000000..8b9aea2 --- /dev/null +++ b/vendor/parquet/src/arrow/arrow_writer/mod.rs @@ -0,0 +1,6019 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Contains writer which writes arrow data into parquet data. + +use crate::column::chunker::ContentDefinedChunker; + +use bytes::Bytes; +use std::io::Write; +use std::slice::Iter; +use std::sync::{Arc, Mutex}; +use std::vec::IntoIter; + +use arrow_array::cast::AsArray; +use arrow_array::{ArrayRef, Int32Array, RecordBatch, RecordBatchWriter}; +use arrow_array::{PrimitiveArray, types::*}; +use arrow_schema::{ + ArrowError, DataType as ArrowDataType, Field, IntervalUnit, SchemaRef, TimeUnit, +}; + +use super::schema::{add_encoded_arrow_schema_to_metadata, decimal_length_from_precision}; + +use crate::arrow::ArrowSchemaConverter; +use crate::arrow::arrow_writer::byte_array::ByteArrayEncoder; +use crate::basic::PageType; +use crate::column::page::{CompressedPage, PageWriteSpec, PageWriter}; +use crate::column::page_encryption::PageEncryptor; +use crate::column::writer::encoder::ColumnValueEncoder; +use crate::column::writer::{ + ColumnCloseResult, ColumnWriter, GenericColumnWriter, get_column_writer, +}; +use crate::data_type::{ByteArray, FixedLenByteArray}; +#[cfg(feature = "encryption")] +use crate::encryption::encrypt::FileEncryptor; +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::{KeyValue, ParquetMetaData, RowGroupMetaData}; +use crate::file::properties::{WriterProperties, WriterPropertiesPtr}; +use crate::file::writer::{SerializedFileWriter, SerializedRowGroupWriter}; +use crate::parquet_thrift::{ThriftCompactOutputProtocol, WriteThrift}; +use crate::schema::types::{ColumnDescPtr, SchemaDescPtr, SchemaDescriptor}; +use levels::{ArrayLevels, calculate_array_levels}; + +mod byte_array; +mod levels; + +#[doc(inline)] +pub use crate::column::page_store::{ + InMemoryPageStore, InMemoryPageStoreFactory, PageKey, PageStore, PageStoreArgs, + PageStoreFactory, +}; + +/// Encodes [`RecordBatch`] to parquet +/// +/// Writes Arrow `RecordBatch`es to a Parquet writer. Multiple [`RecordBatch`] will be encoded +/// to the same row group, up to `max_row_group_size` rows. Any remaining rows will be +/// flushed on close, leading the final row group in the output file to potentially +/// contain fewer than `max_row_group_size` rows +/// +/// # Example: Writing `RecordBatch`es +/// ``` +/// # use std::sync::Arc; +/// # use bytes::Bytes; +/// # use arrow_array::{ArrayRef, Int64Array}; +/// # use arrow_array::RecordBatch; +/// # use parquet::arrow::arrow_writer::ArrowWriter; +/// # use parquet::arrow::arrow_reader::ParquetRecordBatchReader; +/// let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef; +/// let to_write = RecordBatch::try_from_iter([("col", col)]).unwrap(); +/// +/// let mut buffer = Vec::new(); +/// let mut writer = ArrowWriter::try_new(&mut buffer, to_write.schema(), None).unwrap(); +/// writer.write(&to_write).unwrap(); +/// writer.close().unwrap(); +/// +/// let mut reader = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 1024).unwrap(); +/// let read = reader.next().unwrap().unwrap(); +/// +/// assert_eq!(to_write, read); +/// ``` +/// +/// # Memory Usage and Limiting +/// +/// The nature of Parquet requires buffering of an entire row group before it can +/// be flushed to the underlying writer. Data is mostly buffered in its encoded +/// form, reducing memory usage. However, some data such as dictionary keys, +/// large strings or very nested data may still result in non-trivial memory +/// usage. +/// +/// See Also: +/// * [`ArrowWriter::memory_size`]: the current memory usage of the writer. +/// * [`ArrowWriter::in_progress_size`]: Estimated size of the buffered row group, +/// +/// Call [`Self::flush`] to trigger an early flush of a row group based on a +/// memory threshold and/or global memory pressure. However, smaller row groups +/// result in higher metadata overheads, and thus may worsen compression ratios +/// and query performance. +/// +/// ```no_run +/// # use std::io::Write; +/// # use arrow_array::RecordBatch; +/// # use parquet::arrow::ArrowWriter; +/// # let mut writer: ArrowWriter> = todo!(); +/// # let batch: RecordBatch = todo!(); +/// writer.write(&batch).unwrap(); +/// // Trigger an early flush if anticipated size exceeds 1_000_000 +/// if writer.in_progress_size() > 1_000_000 { +/// writer.flush().unwrap(); +/// } +/// ``` +/// +/// ## Type Support +/// +/// The writer supports writing all Arrow [`DataType`]s that have a direct mapping to +/// Parquet types including [`StructArray`] and [`ListArray`]. +/// +/// The following are not supported: +/// +/// * [`IntervalMonthDayNanoArray`]: Parquet does not [support nanosecond intervals]. +/// +/// [`DataType`]: https://docs.rs/arrow/latest/arrow/datatypes/enum.DataType.html +/// [`StructArray`]: https://docs.rs/arrow/latest/arrow/array/struct.StructArray.html +/// [`ListArray`]: https://docs.rs/arrow/latest/arrow/array/type.ListArray.html +/// [`IntervalMonthDayNanoArray`]: https://docs.rs/arrow/latest/arrow/array/type.IntervalMonthDayNanoArray.html +/// [support nanosecond intervals]: https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#interval +/// +/// ## Type Compatibility +/// The writer can write Arrow [`RecordBatch`]s that are logically equivalent. This means that for +/// a given column, the writer can accept multiple Arrow [`DataType`]s that contain the same +/// value type. +/// +/// For example, the following [`DataType`]s are all logically equivalent and can be written +/// to the same column: +/// * String, LargeString, StringView +/// * Binary, LargeBinary, BinaryView +/// +/// The writer can will also accept both native and dictionary encoded arrays if the dictionaries +/// contain compatible values. +/// ``` +/// # use std::sync::Arc; +/// # use arrow_array::{DictionaryArray, LargeStringArray, RecordBatch, StringArray, UInt8Array}; +/// # use arrow_schema::{DataType, Field, Schema}; +/// # use parquet::arrow::arrow_writer::ArrowWriter; +/// let record_batch1 = RecordBatch::try_new( +/// Arc::new(Schema::new(vec![Field::new("col", DataType::LargeUtf8, false)])), +/// vec![Arc::new(LargeStringArray::from_iter_values(vec!["a", "b"]))] +/// ) +/// .unwrap(); +/// +/// let mut buffer = Vec::new(); +/// let mut writer = ArrowWriter::try_new(&mut buffer, record_batch1.schema(), None).unwrap(); +/// writer.write(&record_batch1).unwrap(); +/// +/// let record_batch2 = RecordBatch::try_new( +/// Arc::new(Schema::new(vec![Field::new( +/// "col", +/// DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)), +/// false, +/// )])), +/// vec![Arc::new(DictionaryArray::new( +/// UInt8Array::from_iter_values(vec![0, 1]), +/// Arc::new(StringArray::from_iter_values(vec!["b", "c"])), +/// ))], +/// ) +/// .unwrap(); +/// writer.write(&record_batch2).unwrap(); +/// writer.close(); +/// ``` +pub struct ArrowWriter { + /// Underlying Parquet writer + writer: SerializedFileWriter, + + /// The in-progress row group if any + in_progress: Option, + + /// A copy of the Arrow schema. + /// + /// The schema is used to verify that each record batch written has the correct schema + arrow_schema: SchemaRef, + + /// Creates new [`ArrowRowGroupWriter`] instances as required + row_group_writer_factory: ArrowRowGroupWriterFactory, + + /// The maximum number of rows to write to each row group, or None for unlimited + max_row_group_row_count: Option, + + /// The maximum size in bytes for a row group, or None for unlimited + max_row_group_bytes: Option, + + /// CDC chunkers persisted across row groups (one per leaf column). + cdc_chunkers: Option>, +} + +impl std::fmt::Debug for ArrowWriter { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let buffered_memory = self.in_progress_size(); + f.debug_struct("ArrowWriter") + .field("writer", &self.writer) + .field("in_progress_size", &format_args!("{buffered_memory} bytes")) + .field("in_progress_rows", &self.in_progress_rows()) + .field("arrow_schema", &self.arrow_schema) + .field("max_row_group_row_count", &self.max_row_group_row_count) + .field("max_row_group_bytes", &self.max_row_group_bytes) + .finish() + } +} + +impl ArrowWriter { + /// Try to create a new Arrow writer + /// + /// The writer will fail if: + /// * a `SerializedFileWriter` cannot be created from the ParquetWriter + /// * the Arrow schema contains unsupported datatypes such as Unions + pub fn try_new( + writer: W, + arrow_schema: SchemaRef, + props: Option, + ) -> Result { + let options = ArrowWriterOptions::new().with_properties(props.unwrap_or_default()); + Self::try_new_with_options(writer, arrow_schema, options) + } + + /// Try to create a new Arrow writer with [`ArrowWriterOptions`]. + /// + /// The writer will fail if: + /// * a `SerializedFileWriter` cannot be created from the ParquetWriter + /// * the Arrow schema contains unsupported datatypes such as Unions + pub fn try_new_with_options( + writer: W, + arrow_schema: SchemaRef, + options: ArrowWriterOptions, + ) -> Result { + let mut props = options.properties; + + let schema = if let Some(parquet_schema) = options.schema_descr { + parquet_schema.clone() + } else { + let mut converter = ArrowSchemaConverter::new().with_coerce_types(props.coerce_types()); + if let Some(schema_root) = &options.schema_root { + converter = converter.schema_root(schema_root); + } + + converter.convert(&arrow_schema)? + }; + + if !options.skip_arrow_metadata { + // add serialized arrow schema + add_encoded_arrow_schema_to_metadata(&arrow_schema, &mut props); + } + + let max_row_group_row_count = props.max_row_group_row_count(); + let max_row_group_bytes = props.max_row_group_bytes(); + + let props_ptr = Arc::new(props); + let file_writer = + SerializedFileWriter::new(writer, schema.root_schema_ptr(), Arc::clone(&props_ptr))?; + + let mut row_group_writer_factory = + ArrowRowGroupWriterFactory::new(&file_writer, arrow_schema.clone()); + if let Some(page_store_factory) = options.page_store_factory { + row_group_writer_factory = + row_group_writer_factory.with_page_store_factory(page_store_factory); + } + + let cdc_chunkers = props_ptr + .content_defined_chunking() + .map(|opts| { + file_writer + .schema_descr() + .columns() + .iter() + .map(|desc| ContentDefinedChunker::new(desc, opts)) + .collect::>>() + }) + .transpose()?; + + Ok(Self { + writer: file_writer, + in_progress: None, + arrow_schema, + row_group_writer_factory, + max_row_group_row_count, + max_row_group_bytes, + cdc_chunkers, + }) + } + + /// Returns metadata for any flushed row groups + pub fn flushed_row_groups(&self) -> &[RowGroupMetaData] { + self.writer.flushed_row_groups() + } + + /// Estimated memory usage, in bytes, of this `ArrowWriter` + /// + /// This estimate is formed bu summing the values of + /// [`ArrowColumnWriter::memory_size`] all in progress columns. + pub fn memory_size(&self) -> usize { + match &self.in_progress { + Some(in_progress) => in_progress.writers.iter().map(|x| x.memory_size()).sum(), + None => 0, + } + } + + /// Anticipated encoded size of the in progress row group. + /// + /// This estimate the row group size after being completely encoded is, + /// formed by summing the values of + /// [`ArrowColumnWriter::get_estimated_total_bytes`] for all in progress + /// columns. + pub fn in_progress_size(&self) -> usize { + match &self.in_progress { + Some(in_progress) => in_progress + .writers + .iter() + .map(|x| x.get_estimated_total_bytes()) + .sum(), + None => 0, + } + } + + /// Returns the number of rows buffered in the in progress row group + pub fn in_progress_rows(&self) -> usize { + self.in_progress + .as_ref() + .map(|x| x.buffered_rows) + .unwrap_or_default() + } + + /// Returns the number of bytes written by this instance + pub fn bytes_written(&self) -> usize { + self.writer.bytes_written() + } + + /// Encodes the provided [`RecordBatch`] + /// + /// If this would cause the current row group to exceed [`WriterProperties::max_row_group_row_count`] + /// rows or [`WriterProperties::max_row_group_bytes`] bytes, the contents of `batch` will be + /// written to one or more row groups such that limits are respected. + /// + /// If both limits are `None`, all data is written to a single row group. + /// If one limit is set, that limit is respected. + /// If both limits are set, the lower bound (whichever triggers first) is respected. + /// + /// This will fail if the `batch`'s schema does not match the writer's schema. + pub fn write(&mut self, batch: &RecordBatch) -> Result<()> { + if batch.num_rows() == 0 { + return Ok(()); + } + + let in_progress = match &mut self.in_progress { + Some(in_progress) => in_progress, + x => x.insert( + self.row_group_writer_factory + .create_row_group_writer(self.writer.flushed_row_groups().len())?, + ), + }; + + if let Some(max_rows) = self.max_row_group_row_count { + if in_progress.buffered_rows + batch.num_rows() > max_rows { + let to_write = max_rows - in_progress.buffered_rows; + let a = batch.slice(0, to_write); + let b = batch.slice(to_write, batch.num_rows() - to_write); + self.write(&a)?; + return self.write(&b); + } + } + + // Check byte limit: if we have buffered data, use measured average row size + // to split batch proactively before exceeding byte limit + if let Some(max_bytes) = self.max_row_group_bytes { + if in_progress.buffered_rows > 0 { + let current_bytes = in_progress.get_estimated_total_bytes(); + + if current_bytes >= max_bytes { + self.flush()?; + return self.write(batch); + } + + if let Some(avg_row_bytes) = current_bytes + .checked_div(in_progress.buffered_rows) + .filter(|avg_row_bytes| *avg_row_bytes > 0) + { + // 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 batch.num_rows() > rows_that_fit { + if rows_that_fit > 0 { + let a = batch.slice(0, rows_that_fit); + let b = batch.slice(rows_that_fit, batch.num_rows() - rows_that_fit); + self.write(&a)?; + return self.write(&b); + } else { + self.flush()?; + return self.write(batch); + } + } + } + } + } + + match self.cdc_chunkers.as_mut() { + Some(chunkers) => in_progress.write_with_chunkers(batch, chunkers)?, + None => in_progress.write(batch)?, + } + + let should_flush = self + .max_row_group_row_count + .is_some_and(|max| in_progress.buffered_rows >= max) + || self + .max_row_group_bytes + .is_some_and(|max| in_progress.get_estimated_total_bytes() >= max); + + if should_flush { + self.flush()? + } + Ok(()) + } + + /// Writes the given buf bytes to the internal buffer. + /// + /// It's safe to use this method to write data to the underlying writer, + /// because it will ensure that the buffering and byte‐counting layers are used. + pub fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> { + self.writer.write_all(buf) + } + + /// Flushes underlying writer + pub fn sync(&mut self) -> std::io::Result<()> { + self.writer.flush() + } + + /// Flushes all buffered rows into a new row group + /// + /// Note the underlying writer is not flushed with this call. + /// If this is a desired behavior, please call [`ArrowWriter::sync`]. + pub fn flush(&mut self) -> Result<()> { + let in_progress = match self.in_progress.take() { + Some(in_progress) => in_progress, + None => return Ok(()), + }; + + let mut row_group_writer = self.writer.next_row_group()?; + for chunk in in_progress.close()? { + chunk.append_to_row_group(&mut row_group_writer)?; + } + row_group_writer.close()?; + Ok(()) + } + + /// Additional [`KeyValue`] metadata to be written in addition to those from [`WriterProperties`] + /// + /// This method provide a way to append kv_metadata after write RecordBatch + pub fn append_key_value_metadata(&mut self, kv_metadata: KeyValue) { + self.writer.append_key_value_metadata(kv_metadata) + } + + /// Returns a reference to the underlying writer. + pub fn inner(&self) -> &W { + self.writer.inner() + } + + /// Returns a mutable reference to the underlying writer. + /// + /// **Warning**: if you write directly to this writer, you will skip + /// the `TrackedWrite` buffering and byte‐counting layers. That’ll cause + /// the file footer’s recorded offsets and sizes to diverge from reality, + /// resulting in an unreadable or corrupted Parquet file. + /// + /// If you want to write safely to the underlying writer, use [`Self::write_all`]. + pub fn inner_mut(&mut self) -> &mut W { + self.writer.inner_mut() + } + + /// Flushes any outstanding data and returns the underlying writer. + pub fn into_inner(mut self) -> Result { + self.flush()?; + self.writer.into_inner() + } + + /// Close and finalize the underlying Parquet writer + /// + /// Unlike [`Self::close`] this does not consume self + /// + /// Attempting to write after calling finish will result in an error + pub fn finish(&mut self) -> Result { + self.flush()?; + self.writer.finish() + } + + /// Close and finalize the underlying Parquet writer + pub fn close(mut self) -> Result { + self.finish() + } + + /// Create a new row group writer and return its column writers. + #[deprecated( + since = "56.2.0", + note = "Use `ArrowRowGroupWriterFactory` instead, see `ArrowColumnWriter` for an example" + )] + pub fn get_column_writers(&mut self) -> Result> { + self.flush()?; + let in_progress = self + .row_group_writer_factory + .create_row_group_writer(self.writer.flushed_row_groups().len())?; + Ok(in_progress.writers) + } + + /// Append the given column chunks to the file as a new row group. + #[deprecated( + since = "56.2.0", + note = "Use `SerializedFileWriter` directly instead, see `ArrowColumnWriter` for an example" + )] + pub fn append_row_group(&mut self, chunks: Vec) -> Result<()> { + let mut row_group_writer = self.writer.next_row_group()?; + for chunk in chunks { + chunk.append_to_row_group(&mut row_group_writer)?; + } + row_group_writer.close()?; + Ok(()) + } + + /// Converts this writer into a lower-level [`SerializedFileWriter`] and [`ArrowRowGroupWriterFactory`]. + /// + /// Flushes any outstanding data before returning. + /// + /// This can be useful to provide more control over how files are written, for example + /// to write columns in parallel. See the example on [`ArrowColumnWriter`]. + pub fn into_serialized_writer( + mut self, + ) -> Result<(SerializedFileWriter, ArrowRowGroupWriterFactory)> { + self.flush()?; + Ok((self.writer, self.row_group_writer_factory)) + } +} + +impl RecordBatchWriter for ArrowWriter { + fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> { + self.write(batch).map_err(|e| e.into()) + } + + fn close(self) -> std::result::Result<(), ArrowError> { + self.close()?; + Ok(()) + } +} + +/// Arrow-specific configuration settings for writing parquet files. +/// +/// See [`ArrowWriter`] for how to configure the writer. +#[derive(Debug, Clone, Default)] +pub struct ArrowWriterOptions { + properties: WriterProperties, + skip_arrow_metadata: bool, + schema_root: Option, + schema_descr: Option, + page_store_factory: Option>, +} + +impl ArrowWriterOptions { + /// Creates a new [`ArrowWriterOptions`] with the default settings. + pub fn new() -> Self { + Self::default() + } + + /// Sets the [`WriterProperties`] for writing parquet files. + pub fn with_properties(self, properties: WriterProperties) -> Self { + Self { properties, ..self } + } + + /// Sets the [`PageStoreFactory`] used to buffer completed pages while a row + /// group is being written. + /// + /// The default implementation ([`InMemoryPageStore`]) buffers all completed + /// pages on the heap until the row group is flushed, so peak write memory + /// grows with the row group size. Using this API, pages can be spilled to a + /// file or object storage instead, reducing peak write memory substantially + /// at the expense of an extra write to and read from secondary storage. + /// + /// # Example: spilling pages to a temp file + /// + /// A simple spilling backend uses one temp file per column chunk; `put` + /// appends the page and `take` reads it back. + /// + /// ``` + /// # use std::fs::File; + /// # use std::io::{Read, Seek, SeekFrom, Write}; + /// # use std::sync::Arc; + /// # use bytes::Bytes; + /// # use arrow_array::{ArrayRef, Int64Array, RecordBatch}; + /// # use parquet::arrow::arrow_writer::{ + /// # ArrowWriter, ArrowWriterOptions, PageKey, PageStore, PageStoreArgs, PageStoreFactory, + /// # }; + /// # use parquet::arrow::arrow_reader::ParquetRecordBatchReader; + /// # use parquet::errors::Result; + /// struct TempFilePageStore { + /// file: File, + /// /// Total size of the file + /// end: u64, + /// /// Location of pages: (offset, len) + /// locs: Vec<(u64, usize)>, + /// } + /// + /// impl PageStore for TempFilePageStore { + /// fn put(&mut self, value: Bytes) -> Result { + /// // Append to the end of the file + /// self.file.seek(SeekFrom::Start(self.end))?; + /// self.file.write_all(&value)?; + /// let key = PageKey::new(self.locs.len() as u64); + /// self.locs.push((self.end, value.len())); + /// self.end += value.len() as u64; + /// Ok(key) + /// } + /// + /// fn take(&mut self, key: PageKey) -> Result { + /// let (offset, len) = self.locs[key.get() as usize]; + /// let mut buf = vec![0u8; len]; + /// self.file.seek(SeekFrom::Start(offset))?; + /// self.file.read_exact(&mut buf)?; + /// Ok(Bytes::from(buf)) + /// } + /// } + /// + /// /// Factory for creating [`TempFilePageStore`] + /// #[derive(Debug)] + /// struct TempFilePageStoreFactory; + /// + /// impl PageStoreFactory for TempFilePageStoreFactory { + /// fn create(&self, args: &PageStoreArgs<'_>) -> Result> { + /// // `args` exposes the column index and descriptor (physical/logical + /// // type, path), so a real backend might choose to spill only large columns. + /// let _ = (args.column_index(), args.column_descriptor()); + /// Ok(Box::new(TempFilePageStore { + /// file: tempfile::tempfile()?, // temp file is cleaned on drop + /// end: 0, + /// locs: Vec::new(), + /// })) + /// } + /// } + /// // write 1000 integers + /// let col = Arc::new(Int64Array::from_iter_values(0..1000)) as ArrayRef; + /// let to_write = RecordBatch::try_from_iter([("col", col)]).unwrap(); + /// + /// let options = + /// ArrowWriterOptions::new().with_page_store_factory(Arc::new(TempFilePageStoreFactory)); + /// let mut buffer = Vec::new(); + /// let mut writer = + /// ArrowWriter::try_new_with_options(&mut buffer, to_write.schema(), options).unwrap(); + /// writer.write(&to_write).unwrap(); + /// writer.close().unwrap(); + /// + /// // buffer now holds valid Parquet data, which can be read as normal: + /// let mut reader = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 1024).unwrap(); + /// assert_eq!(to_write, reader.next().unwrap().unwrap()); + /// ``` + pub fn with_page_store_factory(self, page_store_factory: Arc) -> Self { + Self { + page_store_factory: Some(page_store_factory), + ..self + } + } + + /// Skip encoding the embedded arrow metadata (defaults to `false`) + /// + /// Parquet files generated by the [`ArrowWriter`] contain embedded arrow schema + /// by default. + /// + /// Set `skip_arrow_metadata` to true, to skip encoding the embedded metadata. + pub fn with_skip_arrow_metadata(self, skip_arrow_metadata: bool) -> Self { + Self { + skip_arrow_metadata, + ..self + } + } + + /// Set the name of the root parquet schema element (defaults to `"arrow_schema"`) + pub fn with_schema_root(self, schema_root: String) -> Self { + Self { + schema_root: Some(schema_root), + ..self + } + } + + /// Explicitly specify the Parquet schema to be used + /// + /// If omitted (the default), the [`ArrowSchemaConverter`] is used to compute the + /// Parquet [`SchemaDescriptor`]. This may be used When the [`SchemaDescriptor`] is + /// already known or must be calculated using custom logic. + pub fn with_parquet_schema(self, schema_descr: SchemaDescriptor) -> Self { + Self { + schema_descr: Some(schema_descr), + ..self + } + } +} + +/// A single column chunk produced by [`ArrowColumnWriter`]. +/// +/// Holds the serialized page blobs (each page's header ‖ compressed data, in +/// write order) in a [`PageStore`], plus the handles needed to read them back, +/// in order, when the chunk is spliced into the output file. +struct ArrowColumnChunkData { + length: usize, + store: Box, + keys: Vec, + /// Handles to the dictionary page's blobs (header then data) in the store. + /// + /// A dictionary page is produced at most once and bounded by + /// `dict_page_size_limit`, but it must be written *first* in the chunk even + /// though the data pages reach the writer before it (see + /// [`PageWriter::defers_dictionary_ordering`]). Its header and data are `put` + /// into the store like any other page — which keeps the store uniform, and + /// lets an oversized dictionary page spill — and their handles are held apart + /// so they can be emitted ahead of the data pages at splice. + /// Empty for non-dictionary columns. + dictionary_keys: Vec, + /// Serialized length of the dictionary page (0 if there is none), recorded + /// so the data pages can be shifted past it when offsets are rewritten to a + /// dictionary-first layout at splice. + dictionary_len: usize, +} + +impl ArrowColumnChunkData { + fn new(store: Box) -> Self { + Self { + length: 0, + store, + keys: Vec::new(), + dictionary_keys: Vec::new(), + dictionary_len: 0, + } + } + + /// Append a data-page blob to the store, recording its handle in write + /// order. + fn push(&mut self, value: Bytes) -> Result<()> { + let key = self.store.put(value)?; + self.keys.push(key); + Ok(()) + } + + /// Store a dictionary-page blob (header or data) in the page store, + /// recording its handle (emitted first at splice) and accumulating its + /// serialized length. + fn push_dictionary(&mut self, value: Bytes) -> Result<()> { + self.dictionary_len += value.len(); + let key = self.store.put(value)?; + self.dictionary_keys.push(key); + Ok(()) + } + + /// Bytes this chunk currently holds on the heap: whatever the store keeps + /// resident (zero for a spilling backend). + fn memory_size(&self) -> usize { + self.store.memory_size() + } +} + +/// A streaming iterator over one column chunk's buffered page blobs, in final +/// file order: the dictionary page (if any) first, then the data pages. +/// +/// Each blob is taken back out of the [`PageStore`] *as it is +/// consumed* and released immediately afterwards, so splicing a chunk into the +/// output file never materializes more than a single page in memory at a time. +/// This is what keeps the splice phase within the memory bound for a spilling +/// backend (an in-memory store already holds the bytes, so it is unaffected). +struct StreamingColumnChunkPages { + store: Box, + /// Page handles in final file order: the dictionary page first (if any), + /// then the data pages. + keys: IntoIter, +} + +impl StreamingColumnChunkPages { + fn new(data: ArrowColumnChunkData) -> Self { + // The dictionary page must be emitted first, ahead of the data pages, + // even though it was the last page produced. + let keys = if data.dictionary_keys.is_empty() { + data.keys + } else { + let mut keys = Vec::with_capacity(data.dictionary_keys.len() + data.keys.len()); + keys.extend(data.dictionary_keys); + keys.extend(data.keys); + keys + }; + Self { + store: data.store, + keys: keys.into_iter(), + } + } +} + +impl Iterator for StreamingColumnChunkPages { + type Item = Result; + + fn next(&mut self) -> Option { + let key = self.keys.next()?; + Some(self.store.take(key)) + } +} + +/// A shared [`ArrowColumnChunkData`] +/// +/// This allows it to be owned by [`ArrowPageWriter`] whilst allowing access via +/// [`ArrowRowGroupWriter`] on flush, without requiring self-referential borrows +type SharedColumnChunk = Arc>; + +struct ArrowPageWriter { + buffer: SharedColumnChunk, + #[cfg(feature = "encryption")] + page_encryptor: Option, +} + +impl ArrowPageWriter { + /// Create a page writer that buffers completed pages in `store`. + fn new(store: Box) -> Self { + Self { + buffer: Arc::new(Mutex::new(ArrowColumnChunkData::new(store))), + #[cfg(feature = "encryption")] + page_encryptor: None, + } + } + + #[cfg(feature = "encryption")] + pub fn with_encryptor(mut self, page_encryptor: Option) -> Self { + self.page_encryptor = page_encryptor; + self + } + + #[cfg(feature = "encryption")] + fn page_encryptor_mut(&mut self) -> Option<&mut PageEncryptor> { + self.page_encryptor.as_mut() + } + + #[cfg(not(feature = "encryption"))] + fn page_encryptor_mut(&mut self) -> Option<&mut PageEncryptor> { + None + } +} + +impl PageWriter for ArrowPageWriter { + fn write_page(&mut self, page: CompressedPage) -> Result { + let page = match self.page_encryptor_mut() { + Some(page_encryptor) => page_encryptor.encrypt_compressed_page(page)?, + None => page, + }; + + let page_header = page.to_thrift_header()?; + let header = { + let mut header = Vec::with_capacity(1024); + + match self.page_encryptor_mut() { + Some(page_encryptor) => { + page_encryptor.encrypt_page_header(&page_header, &mut header)?; + if page.compressed_page().is_data_page() { + page_encryptor.increment_page(); + } + } + None => { + let mut protocol = ThriftCompactOutputProtocol::new(&mut header); + page_header.write_thrift(&mut protocol)?; + } + }; + + Bytes::from(header) + }; + + let mut buf = self.buffer.try_lock().unwrap(); + + let data = page.compressed_page().buffer().clone(); + let compressed_size = data.len() + header.len(); + + let mut spec = PageWriteSpec::new(); + spec.page_type = page.page_type(); + spec.num_values = page.num_values(); + spec.uncompressed_size = page.uncompressed_size() + header.len(); + spec.offset = buf.length as u64; + spec.compressed_size = compressed_size; + spec.bytes_written = compressed_size as u64; + + buf.length += compressed_size; + if spec.page_type == PageType::DICTIONARY_PAGE { + // Recorded apart from the data pages so it is emitted first at + // splice — see `ArrowColumnChunkData::dictionary_keys`. + buf.push_dictionary(header)?; + buf.push_dictionary(data)?; + } else { + buf.push(header)?; + buf.push(data)?; + } + + Ok(spec) + } + + fn defers_dictionary_ordering(&self) -> bool { + // The Arrow chunk is buffered in full and spliced at row-group flush, so + // data pages may be accepted before the dictionary page and reordered + // then. This lets `GenericColumnWriter` stream dictionary-column data + // pages straight through instead of buffering them in memory. + true + } + + fn buffered_memory_size(&self) -> usize { + // Only what is actually resident: a spilling store reports ~0 here even + // though the chunk's bytes have all passed through it. + self.buffer.try_lock().unwrap().memory_size() + } + + fn close(&mut self) -> Result<()> { + Ok(()) + } +} + +/// A leaf column that can be encoded by [`ArrowColumnWriter`] +#[derive(Debug)] +pub struct ArrowLeafColumn(ArrayLevels); + +/// Computes the [`ArrowLeafColumn`] for a potentially nested [`ArrayRef`] +/// +/// This function can be used along with [`get_column_writers`] to encode +/// individual columns in parallel. See example on [`ArrowColumnWriter`] +pub fn compute_leaves(field: &Field, array: &ArrayRef) -> Result> { + let levels = calculate_array_levels(array, field)?; + Ok(levels.into_iter().map(ArrowLeafColumn).collect()) +} + +/// The data for a single column chunk, see [`ArrowColumnWriter`] +pub struct ArrowColumnChunk { + data: ArrowColumnChunkData, + close: ColumnCloseResult, +} + +impl std::fmt::Debug for ArrowColumnChunk { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ArrowColumnChunk") + .field("length", &self.data.length) + .finish_non_exhaustive() + } +} + +impl ArrowColumnChunk { + /// Returns the [`ColumnCloseResult`] produced when the chunk was closed. + /// + /// Exposes encoding information, collected statistics, and the optional + /// [`ColumnIndexMetaData`](crate::file::page_index::column_index::ColumnIndexMetaData) + /// / [`OffsetIndexMetaData`](crate::file::page_index::offset_index::OffsetIndexMetaData) + /// gathered for the column chunk. + pub fn close(&self) -> &ColumnCloseResult { + &self.close + } + + /// Returns a mutable reference to the [`ColumnCloseResult`]. + /// + /// This allows callers to mutate the close result before the chunk is + /// appended to a row group — for example, clearing `column_index` or + /// `bloom_filter` based on a dynamic rule that inspects the encodings and + /// collected page statistics. + pub fn close_mut(&mut self) -> &mut ColumnCloseResult { + &mut self.close + } + + /// Splices this column's buffered pages into the row group, streaming them + /// back out of the [`PageStore`] one page at a time. + pub fn append_to_row_group( + self, + writer: &mut SerializedRowGroupWriter<'_, W>, + ) -> Result<()> { + let ArrowColumnChunk { data, close } = self; + + // The dictionary page is produced *after* the data pages on this path (so + // they can stream straight through) but must be written *first*, so move + // it ahead of the data pages in the recorded offsets before the splice. + let close = close.update_dictionary_location(data.dictionary_len)?; + + let pages = StreamingColumnChunkPages::new(data); + writer.append_column_from_pages(pages, close) + } +} + +/// Encodes [`ArrowLeafColumn`] to [`ArrowColumnChunk`] +/// +/// `ArrowColumnWriter` instances can be created using an [`ArrowRowGroupWriterFactory`]; +/// +/// Note: This is a low-level interface for applications that require +/// fine-grained control of encoding (e.g. encoding using multiple threads), +/// see [`ArrowWriter`] for a higher-level interface +/// +/// # Example: Encoding two Arrow Array's in Parallel +/// ``` +/// // The arrow schema +/// # use std::sync::Arc; +/// # use arrow_array::*; +/// # use arrow_schema::*; +/// # use parquet::arrow::ArrowSchemaConverter; +/// # use parquet::arrow::arrow_writer::{compute_leaves, ArrowColumnChunk, ArrowLeafColumn, ArrowRowGroupWriterFactory}; +/// # use parquet::file::properties::WriterProperties; +/// # use parquet::file::writer::{SerializedFileWriter, SerializedRowGroupWriter}; +/// # +/// let schema = Arc::new(Schema::new(vec![ +/// Field::new("i32", DataType::Int32, false), +/// Field::new("f32", DataType::Float32, false), +/// ])); +/// +/// // Compute the parquet schema +/// let props = Arc::new(WriterProperties::default()); +/// let parquet_schema = ArrowSchemaConverter::new() +/// .with_coerce_types(props.coerce_types()) +/// .convert(&schema) +/// .unwrap(); +/// +/// // Create parquet writer +/// let root_schema = parquet_schema.root_schema_ptr(); +/// // write to memory in the example, but this could be a File +/// let mut out = Vec::with_capacity(1024); +/// let mut writer = SerializedFileWriter::new(&mut out, root_schema, props.clone()) +/// .unwrap(); +/// +/// // Create a factory for building Arrow column writers +/// let row_group_factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema)); +/// // Create column writers for the 0th row group +/// let col_writers = row_group_factory.create_column_writers(0).unwrap(); +/// +/// // Spawn a worker thread for each column +/// // +/// // Note: This is for demonstration purposes, a thread-pool e.g. rayon or tokio, would be better. +/// // The `map` produces an iterator of type `tuple of (thread handle, send channel)`. +/// let mut workers: Vec<_> = col_writers +/// .into_iter() +/// .map(|mut col_writer| { +/// let (send, recv) = std::sync::mpsc::channel::(); +/// let handle = std::thread::spawn(move || { +/// // receive Arrays to encode via the channel +/// for col in recv { +/// col_writer.write(&col)?; +/// } +/// // once the input is complete, close the writer +/// // to return the newly created ArrowColumnChunk +/// col_writer.close() +/// }); +/// (handle, send) +/// }) +/// .collect(); +/// +/// // Start row group +/// let mut row_group_writer: SerializedRowGroupWriter<'_, _> = writer +/// .next_row_group() +/// .unwrap(); +/// +/// // Create some example input columns to encode +/// let to_write = vec![ +/// Arc::new(Int32Array::from_iter_values([1, 2, 3])) as _, +/// Arc::new(Float32Array::from_iter_values([1., 45., -1.])) as _, +/// ]; +/// +/// // Send the input columns to the workers +/// let mut worker_iter = workers.iter_mut(); +/// for (arr, field) in to_write.iter().zip(&schema.fields) { +/// for leaves in compute_leaves(field, arr).unwrap() { +/// worker_iter.next().unwrap().1.send(leaves).unwrap(); +/// } +/// } +/// +/// // Wait for the workers to complete encoding, and append +/// // the resulting column chunks to the row group (and the file) +/// for (handle, send) in workers { +/// drop(send); // Drop send side to signal termination +/// // wait for the worker to send the completed chunk +/// let chunk: ArrowColumnChunk = handle.join().unwrap().unwrap(); +/// chunk.append_to_row_group(&mut row_group_writer).unwrap(); +/// } +/// // Close the row group which writes to the underlying file +/// row_group_writer.close().unwrap(); +/// +/// let metadata = writer.close().unwrap(); +/// assert_eq!(metadata.file_metadata().num_rows(), 3); +/// ``` +pub struct ArrowColumnWriter { + writer: ArrowColumnWriterImpl, + chunk: SharedColumnChunk, +} + +impl std::fmt::Debug for ArrowColumnWriter { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ArrowColumnWriter").finish_non_exhaustive() + } +} + +enum ArrowColumnWriterImpl { + ByteArray(GenericColumnWriter<'static, ByteArrayEncoder>), + Column(ColumnWriter<'static>), +} + +impl ArrowColumnWriter { + /// Write an [`ArrowLeafColumn`] + pub fn write(&mut self, col: &ArrowLeafColumn) -> Result<()> { + self.write_internal(&col.0) + } + + /// Write with content-defined chunking, inserting page flushes at chunk boundaries. + fn write_with_chunker( + &mut self, + col: &ArrowLeafColumn, + chunker: &mut ContentDefinedChunker, + ) -> Result<()> { + let levels = &col.0; + let chunks = chunker.get_arrow_chunks( + levels.def_level_data().as_ref(), + levels.rep_level_data().as_ref(), + levels.array(), + )?; + + let num_chunks = chunks.len(); + for (i, chunk) in chunks.iter().enumerate() { + let chunk_levels = levels.slice_for_chunk(chunk); + self.write_internal(&chunk_levels)?; + + // Add a page break after each chunk except the last + if i + 1 < num_chunks { + match &mut self.writer { + ArrowColumnWriterImpl::Column(c) => c.add_data_page()?, + ArrowColumnWriterImpl::ByteArray(c) => c.add_data_page()?, + } + } + } + Ok(()) + } + + fn write_internal(&mut self, levels: &ArrayLevels) -> Result<()> { + match &mut self.writer { + ArrowColumnWriterImpl::Column(c) => { + let leaf = levels.array(); + match leaf.as_any_dictionary_opt() { + Some(dictionary) => { + let materialized = + arrow_select::take::take(dictionary.values(), dictionary.keys(), None)?; + write_leaf(c, &materialized, levels)? + } + None => write_leaf(c, leaf, levels)?, + }; + } + ArrowColumnWriterImpl::ByteArray(c) => { + write_primitive(c, levels.array().as_ref(), levels)?; + } + } + Ok(()) + } + + /// Close this column returning the written [`ArrowColumnChunk`] + pub fn close(self) -> Result { + let close = match self.writer { + ArrowColumnWriterImpl::ByteArray(c) => c.close()?, + ArrowColumnWriterImpl::Column(c) => c.close()?, + }; + let chunk = Arc::try_unwrap(self.chunk).ok().unwrap(); + let data = chunk.into_inner().unwrap(); + Ok(ArrowColumnChunk { data, close }) + } + + /// Returns the estimated total memory usage by the writer. + /// + /// This [`Self::get_estimated_total_bytes`] this is an estimate + /// of the current memory usage and not it's anticipated encoded size. + /// + /// This includes: + /// 1. Data buffered in encoded form + /// 2. Data buffered in un-encoded form (e.g. `usize` dictionary keys) + /// + /// This value should be greater than or equal to [`Self::get_estimated_total_bytes`] + pub fn memory_size(&self) -> usize { + match &self.writer { + ArrowColumnWriterImpl::ByteArray(c) => c.memory_size(), + ArrowColumnWriterImpl::Column(c) => c.memory_size(), + } + } + + /// Returns the estimated total encoded bytes for this column writer. + /// + /// This includes: + /// 1. Data buffered in encoded form + /// 2. An estimate of how large the data buffered in un-encoded form would be once encoded + /// + /// This value should be less than or equal to [`Self::memory_size`] + pub fn get_estimated_total_bytes(&self) -> usize { + match &self.writer { + ArrowColumnWriterImpl::ByteArray(c) => c.get_estimated_total_bytes() as _, + ArrowColumnWriterImpl::Column(c) => c.get_estimated_total_bytes() as _, + } + } +} + +/// Encodes [`RecordBatch`] to a parquet row group +/// +/// Note: this structure is created by [`ArrowRowGroupWriterFactory`] internally used to +/// create [`ArrowRowGroupWriter`]s, but it is not exposed publicly. +/// +/// See the example on [`ArrowColumnWriter`] for how to encode columns in parallel +#[derive(Debug)] +struct ArrowRowGroupWriter { + writers: Vec, + schema: SchemaRef, + buffered_rows: usize, +} + +impl ArrowRowGroupWriter { + fn new(writers: Vec, arrow: &SchemaRef) -> Self { + Self { + writers, + schema: arrow.clone(), + buffered_rows: 0, + } + } + + fn write(&mut self, batch: &RecordBatch) -> Result<()> { + self.buffered_rows += batch.num_rows(); + let mut writers = self.writers.iter_mut(); + for (field, column) in self.schema.fields().iter().zip(batch.columns()) { + for leaf in compute_leaves(field.as_ref(), column)? { + writers.next().unwrap().write(&leaf)?; + } + } + Ok(()) + } + + fn write_with_chunkers( + &mut self, + batch: &RecordBatch, + chunkers: &mut [ContentDefinedChunker], + ) -> Result<()> { + self.buffered_rows += batch.num_rows(); + let mut writers = self.writers.iter_mut(); + let mut chunkers = chunkers.iter_mut(); + for (field, column) in self.schema.fields().iter().zip(batch.columns()) { + for leaf in compute_leaves(field.as_ref(), column)? { + writers + .next() + .unwrap() + .write_with_chunker(&leaf, chunkers.next().unwrap())?; + } + } + Ok(()) + } + + /// Returns the estimated total encoded bytes for this row group + fn get_estimated_total_bytes(&self) -> usize { + self.writers + .iter() + .map(|x| x.get_estimated_total_bytes()) + .sum() + } + + fn close(self) -> Result> { + self.writers + .into_iter() + .map(|writer| writer.close()) + .collect() + } +} + +/// Factory that creates new column writers for each row group in the Parquet file. +/// +/// You can create this structure via an [`ArrowWriter::into_serialized_writer`]. +/// See the example on [`ArrowColumnWriter`] for how to encode columns in parallel +#[derive(Debug)] +pub struct ArrowRowGroupWriterFactory { + schema: SchemaDescPtr, + arrow_schema: SchemaRef, + props: WriterPropertiesPtr, + page_store_factory: Arc, + #[cfg(feature = "encryption")] + file_encryptor: Option>, +} + +impl ArrowRowGroupWriterFactory { + /// Create a new [`ArrowRowGroupWriterFactory`] for the provided file writer and Arrow schema + pub fn new( + file_writer: &SerializedFileWriter, + arrow_schema: SchemaRef, + ) -> Self { + let schema = Arc::clone(file_writer.schema_descr_ptr()); + let props = Arc::clone(file_writer.properties()); + Self { + schema, + arrow_schema, + props, + page_store_factory: Arc::new(InMemoryPageStoreFactory), + #[cfg(feature = "encryption")] + file_encryptor: file_writer.file_encryptor(), + } + } + + /// Set the [`PageStoreFactory`] used to allocate the buffer for each column + /// chunk, e.g. to spill completed pages to a temp file or object storage + /// instead of the heap. Defaults to [`InMemoryPageStoreFactory`]. + pub fn with_page_store_factory( + mut self, + page_store_factory: Arc, + ) -> Self { + self.page_store_factory = page_store_factory; + self + } + + fn create_row_group_writer(&self, row_group_index: usize) -> Result { + let writers = self.create_column_writers(row_group_index)?; + Ok(ArrowRowGroupWriter::new(writers, &self.arrow_schema)) + } + + /// Create column writers for a new row group, with the given row group index + pub fn create_column_writers(&self, row_group_index: usize) -> Result> { + let mut writers = Vec::with_capacity(self.arrow_schema.fields.len()); + let mut leaves = self.schema.columns().iter(); + let column_factory = self.column_writer_factory(row_group_index); + for field in &self.arrow_schema.fields { + column_factory.get_arrow_column_writer( + field.data_type(), + &self.props, + &mut leaves, + &mut writers, + )?; + } + Ok(writers) + } + + #[cfg(feature = "encryption")] + fn column_writer_factory(&self, row_group_idx: usize) -> ArrowColumnWriterFactory { + ArrowColumnWriterFactory::new() + .with_page_store_factory(self.page_store_factory.clone()) + .with_file_encryptor(row_group_idx, self.file_encryptor.clone()) + } + + #[cfg(not(feature = "encryption"))] + fn column_writer_factory(&self, _row_group_idx: usize) -> ArrowColumnWriterFactory { + ArrowColumnWriterFactory::new().with_page_store_factory(self.page_store_factory.clone()) + } +} + +/// Returns [`ArrowColumnWriter`]s for each column in a given schema +#[deprecated(since = "57.0.0", note = "Use `ArrowRowGroupWriterFactory` instead")] +pub fn get_column_writers( + parquet: &SchemaDescriptor, + props: &WriterPropertiesPtr, + arrow: &SchemaRef, +) -> Result> { + let mut writers = Vec::with_capacity(arrow.fields.len()); + let mut leaves = parquet.columns().iter(); + let column_factory = ArrowColumnWriterFactory::new(); + for field in &arrow.fields { + column_factory.get_arrow_column_writer( + field.data_type(), + props, + &mut leaves, + &mut writers, + )?; + } + Ok(writers) +} + +/// Creates [`ArrowColumnWriter`] instances +struct ArrowColumnWriterFactory { + /// Allocates the per-column-chunk [`PageStore`] backing each page writer. + page_store_factory: Arc, + #[cfg(feature = "encryption")] + row_group_index: usize, + #[cfg(feature = "encryption")] + file_encryptor: Option>, +} + +impl ArrowColumnWriterFactory { + pub fn new() -> Self { + Self { + page_store_factory: Arc::new(InMemoryPageStoreFactory), + #[cfg(feature = "encryption")] + row_group_index: 0, + #[cfg(feature = "encryption")] + file_encryptor: None, + } + } + + /// Use `page_store_factory` to allocate the buffer for each column chunk. + pub fn with_page_store_factory( + mut self, + page_store_factory: Arc, + ) -> Self { + self.page_store_factory = page_store_factory; + self + } + + #[cfg(feature = "encryption")] + pub fn with_file_encryptor( + mut self, + row_group_index: usize, + file_encryptor: Option>, + ) -> Self { + self.row_group_index = row_group_index; + self.file_encryptor = file_encryptor; + self + } + + #[cfg(feature = "encryption")] + fn create_page_writer( + &self, + column_descriptor: &ColumnDescPtr, + column_index: usize, + ) -> Result> { + let column_path = column_descriptor.path().string(); + let page_encryptor = PageEncryptor::create_if_column_encrypted( + &self.file_encryptor, + self.row_group_index, + column_index, + &column_path, + )?; + let args = PageStoreArgs::new(column_index, column_descriptor); + let store = self.page_store_factory.create(&args)?; + Ok(Box::new( + ArrowPageWriter::new(store).with_encryptor(page_encryptor), + )) + } + + #[cfg(not(feature = "encryption"))] + fn create_page_writer( + &self, + column_descriptor: &ColumnDescPtr, + column_index: usize, + ) -> Result> { + let args = PageStoreArgs::new(column_index, column_descriptor); + let store = self.page_store_factory.create(&args)?; + Ok(Box::new(ArrowPageWriter::new(store))) + } + + /// Gets an [`ArrowColumnWriter`] for the given `data_type`, appending the + /// output ColumnDesc to `leaves` and the column writers to `out` + fn get_arrow_column_writer( + &self, + data_type: &ArrowDataType, + props: &WriterPropertiesPtr, + leaves: &mut Iter<'_, ColumnDescPtr>, + out: &mut Vec, + ) -> Result<()> { + // Instantiate writers for normal columns + let col = |desc: &ColumnDescPtr| -> Result { + let page_writer = self.create_page_writer(desc, out.len())?; + let chunk = page_writer.buffer.clone(); + let writer = get_column_writer(desc.clone(), props.clone(), page_writer); + Ok(ArrowColumnWriter { + chunk, + writer: ArrowColumnWriterImpl::Column(writer), + }) + }; + + // Instantiate writers for byte arrays (e.g. Utf8, Binary, etc) + let bytes = |desc: &ColumnDescPtr| -> Result { + let page_writer = self.create_page_writer(desc, out.len())?; + let chunk = page_writer.buffer.clone(); + let writer = GenericColumnWriter::new(desc.clone(), props.clone(), page_writer); + Ok(ArrowColumnWriter { + chunk, + writer: ArrowColumnWriterImpl::ByteArray(writer), + }) + }; + + match data_type { + _ if data_type.is_primitive() => out.push(col(leaves.next().unwrap())?), + ArrowDataType::FixedSizeBinary(_) | ArrowDataType::Boolean | ArrowDataType::Null => { + out.push(col(leaves.next().unwrap())?) + } + ArrowDataType::LargeBinary + | ArrowDataType::Binary + | ArrowDataType::Utf8 + | ArrowDataType::LargeUtf8 + | ArrowDataType::BinaryView + | ArrowDataType::Utf8View => out.push(bytes(leaves.next().unwrap())?), + ArrowDataType::List(f) + | ArrowDataType::LargeList(f) + | ArrowDataType::FixedSizeList(f, _) + | ArrowDataType::ListView(f) + | ArrowDataType::LargeListView(f) => { + self.get_arrow_column_writer(f.data_type(), props, leaves, out)? + } + ArrowDataType::Struct(fields) => { + for field in fields { + self.get_arrow_column_writer(field.data_type(), props, leaves, out)? + } + } + ArrowDataType::Map(f, _) => match f.data_type() { + ArrowDataType::Struct(f) => { + self.get_arrow_column_writer(f[0].data_type(), props, leaves, out)?; + self.get_arrow_column_writer(f[1].data_type(), props, leaves, out)? + } + _ => unreachable!("invalid map type"), + }, + ArrowDataType::Dictionary(_, value_type) => match value_type.as_ref() { + ArrowDataType::Utf8 + | ArrowDataType::LargeUtf8 + | ArrowDataType::Binary + | ArrowDataType::LargeBinary => out.push(bytes(leaves.next().unwrap())?), + ArrowDataType::Utf8View | ArrowDataType::BinaryView => { + out.push(bytes(leaves.next().unwrap())?) + } + ArrowDataType::FixedSizeBinary(_) => out.push(bytes(leaves.next().unwrap())?), + _ => out.push(col(leaves.next().unwrap())?), + }, + ArrowDataType::RunEndEncoded(_, value_field) => { + self.get_arrow_column_writer(value_field.data_type(), props, leaves, out)? + } + _ => { + return Err(ParquetError::NYI(format!( + "Attempting to write an Arrow type {data_type} to parquet that is not yet implemented" + ))); + } + } + Ok(()) + } +} + +fn write_leaf( + writer: &mut ColumnWriter<'_>, + column: &dyn arrow_array::Array, + levels: &ArrayLevels, +) -> Result { + let indices = levels.non_null_indices(); + + match writer { + // Note: this should match the contents of arrow_to_parquet_type + ColumnWriter::Int32ColumnWriter(typed) => { + match column.data_type() { + ArrowDataType::Null => { + let array = Int32Array::new_null(column.len()); + write_primitive(typed, array.values(), levels) + } + ArrowDataType::Int8 => { + let array: Int32Array = column.as_primitive::().unary(|x| x as i32); + write_primitive(typed, array.values(), levels) + } + ArrowDataType::Int16 => { + let array: Int32Array = column.as_primitive::().unary(|x| x as i32); + write_primitive(typed, array.values(), levels) + } + ArrowDataType::Int32 => { + write_primitive(typed, column.as_primitive::().values(), levels) + } + ArrowDataType::UInt8 => { + let array: Int32Array = column.as_primitive::().unary(|x| x as i32); + write_primitive(typed, array.values(), levels) + } + ArrowDataType::UInt16 => { + let array: Int32Array = column.as_primitive::().unary(|x| x as i32); + write_primitive(typed, array.values(), levels) + } + ArrowDataType::UInt32 => { + // follow C++ implementation and use overflow/reinterpret cast from u32 to i32 which will map + // `(i32::MAX as u32)..u32::MAX` to `i32::MIN..0` + let array = column.as_primitive::(); + write_primitive(typed, array.values().inner().typed_data(), levels) + } + ArrowDataType::Date32 => { + let array = column.as_primitive::(); + write_primitive(typed, array.values(), levels) + } + ArrowDataType::Time32(TimeUnit::Second) => { + let array = column.as_primitive::(); + write_primitive(typed, array.values(), levels) + } + ArrowDataType::Time32(TimeUnit::Millisecond) => { + let array = column.as_primitive::(); + write_primitive(typed, array.values(), levels) + } + ArrowDataType::Date64 => { + // If the column is a Date64, we truncate it + let array: Int32Array = column + .as_primitive::() + .unary(|x| (x / 86_400_000) as _); + + write_primitive(typed, array.values(), levels) + } + ArrowDataType::Decimal32(_, _) => { + let array = column + .as_primitive::() + .unary::<_, Int32Type>(|v| v); + write_primitive(typed, array.values(), levels) + } + ArrowDataType::Decimal64(_, _) => { + // use the int32 to represent the decimal with low precision + let array = column + .as_primitive::() + .unary::<_, Int32Type>(|v| v as i32); + write_primitive(typed, array.values(), levels) + } + ArrowDataType::Decimal128(_, _) => { + // use the int32 to represent the decimal with low precision + let array = column + .as_primitive::() + .unary::<_, Int32Type>(|v| v as i32); + write_primitive(typed, array.values(), levels) + } + ArrowDataType::Decimal256(_, _) => { + // use the int32 to represent the decimal with low precision + let array = column + .as_primitive::() + .unary::<_, Int32Type>(|v| v.as_i128() as i32); + write_primitive(typed, array.values(), levels) + } + d => Err(ParquetError::General(format!("Cannot coerce {d} to I32"))), + } + } + ColumnWriter::BoolColumnWriter(typed) => { + let array = column.as_boolean(); + let values = get_bool_array_slice(array, indices.iter().copied()); + typed.write_batch_internal( + values.as_slice(), + None, + levels.def_level_data().as_ref(), + levels.rep_level_data().as_ref(), + None, + None, + None, + ) + } + ColumnWriter::Int64ColumnWriter(typed) => { + match column.data_type() { + ArrowDataType::Date64 => { + let array = column + .as_primitive::() + .reinterpret_cast::(); + + write_primitive(typed, array.values(), levels) + } + ArrowDataType::Int64 => { + let array = column.as_primitive::(); + write_primitive(typed, array.values(), levels) + } + ArrowDataType::UInt64 => { + let values = column.as_primitive::().values(); + // follow C++ implementation and use overflow/reinterpret cast from u64 to i64 which will map + // `(i64::MAX as u64)..u64::MAX` to `i64::MIN..0` + let array = values.inner().typed_data::(); + write_primitive(typed, array, levels) + } + ArrowDataType::Time64(TimeUnit::Microsecond) => { + let array = column.as_primitive::(); + write_primitive(typed, array.values(), levels) + } + ArrowDataType::Time64(TimeUnit::Nanosecond) => { + let array = column.as_primitive::(); + write_primitive(typed, array.values(), levels) + } + ArrowDataType::Timestamp(unit, _) => match unit { + TimeUnit::Second => { + let array = column.as_primitive::(); + write_primitive(typed, array.values(), levels) + } + TimeUnit::Millisecond => { + let array = column.as_primitive::(); + write_primitive(typed, array.values(), levels) + } + TimeUnit::Microsecond => { + let array = column.as_primitive::(); + write_primitive(typed, array.values(), levels) + } + TimeUnit::Nanosecond => { + let array = column.as_primitive::(); + write_primitive(typed, array.values(), levels) + } + }, + ArrowDataType::Duration(unit) => match unit { + TimeUnit::Second => { + let array = column.as_primitive::(); + write_primitive(typed, array.values(), levels) + } + TimeUnit::Millisecond => { + let array = column.as_primitive::(); + write_primitive(typed, array.values(), levels) + } + TimeUnit::Microsecond => { + let array = column.as_primitive::(); + write_primitive(typed, array.values(), levels) + } + TimeUnit::Nanosecond => { + let array = column.as_primitive::(); + write_primitive(typed, array.values(), levels) + } + }, + ArrowDataType::Decimal64(_, _) => { + let array = column + .as_primitive::() + .reinterpret_cast::(); + write_primitive(typed, array.values(), levels) + } + ArrowDataType::Decimal128(_, _) => { + // use the int64 to represent the decimal with low precision + let array = column + .as_primitive::() + .unary::<_, Int64Type>(|v| v as i64); + write_primitive(typed, array.values(), levels) + } + ArrowDataType::Decimal256(_, _) => { + // use the int64 to represent the decimal with low precision + let array = column + .as_primitive::() + .unary::<_, Int64Type>(|v| v.as_i128() as i64); + write_primitive(typed, array.values(), levels) + } + d => Err(ParquetError::General(format!("Cannot coerce {d} to I64"))), + } + } + ColumnWriter::Int96ColumnWriter(_typed) => { + unreachable!("Currently unreachable because data type not supported") + } + ColumnWriter::FloatColumnWriter(typed) => { + let array = column.as_primitive::(); + write_primitive(typed, array.values(), levels) + } + ColumnWriter::DoubleColumnWriter(typed) => { + let array = column.as_primitive::(); + write_primitive(typed, array.values(), levels) + } + ColumnWriter::ByteArrayColumnWriter(_) => { + unreachable!("should use ByteArrayWriter") + } + ColumnWriter::FixedLenByteArrayColumnWriter(typed) => { + let bytes = match column.data_type() { + ArrowDataType::Interval(interval_unit) => match interval_unit { + IntervalUnit::YearMonth => { + let array = column.as_primitive::(); + get_interval_ym_array_slice(array, indices.iter().copied()) + } + IntervalUnit::DayTime => { + let array = column.as_primitive::(); + get_interval_dt_array_slice(array, indices.iter().copied()) + } + _ => { + return Err(ParquetError::NYI(format!( + "Attempting to write an Arrow interval type {interval_unit:?} to parquet that is not yet implemented" + ))); + } + }, + ArrowDataType::FixedSizeBinary(_) => { + let array = column.as_fixed_size_binary(); + get_fsb_array_slice(array, indices.iter().copied()) + } + ArrowDataType::Decimal32(_, _) => { + let array = column.as_primitive::(); + get_decimal_array_slice(array, indices.iter().copied()) + } + ArrowDataType::Decimal64(_, _) => { + let array = column.as_primitive::(); + get_decimal_array_slice(array, indices.iter().copied()) + } + ArrowDataType::Decimal128(_, _) => { + let array = column.as_primitive::(); + get_decimal_array_slice(array, indices.iter().copied()) + } + ArrowDataType::Decimal256(_, _) => { + let array = column.as_primitive::(); + get_decimal_array_slice(array, indices.iter().copied()) + } + ArrowDataType::Float16 => { + let array = column.as_primitive::(); + get_float_16_array_slice(array, indices.iter().copied()) + } + _ => { + return Err(ParquetError::NYI( + "Attempting to write an Arrow type that is not yet implemented".to_string(), + )); + } + }; + typed.write_batch_internal( + bytes.as_slice(), + None, + levels.def_level_data().as_ref(), + levels.rep_level_data().as_ref(), + None, + None, + None, + ) + } + } +} + +fn write_primitive( + writer: &mut GenericColumnWriter, + values: &E::Values, + levels: &ArrayLevels, +) -> Result { + writer.write_batch_internal( + values, + Some(levels.non_null_indices()), + levels.def_level_data().as_ref(), + levels.rep_level_data().as_ref(), + None, + None, + None, + ) +} + +fn get_bool_array_slice( + array: &arrow_array::BooleanArray, + indices: impl ExactSizeIterator, +) -> Vec { + let mut values = Vec::with_capacity(indices.len()); + for i in indices { + values.push(array.value(i)) + } + values +} + +/// Returns 12-byte values representing 3 values of months, days and milliseconds (4-bytes each). +/// An Arrow YearMonth interval only stores months, thus only the first 4 bytes are populated. +fn get_interval_ym_array_slice( + array: &arrow_array::IntervalYearMonthArray, + indices: impl ExactSizeIterator, +) -> Vec { + chunk_array_slice(12, indices, move |i, chunk| { + let value = array.value(i); + chunk[0..4].copy_from_slice(&value.to_le_bytes()); + }) +} + +/// Returns 12-byte values representing 3 values of months, days and milliseconds (4-bytes each). +/// An Arrow DayTime interval only stores days and millis, thus the first 4 bytes are not populated. +fn get_interval_dt_array_slice( + array: &arrow_array::IntervalDayTimeArray, + indices: impl ExactSizeIterator, +) -> Vec { + chunk_array_slice(12, indices, move |i, chunk| { + let value = array.value(i); + chunk[4..8].copy_from_slice(&value.days.to_le_bytes()); + chunk[8..12].copy_from_slice(&value.milliseconds.to_le_bytes()); + }) +} + +trait NativeDecimalType: DecimalType { + type NativeBytes: AsRef<[u8]>; + + fn to_be_bytes(value: Self::Native) -> Self::NativeBytes; +} +impl NativeDecimalType for Decimal32Type { + type NativeBytes = [u8; Self::BYTE_LENGTH]; + + fn to_be_bytes(value: Self::Native) -> Self::NativeBytes { + value.to_be_bytes() + } +} +impl NativeDecimalType for Decimal64Type { + type NativeBytes = [u8; Self::BYTE_LENGTH]; + + fn to_be_bytes(value: Self::Native) -> Self::NativeBytes { + value.to_be_bytes() + } +} +impl NativeDecimalType for Decimal128Type { + type NativeBytes = [u8; Self::BYTE_LENGTH]; + + fn to_be_bytes(value: Self::Native) -> Self::NativeBytes { + value.to_be_bytes() + } +} +impl NativeDecimalType for Decimal256Type { + type NativeBytes = [u8; Self::BYTE_LENGTH]; + + fn to_be_bytes(value: Self::Native) -> Self::NativeBytes { + value.to_be_bytes() + } +} + +fn get_decimal_array_slice( + array: &PrimitiveArray, + indices: impl ExactSizeIterator, +) -> Vec { + let chunk_size = decimal_length_from_precision(array.precision()); + assert!(chunk_size <= T::BYTE_LENGTH); + + if chunk_size == T::BYTE_LENGTH { + // Special-case that allows inlining memcpy. + chunk_array_slice(chunk_size, indices, move |i, chunk| { + let as_be_bytes = T::to_be_bytes(array.value(i)); + chunk.copy_from_slice(as_be_bytes.as_ref()); + }) + } else { + chunk_array_slice(chunk_size, indices, move |i, chunk| { + let as_be_bytes = T::to_be_bytes(array.value(i)); + let resized_value = &as_be_bytes.as_ref()[(T::BYTE_LENGTH - chunk.len())..]; + chunk.copy_from_slice(resized_value); + }) + } +} + +fn get_float_16_array_slice( + array: &arrow_array::Float16Array, + indices: impl ExactSizeIterator, +) -> Vec { + chunk_array_slice(2, indices, move |i, chunk| { + let value = array.value(i).to_le_bytes(); + chunk.copy_from_slice(&value); + }) +} + +fn get_fsb_array_slice( + array: &arrow_array::FixedSizeBinaryArray, + indices: impl ExactSizeIterator, +) -> Vec { + chunk_array_slice(array.value_size(), indices, move |i, chunk| { + let value = array.value(i); + chunk.copy_from_slice(value); + }) +} + +#[inline] +fn chunk_array_slice( + chunk_size: usize, + indices: impl ExactSizeIterator, + writer: impl Fn(usize, &mut [u8]), +) -> Vec { + let capacity = indices.len() * chunk_size; + // TODO: This could be done with Vec::spare_capacity_mut, + // but [MaybeUninit]::write_copy_of_slice is gated behind MSRV 1.93 + let mut arena = vec![0; capacity]; + for (i, chunk) in indices.zip(arena.chunks_exact_mut(chunk_size)) { + writer(i, chunk); + } + chunk_contiguous_vec(arena, chunk_size) +} + +fn chunk_contiguous_vec(arena: Vec, chunk_size: usize) -> Vec { + let mut values = Vec::with_capacity(arena.len() / chunk_size); + let mut arena = Bytes::from(arena); + while arena.len() >= chunk_size { + let slice = arena.split_to(chunk_size); + values.push(FixedLenByteArray::from(ByteArray::from(slice))); + } + values +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + + use std::fs::File; + + use crate::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder}; + use crate::arrow::{ARROW_SCHEMA_META_KEY, PARQUET_FIELD_ID_META_KEY}; + use crate::column::page::{Page, PageReader}; + use crate::file::metadata::thrift::PageHeader; + use crate::file::page_index::column_index::ColumnIndexMetaData; + use crate::file::reader::SerializedPageReader; + use crate::parquet_thrift::{ReadThrift, ThriftSliceInputProtocol}; + use crate::schema::types::ColumnPath; + use arrow::datatypes::ToByteSlice; + use arrow::datatypes::{DataType, Schema}; + use arrow::error::Result as ArrowResult; + use arrow::util::data_gen::create_random_array; + use arrow::util::pretty::pretty_format_batches; + use arrow::{array::*, buffer::Buffer}; + use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano, NullBuffer, OffsetBuffer, i256}; + use arrow_schema::Fields; + use half::f16; + use num_traits::{FromPrimitive, ToPrimitive}; + use tempfile::tempfile; + + use crate::basic::Encoding; + use crate::data_type::AsBytes; + use crate::file::metadata::{ColumnChunkMetaData, ParquetMetaData, ParquetMetaDataReader}; + use crate::file::properties::{ + BloomFilterPosition, EnabledStatistics, ReaderProperties, WriterVersion, + }; + use crate::file::serialized_reader::ReadOptionsBuilder; + use crate::file::{ + reader::{FileReader, SerializedFileReader}, + statistics::Statistics, + }; + + /// A [`PageStore`] that allocates *sparse, non-contiguous* handles and keeps + /// blobs in a `HashMap` — nothing like the default `Vec`. Used to + /// prove the writer relies only on the opaque-handle contract and never on + /// handles being dense `Vec` indices. Records how many blobs were stored. + #[derive(Debug, Default)] + struct RecordingPageStore { + next: u64, + blobs: HashMap, + puts: Arc, + } + + impl PageStore for RecordingPageStore { + fn put(&mut self, value: Bytes) -> Result { + // Deliberately non-sequential, never-zero handles. + let id = 100 + self.next * 7; + self.next += 1; + self.puts.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + self.blobs.insert(id, value); + Ok(PageKey::new(id)) + } + + fn take(&mut self, key: PageKey) -> Result { + self.blobs + .remove(&key.get()) + .ok_or_else(|| ParquetError::General(format!("missing key {}", key.get()))) + } + } + + #[derive(Debug)] + struct RecordingPageStoreFactory { + puts: Arc, + } + + impl PageStoreFactory for RecordingPageStoreFactory { + fn create(&self, _args: &PageStoreArgs<'_>) -> Result> { + Ok(Box::new(RecordingPageStore { + puts: self.puts.clone(), + ..Default::default() + })) + } + } + + /// A custom [`PageStore`] must produce byte-identical files to the in-memory + /// default, across dictionary and non-dictionary columns and multiple row + /// groups (so multiple store instances are exercised). + #[test] + fn custom_page_store_is_byte_identical_to_default() { + let schema = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Int32, true), + // A low-cardinality string column to exercise the dictionary path. + Field::new("s", DataType::Utf8, true), + ])); + let i = Int32Array::from(vec![Some(1), None, Some(3), Some(4), Some(5), Some(6)]); + let s = StringArray::from(vec![ + Some("a"), + Some("bb"), + Some("a"), + None, + Some("bb"), + Some("ccc"), + ]); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(i), Arc::new(s)]).unwrap(); + + // Small row groups so multiple column chunks (hence multiple store + // instances) are produced. + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .build(); + + let write = |factory: Option>| { + let mut buffer = Vec::new(); + let mut opts = ArrowWriterOptions::new().with_properties(props.clone()); + if let Some(factory) = factory { + opts = opts.with_page_store_factory(factory); + } + let mut writer = + ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + buffer + }; + + let default_bytes = write(None); + + let puts = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let custom_bytes = write(Some(Arc::new(RecordingPageStoreFactory { + puts: puts.clone(), + }))); + + assert!( + puts.load(std::sync::atomic::Ordering::Relaxed) > 0, + "custom PageStore was never written to" + ); + assert_eq!( + default_bytes, custom_bytes, + "a custom PageStore must produce byte-identical output to the default" + ); + } + + /// A dictionary-encoded column written through the deferred-ordering Arrow + /// path must round-trip correctly even with the offset index disabled, when + /// only the chunk-level dictionary/data page offsets are rewritten (there is + /// no offset index to rebuild). Spans multiple data pages so the + /// dictionary-first reordering is exercised. + #[test] + fn dictionary_column_round_trips_with_offset_index_disabled() { + let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, true)])); + + // Low cardinality so the column stays dictionary-encoded; enough rows to + // span several data pages within a single row group. + let values: Vec> = (0..50_000).map(|i| Some(i % 8)).collect(); + let array = Int32Array::from(values.clone()); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap(); + + let props = WriterProperties::builder() + .set_offset_index_disabled(true) + .set_data_page_row_count_limit(4096) + .build(); + let opts = ArrowWriterOptions::new().with_properties(props); + + let mut buffer = Vec::new(); + let mut writer = + ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let reader = ParquetRecordBatchReader::try_new(Bytes::from(buffer), values.len()).unwrap(); + let read: Vec = reader.collect::>().unwrap(); + let read_values: Vec> = read + .iter() + .flat_map(|b| b.column(0).as_primitive::().iter()) + .collect(); + assert_eq!(read_values, values); + } + + /// The dictionary page is routed through the [`PageStore`] like any other + /// page rather than held resident in memory, so a dictionary column chunk's + /// *entire* serialized size — dictionary page included — passes through the + /// store. + #[test] + fn dictionary_page_is_routed_through_the_store() { + /// A store that sums the bytes handed to `put`. + #[derive(Debug, Default)] + struct SizeRecordingPageStore { + blobs: Vec, + bytes_put: Arc, + } + impl PageStore for SizeRecordingPageStore { + fn put(&mut self, value: Bytes) -> Result { + self.bytes_put + .fetch_add(value.len(), std::sync::atomic::Ordering::Relaxed); + let key = PageKey::new(self.blobs.len() as u64); + self.blobs.push(value); + Ok(key) + } + fn take(&mut self, key: PageKey) -> Result { + Ok(std::mem::take(&mut self.blobs[key.get() as usize])) + } + } + #[derive(Debug)] + struct Factory { + bytes_put: Arc, + } + impl PageStoreFactory for Factory { + fn create(&self, _args: &PageStoreArgs<'_>) -> Result> { + Ok(Box::new(SizeRecordingPageStore { + bytes_put: self.bytes_put.clone(), + ..Default::default() + })) + } + } + + let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)])); + // Low cardinality keeps the column dictionary-encoded with a real, + // non-empty dictionary page. + let values: Vec<&str> = (0..2048) + .map(|i| ["alpha", "beta", "gamma", "delta"][i % 4]) + .collect(); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(StringArray::from(values))]) + .unwrap(); + + let bytes_put = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let opts = ArrowWriterOptions::new().with_page_store_factory(Arc::new(Factory { + bytes_put: bytes_put.clone(), + })); + + // A single batch / single column means exactly one row group and one + // store instance, so the bytes it saw map to one column chunk. + let mut buffer = Vec::new(); + let mut writer = + ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let reader = SerializedFileReader::new(Bytes::from(buffer)).unwrap(); + let column = reader.metadata().row_group(0).column(0); + assert!( + column.dictionary_page_offset().is_some(), + "expected the column to be dictionary-encoded" + ); + + // The bytes the store was handed must account for the whole chunk, + // dictionary page included. Holding the dictionary page apart from the + // store would make this fall short by the dictionary page's size. + assert_eq!( + bytes_put.load(std::sync::atomic::Ordering::Relaxed) as i64, + column.compressed_size(), + "the dictionary page must pass through the store like any other page" + ); + } + + #[test] + fn arrow_writer() { + // define schema + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, true), + ]); + + // create some data + let a = Int32Array::from(vec![1, 2, 3, 4, 5]); + let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]); + + // build a record batch + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap(); + + roundtrip(batch, Some(SMALL_SIZE / 2)); + } + + fn get_bytes_after_close(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec { + let mut buffer = vec![]; + + let mut writer = ArrowWriter::try_new(&mut buffer, schema, None).unwrap(); + writer.write(expected_batch).unwrap(); + writer.close().unwrap(); + + buffer + } + + fn get_bytes_by_into_inner(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec { + let mut writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap(); + writer.write(expected_batch).unwrap(); + writer.into_inner().unwrap() + } + + #[test] + fn roundtrip_bytes() { + // define schema + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, true), + ])); + + // create some data + let a = Int32Array::from(vec![1, 2, 3, 4, 5]); + let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]); + + // build a record batch + let expected_batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(a), Arc::new(b)]).unwrap(); + + for buffer in [ + get_bytes_after_close(schema.clone(), &expected_batch), + get_bytes_by_into_inner(schema, &expected_batch), + ] { + let cursor = Bytes::from(buffer); + let mut record_batch_reader = ParquetRecordBatchReader::try_new(cursor, 1024).unwrap(); + + let actual_batch = record_batch_reader + .next() + .expect("No batch found") + .expect("Unable to get batch"); + + assert_eq!(expected_batch.schema(), actual_batch.schema()); + assert_eq!(expected_batch.num_columns(), actual_batch.num_columns()); + assert_eq!(expected_batch.num_rows(), actual_batch.num_rows()); + for i in 0..expected_batch.num_columns() { + let expected_data = expected_batch.column(i).to_data(); + let actual_data = actual_batch.column(i).to_data(); + + assert_eq!(expected_data, actual_data); + } + } + } + + #[test] + fn arrow_writer_non_null() { + // define schema + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + + // create some data + let a = Int32Array::from(vec![1, 2, 3, 4, 5]); + + // build a record batch + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap(); + + roundtrip(batch, Some(SMALL_SIZE / 2)); + } + + #[test] + fn arrow_writer_list() { + // define schema + let schema = Schema::new(vec![Field::new( + "a", + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))), + true, + )]); + + // create some data + let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); + + // Construct a buffer for value offsets, for the nested array: + // [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]] + let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice()); + + // Construct a list array from the above two + let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field( + DataType::Int32, + false, + )))) + .len(5) + .add_buffer(a_value_offsets) + .add_child_data(a_values.into_data()) + .null_bit_buffer(Some(Buffer::from([0b00011011]))) + .build() + .unwrap(); + let a = ListArray::from(a_list_data); + + // build a record batch + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap(); + + assert_eq!(batch.column(0).null_count(), 1); + + // This test fails if the max row group size is less than the batch's length + // see https://github.com/apache/arrow-rs/issues/518 + roundtrip(batch, None); + } + + #[test] + fn arrow_writer_list_non_null() { + // define schema + let schema = Schema::new(vec![Field::new( + "a", + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))), + false, + )]); + + // create some data + let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); + + // Construct a buffer for value offsets, for the nested array: + // [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]] + let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice()); + + // Construct a list array from the above two + let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field( + DataType::Int32, + false, + )))) + .len(5) + .add_buffer(a_value_offsets) + .add_child_data(a_values.into_data()) + .build() + .unwrap(); + let a = ListArray::from(a_list_data); + + // build a record batch + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap(); + + // This test fails if the max row group size is less than the batch's length + // see https://github.com/apache/arrow-rs/issues/518 + assert_eq!(batch.column(0).null_count(), 0); + + roundtrip(batch, None); + } + + #[test] + fn arrow_writer_list_view() { + let list_field = Arc::new(Field::new_list_field(DataType::Int32, false)); + let schema = Schema::new(vec![Field::new( + "a", + DataType::ListView(list_field.clone()), + true, + )]); + + // [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]] + let a = ListViewArray::new( + list_field, + vec![0, 1, 0, 3, 6].into(), + vec![1, 2, 0, 3, 4].into(), + Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])), + Some(vec![true, true, false, true, true].into()), + ); + + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap(); + + assert_eq!(batch.column(0).null_count(), 1); + + roundtrip(batch, None); + } + + #[test] + fn arrow_writer_list_view_non_null() { + let list_field = Arc::new(Field::new_list_field(DataType::Int32, false)); + let schema = Schema::new(vec![Field::new( + "a", + DataType::ListView(list_field.clone()), + false, + )]); + + // [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]] + let a = ListViewArray::new( + list_field, + vec![0, 1, 0, 3, 6].into(), + vec![1, 2, 0, 3, 4].into(), + Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])), + None, + ); + + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap(); + + assert_eq!(batch.column(0).null_count(), 0); + + roundtrip(batch, None); + } + + #[test] + fn arrow_writer_list_view_out_of_order() { + let list_field = Arc::new(Field::new_list_field(DataType::Int32, false)); + let schema = Schema::new(vec![Field::new( + "a", + DataType::ListView(list_field.clone()), + false, + )]); + + // [[1], [2, 3], [], [7, 8, 9, 10], [4, 5, 6]] - out of order offsets + let a = ListViewArray::new( + list_field, + vec![0, 1, 0, 6, 3].into(), + vec![1, 2, 0, 4, 3].into(), + Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])), + None, + ); + + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap(); + + roundtrip(batch, None); + } + + #[test] + fn arrow_writer_large_list_view() { + let list_field = Arc::new(Field::new_list_field(DataType::Int32, false)); + let schema = Schema::new(vec![Field::new( + "a", + DataType::LargeListView(list_field.clone()), + true, + )]); + + // [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]] + let a = LargeListViewArray::new( + list_field, + vec![0i64, 1, 0, 3, 6].into(), + vec![1i64, 2, 0, 3, 4].into(), + Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])), + Some(vec![true, true, false, true, true].into()), + ); + + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap(); + + assert_eq!(batch.column(0).null_count(), 1); + + roundtrip(batch, None); + } + + #[test] + fn arrow_writer_list_view_with_struct() { + // Test ListView containing Struct: ListView> + let struct_fields = Fields::from(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, false), + ]); + let struct_type = DataType::Struct(struct_fields.clone()); + let list_field = Arc::new(Field::new("item", struct_type.clone(), false)); + + let schema = Schema::new(vec![Field::new( + "a", + DataType::ListView(list_field.clone()), + true, + )]); + + // Create struct values + let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]); + let name_array = StringArray::from(vec!["a", "b", "c", "d", "e"]); + let struct_array = StructArray::new( + struct_fields, + vec![Arc::new(id_array), Arc::new(name_array)], + None, + ); + + // Create ListView: [{1, "a"}, {2, "b"}], null, [{3, "c"}, {4, "d"}, {5, "e"}] + let list_view = ListViewArray::new( + list_field, + vec![0, 2, 2].into(), // offsets + vec![2, 0, 3].into(), // sizes + Arc::new(struct_array), + Some(vec![true, false, true].into()), + ); + + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(list_view)]).unwrap(); + + roundtrip(batch, None); + } + + #[test] + fn arrow_writer_binary() { + let string_field = Field::new("a", DataType::Utf8, false); + let binary_field = Field::new("b", DataType::Binary, false); + let schema = Schema::new(vec![string_field, binary_field]); + + let raw_string_values = vec!["foo", "bar", "baz", "quux"]; + let raw_binary_values = [ + b"foo".to_vec(), + b"bar".to_vec(), + b"baz".to_vec(), + b"quux".to_vec(), + ]; + let raw_binary_value_refs = raw_binary_values + .iter() + .map(|x| x.as_slice()) + .collect::>(); + + let string_values = StringArray::from(raw_string_values.clone()); + let binary_values = BinaryArray::from(raw_binary_value_refs); + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![Arc::new(string_values), Arc::new(binary_values)], + ) + .unwrap(); + + roundtrip(batch, Some(SMALL_SIZE / 2)); + } + + #[test] + fn arrow_writer_binary_view() { + let string_field = Field::new("a", DataType::Utf8View, false); + let binary_field = Field::new("b", DataType::BinaryView, false); + let nullable_string_field = Field::new("a", DataType::Utf8View, true); + let schema = Schema::new(vec![string_field, binary_field, nullable_string_field]); + + let raw_string_values = vec!["foo", "bar", "large payload over 12 bytes", "lulu"]; + let raw_binary_values = vec![ + b"foo".to_vec(), + b"bar".to_vec(), + b"large payload over 12 bytes".to_vec(), + b"lulu".to_vec(), + ]; + let nullable_string_values = + vec![Some("foo"), None, Some("large payload over 12 bytes"), None]; + + let string_view_values = StringViewArray::from(raw_string_values); + let binary_view_values = BinaryViewArray::from_iter_values(raw_binary_values); + let nullable_string_view_values = StringViewArray::from(nullable_string_values); + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(string_view_values), + Arc::new(binary_view_values), + Arc::new(nullable_string_view_values), + ], + ) + .unwrap(); + + roundtrip(batch.clone(), Some(SMALL_SIZE / 2)); + roundtrip(batch, None); + } + + #[test] + fn arrow_writer_binary_view_long_value() { + let string_field = Field::new("a", DataType::Utf8View, false); + let binary_field = Field::new("b", DataType::BinaryView, false); + let schema = Schema::new(vec![string_field, binary_field]); + + // There is special case validation for long values (greater than 128) + // 128 encodes as 0x80 0x00 0x00 0x00 in little endian, which should + // trigger the long-string UTF-8 validation branch in the plain decoder. + let long = "a".repeat(128); + let raw_string_values = vec!["foo", long.as_str(), "bar"]; + let raw_binary_values = vec![b"foo".to_vec(), long.as_bytes().to_vec(), b"bar".to_vec()]; + + let string_view_values: ArrayRef = Arc::new(StringViewArray::from(raw_string_values)); + let binary_view_values: ArrayRef = + Arc::new(BinaryViewArray::from_iter_values(raw_binary_values)); + + one_column_roundtrip(Arc::clone(&string_view_values), false); + one_column_roundtrip(Arc::clone(&binary_view_values), false); + + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![string_view_values, binary_view_values], + ) + .unwrap(); + + // Disable dictionary to exercise plain encoding paths in the reader. + for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] { + let props = WriterProperties::builder() + .set_writer_version(version) + .set_dictionary_enabled(false) + .build(); + roundtrip_opts(&batch, props); + } + } + + fn get_decimal_batch(precision: u8, scale: i8) -> RecordBatch { + let decimal_field = Field::new("a", DataType::Decimal128(precision, scale), false); + let schema = Schema::new(vec![decimal_field]); + + let decimal_values = vec![10_000, 50_000, 0, -100] + .into_iter() + .map(Some) + .collect::() + .with_precision_and_scale(precision, scale) + .unwrap(); + + RecordBatch::try_new(Arc::new(schema), vec![Arc::new(decimal_values)]).unwrap() + } + + #[test] + fn arrow_writer_decimal() { + // int32 to store the decimal value + let batch_int32_decimal = get_decimal_batch(5, 2); + roundtrip(batch_int32_decimal, Some(SMALL_SIZE / 2)); + // int64 to store the decimal value + let batch_int64_decimal = get_decimal_batch(12, 2); + roundtrip(batch_int64_decimal, Some(SMALL_SIZE / 2)); + // fixed_length_byte_array to store the decimal value + let batch_fixed_len_byte_array_decimal = get_decimal_batch(30, 2); + roundtrip(batch_fixed_len_byte_array_decimal, Some(SMALL_SIZE / 2)); + } + + #[test] + fn arrow_writer_complex() { + // define schema + let struct_field_d = Arc::new(Field::new("d", DataType::Float64, true)); + let struct_field_f = Arc::new(Field::new("f", DataType::Float32, true)); + let struct_field_g = Arc::new(Field::new_list( + "g", + Field::new_list_field(DataType::Int16, true), + false, + )); + let struct_field_h = Arc::new(Field::new_list( + "h", + Field::new_list_field(DataType::Int16, false), + true, + )); + let struct_field_e = Arc::new(Field::new_struct( + "e", + vec![ + struct_field_f.clone(), + struct_field_g.clone(), + struct_field_h.clone(), + ], + false, + )); + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, true), + Field::new_struct( + "c", + vec![struct_field_d.clone(), struct_field_e.clone()], + false, + ), + ]); + + // create some data + let a = Int32Array::from(vec![1, 2, 3, 4, 5]); + let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]); + let d = Float64Array::from(vec![None, None, None, Some(1.0), None]); + let f = Float32Array::from(vec![Some(0.0), None, Some(333.3), None, Some(5.25)]); + + let g_value = Int16Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); + + // Construct a buffer for value offsets, for the nested array: + // [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]] + let g_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice()); + + // Construct a list array from the above two + let g_list_data = ArrayData::builder(struct_field_g.data_type().clone()) + .len(5) + .add_buffer(g_value_offsets.clone()) + .add_child_data(g_value.to_data()) + .build() + .unwrap(); + let g = ListArray::from(g_list_data); + // The difference between g and h is that h has a null bitmap + let h_list_data = ArrayData::builder(struct_field_h.data_type().clone()) + .len(5) + .add_buffer(g_value_offsets) + .add_child_data(g_value.to_data()) + .null_bit_buffer(Some(Buffer::from([0b00011011]))) + .build() + .unwrap(); + let h = ListArray::from(h_list_data); + + let e = StructArray::from(vec![ + (struct_field_f, Arc::new(f) as ArrayRef), + (struct_field_g, Arc::new(g) as ArrayRef), + (struct_field_h, Arc::new(h) as ArrayRef), + ]); + + let c = StructArray::from(vec![ + (struct_field_d, Arc::new(d) as ArrayRef), + (struct_field_e, Arc::new(e) as ArrayRef), + ]); + + // build a record batch + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![Arc::new(a), Arc::new(b), Arc::new(c)], + ) + .unwrap(); + + roundtrip(batch.clone(), Some(SMALL_SIZE / 2)); + roundtrip(batch, Some(SMALL_SIZE / 3)); + } + + #[test] + fn arrow_writer_complex_mixed() { + // This test was added while investigating https://github.com/apache/arrow-rs/issues/244. + // It was subsequently fixed while investigating https://github.com/apache/arrow-rs/issues/245. + + // define schema + let offset_field = Arc::new(Field::new("offset", DataType::Int32, false)); + let partition_field = Arc::new(Field::new("partition", DataType::Int64, true)); + let topic_field = Arc::new(Field::new("topic", DataType::Utf8, true)); + let schema = Schema::new(vec![Field::new( + "some_nested_object", + DataType::Struct(Fields::from(vec![ + offset_field.clone(), + partition_field.clone(), + topic_field.clone(), + ])), + false, + )]); + + // create some data + let offset = Int32Array::from(vec![1, 2, 3, 4, 5]); + let partition = Int64Array::from(vec![Some(1), None, None, Some(4), Some(5)]); + let topic = StringArray::from(vec![Some("A"), None, Some("A"), Some(""), None]); + + let some_nested_object = StructArray::from(vec![ + (offset_field, Arc::new(offset) as ArrayRef), + (partition_field, Arc::new(partition) as ArrayRef), + (topic_field, Arc::new(topic) as ArrayRef), + ]); + + // build a record batch + let batch = + RecordBatch::try_new(Arc::new(schema), vec![Arc::new(some_nested_object)]).unwrap(); + + roundtrip(batch, Some(SMALL_SIZE / 2)); + } + + #[test] + fn arrow_writer_map() { + // Note: we are using the JSON Arrow reader for brevity + let json_content = r#" + {"stocks":{"long": "$AAA", "short": "$BBB"}} + {"stocks":{"long": null, "long": "$CCC", "short": null}} + {"stocks":{"hedged": "$YYY", "long": null, "short": "$D"}} + "#; + let entries_struct_type = DataType::Struct(Fields::from(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Utf8, true), + ])); + let stocks_field = Field::new( + "stocks", + DataType::Map( + Arc::new(Field::new("entries", entries_struct_type, false)), + false, + ), + true, + ); + let schema = Arc::new(Schema::new(vec![stocks_field])); + let builder = arrow::json::ReaderBuilder::new(schema).with_batch_size(64); + let mut reader = builder.build(std::io::Cursor::new(json_content)).unwrap(); + + let batch = reader.next().unwrap().unwrap(); + roundtrip(batch, None); + } + + #[test] + fn arrow_writer_2_level_struct() { + // tests writing > + let field_c = Field::new("c", DataType::Int32, true); + let field_b = Field::new("b", DataType::Struct(vec![field_c].into()), true); + let type_a = DataType::Struct(vec![field_b.clone()].into()); + let field_a = Field::new("a", type_a, true); + let schema = Schema::new(vec![field_a.clone()]); + + // create data + let c = Int32Array::from(vec![Some(1), None, Some(3), None, None, Some(6)]); + let b_data = ArrayDataBuilder::new(field_b.data_type().clone()) + .len(6) + .null_bit_buffer(Some(Buffer::from([0b00100111]))) + .add_child_data(c.into_data()) + .build() + .unwrap(); + let b = StructArray::from(b_data); + let a_data = ArrayDataBuilder::new(field_a.data_type().clone()) + .len(6) + .null_bit_buffer(Some(Buffer::from([0b00101111]))) + .add_child_data(b.into_data()) + .build() + .unwrap(); + let a = StructArray::from(a_data); + + assert_eq!(a.null_count(), 1); + assert_eq!(a.column(0).null_count(), 2); + + // build a racord batch + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap(); + + roundtrip(batch, Some(SMALL_SIZE / 2)); + } + + #[test] + fn arrow_writer_2_level_struct_non_null() { + // tests writing > + let field_c = Field::new("c", DataType::Int32, false); + let type_b = DataType::Struct(vec![field_c].into()); + let field_b = Field::new("b", type_b.clone(), false); + let type_a = DataType::Struct(vec![field_b].into()); + let field_a = Field::new("a", type_a.clone(), false); + let schema = Schema::new(vec![field_a]); + + // create data + let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]); + let b_data = ArrayDataBuilder::new(type_b) + .len(6) + .add_child_data(c.into_data()) + .build() + .unwrap(); + let b = StructArray::from(b_data); + let a_data = ArrayDataBuilder::new(type_a) + .len(6) + .add_child_data(b.into_data()) + .build() + .unwrap(); + let a = StructArray::from(a_data); + + assert_eq!(a.null_count(), 0); + assert_eq!(a.column(0).null_count(), 0); + + // build a racord batch + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap(); + + roundtrip(batch, Some(SMALL_SIZE / 2)); + } + + #[test] + fn arrow_writer_2_level_struct_mixed_null() { + // tests writing > + let field_c = Field::new("c", DataType::Int32, false); + let type_b = DataType::Struct(vec![field_c].into()); + let field_b = Field::new("b", type_b.clone(), true); + let type_a = DataType::Struct(vec![field_b].into()); + let field_a = Field::new("a", type_a.clone(), false); + let schema = Schema::new(vec![field_a]); + + // create data + let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]); + let b_data = ArrayDataBuilder::new(type_b) + .len(6) + .null_bit_buffer(Some(Buffer::from([0b00100111]))) + .add_child_data(c.into_data()) + .build() + .unwrap(); + let b = StructArray::from(b_data); + // a intentionally has no null buffer, to test that this is handled correctly + let a_data = ArrayDataBuilder::new(type_a) + .len(6) + .add_child_data(b.into_data()) + .build() + .unwrap(); + let a = StructArray::from(a_data); + + assert_eq!(a.null_count(), 0); + assert_eq!(a.column(0).null_count(), 2); + + // build a racord batch + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap(); + + roundtrip(batch, Some(SMALL_SIZE / 2)); + } + + #[test] + fn arrow_writer_2_level_struct_mixed_null_2() { + // tests writing >, where the primitive columns are non-null. + let field_c = Field::new("c", DataType::Int32, false); + let field_d = Field::new("d", DataType::FixedSizeBinary(4), false); + let field_e = Field::new( + "e", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + false, + ); + + let field_b = Field::new( + "b", + DataType::Struct(vec![field_c, field_d, field_e].into()), + false, + ); + let type_a = DataType::Struct(vec![field_b.clone()].into()); + let field_a = Field::new("a", type_a, true); + let schema = Schema::new(vec![field_a.clone()]); + + // create data + let c = Int32Array::from_iter_values(0..6); + let d = FixedSizeBinaryArray::try_from_iter( + ["aaaa", "bbbb", "cccc", "dddd", "eeee", "ffff"].into_iter(), + ) + .expect("four byte values"); + let e = Int32DictionaryArray::from_iter(["one", "two", "three", "four", "five", "one"]); + let b_data = ArrayDataBuilder::new(field_b.data_type().clone()) + .len(6) + .add_child_data(c.into_data()) + .add_child_data(d.into_data()) + .add_child_data(e.into_data()) + .build() + .unwrap(); + let b = StructArray::from(b_data); + let a_data = ArrayDataBuilder::new(field_a.data_type().clone()) + .len(6) + .null_bit_buffer(Some(Buffer::from([0b00100101]))) + .add_child_data(b.into_data()) + .build() + .unwrap(); + let a = StructArray::from(a_data); + + assert_eq!(a.null_count(), 3); + assert_eq!(a.column(0).null_count(), 0); + + // build a record batch + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap(); + + roundtrip(batch, Some(SMALL_SIZE / 2)); + } + + #[test] + fn test_fixed_size_binary_in_dict() { + fn test_fixed_size_binary_in_dict_inner() + where + K: ArrowDictionaryKeyType, + K::Native: FromPrimitive + ToPrimitive + TryFrom, + <::Native as TryFrom>::Error: std::fmt::Debug, + { + let field = Field::new( + "a", + DataType::Dictionary( + Box::new(K::DATA_TYPE), + Box::new(DataType::FixedSizeBinary(4)), + ), + false, + ); + let schema = Schema::new(vec![field]); + + let keys: Vec = vec![ + K::Native::try_from(0u8).unwrap(), + K::Native::try_from(0u8).unwrap(), + K::Native::try_from(1u8).unwrap(), + ]; + let keys = PrimitiveArray::::from_iter_values(keys); + let values = FixedSizeBinaryArray::try_from_iter( + vec![vec![0, 0, 0, 0], vec![1, 1, 1, 1]].into_iter(), + ) + .unwrap(); + + let data = DictionaryArray::::new(keys, Arc::new(values)); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(data)]).unwrap(); + roundtrip(batch, None); + } + + test_fixed_size_binary_in_dict_inner::(); + test_fixed_size_binary_in_dict_inner::(); + test_fixed_size_binary_in_dict_inner::(); + test_fixed_size_binary_in_dict_inner::(); + test_fixed_size_binary_in_dict_inner::(); + test_fixed_size_binary_in_dict_inner::(); + test_fixed_size_binary_in_dict_inner::(); + test_fixed_size_binary_in_dict_inner::(); + } + + #[test] + fn test_empty_dict() { + let struct_fields = Fields::from(vec![Field::new( + "dict", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + false, + )]); + + let schema = Schema::new(vec![Field::new_struct( + "struct", + struct_fields.clone(), + true, + )]); + let dictionary = Arc::new(DictionaryArray::new( + Int32Array::new_null(5), + Arc::new(StringArray::new_null(0)), + )); + + let s = StructArray::new( + struct_fields, + vec![dictionary], + Some(NullBuffer::new_null(5)), + ); + + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(s)]).unwrap(); + roundtrip(batch, None); + } + #[test] + fn arrow_writer_page_size() { + let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + + let mut builder = StringBuilder::with_capacity(100, 329 * 10_000); + + // Generate an array of 10 unique 10 character string + for i in 0..10 { + let value = i + .to_string() + .repeat(10) + .chars() + .take(10) + .collect::(); + + builder.append_value(value); + } + + let array = Arc::new(builder.finish()); + + let batch = RecordBatch::try_new(schema, vec![array]).unwrap(); + + let file = tempfile::tempfile().unwrap(); + + // Set everything very low so we fallback to PLAIN encoding after the first row + let props = WriterProperties::builder() + .set_data_page_size_limit(1) + .set_dictionary_page_size_limit(1) + .set_write_batch_size(1) + .build(); + + let mut writer = + ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)) + .expect("Unable to write file"); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let options = ReadOptionsBuilder::new().with_page_index().build(); + let reader = + SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap(); + + let column = reader.metadata().row_group(0).columns(); + + assert_eq!(column.len(), 1); + + // We should write one row before falling back to PLAIN encoding so there should still be a + // dictionary page. + assert!( + column[0].dictionary_page_offset().is_some(), + "Expected a dictionary page" + ); + + assert!(reader.metadata().offset_index().is_some()); + let offset_indexes = &reader.metadata().offset_index().unwrap()[0]; + + let page_locations = offset_indexes[0].page_locations.clone(); + + // We should fallback to PLAIN encoding after the first row and our max page size is 1 bytes + // so we expect one dictionary encoded page and then a page per row thereafter. + assert_eq!( + page_locations.len(), + 10, + "Expected 10 pages but got {page_locations:#?}" + ); + } + + #[test] + fn arrow_writer_float_nans() { + let f16_field = Field::new("a", DataType::Float16, false); + let f32_field = Field::new("b", DataType::Float32, false); + let f64_field = Field::new("c", DataType::Float64, false); + let schema = Schema::new(vec![f16_field, f32_field, f64_field]); + + let f16_values = (0..MEDIUM_SIZE) + .map(|i| { + Some(if i % 2 == 0 { + f16::NAN + } else { + f16::from_f32(i as f32) + }) + }) + .collect::(); + + let f32_values = (0..MEDIUM_SIZE) + .map(|i| Some(if i % 2 == 0 { f32::NAN } else { i as f32 })) + .collect::(); + + let f64_values = (0..MEDIUM_SIZE) + .map(|i| Some(if i % 2 == 0 { f64::NAN } else { i as f64 })) + .collect::(); + + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(f16_values), + Arc::new(f32_values), + Arc::new(f64_values), + ], + ) + .unwrap(); + + roundtrip(batch, None); + } + + const SMALL_SIZE: usize = 7; + const MEDIUM_SIZE: usize = 63; + + // Write the batch to parquet and read it back out, ensuring + // that what comes out is the same as what was written in + fn roundtrip(expected_batch: RecordBatch, max_row_group_size: Option) -> Vec { + let mut files = vec![]; + for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] { + let mut props = WriterProperties::builder().set_writer_version(version); + + if let Some(size) = max_row_group_size { + props = props.set_max_row_group_row_count(Some(size)) + } + + let props = props.build(); + files.push(roundtrip_opts(&expected_batch, props)) + } + files + } + + // Round trip the specified record batch with the specified writer properties, + // to an in-memory file, and validate the arrays using the specified function. + // Returns the in-memory file. + fn roundtrip_opts_with_array_validation( + expected_batch: &RecordBatch, + props: WriterProperties, + validate: F, + ) -> Bytes + where + F: Fn(&ArrayData, &ArrayData), + { + let mut file = vec![]; + + let mut writer = ArrowWriter::try_new(&mut file, expected_batch.schema(), Some(props)) + .expect("Unable to write file"); + writer.write(expected_batch).unwrap(); + writer.close().unwrap(); + + let file = Bytes::from(file); + let mut record_batch_reader = + ParquetRecordBatchReader::try_new(file.clone(), 1024).unwrap(); + + let actual_batch = record_batch_reader + .next() + .expect("No batch found") + .expect("Unable to get batch"); + + assert_eq!(expected_batch.schema(), actual_batch.schema()); + assert_eq!(expected_batch.num_columns(), actual_batch.num_columns()); + assert_eq!(expected_batch.num_rows(), actual_batch.num_rows()); + for i in 0..expected_batch.num_columns() { + let expected_data = expected_batch.column(i).to_data(); + let actual_data = actual_batch.column(i).to_data(); + validate(&expected_data, &actual_data); + } + + file + } + + fn roundtrip_opts(expected_batch: &RecordBatch, props: WriterProperties) -> Bytes { + roundtrip_opts_with_array_validation(expected_batch, props, |a, b| { + a.validate_full().expect("valid expected data"); + b.validate_full().expect("valid actual data"); + assert_eq!(a, b) + }) + } + + struct RoundTripOptions { + values: ArrayRef, + schema: SchemaRef, + bloom_filter: bool, + bloom_filter_ndv: Option, + bloom_filter_position: BloomFilterPosition, + } + + impl RoundTripOptions { + fn new(values: ArrayRef, nullable: bool) -> Self { + let data_type = values.data_type().clone(); + let schema = Schema::new(vec![Field::new("col", data_type, nullable)]); + Self { + values, + schema: Arc::new(schema), + bloom_filter: false, + bloom_filter_ndv: None, + bloom_filter_position: BloomFilterPosition::AfterRowGroup, + } + } + } + + fn one_column_roundtrip(values: ArrayRef, nullable: bool) -> Vec { + one_column_roundtrip_with_options(RoundTripOptions::new(values, nullable)) + } + + fn one_column_roundtrip_with_schema(values: ArrayRef, schema: SchemaRef) -> Vec { + let mut options = RoundTripOptions::new(values, false); + options.schema = schema; + one_column_roundtrip_with_options(options) + } + + fn one_column_roundtrip_with_options(options: RoundTripOptions) -> Vec { + let RoundTripOptions { + values, + schema, + bloom_filter, + bloom_filter_ndv, + bloom_filter_position, + } = options; + + let encodings = match values.data_type() { + DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary => { + vec![ + Encoding::PLAIN, + Encoding::DELTA_BYTE_ARRAY, + Encoding::DELTA_LENGTH_BYTE_ARRAY, + ] + } + DataType::Int64 + | DataType::Int32 + | DataType::Int16 + | DataType::Int8 + | DataType::UInt64 + | DataType::UInt32 + | DataType::UInt16 + | DataType::UInt8 => vec![ + Encoding::PLAIN, + Encoding::DELTA_BINARY_PACKED, + Encoding::BYTE_STREAM_SPLIT, + ], + DataType::Float32 | DataType::Float64 => { + vec![Encoding::PLAIN, Encoding::BYTE_STREAM_SPLIT] + } + _ => vec![Encoding::PLAIN], + }; + + let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap(); + + let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10]; + + let mut files = vec![]; + for dictionary_size in [0, 1, 1024] { + for encoding in &encodings { + for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] { + for row_group_size in row_group_sizes { + let mut builder = WriterProperties::builder() + .set_writer_version(version) + .set_max_row_group_row_count(Some(row_group_size)) + .set_dictionary_enabled(dictionary_size != 0) + .set_dictionary_page_size_limit(dictionary_size.max(1)) + .set_encoding(*encoding) + .set_bloom_filter_enabled(bloom_filter) + .set_bloom_filter_position(bloom_filter_position); + if let Some(ndv) = bloom_filter_ndv { + builder = builder.set_bloom_filter_max_ndv(ndv); + } + let props = builder.build(); + + files.push(roundtrip_opts(&expected_batch, props)) + } + } + } + } + files + } + + fn values_required(iter: I) -> Vec + where + A: From> + Array + 'static, + I: IntoIterator, + { + let raw_values: Vec<_> = iter.into_iter().collect(); + let values = Arc::new(A::from(raw_values)); + one_column_roundtrip(values, false) + } + + fn values_optional(iter: I) -> Vec + where + A: From>> + Array + 'static, + I: IntoIterator, + { + let optional_raw_values: Vec<_> = iter + .into_iter() + .enumerate() + .map(|(i, v)| if i % 2 == 0 { None } else { Some(v) }) + .collect(); + let optional_values = Arc::new(A::from(optional_raw_values)); + one_column_roundtrip(optional_values, true) + } + + fn required_and_optional(iter: I) + where + A: From> + From>> + Array + 'static, + I: IntoIterator + Clone, + { + values_required::(iter.clone()); + values_optional::(iter); + } + + fn check_bloom_filter( + files: Vec, + file_column: String, + positive_values: Vec, + negative_values: Vec, + ) { + files.into_iter().take(1).for_each(|file| { + let file_reader = SerializedFileReader::new_with_options( + file, + ReadOptionsBuilder::new() + .with_reader_properties( + ReaderProperties::builder() + .set_read_bloom_filter(true) + .build(), + ) + .build(), + ) + .expect("Unable to open file as Parquet"); + let metadata = file_reader.metadata(); + + // Gets bloom filters from all row groups. + let mut bloom_filters: Vec<_> = vec![]; + for (ri, row_group) in metadata.row_groups().iter().enumerate() { + if let Some((column_index, _)) = row_group + .columns() + .iter() + .enumerate() + .find(|(_, column)| column.column_path().string() == file_column) + { + let row_group_reader = file_reader + .get_row_group(ri) + .expect("Unable to read row group"); + if let Some(sbbf) = row_group_reader.get_column_bloom_filter(column_index) { + bloom_filters.push(sbbf.clone()); + } else { + panic!("No bloom filter for column named {file_column} found"); + } + } else { + panic!("No column named {file_column} found"); + } + } + + positive_values.iter().for_each(|value| { + let found = bloom_filters.iter().find(|sbbf| sbbf.check(value)); + assert!( + found.is_some(), + "{}", + format!("Value {:?} should be in bloom filter", value.as_bytes()) + ); + }); + + negative_values.iter().for_each(|value| { + let found = bloom_filters.iter().find(|sbbf| sbbf.check(value)); + assert!( + found.is_none(), + "{}", + format!("Value {:?} should not be in bloom filter", value.as_bytes()) + ); + }); + }); + } + + #[test] + fn all_null_primitive_single_column() { + let values = Arc::new(Int32Array::from(vec![None; SMALL_SIZE])); + one_column_roundtrip(values, true); + } + #[test] + fn null_single_column() { + let values = Arc::new(NullArray::new(SMALL_SIZE)); + one_column_roundtrip(values, true); + // null arrays are always nullable, a test with non-nullable nulls fails + } + + #[test] + fn bool_single_column() { + required_and_optional::( + [true, false].iter().cycle().copied().take(SMALL_SIZE), + ); + } + + #[test] + fn bool_large_single_column() { + let values = Arc::new( + [None, Some(true), Some(false)] + .iter() + .cycle() + .copied() + .take(200_000) + .collect::(), + ); + let schema = Schema::new(vec![Field::new("col", values.data_type().clone(), true)]); + let expected_batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap(); + let file = tempfile::tempfile().unwrap(); + + let mut writer = + ArrowWriter::try_new(file.try_clone().unwrap(), expected_batch.schema(), None) + .expect("Unable to write file"); + writer.write(&expected_batch).unwrap(); + writer.close().unwrap(); + } + + #[test] + fn check_page_offset_index_with_nan() { + let values = Arc::new(Float64Array::from(vec![f64::NAN; 10])); + let schema = Schema::new(vec![Field::new("col", DataType::Float64, true)]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap(); + + let mut out = Vec::with_capacity(1024); + let mut writer = + ArrowWriter::try_new(&mut out, batch.schema(), None).expect("Unable to write file"); + writer.write(&batch).unwrap(); + let file_meta_data = writer.close().unwrap(); + for row_group in file_meta_data.row_groups() { + for column in row_group.columns() { + assert!(column.offset_index_offset().is_some()); + assert!(column.offset_index_length().is_some()); + assert!(column.column_index_offset().is_none()); + assert!(column.column_index_length().is_none()); + } + } + } + + #[test] + fn i8_single_column() { + required_and_optional::(0..SMALL_SIZE as i8); + } + + #[test] + fn i16_single_column() { + required_and_optional::(0..SMALL_SIZE as i16); + } + + #[test] + fn i32_single_column() { + required_and_optional::(0..SMALL_SIZE as i32); + } + + #[test] + fn i64_single_column() { + required_and_optional::(0..SMALL_SIZE as i64); + } + + #[test] + fn u8_single_column() { + required_and_optional::(0..SMALL_SIZE as u8); + } + + #[test] + fn u16_single_column() { + required_and_optional::(0..SMALL_SIZE as u16); + } + + #[test] + fn u32_single_column() { + required_and_optional::(0..SMALL_SIZE as u32); + } + + #[test] + fn u64_single_column() { + required_and_optional::(0..SMALL_SIZE as u64); + } + + #[test] + fn f32_single_column() { + required_and_optional::((0..SMALL_SIZE).map(|i| i as f32)); + } + + #[test] + fn f64_single_column() { + required_and_optional::((0..SMALL_SIZE).map(|i| i as f64)); + } + + // The timestamp array types don't implement From> because they need the timezone + // argument, and they also doesn't support building from a Vec>, so call + // one_column_roundtrip manually instead of calling required_and_optional for these tests. + + #[test] + fn timestamp_second_single_column() { + let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect(); + let values = Arc::new(TimestampSecondArray::from(raw_values)); + + one_column_roundtrip(values, false); + } + + #[test] + fn timestamp_millisecond_single_column() { + let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect(); + let values = Arc::new(TimestampMillisecondArray::from(raw_values)); + + one_column_roundtrip(values, false); + } + + #[test] + fn timestamp_microsecond_single_column() { + let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect(); + let values = Arc::new(TimestampMicrosecondArray::from(raw_values)); + + one_column_roundtrip(values, false); + } + + #[test] + fn timestamp_nanosecond_single_column() { + let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect(); + let values = Arc::new(TimestampNanosecondArray::from(raw_values)); + + one_column_roundtrip(values, false); + } + + #[test] + fn date32_single_column() { + required_and_optional::(0..SMALL_SIZE as i32); + } + + #[test] + fn date64_single_column() { + // Date64 must be a multiple of 86400000, see ARROW-10925 + required_and_optional::( + (0..(SMALL_SIZE as i64 * 86400000)).step_by(86400000), + ); + } + + #[test] + fn time32_second_single_column() { + required_and_optional::(0..SMALL_SIZE as i32); + } + + #[test] + fn time32_millisecond_single_column() { + required_and_optional::(0..SMALL_SIZE as i32); + } + + #[test] + fn time64_microsecond_single_column() { + required_and_optional::(0..SMALL_SIZE as i64); + } + + #[test] + fn time64_nanosecond_single_column() { + required_and_optional::(0..SMALL_SIZE as i64); + } + + #[test] + fn duration_second_single_column() { + required_and_optional::(0..SMALL_SIZE as i64); + } + + #[test] + fn duration_millisecond_single_column() { + required_and_optional::(0..SMALL_SIZE as i64); + } + + #[test] + fn duration_microsecond_single_column() { + required_and_optional::(0..SMALL_SIZE as i64); + } + + #[test] + fn duration_nanosecond_single_column() { + required_and_optional::(0..SMALL_SIZE as i64); + } + + #[test] + fn interval_year_month_single_column() { + required_and_optional::(0..SMALL_SIZE as i32); + } + + #[test] + fn interval_day_time_single_column() { + required_and_optional::(vec![ + IntervalDayTime::new(0, 1), + IntervalDayTime::new(0, 3), + IntervalDayTime::new(3, -2), + IntervalDayTime::new(-200, 4), + ]); + } + + #[test] + #[should_panic( + expected = "Attempting to write an Arrow interval type MonthDayNano to parquet that is not yet implemented" + )] + fn interval_month_day_nano_single_column() { + required_and_optional::(vec![ + IntervalMonthDayNano::new(0, 1, 5), + IntervalMonthDayNano::new(0, 3, 2), + IntervalMonthDayNano::new(3, -2, -5), + IntervalMonthDayNano::new(-200, 4, -1), + ]); + } + + #[test] + fn binary_single_column() { + let one_vec: Vec = (0..SMALL_SIZE as u8).collect(); + let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect(); + let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice()); + + // BinaryArrays can't be built from Vec>, so only call `values_required` + values_required::(many_vecs_iter); + } + + #[test] + fn binary_view_single_column() { + let one_vec: Vec = (0..SMALL_SIZE as u8).collect(); + let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect(); + let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice()); + + // BinaryArrays can't be built from Vec>, so only call `values_required` + values_required::(many_vecs_iter); + } + + #[test] + fn i32_column_bloom_filter_at_end() { + let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32)); + let mut options = RoundTripOptions::new(array, false); + options.bloom_filter = true; + options.bloom_filter_position = BloomFilterPosition::End; + + let files = one_column_roundtrip_with_options(options); + check_bloom_filter( + files, + "col".to_string(), + (0..SMALL_SIZE as i32).collect(), + (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(), + ); + } + + #[test] + fn i32_column_bloom_filter() { + let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32)); + let mut options = RoundTripOptions::new(array, false); + options.bloom_filter = true; + + let files = one_column_roundtrip_with_options(options); + check_bloom_filter( + files, + "col".to_string(), + (0..SMALL_SIZE as i32).collect(), + (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(), + ); + } + + /// Test that bloom filter folding produces correct results even when + /// the configured NDV differs significantly from actual NDV. + /// A large NDV means a larger initial filter that gets folded down; + /// a small NDV means a smaller initial filter. + #[test] + fn i32_column_bloom_filter_fixed_ndv() { + let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32)); + + // NDV much larger than actual distinct values — tests folding a large filter down + let mut options = RoundTripOptions::new(array.clone(), false); + options.bloom_filter = true; + options.bloom_filter_ndv = Some(1_000_000); + + let files = one_column_roundtrip_with_options(options); + check_bloom_filter( + files, + "col".to_string(), + (0..SMALL_SIZE as i32).collect(), + (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(), + ); + + // NDV smaller than actual distinct values — tests the underestimate path + let mut options = RoundTripOptions::new(array, false); + options.bloom_filter = true; + options.bloom_filter_ndv = Some(3); + + let files = one_column_roundtrip_with_options(options); + check_bloom_filter( + files, + "col".to_string(), + (0..SMALL_SIZE as i32).collect(), + (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(), + ); + } + + #[test] + fn binary_column_bloom_filter() { + let one_vec: Vec = (0..SMALL_SIZE as u8).collect(); + let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect(); + let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice()); + + let array = Arc::new(BinaryArray::from_iter_values(many_vecs_iter)); + let mut options = RoundTripOptions::new(array, false); + options.bloom_filter = true; + + let files = one_column_roundtrip_with_options(options); + check_bloom_filter( + files, + "col".to_string(), + many_vecs, + vec![vec![(SMALL_SIZE + 1) as u8]], + ); + } + + #[test] + fn empty_string_null_column_bloom_filter() { + let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect(); + let raw_strs = raw_values.iter().map(|s| s.as_str()); + + let array = Arc::new(StringArray::from_iter_values(raw_strs)); + let mut options = RoundTripOptions::new(array, false); + options.bloom_filter = true; + + let files = one_column_roundtrip_with_options(options); + + let optional_raw_values: Vec<_> = raw_values + .iter() + .enumerate() + .filter_map(|(i, v)| if i % 2 == 0 { None } else { Some(v.as_str()) }) + .collect(); + // For null slots, empty string should not be in bloom filter. + check_bloom_filter(files, "col".to_string(), optional_raw_values, vec![""]); + } + + #[test] + fn large_binary_single_column() { + let one_vec: Vec = (0..SMALL_SIZE as u8).collect(); + let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect(); + let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice()); + + // LargeBinaryArrays can't be built from Vec>, so only call `values_required` + values_required::(many_vecs_iter); + } + + #[test] + fn fixed_size_binary_single_column() { + let mut builder = FixedSizeBinaryBuilder::new(4); + builder.append_value(b"0123").unwrap(); + builder.append_null(); + builder.append_value(b"8910").unwrap(); + builder.append_value(b"1112").unwrap(); + let array = Arc::new(builder.finish()); + + one_column_roundtrip(array, true); + } + + #[test] + fn string_single_column() { + let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect(); + let raw_strs = raw_values.iter().map(|s| s.as_str()); + + required_and_optional::(raw_strs); + } + + #[test] + fn large_string_single_column() { + let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect(); + let raw_strs = raw_values.iter().map(|s| s.as_str()); + + required_and_optional::(raw_strs); + } + + #[test] + fn string_view_single_column() { + let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect(); + let raw_strs = raw_values.iter().map(|s| s.as_str()); + + required_and_optional::(raw_strs); + } + + #[test] + fn null_list_single_column() { + let null_field = Field::new_list_field(DataType::Null, true); + let list_field = Field::new("emptylist", DataType::List(Arc::new(null_field)), true); + + let schema = Schema::new(vec![list_field]); + + // Build [[], null, [null, null]] + let a_values = NullArray::new(2); + let a_value_offsets = arrow::buffer::Buffer::from([0, 0, 0, 2].to_byte_slice()); + let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field( + DataType::Null, + true, + )))) + .len(3) + .add_buffer(a_value_offsets) + .null_bit_buffer(Some(Buffer::from([0b00000101]))) + .add_child_data(a_values.into_data()) + .build() + .unwrap(); + + let a = ListArray::from(a_list_data); + + assert!(a.is_valid(0)); + assert!(!a.is_valid(1)); + assert!(a.is_valid(2)); + + assert_eq!(a.value(0).len(), 0); + assert_eq!(a.value(2).len(), 2); + assert_eq!(a.value(2).logical_nulls().unwrap().null_count(), 2); + + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap(); + roundtrip(batch, None); + } + + #[test] + fn list_single_column() { + let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); + let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice()); + let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field( + DataType::Int32, + false, + )))) + .len(5) + .add_buffer(a_value_offsets) + .null_bit_buffer(Some(Buffer::from([0b00011011]))) + .add_child_data(a_values.into_data()) + .build() + .unwrap(); + + assert_eq!(a_list_data.null_count(), 1); + + let a = ListArray::from(a_list_data); + let values = Arc::new(a); + + one_column_roundtrip(values, true); + } + + #[test] + fn large_list_single_column() { + let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); + let a_value_offsets = arrow::buffer::Buffer::from([0i64, 1, 3, 3, 6, 10].to_byte_slice()); + let a_list_data = ArrayData::builder(DataType::LargeList(Arc::new(Field::new( + "large_item", + DataType::Int32, + true, + )))) + .len(5) + .add_buffer(a_value_offsets) + .add_child_data(a_values.into_data()) + .null_bit_buffer(Some(Buffer::from([0b00011011]))) + .build() + .unwrap(); + + // I think this setup is incorrect because this should pass + assert_eq!(a_list_data.null_count(), 1); + + let a = LargeListArray::from(a_list_data); + let values = Arc::new(a); + + one_column_roundtrip(values, true); + } + + #[test] + fn list_nested_nulls() { + use arrow::datatypes::Int32Type; + let data = vec![ + Some(vec![Some(1)]), + Some(vec![Some(2), Some(3)]), + None, + Some(vec![Some(4), Some(5), None]), + Some(vec![None]), + Some(vec![Some(6), Some(7)]), + ]; + + let list = ListArray::from_iter_primitive::(data.clone()); + one_column_roundtrip(Arc::new(list), true); + + let list = LargeListArray::from_iter_primitive::(data); + one_column_roundtrip(Arc::new(list), true); + } + + #[test] + fn list_utf8_view_selective_padding_roundtrip() { + let item = Arc::new(Field::new_list_field(DataType::Utf8View, true)); + let mut builder = ListBuilder::new(StringViewBuilder::new()).with_field(item); + builder.values().append_value("a"); + builder.values().append_null(); + builder.append(true); + // The null parent list covers selective padding dropping values below + // the list definition level while preserving the preceding item null. + builder.append(false); + // The long string covers the non-inlined Utf8View buffer path. + builder.values().append_value("large payload over 12 bytes"); + builder.append(true); + + one_column_roundtrip(Arc::new(builder.finish()), true); + } + + #[test] + fn struct_single_column() { + let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); + let struct_field_a = Arc::new(Field::new("f", DataType::Int32, false)); + let s = StructArray::from(vec![(struct_field_a, Arc::new(a_values) as ArrayRef)]); + + let values = Arc::new(s); + one_column_roundtrip(values, false); + } + + #[test] + fn list_and_map_coerced_names() { + // Create map and list with non-Parquet naming + let list_field = + Field::new_list("my_list", Field::new("item", DataType::Int32, false), false); + let map_field = Field::new_map( + "my_map", + "entries", + Field::new("keys", DataType::Int32, false), + Field::new("values", DataType::Int32, true), + false, + true, + ); + + let list_array = create_random_array(&list_field, 100, 0.0, 0.0).unwrap(); + let map_array = create_random_array(&map_field, 100, 0.0, 0.0).unwrap(); + + let arrow_schema = Arc::new(Schema::new(vec![list_field, map_field])); + + // Write data to Parquet but coerce names to match spec + let props = Some(WriterProperties::builder().set_coerce_types(true).build()); + let file = tempfile::tempfile().unwrap(); + let mut writer = + ArrowWriter::try_new(file.try_clone().unwrap(), arrow_schema.clone(), props).unwrap(); + + let batch = RecordBatch::try_new(arrow_schema, vec![list_array, map_array]).unwrap(); + writer.write(&batch).unwrap(); + let file_metadata = writer.close().unwrap(); + + let schema = file_metadata.file_metadata().schema(); + // Coerced name of "item" should be "element" + let list_field = &schema.get_fields()[0].get_fields()[0]; + assert_eq!(list_field.get_fields()[0].name(), "element"); + + let map_field = &schema.get_fields()[1].get_fields()[0]; + // Coerced name of "entries" should be "key_value" + assert_eq!(map_field.name(), "key_value"); + // Coerced name of "keys" should be "key" + assert_eq!(map_field.get_fields()[0].name(), "key"); + // Coerced name of "values" should be "value" + assert_eq!(map_field.get_fields()[1].name(), "value"); + + // Double check schema after reading from the file + let reader = SerializedFileReader::new(file).unwrap(); + let file_schema = reader.metadata().file_metadata().schema(); + let fields = file_schema.get_fields(); + let list_field = &fields[0].get_fields()[0]; + assert_eq!(list_field.get_fields()[0].name(), "element"); + let map_field = &fields[1].get_fields()[0]; + assert_eq!(map_field.name(), "key_value"); + assert_eq!(map_field.get_fields()[0].name(), "key"); + assert_eq!(map_field.get_fields()[1].name(), "value"); + } + + #[test] + fn fallback_flush_data_page() { + //tests if the Fallback::flush_data_page clears all buffers correctly + let raw_values: Vec<_> = (0..MEDIUM_SIZE).map(|i| i.to_string()).collect(); + let values = Arc::new(StringArray::from(raw_values)); + let encodings = vec![ + Encoding::DELTA_BYTE_ARRAY, + Encoding::DELTA_LENGTH_BYTE_ARRAY, + ]; + let data_type = values.data_type().clone(); + let schema = Arc::new(Schema::new(vec![Field::new("col", data_type, false)])); + let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap(); + + let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10]; + let data_page_size_limit: usize = 32; + let write_batch_size: usize = 16; + + for encoding in &encodings { + for row_group_size in row_group_sizes { + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_2_0) + .set_max_row_group_row_count(Some(row_group_size)) + .set_dictionary_enabled(false) + .set_encoding(*encoding) + .set_data_page_size_limit(data_page_size_limit) + .set_write_batch_size(write_batch_size) + .build(); + + roundtrip_opts_with_array_validation(&expected_batch, props, |a, b| { + let string_array_a = StringArray::from(a.clone()); + let string_array_b = StringArray::from(b.clone()); + let vec_a: Vec<&str> = string_array_a.iter().map(|v| v.unwrap()).collect(); + let vec_b: Vec<&str> = string_array_b.iter().map(|v| v.unwrap()).collect(); + assert_eq!( + vec_a, vec_b, + "failed for encoder: {encoding:?} and row_group_size: {row_group_size:?}" + ); + }); + } + } + } + + #[test] + fn arrow_writer_string_dictionary() { + // define schema + #[allow(deprecated)] + let schema = Arc::new(Schema::new(vec![Field::new_dict( + "dictionary", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + true, + 42, + true, + )])); + + // create some data + let d: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")] + .iter() + .copied() + .collect(); + + // build a record batch + one_column_roundtrip_with_schema(Arc::new(d), schema); + } + + #[test] + fn arrow_writer_test_type_compatibility() { + fn ensure_compatible_write(array1: T1, array2: T2, expected_result: T1) + where + T1: Array + 'static, + T2: Array + 'static, + { + let schema1 = Arc::new(Schema::new(vec![Field::new( + "a", + array1.data_type().clone(), + false, + )])); + + let file = tempfile().unwrap(); + let mut writer = + ArrowWriter::try_new(file.try_clone().unwrap(), schema1.clone(), None).unwrap(); + + let rb1 = RecordBatch::try_new(schema1.clone(), vec![Arc::new(array1)]).unwrap(); + writer.write(&rb1).unwrap(); + + let schema2 = Arc::new(Schema::new(vec![Field::new( + "a", + array2.data_type().clone(), + false, + )])); + let rb2 = RecordBatch::try_new(schema2, vec![Arc::new(array2)]).unwrap(); + writer.write(&rb2).unwrap(); + + writer.close().unwrap(); + + let mut record_batch_reader = + ParquetRecordBatchReader::try_new(file.try_clone().unwrap(), 1024).unwrap(); + let actual_batch = record_batch_reader.next().unwrap().unwrap(); + + let expected_batch = + RecordBatch::try_new(schema1, vec![Arc::new(expected_result)]).unwrap(); + assert_eq!(actual_batch, expected_batch); + } + + // check compatibility between native and dictionaries + + ensure_compatible_write( + DictionaryArray::new( + UInt8Array::from_iter_values(vec![0]), + Arc::new(StringArray::from_iter_values(vec!["parquet"])), + ), + StringArray::from_iter_values(vec!["barquet"]), + DictionaryArray::new( + UInt8Array::from_iter_values(vec![0, 1]), + Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])), + ), + ); + + ensure_compatible_write( + StringArray::from_iter_values(vec!["parquet"]), + DictionaryArray::new( + UInt8Array::from_iter_values(vec![0]), + Arc::new(StringArray::from_iter_values(vec!["barquet"])), + ), + StringArray::from_iter_values(vec!["parquet", "barquet"]), + ); + + // check compatibility between dictionaries with different key types + + ensure_compatible_write( + DictionaryArray::new( + UInt8Array::from_iter_values(vec![0]), + Arc::new(StringArray::from_iter_values(vec!["parquet"])), + ), + DictionaryArray::new( + UInt16Array::from_iter_values(vec![0]), + Arc::new(StringArray::from_iter_values(vec!["barquet"])), + ), + DictionaryArray::new( + UInt8Array::from_iter_values(vec![0, 1]), + Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])), + ), + ); + + // check compatibility between dictionaries with different value types + ensure_compatible_write( + DictionaryArray::new( + UInt8Array::from_iter_values(vec![0]), + Arc::new(StringArray::from_iter_values(vec!["parquet"])), + ), + DictionaryArray::new( + UInt8Array::from_iter_values(vec![0]), + Arc::new(LargeStringArray::from_iter_values(vec!["barquet"])), + ), + DictionaryArray::new( + UInt8Array::from_iter_values(vec![0, 1]), + Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])), + ), + ); + + // check compatibility between a dictionary and a native array with a different type + ensure_compatible_write( + DictionaryArray::new( + UInt8Array::from_iter_values(vec![0]), + Arc::new(StringArray::from_iter_values(vec!["parquet"])), + ), + LargeStringArray::from_iter_values(vec!["barquet"]), + DictionaryArray::new( + UInt8Array::from_iter_values(vec![0, 1]), + Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])), + ), + ); + + // check compatibility for string types + + ensure_compatible_write( + StringArray::from_iter_values(vec!["parquet"]), + LargeStringArray::from_iter_values(vec!["barquet"]), + StringArray::from_iter_values(vec!["parquet", "barquet"]), + ); + + ensure_compatible_write( + LargeStringArray::from_iter_values(vec!["parquet"]), + StringArray::from_iter_values(vec!["barquet"]), + LargeStringArray::from_iter_values(vec!["parquet", "barquet"]), + ); + + ensure_compatible_write( + StringArray::from_iter_values(vec!["parquet"]), + StringViewArray::from_iter_values(vec!["barquet"]), + StringArray::from_iter_values(vec!["parquet", "barquet"]), + ); + + ensure_compatible_write( + StringViewArray::from_iter_values(vec!["parquet"]), + StringArray::from_iter_values(vec!["barquet"]), + StringViewArray::from_iter_values(vec!["parquet", "barquet"]), + ); + + ensure_compatible_write( + LargeStringArray::from_iter_values(vec!["parquet"]), + StringViewArray::from_iter_values(vec!["barquet"]), + LargeStringArray::from_iter_values(vec!["parquet", "barquet"]), + ); + + ensure_compatible_write( + StringViewArray::from_iter_values(vec!["parquet"]), + LargeStringArray::from_iter_values(vec!["barquet"]), + StringViewArray::from_iter_values(vec!["parquet", "barquet"]), + ); + + // check compatibility for binary types + + ensure_compatible_write( + BinaryArray::from_iter_values(vec![b"parquet"]), + LargeBinaryArray::from_iter_values(vec![b"barquet"]), + BinaryArray::from_iter_values(vec![b"parquet", b"barquet"]), + ); + + ensure_compatible_write( + LargeBinaryArray::from_iter_values(vec![b"parquet"]), + BinaryArray::from_iter_values(vec![b"barquet"]), + LargeBinaryArray::from_iter_values(vec![b"parquet", b"barquet"]), + ); + + ensure_compatible_write( + BinaryArray::from_iter_values(vec![b"parquet"]), + BinaryViewArray::from_iter_values(vec![b"barquet"]), + BinaryArray::from_iter_values(vec![b"parquet", b"barquet"]), + ); + + ensure_compatible_write( + BinaryViewArray::from_iter_values(vec![b"parquet"]), + BinaryArray::from_iter_values(vec![b"barquet"]), + BinaryViewArray::from_iter_values(vec![b"parquet", b"barquet"]), + ); + + ensure_compatible_write( + BinaryViewArray::from_iter_values(vec![b"parquet"]), + LargeBinaryArray::from_iter_values(vec![b"barquet"]), + BinaryViewArray::from_iter_values(vec![b"parquet", b"barquet"]), + ); + + ensure_compatible_write( + LargeBinaryArray::from_iter_values(vec![b"parquet"]), + BinaryViewArray::from_iter_values(vec![b"barquet"]), + LargeBinaryArray::from_iter_values(vec![b"parquet", b"barquet"]), + ); + + // check compatibility for list types + + let list_field_metadata = HashMap::from_iter(vec![( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )]); + let list_field = Field::new_list_field(DataType::Int32, false); + + let values1 = Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4])); + let offsets1 = OffsetBuffer::new(vec![0, 2, 5].into()); + + let values2 = Arc::new(Int32Array::from(vec![5, 6, 7, 8, 9])); + let offsets2 = OffsetBuffer::new(vec![0, 3, 5].into()); + + let values_expected = Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9])); + let offsets_expected = OffsetBuffer::new(vec![0, 2, 5, 8, 10].into()); + + ensure_compatible_write( + // when the initial schema has the metadata ... + ListArray::try_new( + Arc::new( + list_field + .clone() + .with_metadata(list_field_metadata.clone()), + ), + offsets1, + values1, + None, + ) + .unwrap(), + // ... and some intermediate schema doesn't have the metadata + ListArray::try_new(Arc::new(list_field.clone()), offsets2, values2, None).unwrap(), + // ... the write will still go through, and the resulting schema will inherit the initial metadata + ListArray::try_new( + Arc::new( + list_field + .clone() + .with_metadata(list_field_metadata.clone()), + ), + offsets_expected, + values_expected, + None, + ) + .unwrap(), + ); + } + + #[test] + fn arrow_writer_primitive_dictionary() { + // define schema + #[allow(deprecated)] + let schema = Arc::new(Schema::new(vec![Field::new_dict( + "dictionary", + DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::UInt32)), + true, + 42, + true, + )])); + + // create some data + let mut builder = PrimitiveDictionaryBuilder::::new(); + builder.append(12345678).unwrap(); + builder.append_null(); + builder.append(22345678).unwrap(); + builder.append(12345678).unwrap(); + let d = builder.finish(); + + one_column_roundtrip_with_schema(Arc::new(d), schema); + } + + #[test] + fn arrow_writer_decimal32_dictionary() { + let integers = vec![12345, 56789, 34567]; + + let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]); + + let values = Decimal32Array::from(integers.clone()) + .with_precision_and_scale(5, 2) + .unwrap(); + + let array = DictionaryArray::new(keys, Arc::new(values)); + one_column_roundtrip(Arc::new(array.clone()), true); + + let values = Decimal32Array::from(integers) + .with_precision_and_scale(9, 2) + .unwrap(); + + let array = array.with_values(Arc::new(values)); + one_column_roundtrip(Arc::new(array), true); + } + + #[test] + fn arrow_writer_decimal64_dictionary() { + let integers = vec![12345, 56789, 34567]; + + let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]); + + let values = Decimal64Array::from(integers.clone()) + .with_precision_and_scale(5, 2) + .unwrap(); + + let array = DictionaryArray::new(keys, Arc::new(values)); + one_column_roundtrip(Arc::new(array.clone()), true); + + let values = Decimal64Array::from(integers) + .with_precision_and_scale(12, 2) + .unwrap(); + + let array = array.with_values(Arc::new(values)); + one_column_roundtrip(Arc::new(array), true); + } + + #[test] + fn arrow_writer_decimal128_dictionary() { + let integers = vec![12345, 56789, 34567]; + + let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]); + + let values = Decimal128Array::from(integers.clone()) + .with_precision_and_scale(5, 2) + .unwrap(); + + let array = DictionaryArray::new(keys, Arc::new(values)); + one_column_roundtrip(Arc::new(array.clone()), true); + + let values = Decimal128Array::from(integers) + .with_precision_and_scale(12, 2) + .unwrap(); + + let array = array.with_values(Arc::new(values)); + one_column_roundtrip(Arc::new(array), true); + } + + #[test] + fn arrow_writer_decimal256_dictionary() { + let integers = vec![ + i256::from_i128(12345), + i256::from_i128(56789), + i256::from_i128(34567), + ]; + + let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]); + + let values = Decimal256Array::from(integers.clone()) + .with_precision_and_scale(5, 2) + .unwrap(); + + let array = DictionaryArray::new(keys, Arc::new(values)); + one_column_roundtrip(Arc::new(array.clone()), true); + + let values = Decimal256Array::from(integers) + .with_precision_and_scale(12, 2) + .unwrap(); + + let array = array.with_values(Arc::new(values)); + one_column_roundtrip(Arc::new(array), true); + } + + #[test] + fn arrow_writer_string_dictionary_unsigned_index() { + // define schema + #[allow(deprecated)] + let schema = Arc::new(Schema::new(vec![Field::new_dict( + "dictionary", + DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)), + true, + 42, + true, + )])); + + // create some data + let d: UInt8DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")] + .iter() + .copied() + .collect(); + + one_column_roundtrip_with_schema(Arc::new(d), schema); + } + + #[test] + fn u32_min_max() { + // check values roundtrip through parquet + let src = [ + u32::MIN, + 1, + (i32::MAX as u32) - 1, + i32::MAX as u32, + (i32::MAX as u32) + 1, + u32::MAX - 1, + u32::MAX, + ]; + let values = Arc::new(UInt32Array::from_iter_values(src.iter().cloned())); + let files = one_column_roundtrip(values, false); + + for file in files { + // check statistics are valid + let reader = SerializedFileReader::new(file).unwrap(); + let metadata = reader.metadata(); + + let mut row_offset = 0; + for row_group in metadata.row_groups() { + assert_eq!(row_group.num_columns(), 1); + let column = row_group.column(0); + + let num_values = column.num_values() as usize; + let src_slice = &src[row_offset..row_offset + num_values]; + row_offset += column.num_values() as usize; + + let stats = column.statistics().unwrap(); + if let Statistics::Int32(stats) = stats { + assert_eq!( + *stats.min_opt().unwrap() as u32, + *src_slice.iter().min().unwrap() + ); + assert_eq!( + *stats.max_opt().unwrap() as u32, + *src_slice.iter().max().unwrap() + ); + } else { + panic!("Statistics::Int32 missing") + } + } + } + } + + #[test] + fn u64_min_max() { + // check values roundtrip through parquet + let src = [ + u64::MIN, + 1, + (i64::MAX as u64) - 1, + i64::MAX as u64, + (i64::MAX as u64) + 1, + u64::MAX - 1, + u64::MAX, + ]; + let values = Arc::new(UInt64Array::from_iter_values(src.iter().cloned())); + let files = one_column_roundtrip(values, false); + + for file in files { + // check statistics are valid + let reader = SerializedFileReader::new(file).unwrap(); + let metadata = reader.metadata(); + + let mut row_offset = 0; + for row_group in metadata.row_groups() { + assert_eq!(row_group.num_columns(), 1); + let column = row_group.column(0); + + let num_values = column.num_values() as usize; + let src_slice = &src[row_offset..row_offset + num_values]; + row_offset += column.num_values() as usize; + + let stats = column.statistics().unwrap(); + if let Statistics::Int64(stats) = stats { + assert_eq!( + *stats.min_opt().unwrap() as u64, + *src_slice.iter().min().unwrap() + ); + assert_eq!( + *stats.max_opt().unwrap() as u64, + *src_slice.iter().max().unwrap() + ); + } else { + panic!("Statistics::Int64 missing") + } + } + } + } + + #[test] + fn statistics_null_counts_only_nulls() { + // check that null-count statistics for "only NULL"-columns are correct + let values = Arc::new(UInt64Array::from(vec![None, None])); + let files = one_column_roundtrip(values, true); + + for file in files { + // check statistics are valid + let reader = SerializedFileReader::new(file).unwrap(); + let metadata = reader.metadata(); + assert_eq!(metadata.num_row_groups(), 1); + let row_group = metadata.row_group(0); + assert_eq!(row_group.num_columns(), 1); + let column = row_group.column(0); + let stats = column.statistics().unwrap(); + assert_eq!(stats.null_count_opt(), Some(2)); + } + } + + #[test] + fn test_list_of_struct_roundtrip() { + // define schema + let int_field = Field::new("a", DataType::Int32, true); + let int_field2 = Field::new("b", DataType::Int32, true); + + let int_builder = Int32Builder::with_capacity(10); + let int_builder2 = Int32Builder::with_capacity(10); + + let struct_builder = StructBuilder::new( + vec![int_field, int_field2], + vec![Box::new(int_builder), Box::new(int_builder2)], + ); + let mut list_builder = ListBuilder::new(struct_builder); + + // Construct the following array + // [{a: 1, b: 2}], [], null, [null, null], [{a: null, b: 3}], [{a: 2, b: null}] + + // [{a: 1, b: 2}] + let values = list_builder.values(); + values + .field_builder::(0) + .unwrap() + .append_value(1); + values + .field_builder::(1) + .unwrap() + .append_value(2); + values.append(true); + list_builder.append(true); + + // [] + list_builder.append(true); + + // null + list_builder.append(false); + + // [null, null] + let values = list_builder.values(); + values + .field_builder::(0) + .unwrap() + .append_null(); + values + .field_builder::(1) + .unwrap() + .append_null(); + values.append(false); + values + .field_builder::(0) + .unwrap() + .append_null(); + values + .field_builder::(1) + .unwrap() + .append_null(); + values.append(false); + list_builder.append(true); + + // [{a: null, b: 3}] + let values = list_builder.values(); + values + .field_builder::(0) + .unwrap() + .append_null(); + values + .field_builder::(1) + .unwrap() + .append_value(3); + values.append(true); + list_builder.append(true); + + // [{a: 2, b: null}] + let values = list_builder.values(); + values + .field_builder::(0) + .unwrap() + .append_value(2); + values + .field_builder::(1) + .unwrap() + .append_null(); + values.append(true); + list_builder.append(true); + + let array = Arc::new(list_builder.finish()); + + one_column_roundtrip(array, true); + } + + fn row_group_sizes(metadata: &ParquetMetaData) -> Vec { + metadata.row_groups().iter().map(|x| x.num_rows()).collect() + } + + #[test] + fn test_aggregates_records() { + let arrays = [ + Int32Array::from((0..100).collect::>()), + Int32Array::from((0..50).collect::>()), + Int32Array::from((200..500).collect::>()), + ]; + + let schema = Arc::new(Schema::new(vec![Field::new( + "int", + ArrowDataType::Int32, + false, + )])); + + let file = tempfile::tempfile().unwrap(); + + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(200)) + .build(); + + let mut writer = + ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap(); + + for array in arrays { + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap(); + writer.write(&batch).unwrap(); + } + + writer.close().unwrap(); + + let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + assert_eq!(&row_group_sizes(builder.metadata()), &[200, 200, 50]); + + let batches = builder + .with_batch_size(100) + .build() + .unwrap() + .collect::>>() + .unwrap(); + + assert_eq!(batches.len(), 5); + assert!(batches.iter().all(|x| x.num_columns() == 1)); + + let batch_sizes: Vec<_> = batches.iter().map(|x| x.num_rows()).collect(); + + assert_eq!(&batch_sizes, &[100, 100, 100, 100, 50]); + + let values: Vec<_> = batches + .iter() + .flat_map(|x| { + x.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .cloned() + }) + .collect(); + + let expected_values: Vec<_> = [0..100, 0..50, 200..500].into_iter().flatten().collect(); + assert_eq!(&values, &expected_values) + } + + #[test] + fn complex_aggregate() { + // Tests aggregating nested data + let field_a = Arc::new(Field::new("leaf_a", DataType::Int32, false)); + let field_b = Arc::new(Field::new("leaf_b", DataType::Int32, true)); + let struct_a = Arc::new(Field::new( + "struct_a", + DataType::Struct(vec![field_a.clone(), field_b.clone()].into()), + true, + )); + + let list_a = Arc::new(Field::new("list", DataType::List(struct_a), true)); + let struct_b = Arc::new(Field::new( + "struct_b", + DataType::Struct(vec![list_a.clone()].into()), + false, + )); + + let schema = Arc::new(Schema::new(vec![struct_b])); + + // create nested data + let field_a_array = Int32Array::from(vec![1, 2, 3, 4, 5, 6]); + let field_b_array = + Int32Array::from_iter(vec![Some(1), None, Some(2), None, None, Some(6)]); + + let struct_a_array = StructArray::from(vec![ + (field_a.clone(), Arc::new(field_a_array) as ArrayRef), + (field_b.clone(), Arc::new(field_b_array) as ArrayRef), + ]); + + let list_data = ArrayDataBuilder::new(list_a.data_type().clone()) + .len(5) + .add_buffer(Buffer::from_iter(vec![ + 0_i32, 1_i32, 1_i32, 3_i32, 3_i32, 5_i32, + ])) + .null_bit_buffer(Some(Buffer::from_iter(vec![ + true, false, true, false, true, + ]))) + .child_data(vec![struct_a_array.into_data()]) + .build() + .unwrap(); + + let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef; + let struct_b_array = StructArray::from(vec![(list_a.clone(), list_a_array)]); + + let batch1 = + RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)]) + .unwrap(); + + let field_a_array = Int32Array::from(vec![6, 7, 8, 9, 10]); + let field_b_array = Int32Array::from_iter(vec![None, None, None, Some(1), None]); + + let struct_a_array = StructArray::from(vec![ + (field_a, Arc::new(field_a_array) as ArrayRef), + (field_b, Arc::new(field_b_array) as ArrayRef), + ]); + + let list_data = ArrayDataBuilder::new(list_a.data_type().clone()) + .len(2) + .add_buffer(Buffer::from_iter(vec![0_i32, 4_i32, 5_i32])) + .child_data(vec![struct_a_array.into_data()]) + .build() + .unwrap(); + + let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef; + let struct_b_array = StructArray::from(vec![(list_a, list_a_array)]); + + let batch2 = + RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)]) + .unwrap(); + + let batches = &[batch1, batch2]; + + // Verify data is as expected + + let expected = r#" + +-------------------------------------------------------------------------------------------------------+ + | struct_b | + +-------------------------------------------------------------------------------------------------------+ + | {list: [{leaf_a: 1, leaf_b: 1}]} | + | {list: } | + | {list: [{leaf_a: 2, leaf_b: }, {leaf_a: 3, leaf_b: 2}]} | + | {list: } | + | {list: [{leaf_a: 4, leaf_b: }, {leaf_a: 5, leaf_b: }]} | + | {list: [{leaf_a: 6, leaf_b: }, {leaf_a: 7, leaf_b: }, {leaf_a: 8, leaf_b: }, {leaf_a: 9, leaf_b: 1}]} | + | {list: [{leaf_a: 10, leaf_b: }]} | + +-------------------------------------------------------------------------------------------------------+ + "#.trim().split('\n').map(|x| x.trim()).collect::>().join("\n"); + + let actual = pretty_format_batches(batches).unwrap().to_string(); + assert_eq!(actual, expected); + + // Write data + let file = tempfile::tempfile().unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(6)) + .build(); + + let mut writer = + ArrowWriter::try_new(file.try_clone().unwrap(), schema, Some(props)).unwrap(); + + for batch in batches { + writer.write(batch).unwrap(); + } + writer.close().unwrap(); + + // Read Data + // Should have written entire first batch and first row of second to the first row group + // leaving a single row in the second row group + + let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + assert_eq!(&row_group_sizes(builder.metadata()), &[6, 1]); + + let batches = builder + .with_batch_size(2) + .build() + .unwrap() + .collect::>>() + .unwrap(); + + assert_eq!(batches.len(), 4); + let batch_counts: Vec<_> = batches.iter().map(|x| x.num_rows()).collect(); + assert_eq!(&batch_counts, &[2, 2, 2, 1]); + + let actual = pretty_format_batches(&batches).unwrap().to_string(); + assert_eq!(actual, expected); + } + + #[test] + fn test_arrow_writer_metadata() { + let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]); + let file_schema = batch_schema.clone().with_metadata( + vec![("foo".to_string(), "bar".to_string())] + .into_iter() + .collect(), + ); + + let batch = RecordBatch::try_new( + Arc::new(batch_schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _], + ) + .unwrap(); + + let mut buf = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buf, Arc::new(file_schema), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } + + #[test] + fn test_arrow_writer_nullable() { + let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]); + let file_schema = Schema::new(vec![Field::new("int32", DataType::Int32, true)]); + let file_schema = Arc::new(file_schema); + + let batch = RecordBatch::try_new( + Arc::new(batch_schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _], + ) + .unwrap(); + + let mut buf = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let mut read = ParquetRecordBatchReader::try_new(Bytes::from(buf), 1024).unwrap(); + let back = read.next().unwrap().unwrap(); + assert_eq!(back.schema(), file_schema); + assert_ne!(back.schema(), batch.schema()); + assert_eq!(back.column(0).as_ref(), batch.column(0).as_ref()); + } + + #[test] + fn in_progress_accounting() { + // define schema + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + + // create some data + let a = Int32Array::from(vec![1, 2, 3, 4, 5]); + + // build a record batch + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap(); + + let mut writer = ArrowWriter::try_new(vec![], batch.schema(), None).unwrap(); + + // starts empty + assert_eq!(writer.in_progress_size(), 0); + assert_eq!(writer.in_progress_rows(), 0); + assert_eq!(writer.memory_size(), 0); + assert_eq!(writer.bytes_written(), 4); // Initial header + writer.write(&batch).unwrap(); + + // updated on write + let initial_size = writer.in_progress_size(); + assert!(initial_size > 0); + assert_eq!(writer.in_progress_rows(), 5); + let initial_memory = writer.memory_size(); + assert!(initial_memory > 0); + // memory estimate is larger than estimated encoded size + assert!( + initial_size <= initial_memory, + "{initial_size} <= {initial_memory}" + ); + + // updated on second write + writer.write(&batch).unwrap(); + assert!(writer.in_progress_size() > initial_size); + assert_eq!(writer.in_progress_rows(), 10); + assert!(writer.memory_size() > initial_memory); + assert!( + writer.in_progress_size() <= writer.memory_size(), + "in_progress_size {} <= memory_size {}", + writer.in_progress_size(), + writer.memory_size() + ); + + // in progress tracking is cleared, but the overall data written is updated + let pre_flush_bytes_written = writer.bytes_written(); + writer.flush().unwrap(); + assert_eq!(writer.in_progress_size(), 0); + assert_eq!(writer.memory_size(), 0); + assert!(writer.bytes_written() > pre_flush_bytes_written); + + writer.close().unwrap(); + } + + #[test] + fn test_writer_all_null() { + let a = Int32Array::from(vec![1, 2, 3, 4, 5]); + let b = Int32Array::new(vec![0; 5].into(), Some(NullBuffer::new_null(5))); + let batch = RecordBatch::try_from_iter(vec![ + ("a", Arc::new(a) as ArrayRef), + ("b", Arc::new(b) as ArrayRef), + ]) + .unwrap(); + + let mut buf = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let bytes = Bytes::from(buf); + let options = ReadOptionsBuilder::new().with_page_index().build(); + let reader = SerializedFileReader::new_with_options(bytes, options).unwrap(); + let index = reader.metadata().offset_index().unwrap(); + + assert_eq!(index.len(), 1); + assert_eq!(index[0].len(), 2); // 2 columns + assert_eq!(index[0][0].page_locations().len(), 1); // 1 page + assert_eq!(index[0][1].page_locations().len(), 1); // 1 page + } + + #[test] + fn test_disabled_statistics_with_page() { + let file_schema = Schema::new(vec![ + Field::new("a", DataType::Utf8, true), + Field::new("b", DataType::Utf8, true), + ]); + let file_schema = Arc::new(file_schema); + + let batch = RecordBatch::try_new( + file_schema.clone(), + vec![ + Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _, + Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _, + ], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::None) + .set_column_statistics_enabled("a".into(), EnabledStatistics::Page) + .build(); + + let mut buf = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + + let metadata = writer.close().unwrap(); + assert_eq!(metadata.num_row_groups(), 1); + let row_group = metadata.row_group(0); + assert_eq!(row_group.num_columns(), 2); + // Column "a" has both offset and column index, as requested + assert!(row_group.column(0).offset_index_offset().is_some()); + assert!(row_group.column(0).column_index_offset().is_some()); + // Column "b" should only have offset index + assert!(row_group.column(1).offset_index_offset().is_some()); + assert!(row_group.column(1).column_index_offset().is_none()); + + let options = ReadOptionsBuilder::new().with_page_index().build(); + let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap(); + + let row_group = reader.get_row_group(0).unwrap(); + let a_col = row_group.metadata().column(0); + let b_col = row_group.metadata().column(1); + + // Column chunk of column "a" should have chunk level statistics + if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() { + let min = byte_array_stats.min_opt().unwrap(); + let max = byte_array_stats.max_opt().unwrap(); + + assert_eq!(min.as_bytes(), b"a"); + assert_eq!(max.as_bytes(), b"d"); + } else { + panic!("expecting Statistics::ByteArray"); + } + + // The column chunk for column "b" shouldn't have statistics + assert!(b_col.statistics().is_none()); + + let offset_index = reader.metadata().offset_index().unwrap(); + assert_eq!(offset_index.len(), 1); // 1 row group + assert_eq!(offset_index[0].len(), 2); // 2 columns + + let column_index = reader.metadata().column_index().unwrap(); + assert_eq!(column_index.len(), 1); // 1 row group + assert_eq!(column_index[0].len(), 2); // 2 columns + + let a_idx = &column_index[0][0]; + assert!( + matches!(a_idx, ColumnIndexMetaData::BYTE_ARRAY(_)), + "{a_idx:?}" + ); + let b_idx = &column_index[0][1]; + assert!(matches!(b_idx, ColumnIndexMetaData::NONE), "{b_idx:?}"); + } + + #[test] + fn test_disabled_statistics_with_chunk() { + let file_schema = Schema::new(vec![ + Field::new("a", DataType::Utf8, true), + Field::new("b", DataType::Utf8, true), + ]); + let file_schema = Arc::new(file_schema); + + let batch = RecordBatch::try_new( + file_schema.clone(), + vec![ + Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _, + Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _, + ], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::None) + .set_column_statistics_enabled("a".into(), EnabledStatistics::Chunk) + .build(); + + let mut buf = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + + let metadata = writer.close().unwrap(); + assert_eq!(metadata.num_row_groups(), 1); + let row_group = metadata.row_group(0); + assert_eq!(row_group.num_columns(), 2); + // Column "a" should only have offset index + assert!(row_group.column(0).offset_index_offset().is_some()); + assert!(row_group.column(0).column_index_offset().is_none()); + // Column "b" should only have offset index + assert!(row_group.column(1).offset_index_offset().is_some()); + assert!(row_group.column(1).column_index_offset().is_none()); + + let options = ReadOptionsBuilder::new().with_page_index().build(); + let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap(); + + let row_group = reader.get_row_group(0).unwrap(); + let a_col = row_group.metadata().column(0); + let b_col = row_group.metadata().column(1); + + // Column chunk of column "a" should have chunk level statistics + if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() { + let min = byte_array_stats.min_opt().unwrap(); + let max = byte_array_stats.max_opt().unwrap(); + + assert_eq!(min.as_bytes(), b"a"); + assert_eq!(max.as_bytes(), b"d"); + } else { + panic!("expecting Statistics::ByteArray"); + } + + // The column chunk for column "b" shouldn't have statistics + assert!(b_col.statistics().is_none()); + + let column_index = reader.metadata().column_index().unwrap(); + assert_eq!(column_index.len(), 1); // 1 row group + assert_eq!(column_index[0].len(), 2); // 2 columns + + let a_idx = &column_index[0][0]; + assert!(matches!(a_idx, ColumnIndexMetaData::NONE), "{a_idx:?}"); + let b_idx = &column_index[0][1]; + assert!(matches!(b_idx, ColumnIndexMetaData::NONE), "{b_idx:?}"); + } + + #[test] + fn test_arrow_writer_skip_metadata() { + let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]); + let file_schema = Arc::new(batch_schema.clone()); + + let batch = RecordBatch::try_new( + Arc::new(batch_schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _], + ) + .unwrap(); + let skip_options = ArrowWriterOptions::new().with_skip_arrow_metadata(true); + + let mut buf = Vec::with_capacity(1024); + let mut writer = + ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let bytes = Bytes::from(buf); + let reader_builder = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap(); + assert_eq!(file_schema, *reader_builder.schema()); + if let Some(key_value_metadata) = reader_builder + .metadata() + .file_metadata() + .key_value_metadata() + { + assert!( + !key_value_metadata + .iter() + .any(|kv| kv.key.as_str() == ARROW_SCHEMA_META_KEY) + ); + } + } + + #[test] + fn test_arrow_writer_skip_path_in_schema() { + let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]); + let file_schema = Arc::new(batch_schema.clone()); + + let batch = RecordBatch::try_new( + Arc::new(batch_schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _], + ) + .unwrap(); + + // default options should still write path_in_schema + let skip_options = ArrowWriterOptions::new(); + + let mut buf = Vec::with_capacity(1024); + let mut writer = + ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + // override to not write path_in_schema + let skip_options = ArrowWriterOptions::new().with_properties( + WriterProperties::builder() + .set_write_path_in_schema(false) + .build(), + ); + + let mut buf2 = Vec::with_capacity(1024); + let mut writer = + ArrowWriter::try_new_with_options(&mut buf2, file_schema.clone(), skip_options) + .unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + // buf2 should be a bit smaller due to lack of path_in_schema + assert!(buf.len() > buf2.len()); + } + + #[test] + fn mismatched_schemas() { + let batch_schema = Schema::new(vec![Field::new("count", DataType::Int32, false)]); + let file_schema = Arc::new(Schema::new(vec![Field::new( + "temperature", + DataType::Float64, + false, + )])); + + let batch = RecordBatch::try_new( + Arc::new(batch_schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _], + ) + .unwrap(); + + let mut buf = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap(); + + let err = writer.write(&batch).unwrap_err().to_string(); + assert_eq!( + err, + "Arrow: Incompatible type. Field 'temperature' has type Float64, array has type Int32" + ); + } + + #[test] + // https://github.com/apache/arrow-rs/issues/6988 + fn test_roundtrip_empty_schema() { + // create empty record batch with empty schema + let empty_batch = RecordBatch::try_new_with_options( + Arc::new(Schema::empty()), + vec![], + &RecordBatchOptions::default().with_row_count(Some(0)), + ) + .unwrap(); + + // write to parquet + let mut parquet_bytes: Vec = Vec::new(); + let mut writer = + ArrowWriter::try_new(&mut parquet_bytes, empty_batch.schema(), None).unwrap(); + writer.write(&empty_batch).unwrap(); + writer.close().unwrap(); + + // read from parquet + let bytes = Bytes::from(parquet_bytes); + let reader = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap(); + assert_eq!(reader.schema(), &empty_batch.schema()); + let batches: Vec<_> = reader + .build() + .unwrap() + .collect::>>() + .unwrap(); + assert_eq!(batches.len(), 0); + } + + #[test] + fn test_page_stats_not_written_by_default() { + let string_field = Field::new("a", DataType::Utf8, false); + let schema = Schema::new(vec![string_field]); + let raw_string_values = vec!["Blart Versenwald III"]; + let string_values = StringArray::from(raw_string_values.clone()); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap(); + + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_compression(crate::basic::Compression::UNCOMPRESSED) + .build(); + + let file = roundtrip_opts(&batch, props); + + // read file and decode page headers + // Note: use the thrift API as there is no Rust API to access the statistics in the page headers + + // decode first page header + let first_page = &file[4..]; + let mut prot = ThriftSliceInputProtocol::new(first_page); + let hdr = PageHeader::read_thrift(&mut prot).unwrap(); + let stats = hdr.data_page_header.unwrap().statistics; + + assert!(stats.is_none()); + } + + #[test] + fn test_page_stats_when_enabled() { + let string_field = Field::new("a", DataType::Utf8, false); + let schema = Schema::new(vec![string_field]); + let raw_string_values = vec!["Blart Versenwald III", "Andrew Lamb"]; + let string_values = StringArray::from(raw_string_values.clone()); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap(); + + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_write_page_header_statistics(true) + .set_compression(crate::basic::Compression::UNCOMPRESSED) + .build(); + + let file = roundtrip_opts(&batch, props); + + // read file and decode page headers + // Note: use the thrift API as there is no Rust API to access the statistics in the page headers + + // decode first page header + let first_page = &file[4..]; + let mut prot = ThriftSliceInputProtocol::new(first_page); + let hdr = PageHeader::read_thrift(&mut prot).unwrap(); + let stats = hdr.data_page_header.unwrap().statistics; + + let stats = stats.unwrap(); + // check that min/max were actually written to the page + assert!(stats.is_max_value_exact.unwrap()); + assert!(stats.is_min_value_exact.unwrap()); + assert_eq!(stats.max_value.unwrap(), "Blart Versenwald III".as_bytes()); + assert_eq!(stats.min_value.unwrap(), "Andrew Lamb".as_bytes()); + } + + #[test] + fn test_page_stats_truncation() { + let string_field = Field::new("a", DataType::Utf8, false); + let binary_field = Field::new("b", DataType::Binary, false); + let schema = Schema::new(vec![string_field, binary_field]); + + let raw_string_values = vec!["Blart Versenwald III"]; + let raw_binary_values = [b"Blart Versenwald III".to_vec()]; + let raw_binary_value_refs = raw_binary_values + .iter() + .map(|x| x.as_slice()) + .collect::>(); + + let string_values = StringArray::from(raw_string_values.clone()); + let binary_values = BinaryArray::from(raw_binary_value_refs); + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![Arc::new(string_values), Arc::new(binary_values)], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_statistics_truncate_length(Some(2)) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_write_page_header_statistics(true) + .set_compression(crate::basic::Compression::UNCOMPRESSED) + .build(); + + let file = roundtrip_opts(&batch, props); + + // read file and decode page headers + // Note: use the thrift API as there is no Rust API to access the statistics in the page headers + + // decode first page header + let first_page = &file[4..]; + let mut prot = ThriftSliceInputProtocol::new(first_page); + let hdr = PageHeader::read_thrift(&mut prot).unwrap(); + let stats = hdr.data_page_header.unwrap().statistics; + assert!(stats.is_some()); + let stats = stats.unwrap(); + // check that min/max were properly truncated + assert!(!stats.is_max_value_exact.unwrap()); + assert!(!stats.is_min_value_exact.unwrap()); + assert_eq!(stats.max_value.unwrap(), "Bm".as_bytes()); + assert_eq!(stats.min_value.unwrap(), "Bl".as_bytes()); + + // check second page now + let second_page = &prot.as_slice()[hdr.compressed_page_size as usize..]; + let mut prot = ThriftSliceInputProtocol::new(second_page); + let hdr = PageHeader::read_thrift(&mut prot).unwrap(); + let stats = hdr.data_page_header.unwrap().statistics; + assert!(stats.is_some()); + let stats = stats.unwrap(); + // check that min/max were properly truncated + assert!(!stats.is_max_value_exact.unwrap()); + assert!(!stats.is_min_value_exact.unwrap()); + assert_eq!(stats.max_value.unwrap(), "Bm".as_bytes()); + assert_eq!(stats.min_value.unwrap(), "Bl".as_bytes()); + } + + #[test] + fn test_page_encoding_statistics_roundtrip() { + let batch_schema = Schema::new(vec![Field::new( + "int32", + arrow_schema::DataType::Int32, + false, + )]); + + let batch = RecordBatch::try_new( + Arc::new(batch_schema.clone()), + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _], + ) + .unwrap(); + + let mut file: File = tempfile::tempfile().unwrap(); + let mut writer = ArrowWriter::try_new(&mut file, Arc::new(batch_schema), None).unwrap(); + writer.write(&batch).unwrap(); + let file_metadata = writer.close().unwrap(); + + assert_eq!(file_metadata.num_row_groups(), 1); + assert_eq!(file_metadata.row_group(0).num_columns(), 1); + assert!( + file_metadata + .row_group(0) + .column(0) + .page_encoding_stats() + .is_some() + ); + let chunk_page_stats = file_metadata + .row_group(0) + .column(0) + .page_encoding_stats() + .unwrap(); + + // check that the read metadata is also correct + let options = ReadOptionsBuilder::new() + .with_page_index() + .with_encoding_stats_as_mask(false) + .build(); + let reader = SerializedFileReader::new_with_options(file, options).unwrap(); + + let rowgroup = reader.get_row_group(0).expect("row group missing"); + assert_eq!(rowgroup.num_columns(), 1); + let column = rowgroup.metadata().column(0); + assert!(column.page_encoding_stats().is_some()); + let file_page_stats = column.page_encoding_stats().unwrap(); + assert_eq!(chunk_page_stats, file_page_stats); + } + + #[test] + fn test_different_dict_page_size_limit() { + let array = Arc::new(Int64Array::from_iter(0..1024 * 1024)); + let schema = Arc::new(Schema::new(vec![ + Field::new("col0", arrow_schema::DataType::Int64, false), + Field::new("col1", arrow_schema::DataType::Int64, false), + ])); + let batch = + arrow_array::RecordBatch::try_new(schema.clone(), vec![array.clone(), array]).unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_page_size_limit(1024 * 1024) + .set_column_dictionary_page_size_limit(ColumnPath::from("col1"), 1024 * 1024 * 4) + .build(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + let data = Bytes::from(writer.into_inner().unwrap()); + + let mut metadata = ParquetMetaDataReader::new(); + metadata.try_parse(&data).unwrap(); + let metadata = metadata.finish().unwrap(); + let col0_meta = metadata.row_group(0).column(0); + let col1_meta = metadata.row_group(0).column(1); + + let get_dict_page_size = move |meta: &ColumnChunkMetaData| { + let mut reader = + SerializedPageReader::new(Arc::new(data.clone()), meta, 0, None).unwrap(); + let page = reader.get_next_page().unwrap().unwrap(); + match page { + Page::DictionaryPage { buf, .. } => buf.len(), + _ => panic!("expected DictionaryPage"), + } + }; + + assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024); + assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4); + } + + #[test] + fn test_arrow_writer_granular_mode_roundtrip() { + // Granular mode subdivides chunks and writes more pages than the + // default batched path. Make sure the data we write back is + // bit-identical to what went in — page-count assertions elsewhere + // only prove pages were cut, not that the encoded data is correct. + // + // Mix value sizes so that the cumulative-byte-budget cutoff + // lands mid-chunk, exercising both batched and granular paths + // within the same `write_batch_internal` call. + let small = "tiny".to_string(); + let big = "x".repeat(64 * 1024); + let strings: Vec = (0..256) + .map(|i| { + if i % 16 == 0 { + big.clone() + } else { + small.clone() + } + }) + .collect(); + + let schema = Arc::new(Schema::new(vec![Field::new( + "col", + ArrowDataType::Utf8, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StringArray::from(strings.clone())) as _], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(16 * 1024) + .build(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + let data = Bytes::from(writer.into_inner().unwrap()); + + let mut reader = ParquetRecordBatchReader::try_new(data, 1024).unwrap(); + let read = reader.next().unwrap().unwrap(); + assert!(reader.next().is_none(), "expected one batch"); + let col = read + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(col.len(), strings.len()); + for (i, expected) in strings.iter().enumerate() { + assert_eq!( + col.value(i), + expected.as_str(), + "value mismatch at index {i}" + ); + } + } + + #[test] + fn test_arrow_writer_all_null_string_column() { + // The `LevelDataRef::value_count` Uniform branch with + // `value != max_def` (entirely-null chunk) must return 0 so the + // sub-batch sizer short-circuits to batch mode without trying + // to estimate byte budgets for non-existent values. + let num_rows = 1024; + let schema = Arc::new(Schema::new(vec![Field::new( + "col", + ArrowDataType::Utf8, + true, + )])); + let nulls: Vec> = vec![None; num_rows]; + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StringArray::from(nulls)) as _], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(16 * 1024) + .build(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + let data = Bytes::from(writer.into_inner().unwrap()); + + // Re-parse the file: row group has one column, every row is + // null, all data pages report `num_rows / page_count` rows. + let mut metadata = ParquetMetaDataReader::new(); + metadata.try_parse(&data).unwrap(); + let metadata = metadata.finish().unwrap(); + let row_group = metadata.row_group(0); + let col_meta = row_group.column(0); + assert_eq!(row_group.num_rows() as usize, num_rows); + // Statistics record `null_count = num_rows` — proves every value + // was written as null. + if let Some(stats) = col_meta.statistics() { + assert_eq!( + stats.null_count_opt().unwrap_or(0) as usize, + num_rows, + "expected all-null column to report null_count = num_rows" + ); + } + + let mut reader = + SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap(); + let mut total_values = 0u32; + while let Some(page) = reader.get_next_page().unwrap() { + if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) { + total_values += page.num_values(); + } + } + assert_eq!( + total_values as usize, num_rows, + "expected every level position to be represented in some page" + ); + } + + struct WriteBatchesShape { + num_batches: usize, + rows_per_batch: usize, + row_size: usize, + } + + /// Helper function to write batches with the provided `WriteBatchesShape` into an `ArrowWriter` + fn write_batches( + WriteBatchesShape { + num_batches, + rows_per_batch, + row_size, + }: WriteBatchesShape, + props: WriterProperties, + ) -> ParquetRecordBatchReaderBuilder { + let schema = Arc::new(Schema::new(vec![Field::new( + "str", + ArrowDataType::Utf8, + false, + )])); + let file = tempfile::tempfile().unwrap(); + let mut writer = + ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap(); + + for batch_idx in 0..num_batches { + let strings: Vec = (0..rows_per_batch) + .map(|i| format!("{:0>width$}", batch_idx * 10 + i, width = row_size)) + .collect(); + let array = StringArray::from(strings); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap(); + writer.write(&batch).unwrap(); + } + writer.close().unwrap(); + ParquetRecordBatchReaderBuilder::try_new(file).unwrap() + } + + #[test] + // When both limits are None, all data should go into a single row group + fn test_row_group_limit_none_writes_single_row_group() { + let props = WriterProperties::builder() + .set_max_row_group_row_count(None) + .set_max_row_group_bytes(None) + .build(); + + let builder = write_batches( + WriteBatchesShape { + num_batches: 1, + rows_per_batch: 1000, + row_size: 4, + }, + props, + ); + + assert_eq!( + &row_group_sizes(builder.metadata()), + &[1000], + "With no limits, all rows should be in a single row group" + ); + } + + #[test] + // When only max_row_group_size is set, respect the row limit + fn test_row_group_limit_rows_only() { + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(300)) + .set_max_row_group_bytes(None) + .build(); + + let builder = write_batches( + WriteBatchesShape { + num_batches: 1, + rows_per_batch: 1000, + row_size: 4, + }, + props, + ); + + assert_eq!( + &row_group_sizes(builder.metadata()), + &[300, 300, 300, 100], + "Row groups should be split by row count" + ); + } + + #[test] + // When only max_row_group_bytes is set, respect the byte limit + fn test_row_group_limit_bytes_only() { + let props = WriterProperties::builder() + .set_max_row_group_row_count(None) + // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each) + .set_max_row_group_bytes(Some(3500)) + .build(); + + let builder = write_batches( + WriteBatchesShape { + num_batches: 10, + rows_per_batch: 10, + row_size: 100, + }, + props, + ); + + let sizes = row_group_sizes(builder.metadata()); + + assert!( + sizes.len() > 1, + "Should have multiple row groups due to byte limit, got {sizes:?}", + ); + + let total_rows: i64 = sizes.iter().sum(); + assert_eq!(total_rows, 100, "Total rows should be preserved"); + } + + #[test] + // If an in-progress row group is already oversized, it should be flushed before writing more. + fn test_row_group_limit_bytes_flushes_when_current_group_already_too_large() { + let schema = Arc::new(Schema::new(vec![Field::new( + "str", + ArrowDataType::Utf8, + false, + )])); + let file = tempfile::tempfile().unwrap(); + + // Start with no byte limit so we can intentionally build an oversized in-progress row group. + let props = WriterProperties::builder() + .set_max_row_group_row_count(None) + .set_max_row_group_bytes(None) + .build(); + let mut writer = + ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap(); + + let first_array = StringArray::from( + (0..10) + .map(|i| format!("{:0>100}", i)) + .collect::>(), + ); + let first_batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(first_array)]).unwrap(); + writer.write(&first_batch).unwrap(); + assert_eq!(writer.in_progress_rows(), 10); + + // Tighten the limit below the current in-progress bytes to exercise: + // `if current_bytes >= max_bytes { self.flush()?; ... }` + writer.max_row_group_bytes = Some(1); + + let second_array = StringArray::from(vec!["x".to_string()]); + let second_batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(second_array)]).unwrap(); + writer.write(&second_batch).unwrap(); + writer.close().unwrap(); + let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + + assert_eq!( + &row_group_sizes(builder.metadata()), + &[10, 1], + "The second write should flush an oversized in-progress row group first", + ); + } + + #[test] + // When both limits are set, the row limit triggers first + fn test_row_group_limit_both_row_wins_single_batch() { + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(200)) // Will trigger at 200 rows + .set_max_row_group_bytes(Some(1024 * 1024)) // 1MB - won't trigger for small int data + .build(); + + let builder = write_batches( + WriteBatchesShape { + num_batches: 1, + row_size: 4, + rows_per_batch: 1000, + }, + props, + ); + + assert_eq!( + &row_group_sizes(builder.metadata()), + &[200, 200, 200, 200, 200], + "Row limit should trigger before byte limit" + ); + } + + #[test] + // When both limits are set, the row limit triggers first + fn test_row_group_limit_both_row_wins_multiple_batches() { + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(5)) // Will trigger every 5 rows + .set_max_row_group_bytes(Some(9999)) // Won't trigger + .build(); + + let builder = write_batches( + WriteBatchesShape { + num_batches: 10, + rows_per_batch: 10, + row_size: 100, + }, + props, + ); + + assert_eq!( + &row_group_sizes(builder.metadata()), + &[5; 20], + "Row limit should trigger before byte limit" + ); + } + + #[test] + // When both limits are set, the byte limit triggers first + fn test_row_group_limit_both_bytes_wins() { + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(1000)) // Won't trigger for 100 rows + .set_max_row_group_bytes(Some(3500)) // Will trigger at ~30-35 rows + .build(); + + let builder = write_batches( + WriteBatchesShape { + num_batches: 10, + rows_per_batch: 10, + row_size: 100, + }, + props, + ); + + let sizes = row_group_sizes(builder.metadata()); + + assert!( + sizes.len() > 1, + "Byte limit should trigger before row limit, got {sizes:?}", + ); + + assert!( + sizes.iter().all(|&s| s < 1000), + "No row group should hit the row limit" + ); + + let total_rows: i64 = sizes.iter().sum(); + assert_eq!(total_rows, 100, "Total rows should be preserved"); + } + + #[test] + fn arrow_column_chunk_close_mut_drops_column_index() { + use crate::arrow::ArrowSchemaConverter; + use crate::file::writer::SerializedFileWriter; + + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])); + let props = Arc::new( + WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .build(), + ); + let parquet_schema = ArrowSchemaConverter::new() + .with_coerce_types(props.coerce_types()) + .convert(&schema) + .unwrap(); + + let mut buf = Vec::with_capacity(1024); + let mut writer = + SerializedFileWriter::new(&mut buf, parquet_schema.root_schema_ptr(), props.clone()) + .unwrap(); + + let factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema)); + let mut col_writers = factory.create_column_writers(0).unwrap(); + let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64)); + for leaves in compute_leaves(schema.field(0), &arr).unwrap() { + col_writers[0].write(&leaves).unwrap(); + } + let mut chunk = col_writers.pop().unwrap().close().unwrap(); + + // Immutable accessor exposes the close result produced at close time. + assert!( + chunk.close().column_index.is_some(), + "EnabledStatistics::Page should produce a column_index" + ); + + // Mutable accessor lets callers drop the page-level index before append. + chunk.close_mut().column_index = None; + assert!(chunk.close().column_index.is_none()); + + let mut rg = writer.next_row_group().unwrap(); + chunk.append_to_row_group(&mut rg).unwrap(); + rg.close().unwrap(); + let file_meta = writer.close().unwrap(); + + // After dropping column_index, the resulting file records no column + // index offset/length for this chunk. + let cc = file_meta.row_group(0).column(0); + assert!(cc.column_index_range().is_none()); + } + + /// Writes a single-column RecordBatch to an in-memory Parquet buffer. + fn write_column_to_bytes(array: ArrayRef) -> Bytes { + let schema = Arc::new(Schema::new(vec![Field::new( + "col", + array.data_type().clone(), + true, + )])); + let buf = get_bytes_after_close( + schema.clone(), + &RecordBatch::try_new(schema, vec![array]).unwrap(), + ); + Bytes::from(buf) + } + + /// Reads column 0 from a single-row-group Parquet buffer, projecting it with the given schema. + /// Passing a flat schema when the buffer was written from a REE array lets callers decode + /// the physical values without the run-end encoding wrapper. + fn read_column_with_schema(bytes: Bytes, schema: SchemaRef) -> ArrayRef { + let opts = crate::arrow::arrow_reader::ArrowReaderOptions::new().with_schema(schema); + ParquetRecordBatchReaderBuilder::try_new_with_options(bytes, opts) + .unwrap() + .build() + .unwrap() + .next() + .unwrap() + .unwrap() + .column(0) + .clone() + } + + fn ree_write_read_roundtrip(ree: ArrayRef, flat: ArrayRef) { + let flat_schema = Arc::new(Schema::new(vec![Field::new( + "col", + flat.data_type().clone(), + true, + )])); + let ree_bytes = write_column_to_bytes(ree); + let flat_bytes = write_column_to_bytes(flat.clone()); + assert_eq!( + ree_bytes, flat_bytes, + "REE and flat bytes should be identical" + ); + + let decoded_ree = read_column_with_schema(ree_bytes, flat_schema.clone()); + let decoded_flat = read_column_with_schema(flat_bytes, flat_schema); + + assert_eq!(decoded_ree.as_ref(), flat.as_ref()); + assert_eq!(decoded_ree.as_ref(), decoded_flat.as_ref()); + } + + #[test] + fn ree_string() { + let ree: ArrayRef = Arc::new( + [Some("a"), Some("a"), None, Some("b"), Some("b")] + .into_iter() + .collect::(), + ); + let flat: ArrayRef = Arc::new(StringArray::from(vec![ + Some("a"), + Some("a"), + None, + Some("b"), + Some("b"), + ])); + ree_write_read_roundtrip(ree, flat); + } + + #[test] + fn ree_int32() { + let mut b = PrimitiveRunBuilder::::new(); + for v in [Some(1), Some(1), None, Some(2), Some(2)] { + b.append_option(v); + } + let ree: ArrayRef = Arc::new(b.finish()); + let flat: ArrayRef = Arc::new(Int32Array::from(vec![ + Some(1), + Some(1), + None, + Some(2), + Some(2), + ])); + ree_write_read_roundtrip(ree, flat); + } + + #[test] + fn ree_bool() { + // run_ends [3, 5, 7] → [T,T,T, null,null, F,F] + let ree: ArrayRef = Arc::new( + RunArray::try_new( + &Int32Array::from(vec![3, 5, 7]), + &BooleanArray::from(vec![Some(true), None, Some(false)]), + ) + .unwrap(), + ); + let flat: ArrayRef = Arc::new(BooleanArray::from(vec![ + Some(true), + Some(true), + Some(true), + None, + None, + Some(false), + Some(false), + ])); + ree_write_read_roundtrip(ree, flat); + } + + #[test] + fn ree_fixed_size_binary() { + let mk = |vals: &[Option<&[u8]>]| -> FixedSizeBinaryArray { + let mut b = FixedSizeBinaryBuilder::new(2); + for v in vals { + match v { + Some(x) => b.append_value(x).unwrap(), + None => b.append_null(), + } + } + b.finish() + }; + // run_ends [2, 4, 6] → [aa,aa, null,null, bb,bb] + let ree: ArrayRef = Arc::new( + RunArray::try_new( + &Int32Array::from(vec![2, 4, 6]), + &mk(&[Some(b"aa"), None, Some(b"bb")]), + ) + .unwrap(), + ); + let flat: ArrayRef = Arc::new(mk(&[ + Some(b"aa"), + Some(b"aa"), + None, + None, + Some(b"bb"), + Some(b"bb"), + ])); + ree_write_read_roundtrip(ree, flat); + } + + #[test] + fn ree_single_run() { + let ree: ArrayRef = Arc::new(["x", "x", "x"].into_iter().collect::()); + let flat: ArrayRef = Arc::new(StringArray::from(vec!["x", "x", "x"])); + ree_write_read_roundtrip(ree, flat); + } + + #[test] + fn ree_float32() { + // run_ends [2, 4, 5] → [1.0, 1.0, null, null, 2.5] + let ree: ArrayRef = Arc::new( + RunArray::try_new( + &Int32Array::from(vec![2, 4, 5]), + &Float32Array::from(vec![Some(1.0_f32), None, Some(2.5_f32)]), + ) + .unwrap(), + ); + let flat: ArrayRef = Arc::new(Float32Array::from(vec![ + Some(1.0_f32), + Some(1.0_f32), + None, + None, + Some(2.5_f32), + ])); + ree_write_read_roundtrip(ree, flat); + } + + #[test] + fn ree_sliced() { + // A sliced (non-zero offset) REE array: verify that get_physical_index + // correctly accounts for the logical offset when expanding. + // Full array: run_ends [3, 5, 7] → [a,a,a, b,b, c,c] + // After slice(2, 5) the logical view is [a, b, b, c, c]. + let full: ArrayRef = Arc::new( + RunArray::try_new( + &Int32Array::from(vec![3, 5, 7]), + &StringArray::from(vec!["a", "b", "c"]), + ) + .unwrap(), + ); + let sliced = full.slice(2, 5); + let flat: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "b", "c", "c"])); + ree_write_read_roundtrip(sliced, flat); + } + + #[test] + fn ree_struct_with_ree_child() { + // Struct with a REE string field and a REE int field — confirms + // recursion visits every child and each collapses to the right leaf type. + let run_ends = Int32Array::from(vec![2i32, 3, 5]); + + let col_a: ArrayRef = Arc::new( + RunArray::try_new( + &run_ends, + &StringArray::from(vec![Some("foo"), None, Some("bar")]), + ) + .unwrap(), + ); + let col_b: ArrayRef = Arc::new( + RunArray::try_new(&run_ends, &Int32Array::from(vec![Some(1), None, Some(2)])).unwrap(), + ); + + let struct_array: ArrayRef = Arc::new(StructArray::new( + Fields::from(vec![ + Field::new("a", col_a.data_type().clone(), true), + Field::new("b", col_b.data_type().clone(), true), + ]), + vec![col_a, col_b], + None, + )); + + let schema = Arc::new(Schema::new(vec![Field::new( + "row", + struct_array.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![struct_array]).unwrap(); + + let mut buf = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut buf, schema, None).unwrap(); + writer.write(&batch).unwrap(); + let metadata = writer.close().unwrap(); + + let parquet_schema = metadata.file_metadata().schema_descr(); + assert_eq!(parquet_schema.num_columns(), 2); + assert_eq!( + parquet_schema.column(0).physical_type(), + crate::basic::Type::BYTE_ARRAY + ); + assert_eq!(parquet_schema.column(0).path().string(), "row.a"); + assert_eq!( + parquet_schema.column(1).physical_type(), + crate::basic::Type::INT32 + ); + assert_eq!(parquet_schema.column(1).path().string(), "row.b"); + } +} diff --git a/vendor/parquet/src/arrow/async_reader/metadata.rs b/vendor/parquet/src/arrow/async_reader/metadata.rs new file mode 100644 index 0000000..0ab6a62 --- /dev/null +++ b/vendor/parquet/src/arrow/async_reader/metadata.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 crate::arrow::async_reader::AsyncFileReader; +use crate::errors::Result; +use bytes::Bytes; +use futures::future::BoxFuture; +use std::ops::Range; + +/// A data source that can be used with [`ParquetMetaDataReader`] to load [`ParquetMetaData`] +/// +/// Note that implementation is provided for [`AsyncFileReader`]. +/// +/// # Example `MetadataFetch` for a custom async data source +/// +/// ```rust +/// # use parquet::errors::Result; +/// # use parquet::arrow::async_reader::MetadataFetch; +/// # use bytes::Bytes; +/// # use std::ops::Range; +/// # use std::io::SeekFrom; +/// # use futures::future::BoxFuture; +/// # use futures::FutureExt; +/// # use tokio::io::{AsyncReadExt, AsyncSeekExt}; +/// // Adapter that implements the API for reading bytes from an async source (in +/// // this case a tokio::fs::File) +/// struct TokioFileMetadata { +/// file: tokio::fs::File, +/// } +/// impl MetadataFetch for TokioFileMetadata { +/// fn fetch(&mut self, range: Range) -> BoxFuture<'_, Result> { +/// // return a future that fetches data in range +/// async move { +/// let len = (range.end - range.start).try_into().unwrap(); +/// let mut buf = vec![0; len]; // target buffer +/// // seek to the start of the range and read the data +/// self.file.seek(SeekFrom::Start(range.start)).await?; +/// self.file.read_exact(&mut buf).await?; +/// Ok(Bytes::from(buf)) // convert to Bytes +/// } +/// .boxed() // turn into BoxedFuture, using FutureExt::boxed +/// } +/// } +///``` +/// +/// [`ParquetMetaDataReader`]: crate::file::metadata::reader::ParquetMetaDataReader +/// [`ParquetMetaData`]: crate::file::metadata::ParquetMetaData +pub trait MetadataFetch { + /// Return a future that fetches the specified range of bytes asynchronously + /// + /// Note the returned type is a boxed future, often created by + /// [`FutureExt::boxed`]. See the trait documentation for an example + /// + /// [`FutureExt::boxed`]: futures::FutureExt::boxed + fn fetch(&mut self, range: Range) -> BoxFuture<'_, Result>; +} + +impl MetadataFetch for &mut T { + fn fetch(&mut self, range: Range) -> BoxFuture<'_, Result> { + self.get_bytes(range) + } +} + +/// A data source that can be used with [`ParquetMetaDataReader`] to load [`ParquetMetaData`] via suffix +/// requests, without knowing the file size +/// +/// [`ParquetMetaDataReader`]: crate::file::metadata::reader::ParquetMetaDataReader +/// [`ParquetMetaData`]: crate::file::metadata::ParquetMetaData +pub trait MetadataSuffixFetch: MetadataFetch { + /// Return a future that fetches the last `n` bytes asynchronously + /// + /// Note the returned type is a boxed future, often created by + /// [`FutureExt::boxed`]. See the trait documentation for an example + /// + /// [`FutureExt::boxed`]: futures::FutureExt::boxed + fn fetch_suffix(&mut self, suffix: usize) -> BoxFuture<'_, Result>; +} diff --git a/vendor/parquet/src/arrow/async_reader/mod.rs b/vendor/parquet/src/arrow/async_reader/mod.rs new file mode 100644 index 0000000..1d276df --- /dev/null +++ b/vendor/parquet/src/arrow/async_reader/mod.rs @@ -0,0 +1,2068 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! `async` API for reading Parquet files as [`RecordBatch`]es +//! +//! See the [crate-level documentation](crate) for more details. +//! +//! See example on [`ParquetRecordBatchStreamBuilder::new`] + +use std::fmt::Formatter; +use std::io::SeekFrom; +use std::ops::Range; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use bytes::Bytes; +use futures::future::{BoxFuture, FutureExt}; +use futures::stream::Stream; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt}; + +use arrow_array::RecordBatch; +use arrow_schema::{Schema, SchemaRef}; + +use crate::arrow::arrow_reader::{ + ArrowReaderBuilder, ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReader, +}; + +use crate::basic::{BloomFilterAlgorithm, BloomFilterCompression, BloomFilterHash}; +use crate::bloom_filter::{ + SBBF_HEADER_SIZE_ESTIMATE, Sbbf, chunk_read_bloom_filter_header_and_offset, +}; +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::{ParquetMetaData, ParquetMetaDataReader}; + +mod metadata; +pub use metadata::*; + +mod spawn; +pub use spawn::SpawnedReader; + +#[cfg(feature = "object_store")] +mod store; + +use crate::DecodeResult; +use crate::arrow::push_decoder::{ParquetPushDecoder, ParquetPushDecoderBuilder, PushDecoderInput}; +#[allow(deprecated)] +#[cfg(feature = "object_store")] +pub use store::*; + +/// The asynchronous interface used by [`ParquetRecordBatchStream`] to read parquet files +/// +/// Notes: +/// +/// 1. There is a default implementation for types that implement [`AsyncRead`] +/// and [`AsyncSeek`], for example [`tokio::fs::File`]. +/// +/// 2. Implementations for remote storage, such as the `object_store` crate, +/// can implement this interface directly, typically by pairing a store +/// handle with an object path and delegating [`Self::get_bytes`] and +/// [`Self::get_byte_ranges`] to ranged reads. [`SpawnedReader`] can wrap +/// such a reader to perform its I/O on a dedicated runtime, and +/// [`ParquetMetaDataReader::with_arrow_reader_options`] simplifies +/// implementing [`Self::get_metadata`]. +/// +/// # Example: implementing `AsyncFileReader` for the `object_store` crate +/// +/// ```no_run +/// # use std::ops::Range; +/// # use std::sync::Arc; +/// use bytes::Bytes; +/// use futures::future::BoxFuture; +/// use futures::{FutureExt, TryFutureExt}; +/// use object_store::path::Path; +/// use object_store::{GetOptions, GetRange, ObjectStore, ObjectStoreExt}; +/// use parquet::arrow::arrow_reader::ArrowReaderOptions; +/// use parquet::arrow::async_reader::{AsyncFileReader, MetadataSuffixFetch}; +/// use parquet::errors::{ParquetError, Result}; +/// use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader}; +/// +/// fn to_parquet_err(e: object_store::Error) -> ParquetError { +/// ParquetError::External(Box::new(e)) +/// } +/// +/// #[derive(Clone)] +/// struct ObjectStoreReader { +/// store: Arc, +/// path: Path, +/// } +/// +/// impl AsyncFileReader for ObjectStoreReader { +/// fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { +/// self.store +/// .get_range(&self.path, range) +/// .map_err(to_parquet_err) +/// .boxed() +/// } +/// +/// fn get_byte_ranges(&mut self, ranges: Vec>) -> BoxFuture<'_, Result>> { +/// async move { +/// self.store +/// .get_ranges(&self.path, &ranges) +/// .await +/// .map_err(to_parquet_err) +/// } +/// .boxed() +/// } +/// +/// fn get_metadata<'a>( +/// &'a mut self, +/// options: Option<&'a ArrowReaderOptions>, +/// ) -> BoxFuture<'a, Result>> { +/// async move { +/// let metadata = ParquetMetaDataReader::new() +/// .with_arrow_reader_options(options) +/// .load_via_suffix_and_finish(self) +/// .await?; +/// Ok(Arc::new(metadata)) +/// } +/// .boxed() +/// } +/// } +/// +/// /// Supports fetching the parquet footer without knowing the file size, +/// /// via suffix range requests +/// impl MetadataSuffixFetch for &mut ObjectStoreReader { +/// fn fetch_suffix(&mut self, suffix: usize) -> BoxFuture<'_, Result> { +/// let options = GetOptions { +/// range: Some(GetRange::Suffix(suffix as u64)), +/// ..Default::default() +/// }; +/// async move { +/// let resp = self +/// .store +/// .get_opts(&self.path, options) +/// .await +/// .map_err(to_parquet_err)?; +/// resp.bytes().await.map_err(to_parquet_err) +/// } +/// .boxed() +/// } +/// } +/// ``` +/// +/// [`ParquetMetaDataReader::with_arrow_reader_options`]: crate::file::metadata::ParquetMetaDataReader::with_arrow_reader_options +/// +/// [`tokio::fs::File`]: https://docs.rs/tokio/latest/tokio/fs/struct.File.html +pub trait AsyncFileReader: Send { + /// Retrieve the bytes in `range` + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result>; + + /// Retrieve multiple byte ranges. The default implementation will call `get_bytes` sequentially + fn get_byte_ranges(&mut self, ranges: Vec>) -> BoxFuture<'_, Result>> { + async move { + let mut result = Vec::with_capacity(ranges.len()); + + for range in ranges.into_iter() { + let data = self.get_bytes(range).await?; + result.push(data); + } + + Ok(result) + } + .boxed() + } + + /// Return a future which results in the [`ParquetMetaData`] for this Parquet file. + /// + /// This is an asynchronous operation as it may involve reading the file + /// footer and potentially other metadata from disk or a remote source. + /// + /// Reading data from Parquet requires the metadata to understand the + /// schema, row groups, and location of pages within the file. This metadata + /// is stored primarily in the footer of the Parquet file, and can be read using + /// [`ParquetMetaDataReader`]. + /// + /// However, implementations can significantly speed up reading Parquet by + /// supplying cached metadata or pre-fetched metadata via this API. + /// + /// # Parameters + /// * `options`: Optional [`ArrowReaderOptions`] that may contain decryption + /// and other options that affect how the metadata is read. + fn get_metadata<'a>( + &'a mut self, + options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, Result>>; +} + +/// This allows Box to be used as an AsyncFileReader, +impl AsyncFileReader for Box { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + self.as_mut().get_bytes(range) + } + + fn get_byte_ranges(&mut self, ranges: Vec>) -> BoxFuture<'_, Result>> { + self.as_mut().get_byte_ranges(ranges) + } + + fn get_metadata<'a>( + &'a mut self, + options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, Result>> { + self.as_mut().get_metadata(options) + } +} + +impl MetadataSuffixFetch for T { + fn fetch_suffix(&mut self, suffix: usize) -> BoxFuture<'_, Result> { + async move { + self.seek(SeekFrom::End(-(suffix as i64))).await?; + let mut buf = Vec::with_capacity(suffix); + self.take(suffix as _).read_to_end(&mut buf).await?; + Ok(buf.into()) + } + .boxed() + } +} + +impl AsyncFileReader for T { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + async move { + self.seek(SeekFrom::Start(range.start)).await?; + + let to_read = range.end - range.start; + let mut buffer = Vec::with_capacity(to_read.try_into()?); + let read = self.take(to_read).read_to_end(&mut buffer).await?; + if read as u64 != to_read { + return Err(eof_err!("expected to read {} bytes, got {}", to_read, read)); + } + + Ok(buffer.into()) + } + .boxed() + } + + fn get_metadata<'a>( + &'a mut self, + options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, Result>> { + async move { + let metadata_reader = ParquetMetaDataReader::new().with_arrow_reader_options(options); + let parquet_metadata = metadata_reader.load_via_suffix_and_finish(self).await?; + Ok(Arc::new(parquet_metadata)) + } + .boxed() + } +} + +impl ArrowReaderMetadata { + /// Returns a new [`ArrowReaderMetadata`] for this builder + /// + /// See [`ParquetRecordBatchStreamBuilder::new_with_metadata`] for how this can be used + pub async fn load_async( + input: &mut T, + options: ArrowReaderOptions, + ) -> Result { + let metadata = input.get_metadata(Some(&options)).await?; + Self::try_new(metadata, options) + } +} + +#[doc(hidden)] +/// Newtype (wrapper) used within [`ArrowReaderBuilder`] to distinguish sync readers from async +/// +/// Allows sharing the same builder for different readers while keeping the same +/// ParquetRecordBatchStreamBuilder API +pub struct AsyncReader(T); + +/// A builder for reading parquet files from an `async` source as [`ParquetRecordBatchStream`] +/// +/// This can be used to decode a Parquet file in streaming fashion (without +/// downloading the whole file at once) from a remote source, such as an object store. +/// +/// This builder handles reading the parquet file metadata, allowing consumers +/// to use this information to select what specific columns, row groups, etc. +/// they wish to be read by the resulting stream. +/// +/// See examples on [`ParquetRecordBatchStreamBuilder::new`], including how to +/// issue multiple I/O requests in parallel using multiple streams. +/// +/// # See also: +/// * [`ParquetPushDecoderBuilder`] for lower level control over buffering and +/// decoding. +/// * [`ParquetRecordBatchStream::next_row_group`] for I/O prefetching +/// +/// +/// See [`ArrowReaderBuilder`] for additional member functions +pub type ParquetRecordBatchStreamBuilder = ArrowReaderBuilder>; + +impl ParquetRecordBatchStreamBuilder { + /// Create a new [`ParquetRecordBatchStreamBuilder`] for reading from the + /// specified source. + /// + /// # Examples: + /// * [Basic example reading from an async source](#example) + /// * [Configuring options and reading metadata](#example-configuring-options-and-reading-metadata) + /// * [Reading Row Groups in Parallel](#example-reading-row-groups-in-parallel) + /// + /// # Example + /// ``` + /// # #[tokio::main(flavor="current_thread")] + /// # async fn main() { + /// # + /// # use arrow_array::RecordBatch; + /// # use arrow::util::pretty::pretty_format_batches; + /// # use futures::TryStreamExt; + /// # + /// # use parquet::arrow::{ParquetRecordBatchStreamBuilder, ProjectionMask}; + /// # + /// # fn assert_batches_eq(batches: &[RecordBatch], expected_lines: &[&str]) { + /// # let formatted = pretty_format_batches(batches).unwrap().to_string(); + /// # let actual_lines: Vec<_> = formatted.trim().lines().collect(); + /// # assert_eq!( + /// # &actual_lines, expected_lines, + /// # "\n\nexpected:\n\n{:#?}\nactual:\n\n{:#?}\n\n", + /// # expected_lines, actual_lines + /// # ); + /// # } + /// # + /// # let testdata = arrow::util::test_util::parquet_test_data(); + /// # let path = format!("{}/alltypes_plain.parquet", testdata); + /// // Use tokio::fs::File to read data using an async I/O. This can be replaced with + /// // another async I/O reader such as a reader from an object store. + /// let file = tokio::fs::File::open(path).await.unwrap(); + /// + /// // Configure options for reading from the async source + /// let builder = ParquetRecordBatchStreamBuilder::new(file) + /// .await + /// .unwrap(); + /// // Building the stream opens the parquet file (reads metadata, etc) and returns + /// // a stream that can be used to incrementally read the data in batches + /// let stream = builder.build().unwrap(); + /// // In this example, we collect the stream into a Vec + /// // but real applications would likely process the batches as they are read + /// let results = stream.try_collect::>().await.unwrap(); + /// // Demonstrate the results are as expected + /// assert_batches_eq( + /// &results, + /// &[ + /// "+----+----------+-------------+--------------+---------+------------+-----------+------------+------------------+------------+---------------------+", + /// "| id | bool_col | tinyint_col | smallint_col | int_col | bigint_col | float_col | double_col | date_string_col | string_col | timestamp_col |", + /// "+----+----------+-------------+--------------+---------+------------+-----------+------------+------------------+------------+---------------------+", + /// "| 4 | true | 0 | 0 | 0 | 0 | 0.0 | 0.0 | 30332f30312f3039 | 30 | 2009-03-01T00:00:00 |", + /// "| 5 | false | 1 | 1 | 1 | 10 | 1.1 | 10.1 | 30332f30312f3039 | 31 | 2009-03-01T00:01:00 |", + /// "| 6 | true | 0 | 0 | 0 | 0 | 0.0 | 0.0 | 30342f30312f3039 | 30 | 2009-04-01T00:00:00 |", + /// "| 7 | false | 1 | 1 | 1 | 10 | 1.1 | 10.1 | 30342f30312f3039 | 31 | 2009-04-01T00:01:00 |", + /// "| 2 | true | 0 | 0 | 0 | 0 | 0.0 | 0.0 | 30322f30312f3039 | 30 | 2009-02-01T00:00:00 |", + /// "| 3 | false | 1 | 1 | 1 | 10 | 1.1 | 10.1 | 30322f30312f3039 | 31 | 2009-02-01T00:01:00 |", + /// "| 0 | true | 0 | 0 | 0 | 0 | 0.0 | 0.0 | 30312f30312f3039 | 30 | 2009-01-01T00:00:00 |", + /// "| 1 | false | 1 | 1 | 1 | 10 | 1.1 | 10.1 | 30312f30312f3039 | 31 | 2009-01-01T00:01:00 |", + /// "+----+----------+-------------+--------------+---------+------------+-----------+------------+------------------+------------+---------------------+", + /// ], + /// ); + /// # } + /// ``` + /// + /// # Example Configuring Options and Reading Metadata + /// + /// There are many options that control the behavior of the reader, such as + /// `with_batch_size`, `with_projection`, `with_filter`, etc... + /// + /// ``` + /// # #[tokio::main(flavor="current_thread")] + /// # async fn main() { + /// # + /// # use arrow_array::RecordBatch; + /// # use arrow::util::pretty::pretty_format_batches; + /// # use futures::TryStreamExt; + /// # + /// # use parquet::arrow::{ParquetRecordBatchStreamBuilder, ProjectionMask}; + /// # + /// # fn assert_batches_eq(batches: &[RecordBatch], expected_lines: &[&str]) { + /// # let formatted = pretty_format_batches(batches).unwrap().to_string(); + /// # let actual_lines: Vec<_> = formatted.trim().lines().collect(); + /// # assert_eq!( + /// # &actual_lines, expected_lines, + /// # "\n\nexpected:\n\n{:#?}\nactual:\n\n{:#?}\n\n", + /// # expected_lines, actual_lines + /// # ); + /// # } + /// # + /// # let testdata = arrow::util::test_util::parquet_test_data(); + /// # let path = format!("{}/alltypes_plain.parquet", testdata); + /// // As before, use tokio::fs::File to read data using an async I/O. + /// let file = tokio::fs::File::open(path).await.unwrap(); + /// + /// // Configure options for reading from the async source, in this case we set the batch size + /// // to 3 which produces 3 rows at a time. + /// let builder = ParquetRecordBatchStreamBuilder::new(file) + /// .await + /// .unwrap() + /// .with_batch_size(3); + /// + /// // We can also read the metadata to inspect the schema and other metadata + /// // before actually reading the data + /// let file_metadata = builder.metadata().file_metadata(); + /// // Specify that we only want to read the 1st, 2nd, and 6th columns + /// let mask = ProjectionMask::roots(file_metadata.schema_descr(), [1, 2, 6]); + /// + /// let stream = builder.with_projection(mask).build().unwrap(); + /// let results = stream.try_collect::>().await.unwrap(); + /// // Print out the results + /// assert_batches_eq( + /// &results, + /// &[ + /// "+----------+-------------+-----------+", + /// "| bool_col | tinyint_col | float_col |", + /// "+----------+-------------+-----------+", + /// "| true | 0 | 0.0 |", + /// "| false | 1 | 1.1 |", + /// "| true | 0 | 0.0 |", + /// "| false | 1 | 1.1 |", + /// "| true | 0 | 0.0 |", + /// "| false | 1 | 1.1 |", + /// "| true | 0 | 0.0 |", + /// "| false | 1 | 1.1 |", + /// "+----------+-------------+-----------+", + /// ], + /// ); + /// + /// // The results has 8 rows, so since we set the batch size to 3, we expect + /// // 3 batches, two with 3 rows each and the last batch with 2 rows. + /// assert_eq!(results.len(), 3); + /// # } + /// ``` + /// + /// # Example reading Row Groups in Parallel + /// + /// Each [`ParquetRecordBatchStream`] is independent and can be used to read + /// from the same underlying source in parallel. Use + /// [`ParquetRecordBatchStream::next_row_group`] with a single stream to + /// begin prefetching the next Row Group. To read a file in parallel, create + /// a stream for each subset of the file. For example, you can read each + /// row group in parallel by creating a stream for each row group using the + /// [`ParquetRecordBatchStreamBuilder::with_row_groups`] API as shown below + /// + /// ``` + /// # use std::sync::Arc; + /// # use arrow_array::{ArrayRef, Int32Array, RecordBatch}; + /// # use arrow::util::pretty::pretty_format_batches; + /// # use futures::{StreamExt, TryStreamExt}; + /// # use tempfile::NamedTempFile; + /// # use parquet::arrow::{ArrowWriter, ParquetRecordBatchStreamBuilder, ProjectionMask}; + /// # use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; + /// # use parquet::file::metadata::ParquetMetaDataReader; + /// # use parquet::file::properties::{WriterProperties}; + /// # // write to a temporary file with 10 RowGroups and read back with async API + /// # fn write_file() -> parquet::errors::Result { + /// # let mut file = NamedTempFile::new().unwrap(); + /// # let small_batch = RecordBatch::try_from_iter([ + /// # ("id", Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4])) as ArrayRef), + /// # ]).unwrap(); + /// # let props = WriterProperties::builder() + /// # .set_max_row_group_row_count(Some(5)) + /// # .set_write_batch_size(5) + /// # .build(); + /// # let mut writer = ArrowWriter::try_new(&mut file, small_batch.schema(), Some(props))?; + /// # for i in 0..10 { + /// # writer.write(&small_batch)? + /// # }; + /// # writer.close()?; + /// # Ok(file) + /// # } + /// # #[tokio::main(flavor="current_thread")] + /// # async fn main() -> parquet::errors::Result<()> { + /// # let t = write_file()?; + /// # let path = t.path(); + /// // This example uses a tokio::fs::File as the async source, but it + /// // could be any async source such as an object store reader) + /// let mut file = tokio::fs::File::open(path).await?; + /// // To read Row Groups in parallel, create a separate stream builder for each Row Group. + /// // First get the metadata to find the row group information + /// let file_size = file.metadata().await?.len(); + /// let metadata = ParquetMetaDataReader::new().load_and_finish(&mut file, file_size).await?; + /// assert_eq!(metadata.num_row_groups(), 10); // file has 10 row groups with 5 rows each + /// // Create a stream reader for each row group + /// let reader_metadata = ArrowReaderMetadata::try_new( + /// Arc::new(metadata), + /// ArrowReaderOptions::new() + /// )?; + /// let mut streams = vec![]; + /// for row_group_index in 0..10 { + /// // Each stream needs its own source instance to issue + /// // parallel IO requests, so clone the file for each stream + /// let this_file = file.try_clone().await?; + /// let stream = ParquetRecordBatchStreamBuilder::new_with_metadata( + /// this_file, + /// reader_metadata.clone() + /// ) + /// .with_row_groups(vec![row_group_index]) // read only this row group + /// .build()?; + /// streams.push(stream); + /// } + /// // Each reader can now be polled independently and in parallel, for + /// // example using StreamExt::buffered to read from 3 at a time + /// let results = futures::stream::iter(streams) + /// .map(|stream| async move { stream }) + /// .buffered(3) + /// .flatten() + /// .try_collect::>().await?; + /// // read all 50 rows (10 row groups x 5 rows per group) + /// assert_eq!(50, results.iter().map(|s| s.num_rows()).sum::()); + /// # Ok(()) + /// # } + /// ``` + pub async fn new(input: T) -> Result { + Self::new_with_options(input, Default::default()).await + } + + /// Create a new [`ParquetRecordBatchStreamBuilder`] with the provided async source + /// and [`ArrowReaderOptions`]. + pub async fn new_with_options(mut input: T, options: ArrowReaderOptions) -> Result { + let metadata = ArrowReaderMetadata::load_async(&mut input, options).await?; + Ok(Self::new_with_metadata(input, metadata)) + } + + /// Create a [`ParquetRecordBatchStreamBuilder`] from the provided [`ArrowReaderMetadata`] + /// + /// This allows loading metadata once and using it to create multiple builders with + /// potentially different settings, that can be read in parallel. + /// + /// # Example of reading from multiple streams in parallel + /// + /// ``` + /// # use std::fs::metadata; + /// # use std::sync::Arc; + /// # use bytes::Bytes; + /// # use arrow_array::{Int32Array, RecordBatch}; + /// # use arrow_schema::{DataType, Field, Schema}; + /// # use parquet::arrow::arrow_reader::ArrowReaderMetadata; + /// # use parquet::arrow::{ArrowWriter, ParquetRecordBatchStreamBuilder}; + /// # use tempfile::tempfile; + /// # use futures::StreamExt; + /// # #[tokio::main(flavor="current_thread")] + /// # async fn main() { + /// # + /// # let mut file = tempfile().unwrap(); + /// # let schema = Arc::new(Schema::new(vec![Field::new("i32", DataType::Int32, false)])); + /// # let mut writer = ArrowWriter::try_new(&mut file, schema.clone(), None).unwrap(); + /// # let batch = RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1, 2, 3]))]).unwrap(); + /// # writer.write(&batch).unwrap(); + /// # writer.close().unwrap(); + /// // open file with parquet data + /// let mut file = tokio::fs::File::from_std(file); + /// // load metadata once + /// let meta = ArrowReaderMetadata::load_async(&mut file, Default::default()).await.unwrap(); + /// // create two readers, a and b, from the same underlying file + /// // without reading the metadata again + /// let mut a = ParquetRecordBatchStreamBuilder::new_with_metadata( + /// file.try_clone().await.unwrap(), + /// meta.clone() + /// ).build().unwrap(); + /// let mut b = ParquetRecordBatchStreamBuilder::new_with_metadata(file, meta).build().unwrap(); + /// + /// // Can read batches from both readers in parallel + /// assert_eq!( + /// a.next().await.unwrap().unwrap(), + /// b.next().await.unwrap().unwrap(), + /// ); + /// # } + /// ``` + pub fn new_with_metadata(input: T, metadata: ArrowReaderMetadata) -> Self { + Self::new_builder(AsyncReader(input), metadata) + } + + /// Read bloom filter for a column in a row group + /// + /// Returns `None` if the column does not have a bloom filter + /// + /// We should call this function after other forms pruning, such as projection and predicate pushdown. + pub async fn get_row_group_column_bloom_filter( + &mut self, + row_group_idx: usize, + column_idx: usize, + ) -> Result> { + let metadata = self.metadata.row_group(row_group_idx); + let column_metadata = metadata.column(column_idx); + + let offset: u64 = if let Some(offset) = column_metadata.bloom_filter_offset() { + offset + .try_into() + .map_err(|_| ParquetError::General("Bloom filter offset is invalid".to_string()))? + } else { + return Ok(None); + }; + + let buffer = match column_metadata.bloom_filter_length() { + Some(length) => self.input.0.get_bytes(offset..offset + length as u64), + None => self + .input + .0 + .get_bytes(offset..offset + SBBF_HEADER_SIZE_ESTIMATE as u64), + } + .await?; + + let (header, bitset_offset) = + chunk_read_bloom_filter_header_and_offset(offset, buffer.clone())?; + + match header.algorithm { + BloomFilterAlgorithm::BLOCK => { + // this match exists to future proof the singleton algorithm enum + } + } + match header.compression { + BloomFilterCompression::UNCOMPRESSED => { + // this match exists to future proof the singleton compression enum + } + } + match header.hash { + BloomFilterHash::XXHASH => { + // this match exists to future proof the singleton hash enum + } + } + + let bitset = match column_metadata.bloom_filter_length() { + Some(_) => buffer.slice( + (TryInto::::try_into(bitset_offset).unwrap() + - TryInto::::try_into(offset).unwrap()).., + ), + None => { + let bitset_length: u64 = header.num_bytes.try_into().map_err(|_| { + ParquetError::General("Bloom filter length is invalid".to_string()) + })?; + self.input + .0 + .get_bytes(bitset_offset..bitset_offset + bitset_length) + .await? + } + }; + Ok(Some(Sbbf::new(&bitset))) + } + + /// Build a new [`ParquetRecordBatchStream`] + /// + /// See examples on [`ParquetRecordBatchStreamBuilder::new`] + pub fn build(self) -> Result> { + let Self { + input, + metadata, + schema, + fields, + batch_size, + row_groups, + projection, + filter, + selection, + row_selection_policy: selection_strategy, + limit, + offset, + metrics, + max_predicate_cache_size, + } = self; + + // Ensure schema of ParquetRecordBatchStream respects projection, and does + // not store metadata (same as for ParquetRecordBatchReader and emitted RecordBatches) + let projection_len = projection.mask.as_ref().map_or(usize::MAX, |m| m.len()); + let projected_fields = schema + .fields + .filter_leaves(|idx, _| idx < projection_len && projection.leaf_included(idx)); + let projected_schema = Arc::new(Schema::new(projected_fields)); + + let decoder = ParquetPushDecoderBuilder { + input: PushDecoderInput::default(), + metadata, + schema, + fields, + projection, + filter, + selection, + row_selection_policy: selection_strategy, + batch_size, + row_groups, + limit, + offset, + metrics, + max_predicate_cache_size, + } + .build()?; + + let request_state = RequestState::None { input: input.0 }; + + Ok(ParquetRecordBatchStream { + schema: projected_schema, + decoder, + request_state, + }) + } +} + +/// State machine that tracks outstanding requests to fetch data +/// +/// The parameter `T` is the input, typically an `AsyncFileReader` +enum RequestState { + /// No outstanding requests + None { + input: T, + }, + /// There is an outstanding request for data + Outstanding { + /// Ranges that have been requested + ranges: Vec>, + /// Future that will resolve (input, requested_ranges) + /// + /// Note the future owns the reader while the request is outstanding + /// and returns it upon completion + future: BoxFuture<'static, Result<(T, Vec)>>, + }, + Done, +} + +impl RequestState +where + T: AsyncFileReader + Unpin + Send + 'static, +{ + /// Issue a request to fetch `ranges`, returning the Outstanding state + fn begin_request(mut input: T, ranges: Vec>) -> Self { + let ranges_captured = ranges.clone(); + + // Note this must move the input *into* the future + // because the get_byte_ranges future has a lifetime + // (aka can have references internally) and thus must + // own the input while the request is outstanding. + let future = async move { + let data = input.get_byte_ranges(ranges_captured).await?; + Ok((input, data)) + } + .boxed(); + RequestState::Outstanding { ranges, future } + } +} + +impl std::fmt::Debug for RequestState { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self { + RequestState::None { input: _ } => f + .debug_struct("RequestState::None") + .field("input", &"...") + .finish(), + RequestState::Outstanding { ranges, .. } => f + .debug_struct("RequestState::Outstanding") + .field("ranges", &ranges) + .finish(), + RequestState::Done => { + write!(f, "RequestState::Done") + } + } + } +} + +/// An asynchronous [`Stream`]of [`RecordBatch`] constructed using [`ParquetRecordBatchStreamBuilder`] to read parquet files. +/// +/// `ParquetRecordBatchStream` also provides [`ParquetRecordBatchStream::next_row_group`] for fetching row groups, +/// allowing users to decode record batches separately from I/O. +/// +/// # I/O Buffering +/// +/// `ParquetRecordBatchStream` buffers *all* data pages selected after predicates +/// (projection + filtering, etc) and decodes the rows from those buffered pages. +/// +/// For example, if all rows and columns are selected, the entire row group is +/// buffered in memory during decode. This minimizes the number of IO operations +/// required, which is especially important for object stores, where IO operations +/// have latencies in the hundreds of milliseconds +/// +/// See [`ParquetPushDecoderBuilder`] for an API with lower level control over +/// buffering. +/// +/// [`Stream`]: https://docs.rs/futures/latest/futures/stream/trait.Stream.html +pub struct ParquetRecordBatchStream { + /// Output schema of the stream + schema: SchemaRef, + /// Input and Outstanding IO request, if any + request_state: RequestState, + /// Decoding state machine (no IO) + decoder: ParquetPushDecoder, +} + +impl std::fmt::Debug for ParquetRecordBatchStream { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ParquetRecordBatchStream") + .field("request_state", &self.request_state) + .finish() + } +} + +impl ParquetRecordBatchStream { + /// Returns the projected [`SchemaRef`] for reading the parquet file. + /// + /// Note that the schema metadata will be stripped here. See + /// [`ParquetRecordBatchStreamBuilder::schema`] if the metadata is desired. + pub fn schema(&self) -> &SchemaRef { + &self.schema + } +} + +impl ParquetRecordBatchStream +where + T: AsyncFileReader + Unpin + Send + 'static, +{ + /// Fetches the next row group from the stream. + /// + /// Users can continue to call this function to get row groups and decode them concurrently. + /// + /// ## Notes + /// + /// ParquetRecordBatchStream should be used either as a `Stream` or with `next_row_group`; they should not be used simultaneously. + /// + /// ## Returns + /// + /// - `Ok(None)` if the stream has ended. + /// - `Err(error)` if the stream has errored. All subsequent calls will return `Ok(None)`. + /// - `Ok(Some(reader))` which holds all the data for the row group. + pub async fn next_row_group(&mut self) -> Result> { + loop { + // Take ownership of request state to process, leaving self in a + // valid state + let request_state = std::mem::replace(&mut self.request_state, RequestState::Done); + match request_state { + // No outstanding requests, proceed to setup next row group + RequestState::None { input } => { + match self.decoder.try_next_reader()? { + DecodeResult::NeedsData(ranges) => { + self.request_state = RequestState::begin_request(input, ranges); + continue; // poll again (as the input might be ready immediately) + } + DecodeResult::Data(reader) => { + self.request_state = RequestState::None { input }; + return Ok(Some(reader)); + } + DecodeResult::Finished => return Ok(None), + } + } + RequestState::Outstanding { ranges, future } => { + let (input, data) = future.await?; + // Push the requested data to the decoder and try again + self.decoder.push_ranges(ranges, data)?; + self.request_state = RequestState::None { input }; + continue; // try and decode on next iteration + } + RequestState::Done => { + self.request_state = RequestState::Done; + return Ok(None); + } + } + } + } +} + +impl Stream for ParquetRecordBatchStream +where + T: AsyncFileReader + Unpin + Send + 'static, +{ + type Item = Result; + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.poll_next_inner(cx) { + Ok(res) => { + // Successfully decoded a batch, or reached end of stream. + // convert Option to Option> + res.map(|res| Ok(res).transpose()) + } + Err(e) => { + self.request_state = RequestState::Done; + Poll::Ready(Some(Err(e))) + } + } + } +} + +impl ParquetRecordBatchStream +where + T: AsyncFileReader + Unpin + Send + 'static, +{ + /// Inner state machine + /// + /// Note this is separate from poll_next so we can use ? operator to check for errors + /// as it returns `Result>>` + fn poll_next_inner(&mut self, cx: &mut Context<'_>) -> Result>> { + loop { + let request_state = std::mem::replace(&mut self.request_state, RequestState::Done); + match request_state { + RequestState::None { input } => { + // No outstanding requests, proceed to decode the next batch + match self.decoder.try_decode()? { + DecodeResult::NeedsData(ranges) => { + self.request_state = RequestState::begin_request(input, ranges); + continue; // poll again (as the input might be ready immediately) + } + DecodeResult::Data(batch) => { + self.request_state = RequestState::None { input }; + return Ok(Poll::Ready(Some(batch))); + } + DecodeResult::Finished => { + self.request_state = RequestState::Done; + return Ok(Poll::Ready(None)); + } + } + } + RequestState::Outstanding { ranges, mut future } => match future.poll_unpin(cx) { + // Data was ready, push it to the decoder and continue + Poll::Ready(result) => { + let (input, data) = result?; + // Push the requested data to the decoder + self.decoder.push_ranges(ranges, data)?; + self.request_state = RequestState::None { input }; + continue; // next iteration will try to decode the next batch + } + Poll::Pending => { + self.request_state = RequestState::Outstanding { ranges, future }; + return Ok(Poll::Pending); + } + }, + RequestState::Done => { + // Stream is done (error or end), return None + self.request_state = RequestState::Done; + return Ok(Poll::Ready(None)); + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::arrow_reader::tests::test_row_numbers_with_multiple_row_groups_helper; + use crate::arrow::arrow_reader::{ + ArrowPredicateFn, ParquetRecordBatchReaderBuilder, RowFilter, RowSelection, RowSelector, + }; + use crate::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; + use crate::arrow::schema::virtual_type::RowNumber; + use crate::arrow::{ArrowWriter, AsyncArrowWriter, ProjectionMask}; + use crate::file::metadata::PageIndexPolicy; + use crate::file::metadata::ParquetMetaDataReader; + use crate::file::properties::WriterProperties; + use arrow::compute::kernels::cmp::eq; + use arrow::error::Result as ArrowResult; + use arrow_array::builder::{Float32Builder, ListBuilder, StringBuilder}; + use arrow_array::cast::AsArray; + use arrow_array::types::Int32Type; + use arrow_array::{ + Array, ArrayRef, BooleanArray, Int32Array, RecordBatchReader, Scalar, StringArray, + StructArray, UInt64Array, + }; + use arrow_schema::{DataType, Field, Schema}; + use futures::{StreamExt, TryStreamExt}; + use rand::{Rng, rng}; + use std::collections::HashMap; + use std::sync::{Arc, Mutex}; + use tempfile::tempfile; + + #[derive(Clone)] + struct TestReader { + data: Bytes, + metadata: Option>, + requests: Arc>>>, + } + + impl TestReader { + fn new(data: Bytes) -> Self { + Self { + data, + metadata: Default::default(), + requests: Default::default(), + } + } + } + + impl AsyncFileReader for TestReader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + let range = range.clone(); + self.requests + .lock() + .unwrap() + .push(range.start as usize..range.end as usize); + futures::future::ready(Ok(self + .data + .slice(range.start as usize..range.end as usize))) + .boxed() + } + + fn get_metadata<'a>( + &'a mut self, + options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, Result>> { + let metadata_reader = ParquetMetaDataReader::new().with_arrow_reader_options(options); + self.metadata = Some(Arc::new( + metadata_reader.parse_and_finish(&self.data).unwrap(), + )); + futures::future::ready(Ok(self.metadata.clone().unwrap().clone())).boxed() + } + } + + #[tokio::test] + async fn test_async_reader() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_plain.parquet"); + let data = Bytes::from(std::fs::read(path).unwrap()); + + let async_reader = TestReader::new(data.clone()); + + let requests = async_reader.requests.clone(); + let builder = ParquetRecordBatchStreamBuilder::new(async_reader) + .await + .unwrap(); + + let metadata = builder.metadata().clone(); + assert_eq!(metadata.num_row_groups(), 1); + + let mask = ProjectionMask::leaves(builder.parquet_schema(), vec![1, 2]); + let stream = builder + .with_projection(mask.clone()) + .with_batch_size(1024) + .build() + .unwrap(); + + let async_batches: Vec<_> = stream.try_collect().await.unwrap(); + + let sync_batches = ParquetRecordBatchReaderBuilder::try_new(data) + .unwrap() + .with_projection(mask) + .with_batch_size(104) + .build() + .unwrap() + .collect::>>() + .unwrap(); + + assert_eq!(async_batches, sync_batches); + + let requests = requests.lock().unwrap(); + let (offset_1, length_1) = metadata.row_group(0).column(1).byte_range(); + let (offset_2, length_2) = metadata.row_group(0).column(2).byte_range(); + + assert_eq!( + &requests[..], + &[ + offset_1 as usize..(offset_1 + length_1) as usize, + offset_2 as usize..(offset_2 + length_2) as usize + ] + ); + } + + #[tokio::test] + async fn test_async_reader_with_next_row_group() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_plain.parquet"); + let data = Bytes::from(std::fs::read(path).unwrap()); + + let async_reader = TestReader::new(data.clone()); + + let requests = async_reader.requests.clone(); + let builder = ParquetRecordBatchStreamBuilder::new(async_reader) + .await + .unwrap(); + + let metadata = builder.metadata().clone(); + assert_eq!(metadata.num_row_groups(), 1); + + let mask = ProjectionMask::leaves(builder.parquet_schema(), vec![1, 2]); + let mut stream = builder + .with_projection(mask.clone()) + .with_batch_size(1024) + .build() + .unwrap(); + + let mut readers = vec![]; + while let Some(reader) = stream.next_row_group().await.unwrap() { + readers.push(reader); + } + + let async_batches: Vec<_> = readers + .into_iter() + .flat_map(|r| r.map(|v| v.unwrap()).collect::>()) + .collect(); + + let sync_batches = ParquetRecordBatchReaderBuilder::try_new(data) + .unwrap() + .with_projection(mask) + .with_batch_size(104) + .build() + .unwrap() + .collect::>>() + .unwrap(); + + assert_eq!(async_batches, sync_batches); + + let requests = requests.lock().unwrap(); + let (offset_1, length_1) = metadata.row_group(0).column(1).byte_range(); + let (offset_2, length_2) = metadata.row_group(0).column(2).byte_range(); + + assert_eq!( + &requests[..], + &[ + offset_1 as usize..(offset_1 + length_1) as usize, + offset_2 as usize..(offset_2 + length_2) as usize + ] + ); + } + + #[tokio::test] + async fn test_async_reader_with_index() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet"); + let data = Bytes::from(std::fs::read(path).unwrap()); + + let async_reader = TestReader::new(data.clone()); + + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let builder = ParquetRecordBatchStreamBuilder::new_with_options(async_reader, options) + .await + .unwrap(); + + // The builder should have page and offset indexes loaded now + let metadata_with_index = builder.metadata(); + assert_eq!(metadata_with_index.num_row_groups(), 1); + + // Check offset indexes are present for all columns + let offset_index = metadata_with_index.offset_index().unwrap(); + let column_index = metadata_with_index.column_index().unwrap(); + + assert_eq!(offset_index.len(), metadata_with_index.num_row_groups()); + assert_eq!(column_index.len(), metadata_with_index.num_row_groups()); + + let num_columns = metadata_with_index + .file_metadata() + .schema_descr() + .num_columns(); + + // Check page indexes are present for all columns + offset_index + .iter() + .for_each(|x| assert_eq!(x.len(), num_columns)); + column_index + .iter() + .for_each(|x| assert_eq!(x.len(), num_columns)); + + let mask = ProjectionMask::leaves(builder.parquet_schema(), vec![1, 2]); + let stream = builder + .with_projection(mask.clone()) + .with_batch_size(1024) + .build() + .unwrap(); + + let async_batches: Vec<_> = stream.try_collect().await.unwrap(); + + let sync_batches = ParquetRecordBatchReaderBuilder::try_new(data) + .unwrap() + .with_projection(mask) + .with_batch_size(1024) + .build() + .unwrap() + .collect::>>() + .unwrap(); + + assert_eq!(async_batches, sync_batches); + } + + #[tokio::test] + async fn test_async_reader_with_limit() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet"); + let data = Bytes::from(std::fs::read(path).unwrap()); + + let metadata = ParquetMetaDataReader::new() + .parse_and_finish(&data) + .unwrap(); + let metadata = Arc::new(metadata); + + assert_eq!(metadata.num_row_groups(), 1); + + let async_reader = TestReader::new(data.clone()); + + let builder = ParquetRecordBatchStreamBuilder::new(async_reader) + .await + .unwrap(); + + assert_eq!(builder.metadata().num_row_groups(), 1); + + let mask = ProjectionMask::leaves(builder.parquet_schema(), vec![1, 2]); + let stream = builder + .with_projection(mask.clone()) + .with_batch_size(1024) + .with_limit(1) + .build() + .unwrap(); + + let async_batches: Vec<_> = stream.try_collect().await.unwrap(); + + let sync_batches = ParquetRecordBatchReaderBuilder::try_new(data) + .unwrap() + .with_projection(mask) + .with_batch_size(1024) + .with_limit(1) + .build() + .unwrap() + .collect::>>() + .unwrap(); + + assert_eq!(async_batches, sync_batches); + } + + #[tokio::test] + async fn test_async_reader_skip_pages() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet"); + let data = Bytes::from(std::fs::read(path).unwrap()); + + let async_reader = TestReader::new(data.clone()); + + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let builder = ParquetRecordBatchStreamBuilder::new_with_options(async_reader, options) + .await + .unwrap(); + + assert_eq!(builder.metadata().num_row_groups(), 1); + + let selection = RowSelection::from(vec![ + RowSelector::skip(21), // Skip first page + RowSelector::select(21), // Select page to boundary + RowSelector::skip(41), // Skip multiple pages + RowSelector::select(41), // Select multiple pages + RowSelector::skip(25), // Skip page across boundary + RowSelector::select(25), // Select across page boundary + RowSelector::skip(7116), // Skip to final page boundary + RowSelector::select(10), // Select final page + ]); + + let mask = ProjectionMask::leaves(builder.parquet_schema(), vec![9]); + + let stream = builder + .with_projection(mask.clone()) + .with_row_selection(selection.clone()) + .build() + .expect("building stream"); + + let async_batches: Vec<_> = stream.try_collect().await.unwrap(); + + let sync_batches = ParquetRecordBatchReaderBuilder::try_new(data) + .unwrap() + .with_projection(mask) + .with_batch_size(1024) + .with_row_selection(selection) + .build() + .unwrap() + .collect::>>() + .unwrap(); + + assert_eq!(async_batches, sync_batches); + } + + #[tokio::test] + async fn test_fuzz_async_reader_selection() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet"); + let data = Bytes::from(std::fs::read(path).unwrap()); + + let mut rand = rng(); + + for _ in 0..100 { + let mut expected_rows = 0; + let mut total_rows = 0; + let mut skip = false; + let mut selectors = vec![]; + + while total_rows < 7300 { + let row_count: usize = rand.random_range(1..100); + + let row_count = row_count.min(7300 - total_rows); + + selectors.push(RowSelector { row_count, skip }); + + total_rows += row_count; + if !skip { + expected_rows += row_count; + } + + skip = !skip; + } + + let selection = RowSelection::from(selectors); + + let async_reader = TestReader::new(data.clone()); + + let options = + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let builder = ParquetRecordBatchStreamBuilder::new_with_options(async_reader, options) + .await + .unwrap(); + + assert_eq!(builder.metadata().num_row_groups(), 1); + + let col_idx: usize = rand.random_range(0..13); + let mask = ProjectionMask::leaves(builder.parquet_schema(), vec![col_idx]); + + let stream = builder + .with_projection(mask.clone()) + .with_row_selection(selection.clone()) + .build() + .expect("building stream"); + + let async_batches: Vec<_> = stream.try_collect().await.unwrap(); + + let actual_rows: usize = async_batches.into_iter().map(|b| b.num_rows()).sum(); + + assert_eq!(actual_rows, expected_rows); + } + } + + #[tokio::test] + async fn test_async_reader_zero_row_selector() { + //See https://github.com/apache/arrow-rs/issues/2669 + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet"); + let data = Bytes::from(std::fs::read(path).unwrap()); + + let mut rand = rng(); + + let mut expected_rows = 0; + let mut total_rows = 0; + let mut skip = false; + let mut selectors = vec![]; + + selectors.push(RowSelector { + row_count: 0, + skip: false, + }); + + while total_rows < 7300 { + let row_count: usize = rand.random_range(1..100); + + let row_count = row_count.min(7300 - total_rows); + + selectors.push(RowSelector { row_count, skip }); + + total_rows += row_count; + if !skip { + expected_rows += row_count; + } + + skip = !skip; + } + + let selection = RowSelection::from(selectors); + + let async_reader = TestReader::new(data.clone()); + + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let builder = ParquetRecordBatchStreamBuilder::new_with_options(async_reader, options) + .await + .unwrap(); + + assert_eq!(builder.metadata().num_row_groups(), 1); + + let col_idx: usize = rand.random_range(0..13); + let mask = ProjectionMask::leaves(builder.parquet_schema(), vec![col_idx]); + + let stream = builder + .with_projection(mask.clone()) + .with_row_selection(selection.clone()) + .build() + .expect("building stream"); + + let async_batches: Vec<_> = stream.try_collect().await.unwrap(); + + let actual_rows: usize = async_batches.into_iter().map(|b| b.num_rows()).sum(); + + assert_eq!(actual_rows, expected_rows); + } + + #[tokio::test] + async fn test_limit_multiple_row_groups() { + let a = StringArray::from_iter_values(["a", "b", "b", "b", "c", "c"]); + let b = StringArray::from_iter_values(["1", "2", "3", "4", "5", "6"]); + let c = Int32Array::from_iter(0..6); + let data = RecordBatch::try_from_iter([ + ("a", Arc::new(a) as ArrayRef), + ("b", Arc::new(b) as ArrayRef), + ("c", Arc::new(c) as ArrayRef), + ]) + .unwrap(); + + let mut buf = Vec::with_capacity(1024); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .build(); + let mut writer = ArrowWriter::try_new(&mut buf, data.schema(), Some(props)).unwrap(); + writer.write(&data).unwrap(); + writer.close().unwrap(); + + let data: Bytes = buf.into(); + let metadata = ParquetMetaDataReader::new() + .parse_and_finish(&data) + .unwrap(); + + assert_eq!(metadata.num_row_groups(), 2); + + let test = TestReader::new(data); + + let stream = ParquetRecordBatchStreamBuilder::new(test.clone()) + .await + .unwrap() + .with_batch_size(1024) + .with_limit(4) + .build() + .unwrap(); + + let batches: Vec<_> = stream.try_collect().await.unwrap(); + // Expect one batch for each row group + assert_eq!(batches.len(), 2); + + let batch = &batches[0]; + // First batch should contain all rows + assert_eq!(batch.num_rows(), 3); + assert_eq!(batch.num_columns(), 3); + let col2 = batch.column(2).as_primitive::(); + assert_eq!(col2.values(), &[0, 1, 2]); + + let batch = &batches[1]; + // Second batch should trigger the limit and only have one row + assert_eq!(batch.num_rows(), 1); + assert_eq!(batch.num_columns(), 3); + let col2 = batch.column(2).as_primitive::(); + assert_eq!(col2.values(), &[3]); + + let stream = ParquetRecordBatchStreamBuilder::new(test.clone()) + .await + .unwrap() + .with_offset(2) + .with_limit(3) + .build() + .unwrap(); + + let batches: Vec<_> = stream.try_collect().await.unwrap(); + // Expect one batch for each row group + assert_eq!(batches.len(), 2); + + let batch = &batches[0]; + // First batch should contain one row + assert_eq!(batch.num_rows(), 1); + assert_eq!(batch.num_columns(), 3); + let col2 = batch.column(2).as_primitive::(); + assert_eq!(col2.values(), &[2]); + + let batch = &batches[1]; + // Second batch should contain two rows + assert_eq!(batch.num_rows(), 2); + assert_eq!(batch.num_columns(), 3); + let col2 = batch.column(2).as_primitive::(); + assert_eq!(col2.values(), &[3, 4]); + + let stream = ParquetRecordBatchStreamBuilder::new(test.clone()) + .await + .unwrap() + .with_offset(4) + .with_limit(20) + .build() + .unwrap(); + + let batches: Vec<_> = stream.try_collect().await.unwrap(); + // Should skip first row group + assert_eq!(batches.len(), 1); + + let batch = &batches[0]; + // First batch should contain two rows + assert_eq!(batch.num_rows(), 2); + assert_eq!(batch.num_columns(), 3); + let col2 = batch.column(2).as_primitive::(); + assert_eq!(col2.values(), &[4, 5]); + } + + #[tokio::test] + async fn test_batch_size_overallocate() { + let testdata = arrow::util::test_util::parquet_test_data(); + // `alltypes_plain.parquet` only have 8 rows + let path = format!("{testdata}/alltypes_plain.parquet"); + let data = Bytes::from(std::fs::read(path).unwrap()); + + let async_reader = TestReader::new(data.clone()); + + let builder = ParquetRecordBatchStreamBuilder::new(async_reader) + .await + .unwrap(); + + let file_rows = builder.metadata().file_metadata().num_rows() as usize; + + let builder = builder + .with_projection(ProjectionMask::all()) + .with_batch_size(1024); + + // even though the batch size is set to 1024, it should adjust to the max + // number of rows in the file (8) + assert_ne!(1024, file_rows); + assert_eq!(builder.batch_size, file_rows); + + let _stream = builder.build().unwrap(); + } + + #[tokio::test] + async fn test_parquet_record_batch_stream_schema() { + fn get_all_field_names(schema: &Schema) -> Vec<&String> { + schema.flattened_fields().iter().map(|f| f.name()).collect() + } + + // ParquetRecordBatchReaderBuilder::schema differs from + // ParquetRecordBatchReader::schema and RecordBatch::schema in the returned + // schema contents (in terms of custom metadata attached to schema, and fields + // returned). Test to ensure this remains consistent behaviour. + // + // Ensure same for asynchronous versions of the above. + + // Prep data, for a schema with nested fields, with custom metadata + let mut metadata = HashMap::with_capacity(1); + metadata.insert("key".to_string(), "value".to_string()); + + let nested_struct_array = StructArray::from(vec![ + ( + Arc::new(Field::new("d", DataType::Utf8, true)), + Arc::new(StringArray::from(vec!["a", "b"])) as ArrayRef, + ), + ( + Arc::new(Field::new("e", DataType::Utf8, true)), + Arc::new(StringArray::from(vec!["c", "d"])) as ArrayRef, + ), + ]); + let struct_array = StructArray::from(vec![ + ( + Arc::new(Field::new("a", DataType::Int32, true)), + Arc::new(Int32Array::from(vec![-1, 1])) as ArrayRef, + ), + ( + Arc::new(Field::new("b", DataType::UInt64, true)), + Arc::new(UInt64Array::from(vec![1, 2])) as ArrayRef, + ), + ( + Arc::new(Field::new( + "c", + nested_struct_array.data_type().clone(), + true, + )), + Arc::new(nested_struct_array) as ArrayRef, + ), + ]); + + let schema = + Arc::new(Schema::new(struct_array.fields().clone()).with_metadata(metadata.clone())); + let record_batch = RecordBatch::from(struct_array) + .with_schema(schema.clone()) + .unwrap(); + + // Write parquet with custom metadata in schema + let mut file = tempfile().unwrap(); + let mut writer = ArrowWriter::try_new(&mut file, schema.clone(), None).unwrap(); + writer.write(&record_batch).unwrap(); + writer.close().unwrap(); + + let all_fields = ["a", "b", "c", "d", "e"]; + // (leaf indices in mask, expected names in output schema all fields) + let projections = [ + (vec![], vec![]), + (vec![0], vec!["a"]), + (vec![0, 1], vec!["a", "b"]), + (vec![0, 1, 2], vec!["a", "b", "c", "d"]), + (vec![0, 1, 2, 3], vec!["a", "b", "c", "d", "e"]), + ]; + + // Ensure we're consistent for each of these projections + for (indices, expected_projected_names) in projections { + let assert_schemas = |builder: SchemaRef, reader: SchemaRef, batch: SchemaRef| { + // Builder schema should preserve all fields and metadata + assert_eq!(get_all_field_names(&builder), all_fields); + assert_eq!(builder.metadata, metadata); + // Reader & batch schema should show only projected fields, and no metadata + assert_eq!(get_all_field_names(&reader), expected_projected_names); + assert_eq!(reader.metadata, HashMap::default()); + assert_eq!(get_all_field_names(&batch), expected_projected_names); + assert_eq!(batch.metadata, HashMap::default()); + }; + + let builder = + ParquetRecordBatchReaderBuilder::try_new(file.try_clone().unwrap()).unwrap(); + let sync_builder_schema = builder.schema().clone(); + let mask = ProjectionMask::leaves(builder.parquet_schema(), indices.clone()); + let mut reader = builder.with_projection(mask).build().unwrap(); + let sync_reader_schema = reader.schema(); + let batch = reader.next().unwrap().unwrap(); + let sync_batch_schema = batch.schema(); + assert_schemas(sync_builder_schema, sync_reader_schema, sync_batch_schema); + + // asynchronous should be same + let file = tokio::fs::File::from(file.try_clone().unwrap()); + let builder = ParquetRecordBatchStreamBuilder::new(file).await.unwrap(); + let async_builder_schema = builder.schema().clone(); + let mask = ProjectionMask::leaves(builder.parquet_schema(), indices); + let mut reader = builder.with_projection(mask).build().unwrap(); + let async_reader_schema = reader.schema().clone(); + let batch = reader.next().await.unwrap().unwrap(); + let async_batch_schema = batch.schema(); + assert_schemas( + async_builder_schema, + async_reader_schema, + async_batch_schema, + ); + } + } + + #[tokio::test] + async fn test_nested_skip() { + let schema = Arc::new(Schema::new(vec![ + Field::new("col_1", DataType::UInt64, false), + Field::new_list("col_2", Field::new_list_field(DataType::Utf8, true), true), + ])); + + // Default writer properties + let props = WriterProperties::builder() + .set_data_page_row_count_limit(256) + .set_write_batch_size(256) + .set_max_row_group_row_count(Some(1024)); + + // Write data + let mut file = tempfile().unwrap(); + let mut writer = + ArrowWriter::try_new(&mut file, schema.clone(), Some(props.build())).unwrap(); + + let mut builder = ListBuilder::new(StringBuilder::new()); + for id in 0..1024 { + match id % 3 { + 0 => builder.append_value([Some("val_1".to_string()), Some(format!("id_{id}"))]), + 1 => builder.append_value([Some(format!("id_{id}"))]), + _ => builder.append_null(), + } + } + let refs = vec![ + Arc::new(UInt64Array::from_iter_values(0..1024)) as ArrayRef, + Arc::new(builder.finish()) as ArrayRef, + ]; + + let batch = RecordBatch::try_new(schema.clone(), refs).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let selections = [ + RowSelection::from(vec![ + RowSelector::skip(313), + RowSelector::select(1), + RowSelector::skip(709), + RowSelector::select(1), + ]), + RowSelection::from(vec![ + RowSelector::skip(255), + RowSelector::select(1), + RowSelector::skip(767), + RowSelector::select(1), + ]), + RowSelection::from(vec![ + RowSelector::select(255), + RowSelector::skip(1), + RowSelector::select(767), + RowSelector::skip(1), + ]), + RowSelection::from(vec![ + RowSelector::skip(254), + RowSelector::select(1), + RowSelector::select(1), + RowSelector::skip(767), + RowSelector::select(1), + ]), + ]; + + for selection in selections { + let expected = selection.row_count(); + // Read data + let mut reader = ParquetRecordBatchStreamBuilder::new_with_options( + tokio::fs::File::from_std(file.try_clone().unwrap()), + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required), + ) + .await + .unwrap(); + + reader = reader.with_row_selection(selection); + + let mut stream = reader.build().unwrap(); + + let mut total_rows = 0; + while let Some(rb) = stream.next().await { + let rb = rb.unwrap(); + total_rows += rb.num_rows(); + } + assert_eq!(total_rows, expected); + } + } + + #[tokio::test] + async fn empty_offset_index_doesnt_panic_in_read_row_group() { + use tokio::fs::File; + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_plain.parquet"); + let mut file = File::open(&path).await.unwrap(); + let file_size = file.metadata().await.unwrap().len(); + let mut metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .load_and_finish(&mut file, file_size) + .await + .unwrap(); + + metadata.set_offset_index(Some(vec![])); + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let arrow_reader_metadata = ArrowReaderMetadata::try_new(metadata.into(), options).unwrap(); + let reader = + ParquetRecordBatchStreamBuilder::new_with_metadata(file, arrow_reader_metadata) + .build() + .unwrap(); + + let result = reader.try_collect::>().await.unwrap(); + assert_eq!(result.len(), 1); + } + + #[tokio::test] + async fn non_empty_offset_index_doesnt_panic_in_read_row_group() { + use tokio::fs::File; + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_tiny_pages.parquet"); + let mut file = File::open(&path).await.unwrap(); + let file_size = file.metadata().await.unwrap().len(); + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .load_and_finish(&mut file, file_size) + .await + .unwrap(); + + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let arrow_reader_metadata = ArrowReaderMetadata::try_new(metadata.into(), options).unwrap(); + let reader = + ParquetRecordBatchStreamBuilder::new_with_metadata(file, arrow_reader_metadata) + .build() + .unwrap(); + + let result = reader.try_collect::>().await.unwrap(); + assert_eq!(result.len(), 8); + } + + #[tokio::test] + async fn empty_offset_index_doesnt_panic_in_column_chunks() { + use tempfile::TempDir; + use tokio::fs::File; + fn write_metadata_to_local_file( + metadata: ParquetMetaData, + file: impl AsRef, + ) { + use crate::file::metadata::ParquetMetaDataWriter; + use std::fs::File; + let file = File::create(file).unwrap(); + ParquetMetaDataWriter::new(file, &metadata) + .finish() + .unwrap() + } + + fn read_metadata_from_local_file(file: impl AsRef) -> ParquetMetaData { + use std::fs::File; + let file = File::open(file).unwrap(); + ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .parse_and_finish(&file) + .unwrap() + } + + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_plain.parquet"); + let mut file = File::open(&path).await.unwrap(); + let file_size = file.metadata().await.unwrap().len(); + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .load_and_finish(&mut file, file_size) + .await + .unwrap(); + + let tempdir = TempDir::new().unwrap(); + let metadata_path = tempdir.path().join("thrift_metadata.dat"); + write_metadata_to_local_file(metadata, &metadata_path); + let metadata = read_metadata_from_local_file(&metadata_path); + + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let arrow_reader_metadata = ArrowReaderMetadata::try_new(metadata.into(), options).unwrap(); + let reader = + ParquetRecordBatchStreamBuilder::new_with_metadata(file, arrow_reader_metadata) + .build() + .unwrap(); + + // Panics here + let result = reader.try_collect::>().await.unwrap(); + assert_eq!(result.len(), 1); + } + + #[tokio::test] + async fn test_cached_array_reader_sparse_offset_error() { + use futures::TryStreamExt; + + use crate::arrow::arrow_reader::{ArrowPredicateFn, RowFilter, RowSelection, RowSelector}; + use arrow_array::{BooleanArray, RecordBatch}; + + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet"); + let data = Bytes::from(std::fs::read(path).unwrap()); + + let async_reader = TestReader::new(data); + + // Enable page index so the fetch logic loads only required pages + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let builder = ParquetRecordBatchStreamBuilder::new_with_options(async_reader, options) + .await + .unwrap(); + + // Skip the first 22 rows (entire first Parquet page) and then select the + // next 3 rows (22, 23, 24). This means the fetch step will not include + // the first page starting at file offset 0. + let selection = RowSelection::from(vec![RowSelector::skip(22), RowSelector::select(3)]); + + // Trivial predicate on column 0 that always returns `true`. Using the + // same column in both predicate and projection activates the caching + // layer (Producer/Consumer pattern). + let parquet_schema = builder.parquet_schema(); + let proj = ProjectionMask::leaves(parquet_schema, vec![0]); + let always_true = ArrowPredicateFn::new(proj.clone(), |batch: RecordBatch| { + Ok(BooleanArray::from(vec![true; batch.num_rows()])) + }); + let filter = RowFilter::new(vec![Box::new(always_true)]); + + // Build the stream with batch size 8 so the cache reads whole batches + // that straddle the requested row range (rows 0-7, 8-15, 16-23, …). + let stream = builder + .with_batch_size(8) + .with_projection(proj) + .with_row_selection(selection) + .with_row_filter(filter) + .build() + .unwrap(); + + // Collecting the stream should fail with the sparse column chunk offset + // error we want to reproduce. + let _result: Vec<_> = stream.try_collect().await.unwrap(); + } + + #[tokio::test] + async fn test_predicate_cache_disabled() { + let k = Int32Array::from_iter_values(0..10); + let data = RecordBatch::try_from_iter([("k", Arc::new(k) as ArrayRef)]).unwrap(); + + let mut buf = Vec::new(); + // both the page row limit and batch size are set to 1 to create one page per row + let props = WriterProperties::builder() + .set_data_page_row_count_limit(1) + .set_write_batch_size(1) + .set_max_row_group_row_count(Some(10)) + .set_write_page_header_statistics(true) + .build(); + let mut writer = ArrowWriter::try_new(&mut buf, data.schema(), Some(props)).unwrap(); + writer.write(&data).unwrap(); + writer.close().unwrap(); + + let data = Bytes::from(buf); + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .parse_and_finish(&data) + .unwrap(); + let parquet_schema = metadata.file_metadata().schema_descr_ptr(); + + // the filter is not clone-able, so we use a lambda to simplify + let build_filter = || { + let scalar = Int32Array::from_iter_values([5]); + let predicate = ArrowPredicateFn::new( + ProjectionMask::leaves(&parquet_schema, vec![0]), + move |batch| eq(batch.column(0), &Scalar::new(&scalar)), + ); + RowFilter::new(vec![Box::new(predicate)]) + }; + + // select only one of the pages + let selection = RowSelection::from(vec![RowSelector::skip(5), RowSelector::select(1)]); + + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let reader_metadata = ArrowReaderMetadata::try_new(metadata.into(), options).unwrap(); + + // using the predicate cache (default) + let reader_with_cache = TestReader::new(data.clone()); + let requests_with_cache = reader_with_cache.requests.clone(); + let stream = ParquetRecordBatchStreamBuilder::new_with_metadata( + reader_with_cache, + reader_metadata.clone(), + ) + .with_batch_size(1000) + .with_row_selection(selection.clone()) + .with_row_filter(build_filter()) + .build() + .unwrap(); + let batches_with_cache: Vec<_> = stream.try_collect().await.unwrap(); + + // disabling the predicate cache + let reader_without_cache = TestReader::new(data); + let requests_without_cache = reader_without_cache.requests.clone(); + let stream = ParquetRecordBatchStreamBuilder::new_with_metadata( + reader_without_cache, + reader_metadata, + ) + .with_batch_size(1000) + .with_row_selection(selection) + .with_row_filter(build_filter()) + .with_max_predicate_cache_size(0) // disabling it by setting the limit to 0 + .build() + .unwrap(); + let batches_without_cache: Vec<_> = stream.try_collect().await.unwrap(); + + assert_eq!(batches_with_cache, batches_without_cache); + + let requests_with_cache = requests_with_cache.lock().unwrap(); + let requests_without_cache = requests_without_cache.lock().unwrap(); + + // less requests will be made without the predicate cache + assert_eq!(requests_with_cache.len(), 11); + assert_eq!(requests_without_cache.len(), 2); + + // less bytes will be retrieved without the predicate cache + assert_eq!( + requests_with_cache.iter().map(|r| r.len()).sum::(), + 433 + ); + assert_eq!( + requests_without_cache + .iter() + .map(|r| r.len()) + .sum::(), + 92 + ); + } + + #[test] + fn test_row_numbers_with_multiple_row_groups() { + test_row_numbers_with_multiple_row_groups_helper( + false, + |path, selection, _row_filter, batch_size| { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("Could not create runtime"); + runtime.block_on(async move { + let file = tokio::fs::File::open(path).await.unwrap(); + let row_number_field = Arc::new( + Field::new("row_number", DataType::Int64, false) + .with_extension_type(RowNumber), + ); + let options = ArrowReaderOptions::new() + .with_virtual_columns(vec![row_number_field]) + .unwrap(); + let reader = ParquetRecordBatchStreamBuilder::new_with_options(file, options) + .await + .unwrap() + .with_row_selection(selection) + .with_batch_size(batch_size) + .build() + .expect("Could not create reader"); + reader.try_collect::>().await.unwrap() + }) + }, + ); + } + + #[test] + fn test_row_numbers_with_multiple_row_groups_and_filter() { + test_row_numbers_with_multiple_row_groups_helper( + true, + |path, selection, row_filter, batch_size| { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("Could not create runtime"); + runtime.block_on(async move { + let file = tokio::fs::File::open(path).await.unwrap(); + let row_number_field = Arc::new( + Field::new("row_number", DataType::Int64, false) + .with_extension_type(RowNumber), + ); + let options = ArrowReaderOptions::new() + .with_virtual_columns(vec![row_number_field]) + .unwrap(); + let reader = ParquetRecordBatchStreamBuilder::new_with_options(file, options) + .await + .unwrap() + .with_row_selection(selection) + .with_row_filter(row_filter.expect("No row filter")) + .with_batch_size(batch_size) + .build() + .expect("Could not create reader"); + reader.try_collect::>().await.unwrap() + }) + }, + ); + } + + #[tokio::test] + async fn test_nested_lists() -> Result<()> { + // Test case for https://github.com/apache/arrow-rs/issues/8657 + let list_inner_field = Arc::new(Field::new("item", DataType::Float32, true)); + let table_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("vector", DataType::List(list_inner_field.clone()), true), + ])); + + let mut list_builder = + ListBuilder::new(Float32Builder::new()).with_field(list_inner_field.clone()); + list_builder.values().append_slice(&[10.0, 10.0, 10.0]); + list_builder.append(true); + list_builder.values().append_slice(&[20.0, 20.0, 20.0]); + list_builder.append(true); + list_builder.values().append_slice(&[30.0, 30.0, 30.0]); + list_builder.append(true); + list_builder.values().append_slice(&[40.0, 40.0, 40.0]); + list_builder.append(true); + let list_array = list_builder.finish(); + + let data = vec![RecordBatch::try_new( + table_schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4])), + Arc::new(list_array), + ], + )?]; + + let mut buffer = Vec::new(); + let mut writer = AsyncArrowWriter::try_new(&mut buffer, table_schema, None)?; + + for batch in data { + writer.write(&batch).await?; + } + + writer.close().await?; + + let reader = TestReader::new(Bytes::from(buffer)); + let builder = ParquetRecordBatchStreamBuilder::new(reader).await?; + + let predicate = ArrowPredicateFn::new(ProjectionMask::all(), |batch| { + Ok(BooleanArray::from(vec![true; batch.num_rows()])) + }); + + let projection_mask = ProjectionMask::all(); + + let mut stream = builder + .with_row_filter(RowFilter::new(vec![Box::new(predicate)])) + .with_projection(projection_mask) + .build()?; + + while let Some(batch) = stream.next().await { + let _ = batch.unwrap(); // ensure there is no panic + } + + Ok(()) + } +} diff --git a/vendor/parquet/src/arrow/async_reader/spawn.rs b/vendor/parquet/src/arrow/async_reader/spawn.rs new file mode 100644 index 0000000..927e11e --- /dev/null +++ b/vendor/parquet/src/arrow/async_reader/spawn.rs @@ -0,0 +1,227 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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::future::Future; +use std::ops::Range; +use std::sync::Arc; + +use bytes::Bytes; +use futures::future::BoxFuture; +use futures::{FutureExt, TryFutureExt}; +use tokio::runtime::Handle; + +use crate::arrow::arrow_reader::ArrowReaderOptions; +use crate::arrow::async_reader::{AsyncFileReader, MetadataSuffixFetch}; +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::ParquetMetaData; + +/// An [`AsyncFileReader`] that performs I/O on a separate tokio runtime. +/// +/// Tokio is a cooperative scheduler, and relies on tasks yielding in a timely +/// manner to service IO. Therefore, running IO and CPU-bound tasks, such as +/// parquet decoding, on the same tokio runtime can lead to degraded +/// throughput, dropped connections and other issues. For more information see +/// [here]. +/// +/// This wrapper spawns each operation of the inner reader onto the provided +/// runtime [`Handle`], so that the runtime driving the parquet decoding does +/// not also drive the I/O. +/// +/// Note that [`Self::get_metadata`] spawns the entire metadata load, so the +/// footer is also decoded on the provided runtime. +/// +/// The inner reader must be [`Clone`] (typically an `Arc`'d handle to some +/// shared resource) as each spawned task requires a `'static` copy of it. +/// +/// [here]: https://www.influxdata.com/blog/using-rustlangs-async-tokio-runtime-for-cpu-bound-tasks/ +#[derive(Clone, Debug)] +pub struct SpawnedReader { + inner: R, + handle: Handle, +} + +impl SpawnedReader { + /// Creates a new [`SpawnedReader`] that performs the I/O of `inner` on `handle` + pub fn new(inner: R, handle: Handle) -> Self { + Self { inner, handle } + } + + /// Returns the inner reader + pub fn into_inner(self) -> R { + self.inner + } +} + +/// Spawns `fut` on `handle`, propagating panics and mapping task cancellation +/// to [`ParquetError::External`] +fn spawn( + handle: &Handle, + fut: impl Future> + Send + 'static, +) -> BoxFuture<'static, Result> +where + T: Send + 'static, +{ + handle + .spawn(fut) + .map_ok_or_else( + |e| match e.try_into_panic() { + Err(e) => Err(ParquetError::External(Box::new(e))), + Ok(p) => std::panic::resume_unwind(p), + }, + |res| res, + ) + .boxed() +} + +impl AsyncFileReader for SpawnedReader +where + R: AsyncFileReader + Clone + Send + 'static, +{ + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + let mut inner = self.inner.clone(); + spawn(&self.handle, async move { inner.get_bytes(range).await }) + } + + fn get_byte_ranges(&mut self, ranges: Vec>) -> BoxFuture<'_, Result>> { + let mut inner = self.inner.clone(); + spawn( + &self.handle, + async move { inner.get_byte_ranges(ranges).await }, + ) + } + + fn get_metadata<'a>( + &'a mut self, + options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, Result>> { + let mut inner = self.inner.clone(); + let options = options.cloned(); + spawn(&self.handle, async move { + inner.get_metadata(options.as_ref()).await + }) + } +} + +impl MetadataSuffixFetch for &mut SpawnedReader +where + R: AsyncFileReader + Clone + Send + 'static, + for<'a> &'a mut R: MetadataSuffixFetch, +{ + fn fetch_suffix(&mut self, suffix: usize) -> BoxFuture<'_, Result> { + let mut inner = self.inner.clone(); + spawn(&self.handle, async move { + (&mut inner).fetch_suffix(suffix).await + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::ParquetRecordBatchStreamBuilder; + use crate::file::metadata::ParquetMetaDataReader; + use futures::TryStreamExt; + use std::thread::ThreadId; + + /// An in-memory [`AsyncFileReader`] that records the thread each request ran on + #[derive(Clone)] + struct InMemoryReader { + data: Bytes, + threads: Arc>>, + } + + impl InMemoryReader { + fn new(data: Bytes) -> Self { + Self { + data, + threads: Default::default(), + } + } + } + + impl AsyncFileReader for InMemoryReader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + self.threads + .lock() + .unwrap() + .push(std::thread::current().id()); + let data = self.data.slice(range.start as usize..range.end as usize); + futures::future::ready(Ok(data)).boxed() + } + + fn get_metadata<'a>( + &'a mut self, + options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, Result>> { + self.threads + .lock() + .unwrap() + .push(std::thread::current().id()); + let metadata = ParquetMetaDataReader::new() + .with_arrow_reader_options(options) + .parse_and_finish(&self.data); + futures::future::ready(metadata.map(Arc::new)).boxed() + } + } + + #[tokio::test] + async fn test_spawned_reader() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_plain.parquet"); + let data = Bytes::from(std::fs::read(path).unwrap()); + + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .build() + .unwrap(); + + let inner = InMemoryReader::new(data); + let threads = inner.threads.clone(); + let reader = SpawnedReader::new(inner, rt.handle().clone()); + + let builder = ParquetRecordBatchStreamBuilder::new(reader).await.unwrap(); + let batches: Vec<_> = builder.build().unwrap().try_collect().await.unwrap(); + + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 8); + + // All I/O must have run on the spawned runtime, not the current one + let current_id = std::thread::current().id(); + let threads = threads.lock().unwrap(); + assert!(!threads.is_empty()); + assert!(threads.iter().all(|id| *id != current_id)); + + // Runtimes have to be dropped in blocking contexts + tokio::runtime::Handle::current().spawn_blocking(move || drop(rt)); + } + + #[tokio::test] + async fn test_spawned_reader_fails_on_shutdown_runtime() { + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .build() + .unwrap(); + + let inner = InMemoryReader::new(Bytes::from_static(b"PAR1")); + let mut reader = SpawnedReader::new(inner, rt.handle().clone()); + + rt.shutdown_background(); + + let err = reader.get_bytes(0..1).await.unwrap_err().to_string(); + assert!(err.contains("was cancelled"), "{err}"); + } +} diff --git a/vendor/parquet/src/arrow/async_reader/store.rs b/vendor/parquet/src/arrow/async_reader/store.rs new file mode 100644 index 0000000..fe9cb65 --- /dev/null +++ b/vendor/parquet/src/arrow/async_reader/store.rs @@ -0,0 +1,522 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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::{ops::Range, sync::Arc}; + +use crate::arrow::arrow_reader::ArrowReaderOptions; +use crate::arrow::async_reader::{AsyncFileReader, MetadataSuffixFetch}; +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader}; +use bytes::Bytes; +use futures::{FutureExt, TryFutureExt, future::BoxFuture}; +use object_store::ObjectStoreExt; +use object_store::{GetOptions, GetRange}; +use object_store::{ObjectStore, path::Path}; +use tokio::runtime::Handle; +/// Reads Parquet files in object storage using [`ObjectStore`]. +/// +/// ```no_run +/// # use std::io::stdout; +/// # use std::sync::Arc; +/// # use object_store::azure::MicrosoftAzureBuilder; +/// # use object_store::{ObjectStore, ObjectStoreExt}; +/// # use object_store::path::Path; +/// # use parquet::arrow::async_reader::ParquetObjectReader; +/// # use parquet::arrow::ParquetRecordBatchStreamBuilder; +/// # use parquet::schema::printer::print_parquet_metadata; +/// # async fn run() { +/// // Populate configuration from environment +/// let storage_container = Arc::new(MicrosoftAzureBuilder::from_env().build().unwrap()); +/// let location = Path::from("path/to/blob.parquet"); +/// let meta = storage_container.head(&location).await.unwrap(); +/// println!("Found Blob with {}B at {}", meta.size, meta.location); +/// +/// // Show Parquet metadata +/// let reader = ParquetObjectReader::new(storage_container, meta.location).with_file_size(meta.size); +/// let builder = ParquetRecordBatchStreamBuilder::new(reader).await.unwrap(); +/// print_parquet_metadata(&mut stdout(), builder.metadata()); +/// # } +/// ``` +#[deprecated( + since = "59.2.0", + note = "Implement `AsyncFileReader` directly instead; see the example on the `AsyncFileReader` trait documentation and `parquet/examples/object_store.rs`. Use `SpawnedReader` to perform I/O on a dedicated runtime. See https://github.com/apache/arrow-rs/issues/10308" +)] +#[derive(Clone, Debug)] +pub struct ParquetObjectReader { + store: Arc, + path: Path, + file_size: Option, + metadata_size_hint: Option, + preload_column_index: bool, + preload_offset_index: bool, + runtime: Option, +} + +#[allow(deprecated)] +impl ParquetObjectReader { + /// Creates a new [`ParquetObjectReader`] for the provided [`ObjectStore`] and [`Path`]. + pub fn new(store: Arc, path: Path) -> Self { + Self { + store, + path, + file_size: None, + metadata_size_hint: None, + preload_column_index: false, + preload_offset_index: false, + runtime: None, + } + } + + /// Provide a hint as to the size of the parquet file's footer, + /// see [`ParquetMetaDataReader::with_prefetch_hint`] + pub fn with_footer_size_hint(self, hint: usize) -> Self { + Self { + metadata_size_hint: Some(hint), + ..self + } + } + + /// Provide the byte size of this file. + /// + /// If provided, the file size will ensure that only bounded range requests are used. If file + /// size is not provided, the reader will use suffix range requests to fetch the metadata. + /// + /// Providing this size up front is an important optimization to avoid extra calls when the + /// underlying store does not support suffix range requests. + /// + /// The file size can be obtained using [`ObjectStore::list`] or [`ObjectStoreExt::head`]. + pub fn with_file_size(self, file_size: u64) -> Self { + Self { + file_size: Some(file_size), + ..self + } + } + + /// Whether to load the Column Index as part of [`Self::get_metadata`] + /// + /// Note: This setting may be overridden by [`ArrowReaderOptions`] `page_index_policy`. + /// If `page_index_policy` is `Optional` or `Required`, it will take precedence + /// over this preload flag. When it is `Skip` (default), this flag is used. + pub fn with_preload_column_index(self, preload_column_index: bool) -> Self { + Self { + preload_column_index, + ..self + } + } + + /// Whether to load the Offset Index as part of [`Self::get_metadata`] + /// + /// Note: This setting may be overridden by [`ArrowReaderOptions`] `page_index_policy`. + /// If `page_index_policy` is `Optional` or `Required`, it will take precedence + /// over this preload flag. When it is `Skip` (default), this flag is used. + pub fn with_preload_offset_index(self, preload_offset_index: bool) -> Self { + Self { + preload_offset_index, + ..self + } + } + + /// Perform IO on the provided tokio runtime + /// + /// Tokio is a cooperative scheduler, and relies on tasks yielding in a timely manner + /// to service IO. Therefore, running IO and CPU-bound tasks, such as parquet decoding, + /// on the same tokio runtime can lead to degraded throughput, dropped connections and + /// other issues. For more information see [here]. + /// + /// [here]: https://www.influxdata.com/blog/using-rustlangs-async-tokio-runtime-for-cpu-bound-tasks/ + #[deprecated( + since = "59.2.0", + note = "Wrap the reader in a `SpawnedReader` instead, e.g. `SpawnedReader::new(reader, handle)`. See https://github.com/apache/arrow-rs/issues/10308" + )] + pub fn with_runtime(self, handle: Handle) -> Self { + Self { + runtime: Some(handle), + ..self + } + } + + fn spawn(&self, f: F) -> BoxFuture<'_, Result> + where + F: for<'a> FnOnce(&'a Arc, &'a Path) -> BoxFuture<'a, Result> + + Send + + 'static, + O: Send + 'static, + E: Into + Send + 'static, + { + match &self.runtime { + Some(handle) => { + let path = self.path.clone(); + let store = Arc::clone(&self.store); + handle + .spawn(async move { f(&store, &path).await }) + .map_ok_or_else( + |e| match e.try_into_panic() { + Err(e) => Err(ParquetError::External(Box::new(e))), + Ok(p) => std::panic::resume_unwind(p), + }, + |res| res.map_err(|e| e.into()), + ) + .boxed() + } + None => f(&self.store, &self.path).map_err(|e| e.into()).boxed(), + } + } +} + +#[allow(deprecated)] +impl MetadataSuffixFetch for &mut ParquetObjectReader { + fn fetch_suffix(&mut self, suffix: usize) -> BoxFuture<'_, Result> { + let options = GetOptions { + range: Some(GetRange::Suffix(suffix as u64)), + ..Default::default() + }; + self.spawn(|store, path| { + async move { + let resp = store.get_opts(path, options).await?; + Ok::<_, ParquetError>(resp.bytes().await?) + } + .boxed() + }) + } +} + +#[allow(deprecated)] +impl AsyncFileReader for ParquetObjectReader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + self.spawn(|store, path| store.get_range(path, range).boxed()) + } + + fn get_byte_ranges(&mut self, ranges: Vec>) -> BoxFuture<'_, Result>> + where + Self: Send, + { + self.spawn(|store, path| async move { store.get_ranges(path, &ranges).await }.boxed()) + } + + // This method doesn't directly call `self.spawn` because all of the IO that is done down the + // line due to this method call is done through `self.get_bytes` and/or `self.get_byte_ranges`. + // When `self` is passed into `ParquetMetaDataReader::load_and_finish`, it treats it as + // an `impl MetadataFetch` and calls those methods to get data from it. Due to `Self`'s impl of + // `AsyncFileReader`, the calls to `MetadataFetch::fetch` are just delegated to + // `Self::get_bytes`. + fn get_metadata<'a>( + &'a mut self, + options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let metadata_opts = options.map(|o| o.metadata_options().clone()); + let mut metadata = ParquetMetaDataReader::new() + .with_metadata_options(metadata_opts) + .with_column_index_policy(PageIndexPolicy::from(self.preload_column_index)) + .with_offset_index_policy(PageIndexPolicy::from(self.preload_offset_index)) + .with_prefetch_hint(self.metadata_size_hint); + + #[cfg(feature = "encryption")] + if let Some(options) = options { + metadata = metadata.with_decryption_properties( + options.file_decryption_properties.as_ref().map(Arc::clone), + ); + } + + // Override page index policies from ArrowReaderOptions if specified and not Skip. + // When page_index_policy is Skip (default), use the reader's preload flags. + // When page_index_policy is Optional or Required, override the preload flags + // to ensure the specified policy takes precedence. + if let Some(options) = options { + if options.column_index_policy() != PageIndexPolicy::Skip + || options.offset_index_policy() != PageIndexPolicy::Skip + { + metadata = metadata + .with_column_index_policy(options.column_index_policy()) + .with_offset_index_policy(options.offset_index_policy()); + } + } + + let metadata = if let Some(file_size) = self.file_size { + metadata.load_and_finish(self, file_size).await? + } else { + metadata.load_via_suffix_and_finish(self).await? + }; + + Ok(Arc::new(metadata)) + }) + } +} + +#[cfg(test)] +#[allow(deprecated)] +mod tests { + use crate::arrow::async_reader::ArrowReaderOptions; + use crate::file::metadata::PageIndexPolicy; + use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }; + + use futures::TryStreamExt; + + use crate::arrow::ParquetRecordBatchStreamBuilder; + use crate::arrow::async_reader::{AsyncFileReader, ParquetObjectReader}; + use crate::errors::ParquetError; + use arrow::util::test_util::parquet_test_data; + use futures::FutureExt; + use object_store::local::LocalFileSystem; + use object_store::path::Path; + use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt}; + + async fn get_meta_store() -> (ObjectMeta, Arc) { + let res = parquet_test_data(); + let store = LocalFileSystem::new_with_prefix(res).unwrap(); + + let meta = store + .head(&Path::from("alltypes_plain.parquet")) + .await + .unwrap(); + + (meta, Arc::new(store) as Arc) + } + + async fn get_meta_store_with_page_index() -> (ObjectMeta, Arc) { + let res = parquet_test_data(); + let store = LocalFileSystem::new_with_prefix(res).unwrap(); + + let meta = store + .head(&Path::from("alltypes_tiny_pages_plain.parquet")) + .await + .unwrap(); + + (meta, Arc::new(store) as Arc) + } + + #[tokio::test] + async fn test_simple() { + let (meta, store) = get_meta_store().await; + let object_reader = + ParquetObjectReader::new(store, meta.location).with_file_size(meta.size); + + let builder = ParquetRecordBatchStreamBuilder::new(object_reader) + .await + .unwrap(); + let batches: Vec<_> = builder.build().unwrap().try_collect().await.unwrap(); + + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 8); + } + + #[tokio::test] + async fn test_simple_without_file_length() { + let (meta, store) = get_meta_store().await; + let object_reader = ParquetObjectReader::new(store, meta.location); + + let builder = ParquetRecordBatchStreamBuilder::new(object_reader) + .await + .unwrap(); + let batches: Vec<_> = builder.build().unwrap().try_collect().await.unwrap(); + + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 8); + } + + #[tokio::test] + async fn test_not_found() { + let (mut meta, store) = get_meta_store().await; + meta.location = Path::from("I don't exist.parquet"); + + let object_reader = + ParquetObjectReader::new(store, meta.location).with_file_size(meta.size); + // Cannot use unwrap_err as ParquetRecordBatchStreamBuilder: !Debug + match ParquetRecordBatchStreamBuilder::new(object_reader).await { + Ok(_) => panic!("expected failure"), + Err(e) => { + let err = e.to_string(); + assert!(err.contains("I don't exist.parquet not found:"), "{err}",); + } + } + } + + #[tokio::test] + async fn test_runtime_is_used() { + let num_actions = Arc::new(AtomicUsize::new(0)); + + let (a1, a2) = (num_actions.clone(), num_actions.clone()); + let rt = tokio::runtime::Builder::new_multi_thread() + .on_thread_park(move || { + a1.fetch_add(1, Ordering::Relaxed); + }) + .on_thread_unpark(move || { + a2.fetch_add(1, Ordering::Relaxed); + }) + .build() + .unwrap(); + + let (meta, store) = get_meta_store().await; + + let initial_actions = num_actions.load(Ordering::Relaxed); + + let reader = ParquetObjectReader::new(store, meta.location) + .with_file_size(meta.size) + .with_runtime(rt.handle().clone()); + + let builder = ParquetRecordBatchStreamBuilder::new(reader).await.unwrap(); + let batches: Vec<_> = builder.build().unwrap().try_collect().await.unwrap(); + + // Just copied these assert_eqs from the `test_simple` above + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 8); + + assert!(num_actions.load(Ordering::Relaxed) - initial_actions > 0); + + // Runtimes have to be dropped in blocking contexts, so we need to move this one to a new + // blocking thread to drop it. + tokio::runtime::Handle::current().spawn_blocking(move || drop(rt)); + } + + /// Unit test that `ParquetObjectReader::spawn`spawns on the provided runtime + #[tokio::test] + async fn test_runtime_thread_id_different() { + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .build() + .unwrap(); + + let (meta, store) = get_meta_store().await; + + let reader = ParquetObjectReader::new(store, meta.location) + .with_file_size(meta.size) + .with_runtime(rt.handle().clone()); + + let current_id = std::thread::current().id(); + + let other_id = reader + .spawn(|_, _| async move { Ok::<_, ParquetError>(std::thread::current().id()) }.boxed()) + .await + .unwrap(); + + assert_ne!(current_id, other_id); + + tokio::runtime::Handle::current().spawn_blocking(move || drop(rt)); + } + + #[tokio::test] + async fn io_fails_on_shutdown_runtime() { + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .build() + .unwrap(); + + let (meta, store) = get_meta_store().await; + + let mut reader = ParquetObjectReader::new(store, meta.location) + .with_file_size(meta.size) + .with_runtime(rt.handle().clone()); + + rt.shutdown_background(); + + let err = reader.get_bytes(0..1).await.unwrap_err().to_string(); + + assert!(err.to_string().contains("was cancelled")); + } + + #[tokio::test] + async fn test_page_index_policy_skip_uses_preload_true() { + let (meta, store) = get_meta_store_with_page_index().await; + + // Create reader with preload flags set to true + let mut reader = ParquetObjectReader::new(store.clone(), meta.location.clone()) + .with_file_size(meta.size) + .with_preload_column_index(true) + .with_preload_offset_index(true); + + // Create options with page_index_policy set to Skip (default) + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Skip); + + // Get metadata - Skip means use reader's preload flags (true) + let metadata = reader.get_metadata(Some(&options)).await.unwrap(); + + // With preload=true, indexes should be loaded since the test file has them + assert!(metadata.column_index().is_some()); + } + + #[tokio::test] + async fn test_page_index_policy_optional_overrides_preload_false() { + let (meta, store) = get_meta_store_with_page_index().await; + + // Create reader with preload flags set to false + let mut reader = ParquetObjectReader::new(store.clone(), meta.location.clone()) + .with_file_size(meta.size) + .with_preload_column_index(false) + .with_preload_offset_index(false); + + // Create options with page_index_policy set to Optional + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Optional); + + // Get metadata - Optional overrides preload flags and attempts to load indexes + let metadata = reader.get_metadata(Some(&options)).await.unwrap(); + + // With Optional policy, it will TRY to load indexes but won't fail if they don't exist + // The test file has page indexes, so they will be some + assert!(metadata.column_index().is_some()); + } + + #[tokio::test] + async fn test_page_index_policy_optional_vs_skip() { + let (meta, store) = get_meta_store_with_page_index().await; + + // Test 1: preload=false + Skip policy -> uses preload flags (false) + let mut reader1 = ParquetObjectReader::new(store.clone(), meta.location.clone()) + .with_file_size(meta.size) + .with_preload_column_index(false) + .with_preload_offset_index(false); + + let options1 = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Skip); + let metadata1 = reader1.get_metadata(Some(&options1)).await.unwrap(); + + // Test 2: preload=false + Optional policy -> overrides to try loading + let mut reader2 = ParquetObjectReader::new(store.clone(), meta.location.clone()) + .with_file_size(meta.size) + .with_preload_column_index(false) + .with_preload_offset_index(false); + + let options2 = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Optional); + let metadata2 = reader2.get_metadata(Some(&options2)).await.unwrap(); + + // Both should succeed (no panic/error) + // metadata1 (Skip) uses preload=false -> Skip policy + // metadata2 (Optional) overrides preload=false -> Optional policy + assert!(metadata1.column_index().is_none()); + assert!(metadata2.column_index().is_some()); + } + + #[tokio::test] + async fn test_page_index_policy_no_options_uses_preload() { + let (meta, store) = get_meta_store_with_page_index().await; + + // Create reader with preload flags set to true + let mut reader = ParquetObjectReader::new(store, meta.location) + .with_file_size(meta.size) + .with_preload_column_index(true) + .with_preload_offset_index(true); + + // Get metadata without options - should use reader's preload flags + let metadata = reader.get_metadata(None).await.unwrap(); + + // With no options provided, preload flags (true) should be respected + // and converted to Optional policy internally (preload=true -> Optional) + // The test file has page indexes, so they will be some + assert!(metadata.column_index().is_some() && metadata.column_index().is_some()); + } +} diff --git a/vendor/parquet/src/arrow/async_writer/mod.rs b/vendor/parquet/src/arrow/async_writer/mod.rs new file mode 100644 index 0000000..d912441 --- /dev/null +++ b/vendor/parquet/src/arrow/async_writer/mod.rs @@ -0,0 +1,474 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! `async` API for writing [`RecordBatch`]es to Parquet files +//! +//! See the [crate-level documentation](crate) for more details. +//! +//! The `async` API for writing [`RecordBatch`]es is +//! similar to the [`sync` API](ArrowWriter), so please +//! read the documentation there before using this API. +//! +//! Here is an example for using [`AsyncArrowWriter`]: +//! +//! ``` +//! # #[tokio::main(flavor="current_thread")] +//! # async fn main() { +//! # +//! # use std::sync::Arc; +//! # use arrow_array::{ArrayRef, Int64Array, RecordBatch, RecordBatchReader}; +//! # use bytes::Bytes; +//! # use parquet::arrow::{AsyncArrowWriter, arrow_reader::ParquetRecordBatchReaderBuilder}; +//! # +//! let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef; +//! let to_write = RecordBatch::try_from_iter([("col", col)]).unwrap(); +//! +//! let mut buffer = Vec::new(); +//! let mut writer = AsyncArrowWriter::try_new(&mut buffer, to_write.schema(), None).unwrap(); +//! writer.write(&to_write).await.unwrap(); +//! writer.close().await.unwrap(); +//! +//! let buffer = Bytes::from(buffer); +//! let mut reader = ParquetRecordBatchReaderBuilder::try_new(buffer.clone()) +//! .unwrap() +//! .build() +//! .unwrap(); +//! let read = reader.next().unwrap().unwrap(); +//! +//! assert_eq!(to_write, read); +//! # } +//! ``` +//! +//! There is a blanket implementation of [`AsyncFileWriter`] for all types that +//! implement [`AsyncWrite`], so writers such as `tokio::fs::File`, or +//! `object_store::buffered::BufWriter` for writing to object storage, can be +//! passed to [`AsyncArrowWriter`] directly. + +#[cfg(feature = "object_store")] +mod store; +#[allow(deprecated)] +#[cfg(feature = "object_store")] +pub use store::*; + +use crate::{ + arrow::ArrowWriter, + arrow::arrow_writer::ArrowWriterOptions, + errors::{ParquetError, Result}, + file::{ + metadata::{KeyValue, ParquetMetaData, RowGroupMetaData}, + properties::WriterProperties, + }, +}; +use arrow_array::RecordBatch; +use arrow_schema::SchemaRef; +use bytes::Bytes; +use futures::FutureExt; +use futures::future::BoxFuture; +use std::mem; +use tokio::io::{AsyncWrite, AsyncWriteExt}; + +/// The asynchronous interface used by [`AsyncArrowWriter`] to write parquet files. +pub trait AsyncFileWriter: Send { + /// Write the provided bytes to the underlying writer + /// + /// The underlying writer CAN decide to buffer the data or write it immediately. + /// This design allows the writer implementer to control the buffering and I/O scheduling. + /// + /// The underlying writer MAY implement retry logic to prevent breaking users write process. + fn write(&mut self, bs: Bytes) -> BoxFuture<'_, Result<()>>; + + /// Flush any buffered data to the underlying writer and finish writing process. + /// + /// After `complete` returns `Ok(())`, caller SHOULD not call write again. + fn complete(&mut self) -> BoxFuture<'_, Result<()>>; +} + +impl AsyncFileWriter for Box { + fn write(&mut self, bs: Bytes) -> BoxFuture<'_, Result<()>> { + self.as_mut().write(bs) + } + + fn complete(&mut self) -> BoxFuture<'_, Result<()>> { + self.as_mut().complete() + } +} + +impl AsyncFileWriter for T { + fn write(&mut self, bs: Bytes) -> BoxFuture<'_, Result<()>> { + async move { + self.write_all(&bs).await?; + Ok(()) + } + .boxed() + } + + fn complete(&mut self) -> BoxFuture<'_, Result<()>> { + async move { + self.flush().await?; + self.shutdown().await?; + Ok(()) + } + .boxed() + } +} + +/// Encodes [`RecordBatch`] to parquet, outputting to an [`AsyncFileWriter`] +/// +/// ## Memory Usage +/// +/// This writer eagerly writes data as soon as possible to the underlying [`AsyncFileWriter`], +/// permitting fine-grained control over buffering and I/O scheduling. However, the columnar +/// nature of parquet forces data for an entire row group to be buffered in memory, before +/// it can be flushed. Depending on the data and the configured row group size, this buffering +/// may be substantial. +/// +/// Memory usage can be limited by calling [`Self::flush`] to flush the in progress row group, +/// although this will likely increase overall file size and reduce query performance. +/// See [ArrowWriter] for more information. +/// +/// ```no_run +/// # use tokio::fs::File; +/// # use arrow_array::RecordBatch; +/// # use parquet::arrow::AsyncArrowWriter; +/// # async fn test() { +/// let mut writer: AsyncArrowWriter = todo!(); +/// let batch: RecordBatch = todo!(); +/// writer.write(&batch).await.unwrap(); +/// // Trigger an early flush if buffered size exceeds 1_000_000 +/// if writer.in_progress_size() > 1_000_000 { +/// writer.flush().await.unwrap() +/// } +/// # } +/// ``` +pub struct AsyncArrowWriter { + /// Underlying sync writer + sync_writer: ArrowWriter>, + + /// Async writer provided by caller + async_writer: W, +} + +impl AsyncArrowWriter { + /// Try to create a new Async Arrow Writer + pub fn try_new( + writer: W, + arrow_schema: SchemaRef, + props: Option, + ) -> Result { + let options = ArrowWriterOptions::new().with_properties(props.unwrap_or_default()); + Self::try_new_with_options(writer, arrow_schema, options) + } + + /// Try to create a new Async Arrow Writer with [`ArrowWriterOptions`] + pub fn try_new_with_options( + writer: W, + arrow_schema: SchemaRef, + options: ArrowWriterOptions, + ) -> Result { + let sync_writer = ArrowWriter::try_new_with_options(Vec::new(), arrow_schema, options)?; + + Ok(Self { + sync_writer, + async_writer: writer, + }) + } + + /// Returns metadata for any flushed row groups + pub fn flushed_row_groups(&self) -> &[RowGroupMetaData] { + self.sync_writer.flushed_row_groups() + } + + /// Estimated memory usage, in bytes, of this `ArrowWriter` + /// + /// See [ArrowWriter::memory_size] for more information. + pub fn memory_size(&self) -> usize { + self.sync_writer.memory_size() + } + + /// Anticipated encoded size of the in progress row group. + /// + /// See [ArrowWriter::memory_size] for more information. + pub fn in_progress_size(&self) -> usize { + self.sync_writer.in_progress_size() + } + + /// Returns the number of rows buffered in the in progress row group + pub fn in_progress_rows(&self) -> usize { + self.sync_writer.in_progress_rows() + } + + /// Returns the number of bytes written by this instance + pub fn bytes_written(&self) -> usize { + self.sync_writer.bytes_written() + } + + /// Enqueues the provided `RecordBatch` to be written + /// + /// After every sync write by the inner [ArrowWriter], the inner buffer will be + /// checked and flush if at least half full + pub async fn write(&mut self, batch: &RecordBatch) -> Result<()> { + let before = self.sync_writer.flushed_row_groups().len(); + self.sync_writer.write(batch)?; + if before != self.sync_writer.flushed_row_groups().len() { + self.do_write().await?; + } + Ok(()) + } + + /// Flushes all buffered rows into a new row group + pub async fn flush(&mut self) -> Result<()> { + self.sync_writer.flush()?; + self.do_write().await?; + + Ok(()) + } + + /// Append [`KeyValue`] metadata in addition to those in [`WriterProperties`] + /// + /// This method allows to append metadata after [`RecordBatch`]es are written. + pub fn append_key_value_metadata(&mut self, kv_metadata: KeyValue) { + self.sync_writer.append_key_value_metadata(kv_metadata); + } + + /// Close and finalize the writer. + /// + /// All the data in the inner buffer will be force flushed. + /// + /// Unlike [`Self::close`] this does not consume self + /// + /// Attempting to write after calling finish will result in an error + pub async fn finish(&mut self) -> Result { + let metadata = self.sync_writer.finish()?; + + // Force to flush the remaining data. + self.do_write().await?; + self.async_writer.complete().await?; + + Ok(metadata) + } + + /// Close and finalize the writer. + /// + /// All the data in the inner buffer will be force flushed. + pub async fn close(mut self) -> Result { + self.finish().await + } + + /// Consumes the [`AsyncArrowWriter`] and returns the underlying [`AsyncFileWriter`] + /// + /// # Notes + /// + /// This method does **not** flush or finalize the writer, so buffered data + /// will be lost if you have not called [`Self::finish`]. + pub fn into_inner(self) -> W { + self.async_writer + } + + /// Flush the data written by `sync_writer` into the `async_writer` + /// + /// # Notes + /// + /// This method will take the inner buffer from the `sync_writer` and write it into the + /// async writer. After the write, the inner buffer will be empty. + async fn do_write(&mut self) -> Result<()> { + let buffer = mem::take(self.sync_writer.inner_mut()); + + self.async_writer + .write(Bytes::from(buffer)) + .await + .map_err(|e| ParquetError::External(Box::new(e)))?; + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use crate::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow_array::{ArrayRef, BinaryArray, Int32Array, Int64Array, RecordBatchReader}; + use bytes::Bytes; + use std::sync::Arc; + + use super::*; + + fn get_test_reader() -> ParquetRecordBatchReader { + let testdata = arrow::util::test_util::parquet_test_data(); + // This test file is large enough to generate multiple row groups. + let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet"); + let original_data = Bytes::from(std::fs::read(path).unwrap()); + ParquetRecordBatchReaderBuilder::try_new(original_data) + .unwrap() + .build() + .unwrap() + } + + #[tokio::test] + async fn test_async_writer() { + let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef; + let to_write = RecordBatch::try_from_iter([("col", col)]).unwrap(); + + let mut buffer = Vec::new(); + let mut writer = AsyncArrowWriter::try_new(&mut buffer, to_write.schema(), None).unwrap(); + writer.write(&to_write).await.unwrap(); + writer.close().await.unwrap(); + + let buffer = Bytes::from(buffer); + let mut reader = ParquetRecordBatchReaderBuilder::try_new(buffer) + .unwrap() + .build() + .unwrap(); + let read = reader.next().unwrap().unwrap(); + + assert_eq!(to_write, read); + } + + // Read the data from the test file and write it by the async writer and sync writer. + // And then compares the results of the two writers. + #[tokio::test] + async fn test_async_writer_with_sync_writer() { + let reader = get_test_reader(); + + let write_props = WriterProperties::builder() + .set_max_row_group_row_count(Some(64)) + .build(); + + let mut async_buffer = Vec::new(); + let mut async_writer = AsyncArrowWriter::try_new( + &mut async_buffer, + reader.schema(), + Some(write_props.clone()), + ) + .unwrap(); + + let mut sync_buffer = Vec::new(); + let mut sync_writer = + ArrowWriter::try_new(&mut sync_buffer, reader.schema(), Some(write_props)).unwrap(); + for record_batch in reader { + let record_batch = record_batch.unwrap(); + async_writer.write(&record_batch).await.unwrap(); + sync_writer.write(&record_batch).unwrap(); + } + sync_writer.close().unwrap(); + async_writer.close().await.unwrap(); + + assert_eq!(sync_buffer, async_buffer); + } + + #[tokio::test] + async fn test_async_writer_bytes_written() { + let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef; + let to_write = RecordBatch::try_from_iter([("col", col)]).unwrap(); + + let temp = tempfile::tempfile().unwrap(); + + let file = tokio::fs::File::from_std(temp.try_clone().unwrap()); + let mut writer = + AsyncArrowWriter::try_new(file.try_clone().await.unwrap(), to_write.schema(), None) + .unwrap(); + writer.write(&to_write).await.unwrap(); + let _metadata = writer.finish().await.unwrap(); + // After `finish` this should include the metadata and footer + let reported = writer.bytes_written(); + + // Get actual size from file metadata + let actual = file.metadata().await.unwrap().len() as usize; + + assert_eq!(reported, actual); + } + + #[tokio::test] + async fn test_async_writer_file() { + let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef; + let col2 = Arc::new(BinaryArray::from_iter_values(vec![ + vec![0; 500000], + vec![0; 500000], + vec![0; 500000], + ])) as ArrayRef; + let to_write = RecordBatch::try_from_iter([("col", col), ("col2", col2)]).unwrap(); + + let temp = tempfile::tempfile().unwrap(); + + let file = tokio::fs::File::from_std(temp.try_clone().unwrap()); + let mut writer = AsyncArrowWriter::try_new(file, to_write.schema(), None).unwrap(); + writer.write(&to_write).await.unwrap(); + writer.close().await.unwrap(); + + let mut reader = ParquetRecordBatchReaderBuilder::try_new(temp) + .unwrap() + .build() + .unwrap(); + let read = reader.next().unwrap().unwrap(); + + assert_eq!(to_write, read); + } + + #[tokio::test] + async fn in_progress_accounting() { + // define schema + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + + // create some data + let a = Int32Array::from_value(0_i32, 512); + + // build a record batch + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap(); + + let temp = tempfile::tempfile().unwrap(); + let file = tokio::fs::File::from_std(temp.try_clone().unwrap()); + let mut writer = AsyncArrowWriter::try_new(file, batch.schema(), None).unwrap(); + + // starts empty + assert_eq!(writer.in_progress_size(), 0); + assert_eq!(writer.in_progress_rows(), 0); + assert_eq!(writer.bytes_written(), 4); // Initial Parquet header + writer.write(&batch).await.unwrap(); + + // updated on write + let initial_size = writer.in_progress_size(); + assert!(initial_size > 0); + assert_eq!(writer.in_progress_rows(), batch.num_rows()); + let initial_memory = writer.memory_size(); + // memory estimate is larger than estimated encoded size + assert!( + initial_size <= initial_memory, + "{initial_size} <= {initial_memory}" + ); + + // updated on second write + writer.write(&batch).await.unwrap(); + assert!(writer.in_progress_size() > initial_size); + assert_eq!(writer.in_progress_rows(), batch.num_rows() * 2); + assert!(writer.memory_size() > initial_memory); + assert!( + writer.in_progress_size() <= writer.memory_size(), + "in_progress_size {} <= memory_size {}", + writer.in_progress_size(), + writer.memory_size() + ); + + // in progress tracking is cleared, but the overall data written is updated + let pre_flush_bytes_written = writer.bytes_written(); + writer.flush().await.unwrap(); + assert_eq!(writer.in_progress_size(), 0); + assert_eq!(writer.memory_size(), 0); + assert_eq!(writer.in_progress_rows(), 0); + assert!(writer.bytes_written() > pre_flush_bytes_written); + + writer.close().await.unwrap(); + } +} diff --git a/vendor/parquet/src/arrow/async_writer/store.rs b/vendor/parquet/src/arrow/async_writer/store.rs new file mode 100644 index 0000000..ad67426 --- /dev/null +++ b/vendor/parquet/src/arrow/async_writer/store.rs @@ -0,0 +1,173 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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 bytes::Bytes; +use futures::future::BoxFuture; +use std::sync::Arc; + +use crate::arrow::async_writer::AsyncFileWriter; +use crate::errors::{ParquetError, Result}; +use object_store::ObjectStore; +use object_store::buffered::BufWriter; +use object_store::path::Path; +use tokio::io::AsyncWriteExt; + +/// [`ParquetObjectWriter`] for writing to parquet to [`ObjectStore`] +/// +/// This type is deprecated: [`BufWriter`] implements [`AsyncWrite`] and can +/// therefore be passed to [`AsyncArrowWriter`] directly via the blanket +/// [`AsyncFileWriter`] implementation for [`AsyncWrite`] types: +/// +/// ``` +/// # use arrow_array::{ArrayRef, Int64Array, RecordBatch}; +/// # use object_store::buffered::BufWriter; +/// # use object_store::memory::InMemory; +/// # use object_store::path::Path; +/// # use object_store::{ObjectStore, ObjectStoreExt}; +/// # use std::sync::Arc; +/// +/// # use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +/// # use parquet::arrow::AsyncArrowWriter; +/// +/// # #[tokio::main(flavor="current_thread")] +/// # async fn main() { +/// let store = Arc::new(InMemory::new()); +/// +/// let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef; +/// let to_write = RecordBatch::try_from_iter([("col", col)]).unwrap(); +/// +/// let object_store_writer = BufWriter::new(store.clone(), Path::from("test")); +/// let mut writer = +/// AsyncArrowWriter::try_new(object_store_writer, to_write.schema(), None).unwrap(); +/// writer.write(&to_write).await.unwrap(); +/// writer.close().await.unwrap(); +/// +/// let buffer = store +/// .get(&Path::from("test")) +/// .await +/// .unwrap() +/// .bytes() +/// .await +/// .unwrap(); +/// let mut reader = ParquetRecordBatchReaderBuilder::try_new(buffer) +/// .unwrap() +/// .build() +/// .unwrap(); +/// let read = reader.next().unwrap().unwrap(); +/// +/// assert_eq!(to_write, read); +/// # } +/// ``` +/// +/// [`AsyncWrite`]: tokio::io::AsyncWrite +/// [`AsyncArrowWriter`]: crate::arrow::async_writer::AsyncArrowWriter +#[deprecated( + since = "59.2.0", + note = "Pass an `object_store::buffered::BufWriter` to `AsyncArrowWriter` directly instead; see https://github.com/apache/arrow-rs/issues/10308 and `parquet/examples/object_store.rs`." +)] +#[derive(Debug)] +pub struct ParquetObjectWriter { + w: BufWriter, +} + +#[allow(deprecated)] +impl ParquetObjectWriter { + /// Create a new [`ParquetObjectWriter`] that writes to the specified path in the given store. + /// + /// To configure the writer behavior, please build [`BufWriter`] and then use [`Self::from_buf_writer`] + pub fn new(store: Arc, path: Path) -> Self { + Self::from_buf_writer(BufWriter::new(store, path)) + } + + /// Construct a new ParquetObjectWriter via a existing BufWriter. + pub fn from_buf_writer(w: BufWriter) -> Self { + Self { w } + } + + /// Consume the writer and return the underlying BufWriter. + pub fn into_inner(self) -> BufWriter { + self.w + } +} + +#[allow(deprecated)] +impl AsyncFileWriter for ParquetObjectWriter { + fn write(&mut self, bs: Bytes) -> BoxFuture<'_, Result<()>> { + Box::pin(async { + self.w + .put(bs) + .await + .map_err(|err| ParquetError::External(Box::new(err))) + }) + } + + fn complete(&mut self) -> BoxFuture<'_, Result<()>> { + Box::pin(async { + self.w + .shutdown() + .await + .map_err(|err| ParquetError::External(Box::new(err))) + }) + } +} +#[allow(deprecated)] +impl From for ParquetObjectWriter { + fn from(w: BufWriter) -> Self { + Self::from_buf_writer(w) + } +} +#[cfg(test)] +#[allow(deprecated)] +mod tests { + use arrow_array::{ArrayRef, Int64Array, RecordBatch}; + use object_store::memory::InMemory; + use std::sync::Arc; + + use super::*; + use crate::arrow::AsyncArrowWriter; + use crate::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use object_store::ObjectStoreExt; + + #[tokio::test] + async fn test_async_writer() { + let store = Arc::new(InMemory::new()); + + let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef; + let to_write = RecordBatch::try_from_iter([("col", col)]).unwrap(); + + let object_store_writer = ParquetObjectWriter::new(store.clone(), Path::from("test")); + let mut writer = + AsyncArrowWriter::try_new(object_store_writer, to_write.schema(), None).unwrap(); + writer.write(&to_write).await.unwrap(); + writer.close().await.unwrap(); + + let buffer = store + .get(&Path::from("test")) + .await + .unwrap() + .bytes() + .await + .unwrap(); + let mut reader = ParquetRecordBatchReaderBuilder::try_new(buffer) + .unwrap() + .build() + .unwrap(); + let read = reader.next().unwrap().unwrap(); + + assert_eq!(to_write, read); + } +} diff --git a/vendor/parquet/src/arrow/buffer/bit_util.rs b/vendor/parquet/src/arrow/buffer/bit_util.rs new file mode 100644 index 0000000..985943b --- /dev/null +++ b/vendor/parquet/src/arrow/buffer/bit_util.rs @@ -0,0 +1,105 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_buffer::bit_chunk_iterator::UnalignedBitChunk; +use std::ops::Range; + +/// Counts the number of set bits in the provided range +pub fn count_set_bits(bytes: &[u8], range: Range) -> usize { + let unaligned = UnalignedBitChunk::new(bytes, range.start, range.end - range.start); + unaligned.count_ones() +} + +/// Iterates through the set bit positions in `bytes` in reverse order +pub fn iter_set_bits_rev(bytes: &[u8]) -> impl Iterator + '_ { + let bit_length = bytes.len() * 8; + let unaligned = UnalignedBitChunk::new(bytes, 0, bit_length); + let mut chunk_end_idx = bit_length + unaligned.lead_padding() + unaligned.trailing_padding(); + + let iter = unaligned + .prefix() + .into_iter() + .chain(unaligned.chunks().iter().cloned()) + .chain(unaligned.suffix()); + + iter.rev().flat_map(move |chunk| { + let chunk_idx = chunk_end_idx - 64; + chunk_end_idx = chunk_idx; + let mut rev_chunk = chunk.reverse_bits(); + std::iter::from_fn(move || { + if rev_chunk != 0 { + let bit_pos = rev_chunk.trailing_zeros(); + rev_chunk &= rev_chunk - 1; + return Some(chunk_idx + (63 - bit_pos as usize)); + } + None + }) + }) +} + +/// Performs big endian sign extension +pub fn sign_extend_be(b: &[u8]) -> [u8; N] { + assert!(b.len() <= N, "Array too large, expected less than {N}"); + let is_negative = (b[0] & 128u8) == 128u8; + let mut result = if is_negative { [255u8; N] } else { [0u8; N] }; + for (d, s) in result.iter_mut().skip(N - b.len()).zip(b) { + *d = *s; + } + result +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::builder::BooleanBufferBuilder; + use rand::{prelude::*, rng}; + + #[test] + fn test_bit_fns() { + let mut rng = rng(); + let mask_length = rng.random_range(1..1024); + let bools: Vec<_> = std::iter::from_fn(|| Some(rng.next_u32() & 1 == 0)) + .take(mask_length) + .collect(); + + let mut nulls = BooleanBufferBuilder::new(mask_length); + bools.iter().for_each(|b| nulls.append(*b)); + + let actual: Vec<_> = iter_set_bits_rev(nulls.as_slice()).collect(); + let expected: Vec<_> = bools + .iter() + .enumerate() + .rev() + .filter_map(|(x, y)| y.then_some(x)) + .collect(); + assert_eq!(actual, expected); + + assert_eq!(iter_set_bits_rev(&[]).count(), 0); + assert_eq!(count_set_bits(&[], 0..0), 0); + assert_eq!(count_set_bits(&[0xFF], 1..1), 0); + + for _ in 0..20 { + let start = rng.random_range(0..bools.len()); + let end = rng.random_range(start..bools.len()); + + let actual = count_set_bits(nulls.as_slice(), start..end); + let expected = bools[start..end].iter().filter(|x| **x).count(); + + assert_eq!(actual, expected); + } + } +} diff --git a/vendor/parquet/src/arrow/buffer/dictionary_buffer.rs b/vendor/parquet/src/arrow/buffer/dictionary_buffer.rs new file mode 100644 index 0000000..33537b4 --- /dev/null +++ b/vendor/parquet/src/arrow/buffer/dictionary_buffer.rs @@ -0,0 +1,407 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::buffer::offset_buffer::OffsetBuffer; +use crate::arrow::record_reader::buffer::ValuesBuffer; +use crate::errors::{ParquetError, Result}; +use arrow_array::{Array, GenericByteArray, downcast_integer}; +use arrow_array::{ + ArrayRef, FixedSizeBinaryArray, OffsetSizeTrait, + builder::{FixedSizeBinaryDictionaryBuilder, GenericByteDictionaryBuilder}, + cast::AsArray, + make_array, + types::{ArrowDictionaryKeyType, ByteArrayType}, +}; +use arrow_buffer::{ArrowNativeType, Buffer}; +use arrow_data::ArrayDataBuilder; +use arrow_schema::DataType as ArrowType; +use std::sync::Arc; + +/// An array of variable length byte arrays that are potentially dictionary encoded +/// and can be converted into a corresponding [`ArrayRef`] +pub enum DictionaryBuffer { + Dict { keys: Vec, values: ArrayRef }, + Values { values: OffsetBuffer }, +} + +impl DictionaryBuffer { + #[allow(unused)] + pub fn len(&self) -> usize { + match self { + Self::Dict { keys, .. } => keys.len(), + Self::Values { values } => values.len(), + } + } + + /// Returns a mutable reference to a keys array + /// + /// Returns None if the dictionary needs to be recomputed + /// + /// # Panic + /// + /// Panics if the dictionary is too large for `K` + pub fn as_keys(&mut self, dictionary: &ArrayRef) -> Option<&mut Vec> { + assert!(K::from_usize(dictionary.len()).is_some()); + + match self { + Self::Dict { keys, values } => { + // Need to discard fat pointer for equality check + // - https://stackoverflow.com/a/67114787 + // - https://github.com/rust-lang/rust/issues/46139 + let values_ptr = values.as_ref() as *const _ as *const (); + let dict_ptr = dictionary.as_ref() as *const _ as *const (); + if values_ptr == dict_ptr { + Some(keys) + } else if keys.is_empty() { + *values = Arc::clone(dictionary); + Some(keys) + } else { + None + } + } + Self::Values { values } if values.is_empty() => { + *self = Self::Dict { + keys: Default::default(), + values: Arc::clone(dictionary), + }; + match self { + Self::Dict { keys, .. } => Some(keys), + _ => unreachable!(), + } + } + _ => None, + } + } + + /// Returns a mutable reference to a values array + /// + /// If this is currently dictionary encoded, this will convert from the + /// dictionary encoded representation + pub fn spill_values(&mut self) -> Result<&mut OffsetBuffer> { + match self { + Self::Values { values } => Ok(values), + Self::Dict { keys, values } => { + let mut spilled = OffsetBuffer::with_capacity(0); + let data = values.to_data(); + let dict_buffers = data.buffers(); + let dict_offsets = dict_buffers[0].typed_data::(); + let dict_values = dict_buffers[1].as_slice(); + + if values.is_empty() { + // If dictionary is empty, zero pad offsets + spilled.offsets.resize(keys.len() + 1, V::default()); + } else { + // Note: at this point null positions will have arbitrary dictionary keys + // and this will hydrate them to the corresponding byte array. This is + // likely sub-optimal, as we would prefer zero length null "slots", but + // spilling is already a degenerate case and so it is unclear if this is + // worth optimising for, e.g. by keeping a null mask around + spilled.extend_from_dictionary(keys.as_slice(), dict_offsets, dict_values)?; + } + + *self = Self::Values { values: spilled }; + match self { + Self::Values { values } => Ok(values), + _ => unreachable!(), + } + } + } + } + + /// Converts this into an [`ArrayRef`] with the provided `data_type` and `null_buffer` + pub fn into_array( + self, + null_buffer: Option, + data_type: &ArrowType, + ) -> Result { + assert!(matches!(data_type, ArrowType::Dictionary(_, _))); + + match self { + Self::Dict { keys, values } => { + // Validate keys unless dictionary is empty + if !values.is_empty() { + let min = K::from_usize(0).unwrap(); + let max = K::from_usize(values.len()).unwrap(); + + // using copied and fold gets auto-vectorized since rust 1.70 + // all/any would allow early exit on invalid values + // but in the happy case all values have to be checked anyway + if !keys + .as_slice() + .iter() + .copied() + .fold(true, |a, x| a && x >= min && x < max) + { + return Err(general_err!( + "dictionary key beyond bounds of dictionary: 0..{}", + values.len() + )); + } + } + + let ArrowType::Dictionary(_, value_type) = data_type else { + unreachable!() + }; + let values = if let ArrowType::FixedSizeBinary(size) = **value_type { + let binary = values.as_binary::(); + Arc::new(FixedSizeBinaryArray::new( + size, + binary.values().clone(), + binary.nulls().cloned(), + )) as _ + } else { + values + }; + + let builder = ArrayDataBuilder::new(data_type.clone()) + .len(keys.len()) + .add_buffer(Buffer::from_vec(keys)) + .add_child_data(values.into_data()) + .null_bit_buffer(null_buffer); + + let data = match cfg!(debug_assertions) { + true => builder.build().unwrap(), + false => unsafe { builder.build_unchecked() }, + }; + + Ok(make_array(data)) + } + Self::Values { values } => { + let (key_type, value_type) = match data_type { + ArrowType::Dictionary(k, v) => (k, v.as_ref().clone()), + _ => unreachable!(), + }; + + let array = values.into_array(null_buffer, value_type); + pack_values(key_type, &array) + } + } + } +} + +impl ValuesBuffer for DictionaryBuffer { + fn with_capacity(capacity: usize) -> Self { + Self::Values { + values: OffsetBuffer::with_capacity(capacity), + } + } + + fn reserve_exact(&mut self, additional: usize) { + match self { + Self::Dict { keys, .. } => keys.reserve_exact(additional), + Self::Values { values, .. } => values.reserve_exact(additional), + } + } + + fn pad_nulls( + &mut self, + read_offset: usize, + values_read: usize, + levels_read: usize, + valid_mask: &[u8], + ) -> Result<()> { + match self { + Self::Dict { keys, .. } => { + keys.resize(read_offset + levels_read, K::default()); + keys.pad_nulls(read_offset, values_read, levels_read, valid_mask) + } + Self::Values { values, .. } => { + values.pad_nulls(read_offset, values_read, levels_read, valid_mask) + } + } + } +} + +macro_rules! dict_helper { + ($k:ty, $array:ident) => { + match $array.data_type() { + ArrowType::Utf8 => pack_values_impl::<$k, _>($array.as_string::()), + ArrowType::LargeUtf8 => pack_values_impl::<$k, _>($array.as_string::()), + ArrowType::Binary => pack_values_impl::<$k, _>($array.as_binary::()), + ArrowType::LargeBinary => pack_values_impl::<$k, _>($array.as_binary::()), + ArrowType::FixedSizeBinary(_) => { + pack_fixed_values_impl::<$k>($array.as_fixed_size_binary()) + } + _ => unreachable!(), + } + }; +} + +fn pack_values(key_type: &ArrowType, values: &ArrayRef) -> Result { + downcast_integer! { + key_type => (dict_helper, values), + _ => unreachable!(), + } +} + +fn pack_values_impl( + array: &GenericByteArray, +) -> Result { + let mut builder = GenericByteDictionaryBuilder::::with_capacity(array.len(), 1024, 1024); + for x in array { + match x { + Some(x) => builder.append_value(x), + None => builder.append_null(), + } + } + let raw = builder.finish(); + Ok(Arc::new(raw)) +} + +fn pack_fixed_values_impl( + array: &FixedSizeBinaryArray, +) -> Result { + let mut builder = FixedSizeBinaryDictionaryBuilder::::with_capacity( + array.len(), + 1024, + array.value_length(), + ); + for x in array { + match x { + Some(x) => builder.append_value(x), + None => builder.append_null(), + } + } + let raw = builder.finish(); + Ok(Arc::new(raw)) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::compute::cast; + use arrow_array::StringArray; + + #[test] + fn test_dictionary_buffer() { + let dict_type = + ArrowType::Dictionary(Box::new(ArrowType::Int32), Box::new(ArrowType::Utf8)); + + let d1: ArrayRef = Arc::new(StringArray::from(vec!["hello", "world", "", "a", "b"])); + + let mut buffer = DictionaryBuffer::::with_capacity(0); + + // Read some data preserving the dictionary + let values = &[1, 0, 3, 2, 4]; + buffer.as_keys(&d1).unwrap().extend_from_slice(values); + + let mut valid = vec![false, false, true, true, false, true, true, true]; + let valid_buffer = Buffer::from_iter(valid.iter().cloned()); + buffer + .pad_nulls(0, values.len(), valid.len(), valid_buffer.as_slice()) + .unwrap(); + + // Read some data not preserving the dictionary + + let values = buffer.spill_values().unwrap(); + let read_offset = values.len(); + values.try_push("bingo".as_bytes(), false).unwrap(); + values.try_push("bongo".as_bytes(), false).unwrap(); + + valid.extend_from_slice(&[false, false, true, false, true]); + let null_buffer = Buffer::from_iter(valid.iter().cloned()); + buffer + .pad_nulls(read_offset, 2, 5, null_buffer.as_slice()) + .unwrap(); + + assert_eq!(buffer.len(), 13); + let split = std::mem::replace(&mut buffer, DictionaryBuffer::with_capacity(0)); + + let array = split.into_array(Some(null_buffer), &dict_type).unwrap(); + assert_eq!(array.data_type(), &dict_type); + + let strings = cast(&array, &ArrowType::Utf8).unwrap(); + let strings = strings.as_any().downcast_ref::().unwrap(); + assert_eq!( + strings.iter().collect::>(), + vec![ + None, + None, + Some("world"), + Some("hello"), + None, + Some("a"), + Some(""), + Some("b"), + None, + None, + Some("bingo"), + None, + Some("bongo") + ] + ); + + // Can recreate with new dictionary as values is empty + assert!(matches!(&buffer, DictionaryBuffer::Values { .. })); + assert_eq!(buffer.len(), 0); + let d2 = Arc::new(StringArray::from(vec!["bingo", ""])) as ArrayRef; + buffer + .as_keys(&d2) + .unwrap() + .extend_from_slice(&[0, 1, 0, 1]); + + let array = std::mem::replace(&mut buffer, DictionaryBuffer::with_capacity(0)) + .into_array(None, &dict_type) + .unwrap(); + assert_eq!(array.data_type(), &dict_type); + + let strings = cast(&array, &ArrowType::Utf8).unwrap(); + let strings = strings.as_any().downcast_ref::().unwrap(); + assert_eq!( + strings.iter().collect::>(), + vec![Some("bingo"), Some(""), Some("bingo"), Some("")] + ); + + // Can recreate with new dictionary as keys empty + assert!(matches!(&buffer, DictionaryBuffer::Values { .. })); + assert_eq!(buffer.len(), 0); + let d3 = Arc::new(StringArray::from(vec!["bongo"])) as ArrayRef; + buffer.as_keys(&d3).unwrap().extend_from_slice(&[0, 0]); + + // Cannot change dictionary as keys not empty + let d4 = Arc::new(StringArray::from(vec!["bananas"])) as ArrayRef; + assert!(buffer.as_keys(&d4).is_none()); + } + + #[test] + fn test_validates_keys() { + let dict_type = + ArrowType::Dictionary(Box::new(ArrowType::Int32), Box::new(ArrowType::Utf8)); + + let mut buffer = DictionaryBuffer::::with_capacity(0); + let d = Arc::new(StringArray::from(vec!["", "f"])) as ArrayRef; + buffer.as_keys(&d).unwrap().extend_from_slice(&[0, 2, 0]); + + let err = buffer.into_array(None, &dict_type).unwrap_err().to_string(); + assert!( + err.contains("dictionary key beyond bounds of dictionary: 0..2"), + "{}", + err + ); + + let mut buffer = DictionaryBuffer::::with_capacity(0); + let d = Arc::new(StringArray::from(vec![""])) as ArrayRef; + buffer.as_keys(&d).unwrap().extend_from_slice(&[0, 1, 0]); + + let err = buffer.spill_values().unwrap_err().to_string(); + assert!( + err.contains("dictionary key beyond bounds of dictionary: 0..1"), + "{}", + err + ); + } +} diff --git a/vendor/parquet/src/arrow/buffer/mod.rs b/vendor/parquet/src/arrow/buffer/mod.rs new file mode 100644 index 0000000..e7566d9 --- /dev/null +++ b/vendor/parquet/src/arrow/buffer/mod.rs @@ -0,0 +1,23 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Logic for reading data into arrow buffers + +pub mod bit_util; +pub mod dictionary_buffer; +pub mod offset_buffer; +pub mod view_buffer; diff --git a/vendor/parquet/src/arrow/buffer/offset_buffer.rs b/vendor/parquet/src/arrow/buffer/offset_buffer.rs new file mode 100644 index 0000000..ab67694 --- /dev/null +++ b/vendor/parquet/src/arrow/buffer/offset_buffer.rs @@ -0,0 +1,392 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::buffer::bit_util::iter_set_bits_rev; +use crate::arrow::record_reader::buffer::ValuesBuffer; +use crate::errors::{ParquetError, Result}; +use crate::util::utf8::check_valid_utf8; +use arrow_array::{ArrayRef, OffsetSizeTrait, make_array}; +use arrow_buffer::{ArrowNativeType, Buffer}; +use arrow_data::ArrayDataBuilder; +use arrow_schema::DataType as ArrowType; + +/// A buffer of variable-sized byte arrays that can be converted into +/// a corresponding [`ArrayRef`] +#[derive(Debug)] +pub struct OffsetBuffer { + pub offsets: Vec, + pub values: Vec, +} + +impl OffsetBuffer { + /// Create a new `OffsetBuffer` with capacity for at least `capacity` elements + /// + /// Pre-allocates the offsets vector to avoid reallocations during reading. + /// The values vector is not pre-allocated as its size is unpredictable. + pub fn with_capacity(capacity: usize) -> Self { + let mut offsets = Vec::with_capacity(capacity + 1); + offsets.push(I::default()); + Self { + offsets, + values: Vec::new(), + } + } + + /// Returns the number of byte arrays in this buffer + pub fn len(&self) -> usize { + self.offsets.len() - 1 + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// If `validate_utf8` this verifies that the first character of `data` is + /// the start of a UTF-8 codepoint + /// + /// Note: This does not verify that the entirety of `data` is valid + /// UTF-8. This should be done by calling [`Self::check_valid_utf8`] after + /// all data has been written + pub fn try_push(&mut self, data: &[u8], validate_utf8: bool) -> Result<()> { + if validate_utf8 { + if let Some(&b) = data.first() { + // A valid code-point iff it does not start with 0b10xxxxxx + // Bit-magic taken from `std::str::is_char_boundary` + if (b as i8) < -0x40 { + return Err(ParquetError::General( + "encountered non UTF-8 data".to_string(), + )); + } + } + } + + self.values.extend_from_slice(data); + + let index_offset = I::from_usize(self.values.len()) + .ok_or_else(|| general_err!("index overflow decoding byte array"))?; + + self.offsets.push(index_offset); + Ok(()) + } + + /// Extends this buffer with a list of keys + /// + /// For each value `key` in `keys` this will insert + /// `&dict_values[dict_offsets[key]..dict_offsets[key+1]]` + /// + /// Note: This will validate offsets are valid + pub fn extend_from_dictionary( + &mut self, + keys: &[K], + dict_offsets: &[V], + dict_values: &[u8], + ) -> Result<()> { + self.offsets.reserve(keys.len()); + + for key in keys { + let index = key.as_usize(); + if index + 1 >= dict_offsets.len() { + return Err(general_err!( + "dictionary key beyond bounds of dictionary: 0..{}", + dict_offsets.len().saturating_sub(1) + )); + } + let start_offset = dict_offsets[index].as_usize(); + let end_offset = dict_offsets[index + 1].as_usize(); + + // Dictionary values are verified when decoding dictionary page + self.values + .extend_from_slice(&dict_values[start_offset..end_offset]); + let index_offset = I::from_usize(self.values.len()) + .ok_or_else(|| general_err!("index overflow decoding byte array"))?; + self.offsets.push(index_offset); + } + Ok(()) + } + + /// Validates that `&self.values[start_offset..]` is a valid UTF-8 sequence + /// + /// This MUST be combined with validating that the offsets start on a character + /// boundary, otherwise it would be possible for the values array to be a valid UTF-8 + /// sequence, but not the individual string slices it contains + /// + /// [`Self::try_push`] can perform this validation check on insertion + pub fn check_valid_utf8(&self, start_offset: usize) -> Result<()> { + check_valid_utf8(&self.values.as_slice()[start_offset..]) + } + + /// Converts this into an [`ArrayRef`] with the provided `data_type` and `null_buffer` + pub fn into_array(self, null_buffer: Option, data_type: ArrowType) -> ArrayRef { + let array_data_builder = ArrayDataBuilder::new(data_type) + .len(self.len()) + .add_buffer(Buffer::from_vec(self.offsets)) + .add_buffer(Buffer::from_vec(self.values)) + .null_bit_buffer(null_buffer); + + let data = match cfg!(debug_assertions) { + true => array_data_builder.build().unwrap(), + false => unsafe { array_data_builder.build_unchecked() }, + }; + + make_array(data) + } +} + +impl ValuesBuffer for OffsetBuffer { + fn with_capacity(capacity: usize) -> Self { + Self::with_capacity(capacity) + } + + fn reserve_exact(&mut self, additional: usize) { + self.offsets.reserve_exact(additional); + } + + fn pad_nulls( + &mut self, + read_offset: usize, + values_read: usize, + levels_read: usize, + valid_mask: &[u8], + ) -> Result<()> { + if self.offsets.len() != read_offset + values_read + 1 { + return Err(general_err!( + "found inconsistent offsets while padding nulls: expected {} offsets, got {}", + read_offset + values_read + 1, + self.offsets.len() + )); + } + self.offsets + .resize(read_offset + levels_read + 1, I::default()); + + let offsets = &mut self.offsets; + + let mut last_pos = read_offset + levels_read + 1; + let mut last_start_offset = I::from_usize(self.values.len()).unwrap(); + + let values_range = read_offset..read_offset + values_read; + for (value_pos, level_pos) in values_range + .clone() + .rev() + .zip(iter_set_bits_rev(valid_mask)) + { + if level_pos < value_pos || level_pos >= last_pos { + return Err(general_err!("found corrupt level data while padding nulls")); + } + + let end_offset = offsets[value_pos + 1]; + let start_offset = offsets[value_pos]; + + // Fill in any nulls + for x in &mut offsets[level_pos + 1..last_pos] { + *x = end_offset; + } + + if level_pos == value_pos { + return Ok(()); + } + + offsets[level_pos] = start_offset; + last_pos = level_pos; + last_start_offset = start_offset; + } + + // Pad leading nulls up to `last_offset` + for x in &mut offsets[values_range.start + 1..last_pos] { + *x = last_start_offset + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::{Array, LargeStringArray, StringArray}; + + #[test] + fn test_offset_buffer_empty() { + let buffer = OffsetBuffer::::with_capacity(0); + let array = buffer.into_array(None, ArrowType::Utf8); + let strings = array.as_any().downcast_ref::().unwrap(); + assert_eq!(strings.len(), 0); + } + + #[test] + fn test_offset_buffer_append() { + let mut buffer = OffsetBuffer::::with_capacity(0); + buffer.try_push("hello".as_bytes(), true).unwrap(); + buffer.try_push("bar".as_bytes(), true).unwrap(); + buffer + .extend_from_dictionary(&[1, 3, 0, 2], &[0, 2, 4, 5, 6], "abcdef".as_bytes()) + .unwrap(); + + let array = buffer.into_array(None, ArrowType::LargeUtf8); + let strings = array.as_any().downcast_ref::().unwrap(); + assert_eq!( + strings.iter().map(|x| x.unwrap()).collect::>(), + vec!["hello", "bar", "cd", "f", "ab", "e"] + ) + } + + #[test] + fn test_offset_buffer() { + let mut buffer = OffsetBuffer::::with_capacity(0); + for v in ["hello", "world", "cupcakes", "a", "b", "c"] { + buffer.try_push(v.as_bytes(), false).unwrap() + } + let split = std::mem::replace(&mut buffer, OffsetBuffer::with_capacity(0)); + + let array = split.into_array(None, ArrowType::Utf8); + let strings = array.as_any().downcast_ref::().unwrap(); + assert_eq!( + strings.iter().map(|x| x.unwrap()).collect::>(), + vec!["hello", "world", "cupcakes", "a", "b", "c"] + ); + + buffer.try_push("test".as_bytes(), false).unwrap(); + let array = buffer.into_array(None, ArrowType::Utf8); + let strings = array.as_any().downcast_ref::().unwrap(); + assert_eq!( + strings.iter().map(|x| x.unwrap()).collect::>(), + vec!["test"] + ); + } + + #[test] + fn test_offset_buffer_pad_nulls() { + let mut buffer = OffsetBuffer::::with_capacity(0); + let values = ["a", "b", "c", "def", "gh"]; + for v in &values { + buffer.try_push(v.as_bytes(), false).unwrap() + } + + let valid = [ + true, false, false, true, false, true, false, true, true, false, false, + ]; + let valid_mask = Buffer::from_iter(valid.iter().copied()); + + // Both trailing and leading nulls + buffer + .pad_nulls(1, values.len() - 1, valid.len() - 1, valid_mask.as_slice()) + .unwrap(); + + let array = buffer.into_array(Some(valid_mask), ArrowType::Utf8); + let strings = array.as_any().downcast_ref::().unwrap(); + assert_eq!( + strings.iter().collect::>(), + vec![ + Some("a"), + None, + None, + Some("b"), + None, + Some("c"), + None, + Some("def"), + Some("gh"), + None, + None + ] + ); + } + + #[test] + fn test_utf8_validation() { + let valid_2_byte_utf8 = &[0b11001000, 0b10001000]; + std::str::from_utf8(valid_2_byte_utf8).unwrap(); + let valid_3_byte_utf8 = &[0b11101000, 0b10001000, 0b10001000]; + std::str::from_utf8(valid_3_byte_utf8).unwrap(); + let valid_4_byte_utf8 = &[0b11110010, 0b10101000, 0b10101001, 0b10100101]; + std::str::from_utf8(valid_4_byte_utf8).unwrap(); + + let mut buffer = OffsetBuffer::::with_capacity(0); + buffer.try_push(valid_2_byte_utf8, true).unwrap(); + buffer.try_push(valid_3_byte_utf8, true).unwrap(); + buffer.try_push(valid_4_byte_utf8, true).unwrap(); + + // Cannot append string starting with incomplete codepoint + buffer.try_push(&valid_2_byte_utf8[1..], true).unwrap_err(); + buffer.try_push(&valid_3_byte_utf8[1..], true).unwrap_err(); + buffer.try_push(&valid_3_byte_utf8[2..], true).unwrap_err(); + buffer.try_push(&valid_4_byte_utf8[1..], true).unwrap_err(); + buffer.try_push(&valid_4_byte_utf8[2..], true).unwrap_err(); + buffer.try_push(&valid_4_byte_utf8[3..], true).unwrap_err(); + + // Can append data containing an incomplete codepoint + buffer.try_push(&[0b01111111, 0b10111111], true).unwrap(); + + assert_eq!(buffer.len(), 4); + assert_eq!(buffer.values.len(), 11); + + buffer.try_push(valid_3_byte_utf8, true).unwrap(); + + // Should fail due to incomplete codepoint + buffer.check_valid_utf8(0).unwrap_err(); + + // After broken codepoint -> success + buffer.check_valid_utf8(11).unwrap(); + + // Fails if run from middle of codepoint + buffer.check_valid_utf8(12).unwrap_err(); + } + + #[test] + fn test_pad_nulls_corrupt_input_returns_err() { + // Corrupt input must produce a decode error rather than panicking. + + // Offsets inconsistent with `values_read`: only one value was pushed, + // but three are claimed to have been read. + let mut buffer = OffsetBuffer::::with_capacity(0); + buffer.try_push("a".as_bytes(), false).unwrap(); + let valid_mask = Buffer::from_iter([true, false, false]); + let err = buffer + .pad_nulls(0, 3, 3, valid_mask.as_slice()) + .unwrap_err(); + assert!( + err.to_string().contains("inconsistent offsets"), + "unexpected error: {err}" + ); + + // Valid mask has fewer set bits than `values_read`, which previously + // tripped an assertion in the null-padding loop. + let mut buffer = OffsetBuffer::::with_capacity(0); + for v in ["a", "b", "c"] { + buffer.try_push(v.as_bytes(), false).unwrap(); + } + let valid_mask = Buffer::from_iter([true, false, false]); + let err = buffer + .pad_nulls(0, 3, 3, valid_mask.as_slice()) + .unwrap_err(); + assert!( + err.to_string().contains("corrupt level data"), + "unexpected error: {err}" + ); + } + + #[test] + fn test_pad_nulls_empty() { + let mut buffer = OffsetBuffer::::with_capacity(0); + let valid_mask = Buffer::from_iter(std::iter::repeat_n(false, 9)); + buffer.pad_nulls(0, 0, 9, valid_mask.as_slice()).unwrap(); + + let array = buffer.into_array(Some(valid_mask), ArrowType::Utf8); + let strings = array.as_any().downcast_ref::().unwrap(); + + assert_eq!(strings.len(), 9); + assert!(strings.iter().all(|x| x.is_none())) + } +} diff --git a/vendor/parquet/src/arrow/buffer/view_buffer.rs b/vendor/parquet/src/arrow/buffer/view_buffer.rs new file mode 100644 index 0000000..9670d0e --- /dev/null +++ b/vendor/parquet/src/arrow/buffer/view_buffer.rs @@ -0,0 +1,177 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::record_reader::buffer::ValuesBuffer; +use crate::errors::Result; +use arrow_array::{ArrayRef, BinaryViewArray, StringViewArray}; +use arrow_buffer::{Buffer, NullBuffer, ScalarBuffer}; +use arrow_schema::DataType as ArrowType; +use std::sync::Arc; + +/// A buffer of view type byte arrays that can be converted into +/// `GenericByteViewArray` +/// +/// Note this does not reuse `GenericByteViewBuilder` due to the need to call `pad_nulls` +/// and reuse the existing logic for Vec in the parquet crate +#[derive(Debug, Default)] +pub struct ViewBuffer { + pub views: Vec, + pub buffers: Vec, +} + +impl ViewBuffer { + /// Create a new ViewBuffer with capacity for the specified number of views + pub fn with_capacity(capacity: usize) -> Self { + Self { + views: Vec::with_capacity(capacity), + buffers: Vec::new(), + } + } + + pub fn is_empty(&self) -> bool { + self.views.is_empty() + } + + pub fn append_block(&mut self, block: Buffer) -> u32 { + let block_id = self.buffers.len() as u32; + self.buffers.push(block); + block_id + } + + /// Converts this into an [`ArrayRef`] with the provided `data_type` and `null_buffer` + pub fn into_array(self, null_buffer: Option, data_type: &ArrowType) -> ArrayRef { + let len = self.views.len(); + let views = ScalarBuffer::from(self.views); + let nulls = null_buffer.and_then(|b| NullBuffer::from_unsliced_buffer(b, len)); + match data_type { + ArrowType::Utf8View => { + // Safety: views were created correctly, and checked that the data is utf8 when building the buffer + unsafe { Arc::new(StringViewArray::new_unchecked(views, self.buffers, nulls)) } + } + ArrowType::BinaryView => { + // Safety: views were created correctly + unsafe { Arc::new(BinaryViewArray::new_unchecked(views, self.buffers, nulls)) } + } + _ => panic!("Unsupported data type: {data_type}"), + } + } +} + +impl ValuesBuffer for ViewBuffer { + fn with_capacity(capacity: usize) -> Self { + Self::with_capacity(capacity) + } + + fn reserve_exact(&mut self, additional: usize) { + self.views.reserve_exact(additional); + } + + fn pad_nulls( + &mut self, + read_offset: usize, + values_read: usize, + levels_read: usize, + valid_mask: &[u8], + ) -> Result<()> { + self.views + .pad_nulls(read_offset, values_read, levels_read, valid_mask) + } +} + +#[cfg(test)] +mod tests { + + use arrow::array::make_view; + use arrow_array::Array; + + use super::*; + + #[test] + fn test_view_buffer_empty() { + let buffer = ViewBuffer::with_capacity(0); + let array = buffer.into_array(None, &ArrowType::Utf8View); + let strings = array + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(strings.len(), 0); + } + + #[test] + fn test_view_buffer_append_view() { + let mut buffer = ViewBuffer::with_capacity(0); + let data = b"0123456789long string to test string view"; + let string_buffer = Buffer::from(data); + let block_id = buffer.append_block(string_buffer); + + buffer.views.push(make_view(&data[0..1], block_id, 0)); + buffer.views.push(make_view(&data[1..10], block_id, 1)); + buffer.views.push(make_view(&data[10..41], block_id, 10)); + + let array = buffer.into_array(None, &ArrowType::Utf8View); + let string_array = array + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + string_array.iter().collect::>(), + vec![ + Some("0"), + Some("123456789"), + Some("long string to test string view"), + ] + ); + } + + #[test] + fn test_view_buffer_pad_null() { + let mut buffer = ViewBuffer::with_capacity(0); + let data = b"0123456789long string to test string view"; + let string_buffer = Buffer::from(data); + let block_id = buffer.append_block(string_buffer); + + buffer.views.push(make_view(&data[0..1], block_id, 0)); + buffer.views.push(make_view(&data[1..10], block_id, 1)); + buffer.views.push(make_view(&data[10..41], block_id, 10)); + + let valid = [true, false, false, true, false, false, true]; + let valid_mask = Buffer::from_iter(valid.iter().copied()); + + buffer + .pad_nulls(1, 2, valid.len() - 1, valid_mask.as_slice()) + .unwrap(); + + let array = buffer.into_array(Some(valid_mask), &ArrowType::Utf8View); + let strings = array + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!( + strings.iter().collect::>(), + vec![ + Some("0"), + None, + None, + Some("123456789"), + None, + None, + Some("long string to test string view"), + ] + ); + } +} diff --git a/vendor/parquet/src/arrow/decoder/delta_byte_array.rs b/vendor/parquet/src/arrow/decoder/delta_byte_array.rs new file mode 100644 index 0000000..7686a42 --- /dev/null +++ b/vendor/parquet/src/arrow/decoder/delta_byte_array.rs @@ -0,0 +1,138 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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 bytes::Bytes; + +use crate::data_type::Int32Type; +use crate::encodings::decoding::{Decoder, DeltaBitPackDecoder}; +use crate::errors::{ParquetError, Result}; + +/// Decoder for `Encoding::DELTA_BYTE_ARRAY` +pub struct DeltaByteArrayDecoder { + prefix_lengths: Vec, + suffix_lengths: Vec, + data: Bytes, + length_offset: usize, + data_offset: usize, + last_value: Vec, +} + +impl DeltaByteArrayDecoder { + /// Create a new [`DeltaByteArrayDecoder`] with the provided data page + pub fn new(data: Bytes) -> Result { + let mut prefix = DeltaBitPackDecoder::::new(); + prefix.set_data(data.clone(), 0)?; + + let num_prefix = prefix.values_left(); + let mut prefix_lengths = vec![0; num_prefix]; + assert_eq!(prefix.get(&mut prefix_lengths)?, num_prefix); + + let mut suffix = DeltaBitPackDecoder::::new(); + suffix.set_data(data.slice(prefix.get_offset()..), 0)?; + + let num_suffix = suffix.values_left(); + let mut suffix_lengths = vec![0; num_suffix]; + assert_eq!(suffix.get(&mut suffix_lengths)?, num_suffix); + + if num_prefix != num_suffix { + return Err(general_err!(format!( + "inconsistent DELTA_BYTE_ARRAY lengths, prefixes: {num_prefix}, suffixes: {num_suffix}" + ))); + } + + assert_eq!(prefix_lengths.len(), suffix_lengths.len()); + + Ok(Self { + prefix_lengths, + suffix_lengths, + data, + length_offset: 0, + data_offset: prefix.get_offset() + suffix.get_offset(), + last_value: vec![], + }) + } + + /// Returns the number of values remaining + pub fn remaining(&self) -> usize { + self.prefix_lengths.len() - self.length_offset + } + + /// Read up to `len` values, returning the number of values read + /// and calling `f` with each decoded byte slice + /// + /// Will short-circuit and return on error + pub fn read(&mut self, len: usize, mut f: F) -> Result + where + F: FnMut(&[u8]) -> Result<()>, + { + let to_read = len.min(self.remaining()); + + let length_range = self.length_offset..self.length_offset + to_read; + let iter = self.prefix_lengths[length_range.clone()] + .iter() + .zip(&self.suffix_lengths[length_range]); + + let data = self.data.as_ref(); + + for (prefix_length, suffix_length) in iter { + let prefix_length = *prefix_length as usize; + let suffix_length = *suffix_length as usize; + + if self.data_offset + suffix_length > self.data.len() { + return Err(ParquetError::EOF("eof decoding byte array".into())); + } + + self.last_value.truncate(prefix_length); + self.last_value + .extend_from_slice(&data[self.data_offset..self.data_offset + suffix_length]); + f(&self.last_value)?; + + self.data_offset += suffix_length; + } + + self.length_offset += to_read; + Ok(to_read) + } + + /// Skip up to `to_skip` values, returning the number of values skipped + pub fn skip(&mut self, to_skip: usize) -> Result { + let to_skip = to_skip.min(self.prefix_lengths.len() - self.length_offset); + + let length_range = self.length_offset..self.length_offset + to_skip; + let iter = self.prefix_lengths[length_range.clone()] + .iter() + .zip(&self.suffix_lengths[length_range]); + + let data = self.data.as_ref(); + + for (prefix_length, suffix_length) in iter { + let prefix_length = *prefix_length as usize; + let suffix_length = *suffix_length as usize; + + if self.data_offset + suffix_length > self.data.len() { + return Err(ParquetError::EOF("eof decoding byte array".into())); + } + + self.last_value.truncate(prefix_length); + self.last_value + .extend_from_slice(&data[self.data_offset..self.data_offset + suffix_length]); + self.data_offset += suffix_length; + } + self.length_offset += to_skip; + Ok(to_skip) + } +} diff --git a/vendor/parquet/src/arrow/decoder/dictionary_index.rs b/vendor/parquet/src/arrow/decoder/dictionary_index.rs new file mode 100644 index 0000000..7a4b77f --- /dev/null +++ b/vendor/parquet/src/arrow/decoder/dictionary_index.rs @@ -0,0 +1,121 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use bytes::Bytes; + +use crate::encodings::rle::RleDecoder; +use crate::errors::Result; + +/// Decoder for `Encoding::RLE_DICTIONARY` indices +pub struct DictIndexDecoder { + /// Decoder for the dictionary offsets array + decoder: RleDecoder, + + /// We want to decode the offsets in chunks so we will maintain an internal buffer of decoded + /// offsets + index_buf: Box<[i32; 1024]>, + /// Current length of `index_buf` + index_buf_len: usize, + /// Current offset into `index_buf`. If `index_buf_offset` == `index_buf_len` then we've consumed + /// the entire buffer and need to decode another chunk of offsets. + index_offset: usize, + + /// This is a maximum as the null count is not always known, e.g. value data from + /// a v1 data page + max_remaining_values: usize, +} + +impl DictIndexDecoder { + /// Create a new [`DictIndexDecoder`] with the provided data page, the number of levels + /// associated with this data page, and the number of non-null values (if known) + pub fn new(data: Bytes, num_levels: usize, num_values: Option) -> Result { + let bit_width = data[0]; + let mut decoder = RleDecoder::new(bit_width); + decoder.set_data(data.slice(1..))?; + + Ok(Self { + decoder, + index_buf: Box::new([0; 1024]), + index_buf_len: 0, + index_offset: 0, + max_remaining_values: num_values.unwrap_or(num_levels), + }) + } + + /// Read up to `len` values, returning the number of values read + /// and calling `f` with each decoded dictionary index + /// + /// Will short-circuit and return on error + #[inline(always)] + pub fn read Result<()>>(&mut self, len: usize, mut f: F) -> Result { + let total_to_read = len.min(self.max_remaining_values); + + let mut values_read = 0; + + let index_buf = self.index_buf.as_mut(); + while values_read < total_to_read { + if self.index_offset == self.index_buf_len { + // We've consumed the entire index buffer so we need to reload it before proceeding + let read = self.decoder.get_batch(index_buf)?; + if read == 0 { + break; + } + self.index_buf_len = read; + self.index_offset = 0; + } + + let available = self.index_buf_len - self.index_offset; + let n = available.min(total_to_read - values_read); + + f(&index_buf[self.index_offset..self.index_offset + n])?; + + self.index_offset += n; + values_read += n; + } + self.max_remaining_values -= values_read; + + Ok(values_read) + } + + /// Skip up to `to_skip` values, returning the number of values skipped + pub fn skip(&mut self, to_skip: usize) -> Result { + let to_skip = to_skip.min(self.max_remaining_values); + + let mut values_skip = 0; + while values_skip < to_skip { + if self.index_offset == self.index_buf_len { + // Instead of reloading the buffer, just skip in the decoder + let skip = self.decoder.skip(to_skip - values_skip)?; + + if skip == 0 { + break; + } + + self.max_remaining_values -= skip; + values_skip += skip; + } else { + // We still have indices buffered, so skip within the buffer + let skip = (to_skip - values_skip).min(self.index_buf_len - self.index_offset); + + self.index_offset += skip; + self.max_remaining_values -= skip; + values_skip += skip; + } + } + Ok(values_skip) + } +} diff --git a/vendor/parquet/src/arrow/decoder/mod.rs b/vendor/parquet/src/arrow/decoder/mod.rs new file mode 100644 index 0000000..dc1000f --- /dev/null +++ b/vendor/parquet/src/arrow/decoder/mod.rs @@ -0,0 +1,24 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Specialized decoders optimised for decoding to arrow format + +mod delta_byte_array; +mod dictionary_index; + +pub use delta_byte_array::DeltaByteArrayDecoder; +pub use dictionary_index::DictIndexDecoder; diff --git a/vendor/parquet/src/arrow/in_memory_row_group.rs b/vendor/parquet/src/arrow/in_memory_row_group.rs new file mode 100644 index 0000000..6c5f013 --- /dev/null +++ b/vendor/parquet/src/arrow/in_memory_row_group.rs @@ -0,0 +1,317 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::ProjectionMask; +use crate::arrow::array_reader::RowGroups; +use crate::arrow::arrow_reader::RowSelection; +use crate::column::page::{PageIterator, PageReader}; +use crate::errors::ParquetError; +use crate::file::metadata::{ParquetMetaData, RowGroupMetaData}; +use crate::file::page_index::offset_index::OffsetIndexMetaData; +use crate::file::reader::{ChunkReader, Length, SerializedPageReader}; +use bytes::{Buf, Bytes}; +use std::ops::Range; +use std::sync::Arc; + +/// An in-memory collection of column chunks +#[derive(Debug)] +pub(crate) struct InMemoryRowGroup<'a> { + pub(crate) offset_index: Option<&'a [OffsetIndexMetaData]>, + /// Column chunks for this row group + pub(crate) column_chunks: Vec>>, + pub(crate) row_count: usize, + pub(crate) row_group_idx: usize, + pub(crate) metadata: &'a ParquetMetaData, +} + +/// What ranges to fetch for the columns in this row group +#[derive(Debug)] +pub(crate) struct FetchRanges { + /// The byte ranges to fetch + pub(crate) ranges: Vec>, + /// If `Some`, the start offsets of each page for each column chunk + pub(crate) page_start_offsets: Option>>, +} + +impl InMemoryRowGroup<'_> { + /// Returns the byte ranges to fetch for the columns specified in + /// `projection` and `selection`. + /// + /// `cache_mask` indicates which columns, if any, are being cached by + /// [`RowGroupCache`](crate::arrow::array_reader::RowGroupCache). + /// The `selection` for Cached columns is expanded to batch boundaries to simplify + /// accounting for what data is cached. + pub(crate) fn fetch_ranges( + &self, + projection: &ProjectionMask, + selection: Option<&RowSelection>, + batch_size: usize, + cache_mask: Option<&ProjectionMask>, + ) -> FetchRanges { + let metadata = self.metadata.row_group(self.row_group_idx); + if let Some((selection, offset_index)) = selection.zip(self.offset_index) { + let expanded_selection = + selection.expand_to_batch_boundaries(batch_size, self.row_count); + + // If we have a `RowSelection` and an `OffsetIndex` then only fetch + // pages required for the `RowSelection` + // Consider preallocating outer vec: https://github.com/apache/arrow-rs/issues/8667 + let mut page_start_offsets: Vec> = vec![]; + + let ranges = self + .column_chunks + .iter() + .zip(metadata.columns()) + .enumerate() + .filter(|&(idx, (chunk, _chunk_meta))| { + chunk.is_none() && projection.leaf_included(idx) + }) + .flat_map(|(idx, (_chunk, chunk_meta))| { + // If the first page does not start at the beginning of the column, + // then we need to also fetch a dictionary page. + let mut ranges: Vec> = vec![]; + let (start, _len) = chunk_meta.byte_range(); + match offset_index[idx].page_locations.first() { + Some(first) if first.offset as u64 != start => { + ranges.push(start..first.offset as u64); + } + _ => (), + } + + // Expand selection to batch boundaries if needed for caching + // (see doc comment for this function for details on `cache_mask`) + let use_expanded = cache_mask.map(|m| m.leaf_included(idx)).unwrap_or(false); + if use_expanded { + ranges.extend( + expanded_selection.scan_ranges(&offset_index[idx].page_locations), + ); + } else { + ranges.extend(selection.scan_ranges(&offset_index[idx].page_locations)); + } + page_start_offsets.push(ranges.iter().map(|range| range.start).collect()); + + ranges + }) + .collect(); + FetchRanges { + ranges, + page_start_offsets: Some(page_start_offsets), + } + } else { + let ranges = self + .column_chunks + .iter() + .enumerate() + .filter(|&(idx, chunk)| chunk.is_none() && projection.leaf_included(idx)) + .map(|(idx, _chunk)| { + let column = metadata.column(idx); + let (start, length) = column.byte_range(); + start..(start + length) + }) + .collect(); + FetchRanges { + ranges, + page_start_offsets: None, + } + } + } + + /// Fills in `self.column_chunks` with the data fetched from `chunk_data`. + /// + /// This function **must** be called with the data from the ranges returned by + /// `fetch_ranges` and the corresponding page_start_offsets, with the exact same and `selection`. + pub(crate) fn fill_column_chunks( + &mut self, + projection: &ProjectionMask, + page_start_offsets: Option>>, + chunk_data: I, + ) where + I: IntoIterator, + { + let mut chunk_data = chunk_data.into_iter(); + let metadata = self.metadata.row_group(self.row_group_idx); + if let Some(page_start_offsets) = page_start_offsets { + // If we have a `RowSelection` and an `OffsetIndex` then only fetch pages required for the + // `RowSelection` + let mut page_start_offsets = page_start_offsets.into_iter(); + + for (idx, chunk) in self.column_chunks.iter_mut().enumerate() { + if chunk.is_some() || !projection.leaf_included(idx) { + continue; + } + + if let Some(offsets) = page_start_offsets.next() { + let mut chunks = Vec::with_capacity(offsets.len()); + for _ in 0..offsets.len() { + chunks.push(chunk_data.next().unwrap()); + } + + *chunk = Some(Arc::new(ColumnChunkData::Sparse { + length: metadata.column(idx).byte_range().1 as usize, + data: offsets + .into_iter() + .map(|x| x as usize) + .zip(chunks) + .collect(), + })) + } + } + } else { + for (idx, chunk) in self.column_chunks.iter_mut().enumerate() { + if chunk.is_some() || !projection.leaf_included(idx) { + continue; + } + + if let Some(data) = chunk_data.next() { + *chunk = Some(Arc::new(ColumnChunkData::Dense { + offset: metadata.column(idx).byte_range().0 as usize, + data, + })); + } + } + } + } +} + +impl RowGroups for InMemoryRowGroup<'_> { + fn num_rows(&self) -> usize { + self.row_count + } + + /// Return chunks for column i + fn column_chunks(&self, i: usize) -> crate::errors::Result> { + match &self.column_chunks[i] { + None => Err(ParquetError::General(format!( + "Invalid column index {i}, column was not fetched" + ))), + Some(data) => { + let page_locations = self + .offset_index + // filter out empty offset indexes (old versions specified Some(vec![]) when no present) + .filter(|index| !index.is_empty()) + .map(|index| index[i].page_locations.clone()); + let column_chunk_metadata = self.metadata.row_group(self.row_group_idx).column(i); + let page_reader = SerializedPageReader::new( + data.clone(), + column_chunk_metadata, + self.row_count, + page_locations, + )?; + let page_reader = page_reader.add_crypto_context( + self.row_group_idx, + i, + self.metadata, + column_chunk_metadata, + )?; + + let page_reader: Box = Box::new(page_reader); + + Ok(Box::new(ColumnChunkIterator { + reader: Some(Ok(page_reader)), + })) + } + } + } + + fn row_groups(&self) -> Box + '_> { + Box::new(std::iter::once(self.metadata.row_group(self.row_group_idx))) + } + + fn metadata(&self) -> &ParquetMetaData { + self.metadata + } +} + +/// An in-memory column chunk. +/// This allows us to hold either dense column chunks or sparse column chunks and easily +/// access them by offset. +#[derive(Clone, Debug)] +pub(crate) enum ColumnChunkData { + /// Column chunk data representing only a subset of data pages. + /// For example if a row selection (possibly caused by a filter in a query) causes us to read only + /// a subset of the rows in the column. + Sparse { + /// Length of the full column chunk + length: usize, + /// Subset of data pages included in this sparse chunk. + /// + /// Each element is a tuple of (page offset within file, page data). + /// Each entry is a complete page and the list is ordered by offset. + data: Vec<(usize, Bytes)>, + }, + /// Full column chunk and the offset within the original file + Dense { offset: usize, data: Bytes }, +} + +impl ColumnChunkData { + /// Return the data for this column chunk at the given offset + fn get(&self, start: u64) -> crate::errors::Result { + match &self { + ColumnChunkData::Sparse { data, .. } => data + .binary_search_by_key(&start, |(offset, _)| *offset as u64) + .map(|idx| data[idx].1.clone()) + .map_err(|_| { + ParquetError::General(format!( + "Invalid offset in sparse column chunk data: {start}, no matching page found.\ + If you are using a `SelectionStrategyPolicy::Mask`, ensure that the OffsetIndex is provided when \ + creating the InMemoryRowGroup." + )) + }), + ColumnChunkData::Dense { offset, data } => { + let start = start as usize - *offset; + Ok(data.slice(start..)) + } + } + } +} + +impl Length for ColumnChunkData { + /// Return the total length of the full column chunk + fn len(&self) -> u64 { + match &self { + ColumnChunkData::Sparse { length, .. } => *length as u64, + ColumnChunkData::Dense { data, .. } => data.len() as u64, + } + } +} + +impl ChunkReader for ColumnChunkData { + type T = bytes::buf::Reader; + + fn get_read(&self, start: u64) -> crate::errors::Result { + Ok(self.get(start)?.reader()) + } + + fn get_bytes(&self, start: u64, length: usize) -> crate::errors::Result { + Ok(self.get(start)?.slice(..length)) + } +} + +/// Implements [`PageIterator`] for a single column chunk, yielding a single [`PageReader`] +struct ColumnChunkIterator { + reader: Option>>, +} + +impl Iterator for ColumnChunkIterator { + type Item = crate::errors::Result>; + + fn next(&mut self) -> Option { + self.reader.take() + } +} + +impl PageIterator for ColumnChunkIterator {} diff --git a/vendor/parquet/src/arrow/mod.rs b/vendor/parquet/src/arrow/mod.rs new file mode 100644 index 0000000..b89db36 --- /dev/null +++ b/vendor/parquet/src/arrow/mod.rs @@ -0,0 +1,1072 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! API for reading/writing Arrow [`RecordBatch`]es and [`Array`]s to/from +//! Parquet Files. +//! +//! See the [crate-level documentation](crate) for more details on other APIs +//! +//! # Schema Conversion +//! +//! These APIs ensure that data in Arrow [`RecordBatch`]es written to Parquet are +//! read back as [`RecordBatch`]es with the exact same types and values. +//! +//! Parquet and Arrow have different type systems, and there is not +//! always a one to one mapping between the systems. For example, data +//! stored as a Parquet [`BYTE_ARRAY`] can be read as either an Arrow +//! [`BinaryViewArray`] or [`BinaryArray`]. +//! +//! To recover the original Arrow types, the writers in this module add a "hint" to +//! the metadata in the [`ARROW_SCHEMA_META_KEY`] key which records the original Arrow +//! schema. The metadata hint follows the same convention as arrow-cpp based +//! implementations such as `pyarrow`. The reader looks for the schema hint in the +//! metadata to determine Arrow types, and if it is not present, infers the Arrow schema +//! from the Parquet schema. +//! +//! In situations where the embedded Arrow schema is not compatible with the Parquet +//! schema, the Parquet schema takes precedence and no error is raised. +//! See [#1663](https://github.com/apache/arrow-rs/issues/1663) +//! +//! You can also control the type conversion process in more detail using: +//! +//! * [`ArrowSchemaConverter`] control the conversion of Arrow types to Parquet +//! types. +//! +//! * [`ArrowReaderOptions::with_schema`] to explicitly specify your own Arrow schema hint +//! to use when reading Parquet, overriding any metadata that may be present. +//! +//! [`RecordBatch`]: arrow_array::RecordBatch +//! [`Array`]: arrow_array::Array +//! [`BYTE_ARRAY`]: crate::basic::Type::BYTE_ARRAY +//! [`BinaryViewArray`]: arrow_array::BinaryViewArray +//! [`BinaryArray`]: arrow_array::BinaryArray +//! [`ArrowReaderOptions::with_schema`]: arrow_reader::ArrowReaderOptions::with_schema +//! +//! # Example: Writing Arrow `RecordBatch` to Parquet file +//! +//!```rust +//! # use arrow_array::{Int32Array, ArrayRef}; +//! # use arrow_array::RecordBatch; +//! # use parquet::arrow::arrow_writer::ArrowWriter; +//! # use parquet::file::properties::WriterProperties; +//! # use tempfile::tempfile; +//! # use std::sync::Arc; +//! # use parquet::basic::Compression; +//! let ids = Int32Array::from(vec![1, 2, 3, 4]); +//! let vals = Int32Array::from(vec![5, 6, 7, 8]); +//! let batch = RecordBatch::try_from_iter(vec![ +//! ("id", Arc::new(ids) as ArrayRef), +//! ("val", Arc::new(vals) as ArrayRef), +//! ]).unwrap(); +//! +//! let file = tempfile().unwrap(); +//! +//! // WriterProperties can be used to set Parquet file options +//! let props = WriterProperties::builder() +//! .set_compression(Compression::SNAPPY) +//! .build(); +//! +//! let mut writer = ArrowWriter::try_new(file, batch.schema(), Some(props)).unwrap(); +//! +//! writer.write(&batch).expect("Writing batch"); +//! +//! // writer must be closed to write footer +//! writer.close().unwrap(); +//! ``` +//! +//! # Example: Reading Parquet file into Arrow `RecordBatch` +//! +//! ```rust +//! # use std::fs::File; +//! # use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +//! # use std::sync::Arc; +//! # use arrow_array::Int32Array; +//! # use arrow::datatypes::{DataType, Field, Schema}; +//! # use arrow_array::RecordBatch; +//! # use parquet::arrow::arrow_writer::ArrowWriter; +//! # +//! # let ids = Int32Array::from(vec![1, 2, 3, 4]); +//! # let schema = Arc::new(Schema::new(vec![ +//! # Field::new("id", DataType::Int32, false), +//! # ])); +//! # +//! # let file = File::create("data.parquet").unwrap(); +//! # +//! # let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(ids)]).unwrap(); +//! # let batches = vec![batch]; +//! # +//! # let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None).unwrap(); +//! # +//! # for batch in batches { +//! # writer.write(&batch).expect("Writing batch"); +//! # } +//! # writer.close().unwrap(); +//! # +//! let file = File::open("data.parquet").unwrap(); +//! +//! let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); +//! println!("Converted arrow schema is: {}", builder.schema()); +//! +//! let mut reader = builder.build().unwrap(); +//! +//! let record_batch = reader.next().unwrap().unwrap(); +//! +//! println!("Read {} records.", record_batch.num_rows()); +//! ``` +//! +//! # Example: Reading non-uniformly encrypted parquet file into arrow record batch +//! +//! Note: This requires the experimental `encryption` feature to be enabled at compile time. +//! +#![cfg_attr(feature = "encryption", doc = "```rust")] +#![cfg_attr(not(feature = "encryption"), doc = "```ignore")] +//! # use arrow_array::{Int32Array, ArrayRef}; +//! # use arrow_array::{types, RecordBatch}; +//! # use parquet::arrow::arrow_reader::{ +//! # ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReaderBuilder, +//! # }; +//! # use arrow_array::cast::AsArray; +//! # use parquet::file::metadata::ParquetMetaData; +//! # use tempfile::tempfile; +//! # use std::fs::File; +//! # use parquet::encryption::decrypt::FileDecryptionProperties; +//! # let test_data = arrow::util::test_util::parquet_test_data(); +//! # let path = format!("{test_data}/encrypt_columns_and_footer.parquet.encrypted"); +//! # +//! let file = File::open(path).unwrap(); +//! +//! // Define the AES encryption keys required required for decrypting the footer metadata +//! // and column-specific data. If only a footer key is used then it is assumed that the +//! // file uses uniform encryption and all columns are encrypted with the footer key. +//! // If any column keys are specified, other columns without a key provided are assumed +//! // to be unencrypted +//! let footer_key = "0123456789012345".as_bytes(); // Keys are 128 bits (16 bytes) +//! let column_1_key = "1234567890123450".as_bytes(); +//! let column_2_key = "1234567890123451".as_bytes(); +//! +//! let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) +//! .with_column_key("double_field", column_1_key.to_vec()) +//! .with_column_key("float_field", column_2_key.to_vec()) +//! .build() +//! .unwrap(); +//! +//! let options = ArrowReaderOptions::default() +//! .with_file_decryption_properties(decryption_properties); +//! let reader_metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); +//! let file_metadata = reader_metadata.metadata().file_metadata(); +//! assert_eq!(50, file_metadata.num_rows()); +//! +//! let mut reader = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options) +//! .unwrap() +//! .build() +//! .unwrap(); +//! +//! let record_batch = reader.next().unwrap().unwrap(); +//! assert_eq!(50, record_batch.num_rows()); +//! ``` + +experimental!(mod array_reader); +pub mod arrow_reader; +pub mod arrow_writer; +mod buffer; +mod decoder; + +#[cfg(feature = "async")] +pub mod async_reader; +#[cfg(feature = "async")] +pub mod async_writer; + +pub mod push_decoder; + +mod in_memory_row_group; +mod record_reader; + +experimental!(mod schema); + +use std::fmt::Debug; + +pub use self::arrow_writer::ArrowWriter; +#[cfg(feature = "async")] +pub use self::async_reader::ParquetRecordBatchStreamBuilder; +#[cfg(feature = "async")] +pub use self::async_writer::AsyncArrowWriter; +use crate::schema::types::SchemaDescriptor; +use arrow_schema::{FieldRef, Schema}; + +pub use self::schema::{ + ArrowSchemaConverter, FieldLevels, add_encoded_arrow_schema_to_metadata, encode_arrow_schema, + parquet_to_arrow_field_levels, parquet_to_arrow_field_levels_with_virtual, + parquet_to_arrow_schema, parquet_to_arrow_schema_by_columns, virtual_type::*, +}; + +/// Schema metadata key used to store serialized Arrow schema +/// +/// The Arrow schema is encoded using the Arrow IPC format, and then base64 +/// encoded. This is the same format used by arrow-cpp systems, such as pyarrow. +pub const ARROW_SCHEMA_META_KEY: &str = "ARROW:schema"; + +/// The value of this metadata key, if present on [`Field::metadata`], will be used +/// to populate [`BasicTypeInfo::id`] +/// +/// [`Field::metadata`]: arrow_schema::Field::metadata +/// [`BasicTypeInfo::id`]: crate::schema::types::BasicTypeInfo::id +pub const PARQUET_FIELD_ID_META_KEY: &str = "PARQUET:field_id"; + +/// A [`ProjectionMask`] identifies a set of columns within a potentially nested schema to project +/// +/// In particular, a [`ProjectionMask`] can be constructed from a list of leaf column indices +/// or root column indices where: +/// +/// * Root columns are the direct children of the root schema, enumerated in order +/// * Leaf columns are the child-less leaves of the schema as enumerated by a depth-first search +/// +/// For example, the schema +/// +/// ```ignore +/// message schema { +/// REQUIRED boolean leaf_1; +/// REQUIRED GROUP group { +/// OPTIONAL int32 leaf_2; +/// OPTIONAL int64 leaf_3; +/// } +/// } +/// ``` +/// +/// Has roots `["leaf_1", "group"]` and leaves `["leaf_1", "leaf_2", "leaf_3"]` +/// +/// For non-nested schemas, i.e. those containing only primitive columns, the root +/// and leaves are the same +/// +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProjectionMask { + /// If `Some`, a leaf column should be included if the value at + /// the corresponding index is true + /// + /// If `None`, all columns should be included + /// + /// # Examples + /// + /// Given the original parquet schema with leaf columns is `[a, b, c, d]` + /// + /// A mask of `[true, false, true, false]` will result in a schema 2 + /// elements long: + /// * `fields[0]`: `a` + /// * `fields[1]`: `c` + /// + /// A mask of `None` will result in a schema 4 elements long: + /// * `fields[0]`: `a` + /// * `fields[1]`: `b` + /// * `fields[2]`: `c` + /// * `fields[3]`: `d` + mask: Option>, +} + +impl ProjectionMask { + /// Create a [`ProjectionMask`] which selects all columns + pub fn all() -> Self { + Self { mask: None } + } + + /// Create a [`ProjectionMask`] which selects no columns + pub fn none(len: usize) -> Self { + Self { + mask: Some(vec![false; len]), + } + } + + /// Create a [`ProjectionMask`] which selects only the specified leaf columns + /// + /// Note: repeated or out of order indices will not impact the final mask + /// + /// i.e. `[0, 1, 2]` will construct the same mask as `[1, 0, 0, 2]` + pub fn leaves(schema: &SchemaDescriptor, indices: impl IntoIterator) -> Self { + let mut mask = vec![false; schema.num_columns()]; + for leaf_idx in indices { + mask[leaf_idx] = true; + } + Self { mask: Some(mask) } + } + + /// Create a [`ProjectionMask`] which selects only the specified root columns + /// + /// Note: repeated or out of order indices will not impact the final mask + /// + /// i.e. `[0, 1, 2]` will construct the same mask as `[1, 0, 0, 2]` + pub fn roots(schema: &SchemaDescriptor, indices: impl IntoIterator) -> Self { + let num_root_columns = schema.root_schema().get_fields().len(); + let mut root_mask = vec![false; num_root_columns]; + for root_idx in indices { + root_mask[root_idx] = true; + } + + let mask = (0..schema.num_columns()) + .map(|leaf_idx| { + let root_idx = schema.get_column_root_idx(leaf_idx); + root_mask[root_idx] + }) + .collect(); + + Self { mask: Some(mask) } + } + + /// Create a [`ProjectionMask`] which selects only the named columns + /// + /// All leaf columns that fall below a given name will be selected. For example, given + /// the schema + /// ```ignore + /// message schema { + /// OPTIONAL group a (MAP) { + /// REPEATED group key_value { + /// REQUIRED BYTE_ARRAY key (UTF8); // leaf index 0 + /// OPTIONAL group value (MAP) { + /// REPEATED group key_value { + /// REQUIRED INT32 key; // leaf index 1 + /// REQUIRED BOOLEAN value; // leaf index 2 + /// } + /// } + /// } + /// } + /// REQUIRED INT32 b; // leaf index 3 + /// REQUIRED DOUBLE c; // leaf index 4 + /// } + /// ``` + /// `["a.key_value.value", "c"]` would return leaf columns 1, 2, and 4. `["a"]` would return + /// columns 0, 1, and 2. + /// + /// Note: repeated or out of order indices will not impact the final mask. + /// + /// i.e. `["b", "c"]` will construct the same mask as `["c", "b", "c"]`. + /// + /// Also, this will not produce the desired results if a column contains a '.' in its name. + /// Use [`Self::leaves`] or [`Self::roots`] in that case. + pub fn columns<'a>( + schema: &SchemaDescriptor, + names: impl IntoIterator, + ) -> Self { + let mut mask = vec![false; schema.num_columns()]; + for name in names { + let name_path: Vec<&str> = name.split('.').collect(); + for (idx, col) in schema.columns().iter().enumerate() { + let path = col.path().parts(); + // searching for "a.b.c" cannot match "a.b" + if name_path.len() > path.len() { + continue; + } + // now path >= name_path, so check that each element in name_path matches + if name_path.iter().zip(path.iter()).all(|(a, b)| a == b) { + mask[idx] = true; + } + } + } + + Self { mask: Some(mask) } + } + + /// Returns true if the leaf column `leaf_idx` is included by the mask + pub fn leaf_included(&self, leaf_idx: usize) -> bool { + self.mask.as_ref().map(|m| m[leaf_idx]).unwrap_or(true) + } + + /// Union two projection masks + /// + /// Example: + /// ```text + /// mask1 = [true, false, true] + /// mask2 = [false, true, true] + /// union(mask1, mask2) = [true, true, true] + /// ``` + pub fn union(&mut self, other: &Self) { + match (self.mask.as_ref(), other.mask.as_ref()) { + (None, _) | (_, None) => self.mask = None, + (Some(a), Some(b)) => { + debug_assert_eq!(a.len(), b.len()); + let mask = a.iter().zip(b.iter()).map(|(&a, &b)| a || b).collect(); + self.mask = Some(mask); + } + } + } + + /// Intersect two projection masks + /// + /// Example: + /// ```text + /// mask1 = [true, false, true] + /// mask2 = [false, true, true] + /// intersect(mask1, mask2) = [false, false, true] + /// ``` + pub fn intersect(&mut self, other: &Self) { + match (self.mask.as_ref(), other.mask.as_ref()) { + (None, _) => self.mask = other.mask.clone(), + (_, None) => {} + (Some(a), Some(b)) => { + debug_assert_eq!(a.len(), b.len()); + let mask = a.iter().zip(b.iter()).map(|(&a, &b)| a && b).collect(); + self.mask = Some(mask); + } + } + } + + /// Return a new [`ProjectionMask`] that excludes any leaf columns that are + /// part of a nested type, such as struct, list, or map + /// + /// If there are no non-nested columns in the mask, returns `None` + pub(crate) fn without_nested_types(&self, schema: &SchemaDescriptor) -> Option { + let num_leaves = schema.num_columns(); + + // Count how many leaves each root column has + let num_roots = schema.root_schema().get_fields().len(); + let mut root_leaf_counts = vec![0usize; num_roots]; + for leaf_idx in 0..num_leaves { + let root_idx = schema.get_column_root_idx(leaf_idx); + root_leaf_counts[root_idx] += 1; + } + + // Cache only top-level primitive columns. + // Even a one-leaf group is nested; caching it drops parent def levels. + let mut included_leaves = Vec::new(); + for leaf_idx in 0..num_leaves { + if self.leaf_included(leaf_idx) { + let root = schema.get_column_root(leaf_idx); + let root_idx = schema.get_column_root_idx(leaf_idx); + if root_leaf_counts[root_idx] == 1 && root.is_primitive() { + included_leaves.push(leaf_idx); + } + } + } + + if included_leaves.is_empty() { + None + } else { + Some(ProjectionMask::leaves(schema, included_leaves)) + } + } +} + +/// Lookups up the parquet column by name +/// +/// Returns the parquet column index and the corresponding arrow field +pub fn parquet_column<'a>( + parquet_schema: &SchemaDescriptor, + arrow_schema: &'a Schema, + name: &str, +) -> Option<(usize, &'a FieldRef)> { + let (root_idx, field) = arrow_schema.fields.find(name)?; + if field.data_type().is_nested() { + // Nested fields are not supported and require non-trivial logic + // to correctly walk the parquet schema accounting for the + // logical type rules - + // + // For example a ListArray could correspond to anything from 1 to 3 levels + // in the parquet schema + return None; + } + + // This could be made more efficient (#TBD) + let parquet_idx = (0..parquet_schema.columns().len()) + .find(|x| parquet_schema.get_column_root_idx(*x) == root_idx)?; + Some((parquet_idx, field)) +} + +#[cfg(test)] +mod test { + use crate::arrow::ArrowWriter; + use crate::file::metadata::{ + PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions, ParquetMetaDataReader, + ParquetMetaDataWriter, + }; + use crate::file::properties::{EnabledStatistics, WriterProperties}; + use crate::schema::parser::parse_message_type; + use crate::schema::types::SchemaDescriptor; + use arrow_array::{ArrayRef, Int32Array, RecordBatch}; + use bytes::Bytes; + use std::sync::Arc; + + use super::ProjectionMask; + + #[test] + // Reproducer for https://github.com/apache/arrow-rs/issues/6464 + fn test_metadata_read_write_partial_offset() { + let parquet_bytes = create_parquet_file(); + + // read the metadata from the file WITHOUT the page index structures + let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false); + let original_metadata = ParquetMetaDataReader::new() + .with_metadata_options(Some(options)) + .parse_and_finish(&parquet_bytes) + .unwrap(); + + // this should error because the page indexes are not present, but have offsets specified + let metadata_bytes = metadata_to_bytes(&original_metadata); + let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false); + let err = ParquetMetaDataReader::new() + .with_metadata_options(Some(options)) + .with_page_index_policy(PageIndexPolicy::Required) // there are no page indexes in the metadata + .parse_and_finish(&metadata_bytes) + .err() + .unwrap(); + assert_eq!( + err.to_string(), + "EOF: Parquet file too small. Page index range 82..115 overlaps with file metadata 0..357" + ); + } + + #[test] + fn test_metadata_read_write_roundtrip() { + let parquet_bytes = create_parquet_file(); + + // read the metadata from the file + let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false); + let original_metadata = ParquetMetaDataReader::new() + .with_metadata_options(Some(options)) + .parse_and_finish(&parquet_bytes) + .unwrap(); + + // read metadata back from the serialized bytes and ensure it is the same + let metadata_bytes = metadata_to_bytes(&original_metadata); + assert_ne!( + metadata_bytes.len(), + parquet_bytes.len(), + "metadata is subset of parquet" + ); + + let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false); + let roundtrip_metadata = ParquetMetaDataReader::new() + .with_metadata_options(Some(options)) + .parse_and_finish(&metadata_bytes) + .unwrap(); + + assert_eq!(original_metadata, roundtrip_metadata); + } + + #[test] + fn test_metadata_read_write_roundtrip_page_index() { + let parquet_bytes = create_parquet_file(); + + // read the metadata from the file including the page index structures + // (which are stored elsewhere in the footer) + let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false); + let original_metadata = ParquetMetaDataReader::new() + .with_metadata_options(Some(options)) + .with_page_index_policy(PageIndexPolicy::Required) + .parse_and_finish(&parquet_bytes) + .unwrap(); + + // read metadata back from the serialized bytes and ensure it is the same + let metadata_bytes = metadata_to_bytes(&original_metadata); + let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false); + let roundtrip_metadata = ParquetMetaDataReader::new() + .with_metadata_options(Some(options)) + .with_page_index_policy(PageIndexPolicy::Required) + .parse_and_finish(&metadata_bytes) + .unwrap(); + + // Need to normalize the metadata first to remove offsets in data + let original_metadata = normalize_locations(original_metadata); + let roundtrip_metadata = normalize_locations(roundtrip_metadata); + assert_eq!( + format!("{original_metadata:#?}"), + format!("{roundtrip_metadata:#?}") + ); + assert_eq!(original_metadata, roundtrip_metadata); + } + + /// Sets the page index offset locations in the metadata to `None` + /// + /// This is because the offsets are used to find the relative location of the index + /// structures, and thus differ depending on how the structures are stored. + fn normalize_locations(metadata: ParquetMetaData) -> ParquetMetaData { + let mut metadata_builder = metadata.into_builder(); + for rg in metadata_builder.take_row_groups() { + let mut rg_builder = rg.into_builder(); + for col in rg_builder.take_columns() { + rg_builder = rg_builder.add_column_metadata( + col.into_builder() + .set_offset_index_offset(None) + .set_index_page_offset(None) + .set_column_index_offset(None) + .build() + .unwrap(), + ); + } + let rg = rg_builder.build().unwrap(); + metadata_builder = metadata_builder.add_row_group(rg); + } + metadata_builder.build() + } + + /// Write a parquet filed into an in memory buffer + fn create_parquet_file() -> Bytes { + let mut buf = vec![]; + let data = vec![100, 200, 201, 300, 102, 33]; + let array: ArrayRef = Arc::new(Int32Array::from(data)); + let batch = RecordBatch::try_from_iter(vec![("id", array)]).unwrap(); + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .set_write_page_header_statistics(true) + .build(); + + let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + drop(writer); + + Bytes::from(buf) + } + + /// Serializes `ParquetMetaData` into a memory buffer, using `ParquetMetadataWriter + fn metadata_to_bytes(metadata: &ParquetMetaData) -> Bytes { + let mut buf = vec![]; + ParquetMetaDataWriter::new(&mut buf, metadata) + .finish() + .unwrap(); + Bytes::from(buf) + } + + #[test] + fn test_mask_from_column_names() { + let schema = parse_schema( + " + message test_schema { + OPTIONAL group a (MAP) { + REPEATED group key_value { + REQUIRED BYTE_ARRAY key (UTF8); + OPTIONAL group value (MAP) { + REPEATED group key_value { + REQUIRED INT32 key; + REQUIRED BOOLEAN value; + } + } + } + } + REQUIRED INT32 b; + REQUIRED DOUBLE c; + } + ", + ); + + let mask = ProjectionMask::columns(&schema, ["foo", "bar"]); + assert_eq!(mask.mask.unwrap(), vec![false; 5]); + + let mask = ProjectionMask::columns(&schema, []); + assert_eq!(mask.mask.unwrap(), vec![false; 5]); + + let mask = ProjectionMask::columns(&schema, ["a", "c"]); + assert_eq!(mask.mask.unwrap(), [true, true, true, false, true]); + + let mask = ProjectionMask::columns(&schema, ["a.key_value.key", "c"]); + assert_eq!(mask.mask.unwrap(), [true, false, false, false, true]); + + let mask = ProjectionMask::columns(&schema, ["a.key_value.value", "b"]); + assert_eq!(mask.mask.unwrap(), [false, true, true, true, false]); + + let schema = parse_schema( + " + message test_schema { + OPTIONAL group a (LIST) { + REPEATED group list { + OPTIONAL group element (LIST) { + REPEATED group list { + OPTIONAL group element (LIST) { + REPEATED group list { + OPTIONAL BYTE_ARRAY element (UTF8); + } + } + } + } + } + } + REQUIRED INT32 b; + } + ", + ); + + let mask = ProjectionMask::columns(&schema, ["a", "b"]); + assert_eq!(mask.mask.unwrap(), [true, true]); + + let mask = ProjectionMask::columns(&schema, ["a.list.element", "b"]); + assert_eq!(mask.mask.unwrap(), [true, true]); + + let mask = + ProjectionMask::columns(&schema, ["a.list.element.list.element.list.element", "b"]); + assert_eq!(mask.mask.unwrap(), [true, true]); + + let mask = ProjectionMask::columns(&schema, ["b"]); + assert_eq!(mask.mask.unwrap(), [false, true]); + + let schema = parse_schema( + " + message test_schema { + OPTIONAL INT32 a; + OPTIONAL INT32 b; + OPTIONAL INT32 c; + OPTIONAL INT32 d; + OPTIONAL INT32 e; + } + ", + ); + + let mask = ProjectionMask::columns(&schema, ["a", "b"]); + assert_eq!(mask.mask.unwrap(), [true, true, false, false, false]); + + let mask = ProjectionMask::columns(&schema, ["d", "b", "d"]); + assert_eq!(mask.mask.unwrap(), [false, true, false, true, false]); + + let schema = parse_schema( + " + message test_schema { + OPTIONAL INT32 a; + OPTIONAL INT32 b; + OPTIONAL INT32 a; + OPTIONAL INT32 d; + OPTIONAL INT32 e; + } + ", + ); + + let mask = ProjectionMask::columns(&schema, ["a", "e"]); + assert_eq!(mask.mask.unwrap(), [true, false, true, false, true]); + + let schema = parse_schema( + " + message test_schema { + OPTIONAL INT32 a; + OPTIONAL INT32 aa; + } + ", + ); + + let mask = ProjectionMask::columns(&schema, ["a"]); + assert_eq!(mask.mask.unwrap(), [true, false]); + } + + #[test] + fn test_projection_mask_union() { + let mut mask1 = ProjectionMask { + mask: Some(vec![true, false, true]), + }; + let mask2 = ProjectionMask { + mask: Some(vec![false, true, true]), + }; + mask1.union(&mask2); + assert_eq!(mask1.mask, Some(vec![true, true, true])); + + let mut mask1 = ProjectionMask { mask: None }; + let mask2 = ProjectionMask { + mask: Some(vec![false, true, true]), + }; + mask1.union(&mask2); + assert_eq!(mask1.mask, None); + + let mut mask1 = ProjectionMask { + mask: Some(vec![true, false, true]), + }; + let mask2 = ProjectionMask { mask: None }; + mask1.union(&mask2); + assert_eq!(mask1.mask, None); + + let mut mask1 = ProjectionMask { mask: None }; + let mask2 = ProjectionMask { mask: None }; + mask1.union(&mask2); + assert_eq!(mask1.mask, None); + } + + #[test] + fn test_projection_mask_intersect() { + let mut mask1 = ProjectionMask { + mask: Some(vec![true, false, true]), + }; + let mask2 = ProjectionMask { + mask: Some(vec![false, true, true]), + }; + mask1.intersect(&mask2); + assert_eq!(mask1.mask, Some(vec![false, false, true])); + + let mut mask1 = ProjectionMask { mask: None }; + let mask2 = ProjectionMask { + mask: Some(vec![false, true, true]), + }; + mask1.intersect(&mask2); + assert_eq!(mask1.mask, Some(vec![false, true, true])); + + let mut mask1 = ProjectionMask { + mask: Some(vec![true, false, true]), + }; + let mask2 = ProjectionMask { mask: None }; + mask1.intersect(&mask2); + assert_eq!(mask1.mask, Some(vec![true, false, true])); + + let mut mask1 = ProjectionMask { mask: None }; + let mask2 = ProjectionMask { mask: None }; + mask1.intersect(&mask2); + assert_eq!(mask1.mask, None); + } + + #[test] + fn test_projection_mask_without_nested_no_nested() { + // Schema with no nested types + let schema = parse_schema( + " + message test_schema { + OPTIONAL INT32 a; + OPTIONAL INT32 b; + REQUIRED DOUBLE d; + } + ", + ); + + let mask = ProjectionMask::all(); + // All columns are non-nested, but without_nested_types returns a new mask + assert_eq!( + Some(ProjectionMask::leaves(&schema, [0, 1, 2])), + mask.without_nested_types(&schema) + ); + + // select b, c + let mask = ProjectionMask::leaves(&schema, [1, 2]); + assert_eq!(Some(mask.clone()), mask.without_nested_types(&schema)); + } + + #[test] + fn test_projection_mask_without_nested_nested() { + // Schema with nested types (structs) + let schema = parse_schema( + " + message test_schema { + OPTIONAL INT32 a; + OPTIONAL group b { + REQUIRED INT32 b1; + OPTIONAL INT64 b2; + } + OPTIONAL group c (LIST) { + REPEATED group list { + OPTIONAL INT32 element; + } + } + REQUIRED DOUBLE d; + } + ", + ); + + // all leaves --> a, d + let mask = ProjectionMask::all(); + assert_eq!( + Some(ProjectionMask::leaves(&schema, [0, 4])), + mask.without_nested_types(&schema) + ); + + // b1 --> empty (it is nested) + let mask = ProjectionMask::leaves(&schema, [1]); + assert_eq!(None, mask.without_nested_types(&schema)); + + // b2, d --> d + let mask = ProjectionMask::leaves(&schema, [1, 4]); + assert_eq!( + Some(ProjectionMask::leaves(&schema, [4])), + mask.without_nested_types(&schema) + ); + + // element --> empty (it is nested) + let mask = ProjectionMask::leaves(&schema, [3]); + assert_eq!(None, mask.without_nested_types(&schema)); + } + + #[test] + fn test_projection_mask_without_nested_map_only() { + // Example from https://github.com/apache/parquet-format/blob/master/LogicalTypes.md + let schema = parse_schema( + " + message test_schema { + required group my_map (MAP) { + repeated group key_value { + required binary key (STRING); + optional int32 value; + } + } + } + ", + ); + + let mask = ProjectionMask::all(); + assert_eq!(None, mask.without_nested_types(&schema)); + + // key --> empty (it is nested) + let mask = ProjectionMask::leaves(&schema, [0]); + assert_eq!(None, mask.without_nested_types(&schema)); + + // value --> empty (it is nested) + let mask = ProjectionMask::leaves(&schema, [1]); + assert_eq!(None, mask.without_nested_types(&schema)); + } + + #[test] + fn test_projection_mask_without_nested_map_with_non_nested() { + // Example from https://github.com/apache/parquet-format/blob/master/LogicalTypes.md + // with an additional non-nested field + let schema = parse_schema( + " + message test_schema { + REQUIRED INT32 a; + required group my_map (MAP) { + repeated group key_value { + required binary key (STRING); + optional int32 value; + } + } + REQUIRED INT32 b; + } + ", + ); + + // all leaves --> a, b which are the only non nested ones + let mask = ProjectionMask::all(); + assert_eq!( + Some(ProjectionMask::leaves(&schema, [0, 3])), + mask.without_nested_types(&schema) + ); + + // key, value, b --> b (the only non-nested one) + let mask = ProjectionMask::leaves(&schema, [1, 2, 3]); + assert_eq!( + Some(ProjectionMask::leaves(&schema, [3])), + mask.without_nested_types(&schema) + ); + + // key, value --> NONE + let mask = ProjectionMask::leaves(&schema, [1, 2]); + assert_eq!(None, mask.without_nested_types(&schema)); + } + + #[test] + fn test_projection_mask_without_nested_deeply_nested() { + // Map of Maps + let schema = parse_schema( + " + message test_schema { + OPTIONAL group a (MAP) { + REPEATED group key_value { + REQUIRED BYTE_ARRAY key (UTF8); + OPTIONAL group value (MAP) { + REPEATED group key_value { + REQUIRED INT32 key; + REQUIRED BOOLEAN value; + } + } + } + } + REQUIRED INT32 b; + REQUIRED DOUBLE c; + ", + ); + + let mask = ProjectionMask::all(); + assert_eq!( + Some(ProjectionMask::leaves(&schema, [3, 4])), + mask.without_nested_types(&schema) + ); + + // (first) key, c --> c (the only non-nested one) + let mask = ProjectionMask::leaves(&schema, [0, 4]); + assert_eq!( + Some(ProjectionMask::leaves(&schema, [4])), + mask.without_nested_types(&schema) + ); + + // (second) key, value, b --> b (the only non-nested one) + let mask = ProjectionMask::leaves(&schema, [1, 2, 3]); + assert_eq!( + Some(ProjectionMask::leaves(&schema, [3])), + mask.without_nested_types(&schema) + ); + + // key --> NONE (the only non-nested one) + let mask = ProjectionMask::leaves(&schema, [0]); + assert_eq!(None, mask.without_nested_types(&schema)); + } + + #[test] + fn test_projection_mask_without_nested_list() { + // Example from https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#lists + let schema = parse_schema( + " + message test_schema { + required group my_list (LIST) { + repeated group list { + optional binary element (STRING); + } + } + REQUIRED INT32 b; + } + ", + ); + + let mask = ProjectionMask::all(); + assert_eq!( + Some(ProjectionMask::leaves(&schema, [1])), + mask.without_nested_types(&schema), + ); + + // element --> empty (it is nested) + let mask = ProjectionMask::leaves(&schema, [0]); + assert_eq!(None, mask.without_nested_types(&schema)); + + // element, b --> b (it is nested) + let mask = ProjectionMask::leaves(&schema, [0, 1]); + assert_eq!( + Some(ProjectionMask::leaves(&schema, [1])), + mask.without_nested_types(&schema), + ); + } + + #[test] + fn test_projection_mask_without_nested_single_leaf_struct() { + // Regression: a single-leaf struct is still nested. + let schema = parse_schema( + " + message test_schema { + OPTIONAL group address { + REQUIRED BYTE_ARRAY street (UTF8); + } + REQUIRED INT32 id; + } + ", + ); + + // street -> empty; root is a struct + let mask = ProjectionMask::leaves(&schema, [0]); + assert_eq!(None, mask.without_nested_types(&schema)); + + // street, id --> id only + let mask = ProjectionMask::leaves(&schema, [0, 1]); + assert_eq!( + Some(ProjectionMask::leaves(&schema, [1])), + mask.without_nested_types(&schema) + ); + + // all --> id only + let mask = ProjectionMask::all(); + assert_eq!( + Some(ProjectionMask::leaves(&schema, [1])), + mask.without_nested_types(&schema) + ); + } + + /// Converts a schema string into a `SchemaDescriptor` + fn parse_schema(schema: &str) -> SchemaDescriptor { + let parquet_group_type = parse_message_type(schema).unwrap(); + SchemaDescriptor::new(Arc::new(parquet_group_type)) + } +} diff --git a/vendor/parquet/src/arrow/push_decoder/mod.rs b/vendor/parquet/src/arrow/push_decoder/mod.rs new file mode 100644 index 0000000..ba075b4 --- /dev/null +++ b/vendor/parquet/src/arrow/push_decoder/mod.rs @@ -0,0 +1,2411 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`ParquetPushDecoder`]: decodes Parquet data with data provided by the +//! caller (rather than from an underlying reader). + +mod reader_builder; +mod remaining; + +use crate::DecodeResult; +use crate::arrow::arrow_reader::{ + ArrowReaderBuilder, ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReader, +}; +use crate::errors::ParquetError; +use crate::file::metadata::ParquetMetaData; +pub use crate::util::push_buffers::PushBuffers; +use arrow_array::RecordBatch; +use bytes::Bytes; +use reader_builder::{RowBudget, RowGroupReaderBuilder, RowGroupReaderBuilderParts}; +use remaining::{RemainingRowGroups, RemainingRowGroupsParts}; +use std::ops::Range; +use std::sync::Arc; + +/// A builder for [`ParquetPushDecoder`]. +/// +/// To create a new decoder, use [`ParquetPushDecoderBuilder::try_new_decoder`]. +/// +/// You can decode the metadata from a Parquet file using either +/// [`ParquetMetadataReader`] or [`ParquetMetaDataPushDecoder`]. +/// +/// [`ParquetMetadataReader`]: crate::file::metadata::ParquetMetaDataReader +/// [`ParquetMetaDataPushDecoder`]: crate::file::metadata::ParquetMetaDataPushDecoder +/// +/// Note the "input" type is `u64` which represents the length of the Parquet file +/// being decoded. This is needed to initialize the internal buffers that track +/// what data has been provided to the decoder. +/// +/// # Example +/// ``` +/// # use std::ops::Range; +/// # use std::sync::Arc; +/// # use bytes::Bytes; +/// # use arrow_array::record_batch; +/// # use parquet::DecodeResult; +/// # use parquet::arrow::push_decoder::ParquetPushDecoderBuilder; +/// # use parquet::arrow::ArrowWriter; +/// # use parquet::file::metadata::ParquetMetaDataPushDecoder; +/// # let file_bytes = { +/// # let mut buffer = vec![]; +/// # let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap(); +/// # let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap(); +/// # writer.write(&batch).unwrap(); +/// # writer.close().unwrap(); +/// # Bytes::from(buffer) +/// # }; +/// # // mimic IO by returning a function that returns the bytes for a given range +/// # let get_range = |range: &Range| -> Bytes { +/// # let start = range.start as usize; +/// # let end = range.end as usize; +/// # file_bytes.slice(start..end) +/// # }; +/// # let file_length = file_bytes.len() as u64; +/// # let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_length).unwrap(); +/// # metadata_decoder.push_ranges(vec![0..file_length], vec![file_bytes.clone()]).unwrap(); +/// # let DecodeResult::Data(parquet_metadata) = metadata_decoder.try_decode().unwrap() else { panic!("failed to decode metadata") }; +/// # let parquet_metadata = Arc::new(parquet_metadata); +/// // The file length and metadata are required to create the decoder +/// let mut decoder = +/// ParquetPushDecoderBuilder::try_new_decoder(parquet_metadata) +/// .unwrap() +/// // Optionally configure the decoder, e.g. batch size +/// .with_batch_size(1024) +/// // Build the decoder +/// .build() +/// .unwrap(); +/// +/// // In a loop, ask the decoder what it needs next, and provide it with the required data +/// loop { +/// match decoder.try_decode().unwrap() { +/// DecodeResult::NeedsData(ranges) => { +/// // The decoder needs more data. Fetch the data for the given ranges +/// let data = ranges.iter().map(|r| get_range(r)).collect::>(); +/// // Push the data to the decoder +/// decoder.push_ranges(ranges, data).unwrap(); +/// // After pushing the data, we can try to decode again on the next iteration +/// } +/// DecodeResult::Data(batch) => { +/// // Successfully decoded a batch of data +/// assert!(batch.num_rows() > 0); +/// } +/// DecodeResult::Finished => { +/// // The decoder has finished decoding exit the loop +/// break; +/// } +/// } +/// } +/// ``` +/// +/// # Adaptive scans +/// +/// The scan strategy is not fixed once [`build`](Self::build) is called: it +/// can be changed *while decoding*, at row-group boundaries. +/// +/// The important API for this is [`ParquetPushDecoder::try_next_reader`]. +/// Unlike [`try_decode`](ParquetPushDecoder::try_decode), which barrels +/// straight through row-group boundaries, `try_next_reader` returns once per +/// row group — leaving a clean window *between* row groups. At any such +/// boundary, [`ParquetPushDecoder::into_builder`] hands back a +/// `ParquetPushDecoderBuilder` for the row groups not yet decoded. Change any +/// option on it (projection, row filter, row selection policy, …) and +/// [`build`](Self::build) a fresh decoder that resumes from the next row +/// group. This is how a query engine promotes or demotes filters — for +/// example turning a row filter on or off — based on the selectivity observed +/// in the row groups decoded so far. +/// +/// ``` +/// # use std::ops::Range; +/// # use std::sync::Arc; +/// # use bytes::Bytes; +/// # use arrow_array::record_batch; +/// # use parquet::DecodeResult; +/// # use parquet::arrow::ProjectionMask; +/// # use parquet::arrow::push_decoder::ParquetPushDecoderBuilder; +/// # use parquet::arrow::ArrowWriter; +/// # use parquet::file::metadata::ParquetMetaDataPushDecoder; +/// # use parquet::file::properties::WriterProperties; +/// # let file_bytes = { +/// # let batch = record_batch!( +/// # ("a", Int32, [1, 2, 3, 4, 5, 6]), +/// # ("b", Int32, [6, 5, 4, 3, 2, 1]) +/// # ).unwrap(); +/// # // Small row groups so the test file has two of them. +/// # let props = WriterProperties::builder().set_max_row_group_row_count(Some(3)).build(); +/// # let mut buffer = vec![]; +/// # let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), Some(props)).unwrap(); +/// # writer.write(&batch).unwrap(); +/// # writer.close().unwrap(); +/// # Bytes::from(buffer) +/// # }; +/// # let get_range = |r: &Range| file_bytes.slice(r.start as usize..r.end as usize); +/// # let file_length = file_bytes.len() as u64; +/// # let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_length).unwrap(); +/// # metadata_decoder.push_ranges(vec![0..file_length], vec![file_bytes.clone()]).unwrap(); +/// # let DecodeResult::Data(parquet_metadata) = metadata_decoder.try_decode().unwrap() else { panic!() }; +/// # let parquet_metadata = Arc::new(parquet_metadata); +/// let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(parquet_metadata) +/// .unwrap() +/// .build() +/// .unwrap(); +/// +/// // Drive the decoder one row group at a time with `try_next_reader`. +/// loop { +/// match decoder.try_next_reader().unwrap() { +/// DecodeResult::NeedsData(ranges) => { +/// // Fetch and hand over the bytes the decoder asked for. +/// let data = ranges.iter().map(|r| get_range(r)).collect(); +/// decoder.push_ranges(ranges, data).unwrap(); +/// } +/// DecodeResult::Data(reader) => { +/// // Decode this row group's batches. +/// for batch in reader { +/// assert!(batch.unwrap().num_rows() > 0); +/// } +/// // We are now at a row-group boundary. Based on whatever stats +/// // were gathered, optionally change strategy for the row groups +/// // still to come: drop or promote a row filter, narrow or widen +/// // the projection, etc. +/// if decoder.is_at_row_group_boundary() && decoder.row_groups_remaining() > 0 { +/// let builder = decoder.into_builder().unwrap(); +/// // e.g. column "b" turned out not to be needed. +/// let projection = ProjectionMask::columns(builder.parquet_schema(), ["a"]); +/// decoder = builder.with_projection(projection).build().unwrap(); +/// } +/// } +/// DecodeResult::Finished => break, +/// } +/// } +/// ``` +pub type ParquetPushDecoderBuilder = ArrowReaderBuilder; + +/// The `input` of a [`ParquetPushDecoderBuilder`]. +/// +/// The shared [`ArrowReaderBuilder`] is generic over an `input`. The sync and +/// async builders read from a file or async reader; the push decoder has no +/// reader, so its input is the [`PushBuffers`] that caller-pushed bytes +/// accumulate in (empty for a fresh builder). +#[derive(Debug, Default)] +pub struct PushDecoderInput { + /// Bytes pushed into the decoder, awaiting decode. + buffers: PushBuffers, +} + +/// Methods for building a ParquetDecoder. See the base [`ArrowReaderBuilder`] for +/// more options that can be configured. +impl ParquetPushDecoderBuilder { + /// Create a new `ParquetDecoderBuilder` for configuring a Parquet decoder for the given file. + /// + /// See [`ParquetMetadataDecoder`] for a builder that can read the metadata from a Parquet file. + /// + /// [`ParquetMetadataDecoder`]: crate::file::metadata::ParquetMetaDataPushDecoder + /// + /// See example on [`ParquetPushDecoderBuilder`] + pub fn try_new_decoder(parquet_metadata: Arc) -> Result { + Self::try_new_decoder_with_options(parquet_metadata, ArrowReaderOptions::default()) + } + + /// Create a new `ParquetDecoderBuilder` for configuring a Parquet decoder for the given file + /// with the given reader options. + /// + /// This is similar to [`Self::try_new_decoder`] but allows configuring + /// options such as Arrow schema + pub fn try_new_decoder_with_options( + parquet_metadata: Arc, + arrow_reader_options: ArrowReaderOptions, + ) -> Result { + let arrow_reader_metadata = + ArrowReaderMetadata::try_new(parquet_metadata, arrow_reader_options)?; + Ok(Self::new_with_metadata(arrow_reader_metadata)) + } + + /// Create a new `ParquetDecoderBuilder` given [`ArrowReaderMetadata`]. + /// + /// See [`ArrowReaderMetadata::try_new`] for how to create the metadata from + /// the Parquet metadata and reader options. + pub fn new_with_metadata(arrow_reader_metadata: ArrowReaderMetadata) -> Self { + Self::new_builder(PushDecoderInput::default(), arrow_reader_metadata) + } + + /// Provide a preexisting [`PushBuffers`] for the built decoder to read + /// from, so bytes already fetched are not requested again. + pub fn with_buffers(self, buffers: PushBuffers) -> Self { + Self { + input: PushDecoderInput { buffers }, + ..self + } + } + + /// Create a [`ParquetPushDecoder`] with the configured options + pub fn build(self) -> Result { + let Self { + input: PushDecoderInput { buffers }, + metadata: parquet_metadata, + schema, + fields, + batch_size, + row_groups, + projection, + filter, + selection, + limit, + offset, + metrics, + row_selection_policy, + max_predicate_cache_size, + } = self; + + // If no row groups were specified, read all of them + let row_groups = + row_groups.unwrap_or_else(|| (0..parquet_metadata.num_row_groups()).collect()); + let has_predicates = filter + .as_ref() + .is_some_and(|filter| !filter.predicates.is_empty()); + + // Prepare to build RowGroup readers. `buffers` carries any bytes the + // caller already pushed (preserved across `into_builder`); a fresh + // builder supplies an empty `PushBuffers`. + let row_group_reader_builder = RowGroupReaderBuilder::new( + batch_size, + projection, + Arc::clone(&parquet_metadata), + fields, + filter, + metrics, + max_predicate_cache_size, + buffers, + row_selection_policy, + ); + + // Initialize the decoder with the configured options + let remaining_row_groups = RemainingRowGroups::new( + schema, + parquet_metadata, + row_groups, + selection, + RowBudget::new(offset, limit), + has_predicates, + row_group_reader_builder, + ); + + Ok(ParquetPushDecoder { + state: ParquetDecoderState::ReadingRowGroup { + remaining_row_groups: Box::new(remaining_row_groups), + }, + }) + } +} + +/// Reassemble a [`ParquetPushDecoderBuilder`] from a decoder's not-yet-decoded +/// state — the inverse of [`ParquetPushDecoderBuilder::build`]. The rebuilt +/// builder pins the remaining row groups and carries the remaining row +/// selection, offset/limit budget, and buffered bytes. +fn builder_from_remaining(parts: RemainingRowGroupsParts) -> ParquetPushDecoderBuilder { + let RemainingRowGroupsParts { + metadata, + schema, + row_groups, + selection, + offset, + limit, + reader_builder, + } = parts; + let RowGroupReaderBuilderParts { + batch_size, + projection, + fields, + filter, + max_predicate_cache_size, + metrics, + row_selection_policy, + buffers, + } = reader_builder; + + ArrowReaderBuilder { + input: PushDecoderInput::default(), + metadata, + schema, + fields, + batch_size, + // The frontier tracks remaining row groups explicitly, so the rebuilt + // builder always pins them (even if the original left `row_groups` as + // `None` meaning "all"). + row_groups: Some(row_groups), + projection, + filter, + selection, + row_selection_policy, + limit, + offset, + metrics, + max_predicate_cache_size, + } + // Carry the decoder's already-fetched bytes across the rebuild so the new + // decoder does not re-request them. + .with_buffers(buffers) +} + +/// A push based Parquet Decoder +/// +/// See [`ParquetPushDecoderBuilder`] for an example of how to build and use the decoder. +/// +/// [`ParquetPushDecoder`] is a low level API for decoding Parquet data without an +/// underlying reader for performing IO, and thus offers fine grained control +/// over how data is fetched and decoded. +/// +/// When more data is needed to make progress, instead of reading data directly +/// from a reader, the decoder returns [`DecodeResult`] indicating what ranges +/// are needed. Once the caller provides the requested ranges via +/// [`Self::push_ranges`], they try to decode again by calling +/// [`Self::try_decode`]. +/// +/// The decoder's internal state tracks what has been already decoded and what +/// is needed next. +#[derive(Debug)] +pub struct ParquetPushDecoder { + /// The inner state. + /// + /// This state is consumed on every transition and a new state is produced + /// so the Rust compiler can ensure that the state is always valid and + /// transitions are not missed. + state: ParquetDecoderState, +} + +impl ParquetPushDecoder { + /// Attempt to decode the next batch of data, or return what data is needed + /// + /// The the decoder communicates the next state with a [`DecodeResult`] + /// + /// See full example in [`ParquetPushDecoderBuilder`] + /// + /// ```no_run + /// # use parquet::arrow::push_decoder::ParquetPushDecoder; + /// use parquet::DecodeResult; + /// # fn get_decoder() -> ParquetPushDecoder { unimplemented!() } + /// # fn push_data(decoder: &mut ParquetPushDecoder, ranges: Vec>) { unimplemented!() } + /// let mut decoder = get_decoder(); + /// loop { + /// match decoder.try_decode().unwrap() { + /// DecodeResult::NeedsData(ranges) => { + /// // The decoder needs more data. Fetch the data for the given ranges + /// // call decoder.push_ranges(ranges, data) and call again + /// push_data(&mut decoder, ranges); + /// } + /// DecodeResult::Data(batch) => { + /// // Successfully decoded the next batch of data + /// println!("Got batch with {} rows", batch.num_rows()); + /// } + /// DecodeResult::Finished => { + /// // The decoder has finished decoding all data + /// break; + /// } + /// } + /// } + ///``` + pub fn try_decode(&mut self) -> Result, ParquetError> { + let current_state = std::mem::replace(&mut self.state, ParquetDecoderState::Finished); + let (new_state, decode_result) = current_state.try_next_batch()?; + self.state = new_state; + Ok(decode_result) + } + + /// Return a [`ParquetRecordBatchReader`] that reads the next set of rows, or + /// return what data is needed to produce it. + /// + /// This API can be used to get a reader for decoding the next set of + /// RecordBatches while proceeding to begin fetching data for the set (e.g + /// row group) + /// + /// Example + /// ```no_run + /// # use parquet::arrow::push_decoder::ParquetPushDecoder; + /// use parquet::DecodeResult; + /// # fn get_decoder() -> ParquetPushDecoder { unimplemented!() } + /// # fn push_data(decoder: &mut ParquetPushDecoder, ranges: Vec>) { unimplemented!() } + /// let mut decoder = get_decoder(); + /// loop { + /// match decoder.try_next_reader().unwrap() { + /// DecodeResult::NeedsData(ranges) => { + /// // The decoder needs more data. Fetch the data for the given ranges + /// // call decoder.push_ranges(ranges, data) and call again + /// push_data(&mut decoder, ranges); + /// } + /// DecodeResult::Data(reader) => { + /// // spawn a thread to read the batches in parallel + /// // with fetching the next row group / data + /// std::thread::spawn(move || { + /// for batch in reader { + /// let batch = batch.unwrap(); + /// println!("Got batch with {} rows", batch.num_rows()); + /// } + /// }); + /// } + /// DecodeResult::Finished => { + /// // The decoder has finished decoding all data + /// break; + /// } + /// } + /// } + ///``` + pub fn try_next_reader( + &mut self, + ) -> Result, ParquetError> { + let current_state = std::mem::replace(&mut self.state, ParquetDecoderState::Finished); + let (new_state, decode_result) = current_state.try_next_reader()?; + self.state = new_state; + Ok(decode_result) + } + + /// Push data into the decoder for processing + /// + /// This is a convenience wrapper around [`Self::push_ranges`] for pushing a + /// single range of data. + /// + /// Note this can be the entire file or just a part of it. If it is part of the file, + /// the ranges should correspond to the data ranges requested by the decoder. + /// + /// See example in [`ParquetPushDecoderBuilder`] + pub fn push_range(&mut self, range: Range, data: Bytes) -> Result<(), ParquetError> { + self.push_ranges(vec![range], vec![data]) + } + + /// Push data into the decoder for processing + /// + /// This should correspond to the data ranges requested by the decoder + pub fn push_ranges( + &mut self, + ranges: Vec>, + data: Vec, + ) -> Result<(), ParquetError> { + let current_state = std::mem::replace(&mut self.state, ParquetDecoderState::Finished); + self.state = current_state.push_data(ranges, data)?; + Ok(()) + } + + /// Returns the total number of buffered bytes in the decoder + /// + /// This is the sum of the size of all [`Bytes`] that has been pushed to the + /// decoder but not yet consumed. + /// + /// Note that this does not include any overhead of the internal data + /// structures and that since [`Bytes`] are ref counted memory, this may not + /// reflect additional memory usage. + /// + /// This can be used to monitor memory usage of the decoder. + pub fn buffered_bytes(&self) -> u64 { + self.state.buffered_bytes() + } + + /// Clear any staged byte ranges currently buffered for future decode work. + /// + /// This clears byte ranges still owned by the decoder's internal + /// `PushBuffers`. It does not affect any data that has already been handed + /// off to an active [`ParquetRecordBatchReader`]. + pub fn clear_all_ranges(&mut self) { + self.state.clear_all_ranges(); + } + + /// True iff the decoder is at a row-group boundary, where + /// [`Self::into_builder`] can reconfigure the scan. + /// + /// A boundary is "between row groups": the previous row group's + /// [`ParquetRecordBatchReader`] has been fully extracted (via + /// [`Self::try_next_reader`]) or fully drained (via [`Self::try_decode`]), + /// and the next row group has not yet been planned. While + /// [`Self::try_decode`] is iterating an active row group's reader this + /// returns `false`; with [`Self::try_next_reader`] there is a clean + /// window between two consecutive returns where this is `true`. + pub fn is_at_row_group_boundary(&self) -> bool { + self.state.is_at_row_group_boundary() + } + + /// Number of row groups left to decode after the one currently in flight. + /// Useful as a "should I bother reconfiguring the scan?" signal. + pub fn row_groups_remaining(&self) -> usize { + self.state.row_groups_remaining() + } + + /// Returns the row-group index that the next call to + /// [`Self::try_next_reader`] will yield a reader for, after applying + /// any internal skipping (row selection emptiness, exhausted budget, + /// finished state). + /// + /// Safe to call at any time. When called mid-row-group (i.e. while a + /// previously-emitted [`ParquetRecordBatchReader`] is still being + /// drained), the returned index refers to the row group that + /// [`Self::try_next_reader`] will produce *after* the current one. + /// + /// Returns `Ok(None)` when: + /// - the decoder has no more row groups to read, or + /// - every remaining row group would be skipped. + /// + /// This method does not mutate decoder state. It is useful for + /// callers that maintain per-row-group state in lock-step with the + /// decoder (e.g. dynamic row-group pruners) to determine which row + /// group the next reader corresponds to, since + /// [`Self::try_next_reader`] may silently advance past row groups + /// based on filtering and other criteria. + pub fn peek_next_row_group(&self) -> Result, ParquetError> { + self.state.peek_next_row_group() + } + + /// Decompose this decoder back into a [`ParquetPushDecoderBuilder`] for the + /// row groups that have *not* yet been decoded. + /// + /// This is the API for *adaptive* scans. Drive the decoder with + /// [`Self::try_next_reader`]; at any row-group boundary, call + /// `into_builder` to recover a builder, adjust it with the usual + /// [`ParquetPushDecoderBuilder`] setters, and + /// [`build`](ParquetPushDecoderBuilder::build) a fresh decoder that resumes + /// from the next row group: + /// + /// ```no_run + /// # use parquet::arrow::push_decoder::ParquetPushDecoder; + /// # use parquet::arrow::arrow_reader::RowFilter; + /// # fn get_decoder() -> ParquetPushDecoder { unimplemented!() } + /// # fn new_filter() -> RowFilter { unimplemented!() } + /// let mut decoder = get_decoder(); + /// // ... drive `decoder.try_next_reader()` for a few row groups ... + /// if decoder.is_at_row_group_boundary() && decoder.row_groups_remaining() > 0 { + /// decoder = decoder + /// .into_builder() + /// .unwrap() + /// // any builder option can be changed here, e.g. promote a + /// // filter into a row filter based on observed selectivity + /// .with_row_filter(new_filter()) + /// .build() + /// .unwrap(); + /// } + /// ``` + /// + /// The returned builder pins the not-yet-decoded row groups (via + /// [`with_row_groups`](ArrowReaderBuilder::with_row_groups)) and carries the + /// not-yet-consumed row selection and offset/limit budget, so rows from + /// already-decoded row groups are not produced again. Every other option — + /// projection, row filter, row selection policy, batch size, metrics, + /// predicate-cache size — is left exactly as the decoder had it and can be + /// overridden before [`build`](ParquetPushDecoderBuilder::build). + /// + /// # Errors + /// + /// Returns `Err(ParquetError::General)` when the decoder is not at a + /// row-group boundary (check [`Self::is_at_row_group_boundary`] first) or + /// has already finished. The decoder is consumed either way. + /// + /// # Buffered bytes + /// + /// The decoder's buffered bytes are carried across the rebuild: bytes + /// already fetched for row groups the new configuration still reads are + /// not re-requested. Bytes the new configuration no longer needs stay + /// buffered until [`clear_all_ranges`](Self::clear_all_ranges) is called + /// or the rebuilt decoder is dropped. + pub fn into_builder(self) -> Result { + self.state.into_builder() + } +} + +/// Internal state machine for the [`ParquetPushDecoder`] +#[derive(Debug)] +enum ParquetDecoderState { + /// Waiting for data needed to decode the next RowGroup + ReadingRowGroup { + remaining_row_groups: Box, + }, + /// The decoder is actively decoding a RowGroup + DecodingRowGroup { + /// Current active reader + record_batch_reader: Box, + remaining_row_groups: Box, + }, + /// The decoder has finished processing all data + Finished, +} + +impl ParquetDecoderState { + /// If actively reading a RowGroup, return the currently active + /// ParquetRecordBatchReader and advance to the next group. + fn try_next_reader( + self, + ) -> Result<(Self, DecodeResult), ParquetError> { + let mut current_state = self; + loop { + let (next_state, decode_result) = current_state.transition()?; + // if more data is needed to transition, can't proceed further without it + match decode_result { + DecodeResult::NeedsData(ranges) => { + return Ok((next_state, DecodeResult::NeedsData(ranges))); + } + // act next based on state + DecodeResult::Data(()) | DecodeResult::Finished => {} + } + match next_state { + // not ready to read yet, continue transitioning + Self::ReadingRowGroup { .. } => current_state = next_state, + // have a reader ready, so return it and set ourself to ReadingRowGroup + Self::DecodingRowGroup { + record_batch_reader, + remaining_row_groups, + } => { + let result = DecodeResult::Data(*record_batch_reader); + let next_state = Self::ReadingRowGroup { + remaining_row_groups, + }; + return Ok((next_state, result)); + } + Self::Finished => { + return Ok((Self::Finished, DecodeResult::Finished)); + } + } + } + } + + /// Current state --> next state + output + /// + /// This function is called to get the next RecordBatch + /// + /// This structure is used to reduce the indentation level of the main loop + /// in try_build + fn try_next_batch(self) -> Result<(Self, DecodeResult), ParquetError> { + let mut current_state = self; + loop { + let (new_state, decode_result) = current_state.transition()?; + // if more data is needed to transition, can't proceed further without it + match decode_result { + DecodeResult::NeedsData(ranges) => { + return Ok((new_state, DecodeResult::NeedsData(ranges))); + } + // act next based on state + DecodeResult::Data(()) | DecodeResult::Finished => {} + } + match new_state { + // not ready to read yet, continue transitioning + Self::ReadingRowGroup { .. } => current_state = new_state, + // have a reader ready, so decode the next batch + Self::DecodingRowGroup { + mut record_batch_reader, + remaining_row_groups, + } => { + match record_batch_reader.next() { + // Successfully decoded a batch, return it + Some(Ok(batch)) => { + let result = DecodeResult::Data(batch); + let next_state = Self::DecodingRowGroup { + record_batch_reader, + remaining_row_groups, + }; + return Ok((next_state, result)); + } + // No more batches in this row group, move to the next row group + None => { + current_state = Self::ReadingRowGroup { + remaining_row_groups, + } + } + // some error occurred while decoding, so return that + Some(Err(e)) => { + // TODO: preserve ArrowError in ParquetError (rather than convert to a string) + return Err(ParquetError::ArrowError(e.to_string())); + } + } + } + Self::Finished => { + return Ok((Self::Finished, DecodeResult::Finished)); + } + } + } + } + + /// Transition to the next state with a reader (data can be produced), if not end of stream + /// + /// This function is called in a loop until the decoder is ready to return + /// data (has the required pages buffered) or is finished. + fn transition(self) -> Result<(Self, DecodeResult<()>), ParquetError> { + // result returned when there is data ready + let data_ready = DecodeResult::Data(()); + match self { + Self::ReadingRowGroup { + mut remaining_row_groups, + } => { + match remaining_row_groups.try_next_reader()? { + // If we have a next reader, we can transition to decoding it + DecodeResult::Data(record_batch_reader) => { + // Transition to decoding the row group + Ok(( + Self::DecodingRowGroup { + record_batch_reader: Box::new(record_batch_reader), + remaining_row_groups, + }, + data_ready, + )) + } + DecodeResult::NeedsData(ranges) => { + // If we need more data, we return the ranges needed and stay in Reading + // RowGroup state + Ok(( + Self::ReadingRowGroup { + remaining_row_groups, + }, + DecodeResult::NeedsData(ranges), + )) + } + // If there are no more readers, we are finished + DecodeResult::Finished => { + // No more row groups to read, we are finished + Ok((Self::Finished, DecodeResult::Finished)) + } + } + } + // if we are already in DecodingRowGroup, just return data ready + Self::DecodingRowGroup { .. } => Ok((self, data_ready)), + // if finished, just return finished + Self::Finished => Ok((self, DecodeResult::Finished)), + } + } + + /// Push data, and transition state if needed + /// + /// This should correspond to the data ranges requested by the decoder + pub fn push_data( + self, + ranges: Vec>, + data: Vec, + ) -> Result { + match self { + ParquetDecoderState::ReadingRowGroup { + mut remaining_row_groups, + } => { + // Push data to the RowGroupReaderBuilder + remaining_row_groups.push_data(ranges, data); + Ok(ParquetDecoderState::ReadingRowGroup { + remaining_row_groups, + }) + } + // it is ok to get data before we asked for it + ParquetDecoderState::DecodingRowGroup { + record_batch_reader, + mut remaining_row_groups, + } => { + remaining_row_groups.push_data(ranges, data); + Ok(ParquetDecoderState::DecodingRowGroup { + record_batch_reader, + remaining_row_groups, + }) + } + ParquetDecoderState::Finished => Err(ParquetError::General( + "Cannot push data to a finished decoder".to_string(), + )), + } + } + + /// How many bytes are currently buffered in the decoder? + fn buffered_bytes(&self) -> u64 { + match self { + ParquetDecoderState::ReadingRowGroup { + remaining_row_groups, + } => remaining_row_groups.buffered_bytes(), + ParquetDecoderState::DecodingRowGroup { + record_batch_reader: _, + remaining_row_groups, + } => remaining_row_groups.buffered_bytes(), + ParquetDecoderState::Finished => 0, + } + } + + /// Clear any staged ranges currently buffered in the decoder. + fn clear_all_ranges(&mut self) { + match self { + ParquetDecoderState::ReadingRowGroup { + remaining_row_groups, + } => remaining_row_groups.clear_all_ranges(), + ParquetDecoderState::DecodingRowGroup { + record_batch_reader: _, + remaining_row_groups, + } => remaining_row_groups.clear_all_ranges(), + ParquetDecoderState::Finished => {} + } + } + + fn is_at_row_group_boundary(&self) -> bool { + match self { + ParquetDecoderState::ReadingRowGroup { + remaining_row_groups, + } => remaining_row_groups.is_at_row_group_boundary(), + // Mid-row-group: the active reader holds an `ArrayReader` and + // `ReadPlan` keyed to the *current* projection/filter; rebuilding + // would require throwing that work away. + ParquetDecoderState::DecodingRowGroup { .. } => false, + ParquetDecoderState::Finished => false, + } + } + + fn row_groups_remaining(&self) -> usize { + match self { + ParquetDecoderState::ReadingRowGroup { + remaining_row_groups, + } => remaining_row_groups.row_groups_remaining(), + ParquetDecoderState::DecodingRowGroup { + remaining_row_groups, + .. + } => remaining_row_groups.row_groups_remaining(), + ParquetDecoderState::Finished => 0, + } + } + + /// See [`ParquetPushDecoder::peek_next_row_group`] for the public + /// API contract. This inner method delegates to the underlying + /// [`RemainingRowGroups`] for both `ReadingRowGroup` and + /// `DecodingRowGroup`: mid-row-group the answer is the row group + /// `try_next_reader` will produce *after* the active one finishes, + /// which is exactly what `RemainingRowGroups::peek_next_row_group` + /// computes from the queued frontier. + fn peek_next_row_group(&self) -> Result, ParquetError> { + match self { + ParquetDecoderState::ReadingRowGroup { + remaining_row_groups, + } => remaining_row_groups.peek_next_row_group(), + ParquetDecoderState::DecodingRowGroup { + remaining_row_groups, + .. + } => remaining_row_groups.peek_next_row_group(), + ParquetDecoderState::Finished => Ok(None), + } + } + + fn into_builder(self) -> Result { + let remaining_row_groups = match self { + ParquetDecoderState::ReadingRowGroup { + remaining_row_groups, + } => remaining_row_groups, + ParquetDecoderState::DecodingRowGroup { .. } => { + return Err(ParquetError::General( + "into_builder called while a row group is being decoded; \ + check is_at_row_group_boundary() first" + .to_string(), + )); + } + ParquetDecoderState::Finished => { + return Err(ParquetError::General( + "into_builder called on a finished decoder".to_string(), + )); + } + }; + if !remaining_row_groups.is_at_row_group_boundary() { + return Err(ParquetError::General( + "into_builder called mid-row-group; check is_at_row_group_boundary() first" + .to_string(), + )); + } + Ok(builder_from_remaining(remaining_row_groups.into_parts())) + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::DecodeResult; + use crate::arrow::arrow_reader::{ArrowPredicateFn, RowFilter, RowSelection, RowSelector}; + use crate::arrow::push_decoder::{ParquetPushDecoder, ParquetPushDecoderBuilder}; + use crate::arrow::{ArrowWriter, ProjectionMask}; + use crate::errors::ParquetError; + use crate::file::metadata::ParquetMetaDataPushDecoder; + use crate::file::properties::WriterProperties; + use arrow::compute::kernels::cmp::{gt, lt}; + use arrow_array::cast::AsArray; + use arrow_array::types::Int64Type; + use arrow_array::{ArrayRef, Int64Array, RecordBatch, StringViewArray}; + use arrow_select::concat::concat_batches; + use bytes::Bytes; + use std::fmt::Debug; + use std::ops::Range; + use std::sync::{Arc, LazyLock}; + + /// Test decoder struct size (as they are copied around on each transition, they + /// should not grow too large) + #[test] + fn test_decoder_size() { + assert_eq!(std::mem::size_of::(), 24); + } + + /// Decode the entire file at once, simulating a scenario where all data is + /// available in memory + #[test] + fn test_decoder_all_data() { + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .build() + .unwrap(); + + decoder + .push_range(test_file_range(), TEST_FILE_DATA.clone()) + .unwrap(); + + let results = vec![ + // first row group should be decoded without needing more data + expect_data(decoder.try_decode()), + // second row group should be decoded without needing more data + expect_data(decoder.try_decode()), + ]; + expect_finished(decoder.try_decode()); + + let all_output = concat_batches(&TEST_BATCH.schema(), &results).unwrap(); + // Check that the output matches the input batch + assert_eq!(all_output, *TEST_BATCH); + } + + /// Decode the entire file incrementally, simulating a scenario where data is + /// fetched as needed + #[test] + fn test_decoder_incremental() { + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .build() + .unwrap(); + + let mut results = vec![]; + + // First row group, expect a single request + let ranges = expect_needs_data(decoder.try_decode()); + let num_bytes_requested: u64 = ranges.iter().map(|r| r.end - r.start).sum(); + push_ranges_to_decoder(&mut decoder, ranges); + // The decoder should currently only store the data it needs to decode the first row group + assert_eq!(decoder.buffered_bytes(), num_bytes_requested); + results.push(expect_data(decoder.try_decode())); + // the decoder should have consumed the data for the first row group and freed it + assert_eq!(decoder.buffered_bytes(), 0); + + // Second row group, + let ranges = expect_needs_data(decoder.try_decode()); + let num_bytes_requested: u64 = ranges.iter().map(|r| r.end - r.start).sum(); + push_ranges_to_decoder(&mut decoder, ranges); + // The decoder should currently only store the data it needs to decode the second row group + assert_eq!(decoder.buffered_bytes(), num_bytes_requested); + results.push(expect_data(decoder.try_decode())); + // the decoder should have consumed the data for the second row group and freed it + assert_eq!(decoder.buffered_bytes(), 0); + expect_finished(decoder.try_decode()); + + // Check that the output matches the input batch + let all_output = concat_batches(&TEST_BATCH.schema(), &results).unwrap(); + assert_eq!(all_output, *TEST_BATCH); + } + + /// Releasing staged ranges should free speculative buffers without affecting + /// the active row group reader. + #[test] + fn test_decoder_clear_all_ranges() { + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .with_batch_size(100) + .build() + .unwrap(); + + decoder + .push_range(test_file_range(), TEST_FILE_DATA.clone()) + .unwrap(); + assert_eq!(decoder.buffered_bytes(), test_file_len()); + + // The current row group reader is built from the prefetched bytes, but + // the speculative full-file range remains staged in the decoder. + let batch1 = expect_data(decoder.try_decode()); + assert_eq!(batch1, TEST_BATCH.slice(0, 100)); + assert_eq!(decoder.buffered_bytes(), test_file_len()); + + // All of the buffer is released + decoder.clear_all_ranges(); + assert_eq!(decoder.buffered_bytes(), 0); + + // The active reader still owns the current row group's bytes, so it can + // continue decoding without consulting PushBuffers. + let batch2 = expect_data(decoder.try_decode()); + assert_eq!(batch2, TEST_BATCH.slice(100, 100)); + assert_eq!(decoder.buffered_bytes(), 0); + + // Moving to the next row group now requires the decoder to ask for data + // again because the staged speculative ranges were released. + let ranges = expect_needs_data(decoder.try_decode()); + let num_bytes_requested: u64 = ranges.iter().map(|r| r.end - r.start).sum(); + push_ranges_to_decoder(&mut decoder, ranges); + assert_eq!(decoder.buffered_bytes(), num_bytes_requested); + + let batch3 = expect_data(decoder.try_decode()); + assert_eq!(batch3, TEST_BATCH.slice(200, 100)); + assert_eq!(decoder.buffered_bytes(), 0); + + let batch4 = expect_data(decoder.try_decode()); + assert_eq!(batch4, TEST_BATCH.slice(300, 100)); + assert_eq!(decoder.buffered_bytes(), 0); + + expect_finished(decoder.try_decode()); + } + + /// Decode the entire file incrementally, simulating partial reads + #[test] + fn test_decoder_partial() { + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .build() + .unwrap(); + + // First row group, expect a single request for all data needed to read "a" and "b" + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + let batch1 = expect_data(decoder.try_decode()); + let expected1 = TEST_BATCH.slice(0, 200); + assert_eq!(batch1, expected1); + + // Second row group, this time provide the data in two steps + let ranges = expect_needs_data(decoder.try_decode()); + let (ranges1, ranges2) = ranges.split_at(ranges.len() / 2); + assert!(!ranges1.is_empty()); + assert!(!ranges2.is_empty()); + // push first half to simulate partial read + push_ranges_to_decoder(&mut decoder, ranges1.to_vec()); + + // still expect more data + let ranges = expect_needs_data(decoder.try_decode()); + assert_eq!(ranges, ranges2); // should be the remaining ranges + // push empty ranges should be a no-op + push_ranges_to_decoder(&mut decoder, vec![]); + let ranges = expect_needs_data(decoder.try_decode()); + assert_eq!(ranges, ranges2); // should be the remaining ranges + push_ranges_to_decoder(&mut decoder, ranges); + + let batch2 = expect_data(decoder.try_decode()); + let expected2 = TEST_BATCH.slice(200, 200); + assert_eq!(batch2, expected2); + + expect_finished(decoder.try_decode()); + } + + /// Decode multiple columns "a" and "b", expect that the decoder requests + /// only a single request per row group + #[test] + fn test_decoder_selection_does_one_request() { + let builder = + ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap(); + + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + let mut decoder = builder + .with_projection( + ProjectionMask::columns(&schema_descr, ["a", "b"]), // read "a", "b" + ) + .build() + .unwrap(); + + // First row group, expect a single request for all data needed to read "a" and "b" + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + let batch1 = expect_data(decoder.try_decode()); + let expected1 = TEST_BATCH.slice(0, 200).project(&[0, 1]).unwrap(); + assert_eq!(batch1, expected1); + + // Second row group, similarly expect a single request for all data needed to read "a" and "b" + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + let batch2 = expect_data(decoder.try_decode()); + let expected2 = TEST_BATCH.slice(200, 200).project(&[0, 1]).unwrap(); + assert_eq!(batch2, expected2); + + expect_finished(decoder.try_decode()); + } + + /// Decode with a filter that requires multiple requests, but only provide part + /// of the data needed for the filter at a time simulating partial reads. + #[test] + fn test_decoder_single_filter_partial() { + let builder = + ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap(); + + // Values in column "a" range 0..399 + // First filter: "a" > 250 (nothing in Row Group 0, both data pages in Row Group 1) + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + // a > 250 + let row_filter_a = ArrowPredicateFn::new( + // claim to use both a and b so we get two ranges requests for the filter pages + ProjectionMask::columns(&schema_descr, ["a", "b"]), + |batch: RecordBatch| { + let scalar_250 = Int64Array::new_scalar(250); + let column = batch.column(0).as_primitive::(); + gt(column, &scalar_250) + }, + ); + + let mut decoder = builder + .with_projection( + // read only column "a" to test that filter pages are reused + ProjectionMask::columns(&schema_descr, ["a"]), // read "a" + ) + .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)])) + .build() + .unwrap(); + + // First row group, evaluating filters + let ranges = expect_needs_data(decoder.try_decode()); + // only provide half the ranges + let (ranges1, ranges2) = ranges.split_at(ranges.len() / 2); + assert!(!ranges1.is_empty()); + assert!(!ranges2.is_empty()); + push_ranges_to_decoder(&mut decoder, ranges1.to_vec()); + // still expect more data + let ranges = expect_needs_data(decoder.try_decode()); + assert_eq!(ranges, ranges2); // should be the remaining ranges + let ranges = expect_needs_data(decoder.try_decode()); + assert_eq!(ranges, ranges2); // should be the remaining ranges + push_ranges_to_decoder(&mut decoder, ranges2.to_vec()); + + // Since no rows in the first row group pass the filters, there is no + // additional requests to read data pages for "b" here + + // Second row group + 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(251, 149).project(&[0]).unwrap(); + assert_eq!(batch, expected); + + expect_finished(decoder.try_decode()); + } + + /// Decode with a filter where we also skip one of the RowGroups via a RowSelection + #[test] + fn test_decoder_single_filter_and_row_selection() { + let builder = + ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap(); + + // Values in column "a" range 0..399 + // First filter: "a" > 250 (nothing in Row Group 0, last data page in Row Group 1) + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + // a > 250 + let row_filter_a = ArrowPredicateFn::new( + ProjectionMask::columns(&schema_descr, ["a"]), + |batch: RecordBatch| { + let scalar_250 = Int64Array::new_scalar(250); + let column = batch.column(0).as_primitive::(); + gt(column, &scalar_250) + }, + ); + + let mut decoder = builder + .with_projection( + // read only column "a" to test that filter pages are reused + ProjectionMask::columns(&schema_descr, ["b"]), // read "b" + ) + .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)])) + .with_row_selection(RowSelection::from(vec![ + RowSelector::skip(200), // skip first row group + RowSelector::select(100), // first 100 rows of second row group + RowSelector::skip(100), + ])) + .build() + .unwrap(); + + // expect the first row group to be filtered out (no filter is evaluated due to row selection) + + // First row group, first filter (a > 250) + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // Second row group + 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(251, 49).project(&[1]).unwrap(); + assert_eq!(batch, expected); + + expect_finished(decoder.try_decode()); + } + + /// Decode with multiple filters that require multiple requests + #[test] + fn test_decoder_multi_filters() { + // Create a decoder for decoding parquet data (note it does not have any IO / readers) + let builder = + ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap(); + + // Values in column "a" range 0..399 + // Values in column "b" range 400..799 + // First filter: "a" > 175 (last data page in Row Group 0) + // Second filter: "b" < 625 (last data page in Row Group 0 and first DataPage in RowGroup 1) + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + // a > 175 + let row_filter_a = ArrowPredicateFn::new( + ProjectionMask::columns(&schema_descr, ["a"]), + |batch: RecordBatch| { + let scalar_175 = Int64Array::new_scalar(175); + let column = batch.column(0).as_primitive::(); + gt(column, &scalar_175) + }, + ); + + // b < 625 + let row_filter_b = ArrowPredicateFn::new( + ProjectionMask::columns(&schema_descr, ["b"]), + |batch: RecordBatch| { + let scalar_625 = Int64Array::new_scalar(625); + let column = batch.column(0).as_primitive::(); + lt(column, &scalar_625) + }, + ); + + let mut decoder = builder + .with_projection( + ProjectionMask::columns(&schema_descr, ["c"]), // read "c" + ) + .with_row_filter(RowFilter::new(vec![ + Box::new(row_filter_a), + Box::new(row_filter_b), + ])) + .build() + .unwrap(); + + // First row group, first filter (a > 175) + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // first row group, second filter (b < 625) + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // first row group, data pages for "c" + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // expect the first batch to be decoded: rows 176..199, column "c" + let batch1 = expect_data(decoder.try_decode()); + let expected1 = TEST_BATCH.slice(176, 24).project(&[2]).unwrap(); + assert_eq!(batch1, expected1); + + // Second row group, first filter (a > 175) + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // Second row group, second filter (b < 625) + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // Second row group, data pages for "c" + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // expect the second batch to be decoded: rows 200..224, column "c" + let batch2 = expect_data(decoder.try_decode()); + let expected2 = TEST_BATCH.slice(200, 25).project(&[2]).unwrap(); + assert_eq!(batch2, expected2); + + expect_finished(decoder.try_decode()); + } + + /// 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] + fn test_decoder_reuses_filter_pages() { + // Create a decoder for decoding parquet data (note it does not have any IO / readers) + let builder = + ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap(); + + // Values in column "a" range 0..399 + // First filter: "a" > 250 (nothing in Row Group 0, last data page in Row Group 1) + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + // a > 250 + let row_filter_a = ArrowPredicateFn::new( + ProjectionMask::columns(&schema_descr, ["a"]), + |batch: RecordBatch| { + let scalar_250 = Int64Array::new_scalar(250); + let column = batch.column(0).as_primitive::(); + gt(column, &scalar_250) + }, + ); + + let mut decoder = builder + .with_projection( + // read only column "a" to test that filter pages are reused + ProjectionMask::columns(&schema_descr, ["a"]), // read "a" + ) + .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)])) + .build() + .unwrap(); + + // First row group, first filter (a > 175) + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // expect the first row group to be filtered out (no rows match) + + // Second row group, first filter (a > 250) + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // expect that the second row group is decoded: rows 251..399, column "a" + // Note that the filter pages for "a" should be reused and no additional data + // should be requested + let batch = expect_data(decoder.try_decode()); + let expected = TEST_BATCH.slice(251, 149).project(&[0]).unwrap(); + assert_eq!(batch, expected); + + expect_finished(decoder.try_decode()); + } + + #[test] + fn test_decoder_empty_filters() { + let builder = + ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap(); + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + // only read column "c", but with empty filters + let mut decoder = builder + .with_projection( + ProjectionMask::columns(&schema_descr, ["c"]), // read "c" + ) + .with_row_filter(RowFilter::new(vec![ + // empty filters should be ignored + ])) + .build() + .unwrap(); + + // First row group + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // expect the first batch to be decoded: rows 0..199, column "c" + let batch1 = expect_data(decoder.try_decode()); + let expected1 = TEST_BATCH.slice(0, 200).project(&[2]).unwrap(); + assert_eq!(batch1, expected1); + + // Second row group, + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // expect the second batch to be decoded: rows 200..399, column "c" + let batch2 = expect_data(decoder.try_decode()); + let expected2 = TEST_BATCH.slice(200, 200).project(&[2]).unwrap(); + + assert_eq!(batch2, expected2); + + expect_finished(decoder.try_decode()); + } + + /// When filter pushdown is combined with a `LIMIT`, the predicate must + /// not be evaluated for rows beyond the `limit`-th match. + /// + /// Filter `a > 175` produces 24 matches in row group 0 (rows 176..199). + /// With `limit = 10`, only the first 10 matches (rows 176..185) should be + /// emitted, AND the predicate counter should observe that evaluation was + /// short-circuited. + #[test] + fn test_decoder_filter_with_limit_short_circuits_within_row_group() { + 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 rows_filtered = Arc::new(AtomicUsize::new(0)); + let rows_filtered_for_predicate = Arc::clone(&rows_filtered); + + let row_filter_a = ArrowPredicateFn::new( + ProjectionMask::columns(&schema_descr, ["a"]), + move |batch: RecordBatch| { + rows_filtered_for_predicate.fetch_add(batch.num_rows(), Ordering::Relaxed); + let scalar_175 = Int64Array::new_scalar(175); + let column = batch.column(0).as_primitive::(); + gt(column, &scalar_175) + }, + ); + + // Use a small batch size so the row group is evaluated across + // multiple predicate batches; that is the regime where Layer 2's + // short-circuit saves predicate evaluation work. Matching rows are + // 176..199 (24 rows); with batch_size = 10 those span batches 17, 18, + // and 19 (rows 170..199). A limit of 10 should stop filter evaluation + // in the middle of batch 18. + let mut decoder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["a"])) + .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)])) + .with_batch_size(10) + .with_limit(10) + .build() + .unwrap(); + + // First row group: filter columns fetch (predicate is evaluated here) + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // The first 10 matching rows come out: 176..185, column "a" + let batch = expect_data(decoder.try_decode()); + let expected = TEST_BATCH.slice(176, 10).project(&[0]).unwrap(); + assert_eq!(batch, expected); + + // no data for row group 1 should be requested — the limit + // was satisfied by row group 0 and the `Start` state for row group 1 + // short-circuits to `Finished`. + expect_finished(decoder.try_decode()); + + // Row 186 is the 11th match; the scan should stop no later than the + // batch containing it (batch 18 of 10 rows = rows 180..189), so at + // most 190 rows are evaluated. + let evaluated = rows_filtered.load(Ordering::Relaxed); + assert!( + evaluated <= 190, + "predicate evaluated {evaluated} rows; expected ≤ 190 (stop within batch containing 11th match)" + ); + } + + /// Once the limit has been satisfied by a prior row group, subsequent + /// row groups should be skipped entirely — no data request for their + /// filter columns. + #[test] + fn test_decoder_filter_with_limit_skips_later_row_groups() { + let builder = + ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap(); + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + // `a > 175` matches rows 176..199 in row group 0 (24 matches) and + // 200..399 in row group 1 (200 matches). With limit = 5, all matches + // should come from row group 0. + let row_filter_a = ArrowPredicateFn::new( + ProjectionMask::columns(&schema_descr, ["a"]), + |batch: RecordBatch| { + let scalar_175 = Int64Array::new_scalar(175); + let column = batch.column(0).as_primitive::(); + gt(column, &scalar_175) + }, + ); + + let mut decoder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["a"])) + .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)])) + .with_limit(5) + .build() + .unwrap(); + + // Row group 0: fetch filter pages + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // First 5 matches: 176..180 + let batch = expect_data(decoder.try_decode()); + let expected = TEST_BATCH.slice(176, 5).project(&[0]).unwrap(); + assert_eq!(batch, expected); + + // Row group 1 must NOT request data — the limit is already satisfied + // so `Start` in row group 1 short-circuits to `Finished`. + expect_finished(decoder.try_decode()); + } + + /// The predicate short-circuit must account for `self.offset` as well as + /// `self.limit`. The post-predicate `with_offset` step skips that many + /// already-selected rows before `with_limit` counts output rows — so the + /// predicate must retain at least `offset + limit` matches. Without the + /// fix, Layer 2 caps at just `limit` and the later `with_offset` consumes + /// all of them, producing 0 rows instead of `limit`. + /// + /// `a > 175` matches rows 176..199 in row group 0 (24 matches). With + /// `offset = 10, limit = 5`, the expected output is rows 186..190 (the + /// 11th through 15th matches). + #[test] + fn test_decoder_filter_with_offset_and_limit() { + let builder = + ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap(); + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + let row_filter_a = ArrowPredicateFn::new( + ProjectionMask::columns(&schema_descr, ["a"]), + |batch: RecordBatch| { + let scalar_175 = Int64Array::new_scalar(175); + let column = batch.column(0).as_primitive::(); + gt(column, &scalar_175) + }, + ); + + let mut decoder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["a"])) + .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)])) + .with_offset(10) + .with_limit(5) + .build() + .unwrap(); + + 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(186, 5).project(&[0]).unwrap(); + assert_eq!(batch, expected); + + expect_finished(decoder.try_decode()); + } + + /// The limit short-circuit must also be correct when the limited predicate + /// is the last predicate in a multi-predicate chain. + /// + /// `a > 175` first narrows row group 0 to rows 176..199. The final + /// predicate `b < 625` is then evaluated only over those 24 rows, all of + /// which match. With `limit = 10`, the final output should still be rows + /// 176..185, and the second predicate should stop before consuming all 24 + /// selected rows. + #[test] + fn test_decoder_multi_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 first_predicate_rows = Arc::new(AtomicUsize::new(0)); + let second_predicate_rows = Arc::new(AtomicUsize::new(0)); + + let first_predicate_rows_for_filter = Arc::clone(&first_predicate_rows); + let row_filter_a = ArrowPredicateFn::new( + ProjectionMask::columns(&schema_descr, ["a"]), + move |batch: RecordBatch| { + first_predicate_rows_for_filter.fetch_add(batch.num_rows(), Ordering::Relaxed); + let scalar_175 = Int64Array::new_scalar(175); + let column = batch.column(0).as_primitive::(); + gt(column, &scalar_175) + }, + ); + + let second_predicate_rows_for_filter = Arc::clone(&second_predicate_rows); + let row_filter_b = ArrowPredicateFn::new( + ProjectionMask::columns(&schema_descr, ["b"]), + move |batch: RecordBatch| { + second_predicate_rows_for_filter.fetch_add(batch.num_rows(), Ordering::Relaxed); + let scalar_625 = Int64Array::new_scalar(625); + let column = batch.column(0).as_primitive::(); + lt(column, &scalar_625) + }, + ); + + let mut decoder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["c"])) + .with_row_filter(RowFilter::new(vec![ + Box::new(row_filter_a), + Box::new(row_filter_b), + ])) + .with_batch_size(10) + .with_limit(10) + .build() + .unwrap(); + + // Row group 0, first predicate + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // Row group 0, second predicate + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // Final projected data + 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, 10).project(&[2]).unwrap(); + assert_eq!(batch, expected); + + // The overall limit was satisfied by row group 0. + expect_finished(decoder.try_decode()); + + assert_eq!(first_predicate_rows.load(Ordering::Relaxed), 200); + assert!( + second_predicate_rows.load(Ordering::Relaxed) < 24, + "final predicate should short-circuit before consuming all 24 rows selected by the first predicate" + ); + } + + /// When a row selection already exists, limiting the predicate must still + /// preserve alignment with that prior selection. + /// + /// The explicit selection narrows row group 0 to rows 150..199. Applying + /// `a > 175` over that selection yields rows 176..199. With `limit = 10`, + /// the decoder should emit rows 176..185 and stop without evaluating the + /// remaining selected rows. + #[test] + fn test_decoder_filter_with_row_selection_and_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 rows_filtered = Arc::new(AtomicUsize::new(0)); + let rows_filtered_for_predicate = Arc::clone(&rows_filtered); + + let row_filter_a = ArrowPredicateFn::new( + ProjectionMask::columns(&schema_descr, ["a"]), + move |batch: RecordBatch| { + rows_filtered_for_predicate.fetch_add(batch.num_rows(), Ordering::Relaxed); + let scalar_175 = Int64Array::new_scalar(175); + let column = batch.column(0).as_primitive::(); + gt(column, &scalar_175) + }, + ); + + let mut decoder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["a"])) + .with_row_selection(RowSelection::from(vec![ + RowSelector::skip(150), + RowSelector::select(50), + ])) + .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)])) + .with_batch_size(10) + .with_limit(10) + .build() + .unwrap(); + + 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, 10).project(&[0]).unwrap(); + assert_eq!(batch, expected); + + expect_finished(decoder.try_decode()); + + assert!( + rows_filtered.load(Ordering::Relaxed) < 50, + "predicate should short-circuit before consuming all 50 rows from the explicit row selection" + ); + } + + #[test] + fn test_decoder_offset_limit() { + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + // skip entire first row group (200 rows) and first 25 rows of second row group + .with_offset(225) + // and limit to 20 rows + .with_limit(20) + .build() + .unwrap(); + + // First row group should be skipped, + + // Second row group + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // expect the first and only batch to be decoded + let batch1 = expect_data(decoder.try_decode()); + let expected1 = TEST_BATCH.slice(225, 20); + assert_eq!(batch1, expected1); + + expect_finished(decoder.try_decode()); + } + + #[test] + fn test_decoder_try_next_reader_offset_limit() { + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .with_offset(225) + .with_limit(20) + .build() + .unwrap(); + + let ranges = expect_needs_data(decoder.try_next_reader()); + push_ranges_to_decoder(&mut decoder, ranges); + + let reader = expect_data(decoder.try_next_reader()); + let batches = reader + .map(|batch| batch.expect("expected decoded batch")) + .collect::>(); + let output = concat_batches(&TEST_BATCH.schema(), &batches).unwrap(); + assert_eq!(output, TEST_BATCH.slice(225, 20)); + + expect_finished(decoder.try_next_reader()); + } + + #[test] + fn test_decoder_row_group_selection() { + // take only the second row group + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .with_row_groups(vec![1]) + .build() + .unwrap(); + + // First row group should be skipped, + + // Second row group + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // expect the first and only batch to be decoded + let batch1 = expect_data(decoder.try_decode()); + let expected1 = TEST_BATCH.slice(200, 200); + assert_eq!(batch1, expected1); + + expect_finished(decoder.try_decode()); + } + + #[test] + fn test_decoder_row_selection() { + // take only the second row group + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .with_row_selection(RowSelection::from(vec![ + RowSelector::skip(225), // skip first row group and 25 rows of second]) + RowSelector::select(20), // take 20 rows + ])) + .build() + .unwrap(); + + // First row group should be skipped, + + // Second row group + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // expect the first ane only batch to be decoded + let batch1 = expect_data(decoder.try_decode()); + let expected1 = TEST_BATCH.slice(225, 20); + assert_eq!(batch1, expected1); + + expect_finished(decoder.try_decode()); + } + + /// `peek_next_row_group` reports the index of the row group the + /// next `try_next_reader` call will hand back, matching the + /// frontier's internal skip logic. + #[test] + fn test_peek_next_row_group_basic() { + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .build() + .unwrap(); + + // Two row groups (0, 1). At boundary before any read, peek should + // see RG 0. + assert_eq!(decoder.peek_next_row_group().unwrap(), Some(0)); + assert!(decoder.is_at_row_group_boundary()); + + let ranges = expect_needs_data(decoder.try_next_reader()); + push_ranges_to_decoder(&mut decoder, ranges); + let reader = expect_data(decoder.try_next_reader()); + // Once the reader for RG 0 has been handed off, the decoder is + // back at a boundary waiting for RG 1 — peek must reflect that + // (the active reader lives outside the decoder). + assert!(decoder.is_at_row_group_boundary()); + assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1)); + + // Drain RG 0's reader and consume RG 1. + for batch in reader { + let _ = batch.unwrap(); + } + let ranges = expect_needs_data(decoder.try_next_reader()); + push_ranges_to_decoder(&mut decoder, ranges); + let reader = expect_data(decoder.try_next_reader()); + for batch in reader { + let _ = batch.unwrap(); + } + + // No row groups left. + assert_eq!(decoder.peek_next_row_group().unwrap(), None); + } + + /// `peek_next_row_group` honors `with_row_groups` — restricting the + /// scan to a single row group means peek reports only that one and + /// then `None`. + #[test] + fn test_peek_next_row_group_respects_with_row_groups() { + let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .with_row_groups(vec![1]) + .build() + .unwrap(); + + assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1)); + } + + /// When a row-selection segment leaves the next row group with zero + /// selected rows, `peek_next_row_group` mirrors + /// `next_readable_row_group`'s skip: it returns the *following* + /// row group instead of the empty one. + #[test] + fn test_peek_next_row_group_skips_empty_selection() { + // Each row group has 200 rows. Skip all 200 of RG 0 plus 50 of + // RG 1; the next reader will be for RG 1, not RG 0. + let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .with_row_selection(RowSelection::from(vec![ + RowSelector::skip(250), + RowSelector::select(100), + ])) + .build() + .unwrap(); + + assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1)); + } + + /// `peek_next_row_group` returns `None` on a finished decoder. + #[test] + fn test_peek_next_row_group_finished() { + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .with_row_groups(vec![]) + .build() + .unwrap(); + + // No row groups requested ⇒ already finished, no peek. + expect_finished(decoder.try_next_reader()); + assert_eq!(decoder.peek_next_row_group().unwrap(), None); + } + + /// `peek_next_row_group` mirrors `next_readable_row_group`'s + /// budget-skipping logic: with an `OFFSET` that consumes the + /// entire first row group, peek must return the *following* + /// row group, not RG 0 (no predicates, so budget skips apply). + #[test] + fn test_peek_next_row_group_skips_for_budget() { + // Each row group has 200 rows. OFFSET 200 drains RG 0 entirely + // and lands the decoder's first emitted reader at RG 1. + let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .with_offset(200) + .build() + .unwrap(); + + assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1)); + } + + /// OFFSET larger than every remaining row group's row count + /// exhausts the budget, so peek should report `None` — matching + /// `next_readable_row_group`'s behavior of producing `Finished`. + #[test] + fn test_peek_next_row_group_budget_drains_all() { + // 2 row groups × 200 rows = 400 total. OFFSET 500 cannot land + // anywhere — every RG is skipped under the budget. + let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .with_offset(500) + .build() + .unwrap(); + + assert_eq!(decoder.peek_next_row_group().unwrap(), None); + } + + /// `peek_next_row_group` is safe to call while a row group's reader + /// is still being drained via `try_decode`. In `DecodingRowGroup` + /// state it must report the row group `try_next_reader` will yield + /// a reader for *after* the active one — never `None` just because + /// the decoder is mid-row-group. + #[test] + fn test_peek_next_row_group_during_decoding_row_group() { + // Two row groups (200 rows each), small batch size so the + // decoder stays in `DecodingRowGroup` across multiple + // `try_decode` calls within RG 0. + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .with_batch_size(100) + .build() + .unwrap(); + decoder + .push_range(test_file_range(), TEST_FILE_DATA.clone()) + .unwrap(); + + // First batch of RG 0 — after this, the decoder is in + // `DecodingRowGroup` state with the reader retained internally + // and one more 100-row batch still owed for RG 0. Peek must + // therefore look past the active RG to RG 1. + let _batch0 = expect_data(decoder.try_decode()); + assert!(!decoder.is_at_row_group_boundary()); + assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1)); + + // Drain the rest of RG 0; peek still reports RG 1. + let _batch1 = expect_data(decoder.try_decode()); + assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1)); + + // Move into RG 1 — peek now sees no further row groups. + let _batch2 = expect_data(decoder.try_decode()); + assert!(!decoder.is_at_row_group_boundary()); + assert_eq!(decoder.peek_next_row_group().unwrap(), None); + + let _batch3 = expect_data(decoder.try_decode()); + expect_finished(decoder.try_decode()); + } + + /// Peeking is a read-only operation: calling it repeatedly between + /// `try_next_reader` calls must never change which row group the + /// reader path actually produces. Drives the decoder all the way + /// through and asserts each peek/read pair agrees. + #[test] + fn test_peek_next_row_group_does_not_mutate_state() { + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .build() + .unwrap(); + + // Two row groups expected — drive both, asserting peek/read agree. + for expected_rg in [0usize, 1usize] { + assert!(decoder.is_at_row_group_boundary()); + + // Multiple peeks before reading must all agree, and must not + // disturb the upcoming read. + let first_peek = decoder.peek_next_row_group().unwrap(); + let second_peek = decoder.peek_next_row_group().unwrap(); + let third_peek = decoder.peek_next_row_group().unwrap(); + assert_eq!(first_peek, Some(expected_rg)); + assert_eq!(first_peek, second_peek); + assert_eq!(second_peek, third_peek); + + // Now read for real and confirm the decoder hands back exactly + // what peek promised. + let ranges = expect_needs_data(decoder.try_next_reader()); + push_ranges_to_decoder(&mut decoder, ranges); + let reader = expect_data(decoder.try_next_reader()); + for batch in reader { + let _ = batch.unwrap(); + } + } + + // Decoder is drained. Peek must agree with `try_next_reader`'s + // terminal state. + assert_eq!(decoder.peek_next_row_group().unwrap(), None); + expect_finished(decoder.try_next_reader()); + } + + /// `into_builder` between row groups recovers a builder for the + /// not-yet-decoded row groups; rebuilding it with a new row filter + /// applies that filter to the subsequent row groups while leaving the + /// already-decoded row group's results untouched. + /// + /// See the "Adaptive scans" section of [`ParquetPushDecoderBuilder`] for + /// the high-level overview. + #[test] + fn test_into_builder_installs_filter_between_row_groups() { + let schema_descr = test_file_parquet_metadata() + .file_metadata() + .schema_descr_ptr(); + let mut decoder = prefetched_decoder(1024); + + // Reader for row group 0 — no filter. + let reader0 = expect_data(decoder.try_next_reader()); + let batches0: Vec<_> = reader0.collect::>().unwrap(); + let batch0 = concat_batches(&TEST_BATCH.schema(), &batches0).unwrap(); + assert_eq!(batch0, TEST_BATCH.slice(0, 200)); + + // We're between row groups now. Rebuild with a filter on column "a". + assert!(decoder.is_at_row_group_boundary()); + assert_eq!(decoder.row_groups_remaining(), 1); + let filter = + ArrowPredicateFn::new(ProjectionMask::columns(&schema_descr, ["a"]), |batch| { + gt(batch.column(0), &Int64Array::new_scalar(250)) + }); + let mut decoder = decoder + .into_builder() + .unwrap() + .with_row_filter(RowFilter::new(vec![Box::new(filter)])) + .build() + .unwrap(); + + // Reader for row group 1 — filter applied. The rebuilt decoder kept + // the buffered bytes (see `test_into_builder_preserves_buffered_bytes`) + // so no data needs to be re-supplied. Column "a" in RG1 has values + // 200..399; `a > 250` keeps 251..399 = 149 rows. + let reader1 = expect_data(decoder.try_next_reader()); + let batches1: Vec<_> = reader1.collect::>().unwrap(); + let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap(); + assert_eq!(batch1, TEST_BATCH.slice(251, 149)); + expect_finished(decoder.try_next_reader()); + } + + /// `into_builder` is rejected while a row group's reader is being + /// drained (`DecodingRowGroup`); the error points at + /// `is_at_row_group_boundary`. + #[test] + fn test_into_builder_rejected_mid_row_group() { + let mut decoder = prefetched_decoder(50); + + // Decode one batch to land mid-row-group, inside `DecodingRowGroup` + // with an active reader — not a boundary. + expect_data(decoder.try_decode()); + assert!(!decoder.is_at_row_group_boundary()); + + let err = decoder.into_builder().unwrap_err(); + let err_msg = format!("{err}"); + assert!( + err_msg.contains("is_at_row_group_boundary"), + "unexpected error: {err_msg}" + ); + } + + /// `into_builder` is rejected once the decoder has finished. + #[test] + fn test_into_builder_rejected_on_finished_decoder() { + let mut decoder = prefetched_decoder(1024); + expect_data(decoder.try_decode()); + expect_data(decoder.try_decode()); + expect_finished(decoder.try_decode()); + assert!(!decoder.is_at_row_group_boundary()); + + let err = decoder.into_builder().unwrap_err(); + assert!( + format!("{err}").contains("finished"), + "unexpected error: {err}" + ); + } + + /// `try_next_reader` hands the active reader off to the caller and + /// transitions the decoder back to `ReadingRowGroup` — so the caller + /// can call `into_builder` even while still holding the returned + /// reader. (The handed-off reader has no link back to the decoder's + /// projection/filter; it has its own `ArrayReader` and `ReadPlan`.) + #[test] + fn test_into_builder_allowed_while_iterating_handed_off_reader() { + let mut decoder = prefetched_decoder(1024); + + let reader0 = expect_data(decoder.try_next_reader()); + // Decoder no longer owns the reader, so it considers itself + // "between row groups". + assert!(decoder.is_at_row_group_boundary()); + // Recovering the builder consumes the decoder but leaves `reader0` + // valid: iterating it is independent of the decoder's state. + let _builder = decoder.into_builder().unwrap(); + let batches: Vec<_> = reader0.collect::>().unwrap(); + let batch0 = concat_batches(&TEST_BATCH.schema(), &batches).unwrap(); + assert_eq!(batch0, TEST_BATCH.slice(0, 200)); + } + + /// `into_builder` recovers a builder for the *remaining* row groups and + /// carries the not-yet-consumed offset/limit budget, so a rebuilt + /// decoder resumes where the original left off rather than restarting. + #[test] + fn test_into_builder_resumes_remaining_budget() { + // limit = 250 spans both 200-row row groups: all 200 rows of RG0 + // plus the first 50 rows of RG1. + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .with_batch_size(1024) + .with_limit(250) + .build() + .unwrap(); + prefetch_test_file(&mut decoder); + + // RG0 contributes all 200 of its rows. + let reader0 = expect_data(decoder.try_next_reader()); + let batches0: Vec<_> = reader0.collect::>().unwrap(); + let batch0 = concat_batches(&TEST_BATCH.schema(), &batches0).unwrap(); + assert_eq!(batch0, TEST_BATCH.slice(0, 200)); + + // Rebuild without changing anything: the remaining 50-row limit and + // the not-yet-decoded RG1 must carry through (as do the buffers, so + // no data needs re-supplying). + assert!(decoder.is_at_row_group_boundary()); + let mut decoder = decoder.into_builder().unwrap().build().unwrap(); + + let reader1 = expect_data(decoder.try_next_reader()); + let batches1: Vec<_> = reader1.collect::>().unwrap(); + let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap(); + // Only the first 50 rows of RG1 (200..249) — the rest of the limit. + assert_eq!(batch1, TEST_BATCH.slice(200, 50)); + expect_finished(decoder.try_next_reader()); + } + + /// `into_builder` carries the decoder's buffered bytes across the + /// rebuild: the rebuilt decoder keeps them and does not re-request data + /// it already holds. + #[test] + fn test_into_builder_preserves_buffered_bytes() { + let mut decoder = prefetched_decoder(1024); + assert_eq!(decoder.buffered_bytes(), test_file_len()); + + // Drain RG0. + let reader0 = expect_data(decoder.try_next_reader()); + let _: Vec<_> = reader0.collect::>().unwrap(); + // RG1's bytes are still staged inside the decoder. + let buffered = decoder.buffered_bytes(); + assert!(buffered > 0); + + // Rebuilding via into_builder keeps the staged bytes. + let mut decoder = decoder.into_builder().unwrap().build().unwrap(); + assert_eq!(decoder.buffered_bytes(), buffered); + + // RG1's bytes are already buffered, so it decodes without a + // `NeedsData` round-trip. + let reader1 = expect_data(decoder.try_next_reader()); + let batches1: Vec<_> = reader1.collect::>().unwrap(); + let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap(); + assert_eq!(batch1, TEST_BATCH.slice(200, 200)); + expect_finished(decoder.try_next_reader()); + } + + /// Drive the decoder incrementally. Start with a narrow projection, + /// drain RG0, then `into_builder` and widen the projection to all three + /// columns. The rebuilt decoder's `NeedsData` for RG1 must request + /// bytes for *all three* columns, not just the originally-projected + /// "a". The expected ranges are hardcoded because `TEST_BATCH` and the + /// writer settings are static; this pins the layout cleanly without a + /// parallel reference decoder. + #[test] + fn test_into_builder_expand_projection_requests_new_bytes() { + let metadata = test_file_parquet_metadata(); + let schema_descr = metadata.file_metadata().schema_descr_ptr(); + + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(metadata) + .unwrap() + .with_batch_size(1024) + .with_projection(ProjectionMask::columns(&schema_descr, ["a"])) + .build() + .unwrap(); + + // RG0: incrementally satisfy the narrow request — a single + // contiguous range for the "a"-only projection. + let ranges_rg0 = expect_needs_data(decoder.try_next_reader()); + assert_eq!(ranges_rg0, vec![4..1860]); + push_ranges_to_decoder(&mut decoder, ranges_rg0); + + let reader0 = expect_data(decoder.try_next_reader()); + let batches0: Vec<_> = reader0.collect::>().unwrap(); + let batch0 = concat_batches(&batches0[0].schema(), &batches0).unwrap(); + assert_eq!(batch0, TEST_BATCH.slice(0, 200).project(&[0]).unwrap()); + + // Widen the projection at the boundary. + assert!(decoder.is_at_row_group_boundary()); + let mut decoder = decoder + .into_builder() + .unwrap() + .with_projection(ProjectionMask::columns(&schema_descr, ["a", "b", "c"])) + .build() + .unwrap(); + + // RG1 now requests "a", "b", and "c" column chunks: ~1.8KiB each + // for "a" and "b", ~7.5KiB for the StringView column "c". + let ranges_rg1 = expect_needs_data(decoder.try_next_reader()); + assert_eq!(ranges_rg1, vec![11062..12918, 12918..14774, 14774..22230]); + push_ranges_to_decoder(&mut decoder, ranges_rg1); + + let reader1 = expect_data(decoder.try_next_reader()); + let batches1: Vec<_> = reader1.collect::>().unwrap(); + let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap(); + assert_eq!(batch1, TEST_BATCH.slice(200, 200)); + expect_finished(decoder.try_next_reader()); + } + + /// Mirror of [`test_into_builder_expand_projection_requests_new_bytes`]: + /// start with the full projection, drain RG0, then `into_builder` and + /// narrow the projection to just column "a". RG1's `NeedsData` must + /// request only the single "a" column-chunk range, not the three a + /// wide projection would. + #[test] + fn test_into_builder_narrow_projection_requests_fewer_bytes() { + let metadata = test_file_parquet_metadata(); + let schema_descr = metadata.file_metadata().schema_descr_ptr(); + + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(metadata) + .unwrap() + .with_batch_size(1024) + .build() + .unwrap(); + + // RG0 with the default (full) projection — three column ranges. + let ranges_rg0 = expect_needs_data(decoder.try_next_reader()); + assert_eq!(ranges_rg0, vec![4..1860, 1860..3716, 3716..11062]); + push_ranges_to_decoder(&mut decoder, ranges_rg0); + + let reader0 = expect_data(decoder.try_next_reader()); + let batches0: Vec<_> = reader0.collect::>().unwrap(); + let batch0 = concat_batches(&TEST_BATCH.schema(), &batches0).unwrap(); + assert_eq!(batch0, TEST_BATCH.slice(0, 200)); + + // Narrow the projection at the boundary. + assert!(decoder.is_at_row_group_boundary()); + let mut decoder = decoder + .into_builder() + .unwrap() + .with_projection(ProjectionMask::columns(&schema_descr, ["a"])) + .build() + .unwrap(); + + // RG1 now requests column "a" only — a single 1856-byte range. + let ranges_rg1 = expect_needs_data(decoder.try_next_reader()); + assert_eq!(ranges_rg1, vec![11062..12918]); + push_ranges_to_decoder(&mut decoder, ranges_rg1); + + let reader1 = expect_data(decoder.try_next_reader()); + let batches1: Vec<_> = reader1.collect::>().unwrap(); + let batch1 = concat_batches(&batches1[0].schema(), &batches1).unwrap(); + assert_eq!(batch1, TEST_BATCH.slice(200, 200).project(&[0]).unwrap()); + expect_finished(decoder.try_next_reader()); + } + + /// Returns a batch with 400 rows, with 3 columns: "a", "b", "c" + /// + /// Note c is a different types (so the data page sizes will be different) + static TEST_BATCH: LazyLock = LazyLock::new(|| { + let a: ArrayRef = Arc::new(Int64Array::from_iter_values(0..400)); + let b: ArrayRef = Arc::new(Int64Array::from_iter_values(400..800)); + let c: ArrayRef = Arc::new(StringViewArray::from_iter_values((0..400).map(|i| { + if i % 2 == 0 { + format!("string_{i}") + } else { + format!("A string larger than 12 bytes and thus not inlined {i}") + } + }))); + + RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap() + }); + + /// Create a parquet file in memory for testing. + /// + /// See [`TEST_BATCH`] for the data in the file. + /// + /// Each column is written in 4 data pages, each with 100 rows, across 2 + /// row groups. Each column in each row group has two data pages. + /// + /// The data is split across row groups like this + /// + /// Column | Values | Data Page | Row Group + /// -------|------------------------|-----------|----------- + /// a | 0..99 | 1 | 0 + /// a | 100..199 | 2 | 0 + /// a | 200..299 | 1 | 1 + /// a | 300..399 | 2 | 1 + /// + /// b | 400..499 | 1 | 0 + /// b | 500..599 | 2 | 0 + /// b | 600..699 | 1 | 1 + /// b | 700..799 | 2 | 1 + /// + /// c | "string_0".."string_99" | 1 | 0 + /// c | "string_100".."string_199" | 2 | 0 + /// c | "string_200".."string_299" | 1 | 1 + /// c | "string_300".."string_399" | 2 | 1 + static TEST_FILE_DATA: LazyLock = LazyLock::new(|| { + let input_batch = &TEST_BATCH; + let mut output = Vec::new(); + + let writer_options = WriterProperties::builder() + .set_max_row_group_row_count(Some(200)) + .set_data_page_row_count_limit(100) + .build(); + let mut writer = + ArrowWriter::try_new(&mut output, input_batch.schema(), Some(writer_options)).unwrap(); + + // since the limits are only enforced on batch boundaries, write the input + // batch in chunks of 50 + let mut row_remain = input_batch.num_rows(); + while row_remain > 0 { + let chunk_size = row_remain.min(50); + let chunk = input_batch.slice(input_batch.num_rows() - row_remain, chunk_size); + writer.write(&chunk).unwrap(); + row_remain -= chunk_size; + } + writer.close().unwrap(); + Bytes::from(output) + }); + + /// Return the length of [`TEST_FILE_DATA`], in bytes + fn test_file_len() -> u64 { + TEST_FILE_DATA.len() as u64 + } + + /// Return a range that covers the entire [`TEST_FILE_DATA`] + fn test_file_range() -> Range { + 0..test_file_len() + } + + /// Return a slice of the test file data from the given range + pub fn test_file_slice(range: Range) -> Bytes { + let start: usize = range.start.try_into().unwrap(); + let end: usize = range.end.try_into().unwrap(); + TEST_FILE_DATA.slice(start..end) + } + + /// return the metadata for the test file + 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(); + 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, + ranges: Vec>, + ) { + let data = ranges + .iter() + .map(|range| test_file_slice(range.clone())) + .collect::>(); + metadata_decoder.push_ranges(ranges, data).unwrap(); + } + + /// Push the entire test file into `decoder`. + fn prefetch_test_file(decoder: &mut ParquetPushDecoder) { + decoder + .push_range(test_file_range(), TEST_FILE_DATA.clone()) + .unwrap(); + } + + /// Build a decoder over the test file with the given batch size and + /// prefetch the whole file into it. + fn prefetched_decoder(batch_size: usize) -> ParquetPushDecoder { + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .with_batch_size(batch_size) + .build() + .unwrap(); + prefetch_test_file(&mut decoder); + decoder + } + + fn push_ranges_to_decoder(decoder: &mut ParquetPushDecoder, ranges: Vec>) { + let data = ranges + .iter() + .map(|range| test_file_slice(range.clone())) + .collect::>(); + decoder.push_ranges(ranges, data).unwrap(); + } + + /// Expect that the [`DecodeResult`] is a [`DecodeResult::Data`] and return the corresponding element + fn expect_data(result: Result, ParquetError>) -> T { + match result.expect("Expected Ok(DecodeResult::Data(T))") { + DecodeResult::Data(data) => data, + result => panic!("Expected DecodeResult::Data, got {result:?}"), + } + } + + /// Expect that the [`DecodeResult`] is a [`DecodeResult::NeedsData`] and return the corresponding ranges + fn expect_needs_data( + result: Result, ParquetError>, + ) -> Vec> { + match result.expect("Expected Ok(DecodeResult::NeedsData{ranges})") { + DecodeResult::NeedsData(ranges) => ranges, + result => panic!("Expected DecodeResult::NeedsData, got {result:?}"), + } + } + + fn expect_finished(result: Result, ParquetError>) { + match result.expect("Expected Ok(DecodeResult::Finished)") { + DecodeResult::Finished => {} + result => panic!("Expected DecodeResult::Finished, got {result:?}"), + } + } +} diff --git a/vendor/parquet/src/arrow/push_decoder/reader_builder/data.rs b/vendor/parquet/src/arrow/push_decoder/reader_builder/data.rs new file mode 100644 index 0000000..6fbc209 --- /dev/null +++ b/vendor/parquet/src/arrow/push_decoder/reader_builder/data.rs @@ -0,0 +1,233 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`DataRequest`] tracks and holds data needed to construct InMemoryRowGroups + +use crate::arrow::ProjectionMask; +use crate::arrow::arrow_reader::RowSelection; +use crate::arrow::in_memory_row_group::{ColumnChunkData, FetchRanges, InMemoryRowGroup}; +use crate::errors::ParquetError; +use crate::file::metadata::ParquetMetaData; +use crate::file::page_index::offset_index::OffsetIndexMetaData; +use crate::file::reader::ChunkReader; +use crate::util::push_buffers::PushBuffers; +use bytes::Bytes; +use std::ops::Range; +use std::sync::Arc; + +/// Contains in-progress state to construct InMemoryRowGroups +/// +/// See [`DataRequestBuilder`] for creating new requests +#[derive(Debug)] +pub(super) struct DataRequest { + /// Any previously read column chunk data + column_chunks: Vec>>, + /// The ranges of data that are needed next + ranges: Vec>, + /// Optional page start offsets for each requested range. This is used + /// to create the relevant InMemoryRowGroup + page_start_offsets: Option>>, +} + +impl DataRequest { + /// return what ranges are still needed to satisfy this request. Returns an empty vec + /// if all ranges are satisfied + pub fn needed_ranges(&self, buffers: &PushBuffers) -> Vec> { + self.ranges + .iter() + .filter(|&range| !buffers.has_range(range)) + .cloned() + .collect() + } + + /// Returns the chunks from the buffers that satisfy this request + fn get_chunks(&self, buffers: &PushBuffers) -> Result, ParquetError> { + self.ranges + .iter() + .map(|range| { + let length: usize = (range.end - range.start) + .try_into() + .expect("overflow for offset"); + // should have all the data due to the check above + buffers.get_bytes(range.start, length).map_err(|e| { + ParquetError::General(format!( + "Internal Error missing data for range {range:?} in buffers: {e}", + )) + }) + }) + .collect() + } + + /// Create a new InMemoryRowGroup, and fill it with provided data + /// + /// Assumes that all needed data is present in the buffers + /// and clears any explicitly requested ranges + pub fn try_into_in_memory_row_group<'a>( + self, + row_group_idx: usize, + row_count: usize, + parquet_metadata: &'a ParquetMetaData, + projection: &ProjectionMask, + buffers: &mut PushBuffers, + ) -> Result, ParquetError> { + let chunks = self.get_chunks(buffers)?; + + let Self { + column_chunks, + ranges, + page_start_offsets, + } = self; + + // Create an InMemoryRowGroup to hold the column chunks, this is a + // temporary structure used to tell the ArrowReaders what pages are + // needed for decoding + let mut in_memory_row_group = InMemoryRowGroup { + row_count, + column_chunks, + offset_index: get_offset_index(parquet_metadata, row_group_idx), + row_group_idx, + metadata: parquet_metadata, + }; + + in_memory_row_group.fill_column_chunks(projection, page_start_offsets, chunks); + + // Clear the ranges that were explicitly requested + buffers.clear_ranges(&ranges); + + Ok(in_memory_row_group) + } +} + +/// Builder for [`DataRequest`] +pub(super) struct DataRequestBuilder<'a> { + /// The row group index + row_group_idx: usize, + /// The number of rows in the row group + row_count: usize, + /// The batch size to read + batch_size: usize, + /// The parquet metadata + parquet_metadata: &'a ParquetMetaData, + /// The projection mask (which columns to read) + projection: &'a ProjectionMask, + /// Optional row selection to apply + selection: Option<&'a RowSelection>, + /// Optional projection mask if using + /// [`RowGroupCache`](crate::arrow::array_reader::RowGroupCache) + /// for caching decoded columns. + cache_projection: Option<&'a ProjectionMask>, + /// Any previously read column chunks + column_chunks: Option>>>, +} + +impl<'a> DataRequestBuilder<'a> { + pub(super) fn new( + row_group_idx: usize, + row_count: usize, + batch_size: usize, + parquet_metadata: &'a ParquetMetaData, + projection: &'a ProjectionMask, + ) -> Self { + Self { + row_group_idx, + row_count, + batch_size, + parquet_metadata, + projection, + selection: None, + cache_projection: None, + column_chunks: None, + } + } + + /// Set an optional row selection to apply + pub(super) fn with_selection(mut self, selection: Option<&'a RowSelection>) -> Self { + self.selection = selection; + self + } + + /// set columns to cache, if any + pub(super) fn with_cache_projection( + mut self, + cache_projection: Option<&'a ProjectionMask>, + ) -> Self { + self.cache_projection = cache_projection; + self + } + + /// Provide any previously read column chunks + pub(super) fn with_column_chunks( + mut self, + column_chunks: Option>>>, + ) -> Self { + self.column_chunks = column_chunks; + self + } + + pub(crate) fn build(self) -> DataRequest { + let Self { + row_group_idx, + row_count, + batch_size, + parquet_metadata, + projection, + selection, + cache_projection, + column_chunks, + } = self; + + let row_group_meta_data = parquet_metadata.row_group(row_group_idx); + + // If no previously read column chunks are provided, create a new location to hold them + let column_chunks = + column_chunks.unwrap_or_else(|| vec![None; row_group_meta_data.columns().len()]); + + // Create an InMemoryRowGroup to hold the column chunks, this is a + // temporary structure used to tell the ArrowReaders what pages are + // needed for decoding + let row_group = InMemoryRowGroup { + row_count, + column_chunks, + offset_index: get_offset_index(parquet_metadata, row_group_idx), + row_group_idx, + metadata: parquet_metadata, + }; + + let FetchRanges { + ranges, + page_start_offsets, + } = row_group.fetch_ranges(projection, selection, batch_size, cache_projection); + + DataRequest { + // Save any previously read column chunks + column_chunks: row_group.column_chunks, + ranges, + page_start_offsets, + } + } +} + +fn get_offset_index( + parquet_metadata: &ParquetMetaData, + row_group_idx: usize, +) -> Option<&[OffsetIndexMetaData]> { + parquet_metadata + .offset_index() + // filter out empty offset indexes (old versions specified Some(vec![]) when no present) + .filter(|index| !index.is_empty()) + .map(|x| x[row_group_idx].as_slice()) +} diff --git a/vendor/parquet/src/arrow/push_decoder/reader_builder/filter.rs b/vendor/parquet/src/arrow/push_decoder/reader_builder/filter.rs new file mode 100644 index 0000000..219ab06 --- /dev/null +++ b/vendor/parquet/src/arrow/push_decoder/reader_builder/filter.rs @@ -0,0 +1,150 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`FilterInfo`] state machine for evaluating row filters + +use crate::arrow::ProjectionMask; +use crate::arrow::array_reader::{CacheOptionsBuilder, RowGroupCache}; +use crate::arrow::arrow_reader::{ArrowPredicate, RowFilter}; +use std::num::NonZeroUsize; +use std::sync::{Arc, RwLock}; + +/// State machine for evaluating a sequence of predicates. +/// +/// The `FilterInfo` owns the [`RowFilter`] being evaluated and tracks the current +/// predicate to evaluate. +#[derive(Debug)] +pub(super) struct FilterInfo { + /// The predicates to evaluate, in order + /// + /// RowFilter is owned by `FilterInfo` because they may be mutated as part + /// of evaluation. Specifically, [`ArrowPredicate`] requires &mut self for + /// evaluation. + filter: RowFilter, + /// The next filter to be evaluated + next_predicate: NonZeroUsize, + /// Previously computed filter results + cache_info: CacheInfo, +} + +/// Predicate cache +/// +/// Note this is basically the same as CacheOptionsBuilder +/// but it owns the ProjectionMask and RowGroupCache +#[derive(Debug)] +pub(super) struct CacheInfo { + /// The columns to cache in the predicate cache. + /// Normally these are the columns that filters may look at such that + /// if we have a filter like `(a + 10 > 5) AND (a + b = 0)` we cache `a` to avoid re-reading it between evaluating `a + 10 > 5` and `a + b = 0`. + cache_projection: ProjectionMask, + row_group_cache: Arc>, +} + +impl CacheInfo { + pub(super) fn new( + cache_projection: ProjectionMask, + row_group_cache: Arc>, + ) -> Self { + Self { + cache_projection, + row_group_cache, + } + } + + pub(super) fn builder(&self) -> CacheOptionsBuilder<'_> { + CacheOptionsBuilder::new(&self.cache_projection, &self.row_group_cache) + } +} + +pub(super) enum AdvanceResult { + /// Advanced to the next predicate + Continue(FilterInfo), + /// No more predicates returns the row filter and cache info + Done(RowFilter, CacheInfo), +} + +impl FilterInfo { + /// Create a new FilterInfo + pub(super) fn new(filter: RowFilter, cache_info: CacheInfo) -> Self { + Self { + filter, + next_predicate: NonZeroUsize::new(1).expect("1 is always non-zero"), + cache_info, + } + } + + /// Advance to the next predicate + /// + /// Returns + /// * [`AdvanceResult::Continue`] returning the `FilterInfo` if there are + /// more predicate to evaluate. + /// * [`AdvanceResult::Done`] with the inner [`RowFilter`] and [`CacheInfo]` + /// if there are no more predicates + pub(super) fn advance(mut self) -> AdvanceResult { + if self.next_predicate.get() >= self.filter.predicates.len() { + AdvanceResult::Done(self.filter, self.cache_info) + } else { + self.next_predicate = self + .next_predicate + .checked_add(1) + .expect("no usize overflow"); + AdvanceResult::Continue(self) + } + } + + /// Return a mutable reference to the current predicate + pub(super) fn current_mut(&mut self) -> &mut dyn ArrowPredicate { + self.filter + .predicates + .get_mut(self.next_predicate.get() - 1) + // advance ensures next_predicate is always in bounds + .unwrap() + .as_mut() + } + + /// Return the current predicate to evaluate + pub(super) fn current(&self) -> &dyn ArrowPredicate { + self.filter + .predicates + .get(self.next_predicate.get() - 1) + // advance ensures next_predicate is always in bounds + .unwrap() + .as_ref() + } + + /// Returns `true` if the current predicate is the last one in the chain + /// (i.e. the next call to [`Self::advance`] will return + /// [`AdvanceResult::Done`]). + pub(super) fn is_last(&self) -> bool { + self.next_predicate.get() == self.filter.predicates.len() + } + + /// Return a reference to the cache projection + pub(super) fn cache_projection(&self) -> &ProjectionMask { + &self.cache_info.cache_projection + } + + /// Return a cache builder to save the results of predicate evaluation + pub(super) fn cache_builder(&self) -> CacheOptionsBuilder<'_> { + self.cache_info.builder() + } + + /// Returns the inner filter, consuming this FilterInfo + pub(super) fn into_filter(self) -> RowFilter { + self.filter + } +} diff --git a/vendor/parquet/src/arrow/push_decoder/reader_builder/mod.rs b/vendor/parquet/src/arrow/push_decoder/reader_builder/mod.rs new file mode 100644 index 0000000..200974a --- /dev/null +++ b/vendor/parquet/src/arrow/push_decoder/reader_builder/mod.rs @@ -0,0 +1,1053 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +mod data; +mod filter; + +use crate::arrow::ProjectionMask; +use crate::arrow::array_reader::{ArrayReaderBuilder, CacheOptions, RowGroupCache}; +use crate::arrow::arrow_reader::metrics::ArrowReaderMetrics; +use crate::arrow::arrow_reader::selection::{LoadedRowRanges, RowSelectionStrategy}; +use crate::arrow::arrow_reader::{ + ParquetRecordBatchReader, PredicateOptions, ReadPlanBuilder, RowFilter, RowSelection, + RowSelectionPolicy, +}; +use crate::arrow::in_memory_row_group::ColumnChunkData; +use crate::arrow::push_decoder::reader_builder::data::DataRequestBuilder; +use crate::arrow::push_decoder::reader_builder::filter::CacheInfo; +use crate::arrow::schema::ParquetField; +use crate::errors::ParquetError; +use crate::file::metadata::ParquetMetaData; +use crate::file::page_index::offset_index::OffsetIndexMetaData; +use crate::util::push_buffers::PushBuffers; +use bytes::Bytes; +use data::DataRequest; +use filter::AdvanceResult; +use filter::FilterInfo; +use std::ops::Range; +use std::sync::{Arc, RwLock}; + +/// The current row group being read, its read plan, and its offset/limit budget. +#[derive(Debug)] +struct RowGroupInfo { + row_group_idx: usize, + row_count: usize, + plan_builder: ReadPlanBuilder, + budget: RowBudget, +} + +/// This is the inner state machine for reading a single row group. +#[derive(Debug)] +enum RowGroupDecoderState { + Start { + row_group_info: RowGroupInfo, + }, + /// Planning filters, but haven't yet requested data to evaluate them + Filters { + row_group_info: RowGroupInfo, + /// Any previously read column chunk data from prior filters + column_chunks: Option>>>, + filter_info: FilterInfo, + }, + /// Needs data to evaluate current filter + WaitingOnFilterData { + row_group_info: RowGroupInfo, + filter_info: FilterInfo, + data_request: DataRequest, + }, + /// Know what data to actually read, after all predicates + StartData { + row_group_info: RowGroupInfo, + /// Any previously read column chunk data from the filtering phase + column_chunks: Option>>>, + /// Any cached filter results + cache_info: Option, + }, + /// Needs data to proceed with reading the output + WaitingOnData { + row_group_info: RowGroupInfo, + data_request: DataRequest, + /// Any cached filter results + cache_info: Option, + }, + /// Finished (or not yet started) reading this group + Finished, +} + +/// Running offset/limit budget shared across row groups. +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +pub(crate) struct RowBudget { + offset: Option, + limit: Option, +} + +impl RowBudget { + pub(crate) fn new(offset: Option, limit: Option) -> Self { + Self { offset, limit } + } + + pub(crate) fn is_exhausted(self) -> bool { + matches!(self.limit, Some(0)) + } + + /// The offset still to be skipped before the next readable row group. + pub(crate) fn offset(self) -> Option { + self.offset + } + + /// The number of output rows still permitted across the remaining row groups. + pub(crate) fn limit(self) -> Option { + self.limit + } + + /// Returns how many selected rows remain after applying this budget. + pub(crate) fn rows_after(self, rows_before_budget: usize) -> usize { + let rows_after_offset = rows_before_budget.saturating_sub(self.offset.unwrap_or(0)); + match self.limit { + Some(limit) => rows_after_offset.min(limit), + None => rows_after_offset, + } + } + + /// Returns the number of selected rows needed before applying the offset. + fn selected_row_limit(self) -> Option { + self.limit + .map(|limit| limit.saturating_add(self.offset.unwrap_or(0))) + } + + fn apply_to_plan(self, plan_builder: ReadPlanBuilder, row_count: usize) -> BudgetedReadPlan { + let rows_before_budget = plan_builder.num_rows_selected().unwrap_or(row_count); + let plan_builder = plan_builder + .limited(row_count) + .with_offset(self.offset) + .with_limit(self.limit) + .build_limited(); + let rows_after_budget = self.rows_after(rows_before_budget); + + BudgetedReadPlan { + plan_builder, + rows_before_budget, + rows_after_budget, + remaining_budget: self.advance(rows_before_budget, rows_after_budget), + } + } + + /// Advance the budget past one row group. + /// + /// `rows_before_budget` is the number of rows selected before applying the + /// budget, and `rows_after_budget` is the number retained for output from + /// this row group. + pub(crate) fn advance(mut self, rows_before_budget: usize, rows_after_budget: usize) -> Self { + if let Some(offset) = &mut self.offset { + // Reduction is either because of offset or limit, as limit is applied + // after offset has been "exhausted" can just use saturating sub here. + *offset = offset.saturating_sub(rows_before_budget - rows_after_budget); + } + + if rows_after_budget != 0 { + if let Some(limit) = &mut self.limit { + *limit -= rows_after_budget; + } + } + + self + } +} + +#[derive(Debug)] +struct BudgetedReadPlan { + /// Read plan after applying this row group's share of the offset/limit budget. + plan_builder: ReadPlanBuilder, + /// Number of rows selected by row selection and predicates before applying + /// this row group's offset/limit budget. + rows_before_budget: usize, + /// Number of selected rows that remain to be read after applying this row + /// group's offset/limit budget. + rows_after_budget: usize, + /// Budget remaining for later row groups. + remaining_budget: RowBudget, +} + +#[derive(Debug)] +pub(crate) enum RowGroupBuildResult { + /// The active row group is complete without producing a reader. + Finished { + /// Budget remaining after applying this row group's selection. + remaining_budget: RowBudget, + }, + /// More bytes are needed before the active row group can make progress. + NeedsData(Vec>), + /// The active row group produced a reader. + Data { + batch_reader: ParquetRecordBatchReader, + /// Budget remaining after applying this row group's selection. + remaining_budget: RowBudget, + }, +} + +/// Result of a state transition +#[derive(Debug)] +struct NextState { + next_state: RowGroupDecoderState, + /// result to return, if any + /// + /// * `Some`: the processing should stop and return the result + /// * `None`: processing should continue + result: Option, +} + +impl NextState { + /// The next state with no result. + /// + /// This indicates processing should continue + fn again(next_state: RowGroupDecoderState) -> Self { + Self { + next_state, + result: None, + } + } + + /// Create a NextState with a result that should be returned + fn result(next_state: RowGroupDecoderState, result: RowGroupBuildResult) -> Self { + Self { + next_state, + result: Some(result), + } + } +} + +/// Builder for [`ParquetRecordBatchReader`] for a single row group +/// +/// This struct drives the main state machine for decoding each row group -- it +/// determines what data is needed, and then assembles the +/// `ParquetRecordBatchReader` when all data is available. +#[derive(Debug)] +pub(crate) struct RowGroupReaderBuilder { + /// The output batch size + batch_size: usize, + + /// What columns to project (produce in each output batch) + projection: ProjectionMask, + + /// The Parquet file metadata + metadata: Arc, + + /// Top level parquet schema and arrow schema mapping + fields: Option>, + + /// Optional filter + filter: Option, + + /// The size in bytes of the predicate cache to use + /// + /// See [`RowGroupCache`] for details. + max_predicate_cache_size: usize, + + /// The metrics collector + metrics: ArrowReaderMetrics, + + /// Strategy for materialising row selections + row_selection_policy: RowSelectionPolicy, + + /// Current state of the decoder. + /// + /// It is taken when processing, and must be put back before returning + /// it is a bug error if it is not put back after transitioning states. + state: Option, + + /// The underlying data store + buffers: PushBuffers, +} + +/// The parts of a [`RowGroupReaderBuilder`] needed to rebuild it, recovered by +/// [`RowGroupReaderBuilder::into_parts`]. +/// +/// `metadata` is not included: it is a whole-file property carried alongside +/// `schema` in `RemainingRowGroupsParts`. +#[derive(Debug)] +pub(crate) struct RowGroupReaderBuilderParts { + pub batch_size: usize, + pub projection: ProjectionMask, + pub fields: Option>, + pub filter: Option, + pub max_predicate_cache_size: usize, + pub metrics: ArrowReaderMetrics, + pub row_selection_policy: RowSelectionPolicy, + /// Bytes already pushed into the decoder, carried across a rebuild so they + /// are not re-requested. + pub buffers: PushBuffers, +} + +impl RowGroupReaderBuilder { + /// Create a new RowGroupReaderBuilder + #[expect(clippy::too_many_arguments)] + pub(crate) fn new( + batch_size: usize, + projection: ProjectionMask, + metadata: Arc, + fields: Option>, + filter: Option, + metrics: ArrowReaderMetrics, + max_predicate_cache_size: usize, + buffers: PushBuffers, + row_selection_policy: RowSelectionPolicy, + ) -> Self { + Self { + batch_size, + projection, + metadata, + fields, + filter, + metrics, + max_predicate_cache_size, + row_selection_policy, + state: Some(RowGroupDecoderState::Finished), + buffers, + } + } + + /// Decompose into [`RowGroupReaderBuilderParts`] so the builder can be + /// reconstructed. The runtime decode `state` is discarded; `metadata` is + /// recovered from the frontier instead (see `RemainingRowGroups::into_parts`). + pub(crate) fn into_parts(self) -> RowGroupReaderBuilderParts { + // If a new field is added to `RowGroupReaderBuilder`, it must be added here and in `RowGroupReaderBuilderParts`, + // or at least evaluate how it should be handled in the decomposition and reconstruction of the builder. + let Self { + batch_size, + projection, + metadata: _, + fields, + filter, + max_predicate_cache_size, + metrics, + row_selection_policy, + state: _, + buffers, + } = self; + RowGroupReaderBuilderParts { + batch_size, + projection, + fields, + filter, + max_predicate_cache_size, + metrics, + row_selection_policy, + buffers, + } + } + + /// Push new data buffers that can be used to satisfy pending requests + pub fn push_data(&mut self, ranges: Vec>, buffers: Vec) { + self.buffers.push_ranges(ranges, buffers); + } + + /// True iff the inner state is `Finished`. This is the only state in + /// which it is safe to decompose the builder via [`Self::into_parts`], + /// because no `RowGroupInfo`, `FilterInfo`, or in-flight `DataRequest` + /// is referencing the row-group-scoped decode state. + pub(crate) fn is_finished(&self) -> bool { + matches!(self.state, Some(RowGroupDecoderState::Finished)) + } + + /// Returns the total number of buffered bytes available + pub fn buffered_bytes(&self) -> u64 { + self.buffers.buffered_bytes() + } + + /// Clear any staged ranges currently buffered for future decode work. + pub fn clear_all_ranges(&mut self) { + self.buffers.clear_all_ranges(); + } + + /// take the current state, leaving None in its place. + /// + /// Returns an error if there the state wasn't put back after the previous + /// call to [`Self::take_state`]. + /// + /// Any code that calls this method must ensure that the state is put back + /// before returning, otherwise the reader will error next time it is called + fn take_state(&mut self) -> Result { + self.state.take().ok_or_else(|| { + ParquetError::General(String::from( + "Internal Error: RowGroupReader in invalid state", + )) + }) + } + + /// Returns true if this builder is currently decoding a row group. + pub(crate) fn has_active_row_group(&self) -> bool { + !matches!(self.state, Some(RowGroupDecoderState::Finished)) + } + + /// Setup this reader to read the next row group + pub(crate) fn next_row_group( + &mut self, + row_group_idx: usize, + row_count: usize, + selection: Option, + budget: RowBudget, + ) -> Result<(), ParquetError> { + let state = self.take_state()?; + if !matches!(state, RowGroupDecoderState::Finished) { + return Err(ParquetError::General(format!( + "Internal Error: next_row_group called while still reading a row group. Expected Finished state, got {state:?}" + ))); + } + let plan_builder = ReadPlanBuilder::new(self.batch_size) + .with_selection(selection) + .with_row_selection_policy(self.row_selection_policy); + + let row_group_info = RowGroupInfo { + row_group_idx, + row_count, + plan_builder, + budget, + }; + + self.state = Some(RowGroupDecoderState::Start { row_group_info }); + Ok(()) + } + + /// Try to build the next `ParquetRecordBatchReader` for the active row group. + /// + /// Returns [`RowGroupBuildResult::NeedsData`] if more data is needed, + /// [`RowGroupBuildResult::Data`] if a reader is ready, or + /// [`RowGroupBuildResult::Finished`] if the row group completed without + /// producing a reader. + pub(crate) fn try_build(&mut self) -> Result { + loop { + let current_state = self.take_state()?; + // Try to transition the decoder. + match self.try_transition(current_state)? { + // Either produced a batch reader, needed input, or finished + NextState { + next_state, + result: Some(result), + } => { + // put back the next state + self.state = Some(next_state); + return Ok(result); + } + // completed one internal state, maybe can proceed further + NextState { + next_state, + result: None, + } => { + // continue processing + self.state = Some(next_state); + } + } + } + } + + /// Current state --> next state + optional output + /// + /// This is the main state transition function for the row group reader + /// and encodes the row group decoding state machine. + /// + /// # Notes + /// + /// This structure is used to reduce the indentation level of the main loop + /// in try_build + fn try_transition( + &mut self, + current_state: RowGroupDecoderState, + ) -> Result { + let result = match current_state { + RowGroupDecoderState::Start { row_group_info } => { + debug_assert!( + !row_group_info.budget.is_exhausted(), + "RowGroupFrontier should not hand off row groups after the output limit is exhausted" + ); + + let column_chunks = None; // no prior column chunks + + let Some(filter) = self.filter.take() else { + // no filter, start trying to read data immediately + return Ok(NextState::again(RowGroupDecoderState::StartData { + row_group_info, + column_chunks, + cache_info: None, + })); + }; + // no predicates in filter, so start reading immediately + if filter.predicates.is_empty() { + return Ok(NextState::again(RowGroupDecoderState::StartData { + row_group_info, + column_chunks, + cache_info: None, + })); + }; + + // we have predicates to evaluate + let cache_projection = + self.compute_cache_projection(row_group_info.row_group_idx, &filter); + + let cache_info = CacheInfo::new( + cache_projection, + Arc::new(RwLock::new(RowGroupCache::new( + self.batch_size, + self.max_predicate_cache_size, + ))), + ); + + let filter_info = FilterInfo::new(filter, cache_info); + NextState::again(RowGroupDecoderState::Filters { + row_group_info, + filter_info, + column_chunks, + }) + } + // need to evaluate filters + RowGroupDecoderState::Filters { + row_group_info, + column_chunks, + filter_info, + } => { + let RowGroupInfo { + row_group_idx, + row_count, + plan_builder, + budget, + } = row_group_info; + + // If nothing is selected, we are done with this row group + if !plan_builder.selects_any() { + // ruled out entire row group + self.filter = Some(filter_info.into_filter()); + return Ok(NextState::result( + RowGroupDecoderState::Finished, + RowGroupBuildResult::Finished { + remaining_budget: budget, + }, + )); + } + + // Make a request for the data needed to evaluate the current predicate + let predicate = filter_info.current(); + + // need to fetch pages the column needs for decoding, figure + // that out based on the current selection and projection + let data_request = DataRequestBuilder::new( + row_group_idx, + row_count, + self.batch_size, + &self.metadata, + predicate.projection(), // use the predicate's projection + ) + .with_selection(plan_builder.selection()) + // Cached output columns reuse these predicate-stage chunks. Expand their + // selection to cache batch boundaries so a cache miss can safely fetch a + // complete batch from the retained sparse column data. + .with_cache_projection(Some(filter_info.cache_projection())) + .with_column_chunks(column_chunks) + .build(); + + let row_group_info = RowGroupInfo { + row_group_idx, + row_count, + plan_builder, + budget, + }; + + NextState::again(RowGroupDecoderState::WaitingOnFilterData { + row_group_info, + filter_info, + data_request, + }) + } + RowGroupDecoderState::WaitingOnFilterData { + row_group_info, + data_request, + mut filter_info, + } => { + // figure out what ranges we still need + let needed_ranges = data_request.needed_ranges(&self.buffers); + if !needed_ranges.is_empty() { + // still need data + return Ok(NextState::result( + RowGroupDecoderState::WaitingOnFilterData { + row_group_info, + filter_info, + data_request, + }, + RowGroupBuildResult::NeedsData(needed_ranges), + )); + } + + // otherwise we have all the data we need to evaluate the predicate + let RowGroupInfo { + row_group_idx, + row_count, + mut plan_builder, + budget, + } = row_group_info; + + let predicate = filter_info.current(); + + let row_group = data_request.try_into_in_memory_row_group( + row_group_idx, + row_count, + &self.metadata, + predicate.projection(), + &mut self.buffers, + )?; + + let cache_options = filter_info.cache_builder().producer(); + + let array_reader = ArrayReaderBuilder::new(&row_group, &self.metrics) + .with_batch_size(self.batch_size) + .with_cache_options(Some(&cache_options)) + .with_parquet_metadata(&self.metadata) + .build_array_reader(self.fields.as_deref(), predicate.projection())?; + + // Auto resolution and loaded ranges are projection-specific, so restore the + // configured policy before preparing each predicate. + plan_builder = plan_builder.with_row_selection_policy(self.row_selection_policy); + + // Prepare selection execution for pages pruned during fetch. + plan_builder = prepare_selection_for_page_skipping( + plan_builder, + predicate.projection(), + self.row_group_offset_index(row_group_idx), + row_count, + ); + + // When this is the final predicate in the chain and an output + // limit is set, tell the filter evaluation to stop once enough + // matching rows have been accumulated. + let predicate_limit = filter_info + .is_last() + .then(|| budget.selected_row_limit()) + .flatten(); + + // Evaluate the filter via `with_predicate_options`, opting into + // early termination when this is the final predicate and an + // output limit was set. + let mut predicate_options = + PredicateOptions::new(array_reader, filter_info.current_mut()); + if let Some(limit) = predicate_limit { + predicate_options = predicate_options.with_limit(limit, row_count); + } + plan_builder = plan_builder.with_predicate_options(predicate_options)?; + + let row_group_info = RowGroupInfo { + row_group_idx, + row_count, + plan_builder, + budget, + }; + + // Take back the column chunks that were read + let column_chunks = Some(row_group.column_chunks); + + // advance to the next predicate, if any + match filter_info.advance() { + AdvanceResult::Continue(filter_info) => { + NextState::again(RowGroupDecoderState::Filters { + row_group_info, + column_chunks, + filter_info, + }) + } + // done with predicates, proceed to reading data + AdvanceResult::Done(filter, cache_info) => { + // remember we need to put back the filter + assert!(self.filter.is_none()); + self.filter = Some(filter); + NextState::again(RowGroupDecoderState::StartData { + row_group_info, + column_chunks, + cache_info: Some(cache_info), + }) + } + } + } + RowGroupDecoderState::StartData { + row_group_info, + column_chunks, + cache_info, + } => { + let RowGroupInfo { + row_group_idx, + row_count, + plan_builder, + budget, + } = row_group_info; + + let BudgetedReadPlan { + mut plan_builder, + rows_before_budget, + rows_after_budget, + remaining_budget, + } = budget.apply_to_plan(plan_builder, row_count); + + if rows_before_budget == 0 { + // ruled out entire row group + return Ok(NextState::result( + RowGroupDecoderState::Finished, + RowGroupBuildResult::Finished { remaining_budget }, + )); + } + + if rows_after_budget == 0 { + // no rows left after applying limit/offset + return Ok(NextState::result( + RowGroupDecoderState::Finished, + RowGroupBuildResult::Finished { remaining_budget }, + )); + } + + let data_request = DataRequestBuilder::new( + row_group_idx, + row_count, + self.batch_size, + &self.metadata, + &self.projection, + ) + .with_selection(plan_builder.selection()) + .with_column_chunks(column_chunks) + // Final projection fetch shouldn't expand selection for cache + // so don't call with_cache_projection here + .build(); + + plan_builder = plan_builder.with_row_selection_policy(self.row_selection_policy); + + plan_builder = prepare_selection_for_page_skipping( + plan_builder, + &self.projection, + self.row_group_offset_index(row_group_idx), + row_count, + ); + + let row_group_info = RowGroupInfo { + row_group_idx, + row_count, + plan_builder, + budget: remaining_budget, + }; + + NextState::again(RowGroupDecoderState::WaitingOnData { + row_group_info, + data_request, + cache_info, + }) + } + // Waiting on data to proceed with reading the output + RowGroupDecoderState::WaitingOnData { + row_group_info, + data_request, + cache_info, + } => { + let needed_ranges = data_request.needed_ranges(&self.buffers); + if !needed_ranges.is_empty() { + // still need data + return Ok(NextState::result( + RowGroupDecoderState::WaitingOnData { + row_group_info, + data_request, + cache_info, + }, + RowGroupBuildResult::NeedsData(needed_ranges), + )); + } + + // otherwise we have all the data we need to proceed + let RowGroupInfo { + row_group_idx, + row_count, + plan_builder, + budget, + } = row_group_info; + + let row_group = data_request.try_into_in_memory_row_group( + row_group_idx, + row_count, + &self.metadata, + &self.projection, + &mut self.buffers, + )?; + + let plan = plan_builder.build(); + + // if we have any cached results, connect them up + let array_reader_builder = ArrayReaderBuilder::new(&row_group, &self.metrics) + .with_batch_size(self.batch_size) + .with_parquet_metadata(&self.metadata); + let array_reader = if let Some(cache_info) = cache_info.as_ref() { + let cache_options: CacheOptions = cache_info.builder().consumer(); + array_reader_builder + .with_cache_options(Some(&cache_options)) + .build_array_reader(self.fields.as_deref(), &self.projection) + } else { + array_reader_builder + .build_array_reader(self.fields.as_deref(), &self.projection) + }?; + + let reader = ParquetRecordBatchReader::new(array_reader, plan); + NextState::result( + RowGroupDecoderState::Finished, + RowGroupBuildResult::Data { + batch_reader: reader, + remaining_budget: budget, + }, + ) + } + RowGroupDecoderState::Finished => { + return Err(ParquetError::General(String::from( + "Internal Error: try_build called without an active row group", + ))); + } + }; + Ok(result) + } + + /// Which columns should be cached? + /// + /// Returns the columns that are used by the filters *and* then used in the + /// final projection, excluding any nested columns. + fn compute_cache_projection(&self, row_group_idx: usize, filter: &RowFilter) -> ProjectionMask { + let meta = self.metadata.row_group(row_group_idx); + match self.compute_cache_projection_inner(filter) { + Some(projection) => projection, + None => ProjectionMask::none(meta.columns().len()), + } + } + + fn compute_cache_projection_inner(&self, filter: &RowFilter) -> Option { + // Do not compute the projection mask if the predicate cache is disabled + if self.max_predicate_cache_size == 0 { + return None; + } + let mut cache_projection = filter.predicates.first()?.projection().clone(); + for predicate in filter.predicates.iter() { + cache_projection.union(predicate.projection()); + } + cache_projection.intersect(&self.projection); + self.exclude_nested_columns_from_cache(&cache_projection) + } + + /// Exclude leaves belonging to roots that span multiple parquet leaves (i.e. nested columns) + fn exclude_nested_columns_from_cache(&self, mask: &ProjectionMask) -> Option { + mask.without_nested_types(self.metadata.file_metadata().schema_descr()) + } + + /// Get the offset index for the specified row group, if any + fn row_group_offset_index(&self, row_group_idx: usize) -> Option<&[OffsetIndexMetaData]> { + self.metadata + .offset_index() + .filter(|index| !index.is_empty()) + .and_then(|index| index.get(row_group_idx)) + .map(|columns| columns.as_slice()) + } +} + +/// Prepare row selection execution when page pruning produced sparse column data. +/// +/// Some pages can be skipped during row-group construction if they are not read +/// by the selections. This means that the data pages for those rows are never +/// loaded and definition/repetition levels are never read. When using +/// `RowSelections` selection works because `skip_records()` handles this +/// case and skips the page accordingly. +/// +/// However, with the current mask design, all values covered by a mask chunk +/// are decoded before the mask filter is applied. Thus a chunk cannot cross a +/// page that was skipped during row-group construction. +/// +/// A simple example: +/// * the page size is 2, the mask is 100001, row selection should be read(1) skip(4) read(1) +/// * the `ColumnChunkData` would be page1(10), page2(skipped), page3(01) +/// +/// Mask execution records the row ranges loaded for every projected column, so +/// each mask chunk stays within loaded data and `skip_records()` crosses the +/// gaps. This applies both to an explicit mask policy and to Auto when it +/// resolves to mask execution. +fn prepare_selection_for_page_skipping( + plan_builder: ReadPlanBuilder, + projection_mask: &ProjectionMask, + offset_index: Option<&[OffsetIndexMetaData]>, + total_rows: usize, +) -> ReadPlanBuilder { + match plan_builder.resolve_selection_strategy() { + RowSelectionStrategy::Mask => { + let loaded = loaded_row_ranges_for_projection( + plan_builder.selection(), + projection_mask, + offset_index, + total_rows, + ); + plan_builder + .with_row_selection_policy(RowSelectionPolicy::Mask) + .with_loaded_row_ranges(loaded) + } + RowSelectionStrategy::Selectors => { + plan_builder.with_row_selection_policy(RowSelectionPolicy::Selectors) + } + } +} + +/// Computes row ranges for which every projected column has page data loaded. +fn loaded_row_ranges_for_projection( + selection: Option<&RowSelection>, + projection_mask: &ProjectionMask, + offset_index: Option<&[OffsetIndexMetaData]>, + total_rows: usize, +) -> Option { + let selection = selection?; + let columns = offset_index?; + + columns + .iter() + .enumerate() + .filter_map(|(leaf_idx, column)| { + let pages = column.page_locations(); + (projection_mask.leaf_included(leaf_idx) && !pages.is_empty()).then(|| { + RowSelection::from_consecutive_ranges( + selection + .row_ranges_for_selected_pages(pages, total_rows) + .into_iter(), + total_rows, + ) + }) + }) + .reduce(|loaded, column| loaded.intersection(&column)) + .filter(|loaded| loaded.skipped_row_count() != 0) + .map(LoadedRowRanges::from_selection) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::arrow_reader::{RowSelection, RowSelector}; + use crate::file::page_index::offset_index::PageLocation; + + #[test] + // Verify that the size of RowGroupDecoderState does not grow too large + fn test_structure_size() { + assert_eq!(std::mem::size_of::(), 240); + } + + #[test] + fn test_loaded_row_ranges_intersect_column_page_boundaries() { + let column = |first_rows: &[i64]| OffsetIndexMetaData { + page_locations: first_rows + .iter() + .enumerate() + .map(|(idx, first_row_index)| PageLocation { + offset: (idx * 10) as i64, + compressed_page_size: 10, + first_row_index: *first_row_index, + }) + .collect(), + unencoded_byte_array_data_bytes: None, + }; + let columns = vec![column(&[0, 4, 8]), column(&[0, 6, 10])]; + let selection = RowSelection::from(vec![ + RowSelector::skip(1), + RowSelector::select(1), + RowSelector::skip(9), + RowSelector::select(1), + ]); + + let loaded = loaded_row_ranges_for_projection( + Some(&selection), + &ProjectionMask::all(), + Some(&columns), + 12, + ) + .unwrap(); + + assert_eq!(loaded.ranges(), &[0..4, 10..12]); + } + + #[test] + fn test_auto_keeps_mask_when_page_pruning_skips_pages() { + let columns = vec![OffsetIndexMetaData { + page_locations: [0, 2, 4, 6, 8, 10] + .into_iter() + .enumerate() + .map(|(idx, first_row_index)| PageLocation { + offset: (idx * 10) as i64, + compressed_page_size: 10, + first_row_index, + }) + .collect(), + unencoded_byte_array_data_bytes: None, + }]; + let selection = RowSelection::from(vec![ + RowSelector::select(1), + RowSelector::skip(10), + RowSelector::select(1), + ]); + let plan_builder = ReadPlanBuilder::new(12) + .with_selection(Some(selection)) + .with_row_selection_policy(RowSelectionPolicy::Auto { threshold: 32 }); + + let prepared = prepare_selection_for_page_skipping( + plan_builder, + &ProjectionMask::all(), + Some(&columns), + 12, + ); + + assert_eq!(prepared.row_selection_policy(), &RowSelectionPolicy::Mask); + } + + #[test] + fn test_row_budget_offset_limit_across_row_groups() { + let first = + RowBudget::new(Some(225), Some(20)).apply_to_plan(ReadPlanBuilder::new(1024), 200); + assert_eq!(first.rows_before_budget, 200); + assert_eq!(first.rows_after_budget, 0); + assert_eq!(first.remaining_budget, RowBudget::new(Some(25), Some(20))); + assert_eq!(first.plan_builder.num_rows_selected(), Some(0)); + + let second = first + .remaining_budget + .apply_to_plan(ReadPlanBuilder::new(1024), 200); + assert_eq!(second.rows_before_budget, 200); + assert_eq!(second.rows_after_budget, 20); + assert_eq!(second.remaining_budget, RowBudget::new(Some(0), Some(0))); + assert_eq!(second.plan_builder.num_rows_selected(), Some(20)); + } + + #[test] + fn test_row_budget_limit_only() { + let budgeted = + RowBudget::new(None, Some(20)).apply_to_plan(ReadPlanBuilder::new(1024), 200); + assert_eq!(budgeted.rows_before_budget, 200); + assert_eq!(budgeted.rows_after_budget, 20); + assert_eq!(budgeted.remaining_budget, RowBudget::new(None, Some(0))); + assert_eq!(budgeted.plan_builder.num_rows_selected(), Some(20)); + } + + #[test] + fn test_row_budget_empty_selection() { + let empty_selection = RowSelection::from(vec![RowSelector::skip(200)]); + let budgeted = RowBudget::new(Some(10), Some(20)).apply_to_plan( + ReadPlanBuilder::new(1024).with_selection(Some(empty_selection)), + 200, + ); + assert_eq!(budgeted.rows_before_budget, 0); + assert_eq!(budgeted.rows_after_budget, 0); + assert_eq!( + budgeted.remaining_budget, + RowBudget::new(Some(10), Some(20)) + ); + assert_eq!(budgeted.plan_builder.num_rows_selected(), Some(0)); + } +} diff --git a/vendor/parquet/src/arrow/push_decoder/remaining.rs b/vendor/parquet/src/arrow/push_decoder/remaining.rs new file mode 100644 index 0000000..d265818 --- /dev/null +++ b/vendor/parquet/src/arrow/push_decoder/remaining.rs @@ -0,0 +1,392 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::DecodeResult; +use crate::arrow::arrow_reader::{ParquetRecordBatchReader, RowSelection}; +use crate::arrow::push_decoder::reader_builder::{ + RowBudget, RowGroupBuildResult, RowGroupReaderBuilder, RowGroupReaderBuilderParts, +}; +use crate::errors::ParquetError; +use crate::file::metadata::ParquetMetaData; +use arrow_schema::SchemaRef; +use bytes::Bytes; +use std::collections::VecDeque; +use std::ops::Range; +use std::sync::Arc; + +/// Plan for the next queued row group after row-selection slicing. +#[derive(Debug)] +enum QueuedRowGroupDecision { + /// Hand this row group to the builder. + Read(NextRowGroup), + /// Skip this row group, and keep scanning with the updated budget. + Skip { remaining_budget: RowBudget }, +} + +/// Work item handed from [`RowGroupFrontier`] to [`RowGroupReaderBuilder`]. +#[derive(Debug)] +struct NextRowGroup { + row_group_idx: usize, + row_count: usize, + /// This row group's slice of the global selection, or `None` when all rows + /// are selected. + selection: Option, + /// Budget snapshot to apply while decoding this row group. + budget: RowBudget, +} + +#[derive(Debug, Clone)] +struct RowGroupFrontier { + /// Metadata used to resolve row counts for queued row groups. + parquet_metadata: Arc, + /// Row group indices not yet handed to the builder. + row_groups: VecDeque, + /// Cross-row-group cursor for the optional global row selection. + selection: Option, + /// Offset/limit budget before the next readable row group is planned. + budget: RowBudget, + /// If predicates are present, row groups with selected rows must be read so + /// the predicate can decide whether they are actually needed. + has_predicates: bool, +} + +impl RowGroupFrontier { + fn new( + parquet_metadata: Arc, + row_groups: Vec, + selection: Option, + budget: RowBudget, + has_predicates: bool, + ) -> Self { + Self { + parquet_metadata, + row_groups: VecDeque::from(row_groups), + selection, + budget, + has_predicates, + } + } + + fn row_group_num_rows(&self, row_group_idx: usize) -> Result { + self.parquet_metadata + .row_group(row_group_idx) + .num_rows() + .try_into() + .map_err(|e| ParquetError::General(format!("Row count overflow: {e}"))) + } + + fn update_budget_after_row_group(&mut self, budget: RowBudget) { + self.budget = budget; + } + + /// Peek at the next row-group index [`Self::next_readable_row_group`] + /// would hand out, without mutating any state. Returns `None` if every + /// remaining row group would be skipped under the current + /// selection/budget, or if the queue is empty. + /// + /// Runs the real [`Self::next_readable_row_group`] advance logic on a + /// throwaway clone of the frontier, so peek can never drift from the + /// read path. The clone copies the queued row-group indices and optional + /// row-selection (a `Vec`); see + /// [`RemainingRowGroups::peek_next_row_group`]. + fn peek_next_row_group(&self) -> Result, ParquetError> { + Ok(self + .clone() + .next_readable_row_group()? + .map(|next_row_group| next_row_group.row_group_idx)) + } + + fn clear_remaining(&mut self) { + self.selection = None; + self.row_groups.clear(); + } + + /// Plan whether a selected row group should be read or skipped. + /// + /// Selection-only skips are handled before this method is called. This + /// method applies the remaining offset/limit budget and predicate + /// conservatism. + fn plan_selected_row_group( + &self, + next_row_group: NextRowGroup, + selected_rows: usize, + ) -> QueuedRowGroupDecision { + if self.has_predicates { + return QueuedRowGroupDecision::Read(next_row_group); + } + + let rows_after_budget = self.budget.rows_after(selected_rows); + if rows_after_budget != 0 { + return QueuedRowGroupDecision::Read(next_row_group); + } + + QueuedRowGroupDecision::Skip { + remaining_budget: self.budget.advance(selected_rows, rows_after_budget), + } + } + + /// Advance queued row groups until one should be handed to the builder. + fn next_readable_row_group(&mut self) -> Result, ParquetError> { + loop { + let Some(&row_group_idx) = self.row_groups.front() else { + return Ok(None); + }; + if self.budget.is_exhausted() + || self + .selection + .as_ref() + .is_some_and(|selection| selection.row_count() == 0) + { + self.clear_remaining(); + return Ok(None); + } + + let row_count = self.row_group_num_rows(row_group_idx)?; + let (selection, selected_rows) = match self.selection.as_mut() { + Some(selection) => { + let selection = selection.split_off(row_count); + let selected_rows = selection.row_count(); + if selected_rows == 0 { + self.row_groups.pop_front(); + continue; + } + + let selection = if selected_rows == row_count { + None + } else { + Some(selection) + }; + (selection, selected_rows) + } + None => (None, row_count), + }; + + let next_row_group = NextRowGroup { + row_group_idx, + row_count, + selection, + budget: self.budget, + }; + + match self.plan_selected_row_group(next_row_group, selected_rows) { + QueuedRowGroupDecision::Read(next_row_group) => { + self.row_groups.pop_front(); + return Ok(Some(next_row_group)); + } + QueuedRowGroupDecision::Skip { remaining_budget } => { + self.row_groups.pop_front(); + self.budget = remaining_budget; + } + } + } + } +} + +/// State machine that tracks the remaining high level chunks (row groups) of +/// Parquet data left to read. +/// +/// [`RowGroupFrontier`] owns cross-row-group scan state and selects the next +/// work item. [`RowGroupReaderBuilder`] owns decoding for the active row group. +#[derive(Debug)] +pub(crate) struct RemainingRowGroups { + /// The arrow schema of the decoded output. Carried only so + /// [`Self::into_parts`] can hand it to a rebuilt builder; unused while + /// decoding. + schema: SchemaRef, + + /// Cross-row-group scan state for queued work. + frontier: RowGroupFrontier, + + /// State for building the reader for the current row group + row_group_reader_builder: RowGroupReaderBuilder, +} + +/// The state recovered from a [`RemainingRowGroups`] by +/// [`RemainingRowGroups::into_parts`], describing the row groups *not* yet +/// decoded so a builder reconstructed from it resumes where the decoder left off. +#[derive(Debug)] +pub(crate) struct RemainingRowGroupsParts { + /// The arrow schema of the decoded output. + pub schema: SchemaRef, + /// The Parquet file metadata. + pub metadata: Arc, + /// Row groups not yet handed to the reader builder. + pub row_groups: Vec, + /// The not-yet-consumed slice of the global row selection. + pub selection: Option, + /// Offset still to be skipped before the next readable row group. + pub offset: Option, + /// Output rows still permitted across the remaining row groups. + pub limit: Option, + /// Builder-configurable parts of the inner row-group reader builder. + pub reader_builder: RowGroupReaderBuilderParts, +} + +impl RemainingRowGroups { + pub fn new( + schema: SchemaRef, + parquet_metadata: Arc, + row_groups: Vec, + selection: Option, + budget: RowBudget, + has_predicates: bool, + row_group_reader_builder: RowGroupReaderBuilder, + ) -> Self { + Self { + schema, + frontier: RowGroupFrontier::new( + parquet_metadata, + row_groups, + selection, + budget, + has_predicates, + ), + row_group_reader_builder, + } + } + + /// Decompose into [`RemainingRowGroupsParts`]. + /// + /// Must be called at a row-group boundary (see + /// [`Self::is_at_row_group_boundary`]). The inner reader builder's runtime + /// decode state is discarded; its buffered bytes are carried through. + pub(crate) fn into_parts(self) -> RemainingRowGroupsParts { + let Self { + schema, + frontier, + row_group_reader_builder, + } = self; + // `has_predicates` is recomputed by `build()` from the filter. + let RowGroupFrontier { + parquet_metadata, + row_groups, + selection, + budget, + has_predicates: _, + } = frontier; + RemainingRowGroupsParts { + schema, + metadata: parquet_metadata, + row_groups: Vec::from(row_groups), + selection, + offset: budget.offset(), + limit: budget.limit(), + reader_builder: row_group_reader_builder.into_parts(), + } + } + + /// Push new data buffers that can be used to satisfy pending requests + pub fn push_data(&mut self, ranges: Vec>, buffers: Vec) { + self.row_group_reader_builder.push_data(ranges, buffers); + } + + /// Return the total number of bytes buffered so far + pub fn buffered_bytes(&self) -> u64 { + self.row_group_reader_builder.buffered_bytes() + } + + /// Clear any staged ranges currently buffered for future decode work + pub fn clear_all_ranges(&mut self) { + self.row_group_reader_builder.clear_all_ranges(); + } + + /// True iff the inner row-group reader is between row groups (state + /// `Finished`). Forward to [`RowGroupReaderBuilder::is_finished`]. + pub fn is_at_row_group_boundary(&self) -> bool { + self.row_group_reader_builder.is_finished() + } + + /// Number of row groups remaining (not including the one currently + /// being decoded). + pub fn row_groups_remaining(&self) -> usize { + self.frontier.row_groups.len() + } + + /// Peek at the file-level row-group index that the next call to + /// [`Self::try_next_reader`] will produce a reader for, after + /// simulating the same skip logic [`Self::try_next_reader`] applies + /// internally (row-selection emptiness + offset/limit budget). Does + /// not mutate state. + /// + /// Returns `None` when the active row group is still being decoded, + /// when no row groups remain, or when every remaining row group + /// would be skipped under the current selection/budget. + /// + /// Cost: one clone of the queued row-group indices and optional + /// row-selection per call (the frontier is cloned so the real advance + /// logic can run non-destructively). For callers that peek once per + /// row-group boundary this is O(remaining row groups + selectors) per + /// boundary. + pub fn peek_next_row_group(&self) -> Result, ParquetError> { + if self.row_group_reader_builder.has_active_row_group() { + return Ok(None); + } + self.frontier.peek_next_row_group() + } + + /// returns [`ParquetRecordBatchReader`] suitable for reading the next + /// group of rows from the Parquet data, or the list of data ranges still + /// needed to proceed + pub fn try_next_reader( + &mut self, + ) -> Result, ParquetError> { + loop { + if !self.row_group_reader_builder.has_active_row_group() { + // We are done with the previous row group, seek to the next one + // from the frontier, if any. + + match self.frontier.next_readable_row_group()? { + Some(NextRowGroup { + row_group_idx, + row_count, + selection, + budget, + }) => { + self.row_group_reader_builder.next_row_group( + row_group_idx, + row_count, + selection, + budget, + )?; + } + None => return Ok(DecodeResult::Finished), + } + } + + match self.row_group_reader_builder.try_build()? { + RowGroupBuildResult::Finished { remaining_budget } => { + self.frontier + .update_budget_after_row_group(remaining_budget); + // reader is done, proceed to the next row group + } + RowGroupBuildResult::NeedsData(ranges) => { + // need more data to proceed + return Ok(DecodeResult::NeedsData(ranges)); + } + RowGroupBuildResult::Data { + batch_reader, + remaining_budget, + } => { + self.frontier + .update_budget_after_row_group(remaining_budget); + // ready to read the row group + return Ok(DecodeResult::Data(batch_reader)); + } + } + } + } +} diff --git a/vendor/parquet/src/arrow/record_reader/buffer.rs b/vendor/parquet/src/arrow/record_reader/buffer.rs new file mode 100644 index 0000000..7060586 --- /dev/null +++ b/vendor/parquet/src/arrow/record_reader/buffer.rs @@ -0,0 +1,85 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::arrow::buffer::bit_util::iter_set_bits_rev; +use crate::errors::Result; + +/// A buffer that supports padding with nulls +pub trait ValuesBuffer { + /// Create a new buffer with capacity for at least `capacity` logical values + fn with_capacity(capacity: usize) -> Self; + + /// Reserve capacity for `additional` more logical values. + /// + /// Callers should use this once the level decoders have determined how + /// many output slots are actually needed for the next read. + fn reserve_exact(&mut self, additional: usize); + + /// If a column contains nulls, more level data may be read than value data, as null + /// values are not encoded. Therefore, first the levels data is read, the null count + /// determined, and then the corresponding number of values read to a [`ValuesBuffer`]. + /// + /// It is then necessary to move this values data into positions that correspond to + /// the non-null level positions. This is what this method does. + /// + /// It is provided with: + /// + /// - `read_offset` - the offset in [`ValuesBuffer`] to start null padding from + /// - `values_read` - the number of values read + /// - `levels_read` - the number of levels read + /// - `valid_mask` - a packed mask of valid levels + /// + /// Returns an error if the inputs are inconsistent, for example because the + /// decoded data was corrupt. This must not panic on such input. + fn pad_nulls( + &mut self, + read_offset: usize, + values_read: usize, + levels_read: usize, + valid_mask: &[u8], + ) -> Result<()>; +} + +impl ValuesBuffer for Vec { + fn with_capacity(capacity: usize) -> Self { + Vec::with_capacity(capacity) + } + + fn reserve_exact(&mut self, additional: usize) { + Vec::reserve_exact(self, additional) + } + + fn pad_nulls( + &mut self, + read_offset: usize, + values_read: usize, + levels_read: usize, + valid_mask: &[u8], + ) -> Result<()> { + self.resize(read_offset + levels_read, T::default()); + + let values_range = read_offset..read_offset + values_read; + for (value_pos, level_pos) in values_range.rev().zip(iter_set_bits_rev(valid_mask)) { + debug_assert!(level_pos >= value_pos); + if level_pos <= value_pos { + break; + } + self[level_pos] = self[value_pos]; + } + Ok(()) + } +} diff --git a/vendor/parquet/src/arrow/record_reader/definition_levels.rs b/vendor/parquet/src/arrow/record_reader/definition_levels.rs new file mode 100644 index 0000000..a104c93 --- /dev/null +++ b/vendor/parquet/src/arrow/record_reader/definition_levels.rs @@ -0,0 +1,689 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_array::builder::BooleanBufferBuilder; +use arrow_buffer::Buffer; +use arrow_buffer::bit_chunk_iterator::UnalignedBitChunk; +use bytes::Bytes; + +use crate::arrow::buffer::bit_util::count_set_bits; +use crate::basic::Encoding; +use crate::column::reader::decoder::{ + ColumnLevelDecoder, DefinitionLevelDecoder, DefinitionLevelDecoderImpl, +}; +use crate::errors::{ParquetError, Result}; +use crate::schema::types::ColumnDescPtr; + +enum BufferInner { + /// Compute levels and null mask + Full { + levels: Vec, + nulls: BooleanBufferBuilder, + max_level: i16, + }, + /// Only compute null bitmask - requires max level to be 1 + /// + /// This is an optimisation for the common case of a nullable scalar column, as decoding + /// the definition level data is only required when decoding nested structures + /// + Mask { nulls: BooleanBufferBuilder }, +} + +pub struct DefinitionLevelBuffer { + inner: BufferInner, + + /// The length of this buffer + /// + /// Note: `buffer` and `builder` may contain more elements + len: usize, +} + +impl DefinitionLevelBuffer { + pub fn new(desc: &ColumnDescPtr, null_mask_only: bool) -> Self { + let inner = match null_mask_only { + true => { + assert_eq!( + desc.max_def_level(), + 1, + "max definition level must be 1 to only compute null bitmask" + ); + + assert_eq!( + desc.max_rep_level(), + 0, + "max repetition level must be 0 to only compute null bitmask" + ); + + BufferInner::Mask { + nulls: BooleanBufferBuilder::new(0), + } + } + false => BufferInner::Full { + levels: Vec::new(), + nulls: BooleanBufferBuilder::new(0), + max_level: desc.max_def_level(), + }, + }; + + Self { inner, len: 0 } + } + + /// Returns the built level data + pub fn consume_levels(&mut self) -> Option> { + match &mut self.inner { + BufferInner::Full { levels, .. } => Some(std::mem::take(levels)), + BufferInner::Mask { .. } => None, + } + } + + /// Returns the built null bitmask, or None if all values are valid + pub fn consume_bitmask(&mut self) -> Option { + self.len = 0; + let nulls = match &mut self.inner { + BufferInner::Full { nulls, .. } => nulls, + BufferInner::Mask { nulls } => nulls, + }; + + // Always call finish() to reset the builder state for the next batch. + let buffer = nulls.finish().into_inner(); + + // If no bitmap was constructed, return None + if buffer.is_empty() { + return None; + } + + Some(buffer) + } + + pub fn nulls(&self) -> &BooleanBufferBuilder { + match &self.inner { + BufferInner::Full { nulls, .. } => nulls, + BufferInner::Mask { nulls } => nulls, + } + } + + /// Returns the raw definition levels accumulated so far, if available. + /// Only available when the buffer is in Full mode (nested columns). + pub fn levels(&self) -> Option<&[i16]> { + match &self.inner { + BufferInner::Full { levels, .. } => Some(levels.as_slice()), + BufferInner::Mask { .. } => None, + } + } +} + +/// Build a filtered validity bitmap from definition/repetition levels. +/// +/// For each level entry where `d >= include_threshold` (when set) and +/// `r <= max_rep` (when `rep_filter` is provided), appends one bit to +/// `bitmap`: set when `d >= value_level`, unset otherwise. +/// +/// Returns the number of bits appended. +/// +/// This is the shared implementation behind the compact bitmap (leaf +/// readers with selective padding) and the struct validity bitmap. +/// Processing uses 64-level chunks with [`compress`] for word-at-a-time +/// packing. +pub(crate) fn build_filtered_validity_bitmap( + def_levels: &[i16], + rep_filter: Option<(&[i16], i16)>, + include_threshold: Option, + value_level: i16, + bitmap: &mut BooleanBufferBuilder, +) -> usize { + // Fast path: no filtering — every def level produces a bit. + if include_threshold.is_none() && rep_filter.is_none() { + let chunks = def_levels.chunks_exact(u64::BITS as usize); + let remainder = chunks.remainder(); + for chunk in chunks { + let mut word: u64 = 0; + for (i, &d) in chunk.iter().enumerate() { + word |= ((d >= value_level) as u64) << i; + } + bitmap.append_word(word, u64::BITS as usize); + } + for &d in remainder { + bitmap.append(d >= value_level); + } + return def_levels.len(); + } + + // Filtered path: build include mask + value mask per chunk, compress. + let mut item_count: usize = 0; + let chunks = def_levels.chunks_exact(u64::BITS as usize); + let remainder_offset = def_levels.len() - chunks.remainder().len(); + + for (chunk_idx, chunk) in chunks.enumerate() { + let base = chunk_idx * u64::BITS as usize; + let mut include_mask: u64 = 0; + let mut value_mask: u64 = 0; + for (i, &d) in chunk.iter().enumerate() { + let mut include = true; + if let Some(threshold) = include_threshold { + if d < threshold { + include = false; + } + } + if include { + if let Some((reps, max_rep)) = rep_filter { + if reps[base + i] > max_rep { + include = false; + } + } + } + include_mask |= (include as u64) << i; + value_mask |= ((d >= value_level) as u64) << i; + } + + let count = include_mask.count_ones() as usize; + let compact = compress(value_mask, include_mask); + bitmap.append_word(compact, count); + item_count += count; + } + + for idx in remainder_offset..def_levels.len() { + let d = def_levels[idx]; + if let Some(threshold) = include_threshold { + if d < threshold { + continue; + } + } + if let Some((reps, max_rep)) = rep_filter { + if reps[idx] > max_rep { + continue; + } + } + bitmap.append(d >= value_level); + item_count += 1; + } + + item_count +} + +use crate::util::bit_util::compress; + +enum MaybePacked { + Packed(PackedDecoder), + Fallback(DefinitionLevelDecoderImpl), +} + +pub struct DefinitionLevelBufferDecoder { + max_level: i16, + decoder: MaybePacked, +} + +impl DefinitionLevelBufferDecoder { + pub fn new(max_level: i16, packed: bool) -> Self { + let decoder = match packed { + true => MaybePacked::Packed(PackedDecoder::new()), + false => MaybePacked::Fallback(DefinitionLevelDecoderImpl::new(max_level)), + }; + + Self { max_level, decoder } + } +} + +impl ColumnLevelDecoder for DefinitionLevelBufferDecoder { + type Buffer = DefinitionLevelBuffer; + + fn set_data(&mut self, encoding: Encoding, data: Bytes) -> Result<()> { + match &mut self.decoder { + MaybePacked::Packed(d) => d.set_data(encoding, data), + MaybePacked::Fallback(d) => d.set_data(encoding, data)?, + }; + Ok(()) + } +} + +impl DefinitionLevelDecoder for DefinitionLevelBufferDecoder { + fn read_def_levels( + &mut self, + writer: &mut Self::Buffer, + num_levels: usize, + ) -> Result<(usize, usize)> { + match (&mut writer.inner, &mut self.decoder) { + ( + BufferInner::Full { + levels, + nulls, + max_level, + }, + MaybePacked::Fallback(decoder), + ) => { + assert_eq!(self.max_level, *max_level); + + let start = levels.len(); + let (values_read, levels_read) = decoder.read_def_levels(levels, num_levels)?; + + // Safety: slice iterator has a trusted length + unsafe { + nulls + .extend_trusted_len(levels[start..].iter().map(|level| level == max_level)); + } + + Ok((values_read, levels_read)) + } + (BufferInner::Mask { nulls }, MaybePacked::Packed(decoder)) => { + assert_eq!(self.max_level, 1); + + // Fast path: if all requested levels are valid (max definition level), + // we can skip RLE decoding and just append all-ones to the bitmap. + // This is faster than decoding RLE data. + if let Some(count) = decoder.try_consume_all_valid(num_levels)? { + nulls.append_n(count, true); + return Ok((count, count)); // values_read == levels_read when all valid + } + + // Normal path: decode RLE data into the bitmap + let start = nulls.len(); + let levels_read = decoder.read(nulls, num_levels)?; + + let values_read = count_set_bits(nulls.as_slice(), start..start + levels_read); + Ok((values_read, levels_read)) + } + _ => unreachable!("inconsistent null mask"), + } + } + + fn skip_def_levels(&mut self, num_levels: usize) -> Result<(usize, usize)> { + match &mut self.decoder { + MaybePacked::Fallback(decoder) => decoder.skip_def_levels(num_levels), + MaybePacked::Packed(decoder) => decoder.skip(num_levels), + } + } +} + +/// An optimized decoder for decoding [RLE] and [BIT_PACKED] data with a bit width of 1 +/// directly into a bitmask +/// +/// This is significantly faster than decoding the data into `[i16]` and then computing +/// a bitmask from this, as not only can it skip this buffer allocation and construction, +/// but it can exploit properties of the encoded data to reduce work further +/// +/// In particular: +/// +/// * Packed runs are already bitmask encoded and can simply be appended +/// * Runs of 1 or 0 bits can be efficiently appended with byte (or larger) operations +/// +/// [RLE]: https://github.com/apache/parquet-format/blob/master/Encodings.md#run-length-encoding--bit-packing-hybrid-rle--3 +/// [BIT_PACKED]: https://github.com/apache/parquet-format/blob/master/Encodings.md#bit-packed-deprecated-bit_packed--4 +struct PackedDecoder { + data: Bytes, + data_offset: usize, + rle_left: usize, + rle_value: bool, + packed_count: usize, + packed_offset: usize, +} + +impl PackedDecoder { + fn next_rle_block(&mut self) -> Result<()> { + let indicator_value = self.decode_header()?; + if indicator_value & 1 == 1 { + let len = (indicator_value >> 1) as usize; + self.packed_count = len * 8; + self.packed_offset = 0; + } else { + self.rle_left = (indicator_value >> 1) as usize; + let byte = *self.data.as_ref().get(self.data_offset).ok_or_else(|| { + ParquetError::EOF( + "unexpected end of file whilst decoding definition levels rle value".into(), + ) + })?; + + self.data_offset += 1; + self.rle_value = byte != 0; + } + Ok(()) + } + + /// Decodes a VLQ encoded little endian integer and returns it + fn decode_header(&mut self) -> Result { + let mut offset = 0; + let mut v: i64 = 0; + while offset < 10 { + let byte = *self + .data + .as_ref() + .get(self.data_offset + offset) + .ok_or_else(|| { + ParquetError::EOF( + "unexpected end of file whilst decoding definition levels rle header" + .into(), + ) + })?; + + v |= ((byte & 0x7F) as i64) << (offset * 7); + offset += 1; + if byte & 0x80 == 0 { + self.data_offset += offset; + return Ok(v); + } + } + Err(general_err!("too many bytes for VLQ")) + } +} + +impl PackedDecoder { + fn new() -> Self { + Self { + data: Bytes::from(vec![]), + data_offset: 0, + rle_left: 0, + rle_value: false, + packed_count: 0, + packed_offset: 0, + } + } + + fn set_data(&mut self, encoding: Encoding, data: Bytes) { + self.rle_left = 0; + self.rle_value = false; + self.packed_offset = 0; + self.packed_count = match encoding { + Encoding::RLE => 0, + #[allow(deprecated)] + Encoding::BIT_PACKED => data.len() * 8, + _ => unreachable!("invalid level encoding: {}", encoding), + }; + self.data = data; + self.data_offset = 0; + } + + /// Try to consume `len` levels if all are valid (max definition level). + /// + /// Returns `Ok(Some(count))` if successfully consumed `count` all-valid levels. + /// Returns `Ok(None)` if there are any nulls or packed data that prevents fast path. + /// + /// Note: On `None`, the decoder state may have advanced to the next RLE block, + /// but only if `rle_left` was zero (i.e., the block would have been loaded + /// on the next read anyway). + fn try_consume_all_valid(&mut self, len: usize) -> Result> { + // If no active run and no packed data pending, try to parse the next RLE block + if self.rle_left == 0 && self.packed_count == self.packed_offset { + if self.data_offset < self.data.len() { + self.next_rle_block()?; + } else { + // No more data available + return Ok(None); + } + } + + // Fast path only works when we have an active RLE run of true values + // that covers the entire requested length. + if self.rle_left >= len && self.rle_value { + self.rle_left -= len; + return Ok(Some(len)); + } + + // Any other case (null run, packed data, or insufficient length) + // falls back to normal path + Ok(None) + } + + fn read(&mut self, buffer: &mut BooleanBufferBuilder, len: usize) -> Result { + let mut read = 0; + while read != len { + if self.rle_left != 0 { + let to_read = self.rle_left.min(len - read); + buffer.append_n(to_read, self.rle_value); + self.rle_left -= to_read; + read += to_read; + } else if self.packed_count != self.packed_offset { + let to_read = (self.packed_count - self.packed_offset).min(len - read); + let offset = self.data_offset * 8 + self.packed_offset; + buffer.append_packed_range(offset..offset + to_read, self.data.as_ref()); + self.packed_offset += to_read; + read += to_read; + + if self.packed_offset == self.packed_count { + self.data_offset += self.packed_count / 8; + } + } else if self.data_offset == self.data.len() { + break; + } else { + self.next_rle_block()? + } + } + Ok(read) + } + + /// Skips `level_num` definition levels + /// + /// Returns the number of values skipped and the number of levels skipped + fn skip(&mut self, level_num: usize) -> Result<(usize, usize)> { + let mut skipped_value = 0; + let mut skipped_level = 0; + while skipped_level != level_num { + if self.rle_left != 0 { + let to_skip = self.rle_left.min(level_num - skipped_level); + self.rle_left -= to_skip; + skipped_level += to_skip; + if self.rle_value { + skipped_value += to_skip; + } + } else if self.packed_count != self.packed_offset { + let to_skip = + (self.packed_count - self.packed_offset).min(level_num - skipped_level); + let offset = self.data_offset * 8 + self.packed_offset; + let bit_chunk = UnalignedBitChunk::new(self.data.as_ref(), offset, to_skip); + skipped_value += bit_chunk.count_ones(); + self.packed_offset += to_skip; + skipped_level += to_skip; + if self.packed_offset == self.packed_count { + self.data_offset += self.packed_count / 8; + } + } else if self.data_offset == self.data.len() { + break; + } else { + self.next_rle_block()? + } + } + Ok((skipped_value, skipped_level)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use crate::encodings::rle::RleEncoder; + use rand::{Rng, rng}; + + #[test] + fn test_build_validity_bitmap_unfiltered_word_chunk() { + // 65 levels forces the unfiltered path to process one full u64 word + // with append_word, plus a remainder bit. + let def_levels = (0..65) + .map(|i| if i % 3 == 0 { 2 } else { 1 }) + .collect::>(); + let mut bitmap = BooleanBufferBuilder::new(0); + + assert_eq!( + build_filtered_validity_bitmap(&def_levels, None, None, 2, &mut bitmap), + def_levels.len() + ); + + let bitmap = bitmap.finish(); + for (idx, def) in def_levels.iter().enumerate() { + assert_eq!(bitmap.value(idx), *def >= 2); + } + } + + #[test] + fn test_packed_decoder() { + let mut rng = rng(); + let len: usize = rng.random_range(512..1024); + + let mut expected = BooleanBufferBuilder::new(len); + let mut encoder = RleEncoder::new(1, 1024); + for _ in 0..len { + let bool = rng.random_bool(0.8); + encoder.put(bool as u64); + expected.append(bool); + } + assert_eq!(expected.len(), len); + + let encoded = encoder.consume(); + let mut decoder = PackedDecoder::new(); + decoder.set_data(Encoding::RLE, encoded.into()); + + // Decode data in random length intervals + let mut decoded = BooleanBufferBuilder::new(len); + loop { + let remaining = len - decoded.len(); + if remaining == 0 { + break; + } + + let to_read = rng.random_range(1..=remaining); + decoder.read(&mut decoded, to_read).unwrap(); + } + + assert_eq!(decoded.len(), len); + assert_eq!(decoded.as_slice(), expected.as_slice()); + } + + #[test] + fn test_packed_decoder_skip() { + let mut rng = rng(); + let len: usize = rng.random_range(512..1024); + + let mut expected = BooleanBufferBuilder::new(len); + let mut encoder = RleEncoder::new(1, 1024); + + let mut total_value = 0; + for _ in 0..len { + let bool = rng.random_bool(0.8); + encoder.put(bool as u64); + expected.append(bool); + if bool { + total_value += 1; + } + } + assert_eq!(expected.len(), len); + + let encoded = encoder.consume(); + let mut decoder = PackedDecoder::new(); + decoder.set_data(Encoding::RLE, encoded.into()); + + let mut skip_value = 0; + let mut read_value = 0; + let mut skip_level = 0; + let mut read_level = 0; + + loop { + let offset = skip_level + read_level; + let remaining_levels = len - offset; + if remaining_levels == 0 { + break; + } + let to_read_or_skip_level = rng.random_range(1..=remaining_levels); + if rng.random_bool(0.5) { + let (skip_val_num, skip_level_num) = decoder.skip(to_read_or_skip_level).unwrap(); + skip_value += skip_val_num; + skip_level += skip_level_num + } else { + let mut decoded = BooleanBufferBuilder::new(to_read_or_skip_level); + let read_level_num = decoder.read(&mut decoded, to_read_or_skip_level).unwrap(); + read_level += read_level_num; + for i in 0..read_level_num { + assert!(!decoded.is_empty()); + //check each read bit + let read_bit = decoded.get_bit(i); + if read_bit { + read_value += 1; + } + let expect_bit = expected.get_bit(i + offset); + assert_eq!(read_bit, expect_bit); + } + } + } + assert_eq!(read_level + skip_level, len); + assert_eq!(read_value + skip_value, total_value); + } + + #[test] + fn test_try_consume_all_valid() { + // Test with all-valid data (all 1s) - single RLE run + let len = 100; + let mut encoder = RleEncoder::new(1, 1024); + for _ in 0..len { + encoder.put(1); // all valid + } + let encoded = encoder.consume(); + let mut decoder = PackedDecoder::new(); + decoder.set_data(Encoding::RLE, encoded.into()); + + // try_consume_all_valid now parses the RLE block itself, no need to read first + let result = decoder.try_consume_all_valid(len).unwrap(); + assert_eq!(result, Some(len)); + + // Test with all-null data (all 0s) + let mut encoder = RleEncoder::new(1, 1024); + for _ in 0..len { + encoder.put(0); // all null + } + let encoded = encoder.consume(); + let mut decoder = PackedDecoder::new(); + decoder.set_data(Encoding::RLE, encoded.into()); + + // Should return None because rle_value is false (all nulls) + let result = decoder.try_consume_all_valid(len).unwrap(); + assert_eq!(result, None); + + // Test when requesting more than available in current RLE run + let mut encoder = RleEncoder::new(1, 1024); + for _ in 0..10 { + encoder.put(1); // small run of valid + } + for _ in 0..10 { + encoder.put(0); // followed by nulls + } + let encoded = encoder.consume(); + let mut decoder = PackedDecoder::new(); + decoder.set_data(Encoding::RLE, encoded.into()); + + // Request more than the valid run - should return None + // (because we don't look ahead to next block) + let result = decoder.try_consume_all_valid(20).unwrap(); + assert_eq!(result, None); + + // Reset decoder and try requesting within the run + decoder.set_data(Encoding::RLE, { + let mut encoder = RleEncoder::new(1, 1024); + for _ in 0..10 { + encoder.put(1); + } + for _ in 0..10 { + encoder.put(0); + } + encoder.consume().into() + }); + + let result = decoder.try_consume_all_valid(5).unwrap(); + assert_eq!(result, Some(5)); + + // After skipping 5, we should have 5 left in the valid RLE run + let result = decoder.try_consume_all_valid(5).unwrap(); + assert_eq!(result, Some(5)); + + // Now the valid run is exhausted, next call should parse the null run and return None + let result = decoder.try_consume_all_valid(5).unwrap(); + assert_eq!(result, None); + } +} diff --git a/vendor/parquet/src/arrow/record_reader/mod.rs b/vendor/parquet/src/arrow/record_reader/mod.rs new file mode 100644 index 0000000..cdb8084 --- /dev/null +++ b/vendor/parquet/src/arrow/record_reader/mod.rs @@ -0,0 +1,1083 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_buffer::{BooleanBufferBuilder, Buffer}; + +use crate::arrow::record_reader::{ + buffer::ValuesBuffer, + definition_levels::{DefinitionLevelBuffer, DefinitionLevelBufferDecoder}, +}; +use crate::column::reader::decoder::RepetitionLevelDecoderImpl; +use crate::column::{ + page::PageReader, + reader::{ + GenericColumnReader, + decoder::{ColumnValueDecoder, ColumnValueDecoderImpl}, + }, +}; +use crate::data_type::DataType; +use crate::errors::{ParquetError, Result}; +use crate::schema::types::ColumnDescPtr; + +pub(crate) mod buffer; +pub(crate) mod definition_levels; + +/// A `RecordReader` is a stateful column reader that delimits semantic records. +pub type RecordReader = GenericRecordReader::T>, ColumnValueDecoderImpl>; + +pub(crate) type ColumnReader = + GenericColumnReader; + +/// A generic stateful column reader that delimits semantic records +/// +/// This type is hidden from the docs, and relies on private traits with no +/// public implementations. As such this type signature may be changed without +/// breaking downstream users as it can only be constructed through type aliases +pub struct GenericRecordReader { + column_desc: ColumnDescPtr, + + /// Values buffer, lazily initialized on first read to avoid + /// allocating a buffer that may never be used (e.g., after the last batch) + values: Option, + def_levels: Option, + rep_levels: Option>, + column_reader: Option>, + /// Number of buffered levels / null-padded values + num_values: usize, + /// Number of buffered records + num_records: usize, + /// Capacity hint for pre-allocating buffers based on batch size + capacity_hint: usize, + /// Number of values in the values buffer (may differ from num_values when + /// padding_threshold is set, since parent-level padding is excluded). + values_written: usize, + /// Definition-level threshold used for selective null padding. + /// + /// With full padding (`None`), the leaf values buffer has one slot for each + /// decoded definition level. This includes placeholders for null or empty + /// parent lists, which parent `ListArrayReader`s later have to filter out + /// before computing offsets. + /// + /// With selective padding (`Some(threshold)`), the threshold is the nearest + /// enclosing list/map definition level. Entries with `def < threshold` + /// describe a null/empty parent and are skipped entirely. Entries with + /// `def >= threshold` belong to an actual child item slot: real values are + /// copied, and item-level nulls are padded. The companion `compact_bitmap` + /// has the same compact length and becomes the leaf null bitmap. + padding_threshold: Option, + /// Compact bitmap accumulated during selective padding. Each bit + /// corresponds to an item-level entry (def >= threshold): set when the + /// value is real (def >= max_def), unset for item-level nulls. Used both + /// as the valid_mask for `pad_nulls` (via `as_slice()`) and as the null + /// bitmap consumed by the leaf reader (via `consume_compact_bitmap`). + compact_bitmap: Option, +} + +impl GenericRecordReader +where + V: ValuesBuffer, + CV: ColumnValueDecoder, +{ + /// Create a new [`GenericRecordReader`] + /// + /// The capacity is used to pre-allocate internal buffers for full-padding + /// reads, avoiding reallocations when reading fragmented row selections. + pub fn new(desc: ColumnDescPtr, capacity: usize) -> Self { + let def_levels = (desc.max_def_level() > 0) + .then(|| DefinitionLevelBuffer::new(&desc, packed_null_mask(&desc))); + + let rep_levels = (desc.max_rep_level() > 0).then(Vec::new); + + Self { + values: None, // Lazily initialized on first read + def_levels, + rep_levels, + column_reader: None, + column_desc: desc, + num_values: 0, + num_records: 0, + capacity_hint: capacity, + values_written: 0, + padding_threshold: None, + compact_bitmap: None, + } + } + + /// Set the current page reader. + pub fn set_page_reader(&mut self, page_reader: Box) -> Result<()> { + let descr = &self.column_desc; + let values_decoder = CV::new(descr); + + let def_level_decoder = (descr.max_def_level() != 0).then(|| { + DefinitionLevelBufferDecoder::new(descr.max_def_level(), packed_null_mask(descr)) + }); + + let rep_level_decoder = (descr.max_rep_level() != 0) + .then(|| RepetitionLevelDecoderImpl::new(descr.max_rep_level())); + + self.column_reader = Some(GenericColumnReader::new_with_decoders( + self.column_desc.clone(), + page_reader, + values_decoder, + def_level_decoder, + rep_level_decoder, + )); + Ok(()) + } + + /// Try to read `num_records` of column data into internal buffer. + /// + /// # Returns + /// + /// Number of actual records read. + pub fn read_records(&mut self, num_records: usize) -> Result { + if self.column_reader.is_none() { + return Ok(0); + } + + let mut records_read = 0; + + loop { + let records_to_read = num_records - records_read; + records_read += self.read_one_batch(records_to_read)?; + if records_read == num_records || !self.column_reader.as_mut().unwrap().has_next()? { + break; + } + } + Ok(records_read) + } + + /// Try to skip the next `num_records` rows + /// + /// # Returns + /// + /// Number of records skipped + pub fn skip_records(&mut self, num_records: usize) -> Result { + match self.column_reader.as_mut() { + Some(reader) => reader.skip_records(num_records), + None => Ok(0), + } + } + + /// Returns number of records stored in buffer. + #[allow(unused)] + pub fn num_records(&self) -> usize { + self.num_records + } + + /// Return number of values stored in buffer. + /// If the parquet column is not repeated, it should be equals to `num_records`, + /// otherwise it should be larger than or equal to `num_records`. + pub fn num_values(&self) -> usize { + self.num_values + } + + /// Returns definition level data. + /// The implementation has side effects. It will create a new buffer to hold those + /// definition level values that have already been read into memory but not counted + /// as record values, e.g. those from `self.num_values` to `self.values_written`. + pub fn consume_def_levels(&mut self) -> Option> { + self.def_levels.as_mut().and_then(|x| x.consume_levels()) + } + + /// Return repetition level data. + /// The side effect is similar to `consume_def_levels`. + pub fn consume_rep_levels(&mut self) -> Option> { + self.rep_levels.as_mut().map(std::mem::take) + } + + /// Returns currently stored buffer data. + /// The side effect is similar to `consume_def_levels`. + pub fn consume_record_data(&mut self) -> V { + // Take the buffer, leaving None. The next read will lazily allocate a new buffer. + // This avoids allocating a buffer that may never be used (e.g., after the last batch). + self.values.take().unwrap_or_else(|| V::with_capacity(0)) + } + + /// Reset state of record reader. + /// Should be called after consuming data, e.g. `consume_rep_levels`, + /// `consume_rep_levels`, `consume_record_data` and `consume_compact_bitmap`. + pub fn reset(&mut self) { + self.num_values = 0; + self.num_records = 0; + self.values_written = 0; + self.compact_bitmap = None; + } + + /// Returns the maximum definition level for the column being read. + pub fn max_def_level(&self) -> i16 { + self.column_desc.max_def_level() + } + + /// Set the padding threshold. When set, `pad_nulls` only pads entries + /// where `def >= threshold` (item-level nulls within non-null lists), + /// skipping list-level padding entries (def < threshold). + pub fn set_padding_threshold(&mut self, threshold: i16) { + self.padding_threshold = Some(threshold); + } + + /// Returns the number of values in the values buffer. + /// When padding_threshold is None, this equals `num_values` (full padding). + /// When padding_threshold is set, this is the item_count (selective padding). + pub fn values_written(&self) -> usize { + if self.padding_threshold.is_some() { + self.values_written + } else { + self.num_values + } + } + + /// Consume the compact null bitmap built during selective padding. + /// Returns the full bitmap when not using selective padding. + pub fn consume_compact_bitmap(&mut self) -> Option { + if self.padding_threshold.is_some() { + if let Some(levels) = self.def_levels.as_mut() { + levels.consume_bitmask(); + } + self.compact_bitmap + .as_mut() + .map(|b| b.finish().into_inner()) + } else { + self.consume_bitmap() + } + } + + /// Returns bitmap data for nullable columns. + /// For non-nullable columns, the bitmap is discarded. + pub fn consume_bitmap(&mut self) -> Option { + let mask = self + .def_levels + .as_mut() + .and_then(|levels| levels.consume_bitmask()); + + // While we always consume the bitmask here, we only want to return + // the bitmask for nullable arrays. (Marking nulls on a non-nullable + // array may fail validations, even if those nulls are masked off at + // a higher level.) + if self.column_desc.self_type().is_optional() { + mask + } else { + None + } + } + + /// Try to read one batch of data returning the number of records read + fn read_one_batch(&mut self, batch_size: usize) -> Result { + if batch_size == 0 { + return Ok(0); + } + // Update capacity hint to the largest batch size seen. + if batch_size > self.capacity_hint { + self.capacity_hint = batch_size; + } + + let padding_threshold = self.padding_threshold; + let max_def = self.column_desc.max_def_level(); + let mut physical_values_to_read = 0usize; + let mut output_slots_to_read = 0usize; + + let values = self.values.get_or_insert_with(|| V::with_capacity(0)); + let compact_bitmap = &mut self.compact_bitmap; + + if padding_threshold.is_none() { + values.reserve_exact(self.capacity_hint.saturating_sub(self.num_values)); + } + + let (records_read, values_read, levels_read) = self + .column_reader + .as_mut() + .unwrap() + .read_records_with_reservation( + batch_size, + self.def_levels.as_mut(), + self.rep_levels.as_mut(), + values, + |values, values_to_read, levels_to_read, def_levels| { + let output_slots = if let Some(threshold) = padding_threshold { + let def_levels = def_levels.ok_or_else(|| { + general_err!( + "Definition levels should exist when data is less than levels!" + ) + })?; + let all_levels = def_levels.levels().ok_or_else(|| { + general_err!( + "Raw definition levels must be available for selective padding" + ) + })?; + let batch_levels = &all_levels[all_levels.len() - levels_to_read..]; + let bitmap = + compact_bitmap.get_or_insert_with(|| BooleanBufferBuilder::new(0)); + + definition_levels::build_filtered_validity_bitmap( + batch_levels, + None, + Some(threshold), + max_def, + bitmap, + ) + } else { + levels_to_read + }; + + let additional = output_slots_to_read + .saturating_add(output_slots) + .saturating_sub(physical_values_to_read); + values.reserve_exact(additional); + + output_slots_to_read += output_slots; + physical_values_to_read += values_to_read; + Ok(()) + }, + )?; + + if self.padding_threshold.is_some() { + debug_assert_eq!( + physical_values_to_read, values_read, + "reservation accounting must match decoded values" + ); + let item_count = output_slots_to_read; + let bitmap = compact_bitmap.get_or_insert_with(|| BooleanBufferBuilder::new(0)); + + // Pad values to item_count positions (only if there are gaps). + if values_read < item_count { + values.reserve_exact(item_count - values_read); + values.pad_nulls( + self.values_written, + values_read, + item_count, + bitmap.as_slice(), + )?; + } + + debug_assert_eq!( + bitmap.len(), + self.values_written + item_count, + "compact bitmap length must equal total items written" + ); + self.values_written += item_count; + } else if values_read < levels_read { + debug_assert_eq!( + output_slots_to_read, levels_read, + "reservation accounting must match decoded levels" + ); + // Full padding: pad all null positions. + let def_levels = self.def_levels.as_ref().ok_or_else(|| { + general_err!("Definition levels should exist when data is less than levels!") + })?; + + values.reserve_exact(levels_read - values_read); + values.pad_nulls( + self.num_values, + values_read, + levels_read, + def_levels.nulls().as_slice(), + )?; + } + + self.num_records += records_read; + self.num_values += levels_read; + Ok(records_read) + } +} + +/// Returns true if we do not need to unpack the nullability for this column, this is +/// only possible if the max definition level is 1, and corresponds to nulls at the +/// leaf level, as opposed to a nullable parent nested type +fn packed_null_mask(descr: &ColumnDescPtr) -> bool { + descr.max_def_level() == 1 && descr.max_rep_level() == 0 && descr.self_type().is_optional() +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::buffer::Buffer; + + use crate::arrow::arrow_reader::DEFAULT_BATCH_SIZE; + use crate::basic::Encoding; + use crate::data_type::Int32Type; + use crate::schema::parser::parse_message_type; + use crate::schema::types::SchemaDescriptor; + use crate::util::test_common::page_util::{ + DataPageBuilder, DataPageBuilderImpl, InMemoryPageReader, + }; + + use super::RecordReader; + + #[test] + fn test_read_required_records() { + // Construct column schema + let message_type = " + message test_schema { + REQUIRED INT32 leaf; + } + "; + let desc = parse_message_type(message_type) + .map(|t| SchemaDescriptor::new(Arc::new(t))) + .map(|s| s.column(0)) + .unwrap(); + + // Construct record reader + let mut record_reader = RecordReader::::new(desc.clone(), DEFAULT_BATCH_SIZE); + + // First page + + // Records data: + // test_schema + // leaf: 4 + // test_schema + // leaf: 7 + // test_schema + // leaf: 6 + // test_schema + // left: 3 + // test_schema + // left: 2 + { + let values = [4, 7, 6, 3, 2]; + let mut pb = DataPageBuilderImpl::new(desc.clone(), 5, true); + pb.add_values::(Encoding::PLAIN, &values); + let page = pb.consume(); + + let page_reader = Box::new(InMemoryPageReader::new(vec![page])); + record_reader.set_page_reader(page_reader).unwrap(); + assert_eq!(2, record_reader.read_records(2).unwrap()); + assert_eq!(2, record_reader.num_records()); + assert_eq!(2, record_reader.num_values()); + assert_eq!(3, record_reader.read_records(3).unwrap()); + assert_eq!(5, record_reader.num_records()); + assert_eq!(5, record_reader.num_values()); + } + + // Second page + + // Records data: + // test_schema + // leaf: 8 + // test_schema + // leaf: 9 + { + let values = [8, 9]; + let mut pb = DataPageBuilderImpl::new(desc, 2, true); + pb.add_values::(Encoding::PLAIN, &values); + let page = pb.consume(); + + let page_reader = Box::new(InMemoryPageReader::new(vec![page])); + record_reader.set_page_reader(page_reader).unwrap(); + assert_eq!(2, record_reader.read_records(10).unwrap()); + assert_eq!(7, record_reader.num_records()); + assert_eq!(7, record_reader.num_values()); + } + + assert_eq!(record_reader.consume_record_data(), &[4, 7, 6, 3, 2, 8, 9]); + assert_eq!(None, record_reader.consume_def_levels()); + assert_eq!(None, record_reader.consume_bitmap()); + } + + #[test] + fn test_capacity_hint_preserved_across_fragmented_reads() { + let message_type = " + message test_schema { + REQUIRED INT32 leaf; + } + "; + let desc = parse_message_type(message_type) + .map(|t| SchemaDescriptor::new(Arc::new(t))) + .map(|s| s.column(0)) + .unwrap(); + + let mut record_reader = RecordReader::::new(desc.clone(), 8); + let values = [1, 2, 3, 4]; + let mut pb = DataPageBuilderImpl::new(desc, 4, true); + pb.add_values::(Encoding::PLAIN, &values); + let page = pb.consume(); + + let page_reader = Box::new(InMemoryPageReader::new(vec![page])); + record_reader.set_page_reader(page_reader).unwrap(); + + assert_eq!(1, record_reader.read_records(1).unwrap()); + assert_eq!(1, record_reader.read_records(1).unwrap()); + + let record_data = record_reader.consume_record_data(); + assert_eq!(record_data, &[1, 2]); + assert!( + record_data.capacity() >= 8, + "capacity hint should survive fragmented reads, got {}", + record_data.capacity() + ); + } + + #[test] + fn test_read_optional_records() { + // Construct column schema + let message_type = " + message test_schema { + OPTIONAL Group test_struct { + OPTIONAL INT32 leaf; + } + } + "; + + let desc = parse_message_type(message_type) + .map(|t| SchemaDescriptor::new(Arc::new(t))) + .map(|s| s.column(0)) + .unwrap(); + + // Construct record reader + let mut record_reader = RecordReader::::new(desc.clone(), DEFAULT_BATCH_SIZE); + + // First page + + // Records data: + // test_schema + // test_struct + // test_schema + // test_struct + // left: 7 + // test_schema + // test_schema + // test_struct + // leaf: 6 + // test_schema + // test_struct + // leaf: 6 + { + let values = [7, 6, 3]; + //empty, non-empty, empty, non-empty, non-empty + let def_levels = [1i16, 2i16, 0i16, 2i16, 2i16]; + let mut pb = DataPageBuilderImpl::new(desc.clone(), 5, true); + pb.add_def_levels(2, &def_levels); + pb.add_values::(Encoding::PLAIN, &values); + let page = pb.consume(); + + let page_reader = Box::new(InMemoryPageReader::new(vec![page])); + record_reader.set_page_reader(page_reader).unwrap(); + assert_eq!(2, record_reader.read_records(2).unwrap()); + assert_eq!(2, record_reader.num_records()); + assert_eq!(2, record_reader.num_values()); + assert_eq!(3, record_reader.read_records(3).unwrap()); + assert_eq!(5, record_reader.num_records()); + assert_eq!(5, record_reader.num_values()); + } + + // Second page + + // Records data: + // test_schema + // test_schema + // test_struct + // left: 8 + { + let values = [8]; + //empty, non-empty + let def_levels = [0i16, 2i16]; + let mut pb = DataPageBuilderImpl::new(desc, 2, true); + pb.add_def_levels(2, &def_levels); + pb.add_values::(Encoding::PLAIN, &values); + let page = pb.consume(); + + let page_reader = Box::new(InMemoryPageReader::new(vec![page])); + record_reader.set_page_reader(page_reader).unwrap(); + assert_eq!(2, record_reader.read_records(10).unwrap()); + assert_eq!(7, record_reader.num_records()); + assert_eq!(7, record_reader.num_values()); + } + + // Verify result def levels + assert_eq!( + Some(vec![1i16, 2i16, 0i16, 2i16, 2i16, 0i16, 2i16]), + record_reader.consume_def_levels() + ); + + // Verify bitmap + let expected_valid = &[false, true, false, true, true, false, true]; + let expected_buffer = Buffer::from_iter(expected_valid.iter().cloned()); + assert_eq!(Some(expected_buffer), record_reader.consume_bitmap()); + + // Verify result record data + let actual = record_reader.consume_record_data(); + + let expected = &[0, 7, 0, 6, 3, 0, 8]; + assert_eq!(actual.len(), expected.len()); + + // Only validate valid values are equal + let iter = expected_valid.iter().zip(&actual).zip(expected); + for ((valid, actual), expected) in iter { + if *valid { + assert_eq!(actual, expected) + } + } + } + + #[test] + fn test_read_repeated_records() { + // Construct column schema + let message_type = " + message test_schema { + REPEATED Group test_struct { + REPEATED INT32 leaf; + } + } + "; + + let desc = parse_message_type(message_type) + .map(|t| SchemaDescriptor::new(Arc::new(t))) + .map(|s| s.column(0)) + .unwrap(); + + // Construct record reader + let mut record_reader = RecordReader::::new(desc.clone(), DEFAULT_BATCH_SIZE); + + // First page + + // Records data: + // test_schema + // test_struct + // leaf: 4 + // test_schema + // test_schema + // test_struct + // test_struct + // leaf: 7 + // leaf: 6 + // leaf: 3 + // test_struct + // leaf: 2 + { + let values = [4, 7, 6, 3, 2]; + let def_levels = [2i16, 0i16, 1i16, 2i16, 2i16, 2i16, 2i16]; + let rep_levels = [0i16, 0i16, 0i16, 1i16, 2i16, 2i16, 1i16]; + let mut pb = DataPageBuilderImpl::new(desc.clone(), 7, true); + pb.add_rep_levels(2, &rep_levels); + pb.add_def_levels(2, &def_levels); + pb.add_values::(Encoding::PLAIN, &values); + let page = pb.consume(); + + let page_reader = Box::new(InMemoryPageReader::new(vec![page])); + record_reader.set_page_reader(page_reader).unwrap(); + + assert_eq!(1, record_reader.read_records(1).unwrap()); + assert_eq!(1, record_reader.num_records()); + assert_eq!(1, record_reader.num_values()); + assert_eq!(2, record_reader.read_records(3).unwrap()); + assert_eq!(3, record_reader.num_records()); + assert_eq!(7, record_reader.num_values()); + } + + // Second page + + // Records data: + // test_schema + // test_struct + // leaf: 8 + // leaf: 9 + { + let values = [8, 9]; + let def_levels = [2i16, 2i16]; + let rep_levels = [0i16, 2i16]; + let mut pb = DataPageBuilderImpl::new(desc, 2, true); + pb.add_rep_levels(2, &rep_levels); + pb.add_def_levels(2, &def_levels); + pb.add_values::(Encoding::PLAIN, &values); + let page = pb.consume(); + + let page_reader = Box::new(InMemoryPageReader::new(vec![page])); + record_reader.set_page_reader(page_reader).unwrap(); + + assert_eq!(1, record_reader.read_records(10).unwrap()); + assert_eq!(4, record_reader.num_records()); + assert_eq!(9, record_reader.num_values()); + } + + // Verify result def levels + assert_eq!( + Some(vec![2i16, 0i16, 1i16, 2i16, 2i16, 2i16, 2i16, 2i16, 2i16]), + record_reader.consume_def_levels() + ); + + // Verify bitmap + let expected_valid = &[true, false, false, true, true, true, true, true, true]; + let expected_buffer = Buffer::from_iter(expected_valid.iter().cloned()); + assert_eq!(Some(expected_buffer), record_reader.consume_bitmap()); + + // Verify result record data + let actual = record_reader.consume_record_data(); + let expected = &[4, 0, 0, 7, 6, 3, 2, 8, 9]; + assert_eq!(actual.len(), expected.len()); + + // Only validate valid values are equal + let iter = expected_valid.iter().zip(&actual).zip(expected); + for ((valid, actual), expected) in iter { + if *valid { + assert_eq!(actual, expected) + } + } + } + + #[test] + fn test_selective_padding_consumes_full_bitmap() { + let message_type = " + message test_schema { + OPTIONAL GROUP my_list (LIST) { + REPEATED GROUP list { + OPTIONAL INT32 element; + } + } + } + "; + + let desc = parse_message_type(message_type) + .map(|t| SchemaDescriptor::new(Arc::new(t))) + .map(|s| s.column(0)) + .unwrap(); + + let mut record_reader = RecordReader::::new(desc.clone(), DEFAULT_BATCH_SIZE); + record_reader.set_padding_threshold(2); + + let values = [10, 20]; + let def_levels = [3i16, 2i16, 0i16, 2i16, 3i16]; + let rep_levels = [0i16, 1i16, 0i16, 0i16, 1i16]; + let mut pb = DataPageBuilderImpl::new(desc, 5, true); + pb.add_rep_levels(1, &rep_levels); + pb.add_def_levels(3, &def_levels); + pb.add_values::(Encoding::PLAIN, &values); + + let page_reader = Box::new(InMemoryPageReader::new(vec![pb.consume()])); + record_reader.set_page_reader(page_reader).unwrap(); + assert_eq!(3, record_reader.read_records(3).unwrap()); + + let expected_compact = Buffer::from_iter([true, false, false, true]); + assert_eq!( + Some(expected_compact), + record_reader.consume_compact_bitmap() + ); + assert_eq!(None, record_reader.consume_bitmap()); + } + + #[test] + fn test_selective_padding_reserves_compact_capacity() { + let message_type = " + message test_schema { + OPTIONAL GROUP my_list (LIST) { + REPEATED GROUP list { + OPTIONAL INT32 element; + } + } + } + "; + + let desc = parse_message_type(message_type) + .map(|t| SchemaDescriptor::new(Arc::new(t))) + .map(|s| s.column(0)) + .unwrap(); + + let mut record_reader = RecordReader::::new(desc.clone(), DEFAULT_BATCH_SIZE); + record_reader.set_padding_threshold(2); + + let values = [10]; + let mut def_levels = vec![0i16; 30]; + def_levels.extend([3i16, 2i16]); + let rep_levels = vec![0i16; def_levels.len()]; + + let mut pb = DataPageBuilderImpl::new(desc, def_levels.len() as u32, true); + pb.add_rep_levels(1, &rep_levels); + pb.add_def_levels(3, &def_levels); + pb.add_values::(Encoding::PLAIN, &values); + + let page_reader = Box::new(InMemoryPageReader::new(vec![pb.consume()])); + record_reader.set_page_reader(page_reader).unwrap(); + assert_eq!(def_levels.len(), record_reader.read_records(128).unwrap()); + + let expected_compact = Buffer::from_iter([true, false]); + assert_eq!( + Some(expected_compact), + record_reader.consume_compact_bitmap() + ); + + let actual = record_reader.consume_record_data(); + assert_eq!(actual.len(), 2); + assert_eq!(actual[0], 10); + assert!( + actual.capacity() < def_levels.len(), + "selective padding should not reserve one child slot per list-level NULL" + ); + } + + #[test] + fn test_read_more_than_one_batch() { + // Construct column schema + let message_type = " + message test_schema { + REPEATED INT32 leaf; + } + "; + + let desc = parse_message_type(message_type) + .map(|t| SchemaDescriptor::new(Arc::new(t))) + .map(|s| s.column(0)) + .unwrap(); + + // Construct record reader + let mut record_reader = RecordReader::::new(desc.clone(), DEFAULT_BATCH_SIZE); + + { + let values = [100; 5000]; + let def_levels = [1i16; 5000]; + let mut rep_levels = [1i16; 5000]; + for idx in 0..1000 { + rep_levels[idx * 5] = 0i16; + } + + let mut pb = DataPageBuilderImpl::new(desc, 5000, true); + pb.add_rep_levels(1, &rep_levels); + pb.add_def_levels(1, &def_levels); + pb.add_values::(Encoding::PLAIN, &values); + let page = pb.consume(); + + let page_reader = Box::new(InMemoryPageReader::new(vec![page])); + record_reader.set_page_reader(page_reader).unwrap(); + + assert_eq!(1000, record_reader.read_records(1000).unwrap()); + assert_eq!(1000, record_reader.num_records()); + assert_eq!(5000, record_reader.num_values()); + } + } + + #[test] + fn test_row_group_boundary() { + // Construct column schema + let message_type = " + message test_schema { + REPEATED Group test_struct { + REPEATED INT32 leaf; + } + } + "; + + let desc = parse_message_type(message_type) + .map(|t| SchemaDescriptor::new(Arc::new(t))) + .map(|s| s.column(0)) + .unwrap(); + + let values = [1, 2, 3]; + let def_levels = [1i16, 0i16, 1i16, 2i16, 2i16, 1i16, 2i16]; + let rep_levels = [0i16, 0i16, 0i16, 1i16, 2i16, 0i16, 1i16]; + let mut pb = DataPageBuilderImpl::new(desc.clone(), 7, true); + pb.add_rep_levels(2, &rep_levels); + pb.add_def_levels(2, &def_levels); + pb.add_values::(Encoding::PLAIN, &values); + let page = pb.consume(); + + let mut record_reader = RecordReader::::new(desc, DEFAULT_BATCH_SIZE); + let page_reader = Box::new(InMemoryPageReader::new(vec![page.clone()])); + record_reader.set_page_reader(page_reader).unwrap(); + assert_eq!(record_reader.read_records(4).unwrap(), 4); + assert_eq!(record_reader.num_records(), 4); + assert_eq!(record_reader.num_values(), 7); + + assert_eq!(record_reader.read_records(4).unwrap(), 0); + assert_eq!(record_reader.num_records(), 4); + assert_eq!(record_reader.num_values(), 7); + + record_reader.read_records(4).unwrap(); + + let page_reader = Box::new(InMemoryPageReader::new(vec![page])); + record_reader.set_page_reader(page_reader).unwrap(); + + assert_eq!(record_reader.read_records(4).unwrap(), 4); + assert_eq!(record_reader.num_records(), 8); + assert_eq!(record_reader.num_values(), 14); + + assert_eq!(record_reader.read_records(4).unwrap(), 0); + assert_eq!(record_reader.num_records(), 8); + assert_eq!(record_reader.num_values(), 14); + } + + #[test] + fn test_skip_required_records() { + // Construct column schema + let message_type = " + message test_schema { + REQUIRED INT32 leaf; + } + "; + let desc = parse_message_type(message_type) + .map(|t| SchemaDescriptor::new(Arc::new(t))) + .map(|s| s.column(0)) + .unwrap(); + + // Construct record reader + let mut record_reader = RecordReader::::new(desc.clone(), DEFAULT_BATCH_SIZE); + + // First page + + // Records data: + // test_schema + // leaf: 4 + // test_schema + // leaf: 7 + // test_schema + // leaf: 6 + // test_schema + // left: 3 + // test_schema + // left: 2 + { + let values = [4, 7, 6, 3, 2]; + let mut pb = DataPageBuilderImpl::new(desc.clone(), 5, true); + pb.add_values::(Encoding::PLAIN, &values); + let page = pb.consume(); + + let page_reader = Box::new(InMemoryPageReader::new(vec![page])); + record_reader.set_page_reader(page_reader).unwrap(); + assert_eq!(2, record_reader.skip_records(2).unwrap()); + assert_eq!(0, record_reader.num_records()); + assert_eq!(0, record_reader.num_values()); + assert_eq!(3, record_reader.read_records(3).unwrap()); + assert_eq!(3, record_reader.num_records()); + assert_eq!(3, record_reader.num_values()); + } + + // Second page + + // Records data: + // test_schema + // leaf: 8 + // test_schema + // leaf: 9 + { + let values = [8, 9]; + let mut pb = DataPageBuilderImpl::new(desc, 2, true); + pb.add_values::(Encoding::PLAIN, &values); + let page = pb.consume(); + + let page_reader = Box::new(InMemoryPageReader::new(vec![page])); + record_reader.set_page_reader(page_reader).unwrap(); + assert_eq!(2, record_reader.skip_records(10).unwrap()); + assert_eq!(3, record_reader.num_records()); + assert_eq!(3, record_reader.num_values()); + assert_eq!(0, record_reader.read_records(10).unwrap()); + } + + assert_eq!(record_reader.consume_record_data(), &[6, 3, 2]); + assert_eq!(None, record_reader.consume_def_levels()); + assert_eq!(None, record_reader.consume_bitmap()); + } + + #[test] + fn test_skip_optional_records() { + // Construct column schema + let message_type = " + message test_schema { + OPTIONAL Group test_struct { + OPTIONAL INT32 leaf; + } + } + "; + + let desc = parse_message_type(message_type) + .map(|t| SchemaDescriptor::new(Arc::new(t))) + .map(|s| s.column(0)) + .unwrap(); + + // Construct record reader + let mut record_reader = RecordReader::::new(desc.clone(), DEFAULT_BATCH_SIZE); + + // First page + + // Records data: + // test_schema + // test_struct + // test_schema + // test_struct + // leaf: 7 + // test_schema + // test_schema + // test_struct + // leaf: 6 + // test_schema + // test_struct + // leaf: 6 + { + let values = [7, 6, 3]; + //empty, non-empty, empty, non-empty, non-empty + let def_levels = [1i16, 2i16, 0i16, 2i16, 2i16]; + let mut pb = DataPageBuilderImpl::new(desc.clone(), 5, true); + pb.add_def_levels(2, &def_levels); + pb.add_values::(Encoding::PLAIN, &values); + let page = pb.consume(); + + let page_reader = Box::new(InMemoryPageReader::new(vec![page])); + record_reader.set_page_reader(page_reader).unwrap(); + assert_eq!(2, record_reader.skip_records(2).unwrap()); + assert_eq!(0, record_reader.num_records()); + assert_eq!(0, record_reader.num_values()); + assert_eq!(3, record_reader.read_records(3).unwrap()); + assert_eq!(3, record_reader.num_records()); + assert_eq!(3, record_reader.num_values()); + } + + // Second page + + // Records data: + // test_schema + // test_schema + // test_struct + // left: 8 + { + let values = [8]; + //empty, non-empty + let def_levels = [0i16, 2i16]; + let mut pb = DataPageBuilderImpl::new(desc, 2, true); + pb.add_def_levels(2, &def_levels); + pb.add_values::(Encoding::PLAIN, &values); + let page = pb.consume(); + + let page_reader = Box::new(InMemoryPageReader::new(vec![page])); + record_reader.set_page_reader(page_reader).unwrap(); + assert_eq!(2, record_reader.skip_records(10).unwrap()); + assert_eq!(3, record_reader.num_records()); + assert_eq!(3, record_reader.num_values()); + assert_eq!(0, record_reader.read_records(10).unwrap()); + } + + // Verify result def levels + assert_eq!( + Some(vec![0i16, 2i16, 2i16]), + record_reader.consume_def_levels() + ); + + // Verify bitmap + let expected_valid = &[false, true, true]; + let expected_buffer = Buffer::from_iter(expected_valid.iter().cloned()); + assert_eq!(Some(expected_buffer), record_reader.consume_bitmap()); + + // Verify result record data + let actual = record_reader.consume_record_data(); + + let expected = &[0, 6, 3]; + assert_eq!(actual.len(), expected.len()); + + // Only validate valid values are equal + let iter = expected_valid.iter().zip(&actual).zip(expected); + for ((valid, actual), expected) in iter { + if *valid { + assert_eq!(actual, expected) + } + } + } +} diff --git a/vendor/parquet/src/arrow/schema/complex.rs b/vendor/parquet/src/arrow/schema/complex.rs new file mode 100644 index 0000000..161b263 --- /dev/null +++ b/vendor/parquet/src/arrow/schema/complex.rs @@ -0,0 +1,1881 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::collections::HashMap; +use std::sync::Arc; + +use crate::arrow::schema::extension::try_add_extension_type; +use crate::arrow::schema::primitive::convert_primitive; +use crate::arrow::schema::virtual_type::{RowGroupIndex, RowNumber}; +use crate::arrow::{PARQUET_FIELD_ID_META_KEY, ProjectionMask}; +use crate::basic::{ConvertedType, Repetition}; +use crate::errors::ParquetError; +use crate::errors::Result; +use crate::schema::types::{SchemaDescriptor, Type, TypePtr}; +use arrow_schema::{DataType, Field, Fields, SchemaBuilder, extension::ExtensionType}; + +fn get_repetition(t: &Type) -> Repetition { + let info = t.get_basic_info(); + match info.has_repetition() { + true => info.repetition(), + false => Repetition::REQUIRED, + } +} + +/// Representation of a parquet schema element, in terms of arrow schema elements +#[derive(Debug, Clone)] +pub struct ParquetField { + /// The level which represents an insertion into the current list + /// i.e. guaranteed to be > 0 for an element of list type + pub rep_level: i16, + /// The level at which this field is fully defined, + /// i.e. guaranteed to be > 0 for a nullable type or child of a + /// nullable type + pub def_level: i16, + /// Whether this field is nullable + pub nullable: bool, + /// The arrow type of the column data + /// + /// Note: In certain cases the data stored in parquet may have been coerced + /// to a different type and will require conversion on read (e.g. Date64 and Interval) + pub arrow_type: DataType, + /// The type of this field + pub field_type: ParquetFieldType, +} + +impl ParquetField { + /// Converts `self` into an arrow list, with its current type as the field type + /// + /// This is used to convert repeated columns, into their arrow representation + fn into_list(self, name: &str) -> Self { + ParquetField { + rep_level: self.rep_level, + def_level: self.def_level, + nullable: false, + arrow_type: DataType::List(Arc::new(Field::new(name, self.arrow_type.clone(), false))), + field_type: ParquetFieldType::Group { + children: vec![self], + }, + } + } + + /// Converts `self` into an arrow list, with its current type as the field type + /// accept an optional `list_data_type` to specify the type of list to create + /// + /// This is used to convert [deprecated repeated columns] (not in a list), into their arrow representation + /// + /// [deprecated repeated columns]: https://github.com/apache/parquet-format/blob/9fd57b59e0ce1a82a69237dcf8977d3e72a2965d/LogicalTypes.md?plain=1#L649-L650 + fn into_list_with_arrow_list_hint( + self, + parquet_field_type: &Type, + list_data_type: Option, + ) -> Result { + let arrow_field = match &list_data_type { + Some(DataType::List(field_hint)) + | Some(DataType::LargeList(field_hint)) + | Some(DataType::FixedSizeList(field_hint, _)) => Some(field_hint.as_ref()), + Some(_) => { + return Err(general_err!( + "Internal error: should be validated earlier that list_data_type is only a type of list" + )); + } + None => None, + }; + + let arrow_field = convert_field( + parquet_field_type, + &self, + arrow_field, + // Only add the field id to the list and not to the element + false, + )? + .with_nullable(false); + + Ok(ParquetField { + rep_level: self.rep_level, + def_level: self.def_level, + nullable: false, + arrow_type: match list_data_type { + Some(DataType::List(_)) => DataType::List(Arc::new(arrow_field)), + Some(DataType::LargeList(_)) => DataType::LargeList(Arc::new(arrow_field)), + Some(DataType::FixedSizeList(_, len)) => { + DataType::FixedSizeList(Arc::new(arrow_field), len) + } + _ => DataType::List(Arc::new(arrow_field)), + }, + field_type: ParquetFieldType::Group { + children: vec![self], + }, + }) + } + + /// Returns a list of [`ParquetField`] children if this is a group type + pub fn children(&self) -> Option<&[Self]> { + match &self.field_type { + ParquetFieldType::Primitive { .. } => None, + ParquetFieldType::Group { children } => Some(children), + ParquetFieldType::Virtual(_) => None, + } + } +} + +/// Types of virtual columns that can be computed at read time +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum VirtualColumnType { + /// Row number within the file + RowNumber, + /// Row group index + RowGroupIndex, +} + +#[derive(Debug, Clone)] +pub enum ParquetFieldType { + Primitive { + /// The index of the column in parquet + col_idx: usize, + /// The type of the column in parquet + primitive_type: TypePtr, + }, + Group { + children: Vec, + }, + /// Virtual column that doesn't exist in the parquet file + /// but is computed at read time (e.g., row_number) + Virtual(VirtualColumnType), +} + +/// Encodes the context of the parent of the field currently under consideration +struct VisitorContext { + rep_level: i16, + def_level: i16, + /// An optional [`DataType`] sourced from the embedded arrow schema + data_type: Option, + + /// Whether to treat repeated types as list from arrow types + /// when true, if data_type provided it should be DataType::List() (or other list type) + /// and the list field data type would be treated as the hint for the parquet type + /// + /// when false, if data_type provided it will be treated as the hint without unwrapping + /// + /// This is for supporting [deprecated parquet list representation][1] + /// + /// [1]: https://github.com/apache/parquet-format/blob/38818fa0e7efd54b535001a4448030a40619c2a3/LogicalTypes.md?plain=1#L718-L806 + treat_repeated_as_list_arrow_hint: bool, +} + +impl VisitorContext { + /// Compute the resulting definition level, repetition level and nullability + /// for a child field with the given [`Repetition`] + fn levels(&self, repetition: Repetition) -> (i16, i16, bool) { + match repetition { + Repetition::OPTIONAL => (self.def_level + 1, self.rep_level, true), + Repetition::REQUIRED => (self.def_level, self.rep_level, false), + Repetition::REPEATED => (self.def_level + 1, self.rep_level + 1, false), + } + } +} + +/// Walks the parquet schema in a depth-first fashion in order to map it to arrow data structures +/// +/// See [Logical Types] for more information on the conversion algorithm +/// +/// [Logical Types]: https://github.com/apache/parquet-format/blob/master/LogicalTypes.md +struct Visitor { + /// The column index of the next leaf column + next_col_idx: usize, + + /// Mask of columns to include + mask: ProjectionMask, +} + +impl Visitor { + fn visit_primitive( + &mut self, + primitive_type: &TypePtr, + context: VisitorContext, + ) -> Result> { + let col_idx = self.next_col_idx; + self.next_col_idx += 1; + + if !self.mask.leaf_included(col_idx) { + return Ok(None); + } + + let repetition = get_repetition(primitive_type); + let (def_level, rep_level, nullable) = context.levels(repetition); + + let primitive_arrow_data_type = match repetition { + Repetition::REPEATED if context.treat_repeated_as_list_arrow_hint => { + let arrow_field = match &context.data_type { + Some(DataType::List(f)) => Some(f.as_ref()), + Some(DataType::LargeList(f)) => Some(f.as_ref()), + Some(DataType::FixedSizeList(f, _)) => Some(f.as_ref()), + Some(d) => { + return Err(arrow_err!( + "incompatible arrow schema, expected list got {} for repeated primitive field", + d + )); + } + None => None, + }; + + arrow_field.map(|f| f.data_type().clone()) + } + _ => context.data_type.clone(), + }; + + let arrow_type = convert_primitive(primitive_type, primitive_arrow_data_type)?; + + let primitive_field = ParquetField { + rep_level, + def_level, + nullable, + arrow_type, + field_type: ParquetFieldType::Primitive { + primitive_type: primitive_type.clone(), + col_idx, + }, + }; + + Ok(Some(match repetition { + Repetition::REPEATED if context.treat_repeated_as_list_arrow_hint => { + primitive_field.into_list_with_arrow_list_hint(primitive_type, context.data_type)? + } + Repetition::REPEATED => primitive_field.into_list(primitive_type.name()), + _ => primitive_field, + })) + } + + fn visit_struct( + &mut self, + struct_type: &TypePtr, + context: VisitorContext, + ) -> Result> { + // The root type will not have a repetition level + let repetition = get_repetition(struct_type); + let (def_level, rep_level, nullable) = context.levels(repetition); + + let parquet_fields = struct_type.get_fields(); + + // Extract any arrow fields from the hints + let arrow_struct = match repetition { + Repetition::REPEATED if context.treat_repeated_as_list_arrow_hint => { + let arrow_field = match &context.data_type { + Some(DataType::List(f)) => Some(f.as_ref()), + Some(DataType::LargeList(f)) => Some(f.as_ref()), + Some(DataType::FixedSizeList(f, _)) => Some(f.as_ref()), + Some(d) => { + return Err(arrow_err!( + "incompatible arrow schema, expected list got {} for repeated struct field", + d + )); + } + None => None, + }; + + arrow_field.map(|f| f.data_type()) + } + _ => context.data_type.as_ref(), + }; + + let arrow_fields = match &arrow_struct { + Some(DataType::Struct(fields)) => { + if fields.len() != parquet_fields.len() { + return Err(arrow_err!( + "incompatible arrow schema, expected {} struct fields got {}", + parquet_fields.len(), + fields.len() + )); + } + Some(fields) + } + Some(d) => { + return Err(arrow_err!( + "incompatible arrow schema, expected struct got {}", + d + )); + } + None => None, + }; + + let mut child_fields = SchemaBuilder::with_capacity(parquet_fields.len()); + let mut children = Vec::with_capacity(parquet_fields.len()); + + // Perform a DFS of children + for (idx, parquet_field) in parquet_fields.iter().enumerate() { + let data_type = match arrow_fields { + Some(fields) => { + let field = &fields[idx]; + if field.name() != parquet_field.name() { + return Err(arrow_err!( + "incompatible arrow schema, expected field named {} got {}", + parquet_field.name(), + field.name() + )); + } + Some(field.data_type().clone()) + } + None => None, + }; + + let arrow_field = arrow_fields.map(|x| &*x[idx]); + let child_ctx = VisitorContext { + rep_level, + def_level, + data_type, + + // Always true: each child is independently responsible for its own + // repeated-to-list conversion. The parent's flag may be false when + // this struct's own repetition is consumed by an outer visit_list + // backward-compat path, but that only applies to the struct itself, + // not its children. A repeated child's arrow hint will be List<...> + // and needs to be unwrapped accordingly. + treat_repeated_as_list_arrow_hint: true, + }; + + if let Some(child) = self.dispatch(parquet_field, child_ctx)? { + // The child type returned may be different from what is encoded in the arrow + // schema in the event of a mismatch or a projection + child_fields.push(convert_field(parquet_field, &child, arrow_field, true)?); + children.push(child); + } + } + + if children.is_empty() { + return Ok(None); + } + + let struct_field = ParquetField { + rep_level, + def_level, + nullable, + arrow_type: DataType::Struct(child_fields.finish().fields), + field_type: ParquetFieldType::Group { children }, + }; + + Ok(Some(match repetition { + Repetition::REPEATED if context.treat_repeated_as_list_arrow_hint => { + struct_field.into_list_with_arrow_list_hint(struct_type, context.data_type)? + } + Repetition::REPEATED => struct_field.into_list(struct_type.name()), + _ => struct_field, + })) + } + + fn visit_map( + &mut self, + map_type: &TypePtr, + context: VisitorContext, + ) -> Result> { + let rep_level = context.rep_level + 1; + let (def_level, nullable) = match get_repetition(map_type) { + Repetition::REQUIRED => (context.def_level + 1, false), + Repetition::OPTIONAL => (context.def_level + 2, true), + Repetition::REPEATED => return Err(arrow_err!("Map cannot be repeated")), + }; + + if map_type.get_fields().len() != 1 { + return Err(arrow_err!( + "Map field must have exactly one key_value child, found {}", + map_type.get_fields().len() + )); + } + + // Add map entry (key_value) to context + let map_key_value = &map_type.get_fields()[0]; + if map_key_value.get_basic_info().repetition() != Repetition::REPEATED { + return Err(arrow_err!("Child of map field must be repeated")); + } + + // According to the specification the values are optional (#1642). + // In this case, return the keys as a list. + if map_key_value.get_fields().len() == 1 { + return self.visit_list(map_type, context); + } + + if map_key_value.get_fields().len() != 2 { + return Err(arrow_err!( + "Child of map field must have two children, found {}", + map_key_value.get_fields().len() + )); + } + + // Get key and value, and create context for each + let map_key = &map_key_value.get_fields()[0]; + let map_value = &map_key_value.get_fields()[1]; + + match map_key.get_basic_info().repetition() { + Repetition::REPEATED => { + return Err(arrow_err!("Map keys cannot be repeated")); + } + Repetition::REQUIRED | Repetition::OPTIONAL => { + // Relaxed check for having repetition REQUIRED as there exists + // parquet writers and files that do not conform to this standard. + // This allows us to consume a broader range of existing files even + // if they are out of spec. + } + } + + if map_value.get_basic_info().repetition() == Repetition::REPEATED { + return Err(arrow_err!("Map values cannot be repeated")); + } + + // Extract the arrow fields + let (arrow_map, arrow_key, arrow_value, sorted) = match &context.data_type { + Some(DataType::Map(field, sorted)) => match field.data_type() { + DataType::Struct(fields) => { + if fields.len() != 2 { + return Err(arrow_err!( + "Map data type should contain struct with two children, got {}", + fields.len() + )); + } + + (Some(field), Some(&*fields[0]), Some(&*fields[1]), *sorted) + } + d => { + return Err(arrow_err!("Map data type should contain struct got {}", d)); + } + }, + Some(d) => { + return Err(arrow_err!( + "incompatible arrow schema, expected map got {}", + d + )); + } + None => (None, None, None, false), + }; + + let maybe_key = { + let context = VisitorContext { + rep_level, + def_level, + data_type: arrow_key.map(|x| x.data_type().clone()), + // Key is not repeated + treat_repeated_as_list_arrow_hint: false, + }; + + self.dispatch(map_key, context)? + }; + + let maybe_value = { + let context = VisitorContext { + rep_level, + def_level, + data_type: arrow_value.map(|x| x.data_type().clone()), + // Value type can be repeated + treat_repeated_as_list_arrow_hint: true, + }; + + self.dispatch(map_value, context)? + }; + + // Need both columns to be projected + match (maybe_key, maybe_value) { + (Some(key), Some(value)) => { + let key_field = Arc::new( + convert_field(map_key, &key, arrow_key, true)? + // The key is always non-nullable (#5630) + .with_nullable(false), + ); + let value_field = Arc::new(convert_field(map_value, &value, arrow_value, true)?); + let field_metadata = match arrow_map { + Some(field) => field.metadata().clone(), + _ => HashMap::default(), + }; + + let map_field = Field::new_struct( + map_key_value.name(), + [key_field, value_field], + false, // The inner map field is always non-nullable (#1697) + ) + .with_metadata(field_metadata); + + Ok(Some(ParquetField { + rep_level, + def_level, + nullable, + arrow_type: DataType::Map(Arc::new(map_field), sorted), + field_type: ParquetFieldType::Group { + children: vec![key, value], + }, + })) + } + _ => Ok(None), + } + } + + fn visit_list( + &mut self, + list_type: &TypePtr, + context: VisitorContext, + ) -> Result> { + if list_type.is_primitive() { + return Err(arrow_err!( + "{:?} is a list type and can't be processed as primitive.", + list_type + )); + } + + let fields = list_type.get_fields(); + if fields.len() != 1 { + return Err(arrow_err!( + "list type must have a single child, found {}", + fields.len() + )); + } + + let repeated_field = &fields[0]; + if get_repetition(repeated_field) != Repetition::REPEATED { + return Err(arrow_err!("List child must be repeated")); + } + + // If the list is nullable + let (def_level, nullable) = match list_type.get_basic_info().repetition() { + Repetition::REQUIRED => (context.def_level, false), + Repetition::OPTIONAL => (context.def_level + 1, true), + Repetition::REPEATED => return Err(arrow_err!("List type cannot be repeated")), + }; + + let arrow_field = match &context.data_type { + Some(DataType::List(f)) => Some(f.as_ref()), + Some(DataType::LargeList(f)) => Some(f.as_ref()), + Some(DataType::FixedSizeList(f, _)) => Some(f.as_ref()), + Some(DataType::ListView(f)) => Some(f.as_ref()), + Some(DataType::LargeListView(f)) => Some(f.as_ref()), + Some(d) => { + return Err(arrow_err!( + "incompatible arrow schema, expected list got {}", + d + )); + } + None => None, + }; + + if repeated_field.is_primitive() { + // If the repeated field is not a group, then its type is the element type and elements are required. + // + // required/optional group my_list (LIST) { + // repeated int32 element; + // } + // + let context = VisitorContext { + rep_level: context.rep_level, + def_level, + data_type: arrow_field.map(|f| f.data_type().clone()), + treat_repeated_as_list_arrow_hint: false, + }; + + return match self.visit_primitive(repeated_field, context) { + Ok(Some(mut field)) => { + // visit_primitive will infer a non-nullable list, update if necessary + field.nullable = nullable; + Ok(Some(field)) + } + r => r, + }; + } + + // test to see if the repeated field is a struct or one-tuple + let items = repeated_field.get_fields(); + if items.len() != 1 + || (!repeated_field.is_list() + && !repeated_field.has_single_repeated_child() + && (repeated_field.name() == "array" + || repeated_field.name() == format!("{}_tuple", list_type.name()))) + { + // If the repeated field is a group with multiple fields, then its type is the element + // type and elements are required. + // + // If the repeated field is a group with one field and is named either array or uses + // the LIST-annotated group's name with _tuple appended then the repeated type is the + // element type and elements are required. But this rule only applies if the + // repeated field is not annotated, and the single child field is not `repeated`. + let context = VisitorContext { + rep_level: context.rep_level, + def_level, + data_type: arrow_field.map(|f| f.data_type().clone()), + treat_repeated_as_list_arrow_hint: false, + }; + + return match self.visit_struct(repeated_field, context) { + Ok(Some(mut field)) => { + field.nullable = nullable; + Ok(Some(field)) + } + r => r, + }; + } + + // Regular list handling logic + let item_type = &items[0]; + let rep_level = context.rep_level + 1; + let def_level = def_level + 1; + + let new_context = VisitorContext { + def_level, + rep_level, + data_type: arrow_field.map(|f| f.data_type().clone()), + treat_repeated_as_list_arrow_hint: true, + }; + + match self.dispatch(item_type, new_context) { + Ok(Some(item)) => { + let item_field = Arc::new(convert_field(item_type, &item, arrow_field, true)?); + + // Use arrow type as hint for index size + let arrow_type = match context.data_type { + Some(DataType::LargeList(_)) => DataType::LargeList(item_field), + Some(DataType::FixedSizeList(_, len)) => { + DataType::FixedSizeList(item_field, len) + } + Some(DataType::ListView(_)) => DataType::ListView(item_field), + Some(DataType::LargeListView(_)) => DataType::LargeListView(item_field), + _ => DataType::List(item_field), + }; + + Ok(Some(ParquetField { + rep_level, + def_level, + nullable, + arrow_type, + field_type: ParquetFieldType::Group { + children: vec![item], + }, + })) + } + r => r, + } + } + + fn dispatch( + &mut self, + cur_type: &TypePtr, + context: VisitorContext, + ) -> Result> { + if cur_type.is_primitive() { + self.visit_primitive(cur_type, context) + } else { + match cur_type.get_basic_info().converted_type() { + ConvertedType::LIST => self.visit_list(cur_type, context), + ConvertedType::MAP | ConvertedType::MAP_KEY_VALUE => { + self.visit_map(cur_type, context) + } + _ => self.visit_struct(cur_type, context), + } + } + } +} + +/// Converts a virtual Arrow [`Field`] to a [`ParquetField`] +/// +/// Virtual fields don't correspond to any data in the parquet file, +/// but are computed at read time (e.g., row_number) +/// +/// The levels are computed based on the parent context: +/// - If nullable: def_level = parent_def_level + 1 +/// - If required: def_level = parent_def_level +/// - rep_level = parent_rep_level (virtual fields are not repeated) +pub(super) fn convert_virtual_field( + arrow_field: &Field, + parent_rep_level: i16, + parent_def_level: i16, +) -> Result { + let nullable = arrow_field.is_nullable(); + let def_level = if nullable { + parent_def_level + 1 + } else { + parent_def_level + }; + + // Determine the virtual column type based on the extension type name + let extension_name = arrow_field.extension_type_name().ok_or_else(|| { + ParquetError::ArrowError(format!( + "virtual column field '{}' must have an extension type", + arrow_field.name() + )) + })?; + + let virtual_type = match extension_name { + RowNumber::NAME => VirtualColumnType::RowNumber, + RowGroupIndex::NAME => VirtualColumnType::RowGroupIndex, + _ => { + return Err(ParquetError::ArrowError(format!( + "unsupported virtual column type '{}' for field '{}'", + extension_name, + arrow_field.name() + ))); + } + }; + + Ok(ParquetField { + rep_level: parent_rep_level, + def_level, + nullable, + arrow_type: arrow_field.data_type().clone(), + field_type: ParquetFieldType::Virtual(virtual_type), + }) +} + +/// Computes the Arrow [`Field`] for a child column +/// +/// The resulting Arrow [`Field`] will have the type dictated by the Parquet `field`, a name +/// dictated by the `parquet_type`, and any metadata from `arrow_hint` +fn convert_field( + parquet_type: &Type, + field: &ParquetField, + arrow_hint: Option<&Field>, + add_field_id: bool, +) -> Result { + let name = parquet_type.name(); + let data_type = field.arrow_type.clone(); + let nullable = field.nullable; + + match arrow_hint { + Some(hint) => { + // If the inferred type is a dictionary, preserve dictionary metadata + #[allow(deprecated)] + let field = match (&data_type, hint.dict_id(), hint.dict_is_ordered()) { + (DataType::Dictionary(_, _), Some(id), Some(ordered)) => + { + #[allow(deprecated)] + Field::new_dict(name, data_type, nullable, id, ordered) + } + _ => Field::new(name, data_type, nullable), + }; + + Ok(field.with_metadata(hint.metadata().clone())) + } + None => { + let mut ret = Field::new(name, data_type, nullable); + let basic_info = parquet_type.get_basic_info(); + if add_field_id && basic_info.has_id() { + let mut meta = HashMap::with_capacity(1); + meta.insert( + PARQUET_FIELD_ID_META_KEY.to_string(), + basic_info.id().to_string(), + ); + ret.set_metadata(meta); + } + try_add_extension_type(ret, parquet_type) + } + } +} + +/// Computes the [`ParquetField`] for the provided [`SchemaDescriptor`] with `leaf_columns` listing +/// the indexes of leaf columns to project, and `embedded_arrow_schema` the optional +/// [`Fields`] embedded in the parquet metadata +/// +/// Note: This does not support out of order column projection +pub fn convert_schema( + schema: &SchemaDescriptor, + mask: ProjectionMask, + embedded_arrow_schema: Option<&Fields>, +) -> Result> { + let mut visitor = Visitor { + next_col_idx: 0, + mask, + }; + + let context = VisitorContext { + rep_level: 0, + def_level: 0, + data_type: embedded_arrow_schema.map(|fields| DataType::Struct(fields.clone())), + treat_repeated_as_list_arrow_hint: true, + }; + + visitor.dispatch(&schema.root_schema_ptr(), context) +} + +/// Computes the [`ParquetField`] for the provided `parquet_type` +pub fn convert_type(parquet_type: &TypePtr) -> Result { + let mut visitor = Visitor { + next_col_idx: 0, + mask: ProjectionMask::all(), + }; + + let context = VisitorContext { + rep_level: 0, + def_level: 0, + data_type: None, + // We might be inside list + treat_repeated_as_list_arrow_hint: false, + }; + + Ok(visitor.dispatch(parquet_type, context)?.unwrap()) +} + +#[cfg(test)] +mod tests { + use crate::arrow::schema::complex::convert_schema; + use crate::arrow::{PARQUET_FIELD_ID_META_KEY, ProjectionMask}; + use crate::schema::parser::parse_message_type; + use crate::schema::types::SchemaDescriptor; + use arrow_schema::{DataType, Field, Fields}; + use std::sync::Arc; + + trait WithFieldId { + fn with_field_id(self, id: i32) -> Self; + } + impl WithFieldId for arrow_schema::Field { + fn with_field_id(self, id: i32) -> Self { + let mut metadata = self.metadata().clone(); + metadata.insert(PARQUET_FIELD_ID_META_KEY.to_string(), id.to_string()); + self.with_metadata(metadata) + } + } + + fn test_roundtrip(message_type: &str) -> crate::errors::Result<()> { + let parsed_input_schema = Arc::new(parse_message_type(message_type)?); + let schema = SchemaDescriptor::new(parsed_input_schema); + + let converted = convert_schema(&schema, ProjectionMask::all(), None)?.unwrap(); + + let DataType::Struct(schema_fields) = &converted.arrow_type else { + panic!("Expected struct from convert_schema"); + }; + + // Should be able to convert the same thing + let converted_again = + convert_schema(&schema, ProjectionMask::all(), Some(schema_fields))?.unwrap(); + + // Assert that we changed to Utf8 + assert_eq!(converted_again.arrow_type, converted.arrow_type); + + Ok(()) + } + + fn test_expected_type( + message_type: &str, + expected_fields: Fields, + ) -> crate::errors::Result<()> { + test_roundtrip(message_type)?; + + let parsed_input_schema = Arc::new(parse_message_type(message_type)?); + let schema = SchemaDescriptor::new(parsed_input_schema); + + let converted = convert_schema(&schema, ProjectionMask::all(), None)?.unwrap(); + + let DataType::Struct(schema_fields) = &converted.arrow_type else { + panic!("Expected struct from convert_schema"); + }; + + assert_eq!(schema_fields, &expected_fields); + + Ok(()) + } + + /// Taken from the example in [Parquet Format - Nested Types - Lists - Backward-compatibility rules](https://github.com/apache/parquet-format/blob/9fd57b59e0ce1a82a69237dcf8977d3e72a2965d/LogicalTypes.md?plain=1#L766-L769) + #[test] + fn basic_backward_compatible_list_1() -> crate::errors::Result<()> { + test_expected_type( + " + message schema { + optional group my_list (LIST) { + repeated int32 element; + } + } + ", + Fields::from(vec![ + // Rule 1: List (nullable list, non-null elements) + Field::new( + "my_list", + DataType::List(Arc::new(Field::new("element", DataType::Int32, false))), + true, + ), + ]), + ) + } + + /// Taken from the example in [Parquet Format - Nested Types - Lists - Backward-compatibility rules](https://github.com/apache/parquet-format/blob/9fd57b59e0ce1a82a69237dcf8977d3e72a2965d/LogicalTypes.md?plain=1#L771-L777) + #[test] + fn basic_backward_compatible_list_2() -> crate::errors::Result<()> { + test_expected_type( + " + message schema { + optional group my_list (LIST) { + repeated group element { + required binary str (STRING); + required int32 num; + } + } + } + ", + Fields::from(vec![ + // Rule 2: List> (nullable list, non-null elements) + Field::new( + "my_list", + DataType::List(Arc::new(Field::new( + "element", + DataType::Struct(Fields::from(vec![ + Field::new("str", DataType::Utf8, false), + Field::new("num", DataType::Int32, false), + ])), + false, + ))), + true, + ), + ]), + ) + } + + /// Taken from the example in [Parquet Format - Nested Types - Lists - Backward-compatibility rules](https://github.com/apache/parquet-format/blob/9fd57b59e0ce1a82a69237dcf8977d3e72a2965d/LogicalTypes.md?plain=1#L779-L784) + #[test] + fn basic_backward_compatible_list_3() -> crate::errors::Result<()> { + test_expected_type( + " + message schema { + optional group my_list (LIST) { + repeated group array (LIST) { + repeated int32 array; + } + } + } + ", + Fields::from(vec![ + // Rule 3: List> (nullable outer list, non-null elements) + Field::new( + "my_list", + DataType::List(Arc::new(Field::new( + "array", + DataType::List(Arc::new(Field::new("array", DataType::Int32, false))), + false, + ))), + true, + ), + ]), + ) + } + + /// Taken from the example in [Parquet Format - Nested Types - Lists - Backward-compatibility rules](https://github.com/apache/parquet-format/blob/9fd57b59e0ce1a82a69237dcf8977d3e72a2965d/LogicalTypes.md?plain=1#L786-L791) + #[test] + fn basic_backward_compatible_list_4_1() -> crate::errors::Result<()> { + test_expected_type( + " + message schema { + optional group my_list (LIST) { + repeated group array { + required binary str (STRING); + } + } + } + ", + Fields::from(vec![ + // Rule 4: List> (nullable list, non-null elements) + Field::new( + "my_list", + DataType::List(Arc::new(Field::new( + "array", + DataType::Struct(Fields::from(vec![Field::new( + "str", + DataType::Utf8, + false, + )])), + false, + ))), + true, + ), + ]), + ) + } + + /// Taken from the example in [Parquet Format - Nested Types - Lists - Backward-compatibility rules](https://github.com/apache/parquet-format/blob/9fd57b59e0ce1a82a69237dcf8977d3e72a2965d/LogicalTypes.md?plain=1#L793-L798) + #[test] + fn basic_backward_compatible_list_4_2() -> crate::errors::Result<()> { + test_expected_type( + " + message schema { + optional group my_list (LIST) { + repeated group my_list_tuple { + required binary str (STRING); + } + } + } + ", + Fields::from(vec![ + // Rule 4: List> (nullable list, non-null elements) + Field::new( + "my_list", + DataType::List(Arc::new(Field::new( + "my_list_tuple", + DataType::Struct(Fields::from(vec![Field::new( + "str", + DataType::Utf8, + false, + )])), + false, + ))), + true, + ), + ]), + ) + } + + /// Taken from the example in [Parquet Format - Nested Types - Lists - Backward-compatibility rules](https://github.com/apache/parquet-format/blob/9fd57b59e0ce1a82a69237dcf8977d3e72a2965d/LogicalTypes.md?plain=1#L800-L805) + #[test] + fn basic_backward_compatible_list_5() -> crate::errors::Result<()> { + test_expected_type( + " + message schema { + optional group my_list (LIST) { + repeated group element { + optional binary str (STRING); + } + } + } + ", + Fields::from(vec![ + // Rule 5: List (nullable list, nullable elements) + Field::new( + "my_list", + DataType::List(Arc::new(Field::new("str", DataType::Utf8, true))), + true, + ), + ]), + ) + } + + #[test] + fn basic_backward_compatible_map_1() -> crate::errors::Result<()> { + test_expected_type( + " + message schema { + optional group my_map (MAP) { + repeated group map { + required binary str (STRING); + required int32 num; + } + } + } + ", + Fields::from(vec![ + // Map (nullable map, non-null values) + Field::new( + "my_map", + DataType::Map( + Arc::new(Field::new( + "map", + DataType::Struct(Fields::from(vec![ + Field::new("str", DataType::Utf8, false), + Field::new("num", DataType::Int32, false), + ])), + false, + )), + false, + ), + true, + ), + ]), + ) + } + + #[test] + fn basic_backward_compatible_map_2() -> crate::errors::Result<()> { + test_expected_type( + " + message schema { + optional group my_map (MAP_KEY_VALUE) { + repeated group map { + required binary key (STRING); + optional int32 value; + } + } + } + ", + Fields::from(vec![ + // Map (nullable map, nullable values) + Field::new( + "my_map", + DataType::Map( + Arc::new(Field::new( + "map", + DataType::Struct(Fields::from(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Int32, true), + ])), + false, + )), + false, + ), + true, + ), + ]), + ) + } + + #[test] + fn convert_schema_with_nested_list_repeated_primitive() -> crate::errors::Result<()> { + test_roundtrip( + " + message schema { + optional group f1 (LIST) { + repeated group element { + repeated int32 element; + } + } + } + ", + ) + } + + #[test] + fn convert_schema_with_repeated_primitive_keep_field_id() -> crate::errors::Result<()> { + let message_type = " + message schema { + repeated BYTE_ARRAY col_1 = 1; + } + "; + + let parsed_input_schema = Arc::new(parse_message_type(message_type)?); + let schema = SchemaDescriptor::new(parsed_input_schema); + + let converted = convert_schema(&schema, ProjectionMask::all(), None)?.unwrap(); + + let DataType::Struct(schema_fields) = &converted.arrow_type else { + panic!("Expected struct from convert_schema"); + }; + + assert_eq!(schema_fields.len(), 1); + + let expected_schema = DataType::Struct(Fields::from(vec![Arc::new( + arrow_schema::Field::new( + "col_1", + DataType::List(Arc::new( + // No metadata on inner field + arrow_schema::Field::new("col_1", DataType::Binary, false), + )), + false, + ) + // add the field id to the outer list + .with_field_id(1), + )])); + + assert_eq!(converted.arrow_type, expected_schema); + + Ok(()) + } + + #[test] + fn convert_schema_with_repeated_primitive_should_use_inferred_schema() + -> crate::errors::Result<()> { + let message_type = " + message schema { + repeated BYTE_ARRAY col_1 = 1; + } + "; + + let parsed_input_schema = Arc::new(parse_message_type(message_type)?); + let schema = SchemaDescriptor::new(parsed_input_schema); + + let converted = convert_schema(&schema, ProjectionMask::all(), None)?.unwrap(); + + let DataType::Struct(schema_fields) = &converted.arrow_type else { + panic!("Expected struct from convert_schema"); + }; + + assert_eq!(schema_fields.len(), 1); + + let expected_schema = DataType::Struct(Fields::from(vec![Arc::new( + arrow_schema::Field::new( + "col_1", + DataType::List(Arc::new(arrow_schema::Field::new( + "col_1", + DataType::Binary, + false, + ))), + false, + ) + .with_metadata(schema_fields[0].metadata().clone()), + )])); + + assert_eq!(converted.arrow_type, expected_schema); + + let utf8_instead_of_binary = Fields::from(vec![Arc::new( + arrow_schema::Field::new( + "col_1", + DataType::List(Arc::new(arrow_schema::Field::new( + "col_1", + DataType::Utf8, + false, + ))), + false, + ) + .with_metadata(schema_fields[0].metadata().clone()), + )]); + + // Should be able to convert the same thing + let converted_again = convert_schema( + &schema, + ProjectionMask::all(), + Some(&utf8_instead_of_binary), + )? + .unwrap(); + + // Assert that we changed to Utf8 + assert_eq!( + converted_again.arrow_type, + DataType::Struct(utf8_instead_of_binary) + ); + + Ok(()) + } + + #[test] + fn convert_schema_with_repeated_primitive_should_use_inferred_schema_for_list_as_well() + -> crate::errors::Result<()> { + let message_type = " + message schema { + repeated BYTE_ARRAY col_1 = 1; + } + "; + + let parsed_input_schema = Arc::new(parse_message_type(message_type)?); + let schema = SchemaDescriptor::new(parsed_input_schema); + + let converted = convert_schema(&schema, ProjectionMask::all(), None)?.unwrap(); + + let DataType::Struct(schema_fields) = &converted.arrow_type else { + panic!("Expected struct from convert_schema"); + }; + + assert_eq!(schema_fields.len(), 1); + + let expected_schema = DataType::Struct(Fields::from(vec![Arc::new( + arrow_schema::Field::new( + "col_1", + DataType::List(Arc::new(arrow_schema::Field::new( + "col_1", + DataType::Binary, + false, + ))), + false, + ) + .with_metadata(schema_fields[0].metadata().clone()), + )])); + + assert_eq!(converted.arrow_type, expected_schema); + + let utf8_instead_of_binary = Fields::from(vec![Arc::new( + arrow_schema::Field::new( + "col_1", + // Inferring as LargeList instead of List + DataType::LargeList(Arc::new(arrow_schema::Field::new( + "col_1", + DataType::Utf8, + false, + ))), + false, + ) + .with_metadata(schema_fields[0].metadata().clone()), + )]); + + // Should be able to convert the same thing + let converted_again = convert_schema( + &schema, + ProjectionMask::all(), + Some(&utf8_instead_of_binary), + )? + .unwrap(); + + // Assert that we changed to Utf8 + assert_eq!( + converted_again.arrow_type, + DataType::Struct(utf8_instead_of_binary) + ); + + Ok(()) + } + + #[test] + fn convert_schema_with_repeated_struct_and_inferred_schema() -> crate::errors::Result<()> { + test_roundtrip( + " + message schema { + repeated group my_col_1 = 1 { + optional binary my_col_2 = 2; + optional binary my_col_3 = 3; + optional group my_col_4 = 4 { + optional int64 my_col_5 = 5; + optional int32 my_col_6 = 6; + } + } + } + ", + ) + } + + #[test] + fn convert_schema_with_repeated_struct_and_inferred_schema_and_field_id() + -> crate::errors::Result<()> { + let message_type = " + message schema { + repeated group my_col_1 = 1 { + optional binary my_col_2 = 2; + optional binary my_col_3 = 3; + optional group my_col_4 = 4 { + optional int64 my_col_5 = 5; + optional int32 my_col_6 = 6; + } + } + } + "; + + let parsed_input_schema = Arc::new(parse_message_type(message_type)?); + let schema = SchemaDescriptor::new(parsed_input_schema); + + let converted = convert_schema(&schema, ProjectionMask::all(), None)?.unwrap(); + + let DataType::Struct(schema_fields) = &converted.arrow_type else { + panic!("Expected struct from convert_schema"); + }; + + assert_eq!(schema_fields.len(), 1); + + // Should be able to convert the same thing + let converted_again = + convert_schema(&schema, ProjectionMask::all(), Some(schema_fields))?.unwrap(); + + // Assert that we changed to Utf8 + assert_eq!(converted_again.arrow_type, converted.arrow_type); + + Ok(()) + } + + #[test] + fn convert_schema_with_nested_repeated_struct_and_primitives() -> crate::errors::Result<()> { + let message_type = " +message schema { + repeated group my_col_1 = 1 { + optional binary my_col_2 = 2; + repeated BYTE_ARRAY my_col_3 = 3; + repeated group my_col_4 = 4 { + optional int64 my_col_5 = 5; + repeated binary my_col_6 = 6; + } + } +} +"; + + let parsed_input_schema = Arc::new(parse_message_type(message_type)?); + let schema = SchemaDescriptor::new(parsed_input_schema); + + let converted = convert_schema(&schema, ProjectionMask::all(), None)?.unwrap(); + + let DataType::Struct(schema_fields) = &converted.arrow_type else { + panic!("Expected struct from convert_schema"); + }; + + assert_eq!(schema_fields.len(), 1); + + // Build expected schema + let expected_schema = DataType::Struct(Fields::from(vec![Arc::new( + arrow_schema::Field::new( + "my_col_1", + DataType::List(Arc::new(arrow_schema::Field::new( + "my_col_1", + DataType::Struct(Fields::from(vec![ + Arc::new( + arrow_schema::Field::new("my_col_2", DataType::Binary, true) + .with_field_id(2), + ), + Arc::new( + arrow_schema::Field::new( + "my_col_3", + DataType::List(Arc::new(arrow_schema::Field::new( + "my_col_3", + DataType::Binary, + false, + ))), + false, + ) + // add the field id to the outer list + .with_field_id(3), + ), + Arc::new( + arrow_schema::Field::new( + "my_col_4", + DataType::List(Arc::new(arrow_schema::Field::new( + "my_col_4", + DataType::Struct(Fields::from(vec![ + Arc::new( + arrow_schema::Field::new( + "my_col_5", + DataType::Int64, + true, + ) + // add the field id to the outer list + .with_field_id(5), + ), + Arc::new( + arrow_schema::Field::new( + "my_col_6", + DataType::List(Arc::new(arrow_schema::Field::new( + "my_col_6", + DataType::Binary, + false, + ))), + false, + ) + // add the field id to the outer list + .with_field_id(6), + ), + ])), + false, + ))), + false, + ) + // add the field id to the outer list + .with_field_id(4), + ), + ])), + false, + ))), + false, + ) + // add the field id to the outer list + .with_field_id(1), + )])); + + assert_eq!(converted.arrow_type, expected_schema); + + // Test conversion with inferred schema + let converted_again = + convert_schema(&schema, ProjectionMask::all(), Some(schema_fields))?.unwrap(); + + assert_eq!(converted_again.arrow_type, converted.arrow_type); + + // Test conversion with modified schema (change lists to either LargeList or FixedSizeList) + // as well as changing Binary to Utf8 or BinaryView + let modified_schema_fields = Fields::from(vec![Arc::new( + arrow_schema::Field::new( + "my_col_1", + DataType::LargeList(Arc::new(arrow_schema::Field::new( + "my_col_1", + DataType::Struct(Fields::from(vec![ + Arc::new( + arrow_schema::Field::new("my_col_2", DataType::LargeBinary, true) + .with_field_id(2), + ), + Arc::new( + arrow_schema::Field::new( + "my_col_3", + DataType::LargeList(Arc::new(arrow_schema::Field::new( + "my_col_3", + DataType::Utf8, + false, + ))), + false, + ) + // add the field id to the outer list + .with_field_id(3), + ), + Arc::new( + arrow_schema::Field::new( + "my_col_4", + DataType::FixedSizeList( + Arc::new(arrow_schema::Field::new( + "my_col_4", + DataType::Struct(Fields::from(vec![ + Arc::new( + arrow_schema::Field::new( + "my_col_5", + DataType::Int64, + true, + ) + .with_field_id(5), + ), + Arc::new( + arrow_schema::Field::new( + "my_col_6", + DataType::LargeList(Arc::new( + arrow_schema::Field::new( + "my_col_6", + DataType::BinaryView, + false, + ), + )), + false, + ) + // add the field id to the outer list + .with_field_id(6), + ), + ])), + false, + )), + 3, + ), + false, + ) + // add the field id to the outer list + .with_field_id(4), + ), + ])), + false, + ))), + false, + ) + // add the field id to the outer list + .with_field_id(1), + )]); + + let converted_with_modified = convert_schema( + &schema, + ProjectionMask::all(), + Some(&modified_schema_fields), + )? + .unwrap(); + + assert_eq!( + converted_with_modified.arrow_type, + DataType::Struct(modified_schema_fields) + ); + + Ok(()) + } + + /// Backwards-compatibility: LIST with nullable element type - 1 - standard + /// Taken from [Spark](https://github.com/apache/spark/blob/8ab50765cd793169091d983b50d87a391f6ac1f4/sql/core/src/test/scala/org/apache/spark/sql/parquet/ParquetSchemaSuite.scala#L452-L466) + #[test] + fn list_nullable_element_standard() -> crate::errors::Result<()> { + test_expected_type( + " + message root { + optional group f1 (LIST) { + repeated group list { + optional int32 element; + } + } + }", + Fields::from(vec![Field::new( + "f1", + DataType::List(Arc::new(Field::new("element", DataType::Int32, true))), + true, + )]), + ) + } + + /// Backwards-compatibility: LIST with nullable element type - 2 + /// Taken from [Spark](https://github.com/apache/spark/blob/8ab50765cd793169091d983b50d87a391f6ac1f4/sql/core/src/test/scala/org/apache/spark/sql/parquet/ParquetSchemaSuite.scala#L468-L482) + #[test] + fn list_nullable_element_nested() -> crate::errors::Result<()> { + test_expected_type( + " + message root { + optional group f1 (LIST) { + repeated group element { + optional int32 num; + } + } + }", + Fields::from(vec![Field::new( + "f1", + DataType::List(Arc::new(Field::new("num", DataType::Int32, true))), + true, + )]), + ) + } + + /// Backwards-compatibility: LIST with non-nullable element type - 1 - standard + /// Taken from [Spark](https://github.com/apache/spark/blob/8ab50765cd793169091d983b50d87a391f6ac1f4/sql/core/src/test/scala/org/apache/spark/sql/parquet/ParquetSchemaSuite.scala#L484-L495) + #[test] + fn list_required_element_standard() -> crate::errors::Result<()> { + test_expected_type( + " + message root { + optional group f1 (LIST) { + repeated group list { + required int32 element; + } + } + }", + Fields::from(vec![Field::new( + "f1", + DataType::List(Arc::new(Field::new("element", DataType::Int32, false))), + true, + )]), + ) + } + + /// Backwards-compatibility: LIST with non-nullable element type - 2 + /// Taken from [Spark](https://github.com/apache/spark/blob/8ab50765cd793169091d983b50d87a391f6ac1f4/sql/core/src/test/scala/org/apache/spark/sql/parquet/ParquetSchemaSuite.scala#L497-L508) + #[test] + fn list_required_element_nested() -> crate::errors::Result<()> { + test_expected_type( + " + message root { + optional group f1 (LIST) { + repeated group element { + required int32 num; + } + } + }", + Fields::from(vec![Field::new( + "f1", + DataType::List(Arc::new(Field::new("num", DataType::Int32, false))), + true, + )]), + ) + } + + /// Backwards-compatibility: LIST with non-nullable element type - 3 + /// Taken from [Spark](https://github.com/apache/spark/blob/8ab50765cd793169091d983b50d87a391f6ac1f4/sql/core/src/test/scala/org/apache/spark/sql/parquet/ParquetSchemaSuite.scala#L510-L519) + #[test] + fn list_required_element_primitive() -> crate::errors::Result<()> { + test_expected_type( + " + message root { + optional group f1 (LIST) { + repeated int32 element; + } + }", + Fields::from(vec![Field::new( + "f1", + DataType::List(Arc::new(Field::new("element", DataType::Int32, false))), + true, + )]), + ) + } + + /// Backwards-compatibility: LIST with non-nullable element type - 4 + /// Taken from [Spark](https://github.com/apache/spark/blob/8ab50765cd793169091d983b50d87a391f6ac1f4/sql/core/src/test/scala/org/apache/spark/sql/parquet/ParquetSchemaSuite.scala#L521-L540) + #[test] + fn list_required_element_struct() -> crate::errors::Result<()> { + test_expected_type( + " + message root { + optional group f1 (LIST) { + repeated group element { + required binary str (UTF8); + required int32 num; + } + } + }", + Fields::from(vec![Field::new( + "f1", + DataType::List(Arc::new(Field::new( + "element", + DataType::Struct(Fields::from(vec![ + Field::new("str", DataType::Utf8, false), + Field::new("num", DataType::Int32, false), + ])), + false, + ))), + true, + )]), + ) + } + + /// Backwards-compatibility: LIST with non-nullable element type - 5 - parquet-avro style + /// Taken from [Spark](https://github.com/apache/spark/blob/8ab50765cd793169091d983b50d87a391f6ac1f4/sql/core/src/test/scala/org/apache/spark/sql/parquet/ParquetSchemaSuite.scala#L542-L559) + #[test] + fn list_required_element_avro_style() -> crate::errors::Result<()> { + test_expected_type( + " + message root { + optional group f1 (LIST) { + repeated group array { + required binary str (UTF8); + } + } + }", + Fields::from(vec![Field::new( + "f1", + DataType::List(Arc::new(Field::new( + "array", + DataType::Struct(Fields::from(vec![Field::new("str", DataType::Utf8, false)])), + false, + ))), + true, + )]), + ) + } + + /// Backwards-compatibility: LIST with non-nullable element type - 6 - parquet-thrift style + /// Taken from [Spark](https://github.com/apache/spark/blob/8ab50765cd793169091d983b50d87a391f6ac1f4/sql/core/src/test/scala/org/apache/spark/sql/parquet/ParquetSchemaSuite.scala#L561-L578) + #[test] + fn list_required_element_thrift_style() -> crate::errors::Result<()> { + test_expected_type( + " + message root { + optional group f1 (LIST) { + repeated group f1_tuple { + required binary str (UTF8); + } + } + }", + Fields::from(vec![Field::new( + "f1", + DataType::List(Arc::new(Field::new( + "f1_tuple", + DataType::Struct(Fields::from(vec![Field::new("str", DataType::Utf8, false)])), + false, + ))), + true, + )]), + ) + } + + /// Backwards-compatibility: MAP with non-nullable value type - 1 - standard + /// Taken from [Spark](https://github.com/apache/spark/blob/8ab50765cd793169091d983b50d87a391f6ac1f4/sql/core/src/test/scala/org/apache/spark/sql/parquet/ParquetSchemaSuite.scala#L652-L667) + #[test] + fn map_required_value_standard() -> crate::errors::Result<()> { + test_expected_type( + " + message root { + optional group f1 (MAP) { + repeated group key_value { + required int32 key; + required binary value (UTF8); + } + } + }", + Fields::from(vec![Field::new_map( + "f1", + "key_value", + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Utf8, false), + false, + true, + )]), + ) + } + + /// Backwards-compatibility: MAP with non-nullable value type - 2 + /// Taken from [Spark](https://github.com/apache/spark/blob/8ab50765cd793169091d983b50d87a391f6ac1f4/sql/core/src/test/scala/org/apache/spark/sql/parquet/ParquetSchemaSuite.scala#L669-L684) + #[test] + fn map_required_value_map_key_value() -> crate::errors::Result<()> { + test_expected_type( + " + message root { + optional group f1 (MAP_KEY_VALUE) { + repeated group map { + required int32 num; + required binary str (UTF8); + } + } + }", + Fields::from(vec![Field::new_map( + "f1", + "map", + Field::new("num", DataType::Int32, false), + Field::new("str", DataType::Utf8, false), + false, + true, + )]), + ) + } + + /// Backwards-compatibility: MAP with non-nullable value type - 3 - prior to 1.4.x + /// Taken from [Spark](https://github.com/apache/spark/blob/8ab50765cd793169091d983b50d87a391f6ac1f4/sql/core/src/test/scala/org/apache/spark/sql/parquet/ParquetSchemaSuite.scala#L686-L701) + #[test] + fn map_required_value_legacy() -> crate::errors::Result<()> { + test_expected_type( + " + message root { + optional group f1 (MAP) { + repeated group map (MAP_KEY_VALUE) { + required int32 key; + required binary value (UTF8); + } + } + }", + Fields::from(vec![Field::new_map( + "f1", + "map", + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Utf8, false), + false, + true, + )]), + ) + } + + /// Backwards-compatibility: MAP with nullable value type - 1 - standard + /// Taken from [Spark](https://github.com/apache/spark/blob/8ab50765cd793169091d983b50d87a391f6ac1f4/sql/core/src/test/scala/org/apache/spark/sql/parquet/ParquetSchemaSuite.scala#L703-L718) + #[test] + fn map_optional_value_standard() -> crate::errors::Result<()> { + test_expected_type( + " + message root { + optional group f1 (MAP) { + repeated group key_value { + required int32 key; + optional binary value (UTF8); + } + } + }", + Fields::from(vec![Field::new_map( + "f1", + "key_value", + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Utf8, true), + false, + true, + )]), + ) + } + + /// Backwards-compatibility: MAP with nullable value type - 2 + /// Taken from [Spark](https://github.com/apache/spark/blob/8ab50765cd793169091d983b50d87a391f6ac1f4/sql/core/src/test/scala/org/apache/spark/sql/parquet/ParquetSchemaSuite.scala#L720-L735) + #[test] + fn map_optional_value_map_key_value() -> crate::errors::Result<()> { + test_expected_type( + " + message root { + optional group f1 (MAP_KEY_VALUE) { + repeated group map { + required int32 num; + optional binary str (UTF8); + } + } + }", + Fields::from(vec![Field::new_map( + "f1", + "map", + Field::new("num", DataType::Int32, false), + Field::new("str", DataType::Utf8, true), + false, + true, + )]), + ) + } + + /// Backwards-compatibility: MAP with nullable value type - 3 - parquet-avro style + /// Taken from [Spark](https://github.com/apache/spark/blob/8ab50765cd793169091d983b50d87a391f6ac1f4/sql/core/src/test/scala/org/apache/spark/sql/parquet/ParquetSchemaSuite.scala#L737-L752) + #[test] + fn map_optional_value_avro_style() -> crate::errors::Result<()> { + test_expected_type( + " + message root { + optional group f1 (MAP) { + repeated group map (MAP_KEY_VALUE) { + required int32 key; + optional binary value (UTF8); + } + } + }", + Fields::from(vec![Field::new_map( + "f1", + "map", + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Utf8, true), + false, + true, + )]), + ) + } +} diff --git a/vendor/parquet/src/arrow/schema/extension.rs b/vendor/parquet/src/arrow/schema/extension.rs new file mode 100644 index 0000000..4e1ec75 --- /dev/null +++ b/vendor/parquet/src/arrow/schema/extension.rs @@ -0,0 +1,231 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Arrow Extension Type Support for Parquet +//! +//! This module contains mapping code to map Parquet [`LogicalType`]s to/from +//! Arrow [`ExtensionType`]s. +//! +//! Extension types are represented using the metadata from Arrow [`Field`]s +//! with the key "ARROW:extension:name". +//! +//! Extension types are represented using the metadata from Arrow [`Field`]s +//! with the key "ARROW:extension:name". +//! +//! [`ExtensionType`]: arrow_schema::extension::ExtensionType + +use crate::basic::LogicalType; +use crate::errors::ParquetError; +use crate::schema::types::Type; +use arrow_schema::Field; + +/// Adds extension type metadata, if necessary, based on the Parquet field's +/// [`LogicalType`] +/// +/// Some Parquet logical types, such as Variant, do not map directly to an +/// Arrow DataType, and instead are represented by an Arrow ExtensionType. +/// Extension types are attached to Arrow Fields via metadata. +pub(crate) fn try_add_extension_type( + arrow_field: Field, + parquet_type: &Type, +) -> Result { + let Some(parquet_logical_type) = parquet_type.get_basic_info().logical_type_ref() else { + return Ok(arrow_field); + }; + Ok(match parquet_logical_type { + #[cfg(feature = "variant_experimental")] + LogicalType::Variant(_) => { + let mut arrow_field = arrow_field; + arrow_field.try_with_extension_type(parquet_variant_compute::VariantType)?; + arrow_field + } + #[cfg(feature = "arrow_canonical_extension_types")] + LogicalType::Uuid => { + let mut arrow_field = arrow_field; + arrow_field.try_with_extension_type(arrow_schema::extension::Uuid)?; + arrow_field + } + #[cfg(feature = "arrow_canonical_extension_types")] + LogicalType::Json => { + let mut arrow_field = arrow_field; + arrow_field.try_with_extension_type(arrow_schema::extension::Json::default())?; + arrow_field + } + #[cfg(feature = "geospatial")] + LogicalType::Geometry(geometry) => { + // Per Parquet spec: omitted CRS defaults to OGC:CRS84, srid:0 means unset CRS + let crs = match geometry.crs.as_deref() { + None => Some("OGC:CRS84"), + Some("srid:0") => None, + Some(crs) => Some(crs), + }; + let md = parquet_geospatial::WkbMetadata::new(crs, None); + let mut arrow_field = arrow_field; + arrow_field.try_with_extension_type(parquet_geospatial::WkbType::new(Some(md)))?; + arrow_field + } + #[cfg(feature = "geospatial")] + LogicalType::Geography(geography) => { + let algorithm = geography + .algorithm() + .map(|a| a.try_as_edges()) + .transpose()?; + // Per Parquet spec: omitted CRS defaults to OGC:CRS84, srid:0 means unset CRS + let crs = match geography.crs.as_deref() { + None => Some("OGC:CRS84"), + Some("srid:0") => None, + Some(crs) => Some(crs), + }; + let md = parquet_geospatial::WkbMetadata::new(crs, algorithm); + let mut arrow_field = arrow_field; + arrow_field.try_with_extension_type(parquet_geospatial::WkbType::new(Some(md)))?; + arrow_field + } + _ => arrow_field, + }) +} + +/// Returns true if [`try_add_extension_type`] would add an extension type +/// to the specified Parquet field. +/// +/// This is used to preallocate the metadata hashmap size +pub(crate) fn has_extension_type(parquet_type: &Type) -> bool { + let Some(parquet_logical_type) = parquet_type.get_basic_info().logical_type_ref() else { + return false; + }; + match parquet_logical_type { + #[cfg(feature = "variant_experimental")] + LogicalType::Variant { .. } => true, + #[cfg(feature = "arrow_canonical_extension_types")] + LogicalType::Uuid => true, + #[cfg(feature = "arrow_canonical_extension_types")] + LogicalType::Json => true, + #[cfg(feature = "geospatial")] + LogicalType::Geometry { .. } => true, + #[cfg(feature = "geospatial")] + LogicalType::Geography { .. } => true, + _ => false, + } +} + +/// Return the Parquet logical type to use for the specified Arrow Struct field, if any. +#[cfg(feature = "variant_experimental")] +pub(crate) fn logical_type_for_struct(field: &Field) -> Option { + use parquet_variant_compute::VariantType; + if field.has_valid_extension_type::() { + Some(LogicalType::variant(None)) + } else { + None + } +} + +#[cfg(not(feature = "variant_experimental"))] +pub(crate) fn logical_type_for_struct(_field: &Field) -> Option { + None +} + +/// Return the Parquet logical type to use for the specified Arrow fixed size binary field, if any. +#[cfg(feature = "arrow_canonical_extension_types")] +pub(crate) fn logical_type_for_fixed_size_binary(field: &Field) -> Option { + use arrow_schema::extension::Uuid; + // If set, map arrow uuid extension type to parquet uuid logical type. + field + .has_valid_extension_type::() + .then_some(LogicalType::Uuid) +} + +#[cfg(not(feature = "arrow_canonical_extension_types"))] +pub(crate) fn logical_type_for_fixed_size_binary(_field: &Field) -> Option { + None +} + +/// Return the Parquet logical type to use for the specified Arrow string field (Utf8, LargeUtf8) if any +#[cfg(feature = "arrow_canonical_extension_types")] +pub(crate) fn logical_type_for_string(field: &Field) -> Option { + use arrow_schema::extension::Json; + // Use the Json logical type if the canonical Json + // extension type is set on this field. + Some(if field.has_valid_extension_type::() { + LogicalType::Json + } else { + LogicalType::String + }) +} + +#[cfg(not(feature = "arrow_canonical_extension_types"))] +pub(crate) fn logical_type_for_string(_field: &Field) -> Option { + Some(LogicalType::String) +} + +#[cfg(feature = "geospatial")] +pub(crate) fn logical_type_for_binary(field: &Field) -> Option { + use arrow_schema::extension::ExtensionType; + use parquet_geospatial::WkbType; + use parquet_geospatial::WkbTypeHint; + + match field.extension_type_name() { + Some(n) if n == WkbType::NAME => match field.try_extension_type::() { + Ok(wkb_type) => { + // Convert Arrow CRS to Parquet CRS: + // - None → "srid:0" (unset CRS in Parquet) + // - lon/lat CRS (OGC:CRS84, EPSG:4326) → None (default in Parquet) + // - Other CRS → JSON string + let crs = match &wkb_type.metadata().crs { + None => Some("srid:0".to_string()), + Some(_) if wkb_type.metadata().crs_is_lon_lat() => None, + // For string values, use the raw string; for objects, use JSON representation + Some(c) => Some( + c.as_str() + .map(|s| s.to_string()) + .unwrap_or_else(|| c.to_string()), + ), + }; + // Convert Arrow edges to Parquet algorithm: + // - Spherical → None (default for Geography) + // - Other algorithms → Some(algorithm) + let algorithm = wkb_type.metadata().algorithm.and_then(|a| { + use parquet_geospatial::WkbEdges; + match a { + WkbEdges::Spherical => None, // spherical is the default + _ => Some(a.into()), + } + }); + match wkb_type.metadata().type_hint() { + WkbTypeHint::Geometry => Some(LogicalType::geometry(crs)), + WkbTypeHint::Geography => Some(LogicalType::geography(crs, algorithm)), + } + } + Err(_e) => None, + }, + _ => None, + } +} + +#[cfg(not(feature = "geospatial"))] +pub(crate) fn logical_type_for_binary(_field: &Field) -> Option { + None +} + +#[cfg(feature = "geospatial")] +pub(crate) fn logical_type_for_binary_view(field: &Field) -> Option { + logical_type_for_binary(field) +} + +#[cfg(not(feature = "geospatial"))] +pub(crate) fn logical_type_for_binary_view(_field: &Field) -> Option { + None +} diff --git a/vendor/parquet/src/arrow/schema/mod.rs b/vendor/parquet/src/arrow/schema/mod.rs new file mode 100644 index 0000000..7878e7c --- /dev/null +++ b/vendor/parquet/src/arrow/schema/mod.rs @@ -0,0 +1,2399 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Converting Parquet schema <--> Arrow schema: [`ArrowSchemaConverter`] and [parquet_to_arrow_schema] + +use base64::Engine; +use base64::prelude::BASE64_STANDARD; +use std::collections::HashMap; +use std::sync::Arc; + +use arrow_ipc::writer; +use arrow_schema::{DataType, Field, FieldRef, Fields, Schema, TimeUnit}; + +use crate::basic::{ + ConvertedType, LogicalType, Repetition, TimeUnit as ParquetTimeUnit, Type as PhysicalType, +}; +use crate::errors::{ParquetError, Result}; +use crate::file::{metadata::KeyValue, properties::WriterProperties}; +use crate::schema::types::{ColumnDescriptor, SchemaDescriptor, Type}; + +mod complex; +mod extension; +mod primitive; +pub mod virtual_type; + +use super::PARQUET_FIELD_ID_META_KEY; +use crate::arrow::ProjectionMask; +use crate::arrow::schema::extension::{ + has_extension_type, logical_type_for_binary, logical_type_for_binary_view, + logical_type_for_fixed_size_binary, logical_type_for_string, logical_type_for_struct, + try_add_extension_type, +}; +pub(crate) use complex::{ParquetField, ParquetFieldType, VirtualColumnType}; + +/// Convert Parquet schema to Arrow schema including optional metadata +/// +/// Attempts to decode any existing Arrow schema metadata, falling back +/// to converting the Parquet schema column-wise +pub fn parquet_to_arrow_schema( + parquet_schema: &SchemaDescriptor, + key_value_metadata: Option<&Vec>, +) -> Result { + parquet_to_arrow_schema_by_columns(parquet_schema, ProjectionMask::all(), key_value_metadata) +} + +/// Convert parquet schema to arrow schema including optional metadata, +/// only preserving some leaf columns. +pub fn parquet_to_arrow_schema_by_columns( + parquet_schema: &SchemaDescriptor, + mask: ProjectionMask, + key_value_metadata: Option<&Vec>, +) -> Result { + Ok(parquet_to_arrow_schema_and_fields(parquet_schema, mask, key_value_metadata, &[])?.0) +} + +/// Determines the Arrow Schema from a Parquet schema +/// +/// Looks for an Arrow schema metadata "hint" (see +/// [`parquet_to_arrow_field_levels`]), and uses it if present to ensure +/// lossless round trips. +pub(crate) fn parquet_to_arrow_schema_and_fields( + parquet_schema: &SchemaDescriptor, + mask: ProjectionMask, + key_value_metadata: Option<&Vec>, + virtual_columns: &[FieldRef], +) -> Result<(Schema, Option)> { + let mut metadata = parse_key_value_metadata(key_value_metadata).unwrap_or_default(); + let maybe_schema = metadata + .remove(super::ARROW_SCHEMA_META_KEY) + .map(|value| get_arrow_schema_from_metadata(&value)) + .transpose()?; + + // Add the Arrow metadata to the Parquet metadata skipping keys that collide + if let Some(arrow_schema) = &maybe_schema { + arrow_schema.metadata().iter().for_each(|(k, v)| { + metadata.entry(k.clone()).or_insert_with(|| v.clone()); + }); + } + + let hint = maybe_schema.as_ref().map(|s| s.fields()); + let field_levels = + parquet_to_arrow_field_levels_with_virtual(parquet_schema, mask, hint, virtual_columns)?; + let schema = Schema::new_with_metadata(field_levels.fields, metadata); + Ok((schema, field_levels.levels)) +} + +/// Schema information necessary to decode a parquet file as arrow [`Fields`] +/// +/// In particular this stores the dremel-level information necessary to correctly +/// interpret the encoded definition and repetition levels +/// +/// Note: this is an opaque container intended to be used with lower-level APIs +/// within this crate +#[derive(Debug, Clone)] +pub struct FieldLevels { + pub(crate) fields: Fields, + pub(crate) levels: Option, +} + +/// Convert a parquet [`SchemaDescriptor`] to [`FieldLevels`] +/// +/// Columns not included within [`ProjectionMask`] will be ignored. +/// +/// The optional `hint` parameter is the desired Arrow schema. See the +/// [`arrow`] module documentation for more information. +/// +/// [`arrow`]: crate::arrow +/// +/// # Notes: +/// Where a field type in `hint` is compatible with the corresponding parquet type in `schema`, it +/// will be used, otherwise the default arrow type for the given parquet column type will be used. +/// +/// This is to accommodate arrow types that cannot be round-tripped through parquet natively. +/// Depending on the parquet writer, this can lead to a mismatch between a file's parquet schema +/// and its embedded arrow schema. The parquet `schema` must be treated as authoritative in such +/// an event. See [#1663](https://github.com/apache/arrow-rs/issues/1663) for more information +/// +/// Note: this is a low-level API, most users will want to make use of the higher-level +/// [`parquet_to_arrow_schema`] for decoding metadata from a parquet file. +pub fn parquet_to_arrow_field_levels( + schema: &SchemaDescriptor, + mask: ProjectionMask, + hint: Option<&Fields>, +) -> Result { + parquet_to_arrow_field_levels_with_virtual(schema, mask, hint, &[]) +} + +/// Convert a parquet [`SchemaDescriptor`] to [`FieldLevels`] with support for virtual columns +/// +/// Columns not included within [`ProjectionMask`] will be ignored. +/// +/// The optional `hint` parameter is the desired Arrow schema. See the +/// [`arrow`] module documentation for more information. +/// +/// [`arrow`]: crate::arrow +/// +/// # Arguments +/// * `schema` - The Parquet schema descriptor +/// * `mask` - Projection mask to select which columns to include +/// * `hint` - Optional hint for Arrow field types to use instead of defaults +/// * `virtual_columns` - Virtual columns to append to the schema (e.g., row numbers) +/// +/// # Notes: +/// Where a field type in `hint` is compatible with the corresponding parquet type in `schema`, it +/// will be used, otherwise the default arrow type for the given parquet column type will be used. +/// +/// Virtual columns are columns that don't exist in the Parquet file but are generated during reading. +/// They must have extension type names starting with "arrow.virtual.". +/// +/// This is to accommodate arrow types that cannot be round-tripped through parquet natively. +/// Depending on the parquet writer, this can lead to a mismatch between a file's parquet schema +/// and its embedded arrow schema. The parquet `schema` must be treated as authoritative in such +/// an event. See [#1663](https://github.com/apache/arrow-rs/issues/1663) for more information +/// +/// Note: this is a low-level API, most users will want to make use of the higher-level +/// [`parquet_to_arrow_schema`] for decoding metadata from a parquet file. +pub fn parquet_to_arrow_field_levels_with_virtual( + schema: &SchemaDescriptor, + mask: ProjectionMask, + hint: Option<&Fields>, + virtual_columns: &[FieldRef], +) -> Result { + // Validate that all fields are virtual columns + for field in virtual_columns { + if !virtual_type::is_virtual_column(field) { + return Err(ParquetError::General(format!( + "Field '{}' is not a virtual column. Virtual columns must have extension type names starting with 'arrow.virtual.'", + field.name() + ))); + } + } + + // Convert the regular schema first + let mut parquet_field = match complex::convert_schema(schema, mask, hint)? { + Some(field) => field, + None if virtual_columns.is_empty() => { + return Ok(FieldLevels { + fields: Fields::empty(), + levels: None, + }); + } + None => { + // No regular fields, but we have virtual columns - create empty root struct + ParquetField { + rep_level: 0, + def_level: 0, + nullable: false, + arrow_type: DataType::Struct(Fields::empty()), + field_type: ParquetFieldType::Group { + children: Vec::new(), + }, + } + } + }; + + // Append virtual columns if any + if !virtual_columns.is_empty() { + match &mut parquet_field.field_type { + ParquetFieldType::Group { children } => { + // Get the mutable fields from the struct type + let DataType::Struct(ref mut fields) = parquet_field.arrow_type else { + unreachable!("Root field must be a struct"); + }; + + // Convert to mutable Vec to append + let mut fields_vec: Vec = fields.iter().cloned().collect(); + + // Append each virtual column + for virtual_column in virtual_columns { + // Virtual columns can only be added at the root level + assert_eq!( + parquet_field.rep_level, 0, + "Virtual columns can only be added at rep level 0" + ); + assert_eq!( + parquet_field.def_level, 0, + "Virtual columns can only be added at def level 0" + ); + + fields_vec.push(virtual_column.clone()); + let virtual_parquet_field = complex::convert_virtual_field( + virtual_column, + parquet_field.rep_level, + parquet_field.def_level, + )?; + children.push(virtual_parquet_field); + } + + // Update the fields + parquet_field.arrow_type = DataType::Struct(Fields::from(fields_vec)); + } + _ => unreachable!("Root field must be a group"), + } + } + + match &parquet_field.arrow_type { + DataType::Struct(fields) => Ok(FieldLevels { + fields: fields.clone(), + levels: Some(parquet_field), + }), + _ => unreachable!(), + } +} + +/// Try to convert Arrow schema metadata into a schema +fn get_arrow_schema_from_metadata(encoded_meta: &str) -> Result { + let decoded = BASE64_STANDARD.decode(encoded_meta); + match decoded { + Ok(bytes) => { + let slice = if bytes.len() > 8 && bytes[0..4] == [255u8; 4] { + &bytes[8..] + } else { + bytes.as_slice() + }; + match arrow_ipc::root_as_message(slice) { + Ok(message) => message + .header_as_schema() + .map(arrow_ipc::convert::fb_to_schema) + .ok_or_else(|| arrow_err!("the message is not Arrow Schema")), + Err(err) => { + // The flatbuffers implementation returns an error on verification error. + Err(arrow_err!( + "Unable to get root as message stored in {}: {:?}", + super::ARROW_SCHEMA_META_KEY, + err + )) + } + } + } + Err(err) => { + // The C++ implementation returns an error if the schema can't be parsed. + Err(arrow_err!( + "Unable to decode the encoded schema stored in {}, {:?}", + super::ARROW_SCHEMA_META_KEY, + err + )) + } + } +} + +/// Encodes the Arrow schema into the IPC format, and base64 encodes it +pub fn encode_arrow_schema(schema: &Schema) -> String { + let options = writer::IpcWriteOptions::default(); + let mut dictionary_tracker = writer::DictionaryTracker::new(true); + let data_gen = writer::IpcDataGenerator::default(); + let mut serialized_schema = + data_gen.schema_to_bytes_with_dictionary_tracker(schema, &mut dictionary_tracker, &options); + + // manually prepending the length to the schema as arrow uses the legacy IPC format + // TODO: change after addressing ARROW-9777 + let schema_len = serialized_schema.ipc_message.len(); + let mut len_prefix_schema = Vec::with_capacity(schema_len + 8); + len_prefix_schema.append(&mut vec![255u8, 255, 255, 255]); + len_prefix_schema.append((schema_len as u32).to_le_bytes().to_vec().as_mut()); + len_prefix_schema.append(&mut serialized_schema.ipc_message); + + BASE64_STANDARD.encode(&len_prefix_schema) +} + +fn flatten_ree_field(field: &Field) -> Field { + match field.data_type() { + DataType::RunEndEncoded(_, value_field) => field + .clone() + .with_data_type(value_field.data_type().clone()), + _ => field.clone(), + } +} + +/// Mutates writer metadata by storing the encoded Arrow schema hint in +/// [`ARROW_SCHEMA_META_KEY`]. +/// +/// If there is an existing Arrow schema metadata, it is replaced. +/// +/// [`ARROW_SCHEMA_META_KEY`]: crate::arrow::ARROW_SCHEMA_META_KEY +pub fn add_encoded_arrow_schema_to_metadata(schema: &Schema, props: &mut WriterProperties) { + let has_ree = schema + .fields() + .iter() + .any(|f| matches!(f.data_type(), DataType::RunEndEncoded(_, _))); + let flat_schema; + let schema = if has_ree { + let flat_fields: Vec = schema + .fields() + .iter() + .map(|f| flatten_ree_field(f)) + .collect(); + flat_schema = Schema::new_with_metadata(flat_fields, schema.metadata().clone()); + &flat_schema + } else { + schema + }; + let encoded = encode_arrow_schema(schema); + + let schema_kv = KeyValue { + key: super::ARROW_SCHEMA_META_KEY.to_string(), + value: Some(encoded), + }; + + let meta = props + .key_value_metadata + .get_or_insert_with(Default::default); + + // check if ARROW:schema exists, and overwrite it + let schema_meta = meta + .iter() + .enumerate() + .find(|(_, kv)| kv.key.as_str() == super::ARROW_SCHEMA_META_KEY); + match schema_meta { + Some((i, _)) => { + meta.remove(i); + meta.push(schema_kv); + } + None => { + meta.push(schema_kv); + } + } +} + +/// Converter for Arrow schema to Parquet schema +/// +/// See the documentation on the [`arrow`] module for background +/// information on how Arrow schema is represented in Parquet. +/// +/// [`arrow`]: crate::arrow +/// +/// # Example: +/// ``` +/// # use std::sync::Arc; +/// # use arrow_schema::{Field, Schema, DataType}; +/// # use parquet::arrow::ArrowSchemaConverter; +/// use parquet::schema::types::{SchemaDescriptor, Type}; +/// use parquet::basic; // note there are two `Type`s in the following example +/// // create an Arrow Schema +/// let arrow_schema = Schema::new(vec![ +/// Field::new("a", DataType::Int64, true), +/// Field::new("b", DataType::Date32, true), +/// ]); +/// // convert the Arrow schema to a Parquet schema +/// let parquet_schema = ArrowSchemaConverter::new() +/// .convert(&arrow_schema) +/// .unwrap(); +/// +/// let expected_parquet_schema = SchemaDescriptor::new( +/// Arc::new( +/// Type::group_type_builder("arrow_schema") +/// .with_fields(vec![ +/// Arc::new( +/// Type::primitive_type_builder("a", basic::Type::INT64) +/// .build().unwrap() +/// ), +/// Arc::new( +/// Type::primitive_type_builder("b", basic::Type::INT32) +/// .with_converted_type(basic::ConvertedType::DATE) +/// .with_logical_type(Some(basic::LogicalType::Date)) +/// .build().unwrap() +/// ), +/// ]) +/// .build().unwrap() +/// ) +/// ); +/// assert_eq!(parquet_schema, expected_parquet_schema); +/// ``` +#[derive(Debug)] +pub struct ArrowSchemaConverter<'a> { + /// Name of the root schema in Parquet + schema_root: &'a str, + /// Should we coerce Arrow types to compatible Parquet types? + /// + /// See docs on [Self::with_coerce_types]` + coerce_types: bool, +} + +impl Default for ArrowSchemaConverter<'_> { + fn default() -> Self { + Self::new() + } +} + +impl<'a> ArrowSchemaConverter<'a> { + /// Create a new converter + pub fn new() -> Self { + Self { + schema_root: "arrow_schema", + coerce_types: false, + } + } + + /// Should Arrow types be coerced into Parquet native types (default `false`). + /// + /// Setting this option to `true` will result in Parquet files that can be + /// read by more readers, but may lose precision for Arrow types such as + /// [`DataType::Date64`] which have no direct [corresponding Parquet type]. + /// + /// By default, this converter does not coerce to native Parquet types. Enabling type + /// coercion allows for meaningful representations that do not require + /// downstream readers to consider the embedded Arrow schema, and can allow + /// for greater compatibility with other Parquet implementations. However, + /// type coercion also prevents data from being losslessly round-tripped. + /// + /// # Discussion + /// + /// Some Arrow types such as `Date64`, `Timestamp` and `Interval` have no + /// corresponding Parquet logical type. Thus, they can not be losslessly + /// round-tripped when stored using the appropriate Parquet logical type. + /// For example, some Date64 values may be truncated when stored with + /// parquet's native 32 bit date type. + /// + /// For [`List`] and [`Map`] types, some Parquet readers expect certain + /// schema elements to have specific names (earlier versions of the spec + /// were somewhat ambiguous on this point). Type coercion will use the names + /// prescribed by the Parquet specification, potentially losing naming + /// metadata from the Arrow schema. + /// + /// [`List`]: https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#lists + /// [`Map`]: https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#maps + /// [corresponding Parquet type]: https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#date + /// + pub fn with_coerce_types(mut self, coerce_types: bool) -> Self { + self.coerce_types = coerce_types; + self + } + + /// Set the root schema element name (defaults to `"arrow_schema"`). + pub fn schema_root(mut self, schema_root: &'a str) -> Self { + self.schema_root = schema_root; + self + } + + /// Convert the specified Arrow [`Schema`] to the desired Parquet [`SchemaDescriptor`] + /// + /// See example in [`ArrowSchemaConverter`] + pub fn convert(&self, schema: &Schema) -> Result { + let fields = schema + .fields() + .iter() + .map(|field| arrow_to_parquet_type(field, self.coerce_types).map(Arc::new)) + .collect::>()?; + let group = Type::group_type_builder(self.schema_root) + .with_fields(fields) + .build()?; + Ok(SchemaDescriptor::new(Arc::new(group))) + } +} + +fn parse_key_value_metadata( + key_value_metadata: Option<&Vec>, +) -> Option> { + match key_value_metadata { + Some(key_values) => { + let map: HashMap = key_values + .iter() + .filter_map(|kv| { + kv.value + .as_ref() + .map(|value| (kv.key.clone(), value.clone())) + }) + .collect(); + + if map.is_empty() { None } else { Some(map) } + } + None => None, + } +} + +/// Convert parquet column schema to arrow field. +pub fn parquet_to_arrow_field(parquet_column: &ColumnDescriptor) -> Result { + let field = complex::convert_type(&parquet_column.self_type_ptr())?; + let mut ret = Field::new(parquet_column.name(), field.arrow_type, field.nullable); + + let parquet_type = parquet_column.self_type(); + let basic_info = parquet_type.get_basic_info(); + + let mut hash_map_size = 0; + if basic_info.has_id() { + hash_map_size += 1; + } + if has_extension_type(parquet_type) { + hash_map_size += 1; + } + if hash_map_size == 0 { + return Ok(ret); + } + ret.set_metadata(HashMap::with_capacity(hash_map_size)); + if basic_info.has_id() { + ret.metadata_mut().insert( + PARQUET_FIELD_ID_META_KEY.to_string(), + basic_info.id().to_string(), + ); + } + try_add_extension_type(ret, parquet_column.self_type()) +} + +pub fn decimal_length_from_precision(precision: u8) -> usize { + // digits = floor(log_10(2^(8*n - 1) - 1)) // definition in parquet's logical types + // ceil(digits) = log10(2^(8*n - 1) - 1) + // 10^ceil(digits) = 2^(8*n - 1) - 1 + // 10^ceil(digits) + 1 = 2^(8*n - 1) + // log2(10^ceil(digits) + 1) = (8*n - 1) + // log2(10^ceil(digits) + 1) + 1 = 8*n + // (log2(10^ceil(a) + 1) + 1) / 8 = n + (((10.0_f64.powi(precision as i32) + 1.0).log2() + 1.0) / 8.0).ceil() as usize +} + +/// Convert an arrow field to a parquet `Type` +fn arrow_to_parquet_type(field: &Field, coerce_types: bool) -> Result { + const PARQUET_LIST_ELEMENT_NAME: &str = "element"; + const PARQUET_MAP_STRUCT_NAME: &str = "key_value"; + const PARQUET_KEY_FIELD_NAME: &str = "key"; + const PARQUET_VALUE_FIELD_NAME: &str = "value"; + + let name = field.name().as_str(); + let repetition = if field.is_nullable() { + Repetition::OPTIONAL + } else { + Repetition::REQUIRED + }; + let id = field_id(field); + // create type from field + match field.data_type() { + DataType::Null => Type::primitive_type_builder(name, PhysicalType::INT32) + .with_logical_type(Some(LogicalType::Unknown)) + .with_repetition(repetition) + .with_id(id) + .build(), + DataType::Boolean => Type::primitive_type_builder(name, PhysicalType::BOOLEAN) + .with_repetition(repetition) + .with_id(id) + .build(), + DataType::Int8 => Type::primitive_type_builder(name, PhysicalType::INT32) + .with_logical_type(Some(LogicalType::integer(8, true))) + .with_repetition(repetition) + .with_id(id) + .build(), + DataType::Int16 => Type::primitive_type_builder(name, PhysicalType::INT32) + .with_logical_type(Some(LogicalType::integer(16, true))) + .with_repetition(repetition) + .with_id(id) + .build(), + DataType::Int32 => Type::primitive_type_builder(name, PhysicalType::INT32) + .with_repetition(repetition) + .with_id(id) + .build(), + DataType::Int64 => Type::primitive_type_builder(name, PhysicalType::INT64) + .with_repetition(repetition) + .with_id(id) + .build(), + DataType::UInt8 => Type::primitive_type_builder(name, PhysicalType::INT32) + .with_logical_type(Some(LogicalType::integer(8, false))) + .with_repetition(repetition) + .with_id(id) + .build(), + DataType::UInt16 => Type::primitive_type_builder(name, PhysicalType::INT32) + .with_logical_type(Some(LogicalType::integer(16, false))) + .with_repetition(repetition) + .with_id(id) + .build(), + DataType::UInt32 => Type::primitive_type_builder(name, PhysicalType::INT32) + .with_logical_type(Some(LogicalType::integer(32, false))) + .with_repetition(repetition) + .with_id(id) + .build(), + DataType::UInt64 => Type::primitive_type_builder(name, PhysicalType::INT64) + .with_logical_type(Some(LogicalType::integer(64, false))) + .with_repetition(repetition) + .with_id(id) + .build(), + DataType::Float16 => Type::primitive_type_builder(name, PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_repetition(repetition) + .with_id(id) + .with_logical_type(Some(LogicalType::Float16)) + .with_length(2) + .build(), + DataType::Float32 => Type::primitive_type_builder(name, PhysicalType::FLOAT) + .with_repetition(repetition) + .with_id(id) + .build(), + DataType::Float64 => Type::primitive_type_builder(name, PhysicalType::DOUBLE) + .with_repetition(repetition) + .with_id(id) + .build(), + DataType::Timestamp(TimeUnit::Second, _) => { + // Cannot represent seconds in LogicalType + Type::primitive_type_builder(name, PhysicalType::INT64) + .with_repetition(repetition) + .with_id(id) + .build() + } + DataType::Timestamp(time_unit, tz) => { + Type::primitive_type_builder(name, PhysicalType::INT64) + .with_logical_type(Some(LogicalType::timestamp( + // If timezone set, values are normalized to UTC timezone + matches!(tz, Some(z) if !z.as_ref().is_empty()), + match time_unit { + TimeUnit::Second => unreachable!(), + TimeUnit::Millisecond => ParquetTimeUnit::MILLIS, + TimeUnit::Microsecond => ParquetTimeUnit::MICROS, + TimeUnit::Nanosecond => ParquetTimeUnit::NANOS, + }, + ))) + .with_repetition(repetition) + .with_id(id) + .build() + } + DataType::Date32 => Type::primitive_type_builder(name, PhysicalType::INT32) + .with_logical_type(Some(LogicalType::Date)) + .with_repetition(repetition) + .with_id(id) + .build(), + DataType::Date64 => { + if coerce_types { + Type::primitive_type_builder(name, PhysicalType::INT32) + .with_logical_type(Some(LogicalType::Date)) + .with_repetition(repetition) + .with_id(id) + .build() + } else { + Type::primitive_type_builder(name, PhysicalType::INT64) + .with_repetition(repetition) + .with_id(id) + .build() + } + } + DataType::Time32(TimeUnit::Second) => { + // Cannot represent seconds in LogicalType + Type::primitive_type_builder(name, PhysicalType::INT32) + .with_repetition(repetition) + .with_id(id) + .build() + } + DataType::Time32(unit) => Type::primitive_type_builder(name, PhysicalType::INT32) + .with_logical_type(Some(LogicalType::time( + field.metadata().contains_key("adjusted_to_utc"), + match unit { + TimeUnit::Millisecond => ParquetTimeUnit::MILLIS, + u => unreachable!("Invalid unit for Time32: {:?}", u), + }, + ))) + .with_repetition(repetition) + .with_id(id) + .build(), + DataType::Time64(unit) => Type::primitive_type_builder(name, PhysicalType::INT64) + .with_logical_type(Some(LogicalType::time( + field.metadata().contains_key("adjusted_to_utc"), + match unit { + TimeUnit::Microsecond => ParquetTimeUnit::MICROS, + TimeUnit::Nanosecond => ParquetTimeUnit::NANOS, + u => unreachable!("Invalid unit for Time64: {:?}", u), + }, + ))) + .with_repetition(repetition) + .with_id(id) + .build(), + DataType::Duration(_) => Type::primitive_type_builder(name, PhysicalType::INT64) + .with_repetition(repetition) + .with_id(id) + .build(), + DataType::Interval(_) => { + Type::primitive_type_builder(name, PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_converted_type(ConvertedType::INTERVAL) + .with_repetition(repetition) + .with_id(id) + .with_length(12) + .build() + } + DataType::Binary | DataType::LargeBinary => { + Type::primitive_type_builder(name, PhysicalType::BYTE_ARRAY) + .with_repetition(repetition) + .with_id(id) + .with_logical_type(logical_type_for_binary(field)) + .build() + } + DataType::FixedSizeBinary(length) => { + Type::primitive_type_builder(name, PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_repetition(repetition) + .with_id(id) + .with_length(*length) + .with_logical_type(logical_type_for_fixed_size_binary(field)) + .build() + } + DataType::BinaryView => Type::primitive_type_builder(name, PhysicalType::BYTE_ARRAY) + .with_repetition(repetition) + .with_id(id) + .with_logical_type(logical_type_for_binary_view(field)) + .build(), + DataType::Decimal32(precision, scale) + | DataType::Decimal64(precision, scale) + | DataType::Decimal128(precision, scale) + | DataType::Decimal256(precision, scale) => { + // Decimal precision determines the Parquet physical type to use. + // Following the: https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#decimal + let (physical_type, length) = if *precision > 1 && *precision <= 9 { + (PhysicalType::INT32, -1) + } else if *precision <= 18 { + (PhysicalType::INT64, -1) + } else { + ( + PhysicalType::FIXED_LEN_BYTE_ARRAY, + decimal_length_from_precision(*precision) as i32, + ) + }; + Type::primitive_type_builder(name, physical_type) + .with_repetition(repetition) + .with_id(id) + .with_length(length) + .with_logical_type(Some(LogicalType::decimal(*scale as i32, *precision as i32))) + .with_precision(*precision as i32) + .with_scale(*scale as i32) + .build() + } + DataType::Utf8 | DataType::LargeUtf8 => { + Type::primitive_type_builder(name, PhysicalType::BYTE_ARRAY) + .with_logical_type(logical_type_for_string(field)) + .with_repetition(repetition) + .with_id(id) + .build() + } + DataType::Utf8View => Type::primitive_type_builder(name, PhysicalType::BYTE_ARRAY) + .with_logical_type(logical_type_for_string(field)) + .with_repetition(repetition) + .with_id(id) + .build(), + DataType::List(f) + | DataType::FixedSizeList(f, _) + | DataType::LargeList(f) + | DataType::ListView(f) + | DataType::LargeListView(f) => { + let field_ref = if coerce_types && f.name() != PARQUET_LIST_ELEMENT_NAME { + // Ensure proper naming per the Parquet specification + let ff = f.as_ref().clone().with_name(PARQUET_LIST_ELEMENT_NAME); + Arc::new(arrow_to_parquet_type(&ff, coerce_types)?) + } else { + Arc::new(arrow_to_parquet_type(f, coerce_types)?) + }; + + Type::group_type_builder(name) + .with_fields(vec![Arc::new( + Type::group_type_builder("list") + .with_fields(vec![field_ref]) + .with_repetition(Repetition::REPEATED) + .build()?, + )]) + .with_logical_type(Some(LogicalType::List)) + .with_repetition(repetition) + .with_id(id) + .build() + } + DataType::Struct(fields) => { + if fields.is_empty() { + return Err(arrow_err!("Parquet does not support writing empty structs",)); + } + // recursively convert children to types/nodes + let fields = fields + .iter() + .map(|f| arrow_to_parquet_type(f, coerce_types).map(Arc::new)) + .collect::>()?; + Type::group_type_builder(name) + .with_fields(fields) + .with_repetition(repetition) + .with_id(id) + .with_logical_type(logical_type_for_struct(field)) + .build() + } + DataType::Map(field, _) => { + if let DataType::Struct(struct_fields) = field.data_type() { + // If coercing then set inner struct name to "key_value" + let map_struct_name = if coerce_types { + PARQUET_MAP_STRUCT_NAME + } else { + field.name() + }; + + // If coercing then ensure struct fields are named "key" and "value" + let fix_map_field = |name: &str, fld: &Arc| -> Result> { + if coerce_types && fld.name() != name { + let f = fld.as_ref().clone().with_name(name); + Ok(Arc::new(arrow_to_parquet_type(&f, coerce_types)?)) + } else { + Ok(Arc::new(arrow_to_parquet_type(fld, coerce_types)?)) + } + }; + let key_field = fix_map_field(PARQUET_KEY_FIELD_NAME, &struct_fields[0])?; + let val_field = fix_map_field(PARQUET_VALUE_FIELD_NAME, &struct_fields[1])?; + + Type::group_type_builder(name) + .with_fields(vec![Arc::new( + Type::group_type_builder(map_struct_name) + .with_fields(vec![key_field, val_field]) + .with_repetition(Repetition::REPEATED) + .build()?, + )]) + .with_logical_type(Some(LogicalType::Map)) + .with_repetition(repetition) + .with_id(id) + .build() + } else { + Err(arrow_err!( + "DataType::Map should contain a struct field child", + )) + } + } + DataType::Union(_, _) => unimplemented!("See ARROW-8817."), + DataType::Dictionary(_, value) => { + // Dictionary encoding not handled at the schema level + let dict_field = field.clone().with_data_type(value.as_ref().clone()); + arrow_to_parquet_type(&dict_field, coerce_types) + } + DataType::RunEndEncoded(_, value_field) => { + let ree_value_field = field + .clone() + .with_data_type(value_field.data_type().clone()); + arrow_to_parquet_type(&ree_value_field, coerce_types) + } + } +} + +fn field_id(field: &Field) -> Option { + let value = field.metadata().get(super::PARQUET_FIELD_ID_META_KEY)?; + value.parse().ok() // Fail quietly if not a valid integer +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::{collections::HashMap, sync::Arc}; + + use crate::arrow::PARQUET_FIELD_ID_META_KEY; + use crate::file::metadata::KeyValue; + use crate::file::reader::FileReader; + use crate::{ + arrow::{ArrowWriter, arrow_reader::ParquetRecordBatchReaderBuilder}, + schema::{parser::parse_message_type, types::SchemaDescriptor}, + }; + use arrow::datatypes::{DataType, Field, IntervalUnit, TimeUnit}; + + #[test] + fn test_flat_primitives() { + let message_type = " + message test_schema { + REQUIRED BOOLEAN boolean; + REQUIRED INT32 int8 (INT_8); + REQUIRED INT32 int16 (INT_16); + REQUIRED INT32 uint8 (INTEGER(8,false)); + REQUIRED INT32 uint16 (INTEGER(16,false)); + REQUIRED INT32 int32; + REQUIRED INT64 int64; + OPTIONAL DOUBLE double; + OPTIONAL FLOAT float; + OPTIONAL FIXED_LEN_BYTE_ARRAY (2) float16 (FLOAT16); + OPTIONAL BINARY string (UTF8); + OPTIONAL BINARY string_2 (STRING); + OPTIONAL BINARY json (JSON); + } + "; + let parquet_group_type = parse_message_type(message_type).unwrap(); + + let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type)); + let converted_arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap(); + + let arrow_fields = Fields::from(vec![ + Field::new("boolean", DataType::Boolean, false), + Field::new("int8", DataType::Int8, false), + Field::new("int16", DataType::Int16, false), + Field::new("uint8", DataType::UInt8, false), + Field::new("uint16", DataType::UInt16, false), + Field::new("int32", DataType::Int32, false), + Field::new("int64", DataType::Int64, false), + Field::new("double", DataType::Float64, true), + Field::new("float", DataType::Float32, true), + Field::new("float16", DataType::Float16, true), + Field::new("string", DataType::Utf8, true), + Field::new("string_2", DataType::Utf8, true), + json_field(), + ]); + + assert_eq!(&arrow_fields, converted_arrow_schema.fields()); + } + + /// Return the expected Field for a Parquet column annotated with + /// the JSON logical type. + fn json_field() -> Field { + #[cfg(feature = "arrow_canonical_extension_types")] + { + Field::new("json", DataType::Utf8, true) + .with_extension_type(arrow_schema::extension::Json::default()) + } + #[cfg(not(feature = "arrow_canonical_extension_types"))] + { + Field::new("json", DataType::Utf8, true) + } + } + + #[test] + fn test_decimal_fields() { + let message_type = " + message test_schema { + REQUIRED INT32 decimal1 (DECIMAL(4,2)); + REQUIRED INT64 decimal2 (DECIMAL(12,2)); + REQUIRED FIXED_LEN_BYTE_ARRAY (16) decimal3 (DECIMAL(30,2)); + REQUIRED BYTE_ARRAY decimal4 (DECIMAL(33,2)); + REQUIRED BYTE_ARRAY decimal5 (DECIMAL(38,2)); + REQUIRED FIXED_LEN_BYTE_ARRAY (17) decimal6 (DECIMAL(39,2)); + REQUIRED BYTE_ARRAY decimal7 (DECIMAL(39,2)); + } + "; + + let parquet_group_type = parse_message_type(message_type).unwrap(); + + let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type)); + let converted_arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap(); + + let arrow_fields = Fields::from(vec![ + Field::new("decimal1", DataType::Decimal128(4, 2), false), + Field::new("decimal2", DataType::Decimal128(12, 2), false), + Field::new("decimal3", DataType::Decimal128(30, 2), false), + Field::new("decimal4", DataType::Decimal128(33, 2), false), + Field::new("decimal5", DataType::Decimal128(38, 2), false), + Field::new("decimal6", DataType::Decimal256(39, 2), false), + Field::new("decimal7", DataType::Decimal256(39, 2), false), + ]); + assert_eq!(&arrow_fields, converted_arrow_schema.fields()); + } + + #[test] + fn test_byte_array_fields() { + let message_type = " + message test_schema { + REQUIRED BYTE_ARRAY binary; + REQUIRED FIXED_LEN_BYTE_ARRAY (20) fixed_binary; + } + "; + + let parquet_group_type = parse_message_type(message_type).unwrap(); + + let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type)); + let converted_arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap(); + + let arrow_fields = Fields::from(vec![ + Field::new("binary", DataType::Binary, false), + Field::new("fixed_binary", DataType::FixedSizeBinary(20), false), + ]); + assert_eq!(&arrow_fields, converted_arrow_schema.fields()); + } + + #[test] + fn test_duplicate_fields() { + let message_type = " + message test_schema { + REQUIRED BOOLEAN boolean; + REQUIRED INT32 int8 (INT_8); + } + "; + + let parquet_group_type = parse_message_type(message_type).unwrap(); + + let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type)); + let converted_arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap(); + + let arrow_fields = Fields::from(vec![ + Field::new("boolean", DataType::Boolean, false), + Field::new("int8", DataType::Int8, false), + ]); + assert_eq!(&arrow_fields, converted_arrow_schema.fields()); + + let converted_arrow_schema = + parquet_to_arrow_schema_by_columns(&parquet_schema, ProjectionMask::all(), None) + .unwrap(); + assert_eq!(&arrow_fields, converted_arrow_schema.fields()); + } + + #[test] + fn test_parquet_lists() { + let mut arrow_fields = Vec::new(); + + // LIST encoding example taken from parquet-format/LogicalTypes.md + let message_type = " + message test_schema { + REQUIRED GROUP my_list (LIST) { + REPEATED GROUP list { + OPTIONAL BINARY element (UTF8); + } + } + OPTIONAL GROUP my_list (LIST) { + REPEATED GROUP list { + REQUIRED BINARY element (UTF8); + } + } + OPTIONAL GROUP array_of_arrays (LIST) { + REPEATED GROUP list { + REQUIRED GROUP element (LIST) { + REPEATED GROUP list { + REQUIRED INT32 element; + } + } + } + } + OPTIONAL GROUP my_list (LIST) { + REPEATED GROUP element { + REQUIRED BINARY str (UTF8); + } + } + OPTIONAL GROUP my_list (LIST) { + REPEATED INT32 element; + } + OPTIONAL GROUP my_list (LIST) { + REPEATED GROUP element { + REQUIRED BINARY str (UTF8); + REQUIRED INT32 num; + } + } + OPTIONAL GROUP my_list (LIST) { + REPEATED GROUP array { + REQUIRED BINARY str (UTF8); + } + + } + OPTIONAL GROUP my_list (LIST) { + REPEATED GROUP my_list_tuple { + REQUIRED BINARY str (UTF8); + } + } + REPEATED INT32 name; + } + "; + + // // List (list non-null, elements nullable) + // required group my_list (LIST) { + // repeated group list { + // optional binary element (UTF8); + // } + // } + { + arrow_fields.push(Field::new_list( + "my_list", + Field::new("element", DataType::Utf8, true), + false, + )); + } + + // // List (list nullable, elements non-null) + // optional group my_list (LIST) { + // repeated group list { + // required binary element (UTF8); + // } + // } + { + arrow_fields.push(Field::new_list( + "my_list", + Field::new("element", DataType::Utf8, false), + true, + )); + } + + // Element types can be nested structures. For example, a list of lists: + // + // // List> + // optional group array_of_arrays (LIST) { + // repeated group list { + // required group element (LIST) { + // repeated group list { + // required int32 element; + // } + // } + // } + // } + { + let arrow_inner_list = Field::new("element", DataType::Int32, false); + arrow_fields.push(Field::new_list( + "array_of_arrays", + Field::new_list("element", arrow_inner_list, false), + true, + )); + } + + // // List (list nullable, elements non-null) + // optional group my_list (LIST) { + // repeated group element { + // required binary str (UTF8); + // }; + // } + { + arrow_fields.push(Field::new_list( + "my_list", + Field::new("str", DataType::Utf8, false), + true, + )); + } + + // // List (nullable list, non-null elements) + // optional group my_list (LIST) { + // repeated int32 element; + // } + { + arrow_fields.push(Field::new_list( + "my_list", + Field::new("element", DataType::Int32, false), + true, + )); + } + + // // List> (nullable list, non-null elements) + // optional group my_list (LIST) { + // repeated group element { + // required binary str (UTF8); + // required int32 num; + // }; + // } + { + let fields = vec![ + Field::new("str", DataType::Utf8, false), + Field::new("num", DataType::Int32, false), + ]; + arrow_fields.push(Field::new_list( + "my_list", + Field::new_struct("element", fields, false), + true, + )); + } + + // // List> (nullable list, non-null elements) + // optional group my_list (LIST) { + // repeated group array { + // required binary str (UTF8); + // }; + // } + // Special case: group is named array + { + let fields = vec![Field::new("str", DataType::Utf8, false)]; + arrow_fields.push(Field::new_list( + "my_list", + Field::new_struct("array", fields, false), + true, + )); + } + + // // List> (nullable list, non-null elements) + // optional group my_list (LIST) { + // repeated group my_list_tuple { + // required binary str (UTF8); + // }; + // } + // Special case: group named ends in _tuple + { + let fields = vec![Field::new("str", DataType::Utf8, false)]; + arrow_fields.push(Field::new_list( + "my_list", + Field::new_struct("my_list_tuple", fields, false), + true, + )); + } + + // One-level encoding: Only allows required lists with required cells + // repeated value_type name + { + arrow_fields.push(Field::new_list( + "name", + Field::new("name", DataType::Int32, false), + false, + )); + } + + let parquet_group_type = parse_message_type(message_type).unwrap(); + + let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type)); + let converted_arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap(); + let converted_fields = converted_arrow_schema.fields(); + + assert_eq!(arrow_fields.len(), converted_fields.len()); + for i in 0..arrow_fields.len() { + assert_eq!(&arrow_fields[i], converted_fields[i].as_ref(), "{i}"); + } + } + + #[test] + fn test_parquet_list_nullable() { + let mut arrow_fields = Vec::new(); + + let message_type = " + message test_schema { + REQUIRED GROUP my_list1 (LIST) { + REPEATED GROUP list { + OPTIONAL BINARY element (UTF8); + } + } + OPTIONAL GROUP my_list2 (LIST) { + REPEATED GROUP list { + REQUIRED BINARY element (UTF8); + } + } + REQUIRED GROUP my_list3 (LIST) { + REPEATED GROUP list { + REQUIRED BINARY element (UTF8); + } + } + } + "; + + // // List (list non-null, elements nullable) + // required group my_list1 (LIST) { + // repeated group list { + // optional binary element (UTF8); + // } + // } + { + arrow_fields.push(Field::new_list( + "my_list1", + Field::new("element", DataType::Utf8, true), + false, + )); + } + + // // List (list nullable, elements non-null) + // optional group my_list2 (LIST) { + // repeated group list { + // required binary element (UTF8); + // } + // } + { + arrow_fields.push(Field::new_list( + "my_list2", + Field::new("element", DataType::Utf8, false), + true, + )); + } + + // // List (list non-null, elements non-null) + // repeated group my_list3 (LIST) { + // repeated group list { + // required binary element (UTF8); + // } + // } + { + arrow_fields.push(Field::new_list( + "my_list3", + Field::new("element", DataType::Utf8, false), + false, + )); + } + + let parquet_group_type = parse_message_type(message_type).unwrap(); + + let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type)); + let converted_arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap(); + let converted_fields = converted_arrow_schema.fields(); + + assert_eq!(arrow_fields.len(), converted_fields.len()); + for i in 0..arrow_fields.len() { + assert_eq!(&arrow_fields[i], converted_fields[i].as_ref()); + } + } + + #[test] + fn test_parquet_maps() { + let mut arrow_fields = Vec::new(); + + // LIST encoding example taken from parquet-format/LogicalTypes.md + let message_type = " + message test_schema { + REQUIRED group my_map1 (MAP) { + REPEATED group key_value { + REQUIRED binary key (UTF8); + OPTIONAL int32 value; + } + } + OPTIONAL group my_map2 (MAP) { + REPEATED group map { + REQUIRED binary str (UTF8); + REQUIRED int32 num; + } + } + OPTIONAL group my_map3 (MAP_KEY_VALUE) { + REPEATED group map { + REQUIRED binary key (UTF8); + OPTIONAL int32 value; + } + } + REQUIRED group my_map4 (MAP) { + REPEATED group map { + OPTIONAL binary key (UTF8); + REQUIRED int32 value; + } + } + } + "; + + // // Map + // required group my_map (MAP) { + // repeated group key_value { + // required binary key (UTF8); + // optional int32 value; + // } + // } + { + arrow_fields.push(Field::new_map( + "my_map1", + "key_value", + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Int32, true), + false, + false, + )); + } + + // // Map (nullable map, non-null values) + // optional group my_map (MAP) { + // repeated group map { + // required binary str (UTF8); + // required int32 num; + // } + // } + { + arrow_fields.push(Field::new_map( + "my_map2", + "map", + Field::new("str", DataType::Utf8, false), + Field::new("num", DataType::Int32, false), + false, + true, + )); + } + + // // Map (nullable map, nullable values) + // optional group my_map (MAP_KEY_VALUE) { + // repeated group map { + // required binary key (UTF8); + // optional int32 value; + // } + // } + { + arrow_fields.push(Field::new_map( + "my_map3", + "map", + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Int32, true), + false, + true, + )); + } + + // // Map (non-compliant nullable key) + // group my_map (MAP_KEY_VALUE) { + // repeated group map { + // optional binary key (UTF8); + // required int32 value; + // } + // } + { + arrow_fields.push(Field::new_map( + "my_map4", + "map", + Field::new("key", DataType::Utf8, false), // The key is always non-nullable (#5630) + Field::new("value", DataType::Int32, false), + false, + false, + )); + } + + let parquet_group_type = parse_message_type(message_type).unwrap(); + + let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type)); + let converted_arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap(); + let converted_fields = converted_arrow_schema.fields(); + + assert_eq!(arrow_fields.len(), converted_fields.len()); + for i in 0..arrow_fields.len() { + assert_eq!(&arrow_fields[i], converted_fields[i].as_ref()); + } + } + + #[test] + fn test_nested_schema() { + let mut arrow_fields = Vec::new(); + { + let group1_fields = Fields::from(vec![ + Field::new("leaf1", DataType::Boolean, false), + Field::new("leaf2", DataType::Int32, false), + ]); + let group1_struct = Field::new("group1", DataType::Struct(group1_fields), false); + arrow_fields.push(group1_struct); + + let leaf3_field = Field::new("leaf3", DataType::Int64, false); + arrow_fields.push(leaf3_field); + } + + let message_type = " + message test_schema { + REQUIRED GROUP group1 { + REQUIRED BOOLEAN leaf1; + REQUIRED INT32 leaf2; + } + REQUIRED INT64 leaf3; + } + "; + let parquet_group_type = parse_message_type(message_type).unwrap(); + + let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type)); + let converted_arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap(); + let converted_fields = converted_arrow_schema.fields(); + + assert_eq!(arrow_fields.len(), converted_fields.len()); + for i in 0..arrow_fields.len() { + assert_eq!(&arrow_fields[i], converted_fields[i].as_ref()); + } + } + + #[test] + fn test_nested_schema_partial() { + let mut arrow_fields = Vec::new(); + { + let group1_fields = vec![Field::new("leaf1", DataType::Int64, false)].into(); + let group1 = Field::new("group1", DataType::Struct(group1_fields), false); + arrow_fields.push(group1); + + let group2_fields = vec![Field::new("leaf4", DataType::Int64, false)].into(); + let group2 = Field::new("group2", DataType::Struct(group2_fields), false); + arrow_fields.push(group2); + + arrow_fields.push(Field::new("leaf5", DataType::Int64, false)); + } + + let message_type = " + message test_schema { + REQUIRED GROUP group1 { + REQUIRED INT64 leaf1; + REQUIRED INT64 leaf2; + } + REQUIRED GROUP group2 { + REQUIRED INT64 leaf3; + REQUIRED INT64 leaf4; + } + REQUIRED INT64 leaf5; + } + "; + let parquet_group_type = parse_message_type(message_type).unwrap(); + + // Expected partial arrow schema (columns 0, 3, 4): + // required group group1 { + // required int64 leaf1; + // } + // required group group2 { + // required int64 leaf4; + // } + // required int64 leaf5; + + let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type)); + let mask = ProjectionMask::leaves(&parquet_schema, [3, 0, 4, 4]); + let converted_arrow_schema = + parquet_to_arrow_schema_by_columns(&parquet_schema, mask, None).unwrap(); + let converted_fields = converted_arrow_schema.fields(); + + assert_eq!(arrow_fields.len(), converted_fields.len()); + for i in 0..arrow_fields.len() { + assert_eq!(&arrow_fields[i], converted_fields[i].as_ref()); + } + } + + #[test] + fn test_nested_schema_partial_ordering() { + let mut arrow_fields = Vec::new(); + { + let group1_fields = vec![Field::new("leaf1", DataType::Int64, false)].into(); + let group1 = Field::new("group1", DataType::Struct(group1_fields), false); + arrow_fields.push(group1); + + let group2_fields = vec![Field::new("leaf4", DataType::Int64, false)].into(); + let group2 = Field::new("group2", DataType::Struct(group2_fields), false); + arrow_fields.push(group2); + + arrow_fields.push(Field::new("leaf5", DataType::Int64, false)); + } + + let message_type = " + message test_schema { + REQUIRED GROUP group1 { + REQUIRED INT64 leaf1; + REQUIRED INT64 leaf2; + } + REQUIRED GROUP group2 { + REQUIRED INT64 leaf3; + REQUIRED INT64 leaf4; + } + REQUIRED INT64 leaf5; + } + "; + let parquet_group_type = parse_message_type(message_type).unwrap(); + + // Expected partial arrow schema (columns 3, 4, 0): + // required group group1 { + // required int64 leaf1; + // } + // required group group2 { + // required int64 leaf4; + // } + // required int64 leaf5; + + let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type)); + let mask = ProjectionMask::leaves(&parquet_schema, [3, 0, 4]); + let converted_arrow_schema = + parquet_to_arrow_schema_by_columns(&parquet_schema, mask, None).unwrap(); + let converted_fields = converted_arrow_schema.fields(); + + assert_eq!(arrow_fields.len(), converted_fields.len()); + for i in 0..arrow_fields.len() { + assert_eq!(&arrow_fields[i], converted_fields[i].as_ref()); + } + + let mask = + ProjectionMask::columns(&parquet_schema, ["group2.leaf4", "group1.leaf1", "leaf5"]); + let converted_arrow_schema = + parquet_to_arrow_schema_by_columns(&parquet_schema, mask, None).unwrap(); + let converted_fields = converted_arrow_schema.fields(); + + assert_eq!(arrow_fields.len(), converted_fields.len()); + for i in 0..arrow_fields.len() { + assert_eq!(&arrow_fields[i], converted_fields[i].as_ref()); + } + } + + #[test] + fn test_repeated_nested_schema() { + let mut arrow_fields = Vec::new(); + { + arrow_fields.push(Field::new("leaf1", DataType::Int32, true)); + + let inner_group_list = Field::new_list( + "innerGroup", + Field::new_struct( + "innerGroup", + vec![Field::new("leaf3", DataType::Int32, true)], + false, + ), + false, + ); + + let outer_group_list = Field::new_list( + "outerGroup", + Field::new_struct( + "outerGroup", + vec![Field::new("leaf2", DataType::Int32, true), inner_group_list], + false, + ), + false, + ); + arrow_fields.push(outer_group_list); + } + + let message_type = " + message test_schema { + OPTIONAL INT32 leaf1; + REPEATED GROUP outerGroup { + OPTIONAL INT32 leaf2; + REPEATED GROUP innerGroup { + OPTIONAL INT32 leaf3; + } + } + } + "; + let parquet_group_type = parse_message_type(message_type).unwrap(); + + let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type)); + let converted_arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap(); + let converted_fields = converted_arrow_schema.fields(); + + assert_eq!(arrow_fields.len(), converted_fields.len()); + for i in 0..arrow_fields.len() { + assert_eq!(&arrow_fields[i], converted_fields[i].as_ref()); + } + } + + #[test] + fn test_column_desc_to_field() { + let message_type = " + message test_schema { + REQUIRED BOOLEAN boolean; + REQUIRED INT32 int8 (INT_8); + REQUIRED INT32 uint8 (INTEGER(8,false)); + REQUIRED INT32 int16 (INT_16); + REQUIRED INT32 uint16 (INTEGER(16,false)); + REQUIRED INT32 int32; + REQUIRED INT64 int64; + OPTIONAL DOUBLE double; + OPTIONAL FLOAT float; + OPTIONAL FIXED_LEN_BYTE_ARRAY (2) float16 (FLOAT16); + OPTIONAL BINARY string (UTF8); + REPEATED BOOLEAN bools; + OPTIONAL INT32 date (DATE); + OPTIONAL INT32 time_milli (TIME_MILLIS); + OPTIONAL INT64 time_micro (TIME_MICROS); + OPTIONAL INT64 time_nano (TIME(NANOS,false)); + OPTIONAL INT64 ts_milli (TIMESTAMP_MILLIS); + REQUIRED INT64 ts_micro (TIMESTAMP_MICROS); + REQUIRED INT64 ts_nano (TIMESTAMP(NANOS,true)); + REPEATED INT32 int_list; + REPEATED BINARY byte_list; + REPEATED BINARY string_list (UTF8); + REQUIRED INT32 decimal_int32 (DECIMAL(8,2)); + REQUIRED INT64 decimal_int64 (DECIMAL(16,2)); + REQUIRED FIXED_LEN_BYTE_ARRAY (13) decimal_fix_length (DECIMAL(30,2)); + } + "; + let parquet_group_type = parse_message_type(message_type).unwrap(); + + let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type)); + let converted_arrow_fields = parquet_schema + .columns() + .iter() + .map(|c| parquet_to_arrow_field(c).unwrap()) + .collect::>(); + + let arrow_fields = vec![ + Field::new("boolean", DataType::Boolean, false), + Field::new("int8", DataType::Int8, false), + Field::new("uint8", DataType::UInt8, false), + Field::new("int16", DataType::Int16, false), + Field::new("uint16", DataType::UInt16, false), + Field::new("int32", DataType::Int32, false), + Field::new("int64", DataType::Int64, false), + Field::new("double", DataType::Float64, true), + Field::new("float", DataType::Float32, true), + Field::new("float16", DataType::Float16, true), + Field::new("string", DataType::Utf8, true), + Field::new_list( + "bools", + Field::new("bools", DataType::Boolean, false), + false, + ), + Field::new("date", DataType::Date32, true), + Field::new("time_milli", DataType::Time32(TimeUnit::Millisecond), true), + Field::new("time_micro", DataType::Time64(TimeUnit::Microsecond), true), + Field::new("time_nano", DataType::Time64(TimeUnit::Nanosecond), true), + Field::new( + "ts_milli", + DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())), + true, + ), + Field::new( + "ts_micro", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + false, + ), + Field::new( + "ts_nano", + DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())), + false, + ), + Field::new_list( + "int_list", + Field::new("int_list", DataType::Int32, false), + false, + ), + Field::new_list( + "byte_list", + Field::new("byte_list", DataType::Binary, false), + false, + ), + Field::new_list( + "string_list", + Field::new("string_list", DataType::Utf8, false), + false, + ), + Field::new("decimal_int32", DataType::Decimal128(8, 2), false), + Field::new("decimal_int64", DataType::Decimal128(16, 2), false), + Field::new("decimal_fix_length", DataType::Decimal128(30, 2), false), + ]; + + assert_eq!(arrow_fields, converted_arrow_fields); + } + + #[test] + fn test_coerced_map_list() { + // Create Arrow schema with non-Parquet naming + let arrow_fields = vec![ + Field::new_list( + "my_list", + Field::new("item", DataType::Boolean, true), + false, + ), + Field::new_map( + "my_map", + "entries", + Field::new("keys", DataType::Utf8, false), + Field::new("values", DataType::Int32, true), + false, + true, + ), + ]; + let arrow_schema = Schema::new(arrow_fields); + + // Create Parquet schema with coerced names + let message_type = " + message parquet_schema { + REQUIRED GROUP my_list (LIST) { + REPEATED GROUP list { + OPTIONAL BOOLEAN element; + } + } + OPTIONAL GROUP my_map (MAP) { + REPEATED GROUP key_value { + REQUIRED BINARY key (STRING); + OPTIONAL INT32 value; + } + } + } + "; + let parquet_group_type = parse_message_type(message_type).unwrap(); + let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type)); + let converted_arrow_schema = ArrowSchemaConverter::new() + .with_coerce_types(true) + .convert(&arrow_schema) + .unwrap(); + assert_eq!( + parquet_schema.columns().len(), + converted_arrow_schema.columns().len() + ); + + // Create Parquet schema without coerced names + let message_type = " + message parquet_schema { + REQUIRED GROUP my_list (LIST) { + REPEATED GROUP list { + OPTIONAL BOOLEAN item; + } + } + OPTIONAL GROUP my_map (MAP) { + REPEATED GROUP entries { + REQUIRED BINARY keys (STRING); + OPTIONAL INT32 values; + } + } + } + "; + let parquet_group_type = parse_message_type(message_type).unwrap(); + let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type)); + let converted_arrow_schema = ArrowSchemaConverter::new() + .with_coerce_types(false) + .convert(&arrow_schema) + .unwrap(); + assert_eq!( + parquet_schema.columns().len(), + converted_arrow_schema.columns().len() + ); + } + + #[test] + fn test_field_to_column_desc() { + let message_type = " + message arrow_schema { + REQUIRED BOOLEAN boolean; + REQUIRED INT32 int8 (INT_8); + REQUIRED INT32 int16 (INTEGER(16,true)); + REQUIRED INT32 int32; + REQUIRED INT64 int64; + OPTIONAL DOUBLE double; + OPTIONAL FLOAT float; + OPTIONAL FIXED_LEN_BYTE_ARRAY (2) float16 (FLOAT16); + OPTIONAL BINARY string (STRING); + OPTIONAL GROUP bools (LIST) { + REPEATED GROUP list { + OPTIONAL BOOLEAN element; + } + } + REQUIRED GROUP bools_non_null (LIST) { + REPEATED GROUP list { + REQUIRED BOOLEAN element; + } + } + OPTIONAL INT32 date (DATE); + OPTIONAL INT32 time_milli (TIME(MILLIS,false)); + OPTIONAL INT32 time_milli_utc (TIME(MILLIS,true)); + OPTIONAL INT64 time_micro (TIME_MICROS); + OPTIONAL INT64 time_micro_utc (TIME(MICROS, true)); + OPTIONAL INT64 ts_milli (TIMESTAMP_MILLIS); + REQUIRED INT64 ts_micro (TIMESTAMP(MICROS,false)); + REQUIRED INT64 ts_seconds; + REQUIRED INT64 ts_micro_utc (TIMESTAMP(MICROS, true)); + REQUIRED INT64 ts_millis_zero_offset (TIMESTAMP(MILLIS, true)); + REQUIRED INT64 ts_millis_zero_negative_offset (TIMESTAMP(MILLIS, true)); + REQUIRED INT64 ts_micro_non_utc (TIMESTAMP(MICROS, true)); + REQUIRED GROUP struct { + REQUIRED BOOLEAN bools; + REQUIRED INT32 uint32 (INTEGER(32,false)); + REQUIRED GROUP int32 (LIST) { + REPEATED GROUP list { + OPTIONAL INT32 element; + } + } + } + REQUIRED BINARY dictionary_strings (STRING); + REQUIRED INT32 decimal_int32 (DECIMAL(8,2)); + REQUIRED INT64 decimal_int64 (DECIMAL(16,2)); + REQUIRED FIXED_LEN_BYTE_ARRAY (13) decimal_fix_length (DECIMAL(30,2)); + REQUIRED FIXED_LEN_BYTE_ARRAY (16) decimal128 (DECIMAL(38,2)); + REQUIRED FIXED_LEN_BYTE_ARRAY (17) decimal256 (DECIMAL(39,2)); + } + "; + let parquet_group_type = parse_message_type(message_type).unwrap(); + + let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type)); + + let arrow_fields = vec![ + Field::new("boolean", DataType::Boolean, false), + Field::new("int8", DataType::Int8, false), + Field::new("int16", DataType::Int16, false), + Field::new("int32", DataType::Int32, false), + Field::new("int64", DataType::Int64, false), + Field::new("double", DataType::Float64, true), + Field::new("float", DataType::Float32, true), + Field::new("float16", DataType::Float16, true), + Field::new("string", DataType::Utf8, true), + Field::new_list( + "bools", + Field::new("element", DataType::Boolean, true), + true, + ), + Field::new_list( + "bools_non_null", + Field::new("element", DataType::Boolean, false), + false, + ), + Field::new("date", DataType::Date32, true), + Field::new("time_milli", DataType::Time32(TimeUnit::Millisecond), true), + Field::new( + "time_milli_utc", + DataType::Time32(TimeUnit::Millisecond), + true, + ) + .with_metadata(HashMap::from_iter(vec![( + "adjusted_to_utc".to_string(), + "".to_string(), + )])), + Field::new("time_micro", DataType::Time64(TimeUnit::Microsecond), true), + Field::new( + "time_micro_utc", + DataType::Time64(TimeUnit::Microsecond), + true, + ) + .with_metadata(HashMap::from_iter(vec![( + "adjusted_to_utc".to_string(), + "".to_string(), + )])), + Field::new( + "ts_milli", + DataType::Timestamp(TimeUnit::Millisecond, None), + true, + ), + Field::new( + "ts_micro", + DataType::Timestamp(TimeUnit::Microsecond, None), + false, + ), + Field::new( + "ts_seconds", + DataType::Timestamp(TimeUnit::Second, Some("UTC".into())), + false, + ), + Field::new( + "ts_micro_utc", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + false, + ), + Field::new( + "ts_millis_zero_offset", + DataType::Timestamp(TimeUnit::Millisecond, Some("+00:00".into())), + false, + ), + Field::new( + "ts_millis_zero_negative_offset", + DataType::Timestamp(TimeUnit::Millisecond, Some("-00:00".into())), + false, + ), + Field::new( + "ts_micro_non_utc", + DataType::Timestamp(TimeUnit::Microsecond, Some("+01:00".into())), + false, + ), + Field::new_struct( + "struct", + vec![ + Field::new("bools", DataType::Boolean, false), + Field::new("uint32", DataType::UInt32, false), + Field::new_list("int32", Field::new("element", DataType::Int32, true), false), + ], + false, + ), + Field::new_dictionary("dictionary_strings", DataType::Int32, DataType::Utf8, false), + Field::new("decimal_int32", DataType::Decimal128(8, 2), false), + Field::new("decimal_int64", DataType::Decimal128(16, 2), false), + Field::new("decimal_fix_length", DataType::Decimal128(30, 2), false), + Field::new("decimal128", DataType::Decimal128(38, 2), false), + Field::new("decimal256", DataType::Decimal256(39, 2), false), + ]; + let arrow_schema = Schema::new(arrow_fields); + let converted_arrow_schema = ArrowSchemaConverter::new().convert(&arrow_schema).unwrap(); + + assert_eq!( + parquet_schema.columns().len(), + converted_arrow_schema.columns().len() + ); + parquet_schema + .columns() + .iter() + .zip(converted_arrow_schema.columns()) + .for_each(|(a, b)| { + // Only check logical type if it's set on the Parquet side. + // This is because the Arrow conversion always sets logical type, + // even if there wasn't originally one. + // This is not an issue, but is an inconvenience for this test. + match a.logical_type_ref() { + Some(_) => { + assert_eq!(a, b) + } + None => { + assert_eq!(a.name(), b.name()); + assert_eq!(a.physical_type(), b.physical_type()); + assert_eq!(a.converted_type(), b.converted_type()); + } + }; + }); + } + + #[test] + #[should_panic(expected = "Parquet does not support writing empty structs")] + fn test_empty_struct_field() { + let arrow_fields = vec![Field::new( + "struct", + DataType::Struct(Fields::empty()), + false, + )]; + let arrow_schema = Schema::new(arrow_fields); + let converted_arrow_schema = ArrowSchemaConverter::new() + .with_coerce_types(true) + .convert(&arrow_schema); + + converted_arrow_schema.unwrap(); + } + + #[test] + fn test_metadata() { + let message_type = " + message test_schema { + OPTIONAL BINARY string (STRING); + } + "; + let parquet_group_type = parse_message_type(message_type).unwrap(); + + let key_value_metadata = vec![ + KeyValue::new("foo".to_owned(), Some("bar".to_owned())), + KeyValue::new("baz".to_owned(), None), + ]; + + let mut expected_metadata: HashMap = HashMap::new(); + expected_metadata.insert("foo".to_owned(), "bar".to_owned()); + + let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type)); + let converted_arrow_schema = + parquet_to_arrow_schema(&parquet_schema, Some(&key_value_metadata)).unwrap(); + + assert_eq!(converted_arrow_schema.metadata(), &expected_metadata); + } + + #[test] + fn test_arrow_schema_roundtrip() -> Result<()> { + let meta = |a: &[(&str, &str)]| -> HashMap { + a.iter() + .map(|(a, b)| (a.to_string(), b.to_string())) + .collect() + }; + + let schema = Schema::new_with_metadata( + vec![ + Field::new("c1", DataType::Utf8, false) + .with_metadata(meta(&[("Key", "Foo"), (PARQUET_FIELD_ID_META_KEY, "2")])), + Field::new("c2", DataType::Binary, false), + Field::new("c3", DataType::FixedSizeBinary(3), false), + Field::new("c4", DataType::Boolean, false), + Field::new("c5", DataType::Date32, false), + Field::new("c6", DataType::Date64, false), + Field::new("c7", DataType::Time32(TimeUnit::Second), false), + Field::new("c8", DataType::Time32(TimeUnit::Millisecond), false), + Field::new("c13", DataType::Time64(TimeUnit::Microsecond), false), + Field::new("c14", DataType::Time64(TimeUnit::Nanosecond), false), + Field::new("c15", DataType::Timestamp(TimeUnit::Second, None), false), + Field::new( + "c16", + DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())), + false, + ), + Field::new( + "c17", + DataType::Timestamp(TimeUnit::Microsecond, Some("Africa/Johannesburg".into())), + false, + ), + Field::new( + "c18", + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ), + Field::new("c19", DataType::Interval(IntervalUnit::DayTime), false), + Field::new("c20", DataType::Interval(IntervalUnit::YearMonth), false), + Field::new_list( + "c21", + Field::new_list_field(DataType::Boolean, true) + .with_metadata(meta(&[("Key", "Bar"), (PARQUET_FIELD_ID_META_KEY, "5")])), + false, + ) + .with_metadata(meta(&[(PARQUET_FIELD_ID_META_KEY, "4")])), + Field::new( + "c22", + DataType::FixedSizeList( + Arc::new(Field::new_list_field(DataType::Boolean, true)), + 5, + ), + false, + ), + Field::new_list( + "c23", + Field::new_large_list( + "inner", + Field::new_list_field( + DataType::Struct( + vec![ + Field::new("a", DataType::Int16, true), + Field::new("b", DataType::Float64, false), + Field::new("c", DataType::Float32, false), + Field::new("d", DataType::Float16, false), + ] + .into(), + ), + false, + ), + true, + ), + false, + ), + Field::new( + "c24", + DataType::Struct(Fields::from(vec![ + Field::new("a", DataType::Utf8, false), + Field::new("b", DataType::UInt16, false), + ])), + false, + ), + Field::new("c25", DataType::Interval(IntervalUnit::YearMonth), true), + Field::new("c26", DataType::Interval(IntervalUnit::DayTime), true), + // Duration types not supported + // Field::new("c27", DataType::Duration(TimeUnit::Second), false), + // Field::new("c28", DataType::Duration(TimeUnit::Millisecond), false), + // Field::new("c29", DataType::Duration(TimeUnit::Microsecond), false), + // Field::new("c30", DataType::Duration(TimeUnit::Nanosecond), false), + #[allow(deprecated)] + Field::new_dict( + "c31", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + true, + 123, + true, + ) + .with_metadata(meta(&[(PARQUET_FIELD_ID_META_KEY, "6")])), + Field::new("c32", DataType::LargeBinary, true), + Field::new("c33", DataType::LargeUtf8, true), + Field::new_large_list( + "c34", + Field::new_list( + "inner", + Field::new_list_field( + DataType::Struct( + vec![ + Field::new("a", DataType::Int16, true), + Field::new("b", DataType::Float64, true), + ] + .into(), + ), + true, + ), + true, + ), + true, + ), + Field::new("c35", DataType::Null, true), + Field::new("c36", DataType::Decimal128(2, 1), false), + Field::new("c37", DataType::Decimal256(50, 20), false), + Field::new("c38", DataType::Decimal128(18, 12), true), + Field::new_map( + "c39", + "key_value", + Field::new("key", DataType::Utf8, false), + Field::new_list("value", Field::new("element", DataType::Utf8, true), true), + false, // fails to roundtrip keys_sorted + true, + ), + Field::new_map( + "c40", + "my_entries", + Field::new("my_key", DataType::Utf8, false) + .with_metadata(meta(&[(PARQUET_FIELD_ID_META_KEY, "8")])), + Field::new_list( + "my_value", + Field::new_list_field(DataType::Utf8, true) + .with_metadata(meta(&[(PARQUET_FIELD_ID_META_KEY, "10")])), + true, + ) + .with_metadata(meta(&[(PARQUET_FIELD_ID_META_KEY, "9")])), + false, // fails to roundtrip keys_sorted + true, + ) + .with_metadata(meta(&[(PARQUET_FIELD_ID_META_KEY, "7")])), + Field::new_map( + "c41", + "my_entries", + Field::new("my_key", DataType::Utf8, false), + Field::new_list( + "my_value", + Field::new_list_field(DataType::Utf8, true) + .with_metadata(meta(&[(PARQUET_FIELD_ID_META_KEY, "11")])), + true, + ), + false, // fails to roundtrip keys_sorted + false, + ), + Field::new("c42", DataType::Decimal32(5, 2), false), + Field::new("c43", DataType::Decimal64(18, 12), true), + ], + meta(&[("Key", "Value")]), + ); + + // write to an empty parquet file so that schema is serialized + let file = tempfile::tempfile().unwrap(); + let writer = + ArrowWriter::try_new(file.try_clone().unwrap(), Arc::new(schema.clone()), None)?; + writer.close()?; + + // read file back + let arrow_reader = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + + // Check arrow schema + let read_schema = arrow_reader.schema(); + assert_eq!(&schema, read_schema.as_ref()); + + // Walk schema finding field IDs + let mut stack = Vec::with_capacity(10); + let mut out = Vec::with_capacity(10); + + let root = arrow_reader.parquet_schema().root_schema_ptr(); + stack.push((root.name().to_string(), root)); + + while let Some((p, t)) = stack.pop() { + if t.is_group() { + for f in t.get_fields() { + stack.push((format!("{p}.{}", f.name()), f.clone())) + } + } + + let info = t.get_basic_info(); + if info.has_id() { + out.push(format!("{p} -> {}", info.id())) + } + } + out.sort_unstable(); + let out: Vec<_> = out.iter().map(|x| x.as_str()).collect(); + + assert_eq!( + &out, + &[ + "arrow_schema.c1 -> 2", + "arrow_schema.c21 -> 4", + "arrow_schema.c21.list.item -> 5", + "arrow_schema.c31 -> 6", + "arrow_schema.c40 -> 7", + "arrow_schema.c40.my_entries.my_key -> 8", + "arrow_schema.c40.my_entries.my_value -> 9", + "arrow_schema.c40.my_entries.my_value.list.item -> 10", + "arrow_schema.c41.my_entries.my_value.list.item -> 11", + ] + ); + + Ok(()) + } + + #[test] + fn test_read_parquet_field_ids_raw() -> Result<()> { + let meta = |a: &[(&str, &str)]| -> HashMap { + a.iter() + .map(|(a, b)| (a.to_string(), b.to_string())) + .collect() + }; + let schema = Schema::new_with_metadata( + vec![ + Field::new("c1", DataType::Utf8, true) + .with_metadata(meta(&[(PARQUET_FIELD_ID_META_KEY, "1")])), + Field::new("c2", DataType::Utf8, true) + .with_metadata(meta(&[(PARQUET_FIELD_ID_META_KEY, "2")])), + ], + HashMap::new(), + ); + + let writer = ArrowWriter::try_new(vec![], Arc::new(schema.clone()), None)?; + let parquet_bytes = writer.into_inner()?; + + let reader = + crate::file::reader::SerializedFileReader::new(bytes::Bytes::from(parquet_bytes))?; + let schema_descriptor = reader.metadata().file_metadata().schema_descr_ptr(); + + // don't pass metadata so field ids are read from Parquet and not from serialized Arrow schema + let arrow_schema = crate::arrow::parquet_to_arrow_schema(&schema_descriptor, None)?; + + let parq_schema_descr = ArrowSchemaConverter::new() + .with_coerce_types(true) + .convert(&arrow_schema)?; + let parq_fields = parq_schema_descr.root_schema().get_fields(); + assert_eq!(parq_fields.len(), 2); + assert_eq!(parq_fields[0].get_basic_info().id(), 1); + assert_eq!(parq_fields[1].get_basic_info().id(), 2); + + Ok(()) + } + + #[test] + fn test_arrow_schema_roundtrip_lists() -> Result<()> { + let metadata: HashMap = [("Key".to_string(), "Value".to_string())] + .iter() + .cloned() + .collect(); + + let schema = Schema::new_with_metadata( + vec![ + Field::new_list("c21", Field::new("array", DataType::Boolean, true), false), + Field::new( + "c22", + DataType::FixedSizeList( + Arc::new(Field::new("items", DataType::Boolean, false)), + 5, + ), + false, + ), + Field::new_list( + "c23", + Field::new_large_list( + "items", + Field::new_struct( + "items", + vec![ + Field::new("a", DataType::Int16, true), + Field::new("b", DataType::Float64, false), + ], + true, + ), + true, + ), + true, + ), + ], + metadata, + ); + + // write to an empty parquet file so that schema is serialized + let file = tempfile::tempfile().unwrap(); + let writer = + ArrowWriter::try_new(file.try_clone().unwrap(), Arc::new(schema.clone()), None)?; + writer.close()?; + + // read file back + let arrow_reader = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + let read_schema = arrow_reader.schema(); + assert_eq!(&schema, read_schema.as_ref()); + Ok(()) + } + + #[test] + fn test_get_arrow_schema_from_metadata() { + assert!(get_arrow_schema_from_metadata("").is_err()); + } + + #[test] + #[cfg(feature = "arrow_canonical_extension_types")] + fn arrow_uuid_to_parquet_uuid() -> Result<()> { + use arrow_schema::extension::Uuid; + let arrow_schema = Schema::new(vec![ + Field::new("uuid", DataType::FixedSizeBinary(16), false).with_extension_type(Uuid), + ]); + + let parquet_schema = ArrowSchemaConverter::new().convert(&arrow_schema)?; + + assert_eq!( + parquet_schema.column(0).logical_type_ref(), + Some(&LogicalType::Uuid) + ); + + let arrow_schema = parquet_to_arrow_schema(&parquet_schema, None)?; + assert_eq!(arrow_schema.field(0).try_extension_type::()?, Uuid); + + Ok(()) + } + + #[test] + #[cfg(feature = "arrow_canonical_extension_types")] + fn arrow_json_to_parquet_json() -> Result<()> { + use arrow_schema::extension::Json; + let arrow_schema = Schema::new(vec![ + Field::new("json", DataType::Utf8, false).with_extension_type(Json::default()), + ]); + + let parquet_schema = ArrowSchemaConverter::new().convert(&arrow_schema)?; + + assert_eq!( + parquet_schema.column(0).logical_type_ref(), + Some(&LogicalType::Json) + ); + + let arrow_schema = parquet_to_arrow_schema(&parquet_schema, None)?; + assert_eq!( + arrow_schema.field(0).try_extension_type::()?, + Json::default() + ); + + Ok(()) + } + + #[test] + fn test_parquet_to_arrow_field_levels_with_virtual_rejects_non_virtual() { + let message_type = " + message test_schema { + REQUIRED INT32 id; + } + "; + let parquet_schema = Arc::new(parse_message_type(message_type).unwrap()); + let descriptor = SchemaDescriptor::new(parquet_schema); + + // Try to pass a regular field (not a virtual column) + let regular_field = Arc::new(Field::new("regular_column", DataType::Int64, false)); + let result = parquet_to_arrow_field_levels_with_virtual( + &descriptor, + ProjectionMask::all(), + None, + &[regular_field], + ); + + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("is not a virtual column") + ); + } +} diff --git a/vendor/parquet/src/arrow/schema/primitive.rs b/vendor/parquet/src/arrow/schema/primitive.rs new file mode 100644 index 0000000..c708853 --- /dev/null +++ b/vendor/parquet/src/arrow/schema/primitive.rs @@ -0,0 +1,460 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::basic::{ + ConvertedType, IntType, LogicalType, TimeUnit as ParquetTimeUnit, Type as PhysicalType, +}; +use crate::errors::{ParquetError, Result}; +use crate::schema::types::{BasicTypeInfo, Type}; +use arrow_schema::{DECIMAL128_MAX_PRECISION, DataType, IntervalUnit, TimeUnit}; + +/// Converts [`Type`] to [`DataType`] with an optional `arrow_type_hint` +/// provided by the arrow schema +/// +/// Note: the values embedded in the schema are advisory, +pub fn convert_primitive( + parquet_type: &Type, + arrow_type_hint: Option, +) -> Result { + let physical_type = from_parquet(parquet_type)?; + Ok(match arrow_type_hint { + Some(hint) => apply_hint(physical_type, hint), + None => physical_type, + }) +} + +/// Uses an type hint from the embedded arrow schema to aid in faithfully +/// reproducing the data as it was written into parquet +fn apply_hint(parquet: DataType, hint: DataType) -> DataType { + match (&parquet, &hint) { + // Not all time units can be represented as LogicalType / ConvertedType + (DataType::Int32 | DataType::Int64, DataType::Timestamp(_, _)) => hint, + (DataType::Int32, DataType::Time32(_)) => hint, + (DataType::Int64, DataType::Time64(_)) => hint, + (DataType::Int64, DataType::Duration(_)) => hint, + + // Date64 doesn't have a corresponding LogicalType / ConvertedType + (DataType::Int64, DataType::Date64) => hint, + + // Coerce Date32 back to Date64 (#1666) + (DataType::Date32, DataType::Date64) => hint, + + // Timestamps of the same resolution can be converted to a a different timezone. + (DataType::Timestamp(p, _), DataType::Timestamp(h, Some(_))) if p == h => hint, + + // INT96 default to Timestamp(TimeUnit::Nanosecond, None) (see from_parquet below). + // Allow different resolutions to support larger date ranges. + ( + DataType::Timestamp(TimeUnit::Nanosecond, None), + DataType::Timestamp(TimeUnit::Second, _), + ) => hint, + ( + DataType::Timestamp(TimeUnit::Nanosecond, None), + DataType::Timestamp(TimeUnit::Millisecond, _), + ) => hint, + ( + DataType::Timestamp(TimeUnit::Nanosecond, None), + DataType::Timestamp(TimeUnit::Microsecond, _), + ) => hint, + + // Determine offset size + (DataType::Utf8, DataType::LargeUtf8) => hint, + (DataType::Binary, DataType::LargeBinary) => hint, + + // Read as Utf8 + (DataType::Binary, DataType::Utf8) => hint, + (DataType::Binary, DataType::LargeUtf8) => hint, + (DataType::Binary, DataType::Utf8View) => hint, + + // Determine view type + (DataType::Utf8, DataType::Utf8View) => hint, + (DataType::Binary, DataType::BinaryView) => hint, + + // Determine interval time unit (#1666) + (DataType::Interval(_), DataType::Interval(_)) => hint, + + // Promote to Decimal256 or narrow to Decimal32 or Decimal64 + (DataType::Decimal128(_, _), DataType::Decimal32(_, _)) => hint, + (DataType::Decimal128(_, _), DataType::Decimal64(_, _)) => hint, + (DataType::Decimal128(_, _), DataType::Decimal256(_, _)) => hint, + + // Potentially preserve dictionary encoding + (_, DataType::Dictionary(_, value)) => { + // Apply hint to inner type + let hinted = apply_hint(parquet, value.as_ref().clone()); + + // If matches dictionary value - preserve dictionary + // otherwise use hinted inner type + match &hinted == value.as_ref() { + true => hint, + false => hinted, + } + } + _ => parquet, + } +} + +fn from_parquet(parquet_type: &Type) -> Result { + match parquet_type { + Type::PrimitiveType { + physical_type, + basic_info, + type_length, + scale, + precision, + .. + } => match basic_info.logical_type_ref() { + // Any physical type can have the UNKNOWN logical type annotation. Check for that first. + // https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#unknown-always-null + Some(&LogicalType::Unknown) => Ok(DataType::Null), + _ => match physical_type { + PhysicalType::BOOLEAN => Ok(DataType::Boolean), + PhysicalType::INT32 => from_int32(basic_info, *scale, *precision), + PhysicalType::INT64 => from_int64(basic_info, *scale, *precision), + PhysicalType::INT96 => Ok(DataType::Timestamp(TimeUnit::Nanosecond, None)), + PhysicalType::FLOAT => Ok(DataType::Float32), + PhysicalType::DOUBLE => Ok(DataType::Float64), + PhysicalType::BYTE_ARRAY => from_byte_array(basic_info, *precision, *scale), + PhysicalType::FIXED_LEN_BYTE_ARRAY => { + from_fixed_len_byte_array(basic_info, *scale, *precision, *type_length) + } + }, + }, + Type::GroupType { .. } => unreachable!(), + } +} + +fn decimal_type(scale: i32, precision: i32) -> Result { + if precision <= DECIMAL128_MAX_PRECISION as i32 { + decimal_128_type(scale, precision) + } else { + decimal_256_type(scale, precision) + } +} + +fn decimal_128_type(scale: i32, precision: i32) -> Result { + let scale = scale + .try_into() + .map_err(|_| arrow_err!("scale cannot be negative: {}", scale))?; + + let precision = precision + .try_into() + .map_err(|_| arrow_err!("precision cannot be negative: {}", precision))?; + + Ok(DataType::Decimal128(precision, scale)) +} + +fn decimal_256_type(scale: i32, precision: i32) -> Result { + let scale = scale + .try_into() + .map_err(|_| arrow_err!("scale cannot be negative: {}", scale))?; + + let precision = precision + .try_into() + .map_err(|_| arrow_err!("precision cannot be negative: {}", precision))?; + + Ok(DataType::Decimal256(precision, scale)) +} + +#[allow(clippy::manual_range_contains)] +fn check_decimal_length(type_length: i32) -> Result<()> { + if type_length < 1 || type_length > 32 { + return Err(ParquetError::General(format!( + "DECIMAL must be a Fixed Length Byte Array with length 1 to 32, got {type_length}" + ))); + } + Ok(()) +} + +fn from_int32(info: &BasicTypeInfo, scale: i32, precision: i32) -> Result { + match (info.logical_type_ref(), info.converted_type()) { + (None, ConvertedType::NONE) => Ok(DataType::Int32), + (Some(ref t @ LogicalType::Integer(int)), _) => match (int.bit_width, int.is_signed) { + (8, true) => Ok(DataType::Int8), + (16, true) => Ok(DataType::Int16), + (32, true) => Ok(DataType::Int32), + (8, false) => Ok(DataType::UInt8), + (16, false) => Ok(DataType::UInt16), + (32, false) => Ok(DataType::UInt32), + _ => Err(arrow_err!("Cannot create INT32 physical type from {:?}", t)), + }, + (Some(LogicalType::Decimal(decimal)), _) => { + decimal_128_type(decimal.scale, decimal.precision) + } + (Some(LogicalType::Date), _) => Ok(DataType::Date32), + (Some(LogicalType::Time(time)), _) => match time.unit { + ParquetTimeUnit::MILLIS => Ok(DataType::Time32(TimeUnit::Millisecond)), + _ => Err(arrow_err!( + "Cannot create INT32 physical type from {:?}", + time.unit + )), + }, + (None, ConvertedType::UINT_8) => Ok(DataType::UInt8), + (None, ConvertedType::UINT_16) => Ok(DataType::UInt16), + (None, ConvertedType::UINT_32) => Ok(DataType::UInt32), + (None, ConvertedType::INT_8) => Ok(DataType::Int8), + (None, ConvertedType::INT_16) => Ok(DataType::Int16), + (None, ConvertedType::INT_32) => Ok(DataType::Int32), + (None, ConvertedType::DATE) => Ok(DataType::Date32), + (None, ConvertedType::TIME_MILLIS) => Ok(DataType::Time32(TimeUnit::Millisecond)), + (None, ConvertedType::DECIMAL) => decimal_128_type(scale, precision), + (logical, converted) => Err(arrow_err!( + "Unable to convert parquet INT32 logical type {:?} or converted type {}", + logical, + converted + )), + } +} + +fn from_int64(info: &BasicTypeInfo, scale: i32, precision: i32) -> Result { + match (info.logical_type_ref(), info.converted_type()) { + (None, ConvertedType::NONE) => Ok(DataType::Int64), + ( + Some(LogicalType::Integer(IntType { + bit_width: 64, + is_signed, + })), + _, + ) => match is_signed { + true => Ok(DataType::Int64), + false => Ok(DataType::UInt64), + }, + (Some(LogicalType::Time(time)), _) => match time.unit { + ParquetTimeUnit::MILLIS => { + Err(arrow_err!("Cannot create INT64 from MILLIS time unit",)) + } + ParquetTimeUnit::MICROS => Ok(DataType::Time64(TimeUnit::Microsecond)), + ParquetTimeUnit::NANOS => Ok(DataType::Time64(TimeUnit::Nanosecond)), + }, + (Some(LogicalType::Timestamp(timestamp)), _) => Ok(DataType::Timestamp( + match timestamp.unit { + ParquetTimeUnit::MILLIS => TimeUnit::Millisecond, + ParquetTimeUnit::MICROS => TimeUnit::Microsecond, + ParquetTimeUnit::NANOS => TimeUnit::Nanosecond, + }, + if timestamp.is_adjusted_to_u_t_c { + Some("UTC".into()) + } else { + None + }, + )), + (None, ConvertedType::INT_64) => Ok(DataType::Int64), + (None, ConvertedType::UINT_64) => Ok(DataType::UInt64), + (None, ConvertedType::TIME_MICROS) => Ok(DataType::Time64(TimeUnit::Microsecond)), + (None, ConvertedType::TIMESTAMP_MILLIS) => Ok(DataType::Timestamp( + TimeUnit::Millisecond, + Some("UTC".into()), + )), + (None, ConvertedType::TIMESTAMP_MICROS) => Ok(DataType::Timestamp( + TimeUnit::Microsecond, + Some("UTC".into()), + )), + (Some(LogicalType::Decimal(dec)), _) => decimal_128_type(dec.scale, dec.precision), + (None, ConvertedType::DECIMAL) => decimal_128_type(scale, precision), + (logical, converted) => Err(arrow_err!( + "Unable to convert parquet INT64 logical type {:?} or converted type {}", + logical, + converted + )), + } +} + +fn from_byte_array(info: &BasicTypeInfo, precision: i32, scale: i32) -> Result { + match (info.logical_type_ref(), info.converted_type()) { + (Some(LogicalType::String), _) => Ok(DataType::Utf8), + (Some(LogicalType::Json), _) => Ok(DataType::Utf8), + (Some(LogicalType::Bson), _) => Ok(DataType::Binary), + (Some(LogicalType::Enum), _) => Ok(DataType::Binary), + (Some(LogicalType::Geometry { .. }), _) => Ok(DataType::Binary), + (Some(LogicalType::Geography { .. }), _) => Ok(DataType::Binary), + (Some(LogicalType::_Unknown { .. }), _) => Ok(DataType::Binary), + (None, ConvertedType::NONE) => Ok(DataType::Binary), + (None, ConvertedType::JSON) => Ok(DataType::Utf8), + (None, ConvertedType::BSON) => Ok(DataType::Binary), + (None, ConvertedType::ENUM) => Ok(DataType::Binary), + (None, ConvertedType::UTF8) => Ok(DataType::Utf8), + (Some(LogicalType::Decimal(decimal)), _) => decimal_type(decimal.scale, decimal.precision), + (None, ConvertedType::DECIMAL) => decimal_type(scale, precision), + (logical, converted) => Err(arrow_err!( + "Unable to convert parquet BYTE_ARRAY logical type {:?} or converted type {}", + logical, + converted + )), + } +} + +fn from_fixed_len_byte_array( + info: &BasicTypeInfo, + scale: i32, + precision: i32, + type_length: i32, +) -> Result { + match (info.logical_type_ref(), info.converted_type()) { + (Some(LogicalType::Decimal(decimal)), _) => { + check_decimal_length(type_length)?; + // lengths 1..=16 map to Decimal128, 17..=32 to Decimal256 + if type_length <= 16 { + decimal_128_type(decimal.scale, decimal.precision) + } else { + decimal_256_type(decimal.scale, decimal.precision) + } + } + (None, ConvertedType::DECIMAL) => { + check_decimal_length(type_length)?; + if type_length <= 16 { + decimal_128_type(scale, precision) + } else { + decimal_256_type(scale, precision) + } + } + (None, ConvertedType::INTERVAL) => { + if type_length != 12 { + return Err(ParquetError::General(format!( + "INTERVAL must be a Fixed Length Byte Array with length 12, got {type_length}" + ))); + } + // There is currently no reliable way of determining which IntervalUnit + // to return. Thus without the original Arrow schema, the results + // would be incorrect if all 12 bytes of the interval are populated + Ok(DataType::Interval(IntervalUnit::DayTime)) + } + (Some(LogicalType::Float16), _) => { + if type_length == 2 { + Ok(DataType::Float16) + } else { + Err(ParquetError::General( + "FLOAT16 logical type must be Fixed Length Byte Array with length 2" + .to_string(), + )) + } + } + _ => Ok(DataType::FixedSizeBinary(type_length)), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::basic::{DecimalType, Repetition}; + use crate::schema::types::Type; + + // The PrimitiveTypeBuilder rejects bad lengths at construction. To exercise + // the reader-side checks, build a valid type then overwrite its type_length, + // simulating a schema decoded from a file that wasn't produced via the builder. + fn with_type_length(ty: Type, type_length: i32) -> Type { + match ty { + Type::PrimitiveType { + basic_info, + physical_type, + precision, + scale, + .. + } => Type::PrimitiveType { + basic_info, + physical_type, + type_length, + precision, + scale, + }, + _ => unreachable!(), + } + } + + fn flba_decimal_logical(type_length: i32) -> Type { + let valid = Type::primitive_type_builder("c", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_logical_type(Some(LogicalType::Decimal(DecimalType { + precision: 5, + scale: 2, + }))) + .with_length(16) + .with_precision(5) + .with_scale(2) + .build() + .unwrap(); + with_type_length(valid, type_length) + } + + fn flba_decimal_converted(type_length: i32) -> Type { + let valid = Type::primitive_type_builder("c", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::DECIMAL) + .with_length(16) + .with_precision(5) + .with_scale(2) + .build() + .unwrap(); + with_type_length(valid, type_length) + } + + fn flba_interval(type_length: i32) -> Type { + let valid = Type::primitive_type_builder("c", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::INTERVAL) + .with_length(12) + .build() + .unwrap(); + with_type_length(valid, type_length) + } + + fn assert_err_contains(ty: &Type, needle: &str) { + let err = convert_primitive(ty, None).expect_err("expected an error"); + let msg = err.to_string(); + assert!(msg.contains(needle), "expected {needle:?} in error: {msg}"); + } + + #[test] + fn decimal_logical_rejects_invalid_length() { + for bad in [-1, 0, 33] { + assert_err_contains(&flba_decimal_logical(bad), "DECIMAL"); + } + } + + #[test] + fn decimal_converted_rejects_invalid_length() { + for bad in [-1, 0, 33] { + assert_err_contains(&flba_decimal_converted(bad), "DECIMAL"); + } + } + + #[test] + fn decimal_accepts_valid_lengths() { + assert!(matches!( + convert_primitive(&flba_decimal_logical(16), None).unwrap(), + DataType::Decimal128(_, _) + )); + assert!(matches!( + convert_primitive(&flba_decimal_logical(32), None).unwrap(), + DataType::Decimal256(_, _) + )); + } + + #[test] + fn interval_rejects_wrong_length() { + for bad in [0, 11, 13] { + assert_err_contains(&flba_interval(bad), "INTERVAL"); + } + } + + #[test] + fn interval_accepts_length_12() { + assert_eq!( + convert_primitive(&flba_interval(12), None).unwrap(), + DataType::Interval(IntervalUnit::DayTime) + ); + } +} diff --git a/vendor/parquet/src/arrow/schema/virtual_type.rs b/vendor/parquet/src/arrow/schema/virtual_type.rs new file mode 100644 index 0000000..f66352a --- /dev/null +++ b/vendor/parquet/src/arrow/schema/virtual_type.rs @@ -0,0 +1,263 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! RowNumber +//! + +use arrow_schema::{ArrowError, DataType, Field, extension::ExtensionType}; + +/// Prefix for virtual column extension type names. +macro_rules! VIRTUAL_PREFIX { + () => { + "parquet.virtual." + }; +} + +/// The extension type for row group indices +/// +/// Extension name: `parquet.virtual.row_group_index` +/// +/// This virtual column has storage type `Int64` and uses empty string metadata +#[derive(Debug, Default, Clone, Copy, PartialEq)] +pub struct RowGroupIndex; + +impl ExtensionType for RowGroupIndex { + const NAME: &'static str = concat!(VIRTUAL_PREFIX!(), "row_group_index"); + type Metadata = &'static str; + + fn metadata(&self) -> &Self::Metadata { + &"" + } + + fn serialize_metadata(&self) -> Option { + Some(String::default()) + } + + fn deserialize_metadata(metadata: Option<&str>) -> Result { + if metadata.is_some_and(str::is_empty) { + Ok("") + } else { + Err(ArrowError::InvalidArgumentError( + "Virtual column extension type expects an empty string as metadata".to_owned(), + )) + } + } + + fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> { + match data_type { + DataType::Int64 => Ok(()), + data_type => Err(ArrowError::InvalidArgumentError(format!( + "Virtual column data type mismatch, expected Int64, found {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) + } +} + +/// The extension type for row numbers. +/// +/// Extension name: `parquet.virtual.row_number`. +/// +/// This virtual column has storage type `Int64` and uses empty string metadata. +#[derive(Debug, Default, Clone, Copy, PartialEq)] +pub struct RowNumber; + +impl ExtensionType for RowNumber { + const NAME: &'static str = concat!(VIRTUAL_PREFIX!(), "row_number"); + type Metadata = &'static str; + + fn metadata(&self) -> &Self::Metadata { + &"" + } + + fn serialize_metadata(&self) -> Option { + Some(String::default()) + } + + fn deserialize_metadata(metadata: Option<&str>) -> Result { + if metadata.is_some_and(str::is_empty) { + Ok("") + } else { + Err(ArrowError::InvalidArgumentError( + "Virtual column extension type expects an empty string as metadata".to_owned(), + )) + } + } + + fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> { + match data_type { + DataType::Int64 => Ok(()), + data_type => Err(ArrowError::InvalidArgumentError(format!( + "Virtual column data type mismatch, expected Int64, found {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) + } +} + +/// Returns `true` if the field is a virtual column. +/// +/// Virtual columns have extension type names starting with `parquet.virtual.`. +pub fn is_virtual_column(field: &Field) -> bool { + field + .extension_type_name() + .is_some_and(|name| name.starts_with(VIRTUAL_PREFIX!())) +} + +#[cfg(test)] +mod tests { + use arrow_schema::{ + ArrowError, DataType, Field, + extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}, + }; + + use super::*; + + #[test] + fn row_number_valid() -> Result<(), ArrowError> { + let mut field = Field::new("", DataType::Int64, false); + field.try_with_extension_type(RowNumber)?; + field.try_extension_type::()?; + + Ok(()) + } + + #[test] + #[should_panic(expected = "Extension type name missing")] + fn row_number_missing_name() { + let field = Field::new("", DataType::Int64, false).with_metadata( + [(EXTENSION_TYPE_METADATA_KEY.to_owned(), "".to_owned())] + .into_iter() + .collect(), + ); + field.extension_type::(); + } + + #[test] + #[should_panic(expected = "expected Int64, found Int32")] + fn row_number_invalid_type() { + Field::new("", DataType::Int32, false).with_extension_type(RowNumber); + } + + #[test] + #[should_panic(expected = "Virtual column extension type expects an empty string as metadata")] + fn row_number_missing_metadata() { + let field = Field::new("", DataType::Int64, false).with_metadata( + [( + EXTENSION_TYPE_NAME_KEY.to_owned(), + RowNumber::NAME.to_owned(), + )] + .into_iter() + .collect(), + ); + field.extension_type::(); + } + + #[test] + #[should_panic(expected = "Virtual column extension type expects an empty string as metadata")] + fn row_number_invalid_metadata() { + let field = Field::new("", DataType::Int64, false).with_metadata( + [ + ( + EXTENSION_TYPE_NAME_KEY.to_owned(), + RowNumber::NAME.to_owned(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_owned(), + "non-empty".to_owned(), + ), + ] + .into_iter() + .collect(), + ); + field.extension_type::(); + } + + #[test] + fn row_group_index_valid() -> Result<(), ArrowError> { + let mut field = Field::new("", DataType::Int64, false); + field.try_with_extension_type(RowGroupIndex)?; + field.try_extension_type::()?; + + Ok(()) + } + + #[test] + #[should_panic(expected = "Extension type name missing")] + fn row_group_index_missing_name() { + let field = Field::new("", DataType::Int64, false).with_metadata( + [(EXTENSION_TYPE_METADATA_KEY.to_owned(), "".to_owned())] + .into_iter() + .collect(), + ); + field.extension_type::(); + } + + #[test] + #[should_panic(expected = "expected Int64, found Int32")] + fn row_group_index_invalid_type() { + Field::new("", DataType::Int32, false).with_extension_type(RowGroupIndex); + } + + #[test] + #[should_panic(expected = "Virtual column extension type expects an empty string as metadata")] + fn row_group_index_missing_metadata() { + let field = Field::new("", DataType::Int64, false).with_metadata( + [( + EXTENSION_TYPE_NAME_KEY.to_owned(), + RowGroupIndex::NAME.to_owned(), + )] + .into_iter() + .collect(), + ); + field.extension_type::(); + } + + #[test] + #[should_panic(expected = "Virtual column extension type expects an empty string as metadata")] + fn row_group_index_invalid_metadata() { + let field = Field::new("", DataType::Int64, false).with_metadata( + [ + ( + EXTENSION_TYPE_NAME_KEY.to_owned(), + RowGroupIndex::NAME.to_owned(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_owned(), + "non-empty".to_owned(), + ), + ] + .into_iter() + .collect(), + ); + field.extension_type::(); + } +} diff --git a/vendor/parquet/src/basic.rs b/vendor/parquet/src/basic.rs new file mode 100644 index 0000000..b4f18f3 --- /dev/null +++ b/vendor/parquet/src/basic.rs @@ -0,0 +1,2346 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Contains Rust mappings for Thrift definition. This module contains only mappings for thrift +//! enums and unions. Thrift structs are handled elsewhere. +//! Refer to [`parquet.thrift`](https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift) +//! file to see raw definitions. + +use std::io::Write; +use std::str::FromStr; +use std::{fmt, str}; + +pub use crate::compression::{BrotliLevel, GzipLevel, ZstdLevel}; +use crate::file::metadata::HeapSize; +use crate::parquet_thrift::{ + ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol, ThriftCompactOutputProtocol, + WriteThrift, WriteThriftField, validate_list_type, +}; +use crate::{ + thrift_enum, thrift_struct, thrift_union_all_empty, thrift_union_with_unknown, + write_thrift_field, +}; + +use crate::errors::{ParquetError, Result}; + +// ---------------------------------------------------------------------- +// Types from the Thrift definition + +// ---------------------------------------------------------------------- +// Mirrors thrift enum `Type` + +thrift_enum!( +/// Types supported by Parquet. +/// +/// These physical types are intended to be used in combination with the encodings to +/// control the on disk storage format. +/// For example INT16 is not included as a type since a good encoding of INT32 +/// would handle this. +enum Type { + BOOLEAN = 0; + INT32 = 1; + INT64 = 2; + INT96 = 3; // deprecated, only used by legacy implementations. + FLOAT = 4; + DOUBLE = 5; + BYTE_ARRAY = 6; + FIXED_LEN_BYTE_ARRAY = 7; +} +); + +// ---------------------------------------------------------------------- +// Mirrors thrift enum `ConvertedType` + +// TODO(ets): Adding the `NONE` variant to this enum is a bit awkward. We should +// look into removing it and using `Option` instead. +thrift_enum!( +/// Common types (converted types) used by frameworks when using Parquet. +/// +/// This helps map between types in those frameworks to the base types in Parquet. +/// This is only metadata and not needed to read or write the data. +/// +/// This struct was renamed from `LogicalType` in version 4.0.0. +/// If targeting Parquet format 2.4.0 or above, please use [LogicalType] instead. +enum ConvertedType { + /// Not defined in the spec, used internally to indicate no type conversion + NONE = -1; + + /// A BYTE_ARRAY actually contains UTF8 encoded chars. + UTF8 = 0; + + /// A map is converted as an optional field containing a repeated key/value pair. + MAP = 1; + + /// A key/value pair is converted into a group of two fields. + MAP_KEY_VALUE = 2; + + /// A list is converted into an optional field containing a repeated field for its + /// values. + LIST = 3; + + /// An enum is converted into a BYTE_ARRAY field + ENUM = 4; + + /// A decimal value. + /// + /// This may be used to annotate BYTE_ARRAY or FIXED_LEN_BYTE_ARRAY primitive + /// types. The underlying byte array stores the unscaled value encoded as two's + /// complement using big-endian byte order (the most significant byte is the + /// zeroth element). The value of the decimal is the value * 10^{-scale}. + /// + /// This must be accompanied by a (maximum) precision and a scale in the + /// SchemaElement. The precision specifies the number of digits in the decimal + /// and the scale stores the location of the decimal point. For example 1.23 + /// would have precision 3 (3 total digits) and scale 2 (the decimal point is + /// 2 digits over). + DECIMAL = 5; + + /// A date stored as days since Unix epoch, encoded as the INT32 physical type. + DATE = 6; + + /// The total number of milliseconds since midnight. The value is stored as an INT32 + /// physical type. + TIME_MILLIS = 7; + + /// The total number of microseconds since midnight. The value is stored as an INT64 + /// physical type. + TIME_MICROS = 8; + + /// Date and time recorded as milliseconds since the Unix epoch. + /// Recorded as a physical type of INT64. + TIMESTAMP_MILLIS = 9; + + /// Date and time recorded as microseconds since the Unix epoch. + /// The value is stored as an INT64 physical type. + TIMESTAMP_MICROS = 10; + + /// An unsigned 8 bit integer value stored as INT32 physical type. + UINT_8 = 11; + + /// An unsigned 16 bit integer value stored as INT32 physical type. + UINT_16 = 12; + + /// An unsigned 32 bit integer value stored as INT32 physical type. + UINT_32 = 13; + + /// An unsigned 64 bit integer value stored as INT64 physical type. + UINT_64 = 14; + + /// A signed 8 bit integer value stored as INT32 physical type. + INT_8 = 15; + + /// A signed 16 bit integer value stored as INT32 physical type. + INT_16 = 16; + + /// A signed 32 bit integer value stored as INT32 physical type. + INT_32 = 17; + + /// A signed 64 bit integer value stored as INT64 physical type. + INT_64 = 18; + + /// A JSON document embedded within a single UTF8 column. + JSON = 19; + + /// A BSON document embedded within a single BINARY column. + BSON = 20; + + /// An interval of time + /// + /// This type annotates data stored as a FIXED_LEN_BYTE_ARRAY of length 12. + /// This data is composed of three separate little endian unsigned integers. + /// Each stores a component of a duration of time. The first integer identifies + /// the number of months associated with the duration, the second identifies + /// the number of days associated with the duration and the third identifies + /// the number of milliseconds associated with the provided duration. + /// This duration of time is independent of any particular timezone or date. + INTERVAL = 21; +} +); + +// ---------------------------------------------------------------------- +// Mirrors thrift union `TimeUnit` + +thrift_union_all_empty!( +/// Time unit for `Time` and `Timestamp` logical types. +union TimeUnit { + 1: MilliSeconds MILLIS + 2: MicroSeconds MICROS + 3: NanoSeconds NANOS +} +); + +// ---------------------------------------------------------------------- +// Mirrors thrift union `LogicalType` + +thrift_struct!( +pub struct DecimalType { + /// The number of digits in the decimal. + 1: required i32 scale + /// The location of the decimal point. + 2: required i32 precision +} +); + +thrift_struct!( +pub struct TimestampType { + /// Whether the timestamp is adjusted to UTC. + 1: required bool is_adjusted_to_u_t_c + /// The unit of time. + 2: required TimeUnit unit +} +); + +/// Identical to [`TimestampType`] +pub use TimestampType as TimeType; + +thrift_struct!( +pub struct IntType { + /// The number of bits in the integer. + 1: required i8 bit_width + /// Whether the integer is signed. + 2: required bool is_signed +} +); + +thrift_struct!( +pub struct VariantType { + /// The version of the variant specification that the variant was + /// written with. + 1: optional i8 specification_version +} +); + +thrift_struct!( +pub struct GeometryType { + /// A custom CRS. If unset the CRS `OGC:CRS84` should be used, which means that the geometries + /// must be stored in longitude, latitude based on the WGS84 datum. + 1: optional string crs; +} +); + +thrift_struct!( +pub struct GeographyType { + /// A custom CRS. If unset the CRS `OGC:CRS84` should be used. + 1: optional string crs; + /// An optional algorithm can be set to correctly interpret edges interpolation + /// of the geometries. If unset, the `SPHERICAL` algorithm should be used. + 2: optional EdgeInterpolationAlgorithm algorithm; +} +); + +impl GeographyType { + /// Accessor for the `GeographyType::algorithm` field. If this field is not set, this + /// function returns the default value (currently [`EdgeInterpolationAlgorithm::SPHERICAL`] + /// per the Parquet [specification]). + /// + /// [specification]: https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#geography + pub fn algorithm(&self) -> Option { + self.algorithm.or(Some(Default::default())) + } +} + +thrift_union_with_unknown!( +/// Logical types used by version 2.4.0+ of the Parquet format. +/// +/// This is an *entirely new* struct as of version +/// 4.0.0. The struct previously named `LogicalType` was renamed to +/// [`ConvertedType`]. Please see the README.md for more details. +union LogicalType { + /// A UTF8 encoded string. + 1: String + /// A map of key-value pairs. + 2: Map + /// A list of elements. + 3: List + /// A set of predefined values. + 4: Enum + /// A decimal value with a specified scale and precision. + 5: (DecimalType) Decimal + /// A date stored as days since Unix epoch. + 6: Date + /// A time stored as [`TimeUnit`] since midnight. + 7: (TimeType) Time + /// A timestamp stored as [`TimeUnit`] since Unix epoch. + 8: (TimestampType) Timestamp + // 9: reserved for INTERVAL + /// An integer with a specified bit width and signedness. + 10: (IntType) Integer + /// An unknown logical type. + 11: Unknown + /// A JSON document. + 12: Json + /// A BSON document. + 13: Bson + /// A UUID. + 14: Uuid + /// A 16-bit floating point number. + 15: Float16 + /// A Variant value. + 16: (VariantType) Variant + /// A geospatial feature in the Well-Known Binary (WKB) format with linear/planar edges interpolation. + 17: (GeometryType) Geometry + /// A geospatial feature in the WKB format with an explicit (non-linear/non-planar) edges interpolation. + 18: (GeographyType) Geography +} +); + +impl LogicalType { + /// Create a [`LogicalType::Integer`] variant with the given `bit_width` and `is_signed` + pub fn integer(bit_width: i8, is_signed: bool) -> Self { + Self::Integer(IntType { + bit_width, + is_signed, + }) + } + + /// Create a [`LogicalType::Decimal`] variant with the given `scale` and `precision` + pub fn decimal(scale: i32, precision: i32) -> Self { + Self::Decimal(DecimalType { scale, precision }) + } + + /// Create a [`LogicalType::Time`] variant with the given `is_adjusted_to_u_t_c` and `unit` + pub fn time(is_adjusted_to_u_t_c: bool, unit: TimeUnit) -> Self { + Self::Time(TimeType { + is_adjusted_to_u_t_c, + unit, + }) + } + + /// Create a [`LogicalType::Timestamp`] variant with the given `is_adjusted_to_u_t_c` and `unit` + pub fn timestamp(is_adjusted_to_u_t_c: bool, unit: TimeUnit) -> Self { + Self::Timestamp(TimestampType { + is_adjusted_to_u_t_c, + unit, + }) + } + + /// Create a [`LogicalType::Variant`] variant with the given `specification_version` + pub fn variant(specification_version: Option) -> Self { + Self::Variant(VariantType { + specification_version, + }) + } + + /// Create a [`LogicalType::Geometry`] variant with the given `crs` + pub fn geometry(crs: Option) -> Self { + Self::Geometry(GeometryType { crs }) + } + + /// Create a [`LogicalType::Geography`] variant with the given `crs` and `algorithm` + pub fn geography(crs: Option, algorithm: Option) -> Self { + Self::Geography(GeographyType { crs, algorithm }) + } +} + +// ---------------------------------------------------------------------- +// Mirrors thrift enum `FieldRepetitionType` +// + +thrift_enum!( +/// Representation of field types in schema. +enum FieldRepetitionType { + /// This field is required (can not be null) and each row has exactly 1 value. + REQUIRED = 0; + /// The field is optional (can be null) and each row has 0 or 1 values. + OPTIONAL = 1; + /// The field is repeated and can contain 0 or more values. + REPEATED = 2; +} +); + +/// Type alias for thrift `FieldRepetitionType` +pub type Repetition = FieldRepetitionType; + +// ---------------------------------------------------------------------- +// Mirrors thrift enum `Encoding` + +thrift_enum!( +/// Encodings supported by Parquet. +/// +/// Not all encodings are valid for all types. These enums are also used to specify the +/// encoding of definition and repetition levels. +/// +/// By default this crate uses [Encoding::PLAIN], [Encoding::RLE], and [Encoding::RLE_DICTIONARY]. +/// These provide very good encode and decode performance, whilst yielding reasonable storage +/// efficiency and being supported by all major parquet readers. +/// +/// The delta encodings are also supported and will be used if a newer [WriterVersion] is +/// configured, however, it should be noted that these sacrifice encode and decode performance for +/// improved storage efficiency. This performance regression is particularly pronounced in the case +/// of record skipping as occurs during predicate push-down. It is recommended users assess the +/// performance impact when evaluating these encodings. +/// +/// [WriterVersion]: crate::file::properties::WriterVersion +enum Encoding { + /// Default encoding. + /// - BOOLEAN - 1 bit per value. 0 is false; 1 is true. + /// - INT32 - 4 bytes per value. Stored as little-endian. + /// - INT64 - 8 bytes per value. Stored as little-endian. + /// - FLOAT - 4 bytes per value. IEEE. Stored as little-endian. + /// - DOUBLE - 8 bytes per value. IEEE. Stored as little-endian. + /// - BYTE_ARRAY - 4 byte length stored as little endian, followed by bytes. + /// - FIXED_LEN_BYTE_ARRAY - Just the bytes. + PLAIN = 0; + // GROUP_VAR_INT = 1; + /// **Deprecated** dictionary encoding. + /// + /// The values in the dictionary are encoded using PLAIN encoding. + /// Since it is deprecated, RLE_DICTIONARY encoding is used for a data page, and + /// PLAIN encoding is used for dictionary page. + PLAIN_DICTIONARY = 2; + /// Group packed run length encoding. + /// + /// Usable for definition/repetition levels encoding and boolean values. + RLE = 3; + /// **Deprecated** Bit-packed encoding. + /// + /// This can only be used if the data has a known max width. + /// Usable for definition/repetition levels encoding. + /// + /// There are compatibility issues with files using this encoding. + /// The parquet standard specifies the bits to be packed starting from the + /// most-significant bit, several implementations do not follow this bit order. + /// Several other implementations also have issues reading this encoding + /// because of incorrect assumptions about the length of the encoded data. + /// + /// The RLE/bit-packing hybrid is more cpu and memory efficient and should be used instead. + #[deprecated( + since = "51.0.0", + note = "Please see documentation for compatibility issues and use the RLE/bit-packing hybrid encoding instead" + )] + BIT_PACKED = 4; + /// Delta encoding for integers, either INT32 or INT64. + /// + /// Works best on sorted data. + DELTA_BINARY_PACKED = 5; + /// Encoding for byte arrays to separate the length values and the data. + /// + /// The lengths are encoded using DELTA_BINARY_PACKED encoding. + DELTA_LENGTH_BYTE_ARRAY = 6; + /// Incremental encoding for byte arrays. + /// + /// Prefix lengths are encoded using DELTA_BINARY_PACKED encoding. + /// Suffixes are stored using DELTA_LENGTH_BYTE_ARRAY encoding. + DELTA_BYTE_ARRAY = 7; + /// Dictionary encoding. + /// + /// The ids are encoded using the RLE encoding. + RLE_DICTIONARY = 8; + /// Encoding for fixed-width data. + /// + /// K byte-streams are created where K is the size in bytes of the data type. + /// The individual bytes of a value are scattered to the corresponding stream and + /// the streams are concatenated. + /// This itself does not reduce the size of the data but can lead to better compression + /// afterwards. Note that the use of this encoding with FIXED_LEN_BYTE_ARRAY(N) data may + /// perform poorly for large values of N. + BYTE_STREAM_SPLIT = 9; +} +); + +impl FromStr for Encoding { + type Err = ParquetError; + + fn from_str(s: &str) -> Result { + match s { + "PLAIN" | "plain" => Ok(Encoding::PLAIN), + "PLAIN_DICTIONARY" | "plain_dictionary" => Ok(Encoding::PLAIN_DICTIONARY), + "RLE" | "rle" => Ok(Encoding::RLE), + #[allow(deprecated)] + "BIT_PACKED" | "bit_packed" => Ok(Encoding::BIT_PACKED), + "DELTA_BINARY_PACKED" | "delta_binary_packed" => Ok(Encoding::DELTA_BINARY_PACKED), + "DELTA_LENGTH_BYTE_ARRAY" | "delta_length_byte_array" => { + Ok(Encoding::DELTA_LENGTH_BYTE_ARRAY) + } + "DELTA_BYTE_ARRAY" | "delta_byte_array" => Ok(Encoding::DELTA_BYTE_ARRAY), + "RLE_DICTIONARY" | "rle_dictionary" => Ok(Encoding::RLE_DICTIONARY), + "BYTE_STREAM_SPLIT" | "byte_stream_split" => Ok(Encoding::BYTE_STREAM_SPLIT), + _ => Err(general_err!("unknown encoding: {}", s)), + } + } +} + +/// A bitmask representing the [`Encoding`]s employed while encoding a Parquet column chunk. +/// +/// The Parquet [`ColumnMetaData`] struct contains an array that indicates what encodings were +/// used when writing that column chunk. For memory and performance reasons, this crate reduces +/// that array to bitmask, where each bit position represents a different [`Encoding`]. This +/// struct contains that bitmask, and provides methods to interact with the data. +/// +/// # Example +/// ```no_run +/// # use parquet::file::metadata::ParquetMetaDataReader; +/// # use parquet::basic::Encoding; +/// # fn open_parquet_file(path: &str) -> std::fs::File { unimplemented!(); } +/// // read parquet metadata from a file +/// let file = open_parquet_file("some_path.parquet"); +/// let mut reader = ParquetMetaDataReader::new(); +/// reader.try_parse(&file).unwrap(); +/// let metadata = reader.finish().unwrap(); +/// +/// // find the encodings used by the first column chunk in the first row group +/// let col_meta = metadata.row_group(0).column(0); +/// let encodings = col_meta.encodings_mask(); +/// +/// // check to see if a particular encoding was used +/// let used_rle = encodings.is_set(Encoding::RLE); +/// +/// // check to see if all of a set of encodings were used +/// let used_all = encodings.all_set([Encoding::RLE, Encoding::PLAIN].iter()); +/// +/// // convert mask to a Vec +/// let encodings_vec = encodings.encodings().collect::>(); +/// ``` +/// +/// [`ColumnMetaData`]: https://github.com/apache/parquet-format/blob/9fd57b59e0ce1a82a69237dcf8977d3e72a2965d/src/main/thrift/parquet.thrift#L875 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct EncodingMask(i32); + +impl EncodingMask { + /// Highest valued discriminant in the [`Encoding`] enum + const MAX_ENCODING: i32 = Encoding::MAX_DISCRIMINANT; + /// A mask consisting of unused bit positions, used for validation. This includes the never + /// used GROUP_VAR_INT encoding value of `1`. + const ALLOWED_MASK: u32 = + !(1u32 << (EncodingMask::MAX_ENCODING as u32 + 1)).wrapping_sub(1) | 1 << 1; + + /// Attempt to create a new `EncodingMask` from an integer. + /// + /// This will return an error if a bit outside the allowable range is set. + pub fn try_new(val: i32) -> Result { + if val as u32 & Self::ALLOWED_MASK != 0 { + return Err(general_err!("Attempt to create invalid mask: 0x{:x}", val)); + } + Ok(Self(val)) + } + + /// Return an integer representation of this `EncodingMask`. + pub fn as_i32(&self) -> i32 { + self.0 + } + + /// Create a new `EncodingMask` from a collection of [`Encoding`]s. + pub fn new_from_encodings<'a>(encodings: impl Iterator) -> Self { + let mut mask = 0; + for &e in encodings { + mask |= 1 << (e as i32); + } + Self(mask) + } + + /// Mark the given [`Encoding`] as present in this mask. + pub fn insert(&mut self, val: Encoding) { + self.0 |= 1 << (val as i32); + } + + /// Test if a given [`Encoding`] is present in this mask. + pub fn is_set(&self, val: Encoding) -> bool { + self.0 & (1 << (val as i32)) != 0 + } + + /// Test if this mask has only the bit for the given [`Encoding`] set. + pub fn is_only(&self, val: Encoding) -> bool { + self.0 == (1 << (val as i32)) + } + + /// Test if all [`Encoding`]s in a given set are present in this mask. + pub fn all_set<'a>(&self, mut encodings: impl Iterator) -> bool { + encodings.all(|&e| self.is_set(e)) + } + + /// Return an iterator over all [`Encoding`]s present in this mask. + pub fn encodings(&self) -> impl Iterator { + Self::mask_to_encodings_iter(self.0) + } + + fn mask_to_encodings_iter(mask: i32) -> impl Iterator { + (0..=Self::MAX_ENCODING) + .filter(move |i| mask & (1 << i) != 0) + .map(i32_to_encoding) + } +} + +impl HeapSize for EncodingMask { + fn heap_size(&self) -> usize { + 0 // no heap allocations + } +} + +impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for EncodingMask { + fn read_thrift(prot: &mut R) -> Result { + let mut mask = 0; + + // This reads a Thrift `list` and turns it into a bitmask + let list_ident = prot.read_list_begin()?; + // check for enum (encoded as I32) + validate_list_type(ElementType::I32, &list_ident)?; + for _ in 0..list_ident.size { + let val = Encoding::read_thrift(prot)?; + mask |= 1 << val as i32; + } + Ok(Self(mask)) + } +} + +#[allow(deprecated)] +fn i32_to_encoding(val: i32) -> Encoding { + match val { + 0 => Encoding::PLAIN, + 2 => Encoding::PLAIN_DICTIONARY, + 3 => Encoding::RLE, + 4 => Encoding::BIT_PACKED, + 5 => Encoding::DELTA_BINARY_PACKED, + 6 => Encoding::DELTA_LENGTH_BYTE_ARRAY, + 7 => Encoding::DELTA_BYTE_ARRAY, + 8 => Encoding::RLE_DICTIONARY, + 9 => Encoding::BYTE_STREAM_SPLIT, + _ => panic!("Impossible encoding {val}"), + } +} + +// ---------------------------------------------------------------------- +// Mirrors thrift enum `CompressionCodec` + +thrift_enum!( +/// Supported compression algorithms. +/// +/// Codecs added in format version X.Y can be read by readers based on X.Y and later. +/// Codec support may vary between readers based on the format version and +/// libraries available at runtime. +/// +/// See [Compression.md] for a detailed specification of these algorithms. +/// +/// [Compression.md]: https://github.com/apache/parquet-format/blob/master/Compression.md +enum CompressionCodec { + UNCOMPRESSED = 0; + SNAPPY = 1; + GZIP = 2; + LZO = 3; + BROTLI = 4; // Added in 2.4 + LZ4 = 5; // DEPRECATED (Added in 2.4) + ZSTD = 6; // Added in 2.4 + LZ4_RAW = 7; // Added in 2.9 +} +); + +// NOTE: This enum likely belongs in file::properties now, but moving it there would be a +// breaking API change, that's probably not worth the pain. If a new codec is added to the +// Parquet specification, or any other breaking changes are made to this enum, this can be +// revisited. + +/// Supported block compression algorithms. +/// +/// Block compression can yield non-trivial improvements to storage efficiency at the expense +/// of potentially significantly worse encode and decode performance. Many applications, +/// especially those making use of high-throughput and low-cost commodity object storage, +/// may find storage efficiency less important than decode throughput, and therefore may +/// wish to not make use of block compression. +/// +/// The writers in this crate default to no block compression for this reason. +/// +/// Applications that do still wish to use block compression, will find [`Compression::ZSTD`] +/// to provide a good balance of compression, performance, and ecosystem support. Alternatively, +/// [`Compression::LZ4_RAW`] provides much faster decompression speeds, at the cost of typically +/// worse compression ratios. However, it is not as widely supported by the ecosystem, with the +/// Hadoop ecosystem historically favoring the non-standard and now deprecated [`Compression::LZ4`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[allow(non_camel_case_types)] +pub enum Compression { + /// No compression. + UNCOMPRESSED, + /// [Snappy compression](https://en.wikipedia.org/wiki/Snappy_(compression)) + SNAPPY, + /// [Gzip compression](https://www.ietf.org/rfc/rfc1952.txt) + GZIP(GzipLevel), + /// [LZO compression](https://en.wikipedia.org/wiki/Lempel%E2%80%93Ziv%E2%80%93Oberhumer) + LZO, + /// [Brotli compression](https://datatracker.ietf.org/doc/html/rfc7932) + BROTLI(BrotliLevel), + /// [LZ4 compression](https://lz4.org/), [(deprecated)](https://issues.apache.org/jira/browse/PARQUET-2032) + LZ4, + /// [ZSTD compression](https://datatracker.ietf.org/doc/html/rfc8878) + ZSTD(ZstdLevel), + /// [LZ4 compression](https://lz4.org/). + LZ4_RAW, +} + +impl From for Compression { + fn from(value: CompressionCodec) -> Self { + match value { + CompressionCodec::UNCOMPRESSED => Compression::UNCOMPRESSED, + CompressionCodec::SNAPPY => Compression::SNAPPY, + CompressionCodec::GZIP => Compression::GZIP(Default::default()), + CompressionCodec::LZO => Compression::LZO, + CompressionCodec::BROTLI => Compression::BROTLI(Default::default()), + CompressionCodec::LZ4 => Compression::LZ4, + CompressionCodec::ZSTD => Compression::ZSTD(Default::default()), + CompressionCodec::LZ4_RAW => Compression::LZ4_RAW, + } + } +} + +impl From for CompressionCodec { + fn from(value: Compression) -> Self { + match value { + Compression::UNCOMPRESSED => CompressionCodec::UNCOMPRESSED, + Compression::SNAPPY => CompressionCodec::SNAPPY, + Compression::GZIP(_) => CompressionCodec::GZIP, + Compression::LZO => CompressionCodec::LZO, + Compression::BROTLI(_) => CompressionCodec::BROTLI, + Compression::LZ4 => CompressionCodec::LZ4, + Compression::ZSTD(_) => CompressionCodec::ZSTD, + Compression::LZ4_RAW => CompressionCodec::LZ4_RAW, + } + } +} + +fn split_compression_string(str_setting: &str) -> Result<(&str, Option), ParquetError> { + let split_setting = str_setting.split_once('('); + + match split_setting { + Some((codec, level_str)) => { + let level = &level_str[..level_str.len() - 1] + .parse::() + .map_err(|_| { + ParquetError::General(format!("invalid compression level: {level_str}")) + })?; + Ok((codec, Some(*level))) + } + None => Ok((str_setting, None)), + } +} + +fn check_level_is_none(level: &Option) -> Result<(), ParquetError> { + if level.is_some() { + return Err(ParquetError::General( + "compression level is not supported".to_string(), + )); + } + + Ok(()) +} + +fn require_level(codec: &str, level: Option) -> Result { + level.ok_or(ParquetError::General(format!( + "{codec} requires a compression level", + ))) +} + +impl FromStr for Compression { + type Err = ParquetError; + + fn from_str(s: &str) -> std::result::Result { + let (codec, level) = split_compression_string(s)?; + + let c = match codec { + "UNCOMPRESSED" | "uncompressed" => { + check_level_is_none(&level)?; + Compression::UNCOMPRESSED + } + "SNAPPY" | "snappy" => { + check_level_is_none(&level)?; + Compression::SNAPPY + } + "GZIP" | "gzip" => { + let level = require_level(codec, level)?; + Compression::GZIP(GzipLevel::try_new(level)?) + } + "LZO" | "lzo" => { + check_level_is_none(&level)?; + Compression::LZO + } + "BROTLI" | "brotli" => { + let level = require_level(codec, level)?; + Compression::BROTLI(BrotliLevel::try_new(level)?) + } + "LZ4" | "lz4" => { + check_level_is_none(&level)?; + Compression::LZ4 + } + "ZSTD" | "zstd" => { + let level = require_level(codec, level)?; + Compression::ZSTD(ZstdLevel::try_new(level as i32)?) + } + "LZ4_RAW" | "lz4_raw" => { + check_level_is_none(&level)?; + Compression::LZ4_RAW + } + _ => { + return Err(ParquetError::General(format!( + "unsupport compression {codec}" + ))); + } + }; + + Ok(c) + } +} + +// ---------------------------------------------------------------------- +// Mirrors thrift enum `PageType` + +thrift_enum!( +/// Available data pages for Parquet file format. +/// Note that some of the page types may not be supported. +enum PageType { + DATA_PAGE = 0; + INDEX_PAGE = 1; + DICTIONARY_PAGE = 2; + DATA_PAGE_V2 = 3; +} +); + +// ---------------------------------------------------------------------- +// Mirrors thrift enum `BoundaryOrder` + +thrift_enum!( +/// Enum to annotate whether lists of min/max elements inside ColumnIndex +/// are ordered and if so, in which direction. +enum BoundaryOrder { + UNORDERED = 0; + ASCENDING = 1; + DESCENDING = 2; +} +); + +// ---------------------------------------------------------------------- +// Mirrors thrift enum `EdgeInterpolationAlgorithm` + +// this is hand coded to allow for the _Unknown variant (allows this to be forward compatible) + +/// Edge interpolation algorithm for [`LogicalType::Geography`] +#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] +#[repr(i32)] +#[derive(Default)] +pub enum EdgeInterpolationAlgorithm { + /// Edges are interpolated as geodesics on a sphere. + #[default] + SPHERICAL = 0, + /// + VINCENTY = 1, + /// Thomas, Paul D. Spheroidal geodesics, reference systems, & local geometry. US Naval Oceanographic Office, 1970 + THOMAS = 2, + /// Thomas, Paul D. Mathematical models for navigation systems. US Naval Oceanographic Office, 1965. + ANDOYER = 3, + /// Karney, Charles FF. "Algorithms for geodesics." Journal of Geodesy 87 (2013): 43-55 + KARNEY = 4, + /// Unknown algorithm + _Unknown(i32), +} + +#[cfg(feature = "geospatial")] +impl EdgeInterpolationAlgorithm { + /// Converts an [`EdgeInterpolationAlgorithm`] into its corresponding algorithm defined by + /// [`parquet_geospatial::WkbEdges`]. + /// + /// This method will only return an Err if the [`EdgeInterpolationAlgorithm`] is the `_Unknown` + /// variant. + pub fn try_as_edges(&self) -> Result { + match &self { + Self::SPHERICAL => Ok(parquet_geospatial::WkbEdges::Spherical), + Self::VINCENTY => Ok(parquet_geospatial::WkbEdges::Vincenty), + Self::THOMAS => Ok(parquet_geospatial::WkbEdges::Thomas), + Self::ANDOYER => Ok(parquet_geospatial::WkbEdges::Andoyer), + Self::KARNEY => Ok(parquet_geospatial::WkbEdges::Karney), + unknown => Err(general_err!( + "Unknown edge interpolation algorithm: {}", + unknown + )), + } + } +} + +impl fmt::Display for EdgeInterpolationAlgorithm { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.write_fmt(format_args!("{0:?}", self)) + } +} + +#[cfg(feature = "geospatial")] +impl From for EdgeInterpolationAlgorithm { + fn from(value: parquet_geospatial::WkbEdges) -> Self { + match value { + parquet_geospatial::WkbEdges::Spherical => Self::SPHERICAL, + parquet_geospatial::WkbEdges::Vincenty => Self::VINCENTY, + parquet_geospatial::WkbEdges::Thomas => Self::THOMAS, + parquet_geospatial::WkbEdges::Andoyer => Self::ANDOYER, + parquet_geospatial::WkbEdges::Karney => Self::KARNEY, + } + } +} + +impl FromStr for EdgeInterpolationAlgorithm { + type Err = ParquetError; + + fn from_str(s: &str) -> Result { + match s.to_ascii_uppercase().as_str() { + "SPHERICAL" => Ok(EdgeInterpolationAlgorithm::SPHERICAL), + "VINCENTY" => Ok(EdgeInterpolationAlgorithm::VINCENTY), + "THOMAS" => Ok(EdgeInterpolationAlgorithm::THOMAS), + "ANDOYER" => Ok(EdgeInterpolationAlgorithm::ANDOYER), + "KARNEY" => Ok(EdgeInterpolationAlgorithm::KARNEY), + unknown => Err(general_err!( + "Unknown edge interpolation algorithm: {}", + unknown + )), + } + } +} + +impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for EdgeInterpolationAlgorithm { + fn read_thrift(prot: &mut R) -> Result { + let val = prot.read_i32()?; + match val { + 0 => Ok(Self::SPHERICAL), + 1 => Ok(Self::VINCENTY), + 2 => Ok(Self::THOMAS), + 3 => Ok(Self::ANDOYER), + 4 => Ok(Self::KARNEY), + _ => Ok(Self::_Unknown(val)), + } + } +} + +impl WriteThrift for EdgeInterpolationAlgorithm { + const ELEMENT_TYPE: ElementType = ElementType::I32; + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + let val: i32 = match *self { + Self::SPHERICAL => 0, + Self::VINCENTY => 1, + Self::THOMAS => 2, + Self::ANDOYER => 3, + Self::KARNEY => 4, + Self::_Unknown(i) => i, + }; + writer.write_i32(val) + } +} + +write_thrift_field!(EdgeInterpolationAlgorithm, FieldType::I32); + +// ---------------------------------------------------------------------- +// Mirrors thrift union `BloomFilterAlgorithm` + +thrift_union_all_empty!( +/// The algorithm used in Bloom filter. +union BloomFilterAlgorithm { + /// Block-based Bloom filter. + 1: SplitBlockAlgorithm BLOCK; +} +); + +// ---------------------------------------------------------------------- +// Mirrors thrift union `BloomFilterHash` + +thrift_union_all_empty!( +/// The hash function used in Bloom filter. This function takes the hash of a column value +/// using plain encoding. +union BloomFilterHash { + /// xxHash Strategy. + 1: XxHash XXHASH; +} +); + +// ---------------------------------------------------------------------- +// Mirrors thrift union `BloomFilterCompression` + +thrift_union_all_empty!( +/// The compression used in the Bloom filter. +union BloomFilterCompression { + 1: Uncompressed UNCOMPRESSED; +} +); + +// ---------------------------------------------------------------------- +// Mirrors thrift union `ColumnOrder` + +/// Sort order for page and column statistics. +/// +/// Types are associated with sort orders and column stats are aggregated using a sort +/// order, and a sort order should be considered when comparing values with statistics +/// min/max. +/// +/// See reference in +/// +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[allow(non_camel_case_types)] +pub enum SortOrder { + /// Signed (either value or legacy byte-wise) comparison. + SIGNED, + /// Unsigned (depending on physical type either value or byte-wise) comparison. + UNSIGNED, + /// Comparison is undefined. + UNDEFINED, +} + +impl SortOrder { + /// Returns true if this is [`Self::SIGNED`] + pub fn is_signed(&self) -> bool { + matches!(self, Self::SIGNED) + } +} + +/// Column order that specifies what method was used to aggregate min/max values for +/// statistics. +/// +/// If column order is undefined, then it is the legacy behaviour and all values should +/// be compared as signed values/bytes. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[allow(non_camel_case_types)] +pub enum ColumnOrder { + /// Column uses the order defined by its logical or physical type + /// (if there is no logical type), parquet-format 2.4.0+. + TYPE_DEFINED_ORDER(SortOrder), + // The following are not defined in the Parquet spec and should always be last. + /// Undefined column order, means legacy behaviour before parquet-format 2.4.0. + /// Sort order is always SIGNED. + UNDEFINED, + /// An unknown but present ColumnOrder. Statistics with an unknown `ColumnOrder` + /// will be ignored. + UNKNOWN, +} + +impl ColumnOrder { + /// Returns sort order for a physical/logical type. + #[deprecated( + since = "57.1.0", + note = "use `ColumnOrder::sort_order_for_type` instead" + )] + pub fn get_sort_order( + logical_type: Option, + converted_type: ConvertedType, + physical_type: Type, + ) -> SortOrder { + Self::sort_order_for_type(logical_type.as_ref(), converted_type, physical_type) + } + + /// Returns sort order for a physical/logical type. + pub fn sort_order_for_type( + logical_type: Option<&LogicalType>, + converted_type: ConvertedType, + physical_type: Type, + ) -> SortOrder { + match logical_type { + Some(logical) => match logical { + LogicalType::String | LogicalType::Enum | LogicalType::Json | LogicalType::Bson => { + SortOrder::UNSIGNED + } + LogicalType::Integer(int) => match int.is_signed { + true => SortOrder::SIGNED, + false => SortOrder::UNSIGNED, + }, + LogicalType::Map | LogicalType::List => SortOrder::UNDEFINED, + LogicalType::Decimal(_) => SortOrder::SIGNED, + LogicalType::Date => SortOrder::SIGNED, + LogicalType::Time(_) => SortOrder::SIGNED, + LogicalType::Timestamp(_) => SortOrder::SIGNED, + LogicalType::Unknown => SortOrder::UNDEFINED, + LogicalType::Uuid => SortOrder::UNSIGNED, + LogicalType::Float16 => SortOrder::SIGNED, + LogicalType::Variant(_) + | LogicalType::Geometry(_) + | LogicalType::Geography(_) + | LogicalType::_Unknown { .. } => SortOrder::UNDEFINED, + }, + // Fall back to converted type + None => Self::get_converted_sort_order(converted_type, physical_type), + } + } + + fn get_converted_sort_order(converted_type: ConvertedType, physical_type: Type) -> SortOrder { + match converted_type { + // Unsigned byte-wise comparison. + ConvertedType::UTF8 + | ConvertedType::JSON + | ConvertedType::BSON + | ConvertedType::ENUM => SortOrder::UNSIGNED, + + ConvertedType::INT_8 + | ConvertedType::INT_16 + | ConvertedType::INT_32 + | ConvertedType::INT_64 => SortOrder::SIGNED, + + ConvertedType::UINT_8 + | ConvertedType::UINT_16 + | ConvertedType::UINT_32 + | ConvertedType::UINT_64 => SortOrder::UNSIGNED, + + // Signed comparison of the represented value. + ConvertedType::DECIMAL => SortOrder::SIGNED, + + ConvertedType::DATE => SortOrder::SIGNED, + + ConvertedType::TIME_MILLIS + | ConvertedType::TIME_MICROS + | ConvertedType::TIMESTAMP_MILLIS + | ConvertedType::TIMESTAMP_MICROS => SortOrder::SIGNED, + + ConvertedType::INTERVAL => SortOrder::UNDEFINED, + + ConvertedType::LIST | ConvertedType::MAP | ConvertedType::MAP_KEY_VALUE => { + SortOrder::UNDEFINED + } + + // Fall back to physical type. + ConvertedType::NONE => Self::get_default_sort_order(physical_type), + } + } + + /// Returns default sort order based on physical type. + fn get_default_sort_order(physical_type: Type) -> SortOrder { + match physical_type { + // Order: false, true + Type::BOOLEAN => SortOrder::UNSIGNED, + Type::INT32 | Type::INT64 => SortOrder::SIGNED, + Type::INT96 => SortOrder::UNDEFINED, + // Notes to remember when comparing float/double values: + // If the min is a NaN, it should be ignored. + // If the max is a NaN, it should be ignored. + // If the min is +0, the row group may contain -0 values as well. + // If the max is -0, the row group may contain +0 values as well. + // When looking for NaN values, min and max should be ignored. + Type::FLOAT | Type::DOUBLE => SortOrder::SIGNED, + // Unsigned byte-wise comparison + Type::BYTE_ARRAY | Type::FIXED_LEN_BYTE_ARRAY => SortOrder::UNSIGNED, + } + } + + /// Returns sort order associated with this column order. + pub fn sort_order(&self) -> SortOrder { + match *self { + ColumnOrder::TYPE_DEFINED_ORDER(order) => order, + ColumnOrder::UNDEFINED => SortOrder::SIGNED, + ColumnOrder::UNKNOWN => SortOrder::UNDEFINED, + } + } +} + +impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for ColumnOrder { + fn read_thrift(prot: &mut R) -> Result { + let field_ident = prot.read_field_begin(0)?; + if field_ident.field_type == FieldType::Stop { + return Err(general_err!("Received empty union from remote ColumnOrder")); + } + let ret = match field_ident.id { + 1 => { + // NOTE: the sort order needs to be set correctly after parsing. + prot.skip_empty_struct()?; + Self::TYPE_DEFINED_ORDER(SortOrder::SIGNED) + } + _ => { + prot.skip(field_ident.field_type)?; + Self::UNKNOWN + } + }; + let field_ident = prot.read_field_begin(field_ident.id)?; + if field_ident.field_type != FieldType::Stop { + return Err(general_err!( + "Received multiple fields for union from remote ColumnOrder" + )); + } + Ok(ret) + } +} + +impl WriteThrift for ColumnOrder { + const ELEMENT_TYPE: ElementType = ElementType::Struct; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + match *self { + Self::TYPE_DEFINED_ORDER(_) => { + writer.write_field_begin(FieldType::Struct, 1, 0)?; + writer.write_struct_end()?; + } + _ => return Err(general_err!("Attempt to write undefined ColumnOrder")), + } + // write end of struct for this union + writer.write_struct_end() + } +} + +// ---------------------------------------------------------------------- +// Display handlers + +impl fmt::Display for Compression { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{self:?}") + } +} + +impl fmt::Display for SortOrder { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{self:?}") + } +} + +impl fmt::Display for ColumnOrder { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{self:?}") + } +} + +// ---------------------------------------------------------------------- +// LogicalType <=> ConvertedType conversion + +// Note: To prevent type loss when converting from ConvertedType to LogicalType, +// the conversion from ConvertedType -> LogicalType is not implemented. +// Such type loss includes: +// - Not knowing the decimal scale and precision of ConvertedType +// - Time and timestamp nanosecond precision, that is not supported in ConvertedType. + +impl From> for ConvertedType { + fn from(value: Option) -> Self { + match value { + Some(value) => match value { + LogicalType::String => ConvertedType::UTF8, + LogicalType::Map => ConvertedType::MAP, + LogicalType::List => ConvertedType::LIST, + LogicalType::Enum => ConvertedType::ENUM, + LogicalType::Decimal { .. } => ConvertedType::DECIMAL, + LogicalType::Date => ConvertedType::DATE, + LogicalType::Time(time) => match time.unit { + TimeUnit::MILLIS => ConvertedType::TIME_MILLIS, + TimeUnit::MICROS => ConvertedType::TIME_MICROS, + TimeUnit::NANOS => ConvertedType::NONE, + }, + LogicalType::Timestamp(time) => match time.unit { + TimeUnit::MILLIS => ConvertedType::TIMESTAMP_MILLIS, + TimeUnit::MICROS => ConvertedType::TIMESTAMP_MICROS, + TimeUnit::NANOS => ConvertedType::NONE, + }, + LogicalType::Integer(int_type) => match (int_type.bit_width, int_type.is_signed) { + (8, true) => ConvertedType::INT_8, + (16, true) => ConvertedType::INT_16, + (32, true) => ConvertedType::INT_32, + (64, true) => ConvertedType::INT_64, + (8, false) => ConvertedType::UINT_8, + (16, false) => ConvertedType::UINT_16, + (32, false) => ConvertedType::UINT_32, + (64, false) => ConvertedType::UINT_64, + (bit_width, is_signed) => panic!( + "Integer type bit_width={bit_width}, signed={is_signed} is not supported" + ), + }, + LogicalType::Json => ConvertedType::JSON, + LogicalType::Bson => ConvertedType::BSON, + LogicalType::Uuid + | LogicalType::Float16 + | LogicalType::Variant(_) + | LogicalType::Geometry(_) + | LogicalType::Geography(_) + | LogicalType::_Unknown { .. } + | LogicalType::Unknown => ConvertedType::NONE, + }, + None => ConvertedType::NONE, + } + } +} + +// ---------------------------------------------------------------------- +// String conversions for schema parsing. + +impl str::FromStr for Repetition { + type Err = ParquetError; + + fn from_str(s: &str) -> Result { + match s { + "REQUIRED" => Ok(Repetition::REQUIRED), + "OPTIONAL" => Ok(Repetition::OPTIONAL), + "REPEATED" => Ok(Repetition::REPEATED), + other => Err(general_err!("Invalid parquet repetition {}", other)), + } + } +} + +impl str::FromStr for Type { + type Err = ParquetError; + + fn from_str(s: &str) -> Result { + match s { + "BOOLEAN" => Ok(Type::BOOLEAN), + "INT32" => Ok(Type::INT32), + "INT64" => Ok(Type::INT64), + "INT96" => Ok(Type::INT96), + "FLOAT" => Ok(Type::FLOAT), + "DOUBLE" => Ok(Type::DOUBLE), + "BYTE_ARRAY" | "BINARY" => Ok(Type::BYTE_ARRAY), + "FIXED_LEN_BYTE_ARRAY" => Ok(Type::FIXED_LEN_BYTE_ARRAY), + other => Err(general_err!("Invalid parquet type {}", other)), + } + } +} + +impl str::FromStr for ConvertedType { + type Err = ParquetError; + + fn from_str(s: &str) -> Result { + match s { + "NONE" => Ok(ConvertedType::NONE), + "UTF8" => Ok(ConvertedType::UTF8), + "MAP" => Ok(ConvertedType::MAP), + "MAP_KEY_VALUE" => Ok(ConvertedType::MAP_KEY_VALUE), + "LIST" => Ok(ConvertedType::LIST), + "ENUM" => Ok(ConvertedType::ENUM), + "DECIMAL" => Ok(ConvertedType::DECIMAL), + "DATE" => Ok(ConvertedType::DATE), + "TIME_MILLIS" => Ok(ConvertedType::TIME_MILLIS), + "TIME_MICROS" => Ok(ConvertedType::TIME_MICROS), + "TIMESTAMP_MILLIS" => Ok(ConvertedType::TIMESTAMP_MILLIS), + "TIMESTAMP_MICROS" => Ok(ConvertedType::TIMESTAMP_MICROS), + "UINT_8" => Ok(ConvertedType::UINT_8), + "UINT_16" => Ok(ConvertedType::UINT_16), + "UINT_32" => Ok(ConvertedType::UINT_32), + "UINT_64" => Ok(ConvertedType::UINT_64), + "INT_8" => Ok(ConvertedType::INT_8), + "INT_16" => Ok(ConvertedType::INT_16), + "INT_32" => Ok(ConvertedType::INT_32), + "INT_64" => Ok(ConvertedType::INT_64), + "JSON" => Ok(ConvertedType::JSON), + "BSON" => Ok(ConvertedType::BSON), + "INTERVAL" => Ok(ConvertedType::INTERVAL), + other => Err(general_err!("Invalid parquet converted type {}", other)), + } + } +} + +impl str::FromStr for LogicalType { + type Err = ParquetError; + + fn from_str(s: &str) -> Result { + match s { + // The type is a placeholder that gets updated elsewhere + "INTEGER" => Ok(LogicalType::integer(8, false)), + "MAP" => Ok(LogicalType::Map), + "LIST" => Ok(LogicalType::List), + "ENUM" => Ok(LogicalType::Enum), + "DECIMAL" => Ok(LogicalType::decimal(-1, -1)), + "DATE" => Ok(LogicalType::Date), + "TIME" => Ok(LogicalType::time(false, TimeUnit::MILLIS)), + "TIMESTAMP" => Ok(LogicalType::timestamp(false, TimeUnit::MILLIS)), + "STRING" => Ok(LogicalType::String), + "JSON" => Ok(LogicalType::Json), + "BSON" => Ok(LogicalType::Bson), + "UUID" => Ok(LogicalType::Uuid), + "UNKNOWN" => Ok(LogicalType::Unknown), + "INTERVAL" => Err(general_err!( + "Interval parquet logical type not yet supported" + )), + "FLOAT16" => Ok(LogicalType::Float16), + "VARIANT" => Ok(LogicalType::variant(None)), + "GEOMETRY" => Ok(LogicalType::geometry(None)), + "GEOGRAPHY" => Ok(LogicalType::geography( + None, + Some(EdgeInterpolationAlgorithm::SPHERICAL), + )), + other => Err(general_err!("Invalid parquet logical type {}", other)), + } + } +} + +#[cfg(test)] +#[allow(deprecated)] // allow BIT_PACKED encoding for the whole test module +mod tests { + use super::*; + use crate::parquet_thrift::{ThriftSliceInputProtocol, tests::test_roundtrip}; + + #[test] + fn test_display_type() { + assert_eq!(Type::BOOLEAN.to_string(), "BOOLEAN"); + assert_eq!(Type::INT32.to_string(), "INT32"); + assert_eq!(Type::INT64.to_string(), "INT64"); + assert_eq!(Type::INT96.to_string(), "INT96"); + assert_eq!(Type::FLOAT.to_string(), "FLOAT"); + assert_eq!(Type::DOUBLE.to_string(), "DOUBLE"); + assert_eq!(Type::BYTE_ARRAY.to_string(), "BYTE_ARRAY"); + assert_eq!( + Type::FIXED_LEN_BYTE_ARRAY.to_string(), + "FIXED_LEN_BYTE_ARRAY" + ); + } + + #[test] + fn test_from_string_into_type() { + assert_eq!( + Type::BOOLEAN.to_string().parse::().unwrap(), + Type::BOOLEAN + ); + assert_eq!( + Type::INT32.to_string().parse::().unwrap(), + Type::INT32 + ); + assert_eq!( + Type::INT64.to_string().parse::().unwrap(), + Type::INT64 + ); + assert_eq!( + Type::INT96.to_string().parse::().unwrap(), + Type::INT96 + ); + assert_eq!( + Type::FLOAT.to_string().parse::().unwrap(), + Type::FLOAT + ); + assert_eq!( + Type::DOUBLE.to_string().parse::().unwrap(), + Type::DOUBLE + ); + assert_eq!( + Type::BYTE_ARRAY.to_string().parse::().unwrap(), + Type::BYTE_ARRAY + ); + assert_eq!("BINARY".parse::().unwrap(), Type::BYTE_ARRAY); + assert_eq!( + Type::FIXED_LEN_BYTE_ARRAY + .to_string() + .parse::() + .unwrap(), + Type::FIXED_LEN_BYTE_ARRAY + ); + } + + #[test] + fn test_converted_type_roundtrip() { + test_roundtrip(ConvertedType::UTF8); + test_roundtrip(ConvertedType::MAP); + test_roundtrip(ConvertedType::MAP_KEY_VALUE); + test_roundtrip(ConvertedType::LIST); + test_roundtrip(ConvertedType::ENUM); + test_roundtrip(ConvertedType::DECIMAL); + test_roundtrip(ConvertedType::DATE); + test_roundtrip(ConvertedType::TIME_MILLIS); + test_roundtrip(ConvertedType::TIME_MICROS); + test_roundtrip(ConvertedType::TIMESTAMP_MILLIS); + test_roundtrip(ConvertedType::TIMESTAMP_MICROS); + test_roundtrip(ConvertedType::UINT_8); + test_roundtrip(ConvertedType::UINT_16); + test_roundtrip(ConvertedType::UINT_32); + test_roundtrip(ConvertedType::UINT_64); + test_roundtrip(ConvertedType::INT_8); + test_roundtrip(ConvertedType::INT_16); + test_roundtrip(ConvertedType::INT_32); + test_roundtrip(ConvertedType::INT_64); + test_roundtrip(ConvertedType::JSON); + test_roundtrip(ConvertedType::BSON); + test_roundtrip(ConvertedType::INTERVAL); + } + + #[test] + fn test_read_invalid_converted_type() { + let mut prot = ThriftSliceInputProtocol::new(&[0x7eu8]); + let res = ConvertedType::read_thrift(&mut prot); + assert!(res.is_err()); + assert_eq!( + res.unwrap_err().to_string(), + "Parquet error: Unexpected ConvertedType 63" + ); + } + + #[test] + fn test_display_converted_type() { + assert_eq!(ConvertedType::NONE.to_string(), "NONE"); + assert_eq!(ConvertedType::UTF8.to_string(), "UTF8"); + assert_eq!(ConvertedType::MAP.to_string(), "MAP"); + assert_eq!(ConvertedType::MAP_KEY_VALUE.to_string(), "MAP_KEY_VALUE"); + assert_eq!(ConvertedType::LIST.to_string(), "LIST"); + assert_eq!(ConvertedType::ENUM.to_string(), "ENUM"); + assert_eq!(ConvertedType::DECIMAL.to_string(), "DECIMAL"); + assert_eq!(ConvertedType::DATE.to_string(), "DATE"); + assert_eq!(ConvertedType::TIME_MILLIS.to_string(), "TIME_MILLIS"); + assert_eq!(ConvertedType::DATE.to_string(), "DATE"); + assert_eq!(ConvertedType::TIME_MICROS.to_string(), "TIME_MICROS"); + assert_eq!( + ConvertedType::TIMESTAMP_MILLIS.to_string(), + "TIMESTAMP_MILLIS" + ); + assert_eq!( + ConvertedType::TIMESTAMP_MICROS.to_string(), + "TIMESTAMP_MICROS" + ); + assert_eq!(ConvertedType::UINT_8.to_string(), "UINT_8"); + assert_eq!(ConvertedType::UINT_16.to_string(), "UINT_16"); + assert_eq!(ConvertedType::UINT_32.to_string(), "UINT_32"); + assert_eq!(ConvertedType::UINT_64.to_string(), "UINT_64"); + assert_eq!(ConvertedType::INT_8.to_string(), "INT_8"); + assert_eq!(ConvertedType::INT_16.to_string(), "INT_16"); + assert_eq!(ConvertedType::INT_32.to_string(), "INT_32"); + assert_eq!(ConvertedType::INT_64.to_string(), "INT_64"); + assert_eq!(ConvertedType::JSON.to_string(), "JSON"); + assert_eq!(ConvertedType::BSON.to_string(), "BSON"); + assert_eq!(ConvertedType::INTERVAL.to_string(), "INTERVAL"); + assert_eq!(ConvertedType::DECIMAL.to_string(), "DECIMAL") + } + + #[test] + fn test_from_string_into_converted_type() { + assert_eq!( + ConvertedType::NONE + .to_string() + .parse::() + .unwrap(), + ConvertedType::NONE + ); + assert_eq!( + ConvertedType::UTF8 + .to_string() + .parse::() + .unwrap(), + ConvertedType::UTF8 + ); + assert_eq!( + ConvertedType::MAP + .to_string() + .parse::() + .unwrap(), + ConvertedType::MAP + ); + assert_eq!( + ConvertedType::MAP_KEY_VALUE + .to_string() + .parse::() + .unwrap(), + ConvertedType::MAP_KEY_VALUE + ); + assert_eq!( + ConvertedType::LIST + .to_string() + .parse::() + .unwrap(), + ConvertedType::LIST + ); + assert_eq!( + ConvertedType::ENUM + .to_string() + .parse::() + .unwrap(), + ConvertedType::ENUM + ); + assert_eq!( + ConvertedType::DECIMAL + .to_string() + .parse::() + .unwrap(), + ConvertedType::DECIMAL + ); + assert_eq!( + ConvertedType::DATE + .to_string() + .parse::() + .unwrap(), + ConvertedType::DATE + ); + assert_eq!( + ConvertedType::TIME_MILLIS + .to_string() + .parse::() + .unwrap(), + ConvertedType::TIME_MILLIS + ); + assert_eq!( + ConvertedType::TIME_MICROS + .to_string() + .parse::() + .unwrap(), + ConvertedType::TIME_MICROS + ); + assert_eq!( + ConvertedType::TIMESTAMP_MILLIS + .to_string() + .parse::() + .unwrap(), + ConvertedType::TIMESTAMP_MILLIS + ); + assert_eq!( + ConvertedType::TIMESTAMP_MICROS + .to_string() + .parse::() + .unwrap(), + ConvertedType::TIMESTAMP_MICROS + ); + assert_eq!( + ConvertedType::UINT_8 + .to_string() + .parse::() + .unwrap(), + ConvertedType::UINT_8 + ); + assert_eq!( + ConvertedType::UINT_16 + .to_string() + .parse::() + .unwrap(), + ConvertedType::UINT_16 + ); + assert_eq!( + ConvertedType::UINT_32 + .to_string() + .parse::() + .unwrap(), + ConvertedType::UINT_32 + ); + assert_eq!( + ConvertedType::UINT_64 + .to_string() + .parse::() + .unwrap(), + ConvertedType::UINT_64 + ); + assert_eq!( + ConvertedType::INT_8 + .to_string() + .parse::() + .unwrap(), + ConvertedType::INT_8 + ); + assert_eq!( + ConvertedType::INT_16 + .to_string() + .parse::() + .unwrap(), + ConvertedType::INT_16 + ); + assert_eq!( + ConvertedType::INT_32 + .to_string() + .parse::() + .unwrap(), + ConvertedType::INT_32 + ); + assert_eq!( + ConvertedType::INT_64 + .to_string() + .parse::() + .unwrap(), + ConvertedType::INT_64 + ); + assert_eq!( + ConvertedType::JSON + .to_string() + .parse::() + .unwrap(), + ConvertedType::JSON + ); + assert_eq!( + ConvertedType::BSON + .to_string() + .parse::() + .unwrap(), + ConvertedType::BSON + ); + assert_eq!( + ConvertedType::INTERVAL + .to_string() + .parse::() + .unwrap(), + ConvertedType::INTERVAL + ); + assert_eq!( + ConvertedType::DECIMAL + .to_string() + .parse::() + .unwrap(), + ConvertedType::DECIMAL + ) + } + + #[test] + fn test_logical_to_converted_type() { + let logical_none: Option = None; + assert_eq!(ConvertedType::from(logical_none), ConvertedType::NONE); + assert_eq!( + ConvertedType::from(Some(LogicalType::decimal(5, 20))), + ConvertedType::DECIMAL + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::Bson)), + ConvertedType::BSON + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::Json)), + ConvertedType::JSON + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::String)), + ConvertedType::UTF8 + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::Date)), + ConvertedType::DATE + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::time(true, TimeUnit::MILLIS))), + ConvertedType::TIME_MILLIS + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::time(true, TimeUnit::MICROS))), + ConvertedType::TIME_MICROS + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::time(false, TimeUnit::NANOS))), + ConvertedType::NONE + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::timestamp(true, TimeUnit::MILLIS))), + ConvertedType::TIMESTAMP_MILLIS + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::timestamp(false, TimeUnit::MICROS))), + ConvertedType::TIMESTAMP_MICROS + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::timestamp(false, TimeUnit::NANOS))), + ConvertedType::NONE + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::integer(8, false))), + ConvertedType::UINT_8 + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::integer(8, true))), + ConvertedType::INT_8 + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::integer(16, false))), + ConvertedType::UINT_16 + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::integer(16, true))), + ConvertedType::INT_16 + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::integer(32, false))), + ConvertedType::UINT_32 + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::integer(32, true))), + ConvertedType::INT_32 + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::integer(64, false))), + ConvertedType::UINT_64 + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::integer(64, true))), + ConvertedType::INT_64 + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::List)), + ConvertedType::LIST + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::Map)), + ConvertedType::MAP + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::Uuid)), + ConvertedType::NONE + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::Enum)), + ConvertedType::ENUM + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::Float16)), + ConvertedType::NONE + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::variant(None))), + ConvertedType::NONE + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::geometry(None))), + ConvertedType::NONE + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::geography(None, Some(Default::default())))), + ConvertedType::NONE + ); + assert_eq!( + ConvertedType::from(Some(LogicalType::Unknown)), + ConvertedType::NONE + ); + } + + #[test] + fn test_logical_type_roundtrip() { + test_roundtrip(LogicalType::String); + test_roundtrip(LogicalType::Map); + test_roundtrip(LogicalType::List); + test_roundtrip(LogicalType::Enum); + test_roundtrip(LogicalType::decimal(0, 20)); + test_roundtrip(LogicalType::Date); + test_roundtrip(LogicalType::time(true, TimeUnit::MICROS)); + test_roundtrip(LogicalType::time(false, TimeUnit::MILLIS)); + test_roundtrip(LogicalType::time(false, TimeUnit::NANOS)); + test_roundtrip(LogicalType::timestamp(false, TimeUnit::MICROS)); + test_roundtrip(LogicalType::timestamp(true, TimeUnit::MILLIS)); + test_roundtrip(LogicalType::timestamp(true, TimeUnit::NANOS)); + test_roundtrip(LogicalType::integer(8, true)); + test_roundtrip(LogicalType::integer(16, false)); + test_roundtrip(LogicalType::integer(32, true)); + test_roundtrip(LogicalType::integer(64, false)); + test_roundtrip(LogicalType::Json); + test_roundtrip(LogicalType::Bson); + test_roundtrip(LogicalType::Uuid); + test_roundtrip(LogicalType::Float16); + test_roundtrip(LogicalType::variant(Some(1))); + test_roundtrip(LogicalType::variant(None)); + test_roundtrip(LogicalType::geometry(Some("foo".to_owned()))); + test_roundtrip(LogicalType::geometry(None)); + test_roundtrip(LogicalType::geography( + Some("foo".to_owned()), + Some(EdgeInterpolationAlgorithm::ANDOYER), + )); + test_roundtrip(LogicalType::geography( + None, + Some(EdgeInterpolationAlgorithm::KARNEY), + )); + test_roundtrip(LogicalType::geography( + Some("foo".to_owned()), + Some(EdgeInterpolationAlgorithm::SPHERICAL), + )); + test_roundtrip(LogicalType::geography( + None, + Some(EdgeInterpolationAlgorithm::SPHERICAL), + )); + } + + #[test] + fn test_display_repetition() { + assert_eq!(Repetition::REQUIRED.to_string(), "REQUIRED"); + assert_eq!(Repetition::OPTIONAL.to_string(), "OPTIONAL"); + assert_eq!(Repetition::REPEATED.to_string(), "REPEATED"); + } + + #[test] + fn test_from_string_into_repetition() { + assert_eq!( + Repetition::REQUIRED + .to_string() + .parse::() + .unwrap(), + Repetition::REQUIRED + ); + assert_eq!( + Repetition::OPTIONAL + .to_string() + .parse::() + .unwrap(), + Repetition::OPTIONAL + ); + assert_eq!( + Repetition::REPEATED + .to_string() + .parse::() + .unwrap(), + Repetition::REPEATED + ); + } + + #[test] + fn test_display_encoding() { + assert_eq!(Encoding::PLAIN.to_string(), "PLAIN"); + assert_eq!(Encoding::PLAIN_DICTIONARY.to_string(), "PLAIN_DICTIONARY"); + assert_eq!(Encoding::RLE.to_string(), "RLE"); + assert_eq!(Encoding::BIT_PACKED.to_string(), "BIT_PACKED"); + assert_eq!( + Encoding::DELTA_BINARY_PACKED.to_string(), + "DELTA_BINARY_PACKED" + ); + assert_eq!( + Encoding::DELTA_LENGTH_BYTE_ARRAY.to_string(), + "DELTA_LENGTH_BYTE_ARRAY" + ); + assert_eq!(Encoding::DELTA_BYTE_ARRAY.to_string(), "DELTA_BYTE_ARRAY"); + assert_eq!(Encoding::RLE_DICTIONARY.to_string(), "RLE_DICTIONARY"); + } + + #[test] + fn test_compression_conversion() { + assert_eq!( + CompressionCodec::from(Compression::UNCOMPRESSED), + CompressionCodec::UNCOMPRESSED + ); + assert_eq!( + CompressionCodec::from(Compression::SNAPPY), + CompressionCodec::SNAPPY + ); + assert_eq!( + CompressionCodec::from(Compression::GZIP(Default::default())), + CompressionCodec::GZIP + ); + assert_eq!( + CompressionCodec::from(Compression::LZO), + CompressionCodec::LZO + ); + assert_eq!( + CompressionCodec::from(Compression::BROTLI(Default::default())), + CompressionCodec::BROTLI + ); + assert_eq!( + CompressionCodec::from(Compression::LZ4), + CompressionCodec::LZ4 + ); + assert_eq!( + CompressionCodec::from(Compression::ZSTD(Default::default())), + CompressionCodec::ZSTD + ); + assert_eq!( + CompressionCodec::from(Compression::LZ4_RAW), + CompressionCodec::LZ4_RAW + ); + + assert_eq!( + Compression::from(CompressionCodec::UNCOMPRESSED), + Compression::UNCOMPRESSED + ); + assert_eq!( + Compression::from(CompressionCodec::SNAPPY), + Compression::SNAPPY + ); + assert_eq!( + Compression::from(CompressionCodec::GZIP), + Compression::GZIP(Default::default()) + ); + assert_eq!(Compression::from(CompressionCodec::LZO), Compression::LZO); + assert_eq!( + Compression::from(CompressionCodec::BROTLI), + Compression::BROTLI(Default::default()) + ); + assert_eq!(Compression::from(CompressionCodec::LZ4), Compression::LZ4); + assert_eq!( + Compression::from(CompressionCodec::ZSTD), + Compression::ZSTD(Default::default()) + ); + assert_eq!( + Compression::from(CompressionCodec::LZ4_RAW), + Compression::LZ4_RAW + ); + } + + #[test] + fn test_display_compression() { + assert_eq!(Compression::UNCOMPRESSED.to_string(), "UNCOMPRESSED"); + assert_eq!(Compression::SNAPPY.to_string(), "SNAPPY"); + assert_eq!( + Compression::GZIP(Default::default()).to_string(), + "GZIP(GzipLevel(6))" + ); + assert_eq!(Compression::LZO.to_string(), "LZO"); + assert_eq!( + Compression::BROTLI(Default::default()).to_string(), + "BROTLI(BrotliLevel(1))" + ); + assert_eq!(Compression::LZ4.to_string(), "LZ4"); + assert_eq!( + Compression::ZSTD(Default::default()).to_string(), + "ZSTD(ZstdLevel(1))" + ); + } + + #[test] + fn test_display_page_type() { + assert_eq!(PageType::DATA_PAGE.to_string(), "DATA_PAGE"); + assert_eq!(PageType::INDEX_PAGE.to_string(), "INDEX_PAGE"); + assert_eq!(PageType::DICTIONARY_PAGE.to_string(), "DICTIONARY_PAGE"); + assert_eq!(PageType::DATA_PAGE_V2.to_string(), "DATA_PAGE_V2"); + } + + #[test] + fn test_display_sort_order() { + assert_eq!(SortOrder::SIGNED.to_string(), "SIGNED"); + assert_eq!(SortOrder::UNSIGNED.to_string(), "UNSIGNED"); + assert_eq!(SortOrder::UNDEFINED.to_string(), "UNDEFINED"); + } + + #[test] + fn test_display_column_order() { + assert_eq!( + ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::SIGNED).to_string(), + "TYPE_DEFINED_ORDER(SIGNED)" + ); + assert_eq!( + ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED).to_string(), + "TYPE_DEFINED_ORDER(UNSIGNED)" + ); + assert_eq!( + ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNDEFINED).to_string(), + "TYPE_DEFINED_ORDER(UNDEFINED)" + ); + assert_eq!(ColumnOrder::UNDEFINED.to_string(), "UNDEFINED"); + } + + #[test] + fn test_column_order_roundtrip() { + // SortOrder::SIGNED is the default on read. + test_roundtrip(ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::SIGNED)) + } + + #[test] + fn test_column_order_get_logical_type_sort_order() { + // Helper to check the order in a list of values. + // Only logical type is checked. + fn check_sort_order(types: Vec, expected_order: SortOrder) { + for tpe in types { + assert_eq!( + ColumnOrder::get_sort_order(Some(tpe), ConvertedType::NONE, Type::BYTE_ARRAY), + expected_order + ); + } + } + + // Unsigned comparison (physical type does not matter) + let unsigned = vec![ + LogicalType::String, + LogicalType::Json, + LogicalType::Bson, + LogicalType::Enum, + LogicalType::Uuid, + LogicalType::integer(8, false), + LogicalType::integer(16, false), + LogicalType::integer(32, false), + LogicalType::integer(64, false), + ]; + check_sort_order(unsigned, SortOrder::UNSIGNED); + + // Signed comparison (physical type does not matter) + let signed = vec![ + LogicalType::integer(8, true), + LogicalType::integer(16, true), + LogicalType::integer(32, true), + LogicalType::integer(64, true), + LogicalType::decimal(20, 4), + LogicalType::Date, + LogicalType::time(false, TimeUnit::MILLIS), + LogicalType::time(false, TimeUnit::MICROS), + LogicalType::time(true, TimeUnit::NANOS), + LogicalType::timestamp(false, TimeUnit::MILLIS), + LogicalType::timestamp(false, TimeUnit::MICROS), + LogicalType::timestamp(true, TimeUnit::NANOS), + LogicalType::Float16, + ]; + check_sort_order(signed, SortOrder::SIGNED); + + // Undefined comparison + let undefined = vec![ + LogicalType::List, + LogicalType::Map, + LogicalType::variant(None), + LogicalType::geometry(None), + LogicalType::geography(None, Some(Default::default())), + ]; + check_sort_order(undefined, SortOrder::UNDEFINED); + } + + #[test] + fn test_column_order_get_converted_type_sort_order() { + // Helper to check the order in a list of values. + // Only converted type is checked. + fn check_sort_order(types: Vec, expected_order: SortOrder) { + for tpe in types { + assert_eq!( + ColumnOrder::get_sort_order(None, tpe, Type::BYTE_ARRAY), + expected_order + ); + } + } + + // Unsigned comparison (physical type does not matter) + let unsigned = vec![ + ConvertedType::UTF8, + ConvertedType::JSON, + ConvertedType::BSON, + ConvertedType::ENUM, + ConvertedType::UINT_8, + ConvertedType::UINT_16, + ConvertedType::UINT_32, + ConvertedType::UINT_64, + ]; + check_sort_order(unsigned, SortOrder::UNSIGNED); + + // Signed comparison (physical type does not matter) + let signed = vec![ + ConvertedType::INT_8, + ConvertedType::INT_16, + ConvertedType::INT_32, + ConvertedType::INT_64, + ConvertedType::DECIMAL, + ConvertedType::DATE, + ConvertedType::TIME_MILLIS, + ConvertedType::TIME_MICROS, + ConvertedType::TIMESTAMP_MILLIS, + ConvertedType::TIMESTAMP_MICROS, + ]; + check_sort_order(signed, SortOrder::SIGNED); + + // Undefined comparison + let undefined = vec![ + ConvertedType::LIST, + ConvertedType::MAP, + ConvertedType::MAP_KEY_VALUE, + ConvertedType::INTERVAL, + ]; + check_sort_order(undefined, SortOrder::UNDEFINED); + + // Check None logical type + // This should return a sort order for byte array type. + check_sort_order(vec![ConvertedType::NONE], SortOrder::UNSIGNED); + } + + #[test] + fn test_column_order_get_default_sort_order() { + // Comparison based on physical type + assert_eq!( + ColumnOrder::get_default_sort_order(Type::BOOLEAN), + SortOrder::UNSIGNED + ); + assert_eq!( + ColumnOrder::get_default_sort_order(Type::INT32), + SortOrder::SIGNED + ); + assert_eq!( + ColumnOrder::get_default_sort_order(Type::INT64), + SortOrder::SIGNED + ); + assert_eq!( + ColumnOrder::get_default_sort_order(Type::INT96), + SortOrder::UNDEFINED + ); + assert_eq!( + ColumnOrder::get_default_sort_order(Type::FLOAT), + SortOrder::SIGNED + ); + assert_eq!( + ColumnOrder::get_default_sort_order(Type::DOUBLE), + SortOrder::SIGNED + ); + assert_eq!( + ColumnOrder::get_default_sort_order(Type::BYTE_ARRAY), + SortOrder::UNSIGNED + ); + assert_eq!( + ColumnOrder::get_default_sort_order(Type::FIXED_LEN_BYTE_ARRAY), + SortOrder::UNSIGNED + ); + } + + #[test] + fn test_column_order_sort_order() { + assert_eq!( + ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::SIGNED).sort_order(), + SortOrder::SIGNED + ); + assert_eq!( + ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED).sort_order(), + SortOrder::UNSIGNED + ); + assert_eq!( + ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNDEFINED).sort_order(), + SortOrder::UNDEFINED + ); + assert_eq!(ColumnOrder::UNDEFINED.sort_order(), SortOrder::SIGNED); + } + + #[test] + fn test_parse_encoding() { + let mut encoding: Encoding = "PLAIN".parse().unwrap(); + assert_eq!(encoding, Encoding::PLAIN); + encoding = "PLAIN_DICTIONARY".parse().unwrap(); + assert_eq!(encoding, Encoding::PLAIN_DICTIONARY); + encoding = "RLE".parse().unwrap(); + assert_eq!(encoding, Encoding::RLE); + encoding = "BIT_PACKED".parse().unwrap(); + assert_eq!(encoding, Encoding::BIT_PACKED); + encoding = "DELTA_BINARY_PACKED".parse().unwrap(); + assert_eq!(encoding, Encoding::DELTA_BINARY_PACKED); + encoding = "DELTA_LENGTH_BYTE_ARRAY".parse().unwrap(); + assert_eq!(encoding, Encoding::DELTA_LENGTH_BYTE_ARRAY); + encoding = "DELTA_BYTE_ARRAY".parse().unwrap(); + assert_eq!(encoding, Encoding::DELTA_BYTE_ARRAY); + encoding = "RLE_DICTIONARY".parse().unwrap(); + assert_eq!(encoding, Encoding::RLE_DICTIONARY); + encoding = "BYTE_STREAM_SPLIT".parse().unwrap(); + assert_eq!(encoding, Encoding::BYTE_STREAM_SPLIT); + + // test lowercase + encoding = "byte_stream_split".parse().unwrap(); + assert_eq!(encoding, Encoding::BYTE_STREAM_SPLIT); + + // test unknown string + match "plain_xxx".parse::() { + Ok(e) => { + panic!("Should not be able to parse {e:?}"); + } + Err(e) => { + assert_eq!(e.to_string(), "Parquet error: unknown encoding: plain_xxx"); + } + } + } + + #[test] + fn test_parse_compression() { + let mut compress: Compression = "snappy".parse().unwrap(); + assert_eq!(compress, Compression::SNAPPY); + compress = "lzo".parse().unwrap(); + assert_eq!(compress, Compression::LZO); + compress = "zstd(3)".parse().unwrap(); + assert_eq!(compress, Compression::ZSTD(ZstdLevel::try_new(3).unwrap())); + compress = "LZ4_RAW".parse().unwrap(); + assert_eq!(compress, Compression::LZ4_RAW); + compress = "uncompressed".parse().unwrap(); + assert_eq!(compress, Compression::UNCOMPRESSED); + compress = "snappy".parse().unwrap(); + assert_eq!(compress, Compression::SNAPPY); + compress = "gzip(9)".parse().unwrap(); + assert_eq!(compress, Compression::GZIP(GzipLevel::try_new(9).unwrap())); + compress = "lzo".parse().unwrap(); + assert_eq!(compress, Compression::LZO); + compress = "brotli(3)".parse().unwrap(); + assert_eq!( + compress, + Compression::BROTLI(BrotliLevel::try_new(3).unwrap()) + ); + compress = "lz4".parse().unwrap(); + assert_eq!(compress, Compression::LZ4); + + // test unknown compression + let mut err = "plain_xxx".parse::().unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: unknown encoding: plain_xxx" + ); + + // test invalid compress level + err = "gzip(-10)".parse::().unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: unknown encoding: gzip(-10)" + ); + } + + #[test] + fn test_display_boundary_order() { + assert_eq!(BoundaryOrder::ASCENDING.to_string(), "ASCENDING"); + assert_eq!(BoundaryOrder::DESCENDING.to_string(), "DESCENDING"); + assert_eq!(BoundaryOrder::UNORDERED.to_string(), "UNORDERED"); + } + + #[test] + fn test_display_edge_algo() { + assert_eq!( + EdgeInterpolationAlgorithm::SPHERICAL.to_string(), + "SPHERICAL" + ); + assert_eq!(EdgeInterpolationAlgorithm::VINCENTY.to_string(), "VINCENTY"); + assert_eq!(EdgeInterpolationAlgorithm::THOMAS.to_string(), "THOMAS"); + assert_eq!(EdgeInterpolationAlgorithm::ANDOYER.to_string(), "ANDOYER"); + assert_eq!(EdgeInterpolationAlgorithm::KARNEY.to_string(), "KARNEY"); + } + + #[test] + fn test_from_str_edge_algo() { + assert_eq!( + "spHErical".parse::().unwrap(), + EdgeInterpolationAlgorithm::SPHERICAL + ); + assert_eq!( + "vinceNTY".parse::().unwrap(), + EdgeInterpolationAlgorithm::VINCENTY + ); + assert_eq!( + "tHOmas".parse::().unwrap(), + EdgeInterpolationAlgorithm::THOMAS + ); + assert_eq!( + "anDOYEr".parse::().unwrap(), + EdgeInterpolationAlgorithm::ANDOYER + ); + assert_eq!( + "kaRNey".parse::().unwrap(), + EdgeInterpolationAlgorithm::KARNEY + ); + assert!( + "does not exist" + .parse::() + .is_err() + ); + } + + fn encodings_roundtrip(mut encodings: Vec) { + encodings.sort(); + let mask = EncodingMask::new_from_encodings(encodings.iter()); + assert!(mask.all_set(encodings.iter())); + let v = mask.encodings().collect::>(); + assert_eq!(v, encodings); + } + + #[test] + fn test_encoding_roundtrip() { + encodings_roundtrip( + [ + Encoding::RLE, + Encoding::PLAIN, + Encoding::DELTA_BINARY_PACKED, + ] + .into(), + ); + encodings_roundtrip([Encoding::RLE_DICTIONARY, Encoding::PLAIN_DICTIONARY].into()); + encodings_roundtrip([].into()); + let encodings = [ + Encoding::PLAIN, + Encoding::BIT_PACKED, + Encoding::RLE, + Encoding::DELTA_BINARY_PACKED, + Encoding::DELTA_BYTE_ARRAY, + Encoding::DELTA_LENGTH_BYTE_ARRAY, + Encoding::PLAIN_DICTIONARY, + Encoding::RLE_DICTIONARY, + Encoding::BYTE_STREAM_SPLIT, + ]; + encodings_roundtrip(encodings.into()); + } + + #[test] + fn test_invalid_encoding_mask() { + // any set bits higher than the max should trigger an error + let res = EncodingMask::try_new(-1); + assert!(res.is_err()); + let err = res.unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: Attempt to create invalid mask: 0xffffffff" + ); + + // test that GROUP_VAR_INT is disallowed + let res = EncodingMask::try_new(2); + assert!(res.is_err()); + let err = res.unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: Attempt to create invalid mask: 0x2" + ); + } + + #[test] + fn test_encoding_mask_is_only() { + let mask = EncodingMask::new_from_encodings([Encoding::PLAIN].iter()); + assert!(mask.is_only(Encoding::PLAIN)); + + let mask = + EncodingMask::new_from_encodings([Encoding::PLAIN, Encoding::PLAIN_DICTIONARY].iter()); + assert!(!mask.is_only(Encoding::PLAIN)); + } +} diff --git a/vendor/parquet/src/bin/cli/mod.rs b/vendor/parquet/src/bin/cli/mod.rs new file mode 100644 index 0000000..e04b285 --- /dev/null +++ b/vendor/parquet/src/bin/cli/mod.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. + +/// Parses a dot-separated Parquet column path, using `\` as an escape character +/// for literal dots and backslashes. +pub fn parse_column_path(path: &str) -> Result, String> { + let mut parts = Vec::new(); + let mut current = String::new(); + let mut escaped = false; + + for c in path.chars() { + if escaped { + match c { + '.' | '\\' => current.push(c), + _ => { + return Err(format!( + "Invalid escape sequence \\{c} in column path {}", + path + )); + } + } + escaped = false; + continue; + } + + match c { + '\\' => escaped = true, + '.' => { + parts.push(current); + current = String::new(); + } + _ => current.push(c), + } + } + + if escaped { + return Err(format!( + "Column path {} ends with an incomplete escape sequence", + path + )); + } + + parts.push(current); + Ok(parts) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_column_path_splits_unescaped_dots() { + assert_eq!(parse_column_path("a.b.c").unwrap(), vec!["a", "b", "c"]); + } + + #[test] + fn parse_column_path_allows_escaped_dots() { + assert_eq!(parse_column_path(r"a\.b.c").unwrap(), vec!["a.b", "c"]); + } + + #[test] + fn parse_column_path_allows_escaped_backslashes() { + assert_eq!(parse_column_path(r"a\\b.c").unwrap(), vec![r"a\b", "c"]); + } + + #[test] + fn parse_column_path_invalid_escape_sequence() { + let err = parse_column_path(r"a.\b.c").unwrap_err(); + assert!(err.contains("Invalid escape sequence")); + } + + #[test] + fn parse_column_path_incomplete_escape_sequence() { + let err = parse_column_path(r"a.b.c\").unwrap_err(); + assert!(err.contains("incomplete escape sequence")); + } +} diff --git a/vendor/parquet/src/bin/parquet-concat.rs b/vendor/parquet/src/bin/parquet-concat.rs new file mode 100644 index 0000000..a6f1aef --- /dev/null +++ b/vendor/parquet/src/bin/parquet-concat.rs @@ -0,0 +1,138 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Binary that concatenates the column data of one or more parquet files +//! +//! # Install +//! +//! `parquet-concat` can be installed using `cargo`: +//! ``` +//! cargo install parquet --features=cli +//! ``` +//! After this `parquet-concat` should be available: +//! ``` +//! parquet-concat out.parquet a.parquet b.parquet +//! ``` +//! +//! The binary can also be built from the source code and run as follows: +//! ``` +//! cargo run --features=cli --bin parquet-concat out.parquet a.parquet b.parquet +//! ``` +//! +//! Note: this does not currently support preserving the page index or bloom filters +//! + +use clap::Parser; +use parquet::bloom_filter::Sbbf; +use parquet::column::writer::ColumnCloseResult; +use parquet::errors::{ParquetError, Result}; +use parquet::file::metadata::{ColumnChunkMetaData, PageIndexPolicy, ParquetMetaDataReader}; +use parquet::file::properties::WriterProperties; +use parquet::file::reader::ChunkReader; +use parquet::file::writer::SerializedFileWriter; +use std::fs::File; +use std::sync::Arc; + +#[derive(Debug, Parser)] +#[clap(author, version)] +/// Concatenates one or more parquet files +struct Args { + /// Path to output + output: String, + + /// Path to input files + input: Vec, +} + +fn read_bloom_filter(column: &ColumnChunkMetaData, input: &R) -> Option { + Sbbf::read_from_column_chunk(column, input).ok().flatten() +} + +impl Args { + fn run(&self) -> Result<()> { + if self.input.is_empty() { + return Err(ParquetError::General( + "Must provide at least one input file".into(), + )); + } + + let output = File::create(&self.output)?; + + let inputs = self + .input + .iter() + .map(|x| { + let reader = File::open(x)?; + // Enable reading page indexes if present + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Optional) + .parse_and_finish(&reader)?; + Ok((reader, metadata)) + }) + .collect::>>()?; + + let expected = inputs[0].1.file_metadata().schema(); + for (_, metadata) in inputs.iter().skip(1) { + let actual = metadata.file_metadata().schema(); + if expected != actual { + return Err(ParquetError::General(format!( + "inputs must have the same schema, {expected:#?} vs {actual:#?}" + ))); + } + } + + let props = Arc::new(WriterProperties::builder().build()); + let schema = inputs[0].1.file_metadata().schema_descr().root_schema_ptr(); + let mut writer = SerializedFileWriter::new(output, schema, props)?; + + for (input, metadata) in inputs { + let column_indexes = metadata.column_index(); + let offset_indexes = metadata.offset_index(); + + for (rg_idx, rg) in metadata.row_groups().iter().enumerate() { + let rg_column_indexes = column_indexes.and_then(|ci| ci.get(rg_idx)); + let rg_offset_indexes = offset_indexes.and_then(|oi| oi.get(rg_idx)); + let mut rg_out = writer.next_row_group()?; + for (col_idx, column) in rg.columns().iter().enumerate() { + let bloom_filter = read_bloom_filter(column, &input); + let column_index = rg_column_indexes.and_then(|row| row.get(col_idx)).cloned(); + + let offset_index = rg_offset_indexes.and_then(|row| row.get(col_idx)).cloned(); + + let result = ColumnCloseResult { + bytes_written: column.compressed_size() as _, + rows_written: rg.num_rows() as _, + metadata: column.clone(), + bloom_filter, + column_index, + offset_index, + }; + rg_out.append_column(&input, result)?; + } + rg_out.close()?; + } + } + + writer.close()?; + + Ok(()) + } +} + +fn main() -> Result<()> { + Args::parse().run() +} diff --git a/vendor/parquet/src/bin/parquet-fromcsv-help.txt b/vendor/parquet/src/bin/parquet-fromcsv-help.txt new file mode 100644 index 0000000..ac38c56 --- /dev/null +++ b/vendor/parquet/src/bin/parquet-fromcsv-help.txt @@ -0,0 +1,75 @@ + +Usage: parquet [OPTIONS] --schema --input-file --output-file + +Options: + -s, --schema + message schema for output Parquet + + -i, --input-file + input CSV file + + -o, --output-file + output Parquet file + + -f, --input-format + input file format + + [default: csv] + [possible values: csv, tsv] + + -b, --batch-size + batch size + + [env: PARQUET_FROM_CSV_BATCHSIZE=] + [default: 1000] + + -h, --has-header + has header + + -d, --delimiter + field delimiter + + default value: when input_format==CSV: ',' when input_format==TSV: 'TAB' + + -r, --record-terminator + record terminator + + [possible values: lf, crlf, cr] + + -e, --escape-char + escape character + + -q, --quote-char + quote character + + -D, --double-quote + double quote + + [possible values: true, false] + + -C, --csv-compression + compression mode of csv + + [default: UNCOMPRESSED] + + -c, --parquet-compression + compression mode of parquet + + [default: SNAPPY] + + -w, --writer-version + writer version + + -m, --max-row-group-size + max row group size + + --enable-bloom-filter + whether to enable bloom filter writing + + [possible values: true, false] + + --help + display usage help + + -V, --version + Print version diff --git a/vendor/parquet/src/bin/parquet-fromcsv.rs b/vendor/parquet/src/bin/parquet-fromcsv.rs new file mode 100644 index 0000000..4a59385 --- /dev/null +++ b/vendor/parquet/src/bin/parquet-fromcsv.rs @@ -0,0 +1,750 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Binary file to converts csv to Parquet file +//! +//! # Install +//! +//! `parquet-fromcsv` can be installed using `cargo`: +//! +//! ```text +//! cargo install parquet --features=cli +//! ``` +//! +//! After this `parquet-fromcsv` should be available: +//! +//! ```text +//! parquet-fromcsv --schema message_schema_for_parquet.txt input.csv output.parquet +//! ``` +//! +//! The binary can also be built from the source code and run as follows: +//! +//! ```text +//! cargo run --features=cli --bin parquet-fromcsv --schema message_schema_for_parquet.txt \ +//! \ input.csv output.parquet +//! ``` +//! +//! # Options +//! +//! ```text +#![doc = include_str!("./parquet-fromcsv-help.txt")] // Update for this file : Run test test_command_help +//! ``` +//! +//! ## Parquet file options +//! +//! ```text +//! - `-b`, `--batch-size` : Batch size for Parquet +//! - `-c`, `--parquet-compression` : Compression option for Parquet, default is SNAPPY +//! - `-s`, `--schema` : Path to message schema for generated Parquet file +//! - `-o`, `--output-file` : Path to output Parquet file +//! - `-w`, `--writer-version` : Writer version +//! - `-m`, `--max-row-group-size` : Max row group size +//! - `--enable-bloom-filter` : Enable bloom filter during writing +//! ``` +//! +//! ## Input file options +//! +//! ```text +//! - `-i`, `--input-file` : Path to input CSV file +//! - `-f`, `--input-format` : Dialect for input file, `csv` or `tsv`. +//! - `-C`, `--csv-compression` : Compression option for csv, default is UNCOMPRESSED +//! - `-d`, `--delimiter : Field delimiter for CSV file, default depends `--input-format` +//! - `-e`, `--escape` : Escape character for input file +//! - `-h`, `--has-header` : Input has header +//! - `-r`, `--record-terminator` : Record terminator character for input. default is CRLF +//! - `-q`, `--quote-char` : Input quoting character +//! ``` +//! + +use std::{ + fmt::Display, + fs::{File, read_to_string}, + io::Read, + path::{Path, PathBuf}, + sync::Arc, +}; + +use arrow_csv::ReaderBuilder; +use arrow_schema::{ArrowError, Schema}; +use clap::{Parser, ValueEnum}; +use parquet::arrow::arrow_writer::ArrowWriterOptions; +use parquet::{ + arrow::{ArrowWriter, parquet_to_arrow_schema}, + basic::Compression, + errors::ParquetError, + file::properties::{WriterProperties, WriterVersion}, + schema::{parser::parse_message_type, types::SchemaDescriptor}, +}; + +#[derive(Debug)] +enum ParquetFromCsvError { + CommandLineParseError(clap::Error), + IoError(std::io::Error), + ArrowError(ArrowError), + ParquetError(ParquetError), + WithContext(String, Box), +} + +impl From for ParquetFromCsvError { + fn from(e: std::io::Error) -> Self { + Self::IoError(e) + } +} + +impl From for ParquetFromCsvError { + fn from(e: ArrowError) -> Self { + Self::ArrowError(e) + } +} + +impl From for ParquetFromCsvError { + fn from(e: ParquetError) -> Self { + Self::ParquetError(e) + } +} + +impl From for ParquetFromCsvError { + fn from(e: clap::Error) -> Self { + Self::CommandLineParseError(e) + } +} + +impl ParquetFromCsvError { + pub fn with_context>( + inner_error: E, + context: &str, + ) -> ParquetFromCsvError { + let inner = inner_error.into(); + ParquetFromCsvError::WithContext(context.to_string(), Box::new(inner)) + } +} + +impl Display for ParquetFromCsvError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ParquetFromCsvError::CommandLineParseError(e) => write!(f, "{e}"), + ParquetFromCsvError::IoError(e) => write!(f, "{e}"), + ParquetFromCsvError::ArrowError(e) => write!(f, "{e}"), + ParquetFromCsvError::ParquetError(e) => write!(f, "{e}"), + ParquetFromCsvError::WithContext(c, e) => { + writeln!(f, "{e}")?; + write!(f, "context: {c}") + } + } + } +} + +#[derive(Debug, Parser)] +#[clap(author, version, disable_help_flag=true, about("Binary to convert csv to Parquet"), long_about=None)] +struct Args { + /// Path to a text file containing a parquet schema definition + #[clap(short, long, help("message schema for output Parquet"))] + schema: PathBuf, + /// input CSV file path + #[clap(short, long, help("input CSV file"))] + input_file: PathBuf, + /// output Parquet file path + #[clap(short, long, help("output Parquet file"))] + output_file: PathBuf, + /// input file format + #[clap( + value_enum, + short('f'), + long, + help("input file format"), + default_value_t=CsvDialect::Csv + )] + input_format: CsvDialect, + /// batch size + #[clap( + short, + long, + help("batch size"), + default_value_t = 1000, + env = "PARQUET_FROM_CSV_BATCHSIZE" + )] + batch_size: usize, + /// has header line + #[clap(short, long, help("has header"))] + has_header: bool, + /// field delimiter + /// + /// default value: + /// when input_format==CSV: ',' + /// when input_format==TSV: 'TAB' + #[clap(short, long, help("field delimiter"))] + delimiter: Option, + #[clap(value_enum, short, long, help("record terminator"))] + record_terminator: Option, + #[clap(short, long, help("escape character"))] + escape_char: Option, + #[clap(short, long, help("quote character"))] + quote_char: Option, + #[clap(short('D'), long, help("double quote"))] + double_quote: Option, + #[clap(short('C'), long, help("compression mode of csv"), default_value_t=Compression::UNCOMPRESSED)] + #[clap(value_parser=compression_from_str)] + csv_compression: Compression, + #[clap(short('c'), long, help("compression mode of parquet"), default_value_t=Compression::SNAPPY)] + #[clap(value_parser=compression_from_str)] + parquet_compression: Compression, + + #[clap(short, long, help("writer version"))] + #[clap(value_parser=writer_version_from_str)] + writer_version: Option, + #[clap(short, long, help("max row group size"))] + max_row_group_size: Option, + #[clap(long, help("whether to enable bloom filter writing"))] + enable_bloom_filter: Option, + + #[clap(long, action=clap::ArgAction::Help, help("display usage help"))] + help: Option, +} + +fn compression_from_str(cmp: &str) -> Result { + match cmp.to_uppercase().as_str() { + "UNCOMPRESSED" => Ok(Compression::UNCOMPRESSED), + "SNAPPY" => Ok(Compression::SNAPPY), + "GZIP" => Ok(Compression::GZIP(Default::default())), + "LZO" => Ok(Compression::LZO), + "BROTLI" => Ok(Compression::BROTLI(Default::default())), + "LZ4" => Ok(Compression::LZ4), + "ZSTD" => Ok(Compression::ZSTD(Default::default())), + v => Err(format!( + "Unknown compression {v} : possible values UNCOMPRESSED, SNAPPY, GZIP, LZO, BROTLI, LZ4, ZSTD \n\nFor more information try --help" + )), + } +} + +fn writer_version_from_str(cmp: &str) -> Result { + match cmp.to_uppercase().as_str() { + "1" => Ok(WriterVersion::PARQUET_1_0), + "2" => Ok(WriterVersion::PARQUET_2_0), + v => Err(format!("Unknown writer version {v} : possible values 1, 2")), + } +} + +impl Args { + fn schema_path(&self) -> &Path { + self.schema.as_path() + } + fn get_delimiter(&self) -> u8 { + match self.delimiter { + Some(ch) => ch as u8, + None => match self.input_format { + CsvDialect::Csv => b',', + CsvDialect::Tsv => b'\t', + }, + } + } + fn get_terminator(&self) -> Option { + match self.record_terminator { + Some(RecordTerminator::LF) => Some(0x0a), + Some(RecordTerminator::CR) => Some(0x0d), + Some(RecordTerminator::Crlf) => None, + None => match self.input_format { + CsvDialect::Csv => None, + CsvDialect::Tsv => Some(0x0a), + }, + } + } + fn get_escape(&self) -> Option { + self.escape_char.map(|ch| ch as u8) + } + fn get_quote(&self) -> Option { + if self.quote_char.is_none() { + match self.input_format { + CsvDialect::Csv => Some(b'\"'), + CsvDialect::Tsv => None, + } + } else { + self.quote_char.map(|c| c as u8) + } + } +} + +#[derive(Debug, Clone, Copy, ValueEnum, PartialEq)] +enum CsvDialect { + Csv, + Tsv, +} + +#[derive(Debug, Clone, Copy, ValueEnum, PartialEq)] +enum RecordTerminator { + LF, + Crlf, + CR, +} + +fn configure_writer_properties(args: &Args) -> WriterProperties { + let mut properties_builder = + WriterProperties::builder().set_compression(args.parquet_compression); + if let Some(writer_version) = args.writer_version { + properties_builder = properties_builder.set_writer_version(writer_version); + } + if let Some(max_row_group_size) = args.max_row_group_size { + properties_builder = + properties_builder.set_max_row_group_row_count(Some(max_row_group_size)); + } + if let Some(enable_bloom_filter) = args.enable_bloom_filter { + properties_builder = properties_builder.set_bloom_filter_enabled(enable_bloom_filter); + } + properties_builder.build() +} + +fn configure_reader_builder(args: &Args, arrow_schema: Arc) -> ReaderBuilder { + fn configure_reader ReaderBuilder>( + builder: ReaderBuilder, + value: Option, + fun: F, + ) -> ReaderBuilder { + if let Some(val) = value { + fun(builder, val) + } else { + builder + } + } + + let mut builder = ReaderBuilder::new(arrow_schema) + .with_batch_size(args.batch_size) + .with_header(args.has_header) + .with_delimiter(args.get_delimiter()); + + builder = configure_reader( + builder, + args.get_terminator(), + ReaderBuilder::with_terminator, + ); + builder = configure_reader(builder, args.get_escape(), ReaderBuilder::with_escape); + builder = configure_reader(builder, args.get_quote(), ReaderBuilder::with_quote); + + builder +} + +fn convert_csv_to_parquet(args: &Args) -> Result<(), ParquetFromCsvError> { + let schema = read_to_string(args.schema_path()).map_err(|e| { + ParquetFromCsvError::with_context( + e, + &format!("Failed to open schema file {:#?}", args.schema_path()), + ) + })?; + let parquet_schema = Arc::new(parse_message_type(&schema)?); + let desc = SchemaDescriptor::new(parquet_schema); + let arrow_schema = Arc::new(parquet_to_arrow_schema(&desc, None)?); + + // create output parquet writer + let parquet_file = File::create(&args.output_file).map_err(|e| { + ParquetFromCsvError::with_context( + e, + &format!("Failed to create output file {:#?}", &args.output_file), + ) + })?; + + let options = ArrowWriterOptions::new() + .with_properties(configure_writer_properties(args)) + .with_schema_root(desc.name().to_string()); + + let mut arrow_writer = + ArrowWriter::try_new_with_options(parquet_file, arrow_schema.clone(), options) + .map_err(|e| ParquetFromCsvError::with_context(e, "Failed to create ArrowWriter"))?; + + // open input file + let input_file = File::open(&args.input_file).map_err(|e| { + ParquetFromCsvError::with_context( + e, + &format!("Failed to open input file {:#?}", &args.input_file), + ) + })?; + + // open input file decoder + let input_file_decoder = match args.csv_compression { + Compression::UNCOMPRESSED => Box::new(input_file) as Box, + Compression::SNAPPY => Box::new(snap::read::FrameDecoder::new(input_file)) as Box, + Compression::GZIP(_) => { + Box::new(flate2::read::MultiGzDecoder::new(input_file)) as Box + } + Compression::BROTLI(_) => { + Box::new(brotli::Decompressor::new(input_file, 0)) as Box + } + Compression::LZ4 => { + Box::new(lz4_flex::frame::FrameDecoder::new(input_file)) as Box + } + Compression::ZSTD(_) => { + Box::new(zstd::Decoder::new(input_file).map_err(|e| { + ParquetFromCsvError::with_context(e, "Failed to create zstd::Decoder") + })?) as Box + } + d => unimplemented!("compression type {d}"), + }; + + // create input csv reader + let builder = configure_reader_builder(args, arrow_schema); + let reader = builder.build(input_file_decoder)?; + for batch_result in reader { + let batch = batch_result.map_err(|e| { + ParquetFromCsvError::with_context(e, "Failed to read RecordBatch from CSV") + })?; + arrow_writer.write(&batch).map_err(|e| { + ParquetFromCsvError::with_context(e, "Failed to write RecordBatch to parquet") + })?; + } + arrow_writer + .close() + .map_err(|e| ParquetFromCsvError::with_context(e, "Failed to close parquet"))?; + Ok(()) +} + +fn main() -> Result<(), ParquetFromCsvError> { + let args = Args::parse(); + convert_csv_to_parquet(&args) +} + +#[cfg(test)] +mod tests { + use std::{ + io::Write, + path::{Path, PathBuf}, + }; + + use super::*; + use arrow::datatypes::{DataType, Field}; + use brotli::CompressorWriter; + use clap::{CommandFactory, Parser}; + use flate2::write::GzEncoder; + use parquet::basic::{BrotliLevel, GzipLevel, ZstdLevel}; + use parquet::file::reader::{FileReader, SerializedFileReader}; + use snap::write::FrameEncoder; + use tempfile::NamedTempFile; + + #[test] + fn test_command_help() { + let mut cmd = Args::command(); + let dir = std::env::var("CARGO_MANIFEST_DIR").unwrap(); + let mut path_buf = PathBuf::from(dir); + path_buf.push("src"); + path_buf.push("bin"); + path_buf.push("parquet-fromcsv-help.txt"); + let expected = std::fs::read_to_string(path_buf).unwrap(); + let mut buffer_vec = Vec::new(); + let mut buffer = std::io::Cursor::new(&mut buffer_vec); + cmd.write_long_help(&mut buffer).unwrap(); + // Remove Parquet version string from the help text + let mut actual = String::from_utf8(buffer_vec).unwrap(); + let pos = actual.find('\n').unwrap() + 1; + actual = actual[pos..].to_string(); + assert_eq!( + expected, actual, + "help text not match. please update to \n---\n{actual}\n---\n" + ) + } + + fn parse_args(mut extra_args: Vec<&str>) -> Result { + let mut args = vec![ + "test", + "--schema", + "test.schema", + "--input-file", + "infile.csv", + "--output-file", + "out.parquet", + ]; + args.append(&mut extra_args); + let args = Args::try_parse_from(args.iter())?; + Ok(args) + } + + #[test] + fn test_parse_arg_minimum() -> Result<(), ParquetFromCsvError> { + let args = parse_args(vec![])?; + + assert_eq!(args.schema, PathBuf::from(Path::new("test.schema"))); + assert_eq!(args.input_file, PathBuf::from(Path::new("infile.csv"))); + assert_eq!(args.output_file, PathBuf::from(Path::new("out.parquet"))); + // test default values + assert_eq!(args.input_format, CsvDialect::Csv); + assert_eq!(args.batch_size, 1000); + assert!(!args.has_header); + assert_eq!(args.delimiter, None); + assert_eq!(args.get_delimiter(), b','); + assert_eq!(args.record_terminator, None); + assert_eq!(args.get_terminator(), None); // CRLF + assert_eq!(args.quote_char, None); + assert_eq!(args.get_quote(), Some(b'\"')); + assert_eq!(args.double_quote, None); + assert_eq!(args.parquet_compression, Compression::SNAPPY); + Ok(()) + } + + #[test] + fn test_parse_arg_format_variants() -> Result<(), ParquetFromCsvError> { + let args = parse_args(vec!["--input-format", "csv"])?; + assert_eq!(args.input_format, CsvDialect::Csv); + assert_eq!(args.get_delimiter(), b','); + assert_eq!(args.get_terminator(), None); // CRLF + assert_eq!(args.get_quote(), Some(b'\"')); + assert_eq!(args.get_escape(), None); + let args = parse_args(vec!["--input-format", "tsv"])?; + assert_eq!(args.input_format, CsvDialect::Tsv); + assert_eq!(args.get_delimiter(), b'\t'); + assert_eq!(args.get_terminator(), Some(b'\x0a')); // LF + assert_eq!(args.get_quote(), None); // quote none + assert_eq!(args.get_escape(), None); + + let args = parse_args(vec!["--input-format", "csv", "--escape-char", "\\"])?; + assert_eq!(args.input_format, CsvDialect::Csv); + assert_eq!(args.get_delimiter(), b','); + assert_eq!(args.get_terminator(), None); // CRLF + assert_eq!(args.get_quote(), Some(b'\"')); + assert_eq!(args.get_escape(), Some(b'\\')); + + let args = parse_args(vec!["--input-format", "tsv", "--delimiter", ":"])?; + assert_eq!(args.input_format, CsvDialect::Tsv); + assert_eq!(args.get_delimiter(), b':'); + assert_eq!(args.get_terminator(), Some(b'\x0a')); // LF + assert_eq!(args.get_quote(), None); // quote none + assert_eq!(args.get_escape(), None); + + Ok(()) + } + + #[test] + #[should_panic] + fn test_parse_arg_format_error() { + parse_args(vec!["--input-format", "excel"]).unwrap(); + } + + #[test] + fn test_parse_arg_compression_format() { + let args = parse_args(vec!["--parquet-compression", "uncompressed"]).unwrap(); + assert_eq!(args.parquet_compression, Compression::UNCOMPRESSED); + let args = parse_args(vec!["--parquet-compression", "snappy"]).unwrap(); + assert_eq!(args.parquet_compression, Compression::SNAPPY); + let args = parse_args(vec!["--parquet-compression", "gzip"]).unwrap(); + assert_eq!( + args.parquet_compression, + Compression::GZIP(Default::default()) + ); + let args = parse_args(vec!["--parquet-compression", "lzo"]).unwrap(); + assert_eq!(args.parquet_compression, Compression::LZO); + let args = parse_args(vec!["--parquet-compression", "lz4"]).unwrap(); + assert_eq!(args.parquet_compression, Compression::LZ4); + let args = parse_args(vec!["--parquet-compression", "brotli"]).unwrap(); + assert_eq!( + args.parquet_compression, + Compression::BROTLI(Default::default()) + ); + let args = parse_args(vec!["--parquet-compression", "zstd"]).unwrap(); + assert_eq!( + args.parquet_compression, + Compression::ZSTD(Default::default()) + ); + } + + #[test] + fn test_parse_arg_compression_format_fail() { + match parse_args(vec!["--parquet-compression", "zip"]) { + Ok(_) => panic!("unexpected success"), + Err(e) => { + let err = e.to_string(); + assert!(err.contains("error: invalid value 'zip' for '--parquet-compression ': Unknown compression ZIP : possible values UNCOMPRESSED, SNAPPY, GZIP, LZO, BROTLI, LZ4, ZSTD \n\nFor more information try --help"), "{err}") + } + } + } + + fn assert_debug_text(debug_text: &str, name: &str, value: &str) { + let pattern = format!(" {name}: {value}"); + assert!( + debug_text.contains(&pattern), + "\"{debug_text}\" not contains \"{pattern}\"" + ) + } + + #[test] + fn test_configure_reader_builder() { + let args = Args { + schema: PathBuf::from(Path::new("schema.arvo")), + input_file: PathBuf::from(Path::new("test.csv")), + output_file: PathBuf::from(Path::new("out.parquet")), + batch_size: 1000, + input_format: CsvDialect::Csv, + has_header: false, + delimiter: None, + record_terminator: None, + escape_char: None, + quote_char: None, + double_quote: None, + csv_compression: Compression::UNCOMPRESSED, + parquet_compression: Compression::SNAPPY, + writer_version: None, + max_row_group_size: None, + enable_bloom_filter: None, + help: None, + }; + let arrow_schema = Arc::new(Schema::new(vec![ + Field::new("field1", DataType::Utf8, false), + Field::new("field2", DataType::Utf8, false), + Field::new("field3", DataType::Utf8, false), + Field::new("field4", DataType::Utf8, false), + Field::new("field5", DataType::Utf8, false), + ])); + + let reader_builder = configure_reader_builder(&args, arrow_schema); + let builder_debug = format!("{reader_builder:?}"); + assert_debug_text(&builder_debug, "header", "false"); + assert_debug_text(&builder_debug, "delimiter", "Some(44)"); + assert_debug_text(&builder_debug, "quote", "Some(34)"); + assert_debug_text(&builder_debug, "terminator", "None"); + assert_debug_text(&builder_debug, "batch_size", "1000"); + assert_debug_text(&builder_debug, "escape", "None"); + + let args = Args { + schema: PathBuf::from(Path::new("schema.arvo")), + input_file: PathBuf::from(Path::new("test.csv")), + output_file: PathBuf::from(Path::new("out.parquet")), + batch_size: 2000, + input_format: CsvDialect::Tsv, + has_header: true, + delimiter: None, + record_terminator: None, + escape_char: Some('\\'), + quote_char: None, + double_quote: None, + csv_compression: Compression::UNCOMPRESSED, + parquet_compression: Compression::SNAPPY, + writer_version: None, + max_row_group_size: None, + enable_bloom_filter: None, + help: None, + }; + let arrow_schema = Arc::new(Schema::new(vec![ + Field::new("field1", DataType::Utf8, false), + Field::new("field2", DataType::Utf8, false), + Field::new("field3", DataType::Utf8, false), + Field::new("field4", DataType::Utf8, false), + Field::new("field5", DataType::Utf8, false), + ])); + let reader_builder = configure_reader_builder(&args, arrow_schema); + let builder_debug = format!("{reader_builder:?}"); + assert_debug_text(&builder_debug, "header", "true"); + assert_debug_text(&builder_debug, "delimiter", "Some(9)"); + assert_debug_text(&builder_debug, "quote", "None"); + assert_debug_text(&builder_debug, "terminator", "Some(10)"); + assert_debug_text(&builder_debug, "batch_size", "2000"); + assert_debug_text(&builder_debug, "escape", "Some(92)"); + } + + fn test_convert_compressed_csv_to_parquet(csv_compression: Compression) { + let schema = NamedTempFile::new().unwrap(); + let schema_text = r"message my_amazing_schema { + optional int32 id; + optional binary name (STRING); + }"; + schema.as_file().write_all(schema_text.as_bytes()).unwrap(); + + let mut input_file = NamedTempFile::new().unwrap(); + + fn write_tmp_file(w: &mut T) { + for index in 1..2000 { + write!(w, "{index},\"name_{index}\"\r\n").unwrap(); + } + w.flush().unwrap(); + } + + // make sure the input_file's lifetime being long enough + input_file = match csv_compression { + Compression::UNCOMPRESSED => { + write_tmp_file(&mut input_file); + input_file + } + Compression::SNAPPY => { + let mut encoder = FrameEncoder::new(input_file); + write_tmp_file(&mut encoder); + encoder.into_inner().unwrap() + } + Compression::GZIP(level) => { + let mut encoder = GzEncoder::new( + input_file, + flate2::Compression::new(level.compression_level()), + ); + write_tmp_file(&mut encoder); + encoder.finish().unwrap() + } + Compression::BROTLI(level) => { + let mut encoder = + CompressorWriter::new(input_file, 0, level.compression_level(), 0); + write_tmp_file(&mut encoder); + encoder.into_inner() + } + Compression::LZ4 => { + let mut encoder = lz4_flex::frame::FrameEncoder::new(input_file); + write_tmp_file(&mut encoder); + encoder.finish().unwrap() + } + + Compression::ZSTD(level) => { + let mut encoder = zstd::Encoder::new(input_file, level.compression_level()) + .map_err(|e| { + ParquetFromCsvError::with_context(e, "Failed to create zstd::Encoder") + }) + .unwrap(); + write_tmp_file(&mut encoder); + encoder.finish().unwrap() + } + d => unimplemented!("compression type {d}"), + }; + + let output_parquet = NamedTempFile::new().unwrap(); + + let args = Args { + schema: PathBuf::from(schema.path()), + input_file: PathBuf::from(input_file.path()), + output_file: PathBuf::from(output_parquet.path()), + batch_size: 1000, + input_format: CsvDialect::Csv, + has_header: false, + delimiter: None, + record_terminator: None, + escape_char: None, + quote_char: None, + double_quote: None, + csv_compression, + parquet_compression: Compression::SNAPPY, + writer_version: None, + max_row_group_size: None, + // by default we shall test bloom filter writing + enable_bloom_filter: Some(true), + help: None, + }; + convert_csv_to_parquet(&args).unwrap(); + + let file = SerializedFileReader::new(output_parquet.into_file()).unwrap(); + let schema_name = file.metadata().file_metadata().schema().name(); + assert_eq!(schema_name, "my_amazing_schema"); + } + + #[test] + fn test_convert_csv_to_parquet() { + test_convert_compressed_csv_to_parquet(Compression::UNCOMPRESSED); + test_convert_compressed_csv_to_parquet(Compression::SNAPPY); + test_convert_compressed_csv_to_parquet(Compression::GZIP(GzipLevel::try_new(1).unwrap())); + test_convert_compressed_csv_to_parquet(Compression::BROTLI( + BrotliLevel::try_new(2).unwrap(), + )); + test_convert_compressed_csv_to_parquet(Compression::LZ4); + test_convert_compressed_csv_to_parquet(Compression::ZSTD(ZstdLevel::try_new(1).unwrap())); + } +} diff --git a/vendor/parquet/src/bin/parquet-index.rs b/vendor/parquet/src/bin/parquet-index.rs new file mode 100644 index 0000000..241fc20 --- /dev/null +++ b/vendor/parquet/src/bin/parquet-index.rs @@ -0,0 +1,231 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Binary that prints the [page index] of a parquet file +//! +//! # Install +//! +//! `parquet-layout` can be installed using `cargo`: +//! ``` +//! cargo install parquet --features=cli +//! ``` +//! After this `parquet-index` should be available: +//! ``` +//! parquet-index XYZ.parquet COLUMN_PATH +//! ``` +//! +//! The binary can also be built from the source code and run as follows: +//! ``` +//! cargo run --features=cli --bin parquet-index XYZ.parquet COLUMN_PATH +//! +//! [page index]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + +mod cli; + +use clap::Parser; +use parquet::data_type::ByteArray; +use parquet::errors::{ParquetError, Result}; +use parquet::file::page_index::column_index::{ + ByteArrayColumnIndex, ColumnIndexMetaData, PrimitiveColumnIndex, +}; +use parquet::file::page_index::offset_index::{OffsetIndexMetaData, PageLocation}; +use parquet::file::reader::{FileReader, SerializedFileReader}; +use parquet::file::serialized_reader::ReadOptionsBuilder; +use std::fs::File; + +#[derive(Debug, Parser)] +#[clap(author, version, about("Prints the page index of a parquet file"), long_about = None)] +struct Args { + #[clap(help("Path to a parquet file"))] + file: String, + + #[clap(help( + "Dot-separated column path to get index for. Literal dots can be escaped as `\\.`" + ))] + column: String, +} + +impl Args { + fn run(&self) -> Result<()> { + let file = File::open(&self.file)?; + let options = ReadOptionsBuilder::new().with_page_index().build(); + let reader = SerializedFileReader::new_with_options(file, options)?; + + let schema = reader.metadata().file_metadata().schema_descr(); + let column_path = cli::parse_column_path(&self.column).map_err(ParquetError::General)?; + let column_idx = schema + .columns() + .iter() + .position(|x| x.path().parts() == column_path.as_slice()) + .ok_or_else(|| { + ParquetError::General(format!("Failed to find column {}", self.column)) + })?; + + // Column index data for all row groups and columns + let column_index = reader + .metadata() + .column_index() + .ok_or_else(|| ParquetError::General("Column index not found".to_string()))?; + + // Offset index data for all row groups and columns + let offset_index = reader + .metadata() + .offset_index() + .ok_or_else(|| ParquetError::General("Offset index not found".to_string()))?; + + // Iterate through each row group + for (row_group_idx, ((column_indices, offset_indices), row_group)) in column_index + .iter() + .zip(offset_index) + .zip(reader.metadata().row_groups()) + .enumerate() + { + println!("Row Group: {row_group_idx}"); + let offset_index = offset_indices.get(column_idx).ok_or_else(|| { + ParquetError::General(format!( + "No offset index for row group {row_group_idx} column chunk {column_idx}" + )) + })?; + + let row_counts = + compute_row_counts(offset_index.page_locations.as_slice(), row_group.num_rows()); + match &column_indices[column_idx] { + ColumnIndexMetaData::NONE => println!("NO INDEX"), + ColumnIndexMetaData::BOOLEAN(v) => { + print_index::(v, offset_index, &row_counts)? + } + ColumnIndexMetaData::INT32(v) => print_index(v, offset_index, &row_counts)?, + ColumnIndexMetaData::INT64(v) => print_index(v, offset_index, &row_counts)?, + ColumnIndexMetaData::INT96(v) => print_index(v, offset_index, &row_counts)?, + ColumnIndexMetaData::FLOAT(v) => print_index(v, offset_index, &row_counts)?, + ColumnIndexMetaData::DOUBLE(v) => print_index(v, offset_index, &row_counts)?, + ColumnIndexMetaData::BYTE_ARRAY(v) => { + print_bytes_index(v, offset_index, &row_counts)? + } + ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(v) => { + print_bytes_index(v, offset_index, &row_counts)? + } + } + } + Ok(()) + } +} + +/// Computes the number of rows in each page within a column chunk +fn compute_row_counts(offset_index: &[PageLocation], rows: i64) -> Vec { + if offset_index.is_empty() { + return vec![]; + } + + let mut last = offset_index[0].first_row_index; + let mut out = Vec::with_capacity(offset_index.len()); + for o in offset_index.iter().skip(1) { + out.push(o.first_row_index - last); + last = o.first_row_index; + } + out.push(rows - last); + out +} + +/// Prints index information for a single column chunk +fn print_index( + column_index: &PrimitiveColumnIndex, + offset_index: &OffsetIndexMetaData, + row_counts: &[i64], +) -> Result<()> { + if column_index.num_pages() as usize != offset_index.page_locations.len() { + return Err(ParquetError::General(format!( + "Index length mismatch, got {} and {}", + column_index.num_pages(), + offset_index.page_locations.len() + ))); + } + + for (idx, (((min, max), o), row_count)) in column_index + .min_values_iter() + .zip(column_index.max_values_iter()) + .zip(offset_index.page_locations()) + .zip(row_counts) + .enumerate() + { + print!( + "Page {:>5} at offset {:>10} with length {:>10} and row count {:>10}", + idx, o.offset, o.compressed_page_size, row_count + ); + match min { + Some(m) => print!(", min {m:>10}"), + None => print!(", min {:>10}", "NONE"), + } + + match max { + Some(m) => print!(", max {m:>10}"), + None => print!(", max {:>10}", "NONE"), + } + println!() + } + + Ok(()) +} + +fn print_bytes_index( + column_index: &ByteArrayColumnIndex, + offset_index: &OffsetIndexMetaData, + row_counts: &[i64], +) -> Result<()> { + if column_index.num_pages() as usize != offset_index.page_locations.len() { + return Err(ParquetError::General(format!( + "Index length mismatch, got {} and {}", + column_index.num_pages(), + offset_index.page_locations.len() + ))); + } + + for (idx, (((min, max), o), row_count)) in column_index + .min_values_iter() + .zip(column_index.max_values_iter()) + .zip(offset_index.page_locations()) + .zip(row_counts) + .enumerate() + { + print!( + "Page {:>5} at offset {:>10} with length {:>10} and row count {:>10}", + idx, o.offset, o.compressed_page_size, row_count + ); + match min { + Some(m) => match String::from_utf8(m.to_vec()) { + Ok(s) => print!(", min {s:>10}"), + Err(_) => print!(", min {:>10}", ByteArray::from(m)), + }, + None => print!(", min {:>10}", "NONE"), + } + + match max { + Some(m) => match String::from_utf8(m.to_vec()) { + Ok(s) => print!(", max {s:>10}"), + Err(_) => print!(", min {:>10}", ByteArray::from(m)), + }, + None => print!(", max {:>10}", "NONE"), + } + println!() + } + + Ok(()) +} + +fn main() -> Result<()> { + Args::parse().run() +} diff --git a/vendor/parquet/src/bin/parquet-layout.rs b/vendor/parquet/src/bin/parquet-layout.rs new file mode 100644 index 0000000..b92228e --- /dev/null +++ b/vendor/parquet/src/bin/parquet-layout.rs @@ -0,0 +1,175 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Binary that prints the physical layout of a parquet file +//! +//! Alternatives to this binary include [`parquet-cli`] and [`parquet-viewer`]. +//! +//! # Install +//! +//! `parquet-layout` can be installed using `cargo`: +//! ``` +//! cargo install parquet --features=cli +//! ``` +//! After this `parquet-layout` should be available: +//! ``` +//! parquet-layout XYZ.parquet +//! ``` +//! +//! The binary can also be built from the source code and run as follows: +//! ``` +//! cargo run --features=cli --bin parquet-layout XYZ.parquet +//! ``` +//! +//! [`parquet-cli`]: https://github.com/apache/parquet-java/tree/master/parquet-cli +//! [`parquet-viewer`]: https://github.com/xiangpenghao/parquet-viewer + +use std::fs::File; + +use clap::Parser; +use parquet::file::metadata::ParquetMetaDataReader; +use serde::{Serialize, Serializer}; + +use parquet::basic::{CompressionCodec, Encoding}; +use parquet::errors::Result; +use parquet::file::reader::ChunkReader; + +#[derive(Serialize, Debug)] +struct Index { + offset: i64, + length: Option, +} + +#[derive(Serialize, Debug)] +struct Footer { + metadata_size: Option, +} + +#[derive(Serialize, Debug)] +struct ParquetFile { + row_groups: Vec, + footer: Footer, +} + +#[derive(Serialize, Debug)] +struct RowGroup { + columns: Vec, + row_count: i64, +} + +#[derive(Serialize, Debug)] +struct ColumnChunk { + path: String, + has_offset_index: bool, + has_column_index: bool, + has_bloom_filter: bool, + offset_index: Option, + column_index: Option, + bloom_filter: Option, + compression: DebugSerialize, + encodings: Vec>, +} + +#[derive(Debug)] +struct DebugSerialize(T); + +impl Serialize for DebugSerialize { + fn serialize(&self, serializer: S) -> std::result::Result + where + S: Serializer, + { + serializer.serialize_str(&format!("{:?}", &self.0)) + } +} + +fn do_layout(reader: &C) -> Result { + let mut metadata_reader = ParquetMetaDataReader::new(); + metadata_reader.try_parse(reader)?; + let metadata_size = metadata_reader.metadata_size(); + let metadata = metadata_reader.finish()?; + let schema = metadata.file_metadata().schema_descr(); + + let row_groups = (0..metadata.num_row_groups()) + .map(|row_group_idx| { + let row_group = metadata.row_group(row_group_idx); + let columns = row_group + .columns() + .iter() + .zip(schema.columns()) + .map(|(column, column_schema)| { + let compression = DebugSerialize(column.compression_codec()); + let encodings = column.encodings().map(DebugSerialize).collect(); + + Ok(ColumnChunk { + path: column_schema.path().parts().join("."), + has_offset_index: column.offset_index_offset().is_some(), + has_column_index: column.column_index_offset().is_some(), + has_bloom_filter: column.bloom_filter_offset().is_some(), + offset_index: column.offset_index_offset().map(|offset| Index { + offset, + length: column.offset_index_length(), + }), + column_index: column.column_index_offset().map(|offset| Index { + offset, + length: column.column_index_length(), + }), + bloom_filter: column.bloom_filter_offset().map(|offset| Index { + offset, + length: column.bloom_filter_length(), + }), + compression, + encodings, + }) + }) + .collect::>>()?; + + Ok(RowGroup { + columns, + row_count: row_group.num_rows(), + }) + }) + .collect::>>()?; + + Ok(ParquetFile { + row_groups, + footer: Footer { metadata_size }, + }) +} + +#[derive(Debug, Parser)] +#[clap(author, version, about("Prints the physical layout of a parquet file"), long_about = None)] +struct Args { + #[clap(help("Path to a parquet file"))] + file: String, +} + +impl Args { + fn run(&self) -> Result<()> { + let file = File::open(&self.file)?; + let layout = do_layout(&file)?; + + let out = std::io::stdout(); + let writer = out.lock(); + + serde_json::to_writer_pretty(writer, &layout).unwrap(); + Ok(()) + } +} + +fn main() -> Result<()> { + Args::parse().run() +} diff --git a/vendor/parquet/src/bin/parquet-read.rs b/vendor/parquet/src/bin/parquet-read.rs new file mode 100644 index 0000000..fe486e6 --- /dev/null +++ b/vendor/parquet/src/bin/parquet-read.rs @@ -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. + +//! Binary file to read data from a Parquet file. +//! +//! # Install +//! +//! `parquet-read` can be installed using `cargo`: +//! ``` +//! cargo install parquet --features=cli +//! ``` +//! After this `parquet-read` should be available: +//! ``` +//! parquet-read XYZ.parquet +//! ``` +//! +//! The binary can also be built from the source code and run as follows: +//! ``` +//! cargo run --features=cli --bin parquet-read XYZ.parquet +//! ``` +//! +//! Note that `parquet-read` reads full file schema, no projection or filtering is +//! applied. + +use clap::Parser; +use parquet::file::reader::{FileReader, SerializedFileReader}; +use parquet::record::Row; +use std::io::{self, Read}; +use std::{fs::File, path::Path}; + +#[derive(Debug, Parser)] +#[clap(author, version, about("Binary file to read data from a Parquet file"), long_about = None)] +struct Args { + #[clap(help("Path to a parquet file, or - for stdin"))] + file_name: String, + #[clap( + short, + long, + default_value_t = 0_usize, + help("Number of records to read. When not provided or 0, all records are read") + )] + num_records: usize, + #[clap(short, long, help("Print Parquet file in JSON lines format"))] + json: bool, +} + +fn main() { + let args = Args::parse(); + + let filename = args.file_name; + let num_records = args.num_records; + let json = args.json; + + let parquet_reader: Box = if filename == "-" { + let mut buf = Vec::new(); + io::stdin() + .read_to_end(&mut buf) + .expect("Failed to read stdin into a buffer"); + Box::new( + SerializedFileReader::new(bytes::Bytes::from(buf)).expect("Failed to create reader"), + ) + } else { + let path = Path::new(&filename); + let file = File::open(path).expect("Unable to open file"); + Box::new(SerializedFileReader::new(file).expect("Failed to create reader")) + }; + + // Use full schema as projected schema + let mut iter = parquet_reader + .get_row_iter(None) + .expect("Failed to create row iterator"); + + let mut start = 0; + let end = num_records; + let all_records = end == 0; + + while all_records || start < end { + match iter.next() { + Some(row) => print_row(&row.unwrap(), json), + None => break, + }; + start += 1; + } +} + +fn print_row(row: &Row, json: bool) { + if json { + println!("{}", row.to_json_value()) + } else { + println!("{row}"); + } +} diff --git a/vendor/parquet/src/bin/parquet-rewrite.rs b/vendor/parquet/src/bin/parquet-rewrite.rs new file mode 100644 index 0000000..2428c49 --- /dev/null +++ b/vendor/parquet/src/bin/parquet-rewrite.rs @@ -0,0 +1,434 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Binary file to rewrite parquet files. +//! +//! # Install +//! +//! `parquet-rewrite` can be installed using `cargo`: +//! ``` +//! cargo install parquet --features=cli +//! ``` +//! After this `parquet-rewrite` should be available: +//! ``` +//! parquet-rewrite -i XYZ.parquet -o XYZ2.parquet +//! ``` +//! +//! The binary can also be built from the source code and run as follows: +//! ``` +//! cargo run --features=cli --bin parquet-rewrite -- -i XYZ.parquet -o XYZ2.parquet +//! ``` + +use std::fs::File; + +use arrow_array::RecordBatchReader; +use clap::{Parser, ValueEnum, builder::PossibleValue}; +use parquet::{ + arrow::{ArrowWriter, arrow_reader::ParquetRecordBatchReaderBuilder}, + basic::{BrotliLevel, Compression, Encoding, GzipLevel, ZstdLevel}, + file::{ + properties::{BloomFilterPosition, EnabledStatistics, WriterProperties, WriterVersion}, + reader::FileReader, + serialized_reader::SerializedFileReader, + }, +}; + +#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum, Debug)] +enum CompressionArgs { + /// No compression. + None, + + /// Snappy + Snappy, + + /// GZip + Gzip, + + /// LZO + Lzo, + + /// Brotli + Brotli, + + /// LZ4 + Lz4, + + /// Zstd + Zstd, + + /// LZ4 Raw + Lz4Raw, +} + +fn compression_from_args(codec: CompressionArgs, level: Option) -> Compression { + match codec { + CompressionArgs::None => Compression::UNCOMPRESSED, + CompressionArgs::Snappy => Compression::SNAPPY, + CompressionArgs::Gzip => match level { + Some(lvl) => { + Compression::GZIP(GzipLevel::try_new(lvl).expect("invalid gzip compression level")) + } + None => Compression::GZIP(Default::default()), + }, + CompressionArgs::Lzo => Compression::LZO, + CompressionArgs::Brotli => match level { + Some(lvl) => Compression::BROTLI( + BrotliLevel::try_new(lvl).expect("invalid brotli compression level"), + ), + None => Compression::BROTLI(Default::default()), + }, + CompressionArgs::Lz4 => Compression::LZ4, + CompressionArgs::Zstd => match level { + Some(lvl) => Compression::ZSTD( + ZstdLevel::try_new(lvl as i32).expect("invalid zstd compression level"), + ), + None => Compression::ZSTD(Default::default()), + }, + CompressionArgs::Lz4Raw => Compression::LZ4_RAW, + } +} + +#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum, Debug)] +enum EncodingArgs { + /// Default byte encoding. + Plain, + + /// **Deprecated** dictionary encoding. + PlainDictionary, + + /// Group packed run length encoding. + Rle, + + /// **Deprecated** Bit-packed encoding. + BitPacked, + + /// Delta encoding for integers, either INT32 or INT64. + DeltaBinaryPacked, + + /// Encoding for byte arrays to separate the length values and the data. + DeltaLengthByteArray, + + /// Incremental encoding for byte arrays. + DeltaByteArray, + + /// Dictionary encoding. + RleDictionary, + + /// Encoding for fixed-width data. + ByteStreamSplit, +} + +#[allow(deprecated)] +impl From for Encoding { + fn from(value: EncodingArgs) -> Self { + match value { + EncodingArgs::Plain => Self::PLAIN, + EncodingArgs::PlainDictionary => Self::PLAIN_DICTIONARY, + EncodingArgs::Rle => Self::RLE, + EncodingArgs::BitPacked => Self::BIT_PACKED, + EncodingArgs::DeltaBinaryPacked => Self::DELTA_BINARY_PACKED, + EncodingArgs::DeltaLengthByteArray => Self::DELTA_LENGTH_BYTE_ARRAY, + EncodingArgs::DeltaByteArray => Self::DELTA_BYTE_ARRAY, + EncodingArgs::RleDictionary => Self::RLE_DICTIONARY, + EncodingArgs::ByteStreamSplit => Self::BYTE_STREAM_SPLIT, + } + } +} + +#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum, Debug)] +enum EnabledStatisticsArgs { + /// Compute no statistics + None, + + /// Compute chunk-level statistics but not page-level + Chunk, + + /// Compute page-level and chunk-level statistics + Page, +} + +impl From for EnabledStatistics { + fn from(value: EnabledStatisticsArgs) -> Self { + match value { + EnabledStatisticsArgs::None => Self::None, + EnabledStatisticsArgs::Chunk => Self::Chunk, + EnabledStatisticsArgs::Page => Self::Page, + } + } +} + +#[derive(Clone, Copy, Debug)] +enum WriterVersionArgs { + Parquet1_0, + Parquet2_0, +} + +impl ValueEnum for WriterVersionArgs { + fn value_variants<'a>() -> &'a [Self] { + &[Self::Parquet1_0, Self::Parquet2_0] + } + + fn to_possible_value(&self) -> Option { + match self { + WriterVersionArgs::Parquet1_0 => Some(PossibleValue::new("1.0")), + WriterVersionArgs::Parquet2_0 => Some(PossibleValue::new("2.0")), + } + } +} + +impl From for WriterVersion { + fn from(value: WriterVersionArgs) -> Self { + match value { + WriterVersionArgs::Parquet1_0 => Self::PARQUET_1_0, + WriterVersionArgs::Parquet2_0 => Self::PARQUET_2_0, + } + } +} + +#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum, Debug)] +enum BloomFilterPositionArgs { + /// Write Bloom Filters of each row group right after the row group + AfterRowGroup, + + /// Write Bloom Filters at the end of the file + End, +} + +impl From for BloomFilterPosition { + fn from(value: BloomFilterPositionArgs) -> Self { + match value { + BloomFilterPositionArgs::AfterRowGroup => Self::AfterRowGroup, + BloomFilterPositionArgs::End => Self::End, + } + } +} + +#[derive(Debug, Parser)] +#[clap(author, version, about("Read and write parquet file with potentially different settings"), long_about = None)] +struct Args { + /// Path to input parquet file. + #[clap(short, long)] + input: String, + + /// Path to output parquet file. + #[clap(short, long)] + output: String, + + /// Compression used for all columns. + #[clap(long, value_enum)] + compression: Option, + + /// Compression level for gzip/brotli/zstd. + #[clap(long)] + compression_level: Option, + + /// Encoding used for all columns, if dictionary is not enabled. + #[clap(long, value_enum)] + encoding: Option, + + /// Sets flag to enable/disable dictionary encoding for all columns. + #[clap(long)] + dictionary_enabled: Option, + + /// Sets best effort maximum dictionary page size, in bytes. + #[clap(long)] + dictionary_page_size_limit: Option, + + /// Sets maximum number of rows in a row group. + #[clap(long)] + max_row_group_size: Option, + + /// Sets best effort maximum number of rows in a data page. + #[clap(long)] + data_page_row_count_limit: Option, + + /// Sets best effort maximum size of a data page in bytes. + #[clap(long)] + data_page_size_limit: Option, + + /// Sets the max length of min/max statistics in row group and data page + /// header statistics for all columns. + /// + /// Applicable only if statistics are enabled. + #[clap(long)] + statistics_truncate_length: Option, + + /// Sets the max length of min/max statistics in the column index. + /// + /// Applicable only if statistics are enabled. + #[clap(long)] + column_index_truncate_length: Option, + + /// Write statistics to the data page headers? + /// + /// Setting this true will also enable page level statistics. + #[clap(long)] + write_page_header_statistics: Option, + + /// Write path_in_schema to the column metadata. + #[clap(long)] + write_path_in_schema: Option, + + /// Sets whether bloom filter is enabled for all columns. + #[clap(long)] + bloom_filter_enabled: Option, + + /// Sets bloom filter false positive probability (fpp) for all columns. + #[clap(long)] + bloom_filter_fpp: Option, + + /// Sets number of distinct values (ndv) for bloom filter for all columns. + #[clap(long)] + bloom_filter_ndv: Option, + + /// Sets the position of bloom filter + #[clap(long)] + bloom_filter_position: Option, + + /// Sets flag to enable/disable statistics for all columns. + #[clap(long)] + statistics_enabled: Option, + + /// Sets writer version. + #[clap(long)] + writer_version: Option, + + /// Sets write batch size. + #[clap(long)] + write_batch_size: Option, + + /// Sets whether to coerce Arrow types to match Parquet specification + #[clap(long)] + coerce_types: Option, +} + +fn main() { + let args = Args::parse(); + + // read key-value metadata + let parquet_reader = + SerializedFileReader::new(File::open(&args.input).expect("Unable to open input file")) + .expect("Failed to create reader"); + let kv_md = parquet_reader + .metadata() + .file_metadata() + .key_value_metadata() + .cloned(); + + // create actual parquet reader + let parquet_reader = ParquetRecordBatchReaderBuilder::try_new( + File::open(args.input).expect("Unable to open input file"), + ) + .expect("parquet open") + .build() + .expect("parquet open"); + + let mut writer_properties_builder = WriterProperties::builder().set_key_value_metadata(kv_md); + + if let Some(value) = args.compression { + let compression = compression_from_args(value, args.compression_level); + writer_properties_builder = writer_properties_builder.set_compression(compression); + } + + // setup encoding + if let Some(value) = args.encoding { + writer_properties_builder = writer_properties_builder.set_encoding(value.into()); + } + if let Some(value) = args.dictionary_enabled { + writer_properties_builder = writer_properties_builder.set_dictionary_enabled(value); + } + if let Some(value) = args.dictionary_page_size_limit { + writer_properties_builder = writer_properties_builder.set_dictionary_page_size_limit(value); + } + + if let Some(value) = args.max_row_group_size { + writer_properties_builder = + writer_properties_builder.set_max_row_group_row_count(Some(value)); + } + if let Some(value) = args.data_page_row_count_limit { + writer_properties_builder = writer_properties_builder.set_data_page_row_count_limit(value); + } + if let Some(value) = args.data_page_size_limit { + writer_properties_builder = writer_properties_builder.set_data_page_size_limit(value); + } + if let Some(value) = args.dictionary_page_size_limit { + writer_properties_builder = writer_properties_builder.set_dictionary_page_size_limit(value); + } + if let Some(value) = args.statistics_truncate_length { + writer_properties_builder = + writer_properties_builder.set_statistics_truncate_length(Some(value)); + } + if let Some(value) = args.column_index_truncate_length { + writer_properties_builder = + writer_properties_builder.set_column_index_truncate_length(Some(value)); + } + if let Some(value) = args.bloom_filter_enabled { + writer_properties_builder = writer_properties_builder.set_bloom_filter_enabled(value); + + if value { + if let Some(value) = args.bloom_filter_fpp { + writer_properties_builder = writer_properties_builder.set_bloom_filter_fpp(value); + } + if let Some(value) = args.bloom_filter_ndv { + writer_properties_builder = + writer_properties_builder.set_bloom_filter_max_ndv(value); + } + if let Some(value) = args.bloom_filter_position { + writer_properties_builder = + writer_properties_builder.set_bloom_filter_position(value.into()); + } + } + } + if let Some(value) = args.statistics_enabled { + writer_properties_builder = writer_properties_builder.set_statistics_enabled(value.into()); + } + // set this after statistics_enabled + if let Some(value) = args.write_page_header_statistics { + writer_properties_builder = + writer_properties_builder.set_write_page_header_statistics(value); + if value { + writer_properties_builder = + writer_properties_builder.set_statistics_enabled(EnabledStatistics::Page); + } + } + if let Some(value) = args.writer_version { + writer_properties_builder = writer_properties_builder.set_writer_version(value.into()); + } + if let Some(value) = args.coerce_types { + writer_properties_builder = writer_properties_builder.set_coerce_types(value); + } + if let Some(value) = args.write_path_in_schema { + writer_properties_builder = writer_properties_builder.set_write_path_in_schema(value); + } + if let Some(value) = args.write_batch_size { + writer_properties_builder = writer_properties_builder.set_write_batch_size(value); + } + let writer_properties = writer_properties_builder.build(); + let mut parquet_writer = ArrowWriter::try_new( + File::create(&args.output).expect("Unable to open output file"), + parquet_reader.schema(), + Some(writer_properties), + ) + .expect("create arrow writer"); + + for maybe_batch in parquet_reader { + let batch = maybe_batch.expect("reading batch"); + parquet_writer.write(&batch).expect("writing data"); + } + + parquet_writer.close().expect("finalizing file"); +} diff --git a/vendor/parquet/src/bin/parquet-rowcount.rs b/vendor/parquet/src/bin/parquet-rowcount.rs new file mode 100644 index 0000000..07e4bd1 --- /dev/null +++ b/vendor/parquet/src/bin/parquet-rowcount.rs @@ -0,0 +1,69 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Binary file to return the number of rows found from Parquet file(s). +//! +//! # Install +//! +//! `parquet-rowcount` can be installed using `cargo`: +//! ``` +//! cargo install parquet --features=cli +//! ``` +//! After this `parquet-rowcount` should be available: +//! ``` +//! parquet-rowcount XYZ.parquet +//! ``` +//! +//! The binary can also be built from the source code and run as follows: +//! ``` +//! cargo run --features=cli --bin parquet-rowcount XYZ.parquet ABC.parquet ZXC.parquet +//! ``` +//! +//! Note that `parquet-rowcount` reads full file schema, no projection or filtering is +//! applied. + +use clap::Parser; +use parquet::file::reader::{FileReader, SerializedFileReader}; +use std::{fs::File, path::Path}; + +#[derive(Debug, Parser)] +#[clap(author, version, about("Binary file to return the number of rows found from Parquet file(s)"), long_about = None)] +struct Args { + #[clap( + number_of_values(1), + help("List of Parquet files to read from separated by space") + )] + file_paths: Vec, +} + +fn main() { + let args = Args::parse(); + + for filename in args.file_paths { + let path = Path::new(&filename); + let file = File::open(path).expect("Unable to open file"); + let parquet_reader = SerializedFileReader::new(file).expect("Unable to read file"); + let row_group_metadata = parquet_reader.metadata().row_groups(); + let mut total_num_rows = 0; + + for group_metadata in row_group_metadata { + total_num_rows += group_metadata.num_rows(); + } + + eprintln!("File {filename}: rowcount={total_num_rows}"); + } +} diff --git a/vendor/parquet/src/bin/parquet-schema.rs b/vendor/parquet/src/bin/parquet-schema.rs new file mode 100644 index 0000000..bfcb77d --- /dev/null +++ b/vendor/parquet/src/bin/parquet-schema.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. + +//! Binary file to print the schema and metadata of a Parquet file. +//! +//! # Install +//! +//! `parquet-schema` can be installed using `cargo`: +//! ``` +//! cargo install parquet --features=cli +//! ``` +//! After this `parquet-schema` should be available: +//! ``` +//! parquet-schema XYZ.parquet +//! ``` +//! +//! The binary can also be built from the source code and run as follows: +//! ``` +//! cargo run --features=cli --bin parquet-schema XYZ.parquet +//! ``` +//! +//! Note that `verbose` is an optional boolean flag that allows to print schema only, +//! when not provided or print full file metadata when provided. + +use clap::Parser; +use parquet::{ + file::reader::{FileReader, SerializedFileReader}, + schema::printer::{print_file_metadata, print_parquet_metadata}, +}; +use std::{fs::File, path::Path}; + +#[derive(Debug, Parser)] +#[clap(author, version, about("Binary file to print the schema and metadata of a Parquet file"), long_about = None)] +struct Args { + #[clap(help("Path to the parquet file"))] + file_path: String, + #[clap(short, long, help("Enable printing full file metadata"))] + verbose: bool, +} + +fn main() { + let args = Args::parse(); + let filename = args.file_path; + let path = Path::new(&filename); + let file = File::open(path).expect("Unable to open file"); + let verbose = args.verbose; + + match SerializedFileReader::new(file) { + Err(e) => panic!("Error when parsing Parquet file: {e}"), + Ok(parquet_reader) => { + let metadata = parquet_reader.metadata(); + println!("Metadata for file: {}", &filename); + println!(); + if verbose { + print_parquet_metadata(&mut std::io::stdout(), metadata); + } else { + print_file_metadata(&mut std::io::stdout(), metadata.file_metadata()); + } + } + } +} diff --git a/vendor/parquet/src/bin/parquet-show-bloom-filter.rs b/vendor/parquet/src/bin/parquet-show-bloom-filter.rs new file mode 100644 index 0000000..436436d --- /dev/null +++ b/vendor/parquet/src/bin/parquet-show-bloom-filter.rs @@ -0,0 +1,160 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Binary file to read bloom filter data from a Parquet file. +//! +//! # Install +//! +//! `parquet-show-bloom-filter` can be installed using `cargo`: +//! ``` +//! cargo install parquet --features=cli +//! ``` +//! After this `parquet-show-bloom-filter` should be available: +//! ``` +//! parquet-show-bloom-filter XYZ.parquet id a +//! ``` +//! +//! The binary can also be built from the source code and run as follows: +//! ``` +//! cargo run --features=cli --bin parquet-show-bloom-filter -- --file-name XYZ.parquet --column id --values a +//! ``` + +mod cli; + +use clap::Parser; +use parquet::basic::Type; +use parquet::bloom_filter::Sbbf; +use parquet::file::metadata::ColumnChunkMetaData; +use parquet::file::{ + properties::ReaderProperties, + reader::{FileReader, SerializedFileReader}, + serialized_reader::ReadOptionsBuilder, +}; +use std::{fs::File, path::Path}; + +#[derive(Debug, Parser)] +#[clap(author, version, about("Binary file to read bloom filter data from a Parquet file"), long_about = None)] +struct Args { + #[clap(help("Path to the parquet file"))] + file_name: String, + #[clap(help(concat!( + "Dot-separated column path to check bloom filter for. Literal dots can be escaped as `\\.`. ", + "Only string typed columns or columns with an Int32 or Int64 physical type are supported." + ) + ))] + column: String, + #[clap( + help( + "Check if the given values match bloom filter, the values will be parsed to the physical type of the column" + ), + required = true + )] + values: Vec, +} + +fn main() { + let args = Args::parse(); + let file_name = args.file_name; + let path = Path::new(&file_name); + let file = File::open(path).expect("Unable to open file"); + + let file_reader = SerializedFileReader::new_with_options( + file, + ReadOptionsBuilder::new() + .with_reader_properties( + ReaderProperties::builder() + .set_read_bloom_filter(true) + .build(), + ) + .build(), + ) + .expect("Unable to open file as Parquet"); + let metadata = file_reader.metadata(); + let column_path = match cli::parse_column_path(&args.column) { + Ok(path) => path, + Err(err) => { + eprintln!("{}", err); + std::process::exit(1); + } + }; + + for (ri, row_group) in metadata.row_groups().iter().enumerate() { + println!("Row group #{ri}"); + println!("{}", "=".repeat(80)); + if let Some((column_index, column)) = row_group + .columns() + .iter() + .enumerate() + .find(|(_, column)| column.column_path().parts() == column_path.as_slice()) + { + let row_group_reader = file_reader + .get_row_group(ri) + .expect("Unable to read row group"); + if let Some(sbbf) = row_group_reader.get_column_bloom_filter(column_index) { + args.values.iter().for_each(|value| { + match check_filter(sbbf, value, column) { + Ok(present) => { + println!( + "Value {} is {} in bloom filter", + value, + if present { "present" } else { "absent" } + ) + } + Err(err) => { + println!("{err}"); + } + }; + }); + } else { + println!("No bloom filter found for column {}", args.column); + } + } else { + println!( + "No column named {} found, candidate columns are: {}", + args.column, + row_group + .columns() + .iter() + .map(|c| c.column_path().string()) + .collect::>() + .join(", ") + ); + } + } +} + +fn check_filter(sbbf: &Sbbf, value: &String, column: &ColumnChunkMetaData) -> Result { + match column.column_type() { + Type::INT32 => { + let value: i32 = value + .parse() + .map_err(|e| format!("Unable to parse value '{value}' to i32: {e}"))?; + Ok(sbbf.check(&value)) + } + Type::INT64 => { + let value: i64 = value + .parse() + .map_err(|e| format!("Unable to parse value '{value}' to i64: {e}"))?; + Ok(sbbf.check(&value)) + } + Type::BYTE_ARRAY => Ok(sbbf.check(&value.as_str())), + _ => Err(format!( + "Unsupported column type for checking bloom filter: {}", + column.column_type() + )), + } +} diff --git a/vendor/parquet/src/bloom_filter/mod.rs b/vendor/parquet/src/bloom_filter/mod.rs new file mode 100644 index 0000000..8e89ba4 --- /dev/null +++ b/vendor/parquet/src/bloom_filter/mod.rs @@ -0,0 +1,1240 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Bloom filter implementation specific to Parquet, as described +//! in the [spec][parquet-bf-spec]. +//! +//! # Bloom Filter Size +//! +//! Parquet uses the [Split Block Bloom Filter][sbbf-paper] (SBBF) as its bloom filter +//! implementation. For each column upon which bloom filters are enabled, the offset and length of an SBBF +//! is stored in the metadata for each row group in the parquet file. The size of each filter is +//! initialized using a calculation based on the desired number of distinct values (NDV) and false +//! positive probability (FPP). The FPP for a SBBF can be approximated as[1][bf-formulae]: +//! +//! ```text +//! f = (1 - e^(-k * n / m))^k +//! ``` +//! +//! Where, `f` is the FPP, `k` the number of hash functions, `n` the NDV, and `m` the total number +//! of bits in the bloom filter. This can be re-arranged to determine the total number of bits +//! required to achieve a given FPP and NDV: +//! +//! ```text +//! m = -k * n / ln(1 - f^(1/k)) +//! ``` +//! +//! SBBFs use eight hash functions to cleanly fit in SIMD lanes[2][sbbf-paper], therefore +//! `k` is set to 8. The SBBF will spread those `m` bits accross a set of `b` blocks that +//! are each 256 bits, i.e., 32 bytes, in size. The number of blocks is chosen as: +//! +//! ```text +//! b = NP2(m/8) / 32 +//! ``` +//! +//! Where, `NP2` denotes *the next power of two*, and `m` is divided by 8 to be represented as bytes. +//! +//! Here is a table of calculated sizes for various FPP and NDV: +//! +//! | NDV | FPP | b | Size (KB) | +//! |-----------|-----------|---------|-----------| +//! | 10,000 | 0.1 | 256 | 8 | +//! | 10,000 | 0.01 | 512 | 16 | +//! | 10,000 | 0.001 | 1,024 | 32 | +//! | 10,000 | 0.0001 | 1,024 | 32 | +//! | 100,000 | 0.1 | 4,096 | 128 | +//! | 100,000 | 0.01 | 4,096 | 128 | +//! | 100,000 | 0.001 | 8,192 | 256 | +//! | 100,000 | 0.0001 | 16,384 | 512 | +//! | 100,000 | 0.00001 | 16,384 | 512 | +//! | 1,000,000 | 0.1 | 32,768 | 1,024 | +//! | 1,000,000 | 0.01 | 65,536 | 2,048 | +//! | 1,000,000 | 0.001 | 65,536 | 2,048 | +//! | 1,000,000 | 0.0001 | 131,072 | 4,096 | +//! | 1,000,000 | 0.00001 | 131,072 | 4,096 | +//! | 1,000,000 | 0.000001 | 262,144 | 8,192 | +//! +//! # Structure: Filter → Blocks → Words → Bits +//! +//! An SBBF is an array of **blocks**. Each block is 256 bits (32 bytes), +//! divided into eight 32-bit **words**. A word is just a `u32` — an array of +//! 32 individual bits that can each be "set" (1) or "not set" (0). +//! +//! ```text +//! Sbbf (the whole filter) +//! ┌──────────┬──────────┬──────────┬─── ─── ──┬──────────┐ +//! │ Block 0 │ Block 1 │ Block 2 │ ... │ Block N-1│ +//! └──────────┴──────────┴──────────┴─── ─── ──┴──────────┘ +//! │ +//! ▼ +//! One Block = 256 bits = 8 words +//! ┌────────┬────────┬────────┬────────┬────────┬────────┬────────┬────────┐ +//! │ word 0 │ word 1 │ word 2 │ word 3 │ word 4 │ word 5 │ word 6 │ word 7 │ +//! │ (u32) │ (u32) │ (u32) │ (u32) │ (u32) │ (u32) │ (u32) │ (u32) │ +//! └────────┴────────┴────────┴────────┴────────┴────────┴────────┴────────┘ +//! │ +//! ▼ +//! One Word = 32 individual bits +//! ┌─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┬─┐ +//! │0│0│1│0│0│0│0│0│0│0│0│0│0│0│0│0│0│0│0│0│0│0│0│0│0│0│0│0│0│0│0│0│ ← bit 29 is set +//! └─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┘ +//! ``` +//! +//! **Inserting** a value hashes it to a 64-bit number, then: +//! 1. The upper 32 bits pick which **block** (via `Sbbf::hash_to_block_index`). +//! 2. The lower 32 bits pick one bit position in each of the 8 **words** (via `Block::mask`). +//! So each insert sets exactly **8 bits** (one per word) in a single block. +//! +//! **Checking** does the same two steps and returns `true` only if all 8 bits +//! are already set — meaning the value was *probably* inserted (or is a false +//! positive). +//! +//! # Bloom Filter Folding +//! +//! After inserting all values into a bloom filter it can be "folded" to minimize it's size. +//! See [`Sbbf::fold_to_target_fpp`] for details on the algorithm and its mathematical basis. +//! +//! [parquet-bf-spec]: https://github.com/apache/parquet-format/blob/master/BloomFilter.md +//! [sbbf-paper]: https://arxiv.org/pdf/2101.01719 +//! [bf-formulae]: http://tfk.mit.edu/pdf/bloom.pdf + +use crate::basic::{BloomFilterAlgorithm, BloomFilterCompression, BloomFilterHash}; +use crate::data_type::AsBytes; +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::ColumnChunkMetaData; +use crate::file::reader::ChunkReader; +use crate::parquet_thrift::{ + ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol, ThriftCompactOutputProtocol, + ThriftSliceInputProtocol, WriteThrift, WriteThriftField, +}; +use crate::thrift_struct; +use bytes::Bytes; +use std::io::Write; +use twox_hash::XxHash64; + +/// Salt as defined in the [spec](https://github.com/apache/parquet-format/blob/master/BloomFilter.md#technical-approach). +const SALT: [u32; 8] = [ + 0x47b6137b_u32, + 0x44974d91_u32, + 0x8824ad5b_u32, + 0xa2b7289d_u32, + 0x705495c7_u32, + 0x2df1424b_u32, + 0x9efc4947_u32, + 0x5c6bfb31_u32, +]; + +thrift_struct!( +/// Bloom filter header is stored at beginning of Bloom filter data of each column +/// and followed by its bitset. +/// +pub struct BloomFilterHeader { + /// The size of bitset in bytes + 1: required i32 num_bytes; + /// The algorithm for setting bits. + 2: required BloomFilterAlgorithm algorithm; + /// The hash function used for Bloom filter + 3: required BloomFilterHash hash; + /// The compression used in the Bloom filter + 4: required BloomFilterCompression compression; +} +); + +/// A single 256-bit block, the basic unit of the Split Block Bloom Filter. +/// +/// A block is eight contiguous 32-bit **words** (`[u32; 8]`). +/// Each word is an independent bit-array of 32 positions: +/// +/// ```text +/// Block (256 bits total) +/// ┌────────┬────────┬────────┬────────┬────────┬────────┬────────┬────────┐ +/// │ word 0 │ word 1 │ word 2 │ word 3 │ word 4 │ word 5 │ word 6 │ word 7 │ +/// │ 32 bits│ 32 bits│ 32 bits│ 32 bits│ 32 bits│ 32 bits│ 32 bits│ 32 bits│ +/// └────────┴────────┴────────┴────────┴────────┴────────┴────────┴────────┘ +/// ``` +/// +/// When a value is inserted, [`Block::mask`] picks one bit in each word +/// (8 bits total), and those bits are OR'd in. When checking, we verify +/// all 8 bits are set. +#[derive(Debug, Copy, Clone)] +#[repr(transparent)] +struct Block([u32; 8]); +impl Block { + const ZERO: Block = Block([0; 8]); + + /// Produce a block where each of the 8 words has exactly one bit set. + /// + /// For each word `i` the bit position is derived from `x`: + /// + /// ```text + /// y = (x wrapping* SALT[i]) >> 27 // top 5 bits → value in 0..31 + /// word[i] = 1 << y // exactly one bit set per word + /// ``` + /// + /// Because only the top 5 bits survive the shift, each word picks one of + /// 32 possible bit positions. The eight SALT constants spread the choices + /// so different words usually light up different positions. + /// + /// Key property: the mask depends *only* on `x` (a u32) and the fixed + /// SALT constants — it is independent of the filter size. This is why + /// folding preserves bit patterns (see Lemma 2 in tests). + fn mask(x: u32) -> Self { + let mut result = [0_u32; 8]; + for i in 0..8 { + let y = x.wrapping_mul(SALT[i]); // spread bits via multiply + let y = y >> 27; // keep top 5 bits → 0..31 + result[i] = 1 << y; // set exactly that one bit + } + Self(result) + } + + #[inline] + #[cfg(not(target_endian = "little"))] + fn to_ne_bytes(self) -> [u8; 32] { + // SAFETY: [u32; 8] and [u8; 32] have the same size and neither has invalid bit patterns. + unsafe { std::mem::transmute(self.0) } + } + + #[inline] + #[cfg(not(target_endian = "little"))] + fn to_le_bytes(self) -> [u8; 32] { + self.swap_bytes().to_ne_bytes() + } + + #[inline] + #[cfg(not(target_endian = "little"))] + fn swap_bytes(mut self) -> Self { + self.0.iter_mut().for_each(|x| *x = x.swap_bytes()); + self + } + + /// OR the mask bits into this block (`block[i] |= mask[i]`). + /// + /// After insertion the 8 bits chosen by `mask(hash)` are guaranteed set; + /// bits previously set by other hashes are preserved. + fn insert(&mut self, hash: u32) { + let mask = Self::mask(hash); + for i in 0..8 { + self[i] |= mask[i]; + } + } + + /// Check membership: returns `true` when *every* bit from `mask(hash)` is + /// already set in this block (`block[i] & mask[i] != 0` for all 8 words). + /// + /// A `true` result means "probably present" (other inserts may have set + /// the same bits). A `false` is definitive — the value was never inserted. + fn check(&self, hash: u32) -> bool { + let mask = Self::mask(hash); + for i in 0..8 { + if self[i] & mask[i] == 0 { + return false; + } + } + true + } +} + +impl std::ops::Index for Block { + type Output = u32; + + #[inline] + fn index(&self, index: usize) -> &Self::Output { + self.0.index(index) + } +} + +impl std::ops::IndexMut for Block { + #[inline] + fn index_mut(&mut self, index: usize) -> &mut Self::Output { + self.0.index_mut(index) + } +} + +impl std::ops::BitOr for Block { + type Output = Self; + + #[inline] + fn bitor(self, rhs: Self) -> Self { + let mut result = [0u32; 8]; + for (i, item) in result.iter_mut().enumerate() { + *item = self.0[i] | rhs.0[i]; + } + Self(result) + } +} + +impl std::ops::BitOrAssign for Block { + #[inline] + fn bitor_assign(&mut self, rhs: Self) { + for i in 0..8 { + self.0[i] |= rhs.0[i]; + } + } +} + +impl Block { + /// Count the total number of set bits across all 8 words. + /// + /// Computes popcount on each word separately and sums. Keeping the popcount + /// separate from the OR allows the compiler to batch SIMD popcount instructions + /// (e.g., `cnt.16b` on ARM NEON) instead of interleaving them with OR operations. + #[inline] + fn count_ones(self) -> u32 { + // Written as a fold over the array so the compiler sees 8 independent + // popcount operations it can vectorize into cnt.16b + horizontal sum. + self.0.iter().map(|w| w.count_ones()).sum() + } +} + +/// A split block Bloom filter (SBBF). +/// +/// An SBBF partitions its bit space into fixed-size 256-bit (32-byte) blocks, each fitting in a +/// single CPU cache line. Each block contains eight 32-bit words, aligned with SIMD lanes for +/// parallel bit manipulation. When checking membership, only one block is accessed per query, +/// eliminating the cache-miss penalty of standard Bloom filters. +/// +/// ## Sizing and folding +/// +/// Filters are initially sized for a maximum expected number of distinct values (NDV) via +/// [`Sbbf::new_with_ndv_fpp`]. After all values are inserted, the filter is compacted by +/// calling [`Sbbf::fold_to_target_fpp`], which folds the filter down to the smallest size +/// that still meets the target false positive probability. +/// +/// The creation of this structure is based on the [`crate::file::properties::BloomFilterProperties`] +/// struct set via [`crate::file::properties::WriterProperties`] and is thus hidden by default. +#[derive(Debug, Clone)] +pub struct Sbbf(Vec); + +pub(crate) const SBBF_HEADER_SIZE_ESTIMATE: usize = 20; + +/// given an initial offset, and a byte buffer, try to read out a bloom filter header and return +/// both the header and the offset after it (for bitset). +pub(crate) fn chunk_read_bloom_filter_header_and_offset( + offset: u64, + buffer: Bytes, +) -> Result<(BloomFilterHeader, u64), ParquetError> { + let (header, length) = read_bloom_filter_header_and_length(buffer)?; + Ok((header, offset + length)) +} + +/// given a [Bytes] buffer, try to read out a bloom filter header and return both the header and +/// length of the header. +#[inline] +pub(crate) fn read_bloom_filter_header_and_length( + buffer: Bytes, +) -> Result<(BloomFilterHeader, u64), ParquetError> { + read_bloom_filter_header_and_length_from_bytes(buffer.as_ref()) +} + +/// Given a byte slice, try to read out a bloom filter header and return both the header and +/// length of the header. +#[inline] +fn read_bloom_filter_header_and_length_from_bytes( + buffer: &[u8], +) -> Result<(BloomFilterHeader, u64), ParquetError> { + let total_length = buffer.len(); + let mut prot = ThriftSliceInputProtocol::new(buffer); + let header = BloomFilterHeader::read_thrift(&mut prot) + .map_err(|e| ParquetError::General(format!("Could not read bloom filter header: {e}")))?; + Ok((header, (total_length - prot.as_slice().len()) as u64)) +} + +/// The minimum number of bytes for a bloom filter bitset. +pub const BITSET_MIN_LENGTH: usize = 32; +/// The maximum number of bytes for a bloom filter bitset. +pub const BITSET_MAX_LENGTH: usize = 128 * 1024 * 1024; + +#[inline] +fn optimal_num_of_bytes(num_bytes: usize) -> usize { + let num_bytes = num_bytes.min(BITSET_MAX_LENGTH); + let num_bytes = num_bytes.max(BITSET_MIN_LENGTH); + num_bytes.next_power_of_two() +} + +// see http://algo2.iti.kit.edu/documents/cacheefficientbloomfilters-jea.pdf +// given fpp = (1 - e^(-k * n / m)) ^ k +// we have m = - k * n / ln(1 - fpp ^ (1 / k)) +// where k = number of hash functions, m = number of bits, n = number of distinct values +#[inline] +fn num_of_bits_from_ndv_fpp(ndv: u64, fpp: f64) -> usize { + let num_bits = -8.0 * ndv as f64 / (1.0 - fpp.powf(1.0 / 8.0)).ln(); + num_bits as usize +} + +impl Sbbf { + /// Create a new [Sbbf] with given number of distinct values and false positive probability. + /// Will return an error if `fpp` is greater than or equal to 1.0 or less than 0.0. + pub fn new_with_ndv_fpp(ndv: u64, fpp: f64) -> Result { + if !(0.0..1.0).contains(&fpp) { + return Err(ParquetError::General(format!( + "False positive probability must be between 0.0 and 1.0, got {fpp}" + ))); + } + let num_bits = num_of_bits_from_ndv_fpp(ndv, fpp); + Ok(Self::new_with_num_of_bytes(num_bits / 8)) + } + + /// Create a new [Sbbf] with given number of bytes, the exact number of bytes will be adjusted + /// to the next power of two bounded by [BITSET_MIN_LENGTH] and [BITSET_MAX_LENGTH]. + pub fn new_with_num_of_bytes(num_bytes: usize) -> Self { + let num_bytes = optimal_num_of_bytes(num_bytes); + assert_eq!(num_bytes % size_of::(), 0); + let num_blocks = num_bytes / size_of::(); + let bitset = vec![Block::ZERO; num_blocks]; + Self(bitset) + } + + /// Creates a new [Sbbf] from a raw byte slice. + pub fn new(bitset: &[u8]) -> Self { + let data = bitset + .chunks_exact(4 * 8) + .map(|chunk| { + let mut block = Block::ZERO; + for (i, word) in chunk.chunks_exact(4).enumerate() { + block[i] = u32::from_le_bytes(word.try_into().unwrap()); + } + block + }) + .collect::>(); + Self(data) + } + + /// Write the bloom filter data (header and then bitset) to the output. This doesn't + /// flush the writer in order to boost performance of bulk writing all blocks. Caller + /// must remember to flush the writer. + /// This method usually is used in conjunction with [`Self::from_bytes`] for serialization/deserialization. + pub fn write(&self, mut writer: W) -> Result<(), ParquetError> { + let mut protocol = ThriftCompactOutputProtocol::new(&mut writer); + self.header().write_thrift(&mut protocol).map_err(|e| { + ParquetError::General(format!("Could not write bloom filter header: {e}")) + })?; + self.write_bitset(&mut writer)?; + Ok(()) + } + + /// Write the bitset in serialized form to the writer. + #[cfg(not(target_endian = "little"))] + pub fn write_bitset(&self, mut writer: W) -> Result<(), ParquetError> { + for block in &self.0 { + writer + .write_all(block.to_le_bytes().as_slice()) + .map_err(|e| { + ParquetError::General(format!("Could not write bloom filter bit set: {e}")) + })?; + } + Ok(()) + } + + /// Write the bitset in serialized form to the writer. + #[cfg(target_endian = "little")] + pub fn write_bitset(&self, mut writer: W) -> Result<(), ParquetError> { + // Safety: Block is repr(transparent) and [u32; 8] can be reinterpreted as [u8; 32]. + let slice = unsafe { + std::slice::from_raw_parts( + self.0.as_ptr() as *const u8, + self.0.len() * size_of::(), + ) + }; + writer.write_all(slice).map_err(|e| { + ParquetError::General(format!("Could not write bloom filter bit set: {e}")) + })?; + Ok(()) + } + + /// Create and populate [`BloomFilterHeader`] from this bitset for writing to serialized form + fn header(&self) -> BloomFilterHeader { + BloomFilterHeader { + // 8 i32 per block, 4 bytes per i32 + num_bytes: self.0.len() as i32 * 4 * 8, + algorithm: BloomFilterAlgorithm::BLOCK, + hash: BloomFilterHash::XXHASH, + compression: BloomFilterCompression::UNCOMPRESSED, + } + } + + /// Read a new bloom filter from the given offset in the given reader. + pub fn read_from_column_chunk( + column_metadata: &ColumnChunkMetaData, + reader: &R, + ) -> Result, ParquetError> { + let offset: u64 = if let Some(offset) = column_metadata.bloom_filter_offset() { + offset + .try_into() + .map_err(|_| ParquetError::General("Bloom filter offset is invalid".to_string()))? + } else { + return Ok(None); + }; + + let buffer = match column_metadata.bloom_filter_length() { + Some(length) => reader.get_bytes(offset, length as usize), + None => reader.get_bytes(offset, SBBF_HEADER_SIZE_ESTIMATE), + }?; + + let (header, bitset_offset) = + chunk_read_bloom_filter_header_and_offset(offset, buffer.clone())?; + + match header.algorithm { + BloomFilterAlgorithm::BLOCK => { + // this match exists to future proof the singleton algorithm enum + } + } + match header.compression { + BloomFilterCompression::UNCOMPRESSED => { + // this match exists to future proof the singleton compression enum + } + } + match header.hash { + BloomFilterHash::XXHASH => { + // this match exists to future proof the singleton hash enum + } + } + + let bitset = match column_metadata.bloom_filter_length() { + Some(_) => buffer.slice((bitset_offset - offset) as usize..), + None => { + let bitset_length: usize = header.num_bytes.try_into().map_err(|_| { + ParquetError::General("Bloom filter length is invalid".to_string()) + })?; + reader.get_bytes(bitset_offset, bitset_length)? + } + }; + + Ok(Some(Self::new(&bitset))) + } + + /// Map a 64-bit hash to a block index in `[0, num_blocks)`. + /// + /// Uses the "multiply-and-shift" trick (a fast alternative to modulo): + /// + /// ```text + /// upper32 = hash >> 32 // take the top 32 bits of the hash + /// index = (upper32 * N) >> 32 // ∈ [0, N) where N = num_blocks + /// ``` + /// + /// Why this matters for folding (Lemma 1): when N is a power of two and + /// you halve it to N/2, the index also halves: + /// + /// ```text + /// index_N = (upper32 * N) >> 32 + /// index_N/2 = (upper32 * N/2) >> 32 = index_N / 2 (integer division) + /// ``` + /// + /// So the block that held hash `h` in the big filter is at `index / 2` in + /// the half-sized filter — exactly where `fold` ORs it. + #[inline] + fn hash_to_block_index(&self, hash: u64) -> usize { + (((hash >> 32).saturating_mul(self.0.len() as u64)) >> 32) as usize + } + + /// Insert an [AsBytes] value into the filter + pub fn insert(&mut self, value: &T) { + self.insert_hash(hash_as_bytes(value)); + } + + /// Insert a hash into the filter + fn insert_hash(&mut self, hash: u64) { + let block_index = self.hash_to_block_index(hash); + self.0[block_index].insert(hash as u32) + } + + /// Check if an [AsBytes] value is probably present or definitely absent in the filter + pub fn check(&self, value: &T) -> bool { + self.check_hash(hash_as_bytes(value)) + } + + /// Check if a hash is in the filter. May return + /// true for values that was never inserted ("false positive") + /// but will always return false if a hash has not been inserted. + fn check_hash(&self, hash: u64) -> bool { + let block_index = self.hash_to_block_index(hash); + self.0[block_index].check(hash as u32) + } + + /// Return the total in memory size of this bloom filter in bytes + pub(crate) fn estimated_memory_size(&self) -> usize { + self.0.capacity() * std::mem::size_of::() + } + + /// Returns the number of blocks in this bloom filter. + pub fn num_blocks(&self) -> usize { + self.0.len() + } + + /// Fold the bloom filter down to the smallest size that still meets the target FPP + /// (False Positive Percentage). + /// + /// Folds the filter by merging groups of adjacent blocks via bitwise OR, where each + /// fold level halves the number of blocks. The fold count is chosen as the maximum + /// number of folds whose estimated FPP stays within `target_fpp`. The filter stops + /// at a minimum size of 1 block (32 bytes). + /// + /// ## How it works + /// + /// SBBFs use multiplicative hashing for block selection: + /// + /// ```text + /// block_index = ((hash >> 32) * num_blocks) >> 32 + /// ``` + /// + /// A single fold halves the block count: when `num_blocks` is halved, the new index + /// becomes `floor(original_index / 2)`, so blocks `2i` and `2i+1` map to the same + /// position. More generally, `k` folds reduce the block count by `2^k`, merging + /// groups of `2^k` adjacent blocks in a single pass: + /// + /// ```text + /// folded[i] = blocks[i*2^k] | blocks[i*2^k + 1] | ... | blocks[i*2^k + 2^k - 1] + /// ``` + /// + /// This differs from standard Bloom filter folding, which merges the two halves + /// (`B[i] | B[i + m/2]`) because standard filters use modular hashing where + /// `h(x) mod (m/2)` maps indices `i` and `i + m/2` to the same position. + /// + /// ## Correctness + /// + /// Folding **never introduces false negatives**. Every bit that was set in the original + /// filter remains set in the folded filter (via bitwise OR). The only effect is a controlled + /// increase in FPP as set bits from different blocks are merged together. + /// This is was originally proven in [Sailhan & Stehr 2012] for standard bloom filters and is empirically + /// demonstrated for SBBFs in Lemma 1 and Lemma 2 of the tests. + /// + /// ## References + /// + /// [Sailhan & Stehr 2012]: https://doi.org/10.1109/GreenCom.2012.16 + pub fn fold_to_target_fpp(&mut self, target_fpp: f64) { + let num_folds = self.num_folds_for_target_fpp(target_fpp); + if num_folds > 0 { + self.fold_n(num_folds); + } + } + + /// Determine how many folds can be applied without exceeding `target_fpp`. + /// + /// Computes the average per-block fill rate in a single pass (no allocation), + /// then analytically estimates the FPP at each fold level. + /// + /// When two blocks with independent fill rate `f` are OR'd, the expected fill + /// of the merged block is `1 - (1-f)^2`. After `k` folds (merging `2^k` blocks): + /// + /// ```text + /// f_k = 1 - (1 - f)^(2^k) + /// ``` + /// + /// SBBF membership checks perform `k=8` bit checks within one 256-bit block, + /// so the estimated FPP at fold level k is `f_k^8`. + fn num_folds_for_target_fpp(&self, target_fpp: f64) -> u32 { + let len = self.0.len(); + if len < 2 { + return 0; + } + + // Single pass: compute average per-block fill rate. + let total_set_bits: u64 = self.0.iter().map(|b| u64::from(b.count_ones())).sum(); + let avg_fill = total_set_bits as f64 / (len as f64 * 256.0); + + // Empty filter: can fold all the way down. + if avg_fill == 0.0 { + return len.trailing_zeros(); + } + + // Find max folds where estimated FPP stays within target. + // f_k = 1 - (1 - avg_fill)^(2^k), FPP_k = f_k^8 + assert!( + len.is_power_of_two(), + "Number of blocks must be a power of 2 for folding" + ); + let max_folds = len.trailing_zeros(); // log2(len) since len is power of 2 + let one_minus_f = 1.0 - avg_fill; + let mut num_folds = 0u32; + let mut one_minus_fk = one_minus_f; // (1-f)^1 initially + + for _ in 0..max_folds { + // After one more fold: (1-f)^(2^(k+1)) = ((1-f)^(2^k))^2 + one_minus_fk = one_minus_fk * one_minus_fk; + let fk = 1.0 - one_minus_fk; + let estimated_fpp = fk.powi(8); + if estimated_fpp > target_fpp { + break; + } + num_folds += 1; + } + + num_folds + } + + /// Fold the filter `num_folds` times in a single pass. + /// + /// Merges groups of `2^num_folds` adjacent blocks via bitwise OR, producing + /// `len / 2^num_folds` output blocks. The original allocation is reused. + /// + /// # Panics + /// + /// Panics if `num_folds` is 0 or would reduce the filter below 1 block. + fn fold_n(&mut self, num_folds: u32) { + assert!(num_folds > 0, "num_folds must be at least 1"); + let len = self.0.len(); + let group_size = 1usize << num_folds; + assert!( + group_size <= len, + "Cannot fold {num_folds} times: need at least {group_size} blocks, have {len}" + ); + let new_len = len / group_size; + for i in 0..new_len { + let start = i * group_size; + let mut merged = self.0[start]; + for j in 1..group_size { + merged |= self.0[start + j]; + } + self.0[i] = merged; + } + self.0.truncate(new_len); + } + + /// Reads a Sbff from Thrift encoded bytes + /// + /// # Examples + /// + /// ```no_run + /// # use parquet::errors::Result; + /// # use parquet::bloom_filter::Sbbf; + /// # fn main() -> Result<()> { + /// // In a real application, you would read serialized bloom filter bytes from a cache. + /// // This example demonstrates the deserialization process. + /// // Assuming you have bloom filter bytes from a Parquet file: + /// # let serialized_bytes: Vec = vec![]; + /// let bloom_filter = Sbbf::from_bytes(&serialized_bytes)?; + /// // Now you can use the bloom filter to check for values + /// if bloom_filter.check(&"some_value") { + /// println!("Value might be present (or false positive)"); + /// } else { + /// println!("Value is definitely not present"); + /// } + /// # Ok(()) + /// # } + /// ``` + pub fn from_bytes(bytes: &[u8]) -> Result { + let (header, header_len) = read_bloom_filter_header_and_length_from_bytes(bytes)?; + + let bitset_length: u64 = header + .num_bytes + .try_into() + .map_err(|_| ParquetError::General("Bloom filter length is invalid".to_string()))?; + + // Validate that bitset consumes all remaining bytes + if header_len + bitset_length != bytes.len() as u64 { + return Err(ParquetError::General(format!( + "Bloom filter data contains extra bytes: expected {} total bytes, got {}", + header_len + bitset_length, + bytes.len() + ))); + } + + let start = header_len as usize; + let end = (header_len + bitset_length) as usize; + let bitset = bytes + .get(start..end) + .ok_or_else(|| ParquetError::General("Bloom filter bitset is invalid".to_string()))?; + + Ok(Self::new(bitset)) + } +} + +// per spec we use xxHash with seed=0 +const SEED: u64 = 0; + +#[inline] +fn hash_as_bytes(value: &A) -> u64 { + XxHash64::oneshot(SEED, value.as_bytes()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_hash_bytes() { + assert_eq!(hash_as_bytes(""), 17241709254077376921); + } + + #[test] + fn test_mask_set_quick_check() { + for i in 0..1_000_000 { + let result = Block::mask(i); + assert!(result.0.iter().all(|&x| x.is_power_of_two())); + } + } + + #[test] + fn test_block_insert_and_check() { + for i in 0..1_000_000 { + let mut block = Block::ZERO; + block.insert(i); + assert!(block.check(i)); + } + } + + #[test] + fn test_sbbf_insert_and_check() { + let mut sbbf = Sbbf(vec![Block::ZERO; 1_000]); + for i in 0..1_000_000 { + sbbf.insert(&i); + assert!(sbbf.check(&i)); + } + } + + #[test] + fn test_with_fixture() { + // bloom filter produced by parquet-mr/spark for a column of i64 f"a{i}" for i in 0..10 + let bitset: &[u8] = &[ + 200, 1, 80, 20, 64, 68, 8, 109, 6, 37, 4, 67, 144, 80, 96, 32, 8, 132, 43, 33, 0, 5, + 99, 65, 2, 0, 224, 44, 64, 78, 96, 4, + ]; + let sbbf = Sbbf::new(bitset); + for a in 0..10i64 { + let value = format!("a{a}"); + assert!(sbbf.check(&value.as_str())); + } + } + + /// test the assumption that bloom filter header size should not exceed SBBF_HEADER_SIZE_ESTIMATE + /// essentially we are testing that the struct is packed with 4 i32 fields, each can be 1-5 bytes + /// so altogether it'll be 20 bytes at most. + #[test] + fn test_bloom_filter_header_size_assumption() { + let buffer: &[u8; 16] = &[21, 64, 28, 28, 0, 0, 28, 28, 0, 0, 28, 28, 0, 0, 0, 99]; + let ( + BloomFilterHeader { + algorithm, + compression, + hash, + num_bytes, + }, + read_length, + ) = read_bloom_filter_header_and_length(Bytes::copy_from_slice(buffer)).unwrap(); + assert_eq!(read_length, 15); + assert_eq!(algorithm, BloomFilterAlgorithm::BLOCK); + assert_eq!(compression, BloomFilterCompression::UNCOMPRESSED); + assert_eq!(hash, BloomFilterHash::XXHASH); + assert_eq!(num_bytes, 32_i32); + assert_eq!(20, SBBF_HEADER_SIZE_ESTIMATE); + } + + #[test] + fn test_optimal_num_of_bytes() { + for (input, expected) in &[ + (0, 32), + (9, 32), + (31, 32), + (32, 32), + (33, 64), + (99, 128), + (1024, 1024), + (999_000_000, 128 * 1024 * 1024), + ] { + assert_eq!(*expected, optimal_num_of_bytes(*input)); + } + } + + #[test] + fn test_num_of_bits_from_ndv_fpp() { + for (fpp, ndv, num_bits) in &[ + (0.1, 10, 57), + (0.01, 10, 96), + (0.001, 10, 146), + (0.1, 100, 577), + (0.01, 100, 968), + (0.001, 100, 1460), + (0.1, 1000, 5772), + (0.01, 1000, 9681), + (0.001, 1000, 14607), + (0.1, 10000, 57725), + (0.01, 10000, 96815), + (0.001, 10000, 146076), + (0.1, 100000, 577254), + (0.01, 100000, 968152), + (0.001, 100000, 1460769), + (0.1, 1000000, 5772541), + (0.01, 1000000, 9681526), + (0.001, 1000000, 14607697), + (1e-50, 1_000_000_000_000, 14226231280773240832), + ] { + assert_eq!(*num_bits, num_of_bits_from_ndv_fpp(*ndv, *fpp) as u64); + } + } + + #[test] + fn test_fold_n_halves_block_count() { + let mut sbbf = Sbbf::new_with_num_of_bytes(1024); // 32 blocks + assert_eq!(sbbf.num_blocks(), 32); + sbbf.fold_n(1); + assert_eq!(sbbf.num_blocks(), 16); + sbbf.fold_n(1); + assert_eq!(sbbf.num_blocks(), 8); + } + + #[test] + fn test_fold_preserves_inserted_values() { + // Create a large filter, insert values, fold, verify no false negatives + let mut sbbf = Sbbf::new_with_num_of_bytes(32 * 1024); // 32KB = 1024 blocks + let values: Vec = (0..1000).map(|i| format!("value_{i}")).collect(); + for v in &values { + sbbf.insert(v.as_str()); + } + + // Fold several times + let original_blocks = sbbf.num_blocks(); + sbbf.fold_to_target_fpp(0.05); + assert!( + sbbf.num_blocks() < original_blocks, + "should have folded at least once" + ); + + // All inserted values must still be found (no false negatives) + for v in &values { + assert!( + sbbf.check(v.as_str()), + "Value '{}' missing after folding (false negative!)", + v + ); + } + } + + #[test] + fn test_fold_to_target_fpp_stops_before_exceeding_target() { + let mut sbbf = Sbbf::new_with_num_of_bytes(64 * 1024); // 64KB + // Insert enough values to set some bits + for i in 0..5000 { + sbbf.insert(&i); + } + + let target_fpp = 0.01; + sbbf.fold_to_target_fpp(target_fpp); + + // After folding, the estimated FPP should be at or below target + // (the current state should not exceed target — we stopped before that would happen) + let total_bits = (sbbf.num_blocks() * 256) as f64; + let set_bits: u64 = sbbf + .0 + .iter() + .flat_map(|b| b.0.iter()) + .map(|w| w.count_ones() as u64) + .sum(); + let fill = set_bits as f64 / total_bits; + let current_fpp = fill.powi(8); + assert!( + current_fpp <= target_fpp, + "FPP {current_fpp} exceeds target {target_fpp}" + ); + } + + #[test] + fn test_fold_empty_filter_folds_to_minimum() { + // An empty filter has fill=0, so estimated FPP is always 0 — should fold all the way down + let mut sbbf = Sbbf::new_with_num_of_bytes(1024); // 32 blocks + sbbf.fold_to_target_fpp(0.01); + assert_eq!(sbbf.num_blocks(), 1); + } + + #[test] + #[should_panic(expected = "Cannot fold 1 times: need at least 2 blocks, have 1")] + fn test_fold_n_panics_at_minimum_size() { + let mut sbbf = Sbbf::new_with_num_of_bytes(32); // 1 block (minimum) + sbbf.fold_n(1); + } + + #[test] + fn test_sbbf_write_round_trip() { + // Create a bloom filter with a 32-byte bitset (minimum size) + let bitset_bytes = vec![0u8; 32]; + let mut original = Sbbf::new(&bitset_bytes); + + // Insert some test values + let test_values = ["hello", "world", "rust", "parquet", "bloom", "filter"]; + for value in &test_values { + original.insert(value); + } + + // Serialize to bytes + let mut output = Vec::new(); + original.write(&mut output).unwrap(); + + // Validate header was written correctly + let mut protocol = ThriftSliceInputProtocol::new(&output); + let header = BloomFilterHeader::read_thrift(&mut protocol).unwrap(); + assert_eq!(header.num_bytes, bitset_bytes.len() as i32); + assert_eq!(header.algorithm, BloomFilterAlgorithm::BLOCK); + assert_eq!(header.hash, BloomFilterHash::XXHASH); + assert_eq!(header.compression, BloomFilterCompression::UNCOMPRESSED); + + // Deserialize using from_bytes + let reconstructed = Sbbf::from_bytes(&output).unwrap(); + + // Most importantly: verify the bloom filter WORKS correctly after round-trip + // Note: bloom filters can have false positives, but should never have false negatives + // So we can't assert !check(), but we should verify inserted values are found + for value in &test_values { + assert!( + reconstructed.check(value), + "Value '{}' should be present after round-trip", + value + ); + } + } + + /// Prove that folding an SBBF by one level produces the exact same bits + /// as building a fresh filter at the smaller size from scratch. + /// + /// # What is folding? + /// + /// ```text + /// Original (N = 8 blocks): + /// ┌───┬───┬───┬───┬───┬───┬───┬───┐ + /// │ 0 │ 1 │ 2 │ 3 │ 4 │ 5 │ 6 │ 7 │ + /// └─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┘ + /// │ │ │ │ │ │ │ │ + /// └─OR┘ └─OR┘ └─OR┘ └─OR┘ pair-wise OR + /// │ │ │ │ + /// ┌───┴──┬────┴──┬────┴──┬────┴──┐ + /// │ 0|1 │ 2|3 │ 4|5 │ 6|7 │ Folded (N/2 = 4 blocks) + /// └──────┴───────┴───────┴───────┘ + /// ``` + /// + /// # Why folded == fresh (the two lemmas) + /// + /// An SBBF insertion does two things with a 64-bit hash `h`: + /// + /// 1. **Pick a block** — uses the upper 32 bits via `hash_to_block_index` + /// 2. **Set 8 bits in that block** — uses the lower 32 bits via `Block::mask` + /// + /// **Lemma 1 (block index halves):** `hash_to_block_index` uses + /// `(upper32 * N) >> 32`. When N halves, the index halves too: + /// `index_in(N/2) == index_in(N) / 2`. So the hash lands in the same + /// destination block whether you fold or build fresh. + /// + /// **Lemma 2 (mask is size-independent):** `Block::mask(h as u32)` depends + /// only on the lower 32 bits and the fixed SALT constants — the filter + /// size N is not involved. So the same 8 bits get set regardless. + /// + /// Combined: every hash sets the *same bits* in the *same destination + /// block* whether you fold or build fresh → filters are bit-identical. + #[test] + fn test_sbbf_folded_equals_fresh() { + let values = (0..5000).map(|i| format!("elem_{i}")).collect::>(); + let hashes = values + .iter() + .map(|v| hash_as_bytes(v.as_str())) + .collect::>(); + + for num_blocks in [64, 256, 1024] { + let half = num_blocks / 2; + + // Build a filter with N blocks and insert all values. + let mut original = Sbbf::new_with_num_of_bytes(num_blocks * 32); + assert_eq!(original.num_blocks(), num_blocks); + for &h in &hashes { + original.insert_hash(h); + } + + // --- Per-hash verification of the two lemmas --- + for &h in hashes.iter() { + // mask(h as u32) gives the 8-bit pattern that this hash sets + // inside whichever block it lands in. It uses only the lower + // 32 bits of h, so it's the same regardless of filter size. + let mask = Block::mask(h as u32); + + // Lemma 1 check: the block index in the original N-block + // filter, divided by 2, should equal the block index in a + // fresh N/2-block filter. + let orig_idx = original.hash_to_block_index(h); + assert!(orig_idx < num_blocks); + + let fresh_idx = { + let tmp = Sbbf(vec![Block::ZERO; half]); + tmp.hash_to_block_index(h) + }; + let folded_idx = orig_idx / 2; + assert_eq!( + fresh_idx, folded_idx, + "Lemma 1 failed: fresh index {fresh_idx} != folded index {folded_idx}" + ); + + // Lemma 2 check: every bit that mask wants to set is actually + // present in the original block. + // + // mask.0[w] has exactly ONE bit set (see Block::mask: `1 << y`). + // The block at orig_idx has many bits set from many inserts, so + // we can't test equality — we test that the specific mask bit is + // *present*: + // + // block_word & mask_word != 0 + // ⟺ "the one bit in the mask is set in the block" + // + // (Since mask_word has exactly 1 bit, `& mask != 0` is the same + // as `& mask == mask` — but `!= 0` reads more naturally.) + for w in 0..8 { + assert_ne!( + original.0[orig_idx].0[w] & mask.0[w], + 0, + "Lemma 2 failed: mask bit not set in word {w} of block {orig_idx}" + ); + } + } + + // --- Final bit-identical comparison --- + // Fold the original N-block filter down to N/2 blocks. + let mut folded = original.clone(); + folded.fold_n(1); + assert_eq!(folded.num_blocks(), half); + + // Build a fresh N/2-block filter with the same values. + let mut fresh = Sbbf::new_with_num_of_bytes(half * 32); + for &h in &hashes { + fresh.insert_hash(h); + } + + // By lemmas 1 + 2, every block should be bit-identical. + for j in 0..half { + assert_eq!( + folded.0[j].0, fresh.0[j].0, + "Block {j} differs after fold (N={num_blocks} → {half})" + ); + } + } + } + + /// Inductive multi-step folding: folding k times from N blocks produces + /// a filter bit-identical to a fresh N/2^k-block filter. + /// + /// `test_sbbf_folded_equals_fresh` proves the base case (one fold). + /// This test applies folds *repeatedly*, checking after each step: + /// + /// ```text + /// 512 ─fold→ 256 ─fold→ 128 ─…→ 1 (9 folds total) + /// ``` + /// + /// At each intermediate size we build a fresh filter and assert + /// bit-equality, confirming the lemma composes across folds. + #[test] + fn test_multi_step_fold() { + let values = (0..3000).map(|i| format!("x_{i}")).collect::>(); + + // Start with a 512-block filter. + let mut filter = Sbbf::new_with_num_of_bytes(512 * 32); + for v in &values { + filter.insert(v.as_str()); + } + + // Fold one level at a time, comparing against a fresh filter each step. + for expected_blocks in [256, 128, 64, 32, 16, 8, 4, 2, 1] { + filter.fold_n(1); + assert_eq!(filter.num_blocks(), expected_blocks); + + let mut fresh = Sbbf::new_with_num_of_bytes(expected_blocks * 32); + for v in &values { + fresh.insert(v.as_str()); + } + for (fb, rb) in filter.0.iter().zip(fresh.0.iter()) { + assert_eq!(fb.0, rb.0); + } + } + } + + /// test that the fpp estimator's overestimation doesn't cause fold_to_target_fpp + /// to produce significantly oversized filters + /// + /// compare the final size after folding against the theoretical optimal size + #[test] + fn test_fold_size_vs_optimal_fixed_size() { + for (ndv, target_fpp) in [ + (1000, 0.05), + (1000, 0.01), + (5000, 0.05), + (5000, 0.01), + (10000, 0.05), + ] { + let values = (0..ndv).map(|i| format!("d_{i}")).collect::>(); + + let mut folded = Sbbf::new_with_num_of_bytes(128 * 1024); // 128KB + for v in &values { + folded.insert(v.as_str()); + } + folded.fold_to_target_fpp(target_fpp); + + let folded_bytes = folded.num_blocks() * 32; + + let optimal = Sbbf::new_with_ndv_fpp(ndv as u64, target_fpp).unwrap(); + let optimal_bytes = optimal.num_blocks() * 32; + + let ratio = folded_bytes as f64 / optimal_bytes as f64; + + assert_eq!(ratio, 1.0); + } + } + + /// verify that a folded sbbf has the same empirical fpp as a fresh filter of the same size + /// this bridges the bit-identity proof above with the FPP guarantee from the folding paper + /// since the bits are identical, the false-positive rate must be too + /// + /// we measure fpp empirically by probing with values that were never inserted + /// and counting how many are incorrectly marked as present + #[test] + fn test_folded_fpp_matches_fresh_fpp() { + let ndv = 2000; + let num_probes = 50_000; + let inserted = (0..ndv) + .map(|i| format!("ins_{i}")) + .collect::>(); + + // probe values that were NOT inserted (different prefix guarantees no overlap) + let probes = (0..num_probes) + .map(|i| format!("probe_{i}")) + .collect::>(); + + // build a large filter and fold it down several times + let mut folded = Sbbf::new_with_num_of_bytes(512 * 32); // 512 blocks + for v in &inserted { + folded.insert(v.as_str()); + } + + // check FPP at each fold level + for expected_blocks in [256, 128, 64, 32, 16, 8, 4, 2, 1] { + folded.fold_n(1); + assert_eq!(folded.num_blocks(), expected_blocks); + + // build a fresh filter of the same size with the same values + let mut fresh = Sbbf::new_with_num_of_bytes(expected_blocks * 32); + for v in &inserted { + fresh.insert(v.as_str()); + } + + // measure empirical FPP on both + let mut folded_fp = 0u64; + let mut fresh_fp = 0u64; + for p in &probes { + if folded.check(p.as_str()) { + folded_fp += 1; + } + if fresh.check(p.as_str()) { + fresh_fp += 1; + } + } + + // bit-identity means these must be exactly equal + assert_eq!(folded_fp, fresh_fp); + } + } +} diff --git a/vendor/parquet/src/column/chunker/cdc.rs b/vendor/parquet/src/column/chunker/cdc.rs new file mode 100644 index 0000000..b40dd74 --- /dev/null +++ b/vendor/parquet/src/column/chunker/cdc.rs @@ -0,0 +1,2394 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#[cfg(feature = "arrow")] +use crate::column::writer::LevelDataRef; +use crate::errors::{ParquetError, Result}; +use crate::file::properties::CdcOptions; +use crate::schema::types::ColumnDescriptor; + +use super::CdcChunk; +use super::cdc_generated::{GEARHASH_TABLE, NUM_GEARHASH_TABLES}; + +/// CDC (Content-Defined Chunking) divides data into variable-sized chunks based on +/// content rather than fixed-size boundaries. +/// +/// For example, given this sequence of values in a column: +/// +/// ```text +/// File1: [1,2,3, 4,5,6, 7,8,9] +/// chunk1 chunk2 chunk3 +/// ``` +/// +/// If a value is inserted between 3 and 4: +/// +/// ```text +/// File2: [1,2,3,0, 4,5,6, 7,8,9] +/// new-chunk chunk2 chunk3 +/// ``` +/// +/// The chunking process adjusts to maintain stable boundaries across data modifications. +/// Each chunk defines a new parquet data page which is contiguously written to the file. +/// Since each page is compressed independently, the files' contents look like: +/// +/// ```text +/// File1: [Page1][Page2][Page3]... +/// File2: [Page4][Page2][Page3]... +/// ``` +/// +/// When uploaded to a content-addressable storage (CAS) system, the CAS splits the byte +/// stream into content-defined blobs with unique identifiers. Identical blobs are stored +/// only once, so Page2 and Page3 are deduplicated across File1 and File2. +/// +/// ## Implementation +/// +/// Only the parquet writer needs to be aware of content-defined chunking; the reader is +/// unaffected. Each parquet column writer holds a `ContentDefinedChunker` instance +/// depending on the writer's properties. The chunker's state is maintained across the +/// entire column without being reset between pages and row groups. +/// +/// This implements a [FastCDC]-inspired algorithm using gear hashing. The input data is +/// fed byte-by-byte into a rolling hash; when the hash matches a predefined mask, a new +/// chunk boundary candidate is recorded. To reduce the exponential variance of chunk +/// sizes inherent in a single gear hash, the algorithm requires **8 consecutive mask +/// matches** — each against a different pre-computed gear hash table — before committing +/// to a boundary. This [central-limit-theorem normalization] makes the chunk size +/// distribution approximately normal between `min_chunk_size` and `max_chunk_size`. +/// +/// The chunker receives the record-shredded column data (def_levels, rep_levels, values) +/// and iterates over the (def_level, rep_level, value) triplets while adjusting the +/// column-global rolling hash. Whenever the rolling hash matches, the chunker creates a +/// new chunk. For nested data (lists, maps, structs) chunk boundaries are restricted to +/// top-level record boundaries (`rep_level == 0`) so that a nested row is never split +/// across chunks. +/// +/// Note that boundaries are deterministically calculated exclusively based on the data +/// itself, so the same data always produces the same chunks given the same configuration. +/// +/// Ported from the C++ implementation in apache/arrow#45360 +/// (`cpp/src/parquet/chunker_internal.cc`). +/// +/// [FastCDC]: https://www.usenix.org/conference/atc16/technical-sessions/presentation/xia +/// [central-limit-theorem normalization]: https://www.cidrdb.org/cidr2023/papers/p43-low.pdf +#[derive(Debug)] +pub(crate) struct ContentDefinedChunker { + /// Maximum definition level for this column. + max_def_level: i16, + /// Maximum repetition level for this column. + max_rep_level: i16, + /// Definition level at the nearest REPEATED ancestor. + repeated_ancestor_def_level: i16, + + /// Minimum chunk size in bytes. + /// The rolling hash will not be updated until this size is reached for each chunk. + /// All data sent through the hash function counts towards the chunk size, including + /// definition and repetition levels if present. + min_chunk_size: i64, + /// Maximum chunk size in bytes. + /// A new chunk is created whenever the chunk size exceeds this value. The chunk size + /// distribution approximates a normal distribution between `min_chunk_size` and + /// `max_chunk_size`. Note that the parquet writer has a related `data_pagesize` + /// property that controls the maximum size of a parquet data page after encoding. + /// While setting `data_pagesize` smaller than `max_chunk_size` doesn't affect + /// chunking effectiveness, it results in more small parquet data pages. + max_chunk_size: i64, + /// Mask for matching against the rolling hash. + rolling_hash_mask: u64, + + /// Rolling hash state, never reset — initialized once for the entire column. + rolling_hash: u64, + /// Whether the rolling hash has matched the mask since the last chunk boundary. + has_matched: bool, + /// Current run count for the central-limit-theorem normalization. + nth_run: usize, + /// Current chunk size in bytes. + chunk_size: i64, +} + +impl ContentDefinedChunker { + pub fn new(desc: &ColumnDescriptor, options: &CdcOptions) -> Result { + let rolling_hash_mask = Self::calculate_mask( + options.min_chunk_size as i64, + options.max_chunk_size as i64, + options.norm_level, + )?; + Ok(Self { + max_def_level: desc.max_def_level(), + max_rep_level: desc.max_rep_level(), + repeated_ancestor_def_level: desc.repeated_ancestor_def_level(), + min_chunk_size: options.min_chunk_size as i64, + max_chunk_size: options.max_chunk_size as i64, + rolling_hash_mask, + rolling_hash: 0, + has_matched: false, + nth_run: 0, + chunk_size: 0, + }) + } + + /// Calculate the mask used to determine chunk boundaries from the rolling hash. + /// + /// The mask is calculated so that the expected chunk size distribution approximates + /// a normal distribution between min and max chunk sizes. + fn calculate_mask(min_chunk_size: i64, max_chunk_size: i64, norm_level: i32) -> Result { + if min_chunk_size < 0 { + return Err(ParquetError::General( + "min_chunk_size must be non-negative".to_string(), + )); + } + if max_chunk_size <= min_chunk_size { + return Err(ParquetError::General( + "max_chunk_size must be greater than min_chunk_size".to_string(), + )); + } + + let avg_chunk_size = (min_chunk_size + max_chunk_size) / 2; + // Target size after subtracting the min-size skip window and dividing by the + // number of hash tables (for central-limit-theorem normalization). + let target_size = (avg_chunk_size - min_chunk_size) / NUM_GEARHASH_TABLES as i64; + + // floor(log2(target_size)) — equivalent to C++ NumRequiredBits(target_size) - 1 + let mask_bits = if target_size > 0 { + 63 - target_size.leading_zeros() as i32 + } else { + 0 + }; + + let effective_bits = mask_bits - norm_level; + + if !(1..=63).contains(&effective_bits) { + return Err(ParquetError::General(format!( + "The number of bits in the CDC mask must be between 1 and 63, got {effective_bits}" + ))); + } + + // Create the mask by setting the top `effective_bits` bits. + Ok(u64::MAX << (64 - effective_bits)) + } + + /// Feed raw bytes into the rolling hash. + /// + /// The byte count always accumulates toward `chunk_size`, but the actual hash + /// update is skipped until `min_chunk_size` has been reached. This "skip window" + /// is the FastCDC optimization that prevents boundaries from appearing too early + /// in a chunk. + #[inline] + fn roll(&mut self, bytes: &[u8]) { + self.chunk_size += bytes.len() as i64; + if self.chunk_size < self.min_chunk_size { + return; + } + for &b in bytes { + self.rolling_hash = self + .rolling_hash + .wrapping_shl(1) + .wrapping_add(GEARHASH_TABLE[self.nth_run][b as usize]); + self.has_matched = + self.has_matched || ((self.rolling_hash & self.rolling_hash_mask) == 0); + } + } + + /// Feed exactly `N` bytes into the rolling hash (compile-time width). + /// + /// Like [`roll`](Self::roll), but the byte count is known at compile time, + /// allowing the compiler to unroll the inner loop. + #[inline(always)] + fn roll_fixed(&mut self, bytes: &[u8; N]) { + self.chunk_size += N as i64; + if self.chunk_size < self.min_chunk_size { + return; + } + for j in 0..N { + self.rolling_hash = self + .rolling_hash + .wrapping_shl(1) + .wrapping_add(GEARHASH_TABLE[self.nth_run][bytes[j] as usize]); + self.has_matched = + self.has_matched || ((self.rolling_hash & self.rolling_hash_mask) == 0); + } + } + + /// Feed a definition or repetition level (i16) into the rolling hash. + #[inline] + fn roll_level(&mut self, level: i16) { + self.roll_fixed(&level.to_le_bytes()); + } + + /// Check whether a new chunk boundary should be created. + /// + /// A boundary is created when **either** of two conditions holds: + /// + /// 1. **CLT normalization**: The rolling hash has matched the mask (`has_matched`) + /// *and* this is the 8th consecutive such match (`nth_run` reaches + /// `NUM_GEARHASH_TABLES`). Each match advances to the next gear hash table, so + /// 8 independent matches are required. A single hash table would yield + /// exponentially distributed chunk sizes; requiring 8 independent matches + /// approximates a normal (Gaussian) distribution by the central limit theorem. + /// + /// 2. **Hard size limit**: `chunk_size` has reached `max_chunk_size`. This caps + /// chunk size even if the CLT normalization sequence has not completed. + /// + /// Note: when `max_chunk_size` forces a boundary, `nth_run` is **not** reset, so + /// the CLT sequence continues from where it left off in the next chunk. This + /// matches the C++ behavior. + #[inline] + fn need_new_chunk(&mut self) -> bool { + if self.has_matched { + self.has_matched = false; + self.nth_run += 1; + if self.nth_run >= NUM_GEARHASH_TABLES { + self.nth_run = 0; + self.chunk_size = 0; + return true; + } + } + if self.chunk_size >= self.max_chunk_size { + self.chunk_size = 0; + return true; + } + false + } + + /// Compute chunk boundaries for the given column data. + /// + /// The chunking state is maintained across the entire column without being + /// reset between pages and row groups. This enables the chunking process to + /// be continued between different write calls. + /// + /// We go over the (def_level, rep_level, value) triplets one by one while + /// adjusting the column-global rolling hash based on the triplet. Whenever + /// the rolling hash matches a predefined mask it sets `has_matched` to true. + /// + /// After each triplet [`need_new_chunk`](Self::need_new_chunk) is called to + /// evaluate if we need to create a new chunk. + fn calculate( + &mut self, + def_levels: LevelDataRef<'_>, + rep_levels: LevelDataRef<'_>, + num_levels: usize, + mut roll_value: F, + ) -> Vec + where + F: FnMut(&mut Self, usize), + { + let has_def_levels = self.max_def_level > 0; + let has_rep_levels = self.max_rep_level > 0; + + let mut chunks = Vec::new(); + let mut prev_offset: usize = 0; + let mut prev_value_offset: usize = 0; + let mut value_offset: usize = 0; + + if !has_rep_levels && !has_def_levels { + // Fastest path: non-nested, non-null data. + // Every level corresponds to exactly one non-null value, so + // value_offset == level_offset and num_values == num_levels. + // + // Example: required Int32, array = [10, 20, 30] + // level: 0 1 2 + // value_offset: 0 1 2 + for offset in 0..num_levels { + roll_value(self, offset); + if self.need_new_chunk() { + chunks.push(CdcChunk { + level_offset: prev_offset, + num_levels: offset - prev_offset, + value_offset: prev_offset, + num_values: offset - prev_offset, + }); + prev_offset = offset; + } + } + prev_value_offset = prev_offset; + value_offset = num_levels; + } else if !has_rep_levels { + // Non-nested data with nulls. value_offset only increments for + // non-null values (def == max_def), so it diverges from the + // level offset when nulls are present. + // + // Example: optional Int32, array = [1, null, 2, null, 3] + // def_levels: [1, 0, 1, 0, 1] + // level: 0 1 2 3 4 + // value_offset: 0 1 2 (only increments on def==1) + #[allow(clippy::needless_range_loop)] + for offset in 0..num_levels { + let def_level = def_levels + .value_at(offset) + .expect("def_levels required when max_def_level > 0"); + self.roll_level(def_level); + if def_level == self.max_def_level { + // For non-nested data, the leaf array has one slot per + // level (nulls are array elements), so `offset` (the + // level index) is the correct array index for hashing. + roll_value(self, offset); + } + // Check boundary before incrementing value_offset so that + // num_values reflects only entries in the completed chunk. + if self.need_new_chunk() { + chunks.push(CdcChunk { + level_offset: prev_offset, + num_levels: offset - prev_offset, + value_offset: prev_value_offset, + num_values: value_offset - prev_value_offset, + }); + prev_offset = offset; + prev_value_offset = value_offset; + } + if def_level == self.max_def_level { + value_offset += 1; + } + } + } else { + // Nested data with nulls. Two counters are needed: + // + // leaf_offset: index into the leaf values array for hashing, + // incremented for all leaf slots (def >= repeated_ancestor_def_level), + // including null elements. + // + // value_offset: index into non_null_indices for chunk boundaries, + // incremented only for non-null leaf values (def == max_def_level). + // + // These diverge when nullable elements exist inside lists. + // + // Example: List with repeated_ancestor_def_level=2, max_def=3 + // row 0: [1, null, 2] (3 leaf slots, 2 non-null) + // row 1: [3] (1 leaf slot, 1 non-null) + // + // leaf array: [1, null, 2, 3] + // def_levels: [3, 2, 3, 3] + // rep_levels: [0, 1, 1, 0] + // + // level def leaf_offset value_offset action + // ───── ─── ─────────── ──────────── ────────────────────────── + // 0 3 0 0 roll_value(0), value++, leaf++ + // 1 2 1 1 leaf++ only (null element) + // 2 3 2 1 roll_value(2), value++, leaf++ + // 3 3 3 2 roll_value(3), value++, leaf++ + // + // roll_value(2) correctly indexes leaf array position 2 (value "2"). + // Using value_offset=1 would index position 1 (the null slot). + // + // Using value_offset for roll_value would hash the wrong array slot. + let mut leaf_offset: usize = 0; + + for offset in 0..num_levels { + let def_level = def_levels + .value_at(offset) + .expect("def_levels required for nested data"); + let rep_level = rep_levels + .value_at(offset) + .expect("rep_levels required for nested data"); + + self.roll_level(def_level); + self.roll_level(rep_level); + if def_level == self.max_def_level { + roll_value(self, leaf_offset); + } + + // Check boundary before incrementing value_offset so that + // num_values reflects only entries in the completed chunk. + if rep_level == 0 && self.need_new_chunk() { + let levels_to_write = offset - prev_offset; + if levels_to_write > 0 { + chunks.push(CdcChunk { + level_offset: prev_offset, + num_levels: levels_to_write, + value_offset: prev_value_offset, + num_values: value_offset - prev_value_offset, + }); + prev_offset = offset; + prev_value_offset = value_offset; + } + } + if def_level == self.max_def_level { + value_offset += 1; + } + if def_level >= self.repeated_ancestor_def_level { + leaf_offset += 1; + } + } + } + + // Add the last chunk if we have any levels left. + if prev_offset < num_levels { + chunks.push(CdcChunk { + level_offset: prev_offset, + num_levels: num_levels - prev_offset, + value_offset: prev_value_offset, + num_values: value_offset - prev_value_offset, + }); + } + + #[cfg(debug_assertions)] + self.validate_chunks(&chunks, num_levels, value_offset); + + chunks + } + + /// Compute CDC chunk boundaries by dispatching on the Arrow array's data type + /// to feed value bytes into the rolling hash. + #[cfg(feature = "arrow")] + pub(crate) fn get_arrow_chunks( + &mut self, + def_levels: LevelDataRef<'_>, + rep_levels: LevelDataRef<'_>, + array: &dyn arrow_array::Array, + ) -> Result> { + use arrow_array::cast::AsArray; + use arrow_schema::DataType; + + // For nested (list) data, null list entries can own non-zero child + // ranges in the leaf array, so `array.len()` may exceed the number of + // levels. Always drive the loop by the level count; fall back to the + // array length only when there are no levels at all. + let num_levels = match (def_levels.len(), rep_levels.len()) { + (0, 0) => array.len(), + (d, r) => d.max(r), + }; + + macro_rules! fixed_width { + ($N:literal) => {{ + let data = array.to_data(); + let buffer = data.buffers()[0].as_slice(); + let values = &buffer[data.offset() * $N..]; + self.calculate(def_levels, rep_levels, num_levels, |c, i| { + let offset = i * $N; + let slice = &values[offset..offset + $N]; + c.roll_fixed::<$N>(slice.try_into().unwrap()); + }) + }}; + } + + macro_rules! binary_like { + ($a:expr) => {{ + let a = $a; + self.calculate(def_levels, rep_levels, num_levels, |c, i| { + c.roll(a.value(i).as_ref()); + }) + }}; + } + + let dtype = array.data_type(); + let chunks = match dtype { + DataType::Null => self.calculate(def_levels, rep_levels, num_levels, |_, _| {}), + DataType::Boolean => { + let a = array.as_boolean(); + self.calculate(def_levels, rep_levels, num_levels, |c, i| { + c.roll_fixed(&[a.value(i) as u8]); + }) + } + DataType::Int8 | DataType::UInt8 => fixed_width!(1), + DataType::Int16 | DataType::UInt16 | DataType::Float16 => fixed_width!(2), + DataType::Int32 + | DataType::UInt32 + | DataType::Float32 + | DataType::Date32 + | DataType::Time32(_) + | DataType::Interval(arrow_schema::IntervalUnit::YearMonth) + | DataType::Decimal32(_, _) => fixed_width!(4), + DataType::Int64 + | DataType::UInt64 + | DataType::Float64 + | DataType::Date64 + | DataType::Time64(_) + | DataType::Timestamp(_, _) + | DataType::Duration(_) + | DataType::Interval(arrow_schema::IntervalUnit::DayTime) + | DataType::Decimal64(_, _) => fixed_width!(8), + DataType::Interval(arrow_schema::IntervalUnit::MonthDayNano) + | DataType::Decimal128(_, _) => fixed_width!(16), + DataType::Decimal256(_, _) => fixed_width!(32), + DataType::FixedSizeBinary(_) => binary_like!(array.as_fixed_size_binary()), + DataType::Binary => binary_like!(array.as_binary::()), + DataType::LargeBinary => binary_like!(array.as_binary::()), + DataType::Utf8 => binary_like!(array.as_string::()), + DataType::LargeUtf8 => binary_like!(array.as_string::()), + DataType::BinaryView => binary_like!(array.as_binary_view()), + DataType::Utf8View => binary_like!(array.as_string_view()), + DataType::Dictionary(_, _) => { + let dict = array.as_any_dictionary(); + self.get_arrow_chunks(def_levels, rep_levels, dict.keys())? + } + _ => { + return Err(ParquetError::General(format!( + "content-defined chunking is not supported for data type {dtype:?}", + ))); + } + }; + Ok(chunks) + } + + #[cfg(debug_assertions)] + fn validate_chunks(&self, chunks: &[CdcChunk], num_levels: usize, total_values: usize) { + assert!(!chunks.is_empty(), "chunks must be non-empty"); + + let first = &chunks[0]; + assert_eq!(first.level_offset, 0, "first chunk must start at level 0"); + assert_eq!(first.value_offset, 0, "first chunk must start at value 0"); + + let mut sum_levels = first.num_levels; + let mut sum_values = first.num_values; + for i in 1..chunks.len() { + let chunk = &chunks[i]; + let prev = &chunks[i - 1]; + assert!(chunk.num_levels > 0, "chunk must have levels"); + assert_eq!( + chunk.level_offset, + prev.level_offset + prev.num_levels, + "level offsets must be contiguous" + ); + assert_eq!( + chunk.value_offset, + prev.value_offset + prev.num_values, + "value offsets must be contiguous" + ); + sum_levels += chunk.num_levels; + sum_values += chunk.num_values; + } + assert_eq!(sum_levels, num_levels, "chunks must cover all levels"); + assert_eq!(sum_values, total_values, "chunks must cover all values"); + + let last = chunks.last().unwrap(); + assert_eq!( + last.level_offset + last.num_levels, + num_levels, + "last chunk must end at num_levels" + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::basic::Type as PhysicalType; + #[cfg(feature = "arrow")] + use crate::column::writer::LevelDataRef; + use crate::schema::types::{ColumnPath, Type}; + use std::sync::Arc; + + fn make_desc(max_def_level: i16, max_rep_level: i16) -> ColumnDescriptor { + let tp = Type::primitive_type_builder("col", PhysicalType::INT32) + .build() + .unwrap(); + ColumnDescriptor::new( + Arc::new(tp), + max_def_level, + max_rep_level, + ColumnPath::new(vec![]), + ) + } + + #[test] + fn test_calculate_mask_defaults() { + let mask = ContentDefinedChunker::calculate_mask(256 * 1024, 1024 * 1024, 0).unwrap(); + // avg = 640 KiB, target = (640-256)*1024/8 = 49152, log2(49152) = 15 + // mask = u64::MAX << (64 - 15) = top 15 bits set + let expected = u64::MAX << (64 - 15); + assert_eq!(mask, expected); + } + + #[test] + fn test_calculate_mask_with_norm_level() { + let mask = ContentDefinedChunker::calculate_mask(256 * 1024, 1024 * 1024, 1).unwrap(); + let expected = u64::MAX << (64 - 14); + assert_eq!(mask, expected); + } + + #[test] + fn test_calculate_mask_invalid() { + assert!(ContentDefinedChunker::calculate_mask(-1, 100, 0).is_err()); + assert!(ContentDefinedChunker::calculate_mask(100, 50, 0).is_err()); + assert!(ContentDefinedChunker::calculate_mask(100, 100, 0).is_err()); + } + + #[test] + fn test_non_nested_non_null_single_chunk() { + let options = CdcOptions { + min_chunk_size: 8, + max_chunk_size: 1024, + norm_level: 0, + }; + let mut chunker = ContentDefinedChunker::new(&make_desc(0, 0), &options).unwrap(); + + // Write a small amount of data — should produce exactly 1 chunk. + let num_values = 4; + let chunks = chunker.calculate( + LevelDataRef::Absent, + LevelDataRef::Absent, + num_values, + |c, i| { + c.roll_fixed::<4>(&(i as i32).to_le_bytes()); + }, + ); + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0].level_offset, 0); + assert_eq!(chunks[0].value_offset, 0); + assert_eq!(chunks[0].num_levels, 4); + } + + #[test] + fn test_max_chunk_size_forces_boundary() { + let options = CdcOptions { + min_chunk_size: 256, + max_chunk_size: 1024, + norm_level: 0, + }; + let mut chunker = ContentDefinedChunker::new(&make_desc(0, 0), &options).unwrap(); + + // Write enough data to exceed max_chunk_size multiple times. + // Each i32 = 4 bytes, max_chunk_size=1024, so ~256 values per chunk max. + let num_values = 2000; + let chunks = chunker.calculate( + LevelDataRef::Absent, + LevelDataRef::Absent, + num_values, + |c, i| { + c.roll_fixed::<4>(&(i as i32).to_le_bytes()); + }, + ); + + // Should have multiple chunks + assert!(chunks.len() > 1); + + // Verify contiguity + let mut total_levels = 0; + for (i, chunk) in chunks.iter().enumerate() { + assert_eq!(chunk.level_offset, total_levels); + if i < chunks.len() - 1 { + assert!(chunk.num_levels > 0); + } + total_levels += chunk.num_levels; + } + assert_eq!(total_levels, num_values); + } + + #[test] + fn test_deterministic_chunks() { + let options = CdcOptions { + min_chunk_size: 4, + max_chunk_size: 64, + norm_level: 0, + }; + + let roll = |c: &mut ContentDefinedChunker, i: usize| { + c.roll_fixed::<8>(&(i as i64).to_le_bytes()); + }; + + let mut chunker1 = ContentDefinedChunker::new(&make_desc(0, 0), &options).unwrap(); + let chunks1 = chunker1.calculate(LevelDataRef::Absent, LevelDataRef::Absent, 200, roll); + + let mut chunker2 = ContentDefinedChunker::new(&make_desc(0, 0), &options).unwrap(); + let chunks2 = chunker2.calculate(LevelDataRef::Absent, LevelDataRef::Absent, 200, roll); + + assert_eq!(chunks1.len(), chunks2.len()); + for (a, b) in chunks1.iter().zip(chunks2.iter()) { + assert_eq!(a.level_offset, b.level_offset); + assert_eq!(a.num_levels, b.num_levels); + assert_eq!(a.value_offset, b.value_offset); + assert_eq!(a.num_values, b.num_values); + } + } + + #[test] + fn test_nullable_non_nested() { + let options = CdcOptions { + min_chunk_size: 4, + max_chunk_size: 64, + norm_level: 0, + }; + let mut chunker = ContentDefinedChunker::new(&make_desc(1, 0), &options).unwrap(); + + let num_levels = 20; + // def_level=1 means non-null, def_level=0 means null + // Pattern: null at indices 0, 3, 6, 9, 12, 15, 18 → 7 nulls, 13 non-null + let def_levels: Vec = (0..num_levels) + .map(|i| if i % 3 == 0 { 0 } else { 1 }) + .collect(); + let expected_non_null: usize = def_levels.iter().filter(|&&d| d == 1).count(); + + let chunks = chunker.calculate( + LevelDataRef::Materialized(&def_levels), + LevelDataRef::Absent, + num_levels, + |c, i| { + c.roll_fixed::<4>(&(i as i32).to_le_bytes()); + }, + ); + + assert!(!chunks.is_empty()); + let total_levels: usize = chunks.iter().map(|c| c.num_levels).sum(); + let total_values: usize = chunks.iter().map(|c| c.num_values).sum(); + assert_eq!(total_levels, num_levels); + assert_eq!(total_values, expected_non_null); + // With nulls present, total_values < total_levels + assert!(total_values < total_levels); + } +} + +/// Integration tests that exercise CDC through the Arrow writer/reader roundtrip. +/// Ported from the C++ test suite in `chunker_internal_test.cc`. +#[cfg(all(test, feature = "arrow"))] +mod arrow_tests { + use std::borrow::Borrow; + use std::sync::Arc; + + use arrow::util::data_gen::create_random_batch; + use arrow_array::cast::AsArray; + use arrow_array::{Array, ArrayRef, BooleanArray, Int32Array, RecordBatch}; + use arrow_buffer::Buffer; + use arrow_data::ArrayData; + use arrow_schema::{DataType, Field, Fields, Schema}; + + use crate::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use crate::arrow::arrow_writer::ArrowWriter; + use crate::column::writer::LevelDataRef; + use crate::file::properties::{CdcOptions, WriterProperties}; + use crate::file::reader::{FileReader, SerializedFileReader}; + + // --- Constants matching C++ TestCDCSingleRowGroup --- + + const CDC_MIN_CHUNK_SIZE: usize = 4 * 1024; + const CDC_MAX_CHUNK_SIZE: usize = 16 * 1024; + const CDC_PART_SIZE: usize = 128 * 1024; + const CDC_EDIT_SIZE: usize = 128; + const CDC_ROW_GROUP_LENGTH: usize = 1024 * 1024; + + // --- Helpers --- + + /// Deterministic hash function matching the C++ test generator. + fn test_hash(seed: u64, index: u64) -> u64 { + let mut h = (index.wrapping_add(seed)).wrapping_mul(0xc4ceb9fe1a85ec53u64); + h ^= h >> 33; + h = h.wrapping_mul(0xff51afd7ed558ccdu64); + h ^= h >> 33; + h = h.wrapping_mul(0xc4ceb9fe1a85ec53u64); + h ^= h >> 33; + h + } + + /// Generate a deterministic array for any supported data type, matching C++ `GenerateArray`. + fn generate_array(dtype: &DataType, nullable: bool, length: usize, seed: u64) -> ArrayRef { + macro_rules! gen_primitive { + ($array_type:ty, $cast:expr) => {{ + if nullable { + let arr: $array_type = (0..length) + .map(|i| { + let val = test_hash(seed, i as u64); + if val % 10 == 0 { + None + } else { + Some($cast(val)) + } + }) + .collect(); + Arc::new(arr) as ArrayRef + } else { + let arr: $array_type = (0..length) + .map(|i| Some($cast(test_hash(seed, i as u64)))) + .collect(); + Arc::new(arr) as ArrayRef + } + }}; + } + + match dtype { + DataType::Boolean => { + if nullable { + let arr: BooleanArray = (0..length) + .map(|i| { + let val = test_hash(seed, i as u64); + if val % 10 == 0 { + None + } else { + Some(val % 2 == 0) + } + }) + .collect(); + Arc::new(arr) + } else { + let arr: BooleanArray = (0..length) + .map(|i| Some(test_hash(seed, i as u64) % 2 == 0)) + .collect(); + Arc::new(arr) + } + } + DataType::Int32 => gen_primitive!(Int32Array, |v: u64| v as i32), + DataType::Int64 => { + gen_primitive!(arrow_array::Int64Array, |v: u64| v as i64) + } + DataType::Float64 => { + gen_primitive!(arrow_array::Float64Array, |v: u64| (v % 100000) as f64 + / 1000.0) + } + DataType::Utf8 => { + let arr: arrow_array::StringArray = if nullable { + (0..length) + .map(|i| { + let val = test_hash(seed, i as u64); + if val % 10 == 0 { + None + } else { + Some(format!("str_{val}")) + } + }) + .collect() + } else { + (0..length) + .map(|i| Some(format!("str_{}", test_hash(seed, i as u64)))) + .collect() + }; + Arc::new(arr) + } + DataType::Binary => { + let arr: arrow_array::BinaryArray = if nullable { + (0..length) + .map(|i| { + let val = test_hash(seed, i as u64); + if val % 10 == 0 { + None + } else { + Some(format!("bin_{val}").into_bytes()) + } + }) + .collect() + } else { + (0..length) + .map(|i| Some(format!("bin_{}", test_hash(seed, i as u64)).into_bytes())) + .collect() + }; + Arc::new(arr) + } + DataType::FixedSizeBinary(size) => { + let size = *size; + let mut builder = arrow_array::builder::FixedSizeBinaryBuilder::new(size); + for i in 0..length { + let val = test_hash(seed, i as u64); + if nullable && val % 10 == 0 { + builder.append_null(); + } else { + let s = format!("bin_{val}"); + let bytes = s.as_bytes(); + let mut buf = vec![0u8; size as usize]; + let copy_len = bytes.len().min(size as usize); + buf[..copy_len].copy_from_slice(&bytes[..copy_len]); + builder.append_value(&buf).unwrap(); + } + } + Arc::new(builder.finish()) + } + DataType::Date32 => { + gen_primitive!(arrow_array::Date32Array, |v: u64| v as i32) + } + DataType::Timestamp(arrow_schema::TimeUnit::Nanosecond, _) => { + gen_primitive!(arrow_array::TimestampNanosecondArray, |v: u64| v as i64) + } + _ => panic!("Unsupported test data type: {dtype:?}"), + } + } + + /// Generate a RecordBatch with the given schema, matching C++ `GenerateTable`. + fn generate_table(schema: &Arc, length: usize, seed: u64) -> RecordBatch { + let arrays: Vec = schema + .fields() + .iter() + .enumerate() + .map(|(i, field)| { + generate_array( + field.data_type(), + field.is_nullable(), + length, + seed + i as u64 * 10, + ) + }) + .collect(); + RecordBatch::try_new(schema.clone(), arrays).unwrap() + } + + /// Compute the CDC byte width for a data type, matching C++ `bytes_per_record`. + /// Returns 0 for variable-length types. + fn cdc_byte_width(dtype: &DataType) -> usize { + match dtype { + DataType::Boolean => 1, + DataType::Int8 | DataType::UInt8 => 1, + DataType::Int16 | DataType::UInt16 | DataType::Float16 => 2, + DataType::Int32 + | DataType::UInt32 + | DataType::Float32 + | DataType::Date32 + | DataType::Time32(_) => 4, + DataType::Int64 + | DataType::UInt64 + | DataType::Float64 + | DataType::Date64 + | DataType::Time64(_) + | DataType::Timestamp(_, _) + | DataType::Duration(_) => 8, + DataType::Decimal128(_, _) => 16, + DataType::Decimal256(_, _) => 32, + DataType::FixedSizeBinary(n) => *n as usize, + _ => 0, // variable-length + } + } + + /// Compute bytes_per_record for determining part/edit lengths, matching C++. + fn bytes_per_record(dtype: &DataType, nullable: bool) -> usize { + let bw = cdc_byte_width(dtype); + if bw > 0 { + if nullable { bw + 2 } else { bw } + } else { + 16 // variable-length fallback, matching C++ + } + } + + /// Compute the CDC chunk size for an array slice, matching C++ `CalculateCdcSize`. + fn calculate_cdc_size(array: &dyn Array, nullable: bool) -> i64 { + let dtype = array.data_type(); + let bw = cdc_byte_width(dtype); + let result = if bw > 0 { + // Fixed-width: count only non-null values + let valid_count = array.len() - array.null_count(); + (valid_count * bw) as i64 + } else { + // Variable-length: sum of actual byte lengths + match dtype { + DataType::Utf8 => { + let a = array.as_string::(); + (0..a.len()) + .filter(|&i| a.is_valid(i)) + .map(|i| a.value(i).len() as i64) + .sum() + } + DataType::Binary => { + let a = array.as_binary::(); + (0..a.len()) + .filter(|&i| a.is_valid(i)) + .map(|i| a.value(i).len() as i64) + .sum() + } + DataType::LargeBinary => { + let a = array.as_binary::(); + (0..a.len()) + .filter(|&i| a.is_valid(i)) + .map(|i| a.value(i).len() as i64) + .sum() + } + _ => panic!("CDC size calculation not implemented for {dtype:?}"), + } + }; + + if nullable { + // Add 2 bytes per element for definition levels + result + array.len() as i64 * 2 + } else { + result + } + } + + /// Page-level metadata for a single column within a row group. + struct ColumnInfo { + page_lengths: Vec, + has_dictionary_page: bool, + } + + /// Extract per-row-group column info from Parquet data. + fn get_column_info(data: &[u8], column_index: usize) -> Vec { + let reader = SerializedFileReader::new(bytes::Bytes::from(data.to_vec())).unwrap(); + let metadata = reader.metadata(); + let mut result = Vec::new(); + for rg in 0..metadata.num_row_groups() { + let rg_reader = reader.get_row_group(rg).unwrap(); + let col_reader = rg_reader.get_column_page_reader(column_index).unwrap(); + let mut info = ColumnInfo { + page_lengths: Vec::new(), + has_dictionary_page: false, + }; + for page in col_reader { + let page = page.unwrap(); + match page.page_type() { + crate::basic::PageType::DATA_PAGE | crate::basic::PageType::DATA_PAGE_V2 => { + info.page_lengths.push(page.num_values() as i64); + } + crate::basic::PageType::DICTIONARY_PAGE => { + info.has_dictionary_page = true; + } + _ => {} + } + } + result.push(info); + } + result + } + + /// Assert that CDC chunk sizes are within the expected range. + /// Equivalent to C++ `AssertContentDefinedChunkSizes`. + fn assert_cdc_chunk_sizes( + array: &ArrayRef, + info: &ColumnInfo, + nullable: bool, + min_chunk_size: usize, + max_chunk_size: usize, + expect_dictionary_page: bool, + ) { + // Boolean and FixedSizeBinary never produce dictionary pages (matching C++) + let expect_dict = match array.data_type() { + DataType::Boolean | DataType::FixedSizeBinary(_) => false, + _ => expect_dictionary_page, + }; + assert_eq!( + info.has_dictionary_page, + expect_dict, + "dictionary page mismatch for {:?}", + array.data_type() + ); + + let page_lengths = &info.page_lengths; + assert!( + page_lengths.len() > 1, + "CDC should produce multiple pages, got {page_lengths:?}" + ); + + let bw = cdc_byte_width(array.data_type()); + // Only do exact CDC size validation for fixed-width and base binary-like types + if bw > 0 + || matches!( + array.data_type(), + DataType::Utf8 | DataType::Binary | DataType::LargeBinary + ) + { + let mut offset = 0i64; + for (i, &page_len) in page_lengths.iter().enumerate() { + let slice = array.slice(offset as usize, page_len as usize); + let cdc_size = calculate_cdc_size(slice.as_ref(), nullable); + if i < page_lengths.len() - 1 { + assert!( + cdc_size >= min_chunk_size as i64, + "Page {i}: CDC size {cdc_size} < min {min_chunk_size}, pages={page_lengths:?}" + ); + } + assert!( + cdc_size <= max_chunk_size as i64, + "Page {i}: CDC size {cdc_size} > max {max_chunk_size}, pages={page_lengths:?}" + ); + offset += page_len; + } + assert_eq!( + offset, + array.len() as i64, + "page lengths must sum to array length" + ); + } + } + + /// Write batches with CDC options and validate roundtrip. + /// Matches C++ `WriteTableToBuffer`. + fn write_with_cdc_options( + batches: &[&RecordBatch], + min_chunk_size: usize, + max_chunk_size: usize, + max_row_group_rows: Option, + enable_dictionary: bool, + ) -> Vec { + assert!(!batches.is_empty()); + let schema = batches[0].schema(); + let mut builder = WriterProperties::builder() + .set_dictionary_enabled(enable_dictionary) + .set_content_defined_chunking(Some(CdcOptions { + min_chunk_size, + max_chunk_size, + norm_level: 0, + })); + if let Some(max_rows) = max_row_group_rows { + builder = builder.set_max_row_group_row_count(Some(max_rows)); + } + let props = builder.build(); + let mut buf = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), Some(props)).unwrap(); + for batch in batches { + writer.write(batch).unwrap(); + } + writer.close().unwrap(); + + // Roundtrip validation (matching C++ WriteTableToBuffer) + let readback = read_batches(&buf); + let original_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + let readback_rows: usize = readback.iter().map(|b| b.num_rows()).sum(); + assert_eq!(original_rows, readback_rows, "Roundtrip row count mismatch"); + if original_rows > 0 { + let original = concat_batches(batches.iter().copied()); + let roundtrip = concat_batches(&readback); + assert_eq!(original, roundtrip, "Roundtrip validation failed"); + } + + buf + } + + #[test] + fn cdc_all_null_arrow_column_writes_data_pages() { + let array = Arc::new(Int32Array::from(vec![None::; 4096])) as ArrayRef; + let schema = Arc::new(Schema::new(vec![Field::new("f0", DataType::Int32, true)])); + let batch = RecordBatch::try_new(schema, vec![array.clone()]).unwrap(); + + let data = write_with_cdc_options(&[&batch], 64, 256, Some(4096), false); + let info = get_column_info(&data, 0); + + assert_eq!(info.len(), 1); + assert!( + !info[0].page_lengths.is_empty(), + "all-null CDC write should still emit data pages" + ); + assert_eq!( + info[0].page_lengths.iter().sum::(), + array.len() as i64, + "all-null CDC pages should account for every input row" + ); + } + + fn read_batches(data: &[u8]) -> Vec { + let reader = ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::from(data.to_vec())) + .unwrap() + .build() + .unwrap(); + reader.collect::, _>>().unwrap() + } + + fn concat_batches(batches: impl IntoIterator>) -> RecordBatch { + let batches: Vec<_> = batches.into_iter().collect(); + let schema = batches[0].borrow().schema(); + let batches = batches.iter().map(|b| b.borrow()); + arrow_select::concat::concat_batches(&schema, batches).unwrap() + } + + /// LCS-based diff between two sequences of page lengths (ported from C++). + /// Includes the merge-adjacent-diffs post-processing from C++. + fn find_differences(first: &[i64], second: &[i64]) -> Vec<(Vec, Vec)> { + let n = first.len(); + let m = second.len(); + let mut dp = vec![vec![0usize; m + 1]; n + 1]; + for i in 0..n { + for j in 0..m { + if first[i] == second[j] { + dp[i + 1][j + 1] = dp[i][j] + 1; + } else { + dp[i + 1][j + 1] = dp[i + 1][j].max(dp[i][j + 1]); + } + } + } + let mut common = Vec::new(); + let (mut i, mut j) = (n, m); + while i > 0 && j > 0 { + if first[i - 1] == second[j - 1] { + common.push((i - 1, j - 1)); + i -= 1; + j -= 1; + } else if dp[i - 1][j] >= dp[i][j - 1] { + i -= 1; + } else { + j -= 1; + } + } + common.reverse(); + + let mut result = Vec::new(); + let (mut last_i, mut last_j) = (0usize, 0usize); + for (ci, cj) in &common { + if *ci > last_i || *cj > last_j { + result.push((first[last_i..*ci].to_vec(), second[last_j..*cj].to_vec())); + } + last_i = ci + 1; + last_j = cj + 1; + } + if last_i < n || last_j < m { + result.push((first[last_i..].to_vec(), second[last_j..].to_vec())); + } + + // Merge adjacent diffs (matching C++ post-processing) + let mut merged: Vec<(Vec, Vec)> = Vec::new(); + for diff in result { + if let Some(prev) = merged.last_mut() { + if prev.0.is_empty() && diff.1.is_empty() { + prev.0 = diff.0; + continue; + } else if prev.1.is_empty() && diff.0.is_empty() { + prev.1 = diff.1; + continue; + } + } + merged.push(diff); + } + merged + } + + /// Assert exact page length differences between original and modified files. + /// Matches C++ `AssertPageLengthDifferences` (full version). + fn assert_page_length_differences( + original: &ColumnInfo, + modified: &ColumnInfo, + exact_equal_diffs: usize, + exact_larger_diffs: usize, + exact_smaller_diffs: usize, + edit_length: i64, + ) { + let diffs = find_differences(&original.page_lengths, &modified.page_lengths); + let expected = exact_equal_diffs + exact_larger_diffs + exact_smaller_diffs; + + if diffs.len() != expected { + eprintln!("Original: {:?}", original.page_lengths); + eprintln!("Modified: {:?}", modified.page_lengths); + for d in &diffs { + eprintln!(" Diff: {:?} vs {:?}", d.0, d.1); + } + } + assert_eq!( + diffs.len(), + expected, + "Expected {expected} diffs, got {}", + diffs.len() + ); + + let (mut eq, mut larger, mut smaller) = (0usize, 0usize, 0usize); + for (left, right) in &diffs { + let left_sum: i64 = left.iter().sum(); + let right_sum: i64 = right.iter().sum(); + if left_sum == right_sum { + eq += 1; + } else if left_sum < right_sum { + larger += 1; + assert_eq!( + left_sum + edit_length, + right_sum, + "Larger diff mismatch: {left_sum} + {edit_length} != {right_sum}" + ); + } else { + smaller += 1; + assert_eq!( + left_sum, + right_sum + edit_length, + "Smaller diff mismatch: {left_sum} != {right_sum} + {edit_length}" + ); + } + } + + assert_eq!(eq, exact_equal_diffs, "equal diffs count"); + assert_eq!(larger, exact_larger_diffs, "larger diffs count"); + assert_eq!(smaller, exact_smaller_diffs, "smaller diffs count"); + } + + /// Assert page length differences for update cases (simplified version). + /// Matches C++ `AssertPageLengthDifferences` (max_equal_diffs overload). + fn assert_page_length_differences_update( + original: &ColumnInfo, + modified: &ColumnInfo, + max_equal_diffs: usize, + ) { + let diffs = find_differences(&original.page_lengths, &modified.page_lengths); + assert!( + diffs.len() <= max_equal_diffs, + "Expected at most {max_equal_diffs} diffs, got {}", + diffs.len() + ); + for (left, right) in &diffs { + let left_sum: i64 = left.iter().sum(); + let right_sum: i64 = right.iter().sum(); + assert_eq!( + left_sum, right_sum, + "Update diff should not change total row count" + ); + } + } + + // --- FindDifferences tests (ported from C++) --- + + #[test] + fn test_find_differences_basic() { + let diffs = find_differences(&[1, 2, 3, 4, 5], &[1, 7, 8, 4, 5]); + assert_eq!(diffs.len(), 1); + assert_eq!(diffs[0].0, vec![2, 3]); + assert_eq!(diffs[0].1, vec![7, 8]); + } + + #[test] + fn test_find_differences_multiple() { + let diffs = find_differences(&[1, 2, 3, 4, 5, 6, 7], &[1, 8, 9, 4, 10, 6, 11]); + assert_eq!(diffs.len(), 3); + assert_eq!(diffs[0].0, vec![2, 3]); + assert_eq!(diffs[0].1, vec![8, 9]); + assert_eq!(diffs[1].0, vec![5]); + assert_eq!(diffs[1].1, vec![10]); + assert_eq!(diffs[2].0, vec![7]); + assert_eq!(diffs[2].1, vec![11]); + } + + #[test] + fn test_find_differences_different_lengths() { + let diffs = find_differences(&[1, 2, 3], &[1, 2, 3, 4, 5]); + assert_eq!(diffs.len(), 1); + assert!(diffs[0].0.is_empty()); + assert_eq!(diffs[0].1, vec![4, 5]); + } + + #[test] + fn test_find_differences_empty() { + let diffs = find_differences(&[], &[]); + assert!(diffs.is_empty()); + } + + #[test] + fn test_find_differences_changes_at_both_ends() { + let diffs = find_differences(&[1, 2, 3, 4, 5, 6, 7, 8, 9], &[0, 0, 2, 3, 4, 5, 7, 7, 8]); + assert_eq!(diffs.len(), 3); + assert_eq!(diffs[0].0, vec![1]); + assert_eq!(diffs[0].1, vec![0, 0]); + assert_eq!(diffs[1].0, vec![6]); + assert_eq!(diffs[1].1, vec![7]); + assert_eq!(diffs[2].0, vec![9]); + assert!(diffs[2].1.is_empty()); + } + + #[test] + fn test_find_differences_additional() { + let diffs = find_differences( + &[445, 312, 393, 401, 410, 138, 558, 457], + &[445, 312, 393, 393, 410, 138, 558, 457], + ); + assert_eq!(diffs.len(), 1); + assert_eq!(diffs[0].0, vec![401]); + assert_eq!(diffs[0].1, vec![393]); + } + + // --- Parameterized single-row-group tests via macro --- + + macro_rules! cdc_single_rg_tests { + ($mod_name:ident, $dtype:expr, $nullable:expr) => { + mod $mod_name { + use super::*; + + fn config() -> (DataType, bool, usize, usize) { + let dtype: DataType = $dtype; + let nullable: bool = $nullable; + let bpr = bytes_per_record(&dtype, nullable); + let part_length = CDC_PART_SIZE / bpr; + let edit_length = CDC_EDIT_SIZE / bpr; + (dtype, nullable, part_length, edit_length) + } + + fn make_schema(dtype: &DataType, nullable: bool) -> Arc { + Arc::new(Schema::new(vec![Field::new("f0", dtype.clone(), nullable)])) + } + + #[test] + fn delete_once() { + let (dtype, nullable, part_length, edit_length) = config(); + let schema = make_schema(&dtype, nullable); + + let part1 = generate_table(&schema, part_length, 0); + let part2 = generate_table(&schema, edit_length, 1); + let part3 = generate_table(&schema, part_length, part_length as u64); + + let base = concat_batches([&part1, &part2, &part3]); + let modified = concat_batches([&part1, &part3]); + + for enable_dictionary in [false, true] { + let base_data = write_with_cdc_options( + &[&base], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(CDC_ROW_GROUP_LENGTH), + enable_dictionary, + ); + let mod_data = write_with_cdc_options( + &[&modified], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(CDC_ROW_GROUP_LENGTH), + enable_dictionary, + ); + + let base_info = get_column_info(&base_data, 0); + let mod_info = get_column_info(&mod_data, 0); + assert_eq!(base_info.len(), 1); + assert_eq!(mod_info.len(), 1); + + assert_cdc_chunk_sizes( + &base.column(0).clone(), + &base_info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + enable_dictionary, + ); + assert_cdc_chunk_sizes( + &modified.column(0).clone(), + &mod_info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + enable_dictionary, + ); + + assert_page_length_differences( + &base_info[0], + &mod_info[0], + 0, + 0, + 1, + edit_length as i64, + ); + } + } + + #[test] + fn delete_twice() { + let (dtype, nullable, part_length, edit_length) = config(); + let schema = make_schema(&dtype, nullable); + + let part1 = generate_table(&schema, part_length, 0); + let part2 = generate_table(&schema, edit_length, 1); + let part3 = generate_table(&schema, part_length, part_length as u64); + let part4 = generate_table(&schema, edit_length, 2); + let part5 = generate_table(&schema, part_length, 2 * part_length as u64); + + let base = concat_batches([&part1, &part2, &part3, &part4, &part5]); + let modified = concat_batches([&part1, &part3, &part5]); + + for enable_dictionary in [false, true] { + let base_data = write_with_cdc_options( + &[&base], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(CDC_ROW_GROUP_LENGTH), + enable_dictionary, + ); + let mod_data = write_with_cdc_options( + &[&modified], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(CDC_ROW_GROUP_LENGTH), + enable_dictionary, + ); + + let base_info = get_column_info(&base_data, 0); + let mod_info = get_column_info(&mod_data, 0); + assert_eq!(base_info.len(), 1); + assert_eq!(mod_info.len(), 1); + + assert_cdc_chunk_sizes( + &base.column(0).clone(), + &base_info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + enable_dictionary, + ); + assert_cdc_chunk_sizes( + &modified.column(0).clone(), + &mod_info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + enable_dictionary, + ); + + assert_page_length_differences( + &base_info[0], + &mod_info[0], + 0, + 0, + 2, + edit_length as i64, + ); + } + } + + #[test] + fn insert_once() { + let (dtype, nullable, part_length, edit_length) = config(); + let schema = make_schema(&dtype, nullable); + + let part1 = generate_table(&schema, part_length, 0); + let part2 = generate_table(&schema, edit_length, 1); + let part3 = generate_table(&schema, part_length, part_length as u64); + + let base = concat_batches([&part1, &part3]); + let modified = concat_batches([&part1, &part2, &part3]); + + for enable_dictionary in [false, true] { + let base_data = write_with_cdc_options( + &[&base], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(CDC_ROW_GROUP_LENGTH), + enable_dictionary, + ); + let mod_data = write_with_cdc_options( + &[&modified], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(CDC_ROW_GROUP_LENGTH), + enable_dictionary, + ); + + let base_info = get_column_info(&base_data, 0); + let mod_info = get_column_info(&mod_data, 0); + assert_eq!(base_info.len(), 1); + assert_eq!(mod_info.len(), 1); + + assert_cdc_chunk_sizes( + &base.column(0).clone(), + &base_info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + enable_dictionary, + ); + assert_cdc_chunk_sizes( + &modified.column(0).clone(), + &mod_info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + enable_dictionary, + ); + + assert_page_length_differences( + &base_info[0], + &mod_info[0], + 0, + 1, + 0, + edit_length as i64, + ); + } + } + + #[test] + fn insert_twice() { + let (dtype, nullable, part_length, edit_length) = config(); + let schema = make_schema(&dtype, nullable); + + let part1 = generate_table(&schema, part_length, 0); + let part2 = generate_table(&schema, edit_length, 1); + let part3 = generate_table(&schema, part_length, part_length as u64); + let part4 = generate_table(&schema, edit_length, 2); + let part5 = generate_table(&schema, part_length, 2 * part_length as u64); + + let base = concat_batches([&part1, &part3, &part5]); + let modified = concat_batches([&part1, &part2, &part3, &part4, &part5]); + + for enable_dictionary in [false, true] { + let base_data = write_with_cdc_options( + &[&base], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(CDC_ROW_GROUP_LENGTH), + enable_dictionary, + ); + let mod_data = write_with_cdc_options( + &[&modified], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(CDC_ROW_GROUP_LENGTH), + enable_dictionary, + ); + + let base_info = get_column_info(&base_data, 0); + let mod_info = get_column_info(&mod_data, 0); + assert_eq!(base_info.len(), 1); + assert_eq!(mod_info.len(), 1); + + assert_cdc_chunk_sizes( + &base.column(0).clone(), + &base_info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + enable_dictionary, + ); + assert_cdc_chunk_sizes( + &modified.column(0).clone(), + &mod_info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + enable_dictionary, + ); + + assert_page_length_differences( + &base_info[0], + &mod_info[0], + 0, + 2, + 0, + edit_length as i64, + ); + } + } + + #[test] + fn update_once() { + let (dtype, nullable, part_length, edit_length) = config(); + let schema = make_schema(&dtype, nullable); + + let part1 = generate_table(&schema, part_length, 0); + let part2 = generate_table(&schema, edit_length, 1); + let part3 = generate_table(&schema, part_length, part_length as u64); + let part4 = generate_table(&schema, edit_length, 2); + + let base = concat_batches([&part1, &part2, &part3]); + let modified = concat_batches([&part1, &part4, &part3]); + + for enable_dictionary in [false, true] { + let base_data = write_with_cdc_options( + &[&base], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(CDC_ROW_GROUP_LENGTH), + enable_dictionary, + ); + let mod_data = write_with_cdc_options( + &[&modified], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(CDC_ROW_GROUP_LENGTH), + enable_dictionary, + ); + + let base_info = get_column_info(&base_data, 0); + let mod_info = get_column_info(&mod_data, 0); + assert_eq!(base_info.len(), 1); + assert_eq!(mod_info.len(), 1); + + assert_cdc_chunk_sizes( + &base.column(0).clone(), + &base_info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + enable_dictionary, + ); + assert_cdc_chunk_sizes( + &modified.column(0).clone(), + &mod_info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + enable_dictionary, + ); + + assert_page_length_differences_update(&base_info[0], &mod_info[0], 1); + } + } + + #[test] + fn update_twice() { + let (dtype, nullable, part_length, edit_length) = config(); + let schema = make_schema(&dtype, nullable); + + let part1 = generate_table(&schema, part_length, 0); + let part2 = generate_table(&schema, edit_length, 1); + let part3 = generate_table(&schema, part_length, part_length as u64); + let part4 = generate_table(&schema, edit_length, 2); + let part5 = generate_table(&schema, part_length, 2 * part_length as u64); + let part6 = generate_table(&schema, edit_length, 3); + let part7 = generate_table(&schema, edit_length, 4); + + let base = concat_batches([&part1, &part2, &part3, &part4, &part5]); + let modified = concat_batches([&part1, &part6, &part3, &part7, &part5]); + + for enable_dictionary in [false, true] { + let base_data = write_with_cdc_options( + &[&base], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(CDC_ROW_GROUP_LENGTH), + enable_dictionary, + ); + let mod_data = write_with_cdc_options( + &[&modified], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(CDC_ROW_GROUP_LENGTH), + enable_dictionary, + ); + + let base_info = get_column_info(&base_data, 0); + let mod_info = get_column_info(&mod_data, 0); + assert_eq!(base_info.len(), 1); + assert_eq!(mod_info.len(), 1); + + assert_cdc_chunk_sizes( + &base.column(0).clone(), + &base_info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + enable_dictionary, + ); + assert_cdc_chunk_sizes( + &modified.column(0).clone(), + &mod_info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + enable_dictionary, + ); + + assert_page_length_differences_update(&base_info[0], &mod_info[0], 2); + } + } + + #[test] + fn prepend() { + let (dtype, nullable, part_length, edit_length) = config(); + let schema = make_schema(&dtype, nullable); + + let part1 = generate_table(&schema, part_length, 0); + let part2 = generate_table(&schema, edit_length, 1); + let part3 = generate_table(&schema, part_length, part_length as u64); + let part4 = generate_table(&schema, edit_length, 2); + + let base = concat_batches([&part1, &part2, &part3]); + let modified = concat_batches([&part4, &part1, &part2, &part3]); + + for enable_dictionary in [false, true] { + let base_data = write_with_cdc_options( + &[&base], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(CDC_ROW_GROUP_LENGTH), + enable_dictionary, + ); + let mod_data = write_with_cdc_options( + &[&modified], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(CDC_ROW_GROUP_LENGTH), + enable_dictionary, + ); + + let base_info = get_column_info(&base_data, 0); + let mod_info = get_column_info(&mod_data, 0); + assert_eq!(base_info.len(), 1); + assert_eq!(mod_info.len(), 1); + + assert_cdc_chunk_sizes( + &base.column(0).clone(), + &base_info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + enable_dictionary, + ); + assert_cdc_chunk_sizes( + &modified.column(0).clone(), + &mod_info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + enable_dictionary, + ); + + assert!( + mod_info[0].page_lengths.len() >= base_info[0].page_lengths.len(), + "Modified should have same or more pages" + ); + + assert_page_length_differences( + &base_info[0], + &mod_info[0], + 0, + 1, + 0, + edit_length as i64, + ); + } + } + + #[test] + fn append() { + let (dtype, nullable, part_length, edit_length) = config(); + let schema = make_schema(&dtype, nullable); + + let part1 = generate_table(&schema, part_length, 0); + let part2 = generate_table(&schema, edit_length, 1); + let part3 = generate_table(&schema, part_length, part_length as u64); + let part4 = generate_table(&schema, edit_length, 2); + + let base = concat_batches([&part1, &part2, &part3]); + let modified = concat_batches([&part1, &part2, &part3, &part4]); + + for enable_dictionary in [false, true] { + let base_data = write_with_cdc_options( + &[&base], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(CDC_ROW_GROUP_LENGTH), + enable_dictionary, + ); + let mod_data = write_with_cdc_options( + &[&modified], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(CDC_ROW_GROUP_LENGTH), + enable_dictionary, + ); + + let base_info = get_column_info(&base_data, 0); + let mod_info = get_column_info(&mod_data, 0); + assert_eq!(base_info.len(), 1); + assert_eq!(mod_info.len(), 1); + + assert_cdc_chunk_sizes( + &base.column(0).clone(), + &base_info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + enable_dictionary, + ); + assert_cdc_chunk_sizes( + &modified.column(0).clone(), + &mod_info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + enable_dictionary, + ); + + let bp = &base_info[0].page_lengths; + let mp = &mod_info[0].page_lengths; + assert!(mp.len() >= bp.len()); + for i in 0..bp.len() - 1 { + assert_eq!(bp[i], mp[i], "Page {i} should be identical"); + } + assert!(mp[bp.len() - 1] >= bp[bp.len() - 1]); + } + } + + #[test] + fn empty_table() { + let (dtype, nullable, _, _) = config(); + let schema = make_schema(&dtype, nullable); + + let empty = RecordBatch::new_empty(schema); + for enable_dictionary in [false, true] { + let data = write_with_cdc_options( + &[&empty], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(CDC_ROW_GROUP_LENGTH), + enable_dictionary, + ); + let info = get_column_info(&data, 0); + // Empty table: either no row groups or one with no data pages + if !info.is_empty() { + assert!(info[0].page_lengths.is_empty()); + } + } + } + + #[test] + fn array_offsets() { + let (dtype, nullable, part_length, edit_length) = config(); + let schema = make_schema(&dtype, nullable); + + let table = concat_batches([ + &generate_table(&schema, part_length, 0), + &generate_table(&schema, edit_length, 1), + &generate_table(&schema, part_length, part_length as u64), + ]); + + for offset in [0usize, 512, 1024] { + if offset >= table.num_rows() { + continue; + } + let sliced = table.slice(offset, table.num_rows() - offset); + let data = write_with_cdc_options( + &[&sliced], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(CDC_ROW_GROUP_LENGTH), + true, + ); + let info = get_column_info(&data, 0); + assert_eq!(info.len(), 1); + + // Verify CDC actually produced content-defined chunks + assert_cdc_chunk_sizes( + &sliced.column(0).clone(), + &info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + true, + ); + } + } + } + }; + } + + // Instantiate for representative types matching C++ categories + cdc_single_rg_tests!(cdc_bool_non_null, DataType::Boolean, false); + cdc_single_rg_tests!(cdc_i32_non_null, DataType::Int32, false); + cdc_single_rg_tests!(cdc_i64_nullable, DataType::Int64, true); + cdc_single_rg_tests!(cdc_f64_nullable, DataType::Float64, true); + cdc_single_rg_tests!(cdc_utf8_non_null, DataType::Utf8, false); + cdc_single_rg_tests!(cdc_binary_nullable, DataType::Binary, true); + cdc_single_rg_tests!(cdc_fsb16_nullable, DataType::FixedSizeBinary(16), true); + cdc_single_rg_tests!(cdc_date32_non_null, DataType::Date32, false); + cdc_single_rg_tests!( + cdc_timestamp_nullable, + DataType::Timestamp(arrow_schema::TimeUnit::Nanosecond, None), + true + ); + + // --- Multiple row group tests matching C++ TestCDCMultipleRowGroups --- + + mod cdc_multiple_row_groups { + use super::*; + + const PART_LENGTH: usize = 128 * 1024; + const EDIT_LENGTH: usize = 128; + const ROW_GROUP_LENGTH: usize = 64 * 1024; + + fn schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("int32", DataType::Int32, true), + Field::new("float64", DataType::Float64, true), + Field::new("bool", DataType::Boolean, false), + ])) + } + + #[test] + fn insert_once() { + let s = schema(); + let part1 = generate_table(&s, PART_LENGTH, 0); + let part2 = generate_table(&s, PART_LENGTH, 2); + let part3 = generate_table(&s, PART_LENGTH, 4); + let edit1 = generate_table(&s, EDIT_LENGTH, 1); + let edit2 = generate_table(&s, EDIT_LENGTH, 3); + + let base = concat_batches([&part1, &edit1, &part2, &part3]); + let modified = concat_batches([&part1, &edit1, &edit2, &part2, &part3]); + assert_eq!(modified.num_rows(), base.num_rows() + EDIT_LENGTH); + + let base_data = write_with_cdc_options( + &[&base], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(ROW_GROUP_LENGTH), + false, + ); + let mod_data = write_with_cdc_options( + &[&modified], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(ROW_GROUP_LENGTH), + false, + ); + + for col in 0..s.fields().len() { + let base_info = get_column_info(&base_data, col); + let mod_info = get_column_info(&mod_data, col); + + assert_eq!(base_info.len(), 7, "expected 7 row groups for col {col}"); + assert_eq!(mod_info.len(), 7); + + // First two row groups should be identical + assert_eq!(base_info[0].page_lengths, mod_info[0].page_lengths); + assert_eq!(base_info[1].page_lengths, mod_info[1].page_lengths); + + // Middle row groups: 1 larger + 1 smaller diff + for i in 2..mod_info.len() - 1 { + assert_page_length_differences( + &base_info[i], + &mod_info[i], + 0, + 1, + 1, + EDIT_LENGTH as i64, + ); + } + // Last row group: just larger + assert_page_length_differences( + base_info.last().unwrap(), + mod_info.last().unwrap(), + 0, + 1, + 0, + EDIT_LENGTH as i64, + ); + } + } + + #[test] + fn delete_once() { + let s = schema(); + let part1 = generate_table(&s, PART_LENGTH, 0); + let part2 = generate_table(&s, PART_LENGTH, 2); + let part3 = generate_table(&s, PART_LENGTH, 4); + let edit1 = generate_table(&s, EDIT_LENGTH, 1); + let edit2 = generate_table(&s, EDIT_LENGTH, 3); + + let base = concat_batches([&part1, &edit1, &part2, &part3, &edit2]); + let modified = concat_batches([&part1, &part2, &part3, &edit2]); + + let base_data = write_with_cdc_options( + &[&base], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(ROW_GROUP_LENGTH), + false, + ); + let mod_data = write_with_cdc_options( + &[&modified], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(ROW_GROUP_LENGTH), + false, + ); + + for col in 0..s.fields().len() { + let base_info = get_column_info(&base_data, col); + let mod_info = get_column_info(&mod_data, col); + + assert_eq!(base_info.len(), 7); + assert_eq!(mod_info.len(), 7); + + assert_eq!(base_info[0].page_lengths, mod_info[0].page_lengths); + assert_eq!(base_info[1].page_lengths, mod_info[1].page_lengths); + + for i in 2..mod_info.len() - 1 { + assert_page_length_differences( + &base_info[i], + &mod_info[i], + 0, + 1, + 1, + EDIT_LENGTH as i64, + ); + } + assert_page_length_differences( + base_info.last().unwrap(), + mod_info.last().unwrap(), + 0, + 0, + 1, + EDIT_LENGTH as i64, + ); + } + } + + #[test] + fn update_once() { + let s = schema(); + let part1 = generate_table(&s, PART_LENGTH, 0); + let part2 = generate_table(&s, PART_LENGTH, 2); + let part3 = generate_table(&s, PART_LENGTH, 4); + let edit1 = generate_table(&s, EDIT_LENGTH, 1); + let edit2 = generate_table(&s, EDIT_LENGTH, 3); + let edit3 = generate_table(&s, EDIT_LENGTH, 5); + + let base = concat_batches([&part1, &edit1, &part2, &part3, &edit2]); + let modified = concat_batches([&part1, &edit3, &part2, &part3, &edit2]); + + let base_data = write_with_cdc_options( + &[&base], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(ROW_GROUP_LENGTH), + false, + ); + let mod_data = write_with_cdc_options( + &[&modified], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(ROW_GROUP_LENGTH), + false, + ); + + for col in 0..s.fields().len() { + let nullable = s.field(col).is_nullable(); + let base_info = get_column_info(&base_data, col); + let mod_info = get_column_info(&mod_data, col); + + assert_eq!(base_info.len(), 7); + assert_eq!(mod_info.len(), 7); + + // Validate CDC chunk sizes on at least the first row group + assert_cdc_chunk_sizes( + &base.column(col).slice(0, ROW_GROUP_LENGTH), + &base_info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + false, + ); + + assert_eq!(base_info[0].page_lengths, mod_info[0].page_lengths); + assert_eq!(base_info[1].page_lengths, mod_info[1].page_lengths); + + // Row group containing the edit + assert_page_length_differences_update(&base_info[2], &mod_info[2], 1); + + // Remaining row groups should be identical + for i in 3..mod_info.len() { + assert_eq!(base_info[i].page_lengths, mod_info[i].page_lengths); + } + } + } + + #[test] + fn append() { + let s = schema(); + let part1 = generate_table(&s, PART_LENGTH, 0); + let part2 = generate_table(&s, PART_LENGTH, 2); + let part3 = generate_table(&s, PART_LENGTH, 4); + let edit1 = generate_table(&s, EDIT_LENGTH, 1); + let edit2 = generate_table(&s, EDIT_LENGTH, 3); + + let base = concat_batches([&part1, &edit1, &part2, &part3]); + let modified = concat_batches([&part1, &edit1, &part2, &part3, &edit2]); + + let base_data = write_with_cdc_options( + &[&base], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(ROW_GROUP_LENGTH), + false, + ); + let mod_data = write_with_cdc_options( + &[&modified], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + Some(ROW_GROUP_LENGTH), + false, + ); + + for col in 0..s.fields().len() { + let nullable = s.field(col).is_nullable(); + let base_info = get_column_info(&base_data, col); + let mod_info = get_column_info(&mod_data, col); + + assert_eq!(base_info.len(), 7); + assert_eq!(mod_info.len(), 7); + + // Validate CDC chunk sizes on the first row group + assert_cdc_chunk_sizes( + &base.column(col).slice(0, ROW_GROUP_LENGTH), + &base_info[0], + nullable, + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + false, + ); + + // All row groups except last should be identical + for i in 0..base_info.len() - 1 { + assert_eq!(base_info[i].page_lengths, mod_info[i].page_lengths); + } + + // Last row group: pages should be identical except last + let bp = &base_info.last().unwrap().page_lengths; + let mp = &mod_info.last().unwrap().page_lengths; + assert!(mp.len() >= bp.len()); + for i in 0..bp.len() - 1 { + assert_eq!(bp[i], mp[i]); + } + } + } + } + + // --- Direct chunker test (kept from original) --- + + #[test] + fn test_cdc_array_offsets_direct() { + use crate::basic::Type as PhysicalType; + use crate::schema::types::{ColumnDescriptor, ColumnPath, Type}; + + let options = CdcOptions { + min_chunk_size: CDC_MIN_CHUNK_SIZE, + max_chunk_size: CDC_MAX_CHUNK_SIZE, + norm_level: 0, + }; + let desc = { + let tp = Type::primitive_type_builder("col", PhysicalType::INT32) + .build() + .unwrap(); + ColumnDescriptor::new(Arc::new(tp), 0, 0, ColumnPath::new(vec![])) + }; + + let bpr = bytes_per_record(&DataType::Int32, false); + let n = CDC_PART_SIZE / bpr; + let offset = 10usize; + + let array: Int32Array = (0..n).map(|i| test_hash(0, i as u64) as i32).collect(); + let mut chunker = super::ContentDefinedChunker::new(&desc, &options).unwrap(); + let chunks = chunker + .get_arrow_chunks(LevelDataRef::Absent, LevelDataRef::Absent, &array) + .unwrap(); + + let sliced = array.slice(offset, n - offset); + let mut chunker2 = super::ContentDefinedChunker::new(&desc, &options).unwrap(); + let chunks2 = chunker2 + .get_arrow_chunks(LevelDataRef::Absent, LevelDataRef::Absent, &sliced) + .unwrap(); + + let values: Vec = chunks.iter().map(|c| c.num_values).collect(); + let values2: Vec = chunks2.iter().map(|c| c.num_values).collect(); + + assert!(values.len() > 1, "expected multiple chunks, got {values:?}"); + assert_eq!(values.len(), values2.len(), "chunk count must match"); + + assert_eq!( + values[0] - values2[0], + offset, + "offsetted first chunk should be {offset} values shorter" + ); + assert_eq!( + &values[1..], + &values2[1..], + "all chunks after the first must be identical" + ); + } + + /// Regression test for + /// + /// Writing nested list data with CDC enabled panicked with an out-of-bounds + /// slice access when null list entries had non-zero child ranges. + #[test] + fn test_cdc_list_roundtrip() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "_1", + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))), + true, + ), + Field::new( + "_2", + DataType::List(Arc::new(Field::new_list_field(DataType::Boolean, true))), + true, + ), + Field::new( + "_3", + DataType::LargeList(Arc::new(Field::new_list_field(DataType::Utf8, true))), + true, + ), + ])); + let batch = create_random_batch(schema, 10_000, 0.25, 0.75).unwrap(); + write_with_cdc_options( + &[&batch], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + None, + true, + ); + } + + /// Test CDC with deeply nested types: List>, List>> + #[test] + fn test_cdc_deeply_nested_roundtrip() { + let inner_field = Field::new_list_field(DataType::Int32, true); + let inner_type = DataType::List(Arc::new(inner_field)); + let outer_field = Field::new_list_field(inner_type.clone(), true); + let list_list_type = DataType::List(Arc::new(outer_field)); + + let struct_inner_field = Field::new_list_field(DataType::Int32, true); + let struct_inner_type = DataType::List(Arc::new(struct_inner_field)); + let struct_fields = Fields::from(vec![Field::new("a", struct_inner_type, true)]); + let struct_type = DataType::Struct(struct_fields); + let struct_list_field = Field::new_list_field(struct_type, true); + let list_struct_type = DataType::List(Arc::new(struct_list_field)); + + let schema = Arc::new(Schema::new(vec![ + Field::new("list_list", list_list_type, true), + Field::new("list_struct_list", list_struct_type, true), + ])); + let batch = create_random_batch(schema, 10_000, 0.25, 0.75).unwrap(); + write_with_cdc_options( + &[&batch], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + None, + true, + ); + } + + /// Test CDC with list arrays that have non-empty null segments. + /// + /// Per the Arrow columnar format spec: "a null value may correspond to a + /// non-empty segment in the child array". This test constructs such arrays + /// manually and verifies the CDC writer handles them correctly. + #[test] + fn test_cdc_list_non_empty_null_segments() { + // Build List where null entries own non-zero child ranges: + // row 0: [1, 2] offsets[0..2] valid + // row 1: null offsets[2..5] null, but owns 3 child values + // row 2: [6, 7] offsets[5..7] valid + // row 3: null offsets[7..9] null, but owns 2 child values + // row 4: [10] offsets[9..10] valid + let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); + let offsets = Buffer::from_iter([0_i32, 2, 5, 7, 9, 10]); + let null_bitmap = Buffer::from([0b00010101]); // rows 0, 2, 4 valid + + let list_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))); + let list_data = unsafe { + ArrayData::new_unchecked( + list_type.clone(), + 5, + None, + Some(null_bitmap), + 0, + vec![offsets], + vec![values.to_data()], + ) + }; + let list_array = arrow_array::make_array(list_data); + + let schema = Arc::new(Schema::new(vec![Field::new("col", list_type, true)])); + let batch = RecordBatch::try_new(schema, vec![list_array]).unwrap(); + + let buf = write_with_cdc_options( + &[&batch], + CDC_MIN_CHUNK_SIZE, + CDC_MAX_CHUNK_SIZE, + None, + true, + ); + let read = concat_batches(read_batches(&buf)); + let read_list = read.column(0).as_list::(); + assert_eq!(read_list.len(), 5); + assert!(read_list.is_valid(0)); + assert!(read_list.is_null(1)); + assert!(read_list.is_valid(2)); + assert!(read_list.is_null(3)); + assert!(read_list.is_valid(4)); + + let get_vals = |i: usize| -> Vec { + read_list + .value(i) + .as_primitive::() + .values() + .iter() + .copied() + .collect() + }; + assert_eq!(get_vals(0), vec![1, 2]); + assert_eq!(get_vals(2), vec![6, 7]); + assert_eq!(get_vals(4), vec![10]); + } +} diff --git a/vendor/parquet/src/column/chunker/cdc_codegen.py b/vendor/parquet/src/column/chunker/cdc_codegen.py new file mode 100644 index 0000000..3675c92 --- /dev/null +++ b/vendor/parquet/src/column/chunker/cdc_codegen.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python + +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +Produce the given number gearhash tables for rolling hash calculations. + +Each table consists of 256 64-bit integer values and by default 8 tables are +produced. The tables are written to a Rust source file. + +The generated numbers are deterministic "random" numbers created by MD5 hashing +a fixed seed and the table index. This ensures that the tables are the same +across different runs and platforms. The function of generating the numbers is +less important as long as they have sufficiently uniform distribution. + +Reference implementations: +- https://github.com/Borelset/destor/blob/master/src/chunking/fascdc_chunking.c +- https://github.com/nlfiedler/fastcdc-rs/blob/master/examples/table64.rs + +Usage: + python cdc_codegen.py [ntables] + + ntables: Number of gearhash tables to generate (default 8). + + The generated source file is written to ./cdc_generated.rs +""" + +import hashlib +import pathlib +import sys +from io import StringIO + + +template = """\ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// This table should be identical with +// https://github.com/apache/arrow/blob/main/cpp/src/parquet/chunker_internal_generated.h +// Ensure that both tables remain in sync after any changes. + +#[rustfmt::skip] +pub(crate) const NUM_GEARHASH_TABLES: usize = {ntables}; + +#[rustfmt::skip] +pub(crate) const GEARHASH_TABLE: [[u64; 256]; NUM_GEARHASH_TABLES] = [ +{content}]; +""" + + +def generate_hash(n: int, seed: int): + """Produce predictable hash values for a given seed and n using MD5. + + The value can be arbitrary as long as it is deterministic and has a uniform + distribution. The MD5 hash is used to produce a 16 character hexadecimal + string which is then converted to a 64-bit integer. + """ + value = bytes([seed] * 64 + [n] * 64) + hasher = hashlib.md5(value) + return hasher.hexdigest()[:16] + + +def generate_hashtable(seed: int, length=256): + """Generate and render a single gearhash table.""" + table = [generate_hash(n, seed=seed) for n in range(length)] + + out = StringIO() + out.write(f" // seed = {seed}\n") + out.write(" [\n") + for i in range(0, length, 4): + values = [f"0x{value}" for value in table[i : i + 4]] + out.write(f" {', '.join(values)},\n") + out.write(" ]") + + return out.getvalue() + + +def generate_source(ntables=8, relative_path="cdc_generated.rs"): + """Generate a Rust source file with multiple gearhash tables.""" + path = pathlib.Path(__file__).parent / relative_path + tables = [generate_hashtable(seed) for seed in range(ntables)] + content = ",\n".join(tables) + text = template.format(ntables=ntables, content=content) + path.write_text(text) + + +if __name__ == "__main__": + ntables = int(sys.argv[1]) if len(sys.argv) > 1 else 8 + generate_source(ntables) diff --git a/vendor/parquet/src/column/chunker/cdc_generated.rs b/vendor/parquet/src/column/chunker/cdc_generated.rs new file mode 100644 index 0000000..4222e36 --- /dev/null +++ b/vendor/parquet/src/column/chunker/cdc_generated.rs @@ -0,0 +1,558 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#[rustfmt::skip] +pub(crate) const NUM_GEARHASH_TABLES: usize = 8; + +#[rustfmt::skip] +pub(crate) const GEARHASH_TABLE: [[u64; 256]; NUM_GEARHASH_TABLES] = [ + // seed = 0 + [ + 0xf09f35a563783945, 0x0dcc5b3bc5ae410a, 0x63f1ea8d22554270, 0xfbe5ee7bd05a7b61, + 0x3f692ed5e9934aba, 0xaab3755952250eb8, 0xdefb168dc2888fa5, 0x501b36f7c77a7d47, + 0xd2fff45d1989642d, 0x80217c1c600e30a6, 0xb9469ee2e43df7ac, 0x3654b76a61999706, + 0x6ea73dfe5de0c6b6, 0xdfd662e1937a589d, 0x0dbe0cc74b188a68, 0xde45f4e6d73ffc6f, + 0xcdf7a7759e70d87e, 0x5d6a951b8d38c310, 0xdc9423c3813fcf2c, 0x25dc2976e167ffce, + 0xc2555baa1d031c84, 0x115bc3f2230a3ab6, 0xd4b10260f350bede, 0xdfd3501ab447d723, + 0x022e79217edaf167, 0x1635e2255c5a7526, 0xa0a750350cc77102, 0xc027133e05d39f56, + 0xd949459779cf0387, 0xb92f1464f5c688c2, 0xd9ac5f3e8b42f2f3, 0xdf02bb6f5ecaac21, + 0x8156f988fac7bfa4, 0xe4580f97bede2ec8, 0x44fe7d17a76fca32, 0x885f59bd54c2014c, + 0x435e63ec655ffae9, 0x5ebc51930967b1f1, 0x5428c2084ac29e47, 0x9465938fec30e36b, + 0xc7cb3de4977772cd, 0x15692d7c201e8c3a, 0x505ee65cdc4b17f4, 0x7d9839a0a7aead6b, + 0xeef5f5b6a0105291, 0x76c2fb232ce7f5bf, 0x5c13893c1c3ff3a9, 0x65b6b547d4442f98, + 0xb8ad7487c8c96fce, 0x906bcf51c99974f8, 0x2f56e48bb943a48c, 0xbc9ab109f82d3a44, + 0xcd5160cdc8c7e735, 0xbe9acb9df3427732, 0x386b91d477d7fade, 0x36be463621dd5af2, + 0xcbe6a2faffd627a8, 0x9c8fd528463a2f5a, 0xb9b88c6bb802b184, 0xb414b4e665c597c7, + 0xbedb142568209556, 0x5360d81c25429dce, 0x63a69a960a952f37, 0xc900d63899e1b503, + 0x1abc63a8b37c7728, 0xa8b3a8b6409080eb, 0x495e391f662959f6, 0xdf1e136f3e12229b, + 0x33d5fc526b0dd38d, 0x321221ae2abfac63, 0x7fde18351fda7395, 0xed79fe5c3a6aa4c3, + 0x2dd6965a4867d8d4, 0x54813ca20fe8799b, 0x5d59ea6456465c39, 0x0de0c294d1936b81, + 0x4aaf0755002c588c, 0x3530a1857ad04c6d, 0xb8a64f4ce184442b, 0xe0def10bceedfa17, + 0x46e38d0a443757ec, 0x9795a1c645ee16d7, 0x7e531def245eac8a, 0x683b25c43a0716cf, + 0x884583d372da219d, 0x5b06b62c910416e5, 0x54b6902fbebd3dbe, 0x931198d40a761a75, + 0xead7d8e830013590, 0x80b4d5dc99bfaced, 0xf98272c8108a1ad2, 0x1adce054289a0ec6, + 0x7d53a1143c56b465, 0x497fbe4f00c92b52, 0x525e4cc2e81ebd69, 0xc94478e0d5508ff6, + 0xb8a5da83c196d07c, 0x7667a921b65b0603, 0xf236fabbdefe6cd1, 0x53da978d19a92b98, + 0xc604f6e97087124d, 0x2cbd27221924b094, 0x65cd1102c985b1d2, 0x08c0755dc1a97eb4, + 0x5e0419e921c0fef1, 0x282d2c1196f84a29, 0xe21117fcfc5793f7, 0xcf4e985dc38e6c2e, + 0xd521f4f264d55616, 0xde69b04c485f2a10, 0x59410e245305178a, 0xceab1d477c943601, + 0xa9805732d71ee5e9, 0x054cd443896974f6, 0xf2b517717a423a3e, 0x09517937fa9fac95, + 0x4938233e9ca871e3, 0x9132cbaf56f83ec0, 0x4703421ed1dd027d, 0xfd9933f4e6f1ec4e, + 0xf237c7fded2274a8, 0xdf4616efe68cd7b4, 0x5e46de0f39f0a380, 0x3d41e0c6d8e095b0, + 0xc5272f8a5bb2df09, 0x68aa78e8301fb964, 0xbf5b5b52c8e32ae0, 0xbf28ed3df74bdcf7, + 0xd6198f64c833815a, 0x8cd99d2974267544, 0xd90560ea4465ff2c, 0x571d65ad7ad59261, + 0x309453518baa367a, 0xa60538377bc79fb2, 0xace515da1ab4183c, 0xf56d3c8d891d1c5b, + 0x5b0d8370b59def49, 0x775866ce7c83c762, 0x3d76085695c8e18a, 0xba064d1a9af1b114, + 0xc84ef7cd7b98b521, 0x90b9231681c2bc37, 0x37e2b13e6f585b6b, 0x1d0a34e55e0f369f, + 0x86bb8019cf41447c, 0x4b95c6ef55b3f71f, 0x3b6ed1660732b310, 0x617eee603d137f21, + 0xf4f6278b464f3bbc, 0xdfb763b720da205a, 0x353478899b871cb7, 0xe45fbbff574cc41e, + 0x1a94b60847907d72, 0xb10eef051eff67a5, 0xf0e012ec6a284d40, 0xcc1cd1a11b926d7c, + 0xcf9d9c5453e19cad, 0x270febcc0fc0e86b, 0xd6567568778b781e, 0x7323b98965eeb46b, + 0xccecd374567086ff, 0xef7b44bfc497a704, 0xebc479c051a9f0a5, 0xc9b7410e3e00a235, + 0x1d084f7ecdf83dab, 0xc8a9a97e33ba8ba3, 0x8c75318f5b2350d6, 0xaa3cd5d0c684bdda, + 0xa81125fe0901bedf, 0xf7bcd76020edfc93, 0x834ee4c12e75874f, 0xb2bb8a7beb44fa14, + 0x32cd26f50a4f4e4d, 0x0fc5817ca55d959a, 0xd6e4ae2e3ae10718, 0x074abdcceb8d6e38, + 0xc0cc5f4f9b3a9c43, 0x1115d364363595b2, 0x69861db2eb19f2e8, 0x59b8d804cf92bc67, + 0x9bac9785e5e4b863, 0x7fa0e17a41869561, 0x10d3c9633f0c709c, 0x534a03deee6bc44a, + 0x73b1f7201257f581, 0x46fd6a11e2e0706b, 0x494abb554946e67a, 0xb5d6da317864dc8e, + 0x402ded9238f39687, 0xd8fa37d2cbd6d290, 0xcc818293fcb06791, 0x6482ab344806cd4d, + 0x0956e6ee9d8eb60b, 0x01fee622d8465ac8, 0xae7ece370cbd9c35, 0x7ff09e937a177279, + 0xa2c29ee7a33ca5f1, 0x990e8dbee083923b, 0x4a819b72f610863a, 0xddecfad79d3f08be, + 0x627372480fac20a7, 0x802154d6eca2db4c, 0x8fcf02e42f805e55, 0x040a911ff8cea977, + 0xbb544485bc64d0d4, 0xaddde1aeb406d0fb, 0xf6b35fae23dce66f, 0xc07a9fb3645d2f9b, + 0xccd113907e9c0fed, 0xd17af369984fd213, 0x9223823c59a083e7, 0xe19d475606b81013, + 0xe181ac116a90e57a, 0x71f7b6258c6def4c, 0x2246f34b45964f7c, 0xd74aedaea2d31751, + 0xb1add86e5dd305d1, 0xeb9ba881f16d6471, 0xef7600e036f5c6ff, 0x1d50bc9735b8fb85, + 0xe63942bd1f3e2969, 0x9241ba9f8b3f4e72, 0xee8bb2bca07d35b6, 0x55cd55dab522654e, + 0x94d0cfa7c1a6845d, 0x02f9845d559884c3, 0x8ce70ea21063b560, 0xd70998028ef08b74, + 0xdfdb5bbee310876b, 0x4e21b2e348256d16, 0xde007a981c13debc, 0xe51950cbbddabfdd, + 0xd223301dbe9957c1, 0x084b8634cc2cce4b, 0x90e551378aa9d70c, 0x833b533ac633e448, + 0x7891e232882da57f, 0xa1bf26f0163ce2b3, 0xf33a0171eb9c68d5, 0x2e7de18ca69b3fa2, + 0x666fd6f175619199, 0x1239d37edb5feb9f, 0xfa9fc9382e61ff5c, 0x3ca4ad427e3c126f, + 0x37c6dd4c2c31ae6e, 0x1f1bacb619d427b2, 0x7dd09f5d10759afe, 0xc8d941432327d733, + 0x2b389ba25e1d43a7, 0xa4e3030c3740ff21, 0xcc56dae13fd37463, 0x2481457c175b560f, + 0x9deb35bde77c5c41, 0x847aa6ea5549a0c3, 0xcde01bb48b6e7f02, 0x15a28844e64cb211, + ], + // seed = 1 + [ + 0xecfcba92fe5691a3, 0x71377799fea34699, 0xb284c9096fa614e5, 0x54534170f40de6c8, + 0xbbd804d45884fba3, 0x44929a896388c8a1, 0x79b712508e0fa3b1, 0xeb53ab280af31054, + 0x351ea23a6319da7a, 0x2fbe55d9819d85a2, 0x34f4b6568dcd28b1, 0x8c94ea5e5d82967a, + 0x09068d333a46d3c5, 0x762ad4f64cb73381, 0xd5c6db5ef0e22640, 0x36d8ab5a36175680, + 0xd41fe333cdc3525a, 0xa1f51dbdf20ce781, 0x1410a95e786c8be6, 0x96b7499a670c2b41, + 0x3912e1037835d893, 0x272c5bd83e1e9115, 0x2ea7f91cad82a0d6, 0xcd10e85662ce9931, + 0xedad49be8d5e8b74, 0x7ccd8fe0f37d12bc, 0xfac0482005eed593, 0x4513991681f6c8b0, + 0x2804d612eb0ad37d, 0x7cca9e8412b81d34, 0x85ffd6707192b7b8, 0xea0560aeea954411, + 0x0122d28226102bba, 0xf51c47cdbd22fdd1, 0x3707d851183ff17c, 0xaef5a1465f3e902d, + 0xbcb38c2d8736a04f, 0x4025317e864bef15, 0x8d3f66d86e1ea58f, 0xc16759a3d97ed79a, + 0x1c62abdc0659f2f5, 0x23b3eb4e699bd28f, 0x5083c4fceed3ccaf, 0xa65bf34562cc989c, + 0xaa5865932fd79064, 0xf24d08d268c24593, 0x7fbd00a215196999, 0x7812cd366d752964, + 0x62e8dcb27ef3d945, 0xf08b7984e1b946dc, 0x547d23ad9a5c1dcf, 0x496b1fb249b27fb7, + 0xcd692e1db5f3b3ba, 0x41931e39f1e1bc61, 0x286c6a7d7edae82b, 0x17ef6638b6c4ca6e, + 0x609beb5a2576a934, 0xcc5e16fe4a69b83c, 0xbbd14d08b078fc24, 0x2a617680f481cb94, + 0x81dbbd5f86e6d039, 0xeb8205e1fc8ecc3c, 0xe5e3bb576faa8042, 0x5d6f1eb9d9df01b5, + 0x9a47b8739c10fb44, 0x398a7caad7ea7696, 0x9c0fc1d7c46adde6, 0x67cd6de0a51978a6, + 0x68ccc4b77a21cca4, 0x1e067066b82f415c, 0xf7ddade6535e1819, 0xf2185c884291751b, + 0xc322b7381fcbe34f, 0x242f593e88290b9b, 0x8e11ccc0ea5e84a3, 0x40e3a2e3346db8a2, + 0xf18bfc3ad2931a2c, 0x2468397394b00144, 0xeae199cce14e6817, 0x05b462686c75a1ae, + 0xda096cb859c51673, 0xd87aeb967a906bef, 0xaabc74493cb02fe6, 0x74d48fc2e7da143e, + 0x6ec1c8fed3f2c1fd, 0xe01e0704b463f18e, 0xc3d88a4d3a8056e4, 0xd01ae0ffab6c8f3f, + 0x881ba052620ae7c7, 0xcea033aef0a823a5, 0x8d2cad91d83df1e3, 0x18746d205e66dbe9, + 0x3061f8e58d046650, 0xd819c59f0ce2cf8b, 0x144e89e93635e870, 0x3415e88279b21651, + 0xd6f7ab944b86c3fa, 0x45f1dd15d0f67bdc, 0xbf0d97c7f4fa24f4, 0x34a7de520a57fcd2, + 0x4ba86fda03e9e2bc, 0xa7995265a025b552, 0x698f6819d5f51cf7, 0xd07dbe9d8a156981, + 0x2683945373857fc1, 0x116f8a84f96167de, 0x8bc832bd85595ebf, 0xb206519d74fdfafa, + 0xde9519b2e9b5cc5f, 0x16fdd6f2da1d8163, 0x7ba32bd48ef56f11, 0x6f4e4d7ee8b29717, + 0xd31576dde7468aad, 0x023bb08848676045, 0xf6dcc083178160b7, 0x42035f426250e683, + 0x343732993cfed89f, 0x0640a870a22d3d58, 0x65cff80b53b4ae6a, 0x27996fa17ab05215, + 0xfd5db01401b21a04, 0x894508784bc1673c, 0x5bfcf43a2380e27d, 0x4cd6dcc2715583b7, + 0xa43b3763e7d4c902, 0x6da83e12ef0c1257, 0xfe80a602b0335aff, 0x293a7d8f4ff344de, + 0xb4ae7c2b8956bf5a, 0x6b45432d38254b4d, 0xd086acbdf15d9455, 0xa4d19e43f41ea87b, + 0xf01f13ba4bb87fbf, 0xca582cf301a299ff, 0x0ddad3d45298fa7d, 0x0646a130459c3999, + 0xc08e3af3747e2cee, 0xfc7db8aa9ed67295, 0x783b329e7bd79d5f, 0x732dbc607957af7b, + 0x8e446ac19fb26555, 0xff1dfa4d61dc89a5, 0xb6fbc46bd8d011d8, 0x185147ec5779f0d7, + 0x6eb2cf6149a5380f, 0xb0e773df803a1eae, 0xc07706c5519bfce5, 0xc35abcf54fa95f14, + 0x40a01d99a38608ea, 0x776dcd6f603c277f, 0x6ae12389b1d6d0bb, 0x8bd981448df92bb9, + 0x426a6a7ca21a2c16, 0x87efd5b71c1bad26, 0x71fb7fc4cd41de48, 0xdd9033c45619d463, + 0x40eaab322654cef7, 0xe077fffed6f3e3a2, 0x375a4dbef9384447, 0x2066b009d2c4a100, + 0xeca4a5794a068447, 0x2128f64bddf341a1, 0x738b4bb1be90bd61, 0x433772cf3813d52e, + 0x9540c88add8e4474, 0x0b6d5decd21d3519, 0x654ead966745642d, 0xe1bfb03c3b4bdb4c, + 0x0b977a9937515b1f, 0x0a4587509ef63870, 0xe89f0de1d9cfd44a, 0x23a91390272e7f68, + 0xd92defbc9096b8d8, 0x004db87174612539, 0xc88ecaabdd1a71f1, 0x050de38393073346, + 0x8af1426d7964e038, 0xf352c4fef8ad5c87, 0x6f26bc7408e26548, 0x0d41543fd9bf3084, + 0xfc4e07553a840fc6, 0x5ef117de86a555a9, 0x1f11c42dffb5ae1b, 0x4147648f07490fa5, + 0x09b35fd7671b21aa, 0x1453b14f7ccca481, 0x944f6fcce4c9b2ba, 0x5b08dd2e3583dc06, + 0xe0220df78dc9c22d, 0x1c200b9506cbf666, 0x8a0b7465eadb523b, 0xfbcb43a91a1e2d80, + 0xe697f44be3c36a58, 0x2f8a8e48fb7e350d, 0x7baba71b8920d55f, 0x10edc0216105bc96, + 0x52db07c79d7a7a63, 0x1916e8cef9452ac3, 0x5cbbbf21f867b6cc, 0xadd583365a690a4b, + 0x4e4ca2c8bffc2fdb, 0xf5fe3416d2eebcfe, 0x839af8b85e452476, 0x8496c0c54ad44e16, + 0x6c46f1ecad4482bf, 0xb794cad76ae18715, 0x67b762eec7c62985, 0x52dc9e68df5b3a53, + 0x0cc7e444b422a5f9, 0xadbfe90841c112b0, 0xfe37b136f0ca5c34, 0xcfe9e47948a8d73e, + 0xee90572b86a30d91, 0x549e72d8262830aa, 0x3361564b469f32c6, 0x1e6eba9e0d2648e2, + 0x5f8e2b2ac5fcb4eb, 0xe4224fa5f71f7cc6, 0x7357a9230c76757b, 0xcad70f74aaf6b702, + 0xeef28ced23894cc2, 0x753fdd3352aefd68, 0x1fed6ba90bbeb9d2, 0x05316f4ab4034b4b, + 0x3396df022b9f63d6, 0x82d7125a7cfd0935, 0x3519a71caf1f87f0, 0xd1dfb7a5cc3974be, + 0xbfae40ecbdbbcc2a, 0x152c11778e08dd54, 0x4a96566a6c848554, 0x3a84d621c340cdd7, + 0xfd47aa1887e2fb03, 0xa63cae94b2f1d099, 0xed61783f3e5b75e0, 0xefd44864106019be, + 0x145ff78b80b081aa, 0x34670e5fcea9230e, 0x876ef976328db371, 0x4221f3a5269942a6, + 0x95315cbd85c648f4, 0x3ca344dc7c3b1600, 0x38421ea39ff28780, 0x31dbeee967c0435c, + 0x27437c3e268402e7, 0xdd0cf8343312a654, 0x965ab9dad1d8aa29, 0xf871706dd3e23509, + 0xce23d06c7a25e699, 0x1b37d59382b27589, 0x3407f004723d6324, 0x56efb69cdb5deaa1, + 0xf46cdd2b9fd604e0, 0xcad3ca79fdac69bd, 0x7252802a574e63cb, 0xc281fb8acc6ec1d3, + ], + // seed = 2 + [ + 0xdd16cb672ba6979c, 0x3954eaa9ec41ae41, 0x52cb802771d2966d, 0xf57ed8eb0d0294f2, + 0x768be23c71da2219, 0x6131e22d95a84ad3, 0xd849e4e49bb15842, 0x18e8e5c4978cf00d, + 0x3af5e5867ce1f9bd, 0x06c75a9fffe83d63, 0xe8de75a00b58a065, 0x0a773251bc0d755a, + 0x629dc21e54548329, 0x2a168f5e5a883e70, 0x33547375f0996c86, 0xdfcb4c7680451322, + 0x55c1ecaaaa57e397, 0x4546c346c24f5a31, 0x6f8f0401dfabc86c, 0x7760d2d36ee340b4, + 0xf6448e48bdeb229d, 0xba70e1633b4dba65, 0x069cda561e273054, 0xa010b6a84aebf340, + 0x5c23b8229eee34b6, 0xea63c926d90153af, 0x7d7de27b3e43ec1b, 0xea119541eddc3491, + 0xf1259daeddfc724c, 0x2873ca9a67730647, 0xa1e7710dade32607, 0x758de030b61d43fd, + 0xd2c9bcbfa475edb4, 0x18ade47bb8a0aa29, 0xf7a74af0ff1aea88, 0x6f8873274a987162, + 0x6963e8d876f4d282, 0xd435d4fe448c6c5b, 0x93ec80ba404cafff, 0xcf90d24c509e41e7, + 0x5f0fc8a62923e36e, 0x9224878fe458f3a4, 0xd9a039edf1945bcd, 0x0877d1892c288441, + 0x75205491f4b4740b, 0x30f9d2d523a9085b, 0x4b7f4029fa097c99, 0x170bb013745709d4, + 0x7087af537f11ef2e, 0x28c62b88e08fc464, 0x84bbcb3e0bb56271, 0x485a4b099165c681, + 0x357c63357caa9292, 0x819eb7d1aee2d27e, 0xdaa759eb9c0f8c9d, 0x42cdc36729cc3db5, + 0x9489aa852eddbb06, 0x8161e4f85a84e6d4, 0xa964863fdad3eb29, 0xcc095ddbce1a6702, + 0x3ecfadbb8dc2ce58, 0x971316509b95a231, 0xc8f484d1dbc38427, 0xae9c510c463574c0, + 0xdf2b31179600c21a, 0x440de87bada4dfa3, 0xbd8d30f3f6fb7522, 0x84e6d7f678a0e2d0, + 0x0ec4d74323e15975, 0xf6947610dad6d9ab, 0x73a55a95d73fe3a5, 0x3e5f623024d37eda, + 0x8d99a728d95d9344, 0x8b82a7956c4acdc4, 0x7faeaea4385b27f6, 0x540625ff4aa2ff21, + 0x4aa43b3ebd92ce2b, 0x899646a6df2da807, 0x49225115780942d7, 0xe16606636af89525, + 0xb980bcf893888e33, 0xf9ed57695291b0d8, 0x5c6dd14464619afa, 0x50606d69b733d4f3, + 0x7fb1af465b990f97, 0x3fab2634c8bbd936, 0x556da6168838b902, 0x0f15975902a30e1f, + 0xb29d782ae9e1991f, 0xae00e26ff8f7e739, 0xd3da86458bb292d5, 0x4528ee0afb27e4ce, + 0x49882d5ba49fabad, 0x7e873b6a7cf875ee, 0x777edd535113c912, 0x94ed05e7ff149594, + 0x0b8f95fc4211df43, 0x9135c2b42426fef2, 0x411e6c2b47307073, 0x503207d1af0c8cf8, + 0xd76f8619059f9a79, 0x64d24617855dee45, 0xf7bc7a877923196a, 0xd6cc42ed6a65be79, + 0xe3912ff09d4fc574, 0x4192d03b2bc2460a, 0xa0dcc37dad98af85, 0xfc59049b2a5818a4, + 0x2128bae90a5b975f, 0xbe7067ca05ea3294, 0x5bab7e7753064c4f, 0x42cbf0949ef88443, + 0x564df4bbd017492c, 0xf2c2eb500cf80564, 0x5b92e67eb00e92af, 0x8c4103eef59c0341, + 0x83412122b8284998, 0x888daf2da0636b6d, 0x4d54b10303dd07d6, 0x201190e7c1e7b5ed, + 0x3797510bb53a5771, 0x03f7bc598b570b79, 0xdc1e15d67d94f73e, 0x721e8b499ebe02c1, + 0x71f954f606d13fa0, 0x0c7a2e408c168bf0, 0x07df2ef14f69c89d, 0xe295096f46b4baaf, + 0x7a2037916438737e, 0xd1e861aeaf8676ea, 0xb36ebdce368b8108, 0xb7e53b090ddb5d25, + 0x5a606607b390b1aa, 0x475e52994f4a2471, 0xbcc2038ba55b2078, 0x28b8a6b6c80df694, + 0xb5f0130ec972c9a2, 0x7a87cd2a93276b54, 0x4d0eec7ecf92d625, 0xac1a8ce16269a42e, + 0xa4ca0237ca9637b8, 0xd8dc8ff91202b6ff, 0x75b29846799d7678, 0x761b11a5edd9c757, + 0xf2581db294ef3307, 0xe3173c2b6a48e20f, 0xe46fd7d486d65b3c, 0x1352024303580d1f, + 0x2d665dae485c1d6d, 0x4e0905c825d74d3b, 0x14ff470c331c229e, 0xbdc656b8613d8805, + 0x36de38e396345721, 0xaae682c1aa8ff13b, 0x57eb28d7b85a1052, 0xf3145290231d443a, + 0xd0f68095e23cbe39, 0x67f99b3c2570b33d, 0x54575285f3017a83, 0x9b2f7bb03d836a79, + 0xa57b209d303367a9, 0x7ccb545dd0939c79, 0x1392b79a37f4716d, 0x6e81bb91a3c79bcd, + 0x2c2cd80307dddf81, 0xb949e119e2a16cbb, 0x69625382c4c7596f, 0xf19c6d97204fb95c, + 0x1b2ea42a24b6b05e, 0x8976f83cd43d20ac, 0x7149dd3de44c9872, 0xc79f1ae2d2623059, + 0xca17a4f143a414e1, 0x66d7a1a21b6f0185, 0xed2c6198fe73f113, 0x16a5f0295cbe06af, + 0x5f27162e38d98013, 0xf54d9f295bdc0f76, 0x9ba7d562073ef77b, 0xa4a24daaa2cfc571, + 0x49884cf486da43cd, 0x74c641c0e2148a24, 0xbff9dcbff504c482, 0xf8fc2d9403c837ab, + 0x6ccc44828af0bb1e, 0xbcf0d69b4c19dfdb, 0x8fe0d962d47abf8f, 0xa65f1d9d5514271d, + 0x26ff393e62ef6a03, 0xc7153500f283e8fc, 0xea5ed99cdd9d15cd, 0xfc16ac2ba8b48bb7, + 0xf49694b70041c67a, 0xbd35dd30f5d15f72, 0xcf10ad7385f83f98, 0x709e52e27339cdc2, + 0xe9505cb3ec893b71, 0x2ffa610e4a229af7, 0x12e1bc774d1f0e52, 0xe301a3bb7eacccc8, + 0x1fdd3b6dcd877ebf, 0x56a7e8bda59c05aa, 0x99acd421035d6ab4, 0xfd21e401cecd2808, + 0x9a89d23df8b8d46f, 0x4e26b1f1eb297b9c, 0x9df24d973e1eae07, 0xe6cdc74da62a6318, + 0xfc360d74df992db0, 0xf4eca0a739514c98, 0x481c515ba9bf5215, 0xce89cce80f5f3022, + 0xf487a10fc80e4777, 0x235b379a87e41832, 0x76f72e028371f194, 0xd044d4a201325a7d, + 0x47d8e855e0ffbdde, 0x268ae196fe7334b0, 0x123f2b26db46faa8, 0x11741175b86eb083, + 0x72ee185a423e6e31, 0x8da113dfe6f6df89, 0x286b72e338bbd548, 0xa922246204973592, + 0x7237b4f939a6b629, 0x31babda9bedf039a, 0xb2e8f18c6aeec258, 0x0f5f6ce6dd65a45e, + 0x8f9071a0f23e57d3, 0x71307115ba598423, 0xcbe70264c0e1768c, 0x1c23729f955681a8, + 0xfbc829099bc2fc24, 0x9619355cbc37d5d6, 0xea694d4e59b59a74, 0xb41cf8d3a7c4f638, + 0xae1e792df721cd0b, 0x7cd855d28aac11f6, 0xca11ba0efec11238, 0x7c433e554ce261d8, + 0xe3140366f042b6ba, 0x8a59d68642b3b18c, 0x094fcdd5d7bccac2, 0x9517d80356362c37, + 0x4a20a9949c6c74e8, 0xc25bcf1699d3b326, 0xa8893f1d1ed2f340, 0x9b58986e0e8a886e, + 0x29d78c647587ce41, 0x3b210181df471767, 0xd45e8e807627849d, 0x1ec56bc3f2b653e3, + 0x974ff23068558b00, 0xdb72bdac5d34262c, 0x23225143bb206b57, 0xd0a34cfe027cbb7e, + ], + // seed = 3 + [ + 0x39209fb3eb541043, 0xee0cd3754563088f, 0x36c05fc545bf8abe, 0x842cb6381a9d396b, + 0xd5059dcb443ce3bf, 0xe92545a8dfa7097e, 0xb9d47558d8049174, 0xc6389e426f4c2fc0, + 0xd8e0a6e4c0b850d3, 0x7730e54360bd0d0d, 0x6ecb4d4c50d050d5, 0x07a16584d4eb229f, + 0x13305d05f4a92267, 0xb278ddd75db4baec, 0x32381b774138608f, 0x61fe7a7163948057, + 0x460c58a9092efee6, 0x553bf895d9b5ff62, 0x899daf2dabfd0189, 0xf388ab9c1c4b6f70, + 0xd600fe47027ea4cd, 0x16d527ec2b5ef355, 0x5ac1f58ff6908c81, 0xa08d79ff8ee9ffe8, + 0xc1060a80b7a5e117, 0x14b2c23118c60bda, 0x8cc0defbb890df8f, 0xe29540fd94c6d28b, + 0xa604f003f82d5b71, 0xa67583d4eb066d18, 0xd62cbd796322b3fc, 0x070cfe244cdcccf3, + 0x73557c30b3af47e5, 0x2e544e31153a2163, 0x996eef7464d5bead, 0xbc71cb5ab0586cdc, + 0x0bfcb6c1b517ed69, 0x62b4f1fcc82e8ca0, 0x0edbc68f544965c5, 0x40fa39baa24af412, + 0xf39aeb2413dab165, 0x17e6013e7afee738, 0x8109bff1c8d42a9d, 0x3cd99863390989b5, + 0x02021a4cc9c336c8, 0xa06060778cb60aa4, 0xd96591db60bc1e06, 0xd2727175183f4022, + 0xcdc1f1c5bce3e7ce, 0xb393ccc447872a37, 0xdf6efe63257ead3a, 0x20729d0340dbceb6, + 0x9f3d2d26fc0ea0d7, 0xf392e0885189bd79, 0xdf2ee01eb212b8b6, 0x6e103a0c0f97e2c3, + 0x96c604a763bd841b, 0x9fc590c43bba0169, 0xf92dcd5ddc248c40, 0x113a8b54446941dc, + 0x5943eda146b46bb8, 0xbf657901a36a39a7, 0x5a4e0e7ea6568971, 0xb94c635bae9f9117, + 0x2626fb65b3a4ef81, 0xa59bfd5478ce97de, 0x79112ba9cc1a1c63, 0xf41f102f002cf39c, + 0x0a589bcbfb7ff1c8, 0xa1478c53540c4fa1, 0x60d55e72c86dfaca, 0x312e7b6840ea7a39, + 0x8aae72dcccfe1f75, 0xff2f51f55bf0247a, 0x3c2e4b109edb4a90, 0x5c6d73f6525c7637, + 0xe49acb04a199f61c, 0x27860642d966df7f, 0x541ce75fb1e21c30, 0xd9fcd6f90806c7cc, + 0xb87c27bc93a7969b, 0x92f77a1179b8f8dc, 0xb1f29379deb89ed4, 0x7e63ead35808efe7, + 0x13545183d7fa5420, 0x575f593e34cf029d, 0x27f1199fb07344ae, 0xe67f95f7dc741455, + 0x49b478b761ab850b, 0xd7bedf794adfc21e, 0xdc788dcd2dda40ae, 0x14673eb9f4d8ad35, + 0x0cced3c71ecf5eb1, 0xe62d4e6c84471180, 0xdfe1b9e2cb4ada7d, 0x70185a8fce980426, + 0x0ce2db5e8f9553d6, 0x1fedc57bb37b7264, 0xb9310a2e970b3760, 0x989ff8ab9805e87d, + 0x0b912d7eb712d9ee, 0x1fe272830379e67c, 0x16e6a73aff4738fb, 0xeed196d98ba43866, + 0x7088ca12d356cbe2, 0x23539aa43a71eee0, 0xed52f0311fa0f7ad, 0xa12b16233f302eea, + 0xc477786f0870ecb4, 0xd603674717a93920, 0x4abe0ae17fa62a4c, 0xa18f1ad79e4edc8d, + 0xc49fe6db967c6981, 0xcc154d7e3c1271e9, 0xdd075d640013c0c0, 0xc026cd797d10922a, + 0xead7339703f95572, 0x4342f6f11739eb4b, 0x9862f4657d15c197, 0x4f3cb1d4d392f9ff, + 0xe35bffa018b97d03, 0x600c755031939ad3, 0xb8c6557ffea83abf, 0x14c9e7f2f8a122ea, + 0x0a2eb9285ee95a7c, 0x8823fec19840c46f, 0x2c4c445c736ed1d0, 0x83181dff233449f1, + 0x15ed3fca3107bef5, 0x305e9adb688a4c71, 0x7dbef196f68a3e2e, 0x93e47ece3e249187, + 0x8353c5e890ead93c, 0xea8a7ae66abafdf7, 0xf956dbb6becf7f74, 0x9f37c494fbfdb6e4, + 0x11c6cbaa2485dd32, 0x206f336fcca11320, 0x9befe9a59135d8fe, 0x5f3ef8b8db92c7db, + 0xbb305e556ce0ce9a, 0xf26bdafb1305887f, 0xcbf28abe23f08c61, 0x0bc64173b914e00b, + 0x9168da52e983f54a, 0x6ea41d09c3574a3e, 0x78aa44d4a74459ae, 0x2931422878387bf5, + 0x018f64a3a92c2d9c, 0x9be43f6752e66b34, 0xae378890decd1152, 0x07325329a1cb7623, + 0x3b96f4ee3dd9c525, 0x2d6ebcdbe77d61a3, 0x10e32b0e975f510c, 0xffc007b9da959bf9, + 0x38bf66c6559e5d90, 0xbe22bdf0bf8899fe, 0x87807d7a991632a8, 0x149a0d702816766a, + 0x026f723db057e9ab, 0xeeecb83625ec6798, 0xcec2ed5984208148, 0xd985a78e97f03c84, + 0xf96c279e7927b116, 0x99d5027b3204f6e2, 0x13a84878c3d34c55, 0x5cf5ec96229e9676, + 0x0bc36b07e4f8e289, 0xbed33b80a069914d, 0x2fbfbdd1ff4b9396, 0xab352bb6982da90f, + 0x154d219e4fa3f62b, 0x4d087512bb6b9be7, 0xc582e31775ee400e, 0x7dadb002ae8c4a4e, + 0xaae2957375c1aee2, 0x5f36ca643356625b, 0xf87cf8eb76e07fb7, 0x46f432a755e02cc3, + 0x36087e07aba09642, 0xe5642c1e4ebb9939, 0xb9152d22338eefad, 0xf7ba44278a22cf7f, + 0xd3b8013502acd838, 0x7761511da6482659, 0xb0857621638e8e50, 0x552eddb4a8b1d5f5, + 0xc43d9861e812c3ea, 0xd765c2aada47910c, 0x21c935b68f552b19, 0x6256d5641a2b47dc, + 0xab711d8e6c94bc79, 0xa8d0b91a2a01ab81, 0x5e6d66141e8d632a, 0x7638285124d5d602, + 0x794876dbca3e471f, 0x951937d8682670ce, 0x0f99cb1f52ed466a, 0x8c7cd205543b804c, + 0x2fd24d74a9c33783, 0xe5dcb7b7762e5af1, 0x45e6749cca4af77c, 0x540ac7ee61f2259f, + 0x89c505c72802ce86, 0xeab83b9d2d8000d1, 0x9f01d5e76748d005, 0xc740aaef3035b6d0, + 0x49afcd31d582d054, 0xcba5dc4c1efb5ddc, 0xc0a4c07434350ca1, 0xfc8dfaddcc65ee80, + 0x157c9780f6e4b2d9, 0x9762a872e1797617, 0xc4afae2cf3c7e1bd, 0x71cde14591b595d4, + 0x8843c3e0e641f3b9, 0xd92ecd91dce28750, 0x1474e7a1742cb19f, 0xec198e22764fa06b, + 0x39394edb47330c7d, 0x00ba1d925242533d, 0xaed8702536c6fb30, 0x6d3618e531c2967a, + 0x77f7cedcd7cc0411, 0xbc1e2ab82be5b752, 0x07b0cf9223676977, 0x596c693b099edd53, + 0xbb7f570f5b9b2811, 0x96bfdad3c4a6840c, 0x668015e79b60c534, 0x3ad38d72123f1366, + 0x6b994d81d2fcbb09, 0x70885f022c5052d8, 0xc891ee79d9306a7b, 0x2c4df05c0ed02497, + 0x19ebc13816898be2, 0xea7c64df11c392a2, 0xb7663e88dd12e1bd, 0x79f768cb8e154c21, + 0x1fb21b12e945933b, 0xe6a9045643f6906e, 0x544c47acd7e15371, 0xb7709b14f727e3d1, + 0x326ee36a46942971, 0x477f1cf7b0e2d847, 0x88b8f6b82b3b0c24, 0x18bc357b80e3cd5c, + 0x3333de70e4d66e0b, 0x4fd4c5e148583cf6, 0xae1b62f3008c0af3, 0xc49f419b6ab29cf5, + 0x2c29fa65afc3fa28, 0x4b19d93734d03009, 0x7dd6c09e589276ad, 0x1cece97f30de48ad, + ], + // seed = 4 + [ + 0x58bdf4338602e4fb, 0x71a5620b02c926d5, 0x3811c960129c2d9f, 0x29c2fb11fccac567, + 0x0d6b1ea7780f1352, 0xcc4d3ddfae3f87b3, 0xfdd30257362a586b, 0xabc948fde69f25f1, + 0x51b3523469d30f7b, 0xe0f0322724405ace, 0xd3729266d896da1e, 0xb10c37e5147915bf, + 0x8b577039f9fa32a3, 0xe677c6a9cbfb44b3, 0x7317a756ebb51a03, 0xf8e988ef37359485, + 0x600fc1ef3f469ff3, 0xbf0b8f8520444e01, 0x3711168b08b63d73, 0x34146f2944a6cb36, + 0x717feb263862cdde, 0x7185f8347db00412, 0x900798d82127e693, 0x84089e976a473268, + 0x10f8308c0d293719, 0xf62a618d4e5719b8, 0x8bdbd257a1a9516f, 0xf49f666fd7a75110, + 0xbaf45e2db7864339, 0xe4efa1ea0c627697, 0x3e71d4c82a09fe10, 0x54a2a51cf12127bb, + 0xa0592c9f54ba14cd, 0x27dd627a101c7a42, 0x3d2ceb44b3d20d72, 0x7ee1f94a68ca8f5d, + 0x7e8cb8651b006c36, 0xbd9fa7ca3a475259, 0x856de173586a7b34, 0xcedb291b594cb1b5, + 0xa3d6e462fd21cddc, 0x74561d10af9118e4, 0x13a3d389fc2d4b36, 0xeea8594a4a054856, + 0xf56d7474d9ba4b13, 0x25ddce2f6490b2fd, 0x920653ff3a8d830b, 0xcd8c0c9cdac740d1, + 0x2c348a738db9c4a0, 0x2967ccbe8ea44c22, 0x47963f69adb049f8, 0xf9d01eb5b4cf7eb6, + 0x7a5c26eb63a86bd2, 0x62ad8b7a71fa0566, 0xb373213179f250ae, 0x589d4e9a88245a4d, + 0x433dafebe2d558a8, 0x521fbef2c8fe4399, 0x62a31f9ff9ccd46b, 0x51602203eba7c1a6, + 0x9afc8c451b06c99f, 0xb529085bdbaffcea, 0xac251825cc75892b, 0x94976a5bce23d58e, + 0xdd17925b6c71b515, 0x568fd07a57bce92e, 0xefac31200d8bd340, 0x716c3e466b540ef9, + 0x3d2c9e380063c69b, 0x14168f9a3662dd83, 0xd298c7504dbc412f, 0x74490a94f016719f, + 0x0e0da431e1ab80c8, 0xe321f63dc6b169ae, 0xf08671544febc95a, 0x39324450cc394b3b, + 0xea6e3d35f1aa3a70, 0x8ef8a886508ce486, 0xdc1a631ef0a17f06, 0xfda2b3fbcd79e87b, + 0xd75bcae936403b10, 0xf88b5bd9f035f875, 0xc43efec2e3792dd4, 0xe9fac21a9d47cd94, + 0xc2876f0c4b7d47c3, 0xaba156cf49f368b4, 0x5ccda2170fa58bf9, 0xadc92c879ed18df7, + 0x110c1b227354e6c8, 0x298ee7a603249200, 0xde92142ede0e8ee7, 0x88e4a4610644ba9e, + 0xbb62d277e7641d3a, 0xb9be1985b7bf8073, 0x29024e5426cdb0d1, 0xf6aefd01f3092ab8, + 0x2a07087b313133aa, 0x6d71f445d6dfc839, 0x1e2412ff12e5526b, 0xed5cdeba6617b9e1, + 0x20b1d0d5e5f8760e, 0x12ff15705c368260, 0x7bf4338b7c387203, 0x34ff25f00cd06185, + 0x1148c706c518cf28, 0x5c04f0623388f025, 0xcb9d649275d87d79, 0x9b5f0c24fabc42ec, + 0x1a7b5e7964e33858, 0x2a81bbd8efdc6793, 0x8d05431ffe42752e, 0x83915cd511002677, + 0x580ed4d791837b31, 0x5982e041d19ff306, 0xcad0d08fa5d864ca, 0x867bee6efe1afa63, + 0x26467b0320f23009, 0xd842414dfda4ec36, 0x047fcdcbc0a76725, 0xbddb340a3768aeca, + 0xef4ce6fa6e99ab45, 0x88c5b66c7762bf9b, 0x5679f1c51ffb225d, 0xdab79048317d77ee, + 0xf14e9b8a8ba03803, 0xe77f07f7731184c1, 0x4c2aab9a108c1ef5, 0xa137795718e6ad97, + 0x8d6c7cc73350b88b, 0x5c34e2ae74131a49, 0xd4828f579570a056, 0xb7862594da5336fc, + 0x6fd590a4a2bed7a5, 0x138d327de35e0ec1, 0xe8290eb33d585b0b, 0xcee01d52cdf88833, + 0x165c7c76484f160e, 0x7232653da72fc7f6, 0x66600f13445ca481, 0x6bbdf0a01f7b127d, + 0xd7b71d6a1992c73b, 0xcf259d37ae3fda4a, 0xf570c70d05895acf, 0x1e01e6a3e8f60155, + 0x2dacbb83c2bd3671, 0x9c291f5a5bca81af, 0xd976826c68b4ee90, 0x95112eec1f6310a2, + 0x11ebc7f623bc4c9a, 0x18471781b1122b30, 0x48f7c65414b00187, 0x6834b03efa2f5c30, + 0x0875ef5c2c56b164, 0x45248d4f2a60ba71, 0x5a7d466e7f7ba830, 0x2bebe6a5e42c4a1d, + 0xd871d8483db51d10, 0x6ee37decd2fd392f, 0x7d724392010cede3, 0x8e96ef11e1c9bcc8, + 0x804a61d86b89d178, 0xbb1b83ce956055ec, 0xcb44e107410ff64f, 0xc426bb09ee0ba955, + 0x057c08f42c3dd7f1, 0x40ea1ec148602bdf, 0xc24688deeb65d7f1, 0xd8bcc53c768ba4e4, + 0x16e0e3af65c1106c, 0xfc12f7e7d647218b, 0x70d6e1d3ee93cef4, 0x01d2a505c4541ef9, + 0x1ef79e16e764d5c3, 0x0363d14d13870b98, 0xb56ef64345d06b11, 0xe653d557ebb7c346, + 0x8304a8597c2b2706, 0x1536e1322ce7e7bb, 0x525aec08a65af822, 0x91f66d6e98d28e43, + 0xe65af12c0b5c0274, 0xdf6ae56b7d5ea4c2, 0x5cef621cedf3c81c, 0x41e8b1ffd4889944, + 0xb5c0f452c213c3e5, 0x77af86f3e67e499b, 0xe20e76ea5b010704, 0xbdc205ab0c889ec0, + 0xc76d93eb0469cd83, 0x17ac27f65cab0034, 0xd49ec4531fd62133, 0x07a873ea2f1b9984, + 0xbff270dfef0032ee, 0x1764dbe91592f255, 0xe40363126f79e859, 0xa06cad3ab46971f6, + 0x0be596e90dedd875, 0x3387cce5c1658461, 0x44246acf88a9585e, 0xe0ad82b92d5ecb2c, + 0x2177491c9a1600a6, 0x16e7c4aac0f02422, 0x75792eeeec15c4e1, 0x2309cd359d08ee30, + 0x7cd9831dd1b83b0a, 0x374914a7c4ee8cf0, 0x0dd17765c9ac2e54, 0xb7847470ba9a7688, + 0xfba4f4bbe2991173, 0x422b203fc3de040e, 0x63bfcaf2ecf2ab0e, 0x0c5559f3a192946e, + 0xfdf80675c1847695, 0xf5f570accab842c9, 0x65cc5a448767afea, 0x1efeb0a7ee234f2f, + 0x9b05f03d81e7b5d2, 0xe7c31317a8626cf4, 0x620f2a53081d0398, 0x1b6de96cdd9943ae, + 0x8c226a436777d303, 0xa08fbbd50fafb10d, 0x6a64c5ec20104883, 0x9c9c653502c0f671, + 0x678a02b2174f52a0, 0x68e008ba16bbad4b, 0xa317c16d2efb860f, 0xeab2075d17ed714c, + 0x565eeeddf0c4ea15, 0x8ec8e94d242a6c19, 0x139e8e27d9000fae, 0xc977a7ff1b33d2f5, + 0x1d0accca84420346, 0xc9e82602cd436e03, 0x6a2231da53d2ccd3, 0xb44b12d917826e2a, + 0x4f4567c6a74cf0b9, 0xd8e115a42fc6da8f, 0xb6bbe79d95742a74, 0x5686c647f1707dab, + 0xa70d58eb6c008fc5, 0xaaedc2dbe4418026, 0x6661e2267bdcfd3d, 0x4882a6eda7706f9e, + 0xf6c2d2c912dafdd0, 0x2f2298c142fd61f9, 0x31d75afeb17143a8, 0x1f9b96580a2a982f, + 0xa6cd3e5604a8ad49, 0x0dae2a80aad17419, 0xdb9a9d12868124ac, 0x66b6109f80877fac, + 0x9a81d9c703a94029, 0xbd3b381b1e03c647, 0xe88bc07b70f31083, 0x4e17878356a55822, + ], + // seed = 5 + [ + 0xb3c58c2483ad5ead, 0x6570847428cdcf6c, 0x2b38adbf813ac866, 0x8cb9945d37eb9ad3, + 0xf5b409ec3d1aed1c, 0xa35f4bffc9bb5a93, 0x5db89cde3c9e9340, 0xff1225231b2afb2b, + 0x157b0b212b9cc47d, 0xf03faf97a2b2e04d, 0x86fdab8544a20f87, 0xfcb8732744ae5c1c, + 0xd91744c0787986d5, 0x5f8db2a76d65ad05, 0xcff605cbed17a90d, 0xf80284980a3164e7, + 0x59cc24e713fccc7d, 0x268982cada117ce4, 0xcd020e63896e730e, 0xe760dc46e9fe9885, + 0x6aaece8ab49c6b5d, 0x7451194d597aae3e, 0x35d4385900332457, 0xa40fb563a096583d, + 0xa797b612f7f11b76, 0x2fed6eb68e6a2b9b, 0x2f06ee64aeffd943, 0x9dd0e49d9ca45330, + 0x97d48f08bd7f1d8f, 0x1cfa7fe3ebe4d8ee, 0x2a2ba076bd397d42, 0x68c4344f7472f333, + 0xce21ec31987d74b5, 0xb73dabdc91d84088, 0x801aadee592222fe, 0xaf41345398ebc3f5, + 0x8a8f653d7f15ee46, 0xce2d065ff2ba2965, 0x4e05da515da2adb7, 0xa6dbdb8aa25f0fd4, + 0xca9f9666bbd2d5a9, 0x6b917ce50bd46408, 0x1550cc564ba6c84d, 0xb3063ae043506504, + 0x84e5f96bb796653d, 0xe2364798096cf6e3, 0x3b0dfedf6d3a53d0, 0xb7e4c7c77bde8d93, + 0xe99545bac9ab418a, 0xa0e31f96889507bb, 0x883c74f80c346885, 0xf674ae0b039fd341, + 0x8bb6ce2d5e8d1c75, 0x0c48737966a7ed7c, 0x04fcdf897b34c61c, 0xe96ac181bacbd4d6, + 0x5a9c55a6106a9c01, 0x2520f020de4f45d3, 0x935730955e94d208, 0xce5ad4d7f3f67d3b, + 0xa4b6d107fe2d81ca, 0x4f0033f50ae7944e, 0x32c5d28dd8a645a7, 0x57ce018223ef1039, + 0x2cbab15a661ab68e, 0x6de08798c0b5bec2, 0xee197fb2c5c007c6, 0x31b630ac63e7bda2, + 0xab98785aefe9efe3, 0xa36006158a606bf7, 0x7b20376b9f4af635, 0xa40762fdc3c08680, + 0x943b5faffd0ebee2, 0x7f39f41d0b81f06e, 0x7c4b399b116a90f8, 0x24e1662ac92bc9f3, + 0xcf586fc4e8e6c7db, 0xe46e0d047eeb12d7, 0xe8021076e4ea9958, 0x11fc13492e3ca22a, + 0xd61eae01410397e3, 0x7e8c4a58036a8e9f, 0x068a6de267970745, 0x64faab129bef1a41, + 0xb4a6f720943dad01, 0x631491058d73a9d5, 0xdad4fe95eab3ec02, 0x0a8b141c5c3a44f6, + 0x9fc69d4c2b335b98, 0x94d5f84a07d6e4cd, 0x1b73965de143c608, 0x443932c2dda54bcc, + 0x7397818fb0b04cd2, 0xef4ab03a1202b277, 0xf3d2ee459c0c2b92, 0x182d4daf8b058a87, + 0x90e63035d7b51368, 0xba4cd8b9a95d45fd, 0x12a7392c76731090, 0x890d264ec5d082d2, + 0xeeaf5c363da4994e, 0xd6aad756902123fb, 0xb531ebebdb28f191, 0xe71ce659fc59babd, + 0x37c1b94f63f2dcb5, 0xe4e3abeb311f9b96, 0x4a31b72ccb8695d3, 0x52cae1f0629fdce4, + 0xe5b0475e2ed71369, 0x2724e8c3506414fb, 0xbab0367920672deb, 0x0161a781c305449f, + 0x37b70f40f5bb60be, 0xddd1094c50251a01, 0x3b28283afd17224e, 0x06dec0cfe889fc6b, + 0x47608ea95bb4902d, 0xad883ebc12c00e82, 0x9e8d7ae0f7a8df29, 0xa79443e9f7c013a1, + 0xcfa26f68b7c68b71, 0x33ae6cc19bda1f23, 0xd9741e22b407887f, 0xf2bff78066d46b1c, + 0x794123191c9d32d4, 0x56cb6b903764ec76, 0x98775d0ef91e1a5a, 0xae7b713bc15c1db9, + 0x3b4c1a7870ed7a0d, 0x46666965f305cc34, 0x0ea0c3b2e9c6b3cd, 0x4dc387039a143bff, + 0x5f38bb9229ef9477, 0xea5d39ba72af7850, 0x69a5ed0174ce2b6d, 0x06969a36bfe7594d, + 0x0adee8e4065ccaa3, 0x908a581d57113718, 0x64822d6c5a8190ed, 0x8c5068b56ace4e4c, + 0x88ba3b4fb4e30bef, 0xa6ec0b8bb5896cfe, 0x4e23fcc6b47996fd, 0xe18e75b0dd549c7a, + 0xcd90f17e106cf939, 0x1666fdfb2ef7c52f, 0x4fae325f206dd88c, 0xe7bc1160e25b062d, + 0x3cc999cb246db950, 0xc5930a7326cd5c37, 0xb008a48a211367bd, 0xc5559da145a88fd4, + 0x1e3ad46655fac69c, 0x7834266b4841bfd7, 0xa764450fbffc58cc, 0x54d8cf93a939c667, + 0x93c51f11b21b2d9d, 0x0964112082ed65cc, 0x4c2df21213e7fb03, 0xf0405bc877468615, + 0x17b4fc835d116ab4, 0xa6b112ae5f3cb4ef, 0x23cfc8a7fd38a46e, 0x8e0a360dc2774808, + 0x24ca9c8092105ad5, 0xafd3f75524f2e0d5, 0x4f39ed7dbaddc24c, 0xe5e362c7679a7875, + 0x00914a916b07b389, 0xdfe1119b7d5ab5da, 0xabd6ed9940e46161, 0x630ed2044171e22c, + 0xdecc244157dd1601, 0x777e6d5b4b4868d5, 0x9b3530bee67017d8, 0xd2faf08b291fdcb9, + 0x006e99455d6523de, 0xd559b5817f6955b5, 0xefcc1063b0088c61, 0xed73145ae0f00ae7, + 0xab2af402cf5b7421, 0x897767f537644926, 0x26c9c0473ca83695, 0x192e34e1881b2962, + 0xf7cf666ec3b3d020, 0x27f9b79c7404afb7, 0xe533e8bed3010767, 0xe5817838e11d05d3, + 0x65659c531bd36517, 0xd427c5e0a23836fd, 0xf3eab7ea58fa3528, 0x07683adae1289f35, + 0x201d6af7e896dd32, 0xd5da938b9a21ad88, 0x843fb73ad67bc316, 0x1782ec7d5feef21b, + 0x943f66f6ec772877, 0x7e9112e7b26da097, 0xeac8161f8663c2c7, 0xe8600db480a9ebf4, + 0x07807fc90f6eaf5f, 0xe0e4c9deb41abf83, 0xbdf533db271f9c15, 0xb398411b0497afe2, + 0xdebb45ef25448940, 0xe7a5decefcd376c4, 0xaf1ef3c728c83735, 0xb8b83a99355cb15a, + 0x6444a0344f1611e4, 0xe8bb7f5cf3c60179, 0x77ab5c5177e75ff7, 0xc38fd6fa849d585d, + 0x390d57d53029060a, 0xa66327eb7b8b593c, 0x6350a14f6fcd5ac9, 0x2c08125bcd7008b4, + 0x2d00c299a6a6bf8e, 0x6b0039c1f68d1445, 0x0035150c5d06f143, 0xa34d01628cc927e1, + 0xdf5b3164d7b2ede1, 0x8167db1d0583d72e, 0x4e13b341cd2ae8bc, 0xa693d9b1f416e306, + 0xc15ed7ca0bc67609, 0xdc344313c1c4f0af, 0x88b6887ccf772bb4, 0x6326d8f93ca0b20e, + 0x6964fad667dc2f11, 0xe9783dd38fc6d515, 0x359ed258fa022718, 0x27ac934d1f7fd60a, + 0xd68130437294dbcc, 0xaf5f869921f8f416, 0x2b8f149b4ab4bf9f, 0xc41caca607e421cb, + 0x7746976904238ef9, 0x604cb5529b1532f0, 0x1c94cd17c4c4e4ab, 0xe833274b734d6bbe, + 0xe9f1d3ef674539ce, 0x64f56ed68d193c6a, 0xe34192343d8ecfc1, 0xcb162f6c3aa71fe8, + 0x99eaf25f4c0f8fa4, 0x92f11e7361cb8d02, 0xb89170cddff37197, 0x4f86e68a51e071e3, + 0x31abf6afd911a75b, 0x6d20cf259c269333, 0x4150b9f88fcb6513, 0x705063989ebf7451, + 0x559231d927c84410, 0x1ca8ec4b098bc687, 0xebed22405c9180e0, 0xaa815b37d052af59, + ], + // seed = 6 + [ + 0x946ac62246e04460, 0x9cebee264fcbc1ae, 0x8af54943a415652b, 0x2b327ed3b17b8682, + 0x983fde47b3c3847e, 0x10a3013f99a2ad33, 0x6e230bb92d2721ef, 0x1cf8b8369e5c5c50, + 0x7f64017f2b7b3738, 0xd393248a62417fa1, 0x9ff01c0b20a372c5, 0xb0e44abce7e7c220, + 0xcebb9f88d48a815f, 0xdb7df6bd09033886, 0x7844fc82b6fa9091, 0x72d095449863b8ec, + 0xc13e678c89da2c7e, 0x6caf4d5ad231d12f, 0x2e0ab7b5fcf35c49, 0xf410720cb932a70f, + 0xd66ea581f16fce06, 0x175c9f002f57dc98, 0xccbcfd0d32988775, 0xfde4c407d3b0a232, + 0x5db2931ae7e97223, 0x6e07e2173085809f, 0x6e1d1ec0f9cad73c, 0xb2fc251a7f802619, + 0xbc1fc17f04f342de, 0x8de8f21ec658e078, 0x72c0f40cbee53fd6, 0x0678244411fc17a1, + 0x1d5837ca166b9bbd, 0xc8cada003c554345, 0x6a2fe2bfb2e58652, 0xfca9d797a6f7988b, + 0x6699e24ac737948b, 0x69623ffcb05789ba, 0x946429c529d95b75, 0x0d14df0b2a13970f, + 0x593d8592c440dfec, 0x2ee176f3d7e74b94, 0xae003f1da3be9e26, 0x0c7b02c4c0f6764a, + 0x3117e2fa1f632462, 0xf0f23265b6f1eaeb, 0x3111255d9b10c137, 0xc82745e509a00397, + 0xbd1d04037005fea7, 0xe104ab0dd22a9036, 0x51b27ce50851ac7a, 0xb2cb9fb21b471b15, + 0x29d298074c5a3e26, 0x6ebdf2058b737418, 0xc4a974041431b96f, 0x1ec5a30ccb6bdaac, + 0xe818beede9bf4425, 0x4b69b1bce67a5555, 0xf5c35f1eb0d62698, 0xf4509bbd8e99867c, + 0xb17206debd52e1bc, 0x35785668c770b3be, 0xe9343987ff5863bc, 0x2ee768499ac73114, + 0x5132bb3426eeaaf4, 0x471bce2c6833c5ff, 0xbb9a2d5428e6f6f9, 0xd5678943c595792d, + 0xab2a65e7f81e479c, 0xa82407bb23990b31, 0xdae321383984923c, 0x01823bb22648e6f1, + 0xda6e8df4214a8b04, 0x0e172bb88e03d94f, 0x552da6c22e362777, 0x7ce67329fb0e90cb, + 0x7b2d7f287ede7ebf, 0xd44f8222500651bd, 0x4acca1ef58fbb8ab, 0x428ecf058df9656b, + 0xd7e1ec6a8987c185, 0x365be6a54b253246, 0x168849be1e271ee8, 0x6a00f3c4151a8db2, + 0x37602727ca94b33d, 0xf6b50f18504fa9ce, 0x1c10817f6bc872de, 0x4bfe1fe42b0f3638, + 0x135fad4b8ef6143b, 0x1b25ad2bafc25f58, 0x41e37f85cf321f92, 0xfc73f75d9d5b9bea, + 0x9eb3694d1e9cb7e1, 0x601d51f08fa83b90, 0x234a2a9b88366f41, 0x63fe903e16f2c3bf, + 0x1cdbd34fa751c0b0, 0x0ce4fc6747c0558c, 0x51ed72afb8bb49aa, 0x20313ba13ca12c96, + 0x271fa38f9ebd54c1, 0x3696a5ac03a8edde, 0x05602be7df625702, 0x11f1ac73790f7a9f, + 0xa2836c099f0810bd, 0xe5ac2e47caa532fa, 0xd9c000a66d39f681, 0xd93d900e6f3d9d5f, + 0x792c81c65b7900f2, 0x5c5dce790ee20da1, 0x74ff1950edec1aee, 0x71fc85fa1e277d8f, + 0x0e77df17d6546cbc, 0x07debad44816c3b4, 0xbafa721581e92a70, 0x8ab6fbe2ed27bba8, + 0xe83243a20dea304a, 0xaa85a63a84c00a07, 0xde0e79917fc4153a, 0x21bb445e83537896, + 0xeedcac49fc0b433a, 0xffb2926a810ae57a, 0xf724be1f41d28702, 0x79cb95746039bb3b, + 0x5a54fe3742a00900, 0xda4768d64922c04f, 0x420396a84a339dae, 0xa171e26ee5e8724e, + 0x4c8da7c5d289c20a, 0x9ebd79a1a8e94742, 0x39235232b97e9782, 0xb75df0be9bba7d80, + 0x0c1d204dd87d48fc, 0x8f81f3e7177266e8, 0xe4a460b39e78d72b, 0x50b98fa151e65351, + 0xb7cb585c3ee1eddc, 0x11cdad9a76ee1dc4, 0xa38054a78595dc1c, 0x92f09e2ec4978edc, + 0xa8f0061b5efdabaa, 0x04bcc4abc224d230, 0xc58606738e692d46, 0xdd2b27b565952433, + 0x19e6ed1b740beec0, 0xceadd49b2ef9891f, 0x328178c28fe95cad, 0xe5ad4c43afe02848, + 0x03c0cb538cd967c0, 0xec4352526d19a630, 0x4c7e99389d39b031, 0xf65dd05362c2deb6, + 0xd1e70daf6879d28d, 0xbe9f57db6309b265, 0xa4b66f370b872bb7, 0xe26896fbc6ee1fd5, + 0xac705e661bfcf7c5, 0xab4d0d07d7f09940, 0x976417c06aeb6267, 0x8161c684a6bd468c, + 0xf77b6b9976dc4601, 0xc6489b779a39c12c, 0xb2aa58d5681cea1a, 0x043b1b40f8c3e04c, + 0x681fcbfadc845430, 0xab8896c921ba8def, 0x57aaf172606f37b2, 0xc3735048cd5eb8d7, + 0xa7078b96955631bd, 0xdd6b3543aa187f33, 0xc7103ea4a2a697fd, 0x8d7b95f6ff1f7407, + 0xe44f419e84709530, 0xf340caa9132cbb0a, 0x2ba407283143c66c, 0xe1be240ca636c844, + 0x90d32f2877ac08bc, 0x5d26e6294b2c8673, 0x4a6b2f5b27c87a44, 0x961fb9043f76d34f, + 0x0afee02d8d3c55d2, 0x6228e3f48c42e5dc, 0xc338e69ee6593675, 0x853f74b16efb7bdd, + 0xd062f40bdd22e687, 0x647164b9ab4c4190, 0xf94689f67d598369, 0x8e4b29d87a5012d7, + 0xaf02b8b925656fbd, 0x7a722a767179a630, 0xb5c8afe937a75ace, 0xfdb8e8d02d279372, + 0x887ef700cb25fae1, 0xcfe9bd912f72cabe, 0xb1d4dedc24f978de, 0x517522d38319cc2a, + 0x7dd87b2b36aab798, 0x579c4ff3046b5a04, 0xf5c5975c5028b7a7, 0x7094579d1000ec84, + 0xbc8d5b1ea70a5291, 0x161b2d783be8855c, 0xd26d0b0d6d18279f, 0x0be1945f02a78bd5, + 0xb822a5a9e045415b, 0x2fe9d68b1ccc3562, 0xb2e375960033d14f, 0x26aca04e49b4ff22, + 0x732a81c862112aea, 0x8bd901ed6e4260b8, 0xe839532c561ad5b0, 0x8fb6e4d517a79b12, + 0x0dd37f8c0be9b429, 0xc8ad87ad12f1b1b0, 0xc51f3aa62b90318b, 0x031a7e8b86c1cefc, + 0xa95547af2b70fc76, 0x9cb3615c5a98801e, 0xa387e3c3341d7032, 0xa087ea52a1debaef, + 0x16325ec9a2e6e835, 0x587944a484c585eb, 0xc8879033bde22ecc, 0xa39dbfce709c464a, + 0x7acc010f99208774, 0x98dd2973a096c5ad, 0x26458b51139f198c, 0x2f5d19575e8c4f02, + 0x726643f0d38af352, 0x44d879b6d73e6e94, 0xa68a03885c980abe, 0x06048acd161c40c0, + 0xa4dab8f89d405d28, 0x7120c880cb04be18, 0xa062ace22a1cf0cf, 0x3901a9daf29704f4, + 0xff08f3ed989db30a, 0x6d22b13e874c67e9, 0x80c6f35518d73f4d, 0xc23c2a521aac6f29, + 0x2e708fd83aaa42e0, 0x7fc3780f55f1b0fd, 0xabb3075c98cf87f2, 0xb4df3f40f7c61143, + 0x2a04418098a76d75, 0x0d9eeee9509b2d37, 0x6be8ae51f4b59cdc, 0xe746cc7c00e4a2ab, + 0x785bc6df9cac597c, 0x33cb6620ce8adc48, 0xc1ba30739bffcef7, 0x6d95771f18e503f7, + 0xf7be3ae2e62652ff, 0xc8d82ffd2a73c62b, 0x8725a3ba5b110973, 0x67ed6b9c724757ec, + ], + // seed = 7 + [ + 0xc0272d42c19ff3ae, 0x4694228b43ea043b, 0x5709a6ef8a462841, 0xc9210a1e538805c9, + 0x279b171196113ec2, 0x859b769fc2d9e815, 0x0d5d3125a2bf14d3, 0x22bca1cfefa878ba, + 0x481b6bf58037bd83, 0x4933ba8647728d22, 0xf08c7b6b56f6e1b6, 0x374e8af5a15407c7, + 0xa95c4dc3d2487a5c, 0x9b832808ff11e751, 0xf2048507e9da01d5, 0xa9c576189f544a4a, + 0xf6c2a45b2e9d2b41, 0x9b9874c9f10ecc2f, 0x37d9b5f51f8c149e, 0x93aead54c9de9467, + 0x59cf0b4af262da23, 0xe7e9929af18194b2, 0x9df2644e33eb0178, 0xde4122d6f0671938, + 0xf005786c07f4800b, 0xb1fc9d254b5d1039, 0x0bf1088631f6dd7b, 0x665623f0a4b8f0c7, + 0x60f0113a9187db7c, 0xfd7cceda4f0d23a6, 0x26c01e9d89955940, 0x33afa1dfc0f5a6a0, + 0xeb77daf215e9283c, 0xc7575214bf85edb4, 0xeb0d804bf297e616, 0x84bff4ffd564f747, + 0xc4ac33189246f620, 0x43ef61213ecc1005, 0xcbbb0dea6cd96acd, 0x8ed27abfa8cfcb05, + 0x543b61529cb996b6, 0xa5f987ca41ea5e59, 0x3c50e0ac5254cb7a, 0x4192b0446c06d1e6, + 0x3e86592e21b45388, 0xdb766f06fcc6e51e, 0x0448ee36efe632db, 0x663c9db689253e35, + 0x72e0bd4985331dd4, 0xff501b5bf7d94e74, 0xe911ce758e2113a8, 0xec3a8d03a75a6ba4, + 0xaf6b4b72f56edc83, 0xf284857936c0a391, 0x5ba6feff407d46f4, 0x9d689c26de9d6702, + 0x28c04a9083726b5d, 0x2ccf4a627a029730, 0x7b4719500d4f0c71, 0x76470a9a7da250a8, + 0xcc48409404a1c890, 0xccefbdc7ec9a8055, 0xe0db91bff3cc42d3, 0x0532436426141254, + 0xf2ee9325e6f0ff0b, 0x149c20a5fbb28d9d, 0xe71624cd8d2d14d4, 0x8f01d4dc8cc2dd77, + 0x29cf409b333015b7, 0xba8bebd211884dd1, 0xc3396635e8c8db1d, 0x8ed0f6208d0528b8, + 0x0d90b43fdd0ee334, 0xd73c9a3333a044c7, 0xa2595cd208dbdc38, 0xae93cb264f940c09, + 0x8e0538d8afb07a97, 0x19115ec881385ba2, 0xa886f9e6a8039c6a, 0xcd5d62147ce3ecac, + 0xaecdf9e0bb4969f7, 0x2ddd631c53dcad10, 0x73ad1c97b3412054, 0xb08915fa2722efc6, + 0x97966047e5067eb0, 0x337f1675ed91445c, 0xb3a833d150b96a0d, 0x5940a98fe35e5e2e, + 0xfd03cc354ed0d8ff, 0x4e65b98291a8644a, 0x14a259f2852a60b2, 0x7648e3478c1e8e5f, + 0xbc0fbef6d9a919b4, 0xbec4302081346cf1, 0x57d2ce7aa1c7c511, 0x234c209d8f4e1ac3, + 0x87cf80cc933ce443, 0x7c262c616931e94e, 0xc5e33b049cf9eddf, 0x1a80790ed03ae51b, + 0xf2e8b9494f7220cf, 0x124cb59c14fff3ff, 0xa8a06cbfdb86ce18, 0x9068ef1f80b37653, + 0x0c55417b8d90338f, 0xcd579a523f6bcd30, 0xa31bfe2476a8d2a9, 0x1f8d142208094223, + 0x332dc40a5203cfad, 0xf8792fe5b2d33b4c, 0x443bd9668bf9461e, 0xc9019db0ace1409e, + 0x781bea919a113e8b, 0xb0f11d866abfbeec, 0xcfe139a60db0c26a, 0x869ab8721e6aa39e, + 0xdb48a4977717837a, 0x588a5ff151065b18, 0xe4a251ea0028864d, 0x7f0e43ba408a77c3, + 0x65f66dd50a536135, 0x6f49e934d9331c3e, 0xb8d742e0f0fa6b09, 0xe4e9b272deca2348, + 0xaee132ff902f773c, 0x43f658f7c2a0c90a, 0x28cb4dbc76cc53ea, 0x7d92253aa99ac39b, + 0x4fea3d832370baab, 0xb29e36936e51d78e, 0xea10778712321064, 0xff4f21f8ef274be2, + 0x84eff18ddfa0933f, 0xd0ec6a9f86c758a0, 0xaf82e5973c431ae0, 0x352023c00c045425, + 0xad34d7bc4a2f8961, 0xbdb4a02a24d4dee0, 0x354a4846d97447cf, 0x331a8b944d5bc19f, + 0x5ce04f8e17909035, 0x6497581bad8f4aab, 0x07c503bba647111e, 0x85f412ba78e1f7ff, + 0x7f3b920fd20f4cff, 0x424e1a9a4ce34e2f, 0x3035e2d62e1b9f0a, 0xef63114bff7b729a, + 0xe86a05889ab6bb60, 0xee0830cf095585a1, 0x4a54f7fa47d9c94b, 0x17daeece9fcb556a, + 0xc506d3f391834c6f, 0xb3f24be362e1af64, 0xc435e4e23608efdd, 0xeeba9caaa4cc1768, + 0x5a71f306daddc22d, 0x18e5205f41eba1a0, 0x7b29b4d1f6610925, 0x065cb65a0258d9a9, + 0x3e5ac8faa9fd1f95, 0x3b362362c1ea0470, 0xce0e4f6434db7a2e, 0xf327341098de52f2, + 0xcfca3b9e2a1992c3, 0x7483bf9401233e41, 0xbafbac531c6f9281, 0x4b52dd71b2c106f8, + 0xdf73b66e50b5a1f7, 0x237aec0202a20283, 0x23dd5be23dffdf2b, 0xea9730731ee122ef, + 0x5cb3f846014fbcd3, 0xc3b21c8ffdce9201, 0x06a99a02f91a8760, 0x721a81fa8fd7b7a3, + 0x6aafcdddc53cbcd8, 0xd03b464005a93bcc, 0x8212edc1b1669dcb, 0x71f4c31364c31bc7, + 0xfeeec0eba8772307, 0x1948d00a13d88cf1, 0x19064fd6d943ada8, 0x4ec8d31722697bfd, + 0x596d9a953a516609, 0xc4cb4bff53507da2, 0x1d59f3c5be36e4ca, 0xe5b4fc5bf6044c9b, + 0x1bb74e052232f735, 0x04e8a0db611ddd5d, 0x8d04eaa009b421bf, 0xa7878ae0ac0e6d58, + 0x28c1030217cab2b3, 0x827943767e56a883, 0x28fce5fa02d22809, 0xb30c322fffc8c58e, + 0x1ca5a6a9f8066c5b, 0xb24db5f1462b2513, 0x02f653b89b7e5f6c, 0xe31f8fb5d5f78eee, + 0x266acc514ed93501, 0x936879d1c6fddcc4, 0xcd51be3636af1952, 0x3fdbb6fc332c78c8, + 0x9eb656379fa73094, 0x056146cc92fa0f96, 0xed6c4f1836c027c3, 0x021e0bb5d2113f2a, + 0x8983e42ec1c626b3, 0x73ea9bc6513ad9c9, 0x0c904903b24f4247, 0xacbac1e6243e2525, + 0x0b1069a0c230fb06, 0x77d709fca3fc1ce5, 0x87ad0f65020947e6, 0x555302641c53f4e6, + 0x65ea87871fa9aaee, 0x58aaf4ecc1067bb4, 0x1a66c48cc4c65b3f, 0xca96aca48b2ea969, + 0xa68eb70bad14de2b, 0x5ccdb3d7e00a6f6e, 0xe178fbfec73fe72f, 0x2b63d6a16b83e890, + 0x32fdb7a5330fbae0, 0x2ab5803c8d1bf32c, 0xda838388c1527c94, 0x16a50bdc4de24acb, + 0xe561301f134c074a, 0xd7ae63d2816b4db1, 0x036aabd4df0dd741, 0xc5e0db8783435b9d, + 0x9c4386cf0a07f3b2, 0x6a72ac1aa56a13a1, 0x299bbdb04bb20a23, 0x138c1018fda16b81, + 0x0e354f0b3bda49df, 0x9f4c295b23127437, 0xd133ceb2bd561341, 0xd8b4bfd5a526ac29, + 0xcdd0a70ddc1c7bbd, 0x81dce595bf572225, 0x1c6f925c05f6efd7, 0x8ae5097553856ea0, + 0x3aabeaeef248f60d, 0xd9005809d19a69e2, 0x2a3a1a314311cc27, 0x89bb2dc76b2b624a, + 0x50a2a95d0412e289, 0x9def8df564e68581, 0xf49010a9b2e2ea5c, 0x8602ae175d9ff3f0, + 0xbf037e245369a618, 0x8038164365f6e2b5, 0xe2e1f6163b4e8d08, 0x8df9314914f0857e, + ]]; diff --git a/vendor/parquet/src/column/chunker/mod.rs b/vendor/parquet/src/column/chunker/mod.rs new file mode 100644 index 0000000..42631e0 --- /dev/null +++ b/vendor/parquet/src/column/chunker/mod.rs @@ -0,0 +1,40 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Content-defined chunking (CDC) for Parquet data pages. +//! +//! CDC creates data page boundaries based on content rather than fixed sizes, +//! enabling efficient deduplication in content-addressable storage (CAS) systems. +//! See [`CdcOptions`](crate::file::properties::CdcOptions) for configuration. + +mod cdc; +mod cdc_generated; + +pub(crate) use cdc::ContentDefinedChunker; + +/// A chunk of data with level and value offsets for record-shredded nested data. +#[derive(Debug, Clone, Copy)] +pub(crate) struct CdcChunk { + /// The start offset of this chunk inside the given levels. + pub level_offset: usize, + /// The number of levels in this chunk. + pub num_levels: usize, + /// The start index into `non_null_indices` for this chunk. + pub value_offset: usize, + /// The number of `non_null_indices` entries in this chunk. + pub num_values: usize, +} diff --git a/vendor/parquet/src/column/mod.rs b/vendor/parquet/src/column/mod.rs new file mode 100644 index 0000000..9c7e77d --- /dev/null +++ b/vendor/parquet/src/column/mod.rs @@ -0,0 +1,130 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Low level column reader and writer APIs. +//! +//! This API is designed for reading and writing column values, definition and repetition +//! levels directly. +//! +//! # Example of writing and reading data +//! +//! Data has the following format: +//! ```text +//! +---------------+ +//! | values| +//! +---------------+ +//! |[1, 2] | +//! |[3, null, null]| +//! +---------------+ +//! ``` +//! +//! The example uses column writer and reader APIs to write raw values, definition and +//! repetition levels and read them to verify write/read correctness. +//! +//! ```rust,no_run +//! # use std::{fs, path::Path, sync::Arc}; +//! # +//! # use parquet::{ +//! # column::{reader::ColumnReader, writer::ColumnWriter}, +//! # data_type::Int32Type, +//! # file::{ +//! # properties::WriterProperties, +//! # reader::{FileReader, SerializedFileReader}, +//! # writer::SerializedFileWriter, +//! # }, +//! # schema::parser::parse_message_type, +//! # }; +//! +//! let path = Path::new("/path/to/column_sample.parquet"); +//! +//! // Writing data using column writer API. +//! +//! let message_type = " +//! message schema { +//! optional group values (LIST) { +//! repeated group list { +//! optional INT32 element; +//! } +//! } +//! } +//! "; +//! let schema = Arc::new(parse_message_type(message_type).unwrap()); +//! let file = fs::File::create(path).unwrap(); +//! let mut writer = SerializedFileWriter::new(file, schema, Default::default()).unwrap(); +//! +//! let mut row_group_writer = writer.next_row_group().unwrap(); +//! while let Some(mut col_writer) = row_group_writer.next_column().unwrap() { +//! col_writer +//! .typed::() +//! .write_batch(&[1, 2, 3], Some(&[3, 3, 3, 2, 2]), Some(&[0, 1, 0, 1, 1])) +//! .unwrap(); +//! col_writer.close().unwrap(); +//! } +//! row_group_writer.close().unwrap(); +//! +//! writer.close().unwrap(); +//! +//! // Reading data using column reader API. +//! +//! let file = fs::File::open(path).unwrap(); +//! let reader = SerializedFileReader::new(file).unwrap(); +//! let metadata = reader.metadata(); +//! +//! let mut values = vec![]; +//! let mut def_levels = vec![]; +//! let mut rep_levels = vec![]; +//! +//! for i in 0..metadata.num_row_groups() { +//! let row_group_reader = reader.get_row_group(i).unwrap(); +//! let row_group_metadata = metadata.row_group(i); +//! +//! for j in 0..row_group_metadata.num_columns() { +//! let mut column_reader = row_group_reader.get_column_reader(j).unwrap(); +//! match column_reader { +//! // You can also use `get_typed_column_reader` method to extract typed reader. +//! ColumnReader::Int32ColumnReader(ref mut typed_reader) => { +//! let (records, values, levels) = typed_reader.read_records( +//! 8, // maximum records to read +//! Some(&mut def_levels), +//! Some(&mut rep_levels), +//! &mut values, +//! ).unwrap(); +//! assert_eq!(records, 2); +//! assert_eq!(levels, 5); +//! assert_eq!(values, 3); +//! } +//! _ => {} +//! } +//! } +//! } +//! +//! assert_eq!(values, vec![1, 2, 3]); +//! assert_eq!(def_levels, vec![3, 3, 3, 2, 2]); +//! assert_eq!(rep_levels, vec![0, 1, 0, 1, 1]); +//! ``` + +#[cfg(feature = "arrow")] +pub(crate) mod chunker; +pub mod page; +#[cfg(feature = "encryption")] +pub(crate) mod page_encryption; +#[cfg(not(feature = "encryption"))] +#[path = "page_encryption_disabled.rs"] +pub(crate) mod page_encryption; +pub mod page_store; +pub mod reader; +pub mod writer; diff --git a/vendor/parquet/src/column/page.rs b/vendor/parquet/src/column/page.rs new file mode 100644 index 0000000..ed80e27 --- /dev/null +++ b/vendor/parquet/src/column/page.rs @@ -0,0 +1,599 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Contains Parquet Page definitions and page reader interface. + +use bytes::Bytes; + +use crate::basic::{Encoding, PageType}; +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::thrift::{ + DataPageHeader, DataPageHeaderV2, DictionaryPageHeader, PageHeader, +}; +use crate::file::statistics::{Statistics, page_stats_to_thrift}; + +/// Parquet Page definition. +/// +/// List of supported pages. +/// These are 1-to-1 mapped from the equivalent Thrift definitions, except `buf` which +/// used to store uncompressed bytes of the page. +#[derive(Clone, Debug)] +pub enum Page { + /// Data page Parquet format v1. + DataPage { + /// The underlying data buffer + buf: Bytes, + /// Number of values in this page + num_values: u32, + /// Encoding for values in this page + encoding: Encoding, + /// Definition level encoding + def_level_encoding: Encoding, + /// Repetition level encoding + rep_level_encoding: Encoding, + /// Optional statistics for this page + statistics: Option, + }, + /// Data page Parquet format v2. + DataPageV2 { + /// The underlying data buffer + buf: Bytes, + /// Number of values in this page + num_values: u32, + /// Encoding for values in this page + encoding: Encoding, + /// Number of null values in this page + num_nulls: u32, + /// Number of rows in this page + num_rows: u32, + /// Length of definition levels + def_levels_byte_len: u32, + /// Length of repetition levels + rep_levels_byte_len: u32, + /// Is this page compressed + is_compressed: bool, + /// Optional statistics for this page + statistics: Option, + }, + /// Dictionary page. + DictionaryPage { + /// The underlying data buffer + buf: Bytes, + /// Number of values in this page + num_values: u32, + /// Encoding for values in this page + encoding: Encoding, + /// Is dictionary page sorted + is_sorted: bool, + }, +} + +impl Page { + /// Returns [`PageType`] for this page. + pub fn page_type(&self) -> PageType { + match self { + Page::DataPage { .. } => PageType::DATA_PAGE, + Page::DataPageV2 { .. } => PageType::DATA_PAGE_V2, + Page::DictionaryPage { .. } => PageType::DICTIONARY_PAGE, + } + } + + /// Returns whether this page is any version of a data page + pub fn is_data_page(&self) -> bool { + matches!(self, Page::DataPage { .. } | Page::DataPageV2 { .. }) + } + + /// Returns whether this page is a dictionary page + pub fn is_dictionary_page(&self) -> bool { + matches!(self, Page::DictionaryPage { .. }) + } + + /// Returns internal byte buffer reference for this page. + pub fn buffer(&self) -> &Bytes { + match self { + Page::DataPage { buf, .. } => buf, + Page::DataPageV2 { buf, .. } => buf, + Page::DictionaryPage { buf, .. } => buf, + } + } + + /// Returns number of values in this page. + pub fn num_values(&self) -> u32 { + match self { + Page::DataPage { num_values, .. } => *num_values, + Page::DataPageV2 { num_values, .. } => *num_values, + Page::DictionaryPage { num_values, .. } => *num_values, + } + } + + /// Returns this page [`Encoding`]. + pub fn encoding(&self) -> Encoding { + match self { + Page::DataPage { encoding, .. } => *encoding, + Page::DataPageV2 { encoding, .. } => *encoding, + Page::DictionaryPage { encoding, .. } => *encoding, + } + } + + /// Returns optional [`Statistics`]. + pub fn statistics(&self) -> Option<&Statistics> { + match self { + Page::DataPage { statistics, .. } => statistics.as_ref(), + Page::DataPageV2 { statistics, .. } => statistics.as_ref(), + Page::DictionaryPage { .. } => None, + } + } +} + +/// Helper struct to represent pages with potentially compressed buffer (data page v1) or +/// compressed and concatenated buffer (def levels + rep levels + compressed values for +/// data page v2). +/// +/// The difference with `Page` is that `Page` buffer is always uncompressed. +pub struct CompressedPage { + compressed_page: Page, + uncompressed_size: usize, +} + +impl CompressedPage { + /// Creates `CompressedPage` from a page with potentially compressed buffer and + /// uncompressed size. + pub fn new(compressed_page: Page, uncompressed_size: usize) -> Self { + Self { + compressed_page, + uncompressed_size, + } + } + + /// Returns page type. + pub fn page_type(&self) -> PageType { + self.compressed_page.page_type() + } + + /// Returns underlying page with potentially compressed buffer. + pub fn compressed_page(&self) -> &Page { + &self.compressed_page + } + + /// Returns uncompressed size in bytes. + pub fn uncompressed_size(&self) -> usize { + self.uncompressed_size + } + + /// Returns compressed size in bytes. + /// + /// Note that it is assumed that buffer is compressed, but it may not be. In this + /// case compressed size will be equal to uncompressed size. + pub fn compressed_size(&self) -> usize { + self.compressed_page.buffer().len() + } + + /// Number of values in page. + pub fn num_values(&self) -> u32 { + self.compressed_page.num_values() + } + + /// Returns encoding for values in page. + pub fn encoding(&self) -> Encoding { + self.compressed_page.encoding() + } + + /// Returns slice of compressed buffer in the page. + pub fn data(&self) -> &[u8] { + self.compressed_page.buffer() + } + + /// Returns the number of heap bytes this page currently holds. + /// + /// This is the page's compressed buffer (the embedded [`Bytes`]); use it to + /// account for a buffered page's memory footprint rather than reaching for + /// `data().len()` at each call site. + pub fn memory_usage(&self) -> usize { + self.compressed_page.buffer().len() + } + + /// Returns the thrift page header + pub(crate) fn to_thrift_header(&self) -> Result { + let uncompressed_size = self.uncompressed_size(); + let compressed_size = self.compressed_size(); + if uncompressed_size > i32::MAX as usize { + return Err(general_err!( + "Page uncompressed size overflow: {}", + uncompressed_size + )); + } + if compressed_size > i32::MAX as usize { + return Err(general_err!( + "Page compressed size overflow: {}", + compressed_size + )); + } + let num_values = self.num_values(); + let encoding = self.encoding(); + let page_type = self.page_type(); + + let mut page_header = PageHeader { + r#type: page_type, + uncompressed_page_size: uncompressed_size as i32, + compressed_page_size: compressed_size as i32, + // TODO: Add support for crc checksum + crc: None, + data_page_header: None, + index_page_header: None, + dictionary_page_header: None, + data_page_header_v2: None, + }; + + match self.compressed_page { + Page::DataPage { + def_level_encoding, + rep_level_encoding, + ref statistics, + .. + } => { + let data_page_header = DataPageHeader { + num_values: num_values as i32, + encoding, + definition_level_encoding: def_level_encoding, + repetition_level_encoding: rep_level_encoding, + statistics: page_stats_to_thrift(statistics.as_ref()), + }; + page_header.data_page_header = Some(data_page_header); + } + Page::DataPageV2 { + num_nulls, + num_rows, + def_levels_byte_len, + rep_levels_byte_len, + is_compressed, + ref statistics, + .. + } => { + let data_page_header_v2 = DataPageHeaderV2 { + num_values: num_values as i32, + num_nulls: num_nulls as i32, + num_rows: num_rows as i32, + encoding, + definition_levels_byte_length: def_levels_byte_len as i32, + repetition_levels_byte_length: rep_levels_byte_len as i32, + is_compressed: Some(is_compressed), + statistics: page_stats_to_thrift(statistics.as_ref()), + }; + page_header.data_page_header_v2 = Some(data_page_header_v2); + } + Page::DictionaryPage { is_sorted, .. } => { + let dictionary_page_header = DictionaryPageHeader { + num_values: num_values as i32, + encoding, + is_sorted: Some(is_sorted), + }; + page_header.dictionary_page_header = Some(dictionary_page_header); + } + } + Ok(page_header) + } + + /// Update the compressed buffer for a page. + /// This might be required when encrypting page data for example. + /// The size of uncompressed data must not change. + #[cfg(feature = "encryption")] + pub(crate) fn with_new_compressed_buffer(mut self, new_buffer: Bytes) -> Self { + match &mut self.compressed_page { + Page::DataPage { buf, .. } => { + *buf = new_buffer; + } + Page::DataPageV2 { buf, .. } => { + *buf = new_buffer; + } + Page::DictionaryPage { buf, .. } => { + *buf = new_buffer; + } + } + self + } +} + +/// Contains page write metrics. +pub struct PageWriteSpec { + /// The type of page being written + pub page_type: PageType, + /// The total size of the page, before compression + pub uncompressed_size: usize, + /// The compressed size of the page + pub compressed_size: usize, + /// The number of values in the page + pub num_values: u32, + /// The offset of the page in the column chunk + pub offset: u64, + /// The number of bytes written to the underlying sink + pub bytes_written: u64, +} + +impl Default for PageWriteSpec { + fn default() -> Self { + Self::new() + } +} + +impl PageWriteSpec { + /// Creates new spec with default page write metrics. + pub fn new() -> Self { + Self { + page_type: PageType::DATA_PAGE, + uncompressed_size: 0, + compressed_size: 0, + num_values: 0, + offset: 0, + bytes_written: 0, + } + } +} + +/// Contains metadata for a page +#[derive(Clone)] +pub struct PageMetadata { + /// The number of rows within the page if known + pub num_rows: Option, + /// The number of levels within the page if known + pub num_levels: Option, + /// Returns true if the page is a dictionary page + pub is_dict: bool, +} + +impl TryFrom<&crate::file::metadata::thrift::PageHeader> for PageMetadata { + type Error = ParquetError; + + fn try_from( + value: &crate::file::metadata::thrift::PageHeader, + ) -> std::result::Result { + match value.r#type { + PageType::DATA_PAGE => { + let header = value.data_page_header.as_ref().unwrap(); + Ok(PageMetadata { + num_rows: None, + num_levels: Some(header.num_values as _), + is_dict: false, + }) + } + PageType::DICTIONARY_PAGE => Ok(PageMetadata { + num_rows: None, + num_levels: None, + is_dict: true, + }), + PageType::DATA_PAGE_V2 => { + let header = value.data_page_header_v2.as_ref().unwrap(); + Ok(PageMetadata { + num_rows: Some(header.num_rows as _), + num_levels: Some(header.num_values as _), + is_dict: false, + }) + } + other => Err(ParquetError::General(format!( + "page type {other:?} cannot be converted to PageMetadata" + ))), + } + } +} + +/// API for reading pages from a column chunk. +/// This offers a iterator like API to get the next page. +pub trait PageReader: Iterator> + Send { + /// Gets the next page in the column chunk associated with this reader. + /// Returns `None` if there are no pages left. + fn get_next_page(&mut self) -> Result>; + + /// Gets metadata about the next page, returns an error if no + /// column index information + fn peek_next_page(&mut self) -> Result>; + + /// Skips reading the next page, returns an error if no + /// column index information + fn skip_next_page(&mut self) -> Result<()>; + + /// Returns `true` if the next page can be assumed to contain the start of a new record + /// + /// Prior to parquet V2 the specification was ambiguous as to whether a single record + /// could be split across multiple pages, and prior to [(#4327)] the Rust writer would do + /// this in certain situations. However, correctly interpreting the offset index relies on + /// this assumption holding [(#4943)], and so this mechanism is provided for a [`PageReader`] + /// to signal this to the calling context + /// + /// [(#4327)]: https://github.com/apache/arrow-rs/pull/4327 + /// [(#4943)]: https://github.com/apache/arrow-rs/pull/4943 + fn at_record_boundary(&mut self) -> Result { + match self.peek_next_page()? { + // Last page in the column chunk - always a record boundary + None => Ok(true), + // A V2 data page is required by the parquet spec to start at a + // record boundary, so the current page ends at one. V2 pages + // are identified by having `num_rows` set in their header. + Some(metadata) => Ok(metadata.num_rows.is_some()), + } + } +} + +/// API for writing pages in a column chunk. +/// +/// It is reasonable to assume that all pages will be written in the correct order, e.g. +/// dictionary page followed by data pages, or a set of data pages, etc. +pub trait PageWriter: Send { + /// Writes a page into the output stream/sink. + /// Returns `PageWriteSpec` that contains information about written page metrics, + /// including number of bytes, size, number of values, offset, etc. + /// + /// This method is called for every compressed page we write into underlying buffer, + /// either data page or dictionary page. + fn write_page(&mut self, page: CompressedPage) -> Result; + + /// **Unstable, not public API.** This is an internal protocol between + /// [`GenericColumnWriter`] and its in-crate page writers; it is hidden from + /// the rendered docs and may change or be removed without a major version + /// bump. External `PageWriter` implementations should not override it. See + /// the page-spilling cleanup tracked in + /// . + /// + /// Whether this writer resolves the final page layout itself (at flush) + /// rather than committing bytes to their final position as pages arrive. + /// + /// The dictionary page of a column chunk must be written *first*, but it is + /// not finalized until every value has been seen. A writer that commits + /// bytes live (e.g. straight to a file) therefore relies on the column + /// writer buffering the dictionary-encoded data pages in memory until the + /// dictionary page is ready — see [`GenericColumnWriter`]'s `data_pages`. + /// + /// A writer that instead buffers the whole chunk and splices it later (the + /// [`ArrowWriter`] path) can accept data pages *before* the dictionary page + /// and order them itself at flush. Returning `true` tells the column writer + /// to skip that in-memory buffering and stream dictionary-column data pages + /// straight through, bounding the column writer's memory. + /// + /// [`GenericColumnWriter`]: crate::column::writer::GenericColumnWriter + /// [`ArrowWriter`]: crate::arrow::arrow_writer::ArrowWriter + #[doc(hidden)] + fn defers_dictionary_ordering(&self) -> bool { + false + } + + /// **Unstable, not public API.** Companion to + /// [`defers_dictionary_ordering`](Self::defers_dictionary_ordering): an + /// internal hook for the column writer's memory accounting, hidden from the + /// rendered docs and subject to change or removal without a major version + /// bump. External `PageWriter` implementations should not override it. + /// + /// The number of bytes this writer is currently holding **in memory** for + /// pages it has been handed (i.e. completed pages not yet committed to their + /// final destination). + /// + /// Used by the column writer to report its memory footprint. The default is + /// `0`: a writer that streams pages straight to their destination retains + /// nothing. A writer that buffers pages should report what it actually holds + /// on the heap — which, when it spills to a backing store, can be far less + /// than the bytes written. + #[doc(hidden)] + fn buffered_memory_size(&self) -> usize { + 0 + } + + /// Closes resources and flushes underlying sink. + /// Page writer should not be used after this method is called. + fn close(&mut self) -> Result<()>; +} + +/// An iterator over pages of one specific column in a parquet file. +pub trait PageIterator: Iterator>> + Send {} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_page() { + let data_page = Page::DataPage { + buf: Bytes::from(vec![0, 1, 2]), + num_values: 10, + encoding: Encoding::PLAIN, + def_level_encoding: Encoding::RLE, + rep_level_encoding: Encoding::RLE, + statistics: Some(Statistics::int32(Some(1), Some(2), None, Some(1), true)), + }; + assert_eq!(data_page.page_type(), PageType::DATA_PAGE); + assert_eq!(data_page.buffer(), vec![0, 1, 2].as_slice()); + assert_eq!(data_page.num_values(), 10); + assert_eq!(data_page.encoding(), Encoding::PLAIN); + assert_eq!( + data_page.statistics(), + Some(&Statistics::int32(Some(1), Some(2), None, Some(1), true)) + ); + + let data_page_v2 = Page::DataPageV2 { + buf: Bytes::from(vec![0, 1, 2]), + num_values: 10, + encoding: Encoding::PLAIN, + num_nulls: 5, + num_rows: 20, + def_levels_byte_len: 30, + rep_levels_byte_len: 40, + is_compressed: false, + statistics: Some(Statistics::int32(Some(1), Some(2), None, Some(1), true)), + }; + assert_eq!(data_page_v2.page_type(), PageType::DATA_PAGE_V2); + assert_eq!(data_page_v2.buffer(), vec![0, 1, 2].as_slice()); + assert_eq!(data_page_v2.num_values(), 10); + assert_eq!(data_page_v2.encoding(), Encoding::PLAIN); + assert_eq!( + data_page_v2.statistics(), + Some(&Statistics::int32(Some(1), Some(2), None, Some(1), true)) + ); + + let dict_page = Page::DictionaryPage { + buf: Bytes::from(vec![0, 1, 2]), + num_values: 10, + encoding: Encoding::PLAIN, + is_sorted: false, + }; + assert_eq!(dict_page.page_type(), PageType::DICTIONARY_PAGE); + assert_eq!(dict_page.buffer(), vec![0, 1, 2].as_slice()); + assert_eq!(dict_page.num_values(), 10); + assert_eq!(dict_page.encoding(), Encoding::PLAIN); + assert_eq!(dict_page.statistics(), None); + } + + #[test] + fn test_compressed_page() { + let data_page = Page::DataPage { + buf: Bytes::from(vec![0, 1, 2]), + num_values: 10, + encoding: Encoding::PLAIN, + def_level_encoding: Encoding::RLE, + rep_level_encoding: Encoding::RLE, + statistics: Some(Statistics::int32(Some(1), Some(2), None, Some(1), true)), + }; + + let cpage = CompressedPage::new(data_page, 5); + + assert_eq!(cpage.page_type(), PageType::DATA_PAGE); + assert_eq!(cpage.uncompressed_size(), 5); + assert_eq!(cpage.compressed_size(), 3); + assert_eq!(cpage.num_values(), 10); + assert_eq!(cpage.encoding(), Encoding::PLAIN); + assert_eq!(cpage.data(), &[0, 1, 2]); + } + + #[test] + fn test_compressed_page_uncompressed_size_overflow() { + // Test that to_thrift_header fails when uncompressed size exceeds i32::MAX + let data_page = Page::DataPage { + buf: Bytes::from(vec![0, 1, 2]), + num_values: 10, + encoding: Encoding::PLAIN, + def_level_encoding: Encoding::RLE, + rep_level_encoding: Encoding::RLE, + statistics: None, + }; + + // Create a CompressedPage with uncompressed size larger than i32::MAX + let uncompressed_size = (i32::MAX as usize) + 1; + let cpage = CompressedPage::new(data_page, uncompressed_size); + + // Verify that to_thrift_header returns an error + let result = cpage.to_thrift_header(); + assert!(result.is_err()); + + let error_msg = result.unwrap_err().to_string(); + assert!(error_msg.contains("Page uncompressed size overflow")); + } +} diff --git a/vendor/parquet/src/column/page_encryption.rs b/vendor/parquet/src/column/page_encryption.rs new file mode 100644 index 0000000..26df759 --- /dev/null +++ b/vendor/parquet/src/column/page_encryption.rs @@ -0,0 +1,119 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::basic::PageType; +use crate::column::page::CompressedPage; +use crate::encryption::ciphers::BlockEncryptor; +use crate::encryption::encrypt::{FileEncryptor, encrypt_thrift_object}; +use crate::encryption::modules::{ModuleType, create_module_aad}; +use crate::errors::ParquetError; +use crate::errors::Result; +use crate::file::metadata::thrift::PageHeader; +use bytes::Bytes; +use std::io::Write; +use std::sync::Arc; + +#[derive(Debug)] +/// Encrypts page headers and page data for columns +pub(crate) struct PageEncryptor { + file_encryptor: Arc, + block_encryptor: Box, + row_group_index: usize, + column_index: usize, + page_index: usize, +} + +impl PageEncryptor { + /// Create a [`PageEncryptor`] for a column if it should be encrypted + pub fn create_if_column_encrypted( + file_encryptor: &Option>, + row_group_index: usize, + column_index: usize, + column_path: &str, + ) -> Result> { + match file_encryptor { + Some(file_encryptor) if file_encryptor.is_column_encrypted(column_path) => { + let block_encryptor = file_encryptor.get_column_encryptor(column_path)?; + Ok(Some(Self { + file_encryptor: file_encryptor.clone(), + block_encryptor, + row_group_index, + column_index, + page_index: 0, + })) + } + _ => Ok(None), + } + } + + /// Update the page index after a data page has been processed + pub fn increment_page(&mut self) { + self.page_index += 1; + } + + fn encrypt_page(&mut self, page: &CompressedPage) -> Result> { + let module_type = if page.compressed_page().is_data_page() { + ModuleType::DataPage + } else { + ModuleType::DictionaryPage + }; + let aad = create_module_aad( + self.file_encryptor.file_aad(), + module_type, + self.row_group_index, + self.column_index, + Some(self.page_index), + )?; + let encrypted_buffer = self.block_encryptor.encrypt(page.data(), &aad)?; + + Ok(encrypted_buffer) + } + + /// Encrypt compressed column page data + pub fn encrypt_compressed_page(&mut self, page: CompressedPage) -> Result { + let encrypted_page = self.encrypt_page(&page)?; + Ok(page.with_new_compressed_buffer(Bytes::from(encrypted_page))) + } + + /// Encrypt a column page header + pub fn encrypt_page_header( + &mut self, + page_header: &PageHeader, + sink: &mut W, + ) -> Result<()> { + let module_type = match page_header.r#type { + PageType::DATA_PAGE => ModuleType::DataPageHeader, + PageType::DATA_PAGE_V2 => ModuleType::DataPageHeader, + PageType::DICTIONARY_PAGE => ModuleType::DictionaryPageHeader, + _ => { + return Err(general_err!( + "Unsupported page type for page header encryption: {:?}", + page_header.r#type + )); + } + }; + let aad = create_module_aad( + self.file_encryptor.file_aad(), + module_type, + self.row_group_index, + self.column_index, + Some(self.page_index), + )?; + + encrypt_thrift_object(page_header, &mut self.block_encryptor, sink, &aad) + } +} diff --git a/vendor/parquet/src/column/page_encryption_disabled.rs b/vendor/parquet/src/column/page_encryption_disabled.rs new file mode 100644 index 0000000..71f2586 --- /dev/null +++ b/vendor/parquet/src/column/page_encryption_disabled.rs @@ -0,0 +1,44 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::column::page::CompressedPage; +use crate::errors::Result; +use crate::file::metadata::thrift::PageHeader; +use std::io::Write; + +#[derive(Debug)] +/// Dummy PageEncryptor struct that can never be instantiated, +/// provided to support compilation without the encryption feature enabled. +pub(crate) struct PageEncryptor { + _empty: (), +} + +impl PageEncryptor { + pub fn increment_page(&mut self) {} + + pub fn encrypt_compressed_page(&mut self, _page: CompressedPage) -> Result { + unreachable!("The encryption feature is disabled") + } + + pub fn encrypt_page_header( + &mut self, + _page_header: &PageHeader, + _sink: &mut W, + ) -> Result<()> { + unreachable!("The encryption feature is disabled") + } +} diff --git a/vendor/parquet/src/column/page_store.rs b/vendor/parquet/src/column/page_store.rs new file mode 100644 index 0000000..14bb2b9 --- /dev/null +++ b/vendor/parquet/src/column/page_store.rs @@ -0,0 +1,260 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Pluggable storage for completed, serialized page blobs. +//! +//! While a row group is being written the [`ArrowWriter`] must buffer every +//! column's encoded pages, because Parquet requires each column chunk to be +//! contiguous in the file while record batches arrive with all columns interleaved. +//! By default that buffer lives on the heap, so the writer's peak memory grows +//! with the row group size. A [`PageStore`] lets the buffer live somewhere else +//! — a local temp file, object storage, etc. — bounding peak write memory +//! independently of the row group size. +//! +//! [`ArrowWriter`]: crate::arrow::arrow_writer::ArrowWriter + +use std::fmt::Debug; + +use bytes::Bytes; + +use crate::errors::{ParquetError, Result}; +use crate::schema::types::ColumnDescriptor; + +/// An opaque, store-allocated handle to a blob held by a [`PageStore`]. +/// +/// Handles are allocated by the store — densely and sequentially — and are only +/// meaningful to the store that produced them. The caller treats them as opaque +/// tokens. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct PageKey(u64); + +impl PageKey { + /// Create a handle wrapping `raw`. + /// + /// A [`PageStore`] implementation calls this to mint the handle it returns + /// from [`put`](PageStore::put). The value is opaque to the caller, so a + /// store is free to use a dense counter, a packed locator, or anything else + /// it can later resolve in [`take`](PageStore::take). + pub const fn new(raw: u64) -> Self { + Self(raw) + } + + /// The raw value passed to [`new`](Self::new). + pub const fn get(self) -> u64 { + self.0 + } +} + +/// A pluggable store for completed, serialized page blobs. +/// +/// The store is intentionally "dumb": it only maps an opaque [`PageKey`] to a +/// blob of bytes. It knows nothing about pages, dictionaries, ordering, or +/// offsets. The caller keeps the handles it gets back from [`put`](Self::put) +/// and decides what they mean. +/// +/// Each store instance is owned by a single column writer and mutated by one +/// thread at a time (both methods take `&mut self`), so it needs no internal +/// synchronization — hence only `Send`, not `Sync`. +/// +/// The default ([`InMemoryPageStore`]) keeps blobs in memory on the heap. +/// +/// For an example of configuring the Parquet writer to use an alternate +/// `PageStore` see the [`ArrowWriterOptions::with_page_store_factory`] API. +/// +/// [`ArrowWriterOptions::with_page_store_factory`]: crate::arrow::arrow_writer::ArrowWriterOptions::with_page_store_factory +pub trait PageStore: Send { + /// Store `value`, returning a handle that can later be passed to + /// [`take`](Self::take). + fn put(&mut self, value: Bytes) -> Result; + + /// Take back the blob previously stored under `key`. + /// + /// The caller takes ownership of the returned bytes and will **not** request + /// `key` again, so the store may release any resources backing it — eagerly + /// here, or when the store is dropped. + fn take(&mut self, key: PageKey) -> Result; + + /// The number of bytes this store currently holds **in memory** (resident + /// on the heap), used to report the writer's memory footprint. + /// + /// The default is `0`, which is exactly right for a backend that moves + /// every blob off-heap (a temp file, object storage): the bytes it has been + /// handed no longer occupy heap. The in-memory backend overrides this to + /// report its resident blobs. A backend that keeps a partial in-memory + /// buffer should report that buffer's size. + fn memory_size(&self) -> usize { + 0 + } +} + +/// Context for a single [`PageStoreFactory::create`] call. +/// +/// Describes the leaf column chunk the store will buffer. It is held by +/// reference for the duration of the call; a backend reads only what it needs. +/// More fields may be added in future releases without breaking existing +/// implementations — the type is constructed only by the writer, so an +/// implementer only ever receives one and calls its accessors. +pub struct PageStoreArgs<'a> { + column_index: usize, + column_descriptor: &'a ColumnDescriptor, +} + +impl<'a> PageStoreArgs<'a> { + // Constructed only by the Arrow writer; without that feature there is no caller. + #[cfg(feature = "arrow")] + pub(crate) fn new(column_index: usize, column_descriptor: &'a ColumnDescriptor) -> Self { + Self { + column_index, + column_descriptor, + } + } + + /// Index of the leaf column within the row group. + /// + /// A backend may use this to e.g. name spill files or shard across a bounded + /// pool; it carries no ordering or coordination requirement. + pub fn column_index(&self) -> usize { + self.column_index + } + + /// Descriptor for the leaf column: physical/logical type, path, and max + /// definition/repetition levels. + /// + /// Lets a backend tailor buffering to the column — for example spilling only + /// large `BYTE_ARRAY` columns while keeping small fixed-width ones on the + /// heap. + pub fn column_descriptor(&self) -> &ColumnDescriptor { + self.column_descriptor + } +} + +/// Creates a fresh [`PageStore`] for each column chunk. +/// +/// See +/// [`ArrowWriterOptions::with_page_store_factory`](crate::arrow::arrow_writer::ArrowWriterOptions::with_page_store_factory). +pub trait PageStoreFactory: Send + Sync + Debug { + /// Create a new, empty [`PageStore`] for the leaf column described by `args`. + fn create(&self, args: &PageStoreArgs<'_>) -> Result>; +} + +/// The default [`PageStore`], holding blobs on the heap in a `Vec`. +/// +/// Peak memory grows with the row group size; use a spilling backend to bound +/// it. +#[derive(Debug, Default)] +pub struct InMemoryPageStore { + blobs: Vec, + /// Running total of resident blob bytes, kept in step with `put`/`take`. + resident: usize, +} + +impl PageStore for InMemoryPageStore { + fn put(&mut self, value: Bytes) -> Result { + let key = PageKey(self.blobs.len() as u64); + self.resident += value.len(); + self.blobs.push(value); + Ok(key) + } + + fn take(&mut self, key: PageKey) -> Result { + // Replace the slot with an empty `Bytes` so the stored blob is released + // as soon as it is taken, keeping memory bounded while the chunk is + // streamed into the output file. + let blob = self + .blobs + .get_mut(key.0 as usize) + .map(std::mem::take) + .ok_or_else(|| ParquetError::General(format!("invalid page key {}", key.0)))?; + self.resident -= blob.len(); + Ok(blob) + } + + fn memory_size(&self) -> usize { + self.resident + } +} + +/// Factory for [`InMemoryPageStore`] — the default used by +/// [`ArrowWriter`](crate::arrow::arrow_writer::ArrowWriter). +#[derive(Debug, Default)] +pub struct InMemoryPageStoreFactory; + +impl PageStoreFactory for InMemoryPageStoreFactory { + fn create(&self, _args: &PageStoreArgs<'_>) -> Result> { + Ok(Box::new(InMemoryPageStore::default())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn in_memory_round_trips_blobs_in_handle_order() { + let mut store = InMemoryPageStore::default(); + let k0 = store.put(Bytes::from_static(b"hello")).unwrap(); + let k1 = store.put(Bytes::from_static(b"world")).unwrap(); + assert_ne!(k0, k1); + assert_eq!(&store.take(k0).unwrap()[..], b"hello"); + assert_eq!(&store.take(k1).unwrap()[..], b"world"); + } + + #[test] + fn in_memory_take_releases_the_slot() { + let mut store = InMemoryPageStore::default(); + let k = store.put(Bytes::from_static(b"abc")).unwrap(); + assert_eq!(&store.take(k).unwrap()[..], b"abc"); + // A second take yields the emptied placeholder rather than the blob, + // confirming the bytes were released on the first take. + assert!(store.take(k).unwrap().is_empty()); + } + + #[test] + fn in_memory_invalid_key_errors() { + let mut store = InMemoryPageStore::default(); + assert!(store.take(PageKey(99)).is_err()); + } + + #[test] + fn in_memory_reports_resident_bytes() { + let mut store = InMemoryPageStore::default(); + assert_eq!(store.memory_size(), 0); + let k0 = store.put(Bytes::from_static(b"hello")).unwrap(); + let k1 = store.put(Bytes::from_static(b"!")).unwrap(); + assert_eq!(store.memory_size(), 6); + store.take(k0).unwrap(); + assert_eq!(store.memory_size(), 1); + store.take(k1).unwrap(); + assert_eq!(store.memory_size(), 0); + } + + #[test] + fn default_store_memory_size_is_zero() { + // A spilling backend that does not override `memory_size` reports 0, + // reflecting that its blobs no longer occupy the heap. + struct OffHeap; + impl PageStore for OffHeap { + fn put(&mut self, _value: Bytes) -> Result { + Ok(PageKey::new(0)) + } + fn take(&mut self, _key: PageKey) -> Result { + Ok(Bytes::new()) + } + } + assert_eq!(OffHeap.memory_size(), 0); + } +} diff --git a/vendor/parquet/src/column/reader.rs b/vendor/parquet/src/column/reader.rs new file mode 100644 index 0000000..81960ab --- /dev/null +++ b/vendor/parquet/src/column/reader.rs @@ -0,0 +1,1520 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Contains column reader API. + +use bytes::Bytes; + +use super::page::{Page, PageReader}; +use crate::basic::*; +use crate::column::reader::decoder::{ + ColumnValueDecoder, ColumnValueDecoderImpl, DefinitionLevelDecoder, DefinitionLevelDecoderImpl, + RepetitionLevelDecoder, RepetitionLevelDecoderImpl, +}; +use crate::data_type::*; +use crate::errors::{ParquetError, Result}; +use crate::schema::types::ColumnDescPtr; +use crate::util::bit_util::{ceil, num_required_bits, read_num_bytes}; + +pub(crate) mod decoder; + +/// Column reader for a Parquet type. +pub enum ColumnReader { + /// Column reader for boolean type + BoolColumnReader(ColumnReaderImpl), + /// Column reader for int32 type + Int32ColumnReader(ColumnReaderImpl), + /// Column reader for int64 type + Int64ColumnReader(ColumnReaderImpl), + /// Column reader for int96 type + Int96ColumnReader(ColumnReaderImpl), + /// Column reader for float type + FloatColumnReader(ColumnReaderImpl), + /// Column reader for double type + DoubleColumnReader(ColumnReaderImpl), + /// Column reader for byte array type + ByteArrayColumnReader(ColumnReaderImpl), + /// Column reader for fixed length byte array type + FixedLenByteArrayColumnReader(ColumnReaderImpl), +} + +/// Gets a specific column reader corresponding to column descriptor `col_descr`. The +/// column reader will read from pages in `col_page_reader`. +pub fn get_column_reader( + col_descr: ColumnDescPtr, + col_page_reader: Box, +) -> ColumnReader { + match col_descr.physical_type() { + Type::BOOLEAN => { + ColumnReader::BoolColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader)) + } + Type::INT32 => { + ColumnReader::Int32ColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader)) + } + Type::INT64 => { + ColumnReader::Int64ColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader)) + } + Type::INT96 => { + ColumnReader::Int96ColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader)) + } + Type::FLOAT => { + ColumnReader::FloatColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader)) + } + Type::DOUBLE => { + ColumnReader::DoubleColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader)) + } + Type::BYTE_ARRAY => { + ColumnReader::ByteArrayColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader)) + } + Type::FIXED_LEN_BYTE_ARRAY => ColumnReader::FixedLenByteArrayColumnReader( + ColumnReaderImpl::new(col_descr, col_page_reader), + ), + } +} + +/// Gets a typed column reader for the specific type `T`, by "up-casting" `col_reader` of +/// non-generic type to a generic column reader type `ColumnReaderImpl`. +/// +/// Panics if actual enum value for `col_reader` does not match the type `T`. +pub fn get_typed_column_reader(col_reader: ColumnReader) -> ColumnReaderImpl { + T::get_column_reader(col_reader).unwrap_or_else(|| { + panic!( + "Failed to convert column reader into a typed column reader for `{}` type", + T::get_physical_type() + ) + }) +} + +/// Typed value reader for a particular primitive column. +pub type ColumnReaderImpl = GenericColumnReader< + RepetitionLevelDecoderImpl, + DefinitionLevelDecoderImpl, + ColumnValueDecoderImpl, +>; + +/// Reads data for a given column chunk, using the provided decoders: +/// +/// - R: `ColumnLevelDecoder` used to decode repetition levels +/// - D: `ColumnLevelDecoder` used to decode definition levels +/// - V: `ColumnValueDecoder` used to decode value data +pub struct GenericColumnReader { + descr: ColumnDescPtr, + + page_reader: Box, + + /// The total number of values stored in the data page. + num_buffered_values: usize, + + /// The number of values from the current data page that has been decoded into memory + /// so far. + num_decoded_values: usize, + + /// True if the end of the current data page denotes the end of a record + has_record_delimiter: bool, + + /// The decoder for the definition levels if any + def_level_decoder: Option, + + /// The decoder for the repetition levels if any + rep_level_decoder: Option, + + /// The decoder for the values + values_decoder: V, +} + +impl GenericColumnReader +where + V: ColumnValueDecoder, +{ + /// Creates new column reader based on column descriptor and page reader. + pub fn new(descr: ColumnDescPtr, page_reader: Box) -> Self { + let values_decoder = V::new(&descr); + + let def_level_decoder = (descr.max_def_level() != 0) + .then(|| DefinitionLevelDecoderImpl::new(descr.max_def_level())); + + let rep_level_decoder = (descr.max_rep_level() != 0) + .then(|| RepetitionLevelDecoderImpl::new(descr.max_rep_level())); + + Self::new_with_decoders( + descr, + page_reader, + values_decoder, + def_level_decoder, + rep_level_decoder, + ) + } +} + +impl GenericColumnReader +where + R: RepetitionLevelDecoder, + D: DefinitionLevelDecoder, + V: ColumnValueDecoder, +{ + pub(crate) fn new_with_decoders( + descr: ColumnDescPtr, + page_reader: Box, + values_decoder: V, + def_level_decoder: Option, + rep_level_decoder: Option, + ) -> Self { + Self { + descr, + def_level_decoder, + rep_level_decoder, + page_reader, + num_buffered_values: 0, + num_decoded_values: 0, + values_decoder, + has_record_delimiter: false, + } + } + + /// Read up to `max_records` whole records, returning the number of complete + /// records, non-null values and levels decoded. All levels for a given record + /// will be read, i.e. the next repetition level, if any, will be 0 + /// + /// If the max definition level is 0, `def_levels` will be ignored and the number of records, + /// non-null values and levels decoded will all be equal, otherwise `def_levels` will be + /// populated with the number of levels read, with an error returned if it is `None`. + /// + /// If the max repetition level is 0, `rep_levels` will be ignored and the number of records + /// and levels decoded will both be equal, otherwise `rep_levels` will be populated with + /// the number of levels read, with an error returned if it is `None`. + /// + /// `values` will be contiguously populated with the non-null values. Note that if the column + /// is not required, this may be less than either `max_records` or the number of levels read + pub fn read_records( + &mut self, + max_records: usize, + def_levels: Option<&mut D::Buffer>, + rep_levels: Option<&mut R::Buffer>, + values: &mut V::Buffer, + ) -> Result<(usize, usize, usize)> { + self.read_records_with_reservation( + max_records, + def_levels, + rep_levels, + values, + |_, _, _, _| Ok(()), + ) + } + + /// Like [`Self::read_records`], but invokes `reserve_values` after levels + /// are decoded and before values are read. + pub(crate) fn read_records_with_reservation( + &mut self, + max_records: usize, + mut def_levels: Option<&mut D::Buffer>, + mut rep_levels: Option<&mut R::Buffer>, + values: &mut V::Buffer, + mut reserve_values: F, + ) -> Result<(usize, usize, usize)> + where + F: FnMut(&mut V::Buffer, usize, usize, Option<&D::Buffer>) -> Result<()>, + { + let mut total_records_read = 0; + let mut total_levels_read = 0; + let mut total_values_read = 0; + + while total_records_read < max_records && self.has_next()? { + let remaining_records = max_records - total_records_read; + let remaining_levels = self.num_buffered_values - self.num_decoded_values; + + let (records_read, levels_to_read) = match self.rep_level_decoder.as_mut() { + Some(reader) => { + let out = rep_levels + .as_mut() + .ok_or_else(|| general_err!("must specify repetition levels"))?; + + let (mut records_read, levels_read) = + reader.read_rep_levels(out, remaining_records, remaining_levels)?; + + if records_read == 0 && levels_read == 0 { + // The fact that we're still looping implies there must be some levels to read. + return Err(general_err!( + "Insufficient repetition levels read from column" + )); + } + if levels_read == remaining_levels && self.has_record_delimiter { + // Reached end of page, which implies records_read < remaining_records + // as otherwise would have stopped reading before reaching the end + assert!(records_read < remaining_records); // Sanity check + records_read += reader.flush_partial() as usize; + } + (records_read, levels_read) + } + None => { + let min = remaining_records.min(remaining_levels); + (min, min) + } + }; + + let values_to_read = match self.def_level_decoder.as_mut() { + Some(reader) => { + let out = def_levels + .as_mut() + .ok_or_else(|| general_err!("must specify definition levels"))?; + + let (values_read, levels_read) = reader.read_def_levels(out, levels_to_read)?; + + if levels_read != levels_to_read { + return Err(general_err!( + "insufficient definition levels read from column - expected {levels_to_read}, got {levels_read}" + )); + } + + values_read + } + None => levels_to_read, + }; + + let def_levels = def_levels.as_deref(); + reserve_values(values, values_to_read, levels_to_read, def_levels)?; + + let values_read = self.values_decoder.read(values, values_to_read)?; + + if values_read != values_to_read { + return Err(general_err!( + "insufficient values read from column - expected: {values_to_read}, got: {values_read}", + )); + } + + self.num_decoded_values += levels_to_read; + total_records_read += records_read; + total_levels_read += levels_to_read; + total_values_read += values_read; + } + + Ok((total_records_read, total_values_read, total_levels_read)) + } + + /// Skips over `num_records` records, where records are delimited by repetition levels of 0 + /// + /// # Returns + /// + /// Returns the number of records skipped + pub fn skip_records(&mut self, num_records: usize) -> Result { + let mut remaining_records = num_records; + while remaining_records != 0 { + if self.num_buffered_values == self.num_decoded_values { + let metadata = match self.page_reader.peek_next_page()? { + None => return Ok(num_records - remaining_records), + Some(metadata) => metadata, + }; + + // If dictionary, we must read it + if metadata.is_dict { + self.read_dictionary_page()?; + continue; + } + + // If page has less rows than the remaining records to + // be skipped, skip entire page + let rows = metadata.num_rows.or_else(|| { + // If no repetition levels, num_levels == num_rows + self.rep_level_decoder + .is_none() + .then_some(metadata.num_levels)? + }); + + if let Some(rows) = rows { + if rows <= remaining_records { + self.page_reader.skip_next_page()?; + remaining_records -= rows; + continue; + } + } + // because self.num_buffered_values == self.num_decoded_values means + // we need reads a new page and set up the decoders for levels + if !self.read_new_page()? { + return Ok(num_records - remaining_records); + } + } + + // start skip values in page level + + // The number of levels in the current data page + let remaining_levels = self.num_buffered_values - self.num_decoded_values; + + let (records_read, rep_levels_read) = match self.rep_level_decoder.as_mut() { + Some(decoder) => { + let (mut records_read, levels_read) = + decoder.skip_rep_levels(remaining_records, remaining_levels)?; + + if levels_read == remaining_levels && self.has_record_delimiter { + // Reached end of page, which implies records_read < remaining_records + // as otherwise would have stopped reading before reaching the end + assert!(records_read < remaining_records); // Sanity check + records_read += decoder.flush_partial() as usize; + } + + (records_read, levels_read) + } + None => { + // No repetition levels, so each level corresponds to a row + let levels = remaining_levels.min(remaining_records); + (levels, levels) + } + }; + + self.num_decoded_values += rep_levels_read; + remaining_records -= records_read; + + if self.num_buffered_values == self.num_decoded_values { + // Exhausted buffered page - no need to advance other decoders + continue; + } + + let (values_read, def_levels_read) = match self.def_level_decoder.as_mut() { + Some(decoder) => decoder.skip_def_levels(rep_levels_read)?, + None => (rep_levels_read, rep_levels_read), + }; + + if rep_levels_read != def_levels_read { + return Err(general_err!( + "levels mismatch, read {} repetition levels and {} definition levels", + rep_levels_read, + def_levels_read + )); + } + + let values = self.values_decoder.skip_values(values_read)?; + if values != values_read { + return Err(general_err!( + "skipped {} values, expected {}", + values, + values_read + )); + } + } + Ok(num_records - remaining_records) + } + + /// Read the next page as a dictionary page. If the next page is not a dictionary page, + /// this will return an error. + fn read_dictionary_page(&mut self) -> Result<()> { + match self.page_reader.get_next_page()? { + Some(Page::DictionaryPage { + buf, + num_values, + encoding, + is_sorted, + }) => self + .values_decoder + .set_dict(buf, num_values, encoding, is_sorted), + _ => Err(ParquetError::General( + "Invalid page. Expecting dictionary page".to_string(), + )), + } + } + + /// Reads a new page and set up the decoders for levels, values or dictionary. + /// Returns false if there's no page left. + fn read_new_page(&mut self) -> Result { + loop { + match self.page_reader.get_next_page()? { + // No more page to read + None => return Ok(false), + Some(current_page) => { + match current_page { + // 1. Dictionary page: configure dictionary for this page. + Page::DictionaryPage { + buf, + num_values, + encoding, + is_sorted, + } => { + self.values_decoder + .set_dict(buf, num_values, encoding, is_sorted)?; + continue; + } + // 2. Data page v1 + Page::DataPage { + buf, + num_values, + encoding, + def_level_encoding, + rep_level_encoding, + statistics: _, + } => { + self.num_buffered_values = num_values as _; + self.num_decoded_values = 0; + + let max_rep_level = self.descr.max_rep_level(); + let max_def_level = self.descr.max_def_level(); + + let mut offset = 0; + + if max_rep_level > 0 { + let (bytes_read, level_data) = parse_v1_level( + max_rep_level, + num_values, + rep_level_encoding, + buf.slice(offset..), + )?; + offset += bytes_read; + + self.has_record_delimiter = + self.page_reader.at_record_boundary()?; + + self.rep_level_decoder + .as_mut() + .unwrap() + .set_data(rep_level_encoding, level_data)?; + } + + if max_def_level > 0 { + let (bytes_read, level_data) = parse_v1_level( + max_def_level, + num_values, + def_level_encoding, + buf.slice(offset..), + )?; + offset += bytes_read; + + self.def_level_decoder + .as_mut() + .unwrap() + .set_data(def_level_encoding, level_data)?; + } + + self.values_decoder.set_data( + encoding, + buf.slice(offset..), + num_values as usize, + None, + )?; + return Ok(true); + } + // 3. Data page v2 + Page::DataPageV2 { + buf, + num_values, + encoding, + num_nulls, + num_rows: _, + def_levels_byte_len, + rep_levels_byte_len, + is_compressed: _, + statistics: _, + } => { + if num_nulls > num_values { + return Err(general_err!( + "more nulls than values in page, contained {} values and {} nulls", + num_values, + num_nulls + )); + } + + self.num_buffered_values = num_values as _; + self.num_decoded_values = 0; + + // DataPage v2 only supports RLE encoding for repetition + // levels + if self.descr.max_rep_level() > 0 { + // Technically a DataPage v2 should not write a record + // across multiple pages, however, the parquet writer + // used to do this so we preserve backwards compatibility + self.has_record_delimiter = + self.page_reader.at_record_boundary()?; + + self.rep_level_decoder.as_mut().unwrap().set_data( + Encoding::RLE, + buf.slice(..rep_levels_byte_len as usize), + )?; + } + + // DataPage v2 only supports RLE encoding for definition + // levels + if self.descr.max_def_level() > 0 { + self.def_level_decoder.as_mut().unwrap().set_data( + Encoding::RLE, + buf.slice( + rep_levels_byte_len as usize + ..(rep_levels_byte_len + def_levels_byte_len) as usize, + ), + )?; + } + + self.values_decoder.set_data( + encoding, + buf.slice((rep_levels_byte_len + def_levels_byte_len) as usize..), + num_values as usize, + Some((num_values - num_nulls) as usize), + )?; + return Ok(true); + } + }; + } + } + } + } + + /// Check whether there is more data to read from this column, + /// If the current page is fully decoded, this will load the next page + /// (if it exists) into the buffer + #[inline] + pub(crate) fn has_next(&mut self) -> Result { + if self.num_buffered_values == 0 || self.num_buffered_values == self.num_decoded_values { + // TODO: should we return false if read_new_page() = true and + // num_buffered_values = 0? + if !self.read_new_page()? { + Ok(false) + } else { + Ok(self.num_buffered_values != 0) + } + } else { + Ok(true) + } + } +} + +fn parse_v1_level( + max_level: i16, + num_buffered_values: u32, + encoding: Encoding, + buf: Bytes, +) -> Result<(usize, Bytes)> { + match encoding { + Encoding::RLE => { + let i32_size = std::mem::size_of::(); + if i32_size <= buf.len() { + let data_size = read_num_bytes::(i32_size, buf.as_ref()) as usize; + let end = i32_size + .checked_add(data_size) + .ok_or(general_err!("invalid level length"))?; + if end <= buf.len() { + return Ok((end, buf.slice(i32_size..end))); + } + } + Err(general_err!("not enough data to read levels")) + } + #[allow(deprecated)] + Encoding::BIT_PACKED => { + let bit_width = num_required_bits(max_level as u64); + let num_bytes = ceil(num_buffered_values as usize * bit_width as usize, 8); + Ok((num_bytes, buf.slice(..num_bytes))) + } + _ => Err(general_err!("invalid level encoding: {}", encoding)), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use rand::distr::uniform::SampleUniform; + use std::{collections::VecDeque, sync::Arc}; + + use crate::basic::Type as PhysicalType; + use crate::schema::types::{ColumnDescriptor, ColumnPath, Type as SchemaType}; + use crate::util::test_common::page_util::InMemoryPageReader; + use crate::util::test_common::rand_gen::make_pages; + + #[test] + fn test_parse_v1_level_invalid_length() { + // Say length is 10, but buffer is only 4 + let buf = Bytes::from(vec![10, 0, 0, 0]); + let err = parse_v1_level(1, 100, Encoding::RLE, buf).unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: not enough data to read levels" + ); + + // Say length is 4, but buffer is only 3 + let buf = Bytes::from(vec![4, 0, 0]); + let err = parse_v1_level(1, 100, Encoding::RLE, buf).unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: not enough data to read levels" + ); + } + + const NUM_LEVELS: usize = 128; + const NUM_PAGES: usize = 2; + const MAX_DEF_LEVEL: i16 = 5; + const MAX_REP_LEVEL: i16 = 5; + + // Macro to generate test cases + macro_rules! test { + // branch for generating i32 cases + ($test_func:ident, i32, $func:ident, $def_level:expr, $rep_level:expr, + $num_pages:expr, $num_levels:expr, $batch_size:expr, $min:expr, $max:expr) => { + test_internal!( + $test_func, + Int32Type, + get_test_int32_type, + $func, + $def_level, + $rep_level, + $num_pages, + $num_levels, + $batch_size, + $min, + $max + ); + }; + // branch for generating i64 cases + ($test_func:ident, i64, $func:ident, $def_level:expr, $rep_level:expr, + $num_pages:expr, $num_levels:expr, $batch_size:expr, $min:expr, $max:expr) => { + test_internal!( + $test_func, + Int64Type, + get_test_int64_type, + $func, + $def_level, + $rep_level, + $num_pages, + $num_levels, + $batch_size, + $min, + $max + ); + }; + } + + macro_rules! test_internal { + ($test_func:ident, $ty:ident, $pty:ident, $func:ident, $def_level:expr, + $rep_level:expr, $num_pages:expr, $num_levels:expr, $batch_size:expr, + $min:expr, $max:expr) => { + #[test] + fn $test_func() { + let desc = Arc::new(ColumnDescriptor::new( + Arc::new($pty()), + $def_level, + $rep_level, + ColumnPath::new(Vec::new()), + )); + let mut tester = ColumnReaderTester::<$ty>::new(); + tester.$func(desc, $num_pages, $num_levels, $batch_size, $min, $max); + } + }; + } + + test!( + test_read_plain_v1_int32, + i32, + plain_v1, + MAX_DEF_LEVEL, + MAX_REP_LEVEL, + NUM_PAGES, + NUM_LEVELS, + 16, + i32::MIN, + i32::MAX + ); + test!( + test_read_plain_v2_int32, + i32, + plain_v2, + MAX_DEF_LEVEL, + MAX_REP_LEVEL, + NUM_PAGES, + NUM_LEVELS, + 16, + i32::MIN, + i32::MAX + ); + + test!( + test_read_plain_v1_int32_uneven, + i32, + plain_v1, + MAX_DEF_LEVEL, + MAX_REP_LEVEL, + NUM_PAGES, + NUM_LEVELS, + 17, + i32::MIN, + i32::MAX + ); + test!( + test_read_plain_v2_int32_uneven, + i32, + plain_v2, + MAX_DEF_LEVEL, + MAX_REP_LEVEL, + NUM_PAGES, + NUM_LEVELS, + 17, + i32::MIN, + i32::MAX + ); + + test!( + test_read_plain_v1_int32_multi_page, + i32, + plain_v1, + MAX_DEF_LEVEL, + MAX_REP_LEVEL, + NUM_PAGES, + NUM_LEVELS, + 512, + i32::MIN, + i32::MAX + ); + test!( + test_read_plain_v2_int32_multi_page, + i32, + plain_v2, + MAX_DEF_LEVEL, + MAX_REP_LEVEL, + NUM_PAGES, + NUM_LEVELS, + 512, + i32::MIN, + i32::MAX + ); + + // test cases when column descriptor has MAX_DEF_LEVEL = 0 and MAX_REP_LEVEL = 0 + test!( + test_read_plain_v1_int32_required_non_repeated, + i32, + plain_v1, + 0, + 0, + NUM_PAGES, + NUM_LEVELS, + 16, + i32::MIN, + i32::MAX + ); + test!( + test_read_plain_v2_int32_required_non_repeated, + i32, + plain_v2, + 0, + 0, + NUM_PAGES, + NUM_LEVELS, + 16, + i32::MIN, + i32::MAX + ); + + test!( + test_read_plain_v1_int64, + i64, + plain_v1, + 1, + 1, + NUM_PAGES, + NUM_LEVELS, + 16, + i64::MIN, + i64::MAX + ); + test!( + test_read_plain_v2_int64, + i64, + plain_v2, + 1, + 1, + NUM_PAGES, + NUM_LEVELS, + 16, + i64::MIN, + i64::MAX + ); + + test!( + test_read_plain_v1_int64_uneven, + i64, + plain_v1, + 1, + 1, + NUM_PAGES, + NUM_LEVELS, + 17, + i64::MIN, + i64::MAX + ); + test!( + test_read_plain_v2_int64_uneven, + i64, + plain_v2, + 1, + 1, + NUM_PAGES, + NUM_LEVELS, + 17, + i64::MIN, + i64::MAX + ); + + test!( + test_read_plain_v1_int64_multi_page, + i64, + plain_v1, + 1, + 1, + NUM_PAGES, + NUM_LEVELS, + 512, + i64::MIN, + i64::MAX + ); + test!( + test_read_plain_v2_int64_multi_page, + i64, + plain_v2, + 1, + 1, + NUM_PAGES, + NUM_LEVELS, + 512, + i64::MIN, + i64::MAX + ); + + // test cases when column descriptor has MAX_DEF_LEVEL = 0 and MAX_REP_LEVEL = 0 + test!( + test_read_plain_v1_int64_required_non_repeated, + i64, + plain_v1, + 0, + 0, + NUM_PAGES, + NUM_LEVELS, + 16, + i64::MIN, + i64::MAX + ); + test!( + test_read_plain_v2_int64_required_non_repeated, + i64, + plain_v2, + 0, + 0, + NUM_PAGES, + NUM_LEVELS, + 16, + i64::MIN, + i64::MAX + ); + + test!( + test_read_dict_v1_int32_small, + i32, + dict_v1, + MAX_DEF_LEVEL, + MAX_REP_LEVEL, + 2, + 2, + 16, + 0, + 3 + ); + test!( + test_read_dict_v2_int32_small, + i32, + dict_v2, + MAX_DEF_LEVEL, + MAX_REP_LEVEL, + 2, + 2, + 16, + 0, + 3 + ); + + test!( + test_read_dict_v1_int32, + i32, + dict_v1, + MAX_DEF_LEVEL, + MAX_REP_LEVEL, + NUM_PAGES, + NUM_LEVELS, + 16, + 0, + 3 + ); + test!( + test_read_dict_v2_int32, + i32, + dict_v2, + MAX_DEF_LEVEL, + MAX_REP_LEVEL, + NUM_PAGES, + NUM_LEVELS, + 16, + 0, + 3 + ); + + test!( + test_read_dict_v1_int32_uneven, + i32, + dict_v1, + MAX_DEF_LEVEL, + MAX_REP_LEVEL, + NUM_PAGES, + NUM_LEVELS, + 17, + 0, + 3 + ); + test!( + test_read_dict_v2_int32_uneven, + i32, + dict_v2, + MAX_DEF_LEVEL, + MAX_REP_LEVEL, + NUM_PAGES, + NUM_LEVELS, + 17, + 0, + 3 + ); + + test!( + test_read_dict_v1_int32_multi_page, + i32, + dict_v1, + MAX_DEF_LEVEL, + MAX_REP_LEVEL, + NUM_PAGES, + NUM_LEVELS, + 512, + 0, + 3 + ); + test!( + test_read_dict_v2_int32_multi_page, + i32, + dict_v2, + MAX_DEF_LEVEL, + MAX_REP_LEVEL, + NUM_PAGES, + NUM_LEVELS, + 512, + 0, + 3 + ); + + test!( + test_read_dict_v1_int64, + i64, + dict_v1, + MAX_DEF_LEVEL, + MAX_REP_LEVEL, + NUM_PAGES, + NUM_LEVELS, + 16, + 0, + 3 + ); + test!( + test_read_dict_v2_int64, + i64, + dict_v2, + MAX_DEF_LEVEL, + MAX_REP_LEVEL, + NUM_PAGES, + NUM_LEVELS, + 16, + 0, + 3 + ); + + #[test] + fn test_read_batch_values_only() { + test_read_batch_int32(16, 0, 0); + } + + #[test] + fn test_read_batch_values_def_levels() { + test_read_batch_int32(16, MAX_DEF_LEVEL, 0); + } + + #[test] + fn test_read_batch_values_rep_levels() { + test_read_batch_int32(16, 0, MAX_REP_LEVEL); + } + + #[test] + fn test_read_batch_values_def_rep_levels() { + test_read_batch_int32(128, MAX_DEF_LEVEL, MAX_REP_LEVEL); + } + + #[test] + fn test_read_batch_adjust_after_buffering_page() { + // This test covers scenario when buffering new page results in setting number + // of decoded values to 0, resulting on reading `batch_size` of values, but it is + // larger than we can insert into slice (affects values and levels). + // + // Note: values are chosen to reproduce the issue. + // + let primitive_type = get_test_int32_type(); + let desc = Arc::new(ColumnDescriptor::new( + Arc::new(primitive_type), + 1, + 1, + ColumnPath::new(Vec::new()), + )); + + let num_pages = 2; + let num_levels = 4; + let batch_size = 5; + + let mut tester = ColumnReaderTester::::new(); + tester.test_read_batch( + desc, + Encoding::RLE_DICTIONARY, + num_pages, + num_levels, + batch_size, + i32::MIN, + i32::MAX, + false, + ); + } + + // ---------------------------------------------------------------------- + // Helper methods to make pages and test + // + // # Overview + // + // Most of the test functionality is implemented in `ColumnReaderTester`, which + // provides some general data page test methods: + // - `test_read_batch_general` + // - `test_read_batch` + // + // There are also some high level wrappers that are part of `ColumnReaderTester`: + // - `plain_v1` -> call `test_read_batch_general` with data page v1 and plain encoding + // - `plain_v2` -> call `test_read_batch_general` with data page v2 and plain encoding + // - `dict_v1` -> call `test_read_batch_general` with data page v1 + dictionary page + // - `dict_v2` -> call `test_read_batch_general` with data page v2 + dictionary page + // + // And even higher level wrappers that simplify testing of almost the same test cases: + // - `get_test_int32_type`, provides dummy schema type + // - `get_test_int64_type`, provides dummy schema type + // - `test_read_batch_int32`, wrapper for `read_batch` tests, since they are basically + // the same, just different def/rep levels and batch size. + // + // # Page assembly + // + // Page construction and generation of values, definition and repetition levels + // happens in `make_pages` function. + // All values are randomly generated based on provided min/max, levels are calculated + // based on provided max level for column descriptor (which is basically either int32 + // or int64 type in tests) and `levels_per_page` variable. + // + // We use `DataPageBuilder` and its implementation `DataPageBuilderImpl` to actually + // turn values, definition and repetition levels into data pages (either v1 or v2). + // + // Those data pages are then stored as part of `TestPageReader` (we just pass vector + // of generated pages directly), which implements `PageReader` interface. + // + // # Comparison + // + // This allows us to pass test page reader into column reader, so we can test + // functionality of column reader - see `test_read_batch`, where we create column + // reader -> typed column reader, buffer values in `read_batch` method and compare + // output with generated data. + + // Returns dummy Parquet `Type` for primitive field, because most of our tests use + // INT32 physical type. + fn get_test_int32_type() -> SchemaType { + SchemaType::primitive_type_builder("a", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::INT_32) + .with_length(-1) + .build() + .expect("build() should be OK") + } + + // Returns dummy Parquet `Type` for INT64 physical type. + fn get_test_int64_type() -> SchemaType { + SchemaType::primitive_type_builder("a", PhysicalType::INT64) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::INT_64) + .with_length(-1) + .build() + .expect("build() should be OK") + } + + // Tests `read_batch()` functionality for INT32. + // + // This is a high level wrapper on `ColumnReaderTester` that allows us to specify some + // boilerplate code for setting up definition/repetition levels and column descriptor. + fn test_read_batch_int32(batch_size: usize, max_def_level: i16, max_rep_level: i16) { + let primitive_type = get_test_int32_type(); + + let desc = Arc::new(ColumnDescriptor::new( + Arc::new(primitive_type), + max_def_level, + max_rep_level, + ColumnPath::new(Vec::new()), + )); + + let mut tester = ColumnReaderTester::::new(); + tester.test_read_batch( + desc, + Encoding::RLE_DICTIONARY, + NUM_PAGES, + NUM_LEVELS, + batch_size, + i32::MIN, + i32::MAX, + false, + ); + } + + struct ColumnReaderTester + where + T::T: PartialOrd + SampleUniform + Copy, + { + rep_levels: Vec, + def_levels: Vec, + values: Vec, + } + + impl ColumnReaderTester + where + T::T: PartialOrd + SampleUniform + Copy, + { + pub fn new() -> Self { + Self { + rep_levels: Vec::new(), + def_levels: Vec::new(), + values: Vec::new(), + } + } + + // Method to generate and test data pages v1 + fn plain_v1( + &mut self, + desc: ColumnDescPtr, + num_pages: usize, + num_levels: usize, + batch_size: usize, + min: T::T, + max: T::T, + ) { + self.test_read_batch_general( + desc, + Encoding::PLAIN, + num_pages, + num_levels, + batch_size, + min, + max, + false, + ); + } + + // Method to generate and test data pages v2 + fn plain_v2( + &mut self, + desc: ColumnDescPtr, + num_pages: usize, + num_levels: usize, + batch_size: usize, + min: T::T, + max: T::T, + ) { + self.test_read_batch_general( + desc, + Encoding::PLAIN, + num_pages, + num_levels, + batch_size, + min, + max, + true, + ); + } + + // Method to generate and test dictionary page + data pages v1 + fn dict_v1( + &mut self, + desc: ColumnDescPtr, + num_pages: usize, + num_levels: usize, + batch_size: usize, + min: T::T, + max: T::T, + ) { + self.test_read_batch_general( + desc, + Encoding::RLE_DICTIONARY, + num_pages, + num_levels, + batch_size, + min, + max, + false, + ); + } + + // Method to generate and test dictionary page + data pages v2 + fn dict_v2( + &mut self, + desc: ColumnDescPtr, + num_pages: usize, + num_levels: usize, + batch_size: usize, + min: T::T, + max: T::T, + ) { + self.test_read_batch_general( + desc, + Encoding::RLE_DICTIONARY, + num_pages, + num_levels, + batch_size, + min, + max, + true, + ); + } + + // Helper function for the general case of `read_batch()` where `values`, + // `def_levels` and `rep_levels` are always provided with enough space. + #[allow(clippy::too_many_arguments)] + fn test_read_batch_general( + &mut self, + desc: ColumnDescPtr, + encoding: Encoding, + num_pages: usize, + num_levels: usize, + batch_size: usize, + min: T::T, + max: T::T, + use_v2: bool, + ) { + self.test_read_batch( + desc, encoding, num_pages, num_levels, batch_size, min, max, use_v2, + ); + } + + // Helper function to test `read_batch()` method with custom buffers for values, + // definition and repetition levels. + #[allow(clippy::too_many_arguments)] + fn test_read_batch( + &mut self, + desc: ColumnDescPtr, + encoding: Encoding, + num_pages: usize, + num_levels: usize, + batch_size: usize, + min: T::T, + max: T::T, + use_v2: bool, + ) { + let mut pages = VecDeque::new(); + make_pages::( + desc.clone(), + encoding, + num_pages, + num_levels, + min, + max, + &mut self.def_levels, + &mut self.rep_levels, + &mut self.values, + &mut pages, + use_v2, + ); + let max_def_level = desc.max_def_level(); + let max_rep_level = desc.max_rep_level(); + let page_reader = InMemoryPageReader::new(pages); + let column_reader: ColumnReader = get_column_reader(desc, Box::new(page_reader)); + let mut typed_column_reader = get_typed_column_reader::(column_reader); + + let mut values = Vec::new(); + let mut def_levels = Vec::new(); + let mut rep_levels = Vec::new(); + + let mut curr_values_read = 0; + let mut curr_levels_read = 0; + loop { + let (_, values_read, levels_read) = typed_column_reader + .read_records( + batch_size, + Some(&mut def_levels), + Some(&mut rep_levels), + &mut values, + ) + .expect("read_batch() should be OK"); + + curr_values_read += values_read; + curr_levels_read += levels_read; + + if values_read == 0 && levels_read == 0 { + break; + } + } + + assert_eq!(values, self.values, "values content doesn't match"); + + if max_def_level > 0 { + assert_eq!( + def_levels, self.def_levels, + "definition levels content doesn't match" + ); + } + + if max_rep_level > 0 { + assert_eq!( + rep_levels, self.rep_levels, + "repetition levels content doesn't match" + ); + } + + assert!( + curr_levels_read >= curr_values_read, + "expected levels read to be greater than values read" + ); + } + } + + /// Regression test for + /// + /// Reproduces the production scenario: all DataPage v2 pages for a + /// list column (rep_level=1) read without an offset index (i.e. + /// `at_record_boundary` returns false for non-last pages). + /// + /// When a prior operation (here `skip_records(1)`) loads a v2 page, + /// and a subsequent `skip_records` exhausts the remaining levels on + /// that page, the rep level decoder is left with `has_partial=true`. + /// Because `has_record_delimiter` is false, the partial is not + /// flushed during level-based processing. When the next v2 page is + /// then peeked with `num_rows` available, the whole-page-skip + /// shortcut must flush the pending partial first. Otherwise: + /// + /// 1. The skip over-counts (skips N+1 records instead of N), and + /// 2. The stale `has_partial` causes a subsequent `read_records` to + /// produce a "phantom" record with 0 values. + #[test] + fn test_skip_records_v2_page_skip_accounts_for_partial() { + use crate::encodings::levels::LevelEncoder; + + let max_rep_level: i16 = 1; + let max_def_level: i16 = 1; + + // Column descriptor for a list element column (rep=1, def=1) + let primitive_type = SchemaType::primitive_type_builder("element", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(); + let desc = Arc::new(ColumnDescriptor::new( + Arc::new(primitive_type), + max_def_level, + max_rep_level, + ColumnPath::new(vec!["list".to_string(), "element".to_string()]), + )); + + // Helper: build a DataPage v2 for this list column. + let make_v2_page = + |rep_levels: &[i16], def_levels: &[i16], values: &[i32], num_rows: u32| -> Page { + let mut rep_enc = LevelEncoder::v2_streaming(max_rep_level); + rep_enc.put_with_observer(rep_levels, |_, _| {}); + let rep_bytes = rep_enc.consume(); + + let mut def_enc = LevelEncoder::v2_streaming(max_def_level); + def_enc.put_with_observer(def_levels, |_, _| {}); + let def_bytes = def_enc.consume(); + + let val_bytes: Vec = values.iter().flat_map(|v| v.to_le_bytes()).collect(); + + let mut buf = Vec::new(); + buf.extend_from_slice(&rep_bytes); + buf.extend_from_slice(&def_bytes); + buf.extend_from_slice(&val_bytes); + + Page::DataPageV2 { + buf: Bytes::from(buf), + num_values: rep_levels.len() as u32, + encoding: Encoding::PLAIN, + num_nulls: 0, + num_rows, + def_levels_byte_len: def_bytes.len() as u32, + rep_levels_byte_len: rep_bytes.len() as u32, + is_compressed: false, + statistics: None, + } + }; + + // All pages are DataPage v2 (matching the production scenario where + // parquet-rs writes only v2 data pages and no offset index is loaded, + // so at_record_boundary() returns false for non-last pages). + + // Page 1 (v2): 2 records × 2 elements = [10,20], [30,40] + let page1 = make_v2_page(&[0, 1, 0, 1], &[1, 1, 1, 1], &[10, 20, 30, 40], 2); + + // Page 2 (v2): 2 records × 2 elements = [50,60], [70,80] + let page2 = make_v2_page(&[0, 1, 0, 1], &[1, 1, 1, 1], &[50, 60, 70, 80], 2); + + // Page 3 (v2): 1 record × 2 elements = [90,100] + let page3 = make_v2_page(&[0, 1], &[1, 1], &[90, 100], 1); + + // 5 records total: [10,20], [30,40], [50,60], [70,80], [90,100] + let pages = VecDeque::from(vec![page1, page2, page3]); + let page_reader = InMemoryPageReader::new(pages); + let column_reader: ColumnReader = get_column_reader(desc, Box::new(page_reader)); + let mut typed_reader = get_typed_column_reader::(column_reader); + + // Step 1 — skip 1 record: + // Peek page 1: num_rows=2, remaining=1 → rows(2) > remaining(1), + // so the page is LOADED (not whole-page-skipped). + // Level-based skip consumes rep levels [0,1] for record [10,20], + // stopping at the 0 that starts record [30,40]. + let skipped = typed_reader.skip_records(1).unwrap(); + assert_eq!(skipped, 1); + + // Step 2 — skip 2 more records ([30,40] and [50,60]): + // Mid-page in page 1 with 2 remaining levels [0,1] for [30,40]. + // skip_rep_levels(2, 2): the leading 0 does NOT act as a record + // delimiter (has_partial=false, idx==0), so count_records returns + // (true, 0, 2) — all levels consumed, has_partial=true, 0 records. + // + // has_record_delimiter is false → no flush at page boundary. + // Page 1 exhausted → peek page 2 (v2, num_rows=2). + // + // With fix: flush_partial → remaining 2→1, page 2 NOT skipped + // (rows=2 > remaining=1). Load page 2, skip 1 record [50,60]. + // + // Without fix: rows(2) <= remaining(2) → page 2 whole-page-skipped, + // over-counting by 1. has_partial stays true (stale from page 1). + let skipped = typed_reader.skip_records(2).unwrap(); + assert_eq!(skipped, 2); + + // Step 3 — read 1 record: + let mut values = Vec::new(); + let mut def_levels = Vec::new(); + let mut rep_levels = Vec::new(); + + let (records, values_read, levels_read) = typed_reader + .read_records(1, Some(&mut def_levels), Some(&mut rep_levels), &mut values) + .unwrap(); + + // Without the fix: (1, 0, 0) — phantom record from stale has_partial; + // the rep=0 on page 3 "completes" the phantom, yielding 0 values. + // With the fix: (1, 2, 2) — correctly reads record [70, 80]. + assert_eq!(records, 1, "should read exactly 1 record"); + assert_eq!(levels_read, 2, "should read 2 levels for the record"); + assert_eq!(values_read, 2, "should read 2 non-null values"); + assert_eq!(values, vec![70, 80], "should contain 4th record's values"); + assert_eq!(rep_levels, vec![0, 1], "rep levels for a 2-element list"); + assert_eq!(def_levels, vec![1, 1], "def levels (all non-null)"); + } +} diff --git a/vendor/parquet/src/column/reader/decoder.rs b/vendor/parquet/src/column/reader/decoder.rs new file mode 100644 index 0000000..053db81 --- /dev/null +++ b/vendor/parquet/src/column/reader/decoder.rs @@ -0,0 +1,565 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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 bytes::Bytes; + +use crate::basic::{Encoding, EncodingMask}; +use crate::data_type::DataType; +use crate::encodings::{ + decoding::{Decoder, DictDecoder, PlainDecoder, get_decoder}, + rle::RleDecoder, +}; +use crate::errors::{ParquetError, Result}; +use crate::schema::types::ColumnDescPtr; +use crate::util::bit_util::{BitReader, num_required_bits}; + +/// Decodes level data +pub trait ColumnLevelDecoder { + type Buffer; + + /// Set data for this [`ColumnLevelDecoder`] + fn set_data(&mut self, encoding: Encoding, data: Bytes) -> Result<()>; +} + +pub trait RepetitionLevelDecoder: ColumnLevelDecoder { + /// Read up to `max_records` of repetition level data into `out` returning the number + /// of complete records and levels read + /// + /// A record only ends when the data contains a subsequent repetition level of 0, + /// it is therefore left to the caller to delimit the final record in a column + /// + /// # Panics + /// + /// Implementations may panic if `range` overlaps with already written data + fn read_rep_levels( + &mut self, + out: &mut Self::Buffer, + num_records: usize, + num_levels: usize, + ) -> Result<(usize, usize)>; + + /// Skips over up to `num_levels` repetition levels corresponding to `num_records` records, + /// where a record is delimited by a repetition level of 0 + /// + /// Returns the number of records skipped, and the number of levels skipped + /// + /// A record only ends when the data contains a subsequent repetition level of 0, + /// it is therefore left to the caller to delimit the final record in a column + fn skip_rep_levels(&mut self, num_records: usize, num_levels: usize) -> Result<(usize, usize)>; + + /// Flush any partially read or skipped record + fn flush_partial(&mut self) -> bool; +} + +pub trait DefinitionLevelDecoder: ColumnLevelDecoder { + /// Read up to `num_levels` definition levels into `out`. + /// + /// Returns the number of values read, and the number of levels read. + /// + /// # Panics + /// + /// Implementations may panic if `range` overlaps with already written data + fn read_def_levels( + &mut self, + out: &mut Self::Buffer, + num_levels: usize, + ) -> Result<(usize, usize)>; + + /// Skips over `num_levels` definition levels. + /// + /// Returns the number of values skipped, and the number of levels skipped. + fn skip_def_levels(&mut self, num_levels: usize) -> Result<(usize, usize)>; +} + +/// Decodes value data +pub trait ColumnValueDecoder { + type Buffer; + + /// Create a new [`ColumnValueDecoder`] + fn new(col: &ColumnDescPtr) -> Self; + + /// Set the current dictionary page + fn set_dict( + &mut self, + buf: Bytes, + num_values: u32, + encoding: Encoding, + is_sorted: bool, + ) -> Result<()>; + + /// Set the current data page + /// + /// - `encoding` - the encoding of the page + /// - `data` - a point to the page's uncompressed value data + /// - `num_levels` - the number of levels contained within the page, i.e. values including nulls + /// - `num_values` - the number of non-null values contained within the page (V2 page only) + /// + /// Note: data encoded with [`Encoding::RLE`] may not know its exact length, as the final + /// run may be zero-padded. As such if `num_values` is not provided (i.e. `None`), + /// subsequent calls to `ColumnValueDecoder::read` may yield more values than + /// non-null definition levels within the page + fn set_data( + &mut self, + encoding: Encoding, + data: Bytes, + num_levels: usize, + num_values: Option, + ) -> Result<()>; + + /// Read up to `num_values` values into `out` + /// + /// # Panics + /// + /// Implementations may panic if `range` overlaps with already written data + /// + fn read(&mut self, out: &mut Self::Buffer, num_values: usize) -> Result; + + /// Skips over `num_values` values + /// + /// Returns the number of values skipped + fn skip_values(&mut self, num_values: usize) -> Result; +} + +/// Bucket-based storage for decoder instances keyed by `Encoding`. +/// +/// This replaces `HashMap` lookups with direct indexing to avoid hashing overhead in the +/// hot decoding paths. +const ENCODING_SLOTS: usize = Encoding::MAX_DISCRIMINANT as usize + 1; + +/// An implementation of [`ColumnValueDecoder`] for `[T::T]` +pub struct ColumnValueDecoderImpl { + descr: ColumnDescPtr, + + current_encoding: Option, + + /// Cache of decoders for existing encodings. + /// Uses `EncodingMask` and dense storage keyed by encoding discriminant. + decoder_mask: EncodingMask, + decoders: [Option>>; ENCODING_SLOTS], +} + +impl ColumnValueDecoder for ColumnValueDecoderImpl { + type Buffer = Vec; + + fn new(descr: &ColumnDescPtr) -> Self { + Self { + descr: descr.clone(), + current_encoding: None, + decoder_mask: EncodingMask::default(), + decoders: std::array::from_fn(|_| None), + } + } + + fn set_dict( + &mut self, + buf: Bytes, + num_values: u32, + mut encoding: Encoding, + _is_sorted: bool, + ) -> Result<()> { + if encoding == Encoding::PLAIN || encoding == Encoding::PLAIN_DICTIONARY { + encoding = Encoding::RLE_DICTIONARY + } + + if self.decoder_mask.is_set(encoding) { + return Err(general_err!("Column cannot have more than one dictionary")); + } + + if encoding == Encoding::RLE_DICTIONARY { + let mut dictionary = PlainDecoder::::new(self.descr.type_length()); + dictionary.set_data(buf, num_values as usize)?; + + let mut decoder = DictDecoder::new(); + decoder.set_dict(Box::new(dictionary))?; + self.decoders[encoding as usize] = Some(Box::new(decoder)); + self.decoder_mask.insert(encoding); + Ok(()) + } else { + Err(nyi_err!( + "Invalid/Unsupported encoding type for dictionary: {}", + encoding + )) + } + } + + fn set_data( + &mut self, + mut encoding: Encoding, + data: Bytes, + num_levels: usize, + num_values: Option, + ) -> Result<()> { + if encoding == Encoding::PLAIN_DICTIONARY { + encoding = Encoding::RLE_DICTIONARY; + } + + let decoder = if encoding == Encoding::RLE_DICTIONARY { + self.decoders[encoding as usize] + .as_mut() + .expect("Decoder for dict should have been set") + } else { + let slot = encoding as usize; + if self.decoders[slot].is_none() { + let data_decoder = get_decoder::(self.descr.clone(), encoding)?; + self.decoders[slot] = Some(data_decoder); + self.decoder_mask.insert(encoding); + } + self.decoders[slot] + .as_mut() + .expect("decoder should have been inserted") + }; + + decoder.set_data(data, num_values.unwrap_or(num_levels))?; + self.current_encoding = Some(encoding); + Ok(()) + } + + fn read(&mut self, out: &mut Self::Buffer, num_values: usize) -> Result { + let encoding = self + .current_encoding + .expect("current_encoding should be set"); + + let current_decoder = self.decoders[encoding as usize] + .as_mut() + .unwrap_or_else(|| panic!("decoder for encoding {encoding} should be set")); + + // TODO: Push vec into decoder (#5177) + let start = out.len(); + out.resize(start + num_values, T::T::default()); + let read = current_decoder.get(&mut out[start..])?; + out.truncate(start + read); + Ok(read) + } + + fn skip_values(&mut self, num_values: usize) -> Result { + let encoding = self + .current_encoding + .expect("current_encoding should be set"); + + let current_decoder = self.decoders[encoding as usize] + .as_mut() + .unwrap_or_else(|| panic!("decoder for encoding {encoding} should be set")); + + current_decoder.skip(num_values) + } +} + +const SKIP_BUFFER_SIZE: usize = 1024; + +enum LevelDecoder { + Packed(BitReader, u8), + Rle(RleDecoder), +} + +impl LevelDecoder { + fn new(encoding: Encoding, data: Bytes, bit_width: u8) -> Result { + match encoding { + Encoding::RLE => { + let mut decoder = RleDecoder::new(bit_width); + decoder.set_data(data)?; + Ok(Self::Rle(decoder)) + } + #[allow(deprecated)] + Encoding::BIT_PACKED => Ok(Self::Packed(BitReader::new(data), bit_width)), + _ => unreachable!("invalid level encoding: {}", encoding), + } + } + + fn read(&mut self, out: &mut [i16]) -> Result { + match self { + Self::Packed(reader, bit_width) => { + Ok(reader.get_batch::(out, *bit_width as usize)) + } + Self::Rle(reader) => Ok(reader.get_batch(out)?), + } + } +} + +/// An implementation of [`DefinitionLevelDecoder`] for `[i16]` +pub struct DefinitionLevelDecoderImpl { + decoder: Option, + bit_width: u8, + max_level: i16, +} + +impl DefinitionLevelDecoderImpl { + pub fn new(max_level: i16) -> Self { + let bit_width = num_required_bits(max_level as u64); + Self { + decoder: None, + bit_width, + max_level, + } + } +} + +impl ColumnLevelDecoder for DefinitionLevelDecoderImpl { + type Buffer = Vec; + + fn set_data(&mut self, encoding: Encoding, data: Bytes) -> Result<()> { + self.decoder = Some(LevelDecoder::new(encoding, data, self.bit_width)?); + Ok(()) + } +} + +impl DefinitionLevelDecoder for DefinitionLevelDecoderImpl { + fn read_def_levels( + &mut self, + out: &mut Self::Buffer, + num_levels: usize, + ) -> Result<(usize, usize)> { + // TODO: Push vec into decoder (#5177) + let start = out.len(); + out.resize(start + num_levels, 0); + let levels_read = self.decoder.as_mut().unwrap().read(&mut out[start..])?; + out.truncate(start + levels_read); + + let iter = out.iter().skip(start); + let values_read = iter.filter(|x| **x == self.max_level).count(); + Ok((values_read, levels_read)) + } + + fn skip_def_levels(&mut self, num_levels: usize) -> Result<(usize, usize)> { + let mut level_skip = 0; + let mut value_skip = 0; + let mut buf: Vec = vec![]; + while level_skip < num_levels { + let remaining_levels = num_levels - level_skip; + + let to_read = remaining_levels.min(SKIP_BUFFER_SIZE); + buf.resize(to_read, 0); + let (values_read, levels_read) = self.read_def_levels(&mut buf, to_read)?; + if levels_read == 0 { + // Reached end of page + break; + } + + level_skip += levels_read; + value_skip += values_read; + } + + Ok((value_skip, level_skip)) + } +} + +pub(crate) const REPETITION_LEVELS_BATCH_SIZE: usize = 1024; + +/// An implementation of [`RepetitionLevelDecoder`] for `[i16]` +pub struct RepetitionLevelDecoderImpl { + decoder: Option, + bit_width: u8, + buffer: Box<[i16; REPETITION_LEVELS_BATCH_SIZE]>, + buffer_len: usize, + buffer_offset: usize, + has_partial: bool, +} + +impl RepetitionLevelDecoderImpl { + pub fn new(max_level: i16) -> Self { + let bit_width = num_required_bits(max_level as u64); + Self { + decoder: None, + bit_width, + buffer: Box::new([0; REPETITION_LEVELS_BATCH_SIZE]), + buffer_offset: 0, + buffer_len: 0, + has_partial: false, + } + } + + fn fill_buf(&mut self) -> Result<()> { + let read = self.decoder.as_mut().unwrap().read(self.buffer.as_mut())?; + self.buffer_offset = 0; + self.buffer_len = read; + Ok(()) + } + + /// Inspects the buffered repetition levels in the range `self.buffer_offset..self.buffer_len` + /// and returns the number of "complete" records along with the corresponding number of values + /// + /// A "complete" record is one where the buffer contains a subsequent repetition level of 0 + fn count_records(&mut self, records_to_read: usize, num_levels: usize) -> (bool, usize, usize) { + let mut records_read = 0; + + let levels = num_levels.min(self.buffer_len - self.buffer_offset); + let buf = self.buffer.iter().skip(self.buffer_offset); + for (idx, item) in buf.take(levels).enumerate() { + if *item == 0 && (idx != 0 || self.has_partial) { + records_read += 1; + + if records_read == records_to_read { + return (false, records_read, idx); + } + } + } + // Either ran out of space in `num_levels` or data in `self.buffer` + (true, records_read, levels) + } +} + +impl ColumnLevelDecoder for RepetitionLevelDecoderImpl { + type Buffer = Vec; + + fn set_data(&mut self, encoding: Encoding, data: Bytes) -> Result<()> { + self.decoder = Some(LevelDecoder::new(encoding, data, self.bit_width)?); + self.buffer_len = 0; + self.buffer_offset = 0; + Ok(()) + } +} + +impl RepetitionLevelDecoder for RepetitionLevelDecoderImpl { + fn read_rep_levels( + &mut self, + out: &mut Self::Buffer, + num_records: usize, + num_levels: usize, + ) -> Result<(usize, usize)> { + let mut total_records_read = 0; + let mut total_levels_read = 0; + + while total_records_read < num_records && total_levels_read < num_levels { + if self.buffer_len == self.buffer_offset { + self.fill_buf()?; + if self.buffer_len == 0 { + break; + } + } + + let (partial, records_read, levels_read) = self.count_records( + num_records - total_records_read, + num_levels - total_levels_read, + ); + + out.extend_from_slice( + &self.buffer[self.buffer_offset..self.buffer_offset + levels_read], + ); + + total_levels_read += levels_read; + total_records_read += records_read; + self.buffer_offset += levels_read; + self.has_partial = partial; + } + Ok((total_records_read, total_levels_read)) + } + + fn skip_rep_levels(&mut self, num_records: usize, num_levels: usize) -> Result<(usize, usize)> { + let mut total_records_read = 0; + let mut total_levels_read = 0; + + while total_records_read < num_records && total_levels_read < num_levels { + if self.buffer_len == self.buffer_offset { + self.fill_buf()?; + if self.buffer_len == 0 { + break; + } + } + + let (partial, records_read, levels_read) = self.count_records( + num_records - total_records_read, + num_levels - total_levels_read, + ); + + total_levels_read += levels_read; + total_records_read += records_read; + self.buffer_offset += levels_read; + self.has_partial = partial; + } + Ok((total_records_read, total_levels_read)) + } + + fn flush_partial(&mut self) -> bool { + std::mem::take(&mut self.has_partial) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::encodings::rle::RleEncoder; + use rand::{prelude::*, rng}; + + #[test] + fn test_skip_padding() { + let mut encoder = RleEncoder::new(1, 1024); + encoder.put(0); + (0..3).for_each(|_| encoder.put(1)); + let data = Bytes::from(encoder.consume()); + + let mut decoder = RepetitionLevelDecoderImpl::new(1); + decoder.set_data(Encoding::RLE, data.clone()).unwrap(); + let (_, levels) = decoder.skip_rep_levels(100, 4).unwrap(); + assert_eq!(levels, 4); + + // The length of the final bit packed run is ambiguous, so without the correct + // levels limit, it will decode zero padding + let mut decoder = RepetitionLevelDecoderImpl::new(1); + decoder.set_data(Encoding::RLE, data).unwrap(); + let (_, levels) = decoder.skip_rep_levels(100, 6).unwrap(); + assert_eq!(levels, 6); + } + + #[test] + fn test_skip_rep_levels() { + for _ in 0..10 { + let mut rng = rng(); + let total_len = 10000_usize; + let mut encoded: Vec = (0..total_len).map(|_| rng.random_range(0..5)).collect(); + encoded[0] = 0; + let mut encoder = RleEncoder::new(3, 1024); + for v in &encoded { + encoder.put(*v as _) + } + let data = Bytes::from(encoder.consume()); + + let mut decoder = RepetitionLevelDecoderImpl::new(5); + decoder.set_data(Encoding::RLE, data).unwrap(); + + let total_records = encoded.iter().filter(|x| **x == 0).count(); + let mut remaining_records = total_records; + let mut remaining_levels = encoded.len(); + loop { + let skip = rng.random_bool(0.5); + let records = rng.random_range(1..=remaining_records.min(5)); + let (records_read, levels_read) = if skip { + decoder.skip_rep_levels(records, remaining_levels).unwrap() + } else { + let mut decoded = Vec::new(); + let (records_read, levels_read) = decoder + .read_rep_levels(&mut decoded, records, remaining_levels) + .unwrap(); + + assert_eq!( + decoded, + encoded[encoded.len() - remaining_levels..][..levels_read] + ); + (records_read, levels_read) + }; + + remaining_levels = remaining_levels.checked_sub(levels_read).unwrap(); + if remaining_levels == 0 { + assert_eq!(records_read + 1, records); + assert_eq!(records, remaining_records); + break; + } + assert_eq!(records_read, records); + remaining_records -= records; + assert_ne!(remaining_records, 0); + } + } + } +} diff --git a/vendor/parquet/src/column/writer/byte_budget_chunker.rs b/vendor/parquet/src/column/writer/byte_budget_chunker.rs new file mode 100644 index 0000000..56c7b4c --- /dev/null +++ b/vendor/parquet/src/column/writer/byte_budget_chunker.rs @@ -0,0 +1,204 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! See [`ByteBudgetChunker`] for byte-budget-aware mini-batch sizing. + +use crate::basic::Type; +use crate::column::writer::LevelDataRef; +use crate::column::writer::encoder::ColumnValueEncoder; +use crate::file::properties::WriterProperties; +use crate::schema::types::ColumnDescriptor; + +/// Picks byte-budget-aware mini-batch sizes for one column. +/// +/// The parquet column writer checks the data page byte limit only *after* +/// each mini-batch finishes writing. Mini-batches are sized in rows +/// (`write_batch_size`, default 1024), so for BYTE_ARRAY columns whose +/// values are large (e.g. multi-MiB blobs) a single mini-batch can buffer +/// GiB into one page before the limit is consulted. +/// +/// This isolates the per-chunk decision that prevents that: given a chunk's +/// level data and the input values, pick the largest `sub_batch_size` such +/// that one mini-batch fits in one page byte budget. For the overwhelmingly +/// common case (small or fixed-width values) the answer is just `chunk_size` +/// and the decision is O(1) on the column type — only when the input might +/// overflow does the chunker consult the encoder's byte estimate. +pub(crate) struct ByteBudgetChunker { + /// Configured data page byte limit for the column. + page_byte_limit: usize, + /// Max definition level of the column; a level equal to this marks a + /// present (non-null) leaf value. Used to count values per chunk. + max_def_level: i16, + /// `true` when no chunk of `base_batch_size` values can ever overflow + /// `page_byte_limit` regardless of input. Set once at column open from + /// the physical type's known per-value byte size; lets the per-chunk + /// decision short-circuit with no work for every numeric, bool, or + /// narrow `FIXED_LEN_BYTE_ARRAY` column. + static_always_fits: bool, + /// Configured dictionary page byte limit for the column. + dict_page_byte_limit: usize, + /// As [`Self::static_always_fits`] but for the dictionary page: `true` + /// when one `base_batch_size` mini-batch of this fixed-width type cannot + /// overshoot `dict_page_byte_limit` by more than one mini-batch's worth. + static_dict_always_fits: bool, +} + +impl ByteBudgetChunker { + #[inline] + pub(crate) fn new( + descr: &ColumnDescriptor, + props: &WriterProperties, + base_batch_size: usize, + ) -> Self { + let page_byte_limit = props.column_data_page_size_limit(descr.path()); + let dict_page_byte_limit = props.column_dictionary_page_size_limit(descr.path()); + let static_bytes_per_value = match descr.physical_type() { + Type::BOOLEAN => Some(1), + Type::INT32 | Type::FLOAT => Some(std::mem::size_of::()), + Type::INT64 | Type::DOUBLE => Some(std::mem::size_of::()), + Type::INT96 => Some(12), + Type::FIXED_LEN_BYTE_ARRAY => Some(descr.type_length().max(0) as usize), + Type::BYTE_ARRAY => None, + }; + let static_fits = |limit: usize| { + static_bytes_per_value + .map(|b| b.saturating_mul(base_batch_size) <= limit) + .unwrap_or(false) + }; + Self { + page_byte_limit, + max_def_level: descr.max_def_level(), + static_always_fits: static_fits(page_byte_limit), + dict_page_byte_limit, + static_dict_always_fits: static_fits(dict_page_byte_limit), + } + } + + /// Decide how many levels at the start of a chunk belong in one + /// mini-batch, so the mini-batch cannot overflow whichever page is + /// currently accumulating value bytes: the data page when plain-encoding, + /// or the *dictionary* page while dictionary-encoding. A returned value + /// smaller than `chunk_size` triggers granular sub-batching in + /// `write_batch_internal`. + /// + /// While dictionary-encoding, the data page holds only small RLE indices, + /// but the dictionary page accumulates the distinct values themselves — + /// so it is the dictionary page's remaining budget that must bound the + /// mini-batch. The per-mini-batch dictionary spill check would otherwise + /// let one mini-batch of large values balloon the dictionary page. + /// + /// Returns `chunk_size` immediately (no value inspection) when the chunk + /// is empty, or when the column is a fixed-width type whose mini-batches + /// statically cannot overshoot the relevant page. + /// + /// `#[inline]`: this is a tiny per-chunk dispatcher; the actual byte + /// inspection lives in the out-of-line `byte_budget_sub_batch_size`. + #[inline] + pub(crate) fn pick_sub_batch_size( + &self, + encoder: &E, + values: &E::Values, + value_indices: Option<&[usize]>, + chunk_def: LevelDataRef<'_>, + values_offset: usize, + chunk_size: usize, + ) -> usize { + if chunk_size == 0 { + return chunk_size; + } + let budget = if encoder.has_dictionary() { + if self.static_dict_always_fits { + return chunk_size; + } + // Bound the mini-batch by the dictionary page's *remaining* + // budget (it accumulates across mini-batches until it spills). + match encoder.estimated_dict_page_size() { + Some(used) => self.dict_page_byte_limit.saturating_sub(used), + None => return chunk_size, + } + } else { + if self.static_always_fits { + return chunk_size; + } + self.page_byte_limit + }; + self.byte_budget_sub_batch_size::( + values, + value_indices, + chunk_def, + values_offset, + chunk_size, + budget, + ) + } + + /// Inspect value sizes to decide how many of the chunk's values fit in + /// `budget` bytes (the data page or dictionary page remaining budget). + /// + /// `#[inline(never)]` keeps this slow path out of the hot + /// `write_batch_internal` loop; numeric and bool columns never reach it. + #[inline(never)] + fn byte_budget_sub_batch_size( + &self, + values: &E::Values, + value_indices: Option<&[usize]>, + chunk_def: LevelDataRef<'_>, + values_offset: usize, + chunk_size: usize, + budget: usize, + ) -> usize { + // How many of this chunk's levels carry an actual value. For a + // non-nullable, unrepeated column every level is a value, so + // `value_count` is O(1) (`Absent`/`Uniform` def levels); only + // nullable or nested columns pay the O(chunk_size) def-level scan. + let vals_in_chunk = chunk_def.value_count(chunk_size, self.max_def_level); + if vals_in_chunk == 0 { + return chunk_size; + } + // Ask the encoder how many of the next values fit in one page byte + // budget. Dispatch on whether the caller supplied gather indices; + // this mirrors how `write_mini_batch` picks `write_gather` vs + // `write`. + let fit = match value_indices { + Some(idx) => { + let end = (values_offset + vals_in_chunk).min(idx.len()); + let start = values_offset.min(end); + E::count_values_within_byte_budget_gather(values, &idx[start..end], budget) + } + None => { + E::count_values_within_byte_budget(values, values_offset, vals_in_chunk, budget) + } + }; + match fit { + None => chunk_size, + Some(values_per_subbatch) => { + // Convert the value count back into a level count. For a + // non-nullable column this is a no-op; for nullable/nested + // columns scale by the chunk's observed value-to-level + // ratio. + let levels_per_subbatch = if vals_in_chunk == chunk_size { + values_per_subbatch + } else { + (values_per_subbatch * chunk_size) + .div_ceil(vals_in_chunk) + .max(1) + }; + chunk_size.min(levels_per_subbatch.max(1)) + } + } + } +} diff --git a/vendor/parquet/src/column/writer/encoder.rs b/vendor/parquet/src/column/writer/encoder.rs new file mode 100644 index 0000000..b3eaae3 --- /dev/null +++ b/vendor/parquet/src/column/writer/encoder.rs @@ -0,0 +1,581 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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 bytes::Bytes; +use half::f16; + +use crate::basic::{ConvertedType, Encoding, LogicalType, Type}; +use crate::bloom_filter::Sbbf; +use crate::column::writer::{ + compare_greater, fallback_encoding, has_dictionary_support, is_nan, update_max, update_min, +}; +use crate::data_type::DataType; +use crate::data_type::private::ParquetValueType; +use crate::encodings::encoding::{DictEncoder, Encoder, get_encoder}; +use crate::errors::{ParquetError, Result}; +use crate::file::properties::{EnabledStatistics, WriterProperties}; +use crate::geospatial::accumulator::{GeoStatsAccumulator, try_new_geo_stats_accumulator}; +use crate::geospatial::statistics::GeospatialStatistics; +use crate::schema::types::{ColumnDescPtr, ColumnDescriptor}; + +/// A collection of [`ParquetValueType`] encoded by a [`ColumnValueEncoder`] +pub trait ColumnValues { + /// The number of values in this collection + fn len(&self) -> usize; +} + +#[cfg(feature = "arrow")] +impl ColumnValues for dyn arrow_array::Array { + fn len(&self) -> usize { + arrow_array::Array::len(self) + } +} + +impl ColumnValues for [T] { + fn len(&self) -> usize { + self.len() + } +} + +/// The encoded data for a dictionary page +pub struct DictionaryPage { + pub buf: Bytes, + pub num_values: usize, + pub is_sorted: bool, +} + +/// The encoded values for a data page, with optional statistics +pub struct DataPageValues { + pub buf: Bytes, + pub num_values: usize, + pub encoding: Encoding, + pub min_value: Option, + pub max_value: Option, + pub variable_length_bytes: Option, +} + +/// A generic encoder of [`ColumnValues`] to data and dictionary pages used by +/// [super::GenericColumnWriter`] +pub trait ColumnValueEncoder { + /// The underlying value type of [`Self::Values`] + /// + /// Note: this avoids needing to fully qualify `::T` + type T: ParquetValueType; + + /// The values encoded by this encoder + type Values: ColumnValues + ?Sized; + + /// Create a new [`ColumnValueEncoder`] + fn try_new(descr: &ColumnDescPtr, props: &WriterProperties) -> Result + where + Self: Sized; + + /// Write the corresponding values to this [`ColumnValueEncoder`] + fn write(&mut self, values: &Self::Values, offset: usize, len: usize) -> Result<()>; + + /// Write the values at the indexes in `indices` to this [`ColumnValueEncoder`] + fn write_gather(&mut self, values: &Self::Values, indices: &[usize]) -> Result<()>; + + /// Returns the largest `k` such that the first `k` values in + /// `values[offset..offset + len]` encode to at most `byte_budget` + /// bytes — i.e. how many values fit in a single page byte budget. + /// + /// Returns `len` if every value fits. Returns at least 1 if a single + /// value alone exceeds the budget, matching parquet's "at least one + /// value per data page" rule. + /// + /// `None` means "no cheap estimate available"; the caller stays on + /// the batched fast path and lets the post-write + /// `should_add_data_page` check handle bounding. + /// + /// Implementations should short-circuit aggressively: the typical + /// case is "everything fits, return `len`", and the next-most-common + /// case is "one wide value, return 1." The variable-width walk only + /// needs to be precise when the chunk is genuinely near the budget. + fn count_values_within_byte_budget( + _values: &Self::Values, + _offset: usize, + _len: usize, + _byte_budget: usize, + ) -> Option { + None + } + + /// As [`Self::count_values_within_byte_budget`] but using gather + /// `indices` rather than a contiguous range. Returns the number of + /// `indices` that fit, not the maximum index value. + fn count_values_within_byte_budget_gather( + _values: &Self::Values, + _indices: &[usize], + _byte_budget: usize, + ) -> Option { + None + } + + /// Returns the number of buffered values + fn num_values(&self) -> usize; + + /// Returns true if this encoder has a dictionary page + fn has_dictionary(&self) -> bool; + + /// Returns true if the encoder compresses each value against the value + /// immediately before it, within the current page. + /// + /// For such encodings a page boundary is not free: flushing discards the + /// previous value, so the first value of the next page is stored in full. + /// [`GenericColumnWriter::should_add_data_page`] uses this to exempt a + /// page's mandatory first value from the data page byte limit. + /// + /// Per encoding: + /// * `DELTA_BYTE_ARRAY`: true. Each value is stored as the length of the + /// prefix it shares with its predecessor plus the remaining suffix. + /// * Everything else: false, the default. `PLAIN` and + /// `DELTA_LENGTH_BYTE_ARRAY` store a value at the same cost wherever it + /// lands, and a dictionary outlives the pages that index into it, so no + /// page boundary makes a value more expensive. + /// + /// [`GenericColumnWriter::should_add_data_page`]: crate::column::writer::GenericColumnWriter::should_add_data_page + fn compresses_against_previous_value(&self) -> bool { + false + } + + /// Returns the estimated total memory usage of the encoder + /// + fn estimated_memory_size(&self) -> usize; + + /// Returns an estimate of the encoded size of dictionary page size in bytes, or `None` if no dictionary + fn estimated_dict_page_size(&self) -> Option; + + /// Returns an estimate of the encoded data page size in bytes + /// + /// This should include: + /// + + fn estimated_data_page_size(&self) -> usize; + + /// Flush the dictionary page for this column chunk if any. Any subsequent calls to + /// [`Self::write`] will not be dictionary encoded + /// + /// Note: [`Self::flush_data_page`] must be called first, as this will error if there + /// are any pending page values + fn flush_dict_page(&mut self) -> Result>; + + /// Flush the next data page for this column chunk + fn flush_data_page(&mut self) -> Result>; + + /// Flushes bloom filter if enabled and returns it, otherwise returns `None`. Subsequent writes + /// will *not* be tracked by the bloom filter as it is empty since. This should be called once + /// near the end of encoding. + fn flush_bloom_filter(&mut self) -> Option; + + /// Computes [`GeospatialStatistics`], if any, and resets internal state such that any internal + /// accumulator is prepared to accumulate statistics for the next column chunk. + fn flush_geospatial_statistics(&mut self) -> Option>; +} + +pub struct ColumnValueEncoderImpl { + encoder: Box>, + dict_encoder: Option>, + descr: ColumnDescPtr, + num_values: usize, + statistics_enabled: EnabledStatistics, + min_value: Option, + max_value: Option, + bloom_filter: Option, + bloom_filter_target_fpp: f64, + variable_length_bytes: Option, + geo_stats_accumulator: Option>, +} + +impl ColumnValueEncoderImpl { + fn min_max(&self, values: &[T::T], value_indices: Option<&[usize]>) -> Option<(T::T, T::T)> { + match value_indices { + Some(indices) => get_min_max(&self.descr, indices.iter().map(|x| &values[*x])), + None => get_min_max(&self.descr, values.iter()), + } + } + + fn write_slice(&mut self, slice: &[T::T]) -> Result<()> { + if self.statistics_enabled != EnabledStatistics::None + // INTERVAL, Geometry, and Geography have undefined sort order, so don't write min/max stats for them + && self.descr.converted_type() != ConvertedType::INTERVAL + { + if let Some(accumulator) = self.geo_stats_accumulator.as_deref_mut() { + update_geo_stats_accumulator(accumulator, slice.iter()); + } else if let Some((min, max)) = self.min_max(slice, None) { + update_min(&self.descr, &min, &mut self.min_value); + update_max(&self.descr, &max, &mut self.max_value); + } + + if let Some(var_bytes) = T::T::variable_length_bytes(slice) { + *self.variable_length_bytes.get_or_insert(0) += var_bytes; + } + } + + // encode the values into bloom filter if enabled + if let Some(bloom_filter) = &mut self.bloom_filter { + for value in slice { + bloom_filter.insert(value); + } + } + + match &mut self.dict_encoder { + Some(encoder) => encoder.put(slice), + _ => self.encoder.put(slice), + } + } +} + +impl ColumnValueEncoder for ColumnValueEncoderImpl { + type T = T::T; + + type Values = [T::T]; + + fn flush_bloom_filter(&mut self) -> Option { + let mut sbbf = self.bloom_filter.take()?; + sbbf.fold_to_target_fpp(self.bloom_filter_target_fpp); + Some(sbbf) + } + + fn try_new(descr: &ColumnDescPtr, props: &WriterProperties) -> Result { + let dict_supported = props.dictionary_enabled(descr.path()) + && has_dictionary_support(T::get_physical_type(), props); + let dict_encoder = dict_supported.then(|| DictEncoder::new(descr.clone())); + + // Set either main encoder or fallback encoder. + let encoder = get_encoder( + props + .encoding(descr.path()) + .unwrap_or_else(|| fallback_encoding(T::get_physical_type(), props)), + descr, + )?; + + let statistics_enabled = props.statistics_enabled(descr.path()); + + let (bloom_filter, bloom_filter_target_fpp) = create_bloom_filter(props, descr)?; + + let geo_stats_accumulator = try_new_geo_stats_accumulator(descr); + + Ok(Self { + encoder, + dict_encoder, + descr: descr.clone(), + num_values: 0, + statistics_enabled, + bloom_filter, + bloom_filter_target_fpp, + min_value: None, + max_value: None, + variable_length_bytes: None, + geo_stats_accumulator, + }) + } + + fn write(&mut self, values: &[T::T], offset: usize, len: usize) -> Result<()> { + self.num_values += len; + + let slice = values.get(offset..offset + len).ok_or_else(|| { + general_err!( + "Expected to write {} values, but have only {}", + len, + values.len() - offset + ) + })?; + + self.write_slice(slice) + } + + fn write_gather(&mut self, values: &Self::Values, indices: &[usize]) -> Result<()> { + self.num_values += indices.len(); + let slice: Vec<_> = indices.iter().map(|idx| values[*idx].clone()).collect(); + self.write_slice(&slice) + } + + fn count_values_within_byte_budget( + values: &[T::T], + offset: usize, + len: usize, + byte_budget: usize, + ) -> Option { + // Clamp so that a caller-supplied `len` that overruns the input + // (e.g. a level/value mismatch the encoder will reject later) + // returns an estimate instead of panicking here. + let end = (offset + len).min(values.len()); + let start = offset.min(end); + count_within_budget::( + end - start, + byte_budget, + values[start..end].iter().map(Some), + ) + } + + fn count_values_within_byte_budget_gather( + values: &[T::T], + indices: &[usize], + byte_budget: usize, + ) -> Option { + // `values.get` yields `None` for an out-of-range index (defensive + // against a level/value mismatch the encoder rejects later); such a + // position is counted but contributes no bytes. + count_within_budget::( + indices.len(), + byte_budget, + indices.iter().map(|&i| values.get(i)), + ) + } + + fn num_values(&self) -> usize { + self.num_values + } + + fn has_dictionary(&self) -> bool { + self.dict_encoder.is_some() + } + + fn compresses_against_previous_value(&self) -> bool { + // While dictionary encoding is active `self.encoder` is unused: the + // data page holds RLE indices, which carry no cross-value state. + self.dict_encoder.is_none() && self.encoder.encoding() == Encoding::DELTA_BYTE_ARRAY + } + + fn estimated_memory_size(&self) -> usize { + let encoder_size = self.encoder.estimated_memory_size(); + + let dict_encoder_size = self + .dict_encoder + .as_ref() + .map(|encoder| encoder.estimated_memory_size()) + .unwrap_or_default(); + + let bloom_filter_size = self + .bloom_filter + .as_ref() + .map(|bf| bf.estimated_memory_size()) + .unwrap_or_default(); + + encoder_size + dict_encoder_size + bloom_filter_size + } + + fn estimated_dict_page_size(&self) -> Option { + Some(self.dict_encoder.as_ref()?.dict_encoded_size()) + } + + fn estimated_data_page_size(&self) -> usize { + match &self.dict_encoder { + Some(encoder) => encoder.estimated_data_encoded_size(), + _ => self.encoder.estimated_data_encoded_size(), + } + } + + fn flush_dict_page(&mut self) -> Result> { + match self.dict_encoder.take() { + Some(encoder) => { + if self.num_values != 0 { + return Err(general_err!( + "Must flush data pages before flushing dictionary" + )); + } + + let buf = encoder.write_dict()?; + + Ok(Some(DictionaryPage { + buf, + num_values: encoder.num_entries(), + is_sorted: encoder.is_sorted(), + })) + } + _ => Ok(None), + } + } + + fn flush_data_page(&mut self) -> Result> { + let (buf, encoding) = match &mut self.dict_encoder { + Some(encoder) => (encoder.write_indices()?, Encoding::RLE_DICTIONARY), + _ => (self.encoder.flush_buffer()?, self.encoder.encoding()), + }; + + Ok(DataPageValues { + buf, + encoding, + num_values: std::mem::take(&mut self.num_values), + min_value: self.min_value.take(), + max_value: self.max_value.take(), + variable_length_bytes: self.variable_length_bytes.take(), + }) + } + + fn flush_geospatial_statistics(&mut self) -> Option> { + self.geo_stats_accumulator.as_mut().map(|a| a.finish())? + } +} + +fn get_min_max<'a, T, I>(descr: &ColumnDescriptor, mut iter: I) -> Option<(T, T)> +where + T: ParquetValueType + 'a, + I: Iterator, +{ + let first = loop { + let next = iter.next()?; + if !is_nan(descr, next) { + break next; + } + }; + + let mut min = first; + let mut max = first; + for val in iter { + if is_nan(descr, val) { + continue; + } + if compare_greater(descr, min, val) { + min = val; + } + if compare_greater(descr, val, max) { + max = val; + } + } + + // Float/Double statistics have special case for zero. + // + // If computed min is zero, whether negative or positive, + // the spec states that the min should be written as -0.0 + // (negative zero) + // + // For max, it has similar logic but will be written as 0.0 + // (positive zero) + let min = replace_zero(min, descr, -0.0); + let max = replace_zero(max, descr, 0.0); + + Some((min, max)) +} + +#[inline] +fn replace_zero(val: &T, descr: &ColumnDescriptor, replace: f32) -> T { + match T::PHYSICAL_TYPE { + Type::FLOAT if f32::from_le_bytes(val.as_bytes().try_into().unwrap()) == 0.0 => { + T::try_from_le_slice(&f32::to_le_bytes(replace)).unwrap() + } + Type::DOUBLE if f64::from_le_bytes(val.as_bytes().try_into().unwrap()) == 0.0 => { + T::try_from_le_slice(&f64::to_le_bytes(replace as f64)).unwrap() + } + Type::FIXED_LEN_BYTE_ARRAY + if descr.logical_type_ref() == Some(LogicalType::Float16).as_ref() + && f16::from_le_bytes(val.as_bytes().try_into().unwrap()) == f16::NEG_ZERO => + { + T::try_from_le_slice(&f16::to_le_bytes(f16::from_f32(replace))).unwrap() + } + _ => val.clone(), + } +} + +/// Creates a bloom filter sized for the column's configured NDV, returning the filter +/// and the target FPP for folding. +pub(crate) fn create_bloom_filter( + props: &WriterProperties, + descr: &ColumnDescPtr, +) -> Result<(Option, f64)> { + match props.bloom_filter_properties(descr.path()) { + Some(bf_props) => Ok(( + Some(Sbbf::new_with_ndv_fpp(bf_props.ndv(), bf_props.fpp())?), + bf_props.fpp(), + )), + None => Ok((None, 0.0)), + } +} + +fn update_geo_stats_accumulator<'a, T, I>(bounder: &mut dyn GeoStatsAccumulator, iter: I) +where + T: ParquetValueType + 'a, + I: Iterator, +{ + if bounder.is_valid() { + for val in iter { + bounder.update_wkb(val.as_bytes()); + } + } +} + +/// Plain-encoded byte cost of a single value of type `T::T`. +/// +/// Derived from [`ParquetValueType::dict_encoding_size`] (which returns +/// `(per-value overhead, value-bytes)`) so we don't add a parallel +/// per-value-size hook to the trait. Mirrors the dispatch in +/// `KeyStorage::push` (`encodings/encoding/dict_encoder.rs`). +/// +/// Placed at the end of the module deliberately. Inserting it above the +/// `ColumnValueEncoder` trait shifts the trait and `ColumnValueEncoderImpl` +/// within the compiled module enough to perturb downstream code placement, +/// which measurably regresses unrelated arrow-writer string benchmarks +/// (~5-9% on `string` / `string_and_binary_view`). Defining it last keeps +/// the hot encoder code at the offsets it has on `main`. +#[inline] +fn plain_encoded_byte_size(value: &T::T) -> usize { + let (overhead, bytes) = value.dict_encoding_size(); + match ::PHYSICAL_TYPE { + // Plain BYTE_ARRAY = 4-byte length prefix + payload. + Type::BYTE_ARRAY => overhead + bytes, + // Plain FLBA = raw bytes only; `dict_encoding_size`'s length prefix + // is irrelevant here, so the encoder passes `type_length` directly. + Type::FIXED_LEN_BYTE_ARRAY => bytes, + // Numeric/bool are short-circuited by the caller via + // `mem::size_of`, so this is unreachable in practice; fall back to + // `overhead` defensively. + _ => overhead, + } +} + +/// How many leading values fit in `byte_budget` bytes, shared by the two +/// `ColumnValueEncoder::count_values_within_byte_budget*` methods (one walks a +/// contiguous slice, the other gathers by index). +/// +/// `n` is the answer when everything fits; `vals` yields each candidate value, +/// or `None` for a position that should still be counted but contributes no +/// bytes (an out-of-range gather index). The boundary value that crosses the +/// budget is included in the count so the caller's page-flush check trips on +/// this mini-batch rather than leaving a sliver for the next page; this also +/// catches a lone outlier wherever it lands among small values. +/// +/// Defined at the end of the module alongside `plain_encoded_byte_size` for +/// the same reason — see that function's note on code placement and the +/// `string` / `string_and_binary_view` benchmarks. +#[inline] +fn count_within_budget<'a, T: DataType>( + n: usize, + byte_budget: usize, + vals: impl Iterator>, +) -> Option +where + T::T: 'a, +{ + // Fixed-size physical types have a constant per-value byte cost, so the + // answer is one division — no walk needed. + let phys = ::PHYSICAL_TYPE; + if phys != Type::BYTE_ARRAY && phys != Type::FIXED_LEN_BYTE_ARRAY { + let per = std::mem::size_of::().max(1); + return Some((byte_budget / per).max(1).min(n)); + } + // Variable-width: accumulate, exit at the first value past the budget. + let mut cum: usize = 0; + for (i, v) in vals.enumerate() { + if let Some(v) = v { + cum = cum.saturating_add(plain_encoded_byte_size::(v)); + } + if cum > byte_budget { + return Some(i + 1); + } + } + Some(n) +} diff --git a/vendor/parquet/src/column/writer/mod.rs b/vendor/parquet/src/column/writer/mod.rs new file mode 100644 index 0000000..b75dbc2 --- /dev/null +++ b/vendor/parquet/src/column/writer/mod.rs @@ -0,0 +1,5552 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Contains column writer API. + +use bytes::Bytes; +use half::f16; + +use crate::bloom_filter::Sbbf; +use crate::file::page_index::column_index::ColumnIndexMetaData; +use crate::file::page_index::offset_index::OffsetIndexMetaData; +use std::collections::{BTreeSet, VecDeque}; +use std::str; + +use crate::basic::{ + BoundaryOrder, Compression, ConvertedType, Encoding, EncodingMask, IntType, LogicalType, + PageType, Type, +}; +use crate::column::page::{CompressedPage, Page, PageWriteSpec, PageWriter}; +use crate::column::writer::encoder::{ColumnValueEncoder, ColumnValueEncoderImpl, ColumnValues}; +use crate::compression::{Codec, CodecOptionsBuilder, create_codec}; +use crate::data_type::private::ParquetValueType; +use crate::data_type::*; +use crate::encodings::levels::LevelEncoder; +#[cfg(feature = "encryption")] +use crate::encryption::encrypt::get_column_crypto_metadata; +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::{ + ColumnChunkMetaData, ColumnChunkMetaDataBuilder, ColumnIndexBuilder, LevelHistogram, + OffsetIndexBuilder, PageEncodingStats, +}; +use crate::file::properties::{ + EnabledStatistics, WriterProperties, WriterPropertiesPtr, WriterVersion, +}; +use crate::file::statistics::{Statistics, ValueStatistics}; +use crate::schema::types::{ColumnDescPtr, ColumnDescriptor}; + +mod byte_budget_chunker; +pub(crate) mod encoder; + +use byte_budget_chunker::ByteBudgetChunker; + +macro_rules! downcast_writer { + ($e:expr, $i:ident, $b:expr) => { + match $e { + Self::BoolColumnWriter($i) => $b, + Self::Int32ColumnWriter($i) => $b, + Self::Int64ColumnWriter($i) => $b, + Self::Int96ColumnWriter($i) => $b, + Self::FloatColumnWriter($i) => $b, + Self::DoubleColumnWriter($i) => $b, + Self::ByteArrayColumnWriter($i) => $b, + Self::FixedLenByteArrayColumnWriter($i) => $b, + } + }; +} + +/// Column writer for a Parquet type. +/// +/// See [`get_column_writer`] to create instances of this type +pub enum ColumnWriter<'a> { + /// Column writer for boolean type + BoolColumnWriter(ColumnWriterImpl<'a, BoolType>), + /// Column writer for int32 type + Int32ColumnWriter(ColumnWriterImpl<'a, Int32Type>), + /// Column writer for int64 type + Int64ColumnWriter(ColumnWriterImpl<'a, Int64Type>), + /// Column writer for int96 (timestamp) type + Int96ColumnWriter(ColumnWriterImpl<'a, Int96Type>), + /// Column writer for float type + FloatColumnWriter(ColumnWriterImpl<'a, FloatType>), + /// Column writer for double type + DoubleColumnWriter(ColumnWriterImpl<'a, DoubleType>), + /// Column writer for byte array type + ByteArrayColumnWriter(ColumnWriterImpl<'a, ByteArrayType>), + /// Column writer for fixed length byte array type + FixedLenByteArrayColumnWriter(ColumnWriterImpl<'a, FixedLenByteArrayType>), +} + +impl ColumnWriter<'_> { + /// Returns the estimated total memory usage + #[cfg(feature = "arrow")] + pub(crate) fn memory_size(&self) -> usize { + downcast_writer!(self, typed, typed.memory_size()) + } + + /// Returns the estimated total encoded bytes for this column writer + #[cfg(feature = "arrow")] + pub(crate) fn get_estimated_total_bytes(&self) -> u64 { + downcast_writer!(self, typed, typed.get_estimated_total_bytes()) + } + + /// Finalize the currently buffered values as a data page. + /// + /// This is used by content-defined chunking to force a page boundary at + /// content-determined positions. + #[cfg(feature = "arrow")] + pub(crate) fn add_data_page(&mut self) -> Result<()> { + downcast_writer!(self, typed, typed.add_data_page()) + } + + /// Close this [`ColumnWriter`], returning the metadata for the column chunk. + pub fn close(self) -> Result { + downcast_writer!(self, typed, typed.close()) + } +} + +/// Create a specific column writer corresponding to column descriptor `descr`. +pub fn get_column_writer<'a>( + descr: ColumnDescPtr, + props: WriterPropertiesPtr, + page_writer: Box, +) -> ColumnWriter<'a> { + match descr.physical_type() { + Type::BOOLEAN => { + ColumnWriter::BoolColumnWriter(ColumnWriterImpl::new(descr, props, page_writer)) + } + Type::INT32 => { + ColumnWriter::Int32ColumnWriter(ColumnWriterImpl::new(descr, props, page_writer)) + } + Type::INT64 => { + ColumnWriter::Int64ColumnWriter(ColumnWriterImpl::new(descr, props, page_writer)) + } + Type::INT96 => { + ColumnWriter::Int96ColumnWriter(ColumnWriterImpl::new(descr, props, page_writer)) + } + Type::FLOAT => { + ColumnWriter::FloatColumnWriter(ColumnWriterImpl::new(descr, props, page_writer)) + } + Type::DOUBLE => { + ColumnWriter::DoubleColumnWriter(ColumnWriterImpl::new(descr, props, page_writer)) + } + Type::BYTE_ARRAY => { + ColumnWriter::ByteArrayColumnWriter(ColumnWriterImpl::new(descr, props, page_writer)) + } + Type::FIXED_LEN_BYTE_ARRAY => ColumnWriter::FixedLenByteArrayColumnWriter( + ColumnWriterImpl::new(descr, props, page_writer), + ), + } +} + +/// Gets a typed column writer for the specific type `T`, by "up-casting" `col_writer` of +/// non-generic type to a generic column writer type `ColumnWriterImpl`. +/// +/// Panics if actual enum value for `col_writer` does not match the type `T`. +pub fn get_typed_column_writer(col_writer: ColumnWriter) -> ColumnWriterImpl { + T::get_column_writer(col_writer).unwrap_or_else(|| { + panic!( + "Failed to convert column writer into a typed column writer for `{}` type", + T::get_physical_type() + ) + }) +} + +/// Similar to `get_typed_column_writer` but returns a reference. +pub fn get_typed_column_writer_ref<'a, 'b: 'a, T: DataType>( + col_writer: &'b ColumnWriter<'a>, +) -> &'b ColumnWriterImpl<'a, T> { + T::get_column_writer_ref(col_writer).unwrap_or_else(|| { + panic!( + "Failed to convert column writer into a typed column writer for `{}` type", + T::get_physical_type() + ) + }) +} + +/// Similar to `get_typed_column_writer` but returns a reference. +pub fn get_typed_column_writer_mut<'a, 'b: 'a, T: DataType>( + col_writer: &'a mut ColumnWriter<'b>, +) -> &'a mut ColumnWriterImpl<'b, T> { + T::get_column_writer_mut(col_writer).unwrap_or_else(|| { + panic!( + "Failed to convert column writer into a typed column writer for `{}` type", + T::get_physical_type() + ) + }) +} + +/// Metadata for a column chunk of a Parquet file. +/// +/// Note this structure is returned by [`ColumnWriter::close`]. +#[derive(Debug, Clone)] +pub struct ColumnCloseResult { + /// The total number of bytes written + pub bytes_written: u64, + /// The total number of rows written + pub rows_written: u64, + /// Metadata for this column chunk + pub metadata: ColumnChunkMetaData, + /// Optional bloom filter for this column + pub bloom_filter: Option, + /// Optional column index, for filtering + pub column_index: Option, + /// Optional offset index, identifying page locations + pub offset_index: Option, +} + +impl ColumnCloseResult { + /// Rewrite the page offsets for a dictionary-first on-disk layout. + /// + /// A writer that buffers the whole column chunk and splices it later (the + /// Arrow path) may accept the data pages *before* the dictionary page so the + /// data pages can stream straight through, then emit the dictionary page + /// first at splice. The offsets recorded during encoding therefore assume a + /// data-pages-first layout; call this with the serialized length of the + /// dictionary page to move it to offset 0 and shift every data page after + /// it. A `dictionary_len` of 0 (no dictionary page) leaves the result + /// unchanged. + pub fn update_dictionary_location(mut self, dictionary_len: usize) -> Result { + if dictionary_len > 0 { + self.metadata = self + .metadata + .into_builder() + .set_dictionary_page_offset(Some(0)) + .set_data_page_offset(dictionary_len as i64) + .build()?; + if let Some(offset_index) = self.offset_index.as_mut() { + let mut offset = dictionary_len as i64; + for location in offset_index.page_locations.iter_mut() { + location.offset = offset; + offset += location.compressed_page_size as i64; + } + } + } + Ok(self) + } +} + +// Metrics per page +#[derive(Default)] +struct PageMetrics { + num_buffered_values: u32, + num_buffered_rows: u32, + /// Encoded bytes that the data page byte limit does not apply to, + /// because they belong to the page's mandatory first value and cannot be + /// moved elsewhere. Zero unless that value alone exceeded the limit + /// *and* the encoding compresses against the preceding value; see + /// [`ColumnValueEncoder::compresses_against_previous_value`]. + page_size_exemption: usize, + num_page_nulls: u64, + repetition_level_histogram: Option, + definition_level_histogram: Option, +} + +impl PageMetrics { + fn new() -> Self { + Default::default() + } + + /// Initialize the repetition level histogram + fn with_repetition_level_histogram(mut self, max_level: i16) -> Self { + self.repetition_level_histogram = LevelHistogram::try_new(max_level); + self + } + + /// Initialize the definition level histogram + fn with_definition_level_histogram(mut self, max_level: i16) -> Self { + self.definition_level_histogram = LevelHistogram::try_new(max_level); + self + } + + /// Resets the state of this `PageMetrics` to the initial state. + /// If histograms have been initialized their contents will be reset to zero. + fn new_page(&mut self) { + self.num_buffered_values = 0; + self.num_buffered_rows = 0; + self.page_size_exemption = 0; + self.num_page_nulls = 0; + self.repetition_level_histogram + .as_mut() + .map(LevelHistogram::reset); + self.definition_level_histogram + .as_mut() + .map(LevelHistogram::reset); + } +} + +// Metrics per column writer +#[derive(Default)] +struct ColumnMetrics { + total_bytes_written: u64, + total_rows_written: u64, + total_uncompressed_size: u64, + total_compressed_size: u64, + total_num_values: u64, + dictionary_page_offset: Option, + data_page_offset: Option, + min_column_value: Option, + max_column_value: Option, + num_column_nulls: u64, + column_distinct_count: Option, + variable_length_bytes: Option, + repetition_level_histogram: Option, + definition_level_histogram: Option, +} + +impl ColumnMetrics { + fn new() -> Self { + Default::default() + } + + /// Initialize the repetition level histogram + fn with_repetition_level_histogram(mut self, max_level: i16) -> Self { + self.repetition_level_histogram = LevelHistogram::try_new(max_level); + self + } + + /// Initialize the definition level histogram + fn with_definition_level_histogram(mut self, max_level: i16) -> Self { + self.definition_level_histogram = LevelHistogram::try_new(max_level); + self + } + + /// Sum `page_histogram` into `chunk_histogram` + fn update_histogram( + chunk_histogram: &mut Option, + page_histogram: &Option, + ) { + if let (Some(page_hist), Some(chunk_hist)) = (page_histogram, chunk_histogram) { + chunk_hist.add(page_hist); + } + } + + /// Sum the provided PageMetrics histograms into the chunk histograms. Does nothing if + /// page histograms are not initialized. + fn update_from_page_metrics(&mut self, page_metrics: &PageMetrics) { + ColumnMetrics::::update_histogram( + &mut self.definition_level_histogram, + &page_metrics.definition_level_histogram, + ); + ColumnMetrics::::update_histogram( + &mut self.repetition_level_histogram, + &page_metrics.repetition_level_histogram, + ); + } + + /// Sum the provided page variable_length_bytes into the chunk variable_length_bytes + fn update_variable_length_bytes(&mut self, variable_length_bytes: Option) { + if let Some(var_bytes) = variable_length_bytes { + *self.variable_length_bytes.get_or_insert(0) += var_bytes; + } + } +} + +/// Borrowed view of level data, analogous to `&str` for `LevelData`'s `String`. +/// +/// `LevelDataRef` can be constructed from `LevelData` and directly from an existing +/// `&[i16]` without allocating. +/// +/// The variants are different physical representations of the same logical +/// sequence of levels. +#[derive(Debug, Clone, Copy)] +pub(crate) enum LevelDataRef<'a> { + Absent, + Materialized(&'a [i16]), + Uniform { value: i16, count: usize }, +} + +impl<'a> From<&'a [i16]> for LevelDataRef<'a> { + fn from(levels: &'a [i16]) -> Self { + Self::Materialized(levels) + } +} + +impl<'a> From> for LevelDataRef<'a> { + fn from(levels: Option<&'a [i16]>) -> Self { + levels.map_or(Self::Absent, Self::from) + } +} + +impl<'a> LevelDataRef<'a> { + pub(crate) fn len(self) -> usize { + match self { + Self::Absent => 0, + Self::Materialized(values) => values.len(), + Self::Uniform { count, .. } => count, + } + } + + pub(crate) fn first(self) -> Option { + match self { + Self::Absent => None, + Self::Materialized(values) => values.first().copied(), + Self::Uniform { value, count } => (count > 0).then_some(value), + } + } + + #[cfg(feature = "arrow")] + pub(crate) fn value_at(self, idx: usize) -> Option { + match self { + Self::Absent => None, + Self::Materialized(values) => values.get(idx).copied(), + Self::Uniform { value, count } => (idx < count).then_some(value), + } + } + + pub(crate) fn slice(self, offset: usize, len: usize) -> Self { + match self { + Self::Absent => Self::Absent, + Self::Materialized(values) => Self::Materialized(&values[offset..offset + len]), + Self::Uniform { value, .. } => Self::Uniform { value, count: len }, + } + } + + /// Count of positions in this slice that represent an actual value + /// (definition level equal to `max_def`). `Absent` means the column has + /// `max_def == 0` and every position is a value, so the implicit count + /// is the caller-supplied `total`. + pub(crate) fn value_count(self, total: usize, max_def: i16) -> usize { + match self { + Self::Absent => total, + Self::Materialized(values) => values.iter().filter(|&&d| d == max_def).count(), + Self::Uniform { value, count } => { + if value == max_def { + count + } else { + 0 + } + } + } + } +} + +/// Typed column writer for a primitive column. +pub type ColumnWriterImpl<'a, T> = GenericColumnWriter<'a, ColumnValueEncoderImpl>; + +/// Generic column writer for a primitive Parquet column +pub struct GenericColumnWriter<'a, E: ColumnValueEncoder> { + // Column writer properties + descr: ColumnDescPtr, + props: WriterPropertiesPtr, + statistics_enabled: EnabledStatistics, + + page_writer: Box, + codec: Compression, + compressor: Option>, + encoder: E, + + page_metrics: PageMetrics, + // Metrics per column writer + column_metrics: ColumnMetrics, + + /// The order of encodings within the generated metadata does not impact its meaning, + /// but we use a BTreeSet so that the output is deterministic + encodings: BTreeSet, + encoding_stats: Vec, + // Streaming level encoders for definition/repetition levels. + def_levels_encoder: LevelEncoder, + rep_levels_encoder: LevelEncoder, + data_pages: VecDeque, + // column index and offset index + column_index_builder: ColumnIndexBuilder, + offset_index_builder: Option, + + // Below fields used to incrementally check boundary order across data pages. + // We assume they are ascending/descending until proven wrong. + data_page_boundary_ascending: bool, + data_page_boundary_descending: bool, + /// (min, max) + last_non_null_data_page_min_max: Option<(E::T, E::T)>, +} + +impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { + /// Returns a new instance of [`GenericColumnWriter`]. + pub fn new( + descr: ColumnDescPtr, + props: WriterPropertiesPtr, + page_writer: Box, + ) -> Self { + let codec = props.compression(descr.path()); + let codec_options = CodecOptionsBuilder::default().build(); + let compressor = create_codec(codec, &codec_options).unwrap(); + let encoder = E::try_new(&descr, props.as_ref()).unwrap(); + + let statistics_enabled = props.statistics_enabled(descr.path()); + + let mut encodings = BTreeSet::new(); + // Used for level information + encodings.insert(Encoding::RLE); + + let mut page_metrics = PageMetrics::new(); + let mut column_metrics = ColumnMetrics::::new(); + + // Initialize level histograms if collecting page or chunk statistics + if statistics_enabled != EnabledStatistics::None { + page_metrics = page_metrics + .with_repetition_level_histogram(descr.max_rep_level()) + .with_definition_level_histogram(descr.max_def_level()); + column_metrics = column_metrics + .with_repetition_level_histogram(descr.max_rep_level()) + .with_definition_level_histogram(descr.max_def_level()) + } + + // Disable column_index_builder if not collecting page statistics. + let mut column_index_builder = ColumnIndexBuilder::new(descr.physical_type()); + if statistics_enabled != EnabledStatistics::Page { + column_index_builder.to_invalid() + } + + // Disable offset_index_builder if requested by user. + let offset_index_builder = match props.offset_index_disabled() { + false => Some(OffsetIndexBuilder::new()), + _ => None, + }; + + Self { + def_levels_encoder: Self::create_level_encoder(descr.max_def_level(), &props), + rep_levels_encoder: Self::create_level_encoder(descr.max_rep_level(), &props), + descr, + props, + statistics_enabled, + page_writer, + codec, + compressor, + encoder, + data_pages: VecDeque::new(), + page_metrics, + column_metrics, + column_index_builder, + offset_index_builder, + encodings, + encoding_stats: vec![], + data_page_boundary_ascending: true, + data_page_boundary_descending: true, + last_non_null_data_page_min_max: None, + } + } + + #[allow(clippy::too_many_arguments)] + pub(crate) fn write_batch_internal( + &mut self, + values: &E::Values, + value_indices: Option<&[usize]>, + def_levels: LevelDataRef<'_>, + rep_levels: LevelDataRef<'_>, + min: Option<&E::T>, + max: Option<&E::T>, + distinct_count: Option, + ) -> Result { + // Check if number of definition levels is the same as number of repetition levels. + if def_levels.len() != 0 && rep_levels.len() != 0 && def_levels.len() != rep_levels.len() { + return Err(general_err!( + "Inconsistent length of definition and repetition levels: {} != {}", + def_levels.len(), + rep_levels.len() + )); + } + + // We check for DataPage limits only after we have inserted the values. If a user + // writes a large number of values, the DataPage size can be well above the limit. + // + // The purpose of this chunking is to bound this. Even if a user writes large + // number of values, the chunking will ensure that we add data page at a + // reasonable pagesize limit. + + // TODO: find out why we don't account for size of levels when we estimate page + // size. + let num_levels = def_levels.len().max(rep_levels.len()); + let num_levels = if num_levels > 0 { + num_levels + } else { + value_indices.map_or(values.len(), |i| i.len()) + }; + + if let Some(min) = min { + update_min(&self.descr, min, &mut self.column_metrics.min_column_value); + } + if let Some(max) = max { + 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; + } + + let mut values_offset = 0; + let mut levels_offset = 0; + let both_levels_compact = !matches!(def_levels, LevelDataRef::Materialized(_)) + && !matches!(rep_levels, LevelDataRef::Materialized(_)); + let has_levels = !matches!(def_levels, LevelDataRef::Absent) + || !matches!(rep_levels, LevelDataRef::Absent); + // When both level vectors are compact (Uniform or Absent), there is no + // materialized slice to split and the per-mini-batch work is O(1), so we + // can safely use a much larger batch size. + let base_batch_size = if both_levels_compact && has_levels { + self.props.data_page_row_count_limit() + } else { + self.props.write_batch_size() + }; + let chunker = ByteBudgetChunker::new(&self.descr, &self.props, base_batch_size); + while levels_offset < num_levels { + let mut end_offset = num_levels.min(levels_offset + base_batch_size); + + // Split at record boundary + if let LevelDataRef::Materialized(levels) = rep_levels { + while end_offset < levels.len() && levels[end_offset] != 0 { + end_offset += 1; + } + } + + let chunk_size = end_offset - levels_offset; + let chunk_def = def_levels.slice(levels_offset, chunk_size); + let chunk_rep = rep_levels.slice(levels_offset, chunk_size); + + // Key decision point: can we write this whole chunk as one + // mini-batch (the common case — small or fixed-width values, no + // further page-size accounting needed), or must we fall back to + // byte-budget-aware sub-batching to keep a page from overshooting + // `data_page_size_limit`? `pick_sub_batch_size` returns + // `chunk_size` for the former. + let sub_batch_size = chunker.pick_sub_batch_size( + &self.encoder, + values, + value_indices, + chunk_def, + values_offset, + chunk_size, + ); + + if sub_batch_size >= chunk_size { + values_offset += self.write_mini_batch( + values, + values_offset, + value_indices, + chunk_size, + chunk_def, + chunk_rep, + )?; + } else { + values_offset += self.write_granular_chunk( + values, + values_offset, + value_indices, + chunk_size, + chunk_def, + chunk_rep, + sub_batch_size, + )?; + } + levels_offset = end_offset; + } + + // Return total number of values processed. + Ok(values_offset) + } + + /// Writes batch of values, definition levels and repetition levels. + /// Returns number of values processed (written). + /// + /// If definition and repetition levels are provided, we write fully those levels and + /// select how many values to write (this number will be returned), since number of + /// actual written values may be smaller than provided values. + /// + /// If only values are provided, then all values are written and the length of + /// of the values buffer is returned. + /// + /// Definition and/or repetition levels can be omitted, if values are + /// non-nullable and/or non-repeated. + pub fn write_batch( + &mut self, + values: &E::Values, + def_levels: Option<&[i16]>, + rep_levels: Option<&[i16]>, + ) -> Result { + self.write_batch_internal( + values, + None, + LevelDataRef::from(def_levels), + LevelDataRef::from(rep_levels), + None, + None, + None, + ) + } + + /// Writer may optionally provide pre-calculated statistics for use when computing + /// chunk-level statistics + /// + /// NB: [`WriterProperties::statistics_enabled`] must be set to [`EnabledStatistics::Chunk`] + /// for these statistics to take effect. If [`EnabledStatistics::None`] they will be ignored, + /// and if [`EnabledStatistics::Page`] the chunk statistics will instead be computed from the + /// computed page statistics + pub fn write_batch_with_statistics( + &mut self, + values: &E::Values, + def_levels: Option<&[i16]>, + rep_levels: Option<&[i16]>, + min: Option<&E::T>, + max: Option<&E::T>, + distinct_count: Option, + ) -> Result { + self.write_batch_internal( + values, + None, + LevelDataRef::from(def_levels), + LevelDataRef::from(rep_levels), + min, + max, + distinct_count, + ) + } + + /// Returns the estimated total memory usage. + /// + /// Unlike [`Self::get_estimated_total_bytes`] this is an estimate + /// of the current memory usage and not the final anticipated encoded size. + #[cfg(feature = "arrow")] + pub(crate) fn memory_size(&self) -> usize { + // In-flight encoder buffers, plus any completed pages still held on the + // heap: the dictionary-column data pages buffered here (column-at-a-time + // path), plus whatever the page writer keeps resident. A page writer + // that spills completed pages off-heap reports far less than the bytes + // it was handed, so this tracks real memory rather than bytes written. + self.encoder.estimated_memory_size() + + self + .data_pages + .iter() + .map(|page| page.memory_usage()) + .sum::() + + self.page_writer.buffered_memory_size() + } + + /// Returns total number of bytes written by this column writer so far. + /// This value is also returned when column writer is closed. + /// + /// Note: this value does not include any buffered data that has not + /// yet been flushed to a page. + pub fn get_total_bytes_written(&self) -> u64 { + self.column_metrics.total_bytes_written + } + + /// Returns the estimated total encoded bytes for this column writer. + /// + /// Unlike [`Self::get_total_bytes_written`] this includes an estimate + /// of any data that has not yet been flushed to a page, based on it's + /// anticipated encoded size. + #[cfg(feature = "arrow")] + pub(crate) fn get_estimated_total_bytes(&self) -> u64 { + self.data_pages + .iter() + .map(|page| page.data().len() as u64) + .sum::() + + self.column_metrics.total_bytes_written + + self.encoder.estimated_data_page_size() as u64 + + self.encoder.estimated_dict_page_size().unwrap_or_default() as u64 + } + + /// Returns total number of rows written by this column writer so far. + /// This value is also returned when column writer is closed. + pub fn get_total_rows_written(&self) -> u64 { + self.column_metrics.total_rows_written + } + + /// Returns a reference to a [`ColumnDescPtr`] + pub fn get_descriptor(&self) -> &ColumnDescPtr { + &self.descr + } + + /// Finalizes writes and closes the column writer. + /// Returns total bytes written, total rows written and column chunk metadata. + pub fn close(mut self) -> Result { + if self.page_metrics.num_buffered_values > 0 { + self.add_data_page()?; + } + if self.encoder.has_dictionary() { + self.write_dictionary_page()?; + } + self.flush_data_pages()?; + let metadata = self.build_column_metadata()?; + self.page_writer.close()?; + + let boundary_order = match ( + self.data_page_boundary_ascending, + self.data_page_boundary_descending, + ) { + // If the lists are composed of equal elements then will be marked as ascending + // (Also the case if all pages are null pages) + (true, _) => BoundaryOrder::ASCENDING, + (false, true) => BoundaryOrder::DESCENDING, + (false, false) => BoundaryOrder::UNORDERED, + }; + self.column_index_builder.set_boundary_order(boundary_order); + + let column_index = match self.column_index_builder.valid() { + true => Some(self.column_index_builder.build()?), + false => None, + }; + + let offset_index = self.offset_index_builder.map(|b| b.build()); + + Ok(ColumnCloseResult { + bytes_written: self.column_metrics.total_bytes_written, + rows_written: self.column_metrics.total_rows_written, + bloom_filter: self.encoder.flush_bloom_filter(), + metadata, + column_index, + offset_index, + }) + } + + /// Writes a chunk in `sub_batch_size`-level sub-batches, checking the + /// data page byte limit after each. This keeps the page size close to + /// `data_page_size_limit` instead of overshooting it by a whole chunk. + /// + /// For repeated/nested columns sub-batches step from one `rep == 0` + /// boundary to the next so a record never spans data pages, matching + /// the parquet format rule. + /// + /// Returns the total number of values consumed across all sub-batches. + /// + /// `#[inline(never)]` keeps this slow path — only reached for + /// variable-width columns whose values need page splitting — out of + /// the hot `write_batch_internal` loop. + #[allow(clippy::too_many_arguments)] + #[inline(never)] + fn write_granular_chunk( + &mut self, + values: &E::Values, + values_offset: usize, + value_indices: Option<&[usize]>, + chunk_size: usize, + chunk_def: LevelDataRef<'_>, + chunk_rep: LevelDataRef<'_>, + sub_batch_size: usize, + ) -> Result { + // The chunker always sizes a sub-batch to at least one level, so each + // iteration below makes progress (`sub_end > sub_start`). + debug_assert!(sub_batch_size >= 1, "chunker must size at least one level"); + let mut values_consumed = 0; + let mut sub_start = 0; + while sub_start < chunk_size { + let sub_end = match chunk_rep { + LevelDataRef::Materialized(levels) => { + // Pack up to `sub_batch_size` levels per mini-batch, then + // extend to the next record boundary (rep == 0) so a + // record never spans data pages. Packing whole records + // rather than stepping one record at a time avoids + // calling `write_mini_batch` per record: records average + // only a handful of levels, so a record-at-a-time step + // would issue many more mini-batches than necessary. + let mut e = (sub_start + sub_batch_size).min(chunk_size); + while e < chunk_size && levels[e] != 0 { + e += 1; + } + e + } + _ => (sub_start + sub_batch_size).min(chunk_size), + }; + let sub_len = sub_end - sub_start; + let written = self.write_mini_batch( + values, + values_offset + values_consumed, + value_indices, + sub_len, + chunk_def.slice(sub_start, sub_len), + chunk_rep.slice(sub_start, sub_len), + )?; + values_consumed += written; + sub_start = sub_end; + } + Ok(values_consumed) + } + + /// Creates a new streaming level encoder appropriate for the writer version. + fn create_level_encoder(max_level: i16, props: &WriterProperties) -> LevelEncoder { + match props.writer_version() { + WriterVersion::PARQUET_1_0 => LevelEncoder::v1_streaming(max_level), + WriterVersion::PARQUET_2_0 => LevelEncoder::v2_streaming(max_level), + } + } + + /// Writes mini batch of values, definition and repetition levels. + /// This allows fine-grained processing of values and maintaining a reasonable + /// page size. + fn write_mini_batch( + &mut self, + values: &E::Values, + values_offset: usize, + value_indices: Option<&[usize]>, + num_levels: usize, + def_levels: LevelDataRef<'_>, + rep_levels: LevelDataRef<'_>, + ) -> Result { + // Process definition levels and determine how many values to write. + let values_to_write = if self.descr.max_def_level() > 0 { + let max_def = self.descr.max_def_level(); + match def_levels { + LevelDataRef::Absent => { + return Err(general_err!( + "Definition levels are required, because max definition level = {}", + self.descr.max_def_level() + )); + } + LevelDataRef::Materialized(levels) => { + // General path for caller-provided or already-materialized + // level buffers. + let mut values_to_write = 0usize; + let encoder = &mut self.def_levels_encoder; + match self.page_metrics.definition_level_histogram.as_mut() { + Some(histogram) => encoder.put_with_observer(levels, |level, count| { + values_to_write += count * (level == max_def) as usize; + histogram.increment_by(level, count as i64); + }), + None => encoder.put_with_observer(levels, |level, count| { + values_to_write += count * (level == max_def) as usize; + }), + }; + self.page_metrics.num_page_nulls += (levels.len() - values_to_write) as u64; + values_to_write + } + LevelDataRef::Uniform { value, count } => { + // Fast path for all-null, all-valid, or otherwise uniform + // definition levels without materializing a level buffer. + let encoder = &mut self.def_levels_encoder; + match self.page_metrics.definition_level_histogram.as_mut() { + Some(histogram) => { + encoder.put_n_with_observer(value, count, |level, run_len| { + histogram.increment_by(level, run_len as i64); + }) + } + None => encoder.put_n_with_observer(value, count, |_, _| {}), + }; + let values_to_write = count * (value == max_def) as usize; + self.page_metrics.num_page_nulls += (count - values_to_write) as u64; + values_to_write + } + } + } else { + num_levels + }; + + // Process repetition levels and determine how many rows we are about to process. + if self.descr.max_rep_level() > 0 { + // A row could contain more than one value. + let first_level = rep_levels.first().ok_or_else(|| { + general_err!( + "Repetition levels are required, because max repetition level = {}", + self.descr.max_rep_level() + ) + })?; + + if first_level != 0 { + return Err(general_err!( + "Write must start at a record boundary, got non-zero repetition level of {}", + first_level + )); + } + + let mut new_rows = 0u32; + match rep_levels { + LevelDataRef::Absent => unreachable!(), + LevelDataRef::Materialized(levels) => { + let encoder = &mut self.rep_levels_encoder; + match self.page_metrics.repetition_level_histogram.as_mut() { + Some(histogram) => encoder.put_with_observer(levels, |level, count| { + new_rows += (count as u32) * (level == 0) as u32; + histogram.increment_by(level, count as i64); + }), + None => encoder.put_with_observer(levels, |level, count| { + new_rows += (count as u32) * (level == 0) as u32; + }), + }; + } + LevelDataRef::Uniform { value, count } => { + let encoder = &mut self.rep_levels_encoder; + match self.page_metrics.repetition_level_histogram.as_mut() { + Some(histogram) => { + encoder.put_n_with_observer(value, count, |level, run_len| { + new_rows += (run_len as u32) * (level == 0) as u32; + histogram.increment_by(level, run_len as i64); + }) + } + None => encoder.put_n_with_observer(value, count, |level, run_len| { + new_rows += (run_len as u32) * (level == 0) as u32; + }), + }; + } + } + self.page_metrics.num_buffered_rows += new_rows; + } else { + // Each value is exactly one row. + // Equals to the number of values, we count nulls as well. + self.page_metrics.num_buffered_rows += num_levels as u32; + } + + match value_indices { + Some(indices) => { + let indices = &indices[values_offset..values_offset + values_to_write]; + self.encoder.write_gather(values, indices)?; + } + None => self.encoder.write(values, values_offset, values_to_write)?, + } + + let page_was_empty = self.page_metrics.num_buffered_values == 0; + self.page_metrics.num_buffered_values += num_levels as u32; + + if page_was_empty && values_to_write == 1 { + self.set_page_size_exemption(); + } + + if self.should_add_data_page() { + self.add_data_page()?; + } + + if self.should_dict_fallback() { + self.dict_fallback()?; + } + + Ok(values_to_write) + } + + /// Returns true if we need to fall back to non-dictionary encoding. + /// + /// We can only fall back if dictionary encoder is set and we have exceeded dictionary + /// size. + #[inline] + fn should_dict_fallback(&self) -> bool { + match self.encoder.estimated_dict_page_size() { + Some(size) => { + size >= self + .props + .column_dictionary_page_size_limit(self.descr.path()) + } + None => false, + } + } + + /// Exempt a page's mandatory first value from the data page byte limit, + /// when that value alone already exceeds it. + /// + /// Parquet requires every data page to hold at least one value, so such a + /// value cannot be split out no matter how the limit is set. Counting it + /// against the limit makes the limit unsatisfiable, and + /// [`Self::should_add_data_page`] then cuts a page after every single + /// value. + /// + /// For `DELTA_BYTE_ARRAY` that costs more than the extra pages. A value is + /// stored as a suffix of the value before it, and a page boundary resets + /// what "the value before it" refers to, so one value per page means every + /// value is stored in full: a column of large values sharing long prefixes + /// writes exactly the bytes `PLAIN` would + /// ([#10489](https://github.com/apache/arrow-rs/issues/10489)). + /// + /// Only encodings that compress against the preceding value opt in, so + /// `PLAIN` and `DELTA_LENGTH_BYTE_ARRAY` keep their tighter one-value page + /// bound. + /// + /// Known limitation: the caller's trigger keys on a page-opening + /// mini-batch holding exactly one value. Nulls in a chunk make the + /// byte-budget chunker emit multi-level mini-batches, so on nullable + /// columns pages that open with a two-value mini-batch miss the + /// exemption and dedup is only partial; see + /// `test_column_writer_delta_byte_array_nullable_shared_prefix_partial_dedup`. + #[cold] + fn set_page_size_exemption(&mut self) { + if !self.encoder.compresses_against_previous_value() { + return; + } + let size = self.encoder.estimated_data_page_size(); + if size >= self.props.column_data_page_size_limit(self.descr.path()) { + self.page_metrics.page_size_exemption = size; + } + } + + /// Returns true if there is enough data for a data page, false otherwise. + #[inline] + fn should_add_data_page(&self) -> bool { + // This is necessary in the event of a much larger dictionary size than page size + // + // In such a scenario the dictionary decoder may return an estimated encoded + // size in excess of the page size limit, even when there are no buffered values + if self.page_metrics.num_buffered_values == 0 { + return false; + } + + self.page_metrics.num_buffered_rows as usize >= self.props.data_page_row_count_limit() + || self + .encoder + .estimated_data_page_size() + .saturating_sub(self.page_metrics.page_size_exemption) + >= self.props.column_data_page_size_limit(self.descr.path()) + } + + /// Performs dictionary fallback. + /// Prepares and writes dictionary and all data pages into page writer. + fn dict_fallback(&mut self) -> Result<()> { + // At this point we know that we need to fall back. + if self.page_metrics.num_buffered_values > 0 { + self.add_data_page()?; + } + self.write_dictionary_page()?; + self.flush_data_pages()?; + Ok(()) + } + + /// Update the column index and offset index when adding the data page + fn update_column_offset_index( + &mut self, + page_statistics: Option<&ValueStatistics>, + page_variable_length_bytes: Option, + ) { + // update the column index + let null_page = + (self.page_metrics.num_buffered_rows as u64) == self.page_metrics.num_page_nulls; + // a page contains only null values, + // and writers have to set the corresponding entries in min_values and max_values to byte[0] + if null_page && self.column_index_builder.valid() { + self.column_index_builder.append( + null_page, + vec![], + vec![], + self.page_metrics.num_page_nulls as i64, + ); + } else if self.column_index_builder.valid() { + // from page statistics + // If can't get the page statistics, ignore this column/offset index for this column chunk + match &page_statistics { + None => { + self.column_index_builder.to_invalid(); + } + Some(stat) => { + // Check if min/max are still ascending/descending across pages + let new_min = stat.min_opt().unwrap(); + let new_max = stat.max_opt().unwrap(); + if let Some((last_min, last_max)) = &self.last_non_null_data_page_min_max { + if self.data_page_boundary_ascending { + // If last min/max are greater than new min/max then not ascending anymore + let not_ascending = compare_greater(&self.descr, last_min, new_min) + || compare_greater(&self.descr, last_max, new_max); + if not_ascending { + self.data_page_boundary_ascending = false; + } + } + + if self.data_page_boundary_descending { + // If new min/max are greater than last min/max then not descending anymore + let not_descending = compare_greater(&self.descr, new_min, last_min) + || compare_greater(&self.descr, new_max, last_max); + if not_descending { + self.data_page_boundary_descending = false; + } + } + } + self.last_non_null_data_page_min_max = Some((new_min.clone(), new_max.clone())); + + if self.can_truncate_value() { + self.column_index_builder.append( + null_page, + self.truncate_min_value( + self.props.column_index_truncate_length(), + stat.min_bytes_opt().unwrap(), + ) + .0, + self.truncate_max_value( + self.props.column_index_truncate_length(), + stat.max_bytes_opt().unwrap(), + ) + .0, + self.page_metrics.num_page_nulls as i64, + ); + } else { + self.column_index_builder.append( + null_page, + stat.min_bytes_opt().unwrap().to_vec(), + stat.max_bytes_opt().unwrap().to_vec(), + self.page_metrics.num_page_nulls as i64, + ); + } + } + } + } + + // Append page histograms to the `ColumnIndex` histograms + self.column_index_builder.append_histograms( + &self.page_metrics.repetition_level_histogram, + &self.page_metrics.definition_level_histogram, + ); + + // Update the offset index + if let Some(builder) = self.offset_index_builder.as_mut() { + builder.append_row_count(self.page_metrics.num_buffered_rows as i64); + builder.append_unencoded_byte_array_data_bytes(page_variable_length_bytes); + } + } + + /// Determine if we should allow truncating min/max values for this column's statistics + fn can_truncate_value(&self) -> bool { + match self.descr.physical_type() { + // Don't truncate for Float16 and Decimal because their sort order is different + // from that of FIXED_LEN_BYTE_ARRAY sort order. + // So truncation of those types could lead to inaccurate min/max statistics + Type::FIXED_LEN_BYTE_ARRAY + if !matches!( + self.descr.logical_type_ref(), + Some(&LogicalType::Decimal { .. }) | Some(&LogicalType::Float16) + ) => + { + true + } + Type::BYTE_ARRAY => true, + // Truncation only applies for fba/binary physical types + _ => false, + } + } + + /// Returns `true` if this column's logical type is a UTF-8 string. + fn is_utf8(&self) -> bool { + self.get_descriptor().logical_type_ref() == Some(&LogicalType::String) + || self.get_descriptor().converted_type() == ConvertedType::UTF8 + } + + /// Truncates a binary statistic to at most `truncation_length` bytes. + /// + /// If truncation is not possible, returns `data`. + /// + /// The `bool` in the returned tuple indicates whether truncation occurred or not. + /// + /// UTF-8 Note: + /// If the column type indicates UTF-8, and `data` contains valid UTF-8, then the result will + /// also remain valid UTF-8, but may be less tnan `truncation_length` bytes to avoid splitting + /// on non-character boundaries. + fn truncate_min_value(&self, truncation_length: Option, data: &[u8]) -> (Vec, bool) { + truncation_length + .filter(|l| data.len() > *l) + .and_then(|l| + // don't do extra work if this column isn't UTF-8 + if self.is_utf8() { + match str::from_utf8(data) { + Ok(str_data) => truncate_utf8(str_data, l), + Err(_) => Some(data[..l].to_vec()), + } + } else { + Some(data[..l].to_vec()) + } + ) + .map(|truncated| (truncated, true)) + .unwrap_or_else(|| (data.to_vec(), false)) + } + + /// Truncates a binary statistic to at most `truncation_length` bytes, and then increment the + /// final byte(s) to yield a valid upper bound. This may result in a result of less than + /// `truncation_length` bytes if the last byte(s) overflows. + /// + /// If truncation is not possible, returns `data`. + /// + /// The `bool` in the returned tuple indicates whether truncation occurred or not. + /// + /// UTF-8 Note: + /// If the column type indicates UTF-8, and `data` contains valid UTF-8, then the result will + /// also remain valid UTF-8 (but again may be less than `truncation_length` bytes). If `data` + /// does not contain valid UTF-8, then truncation will occur as if the column is non-string + /// binary. + fn truncate_max_value(&self, truncation_length: Option, data: &[u8]) -> (Vec, bool) { + truncation_length + .filter(|l| data.len() > *l) + .and_then(|l| + // don't do extra work if this column isn't UTF-8 + if self.is_utf8() { + match str::from_utf8(data) { + Ok(str_data) => truncate_and_increment_utf8(str_data, l), + Err(_) => increment(data[..l].to_vec()), + } + } else { + increment(data[..l].to_vec()) + } + ) + .map(|truncated| (truncated, true)) + .unwrap_or_else(|| (data.to_vec(), false)) + } + + /// Truncate the min and max values that will be written to a data page + /// header or column chunk Statistics + fn truncate_statistics(&self, statistics: Statistics) -> Statistics { + let backwards_compatible_min_max = self.descr.sort_order().is_signed(); + match statistics { + Statistics::ByteArray(stats) if stats._internal_has_min_max_set() => { + let (min, did_truncate_min) = self.truncate_min_value( + self.props.statistics_truncate_length(), + stats.min_bytes_opt().unwrap(), + ); + let (max, did_truncate_max) = self.truncate_max_value( + self.props.statistics_truncate_length(), + stats.max_bytes_opt().unwrap(), + ); + Statistics::ByteArray( + ValueStatistics::new( + Some(min.into()), + Some(max.into()), + stats.distinct_count(), + stats.null_count_opt(), + backwards_compatible_min_max, + ) + .with_max_is_exact(!did_truncate_max) + .with_min_is_exact(!did_truncate_min), + ) + } + Statistics::FixedLenByteArray(stats) + if (stats._internal_has_min_max_set() && self.can_truncate_value()) => + { + let (min, did_truncate_min) = self.truncate_min_value( + self.props.statistics_truncate_length(), + stats.min_bytes_opt().unwrap(), + ); + let (max, did_truncate_max) = self.truncate_max_value( + self.props.statistics_truncate_length(), + stats.max_bytes_opt().unwrap(), + ); + Statistics::FixedLenByteArray( + ValueStatistics::new( + Some(min.into()), + Some(max.into()), + stats.distinct_count(), + stats.null_count_opt(), + backwards_compatible_min_max, + ) + .with_max_is_exact(!did_truncate_max) + .with_min_is_exact(!did_truncate_min), + ) + } + stats => stats, + } + } + + /// Adds data page. + /// Data page is either buffered in case of dictionary encoding or written directly. + pub(crate) fn add_data_page(&mut self) -> Result<()> { + // Extract encoded values + let values_data = self.encoder.flush_data_page()?; + + let max_def_level = self.descr.max_def_level(); + let max_rep_level = self.descr.max_rep_level(); + + self.column_metrics.num_column_nulls += self.page_metrics.num_page_nulls; + + let page_statistics = match (values_data.min_value, values_data.max_value) { + (Some(min), Some(max)) => { + // Update chunk level statistics + update_min(&self.descr, &min, &mut self.column_metrics.min_column_value); + update_max(&self.descr, &max, &mut self.column_metrics.max_column_value); + + (self.statistics_enabled == EnabledStatistics::Page).then_some( + ValueStatistics::new( + Some(min), + Some(max), + None, + Some(self.page_metrics.num_page_nulls), + false, + ), + ) + } + _ => None, + }; + + // update column and offset index + self.update_column_offset_index( + page_statistics.as_ref(), + values_data.variable_length_bytes, + ); + + // Update histograms and variable_length_bytes in column_metrics + self.column_metrics + .update_from_page_metrics(&self.page_metrics); + self.column_metrics + .update_variable_length_bytes(values_data.variable_length_bytes); + + // From here on, we only need page statistics if they will be written to the page header. + let page_statistics = page_statistics + .filter(|_| self.props.write_page_header_statistics(self.descr.path())) + .map(|stats| self.truncate_statistics(Statistics::from(stats))); + + let compressed_page = match self.props.writer_version() { + WriterVersion::PARQUET_1_0 => { + let mut buffer = vec![]; + + if max_rep_level > 0 { + self.rep_levels_encoder + .flush_to(|data| buffer.extend_from_slice(data)); + } + + if max_def_level > 0 { + self.def_levels_encoder + .flush_to(|data| buffer.extend_from_slice(data)); + } + + buffer.extend_from_slice(&values_data.buf); + let uncompressed_size = buffer.len(); + + if let Some(ref mut cmpr) = self.compressor { + let mut compressed_buf = Vec::with_capacity(uncompressed_size); + cmpr.compress(&buffer[..], &mut compressed_buf)?; + compressed_buf.shrink_to_fit(); + buffer = compressed_buf; + } + + let data_page = Page::DataPage { + buf: buffer.into(), + num_values: self.page_metrics.num_buffered_values, + encoding: values_data.encoding, + def_level_encoding: Encoding::RLE, + rep_level_encoding: Encoding::RLE, + statistics: page_statistics, + }; + + CompressedPage::new(data_page, uncompressed_size) + } + WriterVersion::PARQUET_2_0 => { + let mut rep_levels_byte_len = 0; + let mut def_levels_byte_len = 0; + let mut buffer = vec![]; + + if max_rep_level > 0 { + self.rep_levels_encoder + .flush_to(|data| buffer.extend_from_slice(data)); + rep_levels_byte_len = buffer.len(); + } + + if max_def_level > 0 { + self.def_levels_encoder + .flush_to(|data| buffer.extend_from_slice(data)); + def_levels_byte_len = buffer.len() - rep_levels_byte_len; + } + + let uncompressed_size = + rep_levels_byte_len + def_levels_byte_len + values_data.buf.len(); + + // Data Page v2 compresses values only. + let is_compressed = match self.compressor { + Some(ref mut cmpr) => { + let buffer_len = buffer.len(); + cmpr.compress(&values_data.buf, &mut buffer)?; + let compressed_values_size = buffer.len() - buffer_len; + let threshold = self + .props + .column_data_page_v2_compression_ratio_threshold(self.descr.path()); + if (compressed_values_size as f64) >= (uncompressed_size as f64) * threshold + { + buffer.truncate(buffer_len); + buffer.extend_from_slice(&values_data.buf); + false + } else { + true + } + } + None => { + buffer.extend_from_slice(&values_data.buf); + false + } + }; + + let data_page = Page::DataPageV2 { + buf: buffer.into(), + num_values: self.page_metrics.num_buffered_values, + encoding: values_data.encoding, + num_nulls: self.page_metrics.num_page_nulls as u32, + num_rows: self.page_metrics.num_buffered_rows, + def_levels_byte_len: def_levels_byte_len as u32, + rep_levels_byte_len: rep_levels_byte_len as u32, + is_compressed, + statistics: page_statistics, + }; + + CompressedPage::new(data_page, uncompressed_size) + } + }; + + // Check if we need to buffer data page or flush it to the sink directly. + // + // For dictionary-encoded columns the dictionary page must be written + // first, but it is not final until all values are seen, so completed + // data pages are normally buffered here until `close`. A page writer + // that defers final layout (the Arrow path) instead orders pages itself + // at flush, so we stream the data pages straight through and never let + // them accumulate in memory. + if self.encoder.has_dictionary() && !self.page_writer.defers_dictionary_ordering() { + self.data_pages.push_back(compressed_page); + } else { + self.write_data_page(compressed_page)?; + } + + // Update total number of rows. + self.column_metrics.total_rows_written += self.page_metrics.num_buffered_rows as u64; + self.page_metrics.new_page(); + + Ok(()) + } + + /// Finalises any outstanding data pages and flushes buffered data pages from + /// dictionary encoding into underlying sink. + #[inline] + fn flush_data_pages(&mut self) -> Result<()> { + // Write all outstanding data to a new page. + if self.page_metrics.num_buffered_values > 0 { + self.add_data_page()?; + } + + while let Some(page) = self.data_pages.pop_front() { + self.write_data_page(page)?; + } + + Ok(()) + } + + /// Assembles column chunk metadata. + fn build_column_metadata(&mut self) -> Result { + let total_compressed_size = self.column_metrics.total_compressed_size as i64; + let total_uncompressed_size = self.column_metrics.total_uncompressed_size as i64; + let num_values = self.column_metrics.total_num_values as i64; + let dict_page_offset = self.column_metrics.dictionary_page_offset.map(|v| v as i64); + // If data page offset is not set, then no pages have been written + let data_page_offset = self.column_metrics.data_page_offset.unwrap_or(0) as i64; + + let mut builder = ColumnChunkMetaData::builder(self.descr.clone()) + .set_compression(self.codec) + .set_encodings_mask(EncodingMask::new_from_encodings(self.encodings.iter())) + .set_page_encoding_stats(self.encoding_stats.clone()) + .set_total_compressed_size(total_compressed_size) + .set_total_uncompressed_size(total_uncompressed_size) + .set_num_values(num_values) + .set_data_page_offset(data_page_offset) + .set_dictionary_page_offset(dict_page_offset); + + if self.statistics_enabled != EnabledStatistics::None { + let backwards_compatible_min_max = self.descr.sort_order().is_signed(); + + let statistics = ValueStatistics::::new( + self.column_metrics.min_column_value.clone(), + self.column_metrics.max_column_value.clone(), + self.column_metrics.column_distinct_count, + Some(self.column_metrics.num_column_nulls), + false, + ) + .with_backwards_compatible_min_max(backwards_compatible_min_max) + .into(); + + let statistics = self.truncate_statistics(statistics); + + builder = builder + .set_statistics(statistics) + .set_unencoded_byte_array_data_bytes(self.column_metrics.variable_length_bytes) + .set_repetition_level_histogram( + self.column_metrics.repetition_level_histogram.take(), + ) + .set_definition_level_histogram( + self.column_metrics.definition_level_histogram.take(), + ); + + if let Some(geo_stats) = self.encoder.flush_geospatial_statistics() { + builder = builder.set_geo_statistics(geo_stats); + } + } + + builder = self.set_column_chunk_encryption_properties(builder); + + let metadata = builder.build()?; + Ok(metadata) + } + + /// Writes compressed data page into underlying sink and updates global metrics. + #[inline] + fn write_data_page(&mut self, page: CompressedPage) -> Result<()> { + self.encodings.insert(page.encoding()); + match self.encoding_stats.last_mut() { + Some(encoding_stats) + if encoding_stats.page_type == page.page_type() + && encoding_stats.encoding == page.encoding() => + { + encoding_stats.count += 1; + } + _ => { + // data page type does not change inside a file + // encoding can currently only change from dictionary to non-dictionary once + self.encoding_stats.push(PageEncodingStats { + page_type: page.page_type(), + encoding: page.encoding(), + count: 1, + }); + } + } + let page_spec = self.page_writer.write_page(page)?; + // update offset index + // compressed_size = header_size + compressed_data_size + if let Some(builder) = self.offset_index_builder.as_mut() { + builder + .append_offset_and_size(page_spec.offset as i64, page_spec.compressed_size as i32) + } + self.update_metrics_for_page(page_spec); + Ok(()) + } + + /// Writes dictionary page into underlying sink. + #[inline] + fn write_dictionary_page(&mut self) -> Result<()> { + let compressed_page = { + let mut page = self + .encoder + .flush_dict_page()? + .ok_or_else(|| general_err!("Dictionary encoder is not set"))?; + + let uncompressed_size = page.buf.len(); + + if let Some(ref mut cmpr) = self.compressor { + let mut output_buf = Vec::with_capacity(uncompressed_size); + cmpr.compress(&page.buf, &mut output_buf)?; + page.buf = Bytes::from(output_buf); + } + + let dict_page = Page::DictionaryPage { + buf: page.buf, + num_values: page.num_values as u32, + encoding: self.props.dictionary_page_encoding(), + is_sorted: page.is_sorted, + }; + CompressedPage::new(dict_page, uncompressed_size) + }; + + self.encodings.insert(compressed_page.encoding()); + self.encoding_stats.push(PageEncodingStats { + page_type: PageType::DICTIONARY_PAGE, + encoding: compressed_page.encoding(), + count: 1, + }); + let page_spec = self.page_writer.write_page(compressed_page)?; + self.update_metrics_for_page(page_spec); + // For the directory page, don't need to update column/offset index. + Ok(()) + } + + /// Updates column writer metrics with each page metadata. + #[inline] + fn update_metrics_for_page(&mut self, page_spec: PageWriteSpec) { + self.column_metrics.total_uncompressed_size += page_spec.uncompressed_size as u64; + self.column_metrics.total_compressed_size += page_spec.compressed_size as u64; + self.column_metrics.total_bytes_written += page_spec.bytes_written; + + match page_spec.page_type { + PageType::DATA_PAGE | PageType::DATA_PAGE_V2 => { + self.column_metrics.total_num_values += page_spec.num_values as u64; + if self.column_metrics.data_page_offset.is_none() { + self.column_metrics.data_page_offset = Some(page_spec.offset); + } + } + PageType::DICTIONARY_PAGE => { + assert!( + self.column_metrics.dictionary_page_offset.is_none(), + "Dictionary offset is already set" + ); + self.column_metrics.dictionary_page_offset = Some(page_spec.offset); + } + _ => {} + } + } + + #[inline] + #[cfg(feature = "encryption")] + fn set_column_chunk_encryption_properties( + &self, + builder: ColumnChunkMetaDataBuilder, + ) -> ColumnChunkMetaDataBuilder { + if let Some(encryption_properties) = self.props.file_encryption_properties.as_ref() { + builder.set_column_crypto_metadata(get_column_crypto_metadata( + encryption_properties, + &self.descr, + )) + } else { + builder + } + } + + #[inline] + #[cfg(not(feature = "encryption"))] + fn set_column_chunk_encryption_properties( + &self, + builder: ColumnChunkMetaDataBuilder, + ) -> ColumnChunkMetaDataBuilder { + builder + } +} + +fn update_min(descr: &ColumnDescriptor, val: &T, min: &mut Option) { + update_stat::(descr, val, min, |cur| compare_greater(descr, cur, val)) +} + +fn update_max(descr: &ColumnDescriptor, val: &T, max: &mut Option) { + update_stat::(descr, val, max, |cur| compare_greater(descr, val, cur)) +} + +#[inline] +#[allow(clippy::eq_op)] +fn is_nan(descr: &ColumnDescriptor, val: &T) -> bool { + match T::PHYSICAL_TYPE { + Type::FLOAT | Type::DOUBLE => val != val, + Type::FIXED_LEN_BYTE_ARRAY if descr.logical_type_ref() == Some(&LogicalType::Float16) => { + let val = val.as_bytes(); + let val = f16::from_le_bytes([val[0], val[1]]); + val.is_nan() + } + _ => false, + } +} + +/// Perform a conditional update of `cur`, skipping any NaN values +/// +/// If `cur` is `None`, sets `cur` to `Some(val)`, otherwise calls `should_update` with +/// the value of `cur`, and updates `cur` to `Some(val)` if it returns `true` +fn update_stat( + descr: &ColumnDescriptor, + val: &T, + cur: &mut Option, + should_update: F, +) where + F: Fn(&T) -> bool, +{ + if is_nan(descr, val) { + return; + } + + if cur.as_ref().is_none_or(should_update) { + *cur = Some(val.clone()); + } +} + +/// Evaluate `a > b` according to underlying logical type. +fn compare_greater(descr: &ColumnDescriptor, a: &T, b: &T) -> bool { + match T::PHYSICAL_TYPE { + Type::INT32 | Type::INT64 => { + if let Some(LogicalType::Integer(IntType { + is_signed: false, .. + })) = descr.logical_type_ref() + { + // need to compare unsigned + return compare_greater_unsigned_int(a, b); + } + + match descr.converted_type() { + ConvertedType::UINT_8 + | ConvertedType::UINT_16 + | ConvertedType::UINT_32 + | ConvertedType::UINT_64 => { + return compare_greater_unsigned_int(a, b); + } + _ => {} + }; + } + Type::FIXED_LEN_BYTE_ARRAY | Type::BYTE_ARRAY => { + if let Some(LogicalType::Decimal(_)) = descr.logical_type_ref() { + return compare_greater_byte_array_decimals(a.as_bytes(), b.as_bytes()); + } + if let ConvertedType::DECIMAL = descr.converted_type() { + return compare_greater_byte_array_decimals(a.as_bytes(), b.as_bytes()); + } + if let Some(LogicalType::Float16) = descr.logical_type_ref() { + return compare_greater_f16(a.as_bytes(), b.as_bytes()); + } + } + + _ => {} + } + + // compare independent of logical / converted type + a > b +} + +// ---------------------------------------------------------------------- +// Encoding support for column writer. +// This mirrors parquet-mr default encodings for writes. See: +// https://github.com/apache/parquet-mr/blob/master/parquet-column/src/main/java/org/apache/parquet/column/values/factory/DefaultV1ValuesWriterFactory.java +// https://github.com/apache/parquet-mr/blob/master/parquet-column/src/main/java/org/apache/parquet/column/values/factory/DefaultV2ValuesWriterFactory.java + +/// Returns encoding for a column when no other encoding is provided in writer properties. +fn fallback_encoding(kind: Type, props: &WriterProperties) -> Encoding { + match (kind, props.writer_version()) { + (Type::BOOLEAN, WriterVersion::PARQUET_2_0) => Encoding::RLE, + (Type::INT32, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BINARY_PACKED, + (Type::INT64, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BINARY_PACKED, + (Type::BYTE_ARRAY, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BYTE_ARRAY, + (Type::FIXED_LEN_BYTE_ARRAY, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BYTE_ARRAY, + _ => Encoding::PLAIN, + } +} + +/// Returns true if dictionary is supported for column writer, false otherwise. +fn has_dictionary_support(kind: Type, props: &WriterProperties) -> bool { + match (kind, props.writer_version()) { + // Booleans do not support dict encoding and should use a fallback encoding. + (Type::BOOLEAN, _) => false, + // Dictionary encoding was not enabled in PARQUET 1.0 + (Type::FIXED_LEN_BYTE_ARRAY, WriterVersion::PARQUET_1_0) => false, + (Type::FIXED_LEN_BYTE_ARRAY, WriterVersion::PARQUET_2_0) => true, + _ => true, + } +} + +#[inline] +fn compare_greater_unsigned_int(a: &T, b: &T) -> bool { + a.as_u64().unwrap() > b.as_u64().unwrap() +} + +#[inline] +fn compare_greater_f16(a: &[u8], b: &[u8]) -> bool { + let a = f16::from_le_bytes(a.try_into().unwrap()); + let b = f16::from_le_bytes(b.try_into().unwrap()); + a > b +} + +/// Signed comparison of bytes arrays +fn compare_greater_byte_array_decimals(a: &[u8], b: &[u8]) -> bool { + let a_length = a.len(); + let b_length = b.len(); + + if a_length == 0 || b_length == 0 { + return a_length > 0; + } + + let first_a: u8 = a[0]; + let first_b: u8 = b[0]; + + // We can short circuit for different signed numbers or + // for equal length bytes arrays that have different first bytes. + // The equality requirement is necessary for sign extension cases. + // 0xFF10 should be equal to 0x10 (due to big endian sign extension). + if (0x80 & first_a) != (0x80 & first_b) || (a_length == b_length && first_a != first_b) { + return (first_a as i8) > (first_b as i8); + } + + // When the lengths are unequal and the numbers are of the same + // sign we need to do comparison by sign extending the shorter + // value first, and once we get to equal sized arrays, lexicographical + // unsigned comparison of everything but the first byte is sufficient. + + let extension: u8 = if (first_a as i8) < 0 { 0xFF } else { 0 }; + + if a_length != b_length { + let not_equal = if a_length > b_length { + let lead_length = a_length - b_length; + a[0..lead_length].iter().any(|&x| x != extension) + } else { + let lead_length = b_length - a_length; + b[0..lead_length].iter().any(|&x| x != extension) + }; + + if not_equal { + let negative_values: bool = (first_a as i8) < 0; + let a_longer: bool = a_length > b_length; + return if negative_values { !a_longer } else { a_longer }; + } + } + + (a[1..]) > (b[1..]) +} + +/// Truncate a UTF-8 slice to the longest prefix that is still a valid UTF-8 string, +/// while being less than `length` bytes and non-empty. Returns `None` if truncation +/// is not possible within those constraints. +/// +/// The caller guarantees that data.len() > length. +fn truncate_utf8(data: &str, length: usize) -> Option> { + let split = (1..=length).rfind(|x| data.is_char_boundary(*x))?; + Some(data.as_bytes()[..split].to_vec()) +} + +/// Truncate a UTF-8 slice and increment it's final character. The returned value is the +/// longest such slice that is still a valid UTF-8 string while being less than `length` +/// bytes and non-empty. Returns `None` if no such transformation is possible. +/// +/// The caller guarantees that data.len() > length. +fn truncate_and_increment_utf8(data: &str, length: usize) -> Option> { + // UTF-8 is max 4 bytes, so start search 3 back from desired length + let lower_bound = length.saturating_sub(3); + let split = (lower_bound..=length).rfind(|x| data.is_char_boundary(*x))?; + increment_utf8(data.get(..split)?) +} + +/// Increment the final character in a UTF-8 string in such a way that the returned result +/// is still a valid UTF-8 string. The returned string may be shorter than the input if the +/// last character(s) cannot be incremented (due to overflow or producing invalid code points). +/// Returns `None` if the string cannot be incremented. +/// +/// Note that this implementation will not promote an N-byte code point to (N+1) bytes. +fn increment_utf8(data: &str) -> Option> { + for (idx, original_char) in data.char_indices().rev() { + let original_len = original_char.len_utf8(); + if let Some(next_char) = char::from_u32(original_char as u32 + 1) { + // do not allow increasing byte width of incremented char + if next_char.len_utf8() == original_len { + let mut result = data.as_bytes()[..idx + original_len].to_vec(); + next_char.encode_utf8(&mut result[idx..]); + return Some(result); + } + } + } + + None +} + +/// Try and increment the bytes from right to left. +/// +/// Returns `None` if all bytes are set to `u8::MAX`. +fn increment(mut data: Vec) -> Option> { + for byte in data.iter_mut().rev() { + let (incremented, overflow) = byte.overflowing_add(1); + *byte = incremented; + + if !overflow { + return Some(data); + } + } + + None +} + +#[cfg(test)] +mod tests { + use crate::{ + file::{properties::DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH, writer::SerializedFileWriter}, + schema::parser::parse_message_type, + }; + use core::str; + use rand::distr::uniform::SampleUniform; + use std::{fs::File, sync::Arc}; + + use crate::column::{ + page::PageReader, + reader::{ColumnReaderImpl, get_column_reader, get_typed_column_reader}, + }; + use crate::file::writer::TrackedWrite; + use crate::file::{ + properties::ReaderProperties, reader::SerializedPageReader, writer::SerializedPageWriter, + }; + use crate::schema::types::{ColumnPath, Type as SchemaType}; + use crate::util::test_common::rand_gen::random_numbers_range; + + use super::*; + + #[test] + fn test_column_writer_inconsistent_def_rep_length() { + let page_writer = get_test_page_writer(); + let props = Default::default(); + let mut writer = get_test_column_writer::(page_writer, 1, 1, props); + let res = writer.write_batch(&[1, 2, 3, 4], Some(&[1, 1, 1]), Some(&[0, 0])); + assert!(res.is_err()); + if let Err(err) = res { + assert_eq!( + format!("{err}"), + "Parquet error: Inconsistent length of definition and repetition levels: 3 != 2" + ); + } + } + + #[test] + fn test_column_writer_invalid_def_levels() { + let page_writer = get_test_page_writer(); + let props = Default::default(); + let mut writer = get_test_column_writer::(page_writer, 1, 0, props); + let res = writer.write_batch(&[1, 2, 3, 4], None, None); + assert!(res.is_err()); + if let Err(err) = res { + assert_eq!( + format!("{err}"), + "Parquet error: Definition levels are required, because max definition level = 1" + ); + } + } + + #[test] + fn test_column_writer_invalid_rep_levels() { + let page_writer = get_test_page_writer(); + let props = Default::default(); + let mut writer = get_test_column_writer::(page_writer, 0, 1, props); + let res = writer.write_batch(&[1, 2, 3, 4], None, None); + assert!(res.is_err()); + if let Err(err) = res { + assert_eq!( + format!("{err}"), + "Parquet error: Repetition levels are required, because max repetition level = 1" + ); + } + } + + #[test] + fn test_column_writer_not_enough_values_to_write() { + let page_writer = get_test_page_writer(); + let props = Default::default(); + let mut writer = get_test_column_writer::(page_writer, 1, 0, props); + let res = writer.write_batch(&[1, 2], Some(&[1, 1, 1, 1]), None); + assert!(res.is_err()); + if let Err(err) = res { + assert_eq!( + format!("{err}"), + "Parquet error: Expected to write 4 values, but have only 2" + ); + } + } + + #[test] + fn test_column_writer_write_only_one_dictionary_page() { + let page_writer = get_test_page_writer(); + let props = Default::default(); + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + writer.write_batch(&[1, 2, 3, 4], None, None).unwrap(); + // First page should be correctly written. + writer.add_data_page().unwrap(); + writer.write_dictionary_page().unwrap(); + let err = writer.write_dictionary_page().unwrap_err().to_string(); + assert_eq!(err, "Parquet error: Dictionary encoder is not set"); + } + + #[test] + fn test_column_writer_error_when_writing_disabled_dictionary() { + let page_writer = get_test_page_writer(); + let props = Arc::new( + WriterProperties::builder() + .set_dictionary_enabled(false) + .build(), + ); + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + writer.write_batch(&[1, 2, 3, 4], None, None).unwrap(); + let err = writer.write_dictionary_page().unwrap_err().to_string(); + assert_eq!(err, "Parquet error: Dictionary encoder is not set"); + } + + #[test] + fn test_column_writer_boolean_type_does_not_support_dictionary() { + let page_writer = get_test_page_writer(); + let props = Arc::new( + WriterProperties::builder() + .set_dictionary_enabled(true) + .build(), + ); + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + writer + .write_batch(&[true, false, true, false], None, None) + .unwrap(); + + let r = writer.close().unwrap(); + // PlainEncoder uses bit writer to write boolean values, which all fit into 1 + // byte. + assert_eq!(r.bytes_written, 1); + assert_eq!(r.rows_written, 4); + + let metadata = r.metadata; + assert_eq!( + metadata.encodings().collect::>(), + vec![Encoding::PLAIN, Encoding::RLE] + ); + assert_eq!(metadata.num_values(), 4); // just values + assert_eq!(metadata.dictionary_page_offset(), None); + } + + #[test] + fn test_column_writer_default_encoding_support_bool() { + check_encoding_write_support::( + WriterVersion::PARQUET_1_0, + true, + &[true, false], + None, + &[Encoding::PLAIN, Encoding::RLE], + &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_1_0, + false, + &[true, false], + None, + &[Encoding::PLAIN, Encoding::RLE], + &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_2_0, + true, + &[true, false], + None, + &[Encoding::RLE], + &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE, 1)], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_2_0, + false, + &[true, false], + None, + &[Encoding::RLE], + &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE, 1)], + ); + } + + #[test] + fn test_column_writer_default_encoding_support_int32() { + check_encoding_write_support::( + WriterVersion::PARQUET_1_0, + true, + &[1, 2], + 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), + ], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_1_0, + false, + &[1, 2], + None, + &[Encoding::PLAIN, Encoding::RLE], + &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_2_0, + true, + &[1, 2], + Some(0), + &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY], + &[ + encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1), + encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1), + ], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_2_0, + false, + &[1, 2], + None, + &[Encoding::RLE, Encoding::DELTA_BINARY_PACKED], + &[encoding_stats( + PageType::DATA_PAGE_V2, + Encoding::DELTA_BINARY_PACKED, + 1, + )], + ); + } + + #[test] + fn test_column_writer_default_encoding_support_int64() { + check_encoding_write_support::( + WriterVersion::PARQUET_1_0, + true, + &[1, 2], + 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), + ], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_1_0, + false, + &[1, 2], + None, + &[Encoding::PLAIN, Encoding::RLE], + &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_2_0, + true, + &[1, 2], + Some(0), + &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY], + &[ + encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1), + encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1), + ], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_2_0, + false, + &[1, 2], + None, + &[Encoding::RLE, Encoding::DELTA_BINARY_PACKED], + &[encoding_stats( + PageType::DATA_PAGE_V2, + Encoding::DELTA_BINARY_PACKED, + 1, + )], + ); + } + + #[test] + fn test_column_writer_default_encoding_support_int96() { + check_encoding_write_support::( + WriterVersion::PARQUET_1_0, + true, + &[Int96::from(vec![1, 2, 3])], + 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), + ], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_1_0, + false, + &[Int96::from(vec![1, 2, 3])], + None, + &[Encoding::PLAIN, Encoding::RLE], + &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_2_0, + true, + &[Int96::from(vec![1, 2, 3])], + Some(0), + &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY], + &[ + encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1), + encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1), + ], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_2_0, + false, + &[Int96::from(vec![1, 2, 3])], + None, + &[Encoding::PLAIN, Encoding::RLE], + &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::PLAIN, 1)], + ); + } + + #[test] + fn test_column_writer_default_encoding_support_float() { + check_encoding_write_support::( + WriterVersion::PARQUET_1_0, + true, + &[1.0, 2.0], + 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), + ], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_1_0, + false, + &[1.0, 2.0], + None, + &[Encoding::PLAIN, Encoding::RLE], + &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_2_0, + true, + &[1.0, 2.0], + Some(0), + &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY], + &[ + encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1), + encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1), + ], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_2_0, + false, + &[1.0, 2.0], + None, + &[Encoding::PLAIN, Encoding::RLE], + &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::PLAIN, 1)], + ); + } + + #[test] + fn test_column_writer_default_encoding_support_double() { + check_encoding_write_support::( + WriterVersion::PARQUET_1_0, + true, + &[1.0, 2.0], + 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), + ], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_1_0, + false, + &[1.0, 2.0], + None, + &[Encoding::PLAIN, Encoding::RLE], + &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_2_0, + true, + &[1.0, 2.0], + Some(0), + &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY], + &[ + encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1), + encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1), + ], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_2_0, + false, + &[1.0, 2.0], + None, + &[Encoding::PLAIN, Encoding::RLE], + &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::PLAIN, 1)], + ); + } + + #[test] + fn test_column_writer_default_encoding_support_byte_array() { + check_encoding_write_support::( + WriterVersion::PARQUET_1_0, + true, + &[ByteArray::from(vec![1u8])], + 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), + ], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_1_0, + false, + &[ByteArray::from(vec![1u8])], + None, + &[Encoding::PLAIN, Encoding::RLE], + &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_2_0, + true, + &[ByteArray::from(vec![1u8])], + Some(0), + &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY], + &[ + encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1), + encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1), + ], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_2_0, + false, + &[ByteArray::from(vec![1u8])], + None, + &[Encoding::RLE, Encoding::DELTA_BYTE_ARRAY], + &[encoding_stats( + PageType::DATA_PAGE_V2, + Encoding::DELTA_BYTE_ARRAY, + 1, + )], + ); + } + + #[test] + fn test_column_writer_default_encoding_support_fixed_len_byte_array() { + 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)], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_1_0, + false, + &[ByteArray::from(vec![1u8]).into()], + None, + &[Encoding::PLAIN, Encoding::RLE], + &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_2_0, + true, + &[ByteArray::from(vec![1u8]).into()], + Some(0), + &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY], + &[ + encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1), + encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1), + ], + ); + check_encoding_write_support::( + WriterVersion::PARQUET_2_0, + false, + &[ByteArray::from(vec![1u8]).into()], + None, + &[Encoding::RLE, Encoding::DELTA_BYTE_ARRAY], + &[encoding_stats( + PageType::DATA_PAGE_V2, + Encoding::DELTA_BYTE_ARRAY, + 1, + )], + ); + } + + #[test] + fn test_column_writer_check_metadata() { + let page_writer = get_test_page_writer(); + let props = Default::default(); + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + writer.write_batch(&[1, 2, 3, 4], None, None).unwrap(); + + let r = writer.close().unwrap(); + assert_eq!(r.bytes_written, 20); + assert_eq!(r.rows_written, 4); + + let metadata = r.metadata; + assert_eq!( + metadata.encodings().collect::>(), + vec![Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY] + ); + assert_eq!(metadata.num_values(), 4); + assert_eq!(metadata.compressed_size(), 20); + assert_eq!(metadata.uncompressed_size(), 20); + assert_eq!(metadata.data_page_offset(), 0); + assert_eq!(metadata.dictionary_page_offset(), Some(0)); + if let Some(stats) = metadata.statistics() { + assert_eq!(stats.null_count_opt(), Some(0)); + assert_eq!(stats.distinct_count_opt(), None); + if let Statistics::Int32(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &1); + assert_eq!(stats.max_opt().unwrap(), &4); + } else { + panic!("expecting Statistics::Int32"); + } + } else { + panic!("metadata missing statistics"); + } + } + + #[test] + fn test_column_writer_check_byte_array_min_max() { + let page_writer = get_test_page_writer(); + let props = Default::default(); + let mut writer = get_test_decimals_column_writer::(page_writer, 0, 0, props); + writer + .write_batch( + &[ + ByteArray::from(vec![ + 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 179u8, 172u8, 19u8, + 35u8, 231u8, 90u8, 0u8, 0u8, + ]), + ByteArray::from(vec![ + 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 228u8, 62u8, 146u8, + 152u8, 177u8, 56u8, 0u8, 0u8, + ]), + ByteArray::from(vec![ + 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, + 0u8, + ]), + ByteArray::from(vec![ + 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 41u8, 162u8, 36u8, 26u8, 246u8, + 44u8, 0u8, 0u8, + ]), + ], + None, + None, + ) + .unwrap(); + let metadata = writer.close().unwrap().metadata; + if let Some(stats) = metadata.statistics() { + if let Statistics::ByteArray(stats) = stats { + assert_eq!( + stats.min_opt().unwrap(), + &ByteArray::from(vec![ + 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 179u8, 172u8, 19u8, + 35u8, 231u8, 90u8, 0u8, 0u8, + ]) + ); + assert_eq!( + stats.max_opt().unwrap(), + &ByteArray::from(vec![ + 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 41u8, 162u8, 36u8, 26u8, 246u8, + 44u8, 0u8, 0u8, + ]) + ); + } else { + panic!("expecting Statistics::ByteArray"); + } + } else { + panic!("metadata missing statistics"); + } + } + + #[test] + fn test_column_writer_uint32_converted_type_min_max() { + let page_writer = get_test_page_writer(); + let props = Default::default(); + let mut writer = get_test_unsigned_int_given_as_converted_column_writer::( + page_writer, + 0, + 0, + props, + ); + writer.write_batch(&[0, 1, 2, 3, 4, 5], None, None).unwrap(); + let metadata = writer.close().unwrap().metadata; + if let Some(stats) = metadata.statistics() { + if let Statistics::Int32(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &0,); + assert_eq!(stats.max_opt().unwrap(), &5,); + } else { + panic!("expecting Statistics::Int32"); + } + } else { + panic!("metadata missing statistics"); + } + } + + #[test] + fn test_column_writer_precalculated_statistics() { + let page_writer = get_test_page_writer(); + let props = Arc::new( + WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Chunk) + .build(), + ); + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + writer + .write_batch_with_statistics( + &[1, 2, 3, 4], + None, + None, + Some(&-17), + Some(&9000), + Some(55), + ) + .unwrap(); + + let r = writer.close().unwrap(); + assert_eq!(r.bytes_written, 20); + assert_eq!(r.rows_written, 4); + + let metadata = r.metadata; + assert_eq!( + metadata.encodings().collect::>(), + vec![Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY] + ); + assert_eq!(metadata.num_values(), 4); + assert_eq!(metadata.compressed_size(), 20); + assert_eq!(metadata.uncompressed_size(), 20); + assert_eq!(metadata.data_page_offset(), 0); + assert_eq!(metadata.dictionary_page_offset(), Some(0)); + if let Some(stats) = metadata.statistics() { + assert_eq!(stats.null_count_opt(), Some(0)); + assert_eq!(stats.distinct_count_opt().unwrap_or(0), 55); + if let Statistics::Int32(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &-17); + assert_eq!(stats.max_opt().unwrap(), &9000); + } else { + panic!("expecting Statistics::Int32"); + } + } else { + panic!("metadata missing statistics"); + } + } + + #[test] + fn test_mixed_precomputed_statistics() { + let mut buf = Vec::with_capacity(100); + let mut write = TrackedWrite::new(&mut buf); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let props = Arc::new( + WriterProperties::builder() + .set_write_page_header_statistics(true) + .build(), + ); + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + + writer.write_batch(&[1, 2, 3, 4], None, None).unwrap(); + writer + .write_batch_with_statistics(&[5, 6, 7], None, None, Some(&5), Some(&7), Some(3)) + .unwrap(); + + let r = writer.close().unwrap(); + + let stats = r.metadata.statistics().unwrap(); + assert_eq!(stats.min_bytes_opt().unwrap(), 1_i32.to_le_bytes()); + assert_eq!(stats.max_bytes_opt().unwrap(), 7_i32.to_le_bytes()); + assert_eq!(stats.null_count_opt(), Some(0)); + assert!(stats.distinct_count_opt().is_none()); + + drop(write); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .set_read_page_statistics(true) + .build(); + let reader = SerializedPageReader::new_with_properties( + Arc::new(Bytes::from(buf)), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(); + + let pages = reader.collect::>>().unwrap(); + assert_eq!(pages.len(), 2); + + 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()); + } + + #[test] + fn test_disabled_statistics() { + let mut buf = Vec::with_capacity(100); + let mut write = TrackedWrite::new(&mut buf); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::None) + .set_writer_version(WriterVersion::PARQUET_2_0) + .build(); + let props = Arc::new(props); + + let mut writer = get_test_column_writer::(page_writer, 1, 0, props); + writer + .write_batch(&[1, 2, 3, 4], Some(&[1, 0, 0, 1, 1, 1]), None) + .unwrap(); + + let r = writer.close().unwrap(); + assert!(r.metadata.statistics().is_none()); + + drop(write); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let reader = SerializedPageReader::new_with_properties( + Arc::new(Bytes::from(buf)), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(); + + let pages = reader.collect::>>().unwrap(); + assert_eq!(pages.len(), 2); + + assert_eq!(pages[0].page_type(), PageType::DICTIONARY_PAGE); + assert_eq!(pages[1].page_type(), PageType::DATA_PAGE_V2); + + match &pages[1] { + Page::DataPageV2 { + num_values, + num_nulls, + num_rows, + statistics, + .. + } => { + assert_eq!(*num_values, 6); + assert_eq!(*num_nulls, 2); + assert_eq!(*num_rows, 6); + assert!(statistics.is_none()); + } + _ => unreachable!(), + } + } + + #[test] + fn test_column_writer_empty_column_roundtrip() { + let props = Default::default(); + column_roundtrip::(props, &[], None, None); + } + + #[test] + fn test_column_writer_non_nullable_values_roundtrip() { + let props = Default::default(); + column_roundtrip_random::(props, 1024, i32::MIN, i32::MAX, 0, 0); + } + + #[test] + fn test_column_writer_nullable_non_repeated_values_roundtrip() { + let props = Default::default(); + column_roundtrip_random::(props, 1024, i32::MIN, i32::MAX, 10, 0); + } + + #[test] + fn test_column_writer_nullable_repeated_values_roundtrip() { + let props = Default::default(); + column_roundtrip_random::(props, 1024, i32::MIN, i32::MAX, 10, 10); + } + + #[test] + fn test_column_writer_dictionary_fallback_small_data_page() { + let props = WriterProperties::builder() + .set_dictionary_page_size_limit(32) + .set_data_page_size_limit(32) + .build(); + column_roundtrip_random::(props, 1024, i32::MIN, i32::MAX, 10, 10); + } + + #[test] + fn test_column_writer_small_write_batch_size() { + for i in &[1usize, 2, 5, 10, 11, 1023] { + let props = WriterProperties::builder().set_write_batch_size(*i).build(); + + column_roundtrip_random::(props, 1024, i32::MIN, i32::MAX, 10, 10); + } + } + + #[test] + fn test_column_writer_dictionary_disabled_v1() { + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .build(); + column_roundtrip_random::(props, 1024, i32::MIN, i32::MAX, 10, 10); + } + + #[test] + fn test_column_writer_dictionary_disabled_v2() { + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_2_0) + .set_dictionary_enabled(false) + .build(); + column_roundtrip_random::(props, 1024, i32::MIN, i32::MAX, 10, 10); + } + + #[test] + fn test_column_writer_compression_v1() { + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_compression(Compression::SNAPPY) + .build(); + column_roundtrip_random::(props, 2048, i32::MIN, i32::MAX, 10, 10); + } + + #[test] + fn test_column_writer_compression_v2() { + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_2_0) + .set_compression(Compression::SNAPPY) + .build(); + column_roundtrip_random::(props, 2048, i32::MIN, i32::MAX, 10, 10); + } + + #[test] + fn test_column_writer_v2_compression_ratio_threshold() { + fn write_v2_page(threshold: f64) -> bool { + let mut buf = Vec::with_capacity(4096); + let mut write = TrackedWrite::new(&mut buf); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let props = Arc::new( + WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_2_0) + .set_compression(Compression::SNAPPY) + .set_dictionary_enabled(false) + .set_data_page_v2_compression_ratio_threshold(threshold) + .build(), + ); + + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + let values: Vec = vec![42; 4096]; + writer.write_batch(&values, None, None).unwrap(); + let r = writer.close().unwrap(); + drop(write); + + let reader_props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let reader = SerializedPageReader::new_with_properties( + Arc::new(Bytes::from(buf)), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(reader_props), + ) + .unwrap(); + let pages = reader.collect::>>().unwrap(); + let data_page = pages + .iter() + .find(|p| p.page_type() == PageType::DATA_PAGE_V2) + .expect("expected a v2 data page"); + match data_page { + Page::DataPageV2 { is_compressed, .. } => *is_compressed, + _ => unreachable!(), + } + } + + // Default threshold keeps the compressed buffer for constant data. + assert!(write_v2_page(1.0)); + // A strict threshold (require >1000x reduction) discards it. + assert!(!write_v2_page(0.001)); + } + + #[test] + fn test_column_writer_add_data_pages_with_dict() { + // ARROW-5129: Test verifies that we add data page in case of dictionary encoding + // and no fallback occurred so far. + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let props = Arc::new( + WriterProperties::builder() + .set_data_page_size_limit(10) + .set_write_batch_size(3) // write 3 values at a time + .build(), + ); + let data = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + writer.write_batch(data, None, None).unwrap(); + let r = writer.close().unwrap(); + + drop(write); + + // Read pages and check the sequence + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let mut page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(), + ); + let mut res = Vec::new(); + while let Some(page) = page_reader.get_next_page().unwrap() { + res.push((page.page_type(), page.num_values(), page.buffer().len())); + } + assert_eq!( + res, + vec![ + (PageType::DICTIONARY_PAGE, 10, 40), + (PageType::DATA_PAGE, 9, 10), + (PageType::DATA_PAGE, 1, 3), + ] + ); + assert_eq!( + r.metadata.page_encoding_stats(), + Some(&vec![ + PageEncodingStats { + page_type: PageType::DICTIONARY_PAGE, + encoding: Encoding::PLAIN, + count: 1 + }, + PageEncodingStats { + page_type: PageType::DATA_PAGE, + encoding: Encoding::RLE_DICTIONARY, + count: 2, + } + ]) + ); + } + + #[test] + fn test_column_writer_column_data_page_size_limit() { + let props = Arc::new( + WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_data_page_size_limit(1000) + .set_column_data_page_size_limit(ColumnPath::from("col"), 10) + .set_write_batch_size(3) + .build(), + ); + let data = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; + + let col_values = + write_and_collect_page_values(ColumnPath::from("col"), Arc::clone(&props), data); + let other_values = write_and_collect_page_values(ColumnPath::from("other"), props, data); + + assert_eq!(col_values, vec![3, 3, 3, 1]); + assert_eq!(other_values, vec![10]); + } + + #[test] + fn test_column_writer_caps_page_size_for_large_byte_array_values() { + // Regression: the post-write data page byte limit check only fires + // at mini-batch boundaries, so a 1024-row mini-batch of multi-MiB + // BYTE_ARRAY values used to buffer multiple GiB into a single page + // before the limit was even consulted. With the threshold-based + // granular mode this batch should split into ~one page per value. + let value_size = 64 * 1024; // 64 KiB per value + let page_byte_limit = 16 * 1024; // 16 KiB page limit + let num_rows = 64; + + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_data_page_size_limit(page_byte_limit) + // Default write_batch_size (1024) — without the fix this + // buffers the entire input into a single ~4 MiB page. + .build(); + + let data: Vec<_> = (0..num_rows) + .map(|i| ByteArray::from(vec![i as u8; value_size])) + .collect(); + let pages = write_and_collect_pages::(props, 0, 0, &data, None, None); + + // Every value must end up somewhere. + let total_values: u32 = pages.data_pages.iter().map(|(_, n)| n).sum(); + assert_eq!(total_values as usize, num_rows); + // Without the fix this assertion fired with one ~4 MiB page; the + // threshold splits the input so that no page holds more than a + // single oversized value's worth of bytes. + assert!( + pages.data_pages.len() >= num_rows / 2, + "expected pages to be cut close to one per value, got {:?}", + pages.data_pages, + ); + // Each page must be bounded by roughly one value's worth of bytes; + // parquet allows a single oversized value to occupy a page by + // itself but never lets us pile many of them together. + for (size, _) in &pages.data_pages { + assert!( + *size <= value_size + 64, + "page size {size} exceeds one-value bound ({}B) — pages {:?}", + value_size + 64, + pages.data_pages, + ); + } + } + + #[test] + fn test_column_writer_delta_byte_array_dedups_large_shared_prefix_values() { + // Regression for https://github.com/apache/arrow-rs/issues/10489. + // 16 identical 64 KiB values against a 16 KiB page limit: every value + // is over the limit on its own, and `DELTA_BYTE_ARRAY` should still + // dedup them down to about one value's worth of bytes in total. + let value_size = 64 * 1024; // 64 KiB per value, > the page limit + let page_byte_limit = 16 * 1024; + let num_rows = 16; + + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::DELTA_BYTE_ARRAY) + .set_data_page_size_limit(page_byte_limit) + .set_statistics_enabled(EnabledStatistics::None) + .build(); + + // Identical values: one full value plus `num_rows - 1` zero-length + // suffixes is all this column should cost. + let data: Vec<_> = (0..num_rows) + .map(|_| ByteArray::from(vec![b'a'; value_size])) + .collect(); + let pages = write_and_collect_pages::(props, 0, 0, &data, None, None); + + // Every value must still end up somewhere. + let total_values: u32 = pages.data_pages.iter().map(|(_, n)| n).sum(); + assert_eq!(total_values as usize, num_rows); + + // Before the fix this was `num_rows * value_size` — byte for byte + // what PLAIN produces, i.e. the encoding doing no work at all. + let total_bytes: usize = pages.data_pages.iter().map(|(size, _)| size).sum(); + assert!( + total_bytes < 2 * value_size, + "expected under 2x a single value ({}B) for {num_rows} identical \ + values, got {total_bytes}B across pages {:?}", + 2 * value_size, + pages.data_pages, + ); + } + + #[test] + fn test_column_writer_delta_byte_array_bounds_pages_without_shared_prefix() { + // Companion to the test above: same shape, but the values share no + // prefix, so there is nothing to dedup and pages must stay bounded + // by the value size. This is why the exemption covers one value + // rather than dropping the byte budget altogether. + let value_size = 64 * 1024; + let page_byte_limit = 16 * 1024; + let num_rows = 16; + + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::DELTA_BYTE_ARRAY) + .set_data_page_size_limit(page_byte_limit) + .set_statistics_enabled(EnabledStatistics::None) + .build(); + + // No two values share a prefix: they differ at the first byte. + let data: Vec<_> = (0..num_rows) + .map(|i| ByteArray::from(vec![i as u8; value_size])) + .collect(); + let pages = write_and_collect_pages::(props, 0, 0, &data, None, None); + + let total_values: u32 = pages.data_pages.iter().map(|(_, n)| n).sum(); + assert_eq!(total_values as usize, num_rows); + + // Expect at most two values per page: the exempted first value plus + // one more that trips the budget. + let upper_bound = 2 * value_size + 64; + for (size, n_values) in &pages.data_pages { + assert!( + *size <= upper_bound, + "page size {size} exceeds two-value bound ({upper_bound}B); pages {:?}", + pages.data_pages, + ); + assert!( + *n_values <= 2, + "page holds {n_values} values, expected at most 2; pages {:?}", + pages.data_pages, + ); + } + } + + #[test] + fn test_column_writer_delta_byte_array_nullable_shared_prefix_partial_dedup() { + // Documents the *current* behavior of the first-value exemption on a + // nullable column; this pins a known limitation, not an ideal. + // + // The exemption fires when a page's first mini-batch contains exactly + // one value. For a non-nullable column the byte-budget chunker gives + // an over-limit value a one-level mini-batch, so that always holds. + // One null in the chunk changes the level:value ratio to 17:16, the + // chunker rounds up to two-level mini-batches, and a page whose first + // mini-batch carries two values misses the exemption: it is cut after + // those two values, and its first value is stored in full. + // + // The one mini-batch that pairs the null with a value has a single + // value, so the page it opens does get the exemption and accumulates + // every remaining suffix. The result for 16 identical values with a + // null at index 8 is four two-value pages (each storing one value in + // full), then one exempt page holding the rest: + // + // values per page: [2, 2, 2, 2, 9] (counts include the null level) + // total bytes: ~5 full values, vs ~1 ideally and 16 for PLAIN + // + // If the exemption trigger is ever keyed on values written to the + // page (0 -> 1) instead of mini-batch shape, this test should fail + // with fewer, larger pages — update it to pin the improved layout. + let value_size = 64 * 1024; + let page_byte_limit = 16 * 1024; + let num_values = 16; + + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::DELTA_BYTE_ARRAY) + .set_data_page_size_limit(page_byte_limit) + .set_statistics_enabled(EnabledStatistics::None) + .build(); + + let data: Vec<_> = (0..num_values) + .map(|_| ByteArray::from(vec![b'a'; value_size])) + .collect(); + // 17 levels: a null at index 8, values everywhere else. + let def_levels: Vec = (0..num_values as i16 + 1) + .map(|i| i16::from(i != 8)) + .collect(); + let pages = + write_and_collect_pages::(props, 1, 0, &data, Some(&def_levels), None); + + let per_page_values: Vec = pages.data_pages.iter().map(|(_, n)| *n).collect(); + assert_eq!(per_page_values, vec![2, 2, 2, 2, 9]); + + let total_bytes: usize = pages.data_pages.iter().map(|(size, _)| size).sum(); + assert!( + total_bytes > 4 * value_size && total_bytes < 6 * value_size, + "expected ~5 full values' worth of bytes (partial dedup), \ + got {total_bytes}B across pages {:?}", + pages.data_pages, + ); + } + + #[test] + fn test_column_writer_caps_page_size_for_large_values_in_list() { + // Coverage for the Materialized-rep branch of + // `write_granular_chunk`. The flat-column regression test + // exercises the per-level step; this exercises the + // record-by-record step used when rep levels are present. + // + // Column is `list` (max_def = 1, max_rep = 1) + // with 3 records of 3 large blobs each. The page byte limit is + // smaller than a single blob, so granular mode kicks in, and the + // Materialized-rep arm of `write_granular_chunk` steps from one + // `rep == 0` boundary to the next so a record never spans pages. + let value_size = 32 * 1024; + let page_byte_limit = 16 * 1024; + let values_per_record = 3; + let num_records = 3; + let num_values = values_per_record * num_records; + + // rep levels: 0, 1, 1, 0, 1, 1, 0, 1, 1 + let mut rep_levels = Vec::with_capacity(num_values); + for _ in 0..num_records { + rep_levels.push(0i16); + rep_levels.extend(std::iter::repeat_n(1i16, values_per_record - 1)); + } + let def_levels = vec![1i16; num_values]; + + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_data_page_size_limit(page_byte_limit) + .build(); + + let data: Vec<_> = (0..num_values) + .map(|i| ByteArray::from(vec![i as u8; value_size])) + .collect(); + let pages = write_and_collect_pages::( + props, + 1, + 1, + &data, + Some(&def_levels), + Some(&rep_levels), + ); + let data_pages = pages.data_pages; + + // The Materialized-rep arm groups levels by record, and each + // record's bytes blow the page byte limit on its own, so we get + // exactly one page per record. + assert_eq!( + data_pages.len(), + num_records, + "expected one data page per record, got {data_pages:?}" + ); + for (bytes, n_values) in &data_pages { + assert_eq!( + *n_values as usize, values_per_record, + "each page must hold a whole record's leaves, got {data_pages:?}" + ); + // Each page is one full record (its leaves cannot be split), + // so allow up to `values_per_record` blobs of payload plus a + // small fudge for level encoding overhead. + let upper_bound = values_per_record * (value_size + 16); + assert!( + *bytes <= upper_bound, + "page size {bytes} exceeds whole-record bound ({upper_bound}); pages {data_pages:?}" + ); + } + } + + #[test] + fn test_column_writer_caps_page_size_with_nullable_large_values() { + // Coverage for `LevelDataRef::value_count` on Materialized def + // levels: a nullable column with mixed nulls and large values. + // `value_count` must return the actual non-null count so the + // byte estimate reflects bytes that will actually be written, + // not the level count. + let value_size = 32 * 1024; + let page_byte_limit = 16 * 1024; + let num_levels = 32; + + // Alternating null / non-null: 16 nulls and 16 values. + let def_levels: Vec = (0..num_levels as i16).map(|i| i % 2).collect(); + let num_values = def_levels.iter().filter(|&&d| d == 1).count(); + + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_data_page_size_limit(page_byte_limit) + .build(); + + let data: Vec<_> = (0..num_values) + .map(|i| ByteArray::from(vec![i as u8; value_size])) + .collect(); + let pages = + write_and_collect_pages::(props, 1, 0, &data, Some(&def_levels), None); + let data_pages: Vec<_> = pages.data_pages.iter().map(|(size, _)| *size).collect(); + + // With 16 actual values of 32 KiB each and a 16 KiB page limit, + // every non-null value should get its own page (plus possibly + // adjacent nulls). At minimum, the number of pages must be + // roughly the value count, not 1 (which is what `main` produced). + assert!( + data_pages.len() >= num_values / 2, + "expected at least {} pages for {num_values} large values, got {} pages: {data_pages:?}", + num_values / 2, + data_pages.len(), + ); + // No page contains more than ~one value's worth of payload bytes. + for size in &data_pages { + assert!( + *size <= value_size + 64, + "page size {size} exceeds one-value bound; pages {data_pages:?}" + ); + } + } + + #[test] + fn test_column_writer_dict_enabled_large_values_post_spill() { + // While dictionary encoding is active, `has_dictionary()` short- + // circuits `estimated_value_bytes` — the byte estimate is plain- + // encoded size but dict-encoded pages only store small RLE + // indices, so we'd otherwise shrink pages spuriously. Once the + // dictionary spills (each value is large + unique), plain + // encoding takes over and the byte-budget sub-batch kicks in. + // + // This test makes sure the writer survives that transition and + // produces bounded pages thereafter. + let value_size = 64 * 1024; + let page_byte_limit = 16 * 1024; + let num_rows = 32; + + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(true) + // Force a small dict so it spills quickly even though + // each value here is unique. + .set_dictionary_page_size_limit(1024) + .set_data_page_size_limit(page_byte_limit) + // Small mini-batches so dict fallback happens part-way + // through the input, leaving subsequent mini-batches to + // exercise the post-spill plain-encoding path that the + // page-size fix actually targets. + .set_write_batch_size(4) + .build(); + + let data: Vec<_> = (0..num_rows) + .map(|i| ByteArray::from(vec![i as u8; value_size])) + .collect(); + let pages = write_and_collect_pages::(props, 0, 0, &data, None, None); + let data_pages: Vec<_> = pages.data_pages.iter().map(|(size, _)| *size).collect(); + + // After spill, plain encoding writes one ~64 KiB value per page. + // Without the fix, post-spill writes still buffered all 32 + // values into a single ~2 MiB page. + assert!( + data_pages.len() >= num_rows / 2, + "expected >= {} data pages after dict spill, got {} ({data_pages:?})", + num_rows / 2, + data_pages.len(), + ); + for size in &data_pages { + assert!( + *size <= value_size + 64, + "page size {size} exceeds one-value bound; pages {data_pages:?}" + ); + } + } + + #[test] + fn test_column_writer_caps_dictionary_page_size() { + // A column of large *distinct* values with dictionary encoding on: + // the dictionary page accumulates the values themselves, and its + // spill check runs only once per mini-batch. Without bounding the + // dictionary-encoding mini-batch, one `write_batch_size` mini-batch + // would intern `write_batch_size * value_size` bytes into the + // dictionary page before the check fires (~16 MiB here). The chunker + // must sub-batch the dictionary-encoding phase too. + let value_size = 8 * 1024; + let dict_page_limit = 64 * 1024; + let num_rows = 2048; + + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(true) + .set_dictionary_page_size_limit(dict_page_limit) + .build(); + + let data: Vec<_> = (0..num_rows) + .map(|i| { + // each value distinct, so the dictionary cannot dedup them + let mut v = vec![0u8; value_size]; + v[..8].copy_from_slice(&(i as u64).to_le_bytes()); + ByteArray::from(v) + }) + .collect(); + let pages = write_and_collect_pages::(props, 0, 0, &data, None, None); + let dict_page_size = pages.dict_page_size; + + assert!( + dict_page_size > 0, + "expected the column to dictionary-encode" + ); + // Bounded near the limit (~2x from the post-mini-batch check). Before + // the fix the dictionary page reached num_rows * value_size (~16 MiB, + // 256x the limit). + assert!( + dict_page_size <= 3 * dict_page_limit, + "dictionary page {dict_page_size} exceeds 3x the {dict_page_limit} limit", + ); + } + + #[test] + fn test_column_writer_caps_page_size_for_fixed_len_byte_array() { + // Coverage for `ParquetValueType::byte_size` override on + // `FixedLenByteArray`. With `type_length = 1`, each plain-encoded + // value is one byte, so a 4-byte page byte limit forces the + // sub-batch sizer to write ~4 values per page rather than one + // page for the whole batch. + let page_byte_limit = 4; + let num_values = 128; + + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_data_page_size_limit(page_byte_limit) + .build(); + + let data: Vec<_> = (0..num_values) + .map(|i| { + let mut fla = FixedLenByteArray::default(); + fla.set_data(Bytes::from(vec![i as u8])); + fla + }) + .collect(); + let pages = + write_and_collect_pages::(props, 0, 0, &data, None, None); + let data_pages: Vec<_> = pages.data_pages.iter().map(|(size, _)| *size).collect(); + + // Without the fix this is a single 128-byte page; with the fix + // the byte budget caps each page at ~`page_byte_limit` bytes. + assert!( + data_pages.len() >= num_values / 8, + "expected pages capped by byte budget, got {data_pages:?}" + ); + for size in &data_pages { + assert!( + *size <= page_byte_limit * 4, + "page size {size} larger than expected; pages {data_pages:?}" + ); + } + } + + #[test] + fn test_bool_statistics() { + let stats = statistics_roundtrip::(&[true, false, false, true]); + // Booleans have an unsigned sort order and so are not compatible + // with the deprecated `min` and `max` statistics + assert!(!stats.is_min_max_backwards_compatible()); + if let Statistics::Boolean(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &false); + assert_eq!(stats.max_opt().unwrap(), &true); + } else { + panic!("expecting Statistics::Boolean, got {stats:?}"); + } + } + + #[test] + fn test_int32_statistics() { + let stats = statistics_roundtrip::(&[-1, 3, -2, 2]); + assert!(stats.is_min_max_backwards_compatible()); + if let Statistics::Int32(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &-2); + assert_eq!(stats.max_opt().unwrap(), &3); + } else { + panic!("expecting Statistics::Int32, got {stats:?}"); + } + } + + #[test] + fn test_int64_statistics() { + let stats = statistics_roundtrip::(&[-1, 3, -2, 2]); + assert!(stats.is_min_max_backwards_compatible()); + if let Statistics::Int64(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &-2); + assert_eq!(stats.max_opt().unwrap(), &3); + } else { + panic!("expecting Statistics::Int64, got {stats:?}"); + } + } + + #[test] + fn test_int96_statistics() { + let input = vec![ + Int96::from(vec![1, 20, 30]), + Int96::from(vec![3, 20, 10]), + Int96::from(vec![0, 20, 30]), + Int96::from(vec![2, 20, 30]), + ] + .into_iter() + .collect::>(); + + let stats = statistics_roundtrip::(&input); + assert!(!stats.is_min_max_backwards_compatible()); + if let Statistics::Int96(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &Int96::from(vec![3, 20, 10])); + assert_eq!(stats.max_opt().unwrap(), &Int96::from(vec![2, 20, 30])); + } else { + panic!("expecting Statistics::Int96, got {stats:?}"); + } + } + + #[test] + fn test_float_statistics() { + let stats = statistics_roundtrip::(&[-1.0, 3.0, -2.0, 2.0]); + assert!(stats.is_min_max_backwards_compatible()); + if let Statistics::Float(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &-2.0); + assert_eq!(stats.max_opt().unwrap(), &3.0); + } else { + panic!("expecting Statistics::Float, got {stats:?}"); + } + } + + #[test] + fn test_double_statistics() { + let stats = statistics_roundtrip::(&[-1.0, 3.0, -2.0, 2.0]); + assert!(stats.is_min_max_backwards_compatible()); + if let Statistics::Double(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &-2.0); + assert_eq!(stats.max_opt().unwrap(), &3.0); + } else { + panic!("expecting Statistics::Double, got {stats:?}"); + } + } + + #[test] + fn test_byte_array_statistics() { + let input = ["aawaa", "zz", "aaw", "m", "qrs"] + .iter() + .map(|&s| s.into()) + .collect::>(); + + let stats = statistics_roundtrip::(&input); + assert!(!stats.is_min_max_backwards_compatible()); + if let Statistics::ByteArray(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &ByteArray::from("aaw")); + assert_eq!(stats.max_opt().unwrap(), &ByteArray::from("zz")); + } else { + panic!("expecting Statistics::ByteArray, got {stats:?}"); + } + } + + #[test] + fn test_fixed_len_byte_array_statistics() { + let input = ["aawaa", "zz ", "aaw ", "m ", "qrs "] + .iter() + .map(|&s| ByteArray::from(s).into()) + .collect::>(); + + let stats = statistics_roundtrip::(&input); + assert!(!stats.is_min_max_backwards_compatible()); + if let Statistics::FixedLenByteArray(stats) = stats { + let expected_min: FixedLenByteArray = ByteArray::from("aaw ").into(); + assert_eq!(stats.min_opt().unwrap(), &expected_min); + let expected_max: FixedLenByteArray = ByteArray::from("zz ").into(); + assert_eq!(stats.max_opt().unwrap(), &expected_max); + } else { + panic!("expecting Statistics::FixedLenByteArray, got {stats:?}"); + } + } + + #[test] + fn test_column_writer_check_float16_min_max() { + let input = [ + -f16::ONE, + f16::from_f32(3.0), + -f16::from_f32(2.0), + f16::from_f32(2.0), + ] + .into_iter() + .map(|s| ByteArray::from(s).into()) + .collect::>(); + + let stats = float16_statistics_roundtrip(&input); + assert!(stats.is_min_max_backwards_compatible()); + assert_eq!( + stats.min_opt().unwrap(), + &ByteArray::from(-f16::from_f32(2.0)) + ); + assert_eq!( + stats.max_opt().unwrap(), + &ByteArray::from(f16::from_f32(3.0)) + ); + } + + #[test] + fn test_column_writer_check_float16_nan_middle() { + let input = [f16::ONE, f16::NAN, f16::ONE + f16::ONE] + .into_iter() + .map(|s| ByteArray::from(s).into()) + .collect::>(); + + let stats = float16_statistics_roundtrip(&input); + assert!(stats.is_min_max_backwards_compatible()); + assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::ONE)); + assert_eq!( + stats.max_opt().unwrap(), + &ByteArray::from(f16::ONE + f16::ONE) + ); + } + + #[test] + fn test_float16_statistics_nan_middle() { + let input = [f16::ONE, f16::NAN, f16::ONE + f16::ONE] + .into_iter() + .map(|s| ByteArray::from(s).into()) + .collect::>(); + + let stats = float16_statistics_roundtrip(&input); + assert!(stats.is_min_max_backwards_compatible()); + assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::ONE)); + assert_eq!( + stats.max_opt().unwrap(), + &ByteArray::from(f16::ONE + f16::ONE) + ); + } + + #[test] + fn test_float16_statistics_nan_start() { + let input = [f16::NAN, f16::ONE, f16::ONE + f16::ONE] + .into_iter() + .map(|s| ByteArray::from(s).into()) + .collect::>(); + + let stats = float16_statistics_roundtrip(&input); + assert!(stats.is_min_max_backwards_compatible()); + assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::ONE)); + assert_eq!( + stats.max_opt().unwrap(), + &ByteArray::from(f16::ONE + f16::ONE) + ); + } + + #[test] + fn test_float16_statistics_nan_only() { + let input = [f16::NAN, f16::NAN] + .into_iter() + .map(|s| ByteArray::from(s).into()) + .collect::>(); + + let stats = float16_statistics_roundtrip(&input); + assert!(stats.min_bytes_opt().is_none()); + assert!(stats.max_bytes_opt().is_none()); + assert!(stats.is_min_max_backwards_compatible()); + } + + #[test] + fn test_float16_statistics_zero_only() { + let input = [f16::ZERO] + .into_iter() + .map(|s| ByteArray::from(s).into()) + .collect::>(); + + let stats = float16_statistics_roundtrip(&input); + assert!(stats.is_min_max_backwards_compatible()); + assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::NEG_ZERO)); + assert_eq!(stats.max_opt().unwrap(), &ByteArray::from(f16::ZERO)); + } + + #[test] + fn test_float16_statistics_neg_zero_only() { + let input = [f16::NEG_ZERO] + .into_iter() + .map(|s| ByteArray::from(s).into()) + .collect::>(); + + let stats = float16_statistics_roundtrip(&input); + assert!(stats.is_min_max_backwards_compatible()); + assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::NEG_ZERO)); + assert_eq!(stats.max_opt().unwrap(), &ByteArray::from(f16::ZERO)); + } + + #[test] + fn test_float16_statistics_zero_min() { + let input = [f16::ZERO, f16::ONE, f16::NAN, f16::PI] + .into_iter() + .map(|s| ByteArray::from(s).into()) + .collect::>(); + + let stats = float16_statistics_roundtrip(&input); + assert!(stats.is_min_max_backwards_compatible()); + assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::NEG_ZERO)); + assert_eq!(stats.max_opt().unwrap(), &ByteArray::from(f16::PI)); + } + + #[test] + fn test_float16_statistics_neg_zero_max() { + let input = [f16::NEG_ZERO, f16::NEG_ONE, f16::NAN, -f16::PI] + .into_iter() + .map(|s| ByteArray::from(s).into()) + .collect::>(); + + let stats = float16_statistics_roundtrip(&input); + assert!(stats.is_min_max_backwards_compatible()); + assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(-f16::PI)); + assert_eq!(stats.max_opt().unwrap(), &ByteArray::from(f16::ZERO)); + } + + #[test] + fn test_float_statistics_nan_middle() { + let stats = statistics_roundtrip::(&[1.0, f32::NAN, 2.0]); + assert!(stats.is_min_max_backwards_compatible()); + if let Statistics::Float(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &1.0); + assert_eq!(stats.max_opt().unwrap(), &2.0); + } else { + panic!("expecting Statistics::Float"); + } + } + + #[test] + fn test_float_statistics_nan_start() { + let stats = statistics_roundtrip::(&[f32::NAN, 1.0, 2.0]); + assert!(stats.is_min_max_backwards_compatible()); + if let Statistics::Float(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &1.0); + assert_eq!(stats.max_opt().unwrap(), &2.0); + } else { + panic!("expecting Statistics::Float"); + } + } + + #[test] + fn test_float_statistics_nan_only() { + let stats = statistics_roundtrip::(&[f32::NAN, f32::NAN]); + assert!(stats.min_bytes_opt().is_none()); + assert!(stats.max_bytes_opt().is_none()); + assert!(stats.is_min_max_backwards_compatible()); + assert!(matches!(stats, Statistics::Float(_))); + } + + #[test] + fn test_float_statistics_zero_only() { + let stats = statistics_roundtrip::(&[0.0]); + assert!(stats.is_min_max_backwards_compatible()); + if let Statistics::Float(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &-0.0); + assert!(stats.min_opt().unwrap().is_sign_negative()); + assert_eq!(stats.max_opt().unwrap(), &0.0); + assert!(stats.max_opt().unwrap().is_sign_positive()); + } else { + panic!("expecting Statistics::Float"); + } + } + + #[test] + fn test_float_statistics_neg_zero_only() { + let stats = statistics_roundtrip::(&[-0.0]); + assert!(stats.is_min_max_backwards_compatible()); + if let Statistics::Float(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &-0.0); + assert!(stats.min_opt().unwrap().is_sign_negative()); + assert_eq!(stats.max_opt().unwrap(), &0.0); + assert!(stats.max_opt().unwrap().is_sign_positive()); + } else { + panic!("expecting Statistics::Float"); + } + } + + #[test] + fn test_float_statistics_zero_min() { + let stats = statistics_roundtrip::(&[0.0, 1.0, f32::NAN, 2.0]); + assert!(stats.is_min_max_backwards_compatible()); + if let Statistics::Float(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &-0.0); + assert!(stats.min_opt().unwrap().is_sign_negative()); + assert_eq!(stats.max_opt().unwrap(), &2.0); + } else { + panic!("expecting Statistics::Float"); + } + } + + #[test] + fn test_float_statistics_neg_zero_max() { + let stats = statistics_roundtrip::(&[-0.0, -1.0, f32::NAN, -2.0]); + assert!(stats.is_min_max_backwards_compatible()); + if let Statistics::Float(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &-2.0); + assert_eq!(stats.max_opt().unwrap(), &0.0); + assert!(stats.max_opt().unwrap().is_sign_positive()); + } else { + panic!("expecting Statistics::Float"); + } + } + + #[test] + fn test_double_statistics_nan_middle() { + let stats = statistics_roundtrip::(&[1.0, f64::NAN, 2.0]); + assert!(stats.is_min_max_backwards_compatible()); + if let Statistics::Double(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &1.0); + assert_eq!(stats.max_opt().unwrap(), &2.0); + } else { + panic!("expecting Statistics::Double"); + } + } + + #[test] + fn test_double_statistics_nan_start() { + let stats = statistics_roundtrip::(&[f64::NAN, 1.0, 2.0]); + assert!(stats.is_min_max_backwards_compatible()); + if let Statistics::Double(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &1.0); + assert_eq!(stats.max_opt().unwrap(), &2.0); + } else { + panic!("expecting Statistics::Double"); + } + } + + #[test] + fn test_double_statistics_nan_only() { + let stats = statistics_roundtrip::(&[f64::NAN, f64::NAN]); + assert!(stats.min_bytes_opt().is_none()); + assert!(stats.max_bytes_opt().is_none()); + assert!(matches!(stats, Statistics::Double(_))); + assert!(stats.is_min_max_backwards_compatible()); + } + + #[test] + fn test_double_statistics_zero_only() { + let stats = statistics_roundtrip::(&[0.0]); + assert!(stats.is_min_max_backwards_compatible()); + if let Statistics::Double(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &-0.0); + assert!(stats.min_opt().unwrap().is_sign_negative()); + assert_eq!(stats.max_opt().unwrap(), &0.0); + assert!(stats.max_opt().unwrap().is_sign_positive()); + } else { + panic!("expecting Statistics::Double"); + } + } + + #[test] + fn test_double_statistics_neg_zero_only() { + let stats = statistics_roundtrip::(&[-0.0]); + assert!(stats.is_min_max_backwards_compatible()); + if let Statistics::Double(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &-0.0); + assert!(stats.min_opt().unwrap().is_sign_negative()); + assert_eq!(stats.max_opt().unwrap(), &0.0); + assert!(stats.max_opt().unwrap().is_sign_positive()); + } else { + panic!("expecting Statistics::Double"); + } + } + + #[test] + fn test_double_statistics_zero_min() { + let stats = statistics_roundtrip::(&[0.0, 1.0, f64::NAN, 2.0]); + assert!(stats.is_min_max_backwards_compatible()); + if let Statistics::Double(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &-0.0); + assert!(stats.min_opt().unwrap().is_sign_negative()); + assert_eq!(stats.max_opt().unwrap(), &2.0); + } else { + panic!("expecting Statistics::Double"); + } + } + + #[test] + fn test_double_statistics_neg_zero_max() { + let stats = statistics_roundtrip::(&[-0.0, -1.0, f64::NAN, -2.0]); + assert!(stats.is_min_max_backwards_compatible()); + if let Statistics::Double(stats) = stats { + assert_eq!(stats.min_opt().unwrap(), &-2.0); + assert_eq!(stats.max_opt().unwrap(), &0.0); + assert!(stats.max_opt().unwrap().is_sign_positive()); + } else { + panic!("expecting Statistics::Double"); + } + } + + #[test] + fn test_compare_greater_byte_array_decimals() { + assert!(!compare_greater_byte_array_decimals(&[], &[],),); + assert!(compare_greater_byte_array_decimals(&[1u8,], &[],),); + assert!(!compare_greater_byte_array_decimals(&[], &[1u8,],),); + assert!(compare_greater_byte_array_decimals(&[1u8,], &[0u8,],),); + assert!(!compare_greater_byte_array_decimals(&[1u8,], &[1u8,],),); + assert!(compare_greater_byte_array_decimals(&[1u8, 0u8,], &[0u8,],),); + assert!(!compare_greater_byte_array_decimals( + &[0u8, 1u8,], + &[1u8, 0u8,], + ),); + assert!(!compare_greater_byte_array_decimals( + &[255u8, 35u8, 0u8, 0u8,], + &[0u8,], + ),); + assert!(compare_greater_byte_array_decimals( + &[0u8,], + &[255u8, 35u8, 0u8, 0u8,], + ),); + } + + #[test] + fn test_column_index_with_null_pages() { + // write a single page of all nulls + let page_writer = get_test_page_writer(); + let props = Default::default(); + let mut writer = get_test_column_writer::(page_writer, 1, 0, props); + writer.write_batch(&[], Some(&[0, 0, 0, 0]), None).unwrap(); + + let r = writer.close().unwrap(); + assert!(r.column_index.is_some()); + let col_idx = r.column_index.unwrap(); + let col_idx = match col_idx { + ColumnIndexMetaData::INT32(col_idx) => col_idx, + _ => panic!("wrong stats type"), + }; + // null_pages should be true for page 0 + assert!(col_idx.is_null_page(0)); + // min and max should be empty byte arrays + assert!(col_idx.min_value(0).is_none()); + assert!(col_idx.max_value(0).is_none()); + // null_counts should be defined and be 4 for page 0 + assert!(col_idx.null_count(0).is_some()); + assert_eq!(col_idx.null_count(0), Some(4)); + // there is no repetition so rep histogram should be absent + assert!(col_idx.repetition_level_histogram(0).is_none()); + // definition_level_histogram should be present and should be 0:4, 1:0 + assert!(col_idx.definition_level_histogram(0).is_some()); + assert_eq!(col_idx.definition_level_histogram(0).unwrap(), &[4, 0]); + } + + #[test] + fn test_column_offset_index_metadata() { + // write data + // and check the offset index and column index + let page_writer = get_test_page_writer(); + let props = Default::default(); + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + writer.write_batch(&[1, 2, 3, 4], None, None).unwrap(); + // first page + writer.flush_data_pages().unwrap(); + // second page + writer.write_batch(&[4, 8, 2, -5], None, None).unwrap(); + + let r = writer.close().unwrap(); + let column_index = r.column_index.unwrap(); + let offset_index = r.offset_index.unwrap(); + + assert_eq!(8, r.rows_written); + + // column index + let column_index = match column_index { + ColumnIndexMetaData::INT32(column_index) => column_index, + _ => panic!("wrong stats type"), + }; + assert_eq!(2, column_index.num_pages()); + assert_eq!(2, offset_index.page_locations.len()); + assert_eq!(BoundaryOrder::UNORDERED, column_index.boundary_order); + for idx in 0..2 { + assert!(!column_index.is_null_page(idx)); + assert_eq!(0, column_index.null_count(0).unwrap()); + } + + if let Some(stats) = r.metadata.statistics() { + assert_eq!(stats.null_count_opt(), Some(0)); + assert_eq!(stats.distinct_count_opt(), None); + if let Statistics::Int32(stats) = stats { + // first page is [1,2,3,4] + // second page is [-5,2,4,8] + // note that we don't increment here, as this is a non BinaryArray type. + assert_eq!(stats.min_opt(), column_index.min_value(1)); + assert_eq!(stats.max_opt(), column_index.max_value(1)); + } else { + panic!("expecting Statistics::Int32"); + } + } else { + panic!("metadata missing statistics"); + } + + // page location + assert_eq!(0, offset_index.page_locations[0].first_row_index); + assert_eq!(4, offset_index.page_locations[1].first_row_index); + } + + /// Verify min/max value truncation in the column index works as expected + #[test] + fn test_column_offset_index_metadata_truncating() { + // write data + // and check the offset index and column index + let page_writer = get_test_page_writer(); + let props = WriterProperties::builder() + .set_statistics_truncate_length(None) // disable column index truncation + .build() + .into(); + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + + let mut data = vec![FixedLenByteArray::default(); 3]; + // This is the expected min value - "aaa..." + data[0].set_data(Bytes::from(vec![97_u8; 200])); + // This is the expected max value - "ZZZ..." + data[1].set_data(Bytes::from(vec![112_u8; 200])); + data[2].set_data(Bytes::from(vec![98_u8; 200])); + + writer.write_batch(&data, None, None).unwrap(); + + writer.flush_data_pages().unwrap(); + + let r = writer.close().unwrap(); + let column_index = r.column_index.unwrap(); + let offset_index = r.offset_index.unwrap(); + + let column_index = match column_index { + ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(column_index) => column_index, + _ => panic!("wrong stats type"), + }; + + assert_eq!(3, r.rows_written); + + // column index + assert_eq!(1, column_index.num_pages()); + assert_eq!(1, offset_index.page_locations.len()); + assert_eq!(BoundaryOrder::ASCENDING, column_index.boundary_order); + assert!(!column_index.is_null_page(0)); + assert_eq!(Some(0), column_index.null_count(0)); + + if let Some(stats) = r.metadata.statistics() { + assert_eq!(stats.null_count_opt(), Some(0)); + assert_eq!(stats.distinct_count_opt(), None); + if let Statistics::FixedLenByteArray(stats) = stats { + let column_index_min_value = column_index.min_value(0).unwrap(); + let column_index_max_value = column_index.max_value(0).unwrap(); + + // Column index stats are truncated, while the column chunk's aren't. + assert_ne!(stats.min_bytes_opt().unwrap(), column_index_min_value); + assert_ne!(stats.max_bytes_opt().unwrap(), column_index_max_value); + + assert_eq!( + column_index_min_value.len(), + DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH.unwrap() + ); + assert_eq!(column_index_min_value, &[97_u8; 64]); + assert_eq!( + column_index_max_value.len(), + DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH.unwrap() + ); + + // We expect the last byte to be incremented + assert_eq!( + *column_index_max_value.last().unwrap(), + *column_index_max_value.first().unwrap() + 1 + ); + } else { + panic!("expecting Statistics::FixedLenByteArray"); + } + } else { + panic!("metadata missing statistics"); + } + } + + #[test] + fn test_column_offset_index_truncating_spec_example() { + // write data + // and check the offset index and column index + let page_writer = get_test_page_writer(); + + // Truncate values at 1 byte + let builder = WriterProperties::builder().set_column_index_truncate_length(Some(1)); + let props = Arc::new(builder.build()); + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + + let mut data = vec![FixedLenByteArray::default(); 1]; + // This is the expected min value + data[0].set_data(Bytes::from(String::from("Blart Versenwald III"))); + + writer.write_batch(&data, None, None).unwrap(); + + writer.flush_data_pages().unwrap(); + + let r = writer.close().unwrap(); + let column_index = r.column_index.unwrap(); + let offset_index = r.offset_index.unwrap(); + + let column_index = match column_index { + ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(column_index) => column_index, + _ => panic!("wrong stats type"), + }; + + assert_eq!(1, r.rows_written); + + // column index + assert_eq!(1, column_index.num_pages()); + assert_eq!(1, offset_index.page_locations.len()); + assert_eq!(BoundaryOrder::ASCENDING, column_index.boundary_order); + assert!(!column_index.is_null_page(0)); + assert_eq!(Some(0), column_index.null_count(0)); + + if let Some(stats) = r.metadata.statistics() { + assert_eq!(stats.null_count_opt(), Some(0)); + assert_eq!(stats.distinct_count_opt(), None); + if let Statistics::FixedLenByteArray(_stats) = stats { + let column_index_min_value = column_index.min_value(0).unwrap(); + let column_index_max_value = column_index.max_value(0).unwrap(); + + assert_eq!(column_index_min_value.len(), 1); + assert_eq!(column_index_max_value.len(), 1); + + assert_eq!("B".as_bytes(), column_index_min_value); + assert_eq!("C".as_bytes(), column_index_max_value); + + assert_ne!(column_index_min_value, stats.min_bytes_opt().unwrap()); + assert_ne!(column_index_max_value, stats.max_bytes_opt().unwrap()); + } else { + panic!("expecting Statistics::FixedLenByteArray"); + } + } else { + panic!("metadata missing statistics"); + } + } + + #[test] + fn test_float16_min_max_no_truncation() { + // Even if we set truncation to occur at 1 byte, we should not truncate for Float16 + let builder = WriterProperties::builder().set_column_index_truncate_length(Some(1)); + let props = Arc::new(builder.build()); + let page_writer = get_test_page_writer(); + let mut writer = get_test_float16_column_writer(page_writer, props); + + let expected_value = f16::PI.to_le_bytes().to_vec(); + let data = vec![ByteArray::from(expected_value.clone()).into()]; + writer.write_batch(&data, None, None).unwrap(); + writer.flush_data_pages().unwrap(); + + let r = writer.close().unwrap(); + + // stats should still be written + // ensure bytes weren't truncated for column index + let column_index = r.column_index.unwrap(); + let column_index = match column_index { + ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(column_index) => column_index, + _ => panic!("wrong stats type"), + }; + let column_index_min_bytes = column_index.min_value(0).unwrap(); + let column_index_max_bytes = column_index.max_value(0).unwrap(); + assert_eq!(expected_value, column_index_min_bytes); + assert_eq!(expected_value, column_index_max_bytes); + + // ensure bytes weren't truncated for statistics + let stats = r.metadata.statistics().unwrap(); + if let Statistics::FixedLenByteArray(stats) = stats { + let stats_min_bytes = stats.min_bytes_opt().unwrap(); + let stats_max_bytes = stats.max_bytes_opt().unwrap(); + assert_eq!(expected_value, stats_min_bytes); + assert_eq!(expected_value, stats_max_bytes); + } else { + panic!("expecting Statistics::FixedLenByteArray"); + } + } + + #[test] + fn test_decimal_min_max_no_truncation() { + // Even if we set truncation to occur at 1 byte, we should not truncate for Decimal + let builder = WriterProperties::builder().set_column_index_truncate_length(Some(1)); + let props = Arc::new(builder.build()); + let page_writer = get_test_page_writer(); + let mut writer = + get_test_decimals_column_writer::(page_writer, 0, 0, props); + + let expected_value = vec![ + 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 179u8, 172u8, 19u8, 35u8, + 231u8, 90u8, 0u8, 0u8, + ]; + let data = vec![ByteArray::from(expected_value.clone()).into()]; + writer.write_batch(&data, None, None).unwrap(); + writer.flush_data_pages().unwrap(); + + let r = writer.close().unwrap(); + + // stats should still be written + // ensure bytes weren't truncated for column index + let column_index = r.column_index.unwrap(); + let column_index = match column_index { + ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(column_index) => column_index, + _ => panic!("wrong stats type"), + }; + let column_index_min_bytes = column_index.min_value(0).unwrap(); + let column_index_max_bytes = column_index.max_value(0).unwrap(); + assert_eq!(expected_value, column_index_min_bytes); + assert_eq!(expected_value, column_index_max_bytes); + + // ensure bytes weren't truncated for statistics + let stats = r.metadata.statistics().unwrap(); + if let Statistics::FixedLenByteArray(stats) = stats { + let stats_min_bytes = stats.min_bytes_opt().unwrap(); + let stats_max_bytes = stats.max_bytes_opt().unwrap(); + assert_eq!(expected_value, stats_min_bytes); + assert_eq!(expected_value, stats_max_bytes); + } else { + panic!("expecting Statistics::FixedLenByteArray"); + } + } + + #[test] + fn test_statistics_truncating_byte_array_default() { + let page_writer = get_test_page_writer(); + + // The default truncate length is 64 bytes + let props = WriterProperties::builder().build().into(); + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + + let mut data = vec![ByteArray::default(); 1]; + data[0].set_data(Bytes::from(String::from( + "This string is longer than 64 bytes, so it will almost certainly be truncated.", + ))); + writer.write_batch(&data, None, None).unwrap(); + writer.flush_data_pages().unwrap(); + + let r = writer.close().unwrap(); + + assert_eq!(1, r.rows_written); + + let stats = r.metadata.statistics().expect("statistics"); + if let Statistics::ByteArray(_stats) = stats { + let min_value = _stats.min_opt().unwrap(); + let max_value = _stats.max_opt().unwrap(); + + assert!(!_stats.min_is_exact()); + assert!(!_stats.max_is_exact()); + + let expected_len = 64; + assert_eq!(min_value.len(), expected_len); + assert_eq!(max_value.len(), expected_len); + + let expected_min = + "This string is longer than 64 bytes, so it will almost certainly".as_bytes(); + assert_eq!(expected_min, min_value.as_bytes()); + // note the max value is different from the min value: the last byte is incremented + let expected_max = + "This string is longer than 64 bytes, so it will almost certainlz".as_bytes(); + assert_eq!(expected_max, max_value.as_bytes()); + } else { + panic!("expecting Statistics::ByteArray"); + } + } + + #[test] + fn test_statistics_truncating_byte_array() { + let page_writer = get_test_page_writer(); + + const TEST_TRUNCATE_LENGTH: usize = 1; + + // Truncate values at 1 byte + let builder = + WriterProperties::builder().set_statistics_truncate_length(Some(TEST_TRUNCATE_LENGTH)); + let props = Arc::new(builder.build()); + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + + let mut data = vec![ByteArray::default(); 1]; + // This is the expected min value + data[0].set_data(Bytes::from(String::from("Blart Versenwald III"))); + + writer.write_batch(&data, None, None).unwrap(); + + writer.flush_data_pages().unwrap(); + + let r = writer.close().unwrap(); + + assert_eq!(1, r.rows_written); + + let stats = r.metadata.statistics().expect("statistics"); + assert_eq!(stats.null_count_opt(), Some(0)); + assert_eq!(stats.distinct_count_opt(), None); + if let Statistics::ByteArray(_stats) = stats { + let min_value = _stats.min_opt().unwrap(); + let max_value = _stats.max_opt().unwrap(); + + assert!(!_stats.min_is_exact()); + assert!(!_stats.max_is_exact()); + + assert_eq!(min_value.len(), TEST_TRUNCATE_LENGTH); + assert_eq!(max_value.len(), TEST_TRUNCATE_LENGTH); + + assert_eq!("B".as_bytes(), min_value.as_bytes()); + assert_eq!("C".as_bytes(), max_value.as_bytes()); + } else { + panic!("expecting Statistics::ByteArray"); + } + } + + #[test] + fn test_statistics_truncating_fixed_len_byte_array() { + let page_writer = get_test_page_writer(); + + const TEST_TRUNCATE_LENGTH: usize = 1; + + // Truncate values at 1 byte + let builder = + WriterProperties::builder().set_statistics_truncate_length(Some(TEST_TRUNCATE_LENGTH)); + let props = Arc::new(builder.build()); + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + + let mut data = vec![FixedLenByteArray::default(); 1]; + + const PSEUDO_DECIMAL_VALUE: i128 = 6541894651216648486512564456564654; + const PSEUDO_DECIMAL_BYTES: [u8; 16] = PSEUDO_DECIMAL_VALUE.to_be_bytes(); + + const EXPECTED_MIN: [u8; TEST_TRUNCATE_LENGTH] = [PSEUDO_DECIMAL_BYTES[0]]; // parquet specifies big-endian order for decimals + const EXPECTED_MAX: [u8; TEST_TRUNCATE_LENGTH] = + [PSEUDO_DECIMAL_BYTES[0].overflowing_add(1).0]; + + // This is the expected min value + data[0].set_data(Bytes::from(PSEUDO_DECIMAL_BYTES.as_slice())); + + writer.write_batch(&data, None, None).unwrap(); + + writer.flush_data_pages().unwrap(); + + let r = writer.close().unwrap(); + + assert_eq!(1, r.rows_written); + + let stats = r.metadata.statistics().expect("statistics"); + assert_eq!(stats.null_count_opt(), Some(0)); + assert_eq!(stats.distinct_count_opt(), None); + if let Statistics::FixedLenByteArray(_stats) = stats { + let min_value = _stats.min_opt().unwrap(); + let max_value = _stats.max_opt().unwrap(); + + assert!(!_stats.min_is_exact()); + assert!(!_stats.max_is_exact()); + + assert_eq!(min_value.len(), TEST_TRUNCATE_LENGTH); + assert_eq!(max_value.len(), TEST_TRUNCATE_LENGTH); + + assert_eq!(EXPECTED_MIN.as_slice(), min_value.as_bytes()); + assert_eq!(EXPECTED_MAX.as_slice(), max_value.as_bytes()); + + let reconstructed_min = i128::from_be_bytes([ + min_value.as_bytes()[0], + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + ]); + + let reconstructed_max = i128::from_be_bytes([ + max_value.as_bytes()[0], + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + ]); + + // check that the inner value is correctly bounded by the min/max + println!("min: {reconstructed_min} {PSEUDO_DECIMAL_VALUE}"); + assert!(reconstructed_min <= PSEUDO_DECIMAL_VALUE); + println!("max {reconstructed_max} {PSEUDO_DECIMAL_VALUE}"); + assert!(reconstructed_max >= PSEUDO_DECIMAL_VALUE); + } else { + panic!("expecting Statistics::FixedLenByteArray"); + } + } + + #[test] + fn test_send() { + fn test() {} + test::>(); + } + + #[test] + fn test_increment() { + let v = increment(vec![0, 0, 0]).unwrap(); + assert_eq!(&v, &[0, 0, 1]); + + // Handle overflow + let v = increment(vec![0, 255, 255]).unwrap(); + assert_eq!(&v, &[1, 0, 0]); + + // Return `None` if all bytes are u8::MAX + let v = increment(vec![255, 255, 255]); + assert!(v.is_none()); + } + + #[test] + fn test_increment_utf8() { + let test_inc = |o: &str, expected: &str| { + if let Ok(v) = String::from_utf8(increment_utf8(o).unwrap()) { + // Got the expected result... + assert_eq!(v, expected); + // and it's greater than the original string + assert!(*v > *o); + // Also show that BinaryArray level comparison works here + let mut greater = ByteArray::new(); + greater.set_data(Bytes::from(v)); + let mut original = ByteArray::new(); + original.set_data(Bytes::from(o.as_bytes().to_vec())); + assert!(greater > original); + } else { + panic!("Expected incremented UTF8 string to also be valid."); + } + }; + + // Basic ASCII case + test_inc("hello", "hellp"); + + // 1-byte ending in max 1-byte + test_inc("a\u{7f}", "b"); + + // 1-byte max should not truncate as it would need 2-byte code points + assert!(increment_utf8("\u{7f}\u{7f}").is_none()); + + // UTF8 string + test_inc("❤️🧡💛💚💙💜", "❤️🧡💛💚💙💝"); + + // 2-byte without overflow + test_inc("éééé", "éééê"); + + // 2-byte that overflows lowest byte + test_inc("\u{ff}\u{ff}", "\u{ff}\u{100}"); + + // 2-byte ending in max 2-byte + test_inc("a\u{7ff}", "b"); + + // Max 2-byte should not truncate as it would need 3-byte code points + assert!(increment_utf8("\u{7ff}\u{7ff}").is_none()); + + // 3-byte without overflow [U+800, U+800] -> [U+800, U+801] (note that these + // characters should render right to left). + test_inc("ࠀࠀ", "ࠀࠁ"); + + // 3-byte ending in max 3-byte + test_inc("a\u{ffff}", "b"); + + // Max 3-byte should not truncate as it would need 4-byte code points + assert!(increment_utf8("\u{ffff}\u{ffff}").is_none()); + + // 4-byte without overflow + test_inc("𐀀𐀀", "𐀀𐀁"); + + // 4-byte ending in max unicode + test_inc("a\u{10ffff}", "b"); + + // Max 4-byte should not truncate + assert!(increment_utf8("\u{10ffff}\u{10ffff}").is_none()); + + // Skip over surrogate pair range (0xD800..=0xDFFF) + //test_inc("a\u{D7FF}", "a\u{e000}"); + test_inc("a\u{D7FF}", "b"); + } + + #[test] + fn test_truncate_utf8() { + // No-op + let data = "❤️🧡💛💚💙💜"; + let r = truncate_utf8(data, data.len()).unwrap(); + assert_eq!(r.len(), data.len()); + assert_eq!(&r, data.as_bytes()); + + // We slice it away from the UTF8 boundary + let r = truncate_utf8(data, 13).unwrap(); + assert_eq!(r.len(), 10); + assert_eq!(&r, "❤️🧡".as_bytes()); + + // One multi-byte code point, and a length shorter than it, so we can't slice it + let r = truncate_utf8("\u{0836}", 1); + assert!(r.is_none()); + + // Test truncate and increment for max bounds on UTF-8 statistics + // 7-bit (i.e. ASCII) + let r = truncate_and_increment_utf8("yyyyyyyyy", 8).unwrap(); + assert_eq!(&r, "yyyyyyyz".as_bytes()); + + // 2-byte without overflow + let r = truncate_and_increment_utf8("ééééé", 7).unwrap(); + assert_eq!(&r, "ééê".as_bytes()); + + // 2-byte that overflows lowest byte + let r = truncate_and_increment_utf8("\u{ff}\u{ff}\u{ff}\u{ff}\u{ff}", 8).unwrap(); + assert_eq!(&r, "\u{ff}\u{ff}\u{ff}\u{100}".as_bytes()); + + // max 2-byte should not truncate as it would need 3-byte code points + let r = truncate_and_increment_utf8("߿߿߿߿߿", 8); + assert!(r.is_none()); + + // 3-byte without overflow [U+800, U+800, U+800] -> [U+800, U+801] (note that these + // characters should render right to left). + let r = truncate_and_increment_utf8("ࠀࠀࠀࠀ", 8).unwrap(); + assert_eq!(&r, "ࠀࠁ".as_bytes()); + + // max 3-byte should not truncate as it would need 4-byte code points + let r = truncate_and_increment_utf8("\u{ffff}\u{ffff}\u{ffff}", 8); + assert!(r.is_none()); + + // 4-byte without overflow + let r = truncate_and_increment_utf8("𐀀𐀀𐀀𐀀", 9).unwrap(); + assert_eq!(&r, "𐀀𐀁".as_bytes()); + + // max 4-byte should not truncate + let r = truncate_and_increment_utf8("\u{10ffff}\u{10ffff}", 8); + assert!(r.is_none()); + } + + #[test] + // Check fallback truncation of statistics that should be UTF-8, but aren't + // (see https://github.com/apache/arrow-rs/pull/6870). + fn test_byte_array_truncate_invalid_utf8_statistics() { + let message_type = " + message test_schema { + OPTIONAL BYTE_ARRAY a (UTF8); + } + "; + let schema = Arc::new(parse_message_type(message_type).unwrap()); + + // Create Vec containing non-UTF8 bytes + let data = vec![ByteArray::from(vec![128u8; 32]); 7]; + let def_levels = [1, 1, 1, 1, 0, 1, 0, 1, 0, 1]; + let file: File = tempfile::tempfile().unwrap(); + let props = Arc::new( + WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Chunk) + .set_statistics_truncate_length(Some(8)) + .build(), + ); + + let mut writer = SerializedFileWriter::new(&file, schema, props).unwrap(); + let mut row_group_writer = writer.next_row_group().unwrap(); + + let mut col_writer = row_group_writer.next_column().unwrap().unwrap(); + col_writer + .typed::() + .write_batch(&data, Some(&def_levels), None) + .unwrap(); + col_writer.close().unwrap(); + row_group_writer.close().unwrap(); + let file_metadata = writer.close().unwrap(); + let stats = file_metadata.row_group(0).column(0).statistics().unwrap(); + assert!(!stats.max_is_exact()); + // Truncation of invalid UTF-8 should fall back to binary truncation, so last byte should + // be incremented by 1. + assert_eq!( + stats.max_bytes_opt().map(|v| v.to_vec()), + Some([128, 128, 128, 128, 128, 128, 128, 129].to_vec()) + ); + } + + #[test] + fn test_increment_max_binary_chars() { + let r = increment(vec![0xFF, 0xFE, 0xFD, 0xFF, 0xFF]); + assert_eq!(&r.unwrap(), &[0xFF, 0xFE, 0xFE, 0x00, 0x00]); + + let incremented = increment(vec![0xFF, 0xFF, 0xFF]); + assert!(incremented.is_none()) + } + + #[test] + fn test_no_column_index_when_stats_disabled() { + // https://github.com/apache/arrow-rs/issues/6010 + // Test that column index is not created/written for all-nulls column when page + // statistics are disabled. + let descr = Arc::new(get_test_column_descr::(1, 0)); + let props = Arc::new( + WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::None) + .build(), + ); + let column_writer = get_column_writer(descr, props, get_test_page_writer()); + let mut writer = get_typed_column_writer::(column_writer); + + let data = Vec::new(); + let def_levels = vec![0; 10]; + writer.write_batch(&data, Some(&def_levels), None).unwrap(); + writer.flush_data_pages().unwrap(); + + let column_close_result = writer.close().unwrap(); + assert!(column_close_result.offset_index.is_some()); + assert!(column_close_result.column_index.is_none()); + } + + #[test] + fn test_no_offset_index_when_disabled() { + // Test that offset indexes can be disabled + let descr = Arc::new(get_test_column_descr::(1, 0)); + let props = Arc::new( + WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::None) + .set_offset_index_disabled(true) + .build(), + ); + let column_writer = get_column_writer(descr, props, get_test_page_writer()); + let mut writer = get_typed_column_writer::(column_writer); + + let data = Vec::new(); + let def_levels = vec![0; 10]; + writer.write_batch(&data, Some(&def_levels), None).unwrap(); + writer.flush_data_pages().unwrap(); + + let column_close_result = writer.close().unwrap(); + assert!(column_close_result.offset_index.is_none()); + assert!(column_close_result.column_index.is_none()); + } + + #[test] + fn test_offset_index_overridden() { + // Test that offset indexes are not disabled when gathering page statistics + let descr = Arc::new(get_test_column_descr::(1, 0)); + let props = Arc::new( + WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .set_offset_index_disabled(true) + .build(), + ); + let column_writer = get_column_writer(descr, props, get_test_page_writer()); + let mut writer = get_typed_column_writer::(column_writer); + + let data = Vec::new(); + let def_levels = vec![0; 10]; + writer.write_batch(&data, Some(&def_levels), None).unwrap(); + writer.flush_data_pages().unwrap(); + + let column_close_result = writer.close().unwrap(); + assert!(column_close_result.offset_index.is_some()); + assert!(column_close_result.column_index.is_some()); + } + + #[test] + fn test_boundary_order() -> Result<()> { + let descr = Arc::new(get_test_column_descr::(1, 0)); + // min max both ascending + let column_close_result = write_multiple_pages::( + &descr, + &[ + &[Some(-10), Some(10)], + &[Some(-5), Some(11)], + &[None], + &[Some(-5), Some(11)], + ], + )?; + let boundary_order = column_close_result + .column_index + .unwrap() + .get_boundary_order(); + assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING)); + + // min max both descending + let column_close_result = write_multiple_pages::( + &descr, + &[ + &[Some(10), Some(11)], + &[Some(5), Some(11)], + &[None], + &[Some(-5), Some(0)], + ], + )?; + let boundary_order = column_close_result + .column_index + .unwrap() + .get_boundary_order(); + assert_eq!(boundary_order, Some(BoundaryOrder::DESCENDING)); + + // min max both equal + let column_close_result = write_multiple_pages::( + &descr, + &[&[Some(10), Some(11)], &[None], &[Some(10), Some(11)]], + )?; + let boundary_order = column_close_result + .column_index + .unwrap() + .get_boundary_order(); + assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING)); + + // only nulls + let column_close_result = + write_multiple_pages::(&descr, &[&[None], &[None], &[None]])?; + let boundary_order = column_close_result + .column_index + .unwrap() + .get_boundary_order(); + assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING)); + + // one page + let column_close_result = + write_multiple_pages::(&descr, &[&[Some(-10), Some(10)]])?; + let boundary_order = column_close_result + .column_index + .unwrap() + .get_boundary_order(); + assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING)); + + // one non-null page + let column_close_result = + write_multiple_pages::(&descr, &[&[Some(-10), Some(10)], &[None]])?; + let boundary_order = column_close_result + .column_index + .unwrap() + .get_boundary_order(); + assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING)); + + // min max both unordered + let column_close_result = write_multiple_pages::( + &descr, + &[ + &[Some(10), Some(11)], + &[Some(11), Some(16)], + &[None], + &[Some(-5), Some(0)], + ], + )?; + let boundary_order = column_close_result + .column_index + .unwrap() + .get_boundary_order(); + assert_eq!(boundary_order, Some(BoundaryOrder::UNORDERED)); + + // min max both ordered in different orders + let column_close_result = write_multiple_pages::( + &descr, + &[ + &[Some(1), Some(9)], + &[Some(2), Some(8)], + &[None], + &[Some(3), Some(7)], + ], + )?; + let boundary_order = column_close_result + .column_index + .unwrap() + .get_boundary_order(); + assert_eq!(boundary_order, Some(BoundaryOrder::UNORDERED)); + + Ok(()) + } + + #[test] + fn test_boundary_order_logical_type() -> Result<()> { + // ensure that logical types account for different sort order than underlying + // physical type representation + let f16_descr = Arc::new(get_test_float16_column_descr(1, 0)); + let fba_descr = { + let tpe = SchemaType::primitive_type_builder( + "col", + FixedLenByteArrayType::get_physical_type(), + ) + .with_length(2) + .build()?; + Arc::new(ColumnDescriptor::new( + Arc::new(tpe), + 1, + 0, + ColumnPath::from("col"), + )) + }; + + let values: &[&[Option]] = &[ + &[Some(FixedLenByteArray::from(ByteArray::from(f16::ONE)))], + &[Some(FixedLenByteArray::from(ByteArray::from(f16::ZERO)))], + &[Some(FixedLenByteArray::from(ByteArray::from( + f16::NEG_ZERO, + )))], + &[Some(FixedLenByteArray::from(ByteArray::from(f16::NEG_ONE)))], + ]; + + // f16 descending + let column_close_result = + write_multiple_pages::(&f16_descr, values)?; + let boundary_order = column_close_result + .column_index + .unwrap() + .get_boundary_order(); + assert_eq!(boundary_order, Some(BoundaryOrder::DESCENDING)); + + // same bytes, but fba unordered + let column_close_result = + write_multiple_pages::(&fba_descr, values)?; + let boundary_order = column_close_result + .column_index + .unwrap() + .get_boundary_order(); + assert_eq!(boundary_order, Some(BoundaryOrder::UNORDERED)); + + Ok(()) + } + + #[test] + fn test_interval_stats_should_not_have_min_max() { + let input = [ + vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1], + vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2], + ] + .into_iter() + .map(|s| ByteArray::from(s).into()) + .collect::>(); + + let page_writer = get_test_page_writer(); + let mut writer = get_test_interval_column_writer(page_writer); + writer.write_batch(&input, None, None).unwrap(); + + let metadata = writer.close().unwrap().metadata; + let stats = if let Some(Statistics::FixedLenByteArray(stats)) = metadata.statistics() { + stats.clone() + } else { + panic!("metadata missing statistics"); + }; + assert!(stats.min_bytes_opt().is_none()); + assert!(stats.max_bytes_opt().is_none()); + } + + #[test] + #[cfg(feature = "arrow")] + fn test_column_writer_get_estimated_total_bytes() { + let page_writer = get_test_page_writer(); + let props = Default::default(); + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + assert_eq!(writer.get_estimated_total_bytes(), 0); + + writer.write_batch(&[1, 2, 3, 4], None, None).unwrap(); + writer.add_data_page().unwrap(); + let size_with_one_page = writer.get_estimated_total_bytes(); + assert_eq!(size_with_one_page, 20); + + writer.write_batch(&[5, 6, 7, 8], None, None).unwrap(); + writer.add_data_page().unwrap(); + let size_with_two_pages = writer.get_estimated_total_bytes(); + // different pages have different compressed lengths + assert_eq!(size_with_two_pages, 20 + 21); + } + + fn write_multiple_pages( + column_descr: &Arc, + pages: &[&[Option]], + ) -> Result { + let column_writer = get_column_writer( + column_descr.clone(), + Default::default(), + get_test_page_writer(), + ); + let mut writer = get_typed_column_writer::(column_writer); + + for &page in pages { + let values = page.iter().filter_map(Clone::clone).collect::>(); + let def_levels = page + .iter() + .map(|maybe_value| if maybe_value.is_some() { 1 } else { 0 }) + .collect::>(); + writer.write_batch(&values, Some(&def_levels), None)?; + writer.flush_data_pages()?; + } + + writer.close() + } + + /// Performs write-read roundtrip with randomly generated values and levels. + /// `max_size` is maximum number of values or levels (if `max_def_level` > 0) to write + /// for a column. + fn column_roundtrip_random( + props: WriterProperties, + max_size: usize, + min_value: T::T, + max_value: T::T, + max_def_level: i16, + max_rep_level: i16, + ) where + T::T: PartialOrd + SampleUniform + Copy, + { + let mut num_values: usize = 0; + + let mut buf: Vec = Vec::new(); + let def_levels = if max_def_level > 0 { + random_numbers_range(max_size, 0, max_def_level + 1, &mut buf); + for &dl in &buf[..] { + if dl == max_def_level { + num_values += 1; + } + } + Some(&buf[..]) + } else { + num_values = max_size; + None + }; + + let mut buf: Vec = Vec::new(); + let rep_levels = if max_rep_level > 0 { + random_numbers_range(max_size, 0, max_rep_level + 1, &mut buf); + buf[0] = 0; // Must start on record boundary + Some(&buf[..]) + } else { + None + }; + + let mut values: Vec = Vec::new(); + random_numbers_range(num_values, min_value, max_value, &mut values); + + column_roundtrip::(props, &values[..], def_levels, rep_levels); + } + + /// Performs write-read roundtrip and asserts written values and levels. + fn column_roundtrip( + props: WriterProperties, + values: &[T::T], + def_levels: Option<&[i16]>, + rep_levels: Option<&[i16]>, + ) { + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + + let max_def_level = match def_levels { + Some(buf) => *buf.iter().max().unwrap_or(&0i16), + None => 0i16, + }; + + let max_rep_level = match rep_levels { + Some(buf) => *buf.iter().max().unwrap_or(&0i16), + None => 0i16, + }; + + let mut max_batch_size = values.len(); + if let Some(levels) = def_levels { + max_batch_size = max_batch_size.max(levels.len()); + } + if let Some(levels) = rep_levels { + max_batch_size = max_batch_size.max(levels.len()); + } + + let mut writer = + get_test_column_writer::(page_writer, max_def_level, max_rep_level, Arc::new(props)); + + let values_written = writer.write_batch(values, def_levels, rep_levels).unwrap(); + assert_eq!(values_written, values.len()); + let result = writer.close().unwrap(); + + drop(write); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &result.metadata, + result.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(), + ); + let mut reader = get_test_column_reader::(page_reader, max_def_level, max_rep_level); + + let mut actual_values = Vec::with_capacity(max_batch_size); + let mut actual_def_levels = def_levels.map(|_| Vec::with_capacity(max_batch_size)); + let mut actual_rep_levels = rep_levels.map(|_| Vec::with_capacity(max_batch_size)); + + let (_, values_read, levels_read) = reader + .read_records( + max_batch_size, + actual_def_levels.as_mut(), + actual_rep_levels.as_mut(), + &mut actual_values, + ) + .unwrap(); + + // Assert values, definition and repetition levels. + + assert_eq!(&actual_values[..values_read], values); + match actual_def_levels { + Some(ref vec) => assert_eq!(Some(&vec[..levels_read]), def_levels), + None => assert_eq!(None, def_levels), + } + match actual_rep_levels { + Some(ref vec) => assert_eq!(Some(&vec[..levels_read]), rep_levels), + None => assert_eq!(None, rep_levels), + } + + // Assert written rows. + + if let Some(levels) = actual_rep_levels { + let mut actual_rows_written = 0; + for l in levels { + if l == 0 { + actual_rows_written += 1; + } + } + assert_eq!(actual_rows_written, result.rows_written); + } else if actual_def_levels.is_some() { + assert_eq!(levels_read as u64, result.rows_written); + } else { + assert_eq!(values_read as u64, result.rows_written); + } + } + + /// Performs write of provided values and returns column metadata of those values. + /// Used to test encoding support for column writer. + fn column_write_and_get_metadata( + props: WriterProperties, + values: &[T::T], + ) -> ColumnChunkMetaData { + let page_writer = get_test_page_writer(); + let props = Arc::new(props); + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + writer.write_batch(values, None, None).unwrap(); + writer.close().unwrap().metadata + } + + // Helper function to more compactly create a PageEncodingStats struct. + fn encoding_stats(page_type: PageType, encoding: Encoding, count: i32) -> PageEncodingStats { + PageEncodingStats { + page_type, + encoding, + count, + } + } + + // Function to use in tests for EncodingWriteSupport. This checks that dictionary + // offset and encodings to make sure that column writer uses provided by trait + // encodings. + fn check_encoding_write_support( + version: WriterVersion, + dict_enabled: bool, + data: &[T::T], + dictionary_page_offset: Option, + encodings: &[Encoding], + page_encoding_stats: &[PageEncodingStats], + ) { + let props = WriterProperties::builder() + .set_writer_version(version) + .set_dictionary_enabled(dict_enabled) + .build(); + let meta = column_write_and_get_metadata::(props, data); + assert_eq!(meta.dictionary_page_offset(), dictionary_page_offset); + assert_eq!(meta.encodings().collect::>(), encodings); + assert_eq!(meta.page_encoding_stats().unwrap(), page_encoding_stats); + } + + /// Returns column writer. + fn get_test_column_writer<'a, T: DataType>( + page_writer: Box, + max_def_level: i16, + max_rep_level: i16, + props: WriterPropertiesPtr, + ) -> ColumnWriterImpl<'a, T> { + let descr = Arc::new(get_test_column_descr::(max_def_level, max_rep_level)); + let column_writer = get_column_writer(descr, props, page_writer); + get_typed_column_writer::(column_writer) + } + + fn get_test_column_writer_with_path<'a, T: DataType>( + page_writer: Box, + max_def_level: i16, + max_rep_level: i16, + props: WriterPropertiesPtr, + path: ColumnPath, + ) -> ColumnWriterImpl<'a, T> { + let descr = Arc::new(get_test_column_descr_with_path::( + max_def_level, + max_rep_level, + path, + )); + let column_writer = get_column_writer(descr, props, page_writer); + get_typed_column_writer::(column_writer) + } + + /// Pages collected by [`write_and_collect_pages`]. + struct CollectedPages { + /// `(compressed byte size, value count)` for every data page, in order. + data_pages: Vec<(usize, u32)>, + /// Largest dictionary page seen, or 0 if the column wasn't dict-encoded. + dict_page_size: usize, + } + + /// Writes `data` (with optional def/rep levels) through a raw + /// `ColumnWriterImpl` configured by `props`, then re-reads the file and + /// returns its page layout. Shared by the page-size regression tests so + /// each only has to express its props, input, and assertions. + fn write_and_collect_pages( + props: WriterProperties, + max_def_level: i16, + max_rep_level: i16, + data: &[T::T], + def_levels: Option<&[i16]>, + rep_levels: Option<&[i16]>, + ) -> CollectedPages { + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let mut writer = + get_test_column_writer::(page_writer, max_def_level, max_rep_level, Arc::new(props)); + writer.write_batch(data, def_levels, rep_levels).unwrap(); + let r = writer.close().unwrap(); + drop(write); + + let read_props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let mut page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(read_props), + ) + .unwrap(), + ); + + let mut collected = CollectedPages { + data_pages: Vec::new(), + dict_page_size: 0, + }; + while let Some(page) = page_reader.get_next_page().unwrap() { + match page.page_type() { + PageType::DATA_PAGE | PageType::DATA_PAGE_V2 => { + collected + .data_pages + .push((page.buffer().len(), page.num_values())); + } + PageType::DICTIONARY_PAGE => { + collected.dict_page_size = collected.dict_page_size.max(page.buffer().len()); + } + _ => {} + } + } + collected + } + + /// Returns column reader. + fn get_test_column_reader( + page_reader: Box, + max_def_level: i16, + max_rep_level: i16, + ) -> ColumnReaderImpl { + let descr = Arc::new(get_test_column_descr::(max_def_level, max_rep_level)); + let column_reader = get_column_reader(descr, page_reader); + get_typed_column_reader::(column_reader) + } + + /// Returns descriptor for primitive column. + fn get_test_column_descr( + max_def_level: i16, + max_rep_level: i16, + ) -> ColumnDescriptor { + let path = ColumnPath::from("col"); + let tpe = SchemaType::primitive_type_builder("col", T::get_physical_type()) + // length is set for "encoding support" tests for FIXED_LEN_BYTE_ARRAY type, + // it should be no-op for other types + .with_length(1) + .build() + .unwrap(); + ColumnDescriptor::new(Arc::new(tpe), max_def_level, max_rep_level, path) + } + + fn get_test_column_descr_with_path( + max_def_level: i16, + max_rep_level: i16, + path: ColumnPath, + ) -> ColumnDescriptor { + let name = path.string(); + let tpe = SchemaType::primitive_type_builder(&name, T::get_physical_type()) + // length is set for "encoding support" tests for FIXED_LEN_BYTE_ARRAY type, + // it should be no-op for other types + .with_length(1) + .build() + .unwrap(); + ColumnDescriptor::new(Arc::new(tpe), max_def_level, max_rep_level, path) + } + + fn write_and_collect_page_values( + path: ColumnPath, + props: WriterPropertiesPtr, + data: &[i32], + ) -> Vec { + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let mut writer = + get_test_column_writer_with_path::(page_writer, 0, 0, props, path); + writer.write_batch(data, None, None).unwrap(); + let r = writer.close().unwrap(); + + drop(write); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let mut page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &r.metadata, + r.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(), + ); + + let mut values_per_page = Vec::new(); + while let Some(page) = page_reader.get_next_page().unwrap() { + assert_eq!(page.page_type(), PageType::DATA_PAGE); + values_per_page.push(page.num_values()); + } + + values_per_page + } + + /// Returns page writer that collects pages without serializing them. + fn get_test_page_writer() -> Box { + Box::new(TestPageWriter {}) + } + + struct TestPageWriter {} + + impl PageWriter for TestPageWriter { + fn write_page(&mut self, page: CompressedPage) -> Result { + let mut res = PageWriteSpec::new(); + res.page_type = page.page_type(); + res.uncompressed_size = page.uncompressed_size(); + res.compressed_size = page.compressed_size(); + res.num_values = page.num_values(); + res.offset = 0; + res.bytes_written = page.data().len() as u64; + Ok(res) + } + + fn close(&mut self) -> Result<()> { + Ok(()) + } + } + + /// Write data into parquet using [`get_test_page_writer`] and [`get_test_column_writer`] and returns generated statistics. + fn statistics_roundtrip(values: &[::T]) -> Statistics { + let page_writer = get_test_page_writer(); + let props = Default::default(); + let mut writer = get_test_column_writer::(page_writer, 0, 0, props); + writer.write_batch(values, None, None).unwrap(); + + let metadata = writer.close().unwrap().metadata; + if let Some(stats) = metadata.statistics() { + stats.clone() + } else { + panic!("metadata missing statistics"); + } + } + + /// Returns Decimals column writer. + fn get_test_decimals_column_writer( + page_writer: Box, + max_def_level: i16, + max_rep_level: i16, + props: WriterPropertiesPtr, + ) -> ColumnWriterImpl<'static, T> { + let descr = Arc::new(get_test_decimals_column_descr::( + max_def_level, + max_rep_level, + )); + let column_writer = get_column_writer(descr, props, page_writer); + get_typed_column_writer::(column_writer) + } + + /// Returns descriptor for Decimal type with primitive column. + fn get_test_decimals_column_descr( + max_def_level: i16, + max_rep_level: i16, + ) -> ColumnDescriptor { + let path = ColumnPath::from("col"); + let tpe = SchemaType::primitive_type_builder("col", T::get_physical_type()) + .with_length(16) + .with_logical_type(Some(LogicalType::decimal(2, 3))) + .with_scale(2) + .with_precision(3) + .build() + .unwrap(); + ColumnDescriptor::new(Arc::new(tpe), max_def_level, max_rep_level, path) + } + + fn float16_statistics_roundtrip( + values: &[FixedLenByteArray], + ) -> ValueStatistics { + let page_writer = get_test_page_writer(); + let mut writer = get_test_float16_column_writer(page_writer, Default::default()); + writer.write_batch(values, None, None).unwrap(); + + let metadata = writer.close().unwrap().metadata; + if let Some(Statistics::FixedLenByteArray(stats)) = metadata.statistics() { + stats.clone() + } else { + panic!("metadata missing statistics"); + } + } + + fn get_test_float16_column_writer( + page_writer: Box, + props: WriterPropertiesPtr, + ) -> ColumnWriterImpl<'static, FixedLenByteArrayType> { + let descr = Arc::new(get_test_float16_column_descr(0, 0)); + let column_writer = get_column_writer(descr, props, page_writer); + get_typed_column_writer::(column_writer) + } + + fn get_test_float16_column_descr(max_def_level: i16, max_rep_level: i16) -> ColumnDescriptor { + let path = ColumnPath::from("col"); + let tpe = + SchemaType::primitive_type_builder("col", FixedLenByteArrayType::get_physical_type()) + .with_length(2) + .with_logical_type(Some(LogicalType::Float16)) + .build() + .unwrap(); + ColumnDescriptor::new(Arc::new(tpe), max_def_level, max_rep_level, path) + } + + fn get_test_interval_column_writer( + page_writer: Box, + ) -> ColumnWriterImpl<'static, FixedLenByteArrayType> { + let descr = Arc::new(get_test_interval_column_descr()); + let column_writer = get_column_writer(descr, Default::default(), page_writer); + get_typed_column_writer::(column_writer) + } + + fn get_test_interval_column_descr() -> ColumnDescriptor { + let path = ColumnPath::from("col"); + let tpe = + SchemaType::primitive_type_builder("col", FixedLenByteArrayType::get_physical_type()) + .with_length(12) + .with_converted_type(ConvertedType::INTERVAL) + .build() + .unwrap(); + ColumnDescriptor::new(Arc::new(tpe), 0, 0, path) + } + + /// Returns column writer for UINT32 Column provided as ConvertedType only + fn get_test_unsigned_int_given_as_converted_column_writer<'a, T: DataType>( + page_writer: Box, + max_def_level: i16, + max_rep_level: i16, + props: WriterPropertiesPtr, + ) -> ColumnWriterImpl<'a, T> { + let descr = Arc::new(get_test_converted_type_unsigned_integer_column_descr::( + max_def_level, + max_rep_level, + )); + let column_writer = get_column_writer(descr, props, page_writer); + get_typed_column_writer::(column_writer) + } + + /// Returns column descriptor for UINT32 Column provided as ConvertedType only + fn get_test_converted_type_unsigned_integer_column_descr( + max_def_level: i16, + max_rep_level: i16, + ) -> ColumnDescriptor { + let path = ColumnPath::from("col"); + let tpe = SchemaType::primitive_type_builder("col", T::get_physical_type()) + .with_converted_type(ConvertedType::UINT_32) + .build() + .unwrap(); + ColumnDescriptor::new(Arc::new(tpe), max_def_level, max_rep_level, path) + } + + #[test] + fn test_page_v2_snappy_compression_fallback() { + // Test that PageV2 sets is_compressed to false when Snappy compression increases data size + let page_writer = TestPageWriter {}; + + // Create WriterProperties with PageV2 and Snappy compression + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_2_0) + // Disable dictionary to ensure data is written directly + .set_dictionary_enabled(false) + .set_compression(Compression::SNAPPY) + .build(); + + let mut column_writer = + get_test_column_writer::(Box::new(page_writer), 0, 0, Arc::new(props)); + + // Create small, simple data that Snappy compression will likely increase in size + // due to compression overhead for very small data + let values = vec![ByteArray::from("a")]; + + column_writer.write_batch(&values, None, None).unwrap(); + + let result = column_writer.close().unwrap(); + assert_eq!( + result.metadata.uncompressed_size(), + result.metadata.compressed_size() + ); + } + + struct ColumnRoundTripUniform<'a, T: DataType> { + props: WriterProperties, + values: &'a [T::T], + def_levels: LevelDataRef<'a>, + rep_levels: LevelDataRef<'a>, + max_def_level: i16, + max_rep_level: i16, + expected_values: &'a [T::T], + expected_def_levels: Option<&'a [i16]>, + expected_rep_levels: Option<&'a [i16]>, + } + + impl<'a, T: DataType> ColumnRoundTripUniform<'a, T> + where + T::T: PartialEq + std::fmt::Debug, + { + fn new() -> Self { + Self { + props: Default::default(), + values: &[], + def_levels: LevelDataRef::Absent, + rep_levels: LevelDataRef::Absent, + max_def_level: 0, + max_rep_level: 0, + expected_values: &[], + expected_def_levels: None, + expected_rep_levels: None, + } + } + + fn with_props(mut self, props: WriterProperties) -> Self { + self.props = props; + self + } + + fn with_values(mut self, values: &'a [T::T]) -> Self { + self.values = values; + self + } + + fn with_def_levels(mut self, def_levels: LevelDataRef<'a>) -> Self { + self.def_levels = def_levels; + self + } + + fn with_rep_levels(mut self, rep_levels: LevelDataRef<'a>) -> Self { + self.rep_levels = rep_levels; + self + } + + fn with_max_def_level(mut self, max_def_level: i16) -> Self { + self.max_def_level = max_def_level; + self + } + + fn with_max_rep_level(mut self, max_rep_level: i16) -> Self { + self.max_rep_level = max_rep_level; + self + } + + fn with_expected_values(mut self, expected_values: &'a [T::T]) -> Self { + self.expected_values = expected_values; + self + } + + fn with_expected_def_levels(mut self, expected_def_levels: &'a [i16]) -> Self { + self.expected_def_levels = Some(expected_def_levels); + self + } + + fn with_expected_rep_levels(mut self, expected_rep_levels: &'a [i16]) -> Self { + self.expected_rep_levels = Some(expected_rep_levels); + self + } + + /// Write-then-read roundtrip using `write_batch_internal` with the given + /// [`LevelDataRef`] variants, and assert the read-back matches `expected_*`. + fn run(self) { + let mut file = tempfile::tempfile().unwrap(); + let mut write = TrackedWrite::new(&mut file); + let page_writer = Box::new(SerializedPageWriter::new(&mut write)); + let mut writer = get_test_column_writer::( + page_writer, + self.max_def_level, + self.max_rep_level, + Arc::new(self.props), + ); + + writer + .write_batch_internal( + self.values, + None, + self.def_levels, + self.rep_levels, + None, + None, + None, + ) + .unwrap(); + let result = writer.close().unwrap(); + drop(write); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let page_reader = Box::new( + SerializedPageReader::new_with_properties( + Arc::new(file), + &result.metadata, + result.rows_written as usize, + None, + Arc::new(props), + ) + .unwrap(), + ); + let mut reader = + get_test_column_reader::(page_reader, self.max_def_level, self.max_rep_level); + + let batch_size = self + .expected_def_levels + .map_or(self.expected_values.len(), |l| l.len()); + let mut actual_values = Vec::with_capacity(batch_size); + let mut actual_def = self + .expected_def_levels + .map(|_| Vec::with_capacity(batch_size)); + let mut actual_rep = self + .expected_rep_levels + .map(|_| Vec::with_capacity(batch_size)); + + let (_, values_read, levels_read) = reader + .read_records( + batch_size, + actual_def.as_mut(), + actual_rep.as_mut(), + &mut actual_values, + ) + .unwrap(); + + assert_eq!(&actual_values[..values_read], self.expected_values); + if let Some(ref v) = actual_def { + assert_eq!(&v[..levels_read], self.expected_def_levels.unwrap()); + } + if let Some(ref v) = actual_rep { + assert_eq!(&v[..levels_read], self.expected_rep_levels.unwrap()); + } + } + } + + #[test] + fn test_level_data_ref_value_count() { + // `value_count` is what the byte-budget chunker uses to convert a + // chunk's level span into a leaf-value count. It must work for any + // column shape — flat, nullable, or nested — because the leaf + // values array is decoupled from the rep/def level stream. + let max_def = 2; + // Non-nullable / unrepeated: no def levels materialized — every + // level is a value. + assert_eq!(LevelDataRef::Absent.value_count(64, max_def), 64); + // Uniform run of present values, and of nulls. + assert_eq!( + LevelDataRef::Uniform { + value: max_def, + count: 40 + } + .value_count(40, max_def), + 40 + ); + assert_eq!( + LevelDataRef::Uniform { + value: max_def - 1, + count: 40 + } + .value_count(40, max_def), + 0 + ); + // Materialized def levels (nullable / nested): only levels equal to + // `max_def` are values; empty-list / null levels are not. + let levels = [2i16, 0, 2, 1, 2, 2, 0]; + assert_eq!( + LevelDataRef::Materialized(&levels).value_count(levels.len(), max_def), + 4 + ); + } + + #[test] + fn test_uniform_def_levels_all_null() { + // All-null column: def_level=0 (null) for every slot, no values written. + let max_def_level = 1; + let count = 100; + let expected_def_levels = vec![0i16; count]; + ColumnRoundTripUniform::::new() + .with_def_levels(LevelDataRef::Uniform { value: 0, count }) + .with_max_def_level(max_def_level) + .with_expected_def_levels(&expected_def_levels) + .run(); + } + + #[test] + fn test_uniform_def_levels_all_valid() { + // All-valid column: def_level=max for every slot, all values written. + let max_def_level = 1; + let values: Vec = (0..50).collect(); + let expected_def_levels = vec![max_def_level; values.len()]; + ColumnRoundTripUniform::::new() + .with_values(&values) + .with_def_levels(LevelDataRef::Uniform { + value: max_def_level, + count: values.len(), + }) + .with_max_def_level(max_def_level) + .with_expected_values(&values) + .with_expected_def_levels(&expected_def_levels) + .run(); + } + + #[test] + fn test_uniform_def_and_rep_levels() { + // Simulates a list column where every row is null: + // def=0, rep=0 for each row (one row = one entry with no child values). + let max_def_level = 2; + let max_rep_level = 1; + let count = 200; + let expected_def_levels = vec![0i16; count]; + let expected_rep_levels = vec![0i16; count]; + ColumnRoundTripUniform::::new() + .with_def_levels(LevelDataRef::Uniform { value: 0, count }) + .with_rep_levels(LevelDataRef::Uniform { value: 0, count }) + .with_max_def_level(max_def_level) + .with_max_rep_level(max_rep_level) + .with_expected_def_levels(&expected_def_levels) + .with_expected_rep_levels(&expected_rep_levels) + .run(); + } + + #[test] + fn test_uniform_levels_v1_and_v2() { + // Verify uniform levels work identically for both Parquet writer versions. + for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] { + let props = WriterProperties::builder() + .set_writer_version(version) + .build(); + let max_def = 1; + let count = 100; + let expected_def_levels = vec![0i16; count]; + ColumnRoundTripUniform::::new() + .with_props(props) + .with_def_levels(LevelDataRef::Uniform { value: 0, count }) + .with_max_def_level(max_def) + .with_expected_def_levels(&expected_def_levels) + .run(); + } + } +} diff --git a/vendor/parquet/src/compression.rs b/vendor/parquet/src/compression.rs new file mode 100644 index 0000000..6f3a18f --- /dev/null +++ b/vendor/parquet/src/compression.rs @@ -0,0 +1,947 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Contains codec interface and supported codec implementations. +//! +//! See [`Compression`](crate::basic::Compression) enum for all available compression +//! algorithms. +//! +#[cfg_attr( + feature = "experimental", + doc = r##" +# Example + +```no_run +use parquet::{basic::Compression, compression::{create_codec, CodecOptionsBuilder}}; + +let codec_options = CodecOptionsBuilder::default() + .set_backward_compatible_lz4(false) + .build(); +let mut codec = match create_codec(Compression::SNAPPY, &codec_options) { + Ok(Some(codec)) => codec, + _ => panic!(), +}; + +let data = vec![b'p', b'a', b'r', b'q', b'u', b'e', b't']; +let mut compressed = vec![]; +codec.compress(&data[..], &mut compressed).unwrap(); + +let mut output = vec![]; +codec.decompress(&compressed[..], &mut output, None).unwrap(); + +assert_eq!(output, data); +``` +"## +)] +use crate::basic::Compression as CodecType; +use crate::errors::{ParquetError, Result}; + +/// Parquet compression codec interface. +pub trait Codec: Send { + /// Compresses data stored in slice `input_buf` and appends the compressed result + /// to `output_buf`. + /// + /// Note that you'll need to call `clear()` before reusing the same `output_buf` + /// across different `compress` calls. + fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()>; + + /// Decompresses data stored in slice `input_buf` and appends output to `output_buf`. + /// + /// If the uncompress_size is provided it will allocate the exact amount of memory. + /// Otherwise, it will estimate the uncompressed size, allocating an amount of memory + /// greater or equal to the real uncompress_size. + /// + /// Returns the total number of bytes written. + fn decompress( + &mut self, + input_buf: &[u8], + output_buf: &mut Vec, + uncompress_size: Option, + ) -> Result; +} + +/// Struct to hold `Codec` creation options. +#[derive(Debug, PartialEq, Eq)] +pub struct CodecOptions { + /// Whether or not to fallback to other LZ4 older implementations on error in LZ4_HADOOP. + backward_compatible_lz4: bool, +} + +impl Default for CodecOptions { + fn default() -> Self { + CodecOptionsBuilder::default().build() + } +} + +pub struct CodecOptionsBuilder { + /// Whether or not to fallback to other LZ4 older implementations on error in LZ4_HADOOP. + backward_compatible_lz4: bool, +} + +impl Default for CodecOptionsBuilder { + fn default() -> Self { + Self { + backward_compatible_lz4: true, + } + } +} + +impl CodecOptionsBuilder { + /// Enable/disable backward compatible LZ4. + /// + /// If backward compatible LZ4 is enable, on LZ4_HADOOP error it will fallback + /// to the older versions LZ4 algorithms. That is LZ4_FRAME, for backward compatibility + /// with files generated by older versions of this library, and LZ4_RAW, for backward + /// compatibility with files generated by older versions of parquet-cpp. + /// + /// If backward compatible LZ4 is disabled, on LZ4_HADOOP error it will return the error. + pub fn set_backward_compatible_lz4(mut self, value: bool) -> CodecOptionsBuilder { + self.backward_compatible_lz4 = value; + self + } + + pub fn build(self) -> CodecOptions { + CodecOptions { + backward_compatible_lz4: self.backward_compatible_lz4, + } + } +} + +/// Defines valid compression levels. +pub(crate) trait CompressionLevel { + const MINIMUM_LEVEL: T; + const MAXIMUM_LEVEL: T; + + /// Tests if the provided compression level is valid. + fn is_valid_level(level: T) -> Result<()> { + let compression_range = Self::MINIMUM_LEVEL..=Self::MAXIMUM_LEVEL; + if compression_range.contains(&level) { + Ok(()) + } else { + Err(ParquetError::General(format!( + "valid compression range {}..={} exceeded.", + compression_range.start(), + compression_range.end() + ))) + } + } +} + +/// Given the compression type `codec`, returns a codec used to compress and decompress +/// bytes for the compression type. +/// This returns `None` if the codec type is `UNCOMPRESSED`. +pub fn create_codec(codec: CodecType, _options: &CodecOptions) -> Result>> { + #[allow(unreachable_code, unused_variables)] + match codec { + CodecType::BROTLI(level) => { + #[cfg(any(feature = "brotli", test))] + return Ok(Some(Box::new(BrotliCodec::new(level)))); + Err(ParquetError::General( + "Disabled feature at compile time: brotli".into(), + )) + } + CodecType::GZIP(level) => { + #[cfg(any(feature = "flate2", test))] + return Ok(Some(Box::new(GZipCodec::new(level)))); + Err(ParquetError::General( + "Disabled feature at compile time: flate2".into(), + )) + } + CodecType::SNAPPY => { + #[cfg(any(feature = "snap", test))] + return Ok(Some(Box::new(SnappyCodec::new()))); + Err(ParquetError::General( + "Disabled feature at compile time: snap".into(), + )) + } + CodecType::LZ4 => { + #[cfg(any(feature = "lz4", test))] + return Ok(Some(Box::new(LZ4HadoopCodec::new( + _options.backward_compatible_lz4, + )))); + Err(ParquetError::General( + "Disabled feature at compile time: lz4".into(), + )) + } + CodecType::ZSTD(level) => { + #[cfg(any(feature = "zstd", test))] + return Ok(Some(Box::new(ZSTDCodec::new(level)))); + Err(ParquetError::General( + "Disabled feature at compile time: zstd".into(), + )) + } + CodecType::LZ4_RAW => { + #[cfg(any(feature = "lz4", test))] + return Ok(Some(Box::new(LZ4RawCodec::new()))); + Err(ParquetError::General( + "Disabled feature at compile time: lz4".into(), + )) + } + CodecType::UNCOMPRESSED => Ok(None), + _ => Err(nyi_err!("The codec type {} is not supported yet", codec)), + } +} + +#[cfg(any(feature = "snap", test))] +mod snappy_codec { + use snap::raw::{Decoder, Encoder, decompress_len, max_compress_len}; + + use crate::compression::Codec; + use crate::errors::Result; + + /// Codec for Snappy compression format. + pub struct SnappyCodec { + decoder: Decoder, + encoder: Encoder, + } + + impl SnappyCodec { + /// Creates new Snappy compression codec. + pub(crate) fn new() -> Self { + Self { + decoder: Decoder::new(), + encoder: Encoder::new(), + } + } + } + + impl Codec for SnappyCodec { + fn decompress( + &mut self, + input_buf: &[u8], + output_buf: &mut Vec, + uncompress_size: Option, + ) -> Result { + let len = match uncompress_size { + Some(size) => size, + None => decompress_len(input_buf)?, + }; + let offset = output_buf.len(); + output_buf.resize(offset + len, 0); + self.decoder + .decompress(input_buf, &mut output_buf[offset..]) + .map_err(|e| e.into()) + } + + fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()> { + let output_buf_len = output_buf.len(); + let required_len = max_compress_len(input_buf.len()); + output_buf.resize(output_buf_len + required_len, 0); + let n = self + .encoder + .compress(input_buf, &mut output_buf[output_buf_len..])?; + output_buf.truncate(output_buf_len + n); + Ok(()) + } + } +} +#[cfg(any(feature = "snap", test))] +pub use snappy_codec::*; + +#[cfg(any(feature = "flate2", test))] +mod gzip_codec { + + use std::io::{Read, Write}; + + use flate2::{Compression, read, write}; + + use crate::compression::Codec; + use crate::errors::Result; + + use super::GzipLevel; + + /// Codec for GZIP compression algorithm. + pub struct GZipCodec { + level: GzipLevel, + } + + impl GZipCodec { + /// Creates new GZIP compression codec. + pub(crate) fn new(level: GzipLevel) -> Self { + Self { level } + } + } + + impl Codec for GZipCodec { + fn decompress( + &mut self, + input_buf: &[u8], + output_buf: &mut Vec, + _uncompress_size: Option, + ) -> Result { + let mut decoder = read::MultiGzDecoder::new(input_buf); + decoder.read_to_end(output_buf).map_err(|e| e.into()) + } + + fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()> { + let mut encoder = write::GzEncoder::new(output_buf, Compression::new(self.level.0)); + encoder.write_all(input_buf)?; + encoder.try_finish().map_err(|e| e.into()) + } + } +} +#[cfg(any(feature = "flate2", test))] +pub use gzip_codec::*; + +/// Represents a valid gzip compression level. +/// +/// Defaults to 6. +/// +/// * 0: least compression +/// * 9: most compression (that other software can read) +/// * 10: most compression (incompatible with other software, see below) +/// #### WARNING: +/// Level 10 compression can offer smallest file size, +/// but Parquet files created with it will not be readable +/// by other "standard" paquet readers. +/// +/// Do **NOT** use level 10 if you need other software to +/// be able to read the files. Read below for details. +/// +/// ### IMPORTANT: +/// There's often confusion about the compression levels in `flate2` vs `arrow` +/// as highlighted in issue [#1011](https://github.com/apache/arrow-rs/issues/6282). +/// +/// `flate2` supports two compression backends: `miniz_oxide` and `zlib`. +/// +/// - `zlib` supports levels from 0 to 9. +/// - `miniz_oxide` supports levels from 0 to 10. +/// +/// `arrow` uses `flate` with `rust_backend` feature, +/// which provides `miniz_oxide` as the backend. +/// Therefore 0-10 levels are supported. +/// +/// `flate2` documents this behavior properly with +/// [this commit](https://github.com/rust-lang/flate2-rs/pull/430). +#[derive(Debug, Eq, PartialEq, Hash, Clone, Copy)] +pub struct GzipLevel(u32); + +impl Default for GzipLevel { + fn default() -> Self { + // The default as of miniz_oxide 0.5.1 is 6 for compression level + // (miniz_oxide::deflate::CompressionLevel::DefaultLevel) + Self(6) + } +} + +impl CompressionLevel for GzipLevel { + const MINIMUM_LEVEL: u32 = 0; + const MAXIMUM_LEVEL: u32 = 9; +} + +impl GzipLevel { + /// Attempts to create a gzip compression level. + /// + /// Compression levels must be valid (i.e. be acceptable for [`flate2::Compression`]). + pub fn try_new(level: u32) -> Result { + Self::is_valid_level(level).map(|_| Self(level)) + } + + /// Returns the compression level. + pub fn compression_level(&self) -> u32 { + self.0 + } +} + +#[cfg(any(feature = "brotli", test))] +mod brotli_codec { + + use std::io::{Read, Write}; + + use crate::compression::Codec; + use crate::errors::Result; + + use super::BrotliLevel; + + const BROTLI_DEFAULT_BUFFER_SIZE: usize = 4096; + const BROTLI_DEFAULT_LG_WINDOW_SIZE: u32 = 22; // recommended between 20-22 + + /// Codec for Brotli compression algorithm. + pub struct BrotliCodec { + level: BrotliLevel, + } + + impl BrotliCodec { + /// Creates new Brotli compression codec. + pub(crate) fn new(level: BrotliLevel) -> Self { + Self { level } + } + } + + impl Codec for BrotliCodec { + fn decompress( + &mut self, + input_buf: &[u8], + output_buf: &mut Vec, + uncompress_size: Option, + ) -> Result { + let buffer_size = uncompress_size.unwrap_or(BROTLI_DEFAULT_BUFFER_SIZE); + brotli::Decompressor::new(input_buf, buffer_size) + .read_to_end(output_buf) + .map_err(|e| e.into()) + } + + fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()> { + let mut encoder = brotli::CompressorWriter::new( + output_buf, + BROTLI_DEFAULT_BUFFER_SIZE, + self.level.0, + BROTLI_DEFAULT_LG_WINDOW_SIZE, + ); + encoder.write_all(input_buf)?; + encoder.flush().map_err(|e| e.into()) + } + } +} +#[cfg(any(feature = "brotli", test))] +pub use brotli_codec::*; + +/// Represents a valid brotli compression level. +#[derive(Debug, Eq, PartialEq, Hash, Clone, Copy)] +pub struct BrotliLevel(u32); + +impl Default for BrotliLevel { + fn default() -> Self { + Self(1) + } +} + +impl CompressionLevel for BrotliLevel { + const MINIMUM_LEVEL: u32 = 0; + const MAXIMUM_LEVEL: u32 = 11; +} + +impl BrotliLevel { + /// Attempts to create a brotli compression level. + /// + /// Compression levels must be valid. + pub fn try_new(level: u32) -> Result { + Self::is_valid_level(level).map(|_| Self(level)) + } + + /// Returns the compression level. + pub fn compression_level(&self) -> u32 { + self.0 + } +} + +#[cfg(any(feature = "lz4", test))] +mod lz4_codec { + use std::io::{Read, Write}; + + use crate::compression::Codec; + use crate::errors::{ParquetError, Result}; + + const LZ4_BUFFER_SIZE: usize = 4096; + + /// Codec for LZ4 compression algorithm. + pub struct LZ4Codec {} + + impl LZ4Codec { + /// Creates new LZ4 compression codec. + pub(crate) fn new() -> Self { + Self {} + } + } + + impl Codec for LZ4Codec { + fn decompress( + &mut self, + input_buf: &[u8], + output_buf: &mut Vec, + _uncompress_size: Option, + ) -> Result { + let mut decoder = lz4_flex::frame::FrameDecoder::new(input_buf); + let mut buffer: [u8; LZ4_BUFFER_SIZE] = [0; LZ4_BUFFER_SIZE]; + let mut total_len = 0; + loop { + let len = decoder.read(&mut buffer)?; + if len == 0 { + break; + } + total_len += len; + output_buf.write_all(&buffer[0..len])?; + } + Ok(total_len) + } + + fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()> { + let mut encoder = lz4_flex::frame::FrameEncoder::new(output_buf); + let mut from = 0; + loop { + let to = std::cmp::min(from + LZ4_BUFFER_SIZE, input_buf.len()); + encoder.write_all(&input_buf[from..to])?; + from += LZ4_BUFFER_SIZE; + if from >= input_buf.len() { + break; + } + } + match encoder.finish() { + Ok(_) => Ok(()), + Err(e) => Err(ParquetError::External(Box::new(e))), + } + } + } +} + +#[cfg(all(feature = "experimental", any(feature = "lz4", test)))] +pub use lz4_codec::*; + +#[cfg(any(feature = "zstd", test))] +mod zstd_codec { + use crate::compression::{Codec, ZstdLevel}; + use crate::errors::Result; + use std::io::Cursor; + + /// Codec for Zstandard compression algorithm. + /// + /// Uses `zstd::bulk` API with reusable compressor/decompressor contexts + /// to avoid the overhead of reinitializing contexts for each operation. + pub struct ZSTDCodec { + compressor: zstd::bulk::Compressor<'static>, + decompressor: zstd::bulk::Decompressor<'static>, + } + + impl ZSTDCodec { + /// Creates new Zstandard compression codec. + pub(crate) fn new(level: ZstdLevel) -> Self { + Self { + compressor: zstd::bulk::Compressor::new(level.compression_level()) + .expect("valid zstd compression level"), + decompressor: zstd::bulk::Decompressor::new() + .expect("can create zstd decompressor"), + } + } + } + + impl Codec for ZSTDCodec { + fn decompress( + &mut self, + input_buf: &[u8], + output_buf: &mut Vec, + uncompress_size: Option, + ) -> Result { + let offset = output_buf.len(); + let len = uncompress_size + .or_else(|| { + // Get the decompressed size from the zstd frame header + zstd::zstd_safe::get_frame_content_size(input_buf) + .ok() + .flatten() + .map(|size| size as usize) + }) + .unwrap_or(input_buf.len().saturating_mul(4)); + output_buf.reserve(len); + + let mut cursor = Cursor::new(output_buf); + cursor.set_position(offset as u64); + let len = self + .decompressor + .decompress_to_buffer(input_buf, &mut cursor)?; + Ok(len) + } + + fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()> { + let offset = output_buf.len(); + let len = zstd::zstd_safe::compress_bound(input_buf.len()); + output_buf.reserve(len); + + let mut cursor = Cursor::new(output_buf); + cursor.set_position(offset as u64); + let _written = self.compressor.compress_to_buffer(input_buf, &mut cursor)?; + Ok(()) + } + } +} +#[cfg(any(feature = "zstd", test))] +pub use zstd_codec::*; + +/// Represents a valid zstd compression level. +#[derive(Debug, Eq, PartialEq, Hash, Clone, Copy)] +pub struct ZstdLevel(i32); + +impl CompressionLevel for ZstdLevel { + // zstd binds to C, and hence zstd::compression_level_range() is not const as this calls the + // underlying C library. + const MINIMUM_LEVEL: i32 = 1; + const MAXIMUM_LEVEL: i32 = 22; +} + +impl ZstdLevel { + /// Attempts to create a zstd compression level from a given compression level. + /// + /// Compression levels must be valid (i.e. be acceptable for [`zstd::compression_level_range`]). + pub fn try_new(level: i32) -> Result { + Self::is_valid_level(level).map(|_| Self(level)) + } + + /// Returns the compression level. + pub fn compression_level(&self) -> i32 { + self.0 + } +} + +impl Default for ZstdLevel { + fn default() -> Self { + Self(1) + } +} + +#[cfg(any(feature = "lz4", test))] +mod lz4_raw_codec { + use crate::compression::Codec; + use crate::errors::ParquetError; + use crate::errors::Result; + + /// Codec for LZ4 Raw compression algorithm. + pub struct LZ4RawCodec {} + + impl LZ4RawCodec { + /// Creates new LZ4 Raw compression codec. + pub(crate) fn new() -> Self { + Self {} + } + } + + impl Codec for LZ4RawCodec { + fn decompress( + &mut self, + input_buf: &[u8], + output_buf: &mut Vec, + uncompress_size: Option, + ) -> Result { + let offset = output_buf.len(); + let required_len = match uncompress_size { + Some(uncompress_size) => uncompress_size, + None => { + return Err(ParquetError::General( + "LZ4RawCodec unsupported without uncompress_size".into(), + )); + } + }; + output_buf.resize(offset + required_len, 0); + match lz4_flex::block::decompress_into(input_buf, &mut output_buf[offset..]) { + Ok(n) => { + if n != required_len { + return Err(ParquetError::General( + "LZ4RawCodec uncompress_size is not the expected one".into(), + )); + } + Ok(n) + } + Err(e) => Err(ParquetError::External(Box::new(e))), + } + } + + fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()> { + let offset = output_buf.len(); + let required_len = lz4_flex::block::get_maximum_output_size(input_buf.len()); + output_buf.resize(offset + required_len, 0); + match lz4_flex::block::compress_into(input_buf, &mut output_buf[offset..]) { + Ok(n) => { + output_buf.truncate(offset + n); + Ok(()) + } + Err(e) => Err(ParquetError::External(Box::new(e))), + } + } + } +} +#[cfg(any(feature = "lz4", test))] +pub use lz4_raw_codec::*; + +#[cfg(any(feature = "lz4", test))] +mod lz4_hadoop_codec { + use crate::compression::Codec; + use crate::compression::lz4_codec::LZ4Codec; + use crate::compression::lz4_raw_codec::LZ4RawCodec; + use crate::errors::{ParquetError, Result}; + use std::io; + + /// Size of u32 type. + const SIZE_U32: usize = std::mem::size_of::(); + + /// Length of the LZ4_HADOOP prefix. + const PREFIX_LEN: usize = SIZE_U32 * 2; + + /// Codec for LZ4 Hadoop compression algorithm. + pub struct LZ4HadoopCodec { + /// Whether or not to fallback to other LZ4 implementations on error. + /// Fallback is done to be backward compatible with older versions of this + /// library and older versions parquet-cpp. + backward_compatible_lz4: bool, + } + + impl LZ4HadoopCodec { + /// Creates new LZ4 Hadoop compression codec. + pub(crate) fn new(backward_compatible_lz4: bool) -> Self { + Self { + backward_compatible_lz4, + } + } + } + + /// Try to decompress the buffer as if it was compressed with the Hadoop Lz4Codec. + /// Adapted from pola-rs [compression.rs:try_decompress_hadoop](https://pola-rs.github.io/polars/src/parquet2/compression.rs.html#225) + /// Translated from the apache arrow c++ function [TryDecompressHadoop](https://github.com/apache/arrow/blob/bf18e6e4b5bb6180706b1ba0d597a65a4ce5ca48/cpp/src/arrow/util/compression_lz4.cc#L474). + /// Returns error if decompression failed. + fn try_decompress_hadoop(input_buf: &[u8], output_buf: &mut [u8]) -> io::Result { + // Parquet files written with the Hadoop Lz4Codec use their own framing. + // The input buffer can contain an arbitrary number of "frames", each + // with the following structure: + // - bytes 0..3: big-endian uint32_t representing the frame decompressed size + // - bytes 4..7: big-endian uint32_t representing the frame compressed size + // - bytes 8...: frame compressed data + // + // The Hadoop Lz4Codec source code can be found here: + // https://github.com/apache/hadoop/blob/trunk/hadoop-mapreduce-project/hadoop-mapreduce-client/hadoop-mapreduce-client-nativetask/src/main/native/src/codec/Lz4Codec.cc + let mut input_len = input_buf.len(); + let mut input = input_buf; + let mut read_bytes = 0; + let mut output_len = output_buf.len(); + let mut output: &mut [u8] = output_buf; + while input_len >= PREFIX_LEN { + let mut bytes = [0; SIZE_U32]; + bytes.copy_from_slice(&input[0..4]); + let expected_decompressed_size = u32::from_be_bytes(bytes); + let mut bytes = [0; SIZE_U32]; + bytes.copy_from_slice(&input[4..8]); + let expected_compressed_size = u32::from_be_bytes(bytes); + input = &input[PREFIX_LEN..]; + input_len -= PREFIX_LEN; + + if input_len < expected_compressed_size as usize { + return Err(io::Error::other("Not enough bytes for Hadoop frame")); + } + + if output_len < expected_decompressed_size as usize { + return Err(io::Error::other( + "Not enough bytes to hold advertised output", + )); + } + let decompressed_size = + lz4_flex::decompress_into(&input[..expected_compressed_size as usize], output) + .map_err(|e| ParquetError::External(Box::new(e)))?; + if decompressed_size != expected_decompressed_size as usize { + return Err(io::Error::other("Unexpected decompressed size")); + } + input_len -= expected_compressed_size as usize; + output_len -= expected_decompressed_size as usize; + read_bytes += expected_decompressed_size as usize; + if input_len > expected_compressed_size as usize { + input = &input[expected_compressed_size as usize..]; + output = &mut output[expected_decompressed_size as usize..]; + } else { + break; + } + } + if input_len == 0 { + Ok(read_bytes) + } else { + Err(io::Error::other("Not all input are consumed")) + } + } + + impl Codec for LZ4HadoopCodec { + fn decompress( + &mut self, + input_buf: &[u8], + output_buf: &mut Vec, + uncompress_size: Option, + ) -> Result { + let output_len = output_buf.len(); + let required_len = match uncompress_size { + Some(n) => n, + None => { + return Err(ParquetError::General( + "LZ4HadoopCodec unsupported without uncompress_size".into(), + )); + } + }; + output_buf.resize(output_len + required_len, 0); + match try_decompress_hadoop(input_buf, &mut output_buf[output_len..]) { + Ok(n) => { + if n != required_len { + return Err(ParquetError::General( + "LZ4HadoopCodec uncompress_size is not the expected one".into(), + )); + } + Ok(n) + } + Err(e) if !self.backward_compatible_lz4 => Err(e.into()), + // Fallback done to be backward compatible with older versions of this + // library and older versions of parquet-cpp. + Err(_) => { + // Truncate any inserted element before tryingg next algorithm. + output_buf.truncate(output_len); + match LZ4Codec::new().decompress(input_buf, output_buf, uncompress_size) { + Ok(n) => Ok(n), + Err(_) => { + // Truncate any inserted element before tryingg next algorithm. + output_buf.truncate(output_len); + LZ4RawCodec::new().decompress(input_buf, output_buf, uncompress_size) + } + } + } + } + } + + fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()> { + // Allocate memory to store the LZ4_HADOOP prefix. + let offset = output_buf.len(); + output_buf.resize(offset + PREFIX_LEN, 0); + + // Append LZ4_RAW compressed bytes after prefix. + LZ4RawCodec::new().compress(input_buf, output_buf)?; + + // Prepend decompressed size and compressed size in big endian to be compatible + // with LZ4_HADOOP. + let output_buf = &mut output_buf[offset..]; + let compressed_size = output_buf.len() - PREFIX_LEN; + let compressed_size = compressed_size as u32; + let uncompressed_size = input_buf.len() as u32; + output_buf[..SIZE_U32].copy_from_slice(&uncompressed_size.to_be_bytes()); + output_buf[SIZE_U32..PREFIX_LEN].copy_from_slice(&compressed_size.to_be_bytes()); + + Ok(()) + } + } +} +#[cfg(any(feature = "lz4", test))] +pub use lz4_hadoop_codec::*; + +#[cfg(test)] +mod tests { + use super::*; + + use crate::util::test_common::rand_gen::random_bytes; + + fn test_roundtrip(c: CodecType, data: &[u8], uncompress_size: Option) { + let codec_options = CodecOptionsBuilder::default() + .set_backward_compatible_lz4(false) + .build(); + let mut c1 = create_codec(c, &codec_options).unwrap().unwrap(); + let mut c2 = create_codec(c, &codec_options).unwrap().unwrap(); + + // Compress with c1 + let mut compressed = Vec::new(); + let mut decompressed = Vec::new(); + c1.compress(data, &mut compressed) + .expect("Error when compressing"); + + // Decompress with c2 + let decompressed_size = c2 + .decompress(compressed.as_slice(), &mut decompressed, uncompress_size) + .expect("Error when decompressing"); + assert_eq!(data.len(), decompressed_size); + assert_eq!(data, decompressed.as_slice()); + + decompressed.clear(); + compressed.clear(); + + // Compress with c2 + c2.compress(data, &mut compressed) + .expect("Error when compressing"); + + // Decompress with c1 + let decompressed_size = c1 + .decompress(compressed.as_slice(), &mut decompressed, uncompress_size) + .expect("Error when decompressing"); + assert_eq!(data.len(), decompressed_size); + assert_eq!(data, decompressed.as_slice()); + + decompressed.clear(); + compressed.clear(); + + // Test does not trample existing data in output buffers + let prefix = &[0xDE, 0xAD, 0xBE, 0xEF]; + decompressed.extend_from_slice(prefix); + compressed.extend_from_slice(prefix); + + c2.compress(data, &mut compressed) + .expect("Error when compressing"); + + assert_eq!(&compressed[..4], prefix); + + let decompressed_size = c2 + .decompress(&compressed[4..], &mut decompressed, uncompress_size) + .expect("Error when decompressing"); + + assert_eq!(data.len(), decompressed_size); + assert_eq!(data, &decompressed[4..]); + assert_eq!(&decompressed[..4], prefix); + } + + fn test_codec_with_size(c: CodecType) { + let sizes = vec![100, 10000, 100000]; + for size in sizes { + let data = random_bytes(size); + test_roundtrip(c, &data, Some(data.len())); + } + } + + fn test_codec_without_size(c: CodecType) { + let sizes = vec![100, 10000, 100000]; + for size in sizes { + let data = random_bytes(size); + test_roundtrip(c, &data, None); + } + } + + #[test] + fn test_codec_snappy() { + test_codec_with_size(CodecType::SNAPPY); + test_codec_without_size(CodecType::SNAPPY); + } + + #[test] + fn test_codec_gzip() { + for level in GzipLevel::MINIMUM_LEVEL..=GzipLevel::MAXIMUM_LEVEL { + let level = GzipLevel::try_new(level).unwrap(); + test_codec_with_size(CodecType::GZIP(level)); + test_codec_without_size(CodecType::GZIP(level)); + } + } + + #[test] + fn test_codec_brotli() { + for level in BrotliLevel::MINIMUM_LEVEL..=BrotliLevel::MAXIMUM_LEVEL { + let level = BrotliLevel::try_new(level).unwrap(); + test_codec_with_size(CodecType::BROTLI(level)); + test_codec_without_size(CodecType::BROTLI(level)); + } + } + + #[test] + fn test_codec_lz4() { + test_codec_with_size(CodecType::LZ4); + } + + #[test] + fn test_codec_zstd() { + for level in ZstdLevel::MINIMUM_LEVEL..=ZstdLevel::MAXIMUM_LEVEL { + let level = ZstdLevel::try_new(level).unwrap(); + test_codec_with_size(CodecType::ZSTD(level)); + test_codec_without_size(CodecType::ZSTD(level)); + } + } + + #[test] + fn test_codec_lz4_raw() { + test_codec_with_size(CodecType::LZ4_RAW); + } +} diff --git a/vendor/parquet/src/data_type.rs b/vendor/parquet/src/data_type.rs new file mode 100644 index 0000000..d8c7b92 --- /dev/null +++ b/vendor/parquet/src/data_type.rs @@ -0,0 +1,1416 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Data types that connect Parquet physical types with their Rust-specific +//! representations. +use bytes::Bytes; +use half::f16; +use std::cmp::Ordering; +use std::fmt; +use std::mem; +use std::ops::{Deref, DerefMut}; +use std::str::from_utf8; + +use crate::basic::Type; +use crate::column::reader::{ColumnReader, ColumnReaderImpl}; +use crate::column::writer::{ColumnWriter, ColumnWriterImpl}; +use crate::errors::{ParquetError, Result}; +use crate::util::bit_util::FromBytes; + +/// Rust representation for logical type INT96, value is backed by an array of `u32`. +/// The type only takes 12 bytes, without extra padding. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct Int96 { + value: [u32; 3], +} + +const JULIAN_DAY_OF_EPOCH: i64 = 2_440_588; + +/// Number of seconds in a day +const SECONDS_IN_DAY: i64 = 86_400; +/// Number of milliseconds in a second +const MILLISECONDS: i64 = 1_000; +/// Number of microseconds in a second +const MICROSECONDS: i64 = 1_000_000; +/// Number of nanoseconds in a second +const NANOSECONDS: i64 = 1_000_000_000; + +/// Number of milliseconds in a day +const MILLISECONDS_IN_DAY: i64 = SECONDS_IN_DAY * MILLISECONDS; +/// Number of microseconds in a day +const MICROSECONDS_IN_DAY: i64 = SECONDS_IN_DAY * MICROSECONDS; +/// Number of nanoseconds in a day +const NANOSECONDS_IN_DAY: i64 = SECONDS_IN_DAY * NANOSECONDS; + +impl Int96 { + /// Creates new INT96 type struct with no data set. + pub fn new() -> Self { + Self { value: [0; 3] } + } + + /// Returns underlying data as slice of [`u32`]. + #[inline] + pub fn data(&self) -> &[u32] { + &self.value + } + + /// Sets data for this INT96 type. + #[inline] + pub fn set_data(&mut self, elem0: u32, elem1: u32, elem2: u32) { + self.value = [elem0, elem1, elem2]; + } + + /// Converts this INT96 into an i64 representing the number of SECONDS since EPOCH + /// + /// Will wrap around on overflow + #[inline] + pub fn to_seconds(&self) -> i64 { + let (day, nanos) = self.data_as_days_and_nanos(); + (day as i64 - JULIAN_DAY_OF_EPOCH) + .wrapping_mul(SECONDS_IN_DAY) + .wrapping_add(nanos / 1_000_000_000) + } + + /// Converts this INT96 into an i64 representing the number of MILLISECONDS since EPOCH + /// + /// Will wrap around on overflow + #[inline] + pub fn to_millis(&self) -> i64 { + let (day, nanos) = self.data_as_days_and_nanos(); + (day as i64 - JULIAN_DAY_OF_EPOCH) + .wrapping_mul(MILLISECONDS_IN_DAY) + .wrapping_add(nanos / 1_000_000) + } + + /// Converts this INT96 into an i64 representing the number of MICROSECONDS since EPOCH + /// + /// Will wrap around on overflow + #[inline] + pub fn to_micros(&self) -> i64 { + let (day, nanos) = self.data_as_days_and_nanos(); + (day as i64 - JULIAN_DAY_OF_EPOCH) + .wrapping_mul(MICROSECONDS_IN_DAY) + .wrapping_add(nanos / 1_000) + } + + /// Converts this INT96 into an i64 representing the number of NANOSECONDS since EPOCH + /// + /// Will wrap around on overflow + #[inline] + pub fn to_nanos(&self) -> i64 { + let (day, nanos) = self.data_as_days_and_nanos(); + (day as i64 - JULIAN_DAY_OF_EPOCH) + .wrapping_mul(NANOSECONDS_IN_DAY) + .wrapping_add(nanos) + } + + #[inline] + fn get_days(&self) -> i32 { + self.data()[2] as i32 + } + + #[inline] + fn get_nanos(&self) -> i64 { + ((self.data()[1] as i64) << 32) + self.data()[0] as i64 + } + + #[inline] + fn data_as_days_and_nanos(&self) -> (i32, i64) { + (self.get_days(), self.get_nanos()) + } +} + +impl PartialOrd for Int96 { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for Int96 { + /// Order `Int96` correctly for (deprecated) timestamp types. + /// + /// Note: this is done even though the Int96 type is deprecated and the + /// [spec does not define the sort order] + /// because some engines, notably Spark and Databricks Photon still write + /// Int96 timestamps and rely on their order for optimization. + /// + /// [spec does not define the sort order]: https://github.com/apache/parquet-format/blob/cf943c197f4fad826b14ba0c40eb0ffdab585285/src/main/thrift/parquet.thrift#L1079 + fn cmp(&self, other: &Self) -> Ordering { + match self.get_days().cmp(&other.get_days()) { + Ordering::Equal => self.get_nanos().cmp(&other.get_nanos()), + ord => ord, + } + } +} +impl From> for Int96 { + fn from(buf: Vec) -> Self { + assert_eq!(buf.len(), 3); + let mut result = Self::new(); + result.set_data(buf[0], buf[1], buf[2]); + result + } +} + +impl fmt::Display for Int96 { + #[cold] + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{:?}", self.data()) + } +} + +/// Rust representation for BYTE_ARRAY and FIXED_LEN_BYTE_ARRAY Parquet physical types. +/// Value is backed by a byte buffer. +#[derive(Clone, Default)] +pub struct ByteArray { + data: Option, +} + +// Special case Debug that prints out byte arrays that are valid utf8 as &str's +impl std::fmt::Debug for ByteArray { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let mut debug_struct = f.debug_struct("ByteArray"); + match self.as_utf8() { + Ok(s) => debug_struct.field("data", &s), + Err(_) => debug_struct.field("data", &self.data), + }; + debug_struct.finish() + } +} + +impl PartialOrd for ByteArray { + fn partial_cmp(&self, other: &ByteArray) -> Option { + // sort nulls first (consistent with PartialCmp on Option) + // + // Since ByteBuffer doesn't implement PartialOrd, so can't + // derive an implementation + match (&self.data, &other.data) { + (None, None) => Some(Ordering::Equal), + (None, Some(_)) => Some(Ordering::Less), + (Some(_), None) => Some(Ordering::Greater), + (Some(self_data), Some(other_data)) => { + // compare slices directly + self_data.partial_cmp(&other_data) + } + } + } +} + +impl ByteArray { + /// Creates new byte array with no data set. + #[inline] + pub fn new() -> Self { + ByteArray { data: None } + } + + /// Gets length of the underlying byte buffer. + #[inline] + pub fn len(&self) -> usize { + assert!(self.data.is_some()); + self.data.as_ref().unwrap().len() + } + + /// Checks if the underlying buffer is empty. + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Returns slice of data. + #[inline] + pub fn data(&self) -> &[u8] { + self.data + .as_ref() + .expect("set_data should have been called") + .as_ref() + } + + /// Set data from another byte buffer. + #[inline] + pub fn set_data(&mut self, data: Bytes) { + self.data = Some(data); + } + + /// Returns `ByteArray` instance with slice of values for a data. + #[inline] + pub fn slice(&self, start: usize, len: usize) -> Self { + Self::from( + self.data + .as_ref() + .expect("set_data should have been called") + .slice(start..start + len), + ) + } + + /// Try to convert the byte array to a utf8 slice + pub fn as_utf8(&self) -> Result<&str> { + self.data + .as_ref() + .map(|ptr| ptr.as_ref()) + .ok_or_else(|| general_err!("Can't convert empty byte array to utf8")) + .and_then(|bytes| from_utf8(bytes).map_err(|e| e.into())) + } +} + +impl From> for ByteArray { + fn from(buf: Vec) -> ByteArray { + Self { + data: Some(buf.into()), + } + } +} + +impl<'a> From<&'a [u8]> for ByteArray { + fn from(b: &'a [u8]) -> ByteArray { + let mut v = Vec::new(); + v.extend_from_slice(b); + Self { + data: Some(v.into()), + } + } +} + +impl<'a> From<&'a str> for ByteArray { + fn from(s: &'a str) -> ByteArray { + let mut v = Vec::new(); + v.extend_from_slice(s.as_bytes()); + Self { + data: Some(v.into()), + } + } +} + +impl From for ByteArray { + fn from(value: Bytes) -> Self { + Self { data: Some(value) } + } +} + +impl From for ByteArray { + fn from(value: f16) -> Self { + Self::from(value.to_le_bytes().as_slice()) + } +} + +impl PartialEq for ByteArray { + fn eq(&self, other: &ByteArray) -> bool { + match (&self.data, &other.data) { + (Some(d1), Some(d2)) => d1.as_ref() == d2.as_ref(), + (None, None) => true, + _ => false, + } + } +} + +impl fmt::Display for ByteArray { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{:?}", self.data()) + } +} + +/// Wrapper type for performance reasons, this represents `FIXED_LEN_BYTE_ARRAY` but in all other +/// considerations behaves the same as `ByteArray` +/// +/// # Performance notes: +/// This type is a little unfortunate, without it the compiler generates code that takes quite a +/// big hit on the CPU pipeline. Essentially the previous version stalls awaiting the result of +/// `T::get_physical_type() == Type::FIXED_LEN_BYTE_ARRAY`. +/// +/// Its debatable if this is wanted, it is out of spec for what parquet documents as its base +/// types, although there are code paths in the Rust (and potentially the C++) versions that +/// warrant this. +/// +/// With this wrapper type the compiler generates more targeted code paths matching the higher +/// level logical types, removing the data-hazard from all decoding and encoding paths. +#[repr(transparent)] +#[derive(Clone, Debug, Default)] +pub struct FixedLenByteArray(ByteArray); + +impl PartialEq for FixedLenByteArray { + fn eq(&self, other: &FixedLenByteArray) -> bool { + self.0.eq(&other.0) + } +} + +impl PartialEq for FixedLenByteArray { + fn eq(&self, other: &ByteArray) -> bool { + self.0.eq(other) + } +} + +impl PartialEq for ByteArray { + fn eq(&self, other: &FixedLenByteArray) -> bool { + self.eq(&other.0) + } +} + +impl fmt::Display for FixedLenByteArray { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + self.0.fmt(f) + } +} + +impl PartialOrd for FixedLenByteArray { + fn partial_cmp(&self, other: &FixedLenByteArray) -> Option { + self.0.partial_cmp(&other.0) + } +} + +impl PartialOrd for ByteArray { + fn partial_cmp(&self, other: &FixedLenByteArray) -> Option { + self.partial_cmp(&other.0) + } +} + +impl PartialOrd for FixedLenByteArray { + fn partial_cmp(&self, other: &ByteArray) -> Option { + self.0.partial_cmp(other) + } +} + +impl Deref for FixedLenByteArray { + type Target = ByteArray; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl DerefMut for FixedLenByteArray { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} + +impl From for FixedLenByteArray { + fn from(other: ByteArray) -> Self { + Self(other) + } +} + +impl From> for FixedLenByteArray { + fn from(buf: Vec) -> FixedLenByteArray { + FixedLenByteArray(ByteArray::from(buf)) + } +} + +impl From for ByteArray { + fn from(other: FixedLenByteArray) -> Self { + other.0 + } +} + +/// Rust representation for Decimal values. +/// +/// This is not a representation of Parquet physical type, but rather a wrapper for +/// DECIMAL logical type, and serves as container for raw parts of decimal values: +/// unscaled value in bytes, precision and scale. +#[derive(Clone, Debug)] +pub enum Decimal { + /// Decimal backed by `i32`. + Int32 { + /// The underlying value + value: [u8; 4], + /// The total number of digits in the number + precision: i32, + /// The number of digits to the right of the decimal point + scale: i32, + }, + /// Decimal backed by `i64`. + Int64 { + /// The underlying value + value: [u8; 8], + /// The total number of digits in the number + precision: i32, + /// The number of digits to the right of the decimal point + scale: i32, + }, + /// Decimal backed by byte array. + Bytes { + /// The underlying value + value: ByteArray, + /// The total number of digits in the number + precision: i32, + /// The number of digits to the right of the decimal point + scale: i32, + }, +} + +impl Decimal { + /// Creates new decimal value from `i32`. + pub fn from_i32(value: i32, precision: i32, scale: i32) -> Self { + let bytes = value.to_be_bytes(); + Decimal::Int32 { + value: bytes, + precision, + scale, + } + } + + /// Creates new decimal value from `i64`. + pub fn from_i64(value: i64, precision: i32, scale: i32) -> Self { + let bytes = value.to_be_bytes(); + Decimal::Int64 { + value: bytes, + precision, + scale, + } + } + + /// Creates new decimal value from `ByteArray`. + pub fn from_bytes(value: ByteArray, precision: i32, scale: i32) -> Self { + Decimal::Bytes { + value, + precision, + scale, + } + } + + /// Returns bytes of unscaled value. + pub fn data(&self) -> &[u8] { + match *self { + Decimal::Int32 { ref value, .. } => value, + Decimal::Int64 { ref value, .. } => value, + Decimal::Bytes { ref value, .. } => value.data(), + } + } + + /// Returns decimal precision. + pub fn precision(&self) -> i32 { + match *self { + Decimal::Int32 { precision, .. } => precision, + Decimal::Int64 { precision, .. } => precision, + Decimal::Bytes { precision, .. } => precision, + } + } + + /// Returns decimal scale. + pub fn scale(&self) -> i32 { + match *self { + Decimal::Int32 { scale, .. } => scale, + Decimal::Int64 { scale, .. } => scale, + Decimal::Bytes { scale, .. } => scale, + } + } +} + +impl Default for Decimal { + fn default() -> Self { + Self::from_i32(0, 0, 0) + } +} + +impl PartialEq for Decimal { + fn eq(&self, other: &Decimal) -> bool { + self.precision() == other.precision() + && self.scale() == other.scale() + && self.data() == other.data() + } +} + +/// Converts an instance of data type to a slice of bytes as `u8`. +pub trait AsBytes { + /// Returns slice of bytes for this data type. + fn as_bytes(&self) -> &[u8]; +} + +/// Converts an slice of a data type to a slice of bytes. +pub trait SliceAsBytes: Sized { + /// Returns slice of bytes for a slice of this data type. + fn slice_as_bytes(self_: &[Self]) -> &[u8]; + /// Return the internal representation as a mutable slice + /// + /// # Safety + /// If modified you are _required_ to ensure the internal representation + /// is valid and correct for the actual raw data + unsafe fn slice_as_bytes_mut(self_: &mut [Self]) -> &mut [u8]; +} + +impl AsBytes for [u8] { + fn as_bytes(&self) -> &[u8] { + self + } +} + +macro_rules! gen_as_bytes { + ($source_ty:ident) => { + impl AsBytes for $source_ty { + #[allow(clippy::size_of_in_element_count)] + fn as_bytes(&self) -> &[u8] { + // SAFETY: macro is only used with primitive types that have no padding, so the + // resulting slice always refers to initialized memory. + unsafe { + std::slice::from_raw_parts( + self as *const $source_ty as *const u8, + std::mem::size_of::<$source_ty>(), + ) + } + } + } + + impl SliceAsBytes for $source_ty { + #[inline] + #[allow(clippy::size_of_in_element_count)] + fn slice_as_bytes(self_: &[Self]) -> &[u8] { + // SAFETY: macro is only used with primitive types that have no padding, so the + // resulting slice always refers to initialized memory. + unsafe { + std::slice::from_raw_parts( + self_.as_ptr() as *const u8, + std::mem::size_of_val(self_), + ) + } + } + + #[inline] + #[allow(clippy::size_of_in_element_count)] + unsafe fn slice_as_bytes_mut(self_: &mut [Self]) -> &mut [u8] { + // SAFETY: macro is only used with primitive types that have no padding, so the + // resulting slice always refers to initialized memory. Moreover, self has no + // invalid bit patterns, so all writes to the resulting slice will be valid. + unsafe { + std::slice::from_raw_parts_mut( + self_.as_mut_ptr() as *mut u8, + std::mem::size_of_val(self_), + ) + } + } + } + }; +} + +gen_as_bytes!(i8); +gen_as_bytes!(i16); +gen_as_bytes!(i32); +gen_as_bytes!(i64); +gen_as_bytes!(u8); +gen_as_bytes!(u16); +gen_as_bytes!(u32); +gen_as_bytes!(u64); +gen_as_bytes!(f32); +gen_as_bytes!(f64); + +macro_rules! unimplemented_slice_as_bytes { + ($ty: ty) => { + impl SliceAsBytes for $ty { + fn slice_as_bytes(_self: &[Self]) -> &[u8] { + unimplemented!() + } + + unsafe fn slice_as_bytes_mut(_self: &mut [Self]) -> &mut [u8] { + unimplemented!() + } + } + }; +} + +// TODO - Can Int96 and bool be implemented in these terms? +unimplemented_slice_as_bytes!(Int96); +unimplemented_slice_as_bytes!(bool); +unimplemented_slice_as_bytes!(ByteArray); +unimplemented_slice_as_bytes!(FixedLenByteArray); + +impl AsBytes for bool { + fn as_bytes(&self) -> &[u8] { + // SAFETY: a bool is guaranteed to be either 0x00 or 0x01 in memory, so the memory is + // valid. + unsafe { std::slice::from_raw_parts(self as *const bool as *const u8, 1) } + } +} + +impl AsBytes for Int96 { + fn as_bytes(&self) -> &[u8] { + // SAFETY: Int96::data is a &[u32; 3]. + unsafe { std::slice::from_raw_parts(self.data() as *const [u32] as *const u8, 12) } + } +} + +impl AsBytes for ByteArray { + fn as_bytes(&self) -> &[u8] { + self.data() + } +} + +impl AsBytes for FixedLenByteArray { + fn as_bytes(&self) -> &[u8] { + self.data() + } +} + +impl AsBytes for Decimal { + fn as_bytes(&self) -> &[u8] { + self.data() + } +} + +impl AsBytes for Vec { + fn as_bytes(&self) -> &[u8] { + self.as_slice() + } +} + +impl AsBytes for &str { + fn as_bytes(&self) -> &[u8] { + (self as &str).as_bytes() + } +} + +impl AsBytes for str { + fn as_bytes(&self) -> &[u8] { + (self as &str).as_bytes() + } +} + +pub(crate) mod private { + use bytes::Bytes; + + use crate::encodings::decoding::PlainDecoderDetails; + use crate::util::bit_util::{BitReader, BitWriter, read_num_bytes}; + + use super::{ParquetError, Result, SliceAsBytes}; + use crate::basic::Type; + use crate::file::metadata::HeapSize; + + /// Sealed trait to start to remove specialisation from implementations + /// + /// This is done to force the associated value type to be unimplementable outside of this + /// crate, and thus hint to the type system (and end user) traits are public for the contract + /// and not for extension. + pub trait ParquetValueType: + PartialEq + + std::fmt::Debug + + std::fmt::Display + + Default + + Clone + + super::AsBytes + + super::FromBytes + + SliceAsBytes + + PartialOrd + + Send + + HeapSize + + crate::encodings::decoding::private::GetDecoder + + crate::file::statistics::private::MakeStatistics + { + const PHYSICAL_TYPE: Type; + + /// Encode the value directly from a higher level encoder + fn encode( + values: &[Self], + writer: &mut W, + bit_writer: &mut BitWriter, + ) -> Result<()>; + + /// Establish the data that will be decoded in a buffer + fn set_data(decoder: &mut PlainDecoderDetails, data: Bytes, num_values: usize); + + /// Decode the value from a given buffer for a higher level decoder + fn decode(buffer: &mut [Self], decoder: &mut PlainDecoderDetails) -> Result; + + fn skip(decoder: &mut PlainDecoderDetails, num_values: usize) -> Result; + + /// Return the encoded size for a type + fn dict_encoding_size(&self) -> (usize, usize) { + (std::mem::size_of::(), 1) + } + + /// Return the number of variable length bytes in a given slice of data + /// + /// Returns the sum of lengths for BYTE_ARRAY data, and None for all other data types + fn variable_length_bytes(_: &[Self]) -> Option { + None + } + + /// Return the value as i64 if possible + /// + /// This is essentially the same as `std::convert::TryInto` but can't be + /// implemented for `f32` and `f64`, types that would fail orphan rules + fn as_i64(&self) -> Result { + Err(general_err!("Type cannot be converted to i64")) + } + + /// Return the value as u64 if possible + /// + /// This is essentially the same as `std::convert::TryInto` but can't be + /// implemented for `f32` and `f64`, types that would fail orphan rules + fn as_u64(&self) -> Result { + self.as_i64() + .map_err(|_| general_err!("Type cannot be converted to u64")) + .map(|x| x as u64) + } + + /// Return the value as an Any to allow for downcasts without transmutation + fn as_any(&self) -> &dyn std::any::Any; + + /// Return the value as an mutable Any to allow for downcasts without transmutation + fn as_mut_any(&mut self) -> &mut dyn std::any::Any; + + /// Sets the value of this object from the provided [`Bytes`] + /// + /// Only implemented for `ByteArray` and `FixedLenByteArray`. Will panic for other types. + fn set_from_bytes(&mut self, _data: Bytes) { + unimplemented!(); + } + } + + impl ParquetValueType for bool { + const PHYSICAL_TYPE: Type = Type::BOOLEAN; + + #[inline] + fn encode( + values: &[Self], + _: &mut W, + bit_writer: &mut BitWriter, + ) -> Result<()> { + for value in values { + bit_writer.put_value(*value as u64, 1) + } + Ok(()) + } + + #[inline] + fn set_data(decoder: &mut PlainDecoderDetails, data: Bytes, num_values: usize) { + decoder.bit_reader.replace(BitReader::new(data)); + decoder.num_values = num_values; + } + + #[inline] + fn decode(buffer: &mut [Self], decoder: &mut PlainDecoderDetails) -> Result { + let bit_reader = decoder.bit_reader.as_mut().unwrap(); + let num_values = std::cmp::min(buffer.len(), decoder.num_values); + let values_read = bit_reader.get_batch(&mut buffer[..num_values], 1); + decoder.num_values -= values_read; + Ok(values_read) + } + + fn skip(decoder: &mut PlainDecoderDetails, num_values: usize) -> Result { + let bit_reader = decoder.bit_reader.as_mut().unwrap(); + let num_values = std::cmp::min(num_values, decoder.num_values); + let values_read = bit_reader.skip(num_values, 1); + decoder.num_values -= values_read; + Ok(values_read) + } + + #[inline] + fn as_i64(&self) -> Result { + Ok(*self as i64) + } + + #[inline] + fn as_any(&self) -> &dyn std::any::Any { + self + } + + #[inline] + fn as_mut_any(&mut self) -> &mut dyn std::any::Any { + self + } + } + + macro_rules! impl_from_raw { + ($ty: ty, $physical_ty: expr, $self: ident => $as_i64: block) => { + impl ParquetValueType for $ty { + const PHYSICAL_TYPE: Type = $physical_ty; + + #[inline] + fn encode(values: &[Self], writer: &mut W, _: &mut BitWriter) -> Result<()> { + // SAFETY: Self is one of i32, i64, f32, f64, which have no padding. + let raw = unsafe { + std::slice::from_raw_parts( + values.as_ptr() as *const u8, + std::mem::size_of_val(values), + ) + }; + writer.write_all(raw)?; + + Ok(()) + } + + #[inline] + fn set_data(decoder: &mut PlainDecoderDetails, data: Bytes, num_values: usize) { + decoder.data.replace(data); + decoder.start = 0; + decoder.num_values = num_values; + } + + #[inline] + fn decode(buffer: &mut [Self], decoder: &mut PlainDecoderDetails) -> Result { + let data = decoder.data.as_ref().expect("set_data should have been called"); + let num_values = std::cmp::min(buffer.len(), decoder.num_values); + let bytes_left = data.len() - decoder.start; + let bytes_to_decode = std::mem::size_of::() * num_values; + + if bytes_left < bytes_to_decode { + return Err(eof_err!("Not enough bytes to decode")); + } + + { + // SAFETY: Self has no invalid bit patterns, so writing to the slice + // obtained with slice_as_bytes_mut is always safe. + let raw_buffer = &mut unsafe { Self::slice_as_bytes_mut(buffer) }[..bytes_to_decode]; + raw_buffer.copy_from_slice(data.slice( + decoder.start..decoder.start + bytes_to_decode + ).as_ref()); + }; + decoder.start += bytes_to_decode; + decoder.num_values -= num_values; + + Ok(num_values) + } + + #[inline] + fn skip(decoder: &mut PlainDecoderDetails, num_values: usize) -> Result { + let data = decoder.data.as_ref().expect("set_data should have been called"); + let num_values = num_values.min(decoder.num_values); + let bytes_left = data.len() - decoder.start; + let bytes_to_skip = std::mem::size_of::() * num_values; + + if bytes_left < bytes_to_skip { + return Err(eof_err!("Not enough bytes to skip")); + } + + decoder.start += bytes_to_skip; + decoder.num_values -= num_values; + + Ok(num_values) + } + + #[inline] + fn as_i64(&$self) -> Result { + $as_i64 + } + + #[inline] + fn as_any(&self) -> &dyn std::any::Any { + self + } + + #[inline] + fn as_mut_any(&mut self) -> &mut dyn std::any::Any { + self + } + } + } + } + + impl_from_raw!(i32, Type::INT32, self => { Ok(*self as i64) }); + impl_from_raw!(i64, Type::INT64, self => { Ok(*self) }); + impl_from_raw!(f32, Type::FLOAT, self => { Err(general_err!("Type cannot be converted to i64")) }); + impl_from_raw!(f64, Type::DOUBLE, self => { Err(general_err!("Type cannot be converted to i64")) }); + + impl ParquetValueType for super::Int96 { + const PHYSICAL_TYPE: Type = Type::INT96; + + #[inline] + fn encode( + values: &[Self], + writer: &mut W, + _: &mut BitWriter, + ) -> Result<()> { + for value in values { + let raw = SliceAsBytes::slice_as_bytes(value.data()); + writer.write_all(raw)?; + } + Ok(()) + } + + #[inline] + fn set_data(decoder: &mut PlainDecoderDetails, data: Bytes, num_values: usize) { + decoder.data.replace(data); + decoder.start = 0; + decoder.num_values = num_values; + } + + #[inline] + fn decode(buffer: &mut [Self], decoder: &mut PlainDecoderDetails) -> Result { + // TODO - Remove the duplication between this and the general slice method + let data = decoder + .data + .as_ref() + .expect("set_data should have been called"); + let num_values = std::cmp::min(buffer.len(), decoder.num_values); + let bytes_left = data.len() - decoder.start; + let bytes_to_decode = 12 * num_values; + + if bytes_left < bytes_to_decode { + return Err(eof_err!("Not enough bytes to decode")); + } + + let data_range = data.slice(decoder.start..decoder.start + bytes_to_decode); + let bytes: &[u8] = &data_range; + decoder.start += bytes_to_decode; + + let mut pos = 0; // position in byte array + for item in buffer.iter_mut().take(num_values) { + let elem0 = u32::from_le_bytes(bytes[pos..pos + 4].try_into().unwrap()); + let elem1 = u32::from_le_bytes(bytes[pos + 4..pos + 8].try_into().unwrap()); + let elem2 = u32::from_le_bytes(bytes[pos + 8..pos + 12].try_into().unwrap()); + + item.set_data(elem0, elem1, elem2); + pos += 12; + } + decoder.num_values -= num_values; + + Ok(num_values) + } + + fn skip(decoder: &mut PlainDecoderDetails, num_values: usize) -> Result { + let data = decoder + .data + .as_ref() + .expect("set_data should have been called"); + let num_values = std::cmp::min(num_values, decoder.num_values); + let bytes_left = data.len() - decoder.start; + let bytes_to_skip = 12 * num_values; + + if bytes_left < bytes_to_skip { + return Err(eof_err!("Not enough bytes to skip")); + } + decoder.start += bytes_to_skip; + decoder.num_values -= num_values; + + Ok(num_values) + } + + #[inline] + fn as_any(&self) -> &dyn std::any::Any { + self + } + + #[inline] + fn as_mut_any(&mut self) -> &mut dyn std::any::Any { + self + } + } + + impl HeapSize for super::Int96 { + fn heap_size(&self) -> usize { + 0 // no heap allocations + } + } + + impl ParquetValueType for super::ByteArray { + const PHYSICAL_TYPE: Type = Type::BYTE_ARRAY; + + #[inline] + fn encode( + values: &[Self], + writer: &mut W, + _: &mut BitWriter, + ) -> Result<()> { + for value in values { + let len: u32 = value.len().try_into().unwrap(); + writer.write_all(&len.to_ne_bytes())?; + let raw = value.data(); + writer.write_all(raw)?; + } + Ok(()) + } + + #[inline] + fn set_data(decoder: &mut PlainDecoderDetails, data: Bytes, num_values: usize) { + decoder.data.replace(data); + decoder.start = 0; + decoder.num_values = num_values; + } + + #[inline] + fn decode(buffer: &mut [Self], decoder: &mut PlainDecoderDetails) -> Result { + let data = decoder + .data + .as_mut() + .expect("set_data should have been called"); + let num_values = std::cmp::min(buffer.len(), decoder.num_values); + for val_array in buffer.iter_mut().take(num_values) { + let len: usize = + read_num_bytes::(4, data.slice(decoder.start..).as_ref()) as usize; + decoder.start += std::mem::size_of::(); + + if data.len() < decoder.start + len { + return Err(eof_err!("Not enough bytes to decode")); + } + + val_array.set_data(data.slice(decoder.start..decoder.start + len)); + decoder.start += len; + } + decoder.num_values -= num_values; + + Ok(num_values) + } + + fn variable_length_bytes(values: &[Self]) -> Option { + Some(values.iter().map(|x| x.len() as i64).sum()) + } + + fn skip(decoder: &mut PlainDecoderDetails, num_values: usize) -> Result { + let data = decoder + .data + .as_mut() + .expect("set_data should have been called"); + let num_values = num_values.min(decoder.num_values); + + for _ in 0..num_values { + let len: usize = + read_num_bytes::(4, data.slice(decoder.start..).as_ref()) as usize; + decoder.start += std::mem::size_of::() + len; + } + decoder.num_values -= num_values; + + Ok(num_values) + } + + #[inline] + fn dict_encoding_size(&self) -> (usize, usize) { + (std::mem::size_of::(), self.len()) + } + + #[inline] + fn as_any(&self) -> &dyn std::any::Any { + self + } + + #[inline] + fn as_mut_any(&mut self) -> &mut dyn std::any::Any { + self + } + + #[inline] + fn set_from_bytes(&mut self, data: Bytes) { + self.set_data(data); + } + } + + impl HeapSize for super::ByteArray { + fn heap_size(&self) -> usize { + // note: this is an estimate, not exact, so just return the size + // of the actual data used, don't try to handle the fact that it may + // be shared. + self.data.as_ref().map(|data| data.len()).unwrap_or(0) + } + } + + impl ParquetValueType for super::FixedLenByteArray { + const PHYSICAL_TYPE: Type = Type::FIXED_LEN_BYTE_ARRAY; + + #[inline] + fn encode( + values: &[Self], + writer: &mut W, + _: &mut BitWriter, + ) -> Result<()> { + for value in values { + let raw = value.data(); + writer.write_all(raw)?; + } + Ok(()) + } + + #[inline] + fn set_data(decoder: &mut PlainDecoderDetails, data: Bytes, num_values: usize) { + decoder.data.replace(data); + decoder.start = 0; + decoder.num_values = num_values; + } + + #[inline] + fn decode(buffer: &mut [Self], decoder: &mut PlainDecoderDetails) -> Result { + assert!(decoder.type_length > 0); + + let data = decoder + .data + .as_mut() + .expect("set_data should have been called"); + let num_values = std::cmp::min(buffer.len(), decoder.num_values); + + for item in buffer.iter_mut().take(num_values) { + let len = decoder.type_length as usize; + + if data.len() < decoder.start + len { + return Err(eof_err!("Not enough bytes to decode")); + } + + item.set_data(data.slice(decoder.start..decoder.start + len)); + decoder.start += len; + } + decoder.num_values -= num_values; + + Ok(num_values) + } + + fn skip(decoder: &mut PlainDecoderDetails, num_values: usize) -> Result { + assert!(decoder.type_length > 0); + + let data = decoder + .data + .as_mut() + .expect("set_data should have been called"); + let num_values = std::cmp::min(num_values, decoder.num_values); + for _ in 0..num_values { + let len = decoder.type_length as usize; + + if data.len() < decoder.start + len { + return Err(eof_err!("Not enough bytes to skip")); + } + + decoder.start += len; + } + decoder.num_values -= num_values; + + Ok(num_values) + } + + #[inline] + fn dict_encoding_size(&self) -> (usize, usize) { + (std::mem::size_of::(), self.len()) + } + + #[inline] + fn as_any(&self) -> &dyn std::any::Any { + self + } + + #[inline] + fn as_mut_any(&mut self) -> &mut dyn std::any::Any { + self + } + + #[inline] + fn set_from_bytes(&mut self, data: Bytes) { + self.set_data(data); + } + } + + impl HeapSize for super::FixedLenByteArray { + fn heap_size(&self) -> usize { + self.0.heap_size() + } + } +} + +/// Contains the Parquet physical type information as well as the Rust primitive type +/// presentation. +pub trait DataType: 'static + Send { + /// The physical type of the Parquet data type. + type T: private::ParquetValueType; + + /// Returns Parquet physical type. + fn get_physical_type() -> Type { + ::PHYSICAL_TYPE + } + + /// Returns size in bytes for Rust representation of the physical type. + fn get_type_size() -> usize; + + /// Returns the underlying [`ColumnReaderImpl`] for the given [`ColumnReader`]. + fn get_column_reader(column_writer: ColumnReader) -> Option> + where + Self: Sized; + + /// Returns the underlying [`ColumnWriterImpl`] for the given [`ColumnWriter`]. + fn get_column_writer(column_writer: ColumnWriter<'_>) -> Option> + where + Self: Sized; + + /// Returns a reference to the underlying [`ColumnWriterImpl`] for the given [`ColumnWriter`]. + fn get_column_writer_ref<'a, 'b: 'a>( + column_writer: &'b ColumnWriter<'a>, + ) -> Option<&'b ColumnWriterImpl<'a, Self>> + where + Self: Sized; + + /// Returns a mutable reference to the underlying [`ColumnWriterImpl`] for the given + fn get_column_writer_mut<'a, 'b: 'a>( + column_writer: &'a mut ColumnWriter<'b>, + ) -> Option<&'a mut ColumnWriterImpl<'b, Self>> + where + Self: Sized; +} + +macro_rules! make_type { + ($name:ident, $reader_ident: ident, $writer_ident: ident, $native_ty:ty, $size:expr) => { + #[doc = concat!("Parquet physical type: ", stringify!($name))] + #[derive(Clone)] + pub struct $name {} + + impl DataType for $name { + type T = $native_ty; + + fn get_type_size() -> usize { + $size + } + + fn get_column_reader(column_reader: ColumnReader) -> Option> { + match column_reader { + ColumnReader::$reader_ident(w) => Some(w), + _ => None, + } + } + + fn get_column_writer( + column_writer: ColumnWriter<'_>, + ) -> Option> { + match column_writer { + ColumnWriter::$writer_ident(w) => Some(w), + _ => None, + } + } + + fn get_column_writer_ref<'a, 'b: 'a>( + column_writer: &'a ColumnWriter<'b>, + ) -> Option<&'a ColumnWriterImpl<'b, Self>> { + match column_writer { + ColumnWriter::$writer_ident(w) => Some(w), + _ => None, + } + } + + fn get_column_writer_mut<'a, 'b: 'a>( + column_writer: &'a mut ColumnWriter<'b>, + ) -> Option<&'a mut ColumnWriterImpl<'b, Self>> { + match column_writer { + ColumnWriter::$writer_ident(w) => Some(w), + _ => None, + } + } + } + }; +} + +// Generate struct definitions for all physical types + +make_type!(BoolType, BoolColumnReader, BoolColumnWriter, bool, 1); +make_type!(Int32Type, Int32ColumnReader, Int32ColumnWriter, i32, 4); +make_type!(Int64Type, Int64ColumnReader, Int64ColumnWriter, i64, 8); +make_type!( + Int96Type, + Int96ColumnReader, + Int96ColumnWriter, + Int96, + mem::size_of::() +); +make_type!(FloatType, FloatColumnReader, FloatColumnWriter, f32, 4); +make_type!(DoubleType, DoubleColumnReader, DoubleColumnWriter, f64, 8); +make_type!( + ByteArrayType, + ByteArrayColumnReader, + ByteArrayColumnWriter, + ByteArray, + mem::size_of::() +); +make_type!( + FixedLenByteArrayType, + FixedLenByteArrayColumnReader, + FixedLenByteArrayColumnWriter, + FixedLenByteArray, + mem::size_of::() +); + +impl AsRef<[u8]> for ByteArray { + fn as_ref(&self) -> &[u8] { + self.as_bytes() + } +} + +impl AsRef<[u8]> for FixedLenByteArray { + fn as_ref(&self) -> &[u8] { + self.as_bytes() + } +} + +/// Macro to reduce repetition in making type assertions on the physical type against `T` +macro_rules! ensure_phys_ty { + ($($ty:pat_param)|+ , $($arg:tt)*) => { + match T::get_physical_type() { + $($ty => (),)* + _ => panic!($($arg)*), + }; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_as_bytes() { + // Test Int96 + let i96 = Int96::from(vec![1, 2, 3]); + assert_eq!(i96.as_bytes(), &[1, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0]); + + // Test ByteArray + let ba = ByteArray::from(vec![1, 2, 3]); + assert_eq!(ba.as_bytes(), &[1, 2, 3]); + + // Test Decimal + let decimal = Decimal::from_i32(123, 5, 2); + assert_eq!(decimal.as_bytes(), &[0, 0, 0, 123]); + let decimal = Decimal::from_i64(123, 5, 2); + assert_eq!(decimal.as_bytes(), &[0, 0, 0, 0, 0, 0, 0, 123]); + let decimal = Decimal::from_bytes(ByteArray::from(vec![1, 2, 3]), 5, 2); + assert_eq!(decimal.as_bytes(), &[1, 2, 3]); + } + + #[test] + fn test_int96_from() { + assert_eq!( + Int96::from(vec![1, 12345, 1234567890]).data(), + &[1, 12345, 1234567890] + ); + } + + #[test] + fn test_byte_array_from() { + assert_eq!(ByteArray::from(b"ABC".to_vec()).data(), b"ABC"); + assert_eq!(ByteArray::from("ABC").data(), b"ABC"); + assert_eq!( + ByteArray::from(Bytes::from(vec![1u8, 2u8, 3u8, 4u8, 5u8])).data(), + &[1u8, 2u8, 3u8, 4u8, 5u8] + ); + let buf = vec![6u8, 7u8, 8u8, 9u8, 10u8]; + assert_eq!(ByteArray::from(buf).data(), &[6u8, 7u8, 8u8, 9u8, 10u8]); + } + + #[test] + fn test_decimal_partial_eq() { + assert_eq!(Decimal::default(), Decimal::from_i32(0, 0, 0)); + assert_eq!(Decimal::from_i32(222, 5, 2), Decimal::from_i32(222, 5, 2)); + assert_eq!( + Decimal::from_bytes(ByteArray::from(vec![0, 0, 0, 3]), 5, 2), + Decimal::from_i32(3, 5, 2) + ); + + assert!(Decimal::from_i32(222, 5, 2) != Decimal::from_i32(111, 5, 2)); + assert!(Decimal::from_i32(222, 5, 2) != Decimal::from_i32(222, 6, 2)); + assert!(Decimal::from_i32(222, 5, 2) != Decimal::from_i32(222, 5, 3)); + + assert!(Decimal::from_i64(222, 5, 2) != Decimal::from_i32(222, 5, 2)); + } + + #[test] + fn test_byte_array_ord() { + let ba1 = ByteArray::from(vec![1, 2, 3]); + let ba11 = ByteArray::from(vec![1, 2, 3]); + let ba2 = ByteArray::from(vec![3, 4]); + let ba3 = ByteArray::from(vec![1, 2, 4]); + let ba4 = ByteArray::from(vec![]); + let ba5 = ByteArray::from(vec![2, 2, 3]); + + assert!(ba1 < ba2); + assert!(ba3 > ba1); + assert!(ba1 > ba4); + assert_eq!(ba1, ba11); + assert!(ba5 > ba1); + } +} diff --git a/vendor/parquet/src/encodings/decoding.rs b/vendor/parquet/src/encodings/decoding.rs new file mode 100644 index 0000000..0fb9604 --- /dev/null +++ b/vendor/parquet/src/encodings/decoding.rs @@ -0,0 +1,2444 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Contains all supported decoders for Parquet. + +use bytes::Bytes; +use num_traits::{FromPrimitive, WrappingAdd}; +use std::{cmp, marker::PhantomData, mem}; + +use super::rle::RleDecoder; + +use crate::basic::*; +use crate::data_type::private::ParquetValueType; +use crate::data_type::*; +use crate::encodings::decoding::byte_stream_split_decoder::{ + ByteStreamSplitDecoder, VariableWidthByteStreamSplitDecoder, +}; +use crate::errors::{ParquetError, Result}; +use crate::schema::types::ColumnDescPtr; +use crate::util::bit_util::{self, BitReader, FromBitpacked}; + +mod byte_stream_split_decoder; + +pub(crate) mod private { + use super::*; + + /// A trait that allows getting a [`Decoder`] implementation for a [`DataType`] with + /// the corresponding [`ParquetValueType`]. This is necessary to support + /// [`Decoder`] implementations that may not be applicable for all [`DataType`] + /// and by extension all [`ParquetValueType`] + pub trait GetDecoder { + fn get_decoder>( + descr: ColumnDescPtr, + encoding: Encoding, + ) -> Result>> { + get_decoder_default(descr, encoding) + } + } + + fn get_decoder_default( + descr: ColumnDescPtr, + encoding: Encoding, + ) -> Result>> { + match encoding { + Encoding::PLAIN => Ok(Box::new(PlainDecoder::new(descr.type_length()))), + Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY => Err(general_err!( + "Cannot initialize this encoding through this function" + )), + Encoding::RLE + | Encoding::DELTA_BINARY_PACKED + | Encoding::DELTA_BYTE_ARRAY + | Encoding::DELTA_LENGTH_BYTE_ARRAY => Err(general_err!( + "Encoding {} is not supported for type", + encoding + )), + e => Err(nyi_err!("Encoding {} is not supported", e)), + } + } + + impl GetDecoder for bool { + fn get_decoder>( + descr: ColumnDescPtr, + encoding: Encoding, + ) -> Result>> { + match encoding { + Encoding::RLE => Ok(Box::new(RleValueDecoder::new())), + _ => get_decoder_default(descr, encoding), + } + } + } + + impl GetDecoder for i32 { + fn get_decoder>( + descr: ColumnDescPtr, + encoding: Encoding, + ) -> Result>> { + match encoding { + Encoding::BYTE_STREAM_SPLIT => Ok(Box::new(ByteStreamSplitDecoder::new())), + Encoding::DELTA_BINARY_PACKED => Ok(Box::new(DeltaBitPackDecoder::new())), + _ => get_decoder_default(descr, encoding), + } + } + } + + impl GetDecoder for i64 { + fn get_decoder>( + descr: ColumnDescPtr, + encoding: Encoding, + ) -> Result>> { + match encoding { + Encoding::BYTE_STREAM_SPLIT => Ok(Box::new(ByteStreamSplitDecoder::new())), + Encoding::DELTA_BINARY_PACKED => Ok(Box::new(DeltaBitPackDecoder::new())), + _ => get_decoder_default(descr, encoding), + } + } + } + + impl GetDecoder for f32 { + fn get_decoder>( + descr: ColumnDescPtr, + encoding: Encoding, + ) -> Result>> { + match encoding { + Encoding::BYTE_STREAM_SPLIT => Ok(Box::new(ByteStreamSplitDecoder::new())), + _ => get_decoder_default(descr, encoding), + } + } + } + impl GetDecoder for f64 { + fn get_decoder>( + descr: ColumnDescPtr, + encoding: Encoding, + ) -> Result>> { + match encoding { + Encoding::BYTE_STREAM_SPLIT => Ok(Box::new(ByteStreamSplitDecoder::new())), + _ => get_decoder_default(descr, encoding), + } + } + } + + impl GetDecoder for ByteArray { + fn get_decoder>( + descr: ColumnDescPtr, + encoding: Encoding, + ) -> Result>> { + match encoding { + Encoding::DELTA_BYTE_ARRAY => Ok(Box::new(DeltaByteArrayDecoder::new())), + Encoding::DELTA_LENGTH_BYTE_ARRAY => { + Ok(Box::new(DeltaLengthByteArrayDecoder::new())) + } + _ => get_decoder_default(descr, encoding), + } + } + } + + impl GetDecoder for FixedLenByteArray { + fn get_decoder>( + descr: ColumnDescPtr, + encoding: Encoding, + ) -> Result>> { + match encoding { + Encoding::BYTE_STREAM_SPLIT => Ok(Box::new( + VariableWidthByteStreamSplitDecoder::new(descr.type_length()), + )), + Encoding::DELTA_BYTE_ARRAY => Ok(Box::new(DeltaByteArrayDecoder::new())), + _ => get_decoder_default(descr, encoding), + } + } + } + + impl GetDecoder for Int96 {} +} + +// ---------------------------------------------------------------------- +// Decoders + +/// A Parquet decoder for the data type `T`. +pub trait Decoder: Send { + /// Sets the data to decode to be `data`, which should contain `num_values` of values + /// to decode. + fn set_data(&mut self, data: Bytes, num_values: usize) -> Result<()>; + + /// Consumes values from this decoder and write the results to `buffer`. This will try + /// to fill up `buffer`. + /// + /// Returns the actual number of values decoded, which should be equal to + /// `buffer.len()` unless the remaining number of values is less than + /// `buffer.len()`. + fn get(&mut self, buffer: &mut [T::T]) -> Result; + + /// Consume values from this decoder and write the results to `buffer`, leaving + /// "spaces" for null values. + /// + /// `null_count` is the number of nulls we expect to see in `buffer`, after reading. + /// `valid_bits` stores the valid bit for each value in the buffer. It should contain + /// at least number of bits that equal to `buffer.len()`. + /// + /// Returns the actual number of values decoded. + /// + /// # Panics + /// + /// Panics if `null_count` is greater than `buffer.len()`. + fn get_spaced( + &mut self, + buffer: &mut [T::T], + null_count: usize, + valid_bits: &[u8], + ) -> Result { + assert!(buffer.len() >= null_count); + + // TODO: check validity of the input arguments? + if null_count == 0 { + return self.get(buffer); + } + + let num_values = buffer.len(); + let values_to_read = num_values - null_count; + let values_read = self.get(buffer)?; + if values_read != values_to_read { + return Err(general_err!( + "Number of values read: {}, doesn't match expected: {}", + values_read, + values_to_read + )); + } + let mut values_to_move = values_read; + for i in (0..num_values).rev() { + if bit_util::get_bit(valid_bits, i) { + values_to_move -= 1; + buffer.swap(i, values_to_move); + } + } + + Ok(num_values) + } + + /// Returns the number of values left in this decoder stream. + fn values_left(&self) -> usize; + + /// Returns the encoding for this decoder. + fn encoding(&self) -> Encoding; + + /// Skip the specified number of values in this decoder stream. + fn skip(&mut self, num_values: usize) -> Result; +} + +/// Gets a decoder for the column descriptor `descr` and encoding type `encoding`. +/// +/// NOTE: the primitive type in `descr` MUST match the data type `T`, otherwise +/// disastrous consequence could occur. +pub fn get_decoder( + descr: ColumnDescPtr, + encoding: Encoding, +) -> Result>> { + use self::private::GetDecoder; + T::T::get_decoder(descr, encoding) +} + +// ---------------------------------------------------------------------- +// PLAIN Decoding + +#[derive(Default)] +pub struct PlainDecoderDetails { + // The remaining number of values in the byte array + pub(crate) num_values: usize, + + // The current starting index in the byte array. Not used when `T` is bool. + pub(crate) start: usize, + + // The length for the type `T`. Only used when `T` is `FixedLenByteArrayType` + pub(crate) type_length: i32, + + // The byte array to decode from. Not set if `T` is bool. + pub(crate) data: Option, + + // Read `data` bit by bit. Only set if `T` is bool. + pub(crate) bit_reader: Option, +} + +/// Plain decoding that supports all types. +/// +/// Values are encoded back to back. For native types, data is encoded as little endian. +/// Floating point types are encoded in IEEE. +/// See [`PlainEncoder`](crate::encoding::PlainEncoder) for more information. +pub struct PlainDecoder { + // The binary details needed for decoding + inner: PlainDecoderDetails, + + // To allow `T` in the generic parameter for this struct. This doesn't take any + // space. + _phantom: PhantomData, +} + +impl PlainDecoder { + /// Creates new plain decoder. + pub fn new(type_length: i32) -> Self { + PlainDecoder { + inner: PlainDecoderDetails { + type_length, + num_values: 0, + start: 0, + data: None, + bit_reader: None, + }, + _phantom: PhantomData, + } + } +} + +impl Decoder for PlainDecoder { + #[inline] + fn set_data(&mut self, data: Bytes, num_values: usize) -> Result<()> { + T::T::set_data(&mut self.inner, data, num_values); + Ok(()) + } + + #[inline] + fn values_left(&self) -> usize { + self.inner.num_values + } + + #[inline] + fn encoding(&self) -> Encoding { + Encoding::PLAIN + } + + #[inline] + fn get(&mut self, buffer: &mut [T::T]) -> Result { + T::T::decode(buffer, &mut self.inner) + } + + #[inline] + fn skip(&mut self, num_values: usize) -> Result { + T::T::skip(&mut self.inner, num_values) + } +} + +// ---------------------------------------------------------------------- +// RLE_DICTIONARY/PLAIN_DICTIONARY Decoding + +/// Dictionary decoder. +/// +/// The dictionary encoding builds a dictionary of values encountered in a given column. +/// The dictionary is be stored in a dictionary page per column chunk. +/// See [`DictEncoder`](crate::encoding::DictEncoder) for more information. +pub struct DictDecoder { + // The dictionary, which maps ids to the values + dictionary: Vec, + + // Whether `dictionary` has been initialized + has_dictionary: bool, + + // The decoder for the value ids + rle_decoder: Option, + + // Number of values left in the data stream + num_values: usize, +} + +impl Default for DictDecoder { + fn default() -> Self { + Self::new() + } +} + +impl DictDecoder { + /// Creates new dictionary decoder. + pub fn new() -> Self { + Self { + dictionary: vec![], + has_dictionary: false, + rle_decoder: None, + num_values: 0, + } + } + + /// Decodes and sets values for dictionary using `decoder` decoder. + pub fn set_dict(&mut self, mut decoder: Box>) -> Result<()> { + let num_values = decoder.values_left(); + self.dictionary.resize(num_values, T::T::default()); + let _ = decoder.get(&mut self.dictionary)?; + self.has_dictionary = true; + Ok(()) + } +} + +impl Decoder for DictDecoder { + fn set_data(&mut self, data: Bytes, num_values: usize) -> Result<()> { + // First byte in `data` is bit width + if data.is_empty() { + return Err(eof_err!("Not enough bytes to decode bit_width")); + } + + let bit_width = data.as_ref()[0]; + if bit_width > 32 { + return Err(general_err!( + "Invalid or corrupted RLE bit width {}. Max allowed is 32", + bit_width + )); + } + let mut rle_decoder = RleDecoder::new(bit_width); + rle_decoder.set_data(data.slice(1..))?; + self.num_values = num_values; + self.rle_decoder = Some(rle_decoder); + Ok(()) + } + + fn get(&mut self, buffer: &mut [T::T]) -> Result { + assert!(self.rle_decoder.is_some()); + assert!(self.has_dictionary, "Must call set_dict() first!"); + + let rle = self.rle_decoder.as_mut().unwrap(); + let num_values = cmp::min(buffer.len(), self.num_values); + rle.get_batch_with_dict(&self.dictionary[..], buffer, num_values) + } + + /// Number of values left in this decoder stream + fn values_left(&self) -> usize { + self.num_values + } + + fn encoding(&self) -> Encoding { + Encoding::RLE_DICTIONARY + } + + fn skip(&mut self, num_values: usize) -> Result { + assert!(self.rle_decoder.is_some()); + assert!(self.has_dictionary, "Must call set_dict() first!"); + + let rle = self.rle_decoder.as_mut().unwrap(); + let num_values = cmp::min(num_values, self.num_values); + rle.skip(num_values) + } +} + +// ---------------------------------------------------------------------- +// RLE Decoding + +/// RLE/Bit-Packing hybrid decoding for values. +/// Currently is used only for data pages v2 and supports boolean types. +/// See [`RleValueEncoder`](crate::encoding::RleValueEncoder) for more information. +pub struct RleValueDecoder { + values_left: usize, + decoder: RleDecoder, + _phantom: PhantomData, +} + +impl Default for RleValueDecoder { + fn default() -> Self { + Self::new() + } +} + +impl RleValueDecoder { + pub fn new() -> Self { + Self { + values_left: 0, + decoder: RleDecoder::new(1), + _phantom: PhantomData, + } + } +} + +impl Decoder for RleValueDecoder +where + T::T: FromBitpacked, +{ + #[inline] + fn set_data(&mut self, data: Bytes, num_values: usize) -> Result<()> { + // Only support RLE value reader for boolean values with bit width of 1. + ensure_phys_ty!(Type::BOOLEAN, "RleValueDecoder only supports BoolType"); + + // We still need to remove prefix of i32 from the stream. + const I32_SIZE: usize = mem::size_of::(); + if data.len() < I32_SIZE { + return Err(eof_err!("Not enough bytes to decode")); + } + let data_size = bit_util::read_num_bytes::(I32_SIZE, data.as_ref()) as usize; + if data.len() - I32_SIZE < data_size { + return Err(eof_err!("Not enough bytes to decode")); + } + + self.decoder = RleDecoder::new(1); + self.decoder + .set_data(data.slice(I32_SIZE..I32_SIZE + data_size))?; + self.values_left = num_values; + Ok(()) + } + + #[inline] + fn values_left(&self) -> usize { + self.values_left + } + + #[inline] + fn encoding(&self) -> Encoding { + Encoding::RLE + } + + #[inline] + fn get(&mut self, buffer: &mut [T::T]) -> Result { + let num_values = cmp::min(buffer.len(), self.values_left); + let values_read = self.decoder.get_batch(&mut buffer[..num_values])?; + self.values_left -= values_read; + Ok(values_read) + } + + #[inline] + fn skip(&mut self, num_values: usize) -> Result { + let num_values = cmp::min(num_values, self.values_left); + let values_skipped = self.decoder.skip(num_values)?; + self.values_left -= values_skipped; + Ok(values_skipped) + } +} + +// ---------------------------------------------------------------------- +// DELTA_BINARY_PACKED Decoding + +/// Delta binary packed decoder. +/// Supports INT32 and INT64 types. +/// See [`DeltaBitPackEncoder`](crate::encoding::DeltaBitPackEncoder) for more +/// information. +pub struct DeltaBitPackDecoder { + bit_reader: BitReader, + initialized: bool, + + // Header info + /// The number of values in each block + block_size: usize, + /// The number of values that remain to be read in the current page + values_left: usize, + /// The number of mini-blocks in each block + mini_blocks_per_block: usize, + /// The number of values in each mini block + values_per_mini_block: usize, + + // Per block info + /// The minimum delta in the block + min_delta: T::T, + /// The byte offset of the end of the current block + block_end_offset: usize, + /// The index on the current mini block + mini_block_idx: usize, + /// The bit widths of each mini block in the current block + mini_block_bit_widths: Vec, + /// The number of values remaining in the current mini block + mini_block_remaining: usize, + + /// The first value from the block header if not consumed + first_value: Option, + /// The last value to compute offsets from + last_value: T::T, +} + +impl Default for DeltaBitPackDecoder +where + T::T: Default + FromPrimitive + WrappingAdd + Copy, +{ + fn default() -> Self { + Self::new() + } +} + +impl DeltaBitPackDecoder +where + T::T: Default + FromPrimitive + WrappingAdd + Copy, +{ + /// Creates new delta bit packed decoder. + pub fn new() -> Self { + Self { + bit_reader: BitReader::from(vec![]), + initialized: false, + block_size: 0, + values_left: 0, + mini_blocks_per_block: 0, + values_per_mini_block: 0, + min_delta: Default::default(), + mini_block_idx: 0, + mini_block_bit_widths: vec![], + mini_block_remaining: 0, + block_end_offset: 0, + first_value: None, + last_value: Default::default(), + } + } + + /// Returns the current offset + pub fn get_offset(&self) -> usize { + assert!(self.initialized, "Bit reader is not initialized"); + match self.values_left { + // If we've exhausted this page report the end of the current block + // as we may not have consumed the trailing padding + // + // The max is necessary to handle pages which don't contain more than + // one value and therefore have no blocks, but still contain a page header + 0 => self.bit_reader.get_byte_offset().max(self.block_end_offset), + _ => self.bit_reader.get_byte_offset(), + } + } + + /// Initializes the next block and the first mini block within it + #[inline] + fn next_block(&mut self) -> Result<()> { + let min_delta = self + .bit_reader + .get_zigzag_vlq_int() + .ok_or_else(|| eof_err!("Not enough data to decode 'min_delta'"))?; + + self.min_delta = + T::T::from_i64(min_delta).ok_or_else(|| general_err!("'min_delta' too large"))?; + + self.mini_block_bit_widths.clear(); + self.bit_reader + .get_aligned_bytes(&mut self.mini_block_bit_widths, self.mini_blocks_per_block); + + let mut offset = self.bit_reader.get_byte_offset(); + let mut remaining = self.values_left; + + // Compute the end offset of the current block + for b in &mut self.mini_block_bit_widths { + if remaining == 0 { + // Specification requires handling arbitrary bit widths + // for trailing mini blocks + *b = 0; + } + remaining = remaining.saturating_sub(self.values_per_mini_block); + offset += *b as usize * self.values_per_mini_block / 8; + } + self.block_end_offset = offset; + + if self.mini_block_bit_widths.len() != self.mini_blocks_per_block { + return Err(eof_err!("insufficient mini block bit widths")); + } + + self.mini_block_remaining = self.values_per_mini_block; + self.mini_block_idx = 0; + + Ok(()) + } + + /// Initializes the next mini block + #[inline] + fn next_mini_block(&mut self) -> Result<()> { + if self.mini_block_idx + 1 < self.mini_block_bit_widths.len() { + self.mini_block_idx += 1; + self.mini_block_remaining = self.values_per_mini_block; + Ok(()) + } else { + self.next_block() + } + } + + /// Verify the bit width is smaller then the integer type that it is trying to decode. + #[inline] + fn check_bit_width(&self, bit_width: usize) -> Result<()> { + if bit_width > std::mem::size_of::() * 8 { + return Err(general_err!( + "Invalid delta bit width {} which is larger than expected {} ", + bit_width, + std::mem::size_of::() * 8 + )); + } + Ok(()) + } +} + +impl Decoder for DeltaBitPackDecoder +where + T::T: Default + FromPrimitive + FromBitpacked + WrappingAdd + Copy, +{ + // # of total values is derived from encoding + #[inline] + fn set_data(&mut self, data: Bytes, _index: usize) -> Result<()> { + self.bit_reader = BitReader::new(data); + self.initialized = true; + + // Read header information + self.block_size = self + .bit_reader + .get_vlq_int() + .ok_or_else(|| eof_err!("Not enough data to decode 'block_size'"))? + .try_into() + .map_err(|_| general_err!("invalid 'block_size'"))?; + + self.mini_blocks_per_block = self + .bit_reader + .get_vlq_int() + .ok_or_else(|| eof_err!("Not enough data to decode 'mini_blocks_per_block'"))? + .try_into() + .map_err(|_| general_err!("invalid 'mini_blocks_per_block'"))?; + + if self.mini_blocks_per_block == 0 { + return Err(general_err!("cannot have zero miniblocks per block")); + } + + self.values_left = self + .bit_reader + .get_vlq_int() + .ok_or_else(|| eof_err!("Not enough data to decode 'values_left'"))? + .try_into() + .map_err(|_| general_err!("invalid 'values_left'"))?; + + let first_value = self + .bit_reader + .get_zigzag_vlq_int() + .ok_or_else(|| eof_err!("Not enough data to decode 'first_value'"))?; + + self.first_value = + Some(T::T::from_i64(first_value).ok_or_else(|| general_err!("first value too large"))?); + + if self.block_size % 128 != 0 { + return Err(general_err!( + "'block_size' must be a multiple of 128, got {}", + self.block_size + )); + } + + if self.block_size % self.mini_blocks_per_block != 0 { + return Err(general_err!( + "'block_size' must be a multiple of 'mini_blocks_per_block' got {} and {}", + self.block_size, + self.mini_blocks_per_block + )); + } + + // Reset decoding state + self.mini_block_idx = 0; + self.values_per_mini_block = self.block_size / self.mini_blocks_per_block; + self.mini_block_remaining = 0; + self.mini_block_bit_widths.clear(); + + if self.values_per_mini_block % 32 != 0 { + return Err(general_err!( + "'values_per_mini_block' must be a multiple of 32 got {}", + self.values_per_mini_block + )); + } + + Ok(()) + } + + fn get(&mut self, buffer: &mut [T::T]) -> Result { + assert!(self.initialized, "Bit reader is not initialized"); + if buffer.is_empty() { + return Ok(0); + } + + let mut read = 0; + let to_read = buffer.len().min(self.values_left); + + if let Some(value) = self.first_value.take() { + self.last_value = value; + buffer[0] = value; + read += 1; + self.values_left -= 1; + } + + while read != to_read { + if self.mini_block_remaining == 0 { + self.next_mini_block()?; + } + + let bit_width = self.mini_block_bit_widths[self.mini_block_idx] as usize; + self.check_bit_width(bit_width)?; + let batch_to_read = self.mini_block_remaining.min(to_read - read); + + let batch_read = self + .bit_reader + .get_batch(&mut buffer[read..read + batch_to_read], bit_width); + + if batch_read != batch_to_read { + return Err(general_err!( + "Expected to read {} values from miniblock got {}", + batch_to_read, + batch_read + )); + } + + // At this point we have read the deltas to `buffer` we now need to offset + // these to get back to the original values that were encoded + // + // Optimization: if the bit_width for the miniblock is 0, then we can employ + // a faster decoding method than setting `value[i] = value[i-1] + value[i] + min_delta`. + // Where min_delta is 0 (all values in the miniblock are the same), we can simply + // set all values to `self.last_value`. In the case of non-zero min_delta (values + // in the mini-block form an arithmetic progression) each value can be computed via + // `value[i] = (i + 1) * min_delta + last_value`. In both cases we remove the + // dependence on the preceding value. + // Kudos to @pitrou for the idea https://github.com/apache/arrow/pull/49296 + let min_delta = self.min_delta.as_i64()?; + if bit_width == 0 { + if min_delta == 0 { + buffer[read..read + batch_read].fill(self.last_value); + } else { + // the c++ version multiplies min_delta by the iter index, but doing + // wrapping_mul through T::T was a bit slower. this is still + // faster than before. + let mut delta = self.min_delta; + for v in &mut buffer[read..read + batch_read] { + *v = self.last_value.wrapping_add(&delta); + delta = delta.wrapping_add(&self.min_delta); + } + + self.last_value = buffer[read + batch_read - 1]; + } + } else { + // It is OK for deltas to contain "overflowed" values after encoding, + // e.g. i64::MAX - i64::MIN, so we use `wrapping_add` to "overflow" again and + // restore original value. + if min_delta == 0 { + for v in &mut buffer[read..read + batch_read] { + *v = v.wrapping_add(&self.last_value); + self.last_value = *v; + } + } else { + for v in &mut buffer[read..read + batch_read] { + *v = v + .wrapping_add(&self.min_delta) + .wrapping_add(&self.last_value); + self.last_value = *v; + } + } + } + + read += batch_read; + self.mini_block_remaining -= batch_read; + self.values_left -= batch_read; + } + + Ok(to_read) + } + + fn values_left(&self) -> usize { + self.values_left + } + + fn encoding(&self) -> Encoding { + Encoding::DELTA_BINARY_PACKED + } + + fn skip(&mut self, num_values: usize) -> Result { + let mut skip = 0; + let to_skip = num_values.min(self.values_left); + if to_skip == 0 { + return Ok(0); + } + + // try to consume first value in header. + if let Some(value) = self.first_value.take() { + self.last_value = value; + skip += 1; + self.values_left -= 1; + } + + // See https://github.com/apache/arrow-rs/pull/9794. + // The parquet spec actually allows for miniblock sizes other than 32 or 64, but + // no current writers use anything else. Using values_per_mini_block directly + // for the skip_buffer doesn't allow stack allocation and leads to a significant + // drop in performance. We'll settle for erroring out here and come up with a + // better fix if writers ever start getting creative with block sizes. + let mini_block_batch_size = match self.values_per_mini_block { + 32 => 32, + 64 => 64, + _ => { + return Err(general_err!( + "cannot skip miniblock of size {}", + self.values_per_mini_block + )); + } + }; + + let mut skip_buffer = vec![T::T::default(); mini_block_batch_size]; + while skip < to_skip { + if self.mini_block_remaining == 0 { + self.next_mini_block()?; + } + + let bit_width = self.mini_block_bit_widths[self.mini_block_idx] as usize; + self.check_bit_width(bit_width)?; + let mini_block_to_skip = self.mini_block_remaining.min(to_skip - skip); + + // see commentary in self.get() above regarding optimizations + let min_delta = self.min_delta.as_i64()?; + if bit_width == 0 { + // All remainders are zero: every delta equals min_delta exactly. + // Advance last_value by n * min_delta with no bit reads. + // When min_delta == 0 there is nothing to do: last_value is + // unchanged and no bytes are consumed from the bit reader. + if min_delta != 0 { + let total = min_delta.wrapping_mul(mini_block_to_skip as i64); + let step = T::T::from_i64(total) + .ok_or_else(|| general_err!("delta*n overflow in skip"))?; + self.last_value = self.last_value.wrapping_add(&step); + } + // bit_width=0 payloads occupy zero bytes; no bit_reader advancement needed. + } else { + // bw>0: must decode to track last_value for subsequent get() calls. + let skip_count = self + .bit_reader + .get_batch(&mut skip_buffer[0..mini_block_to_skip], bit_width); + + if skip_count != mini_block_to_skip { + return Err(general_err!( + "Expected to skip {} values from mini block got {}.", + mini_block_to_skip, + skip_count + )); + } + + if min_delta == 0 { + for v in &mut skip_buffer[0..skip_count] { + *v = v.wrapping_add(&self.last_value); + self.last_value = *v; + } + } else { + for v in &mut skip_buffer[0..skip_count] { + *v = v + .wrapping_add(&self.min_delta) + .wrapping_add(&self.last_value); + self.last_value = *v; + } + } + } + + skip += mini_block_to_skip; + self.mini_block_remaining -= mini_block_to_skip; + self.values_left -= mini_block_to_skip; + } + + Ok(to_skip) + } +} + +// ---------------------------------------------------------------------- +// DELTA_LENGTH_BYTE_ARRAY Decoding + +/// Delta length byte array decoder. +/// +/// Only applied to byte arrays to separate the length values and the data, the lengths +/// are encoded using DELTA_BINARY_PACKED encoding. +/// See [`DeltaLengthByteArrayEncoder`](crate::encoding::DeltaLengthByteArrayEncoder) +/// for more information. +pub struct DeltaLengthByteArrayDecoder { + // Lengths for each byte array in `data` + // TODO: add memory tracker to this + lengths: Vec, + + // Current index into `lengths` + current_idx: usize, + + // Concatenated byte array data + data: Option, + + // Offset into `data`, always point to the beginning of next byte array. + offset: usize, + + // Number of values left in this decoder stream + num_values: usize, + + // Placeholder to allow `T` as generic parameter + _phantom: PhantomData, +} + +impl Default for DeltaLengthByteArrayDecoder { + fn default() -> Self { + Self::new() + } +} + +impl DeltaLengthByteArrayDecoder { + /// Creates new delta length byte array decoder. + pub fn new() -> Self { + Self { + lengths: vec![], + current_idx: 0, + data: None, + offset: 0, + num_values: 0, + _phantom: PhantomData, + } + } +} + +impl Decoder for DeltaLengthByteArrayDecoder { + fn set_data(&mut self, data: Bytes, num_values: usize) -> Result<()> { + match T::get_physical_type() { + Type::BYTE_ARRAY => { + let mut len_decoder = DeltaBitPackDecoder::::new(); + len_decoder.set_data(data.clone(), num_values)?; + let num_lengths = len_decoder.values_left(); + self.lengths.resize(num_lengths, 0); + len_decoder.get(&mut self.lengths[..])?; + + self.data = Some(data.slice(len_decoder.get_offset()..)); + self.offset = 0; + self.current_idx = 0; + self.num_values = num_lengths; + Ok(()) + } + _ => Err(general_err!( + "DeltaLengthByteArrayDecoder only support ByteArrayType" + )), + } + } + + fn get(&mut self, buffer: &mut [T::T]) -> Result { + match T::get_physical_type() { + Type::BYTE_ARRAY => { + assert!(self.data.is_some()); + + let data = self.data.as_ref().unwrap(); + let num_values = cmp::min(buffer.len(), self.num_values); + + for item in buffer.iter_mut().take(num_values) { + let len = self.lengths[self.current_idx] as usize; + item.set_from_bytes(data.slice(self.offset..self.offset + len)); + + self.offset += len; + self.current_idx += 1; + } + + self.num_values -= num_values; + Ok(num_values) + } + _ => Err(general_err!( + "DeltaLengthByteArrayDecoder only support ByteArrayType" + )), + } + } + + fn values_left(&self) -> usize { + self.num_values + } + + fn encoding(&self) -> Encoding { + Encoding::DELTA_LENGTH_BYTE_ARRAY + } + + fn skip(&mut self, num_values: usize) -> Result { + match T::get_physical_type() { + Type::BYTE_ARRAY => { + let num_values = cmp::min(num_values, self.num_values); + + let next_offset: i32 = self.lengths + [self.current_idx..self.current_idx + num_values] + .iter() + .sum(); + + self.current_idx += num_values; + self.offset += next_offset as usize; + + self.num_values -= num_values; + Ok(num_values) + } + other_type => Err(general_err!( + "DeltaLengthByteArrayDecoder not support {}, only support byte array", + other_type + )), + } + } +} + +// ---------------------------------------------------------------------- +// DELTA_BYTE_ARRAY Decoding + +/// Delta byte array decoder. +/// +/// Prefix lengths are encoded using `DELTA_BINARY_PACKED` encoding, Suffixes are stored +/// using `DELTA_LENGTH_BYTE_ARRAY` encoding. +/// See [`DeltaByteArrayEncoder`](crate::encoding::DeltaByteArrayEncoder) for more +/// information. +pub struct DeltaByteArrayDecoder { + // Prefix lengths for each byte array + // TODO: add memory tracker to this + prefix_lengths: Vec, + + // The current index into `prefix_lengths`, + current_idx: usize, + + // Decoder for all suffixes, the # of which should be the same as + // `prefix_lengths.len()` + suffix_decoder: Option>, + + // The last byte array, used to derive the current prefix + // Stored as Bytes to avoid clone allocation when creating output + previous_value: Bytes, + + // Number of values left + num_values: usize, + + // Placeholder to allow `T` as generic parameter + _phantom: PhantomData, +} + +impl Default for DeltaByteArrayDecoder { + fn default() -> Self { + Self::new() + } +} + +impl DeltaByteArrayDecoder { + /// Creates new delta byte array decoder. + pub fn new() -> Self { + Self { + prefix_lengths: vec![], + current_idx: 0, + suffix_decoder: None, + previous_value: Bytes::new(), + num_values: 0, + _phantom: PhantomData, + } + } +} + +impl Decoder for DeltaByteArrayDecoder { + fn set_data(&mut self, data: Bytes, num_values: usize) -> Result<()> { + match T::get_physical_type() { + Type::BYTE_ARRAY | Type::FIXED_LEN_BYTE_ARRAY => { + let mut prefix_len_decoder = DeltaBitPackDecoder::::new(); + prefix_len_decoder.set_data(data.clone(), num_values)?; + let num_prefixes = prefix_len_decoder.values_left(); + self.prefix_lengths.resize(num_prefixes, 0); + prefix_len_decoder.get(&mut self.prefix_lengths[..])?; + + let mut suffix_decoder = DeltaLengthByteArrayDecoder::new(); + suffix_decoder + .set_data(data.slice(prefix_len_decoder.get_offset()..), num_values)?; + self.suffix_decoder = Some(suffix_decoder); + self.num_values = num_prefixes; + self.current_idx = 0; + self.previous_value = Bytes::new(); + Ok(()) + } + _ => Err(general_err!( + "DeltaByteArrayDecoder only supports ByteArrayType and FixedLenByteArrayType" + )), + } + } + + fn get(&mut self, buffer: &mut [T::T]) -> Result { + match T::get_physical_type() { + Type::BYTE_ARRAY | Type::FIXED_LEN_BYTE_ARRAY => { + let num_values = cmp::min(buffer.len(), self.num_values); + let mut v: [ByteArray; 1] = [ByteArray::new(); 1]; + for item in buffer.iter_mut().take(num_values) { + // Process suffix + // TODO: this is awkward - maybe we should add a non-vectorized API? + let suffix_decoder = self + .suffix_decoder + .as_mut() + .expect("decoder not initialized"); + suffix_decoder.get(&mut v[..])?; + let suffix = v[0].data(); + + // Extract current prefix length, can be 0 + let prefix_len = usize::try_from(self.prefix_lengths[self.current_idx]) + .map_err(|_| { + general_err!( + "Invalid DELTA_BYTE_ARRAY prefix length {}", + self.prefix_lengths[self.current_idx] + ) + })?; + + if prefix_len > self.previous_value.len() { + return Err(general_err!( + "Invalid DELTA_BYTE_ARRAY prefix length {} exceeds previous value length {}", + prefix_len, + self.previous_value.len() + )); + } + + // Concatenate prefix with suffix + let mut result = Vec::with_capacity(prefix_len + suffix.len()); + result.extend_from_slice(&self.previous_value[0..prefix_len]); + result.extend_from_slice(suffix); + + let data = Bytes::from(result); + item.set_from_bytes(data.clone()); + + self.previous_value = data; + self.current_idx += 1; + } + + self.num_values -= num_values; + Ok(num_values) + } + _ => Err(general_err!( + "DeltaByteArrayDecoder only supports ByteArrayType and FixedLenByteArrayType" + )), + } + } + + fn values_left(&self) -> usize { + self.num_values + } + + fn encoding(&self) -> Encoding { + Encoding::DELTA_BYTE_ARRAY + } + + fn skip(&mut self, num_values: usize) -> Result { + let mut buffer = vec![T::T::default(); num_values]; + self.get(&mut buffer) + } +} + +#[cfg(test)] +mod tests { + use super::{super::encoding::*, *}; + + use std::f32::consts::PI as PI_f32; + use std::f64::consts::PI as PI_f64; + use std::sync::Arc; + + use crate::schema::types::{ColumnDescPtr, ColumnDescriptor, ColumnPath, Type as SchemaType}; + use crate::util::test_common::rand_gen::RandGen; + + #[test] + fn test_delta_byte_array_invalid_prefix_len_returns_error() { + let col_descr = create_test_col_desc_ptr(-1, Type::BYTE_ARRAY); + + let mut encoder = + get_encoder::(Encoding::DELTA_BYTE_ARRAY, &col_descr).unwrap(); + let input = vec![ByteArray::from("a"), ByteArray::from("ab")]; + encoder.put(&input).unwrap(); + let encoded = encoder.flush_buffer().unwrap(); + + // First, decode just the prefix-length stream so we know where the suffix stream starts. + let mut prefix_len_decoder = DeltaBitPackDecoder::::new(); + prefix_len_decoder + .set_data(encoded.clone(), input.len()) + .unwrap(); + let num_prefixes = prefix_len_decoder.values_left(); + let mut prefix_lengths = vec![0; num_prefixes]; + prefix_len_decoder.get(&mut prefix_lengths).unwrap(); + + // check: valid encoding should produce prefix lengths [0, 1] + assert_eq!(prefix_lengths, vec![0, 1]); + + let prefix_stream_end = prefix_len_decoder.get_offset(); + + // Corrupt the prefix-length stream itself: + // replace it with a valid DELTA_BINARY_PACKED stream for [1, 1], + // so the first decoded prefix length becomes impossible because previous_value is empty. + let mut prefix_encoder = get_encoder::( + Encoding::DELTA_BINARY_PACKED, + &create_test_col_desc_ptr(-1, Type::INT32), + ) + .unwrap(); + prefix_encoder.put(&[1i32, 1i32]).unwrap(); + let corrupted_prefix = prefix_encoder.flush_buffer().unwrap(); + + let mut corrupted = Vec::new(); + corrupted.extend_from_slice(corrupted_prefix.as_ref()); + corrupted.extend_from_slice(&encoded[prefix_stream_end..]); + + let mut decoder = DeltaByteArrayDecoder::::new(); + decoder + .set_data(Bytes::from(corrupted), input.len()) + .unwrap(); + + let mut out = vec![ByteArray::new(); input.len()]; + + let err = decoder.get(&mut out).unwrap_err(); + assert!( + err.to_string() + .contains("Invalid DELTA_BYTE_ARRAY prefix length"), + "{}", + err + ); + } + + #[test] + fn test_delta_byte_array_negative_prefix_len_returns_error() { + let col_descr = create_test_col_desc_ptr(-1, Type::BYTE_ARRAY); + + let mut encoder = + get_encoder::(Encoding::DELTA_BYTE_ARRAY, &col_descr).unwrap(); + let input = vec![ByteArray::from("a"), ByteArray::from("ab")]; + encoder.put(&input).unwrap(); + let encoded = encoder.flush_buffer().unwrap(); + + let mut decoder = DeltaByteArrayDecoder::::new(); + decoder.set_data(encoded, input.len()).unwrap(); + + // Force a negative prefix length after decoder initialization + decoder.prefix_lengths[0] = -1; + let mut out = vec![ByteArray::new(); input.len()]; + + let err = decoder.get(&mut out).unwrap_err(); + assert!( + err.to_string() + .contains("Invalid DELTA_BYTE_ARRAY prefix length"), + "{}", + err + ); + } + + #[test] + fn test_get_decoders() { + // supported encodings + create_and_check_decoder::(Encoding::PLAIN, None); + create_and_check_decoder::(Encoding::DELTA_BINARY_PACKED, None); + create_and_check_decoder::(Encoding::DELTA_LENGTH_BYTE_ARRAY, None); + create_and_check_decoder::(Encoding::DELTA_BYTE_ARRAY, None); + create_and_check_decoder::(Encoding::RLE, None); + + // error when initializing + create_and_check_decoder::( + Encoding::RLE_DICTIONARY, + Some(general_err!( + "Cannot initialize this encoding through this function" + )), + ); + create_and_check_decoder::( + Encoding::PLAIN_DICTIONARY, + Some(general_err!( + "Cannot initialize this encoding through this function" + )), + ); + create_and_check_decoder::( + Encoding::DELTA_LENGTH_BYTE_ARRAY, + Some(general_err!( + "Encoding DELTA_LENGTH_BYTE_ARRAY is not supported for type" + )), + ); + create_and_check_decoder::( + Encoding::DELTA_BYTE_ARRAY, + Some(general_err!( + "Encoding DELTA_BYTE_ARRAY is not supported for type" + )), + ); + + // unsupported + #[allow(deprecated)] + create_and_check_decoder::( + Encoding::BIT_PACKED, + Some(nyi_err!("Encoding BIT_PACKED is not supported")), + ); + } + + #[test] + fn test_plain_decode_int32() { + let data = [42, 18, 52]; + let data_bytes = Int32Type::to_byte_array(&data[..]); + let mut buffer = [0; 3]; + test_plain_decode::(Bytes::from(data_bytes), 3, -1, &mut buffer[..], &data[..]); + } + + #[test] + fn test_plain_skip_int32() { + let data = [42, 18, 52]; + let data_bytes = Int32Type::to_byte_array(&data[..]); + test_plain_skip::(Bytes::from(data_bytes), 3, 1, -1, &data[1..]); + } + + #[test] + fn test_plain_skip_all_int32() { + let data = [42, 18, 52]; + let data_bytes = Int32Type::to_byte_array(&data[..]); + test_plain_skip::(Bytes::from(data_bytes), 3, 5, -1, &[]); + } + + #[test] + fn test_plain_decode_int32_spaced() { + let data = [42, 18, 52]; + let expected_data = [0, 42, 0, 18, 0, 0, 52, 0]; + let data_bytes = Int32Type::to_byte_array(&data[..]); + let mut buffer = [0; 8]; + let num_nulls = 5; + let valid_bits = [0b01001010]; + test_plain_decode_spaced::( + Bytes::from(data_bytes), + 3, + -1, + &mut buffer[..], + num_nulls, + &valid_bits, + &expected_data[..], + ); + } + + #[test] + fn test_plain_decode_int64() { + let data = [42, 18, 52]; + let data_bytes = Int64Type::to_byte_array(&data[..]); + let mut buffer = [0; 3]; + test_plain_decode::(Bytes::from(data_bytes), 3, -1, &mut buffer[..], &data[..]); + } + + #[test] + fn test_plain_skip_int64() { + let data = [42, 18, 52]; + let data_bytes = Int64Type::to_byte_array(&data[..]); + test_plain_skip::(Bytes::from(data_bytes), 3, 2, -1, &data[2..]); + } + + #[test] + fn test_plain_skip_all_int64() { + let data = [42, 18, 52]; + let data_bytes = Int64Type::to_byte_array(&data[..]); + test_plain_skip::(Bytes::from(data_bytes), 3, 3, -1, &[]); + } + + #[test] + fn test_plain_decode_float() { + let data = [PI_f32, 2.414, 12.51]; + let data_bytes = FloatType::to_byte_array(&data[..]); + let mut buffer = [0.0; 3]; + test_plain_decode::(Bytes::from(data_bytes), 3, -1, &mut buffer[..], &data[..]); + } + + #[test] + fn test_plain_skip_float() { + let data = [PI_f32, 2.414, 12.51]; + let data_bytes = FloatType::to_byte_array(&data[..]); + test_plain_skip::(Bytes::from(data_bytes), 3, 1, -1, &data[1..]); + } + + #[test] + fn test_plain_skip_all_float() { + let data = [PI_f32, 2.414, 12.51]; + let data_bytes = FloatType::to_byte_array(&data[..]); + test_plain_skip::(Bytes::from(data_bytes), 3, 4, -1, &[]); + } + + #[test] + fn test_plain_skip_double() { + let data = [PI_f64, 2.414f64, 12.51f64]; + let data_bytes = DoubleType::to_byte_array(&data[..]); + test_plain_skip::(Bytes::from(data_bytes), 3, 1, -1, &data[1..]); + } + + #[test] + fn test_plain_skip_all_double() { + let data = [PI_f64, 2.414f64, 12.51f64]; + let data_bytes = DoubleType::to_byte_array(&data[..]); + test_plain_skip::(Bytes::from(data_bytes), 3, 5, -1, &[]); + } + + #[test] + fn test_plain_decode_double() { + let data = [PI_f64, 2.414f64, 12.51f64]; + let data_bytes = DoubleType::to_byte_array(&data[..]); + let mut buffer = [0.0f64; 3]; + test_plain_decode::(Bytes::from(data_bytes), 3, -1, &mut buffer[..], &data[..]); + } + + #[test] + fn test_plain_decode_int96() { + let mut data = [Int96::new(); 4]; + data[0].set_data(11, 22, 33); + data[1].set_data(44, 55, 66); + data[2].set_data(10, 20, 30); + data[3].set_data(40, 50, 60); + let data_bytes = Int96Type::to_byte_array(&data[..]); + let mut buffer = [Int96::new(); 4]; + test_plain_decode::(Bytes::from(data_bytes), 4, -1, &mut buffer[..], &data[..]); + } + + #[test] + fn test_plain_skip_int96() { + let mut data = [Int96::new(); 4]; + data[0].set_data(11, 22, 33); + data[1].set_data(44, 55, 66); + data[2].set_data(10, 20, 30); + data[3].set_data(40, 50, 60); + let data_bytes = Int96Type::to_byte_array(&data[..]); + test_plain_skip::(Bytes::from(data_bytes), 4, 2, -1, &data[2..]); + } + + #[test] + fn test_plain_skip_all_int96() { + let mut data = [Int96::new(); 4]; + data[0].set_data(11, 22, 33); + data[1].set_data(44, 55, 66); + data[2].set_data(10, 20, 30); + data[3].set_data(40, 50, 60); + let data_bytes = Int96Type::to_byte_array(&data[..]); + test_plain_skip::(Bytes::from(data_bytes), 4, 8, -1, &[]); + } + + #[test] + fn test_plain_decode_bool() { + let data = [ + false, true, false, false, true, false, true, true, false, true, + ]; + let data_bytes = BoolType::to_byte_array(&data[..]); + let mut buffer = [false; 10]; + test_plain_decode::(Bytes::from(data_bytes), 10, -1, &mut buffer[..], &data[..]); + } + + #[test] + fn test_plain_skip_bool() { + let data = [ + false, true, false, false, true, false, true, true, false, true, + ]; + let data_bytes = BoolType::to_byte_array(&data[..]); + test_plain_skip::(Bytes::from(data_bytes), 10, 5, -1, &data[5..]); + } + + #[test] + fn test_plain_skip_all_bool() { + let data = [ + false, true, false, false, true, false, true, true, false, true, + ]; + let data_bytes = BoolType::to_byte_array(&data[..]); + test_plain_skip::(Bytes::from(data_bytes), 10, 20, -1, &[]); + } + + #[test] + fn test_plain_decode_byte_array() { + let mut data = vec![ByteArray::new(); 2]; + data[0].set_data(Bytes::from(String::from("hello"))); + data[1].set_data(Bytes::from(String::from("parquet"))); + let data_bytes = ByteArrayType::to_byte_array(&data[..]); + let mut buffer = vec![ByteArray::new(); 2]; + test_plain_decode::( + Bytes::from(data_bytes), + 2, + -1, + &mut buffer[..], + &data[..], + ); + } + + #[test] + fn test_plain_skip_byte_array() { + let mut data = vec![ByteArray::new(); 2]; + data[0].set_data(Bytes::from(String::from("hello"))); + data[1].set_data(Bytes::from(String::from("parquet"))); + let data_bytes = ByteArrayType::to_byte_array(&data[..]); + test_plain_skip::(Bytes::from(data_bytes), 2, 1, -1, &data[1..]); + } + + #[test] + fn test_plain_skip_all_byte_array() { + let mut data = vec![ByteArray::new(); 2]; + data[0].set_data(Bytes::from(String::from("hello"))); + data[1].set_data(Bytes::from(String::from("parquet"))); + let data_bytes = ByteArrayType::to_byte_array(&data[..]); + test_plain_skip::(Bytes::from(data_bytes), 2, 2, -1, &[]); + } + + #[test] + fn test_plain_decode_fixed_len_byte_array() { + let mut data = vec![FixedLenByteArray::default(); 3]; + data[0].set_data(Bytes::from(String::from("bird"))); + data[1].set_data(Bytes::from(String::from("come"))); + data[2].set_data(Bytes::from(String::from("flow"))); + let data_bytes = FixedLenByteArrayType::to_byte_array(&data[..]); + let mut buffer = vec![FixedLenByteArray::default(); 3]; + test_plain_decode::( + Bytes::from(data_bytes), + 3, + 4, + &mut buffer[..], + &data[..], + ); + } + + #[test] + fn test_plain_skip_fixed_len_byte_array() { + let mut data = vec![FixedLenByteArray::default(); 3]; + data[0].set_data(Bytes::from(String::from("bird"))); + data[1].set_data(Bytes::from(String::from("come"))); + data[2].set_data(Bytes::from(String::from("flow"))); + let data_bytes = FixedLenByteArrayType::to_byte_array(&data[..]); + test_plain_skip::(Bytes::from(data_bytes), 3, 1, 4, &data[1..]); + } + + #[test] + fn test_plain_skip_all_fixed_len_byte_array() { + let mut data = vec![FixedLenByteArray::default(); 3]; + data[0].set_data(Bytes::from(String::from("bird"))); + data[1].set_data(Bytes::from(String::from("come"))); + data[2].set_data(Bytes::from(String::from("flow"))); + let data_bytes = FixedLenByteArrayType::to_byte_array(&data[..]); + test_plain_skip::(Bytes::from(data_bytes), 3, 6, 4, &[]); + } + + #[test] + fn test_dict_decoder_empty_data() { + let mut decoder = DictDecoder::::new(); + let err = decoder.set_data(Bytes::new(), 10).unwrap_err(); + assert_eq!(err.to_string(), "EOF: Not enough bytes to decode bit_width"); + } + + fn test_plain_decode( + data: Bytes, + num_values: usize, + type_length: i32, + buffer: &mut [T::T], + expected: &[T::T], + ) { + let mut decoder: PlainDecoder = PlainDecoder::new(type_length); + let result = decoder.set_data(data, num_values); + assert!(result.is_ok()); + let result = decoder.get(buffer); + assert!(result.is_ok()); + assert_eq!(decoder.values_left(), 0); + assert_eq!(buffer, expected); + } + + fn test_plain_skip( + data: Bytes, + num_values: usize, + skip: usize, + type_length: i32, + expected: &[T::T], + ) { + let mut decoder: PlainDecoder = PlainDecoder::new(type_length); + let result = decoder.set_data(data, num_values); + assert!(result.is_ok()); + let skipped = decoder.skip(skip).expect("skipping values"); + + if skip >= num_values { + assert_eq!(skipped, num_values); + + let mut buffer = vec![T::T::default(); 1]; + let remaining = decoder.get(&mut buffer).expect("getting remaining values"); + assert_eq!(remaining, 0); + } else { + assert_eq!(skipped, skip); + let mut buffer = vec![T::T::default(); num_values - skip]; + let remaining = decoder.get(&mut buffer).expect("getting remaining values"); + assert_eq!(remaining, num_values - skip); + assert_eq!(decoder.values_left(), 0); + assert_eq!(buffer, expected); + } + } + + fn test_plain_decode_spaced( + data: Bytes, + num_values: usize, + type_length: i32, + buffer: &mut [T::T], + num_nulls: usize, + valid_bits: &[u8], + expected: &[T::T], + ) { + let mut decoder: PlainDecoder = PlainDecoder::new(type_length); + let result = decoder.set_data(data, num_values); + assert!(result.is_ok()); + let result = decoder.get_spaced(buffer, num_nulls, valid_bits); + assert!(result.is_ok()); + assert_eq!(num_values + num_nulls, result.unwrap()); + assert_eq!(decoder.values_left(), 0); + assert_eq!(buffer, expected); + } + + #[test] + #[should_panic(expected = "RleValueEncoder only supports BoolType")] + fn test_rle_value_encode_int32_not_supported() { + let mut encoder = RleValueEncoder::::new(); + encoder.put(&[1, 2, 3, 4]).unwrap(); + } + + #[test] + #[should_panic(expected = "RleValueDecoder only supports BoolType")] + fn test_rle_value_decode_int32_not_supported() { + let mut decoder = RleValueDecoder::::new(); + decoder.set_data(Bytes::from(vec![5, 0, 0, 0]), 1).unwrap(); + } + + #[test] + fn test_rle_value_decode_missing_size() { + let mut decoder = RleValueDecoder::::new(); + assert!(decoder.set_data(Bytes::from(vec![0]), 1).is_err()); + } + + #[test] + fn test_rle_value_decode_missing_data() { + let mut decoder = RleValueDecoder::::new(); + assert!(decoder.set_data(Bytes::from(vec![5, 0, 0, 0]), 1).is_err()); + } + + #[test] + fn test_rle_value_decode_bool_decode() { + // Test multiple 'put' calls on the same encoder + let data = vec![ + BoolType::gen_vec(-1, 256), + BoolType::gen_vec(-1, 257), + BoolType::gen_vec(-1, 126), + ]; + test_rle_value_decode::(data); + } + + #[test] + #[should_panic(expected = "Bit reader is not initialized")] + fn test_delta_bit_packed_not_initialized_offset() { + // Fail if set_data() is not called before get_offset() + let decoder = DeltaBitPackDecoder::::new(); + decoder.get_offset(); + } + + #[test] + #[should_panic(expected = "Bit reader is not initialized")] + fn test_delta_bit_packed_not_initialized_get() { + // Fail if set_data() is not called before get() + let mut decoder = DeltaBitPackDecoder::::new(); + let mut buffer = vec![]; + decoder.get(&mut buffer).unwrap(); + } + + #[test] + fn test_delta_bit_packed_int32_empty() { + let data = vec![vec![0; 0]]; + test_delta_bit_packed_decode::(data); + } + + #[test] + fn test_delta_bit_packed_int32_repeat() { + let block_data = vec![ + 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, + 6, 7, 8, + ]; + test_delta_bit_packed_decode::(vec![block_data]); + } + + #[test] + fn test_skip_delta_bit_packed_int32_repeat() { + let block_data = vec![ + 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, + 6, 7, 8, + ]; + test_skip::(block_data.clone(), Encoding::DELTA_BINARY_PACKED, 10); + test_skip::(block_data, Encoding::DELTA_BINARY_PACKED, 100); + } + + #[test] + fn test_delta_bit_packed_int32_uneven() { + let block_data = vec![1, -2, 3, -4, 5, 6, 7, 8, 9, 10, 11]; + test_delta_bit_packed_decode::(vec![block_data]); + } + + #[test] + fn test_skip_delta_bit_packed_int32_uneven() { + let block_data = vec![1, -2, 3, -4, 5, 6, 7, 8, 9, 10, 11]; + test_skip::(block_data.clone(), Encoding::DELTA_BINARY_PACKED, 5); + test_skip::(block_data, Encoding::DELTA_BINARY_PACKED, 100); + } + + #[test] + fn test_delta_bit_packed_int32_same_values() { + let block_data = vec![ + 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, + ]; + test_delta_bit_packed_decode::(vec![block_data]); + + let block_data = vec![ + -127, -127, -127, -127, -127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, + ]; + test_delta_bit_packed_decode::(vec![block_data]); + } + + #[test] + fn test_skip_delta_bit_packed_int32_same_values() { + let block_data = vec![ + 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, + ]; + test_skip::(block_data.clone(), Encoding::DELTA_BINARY_PACKED, 5); + test_skip::(block_data, Encoding::DELTA_BINARY_PACKED, 100); + + let block_data = vec![ + -127, -127, -127, -127, -127, -127, -127, -127, -127, -127, -127, -127, -127, -127, + -127, -127, + ]; + test_skip::(block_data.clone(), Encoding::DELTA_BINARY_PACKED, 5); + test_skip::(block_data, Encoding::DELTA_BINARY_PACKED, 100); + } + + #[test] + fn test_delta_bit_packed_int32_min_max() { + let block_data = vec![ + i32::MIN, + i32::MIN, + i32::MIN, + i32::MAX, + i32::MIN, + i32::MAX, + i32::MIN, + i32::MAX, + ]; + test_delta_bit_packed_decode::(vec![block_data]); + } + + #[test] + fn test_skip_delta_bit_packed_int32_min_max() { + let block_data = vec![ + i32::MIN, + i32::MIN, + i32::MIN, + i32::MAX, + i32::MIN, + i32::MAX, + i32::MIN, + i32::MAX, + ]; + test_skip::(block_data.clone(), Encoding::DELTA_BINARY_PACKED, 5); + test_skip::(block_data, Encoding::DELTA_BINARY_PACKED, 100); + } + + #[test] + fn test_skip_delta_bit_packed_bw0_uniform_step_i32() { + // Uniform-step column: every delta equals min_delta, so bw=0 miniblocks. + // Partial skip must advance last_value by n * min_delta (min_delta != 0 path). + let data: Vec = (0..128).map(|i| i * 7).collect(); + test_skip::(data.clone(), Encoding::DELTA_BINARY_PACKED, 50); + test_skip::(data, Encoding::DELTA_BINARY_PACKED, 200); + } + + #[test] + fn test_skip_delta_bit_packed_bw0_uniform_step_i64() { + // Same as above for i64. + let data: Vec = (0..128).map(|i| i * 100).collect(); + test_skip::(data.clone(), Encoding::DELTA_BINARY_PACKED, 50); + test_skip::(data, Encoding::DELTA_BINARY_PACKED, 200); + } + + #[test] + fn test_delta_bit_packed_int32_multiple_blocks() { + // Test multiple 'put' calls on the same encoder + let data = vec![ + Int32Type::gen_vec(-1, 64), + Int32Type::gen_vec(-1, 128), + Int32Type::gen_vec(-1, 64), + ]; + test_delta_bit_packed_decode::(data); + } + + #[test] + fn test_delta_bit_packed_int32_data_across_blocks() { + // Test multiple 'put' calls on the same encoder + let data = vec![Int32Type::gen_vec(-1, 256), Int32Type::gen_vec(-1, 257)]; + test_delta_bit_packed_decode::(data); + } + + #[test] + fn test_delta_bit_packed_int32_with_empty_blocks() { + let data = vec![ + Int32Type::gen_vec(-1, 128), + vec![0; 0], + Int32Type::gen_vec(-1, 64), + ]; + test_delta_bit_packed_decode::(data); + } + + #[test] + fn test_delta_bit_packed_int64_empty() { + let data = vec![vec![0; 0]]; + test_delta_bit_packed_decode::(data); + } + + #[test] + fn test_delta_bit_packed_int64_min_max() { + let block_data = vec![ + i64::MIN, + i64::MAX, + i64::MIN, + i64::MAX, + i64::MIN, + i64::MAX, + i64::MIN, + i64::MAX, + ]; + test_delta_bit_packed_decode::(vec![block_data]); + } + + #[test] + fn test_delta_bit_packed_int64_multiple_blocks() { + // Test multiple 'put' calls on the same encoder + let data = vec![ + Int64Type::gen_vec(-1, 64), + Int64Type::gen_vec(-1, 128), + Int64Type::gen_vec(-1, 64), + ]; + test_delta_bit_packed_decode::(data); + } + + #[test] + fn test_delta_bit_packed_zero_miniblocks() { + // It is invalid for mini_blocks_per_block to be 0 + let data = vec![ + 128, 1, // block_size = 128 + 0, // mini_blocks_per_block = 0 + ]; + let mut decoder = DeltaBitPackDecoder::::new(); + let err = decoder.set_data(data.into(), 0).unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: cannot have zero miniblocks per block" + ); + } + + #[test] + fn test_delta_bit_packed_decoder_sample() { + let data_bytes = vec![ + 128, 1, 4, 3, 58, 28, 6, 0, 0, 0, 0, 8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, + ]; + let mut decoder: DeltaBitPackDecoder = DeltaBitPackDecoder::new(); + decoder.set_data(data_bytes.into(), 3).unwrap(); + // check exact offsets, because when reading partial values we end up with + // some data not being read from bit reader + assert_eq!(decoder.get_offset(), 5); + let mut result = vec![0, 0, 0]; + decoder.get(&mut result).unwrap(); + assert_eq!(decoder.get_offset(), 34); + assert_eq!(result, vec![29, 43, 89]); + } + + #[test] + fn test_delta_bit_packed_padding() { + // Page header + let header = vec![ + // Page Header + + // Block Size - 256 + 128, + 2, + // Miniblocks in block, + 4, + // Total value count - 419 + 128 + 35, + 3, + // First value - 7 + 7, + ]; + + // Block Header + let block1_header = vec![ + 0, // Min delta + 0, 1, 0, 0, // Bit widths + ]; + + // Mini-block 1 - bit width 0 => 0 bytes + // Mini-block 2 - bit width 1 => 8 bytes + // Mini-block 3 - bit width 0 => 0 bytes + // Mini-block 4 - bit width 0 => 0 bytes + let block1 = vec![0xFF; 8]; + + // Block Header + let block2_header = vec![ + 0, // Min delta + 0, 1, 2, 0xFF, // Bit widths, including non-zero padding + ]; + + // Mini-block 1 - bit width 0 => 0 bytes + // Mini-block 2 - bit width 1 => 8 bytes + // Mini-block 3 - bit width 2 => 16 bytes + // Mini-block 4 - padding => no bytes + let block2 = vec![0xFF; 24]; + + let data: Vec = header + .into_iter() + .chain(block1_header) + .chain(block1) + .chain(block2_header) + .chain(block2) + .collect(); + + let length = data.len(); + + let ptr = Bytes::from(data); + let mut reader = BitReader::new(ptr.clone()); + assert_eq!(reader.get_vlq_int().unwrap(), 256); + assert_eq!(reader.get_vlq_int().unwrap(), 4); + assert_eq!(reader.get_vlq_int().unwrap(), 419); + assert_eq!(reader.get_vlq_int().unwrap(), 7); + + // Test output buffer larger than needed and not exact multiple of block size + let mut output = vec![0_i32; 420]; + + let mut decoder = DeltaBitPackDecoder::::new(); + decoder.set_data(ptr.clone(), 0).unwrap(); + assert_eq!(decoder.get(&mut output).unwrap(), 419); + assert_eq!(decoder.get_offset(), length); + + // Test with truncated buffer + decoder.set_data(ptr.slice(..12), 0).unwrap(); + let err = decoder.get(&mut output).unwrap_err().to_string(); + assert!( + err.contains("Expected to read 64 values from miniblock got 8"), + "{}", + err + ); + } + + #[test] + fn test_delta_bit_packed_int32_single_value_large() { + let block_data = vec![3; 10240]; + test_delta_bit_packed_decode::(vec![block_data]); + } + + #[test] + fn test_delta_bit_packed_int32_single_value_skip_large() { + let block_data = vec![3; 10240]; + test_skip::(block_data.clone(), Encoding::DELTA_BINARY_PACKED, 50); + test_skip::(block_data, Encoding::DELTA_BINARY_PACKED, 5000); + } + + #[test] + fn test_delta_bit_packed_int32_increasing_value_large() { + let block_data = (0i32..10240).collect(); + test_delta_bit_packed_decode::(vec![block_data]); + } + + #[test] + fn test_delta_bit_packed_int32_increasing_value_skip_large() { + let block_data = (0i32..10240).collect::>(); + test_skip::(block_data.clone(), Encoding::DELTA_BINARY_PACKED, 50); + test_skip::(block_data, Encoding::DELTA_BINARY_PACKED, 5000); + } + + #[test] + fn test_delta_bit_packed_int32_stepped_value_large() { + let block_data = (0i32..10240).map(|i| i / 2).collect(); + test_delta_bit_packed_decode::(vec![block_data]); + } + + #[test] + fn test_delta_bit_packed_int32_stepped_value_skip_large() { + let block_data = (0i32..10240).map(|i| i / 2).collect::>(); + test_skip::(block_data.clone(), Encoding::DELTA_BINARY_PACKED, 50); + test_skip::(block_data, Encoding::DELTA_BINARY_PACKED, 5000); + } + + #[test] + fn test_delta_bit_packed_int32_mixed_large() { + // should be enough for 4 mini-blocks plus a little so we get some + // mixed mini-blocks + const BLOCK_SIZE: i32 = 133; + let block1_data = (0..BLOCK_SIZE).map(|i| (i * 7) % 11).collect(); + let block2_data = vec![3; BLOCK_SIZE as usize]; + let block3_data = (0..BLOCK_SIZE).map(|i| (i * 5) % 13).collect(); + let block4_data = (0..BLOCK_SIZE).collect(); + let block5_data = (0..BLOCK_SIZE).map(|i| (i * 3) % 17).collect(); + test_delta_bit_packed_decode::(vec![ + block1_data, + block2_data, + block3_data, + block4_data, + block5_data, + ]); + } + + #[test] + fn test_delta_bit_packed_int64_single_value_large() { + let block_data = vec![5; 10240]; + test_delta_bit_packed_decode::(vec![block_data]); + } + + #[test] + fn test_delta_bit_packed_int64_increasing_value_large() { + let block_data = (0i64..10240).collect(); + test_delta_bit_packed_decode::(vec![block_data]); + } + + #[test] + fn test_delta_byte_array_same_arrays() { + let data = vec![ + vec![ByteArray::from(vec![1, 2, 3, 4, 5, 6])], + vec![ + ByteArray::from(vec![1, 2, 3, 4, 5, 6]), + ByteArray::from(vec![1, 2, 3, 4, 5, 6]), + ], + vec![ + ByteArray::from(vec![1, 2, 3, 4, 5, 6]), + ByteArray::from(vec![1, 2, 3, 4, 5, 6]), + ], + ]; + test_delta_byte_array_decode(data); + } + + #[test] + fn test_delta_byte_array_unique_arrays() { + let data = vec![ + vec![ByteArray::from(vec![1])], + vec![ByteArray::from(vec![2, 3]), ByteArray::from(vec![4, 5, 6])], + vec![ + ByteArray::from(vec![7, 8]), + ByteArray::from(vec![9, 0, 1, 2]), + ], + ]; + test_delta_byte_array_decode(data); + } + + #[test] + fn test_delta_byte_array_single_array() { + let data = vec![vec![ByteArray::from(vec![1, 2, 3, 4, 5, 6])]]; + test_delta_byte_array_decode(data); + } + + #[test] + fn test_byte_stream_split_multiple_f32() { + let data = vec![ + vec![ + f32::from_le_bytes([0xAA, 0xBB, 0xCC, 0xDD]), + f32::from_le_bytes([0x00, 0x11, 0x22, 0x33]), + ], + vec![f32::from_le_bytes([0xA3, 0xB4, 0xC5, 0xD6])], + ]; + test_byte_stream_split_decode::(data, -1); + } + + #[test] + fn test_byte_stream_split_f64() { + let data = vec![vec![ + f64::from_le_bytes([0, 1, 2, 3, 4, 5, 6, 7]), + f64::from_le_bytes([8, 9, 10, 11, 12, 13, 14, 15]), + ]]; + test_byte_stream_split_decode::(data, -1); + } + + #[test] + fn test_byte_stream_split_multiple_i32() { + let data = vec![ + vec![ + i32::from_le_bytes([0xAA, 0xBB, 0xCC, 0xDD]), + i32::from_le_bytes([0x00, 0x11, 0x22, 0x33]), + ], + vec![i32::from_le_bytes([0xA3, 0xB4, 0xC5, 0xD6])], + ]; + test_byte_stream_split_decode::(data, -1); + } + + #[test] + fn test_byte_stream_split_i64() { + let data = vec![vec![ + i64::from_le_bytes([0, 1, 2, 3, 4, 5, 6, 7]), + i64::from_le_bytes([8, 9, 10, 11, 12, 13, 14, 15]), + ]]; + test_byte_stream_split_decode::(data, -1); + } + + fn test_byte_stream_split_flba(type_width: usize) { + let data = vec![ + vec![ + FixedLenByteArrayType::r#gen(type_width as i32), + FixedLenByteArrayType::r#gen(type_width as i32), + ], + vec![FixedLenByteArrayType::r#gen(type_width as i32)], + ]; + test_byte_stream_split_decode::(data, type_width as i32); + } + + #[test] + fn test_byte_stream_split_flba5() { + test_byte_stream_split_flba(5); + } + + #[test] + fn test_byte_stream_split_flba16() { + test_byte_stream_split_flba(16); + } + + #[test] + fn test_byte_stream_split_flba19() { + test_byte_stream_split_flba(19); + } + + #[test] + #[should_panic(expected = "Mismatched FixedLenByteArray sizes: 4 != 5")] + fn test_byte_stream_split_flba_mismatch() { + let data = vec![ + vec![ + FixedLenByteArray::from(vec![0xAA, 0xAB, 0xAC, 0xAD, 0xAE]), + FixedLenByteArray::from(vec![0xBA, 0xBB, 0xBC, 0xBD, 0xBE]), + ], + vec![FixedLenByteArray::from(vec![0xCA, 0xCB, 0xCC, 0xCD])], + ]; + test_byte_stream_split_decode::(data, 5); + } + + #[test] + #[should_panic(expected = "Input data length is not a multiple of type width 4")] + fn test_byte_stream_split_flba_bad_input() { + let mut decoder = VariableWidthByteStreamSplitDecoder::::new(4); + decoder + .set_data(Bytes::from(vec![1, 2, 3, 4, 5]), 1) + .unwrap(); + } + + #[test] + fn test_skip_byte_stream_split() { + let block_data = vec![0.3, 0.4, 0.1, 4.10]; + test_skip::(block_data.clone(), Encoding::BYTE_STREAM_SPLIT, 2); + test_skip::( + block_data.into_iter().map(|x| x as f64).collect(), + Encoding::BYTE_STREAM_SPLIT, + 100, + ); + } + + #[test] + fn test_skip_byte_stream_split_ints() { + let block_data = vec![3, 4, 1, 5]; + test_skip::(block_data.clone(), Encoding::BYTE_STREAM_SPLIT, 2); + test_skip::( + block_data.into_iter().map(|x| x as i64).collect(), + Encoding::BYTE_STREAM_SPLIT, + 100, + ); + } + + fn test_rle_value_decode(data: Vec>) { + test_encode_decode::(data, Encoding::RLE, -1); + } + + fn test_delta_bit_packed_decode(data: Vec>) { + test_encode_decode::(data, Encoding::DELTA_BINARY_PACKED, -1); + } + + fn test_byte_stream_split_decode(data: Vec>, type_width: i32) { + test_encode_decode::(data, Encoding::BYTE_STREAM_SPLIT, type_width); + } + + fn test_delta_byte_array_decode(data: Vec>) { + test_encode_decode::(data, Encoding::DELTA_BYTE_ARRAY, -1); + } + + // Input data represents vector of data slices to write (test multiple `put()` calls) + // For example, + // vec![vec![1, 2, 3]] invokes `put()` once and writes {1, 2, 3} + // vec![vec![1, 2], vec![3]] invokes `put()` twice and writes {1, 2, 3} + fn test_encode_decode(data: Vec>, encoding: Encoding, type_width: i32) { + let col_descr = create_test_col_desc_ptr(type_width, T::get_physical_type()); + + // Encode data + let mut encoder = get_encoder::(encoding, &col_descr).expect("get encoder"); + + for v in &data[..] { + encoder.put(&v[..]).expect("ok to encode"); + } + let bytes = encoder.flush_buffer().expect("ok to flush buffer"); + + // Flatten expected data as contiguous array of values + let expected: Vec = data.iter().flat_map(|s| s.clone()).collect(); + + // Decode data and compare with original + let mut decoder = get_decoder::(col_descr, encoding).expect("get decoder"); + + let mut result = vec![T::T::default(); expected.len()]; + decoder + .set_data(bytes, expected.len()) + .expect("ok to set data"); + let mut result_num_values = 0; + while decoder.values_left() > 0 { + result_num_values += decoder + .get(&mut result[result_num_values..]) + .expect("ok to decode"); + } + assert_eq!(result_num_values, expected.len()); + assert_eq!(result, expected); + } + + fn test_skip(data: Vec, encoding: Encoding, skip: usize) { + // Type length should not really matter for encode/decode test, + // otherwise change it based on type + let col_descr = create_test_col_desc_ptr(-1, T::get_physical_type()); + + // Encode data + let mut encoder = get_encoder::(encoding, &col_descr).expect("get encoder"); + + encoder.put(&data).expect("ok to encode"); + + let bytes = encoder.flush_buffer().expect("ok to flush buffer"); + + let mut decoder = get_decoder::(col_descr, encoding).expect("get decoder"); + decoder.set_data(bytes, data.len()).expect("ok to set data"); + + if skip >= data.len() { + let skipped = decoder.skip(skip).expect("ok to skip"); + assert_eq!(skipped, data.len()); + + let skipped_again = decoder.skip(skip).expect("ok to skip again"); + assert_eq!(skipped_again, 0); + } else { + let skipped = decoder.skip(skip).expect("ok to skip"); + assert_eq!(skipped, skip); + + let remaining = data.len() - skip; + + let expected = &data[skip..]; + let mut buffer = vec![T::T::default(); remaining]; + let fetched = decoder.get(&mut buffer).expect("ok to decode"); + assert_eq!(remaining, fetched); + assert_eq!(&buffer, expected); + } + } + + fn create_and_check_decoder(encoding: Encoding, err: Option) { + let descr = create_test_col_desc_ptr(-1, T::get_physical_type()); + let decoder = get_decoder::(descr, encoding); + match err { + Some(parquet_error) => { + assert_eq!( + decoder.err().unwrap().to_string(), + parquet_error.to_string() + ); + } + None => { + assert_eq!(decoder.unwrap().encoding(), encoding); + } + } + } + + // Creates test column descriptor. + fn create_test_col_desc_ptr(type_len: i32, t: Type) -> ColumnDescPtr { + let ty = SchemaType::primitive_type_builder("t", t) + .with_length(type_len) + .build() + .unwrap(); + Arc::new(ColumnDescriptor::new( + Arc::new(ty), + 0, + 0, + ColumnPath::new(vec![]), + )) + } + + fn usize_to_bytes(v: usize) -> [u8; 4] { + (v as u32).to_ne_bytes() + } + + /// A util trait to convert slices of different types to byte arrays + trait ToByteArray { + #[allow(clippy::wrong_self_convention)] + fn to_byte_array(data: &[T::T]) -> Vec; + } + + macro_rules! to_byte_array_impl { + ($ty: ty) => { + impl ToByteArray<$ty> for $ty { + #[allow(clippy::wrong_self_convention)] + fn to_byte_array(data: &[<$ty as DataType>::T]) -> Vec { + <$ty as DataType>::T::slice_as_bytes(data).to_vec() + } + } + }; + } + + to_byte_array_impl!(Int32Type); + to_byte_array_impl!(Int64Type); + to_byte_array_impl!(FloatType); + to_byte_array_impl!(DoubleType); + + impl ToByteArray for BoolType { + #[allow(clippy::wrong_self_convention)] + fn to_byte_array(data: &[bool]) -> Vec { + let mut v = vec![]; + for (i, item) in data.iter().enumerate() { + if i % 8 == 0 { + v.push(0); + } + if *item { + v[i / 8] |= 1 << (i % 8); + } + } + v + } + } + + impl ToByteArray for Int96Type { + #[allow(clippy::wrong_self_convention)] + fn to_byte_array(data: &[Int96]) -> Vec { + let mut v = vec![]; + for d in data { + v.extend_from_slice(d.as_bytes()); + } + v + } + } + + impl ToByteArray for ByteArrayType { + #[allow(clippy::wrong_self_convention)] + fn to_byte_array(data: &[ByteArray]) -> Vec { + let mut v = vec![]; + for d in data { + let buf = d.data(); + let len = &usize_to_bytes(buf.len()); + v.extend_from_slice(len); + v.extend(buf); + } + v + } + } + + impl ToByteArray for FixedLenByteArrayType { + #[allow(clippy::wrong_self_convention)] + fn to_byte_array(data: &[FixedLenByteArray]) -> Vec { + let mut v = vec![]; + for d in data { + let buf = d.data(); + v.extend(buf); + } + v + } + } + + #[test] + // Allow initializing a vector and pushing to it for clarity in this test + #[allow(clippy::vec_init_then_push)] + fn test_delta_bit_packed_invalid_bit_width() { + // Manually craft a buffer with an invalid bit width + let mut buffer = vec![]; + // block_size = 128 + buffer.push(128); + buffer.push(1); + // mini_blocks_per_block = 4 + buffer.push(4); + // num_values = 32 + buffer.push(32); + // first_value = 0 + buffer.push(0); + // min_delta = 0 + buffer.push(0); + // bit_widths, one for each of the 4 mini blocks + buffer.push(33); // Invalid bit width + buffer.push(0); + buffer.push(0); + buffer.push(0); + + let corrupted_buffer = Bytes::from(buffer); + + let mut decoder = DeltaBitPackDecoder::::new(); + decoder.set_data(corrupted_buffer.clone(), 32).unwrap(); + let mut read_buffer = vec![0; 32]; + let err = decoder.get(&mut read_buffer).unwrap_err(); + assert!( + err.to_string() + .contains("Invalid delta bit width 33 which is larger than expected 32"), + "{}", + err + ); + + let mut decoder = DeltaBitPackDecoder::::new(); + decoder.set_data(corrupted_buffer, 32).unwrap(); + let err = decoder.skip(32).unwrap_err(); + assert!( + err.to_string() + .contains("Invalid delta bit width 33 which is larger than expected 32"), + "{}", + err + ); + } +} diff --git a/vendor/parquet/src/encodings/decoding/byte_stream_split_decoder.rs b/vendor/parquet/src/encodings/decoding/byte_stream_split_decoder.rs new file mode 100644 index 0000000..b72ae8f --- /dev/null +++ b/vendor/parquet/src/encodings/decoding/byte_stream_split_decoder.rs @@ -0,0 +1,256 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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::marker::PhantomData; + +use bytes::Bytes; + +use crate::basic::{Encoding, Type}; +use crate::data_type::private::ParquetValueType; +use crate::data_type::{DataType, SliceAsBytes}; +use crate::errors::{ParquetError, Result}; + +use super::Decoder; + +pub struct ByteStreamSplitDecoder { + _phantom: PhantomData, + encoded_bytes: Bytes, + total_num_values: usize, + values_decoded: usize, +} + +impl ByteStreamSplitDecoder { + pub(crate) fn new() -> Self { + Self { + _phantom: PhantomData, + encoded_bytes: Bytes::new(), + total_num_values: 0, + values_decoded: 0, + } + } +} + +// Here we assume src contains the full data (which it must, since we're +// can only know where to split the streams once all data is collected), +// but dst can be just a slice starting from the given index. +// We iterate over the output bytes and fill them in from their strided +// input byte locations. +fn join_streams_const( + src: &[u8], + dst: &mut [u8], + stride: usize, + values_decoded: usize, +) { + let sub_src = &src[values_decoded..]; + for i in 0..dst.len() / TYPE_SIZE { + for j in 0..TYPE_SIZE { + dst[i * TYPE_SIZE + j] = sub_src[i + j * stride]; + } + } +} + +// Like the above, but type_size is not known at compile time. +fn join_streams_variable( + src: &[u8], + dst: &mut [u8], + stride: usize, + type_size: usize, + values_decoded: usize, +) { + let sub_src = &src[values_decoded..]; + for i in 0..dst.len() / type_size { + for j in 0..type_size { + dst[i * type_size + j] = sub_src[i + j * stride]; + } + } +} + +impl Decoder for ByteStreamSplitDecoder { + fn set_data(&mut self, data: Bytes, num_values: usize) -> Result<()> { + self.encoded_bytes = data; + self.total_num_values = num_values; + self.values_decoded = 0; + + Ok(()) + } + + fn get(&mut self, buffer: &mut [::T]) -> Result { + let total_remaining_values = self.values_left(); + let num_values = buffer.len().min(total_remaining_values); + let buffer = &mut buffer[..num_values]; + + // SAFETY: i/f32 and i/f64 has no constraints on their internal representation, so we can modify it as we want + let raw_out_bytes = unsafe { ::T::slice_as_bytes_mut(buffer) }; + let type_size = T::get_type_size(); + let stride = self.encoded_bytes.len() / type_size; + match type_size { + 4 => join_streams_const::<4>( + &self.encoded_bytes, + raw_out_bytes, + stride, + self.values_decoded, + ), + 8 => join_streams_const::<8>( + &self.encoded_bytes, + raw_out_bytes, + stride, + self.values_decoded, + ), + _ => { + return Err(general_err!( + "byte stream split unsupported for data types of size {} bytes", + type_size + )); + } + } + self.values_decoded += num_values; + + Ok(num_values) + } + + fn values_left(&self) -> usize { + self.total_num_values - self.values_decoded + } + + fn encoding(&self) -> Encoding { + Encoding::BYTE_STREAM_SPLIT + } + + fn skip(&mut self, num_values: usize) -> Result { + let to_skip = usize::min(self.values_left(), num_values); + self.values_decoded += to_skip; + Ok(to_skip) + } +} + +pub struct VariableWidthByteStreamSplitDecoder { + _phantom: PhantomData, + encoded_bytes: Bytes, + total_num_values: usize, + values_decoded: usize, + type_width: usize, +} + +impl VariableWidthByteStreamSplitDecoder { + pub(crate) fn new(type_length: i32) -> Self { + Self { + _phantom: PhantomData, + encoded_bytes: Bytes::new(), + total_num_values: 0, + values_decoded: 0, + type_width: type_length as usize, + } + } +} + +impl Decoder for VariableWidthByteStreamSplitDecoder { + fn set_data(&mut self, data: Bytes, num_values: usize) -> Result<()> { + // Rough check that all data elements are the same length + if data.len() % self.type_width != 0 { + return Err(general_err!( + "Input data length is not a multiple of type width {}", + self.type_width + )); + } + + match T::get_physical_type() { + Type::FIXED_LEN_BYTE_ARRAY => { + self.encoded_bytes = data; + self.total_num_values = num_values; + self.values_decoded = 0; + Ok(()) + } + _ => Err(general_err!( + "VariableWidthByteStreamSplitDecoder only supports FixedLenByteArrayType" + )), + } + } + + fn get(&mut self, buffer: &mut [::T]) -> Result { + let total_remaining_values = self.values_left(); + let num_values = buffer.len().min(total_remaining_values); + let buffer = &mut buffer[..num_values]; + let type_size = self.type_width; + + // Since this is FIXED_LEN_BYTE_ARRAY data, we can't use slice_as_bytes_mut. Instead we'll + // have to do some data copies. + let mut tmp_vec = vec![0_u8; num_values * type_size]; + let raw_out_bytes = tmp_vec.as_mut_slice(); + + let stride = self.encoded_bytes.len() / type_size; + match type_size { + 2 => join_streams_const::<2>( + &self.encoded_bytes, + raw_out_bytes, + stride, + self.values_decoded, + ), + 4 => join_streams_const::<4>( + &self.encoded_bytes, + raw_out_bytes, + stride, + self.values_decoded, + ), + 8 => join_streams_const::<8>( + &self.encoded_bytes, + raw_out_bytes, + stride, + self.values_decoded, + ), + 16 => join_streams_const::<16>( + &self.encoded_bytes, + raw_out_bytes, + stride, + self.values_decoded, + ), + _ => join_streams_variable( + &self.encoded_bytes, + raw_out_bytes, + stride, + type_size, + self.values_decoded, + ), + } + self.values_decoded += num_values; + + // create a buffer from the vec so far (and leave a new Vec in its place) + let vec_with_data = std::mem::take(&mut tmp_vec); + // convert Vec to Bytes (which is a ref counted wrapper) + let bytes_with_data = Bytes::from(vec_with_data); + for (i, bi) in buffer.iter_mut().enumerate().take(num_values) { + // Get a view into the data, without also copying the bytes + let data = bytes_with_data.slice(i * type_size..(i + 1) * type_size); + bi.set_from_bytes(data); + } + + Ok(num_values) + } + + fn values_left(&self) -> usize { + self.total_num_values - self.values_decoded + } + + fn encoding(&self) -> Encoding { + Encoding::BYTE_STREAM_SPLIT + } + + fn skip(&mut self, num_values: usize) -> Result { + let to_skip = usize::min(self.values_left(), num_values); + self.values_decoded += to_skip; + Ok(to_skip) + } +} diff --git a/vendor/parquet/src/encodings/encoding/byte_stream_split_encoder.rs b/vendor/parquet/src/encodings/encoding/byte_stream_split_encoder.rs new file mode 100644 index 0000000..0726c6b --- /dev/null +++ b/vendor/parquet/src/encodings/encoding/byte_stream_split_encoder.rs @@ -0,0 +1,231 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::basic::{Encoding, Type}; +use crate::data_type::{AsBytes, DataType, SliceAsBytes}; + +use crate::errors::{ParquetError, Result}; + +use super::Encoder; + +use bytes::{BufMut, Bytes}; +use std::cmp; +use std::marker::PhantomData; + +pub struct ByteStreamSplitEncoder { + buffer: Vec, + _p: PhantomData, +} + +impl ByteStreamSplitEncoder { + pub(crate) fn new() -> Self { + Self { + buffer: Vec::new(), + _p: PhantomData, + } + } +} + +// Here we assume src contains the full data (which it must, since we +// can only know where to split the streams once all data is collected). +// We iterate over the input bytes and write them to their strided output +// byte locations. +fn split_streams_const(src: &[u8], dst: &mut [u8]) { + let stride = src.len() / TYPE_SIZE; + for i in 0..stride { + for j in 0..TYPE_SIZE { + dst[i + j * stride] = src[i * TYPE_SIZE + j]; + } + } +} + +// Like above, but type_size is not known at compile time. +fn split_streams_variable(src: &[u8], dst: &mut [u8], type_size: usize) { + const BLOCK_SIZE: usize = 4; + let stride = src.len() / type_size; + for j in (0..type_size).step_by(BLOCK_SIZE) { + let jrange = cmp::min(BLOCK_SIZE, type_size - j); + for i in 0..stride { + for jj in 0..jrange { + dst[i + (j + jj) * stride] = src[i * type_size + j + jj]; + } + } + } +} + +impl Encoder for ByteStreamSplitEncoder { + fn put(&mut self, values: &[T::T]) -> Result<()> { + self.buffer + .extend(::T::slice_as_bytes(values)); + + ensure_phys_ty!( + Type::FLOAT | Type::DOUBLE | Type::INT32 | Type::INT64, + "ByteStreamSplitEncoder does not support Int96, Boolean, or ByteArray types" + ); + + Ok(()) + } + + fn encoding(&self) -> Encoding { + Encoding::BYTE_STREAM_SPLIT + } + + fn estimated_data_encoded_size(&self) -> usize { + self.buffer.len() + } + + fn flush_buffer(&mut self) -> Result { + let mut encoded = vec![0; self.buffer.len()]; + let type_size = T::get_type_size(); + match type_size { + 4 => split_streams_const::<4>(&self.buffer, &mut encoded), + 8 => split_streams_const::<8>(&self.buffer, &mut encoded), + _ => { + return Err(general_err!( + "byte stream split unsupported for data types of size {} bytes", + type_size + )); + } + } + + self.buffer.clear(); + Ok(encoded.into()) + } + + /// return the estimated memory size of this encoder. + fn estimated_memory_size(&self) -> usize { + self.buffer.capacity() * std::mem::size_of::() + } +} + +pub struct VariableWidthByteStreamSplitEncoder { + buffer: Vec, + type_width: usize, + _p: PhantomData, +} + +impl VariableWidthByteStreamSplitEncoder { + pub(crate) fn new(type_length: i32) -> Self { + Self { + buffer: Vec::new(), + type_width: type_length as usize, + _p: PhantomData, + } + } +} + +fn put_fixed(dst: &mut [u8], values: &[T::T]) { + let mut idx = 0; + values.iter().for_each(|x| { + let bytes = x.as_bytes(); + if bytes.len() != TYPE_SIZE { + panic!( + "Mismatched FixedLenByteArray sizes: {} != {}", + bytes.len(), + TYPE_SIZE + ); + } + dst[idx..(TYPE_SIZE + idx)].copy_from_slice(&bytes[..TYPE_SIZE]); + idx += TYPE_SIZE; + }); +} + +fn put_variable(dst: &mut [u8], values: &[T::T], type_width: usize) { + let mut idx = 0; + values.iter().for_each(|x| { + let bytes = x.as_bytes(); + if bytes.len() != type_width { + panic!( + "Mismatched FixedLenByteArray sizes: {} != {}", + bytes.len(), + type_width + ); + } + dst[idx..idx + type_width].copy_from_slice(bytes); + idx += type_width; + }); +} + +impl Encoder for VariableWidthByteStreamSplitEncoder { + fn put(&mut self, values: &[T::T]) -> Result<()> { + ensure_phys_ty!( + Type::FIXED_LEN_BYTE_ARRAY, + "VariableWidthByteStreamSplitEncoder only supports FixedLenByteArray types" + ); + + // FixedLenByteArray is implemented as ByteArray, so there may be gaps making + // slice_as_bytes untenable + let idx = self.buffer.len(); + let data_len = values.len() * self.type_width; + // Ensure enough capacity for the new data + self.buffer.reserve(values.len() * self.type_width); + // ...and extend the size of buffer to allow direct access + self.buffer.put_bytes(0_u8, data_len); + // Get a slice of the buffer corresponding to the location of the new data + let out_buf = &mut self.buffer[idx..idx + data_len]; + + // Now copy `values` into the buffer. For `type_width` <= 8 use a fixed size when + // performing the copy as it is significantly faster. + match self.type_width { + 2 => put_fixed::(out_buf, values), + 3 => put_fixed::(out_buf, values), + 4 => put_fixed::(out_buf, values), + 5 => put_fixed::(out_buf, values), + 6 => put_fixed::(out_buf, values), + 7 => put_fixed::(out_buf, values), + 8 => put_fixed::(out_buf, values), + _ => put_variable::(out_buf, values, self.type_width), + } + + Ok(()) + } + + fn encoding(&self) -> Encoding { + Encoding::BYTE_STREAM_SPLIT + } + + fn estimated_data_encoded_size(&self) -> usize { + self.buffer.len() + } + + fn flush_buffer(&mut self) -> Result { + let mut encoded = vec![0; self.buffer.len()]; + let type_size = match T::get_physical_type() { + Type::FIXED_LEN_BYTE_ARRAY => self.type_width, + _ => T::get_type_size(), + }; + // split_streams_const() is faster up to type_width == 8 + match type_size { + 2 => split_streams_const::<2>(&self.buffer, &mut encoded), + 3 => split_streams_const::<3>(&self.buffer, &mut encoded), + 4 => split_streams_const::<4>(&self.buffer, &mut encoded), + 5 => split_streams_const::<5>(&self.buffer, &mut encoded), + 6 => split_streams_const::<6>(&self.buffer, &mut encoded), + 7 => split_streams_const::<7>(&self.buffer, &mut encoded), + 8 => split_streams_const::<8>(&self.buffer, &mut encoded), + _ => split_streams_variable(&self.buffer, &mut encoded, type_size), + } + + self.buffer.clear(); + Ok(encoded.into()) + } + + /// return the estimated memory size of this encoder. + fn estimated_memory_size(&self) -> usize { + self.buffer.capacity() * std::mem::size_of::() + } +} diff --git a/vendor/parquet/src/encodings/encoding/dict_encoder.rs b/vendor/parquet/src/encodings/encoding/dict_encoder.rs new file mode 100644 index 0000000..89666bb --- /dev/null +++ b/vendor/parquet/src/encodings/encoding/dict_encoder.rs @@ -0,0 +1,468 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// ---------------------------------------------------------------------- +// Dictionary encoding + +use bytes::Bytes; + +use crate::basic::{Encoding, Type}; +use crate::data_type::DataType; +use crate::data_type::private::ParquetValueType; +use crate::encodings::encoding::{Encoder, PlainEncoder}; +use crate::encodings::rle::RleEncoder; +use crate::errors::Result; +use crate::schema::types::ColumnDescPtr; +use crate::util::bit_util::num_required_bits; +use crate::util::interner::{Interner, Storage}; + +#[derive(Debug)] +struct KeyStorage { + uniques: Vec, + + /// size of unique values (keys) in the dictionary, in bytes. + size_in_bytes: usize, + + type_length: usize, +} + +impl Storage for KeyStorage { + type Key = u64; + type Value = T::T; + + fn get(&self, idx: Self::Key) -> &Self::Value { + &self.uniques[idx as usize] + } + + fn push(&mut self, value: &Self::Value) -> Self::Key { + let (base_size, num_elements) = value.dict_encoding_size(); + + let unique_size = match T::get_physical_type() { + Type::BYTE_ARRAY => base_size + num_elements, + Type::FIXED_LEN_BYTE_ARRAY => self.type_length, + _ => base_size, + }; + self.size_in_bytes += unique_size; + + let key = self.uniques.len() as u64; + self.uniques.push(value.clone()); + key + } + + fn estimated_memory_size(&self) -> usize { + let uniques_heap_bytes = match T::get_physical_type() { + Type::FIXED_LEN_BYTE_ARRAY => self.type_length * self.uniques.len(), + _ => ::variable_length_bytes(&self.uniques) + .unwrap_or(0) as usize, + }; + self.uniques.capacity() * std::mem::size_of::() + uniques_heap_bytes + } +} + +/// Dictionary encoder. +/// The dictionary encoding builds a dictionary of values encountered in a given column. +/// The dictionary page is written first, before the data pages of the column chunk. +/// +/// Dictionary page format: the entries in the dictionary - in dictionary order - +/// using the plain encoding. +/// +/// Data page format: the bit width used to encode the entry ids stored as 1 byte +/// (max bit width = 32), followed by the values encoded using RLE/Bit packed described +/// above (with the given bit width). +pub struct DictEncoder { + interner: Interner>, + + /// The buffered indices + indices: Vec, +} + +impl DictEncoder { + /// Creates new dictionary encoder. + pub fn new(desc: ColumnDescPtr) -> Self { + let storage = KeyStorage { + uniques: vec![], + size_in_bytes: 0, + type_length: desc.type_length() as usize, + }; + + Self { + interner: Interner::new(storage), + indices: vec![], + } + } + + /// Returns true if dictionary entries are sorted, false otherwise. + pub fn is_sorted(&self) -> bool { + // Sorting is not supported currently. + false + } + + /// Returns number of unique values (keys) in the dictionary. + pub fn num_entries(&self) -> usize { + self.interner.storage().uniques.len() + } + + /// Returns size of unique values (keys) in the dictionary, in bytes. + pub fn dict_encoded_size(&self) -> usize { + self.interner.storage().size_in_bytes + } + + /// Writes out the dictionary values with PLAIN encoding in a byte buffer, and return + /// the result. + pub fn write_dict(&self) -> Result { + let mut plain_encoder = PlainEncoder::::new(); + plain_encoder.put(&self.interner.storage().uniques)?; + plain_encoder.flush_buffer() + } + + /// Writes out the dictionary values with RLE encoding in a byte buffer, and return + /// the result. + pub fn write_indices(&mut self) -> Result { + let buffer_len = self.estimated_data_encoded_size(); + let mut buffer = Vec::with_capacity(buffer_len); + buffer.push(self.bit_width()); + + // Write bit width in the first byte + let mut encoder = RleEncoder::new_from_buf(self.bit_width(), buffer); + for index in &self.indices { + encoder.put(*index) + } + self.indices.clear(); + Ok(encoder.consume().into()) + } + + fn put_one(&mut self, value: &T::T) { + self.indices.push(self.interner.intern(value)); + } + + #[inline] + fn bit_width(&self) -> u8 { + num_required_bits(self.num_entries().saturating_sub(1) as u64) + } +} + +impl Encoder for DictEncoder { + fn put(&mut self, values: &[T::T]) -> Result<()> { + self.indices.reserve(values.len()); + for i in values { + self.put_one(i) + } + Ok(()) + } + + // Performance Note: + // As far as can be seen these functions are rarely called and as such we can hint to the + // compiler that they dont need to be folded into hot locations in the final output. + fn encoding(&self) -> Encoding { + Encoding::PLAIN_DICTIONARY + } + + /// Returns an estimate of the data page size in bytes + /// + /// This includes: + /// + + fn estimated_data_encoded_size(&self) -> usize { + let bit_width = self.bit_width(); + RleEncoder::max_buffer_size(bit_width, self.indices.len()) + } + + fn flush_buffer(&mut self) -> Result { + self.write_indices() + } + + /// Returns the estimated total memory usage + /// + /// For this encoder, the indices are unencoded bytes (refer to [`Self::write_indices`]). + fn estimated_memory_size(&self) -> usize { + self.interner.estimated_memory_size() + self.indices.capacity() * std::mem::size_of::() + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + use crate::data_type::{ + ByteArray, ByteArrayType, FixedLenByteArray, FixedLenByteArrayType, Int32Type, + }; + use crate::encodings::encoding::Encoder; + use crate::schema::types::{ColumnDescriptor, ColumnPath, Type as SchemaType}; + + fn make_col_desc() -> ColumnDescPtr { + make_col_desc_with_length::(-1) + } + + fn make_col_desc_with_length(type_length: i32) -> ColumnDescPtr { + let ty = SchemaType::primitive_type_builder("col", T::get_physical_type()) + .with_length(type_length) + .build() + .unwrap(); + Arc::new(ColumnDescriptor::new( + Arc::new(ty), + 0, + 0, + ColumnPath::new(vec![]), + )) + } + + #[test] + fn test_estimated_memory_size_primitive_with_duplicates() { + let mut encoder = DictEncoder::::new(make_col_desc::()); + let empty_size = encoder.estimated_memory_size(); + + // 3 distinct values, repeated to produce 9 indices total. + encoder.put(&[1, 2, 3, 1, 2, 3, 1, 2, 3]).unwrap(); + + let size = encoder.estimated_memory_size(); + + // Must account for the 3 unique dictionary entries. + let dict_entry_size = 3 * std::mem::size_of::(); + assert!( + size >= empty_size + dict_entry_size, + "memory size {size} should grow by at least the dict storage ({dict_entry_size} bytes)" + ); + + // Must also account for the 9 buffered indices. + let indices_size = 9 * std::mem::size_of::(); + assert!( + size >= empty_size + dict_entry_size + indices_size, + "memory size {size} should include indices ({indices_size} bytes)" + ); + } + + #[test] + fn test_estimated_memory_size_primitive_all_distinct() { + let mut encoder = DictEncoder::::new(make_col_desc::()); + let empty_size = encoder.estimated_memory_size(); + + let values: Vec = (0..100).collect(); + encoder.put(&values).unwrap(); + + let size = encoder.estimated_memory_size(); + + // Must account for the 100 unique dictionary entries. + let dict_entry_size = 100 * std::mem::size_of::(); + assert!( + size >= empty_size + dict_entry_size, + "memory size {size} should grow by at least the dict storage ({dict_entry_size} bytes)" + ); + + // Must also account for the 100 buffered indices. + let indices_size = 100 * std::mem::size_of::(); + assert!( + size >= empty_size + dict_entry_size + indices_size, + "memory size {size} should include indices ({indices_size} bytes)" + ); + } + + #[test] + fn test_estimated_memory_size_byte_array_with_duplicates() { + let mut encoder = DictEncoder::::new(make_col_desc::()); + let empty_size = encoder.estimated_memory_size(); + + // 3 distinct byte strings ("foo", "bar", "baz" — 3 bytes each), repeated to produce + // 9 indices total. + let vals: Vec = [ + "foo", "bar", "baz", "foo", "bar", "baz", "foo", "bar", "baz", + ] + .iter() + .map(|s| ByteArray::from(*s)) + .collect(); + encoder.put(&vals).unwrap(); + + let size = encoder.estimated_memory_size(); + + // Must account for the 3 unique dictionary entries, including their heap-allocated bytes. + let dict_entry_size = 3 * std::mem::size_of::() + 3 * 3; // 3 values × 3 bytes each + assert!( + size >= empty_size + dict_entry_size, + "memory size {size} should grow by at least the dict storage ({dict_entry_size} bytes)" + ); + + // Must also account for the 9 buffered indices. + let indices_size = 9 * std::mem::size_of::(); + assert!( + size >= empty_size + dict_entry_size + indices_size, + "memory size {size} should include indices ({indices_size} bytes)" + ); + } + + #[test] + fn test_estimated_memory_size_byte_array_all_distinct() { + let mut encoder = DictEncoder::::new(make_col_desc::()); + let empty_size = encoder.estimated_memory_size(); + + // 100 distinct values: "0".."9" (1 byte each) and "10".."99" (2 bytes each). + let values: Vec = (0..100_u32) + .map(|i| ByteArray::from(i.to_string().into_bytes())) + .collect(); + let bytes_total: usize = values.iter().map(|v| v.len()).sum(); // 10×1 + 90×2 = 190 + encoder.put(&values).unwrap(); + + let size = encoder.estimated_memory_size(); + + // Must account for the 100 unique dictionary entries, including their heap-allocated bytes. + let dict_entry_size = 100 * std::mem::size_of::() + bytes_total; + assert!( + size >= empty_size + dict_entry_size, + "memory size {size} should grow by at least the dict storage ({dict_entry_size} bytes)" + ); + + // Must also account for the 100 buffered indices. + let indices_size = 100 * std::mem::size_of::(); + assert!( + size >= empty_size + dict_entry_size + indices_size, + "memory size {size} should include indices ({indices_size} bytes)" + ); + } + + #[test] + fn test_estimated_memory_size_fixed_len_byte_array_with_duplicates() { + const TYPE_LEN: usize = 3; + let mut encoder = DictEncoder::::new(make_col_desc_with_length::< + FixedLenByteArrayType, + >(TYPE_LEN as i32)); + let empty_size = encoder.estimated_memory_size(); + + // 3 distinct 3-byte values, repeated to produce 9 indices total. + let vals = [ + b"foo", b"bar", b"baz", b"foo", b"bar", b"baz", b"foo", b"bar", b"baz", + ] + .iter() + .map(|b| FixedLenByteArray::from(b.to_vec())) + .collect::>(); + encoder.put(&vals).unwrap(); + + let size = encoder.estimated_memory_size(); + + // Must account for the 3 unique dictionary entries: struct overhead plus the + // fixed-length bytes allocated per entry. + let dict_entry_size = 3 * std::mem::size_of::() + 3 * TYPE_LEN; + assert!( + size >= empty_size + dict_entry_size, + "memory size {size} should grow by at least the dict storage ({dict_entry_size} bytes)" + ); + + // Must also account for the 9 buffered indices. + let indices_size = 9 * std::mem::size_of::(); + assert!( + size >= empty_size + dict_entry_size + indices_size, + "memory size {size} should include indices ({indices_size} bytes)" + ); + } + + #[test] + fn test_estimated_memory_size_fixed_len_byte_array_all_distinct() { + const TYPE_LEN: usize = 3; + let mut encoder = DictEncoder::::new(make_col_desc_with_length::< + FixedLenByteArrayType, + >(TYPE_LEN as i32)); + let empty_size = encoder.estimated_memory_size(); + + // 100 distinct 3-byte values: zero-padded big-endian u8 indices. + let values = (0..100_u8) + .map(|i| FixedLenByteArray::from(vec![0u8, 0u8, i])) + .collect::>(); + encoder.put(&values).unwrap(); + + let size = encoder.estimated_memory_size(); + + // Must account for the 100 unique dictionary entries: struct overhead plus the + // fixed-length bytes allocated per entry. + let dict_entry_size = 100 * std::mem::size_of::() + 100 * TYPE_LEN; + assert!( + size >= empty_size + dict_entry_size, + "memory size {size} should grow by at least the dict storage ({dict_entry_size} bytes)" + ); + + // Must also account for the 100 buffered indices. + let indices_size = 100 * std::mem::size_of::(); + assert!( + size >= empty_size + dict_entry_size + indices_size, + "memory size {size} should include indices ({indices_size} bytes)" + ); + } + + #[test] + fn test_estimated_memory_size_includes_interner_dedup_table() { + // The dedup `HashTable` in `Interner` is preallocated with + // `DEFAULT_DEDUP_CAPACITY` slots at construction, independent of any + // values pushed. + let encoder = DictEncoder::::new(make_col_desc::()); + + let size = encoder.estimated_memory_size(); + + assert!( + size > 0, + "memory size should include the preallocated dedup hash table" + ); + } + + #[test] + fn test_estimated_memory_size_accounts_for_indices_capacity() { + // Exercises the `indices.capacity()` (not `.len()`) accounting. + // After a flush, `indices` is cleared but its capacity is retained; pushing a + // smaller batch afterwards leaves capacity strictly greater than length. + let mut encoder = DictEncoder::::new(make_col_desc::()); + + let big: Vec = vec![0; 64]; + encoder.put(&big).unwrap(); + let _ = encoder.flush_buffer().unwrap(); + + let flushed_size = encoder.estimated_memory_size(); + + // Push a single value — indices.len() == 1 but indices.capacity() >= 64. + // No change on the key storage since the value is already interned. + encoder.put(&[0]).unwrap(); + + let size = encoder.estimated_memory_size(); + + assert_eq!( + size, flushed_size, + "memory size should include retained indices capacity", + ); + } + + #[test] + fn test_estimated_memory_size_accounts_for_uniques_capacity() { + let mut encoder = DictEncoder::::new(make_col_desc::()); + + let values: Vec = (0..64).collect(); + encoder.put(&values).unwrap(); + // Flush indices so they don't mask the uniques accounting in the lower bound. + let _ = encoder.flush_buffer().unwrap(); + + let size1 = encoder.estimated_memory_size(); + + // Push more values to trigger uniques capacity growth. + // The pre-allocated dedup hash table is unlikely to be resized. + let values: Vec = (64..128).collect(); + encoder.put(&values).unwrap(); + // Flush indices so they don't mask the uniques accounting in the lower bound. + let _ = encoder.flush_buffer().unwrap(); + + let size2 = encoder.estimated_memory_size(); + + let min_uniques_bytes = 64 * std::mem::size_of::(); + assert!( + size2 >= size1 + min_uniques_bytes, + "memory size {size2} should grow from {size1} by allocated uniques capacity \ + (at least {min_uniques_bytes} bytes)" + ); + } +} diff --git a/vendor/parquet/src/encodings/encoding/mod.rs b/vendor/parquet/src/encodings/encoding/mod.rs new file mode 100644 index 0000000..eeabcf4 --- /dev/null +++ b/vendor/parquet/src/encodings/encoding/mod.rs @@ -0,0 +1,1197 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Contains all supported encoders for Parquet. + +use std::{cmp, marker::PhantomData}; + +use crate::basic::*; +use crate::data_type::private::ParquetValueType; +use crate::data_type::*; +use crate::encodings::rle::RleEncoder; +use crate::errors::{ParquetError, Result}; +use crate::schema::types::ColumnDescPtr; +use crate::util::bit_util::{BitWriter, num_required_bits}; + +use byte_stream_split_encoder::{ByteStreamSplitEncoder, VariableWidthByteStreamSplitEncoder}; +use bytes::Bytes; +pub use dict_encoder::DictEncoder; + +mod byte_stream_split_encoder; +mod dict_encoder; + +// ---------------------------------------------------------------------- +// Encoders + +/// An Parquet encoder for the data type `T`. +/// +/// Currently this allocates internal buffers for the encoded values. After done putting +/// values, caller should call `flush_buffer()` to get an immutable buffer pointer. +pub trait Encoder: Send { + /// Encodes data from `values`. + fn put(&mut self, values: &[T::T]) -> Result<()>; + + /// Encodes data from `values`, which contains spaces for null values, that is + /// identified by `valid_bits`. + /// + /// Returns the number of non-null values encoded. + #[cfg(test)] + fn put_spaced(&mut self, values: &[T::T], valid_bits: &[u8]) -> Result { + let num_values = values.len(); + let mut buffer = Vec::with_capacity(num_values); + // TODO: this is pretty inefficient. Revisit in future. + for (i, item) in values.iter().enumerate().take(num_values) { + if crate::util::bit_util::get_bit(valid_bits, i) { + buffer.push(item.clone()); + } + } + self.put(&buffer[..])?; + Ok(buffer.len()) + } + + /// Returns the encoding type of this encoder. + fn encoding(&self) -> Encoding; + + /// Returns an estimate of the encoded data, in bytes. + /// Method call must be O(1). + fn estimated_data_encoded_size(&self) -> usize; + + /// Returns an estimate of the memory use of this encoder, in bytes + fn estimated_memory_size(&self) -> usize; + + /// Flushes the underlying byte buffer that's being processed by this encoder, and + /// return the immutable copy of it. This will also reset the internal state. + fn flush_buffer(&mut self) -> Result; +} + +/// Gets a encoder for the particular data type `T` and encoding `encoding`. Memory usage +/// for the encoder instance is tracked by `mem_tracker`. +pub fn get_encoder( + encoding: Encoding, + descr: &ColumnDescPtr, +) -> Result>> { + let encoder: Box> = match encoding { + Encoding::PLAIN => Box::new(PlainEncoder::new()), + Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY => { + return Err(general_err!( + "Cannot initialize this encoding through this function" + )); + } + Encoding::RLE => Box::new(RleValueEncoder::new()), + Encoding::DELTA_BINARY_PACKED => Box::new(DeltaBitPackEncoder::new()), + Encoding::DELTA_LENGTH_BYTE_ARRAY => Box::new(DeltaLengthByteArrayEncoder::new()), + Encoding::DELTA_BYTE_ARRAY => Box::new(DeltaByteArrayEncoder::new()), + 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()), + }, + e => return Err(nyi_err!("Encoding {} is not supported", e)), + }; + Ok(encoder) +} + +// ---------------------------------------------------------------------- +// Plain encoding + +/// Plain encoding that supports all types. +/// Values are encoded back to back. +/// The plain encoding is used whenever a more efficient encoding can not be used. +/// It stores the data in the following format: +/// - BOOLEAN - 1 bit per value, 0 is false; 1 is true. +/// - INT32 - 4 bytes per value, stored as little-endian. +/// - INT64 - 8 bytes per value, stored as little-endian. +/// - FLOAT - 4 bytes per value, stored as IEEE little-endian. +/// - DOUBLE - 8 bytes per value, stored as IEEE little-endian. +/// - BYTE_ARRAY - 4 byte length stored as little endian, followed by bytes. +/// - FIXED_LEN_BYTE_ARRAY - just the bytes are stored. +pub struct PlainEncoder { + buffer: Vec, + bit_writer: BitWriter, + _phantom: PhantomData, +} + +impl Default for PlainEncoder { + fn default() -> Self { + Self::new() + } +} + +impl PlainEncoder { + /// Creates new plain encoder. + pub fn new() -> Self { + Self { + buffer: vec![], + bit_writer: BitWriter::new(256), + _phantom: PhantomData, + } + } +} + +impl Encoder for PlainEncoder { + // Performance Note: + // As far as can be seen these functions are rarely called and as such we can hint to the + // compiler that they dont need to be folded into hot locations in the final output. + #[cold] + fn encoding(&self) -> Encoding { + Encoding::PLAIN + } + + fn estimated_data_encoded_size(&self) -> usize { + self.buffer.len() + self.bit_writer.bytes_written() + } + + #[inline] + fn flush_buffer(&mut self) -> Result { + self.buffer + .extend_from_slice(self.bit_writer.flush_buffer()); + self.bit_writer.clear(); + Ok(std::mem::take(&mut self.buffer).into()) + } + + #[inline] + fn put(&mut self, values: &[T::T]) -> Result<()> { + T::T::encode(values, &mut self.buffer, &mut self.bit_writer)?; + Ok(()) + } + + /// Return the estimated memory size of this encoder. + fn estimated_memory_size(&self) -> usize { + self.buffer.capacity() * std::mem::size_of::() + self.bit_writer.estimated_memory_size() + } +} + +// ---------------------------------------------------------------------- +// RLE encoding + +const DEFAULT_RLE_BUFFER_LEN: usize = 1024; + +/// RLE/Bit-Packing hybrid encoding for values. +/// Currently is used only for data pages v2 and supports boolean types. +pub struct RleValueEncoder { + // Buffer with raw values that we collect, + // when flushing buffer they are encoded using RLE encoder + encoder: Option, + _phantom: PhantomData, +} + +impl Default for RleValueEncoder { + fn default() -> Self { + Self::new() + } +} + +impl RleValueEncoder { + /// Creates new rle value encoder. + pub fn new() -> Self { + Self { + encoder: None, + _phantom: PhantomData, + } + } +} + +impl Encoder for RleValueEncoder { + #[inline] + fn put(&mut self, values: &[T::T]) -> Result<()> { + ensure_phys_ty!(Type::BOOLEAN, "RleValueEncoder only supports BoolType"); + + let rle_encoder = self.encoder.get_or_insert_with(|| { + let mut buffer = Vec::with_capacity(DEFAULT_RLE_BUFFER_LEN); + // Reserve space for length + buffer.extend_from_slice(&[0; 4]); + RleEncoder::new_from_buf(1, buffer) + }); + + for value in values { + let value = value.as_u64()?; + rle_encoder.put(value) + } + Ok(()) + } + + // Performance Note: + // As far as can be seen these functions are rarely called and as such we can hint to the + // compiler that they dont need to be folded into hot locations in the final output. + #[cold] + fn encoding(&self) -> Encoding { + Encoding::RLE + } + + #[inline] + fn estimated_data_encoded_size(&self) -> usize { + match self.encoder { + Some(ref enc) => enc.len(), + None => 0, + } + } + + #[inline] + fn flush_buffer(&mut self) -> Result { + ensure_phys_ty!(Type::BOOLEAN, "RleValueEncoder only supports BoolType"); + let rle_encoder = self + .encoder + .take() + .expect("RLE value encoder is not initialized"); + + // Flush all encoder buffers and raw values + let mut buf = rle_encoder.consume(); + assert!(buf.len() >= 4, "should have had padding inserted"); + + // Note that buf does not have any offset, all data is encoded bytes + let len = (buf.len() - 4) as i32; + buf[..4].copy_from_slice(&len.to_le_bytes()); + + Ok(buf.into()) + } + + /// return the estimated memory size of this encoder. + fn estimated_memory_size(&self) -> usize { + self.encoder + .as_ref() + .map_or(0, |enc| enc.estimated_memory_size()) + } +} + +// ---------------------------------------------------------------------- +// 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. +/// Consists of a header followed by blocks of delta encoded values binary packed. +/// +/// Delta-binary-packing: +/// ```shell +/// [page-header] [block 1], [block 2], ... [block N] +/// ``` +/// +/// Each page header consists of: +/// ```shell +/// [block size] [number of miniblocks in a block] [total value count] [first value] +/// ``` +/// +/// Each block consists of: +/// ```shell +/// [min delta] [list of bitwidths of miniblocks] [miniblocks] +/// ``` +/// +/// Current implementation writes values in `put` method, multiple calls to `put` to +/// existing block or start new block if block size is exceeded. Calling `flush_buffer` +/// writes out all data and resets internal state, including page header. +/// +/// Supports only INT32 and INT64. +pub struct DeltaBitPackEncoder { + page_header_writer: BitWriter, + bit_writer: BitWriter, + total_values: usize, + first_value: i64, + current_value: i64, + block_size: usize, + mini_block_size: usize, + num_mini_blocks: usize, + values_in_block: usize, + deltas: Vec, + _phantom: PhantomData, +} + +impl Default for DeltaBitPackEncoder { + fn default() -> Self { + Self::new() + } +} + +impl DeltaBitPackEncoder { + /// Creates new delta bit packed encoder. + pub fn new() -> Self { + Self::assert_supported_type(); + + // Size miniblocks so that they can be efficiently decoded + let mini_block_size = match T::T::PHYSICAL_TYPE { + Type::INT32 => 32, + Type::INT64 => 64, + _ => unreachable!(), + }; + + let num_mini_blocks = DEFAULT_NUM_MINI_BLOCKS; + let block_size = mini_block_size * num_mini_blocks; + assert_eq!(block_size % 128, 0); + + DeltaBitPackEncoder { + page_header_writer: BitWriter::new(MAX_PAGE_HEADER_WRITER_SIZE), + bit_writer: BitWriter::new(DEFAULT_BIT_WRITER_SIZE), + total_values: 0, + first_value: 0, + current_value: 0, // current value to keep adding deltas + block_size, // can write fewer values than block size for last block + mini_block_size, + num_mini_blocks, + values_in_block: 0, // will be at most block_size + deltas: vec![0; block_size], + _phantom: PhantomData, + } + } + + /// Writes page header for blocks, this method is invoked when we are done encoding + /// values. It is also okay to encode when no values have been provided + fn write_page_header(&mut self) { + // We ignore the result of each 'put' operation, because + // MAX_PAGE_HEADER_WRITER_SIZE is chosen to fit all header values and + // guarantees that writes will not fail. + + // Write the size of each block + self.page_header_writer.put_vlq_int(self.block_size as u64); + // Write the number of mini blocks + self.page_header_writer + .put_vlq_int(self.num_mini_blocks as u64); + // Write the number of all values (including non-encoded first value) + self.page_header_writer + .put_vlq_int(self.total_values as u64); + // Write first value + self.page_header_writer.put_zigzag_vlq_int(self.first_value); + } + + // Write current delta buffer (<= 'block size' values) into bit writer + #[inline(never)] + fn flush_block_values(&mut self) -> Result<()> { + if self.values_in_block == 0 { + return Ok(()); + } + + let mut min_delta = i64::MAX; + for i in 0..self.values_in_block { + min_delta = cmp::min(min_delta, self.deltas[i]); + } + + // Write min delta + self.bit_writer.put_zigzag_vlq_int(min_delta); + + // Slice to store bit width for each mini block + let offset = self.bit_writer.skip(self.num_mini_blocks); + + for i in 0..self.num_mini_blocks { + // Find how many values we need to encode - either block size or whatever + // values left + let n = cmp::min(self.mini_block_size, self.values_in_block); + if n == 0 { + // Decoders should be agnostic to the padding value, we therefore use 0xFF + // when running tests. However, not all implementations may handle this correctly + // so pad with 0 when not running tests + let pad_value = cfg!(test).then(|| 0xFF).unwrap_or(0); + for j in i..self.num_mini_blocks { + self.bit_writer.write_at(offset + j, pad_value); + } + break; + } + + // Compute the max delta in current mini block + let mut max_delta = i64::MIN; + for j in 0..n { + max_delta = cmp::max(max_delta, self.deltas[i * self.mini_block_size + j]); + } + + // Compute bit width to store (max_delta - min_delta) + let bit_width = num_required_bits(self.subtract_u64(max_delta, min_delta)) as usize; + self.bit_writer.write_at(offset + i, bit_width as u8); + + // Encode values in current mini block using min_delta and bit_width + for j in 0..n { + let packed_value = + self.subtract_u64(self.deltas[i * self.mini_block_size + j], min_delta); + self.bit_writer.put_value(packed_value, bit_width); + } + + // Pad the last block (n < mini_block_size) + for _ in n..self.mini_block_size { + self.bit_writer.put_value(0, bit_width); + } + + self.values_in_block -= n; + } + + assert_eq!( + self.values_in_block, 0, + "Expected 0 values in block, found {}", + self.values_in_block + ); + Ok(()) + } +} + +// Implementation is shared between Int32Type and Int64Type, +// see `DeltaBitPackEncoderConversion` below for specifics. +impl Encoder for DeltaBitPackEncoder { + fn put(&mut self, values: &[T::T]) -> Result<()> { + if values.is_empty() { + return Ok(()); + } + + // Define values to encode, initialize state + let mut idx = if self.total_values == 0 { + self.first_value = self.as_i64(values, 0); + self.current_value = self.first_value; + 1 + } else { + 0 + }; + // Add all values (including first value) + self.total_values += values.len(); + + // Write block + while idx < values.len() { + let value = self.as_i64(values, idx); + self.deltas[self.values_in_block] = self.subtract(value, self.current_value); + self.current_value = value; + idx += 1; + self.values_in_block += 1; + if self.values_in_block == self.block_size { + self.flush_block_values()?; + } + } + Ok(()) + } + + // Performance Note: + // As far as can be seen these functions are rarely called and as such we can hint to the + // compiler that they dont need to be folded into hot locations in the final output. + #[cold] + fn encoding(&self) -> Encoding { + Encoding::DELTA_BINARY_PACKED + } + + fn estimated_data_encoded_size(&self) -> usize { + self.bit_writer.bytes_written() + } + + fn flush_buffer(&mut self) -> Result { + // Write remaining values + self.flush_block_values()?; + // Write page header with total values + self.write_page_header(); + + let mut buffer = Vec::new(); + buffer.extend_from_slice(self.page_header_writer.flush_buffer()); + buffer.extend_from_slice(self.bit_writer.flush_buffer()); + + // Reset state + self.page_header_writer.clear(); + self.bit_writer.clear(); + self.total_values = 0; + self.first_value = 0; + self.current_value = 0; + self.values_in_block = 0; + + Ok(buffer.into()) + } + + /// return the estimated memory size of this encoder. + fn estimated_memory_size(&self) -> usize { + self.page_header_writer.estimated_memory_size() + + self.bit_writer.estimated_memory_size() + + self.deltas.capacity() * std::mem::size_of::() + + std::mem::size_of::() + } +} + +/// Helper trait to define specific conversions and subtractions when computing deltas +trait DeltaBitPackEncoderConversion { + // Method should panic if type is not supported, otherwise no-op + fn assert_supported_type(); + + fn as_i64(&self, values: &[T::T], index: usize) -> i64; + + fn subtract(&self, left: i64, right: i64) -> i64; + + fn subtract_u64(&self, left: i64, right: i64) -> u64; +} + +const DELTA_BIT_PACK_TYPE_ERROR: &str = + "DeltaBitPackDecoder only supports Int32Type, UInt32Type, Int64Type, and UInt64Type"; + +impl DeltaBitPackEncoderConversion for DeltaBitPackEncoder { + #[inline] + fn assert_supported_type() { + ensure_phys_ty!(Type::INT32 | Type::INT64, "{}", DELTA_BIT_PACK_TYPE_ERROR); + } + + #[inline] + fn as_i64(&self, values: &[T::T], index: usize) -> i64 { + values[index].as_i64().expect(DELTA_BIT_PACK_TYPE_ERROR) + } + + #[inline] + fn subtract(&self, left: i64, right: i64) -> i64 { + // It is okay for values to overflow, wrapping_sub wrapping around at the boundary + match T::get_physical_type() { + Type::INT32 => (left as i32).wrapping_sub(right as i32) as i64, + Type::INT64 => left.wrapping_sub(right), + _ => panic!("{}", DELTA_BIT_PACK_TYPE_ERROR), + } + } + + #[inline] + fn subtract_u64(&self, left: i64, right: i64) -> u64 { + match T::get_physical_type() { + // Conversion of i32 -> u32 -> u64 is to avoid non-zero left most bytes in int repr + Type::INT32 => (left as i32).wrapping_sub(right as i32) as u32 as u64, + Type::INT64 => left.wrapping_sub(right) as u64, + _ => panic!("{}", DELTA_BIT_PACK_TYPE_ERROR), + } + } +} + +// ---------------------------------------------------------------------- +// DELTA_LENGTH_BYTE_ARRAY encoding + +/// Encoding for byte arrays to separate the length values and the data. +/// The lengths are encoded using DELTA_BINARY_PACKED encoding, data is +/// stored as raw bytes. +pub struct DeltaLengthByteArrayEncoder { + // length encoder + len_encoder: DeltaBitPackEncoder, + // byte array data + data: Vec, + // data size in bytes of encoded values + encoded_size: usize, + _phantom: PhantomData, +} + +impl Default for DeltaLengthByteArrayEncoder { + fn default() -> Self { + Self::new() + } +} + +impl DeltaLengthByteArrayEncoder { + /// Creates new delta length byte array encoder. + pub fn new() -> Self { + Self { + len_encoder: DeltaBitPackEncoder::new(), + data: vec![], + encoded_size: 0, + _phantom: PhantomData, + } + } +} + +impl Encoder for DeltaLengthByteArrayEncoder { + fn put(&mut self, values: &[T::T]) -> Result<()> { + ensure_phys_ty!( + Type::BYTE_ARRAY | Type::FIXED_LEN_BYTE_ARRAY, + "DeltaLengthByteArrayEncoder only supports ByteArrayType" + ); + + let val_it = || { + values + .iter() + .map(|x| x.as_any().downcast_ref::().unwrap()) + }; + + let lengths: Vec = val_it().map(|byte_array| byte_array.len() as i32).collect(); + self.len_encoder.put(&lengths)?; + for byte_array in val_it() { + self.encoded_size += byte_array.len(); + self.data.push(byte_array.clone()); + } + + Ok(()) + } + + // Performance Note: + // As far as can be seen these functions are rarely called and as such we can hint to the + // compiler that they dont need to be folded into hot locations in the final output. + #[cold] + fn encoding(&self) -> Encoding { + Encoding::DELTA_LENGTH_BYTE_ARRAY + } + + fn estimated_data_encoded_size(&self) -> usize { + self.len_encoder.estimated_data_encoded_size() + self.encoded_size + } + + fn flush_buffer(&mut self) -> Result { + ensure_phys_ty!( + Type::BYTE_ARRAY | Type::FIXED_LEN_BYTE_ARRAY, + "DeltaLengthByteArrayEncoder only supports ByteArrayType" + ); + + let mut total_bytes = vec![]; + let lengths = self.len_encoder.flush_buffer()?; + total_bytes.extend_from_slice(&lengths); + self.data.iter().for_each(|byte_array| { + total_bytes.extend_from_slice(byte_array.data()); + }); + self.data.clear(); + self.encoded_size = 0; + + Ok(total_bytes.into()) + } + + /// return the estimated memory size of this encoder. + fn estimated_memory_size(&self) -> usize { + self.len_encoder.estimated_memory_size() + self.data.len() + std::mem::size_of::() + } +} + +// ---------------------------------------------------------------------- +// DELTA_BYTE_ARRAY encoding + +/// Encoding for byte arrays, prefix lengths are encoded using DELTA_BINARY_PACKED +/// encoding, followed by suffixes with DELTA_LENGTH_BYTE_ARRAY encoding. +pub struct DeltaByteArrayEncoder { + prefix_len_encoder: DeltaBitPackEncoder, + suffix_writer: DeltaLengthByteArrayEncoder, + previous: Vec, + _phantom: PhantomData, +} + +impl Default for DeltaByteArrayEncoder { + fn default() -> Self { + Self::new() + } +} + +impl DeltaByteArrayEncoder { + /// Creates new delta byte array encoder. + pub fn new() -> Self { + Self { + prefix_len_encoder: DeltaBitPackEncoder::new(), + suffix_writer: DeltaLengthByteArrayEncoder::new(), + previous: vec![], + _phantom: PhantomData, + } + } +} + +impl Encoder for DeltaByteArrayEncoder { + fn put(&mut self, values: &[T::T]) -> Result<()> { + let mut prefix_lengths: Vec = vec![]; + let mut suffixes: Vec = vec![]; + + let values = values + .iter() + .map(|x| x.as_any()) + .map(|x| match T::get_physical_type() { + Type::BYTE_ARRAY => x.downcast_ref::().unwrap(), + Type::FIXED_LEN_BYTE_ARRAY => x.downcast_ref::().unwrap(), + _ => panic!( + "DeltaByteArrayEncoder only supports ByteArrayType and FixedLenByteArrayType" + ), + }); + + for byte_array in values { + let current = byte_array.data(); + // Maximum prefix length that is shared between previous value and current + // value + let prefix_len = cmp::min(self.previous.len(), current.len()); + let mut match_len = 0; + while match_len < prefix_len && self.previous[match_len] == current[match_len] { + match_len += 1; + } + prefix_lengths.push(match_len as i32); + suffixes.push(byte_array.slice(match_len, byte_array.len() - match_len)); + // Update previous for the next prefix + self.previous.clear(); + self.previous.extend_from_slice(current); + } + self.prefix_len_encoder.put(&prefix_lengths)?; + self.suffix_writer.put(&suffixes)?; + + Ok(()) + } + + // Performance Note: + // As far as can be seen these functions are rarely called and as such we can hint to the + // compiler that they dont need to be folded into hot locations in the final output. + #[cold] + fn encoding(&self) -> Encoding { + Encoding::DELTA_BYTE_ARRAY + } + + fn estimated_data_encoded_size(&self) -> usize { + self.prefix_len_encoder.estimated_data_encoded_size() + + self.suffix_writer.estimated_data_encoded_size() + } + + fn flush_buffer(&mut self) -> Result { + match T::get_physical_type() { + Type::BYTE_ARRAY | Type::FIXED_LEN_BYTE_ARRAY => { + // TODO: investigate if we can merge lengths and suffixes + // without copying data into new vector. + let mut total_bytes = vec![]; + // Insert lengths ... + let lengths = self.prefix_len_encoder.flush_buffer()?; + total_bytes.extend_from_slice(&lengths); + // ... followed by suffixes + let suffixes = self.suffix_writer.flush_buffer()?; + total_bytes.extend_from_slice(&suffixes); + + self.previous.clear(); + Ok(total_bytes.into()) + } + _ => panic!( + "DeltaByteArrayEncoder only supports ByteArrayType and FixedLenByteArrayType" + ), + } + } + + /// return the estimated memory size of this encoder. + fn estimated_memory_size(&self) -> usize { + self.prefix_len_encoder.estimated_memory_size() + + self.suffix_writer.estimated_memory_size() + + (self.previous.capacity() * std::mem::size_of::()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::sync::Arc; + + use crate::encodings::decoding::{Decoder, DictDecoder, PlainDecoder, get_decoder}; + use crate::schema::types::{ColumnDescPtr, ColumnDescriptor, ColumnPath, Type as SchemaType}; + use crate::util::bit_util; + use crate::util::test_common::rand_gen::{RandGen, random_bytes}; + + const TEST_SET_SIZE: usize = 1024; + + #[test] + fn test_get_encoders() { + // 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 + create_and_check_encoder::( + 0, + Encoding::RLE_DICTIONARY, + Some(general_err!( + "Cannot initialize this encoding through this function" + )), + ); + create_and_check_encoder::( + 0, + Encoding::PLAIN_DICTIONARY, + Some(general_err!( + "Cannot initialize this encoding through this function" + )), + ); + + // unsupported + #[allow(deprecated)] + create_and_check_encoder::( + 0, + Encoding::BIT_PACKED, + Some(nyi_err!("Encoding BIT_PACKED is not supported")), + ); + } + + #[test] + fn test_bool() { + BoolType::test(Encoding::PLAIN, TEST_SET_SIZE, -1); + BoolType::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1); + BoolType::test(Encoding::RLE, TEST_SET_SIZE, -1); + } + + #[test] + fn test_i32() { + Int32Type::test(Encoding::PLAIN, TEST_SET_SIZE, -1); + Int32Type::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1); + Int32Type::test(Encoding::DELTA_BINARY_PACKED, TEST_SET_SIZE, -1); + Int32Type::test(Encoding::BYTE_STREAM_SPLIT, TEST_SET_SIZE, -1); + } + + #[test] + fn test_i64() { + Int64Type::test(Encoding::PLAIN, TEST_SET_SIZE, -1); + Int64Type::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1); + Int64Type::test(Encoding::DELTA_BINARY_PACKED, TEST_SET_SIZE, -1); + Int64Type::test(Encoding::BYTE_STREAM_SPLIT, TEST_SET_SIZE, -1); + } + + #[test] + fn test_i96() { + Int96Type::test(Encoding::PLAIN, TEST_SET_SIZE, -1); + Int96Type::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1); + } + + #[test] + fn test_float() { + FloatType::test(Encoding::PLAIN, TEST_SET_SIZE, -1); + FloatType::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1); + FloatType::test(Encoding::BYTE_STREAM_SPLIT, TEST_SET_SIZE, -1); + } + + #[test] + fn test_double() { + DoubleType::test(Encoding::PLAIN, TEST_SET_SIZE, -1); + DoubleType::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1); + DoubleType::test(Encoding::BYTE_STREAM_SPLIT, TEST_SET_SIZE, -1); + } + + #[test] + fn test_byte_array() { + ByteArrayType::test(Encoding::PLAIN, TEST_SET_SIZE, -1); + ByteArrayType::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1); + ByteArrayType::test(Encoding::DELTA_LENGTH_BYTE_ARRAY, TEST_SET_SIZE, -1); + ByteArrayType::test(Encoding::DELTA_BYTE_ARRAY, TEST_SET_SIZE, -1); + } + + #[test] + fn test_fixed_len_byte_array() { + FixedLenByteArrayType::test(Encoding::PLAIN, TEST_SET_SIZE, 100); + FixedLenByteArrayType::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, 100); + FixedLenByteArrayType::test(Encoding::DELTA_BYTE_ARRAY, TEST_SET_SIZE, 100); + FixedLenByteArrayType::test(Encoding::BYTE_STREAM_SPLIT, TEST_SET_SIZE, 100); + } + + #[test] + fn test_dict_encoded_size() { + fn run_test(type_length: i32, values: &[T::T], expected_size: usize) { + let mut encoder = create_test_dict_encoder::(type_length); + assert_eq!(encoder.dict_encoded_size(), 0); + encoder.put(values).unwrap(); + assert_eq!(encoder.dict_encoded_size(), expected_size); + // We do not reset encoded size of the dictionary keys after flush_buffer + encoder.flush_buffer().unwrap(); + assert_eq!(encoder.dict_encoded_size(), expected_size); + } + + // Only 2 variations of values 1 byte each + run_test::(-1, &[true, false, true, false, true], 2); + run_test::(-1, &[1i32, 2i32, 3i32, 4i32, 5i32], 20); + run_test::(-1, &[1i64, 2i64, 3i64, 4i64, 5i64], 40); + run_test::(-1, &[1f32, 2f32, 3f32, 4f32, 5f32], 20); + run_test::(-1, &[1f64, 2f64, 3f64, 4f64, 5f64], 40); + // Int96: len + reference + run_test::( + -1, + &[Int96::from(vec![1, 2, 3]), Int96::from(vec![2, 3, 4])], + 24, + ); + run_test::(-1, &[ByteArray::from("abcd"), ByteArray::from("efj")], 15); + run_test::( + 2, + &[ByteArray::from("ab").into(), ByteArray::from("bc").into()], + 4, + ); + } + + #[test] + fn test_estimated_data_encoded_size() { + fn run_test( + encoding: Encoding, + type_length: i32, + values: &[T::T], + initial_size: usize, + max_size: usize, + flush_size: usize, + ) { + let mut encoder = match encoding { + Encoding::PLAIN_DICTIONARY | Encoding::RLE_DICTIONARY => { + Box::new(create_test_dict_encoder::(type_length)) + } + _ => create_test_encoder::(type_length, encoding), + }; + assert_eq!(encoder.estimated_data_encoded_size(), initial_size); + + encoder.put(values).unwrap(); + assert_eq!(encoder.estimated_data_encoded_size(), max_size); + + encoder.flush_buffer().unwrap(); + assert_eq!(encoder.estimated_data_encoded_size(), flush_size); + } + + // PLAIN + run_test::(Encoding::PLAIN, -1, &[123; 1024], 0, 4096, 0); + + // DICTIONARY + // NOTE: The final size is almost the same because the dictionary entries are + // preserved after encoded values have been written. + run_test::(Encoding::RLE_DICTIONARY, -1, &[123, 1024], 0, 2, 0); + + // DELTA_BINARY_PACKED + run_test::(Encoding::DELTA_BINARY_PACKED, -1, &[123; 1024], 0, 35, 0); + + // RLE + let mut values = vec![]; + values.extend_from_slice(&[true; 16]); + values.extend_from_slice(&[false; 16]); + run_test::(Encoding::RLE, -1, &values, 0, 6, 0); + + // DELTA_LENGTH_BYTE_ARRAY + run_test::( + Encoding::DELTA_LENGTH_BYTE_ARRAY, + -1, + &[ByteArray::from("ab"), ByteArray::from("abc")], + 0, + 5, // only value bytes, length encoder is not flushed yet + 0, + ); + + // DELTA_BYTE_ARRAY + run_test::( + Encoding::DELTA_BYTE_ARRAY, + -1, + &[ByteArray::from("ab"), ByteArray::from("abc")], + 0, + 3, // only suffix bytes, length encoder is not flushed yet + 0, + ); + + // BYTE_STREAM_SPLIT + run_test::(Encoding::BYTE_STREAM_SPLIT, -1, &[0.1, 0.2], 0, 8, 0); + } + + #[test] + fn test_byte_stream_split_example_f32() { + // Test data from https://github.com/apache/parquet-format/blob/2a481fe1aad64ff770e21734533bb7ef5a057dac/Encodings.md#byte-stream-split-byte_stream_split--9 + let mut encoder = create_test_encoder::(0, Encoding::BYTE_STREAM_SPLIT); + let mut decoder = create_test_decoder::(0, Encoding::BYTE_STREAM_SPLIT); + + let input = vec![ + f32::from_le_bytes([0xAA, 0xBB, 0xCC, 0xDD]), + f32::from_le_bytes([0x00, 0x11, 0x22, 0x33]), + f32::from_le_bytes([0xA3, 0xB4, 0xC5, 0xD6]), + ]; + + encoder.put(&input).unwrap(); + let encoded = encoder.flush_buffer().unwrap(); + + assert_eq!( + encoded, + Bytes::from(vec![ + 0xAA_u8, 0x00, 0xA3, 0xBB, 0x11, 0xB4, 0xCC, 0x22, 0xC5, 0xDD, 0x33, 0xD6 + ]) + ); + + let mut decoded = vec![0.0; input.len()]; + decoder.set_data(encoded, input.len()).unwrap(); + decoder.get(&mut decoded).unwrap(); + + assert_eq!(decoded, input); + } + + // See: https://github.com/sunchao/parquet-rs/issues/47 + #[test] + fn test_issue_47() { + let mut encoder = create_test_encoder::(0, Encoding::DELTA_BYTE_ARRAY); + let mut decoder = create_test_decoder::(0, Encoding::DELTA_BYTE_ARRAY); + + let input = vec![ + ByteArray::from("aa"), + ByteArray::from("aaa"), + ByteArray::from("aa"), + ByteArray::from("aaa"), + ]; + + let mut output = vec![ByteArray::default(); input.len()]; + + let mut result = put_and_get(&mut encoder, &mut decoder, &input[..2], &mut output[..2]); + assert!( + result.is_ok(), + "first put_and_get() failed with: {}", + result.unwrap_err() + ); + result = put_and_get(&mut encoder, &mut decoder, &input[2..], &mut output[2..]); + assert!( + result.is_ok(), + "second put_and_get() failed with: {}", + result.unwrap_err() + ); + assert_eq!(output, input); + } + + trait EncodingTester { + fn test(enc: Encoding, total: usize, type_length: i32) { + let result = match enc { + Encoding::PLAIN_DICTIONARY | Encoding::RLE_DICTIONARY => { + Self::test_dict_internal(total, type_length) + } + enc => Self::test_internal(enc, total, type_length), + }; + + assert!( + result.is_ok(), + "Expected result to be OK but got err:\n {}", + result.unwrap_err() + ); + } + + fn test_internal(enc: Encoding, total: usize, type_length: i32) -> Result<()>; + + fn test_dict_internal(total: usize, type_length: i32) -> Result<()>; + } + + impl> EncodingTester for T { + fn test_internal(enc: Encoding, total: usize, type_length: i32) -> Result<()> { + let mut encoder = create_test_encoder::(type_length, enc); + let mut decoder = create_test_decoder::(type_length, enc); + let mut values = >::gen_vec(type_length, total); + let mut result_data = vec![T::T::default(); total]; + + // Test put/get spaced. + let num_bytes = bit_util::ceil(total as i64, 8); + let valid_bits = random_bytes(num_bytes as usize); + let values_written = encoder.put_spaced(&values[..], &valid_bits[..])?; + let data = encoder.flush_buffer()?; + decoder.set_data(data, values_written)?; + let _ = decoder.get_spaced( + &mut result_data[..], + values.len() - values_written, + &valid_bits[..], + )?; + + // Check equality + for i in 0..total { + if bit_util::get_bit(&valid_bits[..], i) { + assert_eq!(result_data[i], values[i]); + } else { + assert_eq!(result_data[i], T::T::default()); + } + } + + let mut actual_total = put_and_get( + &mut encoder, + &mut decoder, + &values[..], + &mut result_data[..], + )?; + assert_eq!(actual_total, total); + assert_eq!(result_data, values); + + // Encode more data after flush and test with decoder + + values = >::gen_vec(type_length, total); + actual_total = put_and_get( + &mut encoder, + &mut decoder, + &values[..], + &mut result_data[..], + )?; + assert_eq!(actual_total, total); + assert_eq!(result_data, values); + + Ok(()) + } + + fn test_dict_internal(total: usize, type_length: i32) -> Result<()> { + let mut encoder = create_test_dict_encoder::(type_length); + let mut values = >::gen_vec(type_length, total); + encoder.put(&values[..])?; + + let mut data = encoder.flush_buffer()?; + let mut decoder = create_test_dict_decoder::(); + let mut dict_decoder = PlainDecoder::::new(type_length); + dict_decoder.set_data(encoder.write_dict()?, encoder.num_entries())?; + decoder.set_dict(Box::new(dict_decoder))?; + let mut result_data = vec![T::T::default(); total]; + decoder.set_data(data, total)?; + let mut actual_total = decoder.get(&mut result_data)?; + + assert_eq!(actual_total, total); + assert_eq!(result_data, values); + + // Encode more data after flush and test with decoder + + values = >::gen_vec(type_length, total); + encoder.put(&values[..])?; + data = encoder.flush_buffer()?; + + let mut dict_decoder = PlainDecoder::::new(type_length); + dict_decoder.set_data(encoder.write_dict()?, encoder.num_entries())?; + decoder.set_dict(Box::new(dict_decoder))?; + decoder.set_data(data, total)?; + actual_total = decoder.get(&mut result_data)?; + + assert_eq!(actual_total, total); + assert_eq!(result_data, values); + + Ok(()) + } + } + + fn put_and_get( + encoder: &mut Box>, + decoder: &mut Box>, + input: &[T::T], + output: &mut [T::T], + ) -> Result { + encoder.put(input)?; + let data = encoder.flush_buffer()?; + decoder.set_data(data, input.len())?; + decoder.get(output) + } + + fn create_and_check_encoder( + type_length: i32, + encoding: Encoding, + err: Option, + ) { + let desc = create_test_col_desc_ptr(type_length, T::get_physical_type()); + let encoder = get_encoder::(encoding, &desc); + match err { + Some(parquet_error) => { + assert_eq!( + encoder.err().unwrap().to_string(), + parquet_error.to_string() + ) + } + None => assert_eq!(encoder.unwrap().encoding(), encoding), + } + } + + // Creates test column descriptor. + fn create_test_col_desc_ptr(type_len: i32, t: Type) -> ColumnDescPtr { + let ty = SchemaType::primitive_type_builder("t", t) + .with_length(type_len) + .build() + .unwrap(); + Arc::new(ColumnDescriptor::new( + Arc::new(ty), + 0, + 0, + ColumnPath::new(vec![]), + )) + } + + fn create_test_encoder(type_len: i32, enc: Encoding) -> Box> { + let desc = create_test_col_desc_ptr(type_len, T::get_physical_type()); + get_encoder(enc, &desc).unwrap() + } + + fn create_test_decoder(type_len: i32, enc: Encoding) -> Box> { + let desc = create_test_col_desc_ptr(type_len, T::get_physical_type()); + get_decoder(desc, enc).unwrap() + } + + fn create_test_dict_encoder(type_len: i32) -> DictEncoder { + let desc = create_test_col_desc_ptr(type_len, T::get_physical_type()); + DictEncoder::::new(desc) + } + + fn create_test_dict_decoder() -> DictDecoder { + DictDecoder::::new() + } +} diff --git a/vendor/parquet/src/encodings/levels.rs b/vendor/parquet/src/encodings/levels.rs new file mode 100644 index 0000000..841afd8 --- /dev/null +++ b/vendor/parquet/src/encodings/levels.rs @@ -0,0 +1,285 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::mem; + +use super::rle::RleEncoder; + +use crate::data_type::AsBytes; +use crate::util::bit_util::num_required_bits; + +/// Encoder for definition/repetition levels. +/// Currently only supports Rle and BitPacked (dev/null) encoding, including v2. +pub enum LevelEncoder { + Rle(RleEncoder), + RleV2(RleEncoder), +} + +impl LevelEncoder { + /// Creates a new streaming level encoder for Data Page v1. + /// + /// This does not require knowing the number of values + /// upfront, making it suitable for incremental encoding where levels are fed in + /// as they arrive via [`put_with_observer`](Self::put_with_observer). + pub fn v1_streaming(max_level: i16) -> Self { + let bit_width = num_required_bits(max_level as u64); + // Reserve space for length header + let buffer = vec![0u8; 4]; + LevelEncoder::Rle(RleEncoder::new_from_buf(bit_width, buffer)) + } + + /// Creates a new streaming RLE level encoder for Data Page v2. + /// + /// This does not require knowing the number of values + /// upfront, making it suitable for incremental encoding where levels are fed in + /// as they arrive via [`put_with_observer`](Self::put_with_observer). + pub fn v2_streaming(max_level: i16) -> Self { + let bit_width = num_required_bits(max_level as u64); + LevelEncoder::RleV2(RleEncoder::new_from_buf(bit_width, Vec::new())) + } + + /// Put/encode levels vector into this level encoder and call + /// `observer(value, count)` for each run of identical values encountered + /// during encoding. + /// + /// Returns number of encoded values that are less than or equal to length + /// of the input buffer. + /// + /// This method does **not** flush the underlying encoder, so it can be called + /// incrementally across multiple batches without forcing run boundaries. + /// The encoder is flushed automatically when [`consume`](Self::consume) is called. + #[inline] + pub fn put_with_observer(&mut self, buffer: &[i16], mut observer: F) -> usize + where + F: FnMut(i16, usize), + { + match *self { + LevelEncoder::Rle(ref mut encoder) | LevelEncoder::RleV2(ref mut encoder) => { + let mut remaining = buffer; + while let Some((&value, rest)) = remaining.split_first() { + encoder.put(value as u64); + // After put(), check if the encoder just entered RLE + // accumulation mode. If so, scan ahead for the rest of + // this run to batch the observer call and bulk-extend. + if encoder.is_accumulating_rle(value as u64) { + let run_len = rest.iter().take_while(|&&v| v == value).count(); + if run_len > 0 { + encoder.extend_run(run_len); + } + observer(value, 1 + run_len); + remaining = &rest[run_len..]; + } else { + observer(value, 1); + remaining = rest; + } + } + buffer.len() + } + } + } + + /// Encode `count` repetitions of a single level value, calling + /// `observer(value, count)` exactly once. + /// + /// This can be used to encode uniform runs without allocating a level + /// buffer. Like [`Self::put_with_observer`], but specialized for a single + /// repeated level value. + /// + /// This is O(1) amortized for RLE-based encoders (after a small warmup). + #[inline] + pub fn put_n_with_observer(&mut self, value: i16, count: usize, mut observer: F) + where + F: FnMut(i16, usize), + { + let encoder = match *self { + LevelEncoder::Rle(ref mut encoder) | LevelEncoder::RleV2(ref mut encoder) => encoder, + }; + + // Feed values individually until the encoder enters RLE accumulation + // mode for this value, or until we've encoded everything. + let mut remaining = count; + while remaining > 0 && !encoder.is_accumulating_rle(value as u64) { + encoder.put(value as u64); + remaining -= 1; + } + // If we're now in accumulation mode, bulk-extend the rest. + if remaining > 0 { + encoder.extend_run(remaining); + } + observer(value, count); + } + + /// Finalizes level encoder, flush all intermediate buffers and return resulting + /// encoded buffer. Returned buffer is already truncated to encoded bytes only. + #[inline] + #[allow(unused)] + pub fn consume(self) -> Vec { + match self { + LevelEncoder::Rle(encoder) => { + let mut encoded_data = encoder.consume(); + // Account for the buffer offset + let encoded_len = encoded_data.len() - mem::size_of::(); + let len = (encoded_len as i32).to_le(); + let len_bytes = len.as_bytes(); + encoded_data[0..len_bytes.len()].copy_from_slice(len_bytes); + encoded_data + } + LevelEncoder::RleV2(encoder) => encoder.consume(), + } + } + + /// Flushes all intermediate buffers, passes the encoded data to `f`, then + /// resets the encoder for reuse while retaining the buffer allocation. + #[inline] + pub fn flush_to(&mut self, f: F) -> R + where + F: FnOnce(&[u8]) -> R, + { + let result = match self { + LevelEncoder::Rle(encoder) => { + let data = encoder.flush_buffer_mut(); + // Patch the 4-byte length header reserved at the start of the buffer + let encoded_len = (data.len() - mem::size_of::()) as i32; + data[..4].copy_from_slice(&encoded_len.to_le_bytes()); + f(data) + } + LevelEncoder::RleV2(encoder) => f(encoder.flush_buffer()), + }; + match self { + LevelEncoder::Rle(encoder) => { + encoder.clear(); + // Re-reserve the 4-byte length header for the next page + encoder.skip(mem::size_of::()); + } + LevelEncoder::RleV2(encoder) => encoder.clear(), + } + result + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Encode `count` repetitions of `value` using `put_with_observer` and + /// return the raw encoded bytes. + fn reference_encode(max_level: i16, values: &[i16]) -> Vec { + let mut enc = LevelEncoder::v2_streaming(max_level); + enc.put_with_observer(values, |_, _| {}); + enc.consume() + } + + #[test] + fn test_put_n_with_observer_large_run() { + // Large count exercises the bulk extend_run path (past the 8-value warmup). + let max_level = 3; + let count = 10_000; + let mut enc = LevelEncoder::v2_streaming(max_level); + enc.put_n_with_observer(2, count, |_, _| {}); + assert_eq!(enc.consume(), reference_encode(max_level, &vec![2; count])); + } + + #[test] + fn test_put_n_with_observer_small_count() { + // Count smaller than the RLE warmup threshold — only the per-element loop runs. + let max_level = 3; + let mut enc = LevelEncoder::v2_streaming(max_level); + enc.put_n_with_observer(1, 5, |_, _| {}); + assert_eq!(enc.consume(), reference_encode(max_level, &[1; 5])); + } + + #[test] + fn test_put_n_with_observer_exact_threshold() { + // Exactly 8 values: the warmup loop completes and extend_run gets 0. + let max_level = 3; + let mut enc = LevelEncoder::v2_streaming(max_level); + enc.put_n_with_observer(3, 8, |_, _| {}); + assert_eq!(enc.consume(), reference_encode(max_level, &[3; 8])); + } + + #[test] + fn test_put_n_with_observer_single_value() { + let max_level = 1; + let mut enc = LevelEncoder::v2_streaming(max_level); + enc.put_n_with_observer(1, 1, |_, _| {}); + assert_eq!(enc.consume(), reference_encode(max_level, &[1])); + } + + #[test] + fn test_put_n_with_observer_zero_count() { + let max_level = 3; + let mut enc = LevelEncoder::v2_streaming(max_level); + enc.put_n_with_observer(2, 0, |_, _| {}); + assert_eq!(enc.consume(), reference_encode(max_level, &[])); + } + + #[test] + fn test_put_n_with_observer_calls_observer_exactly_once() { + let mut enc = LevelEncoder::v2_streaming(3); + let mut calls: Vec<(i16, usize)> = Vec::new(); + enc.put_n_with_observer(2, 500, |val, cnt| calls.push((val, cnt))); + assert_eq!(calls, vec![(2, 500)]); + } + + #[test] + fn test_put_n_with_observer_zero_count_calls_observer() { + let mut enc = LevelEncoder::v2_streaming(3); + let mut calls: Vec<(i16, usize)> = Vec::new(); + enc.put_n_with_observer(1, 0, |val, cnt| calls.push((val, cnt))); + assert_eq!(calls, vec![(1, 0)]); + } + + #[test] + fn test_put_n_with_observer_followed_by_different_value() { + // Two consecutive put_n calls with different values — verifies that + // the encoder correctly transitions between runs. + let max_level = 3; + let mut enc = LevelEncoder::v2_streaming(max_level); + let mut calls: Vec<(i16, usize)> = Vec::new(); + enc.put_n_with_observer(1, 100, |v, c| calls.push((v, c))); + enc.put_n_with_observer(3, 200, |v, c| calls.push((v, c))); + assert_eq!(calls, vec![(1, 100), (3, 200)]); + + let reference = reference_encode(max_level, &[&[1i16; 100][..], &[3i16; 200]].concat()); + assert_eq!(enc.consume(), reference); + } + + #[test] + fn test_put_n_with_observer_interleaved_with_put_with_observer() { + // Mix put_n_with_observer and put_with_observer to verify they compose. + let max_level = 3; + let mut enc = LevelEncoder::v2_streaming(max_level); + enc.put_n_with_observer(2, 50, |_, _| {}); + enc.put_with_observer(&[0, 0, 1, 1, 3], |_, _| {}); + enc.put_n_with_observer(2, 50, |_, _| {}); + + let input = [&[2i16; 50][..], &[0, 0, 1, 1, 3], &[2i16; 50]].concat(); + assert_eq!(enc.consume(), reference_encode(max_level, &input)); + } + + #[test] + fn test_put_n_with_observer_v1_roundtrip() { + // Also verify V1 (Rle variant with length header) works correctly. + let max_level = 3; + let mut enc = LevelEncoder::v1_streaming(max_level); + enc.put_n_with_observer(2, 1000, |_, _| {}); + + let mut ref_enc = LevelEncoder::v1_streaming(max_level); + ref_enc.put_with_observer(&[2; 1000], |_, _| {}); + assert_eq!(enc.consume(), ref_enc.consume()); + } +} diff --git a/vendor/parquet/src/encodings/mod.rs b/vendor/parquet/src/encodings/mod.rs new file mode 100644 index 0000000..894c4fb --- /dev/null +++ b/vendor/parquet/src/encodings/mod.rs @@ -0,0 +1,21 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +pub mod decoding; +pub mod encoding; +pub mod levels; +experimental!(pub(crate) mod rle); diff --git a/vendor/parquet/src/encodings/rle.rs b/vendor/parquet/src/encodings/rle.rs new file mode 100644 index 0000000..5de9dc8 --- /dev/null +++ b/vendor/parquet/src/encodings/rle.rs @@ -0,0 +1,1205 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Rle/Bit-Packing Hybrid Encoding +//! The grammar for this encoding looks like the following (copied verbatim +//! from ): +//! +//! rle-bit-packed-hybrid: `` `` +//! length := length of the `` in bytes stored as 4 bytes little endian +//! encoded-data := ``* +//! run := `` | `` +//! bit-packed-run := `` `` +//! bit-packed-header := varint-encode(`` << 1 | 1) +//! we always bit-pack a multiple of 8 values at a time, so we only store the number of +//! values / 8 +//! bit-pack-count := (number of values in this run) / 8 +//! bit-packed-values := *see 1 below* +//! rle-run := `` `` +//! rle-header := varint-encode( (number of times repeated) << 1) +//! repeated-value := value that is repeated, using a fixed-width of +//! round-up-to-next-byte(bit-width) + +use std::{cmp, mem::size_of}; + +use bytes::Bytes; + +use crate::errors::{ParquetError, Result}; +use crate::util::bit_util::{self, BitReader, BitWriter, FromBitpacked}; + +/// Number of values in one bit-packed group. The Parquet RLE/bit-packing hybrid +/// format always bit-packs values in multiples of this count (see the +/// [format spec](https://github.com/apache/parquet-format/blob/master/Encodings.md#run-length-encoding--bit-packing-hybrid-rle--3): +/// "we always bit-pack a multiple of 8 values at a time"). +const BIT_PACK_GROUP_SIZE: usize = 8; + +/// Maximum groups of `BIT_PACK_GROUP_SIZE` values per bit-packed run. Current value is 64. +const MAX_GROUPS_PER_BIT_PACKED_RUN: usize = 1 << 6; + +/// A RLE/Bit-Packing hybrid encoder. +// TODO: tracking memory usage +pub struct RleEncoder { + // Number of bits needed to encode the value. Must be in the range of [0, 64]. + bit_width: u8, + + // Underlying writer which holds an internal buffer. + bit_writer: BitWriter, + + // Buffered values for bit-packed runs. + buffered_values: [u64; BIT_PACK_GROUP_SIZE], + + // Number of current buffered values. Must be less than BIT_PACK_GROUP_SIZE. + num_buffered_values: usize, + + // The current (also last) value that was written and the count of how many + // times in a row that value has been seen. + current_value: u64, + + // The number of repetitions for `current_value`. If this gets too high we'd + // switch to use RLE encoding. + repeat_count: usize, + + // Number of bit-packed values in the current run. This doesn't include values + // in `buffered_values`. + bit_packed_count: usize, + + // The position of the indicator byte in the `bit_writer`. + indicator_byte_pos: i64, +} + +impl RleEncoder { + #[allow(unused)] + pub fn new(bit_width: u8, buffer_len: usize) -> Self { + let buffer = Vec::with_capacity(buffer_len); + RleEncoder::new_from_buf(bit_width, buffer) + } + + /// Initialize the encoder from existing `buffer` + pub fn new_from_buf(bit_width: u8, buffer: Vec) -> Self { + debug_assert!(bit_width <= 64); + let bit_writer = BitWriter::new_from_buf(buffer); + RleEncoder { + bit_width, + bit_writer, + buffered_values: [0; BIT_PACK_GROUP_SIZE], + num_buffered_values: 0, + current_value: 0, + repeat_count: 0, + bit_packed_count: 0, + indicator_byte_pos: -1, + } + } + + /// Returns the maximum buffer size to encode `num_values` values with + /// `bit_width`. + pub fn max_buffer_size(bit_width: u8, num_values: usize) -> usize { + // The maximum size occurs with the shortest possible runs of BIT_PACK_GROUP_SIZE + let num_runs = bit_util::ceil(num_values, BIT_PACK_GROUP_SIZE); + + // The number of bytes in a run of BIT_PACK_GROUP_SIZE + let bytes_per_run = bit_width as usize; + + // The maximum size if stored as shortest possible bit packed runs of BIT_PACK_GROUP_SIZE + let bit_packed_max_size = num_runs + num_runs * bytes_per_run; + + // The length of `BIT_PACK_GROUP_SIZE` VLQ encoded + let rle_len_prefix = 1; + + // The length of an RLE run of BIT_PACK_GROUP_SIZE + let min_rle_run_size = + rle_len_prefix + bit_util::ceil(bit_width as usize, u8::BITS as usize); + + // The maximum size if stored as shortest possible RLE runs of BIT_PACK_GROUP_SIZE + let rle_max_size = num_runs * min_rle_run_size; + + bit_packed_max_size.max(rle_max_size) + } + + /// Returns `true` if the encoder is currently in RLE accumulation mode + /// for the given value (i.e., `repeat_count >= BIT_PACK_GROUP_SIZE` and `current_value == value`). + /// + /// The encoder enters accumulation mode as soon as the 8th consecutive identical + /// value has been seen: at that point `flush_buffered_values` has committed the + /// RLE decision and cleared the staging buffer, so no more per-element work is + /// needed. Callers may use [`extend_run`](Self::extend_run) to add further + /// repetitions in O(1) once this returns `true`. + #[inline] + pub fn is_accumulating_rle(&self, value: u64) -> bool { + self.repeat_count >= BIT_PACK_GROUP_SIZE && self.current_value == value + } + + /// Extends the current RLE run by `count` additional repetitions. + /// + /// # Preconditions + /// The caller **must** have verified [`is_accumulating_rle`](Self::is_accumulating_rle) + /// returns `true` for the same value before calling this method. + #[inline] + pub fn extend_run(&mut self, count: usize) { + debug_assert!(self.repeat_count >= BIT_PACK_GROUP_SIZE); + self.repeat_count += count; + } + + /// Encodes `value`, which must be representable with `bit_width` bits. + #[inline] + pub fn put(&mut self, value: u64) { + // This function buffers BIT_PACK_GROUP_SIZE values at a time. After seeing that + // many values, it decides whether the current run should be encoded in bit-packed + // or RLE. + if self.current_value == value { + self.repeat_count += 1; + if self.repeat_count > BIT_PACK_GROUP_SIZE { + // A continuation of last value. No need to buffer. + return; + } + } else { + if self.repeat_count >= BIT_PACK_GROUP_SIZE { + // The current RLE run has ended and we've gathered enough. Flush first. + debug_assert_eq!(self.bit_packed_count, 0); + self.flush_rle_run(); + } + self.repeat_count = 1; + self.current_value = value; + } + + self.buffered_values[self.num_buffered_values] = value; + self.num_buffered_values += 1; + if self.num_buffered_values == BIT_PACK_GROUP_SIZE { + // Buffered values are full. Flush them. + debug_assert_eq!(self.bit_packed_count % BIT_PACK_GROUP_SIZE, 0); + self.flush_buffered_values(); + } + } + + #[inline] + #[allow(unused)] + pub fn buffer(&self) -> &[u8] { + self.bit_writer.buffer() + } + + #[inline] + pub fn len(&self) -> usize { + self.bit_writer.bytes_written() + } + + #[allow(unused)] + pub fn is_empty(&self) -> bool { + self.bit_writer.bytes_written() == 0 + } + + #[inline] + pub fn consume(mut self) -> Vec { + self.flush(); + self.bit_writer.consume() + } + + /// Borrow equivalent of the `consume` method. + /// Call `clear()` after invoking this method. + #[inline] + pub fn flush_buffer(&mut self) -> &[u8] { + self.flush(); + self.bit_writer.flush_buffer() + } + + /// Like `flush_buffer`, but returns mutable access to the internal buffer. + /// Call `clear()` after invoking this method. + #[inline] + pub fn flush_buffer_mut(&mut self) -> &mut [u8] { + self.flush(); + self.bit_writer.flush_buffer_mut() + } + + /// Clears the internal state so this encoder can be reused (e.g., after becoming + /// full). + #[inline] + pub fn clear(&mut self) { + self.bit_writer.clear(); + self.num_buffered_values = 0; + self.current_value = 0; + self.repeat_count = 0; + self.bit_packed_count = 0; + self.indicator_byte_pos = -1; + } + + /// Advances the buffer by `num_bytes` zero bytes, delegating to the + /// underlying [`BitWriter::skip`]. + #[inline] + pub fn skip(&mut self, num_bytes: usize) { + self.bit_writer.skip(num_bytes); + } + + /// Flushes all remaining values and return the final byte buffer maintained by the + /// internal writer. + #[inline] + pub fn flush(&mut self) { + if self.bit_packed_count > 0 || self.repeat_count > 0 || self.num_buffered_values > 0 { + let all_repeat = self.bit_packed_count == 0 + && (self.repeat_count == self.num_buffered_values || self.num_buffered_values == 0); + if self.repeat_count > 0 && all_repeat { + self.flush_rle_run(); + } else { + // Buffer the last group of bit-packed values to BIT_PACK_GROUP_SIZE by padding with 0s. + if self.num_buffered_values > 0 { + while self.num_buffered_values < BIT_PACK_GROUP_SIZE { + self.buffered_values[self.num_buffered_values] = 0; + self.num_buffered_values += 1; + } + } + self.bit_packed_count += self.num_buffered_values; + self.flush_bit_packed_run(true); + self.repeat_count = 0; + } + } + } + + fn flush_rle_run(&mut self) { + debug_assert!(self.repeat_count > 0); + let indicator_value = self.repeat_count << 1; + self.bit_writer.put_vlq_int(indicator_value as u64); + self.bit_writer.put_aligned( + self.current_value, + bit_util::ceil(self.bit_width as usize, u8::BITS as usize), + ); + self.num_buffered_values = 0; + self.repeat_count = 0; + } + + fn flush_bit_packed_run(&mut self, end_current_run: bool) { + if self.indicator_byte_pos < 0 { + self.indicator_byte_pos = self.bit_writer.skip(1) as i64; + } + + // Write all buffered values as bit-packed literals + for v in &self.buffered_values[..self.num_buffered_values] { + self.bit_writer.put_value(*v, self.bit_width as usize); + } + self.num_buffered_values = 0; + if end_current_run { + self.finish_bit_packed_run(); + } + } + + // Called when ending a bit-packed run. Writes the indicator byte to the reserved + // position in `bit_writer` + fn finish_bit_packed_run(&mut self) { + let num_groups = self.bit_packed_count / BIT_PACK_GROUP_SIZE; + let indicator_byte = ((num_groups << 1) | 1) as u8; + self.bit_writer + .put_aligned_offset(indicator_byte, 1, self.indicator_byte_pos as usize); + self.indicator_byte_pos = -1; + self.bit_packed_count = 0; + } + + fn flush_buffered_values(&mut self) { + if self.repeat_count >= BIT_PACK_GROUP_SIZE { + // Clear buffered values as they are not needed + self.num_buffered_values = 0; + if self.bit_packed_count > 0 { + // In this case we have chosen to switch to RLE encoding. Close out the + // previous bit-packed run. + debug_assert_eq!(self.bit_packed_count % BIT_PACK_GROUP_SIZE, 0); + self.finish_bit_packed_run(); + } + return; + } + + self.bit_packed_count += self.num_buffered_values; + let num_groups = self.bit_packed_count / BIT_PACK_GROUP_SIZE; + if num_groups + 1 >= MAX_GROUPS_PER_BIT_PACKED_RUN { + // We've reached the maximum value that can be hold in a single bit-packed + // run. + debug_assert!(self.indicator_byte_pos >= 0); + self.flush_bit_packed_run(true); + } else { + self.flush_bit_packed_run(false); + } + self.repeat_count = 0; + } + + /// return the estimated memory size of this encoder. + pub(crate) fn estimated_memory_size(&self) -> usize { + self.bit_writer.estimated_memory_size() + std::mem::size_of::() + } +} + +/// Size, in number of `i32s` of buffer to use for RLE batch reading +const RLE_DECODER_INDEX_BUFFER_SIZE: usize = 1024; + +/// A RLE/Bit-Packing hybrid decoder. +pub struct RleDecoder { + // Number of bits used to encode the value. Must be between [0, 64]. + bit_width: u8, + + // Bit reader loaded with input buffer. + bit_reader: Option, + + // Buffer used when `bit_reader` is not `None`, for batch reading. + index_buf: Option>, + + // The remaining number of values in RLE for this run + rle_left: u32, + + // The remaining number of values in Bit-Packing for this run + bit_packed_left: u32, + + // The current value for the case of RLE mode + current_value: Option, +} + +impl RleDecoder { + pub fn new(bit_width: u8) -> Self { + RleDecoder { + bit_width, + rle_left: 0, + bit_packed_left: 0, + bit_reader: None, + index_buf: None, + current_value: None, + } + } + + #[inline] + pub fn set_data(&mut self, data: Bytes) -> Result<()> { + if let Some(ref mut bit_reader) = self.bit_reader { + bit_reader.reset(data); + } else { + self.bit_reader = Some(BitReader::new(data)); + } + + // Initialize decoder state. The boolean only reports whether the first run contained data, + // and `get`/`get_batch` already interpret that result to drive iteration. We only need + // errors propagated here, so the flag returned is intentionally ignored. + let _ = self.reload()?; + Ok(()) + } + + // These functions inline badly, they tend to inline and then create very large loop unrolls + // that damage L1d-cache occupancy. This results in a ~18% performance drop + #[inline(never)] + #[allow(unused)] + pub fn get(&mut self) -> Result> { + assert!(size_of::() <= size_of::()); + + while self.rle_left == 0 && self.bit_packed_left == 0 { + if !self.reload()? { + return Ok(None); + } + } + + let value = if self.rle_left > 0 { + let current_value = self + .current_value + .ok_or_else(|| general_err!("current_value should be Some"))?; + let rle_value = T::from_u64(current_value); + self.rle_left -= 1; + rle_value + } else { + // self.bit_packed_left > 0 + let bit_reader = self + .bit_reader + .as_mut() + .ok_or_else(|| general_err!("bit_reader should be Some"))?; + let bit_packed_value = bit_reader + .get_value(self.bit_width as usize) + .ok_or_else(|| eof_err!("Not enough data for 'bit_packed_value'"))?; + self.bit_packed_left -= 1; + bit_packed_value + }; + + Ok(Some(value)) + } + + #[inline(never)] + pub fn get_batch(&mut self, buffer: &mut [T]) -> Result { + assert!(size_of::() <= size_of::()); + + let mut values_read = 0; + while values_read < buffer.len() { + if self.rle_left > 0 { + let num_values = cmp::min(buffer.len() - values_read, self.rle_left as usize); + let repeated_value = T::from_u64(self.current_value.unwrap()); + buffer[values_read..values_read + num_values].fill(repeated_value); + self.rle_left -= num_values as u32; + values_read += num_values; + } else if self.bit_packed_left > 0 { + let mut num_values = + cmp::min(buffer.len() - values_read, self.bit_packed_left as usize); + let bit_reader = self + .bit_reader + .as_mut() + .ok_or_else(|| ParquetError::General("bit_reader should be set".into()))?; + + num_values = bit_reader.get_batch::( + &mut buffer[values_read..values_read + num_values], + self.bit_width as usize, + ); + if num_values == 0 { + // Handle writers which truncate the final block + self.bit_packed_left = 0; + continue; + } + self.bit_packed_left -= num_values as u32; + values_read += num_values; + } else if !self.reload()? { + break; + } + } + + Ok(values_read) + } + + #[inline(never)] + pub fn skip(&mut self, num_values: usize) -> Result { + let mut values_skipped = 0; + while values_skipped < num_values { + if self.rle_left > 0 { + let num_values = cmp::min(num_values - values_skipped, self.rle_left as usize); + self.rle_left -= num_values as u32; + values_skipped += num_values; + } else if self.bit_packed_left > 0 { + let mut num_values = + cmp::min(num_values - values_skipped, self.bit_packed_left as usize); + let bit_reader = self + .bit_reader + .as_mut() + .ok_or_else(|| general_err!("bit_reader should be set"))?; + + num_values = bit_reader.skip(num_values, self.bit_width as usize); + if num_values == 0 { + // Handle writers which truncate the final block + self.bit_packed_left = 0; + continue; + } + self.bit_packed_left -= num_values as u32; + values_skipped += num_values; + } else if !self.reload()? { + break; + } + } + + Ok(values_skipped) + } + + #[inline(never)] + pub fn get_batch_with_dict( + &mut self, + dict: &[T], + buffer: &mut [T], + max_values: usize, + ) -> Result + where + T: Default + Clone, + { + debug_assert!(buffer.len() >= max_values); + + let mut values_read = 0; + while values_read < max_values { + let index_buf = self.index_buf.get_or_insert_with(|| Box::new([0; 1024])); + + if self.rle_left > 0 { + let num_values = cmp::min(max_values - values_read, self.rle_left as usize); + let dict_idx = self.current_value.unwrap() as usize; + let dict_value = dict + .get(dict_idx) + .ok_or_else(|| { + general_err!( + "dictionary index out of bounds: the len is {} but the index is {}", + dict.len(), + dict_idx + ) + })? + .clone(); + + buffer[values_read..values_read + num_values].fill(dict_value); + + self.rle_left -= num_values as u32; + values_read += num_values; + } else if self.bit_packed_left > 0 { + let bit_reader = self + .bit_reader + .as_mut() + .ok_or_else(|| general_err!("bit_reader should be set"))?; + + loop { + let to_read = index_buf + .len() + .min(max_values - values_read) + .min(self.bit_packed_left as usize); + + if to_read == 0 { + break; + } + + let num_values = bit_reader + .get_batch::(&mut index_buf[..to_read], self.bit_width as usize); + if num_values == 0 { + // Handle writers which truncate the final block + self.bit_packed_left = 0; + break; + } + { + #[cold] + #[inline(never)] + fn oob(max_idx: u32, dict_len: usize) -> ParquetError { + general_err!( + "dictionary index out of bounds: the len is {} but the index is {}", + dict_len, + max_idx + ) + } + const CHUNK: usize = 16; + let out = &mut buffer[values_read..values_read + num_values]; + let idx = &index_buf[..num_values]; + let dict_len = dict.len(); + let mut out_chunks = out.chunks_exact_mut(CHUNK); + let idx_chunks = idx.chunks_exact(CHUNK); + for (out_chunk, idx_chunk) in out_chunks.by_ref().zip(idx_chunks) { + // u32 max-reduction instead of `.all(|&i| ..)`: `.all` + // short-circuits and blocks autovectorisation. Negative + // i32 cast to u32 becomes a large value so the bounds + // check still rejects it. + let max_idx = idx_chunk.iter().fold(0u32, |acc, &i| acc.max(i as u32)); + if (max_idx as usize) >= dict_len { + return Err(oob(max_idx, dict_len)); + } + for (b, i) in out_chunk.iter_mut().zip(idx_chunk.iter()) { + // SAFETY: all indices checked above to be in bounds + b.clone_from(unsafe { dict.get_unchecked(*i as usize) }); + } + } + for (b, i) in out_chunks + .into_remainder() + .iter_mut() + .zip(idx.chunks_exact(CHUNK).remainder().iter()) + { + let dict_idx = *i as usize; + if dict_idx >= dict_len { + return Err(oob(*i as u32, dict_len)); + } + // SAFETY: bounds checked above + b.clone_from(unsafe { dict.get_unchecked(dict_idx) }); + } + } + self.bit_packed_left -= num_values as u32; + values_read += num_values; + if num_values < to_read { + break; + } + } + } else if !self.reload()? { + break; + } + } + + Ok(values_read) + } + + #[inline] + fn reload(&mut self) -> Result { + let bit_reader = self + .bit_reader + .as_mut() + .ok_or_else(|| general_err!("bit_reader should be set"))?; + + if let Some(indicator_value) = bit_reader.get_vlq_int() { + // fastparquet adds padding to the end of pages. This is not spec-compliant + // but is handled by the C++ implementation + // + if indicator_value == 0 { + return Ok(false); + } + if indicator_value & 1 == 1 { + self.bit_packed_left = ((indicator_value >> 1) * BIT_PACK_GROUP_SIZE as i64) as u32; + } else { + self.rle_left = (indicator_value >> 1) as u32; + let value_width = bit_util::ceil(self.bit_width as usize, u8::BITS as usize); + self.current_value = bit_reader.get_aligned::(value_width); + self.current_value.ok_or_else(|| { + general_err!("parquet_data_error: not enough data for RLE decoding") + })?; + } + Ok(true) + } else { + Ok(false) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use crate::util::bit_util::ceil; + use rand::{self, Rng, SeedableRng, distr::StandardUniform, rng}; + + const MAX_WIDTH: usize = 32; + + #[test] + fn test_rle_decode_int32() { + // Test data: 0-7 with bit width 3 + // 00000011 10001000 11000110 11111010 + let data = vec![0x03, 0x88, 0xC6, 0xFA]; + let mut decoder: RleDecoder = RleDecoder::new(3); + decoder.set_data(data.into()).unwrap(); + let mut buffer = vec![0; BIT_PACK_GROUP_SIZE]; + let expected = vec![0, 1, 2, 3, 4, 5, 6, 7]; + let result = decoder.get_batch::(&mut buffer); + assert!(result.is_ok()); + assert_eq!(buffer, expected); + } + + #[test] + fn test_rle_skip_int32() { + // Test data: 0-7 with bit width 3 + // 00000011 10001000 11000110 11111010 + let data = vec![0x03, 0x88, 0xC6, 0xFA]; + let mut decoder: RleDecoder = RleDecoder::new(3); + decoder.set_data(data.into()).unwrap(); + let expected = vec![2, 3, 4, 5, 6, 7]; + let skipped = decoder.skip(2).expect("skipping values"); + assert_eq!(skipped, 2); + + let mut buffer = vec![0; 6]; + let remaining = decoder + .get_batch::(&mut buffer) + .expect("getting remaining"); + assert_eq!(remaining, 6); + assert_eq!(buffer, expected); + } + + #[test] + fn test_rle_consume_flush_buffer() { + let data = vec![1, 1, 1, 2, 2, 3, 3, 3]; + let mut encoder1 = RleEncoder::new(3, 256); + let mut encoder2 = RleEncoder::new(3, 256); + for value in data { + encoder1.put(value as u64); + encoder2.put(value as u64); + } + let res1 = encoder1.flush_buffer(); + let res2 = encoder2.consume(); + assert_eq!(res1, &res2[..]); + } + + #[test] + fn test_rle_decode_bool() { + // RLE test data: 50 1s followed by 50 0s + // 01100100 00000001 01100100 00000000 + let data1 = vec![0x64, 0x01, 0x64, 0x00]; + + // Bit-packing test data: alternating 1s and 0s, 100 total + // 100 / 8 = 13 groups + // 00011011 10101010 ... 00001010 + let data2 = vec![ + 0x1B, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0x0A, + ]; + + let mut decoder: RleDecoder = RleDecoder::new(1); + decoder.set_data(data1.into()).unwrap(); + let mut buffer = vec![false; 100]; + let mut expected = vec![]; + for i in 0..100 { + if i < 50 { + expected.push(true); + } else { + expected.push(false); + } + } + let result = decoder.get_batch::(&mut buffer); + assert!(result.is_ok()); + assert_eq!(buffer, expected); + + decoder.set_data(data2.into()).unwrap(); + let mut buffer = vec![false; 100]; + let mut expected = vec![]; + for i in 0..100 { + if i % 2 == 0 { + expected.push(false); + } else { + expected.push(true); + } + } + let result = decoder.get_batch::(&mut buffer); + assert!(result.is_ok()); + assert_eq!(buffer, expected); + } + + #[test] + fn test_rle_skip_bool() { + // RLE test data: 50 1s followed by 50 0s + // 01100100 00000001 01100100 00000000 + let data1 = vec![0x64, 0x01, 0x64, 0x00]; + + // Bit-packing test data: alternating 1s and 0s, 100 total + // 100 / 8 = 13 groups + // 00011011 10101010 ... 00001010 + let data2 = vec![ + 0x1B, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0x0A, + ]; + + let mut decoder: RleDecoder = RleDecoder::new(1); + decoder.set_data(data1.into()).unwrap(); + let mut buffer = vec![true; 50]; + let expected = vec![false; 50]; + + let skipped = decoder.skip(50).expect("skipping first 50"); + assert_eq!(skipped, 50); + let remainder = decoder + .get_batch::(&mut buffer) + .expect("getting remaining 50"); + assert_eq!(remainder, 50); + assert_eq!(buffer, expected); + + decoder.set_data(data2.into()).unwrap(); + let mut buffer = vec![false; 50]; + let mut expected = vec![]; + for i in 0..50 { + if i % 2 == 0 { + expected.push(false); + } else { + expected.push(true); + } + } + let skipped = decoder.skip(50).expect("skipping first 50"); + assert_eq!(skipped, 50); + let remainder = decoder + .get_batch::(&mut buffer) + .expect("getting remaining 50"); + assert_eq!(remainder, 50); + assert_eq!(buffer, expected); + } + + #[test] + fn test_rle_decode_with_dict_int32() { + // Test RLE encoding: 3 0s followed by 4 1s followed by 5 2s + // 00000110 00000000 00001000 00000001 00001010 00000010 + let dict = vec![10, 20, 30]; + let data = vec![0x06, 0x00, 0x08, 0x01, 0x0A, 0x02]; + let mut decoder: RleDecoder = RleDecoder::new(3); + decoder.set_data(data.into()).unwrap(); + let mut buffer = vec![0; 12]; + let expected = vec![10, 10, 10, 20, 20, 20, 20, 30, 30, 30, 30, 30]; + let result = decoder.get_batch_with_dict::(&dict, &mut buffer, 12); + assert!(result.is_ok()); + assert_eq!(buffer, expected); + + // Test bit-pack encoding: 345345345455 (2 groups: 8 and 4) + // 011 100 101 011 100 101 011 100 101 100 101 101 + // 00000011 01100011 11000111 10001110 00000011 01100101 00001011 + let dict = vec!["aaa", "bbb", "ccc", "ddd", "eee", "fff"]; + let data = vec![0x03, 0x63, 0xC7, 0x8E, 0x03, 0x65, 0x0B]; + let mut decoder: RleDecoder = RleDecoder::new(3); + decoder.set_data(data.into()).unwrap(); + let mut buffer = vec![""; 12]; + let expected = vec![ + "ddd", "eee", "fff", "ddd", "eee", "fff", "ddd", "eee", "fff", "eee", "fff", "fff", + ]; + let result = + decoder.get_batch_with_dict::<&str>(dict.as_slice(), buffer.as_mut_slice(), 12); + assert!(result.is_ok()); + assert_eq!(buffer, expected); + } + + #[test] + fn test_rle_skip_dict() { + // Test RLE encoding: 3 0s followed by 4 1s followed by 5 2s + // 00000110 00000000 00001000 00000001 00001010 00000010 + let dict = vec![10, 20, 30]; + let data = vec![0x06, 0x00, 0x08, 0x01, 0x0A, 0x02]; + let mut decoder: RleDecoder = RleDecoder::new(3); + decoder.set_data(data.into()).unwrap(); + let mut buffer = vec![0; 10]; + let expected = vec![10, 20, 20, 20, 20, 30, 30, 30, 30, 30]; + let skipped = decoder.skip(2).expect("skipping two values"); + assert_eq!(skipped, 2); + let remainder = decoder + .get_batch_with_dict::(&dict, &mut buffer, 10) + .expect("getting remainder"); + assert_eq!(remainder, 10); + assert_eq!(buffer, expected); + + // Test bit-pack encoding: 345345345455 (2 groups: 8 and 4) + // 011 100 101 011 100 101 011 100 101 100 101 101 + // 00000011 01100011 11000111 10001110 00000011 01100101 00001011 + let dict = vec!["aaa", "bbb", "ccc", "ddd", "eee", "fff"]; + let data = vec![0x03, 0x63, 0xC7, 0x8E, 0x03, 0x65, 0x0B]; + let mut decoder: RleDecoder = RleDecoder::new(3); + decoder.set_data(data.into()).unwrap(); + let mut buffer = vec![""; BIT_PACK_GROUP_SIZE]; + let expected = vec!["eee", "fff", "ddd", "eee", "fff", "eee", "fff", "fff"]; + let skipped = decoder.skip(4).expect("skipping four values"); + assert_eq!(skipped, 4); + let remainder = decoder + .get_batch_with_dict::<&str>( + dict.as_slice(), + buffer.as_mut_slice(), + BIT_PACK_GROUP_SIZE, + ) + .expect("getting remainder"); + assert_eq!(remainder, BIT_PACK_GROUP_SIZE); + assert_eq!(buffer, expected); + } + + fn validate_rle( + values: &[i64], + bit_width: u8, + expected_encoding: Option<&[u8]>, + expected_len: i32, + ) { + let buffer_len = 64 * 1024; + let mut encoder = RleEncoder::new(bit_width, buffer_len); + for v in values { + encoder.put(*v as u64) + } + let buffer: Bytes = encoder.consume().into(); + if expected_len != -1 { + assert_eq!(buffer.len(), expected_len as usize); + } + if let Some(b) = expected_encoding { + assert_eq!(buffer.as_ref(), b); + } + + // Verify read + let mut decoder = RleDecoder::new(bit_width); + decoder.set_data(buffer.clone()).unwrap(); + for v in values { + let val: i64 = decoder + .get() + .expect("get() should be OK") + .expect("get() should return more value"); + assert_eq!(val, *v); + } + + // Verify batch read + decoder.set_data(buffer).unwrap(); + let mut values_read: Vec = vec![0; values.len()]; + decoder + .get_batch(&mut values_read[..]) + .expect("get_batch() should be OK"); + assert_eq!(&values_read[..], values); + } + + #[test] + fn test_rle_specific_sequences() { + let mut expected_buffer = Vec::new(); + let mut values = vec![0; 50]; + values.resize(100, 1); + + expected_buffer.push(50 << 1); + expected_buffer.push(0); + expected_buffer.push(50 << 1); + expected_buffer.push(1); + + for width in 1..9 { + validate_rle(&values[..], width, Some(&expected_buffer[..]), 4); + } + for width in 9..MAX_WIDTH + 1 { + validate_rle( + &values[..], + width as u8, + None, + 2 * (1 + bit_util::ceil(width as i64, u8::BITS as i64) as i32), + ); + } + + // Test 100 0's and 1's alternating + values.clear(); + expected_buffer.clear(); + for i in 0..101 { + values.push(i % 2); + } + let num_groups = bit_util::ceil(100, BIT_PACK_GROUP_SIZE) as u8; + expected_buffer.push((num_groups << 1) | 1); + expected_buffer.resize( + expected_buffer.len() + 100 / BIT_PACK_GROUP_SIZE, + 0b10101010, + ); + + // For the last 4 0 and 1's, padded with 0. + expected_buffer.push(0b00001010); + validate_rle( + &values, + 1, + Some(&expected_buffer[..]), + 1 + num_groups as i32, + ); + for width in 2..MAX_WIDTH + 1 { + let num_values = bit_util::ceil(100, BIT_PACK_GROUP_SIZE) * BIT_PACK_GROUP_SIZE; + validate_rle( + &values, + width as u8, + None, + 1 + bit_util::ceil(width as i64 * num_values as i64, u8::BITS as i64) as i32, + ); + } + } + + // `validate_rle` on `num_vals` with width `bit_width`. If `value` is -1, that value + // is used, otherwise alternating values are used. + fn test_rle_values(bit_width: usize, num_vals: usize, value: i32) { + let mod_val = if bit_width == 64 { + 1 + } else { + 1u64 << bit_width + }; + let mut values: Vec = vec![]; + for v in 0..num_vals { + let val = if value == -1 { + v as i64 % mod_val as i64 + } else { + value as i64 + }; + values.push(val); + } + validate_rle(&values, bit_width as u8, None, -1); + } + + #[test] + fn test_values() { + for width in 1..MAX_WIDTH + 1 { + test_rle_values(width, 1, -1); + test_rle_values(width, 1024, -1); + test_rle_values(width, 1024, 0); + test_rle_values(width, 1024, 1); + } + } + + #[test] + fn test_truncated_rle() { + // The final bit packed run within a page may not be a multiple of 8 values + // Unfortunately the specification stores `(bit-packed-run-len) / 8` + // This means we don't necessarily know how many values are present + // and some writers may not add padding to compensate for this ambiguity + + // Bit pack encode 20 values with a bit width of 8 + let mut data: Vec = vec![ + (3 << 1) | 1, // bit-packed run of 3 * 8 + ]; + data.extend(std::iter::repeat_n(0xFF, 20)); + let data: Bytes = data.into(); + + let mut decoder = RleDecoder::new(8); + decoder.set_data(data.clone()).unwrap(); + + let mut output = vec![0_u16; 100]; + let read = decoder.get_batch(&mut output).unwrap(); + + assert_eq!(read, 20); + assert!(output.iter().take(20).all(|x| *x == 255)); + + // Reset decoder + decoder.set_data(data).unwrap(); + + let dict: Vec = (0..256).collect(); + let mut output = vec![0_u16; 100]; + let read = decoder + .get_batch_with_dict(&dict, &mut output, 100) + .unwrap(); + + assert_eq!(read, 20); + assert!(output.iter().take(20).all(|x| *x == 255)); + } + + #[test] + fn test_rle_padded() { + let values: Vec = vec![0, 1, 1, 3, 1, 0]; + let bit_width = 2; + let buffer_len = RleEncoder::max_buffer_size(bit_width, values.len()); + let mut encoder = RleEncoder::new(bit_width, buffer_len + 1); + for v in &values { + encoder.put(*v as u64) + } + + let mut buffer = encoder.consume(); + buffer.push(0); + + let mut decoder = RleDecoder::new(bit_width); + decoder.set_data(buffer.into()).unwrap(); + + // We don't always reliably know how many non-null values are contained in a page + // and so the decoder must work correctly without a precise value count + let mut actual_values: Vec = vec![0; 12]; + let r = decoder + .get_batch(&mut actual_values) + .expect("get_batch() should be OK"); + + // Should decode BIT_PACK_GROUP_SIZE values despite only encoding 6 as length of + // bit packed run is always a multiple of BIT_PACK_GROUP_SIZE + assert_eq!(r, BIT_PACK_GROUP_SIZE); + assert_eq!(actual_values[..6], values); + assert_eq!(actual_values[6], 0); + assert_eq!(actual_values[7], 0); + } + + /// The encoder enters RLE accumulation mode exactly on the 8th consecutive + /// identical value. + #[test] + fn test_is_accumulating_rle_boundary() { + let bit_width = 2; + let value = 1u64; + + // 7 identical values: not yet accumulating + let mut enc = RleEncoder::new(bit_width, 256); + for _ in 0..7 { + enc.put(value); + } + assert!( + !enc.is_accumulating_rle(value), + "should not be accumulating after 7 values" + ); + + // 8th value tips into accumulation + enc.put(value); + assert!( + enc.is_accumulating_rle(value), + "should be accumulating after 8 values" + ); + + // extend_run from that state and verify the round-trip + enc.extend_run(92); // total: 100 identical values + let encoded = enc.consume(); + + let mut dec = RleDecoder::new(bit_width); + dec.set_data(encoded.into()).unwrap(); + let mut out = vec![0i32; 100]; + let n = dec.get_batch::(&mut out).unwrap(); + assert_eq!(n, 100); + assert!(out.iter().all(|&v| v == value as i32)); + } + + #[test] + fn test_long_run() { + // This writer does not write runs longer than 504 values as this allows + // encoding the run header as a single byte + // + // This tests that the decoder correctly handles longer runs + + let mut writer = BitWriter::new(1024); + let bit_width = 1; + + // Choose a non-multiple of 8 larger than 1024 so that the length + // of the run is ambiguous, as the encoding only stores `num_values / 8` + let num_values = 2002; + + // bit-packed header + let run_bytes = ceil(num_values * bit_width, u8::BITS as usize) as u64; + writer.put_vlq_int((run_bytes << 1) | 1); + for _ in 0..run_bytes { + writer.put_aligned(0xFF_u8, 1); + } + let buffer: Bytes = writer.consume().into(); + + let mut decoder = RleDecoder::new(1); + decoder.set_data(buffer.clone()).unwrap(); + + let mut decoded: Vec = vec![0; num_values]; + let r = decoder.get_batch(&mut decoded).unwrap(); + assert_eq!(r, num_values); + assert_eq!(vec![1; num_values], decoded); + + decoder.set_data(buffer).unwrap(); + let r = decoder + .get_batch_with_dict(&[0, 23], &mut decoded, num_values) + .unwrap(); + assert_eq!(r, num_values); + assert_eq!(vec![23; num_values], decoded); + } + + #[test] + fn test_rle_specific_roundtrip() { + let bit_width = 1; + let values: Vec = vec![0, 1, 1, 1, 1, 0, 0, 0, 0, 1]; + let buffer_len = RleEncoder::max_buffer_size(bit_width, values.len()); + let mut encoder = RleEncoder::new(bit_width, buffer_len); + for v in &values { + encoder.put(*v as u64) + } + let buffer = encoder.consume(); + let mut decoder = RleDecoder::new(bit_width); + decoder.set_data(Bytes::from(buffer)).unwrap(); + let mut actual_values: Vec = vec![0; values.len()]; + decoder + .get_batch(&mut actual_values) + .expect("get_batch() should be OK"); + assert_eq!(actual_values, values); + } + + fn test_round_trip(values: &[i32], bit_width: u8) { + let buffer_len = 64 * 1024; + let mut encoder = RleEncoder::new(bit_width, buffer_len); + for v in values { + encoder.put(*v as u64) + } + + let buffer = Bytes::from(encoder.consume()); + + // Verify read + let mut decoder = RleDecoder::new(bit_width); + decoder.set_data(buffer.clone()).unwrap(); + for v in values { + let val = decoder + .get::() + .expect("get() should be OK") + .expect("get() should return value"); + assert_eq!(val, *v); + } + + // Verify batch read + let mut decoder = RleDecoder::new(bit_width); + decoder.set_data(buffer).unwrap(); + let mut values_read: Vec = vec![0; values.len()]; + decoder + .get_batch(&mut values_read[..]) + .expect("get_batch() should be OK"); + assert_eq!(&values_read[..], values); + } + + #[test] + fn test_random() { + let seed_len = 32; + let niters = 50; + let ngroups = 1000; + let max_group_size = 15; + let mut values = vec![]; + + for _ in 0..niters { + values.clear(); + let rng = rng(); + let seed_vec: Vec = rng + .sample_iter::(&StandardUniform) + .take(seed_len) + .collect(); + let mut seed = [0u8; 32]; + seed.copy_from_slice(&seed_vec[0..seed_len]); + let mut r#gen = rand::rngs::StdRng::from_seed(seed); + + let mut parity = false; + for _ in 0..ngroups { + let mut group_size = r#gen.random_range(1..20); + if group_size > max_group_size { + group_size = 1; + } + for _ in 0..group_size { + values.push(parity as i32); + } + parity = !parity; + } + let bit_width = bit_util::num_required_bits(values.len() as u64); + assert!(bit_width < 64); + test_round_trip(&values[..], bit_width); + } + } +} diff --git a/vendor/parquet/src/encryption/ciphers.rs b/vendor/parquet/src/encryption/ciphers.rs new file mode 100644 index 0000000..beb3acd --- /dev/null +++ b/vendor/parquet/src/encryption/ciphers.rs @@ -0,0 +1,226 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::errors::ParquetError; +use crate::errors::ParquetError::General; +use crate::errors::Result; +use crate::file::metadata::HeapSize; +use ring::aead::{AES_128_GCM, AES_256_GCM, Aad, LessSafeKey, NonceSequence, UnboundKey}; +use ring::rand::{SecureRandom, SystemRandom}; +use std::fmt::Debug; + +const RIGHT_TWELVE: u128 = 0x0000_0000_ffff_ffff_ffff_ffff_ffff_ffff; +pub(crate) const NONCE_LEN: usize = 12; +pub(crate) const TAG_LEN: usize = 16; +pub(crate) const SIZE_LEN: usize = 4; + +pub(crate) trait BlockDecryptor: Debug + Send + Sync + HeapSize { + fn decrypt(&self, length_and_ciphertext: &[u8], aad: &[u8]) -> Result>; + + fn compute_plaintext_tag(&self, aad: &[u8], plaintext: &[u8]) -> Result>; +} + +#[derive(Debug, Clone)] +pub(crate) struct RingGcmBlockDecryptor { + key: LessSafeKey, +} + +impl RingGcmBlockDecryptor { + /// Create a new `RingGcmBlockDecryptor` with a given key. + pub(crate) fn new(key_bytes: &[u8]) -> Result { + let algorithm = if key_bytes.len() == AES_128_GCM.key_len() { + &AES_128_GCM + } else if key_bytes.len() == AES_256_GCM.key_len() { + &AES_256_GCM + } else { + return Err(general_err!( + "Error creating RingGcmBlockDecryptor with unsupported key length: {}", + key_bytes.len() + )); + }; + let key = UnboundKey::new(algorithm, key_bytes) + .map_err(|_| general_err!("Failed to create {:?} key", algorithm))?; + + Ok(Self { + key: LessSafeKey::new(key), + }) + } +} + +impl HeapSize for RingGcmBlockDecryptor { + fn heap_size(&self) -> usize { + // Ring's LessSafeKey doesn't allocate on the heap + 0 + } +} + +impl BlockDecryptor for RingGcmBlockDecryptor { + fn decrypt(&self, length_and_ciphertext: &[u8], aad: &[u8]) -> Result> { + let mut result = Vec::with_capacity(length_and_ciphertext.len() - SIZE_LEN - NONCE_LEN); + result.extend_from_slice(&length_and_ciphertext[SIZE_LEN + NONCE_LEN..]); + + let nonce = ring::aead::Nonce::try_assume_unique_for_key( + &length_and_ciphertext[SIZE_LEN..SIZE_LEN + NONCE_LEN], + )?; + + self.key.open_in_place(nonce, Aad::from(aad), &mut result)?; + + // Truncate result to remove the tag + result.resize(result.len() - TAG_LEN, 0u8); + Ok(result) + } + + fn compute_plaintext_tag(&self, aad: &[u8], plaintext: &[u8]) -> Result> { + let mut plaintext = plaintext.to_vec(); + let nonce = &plaintext[plaintext.len() - NONCE_LEN - TAG_LEN..plaintext.len() - TAG_LEN]; + let nonce = ring::aead::Nonce::try_assume_unique_for_key(nonce)?; + let plaintext_end = plaintext.len() - NONCE_LEN - TAG_LEN; + let tag = self.key.seal_in_place_separate_tag( + nonce, + Aad::from(aad), + &mut plaintext[..plaintext_end], + )?; + Ok(tag.as_ref().to_vec()) + } +} + +pub(crate) trait BlockEncryptor: Debug + Send + Sync { + fn encrypt(&mut self, plaintext: &[u8], aad: &[u8]) -> Result>; +} + +#[derive(Debug, Clone)] +struct CounterNonce { + start: u128, + counter: u128, +} + +impl CounterNonce { + pub fn new(rng: &SystemRandom) -> Result { + let mut buf = [0; 16]; + rng.fill(&mut buf)?; + + // Since this is a random seed value, endianness doesn't matter at all, + // and we can use whatever is platform-native. + let start = u128::from_ne_bytes(buf) & RIGHT_TWELVE; + let counter = start.wrapping_add(1); + + Ok(Self { start, counter }) + } + + /// One accessor for the nonce bytes to avoid potentially flipping endianness + #[inline] + pub fn get_bytes(&self) -> [u8; NONCE_LEN] { + self.counter.to_le_bytes()[0..NONCE_LEN].try_into().unwrap() + } +} + +impl NonceSequence for CounterNonce { + fn advance(&mut self) -> Result { + // If we've wrapped around, we've exhausted this nonce sequence + if (self.counter & RIGHT_TWELVE) == (self.start & RIGHT_TWELVE) { + Err(ring::error::Unspecified) + } else { + // Otherwise, just advance and return the new value + let buf: [u8; NONCE_LEN] = self.get_bytes(); + self.counter = self.counter.wrapping_add(1); + Ok(ring::aead::Nonce::assume_unique_for_key(buf)) + } + } +} + +#[derive(Debug, Clone)] +pub(crate) struct RingGcmBlockEncryptor { + key: LessSafeKey, + nonce_sequence: CounterNonce, +} + +impl RingGcmBlockEncryptor { + /// Create a new `RingGcmBlockEncryptor` with a given key and random nonce. + /// The nonce will advance appropriately with each block encryption and + /// return an error if it wraps around. + pub(crate) fn new(key_bytes: &[u8]) -> Result { + let rng = SystemRandom::new(); + let algorithm = if key_bytes.len() == AES_128_GCM.key_len() { + &AES_128_GCM + } else if key_bytes.len() == AES_256_GCM.key_len() { + &AES_256_GCM + } else { + return Err(general_err!( + "Error creating RingGcmBlockEncryptor with unsupported key length: {}", + key_bytes.len() + )); + }; + + let key = UnboundKey::new(algorithm, key_bytes) + .map_err(|e| general_err!("Error creating {:?} key: {}", algorithm, e))?; + let nonce = CounterNonce::new(&rng)?; + + Ok(Self { + key: LessSafeKey::new(key), + nonce_sequence: nonce, + }) + } +} + +impl BlockEncryptor for RingGcmBlockEncryptor { + fn encrypt(&mut self, plaintext: &[u8], aad: &[u8]) -> Result> { + // Create encrypted buffer. + // Format is: [ciphertext size, nonce, ciphertext, authentication tag] + let ciphertext_length: u32 = (NONCE_LEN + plaintext.len() + TAG_LEN) + .try_into() + .map_err(|err| General(format!("Plaintext data too long. {err:?}")))?; + // Not checking for overflow here because we've already checked for it with ciphertext_length + let mut ciphertext = Vec::with_capacity(SIZE_LEN + ciphertext_length as usize); + ciphertext.extend((ciphertext_length).to_le_bytes()); + + let nonce = self.nonce_sequence.advance()?; + ciphertext.extend(nonce.as_ref()); + ciphertext.extend(plaintext); + + let tag = self.key.seal_in_place_separate_tag( + nonce, + Aad::from(aad), + &mut ciphertext[SIZE_LEN + NONCE_LEN..], + )?; + + ciphertext.extend(tag.as_ref()); + + debug_assert_eq!(SIZE_LEN + ciphertext_length as usize, ciphertext.len()); + + Ok(ciphertext) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_round_trip() { + let key = [0u8; 16]; + let mut encryptor = RingGcmBlockEncryptor::new(&key).unwrap(); + let decryptor = RingGcmBlockDecryptor::new(&key).unwrap(); + + let plaintext = b"hello, world!"; + let aad = b"some aad"; + + let ciphertext = encryptor.encrypt(plaintext, aad).unwrap(); + let decrypted = decryptor.decrypt(&ciphertext, aad).unwrap(); + + assert_eq!(plaintext, decrypted.as_slice()); + } +} diff --git a/vendor/parquet/src/encryption/decrypt.rs b/vendor/parquet/src/encryption/decrypt.rs new file mode 100644 index 0000000..3e42664 --- /dev/null +++ b/vendor/parquet/src/encryption/decrypt.rs @@ -0,0 +1,667 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Configuration and utilities for decryption of files using Parquet Modular Encryption + +use crate::encryption::ciphers::{BlockDecryptor, RingGcmBlockDecryptor, TAG_LEN}; +use crate::encryption::modules::{ModuleType, create_footer_aad, create_module_aad}; +use crate::errors::{ParquetError, Result}; +use crate::file::column_crypto_metadata::ColumnCryptoMetaData; +use crate::file::metadata::HeapSize; +use std::borrow::Cow; +use std::collections::HashMap; +use std::fmt::Formatter; +use std::io::Read; +use std::sync::Arc; + +/// Trait for retrieving an encryption key using the key's metadata +/// +/// # Example +/// +/// This shows how you might use a `KeyRetriever` to decrypt a Parquet file +/// if you have a set of known encryption keys with identifiers, but at read time +/// you may not know which columns were encrypted and which keys were used. +/// +/// In practice, the key metadata might instead store an encrypted key that must +/// be decrypted with a Key Management Server. +/// +/// ``` +/// # use std::collections::HashMap; +/// # use std::sync::{Arc, Mutex}; +/// # use parquet::encryption::decrypt::{FileDecryptionProperties, KeyRetriever}; +/// # use parquet::encryption::encrypt::FileEncryptionProperties; +/// # use parquet::errors::ParquetError; +/// // Define known encryption keys +/// let mut keys = HashMap::new(); +/// keys.insert("kf".to_owned(), b"0123456789012345".to_vec()); +/// keys.insert("kc1".to_owned(), b"1234567890123450".to_vec()); +/// keys.insert("kc2".to_owned(), b"1234567890123451".to_vec()); +/// +/// // Create encryption properties for writing a file, +/// // and specify the key identifiers as the key metadata. +/// let encryption_properties = FileEncryptionProperties::builder(keys.get("kf").unwrap().clone()) +/// .with_footer_key_metadata("kf".into()) +/// .with_column_key_and_metadata("x", keys.get("kc1").unwrap().clone(), "kc1".as_bytes().into()) +/// .with_column_key_and_metadata("y", keys.get("kc2").unwrap().clone(), "kc2".as_bytes().into()) +/// .build()?; +/// +/// // Write an encrypted file with the properties +/// // ... +/// +/// // Define a KeyRetriever that can get encryption keys using their identifiers +/// struct CustomKeyRetriever { +/// keys: Mutex>>, +/// } +/// +/// impl KeyRetriever for CustomKeyRetriever { +/// fn retrieve_key(&self, key_metadata: &[u8]) -> parquet::errors::Result> { +/// // Metadata is bytes, so convert it to a string identifier +/// let key_metadata = std::str::from_utf8(key_metadata).map_err(|e| { +/// ParquetError::General(format!("Could not convert key metadata to string: {e}")) +/// })?; +/// // Lookup the key +/// let keys = self.keys.lock().unwrap(); +/// match keys.get(key_metadata) { +/// Some(key) => Ok(key.clone()), +/// None => Err(ParquetError::General(format!( +/// "Could not retrieve key for metadata {key_metadata:?}" +/// ))), +/// } +/// } +/// } +/// +/// let key_retriever = Arc::new(CustomKeyRetriever { +/// keys: Mutex::new(keys), +/// }); +/// +/// // Create decryption properties for reading an encrypted file. +/// // Note that we don't need to specify which columns are encrypted, +/// // this is determined by the file metadata, and the required keys will be retrieved +/// // dynamically using our key retriever. +/// let decryption_properties = FileDecryptionProperties::with_key_retriever(key_retriever) +/// .build()?; +/// +/// // Read an encrypted file with the decryption properties +/// // ... +/// +/// # Ok::<(), parquet::errors::ParquetError>(()) +/// ``` +pub trait KeyRetriever: Send + Sync { + /// Retrieve a decryption key given the key metadata + fn retrieve_key(&self, key_metadata: &[u8]) -> Result>; +} + +pub(crate) fn read_and_decrypt( + decryptor: &Arc, + input: &mut T, + aad: &[u8], +) -> Result> { + let mut len_bytes = [0; 4]; + input.read_exact(&mut len_bytes)?; + let ciphertext_len = u32::from_le_bytes(len_bytes) as usize; + let mut ciphertext = vec![0; 4 + ciphertext_len]; + input.read_exact(&mut ciphertext[4..])?; + + decryptor.decrypt(&ciphertext, aad.as_ref()) +} + +// CryptoContext is a data structure that holds the context required to +// decrypt parquet modules (data pages, dictionary pages, etc.). +#[derive(Debug, Clone)] +pub(crate) struct CryptoContext { + pub(crate) row_group_idx: usize, + pub(crate) column_ordinal: usize, + pub(crate) page_ordinal: Option, + pub(crate) dictionary_page: bool, + // We have separate data and metadata decryptors because + // in GCM CTR mode, the metadata and data pages use + // different algorithms. + data_decryptor: Arc, + metadata_decryptor: Arc, + file_aad: Vec, +} + +impl CryptoContext { + pub(crate) fn for_column( + file_decryptor: &FileDecryptor, + column_crypto_metadata: &ColumnCryptoMetaData, + row_group_idx: usize, + column_ordinal: usize, + ) -> Result { + let (data_decryptor, metadata_decryptor) = match column_crypto_metadata { + ColumnCryptoMetaData::ENCRYPTION_WITH_FOOTER_KEY => { + // TODO: In GCM-CTR mode will this need to be a non-GCM decryptor? + let data_decryptor = file_decryptor.get_footer_decryptor()?; + let metadata_decryptor = file_decryptor.get_footer_decryptor()?; + (data_decryptor, metadata_decryptor) + } + ColumnCryptoMetaData::ENCRYPTION_WITH_COLUMN_KEY(column_key_encryption) => { + let key_metadata = &column_key_encryption.key_metadata; + let full_column_name; + let column_name = if column_key_encryption.path_in_schema.len() == 1 { + &column_key_encryption.path_in_schema[0] + } else { + full_column_name = column_key_encryption.path_in_schema.join("."); + &full_column_name + }; + let data_decryptor = file_decryptor + .get_column_data_decryptor(column_name, key_metadata.as_deref())?; + let metadata_decryptor = file_decryptor + .get_column_metadata_decryptor(column_name, key_metadata.as_deref())?; + (data_decryptor, metadata_decryptor) + } + }; + + Ok(CryptoContext { + row_group_idx, + column_ordinal, + page_ordinal: None, + dictionary_page: false, + data_decryptor, + metadata_decryptor, + file_aad: file_decryptor.file_aad().clone(), + }) + } + + pub(crate) fn with_page_ordinal(&self, page_ordinal: usize) -> Self { + Self { + row_group_idx: self.row_group_idx, + column_ordinal: self.column_ordinal, + page_ordinal: Some(page_ordinal), + dictionary_page: false, + data_decryptor: self.data_decryptor.clone(), + metadata_decryptor: self.metadata_decryptor.clone(), + file_aad: self.file_aad.clone(), + } + } + + pub(crate) fn create_page_header_aad(&self) -> Result> { + let module_type = if self.dictionary_page { + ModuleType::DictionaryPageHeader + } else { + ModuleType::DataPageHeader + }; + + create_module_aad( + self.file_aad(), + module_type, + self.row_group_idx, + self.column_ordinal, + self.page_ordinal, + ) + } + + pub(crate) fn create_page_aad(&self) -> Result> { + let module_type = if self.dictionary_page { + ModuleType::DictionaryPage + } else { + ModuleType::DataPage + }; + + create_module_aad( + self.file_aad(), + module_type, + self.row_group_idx, + self.column_ordinal, + self.page_ordinal, + ) + } + + pub(crate) fn create_column_index_aad(&self) -> Result> { + create_module_aad( + self.file_aad(), + ModuleType::ColumnIndex, + self.row_group_idx, + self.column_ordinal, + self.page_ordinal, + ) + } + + pub(crate) fn create_offset_index_aad(&self) -> Result> { + create_module_aad( + self.file_aad(), + ModuleType::OffsetIndex, + self.row_group_idx, + self.column_ordinal, + self.page_ordinal, + ) + } + + pub(crate) fn for_dictionary_page(&self) -> Self { + Self { + row_group_idx: self.row_group_idx, + column_ordinal: self.column_ordinal, + page_ordinal: self.page_ordinal, + dictionary_page: true, + data_decryptor: self.data_decryptor.clone(), + metadata_decryptor: self.metadata_decryptor.clone(), + file_aad: self.file_aad.clone(), + } + } + + pub(crate) fn data_decryptor(&self) -> &Arc { + &self.data_decryptor + } + + pub(crate) fn metadata_decryptor(&self) -> &Arc { + &self.metadata_decryptor + } + + pub(crate) fn file_aad(&self) -> &Vec { + &self.file_aad + } +} + +#[derive(Clone, PartialEq)] +struct ExplicitDecryptionKeys { + footer_key: Vec, + column_keys: HashMap>, +} + +impl HeapSize for ExplicitDecryptionKeys { + fn heap_size(&self) -> usize { + self.footer_key.heap_size() + self.column_keys.heap_size() + } +} + +#[derive(Clone)] +enum DecryptionKeys { + Explicit(ExplicitDecryptionKeys), + ViaRetriever(Arc), +} + +impl PartialEq for DecryptionKeys { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (DecryptionKeys::Explicit(keys), DecryptionKeys::Explicit(other_keys)) => { + keys.footer_key == other_keys.footer_key + && keys.column_keys == other_keys.column_keys + } + (DecryptionKeys::ViaRetriever(_), DecryptionKeys::ViaRetriever(_)) => true, + _ => false, + } + } +} + +impl HeapSize for DecryptionKeys { + fn heap_size(&self) -> usize { + match self { + Self::Explicit(keys) => keys.heap_size(), + Self::ViaRetriever(_) => { + // The retriever is a user-defined type we don't control, + // so we can't determine the heap size. + 0 + } + } + } +} + +/// `FileDecryptionProperties` hold keys and AAD data required to decrypt a Parquet file. +/// +/// When reading Arrow data, the `FileDecryptionProperties` should be included in the +/// [`ArrowReaderOptions`](crate::arrow::arrow_reader::ArrowReaderOptions) using +/// [`with_file_decryption_properties`](crate::arrow::arrow_reader::ArrowReaderOptions::with_file_decryption_properties). +/// +/// # Examples +/// +/// Create `FileDecryptionProperties` for a file encrypted with uniform encryption, +/// where all metadata and data are encrypted with the footer key: +/// ``` +/// # use parquet::encryption::decrypt::FileDecryptionProperties; +/// let file_encryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into()) +/// .build()?; +/// # Ok::<(), parquet::errors::ParquetError>(()) +/// ``` +/// +/// Create properties for a file where columns are encrypted with different keys: +/// ``` +/// # use parquet::encryption::decrypt::FileDecryptionProperties; +/// let file_encryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into()) +/// .with_column_key("x", b"1234567890123450".into()) +/// .with_column_key("y", b"1234567890123451".into()) +/// .build()?; +/// # Ok::<(), parquet::errors::ParquetError>(()) +/// ``` +/// +/// Specify additional authenticated data, used to protect against data replacement. +/// This must match the AAD prefix provided when the file was written, otherwise +/// data decryption will fail. +/// ``` +/// # use parquet::encryption::decrypt::FileDecryptionProperties; +/// let file_encryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into()) +/// .with_aad_prefix("example_file".into()) +/// .build()?; +/// # Ok::<(), parquet::errors::ParquetError>(()) +/// ``` +#[derive(Clone, PartialEq)] +pub struct FileDecryptionProperties { + keys: DecryptionKeys, + aad_prefix: Option>, + footer_signature_verification: bool, +} + +impl HeapSize for FileDecryptionProperties { + fn heap_size(&self) -> usize { + self.keys.heap_size() + self.aad_prefix.heap_size() + } +} +impl FileDecryptionProperties { + /// Returns a new [`FileDecryptionProperties`] builder that will use the provided key to + /// decrypt footer metadata. + pub fn builder(footer_key: Vec) -> DecryptionPropertiesBuilder { + DecryptionPropertiesBuilder::new(footer_key) + } + + /// Returns a new [`FileDecryptionProperties`] builder that uses a [`KeyRetriever`] + /// to get decryption keys based on key metadata. + pub fn with_key_retriever( + key_retriever: Arc, + ) -> DecryptionPropertiesBuilderWithRetriever { + DecryptionPropertiesBuilderWithRetriever::new(key_retriever) + } + + /// AAD prefix string uniquely identifies the file and prevents file swapping + pub fn aad_prefix(&self) -> Option<&Vec> { + self.aad_prefix.as_ref() + } + + /// Returns true if footer signature verification is enabled for files with plaintext footers. + pub fn check_plaintext_footer_integrity(&self) -> bool { + self.footer_signature_verification + } + + /// Get the encryption key for decrypting a file's footer, + /// and also column data if uniform encryption is used. + pub fn footer_key(&self, key_metadata: Option<&[u8]>) -> Result>> { + match &self.keys { + DecryptionKeys::Explicit(keys) => Ok(Cow::Borrowed(&keys.footer_key)), + DecryptionKeys::ViaRetriever(retriever) => { + let key = retriever.retrieve_key(key_metadata.unwrap_or_default())?; + Ok(Cow::Owned(key)) + } + } + } + + /// Get the column-specific encryption key for decrypting column data and metadata within a file + pub fn column_key( + &self, + column_name: &str, + key_metadata: Option<&[u8]>, + ) -> Result>> { + match &self.keys { + DecryptionKeys::Explicit(keys) => match keys.column_keys.get(column_name) { + None => Err(general_err!( + "No column decryption key set for encrypted column '{}'", + column_name + )), + Some(key) => Ok(Cow::Borrowed(key)), + }, + DecryptionKeys::ViaRetriever(retriever) => { + let key = retriever.retrieve_key(key_metadata.unwrap_or_default())?; + Ok(Cow::Owned(key)) + } + } + } + + /// Get the column names and associated decryption keys that have been configured. + /// If a key retriever is used rather than explicit decryption keys, the result + /// will be empty. + /// Provided for testing consumer code. + pub fn column_keys(&self) -> (Vec, Vec>) { + let mut column_names: Vec = Vec::new(); + let mut column_keys: Vec> = Vec::new(); + if let DecryptionKeys::Explicit(keys) = &self.keys { + for (key, value) in keys.column_keys.iter() { + column_names.push(key.clone()); + column_keys.push(value.clone()); + } + } + (column_names, column_keys) + } + + /// Whether these decryption properties use a key retriever. + /// When false, explicit keys were provided up front and can + /// be retrieved without providing key metadata, rather than + /// resolved on demand. + pub fn uses_key_retriever(&self) -> bool { + matches!(self.keys, DecryptionKeys::ViaRetriever(_)) + } +} + +impl std::fmt::Debug for FileDecryptionProperties { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "FileDecryptionProperties {{ }}") + } +} + +/// Builder for [`FileDecryptionProperties`] +/// +/// See [`FileDecryptionProperties`] for example usage. +pub struct DecryptionPropertiesBuilder { + footer_key: Vec, + column_keys: HashMap>, + aad_prefix: Option>, + footer_signature_verification: bool, +} + +impl DecryptionPropertiesBuilder { + /// Create a new [`DecryptionPropertiesBuilder`] builder that will use the provided key to + /// decrypt footer metadata. + pub fn new(footer_key: Vec) -> DecryptionPropertiesBuilder { + Self { + footer_key, + column_keys: HashMap::default(), + aad_prefix: None, + footer_signature_verification: true, + } + } + + /// Finalize the builder and return created [`FileDecryptionProperties`] + pub fn build(self) -> Result> { + let keys = DecryptionKeys::Explicit(ExplicitDecryptionKeys { + footer_key: self.footer_key, + column_keys: self.column_keys, + }); + Ok(Arc::new(FileDecryptionProperties { + keys, + aad_prefix: self.aad_prefix, + footer_signature_verification: self.footer_signature_verification, + })) + } + + /// Specify the expected AAD prefix to be used for decryption. + /// This must be set if the file was written with an AAD prefix and the + /// prefix is not stored in the file metadata. + pub fn with_aad_prefix(mut self, value: Vec) -> Self { + self.aad_prefix = Some(value); + self + } + + /// Specify the decryption key to use for a column + pub fn with_column_key(mut self, column_name: &str, decryption_key: Vec) -> Self { + self.column_keys + .insert(column_name.to_string(), decryption_key); + self + } + + /// Specify multiple column decryption keys + pub fn with_column_keys(mut self, column_names: Vec<&str>, keys: Vec>) -> Result { + if column_names.len() != keys.len() { + return Err(general_err!( + "The number of column names ({}) does not match the number of keys ({})", + column_names.len(), + keys.len() + )); + } + for (column_name, key) in column_names.into_iter().zip(keys) { + self.column_keys.insert(column_name.to_string(), key); + } + Ok(self) + } + + /// Disable verification of footer tags for files that use plaintext footers. + /// Signature verification is enabled by default. + pub fn disable_footer_signature_verification(mut self) -> Self { + self.footer_signature_verification = false; + self + } +} + +/// Builder for [`FileDecryptionProperties`] that uses a [`KeyRetriever`] +/// +/// See the [`KeyRetriever`] documentation for example usage. +pub struct DecryptionPropertiesBuilderWithRetriever { + key_retriever: Arc, + aad_prefix: Option>, + footer_signature_verification: bool, +} + +impl DecryptionPropertiesBuilderWithRetriever { + /// Create a new [`DecryptionPropertiesBuilderWithRetriever`] by providing a [`KeyRetriever`] that + /// can be used to get decryption keys based on key metadata. + pub fn new(key_retriever: Arc) -> DecryptionPropertiesBuilderWithRetriever { + Self { + key_retriever, + aad_prefix: None, + footer_signature_verification: true, + } + } + + /// Finalize the builder and return created [`FileDecryptionProperties`] + pub fn build(self) -> Result> { + let keys = DecryptionKeys::ViaRetriever(self.key_retriever); + Ok(Arc::new(FileDecryptionProperties { + keys, + aad_prefix: self.aad_prefix, + footer_signature_verification: self.footer_signature_verification, + })) + } + + /// Specify the expected AAD prefix to be used for decryption. + /// This must be set if the file was written with an AAD prefix and the + /// prefix is not stored in the file metadata. + pub fn with_aad_prefix(mut self, value: Vec) -> Self { + self.aad_prefix = Some(value); + self + } + + /// Disable verification of footer tags for files that use plaintext footers. + /// Signature verification is enabled by default. + pub fn disable_footer_signature_verification(mut self) -> Self { + self.footer_signature_verification = false; + self + } +} + +#[derive(Clone, Debug)] +pub(crate) struct FileDecryptor { + decryption_properties: Arc, + footer_decryptor: Arc, + file_aad: Vec, +} + +impl PartialEq for FileDecryptor { + fn eq(&self, other: &Self) -> bool { + self.decryption_properties == other.decryption_properties && self.file_aad == other.file_aad + } +} + +/// Estimate the size in bytes required for the file decryptor. +/// This is important to track the memory usage of cached Parquet meta data, +/// and is used via [`crate::file::metadata::ParquetMetaData::memory_size`]. +/// Note that when a [`KeyRetriever`] is used, its heap size won't be included +/// and the result will be an underestimate. +/// If the [`FileDecryptionProperties`] are shared between multiple files then the +/// heap size may also be an overestimate. +impl HeapSize for FileDecryptor { + fn heap_size(&self) -> usize { + self.decryption_properties.heap_size() + + (Arc::clone(&self.footer_decryptor) as Arc).heap_size() + + self.file_aad.heap_size() + } +} + +impl FileDecryptor { + pub(crate) fn new( + decryption_properties: &Arc, + footer_key_metadata: Option<&[u8]>, + aad_file_unique: Vec, + aad_prefix: Vec, + ) -> Result { + let file_aad = [aad_prefix.as_slice(), aad_file_unique.as_slice()].concat(); + let footer_key = decryption_properties.footer_key(footer_key_metadata)?; + let footer_decryptor = RingGcmBlockDecryptor::new(&footer_key).map_err(|e| { + general_err!( + "Invalid footer key. {}", + e.to_string().replace("Parquet error: ", "") + ) + })?; + + Ok(Self { + footer_decryptor: Arc::new(footer_decryptor), + decryption_properties: Arc::clone(decryption_properties), + file_aad, + }) + } + + pub(crate) fn get_footer_decryptor(&self) -> Result> { + Ok(self.footer_decryptor.clone()) + } + + /// Verify the signature of the footer + pub(crate) fn verify_plaintext_footer_signature(&self, plaintext_footer: &[u8]) -> Result<()> { + // Plaintext footer format is: [plaintext metadata, nonce, authentication tag] + let tag = &plaintext_footer[plaintext_footer.len() - TAG_LEN..]; + let aad = create_footer_aad(self.file_aad())?; + let footer_decryptor = self.get_footer_decryptor()?; + + let computed_tag = footer_decryptor.compute_plaintext_tag(&aad, plaintext_footer)?; + + if computed_tag != tag { + return Err(general_err!( + "Footer signature verification failed. Computed: {:?}, Expected: {:?}", + computed_tag, + tag + )); + } + Ok(()) + } + + pub(crate) fn get_column_data_decryptor( + &self, + column_name: &str, + key_metadata: Option<&[u8]>, + ) -> Result> { + let column_key = self + .decryption_properties + .column_key(column_name, key_metadata)?; + Ok(Arc::new(RingGcmBlockDecryptor::new(&column_key)?)) + } + + pub(crate) fn get_column_metadata_decryptor( + &self, + column_name: &str, + key_metadata: Option<&[u8]>, + ) -> Result> { + // Once GCM CTR mode is implemented, data and metadata decryptors may be different + self.get_column_data_decryptor(column_name, key_metadata) + } + + pub(crate) fn file_aad(&self) -> &Vec { + &self.file_aad + } +} diff --git a/vendor/parquet/src/encryption/encrypt.rs b/vendor/parquet/src/encryption/encrypt.rs new file mode 100644 index 0000000..d69e3c0 --- /dev/null +++ b/vendor/parquet/src/encryption/encrypt.rs @@ -0,0 +1,437 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Configuration and utilities for Parquet Modular Encryption + +use crate::encryption::ciphers::{ + BlockEncryptor, NONCE_LEN, RingGcmBlockEncryptor, SIZE_LEN, TAG_LEN, +}; +use crate::errors::{ParquetError, Result}; +use crate::file::column_crypto_metadata::{ColumnCryptoMetaData, EncryptionWithColumnKey}; +use crate::parquet_thrift::{ThriftCompactOutputProtocol, WriteThrift}; +use crate::schema::types::{ColumnDescPtr, SchemaDescriptor}; +use ring::rand::{SecureRandom, SystemRandom}; +use std::collections::{HashMap, HashSet}; +use std::io::Write; +use std::sync::Arc; + +#[derive(Debug, Clone, PartialEq)] +struct EncryptionKey { + key: Vec, + key_metadata: Option>, +} + +impl EncryptionKey { + fn new(key: Vec) -> EncryptionKey { + Self { + key, + key_metadata: None, + } + } + + fn with_metadata(mut self, metadata: Vec) -> Self { + self.key_metadata = Some(metadata); + self + } + + fn key(&self) -> &Vec { + &self.key + } +} + +#[derive(Debug, Clone, PartialEq)] +/// Defines how data in a Parquet file should be encrypted +/// +/// The `FileEncryptionProperties` should be included in the [`WriterProperties`](crate::file::properties::WriterProperties) +/// used to write a file by using [`WriterPropertiesBuilder::with_file_encryption_properties`](crate::file::properties::WriterPropertiesBuilder::with_file_encryption_properties). +/// +/// # Examples +/// +/// Create `FileEncryptionProperties` for a file encrypted with uniform encryption, +/// where all metadata and data are encrypted with the footer key: +/// ``` +/// # use parquet::encryption::encrypt::FileEncryptionProperties; +/// let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into()) +/// .build()?; +/// # Ok::<(), parquet::errors::ParquetError>(()) +/// ``` +/// +/// Create properties for a file where columns are encrypted with different keys. +/// Any columns without a key specified will be unencrypted: +/// ``` +/// # use parquet::encryption::encrypt::FileEncryptionProperties; +/// let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into()) +/// .with_column_key("x", b"1234567890123450".into()) +/// .with_column_key("y", b"1234567890123451".into()) +/// .build()?; +/// # Ok::<(), parquet::errors::ParquetError>(()) +/// ``` +/// +/// Specify additional authenticated data, used to protect against data replacement. +/// This should represent the file identity: +/// ``` +/// # use parquet::encryption::encrypt::FileEncryptionProperties; +/// let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into()) +/// .with_aad_prefix("example_file".into()) +/// .build()?; +/// # Ok::<(), parquet::errors::ParquetError>(()) +/// ``` +pub struct FileEncryptionProperties { + encrypt_footer: bool, + footer_key: EncryptionKey, + column_keys: HashMap, + aad_prefix: Option>, + store_aad_prefix: bool, +} + +impl FileEncryptionProperties { + /// Create a new builder for encryption properties with the given footer encryption key + pub fn builder(footer_key: Vec) -> EncryptionPropertiesBuilder { + EncryptionPropertiesBuilder::new(footer_key) + } + + /// Should the footer be encrypted + pub fn encrypt_footer(&self) -> bool { + self.encrypt_footer + } + + /// Retrieval metadata of key used for encryption of footer and (possibly) columns + pub fn footer_key_metadata(&self) -> Option<&Vec> { + self.footer_key.key_metadata.as_ref() + } + + /// Retrieval of key used for encryption of footer and (possibly) columns + pub fn footer_key(&self) -> &Vec { + &self.footer_key.key + } + + /// Get the column names, keys, and metadata for columns to be encrypted + pub fn column_keys(&self) -> (Vec, Vec>, Vec>) { + let mut column_names: Vec = Vec::with_capacity(self.column_keys.len()); + let mut keys: Vec> = Vec::with_capacity(self.column_keys.len()); + let mut meta: Vec> = Vec::with_capacity(self.column_keys.len()); + for (key, value) in self.column_keys.iter() { + column_names.push(key.clone()); + keys.push(value.key.clone()); + if let Some(metadata) = value.key_metadata.as_ref() { + meta.push(metadata.clone()); + } + } + (column_names, keys, meta) + } + + /// AAD prefix string uniquely identifies the file and prevents file swapping + pub fn aad_prefix(&self) -> Option<&Vec> { + self.aad_prefix.as_ref() + } + + /// Should the AAD prefix be stored in the file + pub fn store_aad_prefix(&self) -> bool { + self.store_aad_prefix && self.aad_prefix.is_some() + } + + /// Checks if columns that are to be encrypted are present in schema + pub(crate) fn validate_encrypted_column_names( + &self, + schema: &SchemaDescriptor, + ) -> std::result::Result<(), ParquetError> { + let column_paths = schema + .columns() + .iter() + .map(|c| c.path().string()) + .collect::>(); + let encryption_columns = self + .column_keys + .keys() + .cloned() + .collect::>(); + if !encryption_columns.is_subset(&column_paths) { + let mut columns_missing_in_schema = encryption_columns + .difference(&column_paths) + .cloned() + .collect::>(); + columns_missing_in_schema.sort(); + return Err(ParquetError::General( + format!( + "The following columns with encryption keys specified were not found in the schema: {}", + columns_missing_in_schema.join(", ") + ) + .to_string(), + )); + } + Ok(()) + } +} + +/// Builder for [`FileEncryptionProperties`] +/// +/// See [`FileEncryptionProperties`] for example usage. +pub struct EncryptionPropertiesBuilder { + encrypt_footer: bool, + footer_key: EncryptionKey, + column_keys: HashMap, + aad_prefix: Option>, + store_aad_prefix: bool, +} + +impl EncryptionPropertiesBuilder { + /// Create a new [`EncryptionPropertiesBuilder`] with the given footer encryption key + pub fn new(footer_key: Vec) -> EncryptionPropertiesBuilder { + Self { + footer_key: EncryptionKey::new(footer_key), + column_keys: HashMap::default(), + aad_prefix: None, + encrypt_footer: true, + store_aad_prefix: false, + } + } + + /// Set if the footer should be stored in plaintext (not encrypted). Defaults to false. + pub fn with_plaintext_footer(mut self, plaintext_footer: bool) -> Self { + self.encrypt_footer = !plaintext_footer; + self + } + + /// Set retrieval metadata of key used for encryption of footer and (possibly) columns + pub fn with_footer_key_metadata(mut self, metadata: Vec) -> Self { + self.footer_key = self.footer_key.with_metadata(metadata); + self + } + + /// Set the key used for encryption of a column. Note that if no column keys are configured then + /// all columns will be encrypted with the footer key. + /// If any column keys are configured then only the columns with a key will be encrypted. + pub fn with_column_key(mut self, column_name: &str, key: Vec) -> Self { + self.column_keys + .insert(column_name.to_string(), EncryptionKey::new(key)); + self + } + + /// Set the key used for encryption of a column and its metadata. The Key's metadata field is to + /// enable file readers to recover the key. For example, the metadata can keep a serialized + /// ID of a data key. Note that if no column keys are configured then all columns + /// will be encrypted with the footer key. If any column keys are configured then only the + /// columns with a key will be encrypted. + pub fn with_column_key_and_metadata( + mut self, + column_name: &str, + key: Vec, + metadata: Vec, + ) -> Self { + self.column_keys.insert( + column_name.to_string(), + EncryptionKey::new(key).with_metadata(metadata), + ); + self + } + + /// Set the keys used for encryption of columns. Analogous to + /// with_column_key but for multiple columns. This will add column keys provided to the + /// existing column keys. If column keys were already provided for some columns, the new keys + /// will overwrite the old ones. + pub fn with_column_keys(mut self, column_names: Vec<&str>, keys: Vec>) -> Result { + if column_names.len() != keys.len() { + return Err(general_err!( + "The number of column names ({}) does not match the number of keys ({})", + column_names.len(), + keys.len() + )); + } + for (i, column_name) in column_names.into_iter().enumerate() { + self.column_keys + .insert(column_name.to_string(), EncryptionKey::new(keys[i].clone())); + } + Ok(self) + } + + /// The AAD prefix uniquely identifies the file and allows to differentiate it e.g. from + /// older versions of the file or from other partition files in the same data set (table). + /// These bytes are optionally passed by a writer upon file creation. When not specified, no + /// AAD prefix is used. + pub fn with_aad_prefix(mut self, aad_prefix: Vec) -> Self { + self.aad_prefix = Some(aad_prefix); + self + } + + /// Should the AAD prefix be stored in the file. If false, readers will need to provide the + /// AAD prefix to be able to decrypt data. Defaults to false. + pub fn with_aad_prefix_storage(mut self, store_aad_prefix: bool) -> Self { + self.store_aad_prefix = store_aad_prefix; + self + } + + /// Build the encryption properties + pub fn build(self) -> Result> { + Ok(Arc::new(FileEncryptionProperties { + encrypt_footer: self.encrypt_footer, + footer_key: self.footer_key, + column_keys: self.column_keys, + aad_prefix: self.aad_prefix, + store_aad_prefix: self.store_aad_prefix, + })) + } +} + +#[derive(Debug)] +/// The encryption configuration for a single Parquet file +pub(crate) struct FileEncryptor { + properties: Arc, + aad_file_unique: Vec, + file_aad: Vec, +} + +impl FileEncryptor { + pub(crate) fn new(properties: Arc) -> Result { + // Generate unique AAD for file + let rng = SystemRandom::new(); + let mut aad_file_unique = vec![0u8; 8]; + rng.fill(&mut aad_file_unique)?; + + let file_aad = match properties.aad_prefix.as_ref() { + None => aad_file_unique.clone(), + Some(aad_prefix) => [aad_prefix.clone(), aad_file_unique.clone()].concat(), + }; + + Ok(Self { + properties, + aad_file_unique, + file_aad, + }) + } + + /// Get the encryptor's file encryption properties + pub fn properties(&self) -> &Arc { + &self.properties + } + + /// Combined AAD prefix and suffix for the file generated + pub fn file_aad(&self) -> &[u8] { + &self.file_aad + } + + /// Unique file identifier part of AAD suffix. The full AAD suffix is generated per module by + /// concatenating aad_file_unique, module type, row group ordinal (all except + /// footer), column ordinal (all except footer) and page ordinal (data page and + /// header only). + pub fn aad_file_unique(&self) -> &Vec { + &self.aad_file_unique + } + + /// Returns whether data for the specified column should be encrypted + pub fn is_column_encrypted(&self, column_path: &str) -> bool { + if self.properties.column_keys.is_empty() { + // Uniform encryption + true + } else { + self.properties.column_keys.contains_key(column_path) + } + } + + /// Get the BlockEncryptor for the footer + pub(crate) fn get_footer_encryptor(&self) -> Result> { + Ok(Box::new(RingGcmBlockEncryptor::new( + &self.properties.footer_key.key, + )?)) + } + + /// Get the encryptor for a column. + /// Will return an error if the column is not an encrypted column. + pub(crate) fn get_column_encryptor( + &self, + column_path: &str, + ) -> Result> { + if self.properties.column_keys.is_empty() { + return self.get_footer_encryptor(); + } + match self.properties.column_keys.get(column_path) { + None => Err(general_err!("Column '{}' is not encrypted", column_path)), + Some(column_key) => Ok(Box::new(RingGcmBlockEncryptor::new(column_key.key())?)), + } + } +} + +/// Write an encrypted Thrift serializable object +pub(crate) fn encrypt_thrift_object( + object: &T, + encryptor: &mut Box, + sink: &mut W, + module_aad: &[u8], +) -> Result<()> { + let encrypted_buffer = encrypt_thrift_object_to_vec(object, encryptor, module_aad)?; + sink.write_all(&encrypted_buffer)?; + Ok(()) +} + +pub(crate) fn write_signed_plaintext_thrift_object( + object: &T, + encryptor: &mut Box, + sink: &mut W, + module_aad: &[u8], +) -> Result<()> { + let mut buffer: Vec = vec![]; + { + let mut protocol = ThriftCompactOutputProtocol::new(&mut buffer); + object.write_thrift(&mut protocol)?; + } + sink.write_all(&buffer)?; + buffer = encryptor.encrypt(buffer.as_ref(), module_aad)?; + + // Format of encrypted buffer is: [ciphertext size, nonce, ciphertext, authentication tag] + let nonce = &buffer[SIZE_LEN..SIZE_LEN + NONCE_LEN]; + let tag = &buffer[buffer.len() - TAG_LEN..]; + sink.write_all(nonce)?; + sink.write_all(tag)?; + + Ok(()) +} + +/// Encrypt a Thrift serializable object to a byte vector +pub(crate) fn encrypt_thrift_object_to_vec( + object: &T, + encryptor: &mut Box, + module_aad: &[u8], +) -> Result> { + let mut buffer: Vec = vec![]; + { + let mut unencrypted_protocol = ThriftCompactOutputProtocol::new(&mut buffer); + object.write_thrift(&mut unencrypted_protocol)?; + } + + encryptor.encrypt(buffer.as_ref(), module_aad) +} + +/// Get the crypto metadata for a column from the file encryption properties +pub(crate) fn get_column_crypto_metadata( + properties: &Arc, + column: &ColumnDescPtr, +) -> Option { + if properties.column_keys.is_empty() { + // Uniform encryption + Some(ColumnCryptoMetaData::ENCRYPTION_WITH_FOOTER_KEY) + } else { + properties + .column_keys + .get(&column.path().string()) + .map(|encryption_key| { + // Column is encrypted with a column specific key + ColumnCryptoMetaData::ENCRYPTION_WITH_COLUMN_KEY(EncryptionWithColumnKey { + path_in_schema: column.path().parts().to_vec(), + key_metadata: encryption_key.key_metadata.clone(), + }) + }) + } +} diff --git a/vendor/parquet/src/encryption/mod.rs b/vendor/parquet/src/encryption/mod.rs new file mode 100644 index 0000000..c1f4ca0 --- /dev/null +++ b/vendor/parquet/src/encryption/mod.rs @@ -0,0 +1,112 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! This module implements Parquet Modular Encryption, as described in the +//! [specification](https://github.com/apache/parquet-format/blob/master/Encryption.md). +//! +//! # Example of writing and reading an encrypted Parquet file +//! +//! ``` +//! use arrow::array::{ArrayRef, Float32Array, Int32Array, RecordBatch}; +//! use parquet::arrow::arrow_reader::{ArrowReaderOptions, ParquetRecordBatchReaderBuilder}; +//! use parquet::arrow::ArrowWriter; +//! use parquet::encryption::decrypt::FileDecryptionProperties; +//! use parquet::encryption::encrypt::FileEncryptionProperties; +//! use parquet::errors::Result; +//! use parquet::file::properties::WriterProperties; +//! use std::fs::File; +//! use std::sync::Arc; +//! use tempfile::TempDir; +//! +//! // Define 16 byte AES encryption keys to use. +//! static FOOTER_KEY: &[u8; 16] = b"0123456789012345"; +//! static COLUMN_KEY_1: &[u8; 16] = b"1234567890123450"; +//! static COLUMN_KEY_2: &[u8; 16] = b"1234567890123451"; +//! +//! let temp_dir = TempDir::new()?; +//! let file_path = temp_dir.path().join("encrypted_example.parquet"); +//! +//! // Create file encryption properties, which define how the file is encrypted. +//! // We will specify a key to encrypt the footer metadata, +//! // then separate keys for different columns. +//! // This allows fine-grained control of access to different columns within a Parquet file. +//! // Note that any columns without an encryption key specified will be left un-encrypted. +//! // If only a footer key is specified, then all columns are encrypted with the footer key. +//! let encryption_properties = FileEncryptionProperties::builder(FOOTER_KEY.into()) +//! .with_column_key("x", COLUMN_KEY_1.into()) +//! .with_column_key("y", COLUMN_KEY_2.into()) +//! // We also set an AAD prefix, which is optional. +//! // This contributes to the "additional authenticated data" that is used to verify file +//! // integrity and prevents data being swapped with data encrypted with the same key. +//! .with_aad_prefix(b"example_aad".into()) +//! // Specify that the AAD prefix is stored in the file, so readers don't need +//! // to provide it to read the data, but can optionally provide it if they want to +//! // verify file integrity. +//! .with_aad_prefix_storage(true) +//! .build()?; +//! +//! let writer_properties = WriterProperties::builder() +//! .with_file_encryption_properties(encryption_properties) +//! .build(); +//! +//! // Write the encrypted Parquet file +//! { +//! let file = File::create(&file_path)?; +//! +//! let ids = Int32Array::from(vec![0, 1, 2, 3, 4, 5]); +//! let x_vals = Float32Array::from(vec![0.0, 0.1, 0.2, 0.3, 0.4, 0.5]); +//! let y_vals = Float32Array::from(vec![1.0, 1.1, 1.2, 1.3, 1.4, 1.5]); +//! let batch = RecordBatch::try_from_iter(vec![ +//! ("id", Arc::new(ids) as ArrayRef), +//! ("x", Arc::new(x_vals) as ArrayRef), +//! ("y", Arc::new(y_vals) as ArrayRef), +//! ])?; +//! +//! let mut writer = ArrowWriter::try_new(file, batch.schema(), Some(writer_properties))?; +//! +//! writer.write(&batch)?; +//! writer.close()?; +//! } +//! +//! // In order to read the encrypted Parquet file, we need to know the encryption +//! // keys used to encrypt it. +//! // We don't need to provide the AAD prefix as it was stored in the file metadata, +//! // but we could specify it here if we wanted to verify the file hasn't been tampered with: +//! let decryption_properties = FileDecryptionProperties::builder(FOOTER_KEY.into()) +//! .with_column_key("x", COLUMN_KEY_1.into()) +//! .with_column_key("y", COLUMN_KEY_2.into()) +//! .build()?; +//! +//! let reader_options = +//! ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); +//! +//! // Read the file using the configured decryption properties +//! let file = File::open(&file_path)?; +//! +//! let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, reader_options)?; +//! let record_reader = builder.build()?; +//! for batch in record_reader { +//! let batch = batch?; +//! println!("Read batch: {batch:?}"); +//! } +//! # Ok::<(), parquet::errors::ParquetError>(()) +//! ``` + +pub(crate) mod ciphers; +pub mod decrypt; +pub mod encrypt; +pub(crate) mod modules; diff --git a/vendor/parquet/src/encryption/modules.rs b/vendor/parquet/src/encryption/modules.rs new file mode 100644 index 0000000..418c382 --- /dev/null +++ b/vendor/parquet/src/encryption/modules.rs @@ -0,0 +1,98 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::errors::ParquetError; + +#[derive(PartialEq)] +pub(crate) enum ModuleType { + Footer = 0, + ColumnMetaData = 1, + DataPage = 2, + DictionaryPage = 3, + DataPageHeader = 4, + DictionaryPageHeader = 5, + ColumnIndex = 6, + OffsetIndex = 7, + _BloomFilterHeader = 8, + _BloomFilterBitset = 9, +} + +pub fn create_footer_aad(file_aad: &[u8]) -> crate::errors::Result> { + create_module_aad(file_aad, ModuleType::Footer, 0, 0, None) +} + +pub(crate) fn create_module_aad( + file_aad: &[u8], + module_type: ModuleType, + row_group_idx: usize, + column_ordinal: usize, + page_ordinal: Option, +) -> crate::errors::Result> { + let module_buf = [module_type as u8]; + + if module_buf[0] == (ModuleType::Footer as u8) { + let mut aad = Vec::with_capacity(file_aad.len() + 1); + aad.extend_from_slice(file_aad); + aad.extend_from_slice(module_buf.as_ref()); + return Ok(aad); + } + + if row_group_idx > i16::MAX as usize { + return Err(general_err!( + "Encrypted parquet files can't have more than {} row groups: {}", + i16::MAX, + row_group_idx + )); + } + if column_ordinal > i16::MAX as usize { + return Err(general_err!( + "Encrypted parquet files can't have more than {} columns: {}", + i16::MAX, + column_ordinal + )); + } + + if module_buf[0] != (ModuleType::DataPageHeader as u8) + && module_buf[0] != (ModuleType::DataPage as u8) + { + let mut aad = Vec::with_capacity(file_aad.len() + 5); + aad.extend_from_slice(file_aad); + aad.extend_from_slice(module_buf.as_ref()); + aad.extend_from_slice((row_group_idx as i16).to_le_bytes().as_ref()); + aad.extend_from_slice((column_ordinal as i16).to_le_bytes().as_ref()); + return Ok(aad); + } + + let page_ordinal = + page_ordinal.ok_or_else(|| general_err!("Page ordinal must be set for data pages"))?; + + if page_ordinal > i16::MAX as usize { + return Err(general_err!( + "Encrypted parquet files can't have more than {} pages per column chunk: {}", + i16::MAX, + page_ordinal + )); + } + + let mut aad = Vec::with_capacity(file_aad.len() + 7); + aad.extend_from_slice(file_aad); + aad.extend_from_slice(module_buf.as_ref()); + aad.extend_from_slice((row_group_idx as i16).to_le_bytes().as_ref()); + aad.extend_from_slice((column_ordinal as i16).to_le_bytes().as_ref()); + aad.extend_from_slice((page_ordinal as i16).to_le_bytes().as_ref()); + Ok(aad) +} diff --git a/vendor/parquet/src/errors.rs b/vendor/parquet/src/errors.rs new file mode 100644 index 0000000..73b4d76 --- /dev/null +++ b/vendor/parquet/src/errors.rs @@ -0,0 +1,200 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Common Parquet errors and macros. + +use core::num::TryFromIntError; +use std::error::Error; +use std::string::FromUtf8Error; +use std::{cell, io, result, str}; + +#[cfg(feature = "arrow")] +use arrow_schema::ArrowError; + +/// Parquet error enumeration +// Note: we don't implement PartialEq as the semantics for the +// external variant are not well defined (#4469) +#[derive(Debug)] +#[non_exhaustive] +pub enum ParquetError { + /// General Parquet error. + /// Returned when code violates normal workflow of working with Parquet files. + General(String), + /// "Not yet implemented" Parquet error. + /// Returned when functionality is not yet available. + NYI(String), + /// "End of file" Parquet error. + /// Returned when IO related failures occur, e.g. when there are not enough bytes to + /// decode. + EOF(String), + #[cfg(feature = "arrow")] + /// Arrow error. + /// Returned when reading into arrow or writing from arrow. + ArrowError(String), + /// Error when the requested column index is more than the + /// number of columns in the row group + IndexOutOfBound(usize, usize), + /// An external error variant + External(Box), + /// Returned when a function needs more data to complete properly. The `usize` field indicates + /// the total number of bytes required, not the number of additional bytes. + NeedMoreData(usize), + /// Returned when a function needs more data to complete properly. + /// The `Range` indicates the range of bytes that are needed. + NeedMoreDataRange(std::ops::Range), +} + +impl std::fmt::Display for ParquetError { + fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result { + match &self { + ParquetError::General(message) => { + write!(fmt, "Parquet error: {message}") + } + ParquetError::NYI(message) => write!(fmt, "NYI: {message}"), + ParquetError::EOF(message) => write!(fmt, "EOF: {message}"), + #[cfg(feature = "arrow")] + ParquetError::ArrowError(message) => write!(fmt, "Arrow: {message}"), + ParquetError::IndexOutOfBound(index, bound) => { + write!(fmt, "Index {index} out of bound: {bound}") + } + ParquetError::External(e) => write!(fmt, "External: {e}"), + ParquetError::NeedMoreData(needed) => write!(fmt, "NeedMoreData: {needed}"), + ParquetError::NeedMoreDataRange(range) => { + write!(fmt, "NeedMoreDataRange: {}..{}", range.start, range.end) + } + } + } +} + +impl Error for ParquetError { + fn source(&self) -> Option<&(dyn Error + 'static)> { + match self { + ParquetError::External(e) => Some(e.as_ref()), + _ => None, + } + } +} + +impl From for ParquetError { + fn from(e: TryFromIntError) -> ParquetError { + ParquetError::General(format!("Integer overflow: {e}")) + } +} + +impl From for ParquetError { + fn from(e: io::Error) -> ParquetError { + ParquetError::External(Box::new(e)) + } +} + +#[cfg(any(feature = "snap", test))] +impl From for ParquetError { + fn from(e: snap::Error) -> ParquetError { + ParquetError::External(Box::new(e)) + } +} + +impl From for ParquetError { + fn from(e: cell::BorrowMutError) -> ParquetError { + ParquetError::External(Box::new(e)) + } +} + +impl From for ParquetError { + fn from(e: str::Utf8Error) -> ParquetError { + ParquetError::External(Box::new(e)) + } +} + +impl From for ParquetError { + fn from(e: FromUtf8Error) -> ParquetError { + ParquetError::External(Box::new(e)) + } +} + +#[cfg(feature = "arrow")] +impl From for ParquetError { + fn from(e: ArrowError) -> ParquetError { + ParquetError::External(Box::new(e)) + } +} + +#[cfg(feature = "object_store")] +impl From for ParquetError { + fn from(e: object_store::Error) -> ParquetError { + ParquetError::External(Box::new(e)) + } +} + +#[cfg(feature = "encryption")] +impl From for ParquetError { + fn from(e: ring::error::Unspecified) -> ParquetError { + ParquetError::External(Box::new(e)) + } +} + +/// A specialized `Result` for Parquet errors. +pub type Result = result::Result; + +// ---------------------------------------------------------------------- +// Conversion from `ParquetError` to other types of `Error`s + +impl From for io::Error { + fn from(e: ParquetError) -> Self { + io::Error::other(e) + } +} + +// ---------------------------------------------------------------------- +// Convenient macros for different errors + +macro_rules! general_err { + ($fmt:expr) => (ParquetError::General($fmt.to_owned())); + ($fmt:expr, $($args:expr),*) => (ParquetError::General(format!($fmt, $($args),*))); + ($e:expr, $fmt:expr) => (ParquetError::General($fmt.to_owned(), $e)); + ($e:ident, $fmt:expr, $($args:tt),*) => ( + ParquetError::General(&format!($fmt, $($args),*), $e)); +} + +macro_rules! nyi_err { + ($fmt:expr) => (ParquetError::NYI($fmt.to_owned())); + ($fmt:expr, $($args:expr),*) => (ParquetError::NYI(format!($fmt, $($args),*))); +} + +macro_rules! eof_err { + ($fmt:expr) => (ParquetError::EOF($fmt.to_owned())); + ($fmt:expr, $($args:expr),*) => (ParquetError::EOF(format!($fmt, $($args),*))); +} + +#[cfg(feature = "arrow")] +macro_rules! arrow_err { + ($fmt:expr) => (ParquetError::ArrowError($fmt.to_owned())); + ($fmt:expr, $($args:expr),*) => (ParquetError::ArrowError(format!($fmt, $($args),*))); + ($e:expr, $fmt:expr) => (ParquetError::ArrowError($fmt.to_owned(), $e)); + ($e:ident, $fmt:expr, $($args:tt),*) => ( + ParquetError::ArrowError(&format!($fmt, $($args),*), $e)); +} + +// ---------------------------------------------------------------------- +// Convert parquet error into other errors + +#[cfg(feature = "arrow")] +impl From for ArrowError { + fn from(p: ParquetError) -> Self { + Self::ParquetError(format!("{p}")) + } +} diff --git a/vendor/parquet/src/file/column_crypto_metadata.rs b/vendor/parquet/src/file/column_crypto_metadata.rs new file mode 100644 index 0000000..53d41e5 --- /dev/null +++ b/vendor/parquet/src/file/column_crypto_metadata.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. + +//! Column chunk encryption metadata + +use std::io::Write; + +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::HeapSize; +use crate::parquet_thrift::{ + ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol, ThriftCompactOutputProtocol, + WriteThrift, WriteThriftField, read_thrift_vec, +}; +use crate::{thrift_struct, thrift_union}; + +// define this and ColumnCryptoMetadata here so they're only defined when +// the encryption feature is enabled + +thrift_struct!( +/// Encryption metadata for a column chunk encrypted with a column-specific key +pub struct EncryptionWithColumnKey { + /// Path to the column in the Parquet schema + 1: required list path_in_schema + + /// Path to the column in the Parquet schema + 2: optional binary key_metadata +} +); + +impl HeapSize for EncryptionWithColumnKey { + fn heap_size(&self) -> usize { + self.path_in_schema.heap_size() + self.key_metadata.heap_size() + } +} + +thrift_union!( +/// ColumnCryptoMetadata for a column chunk +union ColumnCryptoMetaData { + 1: ENCRYPTION_WITH_FOOTER_KEY + 2: (EncryptionWithColumnKey) ENCRYPTION_WITH_COLUMN_KEY +} +); + +impl HeapSize for ColumnCryptoMetaData { + fn heap_size(&self) -> usize { + match self { + Self::ENCRYPTION_WITH_FOOTER_KEY => 0, + Self::ENCRYPTION_WITH_COLUMN_KEY(path) => path.heap_size(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::parquet_thrift::tests::test_roundtrip; + + #[test] + fn test_column_crypto_roundtrip() { + test_roundtrip(ColumnCryptoMetaData::ENCRYPTION_WITH_FOOTER_KEY); + + let path_in_schema = vec!["foo".to_owned(), "bar".to_owned(), "really".to_owned()]; + let key_metadata = vec![1u8; 32]; + test_roundtrip(ColumnCryptoMetaData::ENCRYPTION_WITH_COLUMN_KEY( + EncryptionWithColumnKey { + path_in_schema: path_in_schema.clone(), + key_metadata: None, + }, + )); + test_roundtrip(ColumnCryptoMetaData::ENCRYPTION_WITH_COLUMN_KEY( + EncryptionWithColumnKey { + path_in_schema, + key_metadata: Some(key_metadata), + }, + )); + } +} diff --git a/vendor/parquet/src/file/metadata/footer_tail.rs b/vendor/parquet/src/file/metadata/footer_tail.rs new file mode 100644 index 0000000..c33bc7a --- /dev/null +++ b/vendor/parquet/src/file/metadata/footer_tail.rs @@ -0,0 +1,111 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::errors::{ParquetError, Result}; +use crate::file::{FOOTER_SIZE, PARQUET_MAGIC, PARQUET_MAGIC_ENCR_FOOTER}; + +/// Parsed Parquet footer tail (last 8 bytes of a Parquet file) +/// +/// There are 8 bytes at the end of the Parquet footer with the following layout: +/// * 4 bytes for the metadata length +/// * 4 bytes for the magic bytes 'PAR1' or 'PARE' (encrypted footer) +/// +/// ```text +/// +-----+------------------+ +/// | len | 'PAR1' or 'PARE' | +/// +-----+------------------+ +/// ``` +/// +/// # Examples +/// ``` +/// # use parquet::file::metadata::FooterTail; +/// // a non encrypted footer with 28 bytes of metadata +/// let last_8_bytes: [u8; 8] = [0x1C, 0x00, 0x00, 0x00, b'P', b'A', b'R', b'1']; +/// let footer_tail = FooterTail::try_from(last_8_bytes).unwrap(); +/// assert_eq!(footer_tail.metadata_length(), 28); +/// assert_eq!(footer_tail.is_encrypted_footer(), false); +/// ``` +/// +/// ``` +/// # use parquet::file::metadata::FooterTail; +/// // an encrypted footer with 512 bytes of metadata +/// let last_8_bytes = vec![0x00, 0x02, 0x00, 0x00, b'P', b'A', b'R', b'E']; +/// let footer_tail = FooterTail::try_from(&last_8_bytes[..]).unwrap(); +/// assert_eq!(footer_tail.metadata_length(), 512); +/// assert_eq!(footer_tail.is_encrypted_footer(), true); +/// ``` +/// +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct FooterTail { + metadata_length: usize, + encrypted_footer: bool, +} + +impl FooterTail { + /// Try to decode the footer tail from the given 8 bytes + pub fn try_new(slice: &[u8; FOOTER_SIZE]) -> Result { + let magic = &slice[4..]; + let encrypted_footer = if magic == PARQUET_MAGIC_ENCR_FOOTER { + true + } else if magic == PARQUET_MAGIC { + false + } else { + return Err(general_err!("Invalid Parquet file. Corrupt footer")); + }; + // get the metadata length from the footer + let metadata_len = u32::from_le_bytes(slice[..4].try_into().unwrap()); + + Ok(FooterTail { + // u32 won't be larger than usize in most cases + metadata_length: metadata_len.try_into()?, + encrypted_footer, + }) + } + + /// The length of the footer metadata in bytes + pub fn metadata_length(&self) -> usize { + self.metadata_length + } + + /// Whether the footer metadata is encrypted + pub fn is_encrypted_footer(&self) -> bool { + self.encrypted_footer + } +} + +impl TryFrom<[u8; FOOTER_SIZE]> for FooterTail { + type Error = ParquetError; + + fn try_from(value: [u8; FOOTER_SIZE]) -> Result { + Self::try_new(&value) + } +} + +impl TryFrom<&[u8]> for FooterTail { + type Error = ParquetError; + + fn try_from(value: &[u8]) -> Result { + if value.len() != FOOTER_SIZE { + return Err(general_err!( + "Invalid footer length {}, expected {FOOTER_SIZE}", + value.len() + )); + } + let slice: &[u8; FOOTER_SIZE] = value.try_into().unwrap(); + Self::try_new(slice) + } +} diff --git a/vendor/parquet/src/file/metadata/memory.rs b/vendor/parquet/src/file/metadata/memory.rs new file mode 100644 index 0000000..8a5937c --- /dev/null +++ b/vendor/parquet/src/file/metadata/memory.rs @@ -0,0 +1,344 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Memory calculations for [`ParquetMetadata::memory_size`] +//! +//! [`ParquetMetadata::memory_size`]: crate::file::metadata::ParquetMetaData::memory_size +use crate::basic::{BoundaryOrder, ColumnOrder, CompressionCodec, Encoding, PageType}; +use crate::data_type::private::ParquetValueType; +use crate::file::metadata::{ + ColumnChunkMetaData, FileMetaData, KeyValue, PageEncodingStats, ParquetPageEncodingStats, + RowGroupMetaData, SortingColumn, +}; +use crate::file::page_index::column_index::{ + ByteArrayColumnIndex, ColumnIndex, ColumnIndexMetaData, PrimitiveColumnIndex, +}; +use crate::file::page_index::offset_index::{OffsetIndexMetaData, PageLocation}; +use crate::file::statistics::{Statistics, ValueStatistics}; +use std::collections::HashMap; +use std::sync::Arc; + +/// Trait for calculating the size of various containers +pub trait HeapSize { + /// Return the size of any bytes allocated on the heap by this object, + /// including heap memory in those structures + /// + /// Note that the size of the type itself is not included in the result -- + /// instead, that size is added by the caller (e.g. container). + fn heap_size(&self) -> usize; +} + +impl HeapSize for Vec { + fn heap_size(&self) -> usize { + let item_size = std::mem::size_of::(); + // account for the contents of the Vec + (self.capacity() * item_size) + + // add any heap allocations by contents + self.iter().map(|t| t.heap_size()).sum::() + } +} + +impl HeapSize for HashMap { + fn heap_size(&self) -> usize { + let capacity = self.capacity(); + if capacity == 0 { + return 0; + } + + // HashMap doesn't provide a way to get its heap size, so this is an approximation based on + // the behavior of hashbrown::HashMap as at version 0.16.0, and may become inaccurate + // if the implementation changes. + let key_val_size = std::mem::size_of::<(K, V)>(); + // Overhead for the control tags group, which may be smaller depending on architecture + let group_size = 16; + // 1 byte of metadata stored per bucket. + let metadata_size = 1; + + // Compute the number of buckets for the capacity. Based on hashbrown's capacity_to_buckets + let buckets = if capacity < 15 { + let min_cap = match key_val_size { + 0..=1 => 14, + 2..=3 => 7, + _ => 3, + }; + let cap = min_cap.max(capacity); + if cap < 4 { + 4 + } else if cap < 8 { + 8 + } else { + 16 + } + } else { + (capacity.saturating_mul(8) / 7).next_power_of_two() + }; + + group_size + + (buckets * (key_val_size + metadata_size)) + + self.keys().map(|k| k.heap_size()).sum::() + + self.values().map(|v| v.heap_size()).sum::() + } +} + +impl HeapSize for Arc { + fn heap_size(&self) -> usize { + // Arc stores weak and strong counts on the heap alongside an instance of T + 2 * std::mem::size_of::() + std::mem::size_of::() + self.as_ref().heap_size() + } +} + +impl HeapSize for Arc { + fn heap_size(&self) -> usize { + 2 * std::mem::size_of::() + + std::mem::size_of_val(self.as_ref()) + + self.as_ref().heap_size() + } +} + +impl HeapSize for Box { + fn heap_size(&self) -> usize { + std::mem::size_of::() + self.as_ref().heap_size() + } +} + +impl HeapSize for Option { + fn heap_size(&self) -> usize { + self.as_ref().map(|inner| inner.heap_size()).unwrap_or(0) + } +} + +impl HeapSize for String { + fn heap_size(&self) -> usize { + self.capacity() + } +} + +impl HeapSize for FileMetaData { + fn heap_size(&self) -> usize { + #[cfg(feature = "encryption")] + let encryption_heap_size = + self.encryption_algorithm.heap_size() + self.footer_signing_key_metadata.heap_size(); + #[cfg(not(feature = "encryption"))] + let encryption_heap_size = 0; + + self.created_by.heap_size() + + self.key_value_metadata.heap_size() + + self.schema_descr.heap_size() + + self.column_orders.heap_size() + + encryption_heap_size + } +} + +impl HeapSize for KeyValue { + fn heap_size(&self) -> usize { + self.key.heap_size() + self.value.heap_size() + } +} + +impl HeapSize for RowGroupMetaData { + fn heap_size(&self) -> usize { + // don't count schema_descr here because it is already + // counted in FileMetaData + self.columns.heap_size() + self.sorting_columns.heap_size() + } +} + +impl HeapSize for ColumnChunkMetaData { + fn heap_size(&self) -> usize { + #[cfg(feature = "encryption")] + let encryption_heap_size = + self.column_crypto_metadata.heap_size() + self.encrypted_column_metadata.heap_size(); + #[cfg(not(feature = "encryption"))] + let encryption_heap_size = 0; + + // don't count column_descr here because it is already counted in + // FileMetaData + self.encodings.heap_size() + + self.file_path.heap_size() + + self.compression.heap_size() + + self.statistics.heap_size() + + self.encoding_stats.heap_size() + + self.unencoded_byte_array_data_bytes.heap_size() + + self.repetition_level_histogram.heap_size() + + self.definition_level_histogram.heap_size() + + self.geo_statistics.heap_size() + + encryption_heap_size + } +} + +impl HeapSize for Encoding { + fn heap_size(&self) -> usize { + 0 // no heap allocations + } +} + +impl HeapSize for ParquetPageEncodingStats { + fn heap_size(&self) -> usize { + match self { + Self::Full(v) => v.heap_size(), + Self::Mask(_) => 0, + } + } +} + +impl HeapSize for PageEncodingStats { + fn heap_size(&self) -> usize { + self.page_type.heap_size() + self.encoding.heap_size() + } +} + +impl HeapSize for SortingColumn { + fn heap_size(&self) -> usize { + 0 // no heap allocations + } +} +impl HeapSize for CompressionCodec { + fn heap_size(&self) -> usize { + 0 // no heap allocations + } +} + +impl HeapSize for PageType { + fn heap_size(&self) -> usize { + 0 // no heap allocations + } +} + +impl HeapSize for Statistics { + fn heap_size(&self) -> usize { + match self { + Statistics::Boolean(value_statistics) => value_statistics.heap_size(), + Statistics::Int32(value_statistics) => value_statistics.heap_size(), + Statistics::Int64(value_statistics) => value_statistics.heap_size(), + Statistics::Int96(value_statistics) => value_statistics.heap_size(), + Statistics::Float(value_statistics) => value_statistics.heap_size(), + Statistics::Double(value_statistics) => value_statistics.heap_size(), + Statistics::ByteArray(value_statistics) => value_statistics.heap_size(), + Statistics::FixedLenByteArray(value_statistics) => value_statistics.heap_size(), + } + } +} + +impl HeapSize for OffsetIndexMetaData { + fn heap_size(&self) -> usize { + self.page_locations.heap_size() + self.unencoded_byte_array_data_bytes.heap_size() + } +} + +impl HeapSize for ColumnIndexMetaData { + fn heap_size(&self) -> usize { + match self { + Self::NONE => 0, + Self::BOOLEAN(native_index) => native_index.heap_size(), + Self::INT32(native_index) => native_index.heap_size(), + Self::INT64(native_index) => native_index.heap_size(), + Self::INT96(native_index) => native_index.heap_size(), + Self::FLOAT(native_index) => native_index.heap_size(), + Self::DOUBLE(native_index) => native_index.heap_size(), + Self::BYTE_ARRAY(native_index) => native_index.heap_size(), + Self::FIXED_LEN_BYTE_ARRAY(native_index) => native_index.heap_size(), + } + } +} + +impl HeapSize for ColumnIndex { + fn heap_size(&self) -> usize { + self.null_pages.heap_size() + + self.boundary_order.heap_size() + + self.null_counts.heap_size() + + self.definition_level_histograms.heap_size() + + self.repetition_level_histograms.heap_size() + } +} + +impl HeapSize for PrimitiveColumnIndex { + fn heap_size(&self) -> usize { + self.column_index.heap_size() + self.min_values.heap_size() + self.max_values.heap_size() + } +} + +impl HeapSize for ByteArrayColumnIndex { + fn heap_size(&self) -> usize { + self.column_index.heap_size() + + self.min_bytes.heap_size() + + self.min_offsets.heap_size() + + self.max_bytes.heap_size() + + self.max_offsets.heap_size() + } +} + +impl HeapSize for ValueStatistics { + fn heap_size(&self) -> usize { + self.min_opt().map(T::heap_size).unwrap_or(0) + + self.max_opt().map(T::heap_size).unwrap_or(0) + } +} +impl HeapSize for bool { + fn heap_size(&self) -> usize { + 0 // no heap allocations + } +} +impl HeapSize for u8 { + fn heap_size(&self) -> usize { + 0 // no heap allocations + } +} +impl HeapSize for i32 { + fn heap_size(&self) -> usize { + 0 // no heap allocations + } +} +impl HeapSize for i64 { + fn heap_size(&self) -> usize { + 0 // no heap allocations + } +} + +impl HeapSize for f32 { + fn heap_size(&self) -> usize { + 0 // no heap allocations + } +} +impl HeapSize for f64 { + fn heap_size(&self) -> usize { + 0 // no heap allocations + } +} + +impl HeapSize for usize { + fn heap_size(&self) -> usize { + 0 // no heap allocations + } +} + +impl HeapSize for BoundaryOrder { + fn heap_size(&self) -> usize { + 0 // no heap allocations + } +} + +impl HeapSize for PageLocation { + fn heap_size(&self) -> usize { + 0 // no heap allocations + } +} + +impl HeapSize for ColumnOrder { + fn heap_size(&self) -> usize { + 0 // no heap allocations in ColumnOrder + } +} diff --git a/vendor/parquet/src/file/metadata/mod.rs b/vendor/parquet/src/file/metadata/mod.rs new file mode 100644 index 0000000..646438d --- /dev/null +++ b/vendor/parquet/src/file/metadata/mod.rs @@ -0,0 +1,2185 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Parquet metadata API +//! +//! Users should use these structures to interact with Parquet metadata. +//! +//! * [`ParquetMetaData`]: Top level metadata container, read from the Parquet +//! file footer. +//! +//! * [`FileMetaData`]: File level metadata such as schema, row counts and +//! version. +//! +//! * [`RowGroupMetaData`]: Metadata for each Row Group with a File, such as +//! location and number of rows, and column chunks. +//! +//! * [`ColumnChunkMetaData`]: Metadata for each column chunk (primitive leaf) +//! within a Row Group including encoding and compression information, +//! number of values, statistics, etc. +//! +//! # APIs for working with Parquet Metadata +//! +//! The Parquet readers and writers in this crate handle reading and writing +//! metadata into parquet files. To work with metadata directly, +//! the following APIs are available: +//! +//! * [`ParquetMetaDataReader`] for reading metadata from an I/O source (sync and async) +//! * [`ParquetMetaDataPushDecoder`] for decoding from bytes without I/O +//! * [`ParquetMetaDataWriter`] for writing. +//! +//! # Examples +//! +//! Please see [`external_metadata.rs`] +//! +//! [`external_metadata.rs`]: https://github.com/apache/arrow-rs/tree/master/parquet/examples/external_metadata.rs +//! +//! # Metadata Encodings and Structures +//! +//! There are three different encodings of Parquet Metadata in this crate: +//! +//! 1. `bytes`:encoded with the Thrift `TCompactProtocol` as defined in +//! [parquet.thrift] +//! +//! 2. [`format`]: Rust structures automatically generated by the thrift compiler +//! from [parquet.thrift]. These structures are low level and mirror +//! the thrift definitions. +//! +//! 3. [`file::metadata`] (this module): Easier to use Rust structures +//! with a more idiomatic API. Note that, confusingly, some but not all +//! of these structures have the same name as the [`format`] structures. +//! +//! [`file::metadata`]: crate::file::metadata +//! [parquet.thrift]: https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift +//! +//! Graphically, this is how the different structures relate to each other: +//! +//! ```text +//! ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ +//! ┌──────────────┐ │ ┌───────────────────────┐ │ +//! │ │ ColumnIndex │ ││ ParquetMetaData │ +//! └──────────────┘ │ └───────────────────────┘ │ +//! ┌──────────────┐ │ ┌────────────────┐ │┌───────────────────────┐ +//! │ ..0x24.. │ ◀────▶ │ OffsetIndex │ │ ◀────▶ │ ParquetMetaData │ │ +//! └──────────────┘ │ └────────────────┘ │└───────────────────────┘ +//! ... │ ... │ +//! │ ┌──────────────────┐ │ ┌──────────────────┐ +//! bytes │ FileMetaData* │ │ │ FileMetaData* │ │ +//! (thrift encoded) │ └──────────────────┘ │ └──────────────────┘ +//! ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘ +//! +//! format::meta structures file::metadata structures +//! +//! * Same name, different struct +//! ``` +mod footer_tail; +mod memory; +mod options; +mod parser; +mod push_decoder; +pub(crate) mod reader; +pub(crate) mod thrift; +mod writer; + +use crate::basic::{ + BoundaryOrder, ColumnOrder, Compression, CompressionCodec, Encoding, EncodingMask, PageType, + Type, +}; +#[cfg(feature = "encryption")] +use crate::encryption::decrypt::FileDecryptor; +use crate::errors::{ParquetError, Result}; +#[cfg(feature = "encryption")] +use crate::file::column_crypto_metadata::ColumnCryptoMetaData; +pub(crate) use crate::file::metadata::memory::HeapSize; +#[cfg(feature = "encryption")] +use crate::file::metadata::thrift::encryption::EncryptionAlgorithm; +use crate::file::page_index::column_index::{ByteArrayColumnIndex, PrimitiveColumnIndex}; +use crate::file::page_index::{column_index::ColumnIndexMetaData, offset_index::PageLocation}; +use crate::file::statistics::Statistics; +use crate::geospatial::statistics as geo_statistics; +use crate::parquet_thrift::{ + ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol, ThriftCompactOutputProtocol, + WriteThrift, WriteThriftField, +}; +use crate::schema::types::{ + ColumnDescPtr, ColumnDescriptor, ColumnPath, SchemaDescPtr, SchemaDescriptor, + Type as SchemaType, +}; +use crate::thrift_struct; +use crate::{ + data_type::private::ParquetValueType, file::page_index::offset_index::OffsetIndexMetaData, +}; + +pub use footer_tail::FooterTail; +pub use options::{ParquetMetaDataOptions, ParquetStatisticsPolicy}; +pub use push_decoder::ParquetMetaDataPushDecoder; +pub use reader::{PageIndexPolicy, ParquetMetaDataReader}; +use std::io::Write; +use std::ops::Range; +use std::sync::Arc; +pub use writer::ParquetMetaDataWriter; +pub(crate) use writer::ThriftMetadataWriter; + +/// Page level statistics for each column chunk of each row group. +/// +/// This structure is an in-memory representation of multiple [`ColumnIndex`] +/// structures in a parquet file footer, as described in the Parquet [PageIndex +/// documentation]. Each [`ColumnIndex`] holds statistics about all the pages in a +/// particular column chunk. +/// +/// `column_index[row_group_number][column_number]` holds the +/// [`ColumnIndex`] corresponding to column `column_number` of row group +/// `row_group_number`. +/// +/// For example `column_index[2][3]` holds the [`ColumnIndex`] for the fourth +/// column in the third row group of the parquet file. +/// +/// [PageIndex documentation]: https://github.com/apache/parquet-format/blob/master/PageIndex.md +/// [`ColumnIndex`]: crate::file::page_index::column_index::ColumnIndexMetaData +pub type ParquetColumnIndex = Vec>; + +/// [`OffsetIndexMetaData`] for each data page of each row group of each column +/// +/// This structure is the parsed representation of the [`OffsetIndex`] from the +/// Parquet file footer, as described in the Parquet [PageIndex documentation]. +/// +/// `offset_index[row_group_number][column_number]` holds +/// the [`OffsetIndexMetaData`] corresponding to column +/// `column_number`of row group `row_group_number`. +/// +/// [PageIndex documentation]: https://github.com/apache/parquet-format/blob/master/PageIndex.md +/// [`OffsetIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md +pub type ParquetOffsetIndex = Vec>; + +/// Parsed metadata for a single Parquet file +/// +/// This structure is stored in the footer of Parquet files, in the format +/// defined by [`parquet.thrift`]. +/// +/// # Overview +/// The fields of this structure are: +/// * [`FileMetaData`]: Information about the overall file (such as the schema) (See [`Self::file_metadata`]) +/// * [`RowGroupMetaData`]: Information about each Row Group (see [`Self::row_groups`]) +/// * [`ParquetColumnIndex`] and [`ParquetOffsetIndex`]: Optional "Page Index" structures (see [`Self::column_index`] and [`Self::offset_index`]) +/// +/// This structure is read by the various readers in this crate or can be read +/// directly from a file using the [`ParquetMetaDataReader`] struct. +/// +/// See the [`ParquetMetaDataBuilder`] to create and modify this structure. +/// +/// [`parquet.thrift`]: https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift +#[derive(Debug, Clone, PartialEq)] +pub struct ParquetMetaData { + /// File level metadata + file_metadata: FileMetaData, + /// Row group metadata + row_groups: Vec, + /// Page level index for each page in each column chunk + column_index: Option, + /// Offset index for each page in each column chunk + offset_index: Option, + /// Optional file decryptor + #[cfg(feature = "encryption")] + file_decryptor: Option>, +} + +impl ParquetMetaData { + /// Creates Parquet metadata from file metadata and a list of row + /// group metadata + pub fn new(file_metadata: FileMetaData, row_groups: Vec) -> Self { + ParquetMetaData { + file_metadata, + row_groups, + column_index: None, + offset_index: None, + #[cfg(feature = "encryption")] + file_decryptor: None, + } + } + + /// Adds [`FileDecryptor`] to this metadata instance to enable decryption of + /// encrypted data. + #[cfg(feature = "encryption")] + pub(crate) fn with_file_decryptor(&mut self, file_decryptor: Option) { + self.file_decryptor = file_decryptor.map(Box::new); + } + + /// Convert this ParquetMetaData into a [`ParquetMetaDataBuilder`] + pub fn into_builder(self) -> ParquetMetaDataBuilder { + self.into() + } + + /// Returns file metadata as reference. + pub fn file_metadata(&self) -> &FileMetaData { + &self.file_metadata + } + + /// Returns file decryptor as reference. + #[cfg(feature = "encryption")] + pub(crate) fn file_decryptor(&self) -> Option<&FileDecryptor> { + self.file_decryptor.as_deref() + } + + /// Returns number of row groups in this file. + pub fn num_row_groups(&self) -> usize { + self.row_groups.len() + } + + /// Returns row group metadata for `i`th position. + /// Position should be less than number of row groups `num_row_groups`. + pub fn row_group(&self, i: usize) -> &RowGroupMetaData { + &self.row_groups[i] + } + + /// Returns slice of row groups in this file. + pub fn row_groups(&self) -> &[RowGroupMetaData] { + &self.row_groups + } + + /// Returns the column index for this file if loaded + /// + /// Returns `None` if the parquet file does not have a `ColumnIndex` or + /// [ArrowReaderOptions::with_page_index] was set to false. + /// + /// [ArrowReaderOptions::with_page_index]: https://docs.rs/parquet/latest/parquet/arrow/arrow_reader/struct.ArrowReaderOptions.html#method.with_page_index + pub fn column_index(&self) -> Option<&ParquetColumnIndex> { + self.column_index.as_ref() + } + + /// Returns offset indexes in this file, if loaded + /// + /// Returns `None` if the parquet file does not have a `OffsetIndex` or + /// [ArrowReaderOptions::with_page_index] was set to false. + /// + /// [ArrowReaderOptions::with_page_index]: https://docs.rs/parquet/latest/parquet/arrow/arrow_reader/struct.ArrowReaderOptions.html#method.with_page_index + pub fn offset_index(&self) -> Option<&ParquetOffsetIndex> { + self.offset_index.as_ref() + } + + /// Estimate of the bytes allocated to store `ParquetMetadata` + /// + /// # Notes: + /// + /// 1. Includes size of self + /// + /// 2. Includes heap memory for sub fields such as [`FileMetaData`] and + /// [`RowGroupMetaData`]. + /// + /// 3. Includes memory from shared pointers (e.g. [`SchemaDescPtr`]). This + /// means `memory_size` will over estimate the memory size if such pointers + /// are shared. + /// + /// 4. Does not include any allocator overheads + pub fn memory_size(&self) -> usize { + #[cfg(feature = "encryption")] + let encryption_size = self.file_decryptor.heap_size(); + #[cfg(not(feature = "encryption"))] + let encryption_size = 0usize; + + std::mem::size_of::() + + self.file_metadata.heap_size() + + self.row_groups.heap_size() + + self.column_index.heap_size() + + self.offset_index.heap_size() + + encryption_size + } + + /// Override the column index + pub(crate) fn set_column_index(&mut self, index: Option) { + self.column_index = index; + } + + /// Override the offset index + pub(crate) fn set_offset_index(&mut self, index: Option) { + self.offset_index = index; + } +} + +/// A builder for creating / manipulating [`ParquetMetaData`] +/// +/// # Example creating a new [`ParquetMetaData`] +/// +///```no_run +/// # use parquet::file::metadata::{FileMetaData, ParquetMetaData, ParquetMetaDataBuilder, RowGroupMetaData, RowGroupMetaDataBuilder}; +/// # fn get_file_metadata() -> FileMetaData { unimplemented!(); } +/// // Create a new builder given the file metadata +/// let file_metadata = get_file_metadata(); +/// // Create a row group +/// let row_group = RowGroupMetaData::builder(file_metadata.schema_descr_ptr()) +/// .set_num_rows(100) +/// // ... (A real row group needs more than just the number of rows) +/// .build() +/// .unwrap(); +/// // Create the final metadata +/// let metadata: ParquetMetaData = ParquetMetaDataBuilder::new(file_metadata) +/// .add_row_group(row_group) +/// .build(); +/// ``` +/// +/// # Example modifying an existing [`ParquetMetaData`] +/// ```no_run +/// # use parquet::file::metadata::ParquetMetaData; +/// # fn load_metadata() -> ParquetMetaData { unimplemented!(); } +/// // Modify the metadata so only the last RowGroup remains +/// let metadata: ParquetMetaData = load_metadata(); +/// let mut builder = metadata.into_builder(); +/// +/// // Take existing row groups to modify +/// let mut row_groups = builder.take_row_groups(); +/// let last_row_group = row_groups.pop().unwrap(); +/// +/// let metadata = builder +/// .add_row_group(last_row_group) +/// .build(); +/// ``` +pub struct ParquetMetaDataBuilder(ParquetMetaData); + +impl ParquetMetaDataBuilder { + /// Create a new builder from a file metadata, with no row groups + pub fn new(file_meta_data: FileMetaData) -> Self { + Self(ParquetMetaData::new(file_meta_data, vec![])) + } + + /// Create a new builder from an existing ParquetMetaData + pub fn new_from_metadata(metadata: ParquetMetaData) -> Self { + Self(metadata) + } + + /// Adds a row group to the metadata + pub fn add_row_group(mut self, row_group: RowGroupMetaData) -> Self { + self.0.row_groups.push(row_group); + self + } + + /// Sets all the row groups to the specified list + pub fn set_row_groups(mut self, row_groups: Vec) -> Self { + self.0.row_groups = row_groups; + self + } + + /// Takes ownership of the row groups in this builder, and clears the list + /// of row groups. + /// + /// This can be used for more efficient creation of a new ParquetMetaData + /// from an existing one. + pub fn take_row_groups(&mut self) -> Vec { + std::mem::take(&mut self.0.row_groups) + } + + /// Return a reference to the current row groups + pub fn row_groups(&self) -> &[RowGroupMetaData] { + &self.0.row_groups + } + + /// Sets the column index + pub fn set_column_index(mut self, column_index: Option) -> Self { + self.0.column_index = column_index; + self + } + + /// Returns the current column index from the builder, replacing it with `None` + pub fn take_column_index(&mut self) -> Option { + std::mem::take(&mut self.0.column_index) + } + + /// Return a reference to the current column index, if any + pub fn column_index(&self) -> Option<&ParquetColumnIndex> { + self.0.column_index.as_ref() + } + + /// Sets the offset index + pub fn set_offset_index(mut self, offset_index: Option) -> Self { + self.0.offset_index = offset_index; + self + } + + /// Returns the current offset index from the builder, replacing it with `None` + pub fn take_offset_index(&mut self) -> Option { + std::mem::take(&mut self.0.offset_index) + } + + /// Return a reference to the current offset index, if any + pub fn offset_index(&self) -> Option<&ParquetOffsetIndex> { + self.0.offset_index.as_ref() + } + + /// Sets the file decryptor needed to decrypt this metadata. + #[cfg(feature = "encryption")] + pub(crate) fn set_file_decryptor(mut self, file_decryptor: Option) -> Self { + self.0.with_file_decryptor(file_decryptor); + self + } + + /// Creates a new ParquetMetaData from the builder + pub fn build(self) -> ParquetMetaData { + let Self(metadata) = self; + metadata + } +} + +impl From for ParquetMetaDataBuilder { + fn from(meta_data: ParquetMetaData) -> Self { + Self(meta_data) + } +} + +thrift_struct!( +/// A key-value pair for [`FileMetaData`]. +pub struct KeyValue { + 1: required string key + 2: optional string value +} +); + +impl KeyValue { + /// Create a new key value pair + pub fn new(key: String, value: F2) -> KeyValue + where + F2: Into>, + { + KeyValue { + key, + value: value.into(), + } + } +} + +thrift_struct!( +/// PageEncodingStats for a column chunk and data page. +pub struct PageEncodingStats { + 1: required PageType page_type; + 2: required Encoding encoding; + 3: required i32 count; +} +); + +/// Internal representation of the page encoding stats in the [`ColumnChunkMetaData`]. +/// This is not publicly exposed, with different getters defined for each variant. +#[derive(Debug, Clone, PartialEq)] +enum ParquetPageEncodingStats { + /// The full array of stats as defined in the Parquet spec. + Full(Vec), + /// A condensed version of only page encodings seen. + Mask(EncodingMask), +} + +/// Reference counted pointer for [`FileMetaData`]. +pub type FileMetaDataPtr = Arc; + +/// File level metadata for a Parquet file. +/// +/// Includes the version of the file, metadata, number of rows, schema, and column orders +#[derive(Debug, Clone, PartialEq)] +pub struct FileMetaData { + version: i32, + num_rows: i64, + created_by: Option, + key_value_metadata: Option>, + schema_descr: SchemaDescPtr, + column_orders: Option>, + #[cfg(feature = "encryption")] + encryption_algorithm: Option>, + #[cfg(feature = "encryption")] + footer_signing_key_metadata: Option>, +} + +impl FileMetaData { + /// Creates new file metadata. + pub fn new( + version: i32, + num_rows: i64, + created_by: Option, + key_value_metadata: Option>, + schema_descr: SchemaDescPtr, + column_orders: Option>, + ) -> Self { + FileMetaData { + version, + num_rows, + created_by, + key_value_metadata, + schema_descr, + column_orders, + #[cfg(feature = "encryption")] + encryption_algorithm: None, + #[cfg(feature = "encryption")] + footer_signing_key_metadata: None, + } + } + + #[cfg(feature = "encryption")] + pub(crate) fn with_encryption_algorithm( + mut self, + encryption_algorithm: Option, + ) -> Self { + self.encryption_algorithm = encryption_algorithm.map(Box::new); + self + } + + #[cfg(feature = "encryption")] + pub(crate) fn with_footer_signing_key_metadata( + mut self, + footer_signing_key_metadata: Option>, + ) -> Self { + self.footer_signing_key_metadata = footer_signing_key_metadata; + self + } + + /// Returns version of this file. + pub fn version(&self) -> i32 { + self.version + } + + /// Returns number of rows in the file. + pub fn num_rows(&self) -> i64 { + self.num_rows + } + + /// String message for application that wrote this file. + /// + /// This should have the following format: + /// ` version (build )`. + /// + /// ```shell + /// parquet-mr version 1.8.0 (build 0fda28af84b9746396014ad6a415b90592a98b3b) + /// ``` + pub fn created_by(&self) -> Option<&str> { + self.created_by.as_deref() + } + + /// Returns key_value_metadata of this file. + pub fn key_value_metadata(&self) -> Option<&Vec> { + self.key_value_metadata.as_ref() + } + + /// Returns Parquet [`Type`] that describes schema in this file. + /// + /// [`Type`]: crate::schema::types::Type + pub fn schema(&self) -> &SchemaType { + self.schema_descr.root_schema() + } + + /// Returns a reference to schema descriptor. + pub fn schema_descr(&self) -> &SchemaDescriptor { + &self.schema_descr + } + + /// Returns reference counted clone for schema descriptor. + pub fn schema_descr_ptr(&self) -> SchemaDescPtr { + self.schema_descr.clone() + } + + /// Column (sort) order used for `min` and `max` values of each column in this file. + /// + /// Each column order corresponds to one column, determined by its position in the + /// list, matching the position of the column in the schema. + /// + /// When `None` is returned, there are no column orders available, and each column + /// should be assumed to have undefined (legacy) column order. + pub fn column_orders(&self) -> Option<&Vec> { + self.column_orders.as_ref() + } + + /// Returns column order for `i`th column in this file. + /// If column orders are not available, returns undefined (legacy) column order. + pub fn column_order(&self, i: usize) -> ColumnOrder { + self.column_orders + .as_ref() + .map(|data| data[i]) + .unwrap_or(ColumnOrder::UNDEFINED) + } +} + +thrift_struct!( +/// Sort order within a RowGroup of a leaf column +pub struct SortingColumn { + /// The ordinal position of the column (in this row group) + 1: required i32 column_idx + + /// If true, indicates this column is sorted in descending order. + 2: required bool descending + + /// If true, nulls will come before non-null values, otherwise, + /// nulls go at the end. */ + 3: required bool nulls_first +} +); + +/// Reference counted pointer for [`RowGroupMetaData`]. +pub type RowGroupMetaDataPtr = Arc; + +/// Metadata for a row group +/// +/// Includes [`ColumnChunkMetaData`] for each column in the row group, the number of rows +/// the total byte size of the row group, and the [`SchemaDescriptor`] for the row group. +#[derive(Debug, Clone, PartialEq)] +pub struct RowGroupMetaData { + columns: Vec, + num_rows: i64, + sorting_columns: Option>, + total_byte_size: i64, + schema_descr: SchemaDescPtr, + /// We can't infer from file offset of first column since there may empty columns in row group. + file_offset: Option, + /// Ordinal position of this row group in file + ordinal: Option, +} + +impl RowGroupMetaData { + /// Returns builder for row group metadata. + pub fn builder(schema_descr: SchemaDescPtr) -> RowGroupMetaDataBuilder { + RowGroupMetaDataBuilder::new(schema_descr) + } + + /// Number of columns in this row group. + pub fn num_columns(&self) -> usize { + self.columns.len() + } + + /// Returns column chunk metadata for `i`th column. + pub fn column(&self, i: usize) -> &ColumnChunkMetaData { + &self.columns[i] + } + + /// Returns slice of column chunk metadata. + pub fn columns(&self) -> &[ColumnChunkMetaData] { + &self.columns + } + + /// Returns mutable slice of column chunk metadata. + pub fn columns_mut(&mut self) -> &mut [ColumnChunkMetaData] { + &mut self.columns + } + + /// Number of rows in this row group. + pub fn num_rows(&self) -> i64 { + self.num_rows + } + + /// Returns the sort ordering of the rows in this RowGroup if any + pub fn sorting_columns(&self) -> Option<&Vec> { + self.sorting_columns.as_ref() + } + + /// Total byte size of all uncompressed column data in this row group. + pub fn total_byte_size(&self) -> i64 { + self.total_byte_size + } + + /// Total size of all compressed column data in this row group. + pub fn compressed_size(&self) -> i64 { + self.columns.iter().map(|c| c.total_compressed_size).sum() + } + + /// Returns reference to a schema descriptor. + pub fn schema_descr(&self) -> &SchemaDescriptor { + self.schema_descr.as_ref() + } + + /// Returns reference counted clone of schema descriptor. + pub fn schema_descr_ptr(&self) -> SchemaDescPtr { + self.schema_descr.clone() + } + + /// Returns ordinal position of this row group in file. + /// + /// For example if this is the first row group in the file, this will return 0. + /// If this is the second row group in the file, this will return 1. + #[inline(always)] + pub fn ordinal(&self) -> Option { + self.ordinal + } + + /// Returns file offset of this row group in file. + #[inline(always)] + pub fn file_offset(&self) -> Option { + self.file_offset + } + + /// Converts this [`RowGroupMetaData`] into a [`RowGroupMetaDataBuilder`] + pub fn into_builder(self) -> RowGroupMetaDataBuilder { + RowGroupMetaDataBuilder(self) + } +} + +/// Builder for row group metadata. +pub struct RowGroupMetaDataBuilder(RowGroupMetaData); + +impl RowGroupMetaDataBuilder { + /// Creates new builder from schema descriptor. + fn new(schema_descr: SchemaDescPtr) -> Self { + Self(RowGroupMetaData { + columns: Vec::with_capacity(schema_descr.num_columns()), + schema_descr, + file_offset: None, + num_rows: 0, + sorting_columns: None, + total_byte_size: 0, + ordinal: None, + }) + } + + /// Sets number of rows in this row group. + pub fn set_num_rows(mut self, value: i64) -> Self { + self.0.num_rows = value; + self + } + + /// Sets the sorting order for columns + pub fn set_sorting_columns(mut self, value: Option>) -> Self { + self.0.sorting_columns = value; + self + } + + /// Sets total size in bytes for this row group. + pub fn set_total_byte_size(mut self, value: i64) -> Self { + self.0.total_byte_size = value; + self + } + + /// Takes ownership of the the column metadata in this builder, and clears + /// the list of columns. + /// + /// This can be used for more efficient creation of a new RowGroupMetaData + /// from an existing one. + pub fn take_columns(&mut self) -> Vec { + std::mem::take(&mut self.0.columns) + } + + /// Sets column metadata for this row group. + pub fn set_column_metadata(mut self, value: Vec) -> Self { + self.0.columns = value; + self + } + + /// Adds a column metadata to this row group + pub fn add_column_metadata(mut self, value: ColumnChunkMetaData) -> Self { + self.0.columns.push(value); + self + } + + /// Sets ordinal for this row group. + pub fn set_ordinal(mut self, value: i16) -> Self { + self.0.ordinal = Some(value); + self + } + + /// Sets file offset for this row group. + pub fn set_file_offset(mut self, value: i64) -> Self { + self.0.file_offset = Some(value); + self + } + + /// Builds row group metadata. + pub fn build(self) -> Result { + if self.0.schema_descr.num_columns() != self.0.columns.len() { + return Err(general_err!( + "Column length mismatch: {} != {}", + self.0.schema_descr.num_columns(), + self.0.columns.len() + )); + } + + Ok(self.0) + } + + /// Build row group metadata without validation. + pub(super) fn build_unchecked(self) -> RowGroupMetaData { + self.0 + } +} + +/// Metadata for a column chunk. +#[derive(Debug, Clone, PartialEq)] +pub struct ColumnChunkMetaData { + column_descr: ColumnDescPtr, + encodings: EncodingMask, + file_path: Option, + file_offset: i64, + num_values: i64, + compression: CompressionCodec, + total_compressed_size: i64, + total_uncompressed_size: i64, + data_page_offset: i64, + index_page_offset: Option, + dictionary_page_offset: Option, + statistics: Option, + geo_statistics: Option>, + encoding_stats: Option, + bloom_filter_offset: Option, + bloom_filter_length: Option, + offset_index_offset: Option, + offset_index_length: Option, + column_index_offset: Option, + column_index_length: Option, + unencoded_byte_array_data_bytes: Option, + repetition_level_histogram: Option, + definition_level_histogram: Option, + #[cfg(feature = "encryption")] + column_crypto_metadata: Option>, + #[cfg(feature = "encryption")] + encrypted_column_metadata: Option>, + /// When true, indicates the footer is plaintext (not encrypted). + /// This affects how column metadata is serialized when `encrypted_column_metadata` is present. + /// This field is only used at write time and is not needed when reading metadata. + #[cfg(feature = "encryption")] + plaintext_footer_mode: bool, +} + +/// Histograms for repetition and definition levels. +/// +/// Each histogram is a vector of length `max_level + 1`. The value at index `i` is the number of +/// values at level `i`. +/// +/// For example, `vec[0]` is the number of rows with level 0, `vec[1]` is the +/// number of rows with level 1, and so on. +/// +#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)] +pub struct LevelHistogram { + inner: Vec, +} + +impl LevelHistogram { + /// Creates a new level histogram data. + /// + /// Length will be `max_level + 1`. + /// + /// Returns `None` when `max_level == 0` (because histograms are not necessary in this case) + pub fn try_new(max_level: i16) -> Option { + if max_level > 0 { + Some(Self { + inner: vec![0; max_level as usize + 1], + }) + } else { + None + } + } + /// Returns a reference to the the histogram's values. + pub fn values(&self) -> &[i64] { + &self.inner + } + + /// Return the inner vector, consuming self + pub fn into_inner(self) -> Vec { + self.inner + } + + /// Returns the histogram value at the given index. + /// + /// The value of `i` is the number of values with level `i`. For example, + /// `get(1)` returns the number of values with level 1. + /// + /// Returns `None` if the index is out of bounds. + pub fn get(&self, index: usize) -> Option { + self.inner.get(index).copied() + } + + /// Adds the values from the other histogram to this histogram + /// + /// # Panics + /// If the histograms have different lengths + pub fn add(&mut self, other: &Self) { + assert_eq!(self.len(), other.len()); + for (dst, src) in self.inner.iter_mut().zip(other.inner.iter()) { + *dst += src; + } + } + + /// return the length of the histogram + pub fn len(&self) -> usize { + self.inner.len() + } + + /// returns if the histogram is empty + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + /// Sets the values of all histogram levels to 0. + pub fn reset(&mut self) { + for value in self.inner.iter_mut() { + *value = 0; + } + } + + /// Increments the count for a level value by `count`. + #[inline] + pub fn increment_by(&mut self, level: i16, count: i64) { + self.inner[level as usize] += count; + } + + /// Updates histogram values using provided repetition levels + /// + /// # Panics + /// if any of the levels is greater than the length of the histogram ( + /// the argument supplied to [`Self::try_new`]) + #[deprecated(since = "58.2.0", note = "Use `increment_by` instead")] + pub fn update_from_levels(&mut self, levels: &[i16]) { + for &level in levels { + self.increment_by(level, 1); + } + } +} + +impl From> for LevelHistogram { + fn from(inner: Vec) -> Self { + Self { inner } + } +} + +impl From for Vec { + fn from(value: LevelHistogram) -> Self { + value.into_inner() + } +} + +impl HeapSize for LevelHistogram { + fn heap_size(&self) -> usize { + self.inner.heap_size() + } +} + +/// Represents common operations for a column chunk. +impl ColumnChunkMetaData { + /// Returns builder for column chunk metadata. + pub fn builder(column_descr: ColumnDescPtr) -> ColumnChunkMetaDataBuilder { + ColumnChunkMetaDataBuilder::new(column_descr) + } + + /// File where the column chunk is stored. + /// + /// If not set, assumed to belong to the same file as the metadata. + /// This path is relative to the current file. + pub fn file_path(&self) -> Option<&str> { + self.file_path.as_deref() + } + + /// Byte offset of `ColumnMetaData` in `file_path()`. + /// + /// Note that the meaning of this field has been inconsistent between implementations + /// so its use has since been deprecated in the Parquet specification. Modern implementations + /// will set this to `0` to indicate that the `ColumnMetaData` is solely contained in the + /// `ColumnChunk` struct. + pub fn file_offset(&self) -> i64 { + self.file_offset + } + + /// Type of this column. Must be primitive. + pub fn column_type(&self) -> Type { + self.column_descr.physical_type() + } + + /// Path (or identifier) of this column. + pub fn column_path(&self) -> &ColumnPath { + self.column_descr.path() + } + + /// Descriptor for this column. + pub fn column_descr(&self) -> &ColumnDescriptor { + self.column_descr.as_ref() + } + + /// Reference counted clone of descriptor for this column. + pub fn column_descr_ptr(&self) -> ColumnDescPtr { + self.column_descr.clone() + } + + /// All encodings used for this column. + pub fn encodings(&self) -> impl Iterator { + self.encodings.encodings() + } + + /// All encodings used for this column, returned as a bitmask. + pub fn encodings_mask(&self) -> &EncodingMask { + &self.encodings + } + + /// Total number of values in this column chunk. + pub fn num_values(&self) -> i64 { + self.num_values + } + + /// [`Compression`] for this column. + /// + /// This is a default value suitable for passing to [`WriterPropertiesBuilder::set_compression`]. + /// It is constructed from the `codec` field of the Parquet `ColumnMetaData` + /// + /// [`WriterPropertiesBuilder::set_compression`]: crate::file::properties::WriterPropertiesBuilder + pub fn compression(&self) -> Compression { + self.compression.into() + } + + /// Returns the compression codec used when writing this column. + pub fn compression_codec(&self) -> CompressionCodec { + self.compression + } + + /// Returns the total compressed data size of this column chunk. + pub fn compressed_size(&self) -> i64 { + self.total_compressed_size + } + + /// Returns the total uncompressed data size of this column chunk. + pub fn uncompressed_size(&self) -> i64 { + self.total_uncompressed_size + } + + /// Returns the offset for the column data. + pub fn data_page_offset(&self) -> i64 { + self.data_page_offset + } + + /// Returns the offset for the index page. + pub fn index_page_offset(&self) -> Option { + self.index_page_offset + } + + /// Returns the offset for the dictionary page, if any. + pub fn dictionary_page_offset(&self) -> Option { + self.dictionary_page_offset + } + + /// Returns the offset and length in bytes of the column chunk within the file + pub fn byte_range(&self) -> (u64, u64) { + let col_start = match self.dictionary_page_offset() { + Some(dictionary_page_offset) => dictionary_page_offset, + None => self.data_page_offset(), + }; + let col_len = self.compressed_size(); + assert!( + col_start >= 0 && col_len >= 0, + "column start and length should not be negative" + ); + (col_start as u64, col_len as u64) + } + + /// Returns statistics that are set for this column chunk, + /// or `None` if no statistics are available. + pub fn statistics(&self) -> Option<&Statistics> { + self.statistics.as_ref() + } + + /// Returns geospatial statistics that are set for this column chunk, + /// or `None` if no geospatial statistics are available. + pub fn geo_statistics(&self) -> Option<&geo_statistics::GeospatialStatistics> { + self.geo_statistics.as_deref() + } + + /// Returns the page encoding statistics, or `None` if no page encoding statistics + /// are available (or they were converted to a mask). + /// + /// Note: By default, this crate converts page encoding statistics to a mask for performance + /// reasons. To get the full statistics, you must set [`ParquetMetaDataOptions::with_encoding_stats_as_mask`] + /// to `false`. + pub fn page_encoding_stats(&self) -> Option<&Vec> { + match self.encoding_stats.as_ref() { + Some(ParquetPageEncodingStats::Full(stats)) => Some(stats), + _ => None, + } + } + + /// Returns the page encoding statistics reduced to a bitmask, or `None` if statistics are + /// not available (or they were left in their original form). + /// + /// Note: This is the default behavior for this crate. + /// + /// The [`PageEncodingStats`] struct was added to the Parquet specification specifically to + /// enable fast determination of whether all pages in a column chunk are dictionary encoded + /// (see ). + /// Decoding the full page encoding statistics, however, can be very costly, and is not + /// necessary to support the aforementioned use case. As an alternative, this crate can + /// instead distill the list of `PageEncodingStats` down to a bitmask of just the encodings + /// used for data pages + /// (see [`ParquetMetaDataOptions::set_encoding_stats_as_mask`]). + /// To test for an all-dictionary-encoded chunk one could use this bitmask in the following way: + /// + /// ```rust + /// use parquet::basic::Encoding; + /// use parquet::file::metadata::ColumnChunkMetaData; + /// // test if all data pages in the column chunk are dictionary encoded + /// fn is_all_dictionary_encoded(col_meta: &ColumnChunkMetaData) -> bool { + /// // check that dictionary encoding was used + /// col_meta.dictionary_page_offset().is_some() + /// && col_meta.page_encoding_stats_mask().is_some_and(|mask| { + /// // mask should only have one bit set, either for PLAIN_DICTIONARY or + /// // RLE_DICTIONARY + /// mask.is_only(Encoding::PLAIN_DICTIONARY) || mask.is_only(Encoding::RLE_DICTIONARY) + /// }) + /// } + /// ``` + pub fn page_encoding_stats_mask(&self) -> Option<&EncodingMask> { + match self.encoding_stats.as_ref() { + Some(ParquetPageEncodingStats::Mask(stats)) => Some(stats), + _ => None, + } + } + + /// Returns the offset for the bloom filter. + pub fn bloom_filter_offset(&self) -> Option { + self.bloom_filter_offset + } + + /// Returns the offset for the bloom filter. + pub fn bloom_filter_length(&self) -> Option { + self.bloom_filter_length + } + + /// Returns the offset for the column index. + pub fn column_index_offset(&self) -> Option { + self.column_index_offset + } + + /// Returns the offset for the column index length. + pub fn column_index_length(&self) -> Option { + self.column_index_length + } + + /// Returns the range for the offset index if any + pub(crate) fn column_index_range(&self) -> Option> { + let offset = u64::try_from(self.column_index_offset?).ok()?; + let length = u64::try_from(self.column_index_length?).ok()?; + Some(offset..(offset + length)) + } + + /// Returns the offset for the offset index. + pub fn offset_index_offset(&self) -> Option { + self.offset_index_offset + } + + /// Returns the offset for the offset index length. + pub fn offset_index_length(&self) -> Option { + self.offset_index_length + } + + /// Returns the range for the offset index if any + pub(crate) fn offset_index_range(&self) -> Option> { + let offset = u64::try_from(self.offset_index_offset?).ok()?; + let length = u64::try_from(self.offset_index_length?).ok()?; + Some(offset..(offset + length)) + } + + /// Returns the number of bytes of variable length data after decoding. + /// + /// Only set for BYTE_ARRAY columns. This field may not be set by older + /// writers. + pub fn unencoded_byte_array_data_bytes(&self) -> Option { + self.unencoded_byte_array_data_bytes + } + + /// Returns the repetition level histogram. + /// + /// The returned value `vec[i]` is how many values are at repetition level `i`. For example, + /// `vec[0]` indicates how many rows the page contains. + /// This field may not be set by older writers. + pub fn repetition_level_histogram(&self) -> Option<&LevelHistogram> { + self.repetition_level_histogram.as_ref() + } + + /// Returns the definition level histogram. + /// + /// The returned value `vec[i]` is how many values are at definition level `i`. For example, + /// `vec[max_definition_level]` indicates how many non-null values are present in the page. + /// This field may not be set by older writers. + pub fn definition_level_histogram(&self) -> Option<&LevelHistogram> { + self.definition_level_histogram.as_ref() + } + + /// Returns the encryption metadata for this column chunk. + #[cfg(feature = "encryption")] + pub fn crypto_metadata(&self) -> Option<&ColumnCryptoMetaData> { + self.column_crypto_metadata.as_deref() + } + + /// Converts this [`ColumnChunkMetaData`] into a [`ColumnChunkMetaDataBuilder`] + pub fn into_builder(self) -> ColumnChunkMetaDataBuilder { + ColumnChunkMetaDataBuilder::from(self) + } +} + +/// Builder for [`ColumnChunkMetaData`] +/// +/// This builder is used to create a new column chunk metadata or modify an +/// existing one. +/// +/// # Example +/// ```no_run +/// # use parquet::file::metadata::{ColumnChunkMetaData, ColumnChunkMetaDataBuilder}; +/// # fn get_column_chunk_metadata() -> ColumnChunkMetaData { unimplemented!(); } +/// let column_chunk_metadata = get_column_chunk_metadata(); +/// // create a new builder from existing column chunk metadata +/// let builder = ColumnChunkMetaDataBuilder::from(column_chunk_metadata); +/// // clear the statistics: +/// let column_chunk_metadata: ColumnChunkMetaData = builder +/// .clear_statistics() +/// .build() +/// .unwrap(); +/// ``` +pub struct ColumnChunkMetaDataBuilder(ColumnChunkMetaData); + +impl ColumnChunkMetaDataBuilder { + /// Creates new column chunk metadata builder. + /// + /// See also [`ColumnChunkMetaData::builder`] + fn new(column_descr: ColumnDescPtr) -> Self { + Self(ColumnChunkMetaData { + column_descr, + encodings: Default::default(), + file_path: None, + file_offset: 0, + num_values: 0, + compression: CompressionCodec::UNCOMPRESSED, + total_compressed_size: 0, + total_uncompressed_size: 0, + data_page_offset: 0, + index_page_offset: None, + dictionary_page_offset: None, + statistics: None, + geo_statistics: None, + encoding_stats: None, + bloom_filter_offset: None, + bloom_filter_length: None, + offset_index_offset: None, + offset_index_length: None, + column_index_offset: None, + column_index_length: None, + unencoded_byte_array_data_bytes: None, + repetition_level_histogram: None, + definition_level_histogram: None, + #[cfg(feature = "encryption")] + column_crypto_metadata: None, + #[cfg(feature = "encryption")] + encrypted_column_metadata: None, + #[cfg(feature = "encryption")] + plaintext_footer_mode: false, + }) + } + + /// Sets list of encodings for this column chunk. + pub fn set_encodings(mut self, encodings: Vec) -> Self { + self.0.encodings = EncodingMask::new_from_encodings(encodings.iter()); + self + } + + /// Sets the encodings mask for this column chunk. + pub fn set_encodings_mask(mut self, encodings: EncodingMask) -> Self { + self.0.encodings = encodings; + self + } + + /// Sets optional file path for this column chunk. + pub fn set_file_path(mut self, value: String) -> Self { + self.0.file_path = Some(value); + self + } + + /// Sets number of values. + pub fn set_num_values(mut self, value: i64) -> Self { + self.0.num_values = value; + self + } + + /// Sets compression codec given a [`Compression`] configuration value. + pub fn set_compression(mut self, value: Compression) -> Self { + self.0.compression = value.into(); + self + } + + /// Sets compression codec. + pub fn set_compression_codec(mut self, value: CompressionCodec) -> Self { + self.0.compression = value; + self + } + + /// Sets total compressed size in bytes. + pub fn set_total_compressed_size(mut self, value: i64) -> Self { + self.0.total_compressed_size = value; + self + } + + /// Sets total uncompressed size in bytes. + pub fn set_total_uncompressed_size(mut self, value: i64) -> Self { + self.0.total_uncompressed_size = value; + self + } + + /// Sets data page offset in bytes. + pub fn set_data_page_offset(mut self, value: i64) -> Self { + self.0.data_page_offset = value; + self + } + + /// Sets optional dictionary page offset in bytes. + pub fn set_dictionary_page_offset(mut self, value: Option) -> Self { + self.0.dictionary_page_offset = value; + self + } + + /// Sets optional index page offset in bytes. + pub fn set_index_page_offset(mut self, value: Option) -> Self { + self.0.index_page_offset = value; + self + } + + /// Sets statistics for this column chunk. + pub fn set_statistics(mut self, value: Statistics) -> Self { + self.0.statistics = Some(value); + self + } + + /// Sets geospatial statistics for this column chunk. + pub fn set_geo_statistics(mut self, value: Box) -> Self { + self.0.geo_statistics = Some(value); + self + } + + /// Clears the statistics for this column chunk. + pub fn clear_statistics(mut self) -> Self { + self.0.statistics = None; + self + } + + /// Sets page encoding stats for this column chunk. + /// + /// This will overwrite any existing stats, either `Vec` based or bitmask. + pub fn set_page_encoding_stats(mut self, value: Vec) -> Self { + self.0.encoding_stats = Some(ParquetPageEncodingStats::Full(value)); + self + } + + /// Sets page encoding stats mask for this column chunk. + /// + /// This will overwrite any existing stats, either `Vec` based or bitmask. + pub fn set_page_encoding_stats_mask(mut self, value: EncodingMask) -> Self { + self.0.encoding_stats = Some(ParquetPageEncodingStats::Mask(value)); + self + } + + /// Clears the page encoding stats for this column chunk. + pub fn clear_page_encoding_stats(mut self) -> Self { + self.0.encoding_stats = None; + self + } + + /// Sets optional bloom filter offset in bytes. + pub fn set_bloom_filter_offset(mut self, value: Option) -> Self { + self.0.bloom_filter_offset = value; + self + } + + /// Sets optional bloom filter length in bytes. + pub fn set_bloom_filter_length(mut self, value: Option) -> Self { + self.0.bloom_filter_length = value; + self + } + + /// Sets optional offset index offset in bytes. + pub fn set_offset_index_offset(mut self, value: Option) -> Self { + self.0.offset_index_offset = value; + self + } + + /// Sets optional offset index length in bytes. + pub fn set_offset_index_length(mut self, value: Option) -> Self { + self.0.offset_index_length = value; + self + } + + /// Sets optional column index offset in bytes. + pub fn set_column_index_offset(mut self, value: Option) -> Self { + self.0.column_index_offset = value; + self + } + + /// Sets optional column index length in bytes. + pub fn set_column_index_length(mut self, value: Option) -> Self { + self.0.column_index_length = value; + self + } + + /// Sets optional length of variable length data in bytes. + pub fn set_unencoded_byte_array_data_bytes(mut self, value: Option) -> Self { + self.0.unencoded_byte_array_data_bytes = value; + self + } + + /// Sets optional repetition level histogram + pub fn set_repetition_level_histogram(mut self, value: Option) -> Self { + self.0.repetition_level_histogram = value; + self + } + + /// Sets optional repetition level histogram + pub fn set_definition_level_histogram(mut self, value: Option) -> Self { + self.0.definition_level_histogram = value; + self + } + + #[cfg(feature = "encryption")] + /// Set the encryption metadata for an encrypted column + pub fn set_column_crypto_metadata(mut self, value: Option) -> Self { + self.0.column_crypto_metadata = value.map(Box::new); + self + } + + #[cfg(feature = "encryption")] + /// Set the encryption metadata for an encrypted column + pub fn set_encrypted_column_metadata(mut self, value: Option>) -> Self { + self.0.encrypted_column_metadata = value; + self + } + + /// Builds column chunk metadata. + pub fn build(self) -> Result { + Ok(self.0) + } +} + +/// Builder for Parquet [`ColumnIndex`], part of the Parquet [PageIndex] +/// +/// [PageIndex]: https://github.com/apache/parquet-format/blob/master/PageIndex.md +/// [`ColumnIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md +pub struct ColumnIndexBuilder { + column_type: Type, + null_pages: Vec, + min_values: Vec>, + max_values: Vec>, + null_counts: Vec, + boundary_order: BoundaryOrder, + /// contains the concatenation of the histograms of all pages + repetition_level_histograms: Option>, + /// contains the concatenation of the histograms of all pages + definition_level_histograms: Option>, + /// Is the information in the builder valid? + /// + /// Set to `false` if any entry in the page doesn't have statistics for + /// some reason, so statistics for that page won't be written to the file. + /// This might happen if the page is entirely null, or + /// is a floating point column without any non-nan values + /// e.g. + valid: bool, +} + +impl ColumnIndexBuilder { + /// Creates a new column index builder. + pub fn new(column_type: Type) -> Self { + ColumnIndexBuilder { + column_type, + null_pages: Vec::new(), + min_values: Vec::new(), + max_values: Vec::new(), + null_counts: Vec::new(), + boundary_order: BoundaryOrder::UNORDERED, + repetition_level_histograms: None, + definition_level_histograms: None, + valid: true, + } + } + + /// Append statistics for the next page + pub fn append( + &mut self, + null_page: bool, + min_value: Vec, + max_value: Vec, + null_count: i64, + ) { + self.null_pages.push(null_page); + self.min_values.push(min_value); + self.max_values.push(max_value); + self.null_counts.push(null_count); + } + + /// Append the given page-level histograms to the [`ColumnIndex`] histograms. + /// Does nothing if the `ColumnIndexBuilder` is not in the `valid` state. + /// + /// [`ColumnIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + pub fn append_histograms( + &mut self, + repetition_level_histogram: &Option, + definition_level_histogram: &Option, + ) { + if !self.valid { + return; + } + if let Some(rep_lvl_hist) = repetition_level_histogram { + let hist = self.repetition_level_histograms.get_or_insert(Vec::new()); + hist.reserve(rep_lvl_hist.len()); + hist.extend(rep_lvl_hist.values()); + } + if let Some(def_lvl_hist) = definition_level_histogram { + let hist = self.definition_level_histograms.get_or_insert(Vec::new()); + hist.reserve(def_lvl_hist.len()); + hist.extend(def_lvl_hist.values()); + } + } + + /// Set the boundary order of the column index + pub fn set_boundary_order(&mut self, boundary_order: BoundaryOrder) { + self.boundary_order = boundary_order; + } + + /// Mark this column index as invalid + pub fn to_invalid(&mut self) { + self.valid = false; + } + + /// Is the information in the builder valid? + pub fn valid(&self) -> bool { + self.valid + } + + /// Build and get the column index + /// + /// Note: callers should check [`Self::valid`] before calling this method + pub fn build(self) -> Result { + Ok(match self.column_type { + Type::BOOLEAN => { + let index = self.build_page_index()?; + ColumnIndexMetaData::BOOLEAN(index) + } + Type::INT32 => { + let index = self.build_page_index()?; + ColumnIndexMetaData::INT32(index) + } + Type::INT64 => { + let index = self.build_page_index()?; + ColumnIndexMetaData::INT64(index) + } + Type::INT96 => { + let index = self.build_page_index()?; + ColumnIndexMetaData::INT96(index) + } + Type::FLOAT => { + let index = self.build_page_index()?; + ColumnIndexMetaData::FLOAT(index) + } + Type::DOUBLE => { + let index = self.build_page_index()?; + ColumnIndexMetaData::DOUBLE(index) + } + Type::BYTE_ARRAY => { + let index = self.build_byte_array_index()?; + ColumnIndexMetaData::BYTE_ARRAY(index) + } + Type::FIXED_LEN_BYTE_ARRAY => { + let index = self.build_byte_array_index()?; + ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(index) + } + }) + } + + fn build_page_index(self) -> Result> + where + T: ParquetValueType, + { + let min_values: Vec<&[u8]> = self.min_values.iter().map(|v| v.as_slice()).collect(); + let max_values: Vec<&[u8]> = self.max_values.iter().map(|v| v.as_slice()).collect(); + + PrimitiveColumnIndex::try_new( + self.null_pages, + self.boundary_order, + Some(self.null_counts), + self.repetition_level_histograms, + self.definition_level_histograms, + min_values, + max_values, + ) + } + + fn build_byte_array_index(self) -> Result { + let min_values: Vec<&[u8]> = self.min_values.iter().map(|v| v.as_slice()).collect(); + let max_values: Vec<&[u8]> = self.max_values.iter().map(|v| v.as_slice()).collect(); + + ByteArrayColumnIndex::try_new( + self.null_pages, + self.boundary_order, + Some(self.null_counts), + self.repetition_level_histograms, + self.definition_level_histograms, + min_values, + max_values, + ) + } +} + +impl From for ColumnChunkMetaDataBuilder { + fn from(value: ColumnChunkMetaData) -> Self { + ColumnChunkMetaDataBuilder(value) + } +} + +/// Builder for offset index, part of the Parquet [PageIndex]. +/// +/// [PageIndex]: https://github.com/apache/parquet-format/blob/master/PageIndex.md +pub struct OffsetIndexBuilder { + offset_array: Vec, + compressed_page_size_array: Vec, + first_row_index_array: Vec, + unencoded_byte_array_data_bytes_array: Option>, + current_first_row_index: i64, +} + +impl Default for OffsetIndexBuilder { + fn default() -> Self { + Self::new() + } +} + +impl OffsetIndexBuilder { + /// Creates a new offset index builder. + pub fn new() -> Self { + OffsetIndexBuilder { + offset_array: Vec::new(), + compressed_page_size_array: Vec::new(), + first_row_index_array: Vec::new(), + unencoded_byte_array_data_bytes_array: None, + current_first_row_index: 0, + } + } + + /// Append the row count of the next page. + pub fn append_row_count(&mut self, row_count: i64) { + let current_page_row_index = self.current_first_row_index; + self.first_row_index_array.push(current_page_row_index); + self.current_first_row_index += row_count; + } + + /// Append the offset and size of the next page. + pub fn append_offset_and_size(&mut self, offset: i64, compressed_page_size: i32) { + self.offset_array.push(offset); + self.compressed_page_size_array.push(compressed_page_size); + } + + /// Append the unencoded byte array data bytes of the next page. + pub fn append_unencoded_byte_array_data_bytes( + &mut self, + unencoded_byte_array_data_bytes: Option, + ) { + if let Some(val) = unencoded_byte_array_data_bytes { + self.unencoded_byte_array_data_bytes_array + .get_or_insert(Vec::new()) + .push(val); + } + } + + /// Build and get the thrift metadata of offset index + pub fn build(self) -> OffsetIndexMetaData { + let locations = self + .offset_array + .iter() + .zip(self.compressed_page_size_array.iter()) + .zip(self.first_row_index_array.iter()) + .map(|((offset, size), row_index)| PageLocation { + offset: *offset, + compressed_page_size: *size, + first_row_index: *row_index, + }) + .collect::>(); + OffsetIndexMetaData { + page_locations: locations, + unencoded_byte_array_data_bytes: self.unencoded_byte_array_data_bytes_array, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::basic::{PageType, SortOrder}; + use crate::file::metadata::thrift::tests::{ + read_column_chunk, read_column_chunk_with_options, read_row_group, + }; + + #[test] + #[allow(deprecated)] + fn test_level_histogram_update_from_levels_compat() { + let mut histogram = LevelHistogram::try_new(2).unwrap(); + histogram.update_from_levels(&[0, 2, 1, 2, 2]); + assert_eq!(histogram.values(), &[1, 1, 3]); + } + + #[test] + fn test_row_group_metadata_thrift_conversion() { + let schema_descr = get_test_schema_descr(); + + let mut columns = vec![]; + for ptr in schema_descr.columns() { + let column = ColumnChunkMetaData::builder(ptr.clone()).build().unwrap(); + columns.push(column); + } + let row_group_meta = RowGroupMetaData::builder(schema_descr.clone()) + .set_num_rows(1000) + .set_total_byte_size(2000) + .set_column_metadata(columns) + .set_ordinal(1) + .build() + .unwrap(); + + let mut buf = Vec::new(); + let mut writer = ThriftCompactOutputProtocol::new(&mut buf); + row_group_meta.write_thrift(&mut writer).unwrap(); + + let row_group_res = read_row_group(&mut buf, schema_descr).unwrap(); + + assert_eq!(row_group_res, row_group_meta); + } + + #[test] + fn test_row_group_metadata_thrift_conversion_empty() { + let schema_descr = get_test_schema_descr(); + + let row_group_meta = RowGroupMetaData::builder(schema_descr).build(); + + assert!(row_group_meta.is_err()); + if let Err(e) = row_group_meta { + assert_eq!( + format!("{e}"), + "Parquet error: Column length mismatch: 2 != 0" + ); + } + } + + /// Test reading a corrupted Parquet file with 3 columns in its schema but only 2 in its row group + #[test] + fn test_row_group_metadata_thrift_corrupted() { + let schema_descr_2cols = Arc::new(SchemaDescriptor::new(Arc::new( + SchemaType::group_type_builder("schema") + .with_fields(vec![ + Arc::new( + SchemaType::primitive_type_builder("a", Type::INT32) + .build() + .unwrap(), + ), + Arc::new( + SchemaType::primitive_type_builder("b", Type::INT32) + .build() + .unwrap(), + ), + ]) + .build() + .unwrap(), + ))); + + let schema_descr_3cols = Arc::new(SchemaDescriptor::new(Arc::new( + SchemaType::group_type_builder("schema") + .with_fields(vec![ + Arc::new( + SchemaType::primitive_type_builder("a", Type::INT32) + .build() + .unwrap(), + ), + Arc::new( + SchemaType::primitive_type_builder("b", Type::INT32) + .build() + .unwrap(), + ), + Arc::new( + SchemaType::primitive_type_builder("c", Type::INT32) + .build() + .unwrap(), + ), + ]) + .build() + .unwrap(), + ))); + + let row_group_meta_2cols = RowGroupMetaData::builder(schema_descr_2cols.clone()) + .set_num_rows(1000) + .set_total_byte_size(2000) + .set_column_metadata(vec![ + ColumnChunkMetaData::builder(schema_descr_2cols.column(0)) + .build() + .unwrap(), + ColumnChunkMetaData::builder(schema_descr_2cols.column(1)) + .build() + .unwrap(), + ]) + .set_ordinal(1) + .build() + .unwrap(); + let mut buf = Vec::new(); + let mut writer = ThriftCompactOutputProtocol::new(&mut buf); + row_group_meta_2cols.write_thrift(&mut writer).unwrap(); + + let err = read_row_group(&mut buf, schema_descr_3cols) + .unwrap_err() + .to_string(); + assert_eq!( + err, + "Parquet error: Column count mismatch. Schema has 3 columns while Row Group has 2" + ); + } + + #[test] + fn test_column_chunk_metadata_thrift_conversion() { + let column_descr = get_test_schema_descr().column(0); + let col_metadata = ColumnChunkMetaData::builder(column_descr.clone()) + .set_encodings_mask(EncodingMask::new_from_encodings( + [Encoding::PLAIN, Encoding::RLE].iter(), + )) + .set_file_path("file_path".to_owned()) + .set_num_values(1000) + .set_compression_codec(CompressionCodec::SNAPPY) + .set_total_compressed_size(2000) + .set_total_uncompressed_size(3000) + .set_data_page_offset(4000) + .set_dictionary_page_offset(Some(5000)) + .set_page_encoding_stats(vec![ + PageEncodingStats { + page_type: PageType::DATA_PAGE, + encoding: Encoding::PLAIN, + count: 3, + }, + PageEncodingStats { + page_type: PageType::DATA_PAGE, + encoding: Encoding::RLE, + count: 5, + }, + ]) + .set_bloom_filter_offset(Some(6000)) + .set_bloom_filter_length(Some(25)) + .set_offset_index_offset(Some(7000)) + .set_offset_index_length(Some(25)) + .set_column_index_offset(Some(8000)) + .set_column_index_length(Some(25)) + .set_unencoded_byte_array_data_bytes(Some(2000)) + .set_repetition_level_histogram(Some(LevelHistogram::from(vec![100, 100]))) + .set_definition_level_histogram(Some(LevelHistogram::from(vec![0, 200]))) + .build() + .unwrap(); + + let mut buf = Vec::new(); + let mut writer = ThriftCompactOutputProtocol::new(&mut buf); + col_metadata.write_thrift(&mut writer).unwrap(); + let col_chunk_res = read_column_chunk(&mut buf, column_descr.clone()).unwrap(); + + let expected_metadata = ColumnChunkMetaData::builder(column_descr) + .set_encodings_mask(EncodingMask::new_from_encodings( + [Encoding::PLAIN, Encoding::RLE].iter(), + )) + .set_file_path("file_path".to_owned()) + .set_num_values(1000) + .set_compression_codec(CompressionCodec::SNAPPY) + .set_total_compressed_size(2000) + .set_total_uncompressed_size(3000) + .set_data_page_offset(4000) + .set_dictionary_page_offset(Some(5000)) + .set_page_encoding_stats_mask(EncodingMask::new_from_encodings( + [Encoding::PLAIN, Encoding::RLE].iter(), + )) + .set_bloom_filter_offset(Some(6000)) + .set_bloom_filter_length(Some(25)) + .set_offset_index_offset(Some(7000)) + .set_offset_index_length(Some(25)) + .set_column_index_offset(Some(8000)) + .set_column_index_length(Some(25)) + .set_unencoded_byte_array_data_bytes(Some(2000)) + .set_repetition_level_histogram(Some(LevelHistogram::from(vec![100, 100]))) + .set_definition_level_histogram(Some(LevelHistogram::from(vec![0, 200]))) + .build() + .unwrap(); + + assert_eq!(col_chunk_res, expected_metadata); + } + + #[test] + fn test_column_chunk_metadata_thrift_conversion_full_stats() { + let column_descr = get_test_schema_descr().column(0); + let stats = vec![ + PageEncodingStats { + page_type: PageType::DATA_PAGE, + encoding: Encoding::PLAIN, + count: 3, + }, + PageEncodingStats { + page_type: PageType::DATA_PAGE, + encoding: Encoding::RLE, + count: 5, + }, + ]; + let col_metadata = ColumnChunkMetaData::builder(column_descr.clone()) + .set_encodings_mask(EncodingMask::new_from_encodings( + [Encoding::PLAIN, Encoding::RLE].iter(), + )) + .set_num_values(1000) + .set_compression_codec(CompressionCodec::SNAPPY) + .set_total_compressed_size(2000) + .set_total_uncompressed_size(3000) + .set_data_page_offset(4000) + .set_page_encoding_stats(stats) + .build() + .unwrap(); + + let mut buf = Vec::new(); + let mut writer = ThriftCompactOutputProtocol::new(&mut buf); + col_metadata.write_thrift(&mut writer).unwrap(); + + let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false); + let col_chunk_res = + read_column_chunk_with_options(&mut buf, column_descr, Some(&options)).unwrap(); + + assert_eq!(col_chunk_res, col_metadata); + } + + #[test] + fn test_column_chunk_metadata_thrift_conversion_empty() { + let column_descr = get_test_schema_descr().column(0); + + let col_metadata = ColumnChunkMetaData::builder(column_descr.clone()) + .build() + .unwrap(); + + let mut buf = Vec::new(); + let mut writer = ThriftCompactOutputProtocol::new(&mut buf); + col_metadata.write_thrift(&mut writer).unwrap(); + let col_chunk_res = read_column_chunk(&mut buf, column_descr).unwrap(); + + assert_eq!(col_chunk_res, col_metadata); + } + + #[test] + fn test_compressed_size() { + let schema_descr = get_test_schema_descr(); + + let mut columns = vec![]; + for column_descr in schema_descr.columns() { + let column = ColumnChunkMetaData::builder(column_descr.clone()) + .set_total_compressed_size(500) + .set_total_uncompressed_size(700) + .build() + .unwrap(); + columns.push(column); + } + let row_group_meta = RowGroupMetaData::builder(schema_descr) + .set_num_rows(1000) + .set_column_metadata(columns) + .build() + .unwrap(); + + let compressed_size_res: i64 = row_group_meta.compressed_size(); + let compressed_size_exp: i64 = 1000; + + assert_eq!(compressed_size_res, compressed_size_exp); + } + + #[test] + fn test_memory_size() { + let schema_descr = get_test_schema_descr(); + + let columns = schema_descr + .columns() + .iter() + .map(|column_descr| { + ColumnChunkMetaData::builder(column_descr.clone()) + .set_statistics(Statistics::new::(None, None, None, None, false)) + .build() + }) + .collect::>>() + .unwrap(); + let row_group_meta = RowGroupMetaData::builder(schema_descr.clone()) + .set_num_rows(1000) + .set_column_metadata(columns) + .build() + .unwrap(); + let row_group_meta = vec![row_group_meta]; + + let version = 2; + let num_rows = 1000; + let created_by = Some(String::from("test harness")); + let key_value_metadata = Some(vec![KeyValue::new( + String::from("Foo"), + Some(String::from("bar")), + )]); + let column_orders = Some(vec![ + ColumnOrder::UNDEFINED, + ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED), + ]); + let file_metadata = FileMetaData::new( + version, + num_rows, + created_by, + key_value_metadata, + schema_descr.clone(), + column_orders, + ); + + // Now, add in Exact Statistics + let columns_with_stats = schema_descr + .columns() + .iter() + .map(|column_descr| { + ColumnChunkMetaData::builder(column_descr.clone()) + .set_statistics(Statistics::new::( + Some(0), + Some(100), + None, + None, + false, + )) + .build() + }) + .collect::>>() + .unwrap(); + + let row_group_meta_with_stats = RowGroupMetaData::builder(schema_descr) + .set_num_rows(1000) + .set_column_metadata(columns_with_stats) + .build() + .unwrap(); + let row_group_meta_with_stats = vec![row_group_meta_with_stats]; + + let parquet_meta = ParquetMetaDataBuilder::new(file_metadata.clone()) + .set_row_groups(row_group_meta_with_stats) + .build(); + + #[cfg(not(feature = "encryption"))] + let base_expected_size = 2734; + #[cfg(feature = "encryption")] + let base_expected_size = 2902; + + assert_eq!(parquet_meta.memory_size(), base_expected_size); + + let mut column_index = ColumnIndexBuilder::new(Type::BOOLEAN); + column_index.append(false, vec![1u8], vec![2u8, 3u8], 4); + let column_index = column_index.build().unwrap(); + let native_index = match column_index { + ColumnIndexMetaData::BOOLEAN(index) => index, + _ => panic!("wrong type of column index"), + }; + + // Now, add in OffsetIndex + let mut offset_index = OffsetIndexBuilder::new(); + offset_index.append_row_count(1); + offset_index.append_offset_and_size(2, 3); + offset_index.append_unencoded_byte_array_data_bytes(Some(10)); + offset_index.append_row_count(1); + offset_index.append_offset_and_size(2, 3); + offset_index.append_unencoded_byte_array_data_bytes(Some(10)); + let offset_index = offset_index.build(); + + let parquet_meta = ParquetMetaDataBuilder::new(file_metadata) + .set_row_groups(row_group_meta) + .set_column_index(Some(vec![vec![ColumnIndexMetaData::BOOLEAN(native_index)]])) + .set_offset_index(Some(vec![vec![offset_index]])) + .build(); + + #[cfg(not(feature = "encryption"))] + let bigger_expected_size = 3160; + #[cfg(feature = "encryption")] + let bigger_expected_size = 3328; + + // more set fields means more memory usage + assert!(bigger_expected_size > base_expected_size); + assert_eq!(parquet_meta.memory_size(), bigger_expected_size); + } + + #[test] + #[cfg(feature = "encryption")] + fn test_memory_size_with_decryptor() { + use crate::encryption::decrypt::FileDecryptionProperties; + use crate::file::metadata::thrift::encryption::AesGcmV1; + + let schema_descr = get_test_schema_descr(); + + let columns = schema_descr + .columns() + .iter() + .map(|column_descr| ColumnChunkMetaData::builder(column_descr.clone()).build()) + .collect::>>() + .unwrap(); + let row_group_meta = RowGroupMetaData::builder(schema_descr.clone()) + .set_num_rows(1000) + .set_column_metadata(columns) + .build() + .unwrap(); + let row_group_meta = vec![row_group_meta]; + + let version = 2; + let num_rows = 1000; + let aad_file_unique = vec![1u8; 8]; + let aad_prefix = vec![2u8; 8]; + let encryption_algorithm = EncryptionAlgorithm::AES_GCM_V1(AesGcmV1 { + aad_prefix: Some(aad_prefix.clone()), + aad_file_unique: Some(aad_file_unique.clone()), + supply_aad_prefix: Some(true), + }); + let footer_key_metadata = Some(vec![3u8; 8]); + let file_metadata = + FileMetaData::new(version, num_rows, None, None, schema_descr.clone(), None) + .with_encryption_algorithm(Some(encryption_algorithm)) + .with_footer_signing_key_metadata(footer_key_metadata.clone()); + + let parquet_meta_data = ParquetMetaDataBuilder::new(file_metadata.clone()) + .set_row_groups(row_group_meta.clone()) + .build(); + + let base_expected_size = 2042; + assert_eq!(parquet_meta_data.memory_size(), base_expected_size); + + let footer_key = "0123456789012345".as_bytes(); + let column_key = "1234567890123450".as_bytes(); + let mut decryption_properties_builder = + FileDecryptionProperties::builder(footer_key.to_vec()) + .with_aad_prefix(aad_prefix.clone()); + for column in schema_descr.columns() { + decryption_properties_builder = decryption_properties_builder + .with_column_key(&column.path().string(), column_key.to_vec()); + } + let decryption_properties = decryption_properties_builder.build().unwrap(); + let decryptor = FileDecryptor::new( + &decryption_properties, + footer_key_metadata.as_deref(), + aad_file_unique, + aad_prefix, + ) + .unwrap(); + + let parquet_meta_data = ParquetMetaDataBuilder::new(file_metadata.clone()) + .set_row_groups(row_group_meta.clone()) + .set_file_decryptor(Some(decryptor)) + .build(); + + let expected_size_with_decryptor = 3056; + assert!(expected_size_with_decryptor > base_expected_size); + + assert_eq!( + parquet_meta_data.memory_size(), + expected_size_with_decryptor + ); + } + + /// Returns sample schema descriptor so we can create column metadata. + fn get_test_schema_descr() -> SchemaDescPtr { + let schema = SchemaType::group_type_builder("schema") + .with_fields(vec![ + Arc::new( + SchemaType::primitive_type_builder("a", Type::INT32) + .build() + .unwrap(), + ), + Arc::new( + SchemaType::primitive_type_builder("b", Type::INT32) + .build() + .unwrap(), + ), + ]) + .build() + .unwrap(); + + Arc::new(SchemaDescriptor::new(Arc::new(schema))) + } +} diff --git a/vendor/parquet/src/file/metadata/options.rs b/vendor/parquet/src/file/metadata/options.rs new file mode 100644 index 0000000..e91f5bd --- /dev/null +++ b/vendor/parquet/src/file/metadata/options.rs @@ -0,0 +1,305 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Options used to control metadata parsing + +use std::collections::HashSet; +use std::sync::Arc; + +use crate::schema::types::SchemaDescPtr; + +/// Enum to control decoding of some Parquet statistics fields. +/// +/// # Example +/// ```rust +/// use parquet::file::metadata::ParquetStatisticsPolicy; +/// use parquet::file::serialized_reader::ReadOptionsBuilder; +/// use parquet::arrow::arrow_reader::ArrowReaderOptions; +/// +/// // Set arrow options to skip encoding statistics for all columns. +/// let options = +/// ArrowReaderOptions::new().with_encoding_stats_policy(ParquetStatisticsPolicy::SkipAll); +/// +/// // Set serialized reader options to decode encoding statistics for all columns. +/// let options = +/// ReadOptionsBuilder::new().with_encoding_stats_policy(ParquetStatisticsPolicy::KeepAll) +/// .build(); +/// +/// // Set arrow options to skip encoding statistics for all columns, but to decode statistics +/// // for columns 0 and 1. +/// let options = ArrowReaderOptions::new() +/// .with_encoding_stats_policy(ParquetStatisticsPolicy::skip_except(&[0, 1])); +/// ``` +#[derive(Default, Debug, Clone)] +pub enum ParquetStatisticsPolicy { + /// Decode the relevant statistics for all columns. + #[default] + KeepAll, + /// Skip decoding the relevant statistics for all columns. + SkipAll, + /// Skip decoding the relevant statistics for all columns not in the provided set + /// of column indices. + SkipExcept(Arc>), +} + +impl ParquetStatisticsPolicy { + /// Create a `ParquetStatisticsPolicy` to skip all columns except those in `keep`. + /// + /// If `keep` is empty, then this returns [`Self::SkipAll`] + pub fn skip_except(keep: &[usize]) -> Self { + if keep.is_empty() { + Self::SkipAll + } else { + let mut keep_set = HashSet::::with_capacity(keep.len()); + keep_set.extend(keep.iter()); + Self::SkipExcept(Arc::new(keep_set)) + } + } + + /// Returns whether the policy for the given column index is to skip the statistics. + pub(crate) fn is_skip(&self, col_index: usize) -> bool { + match self { + Self::KeepAll => false, + Self::SkipAll => true, + Self::SkipExcept(keep) => !keep.contains(&col_index), + } + } +} + +/// Options that can be set to control what parts of the Parquet file footer +/// metadata will be decoded and made present in the [`ParquetMetaData`] returned +/// by [`ParquetMetaDataReader`] and [`ParquetMetaDataPushDecoder`]. +/// +/// [`ParquetMetaData`]: crate::file::metadata::ParquetMetaData +/// [`ParquetMetaDataReader`]: crate::file::metadata::ParquetMetaDataReader +/// [`ParquetMetaDataPushDecoder`]: crate::file::metadata::ParquetMetaDataPushDecoder +#[derive(Debug, Clone)] +pub struct ParquetMetaDataOptions { + schema_descr: Option, + encoding_stats_as_mask: bool, + encoding_stats_policy: ParquetStatisticsPolicy, + column_stats_policy: ParquetStatisticsPolicy, + size_stats_policy: ParquetStatisticsPolicy, +} + +impl Default for ParquetMetaDataOptions { + fn default() -> Self { + Self { + schema_descr: None, + encoding_stats_as_mask: true, + encoding_stats_policy: ParquetStatisticsPolicy::KeepAll, + column_stats_policy: ParquetStatisticsPolicy::KeepAll, + size_stats_policy: ParquetStatisticsPolicy::KeepAll, + } + } +} + +impl ParquetMetaDataOptions { + /// Return a new default [`ParquetMetaDataOptions`]. + pub fn new() -> Self { + Default::default() + } + + /// Returns an optional [`SchemaDescPtr`] to use when decoding. If this is not `None` then + /// the schema in the footer will be skipped. + pub fn schema(&self) -> Option<&SchemaDescPtr> { + self.schema_descr.as_ref() + } + + /// Provide a schema to use when decoding the metadata. + pub fn set_schema(&mut self, val: SchemaDescPtr) { + self.schema_descr = Some(val); + } + + /// Call [`Self::set_schema`] and return `Self` for chaining. + pub fn with_schema(mut self, val: SchemaDescPtr) -> Self { + self.set_schema(val); + self + } + + /// Returns whether to present the [`encoding_stats`] field of the Parquet `ColumnMetaData` + /// as a bitmask (defaults to `true`). + /// + /// See [`ColumnChunkMetaData::page_encoding_stats_mask`] for an explanation of why this + /// might be desirable. + /// + /// [`ColumnChunkMetaData::page_encoding_stats_mask`]: + /// crate::file::metadata::ColumnChunkMetaData::page_encoding_stats_mask + /// [`encoding_stats`]: + /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L917 + pub fn encoding_stats_as_mask(&self) -> bool { + self.encoding_stats_as_mask + } + + /// Convert [`encoding_stats`] from a vector of [`PageEncodingStats`] to a bitmask. This can + /// speed up metadata decoding while still enabling some use cases served by the full stats. + /// + /// Note that if for a given column both this option and `skip_encoding_stats` are `true`, the + /// stats will be skipped and not be returned as a mask. + /// + /// See [`ColumnChunkMetaData::page_encoding_stats_mask`] for more information. + /// + /// [`PageEncodingStats`]: crate::file::metadata::PageEncodingStats + /// [`ColumnChunkMetaData::page_encoding_stats_mask`]: + /// crate::file::metadata::ColumnChunkMetaData::page_encoding_stats_mask + /// [`encoding_stats`]: + /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L917 + pub fn set_encoding_stats_as_mask(&mut self, val: bool) { + self.encoding_stats_as_mask = val; + } + + /// Call [`Self::set_encoding_stats_as_mask`] and return `Self` for chaining. + pub fn with_encoding_stats_as_mask(mut self, val: bool) -> Self { + self.set_encoding_stats_as_mask(val); + self + } + + /// Returns whether to skip decoding the [`encoding_stats`] in the Parquet `ColumnMetaData` + /// for the column indexed by `col_index`. + /// + /// [`encoding_stats`]: + /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L917 + pub fn skip_encoding_stats(&self, col_index: usize) -> bool { + self.encoding_stats_policy.is_skip(col_index) + } + + /// Sets the decoding policy for [`encoding_stats`] in the Parquet `ColumnMetaData`. + /// + /// The default policy is to decode all `encoding_stats`. + /// + /// This option takes precedence over [`Self::encoding_stats_as_mask`]. + /// + /// [`encoding_stats`]: + /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L917 + pub fn set_encoding_stats_policy(&mut self, policy: ParquetStatisticsPolicy) { + self.encoding_stats_policy = policy; + } + + /// Call [`Self::set_encoding_stats_policy`] and return `Self` for chaining. + pub fn with_encoding_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self { + self.set_encoding_stats_policy(policy); + self + } + + /// Returns whether to skip decoding the [`statistics`] in the Parquet `ColumnMetaData` + /// for the column indexed by `col_index`. + /// + /// [`statistics`]: + /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L912 + pub fn skip_column_stats(&self, col_index: usize) -> bool { + self.column_stats_policy.is_skip(col_index) + } + + /// Sets the decoding policy for [`statistics`] in the Parquet `ColumnMetaData`. + /// + /// The default policy is to decode all `statistics`. + /// + /// [`statistics`]: + /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L912 + pub fn set_column_stats_policy(&mut self, policy: ParquetStatisticsPolicy) { + self.column_stats_policy = policy; + } + + /// Call [`Self::set_column_stats_policy`] and return `Self` for chaining. + pub fn with_column_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self { + self.set_column_stats_policy(policy); + self + } + + /// Returns whether to skip decoding the [`size_statistics`] in the Parquet `ColumnMetaData` + /// for the column indexed by `col_index`. + /// + /// [`size_statistics`]: + /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L936 + pub fn skip_size_stats(&self, col_index: usize) -> bool { + self.size_stats_policy.is_skip(col_index) + } + + /// Sets the decoding policy for [`size_statistics`] in the Parquet `ColumnMetaData`. + /// + /// The default policy is to decode all `size_statistics`. + /// + /// [`size_statistics`]: + /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L936 + pub fn set_size_stats_policy(&mut self, policy: ParquetStatisticsPolicy) { + self.size_stats_policy = policy; + } + + /// Call [`Self::set_size_stats_policy`] and return `Self` for chaining. + pub fn with_size_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self { + self.set_size_stats_policy(policy); + self + } +} + +#[cfg(test)] +mod tests { + use bytes::Bytes; + + use crate::{ + DecodeResult, + file::metadata::{ParquetMetaDataOptions, ParquetMetaDataPushDecoder}, + util::test_common::file_util::get_test_file, + }; + use std::{io::Read, sync::Arc}; + + #[test] + fn test_options_default() { + let options = ParquetMetaDataOptions::default(); + assert!(options.encoding_stats_as_mask()); + } + + #[test] + fn test_provide_schema() { + let mut buf: Vec = Vec::new(); + get_test_file("alltypes_plain.parquet") + .read_to_end(&mut buf) + .unwrap(); + + let data = Bytes::from(buf); + let mut decoder = ParquetMetaDataPushDecoder::try_new(data.len() as u64).unwrap(); + decoder + .push_range(0..data.len() as u64, data.clone()) + .unwrap(); + + let expected = match decoder.try_decode().unwrap() { + DecodeResult::Data(m) => m, + _ => panic!("could not parse metadata"), + }; + let expected_schema = expected.file_metadata().schema_descr_ptr(); + + let mut options = ParquetMetaDataOptions::new(); + options.set_schema(expected_schema); + let options = Arc::new(options); + + let mut decoder = ParquetMetaDataPushDecoder::try_new(data.len() as u64) + .unwrap() + .with_metadata_options(Some(options)); + decoder.push_range(0..data.len() as u64, data).unwrap(); + let metadata = match decoder.try_decode().unwrap() { + DecodeResult::Data(m) => m, + _ => panic!("could not parse metadata"), + }; + + assert_eq!(expected, metadata); + // the schema pointers should be the same + assert!(Arc::ptr_eq( + &expected.file_metadata().schema_descr_ptr(), + &metadata.file_metadata().schema_descr_ptr() + )); + } +} diff --git a/vendor/parquet/src/file/metadata/parser.rs b/vendor/parquet/src/file/metadata/parser.rs new file mode 100644 index 0000000..9df6bcd --- /dev/null +++ b/vendor/parquet/src/file/metadata/parser.rs @@ -0,0 +1,330 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Internal metadata parsing routines +//! +//! These functions parse thrift-encoded metadata from a byte slice +//! into the corresponding Rust structures + +use crate::errors::ParquetError; +use crate::file::metadata::thrift::parquet_metadata_from_bytes; +use crate::file::metadata::{ + ColumnChunkMetaData, PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions, +}; + +use crate::file::page_index::column_index::ColumnIndexMetaData; +use crate::file::page_index::index_reader::{decode_column_index, decode_offset_index}; +use crate::file::page_index::offset_index::OffsetIndexMetaData; +use bytes::Bytes; + +/// Helper struct for metadata parsing +/// +/// This structure parses thrift-encoded bytes into the correct Rust structs, +/// such as [`ParquetMetaData`], handling decryption if necessary. +// +// Note this structure is used to minimize the number of +// places to add `#[cfg(feature = "encryption")]` checks. +pub(crate) use inner::MetadataParser; + +#[cfg(feature = "encryption")] +mod inner { + use std::sync::Arc; + + use super::*; + use crate::encryption::decrypt::FileDecryptionProperties; + use crate::errors::Result; + + /// API for decoding metadata that may be encrypted + #[derive(Debug, Default)] + pub(crate) struct MetadataParser { + // the credentials and keys needed to decrypt metadata + file_decryption_properties: Option>, + // metadata parsing options + metadata_options: Option>, + } + + impl MetadataParser { + pub(crate) fn new() -> Self { + MetadataParser::default() + } + + pub(crate) fn with_file_decryption_properties( + mut self, + file_decryption_properties: Option>, + ) -> Self { + self.file_decryption_properties = file_decryption_properties; + self + } + + pub(crate) fn with_metadata_options( + self, + options: Option>, + ) -> Self { + Self { + metadata_options: options, + ..self + } + } + + pub(crate) fn decode_metadata( + &self, + buf: &[u8], + encrypted_footer: bool, + ) -> Result { + if encrypted_footer || self.file_decryption_properties.is_some() { + crate::file::metadata::thrift::encryption::parquet_metadata_with_encryption( + self.file_decryption_properties.as_ref(), + encrypted_footer, + buf, + self.metadata_options.as_deref(), + ) + } else { + decode_metadata(buf, self.metadata_options.as_deref()) + } + } + } + + pub(super) fn parse_single_column_index( + bytes: &[u8], + metadata: &ParquetMetaData, + column: &ColumnChunkMetaData, + row_group_index: usize, + col_index: usize, + ) -> crate::errors::Result { + use crate::encryption::decrypt::CryptoContext; + match &column.column_crypto_metadata { + Some(crypto_metadata) => { + let file_decryptor = metadata.file_decryptor.as_ref().ok_or_else(|| { + general_err!("Cannot decrypt column index, no file decryptor set") + })?; + let crypto_context = CryptoContext::for_column( + file_decryptor, + crypto_metadata, + row_group_index, + col_index, + )?; + let column_decryptor = crypto_context.metadata_decryptor(); + let aad = crypto_context.create_column_index_aad()?; + let plaintext = column_decryptor.decrypt(bytes, &aad)?; + decode_column_index(&plaintext, column.column_type()) + } + None => decode_column_index(bytes, column.column_type()), + } + } + + pub(super) fn parse_single_offset_index( + bytes: &[u8], + metadata: &ParquetMetaData, + column: &ColumnChunkMetaData, + row_group_index: usize, + col_index: usize, + ) -> crate::errors::Result { + use crate::encryption::decrypt::CryptoContext; + match &column.column_crypto_metadata { + Some(crypto_metadata) => { + let file_decryptor = metadata.file_decryptor.as_ref().ok_or_else(|| { + general_err!("Cannot decrypt offset index, no file decryptor set") + })?; + let crypto_context = CryptoContext::for_column( + file_decryptor, + crypto_metadata, + row_group_index, + col_index, + )?; + let column_decryptor = crypto_context.metadata_decryptor(); + let aad = crypto_context.create_offset_index_aad()?; + let plaintext = column_decryptor.decrypt(bytes, &aad)?; + decode_offset_index(&plaintext) + } + None => decode_offset_index(bytes), + } + } +} + +#[cfg(not(feature = "encryption"))] +mod inner { + use super::*; + use crate::errors::Result; + use std::sync::Arc; + /// parallel implementation when encryption feature is not enabled + /// + /// This has the same API as the encryption-enabled version + #[derive(Debug, Default)] + pub(crate) struct MetadataParser { + // metadata parsing options + metadata_options: Option>, + } + + impl MetadataParser { + pub(crate) fn new() -> Self { + MetadataParser::default() + } + + pub(crate) fn with_metadata_options( + self, + options: Option>, + ) -> Self { + Self { + metadata_options: options, + } + } + + pub(crate) fn decode_metadata( + &self, + buf: &[u8], + encrypted_footer: bool, + ) -> Result { + if encrypted_footer { + Err(general_err!( + "Parquet file has an encrypted footer but the encryption feature is disabled" + )) + } else { + decode_metadata(buf, self.metadata_options.as_deref()) + } + } + } + + pub(super) fn parse_single_column_index( + bytes: &[u8], + _metadata: &ParquetMetaData, + column: &ColumnChunkMetaData, + _row_group_index: usize, + _col_index: usize, + ) -> crate::errors::Result { + decode_column_index(bytes, column.column_type()) + } + + pub(super) fn parse_single_offset_index( + bytes: &[u8], + _metadata: &ParquetMetaData, + _column: &ColumnChunkMetaData, + _row_group_index: usize, + _col_index: usize, + ) -> crate::errors::Result { + decode_offset_index(bytes) + } +} + +/// Decodes [`ParquetMetaData`] from the provided bytes. +/// +/// Typically this is used to decode the metadata from the end of a parquet +/// file. The format of `buf` is the Thrift compact binary protocol, as specified +/// by the [Parquet Spec]. +/// +/// [Parquet Spec]: https://github.com/apache/parquet-format#metadata +pub(crate) fn decode_metadata( + buf: &[u8], + options: Option<&ParquetMetaDataOptions>, +) -> crate::errors::Result { + parquet_metadata_from_bytes(buf, options) +} + +/// Parses column index from the provided bytes and adds it to the metadata. +/// +/// Arguments +/// * `metadata` - The ParquetMetaData to which the parsed column index will be added. +/// * `column_index_policy` - The policy for handling column index parsing (e.g., +/// Required, Optional, Skip). +/// * `bytes` - The byte slice containing the column index data. +/// * `start_offset` - The offset where `bytes` begin in the file. +pub(crate) fn parse_column_index( + metadata: &mut ParquetMetaData, + column_index_policy: PageIndexPolicy, + bytes: &Bytes, + start_offset: u64, +) -> crate::errors::Result<()> { + if column_index_policy == PageIndexPolicy::Skip { + return Ok(()); + } + let index = metadata + .row_groups() + .iter() + .enumerate() + .map(|(rg_idx, x)| { + x.columns() + .iter() + .enumerate() + .map(|(col_idx, c)| match c.column_index_range() { + Some(r) => { + let r_start = usize::try_from(r.start - start_offset)?; + let r_end = usize::try_from(r.end - start_offset)?; + inner::parse_single_column_index( + &bytes[r_start..r_end], + metadata, + c, + rg_idx, + col_idx, + ) + } + None => Ok(ColumnIndexMetaData::NONE), + }) + .collect::>>() + }) + .collect::>>()?; + + metadata.set_column_index(Some(index)); + Ok(()) +} + +pub(crate) fn parse_offset_index( + metadata: &mut ParquetMetaData, + offset_index_policy: PageIndexPolicy, + bytes: &Bytes, + start_offset: u64, +) -> crate::errors::Result<()> { + if offset_index_policy == PageIndexPolicy::Skip { + return Ok(()); + } + let row_groups = metadata.row_groups(); + let mut all_indexes = Vec::with_capacity(row_groups.len()); + for (rg_idx, x) in row_groups.iter().enumerate() { + let mut row_group_indexes = Vec::with_capacity(x.columns().len()); + for (col_idx, c) in x.columns().iter().enumerate() { + let result = match c.offset_index_range() { + Some(r) => { + let r_start = usize::try_from(r.start - start_offset)?; + let r_end = usize::try_from(r.end - start_offset)?; + inner::parse_single_offset_index( + &bytes[r_start..r_end], + metadata, + c, + rg_idx, + col_idx, + ) + } + None => Err(general_err!("missing offset index")), + }; + + match result { + Ok(index) => row_group_indexes.push(index), + Err(e) => { + if offset_index_policy == PageIndexPolicy::Required { + return Err(e); + } else { + // Invalidate and return + metadata.set_column_index(None); + metadata.set_offset_index(None); + return Ok(()); + } + } + } + } + all_indexes.push(row_group_indexes); + } + metadata.set_offset_index(Some(all_indexes)); + Ok(()) +} diff --git a/vendor/parquet/src/file/metadata/push_decoder.rs b/vendor/parquet/src/file/metadata/push_decoder.rs new file mode 100644 index 0000000..7e4beb5 --- /dev/null +++ b/vendor/parquet/src/file/metadata/push_decoder.rs @@ -0,0 +1,748 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::DecodeResult; +#[cfg(feature = "encryption")] +use crate::encryption::decrypt::FileDecryptionProperties; +use crate::errors::{ParquetError, Result}; +use crate::file::FOOTER_SIZE; +use crate::file::metadata::parser::{MetadataParser, parse_column_index, parse_offset_index}; +use crate::file::metadata::{FooterTail, PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions}; +use crate::file::page_index::index_reader::acc_range; +use crate::file::reader::ChunkReader; +use bytes::Bytes; +use std::ops::Range; +use std::sync::Arc; + +/// A push decoder for [`ParquetMetaData`]. +/// +/// This structure implements a push API for decoding Parquet metadata, which +/// decouples IO from the metadata decoding logic (sometimes referred to as +/// [Sans-IO]). +/// +/// See [`ParquetMetaDataReader`] for a pull-based API that incorporates IO and +/// is simpler to use for basic use cases. This decoder is best for customizing +/// your IO operations to minimize bytes read, prefetch data, or use async IO. +/// +/// [Sans-IO]: https://sans-io.readthedocs.io +/// [`ParquetMetaDataReader`]: crate::file::metadata::ParquetMetaDataReader +/// +/// # Example +/// +/// The most basic usage is to feed the decoder with the necessary byte ranges +/// as requested as shown below. This minimizes the number of bytes read, but +/// requires the most IO operations - one to read the footer and then one +/// to read the metadata, and possibly more if page indexes are requested. +/// +#[cfg_attr( + feature = "arrow", + doc = r##" +```rust +# use std::ops::Range; +# use bytes::Bytes; +# use arrow_array::record_batch; +# use parquet::DecodeResult; +# use parquet::arrow::ArrowWriter; +# use parquet::errors::ParquetError; +# use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataPushDecoder}; +# +# fn decode_metadata() -> Result { +# let file_bytes = { +# let mut buffer = vec![0]; +# let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap(); +# let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap(); +# writer.write(&batch).unwrap(); +# writer.close().unwrap(); +# Bytes::from(buffer) +# }; +# // mimic IO by returning a function that returns the bytes for a given range +# let get_range = |range: &Range| -> Bytes { +# let start = range.start as usize; +# let end = range.end as usize; +# file_bytes.slice(start..end) +# }; +# +# let file_len = file_bytes.len() as u64; +// The `ParquetMetaDataPushDecoder` needs to know the file length. +let mut decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap(); +// try to decode the metadata. If more data is needed, the decoder will tell you what ranges +loop { + match decoder.try_decode() { + Ok(DecodeResult::Data(metadata)) => { return Ok(metadata); } // decode successful + Ok(DecodeResult::NeedsData(ranges)) => { + // The decoder needs more data + // + // In this example, we call a function that returns the bytes for each given range. + // In a real application, you would likely read the data from a file or network. + let data = ranges.iter().map(|range| get_range(range)).collect(); + // Push the data into the decoder and try to decode again on the next iteration. + decoder.push_ranges(ranges, data).unwrap(); + } + Ok(DecodeResult::Finished) => { unreachable!("returned metadata in previous match arm") } + Err(e) => return Err(e), + } +} +# } +``` +"## +)] +/// +/// # Example with "prefetching" +/// +/// By default, the [`ParquetMetaDataPushDecoder`] will request only the exact byte +/// ranges it needs. This minimizes the number of bytes read, however it +/// requires at least two IO operations to read the metadata - one to read the +/// footer and then one to read the metadata. +/// +/// If the file has a "Page Index" (see [Self::with_page_index_policy]), three +/// IO operations are required to read the metadata, as the page index is +/// not part of the normal metadata footer. +/// +/// To reduce the number of IO operations in systems with high per operation +/// overhead (e.g. cloud storage), you can "prefetch" the data and then push +/// the data into the decoder before calling [`Self::try_decode`]. If you do +/// not push enough bytes, the decoder will return the ranges that are still +/// needed. +/// +/// This approach can also be used when you have the entire file already in memory +/// for other reasons. +#[cfg_attr( + feature = "arrow", + doc = r##" +```rust +# use std::ops::Range; +# use bytes::Bytes; +# use arrow_array::record_batch; +# use parquet::DecodeResult; +# use parquet::arrow::ArrowWriter; +# use parquet::errors::ParquetError; +# use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataPushDecoder}; +# +# fn decode_metadata() -> Result { +# let file_bytes = { +# let mut buffer = vec![0]; +# let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap(); +# let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap(); +# writer.write(&batch).unwrap(); +# writer.close().unwrap(); +# Bytes::from(buffer) +# }; +# +let file_len = file_bytes.len() as u64; +// For this example, we "prefetch" all the bytes which we have in memory, +// but in a real application, you would likely read a chunk from the end +// for example 1MB. +let prefetched_bytes = file_bytes.clone(); +let mut decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap(); +// push the prefetched bytes into the decoder +decoder.push_ranges(vec![0..file_len], vec![prefetched_bytes]).unwrap(); +// The decoder will now be able to decode the metadata. Note in a real application, +// unless you can guarantee that the pushed data is enough to decode the metadata, +// you still need to call `try_decode` in a loop until it returns `DecodeResult::Data` +// as shown in the previous example + match decoder.try_decode() { + Ok(DecodeResult::Data(metadata)) => { return Ok(metadata); } // decode successful + other => { panic!("expected DecodeResult::Data, got: {other:?}") } + } +# } +``` +"## +)] +/// +/// # Example using [`AsyncRead`] +/// +/// [`ParquetMetaDataPushDecoder`] is designed to work with any data source that can +/// provide byte ranges, including async IO sources. However, it does not +/// implement async IO itself. To use async IO, you simply write an async +/// wrapper around it that reads the required byte ranges and pushes them into the +/// decoder. +#[cfg_attr( + feature = "arrow", + doc = r##" +```rust +# use std::ops::Range; +# use bytes::Bytes; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt}; +# use arrow_array::record_batch; +# use parquet::DecodeResult; +# use parquet::arrow::ArrowWriter; +# use parquet::errors::ParquetError; +# use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataPushDecoder}; +# +// This function decodes Parquet Metadata from anything that implements +// [`AsyncRead`] and [`AsyncSeek`] such as a tokio::fs::File +async fn decode_metadata( + file_len: u64, + mut async_source: impl AsyncRead + AsyncSeek + Unpin +) -> Result { + // We need a ParquetMetaDataPushDecoder to decode the metadata. + let mut decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap(); + loop { + match decoder.try_decode() { + Ok(DecodeResult::Data(metadata)) => { return Ok(metadata); } // decode successful + Ok(DecodeResult::NeedsData(ranges)) => { + // The decoder needs more data + // + // In this example we use the AsyncRead and AsyncSeek traits to read the + // required ranges from the async source. + let mut data = Vec::with_capacity(ranges.len()); + for range in &ranges { + let mut buffer = vec![0; (range.end - range.start) as usize]; + async_source.seek(std::io::SeekFrom::Start(range.start)).await?; + async_source.read_exact(&mut buffer).await?; + data.push(Bytes::from(buffer)); + } + // Push the data into the decoder and try to decode again on the next iteration. + decoder.push_ranges(ranges, data).unwrap(); + } + Ok(DecodeResult::Finished) => { unreachable!("returned metadata in previous match arm") } + Err(e) => return Err(e), + } + } +} +``` +"## +)] +/// [`AsyncRead`]: tokio::io::AsyncRead +#[derive(Debug)] +pub struct ParquetMetaDataPushDecoder { + /// Decoding state + state: DecodeState, + /// policy for loading ColumnIndex (part of the PageIndex) + column_index_policy: PageIndexPolicy, + /// policy for loading OffsetIndex (part of the PageIndex) + offset_index_policy: PageIndexPolicy, + /// Underlying buffers + buffers: crate::util::push_buffers::PushBuffers, + /// Encryption API + metadata_parser: MetadataParser, +} + +impl ParquetMetaDataPushDecoder { + /// Create a new `ParquetMetaDataPushDecoder` with the given file length. + /// + /// By default, this will read page indexes and column indexes. See + /// [`ParquetMetaDataPushDecoder::with_page_index_policy`] for more detail. + /// + /// See examples on [`ParquetMetaDataPushDecoder`]. + pub fn try_new(file_len: u64) -> Result { + if file_len < 8 { + return Err(ParquetError::General(format!( + "Parquet files are at least 8 bytes long, but file length is {file_len}" + ))); + }; + + Ok(Self { + state: DecodeState::ReadingFooter, + column_index_policy: PageIndexPolicy::Optional, + offset_index_policy: PageIndexPolicy::Optional, + buffers: crate::util::push_buffers::PushBuffers::new(file_len), + metadata_parser: MetadataParser::new(), + }) + } + + /// Begin decoding from the given footer tail. + pub(crate) fn try_new_with_footer_tail(file_len: u64, footer_tail: FooterTail) -> Result { + let mut new_self = Self::try_new(file_len)?; + new_self.state = DecodeState::ReadingMetadata(footer_tail); + Ok(new_self) + } + + /// Create a decoder with the given `ParquetMetaData` already known. + /// + /// This can be used to parse and populate the page index structures + /// after the metadata has already been decoded. + pub fn try_new_with_metadata(file_len: u64, metadata: ParquetMetaData) -> Result { + let mut new_self = Self::try_new(file_len)?; + new_self.state = DecodeState::ReadingPageIndex(Box::new(metadata)); + Ok(new_self) + } + + /// Enable or disable reading the page index structures described in + /// "[Parquet page index] Layout to Support Page Skipping". + /// + /// Defaults to [`PageIndexPolicy::Optional`] + /// + /// This requires + /// 1. The Parquet file to have been written with page indexes + /// 2. Additional data to be pushed into the decoder (as the page indexes are not part of the thrift footer) + /// + /// [Parquet page index]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + pub fn with_page_index_policy(mut self, page_index_policy: PageIndexPolicy) -> Self { + self.column_index_policy = page_index_policy; + self.offset_index_policy = page_index_policy; + self + } + + /// Set the policy for reading the ColumnIndex (part of the PageIndex) + pub fn with_column_index_policy(mut self, column_index_policy: PageIndexPolicy) -> Self { + self.column_index_policy = column_index_policy; + self + } + + /// Set the policy for reading the OffsetIndex (part of the PageIndex) + pub fn with_offset_index_policy(mut self, offset_index_policy: PageIndexPolicy) -> Self { + self.offset_index_policy = offset_index_policy; + self + } + + /// Set the options to use when decoding the Parquet metadata. + pub fn with_metadata_options(mut self, options: Option>) -> Self { + self.metadata_parser = self.metadata_parser.with_metadata_options(options); + self + } + + #[cfg(feature = "encryption")] + /// Provide decryption properties for decoding encrypted Parquet files + pub fn with_file_decryption_properties( + mut self, + file_decryption_properties: Option>, + ) -> Self { + self.metadata_parser = self + .metadata_parser + .with_file_decryption_properties(file_decryption_properties); + self + } + + /// Push the data into the decoder's buffer. + /// + /// The decoder does not immediately attempt to decode the metadata + /// after pushing data. Instead, it accumulates the pushed data until you + /// call [`Self::try_decode`]. + /// + /// # Determining required data: + /// + /// To determine what ranges are required to decode the metadata, you can + /// either: + /// + /// 1. Call [`Self::try_decode`] first to get the exact ranges required (see + /// example on [`Self`]) + /// + /// 2. Speculatively push any data that you have available, which may + /// include more than the footer data or requested bytes. + /// + /// Speculatively pushing data can be used when "prefetching" data. See + /// example on [`Self`] + pub fn push_ranges(&mut self, ranges: Vec>, buffers: Vec) -> Result<()> { + if matches!(&self.state, DecodeState::Finished) { + return Err(general_err!( + "ParquetMetaDataPushDecoder: cannot push data after decoding is finished" + )); + } + self.buffers.push_ranges(ranges, buffers); + Ok(()) + } + + /// Pushes a single range of data into the decoder's buffer. + pub fn push_range(&mut self, range: Range, buffer: Bytes) -> Result<()> { + if matches!(&self.state, DecodeState::Finished) { + return Err(general_err!( + "ParquetMetaDataPushDecoder: cannot push data after decoding is finished" + )); + } + self.buffers.push_range(range, buffer); + Ok(()) + } + + /// Clear any staged byte ranges currently buffered for future decode work. + pub fn clear_all_ranges(&mut self) { + self.buffers.clear_all_ranges(); + } + + /// Try to decode the metadata from the pushed data, returning the + /// decoded metadata or an error if not enough data is available. + pub fn try_decode(&mut self) -> Result> { + let file_len = self.buffers.file_len(); + let footer_len = FOOTER_SIZE as u64; + loop { + match std::mem::replace(&mut self.state, DecodeState::Intermediate) { + DecodeState::ReadingFooter => { + // need to have the last 8 bytes of the file to decode the metadata + let footer_start = file_len.saturating_sub(footer_len); + let footer_range = footer_start..file_len; + + if !self.buffers.has_range(&footer_range) { + self.state = DecodeState::ReadingFooter; + return Ok(needs_range(footer_range)); + } + let footer_bytes = self.get_bytes(&footer_range)?; + let footer_tail = FooterTail::try_from(footer_bytes.as_ref())?; + + self.state = DecodeState::ReadingMetadata(footer_tail); + continue; + } + + DecodeState::ReadingMetadata(footer_tail) => { + let metadata_len: u64 = footer_tail.metadata_length() as u64; + let metadata_start = file_len - footer_len - metadata_len; + let metadata_end = metadata_start + metadata_len; + let metadata_range = metadata_start..metadata_end; + + if !self.buffers.has_range(&metadata_range) { + self.state = DecodeState::ReadingMetadata(footer_tail); + return Ok(needs_range(metadata_range)); + } + + let metadata = self.metadata_parser.decode_metadata( + &self.get_bytes(&metadata_range)?, + footer_tail.is_encrypted_footer(), + )?; + // Note: ReadingPageIndex first checks if page indexes are needed + // and is a no-op if not + self.state = DecodeState::ReadingPageIndex(Box::new(metadata)); + continue; + } + + DecodeState::ReadingPageIndex(mut metadata) => { + // First determine if any page indexes are needed based on + // the specified policies + let range = range_for_page_index( + &metadata, + self.column_index_policy, + self.offset_index_policy, + ); + + let Some(page_index_range) = range else { + self.state = DecodeState::Finished; + return Ok(DecodeResult::Data(*metadata)); + }; + + if !self.buffers.has_range(&page_index_range) { + self.state = DecodeState::ReadingPageIndex(metadata); + return Ok(needs_range(page_index_range)); + } + + let buffer = self.get_bytes(&page_index_range)?; + let offset = page_index_range.start; + parse_column_index(&mut metadata, self.column_index_policy, &buffer, offset)?; + parse_offset_index(&mut metadata, self.offset_index_policy, &buffer, offset)?; + self.state = DecodeState::Finished; + return Ok(DecodeResult::Data(*metadata)); + } + + DecodeState::Finished => return Ok(DecodeResult::Finished), + DecodeState::Intermediate => { + return Err(general_err!( + "ParquetMetaDataPushDecoder: internal error, invalid state" + )); + } + } + } + } + + /// Returns the bytes for the given range from the internal buffer + fn get_bytes(&self, range: &Range) -> Result { + let start = range.start; + let raw_len = range.end - range.start; + let len: usize = raw_len.try_into().map_err(|_| { + ParquetError::General(format!( + "ParquetMetaDataPushDecoder: Range length too large to fit in usize: {raw_len}", + )) + })?; + self.buffers.get_bytes(start, len) + } +} + +/// returns a DecodeResults that describes needing the given range +fn needs_range(range: Range) -> DecodeResult { + DecodeResult::NeedsData(vec![range]) +} + +/// Decoding state machine +#[derive(Debug)] +enum DecodeState { + /// Reading the last 8 bytes of the file + ReadingFooter, + /// Reading the metadata thrift structure + ReadingMetadata(FooterTail), + // Actively reading the page index + ReadingPageIndex(Box), + // Decoding is complete + Finished, + /// State left during the `try_decode` method so something valid is present. + /// This state should never be observed. + Intermediate, +} + +/// Returns the byte range needed to read the offset/page indexes, based on the +/// specified policies +/// +/// Returns None if no page indexes are needed +pub fn range_for_page_index( + metadata: &ParquetMetaData, + column_index_policy: PageIndexPolicy, + offset_index_policy: PageIndexPolicy, +) -> Option> { + let mut range = None; + for c in metadata.row_groups().iter().flat_map(|r| r.columns()) { + if column_index_policy != PageIndexPolicy::Skip { + range = acc_range(range, c.column_index_range()); + } + if offset_index_policy != PageIndexPolicy::Skip { + range = acc_range(range, c.offset_index_range()); + } + } + range +} + +// These tests use the arrow writer to create a parquet file in memory +// so they need the arrow feature and the test feature +#[cfg(all(test, feature = "arrow"))] +mod tests { + use super::*; + use crate::arrow::ArrowWriter; + use crate::file::properties::WriterProperties; + use arrow_array::{ArrayRef, Int64Array, RecordBatch, StringViewArray}; + use bytes::Bytes; + use std::fmt::Debug; + use std::ops::Range; + use std::sync::{Arc, LazyLock}; + + /// It is possible to decode the metadata from the entire file at once before being asked + #[test] + fn test_metadata_decoder_all_data() { + let file_len = test_file_len(); + let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap(); + // Push the entire file data into the metadata decoder + push_ranges_to_metadata_decoder(&mut metadata_decoder, vec![test_file_range()]); + + // should be able to decode the metadata without needing more data + let metadata = expect_data(metadata_decoder.try_decode()); + + assert_eq!(metadata.num_row_groups(), 2); + assert_eq!(metadata.row_group(0).num_rows(), 200); + assert_eq!(metadata.row_group(1).num_rows(), 200); + assert!(metadata.column_index().is_some()); + assert!(metadata.offset_index().is_some()); + } + + /// It is possible to feed some, but not all, of the footer into the metadata decoder + /// before asked. This avoids multiple IO requests + #[test] + fn test_metadata_decoder_prefetch_success() { + let file_len = test_file_len(); + let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap(); + // simulate pre-fetching the last 2k bytes of the file without asking the decoder + let prefetch_range = (file_len - 2 * 1024)..file_len; + push_ranges_to_metadata_decoder(&mut metadata_decoder, vec![prefetch_range]); + + // expect the decoder has enough data to decode the metadata + let metadata = expect_data(metadata_decoder.try_decode()); + expect_finished(metadata_decoder.try_decode()); + assert_eq!(metadata.num_row_groups(), 2); + assert_eq!(metadata.row_group(0).num_rows(), 200); + assert_eq!(metadata.row_group(1).num_rows(), 200); + assert!(metadata.column_index().is_some()); + assert!(metadata.offset_index().is_some()); + } + + /// It is possible to pre-fetch some, but not all, of the necessary data + /// data + #[test] + fn test_metadata_decoder_prefetch_retry() { + let file_len = test_file_len(); + let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap(); + // simulate pre-fetching the last 1500 bytes of the file. + // this is enough to read the footer thrift metadata, but not the offset indexes + let prefetch_range = (file_len - 1500)..file_len; + push_ranges_to_metadata_decoder(&mut metadata_decoder, vec![prefetch_range]); + + // expect another request is needed to read the offset indexes (note + // try_decode only returns NeedsData once, whereas without any prefetching it would + // return NeedsData three times) + let ranges = expect_needs_data(metadata_decoder.try_decode()); + push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges); + + // expect the decoder has enough data to decode the metadata + let metadata = expect_data(metadata_decoder.try_decode()); + expect_finished(metadata_decoder.try_decode()); + + assert_eq!(metadata.num_row_groups(), 2); + assert_eq!(metadata.row_group(0).num_rows(), 200); + assert_eq!(metadata.row_group(1).num_rows(), 200); + assert!(metadata.column_index().is_some()); + assert!(metadata.offset_index().is_some()); + } + + #[test] + fn test_metadata_decoder_clear_all_ranges() { + let file_len = test_file_len(); + let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap(); + + metadata_decoder + .push_range(test_file_range(), TEST_FILE_DATA.clone()) + .unwrap(); + assert_eq!(metadata_decoder.buffers.buffered_bytes(), test_file_len()); + + metadata_decoder.clear_all_ranges(); + assert_eq!(metadata_decoder.buffers.buffered_bytes(), 0); + + let ranges = expect_needs_data(metadata_decoder.try_decode()); + assert_eq!(ranges, vec![test_file_len() - 8..test_file_len()]); + } + + /// Decode the metadata incrementally, simulating a scenario where exactly the data needed + /// is read in each step + #[test] + fn test_metadata_decoder_incremental() { + let file_len = TEST_FILE_DATA.len() as u64; + let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap(); + let ranges = expect_needs_data(metadata_decoder.try_decode()); + assert_eq!(ranges.len(), 1); + assert_eq!(ranges[0], test_file_len() - 8..test_file_len()); + push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges); + + // expect the first request to read the footer + let ranges = expect_needs_data(metadata_decoder.try_decode()); + push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges); + + // expect the second request to read the offset indexes + let ranges = expect_needs_data(metadata_decoder.try_decode()); + push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges); + + // expect the third request to read the actual data + let metadata = expect_data(metadata_decoder.try_decode()); + expect_finished(metadata_decoder.try_decode()); + + assert_eq!(metadata.num_row_groups(), 2); + assert_eq!(metadata.row_group(0).num_rows(), 200); + assert_eq!(metadata.row_group(1).num_rows(), 200); + assert!(metadata.column_index().is_some()); + assert!(metadata.offset_index().is_some()); + } + + /// Decode the metadata incrementally, but without reading the page indexes + /// (so only two requests) + #[test] + fn test_metadata_decoder_incremental_no_page_index() { + let file_len = TEST_FILE_DATA.len() as u64; + let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_len) + .unwrap() + .with_page_index_policy(PageIndexPolicy::Skip); + let ranges = expect_needs_data(metadata_decoder.try_decode()); + assert_eq!(ranges.len(), 1); + assert_eq!(ranges[0], test_file_len() - 8..test_file_len()); + push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges); + + // expect the first request to read the footer + let ranges = expect_needs_data(metadata_decoder.try_decode()); + push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges); + + // expect NO second request to read the offset indexes, should just cough up the metadata + let metadata = expect_data(metadata_decoder.try_decode()); + expect_finished(metadata_decoder.try_decode()); + + assert_eq!(metadata.num_row_groups(), 2); + assert_eq!(metadata.row_group(0).num_rows(), 200); + assert_eq!(metadata.row_group(1).num_rows(), 200); + assert!(metadata.column_index().is_none()); // of course, we did not read the column index + assert!(metadata.offset_index().is_none()); // or the offset index + } + + static TEST_BATCH: LazyLock = LazyLock::new(|| { + // Input batch has 400 rows, with 3 columns: "a", "b", "c" + // Note c is a different types (so the data page sizes will be different) + let a: ArrayRef = Arc::new(Int64Array::from_iter_values(0..400)); + let b: ArrayRef = Arc::new(Int64Array::from_iter_values(400..800)); + let c: ArrayRef = Arc::new(StringViewArray::from_iter_values((0..400).map(|i| { + if i % 2 == 0 { + format!("string_{i}") + } else { + format!("A string larger than 12 bytes and thus not inlined {i}") + } + }))); + + RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap() + }); + + /// Create a parquet file in memory for testing. See [`test_file_range`] for details. + static TEST_FILE_DATA: LazyLock = LazyLock::new(|| { + let input_batch = &TEST_BATCH; + let mut output = Vec::new(); + + let writer_options = WriterProperties::builder() + .set_max_row_group_row_count(Some(200)) + .set_data_page_row_count_limit(100) + .build(); + let mut writer = + ArrowWriter::try_new(&mut output, input_batch.schema(), Some(writer_options)).unwrap(); + + // since the limits are only enforced on batch boundaries, write the input + // batch in chunks of 50 + let mut row_remain = input_batch.num_rows(); + while row_remain > 0 { + let chunk_size = row_remain.min(50); + let chunk = input_batch.slice(input_batch.num_rows() - row_remain, chunk_size); + writer.write(&chunk).unwrap(); + row_remain -= chunk_size; + } + writer.close().unwrap(); + Bytes::from(output) + }); + + /// Return the length of the test file in bytes + fn test_file_len() -> u64 { + TEST_FILE_DATA.len() as u64 + } + + /// Return the range of the entire test file + fn test_file_range() -> Range { + 0..test_file_len() + } + + /// Return a slice of the test file data from the given range + pub fn test_file_slice(range: Range) -> Bytes { + let start: usize = range.start.try_into().unwrap(); + let end: usize = range.end.try_into().unwrap(); + TEST_FILE_DATA.slice(start..end) + } + + /// Push the given ranges to the metadata decoder, simulating reading from a file + fn push_ranges_to_metadata_decoder( + metadata_decoder: &mut ParquetMetaDataPushDecoder, + ranges: Vec>, + ) { + let data = ranges + .iter() + .map(|range| test_file_slice(range.clone())) + .collect::>(); + metadata_decoder.push_ranges(ranges, data).unwrap(); + } + + /// Expect that the [`DecodeResult`] is a [`DecodeResult::Data`] and return the corresponding element + fn expect_data(result: Result>) -> T { + match result.expect("Expected Ok(DecodeResult::Data(T))") { + DecodeResult::Data(data) => data, + result => panic!("Expected DecodeResult::Data, got {result:?}"), + } + } + + /// Expect that the [`DecodeResult`] is a [`DecodeResult::NeedsData`] and return the corresponding ranges + fn expect_needs_data(result: Result>) -> Vec> { + match result.expect("Expected Ok(DecodeResult::NeedsData{ranges})") { + DecodeResult::NeedsData(ranges) => ranges, + result => panic!("Expected DecodeResult::NeedsData, got {result:?}"), + } + } + + fn expect_finished(result: Result>) { + match result.expect("Expected Ok(DecodeResult::Finished)") { + DecodeResult::Finished => {} + result => panic!("Expected DecodeResult::Finished, got {result:?}"), + } + } +} diff --git a/vendor/parquet/src/file/metadata/reader.rs b/vendor/parquet/src/file/metadata/reader.rs new file mode 100644 index 0000000..844ae74 --- /dev/null +++ b/vendor/parquet/src/file/metadata/reader.rs @@ -0,0 +1,1399 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#[cfg(feature = "encryption")] +use crate::encryption::decrypt::FileDecryptionProperties; +use crate::errors::{ParquetError, Result}; +use crate::file::FOOTER_SIZE; +use crate::file::metadata::parser::decode_metadata; +use crate::file::metadata::thrift::parquet_schema_from_bytes; +use crate::file::metadata::{ + FooterTail, ParquetMetaData, ParquetMetaDataOptions, ParquetMetaDataPushDecoder, +}; +use crate::file::reader::ChunkReader; +use crate::schema::types::SchemaDescriptor; +use bytes::Bytes; +use std::sync::Arc; +use std::{io::Read, ops::Range}; + +use crate::DecodeResult; +#[cfg(all(feature = "async", feature = "arrow"))] +use crate::arrow::async_reader::{MetadataFetch, MetadataSuffixFetch}; + +/// Reads [`ParquetMetaData`] from a byte stream, with either synchronous or +/// asynchronous I/O. +/// +/// There are two flavors of APIs: +/// * Synchronous: [`Self::try_parse()`], [`Self::try_parse_sized()`], [`Self::parse_and_finish()`], etc. +/// * Asynchronous (requires `async` and `arrow` features): [`Self::try_load()`], etc +/// +/// See the [`ParquetMetaDataPushDecoder`] for an API that does not require I/O. +/// +/// # Format Notes +/// +/// Parquet metadata is not necessarily contiguous in a Parquet file: a portion is stored +/// in the footer (the last bytes of the file), but other portions (such as the +/// PageIndex) can be stored elsewhere. +/// See [`crate::file::metadata::ParquetMetaDataWriter#output-format`] for more details of +/// Parquet metadata. +/// +/// This reader handles reading the footer as well as the non contiguous parts +/// of the metadata (`PageIndex` and `ColumnIndex`). It does not handle reading Bloom Filters. +/// +/// # Example +/// ```no_run +/// # use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader}; +/// # fn open_parquet_file(path: &str) -> std::fs::File { unimplemented!(); } +/// // read parquet metadata including page indexes from a file +/// let file = open_parquet_file("some_path.parquet"); +/// let mut reader = ParquetMetaDataReader::new() +/// .with_page_index_policy(PageIndexPolicy::Required); +/// reader.try_parse(&file).unwrap(); +/// let metadata = reader.finish().unwrap(); +/// assert!(metadata.column_index().is_some()); +/// assert!(metadata.offset_index().is_some()); +/// ``` +#[derive(Default, Debug)] +pub struct ParquetMetaDataReader { + metadata: Option, + column_index: PageIndexPolicy, + offset_index: PageIndexPolicy, + prefetch_hint: Option, + metadata_options: Option>, + // Size of the serialized thrift metadata plus the 8 byte footer. Only set if + // `self.parse_metadata` is called. + metadata_size: Option, + #[cfg(feature = "encryption")] + file_decryption_properties: Option>, +} + +/// Describes the policy for reading page indexes +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum PageIndexPolicy { + /// Do not read the page index. + #[default] + Skip, + /// Read the page index if it exists, otherwise do not error. + Optional, + /// Require the page index to exist, and error if it does not. + Required, +} + +impl From for PageIndexPolicy { + fn from(value: bool) -> Self { + match value { + true => Self::Required, + false => Self::Skip, + } + } +} + +impl ParquetMetaDataReader { + /// Create a new [`ParquetMetaDataReader`] + pub fn new() -> Self { + Default::default() + } + + /// Create a new [`ParquetMetaDataReader`] populated with a [`ParquetMetaData`] struct + /// obtained via other means. + pub fn new_with_metadata(metadata: ParquetMetaData) -> Self { + Self { + metadata: Some(metadata), + ..Default::default() + } + } + + /// Sets the [`PageIndexPolicy`] for the column and offset indexes + pub fn with_page_index_policy(self, policy: PageIndexPolicy) -> Self { + self.with_column_index_policy(policy) + .with_offset_index_policy(policy) + } + + /// Sets the [`PageIndexPolicy`] for the column index + pub fn with_column_index_policy(mut self, policy: PageIndexPolicy) -> Self { + self.column_index = policy; + self + } + + /// Sets the [`PageIndexPolicy`] for the offset index + pub fn with_offset_index_policy(mut self, policy: PageIndexPolicy) -> Self { + self.offset_index = policy; + self + } + + /// Sets the [`ParquetMetaDataOptions`] to use when decoding + pub fn with_metadata_options(mut self, options: Option) -> Self { + self.metadata_options = options.map(Arc::new); + self + } + + /// Provide a hint as to the number of bytes needed to fully parse the [`ParquetMetaData`]. + /// Only used for the asynchronous [`Self::try_load()`] method. + /// + /// By default, the reader will first fetch the last 8 bytes of the input file to obtain the + /// size of the footer metadata. A second fetch will be performed to obtain the needed bytes. + /// After parsing the footer metadata, a third fetch will be performed to obtain the bytes + /// needed to decode the page index structures, if they have been requested. To avoid + /// unnecessary fetches, `prefetch` can be set to an estimate of the number of bytes needed + /// to fully decode the [`ParquetMetaData`], which can reduce the number of fetch requests and + /// reduce latency. Setting `prefetch` too small will not trigger an error, but will result + /// in extra fetches being performed. + pub fn with_prefetch_hint(mut self, prefetch: Option) -> Self { + self.prefetch_hint = prefetch; + self + } + + /// Provide the FileDecryptionProperties to use when decrypting the file. + /// + /// This is only necessary when the file is encrypted. + #[cfg(feature = "encryption")] + pub fn with_decryption_properties( + mut self, + properties: Option>, + ) -> Self { + self.file_decryption_properties = properties; + self + } + + /// Indicates whether this reader has a [`ParquetMetaData`] internally. + pub fn has_metadata(&self) -> bool { + self.metadata.is_some() + } + + /// Return the parsed [`ParquetMetaData`] struct, leaving `None` in its place. + pub fn finish(&mut self) -> Result { + self.metadata + .take() + .ok_or_else(|| general_err!("could not parse parquet metadata")) + } + + /// Given a [`ChunkReader`], parse and return the [`ParquetMetaData`] in a single pass. + /// + /// If `reader` is [`Bytes`] based, then the buffer must contain sufficient bytes to complete + /// the request, and must include the Parquet footer. If page indexes are desired, the buffer + /// must contain the entire file, or [`Self::try_parse_sized()`] should be used. + /// + /// This call will consume `self`. + /// + /// # Example + /// ```no_run + /// # use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader}; + /// # fn open_parquet_file(path: &str) -> std::fs::File { unimplemented!(); } + /// // read parquet metadata including page indexes + /// let file = open_parquet_file("some_path.parquet"); + /// let metadata = ParquetMetaDataReader::new() + /// .with_page_index_policy(PageIndexPolicy::Required) + /// .parse_and_finish(&file).unwrap(); + /// ``` + pub fn parse_and_finish(mut self, reader: &R) -> Result { + self.try_parse(reader)?; + self.finish() + } + + /// Attempts to parse the footer metadata (and optionally page indexes) given a [`ChunkReader`]. + /// + /// If `reader` is [`Bytes`] based, then the buffer must contain sufficient bytes to complete + /// the request, and must include the Parquet footer. If page indexes are desired, the buffer + /// must contain the entire file, or [`Self::try_parse_sized()`] should be used. + pub fn try_parse(&mut self, reader: &R) -> Result<()> { + self.try_parse_sized(reader, reader.len()) + } + + /// Same as [`Self::try_parse()`], but provide the original file size in the case that `reader` + /// is a [`Bytes`] struct that does not contain the entire file. This information is necessary + /// when the page indexes are desired. `reader` must have access to the Parquet footer. + /// + /// Using this function also allows for retrying with a larger buffer. + /// + /// # Errors + /// + /// This function will return [`ParquetError::NeedMoreData`] in the event `reader` does not + /// provide enough data to fully parse the metadata (see example below). The returned error + /// will be populated with a `usize` field indicating the number of bytes required from the + /// tail of the file to completely parse the requested metadata. + /// + /// Other errors returned include [`ParquetError::General`] and [`ParquetError::EOF`]. + /// + /// # Example + /// ```no_run + /// # use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader}; + /// # use parquet::errors::ParquetError; + /// # use crate::parquet::file::reader::Length; + /// # fn get_bytes(file: &std::fs::File, range: std::ops::Range) -> bytes::Bytes { unimplemented!(); } + /// # fn open_parquet_file(path: &str) -> std::fs::File { unimplemented!(); } + /// let file = open_parquet_file("some_path.parquet"); + /// let len = file.len(); + /// // Speculatively read 1 kilobyte from the end of the file + /// let bytes = get_bytes(&file, len - 1024..len); + /// let mut reader = ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required); + /// match reader.try_parse_sized(&bytes, len) { + /// Ok(_) => (), + /// Err(ParquetError::NeedMoreData(needed)) => { + /// // Read the needed number of bytes from the end of the file + /// let bytes = get_bytes(&file, len - needed as u64..len); + /// reader.try_parse_sized(&bytes, len).unwrap(); + /// } + /// _ => panic!("unexpected error") + /// } + /// let metadata = reader.finish().unwrap(); + /// ``` + /// + /// Note that it is possible for the file metadata to be completely read, but there are + /// insufficient bytes available to read the page indexes. [`Self::has_metadata()`] can be used + /// to test for this. In the event the file metadata is present, re-parsing of the file + /// metadata can be skipped by using [`Self::read_page_indexes_sized()`], as shown below. + /// ```no_run + /// # use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader}; + /// # use parquet::errors::ParquetError; + /// # use crate::parquet::file::reader::Length; + /// # fn get_bytes(file: &std::fs::File, range: std::ops::Range) -> bytes::Bytes { unimplemented!(); } + /// # fn open_parquet_file(path: &str) -> std::fs::File { unimplemented!(); } + /// let file = open_parquet_file("some_path.parquet"); + /// let len = file.len(); + /// // Speculatively read 1 kilobyte from the end of the file + /// let mut bytes = get_bytes(&file, len - 1024..len); + /// let mut reader = ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required); + /// // Loop until `bytes` is large enough + /// loop { + /// match reader.try_parse_sized(&bytes, len) { + /// Ok(_) => break, + /// Err(ParquetError::NeedMoreData(needed)) => { + /// // Read the needed number of bytes from the end of the file + /// bytes = get_bytes(&file, len - needed as u64..len); + /// // If file metadata was read only read page indexes, otherwise continue loop + /// if reader.has_metadata() { + /// reader.read_page_indexes_sized(&bytes, len).unwrap(); + /// break; + /// } + /// } + /// _ => panic!("unexpected error") + /// } + /// } + /// let metadata = reader.finish().unwrap(); + /// ``` + pub fn try_parse_sized(&mut self, reader: &R, file_size: u64) -> Result<()> { + self.metadata = match self.parse_metadata(reader) { + Ok(metadata) => Some(metadata), + Err(ParquetError::NeedMoreData(needed)) => { + // If reader is the same length as `file_size` then presumably there is no more to + // read, so return an EOF error. + if file_size == reader.len() || needed as u64 > file_size { + return Err(eof_err!( + "Parquet file too small. Size is {} but need {}", + file_size, + needed + )); + } else { + // Ask for a larger buffer + return Err(ParquetError::NeedMoreData(needed)); + } + } + Err(e) => return Err(e), + }; + + // we can return if page indexes aren't requested + if self.column_index == PageIndexPolicy::Skip && self.offset_index == PageIndexPolicy::Skip + { + return Ok(()); + } + + self.read_page_indexes_sized(reader, file_size) + } + + /// Read the page index structures when a [`ParquetMetaData`] has already been obtained. + /// See [`Self::new_with_metadata()`] and [`Self::has_metadata()`]. + pub fn read_page_indexes(&mut self, reader: &R) -> Result<()> { + self.read_page_indexes_sized(reader, reader.len()) + } + + /// Read the page index structures when a [`ParquetMetaData`] has already been obtained. + /// This variant is used when `reader` cannot access the entire Parquet file (e.g. it is + /// a [`Bytes`] struct containing the tail of the file). + /// See [`Self::new_with_metadata()`] and [`Self::has_metadata()`]. Like + /// [`Self::try_parse_sized()`] this function may return [`ParquetError::NeedMoreData`]. + pub fn read_page_indexes_sized( + &mut self, + reader: &R, + file_size: u64, + ) -> Result<()> { + let Some(metadata) = self.metadata.take() else { + return Err(general_err!( + "Tried to read page indexes without ParquetMetaData metadata" + )); + }; + + let push_decoder = ParquetMetaDataPushDecoder::try_new_with_metadata(file_size, metadata)? + .with_offset_index_policy(self.offset_index) + .with_column_index_policy(self.column_index) + .with_metadata_options(self.metadata_options.clone()); + let mut push_decoder = self.prepare_push_decoder(push_decoder); + + // Get bounds needed for page indexes (if any are present in the file). + let range = match needs_index_data(&mut push_decoder)? { + NeedsIndexData::No(metadata) => { + self.metadata = Some(metadata); + return Ok(()); + } + NeedsIndexData::Yes(range) => range, + }; + + // Check to see if needed range is within `file_range`. Checking `range.end` seems + // redundant, but it guards against `range_for_page_index()` returning garbage. + let file_range = file_size.saturating_sub(reader.len())..file_size; + if !(file_range.contains(&range.start) && file_range.contains(&range.end)) { + // Requested range starts beyond EOF + if range.end > file_size { + return Err(eof_err!( + "Parquet file too small. Range {range:?} is beyond file bounds {file_size}", + )); + } else { + // Ask for a larger buffer + return Err(ParquetError::NeedMoreData( + (file_size - range.start).try_into()?, + )); + } + } + + // Perform extra sanity check to make sure `range` and the footer metadata don't + // overlap. + if let Some(metadata_size) = self.metadata_size { + let metadata_range = file_size.saturating_sub(metadata_size as u64)..file_size; + if range.end > metadata_range.start { + return Err(eof_err!( + "Parquet file too small. Page index range {range:?} overlaps with file metadata {metadata_range:?}", + )); + } + } + + // add the needed ranges to the decoder + let bytes_needed = usize::try_from(range.end - range.start)?; + let bytes = reader.get_bytes(range.start - file_range.start, bytes_needed)?; + + push_decoder.push_range(range, bytes)?; + let metadata = parse_index_data(&mut push_decoder)?; + self.metadata = Some(metadata); + + Ok(()) + } + + /// Given a [`MetadataFetch`], parse and return the [`ParquetMetaData`] in a single pass. + /// + /// This call will consume `self`. + /// + /// See [`Self::with_prefetch_hint`] for a discussion of how to reduce the number of fetches + /// performed by this function. + #[cfg(all(feature = "async", feature = "arrow"))] + pub async fn load_and_finish( + mut self, + fetch: F, + file_size: u64, + ) -> Result { + self.try_load(fetch, file_size).await?; + self.finish() + } + + /// Given a [`MetadataSuffixFetch`], parse and return the [`ParquetMetaData`] in a single pass. + /// + /// This call will consume `self`. + /// + /// See [`Self::with_prefetch_hint`] for a discussion of how to reduce the number of fetches + /// performed by this function. + #[cfg(all(feature = "async", feature = "arrow"))] + pub async fn load_via_suffix_and_finish( + mut self, + fetch: F, + ) -> Result { + self.try_load_via_suffix(fetch).await?; + self.finish() + } + /// Attempts to (asynchronously) parse the footer metadata (and optionally page indexes) + /// given a [`MetadataFetch`]. + /// + /// See [`Self::with_prefetch_hint`] for a discussion of how to reduce the number of fetches + /// performed by this function. + #[cfg(all(feature = "async", feature = "arrow"))] + pub async fn try_load(&mut self, mut fetch: F, file_size: u64) -> Result<()> { + let (metadata, remainder) = self.load_metadata(&mut fetch, file_size).await?; + + self.metadata = Some(metadata); + + // we can return if page indexes aren't requested + if self.column_index == PageIndexPolicy::Skip && self.offset_index == PageIndexPolicy::Skip + { + return Ok(()); + } + + self.load_page_index_with_remainder(fetch, remainder).await + } + + /// Attempts to (asynchronously) parse the footer metadata (and optionally page indexes) + /// given a [`MetadataSuffixFetch`]. + /// + /// See [`Self::with_prefetch_hint`] for a discussion of how to reduce the number of fetches + /// performed by this function. + #[cfg(all(feature = "async", feature = "arrow"))] + pub async fn try_load_via_suffix( + &mut self, + mut fetch: F, + ) -> Result<()> { + let (metadata, remainder) = self.load_metadata_via_suffix(&mut fetch).await?; + + self.metadata = Some(metadata); + + // we can return if page indexes aren't requested + if self.column_index == PageIndexPolicy::Skip && self.offset_index == PageIndexPolicy::Skip + { + return Ok(()); + } + + self.load_page_index_with_remainder(fetch, remainder).await + } + + /// Asynchronously fetch the page index structures when a [`ParquetMetaData`] has already + /// been obtained. See [`Self::new_with_metadata()`]. + #[cfg(all(feature = "async", feature = "arrow"))] + pub async fn load_page_index(&mut self, fetch: F) -> Result<()> { + self.load_page_index_with_remainder(fetch, None).await + } + + #[cfg(all(feature = "async", feature = "arrow"))] + async fn load_page_index_with_remainder( + &mut self, + mut fetch: F, + remainder: Option<(usize, Bytes)>, + ) -> Result<()> { + let Some(metadata) = self.metadata.take() else { + return Err(general_err!("Footer metadata is not present")); + }; + + // in this case we don't actually know what the file size is, so just use u64::MAX + // this is ok since the offsets in the metadata are always valid + let file_size = u64::MAX; + let push_decoder = ParquetMetaDataPushDecoder::try_new_with_metadata(file_size, metadata)? + .with_offset_index_policy(self.offset_index) + .with_column_index_policy(self.column_index) + .with_metadata_options(self.metadata_options.clone()); + let mut push_decoder = self.prepare_push_decoder(push_decoder); + + // Get bounds needed for page indexes (if any are present in the file). + let range = match needs_index_data(&mut push_decoder)? { + NeedsIndexData::No(metadata) => { + self.metadata = Some(metadata); + return Ok(()); + } + NeedsIndexData::Yes(range) => range, + }; + + let bytes = match &remainder { + Some((remainder_start, remainder)) if *remainder_start as u64 <= range.start => { + let remainder_start = *remainder_start as u64; + let offset = usize::try_from(range.start - remainder_start)?; + let end = usize::try_from(range.end - remainder_start)?; + if end > remainder.len() { + return Err(general_err!( + "Corrupted parquet file: index data range ({:?}) exceeds remainder length ({})", + range, + remainder.len() + )); + } + remainder.slice(offset..end) + } + // Note: this will potentially fetch data already in remainder, this keeps things simple + _ => fetch.fetch(range.start..range.end).await?, + }; + + // Sanity check + if bytes.len() as u64 != range.end - range.start { + return Err(general_err!( + "Corrupted parquet file: index data length mismatch, expected {}, got {}", + range.end - range.start, + bytes.len() + )); + } + push_decoder.push_range(range.clone(), bytes)?; + let metadata = parse_index_data(&mut push_decoder)?; + self.metadata = Some(metadata); + Ok(()) + } + + // One-shot parse of footer. + // Side effect: this will set `self.metadata_size` + fn parse_metadata(&mut self, chunk_reader: &R) -> Result { + // check file is large enough to hold footer + let file_size = chunk_reader.len(); + if file_size < (FOOTER_SIZE as u64) { + return Err(ParquetError::NeedMoreData(FOOTER_SIZE)); + } + + let mut footer = [0_u8; FOOTER_SIZE]; + chunk_reader + .get_read(file_size - FOOTER_SIZE as u64)? + .read_exact(&mut footer)?; + + let footer = FooterTail::try_new(&footer)?; + let metadata_len = footer.metadata_length(); + let footer_metadata_len = FOOTER_SIZE + metadata_len; + self.metadata_size = Some(footer_metadata_len); + + if footer_metadata_len as u64 > file_size { + return Err(ParquetError::NeedMoreData(footer_metadata_len)); + } + + let start = file_size - footer_metadata_len as u64; + let bytes = chunk_reader.get_bytes(start, metadata_len)?; + self.decode_footer_metadata(bytes, file_size, footer) + } + + /// Size of the serialized thrift metadata plus the 8 byte footer. Only set if + /// `self.parse_metadata` is called. + pub fn metadata_size(&self) -> Option { + self.metadata_size + } + + /// Return the number of bytes to read in the initial pass. If `prefetch_size` has + /// been provided, then return that value if it is larger than the size of the Parquet + /// file footer (8 bytes). Otherwise returns `8`. + #[cfg(all(feature = "async", feature = "arrow"))] + fn get_prefetch_size(&self) -> usize { + if let Some(prefetch) = self.prefetch_hint { + if prefetch > FOOTER_SIZE { + return prefetch; + } + } + FOOTER_SIZE + } + + #[cfg(all(feature = "async", feature = "arrow"))] + async fn load_metadata( + &self, + fetch: &mut F, + file_size: u64, + ) -> Result<(ParquetMetaData, Option<(usize, Bytes)>)> { + let prefetch = self.get_prefetch_size() as u64; + + if file_size < FOOTER_SIZE as u64 { + return Err(eof_err!("file size of {} is less than footer", file_size)); + } + + // If a size hint is provided, read more than the minimum size + // to try and avoid a second fetch. + // Note: prefetch > file_size is ok since we're using saturating_sub. + let footer_start = file_size.saturating_sub(prefetch); + + let suffix = fetch.fetch(footer_start..file_size).await?; + let suffix_len = suffix.len(); + let fetch_len = (file_size - footer_start) + .try_into() + .expect("footer size should never be larger than u32"); + if suffix_len < fetch_len { + return Err(eof_err!( + "metadata requires {} bytes, but could only read {}", + fetch_len, + suffix_len + )); + } + + let mut footer = [0; FOOTER_SIZE]; + footer.copy_from_slice(&suffix[suffix_len - FOOTER_SIZE..suffix_len]); + + let footer = FooterTail::try_new(&footer)?; + let length = footer.metadata_length(); + + if file_size < (length + FOOTER_SIZE) as u64 { + return Err(eof_err!( + "file size of {} is less than footer + metadata {}", + file_size, + length + FOOTER_SIZE + )); + } + + // Did not fetch the entire file metadata in the initial read, need to make a second request + if length > suffix_len - FOOTER_SIZE { + let metadata_start = file_size - (length + FOOTER_SIZE) as u64; + let meta = fetch + .fetch(metadata_start..(file_size - FOOTER_SIZE as u64)) + .await?; + Ok((self.decode_footer_metadata(meta, file_size, footer)?, None)) + } else { + let metadata_start = (file_size - (length + FOOTER_SIZE) as u64 - footer_start) + .try_into() + .expect("metadata length should never be larger than u32"); + let slice = suffix.slice(metadata_start..suffix_len - FOOTER_SIZE); + Ok(( + self.decode_footer_metadata(slice, file_size, footer)?, + Some((footer_start as usize, suffix.slice(..metadata_start))), + )) + } + } + + #[cfg(all(feature = "async", feature = "arrow"))] + async fn load_metadata_via_suffix( + &self, + fetch: &mut F, + ) -> Result<(ParquetMetaData, Option<(usize, Bytes)>)> { + let prefetch = self.get_prefetch_size(); + + let suffix = fetch.fetch_suffix(prefetch as _).await?; + let suffix_len = suffix.len(); + + if suffix_len < FOOTER_SIZE { + return Err(eof_err!( + "footer metadata requires {} bytes, but could only read {}", + FOOTER_SIZE, + suffix_len + )); + } + + let mut footer = [0; FOOTER_SIZE]; + footer.copy_from_slice(&suffix[suffix_len - FOOTER_SIZE..suffix_len]); + + let footer = FooterTail::try_new(&footer)?; + let length = footer.metadata_length(); + // fake file size as we are only parsing the footer metadata here + // (cant be parsing page indexes without the full file size) + let file_size = (length + FOOTER_SIZE) as u64; + + // Did not fetch the entire file metadata in the initial read, need to make a second request + let metadata_offset = length + FOOTER_SIZE; + if length > suffix_len - FOOTER_SIZE { + let meta = fetch.fetch_suffix(metadata_offset).await?; + + if meta.len() < metadata_offset { + return Err(eof_err!( + "metadata requires {} bytes, but could only read {}", + metadata_offset, + meta.len() + )); + } + + // need to slice off the footer or decryption fails + let meta = meta.slice(0..length); + Ok((self.decode_footer_metadata(meta, file_size, footer)?, None)) + } else { + let metadata_start = suffix_len - metadata_offset; + let slice = suffix.slice(metadata_start..suffix_len - FOOTER_SIZE); + Ok(( + self.decode_footer_metadata(slice, file_size, footer)?, + Some((0, suffix.slice(..metadata_start))), + )) + } + } + + /// Decodes [`ParquetMetaData`] from the provided bytes. + /// + /// Typically, this is used to decode the metadata from the end of a parquet + /// file. The format of `buf` is the Thrift compact binary protocol, as specified + /// by the [Parquet Spec]. + /// + /// It does **NOT** include the 8-byte footer. + /// + /// This method handles using either `decode_metadata` or + /// `decode_metadata_with_encryption` depending on whether the encryption + /// feature is enabled. + /// + /// [Parquet Spec]: https://github.com/apache/parquet-format#metadata + pub(crate) fn decode_footer_metadata( + &self, + buf: Bytes, + file_size: u64, + footer_tail: FooterTail, + ) -> Result { + // The push decoder expects the metadata to be at the end of the file + // (... data ...) + (metadata) + (footer) + // so we need to provide the starting offset of the metadata + // within the file. + let ending_offset = file_size.checked_sub(FOOTER_SIZE as u64).ok_or_else(|| { + general_err!( + "file size {file_size} is smaller than footer size {}", + FOOTER_SIZE + ) + })?; + + let starting_offset = ending_offset.checked_sub(buf.len() as u64).ok_or_else(|| { + general_err!( + "file size {file_size} is smaller than buffer size {} + footer size {}", + buf.len(), + FOOTER_SIZE + ) + })?; + + let range = starting_offset..ending_offset; + + let push_decoder = + ParquetMetaDataPushDecoder::try_new_with_footer_tail(file_size, footer_tail)? + // NOTE: DO NOT enable page indexes here, they are handled separately + .with_page_index_policy(PageIndexPolicy::Skip) + .with_metadata_options(self.metadata_options.clone()); + + let mut push_decoder = self.prepare_push_decoder(push_decoder); + push_decoder.push_range(range, buf)?; + match push_decoder.try_decode()? { + DecodeResult::Data(metadata) => Ok(metadata), + DecodeResult::Finished => Err(general_err!( + "could not parse parquet metadata -- previously finished" + )), + DecodeResult::NeedsData(ranges) => Err(general_err!( + "could not parse parquet metadata, needs ranges {:?}", + ranges + )), + } + } + + /// Prepares a push decoder and runs it to decode the metadata. + #[cfg(feature = "encryption")] + fn prepare_push_decoder( + &self, + push_decoder: ParquetMetaDataPushDecoder, + ) -> ParquetMetaDataPushDecoder { + push_decoder.with_file_decryption_properties( + self.file_decryption_properties + .as_ref() + .map(std::sync::Arc::clone), + ) + } + #[cfg(not(feature = "encryption"))] + fn prepare_push_decoder( + &self, + push_decoder: ParquetMetaDataPushDecoder, + ) -> ParquetMetaDataPushDecoder { + push_decoder + } + + /// Decodes [`ParquetMetaData`] from the provided bytes. + /// + /// Typically this is used to decode the metadata from the end of a parquet + /// file. The format of `buf` is the Thrift compact binary protocol, as specified + /// by the [Parquet Spec]. + /// + /// [Parquet Spec]: https://github.com/apache/parquet-format#metadata + pub fn decode_metadata(buf: &[u8]) -> Result { + decode_metadata(buf, None) + } + + /// Decodes [`ParquetMetaData`] from the provided bytes. + /// + /// Like [`Self::decode_metadata`] but this also accepts + /// metadata parsing options. + pub fn decode_metadata_with_options( + buf: &[u8], + options: Option<&ParquetMetaDataOptions>, + ) -> Result { + decode_metadata(buf, options) + } + + /// Decodes the schema from the Parquet footer in `buf`. Returned as + /// a [`SchemaDescriptor`]. + pub fn decode_schema(buf: &[u8]) -> Result> { + Ok(Arc::new(parquet_schema_from_bytes(buf)?)) + } +} + +/// The bounds needed to read page indexes +// this is an internal enum, so it is ok to allow differences in enum size +#[allow(clippy::large_enum_variant)] +enum NeedsIndexData { + /// no additional data is needed (e.g. the indexes weren't requested) + No(ParquetMetaData), + /// Additional data is needed, with the range that are required + Yes(Range), +} + +/// Determines a single combined range of bytes needed to read the page indexes, +/// or returns the metadata if no additional data is needed (e.g. if no page indexes are requested) +fn needs_index_data(push_decoder: &mut ParquetMetaDataPushDecoder) -> Result { + match push_decoder.try_decode()? { + DecodeResult::NeedsData(ranges) => { + let range = ranges + .into_iter() + .reduce(|a, b| a.start.min(b.start)..a.end.max(b.end)) + .ok_or_else(|| general_err!("Internal error: no ranges provided"))?; + Ok(NeedsIndexData::Yes(range)) + } + DecodeResult::Data(metadata) => Ok(NeedsIndexData::No(metadata)), + DecodeResult::Finished => Err(general_err!("Internal error: decoder was finished")), + } +} + +/// Given a push decoder that has had the needed ranges pushed to it, +/// attempt to decode indexes and return the updated metadata. +fn parse_index_data(push_decoder: &mut ParquetMetaDataPushDecoder) -> Result { + match push_decoder.try_decode()? { + DecodeResult::NeedsData(_) => Err(general_err!( + "Internal error: decoder still needs data after reading required range" + )), + DecodeResult::Data(metadata) => Ok(metadata), + DecodeResult::Finished => Err(general_err!("Internal error: decoder was finished")), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::file::reader::Length; + use crate::util::test_common::file_util::get_test_file; + use std::ops::Range; + + #[test] + fn test_parse_metadata_size_smaller_than_footer() { + let test_file = tempfile::tempfile().unwrap(); + let err = ParquetMetaDataReader::new() + .parse_metadata(&test_file) + .unwrap_err(); + assert!(matches!(err, ParquetError::NeedMoreData(FOOTER_SIZE))); + } + + #[test] + fn test_parse_metadata_corrupt_footer() { + let data = Bytes::from(vec![1, 2, 3, 4, 5, 6, 7, 8]); + let reader_result = ParquetMetaDataReader::new().parse_metadata(&data); + assert_eq!( + reader_result.unwrap_err().to_string(), + "Parquet error: Invalid Parquet file. Corrupt footer" + ); + } + + #[test] + fn test_parse_metadata_invalid_start() { + let test_file = Bytes::from(vec![255, 0, 0, 0, b'P', b'A', b'R', b'1']); + let err = ParquetMetaDataReader::new() + .parse_metadata(&test_file) + .unwrap_err(); + assert!(matches!(err, ParquetError::NeedMoreData(263))); + } + + #[test] + fn test_try_parse() { + let file = get_test_file("alltypes_tiny_pages.parquet"); + let len = file.len(); + + let mut reader = + ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required); + + let bytes_for_range = |range: Range| { + file.get_bytes(range.start, (range.end - range.start).try_into().unwrap()) + .unwrap() + }; + + // read entire file + let bytes = bytes_for_range(0..len); + reader.try_parse(&bytes).unwrap(); + let metadata = reader.finish().unwrap(); + assert!(metadata.column_index.is_some()); + assert!(metadata.offset_index.is_some()); + + // read more than enough of file + let bytes = bytes_for_range(320000..len); + reader.try_parse_sized(&bytes, len).unwrap(); + let metadata = reader.finish().unwrap(); + assert!(metadata.column_index.is_some()); + assert!(metadata.offset_index.is_some()); + + // exactly enough + let bytes = bytes_for_range(323583..len); + reader.try_parse_sized(&bytes, len).unwrap(); + let metadata = reader.finish().unwrap(); + assert!(metadata.column_index.is_some()); + assert!(metadata.offset_index.is_some()); + + // not enough for page index + let bytes = bytes_for_range(323584..len); + // should fail + match reader.try_parse_sized(&bytes, len).unwrap_err() { + // expected error, try again with provided bounds + ParquetError::NeedMoreData(needed) => { + let bytes = bytes_for_range(len - needed as u64..len); + reader.try_parse_sized(&bytes, len).unwrap(); + let metadata = reader.finish().unwrap(); + assert!(metadata.column_index.is_some()); + assert!(metadata.offset_index.is_some()); + } + _ => panic!("unexpected error"), + }; + + // not enough for file metadata, but keep trying until page indexes are read + let mut reader = + ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required); + let mut bytes = bytes_for_range(452505..len); + loop { + match reader.try_parse_sized(&bytes, len) { + Ok(_) => break, + Err(ParquetError::NeedMoreData(needed)) => { + bytes = bytes_for_range(len - needed as u64..len); + if reader.has_metadata() { + reader.read_page_indexes_sized(&bytes, len).unwrap(); + break; + } + } + _ => panic!("unexpected error"), + } + } + let metadata = reader.finish().unwrap(); + assert!(metadata.column_index.is_some()); + assert!(metadata.offset_index.is_some()); + + // not enough for page index but lie about file size + let bytes = bytes_for_range(323584..len); + let reader_result = reader.try_parse_sized(&bytes, len - 323584).unwrap_err(); + assert_eq!( + reader_result.to_string(), + "EOF: Parquet file too small. Range 323583..452504 is beyond file bounds 130649" + ); + + // not enough for file metadata + let mut reader = ParquetMetaDataReader::new(); + let bytes = bytes_for_range(452505..len); + // should fail + match reader.try_parse_sized(&bytes, len).unwrap_err() { + // expected error, try again with provided bounds + ParquetError::NeedMoreData(needed) => { + let bytes = bytes_for_range(len - needed as u64..len); + reader.try_parse_sized(&bytes, len).unwrap(); + reader.finish().unwrap(); + } + _ => panic!("unexpected error"), + }; + + // not enough for file metadata but use try_parse() + let reader_result = reader.try_parse(&bytes).unwrap_err(); + assert_eq!( + reader_result.to_string(), + "EOF: Parquet file too small. Size is 1728 but need 1729" + ); + + // read head of file rather than tail + let bytes = bytes_for_range(0..1000); + let reader_result = reader.try_parse_sized(&bytes, len).unwrap_err(); + assert_eq!( + reader_result.to_string(), + "Parquet error: Invalid Parquet file. Corrupt footer" + ); + + // lie about file size + let bytes = bytes_for_range(452510..len); + let reader_result = reader.try_parse_sized(&bytes, len - 452505).unwrap_err(); + assert_eq!( + reader_result.to_string(), + "EOF: Parquet file too small. Size is 1728 but need 1729" + ); + } +} + +#[cfg(all(feature = "async", feature = "arrow", test))] +mod async_tests { + use super::*; + + use arrow::{array::Int32Array, datatypes::DataType}; + use arrow_array::RecordBatch; + use arrow_schema::{Field, Schema}; + use bytes::Bytes; + use futures::FutureExt; + use futures::future::BoxFuture; + use std::fs::File; + use std::future::Future; + use std::io::{Read, Seek, SeekFrom}; + use std::ops::Range; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use tempfile::NamedTempFile; + + use crate::arrow::ArrowWriter; + use crate::file::properties::WriterProperties; + use crate::file::reader::Length; + use crate::util::test_common::file_util::get_test_file; + + struct MetadataFetchFn(F); + + impl MetadataFetch for MetadataFetchFn + where + F: FnMut(Range) -> Fut + Send, + Fut: Future> + Send, + { + fn fetch(&mut self, range: Range) -> BoxFuture<'_, Result> { + async move { self.0(range).await }.boxed() + } + } + + struct MetadataSuffixFetchFn(F1, F2); + + impl MetadataFetch for MetadataSuffixFetchFn + where + F1: FnMut(Range) -> Fut + Send, + Fut: Future> + Send, + F2: Send, + { + fn fetch(&mut self, range: Range) -> BoxFuture<'_, Result> { + async move { self.0(range).await }.boxed() + } + } + + impl MetadataSuffixFetch for MetadataSuffixFetchFn + where + F1: FnMut(Range) -> Fut + Send, + F2: FnMut(usize) -> Fut + Send, + Fut: Future> + Send, + { + fn fetch_suffix(&mut self, suffix: usize) -> BoxFuture<'_, Result> { + async move { self.1(suffix).await }.boxed() + } + } + + fn read_range(file: &mut File, range: Range) -> Result { + file.seek(SeekFrom::Start(range.start as _))?; + let len = range.end - range.start; + let mut buf = Vec::with_capacity(len.try_into().unwrap()); + file.take(len as _).read_to_end(&mut buf)?; + Ok(buf.into()) + } + + fn read_suffix(file: &mut File, suffix: usize) -> Result { + let file_len = file.len(); + // Don't seek before beginning of file + file.seek(SeekFrom::End(0 - suffix.min(file_len as _) as i64))?; + let mut buf = Vec::with_capacity(suffix); + file.take(suffix as _).read_to_end(&mut buf)?; + Ok(buf.into()) + } + + #[tokio::test] + async fn test_simple() { + let mut file = get_test_file("nulls.snappy.parquet"); + let len = file.len(); + + let expected = ParquetMetaDataReader::new() + .parse_and_finish(&file) + .unwrap(); + let expected = expected.file_metadata().schema(); + let fetch_count = AtomicUsize::new(0); + + let mut fetch = |range| { + fetch_count.fetch_add(1, Ordering::SeqCst); + futures::future::ready(read_range(&mut file, range)) + }; + + let input = MetadataFetchFn(&mut fetch); + let actual = ParquetMetaDataReader::new() + .load_and_finish(input, len) + .await + .unwrap(); + assert_eq!(actual.file_metadata().schema(), expected); + assert_eq!(fetch_count.load(Ordering::SeqCst), 2); + + // Metadata hint too small - below footer size + fetch_count.store(0, Ordering::SeqCst); + let input = MetadataFetchFn(&mut fetch); + let actual = ParquetMetaDataReader::new() + .with_prefetch_hint(Some(7)) + .load_and_finish(input, len) + .await + .unwrap(); + assert_eq!(actual.file_metadata().schema(), expected); + assert_eq!(fetch_count.load(Ordering::SeqCst), 2); + + // Metadata hint too small + fetch_count.store(0, Ordering::SeqCst); + let input = MetadataFetchFn(&mut fetch); + let actual = ParquetMetaDataReader::new() + .with_prefetch_hint(Some(10)) + .load_and_finish(input, len) + .await + .unwrap(); + assert_eq!(actual.file_metadata().schema(), expected); + assert_eq!(fetch_count.load(Ordering::SeqCst), 2); + + // Metadata hint too large + fetch_count.store(0, Ordering::SeqCst); + let input = MetadataFetchFn(&mut fetch); + let actual = ParquetMetaDataReader::new() + .with_prefetch_hint(Some(500)) + .load_and_finish(input, len) + .await + .unwrap(); + assert_eq!(actual.file_metadata().schema(), expected); + assert_eq!(fetch_count.load(Ordering::SeqCst), 1); + + // Metadata hint exactly correct + fetch_count.store(0, Ordering::SeqCst); + let input = MetadataFetchFn(&mut fetch); + let actual = ParquetMetaDataReader::new() + .with_prefetch_hint(Some(428)) + .load_and_finish(input, len) + .await + .unwrap(); + assert_eq!(actual.file_metadata().schema(), expected); + assert_eq!(fetch_count.load(Ordering::SeqCst), 1); + + let input = MetadataFetchFn(&mut fetch); + let err = ParquetMetaDataReader::new() + .load_and_finish(input, 4) + .await + .unwrap_err() + .to_string(); + assert_eq!(err, "EOF: file size of 4 is less than footer"); + + let input = MetadataFetchFn(&mut fetch); + let err = ParquetMetaDataReader::new() + .load_and_finish(input, 20) + .await + .unwrap_err() + .to_string(); + assert_eq!(err, "Parquet error: Invalid Parquet file. Corrupt footer"); + } + + #[tokio::test] + async fn test_suffix() { + let mut file = get_test_file("nulls.snappy.parquet"); + let mut file2 = file.try_clone().unwrap(); + + let expected = ParquetMetaDataReader::new() + .parse_and_finish(&file) + .unwrap(); + let expected = expected.file_metadata().schema(); + let fetch_count = AtomicUsize::new(0); + let suffix_fetch_count = AtomicUsize::new(0); + + let mut fetch = |range| { + fetch_count.fetch_add(1, Ordering::SeqCst); + futures::future::ready(read_range(&mut file, range)) + }; + let mut suffix_fetch = |suffix| { + suffix_fetch_count.fetch_add(1, Ordering::SeqCst); + futures::future::ready(read_suffix(&mut file2, suffix)) + }; + + let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch); + let actual = ParquetMetaDataReader::new() + .load_via_suffix_and_finish(input) + .await + .unwrap(); + assert_eq!(actual.file_metadata().schema(), expected); + assert_eq!(fetch_count.load(Ordering::SeqCst), 0); + assert_eq!(suffix_fetch_count.load(Ordering::SeqCst), 2); + + // Metadata hint too small - below footer size + fetch_count.store(0, Ordering::SeqCst); + suffix_fetch_count.store(0, Ordering::SeqCst); + let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch); + let actual = ParquetMetaDataReader::new() + .with_prefetch_hint(Some(7)) + .load_via_suffix_and_finish(input) + .await + .unwrap(); + assert_eq!(actual.file_metadata().schema(), expected); + assert_eq!(fetch_count.load(Ordering::SeqCst), 0); + assert_eq!(suffix_fetch_count.load(Ordering::SeqCst), 2); + + // Metadata hint too small + fetch_count.store(0, Ordering::SeqCst); + suffix_fetch_count.store(0, Ordering::SeqCst); + let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch); + let actual = ParquetMetaDataReader::new() + .with_prefetch_hint(Some(10)) + .load_via_suffix_and_finish(input) + .await + .unwrap(); + assert_eq!(actual.file_metadata().schema(), expected); + assert_eq!(fetch_count.load(Ordering::SeqCst), 0); + assert_eq!(suffix_fetch_count.load(Ordering::SeqCst), 2); + + dbg!("test"); + // Metadata hint too large + fetch_count.store(0, Ordering::SeqCst); + suffix_fetch_count.store(0, Ordering::SeqCst); + let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch); + let actual = ParquetMetaDataReader::new() + .with_prefetch_hint(Some(500)) + .load_via_suffix_and_finish(input) + .await + .unwrap(); + assert_eq!(actual.file_metadata().schema(), expected); + assert_eq!(fetch_count.load(Ordering::SeqCst), 0); + assert_eq!(suffix_fetch_count.load(Ordering::SeqCst), 1); + + // Metadata hint exactly correct + fetch_count.store(0, Ordering::SeqCst); + suffix_fetch_count.store(0, Ordering::SeqCst); + let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch); + let actual = ParquetMetaDataReader::new() + .with_prefetch_hint(Some(428)) + .load_via_suffix_and_finish(input) + .await + .unwrap(); + assert_eq!(actual.file_metadata().schema(), expected); + assert_eq!(fetch_count.load(Ordering::SeqCst), 0); + assert_eq!(suffix_fetch_count.load(Ordering::SeqCst), 1); + } + + #[cfg(feature = "encryption")] + #[tokio::test] + async fn test_suffix_with_encryption() { + let mut file = get_test_file("uniform_encryption.parquet.encrypted"); + let mut file2 = file.try_clone().unwrap(); + + let mut fetch = |range| futures::future::ready(read_range(&mut file, range)); + let mut suffix_fetch = |suffix| futures::future::ready(read_suffix(&mut file2, suffix)); + + let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch); + + let key_code: &[u8] = "0123456789012345".as_bytes(); + let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec()) + .build() + .unwrap(); + + // just make sure the metadata is properly decrypted and read + let expected = ParquetMetaDataReader::new() + .with_decryption_properties(Some(decryption_properties)) + .load_via_suffix_and_finish(input) + .await + .unwrap(); + assert_eq!(expected.num_row_groups(), 1); + } + + #[tokio::test] + async fn test_page_index() { + let mut file = get_test_file("alltypes_tiny_pages.parquet"); + let len = file.len(); + let fetch_count = AtomicUsize::new(0); + let mut fetch = |range| { + fetch_count.fetch_add(1, Ordering::SeqCst); + futures::future::ready(read_range(&mut file, range)) + }; + + let f = MetadataFetchFn(&mut fetch); + let mut loader = + ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required); + loader.try_load(f, len).await.unwrap(); + assert_eq!(fetch_count.load(Ordering::SeqCst), 3); + let metadata = loader.finish().unwrap(); + assert!(metadata.offset_index().is_some() && metadata.column_index().is_some()); + + // Prefetch just footer exactly + fetch_count.store(0, Ordering::SeqCst); + let f = MetadataFetchFn(&mut fetch); + let mut loader = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .with_prefetch_hint(Some(1729)); + loader.try_load(f, len).await.unwrap(); + assert_eq!(fetch_count.load(Ordering::SeqCst), 2); + let metadata = loader.finish().unwrap(); + assert!(metadata.offset_index().is_some() && metadata.column_index().is_some()); + + // Prefetch more than footer but not enough + fetch_count.store(0, Ordering::SeqCst); + let f = MetadataFetchFn(&mut fetch); + let mut loader = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .with_prefetch_hint(Some(130649)); + loader.try_load(f, len).await.unwrap(); + assert_eq!(fetch_count.load(Ordering::SeqCst), 2); + let metadata = loader.finish().unwrap(); + assert!(metadata.offset_index().is_some() && metadata.column_index().is_some()); + + // Prefetch exactly enough + fetch_count.store(0, Ordering::SeqCst); + let f = MetadataFetchFn(&mut fetch); + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .with_prefetch_hint(Some(130650)) + .load_and_finish(f, len) + .await + .unwrap(); + assert_eq!(fetch_count.load(Ordering::SeqCst), 1); + assert!(metadata.offset_index().is_some() && metadata.column_index().is_some()); + + // Prefetch more than enough but less than the entire file + fetch_count.store(0, Ordering::SeqCst); + let f = MetadataFetchFn(&mut fetch); + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .with_prefetch_hint(Some((len - 1000) as usize)) // prefetch entire file + .load_and_finish(f, len) + .await + .unwrap(); + assert_eq!(fetch_count.load(Ordering::SeqCst), 1); + assert!(metadata.offset_index().is_some() && metadata.column_index().is_some()); + + // Prefetch the entire file + fetch_count.store(0, Ordering::SeqCst); + let f = MetadataFetchFn(&mut fetch); + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .with_prefetch_hint(Some(len as usize)) // prefetch entire file + .load_and_finish(f, len) + .await + .unwrap(); + assert_eq!(fetch_count.load(Ordering::SeqCst), 1); + assert!(metadata.offset_index().is_some() && metadata.column_index().is_some()); + + // Prefetch more than the entire file + fetch_count.store(0, Ordering::SeqCst); + let f = MetadataFetchFn(&mut fetch); + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .with_prefetch_hint(Some((len + 1000) as usize)) // prefetch entire file + .load_and_finish(f, len) + .await + .unwrap(); + assert_eq!(fetch_count.load(Ordering::SeqCst), 1); + assert!(metadata.offset_index().is_some() && metadata.column_index().is_some()); + } + + fn write_parquet_file(offset_index_disabled: bool) -> Result { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + )?; + + let file = NamedTempFile::new().unwrap(); + + // Write properties with page index disabled + let props = WriterProperties::builder() + .set_offset_index_disabled(offset_index_disabled) + .build(); + + let mut writer = ArrowWriter::try_new(file.reopen()?, schema, Some(props))?; + writer.write(&batch)?; + writer.close()?; + + Ok(file) + } + + fn read_and_check(file: &File, policy: PageIndexPolicy) -> Result { + let mut reader = ParquetMetaDataReader::new().with_page_index_policy(policy); + reader.try_parse(file)?; + reader.finish() + } + + #[test] + fn test_page_index_policy() { + // With page index + let f = write_parquet_file(false).unwrap(); + read_and_check(f.as_file(), PageIndexPolicy::Required).unwrap(); + read_and_check(f.as_file(), PageIndexPolicy::Optional).unwrap(); + read_and_check(f.as_file(), PageIndexPolicy::Skip).unwrap(); + + // Without page index + let f = write_parquet_file(true).unwrap(); + let res = read_and_check(f.as_file(), PageIndexPolicy::Required); + assert!(matches!( + res, + Err(ParquetError::General(e)) if e == "missing offset index" + )); + read_and_check(f.as_file(), PageIndexPolicy::Optional).unwrap(); + read_and_check(f.as_file(), PageIndexPolicy::Skip).unwrap(); + } +} diff --git a/vendor/parquet/src/file/metadata/thrift/encryption.rs b/vendor/parquet/src/file/metadata/thrift/encryption.rs new file mode 100644 index 0000000..37e91ba --- /dev/null +++ b/vendor/parquet/src/file/metadata/thrift/encryption.rs @@ -0,0 +1,348 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Encryption support for Thrift serialization + +use crate::{ + encryption::decrypt::{FileDecryptionProperties, FileDecryptor}, + errors::{ParquetError, Result}, + file::{ + column_crypto_metadata::ColumnCryptoMetaData, + metadata::{ + HeapSize, ParquetMetaData, ParquetMetaDataOptions, RowGroupMetaData, + thrift::{parquet_metadata_from_bytes, read_column_metadata, validate_column_metadata}, + }, + }, + parquet_thrift::{ + ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol, + ThriftCompactOutputProtocol, ThriftSliceInputProtocol, WriteThrift, WriteThriftField, + }, + thrift_struct, thrift_union, +}; +use std::io::Write; +use std::sync::Arc; + +thrift_struct!( +pub(crate) struct AesGcmV1 { + /// AAD prefix + 1: optional binary aad_prefix + + /// Unique file identifier part of AAD suffix + 2: optional binary aad_file_unique + + /// In files encrypted with AAD prefix without storing it, + /// readers must supply the prefix + 3: optional bool supply_aad_prefix +} +); + +impl HeapSize for AesGcmV1 { + fn heap_size(&self) -> usize { + self.aad_prefix.heap_size() + + self.aad_file_unique.heap_size() + + self.supply_aad_prefix.heap_size() + } +} + +thrift_struct!( +pub(crate) struct AesGcmCtrV1 { + /// AAD prefix + 1: optional binary aad_prefix + + /// Unique file identifier part of AAD suffix + 2: optional binary aad_file_unique + + /// In files encrypted with AAD prefix without storing it, + /// readers must supply the prefix + 3: optional bool supply_aad_prefix +} +); + +impl HeapSize for AesGcmCtrV1 { + fn heap_size(&self) -> usize { + self.aad_prefix.heap_size() + + self.aad_file_unique.heap_size() + + self.supply_aad_prefix.heap_size() + } +} + +thrift_union!( +union EncryptionAlgorithm { + 1: (AesGcmV1) AES_GCM_V1 + 2: (AesGcmCtrV1) AES_GCM_CTR_V1 +} +); + +impl HeapSize for EncryptionAlgorithm { + fn heap_size(&self) -> usize { + match self { + Self::AES_GCM_V1(gcm) => gcm.heap_size(), + Self::AES_GCM_CTR_V1(gcm_ctr) => gcm_ctr.heap_size(), + } + } +} + +thrift_struct!( +/// Crypto metadata for files with encrypted footer +pub(crate) struct FileCryptoMetaData<'a> { + /// Encryption algorithm. This field is only used for files + /// with encrypted footer. Files with plaintext footer store algorithm id + /// inside footer (FileMetaData structure). + 1: required EncryptionAlgorithm encryption_algorithm + + /// Retrieval metadata of key used for encryption of footer, + /// and (possibly) columns. + 2: optional binary<'a> key_metadata +} +); + +fn row_group_from_encrypted_thrift( + mut rg: RowGroupMetaData, + decryptor: Option<&FileDecryptor>, + options: Option<&ParquetMetaDataOptions>, +) -> Result { + let schema_descr = rg.schema_descr; + + if schema_descr.num_columns() != rg.columns.len() { + return Err(general_err!( + "Column count mismatch. Schema has {} columns while Row Group has {}", + schema_descr.num_columns(), + rg.columns.len() + )); + } + let total_byte_size = rg.total_byte_size; + let num_rows = rg.num_rows; + let mut columns = vec![]; + + for (i, (mut c, d)) in rg + .columns + .drain(0..) + .zip(schema_descr.columns()) + .enumerate() + { + // Read encrypted metadata if it's present and we have a decryptor. + if let (true, Some(decryptor)) = (c.encrypted_column_metadata.is_some(), decryptor) { + let column_decryptor = match c.crypto_metadata() { + None => { + return Err(general_err!( + "No crypto_metadata is set for column '{}', which has encrypted metadata", + d.path().string() + )); + } + Some(ColumnCryptoMetaData::ENCRYPTION_WITH_COLUMN_KEY(crypto_metadata)) => { + let column_name = crypto_metadata.path_in_schema.join("."); + // Try to get the decryptor - if it fails, we don't have the key + match decryptor.get_column_metadata_decryptor( + column_name.as_str(), + crypto_metadata.key_metadata.as_deref(), + ) { + Ok(dec) => dec, + Err(_) => { + // We don't have the key for this column, so we can't decrypt its metadata. + columns.push(c); + continue; + } + } + } + Some(ColumnCryptoMetaData::ENCRYPTION_WITH_FOOTER_KEY) => { + decryptor.get_footer_decryptor()? + } + }; + + let column_aad = crate::encryption::modules::create_module_aad( + decryptor.file_aad(), + crate::encryption::modules::ModuleType::ColumnMetaData, + rg.ordinal.unwrap() as usize, + i, + None, + )?; + + // Take the encrypted column metadata as it is no longer needed. + let encrypted_column_metadata = c.encrypted_column_metadata.take(); + let buf = encrypted_column_metadata.unwrap(); + let decrypted_cc_buf = column_decryptor + .decrypt(&buf, column_aad.as_ref()) + .map_err(|_| { + general_err!( + "Unable to decrypt column '{}', perhaps the column key is wrong?", + d.path().string() + ) + })?; + + // parse decrypted buffer and then replace fields in 'c' + let mut prot = ThriftSliceInputProtocol::new(&decrypted_cc_buf); + let mask = read_column_metadata(&mut prot, &mut c, i, options)?; + validate_column_metadata(mask)?; + + columns.push(c); + } else { + columns.push(c); + } + } + + let sorting_columns = rg.sorting_columns; + let file_offset = rg.file_offset; + let ordinal = rg.ordinal; + + Ok(RowGroupMetaData { + columns, + num_rows, + sorting_columns, + total_byte_size, + schema_descr, + file_offset, + ordinal, + }) +} + +/// Decodes [`ParquetMetaData`] from the provided bytes, handling metadata that may be encrypted. +/// +/// Typically this is used to decode the metadata from the end of a parquet +/// file. The format of `buf` is the Thrift compact binary protocol, as specified +/// by the [Parquet Spec]. Buffer can be encrypted with AES GCM or AES CTR +/// ciphers as specfied in the [Parquet Encryption Spec]. +/// +/// [Parquet Spec]: https://github.com/apache/parquet-format#metadata +/// [Parquet Encryption Spec]: https://parquet.apache.org/docs/file-format/data-pages/encryption/ +pub(crate) fn parquet_metadata_with_encryption( + file_decryption_properties: Option<&Arc>, + encrypted_footer: bool, + buf: &[u8], + options: Option<&ParquetMetaDataOptions>, +) -> Result { + use crate::file::metadata::ParquetMetaDataBuilder; + + let mut buf = buf; + let mut file_decryptor = None; + let decrypted_fmd_buf; + + if encrypted_footer { + let mut prot = ThriftSliceInputProtocol::new(buf); + if let Some(file_decryption_properties) = file_decryption_properties { + let t_file_crypto_metadata: FileCryptoMetaData = + FileCryptoMetaData::read_thrift(&mut prot) + .map_err(|e| general_err!("Could not parse crypto metadata: {}", e))?; + let supply_aad_prefix = match &t_file_crypto_metadata.encryption_algorithm { + EncryptionAlgorithm::AES_GCM_V1(algo) => algo.supply_aad_prefix, + _ => Some(false), + } + .unwrap_or(false); + if supply_aad_prefix && file_decryption_properties.aad_prefix().is_none() { + return Err(general_err!( + "Parquet file was encrypted with an AAD prefix that is not stored in the file, \ + but no AAD prefix was provided in the file decryption properties" + )); + } + let decryptor = get_file_decryptor( + t_file_crypto_metadata.encryption_algorithm, + t_file_crypto_metadata.key_metadata, + file_decryption_properties, + )?; + let footer_decryptor = decryptor.get_footer_decryptor(); + let aad_footer = crate::encryption::modules::create_footer_aad(decryptor.file_aad())?; + + decrypted_fmd_buf = footer_decryptor? + .decrypt(prot.as_slice().as_ref(), aad_footer.as_ref()) + .map_err(|_| { + general_err!( + "Provided footer key and AAD were unable to decrypt parquet footer" + ) + })?; + + buf = &decrypted_fmd_buf; + file_decryptor = Some(decryptor); + } else { + return Err(general_err!( + "Parquet file has an encrypted footer but decryption properties were not provided" + )); + } + } + + let parquet_meta = parquet_metadata_from_bytes(buf, options) + .map_err(|e| general_err!("Could not parse metadata: {}", e))?; + + let ParquetMetaData { + mut file_metadata, + row_groups, + column_index: _, + offset_index: _, + file_decryptor: _, + } = parquet_meta; + + // Take the encryption algorithm and footer signing key metadata as they are no longer + // needed after this. + if let (Some(algo), Some(file_decryption_properties)) = ( + file_metadata.encryption_algorithm.take(), + file_decryption_properties, + ) { + let footer_signing_key_metadata = file_metadata.footer_signing_key_metadata.take(); + + // File has a plaintext footer but encryption algorithm is set + let file_decryptor_value = get_file_decryptor( + *algo, + footer_signing_key_metadata.as_deref(), + file_decryption_properties, + )?; + if file_decryption_properties.check_plaintext_footer_integrity() && !encrypted_footer { + file_decryptor_value.verify_plaintext_footer_signature(buf)?; + } + file_decryptor = Some(file_decryptor_value); + } + + // decrypt column chunk info + let row_groups = row_groups + .into_iter() + .map(|rg| row_group_from_encrypted_thrift(rg, file_decryptor.as_ref(), options)) + .collect::>>()?; + + let metadata = ParquetMetaDataBuilder::new(file_metadata) + .set_row_groups(row_groups) + .set_file_decryptor(file_decryptor) + .build(); + + Ok(metadata) +} + +fn get_file_decryptor( + encryption_algorithm: EncryptionAlgorithm, + footer_key_metadata: Option<&[u8]>, + file_decryption_properties: &Arc, +) -> Result { + match encryption_algorithm { + EncryptionAlgorithm::AES_GCM_V1(algo) => { + let aad_file_unique = algo + .aad_file_unique + .ok_or_else(|| general_err!("AAD unique file identifier is not set"))?; + let aad_prefix = if let Some(aad_prefix) = file_decryption_properties.aad_prefix() { + aad_prefix.clone() + } else { + algo.aad_prefix.map(|v| v.to_vec()).unwrap_or_default() + }; + let aad_file_unique = aad_file_unique.to_vec(); + + FileDecryptor::new( + file_decryption_properties, + footer_key_metadata, + aad_file_unique, + aad_prefix, + ) + } + EncryptionAlgorithm::AES_GCM_CTR_V1(_) => Err(nyi_err!( + "The AES_GCM_CTR_V1 encryption algorithm is not yet supported" + )), + } +} diff --git a/vendor/parquet/src/file/metadata/thrift/mod.rs b/vendor/parquet/src/file/metadata/thrift/mod.rs new file mode 100644 index 0000000..d5a0112 --- /dev/null +++ b/vendor/parquet/src/file/metadata/thrift/mod.rs @@ -0,0 +1,2010 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! This module is the bridge between a Parquet file's thrift encoded metadata +//! and this crate's [Parquet metadata API]. It contains objects and functions used +//! to serialize/deserialize metadata objects into/from the Thrift compact protocol +//! format as defined by the [Parquet specification]. +//! +//! [Parquet metadata API]: crate::file::metadata +//! [Parquet specification]: https://github.com/apache/parquet-format/tree/master + +use std::io::Write; +use std::sync::Arc; + +#[cfg(feature = "encryption")] +pub(crate) mod encryption; + +#[cfg(feature = "encryption")] +use crate::file::{ + column_crypto_metadata::ColumnCryptoMetaData, metadata::thrift::encryption::EncryptionAlgorithm, +}; +use crate::{ + basic::{ + ColumnOrder, CompressionCodec, ConvertedType, Encoding, EncodingMask, LogicalType, + PageType, Repetition, Type, + }, + data_type::{ByteArray, FixedLenByteArray, Int96}, + errors::{ParquetError, Result}, + file::{ + metadata::{ + ColumnChunkMetaData, ColumnChunkMetaDataBuilder, KeyValue, LevelHistogram, + PageEncodingStats, ParquetMetaData, ParquetMetaDataOptions, ParquetPageEncodingStats, + RowGroupMetaData, RowGroupMetaDataBuilder, SortingColumn, + }, + statistics::ValueStatistics, + }, + parquet_thrift::{ + ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol, + ThriftCompactOutputProtocol, ThriftSliceInputProtocol, WriteThrift, WriteThriftField, + read_thrift_vec, validate_list_type, + }, + schema::types::{ + ColumnDescriptor, SchemaDescriptor, TypePtr, num_nodes, parquet_schema_from_array, + }, + thrift_struct, + util::bit_util::FromBytes, + write_thrift_field, +}; + +// this needs to be visible to the schema conversion code +thrift_struct!( +pub(crate) struct SchemaElement<'a> { + /// Data type for this field. Not set if the current element is a non-leaf node + 1: optional Type r#type; + /// If type is FIXED_LEN_BYTE_ARRAY, this is the byte length of the values. + /// Otherwise, if specified, this is the maximum bit length to store any of the values. + /// (e.g. a low cardinality INT col could have this set to 3). Note that this is + /// in the schema, and therefore fixed for the entire file. + 2: optional i32 type_length; + /// Repetition of the field. The root of the schema does not have a repetition_type. + /// All other nodes must have one. + 3: optional Repetition repetition_type; + /// Name of the field in the schema + 4: required string<'a> name; + /// Nested fields. Since thrift does not support nested fields, + /// the nesting is flattened to a single list by a depth-first traversal. + /// The children count is used to construct the nested relationship. + /// This field is not set when the element is a primitive type. + 5: optional i32 num_children; + /// DEPRECATED: When the schema is the result of a conversion from another model. + /// Used to record the original type to help with cross conversion. + /// + /// This is superseded by logical_type. + 6: optional ConvertedType converted_type; + /// DEPRECATED: Used when this column contains decimal data. + /// See the DECIMAL converted type for more details. + /// + /// This is superseded by using the DecimalType annotation in logical_type. + 7: optional i32 scale + 8: optional i32 precision + /// When the original schema supports field ids, this will save the + /// original field id in the parquet schema + 9: optional i32 field_id; + /// The logical type of this SchemaElement + /// + /// LogicalType replaces ConvertedType, but ConvertedType is still required + /// for some logical types to ensure forward-compatibility in format v1. + 10: optional LogicalType logical_type +} +); + +thrift_struct!( +struct Statistics<'a> { + 1: optional binary<'a> max; + 2: optional binary<'a> min; + 3: optional i64 null_count; + 4: optional i64 distinct_count; + 5: optional binary<'a> max_value; + 6: optional binary<'a> min_value; + 7: optional bool is_max_value_exact; + 8: optional bool is_min_value_exact; +} +); + +thrift_struct!( +struct BoundingBox { + 1: required double xmin; + 2: required double xmax; + 3: required double ymin; + 4: required double ymax; + 5: optional double zmin; + 6: optional double zmax; + 7: optional double mmin; + 8: optional double mmax; +} +); + +thrift_struct!( +struct GeospatialStatistics { + 1: optional BoundingBox bbox; + 2: optional list geospatial_types; +} +); + +thrift_struct!( +struct SizeStatistics { + 1: optional i64 unencoded_byte_array_data_bytes; + 2: optional list repetition_level_histogram; + 3: optional list definition_level_histogram; +} +); + +fn convert_geo_stats( + stats: Option, +) -> Option> { + stats.map(|st| { + let bbox = convert_bounding_box(st.bbox); + let geospatial_types: Option> = st.geospatial_types.filter(|v| !v.is_empty()); + Box::new(crate::geospatial::statistics::GeospatialStatistics::new( + bbox, + geospatial_types, + )) + }) +} + +fn convert_bounding_box( + bbox: Option, +) -> Option { + bbox.map(|bb| { + let mut newbb = crate::geospatial::bounding_box::BoundingBox::new( + bb.xmin.into(), + bb.xmax.into(), + bb.ymin.into(), + bb.ymax.into(), + ); + + newbb = match (bb.zmin, bb.zmax) { + (Some(zmin), Some(zmax)) => newbb.with_zrange(zmin.into(), zmax.into()), + // If either None or mismatch, leave it as None and don't error + _ => newbb, + }; + + newbb = match (bb.mmin, bb.mmax) { + (Some(mmin), Some(mmax)) => newbb.with_mrange(mmin.into(), mmax.into()), + // If either None or mismatch, leave it as None and don't error + _ => newbb, + }; + + newbb + }) +} + +/// Create a [`crate::file::statistics::Statistics`] from a thrift [`Statistics`] object. +fn convert_stats( + column_descr: &Arc, + thrift_stats: Option, +) -> Result> { + use crate::file::statistics::Statistics as FStatistics; + Ok(match thrift_stats { + Some(stats) => { + // Generic null count. + let null_count = stats + .null_count + .map(|null_count| { + if null_count < 0 { + return Err(general_err!( + "Statistics null count is negative {}", + null_count + )); + } + Ok(null_count as u64) + }) + .transpose()?; + // Generic distinct count (count of distinct values occurring) + let distinct_count = stats.distinct_count.map(|value| value as u64); + // Whether or not statistics use deprecated min/max fields. + let old_format = stats.min_value.is_none() && stats.max_value.is_none(); + // Generic min value as bytes. + let min = if old_format { + stats.min + } else { + stats.min_value + }; + // Generic max value as bytes. + let max = if old_format { + stats.max + } else { + stats.max_value + }; + + fn check_len(min: &Option<&[u8]>, max: &Option<&[u8]>, len: usize) -> Result<()> { + if let Some(min) = min { + if min.len() < len { + return Err(general_err!("Insufficient bytes to parse min statistic",)); + } + } + if let Some(max) = max { + if max.len() < len { + return Err(general_err!("Insufficient bytes to parse max statistic",)); + } + } + Ok(()) + } + + let physical_type = column_descr.physical_type(); + match physical_type { + Type::BOOLEAN => check_len(&min, &max, 1), + Type::INT32 | Type::FLOAT => check_len(&min, &max, 4), + Type::INT64 | Type::DOUBLE => check_len(&min, &max, 8), + Type::INT96 => check_len(&min, &max, 12), + _ => Ok(()), + }?; + + // Values are encoded using PLAIN encoding definition, except that + // variable-length byte arrays do not include a length prefix. + // + // Instead of using actual decoder, we manually convert values. + let res = match physical_type { + Type::BOOLEAN => FStatistics::boolean( + min.map(|data| data[0] != 0), + max.map(|data| data[0] != 0), + distinct_count, + null_count, + old_format, + ), + Type::INT32 => FStatistics::int32( + min.map(|data| i32::from_le_bytes(data[..4].try_into().unwrap())), + max.map(|data| i32::from_le_bytes(data[..4].try_into().unwrap())), + distinct_count, + null_count, + old_format, + ), + Type::INT64 => FStatistics::int64( + min.map(|data| i64::from_le_bytes(data[..8].try_into().unwrap())), + max.map(|data| i64::from_le_bytes(data[..8].try_into().unwrap())), + distinct_count, + null_count, + old_format, + ), + Type::INT96 => { + // INT96 statistics may not be correct, because comparison is signed + let min = if let Some(data) = min { + if data.len() != 12 { + return Err(general_err!("Incorrect Int96 min statistics")); + } + Some(Int96::try_from_le_slice(data)?) + } else { + None + }; + let max = if let Some(data) = max { + if data.len() != 12 { + return Err(general_err!("Incorrect Int96 max statistics")); + } + Some(Int96::try_from_le_slice(data)?) + } else { + None + }; + FStatistics::int96(min, max, distinct_count, null_count, old_format) + } + Type::FLOAT => FStatistics::float( + min.map(|data| f32::from_le_bytes(data[..4].try_into().unwrap())), + max.map(|data| f32::from_le_bytes(data[..4].try_into().unwrap())), + distinct_count, + null_count, + old_format, + ), + Type::DOUBLE => FStatistics::double( + min.map(|data| f64::from_le_bytes(data[..8].try_into().unwrap())), + max.map(|data| f64::from_le_bytes(data[..8].try_into().unwrap())), + distinct_count, + null_count, + old_format, + ), + Type::BYTE_ARRAY => FStatistics::ByteArray( + ValueStatistics::new( + min.map(ByteArray::from), + max.map(ByteArray::from), + distinct_count, + null_count, + old_format, + ) + .with_max_is_exact(stats.is_max_value_exact.unwrap_or(false)) + .with_min_is_exact(stats.is_min_value_exact.unwrap_or(false)), + ), + Type::FIXED_LEN_BYTE_ARRAY => FStatistics::FixedLenByteArray( + ValueStatistics::new( + min.map(ByteArray::from).map(FixedLenByteArray::from), + max.map(ByteArray::from).map(FixedLenByteArray::from), + distinct_count, + null_count, + old_format, + ) + .with_max_is_exact(stats.is_max_value_exact.unwrap_or(false)) + .with_min_is_exact(stats.is_min_value_exact.unwrap_or(false)), + ), + }; + + Some(res) + } + None => None, + }) +} + +// bit positions for required fields in the Thrift ColumnMetaData struct +const COL_META_TYPE: u16 = 1 << 1; +const COL_META_ENCODINGS: u16 = 1 << 2; +const COL_META_CODEC: u16 = 1 << 4; +const COL_META_NUM_VALUES: u16 = 1 << 5; +const COL_META_TOTAL_UNCOMP_SZ: u16 = 1 << 6; +const COL_META_TOTAL_COMP_SZ: u16 = 1 << 7; +const COL_META_DATA_PAGE_OFFSET: u16 = 1 << 9; + +// a mask where all required fields' bits are set +const COL_META_ALL_REQUIRED: u16 = COL_META_TYPE + | COL_META_ENCODINGS + | COL_META_CODEC + | COL_META_NUM_VALUES + | COL_META_TOTAL_UNCOMP_SZ + | COL_META_TOTAL_COMP_SZ + | COL_META_DATA_PAGE_OFFSET; + +// check mask to see if all required fields are set. return an appropriate error if +// any are missing. +fn validate_column_metadata(mask: u16) -> Result<()> { + if mask != COL_META_ALL_REQUIRED { + if mask & COL_META_ENCODINGS == 0 { + return Err(general_err!("Required field encodings is missing")); + } + + if mask & COL_META_CODEC == 0 { + return Err(general_err!("Required field codec is missing")); + } + if mask & COL_META_NUM_VALUES == 0 { + return Err(general_err!("Required field num_values is missing")); + } + if mask & COL_META_TOTAL_UNCOMP_SZ == 0 { + return Err(general_err!( + "Required field total_uncompressed_size is missing" + )); + } + if mask & COL_META_TOTAL_COMP_SZ == 0 { + return Err(general_err!( + "Required field total_compressed_size is missing" + )); + } + if mask & COL_META_DATA_PAGE_OFFSET == 0 { + return Err(general_err!("Required field data_page_offset is missing")); + } + } + + Ok(()) +} + +fn read_encoding_stats_as_mask<'a>( + prot: &mut ThriftSliceInputProtocol<'a>, +) -> Result { + // read the vector of stats, setting mask bits for data pages + let mut mask = 0i32; + let list_ident = prot.read_list_begin()?; + // check for PageEncodingStats struct + validate_list_type(ElementType::Struct, &list_ident)?; + + for _ in 0..list_ident.size { + let pes = PageEncodingStats::read_thrift(prot)?; + match pes.page_type { + PageType::DATA_PAGE | PageType::DATA_PAGE_V2 => mask |= 1 << pes.encoding as i32, + _ => {} + } + } + EncodingMask::try_new(mask) +} + +// Decode `ColumnMetaData`. Returns a mask of all required fields that were observed. +// This mask can be passed to `validate_column_metadata`. +fn read_column_metadata<'a>( + prot: &mut ThriftSliceInputProtocol<'a>, + column: &mut ColumnChunkMetaData, + col_index: usize, + options: Option<&ParquetMetaDataOptions>, +) -> Result { + // mask for seen required fields in ColumnMetaData + let mut seen_mask = 0u16; + + let mut skip_pes = false; + let mut pes_mask = true; + let mut skip_col_stats = false; + let mut skip_size_stats = false; + + if let Some(opts) = options { + skip_pes = opts.skip_encoding_stats(col_index); + pes_mask = opts.encoding_stats_as_mask(); + skip_col_stats = opts.skip_column_stats(col_index); + skip_size_stats = opts.skip_size_stats(col_index); + } + + // struct ColumnMetaData { + // 1: required Type type + // 2: required list encodings + // 3: required list path_in_schema + // 4: required CompressionCodec codec + // 5: required i64 num_values + // 6: required i64 total_uncompressed_size + // 7: required i64 total_compressed_size + // 8: optional list key_value_metadata + // 9: required i64 data_page_offset + // 10: optional i64 index_page_offset + // 11: optional i64 dictionary_page_offset + // 12: optional Statistics statistics; + // 13: optional list encoding_stats; + // 14: optional i64 bloom_filter_offset; + // 15: optional i32 bloom_filter_length; + // 16: optional SizeStatistics size_statistics; + // 17: optional GeospatialStatistics geospatial_statistics; + // } + let column_descr = &column.column_descr; + + let mut last_field_id = 0i16; + loop { + let field_ident = prot.read_field_begin(last_field_id)?; + if field_ident.field_type == FieldType::Stop { + break; + } + match field_ident.id { + // 1: type is never used, we can use the column descriptor + 1 => { + // read for error handling + Type::read_thrift(&mut *prot)?; + seen_mask |= COL_META_TYPE; + } + 2 => { + column.encodings = EncodingMask::read_thrift(&mut *prot)?; + seen_mask |= COL_META_ENCODINGS; + } + // 3: path_in_schema is redundant + 4 => { + column.compression = CompressionCodec::read_thrift(&mut *prot)?; + seen_mask |= COL_META_CODEC; + } + 5 => { + column.num_values = i64::read_thrift(&mut *prot)?; + seen_mask |= COL_META_NUM_VALUES; + } + 6 => { + column.total_uncompressed_size = i64::read_thrift(&mut *prot)?; + seen_mask |= COL_META_TOTAL_UNCOMP_SZ; + } + 7 => { + column.total_compressed_size = i64::read_thrift(&mut *prot)?; + seen_mask |= COL_META_TOTAL_COMP_SZ; + } + // 8: we don't expose this key value + 9 => { + column.data_page_offset = i64::read_thrift(&mut *prot)?; + seen_mask |= COL_META_DATA_PAGE_OFFSET; + } + 10 => { + column.index_page_offset = Some(i64::read_thrift(&mut *prot)?); + } + 11 => { + column.dictionary_page_offset = Some(i64::read_thrift(&mut *prot)?); + } + 12 if !skip_col_stats => { + column.statistics = + convert_stats(column_descr, Some(Statistics::read_thrift(&mut *prot)?))?; + } + 13 if !skip_pes => { + if pes_mask { + let val = read_encoding_stats_as_mask(&mut *prot)?; + column.encoding_stats = Some(ParquetPageEncodingStats::Mask(val)); + } else { + let val = + read_thrift_vec::(&mut *prot)?; + column.encoding_stats = Some(ParquetPageEncodingStats::Full(val)); + } + } + 14 => { + column.bloom_filter_offset = Some(i64::read_thrift(&mut *prot)?); + } + 15 => { + column.bloom_filter_length = Some(i32::read_thrift(&mut *prot)?); + } + 16 if !skip_size_stats => { + let val = SizeStatistics::read_thrift(&mut *prot)?; + column.unencoded_byte_array_data_bytes = val.unencoded_byte_array_data_bytes; + column.repetition_level_histogram = + val.repetition_level_histogram.map(LevelHistogram::from); + column.definition_level_histogram = + val.definition_level_histogram.map(LevelHistogram::from); + } + 17 => { + let val = GeospatialStatistics::read_thrift(&mut *prot)?; + column.geo_statistics = convert_geo_stats(Some(val)); + } + _ => { + prot.skip(field_ident.field_type)?; + } + }; + last_field_id = field_ident.id; + } + + Ok(seen_mask) +} + +// using ThriftSliceInputProtocol rather than ThriftCompactInputProtocl trait because +// these are all internal and operate on slices. +fn read_column_chunk<'a>( + prot: &mut ThriftSliceInputProtocol<'a>, + column_descr: &Arc, + col_index: usize, + options: Option<&ParquetMetaDataOptions>, +) -> Result { + // create a default initialized ColumnMetaData + let mut col = ColumnChunkMetaDataBuilder::new(column_descr.clone()).build()?; + + // seen flag for file_offset + let mut has_file_offset = false; + + // mask of seen flags for ColumnMetaData + let mut col_meta_mask = 0u16; + + // struct ColumnChunk { + // 1: optional string file_path + // 2: required i64 file_offset = 0 + // 3: optional ColumnMetaData meta_data + // 4: optional i64 offset_index_offset + // 5: optional i32 offset_index_length + // 6: optional i64 column_index_offset + // 7: optional i32 column_index_length + // 8: optional ColumnCryptoMetaData crypto_metadata + // 9: optional binary encrypted_column_metadata + // } + let mut last_field_id = 0i16; + loop { + let field_ident = prot.read_field_begin(last_field_id)?; + if field_ident.field_type == FieldType::Stop { + break; + } + match field_ident.id { + 1 => { + col.file_path = Some(String::read_thrift(&mut *prot)?); + } + 2 => { + col.file_offset = i64::read_thrift(&mut *prot)?; + has_file_offset = true; + } + 3 => { + col_meta_mask = read_column_metadata(&mut *prot, &mut col, col_index, options)?; + } + 4 => { + col.offset_index_offset = Some(i64::read_thrift(&mut *prot)?); + } + 5 => { + col.offset_index_length = Some(i32::read_thrift(&mut *prot)?); + } + 6 => { + col.column_index_offset = Some(i64::read_thrift(&mut *prot)?); + } + 7 => { + col.column_index_length = Some(i32::read_thrift(&mut *prot)?); + } + #[cfg(feature = "encryption")] + 8 => { + let val = ColumnCryptoMetaData::read_thrift(&mut *prot)?; + col.column_crypto_metadata = Some(Box::new(val)); + } + #[cfg(feature = "encryption")] + 9 => { + col.encrypted_column_metadata = Some(<&[u8]>::read_thrift(&mut *prot)?.to_vec()); + } + _ => { + prot.skip(field_ident.field_type)?; + } + }; + last_field_id = field_ident.id; + } + + // the only required field from ColumnChunk + if !has_file_offset { + return Err(general_err!("Required field file_offset is missing")); + }; + + // if encrypted just return. we'll decrypt after finishing the footer and populate the rest. + #[cfg(feature = "encryption")] + if col.encrypted_column_metadata.is_some() { + return Ok(col); + } + + // not encrypted, so make sure all required fields were read + validate_column_metadata(col_meta_mask)?; + + Ok(col) +} + +fn read_row_group( + prot: &mut ThriftSliceInputProtocol, + schema_descr: &Arc, + options: Option<&ParquetMetaDataOptions>, +) -> Result { + // create default initialized RowGroupMetaData + let mut row_group = RowGroupMetaDataBuilder::new(schema_descr.clone()).build_unchecked(); + + // mask values for required fields + const RG_COLUMNS: u8 = 1 << 1; + const RG_TOT_BYTE_SIZE: u8 = 1 << 2; + const RG_NUM_ROWS: u8 = 1 << 3; + const RG_ALL_REQUIRED: u8 = RG_COLUMNS | RG_TOT_BYTE_SIZE | RG_NUM_ROWS; + + let mut mask = 0u8; + + // struct RowGroup { + // 1: required list columns + // 2: required i64 total_byte_size + // 3: required i64 num_rows + // 4: optional list sorting_columns + // 5: optional i64 file_offset + // 6: optional i64 total_compressed_size + // 7: optional i16 ordinal + // } + let mut last_field_id = 0i16; + loop { + let field_ident = prot.read_field_begin(last_field_id)?; + if field_ident.field_type == FieldType::Stop { + break; + } + match field_ident.id { + 1 => { + let list_ident = prot.read_list_begin()?; + // check for list of struct + validate_list_type(ElementType::Struct, &list_ident)?; + if schema_descr.num_columns() != list_ident.size as usize { + return Err(general_err!( + "Column count mismatch. Schema has {} columns while Row Group has {}", + schema_descr.num_columns(), + list_ident.size + )); + } + for i in 0..list_ident.size as usize { + let col = read_column_chunk(prot, &schema_descr.columns()[i], i, options)?; + row_group.columns.push(col); + } + mask |= RG_COLUMNS; + } + 2 => { + row_group.total_byte_size = i64::read_thrift(&mut *prot)?; + mask |= RG_TOT_BYTE_SIZE; + } + 3 => { + row_group.num_rows = i64::read_thrift(&mut *prot)?; + mask |= RG_NUM_ROWS; + } + 4 => { + let val = read_thrift_vec::(&mut *prot)?; + row_group.sorting_columns = Some(val); + } + 5 => { + row_group.file_offset = Some(i64::read_thrift(&mut *prot)?); + } + // 6: we don't expose total_compressed_size + 7 => { + row_group.ordinal = Some(i16::read_thrift(&mut *prot)?); + } + _ => { + prot.skip(field_ident.field_type)?; + } + }; + last_field_id = field_ident.id; + } + + if mask != RG_ALL_REQUIRED { + if mask & RG_COLUMNS == 0 { + return Err(general_err!("Required field columns is missing")); + } + if mask & RG_TOT_BYTE_SIZE == 0 { + return Err(general_err!("Required field total_byte_size is missing")); + } + if mask & RG_NUM_ROWS == 0 { + return Err(general_err!("Required field num_rows is missing")); + } + } + + Ok(row_group) +} + +/// Create a [`SchemaDescriptor`] from thrift input. The input buffer must contain a complete +/// Parquet footer. +pub(crate) fn parquet_schema_from_bytes(buf: &[u8]) -> Result { + let mut prot = ThriftSliceInputProtocol::new(buf); + + let mut last_field_id = 0i16; + loop { + let field_ident = prot.read_field_begin(last_field_id)?; + if field_ident.field_type == FieldType::Stop { + break; + } + match field_ident.id { + 2 => { + // read schema and convert to SchemaDescriptor for use when reading row groups + let val = read_thrift_vec::(&mut prot)?; + let val = parquet_schema_from_array(val)?; + return Ok(SchemaDescriptor::new(val)); + } + _ => prot.skip(field_ident.field_type)?, + } + last_field_id = field_ident.id; + } + Err(general_err!("Input does not contain a schema")) +} + +/// Create [`ParquetMetaData`] from thrift input. Note that this only decodes the file metadata in +/// the Parquet footer. Page indexes will need to be added later. +pub(crate) fn parquet_metadata_from_bytes( + buf: &[u8], + options: Option<&ParquetMetaDataOptions>, +) -> Result { + let mut prot = ThriftSliceInputProtocol::new(buf); + + // begin reading the file metadata + let mut version: Option = None; + let mut num_rows: Option = None; + let mut row_groups: Option> = None; + let mut key_value_metadata: Option> = None; + let mut created_by: Option<&str> = None; + let mut column_orders: Option> = None; + #[cfg(feature = "encryption")] + let mut encryption_algorithm: Option = None; + #[cfg(feature = "encryption")] + let mut footer_signing_key_metadata: Option<&[u8]> = None; + + // this will need to be set before parsing row groups + let mut schema_descr: Option> = None; + + // see if we already have a schema. + if let Some(options) = options { + schema_descr = options.schema().cloned(); + } + + // struct FileMetaData { + // 1: required i32 version + // 2: required list schema; + // 3: required i64 num_rows + // 4: required list row_groups + // 5: optional list key_value_metadata + // 6: optional string created_by + // 7: optional list column_orders; + // 8: optional EncryptionAlgorithm encryption_algorithm + // 9: optional binary footer_signing_key_metadata + // } + let mut last_field_id = 0i16; + loop { + let field_ident = prot.read_field_begin(last_field_id)?; + if field_ident.field_type == FieldType::Stop { + break; + } + match field_ident.id { + 1 => { + version = Some(i32::read_thrift(&mut prot)?); + } + 2 => { + // If schema was passed in, skip parsing it + if schema_descr.is_some() { + prot.skip(field_ident.field_type)?; + } else { + // read schema and convert to SchemaDescriptor for use when reading row groups + let val = + read_thrift_vec::(&mut prot)?; + let val = parquet_schema_from_array(val)?; + schema_descr = Some(Arc::new(SchemaDescriptor::new(val))); + } + } + 3 => { + num_rows = Some(i64::read_thrift(&mut prot)?); + } + 4 => { + if schema_descr.is_none() { + return Err(general_err!("Required field schema is missing")); + } + let schema_descr = schema_descr.as_ref().unwrap(); + let list_ident = prot.read_list_begin()?; + // check for list of struct + validate_list_type(ElementType::Struct, &list_ident)?; + let mut rg_vec = Vec::with_capacity(list_ident.size as usize); + + // Read row groups and handle ordinal assignment + let mut assigner = OrdinalAssigner::new(); + for ordinal in 0..list_ident.size { + let ordinal: i16 = ordinal.try_into().map_err(|_| { + ParquetError::General(format!( + "Row group ordinal {ordinal} exceeds i16 max value", + )) + })?; + let rg = read_row_group(&mut prot, schema_descr, options)?; + rg_vec.push(assigner.ensure(ordinal, rg)?); + } + row_groups = Some(rg_vec); + } + 5 => { + let val = read_thrift_vec::(&mut prot)?; + key_value_metadata = Some(val); + } + 6 => { + created_by = Some(<&str>::read_thrift(&mut prot)?); + } + 7 => { + let val = read_thrift_vec::(&mut prot)?; + column_orders = Some(val); + } + #[cfg(feature = "encryption")] + 8 => { + let val = EncryptionAlgorithm::read_thrift(&mut prot)?; + encryption_algorithm = Some(val); + } + #[cfg(feature = "encryption")] + 9 => { + footer_signing_key_metadata = Some(<&[u8]>::read_thrift(&mut prot)?); + } + _ => { + prot.skip(field_ident.field_type)?; + } + }; + last_field_id = field_ident.id; + } + let Some(version) = version else { + return Err(general_err!("Required field version is missing")); + }; + let Some(num_rows) = num_rows else { + return Err(general_err!("Required field num_rows is missing")); + }; + let Some(row_groups) = row_groups else { + return Err(general_err!("Required field row_groups is missing")); + }; + + let created_by = created_by.map(|c| c.to_owned()); + + // we've tested for `None` by now so this is safe + let schema_descr = schema_descr.unwrap(); + + // need to map read column orders to actual values based on the schema + if column_orders + .as_ref() + .is_some_and(|cos| cos.len() != schema_descr.num_columns()) + { + return Err(general_err!("Column order length mismatch")); + } + // replace default type defined column orders with ones having the correct sort order + // TODO(ets): this could instead be done above when decoding + let column_orders = column_orders.map(|mut cos| { + for (i, column) in schema_descr.columns().iter().enumerate() { + if let ColumnOrder::TYPE_DEFINED_ORDER(_) = cos[i] { + let sort_order = ColumnOrder::sort_order_for_type( + column.logical_type_ref(), + column.converted_type(), + column.physical_type(), + ); + cos[i] = ColumnOrder::TYPE_DEFINED_ORDER(sort_order); + } + } + cos + }); + + #[cfg(not(feature = "encryption"))] + let fmd = crate::file::metadata::FileMetaData::new( + version, + num_rows, + created_by, + key_value_metadata, + schema_descr, + column_orders, + ); + #[cfg(feature = "encryption")] + let fmd = crate::file::metadata::FileMetaData::new( + version, + num_rows, + created_by, + key_value_metadata, + schema_descr, + column_orders, + ) + .with_encryption_algorithm(encryption_algorithm) + .with_footer_signing_key_metadata(footer_signing_key_metadata.map(|v| v.to_vec())); + + Ok(ParquetMetaData::new(fmd, row_groups)) +} + +/// Assign [`RowGroupMetaData::ordinal`] if it is missing. +#[derive(Debug, Default)] +pub(crate) struct OrdinalAssigner { + first_has_ordinal: Option, +} + +impl OrdinalAssigner { + fn new() -> Self { + Default::default() + } + + /// Sets [`RowGroupMetaData::ordinal`] if it is missing. + /// + /// # Arguments + /// - actual_ordinal: The ordinal (index) of the row group being processed + /// in the file metadata. + /// - rg: The [`RowGroupMetaData`] to potentially modify. + /// + /// Ensures: + /// 1. If the first row group has an ordinal, all subsequent row groups must + /// also have ordinals. + /// 2. If the first row group does NOT have an ordinal, all subsequent row + /// groups must also not have ordinals. + fn ensure( + &mut self, + actual_ordinal: i16, + mut rg: RowGroupMetaData, + ) -> Result { + let rg_has_ordinal = rg.ordinal.is_some(); + + // Only set first_has_ordinal if it's None (first row group that arrives) + if self.first_has_ordinal.is_none() { + self.first_has_ordinal = Some(rg_has_ordinal); + } + + // assign ordinal if missing and consistent with first row group + let first_has_ordinal = self.first_has_ordinal.unwrap(); + if !first_has_ordinal && !rg_has_ordinal { + rg.ordinal = Some(actual_ordinal); + } else if first_has_ordinal != rg_has_ordinal { + return Err(general_err!( + "Inconsistent ordinal assignment: first_has_ordinal is set to \ + {} but row-group with actual ordinal {} has rg_has_ordinal set to {}", + first_has_ordinal, + actual_ordinal, + rg_has_ordinal + )); + } + Ok(rg) + } +} + +thrift_struct!( + pub(crate) struct IndexPageHeader {} +); + +thrift_struct!( +pub(crate) struct DictionaryPageHeader { + /// Number of values in the dictionary + 1: required i32 num_values; + + /// Encoding using this dictionary page + 2: required Encoding encoding + + /// If true, the entries in the dictionary are sorted in ascending order + 3: optional bool is_sorted; +} +); + +thrift_struct!( +/// Statistics for the page header. +/// +/// This is a duplicate of the [`Statistics`] struct above. Because the page reader uses +/// the [`Read`] API, we cannot read the min/max values as slices. This should not be +/// a huge problem since this crate no longer reads the page header statistics by default. +/// +/// [`Read`]: crate::parquet_thrift::ThriftReadInputProtocol +pub(crate) struct PageStatistics { + 1: optional binary max; + 2: optional binary min; + 3: optional i64 null_count; + 4: optional i64 distinct_count; + 5: optional binary max_value; + 6: optional binary min_value; + 7: optional bool is_max_value_exact; + 8: optional bool is_min_value_exact; +} +); + +thrift_struct!( +pub(crate) struct DataPageHeader { + 1: required i32 num_values + 2: required Encoding encoding + 3: required Encoding definition_level_encoding; + 4: required Encoding repetition_level_encoding; + 5: optional PageStatistics statistics; +} +); + +impl DataPageHeader { + // reader that skips decoding page statistics + fn read_thrift_without_stats<'a, R>(prot: &mut R) -> Result + where + R: ThriftCompactInputProtocol<'a>, + { + let mut num_values: Option = None; + let mut encoding: Option = None; + let mut definition_level_encoding: Option = None; + let mut repetition_level_encoding: Option = None; + let statistics: Option = None; + let mut last_field_id = 0i16; + loop { + let field_ident = prot.read_field_begin(last_field_id)?; + if field_ident.field_type == FieldType::Stop { + break; + } + match field_ident.id { + 1 => { + let val = i32::read_thrift(&mut *prot)?; + num_values = Some(val); + } + 2 => { + let val = Encoding::read_thrift(&mut *prot)?; + encoding = Some(val); + } + 3 => { + let val = Encoding::read_thrift(&mut *prot)?; + definition_level_encoding = Some(val); + } + 4 => { + let val = Encoding::read_thrift(&mut *prot)?; + repetition_level_encoding = Some(val); + } + _ => { + prot.skip(field_ident.field_type)?; + } + }; + last_field_id = field_ident.id; + } + let Some(num_values) = num_values else { + return Err(general_err!("Required field num_values is missing")); + }; + let Some(encoding) = encoding else { + return Err(general_err!("Required field encoding is missing")); + }; + let Some(definition_level_encoding) = definition_level_encoding else { + return Err(general_err!( + "Required field definition_level_encoding is missing" + )); + }; + let Some(repetition_level_encoding) = repetition_level_encoding else { + return Err(general_err!( + "Required field repetition_level_encoding is missing" + )); + }; + Ok(Self { + num_values, + encoding, + definition_level_encoding, + repetition_level_encoding, + statistics, + }) + } +} + +thrift_struct!( +pub(crate) struct DataPageHeaderV2 { + 1: required i32 num_values + 2: required i32 num_nulls + 3: required i32 num_rows + 4: required Encoding encoding + 5: required i32 definition_levels_byte_length; + 6: required i32 repetition_levels_byte_length; + 7: optional bool is_compressed = true; + 8: optional PageStatistics statistics; +} +); + +impl DataPageHeaderV2 { + // reader that skips decoding page statistics + fn read_thrift_without_stats<'a, R>(prot: &mut R) -> Result + where + R: ThriftCompactInputProtocol<'a>, + { + let mut num_values: Option = None; + let mut num_nulls: Option = None; + let mut num_rows: Option = None; + let mut encoding: Option = None; + let mut definition_levels_byte_length: Option = None; + let mut repetition_levels_byte_length: Option = None; + let mut is_compressed: Option = None; + let statistics: Option = None; + let mut last_field_id = 0i16; + loop { + let field_ident = prot.read_field_begin(last_field_id)?; + if field_ident.field_type == FieldType::Stop { + break; + } + match field_ident.id { + 1 => { + let val = i32::read_thrift(&mut *prot)?; + num_values = Some(val); + } + 2 => { + let val = i32::read_thrift(&mut *prot)?; + num_nulls = Some(val); + } + 3 => { + let val = i32::read_thrift(&mut *prot)?; + num_rows = Some(val); + } + 4 => { + let val = Encoding::read_thrift(&mut *prot)?; + encoding = Some(val); + } + 5 => { + let val = i32::read_thrift(&mut *prot)?; + definition_levels_byte_length = Some(val); + } + 6 => { + let val = i32::read_thrift(&mut *prot)?; + repetition_levels_byte_length = Some(val); + } + 7 => { + is_compressed = Some(field_ident.bool_val()?); + } + _ => { + prot.skip(field_ident.field_type)?; + } + }; + last_field_id = field_ident.id; + } + let Some(num_values) = num_values else { + return Err(general_err!("Required field num_values is missing")); + }; + let Some(num_nulls) = num_nulls else { + return Err(general_err!("Required field num_nulls is missing")); + }; + let Some(num_rows) = num_rows else { + return Err(general_err!("Required field num_rows is missing")); + }; + let Some(encoding) = encoding else { + return Err(general_err!("Required field encoding is missing")); + }; + let Some(definition_levels_byte_length) = definition_levels_byte_length else { + return Err(general_err!( + "Required field definition_levels_byte_length is missing" + )); + }; + let Some(repetition_levels_byte_length) = repetition_levels_byte_length else { + return Err(general_err!( + "Required field repetition_levels_byte_length is missing" + )); + }; + Ok(Self { + num_values, + num_nulls, + num_rows, + encoding, + definition_levels_byte_length, + repetition_levels_byte_length, + is_compressed, + statistics, + }) + } +} + +thrift_struct!( +pub(crate) struct PageHeader { + /// the type of the page: indicates which of the *_header fields is set + 1: required PageType r#type + + /// Uncompressed page size in bytes (not including this header) + 2: required i32 uncompressed_page_size + + /// Compressed (and potentially encrypted) page size in bytes, not including this header + 3: required i32 compressed_page_size + + /// The 32-bit CRC checksum for the page, to be be calculated as follows: + 4: optional i32 crc + + // Headers for page specific data. One only will be set. + 5: optional DataPageHeader data_page_header; + 6: optional IndexPageHeader index_page_header; + 7: optional DictionaryPageHeader dictionary_page_header; + 8: optional DataPageHeaderV2 data_page_header_v2; +} +); + +impl PageHeader { + // reader that skips reading page statistics. obtained by running + // `cargo expand -p parquet --all-features --lib file::metadata::thrift` + // and modifying the impl of `read_thrift` + pub(crate) fn read_thrift_without_stats<'a, R>(prot: &mut R) -> Result + where + R: ThriftCompactInputProtocol<'a>, + { + let mut type_: Option = None; + let mut uncompressed_page_size: Option = None; + let mut compressed_page_size: Option = None; + let mut crc: Option = None; + let mut data_page_header: Option = None; + let mut index_page_header: Option = None; + let mut dictionary_page_header: Option = None; + let mut data_page_header_v2: Option = None; + let mut last_field_id = 0i16; + loop { + let field_ident = prot.read_field_begin(last_field_id)?; + if field_ident.field_type == FieldType::Stop { + break; + } + match field_ident.id { + 1 => { + let val = PageType::read_thrift(&mut *prot)?; + type_ = Some(val); + } + 2 => { + let val = i32::read_thrift(&mut *prot)?; + uncompressed_page_size = Some(val); + } + 3 => { + let val = i32::read_thrift(&mut *prot)?; + compressed_page_size = Some(val); + } + 4 => { + let val = i32::read_thrift(&mut *prot)?; + crc = Some(val); + } + 5 => { + let val = DataPageHeader::read_thrift_without_stats(&mut *prot)?; + data_page_header = Some(val); + } + 6 => { + let val = IndexPageHeader::read_thrift(&mut *prot)?; + index_page_header = Some(val); + } + 7 => { + let val = DictionaryPageHeader::read_thrift(&mut *prot)?; + dictionary_page_header = Some(val); + } + 8 => { + let val = DataPageHeaderV2::read_thrift_without_stats(&mut *prot)?; + data_page_header_v2 = Some(val); + } + _ => { + prot.skip(field_ident.field_type)?; + } + }; + last_field_id = field_ident.id; + } + let Some(type_) = type_ else { + return Err(general_err!("Required field type_ is missing")); + }; + let Some(uncompressed_page_size) = uncompressed_page_size else { + return Err(general_err!( + "Required field uncompressed_page_size is missing" + )); + }; + let Some(compressed_page_size) = compressed_page_size else { + return Err(general_err!( + "Required field compressed_page_size is missing" + )); + }; + Ok(Self { + r#type: type_, + uncompressed_page_size, + compressed_page_size, + crc, + data_page_header, + index_page_header, + dictionary_page_header, + data_page_header_v2, + }) + } +} + +///////////////////////////////////////////////// +// helper functions for writing file meta data + +#[cfg(feature = "encryption")] +fn should_write_column_stats(column_chunk: &ColumnChunkMetaData) -> bool { + // If there is encrypted column metadata present, + // the column is encrypted with a different key to the footer or a plaintext footer is used, + // so the statistics are sensitive and shouldn't be written. + column_chunk.encrypted_column_metadata.is_none() +} + +#[cfg(not(feature = "encryption"))] +fn should_write_column_stats(_column_chunk: &ColumnChunkMetaData) -> bool { + true +} + +// serialize the bits of the column chunk needed for a thrift ColumnMetaData +// struct ColumnMetaData { +// 1: required Type type +// 2: required list encodings +// 3: required list path_in_schema +// 4: required CompressionCodec codec +// 5: required i64 num_values +// 6: required i64 total_uncompressed_size +// 7: required i64 total_compressed_size +// 8: optional list key_value_metadata +// 9: required i64 data_page_offset +// 10: optional i64 index_page_offset +// 11: optional i64 dictionary_page_offset +// 12: optional Statistics statistics; +// 13: optional list encoding_stats; +// 14: optional i64 bloom_filter_offset; +// 15: optional i32 bloom_filter_length; +// 16: optional SizeStatistics size_statistics; +// 17: optional GeospatialStatistics geospatial_statistics; +// } +pub(super) fn serialize_column_meta_data( + column_chunk: &ColumnChunkMetaData, + w: &mut ThriftCompactOutputProtocol, +) -> Result<()> { + use crate::file::statistics::page_stats_to_thrift; + + column_chunk.column_type().write_thrift_field(w, 1, 0)?; + column_chunk + .encodings() + .collect::>() + .write_thrift_field(w, 2, 1)?; + if w.write_path_in_schema() { + let path = column_chunk.column_descr.path().parts(); + let path: Vec<&str> = path.iter().map(|v| v.as_str()).collect(); + path.write_thrift_field(w, 3, 2)?; + column_chunk.compression.write_thrift_field(w, 4, 3)?; + } else { + column_chunk.compression.write_thrift_field(w, 4, 2)?; + } + + column_chunk.num_values.write_thrift_field(w, 5, 4)?; + column_chunk + .total_uncompressed_size + .write_thrift_field(w, 6, 5)?; + column_chunk + .total_compressed_size + .write_thrift_field(w, 7, 6)?; + // no key_value_metadata here + let mut last_field_id = column_chunk.data_page_offset.write_thrift_field(w, 9, 7)?; + if let Some(index_page_offset) = column_chunk.index_page_offset { + last_field_id = index_page_offset.write_thrift_field(w, 10, last_field_id)?; + } + if let Some(dictionary_page_offset) = column_chunk.dictionary_page_offset { + last_field_id = dictionary_page_offset.write_thrift_field(w, 11, last_field_id)?; + } + + if should_write_column_stats(column_chunk) { + // PageStatistics is the same as thrift Statistics, but writable + let stats = page_stats_to_thrift(column_chunk.statistics()); + if let Some(stats) = stats { + last_field_id = stats.write_thrift_field(w, 12, last_field_id)?; + } + if let Some(page_encoding_stats) = column_chunk.page_encoding_stats() { + last_field_id = page_encoding_stats.write_thrift_field(w, 13, last_field_id)?; + } + if let Some(bloom_filter_offset) = column_chunk.bloom_filter_offset { + last_field_id = bloom_filter_offset.write_thrift_field(w, 14, last_field_id)?; + } + if let Some(bloom_filter_length) = column_chunk.bloom_filter_length { + last_field_id = bloom_filter_length.write_thrift_field(w, 15, last_field_id)?; + } + + // SizeStatistics + let size_stats = if column_chunk.unencoded_byte_array_data_bytes.is_some() + || column_chunk.repetition_level_histogram.is_some() + || column_chunk.definition_level_histogram.is_some() + { + let repetition_level_histogram = column_chunk + .repetition_level_histogram() + .map(|hist| hist.clone().into_inner()); + + let definition_level_histogram = column_chunk + .definition_level_histogram() + .map(|hist| hist.clone().into_inner()); + + Some(SizeStatistics { + unencoded_byte_array_data_bytes: column_chunk.unencoded_byte_array_data_bytes, + repetition_level_histogram, + definition_level_histogram, + }) + } else { + None + }; + if let Some(size_stats) = size_stats { + last_field_id = size_stats.write_thrift_field(w, 16, last_field_id)?; + } + + if let Some(geo_stats) = column_chunk.geo_statistics() { + geo_stats.write_thrift_field(w, 17, last_field_id)?; + } + } + + w.write_struct_end() +} + +// temp struct used for writing +pub(super) struct FileMeta<'a> { + pub(super) file_metadata: &'a crate::file::metadata::FileMetaData, + pub(super) row_groups: &'a Vec, + // If true, then write the `path_in_schema` field in the ColumnMetaData struct. + pub(super) write_path_in_schema: bool, +} + +// struct FileMetaData { +// 1: required i32 version +// 2: required list schema; +// 3: required i64 num_rows +// 4: required list row_groups +// 5: optional list key_value_metadata +// 6: optional string created_by +// 7: optional list column_orders; +// 8: optional EncryptionAlgorithm encryption_algorithm +// 9: optional binary footer_signing_key_metadata +// } +impl<'a> WriteThrift for FileMeta<'a> { + const ELEMENT_TYPE: ElementType = ElementType::Struct; + + // needed for last_field_id w/o encryption + #[allow(unused_assignments)] + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + writer.set_write_path_in_schema(self.write_path_in_schema); + + self.file_metadata + .version + .write_thrift_field(writer, 1, 0)?; + + // field 2 is schema. do depth-first traversal of tree, converting to SchemaElement and + // writing along the way. + let root = self.file_metadata.schema_descr().root_schema_ptr(); + let schema_len = num_nodes(&root)?; + writer.write_field_begin(FieldType::List, 2, 1)?; + writer.write_list_begin(ElementType::Struct, schema_len)?; + // recursively write Type nodes as SchemaElements + write_schema(&root, writer)?; + + self.file_metadata + .num_rows + .write_thrift_field(writer, 3, 2)?; + + // this will call RowGroupMetaData::write_thrift + let mut last_field_id = self.row_groups.write_thrift_field(writer, 4, 3)?; + + if let Some(kv_metadata) = self.file_metadata.key_value_metadata() { + last_field_id = kv_metadata.write_thrift_field(writer, 5, last_field_id)?; + } + if let Some(created_by) = self.file_metadata.created_by() { + last_field_id = created_by.write_thrift_field(writer, 6, last_field_id)?; + } + if let Some(column_orders) = self.file_metadata.column_orders() { + last_field_id = column_orders.write_thrift_field(writer, 7, last_field_id)?; + } + #[cfg(feature = "encryption")] + if let Some(algo) = self.file_metadata.encryption_algorithm.as_ref() { + last_field_id = algo.write_thrift_field(writer, 8, last_field_id)?; + } + #[cfg(feature = "encryption")] + if let Some(key) = self.file_metadata.footer_signing_key_metadata.as_ref() { + key.as_slice() + .write_thrift_field(writer, 9, last_field_id)?; + } + + writer.write_struct_end() + } +} + +fn write_schema( + schema: &TypePtr, + writer: &mut ThriftCompactOutputProtocol, +) -> Result<()> { + if !schema.is_group() { + return Err(general_err!("Root schema must be Group type")); + } + write_schema_helper(schema, writer) +} + +fn write_schema_helper( + node: &TypePtr, + writer: &mut ThriftCompactOutputProtocol, +) -> Result<()> { + match node.as_ref() { + crate::schema::types::Type::PrimitiveType { + basic_info, + physical_type, + type_length, + scale, + precision, + } => { + let element = SchemaElement { + r#type: Some(*physical_type), + type_length: if *type_length >= 0 { + Some(*type_length) + } else { + None + }, + repetition_type: Some(basic_info.repetition()), + name: basic_info.name(), + num_children: None, + converted_type: match basic_info.converted_type() { + ConvertedType::NONE => None, + other => Some(other), + }, + scale: if *scale >= 0 { Some(*scale) } else { None }, + precision: if *precision >= 0 { + Some(*precision) + } else { + None + }, + field_id: if basic_info.has_id() { + Some(basic_info.id()) + } else { + None + }, + logical_type: basic_info.logical_type_ref().cloned(), + }; + element.write_thrift(writer) + } + crate::schema::types::Type::GroupType { basic_info, fields } => { + let repetition = if basic_info.has_repetition() { + Some(basic_info.repetition()) + } else { + None + }; + + let element = SchemaElement { + r#type: None, + type_length: None, + repetition_type: repetition, + name: basic_info.name(), + num_children: Some(fields.len().try_into()?), + converted_type: match basic_info.converted_type() { + ConvertedType::NONE => None, + other => Some(other), + }, + scale: None, + precision: None, + field_id: if basic_info.has_id() { + Some(basic_info.id()) + } else { + None + }, + logical_type: basic_info.logical_type_ref().cloned(), + }; + + element.write_thrift(writer)?; + + // Add child elements for a group + for field in fields { + write_schema_helper(field, writer)?; + } + Ok(()) + } + } +} + +// struct RowGroup { +// 1: required list columns +// 2: required i64 total_byte_size +// 3: required i64 num_rows +// 4: optional list sorting_columns +// 5: optional i64 file_offset +// 6: optional i64 total_compressed_size +// 7: optional i16 ordinal +// } +impl WriteThrift for RowGroupMetaData { + const ELEMENT_TYPE: ElementType = ElementType::Struct; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + // this will call ColumnChunkMetaData::write_thrift + self.columns.write_thrift_field(writer, 1, 0)?; + self.total_byte_size.write_thrift_field(writer, 2, 1)?; + let mut last_field_id = self.num_rows.write_thrift_field(writer, 3, 2)?; + if let Some(sorting_columns) = self.sorting_columns() { + last_field_id = sorting_columns.write_thrift_field(writer, 4, last_field_id)?; + } + if let Some(file_offset) = self.file_offset() { + last_field_id = file_offset.write_thrift_field(writer, 5, last_field_id)?; + } + // this is optional, but we'll always write it + last_field_id = self + .compressed_size() + .write_thrift_field(writer, 6, last_field_id)?; + if let Some(ordinal) = self.ordinal() { + ordinal.write_thrift_field(writer, 7, last_field_id)?; + } + writer.write_struct_end() + } +} + +// struct ColumnChunk { +// 1: optional string file_path +// 2: required i64 file_offset = 0 +// 3: optional ColumnMetaData meta_data +// 4: optional i64 offset_index_offset +// 5: optional i32 offset_index_length +// 6: optional i64 column_index_offset +// 7: optional i32 column_index_length +// 8: optional ColumnCryptoMetaData crypto_metadata +// 9: optional binary encrypted_column_metadata +// } +impl WriteThrift for ColumnChunkMetaData { + const ELEMENT_TYPE: ElementType = ElementType::Struct; + + #[allow(unused_assignments)] + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + let mut last_field_id = 0i16; + if let Some(file_path) = self.file_path() { + last_field_id = file_path.write_thrift_field(writer, 1, last_field_id)?; + } + last_field_id = self + .file_offset() + .write_thrift_field(writer, 2, last_field_id)?; + + #[cfg(feature = "encryption")] + let write_meta_data = + self.encrypted_column_metadata.is_none() || self.plaintext_footer_mode; + #[cfg(not(feature = "encryption"))] + let write_meta_data = true; + + // When the footer is encrypted and encrypted_column_metadata is present, + // skip writing the plaintext meta_data field to reduce footer size. + // When the footer is plaintext (plaintext_footer_mode=true), we still write + // meta_data for backward compatibility with readers that expect it, but with + // sensitive fields (statistics, bloom filter info, etc.) stripped out. + if write_meta_data { + writer.write_field_begin(FieldType::Struct, 3, last_field_id)?; + serialize_column_meta_data(self, writer)?; + last_field_id = 3; + } + + if let Some(offset_idx_off) = self.offset_index_offset() { + last_field_id = offset_idx_off.write_thrift_field(writer, 4, last_field_id)?; + } + if let Some(offset_idx_len) = self.offset_index_length() { + last_field_id = offset_idx_len.write_thrift_field(writer, 5, last_field_id)?; + } + if let Some(column_idx_off) = self.column_index_offset() { + last_field_id = column_idx_off.write_thrift_field(writer, 6, last_field_id)?; + } + if let Some(column_idx_len) = self.column_index_length() { + last_field_id = column_idx_len.write_thrift_field(writer, 7, last_field_id)?; + } + #[cfg(feature = "encryption")] + { + if let Some(crypto_metadata) = self.crypto_metadata() { + last_field_id = crypto_metadata.write_thrift_field(writer, 8, last_field_id)?; + } + if let Some(encrypted_meta) = self.encrypted_column_metadata.as_ref() { + encrypted_meta + .as_slice() + .write_thrift_field(writer, 9, last_field_id)?; + } + } + + writer.write_struct_end() + } +} + +// struct GeospatialStatistics { +// 1: optional BoundingBox bbox; +// 2: optional list geospatial_types; +// } +impl WriteThrift for crate::geospatial::statistics::GeospatialStatistics { + const ELEMENT_TYPE: ElementType = ElementType::Struct; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + let mut last_field_id = 0i16; + if let Some(bbox) = self.bounding_box() { + last_field_id = bbox.write_thrift_field(writer, 1, last_field_id)?; + } + if let Some(geo_types) = self.geospatial_types() { + geo_types.write_thrift_field(writer, 2, last_field_id)?; + } + + writer.write_struct_end() + } +} + +// macro cannot handle qualified names +use crate::geospatial::statistics::GeospatialStatistics as RustGeospatialStatistics; +write_thrift_field!(RustGeospatialStatistics, FieldType::Struct); + +// struct BoundingBox { +// 1: required double xmin; +// 2: required double xmax; +// 3: required double ymin; +// 4: required double ymax; +// 5: optional double zmin; +// 6: optional double zmax; +// 7: optional double mmin; +// 8: optional double mmax; +// } +impl WriteThrift for crate::geospatial::bounding_box::BoundingBox { + const ELEMENT_TYPE: ElementType = ElementType::Struct; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + self.get_xmin().write_thrift_field(writer, 1, 0)?; + self.get_xmax().write_thrift_field(writer, 2, 1)?; + self.get_ymin().write_thrift_field(writer, 3, 2)?; + let mut last_field_id = self.get_ymax().write_thrift_field(writer, 4, 3)?; + + if let Some(zmin) = self.get_zmin() { + last_field_id = zmin.write_thrift_field(writer, 5, last_field_id)?; + } + if let Some(zmax) = self.get_zmax() { + last_field_id = zmax.write_thrift_field(writer, 6, last_field_id)?; + } + if let Some(mmin) = self.get_mmin() { + last_field_id = mmin.write_thrift_field(writer, 7, last_field_id)?; + } + if let Some(mmax) = self.get_mmax() { + mmax.write_thrift_field(writer, 8, last_field_id)?; + } + + writer.write_struct_end() + } +} + +// macro cannot handle qualified names +use crate::geospatial::bounding_box::BoundingBox as RustBoundingBox; +write_thrift_field!(RustBoundingBox, FieldType::Struct); + +#[cfg(test)] +pub(crate) mod tests { + use crate::basic::{Encoding, PageType, Type as PhysicalType}; + use crate::errors::Result; + use crate::file::metadata::thrift::{ + BoundingBox, DataPageHeaderV2, DictionaryPageHeader, PageHeader, SchemaElement, + write_schema, + }; + use crate::file::metadata::{ColumnChunkMetaData, ParquetMetaDataOptions, RowGroupMetaData}; + use crate::parquet_thrift::tests::test_roundtrip; + use crate::parquet_thrift::{ + ElementType, ThriftCompactOutputProtocol, ThriftSliceInputProtocol, WriteThrift, + read_thrift_vec, + }; + use crate::schema::types::{ + ColumnDescriptor, ColumnPath, SchemaDescriptor, TypePtr, num_nodes, + parquet_schema_from_array, + }; + use std::sync::Arc; + + // for testing. decode thrift encoded RowGroup + pub(crate) fn read_row_group( + buf: &mut [u8], + schema_descr: Arc, + ) -> Result { + let mut reader = ThriftSliceInputProtocol::new(buf); + crate::file::metadata::thrift::read_row_group(&mut reader, &schema_descr, None) + } + + pub(crate) fn read_column_chunk( + buf: &mut [u8], + column_descr: Arc, + ) -> Result { + read_column_chunk_with_options(buf, column_descr, None) + } + + pub(crate) fn read_column_chunk_with_options( + buf: &mut [u8], + column_descr: Arc, + options: Option<&ParquetMetaDataOptions>, + ) -> Result { + let mut reader = ThriftSliceInputProtocol::new(buf); + crate::file::metadata::thrift::read_column_chunk(&mut reader, &column_descr, 0, options) + } + + pub(crate) fn roundtrip_schema(schema: TypePtr) -> Result { + let num_nodes = num_nodes(&schema)?; + let mut buf = Vec::new(); + let mut writer = ThriftCompactOutputProtocol::new(&mut buf); + + // kick off writing list + writer.write_list_begin(ElementType::Struct, num_nodes)?; + + // write SchemaElements + write_schema(&schema, &mut writer)?; + + let mut prot = ThriftSliceInputProtocol::new(&buf); + let se: Vec = read_thrift_vec(&mut prot)?; + parquet_schema_from_array(se) + } + + pub(crate) fn schema_to_buf(schema: &TypePtr) -> Result> { + let num_nodes = num_nodes(schema)?; + let mut buf = Vec::new(); + let mut writer = ThriftCompactOutputProtocol::new(&mut buf); + + // kick off writing list + writer.write_list_begin(ElementType::Struct, num_nodes)?; + + // write SchemaElements + write_schema(schema, &mut writer)?; + Ok(buf) + } + + pub(crate) fn buf_to_schema_list<'a>(buf: &'a mut Vec) -> Result>> { + let mut prot = ThriftSliceInputProtocol::new(buf.as_mut_slice()); + read_thrift_vec(&mut prot) + } + + fn thrift_bytes(value: &T) -> Vec { + let mut buf = Vec::new(); + let mut writer = ThriftCompactOutputProtocol::new(&mut buf); + value.write_thrift(&mut writer).unwrap(); + buf + } + + fn change_false_bool_field_to_i32(buf: &mut [u8]) { + let pos = buf + .iter() + .rposition(|byte| *byte == 0x12) + .expect("expected BOOL_FALSE field header byte"); + buf[pos] = 0x15; + } + + fn assert_malformed_bool_error(err: crate::errors::ParquetError) { + let msg = err.to_string(); + assert!( + msg.contains("Unexpected struct field type"), + "unexpected error message: {msg}" + ); + } + + #[test] + fn test_bounding_box_roundtrip() { + test_roundtrip(BoundingBox { + xmin: 0.1.into(), + xmax: 10.3.into(), + ymin: 0.001.into(), + ymax: 128.5.into(), + zmin: None, + zmax: None, + mmin: None, + mmax: None, + }); + + test_roundtrip(BoundingBox { + xmin: 0.1.into(), + xmax: 10.3.into(), + ymin: 0.001.into(), + ymax: 128.5.into(), + zmin: Some(11.0.into()), + zmax: Some(1300.0.into()), + mmin: None, + mmax: None, + }); + + test_roundtrip(BoundingBox { + xmin: 0.1.into(), + xmax: 10.3.into(), + ymin: 0.001.into(), + ymax: 128.5.into(), + zmin: Some(11.0.into()), + zmax: Some(1300.0.into()), + mmin: Some(3.7.into()), + mmax: Some(42.0.into()), + }); + } + + #[test] + fn test_convert_stats_preserves_missing_null_count() { + let primitive = + crate::schema::types::Type::primitive_type_builder("col", PhysicalType::INT32) + .build() + .unwrap(); + let column_descr = Arc::new(ColumnDescriptor::new( + Arc::new(primitive), + 0, + 0, + ColumnPath::new(vec![]), + )); + + let none_null_count = super::Statistics { + max: None, + min: None, + null_count: None, + distinct_count: None, + max_value: None, + min_value: None, + is_max_value_exact: None, + is_min_value_exact: None, + }; + let decoded_none = super::convert_stats(&column_descr, Some(none_null_count)) + .unwrap() + .unwrap(); + assert_eq!(decoded_none.null_count_opt(), None); + + let zero_null_count = super::Statistics { + max: None, + min: None, + null_count: Some(0), + distinct_count: None, + max_value: None, + min_value: None, + is_max_value_exact: None, + is_min_value_exact: None, + }; + let decoded_zero = super::convert_stats(&column_descr, Some(zero_null_count)) + .unwrap() + .unwrap(); + assert_eq!(decoded_zero.null_count_opt(), Some(0)); + } + + #[test] + fn test_convert_stats_returns_error_for_overlong_int96_statistics() { + let primitive = + crate::schema::types::Type::primitive_type_builder("col", PhysicalType::INT96) + .build() + .unwrap(); + let column_descr = Arc::new(ColumnDescriptor::new( + Arc::new(primitive), + 0, + 0, + ColumnPath::new(vec![]), + )); + let invalid = (0..13).collect::>(); + + let make_stats = |min, max| super::Statistics { + max, + min, + null_count: Some(0), + distinct_count: None, + max_value: None, + min_value: None, + is_max_value_exact: None, + is_min_value_exact: None, + }; + + let err = super::convert_stats(&column_descr, Some(make_stats(Some(&invalid), None))) + .unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: Incorrect Int96 min statistics" + ); + + let err = super::convert_stats(&column_descr, Some(make_stats(None, Some(&invalid)))) + .unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: Incorrect Int96 max statistics" + ); + } + + #[test] + fn malformed_bool_field_returns_error_not_panic() { + let page_header = PageHeader { + r#type: PageType::DICTIONARY_PAGE, + uncompressed_page_size: 1, + compressed_page_size: 1, + crc: None, + data_page_header: None, + index_page_header: None, + dictionary_page_header: Some(DictionaryPageHeader { + num_values: 1, + encoding: Encoding::PLAIN, + is_sorted: Some(false), + }), + data_page_header_v2: None, + }; + + let mut buf = thrift_bytes(&page_header); + change_false_bool_field_to_i32(&mut buf); + + let mut prot = ThriftSliceInputProtocol::new(&buf); + let err = PageHeader::read_thrift_without_stats(&mut prot) + .expect_err("malformed bool field should return an error"); + assert_malformed_bool_error(err); + } + + #[test] + fn malformed_data_page_v2_bool_field_returns_error_not_panic() { + let data_page_header_v2 = DataPageHeaderV2 { + num_values: 1, + num_nulls: 0, + num_rows: 1, + encoding: Encoding::PLAIN, + definition_levels_byte_length: 0, + repetition_levels_byte_length: 0, + is_compressed: Some(false), + statistics: None, + }; + + let mut buf = thrift_bytes(&data_page_header_v2); + change_false_bool_field_to_i32(&mut buf); + + let mut prot = ThriftSliceInputProtocol::new(&buf); + let err = DataPageHeaderV2::read_thrift_without_stats(&mut prot) + .expect_err("malformed bool field should return an error"); + assert_malformed_bool_error(err); + } +} diff --git a/vendor/parquet/src/file/metadata/writer.rs b/vendor/parquet/src/file/metadata/writer.rs new file mode 100644 index 0000000..cd5d617 --- /dev/null +++ b/vendor/parquet/src/file/metadata/writer.rs @@ -0,0 +1,886 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::file::metadata::thrift::FileMeta; +use crate::file::metadata::{ + ColumnChunkMetaData, ParquetColumnIndex, ParquetOffsetIndex, RowGroupMetaData, +}; +use crate::schema::types::{SchemaDescPtr, SchemaDescriptor}; +use crate::{ + basic::ColumnOrder, + file::metadata::{FileMetaData, ParquetMetaDataBuilder}, +}; +#[cfg(feature = "encryption")] +use crate::{ + encryption::{ + encrypt::{FileEncryptor, encrypt_thrift_object, write_signed_plaintext_thrift_object}, + modules::{ModuleType, create_footer_aad, create_module_aad}, + }, + file::column_crypto_metadata::ColumnCryptoMetaData, + file::metadata::thrift::encryption::{AesGcmV1, EncryptionAlgorithm, FileCryptoMetaData}, +}; +use crate::{errors::Result, file::page_index::column_index::ColumnIndexMetaData}; + +use crate::{ + file::writer::{TrackedWrite, get_file_magic}, + parquet_thrift::WriteThrift, +}; +use crate::{ + file::{ + metadata::{KeyValue, ParquetMetaData}, + page_index::offset_index::OffsetIndexMetaData, + }, + parquet_thrift::ThriftCompactOutputProtocol, +}; +use std::io::Write; +use std::sync::Arc; + +/// Writes `crate::file::metadata` structures to a thrift encoded byte stream +/// +/// See [`ParquetMetaDataWriter`] for background and example. +pub(crate) struct ThriftMetadataWriter<'a, W: Write> { + buf: &'a mut TrackedWrite, + schema_descr: &'a SchemaDescPtr, + row_groups: Vec, + column_indexes: Option>>>, + offset_indexes: Option>>>, + key_value_metadata: Option>, + created_by: Option, + object_writer: MetadataObjectWriter, + writer_version: i32, + write_path_in_schema: bool, +} + +impl<'a, W: Write> ThriftMetadataWriter<'a, W> { + /// Serialize all the offset indexes to `self.buf`, + /// + /// Note: also updates the `ColumnChunk::offset_index_offset` and + /// `ColumnChunk::offset_index_length` to reflect the position and length + /// of the serialized offset indexes. + fn write_offset_indexes( + &mut self, + offset_indexes: &[Vec>], + ) -> Result<()> { + // iter row group + // iter each column + // write offset index to the file + for (row_group_idx, row_group) in self.row_groups.iter_mut().enumerate() { + for (column_idx, column_metadata) in row_group.columns.iter_mut().enumerate() { + if let Some(offset_index) = &offset_indexes[row_group_idx][column_idx] { + let start_offset = self.buf.bytes_written(); + self.object_writer.write_offset_index( + offset_index, + column_metadata, + row_group_idx, + column_idx, + &mut self.buf, + )?; + let end_offset = self.buf.bytes_written(); + // set offset and index for offset index + column_metadata.offset_index_offset = Some(start_offset as i64); + column_metadata.offset_index_length = Some((end_offset - start_offset) as i32); + } + } + } + Ok(()) + } + + /// Serialize all the column indexes to the `self.buf` + /// + /// Note: also updates the `ColumnChunk::column_index_offset` and + /// `ColumnChunk::column_index_length` to reflect the position and length + /// of the serialized column indexes. + fn write_column_indexes( + &mut self, + column_indexes: &[Vec>], + ) -> Result<()> { + // iter row group + // iter each column + // write column index to the file + for (row_group_idx, row_group) in self.row_groups.iter_mut().enumerate() { + for (column_idx, column_metadata) in row_group.columns.iter_mut().enumerate() { + if let Some(column_index) = &column_indexes[row_group_idx][column_idx] { + let start_offset = self.buf.bytes_written(); + // only update column_metadata if the write succeeds + if self.object_writer.write_column_index( + column_index, + column_metadata, + row_group_idx, + column_idx, + &mut self.buf, + )? { + let end_offset = self.buf.bytes_written(); + // set offset and index for offset index + column_metadata.column_index_offset = Some(start_offset as i64); + column_metadata.column_index_length = + Some((end_offset - start_offset) as i32); + } + } + } + } + Ok(()) + } + + /// Serialize the column indexes and transform to `Option` + fn finalize_column_indexes(&mut self) -> Result> { + let column_indexes = std::mem::take(&mut self.column_indexes); + + // Write column indexes to file + if let Some(column_indexes) = column_indexes.as_ref() { + self.write_column_indexes(column_indexes)?; + } + + // check to see if the index is `None` for every row group and column chunk + let all_none = column_indexes + .as_ref() + .is_some_and(|ci| ci.iter().all(|cii| cii.iter().all(|idx| idx.is_none()))); + + // transform from Option>>> to + // Option>> + let column_indexes: Option = if all_none { + None + } else { + column_indexes.map(|ovvi| { + ovvi.into_iter() + .map(|vi| { + vi.into_iter() + .map(|ci| ci.unwrap_or(ColumnIndexMetaData::NONE)) + .collect() + }) + .collect() + }) + }; + + Ok(column_indexes) + } + + /// Serialize the offset indexes and transform to `Option` + fn finalize_offset_indexes(&mut self) -> Result> { + let offset_indexes = std::mem::take(&mut self.offset_indexes); + + // Write offset indexes to file + if let Some(offset_indexes) = offset_indexes.as_ref() { + self.write_offset_indexes(offset_indexes)?; + } + + // check to see if the index is `None` for every row group and column chunk + let all_none = offset_indexes + .as_ref() + .is_some_and(|oi| oi.iter().all(|oii| oii.iter().all(|idx| idx.is_none()))); + + let offset_indexes: Option = if all_none { + None + } else { + // FIXME(ets): this will panic if there's a missing index. + offset_indexes.map(|ovvi| { + ovvi.into_iter() + .map(|vi| vi.into_iter().map(|oi| oi.unwrap()).collect()) + .collect() + }) + }; + + Ok(offset_indexes) + } + + /// Assembles and writes the final metadata to self.buf + pub fn finish(mut self) -> Result { + let num_rows = self.row_groups.iter().map(|x| x.num_rows).sum(); + + // serialize page indexes and transform to the proper form for use in ParquetMetaData + let column_indexes = self.finalize_column_indexes()?; + let offset_indexes = self.finalize_offset_indexes()?; + + // We only include ColumnOrder for leaf nodes. + // Currently only supported ColumnOrder is TypeDefinedOrder so we set this + // for all leaf nodes. + // Even if the column has an undefined sort order, such as INTERVAL, this + // is still technically the defined TYPEORDER so it should still be set. + let column_orders = self + .schema_descr + .columns() + .iter() + .map(|col| { + let sort_order = ColumnOrder::sort_order_for_type( + col.logical_type_ref(), + col.converted_type(), + col.physical_type(), + ); + ColumnOrder::TYPE_DEFINED_ORDER(sort_order) + }) + .collect(); + + // This field is optional, perhaps in cases where no min/max fields are set + // in any Statistics or ColumnIndex object in the whole file. + // But for simplicity we always set this field. + let column_orders = Some(column_orders); + + let (row_groups, unencrypted_row_groups) = self + .object_writer + .apply_row_group_encryption(self.row_groups)?; + + #[cfg(feature = "encryption")] + let (encryption_algorithm, footer_signing_key_metadata) = + self.object_writer.get_plaintext_footer_crypto_metadata(); + #[cfg(feature = "encryption")] + let file_metadata = FileMetaData::new( + self.writer_version, + num_rows, + self.created_by, + self.key_value_metadata, + self.schema_descr.clone(), + column_orders, + ) + .with_encryption_algorithm(encryption_algorithm) + .with_footer_signing_key_metadata(footer_signing_key_metadata); + + #[cfg(not(feature = "encryption"))] + let file_metadata = FileMetaData::new( + self.writer_version, + num_rows, + self.created_by, + self.key_value_metadata, + self.schema_descr.clone(), + column_orders, + ); + + let file_meta = FileMeta { + file_metadata: &file_metadata, + row_groups: &row_groups, + write_path_in_schema: self.write_path_in_schema, + }; + + // Write file metadata + let start_pos = self.buf.bytes_written(); + self.object_writer + .write_file_metadata(&file_meta, &mut self.buf)?; + let end_pos = self.buf.bytes_written(); + + // Write footer + let metadata_len = (end_pos - start_pos) as u32; + + self.buf.write_all(&metadata_len.to_le_bytes())?; + self.buf.write_all(self.object_writer.get_file_magic())?; + + // If row group metadata was encrypted, we replace the encrypted row groups with + // unencrypted metadata before it is returned to users. This allows the metadata + // to be usable for retrieving the row group statistics for example, without users + // needing to decrypt the metadata. + let builder = ParquetMetaDataBuilder::new(file_metadata) + .set_column_index(column_indexes) + .set_offset_index(offset_indexes); + + Ok(match unencrypted_row_groups { + Some(rg) => builder.set_row_groups(rg).build(), + None => builder.set_row_groups(row_groups).build(), + }) + } + + pub fn new( + buf: &'a mut TrackedWrite, + schema_descr: &'a SchemaDescPtr, + row_groups: Vec, + created_by: Option, + writer_version: i32, + write_path_in_schema: bool, + ) -> Self { + Self { + buf, + schema_descr, + row_groups, + column_indexes: None, + offset_indexes: None, + key_value_metadata: None, + created_by, + object_writer: Default::default(), + writer_version, + write_path_in_schema, + } + } + + pub fn with_column_indexes( + mut self, + column_indexes: Vec>>, + ) -> Self { + self.column_indexes = Some(column_indexes); + self + } + + pub fn with_offset_indexes( + mut self, + offset_indexes: Vec>>, + ) -> Self { + self.offset_indexes = Some(offset_indexes); + self + } + + pub fn with_key_value_metadata(mut self, key_value_metadata: Vec) -> Self { + self.key_value_metadata = Some(key_value_metadata); + self + } + + #[cfg(feature = "encryption")] + pub fn with_file_encryptor(mut self, file_encryptor: Option>) -> Self { + self.object_writer = self.object_writer.with_file_encryptor(file_encryptor); + self + } +} + +/// Writes [`ParquetMetaData`] to a byte stream +/// +/// This structure handles the details of writing the various parts of Parquet +/// metadata into a byte stream. It is used to write the metadata into a parquet +/// file and can also write metadata into other locations (such as a store of +/// bytes). +/// +/// # Discussion +/// +/// The process of writing Parquet metadata is tricky because the +/// metadata is not stored as a single inline thrift structure. It can have +/// several "out of band" structures such as the [`OffsetIndex`] and +/// BloomFilters stored in separate structures whose locations are stored as +/// offsets from the beginning of the file. +/// +/// Note: this writer does not directly write BloomFilters. In order to write +/// BloomFilters, write the bloom filters into the buffer before creating the +/// metadata writer. Then set the corresponding `bloom_filter_offset` and +/// `bloom_filter_length` on [`ColumnChunkMetaData`] passed to this writer. +/// +/// # Output Format +/// +/// The format of the metadata is as follows: +/// +/// 1. Optional [`ColumnIndex`] (thrift encoded) +/// 2. Optional [`OffsetIndex`] (thrift encoded) +/// 3. [`FileMetaData`] (thrift encoded) +/// 4. Length of encoded `FileMetaData` (4 bytes, little endian) +/// 5. Parquet Magic Bytes (4 bytes) +/// +/// [`FileMetaData`]: https://github.com/apache/parquet-format/tree/master?tab=readme-ov-file#metadata +/// [`ColumnChunkMetaData`]: crate::file::metadata::ColumnChunkMetaData +/// [`ColumnIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md +/// [`OffsetIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md +/// +/// ```text +/// ┌──────────────────────┐ +/// │ │ +/// │ ... │ +/// │ │ +/// │┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐ │ +/// │ ColumnIndex ◀│─ ─ ─ +/// ││ (Optional) │ │ │ +/// │ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ │ +/// │┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐ │ │ FileMetadata +/// │ OffsetIndex │ contains embedded +/// ││ (Optional) │◀┼ ─ │ offsets to +/// │ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ │ │ ColumnIndex and +/// │╔═══════════════════╗ │ │ OffsetIndex +/// │║ ║ │ │ +/// │║ ║ ┼ ─ │ +/// │║ FileMetadata ║ │ +/// │║ ║ ┼ ─ ─ ┘ +/// │║ ║ │ +/// │╚═══════════════════╝ │ +/// │┌───────────────────┐ │ +/// ││ metadata length │ │ length of FileMetadata (only) +/// │└───────────────────┘ │ +/// │┌───────────────────┐ │ +/// ││ 'PAR1' │ │ Parquet Magic Bytes +/// │└───────────────────┘ │ +/// └──────────────────────┘ +/// Output Buffer +/// ``` +/// +/// # Example +/// ```no_run +/// # use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataWriter}; +/// # fn get_metadata() -> ParquetMetaData { unimplemented!(); } +/// // write parquet metadata to an in-memory buffer +/// let mut buffer = vec![]; +/// let metadata: ParquetMetaData = get_metadata(); +/// let writer = ParquetMetaDataWriter::new(&mut buffer, &metadata); +/// // write the metadata to the buffer +/// writer.finish().unwrap(); +/// assert!(!buffer.is_empty()); +/// ``` +pub struct ParquetMetaDataWriter<'a, W: Write> { + buf: TrackedWrite, + metadata: &'a ParquetMetaData, + write_path_in_schema: bool, +} + +impl<'a, W: Write> ParquetMetaDataWriter<'a, W> { + /// Create a new `ParquetMetaDataWriter` to write to `buf` + /// + /// Note any embedded offsets in the metadata will be written assuming the + /// metadata is at the start of the buffer. If the metadata is being written + /// to a location other than the start of the buffer, see [`Self::new_with_tracked`] + /// + /// See example on the struct level documentation + pub fn new(buf: W, metadata: &'a ParquetMetaData) -> Self { + Self::new_with_tracked(TrackedWrite::new(buf), metadata) + } + + /// Create a new ParquetMetaDataWriter to write to `buf` + /// + /// This method is used when the metadata is being written to a location other + /// than the start of the buffer. + /// + /// See example on the struct level documentation + pub fn new_with_tracked(buf: TrackedWrite, metadata: &'a ParquetMetaData) -> Self { + Self { + buf, + metadata, + write_path_in_schema: true, + } + } + + /// Set whether or not to write the `path_in_schema` field in the Thrift `ColumnMetaData` + /// struct. + pub fn with_write_path_in_schema(self, val: bool) -> Self { + Self { + write_path_in_schema: val, + ..self + } + } + + /// Write the metadata to the buffer + pub fn finish(mut self) -> Result<()> { + let file_metadata = self.metadata.file_metadata(); + + let schema = Arc::new(file_metadata.schema().clone()); + let schema_descr = Arc::new(SchemaDescriptor::new(schema.clone())); + let created_by = file_metadata.created_by().map(str::to_string); + + let row_groups = self.metadata.row_groups.clone(); + + let key_value_metadata = file_metadata.key_value_metadata().cloned(); + + let column_indexes = self.convert_column_indexes(); + let offset_indexes = self.convert_offset_index(); + + let mut encoder = ThriftMetadataWriter::new( + &mut self.buf, + &schema_descr, + row_groups, + created_by, + file_metadata.version(), + self.write_path_in_schema, + ); + + if let Some(column_indexes) = column_indexes { + encoder = encoder.with_column_indexes(column_indexes); + } + + if let Some(offset_indexes) = offset_indexes { + encoder = encoder.with_offset_indexes(offset_indexes); + } + + if let Some(key_value_metadata) = key_value_metadata { + encoder = encoder.with_key_value_metadata(key_value_metadata); + } + encoder.finish()?; + + Ok(()) + } + + fn convert_column_indexes(&self) -> Option>>> { + // TODO(ets): we're converting from ParquetColumnIndex to vec>, + // but then converting back to ParquetColumnIndex in the end. need to unify this. + self.metadata + .column_index() + .map(|row_group_column_indexes| { + (0..self.metadata.row_groups().len()) + .map(|rg_idx| { + let column_indexes = &row_group_column_indexes[rg_idx]; + column_indexes + .iter() + .map(|column_index| Some(column_index.clone())) + .collect() + }) + .collect() + }) + } + + fn convert_offset_index(&self) -> Option>>> { + self.metadata + .offset_index() + .map(|row_group_offset_indexes| { + (0..self.metadata.row_groups().len()) + .map(|rg_idx| { + let offset_indexes = &row_group_offset_indexes[rg_idx]; + offset_indexes + .iter() + .map(|offset_index| Some(offset_index.clone())) + .collect() + }) + .collect() + }) + } +} + +#[derive(Debug, Default)] +struct MetadataObjectWriter { + #[cfg(feature = "encryption")] + file_encryptor: Option>, +} + +impl MetadataObjectWriter { + #[inline] + fn write_thrift_object(object: &impl WriteThrift, sink: impl Write) -> Result<()> { + let mut protocol = ThriftCompactOutputProtocol::new(sink); + object.write_thrift(&mut protocol)?; + Ok(()) + } +} + +/// Implementations of [`MetadataObjectWriter`] methods for when encryption is disabled +#[cfg(not(feature = "encryption"))] +impl MetadataObjectWriter { + /// Write [`FileMetaData`] in Thrift format + /// + /// [`FileMetaData`]: https://github.com/apache/parquet-format/tree/master?tab=readme-ov-file#metadata + fn write_file_metadata(&self, file_metadata: &FileMeta, sink: impl Write) -> Result<()> { + Self::write_thrift_object(file_metadata, sink) + } + + /// Write a column [`OffsetIndex`] in Thrift format + /// + /// [`OffsetIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + fn write_offset_index( + &self, + offset_index: &OffsetIndexMetaData, + _column_chunk: &ColumnChunkMetaData, + _row_group_idx: usize, + _column_idx: usize, + sink: impl Write, + ) -> Result<()> { + Self::write_thrift_object(offset_index, sink) + } + + /// Write a column [`ColumnIndex`] in Thrift format + /// + /// If `column_index` is [`ColumnIndexMetaData::NONE`] the index will not be written and + /// this will return `false`. Returns `true` otherwise. + /// + /// [`ColumnIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + fn write_column_index( + &self, + column_index: &ColumnIndexMetaData, + _column_chunk: &ColumnChunkMetaData, + _row_group_idx: usize, + _column_idx: usize, + sink: impl Write, + ) -> Result { + match column_index { + // Missing indexes may also have the placeholder ColumnIndexMetaData::NONE + ColumnIndexMetaData::NONE => Ok(false), + _ => { + Self::write_thrift_object(column_index, sink)?; + Ok(true) + } + } + } + + /// No-op implementation of row-group metadata encryption + fn apply_row_group_encryption( + &self, + row_groups: Vec, + ) -> Result<(Vec, Option>)> { + Ok((row_groups, None)) + } + + /// Get the "magic" bytes identifying the file type + pub fn get_file_magic(&self) -> &[u8; 4] { + get_file_magic() + } +} + +/// Implementations of [`MetadataObjectWriter`] methods that rely on encryption being enabled +#[cfg(feature = "encryption")] +impl MetadataObjectWriter { + /// Set the file encryptor to use + fn with_file_encryptor(mut self, encryptor: Option>) -> Self { + self.file_encryptor = encryptor; + self + } + + /// Write [`FileMetaData`] in Thrift format, possibly encrypting it if required + /// + /// [`FileMetaData`]: https://github.com/apache/parquet-format/tree/master?tab=readme-ov-file#metadata + fn write_file_metadata(&self, file_metadata: &FileMeta, mut sink: impl Write) -> Result<()> { + match self.file_encryptor.as_ref() { + Some(file_encryptor) if file_encryptor.properties().encrypt_footer() => { + // First write FileCryptoMetadata + let crypto_metadata = Self::file_crypto_metadata(file_encryptor)?; + let mut protocol = ThriftCompactOutputProtocol::new(&mut sink); + crypto_metadata.write_thrift(&mut protocol)?; + + // Then write encrypted footer + let aad = create_footer_aad(file_encryptor.file_aad())?; + let mut encryptor = file_encryptor.get_footer_encryptor()?; + encrypt_thrift_object(file_metadata, &mut encryptor, &mut sink, &aad) + } + Some(file_encryptor) if file_metadata.file_metadata.encryption_algorithm.is_some() => { + let aad = create_footer_aad(file_encryptor.file_aad())?; + let mut encryptor = file_encryptor.get_footer_encryptor()?; + write_signed_plaintext_thrift_object(file_metadata, &mut encryptor, &mut sink, &aad) + } + _ => Self::write_thrift_object(file_metadata, &mut sink), + } + } + + /// Write a column [`OffsetIndex`] in Thrift format, possibly encrypting it if required + /// + /// [`OffsetIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + fn write_offset_index( + &self, + offset_index: &OffsetIndexMetaData, + column_chunk: &ColumnChunkMetaData, + row_group_idx: usize, + column_idx: usize, + sink: impl Write, + ) -> Result<()> { + match &self.file_encryptor { + Some(file_encryptor) => Self::write_thrift_object_with_encryption( + offset_index, + sink, + file_encryptor, + column_chunk, + ModuleType::OffsetIndex, + row_group_idx, + column_idx, + ), + None => Self::write_thrift_object(offset_index, sink), + } + } + + /// Write a column [`ColumnIndex`] in Thrift format, possibly encrypting it if required + /// + /// If `column_index` is [`ColumnIndexMetaData::NONE`] the index will not be written and + /// this will return `false`. Returns `true` otherwise. + /// + /// [`ColumnIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + fn write_column_index( + &self, + column_index: &ColumnIndexMetaData, + column_chunk: &ColumnChunkMetaData, + row_group_idx: usize, + column_idx: usize, + sink: impl Write, + ) -> Result { + match column_index { + // Missing indexes may also have the placeholder ColumnIndexMetaData::NONE + ColumnIndexMetaData::NONE => Ok(false), + _ => { + match &self.file_encryptor { + Some(file_encryptor) => Self::write_thrift_object_with_encryption( + column_index, + sink, + file_encryptor, + column_chunk, + ModuleType::ColumnIndex, + row_group_idx, + column_idx, + )?, + None => Self::write_thrift_object(column_index, sink)?, + } + Ok(true) + } + } + } + + /// If encryption is enabled and configured, encrypt row group metadata. + /// Returns a tuple of the row group metadata to write, + /// and possibly unencrypted metadata to be returned to clients if data was encrypted. + fn apply_row_group_encryption( + &self, + row_groups: Vec, + ) -> Result<(Vec, Option>)> { + match &self.file_encryptor { + Some(file_encryptor) => { + let unencrypted_row_groups = row_groups.clone(); + let encrypted_row_groups = Self::encrypt_row_groups(row_groups, file_encryptor)?; + Ok((encrypted_row_groups, Some(unencrypted_row_groups))) + } + None => Ok((row_groups, None)), + } + } + + /// Get the "magic" bytes identifying the file type + fn get_file_magic(&self) -> &[u8; 4] { + get_file_magic( + self.file_encryptor + .as_ref() + .map(|encryptor| encryptor.properties()), + ) + } + + fn write_thrift_object_with_encryption( + object: &impl WriteThrift, + mut sink: impl Write, + file_encryptor: &FileEncryptor, + column_metadata: &ColumnChunkMetaData, + module_type: ModuleType, + row_group_index: usize, + column_index: usize, + ) -> Result<()> { + let column_path_vec = column_metadata.column_path().as_ref(); + + let joined_column_path; + let column_path = if column_path_vec.len() == 1 { + &column_path_vec[0] + } else { + joined_column_path = column_path_vec.join("."); + &joined_column_path + }; + + if file_encryptor.is_column_encrypted(column_path) { + use crate::encryption::encrypt::encrypt_thrift_object; + + let aad = create_module_aad( + file_encryptor.file_aad(), + module_type, + row_group_index, + column_index, + None, + )?; + let mut encryptor = file_encryptor.get_column_encryptor(column_path)?; + encrypt_thrift_object(object, &mut encryptor, &mut sink, &aad) + } else { + Self::write_thrift_object(object, sink) + } + } + + fn get_plaintext_footer_crypto_metadata( + &self, + ) -> (Option, Option>) { + // Only plaintext footers may contain encryption algorithm and footer key metadata. + if let Some(file_encryptor) = self.file_encryptor.as_ref() { + let encryption_properties = file_encryptor.properties(); + if !encryption_properties.encrypt_footer() { + return ( + Some(Self::encryption_algorithm_from_encryptor(file_encryptor)), + encryption_properties.footer_key_metadata().cloned(), + ); + } + } + (None, None) + } + + fn encryption_algorithm_from_encryptor(file_encryptor: &FileEncryptor) -> EncryptionAlgorithm { + let supply_aad_prefix = file_encryptor + .properties() + .aad_prefix() + .map(|_| !file_encryptor.properties().store_aad_prefix()); + let aad_prefix = if file_encryptor.properties().store_aad_prefix() { + file_encryptor.properties().aad_prefix() + } else { + None + }; + EncryptionAlgorithm::AES_GCM_V1(AesGcmV1 { + aad_prefix: aad_prefix.cloned(), + aad_file_unique: Some(file_encryptor.aad_file_unique().clone()), + supply_aad_prefix, + }) + } + + fn file_crypto_metadata(file_encryptor: &'_ FileEncryptor) -> Result> { + let properties = file_encryptor.properties(); + Ok(FileCryptoMetaData { + encryption_algorithm: Self::encryption_algorithm_from_encryptor(file_encryptor), + key_metadata: properties.footer_key_metadata().map(|v| v.as_slice()), + }) + } + + fn encrypt_row_groups( + row_groups: Vec, + file_encryptor: &Arc, + ) -> Result> { + row_groups + .into_iter() + .enumerate() + .map(|(rg_idx, mut rg)| { + let cols: Result> = rg + .columns + .into_iter() + .enumerate() + .map(|(col_idx, c)| { + Self::encrypt_column_chunk(c, file_encryptor, rg_idx, col_idx) + }) + .collect(); + rg.columns = cols?; + Ok(rg) + }) + .collect() + } + + /// Apply column encryption to column chunk metadata + fn encrypt_column_chunk( + mut column_chunk: ColumnChunkMetaData, + file_encryptor: &Arc, + row_group_index: usize, + column_index: usize, + ) -> Result { + // Column crypto metadata should have already been set when the column was created. + // Here we apply the encryption by encrypting the column metadata if required. + let encryptor = match column_chunk.column_crypto_metadata.as_deref() { + None => None, + Some(ColumnCryptoMetaData::ENCRYPTION_WITH_FOOTER_KEY) => { + let is_footer_encrypted = file_encryptor.properties().encrypt_footer(); + + // When uniform encryption is used the footer is already encrypted, + // so the column chunk does not need additional encryption. + // Except if we're in plaintext footer mode, then we need to encrypt + // the column metadata here. + if !is_footer_encrypted { + Some(file_encryptor.get_footer_encryptor()?) + } else { + None + } + } + Some(ColumnCryptoMetaData::ENCRYPTION_WITH_COLUMN_KEY(col_key)) => { + let column_path = col_key.path_in_schema.join("."); + Some(file_encryptor.get_column_encryptor(&column_path)?) + } + }; + + if let Some(mut encryptor) = encryptor { + use crate::file::metadata::thrift::serialize_column_meta_data; + + let aad = create_module_aad( + file_encryptor.file_aad(), + ModuleType::ColumnMetaData, + row_group_index, + column_index, + None, + )?; + // create temp ColumnMetaData that we can encrypt + let mut buffer: Vec = vec![]; + { + let mut prot = ThriftCompactOutputProtocol::new(&mut buffer); + serialize_column_meta_data(&column_chunk, &mut prot)?; + } + let ciphertext = encryptor.encrypt(&buffer, &aad)?; + column_chunk.encrypted_column_metadata = Some(ciphertext); + // Track whether the footer is plaintext, which affects how we serialize + // the column metadata (we need to write stripped metadata for backward compatibility) + column_chunk.plaintext_footer_mode = !file_encryptor.properties().encrypt_footer(); + } + + Ok(column_chunk) + } +} diff --git a/vendor/parquet/src/file/mod.rs b/vendor/parquet/src/file/mod.rs new file mode 100644 index 0000000..b741256 --- /dev/null +++ b/vendor/parquet/src/file/mod.rs @@ -0,0 +1,113 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! APIs for reading parquet data. +//! +//! Provides access to file and row group readers and writers, record API, metadata, etc. +//! +//! # See Also: +//! * [`SerializedFileReader`] and [`SerializedFileWriter`] for reading / writing parquet +//! * [`metadata`]: for working with metadata such as schema +//! * [`statistics`]: for working with statistics in metadata +//! +//! [`SerializedFileReader`]: serialized_reader::SerializedFileReader +//! [`SerializedFileWriter`]: writer::SerializedFileWriter +//! +//! # Example of writing a new file +//! +//! ```rust,no_run +//! use std::{fs, path::Path, sync::Arc}; +//! +//! use parquet::{ +//! file::{ +//! properties::WriterProperties, +//! writer::SerializedFileWriter, +//! }, +//! schema::parser::parse_message_type, +//! }; +//! +//! let path = Path::new("/path/to/sample.parquet"); +//! +//! let message_type = " +//! message schema { +//! REQUIRED INT32 b; +//! } +//! "; +//! let schema = Arc::new(parse_message_type(message_type).unwrap()); +//! let file = fs::File::create(&path).unwrap(); +//! let mut writer = SerializedFileWriter::new(file, schema, Default::default()).unwrap(); +//! let mut row_group_writer = writer.next_row_group().unwrap(); +//! while let Some(mut col_writer) = row_group_writer.next_column().unwrap() { +//! // ... write values to a column writer +//! col_writer.close().unwrap() +//! } +//! row_group_writer.close().unwrap(); +//! writer.close().unwrap(); +//! +//! let bytes = fs::read(&path).unwrap(); +//! assert_eq!(&bytes[0..4], &[b'P', b'A', b'R', b'1']); +//! ``` +//! # Example of reading an existing file +//! +//! ```rust,no_run +//! use parquet::file::reader::{FileReader, SerializedFileReader}; +//! use std::{fs::File, path::Path}; +//! +//! let path = Path::new("/path/to/sample.parquet"); +//! if let Ok(file) = File::open(&path) { +//! let reader = SerializedFileReader::new(file).unwrap(); +//! +//! let parquet_metadata = reader.metadata(); +//! assert_eq!(parquet_metadata.num_row_groups(), 1); +//! +//! let row_group_reader = reader.get_row_group(0).unwrap(); +//! assert_eq!(row_group_reader.num_columns(), 1); +//! } +//! ``` +//! # Example of reading multiple files +//! +//! ```rust,no_run +//! use parquet::file::reader::SerializedFileReader; +//! use std::convert::TryFrom; +//! +//! let paths = vec![ +//! "/path/to/sample.parquet/part-1.snappy.parquet", +//! "/path/to/sample.parquet/part-2.snappy.parquet" +//! ]; +//! // Create a reader for each file and flat map rows +//! let rows = paths.iter() +//! .map(|p| SerializedFileReader::try_from(*p).unwrap()) +//! .flat_map(|r| r.into_iter()); +//! +//! for row in rows { +//! println!("{}", row.unwrap()); +//! } +//! ``` +#[cfg(feature = "encryption")] +pub mod column_crypto_metadata; +pub mod metadata; +pub mod page_index; +pub mod properties; +pub mod reader; +pub mod serialized_reader; +pub mod statistics; +pub mod writer; + +/// The length of the parquet footer in bytes +pub const FOOTER_SIZE: usize = 8; +const PARQUET_MAGIC: [u8; 4] = *b"PAR1"; +const PARQUET_MAGIC_ENCR_FOOTER: [u8; 4] = *b"PARE"; diff --git a/vendor/parquet/src/file/page_index/column_index.rs b/vendor/parquet/src/file/page_index/column_index.rs new file mode 100644 index 0000000..2f90b3d --- /dev/null +++ b/vendor/parquet/src/file/page_index/column_index.rs @@ -0,0 +1,801 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`ColumnIndexMetaData`] structures holding decoded [`ColumnIndex`] information +//! +//! [`ColumnIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md +//! + +use crate::{ + data_type::{ByteArray, FixedLenByteArray}, + errors::{ParquetError, Result}, + parquet_thrift::{ + ElementType, FieldType, ThriftCompactOutputProtocol, WriteThrift, WriteThriftField, + }, +}; +use std::ops::Deref; + +use crate::{ + basic::BoundaryOrder, + data_type::{Int96, private::ParquetValueType}, + file::page_index::index_reader::ThriftColumnIndex, +}; + +/// Common bits of the column index +#[derive(Debug, Clone, PartialEq)] +pub struct ColumnIndex { + pub(crate) null_pages: Vec, + pub(crate) boundary_order: BoundaryOrder, + pub(crate) null_counts: Option>, + pub(crate) repetition_level_histograms: Option>, + pub(crate) definition_level_histograms: Option>, +} + +impl ColumnIndex { + /// Returns the number of pages + pub fn num_pages(&self) -> u64 { + self.null_pages.len() as u64 + } + + /// Returns the number of null values in the page indexed by `idx` + /// + /// Returns `None` if no null counts have been set in the index + pub fn null_count(&self, idx: usize) -> Option { + self.null_counts.as_ref().map(|nc| nc[idx]) + } + + /// Returns the repetition level histogram for the page indexed by `idx` + pub fn repetition_level_histogram(&self, idx: usize) -> Option<&[i64]> { + if let Some(rep_hists) = self.repetition_level_histograms.as_ref() { + let num_lvls = rep_hists.len() / self.num_pages() as usize; + let start = num_lvls * idx; + Some(&rep_hists[start..start + num_lvls]) + } else { + None + } + } + + /// Returns the definition level histogram for the page indexed by `idx` + pub fn definition_level_histogram(&self, idx: usize) -> Option<&[i64]> { + if let Some(def_hists) = self.definition_level_histograms.as_ref() { + let num_lvls = def_hists.len() / self.num_pages() as usize; + let start = num_lvls * idx; + Some(&def_hists[start..start + num_lvls]) + } else { + None + } + } + + /// Returns whether the page indexed by `idx` consists of all null values + pub fn is_null_page(&self, idx: usize) -> bool { + self.null_pages[idx] + } +} + +/// Column index for primitive types +#[derive(Debug, Clone, PartialEq)] +pub struct PrimitiveColumnIndex { + pub(crate) column_index: ColumnIndex, + pub(crate) min_values: Vec, + pub(crate) max_values: Vec, +} + +impl PrimitiveColumnIndex { + pub(crate) fn try_new( + null_pages: Vec, + boundary_order: BoundaryOrder, + null_counts: Option>, + repetition_level_histograms: Option>, + definition_level_histograms: Option>, + min_bytes: Vec<&[u8]>, + max_bytes: Vec<&[u8]>, + ) -> Result { + let len = null_pages.len(); + + if min_bytes.len() != len || max_bytes.len() != len { + return Err(ParquetError::General(format!( + "ColumnIndex min/max length mismatch: expected {len}, got min={} max={}", + min_bytes.len(), + max_bytes.len() + ))); + } + if let Some(ref nc) = null_counts { + if nc.len() != len { + return Err(ParquetError::General(format!( + "ColumnIndex null_counts length mismatch: expected {len}, got {}", + nc.len() + ))); + } + } + if let Some(ref rep) = repetition_level_histograms { + if len != 0 && rep.len() % len != 0 { + return Err(ParquetError::General( + "Invalid repetition_level_histograms length".to_string(), + )); + } + } + if let Some(ref def) = definition_level_histograms { + if len != 0 && def.len() % len != 0 { + return Err(ParquetError::General( + "Invalid definition_level_histograms length".to_string(), + )); + } + } + + let mut min_values = Vec::with_capacity(len); + let mut max_values = Vec::with_capacity(len); + + for (i, is_null) in null_pages.iter().enumerate().take(len) { + if !is_null { + let min = min_bytes[i]; + min_values.push(T::try_from_le_slice(min)?); + + let max = max_bytes[i]; + max_values.push(T::try_from_le_slice(max)?); + } else { + // need placeholders + min_values.push(Default::default()); + max_values.push(Default::default()); + } + } + + Ok(Self { + column_index: ColumnIndex { + null_pages, + boundary_order, + null_counts, + repetition_level_histograms, + definition_level_histograms, + }, + min_values, + max_values, + }) + } + + pub(super) fn try_from_thrift(index: ThriftColumnIndex) -> Result { + Self::try_new( + index.null_pages, + index.boundary_order, + index.null_counts, + index.repetition_level_histograms, + index.definition_level_histograms, + index.min_values, + index.max_values, + ) + } +} + +impl PrimitiveColumnIndex { + /// Returns an array containing the min values for each page. + /// + /// Values in the returned slice are only valid if [`ColumnIndex::is_null_page()`] + /// is `false` for the same index. + pub fn min_values(&self) -> &[T] { + &self.min_values + } + + /// Returns an array containing the max values for each page. + /// + /// Values in the returned slice are only valid if [`ColumnIndex::is_null_page()`] + /// is `false` for the same index. + pub fn max_values(&self) -> &[T] { + &self.max_values + } + + /// Returns an iterator over the min values. + /// + /// Values may be `None` when [`ColumnIndex::is_null_page()`] is `true`. + pub fn min_values_iter(&self) -> impl Iterator> { + self.min_values.iter().enumerate().map(|(i, min)| { + if self.is_null_page(i) { + None + } else { + Some(min) + } + }) + } + + /// Returns an iterator over the max values. + /// + /// Values may be `None` when [`ColumnIndex::is_null_page()`] is `true`. + pub fn max_values_iter(&self) -> impl Iterator> { + self.max_values.iter().enumerate().map(|(i, min)| { + if self.is_null_page(i) { + None + } else { + Some(min) + } + }) + } + + /// Returns the min value for the page indexed by `idx` + /// + /// It is `None` when all values are null + #[inline] + pub fn min_value(&self, idx: usize) -> Option<&T> { + if self.null_pages[idx] { + None + } else { + Some(&self.min_values[idx]) + } + } + + /// Returns the max value for the page indexed by `idx` + /// + /// It is `None` when all values are null + #[inline] + pub fn max_value(&self, idx: usize) -> Option<&T> { + if self.null_pages[idx] { + None + } else { + Some(&self.max_values[idx]) + } + } +} + +impl Deref for PrimitiveColumnIndex { + type Target = ColumnIndex; + + fn deref(&self) -> &Self::Target { + &self.column_index + } +} + +impl WriteThrift for PrimitiveColumnIndex { + const ELEMENT_TYPE: ElementType = ElementType::Struct; + fn write_thrift( + &self, + writer: &mut ThriftCompactOutputProtocol, + ) -> Result<()> { + self.null_pages.write_thrift_field(writer, 1, 0)?; + + // need to handle min/max manually + let len = self.null_pages.len(); + writer.write_field_begin(FieldType::List, 2, 1)?; + writer.write_list_begin(ElementType::Binary, len)?; + for i in 0..len { + let min = self.min_value(i).map(|m| m.as_bytes()).unwrap_or(&[]); + min.write_thrift(writer)?; + } + writer.write_field_begin(FieldType::List, 3, 2)?; + writer.write_list_begin(ElementType::Binary, len)?; + for i in 0..len { + let max = self.max_value(i).map(|m| m.as_bytes()).unwrap_or(&[]); + max.write_thrift(writer)?; + } + let mut last_field_id = self.boundary_order.write_thrift_field(writer, 4, 3)?; + if let Some(null_counts) = &self.null_counts { + last_field_id = null_counts.write_thrift_field(writer, 5, last_field_id)?; + } + if let Some(repetition_level_histograms) = &self.repetition_level_histograms { + last_field_id = + repetition_level_histograms.write_thrift_field(writer, 6, last_field_id)?; + } + if let Some(definition_level_histograms) = &self.definition_level_histograms { + definition_level_histograms.write_thrift_field(writer, 7, last_field_id)?; + } + writer.write_struct_end() + } +} + +/// Column index for byte arrays (fixed length and variable) +#[derive(Debug, Clone, PartialEq)] +pub struct ByteArrayColumnIndex { + pub(crate) column_index: ColumnIndex, + // raw bytes for min and max values + pub(crate) min_bytes: Vec, + pub(crate) min_offsets: Vec, + pub(crate) max_bytes: Vec, + pub(crate) max_offsets: Vec, +} + +impl ByteArrayColumnIndex { + pub(crate) fn try_new( + null_pages: Vec, + boundary_order: BoundaryOrder, + null_counts: Option>, + repetition_level_histograms: Option>, + definition_level_histograms: Option>, + min_values: Vec<&[u8]>, + max_values: Vec<&[u8]>, + ) -> Result { + let len = null_pages.len(); + + if min_values.len() != len || max_values.len() != len { + return Err(ParquetError::General(format!( + "ColumnIndex min/max length mismatch: expected {len}, got min={} max={}", + min_values.len(), + max_values.len() + ))); + } + if let Some(ref nc) = null_counts { + if nc.len() != len { + return Err(ParquetError::General(format!( + "ColumnIndex null_counts length mismatch: expected {len}, got {}", + nc.len() + ))); + } + } + if let Some(ref rep) = repetition_level_histograms { + if len != 0 && rep.len() % len != 0 { + return Err(ParquetError::General( + "Invalid repetition_level_histograms length".to_string(), + )); + } + } + if let Some(ref def) = definition_level_histograms { + if len != 0 && def.len() % len != 0 { + return Err(ParquetError::General( + "Invalid definition_level_histograms length".to_string(), + )); + } + } + + let min_len = min_values.iter().map(|&v| v.len()).sum(); + let max_len = max_values.iter().map(|&v| v.len()).sum(); + let mut min_bytes = vec![0u8; min_len]; + let mut max_bytes = vec![0u8; max_len]; + + let mut min_offsets = vec![0usize; len + 1]; + let mut max_offsets = vec![0usize; len + 1]; + + let mut min_pos = 0; + let mut max_pos = 0; + + for (i, is_null) in null_pages.iter().enumerate().take(len) { + if !is_null { + let min = min_values[i]; + let dst = &mut min_bytes[min_pos..min_pos + min.len()]; + dst.copy_from_slice(min); + min_offsets[i] = min_pos; + min_pos += min.len(); + + let max = max_values[i]; + let dst = &mut max_bytes[max_pos..max_pos + max.len()]; + dst.copy_from_slice(max); + max_offsets[i] = max_pos; + max_pos += max.len(); + } else { + min_offsets[i] = min_pos; + max_offsets[i] = max_pos; + } + } + + min_offsets[len] = min_pos; + max_offsets[len] = max_pos; + + Ok(Self { + column_index: ColumnIndex { + null_pages, + boundary_order, + null_counts, + repetition_level_histograms, + definition_level_histograms, + }, + min_bytes, + min_offsets, + max_bytes, + max_offsets, + }) + } + + pub(super) fn try_from_thrift(index: ThriftColumnIndex) -> Result { + Self::try_new( + index.null_pages, + index.boundary_order, + index.null_counts, + index.repetition_level_histograms, + index.definition_level_histograms, + index.min_values, + index.max_values, + ) + } + + /// Returns the min value for the page indexed by `idx` + /// + /// It is `None` when all values are null + pub fn min_value(&self, idx: usize) -> Option<&[u8]> { + if self.null_pages[idx] { + None + } else { + let start = self.min_offsets[idx]; + let end = self.min_offsets[idx + 1]; + Some(&self.min_bytes[start..end]) + } + } + + /// Returns the max value for the page indexed by `idx` + /// + /// It is `None` when all values are null + pub fn max_value(&self, idx: usize) -> Option<&[u8]> { + if self.null_pages[idx] { + None + } else { + let start = self.max_offsets[idx]; + let end = self.max_offsets[idx + 1]; + Some(&self.max_bytes[start..end]) + } + } + + /// Returns an iterator over the min values. + /// + /// Values may be `None` when [`ColumnIndex::is_null_page()`] is `true`. + pub fn min_values_iter(&self) -> impl Iterator> { + (0..self.num_pages() as usize).map(|i| self.min_value(i)) + } + + /// Returns an iterator over the max values. + /// + /// Values may be `None` when [`ColumnIndex::is_null_page()`] is `true`. + pub fn max_values_iter(&self) -> impl Iterator> { + (0..self.num_pages() as usize).map(|i| self.max_value(i)) + } +} + +impl Deref for ByteArrayColumnIndex { + type Target = ColumnIndex; + + fn deref(&self) -> &Self::Target { + &self.column_index + } +} + +impl WriteThrift for ByteArrayColumnIndex { + const ELEMENT_TYPE: ElementType = ElementType::Struct; + fn write_thrift( + &self, + writer: &mut ThriftCompactOutputProtocol, + ) -> Result<()> { + self.null_pages.write_thrift_field(writer, 1, 0)?; + + // need to handle min/max manually + let len = self.null_pages.len(); + writer.write_field_begin(FieldType::List, 2, 1)?; + writer.write_list_begin(ElementType::Binary, len)?; + for i in 0..len { + let min = self.min_value(i).unwrap_or(&[]); + min.write_thrift(writer)?; + } + writer.write_field_begin(FieldType::List, 3, 2)?; + writer.write_list_begin(ElementType::Binary, len)?; + for i in 0..len { + let max = self.max_value(i).unwrap_or(&[]); + max.write_thrift(writer)?; + } + let mut last_field_id = self.boundary_order.write_thrift_field(writer, 4, 3)?; + if let Some(null_counts) = &self.null_counts { + last_field_id = null_counts.write_thrift_field(writer, 5, last_field_id)?; + } + if let Some(repetition_level_histograms) = &self.repetition_level_histograms { + last_field_id = + repetition_level_histograms.write_thrift_field(writer, 6, last_field_id)?; + } + if let Some(definition_level_histograms) = &self.definition_level_histograms { + definition_level_histograms.write_thrift_field(writer, 7, last_field_id)?; + } + writer.write_struct_end() + } +} + +// Macro to generate getter functions for ColumnIndexMetaData. +macro_rules! colidx_enum_func { + ($self:ident, $func:ident, $arg:ident) => {{ + match *$self { + Self::BOOLEAN(ref typed) => typed.$func($arg), + Self::INT32(ref typed) => typed.$func($arg), + Self::INT64(ref typed) => typed.$func($arg), + Self::INT96(ref typed) => typed.$func($arg), + Self::FLOAT(ref typed) => typed.$func($arg), + Self::DOUBLE(ref typed) => typed.$func($arg), + Self::BYTE_ARRAY(ref typed) => typed.$func($arg), + Self::FIXED_LEN_BYTE_ARRAY(ref typed) => typed.$func($arg), + _ => panic!(concat!( + "Cannot call ", + stringify!($func), + " on ColumnIndexMetaData::NONE" + )), + } + }}; + ($self:ident, $func:ident) => {{ + match *$self { + Self::BOOLEAN(ref typed) => typed.$func(), + Self::INT32(ref typed) => typed.$func(), + Self::INT64(ref typed) => typed.$func(), + Self::INT96(ref typed) => typed.$func(), + Self::FLOAT(ref typed) => typed.$func(), + Self::DOUBLE(ref typed) => typed.$func(), + Self::BYTE_ARRAY(ref typed) => typed.$func(), + Self::FIXED_LEN_BYTE_ARRAY(ref typed) => typed.$func(), + _ => panic!(concat!( + "Cannot call ", + stringify!($func), + " on ColumnIndexMetaData::NONE" + )), + } + }}; +} + +/// Parsed [`ColumnIndex`] information for a Parquet file. +/// +/// See [`ParquetColumnIndex`] for more information. +/// +/// [`ParquetColumnIndex`]: crate::file::metadata::ParquetColumnIndex +/// [`ColumnIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md +#[derive(Debug, Clone, PartialEq)] +#[allow(non_camel_case_types)] +pub enum ColumnIndexMetaData { + /// Sometimes reading page index from parquet file + /// will only return pageLocations without min_max index, + /// `NONE` represents this lack of index information + NONE, + /// Boolean type index + BOOLEAN(PrimitiveColumnIndex), + /// 32-bit integer type index + INT32(PrimitiveColumnIndex), + /// 64-bit integer type index + INT64(PrimitiveColumnIndex), + /// 96-bit integer type (timestamp) index + INT96(PrimitiveColumnIndex), + /// 32-bit floating point type index + FLOAT(PrimitiveColumnIndex), + /// 64-bit floating point type index + DOUBLE(PrimitiveColumnIndex), + /// Byte array type index + BYTE_ARRAY(ByteArrayColumnIndex), + /// Fixed length byte array type index + FIXED_LEN_BYTE_ARRAY(ByteArrayColumnIndex), +} + +impl ColumnIndexMetaData { + /// Return min/max elements inside ColumnIndex are ordered or not. + pub fn is_sorted(&self) -> bool { + // 0:UNORDERED, 1:ASCENDING ,2:DESCENDING, + if let Some(order) = self.get_boundary_order() { + order != BoundaryOrder::UNORDERED + } else { + false + } + } + + /// Get boundary_order of this page index. + pub fn get_boundary_order(&self) -> Option { + match self { + Self::NONE => None, + Self::BOOLEAN(index) => Some(index.boundary_order), + Self::INT32(index) => Some(index.boundary_order), + Self::INT64(index) => Some(index.boundary_order), + Self::INT96(index) => Some(index.boundary_order), + Self::FLOAT(index) => Some(index.boundary_order), + Self::DOUBLE(index) => Some(index.boundary_order), + Self::BYTE_ARRAY(index) => Some(index.boundary_order), + Self::FIXED_LEN_BYTE_ARRAY(index) => Some(index.boundary_order), + } + } + + /// Returns array of null counts, one per page. + /// + /// Returns `None` if now null counts have been set in the index + pub fn null_counts(&self) -> Option<&Vec> { + match self { + Self::NONE => None, + Self::BOOLEAN(index) => index.null_counts.as_ref(), + Self::INT32(index) => index.null_counts.as_ref(), + Self::INT64(index) => index.null_counts.as_ref(), + Self::INT96(index) => index.null_counts.as_ref(), + Self::FLOAT(index) => index.null_counts.as_ref(), + Self::DOUBLE(index) => index.null_counts.as_ref(), + Self::BYTE_ARRAY(index) => index.null_counts.as_ref(), + Self::FIXED_LEN_BYTE_ARRAY(index) => index.null_counts.as_ref(), + } + } + + /// Returns the number of pages + pub fn num_pages(&self) -> u64 { + colidx_enum_func!(self, num_pages) + } + + /// Returns the number of null values in the page indexed by `idx` + /// + /// Returns `None` if no null counts have been set in the index + pub fn null_count(&self, idx: usize) -> Option { + colidx_enum_func!(self, null_count, idx) + } + + /// Returns the repetition level histogram for the page indexed by `idx` + pub fn repetition_level_histogram(&self, idx: usize) -> Option<&[i64]> { + colidx_enum_func!(self, repetition_level_histogram, idx) + } + + /// Returns the definition level histogram for the page indexed by `idx` + pub fn definition_level_histogram(&self, idx: usize) -> Option<&[i64]> { + colidx_enum_func!(self, definition_level_histogram, idx) + } + + /// Returns whether the page indexed by `idx` consists of all null values + #[inline] + pub fn is_null_page(&self, idx: usize) -> bool { + colidx_enum_func!(self, is_null_page, idx) + } +} + +/// Provides iterators over min and max values of a [`ColumnIndexMetaData`] +pub trait ColumnIndexIterators { + /// Can be one of `bool`, `i32`, `i64`, `Int96`, `f32`, `f64`, [`ByteArray`], + /// or [`FixedLenByteArray`] + type Item; + + /// Return iterator over the min values for the index + fn min_values_iter(colidx: &ColumnIndexMetaData) -> impl Iterator>; + + /// Return iterator over the max values for the index + fn max_values_iter(colidx: &ColumnIndexMetaData) -> impl Iterator>; +} + +macro_rules! column_index_iters { + ($item: ident, $variant: ident, $conv:expr) => { + impl ColumnIndexIterators for $item { + type Item = $item; + + fn min_values_iter( + colidx: &ColumnIndexMetaData, + ) -> impl Iterator> { + if let ColumnIndexMetaData::$variant(index) = colidx { + index.min_values_iter().map($conv) + } else { + panic!(concat!("Wrong type for ", stringify!($item), " iterator")) + } + } + + fn max_values_iter( + colidx: &ColumnIndexMetaData, + ) -> impl Iterator> { + if let ColumnIndexMetaData::$variant(index) = colidx { + index.max_values_iter().map($conv) + } else { + panic!(concat!("Wrong type for ", stringify!($item), " iterator")) + } + } + } + }; +} + +column_index_iters!(bool, BOOLEAN, |v| v.copied()); +column_index_iters!(i32, INT32, |v| v.copied()); +column_index_iters!(i64, INT64, |v| v.copied()); +column_index_iters!(Int96, INT96, |v| v.copied()); +column_index_iters!(f32, FLOAT, |v| v.copied()); +column_index_iters!(f64, DOUBLE, |v| v.copied()); +column_index_iters!(ByteArray, BYTE_ARRAY, |v| v + .map(|v| ByteArray::from(v.to_owned()))); +column_index_iters!(FixedLenByteArray, FIXED_LEN_BYTE_ARRAY, |v| v + .map(|v| FixedLenByteArray::from(v.to_owned()))); + +impl WriteThrift for ColumnIndexMetaData { + const ELEMENT_TYPE: ElementType = ElementType::Struct; + + fn write_thrift( + &self, + writer: &mut ThriftCompactOutputProtocol, + ) -> Result<()> { + match self { + ColumnIndexMetaData::BOOLEAN(index) => index.write_thrift(writer), + ColumnIndexMetaData::INT32(index) => index.write_thrift(writer), + ColumnIndexMetaData::INT64(index) => index.write_thrift(writer), + ColumnIndexMetaData::INT96(index) => index.write_thrift(writer), + ColumnIndexMetaData::FLOAT(index) => index.write_thrift(writer), + ColumnIndexMetaData::DOUBLE(index) => index.write_thrift(writer), + ColumnIndexMetaData::BYTE_ARRAY(index) => index.write_thrift(writer), + ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(index) => index.write_thrift(writer), + _ => Err(general_err!("Cannot serialize NONE index")), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_page_index_min_max_null() { + let column_index = PrimitiveColumnIndex { + column_index: ColumnIndex { + null_pages: vec![false], + boundary_order: BoundaryOrder::ASCENDING, + null_counts: Some(vec![0]), + repetition_level_histograms: Some(vec![1, 2]), + definition_level_histograms: Some(vec![1, 2, 3]), + }, + min_values: vec![-123], + max_values: vec![234], + }; + + assert_eq!(column_index.min_value(0), Some(&-123)); + assert_eq!(column_index.max_value(0), Some(&234)); + assert_eq!(column_index.null_count(0), Some(0)); + assert_eq!(column_index.repetition_level_histogram(0).unwrap(), &[1, 2]); + assert_eq!( + column_index.definition_level_histogram(0).unwrap(), + &[1, 2, 3] + ); + } + + #[test] + fn test_page_index_min_max_null_none() { + let column_index: PrimitiveColumnIndex = PrimitiveColumnIndex:: { + column_index: ColumnIndex { + null_pages: vec![true], + boundary_order: BoundaryOrder::ASCENDING, + null_counts: Some(vec![1]), + repetition_level_histograms: None, + definition_level_histograms: Some(vec![1, 0]), + }, + min_values: vec![Default::default()], + max_values: vec![Default::default()], + }; + + assert_eq!(column_index.min_value(0), None); + assert_eq!(column_index.max_value(0), None); + assert_eq!(column_index.null_count(0), Some(1)); + assert_eq!(column_index.repetition_level_histogram(0), None); + assert_eq!(column_index.definition_level_histogram(0).unwrap(), &[1, 0]); + } + + #[test] + fn test_invalid_column_index() { + let column_index = ThriftColumnIndex { + null_pages: vec![true, false], + min_values: vec![ + &[], + &[], // this shouldn't be empty as null_pages[1] is false + ], + max_values: vec![ + &[], + &[], // this shouldn't be empty as null_pages[1] is false + ], + null_counts: None, + repetition_level_histograms: None, + definition_level_histograms: None, + boundary_order: BoundaryOrder::UNORDERED, + }; + + let err = PrimitiveColumnIndex::::try_from_thrift(column_index).unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: error converting value, expected 4 bytes got 0" + ); + } + + #[test] + fn test_column_index_rejects_mismatched_min_max_lengths() { + // Two pages, but only one min/max entry. The entry itself is valid i32 bytes, + // so this specifically checks that lengths must match the number of pages. + let column_index = ThriftColumnIndex { + null_pages: vec![false, false], + min_values: vec![&[1u8, 0, 0, 0]], + max_values: vec![&[10u8, 0, 0, 0]], + null_counts: None, + repetition_level_histograms: None, + definition_level_histograms: None, + boundary_order: BoundaryOrder::UNORDERED, + }; + + // ColumnIndex arrays must align with the number of pages (null_pages.len()). + let err = PrimitiveColumnIndex::::try_from_thrift(column_index).unwrap_err(); + // Should fail because min/max lengths don’t match null_pages + assert!(err.to_string().contains("length mismatch")); + } +} diff --git a/vendor/parquet/src/file/page_index/index_reader.rs b/vendor/parquet/src/file/page_index/index_reader.rs new file mode 100644 index 0000000..f0e40f7 --- /dev/null +++ b/vendor/parquet/src/file/page_index/index_reader.rs @@ -0,0 +1,107 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Support for reading [`ColumnIndexMetaData`] and [`OffsetIndexMetaData`] from parquet metadata. + +use crate::basic::{BoundaryOrder, Type}; +use crate::data_type::Int96; +use crate::errors::{ParquetError, Result}; +use crate::file::page_index::column_index::{ + ByteArrayColumnIndex, ColumnIndexMetaData, PrimitiveColumnIndex, +}; +use crate::file::page_index::offset_index::OffsetIndexMetaData; +use crate::parquet_thrift::{ + ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol, ThriftCompactOutputProtocol, + ThriftSliceInputProtocol, WriteThrift, WriteThriftField, read_thrift_vec, +}; +use crate::thrift_struct; +use std::io::Write; +use std::ops::Range; + +/// Computes the covering range of two optional ranges +/// +/// For example `acc_range(Some(7..9), Some(1..3)) = Some(1..9)` +pub(crate) fn acc_range(a: Option>, b: Option>) -> Option> { + match (a, b) { + (Some(a), Some(b)) => Some(a.start.min(b.start)..a.end.max(b.end)), + (None, x) | (x, None) => x, + } +} + +pub(crate) fn decode_offset_index(data: &[u8]) -> Result { + let mut prot = ThriftSliceInputProtocol::new(data); + + // Try to read fast-path first. If that fails, fall back to slower but more robust + // decoder. + match OffsetIndexMetaData::try_from_fast(&mut prot) { + Ok(offset_index) => Ok(offset_index), + Err(_) => { + prot = ThriftSliceInputProtocol::new(data); + OffsetIndexMetaData::read_thrift(&mut prot) + } + } +} + +// private struct only used for decoding then discarded +thrift_struct!( +pub(super) struct ThriftColumnIndex<'a> { + 1: required list null_pages + 2: required list<'a> min_values + 3: required list<'a> max_values + 4: required BoundaryOrder boundary_order + 5: optional list null_counts + 6: optional list repetition_level_histograms; + 7: optional list definition_level_histograms; +} +); + +pub(crate) fn decode_column_index( + data: &[u8], + column_type: Type, +) -> Result { + let mut prot = ThriftSliceInputProtocol::new(data); + let index = ThriftColumnIndex::read_thrift(&mut prot)?; + + let index = match column_type { + Type::BOOLEAN => { + ColumnIndexMetaData::BOOLEAN(PrimitiveColumnIndex::::try_from_thrift(index)?) + } + Type::INT32 => { + ColumnIndexMetaData::INT32(PrimitiveColumnIndex::::try_from_thrift(index)?) + } + Type::INT64 => { + ColumnIndexMetaData::INT64(PrimitiveColumnIndex::::try_from_thrift(index)?) + } + Type::INT96 => { + ColumnIndexMetaData::INT96(PrimitiveColumnIndex::::try_from_thrift(index)?) + } + Type::FLOAT => { + ColumnIndexMetaData::FLOAT(PrimitiveColumnIndex::::try_from_thrift(index)?) + } + Type::DOUBLE => { + ColumnIndexMetaData::DOUBLE(PrimitiveColumnIndex::::try_from_thrift(index)?) + } + Type::BYTE_ARRAY => { + ColumnIndexMetaData::BYTE_ARRAY(ByteArrayColumnIndex::try_from_thrift(index)?) + } + Type::FIXED_LEN_BYTE_ARRAY => { + ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(ByteArrayColumnIndex::try_from_thrift(index)?) + } + }; + + Ok(index) +} diff --git a/vendor/parquet/src/file/page_index/mod.rs b/vendor/parquet/src/file/page_index/mod.rs new file mode 100644 index 0000000..71b8290 --- /dev/null +++ b/vendor/parquet/src/file/page_index/mod.rs @@ -0,0 +1,24 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Page Index of "[Column Index] Layout to Support Page Skipping" +//! +//! [Column Index]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + +pub mod column_index; +pub mod index_reader; +pub mod offset_index; diff --git a/vendor/parquet/src/file/page_index/offset_index.rs b/vendor/parquet/src/file/page_index/offset_index.rs new file mode 100644 index 0000000..06d21ef --- /dev/null +++ b/vendor/parquet/src/file/page_index/offset_index.rs @@ -0,0 +1,199 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`OffsetIndexMetaData`] structure holding decoded [`OffsetIndex`] information +//! +//! [`OffsetIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + +use std::io::Write; + +use crate::parquet_thrift::{ + ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol, ThriftCompactOutputProtocol, + WriteThrift, WriteThriftField, read_thrift_vec, validate_list_type, +}; +use crate::{ + errors::{ParquetError, Result}, + thrift_struct, +}; + +thrift_struct!( +/// Page location information for [`OffsetIndexMetaData`] +pub struct PageLocation { + /// Offset of the page in the file + 1: required i64 offset + /// Size of the page, including header. Sum of compressed_page_size and header + 2: required i32 compressed_page_size + /// Index within the RowGroup of the first row of the page. When an + /// OffsetIndex is present, pages must begin on row boundaries + /// (repetition_level = 0). + 3: required i64 first_row_index +} +); + +thrift_struct!( +/// [`OffsetIndex`] information for a column chunk. Contains offsets and sizes for each page +/// in the chunk. Optionally stores fully decoded page sizes for BYTE_ARRAY columns. +/// +/// See [`ParquetOffsetIndex`] for more information. +/// +/// [`ParquetOffsetIndex`]: crate::file::metadata::ParquetOffsetIndex +/// [`OffsetIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md +pub struct OffsetIndexMetaData { + /// Vector of [`PageLocation`] objects, one per page in the chunk. + 1: required list page_locations + /// Optional vector of unencoded page sizes, one per page in the chunk. + /// Only defined for BYTE_ARRAY columns. + 2: optional list unencoded_byte_array_data_bytes +} +); + +impl OffsetIndexMetaData { + /// Vector of [`PageLocation`] objects, one per page in the chunk. + pub fn page_locations(&self) -> &Vec { + &self.page_locations + } + + /// Optional vector of unencoded page sizes, one per page in the chunk. Only defined + /// for BYTE_ARRAY columns. + pub fn unencoded_byte_array_data_bytes(&self) -> Option<&Vec> { + self.unencoded_byte_array_data_bytes.as_ref() + } + + // Fast-path read of offset index. This works because we expect all field deltas to be 1, + // and there's no nesting beyond PageLocation, so no need to save the last field id. Like + // read_page_locations(), this will fail if absolute field id's are used. + pub(super) fn try_from_fast<'a, R: ThriftCompactInputProtocol<'a>>( + prot: &mut R, + ) -> Result { + // Offset index is a struct with 2 fields. First field is an array of PageLocations, + // the second an optional array of i64. + + // read field 1 header, then list header, then vec of PageLocations + let (field_type, delta) = prot.read_field_header()?; + if delta != 1 || field_type != FieldType::List as u8 { + return Err(general_err!("error reading OffsetIndex::page_locations")); + } + + // we have to do this manually because we want to use the fast PageLocation decoder + let list_ident = prot.read_list_begin()?; + validate_list_type(ElementType::Struct, &list_ident)?; + let mut page_locations = Vec::with_capacity(list_ident.size as usize); + for _ in 0..list_ident.size { + page_locations.push(read_page_location(prot)?); + } + + let mut unencoded_byte_array_data_bytes: Option> = None; + + // read second field...if it's Stop we're done + let (mut field_type, delta) = prot.read_field_header()?; + if field_type == FieldType::List as u8 { + if delta != 1 { + return Err(general_err!( + "encountered unknown field while reading OffsetIndex" + )); + } + let vec = read_thrift_vec::(&mut *prot)?; + unencoded_byte_array_data_bytes = Some(vec); + + // this one should be Stop + (field_type, _) = prot.read_field_header()?; + } + + if field_type != FieldType::Stop as u8 { + return Err(general_err!( + "encountered unknown field while reading OffsetIndex" + )); + } + + Ok(Self { + page_locations, + unencoded_byte_array_data_bytes, + }) + } +} + +// hand coding this one because it is very time critical + +// Note: this will fail if the fields are either out of order, or if a suboptimal +// encoder doesn't use field deltas. +fn read_page_location<'a, R: ThriftCompactInputProtocol<'a>>(prot: &mut R) -> Result { + // there are 3 fields, all mandatory, so all field deltas should be 1 + let (field_type, delta) = prot.read_field_header()?; + if delta != 1 || field_type != FieldType::I64 as u8 { + return Err(general_err!("error reading PageLocation::offset")); + } + let offset = prot.read_i64()?; + + let (field_type, delta) = prot.read_field_header()?; + if delta != 1 || field_type != FieldType::I32 as u8 { + return Err(general_err!( + "error reading PageLocation::compressed_page_size" + )); + } + let compressed_page_size = prot.read_i32()?; + + let (field_type, delta) = prot.read_field_header()?; + if delta != 1 || field_type != FieldType::I64 as u8 { + return Err(general_err!("error reading PageLocation::first_row_index")); + } + let first_row_index = prot.read_i64()?; + + // read end of struct...return error if there are unknown fields present + let (field_type, _) = prot.read_field_header()?; + if field_type != FieldType::Stop as u8 { + return Err(general_err!("unexpected field in PageLocation")); + } + + Ok(PageLocation { + offset, + compressed_page_size, + first_row_index, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::parquet_thrift::tests::test_roundtrip; + + #[test] + fn test_offset_idx_roundtrip() { + let page_locations = [ + PageLocation { + offset: 0, + compressed_page_size: 10, + first_row_index: 0, + }, + PageLocation { + offset: 10, + compressed_page_size: 20, + first_row_index: 100, + }, + ] + .to_vec(); + let unenc = [0i64, 100i64].to_vec(); + + test_roundtrip(OffsetIndexMetaData { + page_locations: page_locations.clone(), + unencoded_byte_array_data_bytes: Some(unenc), + }); + test_roundtrip(OffsetIndexMetaData { + page_locations, + unencoded_byte_array_data_bytes: None, + }); + } +} diff --git a/vendor/parquet/src/file/properties.rs b/vendor/parquet/src/file/properties.rs new file mode 100644 index 0000000..19191f6 --- /dev/null +++ b/vendor/parquet/src/file/properties.rs @@ -0,0 +1,2459 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Configuration via [`WriterProperties`] and [`ReaderProperties`] +use crate::basic::{Compression, Encoding}; +use crate::compression::{CodecOptions, CodecOptionsBuilder}; +#[cfg(feature = "encryption")] +use crate::encryption::encrypt::FileEncryptionProperties; +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::{KeyValue, SortingColumn}; +use crate::schema::types::ColumnPath; +use std::str::FromStr; +use std::{collections::HashMap, sync::Arc}; + +/// Default value for [`WriterProperties::data_page_size_limit`] +pub const DEFAULT_PAGE_SIZE: usize = 1024 * 1024; +/// Default value for [`WriterProperties::write_batch_size`] +pub const DEFAULT_WRITE_BATCH_SIZE: usize = 1024; +/// Default value for [`WriterProperties::writer_version`] +pub const DEFAULT_WRITER_VERSION: WriterVersion = WriterVersion::PARQUET_1_0; +/// Default value for [`WriterProperties::compression`] +pub const DEFAULT_COMPRESSION: Compression = Compression::UNCOMPRESSED; +/// Default value for [`WriterProperties::dictionary_enabled`] +pub const DEFAULT_DICTIONARY_ENABLED: bool = true; +/// Default value for [`WriterProperties::dictionary_page_size_limit`] +pub const DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT: usize = DEFAULT_PAGE_SIZE; +/// Default value for [`WriterProperties::data_page_row_count_limit`] +pub const DEFAULT_DATA_PAGE_ROW_COUNT_LIMIT: usize = 20_000; +/// Default value for [`WriterProperties::statistics_enabled`] +pub const DEFAULT_STATISTICS_ENABLED: EnabledStatistics = EnabledStatistics::Page; +/// Default value for [`WriterProperties::write_page_header_statistics`] +pub const DEFAULT_WRITE_PAGE_HEADER_STATISTICS: bool = false; +/// Default value for [`WriterProperties::max_row_group_row_count`] +pub const DEFAULT_MAX_ROW_GROUP_ROW_COUNT: usize = 1024 * 1024; +/// Default value for [`WriterProperties::bloom_filter_position`] +pub const DEFAULT_BLOOM_FILTER_POSITION: BloomFilterPosition = BloomFilterPosition::AfterRowGroup; +/// Default value for [`WriterProperties::created_by`] +pub const DEFAULT_CREATED_BY: &str = concat!("parquet-rs version ", env!("CARGO_PKG_VERSION")); +/// Default value for [`WriterProperties::column_index_truncate_length`] +pub const DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH: Option = Some(64); +/// Default value for [`BloomFilterProperties::fpp()`] +pub const DEFAULT_BLOOM_FILTER_FPP: f64 = 0.05; +/// Default value for [`BloomFilterProperties::ndv()`]. +/// +/// Note: this is only the fallback default used when constructing [`BloomFilterProperties`] +/// directly. When using [`WriterPropertiesBuilder`], columns with bloom filters enabled +/// but without an explicit NDV will have their NDV resolved at build time to +/// [`WriterProperties::max_row_group_row_count`], which may differ from this constant +/// if the user configured a custom row group size. +pub const DEFAULT_BLOOM_FILTER_NDV: u64 = DEFAULT_MAX_ROW_GROUP_ROW_COUNT as u64; +/// Default values for [`WriterProperties::statistics_truncate_length`] +pub const DEFAULT_STATISTICS_TRUNCATE_LENGTH: Option = Some(64); +/// Default value for [`WriterProperties::offset_index_disabled`] +pub const DEFAULT_OFFSET_INDEX_DISABLED: bool = false; +/// Default values for [`WriterProperties::coerce_types`] +pub const DEFAULT_COERCE_TYPES: bool = false; +/// Default value for [`WriterProperties::data_page_v2_compression_ratio_threshold`] +pub const DEFAULT_DATA_PAGE_V2_COMPRESSION_RATIO_THRESHOLD: f64 = 1.0; +/// Default value for [`WriterProperties::write_path_in_schema`] +pub const DEFAULT_WRITE_PATH_IN_SCHEMA: bool = true; +/// Default minimum chunk size for content-defined chunking: 256 KiB. +pub const DEFAULT_CDC_MIN_CHUNK_SIZE: usize = 256 * 1024; +/// Default maximum chunk size for content-defined chunking: 1024 KiB. +pub const DEFAULT_CDC_MAX_CHUNK_SIZE: usize = 1024 * 1024; +/// Default normalization level for content-defined chunking. +pub const DEFAULT_CDC_NORM_LEVEL: i32 = 0; + +/// EXPERIMENTAL: Options for content-defined chunking (CDC). +/// +/// Content-defined chunking is an experimental feature that optimizes parquet +/// files for content addressable storage (CAS) systems by writing data pages +/// according to content-defined chunk boundaries. This allows for more +/// efficient deduplication of data across files, hence more efficient network +/// transfers and storage. +/// +/// Each content-defined chunk is written as a separate parquet data page. The +/// following options control the chunks' size and the chunking process. Note +/// that the chunk size is calculated based on the logical value of the data, +/// before any encoding or compression is applied. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct CdcOptions { + /// Minimum chunk size in bytes, default is 256 KiB. + /// The rolling hash will not be updated until this size is reached for each chunk. + /// Note that all data sent through the hash function is counted towards the chunk + /// size, including definition and repetition levels if present. + pub min_chunk_size: usize, + /// Maximum chunk size in bytes, default is 1024 KiB. + /// The chunker will create a new chunk whenever the chunk size exceeds this value. + /// Note that the parquet writer has a related [`data_page_size_limit`] property that + /// controls the maximum size of a parquet data page after encoding. While setting + /// `data_page_size_limit` to a smaller value than `max_chunk_size` doesn't affect + /// the chunking effectiveness, it results in more small parquet data pages. + /// + /// [`data_page_size_limit`]: WriterPropertiesBuilder::set_data_page_size_limit + pub max_chunk_size: usize, + /// Number of bit adjustment to the gearhash mask in order to center the chunk size + /// around the average size more aggressively, default is 0. + /// Increasing the normalization level increases the probability of finding a chunk, + /// improving the deduplication ratio, but also increasing the number of small chunks + /// resulting in many small parquet data pages. The default value provides a good + /// balance between deduplication ratio and fragmentation. + /// Use norm_level=1 or norm_level=2 to reach a higher deduplication ratio at the + /// expense of fragmentation. Negative values can also be used to reduce the + /// probability of finding a chunk, resulting in larger chunks and fewer data pages. + /// Note that values outside [-3, 3] are not recommended, prefer using the default + /// value of 0 for most use cases. + pub norm_level: i32, +} + +impl Default for CdcOptions { + fn default() -> Self { + Self { + min_chunk_size: DEFAULT_CDC_MIN_CHUNK_SIZE, + max_chunk_size: DEFAULT_CDC_MAX_CHUNK_SIZE, + norm_level: DEFAULT_CDC_NORM_LEVEL, + } + } +} + +/// Parquet writer version. +/// +/// Basic constant, which is not part of the Thrift definition. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[allow(non_camel_case_types)] +pub enum WriterVersion { + /// Parquet format version 1.0 + PARQUET_1_0, + /// Parquet format version 2.0 + PARQUET_2_0, +} + +impl WriterVersion { + /// Returns writer version as `i32`. + pub fn as_num(&self) -> i32 { + match self { + WriterVersion::PARQUET_1_0 => 1, + WriterVersion::PARQUET_2_0 => 2, + } + } +} + +impl FromStr for WriterVersion { + type Err = String; + + fn from_str(s: &str) -> Result { + match s { + "PARQUET_1_0" | "parquet_1_0" => Ok(WriterVersion::PARQUET_1_0), + "PARQUET_2_0" | "parquet_2_0" => Ok(WriterVersion::PARQUET_2_0), + _ => Err(format!("Invalid writer version: {s}")), + } + } +} + +/// Where in the file [`ArrowWriter`](crate::arrow::arrow_writer::ArrowWriter) should +/// write Bloom filters +/// +/// Basic constant, which is not part of the Thrift definition. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BloomFilterPosition { + /// Write Bloom Filters of each row group right after the row group + /// + /// This saves memory by writing it as soon as it is computed, at the cost + /// of data locality for readers + AfterRowGroup, + /// Write Bloom Filters at the end of the file + /// + /// This allows better data locality for readers, at the cost of memory usage + /// for writers. + End, +} + +/// Reference counted writer properties. +pub type WriterPropertiesPtr = Arc; + +/// Resolved state of [`WriterPropertiesBuilder::set_offset_index_disabled`]. +/// +/// When a user disables offset indexes but page-level statistics are enabled, +/// the setting is overridden (offset indexes remain enabled). This enum +/// preserves the user's original intent so that a round-trip through +/// `WriterPropertiesBuilder` does not lose it. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum OffsetIndexSetting { + /// Offset indexes are enabled (the default). + Enabled, + /// User disabled offset indexes and no page-level statistics override it. + Disabled, + /// User disabled offset indexes, but page-level statistics require them, + /// so they remain enabled. + DisabledOverridden, +} + +/// Configuration settings for writing parquet files. +/// +/// Use [`Self::builder`] to create a [`WriterPropertiesBuilder`] to change settings. +/// +/// # Example +/// +/// ```rust +/// # use parquet::{ +/// # basic::{Compression, Encoding}, +/// # file::properties::*, +/// # schema::types::ColumnPath, +/// # }; +/// # +/// // Create properties with default configuration. +/// let props = WriterProperties::default(); +/// +/// // Use properties builder to set certain options and assemble the configuration. +/// let props = WriterProperties::builder() +/// .set_writer_version(WriterVersion::PARQUET_1_0) +/// .set_encoding(Encoding::PLAIN) +/// .set_column_encoding(ColumnPath::from("col1"), Encoding::DELTA_BINARY_PACKED) +/// .set_compression(Compression::SNAPPY) +/// .build(); +/// +/// assert_eq!(props.writer_version(), WriterVersion::PARQUET_1_0); +/// assert_eq!( +/// props.encoding(&ColumnPath::from("col1")), +/// Some(Encoding::DELTA_BINARY_PACKED) +/// ); +/// assert_eq!( +/// props.encoding(&ColumnPath::from("col2")), +/// Some(Encoding::PLAIN) +/// ); +/// ``` +#[derive(Debug, Clone)] +pub struct WriterProperties { + data_page_row_count_limit: usize, + write_batch_size: usize, + max_row_group_row_count: Option, + max_row_group_bytes: Option, + bloom_filter_position: BloomFilterPosition, + writer_version: WriterVersion, + created_by: String, + offset_index_setting: OffsetIndexSetting, + pub(crate) key_value_metadata: Option>, + default_column_properties: ColumnProperties, + column_properties: HashMap, + sorting_columns: Option>, + column_index_truncate_length: Option, + statistics_truncate_length: Option, + coerce_types: bool, + content_defined_chunking: Option, + write_path_in_schema: bool, + #[cfg(feature = "encryption")] + pub(crate) file_encryption_properties: Option>, +} + +impl Default for WriterProperties { + fn default() -> Self { + Self::builder().build() + } +} + +impl WriterProperties { + /// Create a new [`WriterProperties`] with the default settings + /// + /// See [`WriterProperties::builder`] for customising settings + pub fn new() -> Self { + Self::default() + } + + /// Returns a new default [`WriterPropertiesBuilder`] for creating writer + /// properties. + pub fn builder() -> WriterPropertiesBuilder { + WriterPropertiesBuilder::default() + } + + /// Converts this [`WriterProperties`] into a [`WriterPropertiesBuilder`] + /// Used for mutating existing property settings + pub fn into_builder(self) -> WriterPropertiesBuilder { + self.into() + } + + /// Returns data page size limit. + /// + /// Note: this is a best effort limit based on the write batch size + /// + /// For more details see [`WriterPropertiesBuilder::set_data_page_size_limit`] + pub fn data_page_size_limit(&self) -> usize { + self.default_column_properties + .data_page_size_limit() + .unwrap_or(DEFAULT_PAGE_SIZE) + } + + /// Returns data page size limit for a specific column. + /// + /// Takes precedence over [`Self::data_page_size_limit`]. + /// + /// Note: this is a best effort limit based on the write batch size. + pub fn column_data_page_size_limit(&self, col: &ColumnPath) -> usize { + self.column_properties + .get(col) + .and_then(|c| c.data_page_size_limit()) + .or_else(|| self.default_column_properties.data_page_size_limit()) + .unwrap_or(DEFAULT_PAGE_SIZE) + } + + /// Returns dictionary page size limit. + /// + /// Note: this is a best effort limit based on the write batch size + /// + /// For more details see [`WriterPropertiesBuilder::set_dictionary_page_size_limit`] + pub fn dictionary_page_size_limit(&self) -> usize { + self.default_column_properties + .dictionary_page_size_limit() + .unwrap_or(DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT) + } + + /// Returns dictionary page size limit for a specific column. + pub fn column_dictionary_page_size_limit(&self, col: &ColumnPath) -> usize { + self.column_properties + .get(col) + .and_then(|c| c.dictionary_page_size_limit()) + .or_else(|| self.default_column_properties.dictionary_page_size_limit()) + .unwrap_or(DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT) + } + + /// Returns the maximum page row count + /// + /// Note: this is a best effort limit based on the write batch size + /// + /// For more details see [`WriterPropertiesBuilder::set_data_page_row_count_limit`] + pub fn data_page_row_count_limit(&self) -> usize { + self.data_page_row_count_limit + } + + /// Returns configured batch size for writes. + /// + /// When writing a batch of data, this setting allows to split it internally into + /// smaller batches so we can better estimate the size of a page currently being + /// written. + /// + /// For more details see [`WriterPropertiesBuilder::set_write_batch_size`] + pub fn write_batch_size(&self) -> usize { + self.write_batch_size + } + + /// Returns maximum number of rows in a row group, or `usize::MAX` if unlimited. + /// + /// For more details see [`WriterPropertiesBuilder::set_max_row_group_size`] + #[deprecated(since = "58.0.0", note = "Use `max_row_group_row_count` instead")] + pub fn max_row_group_size(&self) -> usize { + self.max_row_group_row_count.unwrap_or(usize::MAX) + } + + /// Returns maximum number of rows in a row group, or `None` if unlimited. + /// + /// For more details see [`WriterPropertiesBuilder::set_max_row_group_row_count`] + pub fn max_row_group_row_count(&self) -> Option { + self.max_row_group_row_count + } + + /// Returns maximum size of a row group in bytes, or `None` if unlimited. + /// + /// For more details see [`WriterPropertiesBuilder::set_max_row_group_bytes`] + pub fn max_row_group_bytes(&self) -> Option { + self.max_row_group_bytes + } + + /// Returns bloom filter position. + /// + /// For more details see [`WriterPropertiesBuilder::set_bloom_filter_position`] + pub fn bloom_filter_position(&self) -> BloomFilterPosition { + self.bloom_filter_position + } + + /// Returns configured writer version. + /// + /// For more details see [`WriterPropertiesBuilder::set_writer_version`] + pub fn writer_version(&self) -> WriterVersion { + self.writer_version + } + + /// Returns `created_by` string. + /// + /// For more details see [`WriterPropertiesBuilder::set_created_by`] + pub fn created_by(&self) -> &str { + &self.created_by + } + + /// Returns `true` if offset index writing is disabled. + /// + /// For more details see [`WriterPropertiesBuilder::set_offset_index_disabled`] + pub fn offset_index_disabled(&self) -> bool { + matches!(self.offset_index_setting, OffsetIndexSetting::Disabled) + } + + /// Returns `key_value_metadata` KeyValue pairs. + /// + /// For more details see [`WriterPropertiesBuilder::set_key_value_metadata`] + pub fn key_value_metadata(&self) -> Option<&Vec> { + self.key_value_metadata.as_ref() + } + + /// Returns sorting columns. + /// + /// For more details see [`WriterPropertiesBuilder::set_sorting_columns`] + pub fn sorting_columns(&self) -> Option<&Vec> { + self.sorting_columns.as_ref() + } + + /// Returns the maximum length of truncated min/max values in the column index. + /// + /// `None` if truncation is disabled, must be greater than 0 otherwise. + /// + /// For more details see [`WriterPropertiesBuilder::set_column_index_truncate_length`] + pub fn column_index_truncate_length(&self) -> Option { + self.column_index_truncate_length + } + + /// Returns the maximum length of truncated min/max values in [`Statistics`]. + /// + /// `None` if truncation is disabled, must be greater than 0 otherwise. + /// + /// For more details see [`WriterPropertiesBuilder::set_statistics_truncate_length`] + /// + /// [`Statistics`]: crate::file::statistics::Statistics + pub fn statistics_truncate_length(&self) -> Option { + self.statistics_truncate_length + } + + /// Returns `true` if type coercion is enabled. + /// + /// For more details see [`WriterPropertiesBuilder::set_coerce_types`] + pub fn coerce_types(&self) -> bool { + self.coerce_types + } + + /// Returns `true` if the `path_in_schema` field of the `ColumnMetaData` Thrift struct + /// should be written. + /// + /// For more details see [`WriterPropertiesBuilder::set_write_path_in_schema`] + pub fn write_path_in_schema(&self) -> bool { + self.write_path_in_schema + } + + /// EXPERIMENTAL: Returns content-defined chunking options, or `None` if CDC is disabled. + /// + /// For more details see [`WriterPropertiesBuilder::set_content_defined_chunking`] + pub fn content_defined_chunking(&self) -> Option<&CdcOptions> { + self.content_defined_chunking.as_ref() + } + + /// Returns the compression ratio threshold at or above which a Data Page v2's + /// compressed values are discarded in favor of writing the values uncompressed. + /// + /// For more details see [`WriterPropertiesBuilder::set_data_page_v2_compression_ratio_threshold`] + pub fn data_page_v2_compression_ratio_threshold(&self) -> f64 { + self.default_column_properties + .data_page_v2_compression_ratio_threshold() + .unwrap_or(DEFAULT_DATA_PAGE_V2_COMPRESSION_RATIO_THRESHOLD) + } + + /// Returns the Data Page v2 compression ratio threshold for a specific column. + /// + /// Takes precedence over [`Self::data_page_v2_compression_ratio_threshold`]. + pub fn column_data_page_v2_compression_ratio_threshold(&self, col: &ColumnPath) -> f64 { + self.column_properties + .get(col) + .and_then(|c| c.data_page_v2_compression_ratio_threshold()) + .or_else(|| { + self.default_column_properties + .data_page_v2_compression_ratio_threshold() + }) + .unwrap_or(DEFAULT_DATA_PAGE_V2_COMPRESSION_RATIO_THRESHOLD) + } + + /// Returns encoding for a data page, when dictionary encoding is enabled. + /// + /// This is not configurable. + #[inline] + pub fn dictionary_data_page_encoding(&self) -> Encoding { + // PLAIN_DICTIONARY encoding is deprecated in writer version 1. + // Dictionary values are encoded using RLE_DICTIONARY encoding. + Encoding::RLE_DICTIONARY + } + + /// Returns encoding for dictionary page, when dictionary encoding is enabled. + /// + /// This is not configurable. + #[inline] + pub fn dictionary_page_encoding(&self) -> Encoding { + // PLAIN_DICTIONARY is deprecated in writer version 1. + // Dictionary is encoded using plain encoding. + Encoding::PLAIN + } + + /// Returns encoding for a column, if set. + /// + /// In case when dictionary is enabled, returns fallback encoding. + /// + /// If encoding is not set, then column writer will choose the best encoding + /// based on the column type. + pub fn encoding(&self, col: &ColumnPath) -> Option { + self.column_properties + .get(col) + .and_then(|c| c.encoding()) + .or_else(|| self.default_column_properties.encoding()) + } + + /// Returns compression codec for a column. + /// + /// For more details see [`WriterPropertiesBuilder::set_column_compression`] + pub fn compression(&self, col: &ColumnPath) -> Compression { + self.column_properties + .get(col) + .and_then(|c| c.compression()) + .or_else(|| self.default_column_properties.compression()) + .unwrap_or(DEFAULT_COMPRESSION) + } + + /// Returns `true` if dictionary encoding is enabled for a column. + /// + /// For more details see [`WriterPropertiesBuilder::set_dictionary_enabled`] + pub fn dictionary_enabled(&self, col: &ColumnPath) -> bool { + self.column_properties + .get(col) + .and_then(|c| c.dictionary_enabled()) + .or_else(|| self.default_column_properties.dictionary_enabled()) + .unwrap_or(DEFAULT_DICTIONARY_ENABLED) + } + + /// Returns which statistics are written for a column. + /// + /// For more details see [`WriterPropertiesBuilder::set_statistics_enabled`] + pub fn statistics_enabled(&self, col: &ColumnPath) -> EnabledStatistics { + self.column_properties + .get(col) + .and_then(|c| c.statistics_enabled()) + .or_else(|| self.default_column_properties.statistics_enabled()) + .unwrap_or(DEFAULT_STATISTICS_ENABLED) + } + + /// Returns `true` if [`Statistics`] are to be written to the page header for a column. + /// + /// For more details see [`WriterPropertiesBuilder::set_write_page_header_statistics`] + /// + /// [`Statistics`]: crate::file::statistics::Statistics + pub fn write_page_header_statistics(&self, col: &ColumnPath) -> bool { + self.column_properties + .get(col) + .and_then(|c| c.write_page_header_statistics()) + .or_else(|| { + self.default_column_properties + .write_page_header_statistics() + }) + .unwrap_or(DEFAULT_WRITE_PAGE_HEADER_STATISTICS) + } + + /// Returns the [`BloomFilterProperties`] for the given column + /// + /// Returns `None` if bloom filter is disabled + /// + /// For more details see [`WriterPropertiesBuilder::set_column_bloom_filter_enabled`] + pub fn bloom_filter_properties(&self, col: &ColumnPath) -> Option<&BloomFilterProperties> { + self.column_properties + .get(col) + .and_then(|c| c.bloom_filter_properties()) + .or_else(|| self.default_column_properties.bloom_filter_properties()) + } + + /// Return file encryption properties + /// + /// For more details see [`WriterPropertiesBuilder::with_file_encryption_properties`] + #[cfg(feature = "encryption")] + pub fn file_encryption_properties(&self) -> Option<&Arc> { + self.file_encryption_properties.as_ref() + } +} + +/// Builder for [`WriterProperties`] Parquet writer configuration. +/// +/// See example on [`WriterProperties`] +#[derive(Debug, Clone)] +pub struct WriterPropertiesBuilder { + data_page_row_count_limit: usize, + write_batch_size: usize, + max_row_group_row_count: Option, + max_row_group_bytes: Option, + bloom_filter_position: BloomFilterPosition, + writer_version: WriterVersion, + created_by: String, + offset_index_disabled: bool, + key_value_metadata: Option>, + default_column_properties: ColumnProperties, + column_properties: HashMap, + sorting_columns: Option>, + column_index_truncate_length: Option, + statistics_truncate_length: Option, + coerce_types: bool, + content_defined_chunking: Option, + write_path_in_schema: bool, + #[cfg(feature = "encryption")] + file_encryption_properties: Option>, +} + +impl Default for WriterPropertiesBuilder { + /// Returns default state of the builder. + fn default() -> Self { + Self { + data_page_row_count_limit: DEFAULT_DATA_PAGE_ROW_COUNT_LIMIT, + write_batch_size: DEFAULT_WRITE_BATCH_SIZE, + max_row_group_row_count: Some(DEFAULT_MAX_ROW_GROUP_ROW_COUNT), + max_row_group_bytes: None, + bloom_filter_position: DEFAULT_BLOOM_FILTER_POSITION, + writer_version: DEFAULT_WRITER_VERSION, + created_by: DEFAULT_CREATED_BY.to_string(), + offset_index_disabled: DEFAULT_OFFSET_INDEX_DISABLED, + key_value_metadata: None, + default_column_properties: Default::default(), + column_properties: HashMap::new(), + sorting_columns: None, + column_index_truncate_length: DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH, + statistics_truncate_length: DEFAULT_STATISTICS_TRUNCATE_LENGTH, + coerce_types: DEFAULT_COERCE_TYPES, + content_defined_chunking: None, + write_path_in_schema: DEFAULT_WRITE_PATH_IN_SCHEMA, + #[cfg(feature = "encryption")] + file_encryption_properties: None, + } + } +} + +impl WriterPropertiesBuilder { + /// Finalizes the configuration and returns immutable writer properties struct. + pub fn build(self) -> WriterProperties { + // Pre-compute offset_index_setting + let offset_index_setting = if self.offset_index_disabled { + let default_page_stats_enabled = self.default_column_properties.statistics_enabled() + == Some(EnabledStatistics::Page); + let column_page_stats_enabled = self.column_properties.iter().any(|path_props| { + path_props.1.statistics_enabled() == Some(EnabledStatistics::Page) + }); + if default_page_stats_enabled || column_page_stats_enabled { + OffsetIndexSetting::DisabledOverridden + } else { + OffsetIndexSetting::Disabled + } + } else { + OffsetIndexSetting::Enabled + }; + + // Resolve bloom filter NDV for columns where it wasn't explicitly set: + // default to max_row_group_row_count so the filter is never undersized. + let default_ndv = self + .max_row_group_row_count + .unwrap_or(DEFAULT_MAX_ROW_GROUP_ROW_COUNT) as u64; + let mut default_column_properties = self.default_column_properties; + default_column_properties.resolve_bloom_filter_ndv(default_ndv); + let mut column_properties = self.column_properties; + for props in column_properties.values_mut() { + props.resolve_bloom_filter_ndv(default_ndv); + } + + WriterProperties { + data_page_row_count_limit: self.data_page_row_count_limit, + write_batch_size: self.write_batch_size, + max_row_group_row_count: self.max_row_group_row_count, + max_row_group_bytes: self.max_row_group_bytes, + bloom_filter_position: self.bloom_filter_position, + writer_version: self.writer_version, + created_by: self.created_by, + offset_index_setting, + key_value_metadata: self.key_value_metadata, + default_column_properties, + column_properties, + sorting_columns: self.sorting_columns, + column_index_truncate_length: self.column_index_truncate_length, + statistics_truncate_length: self.statistics_truncate_length, + coerce_types: self.coerce_types, + content_defined_chunking: self.content_defined_chunking, + write_path_in_schema: self.write_path_in_schema, + #[cfg(feature = "encryption")] + file_encryption_properties: self.file_encryption_properties, + } + } + + // ---------------------------------------------------------------------- + // Writer properties related to a file + + /// Sets the `WriterVersion` written into the parquet metadata (defaults to [`PARQUET_1_0`] + /// via [`DEFAULT_WRITER_VERSION`]) + /// + /// This value can determine what features some readers will support. + /// + /// [`PARQUET_1_0`]: [WriterVersion::PARQUET_1_0] + pub fn set_writer_version(mut self, value: WriterVersion) -> Self { + self.writer_version = value; + self + } + + /// Sets best effort maximum number of rows in a data page (defaults to `20_000` + /// via [`DEFAULT_DATA_PAGE_ROW_COUNT_LIMIT`]). + /// + /// The parquet writer will attempt to limit the number of rows in + /// each `DataPage` to this value. Reducing this value will result + /// in larger parquet files, but may improve the effectiveness of + /// page index based predicate pushdown during reading. + /// + /// Note: this is a best effort limit based on value of + /// [`set_write_batch_size`](Self::set_write_batch_size). + pub fn set_data_page_row_count_limit(mut self, value: usize) -> Self { + self.data_page_row_count_limit = value; + self + } + + /// Sets write batch size (defaults to 1024 via [`DEFAULT_WRITE_BATCH_SIZE`]). + /// + /// For performance reasons, data for each column is written in + /// batches of this size. + /// + /// Additional limits such as such as + /// [`set_data_page_row_count_limit`](Self::set_data_page_row_count_limit) + /// are checked between batches, and thus the write batch size value acts as an + /// upper-bound on the enforcement granularity of other limits. + pub fn set_write_batch_size(mut self, value: usize) -> Self { + self.write_batch_size = value; + self + } + + /// Sets maximum number of rows in a row group (defaults to `1024 * 1024` + /// via [`DEFAULT_MAX_ROW_GROUP_ROW_COUNT`]). + /// + /// # Panics + /// If the value is set to 0. + #[deprecated(since = "58.0.0", note = "Use `set_max_row_group_row_count` instead")] + pub fn set_max_row_group_size(mut self, value: usize) -> Self { + assert!(value > 0, "Cannot have a 0 max row group size"); + self.max_row_group_row_count = Some(value); + self + } + + /// Sets maximum number of rows in a row group, or `None` for unlimited. + /// + /// If both `max_row_group_row_count` and `max_row_group_bytes` are set, + /// the row group with the smaller limit will be produced. + /// + /// # Panics + /// If the value is `Some(0)`. + pub fn set_max_row_group_row_count(mut self, value: Option) -> Self { + assert_ne!(value, Some(0), "Cannot have a 0 max row group row count"); + self.max_row_group_row_count = value; + self + } + + /// Sets maximum size of a row group in bytes, or `None` for unlimited. + /// + /// Row groups are flushed when their estimated encoded size exceeds this threshold. + /// This is similar to the official Java implementation for `parquet.block.size`'s behavior. + /// + /// If both `max_row_group_row_count` and `max_row_group_bytes` are set, + /// the row group with the smaller limit will be produced. + /// + /// # Panics + /// If the value is `Some(0)`. + pub fn set_max_row_group_bytes(mut self, value: Option) -> Self { + assert_ne!(value, Some(0), "Cannot have a 0 max row group bytes"); + self.max_row_group_bytes = value; + self + } + + /// Sets where in the final file Bloom Filters are written (defaults to [`AfterRowGroup`] + /// via [`DEFAULT_BLOOM_FILTER_POSITION`]) + /// + /// [`AfterRowGroup`]: BloomFilterPosition::AfterRowGroup + pub fn set_bloom_filter_position(mut self, value: BloomFilterPosition) -> Self { + self.bloom_filter_position = value; + self + } + + /// Sets "created by" property (defaults to `parquet-rs version ` via + /// [`DEFAULT_CREATED_BY`]). + /// + /// This is a string that will be written into the file metadata + pub fn set_created_by(mut self, value: String) -> Self { + self.created_by = value; + self + } + + /// Sets whether the writing of offset indexes is disabled (defaults to `false` via + /// [`DEFAULT_OFFSET_INDEX_DISABLED`]). + /// + /// If statistics level is set to [`Page`] this setting will be overridden with `false`. + /// + /// Note: As the offset indexes are useful for accessing data by row number, + /// they are always written by default, regardless of whether other statistics + /// are enabled. Disabling this metadata may result in a degradation in read + /// performance, so use this option with care. + /// + /// [`Page`]: EnabledStatistics::Page + pub fn set_offset_index_disabled(mut self, value: bool) -> Self { + self.offset_index_disabled = value; + self + } + + /// Sets "key_value_metadata" property (defaults to `None`). + pub fn set_key_value_metadata(mut self, value: Option>) -> Self { + self.key_value_metadata = value; + self + } + + /// Sets sorting order of rows in the row group if any (defaults to `None`). + pub fn set_sorting_columns(mut self, value: Option>) -> Self { + self.sorting_columns = value; + self + } + + /// Sets the max length of min/max value fields when writing the column + /// [`Index`] (defaults to `Some(64)` via [`DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH`]). + /// + /// This can be used to prevent columns with very long values (hundreds of + /// bytes long) from causing the parquet metadata to become huge. + /// + /// # Notes + /// + /// The column [`Index`] is written when [`Self::set_statistics_enabled`] is + /// set to [`EnabledStatistics::Page`]. + /// + /// * If `Some`, must be greater than 0, otherwise will panic + /// * If `None`, there's no effective limit. + /// + /// [`Index`]: crate::file::page_index::column_index::ColumnIndexMetaData + pub fn set_column_index_truncate_length(mut self, max_length: Option) -> Self { + if let Some(value) = max_length { + assert!( + value > 0, + "Cannot have a 0 column index truncate length. If you wish to disable min/max value truncation, set it to `None`." + ); + } + + self.column_index_truncate_length = max_length; + self + } + + /// Sets the max length of min/max value fields in row group and data page header + /// [`Statistics`] (defaults to `Some(64)` via [`DEFAULT_STATISTICS_TRUNCATE_LENGTH`]). + /// + /// # Notes + /// Row group [`Statistics`] are written when [`Self::set_statistics_enabled`] is + /// set to [`EnabledStatistics::Chunk`] or [`EnabledStatistics::Page`]. Data page header + /// [`Statistics`] are written when [`Self::set_statistics_enabled`] is set to + /// [`EnabledStatistics::Page`]. + /// + /// * If `Some`, must be greater than 0, otherwise will panic + /// * If `None`, there's no effective limit. + /// + /// # See also + /// Truncation of Page Index statistics is controlled separately via + /// [`WriterPropertiesBuilder::set_column_index_truncate_length`] + /// + /// [`Statistics`]: crate::file::statistics::Statistics + pub fn set_statistics_truncate_length(mut self, max_length: Option) -> Self { + if let Some(value) = max_length { + assert!( + value > 0, + "Cannot have a 0 statistics truncate length. If you wish to disable min/max value truncation, set it to `None`." + ); + } + + self.statistics_truncate_length = max_length; + self + } + + /// Should the writer coerce types to parquet native types (defaults to `false` via + /// [`DEFAULT_COERCE_TYPES`]). + /// + /// Leaving this option the default `false` will ensure the exact same data + /// written to parquet using this library will be read. + /// + /// Setting this option to `true` will result in parquet files that can be + /// read by more readers, but potentially lose information in the process. + /// + /// * Types such as [`DataType::Date64`], which have no direct corresponding + /// Parquet type, may be stored with lower precision. + /// + /// * The internal field names of `List` and `Map` types will be renamed if + /// necessary to match what is required by the newest Parquet specification. + /// + /// See [`ArrowToParquetSchemaConverter::with_coerce_types`] for more details + /// + /// [`DataType::Date64`]: arrow_schema::DataType::Date64 + /// [`ArrowToParquetSchemaConverter::with_coerce_types`]: crate::arrow::ArrowSchemaConverter::with_coerce_types + pub fn set_coerce_types(mut self, coerce_types: bool) -> Self { + self.coerce_types = coerce_types; + self + } + + /// EXPERIMENTAL: Should the writer emit the `path_in_schema` element of the + /// `ColumnMetaData` Thrift struct. Defaults to `true` via [`DEFAULT_WRITE_PATH_IN_SCHEMA`]. + /// + /// Because `path_in_schema` is a field on the `ColumnMetaData`, it is repeated + /// `num_columns * num_rowgroups` times. Compounding this is any level of nesting or + /// repetition in the schema. For instance, a top-level list column named `foo` will have + /// a `path_in_schema` of `["foo", "list", "element"]`. A list-of-struct is even worse, + /// because the necessary list wrapping is repeated for each element of the struct. A + /// file with a deeply nested schema and many row groups can have a large percentage of the + /// footer taken up by this field. For example, a file of 38 row groups with a schema containing + /// several lists of structs containing lists had 36% of the footer taken up by `path_in_schema`. + /// Removing this redundant information can greatly speed up footer parsing, which is particularly + /// important in scenarios where one does not wish to read the entire file (e.g. point + /// lookups). + /// + ///
+ /// + /// **WARNING:** + /// Setting this to `false` will break compatibility with Parquet readers that + /// still expect this field to be present. Virtually all Parquet readers (parquet-java, + /// Spark, arrow-cpp, pyarrow, pandas to name a few), with the exception + /// of the one in this crate, expect this field to be present, and will terminate execution + /// if it is not. This will continue to be the case unless/until the Parquet format + /// specification is explicitly changed to allow this field to be missing. As a consquence, + /// users should only set this to `false` if they have verified that any reader(s) they plan + /// to use can tolerate the absence of this field. + /// + /// For more context, see [GH-563]. + /// + ///
+ /// + /// [GH-563]: https://github.com/apache/parquet-format/issues/563 + pub fn set_write_path_in_schema(mut self, write_path_in_schema: bool) -> Self { + self.write_path_in_schema = write_path_in_schema; + self + } + + /// EXPERIMENTAL: Sets content-defined chunking options, or disables CDC with `None`. + /// + /// When enabled, data page boundaries are determined by a rolling hash of the + /// column values, so unchanged data produces identical byte sequences across + /// file versions. This enables efficient deduplication on content-addressable + /// storage systems. + /// + /// Only supported through the Arrow writer interface ([`ArrowWriter`]). + /// + /// # Panics + /// + /// Panics if `min_chunk_size == 0` or `max_chunk_size <= min_chunk_size`. + /// + /// [`ArrowWriter`]: crate::arrow::arrow_writer::ArrowWriter + pub fn set_content_defined_chunking(mut self, options: Option) -> Self { + if let Some(ref options) = options { + assert!( + options.min_chunk_size > 0, + "min_chunk_size must be positive" + ); + assert!( + options.max_chunk_size > options.min_chunk_size, + "max_chunk_size ({}) must be greater than min_chunk_size ({})", + options.max_chunk_size, + options.min_chunk_size + ); + } + self.content_defined_chunking = options; + self + } + + /// Sets the default compression ratio threshold at or above which a Data Page + /// v2's compressed values are discarded in favor of writing the values + /// uncompressed, for all columns (defaults to `1.0` via + /// [`DEFAULT_DATA_PAGE_V2_COMPRESSION_RATIO_THRESHOLD`]). + /// + /// When writing a Data Page v2 with a configured compression codec, the writer + /// first compresses the values and then compares the compressed size to the + /// uncompressed size. If `compressed_size >= uncompressed_size * threshold`, the + /// compressed buffer is discarded and the values are written uncompressed for + /// that page (the page's `is_compressed` flag is set to `false`). + /// + /// The default of `1.0` preserves the historical behavior of only keeping + /// compression when it strictly reduces the size. Setting a value below `1.0` + /// requires a minimum amount of size reduction to keep the compressed page — + /// for example `0.9` requires at least a 10% reduction. Setting a value above + /// `1.0` keeps the compressed buffer even if it's somewhat larger than the + /// uncompressed values. + /// + /// This setting only affects Data Page v2; Data Page v1 always stores the + /// compressor's output regardless of the resulting size. + /// + /// # Panics + /// If `value` is not finite or is not strictly positive. + pub fn set_data_page_v2_compression_ratio_threshold(mut self, value: f64) -> Self { + self.default_column_properties + .set_data_page_v2_compression_ratio_threshold(value); + self + } + + /// Sets FileEncryptionProperties (defaults to `None`) + #[cfg(feature = "encryption")] + pub fn with_file_encryption_properties( + mut self, + file_encryption_properties: Arc, + ) -> Self { + self.file_encryption_properties = Some(file_encryption_properties); + self + } + + // ---------------------------------------------------------------------- + // Setters for any column (global) + + /// Sets default encoding for all columns. + /// + /// If dictionary is not enabled, this is treated as a primary encoding for all + /// columns. In case when dictionary is enabled for any column, this value is + /// considered to be a fallback encoding for that column. + /// + /// # Panics + /// + /// if dictionary encoding is specified, regardless of dictionary + /// encoding flag being set. + pub fn set_encoding(mut self, value: Encoding) -> Self { + self.default_column_properties.set_encoding(value); + self + } + + /// Sets default compression codec for all columns (default to [`UNCOMPRESSED`] via + /// [`DEFAULT_COMPRESSION`]). + /// + /// [`UNCOMPRESSED`]: Compression::UNCOMPRESSED + pub fn set_compression(mut self, value: Compression) -> Self { + self.default_column_properties.set_compression(value); + self + } + + /// Sets default flag to enable/disable dictionary encoding for all columns (defaults to `true` + /// via [`DEFAULT_DICTIONARY_ENABLED`]). + /// + /// Use this method to set dictionary encoding, instead of explicitly specifying + /// encoding in `set_encoding` method. + pub fn set_dictionary_enabled(mut self, value: bool) -> Self { + self.default_column_properties.set_dictionary_enabled(value); + self + } + + /// Sets best effort maximum dictionary page size, in bytes (defaults to `1024 * 1024` + /// via [`DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT`]). + /// + /// The parquet writer will attempt to limit the size of each + /// `DataPage` used to store dictionaries to this many + /// bytes. Reducing this value will result in larger parquet + /// files, but may improve the effectiveness of page index based + /// predicate pushdown during reading. + /// + /// Note: this is a best effort limit based on value of + /// [`set_write_batch_size`](Self::set_write_batch_size). + pub fn set_dictionary_page_size_limit(mut self, value: usize) -> Self { + self.default_column_properties + .set_dictionary_page_size_limit(value); + self + } + + /// Sets best effort maximum size of a data page in bytes (defaults to `1024 * 1024` + /// via [`DEFAULT_PAGE_SIZE`]). + /// + /// The parquet writer will attempt to limit the sizes of each + /// `DataPage` to this many bytes. Reducing this value will result + /// in larger parquet files, but may improve the effectiveness of + /// page index based predicate pushdown during reading. + /// + /// Note: this is a best effort limit based on value of + /// [`set_write_batch_size`](Self::set_write_batch_size). + pub fn set_data_page_size_limit(mut self, value: usize) -> Self { + self.default_column_properties + .set_data_page_size_limit(value); + self + } + + /// Sets default [`EnabledStatistics`] level for all columns (defaults to [`Page`] via + /// [`DEFAULT_STATISTICS_ENABLED`]). + /// + /// [`Page`]: EnabledStatistics::Page + pub fn set_statistics_enabled(mut self, value: EnabledStatistics) -> Self { + self.default_column_properties.set_statistics_enabled(value); + self + } + + /// enable/disable writing [`Statistics`] in the page header + /// (defaults to `false` via [`DEFAULT_WRITE_PAGE_HEADER_STATISTICS`]). + /// + /// Only applicable if [`Page`] level statistics are gathered. + /// + /// Setting this value to `true` can greatly increase the size of the resulting Parquet + /// file while yielding very little added benefit. Most modern Parquet implementations + /// will use the min/max values stored in the [`ParquetColumnIndex`] rather than + /// those in the page header. + /// + /// # Note + /// + /// Prior to version 56.0.0, the `parquet` crate always wrote these + /// statistics (the equivalent of setting this option to `true`). This was + /// changed in 56.0.0 to follow the recommendation in the Parquet + /// specification. See [issue #7580] for more details. + /// + /// [`Statistics`]: crate::file::statistics::Statistics + /// [`ParquetColumnIndex`]: crate::file::metadata::ParquetColumnIndex + /// [`Page`]: EnabledStatistics::Page + /// [issue #7580]: https://github.com/apache/arrow-rs/issues/7580 + pub fn set_write_page_header_statistics(mut self, value: bool) -> Self { + self.default_column_properties + .set_write_page_header_statistics(value); + self + } + + /// Sets if bloom filter should be written for all columns (defaults to `false`). + /// + /// # Notes + /// + /// * If the bloom filter is enabled previously then it is a no-op. + /// + /// * If the bloom filter is not enabled, default values for ndv and fpp + /// value are used used. See [`set_bloom_filter_max_ndv`] and + /// [`set_bloom_filter_fpp`] to further adjust the ndv and fpp. + /// + /// [`set_bloom_filter_max_ndv`]: Self::set_bloom_filter_max_ndv + /// [`set_bloom_filter_fpp`]: Self::set_bloom_filter_fpp + pub fn set_bloom_filter_enabled(mut self, value: bool) -> Self { + self.default_column_properties + .set_bloom_filter_enabled(value); + self + } + + /// Sets the default target bloom filter false positive probability (fpp) + /// for all columns (defaults to `0.05` via [`DEFAULT_BLOOM_FILTER_FPP`]). + /// + /// Implicitly enables bloom writing, as if [`set_bloom_filter_enabled`] had + /// been called. + /// + /// [`set_bloom_filter_enabled`]: Self::set_bloom_filter_enabled + pub fn set_bloom_filter_fpp(mut self, value: f64) -> Self { + self.default_column_properties.set_bloom_filter_fpp(value); + self + } + + /// Sets default maximum expected number of distinct values (ndv) for bloom filter + /// for all columns (defaults to [`DEFAULT_BLOOM_FILTER_NDV`]). + /// + /// The bloom filter is initially sized for this many distinct values at the + /// configured FPP, then folded down after all values are inserted to achieve + /// optimal size. A good heuristic is to set this to the expected number of rows + /// in the row group. + /// + /// Implicitly enables bloom writing, as if [`set_bloom_filter_enabled`] had + /// been called. + /// + /// [`set_bloom_filter_enabled`]: Self::set_bloom_filter_enabled + pub fn set_bloom_filter_max_ndv(mut self, value: u64) -> Self { + self.default_column_properties.set_bloom_filter_ndv(value); + self + } + + /// Deprecated alias for [`Self::set_bloom_filter_max_ndv`]. + #[deprecated(since = "59.0.0", note = "Use `set_bloom_filter_max_ndv` instead")] + pub fn set_bloom_filter_ndv(self, value: u64) -> Self { + self.set_bloom_filter_max_ndv(value) + } + + // ---------------------------------------------------------------------- + // Setters for a specific column + + /// Helper method to get existing or new mutable reference of column properties. + #[inline] + fn get_mut_props(&mut self, col: ColumnPath) -> &mut ColumnProperties { + self.column_properties.entry(col).or_default() + } + + /// Sets encoding for a specific column. + /// + /// Takes precedence over [`Self::set_encoding`]. + /// + /// If dictionary is not enabled, this is treated as a primary encoding for this + /// column. In case when dictionary is enabled for this column, either through + /// global defaults or explicitly, this value is considered to be a fallback + /// encoding for this column. + /// + /// # Panics + /// If user tries to set dictionary encoding here, regardless of dictionary + /// encoding flag being set. + pub fn set_column_encoding(mut self, col: ColumnPath, value: Encoding) -> Self { + self.get_mut_props(col).set_encoding(value); + self + } + + /// Sets compression codec for a specific column. + /// + /// Takes precedence over [`Self::set_compression`]. + pub fn set_column_compression(mut self, col: ColumnPath, value: Compression) -> Self { + self.get_mut_props(col).set_compression(value); + self + } + + /// Sets flag to enable/disable dictionary encoding for a specific column. + /// + /// Takes precedence over [`Self::set_dictionary_enabled`]. + pub fn set_column_dictionary_enabled(mut self, col: ColumnPath, value: bool) -> Self { + self.get_mut_props(col).set_dictionary_enabled(value); + self + } + + /// Sets dictionary page size limit for a specific column. + /// + /// Takes precedence over [`Self::set_dictionary_page_size_limit`]. + pub fn set_column_dictionary_page_size_limit(mut self, col: ColumnPath, value: usize) -> Self { + self.get_mut_props(col) + .set_dictionary_page_size_limit(value); + self + } + + /// Sets data page size limit for a specific column. + /// + /// Takes precedence over [`Self::set_data_page_size_limit`]. + pub fn set_column_data_page_size_limit(mut self, col: ColumnPath, value: usize) -> Self { + self.get_mut_props(col).set_data_page_size_limit(value); + self + } + + /// Sets [`EnabledStatistics`] level for a specific column. + /// + /// Takes precedence over [`Self::set_statistics_enabled`]. + pub fn set_column_statistics_enabled( + mut self, + col: ColumnPath, + value: EnabledStatistics, + ) -> Self { + self.get_mut_props(col).set_statistics_enabled(value); + self + } + + /// Sets whether to write [`Statistics`] in the page header for a specific column. + /// + /// Takes precedence over [`Self::set_write_page_header_statistics`]. + /// + /// [`Statistics`]: crate::file::statistics::Statistics + pub fn set_column_write_page_header_statistics(mut self, col: ColumnPath, value: bool) -> Self { + self.get_mut_props(col) + .set_write_page_header_statistics(value); + self + } + + /// Sets whether a bloom filter should be written for a specific column. + /// + /// Takes precedence over [`Self::set_bloom_filter_enabled`]. + pub fn set_column_bloom_filter_enabled(mut self, col: ColumnPath, value: bool) -> Self { + self.get_mut_props(col).set_bloom_filter_enabled(value); + self + } + + /// Sets the false positive probability for bloom filter for a specific column. + /// + /// Takes precedence over [`Self::set_bloom_filter_fpp`]. + pub fn set_column_bloom_filter_fpp(mut self, col: ColumnPath, value: f64) -> Self { + self.get_mut_props(col).set_bloom_filter_fpp(value); + self + } + + /// Sets the maximum expected number of distinct values for bloom filter for + /// a specific column. + /// + /// Takes precedence over [`Self::set_bloom_filter_max_ndv`]. + pub fn set_column_bloom_filter_max_ndv(mut self, col: ColumnPath, value: u64) -> Self { + self.get_mut_props(col).set_bloom_filter_ndv(value); + self + } + + /// Sets the Data Page v2 compression ratio threshold for a specific column. + /// + /// Takes precedence over [`Self::set_data_page_v2_compression_ratio_threshold`]. + /// + /// # Panics + /// If `value` is not finite or is not strictly positive. + pub fn set_column_data_page_v2_compression_ratio_threshold( + mut self, + col: ColumnPath, + value: f64, + ) -> Self { + self.get_mut_props(col) + .set_data_page_v2_compression_ratio_threshold(value); + self + } + + /// Deprecated alias for [`Self::set_column_bloom_filter_max_ndv`]. + #[deprecated( + since = "59.0.0", + note = "Use `set_column_bloom_filter_max_ndv` instead" + )] + pub fn set_column_bloom_filter_ndv(self, col: ColumnPath, value: u64) -> Self { + self.set_column_bloom_filter_max_ndv(col, value) + } + + /// Sets the [`BloomFilterProperties`] for all columns, implicitly enabling + /// the bloom filter. + /// + /// Both `fpp` and `ndv` from `value` are treated as explicit and will not + /// be overridden by the build-time row-group-size NDV fallback. For + /// dynamic NDV sizing (resolved to `max_row_group_row_count` at build + /// time), use [`Self::set_bloom_filter_enabled`] or + /// [`Self::set_bloom_filter_fpp`] instead. + pub fn set_bloom_filter_properties(mut self, value: BloomFilterProperties) -> Self { + self.default_column_properties + .set_bloom_filter_properties(value); + self + } + + /// Sets the [`BloomFilterProperties`] for a specific column. + /// + /// Takes precedence over [`Self::set_bloom_filter_properties`]. + pub fn set_column_bloom_filter_properties( + mut self, + col: ColumnPath, + value: BloomFilterProperties, + ) -> Self { + self.get_mut_props(col).set_bloom_filter_properties(value); + self + } +} + +impl From for WriterPropertiesBuilder { + fn from(props: WriterProperties) -> Self { + WriterPropertiesBuilder { + data_page_row_count_limit: props.data_page_row_count_limit, + write_batch_size: props.write_batch_size, + max_row_group_row_count: props.max_row_group_row_count, + max_row_group_bytes: props.max_row_group_bytes, + bloom_filter_position: props.bloom_filter_position, + writer_version: props.writer_version, + created_by: props.created_by, + offset_index_disabled: !matches!( + props.offset_index_setting, + OffsetIndexSetting::Enabled + ), + key_value_metadata: props.key_value_metadata, + default_column_properties: props.default_column_properties, + column_properties: props.column_properties, + sorting_columns: props.sorting_columns, + column_index_truncate_length: props.column_index_truncate_length, + statistics_truncate_length: props.statistics_truncate_length, + coerce_types: props.coerce_types, + content_defined_chunking: props.content_defined_chunking, + write_path_in_schema: props.write_path_in_schema, + #[cfg(feature = "encryption")] + file_encryption_properties: props.file_encryption_properties, + } + } +} + +/// Controls the level of statistics to be computed by the writer and stored in +/// the parquet file. +/// +/// Enabling statistics makes the resulting Parquet file larger and requires +/// more time to read the parquet footer. +/// +/// Statistics can be used to improve query performance by pruning row groups +/// and pages during query execution if the query engine supports evaluating the +/// predicate using the statistics. +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +pub enum EnabledStatistics { + /// Compute no statistics. + None, + /// Compute column chunk-level statistics but not page-level. + /// + /// Setting this option will store one set of statistics for each relevant + /// column for each row group. The more row groups written, the more + /// statistics will be stored. + Chunk, + /// Compute page-level and column chunk-level statistics. + /// + /// Setting this option will store one set of statistics for each relevant + /// column for each row group. In addition, this will enable the writing + /// of the column index (the offset index is always written regardless of + /// this setting). See [`ParquetColumnIndex`] for + /// more information. + /// + /// [`ParquetColumnIndex`]: crate::file::metadata::ParquetColumnIndex + Page, +} + +impl FromStr for EnabledStatistics { + type Err = String; + + fn from_str(s: &str) -> Result { + match s { + "NONE" | "none" => Ok(EnabledStatistics::None), + "CHUNK" | "chunk" => Ok(EnabledStatistics::Chunk), + "PAGE" | "page" => Ok(EnabledStatistics::Page), + _ => Err(format!("Invalid statistics arg: {s}")), + } + } +} + +impl Default for EnabledStatistics { + fn default() -> Self { + DEFAULT_STATISTICS_ENABLED + } +} + +/// Controls the bloom filter to be computed by the writer. +/// +/// The bloom filter is initially sized for `ndv` distinct values at the given `fpp`, then +/// automatically folded down after all values are inserted to achieve optimal size while +/// maintaining the target `fpp`. See [`Sbbf::fold_to_target_fpp`] for details on the +/// folding algorithm. +/// +/// # Example +/// +/// ```rust +/// # use parquet::{ +/// # file::properties::{BloomFilterProperties, WriterProperties}, +/// # schema::types::ColumnPath, +/// # }; +/// // Build a BloomFilterProperties via the builder, then apply it to one column. +/// let bf = BloomFilterProperties::builder() +/// .with_fpp(0.01) +/// .with_max_ndv(10_000) +/// .build(); +/// +/// let props = WriterProperties::builder() +/// .set_column_bloom_filter_properties(ColumnPath::from("user_id"), bf.clone()) +/// .build(); +/// +/// assert_eq!( +/// props.bloom_filter_properties(&ColumnPath::from("user_id")), +/// Some(&bf) +/// ); +/// ``` +/// +/// [`Sbbf::fold_to_target_fpp`]: crate::bloom_filter::Sbbf::fold_to_target_fpp +#[derive(Debug, Clone, PartialEq)] +pub struct BloomFilterProperties { + fpp: f64, + ndv: u64, +} + +impl Default for BloomFilterProperties { + fn default() -> Self { + BloomFilterProperties { + fpp: DEFAULT_BLOOM_FILTER_FPP, + ndv: DEFAULT_BLOOM_FILTER_NDV, + } + } +} + +impl BloomFilterProperties { + /// Returns a new [`BloomFilterPropertiesBuilder`] for constructing + /// [`BloomFilterProperties`] with custom values. + pub fn builder() -> BloomFilterPropertiesBuilder { + BloomFilterPropertiesBuilder::new() + } + + /// False positive probability. This should be always between 0 and 1 exclusive. Defaults to [`DEFAULT_BLOOM_FILTER_FPP`]. + /// + /// You should set this value by calling [`WriterPropertiesBuilder::set_bloom_filter_fpp`]. + /// + /// The bloom filter data structure is a trade of between disk and memory space versus fpp, the + /// smaller the fpp, the more memory and disk space is required, thus setting it to a reasonable value + /// e.g. 0.1, 0.05, or 0.001 is recommended. + /// + /// This value also serves as the target FPP for bloom filter folding: after all values + /// are inserted, the filter is folded down to the smallest size that still meets this FPP. + pub fn fpp(&self) -> f64 { + self.fpp + } + + /// Maximum expected number of distinct values. Defaults to [`DEFAULT_BLOOM_FILTER_NDV`]. + /// + /// You should set this value by calling [`WriterPropertiesBuilder::set_bloom_filter_max_ndv`]. + /// + /// When not explicitly set via the builder, this defaults to + /// [`max_row_group_row_count`](WriterProperties::max_row_group_row_count) (resolved at + /// build time). The bloom filter is initially sized for this many distinct values at the + /// given `fpp`, then folded down after insertion to achieve optimal size. A good heuristic + /// is to set this to the expected number of rows in the row group. If fewer distinct values + /// are actually written, the filter will be automatically compacted via folding. + /// + /// Thus the only negative side of overestimating this value is that the bloom filter + /// will use more memory during writing than necessary, but it will not affect the final + /// bloom filter size on disk. + /// + /// If you wish to reduce memory usage during writing and are able to make a reasonable estimate + /// of the number of distinct values in a row group, it is recommended to set this value explicitly + /// rather than relying on the default dynamic sizing based on `max_row_group_row_count`. + /// If you do set this value explicitly it is probably best to set it for each column + /// individually via [`WriterPropertiesBuilder::set_column_bloom_filter_max_ndv`] rather than globally, + /// since different columns may have different numbers of distinct values. + pub fn ndv(&self) -> u64 { + self.ndv + } +} + +/// Builder for [`BloomFilterProperties`]. +/// +/// Use [`BloomFilterProperties::builder`] or [`BloomFilterPropertiesBuilder::new`] +/// as the entry point. +#[derive(Debug, Clone, Default)] +pub struct BloomFilterPropertiesBuilder { + fpp: Option, + ndv: Option, +} + +impl BloomFilterPropertiesBuilder { + /// Returns a new builder with no fields set. + /// + /// Equivalent to [`BloomFilterProperties::builder`]. + pub fn new() -> Self { + Self::default() + } + + /// Sets the target false positive probability. + /// + /// The value must be in `(0.0, 1.0)` exclusively; this is validated at + /// build time by [`Self::build`] / [`Self::try_build`]. When unset, the + /// default is `0.05` (5%, see [`DEFAULT_BLOOM_FILTER_FPP`]). + pub fn with_fpp(mut self, fpp: f64) -> Self { + self.fpp = Some(fpp); + self + } + + /// Sets the maximum expected number of distinct values used to size the + /// bloom filter before folding. + /// + /// When unset, the default is `1_048_576` (see [`DEFAULT_BLOOM_FILTER_NDV`]), + /// which at the default fpp of 5% reserves roughly 1 MiB per column for the + /// filter bitset, derived as follows: + /// + /// ```text + /// ndv = 1,048,576, fpp = 0.05 + /// 0.05^(1/8) ≈ 0.6877 + /// 1 - 0.6877 ≈ 0.3123 + /// ln(0.3123) ≈ -1.164 + /// num_bits = -8 * 1,048,576 / -1.164 ≈ 7,206,000 bits + /// ≈ 900,750 bytes (~900 KB) + /// next_power_of_two(900 KB) = 1 MiB (= 1,048,576 bytes) + /// ``` + pub fn with_max_ndv(mut self, ndv: u64) -> Self { + self.ndv = Some(ndv); + self + } + + /// Builds [`BloomFilterProperties`]. + /// + /// Panics if the configured `fpp` is not in `(0.0, 1.0)` exclusive. + /// Use [`Self::try_build`] for a non-panicking alternative. + pub fn build(self) -> BloomFilterProperties { + self.try_build().unwrap_or_else(|e| panic!("{e}")) + } + + /// Builds [`BloomFilterProperties`], returning an error instead of + /// panicking when the configured `fpp` is not in `(0.0, 1.0)` exclusive. + pub fn try_build(self) -> Result { + let fpp = self.fpp.unwrap_or(DEFAULT_BLOOM_FILTER_FPP); + validate_bloom_filter_fpp(fpp).map_err(ParquetError::General)?; + let ndv = self.ndv.unwrap_or(DEFAULT_BLOOM_FILTER_NDV); + Ok(BloomFilterProperties { fpp, ndv }) + } +} + +/// Single source of truth for the bloom filter fpp range check, shared by +/// [`ColumnProperties::set_bloom_filter_fpp`] (panic path) and +/// [`BloomFilterPropertiesBuilder::try_build`] (Result path). +fn validate_bloom_filter_fpp(fpp: f64) -> std::result::Result<(), String> { + if !(fpp > 0.0 && fpp < 1.0) { + return Err(format!( + "fpp must be between 0.0 and 1.0 exclusive, got {fpp}" + )); + } + Ok(()) +} + +/// Container for column properties that can be changed as part of writer. +/// +/// If a field is `None`, it means that no specific value has been set for this column, +/// so some subsequent or default value must be used. +#[derive(Debug, Clone, Default, PartialEq)] +struct ColumnProperties { + encoding: Option, + codec: Option, + data_page_size_limit: Option, + dictionary_page_size_limit: Option, + dictionary_enabled: Option, + statistics_enabled: Option, + write_page_header_statistics: Option, + /// bloom filter related properties + bloom_filter_properties: Option, + /// Whether the bloom filter NDV was explicitly set by the user + bloom_filter_ndv_is_set: bool, + data_page_v2_compression_ratio_threshold: Option, +} + +impl ColumnProperties { + /// Sets encoding for this column. + /// + /// If dictionary is not enabled, this is treated as a primary encoding for a column. + /// In case when dictionary is enabled for a column, this value is considered to + /// be a fallback encoding. + /// + /// Panics if user tries to set dictionary encoding here, regardless of dictionary + /// encoding flag being set. Use `set_dictionary_enabled` method to enable dictionary + /// for a column. + fn set_encoding(&mut self, value: Encoding) { + if value == Encoding::PLAIN_DICTIONARY || value == Encoding::RLE_DICTIONARY { + panic!("Dictionary encoding can not be used as fallback encoding"); + } + self.encoding = Some(value); + } + + /// Sets compression codec for this column. + fn set_compression(&mut self, value: Compression) { + self.codec = Some(value); + } + + /// Sets data page size limit for this column. + fn set_data_page_size_limit(&mut self, value: usize) { + self.data_page_size_limit = Some(value); + } + + /// Sets whether dictionary encoding is enabled for this column. + fn set_dictionary_enabled(&mut self, enabled: bool) { + self.dictionary_enabled = Some(enabled); + } + + /// Sets dictionary page size limit for this column. + fn set_dictionary_page_size_limit(&mut self, value: usize) { + self.dictionary_page_size_limit = Some(value); + } + + /// Sets the statistics level for this column. + fn set_statistics_enabled(&mut self, enabled: EnabledStatistics) { + self.statistics_enabled = Some(enabled); + } + + /// Sets whether to write statistics in the page header for this column. + fn set_write_page_header_statistics(&mut self, enabled: bool) { + self.write_page_header_statistics = Some(enabled); + } + + /// If `value` is `true`, sets bloom filter properties to default values if not previously set, + /// otherwise it is a no-op. + /// If `value` is `false`, resets bloom filter properties to `None`. + fn set_bloom_filter_enabled(&mut self, value: bool) { + if value && self.bloom_filter_properties.is_none() { + self.bloom_filter_properties = Some(Default::default()) + } else if !value { + self.bloom_filter_properties = None + } + } + + /// Sets the false positive probability for bloom filter for this column, and implicitly enables + /// bloom filter if not previously enabled. + /// + /// # Panics + /// + /// Panics if the `value` is not between 0 and 1 exclusive + fn set_bloom_filter_fpp(&mut self, value: f64) { + if let Err(msg) = validate_bloom_filter_fpp(value) { + panic!("{msg}"); + } + self.bloom_filter_properties + .get_or_insert_with(Default::default) + .fpp = value; + } + + /// Sets the maximum expected number of distinct (unique) values for bloom filter for this + /// column, and implicitly enables bloom filter if not previously enabled. + fn set_bloom_filter_ndv(&mut self, value: u64) { + self.bloom_filter_properties + .get_or_insert_with(Default::default) + .ndv = value; + self.bloom_filter_ndv_is_set = true; + } + + /// Sets the bloom filter properties for this column from a fully-built + /// [`BloomFilterProperties`], implicitly enabling the bloom filter. + /// + /// Both `fpp` and `ndv` from `value` are treated as explicit, so the + /// build-time row-group-size NDV fallback in + /// [`WriterPropertiesBuilder::build`] will not override them. + fn set_bloom_filter_properties(&mut self, value: BloomFilterProperties) { + self.bloom_filter_properties = Some(value); + self.bloom_filter_ndv_is_set = true; + } + + /// Sets the Data Page v2 compression ratio threshold for this column. + /// + /// # Panics + /// If `value` is not finite or is not strictly positive. + fn set_data_page_v2_compression_ratio_threshold(&mut self, value: f64) { + assert!( + value.is_finite() && value > 0.0, + "data_page_v2_compression_ratio_threshold must be a positive finite number, got {value}" + ); + self.data_page_v2_compression_ratio_threshold = Some(value); + } + + /// Returns optional encoding for this column. + fn encoding(&self) -> Option { + self.encoding + } + + /// Returns optional compression codec for this column. + fn compression(&self) -> Option { + self.codec + } + + /// Returns `Some(true)` if dictionary encoding is enabled for this column, if + /// disabled then returns `Some(false)`. If result is `None`, then no setting has + /// been provided. + fn dictionary_enabled(&self) -> Option { + self.dictionary_enabled + } + + /// Returns optional dictionary page size limit for this column. + fn dictionary_page_size_limit(&self) -> Option { + self.dictionary_page_size_limit + } + + /// Returns optional data page size limit for this column. + fn data_page_size_limit(&self) -> Option { + self.data_page_size_limit + } + + /// Returns optional statistics level requested for this column. If result is `None`, + /// then no setting has been provided. + fn statistics_enabled(&self) -> Option { + self.statistics_enabled + } + + /// Returns `Some(true)` if [`Statistics`] are to be written to the page header for this + /// column. + /// + /// [`Statistics`]: crate::file::statistics::Statistics + fn write_page_header_statistics(&self) -> Option { + self.write_page_header_statistics + } + + /// Returns the bloom filter properties, or `None` if not enabled + fn bloom_filter_properties(&self) -> Option<&BloomFilterProperties> { + self.bloom_filter_properties.as_ref() + } + + /// Returns optional Data Page v2 compression ratio threshold for this column. + fn data_page_v2_compression_ratio_threshold(&self) -> Option { + self.data_page_v2_compression_ratio_threshold + } + + /// If bloom filter is enabled and NDV was not explicitly set, resolve it to the + /// given `default_ndv` (typically derived from `max_row_group_row_count`). + fn resolve_bloom_filter_ndv(&mut self, default_ndv: u64) { + if !self.bloom_filter_ndv_is_set { + if let Some(ref mut bf) = self.bloom_filter_properties { + bf.ndv = default_ndv; + } + } + } +} + +/// Reference counted reader properties. +pub type ReaderPropertiesPtr = Arc; + +const DEFAULT_READ_BLOOM_FILTER: bool = false; +const DEFAULT_READ_PAGE_STATS: bool = false; + +/// Configuration settings for reading parquet files. +/// +/// All properties are immutable and `Send` + `Sync`. +/// Use [`ReaderPropertiesBuilder`] to assemble these properties. +/// +/// # Example +/// +/// ```rust +/// use parquet::file::properties::ReaderProperties; +/// +/// // Create properties with default configuration. +/// let props = ReaderProperties::builder().build(); +/// +/// // Use properties builder to set certain options and assemble the configuration. +/// let props = ReaderProperties::builder() +/// .set_backward_compatible_lz4(false) +/// .build(); +/// ``` +pub struct ReaderProperties { + codec_options: CodecOptions, + read_bloom_filter: bool, + read_page_stats: bool, +} + +impl ReaderProperties { + /// Returns builder for reader properties with default values. + pub fn builder() -> ReaderPropertiesBuilder { + ReaderPropertiesBuilder::with_defaults() + } + + /// Returns codec options. + pub(crate) fn codec_options(&self) -> &CodecOptions { + &self.codec_options + } + + /// Returns whether to read bloom filter + pub(crate) fn read_bloom_filter(&self) -> bool { + self.read_bloom_filter + } + + /// Returns whether to read page level statistics + pub(crate) fn read_page_stats(&self) -> bool { + self.read_page_stats + } +} + +/// Builder for parquet file reader configuration. See example on +/// [`ReaderProperties`] +pub struct ReaderPropertiesBuilder { + codec_options_builder: CodecOptionsBuilder, + read_bloom_filter: Option, + read_page_stats: Option, +} + +/// Reader properties builder. +impl ReaderPropertiesBuilder { + /// Returns default state of the builder. + fn with_defaults() -> Self { + Self { + codec_options_builder: CodecOptionsBuilder::default(), + read_bloom_filter: None, + read_page_stats: None, + } + } + + /// Finalizes the configuration and returns immutable reader properties struct. + pub fn build(self) -> ReaderProperties { + ReaderProperties { + codec_options: self.codec_options_builder.build(), + read_bloom_filter: self.read_bloom_filter.unwrap_or(DEFAULT_READ_BLOOM_FILTER), + read_page_stats: self.read_page_stats.unwrap_or(DEFAULT_READ_PAGE_STATS), + } + } + + /// Enable/disable backward compatible LZ4. + /// + /// If backward compatible LZ4 is enable, on LZ4_HADOOP error it will fallback + /// to the older versions LZ4 algorithms. That is LZ4_FRAME, for backward compatibility + /// with files generated by older versions of this library, and LZ4_RAW, for backward + /// compatibility with files generated by older versions of parquet-cpp. + /// + /// If backward compatible LZ4 is disabled, on LZ4_HADOOP error it will return the error. + pub fn set_backward_compatible_lz4(mut self, value: bool) -> Self { + self.codec_options_builder = self + .codec_options_builder + .set_backward_compatible_lz4(value); + self + } + + /// Enable/disable reading bloom filter + /// + /// If reading bloom filter is enabled, bloom filter will be read from the file. + /// If reading bloom filter is disabled, bloom filter will not be read from the file. + /// + /// By default bloom filter is set to be read. + pub fn set_read_bloom_filter(mut self, value: bool) -> Self { + self.read_bloom_filter = Some(value); + self + } + + /// Enable/disable reading page-level statistics + /// + /// If set to `true`, then the reader will decode and populate the [`Statistics`] for + /// each page, if present. + /// If set to `false`, then the reader will skip decoding the statistics. + /// + /// By default statistics will not be decoded. + /// + /// [`Statistics`]: crate::file::statistics::Statistics + pub fn set_read_page_statistics(mut self, value: bool) -> Self { + self.read_page_stats = Some(value); + self + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_writer_version() { + assert_eq!(WriterVersion::PARQUET_1_0.as_num(), 1); + assert_eq!(WriterVersion::PARQUET_2_0.as_num(), 2); + } + + #[test] + fn test_writer_properties_default_settings() { + let props = WriterProperties::default(); + assert_eq!(props.data_page_size_limit(), DEFAULT_PAGE_SIZE); + assert_eq!( + props.dictionary_page_size_limit(), + DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT + ); + assert_eq!(props.write_batch_size(), DEFAULT_WRITE_BATCH_SIZE); + assert_eq!( + props.max_row_group_row_count(), + Some(DEFAULT_MAX_ROW_GROUP_ROW_COUNT) + ); + assert_eq!(props.max_row_group_bytes(), None); + assert_eq!(props.bloom_filter_position(), DEFAULT_BLOOM_FILTER_POSITION); + assert_eq!(props.writer_version(), DEFAULT_WRITER_VERSION); + assert_eq!(props.created_by(), DEFAULT_CREATED_BY); + assert_eq!(props.key_value_metadata(), None); + assert_eq!(props.encoding(&ColumnPath::from("col")), None); + assert_eq!( + props.compression(&ColumnPath::from("col")), + DEFAULT_COMPRESSION + ); + assert_eq!( + props.dictionary_enabled(&ColumnPath::from("col")), + DEFAULT_DICTIONARY_ENABLED + ); + assert_eq!( + props.statistics_enabled(&ColumnPath::from("col")), + DEFAULT_STATISTICS_ENABLED + ); + assert!( + props + .bloom_filter_properties(&ColumnPath::from("col")) + .is_none() + ); + } + + #[test] + fn test_writer_properties_dictionary_encoding() { + // dictionary encoding is not configurable, and it should be the same for both + // writer version 1 and 2. + for version in &[WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] { + let props = WriterProperties::builder() + .set_writer_version(*version) + .build(); + assert_eq!(props.dictionary_page_encoding(), Encoding::PLAIN); + assert_eq!( + props.dictionary_data_page_encoding(), + Encoding::RLE_DICTIONARY + ); + } + } + + #[test] + #[should_panic(expected = "Dictionary encoding can not be used as fallback encoding")] + fn test_writer_properties_panic_when_plain_dictionary_is_fallback() { + // Should panic when user specifies dictionary encoding as fallback encoding. + WriterProperties::builder() + .set_encoding(Encoding::PLAIN_DICTIONARY) + .build(); + } + + #[test] + #[should_panic(expected = "Dictionary encoding can not be used as fallback encoding")] + fn test_writer_properties_panic_when_rle_dictionary_is_fallback() { + // Should panic when user specifies dictionary encoding as fallback encoding. + WriterProperties::builder() + .set_encoding(Encoding::RLE_DICTIONARY) + .build(); + } + + #[test] + #[should_panic(expected = "Dictionary encoding can not be used as fallback encoding")] + fn test_writer_properties_panic_when_dictionary_is_enabled() { + WriterProperties::builder() + .set_dictionary_enabled(true) + .set_column_encoding(ColumnPath::from("col"), Encoding::RLE_DICTIONARY) + .build(); + } + + #[test] + #[should_panic(expected = "Dictionary encoding can not be used as fallback encoding")] + fn test_writer_properties_panic_when_dictionary_is_disabled() { + WriterProperties::builder() + .set_dictionary_enabled(false) + .set_column_encoding(ColumnPath::from("col"), Encoding::RLE_DICTIONARY) + .build(); + } + + #[test] + fn test_writer_properties_builder() { + let props = WriterProperties::builder() + // file settings + .set_writer_version(WriterVersion::PARQUET_2_0) + .set_data_page_size_limit(10) + .set_dictionary_page_size_limit(20) + .set_write_batch_size(30) + .set_max_row_group_row_count(Some(40)) + .set_created_by("default".to_owned()) + .set_key_value_metadata(Some(vec![KeyValue::new( + "key".to_string(), + "value".to_string(), + )])) + // global column settings + .set_encoding(Encoding::DELTA_BINARY_PACKED) + .set_compression(Compression::GZIP(Default::default())) + .set_dictionary_enabled(false) + .set_statistics_enabled(EnabledStatistics::None) + // specific column settings + .set_column_encoding(ColumnPath::from("col"), Encoding::RLE) + .set_column_compression(ColumnPath::from("col"), Compression::SNAPPY) + .set_column_dictionary_enabled(ColumnPath::from("col"), true) + .set_column_statistics_enabled(ColumnPath::from("col"), EnabledStatistics::Chunk) + .set_column_bloom_filter_enabled(ColumnPath::from("col"), true) + .set_column_bloom_filter_max_ndv(ColumnPath::from("col"), 100_u64) + .set_column_bloom_filter_fpp(ColumnPath::from("col"), 0.1) + .build(); + + fn test_props(props: &WriterProperties) { + assert_eq!(props.writer_version(), WriterVersion::PARQUET_2_0); + assert_eq!(props.data_page_size_limit(), 10); + assert_eq!(props.dictionary_page_size_limit(), 20); + assert_eq!(props.write_batch_size(), 30); + assert_eq!(props.max_row_group_row_count(), Some(40)); + assert_eq!(props.created_by(), "default"); + assert_eq!( + props.key_value_metadata(), + Some(&vec![ + KeyValue::new("key".to_string(), "value".to_string(),) + ]) + ); + + assert_eq!( + props.encoding(&ColumnPath::from("a")), + Some(Encoding::DELTA_BINARY_PACKED) + ); + assert_eq!( + props.compression(&ColumnPath::from("a")), + Compression::GZIP(Default::default()) + ); + assert!(!props.dictionary_enabled(&ColumnPath::from("a"))); + assert_eq!( + props.statistics_enabled(&ColumnPath::from("a")), + EnabledStatistics::None + ); + + assert_eq!( + props.encoding(&ColumnPath::from("col")), + Some(Encoding::RLE) + ); + assert_eq!( + props.compression(&ColumnPath::from("col")), + Compression::SNAPPY + ); + assert!(props.dictionary_enabled(&ColumnPath::from("col"))); + assert_eq!( + props.statistics_enabled(&ColumnPath::from("col")), + EnabledStatistics::Chunk + ); + assert_eq!( + props.bloom_filter_properties(&ColumnPath::from("col")), + Some(&BloomFilterProperties { fpp: 0.1, ndv: 100 }) + ); + } + + // Test direct build of properties + test_props(&props); + + // Test that into_builder() gives the same result + let props_into_builder_and_back = props.into_builder().build(); + test_props(&props_into_builder_and_back); + } + + #[test] + fn test_writer_properties_builder_partial_defaults() { + let props = WriterProperties::builder() + .set_encoding(Encoding::DELTA_BINARY_PACKED) + .set_compression(Compression::GZIP(Default::default())) + .set_bloom_filter_enabled(true) + .set_column_encoding(ColumnPath::from("col"), Encoding::RLE) + .build(); + + assert_eq!( + props.encoding(&ColumnPath::from("col")), + Some(Encoding::RLE) + ); + assert_eq!( + props.compression(&ColumnPath::from("col")), + Compression::GZIP(Default::default()) + ); + assert_eq!( + props.dictionary_enabled(&ColumnPath::from("col")), + DEFAULT_DICTIONARY_ENABLED + ); + assert_eq!( + props.bloom_filter_properties(&ColumnPath::from("col")), + Some(&BloomFilterProperties { + fpp: DEFAULT_BLOOM_FILTER_FPP, + ndv: DEFAULT_BLOOM_FILTER_NDV, + }) + ); + } + + #[test] + #[allow(deprecated)] + fn test_writer_properties_deprecated_max_row_group_size_still_works() { + let props = WriterProperties::builder() + .set_max_row_group_size(42) + .build(); + + assert_eq!(props.max_row_group_row_count(), Some(42)); + assert_eq!(props.max_row_group_size(), 42); + } + + #[test] + #[should_panic(expected = "Cannot have a 0 max row group row count")] + fn test_writer_properties_panic_on_zero_row_group_row_count() { + let _ = WriterProperties::builder().set_max_row_group_row_count(Some(0)); + } + + #[test] + #[should_panic(expected = "Cannot have a 0 max row group bytes")] + fn test_writer_properties_panic_on_zero_row_group_bytes() { + let _ = WriterProperties::builder().set_max_row_group_bytes(Some(0)); + } + + #[test] + fn test_writer_properties_bloom_filter_ndv_fpp_set() { + assert_eq!( + WriterProperties::builder() + .build() + .bloom_filter_properties(&ColumnPath::from("col")), + None + ); + assert_eq!( + WriterProperties::builder() + .set_bloom_filter_max_ndv(100) + .build() + .bloom_filter_properties(&ColumnPath::from("col")), + Some(&BloomFilterProperties { + fpp: DEFAULT_BLOOM_FILTER_FPP, + ndv: 100, + }) + ); + assert_eq!( + WriterProperties::builder() + .set_bloom_filter_fpp(0.1) + .build() + .bloom_filter_properties(&ColumnPath::from("col")), + Some(&BloomFilterProperties { + fpp: 0.1, + ndv: DEFAULT_BLOOM_FILTER_NDV, + }) + ); + } + + #[test] + fn test_writer_properties_column_data_page_v2_compression_ratio_threshold() { + let props = WriterProperties::builder() + .set_data_page_v2_compression_ratio_threshold(0.5) + .set_column_data_page_v2_compression_ratio_threshold(ColumnPath::from("col"), 0.1) + .build(); + + assert_eq!(props.data_page_v2_compression_ratio_threshold(), 0.5); + assert_eq!( + props.column_data_page_v2_compression_ratio_threshold(&ColumnPath::from("col")), + 0.1 + ); + assert_eq!( + props.column_data_page_v2_compression_ratio_threshold(&ColumnPath::from("other")), + 0.5 + ); + } + + #[test] + #[should_panic( + expected = "data_page_v2_compression_ratio_threshold must be a positive finite number" + )] + fn test_writer_properties_panic_on_invalid_data_page_v2_compression_ratio_threshold() { + WriterProperties::builder() + .set_data_page_v2_compression_ratio_threshold(0.0) + .build(); + } + + #[test] + #[allow(deprecated)] + fn test_writer_properties_deprecated_bloom_filter_ndv_setters_still_work() { + let col = ColumnPath::from("col"); + let props = WriterProperties::builder() + .set_bloom_filter_ndv(100) + .set_column_bloom_filter_ndv(col.clone(), 200) + .build(); + assert_eq!( + props.bloom_filter_properties(&ColumnPath::from("other")), + Some(&BloomFilterProperties { + fpp: DEFAULT_BLOOM_FILTER_FPP, + ndv: 100, + }) + ); + assert_eq!( + props.bloom_filter_properties(&col), + Some(&BloomFilterProperties { + fpp: DEFAULT_BLOOM_FILTER_FPP, + ndv: 200, + }) + ); + } + + #[test] + fn test_writer_properties_column_dictionary_page_size_limit() { + let props = WriterProperties::builder() + .set_dictionary_page_size_limit(100) + .set_column_dictionary_page_size_limit(ColumnPath::from("col"), 10) + .build(); + + assert_eq!(props.dictionary_page_size_limit(), 100); + assert_eq!( + props.column_dictionary_page_size_limit(&ColumnPath::from("col")), + 10 + ); + assert_eq!( + props.column_dictionary_page_size_limit(&ColumnPath::from("other")), + 100 + ); + } + + #[test] + fn test_writer_properties_column_data_page_size_limit() { + let props = WriterProperties::builder() + .set_data_page_size_limit(100) + .set_column_data_page_size_limit(ColumnPath::from("col"), 10) + .build(); + + assert_eq!(props.data_page_size_limit(), 100); + assert_eq!( + props.column_data_page_size_limit(&ColumnPath::from("col")), + 10 + ); + assert_eq!( + props.column_data_page_size_limit(&ColumnPath::from("other")), + 100 + ); + } + + #[test] + fn test_reader_properties_default_settings() { + let props = ReaderProperties::builder().build(); + + let codec_options = CodecOptionsBuilder::default() + .set_backward_compatible_lz4(true) + .build(); + + assert_eq!(props.codec_options(), &codec_options); + assert!(!props.read_bloom_filter()); + } + + #[test] + fn test_reader_properties_builder() { + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + + let codec_options = CodecOptionsBuilder::default() + .set_backward_compatible_lz4(false) + .build(); + + assert_eq!(props.codec_options(), &codec_options); + } + + #[test] + fn test_parse_writerversion() { + let mut writer_version = "PARQUET_1_0".parse::().unwrap(); + assert_eq!(writer_version, WriterVersion::PARQUET_1_0); + writer_version = "PARQUET_2_0".parse::().unwrap(); + assert_eq!(writer_version, WriterVersion::PARQUET_2_0); + + // test lowercase + writer_version = "parquet_1_0".parse::().unwrap(); + assert_eq!(writer_version, WriterVersion::PARQUET_1_0); + + // test invalid version + match "PARQUET_-1_0".parse::() { + Ok(_) => panic!("Should not be able to parse PARQUET_-1_0"), + Err(e) => { + assert_eq!(e, "Invalid writer version: PARQUET_-1_0"); + } + } + } + + #[test] + fn test_parse_enabledstatistics() { + let mut enabled_statistics = "NONE".parse::().unwrap(); + assert_eq!(enabled_statistics, EnabledStatistics::None); + enabled_statistics = "CHUNK".parse::().unwrap(); + assert_eq!(enabled_statistics, EnabledStatistics::Chunk); + enabled_statistics = "PAGE".parse::().unwrap(); + assert_eq!(enabled_statistics, EnabledStatistics::Page); + + // test lowercase + enabled_statistics = "none".parse::().unwrap(); + assert_eq!(enabled_statistics, EnabledStatistics::None); + + //test invalid statistics + match "ChunkAndPage".parse::() { + Ok(_) => panic!("Should not be able to parse ChunkAndPage"), + Err(e) => { + assert_eq!(e, "Invalid statistics arg: ChunkAndPage"); + } + } + } + + #[test] + fn test_cdc_options_equality() { + let opts = CdcOptions::default(); + assert_eq!(opts, CdcOptions::default()); + + let custom = CdcOptions { + min_chunk_size: 1024, + max_chunk_size: 8192, + norm_level: 1, + }; + assert_eq!(custom, custom); + assert_ne!(opts, custom); + } + + #[test] + fn test_bloom_filter_builder_default() { + let props = BloomFilterProperties::builder().build(); + assert_eq!(props.fpp, DEFAULT_BLOOM_FILTER_FPP); + assert_eq!(props.ndv, DEFAULT_BLOOM_FILTER_NDV); + assert_eq!(props, BloomFilterProperties::default()); + assert_eq!( + BloomFilterPropertiesBuilder::new().build(), + BloomFilterProperties::default() + ); + } + + #[test] + fn test_bloom_filter_builder_explicit_fpp() { + let props = BloomFilterProperties::builder().with_fpp(0.01).build(); + assert_eq!(props.fpp, 0.01); + assert_eq!(props.ndv, DEFAULT_BLOOM_FILTER_NDV); + } + + #[test] + fn test_bloom_filter_builder_explicit_ndv() { + let props = BloomFilterProperties::builder().with_max_ndv(1000).build(); + assert_eq!(props.fpp, DEFAULT_BLOOM_FILTER_FPP); + assert_eq!(props.ndv, 1000); + } + + #[test] + fn test_bloom_filter_builder_validates_fpp() { + for wrong_val in [0.0_f64, 1.0, -0.5, 2.0] { + let result = std::panic::catch_unwind(|| { + BloomFilterProperties::builder().with_fpp(wrong_val).build() + }); + assert!( + result.is_err(), + "with_fpp({wrong_val}).build() should reject value outside (0, 1)" + ); + } + } + + #[test] + fn test_bloom_filter_builder_try_build_validates_fpp() { + for wrong_val in [0.0_f64, 1.0, -0.5, 2.0] { + let result = BloomFilterProperties::builder() + .with_fpp(wrong_val) + .try_build(); + assert!( + result.is_err(), + "try_build() should return Err for fpp outside (0, 1)" + ); + } + + let ok = BloomFilterProperties::builder() + .with_fpp(0.01) + .with_max_ndv(1000) + .try_build() + .expect("valid fpp should yield Ok"); + assert_eq!(ok.fpp, 0.01); + assert_eq!(ok.ndv, 1000); + } + + #[test] + fn test_column_specific_implicit_ndv_uses_row_group_size() { + let custom_row_group_size: usize = 7777; + let col = ColumnPath::from("col"); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(custom_row_group_size)) + .set_column_bloom_filter_enabled(col.clone(), true) + .build(); + let bf = props + .bloom_filter_properties(&col) + .expect("bloom filter should be enabled for col"); + + assert_eq!(bf.ndv, custom_row_group_size as u64); + assert_eq!(bf.fpp, DEFAULT_BLOOM_FILTER_FPP); + } + + #[test] + fn test_set_bloom_filter_properties_applied_globally() { + let bf = BloomFilterProperties::builder() + .with_fpp(0.01) + .with_max_ndv(500) + .build(); + let props = WriterProperties::builder() + .set_bloom_filter_properties(bf.clone()) + .build(); + + assert_eq!( + props.bloom_filter_properties(&ColumnPath::from("a")), + Some(&bf), + ); + assert_eq!( + props.bloom_filter_properties(&ColumnPath::from("b")), + Some(&bf), + ); + } + + #[test] + fn test_set_column_bloom_filter_properties_overrides_global() { + let global = BloomFilterProperties::builder() + .with_fpp(0.01) + .with_max_ndv(500) + .build(); + let tailored = BloomFilterProperties::builder() + .with_fpp(0.02) + .with_max_ndv(1000) + .build(); + + let col = ColumnPath::from("col"); + let props = WriterProperties::builder() + .set_bloom_filter_properties(global.clone()) + .set_column_bloom_filter_properties(col.clone(), tailored.clone()) + .build(); + + assert_eq!(props.bloom_filter_properties(&col), Some(&tailored)); + assert_eq!( + props.bloom_filter_properties(&ColumnPath::from("other")), + Some(&global) + ); + } + + #[test] + fn test_set_bloom_filter_properties_preserve_explicit_ndv() { + let bf = BloomFilterProperties::builder().with_max_ndv(42).build(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(99_999)) + .set_bloom_filter_properties(bf) + .build(); + let result = props + .bloom_filter_properties(&ColumnPath::from("col")) + .expect("bloom filter should be enabled"); + + assert_eq!( + result.ndv, 42, + "explicit ndv must not be overridden by row-group-size fallback" + ); + } +} diff --git a/vendor/parquet/src/file/reader.rs b/vendor/parquet/src/file/reader.rs new file mode 100644 index 0000000..2b3c46f --- /dev/null +++ b/vendor/parquet/src/file/reader.rs @@ -0,0 +1,321 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! File reader API and methods to access file metadata, row group +//! readers to read individual column chunks, or access record +//! iterator. + +use bytes::{Buf, Bytes}; +use std::fs::File; +use std::io::{BufReader, Seek, SeekFrom}; +use std::{io::Read, sync::Arc}; + +use crate::bloom_filter::Sbbf; +use crate::column::page::PageIterator; +use crate::column::{page::PageReader, reader::ColumnReader}; +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::*; +pub use crate::file::serialized_reader::{SerializedFileReader, SerializedPageReader}; +use crate::record::reader::RowIter; +use crate::schema::types::Type as SchemaType; + +use crate::basic::Type; + +use crate::column::reader::ColumnReaderImpl; + +/// Length should return the total number of bytes in the input source. +/// It's mainly used to read the metadata, which is at the end of the source. +#[allow(clippy::len_without_is_empty)] +pub trait Length { + /// Returns the amount of bytes of the inner source. + fn len(&self) -> u64; +} + +/// Generates [`Read`]ers to read chunks of a Parquet data source. +/// +/// The Parquet reader uses [`ChunkReader`] to access Parquet data, allowing +/// multiple decoders to read concurrently from different locations in the same +/// file. +/// +/// The trait functions both as a reader and a factory for readers. +/// * random access via [`Self::get_bytes`] +/// * sequential access via the reader returned via factory method [`Self::get_read`] +/// +/// # Provided Implementations +/// * [`File`] for reading from local file system +/// * [`Bytes`] for reading from an in-memory buffer +/// +/// User provided implementations can implement more sophisticated behaviors +/// such as on-demand buffering or scan sharing. +pub trait ChunkReader: Length + Send + Sync { + /// The concrete type of reader returned by this trait + type T: Read; + + /// Get a [`Read`] instance starting at the provided file offset + /// + /// Returned readers follow the model of [`File::try_clone`] where mutations + /// of one reader affect all readers. Thus subsequent or concurrent calls to + /// [`Self::get_read`] or [`Self::get_bytes`] may cause side-effects on + /// previously returned readers. Callers of `get_read` should take care + /// to avoid race conditions. + fn get_read(&self, start: u64) -> Result; + + /// Get a range of data in memory as [`Bytes`] + /// + /// Similarly to [`Self::get_read`], this method may have side-effects on + /// previously returned readers. + fn get_bytes(&self, start: u64, length: usize) -> Result; +} + +impl Length for File { + fn len(&self) -> u64 { + self.metadata().map(|m| m.len()).unwrap_or(0u64) + } +} + +impl ChunkReader for File { + type T = BufReader; + + fn get_read(&self, start: u64) -> Result { + let mut reader = self.try_clone()?; + reader.seek(SeekFrom::Start(start))?; + Ok(BufReader::new(reader)) + } + + fn get_bytes(&self, start: u64, length: usize) -> Result { + let mut buffer = Vec::with_capacity(length); + let mut reader = self.try_clone()?; + reader.seek(SeekFrom::Start(start))?; + let read = reader.take(length as _).read_to_end(&mut buffer)?; + + if read != length { + return Err(eof_err!( + "Expected to read {} bytes, read only {}", + length, + read + )); + } + Ok(buffer.into()) + } +} + +impl Length for Bytes { + fn len(&self) -> u64 { + self.len() as u64 + } +} + +impl ChunkReader for Bytes { + type T = bytes::buf::Reader; + + fn get_read(&self, start: u64) -> Result { + let start = start as usize; + if start > self.len() { + return Err(eof_err!( + "Expected to read at offset {start}, while file has length {}", + self.len() + )); + } + Ok(self.slice(start..).reader()) + } + + fn get_bytes(&self, start: u64, length: usize) -> Result { + let start = start as usize; + if start > self.len() || start + length > self.len() { + return Err(eof_err!( + "Expected to read {} bytes at offset {}, while file has length {}", + length, + start, + self.len() + )); + } + Ok(self.slice(start..start + length)) + } +} + +// ---------------------------------------------------------------------- +// APIs for file & row group readers + +/// Parquet file reader API. With this, user can get metadata information about the +/// Parquet file, can get reader for each row group, and access record iterator. +pub trait FileReader: Send + Sync { + /// Get metadata information about this file. + fn metadata(&self) -> &ParquetMetaData; + + /// Get the total number of row groups for this file. + fn num_row_groups(&self) -> usize; + + /// Get the `i`th row group reader. Note this doesn't do bound check. + fn get_row_group(&self, i: usize) -> Result>; + + /// Get an iterator over the row in this file, see [`RowIter`] for caveats. + /// + /// Iterator will automatically load the next row group to advance. + /// + /// Projected schema can be a subset of or equal to the file schema, when it is None, + /// full file schema is assumed. + fn get_row_iter(&self, projection: Option) -> Result>; +} + +/// Parquet row group reader API. With this, user can get metadata information about the +/// row group, as well as readers for each individual column chunk. +pub trait RowGroupReader: Send + Sync { + /// Get metadata information about this row group. + fn metadata(&self) -> &RowGroupMetaData; + + /// Get the total number of column chunks in this row group. + fn num_columns(&self) -> usize; + + /// Get page reader for the `i`th column chunk. + fn get_column_page_reader(&self, i: usize) -> Result>; + + /// Get value reader for the `i`th column chunk. + fn get_column_reader(&self, i: usize) -> Result { + let schema_descr = self.metadata().schema_descr(); + let col_descr = schema_descr.column(i); + let col_page_reader = self.get_column_page_reader(i)?; + let col_reader = match col_descr.physical_type() { + Type::BOOLEAN => { + ColumnReader::BoolColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader)) + } + Type::INT32 => { + ColumnReader::Int32ColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader)) + } + Type::INT64 => { + ColumnReader::Int64ColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader)) + } + Type::INT96 => { + ColumnReader::Int96ColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader)) + } + Type::FLOAT => { + ColumnReader::FloatColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader)) + } + Type::DOUBLE => { + ColumnReader::DoubleColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader)) + } + Type::BYTE_ARRAY => ColumnReader::ByteArrayColumnReader(ColumnReaderImpl::new( + col_descr, + col_page_reader, + )), + Type::FIXED_LEN_BYTE_ARRAY => ColumnReader::FixedLenByteArrayColumnReader( + ColumnReaderImpl::new(col_descr, col_page_reader), + ), + }; + Ok(col_reader) + } + + /// Get bloom filter for the `i`th column chunk, if present and the reader was configured + /// to read bloom filters. + fn get_column_bloom_filter(&self, i: usize) -> Option<&Sbbf>; + + /// Get an iterator over the row in this file, see [`RowIter`] for caveats. + /// + /// Projected schema can be a subset of or equal to the file schema, when it is None, + /// full file schema is assumed. + fn get_row_iter(&self, projection: Option) -> Result>; +} + +// ---------------------------------------------------------------------- +// Iterator + +/// Implementation of page iterator for parquet file. +pub struct FilePageIterator { + column_index: usize, + row_group_indices: Box + Send>, + file_reader: Arc, +} + +impl FilePageIterator { + /// Creates a page iterator for all row groups in file. + pub fn new(column_index: usize, file_reader: Arc) -> Result { + let num_row_groups = file_reader.metadata().num_row_groups(); + + let row_group_indices = Box::new(0..num_row_groups); + + Self::with_row_groups(column_index, row_group_indices, file_reader) + } + + /// Create page iterator from parquet file reader with only some row groups. + pub fn with_row_groups( + column_index: usize, + row_group_indices: Box + Send>, + file_reader: Arc, + ) -> Result { + // Check that column_index is valid + let num_columns = file_reader + .metadata() + .file_metadata() + .schema_descr() + .num_columns(); + + if column_index >= num_columns { + return Err(ParquetError::IndexOutOfBound(column_index, num_columns)); + } + + // We don't check iterators here because iterator may be infinite + Ok(Self { + column_index, + row_group_indices, + file_reader, + }) + } +} + +impl Iterator for FilePageIterator { + type Item = Result>; + + fn next(&mut self) -> Option>> { + self.row_group_indices.next().map(|row_group_index| { + self.file_reader + .get_row_group(row_group_index) + .and_then(|r| r.get_column_page_reader(self.column_index)) + }) + } +} + +impl PageIterator for FilePageIterator {} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_bytes_chunk_reader_get_read_out_of_bounds() { + let data = Bytes::from(vec![0, 1, 2, 3]); + let err = data.get_read(5).unwrap_err(); + assert_eq!( + err.to_string(), + "EOF: Expected to read at offset 5, while file has length 4" + ); + } + + #[test] + fn test_bytes_chunk_reader_get_bytes_out_of_bounds() { + let data = Bytes::from(vec![0, 1, 2, 3]); + let err = data.get_bytes(5, 1).unwrap_err(); + assert_eq!( + err.to_string(), + "EOF: Expected to read 1 bytes at offset 5, while file has length 4" + ); + + let err = data.get_bytes(2, 3).unwrap_err(); + assert_eq!( + err.to_string(), + "EOF: Expected to read 3 bytes at offset 2, while file has length 4" + ); + } +} diff --git a/vendor/parquet/src/file/serialized_reader.rs b/vendor/parquet/src/file/serialized_reader.rs new file mode 100644 index 0000000..113e520 --- /dev/null +++ b/vendor/parquet/src/file/serialized_reader.rs @@ -0,0 +1,2874 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Contains implementations of the reader traits FileReader, RowGroupReader and PageReader +//! Also contains implementations of the ChunkReader for files (with buffering) and byte arrays (RAM) + +use crate::basic::{PageType, Type}; +use crate::bloom_filter::Sbbf; +use crate::column::page::{Page, PageMetadata, PageReader}; +use crate::compression::{Codec, create_codec}; +#[cfg(feature = "encryption")] +use crate::encryption::decrypt::{CryptoContext, read_and_decrypt}; +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::thrift::PageHeader; +use crate::file::page_index::offset_index::{OffsetIndexMetaData, PageLocation}; +use crate::file::statistics; +use crate::file::{ + metadata::*, + properties::{ReaderProperties, ReaderPropertiesPtr}, + reader::*, +}; +#[cfg(feature = "encryption")] +use crate::parquet_thrift::ThriftSliceInputProtocol; +use crate::parquet_thrift::{ReadThrift, ThriftReadInputProtocol}; +use crate::record::Row; +use crate::record::reader::RowIter; +use crate::schema::types::{SchemaDescPtr, Type as SchemaType}; +use bytes::Bytes; +use std::collections::VecDeque; +use std::{fs::File, io::Read, path::Path, sync::Arc}; + +impl TryFrom for SerializedFileReader { + type Error = ParquetError; + + fn try_from(file: File) -> Result { + Self::new(file) + } +} + +impl TryFrom<&Path> for SerializedFileReader { + type Error = ParquetError; + + fn try_from(path: &Path) -> Result { + let file = File::open(path)?; + Self::try_from(file) + } +} + +impl TryFrom for SerializedFileReader { + type Error = ParquetError; + + fn try_from(path: String) -> Result { + Self::try_from(Path::new(&path)) + } +} + +impl TryFrom<&str> for SerializedFileReader { + type Error = ParquetError; + + fn try_from(path: &str) -> Result { + Self::try_from(Path::new(&path)) + } +} + +/// Conversion into a [`RowIter`] +/// using the full file schema over all row groups. +impl IntoIterator for SerializedFileReader { + type Item = Result; + type IntoIter = RowIter<'static>; + + fn into_iter(self) -> Self::IntoIter { + RowIter::from_file_into(Box::new(self)) + } +} + +// ---------------------------------------------------------------------- +// Implementations of file & row group readers + +/// A serialized implementation for Parquet [`FileReader`]. +pub struct SerializedFileReader { + chunk_reader: Arc, + metadata: Arc, + props: ReaderPropertiesPtr, +} + +/// A predicate for filtering row groups, invoked with the metadata and index +/// of each row group in the file. Only row groups for which the predicate +/// evaluates to `true` will be scanned +pub type ReadGroupPredicate = Box bool>; + +/// A builder for [`ReadOptions`]. +/// For the predicates that are added to the builder, +/// they will be chained using 'AND' to filter the row groups. +#[derive(Default)] +pub struct ReadOptionsBuilder { + predicates: Vec, + enable_page_index: bool, + props: Option, + metadata_options: ParquetMetaDataOptions, +} + +impl ReadOptionsBuilder { + /// New builder + pub fn new() -> Self { + Self::default() + } + + /// Add a predicate on row group metadata to the reading option, + /// Filter only row groups that match the predicate criteria + pub fn with_predicate(mut self, predicate: ReadGroupPredicate) -> Self { + self.predicates.push(predicate); + self + } + + /// Add a range predicate on filtering row groups if their midpoints are within + /// the Closed-Open range `[start..end) {x | start <= x < end}` + pub fn with_range(mut self, start: i64, end: i64) -> Self { + assert!(start < end); + let predicate = move |rg: &RowGroupMetaData, _: usize| { + let mid = get_midpoint_offset(rg); + mid >= start && mid < end + }; + self.predicates.push(Box::new(predicate)); + self + } + + /// Enable reading the page index structures described in + /// "[Column Index] Layout to Support Page Skipping" + /// + /// [Column Index]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + pub fn with_page_index(mut self) -> Self { + self.enable_page_index = true; + self + } + + /// Set the [`ReaderProperties`] configuration. + pub fn with_reader_properties(mut self, properties: ReaderProperties) -> Self { + self.props = Some(properties); + self + } + + /// Provide a Parquet schema to use when decoding the metadata. The schema in the Parquet + /// footer will be skipped. + pub fn with_parquet_schema(mut self, schema: SchemaDescPtr) -> Self { + self.metadata_options.set_schema(schema); + self + } + + /// Set whether to convert the [`encoding_stats`] in the Parquet `ColumnMetaData` to a bitmask + /// (defaults to `false`). + /// + /// See [`ColumnChunkMetaData::page_encoding_stats_mask`] for an explanation of why this + /// might be desirable. + /// + /// [`encoding_stats`]: + /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L917 + pub fn with_encoding_stats_as_mask(mut self, val: bool) -> Self { + self.metadata_options.set_encoding_stats_as_mask(val); + self + } + + /// Sets the decoding policy for [`encoding_stats`] in the Parquet `ColumnMetaData`. + /// + /// [`encoding_stats`]: + /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L917 + pub fn with_encoding_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self { + self.metadata_options.set_encoding_stats_policy(policy); + self + } + + /// Sets the decoding policy for [`statistics`] in the Parquet `ColumnMetaData`. + /// + /// [`statistics`]: + /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L912 + pub fn with_column_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self { + self.metadata_options.set_column_stats_policy(policy); + self + } + + /// Sets the decoding policy for [`size_statistics`] in the Parquet `ColumnMetaData`. + /// + /// [`size_statistics`]: + /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L936 + pub fn with_size_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self { + self.metadata_options.set_size_stats_policy(policy); + self + } + + /// Seal the builder and return the read options + pub fn build(self) -> ReadOptions { + let props = self + .props + .unwrap_or_else(|| ReaderProperties::builder().build()); + ReadOptions { + predicates: self.predicates, + enable_page_index: self.enable_page_index, + props, + metadata_options: self.metadata_options, + } + } +} + +/// A collection of options for reading a Parquet file. +/// +/// Predicates are currently only supported on row group metadata. +/// All predicates will be chained using 'AND' to filter the row groups. +pub struct ReadOptions { + predicates: Vec, + enable_page_index: bool, + props: ReaderProperties, + metadata_options: ParquetMetaDataOptions, +} + +impl SerializedFileReader { + /// Creates file reader from a Parquet file. + /// Returns an error if the Parquet file does not exist or is corrupt. + pub fn new(chunk_reader: R) -> Result { + let metadata = ParquetMetaDataReader::new().parse_and_finish(&chunk_reader)?; + let props = Arc::new(ReaderProperties::builder().build()); + Ok(Self { + chunk_reader: Arc::new(chunk_reader), + metadata: Arc::new(metadata), + props, + }) + } + + /// Creates file reader from a Parquet file with read options. + /// Returns an error if the Parquet file does not exist or is corrupt. + pub fn new_with_options(chunk_reader: R, options: ReadOptions) -> Result { + let mut metadata_builder = ParquetMetaDataReader::new() + .with_metadata_options(Some(options.metadata_options.clone())) + .parse_and_finish(&chunk_reader)? + .into_builder(); + let mut predicates = options.predicates; + + // Filter row groups based on the predicates + for (i, rg_meta) in metadata_builder.take_row_groups().into_iter().enumerate() { + let mut keep = true; + for predicate in &mut predicates { + if !predicate(&rg_meta, i) { + keep = false; + break; + } + } + if keep { + metadata_builder = metadata_builder.add_row_group(rg_meta); + } + } + + let mut metadata = metadata_builder.build(); + + // If page indexes are desired, build them with the filtered set of row groups + if options.enable_page_index { + let mut reader = ParquetMetaDataReader::new_with_metadata(metadata) + .with_page_index_policy(PageIndexPolicy::Required); + reader.read_page_indexes(&chunk_reader)?; + metadata = reader.finish()?; + } + + Ok(Self { + chunk_reader: Arc::new(chunk_reader), + metadata: Arc::new(metadata), + props: Arc::new(options.props), + }) + } +} + +/// Get midpoint offset for a row group +fn get_midpoint_offset(meta: &RowGroupMetaData) -> i64 { + let col = meta.column(0); + let mut offset = col.data_page_offset(); + if let Some(dic_offset) = col.dictionary_page_offset() { + if offset > dic_offset { + offset = dic_offset + } + }; + offset + meta.compressed_size() / 2 +} + +impl FileReader for SerializedFileReader { + fn metadata(&self) -> &ParquetMetaData { + &self.metadata + } + + fn num_row_groups(&self) -> usize { + self.metadata.num_row_groups() + } + + fn get_row_group(&self, i: usize) -> Result> { + let row_group_metadata = self.metadata.row_group(i); + // Row groups should be processed sequentially. + let props = Arc::clone(&self.props); + let f = Arc::clone(&self.chunk_reader); + Ok(Box::new(SerializedRowGroupReader::new( + f, + row_group_metadata, + self.metadata.offset_index().map(|x| x[i].as_slice()), + props, + )?)) + } + + fn get_row_iter(&self, projection: Option) -> Result> { + RowIter::from_file(projection, self) + } +} + +/// A serialized implementation for Parquet [`RowGroupReader`]. +pub struct SerializedRowGroupReader<'a, R: ChunkReader> { + chunk_reader: Arc, + metadata: &'a RowGroupMetaData, + offset_index: Option<&'a [OffsetIndexMetaData]>, + props: ReaderPropertiesPtr, + bloom_filters: Vec>, +} + +impl<'a, R: ChunkReader> SerializedRowGroupReader<'a, R> { + /// Creates new row group reader from a file, row group metadata and custom config. + pub fn new( + chunk_reader: Arc, + metadata: &'a RowGroupMetaData, + offset_index: Option<&'a [OffsetIndexMetaData]>, + props: ReaderPropertiesPtr, + ) -> Result { + let bloom_filters = if props.read_bloom_filter() { + metadata + .columns() + .iter() + .map(|col| Sbbf::read_from_column_chunk(col, &*chunk_reader)) + .collect::>>()? + } else { + std::iter::repeat_n(None, metadata.columns().len()).collect() + }; + Ok(Self { + chunk_reader, + metadata, + offset_index, + props, + bloom_filters, + }) + } +} + +impl RowGroupReader for SerializedRowGroupReader<'_, R> { + fn metadata(&self) -> &RowGroupMetaData { + self.metadata + } + + fn num_columns(&self) -> usize { + self.metadata.num_columns() + } + + // TODO: fix PARQUET-816 + fn get_column_page_reader(&self, i: usize) -> Result> { + let col = self.metadata.column(i); + + let page_locations = self.offset_index.map(|x| x[i].page_locations.clone()); + + let props = Arc::clone(&self.props); + Ok(Box::new(SerializedPageReader::new_with_properties( + Arc::clone(&self.chunk_reader), + col, + usize::try_from(self.metadata.num_rows())?, + page_locations, + props, + )?)) + } + + /// get bloom filter for the `i`th column + fn get_column_bloom_filter(&self, i: usize) -> Option<&Sbbf> { + self.bloom_filters[i].as_ref() + } + + fn get_row_iter(&self, projection: Option) -> Result> { + RowIter::from_row_group(projection, self) + } +} + +/// Decodes a [`Page`] from the provided `buffer` +pub(crate) fn decode_page( + page_header: PageHeader, + buffer: Bytes, + physical_type: Type, + decompressor: Option<&mut Box>, +) -> Result { + // Verify the 32-bit CRC checksum of the page + #[cfg(feature = "crc")] + if let Some(expected_crc) = page_header.crc { + let crc = crc32fast::hash(&buffer); + if crc != expected_crc as u32 { + return Err(general_err!("Page CRC checksum mismatch")); + } + } + + // When processing data page v2, depending on enabled compression for the + // page, we should account for uncompressed data ('offset') of + // repetition and definition levels. + // + // We always use 0 offset for other pages other than v2, `true` flag means + // that compression will be applied if decompressor is defined + let mut offset: usize = 0; + let mut can_decompress = true; + + if let Some(ref header_v2) = page_header.data_page_header_v2 { + if header_v2.definition_levels_byte_length < 0 + || header_v2.repetition_levels_byte_length < 0 + || header_v2.definition_levels_byte_length + header_v2.repetition_levels_byte_length + > page_header.uncompressed_page_size + { + return Err(general_err!( + "DataPage v2 header contains implausible values \ + for definition_levels_byte_length ({}) \ + and repetition_levels_byte_length ({}) \ + given DataPage header provides uncompressed_page_size ({})", + header_v2.definition_levels_byte_length, + header_v2.repetition_levels_byte_length, + page_header.uncompressed_page_size + )); + } + offset = usize::try_from( + header_v2.definition_levels_byte_length + header_v2.repetition_levels_byte_length, + )?; + // When is_compressed flag is missing the page is considered compressed + can_decompress = header_v2.is_compressed.unwrap_or(true); + } + + let buffer = match decompressor { + Some(decompressor) if can_decompress => { + let uncompressed_page_size = usize::try_from(page_header.uncompressed_page_size)?; + if offset > buffer.len() || offset > uncompressed_page_size { + return Err(general_err!("Invalid page header")); + } + let decompressed_size = uncompressed_page_size - offset; + let mut decompressed = Vec::with_capacity(uncompressed_page_size); + decompressed.extend_from_slice(&buffer[..offset]); + // decompressed size of zero corresponds to a page with no non-null values + // see https://github.com/apache/parquet-format/blob/master/README.md#data-pages + if decompressed_size > 0 { + let compressed = &buffer[offset..]; + decompressor.decompress(compressed, &mut decompressed, Some(decompressed_size))?; + } + + if decompressed.len() != uncompressed_page_size { + return Err(general_err!( + "Actual decompressed size doesn't match the expected one ({} vs {})", + decompressed.len(), + uncompressed_page_size + )); + } + + Bytes::from(decompressed) + } + _ => buffer, + }; + + let result = match page_header.r#type { + PageType::DICTIONARY_PAGE => { + let dict_header = page_header.dictionary_page_header.as_ref().ok_or_else(|| { + ParquetError::General("Missing dictionary page header".to_string()) + })?; + let is_sorted = dict_header.is_sorted.unwrap_or(false); + Page::DictionaryPage { + buf: buffer, + num_values: dict_header.num_values.try_into()?, + encoding: dict_header.encoding, + is_sorted, + } + } + PageType::DATA_PAGE => { + let header = page_header + .data_page_header + .ok_or_else(|| ParquetError::General("Missing V1 data page header".to_string()))?; + Page::DataPage { + buf: buffer, + num_values: header.num_values.try_into()?, + encoding: header.encoding, + def_level_encoding: header.definition_level_encoding, + rep_level_encoding: header.repetition_level_encoding, + statistics: statistics::from_thrift_page_stats(physical_type, header.statistics)?, + } + } + PageType::DATA_PAGE_V2 => { + let header = page_header + .data_page_header_v2 + .ok_or_else(|| ParquetError::General("Missing V2 data page header".to_string()))?; + let is_compressed = header.is_compressed.unwrap_or(true); + Page::DataPageV2 { + buf: buffer, + num_values: header.num_values.try_into()?, + encoding: header.encoding, + num_nulls: header.num_nulls.try_into()?, + num_rows: header.num_rows.try_into()?, + def_levels_byte_len: header.definition_levels_byte_length.try_into()?, + rep_levels_byte_len: header.repetition_levels_byte_length.try_into()?, + is_compressed, + statistics: statistics::from_thrift_page_stats(physical_type, header.statistics)?, + } + } + _ => { + // For unknown page type (e.g., INDEX_PAGE), skip and read next. + return Err(general_err!( + "Page type {:?} is not supported", + page_header.r#type + )); + } + }; + + Ok(result) +} + +enum SerializedPageReaderState { + Values { + /// The current byte offset in the reader + /// Note that offset is u64 (i.e., not usize) to support 32-bit architectures such as WASM + offset: u64, + + /// The length of the chunk in bytes + /// Note that remaining_bytes is u64 (i.e., not usize) to support 32-bit architectures such as WASM + remaining_bytes: u64, + + // If the next page header has already been "peeked", we will cache it and it`s length here + next_page_header: Option>, + + /// The index of the data page within this column chunk + page_index: usize, + + /// Whether the next page is expected to be a dictionary page + require_dictionary: bool, + }, + Pages { + /// Remaining page locations + page_locations: VecDeque, + /// Remaining dictionary location if any + dictionary_page: Option, + /// The total number of rows in this column chunk + total_rows: usize, + /// The index of the data page within this column chunk + page_index: usize, + }, +} + +#[derive(Default)] +struct SerializedPageReaderContext { + /// Controls decoding of page-level statistics + read_stats: bool, + /// Crypto context carrying objects required for decryption + #[cfg(feature = "encryption")] + crypto_context: Option>, +} + +/// A serialized implementation for Parquet [`PageReader`]. +pub struct SerializedPageReader { + /// The chunk reader + reader: Arc, + + /// The compression codec for this column chunk. Only set for non-PLAIN codec. + decompressor: Option>, + + /// Column chunk type. + physical_type: Type, + + state: SerializedPageReaderState, + + context: SerializedPageReaderContext, +} + +impl SerializedPageReader { + /// Creates a new serialized page reader from a chunk reader and metadata + pub fn new( + reader: Arc, + column_chunk_metadata: &ColumnChunkMetaData, + total_rows: usize, + page_locations: Option>, + ) -> Result { + let props = Arc::new(ReaderProperties::builder().build()); + SerializedPageReader::new_with_properties( + reader, + column_chunk_metadata, + total_rows, + page_locations, + props, + ) + } + + /// Stub No-op implementation when encryption is disabled. + #[cfg(all(feature = "arrow", not(feature = "encryption")))] + pub(crate) fn add_crypto_context( + self, + _rg_idx: usize, + _column_idx: usize, + _parquet_meta_data: &ParquetMetaData, + _column_chunk_metadata: &ColumnChunkMetaData, + ) -> Result> { + Ok(self) + } + + /// Adds any necessary crypto context to this page reader, if encryption is enabled. + #[cfg(feature = "encryption")] + pub(crate) fn add_crypto_context( + mut self, + rg_idx: usize, + column_idx: usize, + parquet_meta_data: &ParquetMetaData, + column_chunk_metadata: &ColumnChunkMetaData, + ) -> Result> { + let Some(file_decryptor) = parquet_meta_data.file_decryptor() else { + return Ok(self); + }; + let Some(crypto_metadata) = column_chunk_metadata.crypto_metadata() else { + return Ok(self); + }; + let crypto_context = + CryptoContext::for_column(file_decryptor, crypto_metadata, rg_idx, column_idx)?; + self.context.crypto_context = Some(Arc::new(crypto_context)); + Ok(self) + } + + /// Creates a new serialized page with custom options. + pub fn new_with_properties( + reader: Arc, + meta: &ColumnChunkMetaData, + total_rows: usize, + page_locations: Option>, + props: ReaderPropertiesPtr, + ) -> Result { + let decompressor = create_codec(meta.compression(), props.codec_options())?; + let (start, len) = meta.byte_range(); + + let state = match page_locations { + Some(locations) => { + // If the offset of the first page doesn't match the start of the column chunk + // then the preceding space must contain a dictionary page. + let dictionary_page = match locations.first() { + Some(dict_offset) if dict_offset.offset as u64 != start => Some(PageLocation { + offset: start as i64, + compressed_page_size: (dict_offset.offset as u64 - start) as i32, + first_row_index: 0, + }), + _ => None, + }; + + SerializedPageReaderState::Pages { + page_locations: locations.into(), + dictionary_page, + total_rows, + page_index: 0, + } + } + None => SerializedPageReaderState::Values { + offset: start, + remaining_bytes: len, + next_page_header: None, + page_index: 0, + require_dictionary: meta.dictionary_page_offset().is_some(), + }, + }; + let mut context = SerializedPageReaderContext::default(); + if props.read_page_stats() { + context.read_stats = true; + } + Ok(Self { + reader, + decompressor, + state, + physical_type: meta.column_type(), + context, + }) + } + + /// Similar to `peek_next_page`, but returns the offset of the next page instead of the page metadata. + /// Unlike page metadata, an offset can uniquely identify a page. + /// + /// This is used when we need to read parquet with row-filter, and we don't want to decompress the page twice. + /// This function allows us to check if the next page is being cached or read previously. + #[cfg(test)] + fn peek_next_page_offset(&mut self) -> Result> { + match &mut self.state { + SerializedPageReaderState::Values { + offset, + remaining_bytes, + next_page_header, + page_index, + require_dictionary, + } => { + loop { + if *remaining_bytes == 0 { + return Ok(None); + } + return if let Some(header) = next_page_header.as_ref() { + if let Ok(_page_meta) = PageMetadata::try_from(&**header) { + Ok(Some(*offset)) + } else { + // For unknown page type (e.g., INDEX_PAGE), skip and read next. + *next_page_header = None; + continue; + } + } else { + let mut read = self.reader.get_read(*offset)?; + let (header_len, header) = Self::read_page_header_len( + &self.context, + &mut read, + *page_index, + *require_dictionary, + )?; + *offset += header_len as u64; + *remaining_bytes -= header_len as u64; + let page_meta = if let Ok(_page_meta) = PageMetadata::try_from(&header) { + Ok(Some(*offset)) + } else { + // For unknown page type (e.g., INDEX_PAGE), skip and read next. + continue; + }; + *next_page_header = Some(Box::new(header)); + page_meta + }; + } + } + SerializedPageReaderState::Pages { + page_locations, + dictionary_page, + .. + } => { + if let Some(page) = dictionary_page { + Ok(Some(page.offset as u64)) + } else if let Some(page) = page_locations.front() { + Ok(Some(page.offset as u64)) + } else { + Ok(None) + } + } + } + } + + fn read_page_header_len( + context: &SerializedPageReaderContext, + input: &mut T, + page_index: usize, + dictionary_page: bool, + ) -> Result<(usize, PageHeader)> { + /// A wrapper around a [`std::io::Read`] that keeps track of the bytes read + struct TrackedRead { + inner: R, + bytes_read: usize, + } + + impl Read for TrackedRead { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + let v = self.inner.read(buf)?; + self.bytes_read += v; + Ok(v) + } + } + + let mut tracked = TrackedRead { + inner: input, + bytes_read: 0, + }; + let header = context.read_page_header(&mut tracked, page_index, dictionary_page)?; + Ok((tracked.bytes_read, header)) + } + + fn read_page_header_len_from_bytes( + context: &SerializedPageReaderContext, + buffer: &[u8], + page_index: usize, + dictionary_page: bool, + ) -> Result<(usize, PageHeader)> { + let mut input = std::io::Cursor::new(buffer); + let header = context.read_page_header(&mut input, page_index, dictionary_page)?; + let header_len = input.position() as usize; + Ok((header_len, header)) + } +} + +#[cfg(not(feature = "encryption"))] +impl SerializedPageReaderContext { + fn read_page_header( + &self, + input: &mut T, + _page_index: usize, + _dictionary_page: bool, + ) -> Result { + let mut prot = ThriftReadInputProtocol::new(input); + if self.read_stats { + Ok(PageHeader::read_thrift(&mut prot)?) + } else { + Ok(PageHeader::read_thrift_without_stats(&mut prot)?) + } + } + + fn decrypt_page_data( + &self, + buffer: T, + _page_index: usize, + _dictionary_page: bool, + ) -> Result { + Ok(buffer) + } +} + +#[cfg(feature = "encryption")] +impl SerializedPageReaderContext { + fn read_page_header( + &self, + input: &mut T, + page_index: usize, + dictionary_page: bool, + ) -> Result { + match self.page_crypto_context(page_index, dictionary_page) { + None => { + let mut prot = ThriftReadInputProtocol::new(input); + if self.read_stats { + Ok(PageHeader::read_thrift(&mut prot)?) + } else { + use crate::file::metadata::thrift::PageHeader; + + Ok(PageHeader::read_thrift_without_stats(&mut prot)?) + } + } + Some(page_crypto_context) => { + let data_decryptor = page_crypto_context.data_decryptor(); + let aad = page_crypto_context.create_page_header_aad()?; + + let buf = read_and_decrypt(data_decryptor, input, aad.as_ref()).map_err(|_| { + ParquetError::General(format!( + "Error decrypting page header for column {}, decryption key may be wrong", + page_crypto_context.column_ordinal + )) + })?; + + let mut prot = ThriftSliceInputProtocol::new(buf.as_slice()); + if self.read_stats { + Ok(PageHeader::read_thrift(&mut prot)?) + } else { + Ok(PageHeader::read_thrift_without_stats(&mut prot)?) + } + } + } + } + + fn decrypt_page_data(&self, buffer: T, page_index: usize, dictionary_page: bool) -> Result + where + T: AsRef<[u8]>, + T: From>, + { + let page_crypto_context = self.page_crypto_context(page_index, dictionary_page); + if let Some(page_crypto_context) = page_crypto_context { + let decryptor = page_crypto_context.data_decryptor(); + let aad = page_crypto_context.create_page_aad()?; + let decrypted = decryptor.decrypt(buffer.as_ref(), &aad)?; + Ok(T::from(decrypted)) + } else { + Ok(buffer) + } + } + + fn page_crypto_context( + &self, + page_index: usize, + dictionary_page: bool, + ) -> Option> { + self.crypto_context.as_ref().map(|c| { + Arc::new(if dictionary_page { + c.for_dictionary_page() + } else { + c.with_page_ordinal(page_index) + }) + }) + } +} + +impl Iterator for SerializedPageReader { + type Item = Result; + + fn next(&mut self) -> Option { + self.get_next_page().transpose() + } +} + +fn verify_page_header_len(header_len: usize, remaining_bytes: u64) -> Result<()> { + if header_len as u64 > remaining_bytes { + return Err(eof_err!("Invalid page header")); + } + Ok(()) +} + +fn verify_page_size( + compressed_size: i32, + uncompressed_size: i32, + remaining_bytes: u64, +) -> Result<()> { + // The page's compressed size should not exceed the remaining bytes that are + // available to read. The page's uncompressed size is the expected size + // after decompression, which can never be negative. + if compressed_size < 0 || compressed_size as u64 > remaining_bytes || uncompressed_size < 0 { + return Err(eof_err!("Invalid page header")); + } + Ok(()) +} + +impl PageReader for SerializedPageReader { + fn get_next_page(&mut self) -> Result> { + loop { + let page = match &mut self.state { + SerializedPageReaderState::Values { + offset, + remaining_bytes: remaining, + next_page_header, + page_index, + require_dictionary, + } => { + if *remaining == 0 { + return Ok(None); + } + + let mut read = self.reader.get_read(*offset)?; + let header = if let Some(header) = next_page_header.take() { + *header + } else { + let (header_len, header) = Self::read_page_header_len( + &self.context, + &mut read, + *page_index, + *require_dictionary, + )?; + verify_page_header_len(header_len, *remaining)?; + *offset += header_len as u64; + *remaining -= header_len as u64; + header + }; + verify_page_size( + header.compressed_page_size, + header.uncompressed_page_size, + *remaining, + )?; + let data_len = header.compressed_page_size as usize; + let data_start = *offset; + *offset += data_len as u64; + *remaining -= data_len as u64; + + if header.r#type == PageType::INDEX_PAGE { + continue; + } + + let buffer = self.reader.get_bytes(data_start, data_len)?; + + let buffer = + self.context + .decrypt_page_data(buffer, *page_index, *require_dictionary)?; + + let page = decode_page( + header, + buffer, + self.physical_type, + self.decompressor.as_mut(), + )?; + if page.is_data_page() { + *page_index += 1; + } else if page.is_dictionary_page() { + *require_dictionary = false; + } + page + } + SerializedPageReaderState::Pages { + page_locations, + dictionary_page, + page_index, + .. + } => { + let (front, is_dictionary_page) = match dictionary_page.take() { + Some(front) => (front, true), + None => match page_locations.pop_front() { + Some(front) => (front, false), + None => return Ok(None), + }, + }; + + let page_len = usize::try_from(front.compressed_page_size)?; + let buffer = self.reader.get_bytes(front.offset as u64, page_len)?; + + let (offset, header) = Self::read_page_header_len_from_bytes( + &self.context, + buffer.as_ref(), + *page_index, + is_dictionary_page, + )?; + let bytes = buffer.slice(offset..); + let bytes = + self.context + .decrypt_page_data(bytes, *page_index, is_dictionary_page)?; + + if !is_dictionary_page { + *page_index += 1; + } + decode_page( + header, + bytes, + self.physical_type, + self.decompressor.as_mut(), + )? + } + }; + + return Ok(Some(page)); + } + } + + fn peek_next_page(&mut self) -> Result> { + match &mut self.state { + SerializedPageReaderState::Values { + offset, + remaining_bytes, + next_page_header, + page_index, + require_dictionary, + } => { + loop { + if *remaining_bytes == 0 { + return Ok(None); + } + return if let Some(header) = next_page_header.as_ref() { + if let Ok(page_meta) = (&**header).try_into() { + Ok(Some(page_meta)) + } else { + // For unknown page type (e.g., INDEX_PAGE), skip and read next. + *next_page_header = None; + continue; + } + } else { + let mut read = self.reader.get_read(*offset)?; + let (header_len, header) = Self::read_page_header_len( + &self.context, + &mut read, + *page_index, + *require_dictionary, + )?; + verify_page_header_len(header_len, *remaining_bytes)?; + *offset += header_len as u64; + *remaining_bytes -= header_len as u64; + let page_meta = if let Ok(page_meta) = (&header).try_into() { + Ok(Some(page_meta)) + } else { + // For unknown page type (e.g., INDEX_PAGE), skip and read next. + continue; + }; + *next_page_header = Some(Box::new(header)); + page_meta + }; + } + } + SerializedPageReaderState::Pages { + page_locations, + dictionary_page, + total_rows, + page_index: _, + } => { + if dictionary_page.is_some() { + Ok(Some(PageMetadata { + num_rows: None, + num_levels: None, + is_dict: true, + })) + } else if let Some(page) = page_locations.front() { + let next_rows = page_locations + .get(1) + .map(|x| x.first_row_index as usize) + .unwrap_or(*total_rows); + + Ok(Some(PageMetadata { + num_rows: Some(next_rows - page.first_row_index as usize), + num_levels: None, + is_dict: false, + })) + } else { + Ok(None) + } + } + } + } + + fn skip_next_page(&mut self) -> Result<()> { + match &mut self.state { + SerializedPageReaderState::Values { + offset, + remaining_bytes, + next_page_header, + page_index, + require_dictionary, + } => { + if let Some(buffered_header) = next_page_header.take() { + verify_page_size( + buffered_header.compressed_page_size, + buffered_header.uncompressed_page_size, + *remaining_bytes, + )?; + // The next page header has already been peeked, so just advance the offset + *offset += buffered_header.compressed_page_size as u64; + *remaining_bytes -= buffered_header.compressed_page_size as u64; + } else { + let mut read = self.reader.get_read(*offset)?; + let (header_len, header) = Self::read_page_header_len( + &self.context, + &mut read, + *page_index, + *require_dictionary, + )?; + verify_page_header_len(header_len, *remaining_bytes)?; + verify_page_size( + header.compressed_page_size, + header.uncompressed_page_size, + *remaining_bytes, + )?; + let data_page_size = header.compressed_page_size as u64; + *offset += header_len as u64 + data_page_size; + *remaining_bytes -= header_len as u64 + data_page_size; + } + if *require_dictionary { + *require_dictionary = false; + } else { + *page_index += 1; + } + Ok(()) + } + SerializedPageReaderState::Pages { + page_locations, + dictionary_page, + page_index, + .. + } => { + if dictionary_page.is_some() { + // If a dictionary page exists, consume it by taking it (sets to None) + dictionary_page.take(); + } else { + // If no dictionary page exists, simply pop the data page from page_locations + if page_locations.pop_front().is_some() { + *page_index += 1; + } + } + + Ok(()) + } + } + } + + fn at_record_boundary(&mut self) -> Result { + match &mut self.state { + SerializedPageReaderState::Values { .. } => match self.peek_next_page()? { + None => Ok(true), + // V2 data pages must start at record boundaries per the parquet + // spec, so the current page ends at one. + Some(metadata) => Ok(metadata.num_rows.is_some()), + }, + SerializedPageReaderState::Pages { .. } => Ok(true), + } + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use bytes::Buf; + + use crate::file::page_index::column_index::{ + ByteArrayColumnIndex, ColumnIndexMetaData, PrimitiveColumnIndex, + }; + use crate::file::properties::{EnabledStatistics, WriterProperties}; + + use crate::basic::{self, BoundaryOrder, ColumnOrder, Encoding, SortOrder}; + use crate::column::reader::ColumnReader; + use crate::data_type::private::ParquetValueType; + use crate::data_type::{AsBytes, FixedLenByteArrayType, Int32Type}; + use crate::file::metadata::thrift::DataPageHeaderV2; + use crate::file::writer::SerializedFileWriter; + use crate::record::RowAccessor; + use crate::schema::parser::parse_message_type; + use crate::util::test_common::file_util::{get_test_file, get_test_path}; + + use super::*; + + #[test] + fn test_decode_page_invalid_offset() { + let page_header = PageHeader { + r#type: PageType::DATA_PAGE_V2, + uncompressed_page_size: 10, + compressed_page_size: 10, + data_page_header: None, + index_page_header: None, + dictionary_page_header: None, + crc: None, + data_page_header_v2: Some(DataPageHeaderV2 { + num_nulls: 0, + num_rows: 0, + num_values: 0, + encoding: Encoding::PLAIN, + definition_levels_byte_length: 11, + repetition_levels_byte_length: 0, + is_compressed: None, + statistics: None, + }), + }; + + let buffer = Bytes::new(); + let err = decode_page(page_header, buffer, Type::INT32, None).unwrap_err(); + assert!( + err.to_string() + .contains("DataPage v2 header contains implausible values") + ); + } + + #[test] + fn test_decode_unsupported_page() { + let mut page_header = PageHeader { + r#type: PageType::INDEX_PAGE, + uncompressed_page_size: 10, + compressed_page_size: 10, + data_page_header: None, + index_page_header: None, + dictionary_page_header: None, + crc: None, + data_page_header_v2: None, + }; + let buffer = Bytes::new(); + let err = decode_page(page_header.clone(), buffer.clone(), Type::INT32, None).unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: Page type INDEX_PAGE is not supported" + ); + + page_header.data_page_header_v2 = Some(DataPageHeaderV2 { + num_nulls: 0, + num_rows: 0, + num_values: 0, + encoding: Encoding::PLAIN, + definition_levels_byte_length: 11, + repetition_levels_byte_length: 0, + is_compressed: None, + statistics: None, + }); + let err = decode_page(page_header, buffer, Type::INT32, None).unwrap_err(); + assert!( + err.to_string() + .contains("DataPage v2 header contains implausible values") + ); + } + + #[test] + fn test_cursor_and_file_has_the_same_behaviour() { + let mut buf: Vec = Vec::new(); + get_test_file("alltypes_plain.parquet") + .read_to_end(&mut buf) + .unwrap(); + let cursor = Bytes::from(buf); + let read_from_cursor = SerializedFileReader::new(cursor).unwrap(); + + let test_file = get_test_file("alltypes_plain.parquet"); + let read_from_file = SerializedFileReader::new(test_file).unwrap(); + + let file_iter = read_from_file.get_row_iter(None).unwrap(); + let cursor_iter = read_from_cursor.get_row_iter(None).unwrap(); + + for (a, b) in file_iter.zip(cursor_iter) { + assert_eq!(a.unwrap(), b.unwrap()) + } + } + + #[test] + fn test_file_reader_try_from() { + // Valid file path + let test_file = get_test_file("alltypes_plain.parquet"); + let test_path_buf = get_test_path("alltypes_plain.parquet"); + let test_path = test_path_buf.as_path(); + let test_path_str = test_path.to_str().unwrap(); + + let reader = SerializedFileReader::try_from(test_file); + assert!(reader.is_ok()); + + let reader = SerializedFileReader::try_from(test_path); + assert!(reader.is_ok()); + + let reader = SerializedFileReader::try_from(test_path_str); + assert!(reader.is_ok()); + + let reader = SerializedFileReader::try_from(test_path_str.to_string()); + assert!(reader.is_ok()); + + // Invalid file path + let test_path = Path::new("invalid.parquet"); + let test_path_str = test_path.to_str().unwrap(); + + let reader = SerializedFileReader::try_from(test_path); + assert!(reader.is_err()); + + let reader = SerializedFileReader::try_from(test_path_str); + assert!(reader.is_err()); + + let reader = SerializedFileReader::try_from(test_path_str.to_string()); + assert!(reader.is_err()); + } + + #[test] + fn test_file_reader_into_iter() { + let path = get_test_path("alltypes_plain.parquet"); + let reader = SerializedFileReader::try_from(path.as_path()).unwrap(); + let iter = reader.into_iter(); + let values: Vec<_> = iter.flat_map(|x| x.unwrap().get_int(0)).collect(); + + assert_eq!(values, &[4, 5, 6, 7, 2, 3, 0, 1]); + } + + #[test] + fn test_file_reader_into_iter_project() { + let path = get_test_path("alltypes_plain.parquet"); + let reader = SerializedFileReader::try_from(path.as_path()).unwrap(); + let schema = "message schema { OPTIONAL INT32 id; }"; + let proj = parse_message_type(schema).ok(); + let iter = reader.into_iter().project(proj).unwrap(); + let values: Vec<_> = iter.flat_map(|x| x.unwrap().get_int(0)).collect(); + + assert_eq!(values, &[4, 5, 6, 7, 2, 3, 0, 1]); + } + + #[test] + fn test_reuse_file_chunk() { + // This test covers the case of maintaining the correct start position in a file + // stream for each column reader after initializing and moving to the next one + // (without necessarily reading the entire column). + let test_file = get_test_file("alltypes_plain.parquet"); + let reader = SerializedFileReader::new(test_file).unwrap(); + let row_group = reader.get_row_group(0).unwrap(); + + let mut page_readers = Vec::new(); + for i in 0..row_group.num_columns() { + page_readers.push(row_group.get_column_page_reader(i).unwrap()); + } + + // Now buffer each col reader, we do not expect any failures like: + // General("underlying Thrift error: end of file") + for mut page_reader in page_readers { + assert!(page_reader.get_next_page().is_ok()); + } + } + + #[test] + fn test_file_reader() { + let test_file = get_test_file("alltypes_plain.parquet"); + let reader_result = SerializedFileReader::new(test_file); + assert!(reader_result.is_ok()); + let reader = reader_result.unwrap(); + + // Test contents in Parquet metadata + let metadata = reader.metadata(); + assert_eq!(metadata.num_row_groups(), 1); + + // Test contents in file metadata + let file_metadata = metadata.file_metadata(); + assert!(file_metadata.created_by().is_some()); + assert_eq!( + file_metadata.created_by().unwrap(), + "impala version 1.3.0-INTERNAL (build 8a48ddb1eff84592b3fc06bc6f51ec120e1fffc9)" + ); + assert!(file_metadata.key_value_metadata().is_none()); + assert_eq!(file_metadata.num_rows(), 8); + assert_eq!(file_metadata.version(), 1); + assert_eq!(file_metadata.column_orders(), None); + + // Test contents in row group metadata + let row_group_metadata = metadata.row_group(0); + assert_eq!(row_group_metadata.num_columns(), 11); + assert_eq!(row_group_metadata.num_rows(), 8); + assert_eq!(row_group_metadata.total_byte_size(), 671); + // Check each column order + for i in 0..row_group_metadata.num_columns() { + assert_eq!(file_metadata.column_order(i), ColumnOrder::UNDEFINED); + } + + // Test row group reader + let row_group_reader_result = reader.get_row_group(0); + assert!(row_group_reader_result.is_ok()); + let row_group_reader: Box = row_group_reader_result.unwrap(); + assert_eq!( + row_group_reader.num_columns(), + row_group_metadata.num_columns() + ); + assert_eq!( + row_group_reader.metadata().total_byte_size(), + row_group_metadata.total_byte_size() + ); + + // Test page readers + // TODO: test for every column + let page_reader_0_result = row_group_reader.get_column_page_reader(0); + assert!(page_reader_0_result.is_ok()); + let mut page_reader_0: Box = page_reader_0_result.unwrap(); + let mut page_count = 0; + while let Some(page) = page_reader_0.get_next_page().unwrap() { + let is_expected_page = match page { + Page::DictionaryPage { + buf, + num_values, + encoding, + is_sorted, + } => { + assert_eq!(buf.len(), 32); + assert_eq!(num_values, 8); + assert_eq!(encoding, Encoding::PLAIN_DICTIONARY); + assert!(!is_sorted); + true + } + Page::DataPage { + buf, + num_values, + encoding, + def_level_encoding, + rep_level_encoding, + statistics, + } => { + assert_eq!(buf.len(), 11); + assert_eq!(num_values, 8); + assert_eq!(encoding, Encoding::PLAIN_DICTIONARY); + assert_eq!(def_level_encoding, Encoding::RLE); + #[allow(deprecated)] + let expected_rep_level_encoding = Encoding::BIT_PACKED; + assert_eq!(rep_level_encoding, expected_rep_level_encoding); + assert!(statistics.is_none()); + true + } + _ => false, + }; + assert!(is_expected_page); + page_count += 1; + } + assert_eq!(page_count, 2); + } + + #[test] + fn test_file_reader_datapage_v2() { + let test_file = get_test_file("datapage_v2.snappy.parquet"); + let reader_result = SerializedFileReader::new(test_file); + assert!(reader_result.is_ok()); + let reader = reader_result.unwrap(); + + // Test contents in Parquet metadata + let metadata = reader.metadata(); + assert_eq!(metadata.num_row_groups(), 1); + + // Test contents in file metadata + let file_metadata = metadata.file_metadata(); + assert!(file_metadata.created_by().is_some()); + assert_eq!( + file_metadata.created_by().unwrap(), + "parquet-mr version 1.8.1 (build 4aba4dae7bb0d4edbcf7923ae1339f28fd3f7fcf)" + ); + assert!(file_metadata.key_value_metadata().is_some()); + assert_eq!( + file_metadata.key_value_metadata().to_owned().unwrap().len(), + 1 + ); + + assert_eq!(file_metadata.num_rows(), 5); + assert_eq!(file_metadata.version(), 1); + assert_eq!(file_metadata.column_orders(), None); + + let row_group_metadata = metadata.row_group(0); + + // Check each column order + for i in 0..row_group_metadata.num_columns() { + assert_eq!(file_metadata.column_order(i), ColumnOrder::UNDEFINED); + } + + // Test row group reader + let row_group_reader_result = reader.get_row_group(0); + assert!(row_group_reader_result.is_ok()); + let row_group_reader: Box = row_group_reader_result.unwrap(); + assert_eq!( + row_group_reader.num_columns(), + row_group_metadata.num_columns() + ); + assert_eq!( + row_group_reader.metadata().total_byte_size(), + row_group_metadata.total_byte_size() + ); + + // Test page readers + // TODO: test for every column + let page_reader_0_result = row_group_reader.get_column_page_reader(0); + assert!(page_reader_0_result.is_ok()); + let mut page_reader_0: Box = page_reader_0_result.unwrap(); + let mut page_count = 0; + while let Some(page) = page_reader_0.get_next_page().unwrap() { + let is_expected_page = match page { + Page::DictionaryPage { + buf, + num_values, + encoding, + is_sorted, + } => { + assert_eq!(buf.len(), 7); + assert_eq!(num_values, 1); + assert_eq!(encoding, Encoding::PLAIN); + assert!(!is_sorted); + true + } + Page::DataPageV2 { + buf, + num_values, + encoding, + num_nulls, + num_rows, + def_levels_byte_len, + rep_levels_byte_len, + is_compressed, + statistics, + } => { + assert_eq!(buf.len(), 4); + assert_eq!(num_values, 5); + assert_eq!(encoding, Encoding::RLE_DICTIONARY); + assert_eq!(num_nulls, 1); + assert_eq!(num_rows, 5); + assert_eq!(def_levels_byte_len, 2); + assert_eq!(rep_levels_byte_len, 0); + assert!(is_compressed); + assert!(statistics.is_none()); // page stats are no longer read + true + } + _ => false, + }; + assert!(is_expected_page); + page_count += 1; + } + assert_eq!(page_count, 2); + } + + #[test] + fn test_file_reader_empty_compressed_datapage_v2() { + // this file has a compressed datapage that un-compresses to 0 bytes + let test_file = get_test_file("page_v2_empty_compressed.parquet"); + let reader_result = SerializedFileReader::new(test_file); + assert!(reader_result.is_ok()); + let reader = reader_result.unwrap(); + + // Test contents in Parquet metadata + let metadata = reader.metadata(); + assert_eq!(metadata.num_row_groups(), 1); + + // Test contents in file metadata + let file_metadata = metadata.file_metadata(); + assert!(file_metadata.created_by().is_some()); + assert_eq!( + file_metadata.created_by().unwrap(), + "parquet-cpp-arrow version 14.0.2" + ); + assert!(file_metadata.key_value_metadata().is_some()); + assert_eq!( + file_metadata.key_value_metadata().to_owned().unwrap().len(), + 1 + ); + + assert_eq!(file_metadata.num_rows(), 10); + assert_eq!(file_metadata.version(), 2); + let expected_order = ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::SIGNED); + assert_eq!( + file_metadata.column_orders(), + Some(vec![expected_order].as_ref()) + ); + + let row_group_metadata = metadata.row_group(0); + + // Check each column order + for i in 0..row_group_metadata.num_columns() { + assert_eq!(file_metadata.column_order(i), expected_order); + } + + // Test row group reader + let row_group_reader_result = reader.get_row_group(0); + assert!(row_group_reader_result.is_ok()); + let row_group_reader: Box = row_group_reader_result.unwrap(); + assert_eq!( + row_group_reader.num_columns(), + row_group_metadata.num_columns() + ); + assert_eq!( + row_group_reader.metadata().total_byte_size(), + row_group_metadata.total_byte_size() + ); + + // Test page readers + let page_reader_0_result = row_group_reader.get_column_page_reader(0); + assert!(page_reader_0_result.is_ok()); + let mut page_reader_0: Box = page_reader_0_result.unwrap(); + let mut page_count = 0; + while let Some(page) = page_reader_0.get_next_page().unwrap() { + let is_expected_page = match page { + Page::DictionaryPage { + buf, + num_values, + encoding, + is_sorted, + } => { + assert_eq!(buf.len(), 0); + assert_eq!(num_values, 0); + assert_eq!(encoding, Encoding::PLAIN); + assert!(!is_sorted); + true + } + Page::DataPageV2 { + buf, + num_values, + encoding, + num_nulls, + num_rows, + def_levels_byte_len, + rep_levels_byte_len, + is_compressed, + statistics, + } => { + assert_eq!(buf.len(), 3); + assert_eq!(num_values, 10); + assert_eq!(encoding, Encoding::RLE_DICTIONARY); + assert_eq!(num_nulls, 10); + assert_eq!(num_rows, 10); + assert_eq!(def_levels_byte_len, 2); + assert_eq!(rep_levels_byte_len, 0); + assert!(is_compressed); + assert!(statistics.is_none()); // page stats are no longer read + true + } + _ => false, + }; + assert!(is_expected_page); + page_count += 1; + } + assert_eq!(page_count, 2); + } + + #[test] + fn test_file_reader_empty_datapage_v2() { + // this file has 0 bytes compressed datapage that un-compresses to 0 bytes + let test_file = get_test_file("datapage_v2_empty_datapage.snappy.parquet"); + let reader_result = SerializedFileReader::new(test_file); + assert!(reader_result.is_ok()); + let reader = reader_result.unwrap(); + + // Test contents in Parquet metadata + let metadata = reader.metadata(); + assert_eq!(metadata.num_row_groups(), 1); + + // Test contents in file metadata + let file_metadata = metadata.file_metadata(); + assert!(file_metadata.created_by().is_some()); + assert_eq!( + file_metadata.created_by().unwrap(), + "parquet-mr version 1.13.1 (build db4183109d5b734ec5930d870cdae161e408ddba)" + ); + assert!(file_metadata.key_value_metadata().is_some()); + assert_eq!( + file_metadata.key_value_metadata().to_owned().unwrap().len(), + 2 + ); + + assert_eq!(file_metadata.num_rows(), 1); + assert_eq!(file_metadata.version(), 1); + let expected_order = ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::SIGNED); + assert_eq!( + file_metadata.column_orders(), + Some(vec![expected_order].as_ref()) + ); + + let row_group_metadata = metadata.row_group(0); + + // Check each column order + for i in 0..row_group_metadata.num_columns() { + assert_eq!(file_metadata.column_order(i), expected_order); + } + + // Test row group reader + let row_group_reader_result = reader.get_row_group(0); + assert!(row_group_reader_result.is_ok()); + let row_group_reader: Box = row_group_reader_result.unwrap(); + assert_eq!( + row_group_reader.num_columns(), + row_group_metadata.num_columns() + ); + assert_eq!( + row_group_reader.metadata().total_byte_size(), + row_group_metadata.total_byte_size() + ); + + // Test page readers + let page_reader_0_result = row_group_reader.get_column_page_reader(0); + assert!(page_reader_0_result.is_ok()); + let mut page_reader_0: Box = page_reader_0_result.unwrap(); + let mut page_count = 0; + while let Some(page) = page_reader_0.get_next_page().unwrap() { + let is_expected_page = match page { + Page::DataPageV2 { + buf, + num_values, + encoding, + num_nulls, + num_rows, + def_levels_byte_len, + rep_levels_byte_len, + is_compressed, + statistics, + } => { + assert_eq!(buf.len(), 2); + assert_eq!(num_values, 1); + assert_eq!(encoding, Encoding::PLAIN); + assert_eq!(num_nulls, 1); + assert_eq!(num_rows, 1); + assert_eq!(def_levels_byte_len, 2); + assert_eq!(rep_levels_byte_len, 0); + assert!(is_compressed); + assert!(statistics.is_none()); + true + } + _ => false, + }; + assert!(is_expected_page); + page_count += 1; + } + assert_eq!(page_count, 1); + } + + fn get_serialized_page_reader( + file_reader: &SerializedFileReader, + row_group: usize, + column: usize, + ) -> Result> { + let row_group = { + let row_group_metadata = file_reader.metadata.row_group(row_group); + let props = Arc::clone(&file_reader.props); + let f = Arc::clone(&file_reader.chunk_reader); + SerializedRowGroupReader::new( + f, + row_group_metadata, + file_reader + .metadata + .offset_index() + .map(|x| x[row_group].as_slice()), + props, + )? + }; + + let col = row_group.metadata.column(column); + + let page_locations = row_group + .offset_index + .map(|x| x[column].page_locations.clone()); + + let props = Arc::clone(&row_group.props); + SerializedPageReader::new_with_properties( + Arc::clone(&row_group.chunk_reader), + col, + usize::try_from(row_group.metadata.num_rows())?, + page_locations, + props, + ) + } + + #[test] + fn test_peek_next_page_offset_matches_actual() -> Result<()> { + let test_file = get_test_file("alltypes_plain.parquet"); + let reader = SerializedFileReader::new(test_file)?; + + let mut offset_set = HashSet::new(); + let num_row_groups = reader.metadata.num_row_groups(); + for row_group in 0..num_row_groups { + let num_columns = reader.metadata.row_group(row_group).num_columns(); + for column in 0..num_columns { + let mut page_reader = get_serialized_page_reader(&reader, row_group, column)?; + + while let Ok(Some(page_offset)) = page_reader.peek_next_page_offset() { + match &page_reader.state { + SerializedPageReaderState::Pages { + page_locations, + dictionary_page, + .. + } => { + if let Some(page) = dictionary_page { + assert_eq!(page.offset as u64, page_offset); + } else if let Some(page) = page_locations.front() { + assert_eq!(page.offset as u64, page_offset); + } else { + unreachable!() + } + } + SerializedPageReaderState::Values { + offset, + next_page_header, + .. + } => { + assert!(next_page_header.is_some()); + assert_eq!(*offset, page_offset); + } + } + let page = page_reader.get_next_page()?; + assert!(page.is_some()); + let newly_inserted = offset_set.insert(page_offset); + assert!(newly_inserted); + } + } + } + + Ok(()) + } + + #[test] + fn test_page_iterator() { + let file = get_test_file("alltypes_plain.parquet"); + let file_reader = Arc::new(SerializedFileReader::new(file).unwrap()); + + let mut page_iterator = FilePageIterator::new(0, file_reader.clone()).unwrap(); + + // read first page + let page = page_iterator.next(); + assert!(page.is_some()); + assert!(page.unwrap().is_ok()); + + // reach end of file + let page = page_iterator.next(); + assert!(page.is_none()); + + let row_group_indices = Box::new(0..1); + let mut page_iterator = + FilePageIterator::with_row_groups(0, row_group_indices, file_reader).unwrap(); + + // read first page + let page = page_iterator.next(); + assert!(page.is_some()); + assert!(page.unwrap().is_ok()); + + // reach end of file + let page = page_iterator.next(); + assert!(page.is_none()); + } + + #[test] + fn test_file_reader_key_value_metadata() { + let file = get_test_file("binary.parquet"); + let file_reader = Arc::new(SerializedFileReader::new(file).unwrap()); + + let metadata = file_reader + .metadata + .file_metadata() + .key_value_metadata() + .unwrap(); + + assert_eq!(metadata.len(), 3); + + assert_eq!(metadata[0].key, "parquet.proto.descriptor"); + + assert_eq!(metadata[1].key, "writer.model.name"); + assert_eq!(metadata[1].value, Some("protobuf".to_owned())); + + assert_eq!(metadata[2].key, "parquet.proto.class"); + assert_eq!(metadata[2].value, Some("foo.baz.Foobaz$Event".to_owned())); + } + + #[test] + fn test_file_reader_optional_metadata() { + // file with optional metadata: bloom filters, encoding stats, column index and offset index. + let file = get_test_file("data_index_bloom_encoding_stats.parquet"); + let options = ReadOptionsBuilder::new() + .with_encoding_stats_as_mask(false) + .build(); + let file_reader = Arc::new(SerializedFileReader::new_with_options(file, options).unwrap()); + + let row_group_metadata = file_reader.metadata.row_group(0); + let col0_metadata = row_group_metadata.column(0); + + // test optional bloom filter offset + assert_eq!(col0_metadata.bloom_filter_offset().unwrap(), 192); + + // test page encoding stats + let page_encoding_stats = &col0_metadata.page_encoding_stats().unwrap()[0]; + + assert_eq!(page_encoding_stats.page_type, basic::PageType::DATA_PAGE); + assert_eq!(page_encoding_stats.encoding, Encoding::PLAIN); + assert_eq!(page_encoding_stats.count, 1); + + // test optional column index offset + assert_eq!(col0_metadata.column_index_offset().unwrap(), 156); + assert_eq!(col0_metadata.column_index_length().unwrap(), 25); + + // test optional offset index offset + assert_eq!(col0_metadata.offset_index_offset().unwrap(), 181); + assert_eq!(col0_metadata.offset_index_length().unwrap(), 11); + } + + #[test] + fn test_file_reader_page_stats_mask() { + let file = get_test_file("alltypes_tiny_pages.parquet"); + let options = ReadOptionsBuilder::new() + .with_encoding_stats_as_mask(true) + .build(); + let file_reader = Arc::new(SerializedFileReader::new_with_options(file, options).unwrap()); + + let row_group_metadata = file_reader.metadata.row_group(0); + + // test page encoding stats + let page_encoding_stats = row_group_metadata + .column(0) + .page_encoding_stats_mask() + .unwrap(); + assert!(page_encoding_stats.is_only(Encoding::PLAIN)); + let page_encoding_stats = row_group_metadata + .column(2) + .page_encoding_stats_mask() + .unwrap(); + assert!(page_encoding_stats.is_only(Encoding::PLAIN_DICTIONARY)); + } + + #[test] + fn test_file_reader_page_stats_skipped() { + let file = get_test_file("alltypes_tiny_pages.parquet"); + + // test skipping all + let options = ReadOptionsBuilder::new() + .with_encoding_stats_policy(ParquetStatisticsPolicy::SkipAll) + .with_column_stats_policy(ParquetStatisticsPolicy::SkipAll) + .build(); + let file_reader = Arc::new( + SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap(), + ); + + let row_group_metadata = file_reader.metadata.row_group(0); + for column in row_group_metadata.columns() { + assert!(column.page_encoding_stats().is_none()); + assert!(column.page_encoding_stats_mask().is_none()); + assert!(column.statistics().is_none()); + } + + // test skipping all but one column + let options = ReadOptionsBuilder::new() + .with_encoding_stats_as_mask(true) + .with_encoding_stats_policy(ParquetStatisticsPolicy::skip_except(&[0])) + .with_column_stats_policy(ParquetStatisticsPolicy::skip_except(&[0])) + .build(); + let file_reader = Arc::new( + SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap(), + ); + + let row_group_metadata = file_reader.metadata.row_group(0); + for (idx, column) in row_group_metadata.columns().iter().enumerate() { + assert!(column.page_encoding_stats().is_none()); + assert_eq!(column.page_encoding_stats_mask().is_some(), idx == 0); + assert_eq!(column.statistics().is_some(), idx == 0); + } + } + + #[test] + fn test_file_reader_size_stats_skipped() { + let file = get_test_file("repeated_primitive_no_list.parquet"); + + // test skipping all + let options = ReadOptionsBuilder::new() + .with_size_stats_policy(ParquetStatisticsPolicy::SkipAll) + .build(); + let file_reader = Arc::new( + SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap(), + ); + + let row_group_metadata = file_reader.metadata.row_group(0); + for column in row_group_metadata.columns() { + assert!(column.repetition_level_histogram().is_none()); + assert!(column.definition_level_histogram().is_none()); + assert!(column.unencoded_byte_array_data_bytes().is_none()); + } + + // test skipping all but one column + let options = ReadOptionsBuilder::new() + .with_encoding_stats_as_mask(true) + .with_size_stats_policy(ParquetStatisticsPolicy::skip_except(&[1])) + .build(); + let file_reader = Arc::new( + SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap(), + ); + + let row_group_metadata = file_reader.metadata.row_group(0); + for (idx, column) in row_group_metadata.columns().iter().enumerate() { + assert_eq!(column.repetition_level_histogram().is_some(), idx == 1); + assert_eq!(column.definition_level_histogram().is_some(), idx == 1); + assert_eq!(column.unencoded_byte_array_data_bytes().is_some(), idx == 1); + } + } + + #[test] + fn test_file_reader_with_no_filter() -> Result<()> { + let test_file = get_test_file("alltypes_plain.parquet"); + let origin_reader = SerializedFileReader::new(test_file)?; + // test initial number of row groups + let metadata = origin_reader.metadata(); + assert_eq!(metadata.num_row_groups(), 1); + Ok(()) + } + + #[test] + fn test_file_reader_filter_row_groups_with_predicate() -> Result<()> { + let test_file = get_test_file("alltypes_plain.parquet"); + let read_options = ReadOptionsBuilder::new() + .with_predicate(Box::new(|_, _| false)) + .build(); + let reader = SerializedFileReader::new_with_options(test_file, read_options)?; + let metadata = reader.metadata(); + assert_eq!(metadata.num_row_groups(), 0); + Ok(()) + } + + #[test] + fn test_file_reader_filter_row_groups_with_range() -> Result<()> { + let test_file = get_test_file("alltypes_plain.parquet"); + let origin_reader = SerializedFileReader::new(test_file)?; + // test initial number of row groups + let metadata = origin_reader.metadata(); + assert_eq!(metadata.num_row_groups(), 1); + let mid = get_midpoint_offset(metadata.row_group(0)); + + let test_file = get_test_file("alltypes_plain.parquet"); + let read_options = ReadOptionsBuilder::new().with_range(0, mid + 1).build(); + let reader = SerializedFileReader::new_with_options(test_file, read_options)?; + let metadata = reader.metadata(); + assert_eq!(metadata.num_row_groups(), 1); + + let test_file = get_test_file("alltypes_plain.parquet"); + let read_options = ReadOptionsBuilder::new().with_range(0, mid).build(); + let reader = SerializedFileReader::new_with_options(test_file, read_options)?; + let metadata = reader.metadata(); + assert_eq!(metadata.num_row_groups(), 0); + Ok(()) + } + + #[test] + fn test_file_reader_filter_row_groups_and_range() -> Result<()> { + let test_file = get_test_file("alltypes_tiny_pages.parquet"); + let origin_reader = SerializedFileReader::new(test_file)?; + let metadata = origin_reader.metadata(); + let mid = get_midpoint_offset(metadata.row_group(0)); + + // true, true predicate + let test_file = get_test_file("alltypes_tiny_pages.parquet"); + let read_options = ReadOptionsBuilder::new() + .with_page_index() + .with_predicate(Box::new(|_, _| true)) + .with_range(mid, mid + 1) + .build(); + let reader = SerializedFileReader::new_with_options(test_file, read_options)?; + let metadata = reader.metadata(); + assert_eq!(metadata.num_row_groups(), 1); + assert_eq!(metadata.column_index().unwrap().len(), 1); + assert_eq!(metadata.offset_index().unwrap().len(), 1); + + // true, false predicate + let test_file = get_test_file("alltypes_tiny_pages.parquet"); + let read_options = ReadOptionsBuilder::new() + .with_page_index() + .with_predicate(Box::new(|_, _| true)) + .with_range(0, mid) + .build(); + let reader = SerializedFileReader::new_with_options(test_file, read_options)?; + let metadata = reader.metadata(); + assert_eq!(metadata.num_row_groups(), 0); + assert!(metadata.column_index().is_none()); + assert!(metadata.offset_index().is_none()); + + // false, true predicate + let test_file = get_test_file("alltypes_tiny_pages.parquet"); + let read_options = ReadOptionsBuilder::new() + .with_page_index() + .with_predicate(Box::new(|_, _| false)) + .with_range(mid, mid + 1) + .build(); + let reader = SerializedFileReader::new_with_options(test_file, read_options)?; + let metadata = reader.metadata(); + assert_eq!(metadata.num_row_groups(), 0); + assert!(metadata.column_index().is_none()); + assert!(metadata.offset_index().is_none()); + + // false, false predicate + let test_file = get_test_file("alltypes_tiny_pages.parquet"); + let read_options = ReadOptionsBuilder::new() + .with_page_index() + .with_predicate(Box::new(|_, _| false)) + .with_range(0, mid) + .build(); + let reader = SerializedFileReader::new_with_options(test_file, read_options)?; + let metadata = reader.metadata(); + assert_eq!(metadata.num_row_groups(), 0); + assert!(metadata.column_index().is_none()); + assert!(metadata.offset_index().is_none()); + Ok(()) + } + + #[test] + fn test_file_reader_invalid_metadata() { + let data = [ + 255, 172, 1, 0, 50, 82, 65, 73, 1, 0, 0, 0, 169, 168, 168, 162, 87, 255, 16, 0, 0, 0, + 80, 65, 82, 49, + ]; + let ret = SerializedFileReader::new(Bytes::copy_from_slice(&data)); + assert_eq!( + ret.err().unwrap().to_string(), + "Parquet error: Expected list element type of Struct but got List" + ); + } + + #[test] + // Use java parquet-tools get below pageIndex info + // !``` + // parquet-tools column-index ./data_index_bloom_encoding_stats.parquet + // row group 0: + // column index for column String: + // Boundary order: ASCENDING + // page-0 : + // null count min max + // 0 Hello today + // + // offset index for column String: + // page-0 : + // offset compressed size first row index + // 4 152 0 + ///``` + // + fn test_page_index_reader() { + let test_file = get_test_file("data_index_bloom_encoding_stats.parquet"); + let builder = ReadOptionsBuilder::new(); + //enable read page index + let options = builder.with_page_index().build(); + let reader_result = SerializedFileReader::new_with_options(test_file, options); + let reader = reader_result.unwrap(); + + // Test contents in Parquet metadata + let metadata = reader.metadata(); + assert_eq!(metadata.num_row_groups(), 1); + + let column_index = metadata.column_index().unwrap(); + + // only one row group + assert_eq!(column_index.len(), 1); + let index = if let ColumnIndexMetaData::BYTE_ARRAY(index) = &column_index[0][0] { + index + } else { + unreachable!() + }; + + assert_eq!(index.boundary_order, BoundaryOrder::ASCENDING); + + //only one page group + assert_eq!(index.num_pages(), 1); + + let min = index.min_value(0).unwrap(); + let max = index.max_value(0).unwrap(); + assert_eq!(b"Hello", min.as_bytes()); + assert_eq!(b"today", max.as_bytes()); + + let offset_indexes = metadata.offset_index().unwrap(); + // only one row group + assert_eq!(offset_indexes.len(), 1); + let offset_index = &offset_indexes[0]; + let page_offset = &offset_index[0].page_locations()[0]; + + assert_eq!(4, page_offset.offset); + assert_eq!(152, page_offset.compressed_page_size); + assert_eq!(0, page_offset.first_row_index); + } + + #[test] + fn test_page_index_reader_all_type() { + let test_file = get_test_file("alltypes_tiny_pages_plain.parquet"); + let builder = ReadOptionsBuilder::new(); + //enable read page index + let options = builder.with_page_index().build(); + let reader_result = SerializedFileReader::new_with_options(test_file, options); + let reader = reader_result.unwrap(); + + // Test contents in Parquet metadata + let metadata = reader.metadata(); + assert_eq!(metadata.num_row_groups(), 1); + + let column_index = metadata.column_index().unwrap(); + let row_group_offset_indexes = &metadata.offset_index().unwrap()[0]; + + // only one row group + assert_eq!(column_index.len(), 1); + let row_group_metadata = metadata.row_group(0); + + //col0->id: INT32 UNCOMPRESSED DO:0 FPO:4 SZ:37325/37325/1.00 VC:7300 ENC:BIT_PACKED,RLE,PLAIN ST:[min: 0, max: 7299, num_nulls: 0] + assert!(!&column_index[0][0].is_sorted()); + let boundary_order = &column_index[0][0].get_boundary_order(); + assert!(boundary_order.is_some()); + matches!(boundary_order.unwrap(), BoundaryOrder::UNORDERED); + if let ColumnIndexMetaData::INT32(index) = &column_index[0][0] { + check_native_page_index( + index, + 325, + get_row_group_min_max_bytes(row_group_metadata, 0), + BoundaryOrder::UNORDERED, + ); + assert_eq!(row_group_offset_indexes[0].page_locations.len(), 325); + } else { + unreachable!() + }; + //col1->bool_col:BOOLEAN UNCOMPRESSED DO:0 FPO:37329 SZ:3022/3022/1.00 VC:7300 ENC:BIT_PACKED,RLE,PLAIN ST:[min: false, max: true, num_nulls: 0] + assert!(&column_index[0][1].is_sorted()); + if let ColumnIndexMetaData::BOOLEAN(index) = &column_index[0][1] { + assert_eq!(index.num_pages(), 82); + assert_eq!(row_group_offset_indexes[1].page_locations.len(), 82); + } else { + unreachable!() + }; + //col2->tinyint_col: INT32 UNCOMPRESSED DO:0 FPO:40351 SZ:37325/37325/1.00 VC:7300 ENC:BIT_PACKED,RLE,PLAIN ST:[min: 0, max: 9, num_nulls: 0] + assert!(&column_index[0][2].is_sorted()); + if let ColumnIndexMetaData::INT32(index) = &column_index[0][2] { + check_native_page_index( + index, + 325, + get_row_group_min_max_bytes(row_group_metadata, 2), + BoundaryOrder::ASCENDING, + ); + assert_eq!(row_group_offset_indexes[2].page_locations.len(), 325); + } else { + unreachable!() + }; + //col4->smallint_col: INT32 UNCOMPRESSED DO:0 FPO:77676 SZ:37325/37325/1.00 VC:7300 ENC:BIT_PACKED,RLE,PLAIN ST:[min: 0, max: 9, num_nulls: 0] + assert!(&column_index[0][3].is_sorted()); + if let ColumnIndexMetaData::INT32(index) = &column_index[0][3] { + check_native_page_index( + index, + 325, + get_row_group_min_max_bytes(row_group_metadata, 3), + BoundaryOrder::ASCENDING, + ); + assert_eq!(row_group_offset_indexes[3].page_locations.len(), 325); + } else { + unreachable!() + }; + //col5->smallint_col: INT32 UNCOMPRESSED DO:0 FPO:77676 SZ:37325/37325/1.00 VC:7300 ENC:BIT_PACKED,RLE,PLAIN ST:[min: 0, max: 9, num_nulls: 0] + assert!(&column_index[0][4].is_sorted()); + if let ColumnIndexMetaData::INT32(index) = &column_index[0][4] { + check_native_page_index( + index, + 325, + get_row_group_min_max_bytes(row_group_metadata, 4), + BoundaryOrder::ASCENDING, + ); + assert_eq!(row_group_offset_indexes[4].page_locations.len(), 325); + } else { + unreachable!() + }; + //col6->bigint_col: INT64 UNCOMPRESSED DO:0 FPO:152326 SZ:71598/71598/1.00 VC:7300 ENC:BIT_PACKED,RLE,PLAIN ST:[min: 0, max: 90, num_nulls: 0] + assert!(!&column_index[0][5].is_sorted()); + if let ColumnIndexMetaData::INT64(index) = &column_index[0][5] { + check_native_page_index( + index, + 528, + get_row_group_min_max_bytes(row_group_metadata, 5), + BoundaryOrder::UNORDERED, + ); + assert_eq!(row_group_offset_indexes[5].page_locations.len(), 528); + } else { + unreachable!() + }; + //col7->float_col: FLOAT UNCOMPRESSED DO:0 FPO:223924 SZ:37325/37325/1.00 VC:7300 ENC:BIT_PACKED,RLE,PLAIN ST:[min: -0.0, max: 9.9, num_nulls: 0] + assert!(&column_index[0][6].is_sorted()); + if let ColumnIndexMetaData::FLOAT(index) = &column_index[0][6] { + check_native_page_index( + index, + 325, + get_row_group_min_max_bytes(row_group_metadata, 6), + BoundaryOrder::ASCENDING, + ); + assert_eq!(row_group_offset_indexes[6].page_locations.len(), 325); + } else { + unreachable!() + }; + //col8->double_col: DOUBLE UNCOMPRESSED DO:0 FPO:261249 SZ:71598/71598/1.00 VC:7300 ENC:BIT_PACKED,RLE,PLAIN ST:[min: -0.0, max: 90.89999999999999, num_nulls: 0] + assert!(!&column_index[0][7].is_sorted()); + if let ColumnIndexMetaData::DOUBLE(index) = &column_index[0][7] { + check_native_page_index( + index, + 528, + get_row_group_min_max_bytes(row_group_metadata, 7), + BoundaryOrder::UNORDERED, + ); + assert_eq!(row_group_offset_indexes[7].page_locations.len(), 528); + } else { + unreachable!() + }; + //col9->date_string_col: BINARY UNCOMPRESSED DO:0 FPO:332847 SZ:111948/111948/1.00 VC:7300 ENC:BIT_PACKED,RLE,PLAIN ST:[min: 01/01/09, max: 12/31/10, num_nulls: 0] + assert!(!&column_index[0][8].is_sorted()); + if let ColumnIndexMetaData::BYTE_ARRAY(index) = &column_index[0][8] { + check_byte_array_page_index( + index, + 974, + get_row_group_min_max_bytes(row_group_metadata, 8), + BoundaryOrder::UNORDERED, + ); + assert_eq!(row_group_offset_indexes[8].page_locations.len(), 974); + } else { + unreachable!() + }; + //col10->string_col: BINARY UNCOMPRESSED DO:0 FPO:444795 SZ:45298/45298/1.00 VC:7300 ENC:BIT_PACKED,RLE,PLAIN ST:[min: 0, max: 9, num_nulls: 0] + assert!(&column_index[0][9].is_sorted()); + if let ColumnIndexMetaData::BYTE_ARRAY(index) = &column_index[0][9] { + check_byte_array_page_index( + index, + 352, + get_row_group_min_max_bytes(row_group_metadata, 9), + BoundaryOrder::ASCENDING, + ); + assert_eq!(row_group_offset_indexes[9].page_locations.len(), 352); + } else { + unreachable!() + }; + //col11->timestamp_col: INT96 UNCOMPRESSED DO:0 FPO:490093 SZ:111948/111948/1.00 VC:7300 ENC:BIT_PACKED,RLE,PLAIN ST:[num_nulls: 0, min/max not defined] + //Notice: min_max values for each page for this col not exits. + assert!(!&column_index[0][10].is_sorted()); + if let ColumnIndexMetaData::NONE = &column_index[0][10] { + assert_eq!(row_group_offset_indexes[10].page_locations.len(), 974); + } else { + unreachable!() + }; + //col12->year: INT32 UNCOMPRESSED DO:0 FPO:602041 SZ:37325/37325/1.00 VC:7300 ENC:BIT_PACKED,RLE,PLAIN ST:[min: 2009, max: 2010, num_nulls: 0] + assert!(&column_index[0][11].is_sorted()); + if let ColumnIndexMetaData::INT32(index) = &column_index[0][11] { + check_native_page_index( + index, + 325, + get_row_group_min_max_bytes(row_group_metadata, 11), + BoundaryOrder::ASCENDING, + ); + assert_eq!(row_group_offset_indexes[11].page_locations.len(), 325); + } else { + unreachable!() + }; + //col13->month: INT32 UNCOMPRESSED DO:0 FPO:639366 SZ:37325/37325/1.00 VC:7300 ENC:BIT_PACKED,RLE,PLAIN ST:[min: 1, max: 12, num_nulls: 0] + assert!(!&column_index[0][12].is_sorted()); + if let ColumnIndexMetaData::INT32(index) = &column_index[0][12] { + check_native_page_index( + index, + 325, + get_row_group_min_max_bytes(row_group_metadata, 12), + BoundaryOrder::UNORDERED, + ); + assert_eq!(row_group_offset_indexes[12].page_locations.len(), 325); + } else { + unreachable!() + }; + } + + fn check_native_page_index( + row_group_index: &PrimitiveColumnIndex, + page_size: usize, + min_max: (&[u8], &[u8]), + boundary_order: BoundaryOrder, + ) { + assert_eq!(row_group_index.num_pages() as usize, page_size); + assert_eq!(row_group_index.boundary_order, boundary_order); + assert!(row_group_index.min_values().iter().all(|x| { + x >= &T::try_from_le_slice(min_max.0).unwrap() + && x <= &T::try_from_le_slice(min_max.1).unwrap() + })); + } + + fn check_byte_array_page_index( + row_group_index: &ByteArrayColumnIndex, + page_size: usize, + min_max: (&[u8], &[u8]), + boundary_order: BoundaryOrder, + ) { + assert_eq!(row_group_index.num_pages() as usize, page_size); + assert_eq!(row_group_index.boundary_order, boundary_order); + for i in 0..row_group_index.num_pages() as usize { + let x = row_group_index.min_value(i).unwrap(); + assert!(x >= min_max.0 && x <= min_max.1); + } + } + + fn get_row_group_min_max_bytes(r: &RowGroupMetaData, col_num: usize) -> (&[u8], &[u8]) { + let statistics = r.column(col_num).statistics().unwrap(); + ( + statistics.min_bytes_opt().unwrap_or_default(), + statistics.max_bytes_opt().unwrap_or_default(), + ) + } + + #[test] + fn test_skip_next_page_with_dictionary_page() { + let test_file = get_test_file("alltypes_tiny_pages.parquet"); + let builder = ReadOptionsBuilder::new(); + // enable read page index + let options = builder.with_page_index().build(); + let reader_result = SerializedFileReader::new_with_options(test_file, options); + let reader = reader_result.unwrap(); + + let row_group_reader = reader.get_row_group(0).unwrap(); + + // use 'string_col', Boundary order: UNORDERED, total 352 data pages and 1 dictionary page. + let mut column_page_reader = row_group_reader.get_column_page_reader(9).unwrap(); + + let mut vec = vec![]; + + // Step 1: Peek and ensure dictionary page is correctly identified + let meta = column_page_reader.peek_next_page().unwrap().unwrap(); + assert!(meta.is_dict); + + // Step 2: Call skip_next_page to skip the dictionary page + column_page_reader.skip_next_page().unwrap(); + + // Step 3: Read the next data page after skipping the dictionary page + let page = column_page_reader.get_next_page().unwrap().unwrap(); + assert!(matches!(page.page_type(), basic::PageType::DATA_PAGE)); + + // Step 4: Continue reading remaining data pages and verify correctness + for _i in 0..351 { + // 352 total pages, 1 dictionary page is skipped + let meta = column_page_reader.peek_next_page().unwrap().unwrap(); + assert!(!meta.is_dict); // Verify no dictionary page here + vec.push(meta); + + let page = column_page_reader.get_next_page().unwrap().unwrap(); + assert!(matches!(page.page_type(), basic::PageType::DATA_PAGE)); + } + + // Step 5: Check if all pages are read + assert!(column_page_reader.peek_next_page().unwrap().is_none()); + assert!(column_page_reader.get_next_page().unwrap().is_none()); + + // Step 6: Verify the number of data pages read (should be 351 data pages) + assert_eq!(vec.len(), 351); + } + + #[test] + fn test_skip_page_with_offset_index() { + let test_file = get_test_file("alltypes_tiny_pages_plain.parquet"); + let builder = ReadOptionsBuilder::new(); + //enable read page index + let options = builder.with_page_index().build(); + let reader_result = SerializedFileReader::new_with_options(test_file, options); + let reader = reader_result.unwrap(); + + let row_group_reader = reader.get_row_group(0).unwrap(); + + //use 'int_col', Boundary order: ASCENDING, total 325 pages. + let mut column_page_reader = row_group_reader.get_column_page_reader(4).unwrap(); + + let mut vec = vec![]; + + for i in 0..325 { + if i % 2 == 0 { + vec.push(column_page_reader.get_next_page().unwrap().unwrap()); + } else { + column_page_reader.skip_next_page().unwrap(); + } + } + //check read all pages. + assert!(column_page_reader.peek_next_page().unwrap().is_none()); + assert!(column_page_reader.get_next_page().unwrap().is_none()); + + assert_eq!(vec.len(), 163); + } + + #[test] + fn test_skip_page_without_offset_index() { + let test_file = get_test_file("alltypes_tiny_pages_plain.parquet"); + + // use default SerializedFileReader without read offsetIndex + let reader_result = SerializedFileReader::new(test_file); + let reader = reader_result.unwrap(); + + let row_group_reader = reader.get_row_group(0).unwrap(); + + //use 'int_col', Boundary order: ASCENDING, total 325 pages. + let mut column_page_reader = row_group_reader.get_column_page_reader(4).unwrap(); + + let mut vec = vec![]; + + for i in 0..325 { + if i % 2 == 0 { + vec.push(column_page_reader.get_next_page().unwrap().unwrap()); + } else { + column_page_reader.peek_next_page().unwrap().unwrap(); + column_page_reader.skip_next_page().unwrap(); + } + } + //check read all pages. + assert!(column_page_reader.peek_next_page().unwrap().is_none()); + assert!(column_page_reader.get_next_page().unwrap().is_none()); + + assert_eq!(vec.len(), 163); + } + + #[test] + fn test_peek_page_with_dictionary_page() { + let test_file = get_test_file("alltypes_tiny_pages.parquet"); + let builder = ReadOptionsBuilder::new(); + //enable read page index + let options = builder.with_page_index().build(); + let reader_result = SerializedFileReader::new_with_options(test_file, options); + let reader = reader_result.unwrap(); + let row_group_reader = reader.get_row_group(0).unwrap(); + + //use 'string_col', Boundary order: UNORDERED, total 352 data ages and 1 dictionary page. + let mut column_page_reader = row_group_reader.get_column_page_reader(9).unwrap(); + + let mut vec = vec![]; + + let meta = column_page_reader.peek_next_page().unwrap().unwrap(); + assert!(meta.is_dict); + let page = column_page_reader.get_next_page().unwrap().unwrap(); + assert!(matches!(page.page_type(), basic::PageType::DICTIONARY_PAGE)); + + for i in 0..352 { + let meta = column_page_reader.peek_next_page().unwrap().unwrap(); + // have checked with `parquet-tools column-index -c string_col ./alltypes_tiny_pages.parquet` + // page meta has two scenarios(21, 20) of num_rows expect last page has 11 rows. + if i != 351 { + assert!((meta.num_rows == Some(21)) || (meta.num_rows == Some(20))); + } else { + // last page first row index is 7290, total row count is 7300 + // because first row start with zero, last page row count should be 10. + assert_eq!(meta.num_rows, Some(10)); + } + assert!(!meta.is_dict); + vec.push(meta); + let page = column_page_reader.get_next_page().unwrap().unwrap(); + assert!(matches!(page.page_type(), basic::PageType::DATA_PAGE)); + } + + //check read all pages. + assert!(column_page_reader.peek_next_page().unwrap().is_none()); + assert!(column_page_reader.get_next_page().unwrap().is_none()); + + assert_eq!(vec.len(), 352); + } + + #[test] + fn test_peek_page_with_dictionary_page_without_offset_index() { + let test_file = get_test_file("alltypes_tiny_pages.parquet"); + + let reader_result = SerializedFileReader::new(test_file); + let reader = reader_result.unwrap(); + let row_group_reader = reader.get_row_group(0).unwrap(); + + //use 'string_col', Boundary order: UNORDERED, total 352 data ages and 1 dictionary page. + let mut column_page_reader = row_group_reader.get_column_page_reader(9).unwrap(); + + let mut vec = vec![]; + + let meta = column_page_reader.peek_next_page().unwrap().unwrap(); + assert!(meta.is_dict); + let page = column_page_reader.get_next_page().unwrap().unwrap(); + assert!(matches!(page.page_type(), basic::PageType::DICTIONARY_PAGE)); + + for i in 0..352 { + let meta = column_page_reader.peek_next_page().unwrap().unwrap(); + // have checked with `parquet-tools column-index -c string_col ./alltypes_tiny_pages.parquet` + // page meta has two scenarios(21, 20) of num_rows expect last page has 11 rows. + if i != 351 { + assert!((meta.num_levels == Some(21)) || (meta.num_levels == Some(20))); + } else { + // last page first row index is 7290, total row count is 7300 + // because first row start with zero, last page row count should be 10. + assert_eq!(meta.num_levels, Some(10)); + } + assert!(!meta.is_dict); + vec.push(meta); + let page = column_page_reader.get_next_page().unwrap().unwrap(); + assert!(matches!(page.page_type(), basic::PageType::DATA_PAGE)); + } + + //check read all pages. + assert!(column_page_reader.peek_next_page().unwrap().is_none()); + assert!(column_page_reader.get_next_page().unwrap().is_none()); + + assert_eq!(vec.len(), 352); + } + + #[test] + fn test_fixed_length_index() { + let message_type = " + message test_schema { + OPTIONAL FIXED_LEN_BYTE_ARRAY (11) value (DECIMAL(25,2)); + } + "; + + let schema = parse_message_type(message_type).unwrap(); + let mut out = Vec::with_capacity(1024); + let mut writer = + SerializedFileWriter::new(&mut out, Arc::new(schema), Default::default()).unwrap(); + + let mut r = writer.next_row_group().unwrap(); + let mut c = r.next_column().unwrap().unwrap(); + c.typed::() + .write_batch( + &[vec![0; 11].into(), vec![5; 11].into(), vec![3; 11].into()], + Some(&[1, 1, 0, 1]), + None, + ) + .unwrap(); + c.close().unwrap(); + r.close().unwrap(); + writer.close().unwrap(); + + let b = Bytes::from(out); + let options = ReadOptionsBuilder::new().with_page_index().build(); + let reader = SerializedFileReader::new_with_options(b, options).unwrap(); + let index = reader.metadata().column_index().unwrap(); + + // 1 row group + assert_eq!(index.len(), 1); + let c = &index[0]; + // 1 column + assert_eq!(c.len(), 1); + + match &c[0] { + ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(v) => { + assert_eq!(v.num_pages(), 1); + assert_eq!(v.null_count(0).unwrap(), 1); + assert_eq!(v.min_value(0).unwrap(), &[0; 11]); + assert_eq!(v.max_value(0).unwrap(), &[5; 11]); + } + _ => unreachable!(), + } + } + + #[test] + fn test_multi_gz() { + let file = get_test_file("concatenated_gzip_members.parquet"); + let reader = SerializedFileReader::new(file).unwrap(); + let row_group_reader = reader.get_row_group(0).unwrap(); + match row_group_reader.get_column_reader(0).unwrap() { + ColumnReader::Int64ColumnReader(mut reader) => { + let mut buffer = Vec::with_capacity(1024); + let mut def_levels = Vec::with_capacity(1024); + let (num_records, num_values, num_levels) = reader + .read_records(1024, Some(&mut def_levels), None, &mut buffer) + .unwrap(); + + assert_eq!(num_records, 513); + assert_eq!(num_values, 513); + assert_eq!(num_levels, 513); + + let expected: Vec = (1..514).collect(); + assert_eq!(&buffer, &expected); + } + _ => unreachable!(), + } + } + + #[test] + fn test_byte_stream_split_extended() { + let path = format!( + "{}/byte_stream_split_extended.gzip.parquet", + arrow::util::test_util::parquet_test_data(), + ); + let file = File::open(path).unwrap(); + let reader = Box::new(SerializedFileReader::new(file).expect("Failed to create reader")); + + // Use full schema as projected schema + let mut iter = reader + .get_row_iter(None) + .expect("Failed to create row iterator"); + + let mut start = 0; + let end = reader.metadata().file_metadata().num_rows(); + + let check_row = |row: Result| { + assert!(row.is_ok()); + let r = row.unwrap(); + assert_eq!(r.get_float16(0).unwrap(), r.get_float16(1).unwrap()); + assert_eq!(r.get_float(2).unwrap(), r.get_float(3).unwrap()); + assert_eq!(r.get_double(4).unwrap(), r.get_double(5).unwrap()); + assert_eq!(r.get_int(6).unwrap(), r.get_int(7).unwrap()); + assert_eq!(r.get_long(8).unwrap(), r.get_long(9).unwrap()); + assert_eq!(r.get_bytes(10).unwrap(), r.get_bytes(11).unwrap()); + assert_eq!(r.get_decimal(12).unwrap(), r.get_decimal(13).unwrap()); + }; + + while start < end { + match iter.next() { + Some(row) => check_row(row), + None => break, + }; + start += 1; + } + } + + #[test] + fn test_filtered_rowgroup_metadata() { + let message_type = " + message test_schema { + REQUIRED INT32 a; + } + "; + let schema = Arc::new(parse_message_type(message_type).unwrap()); + let props = Arc::new( + WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .build(), + ); + let mut file: File = tempfile::tempfile().unwrap(); + let mut file_writer = SerializedFileWriter::new(&mut file, schema, props).unwrap(); + let data = [1, 2, 3, 4, 5]; + + // write 5 row groups + for idx in 0..5 { + let data_i: Vec = data.iter().map(|x| x * (idx + 1)).collect(); + let mut row_group_writer = file_writer.next_row_group().unwrap(); + if let Some(mut writer) = row_group_writer.next_column().unwrap() { + writer + .typed::() + .write_batch(data_i.as_slice(), None, None) + .unwrap(); + writer.close().unwrap(); + } + row_group_writer.close().unwrap(); + file_writer.flushed_row_groups(); + } + let file_metadata = file_writer.close().unwrap(); + + assert_eq!(file_metadata.file_metadata().num_rows(), 25); + assert_eq!(file_metadata.num_row_groups(), 5); + + // read only the 3rd row group + let read_options = ReadOptionsBuilder::new() + .with_page_index() + .with_predicate(Box::new(|rgmeta, _| rgmeta.ordinal().unwrap_or(0) == 2)) + .build(); + let reader = + SerializedFileReader::new_with_options(file.try_clone().unwrap(), read_options) + .unwrap(); + let metadata = reader.metadata(); + + // check we got the expected row group + assert_eq!(metadata.num_row_groups(), 1); + assert_eq!(metadata.row_group(0).ordinal(), Some(2)); + + // check we only got the relevant page indexes + assert!(metadata.column_index().is_some()); + assert!(metadata.offset_index().is_some()); + assert_eq!(metadata.column_index().unwrap().len(), 1); + assert_eq!(metadata.offset_index().unwrap().len(), 1); + let col_idx = metadata.column_index().unwrap(); + let off_idx = metadata.offset_index().unwrap(); + let col_stats = metadata.row_group(0).column(0).statistics().unwrap(); + let pg_idx = &col_idx[0][0]; + let off_idx_i = &off_idx[0][0]; + + // test that we got the index matching the row group + match pg_idx { + ColumnIndexMetaData::INT32(int_idx) => { + let min = col_stats.min_bytes_opt().unwrap().get_i32_le(); + let max = col_stats.max_bytes_opt().unwrap().get_i32_le(); + assert_eq!(int_idx.min_value(0), Some(min).as_ref()); + assert_eq!(int_idx.max_value(0), Some(max).as_ref()); + } + _ => panic!("wrong stats type"), + } + + // check offset index matches too + assert_eq!( + off_idx_i.page_locations[0].offset, + metadata.row_group(0).column(0).data_page_offset() + ); + + // read non-contiguous row groups + let read_options = ReadOptionsBuilder::new() + .with_page_index() + .with_predicate(Box::new(|rgmeta, _| rgmeta.ordinal().unwrap_or(0) % 2 == 1)) + .build(); + let reader = + SerializedFileReader::new_with_options(file.try_clone().unwrap(), read_options) + .unwrap(); + let metadata = reader.metadata(); + + // check we got the expected row groups + assert_eq!(metadata.num_row_groups(), 2); + assert_eq!(metadata.row_group(0).ordinal(), Some(1)); + assert_eq!(metadata.row_group(1).ordinal(), Some(3)); + + // check we only got the relevant page indexes + assert!(metadata.column_index().is_some()); + assert!(metadata.offset_index().is_some()); + assert_eq!(metadata.column_index().unwrap().len(), 2); + assert_eq!(metadata.offset_index().unwrap().len(), 2); + let col_idx = metadata.column_index().unwrap(); + let off_idx = metadata.offset_index().unwrap(); + + for (i, col_idx_i) in col_idx.iter().enumerate().take(metadata.num_row_groups()) { + let col_stats = metadata.row_group(i).column(0).statistics().unwrap(); + let pg_idx = &col_idx_i[0]; + let off_idx_i = &off_idx[i][0]; + + // test that we got the index matching the row group + match pg_idx { + ColumnIndexMetaData::INT32(int_idx) => { + let min = col_stats.min_bytes_opt().unwrap().get_i32_le(); + let max = col_stats.max_bytes_opt().unwrap().get_i32_le(); + assert_eq!(int_idx.min_value(0), Some(min).as_ref()); + assert_eq!(int_idx.max_value(0), Some(max).as_ref()); + } + _ => panic!("wrong stats type"), + } + + // check offset index matches too + assert_eq!( + off_idx_i.page_locations[0].offset, + metadata.row_group(i).column(0).data_page_offset() + ); + } + } + + #[test] + fn test_reuse_schema() { + let file = get_test_file("alltypes_plain.parquet"); + let file_reader = SerializedFileReader::new(file.try_clone().unwrap()).unwrap(); + let schema = file_reader.metadata().file_metadata().schema_descr_ptr(); + let expected = file_reader.metadata; + + let options = ReadOptionsBuilder::new() + .with_parquet_schema(schema) + .build(); + let file_reader = SerializedFileReader::new_with_options(file, options).unwrap(); + + assert_eq!(expected.as_ref(), file_reader.metadata.as_ref()); + // Should have used the same schema instance + assert!(Arc::ptr_eq( + &expected.file_metadata().schema_descr_ptr(), + &file_reader.metadata.file_metadata().schema_descr_ptr() + )); + } + + #[test] + fn test_read_unknown_logical_type() { + let file = get_test_file("unknown-logical-type.parquet"); + let reader = SerializedFileReader::new(file).expect("Error opening file"); + + let schema = reader.metadata().file_metadata().schema_descr(); + assert_eq!( + schema.column(0).logical_type_ref(), + Some(&basic::LogicalType::String) + ); + assert_eq!( + schema.column(1).logical_type_ref(), + Some(&basic::LogicalType::_Unknown { field_id: 2555 }) + ); + assert_eq!(schema.column(1).physical_type(), Type::BYTE_ARRAY); + + let mut iter = reader + .get_row_iter(None) + .expect("Failed to create row iterator"); + + let mut num_rows = 0; + while iter.next().is_some() { + num_rows += 1; + } + assert_eq!(num_rows, reader.metadata().file_metadata().num_rows()); + } +} diff --git a/vendor/parquet/src/file/statistics.rs b/vendor/parquet/src/file/statistics.rs new file mode 100644 index 0000000..9682fd5 --- /dev/null +++ b/vendor/parquet/src/file/statistics.rs @@ -0,0 +1,1144 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Contains definitions for working with Parquet statistics. +//! +//! Though some common methods are available on enum, use pattern match to extract +//! actual min and max values from statistics, see below: +//! +//! # Examples +//! ```rust +//! use parquet::file::statistics::Statistics; +//! +//! let stats = Statistics::int32(Some(1), Some(10), None, Some(3), true); +//! assert_eq!(stats.null_count_opt(), Some(3)); +//! assert!(stats.is_min_max_deprecated()); +//! assert!(stats.min_is_exact()); +//! assert!(stats.max_is_exact()); +//! +//! match stats { +//! Statistics::Int32(ref typed) => { +//! assert_eq!(typed.min_opt(), Some(&1)); +//! assert_eq!(typed.max_opt(), Some(&10)); +//! } +//! _ => {} +//! } +//! ``` + +use std::fmt; + +use crate::basic::Type; +use crate::data_type::private::ParquetValueType; +use crate::data_type::*; +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::thrift::PageStatistics; +use crate::util::bit_util::FromBytes; + +pub(crate) mod private { + use super::*; + + pub trait MakeStatistics { + fn make_statistics(statistics: ValueStatistics) -> Statistics + where + Self: Sized; + } + + macro_rules! gen_make_statistics { + ($value_ty:ty, $stat:ident) => { + impl MakeStatistics for $value_ty { + fn make_statistics(statistics: ValueStatistics) -> Statistics + where + Self: Sized, + { + Statistics::$stat(statistics) + } + } + }; + } + + gen_make_statistics!(bool, Boolean); + gen_make_statistics!(i32, Int32); + gen_make_statistics!(i64, Int64); + gen_make_statistics!(Int96, Int96); + gen_make_statistics!(f32, Float); + gen_make_statistics!(f64, Double); + gen_make_statistics!(ByteArray, ByteArray); + gen_make_statistics!(FixedLenByteArray, FixedLenByteArray); +} + +/// Macro to generate methods to create Statistics. +macro_rules! statistics_new_func { + ($func:ident, $vtype:ty, $stat:ident) => { + #[doc = concat!("Creates new statistics for `", stringify!($stat), "` column type.")] + pub fn $func( + min: $vtype, + max: $vtype, + distinct: Option, + nulls: Option, + is_deprecated: bool, + ) -> Self { + Statistics::$stat(ValueStatistics::new( + min, + max, + distinct, + nulls, + is_deprecated, + )) + } + }; +} + +// Macro to generate getter functions for Statistics. +macro_rules! statistics_enum_func { + ($self:ident, $func:ident) => {{ + match *$self { + Statistics::Boolean(ref typed) => typed.$func(), + Statistics::Int32(ref typed) => typed.$func(), + Statistics::Int64(ref typed) => typed.$func(), + Statistics::Int96(ref typed) => typed.$func(), + Statistics::Float(ref typed) => typed.$func(), + Statistics::Double(ref typed) => typed.$func(), + Statistics::ByteArray(ref typed) => typed.$func(), + Statistics::FixedLenByteArray(ref typed) => typed.$func(), + } + }}; +} + +/// Converts Thrift definition into `Statistics`. +pub(crate) fn from_thrift_page_stats( + physical_type: Type, + thrift_stats: Option, +) -> Result> { + Ok(match thrift_stats { + Some(stats) => { + // Generic null count. + let null_count = stats + .null_count + .map(|null_count| { + if null_count < 0 { + return Err(ParquetError::General(format!( + "Statistics null count is negative {null_count}", + ))); + } + Ok(null_count as u64) + }) + .transpose()?; + // Generic distinct count (count of distinct values occurring) + let distinct_count = stats.distinct_count.map(|value| value as u64); + // Whether or not statistics use deprecated min/max fields. + let old_format = stats.min_value.is_none() && stats.max_value.is_none(); + // Generic min value as bytes. + let min = if old_format { + stats.min + } else { + stats.min_value + }; + // Generic max value as bytes. + let max = if old_format { + stats.max + } else { + stats.max_value + }; + + fn check_len(min: &Option>, max: &Option>, len: usize) -> Result<()> { + if let Some(min) = min { + if min.len() < len { + return Err(ParquetError::General( + "Insufficient bytes to parse min statistic".to_string(), + )); + } + } + if let Some(max) = max { + if max.len() < len { + return Err(ParquetError::General( + "Insufficient bytes to parse max statistic".to_string(), + )); + } + } + Ok(()) + } + + match physical_type { + Type::BOOLEAN => check_len(&min, &max, 1), + Type::INT32 | Type::FLOAT => check_len(&min, &max, 4), + Type::INT64 | Type::DOUBLE => check_len(&min, &max, 8), + Type::INT96 => check_len(&min, &max, 12), + _ => Ok(()), + }?; + + // Values are encoded using PLAIN encoding definition, except that + // variable-length byte arrays do not include a length prefix. + // + // Instead of using actual decoder, we manually convert values. + let res = match physical_type { + Type::BOOLEAN => Statistics::boolean( + min.map(|data| data[0] != 0), + max.map(|data| data[0] != 0), + distinct_count, + null_count, + old_format, + ), + Type::INT32 => Statistics::int32( + min.map(|data| i32::from_le_bytes(data[..4].try_into().unwrap())), + max.map(|data| i32::from_le_bytes(data[..4].try_into().unwrap())), + distinct_count, + null_count, + old_format, + ), + Type::INT64 => Statistics::int64( + min.map(|data| i64::from_le_bytes(data[..8].try_into().unwrap())), + max.map(|data| i64::from_le_bytes(data[..8].try_into().unwrap())), + distinct_count, + null_count, + old_format, + ), + Type::INT96 => { + // INT96 statistics may not be correct, because comparison is signed + let min = if let Some(data) = min { + if data.len() != 12 { + return Err(ParquetError::General( + "Incorrect Int96 min statistics".to_string(), + )); + } + Some(Int96::try_from_le_slice(&data)?) + } else { + None + }; + let max = if let Some(data) = max { + if data.len() != 12 { + return Err(ParquetError::General( + "Incorrect Int96 max statistics".to_string(), + )); + } + Some(Int96::try_from_le_slice(&data)?) + } else { + None + }; + Statistics::int96(min, max, distinct_count, null_count, old_format) + } + Type::FLOAT => Statistics::float( + min.map(|data| f32::from_le_bytes(data[..4].try_into().unwrap())), + max.map(|data| f32::from_le_bytes(data[..4].try_into().unwrap())), + distinct_count, + null_count, + old_format, + ), + Type::DOUBLE => Statistics::double( + min.map(|data| f64::from_le_bytes(data[..8].try_into().unwrap())), + max.map(|data| f64::from_le_bytes(data[..8].try_into().unwrap())), + distinct_count, + null_count, + old_format, + ), + Type::BYTE_ARRAY => Statistics::ByteArray( + ValueStatistics::new( + min.map(ByteArray::from), + max.map(ByteArray::from), + distinct_count, + null_count, + old_format, + ) + .with_max_is_exact(stats.is_max_value_exact.unwrap_or(false)) + .with_min_is_exact(stats.is_min_value_exact.unwrap_or(false)), + ), + Type::FIXED_LEN_BYTE_ARRAY => Statistics::FixedLenByteArray( + ValueStatistics::new( + min.map(ByteArray::from).map(FixedLenByteArray::from), + max.map(ByteArray::from).map(FixedLenByteArray::from), + distinct_count, + null_count, + old_format, + ) + .with_max_is_exact(stats.is_max_value_exact.unwrap_or(false)) + .with_min_is_exact(stats.is_min_value_exact.unwrap_or(false)), + ), + }; + + Some(res) + } + None => None, + }) +} + +/// Convert Statistics into Thrift definition. +pub(crate) fn page_stats_to_thrift(stats: Option<&Statistics>) -> Option { + let stats = stats?; + + // record null count if it can fit in i64 + let null_count = stats + .null_count_opt() + .and_then(|value| i64::try_from(value).ok()); + + // record distinct count if it can fit in i64 + let distinct_count = stats + .distinct_count_opt() + .and_then(|value| i64::try_from(value).ok()); + + let mut thrift_stats = PageStatistics { + max: None, + min: None, + null_count, + distinct_count, + max_value: None, + min_value: None, + is_max_value_exact: None, + is_min_value_exact: None, + }; + + // Get min/max if set. + let (min, max, min_exact, max_exact) = ( + stats.min_bytes_opt().map(|x| x.to_vec()), + stats.max_bytes_opt().map(|x| x.to_vec()), + Some(stats.min_is_exact()), + Some(stats.max_is_exact()), + ); + if stats.is_min_max_backwards_compatible() { + // Copy to deprecated min, max values for compatibility with older readers + thrift_stats.min.clone_from(&min); + thrift_stats.max.clone_from(&max); + } + + if !stats.is_min_max_deprecated() { + thrift_stats.min_value = min; + thrift_stats.max_value = max; + } + + thrift_stats.is_min_value_exact = min_exact; + thrift_stats.is_max_value_exact = max_exact; + + Some(thrift_stats) +} + +/// Strongly typed statistics for a column chunk within a row group. +/// +/// This structure is a natively typed, in memory representation of the thrift +/// `Statistics` structure in a Parquet file footer. The statistics stored in +/// this structure can be used by query engines to skip decoding pages while +/// reading parquet data. +/// +/// Page level statistics are stored separately, in [ColumnIndexMetaData]. +/// +/// [ColumnIndexMetaData]: crate::file::page_index::column_index::ColumnIndexMetaData +#[derive(Debug, Clone, PartialEq)] +pub enum Statistics { + /// Statistics for Boolean column + Boolean(ValueStatistics), + /// Statistics for Int32 column + Int32(ValueStatistics), + /// Statistics for Int64 column + Int64(ValueStatistics), + /// Statistics for Int96 column + Int96(ValueStatistics), + /// Statistics for Float column + Float(ValueStatistics), + /// Statistics for Double column + Double(ValueStatistics), + /// Statistics for ByteArray column + ByteArray(ValueStatistics), + /// Statistics for FixedLenByteArray column + FixedLenByteArray(ValueStatistics), +} + +impl From> for Statistics { + fn from(t: ValueStatistics) -> Self { + T::make_statistics(t) + } +} + +impl Statistics { + /// Creates new statistics for a column type + pub fn new( + min: Option, + max: Option, + distinct_count: Option, + null_count: Option, + is_deprecated: bool, + ) -> Self { + Self::from(ValueStatistics::new( + min, + max, + distinct_count, + null_count, + is_deprecated, + )) + } + + statistics_new_func![boolean, Option, Boolean]; + + statistics_new_func![int32, Option, Int32]; + + statistics_new_func![int64, Option, Int64]; + + statistics_new_func![int96, Option, Int96]; + + statistics_new_func![float, Option, Float]; + + statistics_new_func![double, Option, Double]; + + statistics_new_func![byte_array, Option, ByteArray]; + + statistics_new_func![ + fixed_len_byte_array, + Option, + FixedLenByteArray + ]; + + /// Returns `true` if statistics have old `min` and `max` fields set. + /// This means that the column order is likely to be undefined, which, for old files + /// could mean a signed sort order of values. + /// + /// Refer to [`ColumnOrder`](crate::basic::ColumnOrder) and + /// [`SortOrder`](crate::basic::SortOrder) for more information. + pub fn is_min_max_deprecated(&self) -> bool { + statistics_enum_func![self, is_min_max_deprecated] + } + + /// Old versions of parquet stored statistics in `min` and `max` fields, ordered + /// using signed comparison. This resulted in an undefined ordering for unsigned + /// quantities, such as booleans and unsigned integers. + /// + /// These fields were therefore deprecated in favour of `min_value` and `max_value`, + /// which have a type-defined sort order. + /// + /// However, not all readers have been updated. For backwards compatibility, this method + /// returns `true` if the statistics within this have a signed sort order, that is + /// compatible with being stored in the deprecated `min` and `max` fields + pub fn is_min_max_backwards_compatible(&self) -> bool { + statistics_enum_func![self, is_min_max_backwards_compatible] + } + + /// Returns optional value of number of distinct values occurring. + /// When it is `None`, the value should be ignored. + pub fn distinct_count_opt(&self) -> Option { + statistics_enum_func![self, distinct_count] + } + + /// Returns number of null values for the column, if known. + /// Note that this includes all nulls when column is part of the complex type. + /// + /// Note: Versions of this library prior to `58.1.0` returned `0` if the null count + /// was not available. This method now returns `None` in that case. + /// + /// Also, versions of this library prior to `53.1.0` did not store a null count + /// statistic when the null count was `0`. + /// + /// It is unsound to assume that missing nullcount stats mean the column contains no nulls, + /// but code that depends on the old behavior can restore it by defaulting to zero: + /// + /// ```no_run + /// # use parquet::file::statistics::Statistics; + /// # let statistics: Statistics = todo!(); + /// let null_count = statistics.null_count_opt().unwrap_or(0); + /// ``` + pub fn null_count_opt(&self) -> Option { + statistics_enum_func![self, null_count_opt] + } + + /// Returns `true` if the min value is set, and is an exact min value. + pub fn min_is_exact(&self) -> bool { + statistics_enum_func![self, min_is_exact] + } + + /// Returns `true` if the max value is set, and is an exact max value. + pub fn max_is_exact(&self) -> bool { + statistics_enum_func![self, max_is_exact] + } + + /// Returns slice of bytes that represent min value, if min value is known. + pub fn min_bytes_opt(&self) -> Option<&[u8]> { + statistics_enum_func![self, min_bytes_opt] + } + + /// Returns slice of bytes that represent max value, if max value is known. + pub fn max_bytes_opt(&self) -> Option<&[u8]> { + statistics_enum_func![self, max_bytes_opt] + } + + /// Returns physical type associated with statistics. + pub fn physical_type(&self) -> Type { + match self { + Statistics::Boolean(_) => Type::BOOLEAN, + Statistics::Int32(_) => Type::INT32, + Statistics::Int64(_) => Type::INT64, + Statistics::Int96(_) => Type::INT96, + Statistics::Float(_) => Type::FLOAT, + Statistics::Double(_) => Type::DOUBLE, + Statistics::ByteArray(_) => Type::BYTE_ARRAY, + Statistics::FixedLenByteArray(_) => Type::FIXED_LEN_BYTE_ARRAY, + } + } +} + +impl fmt::Display for Statistics { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + Statistics::Boolean(typed) => write!(f, "{typed}"), + Statistics::Int32(typed) => write!(f, "{typed}"), + Statistics::Int64(typed) => write!(f, "{typed}"), + Statistics::Int96(typed) => write!(f, "{typed}"), + Statistics::Float(typed) => write!(f, "{typed}"), + Statistics::Double(typed) => write!(f, "{typed}"), + Statistics::ByteArray(typed) => write!(f, "{typed}"), + Statistics::FixedLenByteArray(typed) => write!(f, "{typed}"), + } + } +} + +/// Typed implementation for [`Statistics`]. +pub type TypedStatistics = ValueStatistics<::T>; + +/// Typed statistics for one column chunk +/// +/// See [`Statistics`] for more details +#[derive(Clone, Eq, PartialEq)] +pub struct ValueStatistics { + min: Option, + max: Option, + // Distinct count could be omitted in some cases + distinct_count: Option, + null_count: Option, + + // Whether or not the min or max values are exact, or truncated. + is_max_value_exact: bool, + is_min_value_exact: bool, + + /// If `true` populate the deprecated `min` and `max` fields instead of + /// `min_value` and `max_value` + is_min_max_deprecated: bool, + + /// If `true` the statistics are compatible with the deprecated `min` and + /// `max` fields. See [`ValueStatistics::is_min_max_backwards_compatible`] + is_min_max_backwards_compatible: bool, +} + +impl ValueStatistics { + /// Creates new typed statistics. + pub fn new( + min: Option, + max: Option, + distinct_count: Option, + null_count: Option, + is_min_max_deprecated: bool, + ) -> Self { + Self { + is_max_value_exact: max.is_some(), + is_min_value_exact: min.is_some(), + min, + max, + distinct_count, + null_count, + is_min_max_deprecated, + is_min_max_backwards_compatible: is_min_max_deprecated, + } + } + + /// Set whether the stored `min` field represents the exact + /// minimum, or just a bound on the minimum value. + /// + /// see [`Self::min_is_exact`] + pub fn with_min_is_exact(self, is_min_value_exact: bool) -> Self { + Self { + is_min_value_exact, + ..self + } + } + + /// Set whether the stored `max` field represents the exact + /// maximum, or just a bound on the maximum value. + /// + /// see [`Self::max_is_exact`] + pub fn with_max_is_exact(self, is_max_value_exact: bool) -> Self { + Self { + is_max_value_exact, + ..self + } + } + + /// Set whether to write the deprecated `min` and `max` fields + /// for compatibility with older parquet writers + /// + /// This should only be enabled if the field is signed, + /// see [`Self::is_min_max_backwards_compatible`] + pub fn with_backwards_compatible_min_max(self, backwards_compatible: bool) -> Self { + Self { + is_min_max_backwards_compatible: backwards_compatible, + ..self + } + } + + /// Returns min value of the statistics, if known. + pub fn min_opt(&self) -> Option<&T> { + self.min.as_ref() + } + + /// Returns max value of the statistics, if known. + pub fn max_opt(&self) -> Option<&T> { + self.max.as_ref() + } + + /// Whether or not min and max values are set. + /// Normally both min/max values will be set to `Some(value)` or `None`. + pub(crate) fn _internal_has_min_max_set(&self) -> bool { + self.min.is_some() && self.max.is_some() + } + + /// Whether or not max value is set, and is an exact value. + pub fn max_is_exact(&self) -> bool { + self.max.is_some() && self.is_max_value_exact + } + + /// Whether or not min value is set, and is an exact value. + pub fn min_is_exact(&self) -> bool { + self.min.is_some() && self.is_min_value_exact + } + + /// Returns optional value of number of distinct values occurring. + pub fn distinct_count(&self) -> Option { + self.distinct_count + } + + /// Returns null count. + pub fn null_count_opt(&self) -> Option { + self.null_count + } + + /// Returns `true` if statistics were created using old min/max fields. + fn is_min_max_deprecated(&self) -> bool { + self.is_min_max_deprecated + } + + /// Old versions of parquet stored statistics in `min` and `max` fields, ordered + /// using signed comparison. This resulted in an undefined ordering for unsigned + /// quantities, such as booleans and unsigned integers. + /// + /// These fields were therefore deprecated in favour of `min_value` and `max_value`, + /// which have a type-defined sort order. + /// + /// However, not all readers have been updated. For backwards compatibility, this method + /// returns `true` if the statistics within this have a signed sort order, that is + /// compatible with being stored in the deprecated `min` and `max` fields + pub fn is_min_max_backwards_compatible(&self) -> bool { + self.is_min_max_backwards_compatible + } +} + +impl ValueStatistics { + /// Returns min value as bytes of the statistics, if min value is known. + pub fn min_bytes_opt(&self) -> Option<&[u8]> { + self.min_opt().map(AsBytes::as_bytes) + } + + /// Returns max value as bytes of the statistics, if max value is known. + pub fn max_bytes_opt(&self) -> Option<&[u8]> { + self.max_opt().map(AsBytes::as_bytes) + } +} + +impl fmt::Display for ValueStatistics { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{{")?; + write!(f, "min: ")?; + match self.min { + Some(ref value) => write!(f, "{value}")?, + None => write!(f, "N/A")?, + } + write!(f, ", max: ")?; + match self.max { + Some(ref value) => write!(f, "{value}")?, + None => write!(f, "N/A")?, + } + write!(f, ", distinct_count: ")?; + match self.distinct_count { + Some(value) => write!(f, "{value}")?, + None => write!(f, "N/A")?, + } + write!(f, ", null_count: ")?; + match self.null_count { + Some(value) => write!(f, "{value}")?, + None => write!(f, "N/A")?, + } + write!(f, ", min_max_deprecated: {}", self.is_min_max_deprecated)?; + write!(f, ", max_value_exact: {}", self.is_max_value_exact)?; + write!(f, ", min_value_exact: {}", self.is_min_value_exact)?; + write!(f, "}}") + } +} + +impl fmt::Debug for ValueStatistics { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "{{min: {:?}, max: {:?}, distinct_count: {:?}, null_count: {:?}, \ + min_max_deprecated: {}, min_max_backwards_compatible: {}, max_value_exact: {}, min_value_exact: {}}}", + self.min, + self.max, + self.distinct_count, + self.null_count, + self.is_min_max_deprecated, + self.is_min_max_backwards_compatible, + self.is_max_value_exact, + self.is_min_value_exact + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_statistics_min_max_bytes() { + let stats = Statistics::int32(Some(-123), Some(234), None, Some(1), false); + assert_eq!(stats.min_bytes_opt(), Some((-123).as_bytes())); + assert_eq!(stats.max_bytes_opt(), Some(234.as_bytes())); + + let stats = Statistics::byte_array( + Some(ByteArray::from(vec![1, 2, 3])), + Some(ByteArray::from(vec![3, 4, 5])), + None, + Some(1), + true, + ); + assert_eq!(stats.min_bytes_opt().unwrap(), &[1, 2, 3]); + assert_eq!(stats.max_bytes_opt().unwrap(), &[3, 4, 5]); + } + + #[test] + #[should_panic(expected = "General(\"Statistics null count is negative -10\")")] + fn test_statistics_negative_null_count() { + let thrift_stats = PageStatistics { + max: None, + min: None, + null_count: Some(-10), + distinct_count: None, + max_value: None, + min_value: None, + is_max_value_exact: None, + is_min_value_exact: None, + }; + + from_thrift_page_stats(Type::INT32, Some(thrift_stats)).unwrap(); + } + + #[test] + fn test_statistics_thrift_none() { + assert_eq!(from_thrift_page_stats(Type::INT32, None).unwrap(), None); + assert_eq!( + from_thrift_page_stats(Type::BYTE_ARRAY, None).unwrap(), + None + ); + } + + #[test] + fn test_statistics_debug() { + let stats = Statistics::int32(Some(1), Some(12), None, Some(12), true); + assert_eq!( + format!("{stats:?}"), + "Int32({min: Some(1), max: Some(12), distinct_count: None, null_count: Some(12), \ + min_max_deprecated: true, min_max_backwards_compatible: true, max_value_exact: true, min_value_exact: true})" + ); + + let stats = Statistics::int32(None, None, None, Some(7), false); + assert_eq!( + format!("{stats:?}"), + "Int32({min: None, max: None, distinct_count: None, null_count: Some(7), \ + min_max_deprecated: false, min_max_backwards_compatible: false, max_value_exact: false, min_value_exact: false})" + ) + } + + #[test] + fn test_statistics_display() { + let stats = Statistics::int32(Some(1), Some(12), None, Some(12), true); + assert_eq!( + format!("{stats}"), + "{min: 1, max: 12, distinct_count: N/A, null_count: 12, min_max_deprecated: true, max_value_exact: true, min_value_exact: true}" + ); + + let stats = Statistics::int64(None, None, None, Some(7), false); + assert_eq!( + format!("{stats}"), + "{min: N/A, max: N/A, distinct_count: N/A, null_count: 7, min_max_deprecated: \ + false, max_value_exact: false, min_value_exact: false}" + ); + + let stats = Statistics::int96( + Some(Int96::from(vec![1, 0, 0])), + Some(Int96::from(vec![2, 3, 4])), + None, + Some(3), + true, + ); + assert_eq!( + format!("{stats}"), + "{min: [1, 0, 0], max: [2, 3, 4], distinct_count: N/A, null_count: 3, \ + min_max_deprecated: true, max_value_exact: true, min_value_exact: true}" + ); + + let stats = Statistics::ByteArray( + ValueStatistics::new( + Some(ByteArray::from(vec![1u8])), + Some(ByteArray::from(vec![2u8])), + Some(5), + Some(7), + false, + ) + .with_max_is_exact(false) + .with_min_is_exact(false), + ); + assert_eq!( + format!("{stats}"), + "{min: [1], max: [2], distinct_count: 5, null_count: 7, min_max_deprecated: false, max_value_exact: false, min_value_exact: false}" + ); + } + + #[test] + fn test_statistics_partial_eq() { + let expected = Statistics::int32(Some(12), Some(45), None, Some(11), true); + + assert!(Statistics::int32(Some(12), Some(45), None, Some(11), true) == expected); + assert!(Statistics::int32(Some(11), Some(45), None, Some(11), true) != expected); + assert!(Statistics::int32(Some(12), Some(44), None, Some(11), true) != expected); + assert!(Statistics::int32(Some(12), Some(45), None, Some(23), true) != expected); + assert!(Statistics::int32(Some(12), Some(45), None, Some(11), false) != expected); + + assert!( + Statistics::int32(Some(12), Some(45), None, Some(11), false) + != Statistics::int64(Some(12), Some(45), None, Some(11), false) + ); + + assert!( + Statistics::boolean(Some(false), Some(true), None, None, true) + != Statistics::double(Some(1.2), Some(4.5), None, None, true) + ); + + assert!( + Statistics::byte_array( + Some(ByteArray::from(vec![1, 2, 3])), + Some(ByteArray::from(vec![1, 2, 3])), + None, + None, + true + ) != Statistics::fixed_len_byte_array( + Some(ByteArray::from(vec![1, 2, 3]).into()), + Some(ByteArray::from(vec![1, 2, 3]).into()), + None, + None, + true, + ) + ); + + assert!( + Statistics::byte_array( + Some(ByteArray::from(vec![1, 2, 3])), + Some(ByteArray::from(vec![1, 2, 3])), + None, + None, + true, + ) != Statistics::ByteArray( + ValueStatistics::new( + Some(ByteArray::from(vec![1, 2, 3])), + Some(ByteArray::from(vec![1, 2, 3])), + None, + None, + true, + ) + .with_max_is_exact(false) + ) + ); + + assert!( + Statistics::fixed_len_byte_array( + Some(FixedLenByteArray::from(vec![1, 2, 3])), + Some(FixedLenByteArray::from(vec![1, 2, 3])), + None, + None, + true, + ) != Statistics::FixedLenByteArray( + ValueStatistics::new( + Some(FixedLenByteArray::from(vec![1, 2, 3])), + Some(FixedLenByteArray::from(vec![1, 2, 3])), + None, + None, + true, + ) + .with_min_is_exact(false) + ) + ); + } + + #[test] + fn test_statistics_from_thrift() { + // Helper method to check statistics conversion. + fn check_stats(stats: Statistics) { + let tpe = stats.physical_type(); + let thrift_stats = page_stats_to_thrift(Some(&stats)); + assert_eq!( + from_thrift_page_stats(tpe, thrift_stats).unwrap(), + Some(stats) + ); + } + + check_stats(Statistics::boolean( + Some(false), + Some(true), + None, + Some(7), + true, + )); + check_stats(Statistics::boolean( + Some(false), + Some(true), + None, + Some(7), + true, + )); + check_stats(Statistics::boolean( + Some(false), + Some(true), + None, + Some(0), + false, + )); + check_stats(Statistics::boolean( + Some(true), + Some(true), + None, + Some(7), + true, + )); + check_stats(Statistics::boolean( + Some(false), + Some(false), + None, + Some(7), + true, + )); + check_stats(Statistics::boolean(None, None, None, Some(7), true)); + + check_stats(Statistics::int32( + Some(-100), + Some(500), + None, + Some(7), + true, + )); + check_stats(Statistics::int32( + Some(-100), + Some(500), + None, + Some(0), + false, + )); + check_stats(Statistics::int32(None, None, None, Some(7), true)); + + check_stats(Statistics::int64( + Some(-100), + Some(200), + None, + Some(7), + true, + )); + check_stats(Statistics::int64( + Some(-100), + Some(200), + None, + Some(0), + false, + )); + check_stats(Statistics::int64(None, None, None, Some(7), true)); + + check_stats(Statistics::float(Some(1.2), Some(3.4), None, Some(7), true)); + check_stats(Statistics::float( + Some(1.2), + Some(3.4), + None, + Some(0), + false, + )); + check_stats(Statistics::float(None, None, None, Some(7), true)); + + check_stats(Statistics::double( + Some(1.2), + Some(3.4), + None, + Some(7), + true, + )); + check_stats(Statistics::double( + Some(1.2), + Some(3.4), + None, + Some(0), + false, + )); + check_stats(Statistics::double(None, None, None, Some(7), true)); + + check_stats(Statistics::byte_array( + Some(ByteArray::from(vec![1, 2, 3])), + Some(ByteArray::from(vec![3, 4, 5])), + None, + Some(7), + true, + )); + check_stats(Statistics::byte_array(None, None, None, Some(7), true)); + + check_stats(Statistics::fixed_len_byte_array( + Some(ByteArray::from(vec![1, 2, 3]).into()), + Some(ByteArray::from(vec![3, 4, 5]).into()), + None, + Some(7), + true, + )); + check_stats(Statistics::fixed_len_byte_array( + None, + None, + None, + Some(7), + true, + )); + } + + #[test] + fn test_count_encoding() { + statistics_count_test(None, None); + statistics_count_test(Some(0), Some(0)); + statistics_count_test(Some(100), Some(2000)); + statistics_count_test(Some(1), None); + statistics_count_test(None, Some(1)); + } + + #[test] + fn test_count_encoding_distinct_too_large() { + // statistics are stored using i64, so test trying to store larger values + let statistics = make_bool_stats(Some(u64::MAX), Some(100)); + let thrift_stats = page_stats_to_thrift(Some(&statistics)).unwrap(); + assert_eq!(thrift_stats.distinct_count, None); // can't store u64 max --> null + assert_eq!(thrift_stats.null_count, Some(100)); + } + + #[test] + fn test_count_encoding_null_too_large() { + // statistics are stored using i64, so test trying to store larger values + let statistics = make_bool_stats(Some(100), Some(u64::MAX)); + let thrift_stats = page_stats_to_thrift(Some(&statistics)).unwrap(); + assert_eq!(thrift_stats.distinct_count, Some(100)); + assert_eq!(thrift_stats.null_count, None); // can' store u64 max --> null + } + + #[test] + fn test_count_decoding_null_invalid() { + let tstatistics = PageStatistics { + null_count: Some(-42), + max: None, + min: None, + distinct_count: None, + max_value: None, + min_value: None, + is_max_value_exact: None, + is_min_value_exact: None, + }; + let err = from_thrift_page_stats(Type::BOOLEAN, Some(tstatistics)).unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: Statistics null count is negative -42" + ); + } + + /// Writes statistics to thrift and reads them back and ensures: + /// - The statistics are the same + /// - The statistics written to thrift are the same as the original statistics + fn statistics_count_test(distinct_count: Option, null_count: Option) { + let statistics = make_bool_stats(distinct_count, null_count); + + let thrift_stats = page_stats_to_thrift(Some(&statistics)).unwrap(); + assert_eq!(thrift_stats.null_count.map(|c| c as u64), null_count); + assert_eq!( + thrift_stats.distinct_count.map(|c| c as u64), + distinct_count + ); + + let round_tripped = from_thrift_page_stats(Type::BOOLEAN, Some(thrift_stats)) + .unwrap() + .unwrap(); + assert_eq!(round_tripped, statistics); + } + + fn make_bool_stats(distinct_count: Option, null_count: Option) -> Statistics { + let min = Some(true); + let max = Some(false); + let is_min_max_deprecated = false; + + // test is about the counts, so we aren't really testing the min/max values + Statistics::Boolean(ValueStatistics::new( + min, + max, + distinct_count, + null_count, + is_min_max_deprecated, + )) + } + + #[test] + fn test_int96_invalid_statistics() { + let mut thrift_stats = PageStatistics { + max: None, + min: Some((0..13).collect()), + null_count: Some(0), + distinct_count: None, + max_value: None, + min_value: None, + is_max_value_exact: None, + is_min_value_exact: None, + }; + + let err = from_thrift_page_stats(Type::INT96, Some(thrift_stats.clone())).unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: Incorrect Int96 min statistics" + ); + + thrift_stats.min = None; + thrift_stats.max = Some((0..13).collect()); + let err = from_thrift_page_stats(Type::INT96, Some(thrift_stats)).unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: Incorrect Int96 max statistics" + ); + } + + // Ensures that we can call ValueStatistics::min_opt from a + // generic function without reyling on a bound to a private trait. + fn generic_statistics_handler(stats: ValueStatistics) -> String { + match stats.min_opt() { + Some(s) => format!("min: {}", s), + None => "min: NA".to_string(), + } + } + + #[test] + fn test_generic_access() { + let stats = Statistics::int32(Some(12), Some(45), None, Some(11), false); + + match stats { + Statistics::Int32(v) => { + let stats_string = generic_statistics_handler(v); + assert_eq!(&stats_string, "min: 12"); + } + _ => unreachable!(), + } + } +} diff --git a/vendor/parquet/src/file/writer.rs b/vendor/parquet/src/file/writer.rs new file mode 100644 index 0000000..b62d888 --- /dev/null +++ b/vendor/parquet/src/file/writer.rs @@ -0,0 +1,2608 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`SerializedFileWriter`]: Low level Parquet writer API + +use crate::bloom_filter::Sbbf; +use crate::file::metadata::thrift::PageHeader; +use crate::file::page_index::column_index::ColumnIndexMetaData; +use crate::file::page_index::offset_index::OffsetIndexMetaData; +use crate::parquet_thrift::{ThriftCompactOutputProtocol, WriteThrift}; +#[cfg(feature = "arrow")] +use bytes::Bytes; +use std::fmt::Debug; +use std::io::{BufWriter, IoSlice, Read}; +use std::{io::Write, sync::Arc}; + +use crate::column::page_encryption::PageEncryptor; +use crate::column::writer::{ColumnCloseResult, ColumnWriterImpl, get_typed_column_writer_mut}; +use crate::column::{ + page::{CompressedPage, PageWriteSpec, PageWriter}, + writer::{ColumnWriter, get_column_writer}, +}; +use crate::data_type::DataType; +#[cfg(feature = "encryption")] +use crate::encryption::encrypt::{ + FileEncryptionProperties, FileEncryptor, get_column_crypto_metadata, +}; +use crate::errors::{ParquetError, Result}; +#[cfg(feature = "encryption")] +use crate::file::PARQUET_MAGIC_ENCR_FOOTER; +use crate::file::properties::{BloomFilterPosition, WriterPropertiesPtr}; +use crate::file::reader::ChunkReader; +use crate::file::{PARQUET_MAGIC, metadata::*}; +use crate::schema::types::{ColumnDescPtr, SchemaDescPtr, SchemaDescriptor, TypePtr}; + +/// A wrapper around a [`Write`] that keeps track of the number +/// of bytes that have been written. The given [`Write`] is wrapped +/// with a [`BufWriter`] to optimize writing performance. +pub struct TrackedWrite { + inner: BufWriter, + bytes_written: usize, +} + +impl TrackedWrite { + /// Create a new [`TrackedWrite`] from a [`Write`] + pub fn new(inner: W) -> Self { + let buf_write = BufWriter::new(inner); + Self { + inner: buf_write, + bytes_written: 0, + } + } + + /// Returns the number of bytes written to this instance + pub fn bytes_written(&self) -> usize { + self.bytes_written + } + + /// Returns a reference to the underlying writer. + pub fn inner(&self) -> &W { + self.inner.get_ref() + } + + /// Returns a mutable reference to the underlying writer. + /// + /// It is inadvisable to directly write to the underlying writer, doing so + /// will likely result in data corruption + pub fn inner_mut(&mut self) -> &mut W { + self.inner.get_mut() + } + + /// Returns the underlying writer. + pub fn into_inner(self) -> Result { + self.inner.into_inner().map_err(|err| { + ParquetError::General(format!("fail to get inner writer: {:?}", err.to_string())) + }) + } +} + +impl Write for TrackedWrite { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + let bytes = self.inner.write(buf)?; + self.bytes_written += bytes; + Ok(bytes) + } + + fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> std::io::Result { + let bytes = self.inner.write_vectored(bufs)?; + self.bytes_written += bytes; + Ok(bytes) + } + + fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> { + self.inner.write_all(buf)?; + self.bytes_written += buf.len(); + + Ok(()) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.inner.flush() + } +} + +/// Callback invoked on closing a column chunk +pub type OnCloseColumnChunk<'a> = Box Result<()> + 'a>; + +/// Callback invoked on closing a row group, arguments are: +/// +/// - the row group metadata +/// - the column index for each column chunk +/// - the offset index for each column chunk +pub type OnCloseRowGroup<'a, W> = Box< + dyn FnOnce( + &'a mut TrackedWrite, + RowGroupMetaData, + Vec>, + Vec>, + Vec>, + ) -> Result<()> + + 'a + + Send, +>; + +// ---------------------------------------------------------------------- +// Serialized impl for file & row group writers + +/// Parquet file writer API. +/// +/// This is a low level API for writing Parquet files directly, and handles +/// tracking the location of file structures such as row groups and column +/// chunks, and writing the metadata and file footer. +/// +/// Data is written to row groups using [`SerializedRowGroupWriter`] and +/// columns using [`SerializedColumnWriter`]. The `SerializedFileWriter` tracks +/// where all the data is written, and assembles the final file metadata. +/// +/// The main workflow should be as following: +/// - Create file writer, this will open a new file and potentially write some metadata. +/// - Request a new row group writer by calling `next_row_group`. +/// - Once finished writing row group, close row group writer by calling `close` +/// - Write subsequent row groups, if necessary. +/// - After all row groups have been written, close the file writer using `close` method. +pub struct SerializedFileWriter { + buf: TrackedWrite, + descr: SchemaDescPtr, + props: WriterPropertiesPtr, + row_groups: Vec, + bloom_filters: Vec>>, + column_indexes: Vec>>, + offset_indexes: Vec>>, + row_group_index: usize, + // kv_metadatas will be appended to `props` when `write_metadata` + kv_metadatas: Vec, + finished: bool, + #[cfg(feature = "encryption")] + file_encryptor: Option>, +} + +impl Debug for SerializedFileWriter { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // implement Debug so this can be used with #[derive(Debug)] + // in client code rather than actually listing all the fields + f.debug_struct("SerializedFileWriter") + .field("descr", &self.descr) + .field("row_group_index", &self.row_group_index) + .field("kv_metadatas", &self.kv_metadatas) + .finish_non_exhaustive() + } +} + +impl SerializedFileWriter { + /// Creates new file writer. + pub fn new(buf: W, schema: TypePtr, properties: WriterPropertiesPtr) -> Result { + let mut buf = TrackedWrite::new(buf); + + let schema_descriptor = SchemaDescriptor::new(schema.clone()); + + #[cfg(feature = "encryption")] + let file_encryptor = Self::get_file_encryptor(&properties, &schema_descriptor)?; + + Self::start_file(&properties, &mut buf)?; + Ok(Self { + buf, + descr: Arc::new(schema_descriptor), + props: properties, + row_groups: vec![], + bloom_filters: vec![], + column_indexes: Vec::new(), + offset_indexes: Vec::new(), + row_group_index: 0, + kv_metadatas: Vec::new(), + finished: false, + #[cfg(feature = "encryption")] + file_encryptor, + }) + } + + #[cfg(feature = "encryption")] + fn get_file_encryptor( + properties: &WriterPropertiesPtr, + schema_descriptor: &SchemaDescriptor, + ) -> Result>> { + if let Some(file_encryption_properties) = properties.file_encryption_properties() { + file_encryption_properties.validate_encrypted_column_names(schema_descriptor)?; + + Ok(Some(Arc::new(FileEncryptor::new(Arc::clone( + file_encryption_properties, + ))?))) + } else { + Ok(None) + } + } + + /// Creates new row group from this file writer. + /// + /// Note: Parquet files are limited to at most 2^15 row groups in a file; and row groups must + /// be written sequentially. + /// + /// Every time the next row group is requested, the previous row group must + /// be finalised and closed using the [`SerializedRowGroupWriter::close`] + /// method or an error will be returned. + pub fn next_row_group(&mut self) -> Result> { + self.assert_previous_writer_closed()?; + let ordinal = self.row_group_index; + + let ordinal: i16 = ordinal.try_into().map_err(|_| { + ParquetError::General(format!( + "Parquet does not support more than {} row groups per file (currently: {})", + i16::MAX, + ordinal + )) + })?; + + self.row_group_index = self + .row_group_index + .checked_add(1) + .expect("SerializedFileWriter::row_group_index overflowed"); + + let bloom_filter_position = self.properties().bloom_filter_position(); + let row_groups = &mut self.row_groups; + let row_bloom_filters = &mut self.bloom_filters; + let row_column_indexes = &mut self.column_indexes; + let row_offset_indexes = &mut self.offset_indexes; + let on_close = move |buf, + mut metadata, + row_group_bloom_filter, + row_group_column_index, + row_group_offset_index| { + row_bloom_filters.push(row_group_bloom_filter); + row_column_indexes.push(row_group_column_index); + row_offset_indexes.push(row_group_offset_index); + // write bloom filters out immediately after the row group if requested + match bloom_filter_position { + BloomFilterPosition::AfterRowGroup => { + write_bloom_filters(buf, row_bloom_filters, &mut metadata)? + } + BloomFilterPosition::End => (), + }; + row_groups.push(metadata); + Ok(()) + }; + + let row_group_writer = SerializedRowGroupWriter::new( + self.descr.clone(), + self.props.clone(), + &mut self.buf, + ordinal, + Some(Box::new(on_close)), + ); + #[cfg(feature = "encryption")] + let row_group_writer = row_group_writer.with_file_encryptor(self.file_encryptor.clone()); + + Ok(row_group_writer) + } + + /// Returns metadata for any flushed row groups + pub fn flushed_row_groups(&self) -> &[RowGroupMetaData] { + &self.row_groups + } + + /// Close and finalize the underlying Parquet writer + /// + /// Unlike [`Self::close`] this does not consume self + /// + /// Attempting to write after calling finish will result in an error + pub fn finish(&mut self) -> Result { + self.assert_previous_writer_closed()?; + let metadata = self.write_metadata()?; + self.buf.flush()?; + Ok(metadata) + } + + /// Closes and finalises file writer, returning the file metadata. + pub fn close(mut self) -> Result { + self.finish() + } + + /// Writes magic bytes at the beginning of the file. + #[cfg(not(feature = "encryption"))] + fn start_file(_properties: &WriterPropertiesPtr, buf: &mut TrackedWrite) -> Result<()> { + buf.write_all(get_file_magic())?; + Ok(()) + } + + /// Writes magic bytes at the beginning of the file. + #[cfg(feature = "encryption")] + fn start_file(properties: &WriterPropertiesPtr, buf: &mut TrackedWrite) -> Result<()> { + let magic = get_file_magic(properties.file_encryption_properties.as_ref()); + + buf.write_all(magic)?; + Ok(()) + } + + /// Assembles and writes metadata at the end of the file. This will take ownership + /// of `row_groups` and the page index structures. + fn write_metadata(&mut self) -> Result { + self.finished = true; + + // write out any remaining bloom filters after all row groups + for row_group in &mut self.row_groups { + write_bloom_filters(&mut self.buf, &mut self.bloom_filters, row_group)?; + } + + let key_value_metadata = match self.props.key_value_metadata() { + Some(kv) => Some(kv.iter().chain(&self.kv_metadatas).cloned().collect()), + None if self.kv_metadatas.is_empty() => None, + None => Some(self.kv_metadatas.clone()), + }; + + // take ownership of metadata + let row_groups = std::mem::take(&mut self.row_groups); + let column_indexes = std::mem::take(&mut self.column_indexes); + let offset_indexes = std::mem::take(&mut self.offset_indexes); + + let write_path_in_schema = self.props.write_path_in_schema(); + let mut encoder = ThriftMetadataWriter::new( + &mut self.buf, + &self.descr, + row_groups, + Some(self.props.created_by().to_string()), + self.props.writer_version().as_num(), + write_path_in_schema, + ); + + #[cfg(feature = "encryption")] + { + encoder = encoder.with_file_encryptor(self.file_encryptor.clone()); + } + + if let Some(key_value_metadata) = key_value_metadata { + encoder = encoder.with_key_value_metadata(key_value_metadata) + } + + encoder = encoder.with_column_indexes(column_indexes); + if !self.props.offset_index_disabled() { + encoder = encoder.with_offset_indexes(offset_indexes); + } + encoder.finish() + } + + #[inline] + fn assert_previous_writer_closed(&self) -> Result<()> { + if self.finished { + return Err(general_err!("SerializedFileWriter already finished")); + } + + if self.row_group_index != self.row_groups.len() { + Err(general_err!("Previous row group writer was not closed")) + } else { + Ok(()) + } + } + + /// Add a [`KeyValue`] to the file writer's metadata + pub fn append_key_value_metadata(&mut self, kv_metadata: KeyValue) { + self.kv_metadatas.push(kv_metadata); + } + + /// Returns a reference to schema descriptor. + pub fn schema_descr(&self) -> &SchemaDescriptor { + &self.descr + } + + /// Returns a reference to schema descriptor Arc. + #[cfg(feature = "arrow")] + pub(crate) fn schema_descr_ptr(&self) -> &SchemaDescPtr { + &self.descr + } + + /// Returns a reference to the writer properties + pub fn properties(&self) -> &WriterPropertiesPtr { + &self.props + } + + /// Returns a reference to the underlying writer. + pub fn inner(&self) -> &W { + self.buf.inner() + } + + /// Writes the given buf bytes to the internal buffer. + /// + /// This can be used to write raw data to an in-progress Parquet file, for + /// example, custom index structures or other payloads. Other Parquet readers + /// will skip this data when reading the files. + /// + /// It's safe to use this method to write data to the underlying writer, + /// because it will ensure that the buffering and byte‐counting layers are used. + pub fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> { + self.buf.write_all(buf) + } + + /// Flushes underlying writer + pub fn flush(&mut self) -> std::io::Result<()> { + self.buf.flush() + } + + /// Returns a mutable reference to the underlying writer. + /// + /// **Warning**: if you write directly to this writer, you will skip + /// the `TrackedWrite` buffering and byte‐counting layers, which can cause + /// the file footer’s recorded offsets and sizes to diverge from reality, + /// resulting in an unreadable or corrupted Parquet file. + /// + /// If you want to write safely to the underlying writer, use [`Self::write_all`]. + pub fn inner_mut(&mut self) -> &mut W { + self.buf.inner_mut() + } + + /// Writes the file footer and returns the underlying writer. + pub fn into_inner(mut self) -> Result { + self.assert_previous_writer_closed()?; + let _ = self.write_metadata()?; + + self.buf.into_inner() + } + + /// Returns the number of bytes written to this instance + pub fn bytes_written(&self) -> usize { + self.buf.bytes_written() + } + + /// Get the file encryptor used by this instance to encrypt data + #[cfg(feature = "encryption")] + pub(crate) fn file_encryptor(&self) -> Option> { + self.file_encryptor.clone() + } +} + +/// Serialize all the bloom filters of the given row group to the given buffer, +/// and returns the updated row group metadata. +fn write_bloom_filters( + buf: &mut TrackedWrite, + bloom_filters: &mut [Vec>], + row_group: &mut RowGroupMetaData, +) -> Result<()> { + // iter row group + // iter each column + // write bloom filter to the file + + let row_group_idx: u16 = row_group + .ordinal() + .expect("Missing row group ordinal") + .try_into() + .map_err(|_| { + ParquetError::General(format!( + "Negative row group ordinal: {})", + row_group.ordinal().unwrap() + )) + })?; + let row_group_idx = row_group_idx as usize; + for (column_idx, column_chunk) in row_group.columns_mut().iter_mut().enumerate() { + if let Some(bloom_filter) = bloom_filters[row_group_idx][column_idx].take() { + let start_offset = buf.bytes_written(); + bloom_filter.write(&mut *buf)?; + let end_offset = buf.bytes_written(); + // set offset and index for bloom filter + *column_chunk = column_chunk + .clone() + .into_builder() + .set_bloom_filter_offset(Some(start_offset as i64)) + .set_bloom_filter_length(Some((end_offset - start_offset) as i32)) + .build()?; + } + } + Ok(()) +} + +/// Parquet row group writer API. +/// +/// Provides methods to access column writers in an iterator-like fashion, order is +/// guaranteed to match the order of schema leaves (column descriptors). +/// +/// All columns should be written sequentially; the main workflow is: +/// - Request the next column using `next_column` method - this will return `None` if no +/// more columns are available to write. +/// - Once done writing a column, close column writer with `close` +/// - Once all columns have been written, close row group writer with `close` +/// method. The close method will return row group metadata and is no-op +/// on already closed row group. +pub struct SerializedRowGroupWriter<'a, W: Write> { + descr: SchemaDescPtr, + props: WriterPropertiesPtr, + buf: &'a mut TrackedWrite, + total_rows_written: Option, + total_bytes_written: u64, + total_uncompressed_bytes: i64, + column_index: usize, + row_group_metadata: Option, + column_chunks: Vec, + bloom_filters: Vec>, + column_indexes: Vec>, + offset_indexes: Vec>, + row_group_index: i16, + file_offset: i64, + on_close: Option>, + #[cfg(feature = "encryption")] + file_encryptor: Option>, +} + +impl<'a, W: Write + Send> SerializedRowGroupWriter<'a, W> { + /// Creates a new `SerializedRowGroupWriter` with: + /// + /// - `schema_descr` - the schema to write + /// - `properties` - writer properties + /// - `buf` - the buffer to write data to + /// - `row_group_index` - row group index in this parquet file. + /// - `file_offset` - file offset of this row group in this parquet file. + /// - `on_close` - an optional callback that will invoked on [`Self::close`] + pub fn new( + schema_descr: SchemaDescPtr, + properties: WriterPropertiesPtr, + buf: &'a mut TrackedWrite, + row_group_index: i16, + on_close: Option>, + ) -> Self { + let num_columns = schema_descr.num_columns(); + let file_offset = buf.bytes_written() as i64; + Self { + buf, + row_group_index, + file_offset, + on_close, + total_rows_written: None, + descr: schema_descr, + props: properties, + column_index: 0, + row_group_metadata: None, + column_chunks: Vec::with_capacity(num_columns), + bloom_filters: Vec::with_capacity(num_columns), + column_indexes: Vec::with_capacity(num_columns), + offset_indexes: Vec::with_capacity(num_columns), + total_bytes_written: 0, + total_uncompressed_bytes: 0, + #[cfg(feature = "encryption")] + file_encryptor: None, + } + } + + #[cfg(feature = "encryption")] + /// Set the file encryptor to use for encrypting row group data and metadata + pub(crate) fn with_file_encryptor( + mut self, + file_encryptor: Option>, + ) -> Self { + self.file_encryptor = file_encryptor; + self + } + + /// Advance `self.column_index` returning the next [`ColumnDescPtr`] if any + fn next_column_desc(&mut self) -> Option { + let ret = self.descr.columns().get(self.column_index)?.clone(); + self.column_index += 1; + Some(ret) + } + + /// Returns [`OnCloseColumnChunk`] for the next writer + fn get_on_close(&mut self) -> (&mut TrackedWrite, OnCloseColumnChunk<'_>) { + let total_bytes_written = &mut self.total_bytes_written; + let total_uncompressed_bytes = &mut self.total_uncompressed_bytes; + let total_rows_written = &mut self.total_rows_written; + let column_chunks = &mut self.column_chunks; + let column_indexes = &mut self.column_indexes; + let offset_indexes = &mut self.offset_indexes; + let bloom_filters = &mut self.bloom_filters; + + let on_close = |r: ColumnCloseResult| { + // Update row group writer metrics + *total_bytes_written += r.bytes_written; + *total_uncompressed_bytes += r.metadata.uncompressed_size(); + column_chunks.push(r.metadata); + bloom_filters.push(r.bloom_filter); + column_indexes.push(r.column_index); + offset_indexes.push(r.offset_index); + + if let Some(rows) = *total_rows_written { + if rows != r.rows_written { + return Err(general_err!( + "Incorrect number of rows, expected {} != {} rows", + rows, + r.rows_written + )); + } + } else { + *total_rows_written = Some(r.rows_written); + } + + Ok(()) + }; + (self.buf, Box::new(on_close)) + } + + /// Returns the next column writer, if available, using the factory function; + /// otherwise returns `None`. + pub(crate) fn next_column_with_factory<'b, F, C>(&'b mut self, factory: F) -> Result> + where + F: FnOnce( + ColumnDescPtr, + WriterPropertiesPtr, + Box, + OnCloseColumnChunk<'b>, + ) -> Result, + { + self.assert_previous_writer_closed()?; + + let encryptor_context = self.get_page_encryptor_context(); + + Ok(match self.next_column_desc() { + Some(column) => { + let props = self.props.clone(); + let (buf, on_close) = self.get_on_close(); + + let page_writer = SerializedPageWriter::new(buf); + let page_writer = + Self::set_page_writer_encryptor(&column, encryptor_context, page_writer)?; + + Some(factory( + column, + props, + Box::new(page_writer), + Box::new(on_close), + )?) + } + None => None, + }) + } + + /// Returns the next column writer, if available; otherwise returns `None`. + /// In case of any IO error or Thrift error, or if row group writer has already been + /// closed returns `Err`. + pub fn next_column(&mut self) -> Result>> { + self.next_column_with_factory(|descr, props, page_writer, on_close| { + let column_writer = get_column_writer(descr, props, page_writer); + Ok(SerializedColumnWriter::new(column_writer, Some(on_close))) + }) + } + + /// Append an encoded column chunk from `reader` directly to the underlying + /// writer. + /// + /// This method can be used for efficiently concatenating or projecting + /// Parquet data, or encoding Parquet data to temporary in-memory buffers. + /// + /// Arguments: + /// - `reader`: a [`ChunkReader`] containing the encoded column data + /// - `close`: the [`ColumnCloseResult`] metadata returned from closing + /// the column writer that wrote the data in `reader`. + /// + /// See Also: + /// 1. [`get_column_writer`] for creating writers that can encode data. + /// 2. [`Self::next_column`] for writing data that isn't already encoded + pub fn append_column( + &mut self, + reader: &R, + close: ColumnCloseResult, + ) -> Result<()> { + // Position a reader at the start of the buffered chunk, then splice the + // bytes through the shared streaming path. + let metadata = &close.metadata; + let src_offset = metadata + .dictionary_page_offset() + .unwrap_or_else(|| metadata.data_page_offset()); + let read = reader.get_read(src_offset as _)?; + self.append_column_from_read(read, close) + } + + /// Splice an already-encoded column chunk into the row group, reading its + /// bytes sequentially from `read`. + /// + /// `read` must be positioned at the start of the chunk (the dictionary page + /// if present, otherwise the first data page — i.e. `src_offset` below) and + /// yield exactly the chunk's compressed bytes. Unlike [`Self::append_column`] + /// this consumes an owned [`Read`], which lets the caller stream the bytes + /// back from a [`PageStore`](crate::column::page_store::PageStore) one page + /// at a time without materializing the whole chunk in memory. + pub(crate) fn append_column_from_read( + &mut self, + read: R, + close: ColumnCloseResult, + ) -> Result<()> { + 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)?; + + if src_length as u64 != write_length { + return Err(general_err!( + "Failed to splice column data, expected {src_length} got {write_length}" + )); + } + + self.finish_appended_column(close, src_offset, write_offset) + } + + /// Splice an already-encoded column chunk into the row group from an + /// in-order sequence of byte buffers (typically its serialized pages). + /// + /// This is a lower-overhead alternative to [`Self::append_column`] / + /// [`Self::append_column_from_read`] for callers that already hold the + /// chunk as owned [`Bytes`]: each buffer is written straight to the output + /// with a single `write_all`, skipping the intermediate copy through + /// [`std::io::copy`]'s fixed-size buffer. + /// + /// `pages` must yield the chunk's compressed bytes in final file order + /// (the dictionary page, if any, first) and together total exactly the + /// compressed size recorded in `close`. + #[cfg(feature = "arrow")] + pub(crate) fn append_column_from_pages( + &mut self, + pages: I, + close: ColumnCloseResult, + ) -> Result<()> + where + I: IntoIterator>, + { + let (src_offset, src_length, write_offset) = self.begin_appended_column(&close)?; + + let mut write_length = 0u64; + for page in pages { + let page = page?; + self.buf.write_all(&page)?; + write_length += page.len() as u64; + } + + if src_length as u64 != write_length { + return Err(general_err!( + "Failed to splice column data, expected {src_length} got {write_length}" + )); + } + + self.finish_appended_column(close, src_offset, write_offset) + } + + /// [`Self::append_column_from_read`] / [`Self::append_column_from_pages`] + /// preamble: validates the writer state and that `close` matches the next + /// expected column. + /// + /// Returns `(src_offset, src_length, write_offset)`: the chunk's start + /// offset and length in the source buffer, and the offset at which it will + /// land in the output file. + fn begin_appended_column(&mut self, close: &ColumnCloseResult) -> Result<(i64, i64, usize)> { + self.assert_previous_writer_closed()?; + let desc = self + .next_column_desc() + .ok_or_else(|| general_err!("exhausted columns in SerializedRowGroupWriter"))?; + + let metadata = &close.metadata; + + if metadata.column_descr() != desc.as_ref() { + return Err(general_err!( + "column descriptor mismatch, expected {:?} got {:?}", + desc, + metadata.column_descr() + )); + } + + let src_offset = metadata + .dictionary_page_offset() + .unwrap_or_else(|| metadata.data_page_offset()); + let src_length = metadata.compressed_size(); + let write_offset = self.buf.bytes_written(); + Ok((src_offset, src_length, write_offset)) + } + + /// [`Self::append_column_from_read`] / [`Self::append_column_from_pages`] + /// epilogue: rewrites the buffer-relative page offsets recorded in `close` + /// to their final positions in the output file and closes the column. + fn finish_appended_column( + &mut self, + mut close: ColumnCloseResult, + src_offset: i64, + write_offset: usize, + ) -> Result<()> { + let metadata = close.metadata; + let src_dictionary_offset = metadata.dictionary_page_offset(); + let src_data_offset = metadata.data_page_offset(); + + let map_offset = |x| x - src_offset + write_offset as i64; + let mut builder = ColumnChunkMetaData::builder(metadata.column_descr_ptr()) + .set_compression_codec(metadata.compression_codec()) + .set_encodings_mask(*metadata.encodings_mask()) + .set_total_compressed_size(metadata.compressed_size()) + .set_total_uncompressed_size(metadata.uncompressed_size()) + .set_num_values(metadata.num_values()) + .set_data_page_offset(map_offset(src_data_offset)) + .set_dictionary_page_offset(src_dictionary_offset.map(map_offset)) + .set_unencoded_byte_array_data_bytes(metadata.unencoded_byte_array_data_bytes()); + + if let Some(rep_hist) = metadata.repetition_level_histogram() { + builder = builder.set_repetition_level_histogram(Some(rep_hist.clone())) + } + if let Some(def_hist) = metadata.definition_level_histogram() { + builder = builder.set_definition_level_histogram(Some(def_hist.clone())) + } + if let Some(statistics) = metadata.statistics() { + builder = builder.set_statistics(statistics.clone()) + } + if let Some(geo_statistics) = metadata.geo_statistics() { + builder = builder.set_geo_statistics(Box::new(geo_statistics.clone())) + } + if let Some(page_encoding_stats) = metadata.page_encoding_stats() { + builder = builder.set_page_encoding_stats(page_encoding_stats.clone()) + } + builder = self.set_column_crypto_metadata(builder, &metadata); + close.metadata = builder.build()?; + + if let Some(offsets) = close.offset_index.as_mut() { + for location in &mut offsets.page_locations { + location.offset = map_offset(location.offset) + } + } + + let (_, on_close) = self.get_on_close(); + on_close(close) + } + + /// Closes this row group writer and returns row group metadata. + pub fn close(mut self) -> Result { + if self.row_group_metadata.is_none() { + self.assert_previous_writer_closed()?; + + let column_chunks = std::mem::take(&mut self.column_chunks); + let row_group_metadata = RowGroupMetaData::builder(self.descr.clone()) + .set_column_metadata(column_chunks) + .set_total_byte_size(self.total_uncompressed_bytes) + .set_num_rows(self.total_rows_written.unwrap_or(0) as i64) + .set_sorting_columns(self.props.sorting_columns().cloned()) + .set_ordinal(self.row_group_index) + .set_file_offset(self.file_offset) + .build()?; + + self.row_group_metadata = Some(Arc::new(row_group_metadata.clone())); + + if let Some(on_close) = self.on_close.take() { + on_close( + self.buf, + row_group_metadata, + self.bloom_filters, + self.column_indexes, + self.offset_indexes, + )? + } + } + + let metadata = self.row_group_metadata.as_ref().unwrap().clone(); + Ok(metadata) + } + + /// Set the column crypto metadata for a column chunk + #[cfg(feature = "encryption")] + fn set_column_crypto_metadata( + &self, + builder: ColumnChunkMetaDataBuilder, + metadata: &ColumnChunkMetaData, + ) -> ColumnChunkMetaDataBuilder { + if let Some(file_encryptor) = self.file_encryptor.as_ref() { + builder.set_column_crypto_metadata(get_column_crypto_metadata( + file_encryptor.properties(), + &metadata.column_descr_ptr(), + )) + } else { + builder + } + } + + /// Get context required to create a [`PageEncryptor`] for a column + #[cfg(feature = "encryption")] + fn get_page_encryptor_context(&self) -> PageEncryptorContext { + PageEncryptorContext { + file_encryptor: self.file_encryptor.clone(), + row_group_index: self.row_group_index as usize, + column_index: self.column_index, + } + } + + /// Set the [`PageEncryptor`] on a page writer if a column is encrypted + #[cfg(feature = "encryption")] + fn set_page_writer_encryptor<'b>( + column: &ColumnDescPtr, + context: PageEncryptorContext, + page_writer: SerializedPageWriter<'b, W>, + ) -> Result> { + let page_encryptor = PageEncryptor::create_if_column_encrypted( + &context.file_encryptor, + context.row_group_index, + context.column_index, + &column.path().string(), + )?; + + Ok(page_writer.with_page_encryptor(page_encryptor)) + } + + /// No-op implementation of setting the column crypto metadata for a column chunk + #[cfg(not(feature = "encryption"))] + fn set_column_crypto_metadata( + &self, + builder: ColumnChunkMetaDataBuilder, + _metadata: &ColumnChunkMetaData, + ) -> ColumnChunkMetaDataBuilder { + builder + } + + #[cfg(not(feature = "encryption"))] + fn get_page_encryptor_context(&self) -> PageEncryptorContext { + PageEncryptorContext {} + } + + /// No-op implementation of setting a [`PageEncryptor`] for when encryption is disabled + #[cfg(not(feature = "encryption"))] + fn set_page_writer_encryptor<'b>( + _column: &ColumnDescPtr, + _context: PageEncryptorContext, + page_writer: SerializedPageWriter<'b, W>, + ) -> Result> { + Ok(page_writer) + } + + #[inline] + fn assert_previous_writer_closed(&self) -> Result<()> { + if self.column_index != self.column_chunks.len() { + Err(general_err!("Previous column writer was not closed")) + } else { + Ok(()) + } + } +} + +/// Context required to create a [`PageEncryptor`] for a column +#[cfg(feature = "encryption")] +struct PageEncryptorContext { + file_encryptor: Option>, + row_group_index: usize, + column_index: usize, +} + +#[cfg(not(feature = "encryption"))] +struct PageEncryptorContext {} + +/// A wrapper around a [`ColumnWriter`] that invokes a callback on [`Self::close`] +pub struct SerializedColumnWriter<'a> { + inner: ColumnWriter<'a>, + on_close: Option>, +} + +impl<'a> SerializedColumnWriter<'a> { + /// Create a new [`SerializedColumnWriter`] from a [`ColumnWriter`] and an + /// optional callback to be invoked on [`Self::close`] + pub fn new(inner: ColumnWriter<'a>, on_close: Option>) -> Self { + Self { inner, on_close } + } + + /// Returns a reference to an untyped [`ColumnWriter`] + pub fn untyped(&mut self) -> &mut ColumnWriter<'a> { + &mut self.inner + } + + /// Returns a reference to a typed [`ColumnWriterImpl`] + pub fn typed(&mut self) -> &mut ColumnWriterImpl<'a, T> { + get_typed_column_writer_mut(&mut self.inner) + } + + /// Close this [`SerializedColumnWriter`] + pub fn close(mut self) -> Result<()> { + let r = self.inner.close()?; + if let Some(on_close) = self.on_close.take() { + on_close(r)? + } + + Ok(()) + } +} + +/// A serialized implementation for Parquet [`PageWriter`]. +/// Writes and serializes pages and metadata into output stream. +/// +/// `SerializedPageWriter` should not be used after calling `close()`. +pub struct SerializedPageWriter<'a, W: Write> { + sink: &'a mut TrackedWrite, + #[cfg(feature = "encryption")] + page_encryptor: Option, +} + +impl<'a, W: Write> SerializedPageWriter<'a, W> { + /// Creates new page writer. + pub fn new(sink: &'a mut TrackedWrite) -> Self { + Self { + sink, + #[cfg(feature = "encryption")] + page_encryptor: None, + } + } + + /// Serializes page header into Thrift. + /// Returns number of bytes that have been written into the sink. + #[inline] + fn serialize_page_header(&mut self, header: PageHeader) -> Result { + let start_pos = self.sink.bytes_written(); + match self.page_encryptor_and_sink_mut() { + Some((page_encryptor, sink)) => { + page_encryptor.encrypt_page_header(&header, sink)?; + } + None => { + let mut protocol = ThriftCompactOutputProtocol::new(&mut self.sink); + header.write_thrift(&mut protocol)?; + } + } + Ok(self.sink.bytes_written() - start_pos) + } +} + +#[cfg(feature = "encryption")] +impl<'a, W: Write> SerializedPageWriter<'a, W> { + /// Set the encryptor to use to encrypt page data + fn with_page_encryptor(mut self, page_encryptor: Option) -> Self { + self.page_encryptor = page_encryptor; + self + } + + fn page_encryptor_mut(&mut self) -> Option<&mut PageEncryptor> { + self.page_encryptor.as_mut() + } + + fn page_encryptor_and_sink_mut( + &mut self, + ) -> Option<(&mut PageEncryptor, &mut &'a mut TrackedWrite)> { + self.page_encryptor.as_mut().map(|pe| (pe, &mut self.sink)) + } +} + +#[cfg(not(feature = "encryption"))] +impl<'a, W: Write> SerializedPageWriter<'a, W> { + fn page_encryptor_mut(&mut self) -> Option<&mut PageEncryptor> { + None + } + + fn page_encryptor_and_sink_mut( + &mut self, + ) -> Option<(&mut PageEncryptor, &mut &'a mut TrackedWrite)> { + None + } +} + +impl PageWriter for SerializedPageWriter<'_, W> { + fn write_page(&mut self, page: CompressedPage) -> Result { + let page = match self.page_encryptor_mut() { + Some(page_encryptor) => page_encryptor.encrypt_compressed_page(page)?, + None => page, + }; + + let page_type = page.page_type(); + let start_pos = self.sink.bytes_written() as u64; + + let page_header = page.to_thrift_header()?; + let header_size = self.serialize_page_header(page_header)?; + + self.sink.write_all(page.data())?; + + let mut spec = PageWriteSpec::new(); + spec.page_type = page_type; + spec.uncompressed_size = page.uncompressed_size() + header_size; + spec.compressed_size = page.compressed_size() + header_size; + spec.offset = start_pos; + spec.bytes_written = self.sink.bytes_written() as u64 - start_pos; + spec.num_values = page.num_values(); + + if let Some(page_encryptor) = self.page_encryptor_mut() { + if page.compressed_page().is_data_page() { + page_encryptor.increment_page(); + } + } + Ok(spec) + } + + fn close(&mut self) -> Result<()> { + self.sink.flush()?; + Ok(()) + } +} + +/// Get the magic bytes at the start and end of the file that identify this +/// as a Parquet file. +#[cfg(feature = "encryption")] +pub(crate) fn get_file_magic( + file_encryption_properties: Option<&Arc>, +) -> &'static [u8; 4] { + match file_encryption_properties.as_ref() { + Some(encryption_properties) if encryption_properties.encrypt_footer() => { + &PARQUET_MAGIC_ENCR_FOOTER + } + _ => &PARQUET_MAGIC, + } +} + +#[cfg(not(feature = "encryption"))] +pub(crate) fn get_file_magic() -> &'static [u8; 4] { + &PARQUET_MAGIC +} + +#[cfg(test)] +mod tests { + use super::*; + + #[cfg(feature = "arrow")] + use arrow_array::RecordBatchReader; + use bytes::Bytes; + use std::fs::File; + + #[cfg(feature = "arrow")] + use crate::arrow::ArrowWriter; + #[cfg(feature = "arrow")] + use crate::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use crate::basic::{ + ColumnOrder, Compression, ConvertedType, Encoding, LogicalType, Repetition, SortOrder, Type, + }; + use crate::column::page::{Page, PageReader}; + use crate::column::reader::get_typed_column_reader; + use crate::compression::{Codec, CodecOptionsBuilder, create_codec}; + use crate::data_type::{BoolType, ByteArrayType, Int32Type}; + use crate::file::page_index::column_index::ColumnIndexMetaData; + use crate::file::properties::EnabledStatistics; + use crate::file::serialized_reader::ReadOptionsBuilder; + use crate::file::statistics::{from_thrift_page_stats, page_stats_to_thrift}; + use crate::file::{ + properties::{ReaderProperties, WriterProperties, WriterVersion}, + reader::{FileReader, SerializedFileReader, SerializedPageReader}, + statistics::Statistics, + }; + use crate::record::{Row, RowAccessor}; + use crate::schema::parser::parse_message_type; + use crate::schema::types; + use crate::schema::types::{ColumnDescriptor, ColumnPath}; + use crate::util::test_common::file_util::get_test_file; + use crate::util::test_common::rand_gen::RandGen; + + #[test] + fn test_row_group_writer_error_not_all_columns_written() { + let file = tempfile::tempfile().unwrap(); + let schema = Arc::new( + types::Type::group_type_builder("schema") + .with_fields(vec![Arc::new( + types::Type::primitive_type_builder("col1", Type::INT32) + .build() + .unwrap(), + )]) + .build() + .unwrap(), + ); + let props = Default::default(); + let mut writer = SerializedFileWriter::new(file, schema, props).unwrap(); + let row_group_writer = writer.next_row_group().unwrap(); + let res = row_group_writer.close(); + assert!(res.is_err()); + if let Err(err) = res { + assert_eq!( + format!("{err}"), + "Parquet error: Column length mismatch: 1 != 0" + ); + } + } + + #[test] + fn test_row_group_writer_num_records_mismatch() { + let file = tempfile::tempfile().unwrap(); + let schema = Arc::new( + types::Type::group_type_builder("schema") + .with_fields(vec![ + Arc::new( + types::Type::primitive_type_builder("col1", Type::INT32) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(), + ), + Arc::new( + types::Type::primitive_type_builder("col2", Type::INT32) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(), + ), + ]) + .build() + .unwrap(), + ); + let props = Default::default(); + let mut writer = SerializedFileWriter::new(file, schema, props).unwrap(); + let mut row_group_writer = writer.next_row_group().unwrap(); + + let mut col_writer = row_group_writer.next_column().unwrap().unwrap(); + col_writer + .typed::() + .write_batch(&[1, 2, 3], None, None) + .unwrap(); + col_writer.close().unwrap(); + + let mut col_writer = row_group_writer.next_column().unwrap().unwrap(); + col_writer + .typed::() + .write_batch(&[1, 2], None, None) + .unwrap(); + + let err = col_writer.close().unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: Incorrect number of rows, expected 3 != 2 rows" + ); + } + + #[test] + fn test_file_writer_empty_file() { + let file = tempfile::tempfile().unwrap(); + + let schema = Arc::new( + types::Type::group_type_builder("schema") + .with_fields(vec![Arc::new( + types::Type::primitive_type_builder("col1", Type::INT32) + .build() + .unwrap(), + )]) + .build() + .unwrap(), + ); + let props = Default::default(); + let writer = SerializedFileWriter::new(file.try_clone().unwrap(), schema, props).unwrap(); + writer.close().unwrap(); + + let reader = SerializedFileReader::new(file).unwrap(); + assert_eq!(reader.get_row_iter(None).unwrap().count(), 0); + } + + #[test] + fn test_file_writer_column_orders_populated() { + let file = tempfile::tempfile().unwrap(); + + let schema = Arc::new( + types::Type::group_type_builder("schema") + .with_fields(vec![ + Arc::new( + types::Type::primitive_type_builder("col1", Type::INT32) + .build() + .unwrap(), + ), + Arc::new( + types::Type::primitive_type_builder("col2", Type::FIXED_LEN_BYTE_ARRAY) + .with_converted_type(ConvertedType::INTERVAL) + .with_length(12) + .build() + .unwrap(), + ), + Arc::new( + types::Type::group_type_builder("nested") + .with_repetition(Repetition::REQUIRED) + .with_fields(vec![ + Arc::new( + types::Type::primitive_type_builder( + "col3", + Type::FIXED_LEN_BYTE_ARRAY, + ) + .with_logical_type(Some(LogicalType::Float16)) + .with_length(2) + .build() + .unwrap(), + ), + Arc::new( + types::Type::primitive_type_builder("col4", Type::BYTE_ARRAY) + .with_logical_type(Some(LogicalType::String)) + .build() + .unwrap(), + ), + ]) + .build() + .unwrap(), + ), + ]) + .build() + .unwrap(), + ); + + let props = Default::default(); + let writer = SerializedFileWriter::new(file.try_clone().unwrap(), schema, props).unwrap(); + writer.close().unwrap(); + + let reader = SerializedFileReader::new(file).unwrap(); + + // only leaves + let expected = vec![ + // INT32 + ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::SIGNED), + // INTERVAL + ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNDEFINED), + // Float16 + ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::SIGNED), + // String + ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED), + ]; + let actual = reader.metadata().file_metadata().column_orders(); + + assert!(actual.is_some()); + let actual = actual.unwrap(); + assert_eq!(*actual, expected); + } + + #[test] + fn test_file_writer_with_metadata() { + let file = tempfile::tempfile().unwrap(); + + let schema = Arc::new( + types::Type::group_type_builder("schema") + .with_fields(vec![Arc::new( + types::Type::primitive_type_builder("col1", Type::INT32) + .build() + .unwrap(), + )]) + .build() + .unwrap(), + ); + let props = Arc::new( + WriterProperties::builder() + .set_key_value_metadata(Some(vec![KeyValue::new( + "key".to_string(), + "value".to_string(), + )])) + .build(), + ); + let writer = SerializedFileWriter::new(file.try_clone().unwrap(), schema, props).unwrap(); + writer.close().unwrap(); + + let reader = SerializedFileReader::new(file).unwrap(); + assert_eq!( + reader + .metadata() + .file_metadata() + .key_value_metadata() + .to_owned() + .unwrap() + .len(), + 1 + ); + } + + #[test] + fn test_file_writer_v2_with_metadata() { + let file = tempfile::tempfile().unwrap(); + let field_logical_type = Some(LogicalType::integer(8, false)); + let field = Arc::new( + types::Type::primitive_type_builder("col1", Type::INT32) + .with_logical_type(field_logical_type.clone()) + .with_converted_type(field_logical_type.into()) + .build() + .unwrap(), + ); + let schema = Arc::new( + types::Type::group_type_builder("schema") + .with_fields(vec![field.clone()]) + .build() + .unwrap(), + ); + let props = Arc::new( + WriterProperties::builder() + .set_key_value_metadata(Some(vec![KeyValue::new( + "key".to_string(), + "value".to_string(), + )])) + .set_writer_version(WriterVersion::PARQUET_2_0) + .build(), + ); + let writer = SerializedFileWriter::new(file.try_clone().unwrap(), schema, props).unwrap(); + writer.close().unwrap(); + + let reader = SerializedFileReader::new(file).unwrap(); + + assert_eq!( + reader + .metadata() + .file_metadata() + .key_value_metadata() + .to_owned() + .unwrap() + .len(), + 1 + ); + + // ARROW-11803: Test that the converted and logical types have been populated + let fields = reader.metadata().file_metadata().schema().get_fields(); + assert_eq!(fields.len(), 1); + assert_eq!(fields[0], field); + } + + #[test] + fn test_file_writer_with_sorting_columns_metadata() { + let file = tempfile::tempfile().unwrap(); + + let schema = Arc::new( + types::Type::group_type_builder("schema") + .with_fields(vec![ + Arc::new( + types::Type::primitive_type_builder("col1", Type::INT32) + .build() + .unwrap(), + ), + Arc::new( + types::Type::primitive_type_builder("col2", Type::INT32) + .build() + .unwrap(), + ), + ]) + .build() + .unwrap(), + ); + let expected_result = Some(vec![SortingColumn { + column_idx: 0, + descending: false, + nulls_first: true, + }]); + let props = Arc::new( + WriterProperties::builder() + .set_key_value_metadata(Some(vec![KeyValue::new( + "key".to_string(), + "value".to_string(), + )])) + .set_sorting_columns(expected_result.clone()) + .build(), + ); + let mut writer = + SerializedFileWriter::new(file.try_clone().unwrap(), schema, props).unwrap(); + let mut row_group_writer = writer.next_row_group().expect("get row group writer"); + + let col_writer = row_group_writer.next_column().unwrap().unwrap(); + col_writer.close().unwrap(); + + let col_writer = row_group_writer.next_column().unwrap().unwrap(); + col_writer.close().unwrap(); + + row_group_writer.close().unwrap(); + writer.close().unwrap(); + + let reader = SerializedFileReader::new(file).unwrap(); + let result: Vec>> = reader + .metadata() + .row_groups() + .iter() + .map(|f| f.sorting_columns()) + .collect(); + // validate the sorting column read match the one written above + assert_eq!(expected_result.as_ref(), result[0]); + } + + #[test] + fn test_file_writer_empty_row_groups() { + let file = tempfile::tempfile().unwrap(); + test_file_roundtrip(file, vec![]); + } + + #[test] + fn test_file_writer_single_row_group() { + let file = tempfile::tempfile().unwrap(); + test_file_roundtrip(file, vec![vec![1, 2, 3, 4, 5]]); + } + + #[test] + fn test_file_writer_multiple_row_groups() { + let file = tempfile::tempfile().unwrap(); + test_file_roundtrip( + file, + vec![ + vec![1, 2, 3, 4, 5], + vec![1, 2, 3], + vec![1], + vec![1, 2, 3, 4, 5, 6], + ], + ); + } + + #[test] + fn test_file_writer_multiple_large_row_groups() { + let file = tempfile::tempfile().unwrap(); + test_file_roundtrip( + file, + vec![vec![123; 1024], vec![124; 1000], vec![125; 15], vec![]], + ); + } + + #[test] + fn test_page_writer_data_pages() { + let pages = [ + Page::DataPage { + buf: Bytes::from(vec![1, 2, 3, 4, 5, 6, 7, 8]), + num_values: 10, + encoding: Encoding::DELTA_BINARY_PACKED, + def_level_encoding: Encoding::RLE, + rep_level_encoding: Encoding::RLE, + statistics: Some(Statistics::int32(Some(1), Some(3), None, Some(7), true)), + }, + Page::DataPageV2 { + buf: Bytes::from(vec![4; 128]), + num_values: 10, + encoding: Encoding::DELTA_BINARY_PACKED, + num_nulls: 2, + num_rows: 12, + def_levels_byte_len: 24, + rep_levels_byte_len: 32, + is_compressed: false, + statistics: Some(Statistics::int32(Some(1), Some(3), None, Some(7), true)), + }, + ]; + + test_page_roundtrip(&pages[..], Compression::SNAPPY, Type::INT32); + test_page_roundtrip(&pages[..], Compression::UNCOMPRESSED, Type::INT32); + } + + #[test] + fn test_page_writer_dict_pages() { + let pages = [ + Page::DictionaryPage { + buf: Bytes::from(vec![1, 2, 3, 4, 5]), + num_values: 5, + encoding: Encoding::RLE_DICTIONARY, + is_sorted: false, + }, + Page::DataPage { + buf: Bytes::from(vec![1, 2, 3, 4, 5, 6, 7, 8]), + num_values: 10, + encoding: Encoding::DELTA_BINARY_PACKED, + def_level_encoding: Encoding::RLE, + rep_level_encoding: Encoding::RLE, + statistics: Some(Statistics::int32(Some(1), Some(3), None, Some(7), true)), + }, + Page::DataPageV2 { + buf: Bytes::from(vec![4; 128]), + num_values: 10, + encoding: Encoding::DELTA_BINARY_PACKED, + num_nulls: 2, + num_rows: 12, + def_levels_byte_len: 24, + rep_levels_byte_len: 32, + is_compressed: false, + statistics: None, + }, + ]; + + test_page_roundtrip(&pages[..], Compression::SNAPPY, Type::INT32); + test_page_roundtrip(&pages[..], Compression::UNCOMPRESSED, Type::INT32); + } + + /// Tests writing and reading pages. + /// Physical type is for statistics only, should match any defined statistics type in + /// pages. + fn test_page_roundtrip(pages: &[Page], codec: Compression, physical_type: Type) { + let mut compressed_pages = vec![]; + let mut total_num_values = 0i64; + let codec_options = CodecOptionsBuilder::default() + .set_backward_compatible_lz4(false) + .build(); + let mut compressor = create_codec(codec, &codec_options).unwrap(); + + for page in pages { + let uncompressed_len = page.buffer().len(); + + let compressed_page = match *page { + Page::DataPage { + ref buf, + num_values, + encoding, + def_level_encoding, + rep_level_encoding, + ref statistics, + } => { + total_num_values += num_values as i64; + let output_buf = compress_helper(compressor.as_mut(), buf); + + Page::DataPage { + buf: Bytes::from(output_buf), + num_values, + encoding, + def_level_encoding, + rep_level_encoding, + statistics: from_thrift_page_stats( + physical_type, + page_stats_to_thrift(statistics.as_ref()), + ) + .unwrap(), + } + } + Page::DataPageV2 { + ref buf, + num_values, + encoding, + num_nulls, + num_rows, + def_levels_byte_len, + rep_levels_byte_len, + ref statistics, + .. + } => { + total_num_values += num_values as i64; + let offset = (def_levels_byte_len + rep_levels_byte_len) as usize; + let cmp_buf = compress_helper(compressor.as_mut(), &buf[offset..]); + let mut output_buf = Vec::from(&buf[..offset]); + output_buf.extend_from_slice(&cmp_buf[..]); + + Page::DataPageV2 { + buf: Bytes::from(output_buf), + num_values, + encoding, + num_nulls, + num_rows, + def_levels_byte_len, + rep_levels_byte_len, + is_compressed: compressor.is_some(), + statistics: from_thrift_page_stats( + physical_type, + page_stats_to_thrift(statistics.as_ref()), + ) + .unwrap(), + } + } + Page::DictionaryPage { + ref buf, + num_values, + encoding, + is_sorted, + } => { + let output_buf = compress_helper(compressor.as_mut(), buf); + + Page::DictionaryPage { + buf: Bytes::from(output_buf), + num_values, + encoding, + is_sorted, + } + } + }; + + let compressed_page = CompressedPage::new(compressed_page, uncompressed_len); + compressed_pages.push(compressed_page); + } + + let mut buffer: Vec = vec![]; + let mut result_pages: Vec = vec![]; + { + let mut writer = TrackedWrite::new(&mut buffer); + let mut page_writer = SerializedPageWriter::new(&mut writer); + + for page in compressed_pages { + page_writer.write_page(page).unwrap(); + } + page_writer.close().unwrap(); + } + { + let reader = bytes::Bytes::from(buffer); + + let t = types::Type::primitive_type_builder("t", physical_type) + .build() + .unwrap(); + + let desc = ColumnDescriptor::new(Arc::new(t), 0, 0, ColumnPath::new(vec![])); + let meta = ColumnChunkMetaData::builder(Arc::new(desc)) + .set_compression_codec(codec.into()) + .set_total_compressed_size(reader.len() as i64) + .set_num_values(total_num_values) + .build() + .unwrap(); + + let props = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .set_read_page_statistics(true) + .build(); + let mut page_reader = SerializedPageReader::new_with_properties( + Arc::new(reader), + &meta, + total_num_values as usize, + None, + Arc::new(props), + ) + .unwrap(); + + while let Some(page) = page_reader.get_next_page().unwrap() { + result_pages.push(page); + } + } + + assert_eq!(result_pages.len(), pages.len()); + for i in 0..result_pages.len() { + assert_page(&result_pages[i], &pages[i]); + } + } + + /// Helper function to compress a slice + fn compress_helper(compressor: Option<&mut Box>, data: &[u8]) -> Vec { + let mut output_buf = vec![]; + if let Some(cmpr) = compressor { + cmpr.compress(data, &mut output_buf).unwrap(); + } else { + output_buf.extend_from_slice(data); + } + output_buf + } + + /// Check if pages match. + fn assert_page(left: &Page, right: &Page) { + assert_eq!(left.page_type(), right.page_type()); + assert_eq!(&left.buffer(), &right.buffer()); + assert_eq!(left.num_values(), right.num_values()); + assert_eq!(left.encoding(), right.encoding()); + assert_eq!( + page_stats_to_thrift(left.statistics()), + page_stats_to_thrift(right.statistics()) + ); + } + + /// Tests roundtrip of i32 data written using `W` and read using `R` + fn test_roundtrip_i32( + file: W, + data: Vec>, + compression: Compression, + ) -> ParquetMetaData + where + W: Write + Send, + R: ChunkReader + From + 'static, + { + test_roundtrip::(file, data, |r| r.get_int(0).unwrap(), compression) + } + + /// Tests roundtrip of data of type `D` written using `W` and read using `R` + /// and the provided `values` function + fn test_roundtrip( + mut file: W, + data: Vec>, + value: F, + compression: Compression, + ) -> ParquetMetaData + where + W: Write + Send, + R: ChunkReader + From + 'static, + D: DataType, + F: Fn(Row) -> D::T, + { + let schema = Arc::new( + types::Type::group_type_builder("schema") + .with_fields(vec![Arc::new( + types::Type::primitive_type_builder("col1", D::get_physical_type()) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(), + )]) + .build() + .unwrap(), + ); + let props = Arc::new( + WriterProperties::builder() + .set_compression(compression) + .build(), + ); + let mut file_writer = SerializedFileWriter::new(&mut file, schema, props).unwrap(); + let mut rows: i64 = 0; + + for (idx, subset) in data.iter().enumerate() { + let row_group_file_offset = file_writer.buf.bytes_written(); + let mut row_group_writer = file_writer.next_row_group().unwrap(); + if let Some(mut writer) = row_group_writer.next_column().unwrap() { + rows += writer + .typed::() + .write_batch(&subset[..], None, None) + .unwrap() as i64; + writer.close().unwrap(); + } + let last_group = row_group_writer.close().unwrap(); + let flushed = file_writer.flushed_row_groups(); + assert_eq!(flushed.len(), idx + 1); + assert_eq!(Some(idx as i16), last_group.ordinal()); + assert_eq!(Some(row_group_file_offset as i64), last_group.file_offset()); + assert_eq!(&flushed[idx], last_group.as_ref()); + } + let file_metadata = file_writer.close().unwrap(); + + let reader = SerializedFileReader::new(R::from(file)).unwrap(); + assert_eq!(reader.num_row_groups(), data.len()); + assert_eq!( + reader.metadata().file_metadata().num_rows(), + rows, + "row count in metadata not equal to number of rows written" + ); + for (i, item) in data.iter().enumerate().take(reader.num_row_groups()) { + let row_group_reader = reader.get_row_group(i).unwrap(); + let iter = row_group_reader.get_row_iter(None).unwrap(); + let res: Vec<_> = iter.map(|row| row.unwrap()).map(&value).collect(); + let row_group_size = row_group_reader.metadata().total_byte_size(); + let uncompressed_size: i64 = row_group_reader + .metadata() + .columns() + .iter() + .map(|v| v.uncompressed_size()) + .sum(); + assert_eq!(row_group_size, uncompressed_size); + assert_eq!(res, *item); + } + file_metadata + } + + /// File write-read roundtrip. + /// `data` consists of arrays of values for each row group. + fn test_file_roundtrip(file: File, data: Vec>) -> ParquetMetaData { + test_roundtrip_i32::(file, data, Compression::UNCOMPRESSED) + } + + #[test] + fn test_bytes_writer_empty_row_groups() { + test_bytes_roundtrip(vec![], Compression::UNCOMPRESSED); + } + + #[test] + fn test_bytes_writer_single_row_group() { + test_bytes_roundtrip(vec![vec![1, 2, 3, 4, 5]], Compression::UNCOMPRESSED); + } + + #[test] + fn test_bytes_writer_multiple_row_groups() { + test_bytes_roundtrip( + vec![ + vec![1, 2, 3, 4, 5], + vec![1, 2, 3], + vec![1], + vec![1, 2, 3, 4, 5, 6], + ], + Compression::UNCOMPRESSED, + ); + } + + #[test] + fn test_bytes_writer_single_row_group_compressed() { + test_bytes_roundtrip(vec![vec![1, 2, 3, 4, 5]], Compression::SNAPPY); + } + + #[test] + fn test_bytes_writer_multiple_row_groups_compressed() { + test_bytes_roundtrip( + vec![ + vec![1, 2, 3, 4, 5], + vec![1, 2, 3], + vec![1], + vec![1, 2, 3, 4, 5, 6], + ], + Compression::SNAPPY, + ); + } + + fn test_bytes_roundtrip(data: Vec>, compression: Compression) { + test_roundtrip_i32::, Bytes>(Vec::with_capacity(1024), data, compression); + } + + #[test] + fn test_boolean_roundtrip() { + let my_bool_values: Vec<_> = (0..2049).map(|idx| idx % 2 == 0).collect(); + test_roundtrip::, Bytes, BoolType, _>( + Vec::with_capacity(1024), + vec![my_bool_values], + |r| r.get_bool(0).unwrap(), + Compression::UNCOMPRESSED, + ); + } + + #[test] + fn test_boolean_compressed_roundtrip() { + let my_bool_values: Vec<_> = (0..2049).map(|idx| idx % 2 == 0).collect(); + test_roundtrip::, Bytes, BoolType, _>( + Vec::with_capacity(1024), + vec![my_bool_values], + |r| r.get_bool(0).unwrap(), + Compression::SNAPPY, + ); + } + + #[test] + fn test_column_offset_index_file() { + let file = tempfile::tempfile().unwrap(); + let file_metadata = test_file_roundtrip(file, vec![vec![1, 2, 3, 4, 5]]); + file_metadata.row_groups().iter().for_each(|row_group| { + row_group.columns().iter().for_each(|column_chunk| { + assert!(column_chunk.column_index_offset().is_some()); + assert!(column_chunk.column_index_length().is_some()); + assert!(column_chunk.offset_index_offset().is_some()); + assert!(column_chunk.offset_index_length().is_some()); + }) + }); + } + + fn test_kv_metadata(initial_kv: Option>, final_kv: Option>) { + let schema = Arc::new( + types::Type::group_type_builder("schema") + .with_fields(vec![Arc::new( + types::Type::primitive_type_builder("col1", Type::INT32) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(), + )]) + .build() + .unwrap(), + ); + let mut out = Vec::with_capacity(1024); + let props = Arc::new( + WriterProperties::builder() + .set_key_value_metadata(initial_kv.clone()) + .build(), + ); + let mut writer = SerializedFileWriter::new(&mut out, schema, props).unwrap(); + let mut row_group_writer = writer.next_row_group().unwrap(); + let column = row_group_writer.next_column().unwrap().unwrap(); + column.close().unwrap(); + row_group_writer.close().unwrap(); + if let Some(kvs) = &final_kv { + for kv in kvs { + writer.append_key_value_metadata(kv.clone()) + } + } + writer.close().unwrap(); + + let reader = SerializedFileReader::new(Bytes::from(out)).unwrap(); + let metadata = reader.metadata().file_metadata(); + let keys = metadata.key_value_metadata(); + + match (initial_kv, final_kv) { + (Some(a), Some(b)) => { + let keys = keys.unwrap(); + assert_eq!(keys.len(), a.len() + b.len()); + assert_eq!(&keys[..a.len()], a.as_slice()); + assert_eq!(&keys[a.len()..], b.as_slice()); + } + (Some(v), None) => assert_eq!(keys.unwrap(), &v), + (None, Some(v)) if !v.is_empty() => assert_eq!(keys.unwrap(), &v), + _ => assert!(keys.is_none()), + } + } + + #[test] + fn test_append_metadata() { + let kv1 = KeyValue::new("cupcakes".to_string(), "awesome".to_string()); + let kv2 = KeyValue::new("bingo".to_string(), "bongo".to_string()); + + test_kv_metadata(None, None); + test_kv_metadata(Some(vec![kv1.clone()]), None); + test_kv_metadata(None, Some(vec![kv2.clone()])); + test_kv_metadata(Some(vec![kv1.clone()]), Some(vec![kv2.clone()])); + test_kv_metadata(Some(vec![]), Some(vec![kv2])); + test_kv_metadata(Some(vec![]), Some(vec![])); + test_kv_metadata(Some(vec![kv1]), Some(vec![])); + test_kv_metadata(None, Some(vec![])); + } + + #[test] + fn test_backwards_compatible_statistics() { + let message_type = " + message test_schema { + REQUIRED INT32 decimal1 (DECIMAL(8,2)); + REQUIRED INT32 i32 (INTEGER(32,true)); + REQUIRED INT32 u32 (INTEGER(32,false)); + } + "; + + let schema = Arc::new(parse_message_type(message_type).unwrap()); + let props = Default::default(); + let mut writer = SerializedFileWriter::new(vec![], schema, props).unwrap(); + let mut row_group_writer = writer.next_row_group().unwrap(); + + for _ in 0..3 { + let mut writer = row_group_writer.next_column().unwrap().unwrap(); + writer + .typed::() + .write_batch(&[1, 2, 3], None, None) + .unwrap(); + writer.close().unwrap(); + } + let metadata = row_group_writer.close().unwrap(); + writer.close().unwrap(); + + // decimal + let s = page_stats_to_thrift(metadata.column(0).statistics()).unwrap(); + assert_eq!(s.min.as_deref(), Some(1_i32.to_le_bytes().as_ref())); + assert_eq!(s.max.as_deref(), Some(3_i32.to_le_bytes().as_ref())); + assert_eq!(s.min_value.as_deref(), Some(1_i32.to_le_bytes().as_ref())); + assert_eq!(s.max_value.as_deref(), Some(3_i32.to_le_bytes().as_ref())); + + // i32 + let s = page_stats_to_thrift(metadata.column(1).statistics()).unwrap(); + assert_eq!(s.min.as_deref(), Some(1_i32.to_le_bytes().as_ref())); + assert_eq!(s.max.as_deref(), Some(3_i32.to_le_bytes().as_ref())); + assert_eq!(s.min_value.as_deref(), Some(1_i32.to_le_bytes().as_ref())); + assert_eq!(s.max_value.as_deref(), Some(3_i32.to_le_bytes().as_ref())); + + // u32 + let s = page_stats_to_thrift(metadata.column(2).statistics()).unwrap(); + assert_eq!(s.min.as_deref(), None); + assert_eq!(s.max.as_deref(), None); + assert_eq!(s.min_value.as_deref(), Some(1_i32.to_le_bytes().as_ref())); + assert_eq!(s.max_value.as_deref(), Some(3_i32.to_le_bytes().as_ref())); + } + + #[test] + fn test_spliced_write() { + let message_type = " + message test_schema { + REQUIRED INT32 i32 (INTEGER(32,true)); + REQUIRED INT32 u32 (INTEGER(32,false)); + } + "; + let schema = Arc::new(parse_message_type(message_type).unwrap()); + let props = Arc::new(WriterProperties::builder().build()); + + let mut file = Vec::with_capacity(1024); + let mut file_writer = SerializedFileWriter::new(&mut file, schema, props.clone()).unwrap(); + + let columns = file_writer.descr.columns(); + let mut column_state: Vec<(_, Option)> = columns + .iter() + .map(|_| (TrackedWrite::new(Vec::with_capacity(1024)), None)) + .collect(); + + let mut column_state_slice = column_state.as_mut_slice(); + let mut column_writers = Vec::with_capacity(columns.len()); + for c in columns { + let ((buf, out), tail) = column_state_slice.split_first_mut().unwrap(); + column_state_slice = tail; + + let page_writer = Box::new(SerializedPageWriter::new(buf)); + let col_writer = get_column_writer(c.clone(), props.clone(), page_writer); + column_writers.push(SerializedColumnWriter::new( + col_writer, + Some(Box::new(|on_close| { + *out = Some(on_close); + Ok(()) + })), + )); + } + + let column_data = [[1, 2, 3, 4], [7, 3, 7, 3]]; + + // Interleaved writing to the column writers + for (writer, batch) in column_writers.iter_mut().zip(column_data) { + let writer = writer.typed::(); + writer.write_batch(&batch, None, None).unwrap(); + } + + // Close the column writers + for writer in column_writers { + writer.close().unwrap() + } + + // Splice column data into a row group + let mut row_group_writer = file_writer.next_row_group().unwrap(); + for (write, close) in column_state { + let buf = Bytes::from(write.into_inner().unwrap()); + row_group_writer + .append_column(&buf, close.unwrap()) + .unwrap(); + } + row_group_writer.close().unwrap(); + file_writer.close().unwrap(); + + // Check data was written correctly + let file = Bytes::from(file); + let test_read = |reader: SerializedFileReader| { + let row_group = reader.get_row_group(0).unwrap(); + + let mut out = Vec::with_capacity(4); + let c1 = row_group.get_column_reader(0).unwrap(); + let mut c1 = get_typed_column_reader::(c1); + c1.read_records(4, None, None, &mut out).unwrap(); + assert_eq!(out, column_data[0]); + + out.clear(); + + let c2 = row_group.get_column_reader(1).unwrap(); + let mut c2 = get_typed_column_reader::(c2); + c2.read_records(4, None, None, &mut out).unwrap(); + assert_eq!(out, column_data[1]); + }; + + let reader = SerializedFileReader::new(file.clone()).unwrap(); + test_read(reader); + + let options = ReadOptionsBuilder::new().with_page_index().build(); + let reader = SerializedFileReader::new_with_options(file, options).unwrap(); + test_read(reader); + } + + #[test] + fn test_disabled_statistics() { + let message_type = " + message test_schema { + REQUIRED INT32 a; + REQUIRED INT32 b; + } + "; + let schema = Arc::new(parse_message_type(message_type).unwrap()); + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::None) + .set_column_statistics_enabled("a".into(), EnabledStatistics::Page) + .set_offset_index_disabled(true) // this should be ignored because of the line above + .build(); + let mut file = Vec::with_capacity(1024); + let mut file_writer = + SerializedFileWriter::new(&mut file, schema, Arc::new(props)).unwrap(); + + let mut row_group_writer = file_writer.next_row_group().unwrap(); + let mut a_writer = row_group_writer.next_column().unwrap().unwrap(); + let col_writer = a_writer.typed::(); + col_writer.write_batch(&[1, 2, 3], None, None).unwrap(); + a_writer.close().unwrap(); + + let mut b_writer = row_group_writer.next_column().unwrap().unwrap(); + let col_writer = b_writer.typed::(); + col_writer.write_batch(&[4, 5, 6], None, None).unwrap(); + b_writer.close().unwrap(); + row_group_writer.close().unwrap(); + + let metadata = file_writer.finish().unwrap(); + assert_eq!(metadata.num_row_groups(), 1); + let row_group = metadata.row_group(0); + assert_eq!(row_group.num_columns(), 2); + // Column "a" has both offset and column index, as requested + assert!(row_group.column(0).offset_index_offset().is_some()); + assert!(row_group.column(0).column_index_offset().is_some()); + // Column "b" should only have offset index + assert!(row_group.column(1).offset_index_offset().is_some()); + assert!(row_group.column(1).column_index_offset().is_none()); + + let err = file_writer.next_row_group().err().unwrap().to_string(); + assert_eq!(err, "Parquet error: SerializedFileWriter already finished"); + + drop(file_writer); + + let options = ReadOptionsBuilder::new().with_page_index().build(); + let reader = SerializedFileReader::new_with_options(Bytes::from(file), options).unwrap(); + + let offset_index = reader.metadata().offset_index().unwrap(); + assert_eq!(offset_index.len(), 1); // 1 row group + assert_eq!(offset_index[0].len(), 2); // 2 columns + + let column_index = reader.metadata().column_index().unwrap(); + assert_eq!(column_index.len(), 1); // 1 row group + assert_eq!(column_index[0].len(), 2); // 2 column + + let a_idx = &column_index[0][0]; + assert!(matches!(a_idx, ColumnIndexMetaData::INT32(_)), "{a_idx:?}"); + let b_idx = &column_index[0][1]; + assert!(matches!(b_idx, ColumnIndexMetaData::NONE), "{b_idx:?}"); + } + + #[test] + fn test_byte_array_size_statistics() { + let message_type = " + message test_schema { + OPTIONAL BYTE_ARRAY a (UTF8); + } + "; + let schema = Arc::new(parse_message_type(message_type).unwrap()); + let data = ByteArrayType::gen_vec(32, 7); + let def_levels = [1, 1, 1, 1, 0, 1, 0, 1, 0, 1]; + let unenc_size: i64 = data.iter().map(|x| x.len() as i64).sum(); + let file: File = tempfile::tempfile().unwrap(); + let props = Arc::new( + WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .build(), + ); + + let mut writer = SerializedFileWriter::new(&file, schema, props).unwrap(); + let mut row_group_writer = writer.next_row_group().unwrap(); + + let mut col_writer = row_group_writer.next_column().unwrap().unwrap(); + col_writer + .typed::() + .write_batch(&data, Some(&def_levels), None) + .unwrap(); + col_writer.close().unwrap(); + row_group_writer.close().unwrap(); + let file_metadata = writer.close().unwrap(); + + assert_eq!(file_metadata.num_row_groups(), 1); + assert_eq!(file_metadata.row_group(0).num_columns(), 1); + + let check_def_hist = |def_hist: &[i64]| { + assert_eq!(def_hist.len(), 2); + assert_eq!(def_hist[0], 3); + assert_eq!(def_hist[1], 7); + }; + + let meta_data = file_metadata.row_group(0).column(0); + + assert!(meta_data.repetition_level_histogram().is_none()); + assert!(meta_data.definition_level_histogram().is_some()); + assert!(meta_data.unencoded_byte_array_data_bytes().is_some()); + assert_eq!( + unenc_size, + meta_data.unencoded_byte_array_data_bytes().unwrap() + ); + check_def_hist(meta_data.definition_level_histogram().unwrap().values()); + + // check that the read metadata is also correct + let options = ReadOptionsBuilder::new().with_page_index().build(); + let reader = SerializedFileReader::new_with_options(file, options).unwrap(); + + let rfile_metadata = reader.metadata().file_metadata(); + assert_eq!( + rfile_metadata.num_rows(), + file_metadata.file_metadata().num_rows() + ); + assert_eq!(reader.num_row_groups(), 1); + let rowgroup = reader.get_row_group(0).unwrap(); + assert_eq!(rowgroup.num_columns(), 1); + let column = rowgroup.metadata().column(0); + assert!(column.definition_level_histogram().is_some()); + assert!(column.repetition_level_histogram().is_none()); + assert!(column.unencoded_byte_array_data_bytes().is_some()); + check_def_hist(column.definition_level_histogram().unwrap().values()); + assert_eq!( + unenc_size, + column.unencoded_byte_array_data_bytes().unwrap() + ); + + // check histogram in column index as well + assert!(reader.metadata().column_index().is_some()); + let column_index = reader.metadata().column_index().unwrap(); + assert_eq!(column_index.len(), 1); + assert_eq!(column_index[0].len(), 1); + let col_idx = if let ColumnIndexMetaData::BYTE_ARRAY(index) = &column_index[0][0] { + assert_eq!(index.num_pages(), 1); + index + } else { + unreachable!() + }; + + assert!(col_idx.repetition_level_histogram(0).is_none()); + assert!(col_idx.definition_level_histogram(0).is_some()); + check_def_hist(col_idx.definition_level_histogram(0).unwrap()); + + assert!(reader.metadata().offset_index().is_some()); + let offset_index = reader.metadata().offset_index().unwrap(); + assert_eq!(offset_index.len(), 1); + assert_eq!(offset_index[0].len(), 1); + assert!(offset_index[0][0].unencoded_byte_array_data_bytes.is_some()); + let page_sizes = offset_index[0][0] + .unencoded_byte_array_data_bytes + .as_ref() + .unwrap(); + assert_eq!(page_sizes.len(), 1); + assert_eq!(page_sizes[0], unenc_size); + } + + #[test] + fn test_too_many_rowgroups() { + let message_type = " + message test_schema { + REQUIRED BYTE_ARRAY a (UTF8); + } + "; + let schema = Arc::new(parse_message_type(message_type).unwrap()); + let file: File = tempfile::tempfile().unwrap(); + let props = Arc::new( + WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::None) + .set_max_row_group_row_count(Some(1)) + .build(), + ); + let mut writer = SerializedFileWriter::new(&file, schema, props).unwrap(); + + // Create 32k empty rowgroups. Should error when i == 32768. + for i in 0..0x8001 { + match writer.next_row_group() { + Ok(mut row_group_writer) => { + assert_ne!(i, 0x8000); + let col_writer = row_group_writer.next_column().unwrap().unwrap(); + col_writer.close().unwrap(); + row_group_writer.close().unwrap(); + } + Err(e) => { + assert_eq!(i, 0x8000); + assert_eq!( + e.to_string(), + "Parquet error: Parquet does not support more than 32767 row groups per file (currently: 32768)" + ); + } + } + } + writer.close().unwrap(); + } + + #[test] + fn test_size_statistics_with_repetition_and_nulls() { + let message_type = " + message test_schema { + OPTIONAL group i32_list (LIST) { + REPEATED group list { + OPTIONAL INT32 element; + } + } + } + "; + // column is: + // row 0: [1, 2] + // row 1: NULL + // row 2: [4, NULL] + // row 3: [] + // row 4: [7, 8, 9, 10] + let schema = Arc::new(parse_message_type(message_type).unwrap()); + let data = [1, 2, 4, 7, 8, 9, 10]; + let def_levels = [3, 3, 0, 3, 2, 1, 3, 3, 3, 3]; + let rep_levels = [0, 1, 0, 0, 1, 0, 0, 1, 1, 1]; + let file = tempfile::tempfile().unwrap(); + let props = Arc::new( + WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .build(), + ); + let mut writer = SerializedFileWriter::new(&file, schema, props).unwrap(); + let mut row_group_writer = writer.next_row_group().unwrap(); + + let mut col_writer = row_group_writer.next_column().unwrap().unwrap(); + col_writer + .typed::() + .write_batch(&data, Some(&def_levels), Some(&rep_levels)) + .unwrap(); + col_writer.close().unwrap(); + row_group_writer.close().unwrap(); + let file_metadata = writer.close().unwrap(); + + assert_eq!(file_metadata.num_row_groups(), 1); + assert_eq!(file_metadata.row_group(0).num_columns(), 1); + + let check_def_hist = |def_hist: &[i64]| { + assert_eq!(def_hist.len(), 4); + assert_eq!(def_hist[0], 1); + assert_eq!(def_hist[1], 1); + assert_eq!(def_hist[2], 1); + assert_eq!(def_hist[3], 7); + }; + + let check_rep_hist = |rep_hist: &[i64]| { + assert_eq!(rep_hist.len(), 2); + assert_eq!(rep_hist[0], 5); + assert_eq!(rep_hist[1], 5); + }; + + // check that histograms are set properly in the write and read metadata + // also check that unencoded_byte_array_data_bytes is not set + let meta_data = file_metadata.row_group(0).column(0); + assert!(meta_data.repetition_level_histogram().is_some()); + assert!(meta_data.definition_level_histogram().is_some()); + assert!(meta_data.unencoded_byte_array_data_bytes().is_none()); + check_def_hist(meta_data.definition_level_histogram().unwrap().values()); + check_rep_hist(meta_data.repetition_level_histogram().unwrap().values()); + + // check that the read metadata is also correct + let options = ReadOptionsBuilder::new().with_page_index().build(); + let reader = SerializedFileReader::new_with_options(file, options).unwrap(); + + let rfile_metadata = reader.metadata().file_metadata(); + assert_eq!( + rfile_metadata.num_rows(), + file_metadata.file_metadata().num_rows() + ); + assert_eq!(reader.num_row_groups(), 1); + let rowgroup = reader.get_row_group(0).unwrap(); + assert_eq!(rowgroup.num_columns(), 1); + let column = rowgroup.metadata().column(0); + assert!(column.definition_level_histogram().is_some()); + assert!(column.repetition_level_histogram().is_some()); + assert!(column.unencoded_byte_array_data_bytes().is_none()); + check_def_hist(column.definition_level_histogram().unwrap().values()); + check_rep_hist(column.repetition_level_histogram().unwrap().values()); + + // check histogram in column index as well + assert!(reader.metadata().column_index().is_some()); + let column_index = reader.metadata().column_index().unwrap(); + assert_eq!(column_index.len(), 1); + assert_eq!(column_index[0].len(), 1); + let col_idx = if let ColumnIndexMetaData::INT32(index) = &column_index[0][0] { + assert_eq!(index.num_pages(), 1); + index + } else { + unreachable!() + }; + + check_def_hist(col_idx.definition_level_histogram(0).unwrap()); + check_rep_hist(col_idx.repetition_level_histogram(0).unwrap()); + + assert!(reader.metadata().offset_index().is_some()); + let offset_index = reader.metadata().offset_index().unwrap(); + assert_eq!(offset_index.len(), 1); + assert_eq!(offset_index[0].len(), 1); + assert!(offset_index[0][0].unencoded_byte_array_data_bytes.is_none()); + } + + #[test] + #[cfg(feature = "arrow")] + fn test_byte_stream_split_extended_roundtrip() { + let path = format!( + "{}/byte_stream_split_extended.gzip.parquet", + arrow::util::test_util::parquet_test_data(), + ); + let file = File::open(path).unwrap(); + + // Read in test file and rewrite to tmp + let parquet_reader = ParquetRecordBatchReaderBuilder::try_new(file) + .expect("parquet open") + .build() + .expect("parquet open"); + + let file = tempfile::tempfile().unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_column_encoding( + ColumnPath::from("float16_byte_stream_split"), + Encoding::BYTE_STREAM_SPLIT, + ) + .set_column_encoding( + ColumnPath::from("float_byte_stream_split"), + Encoding::BYTE_STREAM_SPLIT, + ) + .set_column_encoding( + ColumnPath::from("double_byte_stream_split"), + Encoding::BYTE_STREAM_SPLIT, + ) + .set_column_encoding( + ColumnPath::from("int32_byte_stream_split"), + Encoding::BYTE_STREAM_SPLIT, + ) + .set_column_encoding( + ColumnPath::from("int64_byte_stream_split"), + Encoding::BYTE_STREAM_SPLIT, + ) + .set_column_encoding( + ColumnPath::from("flba5_byte_stream_split"), + Encoding::BYTE_STREAM_SPLIT, + ) + .set_column_encoding( + ColumnPath::from("decimal_byte_stream_split"), + Encoding::BYTE_STREAM_SPLIT, + ) + .build(); + + let mut parquet_writer = ArrowWriter::try_new( + file.try_clone().expect("cannot open file"), + parquet_reader.schema(), + Some(props), + ) + .expect("create arrow writer"); + + for maybe_batch in parquet_reader { + let batch = maybe_batch.expect("reading batch"); + parquet_writer.write(&batch).expect("writing data"); + } + + parquet_writer.close().expect("finalizing file"); + + let reader = SerializedFileReader::new(file).expect("Failed to create reader"); + let filemeta = reader.metadata(); + + // Make sure byte_stream_split encoding was used + let check_encoding = |x: usize, filemeta: &ParquetMetaData| { + assert!( + filemeta + .row_group(0) + .column(x) + .encodings() + .collect::>() + .contains(&Encoding::BYTE_STREAM_SPLIT) + ); + }; + + check_encoding(1, filemeta); + check_encoding(3, filemeta); + check_encoding(5, filemeta); + check_encoding(7, filemeta); + check_encoding(9, filemeta); + check_encoding(11, filemeta); + check_encoding(13, filemeta); + + // Read back tmpfile and make sure all values are correct + let mut iter = reader + .get_row_iter(None) + .expect("Failed to create row iterator"); + + let mut start = 0; + let end = reader.metadata().file_metadata().num_rows(); + + let check_row = |row: Result| { + assert!(row.is_ok()); + let r = row.unwrap(); + assert_eq!(r.get_float16(0).unwrap(), r.get_float16(1).unwrap()); + assert_eq!(r.get_float(2).unwrap(), r.get_float(3).unwrap()); + assert_eq!(r.get_double(4).unwrap(), r.get_double(5).unwrap()); + assert_eq!(r.get_int(6).unwrap(), r.get_int(7).unwrap()); + assert_eq!(r.get_long(8).unwrap(), r.get_long(9).unwrap()); + assert_eq!(r.get_bytes(10).unwrap(), r.get_bytes(11).unwrap()); + assert_eq!(r.get_decimal(12).unwrap(), r.get_decimal(13).unwrap()); + }; + + while start < end { + match iter.next() { + Some(row) => check_row(row), + None => break, + }; + start += 1; + } + } + + #[test] + fn test_rewrite_no_page_indexes() { + let file = get_test_file("alltypes_tiny_pages.parquet"); + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Optional) + .parse_and_finish(&file) + .unwrap(); + + let props = Arc::new(WriterProperties::builder().build()); + let schema = metadata.file_metadata().schema_descr().root_schema_ptr(); + let output = Vec::::new(); + let mut writer = SerializedFileWriter::new(output, schema, props).unwrap(); + + for rg in metadata.row_groups() { + let mut rg_out = writer.next_row_group().unwrap(); + for column in rg.columns() { + let result = ColumnCloseResult { + bytes_written: column.compressed_size() as _, + rows_written: rg.num_rows() as _, + metadata: column.clone(), + bloom_filter: None, + column_index: None, + offset_index: None, + }; + rg_out.append_column(&file, result).unwrap(); + } + rg_out.close().unwrap(); + } + writer.close().unwrap(); + } + + #[test] + fn test_rewrite_missing_column_index() { + // this file has an INT96 column that lacks a column index entry + let file = get_test_file("alltypes_tiny_pages.parquet"); + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Optional) + .parse_and_finish(&file) + .unwrap(); + + let props = Arc::new(WriterProperties::builder().build()); + let schema = metadata.file_metadata().schema_descr().root_schema_ptr(); + let output = Vec::::new(); + let mut writer = SerializedFileWriter::new(output, schema, props).unwrap(); + + let column_indexes = metadata.column_index(); + let offset_indexes = metadata.offset_index(); + + for (rg_idx, rg) in metadata.row_groups().iter().enumerate() { + let rg_column_indexes = column_indexes.and_then(|ci| ci.get(rg_idx)); + let rg_offset_indexes = offset_indexes.and_then(|oi| oi.get(rg_idx)); + let mut rg_out = writer.next_row_group().unwrap(); + for (col_idx, column) in rg.columns().iter().enumerate() { + let column_index = rg_column_indexes.and_then(|row| row.get(col_idx)).cloned(); + let offset_index = rg_offset_indexes.and_then(|row| row.get(col_idx)).cloned(); + let result = ColumnCloseResult { + bytes_written: column.compressed_size() as _, + rows_written: rg.num_rows() as _, + metadata: column.clone(), + bloom_filter: None, + column_index, + offset_index, + }; + rg_out.append_column(&file, result).unwrap(); + } + rg_out.close().unwrap(); + } + writer.close().unwrap(); + } +} diff --git a/vendor/parquet/src/geospatial/accumulator.rs b/vendor/parquet/src/geospatial/accumulator.rs new file mode 100644 index 0000000..3aad106 --- /dev/null +++ b/vendor/parquet/src/geospatial/accumulator.rs @@ -0,0 +1,384 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! This module provides implementations and traits for building [`GeospatialStatistics`] + +use std::sync::{Arc, OnceLock}; + +use crate::{ + basic::LogicalType, errors::ParquetError, geospatial::statistics::GeospatialStatistics, + schema::types::ColumnDescPtr, +}; + +/// Create a new [`GeoStatsAccumulator`] instance if `descr` represents a Geometry or +/// Geography [`LogicalType`] +/// +/// Returns a suitable [`GeoStatsAccumulator`] if `descr` represents a non-geospatial type +/// or `None` otherwise. +pub fn try_new_geo_stats_accumulator( + descr: &ColumnDescPtr, +) -> Option> { + if !matches!( + descr.logical_type_ref(), + Some(LogicalType::Geometry { .. }) | Some(LogicalType::Geography { .. }) + ) { + return None; + } + + Some( + ACCUMULATOR_FACTORY + .get_or_init(|| Arc::new(DefaultGeoStatsAccumulatorFactory::default())) + .new_accumulator(descr), + ) +} + +/// Initialize the global [`GeoStatsAccumulatorFactory`] +/// +/// This may only be done once before any calls to [`try_new_geo_stats_accumulator`]. +/// Clients may use this to implement support for builds of the Parquet crate without +/// geospatial support or to implement support for Geography bounding using external +/// dependencies. +pub fn init_geo_stats_accumulator_factory( + factory: Arc, +) -> Result<(), ParquetError> { + if ACCUMULATOR_FACTORY.set(factory).is_err() { + Err(ParquetError::General( + "Global GeoStatsAccumulatorFactory already set".to_string(), + )) + } else { + Ok(()) + } +} + +/// Global accumulator factory instance +static ACCUMULATOR_FACTORY: OnceLock> = OnceLock::new(); + +/// Factory for [`GeospatialStatistics`] accumulators +/// +/// The GeoStatsAccumulatorFactory is a trait implemented by the global factory that +/// generates new instances of a [`GeoStatsAccumulator`] when constructing new +/// encoders for a Geometry or Geography logical type. +pub trait GeoStatsAccumulatorFactory: Send + Sync { + /// Create a new [`GeoStatsAccumulator`] appropriate for the logical type of a given + /// [`ColumnDescPtr`] + fn new_accumulator(&self, descr: &ColumnDescPtr) -> Box; +} + +/// Dynamic [`GeospatialStatistics`] accumulator +/// +/// The GeoStatsAccumulator is a trait whose implementors can ingest the (non-null) +/// elements of a column and return compliant [`GeospatialStatistics`] (or `None`). +/// When built with geospatial support this will usually be the +/// [`ParquetGeoStatsAccumulator`] +pub trait GeoStatsAccumulator: Send { + /// Returns true if this instance can return [`GeospatialStatistics`] from + /// [`GeoStatsAccumulator::finish`]. + /// + /// This method returns false when this crate is built without geospatial support + /// (i.e., from the [`VoidGeoStatsAccumulator`]) or if the accumulator encountered + /// invalid or unsupported elements for which it cannot compute valid statistics. + fn is_valid(&self) -> bool; + + /// Update with a single slice of WKB-encoded values + /// + /// This method is infallible; however, in the event of improperly encoded values, + /// implementations must ensure that [`GeoStatsAccumulator::finish`] returns `None`. + fn update_wkb(&mut self, wkb: &[u8]); + + /// Compute the final statistics and reset internal state + fn finish(&mut self) -> Option>; +} + +/// Default accumulator for [`GeospatialStatistics`] +/// +/// When this crate is built with geospatial support, this factory constructs a +/// [`ParquetGeoStatsAccumulator`] that ensures Geometry columns are written with +/// statistics when statistics for that column are enabled. Otherwise, this factory +/// returns a [`VoidGeoStatsAccumulator`] that never adds any geospatial statistics. +/// +/// Bounding for Geography columns is not currently implemented by parquet-geospatial +/// and this factory will always return a [`VoidGeoStatsAccumulator`]. +#[derive(Debug, Default)] +pub struct DefaultGeoStatsAccumulatorFactory {} + +impl GeoStatsAccumulatorFactory for DefaultGeoStatsAccumulatorFactory { + fn new_accumulator(&self, _descr: &ColumnDescPtr) -> Box { + #[cfg(feature = "geospatial")] + if let Some(crate::basic::LogicalType::Geometry { .. }) = _descr.logical_type_ref() { + Box::new(ParquetGeoStatsAccumulator::default()) + } else { + Box::new(VoidGeoStatsAccumulator::default()) + } + + #[cfg(not(feature = "geospatial"))] + return Box::new(VoidGeoStatsAccumulator::default()); + } +} + +/// A [`GeoStatsAccumulator`] that never computes any [`GeospatialStatistics`] +#[derive(Debug, Default)] +pub struct VoidGeoStatsAccumulator {} + +impl GeoStatsAccumulator for VoidGeoStatsAccumulator { + fn is_valid(&self) -> bool { + false + } + + fn update_wkb(&mut self, _wkb: &[u8]) {} + + fn finish(&mut self) -> Option> { + None + } +} + +/// A [`GeoStatsAccumulator`] that uses the parquet-geospatial crate to compute Geometry statistics +/// +/// Note that this accumulator only supports Geometry types and will return invalid statistics for +/// non-point Geography input ([`GeoStatsAccumulatorFactory::new_accumulator`] is responsible +/// for ensuring an appropriate accumulator based on the logical type). +#[cfg(feature = "geospatial")] +#[derive(Debug)] +pub struct ParquetGeoStatsAccumulator { + bounder: parquet_geospatial::bounding::GeometryBounder, + invalid: bool, +} + +#[cfg(feature = "geospatial")] +impl Default for ParquetGeoStatsAccumulator { + fn default() -> Self { + Self { + bounder: parquet_geospatial::bounding::GeometryBounder::empty(), + invalid: false, + } + } +} + +#[cfg(feature = "geospatial")] +impl GeoStatsAccumulator for ParquetGeoStatsAccumulator { + fn is_valid(&self) -> bool { + !self.invalid + } + + fn update_wkb(&mut self, wkb: &[u8]) { + if self.bounder.update_wkb(wkb).is_err() { + self.invalid = true; + } + } + + fn finish(&mut self) -> Option> { + use parquet_geospatial::interval::IntervalTrait; + + use crate::geospatial::bounding_box::BoundingBox; + + if self.invalid { + // Reset + self.invalid = false; + self.bounder = parquet_geospatial::bounding::GeometryBounder::empty(); + return None; + } + + let bbox = if self.bounder.x().is_empty() || self.bounder.y().is_empty() { + None + } else { + let mut bbox = BoundingBox::new( + self.bounder.x().lo(), + self.bounder.x().hi(), + self.bounder.y().lo(), + self.bounder.y().hi(), + ); + + if !self.bounder.z().is_empty() { + bbox = bbox.with_zrange(self.bounder.z().lo(), self.bounder.z().hi()); + } + + if !self.bounder.m().is_empty() { + bbox = bbox.with_mrange(self.bounder.m().lo(), self.bounder.m().hi()); + } + + Some(bbox) + }; + + let bounder_geometry_types = self.bounder.geometry_types(); + let geometry_types = if bounder_geometry_types.is_empty() { + None + } else { + Some(bounder_geometry_types) + }; + + // Reset + self.bounder = parquet_geospatial::bounding::GeometryBounder::empty(); + + Some(Box::new(GeospatialStatistics::new(bbox, geometry_types))) + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_void_accumulator() { + let mut accumulator = VoidGeoStatsAccumulator {}; + assert!(!accumulator.is_valid()); + accumulator.update_wkb(&[0x01, 0x02, 0x03]); + assert!(accumulator.finish().is_none()); + } + + #[cfg(feature = "geospatial")] + #[test] + fn test_default_accumulator_geospatial_factory() { + use std::sync::Arc; + + use parquet_geospatial::testing::wkb_point_xy; + + use crate::{ + basic::LogicalType, + geospatial::bounding_box::BoundingBox, + schema::types::{ColumnDescriptor, ColumnPath, Type}, + }; + + // Check that we have a working accumulator for Geometry + let parquet_type = Type::primitive_type_builder("geom", crate::basic::Type::BYTE_ARRAY) + .with_logical_type(Some(LogicalType::geometry(None))) + .build() + .unwrap(); + let column_descr = + ColumnDescriptor::new(Arc::new(parquet_type), 0, 0, ColumnPath::new(vec![])); + let mut accumulator = try_new_geo_stats_accumulator(&Arc::new(column_descr)).unwrap(); + + assert!(accumulator.is_valid()); + accumulator.update_wkb(&wkb_point_xy(1.0, 2.0)); + accumulator.update_wkb(&wkb_point_xy(11.0, 12.0)); + let stats = accumulator.finish().unwrap(); + assert_eq!( + stats.bounding_box().unwrap(), + &BoundingBox::new(1.0, 11.0, 2.0, 12.0) + ); + + // Check that we have a void accumulator for Geography + let parquet_type = Type::primitive_type_builder("geom", crate::basic::Type::BYTE_ARRAY) + .with_logical_type(Some(LogicalType::geography(None, None))) + .build() + .unwrap(); + let column_descr = + ColumnDescriptor::new(Arc::new(parquet_type), 0, 0, ColumnPath::new(vec![])); + let mut accumulator = try_new_geo_stats_accumulator(&Arc::new(column_descr)).unwrap(); + + assert!(!accumulator.is_valid()); + assert!(accumulator.finish().is_none()); + + // Check that we return None if the type is not geometry or goegraphy + let parquet_type = Type::primitive_type_builder("geom", crate::basic::Type::BYTE_ARRAY) + .build() + .unwrap(); + let column_descr = + ColumnDescriptor::new(Arc::new(parquet_type), 0, 0, ColumnPath::new(vec![])); + assert!(try_new_geo_stats_accumulator(&Arc::new(column_descr)).is_none()); + + // We should not be able to initialize a global accumulator after we've initialized at least + // one accumulator + assert!( + init_geo_stats_accumulator_factory(Arc::new( + DefaultGeoStatsAccumulatorFactory::default() + )) + .is_err() + ) + } + + #[cfg(feature = "geospatial")] + #[test] + fn test_geometry_accumulator() { + use parquet_geospatial::testing::{wkb_point_xy, wkb_point_xyzm}; + + use crate::geospatial::bounding_box::BoundingBox; + + let mut accumulator = ParquetGeoStatsAccumulator::default(); + + // A fresh instance should be able to bound input + assert!(accumulator.is_valid()); + accumulator.update_wkb(&wkb_point_xy(1.0, 2.0)); + accumulator.update_wkb(&wkb_point_xy(11.0, 12.0)); + let stats = accumulator.finish().unwrap(); + assert_eq!(stats.geospatial_types().unwrap(), &vec![1]); + assert_eq!( + stats.bounding_box().unwrap(), + &BoundingBox::new(1.0, 11.0, 2.0, 12.0) + ); + + // finish() should have reset the bounder such that the first values + // aren't when computing the next bound of statistics. + assert!(accumulator.is_valid()); + accumulator.update_wkb(&wkb_point_xy(21.0, 22.0)); + accumulator.update_wkb(&wkb_point_xy(31.0, 32.0)); + let stats = accumulator.finish().unwrap(); + assert_eq!(stats.geospatial_types().unwrap(), &vec![1]); + assert_eq!( + stats.bounding_box().unwrap(), + &BoundingBox::new(21.0, 31.0, 22.0, 32.0) + ); + + // When an accumulator encounters invalid input, it reports is_valid() false + // and does not compute subsequent statistics + assert!(accumulator.is_valid()); + accumulator.update_wkb(&wkb_point_xy(41.0, 42.0)); + accumulator.update_wkb("these bytes are not WKB".as_bytes()); + assert!(!accumulator.is_valid()); + assert!(accumulator.finish().is_none()); + + // Subsequent rounds of accumulation should work as expected + assert!(accumulator.is_valid()); + accumulator.update_wkb(&wkb_point_xy(41.0, 42.0)); + accumulator.update_wkb(&wkb_point_xy(51.0, 52.0)); + let stats = accumulator.finish().unwrap(); + assert_eq!(stats.geospatial_types().unwrap(), &vec![1]); + assert_eq!( + stats.bounding_box().unwrap(), + &BoundingBox::new(41.0, 51.0, 42.0, 52.0) + ); + + // When there was no input at all (occurs in the all null case), both geometry + // types and bounding box will be None. This is because Parquet Thrift statistics + // have no mechanism to communicate "empty". (The all null situation may be determined + // from the null count in this case). + assert!(accumulator.is_valid()); + let stats = accumulator.finish().unwrap(); + assert!(stats.geospatial_types().is_none()); + assert!(stats.bounding_box().is_none()); + + // When there was 100% "empty" input (i.e., non-null geometries without + // coordinates), there should be statistics with geometry types but no + // bounding box. + assert!(accumulator.is_valid()); + accumulator.update_wkb(&wkb_point_xy(f64::NAN, f64::NAN)); + let stats = accumulator.finish().unwrap(); + assert_eq!(stats.geospatial_types().unwrap(), &vec![1]); + assert!(stats.bounding_box().is_none()); + + // If Z and/or M are present, they should be reported in the bounding box + assert!(accumulator.is_valid()); + accumulator.update_wkb(&wkb_point_xyzm(1.0, 2.0, 3.0, 4.0)); + accumulator.update_wkb(&wkb_point_xyzm(5.0, 6.0, 7.0, 8.0)); + let stats = accumulator.finish().unwrap(); + assert_eq!(stats.geospatial_types().unwrap(), &vec![3001]); + assert_eq!( + stats.bounding_box().unwrap(), + &BoundingBox::new(1.0, 5.0, 2.0, 6.0) + .with_zrange(3.0, 7.0) + .with_mrange(4.0, 8.0) + ); + } +} diff --git a/vendor/parquet/src/geospatial/bounding_box.rs b/vendor/parquet/src/geospatial/bounding_box.rs new file mode 100644 index 0000000..59a4dfb --- /dev/null +++ b/vendor/parquet/src/geospatial/bounding_box.rs @@ -0,0 +1,224 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Bounding box for GEOMETRY or GEOGRAPHY type in the representation of min/max +//! value pair of coordinates from each axis. +//! +//! Derived from the parquet format spec: +//! +//! + +use crate::file::metadata::HeapSize; + +/// A geospatial instance has at least two coordinate dimensions: X and Y for 2D coordinates of each point. +/// X represents longitude/easting and Y represents latitude/northing. A geospatial instance can optionally +/// have Z and/or M values associated with each point. +/// +/// The Z values introduce the third dimension coordinate, typically used to indicate height or elevation. +/// +/// M values allow tracking a value in a fourth dimension. These can represent: +/// - Linear reference values (e.g., highway milepost) +/// - Timestamps +/// - Other values defined by the CRS +/// +/// The bounding box is defined as min/max value pairs of coordinates from each axis. X and Y values are +/// always present, while Z and M are omitted for 2D geospatial instances. +/// +/// When calculating a bounding box: +/// - Null or NaN values in a coordinate dimension are skipped +/// - If a dimension has only null/NaN values, that dimension is omitted +/// - If either X or Y dimension is missing, no bounding box is produced +/// - Example: POINT (1 NaN) contributes to X but not to Y, Z, or M dimensions +/// +/// Special cases: +/// - For X values only, xmin may exceed xmax. In this case, a point matches if x >= xmin OR x <= xmax +/// - This wraparound can occur when the bounding box crosses the antimeridian line. +/// - In geographic terms: xmin=westernmost, xmax=easternmost, ymin=southernmost, ymax=northernmost +/// +/// For GEOGRAPHY types: +/// - X values must be within [-180, 180] (longitude) +/// - Y values must be within [-90, 90] (latitude) +/// +/// Derived from the parquet format [spec][bounding-box-spec] +/// +/// # Examples +/// +/// ``` +/// use parquet::geospatial::bounding_box::BoundingBox; +/// +/// // 2D bounding box +/// let bbox_2d = BoundingBox::new(0.0, 0.0, 100.0, 100.0); +/// +/// // 3D bounding box with elevation +/// let bbox_3d = BoundingBox::new(0.0, 0.0, 100.0, 100.0) +/// .with_zrange(0.0, 1000.0); +/// +/// // 3D bounding box with elevation and measured value +/// let bbox_3d_m = BoundingBox::new(0.0, 0.0, 100.0, 100.0) +/// .with_zrange(0.0, 1000.0) +/// .with_mrange(0.0, 1000.0); +/// ``` +/// +/// [bounding-box-spec]: https://github.com/apache/parquet-format/blob/master/Geospatial.md#bounding-box +#[derive(Clone, Debug, PartialEq)] +pub struct BoundingBox { + /// X coordinates (longitude or easting): (min, max) + x_range: (f64, f64), + /// Y coordinates (latitude or northing): (min, max) + y_range: (f64, f64), + /// Z coordinates (elevation/height): (min, max), if present + z_range: Option<(f64, f64)>, + /// M coordinates (measured value): (min, max), if present + m_range: Option<(f64, f64)>, +} + +impl BoundingBox { + /// Creates a new bounding box with the specified coordinates. + pub fn new(xmin: f64, xmax: f64, ymin: f64, ymax: f64) -> Self { + Self { + x_range: (xmin, xmax), + y_range: (ymin, ymax), + z_range: None, + m_range: None, + } + } + + /// Updates the bounding box with specified X-coordinate range. + pub fn with_xrange(mut self, xmin: f64, xmax: f64) -> Self { + self.x_range = (xmin, xmax); + self + } + + /// Updates the bounding box with specified Y-coordinate range. + pub fn with_yrange(mut self, ymin: f64, ymax: f64) -> Self { + self.y_range = (ymin, ymax); + self + } + + /// Creates a new bounding box with the specified Z-coordinate range. + pub fn with_zrange(mut self, zmin: f64, zmax: f64) -> Self { + self.z_range = Some((zmin, zmax)); + self + } + + /// Creates a new bounding box with the specified M-coordinate range. + pub fn with_mrange(mut self, mmin: f64, mmax: f64) -> Self { + self.m_range = Some((mmin, mmax)); + self + } + + /// Returns the minimum x-coordinate. + pub fn get_xmin(&self) -> f64 { + self.x_range.0 + } + + /// Returns the maximum x-coordinate. + pub fn get_xmax(&self) -> f64 { + self.x_range.1 + } + + /// Returns the minimum y-coordinate. + pub fn get_ymin(&self) -> f64 { + self.y_range.0 + } + + /// Returns the maximum y-coordinate. + pub fn get_ymax(&self) -> f64 { + self.y_range.1 + } + + /// Returns the minimum z-coordinate, if present. + pub fn get_zmin(&self) -> Option { + self.z_range.map(|z| z.0) + } + + /// Returns the maximum z-coordinate, if present. + pub fn get_zmax(&self) -> Option { + self.z_range.map(|z| z.1) + } + + /// Returns the minimum m-value (measure), if present. + pub fn get_mmin(&self) -> Option { + self.m_range.map(|m| m.0) + } + + /// Returns the maximum m-value (measure), if present. + pub fn get_mmax(&self) -> Option { + self.m_range.map(|m| m.1) + } + + /// Returns `true` if both zmin and zmax are present. + pub fn is_z_valid(&self) -> bool { + self.z_range.is_some() + } + + /// Returns `true` if both mmin and mmax are present. + pub fn is_m_valid(&self) -> bool { + self.m_range.is_some() + } +} + +impl HeapSize for BoundingBox { + fn heap_size(&self) -> usize { + 0 // no heap allocations + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_bounding_box() { + let bbox = BoundingBox::new(0.0, 0.0, 10.0, 10.0); + assert_eq!(bbox.get_xmin(), 0.0); + assert_eq!(bbox.get_xmax(), 0.0); + assert_eq!(bbox.get_ymin(), 10.0); + assert_eq!(bbox.get_ymax(), 10.0); + assert_eq!(bbox.get_zmin(), None); + assert_eq!(bbox.get_zmax(), None); + assert_eq!(bbox.get_mmin(), None); + assert_eq!(bbox.get_mmax(), None); + assert!(!bbox.is_z_valid()); + assert!(!bbox.is_m_valid()); + + // test with zrange + let bbox_z = BoundingBox::new(0.0, 0.0, 10.0, 10.0).with_zrange(5.0, 15.0); + assert_eq!(bbox_z.get_zmin(), Some(5.0)); + assert_eq!(bbox_z.get_zmax(), Some(15.0)); + assert!(bbox_z.is_z_valid()); + assert!(!bbox_z.is_m_valid()); + + // test with mrange + let bbox_m = BoundingBox::new(0.0, 0.0, 10.0, 10.0).with_mrange(10.0, 20.0); + assert_eq!(bbox_m.get_mmin(), Some(10.0)); + assert_eq!(bbox_m.get_mmax(), Some(20.0)); + assert!(!bbox_m.is_z_valid()); + assert!(bbox_m.is_m_valid()); + + // test with zrange and mrange + let bbox_zm = BoundingBox::new(0.0, 0.0, 10.0, 10.0) + .with_zrange(5.0, 15.0) + .with_mrange(10.0, 20.0); + assert_eq!(bbox_zm.get_zmin(), Some(5.0)); + assert_eq!(bbox_zm.get_zmax(), Some(15.0)); + assert_eq!(bbox_zm.get_mmin(), Some(10.0)); + assert_eq!(bbox_zm.get_mmax(), Some(20.0)); + assert!(bbox_zm.is_z_valid()); + assert!(bbox_zm.is_m_valid()); + } +} diff --git a/vendor/parquet/src/geospatial/mod.rs b/vendor/parquet/src/geospatial/mod.rs new file mode 100644 index 0000000..9d55fca --- /dev/null +++ b/vendor/parquet/src/geospatial/mod.rs @@ -0,0 +1,51 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! This module provides functionality for working with geospatial data in Parquet file as defined in the [spec][parquet-geo-spec]. +//! +//! * [`GeospatialStatistics`]: describes the geospatial statistics for a Parquet column. +//! * [`BoundingBox`]: describes the bounding box values for a geospatial column. +//! +//! [`GeospatialStatistics`] describes the geospatial statistics for a Parquet column. +//! * bbox: the [`BoundingBox`] for the geospatial data +//! * geospatial_types: the geospatial types for the geospatial data as specified in [specification][geo-types]. +//! +//! Geospatial bounding box describes the spatial extent of the geospatial data within a Parquet row group. +//! * xmin, xmax: the minimum and maximum longitude values +//! * ymin, ymax: the minimum and maximum latitude values +//! * zmin, zmax: (optional) the minimum and maximum elevation values +//! * mmin, mmax: (optional) the minimum and maximum linear reference values +//! +//! In 2D representation, where x are points: +//! ```text +//! ymax +-----------------------+ +//! | x | +//! | x | +//! | x | +//! | x | +//! ymin +-----------------------+ +//! xmin xmax +//! ``` +//! +//! [`GeospatialStatistics`]: crate::geospatial::statistics::GeospatialStatistics +//! [`BoundingBox`]: crate::geospatial::bounding_box::BoundingBox +//! [parquet-geo-spec]: https://github.com/apache/parquet-format/blob/master/Geospatial.md +//! [geo-types]: https://github.com/apache/parquet-format/blob/master/Geospatial.md#geospatial-types + +pub mod accumulator; +pub mod bounding_box; +pub mod statistics; diff --git a/vendor/parquet/src/geospatial/statistics.rs b/vendor/parquet/src/geospatial/statistics.rs new file mode 100644 index 0000000..edabfea --- /dev/null +++ b/vendor/parquet/src/geospatial/statistics.rs @@ -0,0 +1,74 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Geospatial statistics for Parquet files. +//! +//! This module provides functionality for working with geospatial statistics in Parquet files. +//! It includes support for bounding boxes and geospatial statistics in column chunk metadata. + +use crate::{file::metadata::HeapSize, geospatial::bounding_box::BoundingBox}; + +// ---------------------------------------------------------------------- +// Geospatial Statistics + +/// Represents geospatial statistics for a Parquet column or dataset. +/// +/// This struct contains metadata about the spatial characteristics of geospatial data, +/// including bounding box information and the types of geospatial geometries present. +/// It's used to optimize spatial queries and provide spatial context for data analysis. +/// +/// # Examples +/// +/// ``` +/// use parquet::geospatial::statistics::GeospatialStatistics; +/// use parquet::geospatial::bounding_box::BoundingBox; +/// +/// // Statistics with bounding box +/// let bbox = BoundingBox::new(0.0, 0.0, 100.0, 100.0); +/// let stats = GeospatialStatistics::new(Some(bbox), Some(vec![1, 2, 3])); +/// ``` +#[derive(Clone, Debug, PartialEq, Default)] +pub struct GeospatialStatistics { + bbox: Option, + geospatial_types: Option>, +} + +impl GeospatialStatistics { + /// Creates a new geospatial statistics instance with the specified data. + pub fn new(bbox: Option, geospatial_types: Option>) -> Self { + Self { + bbox, + geospatial_types, + } + } + + /// Optional list of geometry type identifiers, where `None` represents lack of information + pub fn geospatial_types(&self) -> Option<&Vec> { + self.geospatial_types.as_ref() + } + + /// Optional bounding defining the spatial extent, where `None` represents a lack of information. + pub fn bounding_box(&self) -> Option<&BoundingBox> { + self.bbox.as_ref() + } +} + +impl HeapSize for GeospatialStatistics { + fn heap_size(&self) -> usize { + self.bbox.heap_size() + self.geospatial_types.heap_size() + } +} diff --git a/vendor/parquet/src/lib.rs b/vendor/parquet/src/lib.rs new file mode 100644 index 0000000..3acdb61 --- /dev/null +++ b/vendor/parquet/src/lib.rs @@ -0,0 +1,230 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! +//! This crate contains the official Native Rust implementation of +//! [Apache Parquet](https://parquet.apache.org/), part of +//! the [Apache Arrow](https://arrow.apache.org/) project. +//! The crate provides a number of APIs to read and write Parquet files, +//! covering a range of use cases. +//! +//! Please see the [parquet crates.io](https://crates.io/crates/parquet) +//! page for feature flags and tips to improve performance. +//! +//! # Format Overview +//! +//! Parquet is a columnar format, which means that unlike row formats like [CSV], values are +//! iterated along columns instead of rows. Parquet is similar in spirit to [Arrow], but +//! focuses on storage efficiency whereas Arrow prioritizes compute efficiency. +//! +//! Parquet files are partitioned for scalability. Each file contains metadata, +//! along with zero or more "row groups", each row group containing one or +//! more columns. The APIs in this crate reflect this structure. +//! +//! Data in Parquet files is strongly typed and differentiates between logical +//! and physical types (see [`schema`]). In addition, Parquet files may contain +//! other metadata, such as statistics, which can be used to optimize reading +//! (see [`file::metadata`]). +//! For more details about the Parquet format itself, see the [Parquet spec] +//! +//! [Parquet spec]: https://github.com/apache/parquet-format/blob/master/README.md#file-format +//! +//! # APIs +//! +//! This crate exposes a number of APIs for different use-cases. +//! +//! ## Metadata and Schema +//! +//! The [`schema`] module provides APIs to work with Parquet schemas. The +//! [`file::metadata`] module provides APIs to work with Parquet metadata. +//! +//! ## Reading and Writing Arrow (`arrow` feature) +//! +//! The [`arrow`] module supports reading and writing Parquet data to/from +//! Arrow [`RecordBatch`]es. Using Arrow is simple and performant, and allows workloads +//! to leverage the wide range of data transforms provided by the [arrow] crate, and by the +//! ecosystem of [Arrow] compatible systems. +//! +//! Most users will use [`ArrowWriter`] for writing and [`ParquetRecordBatchReaderBuilder`] for +//! reading from synchronous IO sources such as files or in-memory buffers. +//! +//! Lower level APIs include +//! * [`ParquetPushDecoder`] for file grained control over interleaving of IO and CPU. +//! * [`ArrowColumnWriter`] for writing using multiple threads, +//! * [`RowFilter`] to apply filters during decode +//! +//! ### EXPERIMENTAL: Content-Defined Chunking +//! +//! [`ArrowWriter`] supports content-defined chunking (CDC), which creates data page +//! boundaries based on content rather than fixed sizes. CDC enables efficient +//! deduplication in content-addressable storage (CAS) systems: when the same data +//! appears in successive file versions, it will produce identical byte sequences that +//! CAS backends can deduplicate. +//! +//! Enable CDC via [`WriterProperties`]: +//! +//! ```rust +//! # use parquet::file::properties::{WriterProperties, CdcOptions}; +//! let props = WriterProperties::builder() +//! .set_content_defined_chunking(Some(CdcOptions::default())) +//! .build(); +//! ``` +//! +//! See [`CdcOptions`] for chunk size and normalization parameters. +//! +//! [`WriterProperties`]: file::properties::WriterProperties +//! [`CdcOptions`]: file::properties::CdcOptions +//! +//! [`ArrowWriter`]: arrow::arrow_writer::ArrowWriter +//! [`ParquetRecordBatchReaderBuilder`]: arrow::arrow_reader::ParquetRecordBatchReaderBuilder +//! [`ParquetPushDecoder`]: arrow::push_decoder::ParquetPushDecoder +//! [`ArrowColumnWriter`]: arrow::arrow_writer::ArrowColumnWriter +//! [`RowFilter`]: arrow::arrow_reader::RowFilter +//! +//! ## `async` Reading and Writing Arrow (`arrow` feature + `async` feature) +//! +//! The [`async_reader`] and [`async_writer`] modules provide async APIs to +//! read and write [`RecordBatch`]es asynchronously. +//! +//! Most users will use [`AsyncArrowWriter`] for writing and [`ParquetRecordBatchStreamBuilder`] +//! for reading, automatically optimizing IO based on any predicates or projections provided. +//! Object storage services such as S3 can be integrated by implementing +//! [`AsyncFileReader`] on top of a client such as the [object_store] crate, +//! or by passing a writer implementing [`AsyncWrite`] (such as +//! `object_store::buffered::BufWriter`) to [`AsyncArrowWriter`]. +//! +//! [`async_reader`]: arrow::async_reader +//! [`async_writer`]: arrow::async_writer +//! [`AsyncArrowWriter`]: arrow::async_writer::AsyncArrowWriter +//! [`AsyncFileReader`]: arrow::async_reader::AsyncFileReader +//! [`AsyncWrite`]: https://docs.rs/tokio/latest/tokio/io/trait.AsyncWrite.html +//! [`ParquetRecordBatchStreamBuilder`]: arrow::async_reader::ParquetRecordBatchStreamBuilder +//! +//! ## Variant Logical Type (`variant_experimental` feature) +//! +//! The [`variant`] module supports reading and writing Parquet files +//! with the [Variant Binary Encoding] logical type, which can represent +//! semi-structured data such as JSON efficiently. +//! +//! [Variant Binary Encoding]: https://github.com/apache/parquet-format/blob/master/VariantEncoding.md +//! +//! ## Read/Write Parquet Directly +//! +//! Workloads needing finer-grained control, or to avoid a dependence on arrow, +//! can use the APIs in [`mod@file`] directly. These APIs are harder to use +//! as they directly use the underlying Parquet data model, and require knowledge +//! of the Parquet format, including the details of [Dremel] record shredding +//! and [Logical Types]. +//! +//! [arrow]: https://docs.rs/arrow/latest/arrow/index.html +//! [Arrow]: https://arrow.apache.org/ +//! [`RecordBatch`]: https://docs.rs/arrow/latest/arrow/array/struct.RecordBatch.html +//! [CSV]: https://en.wikipedia.org/wiki/Comma-separated_values +//! [Dremel]: https://research.google/pubs/pub36632/ +//! [Logical Types]: https://github.com/apache/parquet-format/blob/master/LogicalTypes.md +//! [object_store]: https://docs.rs/object_store/latest/object_store/ + +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/parquet-format/25f05e73d8cd7f5c83532ce51cb4f4de8ba5f2a2/logo/parquet-logos_1.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/parquet-format/25f05e73d8cd7f5c83532ce51cb4f4de8ba5f2a2/logo/parquet-logos_1.svg" +)] +#![cfg_attr(docsrs, feature(doc_cfg))] +#![warn(missing_docs)] +/// Defines a an item with an experimental public API +/// +/// The module will not be documented, and will only be public if the +/// experimental feature flag is enabled +/// +/// Experimental components have no stability guarantees +#[cfg(feature = "experimental")] +macro_rules! experimental { + ($(#[$meta:meta])* $vis:vis mod $module:ident) => { + #[doc(hidden)] + $(#[$meta])* + pub mod $module; + } +} + +#[cfg(not(feature = "experimental"))] +macro_rules! experimental { + ($(#[$meta:meta])* $vis:vis mod $module:ident) => { + $(#[$meta])* + $vis mod $module; + } +} + +#[cfg(all( + feature = "flate2", + not(any(feature = "flate2-zlib-rs", feature = "flate2-rust_backend")) +))] +compile_error!( + "When enabling `flate2` you must enable one of the features: `flate2-zlib-rs` or `flate2-rust_backend`." +); + +#[macro_use] +pub mod errors; +pub mod basic; + +#[macro_use] +pub mod data_type; + +use std::fmt::Debug; +use std::ops::Range; +// Exported for external use, such as benchmarks +#[cfg(feature = "experimental")] +#[doc(hidden)] +pub use self::encodings::{decoding, encoding}; + +experimental!(#[macro_use] mod util); + +pub use util::utf8; + +#[cfg(feature = "arrow")] +pub mod arrow; +pub mod column; +experimental!(mod compression); +experimental!(mod encodings); +pub mod bloom_filter; + +#[cfg(feature = "encryption")] +experimental!(pub mod encryption); + +pub mod file; +pub mod record; +pub mod schema; + +mod parquet_macros; +mod parquet_thrift; + +/// What data is needed to read the next item from a decoder. +/// +/// This is used to communicate between the decoder and the caller +/// to indicate what data is needed next, or what the result of decoding is. +#[derive(Debug)] +pub enum DecodeResult { + /// The ranges of data necessary to proceed + // TODO: distinguish between minimim needed to make progress and what could be used? + NeedsData(Vec>), + /// The decoder produced an output item + Data(T), + /// The decoder finished processing + Finished, +} + +#[cfg(feature = "variant_experimental")] +pub mod variant; +experimental!(pub mod geospatial); diff --git a/vendor/parquet/src/parquet_macros.rs b/vendor/parquet/src/parquet_macros.rs new file mode 100644 index 0000000..f7ddf57 --- /dev/null +++ b/vendor/parquet/src/parquet_macros.rs @@ -0,0 +1,614 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// These macros are adapted from Jörn Horstmann's thrift macros at +// https://github.com/jhorstmann/compact-thrift +// They allow for pasting sections of the Parquet thrift IDL file +// into a macro to generate rust structures and implementations. + +//! This is a collection of macros used to parse Thrift IDL descriptions of structs, +//! unions, and enums into their corresponding Rust types. These macros will also +//! generate the code necessary to serialize and deserialize to/from the [Thrift compact] +//! protocol. +//! +//! Further details of how to use them (and other aspects of the Thrift serialization process) +//! can be found in [THRIFT.md]. +//! +//! [Thrift compact]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#list-and-set +//! [THRIFT.md]: https://github.com/apache/arrow-rs/blob/main/parquet/THRIFT.md + +#[doc(hidden)] +#[macro_export] +#[allow(clippy::crate_in_macro_def)] +/// Macro used to generate rust enums from a Thrift `enum` definition. +/// +/// Note: +/// - All enums generated with this macro will have `pub` visibility. +/// - When utilizing this macro the Thrift serialization traits and structs need to be in scope. +macro_rules! thrift_enum { + ($(#[$($def_attrs:tt)*])* enum $identifier:ident { $($(#[$($field_attrs:tt)*])* $field_name:ident = $field_value:literal;)* }) => { + $(#[$($def_attrs)*])* + #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] + #[allow(non_camel_case_types)] + #[allow(missing_docs)] + pub enum $identifier { + $($(#[cfg_attr(not(doctest), $($field_attrs)*)])* $field_name = $field_value,)* + } + + impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for $identifier { + #[allow(deprecated)] + fn read_thrift(prot: &mut R) -> Result { + let val = prot.read_i32()?; + match val { + $($field_value => Ok(Self::$field_name),)* + _ => Err(general_err!("Unexpected {} {}", stringify!($identifier), val)), + } + } + } + + impl fmt::Display for $identifier { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{self:?}") + } + } + + impl WriteThrift for $identifier { + const ELEMENT_TYPE: ElementType = ElementType::I32; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + writer.write_i32(*self as i32) + } + } + + impl WriteThriftField for $identifier { + fn write_thrift_field(&self, writer: &mut ThriftCompactOutputProtocol, field_id: i16, last_field_id: i16) -> Result { + writer.write_field_begin(FieldType::I32, field_id, last_field_id)?; + self.write_thrift(writer)?; + Ok(field_id) + } + } + + impl $identifier { + #[allow(deprecated)] + #[doc = "Returns a slice containing every variant of this enum."] + #[allow(dead_code)] + pub const VARIANTS: &'static [Self] = &[ + $(Self::$field_name),* + ]; + + #[allow(deprecated)] + const fn max_discriminant_impl() -> i32 { + let values: &[i32] = &[$($field_value),*]; + let mut max = values[0]; + let mut idx = 1; + while idx < values.len() { + let candidate = values[idx]; + if candidate > max { + max = candidate; + } + idx += 1; + } + max + } + + #[allow(deprecated)] + #[doc = "Returns the largest discriminant value defined for this enum."] + #[allow(dead_code)] + pub const MAX_DISCRIMINANT: i32 = Self::max_discriminant_impl(); + } + } +} + +/// Macro used to generate Rust enums for Thrift unions in which all variants are typed with empty +/// structs. +/// +/// Because the compact protocol does not write any struct type information, these empty structs +/// become a single `0` (end-of-fields marker) upon serialization. Rather than trying to deserialize +/// an empty struct, we can instead simply read the `0` and discard it. +/// +/// The resulting Rust enum will have all unit variants. +/// +/// Note: +/// - All enums generated with this macro will have `pub` visibility. +/// - When utilizing this macro the Thrift serialization traits and structs need to be in scope. +#[doc(hidden)] +#[macro_export] +#[allow(clippy::crate_in_macro_def)] +macro_rules! thrift_union_all_empty { + ($(#[$($def_attrs:tt)*])* union $identifier:ident { $($(#[$($field_attrs:tt)*])* $field_id:literal : $field_type:ident $(< $element_type:ident >)? $field_name:ident $(;)?)* }) => { + $(#[cfg_attr(not(doctest), $($def_attrs)*)])* + #[derive(Clone, Copy, Debug, Eq, PartialEq)] + #[allow(non_camel_case_types)] + #[allow(non_snake_case)] + #[allow(missing_docs)] + pub enum $identifier { + $($(#[cfg_attr(not(doctest), $($field_attrs)*)])* $field_name),* + } + + impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for $identifier { + fn read_thrift(prot: &mut R) -> Result { + let field_ident = prot.read_field_begin(0)?; + if field_ident.field_type == FieldType::Stop { + return Err(general_err!("Received empty union from remote {}", stringify!($identifier))); + } + let ret = match field_ident.id { + $($field_id => { + prot.skip_empty_struct()?; + Self::$field_name + } + )* + _ => { + return Err(general_err!("Unexpected {} {}", stringify!($identifier), field_ident.id)); + } + }; + let field_ident = prot.read_field_begin(field_ident.id)?; + if field_ident.field_type != FieldType::Stop { + return Err(general_err!( + "Received multiple fields for union from remote {}", stringify!($identifier) + )); + } + Ok(ret) + } + } + + impl WriteThrift for $identifier { + const ELEMENT_TYPE: ElementType = ElementType::Struct; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + match *self { + $(Self::$field_name => writer.write_empty_struct($field_id, 0)?,)* + }; + // write end of struct for this union + writer.write_struct_end() + } + } + + impl WriteThriftField for $identifier { + fn write_thrift_field(&self, writer: &mut ThriftCompactOutputProtocol, field_id: i16, last_field_id: i16) -> Result { + writer.write_field_begin(FieldType::Struct, field_id, last_field_id)?; + self.write_thrift(writer)?; + Ok(field_id) + } + } + } +} + +/// Macro used to generate Rust enums for Thrift unions where variants are a mix of unit and +/// tuple types. +/// +/// Use of this macro requires modifying the thrift IDL. For variants with empty structs as their +/// type, delete the typename (i.e. `1: EmptyStruct Var1;` becomes `1: Var1`). For variants with a +/// non-empty type, the typename must be contained within parens (e.g. `1: MyType Var1;` becomes +/// `1: (MyType) Var1;`). +/// +/// Note: +/// - All enums generated with this macro will have `pub` visibility. +/// - This macro allows for specifying lifetime annotations for the resulting `enum` and its fields. +/// - When utilizing this macro the Thrift serialization traits and structs need to be in scope. +#[doc(hidden)] +#[macro_export] +#[allow(clippy::crate_in_macro_def)] +macro_rules! thrift_union { + ($(#[$($def_attrs:tt)*])* union $identifier:ident $(< $lt:lifetime >)? { $($(#[$($field_attrs:tt)*])* $field_id:literal : $( ( $field_type:ident $(< $element_type:ident >)? $(< $field_lt:lifetime >)?) )? $field_name:ident $(;)?)* }) => { + $(#[cfg_attr(not(doctest), $($def_attrs)*)])* + #[derive(Clone, Debug, Eq, PartialEq)] + #[allow(non_camel_case_types)] + #[allow(non_snake_case)] + #[allow(missing_docs)] + pub enum $identifier $(<$lt>)? { + $($(#[cfg_attr(not(doctest), $($field_attrs)*)])* $field_name $( ( $crate::__thrift_union_type!{$field_type $($field_lt)? $($element_type)?} ) )?),* + } + + impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for $identifier $(<$lt>)? { + fn read_thrift(prot: &mut R) -> Result { + let field_ident = prot.read_field_begin(0)?; + if field_ident.field_type == FieldType::Stop { + return Err(general_err!("Received empty union from remote {}", stringify!($identifier))); + } + let ret = match field_ident.id { + $($field_id => { + let val = $crate::__thrift_read_variant!(prot, $field_name $($field_type $($element_type)?)?); + val + })* + _ => { + return Err(general_err!("Unexpected {} {}", stringify!($identifier), field_ident.id)); + } + }; + let field_ident = prot.read_field_begin(field_ident.id)?; + if field_ident.field_type != FieldType::Stop { + return Err(general_err!( + concat!("Received multiple fields for union from remote {}", stringify!($identifier)) + )); + } + Ok(ret) + } + } + + impl $(<$lt>)? WriteThrift for $identifier $(<$lt>)? { + const ELEMENT_TYPE: ElementType = ElementType::Struct; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + match self { + $($crate::__thrift_write_variant_lhs!($field_name $($field_type)?, variant_val) => + $crate::__thrift_write_variant_rhs!($field_id $($field_type)?, writer, variant_val),)* + }; + writer.write_struct_end() + } + } + + impl $(<$lt>)? WriteThriftField for $identifier $(<$lt>)? { + fn write_thrift_field(&self, writer: &mut ThriftCompactOutputProtocol, field_id: i16, last_field_id: i16) -> Result { + writer.write_field_begin(FieldType::Struct, field_id, last_field_id)?; + self.write_thrift(writer)?; + Ok(field_id) + } + } + } +} + +/// Macro used to generate Rust enums for Thrift unions where variants are a mix of unit and +/// tuple types. This version allows for unknown variants for forwards compatibility. +/// +/// Use of this macro requires modifying the thrift IDL. For variants with empty structs as their +/// type, delete the typename (i.e. `1: EmptyStruct Var1;` becomes `1: Var1`). For variants with a +/// non-empty type, the typename must be contained within parens (e.g. `1: MyType Var1;` becomes +/// `1: (MyType) Var1;`). +/// +/// This macro allows for specifying lifetime annotations for the resulting `enum` and its fields. +/// +/// When utilizing this macro the Thrift serialization traits and structs need to be in scope. +#[doc(hidden)] +#[macro_export] +#[allow(clippy::crate_in_macro_def)] +macro_rules! thrift_union_with_unknown { + ($(#[$($def_attrs:tt)*])* union $identifier:ident $(< $lt:lifetime >)? { $($(#[$($field_attrs:tt)*])* $field_id:literal : $( ( $field_type:ident $(< $element_type:ident >)? $(< $field_lt:lifetime >)?) )? $field_name:ident $(;)?)* }) => { + $(#[cfg_attr(not(doctest), $($def_attrs)*)])* + #[derive(Clone, Debug, Eq, PartialEq)] + #[allow(non_camel_case_types)] + #[allow(non_snake_case)] + #[allow(missing_docs)] + pub enum $identifier $(<$lt>)? { + $($(#[cfg_attr(not(doctest), $($field_attrs)*)])* $field_name $( ( $crate::__thrift_union_type!{$field_type $($field_lt)? $($element_type)?} ) )?),*, + _Unknown { + /// The field id encountered when parsing the unknown variant. + field_id: i16, + }, + } + + impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for $identifier $(<$lt>)? { + fn read_thrift(prot: &mut R) -> Result { + let field_ident = prot.read_field_begin(0)?; + if field_ident.field_type == FieldType::Stop { + return Err(general_err!("Received empty union from remote {}", stringify!($identifier))); + } + let ret = match field_ident.id { + $($field_id => { + let val = $crate::__thrift_read_variant!(prot, $field_name $($field_type $($element_type)?)?); + val + })* + _ => { + prot.skip(field_ident.field_type)?; + Self::_Unknown { + field_id: field_ident.id, + } + } + }; + let field_ident = prot.read_field_begin(field_ident.id)?; + if field_ident.field_type != FieldType::Stop { + return Err(general_err!( + concat!("Received multiple fields for union from remote {}", stringify!($identifier)) + )); + } + Ok(ret) + } + } + + impl $(<$lt>)? WriteThrift for $identifier $(<$lt>)? { + const ELEMENT_TYPE: ElementType = ElementType::Struct; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + match self { + $($crate::__thrift_write_variant_lhs!($field_name $($field_type)?, variant_val) => + $crate::__thrift_write_variant_rhs!($field_id $($field_type)?, writer, variant_val),)* + Self::_Unknown{..} => return Err(general_err!("Trying to write unknown variant")), + }; + writer.write_struct_end() + } + } + + impl $(<$lt>)? WriteThriftField for $identifier $(<$lt>)? { + fn write_thrift_field(&self, writer: &mut ThriftCompactOutputProtocol, field_id: i16, last_field_id: i16) -> Result { + writer.write_field_begin(FieldType::Struct, field_id, last_field_id)?; + self.write_thrift(writer)?; + Ok(field_id) + } + } + } +} + +/// Macro used to generate Rust structs from a Thrift `struct` definition. +/// +/// Note: +/// - This macro allows for specifying the visibility of the resulting `struct` and its fields. +/// + The `struct` and all fields will have the same visibility. +/// - This macro allows for specifying lifetime annotations for the resulting `struct` and its fields. +/// - When utilizing this macro the Thrift serialization traits and structs need to be in scope. +#[doc(hidden)] +#[macro_export] +macro_rules! thrift_struct { + ($(#[$($def_attrs:tt)*])* $vis:vis struct $identifier:ident $(< $lt:lifetime >)? { $($(#[$($field_attrs:tt)*])* $field_id:literal : $required_or_optional:ident $field_type:ident $(< $field_lt:lifetime >)? $(< $element_type:ident >)? $field_name:ident $(= $default_value:literal)? $(;)?)* }) => { + $(#[cfg_attr(not(doctest), $($def_attrs)*)])* + #[derive(Clone, Debug, Eq, PartialEq)] + #[allow(non_camel_case_types)] + #[allow(non_snake_case)] + #[allow(missing_docs)] + $vis struct $identifier $(<$lt>)? { + $($(#[cfg_attr(not(doctest), $($field_attrs)*)])* $vis $field_name: $crate::__thrift_required_or_optional!($required_or_optional $crate::__thrift_field_type!($field_type $($field_lt)? $($element_type)?))),* + } + + impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for $identifier $(<$lt>)? { + fn read_thrift(prot: &mut R) -> Result { + $(let mut $field_name: Option<$crate::__thrift_field_type!($field_type $($field_lt)? $($element_type)?)> = None;)* + let mut last_field_id = 0i16; + loop { + let field_ident = prot.read_field_begin(last_field_id)?; + if field_ident.field_type == FieldType::Stop { + break; + } + match field_ident.id { + $($field_id => { + let val = $crate::__thrift_read_field!(prot, field_ident, $field_type $($field_lt)? $($element_type)?); + $field_name = Some(val); + })* + _ => { + prot.skip(field_ident.field_type)?; + } + }; + last_field_id = field_ident.id; + } + $($crate::__thrift_result_required_or_optional!($required_or_optional $field_name);)* + Ok(Self { + $($field_name),* + }) + } + } + + impl $(<$lt>)? WriteThrift for $identifier $(<$lt>)? { + const ELEMENT_TYPE: ElementType = ElementType::Struct; + + #[allow(unused_assignments)] + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + #[allow(unused_mut, unused_variables)] + let mut last_field_id = 0i16; + $($crate::__thrift_write_required_or_optional_field!($required_or_optional $field_name, $field_id, $field_type, self, writer, last_field_id);)* + writer.write_struct_end() + } + } + + impl $(<$lt>)? WriteThriftField for $identifier $(<$lt>)? { + fn write_thrift_field(&self, writer: &mut ThriftCompactOutputProtocol, field_id: i16, last_field_id: i16) -> Result { + writer.write_field_begin(FieldType::Struct, field_id, last_field_id)?; + self.write_thrift(writer)?; + Ok(field_id) + } + } + } +} + +#[doc(hidden)] +#[macro_export] +/// Generate `WriteThriftField` implementation for a struct. +macro_rules! write_thrift_field { + ($identifier:ident $(< $lt:lifetime >)?, $fld_type:expr) => { + impl $(<$lt>)? WriteThriftField for $identifier $(<$lt>)? { + fn write_thrift_field(&self, writer: &mut ThriftCompactOutputProtocol, field_id: i16, last_field_id: i16) -> Result { + writer.write_field_begin($fld_type, field_id, last_field_id)?; + self.write_thrift(writer)?; + Ok(field_id) + } + } + } +} + +#[doc(hidden)] +#[macro_export] +macro_rules! __thrift_write_required_or_optional_field { + (required $field_name:ident, $field_id:literal, $field_type:ident, $self:tt, $writer:tt, $last_id:tt) => { + $crate::__thrift_write_required_field!( + $field_type, + $field_name, + $field_id, + $self, + $writer, + $last_id + ) + }; + (optional $field_name:ident, $field_id:literal, $field_type:ident, $self:tt, $writer:tt, $last_id:tt) => { + $crate::__thrift_write_optional_field!( + $field_type, + $field_name, + $field_id, + $self, + $writer, + $last_id + ) + }; +} + +#[doc(hidden)] +#[macro_export] +macro_rules! __thrift_write_required_field { + (binary, $field_name:ident, $field_id:literal, $self:ident, $writer:ident, $last_id:ident) => { + $writer.write_field_begin(FieldType::Binary, $field_id, $last_id)?; + $writer.write_bytes($self.$field_name)?; + $last_id = $field_id; + }; + ($field_type:ident, $field_name:ident, $field_id:literal, $self:ident, $writer:ident, $last_id:ident) => { + $last_id = $self + .$field_name + .write_thrift_field($writer, $field_id, $last_id)?; + }; +} + +#[doc(hidden)] +#[macro_export] +macro_rules! __thrift_write_optional_field { + (binary, $field_name:ident, $field_id:literal, $self:ident, $writer:tt, $last_id:tt) => { + if $self.$field_name.is_some() { + $writer.write_field_begin(FieldType::Binary, $field_id, $last_id)?; + $writer.write_bytes($self.$field_name.as_ref().unwrap())?; + $last_id = $field_id; + } + }; + ($field_type:ident, $field_name:ident, $field_id:literal, $self:ident, $writer:tt, $last_id:tt) => { + if $self.$field_name.is_some() { + $last_id = $self + .$field_name + .as_ref() + .unwrap() + .write_thrift_field($writer, $field_id, $last_id)?; + } + }; +} + +#[doc(hidden)] +#[macro_export] +macro_rules! __thrift_required_or_optional { + (required $field_type:ty) => { $field_type }; + (optional $field_type:ty) => { Option<$field_type> }; +} + +// Performance note: using `expect` here is about 4% faster on the page index bench, +// but we want to propagate errors. Using `ok_or` is *much* slower. +#[doc(hidden)] +#[macro_export] +macro_rules! __thrift_result_required_or_optional { + (required $field_name:ident) => { + let Some($field_name) = $field_name else { + return Err(general_err!(concat!( + "Required field ", + stringify!($field_name), + " is missing", + ))); + }; + }; + (optional $field_name:ident) => {}; +} + +#[doc(hidden)] +#[macro_export] +macro_rules! __thrift_read_field { + ($prot:tt, $field_ident:tt, list $lt:lifetime binary) => { + read_thrift_vec::<&'a [u8], R>(&mut *$prot)? + }; + ($prot:tt, $field_ident:tt, list $lt:lifetime $element_type:ident) => { + read_thrift_vec::<$element_type, R>(&mut *$prot)? + }; + ($prot:tt, $field_ident:tt, list string) => { + read_thrift_vec::(&mut *$prot)? + }; + ($prot:tt, $field_ident:tt, list $element_type:ident) => { + read_thrift_vec::<$element_type, R>(&mut *$prot)? + }; + ($prot:tt, $field_ident:tt, string $lt:lifetime) => { + <&$lt str>::read_thrift(&mut *$prot)? + }; + ($prot:tt, $field_ident:tt, binary $lt:lifetime) => { + <&$lt [u8]>::read_thrift(&mut *$prot)? + }; + ($prot:tt, $field_ident:tt, $field_type:ident $lt:lifetime) => { + $field_type::read_thrift(&mut *$prot)? + }; + ($prot:tt, $field_ident:tt, string) => { + String::read_thrift(&mut *$prot)? + }; + ($prot:tt, $field_ident:tt, binary) => { + // this one needs to not conflict with `list` + $prot.read_bytes_owned()? + }; + ($prot:tt, $field_ident:tt, double) => { + $crate::parquet_thrift::OrderedF64::read_thrift(&mut *$prot)? + }; + ($prot:tt, $field_ident:tt, bool) => { + $field_ident.bool_val()? + }; + ($prot:tt, $field_ident:tt, $field_type:ident) => { + $field_type::read_thrift(&mut *$prot)? + }; +} + +#[doc(hidden)] +#[macro_export] +macro_rules! __thrift_field_type { + (binary $lt:lifetime) => { &$lt [u8] }; + (string $lt:lifetime) => { &$lt str }; + ($field_type:ident $lt:lifetime) => { $field_type<$lt> }; + (list $lt:lifetime $element_type:ident) => { Vec< $crate::__thrift_field_type!($element_type $lt) > }; + (list string) => { Vec }; + (list $element_type:ident) => { Vec< $crate::__thrift_field_type!($element_type) > }; + (binary) => { Vec }; + (string) => { String }; + (double) => { $crate::parquet_thrift::OrderedF64 }; + ($field_type:ty) => { $field_type }; +} + +#[doc(hidden)] +#[macro_export] +macro_rules! __thrift_union_type { + (binary $lt:lifetime) => { &$lt [u8] }; + (string $lt:lifetime) => { &$lt str }; + ($field_type:ident $lt:lifetime) => { $field_type<$lt> }; + ($field_type:ident) => { $field_type }; + (list $field_type:ident) => { Vec<$field_type> }; +} + +#[doc(hidden)] +#[macro_export] +macro_rules! __thrift_read_variant { + ($prot:tt, $field_name:ident $field_type:ident) => { + Self::$field_name($field_type::read_thrift(&mut *$prot)?) + }; + ($prot:tt, $field_name:ident list $field_type:ident) => { + Self::$field_name(Vec::<$field_type>::read_thrift(&mut *$prot)?) + }; + ($prot:tt, $field_name:ident) => {{ + $prot.skip_empty_struct()?; + Self::$field_name + }}; +} + +#[doc(hidden)] +#[macro_export] +macro_rules! __thrift_write_variant_lhs { + ($field_name:ident $field_type:ident, $val:tt) => { + Self::$field_name($val) + }; + ($field_name:ident, $val:tt) => { + Self::$field_name + }; +} + +#[doc(hidden)] +#[macro_export] +macro_rules! __thrift_write_variant_rhs { + ($field_id:literal $field_type:ident, $writer:tt, $val:ident) => { + $val.write_thrift_field($writer, $field_id, 0)? + }; + ($field_id:literal, $writer:tt, $val:tt) => { + $writer.write_empty_struct($field_id, 0)? + }; +} diff --git a/vendor/parquet/src/parquet_thrift.rs b/vendor/parquet/src/parquet_thrift.rs new file mode 100644 index 0000000..cc6390b --- /dev/null +++ b/vendor/parquet/src/parquet_thrift.rs @@ -0,0 +1,1201 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Structs used for encoding and decoding Parquet Thrift objects. +//! +//! These include: +//! * [`ThriftCompactInputProtocol`]: Trait implemented by Thrift decoders. +//! * [`ThriftSliceInputProtocol`]: Thrift decoder that takes a slice of bytes as input. +//! * [`ThriftReadInputProtocol`]: Thrift decoder that takes a [`Read`] as input. +//! * [`ReadThrift`]: Trait implemented by serializable objects. +//! * [`ThriftCompactOutputProtocol`]: Thrift encoder. +//! * [`WriteThrift`]: Trait implemented by serializable objects. +//! * [`WriteThriftField`]: Trait implemented by serializable objects that are fields in Thrift structs. + +use std::{ + cmp::Ordering, + io::{Read, Write}, +}; + +use crate::{ + errors::{ParquetError, Result}, + write_thrift_field, +}; +use std::io::Error; +use std::num::TryFromIntError; +use std::str::Utf8Error; + +#[derive(Debug)] +pub(crate) enum ThriftProtocolError { + Eof, + IO(Error), + InvalidFieldType(u8), + InvalidElementType(u8), + FieldDeltaOverflow { field_delta: u8, last_field_id: i16 }, + InvalidBoolean(u8), + IntegerOverflow, + Utf8Error, + SkipDepth(FieldType), + SkipUnsupportedType(FieldType), +} + +impl From for ParquetError { + #[inline(never)] + fn from(e: ThriftProtocolError) -> Self { + match e { + ThriftProtocolError::Eof => eof_err!("Unexpected EOF"), + ThriftProtocolError::IO(e) => e.into(), + ThriftProtocolError::InvalidFieldType(value) => match FieldType::try_from(value) { + Ok(fld_type) => general_err!("Unexpected struct field type {:?}", fld_type), + Err(_) => general_err!("Unexpected struct field type {}", value), + }, + ThriftProtocolError::InvalidElementType(value) => { + general_err!("Unexpected list/set element type {}", value) + } + ThriftProtocolError::FieldDeltaOverflow { + field_delta, + last_field_id, + } => general_err!("cannot add {} to {}", field_delta, last_field_id), + ThriftProtocolError::InvalidBoolean(value) => { + general_err!("cannot convert {} into bool", value) + } + ThriftProtocolError::IntegerOverflow => { + general_err!("integer overflow decoding thrift value") + } + ThriftProtocolError::Utf8Error => general_err!("invalid utf8"), + ThriftProtocolError::SkipDepth(field_type) => { + general_err!("cannot parse past {:?}", field_type) + } + ThriftProtocolError::SkipUnsupportedType(field_type) => { + general_err!("cannot skip field type {:?}", field_type) + } + } + } +} + +impl From for ThriftProtocolError { + fn from(_: Utf8Error) -> Self { + // ignore error payload to reduce the size of ThriftProtocolError + Self::Utf8Error + } +} + +impl From for ThriftProtocolError { + fn from(e: Error) -> Self { + Self::IO(e) + } +} + +impl From for ThriftProtocolError { + fn from(_: TryFromIntError) -> Self { + // ignore error payload to reduce the size of ThriftProtocolError + Self::IntegerOverflow + } +} + +pub type ThriftProtocolResult = Result; + +/// Wrapper for thrift `double` fields. This is used to provide +/// an implementation of `Eq` for floats. This implementation +/// uses IEEE 754 total order. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct OrderedF64(f64); + +impl From for OrderedF64 { + fn from(value: f64) -> Self { + Self(value) + } +} + +impl From for f64 { + fn from(value: OrderedF64) -> Self { + value.0 + } +} + +impl Eq for OrderedF64 {} // Marker trait, requires PartialEq + +impl Ord for OrderedF64 { + fn cmp(&self, other: &Self) -> Ordering { + self.0.total_cmp(&other.0) + } +} + +impl PartialOrd for OrderedF64 { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +// Thrift compact protocol types for struct fields. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum FieldType { + Stop = 0, + BooleanTrue = 1, + BooleanFalse = 2, + Byte = 3, + I16 = 4, + I32 = 5, + I64 = 6, + Double = 7, + Binary = 8, + List = 9, + Set = 10, + Map = 11, + Struct = 12, + Uuid = 13, +} + +impl TryFrom for FieldType { + type Error = ThriftProtocolError; + fn try_from(value: u8) -> ThriftProtocolResult { + match value { + 0 => Ok(Self::Stop), + 1 => Ok(Self::BooleanTrue), + 2 => Ok(Self::BooleanFalse), + 3 => Ok(Self::Byte), + 4 => Ok(Self::I16), + 5 => Ok(Self::I32), + 6 => Ok(Self::I64), + 7 => Ok(Self::Double), + 8 => Ok(Self::Binary), + 9 => Ok(Self::List), + 10 => Ok(Self::Set), + 11 => Ok(Self::Map), + 12 => Ok(Self::Struct), + 13 => Ok(Self::Uuid), + _ => Err(ThriftProtocolError::InvalidFieldType(value)), + } + } +} + +impl From for FieldType { + fn from(value: ElementType) -> Self { + match value { + ElementType::Bool => Self::BooleanTrue, + ElementType::Byte => Self::Byte, + ElementType::I16 => Self::I16, + ElementType::I32 => Self::I32, + ElementType::I64 => Self::I64, + ElementType::Double => Self::Double, + ElementType::Binary => Self::Binary, + ElementType::List => Self::List, + ElementType::Set => Self::Set, + ElementType::Map => Self::Map, + ElementType::Struct => Self::Struct, + ElementType::Uuid => Self::Uuid, + } + } +} + +// Thrift compact protocol types for list elements +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum ElementType { + Bool = 2, + Byte = 3, + I16 = 4, + I32 = 5, + I64 = 6, + Double = 7, + Binary = 8, + List = 9, + Set = 10, + Map = 11, + Struct = 12, + Uuid = 13, +} + +impl TryFrom for ElementType { + type Error = ThriftProtocolError; + fn try_from(value: u8) -> ThriftProtocolResult { + match value { + // For historical and compatibility reasons, a reader should be capable to deal with both cases. + // The only valid value in the original spec was 2, but due to an widespread implementation bug + // the defacto standard across large parts of the library became 1 instead. + // As a result, both values are now allowed. + // https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#list-and-set + 1 | 2 => Ok(Self::Bool), + 3 => Ok(Self::Byte), + 4 => Ok(Self::I16), + 5 => Ok(Self::I32), + 6 => Ok(Self::I64), + 7 => Ok(Self::Double), + 8 => Ok(Self::Binary), + 9 => Ok(Self::List), + 10 => Ok(Self::Set), + 11 => Ok(Self::Map), + 12 => Ok(Self::Struct), + 13 => Ok(Self::Uuid), + _ => Err(ThriftProtocolError::InvalidElementType(value)), + } + } +} + +/// Struct used to describe a [thrift struct] field during decoding. +/// +/// [thrift struct]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#struct-encoding +pub(crate) struct FieldIdentifier { + /// The type for the field. + pub(crate) field_type: FieldType, + /// The field's `id`. May be computed from delta or directly decoded. + pub(crate) id: i16, +} + +impl FieldIdentifier { + pub(crate) fn bool_val(&self) -> ThriftProtocolResult { + match self.field_type { + FieldType::BooleanTrue => Ok(true), + FieldType::BooleanFalse => Ok(false), + _ => Err(ThriftProtocolError::InvalidFieldType(self.field_type as u8)), + } + } +} + +/// Struct used to describe a [thrift list]. +/// +/// [thrift list]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#list-and-set +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct ListIdentifier { + /// The type for each element in the list. + pub(crate) element_type: ElementType, + /// Number of elements contained in the list. + pub(crate) size: i32, +} + +/// Low-level object used to deserialize structs encoded with the Thrift [compact] protocol. +/// +/// Implementation of this trait must provide the low-level functions `read_byte`, `read_bytes`, +/// `skip_bytes`, and `read_double`. These primitives are used by the default functions provided +/// here to perform deserialization. +/// +/// [compact]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md +pub(crate) trait ThriftCompactInputProtocol<'a> { + /// Read a single byte from the input. + fn read_byte(&mut self) -> ThriftProtocolResult; + + /// Read a Thrift encoded [binary] from the input. + /// + /// [binary]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#binary-encoding + fn read_bytes(&mut self) -> ThriftProtocolResult<&'a [u8]>; + + fn read_bytes_owned(&mut self) -> ThriftProtocolResult>; + + /// Skip the next `n` bytes of input. + fn skip_bytes(&mut self, n: usize) -> ThriftProtocolResult<()>; + + /// Read a ULEB128 encoded unsigned varint from the input. + fn read_vlq(&mut self) -> ThriftProtocolResult { + // try the happy path first + let byte = self.read_byte()?; + if byte & 0x80 == 0 { + return Ok(byte as u64); + } + let mut in_progress = (byte & 0x7f) as u64; + let mut shift = 7; + loop { + let byte = self.read_byte()?; + in_progress |= ((byte & 0x7F) as u64).wrapping_shl(shift); + if byte & 0x80 == 0 { + return Ok(in_progress); + } + shift += 7; + } + } + + /// Read a zig-zag encoded signed varint from the input. + fn read_zig_zag(&mut self) -> ThriftProtocolResult { + let val = self.read_vlq()?; + Ok((val >> 1) as i64 ^ -((val & 1) as i64)) + } + + /// Read the [`ListIdentifier`] for a Thrift encoded list. + fn read_list_begin(&mut self) -> ThriftProtocolResult { + let header = self.read_byte()?; + // some parquet writers will have an element_type of 0 for an empty list. + // account for that and return a bogus but valid element_type. + if header == 0 { + return Ok(ListIdentifier { + element_type: ElementType::Byte, + size: 0, + }); + } + let element_type = ElementType::try_from(header & 0x0f)?; + + let possible_element_count = (header & 0xF0) >> 4; + let element_count = if possible_element_count != 15 { + // high bits set high if count and type encoded separately + possible_element_count as i32 + } else { + // The list size on the wire is an unsigned varint, but we represent + // it as `i32` (matching Java's `int` and the Thrift schema). + // A varint that decodes above `i32::MAX` is malformed input — reject + // it here at the protocol layer rather than letting the cast wrap + // into a negative size that downstream allocation code has to + // re-validate. + i32::try_from(self.read_vlq()?)? + }; + + Ok(ListIdentifier { + element_type, + size: element_count, + }) + } + + // Full field ids are uncommon. + // Not inlining this method reduces the code size of `read_field_begin`, which then ideally gets + // inlined everywhere. + #[cold] + fn read_full_field_id(&mut self) -> ThriftProtocolResult { + self.read_i16() + } + + /// Read the [`FieldIdentifier`] for a field in a Thrift encoded struct. + fn read_field_begin(&mut self, last_field_id: i16) -> ThriftProtocolResult { + // we can read at least one byte, which is: + // - the type + // - the field delta and the type + let field_type = self.read_byte()?; + if field_type & 0xf == 0 { + return Ok(FieldIdentifier { + field_type: FieldType::Stop, + id: 0, + }); + } + + let field_delta = (field_type & 0xf0) >> 4; + let field_type = FieldType::try_from(field_type & 0xf)?; + + let id = if field_delta != 0 { + last_field_id.checked_add(field_delta as i16).ok_or( + ThriftProtocolError::FieldDeltaOverflow { + field_delta, + last_field_id, + }, + )? + } else { + self.read_full_field_id()? + }; + + Ok(FieldIdentifier { field_type, id }) + } + + /// This is a specialized version of [`Self::read_field_begin`], solely for use in parsing + /// simple structs. This function assumes that the delta field will always be less than 0xf, + /// fields will be in order, and no boolean fields will be read. + /// This also skips validation of the field type. + /// + /// Returns a tuple of `(field_type, field_delta)`. + fn read_field_header(&mut self) -> ThriftProtocolResult<(u8, u8)> { + let field_type = self.read_byte()?; + let field_delta = (field_type & 0xf0) >> 4; + let field_type = field_type & 0xf; + Ok((field_type, field_delta)) + } + + /// Read a boolean list element. This should not be used for struct fields. For the latter, + /// use the [`FieldIdentifier::bool_val`] field. + fn read_bool(&mut self) -> ThriftProtocolResult { + let b = self.read_byte()?; + // Previous versions of the thrift specification said to use 0 and 1 inside collections, + // but that differed from existing implementations. + // The specification was updated in https://github.com/apache/thrift/commit/2c29c5665bc442e703480bb0ee60fe925ffe02e8. + // At least the go implementation seems to have followed the previously documented values. + match b { + 0x01 => Ok(true), + 0x00 | 0x02 => Ok(false), + _ => Err(ThriftProtocolError::InvalidBoolean(b)), + } + } + + /// Read a Thrift [binary] as a UTF-8 encoded string. + /// + /// [binary]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#binary-encoding + fn read_string(&mut self) -> ThriftProtocolResult<&'a str> { + let slice = self.read_bytes()?; + Ok(std::str::from_utf8(slice)?) + } + + /// Read an `i8`. + fn read_i8(&mut self) -> ThriftProtocolResult { + Ok(self.read_byte()? as _) + } + + /// Read an `i16`. + fn read_i16(&mut self) -> ThriftProtocolResult { + Ok(self.read_zig_zag()? as _) + } + + /// Read an `i32`. + fn read_i32(&mut self) -> ThriftProtocolResult { + Ok(self.read_zig_zag()? as _) + } + + /// Read an `i64`. + fn read_i64(&mut self) -> ThriftProtocolResult { + self.read_zig_zag() + } + + /// Read a Thrift `double` as `f64`. + fn read_double(&mut self) -> ThriftProtocolResult; + + /// Skip a ULEB128 encoded varint. + fn skip_vlq(&mut self) -> ThriftProtocolResult<()> { + loop { + let byte = self.read_byte()?; + if byte & 0x80 == 0 { + return Ok(()); + } + } + } + + /// Skip a thrift [binary]. + /// + /// [binary]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#binary-encoding + fn skip_binary(&mut self) -> ThriftProtocolResult<()> { + let len = self.read_vlq()? as usize; + self.skip_bytes(len) + } + + /// Skip a field with type `field_type` recursively until the default + /// maximum skip depth (currently 64) is reached. + fn skip(&mut self, field_type: FieldType) -> ThriftProtocolResult<()> { + const DEFAULT_SKIP_DEPTH: i8 = 64; + self.skip_till_depth(field_type, DEFAULT_SKIP_DEPTH) + } + + /// Empty structs in unions consist of a single byte of 0 for the field stop record. + /// This skips that byte without encuring the cost of processing the [`FieldIdentifier`]. + /// Will return an error if the struct is not actually empty. + fn skip_empty_struct(&mut self) -> Result<()> { + let b = self.read_byte()?; + if b != 0 { + Err(general_err!("Empty struct has fields")) + } else { + Ok(()) + } + } + + /// Skip a field with type `field_type` recursively up to `depth` levels. + fn skip_till_depth(&mut self, field_type: FieldType, depth: i8) -> ThriftProtocolResult<()> { + if depth == 0 { + return Err(ThriftProtocolError::SkipDepth(field_type)); + } + + match field_type { + // boolean field has no data + FieldType::BooleanFalse | FieldType::BooleanTrue => Ok(()), + FieldType::Byte => self.read_i8().map(|_| ()), + FieldType::I16 => self.skip_vlq().map(|_| ()), + FieldType::I32 => self.skip_vlq().map(|_| ()), + FieldType::I64 => self.skip_vlq().map(|_| ()), + FieldType::Double => self.skip_bytes(8).map(|_| ()), + FieldType::Binary => self.skip_binary().map(|_| ()), + // see https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#struct + FieldType::Struct => { + loop { + // we don't need field id for skipping, so always pass 0 for last id + let field_ident = self.read_field_begin(0)?; + if field_ident.field_type == FieldType::Stop { + break; + } + self.skip_till_depth(field_ident.field_type, depth - 1)?; + } + Ok(()) + } + // lists and sets are encoded the same + // see https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#list-and-set + FieldType::List | FieldType::Set => { + let list_ident = self.read_list_begin()?; + let element_type = FieldType::from(list_ident.element_type); + for _ in 0..list_ident.size { + self.skip_till_depth(element_type, depth - 1)?; + } + Ok(()) + } + // see https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#map + FieldType::Map => { + let size = i32::try_from(self.read_vlq()?)?; + if size > 0 { + let kv = self.read_byte()?; + let key_type = FieldType::from(ElementType::try_from(kv >> 4)?); + let val_type = FieldType::from(ElementType::try_from(kv & 0xf)?); + for _ in 0..size { + self.skip_till_depth(key_type, depth - 1)?; + self.skip_till_depth(val_type, depth - 1)?; + } + } + Ok(()) + } + // see https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#universal-unique-identifier-encoding + FieldType::Uuid => self.skip_bytes(16).map(|_| ()), + _ => Err(ThriftProtocolError::SkipUnsupportedType(field_type)), + } + } +} + +/// A high performance Thrift reader that reads from a slice of bytes. +pub(crate) struct ThriftSliceInputProtocol<'a> { + buf: &'a [u8], +} + +impl<'a> ThriftSliceInputProtocol<'a> { + /// Create a new `ThriftSliceInputProtocol` using the bytes in `buf`. + pub fn new(buf: &'a [u8]) -> Self { + Self { buf } + } + + /// Return the current buffer as a slice. + pub fn as_slice(&self) -> &'a [u8] { + self.buf + } +} + +impl<'b, 'a: 'b> ThriftCompactInputProtocol<'b> for ThriftSliceInputProtocol<'a> { + #[inline] + fn read_byte(&mut self) -> ThriftProtocolResult { + let ret = *self.buf.first().ok_or(ThriftProtocolError::Eof)?; + self.buf = &self.buf[1..]; + Ok(ret) + } + + fn read_bytes(&mut self) -> ThriftProtocolResult<&'b [u8]> { + let len = self.read_vlq()? as usize; + let ret = self.buf.get(..len).ok_or(ThriftProtocolError::Eof)?; + self.buf = &self.buf[len..]; + Ok(ret) + } + + fn read_bytes_owned(&mut self) -> ThriftProtocolResult> { + Ok(self.read_bytes()?.to_vec()) + } + + #[inline] + fn skip_bytes(&mut self, n: usize) -> ThriftProtocolResult<()> { + self.buf.get(..n).ok_or(ThriftProtocolError::Eof)?; + self.buf = &self.buf[n..]; + Ok(()) + } + + fn read_double(&mut self) -> ThriftProtocolResult { + let slice = self.buf.get(..8).ok_or(ThriftProtocolError::Eof)?; + self.buf = &self.buf[8..]; + match slice.try_into() { + Ok(slice) => Ok(f64::from_le_bytes(slice)), + Err(_) => unreachable!(), + } + } +} + +/// A Thrift input protocol that wraps a [`Read`] object. +/// +/// Note that this is only intended for use in reading Parquet page headers. This will panic +/// if Thrift `binary` data is encountered because a slice of that data cannot be returned. +pub(crate) struct ThriftReadInputProtocol { + reader: R, +} + +impl ThriftReadInputProtocol { + pub(crate) fn new(reader: R) -> Self { + Self { reader } + } +} + +impl<'a, R: Read> ThriftCompactInputProtocol<'a> for ThriftReadInputProtocol { + #[inline] + fn read_byte(&mut self) -> ThriftProtocolResult { + let mut buf = [0_u8; 1]; + self.reader.read_exact(&mut buf)?; + Ok(buf[0]) + } + + fn read_bytes(&mut self) -> ThriftProtocolResult<&'a [u8]> { + unimplemented!() + } + + fn read_bytes_owned(&mut self) -> ThriftProtocolResult> { + let len = self.read_vlq()? as usize; + let mut v = Vec::with_capacity(len); + std::io::copy(&mut self.reader.by_ref().take(len as u64), &mut v)?; + Ok(v) + } + + fn skip_bytes(&mut self, n: usize) -> ThriftProtocolResult<()> { + std::io::copy( + &mut self.reader.by_ref().take(n as u64), + &mut std::io::sink(), + )?; + Ok(()) + } + + fn read_double(&mut self) -> ThriftProtocolResult { + let mut buf = [0_u8; 8]; + self.reader.read_exact(&mut buf)?; + Ok(f64::from_le_bytes(buf)) + } +} + +/// Trait implemented for objects that can be deserialized from a Thrift input stream. +/// Implementations are provided for Thrift primitive types. +pub(crate) trait ReadThrift<'a, R: ThriftCompactInputProtocol<'a>> { + /// Read an object of type `Self` from the input protocol object. + fn read_thrift(prot: &mut R) -> Result + where + Self: Sized; +} + +impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for bool { + fn read_thrift(prot: &mut R) -> Result { + Ok(prot.read_bool()?) + } +} + +impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for i8 { + fn read_thrift(prot: &mut R) -> Result { + Ok(prot.read_i8()?) + } +} + +impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for i16 { + fn read_thrift(prot: &mut R) -> Result { + Ok(prot.read_i16()?) + } +} + +impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for i32 { + fn read_thrift(prot: &mut R) -> Result { + Ok(prot.read_i32()?) + } +} + +impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for i64 { + fn read_thrift(prot: &mut R) -> Result { + Ok(prot.read_i64()?) + } +} + +impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for OrderedF64 { + fn read_thrift(prot: &mut R) -> Result { + Ok(OrderedF64(prot.read_double()?)) + } +} + +impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for &'a str { + fn read_thrift(prot: &mut R) -> Result { + Ok(prot.read_string()?) + } +} + +impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for String { + fn read_thrift(prot: &mut R) -> Result { + Ok(String::from_utf8(prot.read_bytes_owned()?)?) + } +} + +impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for &'a [u8] { + fn read_thrift(prot: &mut R) -> Result { + Ok(prot.read_bytes()?) + } +} + +/// Read a Thrift encoded [list] from the input protocol object. +/// +/// [list]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#list-and-set +pub(crate) fn read_thrift_vec<'a, T, R>(prot: &mut R) -> Result> +where + R: ThriftCompactInputProtocol<'a>, + T: ReadThrift<'a, R> + WriteThrift, +{ + let list_ident = prot.read_list_begin()?; + validate_list_type(T::ELEMENT_TYPE, &list_ident)?; + let mut res = Vec::with_capacity(list_ident.size as usize); + for _ in 0..list_ident.size { + let val = T::read_thrift(prot)?; + res.push(val); + } + Ok(res) +} + +pub(crate) fn validate_list_type(expected: ElementType, got: &ListIdentifier) -> Result<()> { + if got.element_type != expected { + return Err(general_err!( + "Expected list element type of {:?} but got {:?}", + expected, + got.element_type + )); + } + Ok(()) +} + +///////////////////////// +// thrift compact output + +/// Low-level object used to serialize structs to the Thrift [compact output] protocol. +/// +/// This struct serves as a wrapper around a [`Write`] object, to which thrift encoded data +/// will written. The implementation provides functions to write Thrift primitive types, as well +/// as functions used in the encoding of lists and structs. This should rarely be used directly, +/// but is instead intended for use by implementers of [`WriteThrift`] and [`WriteThriftField`]. +/// +/// [compact output]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md +pub(crate) struct ThriftCompactOutputProtocol { + writer: W, + write_path_in_schema: bool, +} + +impl ThriftCompactOutputProtocol { + /// Create a new `ThriftCompactOutputProtocol` wrapping the byte sink `writer`. + pub(crate) fn new(writer: W) -> Self { + Self { + writer, + write_path_in_schema: true, + } + } + + // TODO(ets): at some point there should probably be a properties object + // to control aspects of thrift output. But since this is the only option to date + // I'm choosing a simpler API. + /// Control the writing of the `path_in_schema` element of the `ColumnMetaData` + pub(crate) fn set_write_path_in_schema(&mut self, val: bool) { + self.write_path_in_schema = val; + } + + /// Indicate whether or not to emit `path_in_schema`. + pub(crate) fn write_path_in_schema(&self) -> bool { + self.write_path_in_schema + } + + /// Write a single byte to the output stream. + fn write_byte(&mut self, b: u8) -> Result<()> { + self.writer.write_all(&[b])?; + Ok(()) + } + + /// Write the given `u64` as a ULEB128 encoded varint. + fn write_vlq(&mut self, val: u64) -> Result<()> { + let mut v = val; + while v > 0x7f { + self.write_byte(v as u8 | 0x80)?; + v >>= 7; + } + self.write_byte(v as u8) + } + + /// Write the given `i64` as a zig-zag encoded varint. + fn write_zig_zag(&mut self, val: i64) -> Result<()> { + let s = (val < 0) as i64; + self.write_vlq((((val ^ -s) << 1) + s) as u64) + } + + /// Used to mark the start of a Thrift struct field of type `field_type`. `last_field_id` + /// is used to compute a delta to the given `field_id` per the compact protocol [spec]. + /// + /// [spec]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#struct-encoding + pub(crate) fn write_field_begin( + &mut self, + field_type: FieldType, + field_id: i16, + last_field_id: i16, + ) -> Result<()> { + let delta = field_id.wrapping_sub(last_field_id); + if delta > 0 && delta <= 0xf { + self.write_byte((delta as u8) << 4 | field_type as u8) + } else { + self.write_byte(field_type as u8)?; + self.write_i16(field_id) + } + } + + /// Used to indicate the start of a list of `element_type` elements. + pub(crate) fn write_list_begin(&mut self, element_type: ElementType, len: usize) -> Result<()> { + if len < 15 { + self.write_byte((len as u8) << 4 | element_type as u8) + } else { + self.write_byte(0xf0u8 | element_type as u8)?; + self.write_vlq(len as _) + } + } + + /// Used to mark the end of a struct. This must be called after all fields of the struct have + /// been written. + pub(crate) fn write_struct_end(&mut self) -> Result<()> { + self.write_byte(0) + } + + /// Serialize a slice of `u8`s. This will encode a length, and then write the bytes without + /// further encoding. + pub(crate) fn write_bytes(&mut self, val: &[u8]) -> Result<()> { + self.write_vlq(val.len() as u64)?; + self.writer.write_all(val)?; + Ok(()) + } + + /// Short-cut method used to encode structs that have no fields (often used in Thrift unions). + /// This simply encodes the field id and then immediately writes the end-of-struct marker. + pub(crate) fn write_empty_struct(&mut self, field_id: i16, last_field_id: i16) -> Result { + self.write_field_begin(FieldType::Struct, field_id, last_field_id)?; + self.write_struct_end()?; + Ok(last_field_id) + } + + /// Write a boolean value. + pub(crate) fn write_bool(&mut self, val: bool) -> Result<()> { + match val { + true => self.write_byte(1), + false => self.write_byte(2), + } + } + + /// Write a zig-zag encoded `i8` value. + pub(crate) fn write_i8(&mut self, val: i8) -> Result<()> { + self.write_byte(val as u8) + } + + /// Write a zig-zag encoded `i16` value. + pub(crate) fn write_i16(&mut self, val: i16) -> Result<()> { + self.write_zig_zag(val as _) + } + + /// Write a zig-zag encoded `i32` value. + pub(crate) fn write_i32(&mut self, val: i32) -> Result<()> { + self.write_zig_zag(val as _) + } + + /// Write a zig-zag encoded `i64` value. + pub(crate) fn write_i64(&mut self, val: i64) -> Result<()> { + self.write_zig_zag(val as _) + } + + /// Write a double value. + pub(crate) fn write_double(&mut self, val: f64) -> Result<()> { + self.writer.write_all(&val.to_le_bytes())?; + Ok(()) + } +} + +/// Trait implemented by objects that are to be serialized to a Thrift [compact output] protocol +/// stream. Implementations are also provided for primitive Thrift types. +/// +/// [compact output]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md +pub(crate) trait WriteThrift { + /// The [`ElementType`] to use when a list of this object is written. + const ELEMENT_TYPE: ElementType; + + /// Serialize this object to the given `writer`. + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()>; +} + +/// Implementation for a vector of thrift serializable objects that implement [`WriteThrift`]. +/// This will write the necessary list header and then serialize the elements one-at-a-time. +impl WriteThrift for Vec +where + T: WriteThrift, +{ + const ELEMENT_TYPE: ElementType = ElementType::List; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + writer.write_list_begin(T::ELEMENT_TYPE, self.len())?; + for item in self { + item.write_thrift(writer)?; + } + Ok(()) + } +} + +impl WriteThrift for bool { + const ELEMENT_TYPE: ElementType = ElementType::Bool; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + writer.write_bool(*self) + } +} + +impl WriteThrift for i8 { + const ELEMENT_TYPE: ElementType = ElementType::Byte; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + writer.write_i8(*self) + } +} + +impl WriteThrift for i16 { + const ELEMENT_TYPE: ElementType = ElementType::I16; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + writer.write_i16(*self) + } +} + +impl WriteThrift for i32 { + const ELEMENT_TYPE: ElementType = ElementType::I32; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + writer.write_i32(*self) + } +} + +impl WriteThrift for i64 { + const ELEMENT_TYPE: ElementType = ElementType::I64; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + writer.write_i64(*self) + } +} + +impl WriteThrift for OrderedF64 { + const ELEMENT_TYPE: ElementType = ElementType::Double; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + writer.write_double(self.0) + } +} + +impl WriteThrift for f64 { + const ELEMENT_TYPE: ElementType = ElementType::Double; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + writer.write_double(*self) + } +} + +impl WriteThrift for &[u8] { + const ELEMENT_TYPE: ElementType = ElementType::Binary; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + writer.write_bytes(self) + } +} + +impl WriteThrift for &str { + const ELEMENT_TYPE: ElementType = ElementType::Binary; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + writer.write_bytes(self.as_bytes()) + } +} + +impl WriteThrift for String { + const ELEMENT_TYPE: ElementType = ElementType::Binary; + + fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { + writer.write_bytes(self.as_bytes()) + } +} + +/// Trait implemented by objects that are fields of Thrift structs. +/// +/// For example, given the Thrift struct definition +/// ```ignore +/// struct MyStruct { +/// 1: required i32 field1 +/// 2: optional bool field2 +/// 3: optional OtherStruct field3 +/// } +/// ``` +/// +/// which becomes in Rust +/// ```no_run +/// # struct OtherStruct {} +/// struct MyStruct { +/// field1: i32, +/// field2: Option, +/// field3: Option, +/// } +/// ``` +/// the impl of `WriteThrift` for `MyStruct` will use the `WriteThriftField` impls for `i32`, +/// `bool`, and `OtherStruct`. +/// +/// ```ignore +/// impl WriteThrift for MyStruct { +/// fn write_thrift(&self, writer: &mut ThriftCompactOutputProtocol) -> Result<()> { +/// let mut last_field_id = 0i16; +/// last_field_id = self.field1.write_thrift_field(writer, 1, last_field_id)?; +/// if self.field2.is_some() { +/// // if field2 is `None` then this assignment won't happen and last_field_id will remain +/// // `1` when writing `field3` +/// last_field_id = self.field2.write_thrift_field(writer, 2, last_field_id)?; +/// } +/// if self.field3.is_some() { +/// // no need to assign last_field_id since this is the final field. +/// self.field3.write_thrift_field(writer, 3, last_field_id)?; +/// } +/// writer.write_struct_end() +/// } +/// } +/// ``` +/// +pub(crate) trait WriteThriftField { + /// Used to write struct fields (which may be primitive or IDL defined types). This will + /// write the field marker for the given `field_id`, using `last_field_id` to compute the + /// field delta used by the Thrift [compact protocol]. On success this will return `field_id` + /// to be used in chaining. + /// + /// [compact protocol]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#struct-encoding + fn write_thrift_field( + &self, + writer: &mut ThriftCompactOutputProtocol, + field_id: i16, + last_field_id: i16, + ) -> Result; +} + +// bool struct fields are written differently to bool values +impl WriteThriftField for bool { + fn write_thrift_field( + &self, + writer: &mut ThriftCompactOutputProtocol, + field_id: i16, + last_field_id: i16, + ) -> Result { + // boolean only writes the field header + match *self { + true => writer.write_field_begin(FieldType::BooleanTrue, field_id, last_field_id)?, + false => writer.write_field_begin(FieldType::BooleanFalse, field_id, last_field_id)?, + } + Ok(field_id) + } +} + +write_thrift_field!(i8, FieldType::Byte); +write_thrift_field!(i16, FieldType::I16); +write_thrift_field!(i32, FieldType::I32); +write_thrift_field!(i64, FieldType::I64); +write_thrift_field!(OrderedF64, FieldType::Double); +write_thrift_field!(f64, FieldType::Double); +write_thrift_field!(String, FieldType::Binary); + +impl WriteThriftField for &[u8] { + fn write_thrift_field( + &self, + writer: &mut ThriftCompactOutputProtocol, + field_id: i16, + last_field_id: i16, + ) -> Result { + writer.write_field_begin(FieldType::Binary, field_id, last_field_id)?; + writer.write_bytes(self)?; + Ok(field_id) + } +} + +impl WriteThriftField for &str { + fn write_thrift_field( + &self, + writer: &mut ThriftCompactOutputProtocol, + field_id: i16, + last_field_id: i16, + ) -> Result { + writer.write_field_begin(FieldType::Binary, field_id, last_field_id)?; + writer.write_bytes(self.as_bytes())?; + Ok(field_id) + } +} + +impl WriteThriftField for Vec +where + T: WriteThrift, +{ + fn write_thrift_field( + &self, + writer: &mut ThriftCompactOutputProtocol, + field_id: i16, + last_field_id: i16, + ) -> Result { + writer.write_field_begin(FieldType::List, field_id, last_field_id)?; + self.write_thrift(writer)?; + Ok(field_id) + } +} + +#[cfg(test)] +pub(crate) mod tests { + use crate::basic::{TimeUnit, Type}; + + use super::*; + use std::fmt::Debug; + + pub(crate) fn test_roundtrip(val: T) + where + T: for<'a> ReadThrift<'a, ThriftSliceInputProtocol<'a>> + WriteThrift + PartialEq + Debug, + { + let mut buf = Vec::::new(); + { + let mut writer = ThriftCompactOutputProtocol::new(&mut buf); + val.write_thrift(&mut writer).unwrap(); + } + + let mut prot = ThriftSliceInputProtocol::new(&buf); + let read_val = T::read_thrift(&mut prot).unwrap(); + assert_eq!(val, read_val); + } + + #[test] + fn test_enum_roundtrip() { + test_roundtrip(Type::BOOLEAN); + test_roundtrip(Type::INT32); + test_roundtrip(Type::INT64); + test_roundtrip(Type::INT96); + test_roundtrip(Type::FLOAT); + test_roundtrip(Type::DOUBLE); + test_roundtrip(Type::BYTE_ARRAY); + test_roundtrip(Type::FIXED_LEN_BYTE_ARRAY); + } + + #[test] + fn test_union_all_empty_roundtrip() { + test_roundtrip(TimeUnit::MILLIS); + test_roundtrip(TimeUnit::MICROS); + test_roundtrip(TimeUnit::NANOS); + } + + #[test] + fn test_decode_empty_list() { + let data = vec![0u8; 1]; + let mut prot = ThriftSliceInputProtocol::new(&data); + let header = prot.read_list_begin().expect("error reading list header"); + assert_eq!(header.size, 0); + assert_eq!(header.element_type, ElementType::Byte); + } + + /// A Thrift list header whose `size` varint decodes above `i32::MAX` + /// must be rejected at the protocol layer rather than wrapping into a + /// negative `i32` and being smuggled into downstream allocation code. + #[test] + fn test_read_list_begin_size_above_i32_max_returns_err() { + // List header: element_type=8 (Binary), 0xF=follow-up varint. + // Varint 80 80 80 80 08 decodes to 0x8000_0000 = i32::MAX + 1. + let mut data: Vec = vec![0xF8]; + data.extend_from_slice(&[0x80, 0x80, 0x80, 0x80, 0x08]); + let mut prot = ThriftSliceInputProtocol::new(&data); + let result = prot.read_list_begin(); + assert!(result.is_err(), "expected error, got {result:?}"); + } + + #[test] + fn test_read_list_wrong_type() { + // list header: 4 elements of `Boolean` + let data = [0x42, 0x01]; + let mut prot = ThriftSliceInputProtocol::new(&data); + // try to read as list + let result = read_thrift_vec::(&mut prot); + println!("{result:?}"); + assert!( + result + .unwrap_err() + .to_string() + .contains("Expected list element type of I32 but got Bool") + ); + } +} diff --git a/vendor/parquet/src/record/api.rs b/vendor/parquet/src/record/api.rs new file mode 100644 index 0000000..104a82e --- /dev/null +++ b/vendor/parquet/src/record/api.rs @@ -0,0 +1,2223 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Contains Row enum that is used to represent record in Rust. + +use std::fmt; + +use chrono::{TimeZone, Utc}; +use half::f16; +use num_bigint::{BigInt, Sign}; +use num_traits::Float; + +use crate::basic::{ConvertedType, LogicalType, Type as PhysicalType}; +use crate::data_type::{ByteArray, Decimal, Int96}; +use crate::errors::{ParquetError, Result}; +use crate::schema::types::ColumnDescPtr; + +#[cfg(any(feature = "json", test))] +use serde_json::Value; + +/// Macro as a shortcut to generate 'not yet implemented' panic error. +macro_rules! nyi { + ($column_descr:ident, $value:ident) => {{ + unimplemented!( + "Conversion for physical type {}, converted type {}, value {:?}", + $column_descr.physical_type(), + $column_descr.converted_type(), + $value + ); + }}; +} + +/// `Row` represents a nested Parquet record. +#[derive(Clone, Debug, PartialEq)] +pub struct Row { + fields: Vec<(String, Field)>, +} + +#[allow(clippy::len_without_is_empty)] +impl Row { + /// Constructs a `Row` from the list of `fields` and returns it. + pub fn new(fields: Vec<(String, Field)>) -> Row { + Row { fields } + } + + /// Get the number of fields in this row. + pub fn len(&self) -> usize { + self.fields.len() + } + + /// Move columns data out of the row. Useful to avoid internal data cloning. + /// + /// # Example + /// + /// ```no_run + /// use std::fs::File; + /// use parquet::record::Row; + /// use parquet::file::reader::{FileReader, SerializedFileReader}; + /// + /// let file = File::open("/path/to/file").unwrap(); + /// let reader = SerializedFileReader::new(file).unwrap(); + /// let row: Row = reader.get_row_iter(None).unwrap().next().unwrap().unwrap(); + /// let columns = row.into_columns(); + /// println!("row columns: {:?}", columns); + /// + /// ``` + pub fn into_columns(self) -> Vec<(String, Field)> { + self.fields + } + + /// Get an iterator to go through all columns in the row. + /// + /// # Example + /// + /// ```no_run + /// use std::fs::File; + /// use parquet::record::Row; + /// use parquet::file::reader::{FileReader, SerializedFileReader}; + /// + /// let file = File::open("/path/to/file").unwrap(); + /// let reader = SerializedFileReader::new(file).unwrap(); + /// let row: Row = reader.get_row_iter(None).unwrap().next().unwrap().unwrap(); + /// for (idx, (name, field)) in row.get_column_iter().enumerate() { + /// println!("column index: {}, column name: {}, column value: {}", idx, name, field); + /// } + /// ``` + pub fn get_column_iter(&self) -> RowColumnIter<'_> { + RowColumnIter { + fields: &self.fields, + curr: 0, + count: self.fields.len(), + } + } + + /// Converts the row into a JSON object. + #[cfg(any(feature = "json", test))] + pub fn to_json_value(&self) -> Value { + Value::Object( + self.fields + .iter() + .map(|(key, field)| (key.to_owned(), field.to_json_value())) + .collect(), + ) + } +} + +/// `RowColumnIter` represents an iterator over column names and values in a Row. +pub struct RowColumnIter<'a> { + fields: &'a Vec<(String, Field)>, + curr: usize, + count: usize, +} + +impl<'a> Iterator for RowColumnIter<'a> { + type Item = (&'a String, &'a Field); + + fn next(&mut self) -> Option { + let idx = self.curr; + if idx >= self.count { + return None; + } + self.curr += 1; + Some((&self.fields[idx].0, &self.fields[idx].1)) + } +} + +/// Trait for type-safe convenient access to fields within a Row. +pub trait RowAccessor { + /// Check if the field at the index is null. + fn is_null(&self, i: usize) -> Result; + /// Try to get a boolean value at the given index. + fn get_bool(&self, i: usize) -> Result; + /// Try to get a byte value at the given index. + fn get_byte(&self, i: usize) -> Result; + /// Try to get a short value at the given index. + fn get_short(&self, i: usize) -> Result; + /// Try to get a int value at the given index. + fn get_int(&self, i: usize) -> Result; + /// Try to get a long value at the given index. + fn get_long(&self, i: usize) -> Result; + /// Try to get a ubyte value at the given index. + fn get_ubyte(&self, i: usize) -> Result; + /// Try to get a ushort value at the given index. + fn get_ushort(&self, i: usize) -> Result; + /// Try to get a uint value at the given index. + fn get_uint(&self, i: usize) -> Result; + /// Try to get a ulong value at the given index. + fn get_ulong(&self, i: usize) -> Result; + /// Try to get a float16 value at the given index. + fn get_float16(&self, i: usize) -> Result; + /// Try to get a float value at the given index. + fn get_float(&self, i: usize) -> Result; + /// Try to get a double value at the given index. + fn get_double(&self, i: usize) -> Result; + /// Try to get a date value at the given index. + fn get_timestamp_millis(&self, i: usize) -> Result; + /// Try to get a date value at the given index. + fn get_timestamp_micros(&self, i: usize) -> Result; + /// Try to get a decimal value at the given index. + fn get_decimal(&self, i: usize) -> Result<&Decimal>; + /// Try to get a string value at the given index. + fn get_string(&self, i: usize) -> Result<&String>; + /// Try to get a bytes value at the given index. + fn get_bytes(&self, i: usize) -> Result<&ByteArray>; + /// Try to get a group value at the given index. + fn get_group(&self, i: usize) -> Result<&Row>; + /// Try to get a list value at the given index. + fn get_list(&self, i: usize) -> Result<&List>; + /// Try to get a map value at the given index. + fn get_map(&self, i: usize) -> Result<&Map>; +} + +/// Trait for formatting fields within a Row. +/// +/// # Examples +/// +/// ``` +/// use std::fs::File; +/// use std::path::Path; +/// use parquet::record::Row; +/// use parquet::record::RowFormatter; +/// use parquet::file::reader::{FileReader, SerializedFileReader}; +/// +/// if let Ok(file) = File::open(&Path::new("test.parquet")) { +/// let reader = SerializedFileReader::new(file).unwrap(); +/// let row = reader.get_row_iter(None).unwrap().next().unwrap().unwrap(); +/// println!("column 0: {}, column 1: {}", row.fmt(0), row.fmt(1)); +/// } +/// ``` +/// +pub trait RowFormatter { + /// The method to format a field at the given index. + fn fmt(&self, i: usize) -> &dyn fmt::Display; +} + +/// Macro to generate type-safe get_xxx methods for primitive types, +/// e.g. `get_bool`, `get_short`. +macro_rules! row_primitive_accessor { + ($METHOD:ident, $VARIANT:ident, $TY:ty) => { + fn $METHOD(&self, i: usize) -> Result<$TY> { + match self.fields.get(i) { + Some((_, Field::$VARIANT(v))) => Ok(*v), + None => Err(ParquetError::IndexOutOfBound(i, self.fields.len())), + _ => Err(general_err!( + "Cannot access {} as {}", + self.fields[i].1.get_type_name(), // Safe access as None is + stringify!($VARIANT) + )), + } + } + }; +} + +/// Macro to generate type-safe get_xxx methods for reference types, +/// e.g. `get_list`, `get_map`. +macro_rules! row_complex_accessor { + ($METHOD:ident, $VARIANT:ident, $TY:ty) => { + fn $METHOD(&self, i: usize) -> Result<&$TY> { + match self.fields.get(i) { + Some((_, Field::$VARIANT(v))) => Ok(v), + None => Err(ParquetError::IndexOutOfBound(i, self.fields.len())), + _ => Err(general_err!( + "Cannot access {} as {}", + self.fields[i].1.get_type_name(), // Safe access as None is + // just checked. + stringify!($VARIANT) + )), + } + } + }; +} + +impl RowFormatter for Row { + /// Get Display reference for a given field. + fn fmt(&self, i: usize) -> &dyn fmt::Display { + if let Some((_, v)) = self.fields.get(i) { + v + } else { + &"" + } + } +} + +impl RowAccessor for Row { + fn is_null(&self, i: usize) -> Result { + match self.fields.get(i) { + Some((_, Field::Null)) => Ok(true), + None => Err(ParquetError::IndexOutOfBound(i, self.len())), + _ => Ok(false), + } + } + + row_primitive_accessor!(get_bool, Bool, bool); + + row_primitive_accessor!(get_byte, Byte, i8); + + row_primitive_accessor!(get_short, Short, i16); + + row_primitive_accessor!(get_int, Int, i32); + + row_primitive_accessor!(get_long, Long, i64); + + row_primitive_accessor!(get_ubyte, UByte, u8); + + row_primitive_accessor!(get_ushort, UShort, u16); + + row_primitive_accessor!(get_uint, UInt, u32); + + row_primitive_accessor!(get_ulong, ULong, u64); + + row_primitive_accessor!(get_float16, Float16, f16); + + row_primitive_accessor!(get_float, Float, f32); + + row_primitive_accessor!(get_double, Double, f64); + + row_primitive_accessor!(get_timestamp_millis, TimestampMillis, i64); + + row_primitive_accessor!(get_timestamp_micros, TimestampMicros, i64); + + row_complex_accessor!(get_decimal, Decimal, Decimal); + + row_complex_accessor!(get_string, Str, String); + + row_complex_accessor!(get_bytes, Bytes, ByteArray); + + row_complex_accessor!(get_group, Group, Row); + + row_complex_accessor!(get_list, ListInternal, List); + + row_complex_accessor!(get_map, MapInternal, Map); +} + +impl fmt::Display for Row { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{{")?; + for (i, (key, value)) in self.fields.iter().enumerate() { + key.fmt(f)?; + write!(f, ": ")?; + value.fmt(f)?; + if i < self.fields.len() - 1 { + write!(f, ", ")?; + } + } + write!(f, "}}") + } +} + +/// `List` represents a list which contains an array of elements. +#[derive(Clone, Debug, PartialEq)] +pub struct List { + elements: Vec, +} + +#[allow(clippy::len_without_is_empty)] +impl List { + /// Get the number of fields in this row + pub fn len(&self) -> usize { + self.elements.len() + } + + /// Get the reference to the elements in this list + pub fn elements(&self) -> &[Field] { + self.elements.as_slice() + } +} + +/// Constructs a `List` from the list of `fields` and returns it. +#[inline] +pub fn make_list(elements: Vec) -> List { + List { elements } +} + +/// Trait for type-safe access of an index for a `List`. +/// Note that the get_XXX methods do not do bound checking. +pub trait ListAccessor { + /// Try getting a `boolean` value at the given index. + fn get_bool(&self, i: usize) -> Result; + /// Try getting a `byte` value at the given index. + fn get_byte(&self, i: usize) -> Result; + /// Try getting an `i16` value at the given index. + fn get_short(&self, i: usize) -> Result; + /// Try getting an `i32` value at the given index. + fn get_int(&self, i: usize) -> Result; + /// Try getting an `i64` value at the given index. + fn get_long(&self, i: usize) -> Result; + /// Try getting a `u8` value at the given index. + fn get_ubyte(&self, i: usize) -> Result; + /// Try getting a `u16` value at the given index. + fn get_ushort(&self, i: usize) -> Result; + /// Try getting a `u32` value at the given index. + fn get_uint(&self, i: usize) -> Result; + /// Try getting a `u64` value at the given index. + fn get_ulong(&self, i: usize) -> Result; + /// Try getting a `f16` value at the given index. + fn get_float16(&self, i: usize) -> Result; + /// Try getting a `f32` value at the given index. + fn get_float(&self, i: usize) -> Result; + /// Try getting a `f64` value at the given index. + fn get_double(&self, i: usize) -> Result; + /// Try getting a `timestamp` as milliseconds value + /// encoded as `i64` at the given index. + fn get_timestamp_millis(&self, i: usize) -> Result; + /// Try getting a `timestamp` as microseconds value + /// encoded as `i64` at the given index. + fn get_timestamp_micros(&self, i: usize) -> Result; + /// Try getting a `decimal` value at the given index. + fn get_decimal(&self, i: usize) -> Result<&Decimal>; + /// Try getting a `string` value at the given index. + fn get_string(&self, i: usize) -> Result<&String>; + /// Try getting a `bytes` value at the given index. + fn get_bytes(&self, i: usize) -> Result<&ByteArray>; + /// Try getting a `group` value at the given index. + fn get_group(&self, i: usize) -> Result<&Row>; + /// Try getting a `list` value at the given index. + fn get_list(&self, i: usize) -> Result<&List>; + /// Try getting a `map` value at the given index. + fn get_map(&self, i: usize) -> Result<&Map>; +} + +/// Macro to generate type-safe get_xxx methods for primitive types, +/// e.g. get_bool, get_short +macro_rules! list_primitive_accessor { + ($METHOD:ident, $VARIANT:ident, $TY:ty) => { + fn $METHOD(&self, i: usize) -> Result<$TY> { + match self.elements[i] { + Field::$VARIANT(v) => Ok(v), + _ => Err(general_err!( + "Cannot access {} as {}", + self.elements[i].get_type_name(), + stringify!($VARIANT) + )), + } + } + }; +} + +/// Macro to generate type-safe get_xxx methods for reference types +/// e.g. get_list, get_map +macro_rules! list_complex_accessor { + ($METHOD:ident, $VARIANT:ident, $TY:ty) => { + fn $METHOD(&self, i: usize) -> Result<&$TY> { + match &self.elements[i] { + Field::$VARIANT(v) => Ok(&v), + _ => Err(general_err!( + "Cannot access {} as {}", + self.elements[i].get_type_name(), + stringify!($VARIANT) + )), + } + } + }; +} + +impl ListAccessor for List { + list_primitive_accessor!(get_bool, Bool, bool); + + list_primitive_accessor!(get_byte, Byte, i8); + + list_primitive_accessor!(get_short, Short, i16); + + list_primitive_accessor!(get_int, Int, i32); + + list_primitive_accessor!(get_long, Long, i64); + + list_primitive_accessor!(get_ubyte, UByte, u8); + + list_primitive_accessor!(get_ushort, UShort, u16); + + list_primitive_accessor!(get_uint, UInt, u32); + + list_primitive_accessor!(get_ulong, ULong, u64); + + list_primitive_accessor!(get_float16, Float16, f16); + + list_primitive_accessor!(get_float, Float, f32); + + list_primitive_accessor!(get_double, Double, f64); + + list_primitive_accessor!(get_timestamp_millis, TimestampMillis, i64); + + list_primitive_accessor!(get_timestamp_micros, TimestampMicros, i64); + + list_complex_accessor!(get_decimal, Decimal, Decimal); + + list_complex_accessor!(get_string, Str, String); + + list_complex_accessor!(get_bytes, Bytes, ByteArray); + + list_complex_accessor!(get_group, Group, Row); + + list_complex_accessor!(get_list, ListInternal, List); + + list_complex_accessor!(get_map, MapInternal, Map); +} + +/// `Map` represents a map which contains a list of key->value pairs. +#[derive(Clone, Debug, PartialEq)] +pub struct Map { + entries: Vec<(Field, Field)>, +} + +#[allow(clippy::len_without_is_empty)] +impl Map { + /// Get the number of fields in this row + pub fn len(&self) -> usize { + self.entries.len() + } + + /// Get the reference to the key-value pairs in this map + pub fn entries(&self) -> &[(Field, Field)] { + self.entries.as_slice() + } +} + +/// Constructs a `Map` from the list of `entries` and returns it. +#[inline] +pub fn make_map(entries: Vec<(Field, Field)>) -> Map { + Map { entries } +} + +/// Trait for type-safe access of an index for a `Map` +pub trait MapAccessor { + /// Get the keys of the map. + fn get_keys<'a>(&'a self) -> Box; + /// Get the values of the map. + fn get_values<'a>(&'a self) -> Box; +} + +struct MapList<'a> { + elements: Vec<&'a Field>, +} + +/// Macro to generate type-safe get_xxx methods for primitive types, +/// e.g. get_bool, get_short +macro_rules! map_list_primitive_accessor { + ($METHOD:ident, $VARIANT:ident, $TY:ty) => { + fn $METHOD(&self, i: usize) -> Result<$TY> { + match self.elements[i] { + Field::$VARIANT(v) => Ok(*v), + _ => Err(general_err!( + "Cannot access {} as {}", + self.elements[i].get_type_name(), + stringify!($VARIANT) + )), + } + } + }; +} + +impl ListAccessor for MapList<'_> { + map_list_primitive_accessor!(get_bool, Bool, bool); + + map_list_primitive_accessor!(get_byte, Byte, i8); + + map_list_primitive_accessor!(get_short, Short, i16); + + map_list_primitive_accessor!(get_int, Int, i32); + + map_list_primitive_accessor!(get_long, Long, i64); + + map_list_primitive_accessor!(get_ubyte, UByte, u8); + + map_list_primitive_accessor!(get_ushort, UShort, u16); + + map_list_primitive_accessor!(get_uint, UInt, u32); + + map_list_primitive_accessor!(get_ulong, ULong, u64); + + map_list_primitive_accessor!(get_float16, Float16, f16); + + map_list_primitive_accessor!(get_float, Float, f32); + + map_list_primitive_accessor!(get_double, Double, f64); + + map_list_primitive_accessor!(get_timestamp_millis, TimestampMillis, i64); + + map_list_primitive_accessor!(get_timestamp_micros, TimestampMicros, i64); + + list_complex_accessor!(get_decimal, Decimal, Decimal); + + list_complex_accessor!(get_string, Str, String); + + list_complex_accessor!(get_bytes, Bytes, ByteArray); + + list_complex_accessor!(get_group, Group, Row); + + list_complex_accessor!(get_list, ListInternal, List); + + list_complex_accessor!(get_map, MapInternal, Map); +} + +impl MapAccessor for Map { + fn get_keys<'a>(&'a self) -> Box { + let map_list = MapList { + elements: self.entries.iter().map(|v| &v.0).collect(), + }; + Box::new(map_list) + } + + fn get_values<'a>(&'a self) -> Box { + let map_list = MapList { + elements: self.entries.iter().map(|v| &v.1).collect(), + }; + Box::new(map_list) + } +} + +/// API to represent a single field in a `Row`. +#[derive(Clone, Debug, PartialEq)] +pub enum Field { + // Primitive types + /// Null value. + Null, + /// Boolean value (`true`, `false`). + Bool(bool), + /// Signed integer INT_8. + Byte(i8), + /// Signed integer INT_16. + Short(i16), + /// Signed integer INT_32. + Int(i32), + /// Signed integer INT_64. + Long(i64), + /// Unsigned integer UINT_8. + UByte(u8), + /// Unsigned integer UINT_16. + UShort(u16), + /// Unsigned integer UINT_32. + UInt(u32), + /// Unsigned integer UINT_64. + ULong(u64), + /// IEEE 16-bit floating point value. + Float16(f16), + /// IEEE 32-bit floating point value. + Float(f32), + /// IEEE 64-bit floating point value. + Double(f64), + /// Decimal value. + Decimal(Decimal), + /// UTF-8 encoded character string. + Str(String), + /// General binary value. + Bytes(ByteArray), + /// Date without a time of day, stores the number of days from the + /// Unix epoch, 1 January 1970. + Date(i32), + + /// The total number of milliseconds since midnight. + TimeMillis(i32), + /// The total number of microseconds since midnight. + TimeMicros(i64), + + /// Milliseconds from the Unix epoch, 1 January 1970. + TimestampMillis(i64), + /// Microseconds from the Unix epoch, 1 January 1970. + TimestampMicros(i64), + + // ---------------------------------------------------------------------- + // Complex types + /// Struct, child elements are tuples of field-value pairs. + Group(Row), + /// List of elements. + ListInternal(List), + /// List of key-value pairs. + MapInternal(Map), +} + +impl Field { + /// Get the type name. + fn get_type_name(&self) -> &'static str { + match *self { + Field::Null => "Null", + Field::Bool(_) => "Bool", + Field::Byte(_) => "Byte", + Field::Short(_) => "Short", + Field::Int(_) => "Int", + Field::Long(_) => "Long", + Field::UByte(_) => "UByte", + Field::UShort(_) => "UShort", + Field::UInt(_) => "UInt", + Field::ULong(_) => "ULong", + Field::Float16(_) => "Float16", + Field::Float(_) => "Float", + Field::Double(_) => "Double", + Field::Decimal(_) => "Decimal", + Field::Date(_) => "Date", + Field::Str(_) => "Str", + Field::Bytes(_) => "Bytes", + Field::TimeMillis(_) => "TimeMillis", + Field::TimeMicros(_) => "TimeMicros", + Field::TimestampMillis(_) => "TimestampMillis", + Field::TimestampMicros(_) => "TimestampMicros", + Field::Group(_) => "Group", + Field::ListInternal(_) => "ListInternal", + Field::MapInternal(_) => "MapInternal", + } + } + + /// Determines if this Row represents a primitive value. + pub fn is_primitive(&self) -> bool { + !matches!( + *self, + Field::Group(_) | Field::ListInternal(_) | Field::MapInternal(_) + ) + } + + /// Converts Parquet BOOLEAN type with logical type into `bool` value. + #[inline] + pub fn convert_bool(_descr: &ColumnDescPtr, value: bool) -> Self { + Field::Bool(value) + } + + /// Converts Parquet INT32 type with converted type into `i32` value. + #[inline] + pub fn convert_int32(descr: &ColumnDescPtr, value: i32) -> Self { + match descr.converted_type() { + ConvertedType::INT_8 => Field::Byte(value as i8), + ConvertedType::INT_16 => Field::Short(value as i16), + ConvertedType::INT_32 | ConvertedType::NONE => Field::Int(value), + ConvertedType::UINT_8 => Field::UByte(value as u8), + ConvertedType::UINT_16 => Field::UShort(value as u16), + ConvertedType::UINT_32 => Field::UInt(value as u32), + ConvertedType::DATE => Field::Date(value), + ConvertedType::TIME_MILLIS => Field::TimeMillis(value), + ConvertedType::DECIMAL => Field::Decimal(Decimal::from_i32( + value, + descr.type_precision(), + descr.type_scale(), + )), + _ => nyi!(descr, value), + } + } + + /// Converts Parquet INT64 type with converted type into `i64` value. + #[inline] + pub fn convert_int64(descr: &ColumnDescPtr, value: i64) -> Self { + match descr.converted_type() { + ConvertedType::INT_64 | ConvertedType::NONE => Field::Long(value), + ConvertedType::UINT_64 => Field::ULong(value as u64), + ConvertedType::TIME_MICROS => Field::TimeMicros(value), + ConvertedType::TIMESTAMP_MILLIS => Field::TimestampMillis(value), + ConvertedType::TIMESTAMP_MICROS => Field::TimestampMicros(value), + ConvertedType::DECIMAL => Field::Decimal(Decimal::from_i64( + value, + descr.type_precision(), + descr.type_scale(), + )), + _ => nyi!(descr, value), + } + } + + /// Converts Parquet INT96 (nanosecond timestamps) type and logical type into + /// `Timestamp` value. + #[inline] + pub fn convert_int96(_descr: &ColumnDescPtr, value: Int96) -> Self { + Field::TimestampMillis(value.to_millis()) + } + + /// Converts Parquet FLOAT type with logical type into `f32` value. + #[inline] + pub fn convert_float(_descr: &ColumnDescPtr, value: f32) -> Self { + Field::Float(value) + } + + /// Converts Parquet DOUBLE type with converted type into `f64` value. + #[inline] + pub fn convert_double(_descr: &ColumnDescPtr, value: f64) -> Self { + Field::Double(value) + } + + /// Converts Parquet BYTE_ARRAY type with converted type into a UTF8 + /// string, decimal, float16, or an array of bytes. + #[inline] + pub fn convert_byte_array(descr: &ColumnDescPtr, value: ByteArray) -> Result { + let field = match descr.physical_type() { + PhysicalType::BYTE_ARRAY => match descr.converted_type() { + ConvertedType::UTF8 | ConvertedType::ENUM | ConvertedType::JSON => { + let value = String::from_utf8(value.data().to_vec()).map_err(|e| { + general_err!( + "Error reading BYTE_ARRAY as String. Bytes: {:?} Error: {:?}", + value.data(), + e + ) + })?; + Field::Str(value) + } + ConvertedType::BSON | ConvertedType::NONE => Field::Bytes(value), + ConvertedType::DECIMAL => Field::Decimal(Decimal::from_bytes( + value, + descr.type_precision(), + descr.type_scale(), + )), + _ => nyi!(descr, value), + }, + PhysicalType::FIXED_LEN_BYTE_ARRAY => match descr.converted_type() { + ConvertedType::DECIMAL => Field::Decimal(Decimal::from_bytes( + value, + descr.type_precision(), + descr.type_scale(), + )), + ConvertedType::NONE if descr.logical_type_ref() == Some(&LogicalType::Float16) => { + if value.len() != 2 { + return Err(general_err!( + "Error reading FIXED_LEN_BYTE_ARRAY as FLOAT16. Length must be 2, got {}", + value.len() + )); + } + let bytes = [value.data()[0], value.data()[1]]; + Field::Float16(f16::from_le_bytes(bytes)) + } + ConvertedType::NONE => Field::Bytes(value), + _ => nyi!(descr, value), + }, + _ => nyi!(descr, value), + }; + Ok(field) + } + + /// Converts the Parquet field into a JSON [`Value`]. + #[cfg(any(feature = "json", test))] + pub fn to_json_value(&self) -> Value { + use base64::Engine; + use base64::prelude::BASE64_STANDARD; + + match &self { + Field::Null => Value::Null, + Field::Bool(b) => Value::Bool(*b), + Field::Byte(n) => Value::Number(serde_json::Number::from(*n)), + Field::Short(n) => Value::Number(serde_json::Number::from(*n)), + Field::Int(n) => Value::Number(serde_json::Number::from(*n)), + Field::Long(n) => Value::Number(serde_json::Number::from(*n)), + Field::UByte(n) => Value::Number(serde_json::Number::from(*n)), + Field::UShort(n) => Value::Number(serde_json::Number::from(*n)), + Field::UInt(n) => Value::Number(serde_json::Number::from(*n)), + Field::ULong(n) => Value::Number(serde_json::Number::from(*n)), + Field::Float16(n) => serde_json::Number::from_f64(f64::from(*n)) + .map(Value::Number) + .unwrap_or(Value::Null), + Field::Float(n) => serde_json::Number::from_f64(f64::from(*n)) + .map(Value::Number) + .unwrap_or(Value::Null), + Field::Double(n) => serde_json::Number::from_f64(*n) + .map(Value::Number) + .unwrap_or(Value::Null), + Field::Decimal(n) => Value::String(convert_decimal_to_string(n)), + Field::Str(s) => Value::String(s.to_owned()), + Field::Bytes(b) => Value::String(BASE64_STANDARD.encode(b.data())), + Field::Date(d) => Value::String(convert_date_to_string(*d)), + Field::TimeMillis(t) => Value::String(convert_time_millis_to_string(*t)), + Field::TimeMicros(t) => Value::String(convert_time_micros_to_string(*t)), + Field::TimestampMillis(ts) => Value::String(convert_timestamp_millis_to_string(*ts)), + Field::TimestampMicros(ts) => Value::String(convert_timestamp_micros_to_string(*ts)), + Field::Group(row) => row.to_json_value(), + Field::ListInternal(fields) => { + Value::Array(fields.elements.iter().map(|f| f.to_json_value()).collect()) + } + Field::MapInternal(map) => Value::Object( + map.entries + .iter() + .map(|(key_field, value_field)| { + let key_val = key_field.to_json_value(); + let key_str = key_val + .as_str() + .map(|s| s.to_owned()) + .unwrap_or_else(|| key_val.to_string()); + (key_str, value_field.to_json_value()) + }) + .collect(), + ), + } + } +} + +impl fmt::Display for Field { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match *self { + Field::Null => write!(f, "null"), + Field::Bool(value) => write!(f, "{value}"), + Field::Byte(value) => write!(f, "{value}"), + Field::Short(value) => write!(f, "{value}"), + Field::Int(value) => write!(f, "{value}"), + Field::Long(value) => write!(f, "{value}"), + Field::UByte(value) => write!(f, "{value}"), + Field::UShort(value) => write!(f, "{value}"), + Field::UInt(value) => write!(f, "{value}"), + Field::ULong(value) => write!(f, "{value}"), + Field::Float16(value) => { + if !value.is_finite() { + write!(f, "{value}") + } else if value.trunc() == value { + write!(f, "{value}.0") + } else { + write!(f, "{value}") + } + } + Field::Float(value) => { + if !(1e-15..=1e19).contains(&value) { + write!(f, "{value:E}") + } else if value.trunc() == value { + write!(f, "{value}.0") + } else { + write!(f, "{value}") + } + } + Field::Double(value) => { + if !(1e-15..=1e19).contains(&value) { + write!(f, "{value:E}") + } else if value.trunc() == value { + write!(f, "{value}.0") + } else { + write!(f, "{value}") + } + } + Field::Decimal(ref value) => { + write!(f, "{}", convert_decimal_to_string(value)) + } + Field::Str(ref value) => write!(f, "\"{value}\""), + Field::Bytes(ref value) => write!(f, "{:?}", value.data()), + Field::Date(value) => write!(f, "{}", convert_date_to_string(value)), + Field::TimeMillis(value) => { + write!(f, "{}", convert_time_millis_to_string(value)) + } + Field::TimeMicros(value) => { + write!(f, "{}", convert_time_micros_to_string(value)) + } + Field::TimestampMillis(value) => { + write!(f, "{}", convert_timestamp_millis_to_string(value)) + } + Field::TimestampMicros(value) => { + write!(f, "{}", convert_timestamp_micros_to_string(value)) + } + Field::Group(ref fields) => write!(f, "{fields}"), + Field::ListInternal(ref list) => { + let elems = &list.elements; + write!(f, "[")?; + for (i, field) in elems.iter().enumerate() { + field.fmt(f)?; + if i < elems.len() - 1 { + write!(f, ", ")?; + } + } + write!(f, "]") + } + Field::MapInternal(ref map) => { + let entries = &map.entries; + write!(f, "{{")?; + for (i, (key, value)) in entries.iter().enumerate() { + key.fmt(f)?; + write!(f, " -> ")?; + value.fmt(f)?; + if i < entries.len() - 1 { + write!(f, ", ")?; + } + } + write!(f, "}}") + } + } + } +} + +/// Helper method to convert Parquet date into a string. +/// Input `value` is a number of days since the epoch in UTC. +/// Date is displayed in local timezone. +#[inline] +fn convert_date_to_string(value: i32) -> String { + static NUM_SECONDS_IN_DAY: i64 = 60 * 60 * 24; + let dt = Utc + .timestamp_opt(value as i64 * NUM_SECONDS_IN_DAY, 0) + .unwrap(); + format!("{}", dt.format("%Y-%m-%d")) +} + +/// Helper method to convert Parquet timestamp into a string. +/// Input `value` is a number of milliseconds since the epoch in UTC. +/// Datetime is displayed in UTC timezone. +#[inline] +fn convert_timestamp_millis_to_string(value: i64) -> String { + let dt = Utc.timestamp_millis_opt(value).unwrap(); + format!("{}", dt.format("%Y-%m-%d %H:%M:%S%.3f %:z")) +} + +/// Helper method to convert Parquet timestamp into a string. +/// Input `value` is a number of microseconds since the epoch in UTC. +/// Datetime is displayed in UTC timezone. +#[inline] +fn convert_timestamp_micros_to_string(value: i64) -> String { + let dt = Utc.timestamp_micros(value).unwrap(); + format!("{}", dt.format("%Y-%m-%d %H:%M:%S%.6f %:z")) +} + +/// Helper method to convert Parquet time (milliseconds since midnight) into a string. +/// Input `value` is a number of milliseconds since midnight. +/// Time is displayed in HH:MM:SS.sss format. +#[inline] +fn convert_time_millis_to_string(value: i32) -> String { + let total_ms = value as u64; + let hours = total_ms / (60 * 60 * 1000); + let minutes = (total_ms % (60 * 60 * 1000)) / (60 * 1000); + let seconds = (total_ms % (60 * 1000)) / 1000; + let millis = total_ms % 1000; + format!("{hours:02}:{minutes:02}:{seconds:02}.{millis:03}") +} + +/// Helper method to convert Parquet time (microseconds since midnight) into a string. +/// Input `value` is a number of microseconds since midnight. +/// Time is displayed in HH:MM:SS.ssssss format. +#[inline] +fn convert_time_micros_to_string(value: i64) -> String { + let total_us = value as u64; + let hours = total_us / (60 * 60 * 1000 * 1000); + let minutes = (total_us % (60 * 60 * 1000 * 1000)) / (60 * 1000 * 1000); + let seconds = (total_us % (60 * 1000 * 1000)) / (1000 * 1000); + let micros = total_us % (1000 * 1000); + format!("{hours:02}:{minutes:02}:{seconds:02}.{micros:06}") +} + +/// Helper method to convert Parquet decimal into a string. +/// We assert that `scale >= 0` and `precision > scale`, but this will be enforced +/// when constructing Parquet schema. +#[inline] +fn convert_decimal_to_string(decimal: &Decimal) -> String { + assert!(decimal.scale() >= 0 && decimal.precision() > decimal.scale()); + + // Specify as signed bytes to resolve sign as part of conversion. + let num = BigInt::from_signed_bytes_be(decimal.data()); + + // Offset of the first digit in a string. + let negative = i32::from(num.sign() == Sign::Minus); + let mut num_str = num.to_string(); + let mut point = num_str.len() as i32 - decimal.scale() - negative; + + // Convert to string form without scientific notation. + if point <= 0 { + // Zeros need to be prepended to the unscaled value. + while point < 0 { + num_str.insert(negative as usize, '0'); + point += 1; + } + num_str.insert_str(negative as usize, "0."); + } else { + // No zeroes need to be prepended to the unscaled value, simply insert decimal + // point. + num_str.insert((point + negative) as usize, '.'); + } + + num_str +} + +#[cfg(test)] +#[allow(clippy::many_single_char_names)] +mod tests { + use super::*; + + use std::f64::consts::PI; + use std::sync::Arc; + + use crate::schema::types::{ColumnDescriptor, ColumnPath, PrimitiveTypeBuilder}; + + /// Creates test column descriptor based on provided type parameters. + macro_rules! make_column_descr { + ($physical_type:expr, $logical_type:expr) => {{ + let tpe = PrimitiveTypeBuilder::new("col", $physical_type) + .with_converted_type($logical_type) + .build() + .unwrap(); + Arc::new(ColumnDescriptor::new( + Arc::new(tpe), + 0, + 0, + ColumnPath::from("col"), + )) + }}; + ($physical_type:expr, $logical_type:expr, $len:expr, $prec:expr, $scale:expr) => {{ + let tpe = PrimitiveTypeBuilder::new("col", $physical_type) + .with_converted_type($logical_type) + .with_length($len) + .with_precision($prec) + .with_scale($scale) + .build() + .unwrap(); + Arc::new(ColumnDescriptor::new( + Arc::new(tpe), + 0, + 0, + ColumnPath::from("col"), + )) + }}; + } + + #[test] + fn test_row_convert_bool() { + // BOOLEAN value does not depend on logical type + let descr = make_column_descr![PhysicalType::BOOLEAN, ConvertedType::NONE]; + + let row = Field::convert_bool(&descr, true); + assert_eq!(row, Field::Bool(true)); + + let row = Field::convert_bool(&descr, false); + assert_eq!(row, Field::Bool(false)); + } + + #[test] + fn test_row_convert_int32() { + let descr = make_column_descr![PhysicalType::INT32, ConvertedType::INT_8]; + let row = Field::convert_int32(&descr, 111); + assert_eq!(row, Field::Byte(111)); + + let descr = make_column_descr![PhysicalType::INT32, ConvertedType::INT_16]; + let row = Field::convert_int32(&descr, 222); + assert_eq!(row, Field::Short(222)); + + let descr = make_column_descr![PhysicalType::INT32, ConvertedType::INT_32]; + let row = Field::convert_int32(&descr, 333); + assert_eq!(row, Field::Int(333)); + + let descr = make_column_descr![PhysicalType::INT32, ConvertedType::UINT_8]; + let row = Field::convert_int32(&descr, -1); + assert_eq!(row, Field::UByte(255)); + + let descr = make_column_descr![PhysicalType::INT32, ConvertedType::UINT_16]; + let row = Field::convert_int32(&descr, 256); + assert_eq!(row, Field::UShort(256)); + + let descr = make_column_descr![PhysicalType::INT32, ConvertedType::UINT_32]; + let row = Field::convert_int32(&descr, 1234); + assert_eq!(row, Field::UInt(1234)); + + let descr = make_column_descr![PhysicalType::INT32, ConvertedType::NONE]; + let row = Field::convert_int32(&descr, 444); + assert_eq!(row, Field::Int(444)); + + let descr = make_column_descr![PhysicalType::INT32, ConvertedType::DATE]; + let row = Field::convert_int32(&descr, 14611); + assert_eq!(row, Field::Date(14611)); + + let descr = make_column_descr![PhysicalType::INT32, ConvertedType::TIME_MILLIS]; + let row = Field::convert_int32(&descr, 14611); + assert_eq!(row, Field::TimeMillis(14611)); + + let descr = make_column_descr![PhysicalType::INT32, ConvertedType::DECIMAL, 0, 8, 2]; + let row = Field::convert_int32(&descr, 444); + assert_eq!(row, Field::Decimal(Decimal::from_i32(444, 8, 2))); + } + + #[test] + fn test_row_convert_int64() { + let descr = make_column_descr![PhysicalType::INT64, ConvertedType::INT_64]; + let row = Field::convert_int64(&descr, 1111); + assert_eq!(row, Field::Long(1111)); + + let descr = make_column_descr![PhysicalType::INT64, ConvertedType::UINT_64]; + let row = Field::convert_int64(&descr, 78239823); + assert_eq!(row, Field::ULong(78239823)); + + let descr = make_column_descr![PhysicalType::INT64, ConvertedType::TIMESTAMP_MILLIS]; + let row = Field::convert_int64(&descr, 1541186529153); + assert_eq!(row, Field::TimestampMillis(1541186529153)); + + let descr = make_column_descr![PhysicalType::INT64, ConvertedType::TIMESTAMP_MICROS]; + let row = Field::convert_int64(&descr, 1541186529153123); + assert_eq!(row, Field::TimestampMicros(1541186529153123)); + + let descr = make_column_descr![PhysicalType::INT64, ConvertedType::TIME_MICROS]; + let row = Field::convert_int64(&descr, 47445123456); + assert_eq!(row, Field::TimeMicros(47445123456)); + + let descr = make_column_descr![PhysicalType::INT64, ConvertedType::NONE]; + let row = Field::convert_int64(&descr, 2222); + assert_eq!(row, Field::Long(2222)); + + let descr = make_column_descr![PhysicalType::INT64, ConvertedType::DECIMAL, 0, 8, 2]; + let row = Field::convert_int64(&descr, 3333); + assert_eq!(row, Field::Decimal(Decimal::from_i64(3333, 8, 2))); + } + + #[test] + fn test_row_convert_int96() { + // INT96 value does not depend on logical type + let descr = make_column_descr![PhysicalType::INT96, ConvertedType::NONE]; + + let value = Int96::from(vec![0, 0, 2454923]); + let row = Field::convert_int96(&descr, value); + assert_eq!(row, Field::TimestampMillis(1238544000000)); + + let value = Int96::from(vec![4165425152, 13, 2454923]); + let row = Field::convert_int96(&descr, value); + assert_eq!(row, Field::TimestampMillis(1238544060000)); + } + + #[test] + fn test_row_convert_float() { + // FLOAT value does not depend on logical type + let descr = make_column_descr![PhysicalType::FLOAT, ConvertedType::NONE]; + let row = Field::convert_float(&descr, 2.31); + assert_eq!(row, Field::Float(2.31)); + } + + #[test] + fn test_row_convert_double() { + // DOUBLE value does not depend on logical type + let descr = make_column_descr![PhysicalType::DOUBLE, ConvertedType::NONE]; + let row = Field::convert_double(&descr, 1.56); + assert_eq!(row, Field::Double(1.56)); + } + + #[test] + fn test_row_convert_byte_array() { + // UTF8 + let descr = make_column_descr![PhysicalType::BYTE_ARRAY, ConvertedType::UTF8]; + let value = ByteArray::from(vec![b'A', b'B', b'C', b'D']); + let row = Field::convert_byte_array(&descr, value); + assert_eq!(row.unwrap(), Field::Str("ABCD".to_string())); + + // ENUM + let descr = make_column_descr![PhysicalType::BYTE_ARRAY, ConvertedType::ENUM]; + let value = ByteArray::from(vec![b'1', b'2', b'3']); + let row = Field::convert_byte_array(&descr, value); + assert_eq!(row.unwrap(), Field::Str("123".to_string())); + + // JSON + let descr = make_column_descr![PhysicalType::BYTE_ARRAY, ConvertedType::JSON]; + let value = ByteArray::from(vec![b'{', b'"', b'a', b'"', b':', b'1', b'}']); + let row = Field::convert_byte_array(&descr, value); + assert_eq!(row.unwrap(), Field::Str("{\"a\":1}".to_string())); + + // NONE + let descr = make_column_descr![PhysicalType::BYTE_ARRAY, ConvertedType::NONE]; + let value = ByteArray::from(vec![1, 2, 3, 4, 5]); + let row = Field::convert_byte_array(&descr, value.clone()); + assert_eq!(row.unwrap(), Field::Bytes(value)); + + // BSON + let descr = make_column_descr![PhysicalType::BYTE_ARRAY, ConvertedType::BSON]; + let value = ByteArray::from(vec![1, 2, 3, 4, 5]); + let row = Field::convert_byte_array(&descr, value.clone()); + assert_eq!(row.unwrap(), Field::Bytes(value)); + + // DECIMAL + let descr = make_column_descr![PhysicalType::BYTE_ARRAY, ConvertedType::DECIMAL, 0, 8, 2]; + let value = ByteArray::from(vec![207, 200]); + let row = Field::convert_byte_array(&descr, value.clone()); + assert_eq!( + row.unwrap(), + Field::Decimal(Decimal::from_bytes(value, 8, 2)) + ); + + // DECIMAL (FIXED_LEN_BYTE_ARRAY) + let descr = make_column_descr![ + PhysicalType::FIXED_LEN_BYTE_ARRAY, + ConvertedType::DECIMAL, + 8, + 17, + 5 + ]; + let value = ByteArray::from(vec![0, 0, 0, 0, 0, 4, 147, 224]); + let row = Field::convert_byte_array(&descr, value.clone()); + assert_eq!( + row.unwrap(), + Field::Decimal(Decimal::from_bytes(value, 17, 5)) + ); + + // FLOAT16 + let descr = { + let tpe = PrimitiveTypeBuilder::new("col", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_logical_type(Some(LogicalType::Float16)) + .with_length(2) + .build() + .unwrap(); + Arc::new(ColumnDescriptor::new( + Arc::new(tpe), + 0, + 0, + ColumnPath::from("col"), + )) + }; + let value = ByteArray::from(f16::PI); + let row = Field::convert_byte_array(&descr, value.clone()); + assert_eq!(row.unwrap(), Field::Float16(f16::PI)); + + // NONE (FIXED_LEN_BYTE_ARRAY) + let descr = make_column_descr![ + PhysicalType::FIXED_LEN_BYTE_ARRAY, + ConvertedType::NONE, + 6, + 0, + 0 + ]; + let value = ByteArray::from(vec![1, 2, 3, 4, 5, 6]); + let row = Field::convert_byte_array(&descr, value.clone()); + assert_eq!(row.unwrap(), Field::Bytes(value)); + } + + #[test] + fn test_convert_date_to_string() { + fn check_date_conversion(y: u32, m: u32, d: u32) { + let datetime = chrono::NaiveDate::from_ymd_opt(y as i32, m, d) + .unwrap() + .and_hms_opt(0, 0, 0) + .unwrap(); + let dt = Utc.from_utc_datetime(&datetime); + let res = convert_date_to_string((dt.timestamp() / 60 / 60 / 24) as i32); + let exp = format!("{}", dt.format("%Y-%m-%d")); + assert_eq!(res, exp); + } + + check_date_conversion(1969, 12, 31); + check_date_conversion(2010, 1, 2); + check_date_conversion(2014, 5, 1); + check_date_conversion(2016, 2, 29); + check_date_conversion(2017, 9, 12); + check_date_conversion(2018, 3, 31); + } + + #[test] + fn test_convert_timestamp_millis_to_string() { + fn check_datetime_conversion( + (y, m, d, h, mi, s, milli): (u32, u32, u32, u32, u32, u32, u32), + exp: &str, + ) { + let datetime = chrono::NaiveDate::from_ymd_opt(y as i32, m, d) + .unwrap() + .and_hms_milli_opt(h, mi, s, milli) + .unwrap(); + let dt = Utc.from_utc_datetime(&datetime); + let res = convert_timestamp_millis_to_string(dt.timestamp_millis()); + assert_eq!(res, exp); + } + + check_datetime_conversion((1969, 9, 10, 1, 2, 3, 4), "1969-09-10 01:02:03.004 +00:00"); + check_datetime_conversion( + (2010, 1, 2, 13, 12, 54, 42), + "2010-01-02 13:12:54.042 +00:00", + ); + check_datetime_conversion((2011, 1, 3, 8, 23, 1, 27), "2011-01-03 08:23:01.027 +00:00"); + check_datetime_conversion((2012, 4, 5, 11, 6, 32, 0), "2012-04-05 11:06:32.000 +00:00"); + check_datetime_conversion( + (2013, 5, 12, 16, 38, 0, 15), + "2013-05-12 16:38:00.015 +00:00", + ); + check_datetime_conversion( + (2014, 11, 28, 21, 15, 12, 59), + "2014-11-28 21:15:12.059 +00:00", + ); + } + + #[test] + fn test_convert_timestamp_micros_to_string() { + fn check_datetime_conversion( + (y, m, d, h, mi, s, micro): (u32, u32, u32, u32, u32, u32, u32), + exp: &str, + ) { + let datetime = chrono::NaiveDate::from_ymd_opt(y as i32, m, d) + .unwrap() + .and_hms_micro_opt(h, mi, s, micro) + .unwrap(); + let dt = Utc.from_utc_datetime(&datetime); + let res = convert_timestamp_micros_to_string(dt.timestamp_micros()); + assert_eq!(res, exp); + } + + check_datetime_conversion( + (1969, 9, 10, 1, 2, 3, 4), + "1969-09-10 01:02:03.000004 +00:00", + ); + check_datetime_conversion( + (2010, 1, 2, 13, 12, 54, 42), + "2010-01-02 13:12:54.000042 +00:00", + ); + check_datetime_conversion( + (2011, 1, 3, 8, 23, 1, 27), + "2011-01-03 08:23:01.000027 +00:00", + ); + check_datetime_conversion( + (2012, 4, 5, 11, 6, 32, 0), + "2012-04-05 11:06:32.000000 +00:00", + ); + check_datetime_conversion( + (2013, 5, 12, 16, 38, 0, 15), + "2013-05-12 16:38:00.000015 +00:00", + ); + check_datetime_conversion( + (2014, 11, 28, 21, 15, 12, 59), + "2014-11-28 21:15:12.000059 +00:00", + ); + } + + #[test] + fn test_convert_float16_to_string() { + assert_eq!(format!("{}", Field::Float16(f16::ONE)), "1.0"); + assert_eq!(format!("{}", Field::Float16(f16::PI)), "3.140625"); + assert_eq!(format!("{}", Field::Float16(f16::MAX)), "65504.0"); + assert_eq!(format!("{}", Field::Float16(f16::NAN)), "NaN"); + assert_eq!(format!("{}", Field::Float16(f16::INFINITY)), "inf"); + assert_eq!(format!("{}", Field::Float16(f16::NEG_INFINITY)), "-inf"); + assert_eq!(format!("{}", Field::Float16(f16::ZERO)), "0.0"); + assert_eq!(format!("{}", Field::Float16(f16::NEG_ZERO)), "-0.0"); + } + + #[test] + fn test_convert_float_to_string() { + assert_eq!(format!("{}", Field::Float(1.0)), "1.0"); + assert_eq!(format!("{}", Field::Float(9.63)), "9.63"); + assert_eq!(format!("{}", Field::Float(1e-15)), "0.000000000000001"); + assert_eq!(format!("{}", Field::Float(1e-16)), "1E-16"); + assert_eq!(format!("{}", Field::Float(1e19)), "10000000000000000000.0"); + assert_eq!(format!("{}", Field::Float(1e20)), "1E20"); + assert_eq!(format!("{}", Field::Float(1.7976931E30)), "1.7976931E30"); + assert_eq!(format!("{}", Field::Float(-1.7976931E30)), "-1.7976931E30"); + } + + #[test] + fn test_convert_double_to_string() { + assert_eq!(format!("{}", Field::Double(1.0)), "1.0"); + assert_eq!(format!("{}", Field::Double(9.63)), "9.63"); + assert_eq!(format!("{}", Field::Double(1e-15)), "0.000000000000001"); + assert_eq!(format!("{}", Field::Double(1e-16)), "1E-16"); + assert_eq!(format!("{}", Field::Double(1e19)), "10000000000000000000.0"); + assert_eq!(format!("{}", Field::Double(1e20)), "1E20"); + assert_eq!( + format!("{}", Field::Double(1.79769313486E308)), + "1.79769313486E308" + ); + assert_eq!( + format!("{}", Field::Double(-1.79769313486E308)), + "-1.79769313486E308" + ); + } + + #[test] + fn test_convert_decimal_to_string() { + // Helper method to compare decimal + fn check_decimal(bytes: Vec, precision: i32, scale: i32, res: &str) { + let decimal = Decimal::from_bytes(ByteArray::from(bytes), precision, scale); + assert_eq!(convert_decimal_to_string(&decimal), res); + } + + // This example previously used to fail in some engines + check_decimal( + vec![0, 0, 0, 0, 0, 0, 0, 0, 13, 224, 182, 179, 167, 100, 0, 0], + 38, + 18, + "1.000000000000000000", + ); + check_decimal( + vec![ + 249, 233, 247, 16, 185, 192, 202, 223, 215, 165, 192, 166, 67, 72, + ], + 36, + 28, + "-12344.0242342304923409234234293432", + ); + check_decimal(vec![0, 0, 0, 0, 0, 4, 147, 224], 17, 5, "3.00000"); + check_decimal(vec![0, 0, 0, 0, 1, 201, 195, 140], 18, 2, "300000.12"); + check_decimal(vec![207, 200], 10, 2, "-123.44"); + check_decimal(vec![207, 200], 10, 8, "-0.00012344"); + } + + #[test] + fn test_row_display() { + // Primitive types + assert_eq!(format!("{}", Field::Null), "null"); + assert_eq!(format!("{}", Field::Bool(true)), "true"); + assert_eq!(format!("{}", Field::Bool(false)), "false"); + assert_eq!(format!("{}", Field::Byte(1)), "1"); + assert_eq!(format!("{}", Field::Short(2)), "2"); + assert_eq!(format!("{}", Field::Int(3)), "3"); + assert_eq!(format!("{}", Field::Long(4)), "4"); + assert_eq!(format!("{}", Field::UByte(1)), "1"); + assert_eq!(format!("{}", Field::UShort(2)), "2"); + assert_eq!(format!("{}", Field::UInt(3)), "3"); + assert_eq!(format!("{}", Field::ULong(4)), "4"); + assert_eq!(format!("{}", Field::Float16(f16::E)), "2.71875"); + assert_eq!(format!("{}", Field::Float(5.0)), "5.0"); + assert_eq!(format!("{}", Field::Float(5.1234)), "5.1234"); + assert_eq!(format!("{}", Field::Double(6.0)), "6.0"); + assert_eq!(format!("{}", Field::Double(6.1234)), "6.1234"); + assert_eq!(format!("{}", Field::Str("abc".to_string())), "\"abc\""); + assert_eq!( + format!("{}", Field::Bytes(ByteArray::from(vec![1, 2, 3]))), + "[1, 2, 3]" + ); + assert_eq!( + format!("{}", Field::Date(14611)), + convert_date_to_string(14611) + ); + assert_eq!( + format!("{}", Field::TimestampMillis(1262391174000)), + convert_timestamp_millis_to_string(1262391174000) + ); + assert_eq!( + format!("{}", Field::TimestampMicros(1262391174000000)), + convert_timestamp_micros_to_string(1262391174000000) + ); + assert_eq!( + format!("{}", Field::Decimal(Decimal::from_i32(4, 8, 2))), + convert_decimal_to_string(&Decimal::from_i32(4, 8, 2)) + ); + + // Complex types + let fields = vec![ + ("x".to_string(), Field::Null), + ("Y".to_string(), Field::Int(2)), + ("z".to_string(), Field::Float(3.1)), + ("a".to_string(), Field::Str("abc".to_string())), + ]; + let row = Field::Group(Row::new(fields)); + assert_eq!(format!("{row}"), "{x: null, Y: 2, z: 3.1, a: \"abc\"}"); + + let row = Field::ListInternal(make_list(vec![ + Field::Int(2), + Field::Int(1), + Field::Null, + Field::Int(12), + ])); + assert_eq!(format!("{row}"), "[2, 1, null, 12]"); + + let row = Field::MapInternal(make_map(vec![ + (Field::Int(1), Field::Float(1.2)), + (Field::Int(2), Field::Float(4.5)), + (Field::Int(3), Field::Float(2.3)), + ])); + assert_eq!(format!("{row}"), "{1 -> 1.2, 2 -> 4.5, 3 -> 2.3}"); + } + + #[test] + fn test_is_primitive() { + // primitives + assert!(Field::Null.is_primitive()); + assert!(Field::Bool(true).is_primitive()); + assert!(Field::Bool(false).is_primitive()); + assert!(Field::Byte(1).is_primitive()); + assert!(Field::Short(2).is_primitive()); + assert!(Field::Int(3).is_primitive()); + assert!(Field::Long(4).is_primitive()); + assert!(Field::UByte(1).is_primitive()); + assert!(Field::UShort(2).is_primitive()); + assert!(Field::UInt(3).is_primitive()); + assert!(Field::ULong(4).is_primitive()); + assert!(Field::Float16(f16::E).is_primitive()); + assert!(Field::Float(5.0).is_primitive()); + assert!(Field::Float(5.1234).is_primitive()); + assert!(Field::Double(6.0).is_primitive()); + assert!(Field::Double(6.1234).is_primitive()); + assert!(Field::Str("abc".to_string()).is_primitive()); + assert!(Field::Bytes(ByteArray::from(vec![1, 2, 3])).is_primitive()); + assert!(Field::TimestampMillis(12345678).is_primitive()); + assert!(Field::TimestampMicros(12345678901).is_primitive()); + assert!(Field::Decimal(Decimal::from_i32(4, 8, 2)).is_primitive()); + + // complex types + assert!( + !Field::Group(Row::new(vec![ + ("x".to_string(), Field::Null), + ("Y".to_string(), Field::Int(2)), + ("z".to_string(), Field::Float(3.1)), + ("a".to_string(), Field::Str("abc".to_string())) + ])) + .is_primitive() + ); + + assert!( + !Field::ListInternal(make_list(vec![ + Field::Int(2), + Field::Int(1), + Field::Null, + Field::Int(12) + ])) + .is_primitive() + ); + + assert!( + !Field::MapInternal(make_map(vec![ + (Field::Int(1), Field::Float(1.2)), + (Field::Int(2), Field::Float(4.5)), + (Field::Int(3), Field::Float(2.3)) + ])) + .is_primitive() + ); + } + + #[test] + fn test_row_primitive_field_fmt() { + // Primitives types + let row = Row::new(vec![ + ("00".to_string(), Field::Null), + ("01".to_string(), Field::Bool(false)), + ("02".to_string(), Field::Byte(3)), + ("03".to_string(), Field::Short(4)), + ("04".to_string(), Field::Int(5)), + ("05".to_string(), Field::Long(6)), + ("06".to_string(), Field::UByte(7)), + ("07".to_string(), Field::UShort(8)), + ("08".to_string(), Field::UInt(9)), + ("09".to_string(), Field::ULong(10)), + ("10".to_string(), Field::Float(11.1)), + ("11".to_string(), Field::Double(12.1)), + ("12".to_string(), Field::Str("abc".to_string())), + ( + "13".to_string(), + Field::Bytes(ByteArray::from(vec![1, 2, 3, 4, 5])), + ), + ("14".to_string(), Field::Date(14611)), + ("15".to_string(), Field::TimestampMillis(1262391174000)), + ("16".to_string(), Field::TimestampMicros(1262391174000000)), + ("17".to_string(), Field::Decimal(Decimal::from_i32(4, 7, 2))), + ("18".to_string(), Field::Float16(f16::PI)), + ]); + + assert_eq!("null", format!("{}", row.fmt(0))); + assert_eq!("false", format!("{}", row.fmt(1))); + assert_eq!("3", format!("{}", row.fmt(2))); + assert_eq!("4", format!("{}", row.fmt(3))); + assert_eq!("5", format!("{}", row.fmt(4))); + assert_eq!("6", format!("{}", row.fmt(5))); + assert_eq!("7", format!("{}", row.fmt(6))); + assert_eq!("8", format!("{}", row.fmt(7))); + assert_eq!("9", format!("{}", row.fmt(8))); + assert_eq!("10", format!("{}", row.fmt(9))); + assert_eq!("11.1", format!("{}", row.fmt(10))); + assert_eq!("12.1", format!("{}", row.fmt(11))); + assert_eq!("\"abc\"", format!("{}", row.fmt(12))); + assert_eq!("[1, 2, 3, 4, 5]", format!("{}", row.fmt(13))); + assert_eq!(convert_date_to_string(14611), format!("{}", row.fmt(14))); + assert_eq!( + convert_timestamp_millis_to_string(1262391174000), + format!("{}", row.fmt(15)) + ); + assert_eq!( + convert_timestamp_micros_to_string(1262391174000000), + format!("{}", row.fmt(16)) + ); + assert_eq!("0.04", format!("{}", row.fmt(17))); + assert_eq!("3.140625", format!("{}", row.fmt(18))); + } + + #[test] + fn test_row_complex_field_fmt() { + // Complex types + let row = Row::new(vec![ + ( + "00".to_string(), + Field::Group(Row::new(vec![ + ("x".to_string(), Field::Null), + ("Y".to_string(), Field::Int(2)), + ])), + ), + ( + "01".to_string(), + Field::ListInternal(make_list(vec![ + Field::Int(2), + Field::Int(1), + Field::Null, + Field::Int(12), + ])), + ), + ( + "02".to_string(), + Field::MapInternal(make_map(vec![ + (Field::Int(1), Field::Float(1.2)), + (Field::Int(2), Field::Float(4.5)), + (Field::Int(3), Field::Float(2.3)), + ])), + ), + ]); + + assert_eq!("{x: null, Y: 2}", format!("{}", row.fmt(0))); + assert_eq!("[2, 1, null, 12]", format!("{}", row.fmt(1))); + assert_eq!("{1 -> 1.2, 2 -> 4.5, 3 -> 2.3}", format!("{}", row.fmt(2))); + } + + #[test] + fn test_row_primitive_accessors() { + // primitives + let row = Row::new(vec![ + ("a".to_string(), Field::Null), + ("b".to_string(), Field::Bool(false)), + ("c".to_string(), Field::Byte(3)), + ("d".to_string(), Field::Short(4)), + ("e".to_string(), Field::Int(5)), + ("f".to_string(), Field::Long(6)), + ("g".to_string(), Field::UByte(3)), + ("h".to_string(), Field::UShort(4)), + ("i".to_string(), Field::UInt(5)), + ("j".to_string(), Field::ULong(6)), + ("k".to_string(), Field::Float(7.1)), + ("l".to_string(), Field::Double(8.1)), + ("m".to_string(), Field::Str("abc".to_string())), + ( + "n".to_string(), + Field::Bytes(ByteArray::from(vec![1, 2, 3, 4, 5])), + ), + ("o".to_string(), Field::Decimal(Decimal::from_i32(4, 7, 2))), + ("p".to_string(), Field::Float16(f16::from_f32(9.1))), + ]); + + assert!(row.is_null(0).unwrap()); + assert!(!row.is_null(1).unwrap()); + assert!(!row.get_bool(1).unwrap()); + assert_eq!(3, row.get_byte(2).unwrap()); + assert_eq!(4, row.get_short(3).unwrap()); + assert_eq!(5, row.get_int(4).unwrap()); + assert_eq!(6, row.get_long(5).unwrap()); + assert_eq!(3, row.get_ubyte(6).unwrap()); + assert_eq!(4, row.get_ushort(7).unwrap()); + assert_eq!(5, row.get_uint(8).unwrap()); + assert_eq!(6, row.get_ulong(9).unwrap()); + assert!((7.1 - row.get_float(10).unwrap()).abs() < f32::EPSILON); + assert!((8.1 - row.get_double(11).unwrap()).abs() < f64::EPSILON); + assert_eq!("abc", row.get_string(12).unwrap()); + assert_eq!(5, row.get_bytes(13).unwrap().len()); + assert_eq!(7, row.get_decimal(14).unwrap().precision()); + assert!((f16::from_f32(9.1) - row.get_float16(15).unwrap()).abs() < f16::EPSILON); + + assert!(matches!( + row.is_null(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_bool(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_byte(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_short(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_int(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_long(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_ubyte(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_ushort(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_uint(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_ulong(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_float(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_double(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_string(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_bytes(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_decimal(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + assert!(matches!( + row.get_float16(16).unwrap_err(), + ParquetError::IndexOutOfBound(16, 16), + )); + } + + #[test] + fn test_row_primitive_invalid_accessors() { + // primitives + let row = Row::new(vec![ + ("a".to_string(), Field::Null), + ("b".to_string(), Field::Bool(false)), + ("c".to_string(), Field::Byte(3)), + ("d".to_string(), Field::Short(4)), + ("e".to_string(), Field::Int(5)), + ("f".to_string(), Field::Long(6)), + ("g".to_string(), Field::UByte(3)), + ("h".to_string(), Field::UShort(4)), + ("i".to_string(), Field::UInt(5)), + ("j".to_string(), Field::ULong(6)), + ("k".to_string(), Field::Float(7.1)), + ("l".to_string(), Field::Double(8.1)), + ("m".to_string(), Field::Str("abc".to_string())), + ( + "n".to_string(), + Field::Bytes(ByteArray::from(vec![1, 2, 3, 4, 5])), + ), + ("o".to_string(), Field::Decimal(Decimal::from_i32(4, 7, 2))), + ("p".to_string(), Field::Float16(f16::from_f32(9.1))), + ]); + + for i in 0..row.len() { + assert!(row.get_group(i).is_err()); + } + } + + #[test] + fn test_row_complex_accessors() { + let row = Row::new(vec![ + ( + "a".to_string(), + Field::Group(Row::new(vec![ + ("x".to_string(), Field::Null), + ("Y".to_string(), Field::Int(2)), + ])), + ), + ( + "b".to_string(), + Field::ListInternal(make_list(vec![ + Field::Int(2), + Field::Int(1), + Field::Null, + Field::Int(12), + ])), + ), + ( + "c".to_string(), + Field::MapInternal(make_map(vec![ + (Field::Int(1), Field::Float(1.2)), + (Field::Int(2), Field::Float(4.5)), + (Field::Int(3), Field::Float(2.3)), + ])), + ), + ]); + + assert_eq!(2, row.get_group(0).unwrap().len()); + assert_eq!(4, row.get_list(1).unwrap().len()); + assert_eq!(3, row.get_map(2).unwrap().len()); + } + + #[test] + fn test_row_complex_invalid_accessors() { + let row = Row::new(vec![ + ( + "a".to_string(), + Field::Group(Row::new(vec![ + ("x".to_string(), Field::Null), + ("Y".to_string(), Field::Int(2)), + ])), + ), + ( + "b".to_string(), + Field::ListInternal(make_list(vec![ + Field::Int(2), + Field::Int(1), + Field::Null, + Field::Int(12), + ])), + ), + ( + "c".to_string(), + Field::MapInternal(make_map(vec![ + (Field::Int(1), Field::Float(1.2)), + (Field::Int(2), Field::Float(4.5)), + (Field::Int(3), Field::Float(2.3)), + ])), + ), + ]); + + assert_eq!( + row.get_float(0).unwrap_err().to_string(), + "Parquet error: Cannot access Group as Float" + ); + assert_eq!( + row.get_float(1).unwrap_err().to_string(), + "Parquet error: Cannot access ListInternal as Float" + ); + assert_eq!( + row.get_float(2).unwrap_err().to_string(), + "Parquet error: Cannot access MapInternal as Float", + ); + } + + #[test] + fn test_list_primitive_accessors() { + // primitives + let list = make_list(vec![Field::Bool(false)]); + assert!(!list.get_bool(0).unwrap()); + + let list = make_list(vec![Field::Byte(3), Field::Byte(4)]); + assert_eq!(4, list.get_byte(1).unwrap()); + + let list = make_list(vec![Field::Short(4), Field::Short(5), Field::Short(6)]); + assert_eq!(6, list.get_short(2).unwrap()); + + let list = make_list(vec![Field::Int(5)]); + assert_eq!(5, list.get_int(0).unwrap()); + + let list = make_list(vec![Field::Long(6), Field::Long(7)]); + assert_eq!(7, list.get_long(1).unwrap()); + + let list = make_list(vec![Field::UByte(3), Field::UByte(4)]); + assert_eq!(4, list.get_ubyte(1).unwrap()); + + let list = make_list(vec![Field::UShort(4), Field::UShort(5), Field::UShort(6)]); + assert_eq!(6, list.get_ushort(2).unwrap()); + + let list = make_list(vec![Field::UInt(5)]); + assert_eq!(5, list.get_uint(0).unwrap()); + + let list = make_list(vec![Field::ULong(6), Field::ULong(7)]); + assert_eq!(7, list.get_ulong(1).unwrap()); + + let list = make_list(vec![Field::Float16(f16::PI)]); + assert!((f16::PI - list.get_float16(0).unwrap()).abs() < f16::EPSILON); + + let list = make_list(vec![ + Field::Float(8.1), + Field::Float(9.2), + Field::Float(10.3), + ]); + assert!((10.3 - list.get_float(2).unwrap()).abs() < f32::EPSILON); + + let list = make_list(vec![Field::Double(PI)]); + assert!((PI - list.get_double(0).unwrap()).abs() < f64::EPSILON); + + let list = make_list(vec![Field::Str("abc".to_string())]); + assert_eq!(&"abc".to_string(), list.get_string(0).unwrap()); + + let list = make_list(vec![Field::Bytes(ByteArray::from(vec![1, 2, 3, 4, 5]))]); + assert_eq!(&[1, 2, 3, 4, 5], list.get_bytes(0).unwrap().data()); + + let list = make_list(vec![Field::Decimal(Decimal::from_i32(4, 5, 2))]); + assert_eq!(&[0, 0, 0, 4], list.get_decimal(0).unwrap().data()); + } + + #[test] + fn test_list_primitive_invalid_accessors() { + // primitives + let list = make_list(vec![Field::Bool(false)]); + assert!(list.get_byte(0).is_err()); + + let list = make_list(vec![Field::Byte(3), Field::Byte(4)]); + assert!(list.get_short(1).is_err()); + + let list = make_list(vec![Field::Short(4), Field::Short(5), Field::Short(6)]); + assert!(list.get_int(2).is_err()); + + let list = make_list(vec![Field::Int(5)]); + assert!(list.get_long(0).is_err()); + + let list = make_list(vec![Field::Long(6), Field::Long(7)]); + assert!(list.get_float(1).is_err()); + + let list = make_list(vec![Field::UByte(3), Field::UByte(4)]); + assert!(list.get_short(1).is_err()); + + let list = make_list(vec![Field::UShort(4), Field::UShort(5), Field::UShort(6)]); + assert!(list.get_int(2).is_err()); + + let list = make_list(vec![Field::UInt(5)]); + assert!(list.get_long(0).is_err()); + + let list = make_list(vec![Field::ULong(6), Field::ULong(7)]); + assert!(list.get_float(1).is_err()); + + let list = make_list(vec![Field::Float16(f16::PI)]); + assert!(list.get_string(0).is_err()); + + let list = make_list(vec![ + Field::Float(8.1), + Field::Float(9.2), + Field::Float(10.3), + ]); + assert!(list.get_double(2).is_err()); + + let list = make_list(vec![Field::Double(PI)]); + assert!(list.get_string(0).is_err()); + + let list = make_list(vec![Field::Str("abc".to_string())]); + assert!(list.get_bytes(0).is_err()); + + let list = make_list(vec![Field::Bytes(ByteArray::from(vec![1, 2, 3, 4, 5]))]); + assert!(list.get_bool(0).is_err()); + + let list = make_list(vec![Field::Decimal(Decimal::from_i32(4, 5, 2))]); + assert!(list.get_bool(0).is_err()); + } + + #[test] + fn test_list_complex_accessors() { + let list = make_list(vec![Field::Group(Row::new(vec![ + ("x".to_string(), Field::Null), + ("Y".to_string(), Field::Int(2)), + ]))]); + assert_eq!(2, list.get_group(0).unwrap().len()); + + let list = make_list(vec![Field::ListInternal(make_list(vec![ + Field::Int(2), + Field::Int(1), + Field::Null, + Field::Int(12), + ]))]); + assert_eq!(4, list.get_list(0).unwrap().len()); + + let list = make_list(vec![Field::MapInternal(make_map(vec![ + (Field::Int(1), Field::Float(1.2)), + (Field::Int(2), Field::Float(4.5)), + (Field::Int(3), Field::Float(2.3)), + ]))]); + assert_eq!(3, list.get_map(0).unwrap().len()); + } + + #[test] + fn test_list_complex_invalid_accessors() { + let list = make_list(vec![Field::Group(Row::new(vec![ + ("x".to_string(), Field::Null), + ("Y".to_string(), Field::Int(2)), + ]))]); + assert_eq!( + list.get_float(0).unwrap_err().to_string(), + "Parquet error: Cannot access Group as Float" + ); + + let list = make_list(vec![Field::ListInternal(make_list(vec![ + Field::Int(2), + Field::Int(1), + Field::Null, + Field::Int(12), + ]))]); + assert_eq!( + list.get_float(0).unwrap_err().to_string(), + "Parquet error: Cannot access ListInternal as Float" + ); + + let list = make_list(vec![Field::MapInternal(make_map(vec![ + (Field::Int(1), Field::Float(1.2)), + (Field::Int(2), Field::Float(4.5)), + (Field::Int(3), Field::Float(2.3)), + ]))]); + assert_eq!( + list.get_float(0).unwrap_err().to_string(), + "Parquet error: Cannot access MapInternal as Float", + ); + } + + #[test] + fn test_map_accessors() { + // a map from int to string + let map = make_map(vec![ + (Field::Int(1), Field::Str("a".to_string())), + (Field::Int(2), Field::Str("b".to_string())), + (Field::Int(3), Field::Str("c".to_string())), + (Field::Int(4), Field::Str("d".to_string())), + (Field::Int(5), Field::Str("e".to_string())), + ]); + + assert_eq!(5, map.len()); + for i in 0..5 { + assert_eq!((i + 1) as i32, map.get_keys().get_int(i).unwrap()); + assert_eq!( + &((i as u8 + b'a') as char).to_string(), + map.get_values().get_string(i).unwrap() + ); + } + } + + #[test] + fn test_to_json_value() { + assert_eq!(Field::Null.to_json_value(), Value::Null); + assert_eq!(Field::Bool(true).to_json_value(), Value::Bool(true)); + assert_eq!(Field::Bool(false).to_json_value(), Value::Bool(false)); + assert_eq!( + Field::Byte(1).to_json_value(), + Value::Number(serde_json::Number::from(1)) + ); + assert_eq!( + Field::Short(2).to_json_value(), + Value::Number(serde_json::Number::from(2)) + ); + assert_eq!( + Field::Int(3).to_json_value(), + Value::Number(serde_json::Number::from(3)) + ); + assert_eq!( + Field::Long(4).to_json_value(), + Value::Number(serde_json::Number::from(4)) + ); + assert_eq!( + Field::UByte(1).to_json_value(), + Value::Number(serde_json::Number::from(1)) + ); + assert_eq!( + Field::UShort(2).to_json_value(), + Value::Number(serde_json::Number::from(2)) + ); + assert_eq!( + Field::UInt(3).to_json_value(), + Value::Number(serde_json::Number::from(3)) + ); + assert_eq!( + Field::ULong(4).to_json_value(), + Value::Number(serde_json::Number::from(4)) + ); + assert_eq!( + Field::Float16(f16::from_f32(5.0)).to_json_value(), + Value::Number(serde_json::Number::from_f64(5.0).unwrap()) + ); + assert_eq!( + Field::Float(5.0).to_json_value(), + Value::Number(serde_json::Number::from_f64(5.0).unwrap()) + ); + assert_eq!( + Field::Float(5.1234).to_json_value(), + Value::Number(serde_json::Number::from_f64(5.1234_f32 as f64).unwrap()) + ); + assert_eq!( + Field::Double(6.0).to_json_value(), + Value::Number(serde_json::Number::from_f64(6.0).unwrap()) + ); + assert_eq!( + Field::Double(6.1234).to_json_value(), + Value::Number(serde_json::Number::from_f64(6.1234).unwrap()) + ); + assert_eq!( + Field::Str("abc".to_string()).to_json_value(), + Value::String(String::from("abc")) + ); + assert_eq!( + Field::Decimal(Decimal::from_i32(4, 8, 2)).to_json_value(), + Value::String(String::from("0.04")) + ); + assert_eq!( + Field::Bytes(ByteArray::from(vec![1, 2, 3])).to_json_value(), + Value::String(String::from("AQID")) + ); + assert_eq!( + Field::TimestampMillis(12345678).to_json_value(), + Value::String("1970-01-01 03:25:45.678 +00:00".to_string()) + ); + assert_eq!( + Field::TimestampMicros(12345678901).to_json_value(), + Value::String("1970-01-01 03:25:45.678901 +00:00".to_string()) + ); + assert_eq!( + Field::TimeMillis(47445123).to_json_value(), + Value::String(String::from("13:10:45.123")) + ); + assert_eq!( + Field::TimeMicros(47445123456).to_json_value(), + Value::String(String::from("13:10:45.123456")) + ); + + let fields = vec![ + ("X".to_string(), Field::Int(1)), + ("Y".to_string(), Field::Double(2.2)), + ("Z".to_string(), Field::Str("abc".to_string())), + ]; + let row = Field::Group(Row::new(fields)); + assert_eq!( + row.to_json_value(), + serde_json::json!({"X": 1, "Y": 2.2, "Z": "abc"}) + ); + + let row = Field::ListInternal(make_list(vec![Field::Int(1), Field::Int(12), Field::Null])); + let array = vec![ + Value::Number(serde_json::Number::from(1)), + Value::Number(serde_json::Number::from(12)), + Value::Null, + ]; + assert_eq!(row.to_json_value(), Value::Array(array)); + + let row = Field::MapInternal(make_map(vec![ + (Field::Str("k1".to_string()), Field::Double(1.2)), + (Field::Str("k2".to_string()), Field::Double(3.4)), + (Field::Str("k3".to_string()), Field::Double(4.5)), + ])); + assert_eq!( + row.to_json_value(), + serde_json::json!({"k1": 1.2, "k2": 3.4, "k3": 4.5}) + ); + } +} + +#[cfg(test)] +#[allow(clippy::many_single_char_names)] +mod api_tests { + use super::{Row, make_list, make_map}; + use crate::record::Field; + + #[test] + fn test_field_visibility() { + let row = Row::new(vec![( + "a".to_string(), + Field::Group(Row::new(vec![ + ("x".to_string(), Field::Null), + ("Y".to_string(), Field::Int(2)), + ])), + )]); + + match row.get_column_iter().next() { + Some(column) => { + assert_eq!("a", column.0); + match column.1 { + Field::Group(r) => { + assert_eq!( + &Row::new(vec![ + ("x".to_string(), Field::Null), + ("Y".to_string(), Field::Int(2)), + ]), + r + ); + } + _ => panic!("Expected the first column to be Field::Group"), + } + } + None => panic!("Expected at least one column"), + } + } + + #[test] + fn test_list_element_access() { + let expected = vec![ + Field::Int(1), + Field::Group(Row::new(vec![ + ("x".to_string(), Field::Null), + ("Y".to_string(), Field::Int(2)), + ])), + ]; + + let list = make_list(expected.clone()); + assert_eq!(expected.as_slice(), list.elements()); + } + + #[test] + fn test_map_entry_access() { + let expected = vec![ + (Field::Str("one".to_owned()), Field::Int(1)), + (Field::Str("two".to_owned()), Field::Int(2)), + ]; + + let map = make_map(expected.clone()); + assert_eq!(expected.as_slice(), map.entries()); + } +} diff --git a/vendor/parquet/src/record/mod.rs b/vendor/parquet/src/record/mod.rs new file mode 100644 index 0000000..f40e914 --- /dev/null +++ b/vendor/parquet/src/record/mod.rs @@ -0,0 +1,32 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Contains record-based API for reading Parquet files. + +mod api; +pub mod reader; +mod record_reader; +mod record_writer; +mod triplet; + +pub use self::{ + api::{ + Field, List, ListAccessor, Map, MapAccessor, Row, RowAccessor, RowColumnIter, RowFormatter, + }, + record_reader::RecordReader, + record_writer::RecordWriter, +}; diff --git a/vendor/parquet/src/record/reader.rs b/vendor/parquet/src/record/reader.rs new file mode 100644 index 0000000..bee49c7 --- /dev/null +++ b/vendor/parquet/src/record/reader.rs @@ -0,0 +1,1980 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Contains implementation of record assembly and converting Parquet types into +//! [`Row`]s. + +use std::{collections::HashMap, fmt, sync::Arc}; + +use crate::basic::{ConvertedType, Repetition}; +use crate::errors::{ParquetError, Result}; +use crate::file::reader::{FileReader, RowGroupReader}; +use crate::record::{ + api::{Field, Row, make_list, make_map}, + triplet::TripletIter, +}; +use crate::schema::types::{ColumnPath, SchemaDescPtr, SchemaDescriptor, Type, TypePtr}; + +/// Default batch size for a reader +const DEFAULT_BATCH_SIZE: usize = 1024; + +/// Tree builder for `Reader` enum. +/// Serves as a container of options for building a reader tree and a builder, and +/// accessing a records iterator [`RowIter`]. +pub struct TreeBuilder { + // Batch size (>= 1) for triplet iterators + batch_size: usize, +} + +impl Default for TreeBuilder { + fn default() -> Self { + Self::new() + } +} + +impl TreeBuilder { + /// Creates new tree builder with default parameters. + pub fn new() -> Self { + Self { + batch_size: DEFAULT_BATCH_SIZE, + } + } + + /// Sets batch size for this tree builder. + pub fn with_batch_size(mut self, batch_size: usize) -> Self { + self.batch_size = batch_size; + self + } + + /// Creates new root reader for provided schema and row group. + pub fn build( + &self, + descr: SchemaDescPtr, + row_group_reader: &dyn RowGroupReader, + ) -> Result { + // Prepare lookup table of column path -> original column index + // This allows to prune columns and map schema leaf nodes to the column readers + let mut paths: HashMap = HashMap::new(); + let row_group_metadata = row_group_reader.metadata(); + + for col_index in 0..row_group_reader.num_columns() { + let col_meta = row_group_metadata.column(col_index); + let col_path = col_meta.column_path().clone(); + paths.insert(col_path, col_index); + } + + // Build child readers for the message type + let mut readers = Vec::new(); + let mut path = Vec::new(); + + for field in descr.root_schema().get_fields() { + let reader = + self.reader_tree(field.clone(), &mut path, 0, 0, &paths, row_group_reader)?; + readers.push(reader); + } + + // Return group reader for message type, + // it is always required with definition level 0 + Ok(Reader::GroupReader(None, 0, readers)) + } + + /// Creates iterator of `Row`s directly from schema descriptor and row group. + pub fn as_iter( + &self, + descr: SchemaDescPtr, + row_group_reader: &dyn RowGroupReader, + ) -> Result { + let num_records = row_group_reader.metadata().num_rows() as usize; + ReaderIter::new(self.build(descr, row_group_reader)?, num_records) + } + + /// Builds tree of readers for the current schema recursively. + fn reader_tree( + &self, + field: TypePtr, + path: &mut Vec, + mut curr_def_level: i16, + mut curr_rep_level: i16, + paths: &HashMap, + row_group_reader: &dyn RowGroupReader, + ) -> Result { + assert!(field.get_basic_info().has_repetition()); + // Update current definition and repetition levels for this type + let repetition = field.get_basic_info().repetition(); + match repetition { + Repetition::OPTIONAL => { + curr_def_level += 1; + } + Repetition::REPEATED => { + curr_def_level += 1; + curr_rep_level += 1; + } + _ => {} + } + + path.push(String::from(field.name())); + let reader = if field.is_primitive() { + let col_path = ColumnPath::new(path.to_vec()); + let orig_index = *paths + .get(&col_path) + .ok_or(general_err!("Path {:?} not found", col_path))?; + let col_descr = row_group_reader + .metadata() + .column(orig_index) + .column_descr_ptr(); + let col_reader = row_group_reader.get_column_reader(orig_index)?; + let column = TripletIter::new(col_descr, col_reader, self.batch_size); + let reader = Reader::PrimitiveReader(field.clone(), Box::new(column)); + if repetition == Repetition::REPEATED { + Reader::RepeatedReader( + field, + curr_def_level - 1, + curr_rep_level - 1, + Box::new(reader), + ) + } else { + reader + } + } else { + match field.get_basic_info().converted_type() { + // List types + ConvertedType::LIST => { + assert_eq!(field.get_fields().len(), 1, "Invalid list type {field:?}"); + + let repeated_field = field.get_fields()[0].clone(); + assert_eq!( + repeated_field.get_basic_info().repetition(), + Repetition::REPEATED, + "Invalid list type {field:?}" + ); + + if Reader::is_element_type(&repeated_field) { + // Support for backward compatible lists + let reader = self.reader_tree( + repeated_field, + path, + curr_def_level, + curr_rep_level, + paths, + row_group_reader, + )?; + + Reader::RepeatedReader( + field, + curr_def_level, + curr_rep_level, + Box::new(reader), + ) + } else { + let child_field = repeated_field.get_fields()[0].clone(); + + path.push(String::from(repeated_field.name())); + + let reader = self.reader_tree( + child_field, + path, + curr_def_level + 1, + curr_rep_level + 1, + paths, + row_group_reader, + )?; + + path.pop(); + + Reader::RepeatedReader( + field, + curr_def_level, + curr_rep_level, + Box::new(reader), + ) + } + } + // Map types (key-value pairs) + ConvertedType::MAP | ConvertedType::MAP_KEY_VALUE => { + assert_eq!(field.get_fields().len(), 1, "Invalid map type: {field:?}"); + assert!( + !field.get_fields()[0].is_primitive(), + "Invalid map type: {field:?}" + ); + + let key_value_type = field.get_fields()[0].clone(); + assert_eq!( + key_value_type.get_basic_info().repetition(), + Repetition::REPEATED, + "Invalid map type: {field:?}" + ); + // Parquet spec allows no value. In that case treat as a list. #1642 + if key_value_type.get_fields().len() != 1 { + // If not a list, then there can only be 2 fields in the struct + assert_eq!( + key_value_type.get_fields().len(), + 2, + "Invalid map type: {field:?}" + ); + } + + path.push(String::from(key_value_type.name())); + + let key_type = &key_value_type.get_fields()[0]; + assert!( + key_type.is_primitive(), + "Map key type is expected to be a primitive type, but found {key_type:?}" + ); + let key_reader = self.reader_tree( + key_type.clone(), + path, + curr_def_level + 1, + curr_rep_level + 1, + paths, + row_group_reader, + )?; + + if key_value_type.get_fields().len() == 1 { + path.pop(); + Reader::RepeatedReader( + field, + curr_def_level, + curr_rep_level, + Box::new(key_reader), + ) + } else { + let value_type = &key_value_type.get_fields()[1]; + let value_reader = self.reader_tree( + value_type.clone(), + path, + curr_def_level + 1, + curr_rep_level + 1, + paths, + row_group_reader, + )?; + + path.pop(); + + Reader::KeyValueReader( + field, + curr_def_level, + curr_rep_level, + Box::new(key_reader), + Box::new(value_reader), + ) + } + } + // A repeated field that is neither contained by a `LIST`- or + // `MAP`-annotated group nor annotated by `LIST` or `MAP` + // should be interpreted as a required list of required + // elements where the element type is the type of the field. + _ if repetition == Repetition::REPEATED => { + let required_field = Type::group_type_builder(field.name()) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(field.get_basic_info().converted_type()) + .with_fields(field.get_fields().to_vec()) + .build()?; + + path.pop(); + + let reader = self.reader_tree( + Arc::new(required_field), + path, + curr_def_level, + curr_rep_level, + paths, + row_group_reader, + )?; + + return Ok(Reader::RepeatedReader( + field, + curr_def_level - 1, + curr_rep_level - 1, + Box::new(reader), + )); + } + // Group types (structs) + _ => { + let mut readers = Vec::new(); + for child in field.get_fields() { + let reader = self.reader_tree( + child.clone(), + path, + curr_def_level, + curr_rep_level, + paths, + row_group_reader, + )?; + readers.push(reader); + } + Reader::GroupReader(Some(field), curr_def_level, readers) + } + } + }; + path.pop(); + + Ok(Reader::option(repetition, curr_def_level, reader)) + } +} + +/// Reader tree for record assembly +pub enum Reader { + /// Primitive reader with type information and triplet iterator + PrimitiveReader(TypePtr, Box), + /// Optional reader with definition level of a parent and a reader + OptionReader(i16, Box), + /// Group (struct) reader with type information, definition level and list of child + /// readers. When it represents message type, type information is None + GroupReader(Option, i16, Vec), + /// Reader for repeated values, e.g. lists, contains type information, definition + /// level, repetition level and a child reader + RepeatedReader(TypePtr, i16, i16, Box), + /// Reader of key-value pairs, e.g. maps, contains type information, definition + /// level, repetition level, child reader for keys and child reader for values + KeyValueReader(TypePtr, i16, i16, Box, Box), +} + +impl Reader { + /// Wraps reader in option reader based on repetition. + fn option(repetition: Repetition, def_level: i16, reader: Reader) -> Self { + if repetition == Repetition::OPTIONAL { + Reader::OptionReader(def_level - 1, Box::new(reader)) + } else { + reader + } + } + + /// Returns true if repeated type is an element type for the list. + /// Used to determine legacy list types. + /// This method is copied from Spark Parquet reader and is based on the reference: + /// + /// #backward-compatibility-rules + fn is_element_type(repeated_type: &Type) -> bool { + // For legacy 2-level list types whose element type is a 2-level list + // + // // ARRAY> (nullable list, non-null elements) + // optional group my_list (LIST) { + // repeated group array (LIST) { + // repeated int32 array; + // }; + // } + // + if repeated_type.is_list() || repeated_type.has_single_repeated_child() { + return false; + } + + // For legacy 2-level list types with primitive element type, e.g.: + // + // // ARRAY (nullable list, non-null elements) + // optional group my_list (LIST) { + // repeated int32 element; + // } + // + repeated_type.is_primitive() || + // For legacy 2-level list types whose element type is a group type with 2 or more + // fields, e.g.: + // + // // ARRAY> (nullable list, non-null elements) + // optional group my_list (LIST) { + // repeated group element { + // required binary str (UTF8); + // required int32 num; + // }; + // } + // + repeated_type.is_group() && repeated_type.get_fields().len() > 1 || + // For legacy 2-level list types generated by parquet-avro (Parquet version < 1.6.0), + // e.g.: + // + // // ARRAY> (nullable list, non-null elements) + // optional group my_list (LIST) { + // repeated group array { + // required binary str (UTF8); + // }; + // } + // + repeated_type.name() == "array" || + // For Parquet data generated by parquet-thrift, e.g.: + // + // // ARRAY> (nullable list, non-null elements) + // optional group my_list (LIST) { + // repeated group my_list_tuple { + // required binary str (UTF8); + // }; + // } + // + repeated_type.name().ends_with("_tuple") + } + + /// Reads current record as `Row` from the reader tree. + /// Automatically advances all necessary readers. + /// This must be called on the root level reader (i.e., for Message type). + /// Otherwise, it will panic. + fn read(&mut self) -> Result { + match *self { + Reader::GroupReader(_, _, ref mut readers) => { + let mut fields = Vec::new(); + for reader in readers { + fields.push((String::from(reader.field_name()), reader.read_field()?)); + } + Ok(Row::new(fields)) + } + _ => panic!("Cannot call read() on {self}"), + } + } + + /// Reads current record as `Field` from the reader tree. + /// Automatically advances all necessary readers. + fn read_field(&mut self) -> Result { + let field = match *self { + Reader::PrimitiveReader(_, ref mut column) => { + if !column.has_next() { + return Err(general_err!("Unexpected end of column data")); + } + let value = column.current_value()?; + column.read_next()?; + value + } + Reader::OptionReader(def_level, ref mut reader) => { + if !reader.has_next() { + return Err(general_err!("Unexpected end of column data")); + } + if reader.current_def_level() > def_level { + reader.read_field()? + } else { + reader.advance_columns()?; + Field::Null + } + } + Reader::GroupReader(_, def_level, ref mut readers) => { + let mut fields = Vec::new(); + for reader in readers { + if reader.repetition() != Repetition::OPTIONAL + || reader.current_def_level() > def_level + { + fields.push((String::from(reader.field_name()), reader.read_field()?)); + } else { + reader.advance_columns()?; + fields.push((String::from(reader.field_name()), Field::Null)); + } + } + let row = Row::new(fields); + Field::Group(row) + } + Reader::RepeatedReader(_, def_level, rep_level, ref mut reader) => { + if !reader.has_next() { + return Err(general_err!("Unexpected end of column data")); + } + let mut elements = Vec::new(); + loop { + if reader.current_def_level() > def_level { + elements.push(reader.read_field()?); + } else { + reader.advance_columns()?; + // If the current definition level is equal to the definition + // level of this repeated type, then the + // result is an empty list and the repetition level + // will always be <= rl. + break; + } + + // This covers case when we are out of repetition levels and should + // close the group, or there are no values left to + // buffer. + if !reader.has_next() || reader.current_rep_level() <= rep_level { + break; + } + } + Field::ListInternal(make_list(elements)) + } + Reader::KeyValueReader(_, def_level, rep_level, ref mut keys, ref mut values) => { + if !keys.has_next() { + return Err(general_err!("Unexpected end of column data")); + } + let mut pairs = Vec::new(); + loop { + if keys.current_def_level() > def_level { + pairs.push((keys.read_field()?, values.read_field()?)); + } else { + keys.advance_columns()?; + values.advance_columns()?; + // If the current definition level is equal to the definition + // level of this repeated type, then the + // result is an empty list and the repetition level + // will always be <= rl. + break; + } + + // This covers case when we are out of repetition levels and should + // close the group, or there are no values left to + // buffer. + if !keys.has_next() || keys.current_rep_level() <= rep_level { + break; + } + } + + Field::MapInternal(make_map(pairs)) + } + }; + Ok(field) + } + + /// Returns field name for the current reader. + fn field_name(&self) -> &str { + match *self { + Reader::PrimitiveReader(ref field, _) => field.name(), + Reader::OptionReader(_, ref reader) => reader.field_name(), + Reader::GroupReader(ref opt, ..) => match opt { + Some(field) => field.name(), + None => panic!("Field is None for group reader"), + }, + Reader::RepeatedReader(ref field, ..) => field.name(), + Reader::KeyValueReader(ref field, ..) => field.name(), + } + } + + /// Returns repetition for the current reader. + fn repetition(&self) -> Repetition { + match *self { + Reader::PrimitiveReader(ref field, _) => field.get_basic_info().repetition(), + Reader::OptionReader(_, ref reader) => reader.repetition(), + Reader::GroupReader(ref opt, ..) => match opt { + Some(field) => field.get_basic_info().repetition(), + None => panic!("Field is None for group reader"), + }, + Reader::RepeatedReader(ref field, ..) => field.get_basic_info().repetition(), + Reader::KeyValueReader(ref field, ..) => field.get_basic_info().repetition(), + } + } + + /// Returns true, if current reader has more values, false otherwise. + /// Method does not advance internal iterator. + fn has_next(&self) -> bool { + match *self { + Reader::PrimitiveReader(_, ref column) => column.has_next(), + Reader::OptionReader(_, ref reader) => reader.has_next(), + Reader::GroupReader(_, _, ref readers) => readers.first().unwrap().has_next(), + Reader::RepeatedReader(_, _, _, ref reader) => reader.has_next(), + Reader::KeyValueReader(_, _, _, ref keys, _) => keys.has_next(), + } + } + + /// Returns current definition level, + /// Method does not advance internal iterator. + fn current_def_level(&self) -> i16 { + match *self { + Reader::PrimitiveReader(_, ref column) => column.current_def_level(), + Reader::OptionReader(_, ref reader) => reader.current_def_level(), + Reader::GroupReader(_, _, ref readers) => match readers.first() { + Some(reader) => reader.current_def_level(), + None => panic!("Current definition level: empty group reader"), + }, + Reader::RepeatedReader(_, _, _, ref reader) => reader.current_def_level(), + Reader::KeyValueReader(_, _, _, ref keys, _) => keys.current_def_level(), + } + } + + /// Returns current repetition level. + /// Method does not advance internal iterator. + fn current_rep_level(&self) -> i16 { + match *self { + Reader::PrimitiveReader(_, ref column) => column.current_rep_level(), + Reader::OptionReader(_, ref reader) => reader.current_rep_level(), + Reader::GroupReader(_, _, ref readers) => match readers.first() { + Some(reader) => reader.current_rep_level(), + None => panic!("Current repetition level: empty group reader"), + }, + Reader::RepeatedReader(_, _, _, ref reader) => reader.current_rep_level(), + Reader::KeyValueReader(_, _, _, ref keys, _) => keys.current_rep_level(), + } + } + + /// Advances leaf columns for the current reader. + fn advance_columns(&mut self) -> Result<()> { + match *self { + Reader::PrimitiveReader(_, ref mut column) => column.read_next().map(|_| ()), + Reader::OptionReader(_, ref mut reader) => reader.advance_columns(), + Reader::GroupReader(_, _, ref mut readers) => { + for reader in readers { + reader.advance_columns()?; + } + Ok(()) + } + Reader::RepeatedReader(_, _, _, ref mut reader) => reader.advance_columns(), + Reader::KeyValueReader(_, _, _, ref mut keys, ref mut values) => { + keys.advance_columns()?; + values.advance_columns() + } + } + } +} + +impl fmt::Display for Reader { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let s = match self { + Reader::PrimitiveReader(..) => "PrimitiveReader", + Reader::OptionReader(..) => "OptionReader", + Reader::GroupReader(..) => "GroupReader", + Reader::RepeatedReader(..) => "RepeatedReader", + Reader::KeyValueReader(..) => "KeyValueReader", + }; + write!(f, "{s}") + } +} + +// ---------------------------------------------------------------------- +// Row iterators + +/// The enum Either with variants That represents a reference and a box of +/// [`FileReader`]. +enum Either<'a> { + Left(&'a dyn FileReader), + Right(Box), +} + +impl Either<'_> { + fn reader(&self) -> &dyn FileReader { + match *self { + Either::Left(r) => r, + Either::Right(ref r) => &**r, + } + } +} + +/// Access parquet data as an iterator of [`Row`] +/// +/// # Caveats +/// +/// Parquet stores data in a columnar fashion using [Dremel] encoding, and is therefore highly +/// optimised for reading data by column, not row. As a consequence applications concerned with +/// performance should prefer the columnar arrow or [ColumnReader] APIs. +/// +/// Additionally the current implementation does not correctly handle repeated fields ([#2394]), +/// and workloads looking to handle such schema should use the other APIs. +/// +/// [#2394]: https://github.com/apache/arrow-rs/issues/2394 +/// [ColumnReader]: crate::file::reader::RowGroupReader::get_column_reader +/// [Dremel]: https://research.google/pubs/pub36632/ +pub struct RowIter<'a> { + descr: SchemaDescPtr, + tree_builder: TreeBuilder, + file_reader: Option>, + current_row_group: usize, + num_row_groups: usize, + row_iter: Option, +} + +impl<'a> RowIter<'a> { + /// Creates a new iterator of [`Row`]s. + fn new( + file_reader: Option>, + row_iter: Option, + descr: SchemaDescPtr, + ) -> Self { + let tree_builder = Self::tree_builder(); + let num_row_groups = match file_reader { + Some(ref r) => r.reader().num_row_groups(), + None => 0, + }; + + Self { + descr, + file_reader, + tree_builder, + num_row_groups, + row_iter, + current_row_group: 0, + } + } + + /// Creates iterator of [`Row`]s for all row groups in a + /// file. + pub fn from_file(proj: Option, reader: &'a dyn FileReader) -> Result { + let either = Either::Left(reader); + let descr = + Self::get_proj_descr(proj, reader.metadata().file_metadata().schema_descr_ptr())?; + + Ok(Self::new(Some(either), None, descr)) + } + + /// Creates iterator of [`Row`]s for a specific row group. + pub fn from_row_group(proj: Option, reader: &'a dyn RowGroupReader) -> Result { + let descr = Self::get_proj_descr(proj, reader.metadata().schema_descr_ptr())?; + let tree_builder = Self::tree_builder(); + let row_iter = tree_builder.as_iter(descr.clone(), reader)?; + + // For row group we need to set `current_row_group` >= `num_row_groups`, because + // we only have one row group and can't buffer more. + Ok(Self::new(None, Some(row_iter), descr)) + } + + /// Creates a iterator of [`Row`]s from a [`FileReader`] using the full file schema. + pub fn from_file_into(reader: Box) -> Self { + let either = Either::Right(reader); + let descr = either + .reader() + .metadata() + .file_metadata() + .schema_descr_ptr(); + + Self::new(Some(either), None, descr) + } + + /// Tries to create a iterator of [`Row`]s using projections. + /// Returns a error if a file reader is not the source of this iterator. + /// + /// The Projected schema can be a subset of or equal to the file schema, + /// when it is None, full file schema is assumed. + pub fn project(self, proj: Option) -> Result { + match self.file_reader { + Some(ref either) => { + let schema = either + .reader() + .metadata() + .file_metadata() + .schema_descr_ptr(); + let descr = Self::get_proj_descr(proj, schema)?; + + Ok(Self::new(self.file_reader, None, descr)) + } + None => Err(general_err!("File reader is required to use projections")), + } + } + + /// Helper method to get schema descriptor for projected schema. + /// If projection is None, then full schema is returned. + #[inline] + fn get_proj_descr(proj: Option, root_descr: SchemaDescPtr) -> Result { + match proj { + Some(projection) => { + // check if projection is part of file schema + let root_schema = root_descr.root_schema(); + if !root_schema.check_contains(&projection) { + return Err(general_err!("Root schema does not contain projection")); + } + Ok(Arc::new(SchemaDescriptor::new(Arc::new(projection)))) + } + None => Ok(root_descr), + } + } + + /// Sets batch size for this row iter. + pub fn with_batch_size(mut self, batch_size: usize) -> Self { + self.tree_builder = self.tree_builder.with_batch_size(batch_size); + self + } + + /// Returns common tree builder, so the same settings are applied to both iterators + /// from file reader and row group. + #[inline] + fn tree_builder() -> TreeBuilder { + TreeBuilder::new() + } +} + +impl Iterator for RowIter<'_> { + type Item = Result; + + fn next(&mut self) -> Option> { + let mut row = None; + if let Some(ref mut iter) = self.row_iter { + row = iter.next(); + } + + while row.is_none() && self.current_row_group < self.num_row_groups { + // We do not expect any failures when accessing a row group, and file reader + // must be set for selecting next row group. + if let Some(ref either) = self.file_reader { + let file_reader = either.reader(); + let row_group_reader = &*file_reader + .get_row_group(self.current_row_group) + .expect("Row group is required to advance"); + + match self + .tree_builder + .as_iter(self.descr.clone(), row_group_reader) + { + Ok(mut iter) => { + row = iter.next(); + + self.current_row_group += 1; + self.row_iter = Some(iter); + } + Err(e) => return Some(Err(e)), + } + } + } + + row + } +} + +/// Internal iterator of [`Row`]s for a reader. +pub struct ReaderIter { + root_reader: Reader, + records_left: usize, +} + +impl ReaderIter { + fn new(mut root_reader: Reader, num_records: usize) -> Result { + // Prepare root reader by advancing all column vectors + root_reader.advance_columns()?; + Ok(Self { + root_reader, + records_left: num_records, + }) + } +} + +impl Iterator for ReaderIter { + type Item = Result; + + fn next(&mut self) -> Option> { + if self.records_left > 0 { + self.records_left -= 1; + Some(self.root_reader.read()) + } else { + None + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use crate::data_type::Int64Type; + use crate::file::reader::SerializedFileReader; + use crate::file::writer::SerializedFileWriter; + use crate::record::api::RowAccessor; + use crate::schema::parser::parse_message_type; + use crate::util::test_common::file_util::{get_test_file, get_test_path}; + use bytes::Bytes; + + // Convenient macros to assemble row, list, map, and group. + + macro_rules! row { + ($($e:tt)*) => { + { + Row::new(vec![$($e)*]) + } + } + } + + macro_rules! list { + ($($e:tt)*) => { + { + Field::ListInternal(make_list(vec![$($e)*])) + } + } + } + + macro_rules! map { + ($($e:tt)*) => { + { + Field::MapInternal(make_map(vec![$($e)*])) + } + } + } + + macro_rules! group { + ( $( $e:expr ), * ) => { + { + Field::Group(row!($( $e ), *)) + } + } + } + + #[test] + fn test_file_reader_rows_nulls() { + let rows = test_file_reader_rows("nulls.snappy.parquet", None).unwrap(); + let expected_rows = vec![ + row![( + "b_struct".to_string(), + group![("b_c_int".to_string(), Field::Null)] + )], + row![( + "b_struct".to_string(), + group![("b_c_int".to_string(), Field::Null)] + )], + row![( + "b_struct".to_string(), + group![("b_c_int".to_string(), Field::Null)] + )], + row![( + "b_struct".to_string(), + group![("b_c_int".to_string(), Field::Null)] + )], + row![( + "b_struct".to_string(), + group![("b_c_int".to_string(), Field::Null)] + )], + row![( + "b_struct".to_string(), + group![("b_c_int".to_string(), Field::Null)] + )], + row![( + "b_struct".to_string(), + group![("b_c_int".to_string(), Field::Null)] + )], + row![( + "b_struct".to_string(), + group![("b_c_int".to_string(), Field::Null)] + )], + ]; + assert_eq!(rows, expected_rows); + } + + #[test] + fn test_file_reader_rows_nonnullable() { + let rows = test_file_reader_rows("nonnullable.impala.parquet", None).unwrap(); + let expected_rows = vec![row![ + ("ID".to_string(), Field::Long(8)), + ("Int_Array".to_string(), list![Field::Int(-1)]), + ( + "int_array_array".to_string(), + list![list![Field::Int(-1), Field::Int(-2)], list![]] + ), + ( + "Int_Map".to_string(), + map![(Field::Str("k1".to_string()), Field::Int(-1))] + ), + ( + "int_map_array".to_string(), + list![ + map![], + map![(Field::Str("k1".to_string()), Field::Int(1))], + map![], + map![] + ] + ), + ( + "nested_Struct".to_string(), + group![ + ("a".to_string(), Field::Int(-1)), + ("B".to_string(), list![Field::Int(-1)]), + ( + "c".to_string(), + group![( + "D".to_string(), + list![list![group![ + ("e".to_string(), Field::Int(-1)), + ("f".to_string(), Field::Str("nonnullable".to_string())) + ]]] + )] + ), + ("G".to_string(), map![]) + ] + ) + ]]; + assert_eq!(rows, expected_rows); + } + + #[test] + fn test_file_reader_rows_nullable() { + let rows = test_file_reader_rows("nullable.impala.parquet", None).unwrap(); + let expected_rows = vec![ + row![ + ("id".to_string(), Field::Long(1)), + ( + "int_array".to_string(), + list![Field::Int(1), Field::Int(2), Field::Int(3)] + ), + ( + "int_array_Array".to_string(), + list![ + list![Field::Int(1), Field::Int(2)], + list![Field::Int(3), Field::Int(4)] + ] + ), + ( + "int_map".to_string(), + map![ + (Field::Str("k1".to_string()), Field::Int(1)), + (Field::Str("k2".to_string()), Field::Int(100)) + ] + ), + ( + "int_Map_Array".to_string(), + list![map![(Field::Str("k1".to_string()), Field::Int(1))]] + ), + ( + "nested_struct".to_string(), + group![ + ("A".to_string(), Field::Int(1)), + ("b".to_string(), list![Field::Int(1)]), + ( + "C".to_string(), + group![( + "d".to_string(), + list![ + list![ + group![ + ("E".to_string(), Field::Int(10)), + ("F".to_string(), Field::Str("aaa".to_string())) + ], + group![ + ("E".to_string(), Field::Int(-10)), + ("F".to_string(), Field::Str("bbb".to_string())) + ] + ], + list![group![ + ("E".to_string(), Field::Int(11)), + ("F".to_string(), Field::Str("c".to_string())) + ]] + ] + )] + ), + ( + "g".to_string(), + map![( + Field::Str("foo".to_string()), + group![( + "H".to_string(), + group![("i".to_string(), list![Field::Double(1.1)])] + )] + )] + ) + ] + ) + ], + row![ + ("id".to_string(), Field::Long(2)), + ( + "int_array".to_string(), + list![ + Field::Null, + Field::Int(1), + Field::Int(2), + Field::Null, + Field::Int(3), + Field::Null + ] + ), + ( + "int_array_Array".to_string(), + list![ + list![Field::Null, Field::Int(1), Field::Int(2), Field::Null], + list![Field::Int(3), Field::Null, Field::Int(4)], + list![], + Field::Null + ] + ), + ( + "int_map".to_string(), + map![ + (Field::Str("k1".to_string()), Field::Int(2)), + (Field::Str("k2".to_string()), Field::Null) + ] + ), + ( + "int_Map_Array".to_string(), + list![ + map![ + (Field::Str("k3".to_string()), Field::Null), + (Field::Str("k1".to_string()), Field::Int(1)) + ], + Field::Null, + map![] + ] + ), + ( + "nested_struct".to_string(), + group![ + ("A".to_string(), Field::Null), + ("b".to_string(), list![Field::Null]), + ( + "C".to_string(), + group![( + "d".to_string(), + list![ + list![ + group![ + ("E".to_string(), Field::Null), + ("F".to_string(), Field::Null) + ], + group![ + ("E".to_string(), Field::Int(10)), + ("F".to_string(), Field::Str("aaa".to_string())) + ], + group![ + ("E".to_string(), Field::Null), + ("F".to_string(), Field::Null) + ], + group![ + ("E".to_string(), Field::Int(-10)), + ("F".to_string(), Field::Str("bbb".to_string())) + ], + group![ + ("E".to_string(), Field::Null), + ("F".to_string(), Field::Null) + ] + ], + list![ + group![ + ("E".to_string(), Field::Int(11)), + ("F".to_string(), Field::Str("c".to_string())) + ], + Field::Null + ], + list![], + Field::Null + ] + )] + ), + ( + "g".to_string(), + map![ + ( + Field::Str("g1".to_string()), + group![( + "H".to_string(), + group![( + "i".to_string(), + list![Field::Double(2.2), Field::Null] + )] + )] + ), + ( + Field::Str("g2".to_string()), + group![("H".to_string(), group![("i".to_string(), list![])])] + ), + (Field::Str("g3".to_string()), Field::Null), + ( + Field::Str("g4".to_string()), + group![( + "H".to_string(), + group![("i".to_string(), Field::Null)] + )] + ), + ( + Field::Str("g5".to_string()), + group![("H".to_string(), Field::Null)] + ) + ] + ) + ] + ) + ], + row![ + ("id".to_string(), Field::Long(3)), + ("int_array".to_string(), list![]), + ("int_array_Array".to_string(), list![Field::Null]), + ("int_map".to_string(), map![]), + ("int_Map_Array".to_string(), list![Field::Null, Field::Null]), + ( + "nested_struct".to_string(), + group![ + ("A".to_string(), Field::Null), + ("b".to_string(), Field::Null), + ("C".to_string(), group![("d".to_string(), list![])]), + ("g".to_string(), map![]) + ] + ) + ], + row![ + ("id".to_string(), Field::Long(4)), + ("int_array".to_string(), Field::Null), + ("int_array_Array".to_string(), list![]), + ("int_map".to_string(), map![]), + ("int_Map_Array".to_string(), list![]), + ( + "nested_struct".to_string(), + group![ + ("A".to_string(), Field::Null), + ("b".to_string(), Field::Null), + ("C".to_string(), group![("d".to_string(), Field::Null)]), + ("g".to_string(), Field::Null) + ] + ) + ], + row![ + ("id".to_string(), Field::Long(5)), + ("int_array".to_string(), Field::Null), + ("int_array_Array".to_string(), Field::Null), + ("int_map".to_string(), map![]), + ("int_Map_Array".to_string(), Field::Null), + ( + "nested_struct".to_string(), + group![ + ("A".to_string(), Field::Null), + ("b".to_string(), Field::Null), + ("C".to_string(), Field::Null), + ( + "g".to_string(), + map![( + Field::Str("foo".to_string()), + group![( + "H".to_string(), + group![( + "i".to_string(), + list![Field::Double(2.2), Field::Double(3.3)] + )] + )] + )] + ) + ] + ) + ], + row![ + ("id".to_string(), Field::Long(6)), + ("int_array".to_string(), Field::Null), + ("int_array_Array".to_string(), Field::Null), + ("int_map".to_string(), Field::Null), + ("int_Map_Array".to_string(), Field::Null), + ("nested_struct".to_string(), Field::Null) + ], + row![ + ("id".to_string(), Field::Long(7)), + ("int_array".to_string(), Field::Null), + ( + "int_array_Array".to_string(), + list![Field::Null, list![Field::Int(5), Field::Int(6)]] + ), + ( + "int_map".to_string(), + map![ + (Field::Str("k1".to_string()), Field::Null), + (Field::Str("k3".to_string()), Field::Null) + ] + ), + ("int_Map_Array".to_string(), Field::Null), + ( + "nested_struct".to_string(), + group![ + ("A".to_string(), Field::Int(7)), + ( + "b".to_string(), + list![Field::Int(2), Field::Int(3), Field::Null] + ), + ( + "C".to_string(), + group![( + "d".to_string(), + list![list![], list![Field::Null], Field::Null] + )] + ), + ("g".to_string(), Field::Null) + ] + ) + ], + ]; + assert_eq!(rows, expected_rows); + } + + #[test] + fn test_file_reader_rows_projection() { + let schema = " + message spark_schema { + REQUIRED DOUBLE c; + REQUIRED INT32 b; + } + "; + let schema = parse_message_type(schema).unwrap(); + let rows = test_file_reader_rows("nested_maps.snappy.parquet", Some(schema)).unwrap(); + let expected_rows = vec![ + row![ + ("c".to_string(), Field::Double(1.0)), + ("b".to_string(), Field::Int(1)) + ], + row![ + ("c".to_string(), Field::Double(1.0)), + ("b".to_string(), Field::Int(1)) + ], + row![ + ("c".to_string(), Field::Double(1.0)), + ("b".to_string(), Field::Int(1)) + ], + row![ + ("c".to_string(), Field::Double(1.0)), + ("b".to_string(), Field::Int(1)) + ], + row![ + ("c".to_string(), Field::Double(1.0)), + ("b".to_string(), Field::Int(1)) + ], + row![ + ("c".to_string(), Field::Double(1.0)), + ("b".to_string(), Field::Int(1)) + ], + ]; + assert_eq!(rows, expected_rows); + } + + #[test] + fn test_iter_columns_in_row() { + let r = row![ + ("c".to_string(), Field::Double(1.0)), + ("b".to_string(), Field::Int(1)) + ]; + let mut result = Vec::new(); + for (name, record) in r.get_column_iter() { + result.push((name, record)); + } + assert_eq!( + vec![ + (&"c".to_string(), &Field::Double(1.0)), + (&"b".to_string(), &Field::Int(1)) + ], + result + ); + } + + #[test] + fn test_into_columns_in_row() { + let r = row![ + ("a".to_string(), Field::Str("My string".to_owned())), + ("b".to_string(), Field::Int(1)) + ]; + assert_eq!( + r.into_columns(), + vec![ + ("a".to_string(), Field::Str("My string".to_owned())), + ("b".to_string(), Field::Int(1)), + ] + ); + } + + #[test] + fn test_file_reader_rows_projection_map() { + let schema = " + message spark_schema { + OPTIONAL group a (MAP) { + REPEATED group key_value { + REQUIRED BYTE_ARRAY key (UTF8); + OPTIONAL group value (MAP) { + REPEATED group key_value { + REQUIRED INT32 key; + REQUIRED BOOLEAN value; + } + } + } + } + } + "; + let schema = parse_message_type(schema).unwrap(); + let rows = test_file_reader_rows("nested_maps.snappy.parquet", Some(schema)).unwrap(); + let expected_rows = vec![ + row![( + "a".to_string(), + map![( + Field::Str("a".to_string()), + map![ + (Field::Int(1), Field::Bool(true)), + (Field::Int(2), Field::Bool(false)) + ] + )] + )], + row![( + "a".to_string(), + map![( + Field::Str("b".to_string()), + map![(Field::Int(1), Field::Bool(true))] + )] + )], + row![( + "a".to_string(), + map![(Field::Str("c".to_string()), Field::Null)] + )], + row![("a".to_string(), map![(Field::Str("d".to_string()), map![])])], + row![( + "a".to_string(), + map![( + Field::Str("e".to_string()), + map![(Field::Int(1), Field::Bool(true))] + )] + )], + row![( + "a".to_string(), + map![( + Field::Str("f".to_string()), + map![ + (Field::Int(3), Field::Bool(true)), + (Field::Int(4), Field::Bool(false)), + (Field::Int(5), Field::Bool(true)) + ] + )] + )], + ]; + assert_eq!(rows, expected_rows); + } + + #[test] + fn test_file_reader_rows_projection_list() { + let schema = " + message spark_schema { + OPTIONAL group a (LIST) { + REPEATED group list { + OPTIONAL group element (LIST) { + REPEATED group list { + OPTIONAL group element (LIST) { + REPEATED group list { + OPTIONAL BYTE_ARRAY element (UTF8); + } + } + } + } + } + } + } + "; + let schema = parse_message_type(schema).unwrap(); + let rows = test_file_reader_rows("nested_lists.snappy.parquet", Some(schema)).unwrap(); + let expected_rows = vec![ + row![( + "a".to_string(), + list![ + list![ + list![Field::Str("a".to_string()), Field::Str("b".to_string())], + list![Field::Str("c".to_string())] + ], + list![Field::Null, list![Field::Str("d".to_string())]] + ] + )], + row![( + "a".to_string(), + list![ + list![ + list![Field::Str("a".to_string()), Field::Str("b".to_string())], + list![Field::Str("c".to_string()), Field::Str("d".to_string())] + ], + list![Field::Null, list![Field::Str("e".to_string())]] + ] + )], + row![( + "a".to_string(), + list![ + list![ + list![Field::Str("a".to_string()), Field::Str("b".to_string())], + list![Field::Str("c".to_string()), Field::Str("d".to_string())], + list![Field::Str("e".to_string())] + ], + list![Field::Null, list![Field::Str("f".to_string())]] + ] + )], + ]; + assert_eq!(rows, expected_rows); + } + + #[test] + fn test_file_reader_rows_invalid_projection() { + let schema = " + message spark_schema { + REQUIRED INT32 key; + REQUIRED BOOLEAN value; + } + "; + let schema = parse_message_type(schema).unwrap(); + let res = test_file_reader_rows("nested_maps.snappy.parquet", Some(schema)); + assert_eq!( + res.unwrap_err().to_string(), + "Parquet error: Root schema does not contain projection" + ); + } + + #[test] + fn test_row_group_rows_invalid_projection() { + let schema = " + message spark_schema { + REQUIRED INT32 key; + REQUIRED BOOLEAN value; + } + "; + let schema = parse_message_type(schema).unwrap(); + let res = test_row_group_rows("nested_maps.snappy.parquet", Some(schema)); + assert_eq!( + res.unwrap_err().to_string(), + "Parquet error: Root schema does not contain projection" + ); + } + + #[test] + fn test_file_reader_rows_nested_map_type() { + let schema = " + message spark_schema { + OPTIONAL group a (MAP) { + REPEATED group key_value { + REQUIRED BYTE_ARRAY key (UTF8); + OPTIONAL group value (MAP) { + REPEATED group key_value { + REQUIRED INT32 key; + } + } + } + } + } + "; + let schema = parse_message_type(schema).unwrap(); + test_file_reader_rows("nested_maps.snappy.parquet", Some(schema)).unwrap(); + } + + #[test] + fn test_file_reader_iter() { + let path = get_test_path("alltypes_plain.parquet"); + let reader = SerializedFileReader::try_from(path.as_path()).unwrap(); + let iter = RowIter::from_file_into(Box::new(reader)); + + let values: Vec<_> = iter.flat_map(|r| r.unwrap().get_int(0)).collect(); + assert_eq!(values, &[4, 5, 6, 7, 2, 3, 0, 1]); + } + + #[test] + fn test_file_reader_iter_projection() { + let path = get_test_path("alltypes_plain.parquet"); + let reader = SerializedFileReader::try_from(path.as_path()).unwrap(); + let schema = "message schema { OPTIONAL INT32 id; }"; + let proj = parse_message_type(schema).ok(); + + let iter = RowIter::from_file_into(Box::new(reader)) + .project(proj) + .unwrap(); + let values: Vec<_> = iter.flat_map(|r| r.unwrap().get_int(0)).collect(); + + assert_eq!(values, &[4, 5, 6, 7, 2, 3, 0, 1]); + } + + #[test] + fn test_file_reader_iter_projection_err() { + let schema = " + message spark_schema { + REQUIRED INT32 key; + REQUIRED BOOLEAN value; + } + "; + let proj = parse_message_type(schema).ok(); + let path = get_test_path("nested_maps.snappy.parquet"); + let reader = SerializedFileReader::try_from(path.as_path()).unwrap(); + let res = RowIter::from_file_into(Box::new(reader)).project(proj); + + assert_eq!( + res.err().unwrap().to_string(), + "Parquet error: Root schema does not contain projection" + ); + } + + #[test] + fn test_tree_reader_handle_repeated_fields_with_no_annotation() { + // Array field `phoneNumbers` does not contain LIST annotation. + // We parse it as struct with `phone` repeated field as array. + let rows = test_file_reader_rows("repeated_no_annotation.parquet", None).unwrap(); + let expected_rows = vec![ + row![ + ("id".to_string(), Field::Int(1)), + ("phoneNumbers".to_string(), Field::Null) + ], + row![ + ("id".to_string(), Field::Int(2)), + ("phoneNumbers".to_string(), Field::Null) + ], + row![ + ("id".to_string(), Field::Int(3)), + ( + "phoneNumbers".to_string(), + group![("phone".to_string(), list![])] + ) + ], + row![ + ("id".to_string(), Field::Int(4)), + ( + "phoneNumbers".to_string(), + group![( + "phone".to_string(), + list![group![ + ("number".to_string(), Field::Long(5555555555)), + ("kind".to_string(), Field::Null) + ]] + )] + ) + ], + row![ + ("id".to_string(), Field::Int(5)), + ( + "phoneNumbers".to_string(), + group![( + "phone".to_string(), + list![group![ + ("number".to_string(), Field::Long(1111111111)), + ("kind".to_string(), Field::Str("home".to_string())) + ]] + )] + ) + ], + row![ + ("id".to_string(), Field::Int(6)), + ( + "phoneNumbers".to_string(), + group![( + "phone".to_string(), + list![ + group![ + ("number".to_string(), Field::Long(1111111111)), + ("kind".to_string(), Field::Str("home".to_string())) + ], + group![ + ("number".to_string(), Field::Long(2222222222)), + ("kind".to_string(), Field::Null) + ], + group![ + ("number".to_string(), Field::Long(3333333333)), + ("kind".to_string(), Field::Str("mobile".to_string())) + ] + ] + )] + ) + ], + ]; + + assert_eq!(rows, expected_rows); + } + + #[test] + fn test_tree_reader_handle_nested_repeated_fields_with_no_annotation() { + // Create schema + let schema = Arc::new( + parse_message_type( + " + message schema { + REPEATED group level1 { + REPEATED group level2 { + REQUIRED group level3 { + REQUIRED INT64 value3; + } + } + REQUIRED INT64 value1; + } + }", + ) + .unwrap(), + ); + + // Write Parquet file to buffer + let mut buffer: Vec = Vec::new(); + let mut file_writer = + SerializedFileWriter::new(&mut buffer, schema, Default::default()).unwrap(); + let mut row_group_writer = file_writer.next_row_group().unwrap(); + + // Write column level1.level2.level3.value3 + let mut column_writer = row_group_writer.next_column().unwrap().unwrap(); + column_writer + .typed::() + .write_batch(&[30, 31, 32], Some(&[2, 2, 2]), Some(&[0, 0, 0])) + .unwrap(); + column_writer.close().unwrap(); + + // Write column level1.value1 + let mut column_writer = row_group_writer.next_column().unwrap().unwrap(); + column_writer + .typed::() + .write_batch(&[10, 11, 12], Some(&[1, 1, 1]), Some(&[0, 0, 0])) + .unwrap(); + column_writer.close().unwrap(); + + // Finalize Parquet file + row_group_writer.close().unwrap(); + file_writer.close().unwrap(); + assert_eq!(&buffer[0..4], b"PAR1"); + + // Read Parquet file from buffer + let file_reader = SerializedFileReader::new(Bytes::from(buffer)).unwrap(); + let rows: Vec<_> = file_reader + .get_row_iter(None) + .unwrap() + .map(|row| row.unwrap()) + .collect(); + + let expected_rows = vec![ + row![( + "level1".to_string(), + list![group![ + ( + "level2".to_string(), + list![group![( + "level3".to_string(), + group![("value3".to_string(), Field::Long(30))] + )]] + ), + ("value1".to_string(), Field::Long(10)) + ]] + )], + row![( + "level1".to_string(), + list![group![ + ( + "level2".to_string(), + list![group![( + "level3".to_string(), + group![("value3".to_string(), Field::Long(31))] + )]] + ), + ("value1".to_string(), Field::Long(11)) + ]] + )], + row![( + "level1".to_string(), + list![group![ + ( + "level2".to_string(), + list![group![( + "level3".to_string(), + group![("value3".to_string(), Field::Long(32))] + )]] + ), + ("value1".to_string(), Field::Long(12)) + ]] + )], + ]; + + assert_eq!(rows, expected_rows); + } + + #[test] + fn test_tree_reader_handle_primitive_repeated_fields_with_no_annotation() { + // In this test the REPEATED fields are primitives + let rows = test_file_reader_rows("repeated_primitive_no_list.parquet", None).unwrap(); + let expected_rows = vec![ + row![ + ( + "Int32_list".to_string(), + Field::ListInternal(make_list([0, 1, 2, 3].map(Field::Int).to_vec())) + ), + ( + "String_list".to_string(), + Field::ListInternal(make_list( + ["foo", "zero", "one", "two"] + .map(|s| Field::Str(s.to_string())) + .to_vec() + )) + ), + ( + "group_of_lists".to_string(), + group![ + ( + "Int32_list_in_group".to_string(), + Field::ListInternal(make_list([0, 1, 2, 3].map(Field::Int).to_vec())) + ), + ( + "String_list_in_group".to_string(), + Field::ListInternal(make_list( + ["foo", "zero", "one", "two"] + .map(|s| Field::Str(s.to_string())) + .to_vec() + )) + ) + ] + ) + ], + row![ + ( + "Int32_list".to_string(), + Field::ListInternal(make_list(vec![])) + ), + ( + "String_list".to_string(), + Field::ListInternal(make_list( + ["three"].map(|s| Field::Str(s.to_string())).to_vec() + )) + ), + ( + "group_of_lists".to_string(), + group![ + ( + "Int32_list_in_group".to_string(), + Field::ListInternal(make_list(vec![])) + ), + ( + "String_list_in_group".to_string(), + Field::ListInternal(make_list( + ["three"].map(|s| Field::Str(s.to_string())).to_vec() + )) + ) + ] + ) + ], + row![ + ( + "Int32_list".to_string(), + Field::ListInternal(make_list(vec![Field::Int(4)])) + ), + ( + "String_list".to_string(), + Field::ListInternal(make_list( + ["four"].map(|s| Field::Str(s.to_string())).to_vec() + )) + ), + ( + "group_of_lists".to_string(), + group![ + ( + "Int32_list_in_group".to_string(), + Field::ListInternal(make_list(vec![Field::Int(4)])) + ), + ( + "String_list_in_group".to_string(), + Field::ListInternal(make_list( + ["four"].map(|s| Field::Str(s.to_string())).to_vec() + )) + ) + ] + ) + ], + row![ + ( + "Int32_list".to_string(), + Field::ListInternal(make_list([5, 6, 7, 8].map(Field::Int).to_vec())) + ), + ( + "String_list".to_string(), + Field::ListInternal(make_list( + ["five", "six", "seven", "eight"] + .map(|s| Field::Str(s.to_string())) + .to_vec() + )) + ), + ( + "group_of_lists".to_string(), + group![ + ( + "Int32_list_in_group".to_string(), + Field::ListInternal(make_list([5, 6, 7, 8].map(Field::Int).to_vec())) + ), + ( + "String_list_in_group".to_string(), + Field::ListInternal(make_list( + ["five", "six", "seven", "eight"] + .map(|s| Field::Str(s.to_string())) + .to_vec() + )) + ) + ] + ) + ], + ]; + assert_eq!(rows, expected_rows); + } + + #[test] + fn test_map_no_value() { + // File schema: + // message schema { + // required group my_map (MAP) { + // repeated group key_value { + // required int32 key; + // optional int32 value; + // } + // } + // required group my_map_no_v (MAP) { + // repeated group key_value { + // required int32 key; + // } + // } + // required group my_list (LIST) { + // repeated group list { + // required int32 element; + // } + // } + // } + let rows = test_file_reader_rows("map_no_value.parquet", None).unwrap(); + + // the my_map_no_v and my_list columns should be equivalent lists by this point + for row in rows { + let cols = row.into_columns(); + assert_eq!(cols[1].1, cols[2].1); + } + } + + 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)?); + let iter = file_reader.get_row_iter(schema)?; + Ok(iter.map(|row| row.unwrap()).collect()) + } + + fn test_row_group_rows(file_name: &str, schema: Option) -> Result> { + let file = get_test_file(file_name); + let file_reader: Box = Box::new(SerializedFileReader::new(file)?); + // Check the first row group only, because files will contain only single row + // group + let row_group_reader = file_reader.get_row_group(0).unwrap(); + let iter = row_group_reader.get_row_iter(schema)?; + Ok(iter.map(|row| row.unwrap()).collect()) + } + + #[test] + fn test_read_old_nested_list() { + let rows = test_file_reader_rows("old_list_structure.parquet", None).unwrap(); + let expected_rows = vec![row![( + "a".to_string(), + Field::ListInternal(make_list( + [ + make_list([1, 2].map(Field::Int).to_vec()), + make_list([3, 4].map(Field::Int).to_vec()) + ] + .map(Field::ListInternal) + .to_vec() + )) + ),]]; + assert_eq!(rows, expected_rows); + } + + fn assert_err_on_overcount(file_name: &str, proj_schema: Option) { + let file = get_test_file(file_name); + let file_reader = SerializedFileReader::new(file).unwrap(); + let metadata = file_reader.metadata(); + let row_group_reader = file_reader.get_row_group(0).unwrap(); + let actual_rows = row_group_reader.metadata().num_rows() as usize; + + let descr = match proj_schema { + Some(schema) => Arc::new(SchemaDescriptor::new(Arc::new(schema))), + None => metadata.file_metadata().schema_descr_ptr(), + }; + let reader = TreeBuilder::new().build(descr, &*row_group_reader).unwrap(); + let iter = ReaderIter::new(reader, actual_rows + 1).unwrap(); + + let rows: Vec> = iter.collect(); + assert_eq!(rows.len(), actual_rows + 1); + for row in &rows[..actual_rows] { + assert!(row.is_ok(), "Expected Ok row, got: {:?}", row); + } + let err = rows[actual_rows].as_ref().unwrap_err(); + assert!( + err.to_string().contains("Unexpected end of column data"), + "Unexpected error message: {}", + err + ); + } + + #[test] + fn test_reader_iter_returns_error_when_num_records_exceeds_data() { + assert_err_on_overcount("nulls.snappy.parquet", None); + } + + #[test] + fn test_reader_iter_returns_error_for_repeated_field_when_num_records_exceeds_data() { + assert_err_on_overcount("repeated_primitive_no_list.parquet", None); + } + + #[test] + fn test_reader_iter_returns_error_for_map_field_when_num_records_exceeds_data() { + let schema = parse_message_type( + "message schema { + REQUIRED group my_map (MAP) { + REPEATED group key_value { + REQUIRED INT32 key; + OPTIONAL INT32 value; + } + } + }", + ) + .unwrap(); + assert_err_on_overcount("map_no_value.parquet", Some(schema)); + } +} diff --git a/vendor/parquet/src/record/record_reader.rs b/vendor/parquet/src/record/record_reader.rs new file mode 100644 index 0000000..75ca4e3 --- /dev/null +++ b/vendor/parquet/src/record/record_reader.rs @@ -0,0 +1,32 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use super::super::errors::ParquetError; +use super::super::file::reader::RowGroupReader; + +/// Read up to `num_records` records from `row_group_reader` into `self`. +/// +/// The type parameter `T` is used to work around the rust orphan rule +/// when implementing on types such as `Vec`. +pub trait RecordReader { + /// Read up to `num_records` records from `row_group_reader` into `self`. + fn read_from_row_group( + &mut self, + row_group_reader: &mut dyn RowGroupReader, + num_records: usize, + ) -> Result<(), ParquetError>; +} diff --git a/vendor/parquet/src/record/record_writer.rs b/vendor/parquet/src/record/record_writer.rs new file mode 100644 index 0000000..56e0aa4 --- /dev/null +++ b/vendor/parquet/src/record/record_writer.rs @@ -0,0 +1,42 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::schema::types::TypePtr; + +use super::super::errors::ParquetError; +use super::super::file::writer::SerializedRowGroupWriter; + +/// Trait describing how to write a record (the implementator) to a row group writer. +/// +/// [`parquet_derive`] crate provides a derive macro [`ParquetRecordWriter`] for this trait +/// for unnested structs. +/// +/// The type parameter `T` is used to work around the rust orphan rule +/// when implementing on types such as `&[T]`. +/// +/// [`parquet_derive`]: https://crates.io/crates/parquet_derive +/// [`ParquetRecordWriter`]: https://docs.rs/parquet_derive/53.0.0/parquet_derive/derive.ParquetRecordWriter.html +pub trait RecordWriter { + /// Writes from `self` into `row_group_writer`. + fn write_to_row_group( + &self, + row_group_writer: &mut SerializedRowGroupWriter, + ) -> Result<(), ParquetError>; + + /// Generated schema used by `row_group_writer` + fn schema(&self) -> Result; +} diff --git a/vendor/parquet/src/record/triplet.rs b/vendor/parquet/src/record/triplet.rs new file mode 100644 index 0000000..b4d39bb --- /dev/null +++ b/vendor/parquet/src/record/triplet.rs @@ -0,0 +1,608 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::basic::Type as PhysicalType; +use crate::column::reader::{ColumnReader, ColumnReaderImpl, get_typed_column_reader}; +use crate::data_type::*; +use crate::errors::{ParquetError, Result}; +use crate::record::api::Field; +use crate::schema::types::ColumnDescPtr; + +/// Macro to generate simple functions that cover all types of triplet iterator. +/// $func is a function of a typed triplet iterator and $token is a either {`ref`} or +/// {`ref`, `mut`} +macro_rules! triplet_enum_func { + ($self:ident, $func:ident, $( $token:tt ),*) => ({ + match *$self { + TripletIter::BoolTripletIter($($token)* typed) => typed.$func(), + TripletIter::Int32TripletIter($($token)* typed) => typed.$func(), + TripletIter::Int64TripletIter($($token)* typed) => typed.$func(), + TripletIter::Int96TripletIter($($token)* typed) => typed.$func(), + TripletIter::FloatTripletIter($($token)* typed) => typed.$func(), + TripletIter::DoubleTripletIter($($token)* typed) => typed.$func(), + TripletIter::ByteArrayTripletIter($($token)* typed) => typed.$func(), + TripletIter::FixedLenByteArrayTripletIter($($token)* typed) => typed.$func() + } + }); +} + +/// High level API wrapper on column reader. +/// Provides per-element access for each primitive column. +#[allow(clippy::enum_variant_names)] +pub enum TripletIter { + BoolTripletIter(TypedTripletIter), + Int32TripletIter(TypedTripletIter), + Int64TripletIter(TypedTripletIter), + Int96TripletIter(TypedTripletIter), + FloatTripletIter(TypedTripletIter), + DoubleTripletIter(TypedTripletIter), + ByteArrayTripletIter(TypedTripletIter), + FixedLenByteArrayTripletIter(TypedTripletIter), +} + +impl TripletIter { + /// Creates new triplet for column reader + pub fn new(descr: ColumnDescPtr, reader: ColumnReader, batch_size: usize) -> Self { + match descr.physical_type() { + PhysicalType::BOOLEAN => { + TripletIter::BoolTripletIter(TypedTripletIter::new(descr, batch_size, reader)) + } + PhysicalType::INT32 => { + TripletIter::Int32TripletIter(TypedTripletIter::new(descr, batch_size, reader)) + } + PhysicalType::INT64 => { + TripletIter::Int64TripletIter(TypedTripletIter::new(descr, batch_size, reader)) + } + PhysicalType::INT96 => { + TripletIter::Int96TripletIter(TypedTripletIter::new(descr, batch_size, reader)) + } + PhysicalType::FLOAT => { + TripletIter::FloatTripletIter(TypedTripletIter::new(descr, batch_size, reader)) + } + PhysicalType::DOUBLE => { + TripletIter::DoubleTripletIter(TypedTripletIter::new(descr, batch_size, reader)) + } + PhysicalType::BYTE_ARRAY => { + TripletIter::ByteArrayTripletIter(TypedTripletIter::new(descr, batch_size, reader)) + } + PhysicalType::FIXED_LEN_BYTE_ARRAY => TripletIter::FixedLenByteArrayTripletIter( + TypedTripletIter::new(descr, batch_size, reader), + ), + } + } + + /// Invokes underlying typed triplet iterator to buffer current value. + /// Should be called once - either before `is_null` or `current_value`. + #[inline] + pub fn read_next(&mut self) -> Result { + triplet_enum_func!(self, read_next, ref, mut) + } + + /// Provides check on values/levels left without invoking the underlying typed triplet + /// iterator. + /// Returns true if more values/levels exist, false otherwise. + /// It is always in sync with `read_next` method. + #[inline] + pub fn has_next(&self) -> bool { + triplet_enum_func!(self, has_next, ref) + } + + /// Returns current definition level for a leaf triplet iterator + #[inline] + pub fn current_def_level(&self) -> i16 { + triplet_enum_func!(self, current_def_level, ref) + } + + /// Returns max definition level for a leaf triplet iterator + #[inline] + pub fn max_def_level(&self) -> i16 { + triplet_enum_func!(self, max_def_level, ref) + } + + /// Returns current repetition level for a leaf triplet iterator + #[inline] + pub fn current_rep_level(&self) -> i16 { + triplet_enum_func!(self, current_rep_level, ref) + } + + /// Returns max repetition level for a leaf triplet iterator + #[inline] + pub fn max_rep_level(&self) -> i16 { + triplet_enum_func!(self, max_rep_level, ref) + } + + /// Returns true, if current value is null. + /// Based on the fact that for non-null value current definition level + /// equals to max definition level. + #[inline] + pub fn is_null(&self) -> bool { + self.current_def_level() < self.max_def_level() + } + + /// Updates non-null value for current row. + pub fn current_value(&self) -> Result { + if self.is_null() { + return Ok(Field::Null); + } + let field = match *self { + TripletIter::BoolTripletIter(ref typed) => { + Field::convert_bool(typed.column_descr(), *typed.current_value()) + } + TripletIter::Int32TripletIter(ref typed) => { + Field::convert_int32(typed.column_descr(), *typed.current_value()) + } + TripletIter::Int64TripletIter(ref typed) => { + Field::convert_int64(typed.column_descr(), *typed.current_value()) + } + TripletIter::Int96TripletIter(ref typed) => { + Field::convert_int96(typed.column_descr(), *typed.current_value()) + } + TripletIter::FloatTripletIter(ref typed) => { + Field::convert_float(typed.column_descr(), *typed.current_value()) + } + TripletIter::DoubleTripletIter(ref typed) => { + Field::convert_double(typed.column_descr(), *typed.current_value()) + } + TripletIter::ByteArrayTripletIter(ref typed) => { + Field::convert_byte_array(typed.column_descr(), typed.current_value().clone())? + } + TripletIter::FixedLenByteArrayTripletIter(ref typed) => Field::convert_byte_array( + typed.column_descr(), + typed.current_value().clone().into(), + )?, + }; + Ok(field) + } +} + +/// Internal typed triplet iterator as a wrapper for column reader +/// (primitive leaf column), provides per-element access. +pub struct TypedTripletIter { + reader: ColumnReaderImpl, + column_descr: ColumnDescPtr, + batch_size: usize, + // type properties + max_def_level: i16, + max_rep_level: i16, + // values and levels + values: Vec, + def_levels: Option>, + rep_levels: Option>, + // current index for the triplet (value, def, rep) + curr_triplet_index: usize, + // how many triplets are left before we need to buffer + triplets_left: usize, + // helper flag to quickly check if we have more values/levels to read + has_next: bool, +} + +impl TypedTripletIter { + /// Creates new typed triplet iterator based on provided column reader. + /// Use batch size to specify the amount of values to buffer from column reader. + fn new(descr: ColumnDescPtr, batch_size: usize, column_reader: ColumnReader) -> Self { + assert!( + batch_size > 0, + "Expected positive batch size, found: {batch_size}" + ); + + let max_def_level = descr.max_def_level(); + let max_rep_level = descr.max_rep_level(); + + let def_levels = if max_def_level == 0 { + None + } else { + Some(vec![0; batch_size]) + }; + let rep_levels = if max_rep_level == 0 { + None + } else { + Some(vec![0; batch_size]) + }; + + Self { + reader: get_typed_column_reader(column_reader), + column_descr: descr, + batch_size, + max_def_level, + max_rep_level, + values: vec![T::T::default(); batch_size], + def_levels, + rep_levels, + curr_triplet_index: 0, + triplets_left: 0, + has_next: false, + } + } + + /// Returns column descriptor reference for the current typed triplet iterator. + #[inline] + pub fn column_descr(&self) -> &ColumnDescPtr { + &self.column_descr + } + + /// Returns maximum definition level for the triplet iterator (leaf column). + #[inline] + fn max_def_level(&self) -> i16 { + self.max_def_level + } + + /// Returns maximum repetition level for the triplet iterator (leaf column). + #[inline] + fn max_rep_level(&self) -> i16 { + self.max_rep_level + } + + /// Returns current value. + /// Method does not advance the iterator, therefore can be called multiple times. + #[inline] + fn current_value(&self) -> &T::T { + assert!( + self.current_def_level() == self.max_def_level(), + "Cannot extract value, max definition level: {}, current level: {}", + self.max_def_level(), + self.current_def_level() + ); + &self.values[self.curr_triplet_index] + } + + /// Returns current definition level. + /// If field is required, then maximum definition level is returned. + #[inline] + fn current_def_level(&self) -> i16 { + if !self.has_next { + return 0; + } + match self.def_levels { + Some(ref vec) => vec[self.curr_triplet_index], + None => self.max_def_level, + } + } + + /// Returns current repetition level. + /// If field is required, then maximum repetition level is returned. + #[inline] + fn current_rep_level(&self) -> i16 { + if !self.has_next { + return 0; + } + match self.rep_levels { + Some(ref vec) => vec[self.curr_triplet_index], + None => self.max_rep_level, + } + } + + /// Quick check if iterator has more values/levels to read. + /// It is updated as a result of `read_next` method, so they are synchronized. + #[inline] + fn has_next(&self) -> bool { + self.has_next + } + + /// Advances to the next triplet. + /// Returns true, if there are more records to read, false there are no records left. + fn read_next(&mut self) -> Result { + self.curr_triplet_index += 1; + + // A loop is required to handle the case of a batch size of 1, as in such a case + // on reaching the end of a record, read_records will return `Ok((1, 0, 0))` + // and therefore not advance `self.triplets_left` + while self.curr_triplet_index >= self.triplets_left { + let (records_read, values_read, levels_read) = { + self.values.clear(); + if let Some(x) = &mut self.def_levels { + x.clear() + } + if let Some(x) = &mut self.rep_levels { + x.clear() + } + + // Buffer triplets + self.reader.read_records( + self.batch_size, + self.def_levels.as_mut(), + self.rep_levels.as_mut(), + &mut self.values, + )? + }; + + // No more values or levels to read + if records_read == 0 && values_read == 0 && levels_read == 0 { + self.curr_triplet_index = 0; + self.has_next = false; + return Ok(false); + } + + // We never read values more than levels + if levels_read == 0 || values_read == levels_read { + // There are no definition levels to read, column is required + // or definition levels match values, so it does not require spacing + self.curr_triplet_index = 0; + self.triplets_left = values_read; + } else if values_read < levels_read { + // Add spacing for triplets. + // The idea is setting values for positions in def_levels when current + // definition level equals to maximum definition level. + // Values and levels are guaranteed to line up, because of + // the column reader method. + + // Note: if values_read == 0, then spacing will not be triggered + let mut idx = values_read; + let def_levels = self.def_levels.as_ref().unwrap(); + self.values.resize(levels_read, T::T::default()); + for i in 0..levels_read { + if def_levels[levels_read - i - 1] == self.max_def_level { + idx -= 1; // This is done to avoid usize becoming a negative value + self.values.swap(levels_read - i - 1, idx); + } + } + self.curr_triplet_index = 0; + self.triplets_left = levels_read; + } else { + return Err(general_err!( + "Spacing of values/levels is wrong, values_read: {}, levels_read: {}", + values_read, + levels_read + )); + } + } + + self.has_next = true; + Ok(true) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use crate::file::reader::{FileReader, SerializedFileReader}; + use crate::schema::types::ColumnPath; + use crate::util::test_common::file_util::get_test_file; + + #[test] + #[should_panic(expected = "Expected positive batch size, found: 0")] + fn test_triplet_zero_batch_size() { + let column_path = ColumnPath::from(vec!["b_struct".to_string(), "b_c_int".to_string()]); + test_column_in_file("nulls.snappy.parquet", 0, &column_path, &[], &[], &[]); + } + + #[test] + fn test_triplet_null_column() { + let path = vec!["b_struct", "b_c_int"]; + let values = vec![]; + let def_levels = vec![1, 1, 1, 1, 1, 1, 1, 1]; + let rep_levels = vec![0, 0, 0, 0, 0, 0, 0, 0]; + test_triplet_iter( + "nulls.snappy.parquet", + path, + &values, + &def_levels, + &rep_levels, + ); + } + + #[test] + fn test_triplet_required_column() { + let path = vec!["ID"]; + let values = vec![Field::Long(8)]; + let def_levels = vec![0]; + let rep_levels = vec![0]; + test_triplet_iter( + "nonnullable.impala.parquet", + path, + &values, + &def_levels, + &rep_levels, + ); + } + + #[test] + fn test_triplet_optional_column() { + let path = vec!["nested_struct", "A"]; + let values = vec![Field::Int(1), Field::Int(7)]; + let def_levels = vec![2, 1, 1, 1, 1, 0, 2]; + let rep_levels = vec![0, 0, 0, 0, 0, 0, 0]; + test_triplet_iter( + "nullable.impala.parquet", + path, + &values, + &def_levels, + &rep_levels, + ); + } + + #[test] + fn test_triplet_optional_list_column() { + let path = vec!["a", "list", "element", "list", "element", "list", "element"]; + let values = vec![ + Field::Str("a".to_string()), + Field::Str("b".to_string()), + Field::Str("c".to_string()), + Field::Str("d".to_string()), + Field::Str("a".to_string()), + Field::Str("b".to_string()), + Field::Str("c".to_string()), + Field::Str("d".to_string()), + Field::Str("e".to_string()), + Field::Str("a".to_string()), + Field::Str("b".to_string()), + Field::Str("c".to_string()), + Field::Str("d".to_string()), + Field::Str("e".to_string()), + Field::Str("f".to_string()), + ]; + let def_levels = vec![7, 7, 7, 4, 7, 7, 7, 7, 7, 4, 7, 7, 7, 7, 7, 7, 4, 7]; + let rep_levels = vec![0, 3, 2, 1, 2, 0, 3, 2, 3, 1, 2, 0, 3, 2, 3, 2, 1, 2]; + test_triplet_iter( + "nested_lists.snappy.parquet", + path, + &values, + &def_levels, + &rep_levels, + ); + } + + #[test] + fn test_triplet_optional_map_column() { + let path = vec!["a", "key_value", "value", "key_value", "key"]; + let values = vec![ + Field::Int(1), + Field::Int(2), + Field::Int(1), + Field::Int(1), + Field::Int(3), + Field::Int(4), + Field::Int(5), + ]; + let def_levels = vec![4, 4, 4, 2, 3, 4, 4, 4, 4]; + let rep_levels = vec![0, 2, 0, 0, 0, 0, 0, 2, 2]; + test_triplet_iter( + "nested_maps.snappy.parquet", + path, + &values, + &def_levels, + &rep_levels, + ); + } + + // Check triplet iterator across different batch sizes + fn test_triplet_iter( + file_name: &str, + column_path: Vec<&str>, + expected_values: &[Field], + expected_def_levels: &[i16], + expected_rep_levels: &[i16], + ) { + // Convert path into column path + let path: Vec = column_path.iter().map(|x| x.to_string()).collect(); + let column_path = ColumnPath::from(path); + + let batch_sizes = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 128, 256]; + for batch_size in batch_sizes { + test_column_in_file( + file_name, + batch_size, + &column_path, + expected_values, + expected_def_levels, + expected_rep_levels, + ); + } + } + + // Check values of a selected column in a file + fn test_column_in_file( + file_name: &str, + batch_size: usize, + column_path: &ColumnPath, + expected_values: &[Field], + expected_def_levels: &[i16], + expected_rep_levels: &[i16], + ) { + let file = get_test_file(file_name); + let file_reader = SerializedFileReader::new(file).unwrap(); + let metadata = file_reader.metadata(); + // Get schema descriptor + let file_metadata = metadata.file_metadata(); + let schema = file_metadata.schema_descr(); + // Get first row group + let row_group_reader = file_reader.get_row_group(0).unwrap(); + + for i in 0..schema.num_columns() { + let descr = schema.column(i); + if descr.path() == column_path { + let reader = row_group_reader.get_column_reader(i).unwrap(); + test_triplet_column( + descr, + reader, + batch_size, + expected_values, + expected_def_levels, + expected_rep_levels, + ); + } + } + } + + // Check values for individual triplet iterator + fn test_triplet_column( + descr: ColumnDescPtr, + reader: ColumnReader, + batch_size: usize, + expected_values: &[Field], + expected_def_levels: &[i16], + expected_rep_levels: &[i16], + ) { + let mut iter = TripletIter::new(descr.clone(), reader, batch_size); + let mut values: Vec = Vec::new(); + let mut def_levels: Vec = Vec::new(); + let mut rep_levels: Vec = Vec::new(); + + assert_eq!(iter.max_def_level(), descr.max_def_level()); + assert_eq!(iter.max_rep_level(), descr.max_rep_level()); + + while let Ok(true) = iter.read_next() { + assert!(iter.has_next()); + if !iter.is_null() { + values.push(iter.current_value().unwrap()); + } + def_levels.push(iter.current_def_level()); + rep_levels.push(iter.current_rep_level()); + } + + assert_eq!(values, expected_values); + assert_eq!(def_levels, expected_def_levels); + assert_eq!(rep_levels, expected_rep_levels); + } + + fn open_triplet_iter(file_name: &str, path: &[&str], batch_size: usize) -> TripletIter { + let column_path = ColumnPath::from(path.iter().map(|x| x.to_string()).collect::>()); + let file = get_test_file(file_name); + let file_reader = SerializedFileReader::new(file).unwrap(); + let metadata = file_reader.metadata(); + let schema = metadata.file_metadata().schema_descr(); + let row_group_reader = file_reader.get_row_group(0).unwrap(); + for i in 0..schema.num_columns() { + let descr = schema.column(i); + if descr.path() == &column_path { + let reader = row_group_reader.get_column_reader(i).unwrap(); + return TripletIter::new(descr.clone(), reader, batch_size); + } + } + panic!("Column {column_path:?} not found in {file_name}"); + } + + #[test] + fn test_current_def_level_safe_after_exhaustion() { + let mut iter = open_triplet_iter("nulls.snappy.parquet", &["b_struct", "b_c_int"], 256); + while let Ok(true) = iter.read_next() {} + assert!(!iter.has_next()); + assert_eq!(iter.current_def_level(), 0); + } + + #[test] + fn test_current_rep_level_safe_after_exhaustion() { + let mut iter = open_triplet_iter( + "nested_lists.snappy.parquet", + &["a", "list", "element", "list", "element", "list", "element"], + 256, + ); + while let Ok(true) = iter.read_next() {} + assert!(!iter.has_next()); + assert_eq!(iter.current_rep_level(), 0); + } +} diff --git a/vendor/parquet/src/schema/mod.rs b/vendor/parquet/src/schema/mod.rs new file mode 100644 index 0000000..415802c --- /dev/null +++ b/vendor/parquet/src/schema/mod.rs @@ -0,0 +1,81 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Parquet schema definitions and methods to print and parse schema. +//! +//! * [`SchemaDescriptor`] describes the data types of the columns stored in a file +//! * [`ColumnDescriptor`]: Describes the schema of a single (leaf) column. +//! * [`ColumnPath`]: Represents the location of a column in the schema (e.g. a nested field) +//! +//! Parquet distinguishes +//! between "logical" and "physical" data types. +//! For instance, strings (logical type) are stored as byte arrays (physical type). +//! Likewise, temporal types like dates, times, timestamps, etc. (logical type) +//! are stored as integers (physical type). +//! +//! [`SchemaDescriptor`]: types::SchemaDescriptor +//! [`ColumnDescriptor`]: types::ColumnDescriptor +//! [`ColumnPath`]: types::ColumnPath +//! +//! # Example +//! +//! ```rust +//! use parquet::{ +//! basic::{ConvertedType, Repetition, Type as PhysicalType}, +//! schema::{parser, printer, types::Type}, +//! }; +//! use std::sync::Arc; +//! +//! // Create the following schema: +//! // +//! // message schema { +//! // OPTIONAL BYTE_ARRAY a (UTF8); +//! // REQUIRED INT32 b; +//! // } +//! +//! let field_a = Type::primitive_type_builder("a", PhysicalType::BYTE_ARRAY) +//! .with_converted_type(ConvertedType::UTF8) +//! .with_repetition(Repetition::OPTIONAL) +//! .build() +//! .unwrap(); +//! +//! let field_b = Type::primitive_type_builder("b", PhysicalType::INT32) +//! .with_repetition(Repetition::REQUIRED) +//! .build() +//! .unwrap(); +//! +//! let schema = Type::group_type_builder("schema") +//! .with_fields(vec![Arc::new(field_a), Arc::new(field_b)]) +//! .build() +//! .unwrap(); +//! +//! let mut buf = Vec::new(); +//! +//! // Print schema into buffer +//! printer::print_schema(&mut buf, &schema); +//! +//! // Parse schema from the string +//! let string_schema = String::from_utf8(buf).unwrap(); +//! let parsed_schema = parser::parse_message_type(&string_schema).unwrap(); +//! +//! assert_eq!(schema, parsed_schema); +//! ``` + +pub mod parser; +pub mod printer; +pub mod types; +pub mod visitor; diff --git a/vendor/parquet/src/schema/parser.rs b/vendor/parquet/src/schema/parser.rs new file mode 100644 index 0000000..071962a --- /dev/null +++ b/vendor/parquet/src/schema/parser.rs @@ -0,0 +1,1093 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Parquet schema parser. +//! Provides methods to parse and validate string message type into Parquet +//! [`Type`]. +//! +//! # Example +//! +//! ```rust +//! use parquet::schema::parser::parse_message_type; +//! +//! let message_type = " +//! message spark_schema { +//! OPTIONAL BYTE_ARRAY a (UTF8); +//! REQUIRED INT32 b; +//! REQUIRED DOUBLE c; +//! REQUIRED BOOLEAN d; +//! OPTIONAL group e (LIST) { +//! REPEATED group list { +//! REQUIRED INT32 element; +//! } +//! } +//! } +//! "; +//! +//! let schema = parse_message_type(message_type).expect("Expected valid schema"); +//! println!("{:?}", schema); +//! ``` + +use std::sync::Arc; + +use crate::basic::{ConvertedType, LogicalType, Repetition, TimeUnit, Type as PhysicalType}; +use crate::errors::{ParquetError, Result}; +use crate::schema::types::{Type, TypePtr}; + +/// Parses message type as string into a Parquet [`Type`] +/// which, for example, could be used to extract individual columns. Returns Parquet +/// general error when parsing or validation fails. +pub fn parse_message_type(message_type: &str) -> Result { + let mut parser = Parser { + tokenizer: &mut Tokenizer::from_str(message_type), + }; + parser.parse_message_type() +} + +/// Tokenizer to split message type string into tokens that are separated using characters +/// defined in `is_schema_delim` method. Tokenizer also preserves delimiters as tokens. +/// Tokenizer provides Iterator interface to process tokens; it also allows to step back +/// to reprocess previous tokens. +struct Tokenizer<'a> { + // List of all tokens for a string + tokens: Vec<&'a str>, + // Current index of vector + index: usize, +} + +impl<'a> Tokenizer<'a> { + // Create tokenizer from message type string + pub fn from_str(string: &'a str) -> Self { + let vec = string + .split_whitespace() + .flat_map(Self::split_token) + .collect(); + Tokenizer { + tokens: vec, + index: 0, + } + } + + // List of all special characters in schema + fn is_schema_delim(c: char) -> bool { + c == ';' || c == '{' || c == '}' || c == '(' || c == ')' || c == '=' || c == ',' + } + + /// Splits string into tokens; input string can already be token or can contain + /// delimiters, e.g. required" -> Vec("required") and + /// "(UTF8);" -> Vec("(", "UTF8", ")", ";") + fn split_token(string: &str) -> Vec<&str> { + let mut buffer: Vec<&str> = Vec::new(); + let mut tail = string; + while let Some(index) = tail.find(Self::is_schema_delim) { + let (h, t) = tail.split_at(index); + if !h.is_empty() { + buffer.push(h); + } + buffer.push(&t[0..1]); + tail = &t[1..]; + } + if !tail.is_empty() { + buffer.push(tail); + } + buffer + } + + // Move pointer to a previous element + fn backtrack(&mut self) { + self.index -= 1; + } +} + +impl<'a> Iterator for Tokenizer<'a> { + type Item = &'a str; + + fn next(&mut self) -> Option<&'a str> { + if self.index < self.tokens.len() { + self.index += 1; + Some(self.tokens[self.index - 1]) + } else { + None + } + } +} + +/// Internal Schema parser. +/// Traverses message type using tokenizer and parses each group/primitive type +/// recursively. +struct Parser<'a> { + tokenizer: &'a mut Tokenizer<'a>, +} + +// Utility function to assert token on validity. +fn assert_token(token: Option<&str>, expected: &str) -> Result<()> { + match token { + Some(value) if value == expected => Ok(()), + Some(other) => Err(general_err!( + "Expected '{}', found token '{}'", + expected, + other + )), + None => Err(general_err!( + "Expected '{}', but no token found (None)", + expected + )), + } +} + +// Utility function to parse i32 or return general error. +#[inline] +fn parse_i32(value: Option<&str>, not_found_msg: &str, parse_fail_msg: &str) -> Result { + value + .ok_or_else(|| general_err!(not_found_msg)) + .and_then(|v| v.parse::().map_err(|_| general_err!(parse_fail_msg))) +} + +// Utility function to parse boolean or return general error. +#[inline] +fn parse_bool(value: Option<&str>, not_found_msg: &str, parse_fail_msg: &str) -> Result { + value + .ok_or_else(|| general_err!(not_found_msg)) + .and_then(|v| { + v.to_lowercase() + .parse::() + .map_err(|_| general_err!(parse_fail_msg)) + }) +} + +// Utility function to parse TimeUnit or return general error. +fn parse_timeunit( + value: Option<&str>, + not_found_msg: &str, + parse_fail_msg: &str, +) -> Result { + value + .ok_or_else(|| general_err!(not_found_msg)) + .and_then(|v| match v.to_uppercase().as_str() { + "MILLIS" => Ok(TimeUnit::MILLIS), + "MICROS" => Ok(TimeUnit::MICROS), + "NANOS" => Ok(TimeUnit::NANOS), + _ => Err(general_err!(parse_fail_msg)), + }) +} + +impl Parser<'_> { + // Entry function to parse message type, uses internal tokenizer. + fn parse_message_type(&mut self) -> Result { + // Check that message type starts with "message". + match self.tokenizer.next() { + Some("message") => { + let name = self + .tokenizer + .next() + .ok_or_else(|| general_err!("Expected name, found None"))?; + Type::group_type_builder(name) + .with_fields(self.parse_child_types()?) + .build() + } + _ => Err(general_err!("Message type does not start with 'message'")), + } + } + + // Parses child types for a current group type. + // This is only invoked on root and group types. + fn parse_child_types(&mut self) -> Result> { + assert_token(self.tokenizer.next(), "{")?; + let mut vec = Vec::new(); + while let Some(value) = self.tokenizer.next() { + if value == "}" { + break; + } else { + self.tokenizer.backtrack(); + vec.push(Arc::new(self.add_type()?)); + } + } + Ok(vec) + } + + fn add_type(&mut self) -> Result { + // Parse repetition + let repetition = self + .tokenizer + .next() + .ok_or_else(|| general_err!("Expected repetition, found None")) + .and_then(|v| v.to_uppercase().parse::())?; + + match self.tokenizer.next() { + Some(group) if group.to_uppercase() == "GROUP" => self.add_group_type(Some(repetition)), + Some(type_string) => { + let physical_type = type_string.to_uppercase().parse::()?; + self.add_primitive_type(repetition, physical_type) + } + None => Err(general_err!("Invalid type, could not extract next token")), + } + } + + fn add_group_type(&mut self, repetition: Option) -> Result { + // Parse name of the group type + let name = self + .tokenizer + .next() + .ok_or_else(|| general_err!("Expected name, found None"))?; + + // Parse logical or converted type if exists + let (logical_type, converted_type) = if let Some("(") = self.tokenizer.next() { + let tpe = self + .tokenizer + .next() + .ok_or_else(|| general_err!("Expected converted type, found None")) + .and_then(|v| { + // Try logical type first + let upper = v.to_uppercase(); + let logical = upper.parse::(); + match logical { + Ok(logical) => { + Ok((Some(logical.clone()), ConvertedType::from(Some(logical)))) + } + Err(_) => Ok((None, upper.parse::()?)), + } + })?; + assert_token(self.tokenizer.next(), ")")?; + tpe + } else { + self.tokenizer.backtrack(); + (None, ConvertedType::NONE) + }; + + // Parse optional id + let id = if let Some("=") = self.tokenizer.next() { + self.tokenizer.next().and_then(|v| v.parse::().ok()) + } else { + self.tokenizer.backtrack(); + None + }; + + let mut builder = Type::group_type_builder(name) + .with_logical_type(logical_type) + .with_converted_type(converted_type) + .with_fields(self.parse_child_types()?) + .with_id(id); + if let Some(rep) = repetition { + builder = builder.with_repetition(rep); + } + builder.build() + } + + fn add_primitive_type( + &mut self, + repetition: Repetition, + physical_type: PhysicalType, + ) -> Result { + // Read type length if the type is FIXED_LEN_BYTE_ARRAY. + let mut length: i32 = -1; + if physical_type == PhysicalType::FIXED_LEN_BYTE_ARRAY { + assert_token(self.tokenizer.next(), "(")?; + length = parse_i32( + self.tokenizer.next(), + "Expected length for FIXED_LEN_BYTE_ARRAY, found None", + "Failed to parse length for FIXED_LEN_BYTE_ARRAY", + )?; + assert_token(self.tokenizer.next(), ")")?; + } + + // Parse name of the primitive type + let name = self + .tokenizer + .next() + .ok_or_else(|| general_err!("Expected name, found None"))?; + + // Parse converted type + let (logical_type, converted_type, precision, scale) = if let Some("(") = + self.tokenizer.next() + { + let (mut logical, mut converted) = self + .tokenizer + .next() + .ok_or_else(|| general_err!("Expected logical or converted type, found None")) + .and_then(|v| { + let upper = v.to_uppercase(); + let logical = upper.parse::(); + match logical { + Ok(logical) => { + Ok((Some(logical.clone()), ConvertedType::from(Some(logical)))) + } + Err(_) => Ok((None, upper.parse::()?)), + } + })?; + + // Parse precision and scale for decimals + let mut precision: i32 = -1; + let mut scale: i32 = -1; + + // Parse the concrete logical type + if let Some(tpe) = &logical { + match tpe { + LogicalType::Decimal { .. } => { + if let Some("(") = self.tokenizer.next() { + precision = parse_i32( + self.tokenizer.next(), + "Expected precision, found None", + "Failed to parse precision for DECIMAL type", + )?; + if let Some(",") = self.tokenizer.next() { + scale = parse_i32( + self.tokenizer.next(), + "Expected scale, found None", + "Failed to parse scale for DECIMAL type", + )?; + assert_token(self.tokenizer.next(), ")")?; + } else { + scale = 0 + } + logical = Some(LogicalType::decimal(scale, precision)); + converted = ConvertedType::from(logical.clone()); + } + } + LogicalType::Time { .. } => { + if let Some("(") = self.tokenizer.next() { + let unit = parse_timeunit( + self.tokenizer.next(), + "Invalid timeunit found", + "Failed to parse timeunit for TIME type", + )?; + if let Some(",") = self.tokenizer.next() { + let is_adjusted_to_u_t_c = parse_bool( + self.tokenizer.next(), + "Invalid boolean found", + "Failed to parse timezone info for TIME type", + )?; + assert_token(self.tokenizer.next(), ")")?; + logical = Some(LogicalType::time(is_adjusted_to_u_t_c, unit)); + converted = ConvertedType::from(logical.clone()); + } else { + // Invalid token for unit + self.tokenizer.backtrack(); + } + } + } + LogicalType::Timestamp { .. } => { + if let Some("(") = self.tokenizer.next() { + let unit = parse_timeunit( + self.tokenizer.next(), + "Invalid timeunit found", + "Failed to parse timeunit for TIMESTAMP type", + )?; + if let Some(",") = self.tokenizer.next() { + let is_adjusted_to_u_t_c = parse_bool( + self.tokenizer.next(), + "Invalid boolean found", + "Failed to parse timezone info for TIMESTAMP type", + )?; + assert_token(self.tokenizer.next(), ")")?; + logical = Some(LogicalType::timestamp(is_adjusted_to_u_t_c, unit)); + converted = ConvertedType::from(logical.clone()); + } else { + // Invalid token for unit + self.tokenizer.backtrack(); + } + } + } + LogicalType::Integer { .. } => { + if let Some("(") = self.tokenizer.next() { + let bit_width = parse_i32( + self.tokenizer.next(), + "Invalid bit_width found", + "Failed to parse bit_width for INTEGER type", + )? as i8; + match physical_type { + PhysicalType::INT32 => match bit_width { + 8 | 16 | 32 => {} + _ => { + return Err(general_err!( + "Incorrect bit width {} for INT32", + bit_width + )); + } + }, + PhysicalType::INT64 => { + if bit_width != 64 { + return Err(general_err!( + "Incorrect bit width {} for INT64", + bit_width + )); + } + } + _ => { + return Err(general_err!( + "Logical type Integer cannot be used with physical type {}", + physical_type + )); + } + } + if let Some(",") = self.tokenizer.next() { + let is_signed = parse_bool( + self.tokenizer.next(), + "Invalid boolean found", + "Failed to parse is_signed for INTEGER type", + )?; + assert_token(self.tokenizer.next(), ")")?; + logical = Some(LogicalType::integer(bit_width, is_signed)); + converted = ConvertedType::from(logical.clone()); + } else { + // Invalid token for unit + self.tokenizer.backtrack(); + } + } + } + _ => {} + } + } else if converted == ConvertedType::DECIMAL { + if let Some("(") = self.tokenizer.next() { + // Parse precision + precision = parse_i32( + self.tokenizer.next(), + "Expected precision, found None", + "Failed to parse precision for DECIMAL type", + )?; + + // Parse scale + scale = if let Some(",") = self.tokenizer.next() { + parse_i32( + self.tokenizer.next(), + "Expected scale, found None", + "Failed to parse scale for DECIMAL type", + )? + } else { + // Scale is not provided, set it to 0. + self.tokenizer.backtrack(); + 0 + }; + + assert_token(self.tokenizer.next(), ")")?; + } else { + self.tokenizer.backtrack(); + } + } + + assert_token(self.tokenizer.next(), ")")?; + (logical, converted, precision, scale) + } else { + self.tokenizer.backtrack(); + (None, ConvertedType::NONE, -1, -1) + }; + + // Parse optional id + let id = if let Some("=") = self.tokenizer.next() { + self.tokenizer.next().and_then(|v| v.parse::().ok()) + } else { + self.tokenizer.backtrack(); + None + }; + assert_token(self.tokenizer.next(), ";")?; + + Type::primitive_type_builder(name, physical_type) + .with_repetition(repetition) + .with_logical_type(logical_type) + .with_converted_type(converted_type) + .with_length(length) + .with_precision(precision) + .with_scale(scale) + .with_id(id) + .build() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_tokenize_empty_string() { + assert_eq!(Tokenizer::from_str("").next(), None); + } + + #[test] + fn test_tokenize_delimiters() { + let mut iter = Tokenizer::from_str(",;{}()="); + assert_eq!(iter.next(), Some(",")); + assert_eq!(iter.next(), Some(";")); + assert_eq!(iter.next(), Some("{")); + assert_eq!(iter.next(), Some("}")); + assert_eq!(iter.next(), Some("(")); + assert_eq!(iter.next(), Some(")")); + assert_eq!(iter.next(), Some("=")); + assert_eq!(iter.next(), None); + } + + #[test] + fn test_tokenize_delimiters_with_whitespaces() { + let mut iter = Tokenizer::from_str(" , ; { } ( ) = "); + assert_eq!(iter.next(), Some(",")); + assert_eq!(iter.next(), Some(";")); + assert_eq!(iter.next(), Some("{")); + assert_eq!(iter.next(), Some("}")); + assert_eq!(iter.next(), Some("(")); + assert_eq!(iter.next(), Some(")")); + assert_eq!(iter.next(), Some("=")); + assert_eq!(iter.next(), None); + } + + #[test] + fn test_tokenize_words() { + let mut iter = Tokenizer::from_str("abc def ghi jkl mno"); + assert_eq!(iter.next(), Some("abc")); + assert_eq!(iter.next(), Some("def")); + assert_eq!(iter.next(), Some("ghi")); + assert_eq!(iter.next(), Some("jkl")); + assert_eq!(iter.next(), Some("mno")); + assert_eq!(iter.next(), None); + } + + #[test] + fn test_tokenize_backtrack() { + let mut iter = Tokenizer::from_str("abc;"); + assert_eq!(iter.next(), Some("abc")); + assert_eq!(iter.next(), Some(";")); + iter.backtrack(); + assert_eq!(iter.next(), Some(";")); + assert_eq!(iter.next(), None); + } + + #[test] + fn test_tokenize_message_type() { + let schema = " + message schema { + required int32 a; + optional binary c (UTF8); + required group d { + required int32 a; + optional binary c (UTF8); + } + required group e (LIST) { + repeated group list { + required int32 element; + } + } + } + "; + let iter = Tokenizer::from_str(schema); + let mut res = Vec::new(); + for token in iter { + res.push(token); + } + assert_eq!( + res, + vec![ + "message", "schema", "{", "required", "int32", "a", ";", "optional", "binary", "c", + "(", "UTF8", ")", ";", "required", "group", "d", "{", "required", "int32", "a", + ";", "optional", "binary", "c", "(", "UTF8", ")", ";", "}", "required", "group", + "e", "(", "LIST", ")", "{", "repeated", "group", "list", "{", "required", "int32", + "element", ";", "}", "}", "}" + ] + ); + } + + #[test] + fn test_assert_token() { + assert!(assert_token(Some("a"), "a").is_ok()); + assert!(assert_token(Some("a"), "b").is_err()); + assert!(assert_token(None, "b").is_err()); + } + + fn parse(schema: &str) -> Result { + let mut iter = Tokenizer::from_str(schema); + Parser { + tokenizer: &mut iter, + } + .parse_message_type() + } + + #[test] + fn test_parse_message_type_invalid() { + assert_eq!( + parse("test").unwrap_err().to_string(), + "Parquet error: Message type does not start with 'message'" + ); + } + + #[test] + fn test_parse_message_type_no_name() { + assert_eq!( + parse("message").unwrap_err().to_string(), + "Parquet error: Expected name, found None" + ); + } + + #[test] + fn test_parse_message_type_fixed_byte_array() { + let schema = " + message schema { + REQUIRED FIXED_LEN_BYTE_ARRAY col; + } + "; + assert_eq!( + parse(schema).unwrap_err().to_string(), + "Parquet error: Expected '(', found token 'col'" + ); + + let schema = " + message schema { + REQUIRED FIXED_LEN_BYTE_ARRAY(16) col; + } + "; + parse(schema).unwrap(); + } + + #[test] + fn test_parse_message_type_integer() { + // Invalid integer syntax + let schema = " + message root { + optional int64 f1 (INTEGER()); + } + "; + assert_eq!( + parse(schema).unwrap_err().to_string(), + "Parquet error: Failed to parse bit_width for INTEGER type" + ); + + // Invalid integer syntax, needs both bit-width and UTC sign + let schema = " + message root { + optional int64 f1 (INTEGER(32,)); + } + "; + assert_eq!( + parse(schema).unwrap_err().to_string(), + "Parquet error: Incorrect bit width 32 for INT64" + ); + + // Invalid integer because of non-numeric bit width + let schema = " + message root { + optional int32 f1 (INTEGER(eight,true)); + } + "; + assert_eq!( + parse(schema).unwrap_err().to_string(), + "Parquet error: Failed to parse bit_width for INTEGER type" + ); + + // Valid types + let schema = " + message root { + optional int32 f1 (INTEGER(8,false)); + optional int32 f2 (INTEGER(8,true)); + optional int32 f3 (INTEGER(16,false)); + optional int32 f4 (INTEGER(16,true)); + optional int32 f5 (INTEGER(32,false)); + optional int32 f6 (INTEGER(32,true)); + optional int64 f7 (INTEGER(64,false)); + optional int64 f7 (INTEGER(64,true)); + } + "; + parse(schema).unwrap(); + } + + #[test] + fn test_parse_message_type_temporal() { + // Invalid timestamp syntax + let schema = " + message root { + optional int64 f1 (TIMESTAMP(); + } + "; + assert_eq!( + parse(schema).unwrap_err().to_string(), + "Parquet error: Failed to parse timeunit for TIMESTAMP type" + ); + + // Invalid timestamp syntax, needs both unit and UTC adjustment + let schema = " + message root { + optional int64 f1 (TIMESTAMP(MILLIS,)); + } + "; + assert_eq!( + parse(schema).unwrap_err().to_string(), + "Parquet error: Failed to parse timezone info for TIMESTAMP type" + ); + + // Invalid timestamp because of unknown unit + let schema = " + message root { + optional int64 f1 (TIMESTAMP(YOCTOS,)); + } + "; + + assert_eq!( + parse(schema).unwrap_err().to_string(), + "Parquet error: Failed to parse timeunit for TIMESTAMP type" + ); + + // Valid types + let schema = " + message root { + optional int32 f1 (DATE); + optional int32 f2 (TIME(MILLIS,true)); + optional int64 f3 (TIME(MICROS,false)); + optional int64 f4 (TIME(NANOS,true)); + optional int64 f5 (TIMESTAMP(MILLIS,true)); + optional int64 f6 (TIMESTAMP(MICROS,true)); + optional int64 f7 (TIMESTAMP(NANOS,false)); + } + "; + parse(schema).unwrap(); + } + + #[test] + fn test_parse_message_type_decimal() { + // It is okay for decimal to omit precision and scale with right syntax. + // Here we test wrong syntax of decimal type + + // Invalid decimal syntax + let schema = " + message root { + optional int32 f1 (DECIMAL(); + } + "; + assert_eq!( + parse(schema).unwrap_err().to_string(), + "Parquet error: Failed to parse precision for DECIMAL type" + ); + + // Invalid decimal, need precision and scale + let schema = " + message root { + optional int32 f1 (DECIMAL()); + } + "; + assert_eq!( + parse(schema).unwrap_err().to_string(), + "Parquet error: Failed to parse precision for DECIMAL type" + ); + + // Invalid decimal because of `,` - has precision, needs scale + let schema = " + message root { + optional int32 f1 (DECIMAL(8,)); + } + "; + assert_eq!( + parse(schema).unwrap_err().to_string(), + "Parquet error: Failed to parse scale for DECIMAL type" + ); + + // Invalid decimal because, we always require either precision or scale to be + // specified as part of converted type + let schema = " + message root { + optional int32 f3 (DECIMAL); + } + "; + assert_eq!( + parse(schema).unwrap_err().to_string(), + "Parquet error: Expected ')', found token ';'" + ); + + // Valid decimal (precision, scale) + let schema = " + message root { + optional int32 f1 (DECIMAL(8, 3)); + optional int32 f2 (DECIMAL(8)); + } + "; + parse(schema).unwrap(); + } + + #[test] + fn test_parse_message_type_compare_1() { + let schema = " + message root { + optional fixed_len_byte_array(5) f1 (DECIMAL(9, 3)); + optional fixed_len_byte_array (16) f2 (DECIMAL (38, 18)); + optional fixed_len_byte_array (2) f3 (FLOAT16); + } + "; + let message = parse(schema).unwrap(); + + let expected = Type::group_type_builder("root") + .with_fields(vec![ + Arc::new( + Type::primitive_type_builder("f1", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_logical_type(Some(LogicalType::decimal(3, 9))) + .with_converted_type(ConvertedType::DECIMAL) + .with_length(5) + .with_precision(9) + .with_scale(3) + .build() + .unwrap(), + ), + Arc::new( + Type::primitive_type_builder("f2", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_logical_type(Some(LogicalType::decimal(18, 38))) + .with_converted_type(ConvertedType::DECIMAL) + .with_length(16) + .with_precision(38) + .with_scale(18) + .build() + .unwrap(), + ), + Arc::new( + Type::primitive_type_builder("f3", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_logical_type(Some(LogicalType::Float16)) + .with_length(2) + .build() + .unwrap(), + ), + ]) + .build() + .unwrap(); + + assert_eq!(message, expected); + } + + #[test] + fn test_parse_message_type_compare_2() { + let schema = " + message root { + required group a0 { + optional group a1 (LIST) { + repeated binary a2 (UTF8); + } + + optional group b1 (LIST) { + repeated group b2 { + optional int32 b3; + optional double b4; + } + } + } + } + "; + let message = parse(schema).unwrap(); + + let expected = Type::group_type_builder("root") + .with_fields(vec![Arc::new( + Type::group_type_builder("a0") + .with_repetition(Repetition::REQUIRED) + .with_fields(vec![ + Arc::new( + Type::group_type_builder("a1") + .with_repetition(Repetition::OPTIONAL) + .with_logical_type(Some(LogicalType::List)) + .with_converted_type(ConvertedType::LIST) + .with_fields(vec![Arc::new( + Type::primitive_type_builder("a2", PhysicalType::BYTE_ARRAY) + .with_repetition(Repetition::REPEATED) + .with_converted_type(ConvertedType::UTF8) + .build() + .unwrap(), + )]) + .build() + .unwrap(), + ), + Arc::new( + Type::group_type_builder("b1") + .with_repetition(Repetition::OPTIONAL) + .with_logical_type(Some(LogicalType::List)) + .with_converted_type(ConvertedType::LIST) + .with_fields(vec![Arc::new( + Type::group_type_builder("b2") + .with_repetition(Repetition::REPEATED) + .with_fields(vec![ + Arc::new( + Type::primitive_type_builder( + "b3", + PhysicalType::INT32, + ) + .build() + .unwrap(), + ), + Arc::new( + Type::primitive_type_builder( + "b4", + PhysicalType::DOUBLE, + ) + .build() + .unwrap(), + ), + ]) + .build() + .unwrap(), + )]) + .build() + .unwrap(), + ), + ]) + .build() + .unwrap(), + )]) + .build() + .unwrap(); + + assert_eq!(message, expected); + } + + #[test] + fn test_parse_message_type_compare_3() { + let schema = " + message root { + required int32 _1 (INT_8); + required int32 _2 (INT_16); + required float _3; + required double _4; + optional int32 _5 (DATE); + optional binary _6 (UTF8); + } + "; + let message = parse(schema).unwrap(); + + let fields = vec![ + Arc::new( + Type::primitive_type_builder("_1", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::INT_8) + .build() + .unwrap(), + ), + Arc::new( + Type::primitive_type_builder("_2", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::INT_16) + .build() + .unwrap(), + ), + Arc::new( + Type::primitive_type_builder("_3", PhysicalType::FLOAT) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(), + ), + Arc::new( + Type::primitive_type_builder("_4", PhysicalType::DOUBLE) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(), + ), + Arc::new( + Type::primitive_type_builder("_5", PhysicalType::INT32) + .with_logical_type(Some(LogicalType::Date)) + .with_converted_type(ConvertedType::DATE) + .build() + .unwrap(), + ), + Arc::new( + Type::primitive_type_builder("_6", PhysicalType::BYTE_ARRAY) + .with_converted_type(ConvertedType::UTF8) + .build() + .unwrap(), + ), + ]; + + let expected = Type::group_type_builder("root") + .with_fields(fields) + .build() + .unwrap(); + assert_eq!(message, expected); + } + + #[test] + fn test_parse_message_type_compare_4() { + let schema = " + message root { + required int32 _1 (INTEGER(8,true)); + required int32 _2 (INTEGER(16,false)); + required float _3; + required double _4; + optional int32 _5 (DATE); + optional int32 _6 (TIME(MILLIS,false)); + optional int64 _7 (TIME(MICROS,true)); + optional int64 _8 (TIMESTAMP(MILLIS,true)); + optional int64 _9 (TIMESTAMP(NANOS,false)); + optional binary _10 (STRING); + } + "; + let message = parse(schema).unwrap(); + + let fields = vec![ + Arc::new( + Type::primitive_type_builder("_1", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .with_logical_type(Some(LogicalType::integer(8, true))) + .build() + .unwrap(), + ), + Arc::new( + Type::primitive_type_builder("_2", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .with_logical_type(Some(LogicalType::integer(16, false))) + .build() + .unwrap(), + ), + Arc::new( + Type::primitive_type_builder("_3", PhysicalType::FLOAT) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(), + ), + Arc::new( + Type::primitive_type_builder("_4", PhysicalType::DOUBLE) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(), + ), + Arc::new( + Type::primitive_type_builder("_5", PhysicalType::INT32) + .with_logical_type(Some(LogicalType::Date)) + .build() + .unwrap(), + ), + Arc::new( + Type::primitive_type_builder("_6", PhysicalType::INT32) + .with_logical_type(Some(LogicalType::time(false, TimeUnit::MILLIS))) + .build() + .unwrap(), + ), + Arc::new( + Type::primitive_type_builder("_7", PhysicalType::INT64) + .with_logical_type(Some(LogicalType::time(true, TimeUnit::MICROS))) + .build() + .unwrap(), + ), + Arc::new( + Type::primitive_type_builder("_8", PhysicalType::INT64) + .with_logical_type(Some(LogicalType::timestamp(true, TimeUnit::MILLIS))) + .build() + .unwrap(), + ), + Arc::new( + Type::primitive_type_builder("_9", PhysicalType::INT64) + .with_logical_type(Some(LogicalType::timestamp(false, TimeUnit::NANOS))) + .build() + .unwrap(), + ), + Arc::new( + Type::primitive_type_builder("_10", PhysicalType::BYTE_ARRAY) + .with_logical_type(Some(LogicalType::String)) + .build() + .unwrap(), + ), + ]; + + let expected = Type::group_type_builder("root") + .with_fields(fields) + .build() + .unwrap(); + assert_eq!(message, expected); + } +} diff --git a/vendor/parquet/src/schema/printer.rs b/vendor/parquet/src/schema/printer.rs new file mode 100644 index 0000000..e43e09b --- /dev/null +++ b/vendor/parquet/src/schema/printer.rs @@ -0,0 +1,1176 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Parquet schema printer. +//! Provides methods to print Parquet file schema and list file metadata. +//! +//! # Example +//! +//! ```rust +//! use parquet::{ +//! file::reader::{FileReader, SerializedFileReader}, +//! schema::printer::{print_file_metadata, print_parquet_metadata, print_schema}, +//! }; +//! use std::{fs::File, path::Path}; +//! +//! // Open a file +//! let path = Path::new("test.parquet"); +//! if let Ok(file) = File::open(&path) { +//! let reader = SerializedFileReader::new(file).unwrap(); +//! let parquet_metadata = reader.metadata(); +//! +//! print_parquet_metadata(&mut std::io::stdout(), &parquet_metadata); +//! print_file_metadata(&mut std::io::stdout(), &parquet_metadata.file_metadata()); +//! +//! print_schema( +//! &mut std::io::stdout(), +//! &parquet_metadata.file_metadata().schema(), +//! ); +//! } +//! ``` + +use std::{fmt, io}; + +use crate::basic::{ + ConvertedType, DecimalType, GeographyType, GeometryType, IntType, LogicalType, TimeUnit, + Type as PhysicalType, VariantType, +}; +use crate::file::metadata::{ColumnChunkMetaData, FileMetaData, ParquetMetaData, RowGroupMetaData}; +use crate::schema::types::Type; + +/// Prints Parquet metadata [`ParquetMetaData`] information. +#[allow(unused_must_use)] +pub fn print_parquet_metadata(out: &mut dyn io::Write, metadata: &ParquetMetaData) { + print_file_metadata(out, metadata.file_metadata()); + writeln!(out); + writeln!(out); + writeln!(out, "num of row groups: {}", metadata.num_row_groups()); + writeln!(out, "row groups:"); + writeln!(out); + for (i, rg) in metadata.row_groups().iter().enumerate() { + writeln!(out, "row group {i}:"); + print_dashes(out, 80); + print_row_group_metadata(out, rg); + } +} + +/// Prints file metadata [`FileMetaData`] information. +#[allow(unused_must_use)] +pub fn print_file_metadata(out: &mut dyn io::Write, file_metadata: &FileMetaData) { + writeln!(out, "version: {}", file_metadata.version()); + writeln!(out, "num of rows: {}", file_metadata.num_rows()); + if let Some(created_by) = file_metadata.created_by().as_ref() { + writeln!(out, "created by: {created_by}"); + } + if let Some(metadata) = file_metadata.key_value_metadata() { + writeln!(out, "metadata:"); + for kv in metadata.iter() { + writeln!( + out, + " {}: {}", + kv.key, + kv.value.as_ref().unwrap_or(&"".to_owned()) + ); + } + } + let schema = file_metadata.schema(); + print_schema(out, schema); +} + +/// Prints Parquet [`Type`] information. +/// +/// # Example +/// +/// ```rust +/// use parquet::{ +/// basic::{ConvertedType, Repetition, Type as PhysicalType}, +/// schema::{printer::print_schema, types::Type}, +/// }; +/// use std::sync::Arc; +/// +/// let field_a = Type::primitive_type_builder("a", PhysicalType::BYTE_ARRAY) +/// .with_id(Some(42)) +/// .with_converted_type(ConvertedType::UTF8) +/// .build() +/// .unwrap(); +/// +/// let field_b = Type::primitive_type_builder("b", PhysicalType::INT32) +/// .with_repetition(Repetition::REQUIRED) +/// .build() +/// .unwrap(); +/// +/// let field_d = Type::primitive_type_builder("d", PhysicalType::INT64) +/// .with_id(Some(99)) +/// .build() +/// .unwrap(); +/// +/// let field_c = Type::group_type_builder("c") +/// .with_id(Some(43)) +/// .with_fields(vec![Arc::new(field_d)]) +/// .build() +/// .unwrap(); +/// +/// let schema = Type::group_type_builder("schema") +/// .with_fields(vec![Arc::new(field_a), Arc::new(field_b), Arc::new(field_c)]) +/// .build() +/// .unwrap(); +/// +/// print_schema(&mut std::io::stdout(), &schema); +/// ``` +/// +/// outputs +/// +/// ```text +/// message schema { +/// OPTIONAL BYTE_ARRAY a [42] (UTF8); +/// REQUIRED INT32 b; +/// message c [43] { +/// OPTIONAL INT64 d [99]; +/// } +/// } +/// ``` +#[allow(unused_must_use)] +pub fn print_schema(out: &mut dyn io::Write, tp: &Type) { + // TODO: better if we can pass fmt::Write to Printer. + // But how can we make it to accept both io::Write & fmt::Write? + let mut s = String::new(); + { + let mut printer = Printer::new(&mut s); + printer.print(tp); + } + writeln!(out, "{s}"); +} + +#[allow(unused_must_use)] +fn print_row_group_metadata(out: &mut dyn io::Write, rg_metadata: &RowGroupMetaData) { + writeln!(out, "total byte size: {}", rg_metadata.total_byte_size()); + writeln!(out, "num of rows: {}", rg_metadata.num_rows()); + writeln!(out); + writeln!(out, "num of columns: {}", rg_metadata.num_columns()); + writeln!(out, "columns: "); + for (i, cc) in rg_metadata.columns().iter().enumerate() { + writeln!(out); + writeln!(out, "column {i}:"); + print_dashes(out, 80); + print_column_chunk_metadata(out, cc); + } +} + +#[allow(unused_must_use)] +fn print_column_chunk_metadata(out: &mut dyn io::Write, cc_metadata: &ColumnChunkMetaData) { + writeln!(out, "column type: {}", cc_metadata.column_type()); + writeln!(out, "column path: {}", cc_metadata.column_path()); + let encoding_strs: Vec<_> = cc_metadata.encodings().map(|e| format!("{e}")).collect(); + writeln!(out, "encodings: {}", encoding_strs.join(" ")); + let file_path_str = cc_metadata.file_path().unwrap_or("N/A"); + writeln!(out, "file path: {file_path_str}"); + writeln!(out, "file offset: {}", cc_metadata.file_offset()); + writeln!(out, "num of values: {}", cc_metadata.num_values()); + writeln!(out, "compression: {}", cc_metadata.compression_codec()); + writeln!( + out, + "total compressed size (in bytes): {}", + cc_metadata.compressed_size() + ); + writeln!( + out, + "total uncompressed size (in bytes): {}", + cc_metadata.uncompressed_size() + ); + writeln!(out, "data page offset: {}", cc_metadata.data_page_offset()); + let index_page_offset_str = match cc_metadata.index_page_offset() { + None => "N/A".to_owned(), + Some(ipo) => ipo.to_string(), + }; + writeln!(out, "index page offset: {index_page_offset_str}"); + let dict_page_offset_str = match cc_metadata.dictionary_page_offset() { + None => "N/A".to_owned(), + Some(dpo) => dpo.to_string(), + }; + writeln!(out, "dictionary page offset: {dict_page_offset_str}"); + let statistics_str = match cc_metadata.statistics() { + None => "N/A".to_owned(), + Some(stats) => stats.to_string(), + }; + writeln!(out, "statistics: {statistics_str}"); + let bloom_filter_offset_str = match cc_metadata.bloom_filter_offset() { + None => "N/A".to_owned(), + Some(bfo) => bfo.to_string(), + }; + writeln!(out, "bloom filter offset: {bloom_filter_offset_str}"); + let bloom_filter_length_str = match cc_metadata.bloom_filter_length() { + None => "N/A".to_owned(), + Some(bfo) => bfo.to_string(), + }; + writeln!(out, "bloom filter length: {bloom_filter_length_str}"); + let offset_index_offset_str = match cc_metadata.offset_index_offset() { + None => "N/A".to_owned(), + Some(oio) => oio.to_string(), + }; + writeln!(out, "offset index offset: {offset_index_offset_str}"); + let offset_index_length_str = match cc_metadata.offset_index_length() { + None => "N/A".to_owned(), + Some(oil) => oil.to_string(), + }; + writeln!(out, "offset index length: {offset_index_length_str}"); + let column_index_offset_str = match cc_metadata.column_index_offset() { + None => "N/A".to_owned(), + Some(cio) => cio.to_string(), + }; + writeln!(out, "column index offset: {column_index_offset_str}"); + let column_index_length_str = match cc_metadata.column_index_length() { + None => "N/A".to_owned(), + Some(cil) => cil.to_string(), + }; + writeln!(out, "column index length: {column_index_length_str}"); + writeln!(out); +} + +#[allow(unused_must_use)] +fn print_dashes(out: &mut dyn io::Write, num: i32) { + for _ in 0..num { + write!(out, "-"); + } + writeln!(out); +} + +const INDENT_WIDTH: i32 = 2; + +/// Struct for printing Parquet message type. +struct Printer<'a> { + output: &'a mut dyn fmt::Write, + indent: i32, +} + +#[allow(unused_must_use)] +impl<'a> Printer<'a> { + fn new(output: &'a mut dyn fmt::Write) -> Self { + Printer { output, indent: 0 } + } + + fn print_indent(&mut self) { + for _ in 0..self.indent { + write!(self.output, " "); + } + } +} + +#[inline] +fn print_timeunit(unit: &TimeUnit) -> &str { + match unit { + TimeUnit::MILLIS => "MILLIS", + TimeUnit::MICROS => "MICROS", + TimeUnit::NANOS => "NANOS", + } +} + +#[inline] +fn print_logical_and_converted( + logical_type: Option<&LogicalType>, + converted_type: ConvertedType, + precision: i32, + scale: i32, +) -> String { + match logical_type { + Some(logical_type) => match logical_type { + LogicalType::Integer(IntType { + bit_width, + is_signed, + }) => { + format!("INTEGER({bit_width},{is_signed})") + } + LogicalType::Decimal(DecimalType { scale, precision }) => { + format!("DECIMAL({precision},{scale})") + } + LogicalType::Timestamp(timestamp) => { + format!( + "TIMESTAMP({},{})", + print_timeunit(×tamp.unit), + timestamp.is_adjusted_to_u_t_c + ) + } + LogicalType::Time(time) => { + format!( + "TIME({},{})", + print_timeunit(&time.unit), + time.is_adjusted_to_u_t_c + ) + } + LogicalType::Date => "DATE".to_string(), + LogicalType::Bson => "BSON".to_string(), + LogicalType::Json => "JSON".to_string(), + LogicalType::String => "STRING".to_string(), + LogicalType::Uuid => "UUID".to_string(), + LogicalType::Enum => "ENUM".to_string(), + LogicalType::List => "LIST".to_string(), + LogicalType::Map => "MAP".to_string(), + LogicalType::Float16 => "FLOAT16".to_string(), + LogicalType::Variant(VariantType { + specification_version, + }) => format!("VARIANT({specification_version:?})"), + LogicalType::Geometry(GeometryType { crs }) => { + if let Some(crs) = crs { + format!("GEOMETRY({crs})") + } else { + "GEOMETRY".to_string() + } + } + LogicalType::Geography(GeographyType { crs, algorithm }) => { + let algorithm = algorithm.unwrap_or_default(); + if let Some(crs) = crs { + format!("GEOGRAPHY({algorithm}, {crs})") + } else { + format!("GEOGRAPHY({algorithm})") + } + } + LogicalType::Unknown => "UNKNOWN".to_string(), + LogicalType::_Unknown { field_id } => format!("_Unknown({field_id})"), + }, + None => { + // Also print converted type if it is available + match converted_type { + ConvertedType::NONE => String::new(), + decimal @ ConvertedType::DECIMAL => { + // For decimal type we should print precision and scale if they + // are > 0, e.g. DECIMAL(9,2) - + // DECIMAL(9) - DECIMAL + let precision_scale = match (precision, scale) { + (p, s) if p > 0 && s > 0 => { + format!("({p},{s})") + } + (p, 0) if p > 0 => format!("({p})"), + _ => String::new(), + }; + format!("{decimal}{precision_scale}") + } + other_converted_type => { + format!("{other_converted_type}") + } + } + } + } +} + +#[allow(unused_must_use)] +impl Printer<'_> { + pub fn print(&mut self, tp: &Type) { + self.print_indent(); + match *tp { + Type::PrimitiveType { + ref basic_info, + physical_type, + type_length, + scale, + precision, + } => { + let phys_type_str = match physical_type { + PhysicalType::FIXED_LEN_BYTE_ARRAY => { + // We need to include length for fixed byte array + format!("{physical_type} ({type_length})") + } + _ => format!("{physical_type}"), + }; + write!( + self.output, + "{} {} {}", + basic_info.repetition(), + phys_type_str, + basic_info.name() + ); + if basic_info.has_id() { + write!(self.output, " [{}]", basic_info.id()); + } + // Also print logical type if it is available + // If there is a logical type, do not print converted type + let logical_type_str = print_logical_and_converted( + basic_info.logical_type_ref(), + basic_info.converted_type(), + precision, + scale, + ); + if !logical_type_str.is_empty() { + write!(self.output, " ({logical_type_str});"); + } else { + write!(self.output, ";"); + } + } + Type::GroupType { + ref basic_info, + ref fields, + } => { + if basic_info.has_repetition() { + write!( + self.output, + "{} group {} ", + basic_info.repetition(), + basic_info.name() + ); + if basic_info.has_id() { + write!(self.output, "[{}] ", basic_info.id()); + } + let logical_str = print_logical_and_converted( + basic_info.logical_type_ref(), + basic_info.converted_type(), + 0, + 0, + ); + if !logical_str.is_empty() { + write!(self.output, "({logical_str}) "); + } + } else { + write!(self.output, "message {} ", basic_info.name()); + if basic_info.has_id() { + write!(self.output, "[{}] ", basic_info.id()); + } + } + writeln!(self.output, "{{"); + + self.indent += INDENT_WIDTH; + for c in fields { + self.print(c); + writeln!(self.output); + } + self.indent -= INDENT_WIDTH; + self.print_indent(); + write!(self.output, "}}"); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::sync::Arc; + + use crate::basic::{Repetition, Type as PhysicalType}; + use crate::errors::Result; + use crate::schema::parser::parse_message_type; + + fn assert_print_parse_message(message: Type) { + let mut s = String::new(); + { + let mut p = Printer::new(&mut s); + p.print(&message); + } + println!("{s}"); + let parsed = parse_message_type(&s).unwrap(); + assert_eq!(message, parsed); + } + + #[test] + fn test_print_primitive_type() { + let types_and_strings = vec![ + ( + Type::primitive_type_builder("field", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::INT_32) + .build() + .unwrap(), + "REQUIRED INT32 field (INT_32);", + ), + ( + Type::primitive_type_builder("field", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::INT_32) + .with_id(Some(42)) + .build() + .unwrap(), + "REQUIRED INT32 field [42] (INT_32);", + ), + ( + Type::primitive_type_builder("field", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(), + "REQUIRED INT32 field;", + ), + ( + Type::primitive_type_builder("field", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .with_id(Some(42)) + .build() + .unwrap(), + "REQUIRED INT32 field [42];", + ), + ]; + types_and_strings.into_iter().for_each(|(field, expected)| { + let mut s = String::new(); + { + let mut p = Printer::new(&mut s); + p.print(&field); + } + assert_eq!(&s, expected) + }); + } + + #[inline] + fn build_primitive_type( + name: &str, + id: Option, + physical_type: PhysicalType, + logical_type: Option, + converted_type: ConvertedType, + repetition: Repetition, + ) -> Result { + Type::primitive_type_builder(name, physical_type) + .with_id(id) + .with_repetition(repetition) + .with_logical_type(logical_type) + .with_converted_type(converted_type) + .build() + } + + #[test] + fn test_print_logical_types() { + let types_and_strings = vec![ + ( + build_primitive_type( + "field", + None, + PhysicalType::INT32, + Some(LogicalType::integer(32, true)), + ConvertedType::NONE, + Repetition::REQUIRED, + ) + .unwrap(), + "REQUIRED INT32 field (INTEGER(32,true));", + ), + ( + build_primitive_type( + "field", + None, + PhysicalType::INT32, + Some(LogicalType::integer(8, false)), + ConvertedType::NONE, + Repetition::OPTIONAL, + ) + .unwrap(), + "OPTIONAL INT32 field (INTEGER(8,false));", + ), + ( + build_primitive_type( + "field", + None, + PhysicalType::INT32, + Some(LogicalType::integer(16, true)), + ConvertedType::INT_16, + Repetition::REPEATED, + ) + .unwrap(), + "REPEATED INT32 field (INTEGER(16,true));", + ), + ( + build_primitive_type( + "field", + Some(42), + PhysicalType::INT32, + Some(LogicalType::integer(16, true)), + ConvertedType::INT_16, + Repetition::REPEATED, + ) + .unwrap(), + "REPEATED INT32 field [42] (INTEGER(16,true));", + ), + ( + build_primitive_type( + "field", + None, + PhysicalType::INT64, + None, + ConvertedType::NONE, + Repetition::REPEATED, + ) + .unwrap(), + "REPEATED INT64 field;", + ), + ( + build_primitive_type( + "field", + None, + PhysicalType::FLOAT, + None, + ConvertedType::NONE, + Repetition::REQUIRED, + ) + .unwrap(), + "REQUIRED FLOAT field;", + ), + ( + build_primitive_type( + "booleans", + None, + PhysicalType::BOOLEAN, + None, + ConvertedType::NONE, + Repetition::OPTIONAL, + ) + .unwrap(), + "OPTIONAL BOOLEAN booleans;", + ), + ( + build_primitive_type( + "booleans", + Some(42), + PhysicalType::BOOLEAN, + None, + ConvertedType::NONE, + Repetition::OPTIONAL, + ) + .unwrap(), + "OPTIONAL BOOLEAN booleans [42];", + ), + ( + build_primitive_type( + "field", + None, + PhysicalType::INT64, + Some(LogicalType::timestamp(true, TimeUnit::MILLIS)), + ConvertedType::NONE, + Repetition::REQUIRED, + ) + .unwrap(), + "REQUIRED INT64 field (TIMESTAMP(MILLIS,true));", + ), + ( + build_primitive_type( + "field", + None, + PhysicalType::INT32, + Some(LogicalType::Date), + ConvertedType::NONE, + Repetition::OPTIONAL, + ) + .unwrap(), + "OPTIONAL INT32 field (DATE);", + ), + ( + build_primitive_type( + "field", + None, + PhysicalType::INT32, + Some(LogicalType::time(false, TimeUnit::MILLIS)), + ConvertedType::TIME_MILLIS, + Repetition::REQUIRED, + ) + .unwrap(), + "REQUIRED INT32 field (TIME(MILLIS,false));", + ), + ( + build_primitive_type( + "field", + Some(42), + PhysicalType::INT32, + Some(LogicalType::time(false, TimeUnit::MILLIS)), + ConvertedType::TIME_MILLIS, + Repetition::REQUIRED, + ) + .unwrap(), + "REQUIRED INT32 field [42] (TIME(MILLIS,false));", + ), + ( + build_primitive_type( + "field", + None, + PhysicalType::BYTE_ARRAY, + None, + ConvertedType::NONE, + Repetition::REQUIRED, + ) + .unwrap(), + "REQUIRED BYTE_ARRAY field;", + ), + ( + build_primitive_type( + "field", + Some(42), + PhysicalType::BYTE_ARRAY, + None, + ConvertedType::NONE, + Repetition::REQUIRED, + ) + .unwrap(), + "REQUIRED BYTE_ARRAY field [42];", + ), + ( + build_primitive_type( + "field", + None, + PhysicalType::BYTE_ARRAY, + None, + ConvertedType::UTF8, + Repetition::REQUIRED, + ) + .unwrap(), + "REQUIRED BYTE_ARRAY field (UTF8);", + ), + ( + build_primitive_type( + "field", + None, + PhysicalType::BYTE_ARRAY, + Some(LogicalType::Json), + ConvertedType::JSON, + Repetition::REQUIRED, + ) + .unwrap(), + "REQUIRED BYTE_ARRAY field (JSON);", + ), + ( + build_primitive_type( + "field", + None, + PhysicalType::BYTE_ARRAY, + Some(LogicalType::Bson), + ConvertedType::BSON, + Repetition::REQUIRED, + ) + .unwrap(), + "REQUIRED BYTE_ARRAY field (BSON);", + ), + ( + build_primitive_type( + "field", + None, + PhysicalType::BYTE_ARRAY, + Some(LogicalType::String), + ConvertedType::NONE, + Repetition::REQUIRED, + ) + .unwrap(), + "REQUIRED BYTE_ARRAY field (STRING);", + ), + ( + build_primitive_type( + "field", + Some(42), + PhysicalType::BYTE_ARRAY, + Some(LogicalType::String), + ConvertedType::NONE, + Repetition::REQUIRED, + ) + .unwrap(), + "REQUIRED BYTE_ARRAY field [42] (STRING);", + ), + ( + build_primitive_type( + "field", + None, + PhysicalType::BYTE_ARRAY, + Some(LogicalType::geometry(None)), + ConvertedType::NONE, + Repetition::REQUIRED, + ) + .unwrap(), + "REQUIRED BYTE_ARRAY field (GEOMETRY);", + ), + ( + build_primitive_type( + "field", + None, + PhysicalType::BYTE_ARRAY, + Some(LogicalType::geometry(Some("non-missing CRS".to_string()))), + ConvertedType::NONE, + Repetition::REQUIRED, + ) + .unwrap(), + "REQUIRED BYTE_ARRAY field (GEOMETRY(non-missing CRS));", + ), + ( + build_primitive_type( + "field", + None, + PhysicalType::BYTE_ARRAY, + Some(LogicalType::geography(None, Some(Default::default()))), + ConvertedType::NONE, + Repetition::REQUIRED, + ) + .unwrap(), + "REQUIRED BYTE_ARRAY field (GEOGRAPHY(SPHERICAL));", + ), + ( + build_primitive_type( + "field", + None, + PhysicalType::BYTE_ARRAY, + Some(LogicalType::geography( + Some("non-missing CRS".to_string()), + Some(Default::default()), + )), + ConvertedType::NONE, + Repetition::REQUIRED, + ) + .unwrap(), + "REQUIRED BYTE_ARRAY field (GEOGRAPHY(SPHERICAL, non-missing CRS));", + ), + ]; + + types_and_strings.into_iter().for_each(|(field, expected)| { + let mut s = String::new(); + { + let mut p = Printer::new(&mut s); + p.print(&field); + } + assert_eq!(&s, expected) + }); + } + + #[inline] + fn decimal_length_from_precision(precision: usize) -> i32 { + let max_val = 10.0_f64.powi(precision as i32) - 1.0; + let bits_unsigned = max_val.log2().ceil(); + let bits_signed = bits_unsigned + 1.0; + (bits_signed / 8.0).ceil() as i32 + } + + #[test] + fn test_print_flba_logical_types() { + let types_and_strings = vec![ + ( + Type::primitive_type_builder("field", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_logical_type(None) + .with_converted_type(ConvertedType::INTERVAL) + .with_length(12) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(), + "REQUIRED FIXED_LEN_BYTE_ARRAY (12) field (INTERVAL);", + ), + ( + Type::primitive_type_builder("field", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_logical_type(Some(LogicalType::Uuid)) + .with_length(16) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(), + "REQUIRED FIXED_LEN_BYTE_ARRAY (16) field (UUID);", + ), + ( + Type::primitive_type_builder("decimal", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_logical_type(Some(LogicalType::decimal(20, 32))) + .with_precision(32) + .with_scale(20) + .with_length(decimal_length_from_precision(32)) + .with_repetition(Repetition::REPEATED) + .build() + .unwrap(), + "REPEATED FIXED_LEN_BYTE_ARRAY (14) decimal (DECIMAL(32,20));", + ), + ( + Type::primitive_type_builder("decimal", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_converted_type(ConvertedType::DECIMAL) + .with_precision(19) + .with_scale(4) + .with_length(decimal_length_from_precision(19)) + .with_repetition(Repetition::OPTIONAL) + .build() + .unwrap(), + "OPTIONAL FIXED_LEN_BYTE_ARRAY (9) decimal (DECIMAL(19,4));", + ), + ( + Type::primitive_type_builder("float16", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_logical_type(Some(LogicalType::Float16)) + .with_length(2) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(), + "REQUIRED FIXED_LEN_BYTE_ARRAY (2) float16 (FLOAT16);", + ), + ]; + + types_and_strings.into_iter().for_each(|(field, expected)| { + let mut s = String::new(); + { + let mut p = Printer::new(&mut s); + p.print(&field); + } + assert_eq!(&s, expected) + }); + } + + #[test] + fn test_print_schema_documentation() { + let mut s = String::new(); + { + let mut p = Printer::new(&mut s); + let field_a = Type::primitive_type_builder("a", PhysicalType::BYTE_ARRAY) + .with_id(Some(42)) + .with_converted_type(ConvertedType::UTF8) + .build() + .unwrap(); + + let field_b = Type::primitive_type_builder("b", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(); + + let field_d = Type::primitive_type_builder("d", PhysicalType::INT64) + .with_id(Some(99)) + .build() + .unwrap(); + + let field_c = Type::group_type_builder("c") + .with_id(Some(43)) + .with_fields(vec![Arc::new(field_d)]) + .build() + .unwrap(); + + let schema = Type::group_type_builder("schema") + .with_fields(vec![ + Arc::new(field_a), + Arc::new(field_b), + Arc::new(field_c), + ]) + .build() + .unwrap(); + p.print(&schema); + } + let expected = "message schema { + OPTIONAL BYTE_ARRAY a [42] (UTF8); + REQUIRED INT32 b; + message c [43] { + OPTIONAL INT64 d [99]; + } +}"; + assert_eq!(&mut s, expected); + } + + #[test] + fn test_print_group_type() { + let mut s = String::new(); + { + let mut p = Printer::new(&mut s); + let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::INT_32) + .build(); + let f2 = Type::primitive_type_builder("f2", PhysicalType::BYTE_ARRAY) + .with_converted_type(ConvertedType::UTF8) + .build(); + let f3 = Type::primitive_type_builder("f3", PhysicalType::BYTE_ARRAY) + .with_logical_type(Some(LogicalType::String)) + .build(); + let f4 = Type::primitive_type_builder("f4", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_repetition(Repetition::REPEATED) + .with_converted_type(ConvertedType::INTERVAL) + .with_length(12) + .build(); + + let struct_fields = vec![ + Arc::new(f1.unwrap()), + Arc::new(f2.unwrap()), + Arc::new(f3.unwrap()), + ]; + let field = Type::group_type_builder("field") + .with_repetition(Repetition::OPTIONAL) + .with_fields(struct_fields) + .build() + .unwrap(); + + let fields = vec![Arc::new(field), Arc::new(f4.unwrap())]; + let message = Type::group_type_builder("schema") + .with_fields(fields) + .build() + .unwrap(); + p.print(&message); + } + let expected = "message schema { + OPTIONAL group field { + REQUIRED INT32 f1 (INT_32); + OPTIONAL BYTE_ARRAY f2 (UTF8); + OPTIONAL BYTE_ARRAY f3 (STRING); + } + REPEATED FIXED_LEN_BYTE_ARRAY (12) f4 (INTERVAL); +}"; + assert_eq!(&mut s, expected); + } + + #[test] + fn test_print_group_type_with_ids() { + let mut s = String::new(); + { + let mut p = Printer::new(&mut s); + let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::INT_32) + .with_id(Some(0)) + .build(); + let f2 = Type::primitive_type_builder("f2", PhysicalType::BYTE_ARRAY) + .with_converted_type(ConvertedType::UTF8) + .with_id(Some(1)) + .build(); + let f3 = Type::primitive_type_builder("f3", PhysicalType::BYTE_ARRAY) + .with_logical_type(Some(LogicalType::String)) + .with_id(Some(1)) + .build(); + let f4 = Type::primitive_type_builder("f4", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_repetition(Repetition::REPEATED) + .with_converted_type(ConvertedType::INTERVAL) + .with_length(12) + .with_id(Some(2)) + .build(); + + let struct_fields = vec![ + Arc::new(f1.unwrap()), + Arc::new(f2.unwrap()), + Arc::new(f3.unwrap()), + ]; + let field = Type::group_type_builder("field") + .with_repetition(Repetition::OPTIONAL) + .with_fields(struct_fields) + .with_id(Some(1)) + .build() + .unwrap(); + + let fields = vec![Arc::new(field), Arc::new(f4.unwrap())]; + let message = Type::group_type_builder("schema") + .with_fields(fields) + .with_id(Some(2)) + .build() + .unwrap(); + p.print(&message); + } + let expected = "message schema [2] { + OPTIONAL group field [1] { + REQUIRED INT32 f1 [0] (INT_32); + OPTIONAL BYTE_ARRAY f2 [1] (UTF8); + OPTIONAL BYTE_ARRAY f3 [1] (STRING); + } + REPEATED FIXED_LEN_BYTE_ARRAY (12) f4 [2] (INTERVAL); +}"; + assert_eq!(&mut s, expected); + } + + #[test] + fn test_print_and_parse_primitive() { + let a2 = Type::primitive_type_builder("a2", PhysicalType::BYTE_ARRAY) + .with_repetition(Repetition::REPEATED) + .with_converted_type(ConvertedType::UTF8) + .build() + .unwrap(); + + let a1 = Type::group_type_builder("a1") + .with_repetition(Repetition::OPTIONAL) + .with_logical_type(Some(LogicalType::List)) + .with_converted_type(ConvertedType::LIST) + .with_fields(vec![Arc::new(a2)]) + .build() + .unwrap(); + + let b3 = Type::primitive_type_builder("b3", PhysicalType::INT32) + .with_repetition(Repetition::OPTIONAL) + .build() + .unwrap(); + + let b4 = Type::primitive_type_builder("b4", PhysicalType::DOUBLE) + .with_repetition(Repetition::OPTIONAL) + .build() + .unwrap(); + + let b2 = Type::group_type_builder("b2") + .with_repetition(Repetition::REPEATED) + .with_converted_type(ConvertedType::NONE) + .with_fields(vec![Arc::new(b3), Arc::new(b4)]) + .build() + .unwrap(); + + let b1 = Type::group_type_builder("b1") + .with_repetition(Repetition::OPTIONAL) + .with_logical_type(Some(LogicalType::List)) + .with_converted_type(ConvertedType::LIST) + .with_fields(vec![Arc::new(b2)]) + .build() + .unwrap(); + + let a0 = Type::group_type_builder("a0") + .with_repetition(Repetition::REQUIRED) + .with_fields(vec![Arc::new(a1), Arc::new(b1)]) + .build() + .unwrap(); + + let message = Type::group_type_builder("root") + .with_fields(vec![Arc::new(a0)]) + .build() + .unwrap(); + + assert_print_parse_message(message); + } + + #[test] + fn test_print_and_parse_nested() { + let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::INT_32) + .build() + .unwrap(); + + let f2 = Type::primitive_type_builder("f2", PhysicalType::BYTE_ARRAY) + .with_repetition(Repetition::OPTIONAL) + .with_converted_type(ConvertedType::UTF8) + .build() + .unwrap(); + + let field = Type::group_type_builder("field") + .with_repetition(Repetition::OPTIONAL) + .with_fields(vec![Arc::new(f1), Arc::new(f2)]) + .build() + .unwrap(); + + let f3 = Type::primitive_type_builder("f3", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_repetition(Repetition::REPEATED) + .with_converted_type(ConvertedType::INTERVAL) + .with_length(12) + .build() + .unwrap(); + + let message = Type::group_type_builder("schema") + .with_fields(vec![Arc::new(field), Arc::new(f3)]) + .build() + .unwrap(); + + assert_print_parse_message(message); + } + + #[test] + fn test_print_and_parse_decimal() { + let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32) + .with_repetition(Repetition::OPTIONAL) + .with_logical_type(Some(LogicalType::decimal(2, 9))) + .with_converted_type(ConvertedType::DECIMAL) + .with_precision(9) + .with_scale(2) + .build() + .unwrap(); + + let f2 = Type::primitive_type_builder("f2", PhysicalType::INT32) + .with_repetition(Repetition::OPTIONAL) + .with_logical_type(Some(LogicalType::decimal(0, 9))) + .with_converted_type(ConvertedType::DECIMAL) + .with_precision(9) + .with_scale(0) + .build() + .unwrap(); + + let message = Type::group_type_builder("schema") + .with_fields(vec![Arc::new(f1), Arc::new(f2)]) + .build() + .unwrap(); + + assert_print_parse_message(message); + } +} diff --git a/vendor/parquet/src/schema/types.rs b/vendor/parquet/src/schema/types.rs new file mode 100644 index 0000000..1f9b859 --- /dev/null +++ b/vendor/parquet/src/schema/types.rs @@ -0,0 +1,2525 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Contains structs and methods to build Parquet schema and schema descriptors. + +use std::vec::IntoIter; +use std::{collections::HashMap, fmt, sync::Arc}; + +use crate::file::metadata::HeapSize; +use crate::file::metadata::thrift::SchemaElement; + +use crate::basic::{ + ColumnOrder, ConvertedType, IntType, LogicalType, Repetition, SortOrder, TimeType, TimeUnit, + Type as PhysicalType, +}; +use crate::errors::{ParquetError, Result}; + +// ---------------------------------------------------------------------- +// Parquet Type definitions + +/// Type alias for `Arc`. +pub type TypePtr = Arc; +/// Type alias for `Arc`. +pub type SchemaDescPtr = Arc; +/// Type alias for `Arc`. +pub type ColumnDescPtr = Arc; + +/// Representation of a Parquet type. +/// +/// Used to describe primitive leaf fields and structs, including top-level schema. +/// +/// Note that the top-level schema is represented using [`Type::GroupType`] whose +/// repetition is `None`. +#[derive(Clone, Debug, PartialEq)] +pub enum Type { + /// Represents a primitive leaf field. + PrimitiveType { + /// Basic information about the type. + basic_info: BasicTypeInfo, + /// Physical type of this primitive type. + physical_type: PhysicalType, + /// Length of this type. + type_length: i32, + /// Scale of this type. + scale: i32, + /// Precision of this type. + precision: i32, + }, + /// Represents a group of fields (similar to struct). + GroupType { + /// Basic information about the type. + basic_info: BasicTypeInfo, + /// Fields of this group type. + fields: Vec, + }, +} + +impl HeapSize for Type { + fn heap_size(&self) -> usize { + match self { + Type::PrimitiveType { basic_info, .. } => basic_info.heap_size(), + Type::GroupType { basic_info, fields } => basic_info.heap_size() + fields.heap_size(), + } + } +} + +impl Type { + /// Creates primitive type builder with provided field name and physical type. + pub fn primitive_type_builder( + name: &str, + physical_type: PhysicalType, + ) -> PrimitiveTypeBuilder<'_> { + PrimitiveTypeBuilder::new(name, physical_type) + } + + /// Creates group type builder with provided column name. + pub fn group_type_builder(name: &str) -> GroupTypeBuilder<'_> { + GroupTypeBuilder::new(name) + } + + /// Returns [`BasicTypeInfo`] information about the type. + pub fn get_basic_info(&self) -> &BasicTypeInfo { + match *self { + Type::PrimitiveType { ref basic_info, .. } => basic_info, + Type::GroupType { ref basic_info, .. } => basic_info, + } + } + + /// Returns this type's field name. + pub fn name(&self) -> &str { + self.get_basic_info().name() + } + + /// Gets the fields from this group type. + /// Note that this will panic if called on a non-group type. + // TODO: should we return `&[&Type]` here? + pub fn get_fields(&self) -> &[TypePtr] { + match *self { + Type::GroupType { ref fields, .. } => &fields[..], + _ => panic!("Cannot call get_fields() on a non-group type"), + } + } + + /// Gets physical type of this primitive type. + /// Note that this will panic if called on a non-primitive type. + pub fn get_physical_type(&self) -> PhysicalType { + match *self { + Type::PrimitiveType { + basic_info: _, + physical_type, + .. + } => physical_type, + _ => panic!("Cannot call get_physical_type() on a non-primitive type"), + } + } + + /// Gets precision of this primitive type. + /// Note that this will panic if called on a non-primitive type. + pub fn get_precision(&self) -> i32 { + match *self { + Type::PrimitiveType { precision, .. } => precision, + _ => panic!("Cannot call get_precision() on non-primitive type"), + } + } + + /// Gets scale of this primitive type. + /// Note that this will panic if called on a non-primitive type. + pub fn get_scale(&self) -> i32 { + match *self { + Type::PrimitiveType { scale, .. } => scale, + _ => panic!("Cannot call get_scale() on non-primitive type"), + } + } + + /// Checks if `sub_type` schema is part of current schema. + /// This method can be used to check if projected columns are part of the root schema. + pub fn check_contains(&self, sub_type: &Type) -> bool { + // Names match, and repetitions match or not set for both + let basic_match = self.get_basic_info().name() == sub_type.get_basic_info().name() + && (self.is_schema() && sub_type.is_schema() + || !self.is_schema() + && !sub_type.is_schema() + && self.get_basic_info().repetition() + == sub_type.get_basic_info().repetition()); + + match *self { + Type::PrimitiveType { .. } if basic_match && sub_type.is_primitive() => { + self.get_physical_type() == sub_type.get_physical_type() + } + Type::GroupType { .. } if basic_match && sub_type.is_group() => { + // build hashmap of name -> TypePtr + let mut field_map = HashMap::new(); + for field in self.get_fields() { + field_map.insert(field.name(), field); + } + + for field in sub_type.get_fields() { + if !field_map + .get(field.name()) + .map(|tpe| tpe.check_contains(field)) + .unwrap_or(false) + { + return false; + } + } + true + } + _ => false, + } + } + + /// Returns `true` if this type is a primitive type, `false` otherwise. + pub fn is_primitive(&self) -> bool { + matches!(*self, Type::PrimitiveType { .. }) + } + + /// Returns `true` if this type is a group type, `false` otherwise. + pub fn is_group(&self) -> bool { + matches!(*self, Type::GroupType { .. }) + } + + /// Returns `true` if this type is the top-level schema type (message type). + pub fn is_schema(&self) -> bool { + match *self { + Type::GroupType { ref basic_info, .. } => !basic_info.has_repetition(), + _ => false, + } + } + + /// Returns `true` if this type is repeated or optional. + /// If this type doesn't have repetition defined, we treat it as required. + pub fn is_optional(&self) -> bool { + self.get_basic_info().has_repetition() + && self.get_basic_info().repetition() != Repetition::REQUIRED + } + + /// Returns `true` if this type is annotated as a list. + pub(crate) fn is_list(&self) -> bool { + if self.is_group() { + let basic_info = self.get_basic_info(); + if let Some(logical_type) = basic_info.logical_type_ref() { + return logical_type == &LogicalType::List; + } + return basic_info.converted_type() == ConvertedType::LIST; + } + false + } + + /// Returns `true` if this type is a group with a single child field that is `repeated`. + pub(crate) fn has_single_repeated_child(&self) -> bool { + if self.is_group() { + let children = self.get_fields(); + return children.len() == 1 + && children[0].get_basic_info().has_repetition() + && children[0].get_basic_info().repetition() == Repetition::REPEATED; + } + false + } +} + +/// A builder for primitive types. All attributes are optional +/// except the name and physical type. +/// Note that if not specified explicitly, `Repetition::OPTIONAL` is used. +pub struct PrimitiveTypeBuilder<'a> { + name: &'a str, + repetition: Repetition, + physical_type: PhysicalType, + converted_type: ConvertedType, + logical_type: Option, + length: i32, + precision: i32, + scale: i32, + id: Option, +} + +impl<'a> PrimitiveTypeBuilder<'a> { + /// Creates new primitive type builder with provided field name and physical type. + pub fn new(name: &'a str, physical_type: PhysicalType) -> Self { + Self { + name, + repetition: Repetition::OPTIONAL, + physical_type, + converted_type: ConvertedType::NONE, + logical_type: None, + length: -1, + precision: -1, + scale: -1, + id: None, + } + } + + /// Sets [`Repetition`] for this field and returns itself. + pub fn with_repetition(self, repetition: Repetition) -> Self { + Self { repetition, ..self } + } + + /// Sets [`ConvertedType`] for this field and returns itself. + pub fn with_converted_type(self, converted_type: ConvertedType) -> Self { + Self { + converted_type, + ..self + } + } + + /// Sets [`LogicalType`] for this field and returns itself. + /// If only the logical type is populated for a primitive type, the converted type + /// will be automatically populated, and can thus be omitted. + pub fn with_logical_type(self, logical_type: Option) -> Self { + Self { + logical_type, + ..self + } + } + + /// Sets type length and returns itself. + /// This is only applied to FIXED_LEN_BYTE_ARRAY and INT96 (INTERVAL) types, because + /// they maintain fixed size underlying byte array. + /// By default, value is `0`. + pub fn with_length(self, length: i32) -> Self { + Self { length, ..self } + } + + /// Sets precision for Parquet DECIMAL physical type and returns itself. + /// By default, it equals to `0` and used only for decimal context. + pub fn with_precision(self, precision: i32) -> Self { + Self { precision, ..self } + } + + /// Sets scale for Parquet DECIMAL physical type and returns itself. + /// By default, it equals to `0` and used only for decimal context. + pub fn with_scale(self, scale: i32) -> Self { + Self { scale, ..self } + } + + /// Sets optional field id and returns itself. + pub fn with_id(self, id: Option) -> Self { + Self { id, ..self } + } + + /// Creates a new `PrimitiveType` instance from the collected attributes. + /// Returns `Err` in case of any building conditions are not met. + pub fn build(self) -> Result { + let mut basic_info = BasicTypeInfo { + name: String::from(self.name), + repetition: Some(self.repetition), + converted_type: self.converted_type, + logical_type: self.logical_type.clone(), + id: self.id, + }; + + // Check length before logical type, since it is used for logical type validation. + if self.physical_type == PhysicalType::FIXED_LEN_BYTE_ARRAY && self.length < 0 { + return Err(general_err!( + "Invalid FIXED_LEN_BYTE_ARRAY length: {} for field '{}'", + self.length, + self.name + )); + } + + if let Some(logical_type) = &self.logical_type { + // If a converted type is populated, check that it is consistent with + // its logical type + if self.converted_type != ConvertedType::NONE { + if ConvertedType::from(self.logical_type.clone()) != self.converted_type { + return Err(general_err!( + "Logical type {:?} is incompatible with converted type {} for field '{}'", + logical_type, + self.converted_type, + self.name + )); + } + } else { + // Populate the converted type for backwards compatibility + basic_info.converted_type = self.logical_type.clone().into(); + } + // Check that logical type and physical type are compatible + match (logical_type, self.physical_type) { + (LogicalType::Map, _) | (LogicalType::List, _) => { + return Err(general_err!( + "{:?} cannot be applied to a primitive type for field '{}'", + logical_type, + self.name + )); + } + (LogicalType::Enum, PhysicalType::BYTE_ARRAY) => {} + (LogicalType::Decimal(decimal), _) => { + // Check that scale and precision are consistent with legacy values + if decimal.scale != self.scale { + return Err(general_err!( + "DECIMAL logical type scale {} must match self.scale {} for field '{}'", + decimal.scale, + self.scale, + self.name + )); + } + if decimal.precision != self.precision { + return Err(general_err!( + "DECIMAL logical type precision {} must match self.precision {} for field '{}'", + decimal.precision, + self.precision, + self.name + )); + } + self.check_decimal_precision_scale()?; + } + (LogicalType::Date, PhysicalType::INT32) => {} + ( + LogicalType::Time(TimeType { + unit: TimeUnit::MILLIS, + .. + }), + PhysicalType::INT32, + ) => {} + (LogicalType::Time(time), PhysicalType::INT64) => { + if time.unit == TimeUnit::MILLIS { + return Err(general_err!( + "Cannot use millisecond unit on INT64 type for field '{}'", + self.name + )); + } + } + (LogicalType::Timestamp(_), PhysicalType::INT64) => {} + (LogicalType::Integer(int), PhysicalType::INT32) if int.bit_width <= 32 => {} + (LogicalType::Integer(int), PhysicalType::INT64) if int.bit_width == 64 => {} + // Null type + (LogicalType::Unknown, _) => {} + (LogicalType::String, PhysicalType::BYTE_ARRAY) => {} + (LogicalType::Json, PhysicalType::BYTE_ARRAY) => {} + (LogicalType::Bson, PhysicalType::BYTE_ARRAY) => {} + (LogicalType::Geometry(_), PhysicalType::BYTE_ARRAY) => {} + (LogicalType::Geography(_), PhysicalType::BYTE_ARRAY) => {} + (LogicalType::Uuid, PhysicalType::FIXED_LEN_BYTE_ARRAY) if self.length == 16 => {} + (LogicalType::Uuid, PhysicalType::FIXED_LEN_BYTE_ARRAY) => { + return Err(general_err!( + "UUID cannot annotate field '{}' because it is not a FIXED_LEN_BYTE_ARRAY(16) field", + self.name + )); + } + (LogicalType::Float16, PhysicalType::FIXED_LEN_BYTE_ARRAY) if self.length == 2 => {} + (LogicalType::Float16, PhysicalType::FIXED_LEN_BYTE_ARRAY) => { + return Err(general_err!( + "FLOAT16 cannot annotate field '{}' because it is not a FIXED_LEN_BYTE_ARRAY(2) field", + self.name + )); + } + // unknown logical type means just use physical type + (LogicalType::_Unknown { .. }, _) => {} + (a, b) => { + return Err(general_err!( + "Cannot annotate {:?} from {} for field '{}'", + a, + b, + self.name + )); + } + } + } + + match self.converted_type { + ConvertedType::NONE => {} + ConvertedType::UTF8 | ConvertedType::BSON | ConvertedType::JSON => { + if self.physical_type != PhysicalType::BYTE_ARRAY { + return Err(general_err!( + "{} cannot annotate field '{}' because it is not a BYTE_ARRAY field", + self.converted_type, + self.name + )); + } + } + ConvertedType::DECIMAL => { + self.check_decimal_precision_scale()?; + } + ConvertedType::DATE + | ConvertedType::TIME_MILLIS + | ConvertedType::UINT_8 + | ConvertedType::UINT_16 + | ConvertedType::UINT_32 + | ConvertedType::INT_8 + | ConvertedType::INT_16 + | ConvertedType::INT_32 => { + if self.physical_type != PhysicalType::INT32 { + return Err(general_err!( + "{} cannot annotate field '{}' because it is not a INT32 field", + self.converted_type, + self.name + )); + } + } + ConvertedType::TIME_MICROS + | ConvertedType::TIMESTAMP_MILLIS + | ConvertedType::TIMESTAMP_MICROS + | ConvertedType::UINT_64 + | ConvertedType::INT_64 => { + if self.physical_type != PhysicalType::INT64 { + return Err(general_err!( + "{} cannot annotate field '{}' because it is not a INT64 field", + self.converted_type, + self.name + )); + } + } + ConvertedType::INTERVAL => { + if self.physical_type != PhysicalType::FIXED_LEN_BYTE_ARRAY || self.length != 12 { + return Err(general_err!( + "INTERVAL cannot annotate field '{}' because it is not a FIXED_LEN_BYTE_ARRAY(12) field", + self.name + )); + } + } + ConvertedType::ENUM => { + if self.physical_type != PhysicalType::BYTE_ARRAY { + return Err(general_err!( + "ENUM cannot annotate field '{}' because it is not a BYTE_ARRAY field", + self.name + )); + } + } + _ => { + return Err(general_err!( + "{} cannot be applied to primitive field '{}'", + self.converted_type, + self.name + )); + } + } + + Ok(Type::PrimitiveType { + basic_info, + physical_type: self.physical_type, + type_length: self.length, + scale: self.scale, + precision: self.precision, + }) + } + + #[inline] + fn check_decimal_precision_scale(&self) -> Result<()> { + match self.physical_type { + PhysicalType::INT32 + | PhysicalType::INT64 + | PhysicalType::BYTE_ARRAY + | PhysicalType::FIXED_LEN_BYTE_ARRAY => (), + _ => { + return Err(general_err!( + "DECIMAL can only annotate INT32, INT64, BYTE_ARRAY and FIXED_LEN_BYTE_ARRAY" + )); + } + } + + // Precision is required and must be a non-zero positive integer. + if self.precision < 1 { + return Err(general_err!( + "Invalid DECIMAL precision: {}", + self.precision + )); + } + + // Scale must be zero or a positive integer less than the precision. + if self.scale < 0 { + return Err(general_err!("Invalid DECIMAL scale: {}", self.scale)); + } + + if self.scale > self.precision { + return Err(general_err!( + "Invalid DECIMAL: scale ({}) cannot be greater than precision \ + ({})", + self.scale, + self.precision + )); + } + + // Check precision and scale based on physical type limitations. + match self.physical_type { + PhysicalType::INT32 => { + if self.precision > 9 { + return Err(general_err!( + "Cannot represent INT32 as DECIMAL with precision {}", + self.precision + )); + } + } + PhysicalType::INT64 => { + if self.precision > 18 { + return Err(general_err!( + "Cannot represent INT64 as DECIMAL with precision {}", + self.precision + )); + } + } + PhysicalType::FIXED_LEN_BYTE_ARRAY => { + let length = self + .length + .checked_mul(8) + .ok_or(general_err!("Invalid length {} for Decimal", self.length))?; + let max_precision = (2f64.powi(length - 1) - 1f64).log10().floor() as i32; + + if self.precision > max_precision { + return Err(general_err!( + "Cannot represent FIXED_LEN_BYTE_ARRAY as DECIMAL with length {} and \ + precision {}. The max precision can only be {}", + self.length, + self.precision, + max_precision + )); + } + } + _ => (), // For BYTE_ARRAY precision is not limited + } + + Ok(()) + } +} + +/// A builder for group types. All attributes are optional except the name. +/// Note that if not specified explicitly, `None` is used as the repetition of the group, +/// which means it is a root (message) type. +pub struct GroupTypeBuilder<'a> { + name: &'a str, + repetition: Option, + converted_type: ConvertedType, + logical_type: Option, + fields: Vec, + id: Option, +} + +impl<'a> GroupTypeBuilder<'a> { + /// Creates new group type builder with provided field name. + pub fn new(name: &'a str) -> Self { + Self { + name, + repetition: None, + converted_type: ConvertedType::NONE, + logical_type: None, + fields: Vec::new(), + id: None, + } + } + + /// Sets [`Repetition`] for this field and returns itself. + pub fn with_repetition(mut self, repetition: Repetition) -> Self { + self.repetition = Some(repetition); + self + } + + /// Sets [`ConvertedType`] for this field and returns itself. + pub fn with_converted_type(self, converted_type: ConvertedType) -> Self { + Self { + converted_type, + ..self + } + } + + /// Sets [`LogicalType`] for this field and returns itself. + pub fn with_logical_type(self, logical_type: Option) -> Self { + Self { + logical_type, + ..self + } + } + + /// Sets a list of fields that should be child nodes of this field. + /// Returns updated self. + pub fn with_fields(self, fields: Vec) -> Self { + Self { fields, ..self } + } + + /// Sets optional field id and returns itself. + pub fn with_id(self, id: Option) -> Self { + Self { id, ..self } + } + + /// Creates a new `GroupType` instance from the gathered attributes. + pub fn build(self) -> Result { + let mut basic_info = BasicTypeInfo { + name: String::from(self.name), + repetition: self.repetition, + converted_type: self.converted_type, + logical_type: self.logical_type.clone(), + id: self.id, + }; + // Populate the converted type if only the logical type is populated + if self.logical_type.is_some() && self.converted_type == ConvertedType::NONE { + basic_info.converted_type = self.logical_type.into(); + } + Ok(Type::GroupType { + basic_info, + fields: self.fields, + }) + } +} + +/// Basic type info. This contains information such as the name of the type, +/// the repetition level, the logical type and the kind of the type (group, primitive). +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct BasicTypeInfo { + name: String, + repetition: Option, + converted_type: ConvertedType, + logical_type: Option, + id: Option, +} + +impl HeapSize for BasicTypeInfo { + fn heap_size(&self) -> usize { + // no heap allocations in any other subfield + self.name.heap_size() + } +} + +impl BasicTypeInfo { + /// Returns field name. + pub fn name(&self) -> &str { + &self.name + } + + /// Returns `true` if type has repetition field set, `false` otherwise. + /// This is mostly applied to group type, because primitive type always has + /// repetition set. + pub fn has_repetition(&self) -> bool { + self.repetition.is_some() + } + + /// Returns [`Repetition`] value for the type. + pub fn repetition(&self) -> Repetition { + assert!(self.repetition.is_some()); + self.repetition.unwrap() + } + + /// Returns [`ConvertedType`] value for the type. + pub fn converted_type(&self) -> ConvertedType { + self.converted_type + } + + /// Returns [`LogicalType`] value for the type. + /// + /// Note that this function will clone the `LogicalType`. If performance is a concern, + /// use [`Self::logical_type_ref`] instead. + #[deprecated( + since = "57.1.0", + note = "use `BasicTypeInfo::logical_type_ref` instead (LogicalType cloning is non trivial)" + )] + pub fn logical_type(&self) -> Option { + // Unlike ConvertedType, LogicalType cannot implement Copy, thus we clone it + self.logical_type.clone() + } + + /// Return a reference to the [`LogicalType`] value for the type. + pub fn logical_type_ref(&self) -> Option<&LogicalType> { + self.logical_type.as_ref() + } + + /// Returns `true` if id is set, `false` otherwise. + pub fn has_id(&self) -> bool { + self.id.is_some() + } + + /// Returns id value for the type. + pub fn id(&self) -> i32 { + assert!(self.id.is_some()); + self.id.unwrap() + } +} + +// ---------------------------------------------------------------------- +// Parquet descriptor definitions + +/// Represents the location of a column in a Parquet schema +/// +/// # Example: refer to column named `'my_column'` +/// ``` +/// # use parquet::schema::types::ColumnPath; +/// let column_path = ColumnPath::from("my_column"); +/// ``` +/// +/// # Example: refer to column named `c` in a nested struct `{a: {b: {c: ...}}}` +/// ``` +/// # use parquet::schema::types::ColumnPath; +/// // form path 'a.b.c' +/// let column_path = ColumnPath::from(vec![ +/// String::from("a"), +/// String::from("b"), +/// String::from("c") +/// ]); +/// ``` +#[derive(Clone, PartialEq, Debug, Eq, Hash)] +pub struct ColumnPath { + parts: Vec, +} + +impl HeapSize for ColumnPath { + fn heap_size(&self) -> usize { + self.parts.heap_size() + } +} + +impl ColumnPath { + /// Creates new column path from vector of field names. + pub fn new(parts: Vec) -> Self { + ColumnPath { parts } + } + + /// Returns string representation of this column path. + /// ```rust + /// use parquet::schema::types::ColumnPath; + /// + /// let path = ColumnPath::new(vec!["a".to_string(), "b".to_string(), "c".to_string()]); + /// assert_eq!(&path.string(), "a.b.c"); + /// ``` + pub fn string(&self) -> String { + self.parts.join(".") + } + + /// Appends more components to end of column path. + /// ```rust + /// use parquet::schema::types::ColumnPath; + /// + /// let mut path = ColumnPath::new(vec!["a".to_string(), "b".to_string(), "c" + /// .to_string()]); + /// assert_eq!(&path.string(), "a.b.c"); + /// + /// path.append(vec!["d".to_string(), "e".to_string()]); + /// assert_eq!(&path.string(), "a.b.c.d.e"); + /// ``` + pub fn append(&mut self, mut tail: Vec) { + self.parts.append(&mut tail); + } + + /// Returns a slice of path components. + pub fn parts(&self) -> &[String] { + &self.parts + } +} + +impl fmt::Display for ColumnPath { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{:?}", self.string()) + } +} + +impl From> for ColumnPath { + fn from(parts: Vec) -> Self { + ColumnPath { parts } + } +} + +impl From<&str> for ColumnPath { + fn from(single_path: &str) -> Self { + let s = String::from(single_path); + ColumnPath::from(s) + } +} + +impl From for ColumnPath { + fn from(single_path: String) -> Self { + let v = vec![single_path]; + ColumnPath { parts: v } + } +} + +impl AsRef<[String]> for ColumnPath { + fn as_ref(&self) -> &[String] { + &self.parts + } +} + +/// Physical type for leaf-level primitive columns. +/// +/// Also includes the maximum definition and repetition levels required to +/// re-assemble nested data. +#[derive(Debug, PartialEq)] +pub struct ColumnDescriptor { + /// The "leaf" primitive type of this column + primitive_type: TypePtr, + + /// The maximum definition level for this column + max_def_level: i16, + + /// The maximum repetition level for this column + max_rep_level: i16, + + /// The definition level at the nearest REPEATED ancestor, or 0 if none. + repeated_ancestor_def_level: i16, + + /// The path of this column. For instance, "a.b.c.d". + path: ColumnPath, +} + +impl HeapSize for ColumnDescriptor { + fn heap_size(&self) -> usize { + // Don't include the heap size of primitive_type, this is already + // accounted for via SchemaDescriptor::schema + self.path.heap_size() + } +} + +impl ColumnDescriptor { + /// Creates new descriptor for leaf-level column. + pub fn new( + primitive_type: TypePtr, + max_def_level: i16, + max_rep_level: i16, + path: ColumnPath, + ) -> Self { + Self::new_with_repeated_ancestor(primitive_type, max_def_level, max_rep_level, path, 0) + } + + pub(crate) fn new_with_repeated_ancestor( + primitive_type: TypePtr, + max_def_level: i16, + max_rep_level: i16, + path: ColumnPath, + repeated_ancestor_def_level: i16, + ) -> Self { + Self { + primitive_type, + max_def_level, + max_rep_level, + repeated_ancestor_def_level, + path, + } + } + + /// Returns maximum definition level for this column. + #[inline] + pub fn max_def_level(&self) -> i16 { + self.max_def_level + } + + /// Returns maximum repetition level for this column. + #[inline] + pub fn max_rep_level(&self) -> i16 { + self.max_rep_level + } + + /// Returns the definition level at the nearest REPEATED ancestor, or 0 if none. + #[inline] + pub fn repeated_ancestor_def_level(&self) -> i16 { + self.repeated_ancestor_def_level + } + + /// Returns [`ColumnPath`] for this column. + pub fn path(&self) -> &ColumnPath { + &self.path + } + + /// Returns self type [`Type`] for this leaf column. + pub fn self_type(&self) -> &Type { + self.primitive_type.as_ref() + } + + /// Returns self type [`TypePtr`] for this leaf + /// column. + pub fn self_type_ptr(&self) -> TypePtr { + self.primitive_type.clone() + } + + /// Returns column name. + pub fn name(&self) -> &str { + self.primitive_type.name() + } + + /// Returns [`ConvertedType`] for this column. + pub fn converted_type(&self) -> ConvertedType { + self.primitive_type.get_basic_info().converted_type() + } + + /// Returns [`LogicalType`] for this column. + /// + /// Note that this function will clone the `LogicalType`. If performance is a concern, + /// use [`Self::logical_type_ref`] instead. + #[deprecated( + since = "57.1.0", + note = "use `ColumnDescriptor::logical_type_ref` instead (LogicalType cloning is non trivial)" + )] + pub fn logical_type(&self) -> Option { + self.primitive_type + .get_basic_info() + .logical_type_ref() + .cloned() + } + + /// Returns a reference to the [`LogicalType`] for this column. + pub fn logical_type_ref(&self) -> Option<&LogicalType> { + self.primitive_type.get_basic_info().logical_type_ref() + } + + /// Returns physical type for this column. + /// Note that it will panic if called on a non-primitive type. + pub fn physical_type(&self) -> PhysicalType { + match self.primitive_type.as_ref() { + Type::PrimitiveType { physical_type, .. } => *physical_type, + _ => panic!("Expected primitive type!"), + } + } + + /// Returns type length for this column. + /// Note that it will panic if called on a non-primitive type. + pub fn type_length(&self) -> i32 { + match self.primitive_type.as_ref() { + Type::PrimitiveType { type_length, .. } => *type_length, + _ => panic!("Expected primitive type!"), + } + } + + /// Returns type precision for this column. + /// Note that it will panic if called on a non-primitive type. + pub fn type_precision(&self) -> i32 { + match self.primitive_type.as_ref() { + Type::PrimitiveType { precision, .. } => *precision, + _ => panic!("Expected primitive type!"), + } + } + + /// Returns type scale for this column. + /// Note that it will panic if called on a non-primitive type. + pub fn type_scale(&self) -> i32 { + match self.primitive_type.as_ref() { + Type::PrimitiveType { scale, .. } => *scale, + _ => panic!("Expected primitive type!"), + } + } + + /// Returns the sort order for this column + pub fn sort_order(&self) -> SortOrder { + ColumnOrder::sort_order_for_type( + self.logical_type_ref(), + self.converted_type(), + self.physical_type(), + ) + } +} + +/// Schema of a Parquet file. +/// +/// Encapsulates the file's schema ([`Type`]) and [`ColumnDescriptor`]s for +/// each primitive (leaf) column. +/// +/// # Example +/// ``` +/// # use std::sync::Arc; +/// use parquet::schema::types::{SchemaDescriptor, Type}; +/// use parquet::basic; // note there are two `Type`s that are different +/// // Schema for a table with two columns: "a" (int64) and "b" (int32, stored as a date) +/// let descriptor = SchemaDescriptor::new( +/// Arc::new( +/// Type::group_type_builder("my_schema") +/// .with_fields(vec![ +/// Arc::new( +/// Type::primitive_type_builder("a", basic::Type::INT64) +/// .build().unwrap() +/// ), +/// Arc::new( +/// Type::primitive_type_builder("b", basic::Type::INT32) +/// .with_converted_type(basic::ConvertedType::DATE) +/// .with_logical_type(Some(basic::LogicalType::Date)) +/// .build().unwrap() +/// ), +/// ]) +/// .build().unwrap() +/// ) +/// ); +/// ``` +#[derive(PartialEq, Clone)] +pub struct SchemaDescriptor { + /// The top-level logical schema (the "message" type). + /// + /// This must be a [`Type::GroupType`] where each field is a root + /// column type in the schema. + schema: TypePtr, + + /// The descriptors for the physical type of each leaf column in this schema + /// + /// Constructed from `schema` in DFS order. + leaves: Vec, + + /// Mapping from a leaf column's index to the root column index that it + /// comes from. + /// + /// For instance: the leaf `a.b.c.d` would have a link back to `a`: + /// ```text + /// -- a <-----+ + /// -- -- b | + /// -- -- -- c | + /// -- -- -- -- d + /// ``` + leaf_to_base: Vec, +} + +impl fmt::Debug for SchemaDescriptor { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + // Skip leaves and leaf_to_base as they only a cache information already found in `schema` + f.debug_struct("SchemaDescriptor") + .field("schema", &self.schema) + .finish() + } +} + +// Need to implement HeapSize in this module as the fields are private +impl HeapSize for SchemaDescriptor { + fn heap_size(&self) -> usize { + self.schema.heap_size() + self.leaves.heap_size() + self.leaf_to_base.heap_size() + } +} + +impl SchemaDescriptor { + /// Creates new schema descriptor from Parquet schema. + pub fn new(tp: TypePtr) -> Self { + const INIT_SCHEMA_DEPTH: usize = 16; + assert!(tp.is_group(), "SchemaDescriptor should take a GroupType"); + // unwrap should be safe since we just asserted tp is a group + let n_leaves = num_leaves(&tp).unwrap(); + let mut leaves = Vec::with_capacity(n_leaves); + let mut leaf_to_base = Vec::with_capacity(n_leaves); + let mut path = Vec::with_capacity(INIT_SCHEMA_DEPTH); + for (root_idx, f) in tp.get_fields().iter().enumerate() { + path.clear(); + build_tree( + f, + root_idx, + 0, + 0, + 0, + &mut leaves, + &mut leaf_to_base, + &mut path, + ); + } + + Self { + schema: tp, + leaves, + leaf_to_base, + } + } + + /// Returns [`ColumnDescriptor`] for a field position. + pub fn column(&self, i: usize) -> ColumnDescPtr { + assert!( + i < self.leaves.len(), + "Index out of bound: {} not in [0, {})", + i, + self.leaves.len() + ); + self.leaves[i].clone() + } + + /// Returns slice of [`ColumnDescriptor`]. + pub fn columns(&self) -> &[ColumnDescPtr] { + &self.leaves + } + + /// Returns number of leaf-level columns. + pub fn num_columns(&self) -> usize { + self.leaves.len() + } + + /// Returns column root [`Type`] for a leaf position. + pub fn get_column_root(&self, i: usize) -> &Type { + let result = self.column_root_of(i); + result.as_ref() + } + + /// Returns column root [`Type`] pointer for a leaf position. + pub fn get_column_root_ptr(&self, i: usize) -> TypePtr { + let result = self.column_root_of(i); + result.clone() + } + + /// Returns the index of the root column for a field position + pub fn get_column_root_idx(&self, leaf: usize) -> usize { + assert!( + leaf < self.leaves.len(), + "Index out of bound: {} not in [0, {})", + leaf, + self.leaves.len() + ); + + *self + .leaf_to_base + .get(leaf) + .unwrap_or_else(|| panic!("Expected a value for index {leaf} but found None")) + } + + fn column_root_of(&self, i: usize) -> &TypePtr { + &self.schema.get_fields()[self.get_column_root_idx(i)] + } + + /// Returns schema as [`Type`]. + pub fn root_schema(&self) -> &Type { + self.schema.as_ref() + } + + /// Returns schema as [`TypePtr`] for cheap cloning. + pub fn root_schema_ptr(&self) -> TypePtr { + self.schema.clone() + } + + /// Returns schema name. + pub fn name(&self) -> &str { + self.schema.name() + } +} + +// walk tree and count nodes +pub(crate) fn num_nodes(tp: &TypePtr) -> Result { + if !tp.is_group() { + return Err(general_err!("Root schema must be Group type")); + } + let mut n_nodes = 1usize; // count root + for f in tp.get_fields().iter() { + count_nodes(f, &mut n_nodes); + } + Ok(n_nodes) +} + +pub(crate) fn count_nodes(tp: &TypePtr, n_nodes: &mut usize) { + *n_nodes += 1; + if let Type::GroupType { fields, .. } = tp.as_ref() { + for f in fields { + count_nodes(f, n_nodes); + } + } +} + +// do a quick walk of the tree to get proper sizing for SchemaDescriptor arrays +fn num_leaves(tp: &TypePtr) -> Result { + if !tp.is_group() { + return Err(general_err!("Root schema must be Group type")); + } + let mut n_leaves = 0usize; + for f in tp.get_fields().iter() { + count_leaves(f, &mut n_leaves); + } + Ok(n_leaves) +} + +fn count_leaves(tp: &TypePtr, n_leaves: &mut usize) { + match tp.as_ref() { + Type::PrimitiveType { .. } => *n_leaves += 1, + Type::GroupType { fields, .. } => { + for f in fields { + count_leaves(f, n_leaves); + } + } + } +} + +#[allow(clippy::too_many_arguments)] +fn build_tree<'a>( + tp: &'a TypePtr, + root_idx: usize, + mut max_rep_level: i16, + mut max_def_level: i16, + mut repeated_ancestor_def_level: i16, + leaves: &mut Vec, + leaf_to_base: &mut Vec, + path_so_far: &mut Vec<&'a str>, +) { + assert!(tp.get_basic_info().has_repetition()); + + path_so_far.push(tp.name()); + match tp.get_basic_info().repetition() { + Repetition::OPTIONAL => { + max_def_level += 1; + } + Repetition::REPEATED => { + max_def_level += 1; + max_rep_level += 1; + repeated_ancestor_def_level = max_def_level; + } + _ => {} + } + + match tp.as_ref() { + Type::PrimitiveType { .. } => { + let mut path: Vec = vec![]; + path.extend(path_so_far.iter().copied().map(String::from)); + let desc = ColumnDescriptor::new_with_repeated_ancestor( + tp.clone(), + max_def_level, + max_rep_level, + ColumnPath::new(path), + repeated_ancestor_def_level, + ); + leaves.push(Arc::new(desc)); + leaf_to_base.push(root_idx); + } + Type::GroupType { fields, .. } => { + for f in fields { + build_tree( + f, + root_idx, + max_rep_level, + max_def_level, + repeated_ancestor_def_level, + leaves, + leaf_to_base, + path_so_far, + ); + path_so_far.pop(); + } + } + } +} + +/// Checks if the logical type is valid. +fn check_logical_type(logical_type: &Option) -> Result<()> { + if let Some(LogicalType::Integer(IntType { bit_width, .. })) = logical_type { + if *bit_width != 8 && *bit_width != 16 && *bit_width != 32 && *bit_width != 64 { + return Err(general_err!( + "Bit width must be 8, 16, 32, or 64 for Integer logical type" + )); + } + } + Ok(()) +} + +// convert thrift decoded array of `SchemaElement` into this crate's representation of +// parquet types. this function consumes `elements`. +pub(crate) fn parquet_schema_from_array<'a>(elements: Vec>) -> Result { + let mut index = 0; + let num_elements = elements.len(); + let mut schema_nodes = Vec::with_capacity(1); // there should only be one element when done + + // turn into iterator so we can take ownership of elements of the vector + let mut elements = elements.into_iter(); + + while index < num_elements { + let t = schema_from_array_helper(&mut elements, num_elements, index)?; + index = t.0; + schema_nodes.push(t.1); + } + if schema_nodes.len() != 1 { + return Err(general_err!( + "Expected exactly one root node, but found {}", + schema_nodes.len() + )); + } + + if !schema_nodes[0].is_group() { + return Err(general_err!("Expected root node to be a group type")); + } + + Ok(schema_nodes.remove(0)) +} + +// recursive helper function for schema conversion +fn schema_from_array_helper<'a>( + elements: &mut IntoIter>, + num_elements: usize, + index: usize, +) -> Result<(usize, TypePtr)> { + // Whether or not the current node is root (message type). + // There is only one message type node in the schema tree. + let is_root_node = index == 0; + + if index >= num_elements { + return Err(general_err!( + "Index out of bound, index = {}, len = {}", + index, + num_elements + )); + } + let element = elements.next().expect("schema vector should not be empty"); + + // Check for empty schema + if let (true, None | Some(0)) = (is_root_node, element.num_children) { + let builder = Type::group_type_builder(element.name); + return Ok((index + 1, Arc::new(builder.build().unwrap()))); + } + + let converted_type = element.converted_type.unwrap_or(ConvertedType::NONE); + + // LogicalType is prefered to ConvertedType, but both may be present. + let logical_type = element.logical_type; + + check_logical_type(&logical_type)?; + + let field_id = element.field_id; + match element.num_children { + // From parquet-format: + // The children count is used to construct the nested relationship. + // This field is not set when the element is a primitive type + // Sometimes parquet-cpp sets num_children field to 0 for primitive types, so we + // have to handle this case too. + None | Some(0) => { + // primitive type + if element.repetition_type.is_none() { + return Err(general_err!( + "Repetition level must be defined for a primitive type" + )); + } + let repetition = element.repetition_type.unwrap(); + if let Some(physical_type) = element.r#type { + let length = element.type_length.unwrap_or(-1); + let scale = element.scale.unwrap_or(-1); + let precision = element.precision.unwrap_or(-1); + let name = element.name; + let builder = Type::primitive_type_builder(name, physical_type) + .with_repetition(repetition) + .with_converted_type(converted_type) + .with_logical_type(logical_type) + .with_length(length) + .with_precision(precision) + .with_scale(scale) + .with_id(field_id); + Ok((index + 1, Arc::new(builder.build()?))) + } else { + let mut builder = Type::group_type_builder(element.name) + .with_converted_type(converted_type) + .with_logical_type(logical_type) + .with_id(field_id); + if !is_root_node { + // Sometimes parquet-cpp and parquet-mr set repetition level REQUIRED or + // REPEATED for root node. + // + // We only set repetition for group types that are not top-level message + // type. According to parquet-format: + // Root of the schema does not have a repetition_type. + // All other types must have one. + builder = builder.with_repetition(repetition); + } + Ok((index + 1, Arc::new(builder.build().unwrap()))) + } + } + Some(n) => { + let repetition = element.repetition_type; + + let mut fields = Vec::with_capacity(usize::try_from(n)?); + let mut next_index = index + 1; + for _ in 0..n { + let child_result = schema_from_array_helper(elements, num_elements, next_index)?; + next_index = child_result.0; + fields.push(child_result.1); + } + + let mut builder = Type::group_type_builder(element.name) + .with_converted_type(converted_type) + .with_logical_type(logical_type) + .with_fields(fields) + .with_id(field_id); + + // Sometimes parquet-cpp and parquet-mr set repetition level REQUIRED or + // REPEATED for root node. + // + // We only set repetition for group types that are not top-level message + // type. According to parquet-format: + // Root of the schema does not have a repetition_type. + // All other types must have one. + if !is_root_node { + let Some(rep) = repetition else { + return Err(general_err!( + "Repetition level must be defined for non-root types" + )); + }; + builder = builder.with_repetition(rep); + } + Ok((next_index, Arc::new(builder.build()?))) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use crate::{ + file::metadata::thrift::tests::{buf_to_schema_list, roundtrip_schema, schema_to_buf}, + schema::parser::parse_message_type, + }; + + // TODO: add tests for v2 types + + #[test] + fn test_primitive_type() { + let mut result = Type::primitive_type_builder("foo", PhysicalType::INT32) + .with_logical_type(Some(LogicalType::integer(32, true))) + .with_id(Some(0)) + .build(); + assert!(result.is_ok()); + + if let Ok(tp) = result { + assert!(tp.is_primitive()); + assert!(!tp.is_group()); + let basic_info = tp.get_basic_info(); + assert_eq!(basic_info.repetition(), Repetition::OPTIONAL); + assert_eq!( + basic_info.logical_type_ref(), + Some(&LogicalType::integer(32, true)) + ); + assert_eq!(basic_info.converted_type(), ConvertedType::INT_32); + assert_eq!(basic_info.id(), 0); + match tp { + Type::PrimitiveType { physical_type, .. } => { + assert_eq!(physical_type, PhysicalType::INT32); + } + _ => panic!(), + } + } + + // Test illegal inputs with logical type + result = Type::primitive_type_builder("foo", PhysicalType::INT64) + .with_repetition(Repetition::REPEATED) + .with_logical_type(Some(LogicalType::integer(8, true))) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: Cannot annotate Integer(IntType { bit_width: 8, is_signed: true }) from INT64 for field 'foo'" + ); + } + + // Test illegal inputs with converted type + result = Type::primitive_type_builder("foo", PhysicalType::INT64) + .with_repetition(Repetition::REPEATED) + .with_converted_type(ConvertedType::BSON) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: BSON cannot annotate field 'foo' because it is not a BYTE_ARRAY field" + ); + } + + result = Type::primitive_type_builder("foo", PhysicalType::INT96) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::DECIMAL) + .with_precision(-1) + .with_scale(-1) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: DECIMAL can only annotate INT32, INT64, BYTE_ARRAY and FIXED_LEN_BYTE_ARRAY" + ); + } + + result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_logical_type(Some(LogicalType::decimal(32, 12))) + .with_precision(-1) + .with_scale(-1) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: DECIMAL logical type scale 32 must match self.scale -1 for field 'foo'" + ); + } + + result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::DECIMAL) + .with_precision(-1) + .with_scale(-1) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: Invalid DECIMAL precision: -1" + ); + } + + result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::DECIMAL) + .with_precision(0) + .with_scale(-1) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: Invalid DECIMAL precision: 0" + ); + } + + result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::DECIMAL) + .with_precision(1) + .with_scale(-1) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!(format!("{e}"), "Parquet error: Invalid DECIMAL scale: -1"); + } + + result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::DECIMAL) + .with_precision(1) + .with_scale(2) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: Invalid DECIMAL: scale (2) cannot be greater than precision (1)" + ); + } + + // It is OK if precision == scale + result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::DECIMAL) + .with_precision(1) + .with_scale(1) + .build(); + assert!(result.is_ok()); + + result = Type::primitive_type_builder("foo", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::DECIMAL) + .with_precision(18) + .with_scale(2) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: Cannot represent INT32 as DECIMAL with precision 18" + ); + } + + result = Type::primitive_type_builder("foo", PhysicalType::INT64) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::DECIMAL) + .with_precision(32) + .with_scale(2) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: Cannot represent INT64 as DECIMAL with precision 32" + ); + } + + result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::DECIMAL) + .with_length(5) + .with_precision(12) + .with_scale(2) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: Cannot represent FIXED_LEN_BYTE_ARRAY as DECIMAL with length 5 and precision 12. The max precision can only be 11" + ); + } + + result = Type::primitive_type_builder("foo", PhysicalType::INT64) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::UINT_8) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: UINT_8 cannot annotate field 'foo' because it is not a INT32 field" + ); + } + + result = Type::primitive_type_builder("foo", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::TIME_MICROS) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: TIME_MICROS cannot annotate field 'foo' because it is not a INT64 field" + ); + } + + result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::INTERVAL) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: INTERVAL cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(12) field" + ); + } + + result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::INTERVAL) + .with_length(1) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: INTERVAL cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(12) field" + ); + } + + result = Type::primitive_type_builder("foo", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::ENUM) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: ENUM cannot annotate field 'foo' because it is not a BYTE_ARRAY field" + ); + } + + result = Type::primitive_type_builder("foo", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::MAP) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: MAP cannot be applied to primitive field 'foo'" + ); + } + + result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::DECIMAL) + .with_length(-1) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: Invalid FIXED_LEN_BYTE_ARRAY length: -1 for field 'foo'" + ); + } + + result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_logical_type(Some(LogicalType::Float16)) + .with_length(2) + .build(); + assert!(result.is_ok()); + + // Can't be other than FIXED_LEN_BYTE_ARRAY for physical type + result = Type::primitive_type_builder("foo", PhysicalType::FLOAT) + .with_repetition(Repetition::REQUIRED) + .with_logical_type(Some(LogicalType::Float16)) + .with_length(2) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: Cannot annotate Float16 from FLOAT for field 'foo'" + ); + } + + // Must have length 2 + result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_logical_type(Some(LogicalType::Float16)) + .with_length(4) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: FLOAT16 cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(2) field" + ); + } + + // Must have length 16 + result = Type::primitive_type_builder("foo", PhysicalType::FIXED_LEN_BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .with_logical_type(Some(LogicalType::Uuid)) + .with_length(15) + .build(); + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!( + format!("{e}"), + "Parquet error: UUID cannot annotate field 'foo' because it is not a FIXED_LEN_BYTE_ARRAY(16) field" + ); + } + + // test unknown logical types are ok + result = Type::primitive_type_builder("foo", PhysicalType::BYTE_ARRAY) + .with_logical_type(Some(LogicalType::_Unknown { field_id: 100 })) + .build(); + assert!(result.is_ok()); + } + + #[test] + fn test_group_type() { + let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32) + .with_converted_type(ConvertedType::INT_32) + .with_id(Some(0)) + .build(); + assert!(f1.is_ok()); + let f2 = Type::primitive_type_builder("f2", PhysicalType::BYTE_ARRAY) + .with_converted_type(ConvertedType::UTF8) + .with_id(Some(1)) + .build(); + assert!(f2.is_ok()); + + let fields = vec![Arc::new(f1.unwrap()), Arc::new(f2.unwrap())]; + + let result = Type::group_type_builder("foo") + .with_repetition(Repetition::REPEATED) + .with_logical_type(Some(LogicalType::List)) + .with_fields(fields) + .with_id(Some(1)) + .build(); + assert!(result.is_ok()); + + let tp = result.unwrap(); + let basic_info = tp.get_basic_info(); + assert!(tp.is_group()); + assert!(!tp.is_primitive()); + assert_eq!(basic_info.repetition(), Repetition::REPEATED); + assert_eq!(basic_info.logical_type_ref(), Some(&LogicalType::List)); + assert_eq!(basic_info.converted_type(), ConvertedType::LIST); + assert_eq!(basic_info.id(), 1); + assert_eq!(tp.get_fields().len(), 2); + assert_eq!(tp.get_fields()[0].name(), "f1"); + assert_eq!(tp.get_fields()[1].name(), "f2"); + } + + #[test] + fn test_column_descriptor() { + let result = test_column_descriptor_helper(); + assert!( + result.is_ok(), + "Expected result to be OK but got err:\n {}", + result.unwrap_err() + ); + } + + fn test_column_descriptor_helper() -> Result<()> { + let tp = Type::primitive_type_builder("name", PhysicalType::BYTE_ARRAY) + .with_converted_type(ConvertedType::UTF8) + .build()?; + + let descr = ColumnDescriptor::new(Arc::new(tp), 4, 1, ColumnPath::from("name")); + + assert_eq!(descr.path(), &ColumnPath::from("name")); + assert_eq!(descr.converted_type(), ConvertedType::UTF8); + assert_eq!(descr.physical_type(), PhysicalType::BYTE_ARRAY); + assert_eq!(descr.max_def_level(), 4); + assert_eq!(descr.max_rep_level(), 1); + assert_eq!(descr.name(), "name"); + assert_eq!(descr.type_length(), -1); + assert_eq!(descr.type_precision(), -1); + assert_eq!(descr.type_scale(), -1); + + Ok(()) + } + + #[test] + fn test_schema_descriptor() { + let result = test_schema_descriptor_helper(); + assert!( + result.is_ok(), + "Expected result to be OK but got err:\n {}", + result.unwrap_err() + ); + } + + // A helper fn to avoid handling the results from type creation + fn test_schema_descriptor_helper() -> Result<()> { + let mut fields = vec![]; + + let inta = Type::primitive_type_builder("a", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::INT_32) + .build()?; + fields.push(Arc::new(inta)); + let intb = Type::primitive_type_builder("b", PhysicalType::INT64) + .with_converted_type(ConvertedType::INT_64) + .build()?; + fields.push(Arc::new(intb)); + let intc = Type::primitive_type_builder("c", PhysicalType::BYTE_ARRAY) + .with_repetition(Repetition::REPEATED) + .with_converted_type(ConvertedType::UTF8) + .build()?; + fields.push(Arc::new(intc)); + + // 3-level list encoding + let item1 = Type::primitive_type_builder("item1", PhysicalType::INT64) + .with_repetition(Repetition::REQUIRED) + .with_converted_type(ConvertedType::INT_64) + .build()?; + let item2 = Type::primitive_type_builder("item2", PhysicalType::BOOLEAN).build()?; + let item3 = Type::primitive_type_builder("item3", PhysicalType::INT32) + .with_repetition(Repetition::REPEATED) + .with_converted_type(ConvertedType::INT_32) + .build()?; + let list = Type::group_type_builder("records") + .with_repetition(Repetition::REPEATED) + .with_converted_type(ConvertedType::LIST) + .with_fields(vec![Arc::new(item1), Arc::new(item2), Arc::new(item3)]) + .build()?; + let bag = Type::group_type_builder("bag") + .with_repetition(Repetition::OPTIONAL) + .with_fields(vec![Arc::new(list)]) + .build()?; + fields.push(Arc::new(bag)); + + let schema = Type::group_type_builder("schema") + .with_repetition(Repetition::REPEATED) + .with_fields(fields) + .build()?; + let descr = SchemaDescriptor::new(Arc::new(schema)); + + let nleaves = 6; + assert_eq!(descr.num_columns(), nleaves); + + // mdef mrep + // required int32 a 0 0 + // optional int64 b 1 0 + // repeated byte_array c 1 1 + // optional group bag 1 0 + // repeated group records 2 1 + // required int64 item1 2 1 + // optional boolean item2 3 1 + // repeated int32 item3 3 2 + let ex_max_def_levels = [0, 1, 1, 2, 3, 3]; + let ex_max_rep_levels = [0, 0, 1, 1, 1, 2]; + + for i in 0..nleaves { + let col = descr.column(i); + assert_eq!(col.max_def_level(), ex_max_def_levels[i], "{i}"); + assert_eq!(col.max_rep_level(), ex_max_rep_levels[i], "{i}"); + } + + assert_eq!(descr.column(0).path().string(), "a"); + assert_eq!(descr.column(1).path().string(), "b"); + assert_eq!(descr.column(2).path().string(), "c"); + assert_eq!(descr.column(3).path().string(), "bag.records.item1"); + assert_eq!(descr.column(4).path().string(), "bag.records.item2"); + assert_eq!(descr.column(5).path().string(), "bag.records.item3"); + + assert_eq!(descr.get_column_root(0).name(), "a"); + assert_eq!(descr.get_column_root(3).name(), "bag"); + assert_eq!(descr.get_column_root(4).name(), "bag"); + assert_eq!(descr.get_column_root(5).name(), "bag"); + + Ok(()) + } + + #[test] + fn test_schema_build_tree_def_rep_levels() { + let message_type = " + message spark_schema { + REQUIRED INT32 a; + OPTIONAL group b { + OPTIONAL INT32 _1; + OPTIONAL INT32 _2; + } + OPTIONAL group c (LIST) { + REPEATED group list { + OPTIONAL INT32 element; + } + } + } + "; + let schema = parse_message_type(message_type).expect("should parse schema"); + let descr = SchemaDescriptor::new(Arc::new(schema)); + // required int32 a + assert_eq!(descr.column(0).max_def_level(), 0); + assert_eq!(descr.column(0).max_rep_level(), 0); + // optional int32 b._1 + assert_eq!(descr.column(1).max_def_level(), 2); + assert_eq!(descr.column(1).max_rep_level(), 0); + // optional int32 b._2 + assert_eq!(descr.column(2).max_def_level(), 2); + assert_eq!(descr.column(2).max_rep_level(), 0); + // repeated optional int32 c.list.element + assert_eq!(descr.column(3).max_def_level(), 3); + assert_eq!(descr.column(3).max_rep_level(), 1); + } + + #[test] + fn test_schema_build_tree_repeated_ancestor_def_level() { + // Flat columns: no REPEATED ancestor → repeated_ancestor_def_level = 0 + let message_type = " + message m { + REQUIRED INT32 a; + OPTIONAL INT32 b; + OPTIONAL group s { + OPTIONAL INT32 x; + } + } + "; + let schema = parse_message_type(message_type).expect("should parse schema"); + let descr = SchemaDescriptor::new(Arc::new(schema)); + assert_eq!(descr.column(0).repeated_ancestor_def_level(), 0); // a + assert_eq!(descr.column(1).repeated_ancestor_def_level(), 0); // b + assert_eq!(descr.column(2).repeated_ancestor_def_level(), 0); // s.x + + // Standard list: OPTIONAL outer, REPEATED group, OPTIONAL element + // repeated_ancestor_def_level is the def_level at the REPEATED group (= 2) + let message_type = " + message m { + OPTIONAL group c (LIST) { + REPEATED group list { + OPTIONAL INT32 element; + } + } + } + "; + let schema = parse_message_type(message_type).expect("should parse schema"); + let descr = SchemaDescriptor::new(Arc::new(schema)); + // c(optional)=1, list(repeated)=2, element(optional)=3 + assert_eq!(descr.column(0).max_def_level(), 3); + assert_eq!(descr.column(0).max_rep_level(), 1); + assert_eq!(descr.column(0).repeated_ancestor_def_level(), 2); + + // Required list: REQUIRED outer, REPEATED group, REQUIRED element + // No OPTIONAL nodes between REPEATED and leaf, so repeated_ancestor_def_level == max_def_level + let message_type = " + message m { + REQUIRED group c (LIST) { + REPEATED group list { + REQUIRED INT32 element; + } + } + } + "; + let schema = parse_message_type(message_type).expect("should parse schema"); + let descr = SchemaDescriptor::new(Arc::new(schema)); + // list(repeated)=1, element(required)=1 + assert_eq!(descr.column(0).max_def_level(), 1); + assert_eq!(descr.column(0).max_rep_level(), 1); + assert_eq!(descr.column(0).repeated_ancestor_def_level(), 1); + + // Nested lists: innermost REPEATED wins + let message_type = " + message m { + OPTIONAL group outer (LIST) { + REPEATED group list { + OPTIONAL group inner (LIST) { + REPEATED group list2 { + OPTIONAL INT32 element; + } + } + } + } + } + "; + let schema = parse_message_type(message_type).expect("should parse schema"); + let descr = SchemaDescriptor::new(Arc::new(schema)); + // outer(opt)=1, list(rep)=2, inner(opt)=3, list2(rep)=4, element(opt)=5 + assert_eq!(descr.column(0).max_def_level(), 5); + assert_eq!(descr.column(0).max_rep_level(), 2); + assert_eq!(descr.column(0).repeated_ancestor_def_level(), 4); + + // Struct inside list: all sibling leaves share the same repeated_ancestor_def_level + let message_type = " + message m { + OPTIONAL group bag (LIST) { + REPEATED group list { + REQUIRED group item { + OPTIONAL INT32 x; + REQUIRED INT32 y; + } + } + } + } + "; + let schema = parse_message_type(message_type).expect("should parse schema"); + let descr = SchemaDescriptor::new(Arc::new(schema)); + // bag(opt)=1, list(rep)=2, item(req)=2, x(opt)=3 + assert_eq!(descr.column(0).repeated_ancestor_def_level(), 2); // bag.list.item.x + // bag(opt)=1, list(rep)=2, item(req)=2, y(req)=2 + assert_eq!(descr.column(1).repeated_ancestor_def_level(), 2); // bag.list.item.y + + // Map type: key (required) and value (optional) under the same REPEATED group + let message_type = " + message m { + OPTIONAL group my_map (MAP) { + REPEATED group key_value { + REQUIRED BYTE_ARRAY key (UTF8); + OPTIONAL INT32 value; + } + } + } + "; + let schema = parse_message_type(message_type).expect("should parse schema"); + let descr = SchemaDescriptor::new(Arc::new(schema)); + // my_map(opt)=1, key_value(rep)=2, key(req)=2 + assert_eq!(descr.column(0).max_def_level(), 2); + assert_eq!(descr.column(0).repeated_ancestor_def_level(), 2); // key: max_def == repeated_ancestor + // my_map(opt)=1, key_value(rep)=2, value(opt)=3 + assert_eq!(descr.column(1).max_def_level(), 3); + assert_eq!(descr.column(1).repeated_ancestor_def_level(), 2); // value: max_def > repeated_ancestor + } + + #[test] + #[should_panic(expected = "Cannot call get_physical_type() on a non-primitive type")] + fn test_get_physical_type_panic() { + let list = Type::group_type_builder("records") + .with_repetition(Repetition::REPEATED) + .build() + .unwrap(); + list.get_physical_type(); + } + + #[test] + fn test_get_physical_type_primitive() { + let f = Type::primitive_type_builder("f", PhysicalType::INT64) + .build() + .unwrap(); + assert_eq!(f.get_physical_type(), PhysicalType::INT64); + + let f = Type::primitive_type_builder("f", PhysicalType::BYTE_ARRAY) + .build() + .unwrap(); + assert_eq!(f.get_physical_type(), PhysicalType::BYTE_ARRAY); + } + + #[test] + fn test_check_contains_primitive_primitive() { + // OK + let f1 = Type::primitive_type_builder("f", PhysicalType::INT32) + .build() + .unwrap(); + let f2 = Type::primitive_type_builder("f", PhysicalType::INT32) + .build() + .unwrap(); + assert!(f1.check_contains(&f2)); + + // OK: different logical type does not affect check_contains + let f1 = Type::primitive_type_builder("f", PhysicalType::INT32) + .with_converted_type(ConvertedType::UINT_8) + .build() + .unwrap(); + let f2 = Type::primitive_type_builder("f", PhysicalType::INT32) + .with_converted_type(ConvertedType::UINT_16) + .build() + .unwrap(); + assert!(f1.check_contains(&f2)); + + // KO: different name + let f1 = Type::primitive_type_builder("f1", PhysicalType::INT32) + .build() + .unwrap(); + let f2 = Type::primitive_type_builder("f2", PhysicalType::INT32) + .build() + .unwrap(); + assert!(!f1.check_contains(&f2)); + + // KO: different type + let f1 = Type::primitive_type_builder("f", PhysicalType::INT32) + .build() + .unwrap(); + let f2 = Type::primitive_type_builder("f", PhysicalType::INT64) + .build() + .unwrap(); + assert!(!f1.check_contains(&f2)); + + // KO: different repetition + let f1 = Type::primitive_type_builder("f", PhysicalType::INT32) + .with_repetition(Repetition::REQUIRED) + .build() + .unwrap(); + let f2 = Type::primitive_type_builder("f", PhysicalType::INT32) + .with_repetition(Repetition::OPTIONAL) + .build() + .unwrap(); + assert!(!f1.check_contains(&f2)); + } + + // function to create a new group type for testing + fn test_new_group_type(name: &str, repetition: Repetition, types: Vec) -> Type { + Type::group_type_builder(name) + .with_repetition(repetition) + .with_fields(types.into_iter().map(Arc::new).collect()) + .build() + .unwrap() + } + + #[test] + fn test_check_contains_group_group() { + // OK: should match okay with empty fields + let f1 = Type::group_type_builder("f").build().unwrap(); + let f2 = Type::group_type_builder("f").build().unwrap(); + assert!(f1.check_contains(&f2)); + assert!(!f1.is_optional()); + + // OK: fields match + let f1 = test_new_group_type( + "f", + Repetition::REPEATED, + vec![ + Type::primitive_type_builder("f1", PhysicalType::INT32) + .build() + .unwrap(), + Type::primitive_type_builder("f2", PhysicalType::INT64) + .build() + .unwrap(), + ], + ); + let f2 = test_new_group_type( + "f", + Repetition::REPEATED, + vec![ + Type::primitive_type_builder("f1", PhysicalType::INT32) + .build() + .unwrap(), + Type::primitive_type_builder("f2", PhysicalType::INT64) + .build() + .unwrap(), + ], + ); + assert!(f1.check_contains(&f2)); + + // OK: subset of fields + let f1 = test_new_group_type( + "f", + Repetition::REPEATED, + vec![ + Type::primitive_type_builder("f1", PhysicalType::INT32) + .build() + .unwrap(), + Type::primitive_type_builder("f2", PhysicalType::INT64) + .build() + .unwrap(), + ], + ); + let f2 = test_new_group_type( + "f", + Repetition::REPEATED, + vec![ + Type::primitive_type_builder("f2", PhysicalType::INT64) + .build() + .unwrap(), + ], + ); + assert!(f1.check_contains(&f2)); + + // KO: different name + let f1 = Type::group_type_builder("f1").build().unwrap(); + let f2 = Type::group_type_builder("f2").build().unwrap(); + assert!(!f1.check_contains(&f2)); + + // KO: different repetition + let f1 = Type::group_type_builder("f") + .with_repetition(Repetition::OPTIONAL) + .build() + .unwrap(); + let f2 = Type::group_type_builder("f") + .with_repetition(Repetition::REPEATED) + .build() + .unwrap(); + assert!(!f1.check_contains(&f2)); + + // KO: different fields + let f1 = test_new_group_type( + "f", + Repetition::REPEATED, + vec![ + Type::primitive_type_builder("f1", PhysicalType::INT32) + .build() + .unwrap(), + Type::primitive_type_builder("f2", PhysicalType::INT64) + .build() + .unwrap(), + ], + ); + let f2 = test_new_group_type( + "f", + Repetition::REPEATED, + vec![ + Type::primitive_type_builder("f1", PhysicalType::INT32) + .build() + .unwrap(), + Type::primitive_type_builder("f2", PhysicalType::BOOLEAN) + .build() + .unwrap(), + ], + ); + assert!(!f1.check_contains(&f2)); + + // KO: different fields + let f1 = test_new_group_type( + "f", + Repetition::REPEATED, + vec![ + Type::primitive_type_builder("f1", PhysicalType::INT32) + .build() + .unwrap(), + Type::primitive_type_builder("f2", PhysicalType::INT64) + .build() + .unwrap(), + ], + ); + let f2 = test_new_group_type( + "f", + Repetition::REPEATED, + vec![ + Type::primitive_type_builder("f3", PhysicalType::INT32) + .build() + .unwrap(), + ], + ); + assert!(!f1.check_contains(&f2)); + } + + #[test] + fn test_check_contains_group_primitive() { + // KO: should not match + let f1 = Type::group_type_builder("f").build().unwrap(); + let f2 = Type::primitive_type_builder("f", PhysicalType::INT64) + .build() + .unwrap(); + assert!(!f1.check_contains(&f2)); + assert!(!f2.check_contains(&f1)); + + // KO: should not match when primitive field is part of group type + let f1 = test_new_group_type( + "f", + Repetition::REPEATED, + vec![ + Type::primitive_type_builder("f1", PhysicalType::INT32) + .build() + .unwrap(), + ], + ); + let f2 = Type::primitive_type_builder("f1", PhysicalType::INT32) + .build() + .unwrap(); + assert!(!f1.check_contains(&f2)); + assert!(!f2.check_contains(&f1)); + + // OK: match nested types + let f1 = test_new_group_type( + "a", + Repetition::REPEATED, + vec![ + test_new_group_type( + "b", + Repetition::REPEATED, + vec![ + Type::primitive_type_builder("c", PhysicalType::INT32) + .build() + .unwrap(), + ], + ), + Type::primitive_type_builder("d", PhysicalType::INT64) + .build() + .unwrap(), + Type::primitive_type_builder("e", PhysicalType::BOOLEAN) + .build() + .unwrap(), + ], + ); + let f2 = test_new_group_type( + "a", + Repetition::REPEATED, + vec![test_new_group_type( + "b", + Repetition::REPEATED, + vec![ + Type::primitive_type_builder("c", PhysicalType::INT32) + .build() + .unwrap(), + ], + )], + ); + assert!(f1.check_contains(&f2)); // should match + assert!(!f2.check_contains(&f1)); // should fail + } + + #[test] + fn test_schema_type_thrift_conversion_err() { + let schema = Type::primitive_type_builder("col", PhysicalType::INT32) + .build() + .unwrap(); + let schema = Arc::new(schema); + let thrift_schema = schema_to_buf(&schema); + assert!(thrift_schema.is_err()); + if let Err(e) = thrift_schema { + assert_eq!( + format!("{e}"), + "Parquet error: Root schema must be Group type" + ); + } + } + + #[test] + fn test_schema_type_thrift_conversion() { + let message_type = " + message conversions { + REQUIRED INT64 id; + OPTIONAL FIXED_LEN_BYTE_ARRAY (2) f16 (FLOAT16); + OPTIONAL group int_array_Array (LIST) { + REPEATED group list { + OPTIONAL group element (LIST) { + REPEATED group list { + OPTIONAL INT32 element; + } + } + } + } + OPTIONAL group int_map (MAP) { + REPEATED group map (MAP_KEY_VALUE) { + REQUIRED BYTE_ARRAY key (UTF8); + OPTIONAL INT32 value; + } + } + OPTIONAL group int_Map_Array (LIST) { + REPEATED group list { + OPTIONAL group g (MAP) { + REPEATED group map (MAP_KEY_VALUE) { + REQUIRED BYTE_ARRAY key (UTF8); + OPTIONAL group value { + OPTIONAL group H { + OPTIONAL group i (LIST) { + REPEATED group list { + OPTIONAL DOUBLE element; + } + } + } + } + } + } + } + } + OPTIONAL group nested_struct { + OPTIONAL INT32 A; + OPTIONAL group b (LIST) { + REPEATED group list { + REQUIRED FIXED_LEN_BYTE_ARRAY (16) element; + } + } + } + } + "; + let expected_schema = parse_message_type(message_type).unwrap(); + let result_schema = roundtrip_schema(Arc::new(expected_schema.clone())).unwrap(); + assert_eq!(result_schema, Arc::new(expected_schema)); + } + + #[test] + fn test_schema_type_thrift_conversion_decimal() { + let message_type = " + message decimals { + OPTIONAL INT32 field0; + OPTIONAL INT64 field1 (DECIMAL (18, 2)); + OPTIONAL FIXED_LEN_BYTE_ARRAY (16) field2 (DECIMAL (38, 18)); + OPTIONAL BYTE_ARRAY field3 (DECIMAL (9)); + } + "; + let expected_schema = parse_message_type(message_type).unwrap(); + let result_schema = roundtrip_schema(Arc::new(expected_schema.clone())).unwrap(); + assert_eq!(result_schema, Arc::new(expected_schema)); + } + + // Tests schema conversion from thrift, when num_children is set to Some(0) for a + // primitive type. + #[test] + fn test_schema_from_thrift_with_num_children_set() { + // schema definition written by parquet-cpp version 1.3.2-SNAPSHOT + let message_type = " + message schema { + OPTIONAL BYTE_ARRAY id (UTF8); + OPTIONAL BYTE_ARRAY name (UTF8); + OPTIONAL BYTE_ARRAY message (UTF8); + OPTIONAL INT32 type (UINT_8); + OPTIONAL INT64 author_time (TIMESTAMP_MILLIS); + OPTIONAL INT64 __index_level_0__; + } + "; + + let expected_schema = Arc::new(parse_message_type(message_type).unwrap()); + let mut buf = schema_to_buf(&expected_schema).unwrap(); + let mut thrift_schema = buf_to_schema_list(&mut buf).unwrap(); + + // Change all of None to Some(0) + for elem in &mut thrift_schema[..] { + if elem.num_children.is_none() { + elem.num_children = Some(0); + } + } + + let result_schema = parquet_schema_from_array(thrift_schema).unwrap(); + assert_eq!(result_schema, expected_schema); + } + + // Sometimes parquet-cpp sets repetition level for the root node, which is against + // the format definition, but we need to handle it by setting it back to None. + #[test] + fn test_schema_from_thrift_root_has_repetition() { + // schema definition written by parquet-cpp version 1.3.2-SNAPSHOT + let message_type = " + message schema { + OPTIONAL BYTE_ARRAY a (UTF8); + OPTIONAL INT32 b (UINT_8); + } + "; + + let expected_schema = Arc::new(parse_message_type(message_type).unwrap()); + let mut buf = schema_to_buf(&expected_schema).unwrap(); + let mut thrift_schema = buf_to_schema_list(&mut buf).unwrap(); + thrift_schema[0].repetition_type = Some(Repetition::REQUIRED); + + let result_schema = parquet_schema_from_array(thrift_schema).unwrap(); + assert_eq!(result_schema, expected_schema); + } + + #[test] + fn test_schema_from_thrift_group_has_no_child() { + let message_type = "message schema {}"; + + let expected_schema = Arc::new(parse_message_type(message_type).unwrap()); + let mut buf = schema_to_buf(&expected_schema).unwrap(); + let mut thrift_schema = buf_to_schema_list(&mut buf).unwrap(); + thrift_schema[0].repetition_type = Some(Repetition::REQUIRED); + + let result_schema = parquet_schema_from_array(thrift_schema).unwrap(); + assert_eq!(result_schema, expected_schema); + } + + #[test] + fn test_parquet_schema_from_array_rejects_negative_num_children() { + let elements = vec![SchemaElement { + r#type: None, + type_length: None, + repetition_type: Some(Repetition::REQUIRED), + name: "schema", + num_children: Some(-1), + converted_type: None, + scale: None, + precision: None, + field_id: None, + logical_type: None, + }]; + let result = parquet_schema_from_array(elements); + assert!(result.unwrap_err().to_string().contains("Integer overflow")); + } +} diff --git a/vendor/parquet/src/schema/visitor.rs b/vendor/parquet/src/schema/visitor.rs new file mode 100644 index 0000000..7a10d3a --- /dev/null +++ b/vendor/parquet/src/schema/visitor.rs @@ -0,0 +1,241 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Utilities to traverse against various parquet type. + +use crate::basic::{ConvertedType, Repetition}; +use crate::errors::ParquetError::General; +use crate::errors::Result; +use crate::schema::types::{Type, TypePtr}; + +/// A utility trait to help user to traverse against parquet type. +pub trait TypeVisitor { + /// Called when a primitive type hit. + fn visit_primitive(&mut self, primitive_type: TypePtr, context: C) -> Result; + + /// Default implementation when visiting a list. + /// + /// It checks list type definition and calls [`Self::visit_list_with_item`] with extracted + /// item type. + /// + /// To fully understand this algorithm, please refer to + /// [parquet doc](https://github.com/apache/parquet-format/blob/master/LogicalTypes.md). + /// + /// For example, a standard list type looks like: + /// + /// ```text + /// required/optional group my_list (LIST) { + // repeated group list { + // required/optional binary element (UTF8); + // } + // } + /// ``` + /// + /// In such a case, [`Self::visit_list_with_item`] will be called with `my_list` as the list + /// type, and `element` as the `item_type` + /// + fn visit_list(&mut self, list_type: TypePtr, context: C) -> Result { + match list_type.as_ref() { + Type::PrimitiveType { .. } => { + panic!("{list_type:?} is a list type and must be a group type") + } + Type::GroupType { + basic_info: _, + fields, + } if fields.len() == 1 => { + let list_item = fields.first().unwrap(); + + match list_item.as_ref() { + Type::PrimitiveType { .. } => { + if list_item.get_basic_info().repetition() == Repetition::REPEATED { + self.visit_list_with_item(list_type.clone(), list_item.clone(), context) + } else { + Err(General( + "Primitive element type of list must be repeated.".to_string(), + )) + } + } + Type::GroupType { + basic_info: _, + fields, + } => { + if fields.len() == 1 + && list_item.name() != "array" + && list_item.name() != format!("{}_tuple", list_type.name()) + { + self.visit_list_with_item( + list_type.clone(), + fields.first().unwrap().clone(), + context, + ) + } else { + self.visit_list_with_item(list_type.clone(), list_item.clone(), context) + } + } + } + } + _ => Err(General( + "Group element type of list can only contain one field.".to_string(), + )), + } + } + + /// Called when a struct type hit. + fn visit_struct(&mut self, struct_type: TypePtr, context: C) -> Result; + + /// Called when a map type hit. + fn visit_map(&mut self, map_type: TypePtr, context: C) -> Result; + + /// A utility method which detects input type and calls corresponding method. + fn dispatch(&mut self, cur_type: TypePtr, context: C) -> Result { + if cur_type.is_primitive() { + self.visit_primitive(cur_type, context) + } else { + match cur_type.get_basic_info().converted_type() { + ConvertedType::LIST => self.visit_list(cur_type, context), + ConvertedType::MAP | ConvertedType::MAP_KEY_VALUE => { + self.visit_map(cur_type, context) + } + _ => self.visit_struct(cur_type, context), + } + } + } + + /// Called by `visit_list`. + fn visit_list_with_item( + &mut self, + list_type: TypePtr, + item_type: TypePtr, + context: C, + ) -> Result; +} + +#[cfg(test)] +mod tests { + use super::TypeVisitor; + use crate::basic::Type as PhysicalType; + use crate::errors::Result; + use crate::schema::parser::parse_message_type; + use crate::schema::types::TypePtr; + use std::sync::Arc; + + struct TestVisitorContext {} + struct TestVisitor { + primitive_visited: bool, + struct_visited: bool, + list_visited: bool, + root_type: TypePtr, + } + + impl TypeVisitor for TestVisitor { + fn visit_primitive( + &mut self, + primitive_type: TypePtr, + _context: TestVisitorContext, + ) -> Result { + assert_eq!( + self.get_field_by_name(primitive_type.name()).as_ref(), + primitive_type.as_ref() + ); + self.primitive_visited = true; + Ok(true) + } + + fn visit_struct( + &mut self, + struct_type: TypePtr, + _context: TestVisitorContext, + ) -> Result { + assert_eq!( + self.get_field_by_name(struct_type.name()).as_ref(), + struct_type.as_ref() + ); + self.struct_visited = true; + Ok(true) + } + + fn visit_map(&mut self, _map_type: TypePtr, _context: TestVisitorContext) -> Result { + unimplemented!() + } + + fn visit_list_with_item( + &mut self, + list_type: TypePtr, + item_type: TypePtr, + _context: TestVisitorContext, + ) -> Result { + assert_eq!( + self.get_field_by_name(list_type.name()).as_ref(), + list_type.as_ref() + ); + assert_eq!("element", item_type.name()); + assert_eq!(PhysicalType::INT32, item_type.get_physical_type()); + self.list_visited = true; + Ok(true) + } + } + + impl TestVisitor { + fn new(root: TypePtr) -> Self { + Self { + primitive_visited: false, + struct_visited: false, + list_visited: false, + root_type: root, + } + } + + fn get_field_by_name(&self, name: &str) -> TypePtr { + self.root_type + .get_fields() + .iter() + .find(|t| t.name() == name) + .cloned() + .unwrap() + } + } + + #[test] + fn test_visitor() { + let message_type = " + message spark_schema { + REQUIRED INT32 a; + OPTIONAL group inner_schema { + REQUIRED INT32 b; + REQUIRED DOUBLE c; + } + + OPTIONAL group e (LIST) { + REPEATED group list { + REQUIRED INT32 element; + } + } + "; + + let parquet_type = Arc::new(parse_message_type(message_type).unwrap()); + + let mut visitor = TestVisitor::new(parquet_type.clone()); + for f in parquet_type.get_fields() { + let c = TestVisitorContext {}; + assert!(visitor.dispatch(f.clone(), c).unwrap()); + } + + assert!(visitor.struct_visited); + assert!(visitor.primitive_visited); + assert!(visitor.list_visited); + } +} diff --git a/vendor/parquet/src/util/bit_pack.rs b/vendor/parquet/src/util/bit_pack.rs new file mode 100644 index 0000000..94ab957 --- /dev/null +++ b/vendor/parquet/src/util/bit_pack.rs @@ -0,0 +1,137 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Vectorised bit-packing utilities + +/// Macro that generates an unpack function taking the number of bits as a const generic +macro_rules! unpack_impl { + ($t:ty, $bytes:literal, $bits:tt) => { + pub fn unpack(input: &[u8], output: &mut [$t; $bits]) { + if NUM_BITS == 0 { + for out in output { + *out = 0; + } + return; + } + + assert!(NUM_BITS <= $bytes * 8); + + let mask = match NUM_BITS { + $bits => <$t>::MAX, + _ => ((1 << NUM_BITS) - 1), + }; + + assert!(input.len() >= NUM_BITS * $bytes); + + let r = |output_idx: usize| { + <$t>::from_le_bytes( + input[output_idx * $bytes..output_idx * $bytes + $bytes] + .try_into() + .unwrap(), + ) + }; + + seq_macro::seq!(i in 0..$bits { + let start_bit = i * NUM_BITS; + let end_bit = start_bit + NUM_BITS; + + let start_bit_offset = start_bit % $bits; + let end_bit_offset = end_bit % $bits; + let start_byte = start_bit / $bits; + let end_byte = end_bit / $bits; + if start_byte != end_byte && end_bit_offset != 0 { + let val = r(start_byte); + let a = val >> start_bit_offset; + let val = r(end_byte); + let b = val << (NUM_BITS - end_bit_offset); + + output[i] = a | (b & mask); + } else { + let val = r(start_byte); + output[i] = (val >> start_bit_offset) & mask; + } + }); + } + }; +} + +/// Macro that generates unpack functions that accept num_bits as a parameter +macro_rules! unpack { + ($name:ident, $t:ty, $bytes:literal, $bits:tt) => { + mod $name { + unpack_impl!($t, $bytes, $bits); + } + + /// Unpack packed `input` into `output` with a bit width of `num_bits` + pub fn $name(input: &[u8], output: &mut [$t; $bits], num_bits: usize) { + // This will get optimised into a jump table + seq_macro::seq!(i in 0..=$bits { + if i == num_bits { + return $name::unpack::(input, output); + } + }); + unreachable!("invalid num_bits {}", num_bits); + } + }; +} + +unpack!(unpack8, u8, 1, 8); +unpack!(unpack16, u16, 2, 16); +unpack!(unpack32, u32, 4, 32); +unpack!(unpack64, u64, 8, 64); + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_basic() { + let input = [0xFF; 4096]; + + for i in 0..=8 { + let mut output = [0; 8]; + unpack8(&input, &mut output, i); + for (idx, out) in output.iter().enumerate() { + assert_eq!(out.trailing_ones() as usize, i, "out[{idx}] = {out}"); + } + } + + for i in 0..=16 { + let mut output = [0; 16]; + unpack16(&input, &mut output, i); + for (idx, out) in output.iter().enumerate() { + assert_eq!(out.trailing_ones() as usize, i, "out[{idx}] = {out}"); + } + } + + for i in 0..=32 { + let mut output = [0; 32]; + unpack32(&input, &mut output, i); + for (idx, out) in output.iter().enumerate() { + assert_eq!(out.trailing_ones() as usize, i, "out[{idx}] = {out}"); + } + } + + for i in 0..=64 { + let mut output = [0; 64]; + unpack64(&input, &mut output, i); + for (idx, out) in output.iter().enumerate() { + assert_eq!(out.trailing_ones() as usize, i, "out[{idx}] = {out}"); + } + } + } +} diff --git a/vendor/parquet/src/util/bit_util.rs b/vendor/parquet/src/util/bit_util.rs new file mode 100644 index 0000000..d5f5945 --- /dev/null +++ b/vendor/parquet/src/util/bit_util.rs @@ -0,0 +1,1474 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::{cmp, mem::size_of}; + +use bytes::Bytes; + +use crate::data_type::{AsBytes, ByteArray, FixedLenByteArray, Int96}; +use crate::errors::{ParquetError, Result}; +use crate::util::bit_pack::{unpack8, unpack16, unpack32, unpack64}; + +#[inline] +fn array_from_slice(bs: &[u8]) -> Result<[u8; N]> { + // Need to slice as may be called with zero-padded values + match bs.get(..N) { + Some(b) => Ok(b.try_into().unwrap()), + None => Err(general_err!( + "error converting value, expected {} bytes got {}", + N, + bs.len() + )), + } +} + +/// Types that can be decoded from plain representations. This includes non-primitive types like +/// `FixedLenByteArray` and also variable length types like `ByteArray`. +pub trait FromBytes: Sized { + type Buffer: AsMut<[u8]> + Default; + fn try_from_le_slice(b: &[u8]) -> Result; + fn from_le_bytes(bs: Self::Buffer) -> Self; +} + +/// Types that can be decoded from bitpacked representations. +/// +/// This is implemented for primitive types and bool that can be +/// directly converted from a u64 value. Types like Int96, ByteArray, +/// and FixedLenByteArray that cannot be represented in 64 bits do not +/// implement this trait. +pub trait FromBitpacked { + /// The maximum number of bits that are allowed to be converted to this type. + /// This is at most the size of the type in bits, but could be less, for example + /// for the boolean type. + const BIT_CAPACITY: usize; + /// How many values are converted by one call to `unpack_batch`. + const BATCH_SIZE: usize; + /// Convert directly from a u64 value by truncation, avoiding byte slice copies. + fn from_u64(v: u64) -> Self; + + /// Converts multiple bitpacked values from `input` to `output`. + /// The `output` slice needs to have space for at least `BATCH_SIZE` elements, + /// otherwise this method will panic. + fn unpack_batch(input: &[u8], output: &mut [Self], num_bits: usize) + where + Self: Sized; +} + +macro_rules! from_le_bytes { + ($($ty: ty),*) => { + $( + impl FromBytes for $ty { + type Buffer = [u8; size_of::()]; + fn try_from_le_slice(b: &[u8]) -> Result { + Ok(Self::from_le_bytes(array_from_slice(b)?)) + } + fn from_le_bytes(bs: Self::Buffer) -> Self { + <$ty>::from_le_bytes(bs) + } + } + )* + }; +} + +macro_rules! from_bitpacked { + ($($ty: ty => $unpack: path),*) => { + $( + impl FromBitpacked for $ty { + const BIT_CAPACITY: usize = std::mem::size_of::<$ty>() * 8; + // this has to match the signature of the unpack* functions + const BATCH_SIZE: usize = std::mem::size_of::<$ty>() * 8; + + #[inline] + fn from_u64(v: u64) -> Self { + v as _ + } + + #[inline] + fn unpack_batch(input: &[u8], output: &mut [Self], num_bits: usize) { + $unpack(input, (&mut output[..Self::BATCH_SIZE]).try_into().unwrap(), num_bits) + } + } + )* + } +} + +macro_rules! from_bitpacked_delegate { + ($($ty: ty => $delegate: ty),*) => { + $( + impl FromBitpacked for $ty { + const BIT_CAPACITY: usize = <$delegate as FromBitpacked>::BIT_CAPACITY; + const BATCH_SIZE: usize = <$delegate as FromBitpacked>::BATCH_SIZE; + + #[inline] + fn from_u64(v: u64) -> Self { + v as _ + } + + #[inline] + fn unpack_batch(input: &[u8], output: &mut [Self], num_bits: usize) { + // Guard against misusages of this macro, due to the const block this will fail + // already at compile-time if the types are not compatible. + const { + assert!( + std::mem::size_of::<$ty>() == std::mem::size_of::<$delegate>() + && std::mem::align_of::<$ty>() == std::mem::align_of::<$delegate>(), + "types need to have the same size and alignment" + ); + } + // Safety: ty and delegate have the same size and alignment, and this macro is only used for types that have transmutable bit patterns. + let output: &mut [$delegate] = unsafe { std::slice::from_raw_parts_mut(output.as_mut_ptr().cast::<$delegate>(), output.len()) }; + <$delegate>::unpack_batch(input, output, num_bits); + } + } + )* + } +} + +from_le_bytes! { u8, u16, u32, u64, i8, i16, i32, i64 } +from_bitpacked!(u8 => unpack8, u16 => unpack16, u32 => unpack32, u64 => unpack64); +from_bitpacked_delegate!(i8 => u8, i16 => u16, i32 => u32, i64 => u64); + +impl FromBitpacked for bool { + const BIT_CAPACITY: usize = 1; + const BATCH_SIZE: usize = ::BATCH_SIZE; + + #[inline] + fn from_u64(v: u64) -> Self { + v != 0 + } + + #[inline] + fn unpack_batch(input: &[u8], output: &mut [Self], num_bits: usize) { + assert!(num_bits == 1); + // Safety: + // we asserted that we will only decode with a bitwidth of 1, + // so the u8 can only be 0 or 1, which are the valid representations of a bool. + let output: &mut [u8] = unsafe { + std::slice::from_raw_parts_mut(output.as_mut_ptr().cast::(), output.len()) + }; + u8::unpack_batch(input, output, num_bits); + } +} + +impl FromBytes for f32 { + type Buffer = [u8; 4]; + fn try_from_le_slice(b: &[u8]) -> Result { + Ok(Self::from_le_bytes(array_from_slice(b)?)) + } + fn from_le_bytes(bs: Self::Buffer) -> Self { + f32::from_le_bytes(bs) + } +} + +impl FromBytes for f64 { + type Buffer = [u8; 8]; + fn try_from_le_slice(b: &[u8]) -> Result { + Ok(Self::from_le_bytes(array_from_slice(b)?)) + } + fn from_le_bytes(bs: Self::Buffer) -> Self { + f64::from_le_bytes(bs) + } +} + +impl FromBytes for bool { + type Buffer = [u8; 1]; + + fn try_from_le_slice(b: &[u8]) -> Result { + Ok(Self::from_le_bytes(array_from_slice(b)?)) + } + fn from_le_bytes(bs: Self::Buffer) -> Self { + bs[0] != 0 + } +} + +impl FromBytes for Int96 { + type Buffer = [u8; 12]; + + fn try_from_le_slice(b: &[u8]) -> Result { + let bs: [u8; 12] = array_from_slice(b)?; + let mut i = Int96::new(); + i.set_data( + u32::try_from_le_slice(&bs[0..4])?, + u32::try_from_le_slice(&bs[4..8])?, + u32::try_from_le_slice(&bs[8..12])?, + ); + Ok(i) + } + + fn from_le_bytes(bs: Self::Buffer) -> Self { + let mut i = Int96::new(); + i.set_data( + u32::try_from_le_slice(&bs[0..4]).unwrap(), + u32::try_from_le_slice(&bs[4..8]).unwrap(), + u32::try_from_le_slice(&bs[8..12]).unwrap(), + ); + i + } +} + +impl FromBytes for ByteArray { + type Buffer = Vec; + + fn try_from_le_slice(b: &[u8]) -> Result { + Ok(b.to_vec().into()) + } + fn from_le_bytes(bs: Self::Buffer) -> Self { + bs.into() + } +} + +impl FromBytes for FixedLenByteArray { + type Buffer = Vec; + + fn try_from_le_slice(b: &[u8]) -> Result { + Ok(b.to_vec().into()) + } + fn from_le_bytes(bs: Self::Buffer) -> Self { + bs.into() + } +} + +/// Reads `size` of bytes from `src`, and reinterprets them as type `ty`, in +/// little-endian order. +/// This is copied and modified from byteorder crate. +pub(crate) fn read_num_bytes(size: usize, src: &[u8]) -> T +where + T: FromBytes, +{ + debug_assert!(size <= src.len()); + let mut buffer = ::Buffer::default(); + buffer.as_mut()[..size].copy_from_slice(&src[..size]); + ::from_le_bytes(buffer) +} + +/// Returns the ceil of value/divisor. +/// +/// This function should be removed after +/// [`int_roundings`](https://github.com/rust-lang/rust/issues/88581) is stable. +#[inline] +pub fn ceil(value: T, divisor: T) -> T { + num_integer::Integer::div_ceil(&value, &divisor) +} + +/// Returns the `num_bits` least-significant bits of `v` +#[inline] +pub fn trailing_bits(v: u64, num_bits: usize) -> u64 { + if num_bits >= 64 { + v + } else { + v & ((1 << num_bits) - 1) + } +} + +/// Returns the minimum number of bits needed to represent the value 'x' +#[inline] +pub fn num_required_bits(x: u64) -> u8 { + 64 - x.leading_zeros() as u8 +} + +static BIT_MASK: [u8; 8] = [1, 2, 4, 8, 16, 32, 64, 128]; + +/// Returns whether bit at position `i` in `data` is set or not +#[inline] +pub fn get_bit(data: &[u8], i: usize) -> bool { + (data[i >> 3] & BIT_MASK[i & 7]) != 0 +} + +/// Writes bit packed values to an in-memory buffer. +/// +/// `BitWriter` is the dual of [`BitReader`] and can write values that are either +/// byte aligned or packed at arbitrary bit widths. It is primarily used by the +/// Parquet RLE/bit-packing hybrid encoder. +/// +/// Bit-packed values are appended to an internal buffer in +/// little-endian bit order: the first value written occupies the +/// least-significant bits of the first byte. Bits that have not yet filled a +/// whole byte are held in an internal accumulator until a byte-aligning +/// operation (such as [`BitWriter::flush`], [`BitWriter::put_aligned`], or +/// [`BitWriter::consume`]) is called. +/// +/// Use [`BitWriter::consume`] to take ownership of the underlying buffer once +/// writing is complete. +/// +/// [`BitReader`]: crate::util::bit_util::BitReader +pub struct BitWriter { + /// Output Buffer + buffer: Vec, + /// Accumulator for in progress values + buffered_values: u64, + /// Current write offset within `buffered_values` + bit_offset: u8, +} + +impl BitWriter { + /// Creates a new [`BitWriter`] backed by an internal buffer of the given + /// initial capacity. + pub fn new(initial_capacity: usize) -> Self { + Self { + buffer: Vec::with_capacity(initial_capacity), + buffered_values: 0, + bit_offset: 0, + } + } + + /// Creates a new [`BitWriter`] that appends to the existing `buffer`. + /// + /// Data written with this writer are appended after existing values. + pub fn new_from_buf(buffer: Vec) -> Self { + Self { + buffer, + buffered_values: 0, + bit_offset: 0, + } + } + + /// Flushes any buffered bits to a byte boundary, then consumes this + /// writer and returns the underlying buffer. + #[inline] + pub fn consume(mut self) -> Vec { + self.flush(); + self.buffer + } + + /// Flushes any buffered bits to a byte boundary and returns a borrowed + /// view of the buffer's contents. + /// + /// This is the borrowing equivalent of [`BitWriter::consume`]. The writer + /// can continue to be used after this call. + #[inline] + pub fn flush_buffer(&mut self) -> &[u8] { + self.flush(); + self.buffer() + } + + /// Like [`BitWriter::flush_buffer`], but returns mutable access to the + /// buffer. + #[inline] + pub fn flush_buffer_mut(&mut self) -> &mut [u8] { + self.flush(); + &mut self.buffer + } + + /// Clears the internal state. + /// + /// Truncates the underlying buffer to length 0 (preserving its capacity) + /// and resets the bit accumulator. + #[inline] + pub fn clear(&mut self) { + self.buffer.clear(); + self.buffered_values = 0; + self.bit_offset = 0; + } + + /// Flushes any buffered bits and aligns the writer to the next byte + /// boundary. + /// + /// Any partial byte currently held in the bit accumulator is appended to + /// the buffer, and the accumulator is reset. Subsequent writes will start + /// at a byte boundary. + #[inline] + pub fn flush(&mut self) { + let num_bytes = ceil(self.bit_offset, 8); + let slice = &self.buffered_values.to_le_bytes()[..num_bytes as usize]; + self.buffer.extend_from_slice(slice); + self.buffered_values = 0; + self.bit_offset = 0; + } + + /// Reserves `num_bytes` bytes of zero-filled space at the current + /// position and returns the byte offset of the start of that region. + /// + /// Internally flushes any buffered bits first so the reservation begins + /// at a byte boundary. Use the returned offset together with + /// [`BitWriter::write_at`] or [`BitWriter::put_aligned_offset`] to fill + /// in the reserved bytes once their contents are known (for example, a + /// length prefix that depends on subsequently encoded data). + #[inline] + pub fn skip(&mut self, num_bytes: usize) -> usize { + self.flush(); + let result = self.buffer.len(); + self.buffer.extend(std::iter::repeat_n(0, num_bytes)); + result + } + + /// Reserves `num_bytes` bytes at the current position and returns a + /// mutable slice over them. + /// + /// Equivalent to [`BitWriter::skip`], but returns the reserved region + /// directly so it can be written into. Useful for filling in a header + /// (such as a length prefix) once the size of the following payload is + /// known. + #[inline] + pub fn get_next_byte_ptr(&mut self, num_bytes: usize) -> &mut [u8] { + let offset = self.skip(num_bytes); + &mut self.buffer[offset..offset + num_bytes] + } + + /// Returns the total number of bytes written so far, including any + /// partial byte still held in the bit accumulator (rounded up). + #[inline] + pub fn bytes_written(&self) -> usize { + self.buffer.len() + ceil(self.bit_offset, 8) as usize + } + + /// Returns a borrowed view of the bytes that have been flushed to the + /// underlying buffer so far. + /// + /// Note that bits currently held in the bit accumulator (i.e. not yet + /// flushed to a byte boundary) are not included. Use + /// [`BitWriter::flush_buffer`] to also flush pending bits before reading. + #[inline] + pub fn buffer(&self) -> &[u8] { + &self.buffer + } + + /// Returns the current offset within the output buffer. + /// + /// This is the index of the next byte that a byte-aligned write would + /// land in (excluding any bits currently held in the bit accumulator). + #[inline] + pub fn byte_offset(&self) -> usize { + self.buffer.len() + } + + /// Overwrites the byte at position `offset` in the underlying buffer + /// with `value`. + /// + /// Typically used together with [`BitWriter::skip`] or + /// [`BitWriter::get_next_byte_ptr`] to back-fill a previously reserved + /// byte once its value is known. + /// + /// # Panics + /// + /// Panics if `offset` is out of bounds for the underlying buffer. + pub fn write_at(&mut self, offset: usize, value: u8) { + self.buffer[offset] = value; + } + + /// Writes the `num_bits` least-significant bits of `v` to the writer in + /// bit-packed form. + /// + /// Values are packed in little-endian bit order: this call appends + /// `num_bits` bits starting at the current bit position. + /// + /// `num_bits` must be no larger than 64. + #[inline] + pub fn put_value(&mut self, v: u64, num_bits: usize) { + debug_assert!(num_bits <= 64); + let num_bits = num_bits as u8; + debug_assert_eq!(v.checked_shr(num_bits as u32).unwrap_or(0), 0); // covers case v >> 64 + + // Add value to buffered_values + self.buffered_values |= v << self.bit_offset; + self.bit_offset += num_bits; + if let Some(remaining) = self.bit_offset.checked_sub(64) { + self.buffer + .extend_from_slice(&self.buffered_values.to_le_bytes()); + self.bit_offset = remaining; + + // Perform checked right shift: v >> offset, where offset < 64, otherwise we + // shift all bits + self.buffered_values = v + .checked_shr((num_bits - self.bit_offset) as u32) + .unwrap_or(0); + } + } + + /// Writes the first `num_bytes` little-endian bytes of `val` to the + /// writer at the next byte boundary. + /// + /// Any buffered bits are first flushed so the value is byte-aligned in + /// the output. If `T` is wider than `num_bytes`, the high-order bytes + /// are silently truncated. + #[inline] + pub fn put_aligned(&mut self, val: T, num_bytes: usize) { + self.flush(); + let slice = val.as_bytes(); + let len = num_bytes.min(slice.len()); + self.buffer.extend_from_slice(&slice[..len]); + } + + /// Writes the first `num_bytes` little-endian bytes of `val` at the + /// given `offset` in the underlying buffer, overwriting any existing + /// data in `offset..offset + num_bytes`. + /// + /// `offset` is measured from the start of the internal buffer. If `T` + /// is wider than `num_bytes`, the high-order bytes are silently + /// truncated. + /// + /// Typically used together with [`BitWriter::skip`] to back-fill a + /// previously reserved region once its contents are known. + /// + /// # Panics + /// + /// Panics if `offset + min(size_of::(), num_bytes)` is out of bounds + /// for the underlying buffer. + #[inline] + pub fn put_aligned_offset(&mut self, val: T, num_bytes: usize, offset: usize) { + let slice = val.as_bytes(); + let len = num_bytes.min(slice.len()); + self.buffer[offset..offset + len].copy_from_slice(&slice[..len]) + } + + /// Writes `v` to the buffer in VLQ (variable-length quantity) encoding, + /// in little-endian byte order. + /// + /// Any buffered bits are first flushed so the encoding starts at a byte + /// boundary. The encoded form is between 1 and [`MAX_VLQ_BYTE_LEN`] + /// bytes long, depending on the magnitude of `v`. + #[inline] + pub fn put_vlq_int(&mut self, mut v: u64) { + while v & 0xFFFFFFFFFFFFFF80 != 0 { + self.put_aligned::(((v & 0x7F) | 0x80) as u8, 1); + v >>= 7; + } + self.put_aligned::((v & 0x7F) as u8, 1); + } + + /// Writes `v` to the buffer in zigzag-VLQ encoding, in little-endian + /// byte order. + /// + /// Zigzag-VLQ is a variant of VLQ encoding where negative and positive + /// numbers are interleaved so that small absolute values produce short + /// encodings regardless of sign. See the [Protocol Buffers encoding + /// documentation](https://developers.google.com/protocol-buffers/docs/encoding) + /// for details. + /// + /// As with [`BitWriter::put_vlq_int`], any buffered bits are first + /// flushed so the encoding starts at a byte boundary. + #[inline] + pub fn put_zigzag_vlq_int(&mut self, v: i64) { + let u: u64 = ((v << 1) ^ (v >> 63)) as u64; + self.put_vlq_int(u) + } + + /// Returns an estimate of the heap memory used by this writer, in bytes. + /// + /// This reflects the capacity of the underlying buffer rather than the + /// number of bytes actually written. + pub fn estimated_memory_size(&self) -> usize { + self.buffer.capacity() * size_of::() + } +} + +/// Maximum byte length for a VLQ encoded integer +/// MAX_VLQ_BYTE_LEN = 5 for i32, and MAX_VLQ_BYTE_LEN = 10 for i64 +pub const MAX_VLQ_BYTE_LEN: usize = 10; + +/// Reads bit packed values from an in-memory buffer. +/// +/// `BitReader` is the dual of [`BitWriter`] and reads values that are either +/// byte aligned or packed at arbitrary bit widths. It is primarily used by the +/// Parquet RLE/bit-packing hybrid decoder. +/// +/// Reads advance an internal cursor; once the buffer is exhausted, the +/// `get_*` methods return `None` rather than panicking. To rewind, use +/// [`BitReader::reset`] with the same (or a different) buffer. +pub struct BitReader { + /// The byte buffer to read from, passed in by client + buffer: Bytes, + + /// Bytes are memcpy'd from `buffer` and values are read from this variable. + /// This is faster than reading values byte by byte directly from `buffer` + /// + /// This is only populated when `self.bit_offset != 0` + buffered_values: u64, + + /// + /// End Start + /// |............|B|B|B|B|B|B|B|B|..............| + /// ^ ^ + /// bit_offset byte_offset + /// + /// Current byte offset in `buffer` + byte_offset: usize, + + /// Current bit offset in `buffered_values` + bit_offset: usize, +} + +impl BitReader { + /// Creates a new [`BitReader`] that reads from `buffer`, starting at + /// bit offset 0. + pub fn new(buffer: Bytes) -> Self { + BitReader { + buffer, + buffered_values: 0, + byte_offset: 0, + bit_offset: 0, + } + } + + /// Resets this reader to read from the start of `buffer`, discarding any + /// previous buffer and position. + /// + /// This is useful for reusing the same `BitReader` instance across + /// multiple input buffers without allocation. + pub fn reset(&mut self, buffer: Bytes) { + self.buffer = buffer; + self.buffered_values = 0; + self.byte_offset = 0; + self.bit_offset = 0; + } + + /// Returns the current byte offset, rounded up to the next whole byte. + /// + /// This is the index of the next byte that a byte-aligned + /// read (such as [`BitReader::get_aligned`]) would consume. + #[inline] + pub fn get_byte_offset(&self) -> usize { + self.byte_offset + ceil(self.bit_offset, 8) + } + + /// Reads a single bit-packed value of `num_bits` bits as a `T` from the + /// stream. + /// + /// The value is read as the low `num_bits` bits of `T`. Bits are consumed + /// from the stream in little-endian bit order. + /// + /// Returns `None` if there are fewer than `num_bits` bits left in the + /// buffer; otherwise `Some(value)`. On `None` the reader's position is + /// left unchanged. + pub fn get_value(&mut self, num_bits: usize) -> Option { + debug_assert!(num_bits <= 64); + debug_assert!(num_bits <= size_of::() * 8); + + if self.byte_offset * 8 + self.bit_offset + num_bits > self.buffer.len() * 8 { + return None; + } + + // If buffer is not byte aligned, `self.buffered_values` will + // have already been populated + if self.bit_offset == 0 { + self.load_buffered_values() + } + + let mut v = + trailing_bits(self.buffered_values, self.bit_offset + num_bits) >> self.bit_offset; + self.bit_offset += num_bits; + + if self.bit_offset >= 64 { + self.byte_offset += 8; + self.bit_offset -= 64; + + // If the new bit_offset is not 0, we need to read the next 64-bit chunk + // to buffered_values and update `v` + if self.bit_offset != 0 { + self.load_buffered_values(); + + v |= trailing_bits(self.buffered_values, self.bit_offset) + .wrapping_shl((num_bits - self.bit_offset) as u32); + } + } + + Some(T::from_u64(v)) + } + + /// Reads up to `batch.len()` bit-packed values of `num_bits` each, into + /// `batch`. + /// + /// Equivalent to repeatedly calling [`BitReader::get_value`] with the same + /// `num_bits`, but faster because it dispatches to SIMD-friendly + /// fixed-width unpacking routines whenever possible. + /// + /// Returns the number of values actually written to `batch`. This will be + /// less than `batch.len()` if the underlying buffer is exhausted before + /// `batch` is filled. + /// + /// # Panics + /// + /// This function panics if + /// - `num_bits` is larger than the bit-capacity of `T` + pub fn get_batch(&mut self, batch: &mut [T], num_bits: usize) -> usize { + debug_assert!(num_bits <= size_of::() * 8); + + let mut values_to_read = batch.len(); + let needed_bits = num_bits * values_to_read; + let remaining_bits = (self.buffer.len() - self.byte_offset) * 8 - self.bit_offset; + if remaining_bits < needed_bits { + values_to_read = remaining_bits / num_bits; + } + + let mut i = 0; + + // First align bit offset to byte offset + if self.bit_offset != 0 { + while i < values_to_read && self.bit_offset != 0 { + batch[i] = self + .get_value(num_bits) + .expect("expected to have more data"); + i += 1; + } + } + + assert_ne!(T::BIT_CAPACITY, 0); + assert!(num_bits <= T::BIT_CAPACITY); + + // Read directly into output buffer + while values_to_read - i >= T::BATCH_SIZE { + T::unpack_batch(&self.buffer[self.byte_offset..], &mut batch[i..], num_bits); + self.byte_offset += num_bits * T::BATCH_SIZE / 8; + i += T::BATCH_SIZE; + } + + // Try to read smaller batches if possible + if size_of::() > 4 && values_to_read - i >= 32 && num_bits <= 32 { + let mut out_buf = [0_u32; 32]; + unpack32(&self.buffer[self.byte_offset..], &mut out_buf, num_bits); + self.byte_offset += 4 * num_bits; + + for out in out_buf { + batch[i] = T::from_u64(out as u64); + i += 1; + } + } + + if size_of::() > 2 && values_to_read - i >= 16 && num_bits <= 16 { + let mut out_buf = [0_u16; 16]; + unpack16(&self.buffer[self.byte_offset..], &mut out_buf, num_bits); + self.byte_offset += 2 * num_bits; + + for out in out_buf { + batch[i] = T::from_u64(out as u64); + i += 1; + } + } + + if size_of::() > 1 && values_to_read - i >= 8 && num_bits <= 8 { + let mut out_buf = [0_u8; 8]; + unpack8(&self.buffer[self.byte_offset..], &mut out_buf, num_bits); + self.byte_offset += num_bits; + + for out in out_buf { + batch[i] = T::from_u64(out as u64); + i += 1; + } + } + + // Read any trailing values + while i < values_to_read { + let value = self + .get_value(num_bits) + .expect("expected to have more data"); + batch[i] = value; + i += 1; + } + + values_to_read + } + + /// Skips `num_values` bit-packed values of `num_bits` bits, advancing the + /// reader past them without decoding. + /// + /// Returns the number of values actually skipped (up to `num_values`). + /// This will be less than `num_values` if the underlying buffer is + /// exhausted. + pub fn skip(&mut self, num_values: usize, num_bits: usize) -> usize { + debug_assert!(num_bits <= 64); + + let needed_bits = num_bits * num_values; + let remaining_bits = (self.buffer.len() - self.byte_offset) * 8 - self.bit_offset; + + let values_to_read = match remaining_bits < needed_bits { + true => remaining_bits / num_bits, + false => num_values, + }; + + let end_bit_offset = self.byte_offset * 8 + values_to_read * num_bits + self.bit_offset; + + self.byte_offset = end_bit_offset / 8; + self.bit_offset = end_bit_offset % 8; + + if self.bit_offset != 0 { + self.load_buffered_values() + } + + values_to_read + } + + /// Reads up to `num_bytes` bytes from the stream, appending them to `buf`, + /// and returns the number of bytes actually appended. + /// + /// The reader is first advanced to the next byte boundary, so any + /// in-progress bit-level read is discarded before the bytes are copied. + pub(crate) fn get_aligned_bytes(&mut self, buf: &mut Vec, num_bytes: usize) -> usize { + // Align to byte offset + self.byte_offset = self.get_byte_offset(); + self.bit_offset = 0; + + let src = &self.buffer[self.byte_offset..]; + let to_read = num_bytes.min(src.len()); + buf.extend_from_slice(&src[..to_read]); + + self.byte_offset += to_read; + + to_read + } + + /// Reads a `num_bytes`-sized value of type `T` from the stream. + /// + /// `T` is interpreted as a little-endian native type. The value is + /// assumed to be byte aligned, so the reader is first advanced to the + /// start of the next byte before reading. + /// + /// Returns `Some(value)` if there are at least `num_bytes` bytes left in + /// the buffer after byte-alignment, and `None` otherwise. On `None` the + /// reader's byte position is still advanced to the alignment boundary. + pub fn get_aligned(&mut self, num_bytes: usize) -> Option { + self.byte_offset = self.get_byte_offset(); + self.bit_offset = 0; + + if self.byte_offset + num_bytes > self.buffer.len() { + return None; + } + + // Advance byte_offset to next unread byte and read num_bytes + let v = read_num_bytes::(num_bytes, &self.buffer[self.byte_offset..]); + self.byte_offset += num_bytes; + + Some(v) + } + + /// Reads a VLQ-encoded (in little-endian order) integer from the stream. + /// + /// The encoded integer must start at the beginning of a byte; the reader + /// is first advanced to the next byte boundary before decoding. + /// + /// Returns `Some(value)` on success, or `None` if the buffer is exhausted + /// before a complete VLQ value is read. + /// + /// # Panics + /// + /// Panics if the encoded integer is longer than [`MAX_VLQ_BYTE_LEN`] + /// bytes (bad input). + pub fn get_vlq_int(&mut self) -> Option { + // Align to byte boundary once, then read bytes directly + self.byte_offset = self.get_byte_offset(); + self.bit_offset = 0; + + let buf = &self.buffer[self.byte_offset..]; + let mut shift = 0; + let mut v: i64 = 0; + + for (i, &byte) in buf.iter().enumerate() { + v |= ((byte & 0x7F) as i64) << shift; + shift += 7; + assert!( + shift <= MAX_VLQ_BYTE_LEN * 7, + "Num of bytes exceed MAX_VLQ_BYTE_LEN ({MAX_VLQ_BYTE_LEN})" + ); + if byte & 0x80 == 0 { + self.byte_offset += i + 1; + return Some(v); + } + } + None + } + + /// Reads a zigzag-VLQ-encoded little-endian integer from the + /// stream. + /// + /// Zigzag-VLQ is a variant of VLQ encoding where negative and positive + /// numbers are interleaved so that small absolute values produce short + /// encodings regardless of sign. See the [Protocol Buffers encoding + /// documentation](https://developers.google.com/protocol-buffers/docs/encoding) + /// for details. + /// + /// As with [`BitReader::get_vlq_int`], the encoded integer must start at + /// the beginning of a byte; the reader is first advanced to the next + /// byte boundary before decoding. + /// + /// Returns `Some(value)` on success, or `None` if the buffer is exhausted + /// before a complete value is read. + #[inline] + pub fn get_zigzag_vlq_int(&mut self) -> Option { + self.get_vlq_int().map(|v| { + let u = v as u64; + (u >> 1) as i64 ^ -((u & 1) as i64) + }) + } + + /// Loads up to the next 8 bytes from `self.buffer` at `self.byte_offset` + /// into `self.buffered_values`. + /// + /// Reads fewer than 8 bytes if there are fewer than 8 bytes left + #[inline] + fn load_buffered_values(&mut self) { + let bytes_to_read = cmp::min(self.buffer.len() - self.byte_offset, 8); + self.buffered_values = + read_num_bytes::(bytes_to_read, &self.buffer[self.byte_offset..]); + } +} + +impl From> for BitReader { + #[inline] + fn from(buffer: Vec) -> Self { + BitReader::new(buffer.into()) + } +} + +/// Parallel bit extract: for each set bit in `mask`, extract the +/// corresponding bit from `value` and pack them contiguously into the low +/// bits of the return value. +/// +/// Equivalent to the x86 BMI2 `PEXT` instruction. When compiled with the +/// `bmi2` target feature enabled (for example `-C target-cpu=x86-64-v3`) +/// this lowers to the hardware `pext` instruction; otherwise it falls back +/// to a portable scalar loop. +/// +/// Replace with `value.compress(mask)` when `uint_gather_scatter_bits` +/// is stabilised: +#[allow(dead_code)] +#[inline] +pub(crate) fn compress(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 `pext` instruction is guaranteed to be available. + unsafe { std::arch::x86_64::_pext_u64(value, mask) } + } + + #[cfg(not(all(target_arch = "x86_64", target_feature = "bmi2")))] + { + let mut mask = mask; + let mut result: u64 = 0; + let mut dest_bit: u64 = 1; + while mask != 0 { + let lowest = mask & mask.wrapping_neg(); + if value & lowest != 0 { + result |= dest_bit; + } + dest_bit <<= 1; + mask ^= lowest; + } + result + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use crate::util::test_common::rand_gen::random_numbers; + use rand::distr::{Distribution, StandardUniform}; + use std::fmt::Debug; + + #[test] + fn test_compress() { + // Reference: gather the `mask`-selected bits of `value` into + // contiguous low bits, least-significant first. + fn reference(value: u64, mut mask: u64) -> u64 { + let mut result = 0u64; + let mut dest = 0u32; + while mask != 0 { + let lowest = mask & mask.wrapping_neg(); + result |= (((value & lowest) != 0) as u64) << dest; + dest += 1; + mask ^= lowest; + } + result + } + + // Hand-picked edge cases. + assert_eq!(compress(0b1010, 0b1111), 0b1010); + assert_eq!(compress(0b1010, 0b1010), 0b11); + assert_eq!(compress(0b1010, 0b0101), 0); + assert_eq!(compress(u64::MAX, 0), 0); + assert_eq!(compress(0, u64::MAX), 0); + assert_eq!(compress(u64::MAX, u64::MAX), u64::MAX); + + // Randomised cross-check against the reference. On a `bmi2` build + // this validates the hardware `pext` path; otherwise it exercises + // the portable fallback. + let values = random_numbers::(1024); + let masks = random_numbers::(1024); + for (&value, &mask) in values.iter().zip(masks.iter()) { + assert_eq!( + compress(value, mask), + reference(value, mask), + "value={value:#x} mask={mask:#x}" + ); + } + } + + #[test] + fn test_ceil() { + assert_eq!(ceil(0, 1), 0); + assert_eq!(ceil(1, 1), 1); + assert_eq!(ceil(1, 2), 1); + assert_eq!(ceil(1, 8), 1); + assert_eq!(ceil(7, 8), 1); + assert_eq!(ceil(8, 8), 1); + assert_eq!(ceil(9, 8), 2); + assert_eq!(ceil(9, 9), 1); + assert_eq!(ceil(10000000000_u64, 10), 1000000000); + assert_eq!(ceil(10_u64, 10000000000), 1); + assert_eq!(ceil(10000000000_u64, 1000000000), 10); + } + + #[test] + fn test_bit_reader_get_byte_offset() { + let buffer = vec![255; 10]; + let mut bit_reader = BitReader::from(buffer); + assert_eq!(bit_reader.get_byte_offset(), 0); // offset (0 bytes, 0 bits) + bit_reader.get_value::(6); + assert_eq!(bit_reader.get_byte_offset(), 1); // offset (0 bytes, 6 bits) + bit_reader.get_value::(10); + assert_eq!(bit_reader.get_byte_offset(), 2); // offset (0 bytes, 16 bits) + bit_reader.get_value::(20); + assert_eq!(bit_reader.get_byte_offset(), 5); // offset (0 bytes, 36 bits) + bit_reader.get_value::(30); + assert_eq!(bit_reader.get_byte_offset(), 9); // offset (8 bytes, 2 bits) + } + + #[test] + fn test_bit_reader_get_value() { + let buffer = vec![255, 0]; + let mut bit_reader = BitReader::from(buffer); + assert_eq!(bit_reader.get_value::(1), Some(1)); + assert_eq!(bit_reader.get_value::(2), Some(3)); + assert_eq!(bit_reader.get_value::(3), Some(7)); + assert_eq!(bit_reader.get_value::(4), Some(3)); + } + + #[test] + fn test_bit_reader_skip() { + let buffer = vec![255, 0]; + let mut bit_reader = BitReader::from(buffer); + let skipped = bit_reader.skip(1, 1); + assert_eq!(skipped, 1); + assert_eq!(bit_reader.get_value::(1), Some(1)); + let skipped = bit_reader.skip(2, 2); + assert_eq!(skipped, 2); + assert_eq!(bit_reader.get_value::(2), Some(3)); + let skipped = bit_reader.skip(4, 1); + assert_eq!(skipped, 4); + assert_eq!(bit_reader.get_value::(4), Some(0)); + let skipped = bit_reader.skip(1, 1); + assert_eq!(skipped, 0); + } + + #[test] + fn test_bit_reader_get_value_boundary() { + let buffer = vec![10, 0, 0, 0, 20, 0, 30, 0, 0, 0, 40, 0]; + let mut bit_reader = BitReader::from(buffer); + assert_eq!(bit_reader.get_value::(32), Some(10)); + assert_eq!(bit_reader.get_value::(16), Some(20)); + assert_eq!(bit_reader.get_value::(32), Some(30)); + assert_eq!(bit_reader.get_value::(16), Some(40)); + } + + #[test] + fn test_bit_reader_skip_boundary() { + let buffer = vec![10, 0, 0, 0, 20, 0, 30, 0, 0, 0, 40, 0]; + let mut bit_reader = BitReader::from(buffer); + assert_eq!(bit_reader.get_value::(32), Some(10)); + assert_eq!(bit_reader.skip(1, 16), 1); + assert_eq!(bit_reader.get_value::(32), Some(30)); + assert_eq!(bit_reader.get_value::(16), Some(40)); + } + + #[test] + fn test_bit_reader_get_aligned() { + // 01110101 11001011 + let buffer = Bytes::from(vec![0x75, 0xCB]); + let mut bit_reader = BitReader::new(buffer.clone()); + assert_eq!(bit_reader.get_value::(3), Some(5)); + assert_eq!(bit_reader.get_aligned::(1), Some(203)); + assert_eq!(bit_reader.get_value::(1), None); + bit_reader.reset(buffer.clone()); + assert_eq!(bit_reader.get_aligned::(3), None); + } + + #[test] + fn test_bit_reader_get_vlq_int() { + // 10001001 00000001 11110010 10110101 00000110 + let buffer: Vec = vec![0x89, 0x01, 0xF2, 0xB5, 0x06]; + let mut bit_reader = BitReader::from(buffer); + assert_eq!(bit_reader.get_vlq_int(), Some(137)); + assert_eq!(bit_reader.get_vlq_int(), Some(105202)); + } + + #[test] + fn test_bit_reader_get_zigzag_vlq_int() { + let buffer: Vec = vec![0, 1, 2, 3]; + let mut bit_reader = BitReader::from(buffer); + assert_eq!(bit_reader.get_zigzag_vlq_int(), Some(0)); + assert_eq!(bit_reader.get_zigzag_vlq_int(), Some(-1)); + assert_eq!(bit_reader.get_zigzag_vlq_int(), Some(1)); + assert_eq!(bit_reader.get_zigzag_vlq_int(), Some(-2)); + } + + #[test] + fn test_num_required_bits() { + assert_eq!(num_required_bits(0), 0); + assert_eq!(num_required_bits(1), 1); + assert_eq!(num_required_bits(2), 2); + assert_eq!(num_required_bits(4), 3); + assert_eq!(num_required_bits(8), 4); + assert_eq!(num_required_bits(10), 4); + assert_eq!(num_required_bits(12), 4); + assert_eq!(num_required_bits(16), 5); + assert_eq!(num_required_bits(u64::MAX), 64); + } + + #[test] + fn test_get_bit() { + // 00001101 + assert!(get_bit(&[0b00001101], 0)); + assert!(!get_bit(&[0b00001101], 1)); + assert!(get_bit(&[0b00001101], 2)); + assert!(get_bit(&[0b00001101], 3)); + + // 01001001 01010010 + assert!(get_bit(&[0b01001001, 0b01010010], 0)); + assert!(!get_bit(&[0b01001001, 0b01010010], 1)); + assert!(!get_bit(&[0b01001001, 0b01010010], 2)); + assert!(get_bit(&[0b01001001, 0b01010010], 3)); + assert!(!get_bit(&[0b01001001, 0b01010010], 4)); + assert!(!get_bit(&[0b01001001, 0b01010010], 5)); + assert!(get_bit(&[0b01001001, 0b01010010], 6)); + assert!(!get_bit(&[0b01001001, 0b01010010], 7)); + assert!(!get_bit(&[0b01001001, 0b01010010], 8)); + assert!(get_bit(&[0b01001001, 0b01010010], 9)); + assert!(!get_bit(&[0b01001001, 0b01010010], 10)); + assert!(!get_bit(&[0b01001001, 0b01010010], 11)); + assert!(get_bit(&[0b01001001, 0b01010010], 12)); + assert!(!get_bit(&[0b01001001, 0b01010010], 13)); + assert!(get_bit(&[0b01001001, 0b01010010], 14)); + assert!(!get_bit(&[0b01001001, 0b01010010], 15)); + } + + #[test] + fn test_skip() { + let mut writer = BitWriter::new(5); + let old_offset = writer.skip(1); + writer.put_aligned(42, 4); + writer.put_aligned_offset(0x10, 1, old_offset); + let result = writer.consume(); + assert_eq!(result.as_ref(), [0x10, 42, 0, 0, 0]); + + writer = BitWriter::new(4); + let result = writer.skip(5); + assert_eq!(result, 0); + assert_eq!(writer.buffer(), &[0; 5]) + } + + #[test] + fn test_get_next_byte_ptr() { + let mut writer = BitWriter::new(5); + { + let first_byte = writer.get_next_byte_ptr(1); + first_byte[0] = 0x10; + } + writer.put_aligned(42, 4); + let result = writer.consume(); + assert_eq!(result.as_ref(), [0x10, 42, 0, 0, 0]); + } + + #[test] + fn test_consume_flush_buffer() { + let mut writer1 = BitWriter::new(3); + let mut writer2 = BitWriter::new(3); + for i in 1..10 { + writer1.put_value(i, 4); + writer2.put_value(i, 4); + } + let res1 = writer1.flush_buffer(); + let res2 = writer2.consume(); + assert_eq!(res1, &res2[..]); + } + + #[test] + fn test_put_get_bool() { + let len = 8; + let mut writer = BitWriter::new(len); + + for i in 0..8 { + writer.put_value(i % 2, 1); + } + + writer.flush(); + { + let buffer = writer.buffer(); + assert_eq!(buffer[0], 0b10101010); + } + + // Write 00110011 + for i in 0..8 { + match i { + 0 | 1 | 4 | 5 => writer.put_value(false as u64, 1), + _ => writer.put_value(true as u64, 1), + } + } + writer.flush(); + { + let buffer = writer.buffer(); + assert_eq!(buffer[0], 0b10101010); + assert_eq!(buffer[1], 0b11001100); + } + + let mut reader = BitReader::from(writer.consume()); + + for i in 0..8 { + let val = reader + .get_value::(1) + .expect("get_value() should return OK"); + assert_eq!(val, i % 2); + } + + for i in 0..8 { + let val = reader + .get_value::(1) + .expect("get_value() should return OK"); + match i { + 0 | 1 | 4 | 5 => assert!(!val), + _ => assert!(val), + } + } + } + + #[test] + fn test_put_value_roundtrip() { + test_put_value_rand_numbers(32, 2); + test_put_value_rand_numbers(32, 3); + test_put_value_rand_numbers(32, 4); + test_put_value_rand_numbers(32, 5); + test_put_value_rand_numbers(32, 6); + test_put_value_rand_numbers(32, 7); + test_put_value_rand_numbers(32, 8); + test_put_value_rand_numbers(64, 16); + test_put_value_rand_numbers(64, 24); + test_put_value_rand_numbers(64, 32); + } + + fn test_put_value_rand_numbers(total: usize, num_bits: usize) { + assert!(num_bits < 64); + let num_bytes = ceil(num_bits, 8); + let mut writer = BitWriter::new(num_bytes * total); + let values: Vec = random_numbers::(total) + .iter() + .map(|v| v & ((1 << num_bits) - 1)) + .collect(); + (0..total).for_each(|i| writer.put_value(values[i], num_bits)); + + let mut reader = BitReader::from(writer.consume()); + (0..total).for_each(|i| { + let v = reader + .get_value::(num_bits) + .expect("get_value() should return OK"); + assert_eq!( + v, values[i], + "[{}]: expected {} but got {}", + i, values[i], v + ); + }); + } + + #[test] + fn test_get_batch() { + const SIZE: &[usize] = &[1, 31, 32, 33, 128, 129]; + for s in SIZE { + for i in 0..=64 { + // `get_batch` is generic over all of these types, so exercise the + // signed integer types and `bool` alongside the unsigned types rather + // than relying on the unsigned coverage alone. + match i { + 0..=8 => { + test_get_batch_helper::(*s, i); + test_get_batch_helper::(*s, i); + } + 9..=16 => { + test_get_batch_helper::(*s, i); + test_get_batch_helper::(*s, i); + } + 17..=32 => { + test_get_batch_helper::(*s, i); + test_get_batch_helper::(*s, i); + } + _ => { + test_get_batch_helper::(*s, i); + test_get_batch_helper::(*s, i); + } + } + // `bool` only supports a bit width of 1. + if i == 1 { + test_get_batch_helper::(*s, i); + } + } + } + } + + fn test_get_batch_helper(total: usize, num_bits: usize) + where + T: FromBitpacked + Default + Clone + Debug + Eq, + { + assert!(num_bits <= 64); + let num_bytes = ceil(num_bits, 8); + let mut writer = BitWriter::new(num_bytes * total); + + let mask = match num_bits { + 64 => u64::MAX, + _ => (1 << num_bits) - 1, + }; + + let values: Vec = random_numbers::(total) + .iter() + .map(|v| v & mask) + .collect(); + + // Generic values used to check against actual values read from `get_batch`. + let expected_values: Vec = values.iter().map(|v| T::from_u64(*v)).collect(); + + (0..total).for_each(|i| writer.put_value(values[i], num_bits)); + + let buf = writer.consume(); + let mut reader = BitReader::from(buf); + let mut batch = vec![T::default(); values.len()]; + let values_read = reader.get_batch::(&mut batch, num_bits); + assert_eq!(values_read, values.len()); + for i in 0..batch.len() { + assert_eq!( + batch[i], + expected_values[i], + "max_num_bits = {}, num_bits = {}, index = {}", + size_of::() * 8, + num_bits, + i + ); + } + } + + #[test] + fn test_put_aligned_roundtrip() { + test_put_aligned_rand_numbers::(4, 3); + test_put_aligned_rand_numbers::(16, 5); + test_put_aligned_rand_numbers::(32, 7); + test_put_aligned_rand_numbers::(32, 9); + test_put_aligned_rand_numbers::(32, 11); + test_put_aligned_rand_numbers::(32, 13); + test_put_aligned_rand_numbers::(32, 17); + test_put_aligned_rand_numbers::(32, 23); + } + + fn test_put_aligned_rand_numbers(total: usize, num_bits: usize) + where + T: Copy + FromBytes + AsBytes + Debug + PartialEq, + StandardUniform: Distribution, + { + assert!(num_bits <= 32); + assert!(total % 2 == 0); + + let aligned_value_byte_width = std::mem::size_of::(); + let value_byte_width = ceil(num_bits, 8); + let mut writer = + BitWriter::new((total / 2) * (aligned_value_byte_width + value_byte_width)); + let values: Vec = random_numbers::(total / 2) + .iter() + .map(|v| v & ((1 << num_bits) - 1)) + .collect(); + let aligned_values = random_numbers::(total / 2); + + for i in 0..total { + let j = i / 2; + if i % 2 == 0 { + writer.put_value(values[j] as u64, num_bits); + } else { + writer.put_aligned::(aligned_values[j], aligned_value_byte_width) + } + } + + let mut reader = BitReader::from(writer.consume()); + for i in 0..total { + let j = i / 2; + if i % 2 == 0 { + let v = reader + .get_value::(num_bits) + .expect("get_value() should return OK"); + assert_eq!( + v, values[j] as u64, + "[{}]: expected {} but got {}", + i, values[j], v + ); + } else { + let v = reader + .get_aligned::(aligned_value_byte_width) + .expect("get_aligned() should return OK"); + assert_eq!( + v, aligned_values[j], + "[{}]: expected {:?} but got {:?}", + i, aligned_values[j], v + ); + } + } + } + + #[test] + fn test_put_vlq_int() { + let total = 64; + let mut writer = BitWriter::new(total * 32); + let values = random_numbers::(total); + (0..total).for_each(|i| writer.put_vlq_int(values[i] as u64)); + + let mut reader = BitReader::from(writer.consume()); + (0..total).for_each(|i| { + let v = reader + .get_vlq_int() + .expect("get_vlq_int() should return OK"); + assert_eq!( + v as u32, values[i], + "[{}]: expected {} but got {}", + i, values[i], v + ); + }); + } + + #[test] + fn test_put_zigzag_vlq_int() { + let total = 64; + let mut writer = BitWriter::new(total * 32); + let values = random_numbers::(total); + (0..total).for_each(|i| writer.put_zigzag_vlq_int(values[i] as i64)); + + let mut reader = BitReader::from(writer.consume()); + (0..total).for_each(|i| { + let v = reader + .get_zigzag_vlq_int() + .expect("get_zigzag_vlq_int() should return OK"); + assert_eq!( + v as i32, values[i], + "[{}]: expected {} but got {}", + i, values[i], v + ); + }); + } + + #[test] + fn test_get_batch_zero_extend() { + let to_read = vec![0xFF; 4]; + let mut reader = BitReader::from(to_read); + + // Create a non-zeroed output buffer + let mut output = [u64::MAX; 32]; + reader.get_batch(&mut output, 1); + + for v in output { + // Values should be read correctly + assert_eq!(v, 1); + } + } +} diff --git a/vendor/parquet/src/util/interner.rs b/vendor/parquet/src/util/interner.rs new file mode 100644 index 0000000..deae372 --- /dev/null +++ b/vendor/parquet/src/util/interner.rs @@ -0,0 +1,93 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::data_type::AsBytes; +use hashbrown::HashTable; + +const DEFAULT_DEDUP_CAPACITY: usize = 4096; + +/// Storage trait for [`Interner`] +pub trait Storage { + type Key: Copy; + + type Value: AsBytes + ?Sized; + + /// Gets an element by its key + fn get(&self, idx: Self::Key) -> &Self::Value; + + /// Adds a new element, returning the key + fn push(&mut self, value: &Self::Value) -> Self::Key; + + /// Return an estimate of the memory used in this storage, in bytes + #[allow(dead_code)] // not used in parquet_derive, so is dead there + fn estimated_memory_size(&self) -> usize; +} + +/// A generic value interner supporting various different [`Storage`] +#[derive(Debug, Default)] +pub struct Interner { + state: ahash::RandomState, + + /// Used to provide a lookup from value to unique value + dedup: HashTable, + + storage: S, +} + +impl Interner { + /// Create a new `Interner` with the provided storage + pub fn new(storage: S) -> Self { + Self { + state: Default::default(), + dedup: HashTable::with_capacity(DEFAULT_DEDUP_CAPACITY), + storage, + } + } + + /// Intern the value, returning the interned key, and if this was a new value + pub fn intern(&mut self, value: &S::Value) -> S::Key { + let hash = self.state.hash_one(value.as_bytes()); + + *self + .dedup + .entry( + hash, + // Compare bytes rather than directly comparing values so NaNs can be interned + |index| value.as_bytes() == self.storage.get(*index).as_bytes(), + |key| self.state.hash_one(self.storage.get(*key).as_bytes()), + ) + .or_insert_with(|| self.storage.push(value)) + .get() + } + + /// Return estimate of the memory used, in bytes + #[allow(dead_code)] // not used in parquet_derive, so is dead there + pub fn estimated_memory_size(&self) -> usize { + self.storage.estimated_memory_size() + self.dedup.allocation_size() + } + + /// Returns the storage for this interner + pub fn storage(&self) -> &S { + &self.storage + } + + /// Unwraps the inner storage + #[cfg(feature = "arrow")] + pub fn into_inner(self) -> S { + self.storage + } +} diff --git a/vendor/parquet/src/util/mod.rs b/vendor/parquet/src/util/mod.rs new file mode 100644 index 0000000..145cdd6 --- /dev/null +++ b/vendor/parquet/src/util/mod.rs @@ -0,0 +1,31 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#[macro_use] +pub mod bit_util; +mod bit_pack; +pub(crate) mod interner; + +pub mod push_buffers; +#[cfg(any(test, feature = "test_common"))] +pub(crate) mod test_common; +pub mod utf8; + +#[cfg(any(test, feature = "test_common"))] +pub use self::test_common::page_util::{ + DataPageBuilder, DataPageBuilderImpl, InMemoryPageIterator, +}; diff --git a/vendor/parquet/src/util/push_buffers.rs b/vendor/parquet/src/util/push_buffers.rs new file mode 100644 index 0000000..32c9ba8 --- /dev/null +++ b/vendor/parquet/src/util/push_buffers.rs @@ -0,0 +1,228 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::errors::ParquetError; +use crate::file::reader::{ChunkReader, Length}; +use bytes::Bytes; +use std::fmt::Display; +use std::ops::Range; + +/// Holds multiple non-contiguous, caller-provided buffers of file data. +/// +/// This is the in-memory buffer used by the push-based Parquet decoders +/// (`ParquetPushDecoder` and `ParquetMetaDataPushDecoder`). It can be +/// constructed up front and handed to a builder so the decoder reuses bytes +/// that have already been fetched. +/// +/// Features: +/// 1. Zero copy +/// 2. non contiguous ranges of bytes +/// +/// # Non Coalescing +/// +/// This buffer does not coalesce (merging adjacent ranges of bytes into a +/// single range). Coalescing at this level would require copying the data but +/// the caller may already have the needed data in a single buffer which would +/// require no copying. +/// +/// Thus, the implementation defers to the caller to coalesce subsequent requests +/// if desired. +#[derive(Debug, Clone, Default)] +pub struct PushBuffers { + /// the virtual "offset" of this buffers (added to any request) + offset: u64, + /// The total length of the file being decoded + file_len: u64, + /// The ranges of data that are available for decoding (not adjusted for offset) + ranges: Vec>, + /// The buffers of data that can be used to decode the Parquet file + buffers: Vec, +} + +impl Display for PushBuffers { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + writeln!( + f, + "Buffers (offset: {}, file_len: {})", + self.offset, self.file_len + )?; + writeln!(f, "Available Ranges (w/ offset):")?; + for range in &self.ranges { + writeln!( + f, + " {}..{} ({}..{}): {} bytes", + range.start, + range.end, + range.start + self.offset, + range.end + self.offset, + range.end - range.start + )?; + } + + Ok(()) + } +} + +impl PushBuffers { + /// Create a new, empty `PushBuffers` for a file of the given length. + /// + /// Use [`PushBuffers::default`] when the file length is unknown or + /// irrelevant (e.g. the push decoder, which tracks ranges by absolute + /// offset and never consults `file_len`). + pub fn new(file_len: u64) -> Self { + Self { + offset: 0, + file_len, + ranges: Vec::new(), + buffers: Vec::new(), + } + } + + /// Push all the ranges and buffers + pub fn push_ranges(&mut self, ranges: Vec>, buffers: Vec) { + assert_eq!( + ranges.len(), + buffers.len(), + "Number of ranges must match number of buffers" + ); + for (range, buffer) in ranges.into_iter().zip(buffers) { + self.push_range(range, buffer); + } + } + + /// Push a new range and its associated buffer + pub fn push_range(&mut self, range: Range, buffer: Bytes) { + assert_eq!( + (range.end - range.start) as usize, + buffer.len(), + "Range length must match buffer length" + ); + self.ranges.push(range); + self.buffers.push(buffer); + } + + /// Returns true if the Buffers contains data for the given range + pub(crate) fn has_range(&self, range: &Range) -> bool { + self.ranges + .iter() + .any(|r| r.start <= range.start && r.end >= range.end) + } + + fn iter(&self) -> impl Iterator, &Bytes)> { + self.ranges.iter().zip(self.buffers.iter()) + } + + /// return the file length of the Parquet file being read + pub(crate) fn file_len(&self) -> u64 { + self.file_len + } + + /// Specify a new offset + fn with_offset(mut self, offset: u64) -> Self { + self.offset = offset; + self + } + + /// Return the total of all buffered ranges + #[cfg(feature = "arrow")] + pub(crate) fn buffered_bytes(&self) -> u64 { + self.ranges.iter().map(|r| r.end - r.start).sum() + } + + /// Clear any range and corresponding buffer that is exactly in the ranges_to_clear + #[cfg(feature = "arrow")] + pub(crate) fn clear_ranges(&mut self, ranges_to_clear: &[Range]) { + let mut new_ranges = Vec::new(); + let mut new_buffers = Vec::new(); + + for (range, buffer) in self.iter() { + if !ranges_to_clear + .iter() + .any(|r| r.start == range.start && r.end == range.end) + { + new_ranges.push(range.clone()); + new_buffers.push(buffer.clone()); + } + } + self.ranges = new_ranges; + self.buffers = new_buffers; + } + + /// Clear all buffered ranges and their corresponding data + pub(crate) fn clear_all_ranges(&mut self) { + self.ranges.clear(); + self.buffers.clear(); + } +} + +impl Length for PushBuffers { + fn len(&self) -> u64 { + self.file_len + } +} + +/// less efficient implementation of Read for Buffers +impl std::io::Read for PushBuffers { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + // Find the range that contains the start offset + let mut found = false; + for (range, data) in self.iter() { + if range.start <= self.offset && range.end >= self.offset + buf.len() as u64 { + // Found the range, figure out the starting offset in the buffer + let start_offset = (self.offset - range.start) as usize; + let end_offset = start_offset + buf.len(); + let slice = data.slice(start_offset..end_offset); + buf.copy_from_slice(slice.as_ref()); + found = true; + break; + } + } + if found { + // If we found the range, we can return the number of bytes read + // advance our offset + self.offset += buf.len() as u64; + Ok(buf.len()) + } else { + Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "No data available in Buffers", + )) + } + } +} + +impl ChunkReader for PushBuffers { + type T = Self; + + fn get_read(&self, start: u64) -> Result { + Ok(self.clone().with_offset(self.offset + start)) + } + + fn get_bytes(&self, start: u64, length: usize) -> Result { + // find the range that contains the start offset + for (range, data) in self.iter() { + if range.start <= start && range.end >= start + length as u64 { + // Found the range, figure out the starting offset in the buffer + let start_offset = (start - range.start) as usize; + return Ok(data.slice(start_offset..start_offset + length)); + } + } + // Signal that we need more data + let requested_end = start + length as u64; + Err(ParquetError::NeedMoreDataRange(start..requested_end)) + } +} diff --git a/vendor/parquet/src/util/test_common/file_util.rs b/vendor/parquet/src/util/test_common/file_util.rs new file mode 100644 index 0000000..6c03135 --- /dev/null +++ b/vendor/parquet/src/util/test_common/file_util.rs @@ -0,0 +1,37 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::{fs, path::PathBuf, str::FromStr}; + +/// Returns path to the test parquet file in 'data' directory +pub fn get_test_path(file_name: &str) -> PathBuf { + let mut pathbuf = PathBuf::from_str(&arrow::util::test_util::parquet_test_data()).unwrap(); + pathbuf.push(file_name); + pathbuf +} + +/// Returns file handle for a test parquet file from 'data' directory +pub fn get_test_file(file_name: &str) -> fs::File { + let path = get_test_path(file_name); + fs::File::open(path.as_path()).unwrap_or_else(|err| { + panic!( + "Test file {} could not be opened, did you do `git submodule update`?: {}", + path.display(), + err + ) + }) +} diff --git a/vendor/parquet/src/util/test_common/mod.rs b/vendor/parquet/src/util/test_common/mod.rs new file mode 100644 index 0000000..8cfc1e6 --- /dev/null +++ b/vendor/parquet/src/util/test_common/mod.rs @@ -0,0 +1,24 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +pub mod page_util; + +#[cfg(test)] +pub mod file_util; + +#[cfg(test)] +pub mod rand_gen; diff --git a/vendor/parquet/src/util/test_common/page_util.rs b/vendor/parquet/src/util/test_common/page_util.rs new file mode 100644 index 0000000..0b73498 --- /dev/null +++ b/vendor/parquet/src/util/test_common/page_util.rs @@ -0,0 +1,253 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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 bytes::Bytes; + +use crate::basic::Encoding; +use crate::column::page::{Page, PageIterator}; +use crate::column::page::{PageMetadata, PageReader}; +use crate::data_type::DataType; +use crate::encodings::encoding::{Encoder, get_encoder}; +use crate::encodings::levels::LevelEncoder; +use crate::errors::Result; +use crate::schema::types::{ColumnDescPtr, ColumnDescriptor, ColumnPath, Type as SchemaType}; +use std::iter::Peekable; +use std::mem; +use std::sync::Arc; + +pub trait DataPageBuilder { + fn add_rep_levels(&mut self, max_level: i16, rep_levels: &[i16]); + fn add_def_levels(&mut self, max_level: i16, def_levels: &[i16]); + fn add_values(&mut self, encoding: Encoding, values: &[T::T]); + fn add_indices(&mut self, indices: Bytes); + fn consume(self) -> Page; +} + +/// A utility struct for building data pages (v1 or v2). Callers must call the +/// following functions in order to populate and obtain a data page: +/// +/// - add_rep_levels() +/// - add_def_levels() +/// - add_values() for normal data page / add_indices() for dictionary data page +/// - consume() +pub struct DataPageBuilderImpl { + encoding: Option, + num_values: u32, + buffer: Vec, + rep_levels_byte_len: u32, + def_levels_byte_len: u32, + datapage_v2: bool, + type_width: i32, +} + +impl DataPageBuilderImpl { + // `num_values` is the number of non-null values to put in the data page. + // `datapage_v2` flag is used to indicate if the generated data page should use V2 + // format or not. + pub fn new(desc: ColumnDescPtr, num_values: u32, datapage_v2: bool) -> Self { + DataPageBuilderImpl { + encoding: None, + num_values, + buffer: vec![], + rep_levels_byte_len: 0, + def_levels_byte_len: 0, + datapage_v2, + type_width: desc.type_length(), + } + } + + // Adds levels to the buffer and return number of encoded bytes + fn add_levels(&mut self, max_level: i16, levels: &[i16]) -> u32 { + if max_level <= 0 { + return 0; + } + let mut level_encoder = LevelEncoder::v1_streaming(max_level); + level_encoder.put_with_observer(levels, |_, _| {}); + let encoded_levels = level_encoder.consume(); + // Actual encoded bytes (without length offset) + let encoded_bytes = &encoded_levels[mem::size_of::()..]; + if self.datapage_v2 { + // Level encoder always initializes with offset of i32, where it stores + // length of encoded data; for data page v2 we explicitly + // store length, therefore we should skip i32 bytes. + self.buffer.extend_from_slice(encoded_bytes); + } else { + self.buffer.extend_from_slice(encoded_levels.as_slice()); + } + encoded_bytes.len() as u32 + } +} + +impl DataPageBuilder for DataPageBuilderImpl { + fn add_rep_levels(&mut self, max_levels: i16, rep_levels: &[i16]) { + self.num_values = rep_levels.len() as u32; + self.rep_levels_byte_len = self.add_levels(max_levels, rep_levels); + } + + fn add_def_levels(&mut self, max_levels: i16, def_levels: &[i16]) { + self.num_values = def_levels.len() as u32; + self.def_levels_byte_len = self.add_levels(max_levels, def_levels); + } + + fn add_values(&mut self, encoding: Encoding, values: &[T::T]) { + assert!( + self.num_values >= values.len() as u32, + "num_values: {}, values.len(): {}", + self.num_values, + values.len() + ); + // Create test column descriptor. + let desc = { + let ty = SchemaType::primitive_type_builder("t", T::get_physical_type()) + .with_length(self.type_width) + .build() + .unwrap(); + Arc::new(ColumnDescriptor::new( + Arc::new(ty), + 0, + 0, + ColumnPath::new(vec![]), + )) + }; + self.encoding = Some(encoding); + let mut encoder: Box> = + get_encoder::(encoding, &desc).expect("get_encoder() should be OK"); + encoder.put(values).expect("put() should be OK"); + let encoded_values = encoder + .flush_buffer() + .expect("consume_buffer() should be OK"); + self.buffer.extend_from_slice(&encoded_values); + } + + fn add_indices(&mut self, indices: Bytes) { + self.encoding = Some(Encoding::RLE_DICTIONARY); + self.buffer.extend_from_slice(&indices); + } + + fn consume(self) -> Page { + if self.datapage_v2 { + Page::DataPageV2 { + buf: Bytes::from(self.buffer), + num_values: self.num_values, + encoding: self.encoding.unwrap(), + num_nulls: 0, /* set to dummy value - don't need this when reading + * data page */ + num_rows: self.num_values, /* num_rows only needs in skip_records, now we not support skip REPEATED field, + * so we can assume num_values == num_rows */ + def_levels_byte_len: self.def_levels_byte_len, + rep_levels_byte_len: self.rep_levels_byte_len, + is_compressed: false, + statistics: None, // set to None, we do not need statistics for tests + } + } else { + Page::DataPage { + buf: Bytes::from(self.buffer), + num_values: self.num_values, + encoding: self.encoding.unwrap(), + def_level_encoding: Encoding::RLE, + rep_level_encoding: Encoding::RLE, + statistics: None, // set to None, we do not need statistics for tests + } + } + } +} + +/// A utility page reader which stores pages in memory. +pub struct InMemoryPageReader> { + page_iter: Peekable

, +} + +impl> InMemoryPageReader

{ + pub fn new(pages: impl IntoIterator) -> Self { + Self { + page_iter: pages.into_iter().peekable(), + } + } +} + +impl + Send> PageReader for InMemoryPageReader

{ + fn get_next_page(&mut self) -> Result> { + Ok(self.page_iter.next()) + } + + fn peek_next_page(&mut self) -> Result> { + if let Some(x) = self.page_iter.peek() { + match x { + Page::DataPage { num_values, .. } => Ok(Some(PageMetadata { + num_rows: None, + num_levels: Some(*num_values as _), + is_dict: false, + })), + Page::DataPageV2 { + num_rows, + num_values, + .. + } => Ok(Some(PageMetadata { + num_rows: Some(*num_rows as _), + num_levels: Some(*num_values as _), + is_dict: false, + })), + Page::DictionaryPage { .. } => Ok(Some(PageMetadata { + num_rows: None, + num_levels: None, + is_dict: true, + })), + } + } else { + Ok(None) + } + } + + fn skip_next_page(&mut self) -> Result<()> { + self.page_iter.next(); + Ok(()) + } +} + +impl + Send> Iterator for InMemoryPageReader

{ + type Item = Result; + + fn next(&mut self) -> Option { + self.get_next_page().transpose() + } +} + +/// A utility page iterator which stores page readers in memory, used for tests. +#[derive(Clone)] +pub struct InMemoryPageIterator>> { + page_reader_iter: I, +} + +impl>> InMemoryPageIterator { + pub fn new(pages: impl IntoIterator, IntoIter = I>) -> Self { + Self { + page_reader_iter: pages.into_iter(), + } + } +} + +impl>> Iterator for InMemoryPageIterator { + type Item = Result>; + + fn next(&mut self) -> Option { + self.page_reader_iter + .next() + .map(|x| Ok(Box::new(InMemoryPageReader::new(x)) as Box)) + } +} + +impl> + Send> PageIterator for InMemoryPageIterator {} diff --git a/vendor/parquet/src/util/test_common/rand_gen.rs b/vendor/parquet/src/util/test_common/rand_gen.rs new file mode 100644 index 0000000..e4ab2f7 --- /dev/null +++ b/vendor/parquet/src/util/test_common/rand_gen.rs @@ -0,0 +1,221 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::basic::Encoding; +use crate::column::page::Page; +use bytes::Bytes; +use rand::{ + Rng, + distr::{Distribution, StandardUniform, uniform::SampleUniform}, + rng, +}; +use std::collections::VecDeque; + +use crate::data_type::*; +use crate::encodings::encoding::{DictEncoder, Encoder}; +use crate::schema::types::ColumnDescPtr; +use crate::util::{DataPageBuilder, DataPageBuilderImpl}; + +/// Random generator of data type `T` values and sequences. +pub trait RandGen { + fn r#gen(len: i32) -> T::T; + + fn gen_vec(len: i32, total: usize) -> Vec { + let mut result = vec![]; + for _ in 0..total { + result.push(Self::r#gen(len)) + } + result + } +} + +impl RandGen for BoolType { + fn r#gen(_: i32) -> bool { + rng().random::() + } +} + +impl RandGen for Int32Type { + fn r#gen(_: i32) -> i32 { + rng().random::() + } +} + +impl RandGen for Int64Type { + fn r#gen(_: i32) -> i64 { + rng().random::() + } +} + +impl RandGen for Int96Type { + fn r#gen(_: i32) -> Int96 { + let mut rng = rng(); + let mut result = Int96::new(); + result.set_data( + rng.random::(), + rng.random::(), + rng.random::(), + ); + result + } +} + +impl RandGen for FloatType { + fn r#gen(_: i32) -> f32 { + rng().random::() + } +} + +impl RandGen for DoubleType { + fn r#gen(_: i32) -> f64 { + rng().random::() + } +} + +impl RandGen for ByteArrayType { + fn r#gen(_: i32) -> ByteArray { + let mut rng = rng(); + let mut result = ByteArray::new(); + let mut value = vec![]; + let len = rng.random_range(0..128); + for _ in 0..len { + value.push(rng.random_range(0..255)); + } + result.set_data(Bytes::from(value)); + result + } +} + +impl RandGen for FixedLenByteArrayType { + fn r#gen(len: i32) -> FixedLenByteArray { + assert!(len >= 0); + let value = random_bytes(len as usize); + ByteArray::from(value).into() + } +} + +pub fn random_bytes(n: usize) -> Vec { + let mut result = vec![]; + let mut rng = rng(); + for _ in 0..n { + result.push(rng.random_range(0..255)); + } + result +} + +pub fn random_numbers(n: usize) -> Vec +where + StandardUniform: Distribution, +{ + let mut rng = rng(); + StandardUniform.sample_iter(&mut rng).take(n).collect() +} + +pub fn random_numbers_range(n: usize, low: T, high: T, result: &mut Vec) +where + T: PartialOrd + SampleUniform + Copy, +{ + let mut rng = rng(); + for _ in 0..n { + result.push(rng.random_range(low..high)); + } +} + +#[allow(clippy::too_many_arguments)] +pub fn make_pages( + desc: ColumnDescPtr, + encoding: Encoding, + num_pages: usize, + levels_per_page: usize, + min: T::T, + max: T::T, + def_levels: &mut Vec, + rep_levels: &mut Vec, + values: &mut Vec, + pages: &mut VecDeque, + use_v2: bool, +) where + T::T: PartialOrd + SampleUniform + Copy, +{ + let mut num_values = 0; + let max_def_level = desc.max_def_level(); + let max_rep_level = desc.max_rep_level(); + + let mut dict_encoder = DictEncoder::::new(desc.clone()); + + for i in 0..num_pages { + let mut num_values_cur_page = 0; + let level_range = i * levels_per_page..(i + 1) * levels_per_page; + + if max_def_level > 0 { + random_numbers_range(levels_per_page, 0, max_def_level + 1, def_levels); + for dl in &def_levels[level_range.clone()] { + if *dl == max_def_level { + num_values_cur_page += 1; + } + } + } else { + num_values_cur_page = levels_per_page; + } + if max_rep_level > 0 { + random_numbers_range(levels_per_page, 0, max_rep_level + 1, rep_levels); + } + random_numbers_range(num_values_cur_page, min, max, values); + + // Generate the current page + + let mut pb = DataPageBuilderImpl::new(desc.clone(), num_values_cur_page as u32, use_v2); + if max_rep_level > 0 { + pb.add_rep_levels(max_rep_level, &rep_levels[level_range.clone()]); + } + if max_def_level > 0 { + pb.add_def_levels(max_def_level, &def_levels[level_range]); + } + + let value_range = num_values..num_values + num_values_cur_page; + match encoding { + Encoding::PLAIN_DICTIONARY | Encoding::RLE_DICTIONARY => { + let _ = dict_encoder.put(&values[value_range.clone()]); + let indices = dict_encoder + .write_indices() + .expect("write_indices() should be OK"); + pb.add_indices(indices); + } + Encoding::PLAIN => { + pb.add_values::(encoding, &values[value_range]); + } + enc => panic!("Unexpected encoding {enc}"), + } + + let data_page = pb.consume(); + pages.push_back(data_page); + num_values += num_values_cur_page; + } + + if encoding == Encoding::PLAIN_DICTIONARY || encoding == Encoding::RLE_DICTIONARY { + let dict = dict_encoder + .write_dict() + .expect("write_dict() should be OK"); + let dict_page = Page::DictionaryPage { + buf: dict, + num_values: dict_encoder.num_entries() as u32, + encoding: Encoding::RLE_DICTIONARY, + is_sorted: false, + }; + pages.push_front(dict_page); + } +} diff --git a/vendor/parquet/src/util/utf8.rs b/vendor/parquet/src/util/utf8.rs new file mode 100644 index 0000000..2a537b5 --- /dev/null +++ b/vendor/parquet/src/util/utf8.rs @@ -0,0 +1,57 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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_valid_utf8`] validation function +use crate::errors::{ParquetError, Result}; + +/// Check that `val` is a valid UTF-8 sequence. +/// +/// If the `simdutf8` feature is enabled, this function will use +/// SIMD-accelerated validation from the [`simdutf8`] crate. Otherwise, it will use +/// [`std::str::from_utf8`]. +/// +/// # Errors +/// +/// Returns `Err::General` with a message compatible with [`std::str::from_utf8`] on failure. +/// +/// # Example +/// ``` +/// use parquet::utf8::check_valid_utf8; +/// assert!(check_valid_utf8(b"hello").is_ok()); +/// assert!(check_valid_utf8(b"hello \xF0\x9F\x98\x8E").is_ok()); +/// // invalid UTF-8 +/// assert!(check_valid_utf8(b"hello \xF0\x9F\x98").is_err()); +/// ``` +/// +/// [`simdutf8`]: https://crates.io/crates/simdutf8 +#[inline(always)] +pub fn check_valid_utf8(val: &[u8]) -> Result<()> { + #[cfg(feature = "simdutf8")] + match simdutf8::basic::from_utf8(val) { + Ok(_) => Ok(()), + Err(_) => { + // Use simdutf8::compat to return details about the decoding error + let e = simdutf8::compat::from_utf8(val).unwrap_err(); + Err(general_err!("encountered non UTF-8 data: {}", e)) + } + } + #[cfg(not(feature = "simdutf8"))] + match std::str::from_utf8(val) { + Ok(_) => Ok(()), + Err(e) => Err(general_err!("encountered non UTF-8 data: {}", e)), + } +} diff --git a/vendor/parquet/src/variant.rs b/vendor/parquet/src/variant.rs new file mode 100644 index 0000000..55df086 --- /dev/null +++ b/vendor/parquet/src/variant.rs @@ -0,0 +1,299 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! ⚠️ Experimental Support for reading and writing [`Variant`]s to / from Parquet files ⚠️ +//! +//! This is a 🚧 Work In Progress +//! +//! Note: Requires the `variant_experimental` feature of the `parquet` crate to be enabled. +//! +//! # Features +//! * Representation of [`Variant`], and [`VariantArray`] for working with +//! Variant values (see [`parquet_variant`] for more details) +//! * Kernels for working with arrays of Variant values +//! such as conversion between `Variant` and JSON, and shredding/unshredding +//! (see [`parquet_variant_compute`] for more details) +//! +//! # Example: Writing a Parquet file with Variant column +//! ```rust +//! # use parquet::variant::{VariantArray, VariantType, VariantArrayBuilder, VariantBuilderExt}; +//! # use std::sync::Arc; +//! # use arrow_array::{Array, ArrayRef, RecordBatch}; +//! # use arrow_schema::{DataType, Field, Schema}; +//! # use parquet::arrow::ArrowWriter; +//! # fn main() -> Result<(), parquet::errors::ParquetError> { +//! // Use the VariantArrayBuilder to build a VariantArray +//! let mut builder = VariantArrayBuilder::new(3); +//! builder.new_object().with_field("name", "Alice").finish(); // row 1: {"name": "Alice"} +//! builder.append_value("such wow"); // row 2: "such wow" (a string) +//! let array = builder.build(); +//! +//! // Since VariantArray is an ExtensionType, it needs to be converted +//! // to an ArrayRef and Field with the appropriate metadata +//! // before it can be written to a Parquet file +//! let field = array.field("data"); +//! let array = ArrayRef::from(array); +//! // create a RecordBatch with the VariantArray +//! let schema = Schema::new(vec![field]); +//! let batch = RecordBatch::try_new(Arc::new(schema), vec![array])?; +//! +//! // Now you can write the RecordBatch to the Parquet file, as normal +//! let file = std::fs::File::create("variant.parquet")?; +//! let mut writer = ArrowWriter::try_new(file, batch.schema(), None)?; +//! writer.write(&batch)?; +//! writer.close()?; +//! +//! # std::fs::remove_file("variant.parquet")?; +//! # Ok(()) +//! # } +//! ``` +//! +//! # Example: Writing JSON into a Parquet file with Variant column +//! ```rust +//! # use std::sync::Arc; +//! # use arrow_array::{ArrayRef, RecordBatch, StringArray}; +//! # use arrow_schema::Schema; +//! # use parquet::variant::{json_to_variant, VariantArray}; +//! # use parquet::arrow::ArrowWriter; +//! # fn main() -> Result<(), parquet::errors::ParquetError> { +//! // Create an array of JSON strings, simulating a column of JSON data +//! let input_array: ArrayRef = Arc::new(StringArray::from(vec![ +//! Some(r#"{"name": "Alice", "age": 30}"#), +//! Some(r#"{"name": "Bob", "age": 25, "address": {"city": "New York"}}"#), +//! None, +//! Some("{}"), +//! ])); +//! +//! // Convert the JSON strings to a VariantArray +//! let array: VariantArray = json_to_variant(&input_array)?; +//! // create a RecordBatch with the VariantArray +//! let schema = Schema::new(vec![array.field("data")]); +//! let batch = RecordBatch::try_new(Arc::new(schema), vec![ArrayRef::from(array)])?; +//! +//! // write the RecordBatch to a Parquet file as normal +//! let file = std::fs::File::create("variant-json.parquet")?; +//! let mut writer = ArrowWriter::try_new(file, batch.schema(), None)?; +//! writer.write(&batch)?; +//! writer.close()?; +//! # std::fs::remove_file("variant-json.parquet")?; +//! # Ok(()) +//! # } +//! ``` +//! +//! # Example: Reading a Parquet file with Variant column +//! +//! Use the [`VariantType`] extension type to find the Variant column: +//! +//! ``` +//! # use std::sync::Arc; +//! # use std::path::PathBuf; +//! # use arrow_array::{ArrayRef, RecordBatch, RecordBatchReader}; +//! # use parquet::variant::{Variant, VariantArray, VariantType}; +//! # use parquet::arrow::arrow_reader::ArrowReaderBuilder; +//! # fn main() -> Result<(), parquet::errors::ParquetError> { +//! # use arrow_array::StructArray; +//! # fn file_path() -> PathBuf { // return a testing file path +//! # PathBuf::from(arrow::util::test_util::parquet_test_data()) +//! # .join("..") +//! # .join("shredded_variant") +//! # .join("case-075.parquet") +//! # } +//! // Read the Parquet file using standard Arrow Parquet reader. +//! // Note this file has 2 columns: "id", "var", and the "var" column +// // contains a variant that looks like this: +// // "Variant(metadata=VariantMetadata(dict={}), value=Variant(type=STRING, value=iceberg))" +//! let file = std::fs::File::open(file_path())?; +//! let mut reader = ArrowReaderBuilder::try_new(file)?.build()?; +//! +//! // You can check if a column contains a Variant using +//! // the VariantType extension type +//! let schema = reader.schema(); +//! let field = schema.field_with_name("var")?; +//! assert!(field.has_valid_extension_type::()); +//! +//! // The reader will yield RecordBatches with a StructArray +//! // to convert them to VariantArray, use VariantArray::try_new +//! let batch = reader.next().unwrap().unwrap(); +//! +//! let col = batch.column_by_name("var").unwrap(); +//! let var_array = VariantArray::try_new(col)?; +//! assert_eq!(var_array.len(), 1); +//! let var_value: Variant = var_array.value(0); +//! assert_eq!(var_value, Variant::from("iceberg")); // the value in case-075.parquet +//! # Ok(()) +//! # } +//! ``` +pub use parquet_variant::*; +pub use parquet_variant_compute::*; + +#[cfg(test)] +mod tests { + use crate::arrow::ArrowWriter; + use crate::arrow::arrow_reader::ArrowReaderBuilder; + use crate::file::metadata::{ParquetMetaData, ParquetMetaDataReader}; + use crate::file::reader::ChunkReader; + use arrow::util::test_util::parquet_test_data; + use arrow_array::{ArrayRef, RecordBatch}; + use arrow_schema::Schema; + use bytes::Bytes; + use parquet_variant::{Variant, VariantBuilderExt}; + use parquet_variant_compute::{VariantArray, VariantArrayBuilder, VariantType}; + use std::path::PathBuf; + use std::sync::Arc; + + #[test] + fn roundtrip_basic() { + roundtrip(variant_array()); + } + + /// Ensure a file with Variant LogicalType, written by another writer in + /// parquet-testing, can be read as a VariantArray + #[test] + fn read_logical_type() { + // Note: case-075 2 columns ("id", "var") + // The variant looks like this: + // "Variant(metadata=VariantMetadata(dict={}), value=Variant(type=STRING, value=iceberg))" + let batch = read_shredded_variant_test_case("case-075.parquet"); + + assert_variant_metadata(&batch, "var"); + let var_column = batch.column_by_name("var").expect("expected var column"); + let var_array = + VariantArray::try_new(&var_column).expect("expected var column to be a VariantArray"); + + // verify the value + assert_eq!(var_array.len(), 1); + assert!(var_array.is_valid(0)); + let var_value = var_array.value(0); + assert_eq!(var_value, Variant::from("iceberg")); + } + + /// Writes a variant to a parquet file and ensures the parquet logical type + /// annotation is correct + #[test] + fn write_logical_type() { + let array = variant_array(); + let batch = variant_array_to_batch(array); + let buffer = write_to_buffer(&batch); + + // read the parquet file's metadata and verify the logical type + let metadata = read_metadata(&Bytes::from(buffer)); + let schema = metadata.file_metadata().schema_descr(); + let fields = schema.root_schema().get_fields(); + assert_eq!(fields.len(), 1); + let field = &fields[0]; + assert_eq!(field.name(), "data"); + // data should have been written with the Variant logical type + assert_eq!( + field.get_basic_info().logical_type_ref(), + Some(&crate::basic::LogicalType::variant(None)) + ); + } + + /// Return a VariantArray with 3 rows: + /// + /// 1. `{"name": "Alice"}` + /// 2. `"such wow"` (a string) + /// 3. `null` + fn variant_array() -> VariantArray { + let mut builder = VariantArrayBuilder::new(3); + // row 1: {"name": "Alice"} + builder.new_object().with_field("name", "Alice").finish(); + // row 2: "such wow" (a string) + builder.append_value("such wow"); + // row 3: null + builder.append_null(); + builder.build() + } + + /// Writes a VariantArray to a parquet file and reads it back, verifying that + /// the data is the same + fn roundtrip(array: VariantArray) { + let source_batch = variant_array_to_batch(array); + assert_variant_metadata(&source_batch, "data"); + + let buffer = write_to_buffer(&source_batch); + let result_batch = read_to_batch(Bytes::from(buffer)); + assert_variant_metadata(&result_batch, "data"); + assert_eq!(result_batch, source_batch); // NB this also checks the schemas + } + + /// creates a RecordBatch with a single column "data" from a VariantArray, + fn variant_array_to_batch(array: VariantArray) -> RecordBatch { + let field = array.field("data"); + let schema = Schema::new(vec![field]); + RecordBatch::try_new(Arc::new(schema), vec![ArrayRef::from(array)]).unwrap() + } + + /// writes a RecordBatch to memory buffer and returns the buffer + fn write_to_buffer(batch: &RecordBatch) -> Vec { + let mut buffer = vec![]; + let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap(); + writer.write(batch).unwrap(); + writer.close().unwrap(); + buffer + } + + /// Reads the Parquet metadata + fn read_metadata(input: &T) -> ParquetMetaData { + let mut reader = ParquetMetaDataReader::new(); + reader.try_parse(input).unwrap(); + reader.finish().unwrap() + } + + /// Reads a RecordBatch from a reader (e.g. Vec or File) + fn read_to_batch(reader: T) -> RecordBatch { + let reader = ArrowReaderBuilder::try_new(reader) + .unwrap() + .build() + .unwrap(); + let mut batches: Vec = reader.collect::, _>>().unwrap(); + assert_eq!(batches.len(), 1); + batches.swap_remove(0) + } + + /// Verifies the variant metadata is present in the schema for the specified + /// field name. + fn assert_variant_metadata(batch: &RecordBatch, field_name: &str) { + let schema = batch.schema(); + let field = schema + .field_with_name(field_name) + .expect("could not find expected field"); + + // explicitly check the metadata so it is clear in the tests what the + // names are + let metadata_value = field + .metadata() + .get("ARROW:extension:name") + .expect("metadata does not exist"); + + assert_eq!(metadata_value, "arrow.parquet.variant"); + + // verify that `VariantType` also correctly finds the metadata + assert!(field.has_valid_extension_type::()); + } + + /// Read the specified test case filename from parquet-testing + /// See parquet-testing/shredded_variant/cases.json for more details + fn read_shredded_variant_test_case(name: &str) -> RecordBatch { + let case_file = PathBuf::from(parquet_test_data()) + .join("..") // go up from data/ to parquet-testing/ + .join("shredded_variant") + .join(name); + let case_file = std::fs::File::open(case_file).unwrap(); + read_to_batch(case_file) + } +} diff --git a/vendor/parquet/tests/arrow_reader/bad_data.rs b/vendor/parquet/tests/arrow_reader/bad_data.rs new file mode 100644 index 0000000..b9cce1a --- /dev/null +++ b/vendor/parquet/tests/arrow_reader/bad_data.rs @@ -0,0 +1,225 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Tests that reading invalid parquet files returns an error + +use arrow::util::test_util::parquet_test_data; +use bytes::Bytes; +use parquet::arrow::arrow_reader::ArrowReaderBuilder; +use parquet::errors::ParquetError; +use parquet::file::metadata::ParquetMetaDataReader; +use std::collections::HashSet; +use std::path::PathBuf; + +static KNOWN_FILES: &[&str] = &[ + "PARQUET-1481.parquet", + "ARROW-GH-41317.parquet", + "ARROW-GH-41321.parquet", + "ARROW-GH-43605.parquet", + "ARROW-RS-GH-6229-DICTHEADER.parquet", + "ARROW-RS-GH-6229-LEVELS.parquet", + "ARROW-GH-45185.parquet", + "ARROW-GH-47662.parquet", + "README.md", + "variants", +]; + +/// Returns the path to 'parquet-testing/bad_data' +fn bad_data_dir() -> PathBuf { + // points to parquet-testing/data + let parquet_testing_data = parquet_test_data(); + PathBuf::from(parquet_testing_data) + .parent() + .expect("was in parquet-testing/data") + .join("bad_data") +} + +#[test] +// Ensure that if we add a new test the files are added to the tests. +fn test_invalid_files() { + let known_files: HashSet<_> = KNOWN_FILES.iter().cloned().collect(); + let mut seen_files = HashSet::new(); + + let files = std::fs::read_dir(bad_data_dir()).unwrap(); + + for file in files { + let file_name = file + .unwrap() + .path() + .file_name() + .unwrap() + .to_str() + .unwrap() + .to_string(); + + // If you see this error, please add a test for the new file following the model below + assert!( + known_files.contains(file_name.as_str()), + "Found new file in bad_data, please add test: {file_name}" + ); + seen_files.insert(file_name); + } + for expected_file in known_files { + assert!( + seen_files.contains(expected_file), + "Expected file not found in bad_data directory: {expected_file}" + ); + } +} + +#[test] +fn test_parquet_1481() { + let err = read_file("PARQUET-1481.parquet").unwrap_err(); + assert_eq!(err.to_string(), "Parquet error: Unexpected Type -7"); +} + +#[test] +fn test_arrow_gh_41321() { + let err = read_file("ARROW-GH-41321.parquet").unwrap_err(); + assert_eq!( + err.to_string(), + "External: Parquet argument error: Parquet error: Invalid or corrupted RLE bit width 254. Max allowed is 32" + ); +} + +#[test] +fn test_arrow_gh_41317() { + let err = read_file("ARROW-GH-41317.parquet").unwrap_err(); + assert_eq!( + err.to_string(), + "Parquet error: Expected list element type of I32 but got I16" + ); +} + +#[test] +fn test_arrow_rs_gh_6229_dict_header() { + let err = read_file("ARROW-RS-GH-6229-DICTHEADER.parquet").unwrap_err(); + assert_eq!( + err.to_string(), + "External: Parquet argument error: Parquet error: Integer overflow: out of range integral type conversion attempted" + ); +} + +#[test] +#[cfg(feature = "snap")] +fn test_arrow_rs_gh_6229_dict_levels() { + let err = read_file("ARROW-RS-GH-6229-LEVELS.parquet").unwrap_err(); + assert_eq!( + err.to_string(), + "External: Parquet argument error: Parquet error: Insufficient repetition levels read from column" + ); +} + +#[test] +#[cfg(feature = "snap")] +fn test_arrow_rs_gh_45185_dict_levels() { + let err = read_file("ARROW-GH-45185.parquet").unwrap_err(); + assert_eq!( + err.to_string(), + "External: Parquet argument error: Parquet error: first repetition level of batch must be 0" + ); +} + +#[test] +fn test_arrow_gh_47662() { + let err = read_file("ARROW-GH-47662.parquet").unwrap_err(); + assert_eq!( + err.to_string(), + "External: Parquet argument error: Parquet error: insufficient values read from column - expected: 100, got: 91" + ); +} + +/// Reads the file and tries to return the total row count +/// Returns an error if the file is invalid +fn read_file(name: &str) -> Result { + let path = bad_data_dir().join(name); + println!("Reading file: {path:?}"); + + let file = std::fs::File::open(&path).unwrap(); + let reader = ArrowReaderBuilder::try_new(file)?.build()?; + + let mut num_rows = 0; + for batch in reader { + let batch = batch?; + num_rows += batch.num_rows(); + } + Ok(num_rows) +} + +#[test] +fn non_standard_delta_blocks() { + let file = Bytes::from_static(include_bytes!("bigdelta.parquet")); + use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; + + let selectors = vec![RowSelector::skip(1000), RowSelector::select(5)]; + + let selection: RowSelection = selectors.into(); + let reader = ArrowReaderBuilder::try_new(file) + .unwrap() + .with_row_selection(selection) + .build() + .unwrap(); + + if let Some(maybe_batch) = reader.into_iter().next() { + // TODO: uncomment if we ever allow skipping miniblocks > 64 elements + //let batch = maybe_batch.expect("skip should succeed"); + //assert_eq!(batch.num_rows(), 5); + assert!( + maybe_batch + .unwrap_err() + .to_string() + .contains("cannot skip miniblock of size 128") + ); + } +} + +#[test] +fn skip_unknown_types() { + // test file contains a FileMetaData with unknown fields with + // types not currently used by Parquet (uuid, set, map). The + // parser should be able to skip these unknown fields without + // erroring. + let file = Bytes::from_static(include_bytes!("new_types.bin")); + ParquetMetaDataReader::decode_metadata(&file).unwrap(); +} + +#[cfg(feature = "async")] +#[tokio::test] +async fn bad_metadata_err() { + use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader}; + + let metadata_buffer = Bytes::from_static(include_bytes!("bad_raw_metadata.bin")); + + let metadata_length = metadata_buffer.len() as u64; + + let mut reader = std::io::Cursor::new(&metadata_buffer); + let mut loader = ParquetMetaDataReader::new(); + loader.try_load(&mut reader, metadata_length).await.unwrap(); + loader = loader.with_page_index_policy(PageIndexPolicy::Skip); + loader.load_page_index(&mut reader).await.unwrap(); + + loader = loader.with_offset_index_policy(PageIndexPolicy::Required); + loader.load_page_index(&mut reader).await.unwrap(); + + loader = loader.with_column_index_policy(PageIndexPolicy::Required); + let err = loader.load_page_index(&mut reader).await.unwrap_err(); + + assert_eq!( + err.to_string(), + "Parquet error: error converting value, expected 4 bytes got 0" + ); +} diff --git a/vendor/parquet/tests/arrow_reader/bad_raw_metadata.bin b/vendor/parquet/tests/arrow_reader/bad_raw_metadata.bin new file mode 100644 index 0000000..47f9aa1 Binary files /dev/null and b/vendor/parquet/tests/arrow_reader/bad_raw_metadata.bin differ diff --git a/vendor/parquet/tests/arrow_reader/bigdelta.parquet b/vendor/parquet/tests/arrow_reader/bigdelta.parquet new file mode 100644 index 0000000..0fb270a Binary files /dev/null and b/vendor/parquet/tests/arrow_reader/bigdelta.parquet differ diff --git a/vendor/parquet/tests/arrow_reader/bloom_filter/async.rs b/vendor/parquet/tests/arrow_reader/bloom_filter/async.rs new file mode 100644 index 0000000..e230b33 --- /dev/null +++ b/vendor/parquet/tests/arrow_reader/bloom_filter/async.rs @@ -0,0 +1,82 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::io::TestReader; +use bytes::Bytes; +use futures::TryStreamExt; +use parquet::{ + arrow::{ArrowWriter, ParquetRecordBatchStreamBuilder}, + file::properties::WriterProperties, +}; + +#[tokio::test] +async fn test_get_row_group_column_bloom_filter_without_length() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/data_index_bloom_encoding_stats.parquet"); + let data = Bytes::from(std::fs::read(path).unwrap()); + test_get_row_group_column_bloom_filter(data, false).await; +} + +#[tokio::test] +async fn test_get_row_group_column_bloom_filter_with_length() { + // convert to new parquet file with bloom_filter_length + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/data_index_bloom_encoding_stats.parquet"); + let data = Bytes::from(std::fs::read(path).unwrap()); + let async_reader = TestReader::new(data.clone()); + let builder = ParquetRecordBatchStreamBuilder::new(async_reader) + .await + .unwrap(); + let schema = builder.schema().clone(); + let stream = builder.build().unwrap(); + let batches = stream.try_collect::>().await.unwrap(); + + let mut parquet_data = Vec::new(); + let props = WriterProperties::builder() + .set_bloom_filter_enabled(true) + .build(); + let mut writer = ArrowWriter::try_new(&mut parquet_data, schema, Some(props)).unwrap(); + for batch in batches { + writer.write(&batch).unwrap(); + } + writer.close().unwrap(); + + // test the new parquet file + test_get_row_group_column_bloom_filter(parquet_data.into(), true).await; +} + +async fn test_get_row_group_column_bloom_filter(data: Bytes, with_length: bool) { + let async_reader = TestReader::new(data.clone()); + + let mut builder = ParquetRecordBatchStreamBuilder::new(async_reader) + .await + .unwrap(); + + let metadata = builder.metadata(); + assert_eq!(metadata.num_row_groups(), 1); + let row_group = metadata.row_group(0); + let column = row_group.column(0); + assert_eq!(column.bloom_filter_length().is_some(), with_length); + + let sbbf = builder + .get_row_group_column_bloom_filter(0, 0) + .await + .unwrap() + .unwrap(); + assert!(sbbf.check(&"Hello")); + assert!(!sbbf.check(&"Hello_Not_Exists")); +} diff --git a/vendor/parquet/tests/arrow_reader/bloom_filter/mod.rs b/vendor/parquet/tests/arrow_reader/bloom_filter/mod.rs new file mode 100644 index 0000000..a09fd4a --- /dev/null +++ b/vendor/parquet/tests/arrow_reader/bloom_filter/mod.rs @@ -0,0 +1,20 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#[cfg(feature = "async")] +mod r#async; +mod sync; diff --git a/vendor/parquet/tests/arrow_reader/bloom_filter/sync.rs b/vendor/parquet/tests/arrow_reader/bloom_filter/sync.rs new file mode 100644 index 0000000..90d0cdc --- /dev/null +++ b/vendor/parquet/tests/arrow_reader/bloom_filter/sync.rs @@ -0,0 +1,74 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fs::File; + +use bytes::Bytes; +use parquet::{ + arrow::{ArrowWriter, arrow_reader::ParquetRecordBatchReaderBuilder}, + file::properties::WriterProperties, +}; + +#[test] +fn test_get_row_group_column_bloom_filter_with_length() { + // convert to new parquet file with bloom_filter_length + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/data_index_bloom_encoding_stats.parquet"); + let file = File::open(path).unwrap(); + let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + let schema = builder.schema().clone(); + let reader = builder.build().unwrap(); + + let mut parquet_data = Vec::new(); + let props = WriterProperties::builder() + .set_bloom_filter_enabled(true) + .build(); + let mut writer = ArrowWriter::try_new(&mut parquet_data, schema, Some(props)).unwrap(); + for batch in reader { + let batch = batch.unwrap(); + writer.write(&batch).unwrap(); + } + writer.close().unwrap(); + + // test the new parquet file + test_get_row_group_column_bloom_filter(parquet_data.into(), true); +} + +#[test] +fn test_get_row_group_column_bloom_filter_without_length() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/data_index_bloom_encoding_stats.parquet"); + let data = Bytes::from(std::fs::read(path).unwrap()); + test_get_row_group_column_bloom_filter(data, false); +} + +fn test_get_row_group_column_bloom_filter(data: Bytes, with_length: bool) { + let builder = ParquetRecordBatchReaderBuilder::try_new(data.clone()).unwrap(); + + let metadata = builder.metadata(); + assert_eq!(metadata.num_row_groups(), 1); + let row_group = metadata.row_group(0); + let column = row_group.column(0); + assert_eq!(column.bloom_filter_length().is_some(), with_length); + + let sbbf = builder + .get_row_group_column_bloom_filter(0, 0) + .unwrap() + .unwrap(); + assert!(sbbf.check(&"Hello")); + assert!(!sbbf.check(&"Hello_Not_Exists")); +} diff --git a/vendor/parquet/tests/arrow_reader/checksum.rs b/vendor/parquet/tests/arrow_reader/checksum.rs new file mode 100644 index 0000000..1a37289 --- /dev/null +++ b/vendor/parquet/tests/arrow_reader/checksum.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. + +//! This file contains an end to end test for verifying checksums when reading parquet files. + +use std::path::PathBuf; + +use arrow::util::test_util::parquet_test_data; +use parquet::arrow::arrow_reader::ArrowReaderBuilder; + +#[test] +fn test_datapage_v1_corrupt_checksum() { + let errors = read_file_batch_errors("datapage_v1-corrupt-checksum.parquet"); + assert_eq!(errors, [ + Err("Parquet argument error: Parquet error: Page CRC checksum mismatch".to_string()), + Ok(()), + Ok(()), + Err("Parquet argument error: Parquet error: Page CRC checksum mismatch".to_string()), + Err("Parquet argument error: Parquet error: Not all children array length are the same!".to_string()) + ]); +} + +#[test] +fn test_datapage_v1_uncompressed_checksum() { + let errors = read_file_batch_errors("datapage_v1-uncompressed-checksum.parquet"); + assert_eq!(errors, [Ok(()), Ok(()), Ok(()), Ok(()), Ok(())]); +} + +#[test] +#[cfg(feature = "snap")] +fn test_datapage_v1_snappy_compressed_checksum() { + let errors = read_file_batch_errors("datapage_v1-snappy-compressed-checksum.parquet"); + assert_eq!(errors, [Ok(()), Ok(()), Ok(()), Ok(()), Ok(())]); +} + +#[test] +fn test_plain_dict_uncompressed_checksum() { + let errors = read_file_batch_errors("plain-dict-uncompressed-checksum.parquet"); + assert_eq!(errors, [Ok(())]); +} +#[test] +#[cfg(feature = "snap")] +fn test_rle_dict_snappy_checksum() { + let errors = read_file_batch_errors("rle-dict-snappy-checksum.parquet"); + assert_eq!(errors, [Ok(())]); +} + +/// Reads a file and returns a vector with one element per record batch. +/// The record batch data is replaced with () and errors are stringified. +fn read_file_batch_errors(name: &str) -> Vec> { + let path = PathBuf::from(parquet_test_data()).join(name); + println!("Reading file: {path:?}"); + let file = std::fs::File::open(&path).unwrap(); + let reader = ArrowReaderBuilder::try_new(file).unwrap().build().unwrap(); + reader + .map(|x| match x { + Ok(_) => Ok(()), + Err(e) => Err(e.to_string()), + }) + .collect() +} diff --git a/vendor/parquet/tests/arrow_reader/int96_stats_roundtrip.rs b/vendor/parquet/tests/arrow_reader/int96_stats_roundtrip.rs new file mode 100644 index 0000000..d6ba8d4 --- /dev/null +++ b/vendor/parquet/tests/arrow_reader/int96_stats_roundtrip.rs @@ -0,0 +1,151 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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 chrono::{DateTime, NaiveDateTime, Utc}; +use parquet::basic::Type; +use parquet::data_type::{Int96, Int96Type}; +use parquet::file::properties::{EnabledStatistics, WriterProperties}; +use parquet::file::reader::{FileReader, SerializedFileReader}; +use parquet::file::statistics::Statistics; +use parquet::file::writer::SerializedFileWriter; +use parquet::schema::parser::parse_message_type; +use rand::seq::SliceRandom; +use std::fs::File; +use std::sync::Arc; +use tempfile::Builder; + +fn datetime_to_int96(dt: &str) -> Int96 { + let naive = NaiveDateTime::parse_from_str(dt, "%Y-%m-%d %H:%M:%S%.f").unwrap(); + let datetime: DateTime = DateTime::from_naive_utc_and_offset(naive, Utc); + let nanos = datetime.timestamp_nanos_opt().unwrap(); + let mut int96 = Int96::new(); + const JULIAN_DAY_OF_EPOCH: i64 = 2_440_588; + const NANOSECONDS_IN_DAY: i64 = 86_400_000_000_000; + let days = nanos / NANOSECONDS_IN_DAY; + let remaining_nanos = nanos % NANOSECONDS_IN_DAY; + let julian_day = (days + JULIAN_DAY_OF_EPOCH) as i32; + let julian_day_u32 = julian_day as u32; + let nanos_low = (remaining_nanos & 0xFFFFFFFF) as u32; + let nanos_high = ((remaining_nanos >> 32) & 0xFFFFFFFF) as u32; + int96.set_data(nanos_low, nanos_high, julian_day_u32); + int96 +} + +fn verify_ordering(data: Vec) { + // Create a temporary file + let tmp = Builder::new() + .prefix("test_int96_stats") + .tempfile() + .unwrap(); + let file_path = tmp.path().to_owned(); + + // Create schema with INT96 field + let message_type = " + message test { + REQUIRED INT96 timestamp; + } + "; + let schema = parse_message_type(message_type).unwrap(); + + // Configure writer properties to enable statistics + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + + let expected_min = data[0]; + let expected_max = data[data.len() - 1]; + + { + let file = File::create(&file_path).unwrap(); + let mut writer = SerializedFileWriter::new(file, schema.into(), Arc::new(props)).unwrap(); + let mut row_group = writer.next_row_group().unwrap(); + let mut col_writer = row_group.next_column().unwrap().unwrap(); + + { + let writer = col_writer.typed::(); + let mut shuffled_data = data.clone(); + shuffled_data.shuffle(&mut rand::rng()); + writer.write_batch(&shuffled_data, None, None).unwrap(); + } + col_writer.close().unwrap(); + row_group.close().unwrap(); + writer.close().unwrap(); + } + + let file = File::open(&file_path).unwrap(); + let reader = SerializedFileReader::new(file).unwrap(); + let metadata = reader.metadata(); + let row_group = metadata.row_group(0); + let column = row_group.column(0); + + let stats = column.statistics().unwrap(); + assert_eq!(stats.physical_type(), Type::INT96); + + if let Statistics::Int96(stats) = stats { + let min = stats.min_opt().unwrap(); + let max = stats.max_opt().unwrap(); + + assert_eq!( + *min, expected_min, + "Min value should be {expected_min} but was {min}" + ); + assert_eq!( + *max, expected_max, + "Max value should be {expected_max} but was {max}" + ); + assert_eq!(stats.null_count_opt(), Some(0)); + } else { + panic!("Expected Int96 statistics"); + } +} + +#[test] +fn test_multiple_dates() { + let data = vec![ + datetime_to_int96("2020-01-01 00:00:00.000"), + datetime_to_int96("2020-02-29 23:59:59.000"), + datetime_to_int96("2020-12-31 23:59:59.000"), + datetime_to_int96("2021-01-01 00:00:00.000"), + datetime_to_int96("2023-06-15 12:30:45.000"), + datetime_to_int96("2024-02-29 15:45:30.000"), + datetime_to_int96("2024-12-25 07:00:00.000"), + datetime_to_int96("2025-01-01 00:00:00.000"), + datetime_to_int96("2025-07-04 20:00:00.000"), + datetime_to_int96("2025-12-31 23:59:59.000"), + ]; + verify_ordering(data); +} + +#[test] +fn test_same_day_different_time() { + let data = vec![ + datetime_to_int96("2020-01-01 00:01:00.000"), + datetime_to_int96("2020-01-01 00:02:00.000"), + datetime_to_int96("2020-01-01 00:03:00.000"), + ]; + verify_ordering(data); +} + +#[test] +fn test_increasing_day_decreasing_time() { + let data = vec![ + datetime_to_int96("2020-01-01 12:00:00.000"), + datetime_to_int96("2020-02-01 11:00:00.000"), + datetime_to_int96("2020-03-01 10:00:00.000"), + ]; + verify_ordering(data); +} diff --git a/vendor/parquet/tests/arrow_reader/invalid_utf8.rs b/vendor/parquet/tests/arrow_reader/invalid_utf8.rs new file mode 100644 index 0000000..1124737 --- /dev/null +++ b/vendor/parquet/tests/arrow_reader/invalid_utf8.rs @@ -0,0 +1,265 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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::sync::Arc; + +use arrow::array::AsArray; +use arrow_array::{ + Array, ArrayRef, GenericBinaryArray, GenericStringArray, OffsetSizeTrait, RecordBatch, + StringViewArray, +}; +use arrow_schema::{ArrowError, DataType}; +use bytes::Bytes; +use parquet::{ + arrow::{ + ArrowWriter, + arrow_reader::{ArrowReaderBuilder, ParquetRecordBatchReader}, + }, + basic::Encoding, + file::properties::WriterProperties, +}; + +#[test] +fn test_invalid_utf8() { + // a parquet file with 1 column with invalid utf8 + let data = vec![ + 80, 65, 82, 49, 21, 6, 21, 22, 21, 22, 92, 21, 2, 21, 0, 21, 2, 21, 0, 21, 4, 21, 0, 18, + 28, 54, 0, 40, 5, 104, 101, 255, 108, 111, 24, 5, 104, 101, 255, 108, 111, 0, 0, 0, 3, 1, + 5, 0, 0, 0, 104, 101, 255, 108, 111, 38, 110, 28, 21, 12, 25, 37, 6, 0, 25, 24, 2, 99, 49, + 21, 0, 22, 2, 22, 102, 22, 102, 38, 8, 60, 54, 0, 40, 5, 104, 101, 255, 108, 111, 24, 5, + 104, 101, 255, 108, 111, 0, 0, 0, 21, 4, 25, 44, 72, 4, 114, 111, 111, 116, 21, 2, 0, 21, + 12, 37, 2, 24, 2, 99, 49, 37, 0, 76, 28, 0, 0, 0, 22, 2, 25, 28, 25, 28, 38, 110, 28, 21, + 12, 25, 37, 6, 0, 25, 24, 2, 99, 49, 21, 0, 22, 2, 22, 102, 22, 102, 38, 8, 60, 54, 0, 40, + 5, 104, 101, 255, 108, 111, 24, 5, 104, 101, 255, 108, 111, 0, 0, 0, 22, 102, 22, 2, 0, 40, + 44, 65, 114, 114, 111, 119, 50, 32, 45, 32, 78, 97, 116, 105, 118, 101, 32, 82, 117, 115, + 116, 32, 105, 109, 112, 108, 101, 109, 101, 110, 116, 97, 116, 105, 111, 110, 32, 111, 102, + 32, 65, 114, 114, 111, 119, 0, 130, 0, 0, 0, 80, 65, 82, 49, + ]; + + let file = Bytes::from(data); + let mut record_batch_reader = ParquetRecordBatchReader::try_new(file, 10).unwrap(); + + let error = record_batch_reader.next().unwrap().unwrap_err(); + + assert!( + error.to_string().contains("invalid utf-8 sequence"), + "{}", + error + ); +} + +#[test] +fn test_invalid_utf8_string_array() { + test_invalid_utf8_string_array_inner::(); +} + +#[test] +fn test_invalid_utf8_large_string_array() { + test_invalid_utf8_string_array_inner::(); +} + +fn test_invalid_utf8_string_array_inner() { + let cases = [ + invalid_utf8_first_char::(), + invalid_utf8_first_char_long_strings::(), + invalid_utf8_later_char::(), + invalid_utf8_later_char_long_strings::(), + invalid_utf8_later_char_really_long_strings::(), + invalid_utf8_later_char_really_long_strings2::(), + ]; + for array in &cases { + for encoding in STRING_ENCODINGS { + // data is not valid utf8 we can not construct a correct StringArray + // safely, so purposely create an invalid StringArray + let array = unsafe { + GenericStringArray::::new_unchecked( + array.offsets().clone(), + array.values().clone(), + array.nulls().cloned(), + ) + }; + let data_type = array.data_type().clone(); + let data = write_to_parquet_with_encoding(Arc::new(array), *encoding); + let err = read_from_parquet(data).unwrap_err(); + let expected_err = "Parquet argument error: Parquet error: encountered non UTF-8 data"; + assert!( + err.to_string().contains(expected_err), + "data type: {data_type}, expected: {expected_err}, got: {err}" + ); + } + } +} + +#[test] +fn test_invalid_utf8_string_view_array() { + let cases = [ + invalid_utf8_first_char::(), + invalid_utf8_first_char_long_strings::(), + invalid_utf8_later_char::(), + invalid_utf8_later_char_long_strings::(), + invalid_utf8_later_char_really_long_strings::(), + invalid_utf8_later_char_really_long_strings2::(), + ]; + + for encoding in STRING_ENCODINGS { + for array in &cases { + let array = arrow_cast::cast(&array, &DataType::BinaryView).unwrap(); + let array = array.as_binary_view(); + + // data is not valid utf8 we can not construct a correct StringArray + // safely, so purposely create an invalid StringViewArray + let array = unsafe { + StringViewArray::new_unchecked( + array.views().clone(), + array.data_buffers().to_vec(), + array.nulls().cloned(), + ) + }; + + let data_type = array.data_type().clone(); + let data = write_to_parquet_with_encoding(Arc::new(array), *encoding); + let err = read_from_parquet(data).unwrap_err(); + let expected_err = "Parquet argument error: Parquet error: encountered non UTF-8 data"; + assert!( + err.to_string().contains(expected_err), + "data type: {data_type}, expected: {expected_err}, got: {err}" + ); + } + } +} + +/// Encodings suitable for string data +const STRING_ENCODINGS: &[Option] = &[ + None, + Some(Encoding::PLAIN), + Some(Encoding::DELTA_LENGTH_BYTE_ARRAY), + Some(Encoding::DELTA_BYTE_ARRAY), +]; + +/// Invalid Utf-8 sequence in the first character +/// +const INVALID_UTF8_FIRST_CHAR: &[u8] = &[0xa0, 0xa1, 0x20, 0x20]; + +/// Invalid Utf=8 sequence in NOT the first character +/// +const INVALID_UTF8_LATER_CHAR: &[u8] = &[0x20, 0x20, 0x20, 0xa0, 0xa1, 0x20, 0x20]; + +/// returns a BinaryArray with invalid UTF8 data in the first character +fn invalid_utf8_first_char() -> GenericBinaryArray { + let valid: &[u8] = b" "; + let invalid = INVALID_UTF8_FIRST_CHAR; + GenericBinaryArray::::from_iter(vec![None, Some(valid), None, Some(invalid)]) +} + +/// Returns a BinaryArray with invalid UTF8 data in the first character of a +/// string larger than 12 bytes which is handled specially when reading +/// `ByteViewArray`s +fn invalid_utf8_first_char_long_strings() -> GenericBinaryArray { + let valid: &[u8] = b" "; + let mut invalid = vec![]; + invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes"); + invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR); + GenericBinaryArray::::from_iter(vec![None, Some(valid), None, Some(&invalid)]) +} + +/// returns a BinaryArray with invalid UTF8 data in a character other than +/// the first (this is checked in a special codepath) +fn invalid_utf8_later_char() -> GenericBinaryArray { + let valid: &[u8] = b" "; + let invalid: &[u8] = INVALID_UTF8_LATER_CHAR; + GenericBinaryArray::::from_iter(vec![None, Some(valid), None, Some(invalid)]) +} + +/// returns a BinaryArray with invalid UTF8 data in a character other than +/// the first in a string larger than 12 bytes which is handled specially +/// when reading `ByteViewArray`s (this is checked in a special codepath) +fn invalid_utf8_later_char_long_strings() -> GenericBinaryArray { + let valid: &[u8] = b" "; + let mut invalid = vec![]; + invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes"); + invalid.extend_from_slice(INVALID_UTF8_LATER_CHAR); + GenericBinaryArray::::from_iter(vec![None, Some(valid), None, Some(&invalid)]) +} + +/// returns a BinaryArray with invalid UTF8 data in a character other than +/// the first in a string larger than 128 bytes which is handled specially +/// when reading `ByteViewArray`s (this is checked in a special codepath) +fn invalid_utf8_later_char_really_long_strings() -> GenericBinaryArray { + let valid: &[u8] = b" "; + let mut invalid = vec![]; + for _ in 0..10 { + // each instance is 38 bytes + invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes"); + } + invalid.extend_from_slice(INVALID_UTF8_LATER_CHAR); + GenericBinaryArray::::from_iter(vec![None, Some(valid), None, Some(&invalid)]) +} + +/// returns a BinaryArray with small invalid UTF8 data followed by a large +/// invalid UTF8 data in a character other than the first in a string larger +fn invalid_utf8_later_char_really_long_strings2() -> GenericBinaryArray { + let valid: &[u8] = b" "; + let mut valid_long = vec![]; + for _ in 0..10 { + // each instance is 38 bytes + valid_long.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes"); + } + let invalid = INVALID_UTF8_LATER_CHAR; + GenericBinaryArray::::from_iter(vec![ + None, + Some(valid), + Some(invalid), + None, + Some(&valid_long), + Some(valid), + ]) +} + +/// writes the array into a single column parquet file with the specified +/// encoding. +/// +/// If no encoding is specified, use default (dictionary) encoding +fn write_to_parquet_with_encoding(array: ArrayRef, encoding: Option) -> Vec { + let batch = RecordBatch::try_from_iter(vec![("c", array)]).unwrap(); + let mut data = vec![]; + let schema = batch.schema(); + let props = encoding.map(|encoding| { + WriterProperties::builder() + // must disable dictionary encoding to actually use encoding + .set_dictionary_enabled(false) + .set_encoding(encoding) + .build() + }); + + { + let mut writer = ArrowWriter::try_new(&mut data, schema, props).unwrap(); + writer.write(&batch).unwrap(); + writer.flush().unwrap(); + writer.close().unwrap(); + }; + data +} + +/// read the parquet file into a record batch +fn read_from_parquet(data: Vec) -> Result, ArrowError> { + let reader = ArrowReaderBuilder::try_new(bytes::Bytes::from(data)) + .unwrap() + .build() + .unwrap(); + + reader.collect() +} diff --git a/vendor/parquet/tests/arrow_reader/io/async_reader.rs b/vendor/parquet/tests/arrow_reader/io/async_reader.rs new file mode 100644 index 0000000..db06dda --- /dev/null +++ b/vendor/parquet/tests/arrow_reader/io/async_reader.rs @@ -0,0 +1,431 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Tests for the async reader ([`ParquetRecordBatchStreamBuilder`]) + +use crate::io::{ + LogEntry, OperationLog, TestParquetFile, filter_a_175_b_625, filter_b_575_625, filter_b_false, + test_file, test_options, +}; +use bytes::Bytes; +use futures::future::BoxFuture; +use futures::{FutureExt, StreamExt}; +use parquet::arrow::arrow_reader::{ArrowReaderOptions, RowSelection, RowSelector}; +use parquet::arrow::async_reader::AsyncFileReader; +use parquet::arrow::{ParquetRecordBatchStreamBuilder, ProjectionMask}; +use parquet::errors::Result; +use parquet::file::metadata::PageIndexPolicy; +use parquet::file::metadata::ParquetMetaData; +use std::ops::Range; +use std::sync::Arc; + +#[tokio::test] +async fn test_read_entire_file() { + // read entire file without any filtering or projection + let test_file = test_file(); + let builder = async_builder(&test_file, test_options()).await; + insta::assert_debug_snapshot!(run( + &test_file, + builder).await, @r#" + [ + "Get Provided Metadata", + "Event: Builder Configured", + "Event: Reader Built", + "Read Multi:", + " Row Group 0, column 'a': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + " Row Group 0, column 'b': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + " Row Group 0, column 'c': MultiPage(dictionary_page: true, data_pages: [0, 1]) (7346 bytes, 1 requests) [data]", + "Read Multi:", + " Row Group 1, column 'a': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + " Row Group 1, column 'b': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + " Row Group 1, column 'c': MultiPage(dictionary_page: true, data_pages: [0, 1]) (7456 bytes, 1 requests) [data]", + ] + "#); +} + +#[tokio::test] +async fn test_read_single_group() { + let test_file = test_file(); + let builder = async_builder(&test_file, test_options()) + .await + // read only second row group + .with_row_groups(vec![1]); + + // Expect to see only IO for Row Group 1. Should see no IO for Row Group 0. + insta::assert_debug_snapshot!(run( + &test_file, + builder).await, @r#" + [ + "Get Provided Metadata", + "Event: Builder Configured", + "Event: Reader Built", + "Read Multi:", + " Row Group 1, column 'a': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + " Row Group 1, column 'b': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + " Row Group 1, column 'c': MultiPage(dictionary_page: true, data_pages: [0, 1]) (7456 bytes, 1 requests) [data]", + ] + "#); +} + +#[tokio::test] +async fn test_read_single_column() { + let test_file = test_file(); + let builder = async_builder(&test_file, test_options()).await; + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + let builder = builder.with_projection(ProjectionMask::columns(&schema_descr, ["b"])); + // Expect to see only IO for column "b". Should see no IO for columns "a" or "c". + insta::assert_debug_snapshot!(run( + &test_file, + builder).await, @r#" + [ + "Get Provided Metadata", + "Event: Builder Configured", + "Event: Reader Built", + "Read Multi:", + " Row Group 0, column 'b': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + "Read Multi:", + " Row Group 1, column 'b': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + ] + "#); +} + +#[tokio::test] +async fn test_read_row_selection() { + // There are 400 total rows spread across 4 data pages (100 rows each) + // select rows 175..225 (i.e. DataPage(1) of row group 0 and DataPage(0) of row group 1) + let test_file = test_file(); + let builder = async_builder(&test_file, test_options()).await; + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + let builder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["a", "b"])) + .with_row_selection(RowSelection::from(vec![ + RowSelector::skip(175), + RowSelector::select(50), + ])); + + // Expect to see only data IO for one page for each column for each row group + insta::assert_debug_snapshot!(run( + &test_file, + builder).await, @r#" + [ + "Get Provided Metadata", + "Event: Builder Configured", + "Event: Reader Built", + "Read Multi:", + " Row Group 0, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", + " Row Group 0, column 'a': DataPage(1) (126 bytes , 1 requests) [data]", + " Row Group 0, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + " Row Group 0, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", + "Read Multi:", + " Row Group 1, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", + " Row Group 1, column 'a': DataPage(0) (113 bytes , 1 requests) [data]", + " Row Group 1, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + " Row Group 1, column 'b': DataPage(0) (113 bytes , 1 requests) [data]", + ] + "#); +} + +#[tokio::test] +async fn test_read_limit() { + // There are 400 total rows spread across 4 data pages (100 rows each) + // a limit of 125 rows should only fetch the first two data pages (DataPage(0) and DataPage(1)) from row group 0 + let test_file = test_file(); + let builder = async_builder(&test_file, test_options()).await; + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + let builder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["a"])) + .with_limit(125); + + insta::assert_debug_snapshot!(run( + &test_file, + builder).await, @r#" + [ + "Get Provided Metadata", + "Event: Builder Configured", + "Event: Reader Built", + "Read Multi:", + " Row Group 0, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", + " Row Group 0, column 'a': DataPage(0) (113 bytes , 1 requests) [data]", + " Row Group 0, column 'a': DataPage(1) (126 bytes , 1 requests) [data]", + ] + "#); +} + +#[tokio::test] +async fn test_read_single_row_filter() { + // Values from column "b" range 400..799 + // filter "b" > 575 and < than 625 + // (last data page in Row Group 0 and first DataPage in Row Group 1) + let test_file = test_file(); + let builder = async_builder(&test_file, test_options()).await; + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + let builder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["a", "b"])) + .with_row_filter(filter_b_575_625(&schema_descr)); + + // Expect to see I/O for column b in both row groups to evaluate filter, + // then a single pages for the "a" column in each row group + insta::assert_debug_snapshot!(run( + &test_file, + builder).await, @r#" + [ + "Get Provided Metadata", + "Event: Builder Configured", + "Event: Reader Built", + "Read Multi:", + " Row Group 0, column 'b': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + "Read Multi:", + " Row Group 0, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", + " Row Group 0, column 'a': DataPage(1) (126 bytes , 1 requests) [data]", + "Read Multi:", + " Row Group 1, column 'b': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + "Read Multi:", + " Row Group 1, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", + " Row Group 1, column 'a': DataPage(0) (113 bytes , 1 requests) [data]", + ] + "#); +} + +#[tokio::test] +async fn test_read_single_row_filter_no_page_index() { + // Values from column "b" range 400..799 + // Apply a filter "b" > 575 and than 625 + // (last data page in Row Group 0 and first DataPage in Row Group 1) + let test_file = test_file(); + let options = test_options().with_page_index_policy(PageIndexPolicy::from(false)); + let builder = async_builder(&test_file, options).await; + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + let builder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["a", "b"])) + .with_row_filter(filter_b_575_625(&schema_descr)); + + // Since we don't have the page index, expect to see: + // 1. I/O for all pages of column b to evaluate the filter + // 2. IO for all pages of column a as the reader doesn't know where the page + // boundaries are so needs to scan them. + insta::assert_debug_snapshot!(run( + &test_file, + builder).await, @r#" + [ + "Get Provided Metadata", + "Event: Builder Configured", + "Event: Reader Built", + "Read Multi:", + " Row Group 0, column 'b': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + "Read Multi:", + " Row Group 0, column 'a': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + "Read Multi:", + " Row Group 1, column 'b': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + "Read Multi:", + " Row Group 1, column 'a': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + ] + "#); +} + +#[tokio::test] +async fn test_read_multiple_row_filter() { + // Values in column "a" range 0..399 + // Values in column "b" range 400..799 + // First filter: "a" > 175 (last data page in Row Group 0) + // Second filter: "b" < 625 (last data page in Row Group 0 and first DataPage in RowGroup 1) + // Read column "c" + let test_file = test_file(); + let builder = async_builder(&test_file, test_options()).await; + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + let builder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["c"])) + .with_row_filter(filter_a_175_b_625(&schema_descr)); + + // Expect that we will see + // 1. IO for all pages of column A (to evaluate the first filter) + // 2. IO for pages of column b that passed the first filter (to evaluate the second filter) + // 3. IO after reader is built only for column c for the rows that passed both filters + insta::assert_debug_snapshot!(run( + &test_file, + builder).await, @r#" + [ + "Get Provided Metadata", + "Event: Builder Configured", + "Event: Reader Built", + "Read Multi:", + " Row Group 0, column 'a': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + "Read Multi:", + " Row Group 0, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + " Row Group 0, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", + "Read Multi:", + " Row Group 0, column 'c': DictionaryPage (7107 bytes, 1 requests) [data]", + " Row Group 0, column 'c': DataPage(1) (126 bytes , 1 requests) [data]", + "Read Multi:", + " Row Group 1, column 'a': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + "Read Multi:", + " Row Group 1, column 'b': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + "Read Multi:", + " Row Group 1, column 'c': DictionaryPage (7217 bytes, 1 requests) [data]", + " Row Group 1, column 'c': DataPage(0) (113 bytes , 1 requests) [data]", + ] + "#); +} + +#[tokio::test] +async fn test_read_single_row_filter_all() { + // Apply a filter that filters out all rows + + let test_file = test_file(); + let builder = async_builder(&test_file, test_options()).await; + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + let builder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["a", "b"])) + .with_row_filter(filter_b_false(&schema_descr)); + + // Expect to see reads for column "b" to evaluate the filter, but no reads + // for column "a" as no rows pass the filter + insta::assert_debug_snapshot!(run( + &test_file, + builder).await, @r#" + [ + "Get Provided Metadata", + "Event: Builder Configured", + "Event: Reader Built", + "Read Multi:", + " Row Group 0, column 'b': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + "Read Multi:", + " Row Group 1, column 'b': MultiPage(dictionary_page: true, data_pages: [0, 1]) (1856 bytes, 1 requests) [data]", + ] + "#); +} + +/// Return a [`ParquetRecordBatchStreamBuilder`] for reading this file +async fn async_builder( + test_file: &TestParquetFile, + options: ArrowReaderOptions, +) -> ParquetRecordBatchStreamBuilder { + let parquet_meta_data = if options.offset_index_policy() != PageIndexPolicy::Skip + || options.column_index_policy() != PageIndexPolicy::Skip + { + Arc::clone(test_file.parquet_metadata()) + } else { + // strip out the page index from the metadata + let metadata = test_file + .parquet_metadata() + .as_ref() + .clone() + .into_builder() + .set_column_index(None) + .set_offset_index(None) + .build(); + Arc::new(metadata) + }; + + let reader = RecordingAsyncFileReader { + bytes: test_file.bytes().clone(), + ops: Arc::clone(test_file.ops()), + parquet_meta_data, + }; + + ParquetRecordBatchStreamBuilder::new_with_options(reader, options) + .await + .unwrap() +} + +/// Build the reader from the specified builder and read all batches from it, +/// and return the operations log. +async fn run( + test_file: &TestParquetFile, + builder: ParquetRecordBatchStreamBuilder, +) -> Vec { + let ops = test_file.ops(); + ops.add_entry(LogEntry::event("Builder Configured")); + let mut stream = builder.build().unwrap(); + ops.add_entry(LogEntry::event("Reader Built")); + while let Some(batch) = stream.next().await { + match batch { + Ok(_) => {} + Err(e) => panic!("Error reading batch: {e}"), + } + } + ops.snapshot() +} + +struct RecordingAsyncFileReader { + bytes: Bytes, + ops: Arc, + parquet_meta_data: Arc, +} + +impl AsyncFileReader for RecordingAsyncFileReader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, parquet::errors::Result> { + let ops = Arc::clone(&self.ops); + let data = self + .bytes + .slice(range.start as usize..range.end as usize) + .clone(); + + // translate to usize from u64 + let logged_range = Range { + start: range.start as usize, + end: range.end as usize, + }; + async move { + ops.add_entry_for_range(&logged_range); + Ok(data) + } + .boxed() + } + + fn get_byte_ranges(&mut self, ranges: Vec>) -> BoxFuture<'_, Result>> { + let ops = Arc::clone(&self.ops); + let datas = ranges + .iter() + .map(|range| { + self.bytes + .slice(range.start as usize..range.end as usize) + .clone() + }) + .collect::>(); + // translate to usize from u64 + let logged_ranges = ranges + .into_iter() + .map(|r| Range { + start: r.start as usize, + end: r.end as usize, + }) + .collect::>(); + + async move { + ops.add_entry_for_ranges(&logged_ranges); + Ok(datas) + } + .boxed() + } + + fn get_metadata<'a>( + &'a mut self, + _options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, Result>> { + let ops = Arc::clone(&self.ops); + let parquet_meta_data = Arc::clone(&self.parquet_meta_data); + async move { + ops.add_entry(LogEntry::GetProvidedMetadata); + Ok(parquet_meta_data) + } + .boxed() + } +} diff --git a/vendor/parquet/tests/arrow_reader/io/mod.rs b/vendor/parquet/tests/arrow_reader/io/mod.rs new file mode 100644 index 0000000..7e50b2c --- /dev/null +++ b/vendor/parquet/tests/arrow_reader/io/mod.rs @@ -0,0 +1,772 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Tests for IO read patterns in the Parquet Reader +//! +//! Each test: +//! 1. Creates a temporary Parquet file with a known row group structure +//! 2. Reads data from that file using the Arrow Parquet Reader, recording the IO operations +//! 3. Asserts the expected IO patterns based on the read operations +//! +//! Note this module contains test infrastructure only. The actual tests are in the +//! sub-modules [`sync_reader`] and [`async_reader`]. +//! +//! Key components: +//! - [`TestParquetFile`] - Represents a Parquet file and its layout +//! - [`OperationLog`] - Records IO operations performed on the file +//! - [`LogEntry`] - Represents a single IO operation in the log + +mod sync_reader; + +#[cfg(feature = "async")] +mod async_reader; + +use arrow::compute::and; +use arrow::compute::kernels::cmp::{gt, lt}; +use arrow_array::cast::AsArray; +use arrow_array::types::Int64Type; +use arrow_array::{ArrayRef, BooleanArray, Int64Array, RecordBatch, StringViewArray}; +use bytes::Bytes; +#[cfg(feature = "async")] +use futures::FutureExt; +#[cfg(feature = "async")] +use futures::future::BoxFuture; +use parquet::arrow::arrow_reader::{ + ArrowPredicateFn, ArrowReaderOptions, ParquetRecordBatchReaderBuilder, RowFilter, +}; +#[cfg(feature = "async")] +use parquet::arrow::async_reader::AsyncFileReader; +use parquet::arrow::{ArrowWriter, ProjectionMask}; +use parquet::data_type::AsBytes; +use parquet::file::FOOTER_SIZE; +use parquet::file::metadata::PageIndexPolicy; +#[cfg(feature = "async")] +use parquet::file::metadata::ParquetMetaDataReader; +use parquet::file::metadata::{FooterTail, ParquetMetaData, ParquetOffsetIndex}; +use parquet::file::page_index::offset_index::PageLocation; +use parquet::file::properties::WriterProperties; +use parquet::schema::types::SchemaDescriptor; +use std::collections::BTreeMap; +use std::fmt::Display; +use std::ops::Range; +use std::sync::{Arc, LazyLock, Mutex}; + +/// Create a new `TestParquetFile` with: +/// 3 columns: "a", "b", "c" +/// +/// 2 row groups, each with 200 rows +/// each data page has 100 rows +/// +/// Values of column "a" are 0..399 +/// Values of column "b" are 400..799 +/// Values of column "c" are alternating strings of length 12 and longer +fn test_file() -> TestParquetFile { + TestParquetFile::new(TEST_FILE_DATA.clone()) +} + +/// Default options for tests +/// +/// Note these tests use the PageIndex to reduce IO +fn test_options() -> ArrowReaderOptions { + ArrowReaderOptions::default().with_page_index_policy(PageIndexPolicy::from(true)) +} + +/// In-memory [`AsyncFileReader`] implementation for tests. +#[cfg(feature = "async")] +#[derive(Clone)] +pub(crate) struct TestReader { + data: Bytes, + metadata: Option>, + requests: Arc>>>, +} + +#[cfg(feature = "async")] +impl TestReader { + pub(crate) fn new(data: Bytes) -> Self { + Self { + data, + metadata: Default::default(), + requests: Default::default(), + } + } + + pub(crate) fn requests(&self) -> Arc>>> { + Arc::clone(&self.requests) + } +} + +#[cfg(feature = "async")] +impl AsyncFileReader for TestReader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, parquet::errors::Result> { + self.requests + .lock() + .unwrap() + .push(range.start as usize..range.end as usize); + futures::future::ready(Ok(self + .data + .slice(range.start as usize..range.end as usize))) + .boxed() + } + + fn get_metadata<'a>( + &'a mut self, + options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, parquet::errors::Result>> { + let mut metadata_reader = ParquetMetaDataReader::new(); + + if let Some(options) = options { + metadata_reader = metadata_reader + .with_column_index_policy(options.column_index_policy()) + .with_offset_index_policy(options.offset_index_policy()); + } + + self.metadata = Some(Arc::new( + metadata_reader.parse_and_finish(&self.data).unwrap(), + )); + futures::future::ready(Ok(self.metadata.clone().unwrap())).boxed() + } +} + +/// Return a row filter that evaluates "b > 575" AND "b < 625" +/// +/// last data page in Row Group 0 and first DataPage in Row Group 1 +fn filter_b_575_625(schema_descr: &SchemaDescriptor) -> RowFilter { + // "b" > 575 and "b" < 625 + let predicate = ArrowPredicateFn::new( + ProjectionMask::columns(schema_descr, ["b"]), + |batch: RecordBatch| { + let scalar_575 = Int64Array::new_scalar(575); + let scalar_625 = Int64Array::new_scalar(625); + let column = batch.column(0).as_primitive::(); + and(>(column, &scalar_575)?, <(column, &scalar_625)?) + }, + ); + RowFilter::new(vec![Box::new(predicate)]) +} + +/// Filter a > 175 and b < 625 +/// First filter: "a" > 175 (last data page in Row Group 0) +/// Second filter: "b" < 625 (last data page in Row Group 0 and first DataPage in RowGroup 1) +fn filter_a_175_b_625(schema_descr: &SchemaDescriptor) -> RowFilter { + // "a" > 175 and "b" < 625 + let predicate_a = ArrowPredicateFn::new( + ProjectionMask::columns(schema_descr, ["a"]), + |batch: RecordBatch| { + let scalar_175 = Int64Array::new_scalar(175); + let column = batch.column(0).as_primitive::(); + gt(column, &scalar_175) + }, + ); + + let predicate_b = ArrowPredicateFn::new( + ProjectionMask::columns(schema_descr, ["b"]), + |batch: RecordBatch| { + let scalar_625 = Int64Array::new_scalar(625); + let column = batch.column(0).as_primitive::(); + lt(column, &scalar_625) + }, + ); + + RowFilter::new(vec![Box::new(predicate_a), Box::new(predicate_b)]) +} + +/// Filter FALSE (no rows) with b +/// Entirely filters out both row groups +/// Note it selects "b" +fn filter_b_false(schema_descr: &SchemaDescriptor) -> RowFilter { + // "false" + let predicate = ArrowPredicateFn::new( + ProjectionMask::columns(schema_descr, ["b"]), + |batch: RecordBatch| { + let result = + BooleanArray::from_iter(std::iter::repeat_n(Some(false), batch.num_rows())); + Ok(result) + }, + ); + RowFilter::new(vec![Box::new(predicate)]) +} + +/// Create a parquet file in memory for testing. See [`test_file`] for details. +static TEST_FILE_DATA: LazyLock = LazyLock::new(|| { + // Input batch has 400 rows, with 3 columns: "a", "b", "c" + // Note c is a different types (so the data page sizes will be different) + let a: ArrayRef = Arc::new(Int64Array::from_iter_values(0..400)); + let b: ArrayRef = Arc::new(Int64Array::from_iter_values(400..800)); + let c: ArrayRef = Arc::new(StringViewArray::from_iter_values((0..400).map(|i| { + if i % 2 == 0 { + format!("string_{i}") + } else { + format!("A string larger than 12 bytes and thus not inlined {i}") + } + }))); + + let input_batch = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + let mut output = Vec::new(); + + let writer_options = WriterProperties::builder() + .set_max_row_group_row_count(Some(200)) + .set_data_page_row_count_limit(100) + .build(); + let mut writer = + ArrowWriter::try_new(&mut output, input_batch.schema(), Some(writer_options)).unwrap(); + + // since the limits are only enforced on batch boundaries, write the input + // batch in chunks of 50 + let mut row_remain = input_batch.num_rows(); + while row_remain > 0 { + let chunk_size = row_remain.min(50); + let chunk = input_batch.slice(input_batch.num_rows() - row_remain, chunk_size); + writer.write(&chunk).unwrap(); + row_remain -= chunk_size; + } + writer.close().unwrap(); + Bytes::from(output) +}); + +/// A test parquet file and its layout. +struct TestParquetFile { + bytes: Bytes, + /// The operation log for IO operations performed on this file + ops: Arc, + /// The (pre-parsed) parquet metadata for this file + #[cfg(feature = "async")] + parquet_metadata: Arc, +} + +impl TestParquetFile { + /// Create a new `TestParquetFile` with the specified temporary directory and path + /// and determines the row group layout. + fn new(bytes: Bytes) -> Self { + // Read the parquet file to determine its layout + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options( + bytes.clone(), + ArrowReaderOptions::default().with_page_index_policy(PageIndexPolicy::from(true)), + ) + .unwrap(); + + let parquet_metadata = Arc::clone(builder.metadata()); + + let offset_index = parquet_metadata + .offset_index() + .expect("Parquet metadata should have a page index"); + + let row_groups = TestRowGroups::new(&parquet_metadata, offset_index); + + // figure out the footer location in the file + let footer_location = bytes.len() - FOOTER_SIZE..bytes.len(); + let footer = bytes.slice(footer_location.clone()); + let footer: &[u8; FOOTER_SIZE] = footer + .as_bytes() + .try_into() // convert to a fixed size array + .unwrap(); + + // figure out the metadata location + let footer = FooterTail::try_new(footer).unwrap(); + let metadata_len = footer.metadata_length(); + let metadata_location = footer_location.start - metadata_len..footer_location.start; + + let ops = Arc::new(OperationLog::new( + footer_location, + metadata_location, + row_groups, + )); + + TestParquetFile { + bytes, + ops, + #[cfg(feature = "async")] + parquet_metadata, + } + } + + /// Return the internal bytes of the parquet file + fn bytes(&self) -> &Bytes { + &self.bytes + } + + /// Return the operation log for this file + fn ops(&self) -> &Arc { + &self.ops + } + + #[cfg(feature = "async")] + fn parquet_metadata(&self) -> &Arc { + &self.parquet_metadata + } +} + +/// Information about a column chunk +#[derive(Debug)] +struct TestColumnChunk { + /// The name of the column + name: String, + + /// The location of the entire column chunk in the file including dictionary pages + /// and data pages. + location: Range, + + /// The offset of the start of of the dictionary page if any + dictionary_page_location: Option, + + /// The location of the data pages in the file + page_locations: Vec, +} + +/// Information about the pages in a single row group +#[derive(Debug)] +struct TestRowGroup { + /// Maps column_name -> Information about the column chunk + columns: BTreeMap, +} + +/// Information about all the row groups in a Parquet file, extracted from its metadata +#[derive(Debug)] +struct TestRowGroups { + /// List of row groups, each containing information about its columns and page locations + row_groups: Vec, +} + +impl TestRowGroups { + fn new(parquet_metadata: &ParquetMetaData, offset_index: &ParquetOffsetIndex) -> Self { + let row_groups = parquet_metadata + .row_groups() + .iter() + .enumerate() + .map(|(rg_index, rg_meta)| { + let columns = rg_meta + .columns() + .iter() + .enumerate() + .map(|(col_idx, col_meta)| { + let column_name = col_meta.column_descr().name().to_string(); + let page_locations = offset_index[rg_index][col_idx].page_locations(); + let dictionary_page_location = col_meta.dictionary_page_offset(); + + // We can find the byte range of the entire column chunk + let (start_offset, length) = col_meta.byte_range(); + let start_offset = start_offset as usize; + let end_offset = start_offset + length as usize; + + TestColumnChunk { + name: column_name.clone(), + location: start_offset..end_offset, + dictionary_page_location, + page_locations: page_locations.clone(), + } + }) + .map(|test_column_chunk| { + // make key=value pairs to insert into the BTreeMap + (test_column_chunk.name.clone(), test_column_chunk) + }) + .collect::>(); + TestRowGroup { columns } + }) + .collect(); + + Self { row_groups } + } + + fn iter(&self) -> impl Iterator { + self.row_groups.iter() + } +} + +/// Type of data read +#[derive(Debug, PartialEq)] +enum PageType { + /// The data page with the specified index + Data { + data_page_index: usize, + }, + Dictionary, + /// Multiple pages read together + Multi { + /// Was the dictionary page included? + dictionary_page: bool, + /// The data pages included + data_page_indices: Vec, + }, +} + +impl Display for PageType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + PageType::Data { data_page_index } => { + write!(f, "DataPage({data_page_index})") + } + PageType::Dictionary => write!(f, "DictionaryPage"), + PageType::Multi { + dictionary_page, + data_page_indices, + } => { + let dictionary_page = if *dictionary_page { + "dictionary_page: true, " + } else { + "" + }; + write!( + f, + "MultiPage({dictionary_page}data_pages: {data_page_indices:?})", + ) + } + } + } +} + +/// Read single logical data object (data page or dictionary page) +/// in one or more requests +#[derive(Debug)] +struct ReadInfo { + row_group_index: usize, + column_name: String, + range: Range, + read_type: PageType, + /// Number of distinct requests (function calls) that were used + num_requests: usize, +} + +impl Display for ReadInfo { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let Self { + row_group_index, + column_name, + range, + read_type, + num_requests, + } = self; + + // If the average read size is less than 10 bytes, assume it is the thrift + // decoder reading the page headers and add an annotation + let annotation = if (range.len() / num_requests) < 10 { + " [header]" + } else { + " [data]" + }; + + // align the read type to 20 characters for better readability, not sure why + // this does not work inline with write! macro below + write!( + f, + "Row Group {row_group_index}, column '{column_name}': {:15} ({:10}, {:8}){annotation}", + // convert to strings so alignment works + format!("{read_type}"), + format!("{} bytes", range.len()), + format!("{num_requests} requests"), + ) + } +} + +/// Store structured entries in the log to make it easier to combine multiple entries +#[derive(Debug)] +enum LogEntry { + /// Read the footer (last 8 bytes) of the parquet file + ReadFooter(Range), + /// Read the metadata of the parquet file + ReadMetadata(Range), + /// Access previously parsed metadata + #[allow(dead_code)] + GetProvidedMetadata, + /// Read a single logical data object + ReadData(ReadInfo), + /// Read one or more logical data objects in a single operation + #[allow(dead_code)] + ReadMultipleData(Vec), + /// Not known where the read came from + Unknown(Range), + /// A user defined event + Event(String), +} + +impl LogEntry { + fn event(event: impl Into) -> Self { + LogEntry::Event(event.into()) + } + + /// Appends a string representation of this log entry to the output vector + fn append_string(&self, output: &mut Vec, indent: usize) { + let indent_str = " ".repeat(indent); + match self { + LogEntry::ReadFooter(range) => { + output.push(format!("{indent_str}Footer: {} bytes", range.len())) + } + LogEntry::ReadMetadata(range) => { + output.push(format!("{indent_str}Metadata: {}", range.len())) + } + LogEntry::GetProvidedMetadata => { + output.push(format!("{indent_str}Get Provided Metadata")) + } + LogEntry::ReadData(read_info) => output.push(format!("{indent_str}{read_info}")), + LogEntry::ReadMultipleData(read_infos) => { + output.push(format!("{indent_str}Read Multi:")); + for read_info in read_infos { + let new_indent = indent + 2; + read_info.append_string(output, new_indent); + } + } + LogEntry::Unknown(range) => { + output.push(format!("{indent_str}UNKNOWN: {range:?} (maybe Page Index)")) + } + LogEntry::Event(event) => output.push(format!("Event: {event}")), + } + } +} + +#[derive(Debug)] +struct OperationLog { + /// The operations performed on the file + ops: Mutex>, + + /// Footer location in the parquet file + footer_location: Range, + + /// Metadata location in the parquet file + metadata_location: Range, + + /// Information about the row group layout in the parquet file, used to + /// translate read operations into human understandable IO operations + /// Path to the parquet file + row_groups: TestRowGroups, +} + +impl OperationLog { + fn new( + footer_location: Range, + metadata_location: Range, + row_groups: TestRowGroups, + ) -> Self { + OperationLog { + ops: Mutex::new(Vec::new()), + metadata_location, + footer_location, + row_groups, + } + } + + /// Add an operation to the log + fn add_entry(&self, entry: LogEntry) { + let mut ops = self.ops.lock().unwrap(); + ops.push(entry); + } + + /// Adds an entry to the operation log for the interesting object that is + /// accessed by the specified range + /// + /// This function checks the ranges in order against possible locations + /// and adds the appropriate operation to the log for the first match found. + fn add_entry_for_range(&self, range: &Range) { + self.add_entry(self.entry_for_range(range)); + } + + /// Adds entries to the operation log for each interesting object that is + /// accessed by the specified range + /// + /// It behaves the same as [`add_entry_for_range`] but for multiple ranges. + #[cfg(feature = "async")] + fn add_entry_for_ranges<'a>(&self, ranges: impl IntoIterator>) { + let entries = ranges + .into_iter() + .map(|range| self.entry_for_range(range)) + .collect::>(); + self.add_entry(LogEntry::ReadMultipleData(entries)); + } + + /// Create an appropriate LogEntry for the specified range + fn entry_for_range(&self, range: &Range) -> LogEntry { + let start = range.start as i64; + let end = range.end as i64; + + // figure out what logical part of the file this range corresponds to + if self.metadata_location.contains(&range.start) + || self.metadata_location.contains(&(range.end - 1)) + { + return LogEntry::ReadMetadata(range.clone()); + } + + if self.footer_location.contains(&range.start) + || self.footer_location.contains(&(range.end - 1)) + { + return LogEntry::ReadFooter(range.clone()); + } + + // Search for the location in each column chunk. + // + // The actual parquet reader must in general decode the page headers + // and determine the byte ranges of the pages. However, for this test + // we assume the following layout: + // + // ```text + // (Dictionary Page) + // (Data Page) + // ... + // (Data Page) + // ``` + // + // We also assume that `self.page_locations` holds the location of all + // data pages, so any read operation that overlaps with a data page + // location is considered a read of that page, and any other read must + // be a dictionary page read. + for (row_group_index, row_group) in self.row_groups.iter().enumerate() { + for (column_name, test_column_chunk) in &row_group.columns { + // Check if the range overlaps with any data page locations + let page_locations = test_column_chunk.page_locations.iter(); + + // What data pages does this range overlap with? + let mut data_page_indices = vec![]; + + for (data_page_index, page_location) in page_locations.enumerate() { + let page_offset = page_location.offset; + let page_end = page_offset + page_location.compressed_page_size as i64; + + // if the range fully contains the page, consider it a read of that page + if start >= page_offset && end <= page_end { + let read_info = ReadInfo { + row_group_index, + column_name: column_name.clone(), + range: range.clone(), + read_type: PageType::Data { data_page_index }, + num_requests: 1, + }; + return LogEntry::ReadData(read_info); + } + + // if the range overlaps with the page, add it to the list of overlapping pages + if start < page_end && end > page_offset { + data_page_indices.push(data_page_index); + } + } + + // was the dictionary page read? + let mut dictionary_page = false; + + // Check if the range overlaps with the dictionary page location + if let Some(dict_page_offset) = test_column_chunk.dictionary_page_location { + let dict_page_end = dict_page_offset + test_column_chunk.location.len() as i64; + if start >= dict_page_offset && end < dict_page_end { + let read_info = ReadInfo { + row_group_index, + column_name: column_name.clone(), + range: range.clone(), + read_type: PageType::Dictionary, + num_requests: 1, + }; + + return LogEntry::ReadData(read_info); + } + + // if the range overlaps with the dictionary page, add it to the list of overlapping pages + if start < dict_page_end && end > dict_page_offset { + dictionary_page = true; + } + } + + // If we can't find a page, but the range overlaps with the + // column chunk location, use the column chunk location + let column_byte_range = &test_column_chunk.location; + if column_byte_range.contains(&range.start) + && column_byte_range.contains(&(range.end - 1)) + { + let read_data_entry = ReadInfo { + row_group_index, + column_name: column_name.clone(), + range: range.clone(), + read_type: PageType::Multi { + data_page_indices, + dictionary_page, + }, + num_requests: 1, + }; + + return LogEntry::ReadData(read_data_entry); + } + } + } + + // If we reach here, the range does not match any known logical part of the file + LogEntry::Unknown(range.clone()) + } + + // Combine entries in the log that are similar to reduce noise in the log. + fn coalesce_entries(&self) { + let mut ops = self.ops.lock().unwrap(); + + // Coalesce entries with the same read type + let prev_ops = std::mem::take(&mut *ops); + for entry in prev_ops { + let Some(last) = ops.last_mut() else { + ops.push(entry); + continue; + }; + + let LogEntry::ReadData(ReadInfo { + row_group_index: last_rg_index, + column_name: last_column_name, + range: last_range, + read_type: last_read_type, + num_requests: last_num_reads, + }) = last + else { + // If the last entry is not a ReadColumnChunk, just push it + ops.push(entry); + continue; + }; + + // If the entry is not a ReadColumnChunk, just push it + let LogEntry::ReadData(ReadInfo { + row_group_index, + column_name, + range, + read_type, + num_requests: num_reads, + }) = &entry + else { + ops.push(entry); + continue; + }; + + // Combine the entries if they are the same and this read is less than 10b. + // + // This heuristic is used to combine small reads (typically 1-2 + // byte) made by the thrift decoder when reading the data/dictionary + // page headers. + if *row_group_index != *last_rg_index + || column_name != last_column_name + || read_type != last_read_type + || (range.start > last_range.end) + || (range.end < last_range.start) + || range.len() > 10 + { + ops.push(entry); + continue; + } + // combine + *last_range = last_range.start.min(range.start)..last_range.end.max(range.end); + *last_num_reads += num_reads; + } + } + + /// return a snapshot of the current operations in the log. + fn snapshot(&self) -> Vec { + self.coalesce_entries(); + let ops = self.ops.lock().unwrap(); + let mut actual = vec![]; + let indent = 0; + ops.iter() + .for_each(|s| s.append_string(&mut actual, indent)); + actual + } +} diff --git a/vendor/parquet/tests/arrow_reader/io/sync_reader.rs b/vendor/parquet/tests/arrow_reader/io/sync_reader.rs new file mode 100644 index 0000000..835029d --- /dev/null +++ b/vendor/parquet/tests/arrow_reader/io/sync_reader.rs @@ -0,0 +1,444 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Tests for the sync reader - [`ParquetRecordBatchReaderBuilder`] + +use crate::io::{ + LogEntry, OperationLog, TestParquetFile, filter_a_175_b_625, filter_b_575_625, filter_b_false, + test_file, test_options, +}; + +use bytes::Bytes; +use parquet::arrow::ProjectionMask; +use parquet::arrow::arrow_reader::{ + ArrowReaderOptions, ParquetRecordBatchReaderBuilder, RowSelection, RowSelector, +}; +use parquet::file::metadata::PageIndexPolicy; +use parquet::file::reader::{ChunkReader, Length}; +use std::io::Read; +use std::sync::Arc; + +#[test] +fn test_read_entire_file() { + // read entire file without any filtering or projection + let test_file = test_file(); + // Expect to see IO for all data pages for each row group and column + let builder = sync_builder(&test_file, test_options()); + insta::assert_debug_snapshot!(run(&test_file, builder), + @r#" + [ + "Footer: 8 bytes", + "Metadata: 1162", + "UNKNOWN: 22230..22877 (maybe Page Index)", + "Event: Builder Configured", + "Event: Reader Built", + "Row Group 0, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 0, column 'a': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 0, column 'a': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 1, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 1, column 'a': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 1, column 'a': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 0, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 0, column 'b': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 0, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 1, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 1, column 'b': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 1, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 0, column 'c': DictionaryPage (7107 bytes, 1 requests) [data]", + "Row Group 0, column 'c': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 0, column 'c': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 1, column 'c': DictionaryPage (7217 bytes, 1 requests) [data]", + "Row Group 1, column 'c': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 1, column 'c': DataPage(1) (126 bytes , 1 requests) [data]", + ] + "#); +} + +#[test] +fn test_read_single_group() { + let test_file = test_file(); + let builder = sync_builder(&test_file, test_options()).with_row_groups(vec![1]); // read only second row group + + // Expect to see only IO for Row Group 1. Should see no IO for Row Group 0. + insta::assert_debug_snapshot!(run(&test_file, builder), + @r#" + [ + "Footer: 8 bytes", + "Metadata: 1162", + "UNKNOWN: 22230..22877 (maybe Page Index)", + "Event: Builder Configured", + "Event: Reader Built", + "Row Group 1, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 1, column 'a': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 1, column 'a': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 1, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 1, column 'b': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 1, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 1, column 'c': DictionaryPage (7217 bytes, 1 requests) [data]", + "Row Group 1, column 'c': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 1, column 'c': DataPage(1) (126 bytes , 1 requests) [data]", + ] + "#); +} + +#[test] +fn test_read_single_column() { + let test_file = test_file(); + let builder = sync_builder(&test_file, test_options()); + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + let builder = builder.with_projection(ProjectionMask::columns(&schema_descr, ["b"])); + // Expect to see only IO for column "b". Should see no IO for columns "a" or "c". + insta::assert_debug_snapshot!(run(&test_file, builder), + @r#" + [ + "Footer: 8 bytes", + "Metadata: 1162", + "UNKNOWN: 22230..22877 (maybe Page Index)", + "Event: Builder Configured", + "Event: Reader Built", + "Row Group 0, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 0, column 'b': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 0, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 1, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 1, column 'b': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 1, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", + ] + "#); +} + +#[test] +fn test_read_single_column_no_page_index() { + let test_file = test_file(); + let options = test_options().with_page_index_policy(PageIndexPolicy::from(false)); + let builder = sync_builder(&test_file, options); + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + let builder = builder.with_projection(ProjectionMask::columns(&schema_descr, ["b"])); + // Expect to see only IO for column "b", should see no IO for columns "a" or "c". + // + // Note that we need to read all data page headers to find the pages for column b + // so there are many more small reads than in the test_read_single_column test above + insta::assert_debug_snapshot!(run(&test_file, builder), + @r#" + [ + "Footer: 8 bytes", + "Metadata: 1162", + "Event: Builder Configured", + "Event: Reader Built", + "Row Group 0, column 'b': DictionaryPage (17 bytes , 17 requests) [header]", + "Row Group 0, column 'b': DictionaryPage (1600 bytes, 1 requests) [data]", + "Row Group 0, column 'b': DataPage(0) (20 bytes , 20 requests) [header]", + "Row Group 0, column 'b': DataPage(0) (93 bytes , 1 requests) [data]", + "Row Group 0, column 'b': DataPage(1) (20 bytes , 20 requests) [header]", + "Row Group 0, column 'b': DataPage(1) (106 bytes , 1 requests) [data]", + "Row Group 1, column 'b': DictionaryPage (17 bytes , 17 requests) [header]", + "Row Group 1, column 'b': DictionaryPage (1600 bytes, 1 requests) [data]", + "Row Group 1, column 'b': DataPage(0) (20 bytes , 20 requests) [header]", + "Row Group 1, column 'b': DataPage(0) (93 bytes , 1 requests) [data]", + "Row Group 1, column 'b': DataPage(1) (20 bytes , 20 requests) [header]", + "Row Group 1, column 'b': DataPage(1) (106 bytes , 1 requests) [data]", + ] + "#); +} + +#[test] +fn test_read_row_selection() { + // There are 400 total rows spread across 4 data pages (100 rows each) + // select rows 175..225 (i.e. DataPage(1) of row group 0 and DataPage(0) of row group 1) + let test_file = test_file(); + let builder = sync_builder(&test_file, test_options()); + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + let builder = builder + .with_projection( + // read both "a" and "b" + ProjectionMask::columns(&schema_descr, ["a", "b"]), + ) + .with_row_selection(RowSelection::from(vec![ + RowSelector::skip(175), + RowSelector::select(50), + ])); + + // Expect to see only data IO for one page for each column for each row group + // Note the data page headers for all pages need to be read to find the correct pages + insta::assert_debug_snapshot!(run(&test_file, builder), + @r#" + [ + "Footer: 8 bytes", + "Metadata: 1162", + "UNKNOWN: 22230..22877 (maybe Page Index)", + "Event: Builder Configured", + "Event: Reader Built", + "Row Group 0, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 0, column 'a': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 0, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 0, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 1, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 1, column 'a': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 1, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 1, column 'b': DataPage(0) (113 bytes , 1 requests) [data]", + ] + "#); +} + +#[test] +fn test_read_limit() { + // There are 400 total rows spread across 4 data pages (100 rows each) + // a limit of 125 rows should only fetch the first two data pages (DataPage(0) and DataPage(1)) from row group 0 + let test_file = test_file(); + let builder = sync_builder(&test_file, test_options()); + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + let builder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["a"])) + .with_limit(125); + + insta::assert_debug_snapshot!(run(&test_file, builder), + @r#" + [ + "Footer: 8 bytes", + "Metadata: 1162", + "UNKNOWN: 22230..22877 (maybe Page Index)", + "Event: Builder Configured", + "Event: Reader Built", + "Row Group 0, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 0, column 'a': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 0, column 'a': DataPage(1) (126 bytes , 1 requests) [data]", + ] + "#); +} + +#[test] +fn test_read_single_row_filter() { + // Values from column "b" range 400..799 + // filter "b" > 575 and < 625 + // (last data page in Row Group 0 and first DataPage in Row Group 1) + let test_file = test_file(); + let builder = sync_builder(&test_file, test_options()); + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + let builder = builder + .with_projection( + // read both "a" and "b" + ProjectionMask::columns(&schema_descr, ["a", "b"]), + ) + // "b" > 575 and "b" < 625 + .with_row_filter(filter_b_575_625(&schema_descr)); + + // Expect to see I/O for column b in both row groups and then reading just a + // single pages for a in each row group + // + // Note there is significant IO that happens during the construction of the + // reader (between "Builder Configured" and "Reader Built") + insta::assert_debug_snapshot!(run(&test_file, builder), + @r#" + [ + "Footer: 8 bytes", + "Metadata: 1162", + "UNKNOWN: 22230..22877 (maybe Page Index)", + "Event: Builder Configured", + "Row Group 0, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 0, column 'b': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 0, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 1, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 1, column 'b': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 1, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", + "Event: Reader Built", + "Row Group 0, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 0, column 'a': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 0, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 0, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 1, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 1, column 'a': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 1, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 1, column 'b': DataPage(0) (113 bytes , 1 requests) [data]", + ] + "#); +} + +#[test] +fn test_read_multiple_row_filter() { + // Values in column "a" range 0..399 + // Values in column "b" range 400..799 + // First filter: "a" > 175 (last data page in Row Group 0) + // Second filter: "b" < 625 (last data page in Row Group 0 and first DataPage in RowGroup 1) + // Read column "c" + let test_file = test_file(); + let builder = sync_builder(&test_file, test_options()); + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + let builder = builder + .with_projection( + ProjectionMask::columns(&schema_descr, ["c"]), // read "c" + ) + // a > 175 and b < 625 + .with_row_filter(filter_a_175_b_625(&schema_descr)); + + // Expect that we will see + // 1. IO for all pages of column A + // 2. IO for pages of column b that passed 1. + // 3. IO after reader is built only for column c + // + // Note there is significant IO that happens during the construction of the + // reader (between "Builder Configured" and "Reader Built") + insta::assert_debug_snapshot!(run(&test_file, builder), + @r#" + [ + "Footer: 8 bytes", + "Metadata: 1162", + "UNKNOWN: 22230..22877 (maybe Page Index)", + "Event: Builder Configured", + "Row Group 0, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 0, column 'a': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 0, column 'a': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 1, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 1, column 'a': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 1, column 'a': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 0, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 0, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 1, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 1, column 'b': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 1, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", + "Event: Reader Built", + "Row Group 0, column 'c': DictionaryPage (7107 bytes, 1 requests) [data]", + "Row Group 0, column 'c': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 1, column 'c': DictionaryPage (7217 bytes, 1 requests) [data]", + "Row Group 1, column 'c': DataPage(0) (113 bytes , 1 requests) [data]", + ] + "#); +} + +#[test] +fn test_read_single_row_filter_all() { + // Apply a filter that entirely filters out rows based on a predicate from one column + // should not read any data pages for any other column + + let test_file = test_file(); + let builder = sync_builder(&test_file, test_options()); + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + let builder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["a", "b"])) + .with_row_filter(filter_b_false(&schema_descr)); + + // Expect to see the Footer and Metadata, then I/O for column b + // in both row groups but then nothing for column "a" + // since the row filter entirely filters out all rows. + // + // Note that all IO that happens during the construction of the reader + // (between "Builder Configured" and "Reader Built") + insta::assert_debug_snapshot!(run(&test_file, builder), + @r#" + [ + "Footer: 8 bytes", + "Metadata: 1162", + "UNKNOWN: 22230..22877 (maybe Page Index)", + "Event: Builder Configured", + "Row Group 0, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 0, column 'b': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 0, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 1, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 1, column 'b': DataPage(0) (113 bytes , 1 requests) [data]", + "Row Group 1, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", + "Event: Reader Built", + ] + "#); +} + +/// Return a [`ParquetRecordBatchReaderBuilder`] for reading this file +fn sync_builder( + test_file: &TestParquetFile, + options: ArrowReaderOptions, +) -> ParquetRecordBatchReaderBuilder { + let reader = RecordingChunkReader { + inner: test_file.bytes().clone(), + ops: Arc::clone(test_file.ops()), + }; + ParquetRecordBatchReaderBuilder::try_new_with_options(reader, options) + .expect("ParquetRecordBatchReaderBuilder") +} + +/// build the reader, and read all batches from it, returning the recorded IO operations +fn run( + test_file: &TestParquetFile, + builder: ParquetRecordBatchReaderBuilder, +) -> Vec { + let ops = test_file.ops(); + ops.add_entry(LogEntry::event("Builder Configured")); + let reader = builder.build().unwrap(); + ops.add_entry(LogEntry::event("Reader Built")); + for batch in reader { + match batch { + Ok(_) => {} + Err(e) => panic!("Error reading batch: {e}"), + } + } + ops.snapshot() +} + +/// Records IO operations on an in-memory chunk reader +struct RecordingChunkReader { + inner: Bytes, + ops: Arc, +} + +impl Length for RecordingChunkReader { + fn len(&self) -> u64 { + self.inner.len() as u64 + } +} + +impl ChunkReader for RecordingChunkReader { + type T = RecordingStdIoReader; + + fn get_read(&self, start: u64) -> parquet::errors::Result { + let reader = RecordingStdIoReader { + start: start as usize, + inner: self.inner.clone(), + ops: Arc::clone(&self.ops), + }; + Ok(reader) + } + + fn get_bytes(&self, start: u64, length: usize) -> parquet::errors::Result { + let start = start as usize; + let range = start..start + length; + self.ops.add_entry_for_range(&range); + Ok(self.inner.slice(start..start + length)) + } +} + +/// Wrapper around a `Bytes` object that implements `Read` +struct RecordingStdIoReader { + /// current offset in the inner `Bytes` that this reader is reading from + start: usize, + inner: Bytes, + ops: Arc, +} + +impl Read for RecordingStdIoReader { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + let remain = self.inner.len() - self.start; + let start = self.start; + let read_length = buf.len().min(remain); + let read_range = start..start + read_length; + + self.ops.add_entry_for_range(&read_range); + + buf.copy_from_slice(self.inner.slice(read_range).as_ref()); + // Update the inner position + self.start += read_length; + Ok(read_length) + } +} diff --git a/vendor/parquet/tests/arrow_reader/large_string_overflow.rs b/vendor/parquet/tests/arrow_reader/large_string_overflow.rs new file mode 100644 index 0000000..7fc99d8 --- /dev/null +++ b/vendor/parquet/tests/arrow_reader/large_string_overflow.rs @@ -0,0 +1,134 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow_array::builder::BinaryBuilder; +use arrow_array::{ArrayRef, RecordBatch}; +use arrow_schema::{DataType, Field, Schema}; + +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use parquet::arrow::arrow_writer::ArrowWriter; +use parquet::basic::Encoding; +use parquet::file::properties::WriterProperties; + +use tempfile::tempfile; + +/// Number of rows written +const ROWS: usize = 1024; + +/// Size of each binary value (~3MB) +/// 1024 * 3MB ≈ 3GB total → guaranteed offset overflow for i32 +const VALUE_SIZE: usize = 3 * 1024 * 1024; + +fn make_large_binary_array() -> ArrayRef { + let mut builder = BinaryBuilder::new(); + + for _ in 0..ROWS { + let data = vec![b'a'; VALUE_SIZE]; + builder.append_value(&data); + } + + Arc::new(builder.finish()) as ArrayRef +} + +fn write_parquet_with_encoding(array: ArrayRef, encoding: Encoding) -> std::fs::File { + let schema = Arc::new(Schema::new(vec![Field::new( + "col", + DataType::Binary, + false, + )])); + + let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap(); + + let file = tempfile().unwrap(); + + let builder = WriterProperties::builder(); + let builder = match encoding { + Encoding::RLE_DICTIONARY => builder.set_dictionary_enabled(true), + _ => builder.set_dictionary_enabled(false).set_encoding(encoding), + }; + + let props = builder.build(); + + let mut writer = ArrowWriter::try_new(file.try_clone().unwrap(), schema, Some(props)).unwrap(); + + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + file +} + +#[test] +// Panics until https://github.com/apache/arrow-rs/issues/7973 is fixed +#[should_panic(expected = "byte array offset overflow")] +fn large_binary_plain_encoding_overflow() { + let array = make_large_binary_array(); + let file = write_parquet_with_encoding(array, Encoding::PLAIN); + + let mut reader = ParquetRecordBatchReaderBuilder::try_new(file) + .unwrap() + .build() + .unwrap(); + + let _ = reader.next().unwrap(); +} + +#[test] +// Panics until https://github.com/apache/arrow-rs/issues/7973 is fixed +#[should_panic(expected = "byte array offset overflow")] +fn large_binary_delta_length_encoding_overflow() { + let array = make_large_binary_array(); + let file = write_parquet_with_encoding(array, Encoding::DELTA_LENGTH_BYTE_ARRAY); + + let mut reader = ParquetRecordBatchReaderBuilder::try_new(file) + .unwrap() + .build() + .unwrap(); + + let _ = reader.next().unwrap(); +} + +#[test] +// Panics until https://github.com/apache/arrow-rs/issues/7973 is fixed +#[should_panic(expected = "byte array offset overflow")] +fn large_binary_delta_byte_array_encoding_overflow() { + let array = make_large_binary_array(); + let file = write_parquet_with_encoding(array, Encoding::DELTA_BYTE_ARRAY); + + let mut reader = ParquetRecordBatchReaderBuilder::try_new(file) + .unwrap() + .build() + .unwrap(); + + let _ = reader.next().unwrap(); +} + +#[test] +// Panics until https://github.com/apache/arrow-rs/issues/7973 is fixed +#[should_panic(expected = "byte array offset overflow")] +fn large_binary_rle_dictionary_encoding_overflow() { + let array = make_large_binary_array(); + let file = write_parquet_with_encoding(array, Encoding::RLE_DICTIONARY); + + let mut reader = ParquetRecordBatchReaderBuilder::try_new(file) + .unwrap() + .build() + .unwrap(); + + let _ = reader.next().unwrap(); +} diff --git a/vendor/parquet/tests/arrow_reader/mod.rs b/vendor/parquet/tests/arrow_reader/mod.rs new file mode 100644 index 0000000..404bebc --- /dev/null +++ b/vendor/parquet/tests/arrow_reader/mod.rs @@ -0,0 +1,1156 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_array::types::{Int8Type, Int32Type}; +use arrow_array::{ + Array, ArrayRef, BinaryArray, BinaryViewArray, BooleanArray, Date32Array, Date64Array, + Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, DictionaryArray, + FixedSizeBinaryArray, Float16Array, Float32Array, Float64Array, Int8Array, Int16Array, + Int32Array, Int64Array, LargeBinaryArray, LargeStringArray, RecordBatch, StringArray, + StringViewArray, StructArray, Time32MillisecondArray, Time32SecondArray, + Time64MicrosecondArray, Time64NanosecondArray, TimestampMicrosecondArray, + TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, UInt8Array, + UInt16Array, UInt32Array, UInt64Array, +}; +use arrow_buffer::i256; +use arrow_schema::{DataType, Field, Schema, TimeUnit}; +use chrono::Datelike; +use chrono::{Duration, TimeDelta}; +use half::f16; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::{ + DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH, EnabledStatistics, WriterProperties, +}; +use std::sync::Arc; +use tempfile::NamedTempFile; + +mod bad_data; +mod bloom_filter; +#[cfg(feature = "crc")] +mod checksum; +mod int96_stats_roundtrip; +mod invalid_utf8; +mod io; +mod large_string_overflow; +#[cfg(feature = "async")] +mod predicate_cache; +mod row_filter; +mod statistics; + +// returns a struct array with columns "int32_col", "float32_col" and "float64_col" with the specified values +fn struct_array(input: Vec<(Option, Option, Option)>) -> ArrayRef { + let int_32: Int32Array = input.iter().map(|(i, _, _)| i).collect(); + let float_32: Float32Array = input.iter().map(|(_, f, _)| f).collect(); + let float_64: Float64Array = input.iter().map(|(_, _, f)| f).collect(); + + let nullable = true; + let struct_array = StructArray::from(vec![ + ( + Arc::new(Field::new("int32_col", DataType::Int32, nullable)), + Arc::new(int_32) as ArrayRef, + ), + ( + Arc::new(Field::new("float32_col", DataType::Float32, nullable)), + Arc::new(float_32) as ArrayRef, + ), + ( + Arc::new(Field::new("float64_col", DataType::Float64, nullable)), + Arc::new(float_64) as ArrayRef, + ), + ]); + Arc::new(struct_array) +} + +/// What data to use +#[derive(Debug, Clone, Copy)] +enum Scenario { + Boolean, + Timestamps, + Dates, + Int, + Int32Range, + UInt, + UInt32Range, + Time32Second, + Time32Millisecond, + Time64Nanosecond, + Time64Microsecond, + /// 7 Rows, for each i8, i16, i32, i64, u8, u16, u32, u64, f32, f64 + /// -MIN, -100, -1, 0, 1, 100, MAX + NumericLimits, + Float16, + Float32, + Float64, + Decimal32, + Decimal64, + Decimal128, + Decimal256, + ByteArray, + Dictionary, + PeriodsInColumnNames, + StructArray, + UTF8, + /// UTF8 with max and min values truncated + TruncatedUTF8, + UTF8View, + BinaryView, +} + +impl Scenario { + // If the test scenario needs to set `set_statistics_truncate_length` to test + // statistics truncation. + fn truncate_stats(&self) -> bool { + matches!(self, Scenario::TruncatedUTF8) + } +} + +fn make_boolean_batch(v: Vec>) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new( + "bool", + DataType::Boolean, + true, + )])); + let array = Arc::new(BooleanArray::from(v)) as ArrayRef; + RecordBatch::try_new(schema, vec![array.clone()]).unwrap() +} + +/// Return record batch with a few rows of data for all of the supported timestamp types +/// values with the specified offset +/// +/// Columns are named: +/// "nanos" --> TimestampNanosecondArray +/// "nanos_timezoned" --> TimestampNanosecondArray with timezone +/// "micros" --> TimestampMicrosecondArray +/// "micros_timezoned" --> TimestampMicrosecondArray with timezone +/// "millis" --> TimestampMillisecondArray +/// "millis_timezoned" --> TimestampMillisecondArray with timezone +/// "seconds" --> TimestampSecondArray +/// "seconds_timezoned" --> TimestampSecondArray with timezone +/// "names" --> StringArray +fn make_timestamp_batch(offset: Duration) -> RecordBatch { + let ts_strings = vec![ + Some("2020-01-01T01:01:01.0000000000001"), + Some("2020-01-01T01:02:01.0000000000001"), + Some("2020-01-01T02:01:01.0000000000001"), + None, + Some("2020-01-02T01:01:01.0000000000001"), + ]; + + let tz_string = "Pacific/Efate"; + + let offset_nanos = offset.num_nanoseconds().expect("non overflow nanos"); + + let ts_nanos = ts_strings + .into_iter() + .map(|t| { + t.map(|t| { + offset_nanos + + t.parse::() + .unwrap() + .and_utc() + .timestamp_nanos_opt() + .unwrap() + }) + }) + .collect::>(); + + let ts_micros = ts_nanos + .iter() + .map(|t| t.as_ref().map(|ts_nanos| ts_nanos / 1000)) + .collect::>(); + + let ts_millis = ts_nanos + .iter() + .map(|t| t.as_ref().map(|ts_nanos| ts_nanos / 1000000)) + .collect::>(); + + let ts_seconds = ts_nanos + .iter() + .map(|t| t.as_ref().map(|ts_nanos| ts_nanos / 1000000000)) + .collect::>(); + + let names = ts_nanos + .iter() + .enumerate() + .map(|(i, _)| format!("Row {i} + {offset}")) + .collect::>(); + + let arr_nanos = TimestampNanosecondArray::from(ts_nanos.clone()); + let arr_nanos_timezoned = TimestampNanosecondArray::from(ts_nanos).with_timezone(tz_string); + let arr_micros = TimestampMicrosecondArray::from(ts_micros.clone()); + let arr_micros_timezoned = TimestampMicrosecondArray::from(ts_micros).with_timezone(tz_string); + let arr_millis = TimestampMillisecondArray::from(ts_millis.clone()); + let arr_millis_timezoned = TimestampMillisecondArray::from(ts_millis).with_timezone(tz_string); + let arr_seconds = TimestampSecondArray::from(ts_seconds.clone()); + let arr_seconds_timezoned = TimestampSecondArray::from(ts_seconds).with_timezone(tz_string); + + let names = names.iter().map(|s| s.as_str()).collect::>(); + let arr_names = StringArray::from(names); + + let schema = Schema::new(vec![ + Field::new("nanos", arr_nanos.data_type().clone(), true), + Field::new( + "nanos_timezoned", + arr_nanos_timezoned.data_type().clone(), + true, + ), + Field::new("micros", arr_micros.data_type().clone(), true), + Field::new( + "micros_timezoned", + arr_micros_timezoned.data_type().clone(), + true, + ), + Field::new("millis", arr_millis.data_type().clone(), true), + Field::new( + "millis_timezoned", + arr_millis_timezoned.data_type().clone(), + true, + ), + Field::new("seconds", arr_seconds.data_type().clone(), true), + Field::new( + "seconds_timezoned", + arr_seconds_timezoned.data_type().clone(), + true, + ), + Field::new("name", arr_names.data_type().clone(), true), + ]); + let schema = Arc::new(schema); + + RecordBatch::try_new( + schema, + vec![ + Arc::new(arr_nanos), + Arc::new(arr_nanos_timezoned), + Arc::new(arr_micros), + Arc::new(arr_micros_timezoned), + Arc::new(arr_millis), + Arc::new(arr_millis_timezoned), + Arc::new(arr_seconds), + Arc::new(arr_seconds_timezoned), + Arc::new(arr_names), + ], + ) + .unwrap() +} + +/// Return record batch with i8, i16, i32, and i64 sequences +/// +/// Columns are named +/// "i8" -> Int8Array +/// "i16" -> Int16Array +/// "i32" -> Int32Array +/// "i64" -> Int64Array +fn make_int_batches(start: i8, end: i8) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("i8", DataType::Int8, true), + Field::new("i16", DataType::Int16, true), + Field::new("i32", DataType::Int32, true), + Field::new("i64", DataType::Int64, true), + ])); + let v8: Vec = (start..end).collect(); + let v16: Vec = (start as _..end as _).collect(); + let v32: Vec = (start as _..end as _).collect(); + let v64: Vec = (start as _..end as _).collect(); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int8Array::from(v8)) as ArrayRef, + Arc::new(Int16Array::from(v16)) as ArrayRef, + Arc::new(Int32Array::from(v32)) as ArrayRef, + Arc::new(Int64Array::from(v64)) as ArrayRef, + ], + ) + .unwrap() +} + +/// Return record batch with Time32Second, Time32Millisecond sequences +fn make_time32_batches(scenario: Scenario, v: Vec) -> RecordBatch { + match scenario { + Scenario::Time32Second => { + let schema = Arc::new(Schema::new(vec![Field::new( + "second", + DataType::Time32(TimeUnit::Second), + true, + )])); + let array = Arc::new(Time32SecondArray::from(v)) as ArrayRef; + RecordBatch::try_new(schema, vec![array]).unwrap() + } + Scenario::Time32Millisecond => { + let schema = Arc::new(Schema::new(vec![Field::new( + "millisecond", + DataType::Time32(TimeUnit::Millisecond), + true, + )])); + let array = Arc::new(Time32MillisecondArray::from(v)) as ArrayRef; + RecordBatch::try_new(schema, vec![array]).unwrap() + } + _ => panic!("Unsupported scenario for Time32"), + } +} + +/// Return record batch with Time64Microsecond, Time64Nanosecond sequences +fn make_time64_batches(scenario: Scenario, v: Vec) -> RecordBatch { + match scenario { + Scenario::Time64Microsecond => { + let schema = Arc::new(Schema::new(vec![Field::new( + "microsecond", + DataType::Time64(TimeUnit::Microsecond), + true, + )])); + let array = Arc::new(Time64MicrosecondArray::from(v)) as ArrayRef; + RecordBatch::try_new(schema, vec![array]).unwrap() + } + Scenario::Time64Nanosecond => { + let schema = Arc::new(Schema::new(vec![Field::new( + "nanosecond", + DataType::Time64(TimeUnit::Nanosecond), + true, + )])); + let array = Arc::new(Time64NanosecondArray::from(v)) as ArrayRef; + RecordBatch::try_new(schema, vec![array]).unwrap() + } + _ => panic!("Unsupported scenario for Time64"), + } +} +/// Return record batch with u8, u16, u32, and u64 sequences +/// +/// Columns are named +/// "u8" -> UInt8Array +/// "u16" -> UInt16Array +/// "u32" -> UInt32Array +/// "u64" -> UInt64Array +fn make_uint_batches(start: u8, end: u8) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("u8", DataType::UInt8, true), + Field::new("u16", DataType::UInt16, true), + Field::new("u32", DataType::UInt32, true), + Field::new("u64", DataType::UInt64, true), + ])); + let v8: Vec = (start..end).collect(); + let v16: Vec = (start as _..end as u16).collect(); + let v32: Vec = (start as _..end as u32).collect(); + let v64: Vec = (start as _..end as u64).collect(); + RecordBatch::try_new( + schema, + vec![ + Arc::new(UInt8Array::from(v8)) as ArrayRef, + Arc::new(UInt16Array::from(v16)) as ArrayRef, + Arc::new(UInt32Array::from(v32)) as ArrayRef, + Arc::new(UInt64Array::from(v64)) as ArrayRef, + ], + ) + .unwrap() +} + +fn make_int32_range(start: i32, end: i32) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)])); + let v = vec![start, end]; + let array = Arc::new(Int32Array::from(v)) as ArrayRef; + RecordBatch::try_new(schema, vec![array.clone()]).unwrap() +} + +fn make_uint32_range(start: u32, end: u32) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("u", DataType::UInt32, true)])); + let v = vec![start, end]; + let array = Arc::new(UInt32Array::from(v)) as ArrayRef; + RecordBatch::try_new(schema, vec![array.clone()]).unwrap() +} + +/// Return record batch with f64 vector +/// +/// Columns are named +/// "f" -> Float64Array +fn make_f64_batch(v: Vec) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("f", DataType::Float64, true)])); + let array = Arc::new(Float64Array::from(v)) as ArrayRef; + RecordBatch::try_new(schema, vec![array.clone()]).unwrap() +} + +fn make_f32_batch(v: Vec) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("f", DataType::Float32, true)])); + let array = Arc::new(Float32Array::from(v)) as ArrayRef; + RecordBatch::try_new(schema, vec![array.clone()]).unwrap() +} + +fn make_f16_batch(v: Vec) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("f", DataType::Float16, true)])); + let array = Arc::new(Float16Array::from(v)) as ArrayRef; + RecordBatch::try_new(schema, vec![array.clone()]).unwrap() +} + +/// Return record batch with decimal32 vector +/// +/// Columns are named +/// "decimal32_col" -> Decimal32Array +fn make_decimal32_batch(v: Vec, precision: u8, scale: i8) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new( + "decimal32_col", + DataType::Decimal32(precision, scale), + true, + )])); + let array = Arc::new( + Decimal32Array::from(v) + .with_precision_and_scale(precision, scale) + .unwrap(), + ) as ArrayRef; + RecordBatch::try_new(schema, vec![array.clone()]).unwrap() +} + +/// Return record batch with decimal64 vector +/// +/// Columns are named +/// "decimal64_col" -> Decimal64Array +fn make_decimal64_batch(v: Vec, precision: u8, scale: i8) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new( + "decimal64_col", + DataType::Decimal64(precision, scale), + true, + )])); + let array = Arc::new( + Decimal64Array::from(v) + .with_precision_and_scale(precision, scale) + .unwrap(), + ) as ArrayRef; + RecordBatch::try_new(schema, vec![array.clone()]).unwrap() +} + +/// Return record batch with decimal128 vector +/// +/// Columns are named +/// "decimal128_col" -> Decimal128Array +fn make_decimal128_batch(v: Vec, precision: u8, scale: i8) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new( + "decimal128_col", + DataType::Decimal128(precision, scale), + true, + )])); + let array = Arc::new( + Decimal128Array::from(v) + .with_precision_and_scale(precision, scale) + .unwrap(), + ) as ArrayRef; + RecordBatch::try_new(schema, vec![array.clone()]).unwrap() +} + +/// Return record batch with decimal256 vector +/// +/// Columns are named +/// "decimal256_col" -> Decimal256Array +fn make_decimal256_batch(v: Vec, precision: u8, scale: i8) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new( + "decimal256_col", + DataType::Decimal256(precision, scale), + true, + )])); + let array = Arc::new( + Decimal256Array::from(v) + .with_precision_and_scale(precision, scale) + .unwrap(), + ) as ArrayRef; + RecordBatch::try_new(schema, vec![array]).unwrap() +} + +/// Return record batch with a few rows of data for all of the supported date +/// types with the specified offset (in days) +/// +/// Columns are named: +/// "date32" --> Date32Array +/// "date64" --> Date64Array +/// "names" --> StringArray +fn make_date_batch(offset: Duration) -> RecordBatch { + let date_strings = vec![ + Some("2020-01-01"), + Some("2020-01-02"), + Some("2020-01-03"), + None, + Some("2020-01-04"), + ]; + + let names = date_strings + .iter() + .enumerate() + .map(|(i, val)| format!("Row {i} + {offset}: {val:?}")) + .collect::>(); + + // Copied from `cast.rs` cast kernel due to lack of temporal kernels + // https://github.com/apache/arrow-rs/issues/527 + const EPOCH_DAYS_FROM_CE: i32 = 719_163; + + let date_seconds = date_strings + .iter() + .map(|t| { + t.map(|t| { + let t = t.parse::().unwrap(); + let t = t + offset; + t.num_days_from_ce() - EPOCH_DAYS_FROM_CE + }) + }) + .collect::>(); + + let date_millis = date_strings + .into_iter() + .map(|t| { + t.map(|t| { + let t = t + .parse::() + .unwrap() + .and_time(chrono::NaiveTime::from_hms_opt(0, 0, 0).unwrap()); + let t = t + offset; + t.and_utc().timestamp_millis() + }) + }) + .collect::>(); + + let arr_date32 = Date32Array::from(date_seconds); + let arr_date64 = Date64Array::from(date_millis); + + let names = names.iter().map(|s| s.as_str()).collect::>(); + let arr_names = StringArray::from(names); + + let schema = Schema::new(vec![ + Field::new("date32", arr_date32.data_type().clone(), true), + Field::new("date64", arr_date64.data_type().clone(), true), + Field::new("name", arr_names.data_type().clone(), true), + ]); + let schema = Arc::new(schema); + + RecordBatch::try_new( + schema, + vec![ + Arc::new(arr_date32), + Arc::new(arr_date64), + Arc::new(arr_names), + ], + ) + .unwrap() +} + +/// returns a batch with two columns (note "service.name" is the name +/// of the column. It is *not* a table named service.name +/// +/// name | service.name +fn make_bytearray_batch( + name: &str, + string_values: Vec<&str>, + binary_values: Vec<&[u8]>, + fixedsize_values: Vec<&[u8; 3]>, + // i64 offset. + large_binary_values: Vec<&[u8]>, +) -> RecordBatch { + let num_rows = string_values.len(); + let name: StringArray = std::iter::repeat_n(Some(name), num_rows).collect(); + let service_string: StringArray = string_values.iter().map(Some).collect(); + let service_binary: BinaryArray = binary_values.iter().map(Some).collect(); + let service_fixedsize: FixedSizeBinaryArray = fixedsize_values + .iter() + .map(|value| Some(value.as_slice())) + .collect::>() + .try_into() + .unwrap(); + let service_large_binary: LargeBinaryArray = large_binary_values.iter().map(Some).collect(); + + let schema = Schema::new(vec![ + Field::new("name", name.data_type().clone(), true), + // note the column name has a period in it! + Field::new("service_string", service_string.data_type().clone(), true), + Field::new("service_binary", service_binary.data_type().clone(), true), + Field::new( + "service_fixedsize", + service_fixedsize.data_type().clone(), + true, + ), + Field::new( + "service_large_binary", + service_large_binary.data_type().clone(), + true, + ), + ]); + let schema = Arc::new(schema); + + RecordBatch::try_new( + schema, + vec![ + Arc::new(name), + Arc::new(service_string), + Arc::new(service_binary), + Arc::new(service_fixedsize), + Arc::new(service_large_binary), + ], + ) + .unwrap() +} + +/// returns a batch with two columns (note "service.name" is the name +/// of the column. It is *not* a table named service.name +/// +/// name | service.name +fn make_names_batch(name: &str, service_name_values: Vec<&str>) -> RecordBatch { + let num_rows = service_name_values.len(); + let name: StringArray = std::iter::repeat_n(Some(name), num_rows).collect(); + let service_name: StringArray = service_name_values.iter().map(Some).collect(); + + let schema = Schema::new(vec![ + Field::new("name", name.data_type().clone(), true), + // note the column name has a period in it! + Field::new("service.name", service_name.data_type().clone(), true), + ]); + let schema = Arc::new(schema); + + RecordBatch::try_new(schema, vec![Arc::new(name), Arc::new(service_name)]).unwrap() +} + +fn make_numeric_limit_batch() -> RecordBatch { + let i8 = Int8Array::from(vec![i8::MIN, 100, -1, 0, 1, -100, i8::MAX]); + let i16 = Int16Array::from(vec![i16::MIN, 100, -1, 0, 1, -100, i16::MAX]); + let i32 = Int32Array::from(vec![i32::MIN, 100, -1, 0, 1, -100, i32::MAX]); + let i64 = Int64Array::from(vec![i64::MIN, 100, -1, 0, 1, -100, i64::MAX]); + let u8 = UInt8Array::from(vec![u8::MIN, 100, 1, 0, 1, 100, u8::MAX]); + let u16 = UInt16Array::from(vec![u16::MIN, 100, 1, 0, 1, 100, u16::MAX]); + let u32 = UInt32Array::from(vec![u32::MIN, 100, 1, 0, 1, 100, u32::MAX]); + let u64 = UInt64Array::from(vec![u64::MIN, 100, 1, 0, 1, 100, u64::MAX]); + let f32 = Float32Array::from(vec![f32::MIN, 100.0, -1.0, 0.0, 1.0, -100.0, f32::MAX]); + let f64 = Float64Array::from(vec![f64::MIN, 100.0, -1.0, 0.0, 1.0, -100.0, f64::MAX]); + let f32_nan = Float32Array::from(vec![f32::NAN, 100.0, -1.0, 0.0, 1.0, -100.0, f32::NAN]); + let f64_nan = Float64Array::from(vec![f64::NAN, 100.0, -1.0, 0.0, 1.0, -100.0, f64::NAN]); + + RecordBatch::try_from_iter(vec![ + ("i8", Arc::new(i8) as _), + ("i16", Arc::new(i16) as _), + ("i32", Arc::new(i32) as _), + ("i64", Arc::new(i64) as _), + ("u8", Arc::new(u8) as _), + ("u16", Arc::new(u16) as _), + ("u32", Arc::new(u32) as _), + ("u64", Arc::new(u64) as _), + ("f32", Arc::new(f32) as _), + ("f64", Arc::new(f64) as _), + ("f32_nan", Arc::new(f32_nan) as _), + ("f64_nan", Arc::new(f64_nan) as _), + ]) + .unwrap() +} + +fn make_utf8_batch(value: Vec>) -> RecordBatch { + let utf8 = StringArray::from(value.clone()); + let large_utf8 = LargeStringArray::from(value); + RecordBatch::try_from_iter(vec![ + ("utf8", Arc::new(utf8) as _), + ("large_utf8", Arc::new(large_utf8) as _), + ]) + .unwrap() +} + +fn make_utf8_view_batch(value: Vec>) -> RecordBatch { + let utf8_view = StringViewArray::from(value); + RecordBatch::try_from_iter(vec![("utf8_view", Arc::new(utf8_view) as _)]).unwrap() +} + +fn make_binary_view_batch(value: Vec>) -> RecordBatch { + let binary_view = BinaryViewArray::from(value); + RecordBatch::try_from_iter(vec![("binary_view", Arc::new(binary_view) as _)]).unwrap() +} + +fn make_dict_batch() -> RecordBatch { + let values = [ + Some("abc"), + Some("def"), + None, + Some("def"), + Some("abc"), + Some("fffff"), + Some("aaa"), + ]; + let dict_i8_array = DictionaryArray::::from_iter(values.iter().cloned()); + let dict_i32_array = DictionaryArray::::from_iter(values.iter().cloned()); + + // Dictionary array of integers + let int64_values = Int64Array::from(vec![0, -100, 100]); + let keys = Int8Array::from_iter([Some(0), Some(1), None, Some(0), Some(0), Some(2), Some(0)]); + let dict_i8_int_array = + DictionaryArray::::try_new(keys, Arc::new(int64_values)).unwrap(); + + RecordBatch::try_from_iter(vec![ + ("string_dict_i8", Arc::new(dict_i8_array) as _), + ("string_dict_i32", Arc::new(dict_i32_array) as _), + ("int_dict_i8", Arc::new(dict_i8_int_array) as _), + ]) + .unwrap() +} + +/// Create data batches for the given scenario. +/// `make_test_file_rg` uses the first batch to inference the schema of the file. +fn create_data_batch(scenario: Scenario) -> Vec { + match scenario { + Scenario::Boolean => { + vec![ + make_boolean_batch(vec![Some(true), Some(false), Some(true), Some(false), None]), + make_boolean_batch(vec![ + Some(false), + Some(false), + Some(false), + Some(false), + Some(false), + ]), + ] + } + Scenario::Timestamps => { + vec![ + make_timestamp_batch(TimeDelta::try_seconds(0).unwrap()), + make_timestamp_batch(TimeDelta::try_seconds(10).unwrap()), + make_timestamp_batch(TimeDelta::try_minutes(10).unwrap()), + make_timestamp_batch(TimeDelta::try_days(10).unwrap()), + ] + } + Scenario::Dates => { + vec![ + make_date_batch(TimeDelta::try_days(0).unwrap()), + make_date_batch(TimeDelta::try_days(10).unwrap()), + make_date_batch(TimeDelta::try_days(300).unwrap()), + make_date_batch(TimeDelta::try_days(3600).unwrap()), + ] + } + Scenario::Int => { + vec![ + make_int_batches(-5, 0), + make_int_batches(-4, 1), + make_int_batches(0, 5), + make_int_batches(5, 10), + ] + } + Scenario::Int32Range => { + vec![make_int32_range(0, 10), make_int32_range(200000, 300000)] + } + Scenario::UInt => { + vec![ + make_uint_batches(0, 5), + make_uint_batches(1, 6), + make_uint_batches(5, 10), + make_uint_batches(250, 255), + ] + } + Scenario::UInt32Range => { + vec![make_uint32_range(0, 10), make_uint32_range(200000, 300000)] + } + Scenario::NumericLimits => { + vec![make_numeric_limit_batch()] + } + Scenario::Float16 => { + vec![ + make_f16_batch( + vec![-5.0, -4.0, -3.0, -2.0, -1.0] + .into_iter() + .map(f16::from_f32) + .collect(), + ), + make_f16_batch( + vec![-4.0, -3.0, -2.0, -1.0, 0.0] + .into_iter() + .map(f16::from_f32) + .collect(), + ), + make_f16_batch( + vec![0.0, 1.0, 2.0, 3.0, 4.0] + .into_iter() + .map(f16::from_f32) + .collect(), + ), + make_f16_batch( + vec![5.0, 6.0, 7.0, 8.0, 9.0] + .into_iter() + .map(f16::from_f32) + .collect(), + ), + ] + } + Scenario::Float32 => { + vec![ + make_f32_batch(vec![-5.0, -4.0, -3.0, -2.0, -1.0]), + make_f32_batch(vec![-4.0, -3.0, -2.0, -1.0, 0.0]), + make_f32_batch(vec![0.0, 1.0, 2.0, 3.0, 4.0]), + make_f32_batch(vec![5.0, 6.0, 7.0, 8.0, 9.0]), + ] + } + Scenario::Float64 => { + vec![ + make_f64_batch(vec![-5.0, -4.0, -3.0, -2.0, -1.0]), + make_f64_batch(vec![-4.0, -3.0, -2.0, -1.0, 0.0]), + make_f64_batch(vec![0.0, 1.0, 2.0, 3.0, 4.0]), + make_f64_batch(vec![5.0, 6.0, 7.0, 8.0, 9.0]), + ] + } + Scenario::Decimal32 => { + // decimal record batch + vec![ + make_decimal32_batch(vec![100, 200, 300, 400, 600], 9, 2), + make_decimal32_batch(vec![-500, 100, 300, 400, 600], 9, 2), + make_decimal32_batch(vec![2000, 3000, 3000, 4000, 6000], 9, 2), + ] + } + Scenario::Decimal64 => { + // decimal record batch + vec![ + make_decimal64_batch(vec![100, 200, 300, 400, 600], 9, 2), + make_decimal64_batch(vec![-500, 100, 300, 400, 600], 9, 2), + make_decimal64_batch(vec![2000, 3000, 3000, 4000, 6000], 9, 2), + ] + } + Scenario::Decimal128 => { + // decimal record batch + vec![ + make_decimal128_batch(vec![100, 200, 300, 400, 600], 9, 2), + make_decimal128_batch(vec![-500, 100, 300, 400, 600], 9, 2), + make_decimal128_batch(vec![2000, 3000, 3000, 4000, 6000], 9, 2), + ] + } + Scenario::Decimal256 => { + // decimal256 record batch + vec![ + make_decimal256_batch( + vec![ + i256::from(100), + i256::from(200), + i256::from(300), + i256::from(400), + i256::from(600), + ], + 9, + 2, + ), + make_decimal256_batch( + vec![ + i256::from(-500), + i256::from(100), + i256::from(300), + i256::from(400), + i256::from(600), + ], + 9, + 2, + ), + make_decimal256_batch( + vec![ + i256::from(2000), + i256::from(3000), + i256::from(3000), + i256::from(4000), + i256::from(6000), + ], + 9, + 2, + ), + ] + } + Scenario::ByteArray => { + // frontends first, then backends. All in order, except frontends 4 and 7 + // are swapped to cause a statistics false positive on the 'fixed size' column. + vec![ + make_bytearray_batch( + "all frontends", + vec![ + "frontend one", + "frontend two", + "frontend three", + "frontend seven", + "frontend five", + ], + vec![ + b"frontend one", + b"frontend two", + b"frontend three", + b"frontend seven", + b"frontend five", + ], + vec![b"fe1", b"fe2", b"fe3", b"fe7", b"fe5"], + vec![ + b"frontend one", + b"frontend two", + b"frontend three", + b"frontend seven", + b"frontend five", + ], + ), + make_bytearray_batch( + "mixed", + vec![ + "frontend six", + "frontend four", + "backend one", + "backend two", + "backend three", + ], + vec![ + b"frontend six", + b"frontend four", + b"backend one", + b"backend two", + b"backend three", + ], + vec![b"fe6", b"fe4", b"be1", b"be2", b"be3"], + vec![ + b"frontend six", + b"frontend four", + b"backend one", + b"backend two", + b"backend three", + ], + ), + make_bytearray_batch( + "all backends", + vec![ + "backend four", + "backend five", + "backend six", + "backend seven", + "backend eight", + ], + vec![ + b"backend four", + b"backend five", + b"backend six", + b"backend seven", + b"backend eight", + ], + vec![b"be4", b"be5", b"be6", b"be7", b"be8"], + vec![ + b"backend four", + b"backend five", + b"backend six", + b"backend seven", + b"backend eight", + ], + ), + ] + } + Scenario::Dictionary => { + vec![make_dict_batch()] + } + Scenario::PeriodsInColumnNames => { + vec![ + // all frontend + make_names_batch( + "HTTP GET / DISPATCH", + vec!["frontend", "frontend", "frontend", "frontend", "frontend"], + ), + // both frontend and backend + make_names_batch( + "HTTP PUT / DISPATCH", + vec!["frontend", "frontend", "backend", "backend", "backend"], + ), + // all backend + make_names_batch( + "HTTP GET / DISPATCH", + vec!["backend", "backend", "backend", "backend", "backend"], + ), + ] + } + Scenario::StructArray => { + let struct_array_data = struct_array(vec![ + (Some(1), Some(6.0), Some(12.0)), + (Some(2), Some(8.5), None), + (None, Some(8.5), Some(14.0)), + ]); + + let schema = Arc::new(Schema::new(vec![Field::new( + "struct", + struct_array_data.data_type().clone(), + true, + )])); + vec![RecordBatch::try_new(schema, vec![struct_array_data]).unwrap()] + } + Scenario::Time32Second => { + vec![ + make_time32_batches(Scenario::Time32Second, vec![18506, 18507, 18508, 18509]), + make_time32_batches(Scenario::Time32Second, vec![18510, 18511, 18512, 18513]), + make_time32_batches(Scenario::Time32Second, vec![18514, 18515, 18516, 18517]), + make_time32_batches(Scenario::Time32Second, vec![18518, 18519, 18520, 18521]), + ] + } + Scenario::Time32Millisecond => { + vec![ + make_time32_batches( + Scenario::Time32Millisecond, + vec![3600000, 3600001, 3600002, 3600003], + ), + make_time32_batches( + Scenario::Time32Millisecond, + vec![3600004, 3600005, 3600006, 3600007], + ), + make_time32_batches( + Scenario::Time32Millisecond, + vec![3600008, 3600009, 3600010, 3600011], + ), + make_time32_batches( + Scenario::Time32Millisecond, + vec![3600012, 3600013, 3600014, 3600015], + ), + ] + } + Scenario::Time64Microsecond => { + vec![ + make_time64_batches( + Scenario::Time64Microsecond, + vec![1234567890123, 1234567890124, 1234567890125, 1234567890126], + ), + make_time64_batches( + Scenario::Time64Microsecond, + vec![1234567890127, 1234567890128, 1234567890129, 1234567890130], + ), + make_time64_batches( + Scenario::Time64Microsecond, + vec![1234567890131, 1234567890132, 1234567890133, 1234567890134], + ), + make_time64_batches( + Scenario::Time64Microsecond, + vec![1234567890135, 1234567890136, 1234567890137, 1234567890138], + ), + ] + } + Scenario::Time64Nanosecond => { + vec![ + make_time64_batches( + Scenario::Time64Nanosecond, + vec![ + 987654321012345, + 987654321012346, + 987654321012347, + 987654321012348, + ], + ), + make_time64_batches( + Scenario::Time64Nanosecond, + vec![ + 987654321012349, + 987654321012350, + 987654321012351, + 987654321012352, + ], + ), + make_time64_batches( + Scenario::Time64Nanosecond, + vec![ + 987654321012353, + 987654321012354, + 987654321012355, + 987654321012356, + ], + ), + make_time64_batches( + Scenario::Time64Nanosecond, + vec![ + 987654321012357, + 987654321012358, + 987654321012359, + 987654321012360, + ], + ), + ] + } + Scenario::UTF8 => { + vec![ + make_utf8_batch(vec![Some("a"), Some("b"), Some("c"), Some("d"), None]), + make_utf8_batch(vec![Some("e"), Some("f"), Some("g"), Some("h"), Some("i")]), + ] + } + Scenario::TruncatedUTF8 => { + // Make utf8 batch with strings longer than 64 bytes + // to check truncation of row group statistics + vec![ + make_utf8_batch(vec![ + Some(&("a".repeat(64) + "1")), + Some(&("b".repeat(64) + "2")), + Some(&("c".repeat(64) + "3")), + None, + Some(&("d".repeat(64) + "4")), + ]), + make_utf8_batch(vec![ + Some(&("e".repeat(64) + "5")), + Some(&("f".repeat(64) + "6")), + Some(&("g".repeat(64) + "7")), + Some(&("h".repeat(64) + "8")), + Some(&("i".repeat(64) + "9")), + ]), + make_utf8_batch(vec![ + Some("j"), + Some("k"), + Some(&("l".repeat(64) + "12")), + Some(&("m".repeat(64) + "13")), + Some(&("n".repeat(64) + "14")), + ]), + ] + } + Scenario::UTF8View => { + // Make utf8_view batch including string length <12 and >12 bytes + // as the internal representation of StringView is differed for strings + // shorter and longer than that length + vec![ + make_utf8_view_batch(vec![Some("a"), Some("b"), Some("c"), Some("d"), None]), + make_utf8_view_batch(vec![Some("a"), Some("e_longerthan12"), None, None, None]), + make_utf8_view_batch(vec![ + Some("e_longerthan12"), + Some("f_longerthan12"), + Some("g_longerthan12"), + Some("h_longerthan12"), + Some("i_longerthan12"), + ]), + ] + } + Scenario::BinaryView => { + vec![ + make_binary_view_batch(vec![Some(b"a"), Some(b"b"), Some(b"c"), Some(b"d"), None]), + make_binary_view_batch(vec![Some(b"a"), Some(b"e_longerthan12"), None, None, None]), + make_binary_view_batch(vec![ + Some(b"e_longerthan12"), + Some(b"f_longerthan12"), + Some(b"g_longerthan12"), + Some(b"h_longerthan12"), + Some(b"i_longerthan12"), + ]), + ] + } + } +} + +/// Create a test parquet file with various data types +async fn make_test_file_rg(scenario: Scenario, row_per_group: usize) -> NamedTempFile { + let mut output_file = tempfile::Builder::new() + .prefix("parquet_pruning") + .suffix(".parquet") + .tempfile() + .expect("tempfile creation"); + + let mut builder = WriterProperties::builder() + .set_max_row_group_row_count(Some(row_per_group)) + .set_bloom_filter_enabled(true) + .set_statistics_enabled(EnabledStatistics::Page); + if scenario.truncate_stats() { + // The same as default `column_index_truncate_length` to check both stats with one value + builder = builder.set_statistics_truncate_length(DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH); + } + let props = builder.build(); + + let batches = create_data_batch(scenario); + let schema = batches[0].schema(); + + let mut writer = ArrowWriter::try_new(&mut output_file, schema, Some(props)).unwrap(); + + for batch in batches { + writer.write(&batch).expect("writing batch"); + } + writer.close().unwrap(); + + output_file +} diff --git a/vendor/parquet/tests/arrow_reader/new_types.bin b/vendor/parquet/tests/arrow_reader/new_types.bin new file mode 100644 index 0000000..00d86fd Binary files /dev/null and b/vendor/parquet/tests/arrow_reader/new_types.bin differ diff --git a/vendor/parquet/tests/arrow_reader/predicate_cache.rs b/vendor/parquet/tests/arrow_reader/predicate_cache.rs new file mode 100644 index 0000000..4029b4e --- /dev/null +++ b/vendor/parquet/tests/arrow_reader/predicate_cache.rs @@ -0,0 +1,402 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Test for predicate cache in Parquet Arrow reader + +use super::io::TestReader; +use arrow::array::Array; +use arrow::array::ArrayRef; +use arrow::array::Int64Array; +use arrow::compute::and; +use arrow::compute::kernels::cmp::{gt, lt}; +use arrow_array::cast::AsArray; +use arrow_array::types::Int64Type; +use arrow_array::{RecordBatch, StringArray, StringViewArray, StructArray}; +use arrow_schema::{DataType, Field}; +use bytes::Bytes; +use futures::StreamExt; +use parquet::arrow::arrow_reader::metrics::ArrowReaderMetrics; +use parquet::arrow::arrow_reader::{ArrowPredicateFn, ArrowReaderOptions, RowFilter}; +use parquet::arrow::arrow_reader::{ArrowReaderBuilder, ParquetRecordBatchReaderBuilder}; +use parquet::arrow::{ArrowWriter, ParquetRecordBatchStreamBuilder, ProjectionMask}; +use parquet::file::properties::WriterProperties; +use std::sync::Arc; +use std::sync::LazyLock; + +#[tokio::test] +async fn test_default_read() { + // The cache is not used without predicates, so we expect 0 records read from cache + let test = ParquetPredicateCacheTest::new().with_expected_records_read_from_cache(0); + let sync_builder = test.sync_builder(); + test.run_sync(sync_builder); + let async_builder = test.async_builder().await; + test.run_async(async_builder).await; +} + +#[tokio::test] +async fn test_async_cache_with_filters() { + let test = ParquetPredicateCacheTest::new().with_expected_records_read_from_cache(49); + let async_builder = test.async_builder().await.add_project_ab_and_filter_b(); + test.run_async(async_builder).await; +} + +#[tokio::test] +async fn test_sync_cache_with_filters() { + let test = ParquetPredicateCacheTest::new() + // The sync reader does not use the cache. See https://github.com/apache/arrow-rs/issues/8000 + .with_expected_records_read_from_cache(0); + + let sync_builder = test.sync_builder().add_project_ab_and_filter_b(); + test.run_sync(sync_builder); +} + +#[tokio::test] +async fn test_cache_disabled_with_filters() { + // expect no records to be read from cache, because the cache is disabled + let test = ParquetPredicateCacheTest::new().with_expected_records_read_from_cache(0); + let sync_builder = test + .sync_builder() + .with_max_predicate_cache_size(0) + .add_project_ab_and_filter_b(); + test.run_sync(sync_builder); + + let async_builder = test + .async_builder() + .await + .with_max_predicate_cache_size(0) + .add_project_ab_and_filter_b(); + test.run_async(async_builder).await; +} + +#[tokio::test] +async fn test_cache_projection_excludes_nested_columns() { + let test = ParquetPredicateCacheTest::new_nested().with_expected_records_read_from_cache(0); + + let sync_builder = test.sync_builder().add_nested_filter(); + test.run_sync(sync_builder); + + let async_builder = test.async_builder().await.add_nested_filter(); + test.run_async(async_builder).await; +} + +/// Regression: cache must match no-cache for a nullable single-leaf struct. +#[tokio::test] +async fn test_async_predicate_on_single_leaf_nullable_struct() { + // Rows: b = NULL, then b.aa = "hello". + let aa: StringArray = StringArray::from(vec!["padding", "hello"]); + let nulls = arrow_buffer::NullBuffer::from(vec![false, true]); + let b = StructArray::new( + vec![Arc::new(Field::new("aa", DataType::Utf8, false))].into(), + vec![Arc::new(aa) as ArrayRef], + Some(nulls), + ); + let input_batch = RecordBatch::try_from_iter([("b", Arc::new(b) as ArrayRef)]).unwrap(); + + let mut output = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut output, input_batch.schema(), None).unwrap(); + writer.write(&input_batch).unwrap(); + writer.close().unwrap(); + let bytes = Bytes::from(output); + + // Since `aa` is required, `b.aa IS NULL` means `b` is NULL. + let build_is_null_filter = |schema_descr: &parquet::schema::types::SchemaDescPtr| -> RowFilter { + let mask = ProjectionMask::leaves(schema_descr, vec![0]); + let predicate = ArrowPredicateFn::new(mask.clone(), |batch: RecordBatch| { + let struct_arr = batch.column(0).as_struct(); + let leaf = struct_arr.column(0); + Ok((0..batch.num_rows()) + .map(|i| struct_arr.is_null(i) || leaf.is_null(i)) + .collect::()) + }); + RowFilter::new(vec![Box::new(predicate)]) + }; + + // Default cache. + let reader = TestReader::new(bytes.clone()); + let async_builder = + ParquetRecordBatchStreamBuilder::new_with_options(reader, ArrowReaderOptions::default()) + .await + .unwrap(); + let schema_descr = async_builder.metadata().file_metadata().schema_descr_ptr(); + let async_builder = async_builder + .with_projection(ProjectionMask::leaves(&schema_descr, vec![0])) + .with_row_filter(build_is_null_filter(&schema_descr)); + let mut stream = async_builder.build().unwrap(); + let mut row_count_cached = 0; + while let Some(batch) = stream.next().await { + row_count_cached += batch.unwrap().num_rows(); + } + + // Cache disabled. + let reader = TestReader::new(bytes.clone()); + let async_builder = + ParquetRecordBatchStreamBuilder::new_with_options(reader, ArrowReaderOptions::default()) + .await + .unwrap(); + let schema_descr = async_builder.metadata().file_metadata().schema_descr_ptr(); + let async_builder = async_builder + .with_projection(ProjectionMask::leaves(&schema_descr, vec![0])) + .with_row_filter(build_is_null_filter(&schema_descr)) + .with_max_predicate_cache_size(0); + let mut stream = async_builder.build().unwrap(); + let mut row_count_uncached = 0; + while let Some(batch) = stream.next().await { + row_count_uncached += batch.unwrap().num_rows(); + } + + assert_eq!( + row_count_uncached, 1, + "control: with cache disabled the predicate must match exactly one row (parent NULL)" + ); + assert_eq!( + row_count_cached, row_count_uncached, + "cached reader must match uncached reader; \ + got {row_count_cached} cached vs {row_count_uncached} uncached. \ + single-leaf struct roots must stay out of the cache." + ); +} + +// -- Begin test infrastructure -- + +/// A test parquet file +struct ParquetPredicateCacheTest { + bytes: Bytes, + expected_records_read_from_cache: usize, +} +impl ParquetPredicateCacheTest { + /// Create a new `TestParquetFile` with: + /// 3 columns: "a", "b", "c" + /// + /// 2 row groups, each with 200 rows + /// each data page has 100 rows + /// + /// Values of column "a" are 0..399 + /// Values of column "b" are 400..799 + /// Values of column "c" are alternating strings of length 12 and longer + fn new() -> Self { + Self { + bytes: TEST_FILE_DATA.clone(), + expected_records_read_from_cache: 0, + } + } + + /// Create a new `TestParquetFile` with + /// 2 columns: + /// + /// * string column `a` + /// * nested struct column `b { aa, bb }` + fn new_nested() -> Self { + Self { + bytes: NESTED_TEST_FILE_DATA.clone(), + expected_records_read_from_cache: 0, + } + } + + /// Set the expected number of records read from the cache + fn with_expected_records_read_from_cache( + mut self, + expected_records_read_from_cache: usize, + ) -> Self { + self.expected_records_read_from_cache = expected_records_read_from_cache; + self + } + + /// Return a [`ParquetRecordBatchReaderBuilder`] for reading this file + fn sync_builder(&self) -> ParquetRecordBatchReaderBuilder { + let reader = self.bytes.clone(); + ParquetRecordBatchReaderBuilder::try_new_with_options(reader, ArrowReaderOptions::default()) + .expect("ParquetRecordBatchReaderBuilder") + } + + /// Return a [`ParquetRecordBatchReaderBuilder`] for reading this file + async fn async_builder(&self) -> ParquetRecordBatchStreamBuilder { + let reader = TestReader::new(self.bytes.clone()); + ParquetRecordBatchStreamBuilder::new_with_options(reader, ArrowReaderOptions::default()) + .await + .unwrap() + } + + /// Build the reader from the specified builder, reading all batches from it, + /// and asserts the + fn run_sync(&self, builder: ParquetRecordBatchReaderBuilder) { + let metrics = ArrowReaderMetrics::enabled(); + + let reader = builder.with_metrics(metrics.clone()).build().unwrap(); + for batch in reader { + match batch { + Ok(_) => {} + Err(e) => panic!("Error reading batch: {e}"), + } + } + self.verify_metrics(metrics) + } + + /// Build the reader from the specified builder, reading all batches from it, + /// and asserts the + async fn run_async(&self, builder: ParquetRecordBatchStreamBuilder) { + let metrics = ArrowReaderMetrics::enabled(); + + let mut stream = builder.with_metrics(metrics.clone()).build().unwrap(); + while let Some(batch) = stream.next().await { + match batch { + Ok(_) => {} + Err(e) => panic!("Error reading batch: {e}"), + } + } + self.verify_metrics(metrics) + } + + fn verify_metrics(&self, metrics: ArrowReaderMetrics) { + let Self { + bytes: _, + expected_records_read_from_cache, + } = self; + + let read_from_cache = metrics + .records_read_from_cache() + .expect("Metrics enabled, so should have metrics"); + + assert_eq!( + &read_from_cache, expected_records_read_from_cache, + "Expected {expected_records_read_from_cache} records read from cache, but got {read_from_cache}" + ); + } +} + +/// Create a parquet file in memory for testing. See [`test_file`] for details. +static TEST_FILE_DATA: LazyLock = LazyLock::new(|| { + // Input batch has 400 rows, with 3 columns: "a", "b", "c" + // Note c is a different types (so the data page sizes will be different) + let a: ArrayRef = Arc::new(Int64Array::from_iter_values(0..400)); + let b: ArrayRef = Arc::new(Int64Array::from_iter_values(400..800)); + let c: ArrayRef = Arc::new(StringViewArray::from_iter_values((0..400).map(|i| { + if i % 2 == 0 { + format!("string_{i}") + } else { + format!("A string larger than 12 bytes and thus not inlined {i}") + } + }))); + + let input_batch = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + let mut output = Vec::new(); + + let writer_options = WriterProperties::builder() + .set_max_row_group_row_count(Some(200)) + .set_data_page_row_count_limit(100) + .build(); + let mut writer = + ArrowWriter::try_new(&mut output, input_batch.schema(), Some(writer_options)).unwrap(); + + // since the limits are only enforced on batch boundaries, write the input + // batch in chunks of 50 + let mut row_remain = input_batch.num_rows(); + while row_remain > 0 { + let chunk_size = row_remain.min(50); + let chunk = input_batch.slice(input_batch.num_rows() - row_remain, chunk_size); + writer.write(&chunk).unwrap(); + row_remain -= chunk_size; + } + writer.close().unwrap(); + Bytes::from(output) +}); + +/// Build a ParquetFile with a +/// +/// * string column `a` +/// * nested struct column `b { aa, bb }` +static NESTED_TEST_FILE_DATA: LazyLock = LazyLock::new(|| { + const NUM_ROWS: usize = 100; + let a: StringArray = (0..NUM_ROWS).map(|i| Some(format!("r{i}"))).collect(); + + let aa: StringArray = (0..NUM_ROWS).map(|i| Some(format!("v{i}"))).collect(); + let bb: StringArray = (0..NUM_ROWS).map(|i| Some(format!("w{i}"))).collect(); + let b = StructArray::from(vec![ + ( + Arc::new(Field::new("aa", DataType::Utf8, true)), + Arc::new(aa) as ArrayRef, + ), + ( + Arc::new(Field::new("bb", DataType::Utf8, true)), + Arc::new(bb) as ArrayRef, + ), + ]); + + let input_batch = RecordBatch::try_from_iter([ + ("a", Arc::new(a) as ArrayRef), + ("b", Arc::new(b) as ArrayRef), + ]) + .unwrap(); + + let mut output = Vec::new(); + let writer_options = None; + let mut writer = + ArrowWriter::try_new(&mut output, input_batch.schema(), writer_options).unwrap(); + writer.write(&input_batch).unwrap(); + writer.close().unwrap(); + Bytes::from(output) +}); + +trait ArrowReaderBuilderExt { + /// Applies the following: + /// 1. a projection selecting the "a" and "b" column + /// 2. a row_filter applied to "b": 575 < "b" < 625 (select 1 data page from each row group) + fn add_project_ab_and_filter_b(self) -> Self; + + /// Adds a row filter that projects the nested leaf column "b.aa" and + /// returns true for all rows. + fn add_nested_filter(self) -> Self; +} + +impl ArrowReaderBuilderExt for ArrowReaderBuilder { + fn add_project_ab_and_filter_b(self) -> Self { + let schema_descr = self.metadata().file_metadata().schema_descr_ptr(); + + // "b" > 575 and "b" < 625 + let row_filter = ArrowPredicateFn::new( + ProjectionMask::columns(&schema_descr, ["b"]), + |batch: RecordBatch| { + let scalar_575 = Int64Array::new_scalar(575); + let scalar_625 = Int64Array::new_scalar(625); + let column = batch.column(0).as_primitive::(); + and(>(column, &scalar_575)?, <(column, &scalar_625)?) + }, + ); + + self.with_projection(ProjectionMask::columns(&schema_descr, ["a", "b"])) + .with_row_filter(RowFilter::new(vec![Box::new(row_filter)])) + } + + fn add_nested_filter(self) -> Self { + let schema_descr = self.metadata().file_metadata().schema_descr_ptr(); + + // Build a RowFilter whose predicate projects a leaf under the nested root `b` + // Leaf indices are depth-first; with schema [a, b.aa, b.bb] we pick index 1 (b.aa) + let nested_leaf_mask = ProjectionMask::leaves(&schema_descr, vec![1]); + + let always_true = ArrowPredicateFn::new(nested_leaf_mask.clone(), |batch: RecordBatch| { + Ok(arrow_array::BooleanArray::from(vec![ + true; + batch.num_rows() + ])) + }); + let row_filter = RowFilter::new(vec![Box::new(always_true)]); + + self.with_projection(nested_leaf_mask) + .with_row_filter(row_filter) + } +} diff --git a/vendor/parquet/tests/arrow_reader/row_filter/async.rs b/vendor/parquet/tests/arrow_reader/row_filter/async.rs new file mode 100644 index 0000000..3adbb05 --- /dev/null +++ b/vendor/parquet/tests/arrow_reader/row_filter/async.rs @@ -0,0 +1,1026 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::io::TestReader; +use std::sync::{Arc, Mutex}; + +use arrow::{ + array::AsArray, + compute::{concat_batches, kernels::cmp::eq, or}, + datatypes::{Int32Type, TimestampNanosecondType}, +}; +use arrow_array::{ + ArrayRef, BooleanArray, Int8Array, Int32Array, Int64Array, RecordBatch, Scalar, StringArray, + StructArray, +}; +use arrow_schema::{DataType, Field, Schema}; +use bytes::Bytes; +use futures::TryStreamExt; +use parquet::{ + arrow::{ + ArrowWriter, ParquetRecordBatchStreamBuilder, ProjectionMask, + arrow_reader::{ + ArrowPredicateFn, ArrowReaderOptions, RowFilter, RowSelection, RowSelectionPolicy, + RowSelector, metrics::ArrowReaderMetrics, + }, + }, + file::{ + metadata::{PageIndexPolicy, ParquetMetaDataReader}, + properties::WriterProperties, + }, +}; + +fn make_two_column_i64_file(values: &[i64], rows_per_page: usize) -> Bytes { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("value", DataType::Int64, false), + ])); + let keys = Int64Array::from(values.to_vec()); + let values = Int64Array::from(values.to_vec()); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(keys) as ArrayRef, Arc::new(values) as ArrayRef], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_write_batch_size(rows_per_page) + .set_data_page_row_count_limit(rows_per_page) + .build(); + let mut buffer = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut buffer, schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + Bytes::from(buffer) +} + +#[tokio::test] +async fn test_row_filter_full_page_skip_is_handled_async() { + let first_value: i64 = 1111; + let last_value: i64 = 9999; + let num_rows: usize = 12; + + // build data with row selection average length 4 + // The result would be (1111 XXXX) ... (4 page in the middle)... (XXXX 9999) + // The Row Selection would be [1111, (skip 10), 9999] + let mut int_values: Vec = (0..num_rows as i64).collect(); + int_values[0] = first_value; + int_values[num_rows - 1] = last_value; + let data = make_two_column_i64_file(&int_values, 2); + + let builder = ParquetRecordBatchStreamBuilder::new_with_options( + TestReader::new(data.clone()), + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required), + ) + .await + .unwrap(); + let schema = builder.parquet_schema().clone(); + let filter_mask = ProjectionMask::leaves(&schema, [0]); + + for policy in [ + RowSelectionPolicy::Auto { threshold: 32 }, + RowSelectionPolicy::Mask, + ] { + for max_predicate_cache_size in [None, Some(0)] { + for with_initial_selection in [false, true] { + let predicate_batch_sizes = Arc::new(Mutex::new(Vec::new())); + let observed_batch_sizes = Arc::clone(&predicate_batch_sizes); + let predicate = + ArrowPredicateFn::new(filter_mask.clone(), move |batch: RecordBatch| { + observed_batch_sizes.lock().unwrap().push(batch.num_rows()); + let column = batch.column(0); + let match_first = eq(column, &Int64Array::new_scalar(first_value))?; + let match_second = eq(column, &Int64Array::new_scalar(last_value))?; + or(&match_first, &match_second) + }); + + let stream = ParquetRecordBatchStreamBuilder::new_with_options( + TestReader::new(data.clone()), + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required), + ) + .await + .unwrap() + .with_row_filter(RowFilter::new(vec![Box::new(predicate)])) + .with_batch_size(12) + .with_row_selection_policy(policy); + let stream = if with_initial_selection { + stream.with_row_selection(RowSelection::from(vec![ + RowSelector::select(1), + RowSelector::skip(10), + RowSelector::select(1), + ])) + } else { + stream + }; + let stream = match max_predicate_cache_size { + Some(size) => stream.with_max_predicate_cache_size(size), + None => stream, + } + .build() + .unwrap(); + + let schema = stream.schema().clone(); + let batches: Vec<_> = stream.try_collect().await.unwrap(); + let batch_sizes = batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(); + assert_eq!( + batch_sizes, + vec![2], + "policy={policy:?}, cache={max_predicate_cache_size:?}, initial={with_initial_selection}" + ); + assert_eq!( + *predicate_batch_sizes.lock().unwrap(), + if with_initial_selection { + vec![2] + } else { + vec![12] + }, + "policy={policy:?}, cache={max_predicate_cache_size:?}, initial={with_initial_selection}" + ); + + let result = concat_batches(&schema, &batches).unwrap(); + assert_eq!(result.num_rows(), 2); + assert_eq!( + result + .column(0) + .as_any() + .downcast_ref::() + .unwrap(), + &Int64Array::from(vec![first_value, last_value]) + ); + } + } + } +} + +#[tokio::test] +async fn test_cached_mask_reads_sparse_pages_without_error() { + let values = (0..60).collect::>(); + let data = make_two_column_i64_file(&values, 20); + + for policy in [ + RowSelectionPolicy::Auto { threshold: 32 }, + RowSelectionPolicy::Mask, + ] { + let metrics = ArrowReaderMetrics::enabled(); + let builder = ParquetRecordBatchStreamBuilder::new_with_options( + TestReader::new(data.clone()), + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required), + ) + .await + .unwrap(); + let schema = builder.parquet_schema().clone(); + let projection = ProjectionMask::leaves(&schema, [0]); + let page_first_rows = builder.metadata().offset_index().unwrap()[0][0] + .page_locations() + .iter() + .map(|page| page.first_row_index) + .collect::>(); + assert_eq!(page_first_rows, vec![0, 20, 40]); + + let predicate = ArrowPredicateFn::new(projection.clone(), |batch: RecordBatch| { + Ok(BooleanArray::from(vec![true; batch.num_rows()])) + }); + // Extending the first mask chunk to the 20-row page boundary would make + // the 8-row cache batch at rows 16..24 cross into the unloaded middle page. + let stream = builder + .with_projection(projection) + .with_row_filter(RowFilter::new(vec![Box::new(predicate)])) + .with_row_selection(RowSelection::from(vec![ + RowSelector::select(1), + RowSelector::skip(39), + RowSelector::select(1), + ])) + .with_batch_size(8) + .with_max_predicate_cache_size(1024) + .with_row_selection_policy(policy) + .with_metrics(metrics.clone()) + .build() + .unwrap(); + + let output_schema = stream.schema().clone(); + let batches: Vec = stream.try_collect().await.unwrap(); + let output = concat_batches(&output_schema, &batches).unwrap(); + assert_eq!( + output + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[0, 40], + "policy={policy:?}" + ); + assert!( + metrics + .records_read_from_cache() + .expect("metrics are enabled") + > 0, + "predicate cache was not exercised for policy={policy:?}" + ); + } +} + +#[tokio::test] +async fn test_mask_coalesces_loaded_ranges_to_batch_size() { + let values = (0..12).collect::>(); + let data = make_two_column_i64_file(&values, 2); + let stream = ParquetRecordBatchStreamBuilder::new_with_options( + TestReader::new(data), + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required), + ) + .await + .unwrap() + .with_row_selection(RowSelection::from(vec![ + RowSelector::select(1), + RowSelector::skip(4), + RowSelector::select(1), + RowSelector::skip(5), + RowSelector::select(1), + ])) + .with_batch_size(2) + .with_row_selection_policy(RowSelectionPolicy::Mask) + .build() + .unwrap(); + + let schema = stream.schema().clone(); + let batches: Vec = stream.try_collect().await.unwrap(); + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![2, 1] + ); + + let batch = concat_batches(&schema, &batches).unwrap(); + assert_eq!( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[0, 5, 11] + ); +} + +#[tokio::test] +async fn test_mask_stops_after_trailing_skipped_pages() { + let values = (0..8).collect::>(); + let data = make_two_column_i64_file(&values, 2); + let stream = ParquetRecordBatchStreamBuilder::new_with_options( + TestReader::new(data), + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required), + ) + .await + .unwrap() + .with_row_selection(RowSelection::from(vec![ + RowSelector::select(1), + RowSelector::skip(4), + RowSelector::select(1), + RowSelector::skip(2), + ])) + .with_batch_size(8) + .with_row_selection_policy(RowSelectionPolicy::Mask) + .build() + .unwrap(); + + let batches: Vec = stream.try_collect().await.unwrap(); + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![2] + ); + assert_eq!( + batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[0, 5] + ); +} + +#[tokio::test] +async fn test_mask_nested_projection_with_different_page_boundaries() { + use arrow_array::{ + Array, UInt32Array, + builder::{Int32Builder, ListBuilder, StringBuilder, StructBuilder}, + }; + use arrow_schema::Fields; + use parquet::file::properties::WriterVersion; + + let num_rows = 14usize; + let long_prefix = "x".repeat(500); + let struct_fields = Fields::from(vec![ + Field::new("x", DataType::Int32, false), + Field::new("y", DataType::Utf8, false), + ]); + let mut list_builder = ListBuilder::new(StructBuilder::from_fields(struct_fields, 0)); + for row in 0..num_rows { + let values = list_builder.values(); + for item in 0..if row % 2 == 0 { 2 } else { 3 } { + values + .field_builder::(0) + .unwrap() + .append_value(row as i32 * 10 + item); + values + .field_builder::(1) + .unwrap() + .append_value(format!("{long_prefix}_{row}_{item}")); + values.append(true); + } + list_builder.append(true); + } + let values = list_builder.finish(); + let schema = Arc::new(Schema::new(vec![Field::new( + "values", + values.data_type().clone(), + false, + )])); + let input = RecordBatch::try_new(schema.clone(), vec![Arc::new(values) as ArrayRef]).unwrap(); + + // The fixed-width leaf fits in one page, while the long string leaf flushes + // frequently and therefore has different page boundaries. + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_2_0) + .set_write_batch_size(2) + .set_dictionary_enabled(false) + .set_data_page_size_limit(512) + .build(); + let mut buffer = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut buffer, schema, Some(props)).unwrap(); + writer.write(&input).unwrap(); + writer.close().unwrap(); + + let builder = ParquetRecordBatchStreamBuilder::new_with_options( + TestReader::new(Bytes::from(buffer)), + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required), + ) + .await + .unwrap(); + let page_first_rows = builder.metadata().offset_index().unwrap()[0] + .iter() + .map(|column| { + column + .page_locations() + .iter() + .map(|page| page.first_row_index) + .collect::>() + }) + .collect::>(); + assert_eq!(page_first_rows.len(), 2); + assert_eq!(page_first_rows[0], vec![0]); + assert_ne!(page_first_rows[0], page_first_rows[1]); + + let variable_width_pages = &page_first_rows[1]; + assert_eq!(variable_width_pages.first(), Some(&0)); + let selected_pages = [0, 6, 13] + .map(|row| variable_width_pages.partition_point(|first_row| *first_row <= row) - 1); + assert!( + selected_pages + .windows(2) + .all(|pages| pages[1] > pages[0] + 1), + "selected rows must leave unloaded pages between them: {page_first_rows:?}" + ); + + let stream = builder + .with_row_selection(RowSelection::from(vec![ + RowSelector::select(1), + RowSelector::skip(5), + RowSelector::select(1), + RowSelector::skip(6), + RowSelector::select(1), + ])) + .with_batch_size(2) + .with_row_selection_policy(RowSelectionPolicy::Mask) + .build() + .unwrap(); + + let output_schema = stream.schema().clone(); + let batches: Vec = stream.try_collect().await.unwrap(); + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![2, 1] + ); + + let output = concat_batches(&output_schema, &batches).unwrap(); + let expected = arrow::compute::take( + input.column(0).as_ref(), + &UInt32Array::from(vec![0, 6, 13]), + None, + ) + .unwrap(); + assert_eq!(output.column(0).to_data(), expected.to_data()); +} + +#[tokio::test] +async fn test_row_filter() { + let a = StringArray::from_iter_values(["a", "b", "b", "b", "c", "c"]); + let b = StringArray::from_iter_values(["1", "2", "3", "4", "5", "6"]); + let data = RecordBatch::try_from_iter([ + ("a", Arc::new(a) as ArrayRef), + ("b", Arc::new(b) as ArrayRef), + ]) + .unwrap(); + + let mut buf = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buf, data.schema(), None).unwrap(); + writer.write(&data).unwrap(); + writer.close().unwrap(); + + let data: Bytes = buf.into(); + let metadata = ParquetMetaDataReader::new() + .parse_and_finish(&data) + .unwrap(); + let parquet_schema = metadata.file_metadata().schema_descr_ptr(); + + let test = TestReader::new(data); + let requests = test.requests(); + + let a_scalar = StringArray::from_iter_values(["b"]); + let a_filter = ArrowPredicateFn::new( + ProjectionMask::leaves(&parquet_schema, vec![0]), + move |batch| eq(batch.column(0), &Scalar::new(&a_scalar)), + ); + + let filter = RowFilter::new(vec![Box::new(a_filter)]); + + let mask = ProjectionMask::leaves(&parquet_schema, vec![0, 1]); + let stream = ParquetRecordBatchStreamBuilder::new(test) + .await + .unwrap() + .with_projection(mask.clone()) + .with_batch_size(1024) + .with_row_filter(filter) + .build() + .unwrap(); + + let batches: Vec<_> = stream.try_collect().await.unwrap(); + assert_eq!(batches.len(), 1); + + let batch = &batches[0]; + assert_eq!(batch.num_columns(), 2); + + // Filter should have kept only rows with "b" in column 0 + assert_eq!( + batch.column(0).as_ref(), + &StringArray::from_iter_values(["b", "b", "b"]) + ); + assert_eq!( + batch.column(1).as_ref(), + &StringArray::from_iter_values(["2", "3", "4"]) + ); + + // Should only have made 2 requests: + // * First request fetches data for evaluating the predicate + // * Second request fetches data for evaluating the projection + assert_eq!(requests.lock().unwrap().len(), 2); +} + +#[tokio::test] +async fn test_two_row_filters() { + let a = StringArray::from_iter_values(["a", "b", "b", "b", "c", "c"]); + let b = StringArray::from_iter_values(["1", "2", "3", "4", "5", "6"]); + let c = Int32Array::from_iter(0..6); + let data = RecordBatch::try_from_iter([ + ("a", Arc::new(a) as ArrayRef), + ("b", Arc::new(b) as ArrayRef), + ("c", Arc::new(c) as ArrayRef), + ]) + .unwrap(); + + let mut buf = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buf, data.schema(), None).unwrap(); + writer.write(&data).unwrap(); + writer.close().unwrap(); + + let data: Bytes = buf.into(); + let metadata = ParquetMetaDataReader::new() + .parse_and_finish(&data) + .unwrap(); + let parquet_schema = metadata.file_metadata().schema_descr_ptr(); + + let test = TestReader::new(data); + let requests = test.requests(); + + let a_scalar = StringArray::from_iter_values(["b"]); + let a_filter = ArrowPredicateFn::new( + ProjectionMask::leaves(&parquet_schema, vec![0]), + move |batch| eq(batch.column(0), &Scalar::new(&a_scalar)), + ); + + let b_scalar = StringArray::from_iter_values(["4"]); + let b_filter = ArrowPredicateFn::new( + ProjectionMask::leaves(&parquet_schema, vec![1]), + move |batch| eq(batch.column(0), &Scalar::new(&b_scalar)), + ); + + let filter = RowFilter::new(vec![Box::new(a_filter), Box::new(b_filter)]); + + let mask = ProjectionMask::leaves(&parquet_schema, vec![0, 2]); + let stream = ParquetRecordBatchStreamBuilder::new(test) + .await + .unwrap() + .with_projection(mask.clone()) + .with_batch_size(1024) + .with_row_filter(filter) + .build() + .unwrap(); + + let batches: Vec<_> = stream.try_collect().await.unwrap(); + assert_eq!(batches.len(), 1); + + let batch = &batches[0]; + assert_eq!(batch.num_rows(), 1); + assert_eq!(batch.num_columns(), 2); + + let col = batch.column(0); + let val = col.as_any().downcast_ref::().unwrap().value(0); + assert_eq!(val, "b"); + + let col = batch.column(1); + let val = col.as_any().downcast_ref::().unwrap().value(0); + assert_eq!(val, 3); + + // Should only have made 3 requests + // * First request fetches data for evaluating the first predicate + // * Second request fetches data for evaluating the second predicate + // * Third request fetches data for evaluating the projection + assert_eq!(requests.lock().unwrap().len(), 3); +} + +#[tokio::test] +async fn test_row_filter_with_index() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet"); + let data = Bytes::from(std::fs::read(path).unwrap()); + + let metadata = ParquetMetaDataReader::new() + .parse_and_finish(&data) + .unwrap(); + let parquet_schema = metadata.file_metadata().schema_descr_ptr(); + + assert_eq!(metadata.num_row_groups(), 1); + + let async_reader = TestReader::new(data.clone()); + + let a_filter = + ArrowPredicateFn::new(ProjectionMask::leaves(&parquet_schema, vec![1]), |batch| { + Ok(batch.column(0).as_boolean().clone()) + }); + + let b_scalar = Int8Array::from(vec![2]); + let b_filter = ArrowPredicateFn::new( + ProjectionMask::leaves(&parquet_schema, vec![2]), + move |batch| eq(batch.column(0), &Scalar::new(&b_scalar)), + ); + + let filter = RowFilter::new(vec![Box::new(a_filter), Box::new(b_filter)]); + + let mask = ProjectionMask::leaves(&parquet_schema, vec![0, 2]); + + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let stream = ParquetRecordBatchStreamBuilder::new_with_options(async_reader, options) + .await + .unwrap() + .with_projection(mask.clone()) + .with_batch_size(1024) + .with_row_filter(filter) + .build() + .unwrap(); + + let batches: Vec = stream.try_collect().await.unwrap(); + + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + + assert_eq!(total_rows, 730); +} + +#[tokio::test] +async fn test_row_filter_nested() { + let a = StringArray::from_iter_values(["a", "b", "b", "b", "c", "c"]); + let b = StructArray::from(vec![ + ( + Arc::new(Field::new("aa", DataType::Utf8, true)), + Arc::new(StringArray::from(vec!["a", "b", "b", "b", "c", "c"])) as ArrayRef, + ), + ( + Arc::new(Field::new("bb", DataType::Utf8, true)), + Arc::new(StringArray::from(vec!["1", "2", "3", "4", "5", "6"])) as ArrayRef, + ), + ]); + let c = Int32Array::from_iter(0..6); + let data = RecordBatch::try_from_iter([ + ("a", Arc::new(a) as ArrayRef), + ("b", Arc::new(b) as ArrayRef), + ("c", Arc::new(c) as ArrayRef), + ]) + .unwrap(); + + let mut buf = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buf, data.schema(), None).unwrap(); + writer.write(&data).unwrap(); + writer.close().unwrap(); + + let data: Bytes = buf.into(); + let metadata = ParquetMetaDataReader::new() + .parse_and_finish(&data) + .unwrap(); + let parquet_schema = metadata.file_metadata().schema_descr_ptr(); + + let test = TestReader::new(data); + let requests = test.requests(); + + let a_scalar = StringArray::from_iter_values(["b"]); + let a_filter = ArrowPredicateFn::new( + ProjectionMask::leaves(&parquet_schema, vec![0]), + move |batch| eq(batch.column(0), &Scalar::new(&a_scalar)), + ); + + let b_scalar = StringArray::from_iter_values(["4"]); + let b_filter = ArrowPredicateFn::new( + ProjectionMask::leaves(&parquet_schema, vec![2]), + move |batch| { + // Filter on the second element of the struct. + let struct_array = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + eq(struct_array.column(0), &Scalar::new(&b_scalar)) + }, + ); + + let filter = RowFilter::new(vec![Box::new(a_filter), Box::new(b_filter)]); + + let mask = ProjectionMask::leaves(&parquet_schema, vec![0, 3]); + let stream = ParquetRecordBatchStreamBuilder::new(test) + .await + .unwrap() + .with_projection(mask.clone()) + .with_batch_size(1024) + .with_row_filter(filter) + .build() + .unwrap(); + + let batches: Vec<_> = stream.try_collect().await.unwrap(); + assert_eq!(batches.len(), 1); + + let batch = &batches[0]; + assert_eq!(batch.num_rows(), 1); + assert_eq!(batch.num_columns(), 2); + + let col = batch.column(0); + let val = col.as_any().downcast_ref::().unwrap().value(0); + assert_eq!(val, "b"); + + let col = batch.column(1); + let val = col.as_any().downcast_ref::().unwrap().value(0); + assert_eq!(val, 3); + + // Should only have made 3 requests + // * First request fetches data for evaluating the first predicate + // * Second request fetches data for evaluating the second predicate + // * Third request fetches data for evaluating the projection + assert_eq!(requests.lock().unwrap().len(), 3); +} + +/// Regression test for adaptive predicate pushdown attempting to read skipped pages. +/// Related issue: https://github.com/apache/arrow-rs/issues/9239 +#[tokio::test] +async fn test_predicate_pushdown_with_skipped_pages() { + use arrow_array::TimestampNanosecondArray; + use arrow_schema::TimeUnit; + + // Time range constants + const TIME_IN_RANGE_START: i64 = 1_704_092_400_000_000_000; + const TIME_IN_RANGE_END: i64 = 1_704_110_400_000_000_000; + const TIME_BEFORE_RANGE: i64 = 1_704_078_000_000_000_000; + + // Create test data: 2 row groups, 300 rows each + // "tag" column: 'a', 'b', 'c' (100 rows each, sorted) + // "time" column: alternating in-range/out-of-range timestamps + let schema = Arc::new(Schema::new(vec![ + Field::new( + "time", + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ), + Field::new("tag", DataType::Utf8, false), + ])); + + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(300)) + .set_data_page_row_count_limit(33) + .build(); + + let mut buffer = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(), Some(props)).unwrap(); + + // Write 2 row groups + for _ in 0..2 { + for (tag_idx, tag) in ["a", "b", "c"].iter().enumerate() { + let times: Vec = (0..100) + .map(|j| { + let row_idx = tag_idx * 100 + j; + if row_idx % 2 == 0 { + TIME_IN_RANGE_START + (j as i64 * 1_000_000) + } else { + TIME_BEFORE_RANGE + (j as i64 * 1_000_000) + } + }) + .collect(); + let tags: Vec<&str> = (0..100).map(|_| *tag).collect(); + + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampNanosecondArray::from(times)) as ArrayRef, + Arc::new(StringArray::from(tags)) as ArrayRef, + ], + ) + .unwrap(); + writer.write(&batch).unwrap(); + } + writer.flush().unwrap(); + } + writer.close().unwrap(); + let buffer = Bytes::from(buffer); + // Read back with various page index policies, should get the same answer with all + for policy in [ + PageIndexPolicy::Skip, + PageIndexPolicy::Optional, + PageIndexPolicy::Required, + ] { + println!("Testing with page index policy: {:?}", policy); + let reader = TestReader::new(buffer.clone()); + let options = ArrowReaderOptions::default().with_page_index_policy(policy); + let builder = ParquetRecordBatchStreamBuilder::new_with_options(reader, options) + .await + .unwrap(); + + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + let num_row_groups = builder.metadata().num_row_groups(); + + // Initial selection: skip middle 100 rows (tag='b') per row group + let mut selectors = Vec::new(); + for _ in 0..num_row_groups { + selectors.push(RowSelector::select(100)); + selectors.push(RowSelector::skip(100)); + selectors.push(RowSelector::select(100)); + } + let selection = RowSelection::from(selectors); + + // Predicate 1: time >= START + let time_gte_predicate = + ArrowPredicateFn::new(ProjectionMask::roots(&schema_descr, [0]), |batch| { + let col = batch.column(0).as_primitive::(); + Ok(BooleanArray::from_iter( + col.iter().map(|t| t.map(|v| v >= TIME_IN_RANGE_START)), + )) + }); + + // Predicate 2: time < END + let time_lt_predicate = + ArrowPredicateFn::new(ProjectionMask::roots(&schema_descr, [0]), |batch| { + let col = batch.column(0).as_primitive::(); + Ok(BooleanArray::from_iter( + col.iter().map(|t| t.map(|v| v < TIME_IN_RANGE_END)), + )) + }); + + let row_filter = RowFilter::new(vec![ + Box::new(time_gte_predicate), + Box::new(time_lt_predicate), + ]); + + // Output projection: Only tag column (time not in output) + let projection = ProjectionMask::roots(&schema_descr, [1]); + + let stream = builder + .with_row_filter(row_filter) + .with_row_selection(selection) + .with_projection(projection) + .build() + .unwrap(); + + // Stream should complete without error and the same results + let batches: Vec = stream.try_collect().await.unwrap(); + + let batch = concat_batches(&batches[0].schema(), &batches).unwrap(); + assert_eq!(batch.num_columns(), 1); + let expected = StringArray::from_iter_values( + std::iter::repeat_n("a", 50) + .chain(std::iter::repeat_n("c", 50)) + .chain(std::iter::repeat_n("a", 50)) + .chain(std::iter::repeat_n("c", 50)), + ); + assert_eq!(batch.column(0).as_string(), &expected); + } +} + +/// Regression test: Auto resolution and loaded row ranges must be refreshed for +/// each predicate's current selection and projection. +/// +/// Scenario: +/// - Dense initial RowSelection (alternating select/skip) covers all pages → Auto resolves to Mask +/// - Predicate 1 evaluates on column A, narrows selection to skip middle pages +/// - Predicate 2's column B is fetched sparsely with the narrowed selection (missing middle pages) +/// - Predicate 2 remains safe with Mask because its loaded row ranges are recomputed for column B +#[tokio::test] +async fn test_multi_predicate_auto_mask_with_sparse_pages() { + // 300 rows, 1 row group, 100 rows per page (3 pages) + let num_rows = 300usize; + let rows_per_page = 100; + + let schema = Arc::new(Schema::new(vec![ + Field::new("filter_col", DataType::Int32, false), + Field::new("value_col", DataType::Int32, false), + ])); + + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(num_rows)) + .set_data_page_row_count_limit(rows_per_page) + .set_write_batch_size(rows_per_page) + .set_dictionary_enabled(false) + .build(); + + // filter_col: 0 for first and last 100 rows, 1 for middle 100 rows + // value_col: just row index + let filter_values: Vec = (0..num_rows as i32) + .map(|i| if (100..200).contains(&i) { 1 } else { 0 }) + .collect(); + let value_values: Vec = (0..num_rows as i32).collect(); + + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(filter_values)) as ArrayRef, + Arc::new(Int32Array::from(value_values)) as ArrayRef, + ], + ) + .unwrap(); + + let mut buffer = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let buffer = Bytes::from(buffer); + + let reader = TestReader::new(buffer); + let options = ArrowReaderOptions::default().with_page_index_policy(PageIndexPolicy::Required); + let builder = ParquetRecordBatchStreamBuilder::new_with_options(reader, options) + .await + .unwrap(); + + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + // Dense initial selection: Select(1), Skip(1) repeated → triggers Mask strategy + // Covers all pages since every page has selected rows + let selectors: Vec = (0..num_rows / 2) + .flat_map(|_| vec![RowSelector::select(1), RowSelector::skip(1)]) + .collect(); + let selection = RowSelection::from(selectors); + + // Predicate 1 on filter_col: keeps only rows where filter_col == 0 + // (first 100 and last 100 rows). After this, middle page is excluded. + let pred1 = ArrowPredicateFn::new(ProjectionMask::roots(&schema_descr, [0]), |batch| { + let col = batch.column(0).as_primitive::(); + Ok(BooleanArray::from_iter( + col.iter().map(|v| v.map(|val| val == 0)), + )) + }); + + // Predicate 2 on value_col: keeps rows where value_col < 250 + // This column is fetched AFTER predicate 1 narrows the selection. + // Its sparse data will be missing the middle page. + let predicate_batch_sizes = Arc::new(Mutex::new(Vec::new())); + let observed_batch_sizes = Arc::clone(&predicate_batch_sizes); + let pred2 = ArrowPredicateFn::new(ProjectionMask::roots(&schema_descr, [1]), move |batch| { + observed_batch_sizes.lock().unwrap().push(batch.num_rows()); + let col = batch.column(0).as_primitive::(); + Ok(BooleanArray::from_iter( + col.iter().map(|v| v.map(|val| val < 250)), + )) + }); + + let row_filter = RowFilter::new(vec![Box::new(pred1), Box::new(pred2)]); + let projection = ProjectionMask::roots(&schema_descr, [0, 1]); + + let stream = builder + .with_row_filter(row_filter) + .with_row_selection(selection) + .with_projection(projection) + .with_max_predicate_cache_size(0) + .build() + .unwrap(); + + // Without loaded-range-aware Mask reads, this panics with an invalid sparse-page offset. + let batches: Vec = stream.try_collect().await.unwrap(); + assert_eq!(*predicate_batch_sizes.lock().unwrap(), vec![100]); + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![75] + ); + + let batch = concat_batches(&batches[0].schema(), &batches).unwrap(); + + // Verify results: rows where filter_col==0 AND value_col<250 AND original alternating selection + // That's even-indexed rows in [0,100) with value<250 → rows 0,2,4,...,98 (50 rows) + // Plus even-indexed rows in [200,250) with value<250 → rows 200,202,...,248 (25 rows) + assert_eq!(batch.num_rows(), 75); +} + +/// Regression test: a Mask cursor built directly from a `BooleanBuffer` +/// (rather than from `Vec`) must also stay within loaded row +/// ranges when page pruning leaves sparse column data. +/// +/// The buffer is a non-byte-aligned `BooleanBuffer::slice(...)` so the direct +/// `BooleanBuffer -> MaskCursor` path is exercised with a bit offset. The +/// selection keeps only the first and last of 12 rows (2 rows per page), so +/// the four middle pages are never fetched. Without loaded-range propagation +/// in `new_mask_from_buffer`, decoding would cross the unloaded pages and fail +/// with an invalid sparse-page offset. +#[tokio::test] +async fn test_mask_from_boolean_buffer_slice_with_sparse_pages() { + use arrow_buffer::BooleanBuffer; + + let num_rows = 12; + let values = (0..num_rows).collect::>(); + let data = make_two_column_i64_file(&values, 2); + + // 5 padding bits force the sliced mask to start mid-byte. + let mut bits = vec![true; 5]; + bits.push(true); // row 0 selected + bits.extend(std::iter::repeat_n(false, 10)); // rows 1..=10 skipped + bits.push(true); // row 11 selected + bits.extend([false, true, false]); // trailing padding outside the slice + let mask = BooleanBuffer::from(bits).slice(5, num_rows as usize); + + for policy in [ + RowSelectionPolicy::Mask, + RowSelectionPolicy::Auto { threshold: 32 }, + ] { + let stream = ParquetRecordBatchStreamBuilder::new_with_options( + TestReader::new(data.clone()), + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required), + ) + .await + .unwrap() + .with_row_selection(RowSelection::from_boolean_buffer(mask.clone())) + .with_batch_size(num_rows as usize) + .with_row_selection_policy(policy) + .build() + .unwrap(); + + let schema = stream.schema().clone(); + let batches: Vec = stream.try_collect().await.unwrap(); + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![2], + "policy={policy:?}" + ); + + let batch = concat_batches(&schema, &batches).unwrap(); + assert_eq!( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[0, 11], + "policy={policy:?}" + ); + } +} diff --git a/vendor/parquet/tests/arrow_reader/row_filter/mod.rs b/vendor/parquet/tests/arrow_reader/row_filter/mod.rs new file mode 100644 index 0000000..a09fd4a --- /dev/null +++ b/vendor/parquet/tests/arrow_reader/row_filter/mod.rs @@ -0,0 +1,20 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#[cfg(feature = "async")] +mod r#async; +mod sync; diff --git a/vendor/parquet/tests/arrow_reader/row_filter/sync.rs b/vendor/parquet/tests/arrow_reader/row_filter/sync.rs new file mode 100644 index 0000000..a224c95 --- /dev/null +++ b/vendor/parquet/tests/arrow_reader/row_filter/sync.rs @@ -0,0 +1,421 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT 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::sync::Arc; + +use arrow::{ + array::AsArray, + compute::{concat_batches, kernels::cmp::eq, or}, +}; +use arrow_array::{ArrayRef, Int32Array, Int64Array, RecordBatch, RecordBatchReader}; +use arrow_schema::{DataType as ArrowDataType, Field, Schema}; +use bytes::Bytes; +use parquet::{ + arrow::{ + ArrowWriter, ProjectionMask, + arrow_reader::{ + ArrowPredicateFn, ArrowReaderOptions, ParquetRecordBatchReaderBuilder, RowFilter, + RowSelection, RowSelectionPolicy, RowSelector, + }, + }, + errors::Result, + file::{ + metadata::PageIndexPolicy, + properties::{WriterProperties, WriterVersion}, + }, +}; + +#[test] +fn test_row_selection_interleaved_skip() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "v", + ArrowDataType::Int32, + false, + )])); + + let values = Int32Array::from(vec![0, 1, 2, 3, 4]); + let batch = RecordBatch::try_from_iter([("v", Arc::new(values) as ArrayRef)]).unwrap(); + + let mut buffer = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(), None).unwrap(); + writer.write(&batch)?; + writer.close()?; + + let selection = RowSelection::from(vec![ + RowSelector::select(1), + RowSelector::skip(2), + RowSelector::select(2), + ]); + + let mut reader = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(buffer))? + .with_batch_size(4) + .with_row_selection(selection) + .build()?; + + let out = reader.next().unwrap()?; + assert_eq!(out.num_rows(), 3); + let values = out + .column(0) + .as_primitive::() + .values(); + assert_eq!(values, &[0, 3, 4]); + assert!(reader.next().is_none()); + Ok(()) +} + +#[test] +fn test_row_selection_mask_sparse_rows() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "v", + ArrowDataType::Int32, + false, + )])); + + let values = Int32Array::from((0..30).collect::>()); + let batch = RecordBatch::try_from_iter([("v", Arc::new(values) as ArrayRef)])?; + + let mut buffer = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(), None)?; + writer.write(&batch)?; + writer.close()?; + + let total_rows = batch.num_rows(); + let ranges = (1..total_rows) + .step_by(2) + .map(|i| i..i + 1) + .collect::>(); + let selection = RowSelection::from_consecutive_ranges(ranges.into_iter(), total_rows); + + let selectors: Vec = selection.clone().into(); + assert!(total_rows < selectors.len() * 8); + + let bytes = Bytes::from(buffer); + + let reader = ParquetRecordBatchReaderBuilder::try_new(bytes.clone())? + .with_batch_size(7) + .with_row_selection(selection) + .build()?; + + let mut collected = Vec::new(); + for batch in reader { + let batch = batch?; + collected.extend_from_slice( + batch + .column(0) + .as_primitive::() + .values(), + ); + } + + let expected: Vec = (1..total_rows).step_by(2).map(|i| i as i32).collect(); + assert_eq!(collected, expected); + Ok(()) +} + +#[test] +fn test_row_filter_full_page_skip_is_handled() { + let first_value: i64 = 1111; + let last_value: i64 = 9999; + let num_rows: usize = 12; + + // build data with row selection average length 4 + // The result would be (1111 XXXX) ... (4 page in the middle)... (XXXX 9999) + // The Row Selection would be [1111, (skip 10), 9999] + let schema = Arc::new(Schema::new(vec![ + Field::new("key", arrow_schema::DataType::Int64, false), + Field::new("value", arrow_schema::DataType::Int64, false), + ])); + + let mut int_values: Vec = (0..num_rows as i64).collect(); + int_values[0] = first_value; + int_values[num_rows - 1] = last_value; + let keys = Int64Array::from(int_values.clone()); + let values = Int64Array::from(int_values.clone()); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(keys) as ArrayRef, Arc::new(values) as ArrayRef], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_write_batch_size(2) + .set_data_page_row_count_limit(2) + .build(); + + let mut buffer = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut buffer, schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let data = Bytes::from(buffer); + + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let builder = + ParquetRecordBatchReaderBuilder::try_new_with_options(data.clone(), options).unwrap(); + let schema = builder.parquet_schema().clone(); + let filter_mask = ProjectionMask::leaves(&schema, [0]); + + let make_predicate = |mask: ProjectionMask| { + ArrowPredicateFn::new(mask, move |batch: RecordBatch| { + let column = batch.column(0); + let match_first = eq(column, &Int64Array::new_scalar(first_value))?; + let match_second = eq(column, &Int64Array::new_scalar(last_value))?; + or(&match_first, &match_second) + }) + }; + + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let predicate = make_predicate(filter_mask.clone()); + + // The batch size is set to 12 to read all rows in one go after filtering + // If the Reader chooses mask to handle filter, it might cause panic because the mid 4 pages may not be decoded. + let reader = ParquetRecordBatchReaderBuilder::try_new_with_options(data.clone(), options) + .unwrap() + .with_row_filter(RowFilter::new(vec![Box::new(predicate)])) + .with_row_selection_policy(RowSelectionPolicy::Auto { threshold: 32 }) + .with_batch_size(12) + .build() + .unwrap(); + + // Predicate pruning used to panic once mask-backed plans removed whole pages. + // Collecting into batches validates that Mask can read around unloaded pages. + let schema = reader.schema().clone(); + let batches = reader.collect::, _>>().unwrap(); + let result = concat_batches(&schema, &batches).unwrap(); + assert_eq!(result.num_rows(), 2); +} + +/// Regression test for +/// +/// When `skip_records` on a list column crosses v2 data page boundaries, +/// the partial record state (`has_partial`) in the repetition level +/// decoder must be flushed before the whole-page-skip shortcut can fire. +/// Without the fix, the list column over-skips by one record, causing +/// struct children to disagree on record counts. +#[test] +fn test_row_selection_list_column_v2_page_boundary_skip() { + use arrow_array::builder::{Int32Builder, ListBuilder}; + + let num_rows = 10usize; + + // Schema: { id: int32, values: list } + // Two top-level columns so that StructArrayReader can detect the + // desync between a simple column (id) and a list column (values). + let schema = Arc::new(Schema::new(vec![ + Field::new("id", ArrowDataType::Int32, false), + Field::new( + "values", + ArrowDataType::List(Arc::new(Field::new("item", ArrowDataType::Int32, true))), + false, + ), + ])); + + // Each row: id = i, values = [i*10, i*10+1] + let ids = Int32Array::from((0..num_rows as i32).collect::>()); + let mut list_builder = ListBuilder::new(Int32Builder::new()); + for i in 0..num_rows as i32 { + list_builder.values().append_value(i * 10); + list_builder.values().append_value(i * 10 + 1); + list_builder.append(true); + } + let values_array = list_builder.finish(); + + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(ids) as ArrayRef, + Arc::new(values_array) as ArrayRef, + ], + ) + .unwrap(); + + // Force v2 data pages with exactly 2 rows per page → 5 pages. + // The default reader (no offset index) puts SerializedPageReader + // in Values state where at_record_boundary() returns false for + // non-last pages, matching the production scenario. + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_2_0) + .set_write_batch_size(2) + .set_data_page_row_count_limit(2) + .build(); + + let mut buffer = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let data = Bytes::from(buffer); + + // 1. Read without row selection — should always succeed + let reader = ParquetRecordBatchReaderBuilder::try_new(data.clone()) + .unwrap() + .build() + .unwrap(); + let batches: Vec<_> = reader.map(|r| r.unwrap()).collect(); + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, num_rows, "full read should return all rows"); + + // 2. Read with row selection: select first and last row. + // This produces the sequence read(1) + skip(8) + read(1). + // + // The skip crosses v2 page boundaries. Without the fix, the + // list column's rep-level decoder has stale has_partial state + // after exhausting the first page's remaining levels, causing + // the whole-page-skip shortcut to over-count by one record. + // + // We must use RowSelectionPolicy::Selectors because the default + // Auto policy would choose the Mask strategy for this small + // selection, which reads all rows then filters (never calling + // skip_records, thereby hiding the bug). + let selection = RowSelection::from(vec![ + RowSelector::select(1), + RowSelector::skip(num_rows - 2), + RowSelector::select(1), + ]); + + let reader = ParquetRecordBatchReaderBuilder::try_new(data) + .unwrap() + .with_row_selection(selection) + .with_row_selection_policy(RowSelectionPolicy::Selectors) + .build() + .unwrap(); + let batches: Vec = reader.map(|r| r.unwrap()).collect(); + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, 2, "selection should return exactly 2 rows"); + + // Verify data correctness: row 0 and row 9 + let result = concat_batches(&schema, &batches).unwrap(); + let id_col = result + .column(0) + .as_primitive::(); + assert_eq!(id_col.value(0), 0); + assert_eq!(id_col.value(1), (num_rows - 1) as i32); + + let list_col = result.column(1).as_list::(); + let first = list_col + .value(0) + .as_primitive::() + .values() + .to_vec(); + assert_eq!(first, vec![0, 1]); + let last = list_col + .value(1) + .as_primitive::() + .values() + .to_vec(); + let n = (num_rows - 1) as i32; + assert_eq!(last, vec![n * 10, n * 10 + 1]); +} + +/// Regression test for +/// +/// When leaf columns inside a `List>` have different page +/// boundaries (due to value-size differences), the `has_partial` state +/// bug causes one leaf to over-skip by one record while the other stays +/// correct. +#[test] +fn test_list_struct_page_boundary_desync_produces_length_mismatch() { + use arrow_array::Array; + use arrow_array::builder::{Int32Builder, ListBuilder, StringBuilder, StructBuilder}; + use arrow_schema::Fields; + + let num_rows = 14usize; + // Long string forces the string column to flush pages much sooner + // than the int32 column, creating different page boundaries. + let long_prefix = "x".repeat(500); + + // Schema: { vals: List> } + let struct_fields = Fields::from(vec![ + Field::new("x", ArrowDataType::Int32, false), + Field::new("y", ArrowDataType::Utf8, false), + ]); + + // Build data: even rows have 2 list elements, odd rows have 3. + // This ensures different physical value counts per record, so + // reading from the wrong position produces a different total. + let mut list_builder = ListBuilder::new(StructBuilder::from_fields(struct_fields, 0)); + for i in 0..num_rows { + let num_elems = if i % 2 == 0 { 2 } else { 3 }; + let sb = list_builder.values(); + for j in 0..num_elems { + sb.field_builder::(0) + .unwrap() + .append_value(i as i32 * 10 + j); + sb.field_builder::(1) + .unwrap() + .append_value(format!("{long_prefix}_{i}_{j}")); + sb.append(true); + } + list_builder.append(true); + } + let vals_array = list_builder.finish(); + + // Derive schema from the actual array type to avoid field-name mismatches. + let schema = Arc::new(Schema::new(vec![Field::new( + "vals", + vals_array.data_type().clone(), + false, + )])); + + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(vals_array) as ArrayRef]).unwrap(); + + // V2 pages + small max_data_page_size. + // Column x (Int32): all 14 rows fit in one page (~140 bytes values). + // Column y (Utf8, 500-byte strings): pages flush after every ~2 rows. + // This creates the page boundary asymmetry needed to trigger the bug. + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_2_0) + .set_write_batch_size(2) + .set_dictionary_enabled(false) + .set_data_page_size_limit(512) + .build(); + + let mut buffer = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let data = Bytes::from(buffer); + + // 1. Read without row selection — should always succeed. + let reader = ParquetRecordBatchReaderBuilder::try_new(data.clone()) + .unwrap() + .build() + .unwrap(); + let batches: Vec<_> = reader.map(|r| r.unwrap()).collect(); + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, num_rows, "full read should return all rows"); + + // 2. Read with selection: select(1) + skip(10) + select(1). + // Without the fix, the string column (y) over-skips by 1, + // reading a record with a different element count than the + // int column (x). The inner StructArrayReader sees arrays + // of different lengths. + let selection = RowSelection::from(vec![ + RowSelector::select(1), + RowSelector::skip(10), + RowSelector::select(1), + ]); + + let reader = ParquetRecordBatchReaderBuilder::try_new(data) + .unwrap() + .with_row_selection(selection) + .with_row_selection_policy(RowSelectionPolicy::Selectors) + .build() + .unwrap(); + let batches: Vec = reader.map(|r| r.unwrap()).collect(); + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, 2, "selection should return exactly 2 rows"); +} diff --git a/vendor/parquet/tests/arrow_reader/statistics.rs b/vendor/parquet/tests/arrow_reader/statistics.rs new file mode 100644 index 0000000..bbb891d --- /dev/null +++ b/vendor/parquet/tests/arrow_reader/statistics.rs @@ -0,0 +1,3226 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! This file contains an end to end test of extracting statistics from parquet files. +//! It writes data into a parquet file, reads statistics and verifies they are correct + +use std::default::Default; +use std::fs::File; +use std::sync::Arc; + +use super::make_test_file_rg; +use super::{Scenario, struct_array}; +use arrow::compute::kernels::cast_utils::Parser; +use arrow::datatypes::{ + Date32Type, Date64Type, TimestampMicrosecondType, TimestampMillisecondType, + TimestampNanosecondType, TimestampSecondType, i256, +}; +use arrow_array::{ + Array, ArrayRef, BinaryArray, BinaryViewArray, BooleanArray, Date32Array, Date64Array, + Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, FixedSizeBinaryArray, + Float16Array, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, Int64Array, + LargeBinaryArray, LargeStringArray, RecordBatch, StringArray, StringViewArray, + Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray, + TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, + TimestampSecondArray, UInt8Array, UInt16Array, UInt32Array, UInt64Array, make_array, + new_null_array, +}; +use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; +use half::f16; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::statistics::StatisticsConverter; +use parquet::arrow::arrow_reader::{ + ArrowReaderBuilder, ArrowReaderOptions, ParquetRecordBatchReaderBuilder, +}; +use parquet::file::metadata::PageIndexPolicy; +use parquet::file::metadata::{ColumnChunkMetaData, RowGroupMetaData}; +use parquet::file::properties::{EnabledStatistics, WriterProperties}; +use parquet::file::statistics::{Statistics, ValueStatistics}; +use parquet::schema::types::{SchemaDescPtr, SchemaDescriptor}; + +#[derive(Debug, Default, Clone)] +struct Int64Case { + /// Number of nulls in the column + null_values: usize, + /// Non null values in the range `[no_null_values_start, + /// no_null_values_end]`, one value for each row + no_null_values_start: i64, + no_null_values_end: i64, + /// Number of rows per row group + row_per_group: usize, + /// if specified, overrides default statistics settings + enable_stats: Option, + /// If specified, the number of values in each page + data_page_row_count_limit: Option, +} + +impl Int64Case { + /// Return a record batch with i64 with Null values + /// The first no_null_values_end - no_null_values_start values + /// are non-null with the specified range, the rest are null + fn make_int64_batches_with_null(&self) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("i64", DataType::Int64, true)])); + + let v64: Vec = + (self.no_null_values_start as _..self.no_null_values_end as _).collect(); + + RecordBatch::try_new( + schema, + vec![make_array( + Int64Array::from_iter( + v64.into_iter() + .map(Some) + .chain(std::iter::repeat_n(None, self.null_values)), + ) + .to_data(), + )], + ) + .unwrap() + } + + // Create a parquet file with the specified settings + pub fn build(&self) -> ParquetRecordBatchReaderBuilder { + let batches = vec![self.make_int64_batches_with_null()]; + build_parquet_file( + self.row_per_group, + self.enable_stats, + self.data_page_row_count_limit, + batches, + ) + } +} + +fn build_parquet_file( + row_per_group: usize, + enable_stats: Option, + data_page_row_count_limit: Option, + batches: Vec, +) -> ParquetRecordBatchReaderBuilder { + let mut output_file = tempfile::Builder::new() + .prefix("parquert_statistics_test") + .suffix(".parquet") + .tempfile() + .expect("tempfile creation"); + + let mut builder = WriterProperties::builder().set_max_row_group_row_count(Some(row_per_group)); + if let Some(enable_stats) = enable_stats { + builder = builder.set_statistics_enabled(enable_stats); + } + if let Some(data_page_row_count_limit) = data_page_row_count_limit { + builder = builder.set_data_page_row_count_limit(data_page_row_count_limit); + } + let props = builder.build(); + + let schema = batches[0].schema(); + + let mut writer = ArrowWriter::try_new(&mut output_file, schema, Some(props)).unwrap(); + + // if we have a datapage limit send the batches in one at a time to give + // the writer a chance to be split into multiple pages + if data_page_row_count_limit.is_some() { + for batch in &batches { + for i in 0..batch.num_rows() { + writer.write(&batch.slice(i, 1)).expect("writing batch"); + } + } + } else { + for batch in &batches { + writer.write(batch).expect("writing batch"); + } + } + + let _file_meta = writer.close().unwrap(); + + let file = output_file.reopen().unwrap(); + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::from(true)); + ArrowReaderBuilder::try_new_with_options(file, options).unwrap() +} + +/// Defines what data to create in a parquet file +#[derive(Debug, Clone, Copy)] +struct TestReader { + /// What data to create in the parquet file + scenario: Scenario, + /// Number of rows per row group + row_per_group: usize, +} + +impl TestReader { + /// Create a parquet file with the specified data, and return a + /// ParquetRecordBatchReaderBuilder opened to that file. + async fn build(self) -> ParquetRecordBatchReaderBuilder { + let TestReader { + scenario, + row_per_group, + } = self; + let file = make_test_file_rg(scenario, row_per_group).await; + + // open the file & get the reader + let file = file.reopen().unwrap(); + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::from(true)); + ArrowReaderBuilder::try_new_with_options(file, options).unwrap() + } +} + +/// Which statistics should we check? +#[derive(Clone, Debug, Copy)] +enum Check { + /// Extract and check row group statistics + RowGroup, + /// Extract and check data page statistics + DataPage, + /// Extract and check both row group and data page statistics. + /// + /// Note if a row group contains a single data page, + /// the statistics for row groups and data pages are the same. + Both, +} + +impl Check { + fn row_group(&self) -> bool { + match self { + Self::RowGroup | Self::Both => true, + Self::DataPage => false, + } + } + + fn data_page(&self) -> bool { + match self { + Self::DataPage | Self::Both => true, + Self::RowGroup => false, + } + } +} + +/// Defines a test case for statistics extraction +struct Test<'a> { + /// The parquet file reader + reader: &'a ParquetRecordBatchReaderBuilder, + expected_min: ArrayRef, + expected_max: ArrayRef, + expected_null_counts: UInt64Array, + expected_row_counts: Option, + expected_max_value_exact: BooleanArray, + expected_min_value_exact: BooleanArray, + /// Which column to extract statistics from + column_name: &'static str, + /// What statistics should be checked? + check: Check, +} + +impl Test<'_> { + fn run(self) { + let converter = StatisticsConverter::try_new( + self.column_name, + self.reader.schema(), + self.reader.parquet_schema(), + ) + .unwrap(); + + self.run_checks(converter); + } + + fn run_with_schema(self, schema: &Schema) { + let converter = + StatisticsConverter::try_new(self.column_name, schema, self.reader.parquet_schema()) + .unwrap(); + + self.run_checks(converter); + } + + fn run_checks(self, converter: StatisticsConverter) { + let Self { + reader, + expected_min, + expected_max, + expected_null_counts, + expected_row_counts, + expected_max_value_exact: expected_max_exact, + expected_min_value_exact: expected_min_exact, + column_name, + check, + } = self; + + let row_groups = reader.metadata().row_groups(); + + if check.data_page() { + let column_page_index = reader + .metadata() + .column_index() + .expect("File should have column page indices"); + + let column_offset_index = reader + .metadata() + .offset_index() + .expect("File should have column offset indices"); + + let row_group_indices: Vec<_> = (0..row_groups.len()).collect(); + + let min = converter + .data_page_mins(column_page_index, column_offset_index, &row_group_indices) + .unwrap(); + assert_eq!( + &min, &expected_min, + "{column_name}: Mismatch with expected data page minimums" + ); + + let max = converter + .data_page_maxes(column_page_index, column_offset_index, &row_group_indices) + .unwrap(); + assert_eq!( + &max, &expected_max, + "{column_name}: Mismatch with expected data page maximum" + ); + + let null_counts = converter + .data_page_null_counts(column_page_index, column_offset_index, &row_group_indices) + .unwrap(); + + assert_eq!( + &null_counts, &expected_null_counts, + "{column_name}: Mismatch with expected data page null counts. \ + Actual: {null_counts:?}. Expected: {expected_null_counts:?}" + ); + + let row_counts = converter + .data_page_row_counts(column_offset_index, row_groups, &row_group_indices) + .unwrap(); + assert_eq!( + row_counts, expected_row_counts, + "{column_name}: Mismatch with expected row counts. \ + Actual: {row_counts:?}. Expected: {expected_row_counts:?}" + ); + } + + if check.row_group() { + let min = converter.row_group_mins(row_groups).unwrap(); + assert_eq!( + &min, &expected_min, + "{column_name}: Mismatch with expected minimums" + ); + + let max = converter.row_group_maxes(row_groups).unwrap(); + assert_eq!( + &max, &expected_max, + "{column_name}: Mismatch with expected maximum" + ); + + let null_counts = converter.row_group_null_counts(row_groups).unwrap(); + assert_eq!( + &null_counts, &expected_null_counts, + "{column_name}: Mismatch with expected null counts. \ + Actual: {null_counts:?}. Expected: {expected_null_counts:?}" + ); + + let row_counts = converter + .row_group_row_counts(reader.metadata().row_groups().iter()) + .unwrap(); + assert_eq!( + row_counts, expected_row_counts, + "{column_name}: Mismatch with expected row counts. \ + Actual: {row_counts:?}. Expected: {expected_row_counts:?}" + ); + + let is_max_value_exact = converter + .row_group_is_max_value_exact(reader.metadata().row_groups().iter()) + .unwrap(); + assert_eq!( + is_max_value_exact, expected_max_exact, + "{column_name}: Mismatch with expected max value exactness. \ + Actual: {is_max_value_exact:?}. Expected: {expected_max_exact:?}" + ); + + let is_min_value_exact = converter + .row_group_is_min_value_exact(reader.metadata().row_groups().iter()) + .unwrap(); + assert_eq!( + is_min_value_exact, expected_min_exact, + "{column_name}: Mismatch with expected min value exactness. \ + Actual: {is_min_value_exact:?}. Expected: {expected_min_exact:?}" + ); + } + } + + /// Run the test and expect a column not found error + fn run_col_not_found(self) { + let Self { + reader, + expected_min: _, + expected_max: _, + expected_null_counts: _, + expected_row_counts: _, + column_name, + .. + } = self; + + let converter = + StatisticsConverter::try_new(column_name, reader.schema(), reader.parquet_schema()); + + assert!(converter.is_err()); + } +} + +// TESTS +// +// Remaining cases +// f64::NAN + +#[tokio::test] +async fn test_max_and_min_value_truncated() { + let reader = TestReader { + scenario: Scenario::TruncatedUTF8, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + // min is truncated to + // 1. `"a".repeat(64)`, original value is `"a".repeat(64) + "1"` + // 2. `"e".repeat(64)`, original value is `"e".repeat(64) + "5"` + // 3. "j", as expected with no truncation + expected_min: Arc::new(StringArray::from(vec![ + &("a".repeat(64)), + &("e".repeat(64)), + "j", + ])), + // max is truncated to + // 1. `"d".repeat(63) + "e"`, original value is `"d".repeat(64) + "4"` + // 2. `"i".repeat(63) + "j"`, original value is `"i".repeat(64) + "6"` + // 3. `"n".repeat(63) + "o"`, original value is `"n".repeat(64) + "14"` + expected_max: Arc::new(StringArray::from(vec![ + "d".repeat(63) + "e", + "i".repeat(63) + "j", + "n".repeat(63) + "o", + ])), + // no nulls + expected_null_counts: UInt64Array::from(vec![1, 0, 0]), + // 3 rows + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5])), + // all max values are truncated + expected_max_value_exact: BooleanArray::from(vec![false, false, false]), + // min values are truncated in the first two row groups + expected_min_value_exact: BooleanArray::from(vec![false, false, true]), + column_name: "utf8", + check: Check::Both, + } + .run() +} + +#[tokio::test] +async fn test_one_row_group_without_null() { + let reader = Int64Case { + null_values: 0, + no_null_values_start: 4, + no_null_values_end: 7, + row_per_group: 20, + ..Default::default() + } + .build(); + + Test { + reader: &reader, + // min is 4 + expected_min: Arc::new(Int64Array::from(vec![4])), + // max is 6 + expected_max: Arc::new(Int64Array::from(vec![6])), + // no nulls + expected_null_counts: UInt64Array::from(vec![0]), + // 3 rows + expected_row_counts: Some(UInt64Array::from(vec![3])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true]), + expected_min_value_exact: BooleanArray::from(vec![true]), + column_name: "i64", + check: Check::Both, + } + .run() +} + +#[tokio::test] +async fn test_one_row_group_with_null_and_negative() { + let reader = Int64Case { + null_values: 2, + no_null_values_start: -1, + no_null_values_end: 5, + row_per_group: 20, + ..Default::default() + } + .build(); + + Test { + reader: &reader, + // min is -1 + expected_min: Arc::new(Int64Array::from(vec![-1])), + // max is 4 + expected_max: Arc::new(Int64Array::from(vec![4])), + // 2 nulls + expected_null_counts: UInt64Array::from(vec![2]), + // 8 rows + expected_row_counts: Some(UInt64Array::from(vec![8])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true]), + expected_min_value_exact: BooleanArray::from(vec![true]), + column_name: "i64", + check: Check::Both, + } + .run() +} + +#[tokio::test] +async fn test_two_row_group_with_null() { + let reader = Int64Case { + null_values: 2, + no_null_values_start: 4, + no_null_values_end: 17, + row_per_group: 10, + ..Default::default() + } + .build(); + + Test { + reader: &reader, + // mins are [4, 14] + expected_min: Arc::new(Int64Array::from(vec![4, 14])), + // maxes are [13, 16] + expected_max: Arc::new(Int64Array::from(vec![13, 16])), + // nulls are [0, 2] + expected_null_counts: UInt64Array::from(vec![0, 2]), + // row counts are [10, 5] + expected_row_counts: Some(UInt64Array::from(vec![10, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "i64", + check: Check::Both, + } + .run() +} + +#[tokio::test] +async fn test_two_row_groups_with_all_nulls_in_one() { + let reader = Int64Case { + null_values: 4, + no_null_values_start: -2, + no_null_values_end: 2, + row_per_group: 5, + ..Default::default() + } + .build(); + + Test { + reader: &reader, + // mins are [-2, null] + expected_min: Arc::new(Int64Array::from(vec![Some(-2), None])), + // maxes are [1, null] + expected_max: Arc::new(Int64Array::from(vec![Some(1), None])), + // nulls are [1, 3] + expected_null_counts: UInt64Array::from(vec![1, 3]), + // row counts are [5, 3] + expected_row_counts: Some(UInt64Array::from(vec![5, 3])), + expected_max_value_exact: BooleanArray::from(vec![true, false]), + expected_min_value_exact: BooleanArray::from(vec![true, false]), + column_name: "i64", + check: Check::Both, + } + .run() +} + +#[tokio::test] +async fn test_multiple_data_pages_nulls_and_negatives() { + let reader = Int64Case { + null_values: 3, + no_null_values_start: -1, + no_null_values_end: 10, + row_per_group: 20, + // limit page row count to 4 + data_page_row_count_limit: Some(4), + enable_stats: Some(EnabledStatistics::Page), + } + .build(); + + // Data layout looks like this: + // + // page 0: [-1, 0, 1, 2] + // page 1: [3, 4, 5, 6] + // page 2: [7, 8, 9, null] + // page 3: [null, null] + Test { + reader: &reader, + expected_min: Arc::new(Int64Array::from(vec![Some(-1), Some(3), Some(7), None])), + expected_max: Arc::new(Int64Array::from(vec![Some(2), Some(6), Some(9), None])), + expected_null_counts: UInt64Array::from(vec![0, 0, 1, 2]), + expected_row_counts: Some(UInt64Array::from(vec![4, 4, 4, 2])), + expected_max_value_exact: BooleanArray::from(vec![true, true, true, false]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, false]), + column_name: "i64", + check: Check::DataPage, + } + .run() +} + +#[tokio::test] +async fn test_data_page_stats_with_all_null_page() { + for data_type in &[ + DataType::Boolean, + DataType::UInt64, + DataType::UInt32, + DataType::UInt16, + DataType::UInt8, + DataType::Int64, + DataType::Int32, + DataType::Int16, + DataType::Int8, + DataType::Float16, + DataType::Float32, + DataType::Float64, + DataType::Date32, + DataType::Date64, + DataType::Time32(TimeUnit::Millisecond), + DataType::Time32(TimeUnit::Second), + DataType::Time64(TimeUnit::Microsecond), + DataType::Time64(TimeUnit::Nanosecond), + DataType::Timestamp(TimeUnit::Second, None), + DataType::Timestamp(TimeUnit::Millisecond, None), + DataType::Timestamp(TimeUnit::Microsecond, None), + DataType::Timestamp(TimeUnit::Nanosecond, None), + DataType::Binary, + DataType::LargeBinary, + DataType::FixedSizeBinary(3), + DataType::Utf8, + DataType::LargeUtf8, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + DataType::Decimal32(8, 2), // as INT32 + DataType::Decimal64(8, 2), // as INT32 + DataType::Decimal64(10, 2), // as INT64 + DataType::Decimal128(8, 2), // as INT32 + DataType::Decimal128(10, 2), // as INT64 + DataType::Decimal128(20, 2), // as FIXED_LEN_BYTE_ARRAY + DataType::Decimal256(8, 2), // as INT32 + DataType::Decimal256(10, 2), // as INT64 + DataType::Decimal256(20, 2), // as FIXED_LEN_BYTE_ARRAY + ] { + let batch = RecordBatch::try_from_iter(vec![("col", new_null_array(data_type, 4))]) + .expect("record batch creation"); + + let reader = build_parquet_file(4, Some(EnabledStatistics::Page), Some(4), vec![batch]); + + let expected_data_type = match data_type { + DataType::Dictionary(_, value_type) => value_type.as_ref(), + _ => data_type, + }; + + // There is one data page with 4 nulls + // The statistics should be present but null + Test { + reader: &reader, + expected_min: new_null_array(expected_data_type, 1), + expected_max: new_null_array(expected_data_type, 1), + expected_null_counts: UInt64Array::from(vec![4]), + expected_row_counts: Some(UInt64Array::from(vec![4])), + expected_max_value_exact: BooleanArray::from(vec![false]), + expected_min_value_exact: BooleanArray::from(vec![false]), + column_name: "col", + check: Check::DataPage, + } + .run() + } +} + +/////////////// MORE GENERAL TESTS ////////////////////// +// . Many columns in a file +// . Differnet data types +// . Different row group sizes + +// Four different integer types +#[tokio::test] +async fn test_int_64() { + // This creates a parquet files of 4 columns named "i8", "i16", "i32", "i64" + let reader = TestReader { + scenario: Scenario::Int, + row_per_group: 5, + } + .build() + .await; + + // since each row has only one data page, the statistics are the same + Test { + reader: &reader, + // mins are [-5, -4, 0, 5] + expected_min: Arc::new(Int64Array::from(vec![-5, -4, 0, 5])), + // maxes are [-1, 0, 4, 9] + expected_max: Arc::new(Int64Array::from(vec![-1, 0, 4, 9])), + // nulls are [0, 0, 0, 0] + expected_null_counts: UInt64Array::from(vec![0, 0, 0, 0]), + // row counts are [5, 5, 5, 5] + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + column_name: "i64", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_int_32() { + // This creates a parquet files of 4 columns named "i8", "i16", "i32", "i64" + let reader = TestReader { + scenario: Scenario::Int, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + // mins are [-5, -4, 0, 5] + expected_min: Arc::new(Int32Array::from(vec![-5, -4, 0, 5])), + // maxes are [-1, 0, 4, 9] + expected_max: Arc::new(Int32Array::from(vec![-1, 0, 4, 9])), + // nulls are [0, 0, 0, 0] + expected_null_counts: UInt64Array::from(vec![0, 0, 0, 0]), + // row counts are [5, 5, 5, 5] + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + column_name: "i32", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_int_16() { + // This creates a parquet files of 4 columns named "i8", "i16", "i32", "i64" + let reader = TestReader { + scenario: Scenario::Int, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + // mins are [-5, -4, 0, 5] + expected_min: Arc::new(Int16Array::from(vec![-5, -4, 0, 5])), // panic here because the actual data is Int32Array + // maxes are [-1, 0, 4, 9] + expected_max: Arc::new(Int16Array::from(vec![-1, 0, 4, 9])), + // nulls are [0, 0, 0, 0] + expected_null_counts: UInt64Array::from(vec![0, 0, 0, 0]), + // row counts are [5, 5, 5, 5] + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + column_name: "i16", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_int_8() { + // This creates a parquet files of 4 columns named "i8", "i16", "i32", "i64" + let reader = TestReader { + scenario: Scenario::Int, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + // mins are [-5, -4, 0, 5] + expected_min: Arc::new(Int8Array::from(vec![-5, -4, 0, 5])), // panic here because the actual data is Int32Array + // maxes are [-1, 0, 4, 9] + expected_max: Arc::new(Int8Array::from(vec![-1, 0, 4, 9])), + // nulls are [0, 0, 0, 0] + expected_null_counts: UInt64Array::from(vec![0, 0, 0, 0]), + // row counts are [5, 5, 5, 5] + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + column_name: "i8", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_float_16() { + // This creates a parquet files of 1 column named f + let reader = TestReader { + scenario: Scenario::Float16, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + // mins are [-5, -4, 0, 5] + expected_min: Arc::new(Float16Array::from(vec![ + f16::from_f32(-5.), + f16::from_f32(-4.), + f16::from_f32(-0.), + f16::from_f32(5.), + ])), + // maxes are [-1, 0, 4, 9] + expected_max: Arc::new(Float16Array::from(vec![ + f16::from_f32(-1.), + f16::from_f32(0.), + f16::from_f32(4.), + f16::from_f32(9.), + ])), + // nulls are [0, 0, 0, 0] + expected_null_counts: UInt64Array::from(vec![0, 0, 0, 0]), + // row counts are [5, 5, 5, 5] + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + column_name: "f", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_float_32() { + // This creates a parquet files of 1 column named f + let reader = TestReader { + scenario: Scenario::Float32, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + // mins are [-5, -4, 0, 5] + expected_min: Arc::new(Float32Array::from(vec![-5., -4., -0., 5.0])), + // maxes are [-1, 0, 4, 9] + expected_max: Arc::new(Float32Array::from(vec![-1., 0., 4., 9.])), + // nulls are [0, 0, 0, 0] + expected_null_counts: UInt64Array::from(vec![0, 0, 0, 0]), + // row counts are [5, 5, 5, 5] + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + column_name: "f", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_float_64() { + // This creates a parquet files of 1 column named f + let reader = TestReader { + scenario: Scenario::Float64, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + // mins are [-5, -4, 0, 5] + expected_min: Arc::new(Float64Array::from(vec![-5., -4., -0., 5.0])), + // maxes are [-1, 0, 4, 9] + expected_max: Arc::new(Float64Array::from(vec![-1., 0., 4., 9.])), + // nulls are [0, 0, 0, 0] + expected_null_counts: UInt64Array::from(vec![0, 0, 0, 0]), + // row counts are [5, 5, 5, 5] + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + column_name: "f", + check: Check::Both, + } + .run(); +} + +// timestamp +#[tokio::test] +async fn test_timestamp() { + // This creates a parquet files of 9 columns named "nanos", "nanos_timezoned", "micros", "micros_timezoned", "millis", "millis_timezoned", "seconds", "seconds_timezoned", "names" + // "nanos" --> TimestampNanosecondArray + // "nanos_timezoned" --> TimestampNanosecondArray + // "micros" --> TimestampMicrosecondArray + // "micros_timezoned" --> TimestampMicrosecondArray + // "millis" --> TimestampMillisecondArray + // "millis_timezoned" --> TimestampMillisecondArray + // "seconds" --> TimestampSecondArray + // "seconds_timezoned" --> TimestampSecondArray + // "names" --> StringArray + // + // The file is created by 4 record batches, each has 5 rows. + // Since the row group size is set to 5, those 4 batches will go into 4 row groups + // This creates a parquet files of 4 columns named "nanos", "nanos_timezoned", "micros", "micros_timezoned", "millis", "millis_timezoned", "seconds", "seconds_timezoned" + let reader = TestReader { + scenario: Scenario::Timestamps, + row_per_group: 5, + } + .build() + .await; + + let tz = "Pacific/Efate"; + + Test { + reader: &reader, + expected_min: Arc::new(TimestampNanosecondArray::from(vec![ + TimestampNanosecondType::parse("2020-01-01T01:01:01"), + TimestampNanosecondType::parse("2020-01-01T01:01:11"), + TimestampNanosecondType::parse("2020-01-01T01:11:01"), + TimestampNanosecondType::parse("2020-01-11T01:01:01"), + ])), + expected_max: Arc::new(TimestampNanosecondArray::from(vec![ + TimestampNanosecondType::parse("2020-01-02T01:01:01"), + TimestampNanosecondType::parse("2020-01-02T01:01:11"), + TimestampNanosecondType::parse("2020-01-02T01:11:01"), + TimestampNanosecondType::parse("2020-01-12T01:01:01"), + ])), + // nulls are [1, 1, 1, 1] + expected_null_counts: UInt64Array::from(vec![1, 1, 1, 1]), + // row counts are [5, 5, 5, 5] + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + column_name: "nanos", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new( + TimestampNanosecondArray::from(vec![ + TimestampNanosecondType::parse("2020-01-01T01:01:01"), + TimestampNanosecondType::parse("2020-01-01T01:01:11"), + TimestampNanosecondType::parse("2020-01-01T01:11:01"), + TimestampNanosecondType::parse("2020-01-11T01:01:01"), + ]) + .with_timezone(tz), + ), + expected_max: Arc::new( + TimestampNanosecondArray::from(vec![ + TimestampNanosecondType::parse("2020-01-02T01:01:01"), + TimestampNanosecondType::parse("2020-01-02T01:01:11"), + TimestampNanosecondType::parse("2020-01-02T01:11:01"), + TimestampNanosecondType::parse("2020-01-12T01:01:01"), + ]) + .with_timezone(tz), + ), + // nulls are [1, 1, 1, 1] + expected_null_counts: UInt64Array::from(vec![1, 1, 1, 1]), + // row counts are [5, 5, 5, 5] + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + column_name: "nanos_timezoned", + check: Check::Both, + } + .run(); + + // micros + Test { + reader: &reader, + expected_min: Arc::new(TimestampMicrosecondArray::from(vec![ + TimestampMicrosecondType::parse("2020-01-01T01:01:01"), + TimestampMicrosecondType::parse("2020-01-01T01:01:11"), + TimestampMicrosecondType::parse("2020-01-01T01:11:01"), + TimestampMicrosecondType::parse("2020-01-11T01:01:01"), + ])), + expected_max: Arc::new(TimestampMicrosecondArray::from(vec![ + TimestampMicrosecondType::parse("2020-01-02T01:01:01"), + TimestampMicrosecondType::parse("2020-01-02T01:01:11"), + TimestampMicrosecondType::parse("2020-01-02T01:11:01"), + TimestampMicrosecondType::parse("2020-01-12T01:01:01"), + ])), + expected_null_counts: UInt64Array::from(vec![1, 1, 1, 1]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + column_name: "micros", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new( + TimestampMicrosecondArray::from(vec![ + TimestampMicrosecondType::parse("2020-01-01T01:01:01"), + TimestampMicrosecondType::parse("2020-01-01T01:01:11"), + TimestampMicrosecondType::parse("2020-01-01T01:11:01"), + TimestampMicrosecondType::parse("2020-01-11T01:01:01"), + ]) + .with_timezone(tz), + ), + expected_max: Arc::new( + TimestampMicrosecondArray::from(vec![ + TimestampMicrosecondType::parse("2020-01-02T01:01:01"), + TimestampMicrosecondType::parse("2020-01-02T01:01:11"), + TimestampMicrosecondType::parse("2020-01-02T01:11:01"), + TimestampMicrosecondType::parse("2020-01-12T01:01:01"), + ]) + .with_timezone(tz), + ), + // nulls are [1, 1, 1, 1] + expected_null_counts: UInt64Array::from(vec![1, 1, 1, 1]), + // row counts are [5, 5, 5, 5] + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + column_name: "micros_timezoned", + check: Check::Both, + } + .run(); + + // millis + Test { + reader: &reader, + expected_min: Arc::new(TimestampMillisecondArray::from(vec![ + TimestampMillisecondType::parse("2020-01-01T01:01:01"), + TimestampMillisecondType::parse("2020-01-01T01:01:11"), + TimestampMillisecondType::parse("2020-01-01T01:11:01"), + TimestampMillisecondType::parse("2020-01-11T01:01:01"), + ])), + expected_max: Arc::new(TimestampMillisecondArray::from(vec![ + TimestampMillisecondType::parse("2020-01-02T01:01:01"), + TimestampMillisecondType::parse("2020-01-02T01:01:11"), + TimestampMillisecondType::parse("2020-01-02T01:11:01"), + TimestampMillisecondType::parse("2020-01-12T01:01:01"), + ])), + expected_null_counts: UInt64Array::from(vec![1, 1, 1, 1]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + column_name: "millis", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new( + TimestampMillisecondArray::from(vec![ + TimestampMillisecondType::parse("2020-01-01T01:01:01"), + TimestampMillisecondType::parse("2020-01-01T01:01:11"), + TimestampMillisecondType::parse("2020-01-01T01:11:01"), + TimestampMillisecondType::parse("2020-01-11T01:01:01"), + ]) + .with_timezone(tz), + ), + expected_max: Arc::new( + TimestampMillisecondArray::from(vec![ + TimestampMillisecondType::parse("2020-01-02T01:01:01"), + TimestampMillisecondType::parse("2020-01-02T01:01:11"), + TimestampMillisecondType::parse("2020-01-02T01:11:01"), + TimestampMillisecondType::parse("2020-01-12T01:01:01"), + ]) + .with_timezone(tz), + ), + // nulls are [1, 1, 1, 1] + expected_null_counts: UInt64Array::from(vec![1, 1, 1, 1]), + // row counts are [5, 5, 5, 5] + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + + column_name: "millis_timezoned", + check: Check::Both, + } + .run(); + + // seconds + Test { + reader: &reader, + expected_min: Arc::new(TimestampSecondArray::from(vec![ + TimestampSecondType::parse("2020-01-01T01:01:01"), + TimestampSecondType::parse("2020-01-01T01:01:11"), + TimestampSecondType::parse("2020-01-01T01:11:01"), + TimestampSecondType::parse("2020-01-11T01:01:01"), + ])), + expected_max: Arc::new(TimestampSecondArray::from(vec![ + TimestampSecondType::parse("2020-01-02T01:01:01"), + TimestampSecondType::parse("2020-01-02T01:01:11"), + TimestampSecondType::parse("2020-01-02T01:11:01"), + TimestampSecondType::parse("2020-01-12T01:01:01"), + ])), + expected_null_counts: UInt64Array::from(vec![1, 1, 1, 1]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + + column_name: "seconds", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new( + TimestampSecondArray::from(vec![ + TimestampSecondType::parse("2020-01-01T01:01:01"), + TimestampSecondType::parse("2020-01-01T01:01:11"), + TimestampSecondType::parse("2020-01-01T01:11:01"), + TimestampSecondType::parse("2020-01-11T01:01:01"), + ]) + .with_timezone(tz), + ), + expected_max: Arc::new( + TimestampSecondArray::from(vec![ + TimestampSecondType::parse("2020-01-02T01:01:01"), + TimestampSecondType::parse("2020-01-02T01:01:11"), + TimestampSecondType::parse("2020-01-02T01:11:01"), + TimestampSecondType::parse("2020-01-12T01:01:01"), + ]) + .with_timezone(tz), + ), + // nulls are [1, 1, 1, 1] + expected_null_counts: UInt64Array::from(vec![1, 1, 1, 1]), + // row counts are [5, 5, 5, 5] + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + + column_name: "seconds_timezoned", + check: Check::Both, + } + .run(); +} + +// timestamp with different row group sizes +#[tokio::test] +async fn test_timestamp_diff_rg_sizes() { + // This creates a parquet files of 9 columns named "nanos", "nanos_timezoned", "micros", "micros_timezoned", "millis", "millis_timezoned", "seconds", "seconds_timezoned", "names" + // "nanos" --> TimestampNanosecondArray + // "nanos_timezoned" --> TimestampNanosecondArray + // "micros" --> TimestampMicrosecondArray + // "micros_timezoned" --> TimestampMicrosecondArray + // "millis" --> TimestampMillisecondArray + // "millis_timezoned" --> TimestampMillisecondArray + // "seconds" --> TimestampSecondArray + // "seconds_timezoned" --> TimestampSecondArray + // "names" --> StringArray + // + // The file is created by 4 record batches (each has a null row), each has 5 rows but then will be split into 3 row groups with size 8, 8, 4 + let reader = TestReader { + scenario: Scenario::Timestamps, + row_per_group: 8, // note that the row group size is 8 + } + .build() + .await; + + let tz = "Pacific/Efate"; + + Test { + reader: &reader, + expected_min: Arc::new(TimestampNanosecondArray::from(vec![ + TimestampNanosecondType::parse("2020-01-01T01:01:01"), + TimestampNanosecondType::parse("2020-01-01T01:11:01"), + TimestampNanosecondType::parse("2020-01-11T01:02:01"), + ])), + expected_max: Arc::new(TimestampNanosecondArray::from(vec![ + TimestampNanosecondType::parse("2020-01-02T01:01:01"), + TimestampNanosecondType::parse("2020-01-11T01:01:01"), + TimestampNanosecondType::parse("2020-01-12T01:01:01"), + ])), + // nulls are [1, 2, 1] + expected_null_counts: UInt64Array::from(vec![1, 2, 1]), + // row counts are [8, 8, 4] + expected_row_counts: Some(UInt64Array::from(vec![8, 8, 4])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "nanos", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new( + TimestampNanosecondArray::from(vec![ + TimestampNanosecondType::parse("2020-01-01T01:01:01"), + TimestampNanosecondType::parse("2020-01-01T01:11:01"), + TimestampNanosecondType::parse("2020-01-11T01:02:01"), + ]) + .with_timezone(tz), + ), + expected_max: Arc::new( + TimestampNanosecondArray::from(vec![ + TimestampNanosecondType::parse("2020-01-02T01:01:01"), + TimestampNanosecondType::parse("2020-01-11T01:01:01"), + TimestampNanosecondType::parse("2020-01-12T01:01:01"), + ]) + .with_timezone(tz), + ), + // nulls are [1, 2, 1] + expected_null_counts: UInt64Array::from(vec![1, 2, 1]), + // row counts are [8, 8, 4] + expected_row_counts: Some(UInt64Array::from(vec![8, 8, 4])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "nanos_timezoned", + check: Check::Both, + } + .run(); + + // micros + Test { + reader: &reader, + expected_min: Arc::new(TimestampMicrosecondArray::from(vec![ + TimestampMicrosecondType::parse("2020-01-01T01:01:01"), + TimestampMicrosecondType::parse("2020-01-01T01:11:01"), + TimestampMicrosecondType::parse("2020-01-11T01:02:01"), + ])), + expected_max: Arc::new(TimestampMicrosecondArray::from(vec![ + TimestampMicrosecondType::parse("2020-01-02T01:01:01"), + TimestampMicrosecondType::parse("2020-01-11T01:01:01"), + TimestampMicrosecondType::parse("2020-01-12T01:01:01"), + ])), + expected_null_counts: UInt64Array::from(vec![1, 2, 1]), + expected_row_counts: Some(UInt64Array::from(vec![8, 8, 4])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "micros", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new( + TimestampMicrosecondArray::from(vec![ + TimestampMicrosecondType::parse("2020-01-01T01:01:01"), + TimestampMicrosecondType::parse("2020-01-01T01:11:01"), + TimestampMicrosecondType::parse("2020-01-11T01:02:01"), + ]) + .with_timezone(tz), + ), + expected_max: Arc::new( + TimestampMicrosecondArray::from(vec![ + TimestampMicrosecondType::parse("2020-01-02T01:01:01"), + TimestampMicrosecondType::parse("2020-01-11T01:01:01"), + TimestampMicrosecondType::parse("2020-01-12T01:01:01"), + ]) + .with_timezone(tz), + ), + // nulls are [1, 2, 1] + expected_null_counts: UInt64Array::from(vec![1, 2, 1]), + // row counts are [8, 8, 4] + expected_row_counts: Some(UInt64Array::from(vec![8, 8, 4])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "micros_timezoned", + check: Check::Both, + } + .run(); + + // millis + Test { + reader: &reader, + expected_min: Arc::new(TimestampMillisecondArray::from(vec![ + TimestampMillisecondType::parse("2020-01-01T01:01:01"), + TimestampMillisecondType::parse("2020-01-01T01:11:01"), + TimestampMillisecondType::parse("2020-01-11T01:02:01"), + ])), + expected_max: Arc::new(TimestampMillisecondArray::from(vec![ + TimestampMillisecondType::parse("2020-01-02T01:01:01"), + TimestampMillisecondType::parse("2020-01-11T01:01:01"), + TimestampMillisecondType::parse("2020-01-12T01:01:01"), + ])), + expected_null_counts: UInt64Array::from(vec![1, 2, 1]), + expected_row_counts: Some(UInt64Array::from(vec![8, 8, 4])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "millis", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new( + TimestampMillisecondArray::from(vec![ + TimestampMillisecondType::parse("2020-01-01T01:01:01"), + TimestampMillisecondType::parse("2020-01-01T01:11:01"), + TimestampMillisecondType::parse("2020-01-11T01:02:01"), + ]) + .with_timezone(tz), + ), + expected_max: Arc::new( + TimestampMillisecondArray::from(vec![ + TimestampMillisecondType::parse("2020-01-02T01:01:01"), + TimestampMillisecondType::parse("2020-01-11T01:01:01"), + TimestampMillisecondType::parse("2020-01-12T01:01:01"), + ]) + .with_timezone(tz), + ), + // nulls are [1, 2, 1] + expected_null_counts: UInt64Array::from(vec![1, 2, 1]), + // row counts are [8, 8, 4] + expected_row_counts: Some(UInt64Array::from(vec![8, 8, 4])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "millis_timezoned", + check: Check::Both, + } + .run(); + + // seconds + Test { + reader: &reader, + expected_min: Arc::new(TimestampSecondArray::from(vec![ + TimestampSecondType::parse("2020-01-01T01:01:01"), + TimestampSecondType::parse("2020-01-01T01:11:01"), + TimestampSecondType::parse("2020-01-11T01:02:01"), + ])), + expected_max: Arc::new(TimestampSecondArray::from(vec![ + TimestampSecondType::parse("2020-01-02T01:01:01"), + TimestampSecondType::parse("2020-01-11T01:01:01"), + TimestampSecondType::parse("2020-01-12T01:01:01"), + ])), + expected_null_counts: UInt64Array::from(vec![1, 2, 1]), + expected_row_counts: Some(UInt64Array::from(vec![8, 8, 4])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "seconds", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new( + TimestampSecondArray::from(vec![ + TimestampSecondType::parse("2020-01-01T01:01:01"), + TimestampSecondType::parse("2020-01-01T01:11:01"), + TimestampSecondType::parse("2020-01-11T01:02:01"), + ]) + .with_timezone(tz), + ), + expected_max: Arc::new( + TimestampSecondArray::from(vec![ + TimestampSecondType::parse("2020-01-02T01:01:01"), + TimestampSecondType::parse("2020-01-11T01:01:01"), + TimestampSecondType::parse("2020-01-12T01:01:01"), + ]) + .with_timezone(tz), + ), + // nulls are [1, 2, 1] + expected_null_counts: UInt64Array::from(vec![1, 2, 1]), + // row counts are [8, 8, 4] + expected_row_counts: Some(UInt64Array::from(vec![8, 8, 4])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "seconds_timezoned", + check: Check::Both, + } + .run(); +} + +// date with different row group sizes +#[tokio::test] +async fn test_dates_32_diff_rg_sizes() { + // This creates a parquet files of 3 columns named "date32", "date64", "names" + // "date32" --> Date32Array + // "date64" --> Date64Array + // "names" --> StringArray + // + // The file is created by 4 record batches (each has a null row), each has 5 rows but then will be split into 2 row groups with size 13, 7 + let reader = TestReader { + scenario: Scenario::Dates, + row_per_group: 13, + } + .build() + .await; + + Test { + reader: &reader, + // mins are [2020-01-01, 2020-10-30] + expected_min: Arc::new(Date32Array::from(vec![ + Date32Type::parse("2020-01-01"), + Date32Type::parse("2020-10-30"), + ])), + // maxes are [2020-10-29, 2029-11-12] + expected_max: Arc::new(Date32Array::from(vec![ + Date32Type::parse("2020-10-29"), + Date32Type::parse("2029-11-12"), + ])), + // nulls are [2, 2] + expected_null_counts: UInt64Array::from(vec![2, 2]), + // row counts are [13, 7] + expected_row_counts: Some(UInt64Array::from(vec![13, 7])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "date32", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_time32_second_diff_rg_sizes() { + let reader = TestReader { + scenario: Scenario::Time32Second, + row_per_group: 4, + } + .build() + .await; + + // Test for Time32Second column + Test { + reader: &reader, + // Assuming specific minimum and maximum values for demonstration + expected_min: Arc::new(Time32SecondArray::from(vec![18506, 18510, 18514, 18518])), + expected_max: Arc::new(Time32SecondArray::from(vec![18509, 18513, 18517, 18521])), + expected_null_counts: UInt64Array::from(vec![0, 0, 0, 0]), // Assuming 1 null per row group for simplicity + expected_row_counts: Some(UInt64Array::from(vec![4, 4, 4, 4])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + column_name: "second", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_time32_millisecond_diff_rg_sizes() { + let reader = TestReader { + scenario: Scenario::Time32Millisecond, + row_per_group: 4, + } + .build() + .await; + + // Test for Time32Millisecond column + Test { + reader: &reader, + // Assuming specific minimum and maximum values for demonstration + expected_min: Arc::new(Time32MillisecondArray::from(vec![ + 3600000, 3600004, 3600008, 3600012, + ])), + expected_max: Arc::new(Time32MillisecondArray::from(vec![ + 3600003, 3600007, 3600011, 3600015, + ])), + expected_null_counts: UInt64Array::from(vec![0, 0, 0, 0]), // Assuming 1 null per row group for simplicity + expected_row_counts: Some(UInt64Array::from(vec![4, 4, 4, 4])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + column_name: "millisecond", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_time64_microsecond_diff_rg_sizes() { + let reader = TestReader { + scenario: Scenario::Time64Microsecond, + row_per_group: 4, + } + .build() + .await; + + // Test for Time64MicroSecond column + Test { + reader: &reader, + // Assuming specific minimum and maximum values for demonstration + expected_min: Arc::new(Time64MicrosecondArray::from(vec![ + 1234567890123, + 1234567890127, + 1234567890131, + 1234567890135, + ])), + expected_max: Arc::new(Time64MicrosecondArray::from(vec![ + 1234567890126, + 1234567890130, + 1234567890134, + 1234567890138, + ])), + expected_null_counts: UInt64Array::from(vec![0, 0, 0, 0]), // Assuming 1 null per row group for simplicity + expected_row_counts: Some(UInt64Array::from(vec![4, 4, 4, 4])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + column_name: "microsecond", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_time64_nanosecond_diff_rg_sizes() { + let reader = TestReader { + scenario: Scenario::Time64Nanosecond, + row_per_group: 4, + } + .build() + .await; + + // Test for Time32Second column + Test { + reader: &reader, + // Assuming specific minimum and maximum values for demonstration + expected_min: Arc::new(Time64NanosecondArray::from(vec![ + 987654321012345, + 987654321012349, + 987654321012353, + 987654321012357, + ])), + expected_max: Arc::new(Time64NanosecondArray::from(vec![ + 987654321012348, + 987654321012352, + 987654321012356, + 987654321012360, + ])), + expected_null_counts: UInt64Array::from(vec![0, 0, 0, 0]), // Assuming 1 null per row group for simplicity + expected_row_counts: Some(UInt64Array::from(vec![4, 4, 4, 4])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + column_name: "nanosecond", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_dates_64_diff_rg_sizes() { + // The file is created by 4 record batches (each has a null row), each has 5 rows but then will be split into 2 row groups with size 13, 7 + let reader = TestReader { + scenario: Scenario::Dates, + row_per_group: 13, + } + .build() + .await; + Test { + reader: &reader, + expected_min: Arc::new(Date64Array::from(vec![ + Date64Type::parse("2020-01-01"), + Date64Type::parse("2020-10-30"), + ])), + expected_max: Arc::new(Date64Array::from(vec![ + Date64Type::parse("2020-10-29"), + Date64Type::parse("2029-11-12"), + ])), + expected_null_counts: UInt64Array::from(vec![2, 2]), + expected_row_counts: Some(UInt64Array::from(vec![13, 7])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "date64", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_uint() { + // This creates a parquet files of 4 columns named "u8", "u16", "u32", "u64" + // "u8" --> UInt8Array + // "u16" --> UInt16Array + // "u32" --> UInt32Array + // "u64" --> UInt64Array + + // The file is created by 4 record batches (each has a null row), each has 5 rows but then will be split into 5 row groups with size 4 + let reader = TestReader { + scenario: Scenario::UInt, + row_per_group: 4, + } + .build() + .await; + + Test { + reader: &reader, + expected_min: Arc::new(UInt8Array::from(vec![0, 1, 4, 7, 251])), + expected_max: Arc::new(UInt8Array::from(vec![3, 4, 6, 250, 254])), + expected_null_counts: UInt64Array::from(vec![0, 0, 0, 0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![4, 4, 4, 4, 4])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true, true]), + column_name: "u8", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new(UInt16Array::from(vec![0, 1, 4, 7, 251])), + expected_max: Arc::new(UInt16Array::from(vec![3, 4, 6, 250, 254])), + expected_null_counts: UInt64Array::from(vec![0, 0, 0, 0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![4, 4, 4, 4, 4])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true, true]), + column_name: "u16", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new(UInt32Array::from(vec![0, 1, 4, 7, 251])), + expected_max: Arc::new(UInt32Array::from(vec![3, 4, 6, 250, 254])), + expected_null_counts: UInt64Array::from(vec![0, 0, 0, 0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![4, 4, 4, 4, 4])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true, true]), + column_name: "u32", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new(UInt64Array::from(vec![0, 1, 4, 7, 251])), + expected_max: Arc::new(UInt64Array::from(vec![3, 4, 6, 250, 254])), + expected_null_counts: UInt64Array::from(vec![0, 0, 0, 0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![4, 4, 4, 4, 4])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true, true]), + column_name: "u64", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_int32_range() { + // This creates a parquet file of 1 column "i" + // file has 2 record batches, each has 2 rows. They will be saved into one row group + let reader = TestReader { + scenario: Scenario::Int32Range, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + expected_min: Arc::new(Int32Array::from(vec![0])), + expected_max: Arc::new(Int32Array::from(vec![300000])), + expected_null_counts: UInt64Array::from(vec![0]), + expected_row_counts: Some(UInt64Array::from(vec![4])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true]), + expected_min_value_exact: BooleanArray::from(vec![true]), + column_name: "i", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_uint32_range() { + // This creates a parquet file of 1 column "u" + // file has 2 record batches, each has 2 rows. They will be saved into one row group + let reader = TestReader { + scenario: Scenario::UInt32Range, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + expected_min: Arc::new(UInt32Array::from(vec![0])), + expected_max: Arc::new(UInt32Array::from(vec![300000])), + expected_null_counts: UInt64Array::from(vec![0]), + expected_row_counts: Some(UInt64Array::from(vec![4])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true]), + expected_min_value_exact: BooleanArray::from(vec![true]), + column_name: "u", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_numeric_limits_unsigned() { + // file has 7 rows, 2 row groups: one with 5 rows, one with 2 rows. + let reader = TestReader { + scenario: Scenario::NumericLimits, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + expected_min: Arc::new(UInt8Array::from(vec![u8::MIN, 100])), + expected_max: Arc::new(UInt8Array::from(vec![100, u8::MAX])), + expected_null_counts: UInt64Array::from(vec![0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 2])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "u8", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new(UInt16Array::from(vec![u16::MIN, 100])), + expected_max: Arc::new(UInt16Array::from(vec![100, u16::MAX])), + expected_null_counts: UInt64Array::from(vec![0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 2])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "u16", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new(UInt32Array::from(vec![u32::MIN, 100])), + expected_max: Arc::new(UInt32Array::from(vec![100, u32::MAX])), + expected_null_counts: UInt64Array::from(vec![0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 2])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "u32", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new(UInt64Array::from(vec![u64::MIN, 100])), + expected_max: Arc::new(UInt64Array::from(vec![100, u64::MAX])), + expected_null_counts: UInt64Array::from(vec![0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 2])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "u64", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_numeric_limits_signed() { + // file has 7 rows, 2 row groups: one with 5 rows, one with 2 rows. + let reader = TestReader { + scenario: Scenario::NumericLimits, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + expected_min: Arc::new(Int8Array::from(vec![i8::MIN, -100])), + expected_max: Arc::new(Int8Array::from(vec![100, i8::MAX])), + expected_null_counts: UInt64Array::from(vec![0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 2])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "i8", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new(Int16Array::from(vec![i16::MIN, -100])), + expected_max: Arc::new(Int16Array::from(vec![100, i16::MAX])), + expected_null_counts: UInt64Array::from(vec![0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 2])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "i16", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new(Int32Array::from(vec![i32::MIN, -100])), + expected_max: Arc::new(Int32Array::from(vec![100, i32::MAX])), + expected_null_counts: UInt64Array::from(vec![0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 2])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "i32", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new(Int64Array::from(vec![i64::MIN, -100])), + expected_max: Arc::new(Int64Array::from(vec![100, i64::MAX])), + expected_null_counts: UInt64Array::from(vec![0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 2])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "i64", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_numeric_limits_float() { + // file has 7 rows, 2 row groups: one with 5 rows, one with 2 rows. + let reader = TestReader { + scenario: Scenario::NumericLimits, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + expected_min: Arc::new(Float32Array::from(vec![f32::MIN, -100.0])), + expected_max: Arc::new(Float32Array::from(vec![100.0, f32::MAX])), + expected_null_counts: UInt64Array::from(vec![0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 2])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "f32", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new(Float64Array::from(vec![f64::MIN, -100.0])), + expected_max: Arc::new(Float64Array::from(vec![100.0, f64::MAX])), + expected_null_counts: UInt64Array::from(vec![0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 2])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "f64", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new(Float32Array::from(vec![-1.0, -100.0])), + expected_max: Arc::new(Float32Array::from(vec![100.0, -100.0])), + expected_null_counts: UInt64Array::from(vec![0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 2])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "f32_nan", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new(Float64Array::from(vec![-1.0, -100.0])), + expected_max: Arc::new(Float64Array::from(vec![100.0, -100.0])), + expected_null_counts: UInt64Array::from(vec![0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 2])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "f64_nan", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_float64() { + // This creates a parquet file of 1 column "f" + // file has 4 record batches, each has 5 rows. They will be saved into 4 row groups + let reader = TestReader { + scenario: Scenario::Float64, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + expected_min: Arc::new(Float64Array::from(vec![-5.0, -4.0, -0.0, 5.0])), + expected_max: Arc::new(Float64Array::from(vec![-1.0, 0.0, 4.0, 9.0])), + expected_null_counts: UInt64Array::from(vec![0, 0, 0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + + column_name: "f", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_float16() { + // This creates a parquet file of 1 column "f" + // file has 4 record batches, each has 5 rows. They will be saved into 4 row groups + let reader = TestReader { + scenario: Scenario::Float16, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + expected_min: Arc::new(Float16Array::from( + vec![-5.0, -4.0, -0.0, 5.0] + .into_iter() + .map(f16::from_f32) + .collect::>(), + )), + expected_max: Arc::new(Float16Array::from( + vec![-1.0, 0.0, 4.0, 9.0] + .into_iter() + .map(f16::from_f32) + .collect::>(), + )), + expected_null_counts: UInt64Array::from(vec![0, 0, 0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true, true]), + + column_name: "f", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_decimal32() { + // This creates a parquet file of 1 column "decimal32_col" with decimal data type and precision 9, scale 2 + // file has 3 record batches, each has 5 rows. They will be saved into 3 row groups + let reader = TestReader { + scenario: Scenario::Decimal32, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + expected_min: Arc::new( + Decimal32Array::from(vec![100, -500, 2000]) + .with_precision_and_scale(9, 2) + .unwrap(), + ), + expected_max: Arc::new( + Decimal32Array::from(vec![600, 600, 6000]) + .with_precision_and_scale(9, 2) + .unwrap(), + ), + expected_null_counts: UInt64Array::from(vec![0, 0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "decimal32_col", + check: Check::Both, + } + .run(); +} +#[tokio::test] +async fn test_decimal64() { + // This creates a parquet file of 1 column "decimal64_col" with decimal data type and precision 9, scale 2 + // file has 3 record batches, each has 5 rows. They will be saved into 3 row groups + let reader = TestReader { + scenario: Scenario::Decimal64, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + expected_min: Arc::new( + Decimal64Array::from(vec![100, -500, 2000]) + .with_precision_and_scale(9, 2) + .unwrap(), + ), + expected_max: Arc::new( + Decimal64Array::from(vec![600, 600, 6000]) + .with_precision_and_scale(9, 2) + .unwrap(), + ), + expected_null_counts: UInt64Array::from(vec![0, 0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "decimal64_col", + check: Check::Both, + } + .run(); +} +#[tokio::test] +async fn test_decimal128() { + // This creates a parquet file of 1 column "decimal128_col" with decimal data type and precision 9, scale 2 + // file has 3 record batches, each has 5 rows. They will be saved into 3 row groups + let reader = TestReader { + scenario: Scenario::Decimal128, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + expected_min: Arc::new( + Decimal128Array::from(vec![100, -500, 2000]) + .with_precision_and_scale(9, 2) + .unwrap(), + ), + expected_max: Arc::new( + Decimal128Array::from(vec![600, 600, 6000]) + .with_precision_and_scale(9, 2) + .unwrap(), + ), + expected_null_counts: UInt64Array::from(vec![0, 0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "decimal128_col", + check: Check::Both, + } + .run(); +} +#[tokio::test] +async fn test_decimal_256() { + // This creates a parquet file of 1 column "decimal256_col" with decimal data type and precicion 9, scale 2 + // file has 3 record batches, each has 5 rows. They will be saved into 3 row groups + let reader = TestReader { + scenario: Scenario::Decimal256, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + expected_min: Arc::new( + Decimal256Array::from(vec![i256::from(100), i256::from(-500), i256::from(2000)]) + .with_precision_and_scale(9, 2) + .unwrap(), + ), + expected_max: Arc::new( + Decimal256Array::from(vec![i256::from(600), i256::from(600), i256::from(6000)]) + .with_precision_and_scale(9, 2) + .unwrap(), + ), + expected_null_counts: UInt64Array::from(vec![0, 0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "decimal256_col", + check: Check::Both, + } + .run(); +} +#[tokio::test] +async fn test_dictionary() { + let reader = TestReader { + scenario: Scenario::Dictionary, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + expected_min: Arc::new(StringArray::from(vec!["abc", "aaa"])), + expected_max: Arc::new(StringArray::from(vec!["def", "fffff"])), + expected_null_counts: UInt64Array::from(vec![1, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 2])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "string_dict_i8", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new(StringArray::from(vec!["abc", "aaa"])), + expected_max: Arc::new(StringArray::from(vec!["def", "fffff"])), + expected_null_counts: UInt64Array::from(vec![1, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 2])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "string_dict_i32", + check: Check::Both, + } + .run(); + + Test { + reader: &reader, + expected_min: Arc::new(Int64Array::from(vec![-100, 0])), + expected_max: Arc::new(Int64Array::from(vec![0, 100])), + expected_null_counts: UInt64Array::from(vec![1, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 2])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "int_dict_i8", + check: Check::Both, + } + .run(); +} + +#[tokio::test] +async fn test_byte() { + // This creates a parquet file of 5 columns + // "name" + // "service_string" + // "service_binary" + // "service_fixedsize" + // "service_large_binary" + + // file has 3 record batches, each has 5 rows. They will be saved into 3 row groups + let reader = TestReader { + scenario: Scenario::ByteArray, + row_per_group: 5, + } + .build() + .await; + + // column "name" + Test { + reader: &reader, + expected_min: Arc::new(StringArray::from(vec![ + "all frontends", + "mixed", + "all backends", + ])), + expected_max: Arc::new(StringArray::from(vec![ + "all frontends", + "mixed", + "all backends", + ])), + expected_null_counts: UInt64Array::from(vec![0, 0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "name", + check: Check::Both, + } + .run(); + + // column "service_string" + Test { + reader: &reader, + expected_min: Arc::new(StringArray::from(vec![ + "frontend five", + "backend one", + "backend eight", + ])), + expected_max: Arc::new(StringArray::from(vec![ + "frontend two", + "frontend six", + "backend six", + ])), + expected_null_counts: UInt64Array::from(vec![0, 0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "service_string", + check: Check::Both, + } + .run(); + + // column "service_binary" + + let expected_service_binary_min_values: Vec<&[u8]> = + vec![b"frontend five", b"backend one", b"backend eight"]; + + let expected_service_binary_max_values: Vec<&[u8]> = + vec![b"frontend two", b"frontend six", b"backend six"]; + + Test { + reader: &reader, + expected_min: Arc::new(BinaryArray::from(expected_service_binary_min_values)), + expected_max: Arc::new(BinaryArray::from(expected_service_binary_max_values)), + expected_null_counts: UInt64Array::from(vec![0, 0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "service_binary", + check: Check::Both, + } + .run(); + + // column "service_fixedsize" + // b"fe1", b"be1", b"be4" + let min_input = vec![vec![102, 101, 49], vec![98, 101, 49], vec![98, 101, 52]]; + // b"fe5", b"fe6", b"be8" + let max_input = vec![vec![102, 101, 55], vec![102, 101, 54], vec![98, 101, 56]]; + + Test { + reader: &reader, + expected_min: Arc::new(FixedSizeBinaryArray::try_from_iter(min_input.into_iter()).unwrap()), + expected_max: Arc::new(FixedSizeBinaryArray::try_from_iter(max_input.into_iter()).unwrap()), + expected_null_counts: UInt64Array::from(vec![0, 0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "service_fixedsize", + check: Check::Both, + } + .run(); + + let expected_service_large_binary_min_values: Vec<&[u8]> = + vec![b"frontend five", b"backend one", b"backend eight"]; + + let expected_service_large_binary_max_values: Vec<&[u8]> = + vec![b"frontend two", b"frontend six", b"backend six"]; + + Test { + reader: &reader, + expected_min: Arc::new(LargeBinaryArray::from( + expected_service_large_binary_min_values, + )), + expected_max: Arc::new(LargeBinaryArray::from( + expected_service_large_binary_max_values, + )), + expected_null_counts: UInt64Array::from(vec![0, 0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "service_large_binary", + check: Check::Both, + } + .run(); +} + +// PeriodsInColumnNames +#[tokio::test] +async fn test_period_in_column_names() { + // This creates a parquet file of 2 columns "name" and "service.name" + // file has 3 record batches, each has 5 rows. They will be saved into 3 row groups + let reader = TestReader { + scenario: Scenario::PeriodsInColumnNames, + row_per_group: 5, + } + .build() + .await; + + // column "name" + Test { + reader: &reader, + expected_min: Arc::new(StringArray::from(vec![ + "HTTP GET / DISPATCH", + "HTTP PUT / DISPATCH", + "HTTP GET / DISPATCH", + ])), + expected_max: Arc::new(StringArray::from(vec![ + "HTTP GET / DISPATCH", + "HTTP PUT / DISPATCH", + "HTTP GET / DISPATCH", + ])), + expected_null_counts: UInt64Array::from(vec![0, 0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "name", + check: Check::Both, + } + .run(); + + // column "service.name" + Test { + reader: &reader, + expected_min: Arc::new(StringArray::from(vec!["frontend", "backend", "backend"])), + expected_max: Arc::new(StringArray::from(vec!["frontend", "frontend", "backend"])), + expected_null_counts: UInt64Array::from(vec![0, 0, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "service.name", + check: Check::Both, + } + .run(); +} + +// Boolean +#[tokio::test] +async fn test_boolean() { + // This creates a parquet files of 1 column named "bool" + // The file is created by 2 record batches each has 5 rows --> 2 row groups + let reader = TestReader { + scenario: Scenario::Boolean, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + expected_min: Arc::new(BooleanArray::from(vec![false, false])), + expected_max: Arc::new(BooleanArray::from(vec![true, false])), + expected_null_counts: UInt64Array::from(vec![1, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "bool", + check: Check::Both, + } + .run(); +} + +// struct array +// BUG +// https://github.com/apache/datafusion/issues/10609 +// Note that: since I have not worked on struct array before, there may be a bug in the test code rather than the real bug in the code +#[ignore] +#[tokio::test] +async fn test_struct() { + // This creates a parquet files of 1 column named "struct" + // The file is created by 1 record batch with 3 rows in the struct array + let reader = TestReader { + scenario: Scenario::StructArray, + row_per_group: 5, + } + .build() + .await; + Test { + reader: &reader, + expected_min: struct_array(vec![(Some(1), Some(6.0), Some(12.0))]), + expected_max: struct_array(vec![(Some(2), Some(8.5), Some(14.0))]), + expected_null_counts: UInt64Array::from(vec![0]), + expected_row_counts: Some(UInt64Array::from(vec![3])), + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "struct", + check: Check::RowGroup, + } + .run(); +} + +// UTF8 +#[tokio::test] +async fn test_utf8() { + let reader = TestReader { + scenario: Scenario::UTF8, + row_per_group: 5, + } + .build() + .await; + + // test for utf8 + Test { + reader: &reader, + expected_min: Arc::new(StringArray::from(vec!["a", "e"])), + expected_max: Arc::new(StringArray::from(vec!["d", "i"])), + expected_null_counts: UInt64Array::from(vec![1, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "utf8", + check: Check::Both, + } + .run(); + + // test for large_utf8 + Test { + reader: &reader, + expected_min: Arc::new(LargeStringArray::from(vec!["a", "e"])), + expected_max: Arc::new(LargeStringArray::from(vec!["d", "i"])), + expected_null_counts: UInt64Array::from(vec![1, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "large_utf8", + check: Check::Both, + } + .run(); +} + +// UTF8View +#[tokio::test] +async fn test_utf8_view() { + let reader = TestReader { + scenario: Scenario::UTF8View, + row_per_group: 5, + } + .build() + .await; + + // test for utf8_view + Test { + reader: &reader, + expected_min: Arc::new(StringViewArray::from(vec!["a", "a", "e_longerthan12"])), + expected_max: Arc::new(StringViewArray::from(vec![ + "d", + "e_longerthan12", + "i_longerthan12", + ])), + expected_null_counts: UInt64Array::from(vec![1, 3, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "utf8_view", + check: Check::Both, + } + .run() +} + +// BinaryView +#[tokio::test] +async fn test_binary_view() { + let reader = TestReader { + scenario: Scenario::BinaryView, + row_per_group: 5, + } + .build() + .await; + + let expected_min: Vec> = vec![Some(b"a"), Some(b"a"), Some(b"e_longerthan12")]; + let expected_max: Vec> = + vec![Some(b"d"), Some(b"e_longerthan12"), Some(b"i_longerthan12")]; + + // test for utf8_view + Test { + reader: &reader, + expected_min: Arc::new(BinaryViewArray::from(expected_min)), + expected_max: Arc::new(BinaryViewArray::from(expected_max)), + expected_null_counts: UInt64Array::from(vec![1, 3, 0]), + expected_row_counts: Some(UInt64Array::from(vec![5, 5, 5])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true, true]), + column_name: "binary_view", + check: Check::Both, + } + .run() +} + +////// Files with missing statistics /////// + +#[tokio::test] +async fn test_missing_statistics() { + let reader = Int64Case { + null_values: 0, + no_null_values_start: 4, + no_null_values_end: 7, + row_per_group: 5, + enable_stats: Some(EnabledStatistics::None), + ..Default::default() + } + .build(); + + Test { + reader: &reader, + expected_min: Arc::new(Int64Array::from(vec![None])), + expected_max: Arc::new(Int64Array::from(vec![None])), + expected_null_counts: UInt64Array::from(vec![None]), + expected_row_counts: Some(UInt64Array::from(vec![3])), // still has row count statistics + expected_max_value_exact: BooleanArray::from(vec![None]), + expected_min_value_exact: BooleanArray::from(vec![None]), + column_name: "i64", + check: Check::Both, + } + .run(); +} + +#[test] +fn missing_null_counts_as_zero() { + let min = None; + let max = None; + let distinct_count = None; + let null_count = None; // NB: no null count + let is_min_max_deprecated = false; + let stats = Statistics::Boolean(ValueStatistics::new( + min, + max, + distinct_count, + null_count, + is_min_max_deprecated, + )); + let (arrow_schema, parquet_schema) = bool_arrow_and_parquet_schema(); + + let column_chunk = ColumnChunkMetaData::builder(parquet_schema.column(0)) + .set_statistics(stats) + .build() + .unwrap(); + let metadata = RowGroupMetaData::builder(parquet_schema.clone()) + .set_column_metadata(vec![column_chunk]) + .build() + .unwrap(); + + let converter = StatisticsConverter::try_new("b", &arrow_schema, &parquet_schema).unwrap(); + + // by default null count should be 0 + assert_eq!( + converter.row_group_null_counts([&metadata]).unwrap(), + UInt64Array::from_iter(vec![Some(0)]) + ); + + // if we disable missing null counts as zero flag null count will be None + let converter = converter.with_missing_null_counts_as_zero(false); + assert_eq!( + converter.row_group_null_counts([&metadata]).unwrap(), + UInt64Array::from_iter(vec![None]) + ); +} + +/// return an Arrow schema and corresponding Parquet SchemaDescriptor for +/// a schema with a single boolean column "b" +fn bool_arrow_and_parquet_schema() -> (SchemaRef, SchemaDescPtr) { + let arrow_schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Boolean, true)])); + use parquet::schema::types::Type as ParquetType; + let parquet_schema = ParquetType::group_type_builder("schema") + .with_fields(vec![Arc::new( + ParquetType::primitive_type_builder("a", parquet::basic::Type::INT32) + .build() + .unwrap(), + )]) + .build() + .unwrap(); + + let parquet_schema = Arc::new(SchemaDescriptor::new(Arc::new(parquet_schema))); + (arrow_schema, parquet_schema) +} + +/////// NEGATIVE TESTS /////// +// column not found +#[tokio::test] +async fn test_column_not_found() { + let reader = TestReader { + scenario: Scenario::Dates, + row_per_group: 5, + } + .build() + .await; + Test { + reader: &reader, + expected_min: Arc::new(Int64Array::from(vec![18262, 18565])), + expected_max: Arc::new(Int64Array::from(vec![18564, 21865])), + expected_null_counts: UInt64Array::from(vec![2, 2]), + expected_row_counts: Some(UInt64Array::from(vec![13, 7])), + // stats are exact + expected_max_value_exact: BooleanArray::from(vec![true, true]), + expected_min_value_exact: BooleanArray::from(vec![true, true]), + column_name: "not_a_column", + check: Check::Both, + } + .run_col_not_found(); +} + +#[tokio::test] +async fn test_column_non_existent() { + // Create a schema with an additional column + // that will not have a matching parquet index + let schema = Arc::new(Schema::new(vec![ + Field::new("i8", DataType::Int8, true), + Field::new("i16", DataType::Int16, true), + Field::new("i32", DataType::Int32, true), + Field::new("i64", DataType::Int64, true), + Field::new("i_do_not_exist", DataType::Int64, true), + ])); + + let reader = TestReader { + scenario: Scenario::Int, + row_per_group: 5, + } + .build() + .await; + + Test { + reader: &reader, + // mins are [-5, -4, 0, 5] + expected_min: Arc::new(Int64Array::from(vec![None, None, None, None])), + // maxes are [-1, 0, 4, 9] + expected_max: Arc::new(Int64Array::from(vec![None, None, None, None])), + // nulls are [0, 0, 0, 0] + expected_null_counts: UInt64Array::from(vec![None, None, None, None]), + // row counts are [5, 5, 5, 5] + expected_row_counts: None, + expected_max_value_exact: BooleanArray::from(vec![None, None, None, None]), + expected_min_value_exact: BooleanArray::from(vec![None, None, None, None]), + column_name: "i_do_not_exist", + check: Check::Both, + } + .run_with_schema(&schema); +} + +/// The following tests were initially in `arrow-rs/parquet/src/arrow/arrow_reader/statistics.rs`. +/// Part of them was moved here to avoid confusion and duplication, +/// including edge conditions and data file validations for only `row_group` statistics. +#[cfg(test)] +mod test { + use super::*; + use arrow::util::test_util::parquet_test_data; + use arrow_array::{ + ArrayRef, BooleanArray, Decimal128Array, Float32Array, Float64Array, Int8Array, Int16Array, + Int32Array, Int64Array, RecordBatch, StringArray, StructArray, TimestampNanosecondArray, + new_empty_array, + }; + use arrow_schema::{DataType, SchemaRef, TimeUnit}; + use bytes::Bytes; + use parquet::arrow::parquet_column; + use parquet::file::metadata::{ParquetMetaData, RowGroupMetaData}; + use std::path::PathBuf; + use std::sync::Arc; + // TODO error cases (with parquet statistics that are mismatched in expected type) + + #[test] + fn roundtrip_empty() { + let all_types = vec![ + DataType::Null, + DataType::Boolean, + DataType::Int8, + DataType::Int16, + DataType::Int32, + DataType::Int64, + DataType::UInt8, + DataType::UInt16, + DataType::UInt32, + DataType::UInt64, + DataType::Float16, + DataType::Float32, + DataType::Float64, + DataType::Timestamp(TimeUnit::Second, None), + DataType::Date32, + DataType::Date64, + // Skip types that don't support statistics + // DataType::Time32(Second), + // DataType::Time64(Second), + // DataType::Duration(Second), + // DataType::Interval(IntervalUnit), + DataType::Binary, + DataType::FixedSizeBinary(0), + DataType::LargeBinary, + DataType::BinaryView, + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + // DataType::List(FieldRef), + // DataType::ListView(FieldRef), + // DataType::FixedSizeList(FieldRef, i32), + // DataType::LargeList(FieldRef), + // DataType::LargeListView(FieldRef), + // DataType::Struct(Fields), + // DataType::Union(UnionFields, UnionMode), + // DataType::Dictionary(Box, Box), + // DataType::Decimal32(u8, i8), + // DataType::Decimal64(u8, i8), + // DataType::Decimal128(u8, i8), + // DataType::Decimal256(u8, i8), + // DataType::Map(FieldRef, bool), + // DataType::RunEndEncoded(FieldRef, FieldRef), + ]; + for data_type in all_types { + let empty_array = new_empty_array(&data_type); + Test { + input: empty_array.clone(), + expected_min: empty_array.clone(), + expected_max: empty_array, + } + .run(); + } + } + + #[test] + fn nan_in_stats() { + // /parquet-testing/data/nan_in_stats.parquet + // row_groups: 1 + // "x": Double({min: Some(1.0), max: Some(NaN), distinct_count: None, null_count: 0, min_max_deprecated: false, min_max_backwards_compatible: false}) + + TestFile::new("nan_in_stats.parquet") + .with_column(ExpectedColumn { + name: "x", + expected_min: Arc::new(Float64Array::from(vec![Some(1.0)])), + expected_max: Arc::new(Float64Array::from(vec![Some(f64::NAN)])), + }) + .run(); + } + + #[test] + fn alltypes_plain() { + // /parquet-testing/data/datapage_v1-snappy-compressed-checksum.parquet + // row_groups: 1 + // (has no statistics) + TestFile::new("alltypes_plain.parquet") + // No column statistics should be read as NULL, but with the right type + .with_column(ExpectedColumn { + name: "id", + expected_min: i32_array([None]), + expected_max: i32_array([None]), + }) + .with_column(ExpectedColumn { + name: "bool_col", + expected_min: bool_array([None]), + expected_max: bool_array([None]), + }) + .run(); + } + + #[test] + fn alltypes_tiny_pages() { + // /parquet-testing/data/alltypes_tiny_pages.parquet + // row_groups: 1 + // "id": Int32({min: Some(0), max: Some(7299), distinct_count: None, null_count: 0, min_max_deprecated: false, min_max_backwards_compatible: false}) + // "bool_col": Boolean({min: Some(false), max: Some(true), distinct_count: None, null_count: 0, min_max_deprecated: false, min_max_backwards_compatible: false}) + // "tinyint_col": Int32({min: Some(0), max: Some(9), distinct_count: None, null_count: 0, min_max_deprecated: false, min_max_backwards_compatible: false}) + // "smallint_col": Int32({min: Some(0), max: Some(9), distinct_count: None, null_count: 0, min_max_deprecated: false, min_max_backwards_compatible: false}) + // "int_col": Int32({min: Some(0), max: Some(9), distinct_count: None, null_count: 0, min_max_deprecated: false, min_max_backwards_compatible: false}) + // "bigint_col": Int64({min: Some(0), max: Some(90), distinct_count: None, null_count: 0, min_max_deprecated: false, min_max_backwards_compatible: false}) + // "float_col": Float({min: Some(0.0), max: Some(9.9), distinct_count: None, null_count: 0, min_max_deprecated: false, min_max_backwards_compatible: false}) + // "double_col": Double({min: Some(0.0), max: Some(90.89999999999999), distinct_count: None, null_count: 0, min_max_deprecated: false, min_max_backwards_compatible: false}) + // "date_string_col": ByteArray({min: Some(ByteArray { data: "01/01/09" }), max: Some(ByteArray { data: "12/31/10" }), distinct_count: None, null_count: 0, min_max_deprecated: false, min_max_backwards_compatible: false}) + // "string_col": ByteArray({min: Some(ByteArray { data: "0" }), max: Some(ByteArray { data: "9" }), distinct_count: None, null_count: 0, min_max_deprecated: false, min_max_backwards_compatible: false}) + // "timestamp_col": Int96({min: None, max: None, distinct_count: None, null_count: 0, min_max_deprecated: true, min_max_backwards_compatible: true}) + // "year": Int32({min: Some(2009), max: Some(2010), distinct_count: None, null_count: 0, min_max_deprecated: false, min_max_backwards_compatible: false}) + // "month": Int32({min: Some(1), max: Some(12), distinct_count: None, null_count: 0, min_max_deprecated: false, min_max_backwards_compatible: false}) + TestFile::new("alltypes_tiny_pages.parquet") + .with_column(ExpectedColumn { + name: "id", + expected_min: i32_array([Some(0)]), + expected_max: i32_array([Some(7299)]), + }) + .with_column(ExpectedColumn { + name: "bool_col", + expected_min: bool_array([Some(false)]), + expected_max: bool_array([Some(true)]), + }) + .with_column(ExpectedColumn { + name: "tinyint_col", + expected_min: i8_array([Some(0)]), + expected_max: i8_array([Some(9)]), + }) + .with_column(ExpectedColumn { + name: "smallint_col", + expected_min: i16_array([Some(0)]), + expected_max: i16_array([Some(9)]), + }) + .with_column(ExpectedColumn { + name: "int_col", + expected_min: i32_array([Some(0)]), + expected_max: i32_array([Some(9)]), + }) + .with_column(ExpectedColumn { + name: "bigint_col", + expected_min: i64_array([Some(0)]), + expected_max: i64_array([Some(90)]), + }) + .with_column(ExpectedColumn { + name: "float_col", + expected_min: f32_array([Some(0.0)]), + expected_max: f32_array([Some(9.9)]), + }) + .with_column(ExpectedColumn { + name: "double_col", + expected_min: f64_array([Some(0.0)]), + expected_max: f64_array([Some(90.89999999999999)]), + }) + .with_column(ExpectedColumn { + name: "date_string_col", + expected_min: utf8_array([Some("01/01/09")]), + expected_max: utf8_array([Some("12/31/10")]), + }) + .with_column(ExpectedColumn { + name: "string_col", + expected_min: utf8_array([Some("0")]), + expected_max: utf8_array([Some("9")]), + }) + // File has no min/max for timestamp_col + .with_column(ExpectedColumn { + name: "timestamp_col", + expected_min: timestamp_nanoseconds_array([None], None), + expected_max: timestamp_nanoseconds_array([None], None), + }) + .with_column(ExpectedColumn { + name: "year", + expected_min: i32_array([Some(2009)]), + expected_max: i32_array([Some(2010)]), + }) + .with_column(ExpectedColumn { + name: "month", + expected_min: i32_array([Some(1)]), + expected_max: i32_array([Some(12)]), + }) + .run(); + } + + #[test] + fn fixed_length_decimal_legacy() { + // /parquet-testing/data/fixed_length_decimal_legacy.parquet + // row_groups: 1 + // "value": FixedLenByteArray({min: Some(FixedLenByteArray(ByteArray { data: Some(ByteBufferPtr { data: b"\0\0\0\0\0\xc8" }) })), max: Some(FixedLenByteArray(ByteArray { data: "\0\0\0\0\t`" })), distinct_count: None, null_count: 0, min_max_deprecated: true, min_max_backwards_compatible: true}) + + TestFile::new("fixed_length_decimal_legacy.parquet") + .with_column(ExpectedColumn { + name: "value", + expected_min: Arc::new( + Decimal128Array::from(vec![Some(200)]) + .with_precision_and_scale(13, 2) + .unwrap(), + ), + expected_max: Arc::new( + Decimal128Array::from(vec![Some(2400)]) + .with_precision_and_scale(13, 2) + .unwrap(), + ), + }) + .run(); + } + + const ROWS_PER_ROW_GROUP: usize = 3; + + /// Writes the input batch into a parquet file, with every every three rows as + /// their own row group, and compares the min/maxes to the expected values + struct Test { + input: ArrayRef, + expected_min: ArrayRef, + expected_max: ArrayRef, + } + + impl Test { + fn run(self) { + let Self { + input, + expected_min, + expected_max, + } = self; + + let input_batch = RecordBatch::try_from_iter([("c1", input)]).unwrap(); + + let schema = input_batch.schema(); + + let metadata = parquet_metadata(schema.clone(), input_batch); + let parquet_schema = metadata.file_metadata().schema_descr(); + + let row_groups = metadata.row_groups(); + + for field in schema.fields() { + if field.data_type().is_nested() { + let lookup = parquet_column(parquet_schema, &schema, field.name()); + assert_eq!(lookup, None); + continue; + } + + let converter = + StatisticsConverter::try_new(field.name(), &schema, parquet_schema).unwrap(); + + assert_eq!(converter.arrow_field(), field.as_ref()); + + let mins = converter.row_group_mins(row_groups.iter()).unwrap(); + assert_eq!( + &mins, + &expected_min, + "Min. Statistics\n\n{}\n\n", + DisplayStats(row_groups) + ); + + let maxes = converter.row_group_maxes(row_groups.iter()).unwrap(); + assert_eq!( + &maxes, + &expected_max, + "Max. Statistics\n\n{}\n\n", + DisplayStats(row_groups) + ); + } + } + } + + #[test] + fn struct_leaf_statistics_from_column_index() { + let leaf_field = Arc::new(Field::new("leaf", DataType::Int32, true)); + let leaf_array: ArrayRef = Arc::new(Int32Array::from(vec![ + Some(3), + Some(1), + None, + Some(9), + Some(4), + Some(6), + ])); + + let amount_array = decimal128_array( + [ + Some(i128::from(100)), + Some(i128::from(-500)), + None, + Some(i128::from(2000)), + Some(i128::from(600)), + Some(i128::from(50)), + ], + 20, + 2, + ); + let amount_field = Arc::new(Field::new("amount", amount_array.data_type().clone(), true)); + + let struct_array = StructArray::from(vec![ + (leaf_field.clone(), leaf_array), + (amount_field.clone(), amount_array), + ]); + let struct_array: ArrayRef = Arc::new(struct_array); + let input_batch = RecordBatch::try_from_iter([("c1", struct_array)]).unwrap(); + + let reader = build_parquet_file( + ROWS_PER_ROW_GROUP * 2, + Some(EnabledStatistics::Page), + Some(ROWS_PER_ROW_GROUP), + vec![input_batch], + ); + + let schema = reader.schema(); + let metadata = reader.metadata(); + let parquet_schema = reader.parquet_schema(); + let row_groups = metadata.row_groups(); + let row_group_indices = [0]; + let column_page_index = metadata + .column_index() + .expect("file should have column page indices"); + let column_offset_index = metadata + .offset_index() + .expect("file should have column offset indices"); + + let DataType::Struct(fields) = schema.field_with_name("c1").unwrap().data_type() else { + unreachable!("c1 must be a struct field") + }; + let leaf_arrow_field = fields[0].as_ref(); + let amount_arrow_field = fields[1].as_ref(); + + assert_eq!(parquet_column(parquet_schema, schema, "c1"), None); + + let leaf_converter = + StatisticsConverter::from_column_index(0, leaf_arrow_field, parquet_schema).unwrap(); + + assert_eq!(leaf_converter.parquet_column_index(), Some(0)); + assert_eq!(leaf_converter.arrow_field(), leaf_arrow_field); + + let leaf_mins = leaf_converter.row_group_mins(row_groups.iter()).unwrap(); + assert_eq!(&leaf_mins, &i32_array([Some(1)])); + + let leaf_maxes = leaf_converter.row_group_maxes(row_groups.iter()).unwrap(); + assert_eq!(&leaf_maxes, &i32_array([Some(9)])); + + let leaf_null_counts = leaf_converter + .row_group_null_counts(row_groups.iter()) + .unwrap(); + assert_eq!(leaf_null_counts, UInt64Array::from(vec![1])); + + let leaf_row_counts = leaf_converter + .row_group_row_counts(row_groups.iter()) + .unwrap(); + assert_eq!(leaf_row_counts, Some(UInt64Array::from(vec![6]))); + + let leaf_page_mins = leaf_converter + .data_page_mins( + column_page_index, + column_offset_index, + row_group_indices.iter(), + ) + .unwrap(); + assert_eq!(&leaf_page_mins, &i32_array([Some(1), Some(4)])); + + let leaf_page_maxes = leaf_converter + .data_page_maxes( + column_page_index, + column_offset_index, + row_group_indices.iter(), + ) + .unwrap(); + assert_eq!(&leaf_page_maxes, &i32_array([Some(3), Some(9)])); + + let leaf_page_null_counts = leaf_converter + .data_page_null_counts( + column_page_index, + column_offset_index, + row_group_indices.iter(), + ) + .unwrap(); + assert_eq!(leaf_page_null_counts, UInt64Array::from(vec![1, 0])); + + let leaf_page_row_counts = leaf_converter + .data_page_row_counts(column_offset_index, row_groups, row_group_indices.iter()) + .unwrap(); + assert_eq!(leaf_page_row_counts, Some(UInt64Array::from(vec![3, 3]))); + + let amount_converter = + StatisticsConverter::from_column_index(1, amount_arrow_field, parquet_schema).unwrap(); + + let amount_mins = amount_converter.row_group_mins(row_groups.iter()).unwrap(); + assert_eq!( + &amount_mins, + &decimal128_array([Some(i128::from(-500))], 20, 2) + ); + + let amount_maxes = amount_converter.row_group_maxes(row_groups.iter()).unwrap(); + assert_eq!( + &amount_maxes, + &decimal128_array([Some(i128::from(2000))], 20, 2) + ); + + let amount_page_mins = amount_converter + .data_page_mins( + column_page_index, + column_offset_index, + row_group_indices.iter(), + ) + .unwrap(); + assert_eq!( + &amount_page_mins, + &decimal128_array([Some(i128::from(-500)), Some(i128::from(50))], 20, 2) + ); + + let amount_page_maxes = amount_converter + .data_page_maxes( + column_page_index, + column_offset_index, + row_group_indices.iter(), + ) + .unwrap(); + assert_eq!( + &amount_page_maxes, + &decimal128_array([Some(i128::from(100)), Some(i128::from(2000))], 20, 2) + ); + + let column_count = parquet_schema.columns().len(); + let err = + StatisticsConverter::from_column_index(column_count, leaf_arrow_field, parquet_schema) + .unwrap_err(); + assert_eq!( + err.to_string(), + format!( + "Arrow: Parquet column index {column_count} out of bounds, column count {column_count}" + ) + ); + } + + fn decimal128_array( + input: impl IntoIterator>, + precision: u8, + scale: i8, + ) -> ArrayRef { + Arc::new( + Decimal128Array::from_iter(input) + .with_precision_and_scale(precision, scale) + .unwrap(), + ) + } + + /// Write the specified batches out as parquet and return the metadata + fn parquet_metadata(schema: SchemaRef, batch: RecordBatch) -> Arc { + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Chunk) + .set_max_row_group_row_count(Some(ROWS_PER_ROW_GROUP)) + .build(); + + let mut buffer = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut buffer, schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let reader = ArrowReaderBuilder::try_new(Bytes::from(buffer)).unwrap(); + reader.metadata().clone() + } + + /// Formats the statistics nicely for display + struct DisplayStats<'a>(&'a [RowGroupMetaData]); + impl std::fmt::Display for DisplayStats<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let row_groups = self.0; + writeln!(f, " row_groups: {}", row_groups.len())?; + for rg in row_groups { + for col in rg.columns() { + if let Some(statistics) = col.statistics() { + writeln!(f, " {}: {:?}", col.column_path(), statistics)?; + } + } + } + Ok(()) + } + } + + struct ExpectedColumn { + name: &'static str, + expected_min: ArrayRef, + expected_max: ArrayRef, + } + + /// Reads statistics out of the specified, and compares them to the expected values + struct TestFile { + file_name: &'static str, + expected_columns: Vec, + } + + impl TestFile { + fn new(file_name: &'static str) -> Self { + Self { + file_name, + expected_columns: Vec::new(), + } + } + + fn with_column(mut self, column: ExpectedColumn) -> Self { + self.expected_columns.push(column); + self + } + + /// Reads the specified parquet file and validates that the expected min/max + /// values for the specified columns are as expected. + fn run(self) { + let path = PathBuf::from(parquet_test_data()).join(self.file_name); + let file = File::open(path).unwrap(); + let reader = ArrowReaderBuilder::try_new(file).unwrap(); + let arrow_schema = reader.schema(); + let metadata = reader.metadata(); + let row_groups = metadata.row_groups(); + let parquet_schema = metadata.file_metadata().schema_descr(); + + for expected_column in self.expected_columns { + let ExpectedColumn { + name, + expected_min, + expected_max, + } = expected_column; + + let converter = + StatisticsConverter::try_new(name, arrow_schema, parquet_schema).unwrap(); + + // test accessors on the converter + let parquet_column_index = + parquet_column(parquet_schema, arrow_schema, name).map(|(idx, _field)| idx); + assert_eq!(converter.parquet_column_index(), parquet_column_index); + assert_eq!(converter.arrow_field().name(), name); + + let actual_min = converter.row_group_mins(row_groups.iter()).unwrap(); + assert_eq!(&expected_min, &actual_min, "column {name}"); + + let actual_max = converter.row_group_maxes(row_groups.iter()).unwrap(); + assert_eq!(&expected_max, &actual_max, "column {name}"); + } + } + } + + fn bool_array(input: impl IntoIterator>) -> ArrayRef { + let array: BooleanArray = input.into_iter().collect(); + Arc::new(array) + } + + fn i8_array(input: impl IntoIterator>) -> ArrayRef { + let array: Int8Array = input.into_iter().collect(); + Arc::new(array) + } + + fn i16_array(input: impl IntoIterator>) -> ArrayRef { + let array: Int16Array = input.into_iter().collect(); + Arc::new(array) + } + + fn i32_array(input: impl IntoIterator>) -> ArrayRef { + let array: Int32Array = input.into_iter().collect(); + Arc::new(array) + } + + fn i64_array(input: impl IntoIterator>) -> ArrayRef { + let array: Int64Array = input.into_iter().collect(); + Arc::new(array) + } + + fn f32_array(input: impl IntoIterator>) -> ArrayRef { + let array: Float32Array = input.into_iter().collect(); + Arc::new(array) + } + + fn f64_array(input: impl IntoIterator>) -> ArrayRef { + let array: Float64Array = input.into_iter().collect(); + Arc::new(array) + } + + fn timestamp_nanoseconds_array( + input: impl IntoIterator>, + timzezone: Option<&str>, + ) -> ArrayRef { + let array: TimestampNanosecondArray = input.into_iter().collect(); + match timzezone { + Some(tz) => Arc::new(array.with_timezone(tz)), + None => Arc::new(array), + } + } + + fn utf8_array<'a>(input: impl IntoIterator>) -> ArrayRef { + let array: StringArray = input + .into_iter() + .map(|s| s.map(|s| s.to_string())) + .collect(); + Arc::new(array) + } +} diff --git a/vendor/parquet/tests/arrow_writer.rs b/vendor/parquet/tests/arrow_writer.rs new file mode 100644 index 0000000..ce65ba0 --- /dev/null +++ b/vendor/parquet/tests/arrow_writer.rs @@ -0,0 +1,564 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Tests for [`ArrowWriter`] + +use std::alloc::{GlobalAlloc, Layout, System}; +use std::cell::Cell; +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::datatypes::{DataType, Field, Schema, SchemaRef}; +use bytes::Bytes; +use parquet::arrow::arrow_writer::{ + ArrowColumnChunk, ArrowRowGroupWriterFactory, ArrowWriterOptions, PageKey, PageStore, + PageStoreArgs, PageStoreFactory, compute_leaves, +}; +use parquet::arrow::{ArrowSchemaConverter, ArrowWriter}; +use parquet::basic::Encoding; +use parquet::errors::Result; +use parquet::file::properties::WriterProperties; +use parquet::file::writer::SerializedFileWriter; + +#[test] +#[should_panic( + expected = "DeltaBitPackDecoder only supports Int32Type, UInt32Type, Int64Type, and UInt64Type" +)] +fn test_delta_bit_pack_type() { + let props = WriterProperties::builder() + .set_column_encoding("col".into(), Encoding::DELTA_BINARY_PACKED) + .build(); + + let record_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "col", + DataType::Float64, + false, + )])), + vec![Arc::new(Float64Array::from_iter_values(vec![1., 2.]))], + ) + .unwrap(); + + let mut buffer = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut buffer, record_batch.schema(), Some(props)).unwrap(); + let _ = writer.write(&record_batch); +} + +// --------------------------------------------------------------------------- +// Heap-memory regression test for the writer's page buffering. +// +// This proves the headline invariant of the pluggable [`PageStore`]: while a +// row group is being written, the heap used to buffer completed pages grows +// with the row group size for the default in-memory store, but stays bounded +// (≈ a few pages per leaf column) once a spilling backend is plugged in. +// +// Peak heap is measured with a thread-local tracking allocator (the same +// pattern used by `parquet/benches/arrow_reader_peak_memory.rs`), so the test +// needs no external profiling dependency. Tracking is thread-local, so the +// measured peak reflects only allocations made on the measuring thread; the +// default `ArrowWriter` is single-threaded, so the writer's buffering all lands +// there. Each measurement resets the peak to the current live baseline and +// reports the delta, so the threads of unrelated tests in this binary do not +// perturb it. +// +// [`PageStore`]: parquet::arrow::arrow_writer::PageStore +// --------------------------------------------------------------------------- + +thread_local! { + static LIVE_BYTES: Cell = const { Cell::new(0) }; + static PEAK_BYTES: Cell = const { Cell::new(0) }; +} + +struct TrackingAllocator { + inner: System, +} + +#[global_allocator] +static GLOBAL: TrackingAllocator = TrackingAllocator { inner: 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)); + }); +} + +#[allow(unsafe_code)] +unsafe impl 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()); + } + 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(); + if new_size > old_size { + add_live_bytes(new_size - old_size); + } else { + subtract_live_bytes(old_size - new_size); + } + } + new_ptr + } +} + +/// Run `f` and return the peak *additional* live heap (bytes) observed on this +/// thread during it — the delta from the live heap when `f` began. +fn peak_heap_bytes(f: impl FnOnce()) -> usize { + let start = LIVE_BYTES.with(Cell::get); + // Reset the peak to the window's baseline so prior allocations don't count. + PEAK_BYTES.with(|peak| peak.set(start)); + f(); + PEAK_BYTES.with(Cell::get).saturating_sub(start) +} + +/// Width of each value in the one "fat" column, in bytes. +const FAT_VALUE_LEN: usize = 4096; +/// Rows per input batch fed to the writer. Kept small so each batch is dropped +/// promptly — only the writer's *buffering* should accumulate, not the input. +const ROWS_PER_BATCH: usize = 64; +/// Number of batches, all funnelled into a single large row group. +const NUM_BATCHES: usize = 64; +/// Total bytes of fat-column payload written (≈ 16 MiB). +const TOTAL_FAT_BYTES: usize = FAT_VALUE_LEN * ROWS_PER_BATCH * NUM_BATCHES; + +/// A wide schema: one fat, high-cardinality binary column (the spill target) +/// plus several tiny integer columns. +fn skewed_schema() -> SchemaRef { + let mut fields = vec![Field::new("fat", DataType::Binary, false)]; + for i in 0..8 { + fields.push(Field::new(format!("small_{i}"), DataType::Int32, false)); + } + Arc::new(Schema::new(fields)) +} + +/// Build one batch of `ROWS_PER_BATCH` rows. The fat column holds unique, +/// high-entropy values (so they neither dictionary-encode nor compress away), +/// derived deterministically from `batch_index`. +fn make_batch(schema: &SchemaRef, batch_index: usize) -> RecordBatch { + let mut fat: Vec = vec![0u8; FAT_VALUE_LEN * ROWS_PER_BATCH]; + // A cheap xorshift fill keyed by the batch index → distinct, incompressible. + let mut state = (batch_index as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1; + for byte in fat.iter_mut() { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + *byte = (state >> 24) as u8; + } + let offsets: Vec = (0..=ROWS_PER_BATCH) + .map(|i| (i * FAT_VALUE_LEN) as i32) + .collect(); + let fat_array = BinaryArray::try_new( + arrow::buffer::OffsetBuffer::new(offsets.into()), + arrow::buffer::Buffer::from_vec(fat), + None, + ) + .unwrap(); + + let mut columns: Vec = vec![Arc::new(fat_array)]; + for c in 0..8 { + let vals: Vec = (0..ROWS_PER_BATCH) + .map(|r| (batch_index * ROWS_PER_BATCH + r + c) as i32) + .collect(); + columns.push(Arc::new(Int32Array::from(vals))); + } + RecordBatch::try_new(schema.clone(), columns).unwrap() +} + +/// Writer properties forcing the whole dataset into a single, uncompressed row +/// group (so the page buffer is the only thing that grows). +fn single_row_group_props() -> WriterProperties { + WriterProperties::builder() + .set_compression(parquet::basic::Compression::UNCOMPRESSED) + // One row group for everything: never auto-flush on row count. + .set_max_row_group_row_count(Some(ROWS_PER_BATCH * NUM_BATCHES * 2)) + .build() +} + +/// Write the full skewed dataset with the given writer options, feeding small +/// batches (each dropped immediately) into one row group. +/// +/// The output is sent to [`io::sink`] so the produced file bytes never live on +/// the heap — the measured peak then reflects only the writer's internal page +/// *buffering*, which is exactly what a [`PageStore`] governs. +fn write_skewed_dataset(options: ArrowWriterOptions) { + let schema = skewed_schema(); + let mut writer = + ArrowWriter::try_new_with_options(std::io::sink(), schema.clone(), options).unwrap(); + for b in 0..NUM_BATCHES { + let batch = make_batch(&schema, b); + writer.write(&batch).unwrap(); + // `batch` dropped here — only the writer's internal buffering persists. + } + writer.close().unwrap(); +} + +/// A spilling [`PageStore`]: one temp file per column chunk. `put` appends the +/// blob and records its `(offset, len)`; `take` seeks and reads it back. The +/// file is unlinked on creation (via [`tempfile::tempfile`]) so it is cleaned up +/// when the store is dropped. This is the canonical "spill completed pages off +/// the heap" backend the design targets. +struct TempFilePageStore { + file: File, + end: u64, + locs: Vec<(u64, usize)>, +} + +impl TempFilePageStore { + fn new() -> Result { + Ok(Self { + file: tempfile::tempfile()?, + end: 0, + locs: Vec::new(), + }) + } +} + +impl PageStore for TempFilePageStore { + fn put(&mut self, value: Bytes) -> Result { + // Always append at the logical end (a prior `take` may have moved the + // OS file cursor). + self.file.seek(SeekFrom::Start(self.end))?; + self.file.write_all(&value)?; + let key = PageKey::new(self.locs.len() as u64); + self.locs.push((self.end, value.len())); + self.end += value.len() as u64; + Ok(key) + } + + fn take(&mut self, key: PageKey) -> Result { + let (offset, len) = self.locs[key.get() as usize]; + let mut buf = vec![0u8; len]; + self.file.seek(SeekFrom::Start(offset))?; + self.file.read_exact(&mut buf)?; + Ok(Bytes::from(buf)) + } +} + +#[derive(Debug, Default)] +struct TempFilePageStoreFactory; + +impl PageStoreFactory for TempFilePageStoreFactory { + fn create(&self, _args: &PageStoreArgs<'_>) -> Result> { + Ok(Box::new(TempFilePageStore::new()?)) + } +} + +/// Rows per batch / batches for the dictionary-column scenario (~4.2M rows). +const DICT_ROWS_PER_BATCH: usize = 8192; +const DICT_NUM_BATCHES: usize = 512; + +/// Write a single, low-cardinality (16 distinct values), high-row-count column +/// as one row group. Such a column stays dictionary-encoded, so its completed +/// data pages would historically pile up in `GenericColumnWriter` until close — +/// the second accumulation point that plain page-buffer spilling does not reach. +fn write_dict_dataset(options: ArrowWriterOptions) { + let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, false)])); + let props = WriterProperties::builder() + .set_compression(parquet::basic::Compression::UNCOMPRESSED) + .set_max_row_group_row_count(Some(DICT_ROWS_PER_BATCH * DICT_NUM_BATCHES * 2)) + .build(); + let options = options.with_properties(props); + let mut writer = + ArrowWriter::try_new_with_options(std::io::sink(), schema.clone(), options).unwrap(); + for b in 0..DICT_NUM_BATCHES { + let vals: Vec = (0..DICT_ROWS_PER_BATCH) + .map(|r| ((b + r) % 16) as i32) + .collect(); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vals))]).unwrap(); + writer.write(&batch).unwrap(); + } + writer.close().unwrap(); +} + +/// All measurements run in one function so they execute sequentially on a single +/// thread — the tracking allocator is thread-local, so running them as separate +/// parallel tests would each see only their own thread's allocations (which is +/// fine), but keeping them together also keeps the in-memory/spill comparison on +/// one consistent baseline. +#[test] +fn page_store_bounds_write_memory() { + let props = single_row_group_props(); + + // Baseline: the default in-memory store buffers the whole row group, so peak + // heap is at least the size of the buffered column data. + let in_memory_peak = peak_heap_bytes(|| { + let opts = ArrowWriterOptions::new().with_properties(props.clone()); + write_skewed_dataset(opts); + }); + + // Spilling: the temp-file store keeps completed pages off the heap, so peak + // heap stays bounded by the in-flight encoder/dictionary buffers plus a page + // or two in flight — independent of the row group size. + let spill_peak = peak_heap_bytes(|| { + let opts = ArrowWriterOptions::new() + .with_properties(props.clone()) + .with_page_store_factory(Arc::new(TempFilePageStoreFactory)); + write_skewed_dataset(opts); + }); + + eprintln!( + "peak heap — in-memory: {:.1} MiB, temp-file spill: {:.1} MiB (total fat payload {:.1} MiB)", + in_memory_peak as f64 / (1024.0 * 1024.0), + spill_peak as f64 / (1024.0 * 1024.0), + TOTAL_FAT_BYTES as f64 / (1024.0 * 1024.0), + ); + + // The in-memory store must hold most of the ~16 MiB of buffered data. + let in_memory_floor = TOTAL_FAT_BYTES * 3 / 4; + assert!( + in_memory_peak >= in_memory_floor, + "expected in-memory peak >= {in_memory_floor} bytes, got {in_memory_peak}" + ); + + // The spilling store must stay near the per-column bound — roughly + // (data_page_size + dict_page_size) per leaf column, ~2 MiB × 9 columns — + // and far below the in-memory baseline. We assert a generous 8 MiB ceiling + // (well under the ~16 MiB row group) to stay robust across platforms. + const SPILL_CEILING: usize = 8 * 1024 * 1024; + assert!( + spill_peak < SPILL_CEILING, + "expected spilling peak < {SPILL_CEILING} bytes (bounded by page/dict size × columns), \ + got {spill_peak}" + ); + assert!( + spill_peak * 2 < in_memory_peak, + "expected spilling peak ({spill_peak}) to be far below the in-memory baseline \ + ({in_memory_peak})" + ); + + // Dictionary-encoded column: completed data pages reach the page writer (and + // thus the store) as they are produced, so spilling bounds them too. + let dict_in_memory = peak_heap_bytes(|| write_dict_dataset(ArrowWriterOptions::new())); + let dict_spill = peak_heap_bytes(|| { + write_dict_dataset( + ArrowWriterOptions::new().with_page_store_factory(Arc::new(TempFilePageStoreFactory)), + ) + }); + eprintln!( + "dict column ({} rows) peak heap — in-memory: {:.2} MiB, temp-file spill: {:.2} MiB", + DICT_ROWS_PER_BATCH * DICT_NUM_BATCHES, + dict_in_memory as f64 / (1024.0 * 1024.0), + dict_spill as f64 / (1024.0 * 1024.0), + ); + assert!( + dict_spill * 2 < dict_in_memory, + "expected dict-column spilling peak ({dict_spill}) to be far below the in-memory \ + baseline ({dict_in_memory}) — dictionary data pages should spill, not accumulate" + ); +} + +/// Number of dictionary-encoded columns written into a single row group. Each +/// produces its own dictionary page, and all of those pages are held at once +/// between every column's `close()` and the row-group splice. Large enough that +/// K × dict_page dominates the single dictionary page a spilling store keeps. +const DICT_NUM_COLUMNS: usize = 16; +/// Distinct values per dictionary column. Moderate cardinality: small enough to +/// stay dictionary-encoded, large enough (with `DICT_VALUE_LEN`) that each +/// dictionary page is ~0.75 MiB. +const DICT_DISTINCT: usize = 2048; +/// Width of each distinct dictionary value, in bytes (→ dictionary page +/// ≈ `DICT_DISTINCT × (DICT_VALUE_LEN + 4)` ≈ 0.75 MiB). +const DICT_VALUE_LEN: usize = 384; +/// Rows per dictionary column (each distinct value repeated twice). +const DICT_ROWS: usize = DICT_DISTINCT * 2; +/// Approximate retained size of one dictionary page (PLAIN: 4-byte length + value). +const DICT_PAGE_BYTES: usize = DICT_DISTINCT * (DICT_VALUE_LEN + 4); + +/// The `DICT_DISTINCT` distinct dictionary values, concatenated. Built once and +/// shared by every column. Constructed outside the measured closure so its bytes +/// sit in the baseline and are not charged to the per-run peak. +fn dict_value_pool() -> Vec { + let mut pool = vec![0u8; DICT_DISTINCT * DICT_VALUE_LEN]; + for d in 0..DICT_DISTINCT { + let slice = &mut pool[d * DICT_VALUE_LEN..(d + 1) * DICT_VALUE_LEN]; + // A deterministic, incompressible blob per distinct value. + let mut state = (d as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1; + for byte in slice.iter_mut() { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + *byte = (state >> 24) as u8; + } + } + pool +} + +/// One binary column of `DICT_ROWS` rows drawn from `pool`, each distinct value +/// repeated twice. +fn dict_column(pool: &[u8]) -> ArrayRef { + let mut data = vec![0u8; DICT_VALUE_LEN * DICT_ROWS]; + let mut offsets: Vec = Vec::with_capacity(DICT_ROWS + 1); + offsets.push(0); + for r in 0..DICT_ROWS { + let distinct = r % DICT_DISTINCT; + data[r * DICT_VALUE_LEN..(r + 1) * DICT_VALUE_LEN] + .copy_from_slice(&pool[distinct * DICT_VALUE_LEN..(distinct + 1) * DICT_VALUE_LEN]); + offsets.push(((r + 1) * DICT_VALUE_LEN) as i32); + } + Arc::new( + BinaryArray::try_new( + arrow::buffer::OffsetBuffer::new(offsets.into()), + arrow::buffer::Buffer::from_vec(data), + None, + ) + .unwrap(), + ) +} + +/// Encode `DICT_NUM_COLUMNS` dictionary columns into a single row group, then +/// splice them in. Columns are encoded and closed one at a time (via the +/// column-writer API), so only the current column's encoder dictionary is +/// resident while it encodes; the pages that accumulate are the completed +/// dictionary pages held in each closed `ArrowColumnChunk` until the splice. +/// +/// With the default in-memory store every closed chunk keeps its ~0.75 MiB +/// dictionary page on the heap, so all `DICT_NUM_COLUMNS` accumulate before the +/// splice (peak ≈ K × dict_page). With a spilling store each dictionary page is +/// pushed off the heap as its column closes, so at most one is ever resident. +/// +/// `page_store_factory` selects the backend (`None` → the default in-memory +/// store). Output goes to [`io::sink`] so produced file bytes never live on the +/// heap and the measured peak reflects only the writer's page buffering. +fn write_dict_columns(page_store_factory: Option>, pool: &[u8]) { + let fields: Vec = (0..DICT_NUM_COLUMNS) + .map(|i| Field::new(format!("k{i}"), DataType::Binary, false)) + .collect(); + let schema = Arc::new(Schema::new(fields)); + + let props = Arc::new( + WriterProperties::builder() + .set_compression(parquet::basic::Compression::UNCOMPRESSED) + // Raise the dictionary-page limit above one dictionary page so every + // column stays dictionary-encoded instead of falling back to PLAIN. + .set_dictionary_page_size_limit(8 * 1024 * 1024) + .build(), + ); + + let parquet_schema = ArrowSchemaConverter::new() + .with_coerce_types(props.coerce_types()) + .convert(&schema) + .unwrap(); + let mut file = + SerializedFileWriter::new(std::io::sink(), parquet_schema.root_schema_ptr(), props) + .unwrap(); + + let mut factory = ArrowRowGroupWriterFactory::new(&file, schema.clone()); + if let Some(f) = page_store_factory { + factory = factory.with_page_store_factory(f); + } + let col_writers = factory.create_column_writers(0).unwrap(); + + // Encode + close each column before starting the next: its input array and + // encoder are dropped at the end of the iteration, leaving only the closed + // chunk (which holds the completed dictionary page) to accumulate. + let mut chunks: Vec = Vec::with_capacity(DICT_NUM_COLUMNS); + for (i, mut writer) in col_writers.into_iter().enumerate() { + let arr = dict_column(pool); + for leaf in compute_leaves(schema.field(i), &arr).unwrap() { + writer.write(&leaf).unwrap(); + } + drop(arr); + chunks.push(writer.close().unwrap()); + } + + // Splice the held chunks into the row group, one page at a time. + let mut row_group_writer = file.next_row_group().unwrap(); + for chunk in chunks { + chunk.append_to_row_group(&mut row_group_writer).unwrap(); + } + row_group_writer.close().unwrap(); + file.close().unwrap(); +} + +/// Writes `DICT_NUM_COLUMNS` moderate-cardinality dictionary columns — each +/// producing a ~0.75 MiB dictionary page — into a single row group via the +/// column-at-a-time API. Columns are encoded and closed one at a time, so the +/// pages that accumulate are the completed dictionary pages held in each closed +/// `ArrowColumnChunk` between every column's `close()` and the row-group splice. +/// +/// With the default in-memory store all K dictionary pages stay resident at once +/// (peak ≈ K × dict_page). With a spilling store each dictionary page is pushed +/// off the heap as its column closes, so at most one is ever resident, keeping +/// the spilling peak far below the in-memory K × dict_page baseline. +#[test] +fn page_store_spills_dictionary_pages() { + // Build the distinct-value pool up front so its bytes sit in the baseline + // and are not charged to either per-run peak below. + let dict_pool = dict_value_pool(); + let dict_in_memory = peak_heap_bytes(|| write_dict_columns(None, &dict_pool)); + let dict_spill = peak_heap_bytes(|| { + write_dict_columns(Some(Arc::new(TempFilePageStoreFactory)), &dict_pool) + }); + let all_dict_pages = DICT_NUM_COLUMNS * DICT_PAGE_BYTES; + eprintln!( + "dict columns ({} cols × ~{:.2} MiB dictionary page = ~{:.1} MiB held at once) peak heap \ + — in-memory: {:.2} MiB, temp-file spill: {:.2} MiB", + DICT_NUM_COLUMNS, + DICT_PAGE_BYTES as f64 / (1024.0 * 1024.0), + all_dict_pages as f64 / (1024.0 * 1024.0), + dict_in_memory as f64 / (1024.0 * 1024.0), + dict_spill as f64 / (1024.0 * 1024.0), + ); + + // The in-memory store must hold most of the K simultaneously-resident + // dictionary pages — confirming the scenario really does accumulate them. + assert!( + dict_in_memory >= all_dict_pages / 2, + "expected in-memory dict peak >= {} bytes (≈ K × dict_page), got {dict_in_memory}", + all_dict_pages / 2 + ); + + // The spilling store keeps at most one dictionary page in flight, so its + // peak stays a small multiple of a single dictionary page — far below the + // in-memory K × dict_page. + assert!( + dict_spill * 2 < dict_in_memory, + "expected dict-column spilling peak ({dict_spill}) to be far below the in-memory \ + baseline ({dict_in_memory}) — the K held dictionary pages should spill, not stay resident" + ); + let spill_ceiling = DICT_PAGE_BYTES * 5; + assert!( + dict_spill < spill_ceiling, + "expected dict-column spilling peak ({dict_spill}) below {spill_ceiling} bytes \ + (a few dictionary pages), not the ~K × dict_page of the in-memory baseline" + ); +} diff --git a/vendor/parquet/tests/arrow_writer_layout.rs b/vendor/parquet/tests/arrow_writer_layout.rs new file mode 100644 index 0000000..15093ed --- /dev/null +++ b/vendor/parquet/tests/arrow_writer_layout.rs @@ -0,0 +1,990 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Tests that the ArrowWriter correctly lays out values into multiple pages + +use arrow::array::{Int32Array, StringArray}; +use arrow::record_batch::RecordBatch; +use arrow_array::builder::{BinaryBuilder, Int32Builder, ListBuilder}; +use arrow_array::types::Int32Type; +use arrow_array::{DictionaryArray, FixedSizeBinaryArray, StringViewArray}; +use bytes::Bytes; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::{ArrowReaderOptions, ParquetRecordBatchReaderBuilder}; +use parquet::basic::{Encoding, PageType}; +use parquet::file::metadata::PageIndexPolicy; +use parquet::file::metadata::ParquetMetaData; +use parquet::file::properties::{EnabledStatistics, ReaderProperties, WriterProperties}; +use parquet::file::reader::SerializedPageReader; +use parquet::schema::types::ColumnPath; +use std::sync::Arc; + +struct Layout { + row_groups: Vec, +} + +struct RowGroup { + columns: Vec, +} + +struct ColumnChunk { + pages: Vec, + dictionary_page: Option, +} + +struct Page { + rows: usize, + compressed_size: usize, + page_header_size: usize, + encoding: Encoding, + page_type: PageType, +} + +struct LayoutTest { + props: WriterProperties, + batches: Vec, + layout: Layout, +} + +fn do_test(test: LayoutTest) { + let mut buf = Vec::with_capacity(1024); + + let mut writer = + ArrowWriter::try_new(&mut buf, test.batches[0].schema(), Some(test.props)).unwrap(); + for batch in test.batches { + writer.write(&batch).unwrap(); + } + writer.close().unwrap(); + let b = Bytes::from(buf); + + // Re-read file to decode column index + let read_options = + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::from(true)); + let reader = + ParquetRecordBatchReaderBuilder::try_new_with_options(b.clone(), read_options).unwrap(); + + assert_layout(&b, reader.metadata().as_ref(), &test.layout); +} + +fn assert_layout(file_reader: &Bytes, meta: &ParquetMetaData, layout: &Layout) { + assert_eq!(meta.row_groups().len(), layout.row_groups.len()); + let iter = meta + .row_groups() + .iter() + .zip(&layout.row_groups) + .zip(meta.offset_index().unwrap()); + + for ((row_group, row_group_layout), offset_index) in iter { + // Check against offset index + assert_eq!(offset_index.len(), row_group_layout.columns.len()); + + for (column_index, column_layout) in offset_index.iter().zip(&row_group_layout.columns) { + assert_eq!( + column_index.page_locations.len(), + column_layout.pages.len(), + "index page count mismatch" + ); + for (idx, (page, page_layout)) in column_index + .page_locations + .iter() + .zip(&column_layout.pages) + .enumerate() + { + assert_eq!( + page.compressed_page_size as usize, + page_layout.compressed_size + page_layout.page_header_size, + "index page {idx} size mismatch" + ); + let next_first_row_index = column_index + .page_locations + .get(idx + 1) + .map(|x| x.first_row_index) + .unwrap_or_else(|| row_group.num_rows()); + + let num_rows = next_first_row_index - page.first_row_index; + assert_eq!( + num_rows as usize, page_layout.rows, + "index page {idx} row count" + ); + } + } + + // Check against page data + assert_eq!( + row_group.columns().len(), + row_group_layout.columns.len(), + "column count mismatch" + ); + + let iter = row_group + .columns() + .iter() + .zip(&row_group_layout.columns) + .enumerate(); + + for (idx, (column, column_layout)) in iter { + let properties = ReaderProperties::builder() + .set_backward_compatible_lz4(false) + .build(); + let page_reader = SerializedPageReader::new_with_properties( + Arc::new(file_reader.clone()), + column, + row_group.num_rows() as usize, + None, + Arc::new(properties), + ) + .unwrap(); + + let pages = page_reader.collect::, _>>().unwrap(); + assert_eq!( + pages.len(), + column_layout.pages.len() + column_layout.dictionary_page.is_some() as usize, + "page {idx} count mismatch" + ); + + let page_layouts = column_layout + .dictionary_page + .iter() + .chain(&column_layout.pages); + + for (page, page_layout) in pages.iter().zip(page_layouts) { + assert_eq!(page.encoding(), page_layout.encoding); + assert_eq!( + page.buffer().len(), + page_layout.compressed_size, + "page {idx} size mismatch" + ); + assert_eq!(page.page_type(), page_layout.page_type); + } + } + } +} + +#[test] +fn test_primitive() { + let array = Arc::new(Int32Array::from_iter_values(0..2000)) as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(1000) + .set_write_batch_size(10) + .set_write_page_header_statistics(true) + .build(); + + // Test spill plain encoding pages + do_test(LayoutTest { + props, + batches: vec![batch.clone()], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + pages: (0..8) + .map(|_| Page { + rows: 250, + page_header_size: 38, + compressed_size: 1000, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + }) + .collect(), + dictionary_page: None, + }], + }], + }, + }); + + // Test spill dictionary + let props = WriterProperties::builder() + .set_dictionary_enabled(true) + .set_dictionary_page_size_limit(1000) + .set_data_page_size_limit(10000) + .set_write_batch_size(10) + .set_write_page_header_statistics(true) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch.clone()], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + pages: vec![ + Page { + rows: 250, + page_header_size: 38, + compressed_size: 258, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }, + Page { + rows: 1750, + page_header_size: 38, + compressed_size: 7000, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + }, + ], + dictionary_page: Some(Page { + rows: 250, + page_header_size: 38, + compressed_size: 1000, + encoding: Encoding::PLAIN, + page_type: PageType::DICTIONARY_PAGE, + }), + }], + }], + }, + }); + + // Test spill dictionary encoded pages + let props = WriterProperties::builder() + .set_dictionary_enabled(true) + .set_dictionary_page_size_limit(10000) + .set_data_page_size_limit(500) + .set_write_batch_size(10) + .set_write_page_header_statistics(true) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch.clone()], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + pages: vec![ + Page { + rows: 400, + page_header_size: 38, + compressed_size: 452, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }, + Page { + rows: 370, + page_header_size: 38, + compressed_size: 472, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }, + Page { + rows: 330, + page_header_size: 38, + compressed_size: 464, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }, + Page { + rows: 330, + page_header_size: 38, + compressed_size: 464, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }, + Page { + rows: 330, + page_header_size: 38, + compressed_size: 464, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }, + Page { + rows: 240, + page_header_size: 38, + compressed_size: 332, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }, + ], + dictionary_page: Some(Page { + rows: 2000, + page_header_size: 38, + compressed_size: 8000, + encoding: Encoding::PLAIN, + page_type: PageType::DICTIONARY_PAGE, + }), + }], + }], + }, + }); + + // Test row count limit + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_row_count_limit(100) + .set_write_batch_size(100) + .set_write_page_header_statistics(true) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + pages: (0..20) + .map(|_| Page { + rows: 100, + page_header_size: 38, + compressed_size: 400, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + }) + .collect(), + dictionary_page: None, + }], + }], + }, + }); +} + +#[test] +fn test_string() { + let array = Arc::new(StringArray::from_iter_values( + (0..2000).map(|x| format!("{x:04}")), + )) as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(1000) + .set_write_batch_size(10) + .set_write_page_header_statistics(true) + .build(); + + // Test spill plain encoding pages + do_test(LayoutTest { + props, + batches: vec![batch.clone()], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + pages: (0..15) + .map(|_| Page { + rows: 130, + page_header_size: 38, + compressed_size: 1040, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + }) + .chain(std::iter::once(Page { + rows: 50, + page_header_size: 37, + compressed_size: 400, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + })) + .collect(), + dictionary_page: None, + }], + }], + }, + }); + + // Test spill dictionary + let props = WriterProperties::builder() + .set_dictionary_enabled(true) + .set_dictionary_page_size_limit(1000) + .set_data_page_size_limit(10000) + .set_write_batch_size(10) + .set_write_page_header_statistics(true) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch.clone()], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + pages: vec![ + Page { + rows: 126, + page_header_size: 38, + compressed_size: 114, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }, + Page { + rows: 1254, + page_header_size: 40, + compressed_size: 10032, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + }, + Page { + rows: 620, + page_header_size: 38, + compressed_size: 4960, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + }, + ], + // The byte-budget chunker sub-batches the dictionary + // phase. The mini-batch deliberately includes the value + // that crosses the 1000-byte limit so the spill triggers + // on this chunk rather than carrying a sliver into the + // next page, giving a 126-row dictionary page at 1008 + // bytes. + dictionary_page: Some(Page { + rows: 126, + page_header_size: 38, + compressed_size: 1008, + encoding: Encoding::PLAIN, + page_type: PageType::DICTIONARY_PAGE, + }), + }], + }], + }, + }); + + // Test spill dictionary encoded pages + let props = WriterProperties::builder() + .set_dictionary_enabled(true) + .set_dictionary_page_size_limit(20000) + .set_data_page_size_limit(500) + .set_write_batch_size(10) + .set_write_page_header_statistics(true) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + pages: vec![ + Page { + rows: 400, + page_header_size: 38, + compressed_size: 452, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }, + Page { + rows: 370, + page_header_size: 38, + compressed_size: 472, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }, + Page { + rows: 330, + page_header_size: 38, + compressed_size: 464, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }, + Page { + rows: 330, + page_header_size: 38, + compressed_size: 464, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }, + Page { + rows: 330, + page_header_size: 38, + compressed_size: 464, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }, + Page { + rows: 240, + page_header_size: 38, + compressed_size: 332, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }, + ], + dictionary_page: Some(Page { + rows: 2000, + page_header_size: 38, + compressed_size: 16000, + encoding: Encoding::PLAIN, + page_type: PageType::DICTIONARY_PAGE, + }), + }], + }], + }, + }); +} + +#[test] +fn test_list() { + let mut list = ListBuilder::new(Int32Builder::new()); + for _ in 0..200 { + let values = list.values(); + for i in 0..8 { + values.append_value(i); + } + list.append(true); + } + let array = Arc::new(list.finish()) as _; + + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_row_count_limit(20) + .set_write_batch_size(3) + .set_write_page_header_statistics(true) + .build(); + + // Test rows not split across pages + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + pages: (0..10) + .map(|_| Page { + rows: 20, + page_header_size: 38, + compressed_size: 672, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + }) + .collect(), + dictionary_page: None, + }], + }], + }, + }); +} + +#[test] +fn test_per_column_data_page_size_limit() { + // Test that per-column page size limits work correctly + // col_a has a small page size limit (500 bytes), col_b uses the default (10000 bytes) + let col_a = Arc::new(Int32Array::from_iter_values(0..2000)) as _; + let col_b = Arc::new(Int32Array::from_iter_values(0..2000)) as _; + let batch = RecordBatch::try_from_iter([("col_a", col_a), ("col_b", col_b)]).unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(10000) + .set_column_data_page_size_limit(ColumnPath::from("col_a"), 500) + .set_write_batch_size(10) + .build(); + + let mut buf = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let b = Bytes::from(buf); + let read_options = + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::from(true)); + let reader = + ParquetRecordBatchReaderBuilder::try_new_with_options(b.clone(), read_options).unwrap(); + let metadata = reader.metadata(); + + // Verify we have one row group with two columns + assert_eq!(metadata.row_groups().len(), 1); + let row_group = &metadata.row_groups()[0]; + assert_eq!(row_group.columns().len(), 2); + + // Get page counts from offset index + let offset_index = metadata.offset_index().unwrap(); + let col_a_page_count = offset_index[0][0].page_locations.len(); + let col_b_page_count = offset_index[0][1].page_locations.len(); + + // col_a should have many more pages than col_b due to smaller page size limit + // col_a: 500 byte limit for 8000 bytes of data -> 16 pages + // col_b: 10000 byte limit for 8000 bytes of data -> 1 page + assert_eq!(col_a_page_count, 16); + assert_eq!(col_b_page_count, 1); +} + +#[test] +fn test_fixed_size_binary() { + // FixedSizeBinary values larger than the data page byte limit. + let value_size = 1024usize; + let num_rows = 64usize; + let values: Vec = (0..num_rows) + .flat_map(|i| vec![i as u8; value_size]) + .collect(); + let array = + Arc::new(FixedSizeBinaryArray::try_new(value_size as i32, values.into(), None).unwrap()) + as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(4096) + .set_write_page_header_statistics(true) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + // 12 pages of 5 values (5 * 1024 = 5120 B, the boundary + // value pushes each page just past the 4096 B limit) plus + // a final page with the remaining 4 values. + pages: (0..12) + .map(|_| Page { + rows: 5, + page_header_size: 157, + compressed_size: 5120, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + }) + .chain(std::iter::once(Page { + rows: 4, + page_header_size: 157, + compressed_size: 4096, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + })) + .collect(), + dictionary_page: None, + }], + }], + }, + }); +} + +#[test] +fn test_dictionary() { + // Arrow `DictionaryArray` input. + let num_rows = 2000; + let dict_values = StringArray::from_iter_values(["alpha", "beta", "gamma", "delta"]); + let keys = Int32Array::from_iter_values((0..num_rows).map(|i| i % 4)); + let array = + Arc::new(DictionaryArray::::try_new(keys, Arc::new(dict_values)).unwrap()) as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(true) + .set_dictionary_page_size_limit(1000) + .set_data_page_size_limit(1000) + .set_write_batch_size(10) + .set_write_page_header_statistics(true) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + pages: vec![Page { + rows: 2000, + page_header_size: 40, + compressed_size: 505, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }], + dictionary_page: Some(Page { + rows: 4, + page_header_size: 38, + compressed_size: 35, + encoding: Encoding::PLAIN, + page_type: PageType::DICTIONARY_PAGE, + }), + }], + }], + }, + }); +} + +#[test] +fn test_large_string() { + // Large `Utf8` values (64 KiB each) with a 16 KiB data page limit. + // + // Each value already exceeds the page byte budget, so the byte-budget + // chunker in `ByteArrayEncoder` (the offsets-buffer scan in + // `count_within_budget_offsets`) cuts one value per page instead of + // buffering the whole ~2 MiB column into a single page. This drives the + // real `ArrowWriter` user path; the lower-level column writer is covered + // by `test_column_writer_caps_page_size_for_large_byte_array_values`. + let value_size = 64 * 1024; + let strings: Vec = (0..32).map(|_| "x".repeat(value_size)).collect(); + let array = Arc::new(StringArray::from(strings)) as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(16 * 1024) + // Disable statistics so page headers stay small and the layout is + // determined purely by the page-splitting logic under test. + .set_statistics_enabled(EnabledStatistics::None) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + // One 64 KiB value per page (4-byte length prefix + value). + pages: (0..32) + .map(|_| Page { + rows: 1, + page_header_size: 21, + compressed_size: 65540, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + }) + .collect(), + dictionary_page: None, + }], + }], + }, + }); +} + +#[test] +fn test_large_string_delta_byte_array_shared_prefix() { + // Regression for https://github.com/apache/arrow-rs/issues/10489, at the + // `ArrowWriter` level the report used. + // + // Same shape as `test_large_string` — 64 KiB values against a 16 KiB + // page limit — but `DELTA_BYTE_ARRAY` and 32 identical values. Expect a + // single page holding all 32 rows and about one value's worth of bytes, + // rather than the 32 pages and ~2 MiB `PLAIN` produces. + let value_size = 64 * 1024; + let strings: Vec = (0..32).map(|_| "x".repeat(value_size)).collect(); + let array = Arc::new(StringArray::from(strings)) as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_encoding(Encoding::DELTA_BYTE_ARRAY) + .set_data_page_size_limit(16 * 1024) + .set_statistics_enabled(EnabledStatistics::None) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + pages: vec![Page { + rows: 32, + page_header_size: 21, + compressed_size: 65696, + encoding: Encoding::DELTA_BYTE_ARRAY, + page_type: PageType::DATA_PAGE, + }], + dictionary_page: None, + }], + }], + }, + }); +} + +#[test] +fn test_large_string_view() { + // Same bytes and expected layout as `test_large_string`, but the input + // is a `Utf8View` array. View arrays expose no contiguous offsets + // buffer, so the arrow writer bounds pages via the view-specific scan + // (`count_within_budget_views`, reading each value's length from the + // low 32 bits of its view word) rather than the offsets scan. This + // confirms that path caps pages identically. + let value_size = 64 * 1024; + let strings: Vec = (0..32).map(|_| "x".repeat(value_size)).collect(); + let array = Arc::new(StringViewArray::from_iter_values( + strings.iter().map(|s| s.as_str()), + )) as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(16 * 1024) + .set_statistics_enabled(EnabledStatistics::None) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + pages: (0..32) + .map(|_| Page { + rows: 1, + page_header_size: 21, + compressed_size: 65540, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + }) + .collect(), + dictionary_page: None, + }], + }], + }, + }); +} + +#[test] +fn test_large_values_in_list() { + // `list` with large leaf values, driving the record-by-record + // (Materialized-rep) arm of the byte-budget chunker through the arrow + // path. Because repetition levels are present, a list element's leaves + // can never span pages, so the chunker steps from one `rep == 0` + // boundary to the next. Three records of three 32 KiB blobs each, with + // a 16 KiB page limit, yield exactly one page per record (a whole + // record's ~96 KiB of leaves stays together even though it blows the + // budget — it cannot be split). The raw-writer analogue is + // `test_column_writer_caps_page_size_for_large_values_in_list`. + let value_size = 32 * 1024; + let mut builder = ListBuilder::new(BinaryBuilder::new()); + let mut byte = 0u8; + for _ in 0..3 { + for _ in 0..3 { + builder.values().append_value(vec![byte; value_size]); + byte = byte.wrapping_add(1); + } + builder.append(true); + } + let array = Arc::new(builder.finish()) as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(16 * 1024) + .set_statistics_enabled(EnabledStatistics::None) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + // One record (3 leaves + rep/def levels) per page. + pages: (0..3) + .map(|_| Page { + rows: 1, + page_header_size: 21, + compressed_size: 98328, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + }) + .collect(), + dictionary_page: None, + }], + }], + }, + }); +} + +#[test] +fn test_nullable_large_values() { + // Nullable `Utf8` column alternating null / 32 KiB value. The byte + // budget must count only the non-null values (the def-level value + // count), not the level count — otherwise the estimate would be wrong + // for sparse columns. With a 16 KiB page limit each 32 KiB value still + // gets its own page (carrying its adjacent leading null), giving 16 + // two-row pages. Mirrors the raw-writer + // `test_column_writer_caps_page_size_with_nullable_large_values`. + let value_size = 32 * 1024; + let big = "x".repeat(value_size); + let values: Vec> = (0..32) + .map(|i| if i % 2 == 1 { Some(big.clone()) } else { None }) + .collect(); + let array = Arc::new(StringArray::from(values)) as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_size_limit(16 * 1024) + .set_statistics_enabled(EnabledStatistics::None) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + // Each page holds one null + one 32 KiB value (2 rows). + pages: (0..16) + .map(|_| Page { + rows: 2, + page_header_size: 21, + compressed_size: 32778, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + }) + .collect(), + dictionary_page: None, + }], + }], + }, + }); +} + +#[test] +fn test_dictionary_spill_large_values() { + // Dictionary encoding is enabled, but each value is large (64 KiB) and + // unique, so the dictionary spills almost immediately (1 KiB dict page + // limit). After the spill, plain encoding takes over and the byte-budget + // sub-batch bounds each page to a single value. The first value is + // interned into the dictionary page (one RLE_DICTIONARY data page + // referencing it); the remaining 31 fall back to PLAIN, one per page. + // Mirrors `test_column_writer_dict_enabled_large_values_post_spill`, + // exercising the same dict→plain transition via the arrow path. + let value_size = 64 * 1024; + let strings: Vec = (0..32) + .map(|i| format!("{i:05}") + &"x".repeat(value_size - 5)) + .collect(); + let array = Arc::new(StringArray::from(strings)) as _; + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(true) + // Tiny dict limit so it spills after the first (already oversized) + // value, leaving the rest to exercise the post-spill plain path. + .set_dictionary_page_size_limit(1024) + .set_data_page_size_limit(16 * 1024) + .set_write_batch_size(4) + .set_statistics_enabled(EnabledStatistics::None) + .build(); + + do_test(LayoutTest { + props, + batches: vec![batch], + layout: Layout { + row_groups: vec![RowGroup { + columns: vec![ColumnChunk { + pages: std::iter::once(Page { + // The single value interned before the dict spilled. + rows: 1, + page_header_size: 17, + compressed_size: 2, + encoding: Encoding::RLE_DICTIONARY, + page_type: PageType::DATA_PAGE, + }) + .chain((0..31).map(|_| Page { + // Post-spill plain-encoded values, one per page. + rows: 1, + page_header_size: 21, + compressed_size: 65540, + encoding: Encoding::PLAIN, + page_type: PageType::DATA_PAGE, + })) + .collect(), + dictionary_page: Some(Page { + rows: 1, + page_header_size: 21, + compressed_size: 65540, + encoding: Encoding::PLAIN, + page_type: PageType::DICTIONARY_PAGE, + }), + }], + }], + }, + }); +} diff --git a/vendor/parquet/tests/encryption/encryption.rs b/vendor/parquet/tests/encryption/encryption.rs new file mode 100644 index 0000000..2b4a438 --- /dev/null +++ b/vendor/parquet/tests/encryption/encryption.rs @@ -0,0 +1,1509 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! This module contains tests for reading encrypted Parquet files with the Arrow API + +use crate::encryption_util; +use crate::encryption_util::{ + AES_128_COLUMN_KEYS, AES_128_COLUMN_NAME_KEYS, AES_128_COLUMN_NAMES, AES_128_FOOTER_KEY, + AES_128_FOOTER_KEY_NAME, AES_128_KEY_NAME_KEY, AES_128_KEY_NAMES, AES_256_COLUMN_KEYS, + AES_256_COLUMN_NAME_KEYS, AES_256_COLUMN_NAMES, AES_256_FOOTER_KEY, AES_256_FOOTER_KEY_NAME, + AES_256_KEY_NAME_KEY, AES_256_KEY_NAMES, BAD_AES_128_FOOTER_KEY, BAD_AES_256_FOOTER_KEY, + TestKeyRetriever, read_and_roundtrip_to_encrypted_file, verify_column_indexes, + verify_encryption_test_file_read, +}; +use arrow::array::*; +use arrow::error::Result as ArrowResult; +use arrow_array::{Int32Array, RecordBatch}; +use arrow_schema::{DataType as ArrowDataType, DataType, Field, Schema}; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::{ + ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReaderBuilder, RowSelection, + RowSelector, +}; +use parquet::data_type::{ByteArray, ByteArrayType}; +use parquet::encryption::decrypt::FileDecryptionProperties; +use parquet::encryption::encrypt::FileEncryptionProperties; +use parquet::errors::ParquetError; +use parquet::file::metadata::{ColumnChunkMetaData, PageIndexPolicy, ParquetMetaData}; +use parquet::file::properties::WriterProperties; +use parquet::file::writer::SerializedFileWriter; +use parquet::schema::parser::parse_message_type; +use std::fs::File; +use std::sync::Arc; + +#[test] +fn test_non_uniform_encryption_plaintext_footer() { + fn non_uniform_encryption_plaintext_footer(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_plaintext_footer.parquet.encrypted", + ); + let file = File::open(path).unwrap(); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + verify_encryption_test_file_read(file, decryption_properties); + } + + // AES-128: there is always a footer key even with a plaintext footer, + // but this is used for signing the footer. + non_uniform_encryption_plaintext_footer(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); + + // AES-256 + non_uniform_encryption_plaintext_footer(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS); +} + +#[test] +fn test_plaintext_footer_signature_verification() { + fn plaintext_footer_signature_verification(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_plaintext_footer.parquet.encrypted", + ); + let file = File::open(path.clone()).unwrap(); + let mut disable_footer_signature_verification_builder = + FileDecryptionProperties::builder(footer_key.to_vec()) + .disable_footer_signature_verification(); + for (column_name, key) in column_keys { + disable_footer_signature_verification_builder = + disable_footer_signature_verification_builder + .with_column_key(column_name, key.to_vec()); + } + let decryption_properties = disable_footer_signature_verification_builder + .build() + .unwrap(); + verify_encryption_test_file_read(file, decryption_properties); + + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + + let file = File::open(path.clone()).unwrap(); + let decryption_properties = builder.build().unwrap(); + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let result = ArrowReaderMetadata::load(&file, options.clone()); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .starts_with("Parquet error: Footer signature verification failed. Computed: [") + ); + } + + plaintext_footer_signature_verification(BAD_AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); + plaintext_footer_signature_verification(BAD_AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS) +} + +#[test] +fn test_non_uniform_encryption_disabled_aad_storage() { + fn non_uniform_encryption_disabled_aad_storage( + footer_key: &[u8], + column_keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer_disable_aad_storage.parquet.encrypted", + ); + + // Can read successfully when providing the correct AAD prefix + let file = File::open(path.clone()).unwrap(); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()) + .with_aad_prefix(b"tester".to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + verify_encryption_test_file_read(file, decryption_properties); + + // Using wrong AAD prefix should fail + let file = File::open(path.clone()).unwrap(); + let mut wrong_aad_builder = FileDecryptionProperties::builder(footer_key.to_vec()) + .with_aad_prefix(b"wrong_aad_prefix".to_vec()); + for (column_name, key) in column_keys { + wrong_aad_builder = wrong_aad_builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = wrong_aad_builder.build().unwrap(); + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let result = ArrowReaderMetadata::load(&file, options.clone()); + assert!(result.is_err()); + std::assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Provided footer key and AAD were unable to decrypt parquet footer" + ); + + // Not providing any AAD prefix should fail as it isn't stored in the file + let mut no_aad_builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + no_aad_builder = no_aad_builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = no_aad_builder.build().unwrap(); + + let file = File::open(path).unwrap(); + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let result = ArrowReaderMetadata::load(&file, options.clone()); + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file was encrypted with an AAD prefix that is not stored in the file, \ + but no AAD prefix was provided in the file decryption properties" + ); + } + + non_uniform_encryption_disabled_aad_storage(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); + non_uniform_encryption_disabled_aad_storage(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS); +} + +#[test] +#[cfg(feature = "snap")] +fn test_plaintext_footer_read_without_decryption() { + crate::encryption_agnostic::read_plaintext_footer_file_without_decryption_properties(); +} + +#[test] +fn test_non_uniform_encryption() { + fn non_uniform_encryption(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let file = File::open(path).unwrap(); + + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + + verify_encryption_test_file_read(file, decryption_properties); + } + + non_uniform_encryption(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); + non_uniform_encryption(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS); +} + +#[test] +fn test_uniform_encryption() { + fn uniform_encryption(footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let file = File::open(path).unwrap(); + + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); + + verify_encryption_test_file_read(file, decryption_properties); + } + + uniform_encryption(AES_128_FOOTER_KEY); + uniform_encryption(AES_256_FOOTER_KEY); +} + +#[test] +fn test_decrypting_without_decryption_properties_fails() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/uniform_encryption.parquet.encrypted"); + let file = File::open(path).unwrap(); + + let options = ArrowReaderOptions::default(); + let result = ArrowReaderMetadata::load(&file, options.clone()); + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" + ); +} + +#[test] +fn test_aes_ctr_encryption() { + fn aes_ctr_encryption(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer_ctr.parquet.encrypted", + ); + let file = File::open(path).unwrap(); + + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + + let options = + ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); + let metadata = ArrowReaderMetadata::load(&file, options); + + match metadata { + Err(ParquetError::NYI(s)) => { + assert!(s.contains("AES_GCM_CTR_V1")); + } + _ => { + panic!("Expected ParquetError::NYI"); + } + }; + } + + aes_ctr_encryption(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); + aes_ctr_encryption(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS); +} + +#[test] +fn test_non_uniform_encryption_plaintext_footer_with_key_retriever() { + fn non_uniform_encryption_plaintext_footer_with_key_retriever( + footer_key: &[u8], + keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_plaintext_footer.parquet.encrypted", + ); + let file = File::open(path).unwrap(); + + let mut key_retriever = TestKeyRetriever::new(); + for (key_name, key) in keys { + key_retriever = key_retriever.with_key((*key_name).to_owned(), (*key).to_vec()); + } + + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read(file, decryption_properties); + } + + non_uniform_encryption_plaintext_footer_with_key_retriever( + AES_128_FOOTER_KEY, + AES_128_KEY_NAME_KEY, + ); + + non_uniform_encryption_plaintext_footer_with_key_retriever( + AES_256_FOOTER_KEY, + AES_256_KEY_NAME_KEY, + ); +} + +#[test] +fn test_roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever() { + fn roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever( + footer_key: &[u8], + footer_key_metadata: &str, + wrong_footer_key: &[u8], + dec_column_keys: &[(&str, &[u8])], + enc_column_keys: &[(&str, &[u8], &str)], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_plaintext_footer.parquet.encrypted", + ); + let file = File::open(path).unwrap(); + + let mut key_retriever = + TestKeyRetriever::new().with_key(footer_key_metadata.to_owned(), footer_key.to_vec()); + + for (key_name, key) in dec_column_keys { + key_retriever = key_retriever.with_key((*key_name).to_owned(), (*key).to_vec()); + } + + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + let options = ArrowReaderOptions::default() + .with_file_decryption_properties(decryption_properties.clone()); + let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); + + // Write data into temporary file with plaintext footer and footer key metadata + let temp_file = tempfile::tempfile().unwrap(); + let mut encryption_properties_builder = + FileEncryptionProperties::builder(footer_key.to_vec()) + .with_footer_key_metadata(footer_key_metadata.into()); + for (column_name, key, metadata) in enc_column_keys { + encryption_properties_builder = encryption_properties_builder + .with_column_key_and_metadata(column_name, key.to_vec(), (*metadata).into()); + } + let encryption_properties = encryption_properties_builder + .with_plaintext_footer(true) + .build() + .unwrap(); + + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap(); + let batch_reader = builder.build().unwrap(); + let batches = batch_reader + .collect::, _>>() + .unwrap(); + + let props = WriterProperties::builder() + .with_file_encryption_properties(encryption_properties) + .build(); + + let mut writer = ArrowWriter::try_new( + temp_file.try_clone().unwrap(), + metadata.schema().clone(), + Some(props), + ) + .unwrap(); + for batch in batches { + writer.write(&batch).unwrap(); + } + + writer.close().unwrap(); + + // Read temporary file with plaintext metadata using key retriever + let options = ArrowReaderOptions::default() + .with_file_decryption_properties(decryption_properties.clone()); + let _ = ArrowReaderMetadata::load(&temp_file, options.clone()).unwrap(); + + // Read temporary file with plaintext metadata using key retriever with invalid key + let mut key_retriever = TestKeyRetriever::new() + .with_key(footer_key_metadata.to_owned(), wrong_footer_key.to_vec()); + + for (key_name, key) in dec_column_keys { + key_retriever = key_retriever.with_key((*key_name).to_owned(), (*key).to_vec()); + } + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let result = ArrowReaderMetadata::load(&temp_file, options.clone()); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .starts_with("Parquet error: Footer signature verification failed. Computed: [") + ); + } + + roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever( + AES_128_FOOTER_KEY, + AES_128_FOOTER_KEY_NAME, + BAD_AES_128_FOOTER_KEY, + &[ + (AES_128_KEY_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_KEY_NAMES[1], AES_128_COLUMN_KEYS[1]), + ], + &[ + ( + AES_128_COLUMN_NAMES[0], + AES_128_COLUMN_KEYS[0], + AES_128_KEY_NAMES[0], + ), + ( + AES_128_COLUMN_NAMES[1], + AES_128_COLUMN_KEYS[1], + AES_128_KEY_NAMES[1], + ), + ], + ); + + roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever( + AES_256_FOOTER_KEY, + AES_256_FOOTER_KEY_NAME, + BAD_AES_256_FOOTER_KEY, + &[ + (AES_256_KEY_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_KEY_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_KEY_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_KEY_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_KEY_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_KEY_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_KEY_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_KEY_NAMES[7], AES_256_COLUMN_KEYS[7]), + ], + &[ + ( + AES_256_COLUMN_NAMES[0], + AES_256_COLUMN_KEYS[0], + AES_256_KEY_NAMES[0], + ), + ( + AES_256_COLUMN_NAMES[1], + AES_256_COLUMN_KEYS[1], + AES_256_KEY_NAMES[1], + ), + ( + AES_256_COLUMN_NAMES[2], + AES_256_COLUMN_KEYS[2], + AES_256_KEY_NAMES[2], + ), + ( + AES_256_COLUMN_NAMES[3], + AES_256_COLUMN_KEYS[3], + AES_256_KEY_NAMES[3], + ), + ( + AES_256_COLUMN_NAMES[4], + AES_256_COLUMN_KEYS[4], + AES_256_KEY_NAMES[4], + ), + ( + AES_256_COLUMN_NAMES[5], + AES_256_COLUMN_KEYS[5], + AES_256_KEY_NAMES[5], + ), + ( + AES_256_COLUMN_NAMES[6], + AES_256_COLUMN_KEYS[6], + AES_256_KEY_NAMES[6], + ), + ( + AES_256_COLUMN_NAMES[7], + AES_256_COLUMN_KEYS[7], + AES_256_KEY_NAMES[7], + ), + ], + ); +} + +#[test] +fn test_non_uniform_encryption_with_key_retriever() { + fn non_uniform_encryption_with_key_retriever(footer_key: &[u8], keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let file = File::open(path).unwrap(); + + let mut key_retriever = TestKeyRetriever::new(); + for (key_name, key) in keys { + key_retriever = key_retriever.with_key((*key_name).to_owned(), (*key).to_vec()); + } + + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read(file, decryption_properties); + } + + non_uniform_encryption_with_key_retriever(AES_128_FOOTER_KEY, AES_128_KEY_NAME_KEY); + + non_uniform_encryption_with_key_retriever(AES_256_FOOTER_KEY, AES_256_KEY_NAME_KEY); +} + +#[test] +fn test_uniform_encryption_with_key_retriever() { + fn uniform_encryption_with_key_retriever(key_name: &str, footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let file = File::open(path).unwrap(); + + let key_retriever = + TestKeyRetriever::new().with_key(key_name.to_owned(), footer_key.to_vec()); + + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read(file, decryption_properties); + } + + uniform_encryption_with_key_retriever(AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY); + uniform_encryption_with_key_retriever(AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY); +} + +fn row_group_sizes(metadata: &ParquetMetaData) -> Vec { + metadata.row_groups().iter().map(|x| x.num_rows()).collect() +} + +#[test] +fn test_uniform_encryption_roundtrip() { + uniform_encryption_roundtrip(false, false).unwrap(); +} + +#[test] +fn test_uniform_encryption_roundtrip_with_dictionary() { + uniform_encryption_roundtrip(false, true).unwrap(); +} + +#[test] +fn test_uniform_encryption_roundtrip_with_page_index() { + uniform_encryption_roundtrip(true, false).unwrap(); +} + +#[test] +fn test_uniform_encryption_roundtrip_with_page_index_and_dictionary() { + uniform_encryption_roundtrip(true, true).unwrap(); +} + +fn uniform_encryption_roundtrip( + page_index: bool, + dictionary_encoding: bool, +) -> parquet::errors::Result<()> { + let x0_arrays = [ + Int32Array::from((0..100).collect::>()), + Int32Array::from((100..150).collect::>()), + ]; + let x1_arrays = [ + Int32Array::from((100..200).collect::>()), + Int32Array::from((200..250).collect::>()), + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("x0", ArrowDataType::Int32, false), + Field::new("x1", ArrowDataType::Int32, false), + ])); + + let file = tempfile::tempfile()?; + + let footer_key = AES_128_FOOTER_KEY; + let file_encryption_properties = + FileEncryptionProperties::builder(footer_key.to_vec()).build()?; + + let props = WriterProperties::builder() + // Ensure multiple row groups + .set_max_row_group_row_count(Some(50)) + // Ensure multiple pages per row group + .set_write_batch_size(20) + .set_data_page_row_count_limit(20) + .set_dictionary_enabled(dictionary_encoding) + .with_file_encryption_properties(file_encryption_properties) + .build(); + + let mut writer = ArrowWriter::try_new(file.try_clone()?, schema.clone(), Some(props))?; + + for (x0, x1) in x0_arrays.into_iter().zip(x1_arrays) { + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(x0), Arc::new(x1)])?; + writer.write(&batch)?; + } + + writer.close()?; + + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()).build()?; + + let options = ArrowReaderOptions::new() + .with_file_decryption_properties(decryption_properties) + .with_page_index_policy(PageIndexPolicy::from(page_index)); + + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options)?; + assert_eq!(&row_group_sizes(builder.metadata()), &[50, 50, 50]); + + let batches = builder + .with_batch_size(100) + .build()? + .collect::>>()?; + + assert_eq!(batches.len(), 2); + assert!(batches.iter().all(|x| x.num_columns() == 2)); + + let batch_sizes: Vec<_> = batches.iter().map(|x| x.num_rows()).collect(); + + assert_eq!(&batch_sizes, &[100, 50]); + + let x0_values: Vec<_> = batches + .iter() + .flat_map(|x| { + x.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .cloned() + }) + .collect(); + + let x1_values: Vec<_> = batches + .iter() + .flat_map(|x| { + x.column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .cloned() + }) + .collect(); + + let expected_x0_values: Vec<_> = (0..150).collect(); + assert_eq!(&x0_values, &expected_x0_values); + + let expected_x1_values: Vec<_> = (100..250).collect(); + assert_eq!(&x1_values, &expected_x1_values); + Ok(()) +} + +#[test] +fn test_uniform_encryption_page_skipping() { + uniform_encryption_page_skipping(false).unwrap(); +} + +#[test] +fn test_uniform_encryption_page_skipping_with_page_index() { + uniform_encryption_page_skipping(true).unwrap(); +} + +fn uniform_encryption_page_skipping(page_index: bool) -> parquet::errors::Result<()> { + let x0_arrays = [ + Int32Array::from((0..100).collect::>()), + Int32Array::from((100..150).collect::>()), + ]; + let x1_arrays = [ + Int32Array::from((100..200).collect::>()), + Int32Array::from((200..250).collect::>()), + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("x0", ArrowDataType::Int32, false), + Field::new("x1", ArrowDataType::Int32, false), + ])); + + let file = tempfile::tempfile()?; + + let footer_key = AES_128_FOOTER_KEY; + let file_encryption_properties = + FileEncryptionProperties::builder(footer_key.to_vec()).build()?; + + let props = WriterProperties::builder() + // Ensure multiple row groups + .set_max_row_group_row_count(Some(50)) + // Ensure multiple pages per row group + .set_write_batch_size(20) + .set_data_page_row_count_limit(20) + .with_file_encryption_properties(file_encryption_properties) + .build(); + + let mut writer = ArrowWriter::try_new(file.try_clone()?, schema.clone(), Some(props))?; + + for (x0, x1) in x0_arrays.into_iter().zip(x1_arrays) { + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(x0), Arc::new(x1)])?; + writer.write(&batch)?; + } + + writer.close()?; + + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()).build()?; + + let options = ArrowReaderOptions::new() + .with_file_decryption_properties(decryption_properties) + .with_page_index_policy(PageIndexPolicy::from(page_index)); + + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options)?; + + let selection = RowSelection::from(vec![ + RowSelector::skip(25), + RowSelector::select(50), + RowSelector::skip(25), + RowSelector::select(25), + RowSelector::skip(25), + ]); + + let batches = builder + .with_row_selection(selection) + .with_batch_size(100) + .build()? + .collect::>>()?; + + assert_eq!(batches.len(), 1); + assert!(batches.iter().all(|x| x.num_columns() == 2)); + + let batch_sizes: Vec<_> = batches.iter().map(|x| x.num_rows()).collect(); + + assert_eq!(&batch_sizes, &[75]); + + let x0_values: Vec<_> = batches + .iter() + .flat_map(|x| { + x.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .cloned() + }) + .collect(); + + let x1_values: Vec<_> = batches + .iter() + .flat_map(|x| { + x.column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .cloned() + }) + .collect(); + + let expected_x0_values: Vec<_> = [25..75, 100..125].into_iter().flatten().collect(); + assert_eq!(&x0_values, &expected_x0_values); + + let expected_x1_values: Vec<_> = [125..175, 200..225].into_iter().flatten().collect(); + assert_eq!(&x1_values, &expected_x1_values); + Ok(()) +} + +#[test] +fn test_write_non_uniform_encryption() { + fn write_non_uniform_encryption( + footer_key: &[u8], + column_names: Vec<&str>, + column_keys: Vec>, + encryption_column_keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let file = File::open(path).unwrap(); + + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .with_column_keys(column_names.to_vec(), column_keys.clone()) + .unwrap() + .build() + .unwrap(); + + let mut builder = FileEncryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in encryption_column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let file_encryption_properties = builder.build().unwrap(); + + read_and_roundtrip_to_encrypted_file( + &file, + decryption_properties, + file_encryption_properties, + ); + } + + write_non_uniform_encryption( + AES_128_FOOTER_KEY, + AES_128_COLUMN_NAMES.to_vec(), + AES_128_COLUMN_KEYS.iter().map(|&s| s.to_vec()).collect(), + AES_128_COLUMN_NAME_KEYS, + ); + + write_non_uniform_encryption( + AES_256_FOOTER_KEY, + AES_256_COLUMN_NAMES.to_vec(), + AES_256_COLUMN_KEYS.iter().map(|&s| s.to_vec()).collect(), + AES_256_COLUMN_NAME_KEYS, + ); +} + +#[test] +fn test_write_uniform_encryption_plaintext_footer() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/encrypt_columns_plaintext_footer.parquet.encrypted"); + let file = File::open(path).unwrap(); + + let footer_key = AES_128_FOOTER_KEY.to_vec(); + let wrong_footer_key = BAD_AES_128_FOOTER_KEY.to_vec(); + let column_1_key = AES_128_COLUMN_KEYS[0].to_vec(); + let column_2_key = AES_128_COLUMN_KEYS[1].to_vec(); + + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .with_column_key(AES_128_COLUMN_NAMES[0], column_1_key.clone()) + .with_column_key(AES_128_COLUMN_NAMES[1], column_2_key.clone()) + .build() + .unwrap(); + + let wrong_decryption_properties = FileDecryptionProperties::builder(wrong_footer_key) + .with_column_key(AES_128_COLUMN_NAMES[0], column_1_key) + .with_column_key(AES_128_COLUMN_NAMES[1], column_2_key) + .build() + .unwrap(); + + let file_encryption_properties = FileEncryptionProperties::builder(footer_key) + .with_plaintext_footer(true) + .build() + .unwrap(); + + // Try writing plaintext footer and then reading it with the correct footer key + read_and_roundtrip_to_encrypted_file( + &file, + Arc::clone(&decryption_properties), + file_encryption_properties.clone(), + ); + + // Try writing plaintext footer and then reading it with the wrong footer key + let temp_file = tempfile::tempfile().unwrap(); + + // read example data + let options = ArrowReaderOptions::default() + .with_file_decryption_properties(decryption_properties.clone()); + let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); + + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap(); + let batch_reader = builder.build().unwrap(); + let batches = batch_reader + .collect::, _>>() + .unwrap(); + + // write example data + let props = WriterProperties::builder() + .with_file_encryption_properties(file_encryption_properties) + .build(); + + let mut writer = ArrowWriter::try_new( + temp_file.try_clone().unwrap(), + metadata.schema().clone(), + Some(props), + ) + .unwrap(); + for batch in batches { + writer.write(&batch).unwrap(); + } + writer.close().unwrap(); + + // Try reading plaintext footer and with the wrong footer key + let options = + ArrowReaderOptions::default().with_file_decryption_properties(wrong_decryption_properties); + let result = ArrowReaderMetadata::load(&temp_file, options.clone()); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .starts_with("Parquet error: Footer signature verification failed. Computed: [") + ); +} + +#[test] +pub fn test_column_statistics_with_plaintext_footer() { + let footer_key = AES_128_FOOTER_KEY.to_vec(); + let column_key = AES_128_COLUMN_KEYS[0].to_vec(); + + // Encrypt with a plaintext footer and column-specific keys + let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) + .with_plaintext_footer(true) + .with_column_key("x", column_key.clone()) + .with_column_key("y", column_key.clone()) + .with_column_key("s", column_key.clone()) + .build() + .unwrap(); + + // Read with only the footer key and the key for one column + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .with_column_key("x", column_key.clone()) + .build() + .unwrap(); + + // Reader can read plaintext stats from the unencrypted column z + // and column x for which the key is provided, but not columns y and s + // for which no key is provided. + write_and_read_stats( + Arc::clone(&encryption_properties), + Some(decryption_properties), + &[true, false, true, false], + ); + + // Read without any decryption properties. + // Reader can only read plaintext stats from the unencrypted column z. + write_and_read_stats(encryption_properties, None, &[false, false, true, false]); +} + +#[test] +pub fn test_column_statistics_with_plaintext_footer_and_uniform_encryption() { + let footer_key = AES_128_FOOTER_KEY.to_vec(); + + // Write with uniform encryption and a plaintext footer. + let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) + .with_plaintext_footer(true) + .build() + .unwrap(); + + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .build() + .unwrap(); + + // Reader can read stats from plaintext footer metadata if a footer key is provided + write_and_read_stats( + Arc::clone(&encryption_properties), + Some(decryption_properties), + &[true, true, true, true], + ); + + // Reader can not read stats from plaintext footer metadata if no key is provided + write_and_read_stats(encryption_properties, None, &[false, false, false, false]); +} + +#[test] +pub fn test_column_statistics_with_encrypted_footer() { + let footer_key = AES_128_FOOTER_KEY.to_vec(); + let column_key = AES_128_COLUMN_KEYS[0].to_vec(); + + // Encrypt with an encrypted footer and column-specific keys + let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) + .with_plaintext_footer(false) + .with_column_key("x", column_key.clone()) + .with_column_key("y", column_key.clone()) + .with_column_key("s", column_key.clone()) + .build() + .unwrap(); + + // Read with only the footer key and the key for one column + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .with_column_key("x", column_key.clone()) + .build() + .unwrap(); + + // Reader can read plaintext stats from the unencrypted column z + // and column x for which the key is provided, but not columns y and s + // for which no key is provided. + write_and_read_stats( + encryption_properties, + Some(decryption_properties), + &[true, false, true, false], + ); +} + +#[test] +pub fn test_column_statistics_with_encrypted_footer_and_uniform_encryption() { + let footer_key = AES_128_FOOTER_KEY.to_vec(); + + // Encrypt with an encrypted footer and uniform encryption + let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) + .with_plaintext_footer(false) + .build() + .unwrap(); + + // Read with the footer key + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .build() + .unwrap(); + + // Reader can read stats for all columns. + write_and_read_stats( + encryption_properties, + Some(decryption_properties), + &[true, true, true, true], + ); +} + +/// Write a file with encryption and then verify whether statistics are readable with the provided decryption properties. +fn write_and_read_stats( + encryption_properties: Arc, + decryption_properties: Option>, + expect_stats: &[bool], +) { + use parquet::basic::Type; + + let int_values = Int32Array::from(vec![8, 3, 4, 19, 5]); + let int_values = Arc::new(int_values); + let string_values: StringArray = vec![ + None, + Some("parquet"), + Some("encryption"), + Some("test"), + None, + ] + .into(); + let string_values = Arc::new(string_values); + + let schema = Arc::new(Schema::new(vec![ + Field::new("x", int_values.data_type().clone(), true), + Field::new("y", int_values.data_type().clone(), true), + Field::new("z", int_values.data_type().clone(), true), + Field::new("s", string_values.data_type().clone(), true), + ])); + let record_batches = vec![ + RecordBatch::try_new( + schema.clone(), + vec![ + int_values.clone(), + int_values.clone(), + int_values.clone(), + string_values.clone(), + ], + ) + .unwrap(), + ]; + + let props = WriterProperties::builder() + .with_file_encryption_properties(encryption_properties) + .set_bloom_filter_enabled(true) + .build(); + + let temp_file = tempfile::tempfile().unwrap(); + let mut writer = ArrowWriter::try_new(&temp_file, schema.clone(), Some(props)).unwrap(); + for batch in record_batches.clone() { + writer.write(&batch).unwrap(); + } + let metadata = writer.close().unwrap(); + + let expected_min = 3i32.to_le_bytes(); + let expected_max = 19i32.to_le_bytes(); + + let check_column_stats = |column: &ColumnChunkMetaData, expect_stats: bool| { + let is_byte_array = column.column_type() == Type::BYTE_ARRAY; + if expect_stats { + assert!(column.page_encoding_stats().is_some()); + assert!(column.statistics().is_some()); + if is_byte_array { + // Size statistics for BYTE_ARRAY columns + assert!(column.unencoded_byte_array_data_bytes().is_some()); + } else { + let column_stats = column.statistics().unwrap(); + assert_eq!(column_stats.min_bytes_opt(), Some(expected_min.as_slice())); + assert_eq!(column_stats.max_bytes_opt(), Some(expected_max.as_slice())); + } + assert!(column.bloom_filter_offset().is_some()); + assert!(column.bloom_filter_length().is_some()); + } else { + assert!(column.statistics().is_none()); + assert!(column.page_encoding_stats().is_none()); + assert!(column.bloom_filter_offset().is_none()); + assert!(column.bloom_filter_length().is_none()); + // Size statistics should also be stripped + if is_byte_array { + assert!(column.unencoded_byte_array_data_bytes().is_none()); + } + } + }; + + // Check column statistics produced at write time are available in full + let row_group = metadata.row_group(0); + for column in row_group.columns().iter() { + check_column_stats(column, true); + } + + // Verify the presence or not of statistics per-column when reading with the provided decryption properties + let mut options = ArrowReaderOptions::default().with_encoding_stats_as_mask(false); + if let Some(decryption_properties) = decryption_properties { + options = options.with_file_decryption_properties(decryption_properties); + } + let reader_metadata = ArrowReaderMetadata::load(&temp_file, options).unwrap(); + let metadata = reader_metadata.metadata(); + let row_group = metadata.row_group(0); + + for (column, stats_expected) in row_group.columns().iter().zip(expect_stats) { + check_column_stats(column, *stats_expected); + } +} + +#[test] +fn test_write_uniform_encryption() { + fn write_uniform_encryption(footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let file = File::open(path).unwrap(); + + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); + + let file_encryption_properties = FileEncryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); + + read_and_roundtrip_to_encrypted_file( + &file, + decryption_properties, + file_encryption_properties, + ); + } + + write_uniform_encryption(AES_128_FOOTER_KEY); + write_uniform_encryption(AES_256_FOOTER_KEY); +} + +#[test] +fn test_write_non_uniform_encryption_column_missmatch() { + fn write_non_uniform_encryption_column_missmatch( + footer_key: &[u8], + column_keys: &[(&str, &[u8])], + encryption_column_keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + + let mut decryption_builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + decryption_builder = decryption_builder.with_column_key(column_name, key.to_vec()); + } + + let decryption_properties = decryption_builder.build().unwrap(); + + let mut file_encryption_builder = FileEncryptionProperties::builder(footer_key.to_vec()) + .with_column_key("other_field", encryption_column_keys[0].1.to_vec()) + .with_column_key("yet_another_field", encryption_column_keys[1].1.to_vec()); + + for (column_name, key) in encryption_column_keys { + file_encryption_builder = + file_encryption_builder.with_column_key(column_name, key.to_vec()); + } + + let file_encryption_properties = file_encryption_builder.build().unwrap(); + + let temp_file = tempfile::tempfile().unwrap(); + + // read example data + let file = File::open(path).unwrap(); + let options = ArrowReaderOptions::default() + .with_file_decryption_properties(decryption_properties.clone()); + let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); + let props = WriterProperties::builder() + .with_file_encryption_properties(file_encryption_properties) + .build(); + + let result = ArrowWriter::try_new( + temp_file.try_clone().unwrap(), + metadata.schema().clone(), + Some(props), + ); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: The following columns with encryption keys specified were not found in the schema: other_field, yet_another_field" + ); + } + + write_non_uniform_encryption_column_missmatch( + AES_128_FOOTER_KEY, + AES_128_COLUMN_NAME_KEYS, + AES_128_COLUMN_NAME_KEYS, + ); + + write_non_uniform_encryption_column_missmatch( + AES_256_FOOTER_KEY, + AES_256_COLUMN_NAME_KEYS, + AES_256_COLUMN_NAME_KEYS, + ); +} + +#[test] +fn test_write_encrypted_column() { + let message_type = " + message test_schema { + OPTIONAL BYTE_ARRAY a (UTF8); + } + "; + let schema = Arc::new(parse_message_type(message_type).unwrap()); + let data = vec![ByteArray::from(b"parquet".to_vec()); 7]; + let def_levels = [1, 1, 1, 1, 0, 1, 0, 1, 0, 1]; + + let num_row_groups = 3; + let num_batches = 3; + let rows_per_batch = def_levels.len(); + let valid_rows_per_batch = def_levels.iter().filter(|&level| *level > 0).count(); + + let file: File = tempfile::tempfile().unwrap(); + + let builder = WriterProperties::builder(); + let footer_key: &[u8] = AES_128_FOOTER_KEY; + let file_encryption_properties = FileEncryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); + + let props = Arc::new( + builder + .with_file_encryption_properties(file_encryption_properties) + .set_data_page_row_count_limit(rows_per_batch) + .build(), + ); + let mut writer = SerializedFileWriter::new(&file, schema, props).unwrap(); + for _ in 0..num_row_groups { + let mut row_group_writer = writer.next_row_group().unwrap(); + let mut col_writer = row_group_writer.next_column().unwrap().unwrap(); + + for _ in 0..num_batches { + col_writer + .typed::() + .write_batch(&data, Some(&def_levels), None) + .unwrap(); + } + + col_writer.close().unwrap(); + row_group_writer.close().unwrap(); + } + + let _file_metadata = writer.close().unwrap(); + + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); + let options = ArrowReaderOptions::default() + .with_file_decryption_properties(decryption_properties.clone()); + let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); + let file_metadata = metadata.metadata().file_metadata(); + + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap(); + let record_reader = builder.build().unwrap(); + + assert_eq!( + file_metadata.num_rows(), + (num_row_groups * num_batches * rows_per_batch) as i64 + ); + assert_eq!(file_metadata.schema_descr().num_columns(), 1); + + assert_eq!(metadata.metadata().num_row_groups(), num_row_groups); + metadata.metadata().row_groups().iter().for_each(|rg| { + assert_eq!(rg.num_columns(), 1); + assert_eq!(rg.num_rows(), (num_batches * rows_per_batch) as i64); + }); + + let mut row_count = 0; + for batch in record_reader { + let batch = batch.unwrap(); + row_count += batch.num_rows(); + + let string_col = batch.column(0).as_string_opt::().unwrap(); + + let mut valid_count = 0; + for x in string_col.iter().flatten() { + valid_count += 1; + assert_eq!(x, "parquet"); + } + assert_eq!( + valid_count, + valid_rows_per_batch * num_batches * num_row_groups + ); + } + + assert_eq!(row_count, file_metadata.num_rows() as usize); +} + +#[test] +fn test_write_encrypted_struct_field() { + let int_32: Int32Array = [Some(1), Some(6)].iter().collect(); + let float_64: Float64Array = [None, Some(8.5)].iter().collect(); + let struct_array = StructArray::from(vec![ + ( + Arc::new(Field::new("int64_col", DataType::Int32, true)), + Arc::new(int_32) as ArrayRef, + ), + ( + Arc::new(Field::new("float64_col", DataType::Float64, true)), + Arc::new(float_64) as ArrayRef, + ), + ]); + let struct_array_data = Arc::new(struct_array); + + let schema = Arc::new(Schema::new(vec![Field::new( + "struct_col", + struct_array_data.data_type().clone(), + true, + )])); + let record_batches = + vec![RecordBatch::try_new(schema.clone(), vec![struct_array_data]).unwrap()]; + + let temp_file = tempfile::tempfile().unwrap(); + + // When configuring encryption keys for struct columns, + // keys need to be specified for each leaf-level Parquet column using the full "." separated + // column path. + let builder = WriterProperties::builder(); + let footer_key = AES_128_FOOTER_KEY.to_vec(); + let column_key_1 = AES_128_COLUMN_KEYS[0].to_vec(); + let column_key_2 = AES_128_COLUMN_KEYS[1].to_vec(); + let file_encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) + .with_column_key("struct_col.int64_col", column_key_1.clone()) + .with_column_key("struct_col.float64_col", column_key_2.clone()) + .build() + .unwrap(); + + let props = builder + .with_file_encryption_properties(file_encryption_properties) + .build(); + let mut writer = + ArrowWriter::try_new(temp_file.try_clone().unwrap(), schema, Some(props)).unwrap(); + for batch in record_batches.clone() { + writer.write(&batch).unwrap(); + } + writer.close().unwrap(); + + let decryption_properties = FileDecryptionProperties::builder(footer_key) + .with_column_key("struct_col.int64_col", column_key_1) + .with_column_key("struct_col.float64_col", column_key_2) + .build() + .unwrap(); + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + + let builder = + ParquetRecordBatchReaderBuilder::try_new_with_options(temp_file, options).unwrap(); + let record_reader = builder.build().unwrap(); + + let read_record_reader = record_reader + .map(|x| x.unwrap()) + .collect::>(); + + // show read batches are equal to written batches + assert_eq!(read_record_reader.len(), record_batches.len()); + for (read_batch, written_batch) in read_record_reader.iter().zip(record_batches.iter()) { + assert_eq!(read_batch.num_columns(), written_batch.num_columns()); + assert_eq!(read_batch.num_rows(), written_batch.num_rows()); + for (read_column, written_column) in read_batch + .columns() + .iter() + .zip(written_batch.columns().iter()) + { + assert_eq!(read_column, written_column); + } + } +} + +/// Test that when per-column encryption is used, +/// unencrypted row group metadata are returned when the writer is closed +/// and statistics can be used. +#[test] +pub fn test_retrieve_row_group_statistics_after_encrypted_write() { + let values = Int32Array::from(vec![8, 3, 4, 19, 5]); + + let schema = Arc::new(Schema::new(vec![Field::new( + "x", + values.data_type().clone(), + true, + )])); + let values = Arc::new(values); + let record_batches = vec![RecordBatch::try_new(schema.clone(), vec![values]).unwrap()]; + + let temp_file = tempfile::tempfile().unwrap(); + + let footer_key = AES_128_FOOTER_KEY.to_vec(); + let column_key = AES_128_COLUMN_KEYS[0].to_vec(); + let file_encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) + .with_column_key("x", column_key.clone()) + .build() + .unwrap(); + + let props = WriterProperties::builder() + .with_file_encryption_properties(file_encryption_properties) + .build(); + let mut writer = ArrowWriter::try_new(temp_file, schema, Some(props)).unwrap(); + + for batch in record_batches.clone() { + writer.write(&batch).unwrap(); + } + let file_metadata = writer.close().unwrap(); + + assert_eq!(file_metadata.num_row_groups(), 1); + let row_group = file_metadata.row_group(0); + assert_eq!(row_group.num_columns(), 1); + let column = row_group.column(0); + let column_stats = column.statistics().unwrap(); + assert_eq!( + column_stats.min_bytes_opt(), + Some(3i32.to_le_bytes().as_slice()) + ); + assert_eq!( + column_stats.max_bytes_opt(), + Some(19i32.to_le_bytes().as_slice()) + ); +} + +#[test] +fn test_decrypt_page_index_uniform() { + fn decrypt_page_index_uniform(footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); + + test_decrypt_page_index(&path, decryption_properties).unwrap(); + } + + decrypt_page_index_uniform(AES_128_FOOTER_KEY); + decrypt_page_index_uniform(AES_256_FOOTER_KEY); +} + +#[test] +fn test_decrypt_page_index_non_uniform() { + fn decrypt_page_index_non_uniform(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + test_decrypt_page_index(&path, decryption_properties).unwrap(); + } + + decrypt_page_index_non_uniform(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); + + decrypt_page_index_non_uniform(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS); +} + +fn test_decrypt_page_index( + path: &str, + decryption_properties: Arc, +) -> Result<(), ParquetError> { + let file = File::open(path)?; + let options = ArrowReaderOptions::default() + .with_file_decryption_properties(decryption_properties) + .with_page_index_policy(PageIndexPolicy::from(true)); + + let arrow_metadata = ArrowReaderMetadata::load(&file, options)?; + + verify_column_indexes(arrow_metadata.metadata()); + + Ok(()) +} + +#[test] +fn test_decryption_properties_uses_key_retriever() { + let key_retriever = TestKeyRetriever::new() + .with_key( + AES_128_FOOTER_KEY_NAME.to_owned(), + AES_128_FOOTER_KEY.to_vec(), + ) + .with_key( + AES_128_KEY_NAMES[0].to_owned(), + AES_128_COLUMN_KEYS[0].to_vec(), + ); + + let properties_with_retriever = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + assert!(properties_with_retriever.uses_key_retriever()); + + let properties_with_keys = FileDecryptionProperties::builder(AES_128_FOOTER_KEY.to_vec()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].to_vec()) + .build() + .unwrap(); + + assert!(!properties_with_keys.uses_key_retriever()); + + let uniform_properties = FileDecryptionProperties::builder(AES_128_FOOTER_KEY.to_vec()) + .build() + .unwrap(); + + assert!(!uniform_properties.uses_key_retriever()); +} diff --git a/vendor/parquet/tests/encryption/encryption_agnostic.rs b/vendor/parquet/tests/encryption/encryption_agnostic.rs new file mode 100644 index 0000000..604155c --- /dev/null +++ b/vendor/parquet/tests/encryption/encryption_agnostic.rs @@ -0,0 +1,146 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Common tests for encryption related functionality + +use arrow_array::cast::AsArray; +use arrow_array::types; +use arrow_schema::ArrowError; +use parquet::arrow::ProjectionMask; +use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ParquetRecordBatchReaderBuilder}; +use std::fs::File; + +pub fn read_plaintext_footer_file_without_decryption_properties() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/encrypt_columns_plaintext_footer.parquet.encrypted"); + let file = File::open(&path).unwrap(); + + let metadata = ArrowReaderMetadata::load(&file, Default::default()).unwrap(); + let file_metadata = metadata.metadata().file_metadata(); + + assert_eq!(file_metadata.num_rows(), 50); + assert_eq!(file_metadata.schema_descr().num_columns(), 8); + assert_eq!( + file_metadata.created_by().unwrap(), + "parquet-cpp-arrow version 19.0.0-SNAPSHOT" + ); + + metadata.metadata().row_groups().iter().for_each(|rg| { + assert_eq!(rg.num_columns(), 8); + assert_eq!(rg.num_rows(), 50); + }); + + // Should be able to read unencrypted columns. Test reading one column. + let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + let mask = ProjectionMask::leaves(builder.parquet_schema(), [1]); + let record_reader = builder.with_projection(mask).build().unwrap(); + + let mut row_count = 0; + for batch in record_reader { + let batch = batch.unwrap(); + row_count += batch.num_rows(); + + let time_col = batch + .column(0) + .as_primitive::(); + for (i, x) in time_col.iter().enumerate() { + assert_eq!(x.unwrap(), i as i32); + } + } + + assert_eq!(row_count, file_metadata.num_rows() as usize); + + // Reading an encrypted column should fail + let file = File::open(&path).unwrap(); + let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + let mask = ProjectionMask::leaves(builder.parquet_schema(), [4]); + let mut record_reader = builder.with_projection(mask).build().unwrap(); + + match record_reader.next() { + Some(Err(ArrowError::ParquetError(s))) => { + assert!(s.contains("Parquet error")); + } + _ => { + panic!("Expected ArrowError::ParquetError"); + } + }; +} + +#[cfg(feature = "async")] +pub async fn read_plaintext_footer_file_without_decryption_properties_async() { + use futures::StreamExt; + use futures::TryStreamExt; + use parquet::arrow::ParquetRecordBatchStreamBuilder; + use parquet::errors::ParquetError; + + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/encrypt_columns_plaintext_footer.parquet.encrypted"); + let mut file = tokio::fs::File::open(&path).await.unwrap(); + + let metadata = ArrowReaderMetadata::load_async(&mut file, Default::default()) + .await + .unwrap(); + let file_metadata = metadata.metadata().file_metadata(); + + assert_eq!(file_metadata.num_rows(), 50); + assert_eq!(file_metadata.schema_descr().num_columns(), 8); + assert_eq!( + file_metadata.created_by().unwrap(), + "parquet-cpp-arrow version 19.0.0-SNAPSHOT" + ); + + metadata.metadata().row_groups().iter().for_each(|rg| { + assert_eq!(rg.num_columns(), 8); + assert_eq!(rg.num_rows(), 50); + }); + + // Should be able to read unencrypted columns. Test reading one column. + let builder = ParquetRecordBatchStreamBuilder::new(file).await.unwrap(); + let mask = ProjectionMask::leaves(builder.parquet_schema(), [1]); + let record_reader = builder.with_projection(mask).build().unwrap(); + let record_batches = record_reader.try_collect::>().await.unwrap(); + + let mut row_count = 0; + for batch in record_batches { + let batch = batch; + row_count += batch.num_rows(); + + let time_col = batch + .column(0) + .as_primitive::(); + for (i, x) in time_col.iter().enumerate() { + assert_eq!(x.unwrap(), i as i32); + } + } + + assert_eq!(row_count, file_metadata.num_rows() as usize); + + // Reading an encrypted column should fail + let file = tokio::fs::File::open(&path).await.unwrap(); + let builder = ParquetRecordBatchStreamBuilder::new(file).await.unwrap(); + let mask = ProjectionMask::leaves(builder.parquet_schema(), [4]); + let mut record_reader = builder.with_projection(mask).build().unwrap(); + + match record_reader.next().await { + Some(Err(ParquetError::ArrowError(s))) => { + assert!(s.contains("Parquet error")); + } + err => { + panic!("Expected ArrowError::ParquetError, got {err:?}"); + } + }; +} diff --git a/vendor/parquet/tests/encryption/encryption_async.rs b/vendor/parquet/tests/encryption/encryption_async.rs new file mode 100644 index 0000000..a2ec2ed --- /dev/null +++ b/vendor/parquet/tests/encryption/encryption_async.rs @@ -0,0 +1,1169 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! This module contains tests for reading encrypted Parquet files with the async Arrow API + +use crate::encryption_util; +use crate::encryption_util::{ + AES_128_COLUMN_KEYS, AES_128_COLUMN_NAME_KEYS, AES_128_COLUMN_NAMES, AES_128_FOOTER_KEY, + AES_128_FOOTER_KEY_NAME, AES_128_KEY_NAME_KEY, AES_256_COLUMN_KEYS, AES_256_COLUMN_NAME_KEYS, + AES_256_COLUMN_NAMES, AES_256_FOOTER_KEY, AES_256_FOOTER_KEY_NAME, AES_256_KEY_NAME_KEY, + BAD_AES_128_FOOTER_KEY, BAD_AES_256_FOOTER_KEY, TestKeyRetriever, read_encrypted_file, + verify_column_indexes, verify_encryption_double_test_data, verify_encryption_test_data, +}; +use arrow_array::RecordBatch; +use arrow_schema::Schema; +use futures::TryStreamExt; +use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; +use parquet::arrow::arrow_writer::{ + ArrowColumnChunk, ArrowColumnWriter, ArrowLeafColumn, ArrowRowGroupWriterFactory, + ArrowWriterOptions, compute_leaves, +}; +use parquet::arrow::{ + ArrowSchemaConverter, ArrowWriter, AsyncArrowWriter, ParquetRecordBatchStreamBuilder, +}; +use parquet::encryption::decrypt::FileDecryptionProperties; +use parquet::encryption::encrypt::FileEncryptionProperties; +use parquet::errors::ParquetError; +use parquet::file::metadata::PageIndexPolicy; +use parquet::file::metadata::ParquetMetaData; +use parquet::file::properties::{WriterProperties, WriterPropertiesBuilder}; +use parquet::file::writer::SerializedFileWriter; +use std::io::Write; +use std::sync::Arc; +use tokio::fs::File; +use tokio::sync::mpsc::{Receiver, Sender}; +use tokio::task::JoinHandle; + +#[tokio::test] +async fn test_non_uniform_encryption_plaintext_footer() { + async fn non_uniform_encryption_plaintext_footer( + footer_key: &[u8], + column_keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_plaintext_footer.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); + } + + // AES-128: there is always a footer key even with a plaintext footer, + // but this is used for signing the footer. + non_uniform_encryption_plaintext_footer(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS).await; + + // AES-256 + non_uniform_encryption_plaintext_footer(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS).await; +} + +#[tokio::test] +async fn test_misspecified_encryption_keys() { + // read file with keys and check for expected error message + async fn check_for_error( + expected_message: &str, + footer_key: &[u8], + column_1_key: &[u8], + column_2_key: &[u8], + additional_column_keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); + + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + + if !column_1_key.is_empty() { + builder = builder.with_column_key(AES_128_COLUMN_NAMES[0], column_1_key.to_vec()); + } + + if !column_2_key.is_empty() { + builder = builder.with_column_key(AES_128_COLUMN_NAMES[1], column_2_key.to_vec()); + } + + for (column_name, key) in additional_column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + + let decryption_properties = builder.build().unwrap(); + + match verify_encryption_test_file_read_async(&mut file, decryption_properties).await { + Ok(_) => { + panic!("did not get expected error") + } + Err(e) => { + assert_eq!(e.to_string(), expected_message); + } + } + } + + // There is always a footer key even with a plaintext footer, + // but this is used for signing the footer. + let footer_key = AES_128_FOOTER_KEY; + let column_1_key = AES_128_COLUMN_KEYS[0]; + let column_2_key = AES_128_COLUMN_KEYS[1]; + let empty_column_key = &[]; + + // Too short footer key + check_for_error( + format!("Parquet error: Invalid footer key. Error creating RingGcmBlockDecryptor with unsupported key length: {}", "bad_pwd".len()).as_str(), + b"bad_pwd", + column_1_key, + column_2_key, + empty_column_key + ) + .await; + + // Wrong footer key + check_for_error( + "Parquet error: Provided footer key and AAD were unable to decrypt parquet footer", + BAD_AES_128_FOOTER_KEY, + column_1_key, + column_2_key, + empty_column_key, + ) + .await; + + // Missing column key + check_for_error( + "Parquet error: No column decryption key set for encrypted column 'double_field'", + footer_key, + "".as_bytes(), + column_2_key, + empty_column_key, + ) + .await; + + // Too short column key + check_for_error( + format!( + "Parquet error: Error creating RingGcmBlockDecryptor with unsupported key length: {}", + "abc".len() + ) + .as_str(), + footer_key, + "abc".as_bytes(), + column_2_key, + empty_column_key, + ) + .await; + + // Wrong column key + check_for_error( + "Parquet error: Unable to decrypt column 'double_field', perhaps the column key is wrong?", + footer_key, + "1123456789012345".as_bytes(), + column_2_key, + empty_column_key, + ) + .await; + + // Mixed up keys + check_for_error( + "Parquet error: Unable to decrypt column 'float_field', perhaps the column key is wrong?", + footer_key, + column_2_key, + column_1_key, + empty_column_key, + ) + .await; + + let aes256_footer_key = AES_256_FOOTER_KEY; + let aes256_column_1_key = AES_256_COLUMN_KEYS[0]; + let aes256_column_2_key = AES_256_COLUMN_KEYS[1]; + let additional_column_keys = &[ + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), + ]; + + // Too short footer key + check_for_error( + format!("Parquet error: Invalid footer key. Error creating RingGcmBlockDecryptor with unsupported key length: {}", "bad_pwd".len()).as_str(), + b"bad_pwd", + aes256_column_1_key, + aes256_column_2_key, + additional_column_keys + ).await; + + // Wrong footer key + check_for_error( + "Parquet error: Provided footer key and AAD were unable to decrypt parquet footer", + BAD_AES_256_FOOTER_KEY, + aes256_column_1_key, + aes256_column_2_key, + additional_column_keys, + ) + .await; + + // Missing column key + check_for_error( + "Parquet error: No column decryption key set for encrypted column 'double_field'", + aes256_footer_key, + "".as_bytes(), + aes256_column_2_key, + additional_column_keys, + ) + .await; + + // Too short column key + check_for_error( + format!( + "Parquet error: Error creating RingGcmBlockDecryptor with unsupported key length: {}", + "abc".len() + ) + .as_str(), + aes256_footer_key, + "abc".as_bytes(), + aes256_column_2_key, + additional_column_keys, + ) + .await; + + // Wrong column key + check_for_error( + "Parquet error: Unable to decrypt column 'double_field', perhaps the column key is wrong?", + aes256_footer_key, + "22345678901234567890123456789012".as_bytes(), + aes256_column_2_key, + additional_column_keys, + ) + .await; + + // Mixed up keys + check_for_error( + "Parquet error: Unable to decrypt column 'float_field', perhaps the column key is wrong?", + aes256_footer_key, + aes256_column_2_key, + aes256_column_1_key, + additional_column_keys, + ) + .await; +} + +#[tokio::test] +#[cfg(feature = "snap")] +async fn test_plaintext_footer_read_without_decryption() { + crate::encryption_agnostic::read_plaintext_footer_file_without_decryption_properties_async() + .await; +} + +#[tokio::test] +async fn test_non_uniform_encryption() { + async fn non_uniform_encryption(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); + + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); + } + + non_uniform_encryption(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS).await; + non_uniform_encryption(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS).await; +} + +#[tokio::test] +async fn test_uniform_encryption() { + async fn uniform_encryption(footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); + + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); + + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); + } + + uniform_encryption(AES_128_FOOTER_KEY).await; + uniform_encryption(AES_256_FOOTER_KEY).await; +} + +#[tokio::test] +async fn test_aes_ctr_encryption() { + async fn aes_ctr_encryption(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer_ctr.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); + + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + + let options = + ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); + let metadata = ArrowReaderMetadata::load_async(&mut file, options).await; + + match metadata { + Err(ParquetError::NYI(s)) => { + assert!(s.contains("AES_GCM_CTR_V1")); + } + _ => { + panic!("Expected ParquetError::NYI"); + } + }; + } + + aes_ctr_encryption(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS).await; + aes_ctr_encryption(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS).await; +} + +#[tokio::test] +async fn test_decrypting_without_decryption_properties_fails() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/uniform_encryption.parquet.encrypted"); + let mut file = File::open(&path).await.unwrap(); + + let options = ArrowReaderOptions::new(); + let result = ArrowReaderMetadata::load_async(&mut file, options).await; + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" + ); +} + +#[tokio::test] +async fn test_write_non_uniform_encryption() { + async fn write_non_uniform_encryption( + footer_key: &[u8], + column_names: Vec<&str>, + column_keys: Vec>, + encryption_column_keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .with_column_keys(column_names.to_vec(), column_keys.clone()) + .unwrap() + .build() + .unwrap(); + + let mut builder = FileEncryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in encryption_column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let file_encryption_properties = builder.build().unwrap(); + + read_and_roundtrip_to_encrypted_file_async( + &path, + decryption_properties, + file_encryption_properties, + ) + .await + .unwrap(); + } + + write_non_uniform_encryption( + AES_128_FOOTER_KEY, + AES_128_COLUMN_NAMES.to_vec(), + AES_128_COLUMN_KEYS.iter().map(|&s| s.to_vec()).collect(), + AES_128_COLUMN_NAME_KEYS, + ) + .await; + + // AES-256 + write_non_uniform_encryption( + AES_256_FOOTER_KEY, + AES_256_COLUMN_NAMES.to_vec(), + AES_256_COLUMN_KEYS.iter().map(|&s| s.to_vec()).collect(), + AES_256_COLUMN_NAME_KEYS, + ) + .await; +} + +/// An [`AsyncFileReader`] reading via an [`ObjectStore`], mirroring the +/// example on the [`AsyncFileReader`] trait documentation +/// +/// [`AsyncFileReader`]: parquet::arrow::async_reader::AsyncFileReader +/// [`ObjectStore`]: object_store::ObjectStore +mod object_store_reader { + use bytes::Bytes; + use futures::future::BoxFuture; + use futures::{FutureExt, TryFutureExt}; + use object_store::path::Path; + use object_store::{GetOptions, GetRange, ObjectStore, ObjectStoreExt}; + use parquet::arrow::arrow_reader::ArrowReaderOptions; + use parquet::arrow::async_reader::{AsyncFileReader, MetadataSuffixFetch}; + use parquet::errors::{ParquetError, Result}; + use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader}; + use std::ops::Range; + use std::sync::Arc; + + fn to_parquet_err(e: object_store::Error) -> ParquetError { + ParquetError::External(Box::new(e)) + } + + #[derive(Clone)] + pub struct ObjectStoreReader { + pub store: Arc, + pub path: Path, + } + + impl AsyncFileReader for ObjectStoreReader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + self.store + .get_range(&self.path, range) + .map_err(to_parquet_err) + .boxed() + } + + fn get_byte_ranges( + &mut self, + ranges: Vec>, + ) -> BoxFuture<'_, Result>> { + async move { + self.store + .get_ranges(&self.path, &ranges) + .await + .map_err(to_parquet_err) + } + .boxed() + } + + fn get_metadata<'a>( + &'a mut self, + options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, Result>> { + async move { + let metadata = ParquetMetaDataReader::new() + .with_arrow_reader_options(options) + .load_via_suffix_and_finish(self) + .await?; + Ok(Arc::new(metadata)) + } + .boxed() + } + } + + impl MetadataSuffixFetch for &mut ObjectStoreReader { + fn fetch_suffix(&mut self, suffix: usize) -> BoxFuture<'_, Result> { + let options = GetOptions { + range: Some(GetRange::Suffix(suffix as u64)), + ..Default::default() + }; + async move { + let resp = self + .store + .get_opts(&self.path, options) + .await + .map_err(to_parquet_err)?; + resp.bytes().await.map_err(to_parquet_err) + } + .boxed() + } + } +} + +#[tokio::test] +async fn test_read_encrypted_file_from_object_store() { + use object_store::local::LocalFileSystem; + use object_store::path::Path; + use object_store_reader::ObjectStoreReader; + use parquet::arrow::async_reader::AsyncFileReader; + use std::sync::Arc; + + let test_data = arrow::util::test_util::parquet_test_data(); + let store = Arc::new(LocalFileSystem::new_with_prefix(test_data).unwrap()); + let path = Path::from("uniform_encryption.parquet.encrypted"); + + let key_code: &[u8] = "0123456789012345".as_bytes(); + let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec()) + .build() + .unwrap(); + let options = ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); + + let mut reader = ObjectStoreReader { store, path }; + let metadata = reader.get_metadata(Some(&options)).await.unwrap(); + let builder = ParquetRecordBatchStreamBuilder::new_with_options(reader, options) + .await + .unwrap(); + let batch_stream = builder.build().unwrap(); + let record_batches: Vec<_> = batch_stream.try_collect().await.unwrap(); + + verify_encryption_test_data(record_batches, &metadata); +} + +#[tokio::test] +async fn test_non_uniform_encryption_plaintext_footer_with_key_retriever() { + async fn non_uniform_encryption_plaintext_footer_with_key_retriever( + footer_key: &[u8], + keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_plaintext_footer.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); + + let mut key_retriever = TestKeyRetriever::new(); + for (key_name, key) in keys { + key_retriever = key_retriever.with_key((*key_name).to_owned(), (*key).to_vec()); + } + + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); + } + + non_uniform_encryption_plaintext_footer_with_key_retriever( + AES_128_FOOTER_KEY, + AES_128_KEY_NAME_KEY, + ) + .await; + + non_uniform_encryption_plaintext_footer_with_key_retriever( + AES_256_FOOTER_KEY, + AES_256_KEY_NAME_KEY, + ) + .await; +} + +#[tokio::test] +async fn test_non_uniform_encryption_with_key_retriever() { + async fn non_uniform_encryption_with_key_retriever(footer_key: &[u8], keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); + + let mut key_retriever = TestKeyRetriever::new(); + for (key_name, key) in keys { + key_retriever = key_retriever.with_key((*key_name).to_owned(), (*key).to_vec()); + } + + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); + } + + non_uniform_encryption_with_key_retriever(AES_128_FOOTER_KEY, AES_128_KEY_NAME_KEY).await; + non_uniform_encryption_with_key_retriever(AES_256_FOOTER_KEY, AES_256_KEY_NAME_KEY).await; +} + +#[tokio::test] +async fn test_uniform_encryption_with_key_retriever() { + async fn uniform_encryption_with_key_retriever(key_name: &str, footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); + + let key_retriever = + TestKeyRetriever::new().with_key(key_name.to_owned(), footer_key.to_vec()); + + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); + } + + uniform_encryption_with_key_retriever(AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY).await; + uniform_encryption_with_key_retriever(AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY).await; +} + +#[tokio::test] +async fn test_decrypt_page_index_uniform() { + async fn decrypt_page_index_uniform(footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); + + test_decrypt_page_index(&path, decryption_properties) + .await + .unwrap(); + } + + decrypt_page_index_uniform(AES_128_FOOTER_KEY).await; + decrypt_page_index_uniform(AES_256_FOOTER_KEY).await; +} + +#[tokio::test] +async fn test_decrypt_page_index_non_uniform() { + async fn decrypt_page_index_non_uniform(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + test_decrypt_page_index(&path, decryption_properties) + .await + .unwrap(); + } + + decrypt_page_index_non_uniform(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS).await; + + decrypt_page_index_non_uniform(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS).await; +} + +async fn test_decrypt_page_index( + path: &str, + decryption_properties: Arc, +) -> Result<(), ParquetError> { + let mut file = File::open(&path).await?; + + let options = ArrowReaderOptions::new() + .with_file_decryption_properties(decryption_properties) + .with_page_index_policy(PageIndexPolicy::from(true)); + + let arrow_metadata = ArrowReaderMetadata::load_async(&mut file, options).await?; + + verify_column_indexes(arrow_metadata.metadata()); + + Ok(()) +} + +async fn verify_encryption_test_file_read_async( + file: &mut tokio::fs::File, + decryption_properties: Arc, +) -> Result<(), ParquetError> { + let options = ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); + + let arrow_metadata = ArrowReaderMetadata::load_async(file, options).await?; + let metadata = arrow_metadata.metadata(); + + let record_reader = ParquetRecordBatchStreamBuilder::new_with_metadata( + file.try_clone().await?, + arrow_metadata.clone(), + ) + .build()?; + let record_batches = record_reader.try_collect::>().await?; + + verify_encryption_test_data(record_batches, metadata); + Ok(()) +} + +async fn read_and_roundtrip_to_encrypted_file_async( + path: &str, + decryption_properties: Arc, + encryption_properties: Arc, +) -> Result<(), ParquetError> { + let temp_file = tempfile::tempfile().unwrap(); + let mut file = File::open(&path).await.unwrap(); + + let options = ArrowReaderOptions::new() + .with_file_decryption_properties(Arc::clone(&decryption_properties)); + let arrow_metadata = ArrowReaderMetadata::load_async(&mut file, options).await?; + let record_reader = ParquetRecordBatchStreamBuilder::new_with_metadata( + file.try_clone().await?, + arrow_metadata.clone(), + ) + .build()?; + let record_batches = record_reader.try_collect::>().await?; + + let props = WriterProperties::builder() + .with_file_encryption_properties(encryption_properties) + .build(); + let options = ArrowWriterOptions::new().with_properties(props); + + let file = tokio::fs::File::from_std(temp_file.try_clone().unwrap()); + let mut writer = + AsyncArrowWriter::try_new_with_options(file, arrow_metadata.schema().clone(), options) + .unwrap(); + for batch in record_batches { + writer.write(&batch).await.unwrap(); + } + writer.close().await.unwrap(); + + let mut file = tokio::fs::File::from_std(temp_file.try_clone().unwrap()); + verify_encryption_test_file_read_async(&mut file, decryption_properties).await +} + +// Type aliases for multithreaded file writing tests +type ColSender = Sender; +type ColumnWriterTask = JoinHandle>; +type RBStreamSerializeResult = Result<(Vec, usize), ParquetError>; + +async fn send_arrays_to_column_writers( + col_array_channels: &[ColSender], + rb: &RecordBatch, + schema: &Arc, +) -> Result<(), ParquetError> { + // Each leaf column has its own channel, increment next_channel for each leaf column sent. + let mut next_channel = 0; + for (array, field) in rb.columns().iter().zip(schema.fields()) { + for c in compute_leaves(field, array)? { + if col_array_channels[next_channel].send(c).await.is_err() { + return Ok(()); + } + next_channel += 1; + } + } + Ok(()) +} + +/// Spawns a tokio task which joins the parallel column writer tasks, +/// and finalizes the row group +fn spawn_rg_join_and_finalize_task( + column_writer_tasks: Vec, + rg_rows: usize, +) -> JoinHandle { + tokio::task::spawn(async move { + let num_cols = column_writer_tasks.len(); + let mut finalized_rg = Vec::with_capacity(num_cols); + for task in column_writer_tasks.into_iter() { + let writer = task + .await + .map_err(|e| ParquetError::General(e.to_string()))??; + finalized_rg.push(writer.close()?); + } + Ok((finalized_rg, rg_rows)) + }) +} + +fn spawn_parquet_parallel_serialization_task( + writer_factory: ArrowRowGroupWriterFactory, + mut data: Receiver, + serialize_tx: Sender>, + schema: Arc, +) -> JoinHandle> { + tokio::spawn(async move { + let max_buffer_rb = 10; + let max_row_group_rows = 10; + let mut row_group_index = 0; + + let column_writers = writer_factory.create_column_writers(row_group_index)?; + + let (mut col_writer_tasks, mut col_array_channels) = + spawn_column_parallel_row_group_writer(column_writers, max_buffer_rb)?; + + let mut current_rg_rows = 0; + + while let Some(mut rb) = data.recv().await { + // This loop allows the "else" block to repeatedly split the RecordBatch to handle the case + // when max_row_group_rows < execution.batch_size as an alternative to a recursive async + // function. + loop { + if current_rg_rows + rb.num_rows() < max_row_group_rows { + send_arrays_to_column_writers(&col_array_channels, &rb, &schema).await?; + current_rg_rows += rb.num_rows(); + break; + } else { + let rows_left = max_row_group_rows - current_rg_rows; + let rb_split = rb.slice(0, rows_left); + send_arrays_to_column_writers(&col_array_channels, &rb_split, &schema).await?; + + // Signal the parallel column writers that the RowGroup is done, join and finalize RowGroup + // on a separate task, so that we can immediately start on the next RG before waiting + // for the current one to finish. + drop(col_array_channels); + + let finalize_rg_task = + spawn_rg_join_and_finalize_task(col_writer_tasks, max_row_group_rows); + + // Do not surface error from closed channel (means something + // else hit an error, and the plan is shutting down). + if serialize_tx.send(finalize_rg_task).await.is_err() { + return Ok(()); + } + + current_rg_rows = 0; + rb = rb.slice(rows_left, rb.num_rows() - rows_left); + + row_group_index += 1; + let column_writers = writer_factory.create_column_writers(row_group_index)?; + (col_writer_tasks, col_array_channels) = + spawn_column_parallel_row_group_writer(column_writers, 100)?; + } + } + } + + drop(col_array_channels); + // Handle leftover rows as final rowgroup, which may be smaller than max_row_group_rows + if current_rg_rows > 0 { + let finalize_rg_task = + spawn_rg_join_and_finalize_task(col_writer_tasks, current_rg_rows); + + // Do not surface error from closed channel (means something + // else hit an error, and the plan is shutting down). + if serialize_tx.send(finalize_rg_task).await.is_err() { + return Ok(()); + } + } + + Ok(()) + }) +} + +fn spawn_column_parallel_row_group_writer( + col_writers: Vec, + max_buffer_size: usize, +) -> Result<(Vec, Vec), ParquetError> { + let num_columns = col_writers.len(); + + let mut col_writer_tasks = Vec::with_capacity(num_columns); + let mut col_array_channels = Vec::with_capacity(num_columns); + for mut col_writer in col_writers.into_iter() { + let (send_array, mut receive_array) = + tokio::sync::mpsc::channel::(max_buffer_size); + col_array_channels.push(send_array); + let handle = tokio::spawn(async move { + while let Some(col) = receive_array.recv().await { + col_writer.write(&col)?; + } + Ok(col_writer) + }); + col_writer_tasks.push(handle); + } + Ok((col_writer_tasks, col_array_channels)) +} + +/// Consume RowGroups serialized by other parallel tasks and concatenate them +/// to the final parquet file +async fn concatenate_parallel_row_groups( + mut parquet_writer: SerializedFileWriter, + mut serialize_rx: Receiver>, +) -> Result { + while let Some(task) = serialize_rx.recv().await { + let result = task.await; + let mut rg_out = parquet_writer.next_row_group()?; + let (serialized_columns, _cnt) = + result.map_err(|e| ParquetError::General(e.to_string()))??; + + for column_chunk in serialized_columns { + column_chunk.append_to_row_group(&mut rg_out)?; + } + rg_out.close()?; + } + + let file_metadata = parquet_writer.close()?; + Ok(file_metadata) +} + +// This test is based on DataFusion's ParquetSink. Motivation is to test +// concurrent writing of encrypted data over multiple row groups using the low-level API. +#[tokio::test] +async fn test_concurrent_encrypted_writing_over_multiple_row_groups() { + // Read example data and set up encryption/decryption properties + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); + let file = std::fs::File::open(path).unwrap(); + + let file_encryption_properties = FileEncryptionProperties::builder(AES_128_FOOTER_KEY.into()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].into()) + .with_column_key(AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1].into()) + .build() + .unwrap(); + let decryption_properties = FileDecryptionProperties::builder(AES_128_FOOTER_KEY.into()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].into()) + .with_column_key(AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1].into()) + .build() + .unwrap(); + + let (record_batches, metadata) = + read_encrypted_file(&file, decryption_properties.clone()).unwrap(); + let schema = metadata.schema(); + + // Create a channel to send RecordBatches to the writer and send row groups + let (record_batch_tx, data) = tokio::sync::mpsc::channel::(100); + let data_generator = tokio::spawn(async move { + for record_batch in record_batches { + record_batch_tx.send(record_batch).await.unwrap(); + } + }); + + let props = Arc::new( + WriterPropertiesBuilder::default() + .with_file_encryption_properties(file_encryption_properties) + .build(), + ); + let parquet_schema = ArrowSchemaConverter::new() + .with_coerce_types(props.coerce_types()) + .convert(schema) + .unwrap(); + + // Create a temporary file to write the encrypted data + let temp_file = tempfile::tempfile().unwrap(); + + let writer = + SerializedFileWriter::new(&temp_file, parquet_schema.root_schema_ptr(), props).unwrap(); + let row_group_writer_factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(schema)); + let max_row_groups = 1; + + let (serialize_tx, serialize_rx) = + tokio::sync::mpsc::channel::>(max_row_groups); + + let launch_serialization_task = spawn_parquet_parallel_serialization_task( + row_group_writer_factory, + data, + serialize_tx, + schema.clone(), + ); + + let _file_metadata = concatenate_parallel_row_groups(writer, serialize_rx) + .await + .unwrap(); + + data_generator.await.unwrap(); + launch_serialization_task.await.unwrap().unwrap(); + + // Check that the file was written correctly + let (read_record_batches, read_metadata) = + read_encrypted_file(&temp_file, decryption_properties.clone()).unwrap(); + + assert_eq!(read_metadata.metadata().file_metadata().num_rows(), 50); + verify_encryption_test_data(read_record_batches, read_metadata.metadata()); +} + +#[tokio::test] +async fn test_multi_threaded_encrypted_writing() { + // Read example data and set up encryption/decryption properties + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); + let file = std::fs::File::open(path).unwrap(); + + let file_encryption_properties = FileEncryptionProperties::builder(AES_128_FOOTER_KEY.into()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].into()) + .with_column_key(AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1].into()) + .build() + .unwrap(); + let decryption_properties = FileDecryptionProperties::builder(AES_128_FOOTER_KEY.into()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].into()) + .with_column_key(AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1].into()) + .build() + .unwrap(); + + let (record_batches, metadata) = + read_encrypted_file(&file, Arc::clone(&decryption_properties)).unwrap(); + let schema = metadata.schema().clone(); + + let props = Arc::new( + WriterPropertiesBuilder::default() + .with_file_encryption_properties(file_encryption_properties) + .build(), + ); + + let parquet_schema = ArrowSchemaConverter::new() + .with_coerce_types(props.coerce_types()) + .convert(&schema) + .unwrap(); + + // Create a temporary file to write the encrypted data + let temp_file = tempfile::tempfile().unwrap(); + let mut writer = + SerializedFileWriter::new(&temp_file, parquet_schema.root_schema_ptr(), props).unwrap(); + let row_group_writer_factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema)); + + let (serialize_tx, mut serialize_rx) = + tokio::sync::mpsc::channel::>(1); + + // Create a channel to send RecordBatches to the writer and send row batches + let (record_batch_tx, mut data) = tokio::sync::mpsc::channel::(100); + let data_generator = tokio::spawn(async move { + for record_batch in record_batches { + record_batch_tx.send(record_batch).await.unwrap(); + } + }); + + // Get column writers + let col_writers = row_group_writer_factory.create_column_writers(0).unwrap(); + + let (col_writer_tasks, col_array_channels) = + spawn_column_parallel_row_group_writer(col_writers, 10).unwrap(); + + // Spawn serialization tasks for incoming RecordBatches + let launch_serialization_task = tokio::spawn(async move { + let Some(rb) = data.recv().await else { + panic!() + }; + send_arrays_to_column_writers(&col_array_channels, &rb, &schema) + .await + .unwrap(); + let finalize_rg_task = spawn_rg_join_and_finalize_task(col_writer_tasks, 10); + + serialize_tx.send(finalize_rg_task).await.unwrap(); + drop(col_array_channels); + }); + + // Append the finalized row groups to the SerializedFileWriter + while let Some(task) = serialize_rx.recv().await { + let (arrow_column_chunks, _) = task.await.unwrap().unwrap(); + let mut row_group_writer = writer.next_row_group().unwrap(); + for chunk in arrow_column_chunks { + chunk.append_to_row_group(&mut row_group_writer).unwrap(); + } + row_group_writer.close().unwrap(); + } + + // Wait for data generator and serialization task to finish + data_generator.await.unwrap(); + launch_serialization_task.await.unwrap(); + let metadata = writer.close().unwrap(); + + // Close the file writer which writes the footer + assert_eq!(metadata.file_metadata().num_rows(), 50); + + // Check that the file was written correctly + let (read_record_batches, read_metadata) = + read_encrypted_file(&temp_file, decryption_properties).unwrap(); + verify_encryption_test_data(read_record_batches, read_metadata.metadata()); + + // Check that file was encrypted + let result = ArrowReaderMetadata::load(&temp_file, ArrowReaderOptions::default()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" + ); +} + +#[tokio::test] +async fn test_multi_threaded_encrypted_writing_deprecated() { + // Read example data and set up encryption/decryption properties + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); + let file = std::fs::File::open(path).unwrap(); + + let file_encryption_properties = FileEncryptionProperties::builder(AES_128_FOOTER_KEY.into()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].into()) + .with_column_key(AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1].into()) + .build() + .unwrap(); + let decryption_properties = FileDecryptionProperties::builder(AES_128_FOOTER_KEY.into()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].into()) + .with_column_key(AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1].into()) + .build() + .unwrap(); + + let (record_batches, metadata) = + read_encrypted_file(&file, Arc::clone(&decryption_properties)).unwrap(); + let to_write: Vec<_> = record_batches + .iter() + .flat_map(|rb| rb.columns().to_vec()) + .collect(); + let schema = metadata.schema().clone(); + + let props = Some( + WriterPropertiesBuilder::default() + .with_file_encryption_properties(file_encryption_properties) + .build(), + ); + + // Create a temporary file to write the encrypted data + let temp_file = tempfile::tempfile().unwrap(); + let mut writer = ArrowWriter::try_new(&temp_file, schema.clone(), props).unwrap(); + + // LOW-LEVEL API: Use low level API to write into a file using multiple threads + + // Get column writers + #[allow(deprecated)] + let col_writers = writer.get_column_writers().unwrap(); + let num_columns = col_writers.len(); + + let (col_writer_tasks, mut col_array_channels) = + spawn_column_parallel_row_group_writer(col_writers, 100).unwrap(); + + // Send the ArrowLeafColumn data to the respective column writer channels + let mut worker_iter = col_array_channels.iter_mut(); + for (array, field) in to_write.iter().zip(schema.fields()) { + for leaves in compute_leaves(field, array).unwrap() { + worker_iter.next().unwrap().send(leaves).await.unwrap(); + } + } + drop(col_array_channels); + + // Wait for all column writers to finish writing + let mut finalized_rg = Vec::with_capacity(num_columns); + for task in col_writer_tasks.into_iter() { + finalized_rg.push(task.await.unwrap().unwrap().close().unwrap()); + } + + // Append the finalized row group to the SerializedFileWriter + #[allow(deprecated)] + writer.append_row_group(finalized_rg).unwrap(); + + // HIGH-LEVEL API: Write RecordBatches into the file using ArrowWriter + + // Write individual RecordBatches into the file + for rb in record_batches { + writer.write(&rb).unwrap() + } + assert!(writer.flush().is_ok()); + + // Close the file writer which writes the footer + let metadata = writer.finish().unwrap(); + assert_eq!(metadata.file_metadata().num_rows(), 100); + + // Check that the file was written correctly + let (read_record_batches, read_metadata) = + read_encrypted_file(&temp_file, decryption_properties).unwrap(); + verify_encryption_double_test_data(read_record_batches, read_metadata.metadata()); + + // Check that file was encrypted + let result = ArrowReaderMetadata::load(&temp_file, ArrowReaderOptions::default()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" + ); +} diff --git a/vendor/parquet/tests/encryption/encryption_disabled.rs b/vendor/parquet/tests/encryption/encryption_disabled.rs new file mode 100644 index 0000000..8b38fd5 --- /dev/null +++ b/vendor/parquet/tests/encryption/encryption_disabled.rs @@ -0,0 +1,63 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; +use std::fs::File; + +#[test] +fn test_read_without_encryption_enabled_fails() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/uniform_encryption.parquet.encrypted"); + let file = File::open(path).unwrap(); + + let options = ArrowReaderOptions::default(); + let result = ArrowReaderMetadata::load(&file, options.clone()); + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file has an encrypted footer but the encryption feature is disabled" + ); +} + +#[tokio::test] +#[cfg(feature = "async")] +async fn test_async_read_without_encryption_enabled_fails() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/uniform_encryption.parquet.encrypted"); + let mut file = tokio::fs::File::open(&path).await.unwrap(); + + let options = ArrowReaderOptions::new(); + let result = ArrowReaderMetadata::load_async(&mut file, options).await; + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file has an encrypted footer but the encryption feature is disabled" + ); +} + +#[test] +#[cfg(feature = "snap")] +fn test_plaintext_footer_read_without_decryption() { + crate::encryption_agnostic::read_plaintext_footer_file_without_decryption_properties(); +} + +#[tokio::test] +#[cfg(all(feature = "async", feature = "snap"))] +async fn test_plaintext_footer_read_without_decryption_async() { + crate::encryption_agnostic::read_plaintext_footer_file_without_decryption_properties_async() + .await; +} diff --git a/vendor/parquet/tests/encryption/encryption_util.rs b/vendor/parquet/tests/encryption/encryption_util.rs new file mode 100644 index 0000000..f0fe666 --- /dev/null +++ b/vendor/parquet/tests/encryption/encryption_util.rs @@ -0,0 +1,399 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_array::cast::AsArray; +use arrow_array::{RecordBatch, types}; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::{ + ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReaderBuilder, +}; +use parquet::encryption::decrypt::{FileDecryptionProperties, KeyRetriever}; +use parquet::encryption::encrypt::FileEncryptionProperties; +use parquet::errors::{ParquetError, Result}; +use parquet::file::metadata::ParquetMetaData; +use parquet::file::properties::WriterProperties; +use ring::aead::AES_256_GCM; +use std::collections::HashMap; +use std::fs::File; +use std::sync::{Arc, Mutex}; + +pub(crate) const AES_128_FOOTER_KEY: &[u8; 16] = b"0123456789012345"; // 128bit/16 +pub(crate) const BAD_AES_128_FOOTER_KEY: &[u8; 16] = b"0000000000000000"; +pub(crate) const AES_128_FOOTER_KEY_NAME: &str = "kf"; +pub(crate) const AES_128_KEY_NAMES: [&str; 2] = ["kc1", "kc2"]; +pub(crate) const AES_128_COLUMN_NAMES: [&str; 2] = ["double_field", "float_field"]; +pub(crate) const AES_128_COLUMN_KEYS: [&[u8; 16]; 2] = [b"1234567890123450", b"1234567890123451"]; + +pub(crate) const AES_128_COLUMN_NAME_KEYS: &[(&str, &[u8]); 2] = &[ + (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), +]; + +pub(crate) const AES_128_KEY_NAME_KEY: &[(&str, &[u8]); 3] = &[ + (AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY), + (AES_128_KEY_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_KEY_NAMES[1], AES_128_COLUMN_KEYS[1]), +]; + +pub(crate) const AES_256_FOOTER_KEY: &[u8; 32] = b"01234567890123456789012345678901"; // 256bit/32 +pub(crate) const BAD_AES_256_FOOTER_KEY: &[u8; 32] = b"00000000000000000000000000000000"; +pub(crate) const AES_256_FOOTER_KEY_NAME: &str = "kf"; +pub(crate) const AES_256_KEY_NAMES: [&str; 8] = + ["kc1", "kc2", "kc3", "kc4", "kc5", "kc6", "kc7", "kc8"]; + +pub(crate) const AES_256_COLUMN_NAMES: [&str; 8] = [ + "double_field", + "float_field", + "boolean_field", + "int32_field", + "ba_field", + "flba_field", + "int64_field.list.element", + "int96_field", +]; +pub(crate) const AES_256_COLUMN_KEYS: [&[u8]; 8] = [ + b"12345678901234567890123456789012", + b"12345678901234567890123456789013", + b"12345678901234567890123456789014", + b"12345678901234567890123456789015", + b"12345678901234567890123456789016", + b"12345678901234567890123456789017", + b"12345678901234567890123456789018", + b"12345678901234567890123456789019", +]; + +pub(crate) const AES_256_COLUMN_NAME_KEYS: &[(&str, &[u8]); 8] = &[ + (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), +]; + +pub(crate) const AES_256_KEY_NAME_KEY: &[(&str, &[u8]); 9] = &[ + (AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY), + (AES_256_KEY_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_KEY_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_KEY_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_KEY_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_KEY_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_KEY_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_KEY_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_KEY_NAMES[7], AES_256_COLUMN_KEYS[7]), +]; + +pub(crate) fn verify_encryption_double_test_data( + record_batches: Vec, + metadata: &ParquetMetaData, +) { + let file_metadata = metadata.file_metadata(); + assert_eq!(file_metadata.num_rows(), 100); + assert_eq!(file_metadata.schema_descr().num_columns(), 8); + + metadata.row_groups().iter().for_each(|rg| { + assert_eq!(rg.num_columns(), 8); + assert_eq!(rg.num_rows(), 50); + }); + + let mut row_count = 0; + let wrap_at = 50; + for batch in record_batches { + let batch = batch; + row_count += batch.num_rows(); + + let bool_col = batch.column(0).as_boolean(); + let time_col = batch + .column(1) + .as_primitive::(); + let list_col = batch.column(2).as_list::(); + let timestamp_col = batch + .column(3) + .as_primitive::(); + let f32_col = batch.column(4).as_primitive::(); + let f64_col = batch.column(5).as_primitive::(); + let binary_col = batch.column(6).as_binary::(); + let fixed_size_binary_col = batch.column(7).as_fixed_size_binary(); + + for (i, x) in bool_col.iter().enumerate() { + assert_eq!(x.unwrap(), i % 2 == 0); + } + for (i, x) in time_col.iter().enumerate() { + assert_eq!(x.unwrap(), (i % wrap_at) as i32); + } + for (i, list_item) in list_col.iter().enumerate() { + let list_item = list_item.unwrap(); + let list_item = list_item.as_primitive::(); + assert_eq!(list_item.len(), 2); + assert_eq!( + list_item.value(0), + (((i % wrap_at) * 2) * 1000000000000) as i64 + ); + assert_eq!( + list_item.value(1), + (((i % wrap_at) * 2 + 1) * 1000000000000) as i64 + ); + } + for x in timestamp_col.iter() { + assert!(x.is_some()); + } + for (i, x) in f32_col.iter().enumerate() { + assert_eq!(x.unwrap(), (i % wrap_at) as f32 * 1.1f32); + } + for (i, x) in f64_col.iter().enumerate() { + assert_eq!(x.unwrap(), (i % wrap_at) as f64 * 1.1111111f64); + } + for (i, x) in binary_col.iter().enumerate() { + assert_eq!(x.is_some(), i % 2 == 0); + if let Some(x) = x { + assert_eq!(&x[0..7], b"parquet"); + } + } + for (i, x) in fixed_size_binary_col.iter().enumerate() { + assert_eq!(x.unwrap(), &[(i % wrap_at) as u8; 10]); + } + } + + assert_eq!(row_count, file_metadata.num_rows() as usize); +} + +/// Verifies data read from an encrypted file from the parquet-testing repository +pub(crate) fn verify_encryption_test_data( + record_batches: Vec, + metadata: &ParquetMetaData, +) { + let file_metadata = metadata.file_metadata(); + assert_eq!(file_metadata.num_rows(), 50); + assert_eq!(file_metadata.schema_descr().num_columns(), 8); + + let mut total_rows = 0; + metadata.row_groups().iter().for_each(|rg| { + assert_eq!(rg.num_columns(), 8); + total_rows += rg.num_rows(); + }); + assert_eq!(total_rows, 50); + + let mut row_count = 0; + for batch in record_batches { + let batch = batch; + + let row_index = |index_in_batch: usize| row_count + index_in_batch; + + let bool_col = batch.column(0).as_boolean(); + let time_col = batch + .column(1) + .as_primitive::(); + let list_col = batch.column(2).as_list::(); + let timestamp_col = batch + .column(3) + .as_primitive::(); + let f32_col = batch.column(4).as_primitive::(); + let f64_col = batch.column(5).as_primitive::(); + let binary_col = batch.column(6).as_binary::(); + let fixed_size_binary_col = batch.column(7).as_fixed_size_binary(); + + for (i, x) in bool_col.iter().enumerate() { + assert_eq!(x.unwrap(), row_index(i) % 2 == 0); + } + for (i, x) in time_col.iter().enumerate() { + assert_eq!(x.unwrap(), row_index(i) as i32); + } + for (i, list_item) in list_col.iter().enumerate() { + let list_item = list_item.unwrap(); + let list_item = list_item.as_primitive::(); + assert_eq!(list_item.len(), 2); + assert_eq!( + list_item.value(0), + ((row_index(i) * 2) * 1000000000000) as i64 + ); + assert_eq!( + list_item.value(1), + ((row_index(i) * 2 + 1) * 1000000000000) as i64 + ); + } + for x in timestamp_col.iter() { + assert!(x.is_some()); + } + for (i, x) in f32_col.iter().enumerate() { + assert_eq!(x.unwrap(), row_index(i) as f32 * 1.1f32); + } + for (i, x) in f64_col.iter().enumerate() { + assert_eq!(x.unwrap(), row_index(i) as f64 * 1.1111111f64); + } + for (i, x) in binary_col.iter().enumerate() { + assert_eq!(x.is_some(), row_index(i) % 2 == 0); + if let Some(x) = x { + assert_eq!(&x[0..7], b"parquet"); + } + } + for (i, x) in fixed_size_binary_col.iter().enumerate() { + assert_eq!(x.unwrap(), &[row_index(i) as u8; 10]); + } + + row_count += batch.num_rows(); + } + + assert_eq!(row_count, file_metadata.num_rows() as usize); +} + +/// Verifies that the column and offset indexes were successfully read from an +/// encrypted test file. +pub(crate) fn verify_column_indexes(metadata: &ParquetMetaData) { + let offset_index = metadata.offset_index().unwrap(); + // 1 row group, 8 columns + assert_eq!(offset_index.len(), 1); + assert_eq!(offset_index[0].len(), 8); + // Check float column, which is encrypted in the non-uniform test file + let float_col_idx = 4; + let offset_index = &offset_index[0][float_col_idx]; + assert_eq!(offset_index.page_locations.len(), 1); + assert!(offset_index.page_locations[0].offset > 0); + + let column_index = metadata.column_index().unwrap(); + assert_eq!(column_index.len(), 1); + assert_eq!(column_index[0].len(), 8); + let column_index = &column_index[0][float_col_idx]; + + match column_index { + parquet::file::page_index::column_index::ColumnIndexMetaData::FLOAT(float_index) => { + assert_eq!(float_index.num_pages(), 1); + assert_eq!(float_index.min_value(0), Some(&0.0f32)); + assert!( + float_index + .max_value(0) + .is_some_and(|max| (max - 53.9).abs() < 1e-6) + ); + } + _ => { + panic!("Expected a float column index for column {float_col_idx}"); + } + }; +} + +pub(crate) fn read_encrypted_file( + file: &File, + decryption_properties: Arc, +) -> std::result::Result<(Vec, ArrowReaderMetadata), ParquetError> { + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let metadata = ArrowReaderMetadata::load(file, options.clone())?; + + let builder = + ParquetRecordBatchReaderBuilder::try_new_with_options(file.try_clone().unwrap(), options)?; + let batch_reader = builder.build()?; + let batches = batch_reader.collect::, _>>()?; + Ok((batches, metadata)) +} + +pub(crate) fn read_and_roundtrip_to_encrypted_file( + file: &File, + decryption_properties: Arc, + encryption_properties: Arc, +) { + // read example data + let (batches, metadata) = + read_encrypted_file(file, Arc::clone(&decryption_properties)).unwrap(); + + // write example data to a temporary file + let temp_file = tempfile::tempfile().unwrap(); + let props = WriterProperties::builder() + .with_file_encryption_properties(encryption_properties) + .build(); + + let mut writer = ArrowWriter::try_new( + temp_file.try_clone().unwrap(), + metadata.schema().clone(), + Some(props), + ) + .unwrap(); + for batch in batches { + writer.write(&batch).unwrap(); + } + + writer.close().unwrap(); + + // check re-written example data + verify_encryption_test_file_read(temp_file, decryption_properties); +} + +pub(crate) fn verify_encryption_test_file_read( + file: File, + decryption_properties: Arc, +) { + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let reader_metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); + let metadata = reader_metadata.metadata(); + + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap(); + let record_reader = builder.build().unwrap(); + let record_batches = record_reader + .map(|x| x.unwrap()) + .collect::>(); + + verify_encryption_test_data(record_batches, metadata); +} + +/// A KeyRetriever to use in Parquet encryption tests, +/// which stores a map from key names/metadata to encryption key bytes. +pub struct TestKeyRetriever { + keys: Mutex>>, +} + +impl TestKeyRetriever { + pub fn new() -> Self { + Self { + keys: Mutex::new(HashMap::default()), + } + } + + pub fn with_key(self, key_name: String, key: Vec) -> Self { + { + let mut keys = self.keys.lock().unwrap(); + keys.insert(key_name, key); + } + self + } +} + +impl KeyRetriever for TestKeyRetriever { + fn retrieve_key(&self, key_metadata: &[u8]) -> Result> { + let key_metadata = std::str::from_utf8(key_metadata).map_err(|e| { + ParquetError::General(format!("Could not convert key metadata to string: {e}")) + })?; + let keys = self.keys.lock().unwrap(); + match keys.get(key_metadata) { + Some(key) => Ok(key.clone()), + None => Err(ParquetError::General(format!( + "Could not retrieve key for metadata {key_metadata:?}" + ))), + } + } +} + +pub fn encrypted_data_path(footer_key: &[u8], file_name: &str) -> String { + let test_data = arrow::util::test_util::parquet_test_data(); + let subpath = if AES_256_GCM.key_len() == footer_key.len() { + "aes256/" + } else { + "" + }; + let path = format!("{test_data}/{subpath}/{file_name}"); + path +} diff --git a/vendor/parquet/tests/encryption/mod.rs b/vendor/parquet/tests/encryption/mod.rs new file mode 100644 index 0000000..7ade1f5 --- /dev/null +++ b/vendor/parquet/tests/encryption/mod.rs @@ -0,0 +1,26 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#[cfg(feature = "encryption")] +mod encryption; +mod encryption_agnostic; +#[cfg(all(feature = "encryption", feature = "async"))] +mod encryption_async; +#[cfg(not(feature = "encryption"))] +mod encryption_disabled; +#[cfg(feature = "encryption")] +mod encryption_util; diff --git a/vendor/parquet/tests/geospatial.rs b/vendor/parquet/tests/geospatial.rs new file mode 100644 index 0000000..388b5c0 --- /dev/null +++ b/vendor/parquet/tests/geospatial.rs @@ -0,0 +1,645 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#[cfg(all(feature = "arrow", feature = "geospatial"))] +mod test { + //! Tests for Geometry and Geography logical types that require the arrow + //! and/or geospatial features enabled + + use std::{fs::File, iter::zip, sync::Arc}; + + use arrow_array::{ArrayRef, BinaryArray, RecordBatch, create_array}; + use arrow_schema::{DataType, Field, Schema, SchemaRef, extension::ExtensionType as _}; + use bytes::Bytes; + use parquet::{ + arrow::{ + ArrowSchemaConverter, ArrowWriter, arrow_reader::ParquetRecordBatchReaderBuilder, + arrow_writer::ArrowWriterOptions, + }, + basic::{EdgeInterpolationAlgorithm, LogicalType}, + column::reader::ColumnReader, + data_type::{ByteArray, ByteArrayType}, + file::{ + metadata::{ParquetMetaData, RowGroupMetaData}, + properties::{EnabledStatistics, WriterProperties}, + reader::{FileReader, SerializedFileReader}, + writer::SerializedFileWriter, + }, + geospatial::{bounding_box::BoundingBox, statistics::GeospatialStatistics}, + schema::types::SchemaDescriptor, + }; + use parquet_geospatial::{WkbEdges, WkbMetadata, WkbType, testing::wkb_point_xy}; + use serde_json::Value; + + fn read_metadata(geospatial_test_file: &str) -> (Arc, SchemaRef) { + let path = format!( + "{}/geospatial/{geospatial_test_file}", + arrow::util::test_util::parquet_test_data(), + ); + let file = File::open(path).unwrap(); + let reader = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + + (reader.metadata().clone(), reader.schema().clone()) + } + + #[test] + fn test_read_logical_type() { + // Some crs values are short strings + let expected_metadata = [ + ( + "crs-default.parquet", + LogicalType::geometry(None), + WkbMetadata::new(Some("OGC:CRS84"), None), // omitted CRS defaults to OGC:CRS84 + ), + ( + "crs-srid.parquet", + LogicalType::geometry(Some("srid:5070".to_string())), + WkbMetadata::new(Some("srid:5070"), None), + ), + ( + "crs-projjson.parquet", + LogicalType::geometry(Some("projjson:projjson_epsg_5070".to_string())), + WkbMetadata::new(Some("projjson:projjson_epsg_5070"), None), + ), + ( + "crs-geography.parquet", + LogicalType::geography(None, None), + WkbMetadata::new(Some("OGC:CRS84"), Some(WkbEdges::Spherical)), // omitted CRS defaults to OGC:CRS84 + ), + ]; + + for (geospatial_file, expected_type, expected_field_meta) in expected_metadata { + let (metadata, schema) = read_metadata(geospatial_file); + let column_descr = metadata.file_metadata().schema_descr().column(1); + let logical_type = column_descr.logical_type_ref().unwrap(); + + assert_eq!(logical_type, &expected_type); + + let field = schema.field(1); + let wkb_type = field.try_extension_type::().unwrap(); + + assert_eq!(wkb_type.metadata().crs, expected_field_meta.crs); + assert_eq!(wkb_type.metadata().algorithm, expected_field_meta.algorithm); + } + + // The crs value may also contain arbitrary values (in this case some JSON + // a bit too lengthy to type out) + let (metadata, schema) = read_metadata("crs-arbitrary-value.parquet"); + let column_descr = metadata.file_metadata().schema_descr().column(1); + let logical_type = column_descr.logical_type_ref().unwrap(); + + if let LogicalType::Geometry(geometry) = logical_type { + let crs = geometry.crs.as_ref(); + let crs_parsed: Value = serde_json::from_str(crs.unwrap()).unwrap(); + assert_eq!(crs_parsed.get("id").unwrap().get("code").unwrap(), 5070); + } else { + panic!("Expected geometry type but got {logical_type:?}"); + } + + let field = schema.field(1); + let wkb_type = field.try_extension_type::().unwrap(); + assert_eq!( + wkb_type.metadata().crs.as_ref().unwrap()["id"]["code"], + 5070 + ); + assert_eq!(wkb_type.metadata().algorithm, None); + } + + #[test] + fn test_read_geospatial_statistics() { + let (metadata, _) = read_metadata("geospatial.parquet"); + + // geospatial.parquet schema: + // optional binary field_id=-1 group (String); + // optional binary field_id=-1 wkt (String); + // optional binary field_id=-1 geometry (Geometry(crs=)); + let fields = metadata.file_metadata().schema().get_fields(); + let logical_type = fields[2].get_basic_info().logical_type_ref().unwrap(); + assert_eq!(logical_type, &LogicalType::geometry(None)); + + let geo_statistics = metadata.row_group(0).column(2).geo_statistics(); + assert!(geo_statistics.is_some()); + + let expected_bbox = BoundingBox::new(10.0, 40.0, 10.0, 40.0) + .with_zrange(30.0, 80.0) + .with_mrange(200.0, 1600.0); + let expected_geospatial_types = vec![ + 1, 2, 3, 4, 5, 6, 7, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 2001, 2002, 2003, 2004, + 2005, 2006, 2007, 3001, 3002, 3003, 3004, 3005, 3006, 3007, + ]; + assert_eq!( + geo_statistics.unwrap().geospatial_types(), + Some(&expected_geospatial_types) + ); + assert_eq!(geo_statistics.unwrap().bounding_box(), Some(&expected_bbox)); + } + + fn read_row_group_metadata(b: Bytes) -> Vec { + let reader = SerializedFileReader::new(b).unwrap(); + reader.metadata().row_groups().to_vec() + } + + fn read_geo_statistics(b: Bytes, column: usize) -> Vec> { + read_row_group_metadata(b) + .iter() + .map(|row_group| row_group.column(column).geo_statistics().cloned()) + .collect() + } + + #[test] + fn test_write_statistics_not_arrow() { + // Four row groups: one all non-null, one with a null, one with all nulls, + // one with invalid WKB + let column_values = vec![ + [wkb_point_xy(1.0, 2.0), wkb_point_xy(11.0, 12.0)].map(ByteArray::from), + ["this is not valid wkb".into(), wkb_point_xy(31.0, 32.0)].map(ByteArray::from), + [wkb_point_xy(21.0, 22.0), vec![]].map(ByteArray::from), + [ByteArray::new(), ByteArray::new()], + ]; + let def_levels = [[1, 1], [1, 1], [1, 0], [0, 0]]; + + // Ensure that nulls are omitted, that completely empty stats are omitted, + // and that invalid WKB results in empty stats + let expected_geometry_types = [Some(vec![1]), None, Some(vec![1]), None]; + let expected_bounding_box = [ + Some(BoundingBox::new(1.0, 11.0, 2.0, 12.0)), + None, + Some(BoundingBox::new(21.0, 21.0, 22.0, 22.0)), + None, + ]; + + let schema = parquet_schema_geometry(); + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Chunk) + .build(); + + let mut buf = Vec::with_capacity(1024); + let mut writer = + SerializedFileWriter::new(&mut buf, schema.root_schema_ptr(), Arc::new(props)).unwrap(); + + for (def_levels, values) in zip(&def_levels, &column_values) { + let mut rg = writer.next_row_group().unwrap(); + let mut col = rg.next_column().unwrap().unwrap(); + col.typed::() + .write_batch(values, Some(def_levels), None) + .unwrap(); + col.close().unwrap(); + rg.close().unwrap(); + } + + writer.close().unwrap(); + + // Check geospatial statistics on file read + let buf_bytes = Bytes::from(buf); + let all_geo_stats = read_geo_statistics(buf_bytes.clone(), 0); + assert_eq!(all_geo_stats.len(), column_values.len()); + assert_eq!(expected_geometry_types.len(), column_values.len()); + assert_eq!(expected_bounding_box.len(), column_values.len()); + + for i in 0..column_values.len() { + if let Some(geo_stats) = all_geo_stats[i].as_ref() { + assert_eq!( + geo_stats.geospatial_types(), + expected_geometry_types[i].as_ref() + ); + assert_eq!(geo_stats.bounding_box(), expected_bounding_box[i].as_ref()); + } else { + assert!(expected_geometry_types[i].is_none()); + assert!(expected_bounding_box[i].is_none()); + } + } + + for (i, rg) in read_row_group_metadata(buf_bytes).iter().enumerate() { + // We should have written Statistics with a null_count + let stats = rg.column(0).statistics().unwrap(); + let expected_null_count: u64 = def_levels[i].iter().map(|l| (*l == 0) as u64).sum(); + assert_eq!(stats.null_count_opt(), Some(expected_null_count)); + + // ...but there should be no min or max value + assert!(stats.min_bytes_opt().is_none()); + assert!(stats.max_bytes_opt().is_none()); + + // There should be no index for this column + assert!(rg.column(0).column_index_length().is_none()); + assert!(rg.column(0).column_index_offset().is_none()); + } + } + + #[test] + fn test_write_statistics_arrow() { + let arrow_schema = Arc::new(Schema::new(vec![Field::new( + "geom", + DataType::Binary, + true, + )])); + + // Check the same cases as for the non-arrow writer. These need checking again because + // the arrow writer uses a different encoder where the code path for skipping nulls + // is independent. + let column_values = [ + wkb_array_xy([Some((1.0, 2.0)), Some((11.0, 12.0))]), + create_array!( + Binary, + [ + "this is not valid wkb".as_bytes(), + &wkb_point_xy(31.0, 32.0) + ] + ), + wkb_array_xy([Some((21.0, 22.0)), None]), + wkb_array_xy([None, None]), + ]; + + let expected_geometry_types = [Some(vec![1]), None, Some(vec![1]), None]; + let expected_bounding_box = [ + Some(BoundingBox::new(1.0, 11.0, 2.0, 12.0)), + None, + Some(BoundingBox::new(21.0, 21.0, 22.0, 22.0)), + None, + ]; + + let schema = parquet_schema_geometry(); + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Chunk) + .build(); + let options = ArrowWriterOptions::new() + .with_parquet_schema(schema) + .with_properties(props); + + let mut buf = Vec::with_capacity(1024); + let mut file_writer = + ArrowWriter::try_new_with_options(&mut buf, arrow_schema.clone(), options).unwrap(); + + for values in &column_values { + let batch = RecordBatch::try_new(arrow_schema.clone(), vec![values.clone()]).unwrap(); + file_writer.write(&batch).unwrap(); + file_writer.flush().unwrap(); + } + + file_writer.close().unwrap(); + + // Check statistics on file read + let buf_bytes = Bytes::from(buf); + let all_geo_stats = read_geo_statistics(buf_bytes.clone(), 0); + assert_eq!(all_geo_stats.len(), column_values.len()); + + for i in 0..column_values.len() { + if let Some(geo_stats) = all_geo_stats[i].as_ref() { + assert_eq!( + geo_stats.geospatial_types(), + expected_geometry_types[i].as_ref() + ); + assert_eq!(geo_stats.bounding_box(), expected_bounding_box[i].as_ref()); + } else { + assert!(expected_geometry_types[i].is_none()); + assert!(expected_bounding_box[i].is_none()); + } + } + + for (i, rg) in read_row_group_metadata(buf_bytes).iter().enumerate() { + // We should have written Statistics with a null_count + let stats = rg.column(0).statistics().unwrap(); + let expected_null_count = column_values[i].null_count(); + assert_eq!(stats.null_count_opt(), Some(expected_null_count as u64)); + + // ...but there should be no min or max value + assert!(stats.min_bytes_opt().is_none()); + assert!(stats.max_bytes_opt().is_none()); + + // There should be no index for this column + assert!(rg.column(0).column_index_length().is_none()); + assert!(rg.column(0).column_index_offset().is_none()); + } + } + + #[test] + fn test_roundtrip_statistics_geospatial() { + let path = format!( + "{}/geospatial/geospatial.parquet", + arrow::util::test_util::parquet_test_data(), + ); + + test_roundtrip_statistics(&path, 2); + } + + #[test] + fn test_roundtrip_geospatial_with_nan() { + let path = format!( + "{}/geospatial/geospatial-with-nan.parquet", + arrow::util::test_util::parquet_test_data(), + ); + + test_roundtrip_statistics(&path, 0); + } + + #[test] + fn test_roundtrip_statistics_crs() { + let path = format!( + "{}/geospatial/crs-default.parquet", + arrow::util::test_util::parquet_test_data(), + ); + + test_roundtrip_statistics(&path, 0); + } + + fn test_roundtrip_statistics(path: &str, column: usize) { + let file_bytes = Bytes::from(std::fs::read(path).unwrap()); + + let reader = SerializedFileReader::new(file_bytes.clone()).unwrap(); + let mut values = Vec::new(); + let mut def_levels = Vec::new(); + + let schema = parquet_schema_geometry(); + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Chunk) + .build(); + + let mut buf = Vec::with_capacity(1024); + let mut writer = + SerializedFileWriter::new(&mut buf, schema.root_schema_ptr(), Arc::new(props)).unwrap(); + + for i in 0..reader.num_row_groups() { + let row_group = reader.get_row_group(i).unwrap(); + values.clear(); + def_levels.clear(); + + let mut row_group_out = writer.next_row_group().unwrap(); + + if let ColumnReader::ByteArrayColumnReader(mut reader) = + row_group.get_column_reader(column).unwrap() + { + reader + .read_records(1000000, Some(&mut def_levels), None, &mut values) + .unwrap(); + + let mut col = row_group_out.next_column().unwrap().unwrap(); + col.typed::() + .write_batch(&values, Some(&def_levels), None) + .unwrap(); + col.close().unwrap(); + row_group_out.close().unwrap(); + } else { + panic!("Unexpected geometry column type"); + } + } + + writer.close().unwrap(); + + let actual_stats = read_geo_statistics(buf.into(), 0); + let expected_stats = read_geo_statistics(file_bytes.clone(), column); + + assert_eq!(actual_stats.len(), expected_stats.len()); + for i in 0..expected_stats.len() { + assert_eq!(actual_stats[i], expected_stats[i], "Row group {i}"); + } + } + + fn parquet_schema_geometry() -> SchemaDescriptor { + let wkb_meta = WkbMetadata::new(None, None); + let wkb_type = WkbType::new(Some(wkb_meta)); + + let field = Field::new("geo", DataType::Binary, true).with_extension_type(wkb_type); + let schema = Schema::new(vec![field]); + + ArrowSchemaConverter::new().convert(&schema).unwrap() + } + + fn wkb_array_xy(coords: impl IntoIterator>) -> ArrayRef { + let array = BinaryArray::from_iter( + coords + .into_iter() + .map(|maybe_xy| maybe_xy.map(|(x, y)| wkb_point_xy(x, y))), + ); + Arc::new(array) + } + + #[test] + fn test_logical_type_to_field_conversion() { + use parquet::arrow::parquet_to_arrow_schema; + use parquet::basic::Type as PhysicalType; + use parquet::schema::types::{SchemaDescriptor, Type}; + + // Test cases: (LogicalType, expected metadata JSON) + let test_cases = [ + // Geometry with default CRS (defaults to OGC:CRS84 per Parquet spec) + (LogicalType::geometry(None), r#"{"crs":"OGC:CRS84"}"#), + // Geometry with srid:0 should result in an unset (omitted) CRS + (LogicalType::geometry(Some("srid:0".to_string())), r#"{}"#), + // Geometry with custom CRSes (authority:code and partial projjson) + ( + LogicalType::geometry(Some("EPSG:4267".to_string())), + r#"{"crs":"EPSG:4267"}"#, + ), + ( + LogicalType::geometry(Some( + r#"{"id":{"authority":"EPSG","code":4326}}"#.to_string(), + )), + r#"{"crs":{"id":{"authority":"EPSG","code":4326}}}"#, + ), + // Geography with default CRS (default OGC:CRS84, spherical edges) + ( + LogicalType::geography(None, None), + r#"{"crs":"OGC:CRS84","edges":"spherical"}"#, + ), + // Geography with explicit edges + ( + LogicalType::geography(None, Some(EdgeInterpolationAlgorithm::SPHERICAL)), + r#"{"crs":"OGC:CRS84","edges":"spherical"}"#, + ), + ( + LogicalType::geography(None, Some(EdgeInterpolationAlgorithm::KARNEY)), + r#"{"crs":"OGC:CRS84","edges":"karney"}"#, + ), + ( + LogicalType::geography(None, Some(EdgeInterpolationAlgorithm::VINCENTY)), + r#"{"crs":"OGC:CRS84","edges":"vincenty"}"#, + ), + ( + LogicalType::geography(None, Some(EdgeInterpolationAlgorithm::ANDOYER)), + r#"{"crs":"OGC:CRS84","edges":"andoyer"}"#, + ), + ( + LogicalType::geography(None, Some(EdgeInterpolationAlgorithm::THOMAS)), + r#"{"crs":"OGC:CRS84","edges":"thomas"}"#, + ), + // Geometry with srid:0 should result in an unset (omitted) CRS + // and spherical edges + ( + LogicalType::geography(Some("srid:0".to_string()), None), + r#"{"edges":"spherical"}"#, + ), + // Geography with custom CRSes (authority:code and partial projjson) + ( + LogicalType::geography(Some("EPSG:4267".to_string()), None), + r#"{"crs":"EPSG:4267","edges":"spherical"}"#, + ), + ( + LogicalType::geography( + Some(r#"{"id":{"authority":"EPSG","code":4326}}"#.to_string()), + None, + ), + r#"{"crs":{"id":{"authority":"EPSG","code":4326}},"edges":"spherical"}"#, + ), + ]; + + for (logical_type, expected_metadata) in test_cases { + // Build a Parquet schema with the given LogicalType + let parquet_schema = SchemaDescriptor::new(Arc::new( + Type::group_type_builder("schema") + .with_fields(vec![Arc::new( + Type::primitive_type_builder("geom", PhysicalType::BYTE_ARRAY) + .with_logical_type(Some(logical_type.clone())) + .build() + .unwrap(), + )]) + .build() + .unwrap(), + )); + + // Convert to Arrow schema + let arrow_schema = parquet_to_arrow_schema(&parquet_schema, None).unwrap(); + let field = arrow_schema.field(0); + + // Check extension type name + let ext_name = field.metadata().get("ARROW:extension:name"); + assert_eq!( + ext_name, + Some(&"geoarrow.wkb".to_string()), + "Extension name mismatch for {logical_type:?}" + ); + + // Check extension metadata + let ext_metadata = field.metadata().get("ARROW:extension:metadata"); + assert_eq!( + ext_metadata, + Some(&expected_metadata.to_string()), + "Extension metadata mismatch for {logical_type:?}" + ); + } + } + + #[test] + fn test_field_to_logical_type_conversion() { + use std::collections::HashMap; + + // Test cases: (extension metadata JSON, expected LogicalType) + let test_cases = [ + // Geometry with no CRS should be GEOMETRY(srid:0) + (r#"{}"#, LogicalType::geometry(Some("srid:0".to_string()))), + // Geometry with string CRS + ( + r#"{"crs":"EPSG:4267"}"#, + LogicalType::geometry(Some("EPSG:4267".to_string())), + ), + // Geometry with PROJJSON CRS + ( + r#"{"crs":{"id":{"authority":"EPSG","code":3857}}}"#, + LogicalType::geometry(Some( + r#"{"id":{"authority":"EPSG","code":3857}}"#.to_string(), + )), + ), + // Geometry with lon/lat CRSes (canonically removed because lon/lat is the + // default Parquet CRS) + (r#"{"crs":"OGC:CRS84"}"#, LogicalType::geometry(None)), + (r#"{"crs":"EPSG:4326"}"#, LogicalType::geometry(None)), + ( + r#"{"crs":{"id":{"authority":"EPSG","code":4326}}}"#, + LogicalType::geometry(None), + ), + ( + r#"{"crs":{"id":{"authority":"EPSG","code":"4326"}}}"#, + LogicalType::geometry(None), + ), + ( + r#"{"crs":{"id":{"authority":"OGC","code":"CRS84"}}}"#, + LogicalType::geometry(None), + ), + // Geography with no CRS, spherical edges + ( + r#"{"edges":"spherical"}"#, + LogicalType::geography(Some("srid:0".to_string()), None), + ), + // Geography with OGC:CRS84 and spherical edges + ( + r#"{"crs":"OGC:CRS84","edges":"spherical"}"#, + LogicalType::geography(None, None), + ), + // Geography with different edge algorithms + ( + r#"{"crs":"OGC:CRS84","edges":"karney"}"#, + LogicalType::geography(None, Some(EdgeInterpolationAlgorithm::KARNEY)), + ), + ( + r#"{"crs":"OGC:CRS84","edges":"vincenty"}"#, + LogicalType::geography(None, Some(EdgeInterpolationAlgorithm::VINCENTY)), + ), + ( + r#"{"crs":"OGC:CRS84","edges":"andoyer"}"#, + LogicalType::geography(None, Some(EdgeInterpolationAlgorithm::ANDOYER)), + ), + ( + r#"{"crs":"OGC:CRS84","edges":"thomas"}"#, + LogicalType::geography(None, Some(EdgeInterpolationAlgorithm::THOMAS)), + ), + // Geography with custom CRS and edges + ( + r#"{"crs":"EPSG:4267","edges":"karney"}"#, + LogicalType::geography( + Some("EPSG:4267".to_string()), + Some(EdgeInterpolationAlgorithm::KARNEY), + ), + ), + // Geography with PROJJSON CRS + ( + r#"{"crs":{"id":{"authority":"EPSG","code":4267}},"edges":"spherical"}"#, + LogicalType::geography( + Some(r#"{"id":{"authority":"EPSG","code":4267}}"#.to_string()), + None, + ), + ), + ]; + + for (ext_metadata, expected_logical_type) in test_cases { + // Create an Arrow Field with raw extension metadata + let metadata = HashMap::from([ + ( + "ARROW:extension:name".to_string(), + "geoarrow.wkb".to_string(), + ), + ( + "ARROW:extension:metadata".to_string(), + ext_metadata.to_string(), + ), + ]); + let field = + Arc::new(Field::new("geom", DataType::Binary, true).with_metadata(metadata)); + let schema = Schema::new(vec![field]); + + // Convert to Parquet schema + let parquet_schema = ArrowSchemaConverter::new().convert(&schema).unwrap(); + + // Check the logical type + let column_descr = parquet_schema.column(0); + let logical_type = column_descr.logical_type_ref(); + + assert_eq!( + logical_type, + Some(&expected_logical_type), + "LogicalType mismatch for extension metadata: {ext_metadata}" + ); + } + } +} diff --git a/vendor/parquet/tests/variant_integration.rs b/vendor/parquet/tests/variant_integration.rs new file mode 100644 index 0000000..e06426a --- /dev/null +++ b/vendor/parquet/tests/variant_integration.rs @@ -0,0 +1,424 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Comprehensive integration tests for Parquet files with Variant columns +//! +//! This test harness reads test case definitions from cases.json, loads expected +//! Variant values from .variant.bin files, reads Parquet files, converts StructArray +//! to VariantArray, and verifies that extracted values match expected results. +//! +//! Inspired by the arrow-go implementation: + +use arrow::util::test_util::parquet_test_data; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use parquet_variant::{Variant, VariantMetadata}; +use parquet_variant_compute::{VariantArray, unshred_variant}; +use serde::Deserialize; +use std::path::Path; +use std::sync::LazyLock; +use std::{fs, path::PathBuf}; + +type Result = std::result::Result; + +/// Creates a test function for a given case number. +/// +/// If an error message is provided, generate an error test case that expects it. +/// +/// Note the index is zero-based, while the case number is one-based +macro_rules! variant_test_case { + ($case_num:literal $(, $expected_error:literal )? ) => { + seq_macro::seq!(N in $case_num..=$case_num { + #[test] + $( #[should_panic(expected = $expected_error)] )? + fn test_variant_integration_case_~N() { + all_cases()[N - 1].run() + } + }); + }; +} + +// Generate test functions for each case +// Notes +// - case 3 is empty in cases.json for some reason +// - cases 40, 42, 87, 127 and 128 are expected to fail always (they include invalid variants) +// - cases 41, 131 and 138 omit the spec-required `value` column. They are read leniently by +// synthesizing an all-null `value` (https://github.com/apache/arrow-rs/issues/10306) +// - the remaining cases are expected to (eventually) pass + +variant_test_case!(1); +variant_test_case!(2); +// case 3 is empty in cases.json 🤷 +// ```json +// { +// "case_number" : 3 +// }, +// ``` +variant_test_case!(3, "parquet_file must be set"); +variant_test_case!(4); +variant_test_case!(5); +variant_test_case!(6); +variant_test_case!(7); +variant_test_case!(8); +variant_test_case!(9); +variant_test_case!(10); +variant_test_case!(11); +variant_test_case!(12); +variant_test_case!(13); +variant_test_case!(14); +variant_test_case!(15); +variant_test_case!(16); +variant_test_case!(17); +variant_test_case!(18); +variant_test_case!(19); +variant_test_case!(20); +variant_test_case!(21); +variant_test_case!(22); +variant_test_case!(23); +variant_test_case!(24); +variant_test_case!(25); +variant_test_case!(26); +variant_test_case!(27); +variant_test_case!(28); +variant_test_case!(29); +variant_test_case!(30); +variant_test_case!(31); +variant_test_case!(32); +variant_test_case!(33); +variant_test_case!(34); +variant_test_case!(35); +variant_test_case!(36); +variant_test_case!(37); +variant_test_case!(38); +variant_test_case!(39); +// Is an error case (should be failing as the expected error message indicates) +variant_test_case!(40, "both value and typed_value are non-null"); +// Schema omits the spec-required `value` column. Read leniently by synthesizing an +// all-null `value` (testArrayMissingValueColumn, https://github.com/apache/arrow-rs/issues/10306) +variant_test_case!(41); +// Is an error case (should be failing as the expected error message indicates) +variant_test_case!(42, "both value and typed_value are non-null"); +// Is an error case (should be failing as the expected error message indicates) +variant_test_case!(43, "Field 'b' appears in both typed_value and value"); +variant_test_case!(44); +variant_test_case!(45); +variant_test_case!(46); +variant_test_case!(47); +variant_test_case!(48); +variant_test_case!(49); +variant_test_case!(50); +variant_test_case!(51); +variant_test_case!(52); +variant_test_case!(53); +variant_test_case!(54); +variant_test_case!(55); +variant_test_case!(56); +variant_test_case!(57); +variant_test_case!(58); +variant_test_case!(59); +variant_test_case!(60); +variant_test_case!(61); +variant_test_case!(62); +variant_test_case!(63); +variant_test_case!(64); +variant_test_case!(65); +variant_test_case!(66); +variant_test_case!(67); +variant_test_case!(68); +variant_test_case!(69); +variant_test_case!(70); +variant_test_case!(71); +variant_test_case!(72); +variant_test_case!(73); +variant_test_case!(74); +variant_test_case!(75); +variant_test_case!(76); +variant_test_case!(77); +variant_test_case!(78); +variant_test_case!(79); +variant_test_case!(80); +variant_test_case!(81); +variant_test_case!(82); +variant_test_case!(83); +// Invalid case, implementations can choose to read the shredded value or error out +variant_test_case!(84); +variant_test_case!(85); +variant_test_case!(86); +// Is an error case (should be failing as the expected error message indicates) +variant_test_case!(87, "Expected object in value field"); +variant_test_case!(88); +variant_test_case!(89); +variant_test_case!(90); +variant_test_case!(91); +variant_test_case!(92); +variant_test_case!(93); +variant_test_case!(94); +variant_test_case!(95); +variant_test_case!(96); +variant_test_case!(97); +variant_test_case!(98); +variant_test_case!(99); +variant_test_case!(100); +variant_test_case!(101); +variant_test_case!(102); +variant_test_case!(103); +variant_test_case!(104); +variant_test_case!(105); +variant_test_case!(106); +variant_test_case!(107); +variant_test_case!(108); +variant_test_case!(109); +variant_test_case!(110); +variant_test_case!(111); +variant_test_case!(112); +variant_test_case!(113); +variant_test_case!(114); +variant_test_case!(115); +variant_test_case!(116); +variant_test_case!(117); +variant_test_case!(118); +variant_test_case!(119); +variant_test_case!(120); +variant_test_case!(121); +variant_test_case!(122); +variant_test_case!(123); +variant_test_case!(124); +// Is an error case (should be failing as the expected error message indicates) +variant_test_case!(125, "Field 'b' appears in both typed_value and value"); +variant_test_case!(126); +// Is an error case (should be failing as the expected error message indicates) +variant_test_case!(127, "Illegal shredded value type: UInt32"); +// Is an error case (should be failing as the expected error message indicates) +variant_test_case!(128, "Expected object in value field"); +variant_test_case!(129); +variant_test_case!(130); +// Schema omits the spec-required `value` column. Read leniently by synthesizing an +// all-null `value` (testMissingValueColumn, https://github.com/apache/arrow-rs/issues/10306) +variant_test_case!(131); +variant_test_case!(132); +variant_test_case!(133); +variant_test_case!(134); +variant_test_case!(135); +variant_test_case!(136); +// Is an error case (should be failing as the expected error message indicates) +variant_test_case!(137, "Illegal shredded value type: FixedSizeBinary(4)"); +// Schema omits the spec-required `value` column. Read leniently by synthesizing an +// all-null `value` (testShreddedObjectMissingValueColumn, https://github.com/apache/arrow-rs/issues/10306) +variant_test_case!(138); + +/// Test case definition structure matching the format from +/// `parquet-testing/parquet_shredded/cases.json` +/// +/// See [README] for details. +/// +/// [README]: https://github.com/apache/parquet-testing/blob/master/shredded_variant/README.md +/// +/// Example JSON +/// ```json +/// { +/// "case_number" : 5, +/// "test" : "testShreddedVariantPrimitives", +/// "parquet_file" : "case-005.parquet", +/// "variant_file" : "case-005_row-0.variant.bin", +/// "variant" : "Variant(metadata=VariantMetadata(dict={}), value=Variant(type=BOOLEAN_FALSE, value=false))" +/// }, +/// ``` +#[allow(dead_code)] // some fields are not used except when printing the struct +#[derive(Debug, Clone, Deserialize)] +struct VariantTestCase { + /// Case number (e.g., 1, 2, 4, etc. - note: case 3 is missing any data) + pub case_number: u32, + /// Test method name (e.g., "testSimpleArray") + pub test: Option, + /// Name of the parquet file (e.g., "case-001.parquet") + pub parquet_file: Option, + + /// Expected variant binary file (e.g., "case-001_row-0.variant.bin") - None for error cases + pub variant_file: Option, + /// Multiple expected variant binary files, for multi row inputs. If there + /// is no variant, there is no file + pub variant_files: Option>>, + /// Expected error message for negative test cases + /// + /// (this is the message from the cases.json file, which is from the Iceberg + /// implementation, so it is not guaranteed to match the actual Rust error message) + pub error_message: Option, + /// Description of the variant value (for debugging) + pub variant_description: Option, +} + +/// Run a single test case +impl VariantTestCase { + /// Run a test case. Panics on unexpected error + fn run(&self) { + println!("{self:#?}"); + + let variant_data = self.load_variants(); + let variant_array = self.load_parquet(); + + // `load_parquet` returns shredded variant values, but the test expectations are provided as + // unshredded variant values. Unshred (failing for invalid input) so we can compare them. + let variant_array = unshred_variant(&variant_array).unwrap(); + + // if this is an error case, the expected error message should be set + if let Some(expected_error) = &self.error_message { + // Unshredding variant array should have already triggered the error + panic!("Expected an error '{expected_error}`, but got no error"); + } + + assert_eq!( + variant_array.len(), + variant_data.len(), + "Number of variants in parquet file does not match expected number" + ); + for (i, expected) in variant_data.iter().enumerate() { + if variant_array.is_null(i) { + assert!( + expected.is_none(), + "Expected null variant at index {i}, but got {:?}", + variant_array.value(i) + ); + continue; + } + let actual = variant_array.value(i); + let expected = variant_data[i] + .as_ref() + .expect("Expected non-null variant data"); + + let expected = expected.as_variant(); + + // compare the variants (is this the right way to compare?) + assert_eq!( + actual, expected, + "Variant data mismatch at index {}\n\nactual\n{actual:#?}\n\nexpected\n{expected:#?}", + i + ); + } + } + + /// Parses the expected variant files, returning a vector of `ExpectedVariant` or None + /// if the corresponding entry in `variant_files` is null + fn load_variants(&self) -> Vec> { + let variant_files: Box>> = + match (&self.variant_files, &self.variant_file) { + (Some(files), None) => Box::new(files.iter().map(|f| f.as_ref())), + (None, Some(file)) => Box::new(std::iter::once(Some(file))), + // error cases may not have any variant files + _ => Box::new(std::iter::empty()), + }; + + // load each file + variant_files + .map(|f| { + let v = ExpectedVariant::try_load(&TEST_CASE_DIR.join(f?)) + .expect("Failed to load expected variant"); + Some(v) + }) + .collect() + } + + /// Load the parquet file, extract the Variant column, and return as a VariantArray + fn load_parquet(&self) -> VariantArray { + let parquet_file = self + .parquet_file + .as_ref() + .expect("parquet_file must be set"); + let path = TEST_CASE_DIR.join(parquet_file); + let file = fs::File::open(&path) + .unwrap_or_else(|e| panic!("cannot open parquet file {path:?}: {e}")); + + let reader = ParquetRecordBatchReaderBuilder::try_new(file) + .and_then(|b| b.build()) + .unwrap_or_else(|e| panic!("Error reading parquet reader for {path:?}: {e}")); + + let mut batches: Vec<_> = reader + .collect::>() + .unwrap_or_else(|e| panic!("Error reading parquet batches for {path:?}: {e}")); + + if batches.is_empty() { + panic!("No parquet batches were found in file {path:?}"); + } + if batches.len() > 1 { + panic!( + "Multiple parquet batches were found in file {path:?}, only single batch supported" + ); + } + let batch = batches.swap_remove(0); + + // The schema is "id", "var" for the id and variant columns + // TODO: support the actual parquet logical type annotation somehow + let var = batch + .column_by_name("var") + .unwrap_or_else(|| panic!("No 'var' column found in parquet file {path:?}")); + + VariantArray::try_new(var).unwrap_or_else(|e| { + panic!("Error converting StructArray to VariantArray for {path:?}: {e}") + }) + } +} + +/// Variant value loaded from .variant.bin file +#[derive(Debug, Clone)] +struct ExpectedVariant { + data: Vec, + data_offset: usize, +} + +impl ExpectedVariant { + fn try_load(path: &Path) -> Result { + // "Each `*.variant.bin` file contains a single variant serialized + // by concatenating the serialized bytes of the variant metadata + // followed by the serialized bytes of the variant value." + let data = fs::read(path).map_err(|e| format!("cannot read variant file {path:?}: {e}"))?; + let metadata = VariantMetadata::try_new(&data) + .map_err(|e| format!("cannot parse variant metadata from {path:?}: {e}"))?; + + let data_offset = metadata.size(); + Ok(Self { data, data_offset }) + } + + fn as_variant(&self) -> Variant<'_, '_> { + let metadata = &self.data[0..self.data_offset]; + let value = &self.data[self.data_offset..]; + Variant::try_new(metadata, value).expect("Invalid variant data") + } +} + +static TEST_CASE_DIR: LazyLock = LazyLock::new(|| { + PathBuf::from(parquet_test_data()) + .join("..") + .join("shredded_variant") +}); + +/// All tests +static ALL_CASES: LazyLock>> = LazyLock::new(|| { + let cases_file = TEST_CASE_DIR.join("cases.json"); + + if !cases_file.exists() { + return Err(format!("cases.json not found at {}", cases_file.display())); + } + + let content = fs::read_to_string(&cases_file) + .map_err(|e| format!("cannot read cases file {cases_file:?}: {e}"))?; + + serde_json::from_str::>(content.as_str()) + .map_err(|e| format!("cannot parse json from {cases_file:?}: {e}")) +}); + +// return a reference to the static ALL_CASES, or panic if loading failed +fn all_cases() -> &'static [VariantTestCase] { + ALL_CASES.as_ref().unwrap() +} diff --git a/vendor/parquet/upstream.json b/vendor/parquet/upstream.json new file mode 100644 index 0000000..09ae824 --- /dev/null +++ b/vendor/parquet/upstream.json @@ -0,0 +1,234 @@ +{ + "archive_sha256": "ff322f54b1a0f9288e614ed1f2d329b380af5476420db19f46ffb865e1163d73", + "crate": "parquet", + "files": { + "CONTRIBUTING.md": "b98075433192fb2ab33b559c58ed18ce8e21dbd6f83c5592204f7faa414befe2", + "Cargo.lock": "949c0aaf19ac3fcb0ba9680ee723f956eb71bb7478549d996d301cc85f7182a0", + "Cargo.toml": "6f4739642a5d71ba4e06d1f899fe5e94b51ea085808488c4620eb1e842306320", + "Cargo.toml.orig": "b72e67e5de4b00263cda12d0526911b16bffa2eb4137cc0db946d8fc262c6e17", + "LICENSE.txt": "cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30", + "NOTICE.txt": "e5b198d7f809311723889814d99a92f808c69e89a847efa7a7b65270f5fdeae6", + "README.md": "b13e919fe8af687cbcbd7a257886a21103d8484376916bbea6785fe17a67a96d", + "THRIFT.md": "3725bf644bd461d9bee6480128c1644d5fa47362a8f95683dda64acf99310c3a", + "benches/arrow_reader.rs": "bd605d7341307e3bf27902bcb3c7ce85ca916143dddde84f3b85e6bfd0b98085", + "benches/arrow_reader_clickbench.rs": "8c237277cd391f956a7209ab0468a75946eac230e98ca25625713b57f68d76fa", + "benches/arrow_reader_peak_memory.rs": "4f674f31d002708e0d0b9e3be02cf0cc35f884a8757f16386577a98da4222b40", + "benches/arrow_reader_row_filter.rs": "fb82ccb35e874d4edd78fd62c7d7d004ce7484447bfda3a490868edc382d036c", + "benches/arrow_reader_row_selection_policy.rs": "0ba474ff3b620b764a8c9fa5a170454fc783aad56930543d38bc39de03abb10a", + "benches/arrow_statistics.rs": "92557f96a1057ec9de7ab4ecf7e74982e8454a221ee4e9180ef63da415d231fa", + "benches/arrow_writer.rs": "e9415e45f98ecda23c57d7d23799dd972d07b691c60068e3b486ba97a6a03d4d", + "benches/bloom_filter.rs": "6076c3aa376cbfd309cc2a03cce56c245564e036d38d1aa9b4354af6bda95db7", + "benches/compression.rs": "e316fdaee780d970862a66fdc62a748c081da7f762dc84df968061e62b6b9d43", + "benches/encoding.rs": "e0dc7244bf8e5adb6eb713384cba72d07ba8ee4ce57803d548de6ba757ec10c5", + "benches/metadata.rs": "4f35d4d4e7bb5b78a8d47ceb2e0f3c06bf626740dfb9d92dc6a9e5d47845dcc2", + "benches/parquet_round_trip.rs": "654200e7bddfc5ab55f98386a090fca55917d70a478b5ab90509ff1e7d68acc0", + "benches/push_decoder.rs": "c940d236dcbdbbfa1ff7efdd63e1e4e95fe78f538d25de1139e7f2bb56c1fb4d", + "benches/row_group_index_reader.rs": "358c2cba7d7dc3af874c8a22bc7a0a2988448d2344019efcba5e0ade6a91ce0b", + "benches/row_selection_cursor.rs": "a5e285d7e7c8f21fa357cc848b0bc484607c0062b9ee7f7c18292709288dd211", + "benches/row_selection_policy_common/assertions.rs": "55747afcce443d62e7456fa32933cf330cb8b8c11a993358841720423d4e7d89", + "benches/row_selection_policy_common/cases.rs": "f79b69c500684797e4f88d8d14d79a17ba3767d389cb2ebd81dfc6e602a7ef37", + "benches/row_selection_policy_common/fixture.rs": "300bf5936a8c2635aeda9e31cfcd1afd06b420a341ab0418b25f41b6a73ce7cd", + "benches/row_selection_policy_common/mod.rs": "9345cbc668524a11e144e3ea163d34d4fa0d287e180e65f67d28fbafff4265b9", + "benches/row_selection_policy_common/model.rs": "7415e57bde8de3a681e17a9a4fe804395c4b6baf2a0960f818f5886792067fdc", + "benches/row_selection_policy_common/register.rs": "bec0cb5652fa16048540f909136ad995f2bd14439b10536fb06caccfbff3115a", + "benches/row_selection_policy_common/runner.rs": "d6a49aae7e3eb32f13246bd293d8fb65b7138ee54bbd0ac2bc86978b3764033e", + "benches/row_selection_policy_common/shapes.rs": "f11865079b19bcb7bbc10bf67d3dac063255bc8c9d9c0e0ceede09e70934e6bf", + "benches/row_selector.rs": "377e5c0c5ee7abbb0d3ddce23cae26eab514d5c93d651e8055ebf5548532e075", + "benches/row_selector_boolean_buffer.rs": "7044948d41774e5d2883476c76ef5bee494751f5617de97620a51047d2edeaf2", + "benches/writer_overhead.rs": "db62402b994fbfd29b37d828dfd7bbca0bc28710384a3c4185a27f09f9e5f778", + "examples/external_metadata.rs": "33c3e8d6b7fbcbb4b910444d669ee955d8141c40966b99d61ae328412d17bb84", + "examples/object_store.rs": "01aafed3f7b501d181aa201393af2f4baa9c074aaa1c8a7ea4fa0339b72afe42", + "examples/read_parquet.rs": "2c98c4376e475b79802aaa6e4810353d07545ca01367bc934921b77262b175e9", + "examples/read_with_rowgroup.rs": "8d0650275b32a0cfd6c510fad334c0b5938f9e971eaf51ed8c3766bd7ae0e0d3", + "examples/write_parquet.rs": "2de9708be17539d99b8f53af09b0de34ad2baef31a0a382c84dcae9880d73f7f", + "pytest/requirements.in": "71720dc5110d9097ebff0b509c377876fd44964f7422d2084fa4ef72b8e7f5c1", + "pytest/requirements.txt": "7c1cec5ae612b752418c99a7df5794b7ab0aed2bd7ae194adae9050742680617", + "pytest/test_parquet_integration.py": "26b854df59d5a03dafaf8b457afdb11a1fbded63172e6eb4b838945204e4be61", + "src/arrow/array_reader/builder.rs": "46c65dc01c1d32de85d07f0d0619c1a0aa2e7659dc0cbbee4901828811c4f3ec", + "src/arrow/array_reader/byte_array.rs": "807d85470436d1bbefbf9080897989e41853b06fa49dc181d65858c325b6a46e", + "src/arrow/array_reader/byte_array_dictionary.rs": "a8eb50f292ce4fde74fdce1d977100f0055320667c009e2788b63440f02253a3", + "src/arrow/array_reader/byte_view_array.rs": "6aefed84348efaee0001b639bcc3f861b1957c41692eae116dae1821b65f00ec", + "src/arrow/array_reader/cached_array_reader.rs": "4f5969919c53d9ffee0e2094021a7d73ee7a28ad205bfd6f1e2a404f655eb05c", + "src/arrow/array_reader/empty_array.rs": "ccb30ef265039e1a750cb9a52d7836eeca55a658310f7f5a06d17dd5641546c8", + "src/arrow/array_reader/fixed_len_byte_array.rs": "5f1f2e772ed95f1026d6ff21cc018084ff16e659fe9157ae2818b443954baf03", + "src/arrow/array_reader/fixed_size_list_array.rs": "2a3d685a2d63a1b8dcd97c1cc1ff4648730759d555d5814201213d176f3ed83b", + "src/arrow/array_reader/list_array.rs": "96a8db9a89f28622a65039b58f039646e6e99db57f147ed505004a95ec150448", + "src/arrow/array_reader/list_view_array.rs": "c5e3845c7234522286873a72a3037ebdb52d52aa9020450e6f2db708cddae87f", + "src/arrow/array_reader/map_array.rs": "710787b1c4c317127992a54458c942ae9976de6908f78ae12fc5b34656216fd2", + "src/arrow/array_reader/mod.rs": "c1fa6975aae757e4cb9776a76136cb75a7f4ee93ccd1a1efc0c52036167e5a42", + "src/arrow/array_reader/null_array.rs": "a3bb3d6c9d5423e0e58223714aa27054881afcc63f096546f9e85dccfe1fae41", + "src/arrow/array_reader/primitive_array.rs": "9df99835aa4971ee68860ff30bb09d9abc945b6ecff68718f02795334cd48c47", + "src/arrow/array_reader/row_group_cache.rs": "b8e04e7bdd70f56bae258ca451f35fecaa7688079fda44ad5f3a7d6ccc0fd628", + "src/arrow/array_reader/row_group_index.rs": "b08e0742b50a4b98b73ed9fccbd585bf1e90310f14d5b1d599e73403fccd1369", + "src/arrow/array_reader/row_number.rs": "9504db4e61235daafdddfe16851c3e9c61ad420f51a2572839a6d039770b1c93", + "src/arrow/array_reader/struct_array.rs": "2bc5cd32c6e65e0ce4a3fa436ec052ec5a7f44befe5410ab38642c59463593d2", + "src/arrow/array_reader/test_util.rs": "353a16ad1a8db9a126ab1f9ed40d162b8210d1153e8ac2e9ba8af9b4fd48045b", + "src/arrow/arrow_reader/filter.rs": "346ecec64d0cb87a679a770ace05f806727b30d724efbee7d4ae64475b8800e0", + "src/arrow/arrow_reader/metrics.rs": "dbd85435fe323fe525d78cb00127099364a4f2de9271bb9d5ff9084674d7d12e", + "src/arrow/arrow_reader/mod.rs": "4d61432d5c5f36e62f39dd96357629f94f46e0a79125e9b328a7b51e37a80fb2", + "src/arrow/arrow_reader/read_plan.rs": "6835543f3418cfd18fd86f81c74684e0db5ad51946cb4d4d5d04a4549d707e71", + "src/arrow/arrow_reader/selection/algebra.rs": "f72c754cd2a370d284a5b9462757db61ad79708f529c3b9e6086366366d09696", + "src/arrow/arrow_reader/selection/boolean.rs": "e0efed1eb87f40fbac31820681aad1315d9f6e8d6f669e292bd33667e545a006", + "src/arrow/arrow_reader/selection/cursor.rs": "e8193daee7ca14c58652ee452a3ea493781453a299d64b2b97c95539ddd7c2de", + "src/arrow/arrow_reader/selection/mod.rs": "cc1e0cf56d31432c461238b490b13a802f10cc9e8c5a1928797d14135cad49e4", + "src/arrow/arrow_reader/selection/ranges.rs": "11feaec66d89c053bce2aeed9938c1d175061a5d74762af44bf85d9f78cdd6c6", + "src/arrow/arrow_reader/selection/selector.rs": "361968fddece9883f2c500b8cdf19eccfee7db9f97c17e8e256ffa0740151298", + "src/arrow/arrow_reader/statistics.rs": "dd023b3ce6f854133de1a2202083148fce882c723878c665d6f076ff225944cc", + "src/arrow/arrow_writer/byte_array.rs": "cf3ff715cf0000e0e2bfe0879f6f026d53ebe3deb9ad3a29415313d28123d9ea", + "src/arrow/arrow_writer/levels.rs": "940b69a818962387ed19701f641c3020a93bd9175dd69c22dcb7faa678001de0", + "src/arrow/arrow_writer/mod.rs": "5a9e199a90d7dc40aef1f0f722affb2c48c60c12e30b73b7cd91f70c6ca01b66", + "src/arrow/async_reader/metadata.rs": "04215161ec9ef03fd383a0d02879c65ddb24eea58188efec7a7c9c759bd8bd6c", + "src/arrow/async_reader/mod.rs": "aaae0c35bba4bf9f9c16dd370e0f4381699df5723a27599406e4b69e6469082b", + "src/arrow/async_reader/spawn.rs": "904b38c23e9a28ac57776f38589689726d4aa1a03426438bf669287ae0865fa7", + "src/arrow/async_reader/store.rs": "ad0aac3a23ce066ed719f36df91b2fb8686f2ab9af6dc7679a5f378a2ffe3180", + "src/arrow/async_writer/mod.rs": "d9f70b03bf5f69f670c75755c6b9044849199d0a61c07b2832f99dea6766aa4e", + "src/arrow/async_writer/store.rs": "dedee9daff3005e2787c1f67df28b1e582214b0cba0e25b0d626a68453e4a186", + "src/arrow/buffer/bit_util.rs": "2293b3619f65a26bf236dc03e498f3efd568f7d6d7657e4805ce69cf575e5677", + "src/arrow/buffer/dictionary_buffer.rs": "0329a6f46a3f8b02a5c439e6d5933ac40160b7824351cc910ad1d1f4fcbebc76", + "src/arrow/buffer/mod.rs": "69e08e0b3347124397dcccd7fc53044629ea93868518aceb667e94040548349c", + "src/arrow/buffer/offset_buffer.rs": "69fdb496298a96abaa95e4bba9c820b31eaf90c7885bb8783f9bed49caebe2b7", + "src/arrow/buffer/view_buffer.rs": "e954bb9639eb16e0b61e460dc42ce30d6a424272784c3c3b298d3de310dd7d90", + "src/arrow/decoder/delta_byte_array.rs": "4a227b591c6f0cf68f60d77f91ac9cf888ea926530421278be42338dbe507497", + "src/arrow/decoder/dictionary_index.rs": "8a2c26327249a4316eec82438c65da00ea3610b8209565eb0cbdf72bbf73fb84", + "src/arrow/decoder/mod.rs": "7d1425a2a0818b7443dc89e0c3a749c4c6e1ab4ff1c4b1141c57deb1604ca29c", + "src/arrow/in_memory_row_group.rs": "c77cf07fbfd6613a2d4c8e23fa2a951d8c91d89b6fc23c792d51f9dc60292c38", + "src/arrow/mod.rs": "f59e61f3c602da05378ebd9116b2d531df8354dbfb0d8639118e70e489280ba2", + "src/arrow/push_decoder/mod.rs": "68aff720ac164e5a496bcea3e6c9ddd309da20cd4761f9738308d690a28769cb", + "src/arrow/push_decoder/reader_builder/data.rs": "3a06dc58a586800e98216dce3523538cc0262338dbcfc013e8e47c778d841380", + "src/arrow/push_decoder/reader_builder/filter.rs": "1fd54d79e140b0d05842caa28927dd534b3781fa4428ee5aa3f553e6c25a8c95", + "src/arrow/push_decoder/reader_builder/mod.rs": "62c3e8add452b4d32ebac71fe72664fe8646ad8d44505966421b26b8154ed2b9", + "src/arrow/push_decoder/remaining.rs": "5e36c6a822c9a9af87b14f0b3136e1151e6b9b5bf1e9d0066c96a830e76bbb23", + "src/arrow/record_reader/buffer.rs": "ce0889931b4a8db732c488560cd8c5391be4508bfc4a08f59be495a17aed03ea", + "src/arrow/record_reader/definition_levels.rs": "c3a37d45fc58896fd5b73250c734f9da35132df062020bd817374b2d0c3f6207", + "src/arrow/record_reader/mod.rs": "75a20fd49f7e2a2a064ffc50c0e1ad3cbd04c07a2d5e3a885b4faf990f7dc550", + "src/arrow/schema/complex.rs": "16bcccb27fdabe3982d37adbbf1e7a3ba33bdb152ceb62aa75e5f5453014f51d", + "src/arrow/schema/extension.rs": "ac9b4f1ff73212fe79801efb7fb12ae2c9ddad52183e09b889378d36e13aab84", + "src/arrow/schema/mod.rs": "8882135cffa3fac4fa0b81ddf6a125ed4db047f473885df0b9061eda693aad6b", + "src/arrow/schema/primitive.rs": "9ebf8d313c3df5720e8571ccaa1d75c19190405228cb52cd5536d1d1c95dae6d", + "src/arrow/schema/virtual_type.rs": "d97d7234a3bd1661b4cd7f5e849b56b505f4c7fef2715762d02a6081595fad22", + "src/basic.rs": "dac039b8b90a6693770f62d87dae7807042742dc02343ea379d73af7703dc88f", + "src/bin/cli/mod.rs": "1c74892a7e19249ecded4157ad008d2aacf91413564a86b616087bb59670fbe4", + "src/bin/parquet-concat.rs": "078be5a34dcf391e10d3c1ffcd9b16ddef0d9b79788849310cfe2287b31b3a56", + "src/bin/parquet-fromcsv-help.txt": "8a96bb57a021849802e61b0b9d07ff459acd7f7e0320330de4e77ae09a8a2e41", + "src/bin/parquet-fromcsv.rs": "f96cdea1adfbaec719efbbe42f4f99f225ad7b6ee66e709ab5eb04514229f6ae", + "src/bin/parquet-index.rs": "b9c537c7aad86bb0b800b9e878fad1420f912d1f7f8199701abfd7e56cbae119", + "src/bin/parquet-layout.rs": "d4d0ce4e78cef1126b10c0a23694a28c7478c3da69243570cb5bc156a3094438", + "src/bin/parquet-read.rs": "a409b3b62675070d43bcddce1c1455bb417f5281a9eebb6dfe148c6f92d22ba9", + "src/bin/parquet-rewrite.rs": "93414b3e32374413dcfa8b9758097c05fb82d596d37c41f644d03f48955169e7", + "src/bin/parquet-rowcount.rs": "3940ee5a39790788dd14a089c3f56528c12ab509c7dd050f2078d58741c0209e", + "src/bin/parquet-schema.rs": "e2e679ead30cdf100c376fb462d109b6e8b0f18d57175cceeacc9eb5b1739e63", + "src/bin/parquet-show-bloom-filter.rs": "43dda4b1cc1018fac4ee26ed562777010abb1cda6dddf88eebbb7c749a6beb08", + "src/bloom_filter/mod.rs": "1205117ba3fc0a702ace9a5c2361c79efc362ca035cde4b089ebabdd0b5214a0", + "src/column/chunker/cdc.rs": "f5c758ce3f8cc021c7bf0e8e170c51f5370ea5a4a4a4e373d2ae052e55684ee9", + "src/column/chunker/cdc_codegen.py": "a14f5291d9f96d893715ed80fda18f5cd50a6574c4b61cb1990d103cc7ac6192", + "src/column/chunker/cdc_generated.rs": "273529ca0d763d22ac8586ce23dd10941f9fef9d5e2f5f2e452f0131d363ebdf", + "src/column/chunker/mod.rs": "b071503ba1208bb8dc3cd0005b31d0ca198f60bae84ffcf98d6bcf7228a86560", + "src/column/mod.rs": "3d05cd64008d9c45cb8223d405368dd48df06e929ff93e623fb4e5bf3fabe00f", + "src/column/page.rs": "328539b6e2ff1c1aa66cd77d6c8af80393181fca1bf33221d4d6db9e954ba09f", + "src/column/page_encryption.rs": "34df686273e31872bbf5f33398a58d632ee378bde4706564fb58ec5e940dfdac", + "src/column/page_encryption_disabled.rs": "71a3f095939dde8451949f3c390a6862f0224db00403759fa80d47981bf14bc8", + "src/column/page_store.rs": "015811d51f70778c59b43aa0515ed9de7794191baa3d660bd91c01300f95f718", + "src/column/reader.rs": "f088a2a70a0adbf1dbd97f29fe401fd049b402c05ff6d6fc07f479c77e000972", + "src/column/reader/decoder.rs": "94483f56915f4a6f215fd02f315fe14e3285206730e2c85ef1ea56cf44a0a1a2", + "src/column/writer/byte_budget_chunker.rs": "8e1d215d2ae074e5258586813adb9ebc56a595db33f6188d32f435d0ec9ef8d0", + "src/column/writer/encoder.rs": "9bf74219546084ee56e8b80389b6339f59c82ec489f07976339823e53c142dc2", + "src/column/writer/mod.rs": "0e8e71d1919b739ee0a1efad93610fb67012b6572192d738bd4de954f4c3894e", + "src/compression.rs": "9a22b05f5faa1f2c6c9152a59d2384d1929b655b4f9543a3fffe36be7d48322c", + "src/data_type.rs": "656c4f97115e4bfbdcdc074a206c624006a46d05e5432a6e95aae0fe1e27fb46", + "src/encodings/decoding.rs": "40dd12aa522feecc910ea8132a84b745105190c044b6814c320146a6b4a5c735", + "src/encodings/decoding/byte_stream_split_decoder.rs": "017114b443e406d29fc66be8fcbc280295a3486802170820c548ed9b3ff50c58", + "src/encodings/encoding/byte_stream_split_encoder.rs": "8ea48728f12212c6fe75f705661e1deb7566a9f9e4cce0f021a411ddd2b635fa", + "src/encodings/encoding/dict_encoder.rs": "5e3ee6b230d6251370811e06e4897a370f38c085768fa819aaf8868edb8eadf6", + "src/encodings/encoding/mod.rs": "4f809107a677d0bcdca4ae514ef97f2869cb6f509da914e033892a2c7abcb84b", + "src/encodings/levels.rs": "7867d8873e8f537ae6c20f80f4230c16da8a9702c9ff7fafe5e5e954bcbfd483", + "src/encodings/mod.rs": "14df7b4bc46628bd874604ce22695832b012e2b59f4ea2912f590c169386821e", + "src/encodings/rle.rs": "4c62e58df99ec5daec7517a9e87ef870bf11d7b8b37cd453963875905fbee532", + "src/encryption/ciphers.rs": "fcc4c962795f7e1992d6838f1cd8163b7fb7d98a579f2cb5a0c5cf6f39106fcd", + "src/encryption/decrypt.rs": "e34fda5c9d3d8a25fad166077be6fc0e40aa382b41fde3b78e5ecae7151c0a5f", + "src/encryption/encrypt.rs": "11e6676d3c39813bd9f43bcebc3008170be6dfffb112a6d7c5033f07bd4c921d", + "src/encryption/mod.rs": "13822f0399262d5082eccaf5d16c2fca76e536bb28ac3077dd06d7ad3bc094e7", + "src/encryption/modules.rs": "7c0bee268abf8efb090c4530e9461cea629b60876877b7c96b050968a72b0d31", + "src/errors.rs": "f2644b92cdad56c7dc292ce8e17bab299881e0d63f98047f34c6acf6e4d50271", + "src/file/column_crypto_metadata.rs": "b986c78dcdceafbeac0af08ab9ccd53618aefd4192477391e807a6e9fe4cddb6", + "src/file/metadata/footer_tail.rs": "5dbb3ca3c621fb8bdce136c76c0ea98607e75d0969bb9277f696f55cd5722839", + "src/file/metadata/memory.rs": "efa6095020512a2539b604c78da0b01918f7b81459fd3b7179438e010c052457", + "src/file/metadata/mod.rs": "34f64fe9fa851edf14745d4ebdb126d00c29ad4dccb10dec90559e6bf943b84d", + "src/file/metadata/options.rs": "0e1d6d4d9f1df6d0eb78170947b2582cf53863b4c412257e800817ee0512b245", + "src/file/metadata/parser.rs": "bbe3eea0e27ab0e599faf62d151014ff25d67db12849d0f4b5f3facf81a73985", + "src/file/metadata/push_decoder.rs": "2989888905d872b0353029759e55b5800672bb8618b2c771f5d4585445ae1b85", + "src/file/metadata/reader.rs": "86d8ee33cb2744a6cd54e64200532a0012205da617668029f87db7cca8f853b7", + "src/file/metadata/thrift/encryption.rs": "7683f56fee986071721bf0cf76da55c8bcb1c1dfb91e37648e94146b1cf9c2d0", + "src/file/metadata/thrift/mod.rs": "bc1b818cc8174c4c5cf15ba0f7005d045b55622ac40b773db719ca6f5e09d4a0", + "src/file/metadata/writer.rs": "1031e71336d642d4711d1ef47bba417f238b975d1647df3eab51338055c07c72", + "src/file/mod.rs": "8ebca6e7555f4437bfdd45c97075fea6d9f4a18c32119bae6e84f155d506f5ac", + "src/file/page_index/column_index.rs": "836c61811abb3610a5b362aa397f162d15a7326ac05f28d8afa78f4115a169d0", + "src/file/page_index/index_reader.rs": "f438b2d254d464608c118b78800360ceea6ede79f443e015c0ccba55d45cf014", + "src/file/page_index/mod.rs": "87717035d2fbd0c8a338994d592e28f9e8e9fa502b4d3369b332ab3156ef3374", + "src/file/page_index/offset_index.rs": "262c11c7a13fda8e30341132872dfae3a30fe191c20ab2608bf0ccd3d9f33364", + "src/file/properties.rs": "658a755243aa7cedf297765c31e721c3a15058cae6e012448fa51030b45eaf7a", + "src/file/reader.rs": "5b921cb5c800ae6158200ffaafeb7e97f52528a606385ce8772c9cd4dd68836b", + "src/file/serialized_reader.rs": "ec64f581c0e20353e7f2854fa50f1a688097ad4e4ff0ab82627295f6d802244f", + "src/file/statistics.rs": "9e22a3c8e10cc3058b41b3262dfbfa2213899e09f52faec13991429957dd2c7b", + "src/file/writer.rs": "6b89a5f8128c98344ad150ac38947ec6851f835fed8695c568f7edbe673f7bad", + "src/geospatial/accumulator.rs": "3abcb6ac5d7beacd888d88254f912ad1aa31dff0c88a399979470f0e5784ea66", + "src/geospatial/bounding_box.rs": "971e8be47f63e53d58581aec3da58ae2b50ff557a02f9467005258bff96e97b5", + "src/geospatial/mod.rs": "e19b53bc377eab2efb21f02e9980056abf6b4a830fe364af025b4fd9a92a31d8", + "src/geospatial/statistics.rs": "824a36f14b7a07fb180208f10f2470fdeae498ad187c1a05fbdafe5538064585", + "src/lib.rs": "a0aed6abdef5dad36e2c1d39047ee23f3b78b1aa99e78fb5b95680655b892c00", + "src/parquet_macros.rs": "a3df3e4bb35f839285ba885dddceef24652bd9af8ccf29d7eec216e0f8a374fd", + "src/parquet_thrift.rs": "efdde5da47d94685096ae9c197bc657d15269cdd53f53ea212aff2cd4122075c", + "src/record/api.rs": "4fcb1670b8543d82dc9ceb0ad804031a453463669210da0f81b3a0b6ef838ad9", + "src/record/mod.rs": "b4f026402e005579250d97cb6238999e1ac22b87e2c075a3eaf311f169b35cdc", + "src/record/reader.rs": "86fd5cadd58f9b74bf1b669b17a7abf7454fd58c16d3c2cc2be1f845a8c4bfda", + "src/record/record_reader.rs": "cffc96f01028676ada1a555debafe4ba5b15d47253b7b1583b2455ee2b3db932", + "src/record/record_writer.rs": "41efd6c6aedeee01a3e8097c9216fcc337510ba608f3e6318658f3b916258c20", + "src/record/triplet.rs": "1ecc94fb17318aa998173e9e2d38876c4211990aefd8a2d29b97aad8ba53de00", + "src/schema/mod.rs": "a365a4b256a32ad1779e23f04252803c9da2f6579b77dc38d38c43fe7ec99fd4", + "src/schema/parser.rs": "fcce3efcc9db42a6f619235158dae4151e7377082ecdee3244a13631470ed0a3", + "src/schema/printer.rs": "245cfeffb454d273ece46f545535bd46541954db40eddf0a0f1d8aa5c0961847", + "src/schema/types.rs": "2ad50a0f758b46245448762ffe6be8af6f682b5ae30ff4ff4a6c1f2567d6291b", + "src/schema/visitor.rs": "452ff71b92d6fd42737c4da8004d0f98a6ecd805b4bf44633a19f0d3a7d31812", + "src/util/bit_pack.rs": "ee166cc8e9c93c16ee1e54f4c5ef60f387948d5fea248b61c1bf45c953a5d5af", + "src/util/bit_util.rs": "ef93cc39a384353f70a311efc916f86d94c35df5daed24b6e2507409b938ca56", + "src/util/interner.rs": "56ae30a7571fe47f36f077fd8df8b9eb0b654efa01e75ea66bd8238a37f40f79", + "src/util/mod.rs": "129f0cc534ba8ccb25f962fa068b197838d5b96258f0d27c1f057ccb9885b7e9", + "src/util/push_buffers.rs": "ddf81c09a6e54e849a3878242e39ac4d77c8c76dab21eff1613feb9be476f6d5", + "src/util/test_common/file_util.rs": "f44d07d280c9e473da2780d4d0860fc45d2e72e48afcf828eaed13efb3c17a0f", + "src/util/test_common/mod.rs": "e1193f2ddb43d626d01129a5fb11ef550d93bc129e1652a9df920a91bcada13d", + "src/util/test_common/page_util.rs": "5c31212e6c6d51075b8e921302b3f68a735c359ee8b11699283669914999b603", + "src/util/test_common/rand_gen.rs": "7fe918eceb3f21a341d483cfa8c54186e9e7536020ee01e552874be7dc6e034f", + "src/util/utf8.rs": "bfa5b876a43bc5d124ece8153a18ef08da244e1621961ba8e0fc03147c881ec0", + "src/variant.rs": "6494ee5e763ada61a3fb912c968b4798ca9e246a1c2430315060d58656d65b09", + "tests/arrow_reader/bad_data.rs": "074b8080bf53504490db79fb53bd87c49d546bbc6dabe85c50803b5fcf425e1c", + "tests/arrow_reader/bad_raw_metadata.bin": "6b5ea4cf5c2c2098f59d658c65188089ab0f4db4434d3ad1176e9ec0f93b092c", + "tests/arrow_reader/bigdelta.parquet": "5c2d5f5125617fe57a8d65269e66bea5fc5680374f9f80fdf93aa48821ff0046", + "tests/arrow_reader/bloom_filter/async.rs": "69360b25596c6f17e1000b778c8fc0c39473fc87fc192d80eb476fa616b4ea5a", + "tests/arrow_reader/bloom_filter/mod.rs": "d8828e6555f080b2535459e17602db3dd28f1cf8421aec316f765a6d9b2d1bc9", + "tests/arrow_reader/bloom_filter/sync.rs": "dd82259a72a3081cc93515a79a7ea13c70ec42d697c293865a068b99e0f82ea7", + "tests/arrow_reader/checksum.rs": "93a87027886fd1460dbbcf349ac32fd82efca2f3bd1c961a9e7038d7555a6044", + "tests/arrow_reader/int96_stats_roundtrip.rs": "46c4ca1fe3fcc9e6f37c0e642c34ef3da81b0b14d9423d7617e9462c36d3696b", + "tests/arrow_reader/invalid_utf8.rs": "2c7f0e167b5df8eb6ef1801c2e6413cf0f0cf3725e08510cd201dc43d0895213", + "tests/arrow_reader/io/async_reader.rs": "28692a1e9b1676284cb48ad5022bbd8c1e6c763f62e1fd9d0da1ee4ab866329c", + "tests/arrow_reader/io/mod.rs": "0f8435cdde1e8fc064bb5fc78afb378b5f9f69b68aaecd0e90ba2de8f8d96b55", + "tests/arrow_reader/io/sync_reader.rs": "11c5f3b4a01846a252e3edc4b657e7f00b6a351d63506c19cc33d432c379af6a", + "tests/arrow_reader/large_string_overflow.rs": "a6a4dbecbd49a9c9c9a9056c588044d3b6fcebd48dbcbbaa1de59dcf7b01b8a9", + "tests/arrow_reader/mod.rs": "9fddc48e6cb389819b2b890be0281d403ae234c9d195862d187d2dd8894f7e90", + "tests/arrow_reader/new_types.bin": "1b99e7c5fb79cc7cf744306ac81bdea617e6071686eb0afe7935fc28b433f638", + "tests/arrow_reader/predicate_cache.rs": "56b18ebd6369b62c0fe9f08d587f361d4b1edcc48c863a0b4854c2b17c7d77ca", + "tests/arrow_reader/row_filter/async.rs": "4f9c7c28a2ab65ee8f8ccc23e5ed1b91f275f5ade00ce2ff674552534860d7d3", + "tests/arrow_reader/row_filter/mod.rs": "d8828e6555f080b2535459e17602db3dd28f1cf8421aec316f765a6d9b2d1bc9", + "tests/arrow_reader/row_filter/sync.rs": "40c2546686a59de3bcd10fb633753eed307bcdcd3854d592ac637e0b32fe98c3", + "tests/arrow_reader/statistics.rs": "d17e32df1314945a3ba0e31568e7f4995aca0b066b1c44b266aab858ca889ea2", + "tests/arrow_writer.rs": "d9f36bc3d40fca91c7f365def4096e28b6fde5cf837c4bfce10a365ca7ad24ba", + "tests/arrow_writer_layout.rs": "6a750c563c58430d2c09b44c337c1949ed8ea512a18dc344d807e8926a2ae580", + "tests/encryption/encryption.rs": "1671702fd51181bbd61e6250dfb3bbf7b34812ddf57d71bfa67fb044288cba7f", + "tests/encryption/encryption_agnostic.rs": "6c2938dad5aee6c1a7ce034f04268f389af89d885f8927faee31be0544552f42", + "tests/encryption/encryption_async.rs": "371033e2736723e887fa4b1ef5c9851244fe32f537d71894ccaf3d0f83bd87d9", + "tests/encryption/encryption_disabled.rs": "c684f55f3c5d779d862a3fe1ea0fa811acfbee98bd8e4ba5dd9d276590468757", + "tests/encryption/encryption_util.rs": "cecae13fe258bf4b8907dac1c949b9e45249d01ebb37bdced7cb6b3fb2844b19", + "tests/encryption/mod.rs": "6112316d394e4288f8f8b80572f00c21f911191c331ca10e77079b513ce5f7b6", + "tests/geospatial.rs": "a33f1c8d49ae33684ad7b2c461df8ce82ff919547eafde93f46fa824c0e2f0aa", + "tests/variant_integration.rs": "a99ea247998343659311d76b806f77a5ab65bafcf689d0323cab5b39b0da0610" + }, + "version": "59.3.0" +} diff --git a/workflows/native-gates.toml b/workflows/native-gates.toml index 1500b00..6d5ddf3 100644 --- a/workflows/native-gates.toml +++ b/workflows/native-gates.toml @@ -1,9 +1,9 @@ command_seconds = 240 -environment = ["PATH", "RUSTUP_TOOLCHAIN", "RUSTFLAGS", "CARGO_TARGET_DIR", "RUSTC_WRAPPER"] +environment = ["PATH", "RUSTUP_TOOLCHAIN", "RUSTFLAGS", "CARGO_TARGET_DIR", "RUSTC_WRAPPER", "CARGO_PROFILE_DEV_DEBUG", "CARGO_PROFILE_TEST_DEBUG", "CARGO_BUILD_JOBS"] [[gates]] name = "toolkit format" -argv = ["cargo", "fmt", "--all", "--", "--check"] +argv = ["cargo", "fmt", "--package", "allstack", "--", "--check"] [[gates]] name = "native skills and runtime"