diff --git a/.github/workflows/TagBot.yml b/.github/workflows/TagBot.yml index ad381d1f..99954fd9 100644 --- a/.github/workflows/TagBot.yml +++ b/.github/workflows/TagBot.yml @@ -25,9 +25,22 @@ jobs: TagBot: if: github.event_name == 'workflow_dispatch' || github.actor == 'JuliaTagBot' runs-on: ubuntu-latest - timeout-minutes: 5 + timeout-minutes: 15 steps: - - uses: JuliaRegistries/TagBot@v1 + - name: Tag Arrow.jl + uses: JuliaRegistries/TagBot@v1 with: token: ${{ secrets.GITHUB_TOKEN }} ssh: ${{ secrets.DOCUMENTER_KEY }} + - name: Tag ArrowTypes.jl + uses: JuliaRegistries/TagBot@v1 + with: + token: ${{ secrets.GITHUB_TOKEN }} + ssh: ${{ secrets.DOCUMENTER_KEY }} + subdir: src/ArrowTypes + - name: Tag ArrowStrings.jl + uses: JuliaRegistries/TagBot@v1 + with: + token: ${{ secrets.GITHUB_TOKEN }} + ssh: ${{ secrets.DOCUMENTER_KEY }} + subdir: src/ArrowStrings diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 35bb2352..6fee3d42 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -61,7 +61,9 @@ jobs: run: | git config user.name "github-actions[bot]" git config user.email "github-actions[bot]@users.noreply.github.com" - RELEASE_DEFAULT=0 dev/release/release_rc.sh 100 + RELEASE_DEFAULT=0 dev/release/release_rc.sh 100 > release-rc-output.txt + cat release-rc-output.txt + grep -F "The vote will be open for at least 72 hours." release-rc-output.txt - uses: actions/cache@v6 with: path: ~/.julia/artifacts @@ -75,20 +77,15 @@ jobs: -e 's/"$//g') VERIFY_DEFAULT=0 dev/release/verify_rc.sh ${version} 100 test: - name: ${{ matrix.pkg.name }} - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.nthreads }} threads - ${{ github.event_name }} + name: Arrow.jl - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.nthreads }} threads - ${{ github.event_name }} runs-on: ${{ matrix.os }} timeout-minutes: 45 strategy: fail-fast: false matrix: - pkg: - - name: Arrow.jl - dir: '.' - - name: ArrowTypes.jl - dir: './src/ArrowTypes' version: - 'min' - - 'lts' + - '1.11' - '1' # automatically expands to the latest stable 1.x release of Julia - 'pre' # expands to latest alpha, beta or RC, if available, otherwise same as `1` os: @@ -102,6 +99,7 @@ jobs: - uses: julia-actions/setup-julia@v3 with: version: ${{ matrix.version }} + project: '.' - uses: actions/cache@v6 env: cache-name: cache-artifacts @@ -112,21 +110,64 @@ jobs: ${{ runner.os }}-test-${{ env.cache-name }}- ${{ runner.os }}-test- ${{ runner.os }}- - - uses: julia-actions/julia-buildpkg@v1.6 - with: - project: ${{ matrix.pkg.dir }} - - name: Dev local ArrowTypes for Arrow.jl tests - if: matrix.pkg.name == 'Arrow.jl' + # Before the build: ArrowStrings is unregistered and ArrowTypes 2.4 is in + # this monorepo; Julia 1.10's Pkg does not read the `[sources]` entries + # that point at them, so the resolve would fail without this step. + - name: Dev local subpackages shell: julia --project=. {0} run: | using Pkg - Pkg.develop(PackageSpec(path="src/ArrowTypes")) + Pkg.develop([PackageSpec(path="src/ArrowStrings"), PackageSpec(path="src/ArrowTypes")]) + - uses: julia-actions/julia-buildpkg@v1.6 - uses: julia-actions/julia-runtest@v1 env: JULIA_NUM_THREADS: ${{ matrix.nthreads }} + - uses: julia-actions/julia-processcoverage@v1 + with: + directories: src,ext + - uses: codecov/codecov-action@v7 + with: + files: lcov.info + # The subpackages' own suites: one job each, on the package's minimum + # supported Julia. The full Arrow.jl matrix above exercises both packages' + # code on every version, OS, and thread count. + test-subpackages: + name: ${{ matrix.pkg.name }} - Julia min - ${{ github.event_name }} + runs-on: ubuntu-latest + timeout-minutes: 20 + strategy: + fail-fast: false + matrix: + pkg: + - name: ArrowTypes.jl + dir: './src/ArrowTypes' + - name: ArrowStrings.jl + dir: './src/ArrowStrings' + steps: + - uses: actions/checkout@v7 + - uses: julia-actions/setup-julia@v3 + with: + version: 'min' + project: ${{ matrix.pkg.dir }} + - uses: actions/cache@v6 + env: + cache-name: cache-artifacts + with: + path: ~/.julia/artifacts + key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }} + restore-keys: | + ${{ runner.os }}-test-${{ env.cache-name }}- + ${{ runner.os }}-test- + ${{ runner.os }}- + - uses: julia-actions/julia-buildpkg@v1.6 + with: + project: ${{ matrix.pkg.dir }} + - uses: julia-actions/julia-runtest@v1 with: project: ${{ matrix.pkg.dir }} - uses: julia-actions/julia-processcoverage@v1 + with: + directories: ${{ matrix.pkg.dir }}/src - uses: codecov/codecov-action@v7 with: files: lcov.info @@ -163,19 +204,106 @@ jobs: shell: julia --project=monorepo {0} run: | using Pkg; - Pkg.develop([PackageSpec(path="."), PackageSpec(path="src/ArrowTypes")]) + Pkg.develop([PackageSpec(path="."), PackageSpec(path="src/ArrowTypes"), PackageSpec(path="src/ArrowStrings")]) - name: Run monorepo tests continue-on-error: false run: > julia --color=yes --project=monorepo -e 'using Pkg; Pkg.test("Arrow")' + test_min_arrowtypes: + name: Arrow.jl - ArrowTypes 2.0 compatibility + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@v7 + - uses: julia-actions/setup-julia@v3 + with: + version: '1.10' + project: '.' + - uses: actions/cache@v6 + with: + path: ~/.julia/artifacts + key: min-arrowtypes-${{ runner.os }}-${{ hashFiles('Project.toml') }} + restore-keys: min-arrowtypes-${{ runner.os }}- + - name: Resolve the minimum supported ArrowTypes version + shell: julia --project=. {0} + run: | + using Pkg + Pkg.develop(PackageSpec(path="src/ArrowStrings")) + Pkg.add(PackageSpec(name="ArrowTypes", version=v"2.0.0")) + Pkg.pin(PackageSpec(name="ArrowTypes")) + - name: Test custom-type file and stream round trips + run: julia --startup-file=no --history-file=no --project=. test/min_arrowtypes_compat.jl + test_windows: + name: ${{ matrix.pkg.name }} - Julia latest - Windows - 2 threads + runs-on: windows-latest + timeout-minutes: 45 + strategy: + fail-fast: false + matrix: + pkg: + - name: Arrow.jl + dir: '.' + - name: ArrowTypes.jl + dir: './src/ArrowTypes' + - name: ArrowStrings.jl + dir: './src/ArrowStrings' + steps: + - uses: actions/checkout@v7 + - uses: julia-actions/setup-julia@v3 + with: + version: '1' + project: ${{ matrix.pkg.dir }} + - uses: actions/cache@v6 + with: + path: ~/.julia/artifacts + key: windows-test-${{ hashFiles('**/Project.toml') }} + restore-keys: windows-test- + - name: Prepare Arrow.jl development dependencies + if: matrix.pkg.name == 'Arrow.jl' + shell: julia --project=. {0} + run: | + using Pkg + Pkg.develop([PackageSpec(path="src/ArrowStrings"), PackageSpec(path="src/ArrowTypes")]) + - uses: julia-actions/julia-buildpkg@v1.6 + with: + project: ${{ matrix.pkg.dir }} + - uses: julia-actions/julia-runtest@v1 + env: + JULIA_NUM_THREADS: 2 + with: + project: ${{ matrix.pkg.dir }} + trim: + name: ArrowCore and ArrowStrings JuliaC trim-safe compile + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@v7 + - uses: julia-actions/setup-julia@v3 + with: + version: '1.12.6' + - uses: actions/cache@v6 + with: + path: ~/.julia/artifacts + key: trim-${{ runner.os }}-${{ hashFiles('test/trim/Project.toml') }} + restore-keys: trim-${{ runner.os }}- + - name: Compile and run the trim-safe ArrowCore and ArrowStrings workload + run: julia --startup-file=no --history-file=no test/trim_compile_tests.jl docs: name: Documentation runs-on: ubuntu-latest timeout-minutes: 5 steps: - uses: actions/checkout@v7 - - uses: julia-actions/julia-buildpkg@latest - - uses: julia-actions/julia-docdeploy@latest + - uses: julia-actions/setup-julia@v3 + with: + version: '1' + - name: Prepare Arrow.jl development dependencies for the docs + shell: julia --project=docs {0} + run: | + using Pkg + Pkg.develop([PackageSpec(path="."), PackageSpec(path="src/ArrowStrings"), PackageSpec(path="src/ArrowTypes")]) + - uses: julia-actions/julia-buildpkg@v1.7.0 + - uses: julia-actions/julia-docdeploy@v1.3.1 env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }} @@ -186,11 +314,9 @@ jobs: - uses: julia-actions/setup-julia@v3 - uses: actions/checkout@v7 - name: Install JuliaFormatter and format - # This will use the latest version by default but you can set the version like so: - # julia -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter", version="0.13.0"))' run: | - julia -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter"))' - julia -e 'using JuliaFormatter; format(".", verbose=true)' + julia -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter", version="2.12.4"))' + julia -e 'using JuliaFormatter; format(".", verbose=true)' - name: Format check run: | julia -e ' diff --git a/.github/workflows/ci_nightly.yml b/.github/workflows/ci_nightly.yml index 53b2432b..cac807de 100644 --- a/.github/workflows/ci_nightly.yml +++ b/.github/workflows/ci_nightly.yml @@ -33,6 +33,8 @@ jobs: dir: '.' - name: ArrowTypes.jl dir: './src/ArrowTypes' + - name: ArrowStrings.jl + dir: './src/ArrowStrings' version: - 'nightly' os: @@ -64,7 +66,16 @@ jobs: JULIA_NUM_THREADS: ${{ matrix.nthreads }} with: project: ${{ matrix.pkg.dir }} - - uses: julia-actions/julia-processcoverage@v1 + - name: Process Arrow.jl coverage + if: matrix.pkg.name == 'Arrow.jl' + uses: julia-actions/julia-processcoverage@v1 + with: + directories: src,ext + - name: Process subpackage coverage + if: matrix.pkg.name != 'Arrow.jl' + uses: julia-actions/julia-processcoverage@v1 + with: + directories: ${{ matrix.pkg.dir }}/src - uses: codecov/codecov-action@v7 with: files: lcov.info @@ -101,7 +112,7 @@ jobs: shell: julia --project=monorepo {0} run: | using Pkg; - Pkg.develop([PackageSpec(path="."), PackageSpec(path="src/ArrowTypes")]) + Pkg.develop([PackageSpec(path="."), PackageSpec(path="src/ArrowTypes"), PackageSpec(path="src/ArrowStrings")]) - name: Run monorepo tests continue-on-error: false run: > diff --git a/.github/workflows/conformance.yml b/.github/workflows/conformance.yml new file mode 100644 index 00000000..ca867627 --- /dev/null +++ b/.github/workflows/conformance.yml @@ -0,0 +1,54 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +name: Conformance +on: + workflow_dispatch: + schedule: + - cron: '30 1 * * 1' + pull_request: + paths: + - 'Project.toml' + - 'conformance/**' + - 'src/**' + - 'test/**' + - '.github/workflows/conformance.yml' + +permissions: + contents: read + +concurrency: + group: conformance-${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + conformance: + name: Arrow corpus and external oracles + runs-on: ubuntu-latest + timeout-minutes: 120 + steps: + - uses: actions/checkout@v7 + - uses: julia-actions/setup-julia@v3 + with: + version: '1.12.6' + - uses: actions/cache@v6 + with: + path: ~/.julia/artifacts + key: conformance-${{ runner.os }}-${{ hashFiles('conformance/host/Project.toml') }} + restore-keys: conformance-${{ runner.os }}- + - name: Run the Arrow corpus and PyArrow/Nanoarrow oracles + run: julia --startup-file=no --history-file=no conformance/run.jl --rebuild diff --git a/.github/workflows/fuzz.yml b/.github/workflows/fuzz.yml new file mode 100644 index 00000000..808d0825 --- /dev/null +++ b/.github/workflows/fuzz.yml @@ -0,0 +1,94 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +name: Deterministic fuzz + +on: + schedule: + - cron: '15 3 * * 3' + workflow_dispatch: + inputs: + seed: + description: Master UInt64 seed (decimal or 0x-prefixed hex) + required: false + default: '0x9f5a37c241de880b' + cases: + description: Differential cases + required: false + default: '512' + mutations: + description: Byte mutations + required: false + default: '20000' + +permissions: + contents: read + +concurrency: + group: fuzz-${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: false + +jobs: + fuzz: + name: Extended differential and mutation fuzz - stable Julia + runs-on: ubuntu-latest + timeout-minutes: 90 + env: + FUZZ_SEED: ${{ github.event_name == 'schedule' && github.run_id || inputs.seed || '0x9f5a37c241de880b' }} + FUZZ_CASES: ${{ inputs.cases || '512' }} + FUZZ_MUTATIONS: ${{ inputs.mutations || '20000' }} + FUZZ_REPRO_DIR: fuzz-reproductions + steps: + - uses: actions/checkout@v7 + - uses: julia-actions/setup-julia@v3 + with: + version: '1' + - uses: actions/cache@v6 + with: + path: ~/.julia/artifacts + key: fuzz-${{ runner.os }}-${{ hashFiles('**/Project.toml') }} + restore-keys: fuzz-${{ runner.os }}- + - name: Prepare development dependencies + shell: julia --project=. {0} + run: | + using Pkg + Pkg.develop([PackageSpec(path="src/ArrowStrings"), PackageSpec(path="src/ArrowTypes")]) + - name: Report replay coordinates + run: | + echo "event=$GITHUB_EVENT_NAME" + echo "master_seed=$FUZZ_SEED" + if [ "$GITHUB_EVENT_NAME" = "schedule" ] && [ "$FUZZ_SEED" != "$GITHUB_RUN_ID" ]; then + echo "scheduled fuzz seed must equal the workflow run ID" >&2 + exit 1 + fi + - name: Run deterministic differential and mutation fuzzing + run: >- + timeout --signal=TERM --kill-after=30s 80m + julia --color=yes --startup-file=no --history-file=no --project=. + test/fuzz.jl + --seed "$FUZZ_SEED" + --cases "$FUZZ_CASES" + --mutations "$FUZZ_MUTATIONS" + --determinism-every 256 + --repro-dir "$FUZZ_REPRO_DIR" + - name: Upload failure reproductions + if: failure() && hashFiles('fuzz-reproductions/**') != '' + uses: actions/upload-artifact@v6 + with: + name: deterministic-fuzz-reproductions + path: fuzz-reproductions + if-no-files-found: ignore + retention-days: 14 diff --git a/.gitignore b/.gitignore index 4c10e7e3..a722be11 100644 --- a/.gitignore +++ b/.gitignore @@ -17,10 +17,10 @@ Manifest.toml Manifest-v*.toml -.DS_STORE +.DS_Store *.jl.cov *.jl.*.cov *.jl.mem +docs/build/ test/_scrap.jl -.DS_STORE diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 00000000..1745beba --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,197 @@ + + +# Changelog + +This file records user-visible changes to Arrow.jl. See the +[3.0 migration guide](docs/src/migration.md) for update instructions. + +## 3.0.0 - Unreleased + +Arrow.jl 3.0 is a breaking major release. It replaces the internal reader, +writer, validation, scan, and C interface engines. + +### Breaking changes + +- Julia 1.10 is now the minimum supported Julia version. +- `Arrow.Table` materializes plain Julia vectors. It no longer returns lazy + `ArrowVector` views into the source. +- `Arrow.Table(...; convert=false)` and the related lazy conversion path were + removed. +- Positional byte-window arguments and the multi-input `Table` and `Stream` + constructors were removed. Each read accepts one complete source. +- `Arrow.write` is an eager, whole-buffer writer. `Arrow.Writer` (incremental + writes, both formats) and `Arrow.append` (IPC streams) are reimplemented on + the 3.0 core with one semantic change: the first table written fixes the + schema, and later tables must conform — no inference crosses tables. +- Writing to an `IO` now emits the IPC file format by default. Pass + `file=false` for the stream format. `Arrow.tobuffer` still emits the stream + format, matching its 2.x bytes. +- The `alignment`, `dictencode`, `dictencodenested`, `denseunions`, + `largelists`, `maxdepth`, and `ntasks` writer keywords no longer have any + effect. `Arrow.write` accepts them with a one-time warning each and ignores + them; wrap columns in `Arrow.DictEncode` to dictionary-encode. +- The `compress` keyword accepts `nothing`, `:lz4`, or `:zstd`; initialized + compressor objects are no longer accepted. +- The `Arrow.ToArrow` compatibility binding was removed. Import + `ArrowTypes.ToArrow` directly when an explicit lazy conversion view is + needed. Normal writes apply `ArrowTypes.ArrowType` and `ArrowTypes.toarrow` + automatically. +- The DataAPI.jl metadata interface replaces the Arrow 2.x metadata + accessors. `Arrow.getmetadata(table)` remains as a compatibility method + over it; the per-column `getmetadata(column)` form is gone because columns + are plain vectors — use `DataAPI.colmetadata` instead. +- The package now has a narrow export surface. Use names such as + `Arrow.Table`, `Arrow.Stream`, `Arrow.write`, and `Arrow.DictEncode` through + the `Arrow` namespace. Only `release!` and `ArrowTypes` are exported. + Packages that define custom mappings should still depend on and import + ArrowTypes.jl directly. +- Big-endian IPC and delta-dictionary messages are rejected. +- Arrow 3.0 requires ArrowTypes.jl 2.x, Tables.jl 1.14 (the first release + with `Tables.Scan`), and the ArrowStrings.jl 1.0 release. + +### Added + +- `Tables.Scan` projection, filter, limit, and offset pushdown. +- Typed `Tables.Scan` select overrides for composite rows: a + `:column => NamedTuple{...}` select item reads a Struct column as typed + rows (recursively, including `Vector{...}` targets for list columns). +- A TimeZones.jl extension. When TimeZones.jl is loaded, second- and + millisecond-unit timestamps that declare a timezone read as + `ZonedDateTime` (the Arrow 2.x behavior) and round-trip through rewrite, + and a fresh single-zone `ZonedDateTime` column writes as a + timezone-declared millisecond timestamp; without the extension those + columns read as naive UTC `DateTime` values. +- The Arrow 2.x compatibility surface: exported `ArrowTypes`, + `Arrow.getmetadata(table)`, `Arrow.tobuffer`, and the curried + `table |> Arrow.write(sink)` form. +- An incremental writer: `Arrow.Writer(sink; file=true)` publishes each + written table's batches immediately, holding only the current table in + memory, for both IPC formats. `Arrow.append(sink, table)` extends an + existing IPC stream in place. `Arrow.Writer(sink; dictreplacement=true)` + declares the DictionaryReplacement feature up front so a stream's pools + can be replaced by later writes or appends; declared features are demands + on readers, so nothing is declared speculatively. +- The IPC reader accepts vtables that understate a wider field's extent, + as the reference implementation does. Arrow 2.x's FlatBuffers builder + produced them for every dictionary-encoded field carrying metadata, so + Arrow 3.0 now reads every such 2.x file (including all CategoricalArrays + columns) that earlier 3.0 development builds refused. +- `Arrow.AbstractArrowSource` for sparse byte-range reads. +- A CloudStore.jl extension for remote object reads. +- Footer statistics that can prune record batches before their data is read. +- Arrow C data and C stream import and export. +- Arrow StringView and BinaryView support. ArrowStrings.jl provides a reusable + zero-copy StringView representation for Arrow.jl and compatible producers. +- Fresh Julia columns with a heterogeneous declared `Union` element type are + synthesized as canonical dense Arrow Union arrays. Each child uses the + recursive core or ArrowTypes.jl mapping supported at that nesting depth. +- Recursive ArrowTypes.jl custom-type lowering and extension-type restoration + for top-level values and values nested in lists, tuples and fixed-size lists, + structs, maps, dictionary-encoded values, and freshly synthesized + heterogeneous Unions. +- Structural, semantic, and optional full-content validation tiers. +- Resource limits for untrusted IPC metadata and buffers. +- Apache Arrow gold-corpus tests, external IPC oracle tests, C interface oracle + tests, deterministic differential and bounded mutation fuzzing, and a JuliaC + `--trim=safe` compile gate. + +### Changed + +- `Arrow.Stream` yields materialized `Arrow.Table` batches. +- A file-format path is memory-mapped by default. Call `Arrow.release!` to release + the map at a known time. Materialized table columns remain usable afterward. +- A table read from Arrow retains compatible schema details when it is written + again, including temporal units, dictionary encoding and category order, + list widths, composite descriptors, nullability, and ordered duplicate + metadata. Fresh heterogeneous Julia Union columns can be synthesized. + Registered ArrowTypes.jl public-domain values can also reconstruct a retained Union + from their writer-side type evidence, including dense or sparse mode and type + IDs. An unregistered retained Union still fails clearly after facade + materialization discards its original routing. + Nested Dictionaries fail clearly after their pool data is lost. +- Writing is validated before bytes are published to the output sink. +- Fresh `Union{Missing, NamedTuple}` columns use Struct parent validity while + preserving each child's declared nullability and type. +- Custom values are lowered recursively through `ArrowTypes.ArrowType` and + `ArrowTypes.toarrow`. Extension names and metadata are written, and reads use + `ArrowTypes.JuliaType`, `ArrowTypes.fromarrow`, and + `ArrowTypes.fromarrowstruct` to restore registered logical types. An unknown + extension name warns and returns its storage value. +- A fresh abstract column with no mapping of its own uses concrete subtype + evidence across the complete column. Concrete subtype extensions are kept, + and heterogeneous subtype evidence forms an explicit bounded Union instead + of silently erasing subtype metadata. +- Registered logical types whose storage is a Union preserve external child + order, child labels, type IDs, dense or sparse mode, nested descriptor details, + and outer-null routing on rewrite. Sparse children use canonical hidden + placeholders outside their active rows. An abstract registered target accepts + an extensionless concrete subtype or one with the retained parent identity; a + different explicit identity fails closed. A logical storage Union that already + uses `Missing` cannot also add an outer missing state because the two states + have no distinguishable Arrow representation. +- Hidden retained composite slots are built directly from their Field and + logical length. Null-only fixed-size-list descendants do not allocate one + Julia placeholder per hidden element, including partly missing fresh or + retained columns and inactive sparse-Union rows. Sparse Union children also + distinguish inactive physical slots from selected values, so inactive null + storage does not weaken selected-value validation. +- Recursive ArrowTypes.jl storage schemas, recursive value containers, and + custom mapping nesting beyond 64 levels fail with `ArgumentError`. The + removed `maxdepth` keyword does not make writer recursion unbounded. +- Declared writer Unions support up to 32 branches. Runtime writer or storage + inference accepts at most 8 distinct types across the complete column. This + covers abstract ArrowTypes storage, abstract or `Any` dictionary values, and + abstract retained ArrowTypes targets, and bounds per-type schema planning and + compiler work. +- ArrowTypes logical resolution uses exact fixed-size-list tuple signatures + through arity 1024 and a compact tuple-family signature above that limit, + preventing fixed-size-list sizes from causing proportional type allocation. +- Extension Struct signatures are exact through 1024 children only when child + names are unique, contain no embedded NUL, already exist as Julia `Symbol`s, + and pass the 4096-byte per-name and 64-KiB total limits. Otherwise labelled + Structs remain unknown extensions and return ordered `Pair` storage. +- One bounded ArrowTypes.jl Tuple compatibility exception may intern only a + complete canonical positional sequence `"1"`, `"2"`, …, `string(N)` through + `N = 1024`. Unknown extension labels return before that check; arbitrary or + partly positional Struct names are not interned. +- Any writer-side `ArrowType` result that is a concrete tuple with more than + 1024 fields is rejected before the writer specializes on the oversized + storage shape. This includes ArrowTypes.jl's default mapping for a tuple + value. +- Core keeps schema names as strings. The Tables.jl facade preflights per-name, + novel-name-count, and novel-name-byte limits before it interns top-level + column names. Unknown ArrowTypes extension labels do not create Julia + symbols. Unsupported-label warnings are deduplicated by label and capped at + 16 distinct labels plus one suppression notice per table materialization. A + novel `JuliaLang.Symbol` IPC payload now raises `ValidationError` instead of + being interned, so an input that an earlier Arrow.jl release read successfully + can now fail. +- Registered public-domain values can rebuild compatible retained binary, list, + date-like, duration, wide-decimal, and interval descriptors. Retained widths, + sizes, units, child fields, and sorted Map claims are checked before output. +- Empty and typed all-missing registered columns use their concrete declared + element types as schema evidence instead of bypassing retained-field checks. +- Scan filters over fields that contain registered ArrowTypes.jl logical types + at any depth evaluate over the public materialized values. The mapping + interface does not require storage lowering to preserve Julia comparison + semantics. +- Removed the JSON3/StructTypes-based arrow-JSON test integration and the + associated support claims. Arrow 2.x used those packages only in tests; no + public API existed. diff --git a/CONTEXT.md b/CONTEXT.md new file mode 100644 index 00000000..3bb11fd8 --- /dev/null +++ b/CONTEXT.md @@ -0,0 +1,54 @@ + + +# Arrow.jl + +Arrow.jl maps between the Arrow columnar format and Julia values. These terms name the +domains and verification paths used by the package. + +## Language + +**Storage domain**: +The Julia values that directly represent an Arrow physical layout before facade or +ArrowTypes conversion. +_Avoid_: Raw values, wire values + +**Public domain**: +The Julia values exposed by `Arrow.Table` and `Arrow.Stream` after facade and +ArrowTypes conversion. +_Avoid_: Converted values, logical values + +**Scan plan**: +A `Tables.Scan` request resolved against one Arrow schema, including its storage-domain +filter, selected output, row window, and remaining public-domain conversion. +_Avoid_: Lowered scan + +**Column construction**: +The mapping from a Julia column and optional retained Arrow field to one Arrow field and +its array data. +_Avoid_: Column conversion, column building + +**Acceptance battery**: +An assertion-dense executable that verifies one Arrow adapter through its supported +interface. +_Avoid_: Integration script, smoke test + +**Conformance oracle**: +An independent Arrow implementation or corpus used to judge interoperability. +_Avoid_: Golden test, reference test diff --git a/NOTICE b/NOTICE index 147b1901..fc6b69b3 100644 --- a/NOTICE +++ b/NOTICE @@ -1,5 +1,5 @@ Apache Arrow Julia -Copyright 2016-2025 The Apache Software Foundation +Copyright 2016-2026 The Apache Software Foundation This product includes software developed at The Apache Software Foundation (https://www.apache.org/). diff --git a/Project.toml b/Project.toml index e9fc73f0..b0314098 100644 --- a/Project.toml +++ b/Project.toml @@ -17,38 +17,50 @@ name = "Arrow" uuid = "69666777-d1a9-59fb-9406-91d4454c9d45" authors = ["quinnj "] -version = "2.8.1" +version = "3.0.0-DEV" [deps] +ArrowStrings = "c38d8858-22a2-449e-9eca-ef92ec15f353" ArrowTypes = "31f734f8-188a-4ce0-8406-c8a06bd891cd" -BitIntegers = "c3b6d118-76ef-56ca-8cc7-ebb389d030a1" +Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f" CodecLz4 = "5ba52731-8f18-5e0d-9241-30f10d1ec561" CodecZstd = "6b39b394-51ab-5f42-8807-6242bab2b4c2" -ConcurrentUtilities = "f0e56b4a-5159-44fe-b623-3e5288b988bb" DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56" Mmap = "a63ad114-7e13-5084-954f-fe012c677804" -PooledArrays = "2dfb63ee-cc39-5dd5-95bd-886bf059d720" -SentinelArrays = "91c51154-3ec4-41a3-a24f-3f23e20d615c" -StringViews = "354b36f9-a18e-4713-926e-db85100087ba" +PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" -TimeZones = "f269a46b-ccf7-5d73-abea-4c690281aa53" TranscodingStreams = "3bb67fe8-82b1-5028-8e26-92a6c54297fa" -UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4" + +# ArrowStrings and ArrowTypes live in this repository and are released and +# registered independently; the source entries resolve their subdirectories. +# Pkg reads `[sources]` from Julia 1.11; on 1.10 the CI workflow devs both explicitly. +[sources] +ArrowStrings = {path = "src/ArrowStrings"} +ArrowTypes = {path = "src/ArrowTypes"} + +[weakdeps] +CloudStore = "3365d9ee-d53b-4a56-812d-5344d5b716d7" +TimeZones = "f269a46b-ccf7-5d73-abea-4c690281aa53" + +[extensions] +ArrowCloudStoreExt = "CloudStore" +ArrowTimeZonesExt = "TimeZones" [compat] -ArrowTypes = "1.1,2" -BitIntegers = "0.2, 0.3" +ArrowStrings = "1" +ArrowTypes = "2" +Base64 = "1.10" +CloudStore = "1.6" CodecLz4 = "0.4" -CodecZstd = "0.7, 0.8" -ConcurrentUtilities = "2" DataAPI = "1" +Dates = "1.10" +CodecZstd = "0.8" EnumX = "1" -PooledArrays = "0.5, 1.0" -SentinelArrays = "1" -StringViews = "1, 2" -Tables = "1.1" +Mmap = "1.10" +PrecompileTools = "1" +Tables = "1.14" TimeZones = "1" TranscodingStreams = "0.9.12, 0.10, 0.11" -julia = "1.9" +julia = "1.10" diff --git a/README.md b/README.md index 98bc9fd9..29f0b31f 100644 --- a/README.md +++ b/README.md @@ -17,55 +17,99 @@ under the License. --> -# Arrow +# Arrow.jl -[![docs](https://img.shields.io/badge/docs-latest-blue&logo=julia)](https://arrow.apache.org/julia/) -[![CI](https://github.com/apache/arrow-julia/workflows/CI/badge.svg)](https://github.com/apache/arrow-julia/actions?query=workflow%3ACI) -[![codecov](https://app.codecov.io/gh/apache/arrow-julia/branch/main/graph/badge.svg)](https://app.codecov.io/gh/apache/arrow-julia) +[![Documentation](https://img.shields.io/badge/docs-latest-blue?logo=julia)](https://arrow.apache.org/julia/) +[![CI](https://github.com/apache/arrow-julia/actions/workflows/ci.yml/badge.svg)](https://github.com/apache/arrow-julia/actions/workflows/ci.yml) +[![Codecov](https://codecov.io/gh/apache/arrow-julia/branch/main/graph/badge.svg)](https://codecov.io/gh/apache/arrow-julia) -[![deps](https://juliahub.com/docs/Arrow/deps.svg)](https://juliahub.com/ui/Packages/Arrow/QnF3w?t=2) -[![version](https://juliahub.com/docs/Arrow/version.svg)](https://juliahub.com/ui/Packages/Arrow/QnF3w) -[![pkgeval](https://juliahub.com/docs/Arrow/pkgeval.svg)](https://juliahub.com/ui/Packages/Arrow/QnF3w) +Arrow.jl is a pure Julia implementation of the +[Apache Arrow](https://arrow.apache.org) columnar data standard. It reads and +writes Arrow IPC files and streams. It also supports the Arrow C data and C +stream interfaces, Tables.jl, compressed buffers, and selective byte-range +reads. -This is a pure Julia implementation of the [Apache Arrow](https://arrow.apache.org) data standard. This package provides Julia `AbstractVector` objects for -referencing data that conforms to the Arrow standard. This allows users to seamlessly interface Arrow formatted data with a great deal of existing Julia code. - -Please see this [document](https://arrow.apache.org/docs/format/Columnar.html#physical-memory-layout) for a description of the Arrow memory layout. +> [!IMPORTANT] +> This is the Arrow.jl 3.0 development branch. Arrow 3.0 is not registered +> yet, and it requires the first registered ArrowStrings.jl release. A +> checkout uses the in-repository `src/ArrowStrings` and `src/ArrowTypes` +> packages: Julia 1.11+ resolves them through `[sources]`; on Julia 1.10 run +> the `Pkg.develop` commands below. ## Installation -The package can be installed by typing in the following in a Julia REPL: +Install the latest registered release from the Julia REPL: ```julia -julia> using Pkg; Pkg.add("Arrow") +import Pkg +Pkg.add("Arrow") ``` -## Local Development +## Quick start + +```julia +using Arrow -When developing on Arrow.jl it is recommended that you run the following to ensure that any -changes to ArrowTypes.jl are immediately available to Arrow.jl without requiring a release: +data = (id = [1, 2, 3], name = ["Ada", "Babbage", missing]) +Arrow.write("data.arrow", data) -```sh -julia --project -e 'using Pkg; Pkg.develop(path="src/ArrowTypes")' +table = Arrow.Table("data.arrow") +propertynames(table) # [:id, :name] +isequal(collect(table.name), ["Ada", "Babbage", missing]) # true ``` -## Format Support +`Arrow.Table` accepts a path, an `IO`, IPC bytes, or an +`Arrow.AbstractArrowSource`. `Arrow.Stream` iterates one record batch at a time. +`Arrow.write` accepts any Tables.jl source. + +Arrow 3.0 includes: -This implementation supports the 1.0 version of the specification, including support for: - * All primitive data types - * All nested data types - * Dictionary encodings and messages - * Extension types - * Streaming, file, record batch, and replacement and isdelta dictionary messages +- IPC file and stream reads and writes. +- Incremental file and stream writing, plus IPC stream append. +- LZ4 frame and Zstandard buffer compression. +- Dictionary encoding. +- `Tables.Scan` projection, filter, limit, and offset pushdown. +- Sparse byte-range reads, including a CloudStore.jl extension. +- Arrow C data and C stream import and export. +- Recursive ArrowTypes.jl mappings for custom and extension types. +- Structural, semantic, and optional full-content validation. -It currently doesn't include support for: - * Tensors or sparse tensors - * Flight RPC - * C data interface +Arrow 3.0 is a breaking rewrite. Read the +[migration guide](docs/src/migration.md) before you update from Arrow 2.x. +See the [changelog](CHANGELOG.md) for the full release summary. The +[user manual](docs/src/manual.md) and +[API reference](docs/src/reference.md) describe the supported public API. -Third-party data formats: - * CSV, parquet and avro support via the existing [CSV.jl](https://github.com/JuliaData/CSV.jl), [Parquet.jl](https://github.com/JuliaIO/Parquet.jl) and [Avro.jl](https://github.com/JuliaData/Avro.jl) packages - * Other Tables.jl-compatible packages automatically supported ([DataFrames.jl](https://github.com/JuliaData/DataFrames.jl), [JSONTables.jl](https://github.com/JuliaData/JSONTables.jl), [JuliaDB.jl](https://github.com/JuliaData/JuliaDB.jl), [SQLite.jl](https://github.com/JuliaDatabases/SQLite.jl), [MySQL.jl](https://github.com/JuliaDatabases/MySQL.jl), [JDBC.jl](https://github.com/JuliaDatabases/JDBC.jl), [ODBC.jl](https://github.com/JuliaDatabases/ODBC.jl), [XLSX.jl](https://github.com/felipenoris/XLSX.jl), etc.) - * No current Julia packages support ORC +## Development + +In a checkout of this branch, prepare the local subpackages, then run the +tests: + +```julia +import Pkg +Pkg.activate(".") +Pkg.develop(path="src/ArrowStrings") +Pkg.develop(path="src/ArrowTypes") +Pkg.test() +``` -See the [full documentation](https://arrow.apache.org/julia/) for details on reading and writing arrow data. +The repository also has Apache Arrow gold-corpus checks, PyArrow and +Nanoarrow IPC oracle checks, and PyArrow C interface checks. Run all of them +with `julia conformance/run.jl`. Docker and network access for the first image +build are required. + +Run `julia --project=. test/fuzz.jl --cases 16 --mutations 64` for the +deterministic PR-sized fuzz suite. The scheduled workflow runs the extended +512-case and 20,000-mutation suite with a new reproducible master seed for each +scheduled run. It repeats the first full route sweep and every 256th mutation +to detect unstable outcomes. If the runner records or times out on a case, the +workflow uploads its replay coordinates, mutated bytes when available, the +resolved package environment, and a location-independent `replay.sh` wrapper. + +The Arrow 3.0 rewrite used Anthropic Claude Code and OpenAI Codex for code +generation, test generation, and review. Apache Arrow maintainers remain +responsible for understanding, reviewing, testing, and approving the code and +each release. + +See [the engine design](docs/dev/core-README.md) for the source layout and +internal contracts. diff --git a/bench/bench_pyarrow.py b/bench/bench_pyarrow.py new file mode 100644 index 00000000..675398fa --- /dev/null +++ b/bench/bench_pyarrow.py @@ -0,0 +1,118 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# PyArrow serialize/deserialize timing over the same logical workloads. +# Reads are pyarrow-idiomatic (memory-mapped read_all): pyarrow defers +# per-element materialization, so its read numbers measure wrapping, not +# element conversion — the driver's report states this asymmetry. +# Usage: run through bench/run.jl, which executes this file with the +# conformance image's oracle interpreter (/opt/pyarrow/bin/python). + +import sys, time, os +import pyarrow as pa +import pyarrow.ipc as ipc + +ROWS_PRIMITIVE = 10_000_000 +ROWS_STRINGS = 2_000_000 +ROWS_LISTS = 1_000_000 +ROWS_DICT = 2_000_000 + + +def wl_primitive(): + n = ROWS_PRIMITIVE + return pa.table({ + "a": pa.array(range(1, n + 1), type=pa.int64()), + "b": pa.array((float(i) for i in range(1, n + 1)), + type=pa.float64(), size=n), + }) + + +def wl_nullable(): + n = ROWS_PRIMITIVE + return pa.table({ + "a": pa.array((None if i % 7 == 0 else i for i in range(1, n + 1)), + type=pa.int64(), size=n), + }) + + +def wl_strings(): + n = ROWS_STRINGS + return pa.table({"s": pa.array("value-%d" % (i % 1000) + for i in range(1, n + 1))}) + + +def wl_lists(): + n = ROWS_LISTS + return pa.table({"l": pa.array([[i, i + 1, i + 2] + for i in range(1, n + 1)], + type=pa.list_(pa.int64()))}) + + +def wl_dictpool(): + # Plain strings: dictionary_encode runs INSIDE the write timer so all + # three legs time pool construction + dictionary write. + n = ROWS_DICT + return pa.table({"d": pa.array("cat-%d" % (i % 32) + for i in range(1, n + 1))}) + + +WORKLOADS = [ + ("primitive", wl_primitive), + ("nullable", wl_nullable), + ("strings", wl_strings), + ("lists", wl_lists), + ("dictpool", wl_dictpool), +] + + +def bench(f, runs=3): + f() + ts = [] + for _ in range(runs): + t0 = time.perf_counter() + f() + ts.append(time.perf_counter() - t0) + ts.sort() + return ts[len(ts) // 2] + + +def main(outdir): + for name, make in WORKLOADS: + tbl = make() + path = os.path.join(outdir, "pyarrow-%s.arrow" % name) + + def write(): + out = tbl + if name == "dictpool": + out = pa.table({"d": tbl["d"].combine_chunks() + .dictionary_encode()}) + with ipc.new_file(path, out.schema) as w: + w.write_table(out) + + twrite = bench(write) + size = os.path.getsize(path) + + def read(): + with pa.memory_map(path) as src: + ipc.open_file(src).read_all() + + tread = bench(read) + for op, secs in (("write", twrite), ("read", tread)): + print('{"impl":"pyarrow","workload":"%s","op":"%s",' + '"seconds":%r,"bytes":%d}' % (name, op, secs, size)) + + +main(sys.argv[1]) diff --git a/bench/bench_rewrite.jl b/bench/bench_rewrite.jl new file mode 100644 index 00000000..62b9f222 --- /dev/null +++ b/bench/bench_rewrite.jl @@ -0,0 +1,48 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Arrow.jl 3.0 (this repo) serialize/deserialize timing. Emits one JSON +# object per line: {"impl","workload","op","seconds","bytes"}. +# Usage: julia --project=. bench/bench_rewrite.jl + +using Arrow, Tables +include(joinpath(@__DIR__, "workloads.jl")) + +function main(outdir::String) + for (name, make) in BENCH_WORKLOADS + tbl = make() + path = joinpath(outdir, "rewrite-$name.arrow") + # write: table -> file bytes on disk + twrite = bench_time(() -> Arrow.write(path, tbl)) + sz = filesize(path) + # read: file -> fully materialized public columns + tread = bench_time() do + t = Arrow.Table(path) + for nm in Tables.columnnames(t) + length(Tables.getcolumn(t, nm)) + end + Arrow.release!(t) + end + for (op, secs) in (("write", twrite), ("read", tread)) + println( + "{\"impl\":\"rewrite\",\"workload\":\"$name\"," * + "\"op\":\"$op\",\"seconds\":$secs,\"bytes\":$sz}", + ) + end + end +end + +main(ARGS[1]) diff --git a/bench/run.jl b/bench/run.jl new file mode 100644 index 00000000..d3219f93 --- /dev/null +++ b/bench/run.jl @@ -0,0 +1,133 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Serialize/deserialize benchmark driver: Arrow.jl vs PyArrow over +# identical logical workloads. +# +# julia --project=. bench/run.jl [workdir] +# +# Legs run in their own processes (PyArrow inside the conformance image +# when docker is available — skipped cleanly otherwise). Results print as +# a markdown table of seconds and throughput. +# +# READ SEMANTICS DIFFER BY DESIGN, so read rows are not like-for-like: +# rewrite = full structural+semantic validation + materialized Julia +# Vectors (the facade contract) +# pyarrow = memory-mapped wrap only; per-element work is deferred +# Write rows ARE like-for-like: table in memory -> IPC file on disk. + +function _runleg(cmd::Cmd, out::String) + open(out, "w") do io + run(pipeline(cmd; stdout=io)) + end + return nothing +end + +# The workload names each leg must report (bench/workloads.jl owns their +# definitions; the driver cannot include that file — it references Arrow). +const WORKLOAD_NAMES = ("primitive", "nullable", "strings", "lists", "dictpool") + +function main(workdir::String) + mkpath(workdir) + here = @__DIR__ + repo = dirname(here) + legs = Tuple{String,String}[] + + rewriteout = joinpath(workdir, "rewrite.jsonl") + rewritescript = joinpath(here, "bench_rewrite.jl") + _runleg( + `$(Base.julia_cmd()) --startup-file=no --project=$repo $rewritescript $workdir`, + rewriteout, + ) + push!(legs, ("rewrite", rewriteout)) + + pyout = joinpath(workdir, "pyarrow.jsonl") + # The pyarrow leg runs the conformance image's oracle interpreter + # (build it once with `julia conformance/run.jl`). + havedocker = + Sys.which("docker") !== nothing && try + success( + pipeline( + `docker image inspect arrow-julia-conformance:latest`; + stdout=devnull, + stderr=devnull, + ), + ) + catch + false + end + if havedocker + image = "arrow-julia-conformance:latest" + python = "/opt/pyarrow/bin/python" + _runleg( + `docker run --rm -v $workdir:/bench -v $here:/src $image $python /src/bench_pyarrow.py /bench`, + pyout, + ) + push!(legs, ("pyarrow", pyout)) + else + println("(pyarrow leg skipped: conformance docker image not available)") + end + + # Minimal JSONL field extraction; the emitters write flat one-line + # objects with known keys. + results = Dict{Tuple{String,String,String},Tuple{Float64,Int64}}() + for (_, file) in legs, line in eachline(file) + isempty(strip(line)) && continue + g(k) = match(Regex("\"$k\":\"?([^\",}]+)"), line).captures[1] + key = (g("impl"), g("workload"), g("op")) + haskey(results, key) && error("duplicate benchmark record for $key") + results[key] = (parse(Float64, g("seconds")), parse(Int64, g("bytes"))) + end + # A leg that exits 0 with partial output must refuse, not print a + # plausible table. + impls = [name for (name, _) in legs] + for impl in impls, wl in WORKLOAD_NAMES, op in ("write", "read") + haskey(results, (impl, wl, op)) || + error("missing benchmark record for $((impl, wl, op))") + end + + println() + println("READ ROWS ARE NOT LIKE-FOR-LIKE: rewrite = validate + fully") + println("materialized Julia Vectors; pyarrow = memory-mapped wrap only,") + println("all per-element work deferred. Write rows are like-for-like.") + println() + println( + "| workload | op | " * + join(impls, " | ") * + " | MB/s (" * + join(impls, " / ") * + ") |", + ) + println("|---|---|" * repeat("---|", length(impls) + 1)) + for wl in WORKLOAD_NAMES, op in ("write", "read") + secs = [get(results, (impl, wl, op), (NaN, 0))[1] for impl in impls] + mbs = [ + begin + s, b = get(results, (impl, wl, op), (NaN, 0)) + isnan(s) ? "-" : string(round(b / s / 1e6; digits=0)) + end for impl in impls + ] + println( + "| $wl | $op | " * + join([isnan(s) ? "-" : string(round(s; digits=4)) for s in secs], " | ") * + " | " * + join(mbs, " / ") * + " |", + ) + end +end + +main(isempty(ARGS) ? mktempdir() : abspath(ARGS[1])) diff --git a/bench/workloads.jl b/bench/workloads.jl new file mode 100644 index 00000000..a24a4a73 --- /dev/null +++ b/bench/workloads.jl @@ -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. + +# Shared workload definitions for the serialize/deserialize benchmarks. +# Deterministic arithmetic data — every implementation +# builds the same logical tables, so file sizes and work agree. + +const BENCH_ROWS_PRIMITIVE = 10_000_000 +const BENCH_ROWS_STRINGS = 2_000_000 +const BENCH_ROWS_LISTS = 1_000_000 +const BENCH_ROWS_DICT = 2_000_000 + +function workload_primitive() + n = BENCH_ROWS_PRIMITIVE + return (a=collect(Int64, 1:n), b=collect(Float64, 1:n)) +end + +function workload_nullable() + n = BENCH_ROWS_PRIMITIVE + a = Vector{Union{Missing,Int64}}(undef, n) + for i = 1:n + a[i] = i % 7 == 0 ? missing : Int64(i) + end + return (a=a,) +end + +function workload_strings() + n = BENCH_ROWS_STRINGS + return (s=[string("value-", i % 1000) for i = 1:n],) +end + +function workload_lists() + n = BENCH_ROWS_LISTS + return (l=[Int64[i, i + 1, i + 2] for i = 1:n],) +end + +function workload_dictpool() + # The leg times pool construction + dictionary write from plain + # strings, matching the PyArrow leg's timed dictionary_encode + write. + n = BENCH_ROWS_DICT + return (d=Arrow.DictEncode([string("cat-", i % 32) for i = 1:n]),) +end + +const BENCH_WORKLOADS = ( + ("primitive", workload_primitive), + ("nullable", workload_nullable), + ("strings", workload_strings), + ("lists", workload_lists), + ("dictpool", workload_dictpool), +) + +"Median-of-k timing after one warmup run." +function bench_time(f::F; runs::Int=3) where {F} + f() + times = Float64[] + for _ = 1:runs + t0 = time_ns() + f() + push!(times, (time_ns() - t0) / 1e9) + end + return sort!(times)[cld(length(times), 2)] +end diff --git a/codecov.yaml b/codecov.yaml index a7b3fb59..d9b56627 100644 --- a/codecov.yaml +++ b/codecov.yaml @@ -18,5 +18,5 @@ --- codecov: notify: - # Wait for all "test" matrix jobs - after_n_builds: 6 + # 48 base matrix jobs minus the two ArrowTypes/min/macOS thread variants. + after_n_builds: 46 diff --git a/conformance/ConformanceSupport.jl b/conformance/ConformanceSupport.jl new file mode 100644 index 00000000..33cbbea0 --- /dev/null +++ b/conformance/ConformanceSupport.jl @@ -0,0 +1,1207 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Private support for the standalone Arrow conformance adapters. + +This module owns integration-JSON conversion, value-level comparison, skip +policy, verdict construction, and common reporting. The corpus, IPC oracle, +and C-interface oracle remain thin adapters that select cases and routes. +""" +module ConformanceSupport + +using CodecZlib +using JSON +import Arrow + +const AC = Arrow.ArrowCore +const ArrowCore = Arrow.ArrowCore + +include(joinpath(@__DIR__, "arrowjson.jl")) + +const DEFAULT_CORPUS = get(ENV, "ARROW_TESTING_DIR", "") + +# Families declared out of scope for every conformance route. Everything else +# must pass or it is a failure. +const FAMILY_SKIP = Dict{String,String}( + "0.17.1" => "V4 experimental compression marker era; superseded by 2.0.0-compression", +) + +# These families contain legacy gold IPC that our reader intentionally does +# not accept. Their integration JSON and IPC written by the current Arrow.jl +# implementation remain valid inputs to every other corpus and oracle route. +const GOLD_IPC_SKIP = Dict{String,String}( + "1.0.0-bigendian" => "big-endian streams are not supported (no endianness normalization)", + "0.14.1" => "pre-1.0 legacy framing (four-byte prefix) is not accepted by design", +) + +familyskipreason(family::AbstractString, context::AbstractString="") = + get(FAMILY_SKIP, family, get(FAMILY_SKIP, context, "")) + +goldipcskipreason(family::AbstractString, context::AbstractString="") = + get(GOLD_IPC_SKIP, family, get(GOLD_IPC_SKIP, context, "")) + +# --- value-level comparison ----------------------------------------------------- + +# Compared values normally need no physical identity: dictionary ids, pool +# unification, and integer JSON spellings are representation choices. A Union's +# selected type id is different. It chooses one declared child and remains part +# of the public-domain value even when two children materialize as equal Julia values. +struct _PublicUnion + typeid::Int8 + value::Any +end + +# Core materialization intentionally replaces the physical integration-JSON +# columns during comparison. Validate the JSON envelope first so that arity, +# names, scalar kinds, and per-slot array lengths cannot disappear at that +# boundary. Array lengths describe JSON entries, not legal byte-level buffer +# padding. Nested child counts remain independent so null List/Map segments can +# retain any spec-valid physical span. +_isjsoninteger(value) = value isa Integer && !(value isa Bool) + +function _jsonkind(value) + value === nothing && return "null" + value isa Bool && return "Bool" + value isa Integer && return "integer" + value isa AbstractFloat && return "float" + value isa AbstractString && return "string" + value isa AbstractVector && return "array" + value isa AbstractDict && return "object" + return string(typeof(value)) +end + +function _shapediff!(diffs::Vector{String}, path::String, detail::AbstractString) + length(diffs) > 20 || push!(diffs, "$path: $detail") + return nothing +end + +function _requireproperty!( + object::AbstractDict, + key::String, + predicate, + expectedkind::String, + path::String, + diffs::Vector{String}, +) + if !haskey(object, key) + _shapediff!(diffs, "$path.$key", "missing required property") + return false + end + predicate(object[key]) && return true + _shapediff!( + diffs, + "$path.$key", + "expected $expectedkind, got $(_jsonkind(object[key]))", + ) + return false +end + +function _validatemetadata!(owner::AbstractDict, path::String, diffs::Vector{String}) + haskey(owner, "metadata") || return + metadata = owner["metadata"] + metadata === nothing && return + if !(metadata isa AbstractVector) + _shapediff!( + diffs, + "$path.metadata", + "expected array or null, got $(_jsonkind(metadata))", + ) + return + end + for (index, entry) in enumerate(metadata) + entrypath = "$path.metadata[$index]" + if !(entry isa AbstractDict) + _shapediff!(diffs, entrypath, "expected object, got $(_jsonkind(entry))") + continue + end + _requireproperty!( + entry, + "key", + x -> x isa AbstractString, + "string", + entrypath, + diffs, + ) + _requireproperty!( + entry, + "value", + x -> x isa AbstractString, + "string", + entrypath, + diffs, + ) + end + return +end + +function _validatejsontype!(type, path::String, diffs::Vector{String}) + if !(type isa AbstractDict) + _shapediff!(diffs, path, "expected object, got $(_jsonkind(type))") + return + end + _requireproperty!(type, "name", x -> x isa AbstractString, "string", path, diffs) || + return + name = type["name"] + if name == "int" + _requireproperty!(type, "bitWidth", _isjsoninteger, "integer", path, diffs) + _requireproperty!(type, "isSigned", x -> x isa Bool, "Bool", path, diffs) + elseif name == "floatingpoint" + _requireproperty!( + type, + "precision", + x -> x isa AbstractString, + "string", + path, + diffs, + ) + elseif name == "fixedsizebinary" + _requireproperty!(type, "byteWidth", _isjsoninteger, "integer", path, diffs) + elseif name == "decimal" + _requireproperty!(type, "precision", _isjsoninteger, "integer", path, diffs) + _requireproperty!(type, "scale", _isjsoninteger, "integer", path, diffs) + haskey(type, "bitWidth") && + !_isjsoninteger(type["bitWidth"]) && + _shapediff!( + diffs, + "$path.bitWidth", + "expected integer, got $(_jsonkind(type["bitWidth"]))", + ) + elseif name in ("date", "duration", "interval") + _requireproperty!(type, "unit", x -> x isa AbstractString, "string", path, diffs) + elseif name == "time" + _requireproperty!(type, "unit", x -> x isa AbstractString, "string", path, diffs) + _requireproperty!(type, "bitWidth", _isjsoninteger, "integer", path, diffs) + elseif name == "timestamp" + _requireproperty!(type, "unit", x -> x isa AbstractString, "string", path, diffs) + haskey(type, "timezone") && + !(type["timezone"] isa AbstractString) && + _shapediff!( + diffs, + "$path.timezone", + "expected string, got $(_jsonkind(type["timezone"]))", + ) + elseif name == "fixedsizelist" + _requireproperty!(type, "listSize", _isjsoninteger, "integer", path, diffs) + elseif name == "map" + _requireproperty!(type, "keysSorted", x -> x isa Bool, "Bool", path, diffs) + elseif name == "union" + _requireproperty!(type, "mode", x -> x isa AbstractString, "string", path, diffs) + if _requireproperty!( + type, + "typeIds", + x -> x isa AbstractVector, + "array", + path, + diffs, + ) + for (index, id) in enumerate(type["typeIds"]) + _isjsoninteger(id) || _shapediff!( + diffs, + "$path.typeIds[$index]", + "expected integer, got $(_jsonkind(id))", + ) + end + end + elseif !( + name in ( + "null", + "bool", + "utf8", + "largeutf8", + "binary", + "largebinary", + "utf8view", + "binaryview", + "list", + "largelist", + "listview", + "largelistview", + "struct", + "runendencoded", + ) + ) + _shapediff!(diffs, "$path.name", "unknown Arrow type $(repr(name))") + end + return +end + +const _NO_CHILD_TYPES = Set(( + "null", + "bool", + "int", + "floatingpoint", + "utf8", + "largeutf8", + "binary", + "largebinary", + "utf8view", + "binaryview", + "fixedsizebinary", + "decimal", + "date", + "time", + "timestamp", + "duration", + "interval", +)) +const _ONE_CHILD_TYPES = + Set(("list", "largelist", "listview", "largelistview", "fixedsizelist", "map")) + +function _validatefieldarity!( + field::AbstractDict, + children::AbstractVector, + path::String, + diffs::Vector{String}, +) + type = get(field, "type", nothing) + type isa AbstractDict || return + name = get(type, "name", nothing) + name isa AbstractString || return + expected = + name in _NO_CHILD_TYPES ? 0 : + name in _ONE_CHILD_TYPES ? 1 : name == "runendencoded" ? 2 : nothing + if name == "union" && get(type, "typeIds", nothing) isa AbstractVector + expected = length(type["typeIds"]) + end + expected === nothing || + length(children) == expected || + _shapediff!( + diffs, + "$path.children", + "$name requires $expected schema children, got $(length(children))", + ) + if name == "map" && length(children) == 1 && children[1] isa AbstractDict + entries = children[1] + entrytype = get(entries, "type", nothing) + get(entrytype isa AbstractDict ? entrytype : Dict(), "name", nothing) == "struct" || + _shapediff!(diffs, "$path.children[1].type", "map entries must be a struct") + entrychildren = get(entries, "children", Any[]) + entrychildren isa AbstractVector && length(entrychildren) == 2 || _shapediff!( + diffs, + "$path.children[1].children", + "map entries struct must have key and value children", + ) + end + return +end + +function _validatefield!( + field, + path::String, + dictionaries::Dict{Int64,Tuple{Any,String}}, + diffs::Vector{String}, +) + if !(field isa AbstractDict) + _shapediff!(diffs, path, "expected object, got $(_jsonkind(field))") + return + end + _requireproperty!(field, "name", x -> x isa AbstractString, "string", path, diffs) + _requireproperty!(field, "nullable", x -> x isa Bool, "Bool", path, diffs) + if _requireproperty!(field, "type", x -> x isa AbstractDict, "object", path, diffs) + _validatejsontype!(field["type"], "$path.type", diffs) + end + _validatemetadata!(field, path, diffs) + + children = get(field, "children", Any[]) + if !(children isa AbstractVector) + _shapediff!(diffs, "$path.children", "expected array, got $(_jsonkind(children))") + children = Any[] + end + _validatefieldarity!(field, children, path, diffs) + for (index, child) in enumerate(children) + _validatefield!(child, "$path.children[$index]", dictionaries, diffs) + end + + haskey(field, "dictionary") || return + dictionary = field["dictionary"] + if !(dictionary isa AbstractDict) + _shapediff!( + diffs, + "$path.dictionary", + "expected object, got $(_jsonkind(dictionary))", + ) + return + end + idok = _requireproperty!( + dictionary, + "id", + _isjsoninteger, + "integer", + "$path.dictionary", + diffs, + ) + if _requireproperty!( + dictionary, + "indexType", + x -> x isa AbstractDict, + "object", + "$path.dictionary", + diffs, + ) + indextype = dictionary["indexType"] + _validatejsontype!(indextype, "$path.dictionary.indexType", diffs) + get(indextype, "name", nothing) == "int" || _shapediff!( + diffs, + "$path.dictionary.indexType.name", + "dictionary index type must be int", + ) + end + _requireproperty!( + dictionary, + "isOrdered", + x -> x isa Bool, + "Bool", + "$path.dictionary", + diffs, + ) + idok || return + id = try + Int64(dictionary["id"]) + catch + _shapediff!(diffs, "$path.dictionary.id", "integer is outside Int64 range") + return + end + get!(dictionaries, id, (field, path)) + return +end + +const _COLUMN_ARRAY_PROPERTIES = ("DATA", "OFFSET", "SIZE", "TYPE_ID", "VIEWS") +const _DATA_TYPES = Set(( + "bool", + "int", + "floatingpoint", + "decimal", + "date", + "time", + "timestamp", + "duration", + "interval", + "fixedsizebinary", +)) + +function _validatearraylength!( + column::AbstractDict, + key::String, + expected::Int, + path::String, + diffs::Vector{String}; + required::Bool=true, +) + if !haskey(column, key) + required && _shapediff!(diffs, "$path.$key", "missing required array") + return + end + values = column[key] + if !(values isa AbstractVector) + _shapediff!(diffs, "$path.$key", "expected array, got $(_jsonkind(values))") + return + end + length(values) == expected || _shapediff!( + diffs, + "$path.$key", + "expected $expected entries from count, got $(length(values))", + ) + return +end + +function _validatecolumnarrays!( + field::AbstractDict, + column::AbstractDict, + count, + path::String, + diffs::Vector{String}; + dictionaryvalue::Bool=false, +) + count isa Integer || return + count < 0 && (_shapediff!(diffs, "$path.count", "must be non-negative"); return) + n = try + Int(count) + catch + _shapediff!(diffs, "$path.count", "integer is outside Int range") + return + end + if haskey(column, "VALIDITY") + _validatearraylength!(column, "VALIDITY", n, path, diffs) + end + + type = get(field, "type", nothing) + type isa AbstractDict || return + name = + !dictionaryvalue && haskey(field, "dictionary") ? "dictionary" : + get(type, "name", nothing) + name isa AbstractString || return + lengths = Dict{String,Int}() + if name == "dictionary" || name in _DATA_TYPES + lengths["DATA"] = n + elseif name in ("utf8", "largeutf8", "binary", "largebinary") + lengths["DATA"] = n + lengths["OFFSET"] = n + 1 + elseif name in ("utf8view", "binaryview") + lengths["VIEWS"] = n + if haskey(column, "VARIADIC_DATA_BUFFERS") && + !(column["VARIADIC_DATA_BUFFERS"] isa AbstractVector) + _shapediff!( + diffs, + "$path.VARIADIC_DATA_BUFFERS", + "expected array, got $(_jsonkind(column["VARIADIC_DATA_BUFFERS"]))", + ) + end + elseif name in ("list", "largelist", "map") + lengths["OFFSET"] = n + 1 + elseif name in ("listview", "largelistview") + lengths["OFFSET"] = n + lengths["SIZE"] = n + elseif name == "union" + lengths["TYPE_ID"] = n + get(type, "mode", nothing) == "DENSE" && (lengths["OFFSET"] = n) + end + for (key, expected) in lengths + _validatearraylength!(column, key, expected, path, diffs) + end + for key in _COLUMN_ARRAY_PROPERTIES + haskey(column, key) && + !haskey(lengths, key) && + _shapediff!(diffs, "$path.$key", "unexpected array for $name") + end + if name != "utf8view" && name != "binaryview" && haskey(column, "VARIADIC_DATA_BUFFERS") + _shapediff!(diffs, "$path.VARIADIC_DATA_BUFFERS", "unexpected array for $name") + end + return +end + +function _validatecolumn!( + field::AbstractDict, + column, + path::String, + diffs::Vector{String}; + checkname::Bool=true, + dictionaryvalue::Bool=false, +) + if !(column isa AbstractDict) + _shapediff!(diffs, path, "expected object, got $(_jsonkind(column))") + return nothing + end + nameok = + _requireproperty!(column, "name", x -> x isa AbstractString, "string", path, diffs) + fieldname = get(field, "name", nothing) + if checkname && nameok && fieldname isa AbstractString && fieldname != column["name"] + _shapediff!( + diffs, + "$path.name", + "$(repr(column["name"])) does not match schema field $(repr(fieldname))", + ) + end + countok = _requireproperty!(column, "count", _isjsoninteger, "integer", path, diffs) + countok && + _validatecolumnarrays!(field, column, column["count"], path, diffs; dictionaryvalue) + children = get(column, "children", Any[]) + if !(children isa AbstractVector) + _shapediff!(diffs, "$path.children", "expected array, got $(_jsonkind(children))") + children = Any[] + end + fieldchildren = + !dictionaryvalue && haskey(field, "dictionary") ? Any[] : + get(field, "children", Any[]) + if fieldchildren isa AbstractVector + length(children) == length(fieldchildren) || _shapediff!( + diffs, + "$path.children", + "expected $(length(fieldchildren)) child columns, got $(length(children))", + ) + for index = 1:min(length(children), length(fieldchildren)) + childfield = fieldchildren[index] + childfield isa AbstractDict || continue + _validatecolumn!(childfield, children[index], "$path.children[$index]", diffs) + end + end + return countok ? column["count"] : nothing +end + +function _validatedocument!(document, path::String, diffs::Vector{String}) + if !(document isa AbstractDict) + _shapediff!(diffs, path, "expected object, got $(_jsonkind(document))") + return + end + if !_requireproperty!( + document, + "schema", + x -> x isa AbstractDict, + "object", + path, + diffs, + ) + return + end + schema = document["schema"] + _validatemetadata!(schema, "$path.schema", diffs) + if !_requireproperty!( + schema, + "fields", + x -> x isa AbstractVector, + "array", + "$path.schema", + diffs, + ) + return + end + fields = schema["fields"] + dictionaryfields = Dict{Int64,Tuple{Any,String}}() + for (index, field) in enumerate(fields) + _validatefield!(field, "$path.schema.fields[$index]", dictionaryfields, diffs) + end + + if _requireproperty!( + document, + "batches", + x -> x isa AbstractVector, + "array", + path, + diffs, + ) + for (batchindex, batch) in enumerate(document["batches"]) + batchpath = "$path.batches[$batchindex]" + if !(batch isa AbstractDict) + _shapediff!(diffs, batchpath, "expected object, got $(_jsonkind(batch))") + continue + end + countok = _requireproperty!( + batch, + "count", + _isjsoninteger, + "integer", + batchpath, + diffs, + ) + if !_requireproperty!( + batch, + "columns", + x -> x isa AbstractVector, + "array", + batchpath, + diffs, + ) + continue + end + columns = batch["columns"] + length(columns) == length(fields) || _shapediff!( + diffs, + "$batchpath.columns", + "expected $(length(fields)) columns, got $(length(columns))", + ) + for index = 1:min(length(columns), length(fields)) + field = fields[index] + field isa AbstractDict || continue + columncount = _validatecolumn!( + field, + columns[index], + "$batchpath.columns[$index]", + diffs, + ) + if countok && columncount !== nothing && columncount != batch["count"] + _shapediff!( + diffs, + "$batchpath.columns[$index].count", + "$columncount does not match batch count $(batch["count"])", + ) + end + end + end + end + + entries = get(document, "dictionaries", Any[]) + if !(entries isa AbstractVector) + _shapediff!( + diffs, + "$path.dictionaries", + "expected array, got $(_jsonkind(entries))", + ) + return + end + seen = Set{Int64}() + for (entryindex, entry) in enumerate(entries) + entrypath = "$path.dictionaries[$entryindex]" + if !(entry isa AbstractDict) + _shapediff!(diffs, entrypath, "expected object, got $(_jsonkind(entry))") + continue + end + idok = _requireproperty!(entry, "id", _isjsoninteger, "integer", entrypath, diffs) + dataok = _requireproperty!( + entry, + "data", + x -> x isa AbstractDict, + "object", + entrypath, + diffs, + ) + (idok && dataok) || continue + id = try + Int64(entry["id"]) + catch + _shapediff!(diffs, "$entrypath.id", "integer is outside Int64 range") + continue + end + id in seen && _shapediff!(diffs, "$entrypath.id", "duplicate dictionary id $id") + push!(seen, id) + if !haskey(dictionaryfields, id) + _shapediff!( + diffs, + "$entrypath.id", + "dictionary id $id is not used by the schema", + ) + continue + end + data = entry["data"] + countok = _requireproperty!( + data, + "count", + _isjsoninteger, + "integer", + "$entrypath.data", + diffs, + ) + if !_requireproperty!( + data, + "columns", + x -> x isa AbstractVector, + "array", + "$entrypath.data", + diffs, + ) + continue + end + columns = data["columns"] + length(columns) == 1 || _shapediff!( + diffs, + "$entrypath.data.columns", + "expected one dictionary value column, got $(length(columns))", + ) + isempty(columns) && continue + field, _ = dictionaryfields[id] + columncount = _validatecolumn!( + field, + columns[1], + "$entrypath.data.columns[1]", + diffs; + checkname=false, + dictionaryvalue=true, + ) + if countok && columncount !== nothing && columncount != data["count"] + _shapediff!( + diffs, + "$entrypath.data.columns[1].count", + "$columncount does not match dictionary batch count $(data["count"])", + ) + end + end + for (id, (_, fieldpath)) in dictionaryfields + id in seen || _shapediff!( + diffs, + "$fieldpath.dictionary.id", + "dictionary id $id has no value batch", + ) + end + return +end + +function _eq(a, b, path::String, diffs::Vector{String}; strict::Bool=false) + if a isa _PublicUnion || b isa _PublicUnion + if !(a isa _PublicUnion && b isa _PublicUnion) + push!(diffs, "$path: Union value vs non-Union value") + elseif a.typeid != b.typeid + push!(diffs, "$path: Union type id $(a.typeid) vs $(b.typeid)") + else + _eq(a.value, b.value, path * ".value", diffs; strict) + end + elseif a isa AbstractDict && b isa AbstractDict + ka, kb = Set(keys(a)), Set(keys(b)) + for k in union(ka, kb) + childpath = path * "." * String(k) + if !haskey(a, k) + push!(diffs, "$childpath: missing from actual document") + elseif !haskey(b, k) + push!(diffs, "$childpath: missing from expected document") + else + _eq(a[k], b[k], childpath, diffs; strict) + end + length(diffs) > 20 && return + end + elseif a isa AbstractVector && b isa AbstractVector + length(a) == length(b) || + (push!(diffs, "$path: length $(length(a)) vs $(length(b))"); return) + for (i, (x, y)) in enumerate(zip(a, b)) + _eq(x, y, path * "[$i]", diffs; strict) + length(diffs) > 20 && return + end + elseif strict + equal = + a isa Bool || b isa Bool ? (a isa Bool && b isa Bool && a == b) : + _isjsoninteger(a) || _isjsoninteger(b) ? + (_isjsoninteger(a) && _isjsoninteger(b) && a == b) : + a isa AbstractFloat || b isa AbstractFloat ? + (a isa AbstractFloat && b isa AbstractFloat && isequal(a, b)) : isequal(a, b) + equal || push!( + diffs, + "$path: $(repr(a)) ($(_jsonkind(a))) vs $(repr(b)) ($(_jsonkind(b)))", + ) + elseif a isa AbstractFloat || b isa AbstractFloat + # EXACT equality (± zero unified, NaN equal): a tolerance here would + # bless changed values. Core materialization has already canonicalized + # sub-double columns through their declared physical precision. + equal = + a isa AbstractFloat && b isa AbstractFloat && ((isnan(a) && isnan(b)) || a == b) + equal || push!(diffs, "$path: $(repr(a)) vs $(repr(b))") + elseif a isa Bool || b isa Bool + (a isa Bool && b isa Bool && a == b) || + push!(diffs, "$path: $(repr(a)) vs $(repr(b))") + elseif a isa Integer || b isa Integer + (a isa Integer && b isa Integer && Int128(a) == Int128(b)) || + push!(diffs, "$path: $(repr(a)) vs $(repr(b))") + else + isequal(a, b) || push!(diffs, "$path: $(repr(a)) vs $(repr(b))") + end + return +end + +# Dictionary ids are adapter bookkeeping, and the gold corpus is not id-stable +# across its own representations. Canonicalize ids by dictionary-typed field +# position. `_logicaldocument` rebuilds the pools in the same traversal order. +function _normalizedictionaryids!(fields) + nextid = Ref(0) + function renumber!(f) + f isa AbstractDict || return + if get(f, "dictionary", nothing) isa AbstractDict + f["dictionary"]["id"] = nextid[] + nextid[] += 1 + end + foreach(renumber!, get(f, "children", Any[])) + end + foreach(renumber!, fields) + return fields +end + +# Map entries-struct names are not round-trip stable in the corpus itself. +# Compare map children structurally by position. +function _normalizemapschema!(fields) + for field in fields + field isa AbstractDict || continue + if get(get(field, "type", Dict()), "name", "") == "map" && haskey(field, "children") + for child in field["children"] + child["name"] = "entries" + for (index, keyvalue) in enumerate(get(child, "children", Any[])) + keyvalue["name"] = index == 1 ? "key" : "value" + end + end + end + haskey(field, "children") && _normalizemapschema!(field["children"]) + end + return +end + +function _normalizemetadata!(value) + if value isa AbstractDict + if haskey(value, "metadata") + metadata = value["metadata"] + if metadata === nothing || (metadata isa AbstractVector && isempty(metadata)) + delete!(value, "metadata") + elseif metadata isa AbstractVector + value["metadata"] = sort( + metadata; + by=keyvalue -> (String(keyvalue["key"]), String(keyvalue["value"])), + ) + end + end + foreach(_normalizemetadata!, values(value)) + elseif value isa AbstractVector + foreach(_normalizemetadata!, value) + end + return value +end + +# Decimal bitWidth is optional with a default of 128 in integration JSON. +function _normalizedecimals!(value) + if value isa AbstractDict + if get(value, "name", "") == "decimal" && haskey(value, "precision") + width = get(value, "bitWidth", 128) + _isjsoninteger(width) && width == 128 && delete!(value, "bitWidth") + end + foreach(_normalizedecimals!, values(value)) + elseif value isa AbstractVector + foreach(_normalizedecimals!, value) + end + return value +end + +# Normalize only the schema. `_logicaldocument` replaces the raw batch and +# dictionary envelopes after it materializes their storage-domain values. +function _normalize!(doc::AbstractDict) + schema = get(doc, "schema", Dict()) + fields = get(schema, "fields", Any[]) + _normalizedictionaryids!(fields) + _normalizemapschema!(fields) + _normalizemetadata!(schema) + _normalizedecimals!(schema) + return doc +end + +_publicvalue(::Missing) = nothing +_publicvalue(value::Pair) = Any[_publicvalue(first(value)), _publicvalue(last(value))] +_publicvalue(value::NamedTuple) = + Any[Any[String(name), _publicvalue(item)] for (name, item) in pairs(value)] +_publicvalue(value::Tuple) = Any[_publicvalue(item) for item in value] +_publicvalue(value::AbstractVector) = Any[_publicvalue(item) for item in value] +_publicvalue(value) = value + +_typecontainsunion(::AC.ArrowType) = false +_typecontainsunion(::AC.UnionType) = true +_typecontainsunion(type::AC.DictionaryType) = _typecontainsunion(type.valuetype) +_fieldcontainsunion(field) = + _typecontainsunion(field.type) || any(_fieldcontainsunion, field.children) + +# Core's ordinary materializer intentionally exposes only the selected Union +# child's Julia value. Conformance comparison additionally needs the selected +# Arrow type id. Walk only fields that contain a Union, preserving that id at +# every depth while applying the same logical masking rules as Core: storage +# below a null parent does not participate in the value. +function _publicvalueat(field, data, index::Int64) + 1 <= index <= data.len || throw(BoundsError(data, index)) + type = data.type + + if type isa AC.UnionType + typeid = + AC.loadat(AC.rolebuffer(data, AC.TYPE_IDS), Int8, AC._slotindex0(data, index)) + childfield, childdata, childindex = AC._union_child(field, data, index) + return _PublicUnion(typeid, _publicvalueat(childfield, childdata, childindex)) + elseif type isa AC.RunEndEncodedType + run = AC._ree_runindex(data, index) + return _publicvalueat(field.children[2], data.children[2], run) + elseif type isa AC.NullType + return nothing + end + + AC.isvalid_at(data, index) || return nothing + + if type isa AC.DictionaryType + width = AC.primwidth(type.indextype) + dictionaryindex = AC._load_int( + AC.rolebuffer(data, AC.DATA), + type.indextype, + AC._slotbyteoff(data, index, width), + ) + dictionary = data.dictionary + dictionary === nothing && + throw(AC.ValidationError("dictionary-encoded array without a dictionary")) + valuefield = AC.dictvaluefield(field, type) + return _publicvalueat( + valuefield, + dictionary, + AC.checked_add(Int64(dictionaryindex), Int64(1)), + ) + elseif type isa AC.ListType + low, high = AC._offsets_at(data, index, type.large) + childfield, childdata = field.children[1], data.children[1] + values = Vector{Any}(undef, Int(high - low)) + for item = 1:length(values) + values[item] = + _publicvalueat(childfield, childdata, AC.checked_add(low, Int64(item))) + end + return values + elseif type isa AC.FixedSizeListType + childfield, childdata = field.children[1], data.children[1] + base = AC.checked_mul(AC._slotindex0(data, index), Int64(type.listsize)) + values = Vector{Any}(undef, type.listsize) + for item = 1:(type.listsize) + values[item] = + _publicvalueat(childfield, childdata, AC.checked_add(base, Int64(item))) + end + return values + elseif type isa AC.StructType + childindex = AC.checked_add(data.offset, index) + return Any[ + Any[String(childfield.name), _publicvalueat(childfield, childdata, childindex)] + for (childfield, childdata) in zip(field.children, data.children) + ] + elseif type isa AC.MapType + low, high = AC._offsets_at(data, index, false) + entriesfield, entriesdata = field.children[1], data.children[1] + keyfield, valuefield = entriesfield.children + keydata, valuedata = entriesdata.children + values = Vector{Any}(undef, Int(high - low)) + for item = 1:length(values) + entryindex = + AC.checked_add(entriesdata.offset, AC.checked_add(low, Int64(item))) + values[item] = Any[ + _publicvalueat(keyfield, keydata, entryindex), + _publicvalueat(valuefield, valuedata, entryindex), + ] + end + return values + elseif type isa AC.ListViewType + offset, size = AC._listview_range(type, data, index) + childfield, childdata = field.children[1], data.children[1] + values = Vector{Any}(undef, Int(size)) + for item = 1:length(values) + values[item] = + _publicvalueat(childfield, childdata, AC.checked_add(offset, Int64(item))) + end + return values + end + + return _publicvalue(AC.getvalue(field, data, index)) +end + +function _logicalcolumn(field, data; name=field.name) + values = if _fieldcontainsunion(field) + Any[_publicvalueat(field, data, Int64(index)) for index = 1:(data.len)] + else + Any[_publicvalue(value) for value in AC.materialize(field, data)] + end + return Dict{String,Any}( + "name" => String(name), + "count" => Int(data.len), + "VALUES" => values, + ) +end + +function _logicaldictionaries(fields, dictids, dictionaries) + result = Any[] + function visit(field) + if field.type isa AC.DictionaryType + id = length(result) + values = dictionaries[dictids[field]] + valuefield = AC.dictvaluefield(field, field.type) + push!( + result, + Dict{String,Any}( + "id" => id, + "data" => Dict{String,Any}( + "count" => Int(values.len), + "columns" => + Any[_logicalcolumn(valuefield, values; name="DICT")], + ), + ), + ) + end + foreach(visit, field.children) + return + end + foreach(visit, fields) + return result +end + +# Integration JSON exposes physical buffers, but every conformance route is a +# value comparison. Materializing through Core gives one canonical form for +# primitive and nested values. In particular, a null List or Map may own an +# arbitrary non-empty child segment, and a null Struct or FixedSizeList hides +# the corresponding child slots. A recursive DATA mask cannot normalize those +# legal layouts because their child lengths and offsets can differ. +function _logicaldocument(doc::AbstractDict, normalized::AbstractDict) + schema, batches, dictids, dictionaries = ArrowJSON._fromjson(doc) + result = normalized + result["batches"] = Any[ + Dict{String,Any}( + "count" => Int(batch.nrows), + "columns" => Any[ + _logicalcolumn(field, data) for + (field, data) in zip(schema.fields, batch.columns) + ], + ) for batch in batches + ] + pools = _logicaldictionaries(schema.fields, dictids, dictionaries) + if isempty(pools) + delete!(result, "dictionaries") + else + result["dictionaries"] = pools + end + return result +end + +function _normalizedschema(schema) + rendered = ArrowJSON.tojson(schema, AC.RecordBatch[])["schema"] + _normalizemapschema!(get(rendered, "fields", Any[])) + _normalizemetadata!(rendered) + _normalizedecimals!(rendered) + return rendered +end + +function _logicalwindow(field, data, offset::Int, rows::Int) + if _fieldcontainsunion(field) + return Any[_publicvalueat(field, data, Int64(offset + index)) for index = 1:rows] + end + return Any[_publicvalue(AC.getvalue(field, data, offset + index)) for index = 1:rows] +end + +""" +Compare two Core record-batch row windows by schema and normalized values. + +This is the canonical comparison for nonzero-offset C Data slices. Rendering +such an array as standalone integration JSON loses its physical base offset; +for nested arrays it can also expose child rows outside the logical window. +The actual and expected offsets are zero-based row offsets within each batch. +""" +function corebatchdiffs( + actual, + expected; + actualoffset::Integer=0, + expectedoffset::Integer=0, + rows::Integer=actual.nrows - actualoffset, +) + ao, eo, n = Int(actualoffset), Int(expectedoffset), Int(rows) + ao >= 0 || throw(ArgumentError("actual row offset must be non-negative")) + eo >= 0 || throw(ArgumentError("expected row offset must be non-negative")) + n >= 0 || throw(ArgumentError("row count must be non-negative")) + ao + n <= actual.nrows || throw(BoundsError(actual, (ao + 1):(ao + n))) + eo + n <= expected.nrows || throw(BoundsError(expected, (eo + 1):(eo + n))) + + diffs = String[] + _eq( + _normalizedschema(actual.schema), + _normalizedschema(expected.schema), + ".schema", + diffs; + strict=true, + ) + isempty(diffs) || return diffs + for (index, (actualfield, actualdata, expectedfield, expecteddata)) in enumerate( + zip(actual.schema.fields, actual.columns, expected.schema.fields, expected.columns), + ) + _eq( + _logicalwindow(actualfield, actualdata, ao, n), + _logicalwindow(expectedfield, expecteddata, eo, n), + ".columns[$index]", + diffs, + ) + length(diffs) > 20 && break + end + return diffs +end + +""" +Return logical-value differences between two integration-JSON documents. + +The operation owns schema normalization and Core materialization. It compares +schema before materialization so a missing required schema property remains a +precise difference instead of becoming a parser error. Callers keep their +documents unchanged and do not depend on the comparison implementation. +""" +function documentdiffs(actual, expected) + diffs = String[] + _validatedocument!(actual, ".actual", diffs) + _validatedocument!(expected, ".expected", diffs) + isempty(diffs) || return diffs + gotnormalized = _normalize!(deepcopy(actual)) + wantnormalized = _normalize!(deepcopy(expected)) + _eq( + get(gotnormalized, "schema", nothing), + get(wantnormalized, "schema", nothing), + ".schema", + diffs, + strict=true, + ) + isempty(diffs) || return diffs + got = _logicaldocument(actual, gotnormalized) + want = _logicaldocument(expected, wantnormalized) + _eq(got, want, "", diffs) + return diffs +end + +const VERDICT_STATUSES = (:pass, :fail, :skip) + +struct Verdict + family::String + check::String + status::Symbol + detail::String + + function Verdict(family, check, status::Symbol, detail="") + status in VERDICT_STATUSES || + throw(ArgumentError("invalid conformance status: $(repr(status))")) + return new(String(family), String(check), status, String(detail)) + end +end + +function errorverdict(family, check, err) + if err isa InterruptException || err isa OutOfMemoryError + throw(err) + end + detail = sprint(showerror, err) + return Verdict(family, check, :fail, first(detail, 200)) +end + +""" +Run one document producer and convert its result or error to one verdict. +""" +function documentcheck(producer::F, family, check, expected) where {F} + try + diffs = documentdiffs(producer(), expected) + return Verdict( + family, + check, + isempty(diffs) ? :pass : :fail, + isempty(diffs) ? "" : first(diffs), + ) + catch err + return errorverdict(family, check, err) + end +end + +function readjson(path::AbstractString) + bytes = read(path) + endswith(path, ".gz") && (bytes = transcode(GzipDecompressor, bytes)) + return JSON.parse(String(bytes)) +end + +function streamdocument(bytes::Vector{UInt8}) + stream = Arrow.readstream(bytes) + # Use the reader's id table so shared and nested pool ids survive. + return ArrowJSON.tojson(stream.schema, stream.batches; dictids=stream.fielddictids) +end + +function filedocument(bytes::Vector{UInt8}) + file = Arrow.readfile(bytes) + batches = AC.RecordBatch[file[i] for i = 1:length(file)] + return ArrowJSON.tojson(file.schema, batches; dictids=file.fielddictids) +end + +function report( + header::AbstractString, + verdicts::AbstractVector{Verdict}; + io::IO=stdout, + checkwidth::Integer=24, +) + isempty(verdicts) && throw( + ArgumentError("$header produced no conformance verdicts; refusing an empty pass"), + ) + npass = count(v -> v.status == :pass, verdicts) + nfail = count(v -> v.status == :fail, verdicts) + nskip = count(v -> v.status == :skip, verdicts) + println(io, "$header: $npass pass, $nfail fail, $nskip skip") + println(io) + for item in verdicts + item.status == :pass && continue + tag = item.status == :fail ? "FAIL" : "skip" + println( + io, + rpad(tag, 5), + rpad(item.family, 58), + rpad(item.check, checkwidth), + item.detail, + ) + end + return nfail +end + +end # module ConformanceSupport diff --git a/conformance/Dockerfile b/conformance/Dockerfile new file mode 100644 index 00000000..eb47c165 --- /dev/null +++ b/conformance/Dockerfile @@ -0,0 +1,95 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# The Arrow.jl conformance image: everything the conformance suites need, +# in one container, so that "docker present" is the only host requirement. +# +# * Julia (the base image) — the suites run IN the container, because the +# C data and C stream interfaces hand pointers across an in-process +# boundary; only IPC bytes could cross a container boundary. +# * A Python with pyarrow and nanoarrow — the oracles. PythonCall binds to +# it (JULIA_PYTHONCALL_EXE; the CondaPkg backend is disabled). +# * The apache/arrow-testing gold corpus. +# * A warm Julia depot with the registered dependencies, in a named volume +# at run time so precompilation caches survive between runs. +# +# The repository is bind-mounted at /work at run time (never copied in), so +# the image is independent of the source tree; conformance/run.jl builds it +# once and drives it. Build context: the conformance/ directory. +# +# docker build -t arrow-julia-conformance:latest -f conformance/Dockerfile conformance/ + +# Pinned direct tool and corpus revisions make conformance runs repeatable +# against those inputs. The base-image digest, apt packages, and transitive +# Julia/Python dependencies are not fully locked. The package itself supports +# Julia 1.10 and later; Julia 1.12 is also the trim-test floor. +FROM julia:1.12.6-bookworm + +ARG PIP_VERSION=26.2.1 +ARG PYARROW_VERSION=25.0.1 +ARG NANOARROW_VERSION=0.9.0 +ARG ARROW_TESTING_REV=9ff285c88565f0f6abc855918c6a342e70e4909c + +# python3-dev brings the shared libpython that PythonCall embeds; Debian's +# interpreter binary is statically linked and does not ship it otherwise. +RUN apt-get update \ + && apt-get install -y --no-install-recommends python3 python3-venv python3-dev git ca-certificates \ + && rm -rf /var/lib/apt/lists/* + +# The oracle interpreter. A venv keeps the Debian system Python untouched. +RUN python3 -m venv /opt/pyarrow \ + && /opt/pyarrow/bin/pip install --no-cache-dir "pip==${PIP_VERSION}" \ + && /opt/pyarrow/bin/pip install --no-cache-dir \ + "pyarrow==${PYARROW_VERSION}" "nanoarrow==${NANOARROW_VERSION}" + +RUN git init /opt/arrow-testing \ + && git -C /opt/arrow-testing remote add origin https://github.com/apache/arrow-testing.git \ + && git -C /opt/arrow-testing fetch --depth 1 origin "${ARROW_TESTING_REV}" \ + && git -C /opt/arrow-testing checkout --detach FETCH_HEAD + +ENV ARROW_TESTING_DIR=/opt/arrow-testing \ + ARROW_ORACLE_PYTHON=/opt/pyarrow/bin/python \ + ARROW_CDATA_ORACLE_PYTHON=/opt/pyarrow/bin/python \ + JULIA_PYTHONCALL_EXE=/opt/pyarrow/bin/python \ + JULIA_CONDAPKG_BACKEND=Null \ + JULIA_DEPOT_PATH=/opt/julia-depot \ + JULIA_PROJECT=/opt/env +# JULIA_PROJECT serves interactive `docker exec ... julia`; the driver passes +# --project explicitly. + +# Warm the depot: every registered dependency the package and the conformance +# environment resolve to (keep in step with Project.toml and +# conformance/Project.toml — the run-time instantiate is authoritative and +# fetches any delta), plus PythonCall precompiled against the oracle Python. +RUN julia -e 'using Pkg; Pkg.activate("/opt/warm"); \ + Pkg.add([PackageSpec(name="CodecLz4", version="0.4.6"), \ + PackageSpec(name="CodecZstd", version="0.8.7"), \ + PackageSpec(name="DataAPI", version="1.16.0"), \ + PackageSpec(name="EnumX", version="1.0.7"), \ + PackageSpec(name="TranscodingStreams", version="0.11.3"), \ + PackageSpec(name="JSON", version="1.7.1"), \ + PackageSpec(name="CodecZlib", version="0.7.9"), \ + PackageSpec(name="Tables", version="1.14.0"), \ + PackageSpec(name="PythonCall", version="0.9.35")]); \ + Pkg.precompile()' + +# The suite environment is created from the mounted repository's +# conformance/Project.toml by conformance/run.jl (developing /work into it), +# so a Manifest never lives in the repository checkout. +RUN mkdir -p /opt/env /work +WORKDIR /work +# The driver keeps its container alive the same way. +CMD ["sleep", "infinity"] diff --git a/conformance/Project.toml b/conformance/Project.toml new file mode 100644 index 00000000..4b4108a3 --- /dev/null +++ b/conformance/Project.toml @@ -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. + + +[deps] +Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45" +ArrowStrings = "c38d8858-22a2-449e-9eca-ef92ec15f353" +ArrowTypes = "31f734f8-188a-4ce0-8406-c8a06bd891cd" +CodecLz4 = "5ba52731-8f18-5e0d-9241-30f10d1ec561" +CodecZlib = "944b1d66-785c-5afd-91f1-9de20f533193" +CodecZstd = "6b39b394-51ab-5f42-8807-6242bab2b4c2" +DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a" +EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56" +JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" +PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" +Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" +TranscodingStreams = "3bb67fe8-82b1-5028-8e26-92a6c54297fa" + +[compat] +ArrowStrings = "=1.0.0" +ArrowTypes = "=2.4.0" +CodecLz4 = "=0.4.6" +CodecZlib = "=0.7.9" +CodecZstd = "=0.8.7" +DataAPI = "=1.16.0" +EnumX = "=1.0.7" +JSON = "=1.7.1" +PythonCall = "=0.9.35" +Tables = "=1.14.0" +TranscodingStreams = "=0.11.3" +julia = "1.12.6" diff --git a/conformance/arrowjson.jl b/conformance/arrowjson.jl new file mode 100644 index 00000000..2ebad812 --- /dev/null +++ b/conformance/arrowjson.jl @@ -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. + +# ============================================================================= +# Arrow integration JSON ("arrowjson") <-> ArrowCore. +# +# The integration JSON format is the cross-implementation conformance +# interchange used by apache/arrow-testing gold files and by archery. This +# module maps it to and from Core `Schema`/`RecordBatch` values so the corpus +# can be tested in every direction (JSON -> Core -> IPC vs gold bytes; gold +# IPC -> Core -> JSON vs gold JSON) and so our own writer output can be +# expressed as JSON for third-party consumers. +# +# Value conventions, pinned from real gold files (arrow-testing +# 1.0.0-littleendian and cpp-21.0.0): +# * every batch column: {name, count, VALIDITY?, DATA?/OFFSET?/... , children?} +# * VALIDITY is 0/1 ints; DATA for 8..32-bit ints, floats, bool are JSON +# scalars; 64-bit ints, decimals, large offsets/sizes are STRINGS +# * binary/fixedsizebinary/large binary DATA are UPPERCASE HEX strings +# * views: VIEWS entries {SIZE, INLINED} (hex for binary, utf8 for strings) +# or {SIZE, PREFIX_HEX, BUFFER_INDEX, OFFSET}, plus VARIADIC_DATA_BUFFERS +# as hex strings +# * list-view: OFFSET + SIZE per slot (strings when 64-bit) +# * REE: no buffers; children run_ends + values +# * unions: TYPE_ID (+ OFFSET for dense); no VALIDITY +# * interval MONTH_DAY_NANO: {months, days, nanoseconds} objects; +# DAY_TIME: {days, milliseconds}; YEAR_MONTH: ints +# * dictionaries: top-level `dictionaries: [{id, data: {count, columns}}]`, +# field carries `dictionary: {id, indexType, isOrdered}` and its `type` is +# the VALUE type +# * metadata: [{key, value}] lists on schema and fields +# ============================================================================= + +module ArrowJSON + +using JSON +using ..ArrowCore +const AC = ArrowCore + +# --- type descriptors ------------------------------------------------------- + +_timeunit(s) = + s == "SECOND" ? AC.SECOND : + s == "MILLISECOND" ? AC.MILLISECOND : + s == "MICROSECOND" ? AC.MICROSECOND : + s == "NANOSECOND" ? AC.NANOSECOND : error("unknown time unit $s") +_timeunitname(u) = + u == AC.SECOND ? "SECOND" : + u == AC.MILLISECOND ? "MILLISECOND" : u == AC.MICROSECOND ? "MICROSECOND" : "NANOSECOND" + +function fromjsontype(t::AbstractDict)::ArrowType + n = t["name"] + n == "null" && return NullType() + n == "bool" && return BoolType() + n == "int" && return IntType(Int(t["bitWidth"]), Bool(t["isSigned"])) + n == "floatingpoint" && return FloatType( + t["precision"] == "HALF" ? 16 : t["precision"] == "SINGLE" ? 32 : 64, + ) + n == "utf8" && return Utf8Type(false) + n == "largeutf8" && return Utf8Type(true) + n == "binary" && return BinaryType(false) + n == "largebinary" && return BinaryType(true) + n == "utf8view" && return ViewType(true) + n == "binaryview" && return ViewType(false) + n == "fixedsizebinary" && return FixedSizeBinaryType(Int(t["byteWidth"])) + n == "decimal" && return DecimalType( + Int(t["precision"]), + Int(t["scale"]), + Int(get(t, "bitWidth", 128)), + ) + n == "date" && return DateType(t["unit"] == "DAY" ? AC.DAY : AC.MILLISECOND_DATE) + n == "time" && return TimeType(_timeunit(t["unit"]), Int(t["bitWidth"])) + n == "timestamp" && return TimestampType( + _timeunit(t["unit"]), + haskey(t, "timezone") ? String(t["timezone"]) : nothing, + ) + n == "duration" && return DurationType(_timeunit(t["unit"])) + n == "interval" && return IntervalType( + t["unit"] == "YEAR_MONTH" ? AC.YEAR_MONTH : + t["unit"] == "DAY_TIME" ? AC.DAY_TIME : AC.MONTH_DAY_NANO, + ) + n == "list" && return ListType(false) + n == "largelist" && return ListType(true) + n == "listview" && return ListViewType(false) + n == "largelistview" && return ListViewType(true) + n == "fixedsizelist" && return FixedSizeListType(Int(t["listSize"])) + n == "struct" && return StructType() + n == "map" && return MapType(Bool(get(t, "keysSorted", false))) + n == "union" && return UnionType( + t["mode"] == "SPARSE" ? AC.SparseMode : AC.DenseMode, + Int8[Int8(x) for x in t["typeIds"]], + ) + n == "runendencoded" && return RunEndEncodedType() + error("arrowjson: unmapped type $n") +end + +function tojsontype(t::ArrowType) + t isa NullType && return Dict("name" => "null") + t isa BoolType && return Dict("name" => "bool") + t isa IntType && + return Dict("name" => "int", "bitWidth" => t.bits, "isSigned" => t.signed) + t isa FloatType && return Dict( + "name" => "floatingpoint", + "precision" => t.bits == 16 ? "HALF" : t.bits == 32 ? "SINGLE" : "DOUBLE", + ) + t isa Utf8Type && return Dict("name" => t.large ? "largeutf8" : "utf8") + t isa BinaryType && return Dict("name" => t.large ? "largebinary" : "binary") + t isa ViewType && return Dict("name" => t.utf8 ? "utf8view" : "binaryview") + t isa FixedSizeBinaryType && + return Dict("name" => "fixedsizebinary", "byteWidth" => t.nbytes) + t isa DecimalType && return Dict( + "name" => "decimal", + "precision" => t.precision, + "scale" => t.scale, + "bitWidth" => t.bits, + ) + t isa DateType && + return Dict("name" => "date", "unit" => t.unit == AC.DAY ? "DAY" : "MILLISECOND") + t isa TimeType && + return Dict("name" => "time", "unit" => _timeunitname(t.unit), "bitWidth" => t.bits) + if t isa TimestampType + d = Dict{String,Any}("name" => "timestamp", "unit" => _timeunitname(t.unit)) + t.timezone === nothing || (d["timezone"] = t.timezone) + return d + end + t isa DurationType && return Dict("name" => "duration", "unit" => _timeunitname(t.unit)) + t isa IntervalType && return Dict( + "name" => "interval", + "unit" => + t.unit == AC.YEAR_MONTH ? "YEAR_MONTH" : + t.unit == AC.DAY_TIME ? "DAY_TIME" : "MONTH_DAY_NANO", + ) + t isa ListType && return Dict("name" => t.large ? "largelist" : "list") + t isa ListViewType && return Dict("name" => t.large ? "largelistview" : "listview") + t isa FixedSizeListType && + return Dict("name" => "fixedsizelist", "listSize" => t.listsize) + t isa StructType && return Dict("name" => "struct") + t isa MapType && return Dict("name" => "map", "keysSorted" => t.keyssorted) + t isa UnionType && return Dict( + "name" => "union", + "mode" => t.mode == AC.SparseMode ? "SPARSE" : "DENSE", + "typeIds" => Int.(t.typeids), + ) + t isa RunEndEncodedType && return Dict("name" => "runendencoded") + error("arrowjson: unmapped descriptor $(AC.descriptorname(t))") +end + +_metadict(m) = + m === nothing ? nothing : + Dict{String,String}(String(kv["key"]) => String(kv["value"]) for kv in m) +_metalist(m) = + m === nothing ? nothing : + [Dict("key" => k, "value" => v) for (k, v) in sort!(collect(m); by=first)] + +""" +Parse one JSON field into a Core `Field`. Dictionary-encoded fields become +`DictionaryType(indextype, valuetype, ordered)`; the JSON id is recorded in +`dictids` (an adapter-side table, exactly as the IPC adapter keeps ids). +""" +function fromjsonfield(f::AbstractDict, dictids::IdDict{Field,Int64})::Field + children = Field[fromjsonfield(c, dictids) for c in get(f, "children", Any[])] + t = fromjsontype(f["type"]) + meta = _metadict(get(f, "metadata", nothing)) + if haskey(f, "dictionary") + d = f["dictionary"] + idx = fromjsontype(d["indexType"])::IntType + cf = Field( + String(f["name"]), + DictionaryType(idx, t, Bool(get(d, "isOrdered", false))); + nullable=Bool(f["nullable"]), + metadata=meta, + children=children, + ) + dictids[cf] = Int64(d["id"]) + return cf + end + return Field( + String(f["name"]), + t; + nullable=Bool(f["nullable"]), + metadata=meta, + children=children, + ) +end + +function tojsonfield(f::Field, dictids::IdDict{Field,Int64}) + t = f.type + d = Dict{String,Any}( + "name" => f.name, + "nullable" => f.nullable, + "children" => Any[tojsonfield(c, dictids) for c in f.children], + ) + if t isa DictionaryType + d["type"] = tojsontype(t.valuetype) + d["dictionary"] = Dict( + "id" => dictids[f], + "indexType" => tojsontype(t.indextype), + "isOrdered" => t.ordered, + ) + else + d["type"] = tojsontype(t) + end + f.metadata === nothing || (d["metadata"] = _metalist(f.metadata)) + return d +end + +# --- values: JSON -> buffers -------------------------------------------------- + +_hex(bytes) = uppercase(bytes2hex(bytes)) +_unhex(s::AbstractString) = hex2bytes(s) +_i64(x) = x isa AbstractString ? parse(Int64, x) : Int64(x) + +function _validity(col, n::Int) + v = get(col, "VALIDITY", nothing) + (v === nothing || n == 0) && return BufferSlice() + bytes = zeros(UInt8, cld(n, 8)) + for i = 1:n + v[i] != 0 && (bytes[(i - 1) ÷ 8 + 1] |= UInt8(1) << ((i - 1) % 8)) + end + return AC._databuffer(bytes) +end + +function _bitmap(vals::AbstractVector{Bool}) + n = length(vals) + bytes = zeros(UInt8, cld(n, 8)) + for i = 1:n + vals[i] && (bytes[(i - 1) ÷ 8 + 1] |= UInt8(1) << ((i - 1) % 8)) + end + return AC._databuffer(bytes) +end + +function _intdata(t::IntType, data) + T = + t.signed ? + (t.bits == 8 ? Int8 : t.bits == 16 ? Int16 : t.bits == 32 ? Int32 : Int64) : + (t.bits == 8 ? UInt8 : t.bits == 16 ? UInt16 : t.bits == 32 ? UInt32 : UInt64) + vals = T[T(x isa AbstractString ? parse(T, x) : x) for x in data] + return AC._databuffer(vals) +end + +_decimalint(s, bits) = + bits == 32 ? Int32(parse(Int128, s)) : + bits == 64 ? Int64(parse(Int128, s)) : + bits == 128 ? parse(Int128, s) : error("unsupported decimal width $bits") + +function _decimal256bytes(s::AbstractString) + value = parse(BigInt, s) + half = BigInt(1) << 255 + -half <= value < half || + throw(ArgumentError("decimal256 value is outside the signed 256-bit range: $s")) + unsigned = value < 0 ? value + (BigInt(1) << 256) : value + bytes = Vector{UInt8}(undef, 32) + for i = 1:32 + bytes[i] = UInt8(unsigned & 0xff) + unsigned >>= 8 + end + return bytes +end + +function _decimal256value(bytes::AbstractVector{UInt8}) + length(bytes) == 32 || + throw(ArgumentError("a decimal256 value must contain exactly 32 bytes")) + value = BigInt(0) + for byte in Iterators.reverse(bytes) + value = (value << 8) | byte + end + (bytes[end] & 0x80) == 0 || (value -= BigInt(1) << 256) + return value +end + +""" +Build one Core `ArrayData` from a JSON column. `f` supplies the layout; +`dicts` resolves dictionary ids to already-built pools. +""" +function fromjsoncolumn( + f::Field, + col::AbstractDict, + dicts::Dict{Int64,ArrayData}, + dictids::IdDict{Field,Int64}, +)::ArrayData + t = f.type + n = Int(col["count"]) + data = get(col, "DATA", nothing) + validity = _validity(col, n) + nulls = + get(col, "VALIDITY", nothing) === nothing ? 0 : count(==(0), col["VALIDITY"][1:n]) + if t isa DictionaryType + idx = _intdata(t.indextype, data) + return ArrayData( + t, + n, + [validity, idx]; + dictionary=dicts[dictids[f]], + nullcount=nulls, + ) + elseif t isa NullType + return ArrayData(t, n, BufferSlice[]; nullcount=n) + elseif t isa BoolType + return ArrayData( + t, + n, + [validity, _bitmap(Bool[Bool(x) for x in data])]; + nullcount=nulls, + ) + elseif t isa IntType + return ArrayData(t, n, [validity, _intdata(t, data)]; nullcount=nulls) + elseif t isa FloatType + vals = + t.bits == 16 ? Float16[Float16(x) for x in data] : + t.bits == 32 ? Float32[Float32(x) for x in data] : + Float64[Float64(x) for x in data] + return ArrayData(t, n, [validity, AC._databuffer(vals)]; nullcount=nulls) + elseif t isa DecimalType + raw = if t.bits == 256 + bytes = UInt8[] + sizehint!(bytes, 32 * length(data)) + for value in data + append!(bytes, _decimal256bytes(String(value))) + end + bytes + else + vals = [_decimalint(String(x), t.bits) for x in data] + t.bits == 32 ? Int32.(vals) : t.bits == 64 ? Int64.(vals) : Int128.(vals) + end + return ArrayData(t, n, [validity, AC._databuffer(raw)]; nullcount=nulls) + elseif t isa DateType + vals = + t.unit == AC.DAY ? Int32[Int32(_i64(x)) for x in data] : + Int64[_i64(x) for x in data] + return ArrayData(t, n, [validity, AC._databuffer(vals)]; nullcount=nulls) + elseif t isa TimeType + vals = + t.bits == 32 ? Int32[Int32(_i64(x)) for x in data] : + Int64[_i64(x) for x in data] + return ArrayData(t, n, [validity, AC._databuffer(vals)]; nullcount=nulls) + elseif t isa TimestampType || t isa DurationType + return ArrayData( + t, + n, + [validity, AC._databuffer(Int64[_i64(x) for x in data])]; + nullcount=nulls, + ) + elseif t isa IntervalType + raw = if t.unit == AC.YEAR_MONTH + reinterpret(UInt8, Int32[Int32(_i64(x)) for x in data]) + elseif t.unit == AC.DAY_TIME + reinterpret( + UInt8, + Int32[Int32(_i64(v)) for x in data for v in (x["days"], x["milliseconds"])], + ) + else + out = UInt8[] + for x in data + append!( + out, + reinterpret( + UInt8, + Int32[Int32(_i64(x["months"])), Int32(_i64(x["days"]))], + ), + ) + append!(out, reinterpret(UInt8, Int64[_i64(x["nanoseconds"])])) + end + out + end + return ArrayData( + t, + n, + [validity, AC._databuffer(collect(UInt8, raw))]; + nullcount=nulls, + ) + elseif t isa FixedSizeBinaryType + bytes = UInt8[] + for x in data + b = _unhex(x) + length(b) == t.nbytes || error("fixedsizebinary width mismatch") + append!(bytes, b) + end + return ArrayData(t, n, [validity, AC._databuffer(bytes)]; nullcount=nulls) + elseif t isa Utf8Type || t isa BinaryType + offs = [_i64(x) for x in col["OFFSET"]] + bytes = UInt8[] + for x in data + append!(bytes, t isa Utf8Type ? codeunits(String(x)) : _unhex(x)) + end + offbuf = t.large ? AC._databuffer(Int64.(offs)) : AC._databuffer(Int32.(offs)) + return ArrayData(t, n, [validity, offbuf, AC._databuffer(bytes)]; nullcount=nulls) + elseif t isa ViewType + views = UInt8[] + for v in col["VIEWS"] + sz = Int32(v["SIZE"]) + append!(views, reinterpret(UInt8, Int32[sz])) + if haskey(v, "INLINED") + inl = + t.utf8 ? collect(codeunits(String(v["INLINED"]))) : _unhex(v["INLINED"]) + append!(views, inl) + append!(views, zeros(UInt8, 12 - length(inl))) + else + append!(views, _unhex(v["PREFIX_HEX"])) + append!( + views, + reinterpret(UInt8, Int32[Int32(v["BUFFER_INDEX"]), Int32(v["OFFSET"])]), + ) + end + end + bufs = BufferSlice[validity, AC._databuffer(views)] + for h in get(col, "VARIADIC_DATA_BUFFERS", Any[]) + b = _unhex(h) + push!(bufs, isempty(b) ? BufferSlice() : AC._databuffer(b)) + end + return ArrayData(t, n, bufs; nullcount=nulls) + elseif t isa ListType || t isa MapType + offs = [_i64(x) for x in col["OFFSET"]] + offbuf = + (t isa ListType && t.large) ? AC._databuffer(Int64.(offs)) : + AC._databuffer(Int32.(offs)) + child = fromjsoncolumn(f.children[1], col["children"][1], dicts, dictids) + return ArrayData(t, n, [validity, offbuf]; children=[child], nullcount=nulls) + elseif t isa ListViewType + offs = [_i64(x) for x in col["OFFSET"]] + sizes = [_i64(x) for x in col["SIZE"]] + ob = t.large ? AC._databuffer(Int64.(offs)) : AC._databuffer(Int32.(offs)) + sb = t.large ? AC._databuffer(Int64.(sizes)) : AC._databuffer(Int32.(sizes)) + child = fromjsoncolumn(f.children[1], col["children"][1], dicts, dictids) + return ArrayData(t, n, [validity, ob, sb]; children=[child], nullcount=nulls) + elseif t isa FixedSizeListType + child = fromjsoncolumn(f.children[1], col["children"][1], dicts, dictids) + return ArrayData(t, n, [validity]; children=[child], nullcount=nulls) + elseif t isa StructType + children = ArrayData[ + fromjsoncolumn(cf, cc, dicts, dictids) for + (cf, cc) in zip(f.children, col["children"]) + ] + return ArrayData(t, n, [validity]; children=children, nullcount=nulls) + elseif t isa UnionType + ids = AC._databuffer(Int8[Int8(x) for x in col["TYPE_ID"]]) + children = ArrayData[ + fromjsoncolumn(cf, cc, dicts, dictids) for + (cf, cc) in zip(f.children, col["children"]) + ] + if t.mode == AC.DenseMode + offs = AC._databuffer(Int32[Int32(x) for x in col["OFFSET"]]) + return ArrayData(t, n, [ids, offs]; children=children, nullcount=0) + end + return ArrayData(t, n, [ids]; children=children, nullcount=0) + elseif t isa RunEndEncodedType + children = ArrayData[ + fromjsoncolumn(cf, cc, dicts, dictids) for + (cf, cc) in zip(f.children, col["children"]) + ] + return ArrayData(t, n, BufferSlice[]; children=children, nullcount=0) + end + error("arrowjson: unmapped layout $(AC.descriptorname(t))") +end + +# --- values: Core -> JSON -------------------------------------------------------- + +_validitylist(d::ArrayData) = Int[AC.isvalid_at(d, i) ? 1 : 0 for i = 1:d.len] + +function _rawvals(d::ArrayData, ::Type{T}) where {T} + b = AC.rolebuffer(d, AC.DATA) + return T[AC.loadat(b, T, AC._slotbyteoff(d, Int64(i), sizeof(T))) for i = 1:d.len] +end + +function _offsetlist(d::ArrayData, wide::Bool) + b = AC.rolebuffer(d, AC.OFFSETS) + n = d.len + # A zero-length UNSLICED array may carry Core's canonical empty offsets + # buffer (nanoarrow and C++ write that form); the JSON column still + # declares the single terminal zero. A sliced empty array must carry the + # physical prefix its offset addresses, so it takes the loadat path (and + # its absence stays a visible error, not a fabricated zero). + b.len == 0 && n == 0 && d.offset == 0 && return wide ? Int64[0] : Int32[0] + if wide + return Int64[AC.loadat(b, Int64, (d.offset + i) * 8) for i = 0:n] + end + return Int32[AC.loadat(b, Int32, (d.offset + i) * 4) for i = 0:n] +end + +""" +Render one Core column as an integration-JSON column object. Values are +read through raw buffers (not `getvalue`) so null slots keep their physical +DATA — the gold files carry data under nulls and diff tools compare it. +""" +function tojsoncolumn(f::Field, d::ArrayData) + t = f.type + n = Int(d.len) + col = Dict{String,Any}("name" => f.name, "count" => n) + hasvalidity(t) = !(t isa NullType || t isa UnionType || t isa RunEndEncodedType) + hasvalidity(t) && (col["VALIDITY"] = _validitylist(d)) + if t isa DictionaryType + it = t.indextype + col["DATA"] = _intjson(it, d) + elseif t isa NullType + # nothing + elseif t isa BoolType + b = AC.rolebuffer(d, AC.DATA) + col["DATA"] = Bool[AC.getbit(b, AC._slotindex0(d, Int64(i))) for i = 1:n] + elseif t isa IntType + col["DATA"] = _intjson(t, d) + elseif t isa FloatType + col["DATA"] = + t.bits == 16 ? Float64.(_rawvals(d, Float16)) : + t.bits == 32 ? _rawvals(d, Float32) : _rawvals(d, Float64) + elseif t isa DecimalType + vals = if t.bits == 256 + buffer = AC.rolebuffer(d, AC.DATA) + [ + _decimal256value( + AC.slicebytes( + AC.subslice(buffer, AC._slotbyteoff(d, Int64(i), 32), 32), + ), + ) for i = 1:n + ] + else + t.bits == 32 ? _rawvals(d, Int32) : + t.bits == 64 ? _rawvals(d, Int64) : _rawvals(d, Int128) + end + col["DATA"] = string.(vals) + elseif t isa DateType + col["DATA"] = t.unit == AC.DAY ? _rawvals(d, Int32) : string.(_rawvals(d, Int64)) + elseif t isa TimeType + col["DATA"] = t.bits == 32 ? _rawvals(d, Int32) : string.(_rawvals(d, Int64)) + elseif t isa TimestampType || t isa DurationType + col["DATA"] = string.(_rawvals(d, Int64)) + elseif t isa IntervalType + b = AC.rolebuffer(d, AC.DATA) + if t.unit == AC.YEAR_MONTH + col["DATA"] = _rawvals(d, Int32) + elseif t.unit == AC.DAY_TIME + col["DATA"] = [ + Dict( + "days" => AC.loadat(b, Int32, AC._slotbyteoff(d, Int64(i), 8)), + "milliseconds" => + AC.loadat(b, Int32, AC._slotbyteoff(d, Int64(i), 8) + 4), + ) for i = 1:n + ] + else + col["DATA"] = [ + Dict( + "months" => AC.loadat(b, Int32, AC._slotbyteoff(d, Int64(i), 16)), + "days" => AC.loadat(b, Int32, AC._slotbyteoff(d, Int64(i), 16) + 4), + "nanoseconds" => string( + AC.loadat(b, Int64, AC._slotbyteoff(d, Int64(i), 16) + 8), + ), + ) for i = 1:n + ] + end + elseif t isa FixedSizeBinaryType + b = AC.rolebuffer(d, AC.DATA) + col["DATA"] = [ + _hex( + AC.slicebytes( + AC.subslice(b, AC._slotbyteoff(d, Int64(i), t.nbytes), t.nbytes), + ), + ) for i = 1:n + ] + elseif t isa Utf8Type || t isa BinaryType + offs = _offsetlist(d, t.large) + col["OFFSET"] = t.large ? string.(offs) : offs + b = AC.rolebuffer(d, AC.DATA) + col["DATA"] = [ + begin + lo, hi = Int64(offs[i]), Int64(offs[i + 1]) + bytes = hi > lo ? AC.slicebytes(AC.subslice(b, lo, hi - lo)) : UInt8[] + t isa Utf8Type ? String(bytes) : _hex(bytes) + end for i = 1:n + ] + elseif t isa ViewType + views = AC.rolebuffer(d, AC.VIEWS) + entries = Any[] + for i = 1:n + base = AC._viewbase(d, Int64(i)) + sz = AC.loadat(views, Int32, base) + if sz <= AC.VIEW_INLINE_MAX + inl = AC.slicebytes(AC.subslice(views, base + 4, Int64(sz))) + push!( + entries, + Dict("SIZE" => sz, "INLINED" => t.utf8 ? String(inl) : _hex(inl)), + ) + else + push!( + entries, + Dict( + "SIZE" => sz, + "PREFIX_HEX" => + _hex(AC.slicebytes(AC.subslice(views, base + 4, 4))), + "BUFFER_INDEX" => AC.loadat(views, Int32, base + 8), + "OFFSET" => AC.loadat(views, Int32, base + 12), + ), + ) + end + end + col["VIEWS"] = entries + col["VARIADIC_DATA_BUFFERS"] = [_hex(AC.slicebytes(b)) for b in d.buffers[3:end]] + elseif t isa ListType || t isa MapType + wide = t isa ListType && t.large + offs = _offsetlist(d, wide) + col["OFFSET"] = wide ? string.(offs) : offs + col["children"] = Any[tojsoncolumn(f.children[1], d.children[1])] + elseif t isa ListViewType + ob = AC.rolebuffer(d, AC.ELEMENT_OFFSETS) + sb = AC.rolebuffer(d, AC.SIZES) + w = t.large ? 8 : 4 + offs = [ + t.large ? AC.loadat(ob, Int64, AC._slotbyteoff(d, Int64(i), w)) : + AC.loadat(ob, Int32, AC._slotbyteoff(d, Int64(i), w)) for i = 1:n + ] + sizes = [ + t.large ? AC.loadat(sb, Int64, AC._slotbyteoff(d, Int64(i), w)) : + AC.loadat(sb, Int32, AC._slotbyteoff(d, Int64(i), w)) for i = 1:n + ] + col["OFFSET"] = t.large ? string.(offs) : offs + col["SIZE"] = t.large ? string.(sizes) : sizes + col["children"] = Any[tojsoncolumn(f.children[1], d.children[1])] + elseif t isa FixedSizeListType || t isa StructType + col["children"] = + Any[tojsoncolumn(cf, cd) for (cf, cd) in zip(f.children, d.children)] + elseif t isa UnionType + ids = AC.rolebuffer(d, AC.TYPE_IDS) + col["TYPE_ID"] = Int[AC.loadat(ids, Int8, AC._slotindex0(d, Int64(i))) for i = 1:n] + if t.mode == AC.DenseMode + ob = AC.rolebuffer(d, AC.ELEMENT_OFFSETS) + col["OFFSET"] = + Int32[AC.loadat(ob, Int32, AC._slotbyteoff(d, Int64(i), 4)) for i = 1:n] + end + col["children"] = + Any[tojsoncolumn(cf, cd) for (cf, cd) in zip(f.children, d.children)] + elseif t isa RunEndEncodedType + col["children"] = + Any[tojsoncolumn(cf, cd) for (cf, cd) in zip(f.children, d.children)] + else + error("arrowjson: unmapped layout $(AC.descriptorname(t))") + end + return col +end + +function _intjson(t::IntType, d::ArrayData) + if t.bits == 64 + return t.signed ? string.(_rawvals(d, Int64)) : string.(_rawvals(d, UInt64)) + end + return t.signed ? + ( + t.bits == 8 ? _rawvals(d, Int8) : + t.bits == 16 ? _rawvals(d, Int16) : _rawvals(d, Int32) + ) : + ( + t.bits == 8 ? _rawvals(d, UInt8) : + t.bits == 16 ? _rawvals(d, UInt16) : _rawvals(d, UInt32) + ) +end + +# --- documents --------------------------------------------------------------------- + +""" + fromjson(doc) -> (schema::Schema, batches::Vector{RecordBatch}, dictids) + +Parse an integration-JSON document (already `JSON.parse`d) into Core values. +Dictionaries are built first (in id order) so batch columns can reference +them; the returned `dictids` maps each dictionary-typed Field to its JSON id +for writers that must preserve ids. +""" +function fromjson(doc::AbstractDict) + sch, batches, dictids, _ = _fromjson(doc) + return sch, batches, dictids +end + +# Internal form used by the value comparator. The dictionary arrays are part +# of the logical document even when no batch index currently references every +# pool value, so the comparator must retain them separately from the batches. +function _fromjson(doc::AbstractDict) + dictids = IdDict{Field,Int64}() + fields = Field[fromjsonfield(f, dictids) for f in doc["schema"]["fields"]] + sch = Schema( + fields; + metadata=_metadict(get(doc["schema"], "metadata", nothing)), + endianness=AC.LittleEndian, + ) + dicts = Dict{Int64,ArrayData}() + # dictionaries may depend on other dictionaries (nested); resolve by + # repeated passes until all build + pending = collect(get(doc, "dictionaries", Any[])) + valuefield = Dict{Int64,Field}() + function walk(f::Field) + if f.type isa DictionaryType + valuefield[dictids[f]] = AC.dictvaluefield(f, f.type) + end + foreach(walk, f.children) + f.type isa DictionaryType && walk(AC.dictvaluefield(f, f.type)) + end + foreach(walk, fields) + while !isempty(pending) + progressed = false + for (k, entry) in enumerate(pending) + id = Int64(entry["id"]) + vf = valuefield[id] + try + col = entry["data"]["columns"][1] + dicts[id] = fromjsoncolumn(vf, col, dicts, dictids) + deleteat!(pending, k) + progressed = true + break + catch e + e isa KeyError || rethrow() + end + end + progressed || error("arrowjson: unresolvable dictionary dependencies") + end + batches = AC.RecordBatch[] + for b in doc["batches"] + cols = ArrayData[ + fromjsoncolumn(f, c, dicts, dictids) for (f, c) in zip(fields, b["columns"]) + ] + push!(batches, AC.RecordBatch(sch, cols, Int(b["count"]))) + end + return sch, batches, dictids, dicts +end + +""" + tojson(schema, batches; dictids) -> Dict + +Render Core values as an integration-JSON document. Dictionary pools are +emitted once per id from the first batch that carries them (the file-format +convention; replacement streams need per-batch dictionaries and are outside +this writer). +""" +function tojson( + sch::Schema, + batches::AbstractVector{AC.RecordBatch}; + dictids::IdDict{Field,Int64}=IdDict{Field,Int64}(), +) + if isempty(dictids) + next = Int64(0) + function assign(f::Field) + if f.type isa DictionaryType + dictids[f] = next + next += 1 + end + foreach(assign, f.children) + end + foreach(assign, sch.fields) + end + doc = Dict{String,Any}() + schemadoc = + Dict{String,Any}("fields" => Any[tojsonfield(f, dictids) for f in sch.fields]) + sch.metadata === nothing || (schemadoc["metadata"] = _metalist(sch.metadata)) + doc["schema"] = schemadoc + dictdocs = Any[] + seen = Set{Int64}() + function collectpools(f::Field, d::ArrayData) + if f.type isa DictionaryType + id = dictids[f] + if !(id in seen) + push!(seen, id) + vf = AC.dictvaluefield(f, f.type) + pool = d.dictionary::ArrayData + push!( + dictdocs, + Dict( + "id" => id, + "data" => Dict( + "count" => Int(pool.len), + "columns" => Any[tojsoncolumn(vf, pool)], + ), + ), + ) + collectpools(vf, pool) + end + return + end + for (cf, cd) in zip(f.children, d.children) + collectpools(cf, cd) + end + end + for b in batches, (f, d) in zip(sch.fields, b.columns) + collectpools(f, d) + end + isempty(dictdocs) || (doc["dictionaries"] = dictdocs) + doc["batches"] = Any[ + Dict( + "count" => Int(b.nrows), + "columns" => + Any[tojsoncolumn(f, d) for (f, d) in zip(sch.fields, b.columns)], + ) for b in batches + ] + return doc +end + +end # module ArrowJSON diff --git a/conformance/cdata_oracle.jl b/conformance/cdata_oracle.jl new file mode 100644 index 00000000..0d84c43b --- /dev/null +++ b/conformance/cdata_oracle.jl @@ -0,0 +1,480 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# ============================================================================= +# C Data / C Stream oracle: OUR C-interface structures through pyarrow, in +# one process, over the whole gold data matrix. +# +# julia conformance/run.jl cdata +# +# `oracle.jl` proves our IPC BYTES against pyarrow and nanoarrow. This suite +# proves our C DATA INTERFACE and C STREAM INTERFACE the same way: pyarrow +# runs in-process (PythonCall) so real ArrowSchema/ArrowArray/ArrowArrayStream +# pointers cross the boundary in both directions with real ownership moves. +# The gold corpus supplies the data (every layout the format defines), parsed +# to Core through the corpus's JSON reader; comparison is the corpus's own +# value-level document comparison. For every gold family: +# +# ours→pyarrow(C)→ours each batch is exported as one struct-typed +# ArrowArray + ArrowSchema; pyarrow imports it as a +# RecordBatch, FULLY validates it (its independent +# judgment of our export), and exports it back; our +# importer reads pyarrow's structures. Values must +# equal the gold JSON. Proves both directions of the +# C Data interface, including field names and +# nullability. (Metadata VALUES ride along; exact +# metadata ORDER and duplicate keys are the +# synthetic sentinel's job below — the corpus +# comparison normalizes them away.) +# pyarrow-native→ours pyarrow rebuilds the batch through its OWN IPC +# reader (its allocator, its buffer choices, its +# dictionary memo) and exports that; we import it. +# Proves our importer against pyarrow-produced memory, +# not just our own memory reflected back. +# pyarrow slice→ours pyarrow exports a SLICED batch (nonzero per-node +# offsets); we import it and compare against our own +# logical slice of the source. Proves offset handling +# on both sides. +# ours→pyarrow(stream)→ours +# the whole family as one C stream: pyarrow imports +# our ArrowArrayStream as a RecordBatchReader and +# re-exports it; we pull batches through pyarrow's +# stream (a re-entrant Julia→C→Julia pull path). +# +# The suite ends by draining the export registries: every C structure handed +# to pyarrow must have been released back exactly once. +# +# The interpreter is ARROW_CDATA_ORACLE_PYTHON — a Python with pyarrow +# importable; the conformance image (conformance/Dockerfile) sets it. The +# parent process re-launches this file as a child with PythonCall bound to +# that interpreter (PythonCall reads its interpreter at load, so the parent +# never loads it). +# ============================================================================= + +import Arrow + +if !isdefined(@__MODULE__, :ConformanceSupport) + include(joinpath(@__DIR__, "ConformanceSupport.jl")) +end +using .ConformanceSupport: + ArrowJSON, + DEFAULT_CORPUS, + Verdict, + corebatchdiffs, + documentcheck, + errorverdict, + readjson + +const AC = Arrow.ArrowCore + +const _CDATA_ORACLE_CHILD = "--child" +const _CDATA_ORACLE_STANDALONE = abspath(PROGRAM_FILE) == abspath(@__FILE__) + +# --- parent: environment preparation + relaunch ------------------------------ + +function _oracle_python() + py = get(ENV, "ARROW_CDATA_ORACLE_PYTHON", "") + isempty(py) && error( + "ARROW_CDATA_ORACLE_PYTHON is not set: run this suite " * + "through `julia conformance/run.jl cdata` (the " * + "conformance image), or point ARROW_CDATA_ORACLE_PYTHON at a Python with pyarrow", + ) + return py +end + +function _oracle_parent(args) + py = _oracle_python() + ver = readchomp(`$py -c "import pyarrow; print(pyarrow.__version__)"`) + println("cdata oracle: pyarrow $ver at $py") + env = copy(ENV) + env["JULIA_PYTHONCALL_EXE"] = py + env["JULIA_CONDAPKG_BACKEND"] = "Null" + cmd = `$(Base.julia_cmd()) --project=$(dirname(Base.active_project())) + --startup-file=no $(@__FILE__) $_CDATA_ORACLE_CHILD $args` + proc = run(ignorestatus(setenv(cmd, env))) + exit(proc.exitcode) +end + +function _run_cdata_oracle(args) + corpus = isempty(args) ? DEFAULT_CORPUS : args[1] + verdicts = runcdataoracle(corpus) + nfail = ConformanceSupport.report("C Data / C Stream oracle", verdicts) + println() + println("PASS families by check:") + for check in unique(v.check for v in verdicts if v.check != "all") + n = count(v -> v.check == check && v.status == :pass, verdicts) + m = count(v -> v.check == check && v.status != :skip, verdicts) + println(" ", rpad(check, 28), n, "/", m) + end + exit(nfail == 0 ? 0 : 1) +end + +# Everything below is only defined in the PythonCall child. +if _CDATA_ORACLE_STANDALONE && length(ARGS) >= 1 && ARGS[1] == _CDATA_ORACLE_CHILD + using PythonCall + + const pa = pyimport("pyarrow") + const paipc = pyimport("pyarrow.ipc") + const CS = Arrow.CArrowSchema + const CA = Arrow.CArrowArray + const CAS = Arrow.CArrowArrayStream + + # One RecordBatch is one struct-typed column: the schema-level metadata rides + # the struct Field, the fields are its children, the columns its child arrays. + function _structwrap(sch::AC.Schema, columns, nrows::Integer) + f = AC.Field( + "", + AC.StructType(); + nullable=false, + metadata=sch.metadata, + children=collect(AC.Field, sch.fields), + ) + d = AC.ArrayData( + AC.StructType(), + nrows, + [AC.BufferSlice()]; + children=collect(AC.ArrayData, columns), + nullcount=0, + ) + return f, d + end + + function _unwrap(f2::AC.Field, d2::AC.ArrayData) + sch2 = AC.Schema( + collect(AC.Field, f2.children); + metadata=f2.metadata, + endianness=AC.LittleEndian, + ) + return sch2, AC.RecordBatch(sch2, collect(AC.ArrayData, d2.children), d2.len) + end + + # Export one struct-wrapped batch to pyarrow; pyarrow validates fully. + function _to_pyarrow(sch::AC.Schema, b::AC.RecordBatch) + f, d = _structwrap(sch, b.columns, b.nrows) + sp, ap = Arrow.to_c_data(f, d) + pyb = pa.RecordBatch._import_from_c(UInt(ap), UInt(sp)) + unsafe_load(sp).release == C_NULL || + error("pyarrow did not mark the imported schema released") + unsafe_load(ap).release == C_NULL || + error("pyarrow did not mark the imported array released") + pyb.validate(full=true) + return pyb + end + + # Import a pyarrow RecordBatch through the C interface into Core. + function _from_pyarrow(pyb) + aref = Ref{CA}() + sref = Ref{CS}() + return GC.@preserve aref sref begin + ap = Base.unsafe_convert(Ptr{CA}, aref) + sp = Base.unsafe_convert(Ptr{CS}, sref) + pyb._export_to_c(UInt(ap), UInt(sp)) + f2, d2 = Arrow.from_c_data(sp, ap) + _unwrap(f2, d2) + end + end + + # pyarrow rebuilds a batch through its own IPC reader: pyarrow-owned memory. + function _pyarrow_rebuild(pyb) + sink = pa.BufferOutputStream() + w = paipc.new_stream(sink, pyb.schema) + w.write_batch(pyb) + w.close() + r = paipc.open_stream(sink.getvalue()) + return r.read_next_batch() + end + + _releasebatch!(b::AC.RecordBatch) = + isempty(b.columns) || Arrow.release!(b.columns[1].owner::Arrow.ForeignOwner) + + function runcdatafamily(dir::String, family::String, verdicts::Vector{Verdict}) + gold = readjson(joinpath(dir, family * ".json.gz")) + sch, batches, _ = ArrowJSON.fromjson(gold) + + # 1. ours → pyarrow (validate full) → ours ; 2. pyarrow-native → ours + check1 = "ours→pyarrow(C)→ours" + check2 = "pyarrow-native→ours" + sch1 = sch + b1s = AC.RecordBatch[] + sch2 = sch + b2s = AC.RecordBatch[] + ok1 = ok2 = true + for b in batches + local pyb + try + pyb = _to_pyarrow(sch, b) + catch e + ok1 && push!(verdicts, errorverdict(family, check1, e)) + ok2 && push!(verdicts, errorverdict(family, check2, e)) + ok1 = ok2 = false + break + end + if ok1 + try + sch1, b1 = _from_pyarrow(pyb) + push!(b1s, b1) + catch e + push!(verdicts, errorverdict(family, check1, e)) + ok1 = false + end + end + if ok2 + try + native = _pyarrow_rebuild(pyb) + sch2, b2 = _from_pyarrow(native) + push!(b2s, b2) + PythonCall.pydel!(native) + catch e + push!(verdicts, errorverdict(family, check2, e)) + ok2 = false + end + end + PythonCall.pydel!(pyb) + end + ok1 && push!(verdicts, documentcheck(family, check1, gold) do + ArrowJSON.tojson(sch1, b1s) + end) + ok2 && push!(verdicts, documentcheck(family, check2, gold) do + ArrowJSON.tojson(sch2, b2s) + end) + foreach(_releasebatch!, b1s) + foreach(_releasebatch!, b2s) + + # 3. pyarrow slice → ours, against our own logical slice + check3 = "pyarrow slice→ours" + sliceable = [b for b in batches if b.nrows >= 3] + if isempty(sliceable) + push!(verdicts, Verdict(family, check3, :skip, "no batch with ≥3 rows")) + else + try + diffs = String[] + for b in sliceable + off, len = 1, b.nrows - 2 + pyb = _to_pyarrow(sch, b) + sliced = pyb.slice(off, len) + schs, bs = _from_pyarrow(sliced) + PythonCall.pydel!(sliced) + PythonCall.pydel!(pyb) + append!(diffs, corebatchdiffs(bs, b; expectedoffset=off, rows=len)) + _releasebatch!(bs) + isempty(diffs) || break + end + push!( + verdicts, + Verdict( + family, + check3, + isempty(diffs) ? :pass : :fail, + isempty(diffs) ? "" : first(diffs), + ), + ) + catch e + push!(verdicts, errorverdict(family, check3, e)) + end + end + + # 4. ours → pyarrow RecordBatchReader → ours, over the C stream interface + check4 = "ours→pyarrow(stream)→ours" + try + outref = Ref{CAS}() + inref = Ref{CAS}() + sch4, b4s = GC.@preserve outref inref begin + outp = Base.unsafe_convert(Ptr{CAS}, outref) + Arrow.export_stream!(outp, sch, batches) + reader = pa.RecordBatchReader._import_from_c(UInt(outp)) + unsafe_load(outp).release == C_NULL || + error("pyarrow did not mark the imported stream released") + inp = Base.unsafe_convert(Ptr{CAS}, inref) + reader._export_to_c(UInt(inp)) + s = Arrow.from_c_stream(inp) + got = AC.RecordBatch[] + while (b = AC.nextbatch!(s)) !== nothing + push!(got, b) + end + Arrow.release!(s) + PythonCall.pydel!(reader) + s.schema, got + end + push!(verdicts, documentcheck(family, check4, gold) do + ArrowJSON.tojson(sch4, b4s) + end) + foreach(_releasebatch!, b4s) + catch e + push!(verdicts, errorverdict(family, check4, e)) + end + return + end + + # Metadata SEQUENCE fidelity: the corpus comparison normalizes metadata (it + # sorts pairs and the integration JSON collapses duplicate keys into a Dict), + # so it cannot see order or duplicate-key corruption. This synthetic sentinel + # compares the ORDERED pair sequences exactly — schema level, a leaf field, + # a nested child, and a dictionary field — through both the C data path and + # the C stream path. + const SENTINEL_META = ["z" => "1", "a" => "2", "z" => "3", "m" => ""] + + function _metadata_sentinel!(verdicts::Vector{Verdict}) + lf, ld = AC.fromjulia("leaf", Int64[1, 2, 3]) + leaf = AC.Field(lf.name, lf.type; nullable=lf.nullable, metadata=SENTINEL_META) + cf, cd = AC.fromjulia("item", Int64[7, 8, 9]) + child = AC.Field( + cf.name, + cf.type; + nullable=cf.nullable, + metadata=reverse(SENTINEL_META), + ) + lstf, lstd = AC.fromjulia("lst", [Int64[7], Int64[8], Int64[9]]) + lst = AC.Field( + lstf.name, + lstf.type; + nullable=lstf.nullable, + metadata=SENTINEL_META, + children=[child], + ) + df0, dd = AC.fromjulia_dict("dict", ["lo", "hi"], [0, 1, 0]) + dict = AC.Field( + df0.name, + df0.type; + nullable=df0.nullable, + metadata=SENTINEL_META, + children=collect(AC.Field, df0.children), + ) + sch = AC.Schema(AC.Field[leaf, lst, dict]; metadata=reverse(SENTINEL_META)) + b = AC.RecordBatch(sch, AC.ArrayData[ld, lstd, dd], 3) + seqs(schema) = Any[ + collect(schema.metadata), + [collect(f.metadata) for f in schema.fields]..., + collect(schema.fields[2].children[1].metadata), + ] + want = seqs(sch) + for (check, roundtrip) in ( + ("metadata sequence (C data)", () -> begin + pyb = _to_pyarrow(sch, b) + s2, b2 = _from_pyarrow(pyb) + PythonCall.pydel!(pyb) + _releasebatch!(b2) + s2 + end), + ( + "metadata sequence (C stream)", + () -> begin + outref = Ref{CAS}() + inref = Ref{CAS}() + GC.@preserve outref inref begin + outp = Base.unsafe_convert(Ptr{CAS}, outref) + Arrow.export_stream!(outp, sch, AC.RecordBatch[b]) + reader = pa.RecordBatchReader._import_from_c(UInt(outp)) + inp = Base.unsafe_convert(Ptr{CAS}, inref) + reader._export_to_c(UInt(inp)) + st = Arrow.from_c_stream(inp) + got = AC.nextbatch!(st) + got === nothing || _releasebatch!(got) + Arrow.release!(st) + PythonCall.pydel!(reader) + st.schema + end + end, + ), + ) + try + got = seqs(roundtrip()) + push!( + verdicts, + Verdict( + "(sentinel)", + check, + got == want ? :pass : :fail, + got == want ? "" : + "ordered metadata sequences differ: $got vs $want", + ), + ) + catch e + push!(verdicts, errorverdict("(sentinel)", check, e)) + end + end + return + end + + # The C interfaces carry data, not IPC framing, so a family's JSON is the + # same test whichever corpus version directory it lives in: run each family + # once, from the newest directory that has it. + function _familydirs(corpus::String) + root = joinpath(corpus, "data", "arrow-ipc-stream", "integration") + isdir(root) || error("corpus not found at $root (set ARROW_TESTING_DIR)") + vdirs = filter(d -> isdir(joinpath(root, d)), readdir(root)) + sort!(vdirs; by=d -> (startswith(d, "cpp-"), d), rev=true) + chosen = Pair{String,String}[] # family => version dir + seen = Set{String}() + for v in vdirs + dir = joinpath(root, v) + for f in sort!(readdir(dir)) + endswith(f, ".json.gz") || continue + fam = replace(f, r"\.json\.gz$" => "") + fam in seen && continue + push!(seen, fam) + push!(chosen, fam => v) + end + end + return root, chosen + end + + function runcdataoracle(corpus::String=DEFAULT_CORPUS) + root, chosen = _familydirs(corpus) + println( + "cdata oracle: pyarrow ", + pa.__version__, + " over ", + length(chosen), + " families", + ) + verdicts = Verdict[] + for (fam, v) in chosen + before = length(verdicts) + runcdatafamily(joinpath(root, v), fam, verdicts) + for i = (before + 1):length(verdicts) + vd = verdicts[i] + verdicts[i] = Verdict(v * "/" * vd.family, vd.check, vd.status, vd.detail) + end + end + isempty(verdicts) && + error("arrow-testing corpus contains no runnable C Data families under $root") + _metadata_sentinel!(verdicts) + # Every structure handed to pyarrow must have come back exactly once. + PythonCall.GC.gc() + GC.gc() + GC.gc() + Arrow.reap!() + leaked = length(Arrow.EXPORT_REGISTRY) + Arrow._stream_registry_count() + push!( + verdicts, + Verdict( + "(all)", + "export registries drained", + leaked == 0 ? :pass : :fail, + leaked == 0 ? "" : "$leaked export root(s) still registered", + ), + ) + return verdicts + end +end # child definitions + +if _CDATA_ORACLE_STANDALONE + if length(ARGS) >= 1 && ARGS[1] == _CDATA_ORACLE_CHILD + _run_cdata_oracle(ARGS[2:end]) + else + _oracle_parent(ARGS) + end +end diff --git a/conformance/corpus.jl b/conformance/corpus.jl new file mode 100644 index 00000000..f13b6215 --- /dev/null +++ b/conformance/corpus.jl @@ -0,0 +1,165 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# ============================================================================= +# Corpus conformance: the apache/arrow-testing integration gold files. +# +# julia conformance/run.jl corpus # in the conformance image +# +# For every gold family (a `.json.gz` with sibling `.stream` and +# `.arrow_file`), run the four checks that make up cross-implementation +# conformance and report a per-file verdict: +# +# JSON→Core→JSON parse the gold JSON into Core, render it back, compare +# documents (proves the JSON mapping is lossless); +# gold stream→JSON read the gold .stream through our IPC reader, render as +# JSON, compare to the gold JSON (proves READ conformance); +# gold file→JSON same through the file reader (footer path); +# JSON→our IPC→gold values +# write the gold JSON's Core values with OUR writer (stream +# + file), read back through OUR reader, compare values to +# the gold JSON (proves WRITE round-trip). Byte-identity +# with the gold IPC is NOT required — writers legitimately +# differ in padding, dictionary ordering, and metadata. +# +# Comparison is value-level over the JSON documents (schema, dictionaries, +# batches) with numeric normalization (floats compared EXACTLY after +# half/single columns are canonicalized through their physical precision; +# 64-bit strings vs numbers unified). Skips are explicit and categorized so +# the report reads as coverage, not silence. +# ============================================================================= + +import Arrow + +if !isdefined(@__MODULE__, :ConformanceSupport) + include(joinpath(@__DIR__, "ConformanceSupport.jl")) +end +using .ConformanceSupport: + ArrowJSON, + DEFAULT_CORPUS, + Verdict, + documentcheck, + familyskipreason, + filedocument, + goldipcskipreason, + readjson, + streamdocument + +function runfamily(dir::String, family::String, verdicts::Vector{Verdict}) + why = familyskipreason(family, basename(dir)) + if !isempty(why) + push!(verdicts, Verdict(family, "all", :skip, why)) + return + end + + gold = readjson(joinpath(dir, family * ".json.gz")) + + # 1. JSON -> Core -> JSON + check = "json→core→json" + push!(verdicts, documentcheck(family, check, gold) do + sch, batches, dictids = ArrowJSON.fromjson(gold) + ArrowJSON.tojson(sch, batches; dictids=dictids) + end) + + # 2. gold stream -> JSON ; 3. gold file -> JSON + for (check, path, reader) in ( + ("gold stream→json", joinpath(dir, family * ".stream"), streamdocument), + ("gold file→json", joinpath(dir, family * ".arrow_file"), filedocument), + ) + why = goldipcskipreason(family, basename(dir)) + if !isempty(why) + push!(verdicts, Verdict(family, check, :skip, why)) + continue + end + if !isfile(path) + push!(verdicts, Verdict(family, check, :skip, "no gold file")) + continue + end + push!(verdicts, documentcheck(family, check, gold) do + reader(read(path)) + end) + end + + # 4. JSON -> our IPC (stream + file) -> our reader -> JSON vs gold + for (check, writer, reader) in ( + ( + "json→our stream→json", + (s, b, ids) -> Arrow.writestream(s, b; dictids=ids), + streamdocument, + ), + ( + "json→our file→json", + (s, b, ids) -> Arrow.writefile(s, b; dictids=ids), + filedocument, + ), + ) + push!(verdicts, documentcheck(family, check, gold) do + sch, batches, dictids = ArrowJSON.fromjson(gold) + reader(writer(sch, batches, dictids)) + end) + end + return +end + +function runcorpus(corpus::String=DEFAULT_CORPUS; versions=nothing) + isempty(corpus) && error( + "ARROW_TESTING_DIR is not set: run this suite " * + "through `julia conformance/run.jl corpus`", + ) + root = joinpath(corpus, "data", "arrow-ipc-stream", "integration") + isdir(root) || error("corpus not found at $root (set ARROW_TESTING_DIR)") + verdicts = Verdict[] + vdirs = + versions === nothing ? filter(d -> isdir(joinpath(root, d)), readdir(root)) : + versions + for v in sort(vdirs) + dir = joinpath(root, v) + families = sort!( + unique!([ + replace(f, r"\.json\.gz$" => "") for + f in readdir(dir) if endswith(f, ".json.gz") + ]), + ) + for fam in families + before = length(verdicts) + runfamily(dir, fam, verdicts) + for i = (before + 1):length(verdicts) + vd = verdicts[i] + verdicts[i] = Verdict(v * "/" * vd.family, vd.check, vd.status, vd.detail) + end + end + end + (isempty(verdicts) || all(v -> v.status == :skip, verdicts)) && + error("arrow-testing corpus contains no runnable cases under $root") + return verdicts +end + +report(verdicts::Vector{Verdict}; io=stdout) = + ConformanceSupport.report("arrow-testing corpus", verdicts; io=io) + +if abspath(PROGRAM_FILE) == abspath(@__FILE__) + corpus = isempty(ARGS) ? DEFAULT_CORPUS : ARGS[1] + verdicts = runcorpus(corpus) + nfail = report(verdicts) + println() + println("PASS families by check:") + for check in unique(v.check for v in verdicts if v.check != "all") + n = count(v -> v.check == check && v.status == :pass, verdicts) + m = count(v -> v.check == check && v.status != :skip, verdicts) + println(" ", rpad(check, 24), n, "/", m) + end + exit(nfail == 0 ? 0 : 1) +end diff --git a/conformance/host/Project.toml b/conformance/host/Project.toml new file mode 100644 index 00000000..597a8b27 --- /dev/null +++ b/conformance/host/Project.toml @@ -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. + +# The HOST side of the conformance driver: only what conformance/run.jl +# needs to build the image and drive the container. run.jl activates and +# instantiates this environment itself. (conformance/Project.toml is the +# in-container suite environment.) +[deps] +Harbor = "af79dbb9-1a80-47ad-8928-192a4af69376" + +[compat] +Harbor = "=1.1.0" +julia = "1.10" diff --git a/conformance/oracle.jl b/conformance/oracle.jl new file mode 100644 index 00000000..e749b35c --- /dev/null +++ b/conformance/oracle.jl @@ -0,0 +1,388 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# ============================================================================= +# Oracle round-trips: OUR IPC bytes through pyarrow and nanoarrow. +# +# julia conformance/run.jl oracle # in the conformance image +# +# The gold corpus proves us against files C++ wrote years ago; this suite +# proves us against implementations running today. The corpus supplies the +# data matrix (every layout the format defines), the conformance image +# (conformance/Dockerfile) supplies the oracles — a Python with pyarrow and +# nanoarrow named by ARROW_ORACLE_PYTHON (ORACLE_WORKDIR keeps the generated +# cases and oracle output in a fixed directory instead of a temp dir) — and +# for every gold family we run: +# +# ours→pyarrow stream parse the gold JSON into Core, write OUR stream +# bytes; pyarrow reads them, structurally validates, +# rewrites its own stream; OUR reader reads that +# back and the values must equal the gold JSON. +# Proves pyarrow accepts our bytes and we accept +# pyarrow's, value-losslessly. +# ours→pyarrow file the same through the file format (footer path). +# ours→nanoarrow stream nanoarrow's IPC reader consumes our stream and +# hands the arrays over the C-stream capsule; +# nanoarrow's writer (or pyarrow's, on releases +# without one) produces the return stream. Proves +# the nanoarrow reader accepts our bytes. +# +lz4 / +zstd compressed-body variants of primitive, +# nested-dictionary, and binary-view families, +# proving our compressed framing against C++. +# +stats one primitive file with Arrow.jl statistics schema +# metadata; pyarrow must accept it and preserve values. +# +# Comparison is the corpus's own value-level document comparison, so id +# reassignment, pool unification, and padding differences by the oracle +# writers are already normalized away. +# ============================================================================= + +using JSON +import Arrow + +if !isdefined(@__MODULE__, :ConformanceSupport) + include(joinpath(@__DIR__, "ConformanceSupport.jl")) +end +using .ConformanceSupport: + ArrowJSON, + DEFAULT_CORPUS, + Verdict, + documentcheck, + familyskipreason, + filedocument, + readjson, + streamdocument + +# The oracle interpreter: a Python with pyarrow (and nanoarrow) importable. +# The conformance image sets it; on a host, point it at any such interpreter. +function _oraclepython() + py = get(ENV, "ARROW_ORACLE_PYTHON", "") + isempty(py) && error( + "ARROW_ORACLE_PYTHON is not set: run this suite through " * + "`julia conformance/run.jl oracle` (the conformance " * + "image), or point ARROW_ORACLE_PYTHON at a Python with pyarrow and nanoarrow", + ) + return py +end + +# The in-container driver. One process over all cases: reads each of our +# streams/files, structurally validates them, and writes the return bytes plus a +# results.json of per-case statuses (an oracle refusing our bytes is a +# finding, not a crash). +const PYDRIVER = raw""" +import json, os, sys +import pyarrow as pa +import pyarrow.ipc as ipc +import nanoarrow as na +import nanoarrow.ipc as naipc + +def na_stream(path): + # nanoarrow's IPC entry point moved across releases; probe. + if hasattr(na, "ArrayStream") and hasattr(na.ArrayStream, "from_path"): + return na.ArrayStream.from_path(path) + return na.ArrayStream(naipc.InputStream.from_path(path)) + +def classify(e): + # Raw error text only — the Julia side decides skip-vs-fail against an + # explicit whitelist of known oracle capability gaps, so a NEW + # interoperability failure can never classify itself into a skip. + return f"{type(e).__name__}: {e}"[:200] + +def rewrite(batches, schema, open_sink): + # Per-batch rewrite: read_all()/write_table merges chunks and drops + # zero-length batches, which breaks batch-boundary comparison against + # the gold JSON. Structural validate only: full validation enforces the + # advisory contracts (date64 divisibility, decimal precision) that the + # gold corpus itself violates and default C++ reads accept. + for b in batches: + b.validate() + with open_sink(schema) as w: + for b in batches: + w.write_batch(b) + +work = sys.argv[1] +cases = json.load(open(os.path.join(work, "cases.json"))) +outdir = os.path.join(work, "out") +os.makedirs(outdir, exist_ok=True) +results = {"pyarrow": pa.__version__, + "nanoarrow": na.__version__, + "cases": {}} +for case in cases: + name = case["name"] + r = {} + spath = os.path.join(work, "cases", name + ".stream") + fpath = os.path.join(work, "cases", name + ".arrow") + try: + reader = ipc.open_stream(spath) + rewrite(list(reader), reader.schema, lambda s: ipc.new_stream( + os.path.join(outdir, name + ".pyarrow.stream"), s)) + r["pyarrow_stream"] = "ok" + except Exception as e: + r["pyarrow_stream"] = classify(e) + try: + f = ipc.open_file(fpath) + rewrite([f.get_batch(i) for i in range(f.num_record_batches)], + f.schema, lambda s: ipc.new_file( + os.path.join(outdir, name + ".pyarrow.arrow"), s)) + r["pyarrow_file"] = "ok" + except Exception as e: + r["pyarrow_file"] = classify(e) + try: + reader = pa.RecordBatchReader.from_stream(na_stream(spath)) + batches = list(reader) # nanoarrow reads OUR bytes + outpath = os.path.join(outdir, name + ".nanoarrow.stream") + try: # prefer nanoarrow's own writer for the return trip + back = pa.RecordBatchReader.from_batches(reader.schema, batches) + with naipc.StreamWriter.from_path(outpath) as w: + w.write_stream(na.c_array_stream(back)) + except Exception: + rewrite(batches, reader.schema, + lambda s: ipc.new_stream(outpath, s)) + r["nanoarrow_stream"] = "ok" + except Exception as e: + r["nanoarrow_stream"] = classify(e) + results["cases"][name] = r +json.dump(results, open(os.path.join(work, "results.json"), "w")) +print(f"driver: {len(cases)} cases") +""" + +struct OracleCase + name::String # __[+codec], also the file stem + label::String # /[+codec], for the report + goldpath::String +end + +# These families cover primitive buffers, dictionary batches with nested child +# buffers, and variadic view buffers without multiplying every corpus case. +const COMPRESSED_ORACLE_FAMILIES = + ("generated_primitive", "generated_nested_dictionary", "generated_binary_view") +const STATISTICS_ORACLE_FAMILY = "generated_primitive" +const TARGETED_ORACLE_VARIANTS = ( + ( + (family, codec) for family in COMPRESSED_ORACLE_FAMILIES for codec in (:lz4, :zstd) + )..., + (STATISTICS_ORACLE_FAMILY, :stats), +) + +function _requiretargetedoracles(cases) + missing = Tuple{String,Symbol}[] + for (family, variant) in TARGETED_ORACLE_VARIANTS + suffix = "__" * family * "+" * String(variant) + any(case -> endswith(case.name, suffix), cases) || push!(missing, (family, variant)) + end + isempty(missing) || + error("pinned corpus did not prepare targeted oracle cases: $missing") + return nothing +end + +""" +Write OUR stream + file bytes for every non-skipped gold family (plus targeted +compressed families) into workdir/cases, and the case manifest the python +driver walks. +""" +function preparecases(corpus::String, workdir::String; require_targeted::Bool=false) + root = joinpath(corpus, "data", "arrow-ipc-stream", "integration") + isdir(root) || error("corpus not found at $root (set ARROW_TESTING_DIR)") + casedir = joinpath(workdir, "cases") + mkpath(casedir) + cases = OracleCase[] + skips = Tuple{String,String}[] + for v in sort(filter(d -> isdir(joinpath(root, d)), readdir(root))) + dir = joinpath(root, v) + families = sort!( + unique!([ + replace(f, r"\.json\.gz$" => "") for + f in readdir(dir) if endswith(f, ".json.gz") + ]), + ) + for fam in families + why = familyskipreason(fam, v) + isempty(why) || (push!(skips, (v * "/" * fam, why)); continue) + goldpath = joinpath(dir, fam * ".json.gz") + sch, batches, dictids = ArrowJSON.fromjson(readjson(goldpath)) + variants = fam in COMPRESSED_ORACLE_FAMILIES ? (:none, :lz4, :zstd) : (:none,) + for compress in variants + suffix = compress == :none ? "" : "+" * String(compress) + name = v * "__" * fam * suffix + write( + joinpath(casedir, name * ".stream"), + Arrow.writestream(sch, batches; compress=compress, dictids=dictids), + ) + write( + joinpath(casedir, name * ".arrow"), + Arrow.writefile(sch, batches; compress=compress, dictids=dictids), + ) + push!(cases, OracleCase(name, v * "/" * fam * suffix, goldpath)) + end + if fam == STATISTICS_ORACLE_FAMILY + name = v * "__" * fam * "+stats" + write( + joinpath(casedir, name * ".stream"), + Arrow.writestream(sch, batches; dictids), + ) + write(joinpath(casedir, name * ".arrow"), Arrow.statsfile(sch, batches)) + push!(cases, OracleCase(name, v * "/" * fam * "+stats", goldpath)) + end + end + end + require_targeted && _requiretargetedoracles(cases) + open(joinpath(workdir, "cases.json"), "w") do io + JSON.print(io, [Dict("name" => c.name) for c in cases]) + end + isempty(cases) && + error("arrow-testing corpus contains no runnable oracle cases under $root") + return cases, skips +end + +""" +Run the python driver over workdir with the oracle interpreter and return +the parsed results.json (an oracle refusing our bytes is a finding, not a +crash: the driver records per-case statuses and exits 0). +""" +function runoracles(workdir::String) + driver = joinpath(workdir, "driver.py") + write(driver, PYDRIVER) + py = _oraclepython() + Base.run(`$py $driver $workdir`) + return JSON.parsefile(joinpath(workdir, "results.json")) +end + +const ORACLE_CHECKS = ( + ("ours→pyarrow stream", "pyarrow_stream", ".pyarrow.stream", streamdocument), + ("ours→pyarrow file", "pyarrow_file", ".pyarrow.arrow", filedocument), + ("ours→nanoarrow stream", "nanoarrow_stream", ".nanoarrow.stream", streamdocument), +) + +# The ONLY oracle errors this suite treats as skips: known capability gaps, +# whitelisted by check, case, and error text. Anything else — including a +# feature error on a case not listed here — is a failure to investigate. +const ORACLE_EXPECTED_GAPS = ( + ( + "nanoarrow_stream", + n -> endswith(n, "+lz4") || endswith(n, "+zstd"), + "unsupported feature COMPRESSED_BODY", + ), + ( + "nanoarrow_stream", + n -> occursin("generated_binary_view", n), + "BinaryView not yet supported", + ), + ( + "nanoarrow_stream", + n -> occursin("generated_list_view", n), + "ListView/LargeListView not yet supported", + ), + ( + "nanoarrow_stream", + n -> occursin("generated_run_end_encoded", n), + "RunEndEncoded not yet supported", + ), +) + +_expectedgap(key::String, name::String, status::String) = any( + k == key && pred(name) && occursin(text, status) for + (k, pred, text) in ORACLE_EXPECTED_GAPS +) + +function runoracle( + corpus::String=DEFAULT_CORPUS; + workdir::String=get(ENV, "ORACLE_WORKDIR", mktempdir(prefix="arrow-oracle-")), +) + cases, skips = preparecases(corpus, workdir; require_targeted=true) + println("oracle: ", length(cases), " cases prepared in ", workdir) + results = runoracles(workdir) + return compareresults(cases, skips, results, workdir), results +end + +"Remove exactly Arrow.jl's statistics key from one decoded IPC document." +function _withoutstatsmetadata(document) + stripped = deepcopy(document) + schema = get(stripped, "schema", nothing) + schema isa AbstractDict || error("statistics oracle document has no schema object") + metadata = get(schema, "metadata", nothing) + metadata isa AbstractVector || + error("statistics oracle document has no schema metadata") + matches = count( + entry -> + entry isa AbstractDict && get(entry, "key", nothing) == Arrow.STATS_KEY, + metadata, + ) + matches == 1 || error("statistics oracle document has $matches statistics keys") + kept = [ + entry for entry in metadata if + !(entry isa AbstractDict && get(entry, "key", nothing) == Arrow.STATS_KEY) + ] + isempty(kept) ? delete!(schema, "metadata") : (schema["metadata"] = kept) + return stripped +end + +""" +Judge the oracle outputs: for every case the driver reported "ok", read the +return bytes with OUR reader and compare values against the gold JSON. +""" +function compareresults(cases::Vector{OracleCase}, skips, results, workdir::String) + verdicts = Verdict[] + for (label, why) in skips + push!(verdicts, Verdict(label, "all", :skip, why)) + end + for case in cases + gold = readjson(case.goldpath) + statsinput = if endswith(case.name, "+stats") + filedocument(read(joinpath(workdir, "cases", case.name * ".arrow"))) + end + if statsinput !== nothing + push!(verdicts, documentcheck(case.label, "ours stats file→gold", gold) do + _withoutstatsmetadata(statsinput) + end) + end + r = get(results["cases"], case.name, Dict{String,Any}()) + for (check, key, suffix, reader) in ORACLE_CHECKS + status = get(r, key, "driver produced no result") + if status != "ok" + # Capability errors skip only through the explicit whitelist. + # A missing required oracle or result is a conformance failure. + kind = _expectedgap(key, case.name, status) ? :skip : :fail + push!(verdicts, Verdict(case.label, check, kind, status)) + continue + end + # The targeted file route adds the statistics key that the gold + # JSON does not contain. Compare PyArrow's return against our + # exact input document there, proving both values and unknown + # schema metadata survive. Its ordinary stream routes still use + # the metadata-free gold document. + expected = + key == "pyarrow_file" && endswith(case.name, "+stats") ? statsinput : gold + push!(verdicts, documentcheck(case.label, check, expected) do + reader(read(joinpath(workdir, "out", case.name * suffix))) + end) + end + end + return verdicts +end + +function oraclereport(verdicts::Vector{Verdict}, results; io=stdout) + header = + "oracle round-trips (pyarrow $(get(results, "pyarrow", "?")), " * + "nanoarrow $(get(results, "nanoarrow", "?")))" + return ConformanceSupport.report(header, verdicts; io=io) +end + +if abspath(PROGRAM_FILE) == abspath(@__FILE__) + corpus = isempty(ARGS) ? DEFAULT_CORPUS : ARGS[1] + verdicts, results = runoracle(corpus) + nfail = oraclereport(verdicts, results) + exit(nfail == 0 ? 0 : 1) +end diff --git a/conformance/run.jl b/conformance/run.jl new file mode 100644 index 00000000..de020353 --- /dev/null +++ b/conformance/run.jl @@ -0,0 +1,128 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 conformance driver: run the suites inside the conformance image. +# +# julia conformance/run.jl [suite ...] [--rebuild] +# +# Suites: `corpus` (arrow-testing gold files), `oracle` (our IPC bytes +# through pyarrow and nanoarrow), `cdata` (C Data / C Stream through an +# in-process pyarrow); no argument runs all three. Docker is the only host +# requirement (plus network the first time, to fetch Harbor.jl and build the +# image): the image (conformance/Dockerfile) carries Julia, the oracle +# Python, and the gold corpus, and the repository is bind-mounted at /work. The image is built once (`--rebuild` forces it) and +# a named volume keeps the Julia depot — precompilation caches — between +# runs. Harbor.jl manages the container; suite output streams live. +# +# This script needs no `--project`: it activates and instantiates its own +# host environment (conformance/host/, Harbor only). conformance/Project.toml +# is the IN-CONTAINER suite environment and is prepared inside the image. +# +# Exit code: 0 iff every requested suite passed. +# ============================================================================= + +import Pkg +Pkg.activate(joinpath(@__DIR__, "host"); io=devnull) +Pkg.instantiate(; io=devnull) +using Harbor + +const IMAGE = "arrow-julia-conformance:latest" +const DEPOT_VOLUME = "arrow-julia-conformance-depot" +const REPO = normpath(joinpath(@__DIR__, "..")) +const SUITES = Dict( + "corpus" => "conformance/corpus.jl", + "oracle" => "conformance/oracle.jl", + "cdata" => "conformance/cdata_oracle.jl", +) +const SUITE_ORDER = ["corpus", "oracle", "cdata"] + +_haveimage() = + success(pipeline(`docker image inspect $IMAGE`; stdout=devnull, stderr=devnull)) + +function buildimage() + println("conformance: building $IMAGE (once; --rebuild forces)") + Base.run(`docker build -t $IMAGE -f $(joinpath(REPO, "conformance", "Dockerfile")) + $(joinpath(REPO, "conformance"))`) + return nothing +end + +# `docker exec` directly (not Harbor.exec, which captures output): the suites +# run for minutes and their progress belongs on the terminal. +function execstream(container, cmd::Vector{String}) + proc = Base.run(ignorestatus(`docker exec $(container.id) $cmd`)) + return proc.exitcode +end + +function main(args) + rebuild = "--rebuild" in args + requested = filter(a -> a != "--rebuild", args) + isempty(requested) && (requested = SUITE_ORDER) + for s in requested + haskey(SUITES, s) || + error("unknown suite $(repr(s)); choose from $(join(SUITE_ORDER, ", "))") + end + (rebuild || !_haveimage()) && buildimage() + + container = Harbor.run!( + IMAGE; + command=["sleep", "infinity"], + detach=true, + volumes=Dict("/work" => REPO, "/opt/julia-depot" => DEPOT_VOLUME), + ) + results = Dict{String,Int}() + try + # The suite environment: the repository's conformance project with + # the mounted checkout developed in. + # Cheap when the depot volume is warm; fetches only what changed. + println("conformance: preparing the suite environment") + rc = execstream( + container, + [ + "julia", + "--project=/opt/env", + "--startup-file=no", + "-e", + """using Pkg + cp("/work/conformance/Project.toml", "/opt/env/Project.toml"; force=true) + Pkg.develop([ + PackageSpec(path="/work/src/ArrowStrings"), + PackageSpec(path="/work/src/ArrowTypes"), + PackageSpec(path="/work"), + ]) + Pkg.instantiate(); Pkg.precompile()""", + ], + ) + rc == 0 || error("suite environment preparation failed (exit $rc)") + for s in requested + println() + println("conformance: ===== $s =====") + results[s] = execstream( + container, + ["julia", "--project=/opt/env", "--startup-file=no", SUITES[s]], + ) + end + finally + Harbor.cleanup!(container) + end + println() + for s in requested + println(rpad(s, 8), results[s] == 0 ? "PASS" : "FAIL (exit $(results[s]))") + end + exit(all(==(0), values(results)) ? 0 : 1) +end + +main(ARGS) diff --git a/dev/release/README.md b/dev/release/README.md index 0ed3d9b2..7181380f 100644 --- a/dev/release/README.md +++ b/dev/release/README.md @@ -45,6 +45,12 @@ $ dev/release/release_rc.sh 1 The argument of `release_rc.sh` is the RC number. If RC1 has a problem, we'll increment the RC number such as RC2, RC3 and so on. +Keep the vote open for at least 72 hours. A release requires at least three +binding `+1` votes from PMC members and more binding positive votes than +binding negative votes. See the [ASF voting +rules](https://www.apache.org/foundation/voting.html) and the [Arrow release +guide](https://arrow.apache.org/docs/dev/developers/release.html). + Requirements to run `release_rc.sh`: * You must be an Apache Arrow committer or PMC member @@ -103,6 +109,16 @@ If ArrowTypes is also registered, we also need to post a comment on the commit w @JuliaRegistrator register subdir=src/ArrowTypes ``` +To register or release ArrowStrings, post the following comment on the commit +that contains its version in `src/ArrowStrings/Project.toml`: + +```markdown +@JuliaRegistrator register subdir=src/ArrowStrings +``` + +ArrowStrings must be registered and released before an Arrow.jl release that +depends on it can be registered. + ### Verify We have a script to verify a RC. diff --git a/dev/release/rat_exclude_files.txt b/dev/release/rat_exclude_files.txt index 6e32d072..af00dc93 100644 --- a/dev/release/rat_exclude_files.txt +++ b/dev/release/rat_exclude_files.txt @@ -19,4 +19,3 @@ Manifest.toml dev/release/apache-rat-*.jar dev/release/filtered_rat.txt dev/release/rat.xml -test/arrowjson/*.json diff --git a/dev/release/release_rc.sh b/dev/release/release_rc.sh index bae1997b..a5996df7 100755 --- a/dev/release/release_rc.sh +++ b/dev/release/release_rc.sh @@ -136,7 +136,7 @@ The source release rc${rc} is hosted at [2]. Please download, verify checksums and signatures, run the unit tests, and vote on the release. See [3] for how to validate a release candidate. -The vote will be open for at least 24 hours. +The vote will be open for at least 72 hours. [ ] +1 Release this as Apache Arrow Julia ${version} [ ] +0 diff --git a/docs/Project.toml b/docs/Project.toml index ebc6a27f..f431466d 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -17,7 +17,16 @@ [deps] Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45" +ArrowStrings = "c38d8858-22a2-449e-9eca-ef92ec15f353" +ArrowTypes = "31f734f8-188a-4ce0-8406-c8a06bd891cd" +DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" [compat] -Documenter = "1" +ArrowStrings = "1" +ArrowTypes = "2" +DataAPI = "1" +Documenter = "1.17" +Tables = "1.14" +julia = "1.10" diff --git a/docs/dev/DESIGN-incremental-writer.md b/docs/dev/DESIGN-incremental-writer.md new file mode 100644 index 00000000..35ec1322 --- /dev/null +++ b/docs/dev/DESIGN-incremental-writer.md @@ -0,0 +1,204 @@ + + +# Design: the Arrow 3.x incremental writer + +Status: PROPOSAL (not yet reviewed). Motivating issues: #237, #244, #247, +#413 (write memory), #485/#501 (Writer/append removal fallout), Pioneer.jl +and SpineOpt.jl migration blockers. + +## Problem + +Arrow 3.0's writer is eager: `Arrow.write` materializes every partition, +constructs every column across all partitions at once, validates, and then +publishes the complete IPC output. That buys whole-table schema inference +(cross-partition pool unification, Union routes, abstract narrowing) and +validation-before-publish, but it removes the two 2.x capabilities people +actually used: writing batches as they are produced (`Arrow.Writer`), and +adding batches to an existing sink (`Arrow.append`). It also makes every +write hold the whole table plus the whole output in memory. + +## What the internals already give us + +Both IPC writers are batch-sequential loops over a byte accumulator: + +- `writestream`: `_schemamessage!` once → per batch `_dictionarymessage!` + (replacement supported: a pool identity change re-emits under the same id) + + `_recordmessage!` → EOS marker. +- `writefile`: leading magic → the same message loop with per-message Block + bookkeeping (offsets measured against the accumulator length) → EOS → + Footer (schema + dictionary/record Block vectors) + length + magic. + +Nothing in the loop needs the batch *vector*; only three things are +whole-table today: the `out::Vector{UInt8}` accumulator, the upfront +`_validatewriterbatches(sch, batches, ids)` sweep, and the facade's +materialize-all-partitions phase in `_writebytes`. + +The second ingredient already exists too: **retained-schema column +construction**. `_constructcolumn(name, parts; retained=field)` builds a +column against a FIXED `Field` and fails closed with clear errors when the +data does not conform — the engine a fixed-schema incremental writer needs, +already tested by every facade rewrite. + +## Proposal + +### Layer 1 (internal): `IPCWriteState` + +Extract the shared loop into a state struct in `ipc_write.jl`: + +```julia +mutable struct IPCWriteState + io::IO # the sink; batch bytes publish per batch + file::Bool # footer bookkeeping on close? + schema::Schema + fielddictids::IdDict{Field,Int64} + ids # assigndictids result + codec::UInt8 + state::Union{Nothing,EncodeState} + current::Dict{Int64,ArrayData} # stream: last pool per id + emitted::Set{Int64} # file: ids already written + dictblocks::Vector{NTuple{3,Int64}} # file only + recordblocks::Vector{NTuple{3,Int64}} + written::Int64 # bytes published (replaces length(out)) + closed::Bool +end + +beginwrite!(io, sch; file, compress, dictids) -> IPCWriteState +writebatch!(st, batch::AC.RecordBatch) # validate THIS batch, stage its + # bytes, publish once +finishwrite!(st) # EOS; file: footer + magic +``` + +Each `writebatch!` stages one batch's messages into a scratch buffer, +validates before touching `io`, then publishes — the eager writer's +validate-before-publish invariant holds per batch instead of per file. +`writestream`/`writefile` become `beginwrite!` + a loop + `finishwrite!` +over an in-memory sink, so the eager paths are byte-identical and the whole +existing battery/conformance/oracle surface proves the refactor. + +Dictionary rules fall out of the existing checks: + +- Stream format: a later batch with a different pool identity re-emits a + replacement dictionary message (already the loop's behavior). +- File format: one dictionary batch per id; `writebatch!` refuses a pool + identity change with the same wording `writefile` uses today. (A future + option is accumulating pools and emitting dictionary blocks just before + the footer — the Footer's Block index makes position irrelevant to + conforming readers — but that changes pool-growth semantics and is out of + scope for v1.) + +### Layer 2 (facade): `Arrow.Writer` + +The current `Arrow.Writer` tombstone becomes the real thing again: + +```julia +w = Arrow.Writer(sink; file=true, compress=nothing, + metadata=nothing, colmetadata=nothing) +Arrow.write(w, table) # one record batch per Tables.partitions partition +close(w) # idempotent; EOS or footer; then close(io) iff we opened it +Arrow.Writer(sink; kw...) do w ... end +``` + +- The FIRST `Arrow.write(w, table)` runs today's full column construction on + that table alone, fixing the `Schema`, and emits the schema message plus + the first batches. +- Every LATER write constructs each column with `retained=` that fixed + schema — exactly the rewrite path — so conformance failures reuse the + existing, well-worded errors. +- Nullability is declared, not data-driven, matching the 3.0 write rule: + a field is nullable iff the first table's column eltype admits `Missing`. + A later batch with missing values under a non-nullable field is refused + with a fix-forward message ("make the first table's column eltype + missing-capable, or write an empty typed table first"). Writing a + zero-row table with a fully typed `Tables.Schema` as the first write is + the supported way to pin a schema explicitly — no new kwarg needed. +- `sink` is a path (opened at construction; the handle owns and closes it) + or an `IO` (borrowed; `close(w)` finishes the IPC output but leaves the + `IO` open). A mid-stream error leaves a torn sink — documented, like any + incremental format writer; `Arrow.write` remains the atomic option. +- Single-owner cursor, like the read side: overlapping `write`/`close` + calls on one handle throw; an `@atomic closed` flag backs idempotent + close and use-after-close errors. +- Memory: O(current table) plus retained pool references. This is the + documented answer to #237/#244/#247/#413. + +### Layer 3 (facade): `Arrow.append` + +```julia +Arrow.append(path_or_io, table; compress=nothing) +``` + +v1 supports the STREAM format only — the format where append is natural and +the format 2.x users actually appended to (SpineOpt, Agents.jl): + +1. Read the existing schema: open the sink, `readstream` the schema message + (and dictionary ids) — validation of the existing prefix is the reader's + normal framing pass over what it consumes; append does not re-verify + every historical batch body. +2. Construct the new table's columns with `retained=` the existing schema. +3. Seek to the trailing EOS marker (last 8 bytes), truncate/overwrite it, + `writebatch!` the new batches (dictionary replacement messages allowed), + re-emit EOS. + +A file-format sink gets a clean refusal pointing at `Arrow.Writer` or a +whole-file rewrite. (File-format append — footer read, truncate, extend the +Block lists, rewrite the footer — is mechanically possible with +`verify_footer` + the Block machinery, but it rewrites trailing bytes in +place and v1 should not normalize that risk. Revisit on demand.) + +## What this deliberately does not do + +- No cross-batch schema inference: the first write (or an explicit typed + empty first write) is the schema authority. This is the semantic price of + incrementality, stated loudly in the docstrings: the eager + `Arrow.write` can infer a Union across partitions; the Writer cannot. +- No `ntasks`-style pipelining: `writebatch!` is serial. The 2.x + concurrency was a large share of its bug surface; add measured + parallelism later if benchmarks demand it. +- No delta dictionaries (the 3.0 reader rejects them by design) and no + file-format pool growth. +- No append-to-file-format in v1. + +## Testing + +- Refactor equivalence: eager `writestream`/`writefile` outputs are + byte-identical before/after Layer 1 (golden comparison in the battery). +- Writer suites: multi-write stream and file outputs are byte-identical to + the equivalent eager `Tables.partitioner` write for conforming inputs; + schema-conformance refusals (type drift, missing under non-nullable, + file-format pool change) assert the error wording; torn-sink behavior on + a mid-write error; do-block; close idempotence; use-after-close. +- Append suites: append to a 3.0-written stream, a 2.x-written stream + fixture (fixtures2x), and an appended-to-appended stream; read back with + `Arrow.Stream` and pyarrow (conformance oracle case); dictionary + replacement across appends; refusal on schema drift and on file-format + sinks. +- Fuzz: extend the deterministic fuzzer with a Writer route (random batch + splits of each corpus table must read back equal to the eager write). + +## Sizing and sequencing + +Layer 1 is a contained refactor of two functions that already have the +right loop shape. Layer 2's hard half (fixed-schema construction) already +exists as the retained path; the new code is the handle, first-write +bootstrapping, and errors. Layer 3 reuses both. Ship order: Layer 1+2 +(Writer, both formats) before the 3.0 release — it converts the release's +one structural regression into a headline feature and unblocks Pioneer; +Layer 3 (stream append) with it if review goes smoothly, else in 3.0.x +(SpineOpt/Agents migrate to `Arrow.Writer` either way). diff --git a/docs/dev/DESIGN-scan-ranges-trim.md b/docs/dev/DESIGN-scan-ranges-trim.md new file mode 100644 index 00000000..61a28866 --- /dev/null +++ b/docs/dev/DESIGN-scan-ranges-trim.md @@ -0,0 +1,404 @@ + + +# Design: Tables.Scan pushdown, cloud byte-range reads, and the trim contract + +Implemented in `src/scan.jl` and `src/source.jl` and exposed through the +facade (`Arrow.Table(source; scan=…)` over an `AbstractArrowSource`); §5 +lists what is and is not built. The three pieces share one mechanism: **a bound column set drives +both what gets decoded and what gets fetched, and every request value is +plain data intended to remain visible to the trim verifier.** Section 4 +separates that design intent from what the trim harness actually compiles. + +--- + +## 1. Tables.Scan support + +`Tables.Scan` (from the Tables.jl revision pinned in `Project.toml` until its +first release) is a plain-data scan request — +select/rename/type items, a closed predicate algebra (`Cmp`/`In`/`IsNull`/ +`StrPred`/`And`/`Or`/`Not`, with `OpNode` as the growth channel), `limit`/ +`offset`. A source accepts a `Scan` as a keyword and pushes down what it can +while materializing. Arrow resolves each request once against the source +schema. Direct handles pass the bound request to the storage-only +`_applyscan(handle, bound)`. The public facade passes the complete plan to +`_applyfacadescan`, which owns route-aware materialization, public conversion, +and result wrapping as one operation. Key contract points this design leans +on: + +- a pushed result must equal the executor's over the same scan, row for + row (the differential battery pins this); +- `limit`/`offset` may be consumed **only** over exactly filtered rows; +- no `Function` values in `Scan` or predicate data — the request algebra is + closed and value-only. + +### What Arrow can push, by axis + +| Axis | Mechanism | Exactness | +|---|---|---| +| `select` | decode only (selected ∪ filter-referenced) columns: a registry-driven `skipfield!` advances the node/buffer cursor past unselected fields without body slicing, content validation, or materialization. Complete node/buffer metadata is still validated first. Nested subtrees skip with their parent; no body range is requested for an unselected dictionary column. | exact as IO/decode reduction (see below for who projects) | +| `limit`/`offset` | Without a filter `RecordBatch.length` is wire metadata: whole batches before `offset` and after `offset+limit` are never decoded (ranged reads still fetch candidate RecordBatch metadata because Footer Blocks have no row counts, but request no body range for excluded batches). With a filter the window composes over the qualifying rows of each decoded batch and decoding stops once it is full. Tail reads and configured coalescing may physically over-read otherwise unrequested bytes. | exact | +| `filter` | two tiers: (a) **statistics pruning** — per-batch min/max/null-count, when present (§3 of this doc), prune batches that cannot satisfy the predicate; (b) **mask at materialization** — evaluate the predicate over each batch's decoded columns with the generic evaluator (`Tables.filtermask`) and keep only the qualifying rows when building output columns. | (a) inexact — a pruned batch is provably empty; (b) exact — enables limit/offset pushdown with filters | +| `types` (`ref => T`) | An elementwise conversion request, not a parse seed. Direct handle scans apply it to storage-domain columns. The facade applies it after ArrowTypes and native public conversion. See §4 for the known-schema role in trim mode. | exact | + +### The pushdown shape + +`_ScanPlan` resolves a facade request once. It keeps the public-domain bound +plan and, when every filter literal has an exact storage representation, one +lowered storage-domain bound plan. `_runboundscan` is the private batch kernel. +`_applyscan` closes it over ordinary storage materialization for direct +handles. File and ranged `_applyfacadescan` methods close it over facade +routing. The stream method routes decoded columns and uses `_executeplan`. +Every facade method consumes its private route markers before returning: + + plan(scan, fields) = + resolve once against schema names → + lower filter literals exactly for storage, or mark public fallback → + strip public type overrides from the storage plan + + Temporal literals lower only when the public conversion preserves the + requested operator over the complete physical domain. Date64 and + millisecond Timestamp mappings preserve equality but not order; + Timestamp-second and Time mappings alias physical values. Duration + literals may use the column unit or a coarser fixed unit, never a finer + one. `In` applies the equality rule to Tuple and Array; Set additionally + requires the canonical public type to preserve `isequal` and hashing. + Custom membership objects force public-domain fallback. + +Tables.jl does not yet expose execution for an existing `BoundScan`. +`_executeplan` is the narrow local adapter for public and stream fallback. It +uses Tables.jl's predicate and allocation seams and is pinned to the generic +executor by the differential battery; it avoids reconstructing and resolving a +second `Scan`. + + _runboundscan(f, bound, materializecolumn) = + decode set = selected ∪ filtercols (source order, source names) → + batch set = limit/offset window from wire row counts (no filter), + ∩ stats-surviving batches (filter present, stats present) → + per decoded batch (`_ScanSink`): + materialize the decode set through the operation's fixed policy → + rows = filter ? qualifying rows after this batch's share of + offset/limit (saturating, the executor's rule) + : the metadata window's rows → + keep each selected column's rows, in selection order, under its + output name; stop decoding once the window is full → + return table over the selection + + _applyscan(handle, bound) = + _runboundscan(handle, bound, storage materializer) + + _applyfacadescan(file-or-ranged-source, plan) = + assert plan.storage has no public type overrides → + _runboundscan(source, plan.storage, route-aware materializer) → + lift ArrowTypes routes and native public values → + apply plan.public type overrides and wrap the public Table + + _applyfacadescan(stream, plan) = + assert plan.storage has no public type overrides → + materialize route-aware columns after stream decode → + _executeplan(columns, plan.storage) → + lift routes, apply plan.public overrides, and wrap the public Table + +The filter is evaluated by the generic evaluator (`Tables.filtermask`) over +the batch's decoded columns, so Arrow's pushdown and the executor share one +three-valued semantics by construction (a `missing` predicate excludes the +row). Unselected columns cost zero decode and add zero planned body bytes. +Tail reads and coalescing may still over-read them under §2's explicit +policy. + +Properties the composition holds (all implemented in `src/scan.jl`): + +- **Positional filter references resolve to source names** before + evaluation (`Tables.resolve` normalizes them): a bound `col(3)` means + source column 3, and the decode-set table the evaluator sees carries + source names. +- **Wire row counts are trusted only after metadata validation.** Before a + `RecordBatch.length` drives a window, it is range-checked and matched to + every top-level FieldNode length. Exact node/buffer counts and buffer + geometry are also checked from metadata alone. Aggregate scan row counts + must fit Tables' `Int` row-count API: a planned result through + `typemax(Int)` is accepted and a larger one rejected. An offset-only + window represents `limit=nothing` explicitly; it does not use a finite + sentinel that can omit later batches. +- **One scan has one allocation budget.** Standalone lazy `file[i]` calls + retain their documented per-call budgets. A scan that visits many batches + shares one budget and codec state across footer work, metadata, fetched + range payloads, decompression, Core and ArrowTypes materialization, scan + slices/joins/overrides, and any full-object public-domain fallback. +- **Windows saturate.** `offset`/`limit` compose with `min` arithmetic over + row counts, as the executor does, so `offset + limit` never overflows. +- **Type overrides run once in the correct value domain.** Direct handle + scans apply them to storage-domain columns. The facade keeps them out of + its storage plan and applies them after public conversion. +- **Private Union routes stay local.** Direct `_applyscan` cannot select the + ArrowTypes route-aware materializer. `_applyfacadescan` creates and consumes + identical-storage Union child markers inside one operation, before a result + crosses back into `table.jl`. +- **Zero-field sources** consume filter and window from row counts alone + (`_zerofieldwindow`): a row-invariant predicate evaluates once, and no + per-row mask is allocated from an untrusted row count. +- **Advisory nullability does not weaken the skip boundary.** The reader + admits nulls under `nullable=false`; `validate_full` rejects them. If such + a null exists only in an unread batch, it does not widen the selected + result's concrete Julia vector type. Row values and conforming schemas + still match the generic executor exactly. + +The statistics fold resolves dictionary indices through the pool before it +computes logical null/min/max values; the row evaluator compares decoded +values. Stream-format sources are scanned by the facade after decode — the +eager stream reader has already decoded by then; stream pushdown would need +an incremental framer. + +--- + +## 2. Cloud byte-range reads + +### Why the format already supports this + +The IPC **file** format is: magic · messages · Footer(schema, dictionary +Blocks, record Blocks) · footer-length · magic. Every Block carries +`(offset, metaDataLength, bodyLength)`; every RecordBatch header carries a +per-buffer `(offset, length)` table within its body; the registry walk maps +buffer indices → fields deterministically (the exact mechanism +`decodefield`/`encodefield!`/`skipfield!` share). So the fetch plan for +"columns X, Y of batches 3..7" is pure arithmetic over two small metadata +reads. Compressed buffers are self-contained (per-buffer prefix + frame), so +they range-fetch identically. The **stream** format has no footer and stays +sequential — cloud-native access is a file-format feature, stated plainly. + +### The fetch protocol + +1. **Tail fetch** (one range request, cached on the handle): the last + `tailbytes` (default 64 KiB). Its trailing magic decides file vs stream + format (a stream object is read whole instead), and it covers + footer-length + magic + the whole Footer in almost every real file; if + `footerlen + 10 > tailbytes`, one exact follow-up fetch (also cached). + → schema, Block indexes, (§3) statistics — everything pruning needs. + The leading magic is not fetched: the Footer is the sole authority. +2. **Statistics prune** from the Footer metadata — zero additional fetches. +3. **Block metadata fetches**: dictionary metadata plus + `(offset, metaDataLength)` for each statistics-surviving record batch, + coalesced across nearby spans. RecordBatch row counts are here, not in the + Footer, so `limit`/`offset` windowing happens after this pass. +4. **Window and buffer-range plan**: row counts choose the exact batch/body + window when there is no filter; the bound column set then maps to a buffer + index set (subtree- + inclusive; dictionary Blocks for selected dictionary columns) → byte + ranges → **coalesce** ranges with gaps below `coalesce_gap` (default + ~256 KiB — a gap fetch is usually cheaper than a request round-trip; + both knobs are options, not constants). +5. **Body fetches** — the selected dictionary bodies and the selected record + buffers in ONE round: each coalesced range lands in its own owned heap + region; the dictionaries decode first from the shared spans. Decode + resolves each declared buffer `(offset, len)` to its containing fetched + range and subslices — the message-body-authority + invariant becomes *"every buffer must fall inside a fetched range that + was itself derived from the verified buffer table"*: same trust story, + sparse backing. + +Request-count model (what actually matters against cloud latency): `1` tail +(plus one sequential exact Footer follow-up when the tail is too small) ++ `⌈candidate metadata spans after coalescing⌉` + `⌈coalesced body ranges⌉`. +The common path has three sequential request rounds. A Footer follow-up makes +four. The source can issue a metadata or body round's planned ranges +concurrently. +For a 40-column file reading 3 columns of every batch, this moves roughly +`3/40` of the body bytes plus metadata. Statistics-pruned batches contribute +no requested metadata/body range. Coalescing is an explicit over-read policy, +so a requested span may cross otherwise unneeded bytes when the configured +gap permits it. + +### The interface (no HTTP/CloudStore deps in Arrow) + +Arrow defines a minimal source contract (`src/source.jl`) and owns the +planner; transports live in extensions: + + abstract type AbstractArrowSource end + sourcelength(src)::Integer # total object length, known up front + readrange(src, offset, len)::Vector{UInt8} # one exact range, 0-based offset + concurrentreads(src)::Int # default 1 + # Arrow reads a round's planned ranges through readrange with a + # handle-wide worker pool, capped by Limits.max_concurrent_reads, + # storing every result by request index — + # a source's completion order can never permute payloads, and the + # reads in flight are bounded whatever the span count. + +- The entry point is `Arrow.Table(src; scan=…)`, which builds the internal + `SourceFile(src; limits, tailbytes, coalesce_gap)` handle and runs + `Tables.scan(sf, scan)`; the source's length is read once, at handle + construction, and the tail once per handle. `Arrow.Table(src)` without a + scan, with a scan that cannot be pushed down, over a zero-field file, or + over a stream-format object reads the object whole (one request beyond + the cached tail), as does `Arrow.Stream(src)`. +- One `SourceFile` samples `concurrentreads` once and shares one semaphore + across its operations. Separate handles are independent; a transport that + requires one wider cap enforces it in `readrange`. +- The abstract type is a dynamic call under `--trim` for a source type the + trimmed app never mentions; an app that names its concrete source type + resolves statically. The planner itself is arithmetic over `Int64`s. +- Extension: `ext/ArrowCloudStoreExt.jl` (loaded with CloudStore.jl) makes + a `CloudStore.Object` a source — its known `size` is the length, one + range is one HTTP `Range` GET pinned to the object's ETag with + `If-Match`, and `concurrentreads` is 16 — and adds `Arrow.Table(::CloudStore.Object; …)` and + `Arrow.Stream(::CloudStore.Object; …)`. An HTTP transport is the same two + methods. Zero new hard deps. +- The differential test: sparse fetch ≡ whole-file read, plus + request-count/byte-count/range assertions on a counting test source + (`test/scan_battery.jl`), and the CloudStore extension end to end against + a local S3-compatible server (`test/cloudstore_tests.jl`). +- Explicitly out of scope v1, documented: caching/prefetch policy beyond + coalescing, retries (the source's job), writers over ranges, stream + format (read whole), mutation detection (ETag pinning is the extension's + concern). + +The ranged reader deliberately uses the Footer schema as its sole schema +authority. It does not parse or cross-check the leading schema message or +optional EOS marker, although a head, tail, or coalesced request can physically +over-read those or other unrequested bytes. The complete Footer Block index is +bounded and checked for overlap. Required features and message limits are +global. Per-record metadata/body/buffer limits stay lazy like `ArrowFile`: +dictionary blocks and statistics-surviving record candidates are checked, while +statistics-pruned record metadata causes no dedicated range request and is not +validated (a tail or coalesced request may still over-read it). Message kind, +version, legacy-compression state, complete node/buffer metadata, +layout-derived buffer minima, child extents, and fixed or fully-covered +null-count contracts for planned subtrees, planned codec, and required +dictionary presence are validated before any planned body range is requested. +Skipped buffer contents remain unvalidated by design. + +--- + +## 3. Per-batch statistics (the pruning fuel) + +Arrow's format has no per-batch statistics on the wire; the ecosystem's +"statistics schema" standardizes the **value layout** for exchanging +statistics as Arrow data, but placement in IPC files is not (yet) +standardized upstream. This convention is deliberately conservative: + +- **Placement (our convention, upgradeable)**: one schema-level custom + metadata key, e.g. `JuliaArrow:batch_statistics.v1`, carried in the + **Footer's** schema copy so the tail fetch alone powers pruning. +- **Value layout**: follow the official statistics-schema array layout, + serialized as one embedded IPC stream (statistics ARE Arrow data); per + record batch × flattened RecordBatch FieldNode index: row_count, + null_count, min, max (the `ARROW:*:exact` keys). Top-level fields after + nested fields + therefore do not use their top-level ordinal as the `column` value. + Using the official layout keeps us convention-compatible if upstream + standardizes placement later — we then emit both keys for a deprecation + cycle and read either. +- Writer: `withstatistics` / `statsfile` eagerly compute the embedded + stream for already-encoded batches. An opt-in `statistics=true` keyword + computing the same fold state during encode is not implemented; file + format only. An append path would have to recompute or drop + — dropping with a warning is the honest v1. +- Reader: prune under `Cmp`/`In`/`IsNull` (and `StrPred` prefix ranges for + `startswith`) with one-sided may-contain logic — a batch survives unless + the predicate is provably false for ALL rows; the exact row filter always + runs after pruning (pruning is inexact by design). Missing or MALFORMED + statistics degrade to "no pruning", never to an error. The embedded + stream and Base64 output share the enclosing scan allocation budget; + exhausting that cumulative caller limit remains a scan error instead of + being mistaken for malformed optional metadata. + Float comparisons use the predicate's IEEE operators; any NaN disables + bounds, and signed zero is not ordered with `isless`. Dictionary folds + count null pool results as logical nulls. +- **Trust model, stated plainly**: statistics are + trusted-for-completeness, exactly like Parquet row-group stats. The exact + row filter protects one direction only — batches kept by lying stats still + filter row-exactly. The other direction has no net: stats + that under-report a range cause false EXCLUSION. Excluded batches are not + decoded and cause no dedicated metadata/body range request, so their + qualifying rows are silently lost. Tail/coalescing over-read does not restore + them. Wide (conservative) lies cost pruning, never correctness; narrow lies + lose rows. The acceptance battery pins all three behaviors. + +--- + +## 4. The trim contract (staying on the radar, explicitly) + +**Trimmability is a standing production gate, not an aspiration.** The +`--trim=safe` harness (0 errors / 0 warnings / binary exit 0) compiles +`ArrowCore` plus its value-domain and typed-value workloads and the C-data +seams, plus a representative ArrowStrings workload. It does **not** yet compile +a scan-and-materialize app, so §1–§3 are designed for trim but not yet gated by +it. The rules in `core-README.md` ("Trim-compile support") constrain their form: + +- `Tables.Scan` is trim-aligned by its own charter (no `Function` fields; + closed algebra). The evaluator uses the same closed-set `isa` ladder + pattern as `layoutspec_of`; `OpNode` rejection keeps the set closed. + `resolve` is plain data → plain data. +- The range planner is arithmetic over `Int64`s; a trimmed app that names + its concrete `AbstractArrowSource` type resolves the source calls + statically. No dynamic registry on the hot path. +- **Two-tier public API**: the runtime-tagged + core is inherently trim-safe — descriptors are values, accessors use + literal load widths, struct scalars are `Vector{Pair{String,Any}}`. So: + - **Tier 1 (trim target)**: the value-domain entry points — + open/scan/materialize returning value-domain data, plus C-data/stream + interop and the typed `getvalue(::Type{T}, …)`/`materialize(::Type{T}, …)` + path. A harness compiling a scan-and-materialize app at 0/0/exit-0, + kept permanently in CI, is what would make the scan half guaranteed. + - **Tier 2 (dynamic, ergonomic)**: the typed facade (`Arrow.Table` + property access, NamedTuple rows) — explicitly NOT trim-guaranteed. + - **The known-schema bridge**: `Scan`'s `ref => Type` overrides ARE the + known-schema declaration. In a trimmed app, a scan with concrete type + pins can drive the typed-column path whose element types are statically + known (`Vector{Int64}`, `Vector{Union{Missing,Float64}}`, …) through + closed-width branches — "provide a known schema and get typed columns, + trimmed" falls out of the same plain-data request, no second schema + surface needed. + +--- + +## 5. Status + +Implemented (`src/scan.jl`, `src/table.jl`): + +- **Scan pushdown**: `skipfield!`, one `_ScanPlan` compilation per facade + request, storage-only `_applyscan(::ArrowFile, bound)`, and the closed + file/ranged `_applyfacadescan(source, plan)` route consuming the bound scan + exactly through `_ScanSink` (per-batch filter evaluation, composed + limit/offset with early stop, selection, renames, direct-handle type + overrides at column construction, and facade conversion before return), + zero-field scans, and the + differential battery with corruption-backed never-decoded proofs. `Arrow.Table(source; scan=…)` + routes through it on file-format and ranged inputs; stream-format inputs + route inside their facade operation and scan post-decode with identical + results. +- **Byte-range sources**: the `AbstractArrowSource` contract, the + `SourceFile` fetch protocol, the coalescing planner, `SparseBody` decode, + the CloudStore.jl extension, and counting-source proofs (zero planned body ranges for skipped columns, + window-excluded batches, and unneeded dictionary bodies, with exact + request-log checks under the fixtures' tail/coalescing settings). +- **Statistics**: `withstatistics`/`statsfile` fold the official statistics + value layout into `JuliaArrow:batch_statistics.v1` (footer schema + metadata, base64-wrapped IPC stream, one statistics batch per data batch, + serialized through this writer); `_maypass` may-contain pruning is wired + into both applies (ranged pruning happens before the block-metadata pass, + so pruned batches cause no dedicated metadata/body request; configured + tail/coalescing may over-read them); acceptance pins exactness, + degradation, and both lie directions. + +Not implemented: an HTTP transport extension (the source contract is the +extension point; the CloudStore extension is the model), an encode-time +`statistics=true` writer keyword, upstream-placement tracking +for statistics, and the scan-and-materialize trim harness. Scan pushdown +depends on the pinned Tables.jl development revision until that API is +released. diff --git a/docs/dev/core-README.md b/docs/dev/core-README.md new file mode 100644 index 00000000..c57a6144 --- /dev/null +++ b/docs/dev/core-README.md @@ -0,0 +1,523 @@ + + +# The Arrow.jl engine: ArrowCore and its adapters + +Arrow.jl is a runtime-tagged, C-data-shaped core (`ArrowCore`, private) with +the IPC reader/writer, the C data and C stream interfaces, and `Tables.Scan` +pushdown as peers over it, and `Arrow.Table`/`Arrow.Stream`/`Arrow.write` as +the public facade on top. `ArrowCore` depends only on Base and the Mmap +standard library. This document is the design rationale and the exact +scope of every layer. + +## Files + +| File | Purpose | +|---|---| +| `src/ArrowCore.jl` | Ownership regions with one revocation cell, runtime descriptors, `Field`/`Schema`, `ArrayData`, the layout registry, staged validation, dynamic and typed accessors, bulk fixed-width extraction, minimal builders, `RecordBatch`, and `RecordBatchSource` | +| `src/metadata/` | FlatBuffers metadata bindings and the shape verifier, both generated by `tools/fbsgen.jl` from the vendored `src/metadata/fbs/*.fbs`, over the schema-blind `VerifierRuntime.jl` | +| `src/FlatBuffers/` | The vendored FlatBuffers runtime (table reads, builder) | +| `src/ipc_read.jl` | Checked IPC stream framing, resource limits, metadata-to-Core mapping, dictionary state, one registry-driven decoder, per-buffer decompression | +| `src/ipc_write.jl` | The write half over the same registry: Core-to-metadata mapping, one generic registry-driven encoder, replacement-on-change dictionary batches, per-buffer compression, the file format (Block index + Footer), and the lazy random-access `ArrowFile` reader | +| `src/cdata.jl` | C data and C stream interfaces both directions: zero-copy ownership, move semantics, exactly-once release, field and schema metadata transport | +| `src/source.jl` | The `AbstractArrowSource` byte-range source interface (`sourcelength`, `readrange`, `concurrentreads`) | +| `src/scan.jl` | The private scan-plan module: one-time binding and lowering, exact execution and schema projection, sparse byte-range reads over `SourceFile`, and embedded per-batch statistics. `table.jl` includes it at the public/storage type seam. | +| `src/table.jl` | The read facade, including the shared public/storage type seam used by scan lowering and retained construction | +| `src/columnconstruction.jl` | The deep column-construction module: fresh inference, retained-schema reconstruction, recursive ArrowTypes lowering, shared dictionary pools, partition agreement, and field metadata behind `_constructcolumn` | +| `src/write.jl` | The write facade: partition binding, retained-field alignment, schema and batch assembly, compression selection, and IPC emission | +| `ext/ArrowCloudStoreExt.jl` | CloudStore.jl objects as sources: HTTP `Range` reads, concurrent per planned range | +| `src/ArrowStrings/` | ArrowStrings.jl — the shared inline-else-view string representation (`ArrowString`, `StringVector` = Utf8View memory); a separate package, registered on its own like ArrowTypes, that Arrow depends on through a `[sources]` path entry until its first release | +| `src/ArrowTypes/` | ArrowTypes.jl — the separate custom-type interface package; the facade applies its lowering and extension hooks recursively | +| `test/support/AcceptanceSupport.jl` | One explicit private dependency seam for the four stable adapter acceptance suites; the C Data stress child uses its own narrow support module | +| `test/support/SeededFuzz.jl` | Version-stable differential, layout-family, ranged-read, statistics-pruning, and bounded mutation fuzzing with exact replay artifacts | +| `test/` | Core and facade tests, ArrowTypes compatibility, support contracts, the four adapter acceptance batteries, frozen 2.x-written fixtures, and the `--trim=safe` gate | +| `conformance/ConformanceSupport.jl` | Strict integration-JSON shape validation, logical-value canonicalization and comparison, skip policy, verdict construction, and reporting shared by every conformance adapter | +| `conformance/` | The arrow-testing gold-corpus runner, the integration-JSON implementation, the pyarrow/nanoarrow IPC oracle, and the in-process pyarrow C Data / C Stream oracle | +| `bench/` | The serialize/deserialize benchmark harness (this package and PyArrow) | +| `docs/dev/DESIGN-scan-ranges-trim.md` | The scan pushdown, ranged-fetch, and statistics design | + +## Run it + +On Julia 1.10, prepare a fresh checkout once before you run the commands below. +Julia 1.10 does not read the repository's `[sources]` entries. + +```bash +julia --project=. -e 'using Pkg; Pkg.develop([PackageSpec(path="src/ArrowStrings"), PackageSpec(path="src/ArrowTypes")])' +``` + +```bash +julia --project=. -e 'using Pkg; Pkg.test()' # core + facade + batteries +julia --startup-file=no test/trim_compile_tests.jl # JuliaC --trim=safe gate +julia --startup-file=no conformance/run.jl # all conformance suites, in the docker image +julia --startup-file=no conformance/run.jl corpus # one suite: the gold corpus … +julia --startup-file=no conformance/run.jl oracle # … IPC bytes through pyarrow + nanoarrow … +julia --startup-file=no conformance/run.jl cdata # … C Data + C Stream through an in-process pyarrow +julia --project=. test/fuzz.jl --cases 16 --mutations 64 # deterministic PR-sized fuzz suite +julia --project=. test/fuzz.jl --cases 512 --mutations 20000 --determinism-every 256 --repro-dir fuzz-reproductions # extended suite +julia --project=. bench/run.jl # benchmarks +julia tools/fbsgen.jl src/metadata/fbs src/metadata # regenerate bindings + verifier +``` + +For a chosen master seed, a fixed SplitMix64 stream derives every case. The +master seed and case index identify the same input on every supported Julia +version. PR tests and manual runs default to one fixed seed. Each scheduled run +uses its workflow run ID as a new reproducible seed. Active-case artifacts are +written before parser work. They record both seeds, the source revision, exact +generated and saved-byte commands, and the active Project and Manifest with a +restore command. Mutation replays compare two full outcomes. The scheduled +workflow also repeats the first complete mutation-route sweep and every 256th +later mutation. It stops the fuzz process before the job timeout, then uploads +the active artifacts even for a hang or forced stop. After download, run the +artifact's `replay.sh` with a clean Arrow.jl checkout path. The wrapper resolves +its own artifact directory, temporarily installs and instantiates the recorded +Project and Manifest, runs the replay, and restores the checkout files. + +`Tables.Scan` pushdown needs the Tables.jl revision pinned in `Project.toml` +developed into the project environment (the conformance image clones it). The conformance +suites run inside one docker image (`conformance/Dockerfile`: Julia, a +Python with pyarrow and nanoarrow that PythonCall binds to, the +apache/arrow-testing corpus, the pinned Tables revision, a warm depot in a named +volume) driven by `conformance/run.jl` through Harbor.jl. The driver +instantiates its own tiny host environment (`conformance/host/`, Harbor.jl +only) on first run, so docker — and network for that first run — are the +only host requirements; `conformance/Project.toml` is the IN-CONTAINER suite +environment. The C interfaces hand pointers across an in-process +boundary, which is why the suites run in the container rather than against +one. + +## Design in one table + +| Principle | Where it lives | +|---|---| +| Ownership is an object; bad spans fail before access | `OwnerRegion`, checked `BufferSlice` construction, bounds-checked `loadat`. Foreign C extents are trusted declarations. | +| Validity is reachability plus one revocation bit | The memory model below. | +| Logical parameters are values, not type parameters | `TimestampType(unit, timezone)`, `DecimalType(precision, scale, bitwidth)`, and every other descriptor keep schema data out of Julia types. | +| One structural registry plus bounded per-layout methods | `layoutspec` defines buffer roles, child arity, offset width, and variadic status; access and semantic rules are grouped methods. | +| Staged validation | Structural, then semantic, then the opt-in full tier; each public stage composes the earlier ones. Data-intrinsic semantic results are cached; Field contracts run every time. | +| Framing enforces limits before allocation | The IPC framer enforces metadata, body, message, and allocation limits; the generated verifier bounds objects, depth, and copy reserve; the decode cursor enforces array and buffer limits before the related work. | +| The message body is the decode authority | Every declared batch buffer is a checked `subslice` of its own message body; cursor completion and non-overlap checks reject skewed buffer tables. | +| IPC ids are adapter state | `corefield` records ids in identity-keyed adapter tables; `DictionaryType` holds the value type and `ArrayData.dictionary` the value array; neither stores an id. | +| C Data is a direct mapping over `ArrayData` | `to_c_data`/`from_c_data` use per-structure callbacks and controls, separate schema/array roots that keep sources reachable, transitive release, and explicit reaping. | +| Function-barrier bulk access | `materialize` resolves the layout once and loops; scalar `getvalue` pays one dynamic dispatch. Static claims through `getvalue(::Type{T}, …)`/`materialize(::Type{T}, …)` resolve statically. | + +## Memory model + +Buffer validity is GC reachability, plus one revocation bit. An +`OwnerRegion` is a `(ptr, len, alignment, root, cell)` record: `root` is an +opaque GC anchor (the wrapped `Vector`, the Mmap-stdlib array, or an +adapter's owner object), so holding any slice of a region keeps the backing +memory alive by construction; `cell` is a `ReleaseCell` shared by every +region over one underlying lifetime. Loads are a bounds check, one monotonic +closed-flag load, and the raw read — no lock, no guard, no state machine on +the hot path. + +`release!` is the deterministic release: it revokes every region sharing the +cell (later raw access throws `InvalidStateException`) and then runs the +cell's release action exactly once. An mmap region unmaps NOW (the eager +path exists for hosts where a GC-timed unmap is not enough — deleting a +still-mapped file on Windows being the canonical case); an imported C-data +tree runs the producer's release callback; a borrowed heap region is revoked +with no eager action (running a borrowed object's finalizers is not ours to +do). `release!` is idempotent and is not a data-race shield for accesses +concurrent WITH the close — quiescing readers first is the caller's +contract, as with `Base.close` on a shared IO. Every buffer imported from +one C-data tree is backed by regions sharing one cell, so closing any of +those regions, or `release!` on the import's owner, revokes all siblings +before the single producer release. + +What the model does not do: nothing prevents external writes to or +truncation of a mapped file while the mapping or cached validation results +are in use (no userspace scheme can). On systems that prohibit deleting an +active mapping, `release!` (or collection) must complete before the path can +be deleted. + +Exactly-once release lives in the adapters' owners: the C-data +`ForeignOwner` and C-stream `StreamOwner` carry one `@atomic` flag each, a +finalizer, and an explicit `release!`; their revocation cells route through +the same flag. The export registries root exported columns and streams until +the consumer releases them and cleanup drops the root. + +## Scope and limits, layer by layer + +### Core + +Accessors and validation cover every format-1.5 layout: integer, floating +point, Boolean, decimal (32/64 as integers, 128/256 as raw native-endian +bytes), date, time, timestamp, duration, all interval variants, UTF-8 and +binary with 32-bit or 64-bit offsets, Utf8View and BinaryView (16-byte +entries, inline and out-of-line, variadic data buffers, the spec's +prefix-must-match rule), fixed-size binary, list, fixed-size list, ListView +and LargeListView (per-slot offsets and sizes, unordered and overlapping +ranges, invariants binding null slots too), struct, map, sparse and dense +union over the full Int8 id domain, run-end encoding (signed 16/32/64 +no-null strictly-ascending run ends, binary-search access, logical nulls +through the values child, parent null count zero or unknown), dictionary, +and null arrays. Logical parent offsets and nested slices are supported. + +Struct scalars on the dynamic path are ordered `Vector{Pair{String,Any}}`, +so names stay in the value domain and duplicate, empty, or +non-Symbol-compatible names are representable. Core keeps every schema field +and child name as a `String`; it does not intern schema names. The typed path +(`getvalue(::Type{T}, field, data, i)`, `materialize(::Type{T}, field, +data)`) is a caller-asserted element domain: exact match only (no +conversion; `Missing <: T` admits nulls; `Any` is the dynamic path), +composites recurse (List → `Vector{E}`, Struct → `Vector{Pair}` or a +NamedTuple with name checks, Dictionary → pool values, run-end encoding +transparent), unions refuse every static claim, and claims are checked +against the descriptor before any element is read. Closed fixed-width +claims over Int/Float/temporal/Decimal32/64 columns use one bounds-checked +bulk byte copy plus a bitmap null punch. + +`validate_full` adds UTF-8 well-formedness for Utf8 and Utf8View, the +advisory nullability contract, Date64 day divisibility, time-of-day range, +decimal precision, and canonical bit-packed form (zeroed trailing bits and +padding in validity/Bool buffers; unsliced arrays only, since sliced windows +legitimately share bitmap bytes). On-wire buffer padding is a writer +guarantee, not a reader requirement. Map validation checks physical layout +and reachable Field nullability; it does not check key uniqueness, +hashability, or ordering — `keysSorted` is a producer declaration. +Timestamp validation checks the unit domain and timezone-string UTF-8; it +does not resolve names against a timezone database. `RecordBatch` buffers +must be host-native endian (the IPC adapters refuse big-endian input; no +adapter normalizes). Julia vectors wrapped zero-copy by the builders or `heapregion` are +scoped borrows: they must not be resized or mutated while their `ArrayData` +or cached validation results are in use. + +`fromviewentries` wraps a vector of Arrow view entries (ArrowStrings' +`ArrowStringPayload`, or any 16-byte isbits type with that layout) and +its data buffers as a Utf8View column, zero-copy — the payload vector IS the +views buffer and every data buffer is retained by identity; only the +validity bitmap is built, and long-entry geometry (offsets inside their +buffer, prefixes matching the data) is checked by semantic/full validation, +not at construction. The facade's `Arrow.write` routes +`ArrowStrings.StringVector` columns through it. + +### IPC + +The reader maps every layout above, including nested dictionary encoding +(read and written in dependency order). It accepts V4 and V5 metadata on +little-endian hosts, supports feature-gated full dictionary replacement, +preserves old dictionary snapshots, and rejects delta dictionaries. It +requires the eight-byte continuation-marker framing (the pre-0.15 four-byte +prefix is not accepted). Compression uses the V5 `BodyCompression` field +for LZ4_FRAME and ZSTD; it accepts the `COMPRESSED_BODY` schema feature and +also accepts V5 compressed streams from Arrow.jl 2.x that omit it; it +rejects `BodyCompression` under V4 and the pre-1.0 experimental V4 +compression marker. Big-endian streams are refused (no endianness +normalization). + +Compatible fields that share one IPC dictionary id share one immutable pool +object; eager stream decoding fully validates each pool snapshot once and +reuses that identity certificate for structural, intrinsic, and +Field-contract validation while still checking each field's index array +independently — validation work is linear in the encoded indices plus +distinct pool data. The reader runs structural and semantic validation +before exposing a batch and does not opt into `validate_full`; the +generated verifier does validate FlatBuffer strings. The framer refuses a +non-little-endian host before any generated getter runs. + +`readstream` decodes a borrowed `Vector{UInt8}` eagerly (raw batch buffers +are zero-copy views; positively compressed buffers are exact-sized owned +copies) behind the `RecordBatchSource` pull interface; the caller must not +mutate or resize the vector while the stream or its batches live. +`IPCStream` is a single-owner cursor — overlapping `nextbatch!` calls throw +`ConcurrencyViolationError`. `max_total_allocated_bytes` is one +reader-wide, conservative budget for metadata copies, metadata-directed +Julia containers, decompressed outputs, Core materialization, ArrowTypes +route containers, and facade copies. `IPCStream` carries the remaining budget +into `Arrow.Table` and `Arrow.Stream`; it is not a measurement of custom user +hook allocations or every Julia allocation. Package-owned vector reserves +include conservative backing capacity because Julia can round the requested +payload to a larger allocation class. Schema and Field metadata stay as ordered pair +vectors, so duplicate keys and their original order survive IPC reads and +rewrites. + +The writer covers the same layouts with one registry-driven encoder, the +declared inverse of `decodefield`. It writes V5 stream bytes and the file +format (magics, Block indexes, Footer) with per-buffer LZ4_FRAME/ZSTD +compression behind the Int64 prefix and the `-1` stored-raw fallback. +Dictionary handling is replacement-on-change (one batch per pool snapshot; +`Feature.DICTIONARY_REPLACEMENT` declared when a replacement is emitted; +`COMPRESSED_BODY` when a compressed batch is). Files declare the +compression feature in both schema copies and refuse pools that change +identity across batches. Every column is semantically validated before its +bytes are published. The writer is eager and sequential (byte vectors, +buffer contents copied into message bodies); arrays with a nonzero element +offset are refused (materialize first); each schema position must be a +distinct `Field` object; fresh dictionary ids are assigned per field, and a +caller-supplied id table makes shared ids write as one shared dictionary +batch with value-schema compatibility, one nested-id topology per repeated +id, and one pool per id within each batch enforced before bytes are +emitted. Canonical empty offset arrays materialize their terminal zero on +the wire (the reader accepts the omitted form other writers emit). + +`readfile` verifies both magics, the leading and footer schemas, cumulative +footer work, and every Block's frame, message kind, wire-buffer extents, and +overlap before optional-EOS classification; `ArrowFile` decodes record +batches lazily by footer index — each `getindex` runs with a fresh +allocation budget and codec contexts over the shared, eagerly-decoded +dictionary set, so concurrent reads need no coordination. + +### Scan pushdown and ranged reads + +`Tables.scan(::ArrowFile, scan)` and file/ranged `Arrow.Table(source; scan=…)` +share the same batch kernel through separate closed direct and facade +operations. They decode only the selected and filter-referenced columns, +prune whole batches through the embedded statistics (one-sided: a pruned +batch is provably empty), and consume the +scan exactly: without a filter `limit`/`offset` are metadata arithmetic and +batches outside the window are never decoded; with one the `_ScanSink` +evaluates the filter per batch through the generic evaluator +(`Tables.filtermask`), composes the window over the qualifying rows, and +stops decoding once it is full. Each request is resolved once. Projection +and renames are applied at column construction. Direct handle scans apply +type overrides in the storage domain. The facade applies them after public +conversion. The direct `_applyscan` seam is storage-only. The facade's +`_applyfacadescan` operation owns route-aware ArrowTypes Union materialization, +public conversion, and wrapping, so its private child markers cannot cross +back into `table.jl`. Stream facade scans keep the same route local but use +`_executeplan` after decode. `SourceFile` runs the same plan over an +`AbstractArrowSource`: +it uses the Footer (normally from one cached tail read, with one exact cached +follow-up when the Footer escapes that window) as its sole schema authority, +validates the full Block index and the complete metadata plan for every +statistics-surviving record before +requesting a body range, and requests per-buffer body ranges for exactly +the decode set, coalesced under `coalesce_gap`. It does not fetch the +leading magic, parse or cross-check the leading schema message, or inspect +the optional EOS marker; tail reads and coalescing may physically +over-read any unrequested bytes. Embedded batch statistics use the official +Arrow statistics value layout under the `JuliaArrow:batch_statistics.v1` +placement key (placement is scoped out of the upstream spec) and are +trusted for completeness: conservative lies cost pruning, narrow lies can +lose rows. Scan pushdown over duplicate column names is refused. + +### C data and C stream interfaces + +Every Core layout crosses the boundary: Boolean, integer, floating point, +null, decimal (32/64/128/256 in the `d:` form), date, time, timestamp (with +and without timezone), duration, all three interval units, UTF-8 and binary +(both offset widths), fixed-size binary, list, large list, fixed-size list, +struct, map, sparse and dense union (ids in the format string), dictionary, +Utf8View and BinaryView (`vu`/`vz`, with the C-Data-only trailing int64 +buffer of variadic data-buffer lengths), ListView/LargeListView +(`+vl`/`+vL`), and run-end encoding (`+r`). Field metadata and schema-level +metadata cross both directions (schema metadata rides the stream's +struct-typed schema node; dictionary field metadata rides the wrapper node, +matching the C++ bridge, and import concatenates wrapper and dependent +pairs losslessly). + +Foreign allocation extents cannot be verified by the ABI and are trusted +declarations; the producer must keep declared storage alive and unchanged +until Core releases it. Import checks the pointer tables, counts, descriptor +shape, and checked geometry that the ABI does expose. Import and export +apply the semantic validation tier — the same default as the IPC reader and +writer; `validate_full` is the caller's opt-in on either side. Field names +containing an embedded NUL are refused (C strings are NUL-terminated), +imported names must be valid UTF-8, and C strings longer than 1 MiB without +a terminator are refused instead of scanned. The format parser accepts only +the specified decimal integer grammar and bounds decimal descriptors and +union ids before recursive or geometry work. Empty offset layouts export one +non-NULL terminal zero offset for strict cross-implementation parity. The C +timestamp format has one empty-timezone spelling, so a Core empty string +canonicalizes to `nothing` when imported again. + +Release callbacks use producer-owned canonical child and dictionary +topology, so cleanup does not depend on caller-mutated public counts or +pointer tables; they inspect canonical descendants' public release fields +to honor consumer moves. A callback transaction that fails before commit +restores its node to LIVE and returns at the void C boundary; a later +explicit call resumes it without repeating completed children. Callbacks +for one exported tree are serialized and legal only on Julia-attached +threads (they call Julia and take a `ReentrantLock`); there is no lock-free +foreign-thread trampoline. `reap!` performs an explicit registry scan; +there is no background reaper. Schema and array trees have independent +aggregate lifetimes and per-node control blocks. + +`export_stream!` fills a caller-owned `ArrowArrayStream` that streams +batches as struct-typed arrays; each `get_schema`/`get_next` result is an +ordinary export root, producer-side failures surface through +`get_last_error` (EINVAL plus a NUL-terminated message owned by the stream +until replaced or released), and the stream's own root drops at its release +callback. `from_c_stream` moves a producer's stream, reads the schema once, +pulls batches whose trees each own one `ForeignOwner`, and surfaces +producer errors as exceptions carrying the producer's message. Stream +callbacks call into Julia, so they are legal only from Julia-attached +threads and calls on one stream must not overlap (the C stream spec itself +declares the structure not thread-safe). + +The ABI layout gates include 32-bit expectations; only the 64-bit branch +is exercised (on the available hosts), the 32-bit branch is inspected. + +### Facade + +`Arrow.Table` materializes columns into plain Julia vectors (closed +fixed-width claims through Core's bulk typed path, everything else through +the dynamic accessors, then facade conversions: Dates types in both +directions, with sub-millisecond timestamps staying raw integers rather than +silently truncating). `Arrow.Stream` iterates record batches as one Table +each. `Arrow.write` accepts any Tables.jl source (partitions become record +batches), `DictEncode` marks a column for pooling, retained-schema rewrites +of a `Table`/`Stream` recursively preserve every descriptor that materialized +values can reconstruct, plus nullability and ordered metadata. Top-level +dictionary pools retain order, unused and duplicate entries, null entries, +and index width; multi-partition dictionary columns share one pool object. +Unregistered Union routing and nested dictionary pools are no longer present +after facade materialization, so those retained rewrites fail closed. +Registered ArrowTypes.jl public-domain values keep writer-side type evidence and can +reconstruct retained Union routing, including retained dense or sparse mode and +type IDs. Sparse children receive canonical hidden placeholders outside their +active rows. An abstract registered target accepts an extensionless concrete +subtype or one with the retained parent identity; a different explicit identity +fails closed. A nullable Dictionary with an unknown extension also fails +closed because materialized `missing` values cannot retain the difference +between a valid null-pool index and a null index. View buffer +topology, ListView overlap, and exact run segmentation rebuild canonically. +DataAPI metadata reads through. The facade applies ArrowTypes.jl lowering and +extension restoration recursively to top-level and nested values. There is no +lazy typed-view layer, no parallel writer pipeline, and no append-as-resume. +One column-scoped construction context caches `ArrowType` per Julia type, +extension shape and Union decomposition per Julia type, runtime Union branch +routes, and `JuliaType` per retained `Field`. +Tables column names are the explicit process-global `Symbol` boundary. Before +interning any novel top-level name, the facade preflights the complete schema: +4096 UTF-8 bytes per name, at most 65,536 novel names per table +materialization, and at most 1 MiB of novel-name bytes in total. Failure occurs +before partial interning. Nested Core names remain strings. +On the ordinary resolved paths, `toarrow` runs once for each value that reaches +lowering; dictionary categories are pooled before they are lowered. The +ArrowTypes 2.x fallback for an unresolved all-missing abstract storage type is +kept for compatibility. Recursive custom schemas and value containers, and +custom mapping nesting beyond 64 levels, fail with `ArgumentError`. +Fresh unresolved abstract declarations collect concrete subtype evidence once +for the complete column. This preserves subtype extensions and uses an explicit +bounded Union when the observed subtype Fields differ. +Declared writer Unions have at most 32 branches. Runtime writer or storage +inference has at most 8 distinct types across the complete column. This covers +abstract ArrowTypes storage, abstract or `Any` dictionary values, and abstract +retained ArrowTypes targets. The lower inferred limit bounds per-type trait and +candidate compilation; an explicit declared Union remains the schema authority +for wider intentional type sets. +Fixed-size-list storage signatures are exact through arity 1024. Larger +descriptors use compact `Tuple{Vararg{T}}` signatures so logical type resolution +cannot allocate in proportion to an untrusted list size. Extension Struct +signatures are exact through 1024 children only when child names are unique, +contain no embedded NUL, already exist as Julia `Symbol`s, are at most 4096 +UTF-8 bytes each, and use at most 64 KiB in total. Otherwise the labelled Struct +remains unknown and materializes as ordered `Pair` storage. The bounded +ArrowTypes.jl Tuple compatibility exception may intern names only when the +complete child sequence is exactly `"1"`, `"2"`, …, `string(N)` for +`N ≤ 1024`; it can therefore add only `Symbol("1")` through `Symbol("1024")`. +Unknown extension labels return before preflight. Arbitrary or partly positional +Struct names never take this exception. +Unknown extension labels are probed with a non-interning Julia symbol lookup; +only an existing symbol can reach `JuliaType(Val(...))`. Unsupported-extension +warnings are deduplicated by the complete label. Each table materialization +emits at most one warning for each of 16 distinct labels, then one suppression +notice for further distinct labels. A warning displays at most 128 UTF-8 bytes +of its label. The built-in `JuliaLang.Symbol` mapping likewise rejects a storage +string that is not already interned, rather than growing process-global symbol +state from input. +Writer-side `ArrowType` results are checked at the trait cache boundary. A +returned concrete tuple above arity 1024 is rejected before downstream writer +specialization, including the default mapping for a tuple value. A custom trait +method itself is trusted Julia code. +Hidden retained composite data is built from Field plus length. Null-only +fixed-size lists recurse without per-slot Julia placeholders, including wholly +or partly inactive sparse-Union children. + +Retained descriptor matching recognizes only storage families whose builders +can enforce the original schema exactly. These include sequence layouts, +opaque binary and wide-decimal byte storage, interval NamedTuple storage, and +compatible temporal units. The retained builders enforce fixed widths, list +sizes, interval shapes, exact temporal conversion, and sorted Map claims. +Concrete declared element types remain planning evidence when a column is empty +or contains only missing values. + +## Trim-compile support (JuliaC `--trim=safe`) + +`test/trim_compile_tests.jl` compiles `test/trim_entrypoint.jl` with +JuliaC's `--trim=safe` and requires **zero verifier errors, zero verifier +warnings, and a produced binary that runs to exit 0**. The workload covers +regions, mmap, C-data export/import/release, dynamic values, typed values +(a `from_c_data` → `materialize(Int64, …)` scenario among them), and +validation errors. It also covers ArrowStrings construction, inline and view +access, missing values, comparison, and materialization. The rules that keep a +runtime-tagged core there: + +- **Closed-set dispatch ladders.** Dispatch on an abstract-typed field is + dynamic; the descriptor set is closed (it IS the layout registry), so + `@inline` `isa` ladders (`layoutspec_of`, `_value_of`, `_materialize_of`, + `typeequal`, `descriptorname`, `_validate_descriptor_of`, and the C-data + `formatstring_of`) devirtualize every generic entry point. Multiple + dispatch stays the per-layout extension surface underneath. Plain + forwards do not work: the verifier reports the abstract call site as + unresolved rather than enumerating the closed method table. +- **Narrow after `||`-checks.** An `isa` test inside an `||` condition does + not narrow the binding; a typeassert after it (`rt::IntType`) is what lets + `primwidth`/`_load_int` resolve. +- **Literal load widths.** `loadat(b, T, off)` with a runtime `T::DataType` + leaves the raw-load path unresolved; accessors branch to literal widths. +- **A recursion cycle's non-inlined edge is a compiled function with + all-concrete argument types.** Dynamic recursion routes through the public + `getvalue(Field, ArrayData, Int64)`; the typed path splits its edge into + an `@inline` scalar fast ladder (scalar children SROA into the parent + loop) plus a compiled shell for composites, with `::T` asserts pinning + claim-typed returns. `@generated` struct rows keep every field's claim a + literal type past the arity-4 `ntuple` cliff. +- **CAS, not `Core.modifyfield!`.** The verifier has not implemented the + read-modify-write builtin (each `@atomic x.f += 1` is a warning), while + `@atomicreplace` verifies clean. +- **`Ptr{Cvoid}` finalizers and cfunctions.** Base's generic + `finalizer(f, o)` is `@nospecialize`d and unresolvable; the typed pointer + form (`finalizer(@cfunction(...), o)`) is an ordinary ccall. Release + actions are runtime Ptr-ABI cfunctions (an `Any`-argument cfunction is + rejected), never stored in module-level `const`s (raw-pointer consts are + precompile-poison). +- **Concrete containers at the boundary.** Struct scalars are + `Vector{Pair{String,Any}}` on the dynamic path; lists materialize as + `Vector{Any}` without a runtime-narrowing comprehension. +- **Kwcall discipline.** A `Union`-typed keyword argument makes the kwcall + tuple imprecise — branch on presence instead; abstract-typed keyword calls + need positional twins; boxed closure captures (reassigned-under-`try` + locals) are rejected — single-assign before `try`. +- **Beware splatting Base conveniences.** `write(filename, x)` and + `open(...) do` route through vararg-splatting internals; `mktempdir`'s + cleanup registry parks the trimmed runtime's scheduler. +- Heterogeneous NamedTuple ingestion (`batch(nt)`, `fromjulia_struct`) is + runtime-schema builder work outside the trim-safe surface. + +## Compression + +Both IPC directions implement spec buffer compression for **LZ4_FRAME and +ZSTD** through the direct CodecLz4/CodecZstd dependencies over +TranscodingStreams. Each reader lazily creates raw native codec contexts and +closes them on every `readstream` exit path; each writer owns one lazily +initialized compressor per codec and finalizes it on every writer exit +path; there are no global pools. The write side emits the Int64 +uncompressed-length prefix per buffer and stores incompressible payloads raw +behind the `-1` sentinel. The read side checks the prefix and the sentinel; +a zero-byte wire buffer may omit the prefix; a nonzero compressed buffer, +including declared length zero, must contain a valid frame. Declared sizes +are bounded and charged to the shared reader budget before one exact-sized +output vector is allocated; the codecs decode directly from the wire slice +(its region rooted across the native call with `GC.@preserve`) with no +payload copy and no growable output; the LZ4 loop requires one complete +frame, exact input consumption, and exact output size, and the ZSTD +one-shot decode uses the same exact destination. diff --git a/docs/make.jl b/docs/make.jl index b26c44bc..bda7864b 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -22,16 +22,22 @@ makedocs(; modules=[Arrow], repo=Remotes.GitHub("apache", "arrow-julia"), sitename="Arrow.jl", + # The reference page documents the public surface explicitly; internal + # helpers carry docstrings for maintainers and are not part of the site. + checkdocs=:public, + checkdocs_ignored_modules=[Arrow.ArrowCore, Arrow.FlatBuffers, Arrow.Meta], format=Documenter.HTML(; - prettyurls=get(ENV, "CI", "false") == "true", + prettyurls=true, canonical="https://arrow.apache.org/julia/", assets=String[], ), pages=[ "Home" => "index.md", "User Manual" => "manual.md", + "Migrating from 2.x" => "migration.md", "API Reference" => "reference.md", ], + pagesonly=true, ) deploydocs(; repo="github.com/apache/arrow-julia", devbranch="main", branch="asf-site") diff --git a/docs/src/index.md b/docs/src/index.md index 4f900590..48a05fd7 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -21,8 +21,23 @@ # Arrow.jl +A pure Julia implementation of the [Apache Arrow](https://arrow.apache.org) +columnar format: the IPC stream and file formats (read and write, with +memory-mapped and byte-range reads, scan pushdown, and compression), the C +data and C stream interfaces for in-process exchange with other +implementations, recursive ArrowTypes.jl mappings for custom values, and +Tables.jl integration throughout. + +```@example quick_start +using Arrow +path = joinpath(mktempdir(), "data.arrow") +Arrow.write(path, (a = [1, 2, 3], b = ["x", "y", missing])) +tbl = Arrow.Table(path) +tbl.b +``` + ```@contents -Pages = ["manual.md", "reference.md"] +Pages = ["manual.md", "migration.md", "reference.md"] Depth = 3 ``` diff --git a/docs/src/manual.md b/docs/src/manual.md index 5a3330fd..aac9875a 100644 --- a/docs/src/manual.md +++ b/docs/src/manual.md @@ -1,4 +1,6 @@ -```@raw html +```@meta +CurrentModule = Arrow +``` -``` # User Manual -The goal of this documentation is to provide a brief introduction to the arrow data format, then provide a walk-through of the functionality provided in the Arrow.jl Julia package, with an aim to expose a little of the machinery "under the hood" to help explain how things work and how that influences real-world use-cases for the arrow data format. - -The best place to learn about the Apache arrow project is [the website itself](https://arrow.apache.org/), specifically the data format [specification](https://arrow.apache.org/docs/format/Columnar.html). Put briefly, the arrow project provides a formal specification for how columnar, "table" data can be laid out efficiently in memory to standardize and maximize the ability to share data across languages/platforms. In the current [apache/arrow GitHub repository](https://github.com/apache/arrow), language implementations exist for C++, Java, Go, Javascript, Rust, to name a few. Other database vendors and data processing frameworks/applications have also built support for the arrow format, allowing for a wide breadth of possibility for applications to "speak the data language" of arrow. - -The [Arrow.jl](https://github.com/apache/arrow-julia) Julia package is another implementation, allowing the ability to both read and write data in the arrow format. As a data format, arrow specifies an exact memory layout to be used for columnar table data, and as such, "reading" involves custom Julia objects ([`Arrow.Table`](@ref) and [`Arrow.Stream`](@ref)), which read the *metadata* of an "arrow memory blob", then *wrap* the array data contained therein, having learned the type and size, amongst other properties, from the metadata. Let's take a closer look at what this "reading" of arrow memory really means/looks like. +[Apache Arrow](https://arrow.apache.org) specifies a columnar memory layout +and an interprocess (IPC) serialization format for it, so that table data +can be shared across languages and processes without conversion. Arrow.jl +is a pure Julia implementation of that +[specification](https://arrow.apache.org/docs/format/Columnar.html): it +reads and writes the IPC stream and file formats, exchanges in-memory data +with other implementations through the C data and C stream interfaces, and +presents everything to Julia through the [Tables.jl](https://tables.juliadata.org) +interface (Tables.jl 1.14 or later, for `Tables.Scan`). -## Support for generic path-like types - -Arrow.jl attempts to support any path-like type wherever a function takes a path as an argument. The Arrow.jl API should generically work as long as the type supports: - -- `Base.open(path, mode)::I where I <: IO` +```julia +using Arrow, Tables -When a custom `IO` subtype is returned (`I`) then the following methods also need to be defined: +Arrow.write("data.arrow", (a = [1, 2, 3], b = ["x", "y", missing])) +tbl = Arrow.Table("data.arrow") +tbl.a # Vector{Int64} +tbl.b # Vector{Union{Missing, String}} +``` -- `Base.read(io::I, ::Type{UInt8})` or `Base.read(io::I)` -- `Base.write(io::I, x)` +## Reading -## Reading arrow data +### `Arrow.Table` -After installing the Arrow.jl Julia package (via `] add Arrow`), and if you have some arrow data, let's say a file named `data.arrow` generated from the [`pyarrow`](https://arrow.apache.org/docs/python/) library (a Python library for interfacing with arrow data), you can then read that arrow data into a Julia session by doing: +[`Arrow.Table`](@ref) reads an IPC source into columns. The source may be a +file path, an `IO`, a `Vector{UInt8}` of IPC bytes, or a byte-range fetcher +([Reading remote and partial files](@ref)). Both IPC formats are accepted +and detected automatically: the *file* format (`ARROW1` magic, random +access, optional footer statistics) and the *stream* format. ```julia -using Arrow - -table = Arrow.Table("data.arrow") +tbl = Arrow.Table("data.arrow") # a file-format path: memory-mapped +tbl = Arrow.Table(io) # an IO: read to the end +tbl = Arrow.Table(bytes) # IPC bytes already in memory ``` -### `Arrow.Table` - -The type of `table` in this example will be an `Arrow.Table`. When "reading" the arrow data, `Arrow.Table` first ["mmapped"](https://en.wikipedia.org/wiki/Mmap) the `data.arrow` file, which is an important technique for dealing with data larger than available RAM on a system. By "mmapping" a file, the OS doesn't actually load the entire file contents into RAM at the same time, but file contents are "swapped" into RAM as different regions of a file are requested. Once "mmapped", `Arrow.Table` then inspected the metadata in the file to determine the number of columns, their names and types, at which byte offset each column begins in the file data, and even how many "batches" are included in this file (arrow tables may be partitioned into one or more "record batches" each containing portions of the data). Armed with all the appropriate metadata, `Arrow.Table` then created custom array objects ([`Arrow.ArrowVector`](@ref)), which act as "views" into the raw arrow memory bytes. This is a significant point in that no extra memory is allocated for "data" when reading arrow data. This is in contrast to if we wanted to read data from a csv file as columns into Julia structures; we would need to allocate those array structures ourselves, then parse the file, "filling in" each element of the array with the data we parsed from the file. Arrow data, on the other hand, is *already laid out in memory or on disk* in a binary format, and as long as we have the metadata to interpret the raw bytes, we can figure out whether to treat those bytes as a `Vector{Float64}`, etc. A sample of the kinds of arrow array types you might see when deserializing arrow data, include: +`Arrow.Table` satisfies the Tables.jl columns interface, so it works with +every Tables.jl-aware sink and consumer: -* [`Arrow.Primitive`](@ref): the most common array type for simple, fixed-size elements like integers, floats, time types, and decimals -* [`Arrow.List`](@ref): an array type where its own elements are also arrays of some kind, like string columns, where each element can be thought of as an array of characters -* [`Arrow.FixedSizeList`](@ref): similar to the `List` type, but where each array element has a fixed number of elements itself; you can think of this like a `Vector{NTuple{N, T}}`, where `N` is the fixed-size width -* [`Arrow.Map`](@ref): an array type where each element is like a Julia `Dict`; a list of key value pairs like a `Vector{Dict}` -* [`Arrow.Struct`](@ref): an array type where each element is an instance of a custom struct, i.e. an ordered collection of named & typed fields, kind of like a `Vector{NamedTuple}` -* [`Arrow.DenseUnion`](@ref): an array type where elements may be of several different types, stored compactly; can be thought of like `Vector{Union{A, B}}` -* [`Arrow.SparseUnion`](@ref): another array type where elements may be of several different types, but stored as if made up of identically lengthed child arrays for each possible type (less memory efficient than `DenseUnion`) -* [`Arrow.DictEncoded`](@ref): a special array type where values are "dictionary encoded", meaning the list of unique, possible values for an array are stored internally in an "encoding pool", whereas each stored element of the array is just an integer "code" to index into the encoding pool for the actual value. +```julia +Tables.columnnames(tbl) # [:a, :b] +Tables.schema(tbl) +tbl.a # property access = a column +Tables.getcolumn(tbl, :b) +length(tbl) # number of rows +``` -And while these custom array types do subtype `AbstractArray`, there is no current support for `setindex!`. Remember, these arrays are "views" into the raw arrow bytes, so for array types other than `Arrow.Primitive`, it gets pretty tricky to allow manipulating those raw arrow bytes. Nevetheless, it's as simple as calling `copy(x)` where `x` is any `ArrowVector` type, and a normal Julia `Vector` type will be fully materialized (which would then allow mutating/manipulating values). +Pass `tbl` directly to any Tables.jl-compatible sink. -So, what can you do with an `Arrow.Table` full of data? Quite a bit actually! +Columns are **materialized**: each column is a plain Julia `Vector` with a +concrete element type determined by the Arrow schema (see [Type +mapping when reading](@ref)). A `Table` therefore does not borrow the source bytes after +it is constructed, and its columns behave like any other Julia vectors. -Because `Arrow.Table` implements the [Tables.jl](https://juliadata.github.io/Tables.jl/stable/) interface, it opens up a world of integrations for using arrow data. A few examples include: +### Memory mapping and `release!` -* `df = DataFrame(Arrow.Table(file))`: Build a [`DataFrame`](https://juliadata.github.io/DataFrames.jl/stable/), using the arrow vectors themselves; this allows utilizing a host of DataFrames.jl functionality directly on arrow data; grouping, joining, selecting, etc. -* `df = copy(DataFrame(Arrow.Table(file)))`: Build a [`DataFrame`](https://juliadata.github.io/DataFrames.jl/stable/), where the columns are regular in-memory vectors (specifically, `Base.Vector`s and/or `PooledVector`s). This requires that you have enough memory to load the entire `DataFrame` into memory. -* `Tables.datavaluerows(Arrow.Table(file)) |> @map(...) |> @filter(...) |> DataFrame`: use [`Query.jl`'s](https://www.queryverse.org/Query.jl/stable/standalonequerycommands/) row-processing utilities to map, group, filter, mutate, etc. directly over arrow data. -* `Arrow.Table(file) |> SQLite.load!(db, "arrow_table")`: load arrow data directly into an sqlite database/table, where sql queries can be executed on the data -* `Arrow.Table(file) |> CSV.write("arrow.csv")`: write arrow data out to a csv file +A path to a file-format source is memory-mapped by default (`mmap=true`), so +reading a large file does not copy it into RAM up front; pass `mmap=false` +to read it into memory instead (a stream-format path is always read into +memory). A memory map is released when the last reference to it +is garbage collected. To release it deterministically — required on Windows +before a still-mapped file can be deleted, and useful anywhere for prompt +resource release — call [`Arrow.release!`](@ref): -A full list of Julia packages leveraging the Tables.jl inteface can be found [here](https://github.com/JuliaData/Tables.jl/blob/master/INTEGRATIONS.md). +```julia +tbl = Arrow.Table("data.arrow") +# ... use tbl ... +Arrow.release!(tbl) # unmaps NOW; tbl's columns remain usable +rm("data.arrow") +``` -Apart from letting other packages have all the fun, an `Arrow.Table` itself can be plenty useful. For example, with `tbl = Arrow.Table(file)`: -* `tbl[1]`: retrieve the first column via indexing; the number of columns can be queried via `length(tbl)` -* `tbl[:col1]` or `tbl.col1`: retrieve the column named `col1`, either via indexing with the column name given as a `Symbol`, or via "dot-access" -* `for col in tbl`: iterate through columns in the table -* `AbstractDict` methods like `haskey(tbl, :col1)`, `get(tbl, :col1, nothing)`, `keys(tbl)`, or `values(tbl)` +`release!` is idempotent. Because a `Table`'s columns are copies, a released +`Table` remains fully usable; a released [`Arrow.Stream`](@ref) refuses +further iteration cleanly. Every yielded `Table` shares the source lifetime: +releasing a batch closes its parent `Stream`, while the batch's materialized +columns remain usable. -### Arrow types +### `Arrow.Stream` -In the arrow data format, specific logical types are supported, a list of which can be found [here](https://arrow.apache.org/docs/status.html#data-types). These include booleans, integers of various bit widths, floats, decimals, time types, and binary/string. While most of these map naturally to types builtin to Julia itself, there are a few cases where the definitions are slightly different, and in these cases, by default, they are converted to more "friendly" Julia types (this auto conversion can be avoided by passing `convert=false` to `Arrow.Table`, like `Arrow.Table(file; convert=false)`). Examples of arrow to julia type mappings include: +[`Arrow.Stream`](@ref) iterates a source one record batch at a time; each +iteration yields an `Arrow.Table` for that batch — the tool for pipelines +that process batches independently: -* `Date`, `Time`, `Timestamp`, and `Duration` all have natural Julia defintions in `Dates.Date`, `Dates.Time`, `TimeZones.ZonedDateTime`, and `Dates.Period` subtypes, respectively. -* `Char` and `Symbol` Julia types are mapped to arrow string types, with additional metadata of the original Julia type; this allows deserializing directly to `Char` and `Symbol` in Julia, while other language implementations will see these columns as just strings -* Similarly to the above, the `UUID` Julia type is mapped to a 128-bit `FixedSizeBinary` arrow type. -* `Decimal128` and `Decimal256` have no corresponding builtin Julia types, so they're deserialized using a compatible type definition in Arrow.jl itself: `Arrow.Decimal` +```julia +for batch in Arrow.Stream("big.arrow") + process(batch) # batch isa Arrow.Table +end +``` +A `Stream` satisfies `Tables.partitions` (each batch is one partition), so +partition-aware sinks see the source's batch structure — +`Arrow.write(sink, Arrow.Stream(...))` writes one record batch per input +batch. + +Memory: only the memory-mapped **file-format** path avoids holding the whole +source. Batches decode from the map one at a time, so a loop over such a +`Stream` holds one batch of columns plus the file's dictionaries. That is the +way to process a file larger than RAM. A +`Stream` security budget (the `max_total_allocated_bytes` field of +[`Arrow.Limits`](@ref)) stays cumulative across every batch it yields, even +after the consumer drops that batch. Raise the limit explicitly for a trusted +large file whose total decoded allocation exceeds the default; the same +`limits` keyword applies to [`Arrow.Table`](@ref): -Note that when `convert=false` is passed, data will be returned in Arrow.jl-defined types that exactly match the arrow definitions of those types; the authoritative source for how each type represents its data can be found in the arrow [`Schema.fbs`](https://github.com/apache/arrow/blob/master/format/Schema.fbs) file. +```julia +trusted_limits = Arrow.Limits(max_total_allocated_bytes = 2 * 1024^3) +stream = Arrow.Stream("big.arrow"; limits = trusted_limits) +``` -One note on performance: when writing `TimeZones.ZonedDateTime` columns to the arrow format (via `Arrow.write`), it is preferrable to "wrap" the columns in `Arrow.ToTimestamp(col)`, as long -as the column has `ZonedDateTime` elements that all share a common timezone. This ensures the writing process can know "upfront" which timezone will be encoded and is thus much more -efficient and performant. +A **stream-format** source is read to the end and every batch is decoded when +the `Stream` is constructed. A file-format `IO` or byte-vector input is +also read to the end (the whole source is held in memory), but its record +batches are still decoded lazily, one per iteration. `Arrow.write` +materializes every partition before writing (see +[Writing](@ref "manual-writing")), so it does not bound memory either. -#### Custom types +### Metadata -To support writing your custom Julia struct, Arrow.jl utilizes the format's mechanism for "extension types" by allowing the storing of Julia type name and metadata in the field metadata. To "hook in" to this machinery, custom types can utilize the interface methods defined in the `Arrow.ArrowTypes` submodule. For example: +Schema-level and per-column key/value metadata carried in the IPC schema is +readable through the [DataAPI.jl](https://github.com/JuliaData/DataAPI.jl) +metadata interface. Code that imports DataAPI must add it as a direct +dependency: ```julia -using Arrow - -struct Person - id::Int - name::String -end +import Pkg +Pkg.add("DataAPI") +``` -# overload interface method for custom type Person; return a symbol as the "name" -# this instructs Arrow.write what "label" to include with a column with this custom type -const NAME = Symbol("JuliaLang.MyPackage.Person") -ArrowTypes.arrowname(::Type{Person}) = NAME -# overload JuliaType on `Val{NAME}`, which is like a dispatchable string -# return our custom *type* Person; this enables Arrow.Table to know how the "label" -# on a custom column should be mapped to a Julia type and deserialized -ArrowTypes.JuliaType(::Val{NAME}) = Person +```@example metadata_api +using Arrow, DataAPI -table = (col1=[Person(1, "Bob"), Person(2, "Jane")],) io = IOBuffer() -Arrow.write(io, table) -seekstart(io) -table2 = Arrow.Table(io) +Arrow.write(io, (a = [1, 2],); + metadata = ["source" => "docs"], + colmetadata = Dict(:a => ["unit" => "count"])) +metadata_table = Arrow.Table(take!(io)) + +( + schema_keys = collect(DataAPI.metadatakeys(metadata_table)), + source = DataAPI.metadata(metadata_table, "source"), + column_keys = collect(DataAPI.colmetadatakeys(metadata_table, :a)), + unit = DataAPI.colmetadata(metadata_table, :a, "unit"), +) ``` -In this example, we're writing our `table`, which is a NamedTuple with one column named `col1`, which has two -elements which are instances of our custom `Person` struct. We overload `Arrowtypes.arrowname` so that -Arrow.jl knows how to serialize our `Person` struct. We then overload `ArrowTypes.JuliaType` so the deserialization process knows how to map from our type label back to our `Person` struct type. We can then write our data in the arrow format to an in-memory `IOBuffer`, then read the table back in using `Arrow.Table`. -The table we get back will be an `Arrow.Table`, with a single `Arrow.Struct` column with element type `Person`. +Arrow IPC permits duplicate field names. Use integer column positions for +such a table. Name-based property access, `Tables.getcolumn`, and DataAPI +column metadata access are ambiguous and throw an `ArgumentError`. Scan +pushdown also refuses duplicate names. Positional access preserves every +column and its metadata: -Note that without calling `Arrowtypes.JuliaType`, we may get into a weird limbo state where we've written -our table with `Person` structs out as a table, but when reading back in, Arrow.jl doesn't know what a `Person` is; -deserialization won't fail, but we'll just get a `Namedtuple{(:id, :name), Tuple{Int, String}}` back instead of `Person`. - -While this example is very simple, it shows the basics to allow a custom type to be serialized/deserialized. But the `ArrowTypes` module offers even more powerful functionality for "hooking" non-native arrow types into the serialization/deserialization processes. Let's walk through a couple more examples; if you've had enough custom type shenanigans, feel free to skip to the next section. +```julia +Tables.getcolumn(tbl, 1) +DataAPI.colmetadata(tbl, 1, "key") +``` -Let's take a look at how Arrow.jl allows serializing the `nothing` value, which is often referred to as the "software engineer's NULL" in Julia. While Arrow.jl treats `missing` as the default arrow NULL value, `nothing` is pretty similar, but we'd still like to treat it separately if possible. Here's how we enable serialization/deserialization in the `ArrowTypes` module: +### Type mapping when reading + +Scalar Arrow types map to Julia element types by a closed rule over the +schema — never by inspecting values, so an all-`missing` or zero-row scalar +column has the same element type as a populated one. A nullable Arrow field +maps to `Union{Missing, T}`. + +| Arrow type | Julia element type | +|---|---| +| Int8…Int64, UInt8…UInt64 | the same-width `Integer` | +| Float16/32/64 | `Float16`/`Float32`/`Float64` | +| Bool | `Bool` | +| Utf8, LargeUtf8, Utf8View | `String` | +| Binary, LargeBinary, BinaryView, FixedSizeBinary | `Vector{UInt8}` | +| Date32 | `Dates.Date` | +| Date64 | `Dates.DateTime` | +| Timestamp (second, millisecond) | `Dates.DateTime` (UTC instant; a declared timezone stays in the retained schema — loading TimeZones.jl reads these as `ZonedDateTime` instead) | +| Timestamp (microsecond, nanosecond) | `Int64` (raw storage — `DateTime` cannot represent it) | +| Time32/Time64 | `Dates.Time` | +| Duration | `Dates.Second`/`Millisecond`/`Microsecond`/`Nanosecond` by unit | +| Decimal32/64 | `Int32`/`Int64` (unscaled integer storage) | +| Decimal128/256 | `Vector{UInt8}` (raw native-endian storage) | +| Interval | `Int32` (year-month) or a `NamedTuple` (day-time, month-day-nano) | +| Dictionary-encoded scalar | the mapping of the *value* type (indices are resolved) | + +Composite layouts (list, struct, map, union) are read on the dynamic path +(each row is built as a Julia value), and their column element type is +derived from the schema — the declared row container — so it too is the +same for a zero-row, an all-`missing`, and a populated column, with two +exceptions: the heterogeneous-union case in the Union row, and nulls the +field did not declare (below). The two wrapper layouts are transparent: +a dictionary-encoded or run-end-encoded column takes the route and the +element type of its value child (a closed scalar child keeps the typed +path): + +| Arrow type | Element type | +|---|---| +| List, LargeList, FixedSizeList, ListView, LargeListView | `Vector{Any}` (rows are vectors of the child's values) | +| Struct | `Vector{Pair{String,Any}}` (rows are ordered name => value pairs) | +| Map | `Vector{Pair{Any,Any}}` | +| Union | the join of the children's element types when it is concrete (a homogeneous union reads as that type); otherwise `Any`, narrowed from the rows | +| Run-end encoded | the *values* child's element type (runs are expanded) | +| Null | `Missing` | + +Three refinements. A list-family field with a registered ArrowTypes.jl child +uses a row vector with that restored child element type; this is the same for +a zero-row and a populated column. The `Dates` conversions above apply at the top level and +through dictionary encoding; a temporal type nested under a run-end-encoded +or union wrapper stays in its raw integer storage. And field nullability is +*advisory* in Arrow (the reference implementation and the conformance +corpus accept a null under a `nullable=false` field), so a column that +holds a null its field did not declare reads as `Union{Missing, T}` rather +than failing; a conforming column keeps its declared, `Missing`-free type. +Sub-millisecond timestamps stay as raw integers everywhere rather than +silently truncating into `DateTime`; the same rule applies when writing. + +### Scan pushdown + +`Arrow.Table` accepts a `Tables.Scan` — a plain-data description of which +columns to keep, which rows qualify, and how many — and pushes it down into +the reader: ```julia -ArrowTypes.ArrowKind(::Type{Nothing}) = ArrowTypes.NullKind() -ArrowTypes.ArrowType(::Type{Nothing}) = Missing -ArrowTypes.toarrow(::Nothing) = missing -const NOTHING = Symbol("JuliaLang.Nothing") -ArrowTypes.arrowname(::Type{Nothing}) = NOTHING -ArrowTypes.JuliaType(::Val{NOTHING}) = Nothing -ArrowTypes.fromarrow(::Type{Nothing}, ::Missing) = nothing +using Tables: Scan, col, colcmp, colin, isnull + +scan = Scan(select = (:id, :amount), + filter = (col(:amount) > 100) & !isnull(col(:id)), + limit = 1_000) +tbl = Arrow.Table("orders.arrow"; scan = scan) ``` -Let's walk through what's going on here, line-by-line: - * `ArrowKind` overload: `ArrowKind`s are generic "categories" of types supported by the arrow format, like `PrimitiveKind`, `ListKind`, etc. They each correspond to a different data layout strategy supported in the arrow format. Here, we define `nothing`'s kind to be `NullKind`, which means no actual memory is needed for storage, it's strictly a "metadata" type where we store the type and # of elements. In our `Person` example, we didn't need to overload this since types declared like `struct T` or `mutable struct T` are defined as `ArrowTypes.StructKind` by default - * `ArrowType` overload: here we're signaling that our type (`Nothing`) maps to the natively supported arrow type of `Missing`; this is important for the serializer so it knows which arrow type it will be serializing. Again, we didn't need to overload this for `Person` since the serializer knows how to serialize custom structs automatically by using reflection methods like `fieldnames(T)` and `getfield(x, i)`. - * `ArrowTypes.toarrow` overload: this is a sister method to `ArrowType`; we said our type will map to the `Missing` arrow type, so here we actually define ___how___ it converts to the arrow type; and in this case, it just returns `missing`. This is yet another method that didn't show up for `Person`; why? Well, as we noted in `ArrowType`, the serializer already knows how to serialize custom structs by using all their fields; if, for some reason, we wanted to omit some fields or otherwise transform things, then we could define corresponding `ArrowType` and `toarrow` methods - * `arrowname` overload: similar to our `Person` example, we need to instruct the serializer how to label our custom type in the arrow type metadata; here we give it the symbol `Symbol("JuliaLang.Nothing")`. Note that while this will ultimately allow us to disambiguate `nothing` from `missing` when reading arrow data, if we pass this data to other language implementations, they will only treat the data as `missing` since they (probably) won't know how to "understand" the `JuliaLang.Nothing` type label - * `JuliaType` overload: again, like our `Person` example, we instruct the deserializer that when it encounters the `JuliaLang.Nothing` type label, it should treat those values as `Nothing` type. - * And finally, `fromarrow` overload: this allows specifying how the native-arrow data should be converted back to our custom type. `fromarrow(T, x...)` by default will call `T(x...)`, which is why we didn't need this overload for `Person`, but in this example, `Nothing(missing)` won't work, so we define our own custom conversion. +* `select`: a reference or tuple of select items (`ref`, `ref => name`, + `ref => Type`, `ref => Type => name`; refs are `Symbol`, `String`, `Int`, + `Regex`, `Tables.Not`, `Tables.All`). Only the selected columns and the + columns the filter references are decoded. For every decoded batch, complete + node and buffer metadata is validated first. Other columns' body buffers are + not sliced, decompressed, content-validated, or materialized. +* `filter`: an expression over `Tables.col` — comparisons against literals + (`>`, `>=`, `<`, `<=`, `colcmp`), `colin`, `isnull`, string + predicates, combined with `&`, `|`, `!`. A row is kept iff the predicate + is exactly `true` (`missing` excludes, SQL-style). +* `limit`/`offset`: applied to qualifying rows. + +On the file format, batches whose footer statistics prove no row can match +receive no dedicated metadata or body request and are not decoded; the filter +is evaluated batch by batch and `limit`/`offset` compose exactly over the +qualifying rows. Without a filter, whole batches outside the window are never +decoded, and with one decoding stops as soon as the window is full. On the stream format +the scan is applied after decode with identical results. + +A scan whose filter has no semantics-preserving storage representation falls +back to reading the whole source and evaluating over converted public values. +This includes an inexact literal and a temporal conversion that aliases values +or wraps ordering. The detailed rules: + +* Temporal lowering: Date64 and millisecond Timestamp equality can lower, but + their ordered comparisons stay public; Timestamp-second and Time predicates + also stay public. Duration lowering accepts a literal in the column unit or + a coarser fixed unit, but a finer unit stays public because Julia can + overflow while promoting stored values. +* Set membership: temporal Tuple and Array membership follows the same + equality rules. Set members must also have the column's canonical public + type, which preserves `isequal` and hashing. A custom membership object + stays public because Arrow cannot transform it without changing its `in` + semantics. +* Empty projection: `select = ()` stays on the ranged path. It preserves the + selected row count without fetching output column bodies. + +Filters over a field that contains a registered ArrowTypes.jl extension value +at any depth are evaluated over the restored public values. The ArrowTypes +interface does not require `toarrow` to preserve Julia comparison semantics, +so lowering an arbitrary custom literal for physical pushdown would be +incorrect. Projection and the other scan operations still apply normally. + +Batch pruning uses per-batch statistics (row count, null count, min, max +per column) carried in the file's footer schema metadata under the key +`JuliaArrow:batch_statistics.v1`, in the value layout of the Arrow +project's statistics schema — other readers see ordinary metadata. +[`Arrow.write`](@ref) does not embed them; files that carry them prune, +files that do not are simply scanned batch by batch. + +### Reading remote and partial files + +The file format is random-access: the footer says where every batch and +buffer lives, so a reader that can fetch byte ranges — from object storage, +over HTTP, or from a local file it prefers not to map whole — needs only +the ranges its scan touches. [`Arrow.AbstractArrowSource`](@ref) is that +contract: a byte-addressable object of known length, read through +[`Arrow.sourcelength`](@ref) and [`Arrow.readrange`](@ref). Given one, +`Arrow.Table` with a scan first fetches and caches a tail window. That window +normally contains the complete Footer. If it does not, one exact cached +follow-up fetch retrieves the Footer. The reader then keeps only the batches +that the Footer's statistics and the scan's window allow, fetches those +batches' metadata, and requests ranges that cover the selected and +filter-referenced buffers. The common path has three sequential request rounds: +tail/Footer, metadata, and body. An oversized Footer adds a fourth round. Each +metadata or body round can contain multiple coalesced requests, which a +concurrent source can issue in parallel. Tail reads and coalescing may +physically over-read unrequested bytes. + +With [CloudStore.jl](https://github.com/JuliaServices/CloudStore.jl) +loaded, a `CloudStore.Object` (S3 or Azure Blob Storage) is such a source +directly, and its planned ranges are requested concurrently: + +```julia +using Arrow, Tables, CloudStore +obj = CloudStore.Object(bucket, "events/2024-05.arrow"; credentials) +tbl = Arrow.Table(obj; scan = Scan(select = (:id,), filter = col(:day) > 20)) +``` -Let's run through one more complex example, just for fun and to really see how far the system can be pushed: +Any other transport is two methods away: ```julia -using Intervals -table = (col = [ - Interval{Closed,Unbounded}(1,nothing), -],) -const NAME = Symbol("JuliaLang.Intervals.Interval") -ArrowTypes.arrowname(::Type{Interval{T, L, R}}) where {T, L, R} = NAME -const LOOKUP = Dict( - "Closed" => Closed, - "Unbounded" => Unbounded -) -ArrowTypes.arrowmetadata(::Type{Interval{T, L, R}}) where {T, L, R} = string(L, ".", R) -function ArrowTypes.JuliaType(::Val{NAME}, ::Type{NamedTuple{names, types}}, meta) where {names, types} - L, R = split(meta, ".") - return Interval{fieldtype(types, 1), LOOKUP[L], LOOKUP[R]} +struct HTTPSource <: Arrow.AbstractArrowSource + url::String + length::Int64 end -ArrowTypes.fromarrow(::Type{Interval{T, L, R}}, first, last) where {T, L, R} = Interval{L, R}(first, R == Unbounded ? nothing : last) -io = Arrow.tobuffer(table) -tbl = Arrow.Table(io) +Arrow.sourcelength(s::HTTPSource) = s.length +Arrow.readrange(s::HTTPSource, offset, len) = fetchbytes(s.url, offset, len) # a Range GET + +tbl = Arrow.Table(HTTPSource(url, objectsize); scan = Scan(select = (:id,))) ``` -Again, let's break down what's going on here: - * Here we're trying to save an `Interval` type in the arrow format; this type is unique in that it has two type parameters (`Closed` and `Unbounded`) that are not inferred/based on fields, but are just "type tags" on the type itself - * Note that we define a generic `arrowname` method on all `Interval`s, regardless of type parameters. We just want to let arrow know which general type we're dealing with here - * Next we use a new method `ArrowTypes.arrowmetadata` to encode the two non-field-based type parameters as a string with a dot delimiter; we encode this information here because remember, we have to match our `arrowname` Symbol typename in our `JuliaType(::Val(name))` definition in order to dispatch correctly; if we encoded the type parameters in `arrowname`, we would need separate `arrowname` definitions for each unique combination of those two type parameters, and corresponding `JuliaType` definitions for each as well; yuck. Instead, we let `arrowname` be generic to our type, and store the type parameters *for this specific column* using `arrowmetadata` - * Now in `JuliaType`, note we're using the 3-argument overload; we want the `NamedTuple` type that is the native arrow type our `Interval` is being serialized as; we use this to retrieve the 1st type parameter for our `Interval`, which is simply the type of the two `first` and `last` fields. Then we use the 3rd argument, which is whatever string we returned from `arrowmetadata`. We call `L, R = split(meta, ".")` to parse the two type parameters (in this case `Closed` and `Unbounded`), then do a lookup on those strings from a predefined `LOOKUP` Dict that matches the type parameter name as string to the actual type. We then have all the information to recreate the full `Interval` type. Neat! - * The one final wrinkle is in our `fromarrow` method; `Interval`s that are `Unbounded`, actually take `nothing` as the 2nd argument. So letting the default `fromarrow` definition call `Interval{T, L, R}(first, last)`, where `first` and `last` are both integers isn't going to work. Instead, we check if the `R` type parameter is `Unbounded` and if so, pass `nothing` as the 2nd arg, otherwise we can pass `last`. +Overriding [`Arrow.concurrentreads`](@ref) lets Arrow issue planned ranges +concurrently through `readrange`. One reader samples that value once, clamps +it by the `max_concurrent_reads` field of [`Arrow.Limits`](@ref), and shares +the cap across all operations on that read. Separate reads are +independent; enforce a transport-wide cap inside `readrange` when required. +Results are placed by request, and the default reads one at a time. +Without a scan the whole object is read, as is a stream-format object (no +footer) and a scan that cannot be pushed down. Arrow.jl has no HTTP or +cloud dependency of its own. -This stuff can definitely make your eyes glaze over if you stare at it long enough. As always, don't hesitate to reach out for quick questions on the [#data](https://julialang.slack.com/messages/data/) slack channel, or [open a new issue](https://github.com/apache/arrow-julia/issues/new) detailing what you're trying to do. +## [Writing](@id manual-writing) -### `Arrow.Stream` +### `Arrow.write` -In addition to `Arrow.Table`, the Arrow.jl package also provides `Arrow.Stream` for processing arrow data. While `Arrow.Table` will iterate all record batches in an arrow file/stream, concatenating columns, `Arrow.Stream` provides a way to *iterate* through record batches, one at a time. Each iteration yields an `Arrow.Table` instance, with columns/data for a single record batch. This allows, if so desired, "batch processing" of arrow data, one record batch at a time, instead of creating a single long table via `Arrow.Table`. +[`Arrow.write`](@ref) writes any Tables.jl-compatible source to a path or +an `IO`: -### Custom application metadata +```julia +Arrow.write("out.arrow", tbl) # file format (ARROW1 + footer) +Arrow.write(io, tbl; file = false) # stream format +Arrow.write("out.arrow", tbl; compress = :zstd) # or :lz4 +Arrow.write("out.arrow", tbl; + metadata = ["source" => "sensor-7"], + colmetadata = Dict(:temp => ["unit" => "C"])) +``` -The Arrow format allows data producers to [attach custom metadata](https://arrow.apache.org/docs/format/Columnar.html#custom-application-metadata) to various Arrow objects. +Each `Tables.partitions` partition of the source becomes one record batch, +so `Arrow.write(sink, Arrow.Stream(path))` and +`Arrow.write(sink, Tables.partitioner(...))` preserve batch structure. The +writer is eager and whole-buffer: every partition is materialized, encoded +and validated in memory, then the complete IPC bytes are written to the +sink once — it holds the whole table, so it is not a bounded-memory path +for sources larger than RAM. Use [`Arrow.Writer`](@ref) when tables must be +published incrementally. -Arrow.jl provides a convenient accessor for this metadata via [`Arrow.getmetadata`](@ref). `Arrow.getmetadata(t::Arrow.Table)` will return an immutable `AbstractDict{String,String}` that represents the [`custom_metadata` of the table's associated `Schema`](https://github.com/apache/arrow/blob/85d8175ea24b4dd99f108a673e9b63996d4f88cc/format/Schema.fbs#L515) (or `nothing` if no such metadata exists), while `Arrow.getmetadata(c::Arrow.ArrowVector)` will return a similar representation of [the column's associated `Field` `custom_metadata`](https://github.com/apache/arrow/blob/85d8175ea24b4dd99f108a673e9b63996d4f88cc/format/Schema.fbs#L480) (or `nothing` if no such metadata exists). +`compress` applies per-buffer LZ4 frame or Zstandard compression as +defined by the IPC specification (buffers that do not shrink are stored +raw). Compressed files are readable by every implementation that supports +IPC compression. -To attach custom schema/column metadata to Arrow tables at serialization time, see the `metadata` and `colmetadata` keyword arguments to [`Arrow.write`](@ref). +### Incremental writing and stream append -## Writing arrow data +[`Arrow.Writer`](@ref) publishes each table's batches before the next table +is supplied, so it holds only the current table in memory: -Ok, so that's a pretty good rundown of *reading* arrow data, but how do you *produce* arrow data? Enter `Arrow.write`. +```julia +Arrow.Writer("out.arrow") do writer + for table in tables + Arrow.write(writer, table) + end +end +``` -### `Arrow.write` +The first table fixes the schema. Later tables must have the same columns in +the same order and compatible types. Closing the writer finishes the output; +an abandoned writer can leave a torn stream or a file without its footer. -With `Arrow.write`, you provide either an `io::IO` argument or a [`file_path`](#support-for-generic-path-like-types) to write the arrow data to, as well as a Tables.jl-compatible source that contains the data to be written. +[`Arrow.append`](@ref) extends an existing IPC stream in place. It refuses the +IPC file format, whose footer must be written when the file is closed. Changed +dictionary pools require a stream created with +`Arrow.Writer(...; file=false, dictreplacement=true)`. -What are some examples of Tables.jl-compatible sources? A few examples include: -* `Arrow.write(io, df::DataFrame)`: A `DataFrame` is a collection of indexable columns -* `Arrow.write(io, CSV.File(file))`: read data from a csv file and write out to arrow format -* `Arrow.write(io, DBInterface.execute(db, sql_query))`: Execute an SQL query against a database via the [`DBInterface.jl`](https://github.com/JuliaDatabases/DBInterface.jl) interface, and write the query resultset out directly in the arrow format. Packages that implement DBInterface include [SQLite.jl](https://juliadatabases.github.io/SQLite.jl/stable/), [MySQL.jl](https://juliadatabases.github.io/MySQL.jl/dev/), and [ODBC.jl](http://juliadatabases.github.io/ODBC.jl/latest/). -* `df |> @map(...) |> Arrow.write(io)`: Write the results of a [Query.jl](https://www.queryverse.org/Query.jl/stable/) chain of operations directly out as arrow data -* `jsontable(json) |> Arrow.write(io)`: Treat a json array of objects or object of arrays as a "table" and write it out as arrow data using the [JSONTables.jl](https://github.com/JuliaData/JSONTables.jl) package -* `Arrow.write(io, (col1=data1, col2=data2, ...))`: a `NamedTuple` of `AbstractVector`s or an `AbstractVector` of `NamedTuple`s are both considered tables by default, so they can be quickly constructed for easy writing of arrow data if you already have columns of data +### Dictionary encoding -And these are just a few examples of the numerous [integrations](https://github.com/JuliaData/Tables.jl/blob/master/INTEGRATIONS.md). +Wrap a column in [`Arrow.DictEncode`](@ref) to write it dictionary-encoded +(a category pool plus integer indices), which is what a categorical or +low-cardinality string column wants: -In addition to just writing out a single "table" of data as a single arrow record batch, `Arrow.write` also supports writing out multiple record batches when the input supports the `Tables.partitions` functionality. One immediate, though perhaps not incredibly useful example, is `Arrow.Stream`. `Arrow.Stream` implements `Tables.partitions` in that it iterates "tables" (specifically `Arrow.Table`), and as such, `Arrow.write` will iterate an `Arrow.Stream`, and write out each `Arrow.Table` as a separate record batch. Another important point for why this example works is because an `Arrow.Stream` iterates `Arrow.Table`s that all have the same schema. This is important because when writing arrow data, a "schema" message is always written first, with all subsequent record batches written with data matching the initial schema. +```julia +Arrow.write("out.arrow", (region = Arrow.DictEncode(regions), sales = sales)) +``` -In addition to inputs that support `Tables.partitions`, note that the Tables.jl itself provides the `Tables.partitioner` function, which allows providing your own separate instances of similarly-schema-ed tables as "partitions", like: +Reading a dictionary-encoded column resolves the indices: the column comes +back as its value type. When a `Table` read from Arrow is written again, +its dictionary encoding is preserved. Fresh `DictEncode` output coalesces +categories by exact Arrow storage identity. A retained rewrite preserves the +source pool prefix, including unused or duplicate physical categories, because +pool order and index meaning are part of the encoded data. + +### Type mapping when writing + +Writing maps Julia element types to Arrow types. The *core* domain — what +can appear at any nesting depth — is: + +| Julia element type | Arrow type | +|---|---| +| `Int8`…`Int64`, `UInt8`…`UInt64` | the same-width integer | +| `Float16/32/64` | the same-width float | +| `Bool` | Bool | +| `String` | Utf8 | +| `Vector{T}` for core `T` (including `Vector{UInt8}`) | List of the mapping of `T` | +| `Union{Missing, T}` for core `T` | the mapping of `T`, nullable | +| `Union{T1, T2, ...}` with two or more non-`Missing` writable member types | canonical dense Union with one child per declared member type | +| `Missing` | Null | + +The dense Union mapping applies to freshly supplied Julia data. Each member +uses the same recursive core or ArrowTypes.jl mapping that it would use at +that nesting depth. A `Missing` member is represented by a Null child. A +Union may have at most 32 declared members. + +An abstract or `Any` element type does not declare those member types. When +Arrow.jl must infer writer or storage types from runtime values, it accepts at +most 8 distinct types across the complete column and all of its partitions. +This rule includes abstract ArrowTypes.jl storage, dictionary pools, and +retained registered ArrowTypes.jl columns. Use an explicit declared `Union` +when a column intentionally has more types. This separate limit bounds +schema-planning and compiler work for runtime-generated parametric types; it +does not reduce the 32-member declared Union limit. + +If an abstract declaration has no `ArrowType` mapping or extension identity of +its own, Arrow uses its observed concrete subtypes as writer evidence. This +keeps a concrete subtype's extension metadata. Multiple observed subtypes form +an explicit Union under the same 8-type inference limit. An empty abstract +column still needs declared schema evidence because it has no runtime subtype. + +At the *top level* of a column the facade adds: + +| Julia element type | Arrow type | +|---|---| +| any other `AbstractString` (e.g. `SubString`) | Utf8 | +| `Dates.Date` | Date32 | +| `Dates.DateTime` | Timestamp (millisecond) | +| `Dates.Time` | Time64 (nanosecond) | +| `Dates.Second/Millisecond/Microsecond/Nanosecond` | Duration of that unit | +| `NamedTuple` whose fields are core columns | Struct; `Union{Missing, T}` adds parent validity while child nullability stays declared | +| `Arrow.DictEncode` over a writable column | Dictionary of the recursive mapping of its values | +| `ArrowStrings.StringVector` | Utf8View, **zero-copy** — the column's memory is the Arrow array (see below) | + +These native facade conversions do not recurse through list or struct shapes: +a `Vector{Date}` inside a list, or a `Date` or `SubString` field of a +`NamedTuple`, is refused with an `ArgumentError` naming the element type. A +top-level `DictEncode` pool uses the same native column mapping as an ordinary +top-level column. The ArrowTypes.jl mappings described below do recurse. A +column with element type `Any` is narrowed once (recovering list columns of a +common element type) and refused if it cannot be narrowed to a writable type. +When the source is an `Arrow.Table` or `Arrow.Stream`, the writer retains the +compatible Arrow descriptor tree. Temporal units, byte and list widths, +Struct, Map, Run-End Encoding, nullability, field metadata, schema metadata, +and top-level dictionary index types and category order survive a read/write +round trip. Buffer sharing, overlapping ListView ranges, and exact run +segmentation are rebuilt into a canonical form without changing logical +values. A retained Map that declares sorted keys is rewritten only when each +row remains sorted. + +A fresh Julia column with a heterogeneous declared `Union` element type is +synthesized as a canonical dense Arrow Union. This is distinct from rewriting +a retained Union. Once an unregistered Arrow Union is materialized, its +original child type IDs and offsets are no longer present in the Julia values. +Writing that retained Union from an `Arrow.Table` fails with a clear +`ArgumentError` instead of inventing new routing under the old schema. + +A registered ArrowTypes.jl target is different. Its public-domain values keep +writer-side type evidence. When that type lowers to Union storage, Arrow.jl can +route each value back through the retained children and preserve external child +order, labels, type IDs, dense or sparse mode, nullability, and metadata. Sparse +children use canonical hidden placeholder values outside their active rows. +This also works when `JuliaType` returns an abstract read target and concrete +writer subtypes provide the storage mapping. A concrete subtype may omit an +extension identity or use the retained parent's identity. A different explicit +extension name or metadata is rejected instead of being silently relabeled. An +outer missing value uses a separate unmarked Null child. +If the logical storage Union already contains `Missing`, adding an outer missing +state is rejected because Arrow cannot distinguish the two states. + +Registered values can also rebuild compatible retained storage descriptors. +This includes binary and binary-view widths, fixed-size binary, list and +fixed-size-list layouts, date and duration units, wide decimals, and interval +layouts. The retained descriptor remains authoritative. Byte widths, list +sizes, temporal exactness, child fields, and interval storage shapes are +checked before output is written. Concrete declared element types provide the +same schema evidence for empty and typed all-missing columns, so those columns +cannot bypass retained-schema checks. + +A nested dictionary has the same fail-closed rule as an unregistered Union +because its pool is not retained. Top-level dictionaries, including +dictionaries of composite values, are retained on a full read. A scan result +may not carry the hidden source pool; an ordered dictionary then fails instead +of inventing category order. A nullable `Dictionary` with an unknown +extension also fails closed: after materialization, `missing` cannot say whether +the source row was a valid index into the Null pool or a null dictionary index. + +### Custom and extension types + +Arrow 3.0 applies [ArrowTypes.jl](https://github.com/apache/arrow-julia/tree/main/src/ArrowTypes) +mappings automatically. Import ArrowTypes.jl directly and declare it as a +dependency of the package that owns the custom type. `Arrow.ArrowTypes` +remains available as a qualified compatibility binding, and `ArrowTypes` +stays exported as an Arrow 2.x compatibility exception. + +Define `ArrowType` and `toarrow` to lower a custom value to a supported storage +type. Define an extension name and the read hooks when the logical type must +round-trip: + +```@example arrowtypes_account_id +using Arrow +import ArrowTypes -```julia -# treat 2 separate NamedTuples of vectors with same schema as 1 table, 2 partitions -tbl_parts = Tables.partitioner([(col1=data1, col2=data2), (col1=data3, col2=data4)]) -Arrow.write(io, tbl_parts) +struct AccountID + value::Int64 +end -# treat an array of csv files with same schema where each file is a partition -# in this form, a function `CSV.File` is applied to each element of 2nd argument -csv_parts = Tables.partitioner(CSV.File, csv_files) -Arrow.write(io, csv_parts) -``` +const ACCOUNT_ID = Symbol("JuliaLang.Example.AccountID") -### `Arrow.Writer` +ArrowTypes.ArrowType(::Type{AccountID}) = Int64 +ArrowTypes.toarrow(id::AccountID) = id.value +ArrowTypes.arrowname(::Type{AccountID}) = ACCOUNT_ID +ArrowTypes.JuliaType(::Val{ACCOUNT_ID}, ::Type{Int64}, metadata) = AccountID +ArrowTypes.fromarrow(::Type{AccountID}, value::Int64) = AccountID(value) -With `Arrow.Writer`, you instantiate an `Arrow.Writer` object, write sources using it, and then close it. This allows for incrmental writes to the same sink. It is similar to `Arrow.append` without having to close and re-open the sink in between writes and without the limitation of only supporting the IPC stream format. +io = IOBuffer() +Arrow.write(io, (id = AccountID.(1:3),); file = false) +table = Arrow.Table(take!(io)) +values = getfield.(table.id, :value) +@assert values == [1, 2, 3] +values +``` -### Multithreaded writing +Arrow applies `ArrowType` and `toarrow` recursively to top-level values and to +values nested in lists, tuples and fixed-size lists, structs, maps, +dictionary-encoded values, and freshly synthesized heterogeneous Unions. It +writes `arrowname` and `arrowmetadata` as standard extension metadata. On read, +it resolves the logical type with `JuliaType` and restores each value with +`fromarrow` or `fromarrowstruct`. +When an extension name has no registered mapping in the current process, +Arrow warns and returns its ordinary storage values. A plain concrete struct +whose fields are supported can use ArrowTypes.jl's default `StructKind` +mapping; without extension hooks, it reads back as ordinary Struct storage. +Use the lowering interface to select a different stable storage +representation. An `ArrowKind` override alone does not select an arbitrary +Arrow 3.0 physical layout. + +Arrow does not intern an unknown on-wire extension name merely to probe +`JuliaType(Val(...))`. It dispatches only when that name is already a Julia +`Symbol`. Unsupported-name warnings are deduplicated by the complete extension +label. Each table materialization emits at most one warning for each of 16 +distinct labels, then one suppression notice for further distinct labels. A +warning shows at most 128 UTF-8 bytes of its label. The built-in +`JuliaLang.Symbol` mapping also lifts only a payload that is already interned. A +novel payload produces `ValidationError` instead of adding permanent +process-global symbol state. + +Recursive custom storage schemas and recursive value containers are rejected +with `ArgumentError`. Custom mappings may nest to 64 levels. Deeper mappings +are rejected before the writer can recurse without a bound. + +Arrow also bounds exact Julia type construction for composite descriptors. +It passes the exact `NTuple{N,T}` storage type to `JuliaType` when `N` is at +most 1024. For a larger fixed-size-list descriptor, it passes the compact +`Tuple{Vararg{T}}` family instead. Generic registrations still resolve. A +registration that requires the exact oversized arity remains unknown, so Arrow +returns ordinary storage values. An extension-labelled Struct receives its +exact `NamedTuple` storage signature through 1024 children only when its child +names are unique, contain no embedded NUL, already exist as Julia `Symbol`s, +are at most 4096 UTF-8 bytes each, and use at most 64 KiB in total. Otherwise +the labelled Struct remains an unknown extension and reads as ordered `Pair` +storage. One bounded compatibility exception preserves ArrowTypes.jl Tuple +storage: when the complete child-name sequence is exactly `"1"`, `"2"`, …, +`string(N)` for `N ≤ 1024`, Arrow may intern only that fixed finite set of +positional names. The exception does not apply to an unknown extension label or +to arbitrary or partly positional Struct names. +If writer-side `ArrowType` resolution returns a concrete tuple with more than +1024 fields, the writer rejects that storage type before specializing on it. +This also applies to ArrowTypes.jl's default mapping for a tuple value. Custom +trait code still runs as ordinary trusted Julia code; Arrow cannot recover if +the hook crashes while it constructs its return value. + +Retained composite rows hidden by a nullable ancestor are constructed from the +retained Field and the required logical length. Arrow does not manufacture a +Julia object for every hidden child slot. This keeps Null-only fixed-size-list +and inactive sparse-Union storage bounded by the bytes the output requires. + +### ArrowStrings columns + +[ArrowStrings.jl](https://github.com/apache/arrow-julia/tree/main/src/ArrowStrings) +(a separate package that lives in this repository) defines +`ArrowString`, a 16-byte string value that *is* an Arrow StringView entry +(inline up to 12 bytes, otherwise a prefix plus buffer index and offset), +and `StringVector`, a column of them over a set of byte buffers — +which *is* an Arrow Utf8View array's memory. A parser or other producer can +build this representation directly. `Arrow.write` then wraps its payload +vector and buffers as the Arrow column without repacking them or +materializing a `String`. + +## Validation + +Every batch decoded by a read, and every batch emitted by a write, is validated: +buffer arity and byte lengths against the schema (structural), and offset +monotonicity, dictionary index domains, union type ids, and the other +data-intrinsic invariants (semantic). Metadata is verified by a generated +FlatBuffers shape verifier before any of it is used. Resource limits (metadata +size, body size, allocation budget, nesting depth) are enforced before any +metadata-directed or package-controlled read-facade allocation. The one +cumulative budget covers byte-range fetch and assembly, IPC decode, Core and +ArrowTypes containers, scan slices and joins, and final table columns. Custom +user hook allocations remain the hook author's responsibility. Vector +reserves include a conservative backing-store capacity, not only the logical +element payload, because supported Julia versions can round that capacity to +the next allocation class. + +Core schema field and child names remain `String` values and are not interned. +The Tables.jl facade is the explicit `Symbol` boundary because Tables column +names are Symbols. Before it interns any novel top-level field name, Arrow +preflights the complete schema: each name is at most 4096 UTF-8 bytes, one table +materialization may add at most 65,536 novel names, and their combined UTF-8 +size is at most 1 MiB. A schema outside those limits produces +`ValidationError` without partially interning its novel names. + +A ranged scan validates the Footer's complete Block index up front and the +complete node and buffer metadata for every statistics-surviving record before +it requests body ranges. It does not request or decode statistics-pruned record +metadata or skipped body content. Corrupt or hostile input that reaches these +validation stages produces a clean `ValidationError` rather than a crash or an +unbounded allocation. + +Content-policy checks that the reference implementation treats as +advisory — UTF-8 well-formedness of string bytes, the `nullable=false` +declaration on a field, canonical zero padding of bitmaps — are not +enforced by default, matching the behavior of the other Arrow +implementations on the ecosystem's own conformance files. + +## The C data interface + +Arrow's [C data interface](https://arrow.apache.org/docs/format/CDataInterface.html) +and [C stream interface](https://arrow.apache.org/docs/format/CStreamInterface.html) +move columns and record batches between implementations in the same +process without copying: a pair of C structs (`ArrowSchema`, `ArrowArray`) +or a stream struct (`ArrowArrayStream`) is filled by a producer and read by +a consumer, and ownership is transferred with a release callback. + +Arrow.jl exposes a low-level interface over `Arrow.Field`, `Arrow.ArrayData`, +`Arrow.Schema`, and `Arrow.RecordBatch`. These names and the exact ABI structs +`Arrow.CArrowSchema`, `Arrow.CArrowArray`, and `Arrow.CArrowArrayStream` are +public but not exported. Qualify them with `Arrow.`. `Arrow.fromjulia` and +`Arrow.batch` build engine values from Julia columns, and `Arrow.materialize` +converts imported column data back to a Julia vector. + +* `Arrow.to_c_data(field, data) -> (schemaptr, arrayptr)` exports one + column; the structs stay valid until the consumer calls their `release` + callbacks, and `Arrow.reap!()` reclaims the export bookkeeping afterward. +* `Arrow.from_c_data(schemaptr, arrayptr) -> (field, data)` imports one + column, *moving* the array (its source `release` is nulled, as the spec + requires). The imported buffers stay valid as long as the returned data is + reachable; `Arrow.release!` on the import's `ForeignOwner` (or on the + owner region behind any imported buffer, `buffer.region`) revokes every + imported buffer and runs the producer's release callback exactly once. +* `Arrow.export_stream!(streamptr, schema, batches)` fills a caller-owned + `ArrowArrayStream`; `Arrow.from_c_stream(streamptr)` imports one and + yields record batches through `Arrow.nextbatch!`. + +For example, handing a column to PyArrow in-process through PythonCall: -By default, `Arrow.write` will use multiple threads to write multiple -record batches simultaneously (e.g. if julia is started with `julia -t 8` or the `JULIA_NUM_THREADS` environment variable is set). The number of concurrent tasks to use when writing can be controlled by passing the `ntasks` keyword argument to `Arrow.write`. Passing `ntasks=1` avoids any multithreading when writing. +```julia +using Arrow, PythonCall +pa = pyimport("pyarrow") -### Compression +f, d = Arrow.fromjulia("x", [1, 2, missing, 4]) +sp, ap = Arrow.to_c_data(f, d) +pyarr = pa.Array._import_from_c(UInt(ap), UInt(sp)) # PyArrow now owns the structs +pyarr.to_pylist() # [1, 2, None, 4] +``` -Compression is supported when writing via the `compress` keyword argument. Possible values include `:lz4`, `:zstd`, or your own initialized `LZ4FrameCompressor` or `ZstdCompressor` objects; will cause all buffers in each record batch to use the respective compression encoding or compressor. +The conformance suite under `conformance/` round-trips the supported layout +families through PyArrow over this path in both directions. It records +explicit skips for invalid fixtures, layouts an external oracle cannot +construct or import, and features that Arrow.jl intentionally rejects. + +## Compiling with JuliaC `--trim` + +Arrow.jl's engine is designed to compile under JuliaC's `--trim=safe`: +type descriptors are runtime values, layout dispatch goes through closed +`isa` ladders, and the value-domain entry points (reading, C data +import/export, and the typed accessors `Arrow.ArrowCore.materialize(::Type{T}, +field, data)`) are statically resolvable. The repository's +`test/trim_compile_tests.jl` gate holds that at zero verifier errors and +warnings. The same gate exercises ArrowStrings construction, inline and view +access, missing values, comparison, and materialization. Arrow.jl itself +supports Julia 1.10 and later; JuliaC's `--trim` needs Julia 1.12, so the gate +runs only there. The dynamic facade conveniences (property access on +`Arrow.Table`, `NamedTuple` rows) are not part of that guarantee. + +## Updating from Arrow.jl 2.x + +Arrow.jl 3.0 changes the storage model and removes some advanced Arrow 2.x +write features. Read [Migrating from Arrow.jl 2.x](@ref) before you update. diff --git a/docs/src/migration.md b/docs/src/migration.md new file mode 100644 index 00000000..f2d14fe6 --- /dev/null +++ b/docs/src/migration.md @@ -0,0 +1,348 @@ +```@meta +CurrentModule = Arrow +``` + + +# Migrating from Arrow.jl 2.x + +Arrow.jl 3.0 is a breaking rewrite. The common names remain the same: +[`Arrow.Table`](@ref), [`Arrow.Stream`](@ref), [`Arrow.write`](@ref), and +[`Arrow.DictEncode`](@ref). The storage model and some advanced write +features changed. + +## Requirements + +Arrow 3.0 requires Julia 1.10 or later, ArrowTypes.jl 2.x, Tables.jl 1.14 +(the first release that provides `Tables.Scan`), and ArrowStrings.jl 1.0. + +## Materialized columns + +Arrow 2.x returned lazy `ArrowVector` objects that viewed the source buffers. +Arrow 3.0 returns plain Julia vectors: + +```julia +table = Arrow.Table("data.arrow") +table.id isa Vector +``` + +There is no `convert=false` mode. Code that checks for an `ArrowVector` type +must use the Tables.jl interface or the column's normal `AbstractVector` +interface instead. + +Concrete Arrow 2.x array types such as `Arrow.Primitive`, `Arrow.List`, and +`Arrow.DictEncoded` are no longer part of the read result. Use +`Tables.schema(table)`, `eltype(column)`, and the standard array interface. +Use `Tables.getcolumn(table, name_or_index)` in place of Arrow 2.x table +indexing. + +Decimal read types changed. Arrow 2.x read decimal columns as `Arrow.Decimal` +values that wrap a BitIntegers.jl `Int128` or `Int256` unscaled integer. Arrow +3.0 reads Decimal32 and Decimal64 as unscaled `Int32` and `Int64` values, and +Decimal128 and Decimal256 as raw native-endian byte vectors. See +[Type mapping when reading](@ref) for the complete table. + +Struct row types changed. Arrow 2.x read a plain (non-extension) Struct +column as `NamedTuple` rows. Arrow 3.0 reads it as `Vector{Pair{String,Any}}` +rows, because Arrow field names are not always valid or unique `Symbol`s. Ask +for typed rows with a `Tables.Scan` select type override; it applies +recursively, so a list-of-struct column takes a `Vector{...}` target: + +```julia +NT = @NamedTuple{a::Float64, b::String} +table = Arrow.Table("data.arrow"; scan=Tables.Scan(select=(:x, :c => NT))) +eltype(table.c) == NT # true +``` + +Code that must read struct columns under both Arrow 2.x and 3.0 can branch on +what it received instead of on the package version: + +```julia +rows = table.c isa AbstractVector{<:NamedTuple} ? table.c : + [r === missing ? missing : NT(Tuple(last(kv) for kv in r)) for r in table.c] +``` + +Timezone-aware timestamps changed. Arrow 2.x depended on TimeZones.jl and +read a timestamp column with a declared timezone as `ZonedDateTime` values. +Arrow 3.0 does not depend on TimeZones.jl: by default such a column reads as +naive UTC `DateTime` values (the stored instants), and the declared zone is +kept in the retained schema for rewrite. Loading TimeZones.jl activates +Arrow's extension and restores the `ZonedDateTime` behavior for second- and +millisecond-unit columns, on both sides: those columns read as +`ZonedDateTime`, and a fresh `ZonedDateTime` column writes as a +timezone-declared millisecond timestamp. One written column carries one +zone; convert mixed-zone values with `astimezone` first. Micro- and +nanosecond timestamps read as raw `Int64` storage either way; neither +`DateTime` nor `ZonedDateTime` can hold them exactly, and Arrow 3.0 never +truncates silently (Arrow 2.x truncated with a warning). + +The old positional byte-window arguments and multi-input constructors were +removed. Pass one complete path, `IO`, byte vector, or byte-range source to +`Arrow.Table` or `Arrow.Stream`. Slice an in-memory byte vector before the call +when needed. Read independent IPC sources separately and combine their tables +with a Tables.jl-aware consumer. + +The source can still be memory-mapped while it is read. The returned columns +do not borrow the map. Call [`Arrow.release!`](@ref) when you need to release a +mapped file at a known time: + +```julia +table = Arrow.Table("data.arrow") +Arrow.release!(table) +rm("data.arrow") + +# The materialized columns are still valid. +sum(table.id) +``` + +## Removed input support + +Arrow 2.x read delta dictionary batches. Arrow 3.0 rejects delta dictionary +batches and big-endian IPC input, so a file that Arrow 2.x accepted can now +fail with a validation error. + +## Writing + +`Arrow.write` is eager. It materializes and validates all input partitions, +builds the complete IPC output, and then writes it to the sink. Arrow 3.0 does +not provide these Arrow 2.x features: + +| Arrow 2.x feature | Arrow 3.0 action | +|---|---| +| `Arrow.Writer` | Still works, reimplemented: an incremental writer for both formats. The FIRST table written fixes the schema; later tables must conform (no cross-table inference). See [`Arrow.Writer`](@ref). | +| `Arrow.append` | Still works for the IPC STREAM format, reimplemented. The file format refuses: produce it incrementally with `Arrow.Writer` or rewrite it. A 2.x stream whose dictionary pools carry null slots needs one rewrite through `Arrow.Writer(path; file=false, dictreplacement=true)` before it accepts pool-changing appends. | +| `table \|> Arrow.write(sink)` | Still works: the curried form is kept. | +| `Arrow.tobuffer(table)` | Still works, and still emits the IPC stream format. | +| `ntasks` | Accepted and ignored with a one-time warning. Encoding is not task-parallel in 3.0, including in `Arrow.Writer`. | +| `Arrow.ToArrow` | Import `ArrowTypes.ToArrow` directly when an explicit lazy conversion view is needed. Normal writes apply the mapping automatically. | + +The output default for an `IO` changed. Arrow 2.x wrote the stream format to +an `IO` by default. Arrow 3.0 uses `file=true` for both paths and `IO` sinks. +Pass `file=false` when you need the IPC stream format: + +```julia +io = IOBuffer() +Arrow.write(io, table; file=false) +stream_bytes = take!(io) +``` + +`Arrow.tobuffer(table)` produces exactly those bytes, seeked to the start. +Use `file=true` when you need an IPC file with footer-based random access. + +Arrow 3.0 keeps `file`, `compress`, `metadata`, and `colmetadata`. Compression +is selected with `:lz4` or `:zstd`; passing an initialized compressor object +is no longer supported. The Arrow 2.x `alignment`, `dictencode`, +`dictencodenested`, `denseunions`, `largelists`, `maxdepth`, and `ntasks` +writer keywords no longer have any effect: `Arrow.write` accepts them with a +one-time warning each and ignores them. Wrap only the columns that need +dictionary encoding in [`Arrow.DictEncode`](@ref). + +Arrow 3.0 again consumes the ArrowTypes.jl mapping interface. Package authors +should depend on and import ArrowTypes.jl directly. Define `ArrowType` and +`toarrow` to lower a custom value to supported storage. Add an extension name +and the read hooks when the logical type must round-trip: + +```julia +import ArrowTypes + +struct AccountID + value::Int64 +end + +const ACCOUNT_ID = Symbol("JuliaLang.Example.AccountID") + +ArrowTypes.ArrowType(::Type{AccountID}) = Int64 +ArrowTypes.toarrow(id::AccountID) = id.value +ArrowTypes.arrowname(::Type{AccountID}) = ACCOUNT_ID +ArrowTypes.JuliaType(::Val{ACCOUNT_ID}, ::Type{Int64}, metadata) = AccountID +ArrowTypes.fromarrow(::Type{AccountID}, value::Int64) = AccountID(value) + +Arrow.write("accounts.arrow", (id = AccountID.(1:3),)) +table = Arrow.Table("accounts.arrow") +getfield.(table.id, :value) == [1, 2, 3] # true +``` + +The lowering and restoration apply recursively to top-level values and values +nested in lists, tuples and fixed-size lists, structs, maps, +dictionary-encoded values, and freshly synthesized heterogeneous Unions. +Arrow writes `arrowname` and `arrowmetadata` as standard extension metadata. +On read, it uses `JuliaType` and then `fromarrow` or `fromarrowstruct`. If the +current process has no mapping for an extension name, Arrow warns and returns +the ordinary storage values instead. +Defining `ArrowKind` alone is not a supported way to select an Arrow 3.0 +physical layout. Use the `ArrowType` and `toarrow` lowering interface. + +[`Arrow.DictEncode`](@ref) remains the opt-in wrapper for a newly written +dictionary-encoded column. Its pool values use the same recursive ArrowTypes.jl +mapping. A fully read top-level dictionary also retains its pool for rewrite. + +## Supported write types + +The writer accepts fixed-width integers and floats, `Bool`, strings, supported +`Dates` values, lists of supported core values, top-level `NamedTuple` struct +columns, and fresh heterogeneous Julia Union columns whose members are +writable at that nesting depth. Fresh heterogeneous Unions use the canonical +dense Arrow Union layout. See [Type mapping when writing](@ref) for the +complete table. + +A top-level `NamedTuple` column may use `Union{Missing, T}`. Arrow writes the +outer missing state in the Struct validity bitmap and keeps each child's +declared nullability unchanged. + +Declared Unions may contain up to 32 members. Runtime writer or storage +inference is limited to 8 distinct types across all partitions. This includes +abstract ArrowTypes.jl storage, abstract or `Any` dictionary values, and +abstract retained ArrowTypes.jl targets. Declare the intended Union when a +column needs more runtime types. + +A fresh unresolved abstract declaration keeps concrete subtype writer evidence +across all partitions. A concrete subtype's extension metadata is not discarded. +Heterogeneous subtype evidence uses an explicit Union and the same 8-type +inference limit. + +A plain concrete struct whose fields are supported can use ArrowTypes.jl's +default `StructKind` mapping. Without extension hooks, it reads back as +ordinary Struct storage rather than the original Julia type. To select a +different stable representation, map the struct to a supported storage type +with `ArrowTypes.ArrowType` and `ArrowTypes.toarrow`, or convert it to a +`NamedTuple` or separate columns. An `ArrowKind` override alone does not select +an arbitrary Arrow 3.0 physical layout. `Arrow.ToTimestamp` was removed; define +an ArrowTypes.jl lowering or convert zoned values before writing. + +## Schema retention + +When the source is an `Arrow.Table` or `Arrow.Stream`, the writer keeps +compatible details from the source schema. These details include temporal +units, byte and list widths, composite descriptors, nullability, field +metadata, schema metadata, and top-level dictionary index types and category +order. A column that was replaced with an incompatible Julia type is rejected +instead of being silently written under the old schema. + +A fresh Julia column with a heterogeneous declared `Union` element type is +synthesized as a canonical dense Arrow Union. Materialization of an existing +unregistered Arrow Union discards its child type IDs and offsets. Such a +retained Union still fails clearly when rewritten from an `Arrow.Table`; the +writer does not invent new routing under the retained schema. A registered +ArrowTypes.jl logical type whose storage is a Union retains enough writer-side +type evidence to reconstruct the original child routing, including dense or +sparse mode and type IDs. Sparse children use canonical hidden placeholder +values outside their active rows. When `JuliaType` returns an abstract target, +a concrete writer subtype may omit an extension identity or use the retained +parent's identity. A different explicit identity is rejected instead of being +silently relabeled. A retained nested Dictionary also fails because its pool is +lost. Top-level dictionaries of scalar or composite values are supported after +a full read. A scan result can lack the hidden source pool, so an ordered +dictionary from such a result also fails clearly. A nullable `Dictionary` +with an unknown extension fails +closed because materialization cannot retain valid-index versus null-index +provenance. Exact buffer sharing, overlapping ListView ranges, and Run-End +Encoding segmentation are not retained; the writer emits a canonical layout +with the same public-domain values and schema type. + +Registered public-domain values can rebuild compatible retained binary, list, +date-like, duration, wide-decimal, and interval descriptors. The retained +descriptor controls widths, sizes, units, and child fields. Values that do not +meet those exact constraints fail with `ArgumentError`. A retained Map with +`keysSorted=true` also rejects a replacement row whose keys are not sorted. + +The removed `maxdepth` keyword is not replaced by an unbounded writer. +Recursive ArrowTypes.jl storage schemas and recursive value containers throw +`ArgumentError`, as does custom mapping nesting beyond the fixed depth of 64. +To keep composite descriptors from allocating in proportion to untrusted +schema width, `JuliaType` receives exact `NTuple{N,T}` fixed-list storage when +`N` is at most 1024 and the compact `Tuple{Vararg{T}}` family above that limit. +Exact-arity registrations above the limit remain unregistered. Extension +Structs receive an exact `NamedTuple` storage signature through 1024 children +only when child names are unique, contain no embedded NUL, already exist as +Julia `Symbol`s, are at most 4096 UTF-8 bytes each, and use at most 64 KiB in +total. Otherwise the labelled Struct remains unknown and reads as ordered +`Pair` storage. To preserve ArrowTypes.jl Tuple storage, one bounded exception +may intern positional child names only when the complete sequence is exactly +`"1"`, `"2"`, …, `string(N)` for `N ≤ 1024`. Unknown extension labels return +before this check. Arbitrary or partly positional Struct names are not interned. +Any writer-side `ArrowType` result that is a concrete tuple above the same limit +is rejected before writer specialization. This includes ArrowTypes.jl's default +mapping for a tuple value. A custom hook itself remains trusted Julia code and +must return normally for Arrow to validate its result. + +Retained null-parent child slots now use direct length-based construction. +Null-only fixed-size lists and inactive sparse-Union children no longer expand +into one Julia placeholder per hidden element. + +## Names and imports + +Arrow 3.0 uses a small export surface. Qualify package functions: + +```julia +using Arrow + +table = Arrow.Table("data.arrow") +Arrow.write("copy.arrow", table) +``` + +`ArrowTypes` stays exported, so `using Arrow` keeps the bare `ArrowTypes` +binding working. Packages that define mappings should still depend on and +import ArrowTypes.jl directly rather than reach it through Arrow. + +Core schema names now stay as `String` values. `Arrow.Table` converts only +top-level Tables.jl column names to `Symbol`, after it preflights the complete +schema against a 4096-byte per-name limit, a 65,536-novel-name limit, and a +1-MiB novel-name byte budget. Unknown ArrowTypes extension labels are checked +without interning them. The built-in `JuliaLang.Symbol` extension rejects a +novel IPC payload with `ValidationError` instead of interning input-controlled +process-global state. This is an intentional behavior change: an input that a +prior Arrow.jl release read by interning its payload can now fail. + +Arrow 3.0 uses DataAPI.jl metadata methods. `Arrow.getmetadata(table)` +remains as a compatibility method over them, returning the table's key-value +metadata as a `Dict{String,String}` or `nothing`. The Arrow 2.x per-column +form `getmetadata(column)` is gone — columns are plain vectors — so use +`DataAPI.colmetadata` instead. Add DataAPI.jl as a direct dependency of code +that imports it: + +```julia +import Pkg +Pkg.add("DataAPI") +``` + +```julia +using Arrow, DataAPI + +collect(DataAPI.metadatakeys(table)) +DataAPI.metadata(table, "key") +DataAPI.colmetadata(table, :column, "key") +``` + +## New features + +Arrow 3.0 adds: + +- `Tables.Scan` pushdown for projection, filters, limits, and offsets. +- Recursive ArrowTypes.jl custom and extension-type mappings. Filters over a + field that contains a registered logical type at any depth evaluate over the + public materialized values because the interface does not require its + storage lowering to preserve Julia comparison semantics. +- Sparse byte-range reads through [`Arrow.AbstractArrowSource`](@ref). +- A CloudStore.jl extension for object storage. +- Arrow C data and C stream import and export. +- Arrow StringView and BinaryView support. +- Stronger validation and resource limits for untrusted IPC input. + +See the [User Manual](@ref) for examples and the [API Reference](@ref) for the +supported entry points. diff --git a/docs/src/reference.md b/docs/src/reference.md index 0409d76f..f70966ca 100644 --- a/docs/src/reference.md +++ b/docs/src/reference.md @@ -21,15 +21,69 @@ # API Reference -```@autodocs -Modules = [Arrow] -Order = [:type, :function] +## Read + +```@docs +Arrow.Table +Arrow.Stream +Arrow.Limits +Arrow.release!(::Arrow.Table) +Arrow.getmetadata +``` + +## Write + +```@docs +Arrow.write +Arrow.Writer +Arrow.write(::Arrow.Writer, ::Any) +Arrow.append +Arrow.tobuffer +Arrow.DictEncode +``` + +## Byte-range reads + +```@docs +Arrow.AbstractArrowSource +Arrow.sourcelength +Arrow.readrange +Arrow.concurrentreads +``` + +## The C data and C stream interfaces + +The low-level value and ABI types are public but not exported. + +```@docs +Arrow.Field +Arrow.Schema +Arrow.ArrayData +Arrow.RecordBatch +Arrow.fromjulia +Arrow.batch +Arrow.materialize +Arrow.CArrowSchema +Arrow.CArrowArray +Arrow.CArrowArrayStream ``` -## Internals: `Arrow.FlatBuffers` +```@docs +Arrow.to_c_data +Arrow.from_c_data +Arrow.export_stream! +Arrow.from_c_stream +Arrow.ForeignOwner +Arrow.ImportedStream +Arrow.nextbatch! +Arrow.release!(::Arrow.ForeignOwner) +Arrow.release!(::Arrow.ImportedStream) +Arrow.reap! +``` -The `FlatBuffers` module is not part of Arrow.jl's public API, and these functions may change without notice. +## Errors -```@autodocs -Modules = [Arrow.FlatBuffers] +```@docs +Arrow.AllocationLimitError +Arrow.ValidationError ``` diff --git a/ext/ArrowCloudStoreExt.jl b/ext/ArrowCloudStoreExt.jl new file mode 100644 index 00000000..5d1bfb1d --- /dev/null +++ b/ext/ArrowCloudStoreExt.jl @@ -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. + +# CloudStore.jl objects as Arrow byte-range sources: `Arrow.Table(obj; +# scan=…)` reads only the selected columns' bytes from S3, Azure Blob +# Storage through HTTP `Range` requests. +module ArrowCloudStoreExt + +using Arrow +using CloudStore: CloudStore, Object + +""" + CloudObjectSource(obj::CloudStore.Object) <: Arrow.AbstractArrowSource + +A `CloudStore.Object` as an [`Arrow.AbstractArrowSource`](@ref): the length +is the object's known size, one range is one HTTP `Range` GET pinned to the +object's ETag with `If-Match` (an overwritten key fails the read instead of +mixing versions across requests), and Arrow issues up to +`CONCURRENT_RANGE_READS` of a round's ranges at once. `Arrow.Table(obj; …)` +and `Arrow.Stream(obj; …)` construct one implicitly. +""" +struct CloudObjectSource{O<:Object} <: Arrow.AbstractArrowSource + obj::O +end + +# Independent HTTP range GETs are latency-bound, so a round's requests are +# worth overlapping; the bound keeps a fragmented object from becoming a +# request storm. +const CONCURRENT_RANGE_READS = 16 + +Arrow.sourcelength(s::CloudObjectSource) = Int64(s.obj.size) +Arrow.concurrentreads(::CloudObjectSource) = CONCURRENT_RANGE_READS + +# The object's ETag as an `If-Match` value (the header wants it quoted). +function _ifmatch(etag::AbstractString) + isempty(etag) && return nothing + return startswith(etag, '"') ? String(etag) : string('"', etag, '"') +end + +function Arrow.readrange(s::CloudObjectSource, off, len) + len == 0 && return UInt8[] + last = off + len - 1 + obj = s.obj + headers = ["Range" => "bytes=$(off)-$(last)"] + etag = _ifmatch(obj.eTag) + etag === nothing || push!(headers, "If-Match" => etag) + bytes = CloudStore.get( + obj.store, + obj.key; + credentials=obj.credentials, + headers=headers, + allowMultipart=false, + objectMaxSize=Int(len), + ) + return bytes isa Vector{UInt8} ? bytes : Vector{UInt8}(bytes) +end + +Arrow.Table(obj::Object; kw...) = Arrow.Table(CloudObjectSource(obj); kw...) +Arrow.Stream(obj::Object; kw...) = Arrow.Stream(CloudObjectSource(obj); kw...) + +end # module diff --git a/ext/ArrowTimeZonesExt.jl b/ext/ArrowTimeZonesExt.jl new file mode 100644 index 00000000..fbabd5dd --- /dev/null +++ b/ext/ArrowTimeZonesExt.jl @@ -0,0 +1,129 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +Loading TimeZones.jl restores Arrow 2.x reads for timezone-aware timestamps: +a second- or millisecond-unit timestamp column that declares a timezone +materializes as `TimeZones.ZonedDateTime` (the stored value is the UTC +instant; the declared zone is the display zone) instead of a naive UTC +`Dates.DateTime`. + +Micro- and nanosecond timestamps keep their raw `Int64` storage values with +or without this extension: neither `DateTime` nor `ZonedDateTime` can hold +them exactly, and Arrow 3.0 never truncates silently. A declared zone that +TimeZones.jl cannot parse falls back to the naive read with a one-time +warning. Scan filters that compare `ZonedDateTime` literals evaluate in the +public value domain (no storage-domain pushdown). +""" +module ArrowTimeZonesExt + +import Arrow +import Dates +import TimeZones + +const AC = Arrow.ArrowCore + +# Parse an Arrow timezone string: an IANA name ("America/Denver", legacy +# aliases included) or a fixed offset ("+07:00"). `nothing` means the zone is +# unusable and the caller keeps the naive read. +function _timezone(tz::AbstractString) + try + return TimeZones.TimeZone(tz, TimeZones.Class(:ALL)) + catch + end + try + return TimeZones.FixedTimeZone(tz) + catch + end + @warn "Arrow timestamp declares timezone $(repr(String(tz))) that " * + "TimeZones.jl cannot parse; reading the column as naive UTC values" _id = + Symbol(:arrow_bad_timezone_, tz) maxlog = 1 + return nothing +end + +# The two hooks `Arrow._zonedext` routes to. Both return `nothing` for an +# unusable zone so the eltype decision and the conversion agree. +zonedtype(tz::AbstractString) = + _timezone(tz) === nothing ? nothing : TimeZones.ZonedDateTime + +# Lower one ZonedDateTime to a timestamp's storage domain exactly: the UTC +# instant in the descriptor's unit, or `nothing` when the value is not a +# ZonedDateTime or would not round-trip (a second-unit column cannot hold a +# sub-second instant). The declared zone deliberately plays no role: storage +# is the UTC instant, so comparison semantics survive lowering. +function zonedstorage(unit::AC.TimeUnit, value) + value isa TimeZones.ZonedDateTime || return nothing + ms = Dates.value(Dates.DateTime(value, TimeZones.UTC)) - Dates.UNIXEPOCH + unit == AC.MILLISECOND && return ms + (unit == AC.SECOND && ms % 1000 == 0) && return ms ÷ 1000 + return nothing +end + +function zonedcolumn(t::AC.TimestampType, col, budget) + zone = _timezone(t.timezone) + zone === nothing && return nothing + scale = t.unit == AC.SECOND ? Int64(1000) : Int64(1) + return Arrow._mapcol( + TimeZones.ZonedDateTime, + x -> TimeZones.ZonedDateTime( + Dates.DateTime(Dates.UTM(Int64(x) * scale + Dates.UNIXEPOCH)), + zone; + from_utc=true, + ), + col, + budget, + ) +end + +# The fresh-write hooks `Arrow._zonedwritertype`/`Arrow._zonednativepart` +# route to. A fresh ZonedDateTime column writes as a timezone-declared +# millisecond timestamp, matching Arrow 2.x, instead of falling through to +# reflected-struct lowering of the zone's whole transition table. +iszonedtype(T::Type) = T === TimeZones.ZonedDateTime + +# One column carries one zone (the descriptor has one timezone slot), so +# mixed zones refuse with a normalization hint. An empty or all-missing +# column has no zone evidence and declares UTC. +function zonednativepart(name::String, v::AbstractVector, T::Type) + T === TimeZones.ZonedDateTime || return nothing + zone = nothing + for x in v + x === missing && continue + z = TimeZones.timezone(x) + if zone === nothing + zone = z + elseif z != zone + throw( + ArgumentError( + "column $name holds ZonedDateTime values in more than " * + "one timezone ($(TimeZones.name(zone)) and " * + "$(TimeZones.name(z))); convert them to one zone with " * + "astimezone before writing", + ), + ) + end + end + tzname = zone === nothing ? "UTC" : String(TimeZones.name(zone)) + return Arrow._constructtemporalpart( + name, + v, + AC.TimestampType(AC.MILLISECOND, tzname), + x -> Int64(Dates.value(Dates.DateTime(x, TimeZones.UTC)) - Dates.UNIXEPOCH), + ) +end + +end # module ArrowTimeZonesExt diff --git a/src/Arrow.jl b/src/Arrow.jl index 6f3ccdf8..5c4e05aa 100644 --- a/src/Arrow.jl +++ b/src/Arrow.jl @@ -15,130 +15,211 @@ # limitations under the License. """ - Arrow.jl + Arrow.jl — a pure Julia implementation of the Apache Arrow columnar format. + +Public surface: `Arrow.Table` reads the IPC stream and file formats (paths, +`IO`, byte vectors, or an `AbstractArrowSource` — a byte-range-addressable +object such as one in cloud storage) as Tables.jl tables, with `Tables.Scan` +pushdown; `Arrow.Stream` iterates a source one record batch at a time; +`Arrow.write` writes any Tables.jl source; `release!` releases mapped or +foreign storage deterministically. + +Layering ([docs/dev/core-README.md](https://github.com/apache/arrow-julia/blob/main/docs/dev/core-README.md) +documents each layer in depth): + +- `ArrowCore` (private): ownership regions, layout registry, `ArrayData`, + staged validation, accessors — the trim-friendly, dependency-free core. +- `Meta`: FlatBuffers metadata bindings AND the shape verifier, both + GENERATED from the vendored spec schemas (`src/metadata/fbs/`, generator + `tools/fbsgen.jl`) over the vendored `FlatBuffers` runtime. +- IPC adapters (`ipc_read.jl`, `ipc_write.jl`): stream and file formats, + framing, resource limits, compression, dictionary lifecycles. +- `cdata.jl`: the C data and C stream interfaces, import and export, with + lifecycle accounting. +- `scan.jl`: `Tables.Scan` pushdown over byte ranges plus footer-carried + statistics pruning. +- `table.jl`, `write.jl`: the read facade and write orchestration. +- `columnconstruction.jl`: column-wide construction policy, including + retained schemas, dictionaries, and recursive ArrowTypes.jl lowering. +""" +module Arrow -A pure Julia implementation of the [apache arrow](https://arrow.apache.org/) memory format specification. +using Tables +import Base64 +import DataAPI +import ArrowStrings +import ArrowTypes +import Dates +import Mmap +import CodecLz4 +import CodecZstd +import TranscodingStreams +using CodecLz4: LZ4FrameCompressor +using CodecZstd: ZstdCompressor + +const CLZ4 = CodecLz4 +const CZSTD = CodecZstd +const ZSTD = CZSTD.LibZstd +const TS = TranscodingStreams + +isdefined(Tables, :Scan) || error( + "Arrow 3.0's scan support needs Tables.jl's `Tables.Scan` " * + "interface; upgrade Tables.jl to 1.14 or later", +) + +include(joinpath("FlatBuffers", "FlatBuffers.jl")) +const FB = FlatBuffers + +# Generated metadata bindings + shape verifier (regenerate with +# `julia tools/fbsgen.jl src/metadata/fbs src/metadata`). +module Meta +using EnumX +using ..FlatBuffers +include(joinpath("metadata", "Schema.jl")) +include(joinpath("metadata", "File.jl")) +include(joinpath("metadata", "Message.jl")) +include(joinpath("metadata", "VerifierRuntime.jl")) +include(joinpath("metadata", "Verifier.jl")) +end -This implementation supports the 1.0 version of the specification, including support for: - * All primitive data types - * All nested data types - * Dictionary encodings, nested dictionary encodings, and messages - * Extension types - * Streaming, file, record batch, and replacement and isdelta dictionary messages - * Buffer compression/decompression via the standard LZ4 frame and Zstd formats +include("ArrowCore.jl") +using .ArrowCore +import .ArrowCore: release! +const AC = ArrowCore -It currently doesn't include support for: - * Tensors or sparse tensors - * Flight RPC - * C data interface +include("ipc_read.jl") +include("ipc_write.jl") +include("cdata.jl") +include("source.jl") +include("arrowtypes.jl") -Third-party data formats: - * csv and parquet support via the existing [CSV.jl](https://github.com/JuliaData/CSV.jl) and [Parquet.jl](https://github.com/JuliaIO/Parquet.jl) packages - * Other [Tables.jl](https://github.com/JuliaData/Tables.jl)-compatible packages automatically supported ([DataFrames.jl](https://github.com/JuliaData/DataFrames.jl), [JSONTables.jl](https://github.com/JuliaData/JSONTables.jl), [JuliaDB.jl](https://github.com/JuliaData/JuliaDB.jl), [SQLite.jl](https://github.com/JuliaDatabases/SQLite.jl), [MySQL.jl](https://github.com/JuliaDatabases/MySQL.jl), [JDBC.jl](https://github.com/JuliaDatabases/JDBC.jl), [ODBC.jl](https://github.com/JuliaDatabases/ODBC.jl), [XLSX.jl](https://github.com/felipenoris/XLSX.jl), etc.) - * No current Julia packages support ORC or Avro data formats +# The public facade. `table.jl` includes the private scan-plan module after +# its public/storage type seam is defined; `write.jl` includes the private +# column-construction module after the `DictEncode` marker is defined. +include("table.jl") +include("write.jl") -See docs for official Arrow.jl API with the [User Manual](@ref) and reference docs for [`Arrow.Table`](@ref), [`Arrow.write`](@ref), and [`Arrow.Stream`](@ref). -""" -module Arrow +@doc """ + Arrow.Field -using Base.Iterators -using Mmap -import Dates -using DataAPI, - Tables, - SentinelArrays, - PooledArrays, - CodecLz4, - CodecZstd, - TimeZones, - BitIntegers, - ConcurrentUtilities, - StringViews +A low-level Arrow column descriptor used by the C data interface. Obtain a +field with [`Arrow.fromjulia`](@ref) or [`Arrow.from_c_data`](@ref). +""" Field -export ArrowTypes +@doc """ + Arrow.Schema -using Base: @propagate_inbounds -import Base: == +A low-level ordered collection of [`Arrow.Field`](@ref) values plus optional +schema metadata. The C stream interface uses it to describe each batch. +""" Schema -const FILE_FORMAT_MAGIC_BYTES = b"ARROW1" -const CONTINUATION_INDICATOR_BYTES = 0xffffffff +@doc """ + Arrow.ArrayData -# vendored flatbuffers code for now -include("FlatBuffers/FlatBuffers.jl") -using .FlatBuffers +A low-level Arrow array: buffers, children, an optional dictionary, and a +logical length. Obtain it with [`Arrow.fromjulia`](@ref) or +[`Arrow.from_c_data`](@ref), and convert it with [`Arrow.materialize`](@ref). +""" ArrayData -include("metadata/Flatbuf.jl") -using .Flatbuf -const Meta = Flatbuf +@doc """ + Arrow.RecordBatch -using ArrowTypes -include("utils.jl") -include("arraytypes/arraytypes.jl") -include("eltypes.jl") -include("table.jl") -include("write.jl") -include("append.jl") -include("show.jl") - -const ZSTD_COMPRESSOR = Lockable{ZstdCompressor}[] -const ZSTD_DECOMPRESSOR = Lockable{ZstdDecompressor}[] -const LZ4_FRAME_COMPRESSOR = Lockable{LZ4FrameCompressor}[] -const LZ4_FRAME_DECOMPRESSOR = Lockable{LZ4FrameDecompressor}[] - -function init_zstd_compressor() - zstd = ZstdCompressor(; level=3) - CodecZstd.TranscodingStreams.initialize(zstd) - return Lockable(zstd) -end +A low-level [`Arrow.Schema`](@ref) and an equal-length `Arrow.ArrayData` +column for each field. [`Arrow.batch`](@ref) builds one from Julia vectors. +""" RecordBatch -function init_zstd_decompressor() - zstd = ZstdDecompressor() - CodecZstd.TranscodingStreams.initialize(zstd) - return Lockable(zstd) -end +@doc """ + Arrow.fromjulia(name, values) -> (Arrow.Field, Arrow.ArrayData) -function init_lz4_frame_compressor() - lz4 = LZ4FrameCompressor(; compressionlevel=4) - CodecLz4.TranscodingStreams.initialize(lz4) - return Lockable(lz4) -end +Build the low-level C-interchange representation of one supported Julia +vector. Do not resize or mutate zero-copy input buffers while the result is in +use. +""" fromjulia -function init_lz4_frame_decompressor() - lz4 = LZ4FrameDecompressor() - CodecLz4.TranscodingStreams.initialize(lz4) - return Lockable(lz4) -end +@doc """ + Arrow.batch(columns::NamedTuple) -> Arrow.RecordBatch -function access_threaded(f, v::Vector) - tid = Threads.threadid() - 0 < tid <= length(v) || _length_assert() - if @inbounds isassigned(v, tid) - @inbounds x = v[tid] - else - x = f() - @inbounds v[tid] = x - end - return x -end -@noinline _length_assert() = @assert false "0 < tid <= v" - -zstd_compressor() = access_threaded(init_zstd_compressor, ZSTD_COMPRESSOR) -zstd_decompressor() = access_threaded(init_zstd_decompressor, ZSTD_DECOMPRESSOR) -lz4_frame_compressor() = access_threaded(init_lz4_frame_compressor, LZ4_FRAME_COMPRESSOR) -lz4_frame_decompressor() = - access_threaded(init_lz4_frame_decompressor, LZ4_FRAME_DECOMPRESSOR) - -function __init__() - nt = @static if isdefined(Base.Threads, :maxthreadid) - Threads.maxthreadid() - else - Threads.nthreads() - end - resize!(empty!(LZ4_FRAME_COMPRESSOR), nt) - resize!(empty!(ZSTD_COMPRESSOR), nt) - resize!(empty!(LZ4_FRAME_DECOMPRESSOR), nt) - resize!(empty!(ZSTD_DECOMPRESSOR), nt) - return +Build a low-level record batch from a named tuple of supported Julia vectors. +""" batch + +@doc """ + Arrow.materialize(field, data) -> Vector + Arrow.materialize(T, field, data) -> Vector{T} + +Convert low-level `Arrow.ArrayData` to native Julia values. The typed form +checks that `T` agrees with the Arrow descriptor before conversion. +""" materialize + +""" + Arrow.nextbatch!(source) -> Union{Nothing, RecordBatch} + +Pull the next record batch from a record-batch source — an IPC stream +(`readstream`), an imported C stream (`from_c_stream`) — or `nothing` at end +of stream. Sources are single-owner cursors: overlapping calls on one source +are an error. +""" +AC.nextbatch! + +""" + Arrow.ValidationError + +Thrown by every validation tier — structural, semantic, and the opt-in +`validate_full` — when a descriptor or array violates the Arrow format. +""" +AC.ValidationError + +# Julia 1.11 added `public`. Keep the source parseable on the supported +# Julia 1.10 floor while giving tooling an exact, non-exporting API boundary +# on newer Julia versions. +@static if VERSION >= v"1.11" + Core.eval( + @__MODULE__, + Expr( + :public, + :Table, + :Stream, + :write, + :Writer, + :append, + :tobuffer, + :getmetadata, + :DictEncode, + :AbstractArrowSource, + :sourcelength, + :readrange, + :concurrentreads, + :Field, + :Schema, + :ArrayData, + :RecordBatch, + :fromjulia, + :batch, + :materialize, + :CArrowSchema, + :CArrowArray, + :CArrowArrayStream, + :to_c_data, + :from_c_data, + :export_stream!, + :from_c_stream, + :ForeignOwner, + :ImportedStream, + :nextbatch!, + :reap!, + :Limits, + :AllocationLimitError, + :ValidationError, + ), + ) end -end # module Arrow +export release! +# The one name Arrow 2.x exported. Packages that define custom-type mappings +# reach the interface as `Arrow.ArrowTypes` or through `using Arrow`. +export ArrowTypes + +# Precompile the common write/read/scan shapes last, once every layer above +# is defined. +include("precompile.jl") + +end # module Arrow diff --git a/src/ArrowCore.jl b/src/ArrowCore.jl new file mode 100644 index 00000000..56baa92d --- /dev/null +++ b/src/ArrowCore.jl @@ -0,0 +1,3835 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" + ArrowCore + +The runtime-tagged, C-data-shaped core of Arrow.jl. It depends only on Base +and the Mmap standard library; the IPC, C-data, scan, and facade layers in +`src/` sit on top of it. + +Design rules: + +1. One physical data model. `ArrayData` = layout + buffers + children + + dictionary, mirroring the Arrow C data interface's `ArrowArray`. Logical + type parameters such as timezone and precision/scale are fields on + `ArrowType` descriptors. Names and nullability are fields on `Field`. + None parameterize the Core storage types. Struct materialization returns + `Vector{Pair{String,Any}}` on the dynamic path; static element types are + the caller's claim through `getvalue(::Type{T}, ...)`/`materialize(::Type{T}, ...)`. + +2. Memory validity is GC reachability, plus one revocation bit. Every + buffer is a `BufferSlice` into an `OwnerRegion` — a (pointer, length, + alignment, root, cell) record whose `root` anchors the backing storage + and whose `ReleaseCell` supports `release!`: regions sharing one + underlying lifetime share one cell, so a release revokes every sibling and + runs the release action (mmap unmap, foreign release callback) exactly + once, and later access is a clean error. Slices are bounds-checked + against the region at construction; loads are a final bounds check, one + monotonic closed-flag load, and the raw read. Foreign extents are + trusted declarations, and mapped files are exposed to external changes. + +3. One structural layout registry. `layoutspec(type)` returns the buffer + roles / child arity / offset width for each of the format-1.5 layouts. + Generic code (buffer walking, structural validation, the IPC adapter's + node/buffer accounting in src/ipc_read.jl) is driven by the registry; + per-layout SEMANTICS (element access, semantic validation) are ordinary + methods grouped per layout below. Adding a layout means one registry + entry plus bounded method groups in the layers that support it. + +4. Validation is staged: structural checks are O(buffers) and run at + construction/adaptation time. Data-intrinsic semantic checks are O(n) + when an adapter or caller requests them; a successful result is cached. + Benign concurrent callers may repeat the same scan. Field-dependent + dictionary contracts run on every validation call. Advisory contracts — + Field.nullable enforcement, Date64 day divisibility, time-of-day range, + decimal precision, and body UTF-8 — are opt-in via `validate_full`: the + ecosystem's gold files violate them and the reference implementation + reads those files. Framing-stage checks (checked spans, metadata + verification, resource limits before metadata-directed allocation) + belong to the adapters. + +The registry, staged validation, element access, and materialization cover +every format-1.5 layout, including binary views, list views, and run-end +encoding. `validate_full` additionally enforces canonical bit-packed form +(zeroed trailing bits and padding); on-wire buffer padding is a writer +guarantee, not a reader requirement — the spec permits unpadded buffers and +this reader accepts them. Core has no codec dependency; the IPC adapter +implements compression. Bulk access is `materialize` (dynamic, one function +barrier per layout) and `materialize(::Type{T}, …)` (a static element claim, +resolved without boxing except at the dictionary/run-end wrapper edge). +""" +module ArrowCore + +using Base: Checked +import Mmap +const checked_add = Checked.checked_add +const checked_sub = Checked.checked_sub +const checked_mul = Checked.checked_mul + +export OwnerRegion, + BufferSlice, + heapregion, + mmapregion, + release!, + ReleaseCell, + ReleaseCounter, + increment!, + ArrowType, + NullType, + BoolType, + IntType, + FloatType, + DecimalType, + FixedSizeBinaryType, + BinaryType, + Utf8Type, + DateType, + TimeType, + TimestampType, + DurationType, + IntervalType, + ListType, + FixedSizeListType, + StructType, + MapType, + UnionType, + DictionaryType, + ViewType, + ListViewType, + RunEndEncodedType, + Field, + Schema, + ArrayData, + RecordBatch, + RecordBatchSource, + nextbatch!, + LayoutSpec, + layoutspec, + BufferRole, + validate_structural, + validate_semantic, + validate_full, + ValidationError, + nullcount, + getvalue, + materialize, + fromjulia, + fromviewentries, + batch + +# --------------------------------------------------------------------------- +# §1 Memory: regions as GC anchors (constrained model) +# --------------------------------------------------------------------------- +# +# Buffer validity is GC REACHABILITY — Julia's native memory-safety +# contract — plus one explicit revocation layer. A region's `root` is +# whatever keeps the memory alive (the wrapped Julia array, the Mmap-stdlib +# array, a C-data adapter's owner object); views hold their region, the +# region holds its root, so memory a view can reach is memory that is +# valid. `release!` is the deterministic release path on top: one +# `ReleaseCell` per underlying lifetime revokes every region over it and +# runs the eager release action exactly once (unmap now; run the foreign +# release now), turning use-after-release into `InvalidStateException` +# instead of undefined behavior. What stays out of scope, so the contract +# is informed: +# * Data-race shielding for loads concurrent WITH release!: quiescing +# readers first is the caller's contract, as with `Base.close` on a +# shared IO. Loads take no locks. +# * External-truncation protection: a shared mapping's pages can vanish +# under any implementation. Same exposure as every mmap-based reader. + +"An atomic counter for the test suite's exactly-once release bookkeeping." +mutable struct ReleaseCounter + @atomic n::Int +end +ReleaseCounter() = ReleaseCounter(0) +Base.getindex(c::ReleaseCounter) = @atomic c.n +# CAS loop rather than `@atomic c.n += 1`: the atomic read-modify-write +# builtin is not implemented in JuliaC's trim verifier, while +# compare-and-swap is; contention here is negligible. +function increment!(c::ReleaseCounter) + while true + old = @atomic c.n + _, ok = @atomicreplace c.n old => old + 1 + ok && return old + 1 + end +end + +""" + ReleaseCell(action::Ptr{Cvoid}, arg) + ReleaseCell() + +The revocation state one release action guards. Every `OwnerRegion` carries +a cell; regions that share one underlying lifetime (all buffers imported +from one C-data tree) share ONE cell, so closing any of them revokes every +sibling before the single release action runs. The action is a +`@cfunction(f, Cvoid, (Ptr{Cvoid},))` trampoline receiving +`pointer_from_objref(arg)` — pure C ABI, the same idiom the C-data +adapter's release callbacks use and the form the trim verifier accepts +(an `Any`-argument cfunction is not) — or `C_NULL` when the backing +storage is a borrow with no eager action (a heap vector — running a +borrowed object's finalizers is not ours to do). `arg` must be a mutable +heap object; the cell's reference keeps it alive across the call. Build +trampolines at runtime, never in a module-level const (a serialized +cfunction pointer is garbage after precompile reload). +""" +mutable struct ReleaseCell + @atomic closed::Bool + const action::Ptr{Cvoid} + const arg::Any +end +ReleaseCell(action::Ptr{Cvoid}, arg) = ReleaseCell(false, action, arg) +ReleaseCell() = ReleaseCell(false, Ptr{Cvoid}(C_NULL), nothing) + +""" + release!(cell::ReleaseCell) + +Revoke every region sharing the cell — later raw access throws +`InvalidStateException` — and run the cell's release action exactly once. +Idempotent. Not a data-race shield for accesses concurrent WITH the release; +quiescing readers first is the caller's contract, as with `Base.close` on +a shared IO. +""" +function release!(cell::ReleaseCell) + (@atomicswap :acquire_release cell.closed = true) && return nothing + if cell.action != C_NULL + arg = cell.arg + GC.@preserve arg ccall(cell.action, Cvoid, (Ptr{Cvoid},), pointer_from_objref(arg)) + end + return nothing +end + +""" + OwnerRegion + +One contiguous memory region, the object that keeps it alive, and the +[`ReleaseCell`](@ref) that can revoke it. The region is valid while it is +reachable — `root` anchors the backing storage (a borrowed Julia array, the +Mmap-stdlib array, or an adapter's owner object) — or until `release!` runs +its cell's release action, after which every raw access through `sliceptr` +throws. Slices reject geometry outside the declared `len` at construction; +loads retain a final bounds check before the raw read. + +The scoped-borrow contract for wrapped Julia arrays: the caller must not +mutate or resize the array while the region or any cached validation result +remains in use. Mutation can invalidate a semantic certificate; resizing can +reallocate the storage and invalidate its pointer. +""" +struct OwnerRegion + ptr::Ptr{UInt8} + len::Int64 + alignment::Base.Int # guaranteed ptr alignment, capped at 64; loads consult it + root::Any # GC anchor; never dispatched on, only stored + cell::ReleaseCell + + function OwnerRegion( + ptr::Ptr{UInt8}, + len::Integer; + root=nothing, + cell::ReleaseCell=ReleaseCell(), + ) + len >= 0 || throw(ArgumentError("region length must be non-negative")) + n = Int64(len) + (ptr != C_NULL || n == 0) || + throw(ArgumentError("a non-empty region requires a non-NULL pointer")) + (root !== nothing || n == 0) || + throw(ArgumentError("a non-empty region requires a GC root")) + # BufferSlice bounds are only meaningful if every declared byte also + # has a representable pointer address. Reject a foreign extent whose + # final byte would wrap native pointer arithmetic. + if n > 0 + lastaddr = UInt128(UInt(ptr)) + UInt128(n - 1) + lastaddr <= UInt128(typemax(UInt)) || + throw(ArgumentError("region extent wraps the native address space")) + end + # OR-ing in 64 caps the detected alignment at 64 bytes without a + # branch; a NULL (necessarily empty) region reports the cap, since + # nothing loads from it. + align = ptr == C_NULL ? 64 : (1 << trailing_zeros(UInt(ptr) | UInt(64))) + return new(ptr, n, align, root, cell) + end +end + +""" + release!(r::OwnerRegion) + +Deterministically release the region's backing storage through its +`ReleaseCell`: every region sharing the cell is revoked (later raw +access throws `InvalidStateException`) and the cell's release action runs +exactly once — an mmap region unmaps NOW (the eager path exists for hosts +where a GC-timed unmap is not enough, deleting a still-mapped file on +Windows being the canonical case); an imported C-data tree runs the +producer's release callback; a borrowed heap region is revoked with no +eager action. Idempotent. +""" +function release!(r::OwnerRegion) + return release!(r.cell) +end + +""" + heapregion(v::Vector{T}) -> OwnerRegion + +Borrow a Julia array as a region (zero-copy). The array is the `root`. +""" +function heapregion(v::Vector{T}) where {T} + isbitstype(T) || throw(ArgumentError("heapregion requires an isbits element type")) + return OwnerRegion(Ptr{UInt8}(pointer(v)), sizeof(v); root=v) +end + +# `_mmaproot(arr)`: the object Mmap registered its unmap finalizer on — the +# array's backing `Memory` from Julia 1.11 (`finalize(arr)` is then a +# no-op), the array itself before. The release cell targets that object so +# `release!` truly unmaps now. +@static if VERSION >= v"1.11" + _mmaproot(arr::Vector{UInt8}) = arr.ref.mem + function _release_mmap(p::Ptr{Cvoid})::Cvoid + finalize(unsafe_pointer_to_objref(p)::Memory{UInt8}) + return nothing + end +else + _mmaproot(arr::Vector{UInt8}) = arr + function _release_mmap(p::Ptr{Cvoid})::Cvoid + finalize(unsafe_pointer_to_objref(p)::Vector{UInt8}) + return nothing + end +end + +# Map from a caller-owned open `io` (the caller closes it; the mapping +# outlives the descriptor); `label` names the mapping in errors. +function _mmapregion(io::IO, label) + arr = Mmap.mmap(io, Vector{UInt8}) + isempty(arr) && throw(ArgumentError("cannot map empty file: $label")) + cell = ReleaseCell(@cfunction(_release_mmap, Cvoid, (Ptr{Cvoid},)), _mmaproot(arr)) + return OwnerRegion(Ptr{UInt8}(pointer(arr)), length(arr); root=arr, cell=cell) +end + +""" + mmapregion(path) -> OwnerRegion + +Map a file read-only via the Mmap STDLIB (cross-platform) and wrap the +mapped array as the region's `root`. The stdlib's own machinery unmaps when +the array is collected — validity is reachability, like every other region. + +The caller must prevent external writes or truncation of the mapped file +while the region or any cached validation result remains in use: a shared +mapping cannot keep a semantic certificate valid when another process +changes its bytes, and truncation can make an in-range load fault. +""" +function mmapregion(path::AbstractString) + io = open(path, "r") + try + # The mapping outlives the descriptor. + return _mmapregion(io, path) + finally + close(io) + end +end + +# --- BufferSlice ------------------------------------------------------------ + +""" + BufferSlice + +A bounds-checked window into an `OwnerRegion` — the only currency for Arrow +buffer data in this module. Constructing a slice validates +`offset + len <= region.len` with checked arithmetic, so downstream code can +assume every slice is in-bounds and concentrate on layout logic. + +An all-default `BufferSlice()` is the canonical empty buffer (used for +absent validity bitmaps, empty data buffers). +""" +struct BufferSlice + region::Union{Nothing,OwnerRegion} + offset::Int64 + len::Int64 + function BufferSlice(region::OwnerRegion, offset::Integer, len::Integer) + offset >= 0 || throw(ArgumentError("negative buffer offset")) + len >= 0 || throw(ArgumentError("negative buffer length")) + checked_add(Int64(offset), Int64(len)) <= region.len || throw( + ArgumentError( + "buffer [offset=$offset len=$len] exceeds region of $(region.len) bytes", + ), + ) + return new(region, Int64(offset), Int64(len)) + end + BufferSlice() = new(nothing, 0, 0) +end + +Base.length(b::BufferSlice) = b.len +isempty_buffer(b::BufferSlice) = b.len == 0 +@inline function sliceptr(b::BufferSlice) + b.region === nothing && return Ptr{UInt8}(0) + r = b.region::OwnerRegion + (@atomic :monotonic r.cell.closed) && + throw(InvalidStateException("the backing region was released", :closed)) + return r.ptr + b.offset +end + +"Sub-slice with checked arithmetic (relative bounds against the parent slice)." +function subslice(b::BufferSlice, offset::Integer, len::Integer) + offset >= 0 || throw(ArgumentError("negative subslice offset")) + len >= 0 || throw(ArgumentError("negative subslice length")) + b.region === nothing && (len == 0 && offset == 0) && return b + b.region === nothing && throw(ArgumentError("cannot subslice the empty buffer")) + checked_add(Int64(offset), Int64(len)) <= b.len || throw( + ArgumentError("subslice [offset=$offset len=$len] exceeds slice of $(b.len) bytes"), + ) + return BufferSlice(b.region, checked_add(b.offset, Int64(offset)), Int64(len)) +end + +""" +Load a `T` at byte offset `byteoff` (0-based) within the slice. Handles the +misaligned case with a byte-wise load: alignment is a property of the region +(Arrow controls only its own allocations; mmap and foreign pointers can be +anything), so the branch lives here, in one place, instead of as a copy +workaround scattered through per-type code. +""" +@inline function loadat(b::BufferSlice, ::Type{T}, byteoff::Int64) where {T} + # Raw Arrow bytes may only materialize pointer-free values. Loading a + # struct with managed references would treat attacker-controlled bytes as + # GC pointers and can crash Julia before it can report an ordinary error. + isbitstype(T) || throw(ArgumentError("loadat requires an isbits type, got $T")) + # Bounds: byteoff + sizeof(T) <= len. byteoff is computed by callers from + # validated element indices, but re-check cheaply: this is the last line + # of defense before a raw pointer dereference. + width = Int64(sizeof(T)) + # Express this as subtraction, not `byteoff + width <= len`: a hostile + # byte offset near typemax(Int64) must not wrap through the last bounds + # check and reach pointer arithmetic. + (byteoff >= 0 && width <= b.len && byteoff <= b.len - width) || + throw(BoundsError(b, byteoff)) + b.region === nothing && throw(ArgumentError("empty buffer has no data")) + # Raw pointers do not keep Julia owners alive. Preserve the slice through + # the full dereference so its region and opaque root remain reachable. + GC.@preserve b begin + p = sliceptr(b) + byteoff + required = datatype_alignment(T) + relative = checked_add(b.offset, byteoff) + region = b.region::OwnerRegion + if region.alignment >= required && relative % required == 0 + return unsafe_load(Ptr{T}(p)) + else + return _load_unaligned(T, p) + end + end +end + +datatype_alignment(::Type{T}) where {T} = Base.datatype_alignment(T) + +@inline function _load_unaligned(::Type{T}, p::Ptr{UInt8}) where {T} + bytes = ntuple(i -> unsafe_load(p + (i - 1)), Val(sizeof(T))) + return reinterpret_bytes(T, bytes) +end +@inline reinterpret_bytes(::Type{T}, bytes::NTuple{N,UInt8}) where {T,N} = + reinterpret(T, bytes) + +"Copy the slice into a fresh `Vector{UInt8}`." +function slicebytes(b::BufferSlice) + b.len == 0 && return UInt8[] + b.region === nothing && throw(ArgumentError("empty buffer has no data")) + out = Vector{UInt8}(undef, b.len) + # Preserve both owners across the raw copy. Neither pointer roots its + # source or destination allocation. + GC.@preserve b out begin + unsafe_copyto!(pointer(out), sliceptr(b), b.len) + end + return out +end + +"Read one bit (0-based bit index) from a validity/values bitmap slice." +@inline function getbit(b::BufferSlice, i::Int64) + byte = loadat(b, UInt8, i >> 3) + return (byte >> (i & 7)) & 0x01 == 0x01 +end + +# --------------------------------------------------------------------------- +# §2 Type system: runtime descriptors, Field, Schema +# --------------------------------------------------------------------------- + +""" + ArrowType + +Abstract supertype of the runtime logical-type descriptors. These are small +immutable structs whose *fields* carry the logical type parameters +(timestamp unit and timezone, decimal precision/scale/width, ...). Two +timestamp columns with different timezones have the SAME Julia type — +schema diversity costs data, not method instances. +""" +abstract type ArrowType end + +""" +Read-only, defensively-copied vector storage for the frozen data model. +Its type does not encode the length, so schema width and nesting depth do +not create a new family of container types. The backing field is internal; +normal mutation APIs such as `setindex!` and `push!` are unavailable. +""" +struct FrozenVector{T} <: AbstractVector{T} + _data::Vector{T} + FrozenVector{T}(data::Vector{T}, ::Nothing) where {T} = new{T}(data) +end +FrozenVector{T}(xs::FrozenVector{T}) where {T} = xs +FrozenVector{T}(xs) where {T} = FrozenVector{T}(collect(T, xs), nothing) +Base.size(v::FrozenVector) = size(getfield(v, :_data)) +Base.length(v::FrozenVector) = length(getfield(v, :_data)) +Base.getindex(v::FrozenVector, i::Int) = getfield(v, :_data)[i] +Base.IndexStyle(::Type{<:FrozenVector}) = IndexLinear() + +@enum TimeUnit::UInt8 SECOND MILLISECOND MICROSECOND NANOSECOND +@enum DateUnit::UInt8 DAY MILLISECOND_DATE +@enum IntervalUnit::UInt8 YEAR_MONTH DAY_TIME MONTH_DAY_NANO +@enum UnionMode::UInt8 SparseMode DenseMode +@enum Endianness::UInt8 LittleEndian BigEndian +_native_endianness() = Base.ENDIAN_BOM == 0x04030201 ? LittleEndian : BigEndian + +"Null: every slot null; no buffers." +struct NullType <: ArrowType end +"Boolean: bit-packed values." +struct BoolType <: ArrowType end +"Integer of `bits` width, `signed` or not." +struct IntType <: ArrowType + bits::Int # 8/16/32/64 — the spec's Int; wider is NOT valid + signed::Bool +end +"IEEE floating point of `bits` width (16/32/64)." +struct FloatType <: ArrowType + bits::Int # 16/32/64 +end +"Decimal of `precision` digits and `scale`, stored in `bits` (32/64/128/256)." +struct DecimalType <: ArrowType + precision::Int + scale::Int + bits::Int # 32/64/128/256 (format 1.5) +end +"Fixed-width binary of `nbytes` per slot." +struct FixedSizeBinaryType <: ArrowType + nbytes::Int +end +"Variable-length binary; `large` selects Int64 offsets." +struct BinaryType <: ArrowType + large::Bool # Int64 offsets when true +end +"UTF-8 string; `large` selects Int64 offsets." +struct Utf8Type <: ArrowType + large::Bool +end +"Date in days (Int32) or milliseconds (Int64) since the epoch." +struct DateType <: ArrowType + unit::DateUnit # DAY => Int32 storage, MILLISECOND => Int64 +end +"Time of day in `unit`, stored in `bits` (32 for s/ms, 64 for us/ns)." +struct TimeType <: ArrowType + unit::TimeUnit + bits::Int # 32 (s/ms) or 64 (us/ns) +end +"Timestamp in `unit` since the epoch, with an optional `timezone`." +struct TimestampType <: ArrowType + unit::TimeUnit + timezone::Union{Nothing,String} # a VALUE — one method instance total +end +"Elapsed time in `unit`, Int64 storage." +struct DurationType <: ArrowType + unit::TimeUnit +end +"Calendar interval in `unit` (year-month, day-time, or month-day-nano)." +struct IntervalType <: ArrowType + unit::IntervalUnit # includes MONTH_DAY_NANO (format 1.2) +end +"Variable-length list; `large` selects Int64 offsets." +struct ListType <: ArrowType + large::Bool +end +"List of exactly `listsize` child slots per parent slot." +struct FixedSizeListType <: ArrowType + listsize::Int +end +"Struct: named children declared by `Field.children`." +struct StructType <: ArrowType end +"Map: a list of key/value struct entries; `keyssorted` per entry." +struct MapType <: ArrowType + keyssorted::Bool +end +"Union in sparse or dense `mode` over the children's declared `typeids`." +struct UnionType <: ArrowType + mode::UnionMode + typeids::FrozenVector{Int8} # declared type-id domain, child order +end +UnionType(mode::UnionMode, typeids) = UnionType(mode, FrozenVector{Int8}(typeids)) +"Dictionary-encoded: `indextype` is the physical index; values live in `ArrayData.dictionary`." +struct DictionaryType <: ArrowType + indextype::IntType + valuetype::ArrowType + ordered::Bool +end +"Utf8View / BinaryView (format 1.4): 16-byte entries plus variadic data buffers." +struct ViewType <: ArrowType + utf8::Bool +end +"ListView / LargeListView (format 1.4): per-slot child offsets and sizes." +struct ListViewType <: ArrowType + large::Bool +end +"Run-end encoded (format 1.3): signed run ends and values of any Arrow type." +struct RunEndEncodedType <: ArrowType end + +""" + Field + +One column/child descriptor: name, logical type, nullability, metadata, and +child fields. Dictionary columns are `DictionaryType` here; the IPC-level +dictionary *id* is NOT a Field concern — it is IPC bookkeeping and lives in +the adapter (Core dictionaries are object references; the id↔dictionary +table is the adapter's). +""" +struct Field + name::String + type::ArrowType + nullable::Bool + metadata::Union{Nothing,FrozenVector{Pair{String,String}}} + children::FrozenVector{Field} +end +_freezemetadata(::Nothing) = nothing +_freezemetadata(metadata::FrozenVector{Pair{String,String}}) = metadata +function _freezemetadata(metadata::Union{AbstractVector,Tuple}) + all(kv -> kv isa Pair, metadata) || + throw(ArgumentError("metadata sequences must contain Pair values")) + return FrozenVector{Pair{String,String}}( + String(first(kv)) => String(last(kv)) for kv in metadata + ) +end +_freezemetadata(metadata) = + FrozenVector{Pair{String,String}}(String(k) => String(v) for (k, v) in pairs(metadata)) +Field(name, type; nullable=true, metadata=nothing, children=()) = Field( + String(name), + type, + Bool(nullable), + _freezemetadata(metadata), + FrozenVector{Field}(children), +) +# Narrower than the struct's implicit (Any...) convert constructor so this +# ADDS a positional-with-conversion method instead of overwriting it (which +# precompilation forbids); exact-typed calls still take the implicit one. +Field(name::AbstractString, type::ArrowType, nullable, metadata, children) = + Field(name, type; nullable=nullable, metadata=metadata, children=children) + +""" + Schema + +An ordered set of top-level `Field`s plus optional schema-level metadata: +the shape of every record batch that carries it. +""" +struct Schema + fields::FrozenVector{Field} + metadata::Union{Nothing,FrozenVector{Pair{String,String}}} + endianness::Endianness +end +Schema(fields; metadata=nothing, endianness=_native_endianness()) = + Schema(FrozenVector{Field}(fields), _freezemetadata(metadata), endianness) + +# --------------------------------------------------------------------------- +# §3 Layout registry (structural facts only) +# --------------------------------------------------------------------------- + +# OFFSETS are RANGE offsets (len+1 entries bounding variable-size slots); +# ELEMENT_OFFSETS are per-element child positions (len entries — dense union). +# The distinction is structural, so it lives in the registry, not in +# per-layout special cases inside the validator. +"The role of one buffer in a layout's buffer sequence (see `LayoutSpec`)." +@enum BufferRole::UInt8 VALIDITY DATA OFFSETS ELEMENT_OFFSETS SIZES VIEWS TYPE_IDS + +""" + LayoutSpec + +The STRUCTURAL facts for one physical layout: which buffers it has (in +order), how many children, its offset width, whether the trailing data +buffers are variadic (view layouts). This is everything generic code needs +to walk a layout — and nothing more. Semantics (what the bytes mean, how to +access element `i`) are per-layout methods, not registry rows. + +`childcount == -1` means "declared by Field.children" (struct/union); +`fixedwidth` is bytes-per-element for fixed-stride DATA buffers, 0 when the +data buffer is byte-addressed (varbinary) or absent, and -1 for bit-packed. +""" +struct LayoutSpec + buffers::FrozenVector{BufferRole} + childcount::Int + offsetwidth::Int # 0, 4, or 8 — width of the OFFSETS buffer entries + fixedwidth::Int + variadic::Bool +end +# Narrower than the implicit convert constructor (see Field above). +LayoutSpec(buffers::AbstractVector, childcount, offsetwidth, fixedwidth, variadic) = + LayoutSpec( + FrozenVector{BufferRole}(buffers), + childcount, + offsetwidth, + fixedwidth, + variadic, + ) + +primwidth(t::IntType) = t.bits ÷ 8 +primwidth(t::FloatType) = t.bits ÷ 8 +primwidth(t::DecimalType) = t.bits ÷ 8 +primwidth(t::DateType) = t.unit == DAY ? 4 : 8 +primwidth(t::TimeType) = t.bits ÷ 8 +primwidth(::TimestampType) = 8 +primwidth(::DurationType) = 8 +primwidth(t::IntervalType) = t.unit == YEAR_MONTH ? 4 : t.unit == DAY_TIME ? 8 : 16 +primwidth(t::FixedSizeBinaryType) = t.nbytes + +# The buffer-role sequences are shared constants: `layoutspec` runs on +# every buffer-by-role lookup, so a per-call vector would allocate in the +# element accessors' inner loops. +const NO_BUFFERS = FrozenVector{BufferRole}(()) +const VALIDITY_ONLY = FrozenVector{BufferRole}((VALIDITY,)) +const VALIDITY_DATA = FrozenVector{BufferRole}((VALIDITY, DATA)) +const VALIDITY_OFFSETS = FrozenVector{BufferRole}((VALIDITY, OFFSETS)) +const VALIDITY_OFFSETS_DATA = FrozenVector{BufferRole}((VALIDITY, OFFSETS, DATA)) +const VALIDITY_VIEWS = FrozenVector{BufferRole}((VALIDITY, VIEWS)) +const VALIDITY_ELEMENT_OFFSETS_SIZES = + FrozenVector{BufferRole}((VALIDITY, ELEMENT_OFFSETS, SIZES)) +const TYPE_IDS_ONLY = FrozenVector{BufferRole}((TYPE_IDS,)) +const TYPE_IDS_ELEMENT_OFFSETS = FrozenVector{BufferRole}((TYPE_IDS, ELEMENT_OFFSETS)) + +""" + layoutspec(t::ArrowType) -> LayoutSpec + +The structural facts of `t`'s physical layout: one method per descriptor +type (the registry's extension point); generic code reaches it through the +closed-set ladder `layoutspec_of`. +""" +layoutspec(::NullType) = LayoutSpec(NO_BUFFERS, 0, 0, 0, false) +layoutspec(::BoolType) = LayoutSpec(VALIDITY_DATA, 0, 0, -1, false) +layoutspec(t::IntType) = LayoutSpec(VALIDITY_DATA, 0, 0, primwidth(t), false) +layoutspec(t::FloatType) = LayoutSpec(VALIDITY_DATA, 0, 0, primwidth(t), false) +layoutspec(t::DecimalType) = LayoutSpec(VALIDITY_DATA, 0, 0, primwidth(t), false) +layoutspec(t::DateType) = LayoutSpec(VALIDITY_DATA, 0, 0, primwidth(t), false) +layoutspec(t::TimeType) = LayoutSpec(VALIDITY_DATA, 0, 0, primwidth(t), false) +layoutspec(t::TimestampType) = LayoutSpec(VALIDITY_DATA, 0, 0, 8, false) +layoutspec(t::DurationType) = LayoutSpec(VALIDITY_DATA, 0, 0, 8, false) +layoutspec(t::IntervalType) = LayoutSpec(VALIDITY_DATA, 0, 0, primwidth(t), false) +layoutspec(t::FixedSizeBinaryType) = LayoutSpec(VALIDITY_DATA, 0, 0, t.nbytes, false) +layoutspec(t::BinaryType) = LayoutSpec(VALIDITY_OFFSETS_DATA, 0, t.large ? 8 : 4, 0, false) +layoutspec(t::Utf8Type) = LayoutSpec(VALIDITY_OFFSETS_DATA, 0, t.large ? 8 : 4, 0, false) +layoutspec(t::ListType) = LayoutSpec(VALIDITY_OFFSETS, 1, t.large ? 8 : 4, 0, false) +layoutspec(::FixedSizeListType) = LayoutSpec(VALIDITY_ONLY, 1, 0, 0, false) +layoutspec(::StructType) = LayoutSpec(VALIDITY_ONLY, -1, 0, 0, false) +layoutspec(::MapType) = LayoutSpec(VALIDITY_OFFSETS, 1, 4, 0, false) +layoutspec(t::UnionType) = + t.mode == SparseMode ? LayoutSpec(TYPE_IDS_ONLY, -1, 0, 0, false) : + LayoutSpec(TYPE_IDS_ELEMENT_OFFSETS, -1, 4, 0, false) +layoutspec(t::DictionaryType) = + LayoutSpec(VALIDITY_DATA, 0, 0, primwidth(t.indextype), false) +layoutspec(::ViewType) = LayoutSpec(VALIDITY_VIEWS, 0, 0, 16, true) +# ListView has one offset and one size per parent slot. These are not +# the length+1 monotone range offsets used by List/Utf8/Binary. +layoutspec(t::ListViewType) = + LayoutSpec(VALIDITY_ELEMENT_OFFSETS_SIZES, 1, t.large ? 8 : 4, 0, false) +# REE: no top-level validity; run_ends and values are CHILDREN, not buffers. +layoutspec(::RunEndEncodedType) = LayoutSpec(NO_BUFFERS, 2, 0, 0, false) + +""" + layoutspec_of(t::ArrowType) -> LayoutSpec + +The closed-set dispatch ladder over the runtime descriptors. This is the +trim-compile story for a runtime-tagged core: dispatch +on an abstract-typed field is dynamic, which JuliaC `--trim=safe` rejects — +but the descriptor set is CLOSED (it is the layout registry), so one +`isa` ladder devirtualizes every generic call site statically. Multiple +dispatch remains the extension surface (each branch calls the ordinary +`layoutspec` method); the ladder is only the entry point generic code uses +when the descriptor's concrete type is unknown. Branches are ordered by +expected frequency. +""" +@inline function layoutspec_of(t::ArrowType)::LayoutSpec + t isa IntType && return layoutspec(t) + t isa FloatType && return layoutspec(t) + t isa Utf8Type && return layoutspec(t) + t isa BoolType && return layoutspec(t) + t isa ListType && return layoutspec(t) + t isa StructType && return layoutspec(t) + t isa DictionaryType && return layoutspec(t) + t isa TimestampType && return layoutspec(t) + t isa DateType && return layoutspec(t) + t isa TimeType && return layoutspec(t) + t isa DurationType && return layoutspec(t) + t isa BinaryType && return layoutspec(t) + t isa FixedSizeBinaryType && return layoutspec(t) + t isa FixedSizeListType && return layoutspec(t) + t isa MapType && return layoutspec(t) + t isa UnionType && return layoutspec(t) + t isa DecimalType && return layoutspec(t) + t isa IntervalType && return layoutspec(t) + t isa NullType && return layoutspec(t) + t isa ViewType && return layoutspec(t) + t isa ListViewType && return layoutspec(t) + t isa RunEndEncodedType && return layoutspec(t) + throw(ArgumentError("unregistered ArrowType")) +end + +# Junk descriptors get a clean error instead of a MethodError wherever the +# raw method table is called directly. +layoutspec(::Any) = throw(ArgumentError("unregistered ArrowType")) + +# --------------------------------------------------------------------------- +# §4 ArrayData +# --------------------------------------------------------------------------- + +""" + ArrayData + +The one physical array representation (≅ C `ArrowArray`): buffers + children ++ optional dictionary, plus logical length and a lazily-computed, cached +null count. Mutable only for the two caches (`nullcount`, `semachecked`); +everything user-visible is immutable after construction. + +`offset` (element offset into the buffers) is carried for C-data import +compatibility; the accessors below apply it uniformly. +""" +mutable struct ArrayData + const type::ArrowType + const len::Int64 + const offset::Int64 + const buffers::FrozenVector{BufferSlice} + const children::FrozenVector{ArrayData} + const dictionary::Union{Nothing,ArrayData} + const owner::Any # adapter lifetime anchor, if needed + @atomic nullcount::Int64 # -1 = unknown, computed on demand + @atomic semachecked::Bool # data-intrinsic semantic checks passed +end + +function ArrayData( + type::ArrowType, + len::Integer, + buffers; + offset::Integer=0, + children=(), + dictionary::Union{Nothing,ArrayData}=nothing, + owner=nothing, + nullcount::Integer=-1, +) + return _arraydata(type, len, buffers, offset, children, dictionary, owner, nullcount) +end + +# Positional twin of the keyword constructor: Julia's kwcall machinery does +# not statically resolve over an abstract-typed leading argument, so +# trim-verified adapters (the C-data import walk) construct through this +# single generic method instead. +function _arraydata( + @nospecialize(type::ArrowType), + len::Integer, + buffers, + offset::Integer, + children, + dictionary::Union{Nothing,ArrayData}, + owner, + nullcount::Integer, +) + len >= 0 || throw(ArgumentError("negative array length")) + offset >= 0 || throw(ArgumentError("negative array offset")) + -1 <= nullcount <= len || + throw(ArgumentError("null count must be -1 or in [0, length]")) + return ArrayData( + type, + Int64(len), + Int64(offset), + FrozenVector{BufferSlice}(buffers), + FrozenVector{ArrayData}(children), + dictionary, + owner, + Int64(nullcount), + false, + ) +end + +Base.length(d::ArrayData) = d.len + +# Adapter-private certificate set. An entry means that one immutable +# dictionary pool snapshot already passed structural, intrinsic semantic, and +# Field-contract validation under the adapter's canonical value Field. +const _ValidatedDictionaries = IdDict{ArrayData,Nothing} +@inline _dictionary_validated(::Nothing, ::ArrayData) = false +@inline _dictionary_validated(memo::_ValidatedDictionaries, d::ArrayData) = haskey(memo, d) + +@inline _slotindex0(d::ArrayData, i::Int64) = + checked_add(d.offset, checked_sub(i, Int64(1))) +@inline _slotbyteoff(d::ArrayData, i::Int64, width::Integer) = + checked_mul(_slotindex0(d, i), Int64(width)) + +# Buffer-by-role lookup, driven by the registry. Structural validation +# guarantees position/arity, so adapters and accessors never hand-count. +function rolebuffer(d::ArrayData, role::BufferRole) + spec = layoutspec_of(d.type) + idx = findfirst(==(role), spec.buffers) + idx === nothing && + throw(ArgumentError("layout $(descriptorname(d.type)) has no $role buffer")) + return d.buffers[idx] +end + +validitybuffer(d::ArrayData) = rolebuffer(d, VALIDITY) + +""" + isvalid_at(d, i) + +Element validity for 1-based logical index `i`. An empty validity slice +means "no nulls recorded" — every element valid (the spec's empty-bitmap +convention). Layouts with no validity buffer (null, union, REE) answer +through their own accessors. +""" +@inline function isvalid_at(d::ArrayData, i::Integer) + v = validitybuffer(d) + isempty_buffer(v) && return true + return getbit(v, _slotindex0(d, Int64(i))) +end + +""" + nullcount(d) -> Int64 + +Cached lazy null count (the polars `unset_bit_count_cache` idea). The benign +race — two tasks computing the same value and both storing it — is +harmless; `:monotonic` ordering is all the cache needs. +""" +function nullcount(d::ArrayData) + nc = @atomic :monotonic d.nullcount + nc >= 0 && return nc + nc = _count_nulls(d) + @atomic :monotonic d.nullcount = nc + return nc +end + +function _count_nulls(d::ArrayData) + d.type isa NullType && return d.len + spec = layoutspec_of(d.type) + isempty(spec.buffers) && return Int64(0) + spec.buffers[1] == VALIDITY || return Int64(0) # unions: no top-level nulls + v = d.buffers[1] + isempty_buffer(v) && return Int64(0) + n = Int64(0) + for i = 1:(d.len) + n += !getbit(v, _slotindex0(d, Int64(i))) + end + return n +end + +# --------------------------------------------------------------------------- +# §5 Staged validation +# --------------------------------------------------------------------------- + +struct ValidationError <: Exception + msg::String +end + +function expected_validity_bytes(len::Int64) + len >= 0 || throw(ArgumentError("negative bitmap length")) + # `len + 7` can overflow a hostile logical extent; the requirement is + # then unsatisfiable, which is a validation failure, not a raw + # `OverflowError`. + len <= typemax(Int64) - 7 || + throw(ValidationError("required bitmap bytes overflow Int64 for $len slots")) + return (len + 7) >> 3 +end + +# Structural requirements multiply untrusted logical extents. When the +# product overflows Int64 no real buffer or child can satisfy it, so it +# surfaces as the same ValidationError an undersized one gets, not as a +# raw `OverflowError`. Both operands are non-negative at every call site. +@inline function _required_extent(count::Int64, width::Int64, what::AbstractString) + count == 0 || + width <= div(typemax(Int64), count) || + throw(ValidationError("required $what overflows Int64 ($count × $width)")) + return count * width +end + +# Runtime descriptor equality must compare values, not only Julia types. +# One more closed-set ladder (no fieldcount/getfield reflection — that is +# dynamic and trim-hostile); each branch compares its descriptor's fields +# explicitly. +# Non-recursive: the spec forbids dictionary-encoded dictionary VALUES, so +# a DictionaryType's valuetype is never itself a DictionaryType (enforced in +# _validate_descriptor) and one nested ladder suffices — which is what lets +# both levels inline at abstract call sites for trim. +@inline function _typeequal_nondict(a::ArrowType, b::ArrowType) + a isa IntType && return b isa IntType && a.bits == b.bits && a.signed == b.signed + a isa FloatType && return b isa FloatType && a.bits == b.bits + a isa Utf8Type && return b isa Utf8Type && a.large == b.large + a isa BoolType && return b isa BoolType + a isa ListType && return b isa ListType && a.large == b.large + a isa StructType && return b isa StructType + a isa TimestampType && + return b isa TimestampType && a.unit == b.unit && a.timezone == b.timezone + a isa DateType && return b isa DateType && a.unit == b.unit + a isa TimeType && return b isa TimeType && a.unit == b.unit && a.bits == b.bits + a isa DurationType && return b isa DurationType && a.unit == b.unit + a isa BinaryType && return b isa BinaryType && a.large == b.large + a isa FixedSizeBinaryType && return b isa FixedSizeBinaryType && a.nbytes == b.nbytes + a isa FixedSizeListType && return b isa FixedSizeListType && a.listsize == b.listsize + a isa MapType && return b isa MapType && a.keyssorted == b.keyssorted + a isa UnionType && return b isa UnionType && a.mode == b.mode && a.typeids == b.typeids + a isa DecimalType && return b isa DecimalType && + a.precision == b.precision && + a.scale == b.scale && + a.bits == b.bits + a isa IntervalType && return b isa IntervalType && a.unit == b.unit + a isa NullType && return b isa NullType + a isa ViewType && return b isa ViewType && a.utf8 == b.utf8 + a isa ListViewType && return b isa ListViewType && a.large == b.large + a isa RunEndEncodedType && return b isa RunEndEncodedType + return false +end + +@inline function typeequal(a::ArrowType, b::ArrowType) + if a isa DictionaryType + return b isa DictionaryType && + a.indextype.bits == b.indextype.bits && + a.indextype.signed == b.indextype.signed && + _typeequal_nondict(a.valuetype, b.valuetype) && + a.ordered == b.ordered + end + b isa DictionaryType && return false + return _typeequal_nondict(a, b) +end + +""" + descriptorname(t::ArrowType) -> Symbol + +Closed-set name ladder for error messages: `nameof(typeof(x))` on an +abstract-typed value is itself a dynamic call, so diagnostics use this +instead. +""" +@inline function descriptorname(t::ArrowType)::Symbol + t isa IntType && return :IntType + t isa FloatType && return :FloatType + t isa Utf8Type && return :Utf8Type + t isa BoolType && return :BoolType + t isa ListType && return :ListType + t isa StructType && return :StructType + t isa DictionaryType && return :DictionaryType + t isa TimestampType && return :TimestampType + t isa DateType && return :DateType + t isa TimeType && return :TimeType + t isa DurationType && return :DurationType + t isa BinaryType && return :BinaryType + t isa FixedSizeBinaryType && return :FixedSizeBinaryType + t isa FixedSizeListType && return :FixedSizeListType + t isa MapType && return :MapType + t isa UnionType && return :UnionType + t isa DecimalType && return :DecimalType + t isa IntervalType && return :IntervalType + t isa NullType && return :NullType + t isa ViewType && return :ViewType + t isa ListViewType && return :ListViewType + t isa RunEndEncodedType && return :RunEndEncodedType + return :UnknownArrowType +end + +_validate_descriptor(::Utf8Type) = nothing +_validate_descriptor(::BoolType) = nothing +_validate_descriptor(::ListType) = nothing +_validate_descriptor(::StructType) = nothing +_validate_descriptor(::BinaryType) = nothing +_validate_descriptor(::MapType) = nothing +_validate_descriptor(::NullType) = nothing +_validate_descriptor(::ViewType) = nothing +_validate_descriptor(::ListViewType) = nothing +_validate_descriptor(::RunEndEncodedType) = nothing +_validate_descriptor(::Any) = throw(ArgumentError("unregistered ArrowType")) + +@inline function _validate_descriptor_of(t::ArrowType) + t isa IntType && return _validate_descriptor(t) + t isa FloatType && return _validate_descriptor(t) + t isa Utf8Type && return _validate_descriptor(t) + t isa BoolType && return _validate_descriptor(t) + t isa ListType && return _validate_descriptor(t) + t isa StructType && return _validate_descriptor(t) + t isa DictionaryType && return _validate_descriptor(t) + t isa TimestampType && return _validate_descriptor(t) + t isa DateType && return _validate_descriptor(t) + t isa TimeType && return _validate_descriptor(t) + t isa DurationType && return _validate_descriptor(t) + t isa BinaryType && return _validate_descriptor(t) + t isa FixedSizeBinaryType && return _validate_descriptor(t) + t isa FixedSizeListType && return _validate_descriptor(t) + t isa MapType && return _validate_descriptor(t) + t isa UnionType && return _validate_descriptor(t) + t isa DecimalType && return _validate_descriptor(t) + t isa IntervalType && return _validate_descriptor(t) + t isa NullType && return _validate_descriptor(t) + t isa ViewType && return _validate_descriptor(t) + t isa ListViewType && return _validate_descriptor(t) + t isa RunEndEncodedType && return _validate_descriptor(t) + throw(ArgumentError("unregistered ArrowType")) +end +_validate_descriptor(t::IntType) = + t.bits in (8, 16, 32, 64) || + throw(ValidationError("integer bit width must be 8, 16, 32, or 64")) +_validate_descriptor(t::FloatType) = + t.bits in (16, 32, 64) || + throw(ValidationError("floating-point bit width must be 16, 32, or 64")) +function _validate_descriptor(t::DecimalType) + maxprecision = + t.bits == 32 ? 9 : t.bits == 64 ? 18 : t.bits == 128 ? 38 : t.bits == 256 ? 76 : 0 + maxprecision != 0 || + throw(ValidationError("decimal bit width must be 32, 64, 128, or 256")) + 1 <= t.precision <= maxprecision || throw( + ValidationError( + "decimal precision $(t.precision) is invalid for $(t.bits)-bit storage", + ), + ) + typemin(Int32) <= t.scale <= typemax(Int32) || throw( + ValidationError("decimal scale $(t.scale) does not fit the Arrow Int32 wire field"), + ) + return nothing +end +_validate_descriptor(t::FixedSizeBinaryType) = + 0 <= t.nbytes <= typemax(Int32) || + throw(ValidationError("fixed-size-binary width must be in [0, $(typemax(Int32))]")) +_validate_descriptor(t::DateType) = + t.unit in (DAY, MILLISECOND_DATE) || + throw(ValidationError("invalid Arrow date unit $(repr(t.unit))")) +function _validate_descriptor(t::TimeType) + t.unit in (SECOND, MILLISECOND, MICROSECOND, NANOSECOND) || + throw(ValidationError("invalid Arrow time unit $(repr(t.unit))")) + valid = t.unit in (SECOND, MILLISECOND) ? t.bits == 32 : t.bits == 64 + valid || throw( + ValidationError("time unit $(t.unit) is incompatible with $(t.bits)-bit storage"), + ) + return nothing +end +function _validate_descriptor(t::TimestampType) + t.unit in (SECOND, MILLISECOND, MICROSECOND, NANOSECOND) || + throw(ValidationError("invalid Arrow timestamp unit $(repr(t.unit))")) + (t.timezone === nothing || isvalid(t.timezone)) || + throw(ValidationError("timestamp timezone is not valid UTF-8")) + return nothing +end +_validate_descriptor(t::DurationType) = + t.unit in (SECOND, MILLISECOND, MICROSECOND, NANOSECOND) || + throw(ValidationError("invalid Arrow duration unit $(repr(t.unit))")) +_validate_descriptor(t::IntervalType) = + t.unit in (YEAR_MONTH, DAY_TIME, MONTH_DAY_NANO) || + throw(ValidationError("invalid Arrow interval unit $(repr(t.unit))")) +_validate_descriptor(t::FixedSizeListType) = + 0 <= t.listsize <= typemax(Int32) || + throw(ValidationError("fixed-size-list size must be in [0, $(typemax(Int32))]")) +_validate_descriptor(t::UnionType) = + t.mode in (SparseMode, DenseMode) || + throw(ValidationError("invalid Arrow union mode $(repr(t.mode))")) +function _validate_descriptor(t::DictionaryType) + _validate_descriptor(t.indextype) + # The spec forbids dictionary-encoded dictionary values; enforcing it + # here is also what keeps descriptor equality non-recursive (typeequal). + t.valuetype isa DictionaryType && + throw(ValidationError("dictionary values cannot themselves be dictionary-encoded")) + _validate_descriptor_of(t.valuetype) + return nothing +end + +function _validate_metadata(metadata, what::AbstractString) + metadata === nothing && return nothing + for (key, value) in metadata + isvalid(key) || throw(ValidationError("$what metadata key is not valid UTF-8")) + isvalid(value) || throw(ValidationError("$what metadata value is not valid UTF-8")) + end + return nothing +end + +function _validate_schema(s::Schema) + s.endianness in (LittleEndian, BigEndian) || + throw(ValidationError("invalid Arrow schema endianness $(repr(s.endianness))")) + s.endianness == _native_endianness() || throw( + ValidationError( + "non-native Arrow schema endianness must be normalized before Core access", + ), + ) + _validate_metadata(s.metadata, "schema") + return s +end + +""" + validate_structural(field, data) + +Stage-2 validation: O(buffers), registry-driven, run at +construction/adaptation time. Checks buffer arity against the layout, and +every buffer's byte length against what the logical length requires — with +checked arithmetic, because these lengths come from untrusted metadata. +Recurses into children and the dictionary. + +BufferSlice construction has already bounded every slice inside its region, +so this stage never touches memory — it is pure arithmetic on declared +sizes. (The framing stage — resource limits before metadata-directed decode +allocation and checked message-body spans — belongs to the adapters; see +src/ipc_read.jl.) +""" +validate_structural(f::Field, d::ArrayData) = _validate_structural(f, d, nothing) + +function _validate_structural( + f::Field, + d::ArrayData, + validated_dictionaries::Union{Nothing,_ValidatedDictionaries}, +) + isvalid(f.name) || throw(ValidationError("field name is not valid UTF-8")) + _validate_metadata(f.metadata, "field") + typeequal(f.type, d.type) || throw( + ValidationError( + "field/type mismatch: $(descriptorname(f.type)) vs $(descriptorname(d.type))", + ), + ) + _validate_descriptor_of(d.type) + spec = layoutspec_of(d.type) + nfixed = length(spec.buffers) + buffers_ok = spec.variadic ? length(d.buffers) >= nfixed : length(d.buffers) == nfixed + buffers_ok || throw( + ValidationError( + "$(descriptorname(d.type)): expected $(spec.variadic ? "at least " : "")$nfixed buffers, got $(length(d.buffers))", + ), + ) + # `len` and `offset` are individually non-negative (constructor + # invariant), but hostile metadata can still declare a sum past Int64 — + # surface that as ValidationError, not a raw OverflowError. + d.len <= typemax(Int64) - d.offset || throw( + ValidationError("array length $(d.len) plus offset $(d.offset) overflows Int64"), + ) + total::Int64 = d.len + d.offset + declared_nulls = @atomic :monotonic d.nullcount + for (i, role) in enumerate(spec.buffers) + b = d.buffers[i] + if role == VALIDITY + if isempty_buffer(b) + declared_nulls > 0 && throw( + ValidationError("absent validity bitmap with positive null count"), + ) + continue + end + b.len >= expected_validity_bytes(total) || throw( + ValidationError( + "validity bitmap too small: $(b.len) bytes for $total slots", + ), + ) + elseif role == DATA + if spec.fixedwidth > 0 + need = _required_extent(total, Int64(spec.fixedwidth), "data buffer bytes") + b.len >= need || + throw(ValidationError("data buffer too small: $(b.len) < $need bytes")) + elseif spec.fixedwidth == -1 # bit-packed (Bool) + b.len >= expected_validity_bytes(total) || + throw(ValidationError("bit-packed data buffer too small")) + end + # fixedwidth == 0 (varbinary DATA): bounded by offsets in the + # semantic stage — nothing structural to require here. + elseif role == OFFSETS + # Canonical empty offset-based arrays may omit the offsets + # buffer. A sliced empty array (`offset > 0`) still needs the + # physical prefix that its offset addresses. + isempty_buffer(b) && d.len == 0 && d.offset == 0 && continue + # `total + 1` offset slots: the +1 itself can overflow a hostile + # extent. + total < typemax(Int64) || + throw(ValidationError("required offsets buffer bytes overflow Int64")) + need = _required_extent( + total + Int64(1), + Int64(spec.offsetwidth), + "offsets buffer bytes", + ) + b.len >= need || + throw(ValidationError("offsets buffer too small: $(b.len) < $need bytes")) + elseif role == ELEMENT_OFFSETS + need = _required_extent(total, Int64(spec.offsetwidth), "element-offsets bytes") + b.len >= need || throw( + ValidationError("element-offsets buffer too small: $(b.len) < $need bytes"), + ) + elseif role == SIZES + need = _required_extent(total, Int64(spec.offsetwidth), "sizes buffer bytes") + b.len >= need || throw(ValidationError("sizes buffer too small")) + elseif role == TYPE_IDS + b.len >= total || throw(ValidationError("type_ids buffer too small")) + elseif role == VIEWS + need = _required_extent(total, Int64(16), "views buffer bytes") + b.len >= need || throw(ValidationError("views buffer too small")) + end + end + # Child arity: registry-declared, or Field-declared for struct/union. + expected_children = spec.childcount == -1 ? length(f.children) : spec.childcount + if !(d.type isa DictionaryType) + length(f.children) == expected_children || throw( + ValidationError( + "$(descriptorname(d.type)): expected $expected_children child fields, got $(length(f.children))", + ), + ) + end + length(d.children) == expected_children || throw( + ValidationError( + "$(descriptorname(d.type)): expected $expected_children children, got $(length(d.children))", + ), + ) + for (cf, cd) in zip(childfields(f), d.children) + _validate_structural(cf, cd, validated_dictionaries) + end + if d.type isa DictionaryType + d.dictionary === nothing && + throw(ValidationError("dictionary-encoded array without a dictionary")) + dictionary = d.dictionary::ArrayData + _dictionary_validated(validated_dictionaries, dictionary) || _validate_structural( + dictvaluefield(f, d.type), + dictionary, + validated_dictionaries, + ) + elseif d.dictionary !== nothing + throw(ValidationError("dictionary values attached to a non-dictionary array")) + end + fslt = d.type + if fslt isa FixedSizeListType + need = _required_extent(total, Int64(fslt.listsize), "fixed-size-list child length") + length(d.children[1]) >= need || throw( + ValidationError( + "fixed-size-list child too short: $(length(d.children[1])) < $need", + ), + ) + end + # Struct and sparse-union children are parent-length arrays indexed at + # parent.offset + i (each child then applies its own offset), so every + # child must cover offset+len slots. + if d.type isa StructType || (d.type isa UnionType && d.type.mode == SparseMode) + for (ci, child) in enumerate(d.children) + length(child) >= total || throw( + ValidationError( + "child $ci too short for parent extent: $(length(child)) < $total", + ), + ) + end + end + ut = d.type + if ut isa UnionType + length(ut.typeids) == length(f.children) || + throw(ValidationError("union type-id count must equal child count")) + length(unique(ut.typeids)) == length(ut.typeids) || + throw(ValidationError("union type ids must be unique")) + all(>=(0), ut.typeids) || + throw(ValidationError("union type ids must be in [0, 127]")) + end + if d.type isa MapType + entries = f.children[1] + entries.type isa StructType || + throw(ValidationError("map child must be an entries struct")) + !entries.nullable || + throw(ValidationError("map entries field must be non-nullable")) + length(entries.children) == 2 || + throw(ValidationError("map entries struct must have key and value children")) + !entries.children[1].nullable || + throw(ValidationError("map keys must be non-nullable")) + end + if d.type isa RunEndEncodedType + runfield, valuefield = f.children + runfield.name == "run_ends" && valuefield.name == "values" || + throw(ValidationError("REE children must be named run_ends and values")) + runtype0 = runfield.type + runtype0 isa IntType || + throw(ValidationError("REE run ends must be signed int16, int32, or int64")) + runtype = runtype0::IntType + (runtype.signed && runtype.bits in (16, 32, 64)) || + throw(ValidationError("REE run ends must be signed int16, int32, or int64")) + !runfield.nullable || throw(ValidationError("REE run ends must be non-nullable")) + # The parent has no validity bitmap, so a positive count is + # malformed; unknown (-1) is legal for any layout. + declared_nulls <= 0 || throw(ValidationError("REE parent null count must be zero")) + length(d.children[1]) == length(d.children[2]) || + throw(ValidationError("REE run-end and value child lengths must match")) + d.len == 0 || + length(d.children[1]) > 0 || + throw( + ValidationError("a nonempty REE array requires at least one physical run"), + ) + maxrunend = + runtype.bits == 16 ? Int64(typemax(Int16)) : + runtype.bits == 32 ? Int64(typemax(Int32)) : typemax(Int64) + total <= maxrunend || throw( + ValidationError( + "REE logical extent $total exceeds the $(runtype.bits)-bit run-end range", + ), + ) + end + return d +end + +# Child Fields for traversal. For list-ish layouts the child field is the +# Field's single declared child; dictionary values reuse the field with the +# value type. +childfields(f::Field) = f.children +# Dictionary values have their own nullability. The index field's +# nullable flag describes only the indices and cannot constrain the pool. +dictvaluefield(f::Field, t::DictionaryType) = + Field(f.name, t.valuetype; nullable=true, children=f.children) + +const MILLISECONDS_PER_DAY = Int64(86_400_000) + +# Date64 whole-day divisibility and Decimal precision are ADVISORY in +# practice: the spec phrases Date64 as "evenly divisible by 86400000" and +# precision as "total number of decimal digits", but the reference C++ +# implementation neither enforces them on read nor honors them on write — the +# apache/arrow-testing gold corpus itself carries Date64 values off day +# boundaries and decimal(3,2) values with five digits. Rejecting those in +# `validate_semantic` would make a conforming reader refuse canonical data, +# so both checks live in the opt-in `validate_full` tier +# (`_validate_full_content`), where strict callers can still demand them. +function _validate_advisory_values(t::DateType, d::ArrayData) + t.unit == MILLISECOND_DATE || return nothing + data = rolebuffer(d, DATA) + for i = 1:(d.len) + isvalid_at(d, i) || continue + value = loadat(data, Int64, _slotbyteoff(d, Int64(i), 8)) + value % MILLISECONDS_PER_DAY == 0 || + throw(ValidationError("Date64 value $value is not a whole day in milliseconds")) + end + return nothing +end + +function _decimal_limb( + t::DecimalType, + data::BufferSlice, + byteoff::Int64, + nlimbs::Int, + limb::Int, +)::UInt64 + limb <= nlimbs || return UInt64(0) + source_limb = _native_endianness() == LittleEndian ? limb : nlimbs - limb + 1 + base = checked_add(byteoff, Int64(8 * (source_limb - 1))) + return t.bits == 32 ? UInt64(loadat(data, UInt32, base)) : loadat(data, UInt64, base) +end + +function _decimal_fits_precision(t::DecimalType, data::BufferSlice, byteoff::Int64) + # Core accepts only native-endian array buffers. Arrow decimal storage is + # a two's-complement integer, so put native chunks into least-significant + # limb order before comparing its magnitude with 10^p. Work in fixed + # UInt256-style arithmetic so Core stays Base-only and Decimal256 does not + # require BigInt allocations or BitIntegers. (No closures here: captured + # and reassigned locals box, which defeats trim verification.) + nlimbs = cld(t.bits, 64) + l1 = _decimal_limb(t, data, byteoff, nlimbs, 1) + l2 = _decimal_limb(t, data, byteoff, nlimbs, 2) + l3 = _decimal_limb(t, data, byteoff, nlimbs, 3) + l4 = _decimal_limb(t, data, byteoff, nlimbs, 4) + signbit = t.bits == 32 ? UInt64(1) << 31 : UInt64(1) << 63 + negative = + ((nlimbs == 1 ? l1 : nlimbs == 2 ? l2 : nlimbs == 3 ? l3 : l4) & signbit) != 0 + if negative && t.bits == 32 + l1 |= typemax(UInt64) << 32 + end + m1 = nlimbs >= 1 ? (negative ? ~l1 : l1) : UInt64(0) + m2 = nlimbs >= 2 ? (negative ? ~l2 : l2) : UInt64(0) + m3 = nlimbs >= 3 ? (negative ? ~l3 : l3) : UInt64(0) + m4 = nlimbs >= 4 ? (negative ? ~l4 : l4) : UInt64(0) + if negative + m1 += UInt64(1) + c = m1 == 0 + m2 += c ? UInt64(1) : UInt64(0) + c &= m2 == 0 + m3 += c ? UInt64(1) : UInt64(0) + c &= m3 == 0 + m4 += c ? UInt64(1) : UInt64(0) + end + + L1, L2, L3, L4 = UInt64(1), UInt64(0), UInt64(0), UInt64(0) + for _ = 1:(t.precision) + p1 = UInt128(L1) * 10 + p2 = UInt128(L2) * 10 + (p1 >> 64) + p3 = UInt128(L3) * 10 + (p2 >> 64) + p4 = UInt128(L4) * 10 + (p3 >> 64) + L1 = UInt64(p1 & UInt128(typemax(UInt64))) + L2 = UInt64(p2 & UInt128(typemax(UInt64))) + L3 = UInt64(p3 & UInt128(typemax(UInt64))) + L4 = UInt64(p4 & UInt128(typemax(UInt64))) + end + m4 < L4 && return true + m4 > L4 && return false + m3 < L3 && return true + m3 > L3 && return false + m2 < L2 && return true + m2 > L2 && return false + return m1 < L1 +end + +function _validate_advisory_values(t::DecimalType, d::ArrayData) + data = rolebuffer(d, DATA) + width = Int64(primwidth(t)) + for i = 1:(d.len) + isvalid_at(d, i) || continue + byteoff = _slotbyteoff(d, Int64(i), width) + _decimal_fits_precision(t, data, byteoff) || throw( + ValidationError( + "Decimal value at element $i does not fit precision $(t.precision)", + ), + ) + end + return nothing +end + +# Time-of-day range is advisory for the same reason as Date64 divisibility: +# the 1.0.0 gold corpus carries out-of-range Time32 values that C++ reads. +function _validate_advisory_values(t::TimeType, d::ArrayData) + units_per_day = + t.unit == SECOND ? Int64(86_400) : + t.unit == MILLISECOND ? MILLISECONDS_PER_DAY : + t.unit == MICROSECOND ? Int64(86_400_000_000) : Int64(86_400_000_000_000) + data = rolebuffer(d, DATA) + for i = 1:(d.len) + isvalid_at(d, i) || continue + value = if t.bits == 32 + Int64(loadat(data, Int32, _slotbyteoff(d, Int64(i), Int64(4)))) + else + loadat(data, Int64, _slotbyteoff(d, Int64(i), Int64(8))) + end + 0 <= value < units_per_day || throw( + ValidationError( + "Time value $value is outside [0, $units_per_day) for $(t.unit)", + ), + ) + end + return nothing +end + +""" + validate_semantic(field, data) + +Stage-3 validation. This public stage composes structural validation before +any content access, so callers cannot accidentally certify malformed buffer +geometry by skipping `validate_structural`. Data-intrinsic checks are cached +on the ArrayData (`semachecked`); benign concurrent callers may repeat the +same scan. Field-dependent dictionary contracts run on every call because +the same data can be checked against another Field. Per-slot nullability +enforcement is advisory and lives in the opt-in `validate_full` tier. +""" +function validate_semantic(f::Field, d::ArrayData) + return _validate_semantic(f, d, nothing) +end + +function _validate_semantic( + f::Field, + d::ArrayData, + validated_dictionaries::Union{Nothing,_ValidatedDictionaries}, +) + _validate_structural(f, d, validated_dictionaries) + _validate_semantic_intrinsic(f, d, validated_dictionaries) + _validate_field_contracts(f, d, validated_dictionaries) + return d +end + +function _validate_semantic_intrinsic( + f::Field, + d::ArrayData, + validated_dictionaries::Union{Nothing,_ValidatedDictionaries}, +) + t = d.type + if !(@atomic :monotonic d.semachecked) + spec = layoutspec_of(t) + oi = findfirst(==(OFFSETS), spec.buffers) + if oi !== nothing && spec.offsetwidth != 0 + wide = spec.offsetwidth == 8 + offs = d.buffers[oi] + if !(isempty_buffer(offs) && d.len == 0 && d.offset == 0) + databytes = if t isa Utf8Type || t isa BinaryType + di = findfirst(==(DATA), spec.buffers) + d.buffers[di].len + else + isempty(d.children) ? Int64(0) : Int64(length(d.children[1])) + end + prev = _load_offset(offs, wide, d.offset) + prev >= 0 || throw(ValidationError("negative first offset")) + for i = 1:(d.len) + cur = _load_offset(offs, wide, checked_add(d.offset, Int64(i))) + cur >= prev || throw( + ValidationError("offsets not monotonically non-decreasing at $i"), + ) + prev = cur + end + prev <= databytes || throw( + ValidationError("final offset $prev exceeds data extent $databytes"), + ) + end + end + if t isa DictionaryType + dictlen = length(d.dictionary) + data = rolebuffer(d, DATA) + w = primwidth(t.indextype) + for i = 1:(d.len) + isvalid_at(d, i) || continue + idx = _load_int(data, t.indextype, _slotbyteoff(d, Int64(i), w)) + 0 <= idx < dictlen || throw( + ValidationError("dictionary index $idx out of bounds [0, $dictlen)"), + ) + end + end + if t isa UnionType + ids = rolebuffer(d, TYPE_IDS) + # Dense-mode-only bookkeeping; sparse unions never read it. + lastoffset = t.mode == DenseMode ? fill(Int64(-1), length(d.children)) : Int64[] + for i = 1:(d.len) + tid = loadat(ids, Int8, _slotindex0(d, Int64(i))) + pos = findfirst(==(tid), t.typeids) + pos === nothing && + throw(ValidationError("union type id $tid not in declared domain")) + if t.mode == DenseMode + off = loadat( + rolebuffer(d, ELEMENT_OFFSETS), + Int32, + _slotbyteoff(d, Int64(i), 4), + ) + 0 <= off < length(d.children[pos]) || throw( + ValidationError( + "dense union offset $off out of bounds for child $pos", + ), + ) + Int64(off) >= lastoffset[pos] || throw( + ValidationError( + "dense union offsets must be nondecreasing within child $pos", + ), + ) + lastoffset[pos] = Int64(off) + end + end + end + t isa ViewType && _validate_view_values(t, d) + t isa ListViewType && _validate_listview_values(t, d) + t isa RunEndEncodedType && _validate_ree_values(d) + actual_nulls = _count_nulls(d) + declared_nulls = @atomic :monotonic d.nullcount + if declared_nulls >= 0 && declared_nulls != actual_nulls + throw( + ValidationError( + "declared null count $declared_nulls does not match bitmap count $actual_nulls", + ), + ) + elseif declared_nulls < 0 + # Promote the unknown sentinel to the counted value: later + # structural checks on this same array then enforce the + # absent-bitmap rule against a known count instead of skipping it. + @atomic :monotonic d.nullcount = actual_nulls + end + @atomic :monotonic d.semachecked = true + end + # The cache flag is per-node, so recurse even when this node is cached: + # each child short-circuits on its own flag. + for (cf, cd) in zip(childfields(f), d.children) + _validate_semantic_intrinsic(cf, cd, validated_dictionaries) + end + if t isa DictionaryType + dictionary = d.dictionary::ArrayData + _dictionary_validated(validated_dictionaries, dictionary) || + _validate_semantic_intrinsic( + dictvaluefield(f, t), + dictionary, + validated_dictionaries, + ) + end + return d +end + +# -- view layouts (format 1.4) ---------------------------------------------- + +# One 16-byte view entry: length, then either 12 inline bytes (length <= 12, +# zero-padded) or prefix + buffer index + offset into one of the variadic +# data buffers that follow validity and views. +const VIEW_INLINE_MAX = Int32(12) + +@inline _viewbase(d::ArrayData, i::Int64) = _slotbyteoff(d, i, 16) + +function _viewdatabuffer(d::ArrayData, bufidx::Int32) + nvariadic = length(d.buffers) - 2 + 0 <= bufidx < nvariadic || + throw(ValidationError("view buffer index $bufidx outside [0, $nvariadic)")) + return d.buffers[3 + Int(bufidx)] +end + +""" +Semantic checks for Utf8View/BinaryView: non-null long entries must point +inside their indicated variadic buffer, and the inline prefix MUST be a copy +of the referenced data's first four bytes (the spec's comparison-fast-path +contract). Null entries' bytes are unrestricted by the spec, so only valid +slots are checked. Canonical zero-padding of short entries' unused inline +bytes is a writer recommendation no validation tier enforces: +`validate_full`'s canonical-form checks cover bit-packed buffers only. +""" +function _validate_view_values(t::ViewType, d::ArrayData) + views = rolebuffer(d, VIEWS) + for i = 1:(d.len) + isvalid_at(d, i) || continue + base = _viewbase(d, Int64(i)) + len = loadat(views, Int32, base) + len >= 0 || throw(ValidationError("negative view length $len")) + len <= VIEW_INLINE_MAX && continue + bufidx = loadat(views, Int32, checked_add(base, Int64(8))) + off = Int64(loadat(views, Int32, checked_add(base, Int64(12)))) + data = _viewdatabuffer(d, bufidx) + off >= 0 || throw(ValidationError("negative view offset $off")) + # Both operands are 32-bit loads, so the sum cannot overflow Int64. + checked_add(off, Int64(len)) <= data.len || throw( + ValidationError("view range [$off, $(off + len)) escapes data buffer $bufidx"), + ) + for k = 0:3 + loadat(views, UInt8, checked_add(base, Int64(4 + k))) == + loadat(data, UInt8, checked_add(off, Int64(k))) || + throw(ValidationError("view prefix does not match referenced data")) + end + end + return nothing +end + +@inline function _listview_range(t::ListViewType, d::ArrayData, i::Int64) + wide = t.large + slot = _slotindex0(d, i) + offs = rolebuffer(d, ELEMENT_OFFSETS) + sizes = rolebuffer(d, SIZES) + off = + wide ? loadat(offs, Int64, checked_mul(slot, Int64(8))) : + Int64(loadat(offs, Int32, checked_mul(slot, Int64(4)))) + sz = + wide ? loadat(sizes, Int64, checked_mul(slot, Int64(8))) : + Int64(loadat(sizes, Int32, checked_mul(slot, Int64(4)))) + return off, sz +end + +""" +Semantic checks for ListView/LargeListView. The spec's invariants bind EVERY +slot, null included: `0 <= offsets[i]`, `0 <= sizes[i]`, and +`offsets[i] + sizes[i] <= child length`. Out-of-order and overlapping ranges +are legal — that is the layout's point. +""" +function _validate_listview_values(t::ListViewType, d::ArrayData) + childlen = Int64(length(d.children[1])) + for i = 1:(d.len) + off, sz = _listview_range(t, d, Int64(i)) + (off >= 0 && sz >= 0) || + throw(ValidationError("list-view offset and size must be non-negative")) + # Subtraction form: `off + sz` on two hostile 64-bit loads can + # overflow Int64 (which must read as invalid data, not a raw + # `OverflowError`); the message widens to Int128 so the true + # endpoint prints either way. + (off <= childlen && sz <= childlen - off) || throw( + ValidationError( + "list-view range [$off, $(Int128(off) + Int128(sz))) escapes " * + "child length $childlen", + ), + ) + end + return nothing +end + +""" +Semantic checks for run-end encoding: a signed 16/32/64-bit run-ends child +with no nulls, equal-length children (one value per run), run ends positive +and strictly ascending, and the last run end covering every logical slot +(`>= offset + length` — equality holds for unsliced arrays). The REE parent +has no validity bitmap and its null count field is zero or the unknown +sentinel `-1`; logical nulls live in the values child's runs. The structural +stage has already established the run-ends descriptor (signed 16/32/64-bit), +the equal child lengths, and the parent null count. +""" +function _validate_ree_values(d::ArrayData) + runs, values = d.children[1], d.children[2] + nullcount(runs) == 0 || throw(ValidationError("a run end cannot be null")) + total = checked_add(d.offset, d.len) + data = rolebuffer(runs, DATA) + # The typeassert re-concretizes the structurally established descriptor — + # without it `primwidth`/`_load_int` see `ArrowType` and the trim + # verifier reports unresolved calls. + rti = runs.type::IntType + w = primwidth(rti) + prev = Int64(0) + for i = 1:(runs.len) + re = _load_int(data, rti, _slotbyteoff(runs, Int64(i), w)) + re > prev || + throw(ValidationError("run ends must be positive and strictly ascending")) + prev = re + end + d.len == 0 || + prev >= total || + throw(ValidationError("run ends cover $prev of $total logical slots")) + return nothing +end + +""" +The run whose end first reaches 1-based logical position `offset + i` — +binary search over the run-ends child, the REE random-access primitive. +""" +function _ree_runindex(d::ArrayData, i::Int64) + runs = d.children[1] + rt = runs.type::IntType + data = rolebuffer(runs, DATA) + w = primwidth(rt) + target = checked_add(d.offset, i) + lo, hi = Int64(1), runs.len + while lo < hi + mid = (lo + hi) >>> 1 + re = _load_int(data, rt, _slotbyteoff(runs, mid, w)) + re >= target ? (hi = mid) : (lo = mid + 1) + end + return lo +end + +function _logical_null_at(f::Field, d::ArrayData, i::Int64) + t = d.type + t isa NullType && return true + if t isa RunEndEncodedType + run = _ree_runindex(d, i) + return _logical_null_at(f.children[2], d.children[2], run) + end + if t isa UnionType + cf, cd, childi = _union_child(f, d, i) + return _logical_null_at(cf, cd, childi) + end + spec = layoutspec_of(t) + return !isempty(spec.buffers) && spec.buffers[1] == VALIDITY && !isvalid_at(d, i) +end + +function _union_child(f::Field, d::ArrayData, i::Int64) + t = d.type::UnionType + tid = loadat(rolebuffer(d, TYPE_IDS), Int8, _slotindex0(d, i)) + pos = findfirst(==(tid), t.typeids) + pos === nothing && throw(ValidationError("union type id $tid not in declared domain")) + childi = if t.mode == DenseMode + off = loadat(rolebuffer(d, ELEMENT_OFFSETS), Int32, _slotbyteoff(d, i, 4)) + checked_add(Int64(off), Int64(1)) + else + checked_add(d.offset, i) + end + return f.children[pos], d.children[pos], childi +end + +function _validate_field_contract_at(f::Field, d::ArrayData, i::Int64) + t = d.type + if t isa UnionType + if !f.nullable && _logical_null_at(f, d, i) + throw( + ValidationError( + "non-nullable field $(repr(f.name)) contains a null at element $i", + ), + ) + end + # A union has no parent validity bitmap. Its selected child supplies + # both the value and any logical null, so validate that child even + # when the union Field itself permits nulls. Unselected child slots + # are not part of this selected union value and must remain ignored. + cf, cd, childi = _union_child(f, d, i) + _validate_field_contract_at(cf, cd, childi) + return nothing + end + if t isa RunEndEncodedType + # Same bitmap-less shape as unions: the selected VALUES run supplies + # the value and any logical null. The runs child was already checked + # whole (no nulls, ascending) by the intrinsic stage. + if !f.nullable && _logical_null_at(f, d, i) + throw( + ValidationError( + "non-nullable field $(repr(f.name)) contains a null at element $i", + ), + ) + end + _validate_field_contract_at(f.children[2], d.children[2], _ree_runindex(d, i)) + return nothing + end + + slotnull = t isa NullType || !isvalid_at(d, i) + if slotnull + f.nullable || throw( + ValidationError( + "non-nullable field $(repr(f.name)) contains a null at element $i", + ), + ) + # Child storage below a null parent value is unspecified. In + # particular, null Struct/FixedSizeList slots and null List/Map + # ranges mask nulls in otherwise non-nullable child Fields. + return nothing + end + + if t isa StructType + childi = checked_add(d.offset, i) + for (cf, cd) in zip(f.children, d.children) + _validate_field_contract_at(cf, cd, childi) + end + elseif t isa FixedSizeListType + base = checked_mul(_slotindex0(d, i), Int64(t.listsize)) + cf, cd = f.children[1], d.children[1] + for j = 1:(t.listsize) + _validate_field_contract_at(cf, cd, checked_add(base, Int64(j))) + end + elseif t isa Union{ListType,MapType} + lo, hi = _offsets_at(d, i, layoutspec(t).offsetwidth == 8) + lo == hi && return nothing + cf, cd = f.children[1], d.children[1] + for childi = checked_add(lo, Int64(1)):hi + _validate_field_contract_at(cf, cd, childi) + end + elseif t isa ListViewType + off, sz = _listview_range(t, d, i) + sz == 0 && return nothing + cf, cd = f.children[1], d.children[1] + for childi = checked_add(off, Int64(1)):checked_add(off, sz) + _validate_field_contract_at(cf, cd, childi) + end + end + return nothing +end + +function _validate_dictionary_contracts( + f::Field, + d::ArrayData, + validated_dictionaries::Union{Nothing,_ValidatedDictionaries}, +) + if d.type isa DictionaryType + # Dictionary values form an independent array. Index nullability never + # constrains pool nullability, but nested Field contracts inside the + # pool still apply to every pool value, even when the dictionary array + # itself is nested below a masked parent. + dictionary = d.dictionary::ArrayData + _dictionary_validated(validated_dictionaries, dictionary) || + _validate_field_contracts( + dictvaluefield(f, d.type), + dictionary, + validated_dictionaries, + ) + end + for (cf, cd) in zip(f.children, d.children) + _validate_dictionary_contracts(cf, cd, validated_dictionaries) + end + return nothing +end + +# `Field.nullable` is ADVISORY schema metadata in the ecosystem: the +# reference C++ implementation neither enforces it on read nor rejects a +# non-nullable field whose data holds nulls, and the apache/arrow-testing +# gold corpus carries exactly that (a `nullable=false` union whose selected +# child is null). The semantic stage therefore validates only the +# structurally-load-bearing dictionary contracts; the per-slot nullability +# walk (`_validate_field_contract_at`) runs in the opt-in `validate_full` +# tier for callers who want the declaration enforced. +function _validate_field_contracts( + f::Field, + d::ArrayData, + validated_dictionaries::Union{Nothing,_ValidatedDictionaries}, +) + _validate_dictionary_contracts(f, d, validated_dictionaries) + return nothing +end + +function _validate_nullability(f::Field, d::ArrayData) + for i = 1:(d.len) + _validate_field_contract_at(f, d, Int64(i)) + end + # Dictionary pools are independent arrays: their nested Field contracts + # apply to every pool value regardless of which indices reference them + # (and regardless of masking above the dictionary array), so each pool + # gets its own root walk. + _validate_pool_nullability(f, d) + return nothing +end + +function _validate_pool_nullability(f::Field, d::ArrayData) + if d.type isa DictionaryType + _validate_nullability(dictvaluefield(f, d.type), d.dictionary::ArrayData) + end + for (cf, cd) in zip(f.children, d.children) + _validate_pool_nullability(cf, cd) + end + return nothing +end + +""" + validate_full(field, data) + +Stage-4 (opt-in) content validation. It composes semantic (and therefore +structural) validation before the more expensive whole-content checks. +""" +function validate_full(f::Field, d::ArrayData) + validate_semantic(f, d) + # The nullability walk enters ONCE at the root: it routes through + # unions/REE and applies parent-null masking itself, so recursing it per + # child would flag masked slots that are not part of any public-domain value. + _validate_nullability(f, d) + _validate_full_content(f, d) + return d +end + +# Canonical bit-packed form: the spec recommends writers zero the unused +# trailing bits of the final byte and any padding bytes, and forbids readers +# from relying on either — so enforcement is full-tier only. Sliced arrays +# are exempt: trailing bits inside a shared bitmap window can legitimately +# belong to a sibling slice. +function _validate_canonical_bits(d::ArrayData) + d.offset == 0 && d.len > 0 || return nothing + spec = layoutspec_of(d.type) + for (idx, role) in enumerate(spec.buffers) + role == VALIDITY || (role == DATA && d.type isa BoolType) || continue + b = d.buffers[idx] + nbytes = Int64(cld(d.len, 8)) + # Absent or short bitmaps are the structural tier's concern. + b.len >= nbytes || continue + tail = d.len % 8 + if tail != 0 + mask = UInt8(0xff) << tail + loadat(b, UInt8, nbytes - 1) & mask == 0x00 || throw( + ValidationError( + "canonical form requires zeroed unused bits in the final " * + "byte of a bit-packed buffer", + ), + ) + end + for i = nbytes:(b.len - 1) + loadat(b, UInt8, i) == 0x00 || throw( + ValidationError( + "canonical form requires zeroed padding in bit-packed buffers", + ), + ) + end + end + return nothing +end + +# Closed-set ladder over the three advisory-check descriptors (same +# devirtualization story as layoutspec_of). +@inline function _validate_advisory_values_of(d::ArrayData) + t = d.type + t isa DateType && return _validate_advisory_values(t, d) + t isa TimeType && return _validate_advisory_values(t, d) + t isa DecimalType && return _validate_advisory_values(t, d) + return nothing +end + +function _validate_full_content(f::Field, d::ArrayData) + _validate_advisory_values_of(d) + _validate_canonical_bits(d) + if d.type isa Utf8Type || (d.type isa ViewType && d.type.utf8) + for i = 1:(d.len) + isvalid_at(d, i) || continue + s = getvalue(f, d, i)::String + # A malformed byte sequence iterates as invalid Chars; checking + # every Char is the stdlib idiom for whole-string validity. + all(isvalid, s) || throw(ValidationError("invalid UTF-8 at element $i")) + end + end + for (cf, cd) in zip(childfields(f), d.children) + _validate_full_content(cf, cd) + end + if d.type isa DictionaryType + _validate_full_content(dictvaluefield(f, d.type), d.dictionary) + end + return nothing +end + +# --------------------------------------------------------------------------- +# §6 Element access (per-layout semantics; the facade's raw material) +# --------------------------------------------------------------------------- + +# Julia storage type for each descriptor: what getvalue returns for valid +# elements. Timestamps etc. return their raw storage integers here — the +# *facade* owns the Dates conversion layer; keeping Core conversion-free is +# what lets the C-data and IPC adapters share it unchanged. +juliatype(::BoolType) = Bool +juliatype(t::IntType) = + t.signed ? (t.bits == 8 ? Int8 : t.bits == 16 ? Int16 : t.bits == 32 ? Int32 : Int64) : + (t.bits == 8 ? UInt8 : t.bits == 16 ? UInt16 : t.bits == 32 ? UInt32 : UInt64) +juliatype(t::FloatType) = t.bits == 16 ? Float16 : t.bits == 32 ? Float32 : Float64 +juliatype(::TimestampType) = Int64 +juliatype(::DurationType) = Int64 +juliatype(t::DateType) = t.unit == DAY ? Int32 : Int64 +juliatype(t::TimeType) = t.bits == 32 ? Int32 : Int64 +juliatype(::Utf8Type) = String +juliatype(::BinaryType) = Vector{UInt8} +juliatype(t::FixedSizeBinaryType) = Vector{UInt8} +juliatype(t::ViewType) = t.utf8 ? String : Vector{UInt8} +juliatype(t::DecimalType) = t.bits == 32 ? Int32 : t.bits == 64 ? Int64 : Vector{UInt8} +juliatype(t::IntervalType) = + t.unit == YEAR_MONTH ? Int32 : + t.unit == DAY_TIME ? NamedTuple{(:days, :millis),Tuple{Int32,Int32}} : + NamedTuple{(:months, :days, :nanos),Tuple{Int32,Int32,Int64}} +juliatype(::StructType) = Vector{Pair{String,Any}} +juliatype(::MapType) = Vector{Pair{Any,Any}} + +@inline function _load_int(b::BufferSlice, t::IntType, byteoff::Int64)::Int64 + # Literal load widths avoid a runtime DataType in the raw-load path, which + # produces code that the trim verifier cannot resolve. + if t.signed + t.bits == 64 && return loadat(b, Int64, byteoff) + t.bits == 32 && return Int64(loadat(b, Int32, byteoff)) + t.bits == 16 && return Int64(loadat(b, Int16, byteoff)) + return Int64(loadat(b, Int8, byteoff)) + else + if t.bits == 64 + u = loadat(b, UInt64, byteoff) + # Explicit range check: a stored UInt64 above typemax(Int64) is + # out of domain for every Int64-typed consumer (dictionary + # lengths and run ends are Int64), and a bare `Int64(u)` would + # leak an `InexactError` through the validation tier. + u <= UInt64(typemax(Int64)) || + throw(ValidationError("unsigned 64-bit value $u exceeds the Int64 range")) + return Int64(u) + end + t.bits == 32 && return Int64(loadat(b, UInt32, byteoff)) + t.bits == 16 && return Int64(loadat(b, UInt16, byteoff)) + return Int64(loadat(b, UInt8, byteoff)) + end +end + +""" + getvalue(field, data, i) -> Union{Missing, value} + +Read logical element `i` (1-based). Layout dispatch is the closed-set `isa` +ladder over the runtime descriptor — a type-test chain per call, not a +dynamic dispatch. This is Core's honest contract: +scalar access through the erased representation pays a boundary cost; +`materialize` resolves the layout once and loops through a function +barrier. +""" +function getvalue(f::Field, d::ArrayData, i::Integer) + 1 <= i <= d.len || throw(BoundsError(d, i)) + return _value_of(d.type, f, d, Int64(i)) +end + +# Reader adapters can attach one cumulative allocation budget without making +# the dependency-free Core know the adapter's budget type. The ordinary +# methods below remain the trim-safe API. Budgeted overloads only preflight +# package-owned output containers, then call the same audited extraction. +function _charge_materialization! end +function _materialization_remaining end +function _materialization_limit_exceeded! end + +@inline function _materializedvectorbytes(::Type{T}, n::Integer) where {T} + n >= 0 || throw(ValidationError("materialized vector length is negative")) + payload = checked_mul(Int64(n), Int64(Base.elsize(Vector{T}))) + # Julia stores one selector byte per element beside an isbits-Union + # vector's ordinary payload. + Base.isbitsunion(T) && (payload = checked_add(payload, Int64(n))) + # Supported Julia runtimes allocate at most one 64-byte object for an + # empty Vector after warm-up. Charging the nonempty 128-byte header model + # per empty nested value rejects compact Arrow columns by hundreds of + # megabytes even though their materialized empty containers are small. + payload == 0 && return Int64(64) + # Julia 1.11's `Memory` backing store may round a just-over-half-full + # allocation to the next size class. The measured worst case approaches + # twice the requested payload. Reserve that full capacity plus both the + # Vector and Memory headers; using the logical payload alone can let one + # package-owned allocation exceed the caller's budget by almost 2x. + capacity = checked_mul(Int64(2), payload) + return checked_add(Int64(128), capacity) +end + +@inline function _materializedbitvectorbytes(n::Integer) + n >= 0 || throw(ValidationError("materialized bit-vector length is negative")) + chunks = cld(Int64(n), Int64(64)) + return checked_add( + _materializedvectorbytes(UInt64, chunks), + _materializedobjectbytes(sizeof(BitVector)), + ) +end + +# Conservative Julia heap-object size for a newly boxed inline value: an +# 8-byte object tag plus the payload, rounded to the 16-byte GC size class. +@inline function _materializedobjectbytes(payload::Integer) + payload >= 0 || throw(ValidationError("materialized object size is negative")) + total = checked_add(Int64(payload), Int64(8)) + return checked_mul(Int64(16), cld(total, Int64(16))) +end + +mutable struct _MaterializationEstimate{B} + budget::B + bytes::Int64 + limit::Int64 +end + +_MaterializationEstimate(budget) = + _MaterializationEstimate(budget, Int64(0), _materialization_remaining(budget)) + +@inline function _addestimate!( + estimate::_MaterializationEstimate, + amount::Int64, + what::AbstractString, +) + amount >= 0 || throw(ValidationError("materialization estimate is negative")) + amount <= estimate.limit - estimate.bytes || + _materialization_limit_exceeded!(estimate.budget, what) + estimate.bytes += amount + return nothing +end + +@inline _reservevector!( + estimate::_MaterializationEstimate, + ::Type{T}, + n::Integer, + what::AbstractString, +) where {T} = _addestimate!(estimate, _materializedvectorbytes(T, n), what) + +@inline _reserveobject!( + estimate::_MaterializationEstimate, + payload::Integer, + what::AbstractString, +) = _addestimate!(estimate, _materializedobjectbytes(payload), what) + +@inline function _commitestimate!(estimate::_MaterializationEstimate, what::AbstractString) + _charge_materialization!(estimate.budget, estimate.bytes, what) + return nothing +end + +_value(::Any, ::Field, ::ArrayData, ::Int64) = + throw(ArgumentError("unregistered ArrowType")) + +# -- primitives ------------------------------------------------------------- + +# Primitive accessors branch to LITERAL load widths: `loadat(b, T, off)` with +# a runtime `T::DataType` leaves the raw-load path unresolved under trim +# verification, and a concrete branch is faster anyway. +function _value(t::IntType, f::Field, d::ArrayData, i::Int64) + isvalid_at(d, i) || return missing + b = rolebuffer(d, DATA) + if t.signed + t.bits == 64 && return loadat(b, Int64, _slotbyteoff(d, i, 8)) + t.bits == 32 && return loadat(b, Int32, _slotbyteoff(d, i, 4)) + t.bits == 16 && return loadat(b, Int16, _slotbyteoff(d, i, 2)) + return loadat(b, Int8, _slotbyteoff(d, i, 1)) + else + t.bits == 64 && return loadat(b, UInt64, _slotbyteoff(d, i, 8)) + t.bits == 32 && return loadat(b, UInt32, _slotbyteoff(d, i, 4)) + t.bits == 16 && return loadat(b, UInt16, _slotbyteoff(d, i, 2)) + return loadat(b, UInt8, _slotbyteoff(d, i, 1)) + end +end + +function _value(t::FloatType, f::Field, d::ArrayData, i::Int64) + isvalid_at(d, i) || return missing + b = rolebuffer(d, DATA) + t.bits == 64 && return loadat(b, Float64, _slotbyteoff(d, i, 8)) + t.bits == 32 && return loadat(b, Float32, _slotbyteoff(d, i, 4)) + return loadat(b, Float16, _slotbyteoff(d, i, 2)) +end + +function _value(t::Union{TimestampType,DurationType}, f::Field, d::ArrayData, i::Int64) + isvalid_at(d, i) || return missing + return loadat(rolebuffer(d, DATA), Int64, _slotbyteoff(d, i, 8)) +end + +function _value(t::DateType, f::Field, d::ArrayData, i::Int64) + isvalid_at(d, i) || return missing + b = rolebuffer(d, DATA) + return t.unit == DAY ? loadat(b, Int32, _slotbyteoff(d, i, 4)) : + loadat(b, Int64, _slotbyteoff(d, i, 8)) +end + +function _value(t::TimeType, f::Field, d::ArrayData, i::Int64) + isvalid_at(d, i) || return missing + b = rolebuffer(d, DATA) + return t.bits == 32 ? loadat(b, Int32, _slotbyteoff(d, i, 4)) : + loadat(b, Int64, _slotbyteoff(d, i, 8)) +end + +function _value(t::DecimalType, f::Field, d::ArrayData, i::Int64) + isvalid_at(d, i) || return missing + w = primwidth(t) + # 128/256-bit decimals surface as raw native-endian bytes (BigInt/Int256 + # conversion is the facade's); 32/64 as integers. Core RecordBatches + # accept native-endian buffers only. + if t.bits == 32 + return loadat(rolebuffer(d, DATA), Int32, _slotbyteoff(d, i, w)) + elseif t.bits == 64 + return loadat(rolebuffer(d, DATA), Int64, _slotbyteoff(d, i, w)) + else + b = rolebuffer(d, DATA) + off = _slotbyteoff(d, i, w) + return [loadat(b, UInt8, checked_add(off, Int64(k))) for k = 0:(w - 1)] + end +end + +function _value(t::IntervalType, f::Field, d::ArrayData, i::Int64) + isvalid_at(d, i) || return missing + b = rolebuffer(d, DATA) + if t.unit == YEAR_MONTH + return loadat(b, Int32, _slotbyteoff(d, i, 4)) + elseif t.unit == DAY_TIME + off = _slotbyteoff(d, i, 8) + return ( + days=loadat(b, Int32, off), + millis=loadat(b, Int32, checked_add(off, Int64(4))), + ) + else # MONTH_DAY_NANO + off = _slotbyteoff(d, i, 16) + return ( + months=loadat(b, Int32, off), + days=loadat(b, Int32, checked_add(off, Int64(4))), + nanos=loadat(b, Int64, checked_add(off, Int64(8))), + ) + end +end + +function _value(::BoolType, f::Field, d::ArrayData, i::Int64) + isvalid_at(d, i) || return missing + return getbit(rolebuffer(d, DATA), _slotindex0(d, i)) +end + +function _value(::NullType, f::Field, d::ArrayData, i::Int64) + return missing +end + +function _value(t::FixedSizeBinaryType, f::Field, d::ArrayData, i::Int64) + isvalid_at(d, i) || return missing + b = rolebuffer(d, DATA) + off = _slotbyteoff(d, i, t.nbytes) + return slicebytes(subslice(b, off, t.nbytes)) +end + +# -- varbinary -------------------------------------------------------------- + +"Concrete-width offset load: `idx0` is the 0-based entry index." +@inline function _load_offset(offs::BufferSlice, wide::Bool, idx0::Int64)::Int64 + return wide ? loadat(offs, Int64, checked_mul(idx0, Int64(8))) : + Int64(loadat(offs, Int32, checked_mul(idx0, Int64(4)))) +end + +@inline function _offsets_at(d::ArrayData, i::Int64, wide::Bool) + offs = rolebuffer(d, OFFSETS) + lo = _load_offset(offs, wide, d.offset + i - 1) + hi = _load_offset(offs, wide, d.offset + i) + return lo, hi +end + +function _value(t::Union{Utf8Type,BinaryType}, f::Field, d::ArrayData, i::Int64) + isvalid_at(d, i) || return missing + lo, hi = _offsets_at(d, i, layoutspec(t).offsetwidth == 8) + data = rolebuffer(d, DATA) + n = hi - lo + n == 0 && return t isa Utf8Type ? "" : UInt8[] + # Semantic validation bounded final offsets against the data extent, but + # subslice re-checks: accessors stay safe even when a caller skipped + # validate_semantic (they just pay per-access checking). + bytes = slicebytes(subslice(data, lo, n)) + return t isa Utf8Type ? String(bytes) : bytes +end + +# -- nested ----------------------------------------------------------------- + +# Dynamic recursion mirrors `_typedchild`: scalar children inline into their +# parent's loop, while composite children cross one compiled function barrier. +# Calling public `getvalue` for every scalar child allocates a dispatch box per +# element; recursively inlining every composite instead breaks trim inference. +@inline function _dynamicchild(f::Field, d::ArrayData, i::Int64) + 1 <= i <= d.len || throw(BoundsError(d, i)) + t = d.type + t isa IntType && return _value(t, f, d, i) + t isa FloatType && return _value(t, f, d, i) + t isa BoolType && return _value(t, f, d, i) + t isa Utf8Type && return _value(t, f, d, i) + t isa BinaryType && return _value(t, f, d, i) + t isa FixedSizeBinaryType && return _value(t, f, d, i) + t isa TimestampType && return _value(t, f, d, i) + t isa DateType && return _value(t, f, d, i) + t isa TimeType && return _value(t, f, d, i) + t isa DurationType && return _value(t, f, d, i) + t isa DecimalType && return _value(t, f, d, i) + t isa IntervalType && return _value(t, f, d, i) + t isa ViewType && return _value(t, f, d, i) + t isa NullType && return _value(t, f, d, i) + return _dynamicchildbox(f, d, i) +end + +function _dynamicchildbox(f::Field, d::ArrayData, i::Int64) + return _value_of(d.type, f, d, i) +end + +function _value(t::ListType, f::Field, d::ArrayData, i::Int64) + isvalid_at(d, i) || return missing + lo, hi = _offsets_at(d, i, layoutspec(t).offsetwidth == 8) + child, cf = d.children[1], f.children[1] + # Explicit Vector{Any}: an Any-first comprehension re-narrows its result + # at runtime, which is both trim-hostile and wasted work — typed element + # containers are the facade's job. + out = Vector{Any}(undef, Int(hi - lo)) + for k = 1:Int(hi - lo) + out[k] = _dynamicchild(cf, child, checked_add(lo, Int64(k))) + end + return out +end + +function _value(t::FixedSizeListType, f::Field, d::ArrayData, i::Int64) + isvalid_at(d, i) || return missing + child, cf = d.children[1], f.children[1] + base = checked_mul(_slotindex0(d, i), Int64(t.listsize)) + out = Vector{Any}(undef, t.listsize) + for j = 1:(t.listsize) + out[j] = _dynamicchild(cf, child, checked_add(base, Int64(j))) + end + return out +end + +function _value(::StructType, f::Field, d::ArrayData, i::Int64) + # Core's struct scalar is an ordered Vector{Pair{String,Any}} — always. + # A NamedTuple carries its names in the TYPE domain, so building one from + # runtime schema names is intrinsically dynamic (and cannot represent + # Arrow's duplicate/empty/non-Symbol names at all). The typed NamedTuple + # surface belongs to the facade and to callers' static claims through + # `getvalue(::Type{T}, ...)`; Core stays concrete and trim-clean. + isvalid_at(d, i) || return missing + childindex = checked_add(d.offset, i) + n = length(f.children) + out = Vector{Pair{String,Any}}(undef, n) + for j = 1:n + out[j] = Pair{String,Any}( + f.children[j].name, + _dynamicchild(f.children[j], d.children[j], childindex), + ) + end + return out +end + +function _value(t::MapType, f::Field, d::ArrayData, i::Int64) + # Map = List>; reuse the list walk and pair up. + isvalid_at(d, i) || return missing + lo, hi = _offsets_at(d, i, false) + entries, ef = d.children[1], f.children[1] + kf, vf = ef.children[1], ef.children[2] + kd, vd = entries.children[1], entries.children[2] + out = Vector{Pair{Any,Any}}(undef, Int(hi - lo)) + for k = 1:Int(hi - lo) + entryindex = checked_add(entries.offset, checked_add(lo, Int64(k))) + out[k] = Pair{Any,Any}( + _dynamicchild(kf, kd, entryindex), + _dynamicchild(vf, vd, entryindex), + ) + end + return out +end + +function _value(t::UnionType, f::Field, d::ArrayData, i::Int64) + tid = loadat(rolebuffer(d, TYPE_IDS), Int8, _slotindex0(d, i)) + pos = findfirst(==(tid), t.typeids) + pos === nothing && throw(ValidationError("union type id $tid not in declared domain")) + child, cf = d.children[pos], f.children[pos] + if t.mode == DenseMode + off = loadat(rolebuffer(d, ELEMENT_OFFSETS), Int32, _slotbyteoff(d, i, 4)) + return _dynamicchild(cf, child, checked_add(Int64(off), Int64(1))) + else + return _dynamicchild(cf, child, checked_add(d.offset, i)) + end +end + +function _value(t::DictionaryType, f::Field, d::ArrayData, i::Int64) + isvalid_at(d, i) || return missing + w = primwidth(t.indextype) + idx = _load_int(rolebuffer(d, DATA), t.indextype, _slotbyteoff(d, i, w)) + dict = d.dictionary + dict === nothing && + throw(ValidationError("dictionary-encoded array without a dictionary")) + return _dynamicchild(dictvaluefield(f, t), dict, checked_add(Int64(idx), Int64(1))) +end + +function _value(t::ViewType, f::Field, d::ArrayData, i::Int64) + isvalid_at(d, i) || return missing + views = rolebuffer(d, VIEWS) + base = _viewbase(d, i) + len = loadat(views, Int32, base) + len >= 0 || throw(ValidationError("negative view length $len")) + n = Int64(len) + # Semantic validation certified geometry and prefixes; subslice re-checks + # bounds so unvalidated access still cannot escape a buffer. + bytes = if len <= VIEW_INLINE_MAX + slicebytes(subslice(views, checked_add(base, Int64(4)), n)) + else + bufidx = loadat(views, Int32, checked_add(base, Int64(8))) + off = Int64(loadat(views, Int32, checked_add(base, Int64(12)))) + off >= 0 || throw(ValidationError("negative view offset $off")) + slicebytes(subslice(_viewdatabuffer(d, bufidx), off, n)) + end + return t.utf8 ? String(bytes) : bytes +end + +function _value(t::ListViewType, f::Field, d::ArrayData, i::Int64) + isvalid_at(d, i) || return missing + off, sz = _listview_range(t, d, i) + (off >= 0 && sz >= 0) || + throw(ValidationError("list-view offset and size must be non-negative")) + child, cf = d.children[1], f.children[1] + out = Vector{Any}(undef, Int(sz)) + for k = 1:Int(sz) + out[k] = _dynamicchild(cf, child, checked_add(off, Int64(k))) + end + return out +end + +_value(::RunEndEncodedType, f::Field, d::ArrayData, i::Int64) = + _dynamicchild(f.children[2], d.children[2], _ree_runindex(d, i)) + +# Preflight the containers allocated by one dynamic Core value. This walk uses +# only validated offsets and scalar buffer loads. In particular, a hostile +# FixedSizeList width is charged before the child loop starts. +function _prechargevalue!( + ::Type{T}, + t::Union{IntType,FloatType}, + f::Field, + d::ArrayData, + i::Int64, + estimate, +) where {T} + isvalid_at(d, i) || return nothing + T === Any && _reserveobject!(estimate, primwidth(t), "boxed numeric value") + return nothing +end + +function _prechargevalue!( + ::Type{T}, + ::Union{TimestampType,DurationType}, + f::Field, + d::ArrayData, + i::Int64, + estimate, +) where {T} + isvalid_at(d, i) || return nothing + T === Any && _reserveobject!(estimate, 8, "boxed temporal value") + return nothing +end + +function _prechargevalue!( + ::Type{T}, + t::DateType, + f::Field, + d::ArrayData, + i::Int64, + estimate, +) where {T} + isvalid_at(d, i) || return nothing + T === Any && _reserveobject!(estimate, t.unit == DAY ? 4 : 8, "boxed date value") + return nothing +end + +function _prechargevalue!( + ::Type{T}, + t::TimeType, + f::Field, + d::ArrayData, + i::Int64, + estimate, +) where {T} + isvalid_at(d, i) || return nothing + T === Any && _reserveobject!(estimate, t.bits == 32 ? 4 : 8, "boxed time value") + return nothing +end + +function _prechargevalue!( + ::Type{T}, + t::IntervalType, + f::Field, + d::ArrayData, + i::Int64, + estimate, +) where {T} + isvalid_at(d, i) || return nothing + payload = t.unit == YEAR_MONTH ? 4 : t.unit == DAY_TIME ? 8 : 16 + T === Any && _reserveobject!(estimate, payload, "boxed interval value") + return nothing +end + +_prechargevalue!( + ::Type, + ::Union{BoolType,NullType}, + ::Field, + ::ArrayData, + ::Int64, + estimate, +) = nothing + +function _prechargevalue!( + ::Type{T}, + t::DecimalType, + f::Field, + d::ArrayData, + i::Int64, + estimate, +) where {T} + isvalid_at(d, i) || return nothing + if t.bits > 64 + _reservevector!(estimate, UInt8, primwidth(t), "decimal value") + elseif T === Any + _reserveobject!(estimate, primwidth(t), "boxed decimal value") + end + return nothing +end + +function _prechargevalue!( + ::Type, + t::FixedSizeBinaryType, + f::Field, + d::ArrayData, + i::Int64, + estimate, +) + isvalid_at(d, i) || return nothing + _reservevector!(estimate, UInt8, t.nbytes, "fixed-size binary value") + return nothing +end + +function _prechargevalue!( + ::Type, + t::Union{Utf8Type,BinaryType}, + f::Field, + d::ArrayData, + i::Int64, + estimate, +) + isvalid_at(d, i) || return nothing + lo, hi = _offsets_at(d, i, layoutspec(t).offsetwidth == 8) + # `_value` returns Base's singleton empty String without allocating. Keep + # that exact fast path out of the per-value reserve. Empty Binary values + # still allocate a fresh UInt8 vector and remain charged below. + t isa Utf8Type && hi == lo && return nothing + _reservevector!(estimate, UInt8, hi - lo, "binary value") + return nothing +end + +function _prechargevalue!(::Type, t::ViewType, f::Field, d::ArrayData, i::Int64, estimate) + isvalid_at(d, i) || return nothing + len = loadat(rolebuffer(d, VIEWS), Int32, _viewbase(d, i)) + len >= 0 || throw(ValidationError("negative view length $len")) + _reservevector!(estimate, UInt8, len, "view value") + return nothing +end + +function _prechargevalue!( + ::Type{T}, + t::ListType, + f::Field, + d::ArrayData, + i::Int64, + estimate, +) where {T} + isvalid_at(d, i) || return nothing + lo, hi = _offsets_at(d, i, layoutspec(t).offsetwidth == 8) + _prechargelist!(T, f, d, lo, hi - lo, estimate) + return nothing +end + +function _prechargevalue!( + ::Type{T}, + t::ListViewType, + f::Field, + d::ArrayData, + i::Int64, + estimate, +) where {T} + isvalid_at(d, i) || return nothing + off, n = _listview_range(t, d, i) + (off >= 0 && n >= 0) || + throw(ValidationError("list-view offset and size must be non-negative")) + _prechargelist!(T, f, d, off, n, estimate) + return nothing +end + +function _prechargelist!( + ::Type{T}, + f::Field, + d::ArrayData, + off::Int64, + n::Int64, + estimate, +) where {T} + child, cf = d.children[1], f.children[1] + if T === Any + _reservevector!(estimate, Any, n, "list value") + for k = 1:Int(n) + _prechargechild!(Any, cf, child, checked_add(off, Int64(k)), estimate) + end + else + CE = eltype(Base.nonmissingtype(T)) + _reservevector!(estimate, CE, n, "typed list value") + for k = 1:Int(n) + _prechargechild!(CE, cf, child, checked_add(off, Int64(k)), estimate) + end + end + return nothing +end + +function _prechargevalue!( + ::Type{T}, + t::FixedSizeListType, + f::Field, + d::ArrayData, + i::Int64, + estimate, +) where {T} + isvalid_at(d, i) || return nothing + child, cf = d.children[1], f.children[1] + base = checked_mul(_slotindex0(d, i), Int64(t.listsize)) + if T === Any + CE = Any + _reservevector!(estimate, CE, t.listsize, "fixed-size list value") + else + CE = eltype(Base.nonmissingtype(T)) + _reservevector!(estimate, CE, t.listsize, "typed fixed-size list value") + end + for k = 1:(t.listsize) + _prechargechild!(CE, cf, child, checked_add(base, Int64(k)), estimate) + end + return nothing +end + +function _prechargevalue!( + ::Type{T}, + ::StructType, + f::Field, + d::ArrayData, + i::Int64, + estimate, +) where {T} + isvalid_at(d, i) || return nothing + E = T === Any ? Any : Base.nonmissingtype(T) + childindex = checked_add(d.offset, i) + if E === Any || E === Vector{Pair{String,Any}} + _reservevector!(estimate, Pair{String,Any}, length(f.children), "struct value") + for k in eachindex(f.children) + _prechargechild!(Any, f.children[k], d.children[k], childindex, estimate) + end + return nothing + end + E <: NamedTuple || return nothing + Base.allocatedinline(T) || + _reserveobject!(estimate, sizeof(E), "boxed typed struct value") + _prechargetypedstruct!(E, f, d, childindex, estimate) + return nothing +end + +@generated function _prechargetypedstruct!( + ::Type{E}, + f::Field, + d::ArrayData, + childindex::Int64, + estimate, +) where {E<:NamedTuple} + calls = Expr[ + :(_prechargechild!( + $(fieldtype(E, j)), + f.children[$j], + d.children[$j], + childindex, + estimate, + )) for j = 1:fieldcount(E) + ] + return quote + $(calls...) + return nothing + end +end + +function _prechargevalue!(::Type, t::MapType, f::Field, d::ArrayData, i::Int64, estimate) + isvalid_at(d, i) || return nothing + lo, hi = _offsets_at(d, i, false) + _reservevector!(estimate, Pair{Any,Any}, hi - lo, "map value") + entries, ef = d.children[1], f.children[1] + for k = 1:Int(hi - lo) + entryindex = checked_add(entries.offset, checked_add(lo, Int64(k))) + for j = 1:2 + _prechargechild!(Any, ef.children[j], entries.children[j], entryindex, estimate) + end + end + return nothing +end + +function _prechargevalue!(::Type, t::UnionType, f::Field, d::ArrayData, i::Int64, estimate) + cf, child, childindex = _union_child(f, d, i) + _prechargechild!(Any, cf, child, childindex, estimate) + return nothing +end + +function _prechargevalue!( + ::Type{T}, + t::DictionaryType, + f::Field, + d::ArrayData, + i::Int64, + estimate, +) where {T} + isvalid_at(d, i) || return nothing + w = primwidth(t.indextype) + index = _load_int(rolebuffer(d, DATA), t.indextype, _slotbyteoff(d, i, w)) + dictionary = d.dictionary + dictionary === nothing && + throw(ValidationError("dictionary-encoded array without a dictionary")) + _prechargechild!( + T, + dictvaluefield(f, t), + dictionary, + checked_add(Int64(index), Int64(1)), + estimate, + ) + return nothing +end + +function _prechargevalue!( + ::Type{T}, + ::RunEndEncodedType, + f::Field, + d::ArrayData, + i::Int64, + estimate, +) where {T} + _prechargechild!(T, f.children[2], d.children[2], _ree_runindex(d, i), estimate) + return nothing +end + +_prechargevalue!(::Type, ::ArrowType, ::Field, ::ArrayData, ::Int64, estimate) = nothing + +# Like the value path, preflight keeps leaves inline and gives recursive +# composites one concrete compiled edge. The walk itself must not create an +# uncharged dispatch allocation for every child it inspects. +@inline function _prechargechild!( + ::Type{T}, + f::Field, + d::ArrayData, + i::Int64, + estimate, +) where {T} + 1 <= i <= d.len || throw(BoundsError(d, i)) + t = d.type + t isa IntType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa FloatType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa BoolType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa Utf8Type && return _prechargevalue!(T, t, f, d, i, estimate) + t isa BinaryType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa FixedSizeBinaryType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa TimestampType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa DateType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa TimeType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa DurationType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa DecimalType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa IntervalType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa ViewType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa NullType && return _prechargevalue!(T, t, f, d, i, estimate) + return _prechargechildbox!(T, f, d, i, estimate) +end + +function _prechargechildbox!( + ::Type{T}, + f::Field, + d::ArrayData, + i::Int64, + estimate, +) where {T} + return _prechargevalue_of!(T, d.type, f, d, i, estimate) +end + +@inline function _prechargevalue_of!( + ::Type{T}, + t::ArrowType, + f::Field, + d::ArrayData, + i::Int64, + estimate, +) where {T} + t isa IntType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa FloatType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa Utf8Type && return _prechargevalue!(T, t, f, d, i, estimate) + t isa BoolType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa ListType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa StructType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa DictionaryType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa TimestampType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa DateType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa TimeType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa DurationType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa BinaryType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa FixedSizeBinaryType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa FixedSizeListType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa MapType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa UnionType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa DecimalType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa IntervalType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa NullType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa ViewType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa ListViewType && return _prechargevalue!(T, t, f, d, i, estimate) + t isa RunEndEncodedType && return _prechargevalue!(T, t, f, d, i, estimate) + return _prechargevalue!(T, t, f, d, i, estimate) +end + +function getvalue(f::Field, d::ArrayData, i::Integer, budget) + 1 <= i <= d.len || throw(BoundsError(d, i)) + estimate = _MaterializationEstimate(budget) + _prechargevalue_of!(Any, d.type, f, d, Int64(i), estimate) + _commitestimate!(estimate, "materialized value") + return _value_of(d.type, f, d, Int64(i)) +end + +""" + materialize(field, data) -> Vector + +Bulk conversion to native Julia values: resolve the layout ONCE, then run +a specialized loop behind a function barrier — the loop body compiles per +LAYOUT (a small closed set), never per schema. +""" +materialize(f::Field, d::ArrayData) = _materialize_of(d.type, f, d) + +function materialize(f::Field, d::ArrayData, budget) + estimate = _MaterializationEstimate(budget) + _reservevector!(estimate, Any, d.len, "materialized column") + for i = 1:(d.len) + _prechargevalue_of!(Any, d.type, f, d, Int64(i), estimate) + end + _commitestimate!(estimate, "materialized column") + return materialize(f, d) +end + +# The same closed-set ladder as `layoutspec_of`, for element access and the +# materialize function barrier: generic entry devirtualizes here; per-layout +# `_value` methods stay the extension surface. +@inline function _value_of(t::ArrowType, f::Field, d::ArrayData, i::Int64) + t isa IntType && return _value(t, f, d, i) + t isa FloatType && return _value(t, f, d, i) + t isa Utf8Type && return _value(t, f, d, i) + t isa BoolType && return _value(t, f, d, i) + t isa ListType && return _value(t, f, d, i) + t isa StructType && return _value(t, f, d, i) + t isa DictionaryType && return _value(t, f, d, i) + t isa TimestampType && return _value(t, f, d, i) + t isa DateType && return _value(t, f, d, i) + t isa TimeType && return _value(t, f, d, i) + t isa DurationType && return _value(t, f, d, i) + t isa BinaryType && return _value(t, f, d, i) + t isa FixedSizeBinaryType && return _value(t, f, d, i) + t isa FixedSizeListType && return _value(t, f, d, i) + t isa MapType && return _value(t, f, d, i) + t isa UnionType && return _value(t, f, d, i) + t isa DecimalType && return _value(t, f, d, i) + t isa IntervalType && return _value(t, f, d, i) + t isa NullType && return _value(t, f, d, i) + t isa ViewType && return _value(t, f, d, i) + t isa ListViewType && return _value(t, f, d, i) + t isa RunEndEncodedType && return _value(t, f, d, i) + throw(ArgumentError("unregistered ArrowType")) +end + +@inline function _materialize_of(t::ArrowType, f::Field, d::ArrayData) + t isa IntType && return _materialize_loop(t, f, d) + t isa FloatType && return _materialize_loop(t, f, d) + t isa Utf8Type && return _materialize_loop(t, f, d) + t isa BoolType && return _materialize_loop(t, f, d) + t isa ListType && return _materialize_loop(t, f, d) + t isa StructType && return _materialize_loop(t, f, d) + t isa DictionaryType && return _materialize_loop(t, f, d) + t isa TimestampType && return _materialize_loop(t, f, d) + t isa DateType && return _materialize_loop(t, f, d) + t isa TimeType && return _materialize_loop(t, f, d) + t isa DurationType && return _materialize_loop(t, f, d) + t isa BinaryType && return _materialize_loop(t, f, d) + t isa FixedSizeBinaryType && return _materialize_loop(t, f, d) + t isa FixedSizeListType && return _materialize_loop(t, f, d) + t isa MapType && return _materialize_loop(t, f, d) + t isa UnionType && return _materialize_loop(t, f, d) + t isa DecimalType && return _materialize_loop(t, f, d) + t isa IntervalType && return _materialize_loop(t, f, d) + t isa NullType && return _materialize_loop(t, f, d) + t isa ViewType && return _materialize_loop(t, f, d) + t isa ListViewType && return _materialize_loop(t, f, d) + t isa RunEndEncodedType && return _materialize_loop(t, f, d) + throw(ArgumentError("unregistered ArrowType")) +end + +function _materialize_loop(t::T, f::Field, d::ArrayData) where {T<:ArrowType} + out = Vector{Any}(undef, d.len) + for i = 1:(d.len) + out[i] = _value(t, f, d, Int64(i)) + end + # Vector{Any} by design: result-element typing is the caller's claim + # through `materialize(::Type{T}, ...)` (which the facade uses for closed + # claims), and a runtime narrow is trim-hostile. Tests compare with + # ==/isequal, which is eltype-agnostic. + return out +end + +# --------------------------------------------------------------------------- +# Typed element access: the caller asserts the element domain. With a +# concrete static schema at the call site every load +# resolves statically — the trim-compile contract dynamic access cannot +# offer. The type is a CLAIM about the same value domain the dynamic +# accessors return (storage integers for temporal, `Vector{Pair}` rows for +# struct/map): the read produces exactly that type or refuses with a clear +# error. `Any` is the dynamic path unchanged. +# +# Recursion architecture: it mirrors the dynamic path exactly. Recursive +# edges route through `_typedchild` — a COMPILED function whose argument +# types are all concrete (like public `getvalue` on the dynamic side) — so +# the cycle's one non-inlined call is fully resolvable; the `@inline` +# ladder and leaf methods flatten into it. An `@inline` ladder call +# carrying an abstract descriptor as the recursive edge is unresolvable +# under trim (no standalone specialization exists). +# --------------------------------------------------------------------------- + +""" + getvalue(::Type{T}, field, data, i) -> T + +Statically typed element access: `T` asserts the element domain (what the +dynamic accessors return for this layout — see `juliatype` for the leaf +layouts), with `Missing <: T` required to admit nulls. The claim checks against the +DESCRIPTOR up front — an empty or all-null column certifies nothing. +Composites recurse: a `List` column reads as +`Vector{Vector{Int64}}`, and a `Struct` column may read as a `NamedTuple` +row type whose names match the child fields in order. Mismatches refuse +with `ArgumentError` — values are never converted. `T === Any` delegates +to the dynamic path. +""" +function getvalue(::Type{T}, f::Field, d::ArrayData, i::Integer) where {T} + T === Any && return getvalue(f, d, i) + _checkclaim(T, f, d) + 1 <= i <= d.len || throw(BoundsError(d, i)) + return _typedvalue_of(T, d.type, f, d, Int64(i))::T +end + +""" + materialize(::Type{T}, field, data) -> Vector{T} + +The bulk form of the typed [`getvalue`](@ref): every element under the +static claim `T` (see `getvalue(::Type{T}, field, data, i)` for the claim +rules), through one typed loop per layout. +""" +function materialize(::Type{T}, f::Field, d::ArrayData) where {T} + T === Any && return materialize(f, d) + _checkclaim(T, f, d) + return _typedmaterialize_of(T, d.type, f, d)::Vector{T} +end + +function materialize(::Type{T}, f::Field, d::ArrayData, budget) where {T} + T === Any && return materialize(f, d, budget) + _checkclaim(T, f, d) + estimate = _MaterializationEstimate(budget) + _reservevector!(estimate, T, d.len, "typed materialized column") + for i = 1:(d.len) + _prechargevalue_of!(T, d.type, f, d, Int64(i), estimate) + end + # Nullable fixed-width bulk extraction first builds Vector{E}, then its + # public Vector{Union{Missing,E}}. Account for that private scratch copy. + E = Base.nonmissingtype(T) + if Missing <: T && + isbitstype(E) && + d.type isa + Union{IntType,FloatType,TimestampType,DateType,TimeType,DurationType,DecimalType} && + E === juliatype(d.type) && + primwidth(d.type) == Int64(sizeof(E)) + _reservevector!(estimate, E, d.len, "typed materialization scratch") + end + _commitestimate!(estimate, "typed materialized column") + return _typedmaterialize_of(T, d.type, f, d)::Vector{T} +end + +# The message names the layout via `nameof` (generic struct/type `show` is +# trim-hostile, and an abstract descriptor argument would leave the throw +# helper unresolvable); `juliatype(t)` tells a caller the expected claim. +# Closed-set ladder to literal strings: `nameof(typeof(t))` on an abstract +# descriptor is itself an unresolvable call under trim. +@inline function _layoutname(t::ArrowType) + t isa IntType && return "IntType" + t isa FloatType && return "FloatType" + t isa BoolType && return "BoolType" + t isa Utf8Type && return "Utf8Type" + t isa BinaryType && return "BinaryType" + t isa FixedSizeBinaryType && return "FixedSizeBinaryType" + t isa TimestampType && return "TimestampType" + t isa DateType && return "DateType" + t isa TimeType && return "TimeType" + t isa DurationType && return "DurationType" + t isa ViewType && return "ViewType" + t isa DecimalType && return "DecimalType" + t isa IntervalType && return "IntervalType" + t isa MapType && return "MapType" + t isa StructType && return "StructType" + t isa ListType && return "ListType" + t isa ListViewType && return "ListViewType" + t isa FixedSizeListType && return "FixedSizeListType" + t isa DictionaryType && return "DictionaryType" + t isa RunEndEncodedType && return "RunEndEncodedType" + t isa UnionType && return "UnionType" + t isa NullType && return "NullType" + return "ArrowType" +end +@noinline _typedrefuse(::Type{E}, kind::String, f::Field) where {E} = throw( + ArgumentError( + "field $(f.name) materializes $(kind)-layout " * + "values; the claimed static element type does not match", + ), +) +@noinline _typednullrefuse(f::Field) = throw( + ArgumentError( + "field $(f.name) holds a null but the static " * + "element type does not admit missing", + ), +) +@inline _typedmissing(::Type{T}, f::Field) where {T} = + Missing <: T ? missing : _typednullrefuse(f) + +""" +Descriptor-level claim preflight: `T` must match the element domain the +schema DECLARES — acceptance never depends on which values a batch +happens to contain (an empty or all-null column certifies nothing). +Shapes, field counts, and names check ONCE here; the element loop stays +check-free. Compiled (not `@inline`): its recursion keeps the claim +intact through transparent wrappers, and a compiled concrete-arg edge is +what makes that cycle trim-resolvable. +""" +function _checkclaim(::Type{T}, f::Field, d::ArrayData)::Nothing where {T} + t = d.type + if t isa ListType || t isa ListViewType || t isa FixedSizeListType + E = Base.nonmissingtype(T) + E <: Vector || _typedrefuse(E, _layoutname(t), f) + length(f.children) == 1 && length(d.children) == 1 || + _typedrefuse(E, _layoutname(t), f) + return _checkclaim(eltype(E), f.children[1], d.children[1]) + end + if t isa StructType + E = Base.nonmissingtype(T) + E === Vector{Pair{String,Any}} && return nothing + # An EXACT NamedTuple shape only: a Union or UnionAll of row types + # satisfies `<: NamedTuple` but has no field reflection — it must + # refuse here, not leak a generation error. + (E isa DataType && E <: NamedTuple) || _typedrefuse(E, _layoutname(t), f) + (fieldcount(E) == length(f.children) && fieldcount(E) == length(d.children)) || + _typedrefuse(E, _layoutname(t), f) + return _checkstructclaim(E, f, d) + end + if t isa DictionaryType + dict = d.dictionary + dict === nothing && + throw(ValidationError("dictionary-encoded array without a dictionary")) + return _checkclaim(T, dictvaluefield(f, t), dict) + end + if t isa RunEndEncodedType + (length(f.children) == 2 && length(d.children) == 2) || + _typedrefuse(Base.nonmissingtype(T), _layoutname(t), f) + return _checkclaim(T, f.children[2], d.children[2]) + end + t isa UnionType && _typedrefuse(Base.nonmissingtype(T), _layoutname(t), f) + if t isa NullType + Missing <: T || _typednullrefuse(f) + return nothing + end + E = Base.nonmissingtype(T) + E === _juliatype_of(t) || _typedrefuse(E, _layoutname(t), f) + return nothing +end + +# Generated so every field index is a LITERAL: `fieldtype(E, j)` with a +# runtime `j` yields an abstract `Type` and poisons the recursion, and the +# name strings bake in at generation (no per-call conversion at all). +@generated function _checkstructclaim( + ::Type{E}, + f::Field, + d::ArrayData, +)::Nothing where {E<:NamedTuple} + checks = Expr[] + for j = 1:fieldcount(E) + push!( + checks, + :( + $(String(fieldnames(E)[j])) == f.children[$j].name || + _typedrefuse(E, _layoutname(d.type), f) + ), + ) + push!(checks, :(_checkclaim($(fieldtype(E, j)), f.children[$j], d.children[$j]))) + end + return quote + $(checks...) + return nothing + end +end + +# Closed-set ladder for the preflight's leaf claims — scalars plus the +# `Map` row vector (an abstract +# `juliatype(t::ArrowType)` call would defeat trim resolution). +@inline function _juliatype_of(t::ArrowType) + t isa IntType && return juliatype(t) + t isa FloatType && return juliatype(t) + t isa BoolType && return juliatype(t) + t isa Utf8Type && return juliatype(t) + t isa BinaryType && return juliatype(t) + t isa FixedSizeBinaryType && return juliatype(t) + t isa TimestampType && return juliatype(t) + t isa DateType && return juliatype(t) + t isa TimeType && return juliatype(t) + t isa DurationType && return juliatype(t) + t isa ViewType && return juliatype(t) + t isa DecimalType && return juliatype(t) + t isa IntervalType && return juliatype(t) + t isa MapType && return juliatype(t) + throw(ArgumentError("unregistered ArrowType")) +end + +# The typed recursion edge, split for two masters. The logical-bounds guard +# matches what the dynamic path gets from public `getvalue`: unvalidated +# geometry must not read backing values past a child's logical length. +# SCALAR leaves inline into the parent's loop so the buffer-slice +# temporaries stay stack-allocated (a compiled boundary per child read +# would allocate); COMPOSITE children route to `_typedchildbox`, a compiled +# shell whose argument types are all concrete — the resolvable edge trim +# requires. The generic ladder flattens reliably into a dedicated shell but +# not into arbitrary hoisted contexts, so the shell is the only caller. The +# `::T` asserts pin inference to the claim where the same-claim wrapper +# cycle (Dictionary/REE) would widen to Any. +@inline function _typedchild(::Type{T}, f::Field, d::ArrayData, i::Int64) where {T} + 1 <= i <= d.len || throw(BoundsError(d, i)) + t = d.type + t isa IntType && return _typedvalue(T, t, f, d, i)::T + t isa FloatType && return _typedvalue(T, t, f, d, i)::T + t isa BoolType && return _typedvalue(T, t, f, d, i)::T + t isa Utf8Type && return _typedvalue(T, t, f, d, i)::T + t isa BinaryType && return _typedvalue(T, t, f, d, i)::T + t isa FixedSizeBinaryType && return _typedvalue(T, t, f, d, i)::T + t isa TimestampType && return _typedvalue(T, t, f, d, i)::T + t isa DateType && return _typedvalue(T, t, f, d, i)::T + t isa TimeType && return _typedvalue(T, t, f, d, i)::T + t isa DurationType && return _typedvalue(T, t, f, d, i)::T + t isa DecimalType && return _typedvalue(T, t, f, d, i)::T + t isa IntervalType && return _typedvalue(T, t, f, d, i)::T + t isa ViewType && return _typedvalue(T, t, f, d, i)::T + t isa NullType && return _typedvalue(T, t, f, d, i)::T + return _typedchildbox(T, f, d, i)::T +end + +function _typedchildbox(::Type{T}, f::Field, d::ArrayData, i::Int64) where {T} + return _typedvalue_of(T, d.type, f, d, i)::T +end + +# The same closed-set ladder as `_value_of`, with the claimed type threaded. +@inline function _typedvalue_of( + ::Type{T}, + t::ArrowType, + f::Field, + d::ArrayData, + i::Int64, +) where {T} + t isa IntType && return _typedvalue(T, t, f, d, i) + t isa FloatType && return _typedvalue(T, t, f, d, i) + t isa Utf8Type && return _typedvalue(T, t, f, d, i) + t isa BoolType && return _typedvalue(T, t, f, d, i) + t isa ListType && return _typedvalue(T, t, f, d, i) + t isa StructType && return _typedvalue(T, t, f, d, i) + # Wrapper branches keep the claim intact, so they alone can recurse + # with an UNCHANGED signature: the ::T assert stops that cycle from + # widening every other branch to Any in fresh-process inference; the + # wrapper read itself pays one box. + t isa DictionaryType && return _typedvalue(T, t, f, d, i)::T + t isa TimestampType && return _typedvalue(T, t, f, d, i) + t isa DateType && return _typedvalue(T, t, f, d, i) + t isa TimeType && return _typedvalue(T, t, f, d, i) + t isa DurationType && return _typedvalue(T, t, f, d, i) + t isa BinaryType && return _typedvalue(T, t, f, d, i) + t isa FixedSizeBinaryType && return _typedvalue(T, t, f, d, i) + t isa FixedSizeListType && return _typedvalue(T, t, f, d, i) + t isa MapType && return _typedvalue(T, t, f, d, i) + t isa UnionType && return _typedvalue(T, t, f, d, i) + t isa DecimalType && return _typedvalue(T, t, f, d, i) + t isa IntervalType && return _typedvalue(T, t, f, d, i) + t isa NullType && return _typedvalue(T, t, f, d, i) + t isa ViewType && return _typedvalue(T, t, f, d, i) + t isa ListViewType && return _typedvalue(T, t, f, d, i) + t isa RunEndEncodedType && return _typedvalue(T, t, f, d, i)::T + throw(ArgumentError("unregistered ArrowType")) +end + +# Closed leaf claims (scalars plus the `Map` row vector): the claim must +# equal the layout's `juliatype` +# exactly; the audited dynamic extraction runs and the assert makes the +# result statically typed (and free when the claim is right). +function _typedvalue( + ::Type{T}, + t::Union{ + IntType, + FloatType, + BoolType, + Utf8Type, + BinaryType, + FixedSizeBinaryType, + TimestampType, + DateType, + TimeType, + DurationType, + ViewType, + DecimalType, + IntervalType, + MapType, + }, + f::Field, + d::ArrayData, + i::Int64, +) where {T} + isvalid_at(d, i) || return _typedmissing(T, f) + E = Base.nonmissingtype(T) + E === juliatype(t) || _typedrefuse(E, _layoutname(t), f) + # Invalid int/float widths fall through juliatype's 64-bit fallback: + # refuse them here (managed, fails closed) rather than let the raw + # extraction's own width ladder produce a mistyped value. + (t isa IntType || t isa FloatType) && + primwidth(t) != Int64(sizeof(E)) && + _typedrefuse(E, _layoutname(t), f) + return _value(t, f, d, i)::E +end + +function _typedvalue( + ::Type{T}, + t::Union{ListType,ListViewType}, + f::Field, + d::ArrayData, + i::Int64, +) where {T} + isvalid_at(d, i) || return _typedmissing(T, f) + E = Base.nonmissingtype(T) + E <: Vector || _typedrefuse(E, _layoutname(t), f) + if t isa ListType + lo, hi = _offsets_at(d, i, layoutspec(t).offsetwidth == 8) + off = lo + n = hi - lo + else + off, n = _listview_range(t, d, i) + (off >= 0 && n >= 0) || + throw(ValidationError("list-view offset and size must be non-negative")) + end + child, cf = d.children[1], f.children[1] + CE = eltype(E) + out = Vector{CE}(undef, Int(n)) + for k = 1:Int(n) + out[k] = _typedchild(CE, cf, child, checked_add(off, Int64(k))) + end + return out +end + +function _typedvalue( + ::Type{T}, + t::FixedSizeListType, + f::Field, + d::ArrayData, + i::Int64, +) where {T} + isvalid_at(d, i) || return _typedmissing(T, f) + E = Base.nonmissingtype(T) + E <: Vector || _typedrefuse(E, _layoutname(t), f) + child, cf = d.children[1], f.children[1] + base = checked_mul(_slotindex0(d, i), Int64(t.listsize)) + CE = eltype(E) + out = Vector{CE}(undef, t.listsize) + for j = 1:(t.listsize) + out[j] = _typedchild(CE, cf, child, checked_add(base, Int64(j))) + end + return out +end + +function _typedvalue(::Type{T}, t::StructType, f::Field, d::ArrayData, i::Int64) where {T} + isvalid_at(d, i) || return _typedmissing(T, f) + E = Base.nonmissingtype(T) + E === Vector{Pair{String,Any}} && return _value(t, f, d, i)::E + E <: NamedTuple || _typedrefuse(E, _layoutname(t), f) + return _structrow(E, f, d, checked_add(d.offset, i)) +end + +# Generated so every field's claim is a LITERAL type and the row build is +# a flat tuple expression: an `ntuple(Val(N))` closure erases per-field +# types to `NTuple{N,Any}` at arity >= 4, and index-recursion trips the +# inference recursion limiter. The preflight already checked names. +@generated function _structrow( + ::Type{E}, + f::Field, + d::ArrayData, + childindex::Int64, +) where {E<:NamedTuple} + vals = Expr[ + :(_typedchild($(fieldtype(E, j)), f.children[$j], d.children[$j], childindex)) + for j = 1:fieldcount(E) + ] + return :(E(($(vals...),))) +end + +function _typedvalue( + ::Type{T}, + t::DictionaryType, + f::Field, + d::ArrayData, + i::Int64, +) where {T} + isvalid_at(d, i) || return _typedmissing(T, f) + w = primwidth(t.indextype) + idx = _load_int(rolebuffer(d, DATA), t.indextype, _slotbyteoff(d, i, w)) + dict = d.dictionary + dict === nothing && + throw(ValidationError("dictionary-encoded array without a dictionary")) + return _typedchild(T, dictvaluefield(f, t), dict, checked_add(Int64(idx), Int64(1))) +end + +_typedvalue(::Type{T}, t::RunEndEncodedType, f::Field, d::ArrayData, i::Int64) where {T} = + _typedchild(T, f.children[2], d.children[2], _ree_runindex(d, i)) + +_typedvalue(::Type{T}, ::NullType, f::Field, ::ArrayData, ::Int64) where {T} = + _typedmissing(T, f) + +# Union rows take the WINNING child's runtime type: no static claim can +# hold across children, so only the dynamic path reads unions. +_typedvalue(::Type{T}, t::UnionType, f::Field, ::ArrayData, ::Int64) where {T} = + _typedrefuse(Base.nonmissingtype(T), _layoutname(t), f) + +@inline function _typedmaterialize_of( + ::Type{T}, + t::ArrowType, + f::Field, + d::ArrayData, +) where {T} + t isa IntType && return _typedmaterialize_loop(T, t, f, d) + t isa FloatType && return _typedmaterialize_loop(T, t, f, d) + t isa Utf8Type && return _typedmaterialize_loop(T, t, f, d) + t isa BoolType && return _typedmaterialize_loop(T, t, f, d) + t isa ListType && return _typedmaterialize_loop(T, t, f, d) + t isa StructType && return _typedmaterialize_loop(T, t, f, d) + t isa DictionaryType && return _typedmaterialize_loop(T, t, f, d) + t isa TimestampType && return _typedmaterialize_loop(T, t, f, d) + t isa DateType && return _typedmaterialize_loop(T, t, f, d) + t isa TimeType && return _typedmaterialize_loop(T, t, f, d) + t isa DurationType && return _typedmaterialize_loop(T, t, f, d) + t isa BinaryType && return _typedmaterialize_loop(T, t, f, d) + t isa FixedSizeBinaryType && return _typedmaterialize_loop(T, t, f, d) + t isa FixedSizeListType && return _typedmaterialize_loop(T, t, f, d) + t isa MapType && return _typedmaterialize_loop(T, t, f, d) + t isa UnionType && return _typedmaterialize_loop(T, t, f, d) + t isa DecimalType && return _typedmaterialize_loop(T, t, f, d) + t isa IntervalType && return _typedmaterialize_loop(T, t, f, d) + t isa NullType && return _typedmaterialize_loop(T, t, f, d) + t isa ViewType && return _typedmaterialize_loop(T, t, f, d) + t isa ListViewType && return _typedmaterialize_loop(T, t, f, d) + t isa RunEndEncodedType && return _typedmaterialize_loop(T, t, f, d) + throw(ArgumentError("unregistered ArrowType")) +end + +function _typedmaterialize_loop( + ::Type{T}, + t::TT, + f::Field, + d::ArrayData, +) where {T,TT<:ArrowType} + bulk = _bulkmaterialize(T, t, f, d) + bulk === nothing || return bulk::Vector{T} + out = Vector{T}(undef, d.len) + for i = 1:(d.len) + out[i] = _typedvalue(T, t, f, d, Int64(i)) + end + return out +end + +# --------------------------------------------------------------------------- +# Bulk fixed-width extraction: for closed isbits claims over plain +# fixed-width layouts, one bounds-checked byte copy replaces ten million +# per-element calls (the benchmark-dominant cost of materializing reads). +# Nulls punch in afterward from the validity bitmap. Everything else +# (strings, composites, bitmaps, decimal-as-bytes) keeps the element loop. +# --------------------------------------------------------------------------- + +_bulkmaterialize(::Type{T}, ::ArrowType, ::Field, ::ArrayData) where {T} = nothing + +function _bulkmaterialize( + ::Type{T}, + t::Union{IntType,FloatType,TimestampType,DateType,TimeType,DurationType,DecimalType}, + f::Field, + d::ArrayData, +) where {T} + E = Base.nonmissingtype(T) + isbitstype(E) || return nothing + E === juliatype(t) || return nothing + w = Int64(sizeof(E)) + # The claim's byte size must equal the DESCRIPTOR's layout width: an + # invalid 24-bit descriptor falls through juliatype's fallback to a + # 64-bit Julia type, and copying at the claim's width would misread — + # such descriptors take the element loop (and validation refuses them). + primwidth(t) == w || return nothing + n = d.len + # The typed path serves unvalidated data too: subslice re-checks the + # extraction window against the buffer's declared bounds. + src = subslice(rolebuffer(d, DATA), checked_mul(d.offset, w), checked_mul(n, w)) + vals = Vector{E}(undef, n) + if n > 0 + GC.@preserve vals d begin + unsafe_copyto!(Ptr{UInt8}(pointer(vals)), sliceptr(src), Int(src.len)) + end + end + # The BITMAP is the validity authority, exactly as per-element access: + # a caller-supplied null-count cache is only certified after semantic + # validation, and this path explicitly serves unvalidated data. + v = validitybuffer(d) + # When the claim admits no Missing, E === T, so `vals` already has the + # public element type and needs no copy. + if !(Missing <: T) + if !isempty_buffer(v) + for i = 1:n + isvalid_at(d, Int64(i)) || _typednullrefuse(f) + end + end + return vals + end + out = Vector{T}(undef, n) + copyto!(out, vals) + isempty_buffer(v) && return out + for i = 1:n + isvalid_at(d, Int64(i)) || (out[i] = missing) + end + return out +end + +# --------------------------------------------------------------------------- +# §7 Builders: Julia data -> (Field, ArrayData) +# --------------------------------------------------------------------------- + +# The write-side counterpart, kept intentionally small: enough construction +# machinery to build every implemented layout without an IPC file in the +# loop. These are "zero-copy wrap + bitmap build" fast paths; the +# append-oriented builder layer is the facade's. + +@inline function _setbitmapbit!(bytes::Vector{UInt8}, i::Int) + bytes[1 + ((i - 1) >> 3)] |= UInt8(1) << ((i - 1) & 7) + return nothing +end + +function _bitmapbuffer(present::AbstractVector{Bool}) + any(!, present) || return BufferSlice() # no nulls -> canonical empty + bytes = zeros(UInt8, expected_validity_bytes(Int64(length(present)))) + for (i, p) in enumerate(present) + p && _setbitmapbit!(bytes, i) + end + return BufferSlice(heapregion(bytes), 0, length(bytes)) +end + +_databuffer(v::Vector{T}) where {T} = BufferSlice(heapregion(v), 0, sizeof(v)) + +"Build Boolean data and validity buffers directly in Arrow's bit-packed form." +function _build_bool(name, v::Vector{T}; nullable::Bool) where {T<:Union{Missing,Bool}} + bytes = zeros(UInt8, expected_validity_bytes(Int64(length(v)))) + nc = 0 + for (i, x) in enumerate(v) + if x === missing + nc += 1 + elseif x + _setbitmapbit!(bytes, i) + end + end + + validity = if nc == 0 + BufferSlice() + else + validbytes = zeros(UInt8, length(bytes)) + for (i, x) in enumerate(v) + x === missing || _setbitmapbit!(validbytes, i) + end + _databuffer(validbytes) + end + t = BoolType() + return Field(name, t; nullable=nullable), + ArrayData(t, length(v), [validity, _databuffer(bytes)]; nullcount=nc) +end + +"Build non-null Boolean data through a caller-supplied direct bit fill." +function _build_bool(name, len::Int, fillbits!::F; nullable::Bool) where {F} + len >= 0 || throw(ArgumentError("Boolean array length is negative")) + bytes = zeros(UInt8, expected_validity_bytes(Int64(len))) + fillbits!(bytes) + t = BoolType() + return Field(name, t; nullable=nullable), + ArrayData(t, len, [BufferSlice(), _databuffer(bytes)]; nullcount=0) +end + +arrowtype_for(::Type{Bool}) = BoolType() +arrowtype_for(::Type{T}) where {T<:Signed} = IntType(8 * sizeof(T), true) +arrowtype_for(::Type{T}) where {T<:Unsigned} = IntType(8 * sizeof(T), false) +arrowtype_for(::Type{Float16}) = FloatType(16) +arrowtype_for(::Type{Float32}) = FloatType(32) +arrowtype_for(::Type{Float64}) = FloatType(64) + +""" + fromjulia(name, v) -> (Field, ArrayData) + +Adapt a Julia vector to Core form. `Vector{T}` for fixed-width isbits `T` +other than `Bool` is a ZERO-COPY wrap (the vector becomes the region's root; +scoped-borrow contract: don't resize/mutate while in use). `Bool` +(bit-packed), `Union{T,Missing}`, and String inputs build fresh buffers. +""" +function fromjulia(name, v::Vector{T}) where {T} + if T === Union{} + throw( + ArgumentError( + "fromjulia: bottom element type Union{} has no Arrow type; " * + "give the empty vector a declared element type", + ), + ) + elseif T <: + Union{Int8,Int16,Int32,Int64,UInt8,UInt16,UInt32,UInt64,Float16,Float32,Float64} + t = arrowtype_for(T) + return Field(name, t; nullable=false), + ArrayData(t, length(v), [BufferSlice(), _databuffer(v)]; nullcount=0) + elseif T == Bool + return _build_bool(name, v; nullable=false) + elseif T == String + return _build_strings(name, v) + elseif T <: Union{ + Missing, + Int8, + Int16, + Int32, + Int64, + UInt8, + UInt16, + UInt32, + UInt64, + Float16, + Float32, + Float64, + Bool, + } + return _build_nullable_primitive(name, v) + elseif T <: Union{Missing,String} + return _build_strings(name, v) + elseif T <: AbstractVector || T <: Union{Missing,<:AbstractVector} + return _build_list(name, v) + else + throw(ArgumentError("fromjulia: unsupported element type $T")) + end +end + +function _build_nullable_primitive(name, v::Vector{T}; nullable::Bool=true) where {T} + S = Base.nonmissingtype(T) + if S === Union{} + t = NullType() + return Field(name, t; nullable=true), + ArrayData(t, length(v), BufferSlice[]; nullcount=length(v)) + end + S === Bool && return _build_bool(name, v; nullable=nullable) + t = arrowtype_for(S) + present = [x !== missing for x in v] + validity = _bitmapbuffer(present) + vals = S[x === missing ? zero(S) : S(x) for x in v] + data = _databuffer(vals) + nc = count(!, present) + # Nullability is the DECLARED element type's, not the observed count's: + # a Union{Missing,T} column with no missing values is still nullable. + return Field(name, t; nullable=nullable), + ArrayData(t, length(v), [validity, data]; nullcount=nc) +end + +function _build_strings(name, v::Vector) + t = Utf8Type(false) + present = [x !== missing for x in v] + offsets = Vector{Int32}(undef, length(v) + 1) + offsets[1] = 0 + nbytes = 0 + for (i, x) in enumerate(v) + nbytes += x === missing ? 0 : ncodeunits(x) + nbytes <= typemax(Int32) || + throw(ArgumentError("column $name exceeds the Int32 offset range")) + offsets[i + 1] = Int32(nbytes) + end + bytes = Vector{UInt8}(undef, nbytes) + pos = 1 + for x in v + x === missing && continue + n = ncodeunits(x) + copyto!(bytes, pos, codeunits(x), 1, n) + pos += n + end + nc = count(!, present) + data = nbytes == 0 ? BufferSlice() : BufferSlice(heapregion(bytes), 0, nbytes) + return Field(name, t; nullable=eltype(v) >: Missing), + ArrayData( + t, + length(v), + [_bitmapbuffer(present), _databuffer(offsets), data]; + nullcount=nc, + ) +end + +function _build_list(name, v::Vector) + present = [x !== missing for x in v] + offsets = Vector{Int32}(undef, length(v) + 1) + offsets[1] = 0 + total = 0 + for (i, x) in enumerate(v) + total += x === missing ? 0 : length(x) + total <= typemax(Int32) || + throw(ArgumentError("column $name exceeds the Int32 offset range")) + offsets[i + 1] = Int32(total) + end + nonmissing = [x for x in v if x !== missing] + childtype = eltype(Base.nonmissingtype(eltype(v))) + flat = isempty(nonmissing) ? childtype[] : reduce(vcat, nonmissing) + cf, cd = fromjulia("item", collect(flat)) + nc = count(!, present) + t = ListType(false) + return Field(name, t; nullable=eltype(v) >: Missing, children=[cf]), + ArrayData( + t, + length(v), + [_bitmapbuffer(present), _databuffer(offsets)]; + children=[cd], + nullcount=nc, + ) +end + +""" + fromjulia_struct(name, nt::NamedTuple) -> (Field, ArrayData) + +Build a struct column from equal-length child vectors (no top-level nulls). +""" +function fromjulia_struct(name, nt::NamedTuple) + pairs = [fromjulia(String(k), v) for (k, v) in Base.pairs(nt)] + n = length(first(values(nt))) + all(length(v) == n for v in values(nt)) || + throw(ArgumentError("struct children must have equal lengths")) + t = StructType() + return Field(name, t; nullable=false, children=[p[1] for p in pairs]), + ArrayData(t, n, [BufferSlice()]; children=[p[2] for p in pairs], nullcount=0) +end + +""" + fromjulia_dict(name, values, indices0) -> (Field, ArrayData) + +Build a dictionary-encoded column from a value pool and 0-based Int32 +indices (`missing` for null slots). +""" +function fromjulia_dict(name, pool::Vector, indices0::Vector) + vf, vd = fromjulia(name, pool) + t = DictionaryType(IntType(32, true), vf.type, false) + present = [x !== missing for x in indices0] + inds = Int32[x === missing ? Int32(0) : Int32(x) for x in indices0] + nc = count(!, present) + return Field(name, t; nullable=nc > 0, children=vf.children), + ArrayData( + t, + length(indices0), + [_bitmapbuffer(present), _databuffer(inds)]; + dictionary=vd, + nullcount=nc, + ) +end + +""" + fromviewentries(name, payloads::Vector{P}, buffers::Vector{Vector{UInt8}}; nullable=true) -> (Field, ArrayData) + fromviewentries(name, payloads::Vector{P}, buf, extra; nullable=true) + +Wrap a vector of Arrow view entries as a Utf8View column, ZERO-COPY. `P` is +any 16-byte isbits type whose values are Arrow StringView entries — the +representation ArrowStrings' `ArrowString` columns use: + + bytes 0..3 Int32 content length (negative marks a null slot) + bytes 4..15 the content, zero-padded (length ≤ 12) + bytes 4..7 the content's 4-byte prefix (length > 12) + bytes 8..11 Int32 buffer index into `buffers` (0-based) + bytes 12..15 Int32 0-based byte offset within that buffer + +`payloads` becomes the views buffer and `buffers` the variadic data buffers, +in order, without copying (the two-buffer form is `[buf, extra]`; a buffer +may be empty, and an all-inline column may have none at all). + +The only work is the validity bitmap: a slot whose length is negative is +null; the spec leaves a null slot's entry bytes unspecified, and neither +this reader's nor the reference implementation's validation reads them. +Long-entry geometry (offsets inside their buffer, prefixes matching the +data) is checked where every builder's is — by +`validate_semantic`/`validate_full` — not here. The scoped-borrow rule of +every zero-copy wrap applies to every vector passed in. +""" +fromviewentries( + name, + payloads::Vector{P}, + buf::Vector{UInt8}, + extra::Vector{UInt8}; + nullable::Bool=true, +) where {P} = fromviewentries(name, payloads, Vector{UInt8}[buf, extra]; nullable=nullable) + +function fromviewentries( + name, + payloads::Vector{P}, + buffers::Vector{Vector{UInt8}}; + nullable::Bool=true, +) where {P} + isbitstype(P) && sizeof(P) == 16 || + throw(ArgumentError("view-entry payloads must be a 16-byte isbits type")) + # The entry words are values assembled by shifts; Arrow's byte layout is + # what those values spell out on a little-endian host, and Core reads + # view entries host-natively. + _native_endianness() == LittleEndian || + throw(ArgumentError("fromviewentries requires a little-endian host")) + n = length(payloads) + present = Vector{Bool}(undef, n) + nnull = 0 + GC.@preserve payloads begin + src = Ptr{Int32}(pointer(payloads)) + for i = 1:n + ok = unsafe_load(src, 4 * i - 3) >= 0 # entry i's length word + present[i] = ok + nnull += !ok + end + end + t = ViewType(true) + slices = BufferSlice[_bitmapbuffer(present), _databuffer(payloads)] + for b in buffers + push!(slices, _databuffer(b)) + end + return Field(name, t; nullable=nullable), ArrayData(t, n, slices; nullcount=nnull) +end + +# --------------------------------------------------------------------------- +# §8 RecordBatch + source protocol +# --------------------------------------------------------------------------- + +""" + RecordBatch + +Schema + equal-length columns: the interchange unit between Core and every +adapter. The IPC and C-stream adapters produce and consume batches; chunked +columns are a facade convenience over them. +""" +struct RecordBatch + schema::Schema + columns::FrozenVector{ArrayData} + nrows::Int64 + function RecordBatch( + schema::Schema, + columns, + nrows::Integer, + validated_dictionaries::Union{Nothing,_ValidatedDictionaries}=nothing, + ) + _validate_schema(schema) + cols = FrozenVector{ArrayData}(columns) + n = Int64(nrows) + n >= 0 || throw(ArgumentError("negative row count")) + length(schema.fields) == length(cols) || + throw(ArgumentError("schema/column count mismatch")) + for (f, c) in zip(schema.fields, cols) + length(c) == n || throw(ArgumentError("unequal column lengths")) + _validate_structural(f, c, validated_dictionaries) + end + return new(schema, cols, n) + end +end +RecordBatch(schema::Schema, columns) = + RecordBatch(schema, columns, isempty(columns) ? 0 : length(first(columns))) + +"Build a low-level `RecordBatch` from a `NamedTuple` of Julia vectors." +function batch(nt::NamedTuple) + pairs = [fromjulia(String(k), v) for (k, v) in Base.pairs(nt)] + sch = Schema([p[1] for p in pairs]) + return RecordBatch(sch, [p[2] for p in pairs]) +end + +""" + RecordBatchSource + +The shared pull-iteration protocol: implement +`nextbatch!(src) -> Union{Nothing,RecordBatch}` and `schema(src)`. The IPC +reader and C-stream importer present this shape, so a writer or dataset +layer need not know which adapter produced the stream. +""" +abstract type RecordBatchSource end +function nextbatch! end +schema(src::RecordBatchSource) = + error("RecordBatchSource implementations must define schema(src)") + +end # module ArrowCore diff --git a/src/ArrowStrings/LICENSE.md b/src/ArrowStrings/LICENSE.md new file mode 100644 index 00000000..d6456956 --- /dev/null +++ b/src/ArrowStrings/LICENSE.md @@ -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/src/ArrowStrings/NOTICE b/src/ArrowStrings/NOTICE new file mode 100644 index 00000000..fc6b69b3 --- /dev/null +++ b/src/ArrowStrings/NOTICE @@ -0,0 +1,5 @@ +Apache Arrow Julia +Copyright 2016-2026 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (https://www.apache.org/). diff --git a/src/ArrowStrings/Project.toml b/src/ArrowStrings/Project.toml new file mode 100644 index 00000000..5b40c9f9 --- /dev/null +++ b/src/ArrowStrings/Project.toml @@ -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. + +name = "ArrowStrings" +uuid = "c38d8858-22a2-449e-9eca-ef92ec15f353" +authors = ["quinnj "] +version = "1.0.0" + +[compat] +julia = "1.10" diff --git a/src/ArrowStrings/README.md b/src/ArrowStrings/README.md new file mode 100644 index 00000000..0da64f14 --- /dev/null +++ b/src/ArrowStrings/README.md @@ -0,0 +1,92 @@ + + +# ArrowStrings.jl + +An inline-else-view string representation designed for Arrow.jl and compatible +parsers. It is a separate package, released and registered independently from +this subdirectory like `ArrowTypes.jl`, so producers and consumers can use the +representation without depending on each other. Until its first release, +Arrow.jl resolves it through a `[sources]` path entry. + +## Installation + +ArrowStrings is not yet registered in General. Until then, run this from the +repository root: + +```julia +import Pkg +Pkg.develop(path="src/ArrowStrings") +``` + +After registration, install a released version from the Julia REPL: + +```julia +import Pkg +Pkg.add("ArrowStrings") +``` + +## API + +* `ArrowString <: AbstractString` — a string value whose 16-byte payload **is** an + Arrow StringView entry: strings of up to 12 bytes are stored inline; + longer strings are a 4-byte prefix plus `(Int32 buffer index, Int32 + offset)` into a byte buffer. Byte access, `==`, `cmp`/`isless`, `hash`, + and iteration never allocate and agree with `String`; `String(s)` copies + out. +* `StringVector{ELT}` — a column of them: a payload vector plus the + byte buffers the views point into. That is an Arrow Utf8View array's + memory (views buffer + variadic data buffers), so Arrow.jl can write the + column without repacking its payloads or data buffers. + `ELT` is `ArrowString` or `Union{Missing, ArrowString}`; `getindex` + allocates nothing; `materialize` copies out to `Vector{String}` (or + `Vector{Union{String,Missing}}` for the nullable `ELT`). +* `ArrowBytes <: AbstractVector{UInt8}` and `BytesVector{ELT}` — the bytes + counterparts, over the same 16-byte payload: Arrow's BinaryView entry is + byte for byte the StringView layout, so one payload type serves both. + `ArrowBytes` gives allocation-free byte access with `==` and `hash` + agreeing with `Vector{UInt8}` (they work as `Dict` keys next to byte + vectors); `Vector{UInt8}(b)` copies one value out, and + `materialize(::BytesVector)` copies a column out to + `Vector{Vector{UInt8}}`. `ELT` is `ArrowBytes` or + `Union{Missing, ArrowBytes}`. + +Everything depends only on Base and uses concrete types. CI compiles and runs +representative construction, access, comparison, and materialization under +JuliaC `--trim=safe`. Buffers must stay under 2 GiB (Arrow's `Int32` view +words). Construction validates payload geometry and prefixes; a builder that +produced every payload itself from bounds it already checked — a parser whose +offsets were validated on read — can pass `Val(:trusted)` as a trailing +constructor argument to skip that re-validation, vouching for the invariants +itself. Do not resize or mutate the payload vector or any referenced buffer +while a column is in use. + +```julia +using ArrowStrings +buf = Vector{UInt8}(codeunits("id,name\n1,abcd\n2,a much longer value\n")) +payloads = [ArrowStrings.inline_payload(buf, 11, 4), + ArrowStrings.view_payload(buf, 18, 19, 0, 17)] +col = StringVector{ArrowString}(payloads, buf, UInt8[]) +col[2] == "a much longer value" # true, no allocation +``` + +The Arrow 3.0 rewrite and this package used Anthropic Claude Code and OpenAI +Codex for code generation, test generation, and review. Apache Arrow +maintainers remain responsible for understanding, reviewing, testing, and +approving the code and each release. diff --git a/src/ArrowStrings/src/ArrowStrings.jl b/src/ArrowStrings/src/ArrowStrings.jl new file mode 100644 index 00000000..1114d85e --- /dev/null +++ b/src/ArrowStrings/src/ArrowStrings.jl @@ -0,0 +1,847 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" + ArrowStrings + +An inline-else-view string representation for Arrow.jl and compatible parsers: +`ArrowString`, a string value whose 16-byte payload IS an Arrow StringView +entry, and `StringVector`, a column of them over a set of byte buffers that IS +an Arrow Utf8View array's memory. A column parsed into this representation can +be written as an Arrow column without repacking its payloads or data buffers. + +Every string is one 16-byte payload (`ArrowStringPayload`, two `UInt64` +words `a` and `b`, packed by explicit shifts so the layout is +endianness-independent): + + a bits 0..31 content length as Int32 (-1 marks a missing value) + bits 32..63 content bytes 1..4 — the full bytes when the string is + inline (length ≤ 12), the 4-byte PREFIX when it is a view + (one word compare — length + prefix — rejects most + unequal views) + b length ≤ 12 content bytes 5..12, zero-padded + length > 12 bits 0..31 an Int32 BUFFER INDEX and bits 32..63 an Int32 + 0-based byte OFFSET of the content within that buffer + +That is byte for byte the Arrow StringView layout (12-byte inline, 4-byte +prefix, int32 buffer index + int32 offset). Arrow's Int32 words are why a +buffer must stay under 2 GiB. Byte access, comparison, hashing, and iteration +never allocate; `String(s)` copies out; `materialize(v)` copies a whole +column out to `Vector{String}`. Everything here depends only on Base and uses +concrete types throughout. + +The same payload is Arrow's BinaryView entry, so the package also carries the +bytes counterparts: `ArrowBytes`, an opaque binary value over one payload, and +`BytesVector`, a column of them that IS a BinaryView array's memory. + +Lifetime: an `ArrowString` view pins its buffer (`data`), and a +`StringVector` pins all of its buffers, exactly like any zero-copy +string view; a consumer that must outlive the source materializes. +""" +module ArrowStrings + +export ArrowString, StringVector, ArrowStringPayload, ArrowBytes, BytesVector + +# Mark the supported non-exported surface with public. It goes through +# Core.eval because Julia < 1.11 cannot parse the public keyword at all. +@static if VERSION >= v"1.11" + Core.eval( + @__MODULE__, + Expr( + :public, + :inline_payload, + :view_payload, + :rebase_payload, + :payloadlength, + :payloadbufidx, + :payloadoffset, + :payloadpos, + :materialize, + :PAYLOAD_MISSING, + :INLINE_MAX, + ), + ) +end + +""" + ArrowStringPayload + +One 16-byte inline-else-view string entry — an Arrow StringView entry (see +[`ArrowStrings`](@ref) for the word layout). `payloadlength(p)` is the length +(negative for missing); for a view, `payloadbufidx(p)` and `payloadoffset(p)` are its +buffer index and 0-based offset, and `payloadpos(p)` the 1-based position. +""" +struct ArrowStringPayload + a::UInt64 + b::UInt64 +end + +"The payload of a missing value: length -1." +const PAYLOAD_MISSING = ArrowStringPayload(UInt64(0xffffffff), zero(UInt64)) +"Longest string stored inline; longer strings are views into a buffer." +const INLINE_MAX = 12 +const EMPTY_BYTES = UInt8[] + +"Return the content byte length; a negative result marks a missing value." +@inline payloadlength(p::ArrowStringPayload) = reinterpret(Int32, p.a % UInt32) +"Return the zero-based referenced-buffer index of a view payload." +@inline payloadbufidx(p::ArrowStringPayload) = reinterpret(Int32, p.b % UInt32) +"Return the zero-based byte offset of a view payload within its buffer." +@inline payloadoffset(p::ArrowStringPayload) = reinterpret(Int32, (p.b >> 32) % UInt32) +"Return the one-based byte position of a view payload within its buffer." +@inline payloadpos(p::ArrowStringPayload) = Int(payloadoffset(p)) + 1 +@inline _viewword(bufidx::Integer, offset0::Integer) = + UInt64(bufidx % UInt32) | (UInt64(offset0 % UInt32) << 32) + +@inline function _checkrange(src::AbstractVector, pos::Int, len::Int, label::String) + Base.require_one_based_indexing(src) + n = length(src) + # pos == n+1 is legal for len == 0: an empty string at the end of a buffer. + (1 <= pos <= n + 1 && len <= n - pos + 1) || + throw(BoundsError("$label: range starts at $pos with length $len in $n bytes")) + return nothing +end + +@inline function _inline_payload_loop(src::AbstractVector{UInt8}, pos::Int, len::Int) + a = UInt64(len % UInt32) + b = zero(UInt64) + @inbounds for i = 1:min(len, 4) + a |= UInt64(src[pos + i - 1]) << (32 + 8 * (i - 1)) + end + @inbounds for i = 5:len + b |= UInt64(src[pos + i - 1]) << (8 * (i - 5)) + end + return ArrowStringPayload(a, b) +end + +""" + inline_payload(src::AbstractVector{UInt8}, pos::Int, len::Int) -> ArrowStringPayload + +The payload of the `len` (≤ 12) bytes of `src` starting at 1-based `pos`, +stored inline. The requested byte range is checked before it is read. +`src` must use one-based indexing. +""" +@inline function inline_payload(src::AbstractVector{UInt8}, pos::Int, len::Int) + 0 <= len <= INLINE_MAX || + throw(ArgumentError("inline_payload: length $len is not in 0:$INLINE_MAX")) + _checkrange(src, pos, len, "inline_payload") + return _inline_payload_loop(src, pos, len) +end + +# Two overlapping little-endian loads gather up to 12 content bytes +# branch-free; the byte-loop fallback only runs within 11 bytes of the +# buffer's end. +@inline function inline_payload(src::Vector{UInt8}, pos::Int, len::Int) + 0 <= len <= INLINE_MAX || + throw(ArgumentError("inline_payload: length $len is not in 0:$INLINE_MAX")) + _checkrange(src, pos, len, "inline_payload") + if pos <= length(src) - 11 + GC.@preserve src begin + p = pointer(src, pos) + lo = ltoh(unsafe_load(Ptr{UInt64}(p))) # content bytes 1..8 + hi = ltoh(unsafe_load(Ptr{UInt64}(p + 4))) # content bytes 5..12 + end + m4 = len >= 4 ? 0x00000000ffffffff : (UInt64(1) << (8 * len)) - 1 + nb = max(len - 4, 0) + m8 = nb >= 8 ? typemax(UInt64) : (UInt64(1) << (8 * nb)) - 1 + return ArrowStringPayload(UInt64(len % UInt32) | ((lo & m4) << 32), hi & m8) + end + return _inline_payload_loop(src, pos, len) +end + +""" + view_payload(src::AbstractVector{UInt8}, srcpos::Int, len::Int, bufidx, offset0) -> ArrowStringPayload + +The payload of a view. The first four content bytes sit at 1-based `srcpos` +in `src`; the complete `len` (> 12) bytes are addressed by buffer index +`bufidx` and 0-based byte `offset0`. The prefix range is checked. A length or +word that does not fit Arrow's Int32 is refused because an oversized value +would wrap into the null marker. `src` must use one-based indexing. +""" +@inline function view_payload( + src::AbstractVector{UInt8}, + srcpos::Int, + len::Int, + bufidx::Integer, + offset0::Integer, +) + INLINE_MAX < len <= typemax(Int32) || throw( + ArgumentError( + "view_payload: length $len is not in $(INLINE_MAX + 1):$(typemax(Int32))", + ), + ) + (0 <= offset0 <= typemax(Int32) && 0 <= bufidx <= typemax(Int32)) || throw( + ArgumentError( + "ArrowString view (buffer $bufidx, offset $offset0) " * + "does not fit Arrow's Int32 view words; buffers must stay under 2 GiB", + ), + ) + _checkrange(src, srcpos, 4, "view_payload") + pre = zero(UInt32) + @inbounds for i = 0:3 + pre |= UInt32(src[srcpos + i]) << (8 * i) + end + a = UInt64(len % UInt32) | (UInt64(pre) << 32) + return ArrowStringPayload(a, _viewword(bufidx, offset0)) +end + +""" + rebase_payload(p::ArrowStringPayload, base::Integer) -> ArrowStringPayload + +The same view entry re-pointed `base` bytes further into its buffer — what +concatenating buffers (a chunk's buffer appended to a column's) needs. +""" +@inline function rebase_payload(p::ArrowStringPayload, base::Integer) + payloadlength(p) > INLINE_MAX || + throw(ArgumentError("rebase_payload requires an out-of-line view payload")) + # base may not fit Int and the sum may overflow; both are user-visible + # errors, not exceptions to propagate. + shift = try + Int(base) + catch + throw(ArgumentError("ArrowString view offset adjustment $base does not fit Int")) + end + off = try + Base.checked_add(Int(payloadoffset(p)), shift) + catch + throw(ArgumentError("rebased ArrowString view offset overflow")) + end + 0 <= off <= typemax(Int32) || throw( + ArgumentError( + "rebased ArrowString view offset $off does not fit " * + "Arrow's Int32 view offset; buffers must stay under 2 GiB", + ), + ) + return ArrowStringPayload(p.a, _viewword(payloadbufidx(p), off)) +end + +@inline _inlinebyte(p::ArrowStringPayload, i::Int) = + i <= 4 ? (p.a >> (32 + 8 * (i - 1))) % UInt8 : (p.b >> (8 * (i - 5))) % UInt8 + +""" +Check one payload against its data buffer: missing marker exactness, inline +zero padding, view range, and prefix agreement. `missingok=true` accepts only +the exact `PAYLOAD_MISSING` sentinel. +""" +function _validate_payload( + p::ArrowStringPayload, + data::Vector{UInt8}; + missingok::Bool=false, +) + len = payloadlength(p) + if len < 0 + missingok && p == PAYLOAD_MISSING || + throw(ArgumentError("invalid missing ArrowString payload")) + return nothing + end + if len <= INLINE_MAX + for i = (Int(len) + 1):INLINE_MAX + iszero(_inlinebyte(p, i)) || + throw(ArgumentError("inline ArrowString payload has nonzero padding")) + end + return nothing + end + off = Int(payloadoffset(p)) + 0 <= off <= length(data) || throw( + ArgumentError( + "ArrowString view offset $off is outside a $(length(data))-byte buffer", + ), + ) + Int(len) <= length(data) - off || throw( + ArgumentError( + "ArrowString view of length $len at offset $off escapes a $(length(data))-byte buffer", + ), + ) + @inbounds for i = 0:3 + _inlinebyte(p, i + 1) == data[off + i + 1] || + throw(ArgumentError("ArrowString view prefix does not match its data buffer")) + end + return nothing +end + +""" + ArrowString <: AbstractString + +A string value: its 16-byte payload plus the byte vector a view's content +lives in (a shared empty vector for inline values). Byte access, direct +comparisons, hashing, and iteration do not allocate; they use the inline +bytes or the retained buffer. Hashing and ordering agree with `String`. +`String(s)` copies out. The two-argument constructor validates the payload +and throws `ArgumentError`; `Val(:unchecked)` is the internal bypass. +""" +struct ArrowString <: AbstractString + p::ArrowStringPayload + data::Vector{UInt8} # dereferenced only when the payload is a view + function ArrowString(p::ArrowStringPayload, data::Vector{UInt8}) + _validate_payload(p, data) + return new(p, data) + end + ArrowString(p::ArrowStringPayload, data::Vector{UInt8}, ::Val{:unchecked}) = + new(p, data) +end + +@inline _unchecked_arrowstring(p::ArrowStringPayload, data::Vector{UInt8}) = + ArrowString(p, data, Val(:unchecked)) + +Base.ncodeunits(s::ArrowString) = Int(payloadlength(s.p)) +Base.codeunit(::ArrowString) = UInt8 +Base.@propagate_inbounds function Base.codeunit(s::ArrowString, i::Int) + @boundscheck 1 <= i <= ncodeunits(s) || throw(BoundsError(s, i)) + len = payloadlength(s.p) + if len <= INLINE_MAX + return i <= 4 ? (s.p.a >> (32 + 8 * (i - 1))) % UInt8 : + (s.p.b >> (8 * (i - 5))) % UInt8 + else + return @inbounds s.data[payloadpos(s.p) + i - 1] + end +end + +function Base.isvalid(s::ArrowString, i::Int) + 1 <= i <= ncodeunits(s) || return false + @inbounds b = codeunit(s, i) + b & 0xc0 == 0x80 || return true + i > 1 || return true + @inbounds b = codeunit(s, i - 1) + 0xc0 <= b <= 0xf7 && return false + b & 0xc0 == 0x80 && i > 2 || return true + @inbounds b = codeunit(s, i - 2) + 0xe0 <= b <= 0xf7 && return false + b & 0xc0 == 0x80 && i > 3 || return true + @inbounds b = codeunit(s, i - 3) + return !(0xf0 <= b <= 0xf7) +end + +# UTF-8 iteration mirroring `String`'s tolerant behavior: Julia `Char`s ARE the +# UTF-8 bytes left-aligned in 32 bits, and a malformed sequence yields the bytes +# consumed so far as an (invalid) Char. +function Base.iterate(s::ArrowString, i::Int=1) + i > ncodeunits(s) && return nothing + @inbounds b1 = codeunit(s, i) + b1 < 0x80 && return (reinterpret(Char, UInt32(b1) << 24), i + 1) + l = b1 < 0xc0 ? 1 : b1 < 0xe0 ? 2 : b1 < 0xf0 ? 3 : b1 < 0xf8 ? 4 : 1 + n = ncodeunits(s) + c = UInt32(b1) << 24 + j = 1 + @inbounds while j < l && i + j <= n + nb = codeunit(s, i + j) + (nb & 0xc0) == 0x80 || break + c |= UInt32(nb) << (24 - 8 * j) + j += 1 + end + return (reinterpret(Char, c), i + j) +end + +# Base's generic AbstractString length is isvalid-count-based, which undercounts +# malformed inputs (String yields each bare continuation byte as its own invalid +# Char). Count by iteration so length/collect agree with String. +function Base.length(s::ArrowString) + n = 0 + for _ in s + n += 1 + end + return n +end + +function Base.:(==)(x::ArrowString, y::ArrowString) + n = ncodeunits(x) + n == ncodeunits(y) || return false + if n <= INLINE_MAX + return x.p.a == y.p.a && x.p.b == y.p.b # payload holds the full content + end + x.p.a == y.p.a || return false # length + 4-byte prefix reject + GC.@preserve x y begin + return ccall( + :memcmp, + Cint, + (Ptr{UInt8}, Ptr{UInt8}, Csize_t), + pointer(x.data, payloadpos(x.p)), + pointer(y.data, payloadpos(y.p)), + n, + ) == 0 + end +end +# Direct byte comparison against String — Base's generic AbstractString == +# decodes chars, an order of magnitude slower on this hot path (filtering and +# grouping compare ArrowString columns against String literals constantly). +function Base.:(==)(x::ArrowString, y::Union{String,SubString{String}}) + n = ncodeunits(x) + n == ncodeunits(y) || return false + GC.@preserve x y begin + py = pointer(y) + if n <= INLINE_MAX + @inbounds for i = 1:n + codeunit(x, i) == unsafe_load(py, i) || return false + end + return true + end + return ccall( + :memcmp, + Cint, + (Ptr{UInt8}, Ptr{UInt8}, Csize_t), + pointer(x.data, payloadpos(x.p)), + py, + n, + ) == 0 + end +end +Base.:(==)(y::Union{String,SubString{String}}, x::ArrowString) = x == y + +# Ordering: memcmp over the bytes, exactly like String's `cmp` (Base's generic +# AbstractString fallback iterates chars — 15-45x slower on sortperm). +# Inline×inline compares in registers; view×view goes straight to memcmp on +# the retained buffers; only the mixed case materializes a stack scratch. +# Raw payload words with content byte k at byte k (byte 1 = LSB of w1) — +# bit-defined, so endian-independent. +@inline _payload_words(s::ArrowString) = + ((s.p.a >> 32) | ((s.p.b & 0xffffffff) << 32), s.p.b >> 32) +@inline _payload_scratch(s::ArrowString) = map(htol, _payload_words(s)) +function Base.cmp(x::ArrowString, y::ArrowString) + nx, ny = ncodeunits(x), ncodeunits(y) + if (nx <= INLINE_MAX) & (ny <= INLINE_MAX) + # Inline×inline compares in registers. Payload words are zero-padded + # past each length, so the first differing big-endian word decides on + # its first differing byte; all-equal words mean the shorter string is + # a prefix of the longer. That is memcmp(min bytes) plus the length + # tiebreak. `&` (not `&&`) keeps this to one branch. + w1x, w2x = _payload_words(x) + w1y, w2y = _payload_words(y) + a, b = bswap(w1x), bswap(w1y) + a == b || return a < b ? -1 : 1 + a, b = bswap(w2x), bswap(w2y) + a == b || return a < b ? -1 : 1 + return cmp(nx, ny) + end + n = min(nx, ny) + rx = Ref(_payload_scratch(x)) + ry = Ref(_payload_scratch(y)) + GC.@preserve x y rx ry begin + px = + nx <= INLINE_MAX ? + Ptr{UInt8}(Base.unsafe_convert(Ptr{Tuple{UInt64,UInt64}}, rx)) : + pointer(x.data, payloadpos(x.p)) + py = + ny <= INLINE_MAX ? + Ptr{UInt8}(Base.unsafe_convert(Ptr{Tuple{UInt64,UInt64}}, ry)) : + pointer(y.data, payloadpos(y.p)) + c = ccall(:memcmp, Cint, (Ptr{UInt8}, Ptr{UInt8}, Csize_t), px, py, n) + end + return c < 0 ? -1 : c > 0 ? 1 : cmp(nx, ny) +end +function Base.cmp(x::ArrowString, y::Union{String,SubString{String}}) + nx, ny = ncodeunits(x), ncodeunits(y) + n = min(nx, ny) + rx = Ref(_payload_scratch(x)) + GC.@preserve x y rx begin + px = + nx <= INLINE_MAX ? + Ptr{UInt8}(Base.unsafe_convert(Ptr{Tuple{UInt64,UInt64}}, rx)) : + pointer(x.data, payloadpos(x.p)) + c = ccall(:memcmp, Cint, (Ptr{UInt8}, Ptr{UInt8}, Csize_t), px, pointer(y), n) + end + return c < 0 ? -1 : c > 0 ? 1 : cmp(nx, ny) +end +Base.cmp(y::Union{String,SubString{String}}, x::ArrowString) = -cmp(x, y) +Base.isless(x::ArrowString, y::ArrowString) = cmp(x, y) < 0 +Base.isless(x::ArrowString, y::Union{String,SubString{String}}) = cmp(x, y) < 0 +Base.isless(y::Union{String,SubString{String}}, x::ArrowString) = cmp(y, x) < 0 + +# hash contract: hash(s) == hash(String(s)) — ArrowStrings are Dict keys +# next to Strings. Base hashes a String's bytes through one C routine; we run +# the same routine over the bytes we already have: the retained buffer for +# views, a stack copy of the payload words for inline strings. No String +# allocation on either path. The routine differs across hashing generations: +# ≤ 1.12 memhash(bytes, n, seed) + seed with seed = h + memhash_seed +# ≥ 1.13 hash_bytes(ptr, n, UInt64(h), HASH_SECRET) % UInt (rapidhash) +# The gate is on the API that exists, not the version number. +@static if isdefined(Base, :hash_bytes) && isdefined(Base, :HASH_SECRET) + @inline _stringhash(p::Ptr{UInt8}, n::Int, h::UInt) = + Base.hash_bytes(p, n, UInt64(h), Base.HASH_SECRET) % UInt +else + @inline function _stringhash(p::Ptr{UInt8}, n::Int, h::UInt) + h += Base.memhash_seed + return ccall(Base.memhash, UInt, (Ptr{UInt8}, Csize_t, UInt32), p, n, h % UInt32) + + h + end +end + +function Base.hash(s::ArrowString, h::UInt) + n = ncodeunits(s) + if n > INLINE_MAX + GC.@preserve s begin + return _stringhash(pointer(s.data, payloadpos(s.p)), n, h) + end + end + # inline: bytes 1-4 are the high 32 bits of `a`, bytes 5-12 are `b` — + # pack them contiguously into two little-endian words: word 1 = bytes 1-8 + # = (a>>32) | (low 32 bits of b) << 32, word 2 = bytes 9-12 = b >> 32 + w1 = (s.p.a >> 32) | ((s.p.b & 0xffffffff) << 32) + w2 = s.p.b >> 32 + scratch = (htol(w1), htol(w2)) + r = Ref(scratch) + GC.@preserve r begin + p = Ptr{UInt8}(Base.unsafe_convert(Ptr{Tuple{UInt64,UInt64}}, r)) + return _stringhash(p, n, h) + end +end + +function Base.String(s::ArrowString) + n = ncodeunits(s) + if n > INLINE_MAX + # view: one memcpy out of the retained buffer + GC.@preserve s begin + return unsafe_string(pointer(s.data, payloadpos(s.p)), n) + end + end + out = Vector{UInt8}(undef, n) + @inbounds for i = 1:n + out[i] = codeunit(s, i) + end + return String(out) +end +Base.convert(::Type{String}, s::ArrowString) = String(s) +Base.Symbol(s::ArrowString) = Symbol(String(s)) +Base.promote_rule(::Type{ArrowString}, ::Type{String}) = String + +function Base.write(io::IO, s::ArrowString) + n = 0 + @inbounds for i = 1:ncodeunits(s) + n += write(io, codeunit(s, i)) + end + return n +end +Base.print(io::IO, s::ArrowString) = (write(io, s); nothing) + +# Validate one column's payloads against its buffers — the shared body of the +# StringVector and BytesVector checked constructors. +function _validate_column( + payloads::Vector{ArrowStringPayload}, + buffers::Vector{Vector{UInt8}}, + missingok::Bool, + what::String, +) + for p in payloads + len = payloadlength(p) + # Inline and missing payloads reference no buffer; pass the + # shared empty vector so only the padding/marker checks run. + if len < 0 + _validate_payload(p, EMPTY_BYTES; missingok=missingok) + elseif len <= INLINE_MAX + _validate_payload(p, EMPTY_BYTES) + else + bufidx = Int(payloadbufidx(p)) + 0 <= bufidx < length(buffers) || throw( + ArgumentError( + "$what view buffer index $bufidx is outside 0:$(length(buffers) - 1)", + ), + ) + _validate_payload(p, buffers[bufidx + 1]) + end + end + return nothing +end + +""" + StringVector{ELT}(payloads, buffers::Vector{Vector{UInt8}}) + StringVector{ELT}(payloads, buf::Vector{UInt8}, extra::Vector{UInt8}) + StringVector{ELT}(payloads, buffers, Val(:trusted)) + +A string column: one payload per element and the byte buffers that view +payloads point into (`buffers[bufidx + 1]` for an entry's buffer index). +`ELT` is `ArrowString` for a column with no missing values, or +`Union{Missing, ArrowString}`. `getindex` returns an `ArrowString` (or +`missing`) with NO allocation; `materialize` copies out to `Vector{String}`. + +Construction validates every payload, including missing markers, inline +padding, buffer indices, byte ranges, and long-string prefixes. This makes +later zero-copy access safe. Do not resize or mutate the payload vector or any +buffer while the column is in use. + +The `Val(:trusted)` constructor skips that validation. It is for builders +that produced every payload themselves from bounds they already checked — a +parser whose offsets were validated as they were read, for example — where +re-validating each entry would double the column's construction cost. The +caller vouches for every invariant the checked constructors enforce; a payload +that violates them makes later access read out of bounds. Payloads that come +from anywhere else go through a checked constructor. + +This is an Arrow Utf8View array's memory: `payloads` is its views buffer and +`buffers` its variadic data buffers, so Arrow.jl can write the column without +repacking either one. The two-buffer constructor is the CSV shape: buffer 0 +the input, buffer 1 the column's `extra` buffer of unescaped values. +""" +struct StringVector{ELT} <: AbstractVector{ELT} + payloads::Vector{ArrowStringPayload} + buffers::Vector{Vector{UInt8}} + function StringVector{ELT}( + payloads::Vector{ArrowStringPayload}, + buffers::Vector{Vector{UInt8}}, + ) where {ELT} + _check_string_elt(ELT) + _validate_column(payloads, buffers, Missing <: ELT, "ArrowString") + return new{ELT}(payloads, buffers) + end + function StringVector{ELT}( + payloads::Vector{ArrowStringPayload}, + buffers::Vector{Vector{UInt8}}, + ::Val{:trusted}, + ) where {ELT} + _check_string_elt(ELT) + return new{ELT}(payloads, buffers) + end +end +_check_string_elt(ELT) = + (ELT === ArrowString || ELT === Union{Missing,ArrowString}) || throw( + ArgumentError( + "StringVector element type must be ArrowString or Union{Missing,ArrowString}", + ), + ) +function StringVector{ELT}( + payloads::Vector{ArrowStringPayload}, + buf::Vector{UInt8}, + extra::Vector{UInt8}, +) where {ELT} + return StringVector{ELT}(payloads, Vector{UInt8}[buf, extra]) +end + +function Base.size(v::StringVector) + return size(v.payloads) +end +Base.@propagate_inbounds @inline function Base.getindex( + v::StringVector{ELT}, + i::Int, +) where {ELT} + @boundscheck checkbounds(v.payloads, i) + @inbounds p = v.payloads[i] + len = payloadlength(p) + len < 0 && return missing + len <= INLINE_MAX && return _unchecked_arrowstring(p, EMPTY_BYTES) + return _unchecked_arrowstring(p, v.buffers[payloadbufidx(p) + 1]) +end +# All-present columns skip the missing branch entirely — the concrete return +# type is what lets access compile down to zero allocations. +Base.@propagate_inbounds @inline function Base.getindex( + v::StringVector{ArrowString}, + i::Int, +) + @boundscheck checkbounds(v.payloads, i) + @inbounds p = v.payloads[i] + len = payloadlength(p) + len <= INLINE_MAX && return _unchecked_arrowstring(p, EMPTY_BYTES) + return _unchecked_arrowstring(p, v.buffers[payloadbufidx(p) + 1]) +end + +""" + materialize(v::StringVector) -> Vector{String} or Vector{Union{String,Missing}} + +Copy every element out to a plain `String`, detaching the result from the +column's buffers. +""" +function materialize(v::StringVector{ELT}) where {ELT} + out = Vector{ELT === ArrowString ? String : Union{String,Missing}}(undef, length(v)) + scratch = Vector{UInt8}(undef, 16) # inline payloads reconstruct via two word stores + GC.@preserve scratch begin + q = pointer(scratch) + @inbounds for i in eachindex(v.payloads) + p = v.payloads[i] + len = payloadlength(p) + if len < 0 + out[i] = missing + elseif len <= INLINE_MAX + unsafe_store!(Ptr{UInt64}(q), htol((p.a >> 32) | (p.b << 32))) + unsafe_store!(Ptr{UInt64}(q + 8), htol(p.b >> 32)) + out[i] = unsafe_string(q, len) + else + src = v.buffers[payloadbufidx(p) + 1] + GC.@preserve src begin + out[i] = unsafe_string(pointer(src, payloadpos(p)), len) + end + end + end + end + return out +end + +# ---- bytes columns: the same payload machinery for opaque binary values ---- + +""" + ArrowBytes <: AbstractVector{UInt8} + +A binary value: its 16-byte payload plus the byte vector a view's content +lives in — the bytes counterpart of [`ArrowString`](@ref). The payload layout +is Arrow's BinaryView entry, which is byte for byte the StringView layout, so +[`ArrowStringPayload`](@ref) serves both. Byte access, comparison, and hashing +do not allocate; hashing and equality agree with `Vector{UInt8}` through the +generic `AbstractArray` definitions. `Vector{UInt8}(b)` copies out. The +two-argument constructor validates the payload and throws `ArgumentError`; +`Val(:unchecked)` is the internal bypass. +""" +struct ArrowBytes <: AbstractVector{UInt8} + p::ArrowStringPayload + data::Vector{UInt8} # dereferenced only when the payload is a view + function ArrowBytes(p::ArrowStringPayload, data::Vector{UInt8}) + _validate_payload(p, data) + return new(p, data) + end + ArrowBytes(p::ArrowStringPayload, data::Vector{UInt8}, ::Val{:unchecked}) = new(p, data) +end + +@inline _unchecked_arrowbytes(p::ArrowStringPayload, data::Vector{UInt8}) = + ArrowBytes(p, data, Val(:unchecked)) + +Base.size(b::ArrowBytes) = (Int(payloadlength(b.p)),) +Base.IndexStyle(::Type{ArrowBytes}) = IndexLinear() +Base.@propagate_inbounds function Base.getindex(b::ArrowBytes, i::Int) + @boundscheck checkbounds(b, i) + len = payloadlength(b.p) + len <= INLINE_MAX && return _inlinebyte(b.p, i) + return @inbounds b.data[payloadpos(b.p) + i - 1] +end + +# One word compares length + first four bytes; equal inline payloads then +# compare in registers, equal-prefix views memcmp their retained buffers. +function Base.:(==)(x::ArrowBytes, y::ArrowBytes) + x.p.a == y.p.a || return false + n = payloadlength(x.p) + n <= INLINE_MAX && return x.p.b == y.p.b + GC.@preserve x y begin + return ccall( + :memcmp, + Cint, + (Ptr{UInt8}, Ptr{UInt8}, Csize_t), + pointer(x.data, payloadpos(x.p)), + pointer(y.data, payloadpos(y.p)), + n, + ) == 0 + end +end + +function Base.Vector{UInt8}(b::ArrowBytes) + n = Int(payloadlength(b.p)) + out = Vector{UInt8}(undef, n) + if n > INLINE_MAX + GC.@preserve b out begin + unsafe_copyto!(pointer(out), pointer(b.data, payloadpos(b.p)), n) + end + else + @inbounds for i = 1:n + out[i] = _inlinebyte(b.p, i) + end + end + return out +end +Base.convert(::Type{Vector{UInt8}}, b::ArrowBytes) = Vector{UInt8}(b) + +""" + BytesVector{ELT}(payloads, buffers::Vector{Vector{UInt8}}) + BytesVector{ELT}(payloads, buffers, Val(:trusted)) + +A binary column: one payload per element and the byte buffers that view +payloads point into — the bytes counterpart of [`StringVector`](@ref), and an +Arrow BinaryView array's memory. `ELT` is `ArrowBytes` or +`Union{Missing, ArrowBytes}`. `getindex` returns an `ArrowBytes` (or +`missing`) with NO allocation; `materialize` copies out to +`Vector{Vector{UInt8}}`. + +Construction validates every payload exactly as [`StringVector`](@ref) does, +and the `Val(:trusted)` constructor skips that validation under the same +contract. Do not resize or mutate the payload vector or any buffer while the +column is in use. +""" +struct BytesVector{ELT} <: AbstractVector{ELT} + payloads::Vector{ArrowStringPayload} + buffers::Vector{Vector{UInt8}} + function BytesVector{ELT}( + payloads::Vector{ArrowStringPayload}, + buffers::Vector{Vector{UInt8}}, + ) where {ELT} + _check_bytes_elt(ELT) + _validate_column(payloads, buffers, Missing <: ELT, "ArrowBytes") + return new{ELT}(payloads, buffers) + end + function BytesVector{ELT}( + payloads::Vector{ArrowStringPayload}, + buffers::Vector{Vector{UInt8}}, + ::Val{:trusted}, + ) where {ELT} + _check_bytes_elt(ELT) + return new{ELT}(payloads, buffers) + end +end +_check_bytes_elt(ELT) = + (ELT === ArrowBytes || ELT === Union{Missing,ArrowBytes}) || throw( + ArgumentError( + "BytesVector element type must be ArrowBytes or Union{Missing,ArrowBytes}", + ), + ) + +Base.size(v::BytesVector) = size(v.payloads) +Base.@propagate_inbounds @inline function Base.getindex( + v::BytesVector{ELT}, + i::Int, +) where {ELT} + @boundscheck checkbounds(v.payloads, i) + @inbounds p = v.payloads[i] + len = payloadlength(p) + len < 0 && return missing + len <= INLINE_MAX && return _unchecked_arrowbytes(p, EMPTY_BYTES) + return _unchecked_arrowbytes(p, v.buffers[payloadbufidx(p) + 1]) +end +# All-present columns skip the missing branch entirely — the concrete return +# type is what lets access compile down to zero allocations. +Base.@propagate_inbounds @inline function Base.getindex(v::BytesVector{ArrowBytes}, i::Int) + @boundscheck checkbounds(v.payloads, i) + @inbounds p = v.payloads[i] + len = payloadlength(p) + len <= INLINE_MAX && return _unchecked_arrowbytes(p, EMPTY_BYTES) + return _unchecked_arrowbytes(p, v.buffers[payloadbufidx(p) + 1]) +end + +""" + materialize(v::BytesVector) -> Vector{Vector{UInt8}} or Vector{Union{Vector{UInt8},Missing}} + +Copy every element out to a plain `Vector{UInt8}`, detaching the result from +the column's buffers. +""" +function materialize(v::BytesVector{ELT}) where {ELT} + out = Vector{ELT === ArrowBytes ? Vector{UInt8} : Union{Vector{UInt8},Missing}}( + undef, + length(v), + ) + @inbounds for i in eachindex(v.payloads) + p = v.payloads[i] + len = payloadlength(p) + if len < 0 + out[i] = missing + else + b = Vector{UInt8}(undef, len) + if len <= INLINE_MAX + for j = 1:len + b[j] = _inlinebyte(p, j) + end + else + src = v.buffers[payloadbufidx(p) + 1] + GC.@preserve src b begin + unsafe_copyto!(pointer(b), pointer(src, payloadpos(p)), len) + end + end + out[i] = b + end + end + return out +end + +end # module diff --git a/src/ArrowStrings/test/Project.toml b/src/ArrowStrings/test/Project.toml new file mode 100644 index 00000000..81a4d260 --- /dev/null +++ b/src/ArrowStrings/test/Project.toml @@ -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. + +[deps] +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/src/ArrowStrings/test/runtests.jl b/src/ArrowStrings/test/runtests.jl new file mode 100644 index 00000000..7c579513 --- /dev/null +++ b/src/ArrowStrings/test/runtests.jl @@ -0,0 +1,499 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +using Test, Random +using ArrowStrings +const AS = ArrowStrings + +# --- helpers (top-level so allocation probes measure the loop, not closures) -- + +struct ZeroBasedBytes <: AbstractVector{UInt8} + data::Vector{UInt8} +end +Base.size(v::ZeroBasedBytes) = size(v.data) +Base.axes(v::ZeroBasedBytes) = (0:(length(v.data) - 1),) +Base.IndexStyle(::Type{ZeroBasedBytes}) = IndexLinear() +function Base.getindex(v::ZeroBasedBytes, i::Int) + checkbounds(v, i) + return v.data[i + 1] +end + +function asfrombytes(bytes::Vector{UInt8}) + n = length(bytes) + n <= AS.INLINE_MAX && return ArrowString(AS.inline_payload(bytes, 1, n), AS.EMPTY_BYTES) + return ArrowString(AS.view_payload(bytes, 1, n, 0, 0), bytes) +end + +function asscratchbytes(s::ArrowString) + r = Ref(AS._payload_scratch(s)) + out = Vector{UInt8}(undef, 16) + GC.@preserve r begin + p = Ptr{UInt8}(Base.unsafe_convert(Ptr{Tuple{UInt64,UInt64}}, r)) + unsafe_copyto!(pointer(out), p, 16) + end + return out +end + +function foldashash(v, h::UInt) + @inbounds for x in v + h = hash(x, h) + end + return h +end + +function foldascmp(v) + s = 0 + @inbounds for i = 2:length(v) + s += cmp(v[i - 1], v[i]) + (v[i - 1] == v[i]) + end + return s +end + +function sumncodeunits(c::StringVector{ArrowString}) + t = 0 + for i in eachindex(c) + t += ncodeunits(c[i]) + end + return t +end + +# Allocation is measured INSIDE type-stable top-level functions: an +# `@allocated` in testset scope can charge the closure's own boxing (16 +# bytes on Julia 1.10) to the kernel it is measuring. +allocated_hash(v::Vector{ArrowString}) = @allocated(foldashash(v, UInt(9))) +allocated_cmp(v::Vector{ArrowString}) = @allocated(foldascmp(v)) +allocated_access(c::StringVector{ArrowString}) = @allocated(sumncodeunits(c)) + +# A column from Strings, laid out the CSV way: inline when it fits, else a +# view into buffer 0 (`buf`), or into buffer 1 (`extra`) when `inextra(i)`. +function column(strings::Vector; inextra=i -> false) + buf = UInt8[] + extra = UInt8[] + payloads = ArrowStringPayload[] + for (i, s) in enumerate(strings) + if s === missing + push!(payloads, AS.PAYLOAD_MISSING) + continue + end + bytes = Vector{UInt8}(codeunits(s)) + n = length(bytes) + if n <= AS.INLINE_MAX + push!(payloads, AS.inline_payload(bytes, 1, n)) + else + target = inextra(i) ? extra : buf + off0 = length(target) + append!(target, bytes) + push!(payloads, AS.view_payload(target, off0 + 1, n, inextra(i) ? 1 : 0, off0)) + end + end + ELT = any(ismissing, strings) ? Union{Missing,ArrowString} : ArrowString + return StringVector{ELT}(payloads, buf, extra) +end + +@testset "ArrowStrings" begin + @testset "payload layout is an Arrow StringView entry" begin + # inline: bytes 0..3 length, 4..15 content zero-padded + p = AS.inline_payload(Vector{UInt8}(codeunits("hello")), 1, 5) + words = reinterpret(UInt8, [htol(p.a), htol(p.b)]) + @test words[1:4] == reinterpret(UInt8, Int32[5]) + @test words[5:9] == codeunits("hello") + @test all(iszero, words[10:16]) + @test AS.payloadlength(p) == 5 + # view: bytes 0..3 length, 4..7 prefix, 8..11 buffer index, 12..15 offset + data = Vector{UInt8}(codeunits("xxthirteen-bytesyy")) + q = AS.view_payload(data, 3, 13, 1, 2) + words = reinterpret(UInt8, [htol(q.a), htol(q.b)]) + @test words[1:4] == reinterpret(UInt8, Int32[13]) + @test words[5:8] == codeunits("thir") + @test words[9:12] == reinterpret(UInt8, Int32[1]) + @test words[13:16] == reinterpret(UInt8, Int32[2]) + @test AS.payloadlength(q) == 13 && + AS.payloadbufidx(q) == 1 && + AS.payloadoffset(q) == 2 && + AS.payloadpos(q) == 3 + # missing + @test AS.payloadlength(AS.PAYLOAD_MISSING) == -1 + # rebase moves the offset only + r = AS.rebase_payload(q, 100) + @test AS.payloadbufidx(r) == 1 && AS.payloadoffset(r) == 102 && r.a == q.a + # words that do not fit Arrow's Int32 refuse; the boundaries are accepted + @test_throws ArgumentError AS.view_payload( + data, + 3, + 13, + 0, + Int64(typemax(Int32)) + 1, + ) + @test_throws ArgumentError AS.view_payload(data, 3, 13, -1, 0) + @test_throws ArgumentError AS.rebase_payload(q, Int64(typemax(Int32))) + @test AS.payloadoffset(AS.view_payload(data, 3, 13, 0, Int64(typemax(Int32)))) == + typemax(Int32) + @test AS.payloadbufidx(AS.view_payload(data, 3, 13, Int64(typemax(Int32)), 0)) == + typemax(Int32) + # ... and so does a length outside (12, typemax(Int32)] — an oversized + # length would otherwise wrap into the null marker + @test AS.payloadlength(AS.view_payload(data, 3, Int(typemax(Int32)), 0, 0)) == + typemax(Int32) + @static if Sys.WORD_SIZE > 32 + @test_throws ArgumentError AS.view_payload( + data, + 3, + Int(typemax(Int32)) + 1, + 0, + 0, + ) + end + @test_throws ArgumentError AS.view_payload(data, 3, 12, 0, 0) + @test_throws ArgumentError AS.view_payload(data, 3, -1, 0, 0) + @test_throws ArgumentError AS.inline_payload(data, 1, 13) + @test_throws ArgumentError AS.inline_payload(data, 1, -1) + @test AS.payloadlength(AS.inline_payload(UInt8[], 1, 0)) == 0 + @test_throws BoundsError AS.inline_payload(UInt8[], 0, 0) + @test_throws BoundsError AS.inline_payload(UInt8[], 2, 0) + @test_throws BoundsError AS.inline_payload(UInt8[], 1, 1) + @test_throws BoundsError AS.inline_payload(data, length(data), 2) + @test_throws BoundsError AS.view_payload(data, 0, 13, 0, 0) + @test_throws BoundsError AS.view_payload(data, length(data) - 2, 13, 0, 0) + zero_based = ZeroBasedBytes(copy(data)) + @test_throws ArgumentError AS.inline_payload(zero_based, 1, 4) + @test_throws ArgumentError AS.inline_payload(zero_based, 0, 4) + @test_throws ArgumentError AS.view_payload(zero_based, 1, 13, 0, 0) + @test_throws ArgumentError AS.view_payload(zero_based, 0, 13, 0, 0) + @test_throws ArgumentError AS.rebase_payload(p, 1) + @test_throws ArgumentError AS.rebase_payload(AS.PAYLOAD_MISSING, 1) + # the byte-loop inline fallback (near the end of the buffer) agrees with + # the two-load fast path + long = Vector{UInt8}(codeunits("abcdefghijklmnopqrstuvwxyz")) + for len = 0:12, pos = 1:(length(long) - len + 1) + fast = AS.inline_payload(long, pos, len) + tail = long[pos:(pos + len - 1)] + slow = AS.inline_payload(tail, 1, len) # pos + 11 > length ⇒ byte loop + @test fast == slow + end + end + + @testset "ArrowString: equality, hashing, ordering agree with String" begin + # inline/view boundary: 12 bytes inline, 13 views the buffer + col = column(["x"^12, "y"^13]) + @test col isa StringVector{ArrowString} + @test col[1] == "x"^12 && col[2] == "y"^13 + @test ncodeunits(col[1]) == 12 && ncodeunits(col[2]) == 13 + @test String(col[1]) == "x"^12 && String(col[2]) == "y"^13 + @test col[1] == "x"^12 && "x"^12 == col[1] + @test isequal(col[1], "x"^12) && hash(col[1]) == hash("x"^12) + d = Dict("x"^12 => 1) + @test d[col[1]] == 1 + d2 = Dict(col[2] => 2) + @test d2["y" ^ 13] == 2 + @test sort([col[2], col[1]]) == [col[1], col[2]] + @test cmp(col[1], col[2]) == cmp("x"^12, "y"^13) + + # Exhaust the payload-length byte with data that includes NUL, invalid + # UTF-8, all-one bytes, and sentinel-like runs; odd out-of-line lengths + # use buffer index 1, even ones buffer index 0 — hashing must not care. + pattern = UInt8[ + 0x00, + 0xff, + 0xff, + 0xff, + 0xff, + 0xff, + 0xff, + 0xff, + 0xff, + 0xff, + 0xff, + 0xff, + 0x80, + 0xc0, + 0x7f, + 0x41, + 0xfe, + ] + strings = String[] + payloads = ArrowString[] + for n = 0:255 + bytes = + UInt8[xor(pattern[mod1(i, length(pattern))], UInt8(i % 251)) for i = 1:n] + push!(strings, String(copy(bytes))) + if n <= AS.INLINE_MAX + push!(payloads, ArrowString(AS.inline_payload(bytes, 1, n), AS.EMPTY_BYTES)) + else + data = vcat(UInt8[0x11], bytes, UInt8[0x22]) + bufidx = isodd(n) ? 1 : 0 + push!(payloads, ArrowString(AS.view_payload(data, 2, n, bufidx, 1), data)) + end + end + seeds = UInt[0, 1, 7, typemax(UInt), 0x0123456789abcdef] + @test all( + hash(payloads[i], h) == hash(strings[i], h) for + i in eachindex(payloads), h in seeds + ) + @test all( + [codeunit(payloads[i], j) for j = 1:ncodeunits(payloads[i])] == collect(codeunits(strings[i])) + for i in eachindex(payloads) + ) + @test all( + begin + n = ncodeunits(payloads[i]) + bytes = asscratchbytes(payloads[i]) + bytes[1:n] == collect(codeunits(strings[i])) && + all(iszero, bytes[(n + 1):end]) + end for i = 1:(AS.INLINE_MAX + 1) + ) + @test all( + cmp(payloads[i], payloads[j]) == cmp(strings[i], strings[j]) && + isless(payloads[i], payloads[j]) == isless(strings[i], strings[j]) && + cmp(payloads[i], strings[j]) == cmp(strings[i], strings[j]) && + cmp(strings[i], payloads[j]) == cmp(strings[i], strings[j]) && + (payloads[i] == payloads[j]) == (strings[i] == strings[j]) && + (payloads[i] == strings[j]) == (strings[i] == strings[j]) for + i in eachindex(payloads), j in eachindex(payloads) + ) + @test sortperm(payloads) == sortperm(strings) + valid = ["a", "abcdefgh1234", "abcdefgh12345", "α", "漢字", "z"^40, "a\0b"] + validcs = [asfrombytes(Vector{UInt8}(codeunits(s))) for s in valid] + substrings = [ + SubString("!" * s * "?", 2, prevind("!" * s * "?", lastindex("!" * s * "?"))) for s in valid + ] + @test all( + cmp(validcs[i], substrings[j]) == cmp(valid[i], String(substrings[j])) && + cmp(substrings[j], validcs[i]) == cmp(String(substrings[j]), valid[i]) && + (validcs[i] == substrings[j]) == (valid[i] == String(substrings[j])) for + i in eachindex(validcs), j in eachindex(substrings) + ) + @test isless(first(validcs), missing) == isless(first(valid), missing) + @test isless(missing, first(validcs)) == isless(missing, first(valid)) + # hashing and comparing across every inline/view mix never allocates + allocated_hash(payloads) + @test allocated_hash(payloads) == 0 + allocated_cmp(payloads) + @test allocated_cmp(payloads) == 0 + # Symbol, promotion, print + @test Symbol(asfrombytes(Vector{UInt8}(codeunits("αβγδεζηθικλμ")))) == :αβγδεζηθικλμ + @test promote_type(ArrowString, String) === String + io = IOBuffer() + print(io, col[2]) + @test String(take!(io)) == "y"^13 + @test convert(String, col[1]) == "x"^12 + end + + @testset "ArrowString: iteration and character indexing match String" begin + rng = MersenneTwister(99) + for _ = 1:200 + n = rand(rng, 0:24) + bytes = rand(rng, UInt8, n) + s = String(copy(bytes)) + v = asfrombytes(copy(bytes)) + @test collect(v) == collect(s) + @test v == s && hash(v) == hash(s) + @test length(v) == length(s) + end + # Character-index APIs use the same tolerant partition as String: a bare + # continuation byte is its own invalid Char and starts at a valid index; a + # continuation consumed by a preceding lead byte does not. + invalidcases = ( + UInt8[0x80], + UInt8[0x61, 0x80, 0x62], + UInt8[0xc2], + UInt8[0xc2, 0x41], + UInt8[0xe0, 0x80], + UInt8[0xf0, 0x80, 0x41], + UInt8[0xc2, 0x80], + Vector{UInt8}(codeunits("thirteen-bytes-and-then-α-β")), + ) + result(f) = + try + (:value, f()) + catch e + (:error, typeof(e)) + end + for bytes in invalidcases + s = String(copy(bytes)) + v = asfrombytes(copy(bytes)) + @test collect(eachindex(v)) == collect(eachindex(s)) + @test lastindex(v) == lastindex(s) + for i = 0:(length(bytes) + 1) + @test isvalid(v, i) == isvalid(s, i) + @test result(() -> thisind(v, i)) == result(() -> thisind(s, i)) + @test result(() -> nextind(v, i)) == result(() -> nextind(s, i)) + @test result(() -> prevind(v, i)) == result(() -> prevind(s, i)) + @test result(() -> v[i]) == result(() -> s[i]) + end + for i = 1:length(bytes), j = i:length(bytes) + @test result(() -> String(SubString(v, i, j))) == + result(() -> String(SubString(s, i, j))) + end + end + end + + @testset "StringVector: buffers, missing, materialize, allocation" begin + strings = ["value$(i)_" * "p"^(i % 20) for i = 1:1000] + col = column(strings; inextra=i -> i % 3 == 0) + @test col isa StringVector{ArrowString} + @test length(col) == 1000 && length(col.buffers) == 2 + @test collect(String, col) == strings + # long values landed in the buffer their index says + for (i, s) in enumerate(strings) + ncodeunits(s) > AS.INLINE_MAX || continue + @test AS.payloadbufidx(col.payloads[i]) == (i % 3 == 0 ? 1 : 0) + end + allocated_access(col) + @test allocated_access(col) == 0 + m = AS.materialize(col) + @test m isa Vector{String} && m == strings + + withmissing = Any["a", missing, "twelve-bytes", "a much longer value", missing] + mcol = column(withmissing; inextra=i -> i == 4) + @test mcol isa StringVector{Union{Missing,ArrowString}} + @test isequal(collect(mcol), withmissing) + @test AS.materialize(mcol) isa Vector{Union{String,Missing}} + @test isequal(AS.materialize(mcol), withmissing) + @test eltype(mcol) === Union{Missing,ArrowString} + + # N buffers: an Arrow Utf8View column may spread views over any number of + # data buffers; the buffer index selects among them + b0 = Vector{UInt8}(codeunits("--first-buffer-value--")) + b1 = Vector{UInt8}(codeunits("second-buffer-value!!")) + b2 = Vector{UInt8}(codeunits("xxthird buffer, longer value")) + payloads = ArrowStringPayload[ + AS.view_payload(b0, 3, 18, 0, 2), + AS.view_payload(b1, 1, 19, 1, 0), + AS.view_payload(b2, 3, 26, 2, 2), + AS.inline_payload(b1, 1, 6), + ] + ncol = StringVector{ArrowString}(payloads, Vector{UInt8}[b0, b1, b2]) + @test collect(String, ncol) == [ + "first-buffer-value", + "second-buffer-value", + "third buffer, longer value", + "second", + ] + # Construction rejects every geometry that could make later + # zero-copy access leave a buffer. + @test_throws ArgumentError StringVector{ArrowString}( + [AS.view_payload(b0, 3, 18, 7, 2)], + Vector{UInt8}[b0], + ) + @test_throws ArgumentError StringVector{ArrowString}( + [AS.view_payload(b0, 3, 18, 0, 100)], + Vector{UInt8}[b0], + ) + @test_throws ArgumentError StringVector{ArrowString}( + [AS.PAYLOAD_MISSING], + Vector{UInt8}[], + ) + @test_throws ArgumentError StringVector{Int}(ArrowStringPayload[], Vector{UInt8}[]) + badmissing = ArrowStringPayload(UInt64(0xfffffffe), zero(UInt64)) + @test_throws ArgumentError StringVector{Union{Missing,ArrowString}}( + [badmissing], + Vector{UInt8}[], + ) + badprefix = ArrowStringPayload(payloads[1].a ⊻ (UInt64(1) << 32), payloads[1].b) + @test_throws ArgumentError StringVector{ArrowString}([badprefix], Vector{UInt8}[b0]) + short = AS.inline_payload(b1, 1, 5) + badpadding = ArrowStringPayload(short.a, short.b | (UInt64(1) << 8)) + @test_throws ArgumentError ArrowString(badpadding, AS.EMPTY_BYTES) + @test_throws ArgumentError ArrowString(AS.view_payload(b0, 3, 18, 0, 100), b0) + # rebase: appending one buffer to another re-points its entries + combined = vcat(b0, b1) + rebased = AS.rebase_payload(payloads[2], length(b0)) + @test String(ArrowString(rebased, combined)) == "second-buffer-value" + end + + @testset "StringVector: trusted constructor" begin + data = Vector{UInt8}(codeunits("hello, a value that is longer than twelve bytes")) + payloads = ArrowStringPayload[ + AS.inline_payload(data, 1, 5), + AS.view_payload(data, 8, length(data) - 7, 0, 7), + ] + checked = StringVector{ArrowString}(payloads, Vector{UInt8}[data]) + trusted = StringVector{ArrowString}(payloads, Vector{UInt8}[data], Val(:trusted)) + @test collect(String, trusted) == collect(String, checked) + # the element-type gate still applies + @test_throws ArgumentError StringVector{Int}( + ArrowStringPayload[], + Vector{UInt8}[], + Val(:trusted), + ) + # a payload the checked constructor rejects is accepted untouched: the + # caller vouches instead (nonzero inline padding is safe to observe, + # so it is the probe) + short = AS.inline_payload(data, 1, 5) + badpadding = ArrowStringPayload(short.a, short.b | (UInt64(1) << 8)) + @test_throws ArgumentError StringVector{ArrowString}( + [badpadding], + Vector{UInt8}[data], + ) + tv = StringVector{ArrowString}([badpadding], Vector{UInt8}[data], Val(:trusted)) + @test tv isa StringVector{ArrowString} && length(tv) == 1 + mp = [AS.PAYLOAD_MISSING] + mv = StringVector{Union{Missing,ArrowString}}(mp, Vector{UInt8}[], Val(:trusted)) + @test mv[1] === missing + end + + @testset "ArrowBytes and BytesVector" begin + vals = Vector{UInt8}[rand(UInt8, n) for n in (0, 3, 12, 13, 40)] + data = vcat(vals...) + offs = cumsum([0; map(length, vals)]) + payloads = ArrowStringPayload[] + for (i, v) in enumerate(vals) + n = length(v) + if n <= AS.INLINE_MAX + push!(payloads, AS.inline_payload(data, offs[i] + 1, n)) + else + push!(payloads, AS.view_payload(data, offs[i] + 1, n, 0, offs[i])) + end + end + col = BytesVector{ArrowBytes}(payloads, Vector{UInt8}[data]) + @test col isa BytesVector{ArrowBytes} && length(col) == length(vals) + @test all(col[i] == vals[i] for i in eachindex(vals)) + @test all(isequal(col[i], vals[i]) for i in eachindex(vals)) + # hash agrees with Vector{UInt8} (generic AbstractArray hashing), so + # ArrowBytes work as Dict keys next to byte vectors + @test all(hash(col[i]) == hash(vals[i]) for i in eachindex(vals)) + @test Vector{UInt8}(col[2]) == vals[2] && Vector{UInt8}(col[5]) == vals[5] + @test col[5][7] == vals[5][7] && col[2][1] == vals[2][1] + @test cmp(col[2], col[2]) == 0 && cmp(collect(col[2]), vals[2]) == 0 + m = AS.materialize(col) + @test m isa Vector{Vector{UInt8}} && m == vals + # equality: one-word length+prefix reject, then memcmp on the tail + long1 = Vector{UInt8}([collect(0x01:0x10); 0xaa]) + long2 = Vector{UInt8}([collect(0x01:0x10); 0xbb]) + ab1 = ArrowBytes(AS.view_payload(long1, 1, 17, 0, 0), long1) + ab2 = ArrowBytes(AS.view_payload(long2, 1, 17, 0, 0), long2) + @test ab1 != ab2 + @test ab1 == ArrowBytes(AS.view_payload(long1, 1, 17, 0, 0), long1) + + # missing markers, materialize, and the trusted constructor + mp = ArrowStringPayload[payloads[4], AS.PAYLOAD_MISSING] + mcol = BytesVector{Union{Missing,ArrowBytes}}(mp, Vector{UInt8}[data]) + @test mcol[1] == vals[4] && mcol[2] === missing + @test isequal(AS.materialize(mcol), [vals[4], missing]) + @test AS.materialize(mcol) isa Vector{Union{Vector{UInt8},Missing}} + tcol = + BytesVector{Union{Missing,ArrowBytes}}(mp, Vector{UInt8}[data], Val(:trusted)) + @test tcol[2] === missing && tcol[1] == vals[4] + # checked constructors reject bad geometry exactly like the string side + @test_throws ArgumentError BytesVector{ArrowBytes}( + [AS.view_payload(data, 1, 20, 0, 100)], + Vector{UInt8}[data], + ) + @test_throws ArgumentError BytesVector{Int}(ArrowStringPayload[], Vector{UInt8}[]) + @test_throws ArgumentError ArrowBytes(AS.view_payload(data, 1, 20, 0, 100), data) + end +end diff --git a/src/ArrowTypes/LICENSE.md b/src/ArrowTypes/LICENSE.md index 136e1c6a..d6456956 100644 --- a/src/ArrowTypes/LICENSE.md +++ b/src/ArrowTypes/LICENSE.md @@ -1,15 +1,202 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. \ No newline at end of file + + 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/src/ArrowTypes/NOTICE b/src/ArrowTypes/NOTICE new file mode 100644 index 00000000..fc6b69b3 --- /dev/null +++ b/src/ArrowTypes/NOTICE @@ -0,0 +1,5 @@ +Apache Arrow Julia +Copyright 2016-2026 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (https://www.apache.org/). diff --git a/src/ArrowTypes/Project.toml b/src/ArrowTypes/Project.toml index 0166f602..1f2d1de8 100644 --- a/src/ArrowTypes/Project.toml +++ b/src/ArrowTypes/Project.toml @@ -18,11 +18,19 @@ name = "ArrowTypes" uuid = "31f734f8-188a-4ce0-8406-c8a06bd891cd" authors = ["quinnj "] -version = "2.3.0" +version = "2.4.0" [deps] Sockets = "6462fe0b-24de-5631-8697-dd941f90decc" UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4" +[extras] +OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[targets] +test = ["OffsetArrays", "Test"] + [compat] +OffsetArrays = "0.11, 1" julia = "1.0" diff --git a/src/ArrowTypes/README.md b/src/ArrowTypes/README.md new file mode 100644 index 00000000..93ddecb6 --- /dev/null +++ b/src/ArrowTypes/README.md @@ -0,0 +1,43 @@ + + +# ArrowTypes.jl + +The type-mapping interface for serializing custom Julia types in the Arrow +format, kept as its own package (registered separately from this +subdirectory of the arrow-julia repository) so that packages can declare +how their types map to Arrow without depending on Arrow.jl itself. + +* `ArrowTypes.ArrowKind(T)` — the general category of Arrow type a Julia + type is treated as (`PrimitiveKind`, `ListKind`, `StructKind`, …). A + consumer decides which categories and layouts it supports. +* `ArrowTypes.ArrowType(T)` / `ArrowTypes.toarrow(x)` — the natively + supported type a value is converted to for serialization. +* `ArrowTypes.arrowname(T)`, `ArrowTypes.arrowmetadata(T)`, + `ArrowTypes.JuliaType(Val(name), S, meta)` and `ArrowTypes.fromarrow` — + the round trip back to the custom type through Arrow extension-type + metadata. + +Arrow.jl 2.x and 3.x consume this stable 2.x interface. Arrow.jl 3.x applies +`ArrowType` and `toarrow` recursively for custom values; defining `ArrowKind` +alone does not select an arbitrary Arrow.jl 3.x physical layout. Packages that +own custom types should depend on and import ArrowTypes.jl directly. Arrow.jl +keeps `ArrowTypes` exported as a compatibility exception, and +`Arrow.ArrowTypes` remains available as the qualified binding. See the +docstrings of the functions above for the contract. diff --git a/src/ArrowTypes/src/ArrowTypes.jl b/src/ArrowTypes/src/ArrowTypes.jl index fe2223f4..29aa166c 100644 --- a/src/ArrowTypes/src/ArrowTypes.jl +++ b/src/ArrowTypes/src/ArrowTypes.jl @@ -15,8 +15,10 @@ # limitations under the License. """ -The ArrowTypes module provides the [`ArrowTypes.Arrowtype`](@ref) interface trait that objects can define -in order to signal how they should be serialized in the arrow format. +The ArrowTypes module defines traits and hooks that packages can implement to +describe how their Julia types map to the Arrow format. Arrow consumers, +including Arrow.jl 2.x and 3.x, use this interface without requiring the +package that owns a custom type to depend on Arrow.jl. """ module ArrowTypes @@ -41,24 +43,25 @@ export ArrowKind, """ ArrowTypes.ArrowKind(T) -For a give type `T`, define it's "arrow type kind", or the general category of arrow types it should be treated as. Must be one of: +For a given type `T`, define its "arrow type kind", or the general category of arrow types it should be treated as. Must be one of: * [`ArrowTypes.NullKind`](@ref): `Missing` is the only type defined as `NullKind` - * [`ArrowTypes.PrimitiveKind`](@ref): `<:Integer`, `<:AbstractFloat`, along with `Arrow.Decimal`, and the various `Arrow.ArrowTimeType` subtypes + * [`ArrowTypes.PrimitiveKind`](@ref): `<:Integer`, `<:AbstractFloat`, along with decimal and temporal types * [`ArrowTypes.BoolKind`](@ref): only `Bool` * [`ArrowTypes.ListKind`](@ref): any `AbstractString` or `AbstractArray` - * [`ArrowTypes.FixedSizeList`](@ref): `NTuple{N, T}` + * [`ArrowTypes.FixedSizeListKind`](@ref): `NTuple{N, T}` * [`ArrowTypes.MapKind`](@ref): any `AbstractDict` * [`ArrowTypes.StructKind`](@ref): any `NamedTuple` or plain struct (mutable or otherwise) * [`ArrowTypes.UnionKind`](@ref): any `Union` - * [`ArrowTypes.DictEncodedKind`](@ref): array types that implement the `DataAPI.refpool` interface - -The list of `ArrowKind`s listed above translate to different ways to physically store data as supported by the arrow data format. -See the docs for each for an idea of whether they might be an appropriate fit for a custom type. -Note that custom types need to satisfy any additional "interface methods" as required by the various `ArrowKind` -types. By default, if a type in julia is declared like `primitive type ...` it is considered a `PrimitiveKind` -and if `struct` or `mutable struct` it's considered a `StructKind`. Also note that types will rarely need to define `ArrowKind`; -much more common is to define `ArrowType(T)` and `toarrow(x::T)` to transform `T` to a natively supported arrow type, which will -already have its `ArrowKind` defined. + * [`ArrowTypes.DictEncodedKind`](@ref): a kind a consumer may map its own dictionary-encoded array types to; ArrowTypes defines no such mapping itself + +The `ArrowKind`s describe general Arrow storage categories. Each consumer +decides which categories, layouts, and extra interface methods it supports. +By default, a Julia `primitive type` is a `PrimitiveKind`, and a `struct` or +`mutable struct` is a `StructKind`. Types rarely need to define `ArrowKind`. +It is more common to define `ArrowType(T)` and `toarrow(x::T)` to lower `T` to +a natively supported Arrow type, which already has an `ArrowKind`. In +particular, an `ArrowKind` override alone does not select an arbitrary +physical layout in a consumer. """ abstract type ArrowKind end @@ -72,9 +75,9 @@ Interface method to define the natively supported arrow type `S` that a given ty Useful when a custom type wants a "serialization hook" or otherwise needs to be transformed/converted into a natively supported arrow type for serialization. If a type defines `ArrowType`, it must also define a corresponding [`ArrowTypes.toarrow(x::T)`](@ref) method which does the actual conversion from `T` to `S`. -Note that custom structs defined like `struct T` or `mutable struct T` are natively supported in serialization, so unless -_additional_ transformation/customization is desired, a custom type `T` can serialize with no `ArrowType` definition (by default, -each field of a struct is serialized, using the results of `fieldnames(T)` and `getfield(x, i)`). +Some consumers can serialize plain structs by discovering their fields; check +the consumer's supported write types. Define `ArrowType` and `toarrow` when a +custom type needs a different storage representation. Note that defining these methods only deal with custom _serialization_ to the arrow format; to be able to _deserialize_ custom types at all, see the docs for [`ArrowTypes.arrowname`](@ref), [`ArrowTypes.arrowmetadata`](@ref), [`ArrowTypes.JuliaType`](@ref), and [`ArrowTypes.fromarrow`](@ref). @@ -142,15 +145,18 @@ arrowmetadata(::Type{Any}) = EMPTY_STRING ArrowTypes.JuliaType(::Val{Symbol(name)}, ::Type{S}, arrowmetadata::String) = T Interface method to define the custom Julia logical type `T` that a serialized metadata label should be converted to when -deserializing. When reading arrow data, and a logical type label is encountered for a column, it will call -`ArrowTypes.JuliaType(Val(Symbol(name)), S, arrowmetadata)` to see if a Julia type has been "registered" for deserialization. The `name` +deserializing. To resolve a logical type label, a reader calls +`ArrowTypes.JuliaType(Val(name_symbol), S, arrowmetadata)` with the corresponding existing Julia `Symbol`. The `name` used when defining the method *must* correspond to the same `name` when defining `ArrowTypes.arrowname(::Type{T}) = Symbol(name)`. The use of `Val(Symbol(...))` is to allow overloading a method on a specific logical type label. The `S` 2nd argument passed to `JuliaType` is the native arrow serialized type. This can be useful for parametric Julia types that wish to correctly parameterize their custom type based on what was serialized. The 3rd argument `arrowmetadata` is any metadata that was stored when the logical type was serialized as the result of calling `ArrowTypes.arrowmetadata(T)`. Note the 2nd and 3rd arguments are optional when overloading if unneeded. -When defining [`ArrowTypes.arrowname`](@ref) and `ArrowTypes.JuliaType`, you may also want to implement [`ArrowTypes.fromarrow`] +Readers must not intern an unknown, input-controlled name only to probe this interface: resolve the label with a non-interning +lookup and call `JuliaType(Val(...))` only when the corresponding `Symbol` already exists, preserving ordinary Arrow storage +values otherwise. +When defining [`ArrowTypes.arrowname`](@ref) and `ArrowTypes.JuliaType`, you may also want to implement [`ArrowTypes.fromarrow`](@ref) in order to customize how a custom type `T` should be constructed from the native arrow data type. See its docs for more details. """ function JuliaType end @@ -188,7 +194,7 @@ fromarrow(::Type{Union{Missing,T}}, x::T) where {T} = x fromarrow(::Type{Union{Missing,T}}, x::T) where {T<:NamedTuple} = x # ambiguity fix fromarrow(::Type{Union{Missing,T}}, x) where {T} = fromarrow(T, x) -"NullKind data is actually not physically stored since the data is constant; just the length is needed" +"NullKind data is not stored physically — the value is constant, so only the length is needed." struct NullKind <: ArrowKind end ArrowKind(::Type{Missing}) = NullKind() @@ -213,11 +219,11 @@ arrowname(::Type{Char}) = CHAR JuliaType(::Val{CHAR}) = Char fromarrow(::Type{Char}, x::UInt32) = Char(x) -"BoolKind data is stored with values packed down to individual bits; so instead of a traditional Bool being 1 byte/8 bits, 8 Bool values would be packed into a single byte" +"BoolKind data is bit-packed: 8 values per byte, not 1 byte per value." struct BoolKind <: ArrowKind end ArrowKind(::Type{Bool}) = BoolKind() -"ListKind data are stored in two separate buffers; one buffer contains all the original data elements flattened into one long buffer; the 2nd buffer contains an offset into the 1st buffer for how many elements make up the original array element" +"ListKind data is stored in two buffers: one holds every element flattened; the other holds the start offset of each list, so element i spans offsets[i]:offsets[i+1]-1." struct ListKind{stringtype} <: ArrowKind end ListKind() = ListKind{false}() @@ -225,7 +231,7 @@ isstringtype(::ListKind{stringtype}) where {stringtype} = stringtype isstringtype(::Type{ListKind{stringtype}}) where {stringtype} = stringtype ArrowKind(::Type{<:AbstractString}) = ListKind{true}() -# Treate Base.CodeUnits as Binary arrow type +# Treat Base.CodeUnits as a Binary arrow type ArrowKind(::Type{<:Base.CodeUnits}) = ListKind{true}() fromarrow(::Type{T}, ptr::Ptr{UInt8}, len::Int) where {T} = @@ -239,6 +245,8 @@ toarrow(x::Symbol) = String(x) const SYMBOL = Symbol("JuliaLang.Symbol") arrowname(::Type{Symbol}) = SYMBOL JuliaType(::Val{SYMBOL}) = Symbol +# Arrow.jl's IPC reader checks that a Symbol payload is already interned before +# it calls this compatibility hook. Direct ArrowTypes callers remain trusted. _symbol(ptr, len) = ccall(:jl_symbol_n, Ref{Symbol}, (Ptr{UInt8}, Int), ptr, len) fromarrow(::Type{Symbol}, ptr::Ptr{UInt8}, len::Int) = _symbol(ptr, len) @@ -285,6 +293,8 @@ arrowname(::Type{IPv6}) = IPV6_SYMBOL JuliaType(::Val{IPV6_SYMBOL}) = IPv6 fromarrow(::Type{IPv6}, x::NTuple{16,UInt8}) = IPv6(_cast(UInt128, x)) +# Not a plain reinterpret: the Ref round trip keeps the UInt128 ↔ +# NTuple{16,UInt8} conversion allocation- and trim-safe without a bitcast. function _cast(::Type{Y}, x)::Y where {Y} y = Ref{Y}() _unsafe_cast!(y, Ref(x), 1) @@ -306,6 +316,16 @@ struct StructKind <: ArrowKind end ArrowKind(::Type{<:NamedTuple}) = StructKind() +""" + ArrowTypes.fromarrowstruct(::Type{T}, ::Val{fnames}, x...) => T + +Optional `StructKind` deserialization hook. The field values `x` are passed +together with their serialized field names `fnames` (a tuple of Symbols), so +`T` can be reconstructed agnostic to the field order the serializer used. +When a method is defined for `T`, it takes precedence over +[`ArrowTypes.fromarrow`](@ref) in `StructKind` deserialization; the default +forwards to `fromarrow(T, x...)`. +""" @inline fromarrowstruct(T::Type, ::Val, x...) = fromarrow(T, x...) fromarrow( @@ -335,7 +355,7 @@ function fromarrow(::Type{VersionNumber}, v::NamedTuple) VersionNumber(v.major, v.minor, v.patch, v.prerelease, v.build) end -"MapKind data are stored similarly to ListKind, where elements are flattened, and a 2nd offsets buffer contains the individual list element length data" +"MapKind data is stored like ListKind: flattened key-value entries plus an offsets buffer delimiting each map." struct MapKind <: ArrowKind end ArrowKind(::Type{<:AbstractDict}) = MapKind() @@ -345,7 +365,7 @@ struct UnionKind <: ArrowKind end ArrowKind(::Union) = UnionKind() -"DictEncodedKind store a small pool of unique values in one buffer, with a full-length buffer of integer offsets into the small value pool" +"`DictEncodedKind` stores a category pool in one buffer and a full-length buffer of integer indices into it. A category pool may contain unused or duplicate physical entries." struct DictEncodedKind <: ArrowKind end """ @@ -392,7 +412,13 @@ function promoteunion(T, S) return isabstracttype(new) ? Union{T,S} : new end -# lazily call toarrow(x) on getindex for each x in data +""" + ArrowTypes.ToArrow(x) -> AbstractVector + +A lazy view over `x` that applies [`ArrowTypes.toarrow`](@ref) on `getindex`, +with a concrete element type. Returns `x` itself when its element type is +already a concrete natively supported arrow type indexed from 1. +""" struct ToArrow{T,A} <: AbstractVector{T} data::A end @@ -423,6 +449,9 @@ end Base.IndexStyle(::Type{<:ToArrow}) = Base.IndexLinear() Base.size(x::ToArrow) = (length(x.data),) Base.eltype(::Type{TA}) where {T,A,TA<:ToArrow{T,A}} = T +# A conversion failure means "try the other Union branch"; anything else +# (including InterruptException) must propagate. +const _CONVERSION_ERRORS = Union{MethodError,InexactError,OverflowError,TypeError} function _convert(::Type{T}, x) where {T} if x isa T return x @@ -430,10 +459,10 @@ function _convert(::Type{T}, x) where {T} # T was a promoted Union and x is not already one of # the concrete Union types, so we need to just try # to convert, recursively, to one of the Union types - # unfortunately not much we can do more efficiently here try return _convert(T.a, x) - catch + catch err + err isa _CONVERSION_ERRORS || rethrow() return _convert(T.b, x) end else diff --git a/src/ArrowTypes/test/tests.jl b/src/ArrowTypes/test/tests.jl index 22d8dd0e..083c2456 100644 --- a/src/ArrowTypes/test/tests.jl +++ b/src/ArrowTypes/test/tests.jl @@ -22,6 +22,21 @@ struct Person name::String end +# Keep these fixtures at top level. Julia 1.0 does not permit type definitions +# inside the local scope introduced by `@testset`. +struct DateTimeTZ + instant::Int64 + tz::String +end + +struct Timestamp{TZ} + x::Int64 +end + +ArrowTypes.ArrowType(::Type{DateTimeTZ}) = Timestamp +ArrowTypes.toarrow(x::DateTimeTZ) = Timestamp{Symbol(x.tz)}(x.instant) +ArrowTypes.default(::Type{DateTimeTZ}) = DateTimeTZ(0, "UTC") + @testset "ArrowTypes" begin @test ArrowTypes.ArrowKind(MyInt) == ArrowTypes.PrimitiveKind() @test ArrowTypes.ArrowKind(Person) == ArrowTypes.StructKind() @@ -203,19 +218,6 @@ end @test x == [1, 3.14] @testset "respect non-missing concrete type" begin - struct DateTimeTZ - instant::Int64 - tz::String - end - - struct Timestamp{TZ} - x::Int64 - end - - ArrowTypes.ArrowType(::Type{DateTimeTZ}) = Timestamp - ArrowTypes.toarrow(x::DateTimeTZ) = Timestamp{Symbol(x.tz)}(x.instant) - ArrowTypes.default(::Type{DateTimeTZ}) = DateTimeTZ(0, "UTC") - T = Union{DateTimeTZ,Missing} @test !ArrowTypes.concrete_or_concreteunion(ArrowTypes.ArrowType(T)) @test eltype(ArrowTypes.ToArrow(T[missing])) == Union{Timestamp{:UTC},Missing} diff --git a/src/FlatBuffers/builder.jl b/src/FlatBuffers/builder.jl index 1ca11986..c3649994 100644 --- a/src/FlatBuffers/builder.jl +++ b/src/FlatBuffers/builder.jl @@ -14,11 +14,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -const fileIdentifierLength = 4 - """ Scalar -A Union of the Julia types `T <: Number` that are allowed in FlatBuffers schema +The Julia scalar types allowed in a FlatBuffers schema: the fixed-width +integers and floats, `Bool`, and `Enum`. """ const Scalar = Union{Bool,Int8,Int16,Int32,Int64,UInt8,UInt16,UInt32,UInt64,Float32,Float64,Enum} @@ -39,7 +38,6 @@ mutable struct Builder head::UOffsetT nested::Bool finished::Bool - sharedstrings::Dict{String,UOffsetT} end bytes(b::Builder) = getfield(b, :bytes) @@ -53,21 +51,8 @@ Builder(size=0) = Builder( UOffsetT(size), false, false, - Dict{String,UOffsetT}(), ) -function reset!(b::Builder) - empty!(b.bytes) - empty!(b.vtable) - emtpy!(b.vtables) - empty!(b.sharedstrings) - b.minalign = 1 - b.nested = false - b.finished = false - b.head = 0 - return -end - Base.write(sink::Builder, o, x::Union{Bool,UInt8}) = sink.bytes[o + 1] = UInt8(x) function Base.write(sink::Builder, off, x::T) where {T} off += 1 @@ -80,9 +65,8 @@ Base.write(b::Builder, o, x::Float64) = write(b, o, reinterpret(UInt64, x)) Base.write(b::Builder, o, x::Enum) = write(b, o, basetype(x)(x)) """ -`finishedbytes` returns a pointer to the written data in the byte buffer. -Panics if the builder is not in a finished state (which is caused by calling -`finish!()`). +`finishedbytes` returns a view of the written data in the byte buffer. It +throws unless `finish!` has put the builder in a finished state. """ function finishedbytes(b::Builder) assertfinished(b) @@ -99,7 +83,7 @@ function startobject!(b::Builder, numfields) end """ -WriteVtable serializes the vtable for the current object, if applicable. +`writevtable!` serializes the vtable for the current object, if applicable. Before writing out the vtable, this checks pre-existing vtables for equality to this one. If an equal vtable is found, point the object to the existing @@ -134,10 +118,7 @@ function writevtable!(b::Builder) end # Search backwards through existing vtables, because similar vtables - # are likely to have been recently appended. See - # BenchmarkVtableDeduplication for a case in which this heuristic - # saves about 30% of the time used in writing objects with duplicate - # tables. + # are likely to have been recently appended. for i = length(b.vtables):-1:1 # Find the other vtable, which is associated with `i`: vt2Offset = b.vtables[i] @@ -146,7 +127,8 @@ function writevtable!(b::Builder) metadata = VtableMetadataFields * sizeof(VOffsetT) vt2End = vt2Start + vt2Len - vt2 = view(b.bytes, (vt2Start + metadata + 1):vt2End) #TODO: might need a +1 on the start of range here + # The field entries only (past the two metadata VOffsetTs). + vt2 = view(b.bytes, (vt2Start + metadata + 1):vt2End) # Compare the other vtable to the one under consideration. # If they are equal, store the offset and break: @@ -205,7 +187,7 @@ function writevtable!(b::Builder) end """ -`endobject` writes data necessary to finish object construction. +`endobject!` writes data necessary to finish object construction. """ function endobject!(b::Builder) assertnested(b) @@ -270,7 +252,7 @@ function prependoffsetslot!(b::Builder, o::Int, x::T, d) where {T} end """ -`startvector` initializes bookkeeping for writing a new vector. +`startvector!` initializes bookkeeping for writing a new vector. A vector has the following format: @@ -285,7 +267,7 @@ function startvector!(b::Builder, elemSize, numElems, alignment) end """ -`endvector` writes data necessary to finish vector construction. +`endvector!` writes data necessary to finish vector construction. """ function endvector!(b::Builder, vectorNumElems) assertnested(b) @@ -294,12 +276,6 @@ function endvector!(b::Builder, vectorNumElems) return offset(b) end -function createsharedstring!(b::Builder, s::AbstractString) - get!(b.sharedstrings, s) do - createstring!(b, s) - end -end - """ `createstring!` writes a null-terminated string as a vector. """ @@ -317,8 +293,6 @@ function createstring!(b::Builder, s::Union{AbstractString,AbstractVector{UInt8} return endvector!(b, sizeof(s)) end -createbytevector(b::Builder, v) = createstring!(b, v) - function assertnested(b::Builder) # If you get this assert, you're in an object while trying to write # data that belongs outside of an object. @@ -361,7 +335,7 @@ end If value `x` equals default `d`, then the slot will be set to zero and no other data will be written. """ -function prependslot!(b::Builder, o::Int, x::T, d, sh=false) where {T<:Scalar} +function prependslot!(b::Builder, o::Int, x::T, d) where {T<:Scalar} if x != T(d) prepend!(b, x) slot!(b, o) @@ -369,22 +343,6 @@ function prependslot!(b::Builder, o::Int, x::T, d, sh=false) where {T<:Scalar} return end -""" -`prependstructslot!` prepends a struct onto the object at vtable slot `o`. -Structs are stored inline, so nothing additional is being added. -In generated code, `d` is always 0. -""" -function prependstructslot!(b::Builder, voffset, x, d) - if x != d - assertnested(b) - if x != offset(b) - throw(ArgumentError("inline data write outside of object")) - end - slot!(b, voffset) - end - return -end - """ `slot!` sets the vtable key `voffset` to the current location in the buffer. """ @@ -392,23 +350,6 @@ function slot!(b::Builder, slotnum) b.vtable[slotnum + 1] = offset(b) end -# FinishWithFileIdentifier finalizes a buffer, pointing to the given `rootTable`. -# as well as applys a file identifier -function finishwithfileidentifier(b::Builder, rootTable, fid) - if length(fid) != fileIdentifierLength - error("incorrect file identifier length") - end - # In order to add a file identifier to the flatbuffer message, we need - # to prepare an alignment and file identifier length - prep!(b, b.minalign, sizeof(Int32) + fileIdentifierLength) - for i = fileIdentifierLength:-1:1 - # place the file identifier - place!(b, fid[i]) - end - # finish - finish!(b, rootTable) -end - """ `finish!` finalizes a buffer, pointing to the given `rootTable`. """ @@ -427,7 +368,8 @@ function vtableEqual(a::Vector{UOffsetT}, objectStart, b::AbstractVector{UInt8}) end for i = 0:(length(a) - 1) - x = read(IOBuffer(view(b, (i * sizeof(VOffsetT) + 1):length(b))), VOffsetT) + base = i * sizeof(VOffsetT) + x = VOffsetT(b[base + 1]) | (VOffsetT(b[base + 2]) << 8) # Skip vtable entries that indicate a default value. x == 0 && a[i + 1] == 0 && continue diff --git a/src/FlatBuffers/table.jl b/src/FlatBuffers/table.jl index a17d1cd0..0a42b6b9 100644 --- a/src/FlatBuffers/table.jl +++ b/src/FlatBuffers/table.jl @@ -21,8 +21,10 @@ The object containing the flatbuffer and positional information specific to the The `vtable` containing the offsets for specific members precedes `pos`. The actual values in the table follow `pos` offset and size of the vtable. +Concrete subtypes declare: + - `bytes::Vector{UInt8}`: the flatbuffer itself -- `pos::Integer`: the base position in `bytes` of the table +- `pos::Base.Int`: the base position in `bytes` of the table """ abstract type Table end abstract type Struct end @@ -35,9 +37,6 @@ pos(x::TableOrStruct) = getfield(x, :pos) ==(a::T, b::T) where {T<:TableOrStruct} = all(getproperty(a, p) == getproperty(b, p) for p in propertynames(a)) -(::Type{T})(b::Builder) where {T<:TableOrStruct} = - T(b.bytes[(b.head + 1):end], get(b, b.head, Int32)) - getrootas(::Type{T}, bytes::Vector{UInt8}, offset) where {T<:Table} = init(T, bytes, offset + readbuffer(bytes, offset, UOffsetT)) init(::Type{T}, bytes::Vector{UInt8}, pos::Integer) where {T<:TableOrStruct} = T(bytes, pos) @@ -61,10 +60,6 @@ end "`indirect` retrieves the relative offset stored at `offset`." indirect(t::Table, off) = off + get(t, off, UOffsetT) -getvalue(t, o, ::Type{Nothing}) = nothing -getvalue(t, o, ::Type{T}) where {T<:Scalar} = get(t, pos(t) + o, T) -getvalue(t, o, ::Type{T}) where {T<:Enum} = T(get(t, pos(t) + o, enumtype(T))) - function Base.String(t::Table, off) off += get(t, off, UOffsetT) start = off + sizeof(UOffsetT) @@ -72,13 +67,6 @@ function Base.String(t::Table, off) return unsafe_string(pointer(bytes(t), start + 1), len) end -function bytevector(t::Table, off) - off += get(t, off, UOffsetT) - start = off + sizeof(UOffsetT) - len = get(t, off, UOffsetT) - return view(bytes(t), (start + 1):(start + len + 1)) -end - """ `vectorlen` retrieves the length of the vector whose offset is stored at `off` in this object. @@ -128,49 +116,7 @@ Base.@propagate_inbounds function Base.getindex(A::Array{T,S}, i::Integer) where end end -Base.@propagate_inbounds function Base.setindex!(A::Array{T,S}, v, i::Integer) where {T,S} - if T === S - return setindex!(A.data, v, i) - else - error("setindex! not supported for reference/table types") - end -end - function union(t::Table, off) off += pos(t) return off + get(t, off, UOffsetT) end - -function union!(t::Table, t2::Table, off) - off += pos(t) - t2.pos = off + get(t, off, UOffsetT) - t2.bytes = bytes(t) - return -end - -""" -GetVOffsetTSlot retrieves the VOffsetT that the given vtable location -points to. If the vtable value is zero, the default value `d` -will be returned. -""" -function getoffsetslot(t::Table, slot, d) - off = offset(t, slot) - if off == 0 - return d - end - return off -end - -""" -`getslot` retrieves the `T` that the given vtable location -points to. If the vtable value is zero, the default value `d` -will be returned. -""" -function getslot(t::Table, slot, d::T) where {T} - off = offset(t, slot) - if off == 0 - return d - end - - return get(t, pos(t) + off, T) -end diff --git a/src/append.jl b/src/append.jl deleted file mode 100644 index 1a5119e6..00000000 --- a/src/append.jl +++ /dev/null @@ -1,315 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT 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.append(io::IO, tbl) - Arrow.append(file::String, tbl) - tbl |> Arrow.append(file) - -Append any [Tables.jl](https://github.com/JuliaData/Tables.jl)-compatible `tbl` -to an existing arrow formatted file or IO. The existing arrow data must be in -IPC stream format. Note that appending to the "feather formatted file" is _not_ -allowed, as this file format doesn't support appending. That means files written -like `Arrow.write(filename::String, tbl)` _cannot_ be appended to; instead, you -should write like `Arrow.write(filename::String, tbl; file=false)`. - -When an IO object is provided to be written on to, it must support seeking. For -example, a file opened in `r+` mode or an `IOBuffer` that is readable, writable -and seekable can be appended to, but not a network stream. - -Multiple record batches will be written based on the number of -`Tables.partitions(tbl)` that are provided; by default, this is just -one for a given table, but some table sources support automatic -partitioning. Note you can turn multiple table objects into partitions -by doing `Tables.partitioner([tbl1, tbl2, ...])`, but note that -each table must have the exact same `Tables.Schema`. - -By default, `Arrow.append` will use multiple threads to write multiple -record batches simultaneously (e.g. if julia is started with `julia -t 8` -or the `JULIA_NUM_THREADS` environment variable is set). - -Supported keyword arguments to `Arrow.append` include: - * `alignment::Int=8`: specify the number of bytes to align buffers to when written in messages; strongly recommended to only use alignment values of 8 or 64 for modern memory cache line optimization - * `colmetadata=nothing`: the metadata that should be written as the table's columns' `custom_metadata` fields; must either be `nothing` or an `AbstractDict` of `column_name::Symbol => column_metadata` where `column_metadata` is an iterable of `<:AbstractString` pairs. - * `dictencode::Bool=false`: whether all columns should use dictionary encoding when being written; to dict encode specific columns, wrap the column/array in `Arrow.DictEncode(col)` - * `dictencodenested::Bool=false`: whether nested data type columns should also dict encode nested arrays/buffers; other language implementations [may not support this](https://arrow.apache.org/docs/status.html) - * `denseunions::Bool=true`: whether Julia `Vector{<:Union}` arrays should be written using the dense union layout; passing `false` will result in the sparse union layout - * `largelists::Bool=false`: causes list column types to be written with Int64 offset arrays; mainly for testing purposes; by default, Int64 offsets will be used only if needed - * `maxdepth::Int=$DEFAULT_MAX_DEPTH`: deepest allowed nested serialization level; this is provided by default to prevent accidental infinite recursion with mutually recursive data structures - * `metadata=Arrow.getmetadata(tbl)`: the metadata that should be written as the table's schema's `custom_metadata` field; must either be `nothing` or an iterable of `<:AbstractString` pairs. - * `ntasks::Int`: number of concurrent threaded tasks to allow while writing input partitions out as arrow record batches; default is no limit; to disable multithreaded writing, pass `ntasks=1` - * `convert::Bool`: whether certain arrow primitive types in the schema of `file` should be converted to Julia defaults for matching them to the schema of `tbl`; by default, `convert=true`. - * `file::Bool`: applicable when an `IO` is provided, whether it is a file; by default `file=false`. -""" -function append end - -append(io_or_file; kw...) = x -> append(io_or_file, x; kw...) - -function append(file::String, tbl; kwargs...) - open(file, isfile(file) ? "r+" : "w+") do io - append(io, tbl; file=true, kwargs...) - end - - return file -end - -function append( - io::IO, - tbl; - metadata=getmetadata(tbl), - colmetadata=nothing, - largelists::Bool=false, - denseunions::Bool=true, - dictencode::Bool=false, - dictencodenested::Bool=false, - alignment::Int=8, - maxdepth::Int=DEFAULT_MAX_DEPTH, - ntasks=Inf, - convert::Bool=true, - file::Bool=false, -) - if ntasks < 1 - throw( - ArgumentError( - "ntasks keyword argument must be > 0; pass `ntasks=1` to disable multithreaded writing", - ), - ) - end - - startpos = position(io) - seekend(io) - len = position(io) - startpos - seek(io, startpos) # leave the stream position unchanged - - if len == 0 # empty file, not initialized, we can just write to it - kwargs = Dict{Symbol,Any}( - :largelists => largelists, - :denseunions => denseunions, - :dictencode => dictencode, - :dictencodenested => dictencodenested, - :alignment => alignment, - :maxdepth => maxdepth, - :metadata => metadata, - :colmetadata => colmetadata, - ) - if isa(ntasks, Integer) - kwargs[:ntasks] = ntasks - end - write(io, tbl; kwargs...) - else - isstream, arrow_schema, compress = stream_properties(io; convert=convert) - if !isstream - throw(ArgumentError("append is supported only to files in arrow stream format")) - end - if compress isa Symbol && compress !== :lz4 && compress !== :zstd - throw( - ArgumentError( - "unsupported compress keyword argument value: $compress. Valid values include `:lz4` or `:zstd`", - ), - ) - end - append( - io, - tbl, - arrow_schema, - compress, - largelists, - denseunions, - dictencode, - dictencodenested, - alignment, - maxdepth, - ntasks, - metadata, - colmetadata, - ) - end - - return io -end - -function append( - io::IO, - source, - arrow_schema, - compress, - largelists, - denseunions, - dictencode, - dictencodenested, - alignment, - maxdepth, - ntasks, - meta, - colmeta, -) - seekend(io) - skip(io, -8) # overwrite last 8 bytes of last empty message footer - - sch = Ref{Tables.Schema}(arrow_schema) - sync = OrderedSynchronizer() - msgs = Channel{Message}(ntasks) - dictencodings = Dict{Int64,Any}() # Lockable{DictEncoding} - # build messages - blocks = (Block[], Block[]) - # start message writing from channel - threaded = ntasks > 1 - tsk = - threaded ? (@wkspawn for msg in msgs - Base.write(io, msg, blocks, sch, alignment) - end) : (@async for msg in msgs - Base.write(io, msg, blocks, sch, alignment) - end) - anyerror = Threads.Atomic{Bool}(false) - errorref = Ref{Any}() - @sync for (i, tbl) in enumerate(Tables.partitions(source)) - if anyerror[] - @error "error writing arrow data on partition = $(errorref[][3])" exception = - (errorref[][1], errorref[][2]) - error("fatal error writing arrow data") - end - @debug "processing table partition i = $i" - tbl_cols = Tables.columns(tbl) - tbl_schema = Tables.schema(tbl_cols) - - if !is_equivalent_schema(arrow_schema, tbl_schema) - throw(ArgumentError("Table schema does not match existing arrow file schema")) - end - - if threaded - @wkspawn process_partition( - tbl_cols, - dictencodings, - largelists, - compress, - denseunions, - dictencode, - dictencodenested, - maxdepth, - sync, - msgs, - alignment, - i, - sch, - errorref, - anyerror, - meta, - colmeta, - ) - else - @async process_partition( - tbl_cols, - dictencodings, - largelists, - compress, - denseunions, - dictencode, - dictencodenested, - maxdepth, - sync, - msgs, - alignment, - i, - sch, - errorref, - anyerror, - meta, - colmeta, - ) - end - end - if anyerror[] - @error "error writing arrow data on partition = $(errorref[][3])" exception = - (errorref[][1], errorref[][2]) - error("fatal error writing arrow data") - end - # close our message-writing channel, no further put!-ing is allowed - close(msgs) - # now wait for our message-writing task to finish writing - wait(tsk) - - Base.write( - io, - Message(UInt8[], nothing, 0, true, false, Meta.Schema), - blocks, - sch, - alignment, - ) - - return io -end - -function stream_properties(io::IO; convert::Bool=true) - startpos = position(io) - buff = similar(FILE_FORMAT_MAGIC_BYTES) - start_magic = read!(io, buff) == FILE_FORMAT_MAGIC_BYTES - seekend(io) - len = position(io) - startpos - skip(io, -length(FILE_FORMAT_MAGIC_BYTES)) - end_magic = read!(io, buff) == FILE_FORMAT_MAGIC_BYTES - seek(io, startpos) # leave the stream position unchanged - - isstream = !(len > 24 && start_magic && end_magic) - if isstream - stream = Stream(io, convert=convert) - for table in stream - # no need to scan further once we get compression information - (stream.compression[] !== nothing) && break - end - seek(io, startpos) # leave the stream position unchanged - return isstream, Tables.Schema(stream.names, stream.types), stream.compression[] - else - return isstream, nothing, nothing - end -end - -function is_equivalent_schema(sch1::Tables.Schema, sch2::Tables.Schema) - (sch1.names == sch2.names) || (return false) - for (t1, t2) in zip(sch1.types, sch2.types) - tt1 = Base.nonmissingtype(t1) - tt2 = Base.nonmissingtype(t2) - if t1 == t2 - continue - elseif tt1 <: AbstractVector && tt2 <: AbstractVector && eltype(tt1) == eltype(tt2) - continue - elseif isstructtype(tt1) && isstructtype(tt2) - is_equivalent_type_by_field(tt1, tt2) - else - return false - end - end - true -end - -function is_equivalent_type_by_field(T1, T2) - n1 = fieldcount(T1) - n2 = fieldcount(T2) - n1 != n2 && return false - - for i = 1:n1 - fieldname(T1, i) == fieldname(T2, i) || return false - - if fieldtype(T1, i) == fieldtype(T2, i) - continue - elseif isstructtype(T1) && isstructtype(T2) - is_equivalent_type_by_field(T1, T2) || continue - else - return false - end - end - true -end diff --git a/src/arraytypes/arraytypes.jl b/src/arraytypes/arraytypes.jl deleted file mode 100644 index 58bab082..00000000 --- a/src/arraytypes/arraytypes.jl +++ /dev/null @@ -1,274 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT 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.ArrowVector - -An abstract type that subtypes `AbstractVector`. Each specific arrow array type -subtypes `ArrowVector`. See [`BoolVector`](@ref), [`Primitive`](@ref), [`List`](@ref), -[`Map`](@ref), [`FixedSizeList`](@ref), [`Struct`](@ref), [`DenseUnion`](@ref), -[`SparseUnion`](@ref), and [`DictEncoded`](@ref) for more details. -""" -abstract type ArrowVector{T} <: AbstractVector{T} end - -Base.IndexStyle(::Type{A}) where {A<:ArrowVector} = Base.IndexLinear() -Base.similar(::Type{A}, dims::Dims) where {T,A<:ArrowVector{T}} = Vector{T}(undef, dims) -validitybitmap(x::ArrowVector) = x.validity -nullcount(x::ArrowVector) = validitybitmap(x).nc -getmetadata(x::ArrowVector) = x.metadata -Base.deleteat!(x::T, inds) where {T<:ArrowVector} = throw( - ArgumentError("`$T` does not support `deleteat!`; arrow data is by nature immutable"), -) - -function toarrowvector( - x, - i=1, - de=Dict{Int64,Any}(), - ded=DictEncoding[], - meta=getmetadata(x); - compression::Union{Nothing,Symbol,LZ4FrameCompressor,ZstdCompressor}=nothing, - kw..., -) - @debug "converting top-level column to arrow format: col = $(typeof(x)), compression = $compression, kw = $(values(kw))" - @debug x - A = arrowvector(x, i, 0, 0, de, ded, meta; compression=compression, kw...) - if compression isa LZ4FrameCompressor - A = compress(Meta.CompressionType.LZ4_FRAME, compression, A) - elseif compression isa ZstdCompressor - A = compress(Meta.CompressionType.ZSTD, compression, A) - elseif compression isa Symbol && compression == :lz4 - comp = lz4_frame_compressor() - A = Base.@lock comp begin - compress(Meta.CompressionType.LZ4_FRAME, comp[], A) - end - elseif compression isa Symbol && compression == :zstd - comp = zstd_compressor() - A = Base.@lock comp begin - compress(Meta.CompressionType.ZSTD, comp[], A) - end - end - @debug "converted top-level column to arrow format: $(typeof(A))" - @debug A - return A -end - -function arrowvector( - x, - i, - nl, - fi, - de, - ded, - meta; - dictencoding::Bool=false, - dictencode::Bool=false, - maxdepth::Int=DEFAULT_MAX_DEPTH, - kw..., -) - if nl > maxdepth - error( - "reached nested serialization level ($nl) deeper than provided max depth argument ($(maxdepth)); to increase allowed nesting level, pass `maxdepth=X`", - ) - end - T = maybemissing(eltype(x)) - if !(x isa DictEncode) && !dictencoding && (dictencode || DataAPI.refarray(x) !== x) - x = DictEncode(x, dictencodeid(i, nl, fi)) - elseif x isa DictEncoded - return arrowvector( - DictEncodeType, - x, - i, - nl, - fi, - de, - ded, - meta; - dictencode=dictencode, - kw..., - ) - elseif !(x isa DictEncode) - x = ToArrow(x) - end - S = maybemissing(eltype(x)) - if ArrowTypes.hasarrowname(T) - meta = _arrowtypemeta( - _normalizemeta(meta), - String(ArrowTypes.arrowname(T)), - String(ArrowTypes.arrowmetadata(T)), - ) - end - return arrowvector( - S, - x, - i, - nl, - fi, - de, - ded, - meta; - dictencode=dictencode, - maxdepth=maxdepth, - kw..., - ) -end - -_normalizemeta(::Nothing) = nothing -_normalizemeta(meta) = toidict(String(k) => String(v) for (k, v) in meta) - -_normalizecolmeta(::Nothing) = nothing -_normalizecolmeta(colmeta) = toidict( - Symbol(k) => toidict(String(v1) => String(v2) for (v1, v2) in v) for (k, v) in colmeta -) - -function _arrowtypemeta(::Nothing, n, m) - return toidict(("ARROW:extension:name" => n, "ARROW:extension:metadata" => m)) -end - -function _arrowtypemeta(meta, n, m) - dict = Dict(meta) - dict["ARROW:extension:name"] = n - dict["ARROW:extension:metadata"] = m - return toidict(dict) -end - -# now we check for ArrowType converions and dispatch on ArrowKind -function arrowvector(::Type{S}, x, i, nl, fi, de, ded, meta; kw...) where {S} - meta = _normalizemeta(meta) - return arrowvector(ArrowKind(S), x, i, nl, fi, de, ded, meta; kw...) -end - -struct NullVector{T} <: ArrowVector{T} - data::MissingVector - metadata::Union{Nothing,Base.ImmutableDict{String,String}} -end -Base.size(v::NullVector) = (length(v.data),) -Base.getindex(v::NullVector{T}, i::Int) where {T} = - ArrowTypes.fromarrow(T, getindex(v.data, i)) - -arrowvector(::NullKind, x, i, nl, fi, de, ded, meta; kw...) = NullVector{eltype(x)}( - MissingVector(length(x)), - isnothing(meta) ? nothing : toidict(meta), -) -compress(Z::Meta.CompressionType.T, comp, v::NullVector) = - Compressed{Z,NullVector}(v, CompressedBuffer[], length(v), length(v), Compressed[]) - -function makenodesbuffers!( - col::NullVector, - fieldnodes, - fieldbuffers, - bufferoffset, - alignment, -) - push!(fieldnodes, FieldNode(length(col), length(col))) - @debug "made field node: nodeidx = $(length(fieldnodes)), col = $(typeof(col)), len = $(fieldnodes[end].length), nc = $(fieldnodes[end].null_count)" - return bufferoffset -end - -function writebuffer(io, col::NullVector, alignment) - return -end - -""" - Arrow.ValidityBitmap - -A bit-packed array type where each bit corresponds to an element in an -[`ArrowVector`](@ref), indicating whether that element is "valid" (bit == 1), -or not (bit == 0). Used to indicate element missingness (whether it's null). - -If the null count of an array is zero, the `ValidityBitmap` will be "empty" -and all elements are treated as "valid"/non-null. -""" -struct ValidityBitmap <: ArrowVector{Bool} - bytes::Vector{UInt8} # arrow memory blob - pos::Int # starting byte of validity bitmap - ℓ::Int # # of _elements_ (not bytes!) in bitmap (because bitpacking) - nc::Int # null count -end - -Base.size(p::ValidityBitmap) = (p.ℓ,) -nullcount(x::ValidityBitmap) = x.nc - -function ValidityBitmap(x) - T = eltype(x) - if !(T >: Missing) - return ValidityBitmap(UInt8[], 1, length(x), 0) - end - len = length(x) - blen = cld(len, 8) - bytes = Vector{UInt8}(undef, blen) - st = iterate(x) - nc = 0 - b = 0xff - j = k = 1 - for y in x - if y === missing - nc += 1 - b = setbit(b, false, j) - end - j += 1 - if j == 9 - @inbounds bytes[k] = b - b = 0xff - j = 1 - k += 1 - end - end - if j > 1 - bytes[k] = b - end - return ValidityBitmap(nc == 0 ? UInt8[] : bytes, 1, nc == 0 ? 0 : len, nc) -end - -@propagate_inbounds function Base.getindex(p::ValidityBitmap, i::Integer) - # no boundscheck because parent array should do it - # if a validity bitmap is empty, it either means: - # 1) the parent array null_count is 0, so all elements are valid - # 2) parent array is also empty, so "all" elements are valid - p.nc == 0 && return true - # translate element index to bitpacked byte index - a, b = divrem(i - 1, 8) .+ (1, 1) - @inbounds byte = p.bytes[p.pos + a - 1] - # check individual bit of byte - return getbit(byte, b) -end - -@propagate_inbounds function Base.setindex!(p::ValidityBitmap, v, i::Integer) - x = convert(Bool, v) - p.ℓ == 0 && !x && throw(BoundsError(p, i)) - a, b = fldmod1(i, 8) - @inbounds byte = p.bytes[p.pos + a - 1] - @inbounds p.bytes[p.pos + a - 1] = setbit(byte, x, b) - return v -end - -function writebitmap(io, col::ArrowVector, alignment) - v = col.validity - @debug "writing validity bitmap: nc = $(v.nc), n = $(cld(v.ℓ, 8))" - v.nc == 0 && return 0 - n = Base.write(io, view(v.bytes, (v.pos):(v.pos + cld(v.ℓ, 8) - 1))) - return n + writezeros(io, paddinglength(n, alignment)) -end - -include("compressed.jl") -include("primitive.jl") -include("bool.jl") -include("list.jl") -include("fixedsizelist.jl") -include("map.jl") -include("struct.jl") -include("unions.jl") -include("dictencoding.jl") -include("views.jl") diff --git a/src/arraytypes/bool.jl b/src/arraytypes/bool.jl deleted file mode 100644 index 29c1505a..00000000 --- a/src/arraytypes/bool.jl +++ /dev/null @@ -1,117 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT 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.BoolVector - -A bit-packed array type, similar to [`ValidityBitmap`](@ref), but which -holds boolean values, `true` or `false`. -""" -struct BoolVector{T} <: ArrowVector{T} - arrow::Vector{UInt8} # need to hold a reference to arrow memory blob - pos::Int - validity::ValidityBitmap - ℓ::Int64 - metadata::Union{Nothing,Base.ImmutableDict{String,String}} -end - -Base.size(p::BoolVector) = (p.ℓ,) - -@propagate_inbounds function Base.getindex(p::BoolVector{T}, i::Integer) where {T} - @boundscheck checkbounds(p, i) - if T >: Missing - @inbounds !p.validity[i] && return missing - end - a, b = fldmod1(i, 8) - @inbounds byte = p.arrow[p.pos + a - 1] - # check individual bit of byte - return ArrowTypes.fromarrow(T, getbit(byte, b)) -end - -@propagate_inbounds function Base.setindex!(p::BoolVector, v, i::Integer) - @boundscheck checkbounds(p, i) - x = convert(Bool, v) - a, b = fldmod1(i, 8) - @inbounds byte = p.arrow[p.pos + a - 1] - @inbounds p.arrow[p.pos + a - 1] = setbit(byte, x, b) - return v -end - -arrowvector(::BoolKind, x::BoolVector, i, nl, fi, de, ded, meta; kw...) = x - -function arrowvector(::BoolKind, x, i, nl, fi, de, ded, meta; kw...) - validity = ValidityBitmap(x) - len = length(x) - blen = cld(len, 8) - bytes = Vector{UInt8}(undef, blen) - b = 0xff - j = k = 1 - for y in x - if y === false - b = setbit(b, false, j) - end - j += 1 - if j == 9 - @inbounds bytes[k] = b - b = 0xff - j = 1 - k += 1 - end - end - if j > 1 - bytes[k] = b - end - return BoolVector{eltype(x)}(bytes, 1, validity, len, meta) -end - -function compress(Z::Meta.CompressionType.T, comp, p::P) where {P<:BoolVector} - len = length(p) - nc = nullcount(p) - validity = compress(Z, comp, p.validity) - data = compress(Z, comp, view(p.arrow, (p.pos):(p.pos + cld(p.ℓ, 8) - 1))) - return Compressed{Z,P}(p, [validity, data], len, nc, Compressed[]) -end - -function makenodesbuffers!( - col::BoolVector, - fieldnodes, - fieldbuffers, - bufferoffset, - alignment, -) - len = length(col) - nc = nullcount(col) - push!(fieldnodes, FieldNode(len, nc)) - @debug "made field node: nodeidx = $(length(fieldnodes)), col = $(typeof(col)), len = $(fieldnodes[end].length), nc = $(fieldnodes[end].null_count)" - # validity bitmap - blen = nc == 0 ? 0 : bitpackedbytes(len, alignment) - push!(fieldbuffers, Buffer(bufferoffset, blen)) - @debug "made field buffer: bufferidx = $(length(fieldbuffers)), offset = $(fieldbuffers[end].offset), len = $(fieldbuffers[end].length), padded = $(padding(fieldbuffers[end].length, alignment))" - # adjust buffer offset, make primitive array buffer - bufferoffset += blen - blen = bitpackedbytes(len, alignment) - push!(fieldbuffers, Buffer(bufferoffset, blen)) - @debug "made field buffer: bufferidx = $(length(fieldbuffers)), offset = $(fieldbuffers[end].offset), len = $(fieldbuffers[end].length), padded = $(padding(fieldbuffers[end].length, alignment))" - return bufferoffset + blen -end - -function writebuffer(io, col::BoolVector, alignment) - @debug "writebuffer: col = $(typeof(col))" - @debug col - writebitmap(io, col, alignment) - n = Base.write(io, view(col.arrow, (col.pos):(col.pos + cld(col.ℓ, 8) - 1))) - return n + writezeros(io, paddinglength(n, alignment)) -end diff --git a/src/arraytypes/compressed.jl b/src/arraytypes/compressed.jl deleted file mode 100644 index 070ca33b..00000000 --- a/src/arraytypes/compressed.jl +++ /dev/null @@ -1,98 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -struct CompressedBuffer - data::Vector{UInt8} - uncompressedlength::Int64 -end - -""" - Arrow.Compressed - -Represents the compressed version of an [`ArrowVector`](@ref). -Holds a reference to the original column. May have `Compressed` -children for nested array types. -""" -struct Compressed{Z,A} - data::A - buffers::Vector{CompressedBuffer} - len::Int64 - nullcount::Int64 - children::Vector{Compressed} -end - -Base.length(c::Compressed) = c.len -Base.eltype(::Type{C}) where {Z,A,C<:Compressed{Z,A}} = eltype(A) -getmetadata(x::Compressed) = getmetadata(x.data) -compressiontype(c::Compressed{Z}) where {Z} = Z - -function compress(Z::Meta.CompressionType.T, comp, x::Array) - GC.@preserve x begin - y = unsafe_wrap(Array, convert(Ptr{UInt8}, pointer(x)), sizeof(x)) - return CompressedBuffer(transcode(comp, y), length(y)) - end -end - -compress(Z::Meta.CompressionType.T, comp, x) = compress(Z, comp, convert(Array, x)) - -compress(Z::Meta.CompressionType.T, comp, v::ValidityBitmap) = - v.nc == 0 ? CompressedBuffer(UInt8[], 0) : - compress(Z, comp, view(v.bytes, (v.pos):(v.pos + cld(v.ℓ, 8) - 1))) - -function makenodesbuffers!( - col::Compressed, - fieldnodes, - fieldbuffers, - bufferoffset, - alignment, -) - push!(fieldnodes, FieldNode(col.len, col.nullcount)) - @debug "made field node: nodeidx = $(length(fieldnodes)), col = $(typeof(col)), len = $(fieldnodes[end].length), nc = $(fieldnodes[end].null_count)" - for buffer in col.buffers - blen = length(buffer.data) == 0 ? 0 : 8 + length(buffer.data) - push!(fieldbuffers, Buffer(bufferoffset, blen)) - @debug "made field buffer: bufferidx = $(length(fieldbuffers)), offset = $(fieldbuffers[end].offset), len = $(fieldbuffers[end].length), padded = $(padding(fieldbuffers[end].length, alignment))" - bufferoffset += padding(blen, alignment) - end - for child in col.children - bufferoffset = - makenodesbuffers!(child, fieldnodes, fieldbuffers, bufferoffset, alignment) - end - return bufferoffset -end - -function writearray(io, b::CompressedBuffer) - if length(b.data) > 0 - n = Base.write(io, b.uncompressedlength) - @debug "writing compressed buffer: uncompressedlength = $(b.uncompressedlength), n = $(length(b.data))" - @debug b.data - return n + Base.write(io, b.data) - end - return 0 -end - -function writebuffer(io, col::Compressed, alignment) - @debug "writebuffer: col = $(typeof(col))" - @debug col - for buffer in col.buffers - n = writearray(io, buffer) - writezeros(io, paddinglength(n, alignment)) - end - for child in col.children - writebuffer(io, child, alignment) - end - return -end diff --git a/src/arraytypes/dictencoding.jl b/src/arraytypes/dictencoding.jl deleted file mode 100644 index 3e3576c5..00000000 --- a/src/arraytypes/dictencoding.jl +++ /dev/null @@ -1,418 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT 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.DictEncoding - -Represents the "pool" of possible values for a [`DictEncoded`](@ref) -array type. Whether the order of values is significant can be checked -by looking at the `isOrdered` boolean field. - -The `S` type parameter, while not tied directly to any field, is the -signed integer "index type" of the parent DictEncoded. We keep track -of this in the DictEncoding in order to validate the length of the pool -doesn't exceed the index type limit. The general workflow of writing arrow -data means the initial schema will typically be based off the data in the -first record batch, and subsequent record batches need to match the same -schema exactly. For example, if a non-first record batch dict encoded column -were to cause a DictEncoding pool to overflow on unique values, a fatal error -should be thrown. -""" -mutable struct DictEncoding{T,S,A} <: ArrowVector{T} - id::Int64 - data::A - isOrdered::Bool - metadata::Union{Nothing,Base.ImmutableDict{String,String}} -end - -indextype(::Type{DictEncoding{T,S,A}}) where {T,S,A} = S -indextype(::T) where {T<:DictEncoding} = indextype(T) - -Base.size(d::DictEncoding) = size(d.data) - -@propagate_inbounds function Base.getindex(d::DictEncoding{T}, i::Integer) where {T} - @boundscheck checkbounds(d, i) - return @inbounds ArrowTypes.fromarrow(T, d.data[i]) -end - -# convenience wrapper to signal that an input column should be -# dict encoded when written to the arrow format -struct DictEncodeType{T} end -getT(::Type{DictEncodeType{T}}) where {T} = T - -""" - Arrow.DictEncode(::AbstractVector, id::Integer=nothing) - -Signals that a column/array should be dictionary encoded when serialized -to the arrow streaming/file format. An optional `id` number may be provided -to signal that multiple columns should use the same pool when being -dictionary encoded. -""" -struct DictEncode{T,A} <: AbstractVector{DictEncodeType{T}} - id::Int64 - data::A -end - -DictEncode(x::A, id=-1) where {A} = DictEncode{eltype(A),A}(id, x) -Base.IndexStyle(::Type{<:DictEncode}) = Base.IndexLinear() -Base.size(x::DictEncode) = (length(x.data),) -Base.iterate(x::DictEncode, st...) = iterate(x.data, st...) -Base.getindex(x::DictEncode, i::Int) = getindex(x.data, i) -ArrowTypes.ArrowKind(::Type{<:DictEncodeType}) = DictEncodedKind() -Base.copy(x::DictEncode) = DictEncode(x.data, x.id) - -""" - Arrow.DictEncoded - -A dictionary encoded array type (similar to a `PooledArray`). Behaves just -like a normal array in most respects; internally, possible values are stored -in the `encoding::DictEncoding` field, while the `indices::Vector{<:Integer}` -field holds the "codes" of each element for indexing into the encoding pool. -Any column/array can be dict encoding when serializing to the arrow format -either by passing the `dictencode=true` keyword argument to [`Arrow.write`](@ref) -(which causes _all_ columns to be dict encoded), or wrapping individual columns/ -arrays in [`Arrow.DictEncode(x)`](@ref). -""" -struct DictEncoded{T,S,A} <: ArrowVector{T} - arrow::Vector{UInt8} # need to hold a reference to arrow memory blob - validity::ValidityBitmap - indices::Vector{S} - encoding::DictEncoding{T,S,A} - metadata::Union{Nothing,Base.ImmutableDict{String,String}} -end - -DictEncoded( - b::Vector{UInt8}, - v::ValidityBitmap, - inds::Vector{S}, - encoding::DictEncoding{T,S,A}, - meta, -) where {S,T,A} = DictEncoded{T,S,A}(b, v, inds, encoding, meta) - -Base.size(d::DictEncoded) = size(d.indices) - -isdictencoded(d::DictEncoded) = true -isdictencoded(x) = false -isdictencoded(c::Compressed{Z,A}) where {Z,A<:DictEncoded} = true - -function signedtype(n::Integer) - typs = (Int8, Int16, Int32, Int64) - typs[something(findfirst(n .≤ typemax.(typs)), 4)] -end - -signedtype(::Type{UInt8}) = Int8 -signedtype(::Type{UInt16}) = Int16 -signedtype(::Type{UInt32}) = Int32 -signedtype(::Type{UInt64}) = Int64 -signedtype(::Type{T}) where {T<:Signed} = T - -indtype(d::DictEncoded{T,S,A}) where {T,S,A} = S -indtype(c::Compressed{Z,A}) where {Z,A<:DictEncoded} = indtype(c.data) - -dictencodeid(colidx, nestedlevel, fieldid) = - (Int64(nestedlevel) << 48) | (Int64(fieldid) << 32) | Int64(colidx) - -getid(d::DictEncoded) = d.encoding.id -getid(c::Compressed{Z,A}) where {Z,A<:DictEncoded} = c.data.encoding.id - -function arrowvector( - ::DictEncodedKind, - x::DictEncoded, - i, - nl, - fi, - de, - ded, - meta; - dictencode::Bool=false, - dictencodenested::Bool=false, - kw..., -) - id = x.encoding.id - # XXX This is a race condition if two workers hit this block at the same time, then they'll create - # distinct locks - if !haskey(de, id) - de[id] = Lockable(x.encoding) - else - encodinglockable = de[id] - Base.@lock encodinglockable begin - encoding = encodinglockable.value - # in this case, we just need to check if any values in our local pool need to be delta dicationary serialized - deltas = setdiff(x.encoding, encoding) - if !isempty(deltas) - ET = indextype(encoding) - if length(deltas) + length(encoding) > typemax(ET) - error( - "fatal error serializing dict encoded column with ref index type of $ET; subsequent record batch unique values resulted in $(length(deltas) + length(encoding)) unique values, which exceeds possible index values in $ET", - ) - end - data = arrowvector( - deltas, - i, - nl, - fi, - de, - ded, - nothing; - dictencode=dictencodenested, - dictencodenested=dictencodenested, - dictencoding=true, - kw..., - ) - push!( - ded, - DictEncoding{eltype(data),ET,typeof(data)}( - id, - data, - false, - getmetadata(data), - ), - ) - if typeof(encoding.data) <: ChainedVector - append!(encoding.data, data) - else - data2 = ChainedVector([encoding.data, data]) - encoding = DictEncoding{eltype(data2),ET,typeof(data2)}( - id, - data2, - false, - getmetadata(encoding), - ) - de[id] = Lockable(encoding) - end - end - end - end - return x -end - -function arrowvector( - ::DictEncodedKind, - x, - i, - nl, - fi, - de, - ded, - meta; - dictencode::Bool=false, - dictencodenested::Bool=false, - kw..., -) - @assert x isa DictEncode - id = x.id == -1 ? dictencodeid(i, nl, fi) : x.id - x = x.data - len = length(x) - validity = ValidityBitmap(x) - # XXX This is a race condition if two workers hit this block at the same time, then they'll create - # distinct locks - if !haskey(de, id) - # dict encoding doesn't exist yet, so create for 1st time - if DataAPI.refarray(x) === x || DataAPI.refpool(x) === nothing - # need to encode ourselves - x = PooledArray(x; signed=true, compress=true) - inds = refa = DataAPI.refarray(x) - pool = DataAPI.refpool(x) - else - pool = DataAPI.refpool(x) - refa = DataAPI.refarray(x) - inds = copyto!(similar(Vector{signedtype(length(pool))}, length(refa)), refa) - end - # adjust to "offset" instead of index - inds .-= firstindex(refa) - data = arrowvector( - pool, - i, - nl, - fi, - de, - ded, - nothing; - dictencode=dictencodenested, - dictencodenested=dictencodenested, - dictencoding=true, - kw..., - ) - encoding = DictEncoding{eltype(data),eltype(inds),typeof(data)}( - id, - data, - false, - getmetadata(data), - ) - de[id] = Lockable(encoding) - else - # encoding already exists - # compute inds based on it - # if value doesn't exist in encoding, push! it - # also add to deltas updates - encodinglockable = de[id] - Base.@lock encodinglockable begin - encoding = encodinglockable.value - len = length(x) - ET = indextype(encoding) - pool = Dict{Union{eltype(encoding),eltype(x)},ET}( - a => (b - 1) for (b, a) in enumerate(encoding) - ) - deltas = eltype(x)[] - inds = Vector{ET}(undef, len) - for (j, val) in enumerate(x) - @inbounds inds[j] = get!(pool, val) do - push!(deltas, val) - return length(pool) - end - end - if !isempty(deltas) - if length(deltas) + length(encoding) > typemax(ET) - error( - "fatal error serializing dict encoded column with ref index type of $ET; subsequent record batch unique values resulted in $(length(deltas) + length(encoding)) unique values, which exceeds possible index values in $ET", - ) - end - data = arrowvector( - deltas, - i, - nl, - fi, - de, - ded, - nothing; - dictencode=dictencodenested, - dictencodenested=dictencodenested, - dictencoding=true, - kw..., - ) - push!( - ded, - DictEncoding{eltype(data),ET,typeof(data)}( - id, - data, - false, - getmetadata(data), - ), - ) - if typeof(encoding.data) <: ChainedVector - append!(encoding.data, data) - else - data2 = ChainedVector([encoding.data, data]) - encoding = DictEncoding{eltype(data2),ET,typeof(data2)}( - id, - data2, - false, - getmetadata(encoding), - ) - de[id] = Lockable(encoding) - end - end - end - end - if meta !== nothing && getmetadata(encoding) !== nothing - meta = toidict(merge!(Dict(meta), Dict(getmetadata(encoding)))) - elseif getmetadata(encoding) !== nothing - meta = getmetadata(encoding) - end - return DictEncoded(UInt8[], validity, inds, encoding, meta) -end - -@propagate_inbounds function Base.getindex(d::DictEncoded, i::Integer) - @boundscheck checkbounds(d, i) - @inbounds valid = d.validity[i] - !valid && return missing - @inbounds idx = d.indices[i] - return @inbounds d.encoding[idx + 1] -end - -@propagate_inbounds function Base.setindex!(d::DictEncoded{T}, v, i::Integer) where {T} - @boundscheck checkbounds(d, i) - if v === missing - @inbounds d.validity[i] = false - else - ix = findfirst(d.encoding.data, v) - if ix === nothing - push!(d.encoding.data, v) - @inbounds d.indices[i] = length(d.encoding.data) - 1 - else - @inbounds d.indices[i] = ix - 1 - end - end - return v -end - -function Base.copy(x::DictEncoded{T,S}) where {T,S} - pool = copy(x.encoding.data) - valid = x.validity - inds = x.indices - refs = copy(inds) - @inbounds for i = 1:length(inds) - refs[i] = refs[i] + one(S) - end - return PooledArray( - PooledArrays.RefArray(refs), - Dict{T,S}(val => i for (i, val) in enumerate(pool)), - pool, - ) -end - -function compress(Z::Meta.CompressionType.T, comp, x::A) where {A<:DictEncoded} - len = length(x) - nc = nullcount(x) - validity = compress(Z, comp, x.validity) - inds = compress(Z, comp, x.indices) - return Compressed{Z,A}(x, [validity, inds], len, nc, Compressed[]) -end - -function DataAPI.levels(x::DictEncoded) - rp = DataAPI.refpool(x) # may contain missing values - Missing <: eltype(rp) || return rp - convert(AbstractArray{nonmissingtype(eltype(rp))}, deleteat!(rp, ismissing.(rp))) -end - -function makenodesbuffers!( - col::DictEncoded{T,S}, - fieldnodes, - fieldbuffers, - bufferoffset, - alignment, -) where {T,S} - len = length(col) - nc = nullcount(col) - push!(fieldnodes, FieldNode(len, nc)) - @debug "made field node: nodeidx = $(length(fieldnodes)), col = $(typeof(col)), len = $(fieldnodes[end].length), nc = $(fieldnodes[end].null_count)" - # validity bitmap - blen = nc == 0 ? 0 : bitpackedbytes(len, alignment) - push!(fieldbuffers, Buffer(bufferoffset, blen)) - @debug "made field buffer: bufferidx = $(length(fieldbuffers)), offset = $(fieldbuffers[end].offset), len = $(fieldbuffers[end].length), padded = $(padding(fieldbuffers[end].length, alignment))" - bufferoffset += blen - # indices - blen = sizeof(S) * len - push!(fieldbuffers, Buffer(bufferoffset, blen)) - @debug "made field buffer: bufferidx = $(length(fieldbuffers)), offset = $(fieldbuffers[end].offset), len = $(fieldbuffers[end].length), padded = $(padding(fieldbuffers[end].length, alignment))" - bufferoffset += padding(blen, alignment) - return bufferoffset -end - -DataAPI.refarray(x::DictEncoded{T,S}) where {T,S} = x.indices .+ one(S) - -DataAPI.refpool(x::DictEncoded) = copy(x.encoding.data) - -function writebuffer(io, col::DictEncoded, alignment) - @debug "writebuffer: col = $(typeof(col))" - @debug col - writebitmap(io, col, alignment) - # write indices - n = writearray(io, col.indices) - @debug "writing array: col = $(typeof(col.indices)), n = $n, padded = $(padding(n, alignment))" - writezeros(io, paddinglength(n, alignment)) - return -end diff --git a/src/arraytypes/fixedsizelist.jl b/src/arraytypes/fixedsizelist.jl deleted file mode 100644 index 2558dd54..00000000 --- a/src/arraytypes/fixedsizelist.jl +++ /dev/null @@ -1,203 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT 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.FixedSizeList - -An `ArrowVector` where each element is a "fixed size" list of some kind, like a `NTuple{N, T}`. -""" -struct FixedSizeList{T,A<:AbstractVector} <: ArrowVector{T} - arrow::Vector{UInt8} # need to hold a reference to arrow memory blob - validity::ValidityBitmap - data::A - ℓ::Int - metadata::Union{Nothing,Base.ImmutableDict{String,String}} -end - -Base.size(l::FixedSizeList) = (l.ℓ,) - -@propagate_inbounds function Base.getindex(l::FixedSizeList{T}, i::Integer) where {T} - @boundscheck checkbounds(l, i) - S = Base.nonmissingtype(T) - X = ArrowTypes.ArrowKind(ArrowTypes.ArrowType(S)) - N = ArrowTypes.getsize(X) - Y = ArrowTypes.gettype(X) - if X !== T && !(l.validity[i]) - return missing - else - off = (i - 1) * N - if X === T && isbitstype(Y) - tup = _unsafe_load_tuple(NTuple{N,Y}, l.data, off + 1) - else - tup = ntuple(j -> l.data[off + j], N) - end - return ArrowTypes.fromarrow(T, tup) - end -end - -function _unsafe_load_tuple( - ::Type{NTuple{N,T}}, - bytes::Vector{UInt8}, - i::Integer, -) where {N,T} - x = Ref(bytes, i) - y = Ref{NTuple{N,T}}() - ArrowTypes._unsafe_cast!(y, x, N) - return y[] -end - -@propagate_inbounds function Base.setindex!(l::FixedSizeList{T}, v::T, i::Integer) where {T} - @boundscheck checkbounds(l, i) - if v === missing - @inbounds l.validity[i] = false - else - N = ArrowTypes.getsize( - ArrowTypes.ArrowKind(ArrowTypes.ArrowType(Base.nonmissingtype(T))), - ) - off = (i - 1) * N - foreach(1:N) do j - @inbounds l.data[off + j] = v[j] - end - end - return v -end - -# lazy equal-spaced flattener -struct ToFixedSizeList{T,N,A} <: AbstractVector{T} - data::A # A is AbstractVector of (AbstractVector or AbstractString) -end - -origtype(::ToFixedSizeList{T,N,A}) where {T,N,A} = eltype(A) - -function ToFixedSizeList(input) - NT = ArrowTypes.ArrowKind(Base.nonmissingtype(eltype(input))) # typically NTuple{N, T} - return ToFixedSizeList{ArrowTypes.gettype(NT),ArrowTypes.getsize(NT),typeof(input)}( - input, - ) -end - -Base.IndexStyle(::Type{<:ToFixedSizeList}) = Base.IndexLinear() -Base.size(x::ToFixedSizeList{T,N}) where {T,N} = (N * length(x.data),) - -Base.@propagate_inbounds function Base.getindex( - A::ToFixedSizeList{T,N}, - i::Integer, -) where {T,N} - @boundscheck checkbounds(A, i) - a, b = fldmod1(i, N) - @inbounds x = A.data[a] - return @inbounds x === missing ? ArrowTypes.default(T) : x[b] -end - -# efficient iteration -@inline function Base.iterate( - A::ToFixedSizeList{T,N}, - (i, chunk, chunk_i, len)=(1, 1, 1, length(A)), -) where {T,N} - i > len && return nothing - @inbounds y = A.data[chunk] - @inbounds x = y === missing ? ArrowTypes.default(T) : y[chunk_i] - if chunk_i == N - chunk += 1 - chunk_i = 1 - else - chunk_i += 1 - end - return x, (i + 1, chunk, chunk_i, len) -end - -arrowvector(::FixedSizeListKind, x::FixedSizeList, i, nl, fi, de, ded, meta; kw...) = x - -function arrowvector( - ::FixedSizeListKind{N,T}, - x, - i, - nl, - fi, - de, - ded, - meta; - kw..., -) where {N,T} - len = length(x) - validity = ValidityBitmap(x) - flat = ToFixedSizeList(x) - if eltype(flat) == UInt8 - data = flat - S = origtype(flat) - else - data = arrowvector(flat, i, nl + 1, fi, de, ded, nothing; kw...) - S = withmissing(eltype(x), NTuple{N,eltype(data)}) - end - return FixedSizeList{S,typeof(data)}(UInt8[], validity, data, len, meta) -end - -function compress(Z::Meta.CompressionType.T, comp, x::FixedSizeList{T,A}) where {T,A} - len = length(x) - nc = nullcount(x) - validity = compress(Z, comp, x.validity) - buffers = [validity] - children = Compressed[] - if eltype(A) == UInt8 - push!(buffers, compress(Z, comp, x.data)) - else - push!(children, compress(Z, comp, x.data)) - end - return Compressed{Z,typeof(x)}(x, buffers, len, nc, children) -end - -function makenodesbuffers!( - col::FixedSizeList{T,A}, - fieldnodes, - fieldbuffers, - bufferoffset, - alignment, -) where {T,A} - len = length(col) - nc = nullcount(col) - push!(fieldnodes, FieldNode(len, nc)) - @debug "made field node: nodeidx = $(length(fieldnodes)), col = $(typeof(col)), len = $(fieldnodes[end].length), nc = $(fieldnodes[end].null_count)" - # validity bitmap - blen = nc == 0 ? 0 : bitpackedbytes(len, alignment) - push!(fieldbuffers, Buffer(bufferoffset, blen)) - @debug "made field buffer: bufferidx = $(length(fieldbuffers)), offset = $(fieldbuffers[end].offset), len = $(fieldbuffers[end].length), padded = $(padding(fieldbuffers[end].length, alignment))" - bufferoffset += blen - if eltype(A) === UInt8 - blen = ArrowTypes.getsize(ArrowTypes.ArrowKind(Base.nonmissingtype(T))) * len - push!(fieldbuffers, Buffer(bufferoffset, blen)) - @debug "made field buffer: bufferidx = $(length(fieldbuffers)), offset = $(fieldbuffers[end].offset), len = $(fieldbuffers[end].length), padded = $(padding(fieldbuffers[end].length, alignment))" - bufferoffset += padding(blen, alignment) - else - bufferoffset = - makenodesbuffers!(col.data, fieldnodes, fieldbuffers, bufferoffset, alignment) - end - return bufferoffset -end - -function writebuffer(io, col::FixedSizeList{T,A}, alignment) where {T,A} - @debug "writebuffer: col = $(typeof(col))" - @debug col - writebitmap(io, col, alignment) - # write values array - if eltype(A) === UInt8 - n = writearray(io, UInt8, col.data) - @debug "writing array: col = $(typeof(col.data)), n = $n, padded = $(padding(n, alignment))" - writezeros(io, paddinglength(n, alignment)) - else - writebuffer(io, col.data, alignment) - end - return -end diff --git a/src/arraytypes/list.jl b/src/arraytypes/list.jl deleted file mode 100644 index 41ac66f9..00000000 --- a/src/arraytypes/list.jl +++ /dev/null @@ -1,260 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -struct Offsets{T<:Union{Int32,Int64}} <: ArrowVector{Tuple{T,T}} - arrow::Vector{UInt8} # need to hold a reference to arrow memory blob - offsets::Vector{T} -end - -Base.size(o::Offsets) = (length(o.offsets) - 1,) - -@propagate_inbounds function Base.getindex(o::Offsets, i::Integer) - @boundscheck checkbounds(o, i) - @inbounds lo = o.offsets[i] + 1 - @inbounds hi = o.offsets[i + 1] - return lo, hi -end - -""" - Arrow.List - -An `ArrowVector` where each element is a variable sized list of some kind, like an `AbstractVector` or `AbstractString`. -""" -struct List{T,O,A} <: ArrowVector{T} - arrow::Vector{UInt8} # need to hold a reference to arrow memory blob - validity::ValidityBitmap - offsets::Offsets{O} - data::A - ℓ::Int - metadata::Union{Nothing,Base.ImmutableDict{String,String}} -end - -Base.size(l::List) = (l.ℓ,) - -@propagate_inbounds function Base.getindex(l::List{T}, i::Integer) where {T} - @boundscheck checkbounds(l, i) - @inbounds lo, hi = l.offsets[i] - S = Base.nonmissingtype(T) - K = ArrowTypes.ArrowKind(ArrowTypes.ArrowType(S)) - # special-case Base.CodeUnits for ArrowTypes compat - if ArrowTypes.isstringtype(K) || S <: Base.CodeUnits - if S !== T - if S <: Base.CodeUnits - return l.validity[i] ? - Base.CodeUnits(unsafe_string(pointer(l.data, lo), hi - lo + 1)) : - missing - else - return l.validity[i] ? - ArrowTypes.fromarrow(T, pointer(l.data, lo), hi - lo + 1) : missing - end - else - if S <: Base.CodeUnits - return Base.CodeUnits(unsafe_string(pointer(l.data, lo), hi - lo + 1)) - else - return ArrowTypes.fromarrow(T, pointer(l.data, lo), hi - lo + 1) - end - end - elseif S !== T - return l.validity[i] ? ArrowTypes.fromarrow(T, view(l.data, lo:hi)) : missing - else - return ArrowTypes.fromarrow(T, view(l.data, lo:hi)) - end -end - -# @propagate_inbounds function Base.setindex!(l::List{T}, v, i::Integer) where {T} - -# end - -# internal interface definitions to be able to treat AbstractString/CodeUnits similarly -_ncodeunits(x::AbstractString) = ncodeunits(x) -_codeunits(x::AbstractString) = codeunits(x) -_ncodeunits(x::Base.CodeUnits) = length(x) -_codeunits(x::Base.CodeUnits) = x - -# an AbstractVector version of Iterators.flatten -# code based on SentinelArrays.ChainedVector -struct ToList{T,stringtype,A,I} <: AbstractVector{T} - data::Vector{A} # A is AbstractVector or AbstractString - inds::Vector{I} -end - -origtype(::ToList{T,S,A,I}) where {T,S,A,I} = A -liststringtype(::Type{ToList{T,S,A,I}}) where {T,S,A,I} = S -function liststringtype(::List{T,O,A}) where {T,O,A} - ST = Base.nonmissingtype(T) - K = ArrowTypes.ArrowKind(ST) - return liststringtype(A) || ArrowTypes.isstringtype(K) || ST <: Base.CodeUnits # add the CodeUnits check for ArrowTypes compat for now -end -liststringtype(T) = false - -function ToList(input; largelists::Bool=false) - AT = eltype(input) - ST = Base.nonmissingtype(AT) - K = ArrowTypes.ArrowKind(ST) - stringtype = ArrowTypes.isstringtype(K) || ST <: Base.CodeUnits # add the CodeUnits check for ArrowTypes compat for now - T = stringtype ? UInt8 : eltype(ST) - len = stringtype ? _ncodeunits : length - data = AT[] - I = largelists ? Int64 : Int32 - inds = I[0] - sizehint!(data, length(input)) - sizehint!(inds, length(input)) - totalsize = I(0) - for x in input - if x === missing - push!(data, missing) - else - push!(data, x) - totalsize += len(x) - if I === Int32 && totalsize > 2147483647 - I = Int64 - inds = convert(Vector{Int64}, inds) - end - end - push!(inds, totalsize) - end - return ToList{T,stringtype,AT,I}(data, inds) -end - -Base.IndexStyle(::Type{<:ToList}) = Base.IndexLinear() -Base.size(x::ToList{T,S,A,I}) where {T,S,A,I} = (isempty(x.inds) ? zero(I) : x.inds[end],) - -function Base.pointer(A::ToList{UInt8}, i::Integer) - chunk = searchsortedfirst(A.inds, i) - chunk = chunk > length(A.inds) ? 1 : (chunk - 1) - return pointer(A.data[chunk]) -end - -@inline function index(A::ToList, i::Integer) - chunk = searchsortedfirst(A.inds, i) - return chunk - 1, i - (@inbounds A.inds[chunk - 1]) -end - -Base.@propagate_inbounds function Base.getindex( - A::ToList{T,stringtype}, - i::Integer, -) where {T,stringtype} - @boundscheck checkbounds(A, i) - chunk, ix = index(A, i) - @inbounds x = A.data[chunk] - return @inbounds stringtype ? _codeunits(x)[ix] : x[ix] -end - -Base.@propagate_inbounds function Base.setindex!( - A::ToList{T,stringtype}, - v, - i::Integer, -) where {T,stringtype} - @boundscheck checkbounds(A, i) - chunk, ix = index(A, i) - @inbounds x = A.data[chunk] - if stringtype - _codeunits(x)[ix] = v - else - x[ix] = v - end - return v -end - -# efficient iteration -@inline function Base.iterate(A::ToList{T,stringtype}) where {T,stringtype} - length(A) == 0 && return nothing - i = 1 - chunk = 2 - chunk_i = 1 - chunk_len = A.inds[chunk] - while i > chunk_len - chunk += 1 - chunk_len = A.inds[chunk] - end - val = A.data[chunk - 1] - x = stringtype ? _codeunits(val)[1] : val[1] - # find next valid index - i += 1 - if i > chunk_len - while true - chunk += 1 - chunk > length(A.inds) && break - chunk_len = A.inds[chunk] - i <= chunk_len && break - end - else - chunk_i += 1 - end - return x, (i, chunk, chunk_i, chunk_len, length(A)) -end - -@inline function Base.iterate( - A::ToList{T,stringtype}, - (i, chunk, chunk_i, chunk_len, len), -) where {T,stringtype} - i > len && return nothing - @inbounds val = A.data[chunk - 1] - @inbounds x = stringtype ? _codeunits(val)[chunk_i] : val[chunk_i] - i += 1 - if i > chunk_len - chunk_i = 1 - while true - chunk += 1 - chunk > length(A.inds) && break - @inbounds chunk_len = A.inds[chunk] - i <= chunk_len && break - end - else - chunk_i += 1 - end - return x, (i, chunk, chunk_i, chunk_len, len) -end - -arrowvector(::ListKind, x::List, i, nl, fi, de, ded, meta; kw...) = x - -function arrowvector(::ListKind, x, i, nl, fi, de, ded, meta; largelists::Bool=false, kw...) - len = length(x) - validity = ValidityBitmap(x) - flat = ToList(x; largelists=largelists) - offsets = Offsets(UInt8[], flat.inds) - if liststringtype(typeof(flat)) && eltype(flat) == UInt8 # binary or utf8string - data = flat - T = origtype(flat) - else - data = - arrowvector(flat, i, nl + 1, fi, de, ded, nothing; largelists=largelists, kw...) - T = withmissing(eltype(x), Vector{eltype(data)}) - end - return List{T,eltype(flat.inds),typeof(data)}( - UInt8[], - validity, - offsets, - data, - len, - meta, - ) -end - -function compress(Z::Meta.CompressionType.T, comp, x::List{T,O,A}) where {T,O,A} - len = length(x) - nc = nullcount(x) - validity = compress(Z, comp, x.validity) - offsets = compress(Z, comp, x.offsets.offsets) - buffers = [validity, offsets] - children = Compressed[] - if liststringtype(x) - push!(buffers, compress(Z, comp, x.data)) - else - push!(children, compress(Z, comp, x.data)) - end - return Compressed{Z,typeof(x)}(x, buffers, len, nc, children) -end diff --git a/src/arraytypes/map.jl b/src/arraytypes/map.jl deleted file mode 100644 index 42160732..00000000 --- a/src/arraytypes/map.jl +++ /dev/null @@ -1,146 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" - Arrow.Map - -An `ArrowVector` where each element is a "map" of some kind, like a `Dict`. -""" -struct Map{T,O,A} <: ArrowVector{T} - validity::ValidityBitmap - offsets::Offsets{O} - data::A - ℓ::Int - metadata::Union{Nothing,Base.ImmutableDict{String,String}} -end - -Base.size(l::Map) = (l.ℓ,) - -@propagate_inbounds function Base.getindex(l::Map{T}, i::Integer) where {T} - @boundscheck checkbounds(l, i) - @inbounds lo, hi = l.offsets[i] - if Base.nonmissingtype(T) !== T - return l.validity[i] ? - ArrowTypes.fromarrow( - T, - Dict(x.key => x.value for x in view(l.data, lo:hi)), - ) : missing - else - return ArrowTypes.fromarrow(T, Dict(x.key => x.value for x in view(l.data, lo:hi))) - end -end - -keyvalues(KT, ::Missing) = missing -keyvalues(KT, x::AbstractDict) = [KT(k, v) for (k, v) in pairs(x)] - -keyvaluetypes(::Type{NamedTuple{(:key, :value),Tuple{K,V}}}) where {K,V} = (K, V) - -arrowvector(::MapKind, x::Map, i, nl, fi, de, ded, meta; kw...) = x - -function arrowvector(::MapKind, x, i, nl, fi, de, ded, meta; largelists::Bool=false, kw...) - len = length(x) - validity = ValidityBitmap(x) - ET = eltype(x) - DT = Base.nonmissingtype(ET) - KDT, VDT = keytype(DT), valtype(DT) - ArrowTypes.concrete_or_concreteunion(KDT) || throw( - ArgumentError( - "`keytype(d)` must be concrete to serialize map-like `d`, but `keytype(d) == $KDT`", - ), - ) - ArrowTypes.concrete_or_concreteunion(VDT) || throw( - ArgumentError( - "`valtype(d)` must be concrete to serialize map-like `d`, but `valtype(d) == $VDT`", - ), - ) - KT = KeyValue{KDT,VDT} - VT = Vector{KT} - T = DT !== ET ? Union{Missing,VT} : VT - flat = ToList(T[keyvalues(KT, y) for y in x]; largelists=largelists) - offsets = Offsets(UInt8[], flat.inds) - data = arrowvector(flat, i, nl + 1, fi, de, ded, nothing; largelists=largelists, kw...) - K, V = keyvaluetypes(eltype(data)) - return Map{withmissing(ET, Dict{K,V}),eltype(flat.inds),typeof(data)}( - validity, - offsets, - data, - len, - meta, - ) -end - -function compress(Z::Meta.CompressionType.T, comp, x::A) where {A<:Map} - len = length(x) - nc = nullcount(x) - validity = compress(Z, comp, x.validity) - offsets = compress(Z, comp, x.offsets.offsets) - buffers = [validity, offsets] - children = Compressed[] - push!(children, compress(Z, comp, x.data)) - return Compressed{Z,A}(x, buffers, len, nc, children) -end - -function makenodesbuffers!( - col::Union{Map{T,O,A},List{T,O,A}}, - fieldnodes, - fieldbuffers, - bufferoffset, - alignment, -) where {T,O,A} - len = length(col) - nc = nullcount(col) - push!(fieldnodes, FieldNode(len, nc)) - @debug "made field node: nodeidx = $(length(fieldnodes)), col = $(typeof(col)), len = $(fieldnodes[end].length), nc = $(fieldnodes[end].null_count)" - # validity bitmap - blen = nc == 0 ? 0 : bitpackedbytes(len, alignment) - push!(fieldbuffers, Buffer(bufferoffset, blen)) - @debug "made field buffer: bufferidx = $(length(fieldbuffers)), offset = $(fieldbuffers[end].offset), len = $(fieldbuffers[end].length), padded = $(padding(fieldbuffers[end].length, alignment))" - # adjust buffer offset, make array buffer - bufferoffset += blen - blen = sizeof(O) * (len + 1) - push!(fieldbuffers, Buffer(bufferoffset, blen)) - @debug "made field buffer: bufferidx = $(length(fieldbuffers)), offset = $(fieldbuffers[end].offset), len = $(fieldbuffers[end].length), padded = $(padding(fieldbuffers[end].length, alignment))" - bufferoffset += padding(blen, alignment) - if liststringtype(col) - blen = length(col.data) - push!(fieldbuffers, Buffer(bufferoffset, blen)) - @debug "made field buffer: bufferidx = $(length(fieldbuffers)), offset = $(fieldbuffers[end].offset), len = $(fieldbuffers[end].length), padded = $(padding(fieldbuffers[end].length, alignment))" - bufferoffset += padding(blen, alignment) - else - bufferoffset = - makenodesbuffers!(col.data, fieldnodes, fieldbuffers, bufferoffset, alignment) - end - return bufferoffset -end - -function writebuffer(io, col::Union{Map{T,O,A},List{T,O,A}}, alignment) where {T,O,A} - @debug "writebuffer: col = $(typeof(col))" - @debug col - writebitmap(io, col, alignment) - # write offsets - n = writearray(io, O, col.offsets.offsets) - @debug "writing array: col = $(typeof(col.offsets.offsets)), n = $n, padded = $(padding(n, alignment))" - writezeros(io, paddinglength(n, alignment)) - # write values array - if liststringtype(col) - n = writearray(io, UInt8, col.data) - @debug "writing array: col = $(typeof(col.data)), n = $n, padded = $(padding(n, alignment))" - writezeros(io, paddinglength(n, alignment)) - else - writebuffer(io, col.data, alignment) - end - return -end diff --git a/src/arraytypes/primitive.jl b/src/arraytypes/primitive.jl deleted file mode 100644 index 7d86bfe0..00000000 --- a/src/arraytypes/primitive.jl +++ /dev/null @@ -1,112 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" - Arrow.Primitive - -An `ArrowVector` where each element is a "fixed size" scalar of some kind, like an integer, float, decimal, or time type. -""" -struct Primitive{T,A} <: ArrowVector{T} - arrow::Vector{UInt8} # need to hold a reference to arrow memory blob - validity::ValidityBitmap - data::A - ℓ::Int64 - metadata::Union{Nothing,Base.ImmutableDict{String,String}} -end - -Primitive(::Type{T}, b::Vector{UInt8}, v::ValidityBitmap, data::A, l, meta) where {T,A} = - Primitive{T,A}(b, v, data, l, meta) - -Base.size(p::Primitive) = (p.ℓ,) - -function Base.copy(p::Primitive{T,A}) where {T,A} - if nullcount(p) == 0 && T === eltype(A) - return copy(p.data) - else - return convert(Array, p) - end -end - -@propagate_inbounds function Base.getindex(p::Primitive{T}, i::Integer) where {T} - @boundscheck checkbounds(p, i) - if T >: Missing - return @inbounds (p.validity[i] ? ArrowTypes.fromarrow(T, p.data[i]) : missing) - else - return @inbounds ArrowTypes.fromarrow(T, p.data[i]) - end -end - -@propagate_inbounds function Base.setindex!(p::Primitive{T}, v, i::Integer) where {T} - @boundscheck checkbounds(p, i) - if T >: Missing - if v === missing - @inbounds p.validity[i] = false - else - @inbounds p.data[i] = convert(Base.nonmissingtype(T), v) - end - else - @inbounds p.data[i] = convert(Base.nonmissingtype(T), v) - end - return v -end - -arrowvector(::PrimitiveKind, x::Primitive, i, nl, fi, de, ded, meta; kw...) = x - -function arrowvector(::PrimitiveKind, x, i, nl, fi, de, ded, meta; kw...) - validity = ValidityBitmap(x) - return Primitive(eltype(x), UInt8[], validity, x, length(x), meta) -end - -function compress(Z::Meta.CompressionType.T, comp, p::P) where {P<:Primitive} - len = length(p) - nc = nullcount(p) - validity = compress(Z, comp, p.validity) - data = compress(Z, comp, p.data) - return Compressed{Z,P}(p, [validity, data], len, nc, Compressed[]) -end - -function makenodesbuffers!( - col::Primitive{T}, - fieldnodes, - fieldbuffers, - bufferoffset, - alignment, -) where {T} - len = length(col) - nc = nullcount(col) - push!(fieldnodes, FieldNode(len, nc)) - @debug "made field node: nodeidx = $(length(fieldnodes)), col = $(typeof(col)), len = $(fieldnodes[end].length), nc = $(fieldnodes[end].null_count)" - # validity bitmap - blen = nc == 0 ? 0 : bitpackedbytes(len, alignment) - push!(fieldbuffers, Buffer(bufferoffset, blen)) - @debug "made field buffer: bufferidx = $(length(fieldbuffers)), offset = $(fieldbuffers[end].offset), len = $(fieldbuffers[end].length), padded = $(padding(fieldbuffers[end].length, alignment))" - # adjust buffer offset, make primitive array buffer - bufferoffset += blen - blen = len * sizeof(Base.nonmissingtype(T)) - push!(fieldbuffers, Buffer(bufferoffset, blen)) - @debug "made field buffer: bufferidx = $(length(fieldbuffers)), offset = $(fieldbuffers[end].offset), len = $(fieldbuffers[end].length), padded = $(padding(fieldbuffers[end].length, alignment))" - return bufferoffset + padding(blen, alignment) -end - -function writebuffer(io, col::Primitive{T}, alignment) where {T} - @debug "writebuffer: col = $(typeof(col))" - @debug col - writebitmap(io, col, alignment) - n = writearray(io, Base.nonmissingtype(T), col.data) - @debug "writing array: col = $(typeof(col.data)), n = $n, padded = $(padding(n, alignment))" - writezeros(io, paddinglength(n, alignment)) - return -end diff --git a/src/arraytypes/struct.jl b/src/arraytypes/struct.jl deleted file mode 100644 index 23a8b641..00000000 --- a/src/arraytypes/struct.jl +++ /dev/null @@ -1,162 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT 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.Struct - -An `ArrowVector` where each element is a "struct" of some kind with ordered, named fields, like a `NamedTuple{names, types}` or regular julia `struct`. -""" -struct Struct{T,S,fnames} <: ArrowVector{T} - validity::ValidityBitmap - data::S # Tuple of ArrowVector - ℓ::Int - metadata::Union{Nothing,Base.ImmutableDict{String,String}} -end - -Base.size(s::Struct) = (s.ℓ,) - -isnamedtuple(::Type{<:NamedTuple}) = true -isnamedtuple(T) = false -istuple(::Type{<:Tuple}) = true -istuple(T) = false - -if isdefined(ArrowTypes, :fromarrowstruct) - # https://github.com/apache/arrow-julia/pull/493 - @inline function _fromarrowstruct(T::Type, v::Val, x...) - return ArrowTypes.fromarrowstruct(T, v, x...) - end -else - @inline function _fromarrowstruct(T::Type, ::Val, x...) - return ArrowTypes.fromarrow(T, x...) - end -end - -@propagate_inbounds function Base.getindex( - s::Struct{T,S,fnames}, - i::Integer, -) where {T,S,fnames} - @boundscheck checkbounds(s, i) - NT = Base.nonmissingtype(T) - NT !== T && (s.validity[i] || return missing) - vals = ntuple(j -> s.data[j][i], fieldcount(S)) - if isnamedtuple(NT) || istuple(NT) - return NT(vals) - else - return _fromarrowstruct(NT, Val{fnames}(), vals...) - end -end - -# @propagate_inbounds function Base.setindex!(s::Struct{T}, v::T, i::Integer) where {T} -# @boundscheck checkbounds(s, i) -# if v === missing -# @inbounds s.validity[i] = false -# else -# NT = Base.nonmissingtype(T) -# N = fieldcount(NT) -# foreach(1:N) do j -# @inbounds s.data[j][i] = getfield(v, j) -# end -# end -# return v -# end - -struct ToStruct{T,i,A} <: AbstractVector{T} - data::A # eltype is NamedTuple or some struct -end - -ToStruct(x::A, j::Integer) where {A} = - ToStruct{fieldtype(Base.nonmissingtype(eltype(A)), j),j,A}(x) - -Base.IndexStyle(::Type{<:ToStruct}) = Base.IndexLinear() -Base.size(x::ToStruct) = (length(x.data),) - -Base.@propagate_inbounds function Base.getindex(A::ToStruct{T,j}, i::Integer) where {T,j} - @boundscheck checkbounds(A, i) - @inbounds x = A.data[i] - return x === missing ? ArrowTypes.default(T) : getfield(x, j) -end - -arrowvector(::StructKind, x::Struct, i, nl, fi, de, ded, meta; kw...) = x - -namedtupletype(::Type{NamedTuple{names,types}}, data) where {names,types} = - NamedTuple{names,Tuple{(eltype(x) for x in data)...}} -namedtupletype(::Type{T}, data) where {T} = - NamedTuple{fieldnames(T),Tuple{(eltype(x) for x in data)...}} -namedtupletype(::Type{T}, data) where {T<:Tuple} = - NamedTuple{map(Symbol, fieldnames(T)),Tuple{(eltype(x) for x in data)...}} - -function arrowvector(::StructKind, x, i, nl, fi, de, ded, meta; kw...) - len = length(x) - validity = ValidityBitmap(x) - T = Base.nonmissingtype(eltype(x)) - data = Tuple( - arrowvector(ToStruct(x, j), i, nl + 1, j, de, ded, nothing; kw...) for - j = 1:fieldcount(T) - ) - NT = namedtupletype(T, data) - return Struct{withmissing(eltype(x), NT),typeof(data),fieldnames(NT)}( - validity, - data, - len, - meta, - ) -end - -function compress(Z::Meta.CompressionType.T, comp, x::A) where {A<:Struct} - len = length(x) - nc = nullcount(x) - validity = compress(Z, comp, x.validity) - buffers = [validity] - children = Compressed[] - for y in x.data - push!(children, compress(Z, comp, y)) - end - return Compressed{Z,A}(x, buffers, len, nc, children) -end - -function makenodesbuffers!( - col::Struct{T}, - fieldnodes, - fieldbuffers, - bufferoffset, - alignment, -) where {T} - len = length(col) - nc = nullcount(col) - push!(fieldnodes, FieldNode(len, nc)) - @debug "made field node: nodeidx = $(length(fieldnodes)), col = $(typeof(col)), len = $(fieldnodes[end].length), nc = $(fieldnodes[end].null_count)" - # validity bitmap - blen = nc == 0 ? 0 : bitpackedbytes(len, alignment) - push!(fieldbuffers, Buffer(bufferoffset, blen)) - @debug "made field buffer: bufferidx = $(length(fieldbuffers)), offset = $(fieldbuffers[end].offset), len = $(fieldbuffers[end].length), padded = $(padding(fieldbuffers[end].length, alignment))" - bufferoffset += blen - for child in col.data - bufferoffset = - makenodesbuffers!(child, fieldnodes, fieldbuffers, bufferoffset, alignment) - end - return bufferoffset -end - -function writebuffer(io, col::Struct, alignment) - @debug "writebuffer: col = $(typeof(col))" - @debug col - writebitmap(io, col, alignment) - # write values arrays - for child in col.data - writebuffer(io, child, alignment) - end - return -end diff --git a/src/arraytypes/unions.jl b/src/arraytypes/unions.jl deleted file mode 100644 index ab4673d9..00000000 --- a/src/arraytypes/unions.jl +++ /dev/null @@ -1,345 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# Union arrays -# need a custom representation of Union types since arrow unions -# are ordered, and possibly indirected via separate typeIds array -# here, T is Meta.UnionMode.Dense or Meta.UnionMode.Sparse, -# typeIds is a NTuple{N, Int32}, and U is a Tuple{...} of the -# unioned types -struct UnionT{T,typeIds,U} end - -unionmode(::Type{UnionT{T,typeIds,U}}) where {T,typeIds,U} = T -typeids(::Type{UnionT{T,typeIds,U}}) where {T,typeIds,U} = typeIds -Base.eltype(::Type{UnionT{T,typeIds,U}}) where {T,typeIds,U} = U -uniontypewith(::Type{UnionT{T,typeIds,U}}, ::Type{U2}) where {T,typeIds,U,U2<:Tuple} = - UnionT{T,typeIds,U2} - -ArrowTypes.ArrowKind(::Type{<:UnionT}) = ArrowTypes.UnionKind() - -# iterate a Julia Union{...} type, producing an array of unioned types -function eachunion(U::Union, elems=nothing) - if elems === nothing - return eachunion(U.b, Type[U.a]) - else - push!(elems, U.a) - return eachunion(U.b, elems) - end -end - -function eachunion(T, elems) - push!(elems, T) - return elems -end - -# produce typeIds, offsets, data tuple for DenseUnion -isatypeid(x::T, ::Type{types}) where {T,types} = isatypeid(x, fieldtype(types, 1), types, 1) -isatypeid(x::T, ::Type{S}, ::Type{types}, i) where {T,S,types} = - x isa S ? i : isatypeid(x, fieldtype(types, i + 1), types, i + 1) - -""" - Arrow.DenseUnion - -An `ArrowVector` where the type of each element is one of a fixed set of types, meaning its eltype is like a julia `Union{type1, type2, ...}`. -An `Arrow.DenseUnion`, in comparison to `Arrow.SparseUnion`, stores elements in a set of arrays, one array per possible type, and an "offsets" -array, where each offset element is the index into one of the typed arrays. This allows a sort of "compression", where no extra space is -used/allocated to store all the elements. -""" -struct DenseUnion{T,U,S} <: ArrowVector{T} - arrow::Vector{UInt8} # need to hold a reference to arrow memory blob - arrow2::Vector{UInt8} # if arrow blob is compressed, need a 2nd reference for uncompressed offsets bytes - typeIds::Vector{UInt8} - offsets::Vector{Int32} - data::S # Tuple of ArrowVector - metadata::Union{Nothing,Base.ImmutableDict{String,String}} -end - -Base.size(s::DenseUnion) = size(s.typeIds) -nullcount(x::DenseUnion) = 0 # DenseUnion has no validity bitmap; only children do - -@propagate_inbounds function Base.getindex( - s::DenseUnion{T,UnionT{M,typeIds,U}}, - i::Integer, -) where {T,M,typeIds,U} - @boundscheck checkbounds(s, i) - @inbounds typeId = s.typeIds[i] - @inbounds off = s.offsets[i] - @inbounds x = s.data[typeId + 1][off + 1] - return ArrowTypes.fromarrow(fieldtype(U, typeId + 1), x) -end - -# @propagate_inbounds function Base.setindex!(s::DenseUnion{UnionT{T, typeIds, U}}, v, i::Integer) where {T, typeIds, U} -# @boundscheck checkbounds(s, i) -# @inbounds typeId = s.typeIds[i] -# typeids = typeIds === nothing ? (0:(fieldcount(U) - 1)) : typeIds -# vtypeId = Int8(typeids[isatypeid(v, U)]) -# if typeId == vtypeId -# @inbounds off = s.offsets[i] -# @inbounds s.data[typeId +1][off + 1] = v -# else -# throw(ArgumentError("type of item to set $(typeof(v)) must match existing item $(fieldtype(U, typeid))")) -# end -# return v -# end - -# convenience wrappers for signaling that an array shoudld be written -# as with dense/sparse union arrow buffers -struct DenseUnionVector{T,U} <: AbstractVector{UnionT{Meta.UnionMode.Dense,nothing,U}} - itr::T -end - -DenseUnionVector(x::T) where {T} = DenseUnionVector{T,Tuple{eachunion(eltype(x))...}}(x) -Base.IndexStyle(::Type{<:DenseUnionVector}) = Base.IndexLinear() -Base.size(x::DenseUnionVector) = (length(x.itr),) -Base.iterate(x::DenseUnionVector, st...) = iterate(x.itr, st...) -Base.getindex(x::DenseUnionVector, i::Int) = getindex(x.itr, i) - -function todense(::Type{UnionT{T,typeIds,U}}, x) where {T,typeIds,U} - typeids = typeIds === nothing ? (0:(fieldcount(U) - 1)) : typeIds - len = length(x) - types = Vector{UInt8}(undef, len) - offsets = Vector{Int32}(undef, len) - data = Tuple( - Vector{i == 1 ? Union{Missing,fieldtype(U, i)} : fieldtype(U, i)}(undef, 0) for - i = 1:fieldcount(U) - ) - for (i, y) in enumerate(x) - typeid = y === missing ? 0x00 : UInt8(typeids[isatypeid(y, U)]) - @inbounds types[i] = typeid - @inbounds offsets[i] = length(data[typeid + 1]) - push!(data[typeid + 1], y) - end - return types, offsets, data -end - -struct SparseUnionVector{T,U} <: AbstractVector{UnionT{Meta.UnionMode.Sparse,nothing,U}} - itr::T -end - -SparseUnionVector(x::T) where {T} = SparseUnionVector{T,Tuple{eachunion(eltype(x))...}}(x) -Base.IndexStyle(::Type{<:SparseUnionVector}) = Base.IndexLinear() -Base.size(x::SparseUnionVector) = (length(x.itr),) -Base.iterate(x::SparseUnionVector, st...) = iterate(x.itr, st...) -Base.getindex(x::SparseUnionVector, i::Int) = getindex(x.itr, i) - -# sparse union child array producer -# for sparse unions, we split the parent array into -# N children arrays, each having the same length as the parent -# but with one child array per unioned type; each child -# should include the elements from parent of its type -# and other elements can be missing/default -function sparsetypeids(::Type{UnionT{T,typeIds,U}}, x) where {T,typeIds,U} - typeids = typeIds === nothing ? (0:(fieldcount(U) - 1)) : typeIds - len = length(x) - types = Vector{UInt8}(undef, len) - for (i, y) in enumerate(x) - typeid = y === missing ? 0x00 : UInt8(typeids[isatypeid(y, U)]) - @inbounds types[i] = typeid - end - return types -end - -struct ToSparseUnion{T,A} <: AbstractVector{T} - data::A -end - -ToSparseUnion(::Type{T}, data::A) where {T,A} = ToSparseUnion{T,A}(data) - -Base.IndexStyle(::Type{<:ToSparseUnion}) = Base.IndexLinear() -Base.size(x::ToSparseUnion) = (length(x.data),) - -Base.@propagate_inbounds function Base.getindex(A::ToSparseUnion{T}, i::Integer) where {T} - @boundscheck checkbounds(A, i) - @inbounds x = A.data[i] - return @inbounds x isa T ? x : ArrowTypes.default(T) -end - -function compress(Z::Meta.CompressionType.T, comp, x::A) where {A<:DenseUnion} - len = length(x) - nc = nullcount(x) - typeIds = compress(Z, comp, x.typeIds) - offsets = compress(Z, comp, x.offsets) - buffers = [typeIds, offsets] - children = Compressed[] - for y in x.data - push!(children, compress(Z, comp, y)) - end - return Compressed{Z,A}(x, buffers, len, nc, children) -end - -""" - Arrow.SparseUnion - -An `ArrowVector` where the type of each element is one of a fixed set of types, meaning its eltype is like a julia `Union{type1, type2, ...}`. -An `Arrow.SparseUnion`, in comparison to `Arrow.DenseUnion`, stores elements in a set of arrays, one array per possible type, and each typed -array has the same length as the full array. This ends up with "wasted" space, since only one slot among the typed arrays is valid per full -array element, but can allow for certain optimizations when each typed array has the same length. -""" -struct SparseUnion{T,U,S} <: ArrowVector{T} - arrow::Vector{UInt8} # need to hold a reference to arrow memory blob - typeIds::Vector{UInt8} - data::S # Tuple of ArrowVector - metadata::Union{Nothing,Base.ImmutableDict{String,String}} -end - -Base.size(s::SparseUnion) = size(s.typeIds) -nullcount(x::SparseUnion) = 0 - -@propagate_inbounds function Base.getindex( - s::SparseUnion{T,UnionT{M,typeIds,U}}, - i::Integer, -) where {T,M,typeIds,U} - @boundscheck checkbounds(s, i) - @inbounds typeId = s.typeIds[i] - @inbounds x = s.data[typeId + 1][i] - return ArrowTypes.fromarrow(fieldtype(U, typeId + 1), x) -end - -# @propagate_inbounds function Base.setindex!(s::SparseUnion{UnionT{T, typeIds, U}}, v, i::Integer) where {T, typeIds, U} -# @boundscheck checkbounds(s, i) -# typeids = typeIds === nothing ? (0:(fieldcount(U) - 1)) : typeIds -# vtypeId = Int8(typeids[isatypeid(v, U)]) -# @inbounds s.typeIds[i] = vtypeId -# @inbounds s.data[vtypeId + 1][i] = v -# return v -# end - -arrowvector(U::Union, x, i, nl, fi, de, ded, meta; denseunions::Bool=true, kw...) = - arrowvector( - denseunions ? DenseUnionVector(x) : SparseUnionVector(x), - i, - nl, - fi, - de, - ded, - meta; - denseunions=denseunions, - kw..., - ) - -arrowvector( - ::UnionKind, - x::Union{DenseUnion,SparseUnion}, - i, - nl, - fi, - de, - ded, - meta; - kw..., -) = x - -function arrowvector(::UnionKind, x, i, nl, fi, de, ded, meta; kw...) - UT = eltype(x) - if unionmode(UT) == Meta.UnionMode.Dense - x = x isa DenseUnionVector ? x.itr : x - typeids, offsets, data = todense(UT, x) - data2 = map( - y -> arrowvector(y[2], i, nl + 1, y[1], de, ded, nothing; kw...), - enumerate(data), - ) - UT2 = uniontypewith(UT, Tuple{(eltype(x) for x in data2)...}) - return DenseUnion{Union{(eltype(x) for x in data2)...},UT2,typeof(data2)}( - UInt8[], - UInt8[], - typeids, - offsets, - data2, - meta, - ) - else - x = x isa SparseUnionVector ? x.itr : x - typeids = sparsetypeids(UT, x) - data3 = Tuple( - arrowvector( - ToSparseUnion(fieldtype(eltype(UT), j), x), - i, - nl + 1, - j, - de, - ded, - nothing; - kw..., - ) for j = 1:fieldcount(eltype(UT)) - ) - UT2 = uniontypewith(UT, Tuple{(eltype(x) for x in data3)...}) - return SparseUnion{Union{(eltype(x) for x in data3)...},UT2,typeof(data3)}( - UInt8[], - typeids, - data3, - meta, - ) - end -end - -function compress(Z::Meta.CompressionType.T, comp, x::A) where {A<:SparseUnion} - len = length(x) - nc = nullcount(x) - typeIds = compress(Z, comp, x.typeIds) - buffers = [typeIds] - children = Compressed[] - for y in x.data - push!(children, compress(Z, comp, y)) - end - return Compressed{Z,A}(x, buffers, len, nc, children) -end - -function makenodesbuffers!( - col::Union{DenseUnion,SparseUnion}, - fieldnodes, - fieldbuffers, - bufferoffset, - alignment, -) - len = length(col) - nc = nullcount(col) - push!(fieldnodes, FieldNode(len, nc)) - @debug "made field node: nodeidx = $(length(fieldnodes)), col = $(typeof(col)), len = $(fieldnodes[end].length), nc = $(fieldnodes[end].null_count)" - # typeIds buffer - push!(fieldbuffers, Buffer(bufferoffset, len)) - @debug "made field buffer: bufferidx = $(length(fieldbuffers)), offset = $(fieldbuffers[end].offset), len = $(fieldbuffers[end].length), padded = $(padding(fieldbuffers[end].length, alignment))" - bufferoffset += padding(len, alignment) - if col isa DenseUnion - # offsets buffer - blen = sizeof(Int32) * len - push!(fieldbuffers, Buffer(bufferoffset, blen)) - @debug "made field buffer: bufferidx = $(length(fieldbuffers)), offset = $(fieldbuffers[end].offset), len = $(fieldbuffers[end].length), padded = $(padding(fieldbuffers[end].length, alignment))" - bufferoffset += padding(blen, alignment) - end - for child in col.data - bufferoffset = - makenodesbuffers!(child, fieldnodes, fieldbuffers, bufferoffset, alignment) - end - return bufferoffset -end - -function writebuffer(io, col::Union{DenseUnion,SparseUnion}, alignment) - @debug "writebuffer: col = $(typeof(col))" - @debug col - # typeIds buffer - n = writearray(io, UInt8, col.typeIds) - @debug "writing array: col = $(typeof(col.typeIds)), n = $n, padded = $(padding(n, alignment))" - writezeros(io, paddinglength(n, alignment)) - if col isa DenseUnion - n = writearray(io, Int32, col.offsets) - @debug "writing array: col = $(typeof(col.offsets)), n = $n, padded = $(padding(n, alignment))" - writezeros(io, paddinglength(n, alignment)) - end - for child in col.data - writebuffer(io, child, alignment) - end - return -end diff --git a/src/arraytypes/views.jl b/src/arraytypes/views.jl deleted file mode 100644 index 0a43f6fc..00000000 --- a/src/arraytypes/views.jl +++ /dev/null @@ -1,80 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -struct ViewElement - length::Int32 - prefix::Int32 - bufindex::Int32 - offset::Int32 -end - -""" - Arrow.View - -An `ArrowVector` where each element is a variable sized list of some kind, like an `AbstractVector` or `AbstractString`. -""" -struct View{T} <: ArrowVector{T} - arrow::Vector{UInt8} # need to hold a reference to arrow memory blob - validity::ValidityBitmap - data::Vector{ViewElement} - inline::Vector{UInt8} # `data` field reinterpreted as a byte array - buffers::Vector{Vector{UInt8}} # holds non-inlined data - ℓ::Int - metadata::Union{Nothing,Base.ImmutableDict{String,String}} -end - -Base.size(l::View) = (l.ℓ,) - -@propagate_inbounds function Base.getindex(l::View{T}, i::Integer) where {T} - @boundscheck checkbounds(l, i) - @inbounds v = l.data[i] - S = Base.nonmissingtype(T) - if S <: Base.CodeUnits - # BinaryView - return !l.validity[i] ? missing : - v.length < 13 ? - Base.CodeUnits( - StringView( - @view l.inline[(((i - 1) * 16) + 5):(((i - 1) * 16) + 5 + v.length - 1)] - ), - ) : - Base.CodeUnits( - StringView( - @view l.buffers[v.bufindex + 1][(v.offset + 1):(v.offset + v.length)] - ), - ) - else - # Utf8View - return !l.validity[i] ? missing : - v.length < 13 ? - ArrowTypes.fromarrow( - T, - StringView( - @view l.inline[(((i - 1) * 16) + 5):(((i - 1) * 16) + 5 + v.length - 1)] - ), - ) : - ArrowTypes.fromarrow( - T, - StringView( - @view l.buffers[v.bufindex + 1][(v.offset + 1):(v.offset + v.length)] - ), - ) - end -end - -# @propagate_inbounds function Base.setindex!(l::List{T}, v, i::Integer) where {T} - -# end diff --git a/src/arrowtypes.jl b/src/arrowtypes.jl new file mode 100644 index 00000000..5faaad13 --- /dev/null +++ b/src/arrowtypes.jl @@ -0,0 +1,1303 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# ArrowTypes is a facade concern. ArrowCore stays dependency- and +# conversion-free; this layer interprets the standard extension metadata on +# materialized values and supplies the shared metadata helpers used by the +# writer. An unregistered extension type stays an ordinary Arrow storage +# value; the format requires readers to accept that. + +const _EXTENSION_NAME_KEY = "ARROW:extension:name" +const _EXTENSION_METADATA_KEY = "ARROW:extension:metadata" +const _MAX_ARROWTYPE_EXACT_ARITY = 1024 +const _MAX_ARROWTYPE_UNION_BRANCHES = 32 +const _MAX_ARROWTYPE_SCHEMA_NAME_BYTES = 4096 +const _MAX_ARROWTYPE_STRUCT_NAME_BYTES = 64 * 1024 +const _MAX_EXTENSION_WARNING_BYTES = 128 +const _MAX_EXTENSION_WARNINGS = 16 + +# Look a Symbol up without creating one: interning an input-controlled name +# would permanently allocate process-global state. +function _existingjlsymbol(name::AbstractString) + occursin('\0', name) && return nothing + ncodeunits(name) <= _MAX_ARROWTYPE_SCHEMA_NAME_BYTES || return nothing + pointer = ccall(:jl_symbol_lookup, Ptr{Cvoid}, (Cstring,), name) + return pointer == C_NULL ? nothing : Symbol(name) +end + +function _boundedfixedliststoragetype(listsize::Int, element) + # Julia materializes an NTuple type's parameter list in O(N) space. Keep + # the exact storage signature for ordinary fixed lists, but use the compact + # tuple-family type for a hostile or unusually large descriptor. + return listsize <= _MAX_ARROWTYPE_EXACT_ARITY ? NTuple{listsize,element} : + Tuple{Vararg{element}} +end + +function _fieldmetavalue(f::AC.Field, key::String, default=nothing) + value = default + f.metadata === nothing && return value + for kv in f.metadata + first(kv) == key && (value = last(kv)) + end + return value +end + +function _arrowtypesextension(f::AC.Field) + name = _fieldmetavalue(f, _EXTENSION_NAME_KEY) + name === nothing && return nothing + metadata = _fieldmetavalue(f, _EXTENSION_METADATA_KEY, "") + return String(name), String(metadata) +end + +function _withoutownextension(f::AC.Field, t::AC.ArrowType=f.type; nullable=f.nullable) + metadata = + f.metadata === nothing ? nothing : + Pair{String,String}[ + String(first(kv)) => String(last(kv)) for kv in f.metadata if + first(kv) != _EXTENSION_NAME_KEY && first(kv) != _EXTENSION_METADATA_KEY + ] + return AC.Field(f.name, t; nullable=nullable, metadata=metadata, children=f.children) +end + +_withmissingtype(T, nullable::Bool) = nullable && T !== Missing ? Union{Missing,T} : T +_nonmissingstoragetype(T) = T === Missing ? Missing : Base.nonmissingtype(T) + +@inline function _ispositionalchildname(name::String, index::Int) + bytes = codeunits(name) + length(bytes) == ndigits(index) || return false + value = 0 + for byte in bytes + 0x30 <= byte <= 0x39 || return false + value = 10 * value + Int(byte - 0x30) + end + return value == index +end + +function _arrowtypesstructnames( + f::AC.Field, + budget::Union{Nothing,AllocationBudget}=nothing, +) + length(f.children) <= _MAX_ARROWTYPE_EXACT_ARITY || throw( + ArgumentError( + "extension struct $(f.name) has more than " * + "$_MAX_ARROWTYPE_EXACT_ARITY children and cannot form an exact Julia " * + "NamedTuple storage type", + ), + ) + nchildren = length(f.children) + _chargevector!(budget, Symbol, nchildren, "ArrowTypes struct-name workspace") + _chargedict!(budget, String, Nothing, nchildren, "ArrowTypes struct-name duplicate set") + names = Symbol[] + seen = Set{String}() + sizehint!(names, nchildren) + sizehint!(seen, nchildren) + totalbytes = 0 + # ArrowTypes represents positional Tuple storage as Struct children named + # "1", "2", ... . Preserve that established interface without opening an + # unbounded input-to-Symbol path: at most the fixed set 1:1024 can be added. + positional = all( + index -> _ispositionalchildname(f.children[index].name, index), + eachindex(f.children), + ) + for child in f.children + name = child.name + occursin('\0', name) && throw( + ArgumentError( + "extension struct $(f.name) has a child name with an embedded NUL", + ), + ) + namebytes = ncodeunits(name) + namebytes <= _MAX_ARROWTYPE_SCHEMA_NAME_BYTES || throw( + ArgumentError( + "extension struct $(f.name) has a child name longer than " * + "$_MAX_ARROWTYPE_SCHEMA_NAME_BYTES bytes", + ), + ) + totalbytes <= _MAX_ARROWTYPE_STRUCT_NAME_BYTES - namebytes || throw( + ArgumentError( + "extension struct $(f.name) child names exceed the supported " * + "$_MAX_ARROWTYPE_STRUCT_NAME_BYTES-byte total", + ), + ) + totalbytes += namebytes + name in seen && throw( + ArgumentError( + "extension struct $(f.name) has duplicate child names and cannot " * + "form a Julia NamedTuple storage type", + ), + ) + push!(seen, name) + symbol = _existingjlsymbol(name) + if symbol === nothing + positional || throw( + ArgumentError( + "extension struct $(f.name) has an unregistered child name " * + "$(repr(name)); exact NamedTuple lifting would permanently intern it", + ), + ) + _chargeobject!(budget, namebytes, "ArrowTypes positional child-name symbol") + symbol = Symbol(name) + end + push!(names, symbol) + end + _chargeobject!( + budget, + AC.checked_mul(Int64(Base.elsize(Vector{Symbol})), Int64(nchildren)), + "ArrowTypes struct-name tuple", + ) + return Tuple(names) +end + +"ArrowTypes' Julia storage type for one non-composite Field." +function _arrowtypesprimitivebasetype(f::AC.Field) + t = f.type + t isa AC.NullType && return Missing + t isa AC.IntType && return AC.juliatype(t) + t isa AC.FloatType && return AC.juliatype(t) + t isa AC.BoolType && return Bool + t isa AC.Utf8Type && return String + t isa AC.BinaryType && return Vector{UInt8} + t isa AC.FixedSizeBinaryType && return Vector{UInt8} + t isa AC.ViewType && return t.utf8 ? String : Vector{UInt8} + if t isa AC.DateType + return t.unit == AC.DAY ? Dates.Date : Dates.DateTime + end + if t isa AC.TimestampType + return t.unit == AC.SECOND || t.unit == AC.MILLISECOND ? Dates.DateTime : Int64 + end + t isa AC.TimeType && return Dates.Time + if t isa AC.DurationType + return t.unit == AC.SECOND ? Dates.Second : + t.unit == AC.MILLISECOND ? Dates.Millisecond : + t.unit == AC.MICROSECOND ? Dates.Microsecond : Dates.Nanosecond + end + if t isa AC.DecimalType + return t.bits == 32 ? Int32 : t.bits == 64 ? Int64 : Vector{UInt8} + end + if t isa AC.IntervalType + return t.unit == AC.YEAR_MONTH ? Int32 : + t.unit == AC.DAY_TIME ? NamedTuple{(:days, :millis),Tuple{Int32,Int32}} : + NamedTuple{(:months, :days, :nanos),Tuple{Int32,Int32,Int64}} + end + return Any +end + +function _extensionwarningname( + name::AbstractString, + budget::Union{Nothing,AllocationBudget}=nothing, +) + ncodeunits(name) <= _MAX_EXTENSION_WARNING_BYTES && return name + bytes = 0 + lastindex = 0 + for index in eachindex(name) + width = nextind(name, index) - index + bytes + width > _MAX_EXTENSION_WARNING_BYTES && break + bytes += width + lastindex = index + end + lastindex == 0 && return "…" + _chargeobject!(budget, bytes + ncodeunits("…"), "ArrowTypes extension-warning display") + return string(SubString(name, firstindex(name), lastindex), '…') +end + +function _warnunsupportedextension(ctx, name::String, f::AC.Field) + ctx.warn || return nothing + name in ctx.warned && return nothing + if ctx.warnings >= _MAX_EXTENSION_WARNINGS + if !ctx.warningsuppressed + ctx.warningsuppressed = true + @warn "additional unsupported Arrow extension warnings were suppressed" + end + return nothing + end + _chargedictentry!(ctx.budget, String, Nothing, "ArrowTypes warning cache") + push!(ctx.warned, name) + ctx.warnings += 1 + displayname = _extensionwarningname(name, ctx.budget) + descriptor = AC.descriptorname(f.type) + # Charge only the message this adapter builds. The logging framework's + # own records and rendering are outside this budget. + _chargeobject!( + ctx.budget, + AC.checked_add( + Int64(128), + AC.checked_mul(Int64(4), Int64(ncodeunits(displayname))), + ), + "ArrowTypes extension-warning message", + ) + @warn "unsupported ARROW:extension:name type: $(repr(displayname)), storage descriptor = $descriptor" + return nothing +end + +function _hasarrowtypesextension(f::AC.Field) + _arrowtypesextension(f) === nothing || return true + return any(_hasarrowtypesextension, f.children) +end + +# Core's ordinary materializer deliberately returns only storage-domain +# values. Two facade-only facts still have to survive until ArrowTypes lifting: +# a Union's selected child, and whether a Dictionary row was a valid +# index into its null pool rather than a null index. Keep those facts on this +# private route; the facade consumes every marker before returning a column. +struct _ArrowTypesRoutedUnion{T} + child::Int + value::T +end + +struct _ArrowTypesRoutedNull end + +mutable struct _ArrowTypesRoutePlan + extensions::Dict{AC.Field,Bool} + routes::Dict{AC.Field,Bool} + dictionaryfields::Dict{AC.Field,AC.Field} + budget::Union{Nothing,AllocationBudget} +end +function _ArrowTypesRoutePlan( + budget::Union{Nothing,AllocationBudget}=nothing, + extensions::Union{Nothing,Dict{AC.Field,Bool}}=nothing, +) + _chargeobject!(budget, sizeof(_ArrowTypesRoutePlan), "ArrowTypes route-plan owner") + if extensions === nothing + _chargeemptydict!( + budget, + AC.Field, + Bool, + "ArrowTypes extension-presence cache container", + ) + extensions = Dict{AC.Field,Bool}() + end + _chargeemptydict!(budget, AC.Field, Bool, "ArrowTypes route cache container") + _chargeemptydict!( + budget, + AC.Field, + AC.Field, + "ArrowTypes routed dictionary-field cache container", + ) + return _ArrowTypesRoutePlan( + extensions, + Dict{AC.Field,Bool}(), + Dict{AC.Field,AC.Field}(), + budget, + ) +end + +function _hasarrowtypesextension(f::AC.Field, plan::_ArrowTypesRoutePlan) + return _memoized!( + plan.extensions, + f, + plan.budget, + "ArrowTypes extension-presence cache", + ) do + _arrowtypesextension(f) === nothing || return true + return any(child -> _hasarrowtypesextension(child, plan), f.children) + end +end + +function _routedictionaryfield!( + plan::_ArrowTypesRoutePlan, + f::AC.Field, + t::AC.DictionaryType, +) + return _memoized!( + plan.dictionaryfields, + f, + plan.budget, + "ArrowTypes routed dictionary-field cache", + ) do + _arrowtypesdictvaluefield(f, t; retainmetadata=true, budget=plan.budget) + end +end + +function _needsarrowtypesroute(f::AC.Field, plan::_ArrowTypesRoutePlan) + return _memoized!(plan.routes, f, plan.budget, "ArrowTypes route cache") do + t = f.type + if t isa AC.DictionaryType + ownnull = t.valuetype isa AC.NullType && _arrowtypesextension(f) !== nothing + return ownnull || + _needsarrowtypesroute(_routedictionaryfield!(plan, f, t), plan) + elseif t isa AC.RunEndEncodedType + return length(f.children) == 2 && _needsarrowtypesroute(f.children[2], plan) + elseif t isa AC.UnionType && _hasarrowtypesextension(f, plan) + return true + end + return any(child -> _needsarrowtypesroute(child, plan), f.children) + end +end + +function _arrowtypesroutedvalue( + f::AC.Field, + d::AC.ArrayData, + i::Int64, + plan::_ArrowTypesRoutePlan, + routeallunions::Bool=false, + budget::Union{Nothing,AllocationBudget}=nothing, +) + (routeallunions || _needsarrowtypesroute(f, plan)) || + return budget === nothing ? AC.getvalue(f, d, i) : AC.getvalue(f, d, i, budget) + t = f.type + if t isa AC.DictionaryType + AC.isvalid_at(d, i) || return missing + w = AC.primwidth(t.indextype) + index = + AC._load_int(AC.rolebuffer(d, AC.DATA), t.indextype, AC._slotbyteoff(d, i, w)) + pool = d.dictionary + pool === nothing && throw(AC.ValidationError("dictionary array has no value pool")) + value = _arrowtypesroutedvalue( + _routedictionaryfield!(plan, f, t), + pool, + AC.checked_add(Int64(index), Int64(1)), + plan, + routeallunions, + budget, + ) + # A NullType pool's physical value is always `missing`. The valid index + # above makes it a logical extension value; an invalid outer index has + # already returned ordinary `missing` and must remain distinct. + return t.valuetype isa AC.NullType && _arrowtypesextension(f) !== nothing ? + _ArrowTypesRoutedNull() : value + elseif t isa AC.RunEndEncodedType + length(f.children) == 2 || + return budget === nothing ? AC.getvalue(f, d, i) : AC.getvalue(f, d, i, budget) + return _arrowtypesroutedvalue( + f.children[2], + d.children[2], + AC._ree_runindex(d, i), + plan, + routeallunions, + budget, + ) + elseif t isa AC.UnionType + childfield, childdata, childindex = AC._union_child(f, d, i) + child = findfirst(x -> x === childfield, f.children) + child === nothing && throw(AC.ValidationError("union selected an undeclared child")) + value = _arrowtypesroutedvalue( + childfield, + childdata, + childindex, + plan, + routeallunions, + budget, + ) + routed = _ArrowTypesRoutedUnion(child, value) + _chargeobject!(budget, sizeof(typeof(routed)), "ArrowTypes routed union value") + return routed + elseif t isa AC.ListType + AC.isvalid_at(d, i) || return missing + lo, hi = AC._offsets_at(d, i, AC.layoutspec(t).offsetwidth == 8) + childfield, childdata = f.children[1], d.children[1] + _chargevector!(budget, Any, hi - lo, "ArrowTypes routed list value") + out = Vector{Any}(undef, Int(hi - lo)) + for k = 1:length(out) + out[k] = _arrowtypesroutedvalue( + childfield, + childdata, + AC.checked_add(lo, Int64(k)), + plan, + routeallunions, + budget, + ) + end + return out + elseif t isa AC.ListViewType + AC.isvalid_at(d, i) || return missing + off, size = AC._listview_range(t, d, i) + childfield, childdata = f.children[1], d.children[1] + _chargevector!(budget, Any, size, "ArrowTypes routed list-view value") + out = Vector{Any}(undef, Int(size)) + for k = 1:length(out) + out[k] = _arrowtypesroutedvalue( + childfield, + childdata, + AC.checked_add(off, Int64(k)), + plan, + routeallunions, + budget, + ) + end + return out + elseif t isa AC.FixedSizeListType + AC.isvalid_at(d, i) || return missing + childfield, childdata = f.children[1], d.children[1] + base = AC.checked_mul(AC._slotindex0(d, i), Int64(t.listsize)) + _chargevector!(budget, Any, t.listsize, "ArrowTypes routed fixed-list value") + out = Vector{Any}(undef, t.listsize) + for k = 1:(t.listsize) + out[k] = _arrowtypesroutedvalue( + childfield, + childdata, + AC.checked_add(base, Int64(k)), + plan, + routeallunions, + budget, + ) + end + return out + elseif t isa AC.StructType + AC.isvalid_at(d, i) || return missing + childindex = AC.checked_add(d.offset, i) + _chargevector!( + budget, + Pair{String,Any}, + length(f.children), + "ArrowTypes routed struct value", + ) + out = Vector{Pair{String,Any}}(undef, length(f.children)) + for k in eachindex(f.children) + out[k] = + f.children[k].name => _arrowtypesroutedvalue( + f.children[k], + d.children[k], + childindex, + plan, + routeallunions, + budget, + ) + end + return out + elseif t isa AC.MapType + AC.isvalid_at(d, i) || return missing + lo, hi = AC._offsets_at(d, i, false) + entriesfield, entriesdata = f.children[1], d.children[1] + keyfield, valuefield = entriesfield.children + keydata, valuedata = entriesdata.children + _chargevector!(budget, Pair{Any,Any}, hi - lo, "ArrowTypes routed map value") + out = Vector{Pair{Any,Any}}(undef, Int(hi - lo)) + for k = 1:length(out) + entryindex = AC.checked_add(entriesdata.offset, AC.checked_add(lo, Int64(k))) + out[k] = Pair{Any,Any}( + _arrowtypesroutedvalue( + keyfield, + keydata, + entryindex, + plan, + routeallunions, + budget, + ), + _arrowtypesroutedvalue( + valuefield, + valuedata, + entryindex, + plan, + routeallunions, + budget, + ), + ) + end + return out + end + return budget === nothing ? AC.getvalue(f, d, i) : AC.getvalue(f, d, i, budget) +end + +"Materialize only when facade lifting needs storage provenance from the layout." +function _arrowtypesroutedcolumn(f::AC.Field, d::AC.ArrayData, plan::_ArrowTypesRoutePlan) + _needsarrowtypesroute(f, plan) || return nothing + budget = plan.budget + _chargevector!(budget, Any, d.len, "ArrowTypes routed column") + return Any[ + _arrowtypesroutedvalue(f, d, Int64(i), plan, false, budget) for i = 1:(d.len) + ] +end + +_arrowtypesroutedcolumn( + f::AC.Field, + d::AC.ArrayData, + budget::Union{Nothing,AllocationBudget}=nothing, +) = _arrowtypesroutedcolumn(f, d, _ArrowTypesRoutePlan(budget)) + +"Per read/write operation ArrowTypes resolution; each Field calls JuliaType at most once." +mutable struct _ArrowTypesContext + targets::Dict{AC.Field,Any} + logicaltypes::Dict{AC.Field,Any} + publictypes::Dict{AC.Field,Any} + extensions::Dict{AC.Field,Bool} + structnames::Dict{AC.Field,Tuple} + storage::IdDict{Type,Any} + dictionaryfields::Dict{AC.Field,AC.Field} + metadata_dictionaryfields::Dict{AC.Field,AC.Field} + public_dictionaryfields::Dict{AC.Field,AC.Field} + routeplan::Union{Nothing,_ArrowTypesRoutePlan} + warn::Bool + warned::Set{String} + warnings::Int + warningsuppressed::Bool + budget::Union{Nothing,AllocationBudget} +end +function _ArrowTypesContext(; + warn::Bool=true, + budget::Union{Nothing,AllocationBudget}=nothing, +) + _chargeobject!(budget, sizeof(_ArrowTypesContext), "ArrowTypes context owner") + _chargeemptydict!(budget, AC.Field, Any, "ArrowTypes target cache container") + _chargeemptydict!(budget, AC.Field, Any, "ArrowTypes logical-type cache container") + _chargeemptydict!(budget, AC.Field, Any, "ArrowTypes public-type cache container") + _chargeemptydict!( + budget, + AC.Field, + Bool, + "ArrowTypes extension-presence cache container", + ) + _chargeemptydict!(budget, AC.Field, Tuple, "ArrowTypes struct-name cache container") + _chargeemptydict!(budget, Type, Any, "ArrowTypes storage-type cache container") + _chargeemptydict!( + budget, + AC.Field, + AC.Field, + "ArrowTypes dictionary-field cache container", + ) + _chargeemptydict!( + budget, + AC.Field, + AC.Field, + "ArrowTypes metadata dictionary-field cache container", + ) + _chargeemptydict!( + budget, + AC.Field, + AC.Field, + "ArrowTypes public dictionary-field cache container", + ) + _chargeemptydict!(budget, String, Nothing, "ArrowTypes warning cache container") + _chargeobject!(budget, sizeof(Set{String}), "ArrowTypes warning-set owner") + return _ArrowTypesContext( + Dict{AC.Field,Any}(), + Dict{AC.Field,Any}(), + Dict{AC.Field,Any}(), + Dict{AC.Field,Bool}(), + Dict{AC.Field,Tuple}(), + IdDict{Type,Any}(), + Dict{AC.Field,AC.Field}(), + Dict{AC.Field,AC.Field}(), + Dict{AC.Field,AC.Field}(), + nothing, + warn, + Set{String}(), + 0, + false, + budget, + ) +end + +function _hasarrowtypesextension(f::AC.Field, ctx::_ArrowTypesContext) + return _memoized!( + ctx.extensions, + f, + ctx.budget, + "ArrowTypes extension-presence cache", + ) do + _arrowtypesextension(f) === nothing || return true + return any(child -> _hasarrowtypesextension(child, ctx), f.children) + end +end + +function _arrowtypesstructnames(ctx::_ArrowTypesContext, f::AC.Field) + return _memoized!(ctx.structnames, f, ctx.budget, "ArrowTypes struct-name cache") do + _arrowtypesstructnames(f, ctx.budget) + end +end + +function _arrowtypesrouteplan!(ctx::_ArrowTypesContext) + plan = ctx.routeplan + plan === nothing || return plan + plan = _ArrowTypesRoutePlan(ctx.budget, ctx.extensions) + ctx.routeplan = plan + return plan +end + +Base.@noinline function _arrowtypesrawstoragetype(T::Type) + Base.@nospecialize T + return Base.inferencebarrier(ArrowTypes.ArrowType(Base.inferencebarrier(T))) +end + +function _checkarrowtypesstorageshape!( + owner::Type, + storage, + seen::Base.IdSet{Any}, + depth::Int, +) + Base.@nospecialize owner storage + storage isa Type || return nothing + storage in seen && return nothing + depth <= 64 || throw( + ArgumentError( + "ArrowTypes.ArrowType($owner) returned a storage type nested beyond " * + "the supported depth 64", + ), + ) + push!(seen, storage) + if storage isa Union + variants = Base.uniontypes(storage) + length(variants) <= _MAX_ARROWTYPE_UNION_BRANCHES || throw( + ArgumentError( + "ArrowTypes.ArrowType($owner) returned a Union with more than " * + "$_MAX_ARROWTYPE_UNION_BRANCHES branches", + ), + ) + for variant in variants + _checkarrowtypesstorageshape!(owner, variant, seen, depth + 1) + end + elseif storage isa DataType + if storage <: Tuple && isconcretetype(storage) + nfields = fieldcount(storage) + nfields <= _MAX_ARROWTYPE_EXACT_ARITY || throw( + ArgumentError( + "ArrowTypes.ArrowType($owner) returned a concrete Tuple with " * + "$nfields fields; the supported limit is " * + "$_MAX_ARROWTYPE_EXACT_ARITY", + ), + ) + end + for parameter in storage.parameters + parameter isa Type || continue + _checkarrowtypesstorageshape!(owner, parameter, seen, depth + 1) + end + end + return nothing +end + +function _validatedarrowtypesstoragetype(T::Type, storage) + Base.@nospecialize T storage + storage isa Type || throw( + ArgumentError( + "ArrowTypes.ArrowType($T) must return a Julia type, not $(repr(storage))", + ), + ) + _checkarrowtypesstorageshape!(T, storage, Base.IdSet{Any}(), 0) + return storage +end + +Base.@noinline function _arrowtypesstoragetype!(ctx::_ArrowTypesContext, T::Type) + Base.@nospecialize T + return _memoized!(ctx.storage, T, ctx.budget, "ArrowTypes storage-type cache") do + _validatedarrowtypesstoragetype(T, _arrowtypesrawstoragetype(T)) + end +end + +function _arrowtypesstoragebasetype(ctx::_ArrowTypesContext, f::AC.Field) + t = f.type + if t isa AC.DictionaryType + valuefield = _arrowtypesdictvaluefield(ctx, f, t) + return _arrowtypesstoragebasetype(ctx, valuefield) + end + if t isa AC.RunEndEncodedType + length(f.children) == 2 || return Any + return Base.nonmissingtype(_arrowtypeslogicaleltype(ctx, f.children[2])) + end + if t isa Union{AC.ListType,AC.ListViewType} + length(f.children) == 1 || return Vector{Any} + return Vector{_arrowtypeslogicaleltype(ctx, f.children[1])} + end + if t isa AC.FixedSizeListType + length(f.children) == 1 || return NTuple{0,Any} + element = _arrowtypeslogicaleltype(ctx, f.children[1]) + # Generic JuliaType registrations still resolve from the bounded + # signature. A hook that requires an exact oversized arity safely + # remains unregistered. + return _boundedfixedliststoragetype(t.listsize, element) + end + if t isa AC.StructType + length(f.children) > _MAX_ARROWTYPE_EXACT_ARITY && return Vector{Pair{String,Any}} + names = _arrowtypesstructnames(ctx, f) + _chargevector!( + ctx.budget, + Type, + length(f.children), + "ArrowTypes struct type workspace", + ) + types = Type[_arrowtypeslogicaleltype(ctx, child) for child in f.children] + return Core.apply_type(NamedTuple, names, Core.apply_type(Tuple, types...)) + end + if t isa AC.MapType + length(f.children) == 1 || return Dict{Any,Any} + entries = f.children[1] + length(entries.children) == 2 || return Dict{Any,Any} + K = _arrowtypeslogicaleltype(ctx, entries.children[1]) + V = _arrowtypeslogicaleltype(ctx, entries.children[2]) + return Dict{K,V} + end + if t isa AC.UnionType + isempty(f.children) && return Union{} + T = _arrowtypeslogicaleltype(ctx, f.children[1]) + for child in Iterators.drop(f.children, 1) + T = Union{T,_arrowtypeslogicaleltype(ctx, child)} + end + return T + end + return _arrowtypesprimitivebasetype(f) +end + +function _arrowtypesfallbackstoragebasetype(ctx::_ArrowTypesContext, f::AC.Field) + t = f.type + if t isa AC.DictionaryType + valuefield = _arrowtypesdictvaluefield(ctx, f, t) + return _arrowtypesfallbackstoragebasetype(ctx, valuefield) + end + if t isa AC.RunEndEncodedType + return length(f.children) == 2 ? + _arrowtypesfallbackstoragebasetype(ctx, f.children[2]) : Any + end + t isa Union{AC.ListType,AC.ListViewType} && return Vector{Any} + t isa AC.FixedSizeListType && return Tuple{Vararg{Any}} + t isa AC.StructType && return Vector{Pair{String,Any}} + t isa AC.MapType && return Dict{Any,Any} + t isa AC.UnionType && return Any + return _arrowtypesprimitivebasetype(f) +end + +function _preflightarrowtypesstructnames(ctx::_ArrowTypesContext, f::AC.Field) + f.type isa AC.StructType && _arrowtypesstructnames(ctx, f) + foreach(child -> _preflightarrowtypesstructnames(ctx, child), f.children) + return nothing +end + +function _arrowtypestarget(ctx::_ArrowTypesContext, f::AC.Field) + return _memoized!(ctx.targets, f, ctx.budget, "ArrowTypes target cache") do + ext = _arrowtypesextension(f) + # Most Fields have no extension label. Do not synthesize their full + # ArrowTypes storage shape merely to discover that no JuliaType hook + # can apply. In particular, forming NTuple{N,T} for an unlabelled + # FixedSizeList must not allocate in proportion to an untrusted N. + ext === nothing && return (false, nothing, nothing) + name, metadata = ext + sym = _existingjlsymbol(name) + if sym === nothing + storage = _arrowtypesfallbackstoragebasetype(ctx, f) + _warnunsupportedextension(ctx, name, f) + return (true, nothing, storage) + end + try + _preflightarrowtypesstructnames(ctx, f) + catch err + err isa ArgumentError || rethrow() + storage = _arrowtypesfallbackstoragebasetype(ctx, f) + _warnunsupportedextension(ctx, name, f) + return (true, nothing, storage) + end + storage = _arrowtypesstoragebasetype(ctx, f) + target = ArrowTypes.JuliaType(Val(sym), _nonmissingstoragetype(storage), metadata) + target === nothing || + target isa Type || + throw( + ArgumentError( + "ArrowTypes.JuliaType for extension $(repr(name)) on field " * + "$(f.name) must return a Julia type or nothing, not " * + "$(repr(target))", + ), + ) + target === nothing && _warnunsupportedextension(ctx, name, f) + return (true, target, storage) + end +end + +function _arrowtypeslogicalnullable(f::AC.Field, target, storage=nothing) + if f.type isa AC.UnionType + storage === nothing && (storage = ArrowTypes.ArrowType(target)) + end + # A nullable logical column adds exactly one Null child around the + # target's storage Union, so one extra field child means "outer nullable"; + # equal counts mean not nullable, and anything else is a mismatch. + if storage isa Union + fieldbranches = length(f.children) + storagebranches = length(Base.uniontypes(storage)) + fieldbranches == storagebranches && return false + if fieldbranches == storagebranches + 1 + Missing <: storage && throw( + ArgumentError( + "registered ArrowTypes target $target uses Missing as a storage " * + "Union branch, so retained field $(f.name) cannot add an " * + "indistinguishable outer missing branch", + ), + ) + return true + end + throw( + ArgumentError( + "registered ArrowTypes target $target has $storagebranches storage " * + "Union branches, but retained field $(f.name) has $fieldbranches children", + ), + ) + end + return f.nullable +end + +_arrowtypeslogicalnullable(ctx::_ArrowTypesContext, f::AC.Field, target) = + _arrowtypeslogicalnullable( + f, + target, + f.type isa AC.UnionType ? _arrowtypesstoragetype!(ctx, target) : nothing, + ) + +function _arrowtypeslogicaleltype(ctx::_ArrowTypesContext, f::AC.Field) + return _memoized!(ctx.logicaltypes, f, ctx.budget, "ArrowTypes logical-type cache") do + _, target, storage = _arrowtypestarget(ctx, f) + storage === nothing && (storage = _arrowtypesstoragebasetype(ctx, f)) + nullable = + target === nothing ? f.nullable : _arrowtypeslogicalnullable(ctx, f, target) + _withmissingtype(target === nothing ? storage : target, nullable) + end +end + +function _arrowtypespubliceltype(ctx::_ArrowTypesContext, f::AC.Field) + return _memoized!(ctx.publictypes, f, ctx.budget, "ArrowTypes public-type cache") do + _, target, storage = _arrowtypestarget(ctx, f) + if target !== nothing + nullable = + !(f.type isa AC.NullType) && _arrowtypeslogicalnullable(ctx, f, target) + return _withmissingtype(target, nullable) + end + t = f.type + if t isa AC.DictionaryType + valuefield = _arrowtypespublicdictvaluefield(ctx, f, t) + return _withmissingtype( + Base.nonmissingtype(_arrowtypespubliceltype(ctx, valuefield)), + f.nullable, + ) + elseif t isa AC.RunEndEncodedType + length(f.children) == 2 || return Any + return _arrowtypespubliceltype(ctx, f.children[2]) + elseif t isa Union{AC.ListType,AC.ListViewType,AC.FixedSizeListType} + length(f.children) == 1 || return _withmissingtype(Vector{Any}, f.nullable) + return _withmissingtype( + Vector{_arrowtypespubliceltype(ctx, f.children[1])}, + f.nullable, + ) + elseif t isa AC.StructType + foreach(child -> _arrowtypespubliceltype(ctx, child), f.children) + return _withmissingtype(Vector{Pair{String,Any}}, f.nullable) + elseif t isa AC.MapType + foreach(child -> _arrowtypespubliceltype(ctx, child), f.children) + return _withmissingtype(Vector{Pair{Any,Any}}, f.nullable) + elseif t isa AC.UnionType + isempty(f.children) && return Union{} + T = _arrowtypespubliceltype(ctx, f.children[1]) + for child in Iterators.drop(f.children, 1) + T = Base.promote_typejoin(T, _arrowtypespubliceltype(ctx, child)) + end + return T + end + storage === nothing && (storage = _arrowtypesstoragebasetype(ctx, f)) + return _withmissingtype(storage, f.nullable) + end +end + +function _typedvalues(::Type{Any}, values, budget=nothing) + _chargevector!(budget, Any, length(values), "ArrowTypes public column") + # map(identity, …) keeps the input's own narrowed eltype for an Any + # target instead of widening every column to Vector{Any}. + return map(identity, values) +end +function _typedvalues(::Type{T}, values, budget=nothing) where {T} + _chargevector!(budget, T, length(values), "ArrowTypes typed column") + return T[x for x in values] +end + +function _arrowtypesscalar(t::AC.ArrowType, x) + x === missing && return missing + if t isa AC.DateType + return t.unit == AC.DAY ? Dates.Date(Dates.UTD(Int64(x) + _EPOCH_DAYS)) : + Dates.DateTime(Dates.UTM(Int64(x) + Dates.UNIXEPOCH)) + end + if t isa AC.TimestampType + # DateTime cannot hold micro/nanosecond precision, so those units + # stay raw Int64 rather than lose information. + t.unit == AC.SECOND && + return Dates.DateTime(Dates.UTM(Int64(x) * 1000 + Dates.UNIXEPOCH)) + t.unit == AC.MILLISECOND && + return Dates.DateTime(Dates.UTM(Int64(x) + Dates.UNIXEPOCH)) + return x + end + if t isa AC.TimeType + scale = + t.unit == AC.SECOND ? Int64(1_000_000_000) : + t.unit == AC.MILLISECOND ? Int64(1_000_000) : + t.unit == AC.MICROSECOND ? Int64(1_000) : Int64(1) + return Dates.Time(Dates.Nanosecond(Int64(x) * scale)) + end + if t isa AC.DurationType + P = + t.unit == AC.SECOND ? Dates.Second : + t.unit == AC.MILLISECOND ? Dates.Millisecond : + t.unit == AC.MICROSECOND ? Dates.Microsecond : Dates.Nanosecond + return P(Int64(x)) + end + return x +end + +function _arrowtypesdictvaluefield( + f::AC.Field, + t::AC.DictionaryType; + retainmetadata::Bool=false, + nullable::Bool=true, + budget::Union{Nothing,AllocationBudget}=nothing, +) + metadata = retainmetadata ? f.metadata : nothing + # Field and FrozenVector are immutable. Use Core's exact positional + # constructor so every reader view shares the already-frozen child tree. + _chargeobject!(budget, sizeof(AC.Field), "ArrowTypes derived dictionary Field") + return AC.Field(f.name, t.valuetype, nullable, metadata, f.children) +end + +function _arrowtypesdictvaluefield( + ctx::_ArrowTypesContext, + f::AC.Field, + t::AC.DictionaryType; + retainmetadata::Bool=false, +) + cache = retainmetadata ? ctx.metadata_dictionaryfields : ctx.dictionaryfields + what = + retainmetadata ? "ArrowTypes metadata dictionary-field cache" : + "ArrowTypes dictionary-field cache" + return _memoized!(cache, f, ctx.budget, what) do + _arrowtypesdictvaluefield(f, t; retainmetadata, budget=ctx.budget) + end +end + +function _arrowtypespublicdictvaluefield( + ctx::_ArrowTypesContext, + f::AC.Field, + t::AC.DictionaryType, +) + return _memoized!( + ctx.public_dictionaryfields, + f, + ctx.budget, + "ArrowTypes public dictionary-field cache", + ) do + _arrowtypesdictvaluefield(f, t; nullable=false, budget=ctx.budget) + end +end + +"Convert a child while preserving the container shape required by its parent." +function _arrowtypesnestedvalue( + ctx::_ArrowTypesContext, + f::AC.Field, + x; + extension_shape::Bool, +) + _, target, _ = _arrowtypestarget(ctx, f) + # Dictionary needs a private marker to distinguish a valid index + # into its null pool from a null dictionary index. Consume that marker at + # every recursive Field seam, including when the extension label is not + # registered in this process. No private routing value may reach a public + # row container. + (target !== nothing || x isa _ArrowTypesRoutedNull) && + return _arrowtypesvalue(ctx, f, x) + return _arrowtypesstoragevalue(ctx, f, x; extension_shape=extension_shape) +end + +function _arrowtypesnestedeltype( + ctx::_ArrowTypesContext, + f::AC.Field; + extension_shape::Bool, +) + _, target, _ = _arrowtypestarget(ctx, f) + if target !== nothing || extension_shape + return _arrowtypeslogicaleltype(ctx, f) + end + return _arrowtypespubliceltype(ctx, f) +end + +"Convert children, preserving the reader's unmarked row containers unless requested." +function _arrowtypesstoragevalue( + ctx::_ArrowTypesContext, + f::AC.Field, + x; + extension_shape::Bool, +) + t = f.type + x === missing && return missing + t isa AC.DictionaryType && return _arrowtypesstoragevalue( + ctx, + _arrowtypesdictvaluefield(ctx, f, t), + x; + extension_shape=extension_shape, + ) + if t isa AC.RunEndEncodedType + length(f.children) == 2 || return x + return _arrowtypesnestedvalue( + ctx, + f.children[2], + x; + extension_shape=extension_shape, + ) + end + if t isa Union{AC.ListType,AC.ListViewType,AC.FixedSizeListType} + length(f.children) == 1 || return x + child = f.children[1] + _chargevector!(ctx.budget, Any, length(x), "ArrowTypes converted list value") + vals = Any[ + _arrowtypesnestedvalue(ctx, child, y; extension_shape=extension_shape) for + y in x + ] + if t isa AC.FixedSizeListType && extension_shape + _chargevector!(ctx.budget, Any, length(vals), "ArrowTypes tuple storage") + return Tuple(vals) + end + return _typedvalues( + _arrowtypesnestedeltype(ctx, child; extension_shape=extension_shape), + vals, + ctx.budget, + ) + end + if t isa AC.StructType + length(x) == length(f.children) || throw( + AC.ValidationError("extension struct value width does not match its Field"), + ) + _chargevector!(ctx.budget, Any, length(f.children), "ArrowTypes struct workspace") + vals = Vector{Any}(undef, length(f.children)) + for (i, child) in enumerate(f.children) + kv = x[i] + kv isa Pair || + throw(AC.ValidationError("extension struct rows must contain Pairs")) + first(kv) == child.name || throw( + AC.ValidationError("extension struct child order does not match its Field"), + ) + vals[i] = _arrowtypesnestedvalue( + ctx, + child, + last(kv); + extension_shape=extension_shape, + ) + end + if extension_shape + NT = _arrowtypesstoragebasetype(ctx, f) + if NT <: NamedTuple + _chargevector!( + ctx.budget, + Any, + length(vals), + "ArrowTypes NamedTuple storage", + ) + return NT(Tuple(vals)) + end + _chargevector!( + ctx.budget, + Pair{String,Any}, + length(vals), + "ArrowTypes converted struct value", + ) + return Pair{String,Any}[f.children[i].name => vals[i] for i in eachindex(vals)] + end + _chargevector!( + ctx.budget, + Pair{String,Any}, + length(vals), + "ArrowTypes converted struct value", + ) + return Pair{String,Any}[f.children[i].name => vals[i] for i in eachindex(vals)] + end + if t isa AC.MapType + length(f.children) == 1 || return x + entries = f.children[1] + length(entries.children) == 2 || return x + kf, vf = entries.children + _chargevector!( + ctx.budget, + Pair{Any,Any}, + length(x), + "ArrowTypes converted map value", + ) + vals = Pair{Any,Any}[ + _arrowtypesnestedvalue(ctx, kf, first(kv); extension_shape=extension_shape) => _arrowtypesnestedvalue( + ctx, + vf, + last(kv); + extension_shape=extension_shape, + ) for kv in x + ] + if extension_shape + D = _arrowtypesstoragebasetype(ctx, f) + if D <: Dict + _chargedict!( + ctx.budget, + keytype(D), + valtype(D), + length(vals), + "ArrowTypes Dict storage", + ) + end + out = D() + D <: Dict && sizehint!(out, length(vals)) + for kv in vals + out[first(kv)] = last(kv) + end + return out + end + return vals + end + if t isa AC.UnionType + if x isa _ArrowTypesRoutedUnion + 1 <= x.child <= length(f.children) || + throw(AC.ValidationError("routed union child is outside its Field")) + return _arrowtypesnestedvalue( + ctx, + f.children[x.child], + x.value; + extension_shape=extension_shape, + ) + end + for child in f.children + T = _arrowtypesnestedeltype(ctx, child; extension_shape=extension_shape) + x isa T && return _arrowtypesnestedvalue( + ctx, + child, + x; + extension_shape=extension_shape, + ) + end + return x + end + # Only a marked parent's ArrowTypes storage shape asks for native scalars + # here. An unmarked composite keeps its temporal children in the raw + # storage domain, even when a marked sibling triggered this recursion. + return extension_shape ? _arrowtypesscalar(t, x) : x +end + +function _arrowtypesutf8storage(ctx::_ArrowTypesContext, f::AC.Field) + t = f.type + t isa AC.DictionaryType && + return _arrowtypesutf8storage(ctx, _arrowtypesdictvaluefield(ctx, f, t)) + if t isa AC.RunEndEncodedType + return length(f.children) == 2 && _arrowtypesutf8storage(ctx, f.children[2]) + end + return t isa AC.Utf8Type || (t isa AC.ViewType && t.utf8) +end + +function _arrowtypesfromarrow(ctx::_ArrowTypesContext, T, f::AC.Field, storage) + t = f.type + if _arrowtypesutf8storage(ctx, f) && storage isa AbstractString + if T === Symbol + symbol = _existingjlsymbol(storage) + if symbol === nothing + displayname = _extensionwarningname(storage, ctx.budget) + _chargeobject!( + ctx.budget, + AC.checked_add( + Int64(160), + AC.checked_mul(Int64(4), Int64(ncodeunits(displayname))), + ), + "ArrowTypes Symbol validation message", + ) + throw( + AC.ValidationError( + "JuliaLang.Symbol payload $(repr(displayname)) is not already " * + "interned; automatic lifting would permanently allocate " * + "process-global state", + ), + ) + end + return symbol + end + bytes = codeunits(storage) + GC.@preserve storage bytes begin + return ArrowTypes.fromarrow(T, pointer(bytes), length(bytes)) + end + elseif t isa AC.DictionaryType + # The extension label lives on the dictionary field, but the shape + # passed to ArrowTypes is the dictionary VALUE shape. In particular, a + # Struct value must use `fromarrowstruct` with its declared child names; + # the generic `fromarrow(T, storage)` path can bind reordered storage + # children to the logical constructor by position. + return _arrowtypesfromarrow(ctx, T, _arrowtypesdictvaluefield(ctx, f, t), storage) + elseif t isa AC.RunEndEncodedType + return ArrowTypes.fromarrow(T, storage) + elseif t isa AC.StructType + if T <: NamedTuple || T <: Tuple + return T(Tuple(storage)) + end + names = _arrowtypesstructnames(ctx, f) + values = Tuple(storage) + if isdefined(ArrowTypes, :fromarrowstruct) + fromstruct = getfield(ArrowTypes, :fromarrowstruct) + applicable(fromstruct, T, Val(names), values...) && + return fromstruct(T, Val(names), values...) + end + return ArrowTypes.fromarrow(T, values...) + end + return ArrowTypes.fromarrow(T, storage) +end + +function _arrowtypesvalue(ctx::_ArrowTypesContext, f::AC.Field, x) + routednull = x isa _ArrowTypesRoutedNull + routednull && (x = missing) + has_label, target, _ = _arrowtypestarget(ctx, f) + if target === nothing + return _arrowtypesstoragevalue(ctx, f, x; extension_shape=false) + end + if f.type isa AC.UnionType && x isa _ArrowTypesRoutedUnion + child = f.children[x.child] + storage = _arrowtypesstoragetype!(ctx, target) + # A nullable logical column adds one unmarked Null child around a + # target whose own storage Union has no Missing branch. Preserve that + # outer missing instead of passing it to the target's fromarrow hook. + child.type isa AC.NullType && + _arrowtypesextension(child) === nothing && + !(storage isa Union && Missing <: storage) && + return missing + end + # A validity null remains `missing`; NullType is different — it is the + # physical storage of public-domain values such as `nothing` and must lift. + x === missing && !routednull && !(f.type isa AC.NullType) && return missing + storage = _arrowtypesstoragevalue(ctx, f, x; extension_shape=true) + return _arrowtypesfromarrow(ctx, target, f, storage) +end + +"Interpret extension labels recursively over an already materialized column." +function _arrowtypescolumn( + f::AC.Field, + col::AbstractVector, + ctx::_ArrowTypesContext=_ArrowTypesContext(), +) + has_label, target, _ = _arrowtypestarget(ctx, f) + has_registered_child = any( + child -> + _arrowtypestarget(ctx, child)[2] !== nothing || + _hasarrowtypesextension(child, ctx), + f.children, + ) + has_routed_value = + any(x -> x isa _ArrowTypesRoutedNull || x isa _ArrowTypesRoutedUnion, col) + if target === nothing && has_label && !has_registered_child && !has_routed_value + # No hook or routed provenance is available to consume. Preserve the + # ordinary storage column without building and discarding a per-row + # ArrowTypes workspace. + return _publiccolumn(f, _postconvert(f.type, col, ctx.budget), ctx.budget) + end + _chargevector!(ctx.budget, Any, length(col), "ArrowTypes lifting workspace") + values = Any[_arrowtypesvalue(ctx, f, x) for x in col] + if target !== nothing + # NullType uses the physical null slots as its storage values. A + # registered logical type such as `Nothing` lifts those slots to real + # values, so physical field nullability must not add `Missing` back. + nullable = + !(f.type isa AC.NullType) && + (_arrowtypeslogicalnullable(ctx, f, target) || any(ismissing, values)) + T = _withmissingtype(target, nullable) + return _typedvalues(T, values, ctx.budget) + end + # An unknown top-level extension remains the ordinary storage column. + # Recursive registered children have already been lifted in `values`. + if has_label && !has_registered_child + # An unknown Dictionary extension still took the private route to + # retain index validity. Use the consumed values so its marker cannot + # escape; other unknown extensions keep their original storage column. + storage = has_routed_value ? values : col + return _publiccolumn(f, _postconvert(f.type, storage, ctx.budget), ctx.budget) + end + T = _arrowtypespubliceltype(ctx, f) + any(ismissing, values) && !(Missing <: T) && (T = Union{Missing,T}) + return _typedvalues(T, values, ctx.budget) +end diff --git a/src/cdata.jl b/src/cdata.jl new file mode 100644 index 00000000..d876d7e3 --- /dev/null +++ b/src/cdata.jl @@ -0,0 +1,2108 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 C data interface and C stream interface adapter over ArrowCore. +# +# `ArrayData` has the shape of the C `ArrowArray` (buffers + children + +# dictionary + length/null_count/offset), so export is struct filling and +# import is struct reading. +# +# * Export: ONE release callback per C structure (never per buffer). A +# parent callback releases each child/dictionary that has not been moved; +# a moved child keeps the shared export allocation alive until its own +# callback runs. `private_data` points to a per-node malloc'd, +# never-GC-scanned CONTROL BLOCK holding an exactly-once state and the +# registry key. The Julia-side owner (which roots the Core columns and +# every malloc'd C struct) stays in a global EXPORT REGISTRY until +# release — a raw pointer in private_data roots nothing by itself. The +# @cfunction release callback recursively marks the C tree released. +# Callback traversal uses producer-owned canonical child/dictionary +# topology, not the caller-visible counts and pointer tables. It still +# reads each canonical descendant's public release field so conforming +# moves are honored. A reaper pass (`reap!`) finds aggregates whose last +# outstanding node was released, frees their mallocs, and drops the +# registry root. Dropping that root is what lets the source columns — +# and, through their `OwnerRegion` roots, the buffer memory itself — +# become collectable again. Callback contract: releases for one tree are +# serialized and run only on Julia-attached threads. +# +# * Import: the moved ArrowArray becomes ONE ForeignOwner shared by every +# child/dictionary BufferSlice (a single release for the whole tree — +# per-buffer owners would double-release). Buffer extents are DECLARED, +# not verified: the ABI cannot prove allocation sizes, so extents are +# computed from length/offset/layout, and offsets buffers are read +# (bounded by their computed size) to size the data buffers they govern. +# Failed imports release the moved structure exactly once before +# throwing. Per spec, moving marks the source released (release = NULL). +# Validity is reachability: every imported region's `root` is the +# ForeignOwner, so the producer's memory outlives every slice by +# construction. Every region over one import shares one `ReleaseCell`, so +# `release!` on the owner or on any of its regions revokes all siblings +# (later access is a clean `InvalidStateException`) and then runs the +# producer's release exactly once. +# +# * Streams: `ArrowArrayStream` maps in both directions with one +# independently-owned export root per result and exception-safe +# move/release handoffs. +# ============================================================================= + +# --------------------------------------------------------------------------- +# ABI structs (field-exact per https://arrow.apache.org/docs/format/CDataInterface.html) +# --------------------------------------------------------------------------- + +""" + CArrowSchema + +The field-exact Julia representation of the Arrow C data interface's +`ArrowSchema` ABI struct. +""" +struct CArrowSchema + format::Ptr{UInt8} + name::Ptr{UInt8} + metadata::Ptr{UInt8} + flags::Int64 + n_children::Int64 + children::Ptr{Ptr{CArrowSchema}} + dictionary::Ptr{CArrowSchema} + release::Ptr{Cvoid} + private_data::Ptr{Cvoid} +end + +""" + CArrowArray + +The field-exact Julia representation of the Arrow C data interface's +`ArrowArray` ABI struct. +""" +struct CArrowArray + length::Int64 + null_count::Int64 + offset::Int64 + n_buffers::Int64 + n_children::Int64 + buffers::Ptr{Ptr{Cvoid}} + children::Ptr{Ptr{CArrowArray}} + dictionary::Ptr{CArrowArray} + release::Ptr{Cvoid} + private_data::Ptr{Cvoid} +end + +const ARROW_FLAG_NULLABLE = Int64(2) +const ARROW_FLAG_DICTIONARY_ORDERED = Int64(1) +const ARROW_FLAG_MAP_KEYS_SORTED = Int64(4) +const ARROW_FLAG_ALL_SUPPORTED = + ARROW_FLAG_NULLABLE | ARROW_FLAG_DICTIONARY_ORDERED | ARROW_FLAG_MAP_KEYS_SORTED + +# --------------------------------------------------------------------------- +# Format strings <-> Core descriptors (parity with Core's accessor set) +# --------------------------------------------------------------------------- + +_tuchar(u) = + u == AC.SECOND ? "s" : u == AC.MILLISECOND ? "m" : u == AC.MICROSECOND ? "u" : "n" + +formatstring(t::IntType) = ( + t.signed ? Dict(8 => "c", 16 => "s", 32 => "i", 64 => "l") : + Dict(8 => "C", 16 => "S", 32 => "I", 64 => "L") +)[t.bits] +formatstring(t::FloatType) = Dict(16 => "e", 32 => "f", 64 => "g")[t.bits] +formatstring(::BoolType) = "b" +formatstring(::NullType) = "n" +formatstring(t::Utf8Type) = t.large ? "U" : "u" +formatstring(t::BinaryType) = t.large ? "Z" : "z" +formatstring(t::FixedSizeBinaryType) = "w:$(t.nbytes)" +formatstring(t::DecimalType) = + t.bits == 128 ? "d:$(t.precision),$(t.scale)" : "d:$(t.precision),$(t.scale),$(t.bits)" +formatstring(t::DateType) = t.unit == AC.DAY ? "tdD" : "tdm" +formatstring(t::TimeType) = "tt" * _tuchar(t.unit) +formatstring(t::TimestampType) = "ts" * _tuchar(t.unit) * ":" * something(t.timezone, "") +formatstring(t::DurationType) = "tD" * _tuchar(t.unit) +formatstring(t::IntervalType) = + t.unit == AC.YEAR_MONTH ? "tiM" : t.unit == AC.DAY_TIME ? "tiD" : "tin" +formatstring(t::ListType) = t.large ? "+L" : "+l" +formatstring(t::FixedSizeListType) = "+w:$(t.listsize)" +formatstring(::StructType) = "+s" +formatstring(::MapType) = "+m" +formatstring(t::UnionType) = + (t.mode == AC.SparseMode ? "+us:" : "+ud:") * join(Int.(t.typeids), ",") +formatstring(t::ViewType) = t.utf8 ? "vu" : "vz" +formatstring(t::ListViewType) = t.large ? "+vL" : "+vl" +formatstring(::RunEndEncodedType) = "+r" + +# Closed-set ladder (same devirtualization story as AC.layoutspec_of): the +# export walk reaches this with an abstract-typed Field slot. +@inline function formatstring_of(t::ArrowType)::String + t isa IntType && return formatstring(t) + t isa FloatType && return formatstring(t) + t isa BoolType && return formatstring(t) + t isa NullType && return formatstring(t) + t isa Utf8Type && return formatstring(t) + t isa BinaryType && return formatstring(t) + t isa FixedSizeBinaryType && return formatstring(t) + t isa DecimalType && return formatstring(t) + t isa DateType && return formatstring(t) + t isa TimeType && return formatstring(t) + t isa TimestampType && return formatstring(t) + t isa DurationType && return formatstring(t) + t isa IntervalType && return formatstring(t) + t isa ListType && return formatstring(t) + t isa FixedSizeListType && return formatstring(t) + t isa StructType && return formatstring(t) + t isa MapType && return formatstring(t) + t isa UnionType && return formatstring(t) + t isa ViewType && return formatstring(t) + t isa ListViewType && return formatstring(t) + t isa RunEndEncodedType && return formatstring(t) + # per spec: the index format on the node; values on schema.dictionary + t isa DictionaryType && return formatstring_of(t.indextype) + throw(ArgumentError("unregistered ArrowType")) +end + +_formaterror(fmt) = throw(ValidationError("unsupported C format string \"$fmt\"")) + +function _parseformatint(fmt, s, what; low=0, high=typemax(Int32)) + bytes = codeunits(s) + isempty(bytes) && throw(ValidationError("invalid $what in C format string \"$fmt\"")) + firstdigit = 1 + if bytes[1] == UInt8('-') + low < 0 || throw(ValidationError("invalid $what in C format string \"$fmt\"")) + length(bytes) > 1 || + throw(ValidationError("invalid $what in C format string \"$fmt\"")) + firstdigit = 2 + end + for i = firstdigit:length(bytes) + UInt8('0') <= bytes[i] <= UInt8('9') || + throw(ValidationError("invalid $what in C format string \"$fmt\"")) + end + n = tryparse(Int64, s) + (n === nothing || !(low <= n <= high)) && + throw(ValidationError("invalid $what in C format string \"$fmt\"")) + return Int(n) +end + +_parsetimeunit(fmt, c) = + c == UInt8('s') ? AC.SECOND : + c == UInt8('m') ? AC.MILLISECOND : + c == UInt8('u') ? AC.MICROSECOND : c == UInt8('n') ? AC.NANOSECOND : _formaterror(fmt) + +function _parseunionids(fmt, body) + ids = Int8[] + isempty(body) && return ids + # A valid Int8-domain union has at most 128 children. Count separators + # without splitting so an overlong malformed string cannot direct a large + # temporary allocation before it is rejected. + nids = 1 + for b in codeunits(body) + b == UInt8(',') || continue + nids += 1 + nids <= 128 || + throw(ValidationError("union C format string declares more than 128 type ids")) + end + sizehint!(ids, nids) + seen = UInt128(0) + value = 0 + have_digit = false + for b in codeunits(body) + if UInt8('0') <= b <= UInt8('9') + have_digit = true + value = 10 * value + Int(b - UInt8('0')) + value <= 127 || throw(ValidationError("union type ids must be in [0, 127]")) + elseif b == UInt8(',') + have_digit || + throw(ValidationError("invalid union type id in C format string \"$fmt\"")) + bit = UInt128(1) << value + seen & bit == 0 || throw(ValidationError("union type ids must be unique")) + push!(ids, Int8(value)) + seen |= bit + value = 0 + have_digit = false + else + throw(ValidationError("invalid union type id in C format string \"$fmt\"")) + end + end + have_digit || + throw(ValidationError("invalid union type id in C format string \"$fmt\"")) + bit = UInt128(1) << value + seen & bit == 0 || throw(ValidationError("union type ids must be unique")) + push!(ids, Int8(value)) + return ids +end + +function parseformat(fmt::AbstractString, flags::Int64=0)::ArrowType + isvalid(fmt) || throw(ValidationError("C format string is not valid UTF-8")) + occursin('\0', fmt) && + throw(ValidationError("C format string cannot contain embedded NUL characters")) + fmt = String(fmt) + fmt == "b" && return BoolType() + fmt == "n" && return NullType() + fmt == "u" && return Utf8Type(false) + fmt == "U" && return Utf8Type(true) + fmt == "z" && return BinaryType(false) + fmt == "Z" && return BinaryType(true) + fmt == "+l" && return ListType(false) + fmt == "+L" && return ListType(true) + fmt == "vu" && return ViewType(true) + fmt == "vz" && return ViewType(false) + fmt == "+vl" && return ListViewType(false) + fmt == "+vL" && return ListViewType(true) + fmt == "+r" && return RunEndEncodedType() + fmt == "+s" && return StructType() + fmt == "+m" && return MapType((flags & ARROW_FLAG_MAP_KEYS_SORTED) != 0) + fmt == "e" && return FloatType(16) + fmt == "f" && return FloatType(32) + fmt == "g" && return FloatType(64) + fmt == "tdD" && return DateType(AC.DAY) + fmt == "tdm" && return DateType(AC.MILLISECOND_DATE) + fmt == "tiM" && return IntervalType(AC.YEAR_MONTH) + fmt == "tiD" && return IntervalType(AC.DAY_TIME) + fmt == "tin" && return IntervalType(AC.MONTH_DAY_NANO) + m = Dict( + "c" => (8, true), + "C" => (8, false), + "s" => (16, true), + "S" => (16, false), + "i" => (32, true), + "I" => (32, false), + "l" => (64, true), + "L" => (64, false), + ) + haskey(m, fmt) && return IntType(m[fmt]...) + if ncodeunits(fmt) == 3 && startswith(fmt, "tt") + u = _parsetimeunit(fmt, codeunit(fmt, 3)) + return TimeType(u, u == AC.SECOND || u == AC.MILLISECOND ? 32 : 64) + end + ncodeunits(fmt) == 3 && + startswith(fmt, "tD") && + return DurationType(_parsetimeunit(fmt, codeunit(fmt, 3))) + if startswith(fmt, "ts") && ncodeunits(fmt) >= 4 && codeunit(fmt, 4) == UInt8(':') + u = _parsetimeunit(fmt, codeunit(fmt, 3)) + tz = SubString(fmt, 5) + return TimestampType(u, isempty(tz) ? nothing : String(tz)) + end + if startswith(fmt, "w:") + return FixedSizeBinaryType(_parseformatint(fmt, fmt[3:end], "byte width")) + end + if startswith(fmt, "+w:") + return FixedSizeListType(_parseformatint(fmt, fmt[4:end], "list size")) + end + if startswith(fmt, "d:") + parts = split(fmt[3:end], ","; limit=4, keepempty=true) + 2 <= length(parts) <= 3 || + throw(ValidationError("invalid decimal C format string \"$fmt\"")) + precision = _parseformatint(fmt, parts[1], "decimal precision") + scale = _parseformatint(fmt, parts[2], "decimal scale"; low=typemin(Int32)) + bits = + length(parts) == 3 ? _parseformatint(fmt, parts[3], "decimal bit width") : 128 + t = DecimalType(precision, scale, bits) + AC._validate_descriptor(t) + return t + end + startswith(fmt, "+us:") && + return UnionType(AC.SparseMode, _parseunionids(fmt, fmt[5:end])) + startswith(fmt, "+ud:") && + return UnionType(AC.DenseMode, _parseunionids(fmt, fmt[5:end])) + _formaterror(fmt) +end + +# --------------------------------------------------------------------------- +# Export: Core -> C structs, per-node controls + registry + explicit reaper +# --------------------------------------------------------------------------- + +# Per-node control block layout (malloc'd, never GC-scanned): +# offset 0: UInt8 state (0 = live, 1 = releasing, 2 = released) +# offset 8: Int64 registry key +const CONTROL_BLOCK_BYTES = 16 + +""" +Everything one export tree must keep alive and eventually free: the Core +columns (whose OwnerRegions root the actual buffers), every malloc'd C struct +and string, and every per-node control block. Held in EXPORT_REGISTRY under +their shared aggregate key until all non-moved and moved nodes have been +released and the reaper runs. Rooting the columns here is the entire +source-liveness story: raw C pointers handed to a consumer stay valid because +the registry keeps this object and all source regions reachable. +""" +mutable struct ExportedRoot + roots::Vector{Any} # ArrayData/Field/Schema kept reachable + mallocs::Vector{Ptr{Cvoid}} # every Libc.malloc'd allocation, freed on reap + key::Int64 + remaining::Int64 # exported C nodes whose callback has not run + schema_topology::Dict{Ptr{Cvoid},Tuple{Vector{Ptr{CArrowSchema}},Ptr{CArrowSchema}}} + array_topology::Dict{Ptr{Cvoid},Tuple{Vector{Ptr{CArrowArray}},Ptr{CArrowArray}}} +end + +const EXPORT_REGISTRY = Dict{Int64,ExportedRoot}() +const REGISTRY_LOCK = ReentrantLock() +const NEXT_KEY = Ref{Int64}(0) +function _claim_array_node(a::Ptr{CArrowArray}, claimed_slot) + a == C_NULL && return nothing + return lock(REGISTRY_LOCK) do + arr = unsafe_load(a) + arr.release == C_NULL && return nothing + p = arr.private_data + p == C_NULL && return nothing + key = unsafe_load(Ptr{Int64}(p + 8)) + root = get(EXPORT_REGISTRY, key, nothing) + root === nothing && error("C Data export root disappeared during release") + topology = get(root.array_topology, p, nothing) + topology === nothing && error("C Data array topology disappeared during release") + flag = unsafe_load(Ptr{UInt8}(p)) + flag == 0x00 || return nothing + claimed = (p, topology) + claimed_slot[] = claimed + unsafe_store!(Ptr{UInt8}(p), 0x01) + return claimed + end +end + +function _claim_schema_node(s::Ptr{CArrowSchema}, claimed_slot) + s == C_NULL && return nothing + return lock(REGISTRY_LOCK) do + sch = unsafe_load(s) + sch.release == C_NULL && return nothing + p = sch.private_data + p == C_NULL && return nothing + key = unsafe_load(Ptr{Int64}(p + 8)) + root = get(EXPORT_REGISTRY, key, nothing) + root === nothing && error("C Data export root disappeared during release") + topology = get(root.schema_topology, p, nothing) + topology === nothing && error("C Data schema topology disappeared during release") + flag = unsafe_load(Ptr{UInt8}(p)) + flag == 0x00 || return nothing + claimed = (p, topology) + claimed_slot[] = claimed + unsafe_store!(Ptr{UInt8}(p), 0x01) + return claimed + end +end + +function _finish_node!(p, control::Ptr{Cvoid}, claimed_slot, committed_slot) + # This locked block is the callback's final access to export-owned memory. + # The reaper observes zero only after every non-moved descendant callback, + # and every independently moved node callback, has completed. Scanning in + # reap! keeps allocation and registry removal out of the C callback. + lock(REGISTRY_LOCK) do + unsafe_load(Ptr{UInt8}(control)) == 0x01 || + error("C Data node is not in releasing state") + key = unsafe_load(Ptr{Int64}(control + 8)) + root = get(EXPORT_REGISTRY, key, nothing) + root === nothing && error("C Data export root disappeared during release") + root.remaining > 0 || error("C Data export node counter underflow") + oldremaining = root.remaining + oldrelease = unsafe_load(p).release + try + root.remaining = oldremaining - 1 + unsafe_store!(Ptr{UInt8}(control), 0x02) + _store_field!(p, Val(:release), Ptr{Cvoid}(C_NULL)) + # The outer catch must not touch `control` once remaining is + # zero: a reaper may free it as soon as this lock is released. + # Transfer the completed claim while the lock still excludes + # cleanup. A later exception observes a committed callback. + committed_slot[] = true + claimed_slot[] = nothing + catch + if !committed_slot[] + # Nothing can reap this root while the registry lock is held. + # Restore the whole commit before the outer transaction + # returns the node from RELEASING to LIVE. + root.remaining = oldremaining + unsafe_store!(Ptr{UInt8}(control), 0x01) + _store_field!(p, Val(:release), oldrelease) + end + rethrow() + end + end + return nothing +end + +function _reset_node_claim!(control::Ptr{Cvoid}) + lock(REGISTRY_LOCK) do + flag = unsafe_load(Ptr{UInt8}(control)) + flag == 0x01 || return nothing + unsafe_store!(Ptr{UInt8}(control), 0x00) + end + return nothing +end + +function _release_array_children!(topology) + children, dictionary = topology + # Read the child's release pointer under the registry lock, then call it + # unlocked: a moved child's callback belongs to the consumer and must not + # run while we hold the lock. + for child in children + release = lock(REGISTRY_LOCK) do + unsafe_load(child).release + end + if release != C_NULL + ccall(release, Cvoid, (Ptr{CArrowArray},), child) + lock(REGISTRY_LOCK) do + unsafe_load(child).release == C_NULL || + error("C Data child array release did not complete") + end + end + end + if dictionary != C_NULL + release = lock(REGISTRY_LOCK) do + unsafe_load(dictionary).release + end + if release != C_NULL + ccall(release, Cvoid, (Ptr{CArrowArray},), dictionary) + lock(REGISTRY_LOCK) do + unsafe_load(dictionary).release == C_NULL || + error("C Data dictionary array release did not complete") + end + end + end + return nothing +end + +function _release_schema_children!(topology) + children, dictionary = topology + # Read the child's release pointer under the registry lock, then call it + # unlocked: a moved child's callback belongs to the consumer and must not + # run while we hold the lock. + for child in children + release = lock(REGISTRY_LOCK) do + unsafe_load(child).release + end + if release != C_NULL + ccall(release, Cvoid, (Ptr{CArrowSchema},), child) + lock(REGISTRY_LOCK) do + unsafe_load(child).release == C_NULL || + error("C Data child schema release did not complete") + end + end + end + if dictionary != C_NULL + release = lock(REGISTRY_LOCK) do + unsafe_load(dictionary).release + end + if release != C_NULL + ccall(release, Cvoid, (Ptr{CArrowSchema},), dictionary) + lock(REGISTRY_LOCK) do + unsafe_load(dictionary).release == C_NULL || + error("C Data dictionary schema release did not complete") + end + end + end + return nothing +end + +function _release_array(a::Ptr{CArrowArray}) + committed_slot = Ref(false) + claimed_slot = Ref{ + Union{Nothing,Tuple{Ptr{Cvoid},Tuple{Vector{Ptr{CArrowArray}},Ptr{CArrowArray}}}}, + }( + nothing, + ) + try + claimed = _claim_array_node(a, claimed_slot) + claimed === nothing && return nothing + control, topology = claimed + _release_array_children!(topology) + _finish_node!(a, control, claimed_slot, committed_slot) + catch + # A C release callback has no error channel, and a Julia exception + # must not unwind through the C ABI. Completed descendants are + # already NULL. Restore this node to LIVE so a later explicit call + # can resume without double release. + if !committed_slot[] + claimed = claimed_slot[] + claimed === nothing || _reset_node_claim!(claimed[1]) + end + end + return nothing +end + +function _release_schema(s::Ptr{CArrowSchema}) + committed_slot = Ref(false) + claimed_slot = Ref{ + Union{Nothing,Tuple{Ptr{Cvoid},Tuple{Vector{Ptr{CArrowSchema}},Ptr{CArrowSchema}}}}, + }( + nothing, + ) + try + claimed = _claim_schema_node(s, claimed_slot) + claimed === nothing && return nothing + control, topology = claimed + _release_schema_children!(topology) + _finish_node!(s, control, claimed_slot, committed_slot) + catch + if !committed_slot[] + claimed = claimed_slot[] + claimed === nothing || _reset_node_claim!(claimed[1]) + end + end + return nothing +end + +# Store one field of a C struct in place (structs are immutable in Julia; +# the C memory is not). +@generated function _store_field!(p::Ptr{T}, ::Val{name}, v) where {T,name} + i = findfirst(==(name), fieldnames(T)) + off = fieldoffset(T, i) + FT = fieldtype(T, i) + return :(unsafe_store!(Ptr{$FT}(Ptr{Cvoid}(p) + $off), convert($FT, v)); nothing) +end +_store_field!(p, name::Symbol, v) = _store_field!(p, Val(name), v) + +function _malloc!( + root::ExportedRoot, + n::Integer, + (register!)=push!, + (deallocate!)=Libc.free, +) + n >= 0 || throw(ArgumentError("negative export allocation size")) + n64 = Int64(n) + # Reserve the ledger slot before acquiring native memory. After malloc, + # either registration owns the pointer or the local catch deallocates it. + sizehint!(root.mallocs, AC.checked_add(length(root.mallocs), 1)) + oldlen = length(root.mallocs) + p = Ptr{Cvoid}(C_NULL) + owned = false + try + p = Libc.malloc(max(n64, Int64(1))) + p == C_NULL && throw(OutOfMemoryError()) + owned = true + register!(root.mallocs, p) + owned = false + catch + if owned + if length(root.mallocs) == oldlen + deallocate!(p) + owned = false + elseif length(root.mallocs) == oldlen + 1 && root.mallocs[end] == p + owned = false + else + error("export malloc registration left an invalid ledger state") + end + end + rethrow() + end + return Ptr{Cvoid}(p) +end + +""" +Encode field metadata per the C data interface: int32 pair count, then +per pair an int32 key length, key bytes, int32 value length, value bytes +(native endian, not NUL-terminated). NULL when there is no metadata. +""" +function _cmetadata!( + root::ExportedRoot, + metadata::Union{Nothing,AC.FrozenVector{Pair{String,String}}}, +)::Ptr{UInt8} + metadata === nothing && return Ptr{UInt8}(C_NULL) + n = length(metadata) + n == 0 && return Ptr{UInt8}(C_NULL) + buf = UInt8[] + append!(buf, reinterpret(UInt8, Int32[Int32(n)])) + for kv in metadata + k = first(kv) + v = last(kv) + append!(buf, reinterpret(UInt8, Int32[Int32(sizeof(k))])) + append!(buf, codeunits(k)) + append!(buf, reinterpret(UInt8, Int32[Int32(sizeof(v))])) + append!(buf, codeunits(v)) + end + p = Ptr{UInt8}(_malloc!(root, length(buf))) + GC.@preserve buf unsafe_copyto!(p, pointer(buf), length(buf)) + return p +end + +function _cstring!(root::ExportedRoot, s::AbstractString) + isvalid(s) || throw(ValidationError("C Data strings must be valid UTF-8")) + occursin('\0', s) && + throw(ValidationError("C Data strings cannot contain embedded NUL characters")) + n = ncodeunits(s) + p = Ptr{UInt8}(_malloc!(root, AC.checked_add(Int64(n), Int64(1)))) + for (i, b) in enumerate(codeunits(s)) + unsafe_store!(p, b, i) + end + unsafe_store!(p, 0x00, n + 1) + return p +end + +function _newcontrol!(root::ExportedRoot) + control = _malloc!(root, CONTROL_BLOCK_BYTES) + unsafe_store!(Ptr{UInt8}(control), 0x00) + unsafe_store!(Ptr{Int64}(control + 8), root.key) + root.remaining = AC.checked_add(root.remaining, Int64(1)) + return control +end + +function _export_schema!( + root::ExportedRoot, + f::Field, + release::Ptr{Cvoid}, +)::Ptr{CArrowSchema} + p = Ptr{CArrowSchema}(_malloc!(root, sizeof(CArrowSchema))) + childfields = f.type isa DictionaryType ? Field[] : f.children + nchildren = length(childfields) + canonical_children = Ptr{CArrowSchema}[] + childptrs = Ptr{Ptr{CArrowSchema}}(C_NULL) + if nchildren > 0 + childptrs = Ptr{Ptr{CArrowSchema}}( + _malloc!(root, AC.checked_mul(Int64(nchildren), Int64(sizeof(Ptr)))), + ) + for (i, cf) in enumerate(childfields) + child = _export_schema!(root, cf, release) + push!(canonical_children, child) + unsafe_store!(childptrs, child, i) + end + end + dict = Ptr{CArrowSchema}(C_NULL) + if f.type isa DictionaryType + dict = _export_schema!(root, AC.dictvaluefield(f, f.type), release) + end + flags = f.nullable ? ARROW_FLAG_NULLABLE : Int64(0) + f.type isa DictionaryType && f.type.ordered && (flags |= ARROW_FLAG_DICTIONARY_ORDERED) + f.type isa MapType && f.type.keyssorted && (flags |= ARROW_FLAG_MAP_KEYS_SORTED) + control = _newcontrol!(root) + unsafe_store!( + p, + CArrowSchema( + _cstring!(root, formatstring_of(f.type)), + _cstring!(root, f.name), + # Field metadata rides the OUTER node for every field, dictionary + # wrappers included — the C++ bridge exports field.metadata() on + # the wrapper and only TYPE metadata (extensions) on the dependent + # value node, and PyArrow imports only the wrapper's pairs. + _cmetadata!(root, f.metadata), + flags, + nchildren, + childptrs, + dict, + release, + control, + ), + ) + root.schema_topology[control] = (canonical_children, dict) + return p +end + +function _export_array!( + root::ExportedRoot, + d::ArrayData, + release::Ptr{Cvoid}, +)::Ptr{CArrowArray} + p = Ptr{CArrowArray}(_malloc!(root, sizeof(CArrowArray))) + spec = AC.layoutspec_of(d.type) + ncore = length(d.buffers) + # C Data appends one int64 buffer of variadic data-buffer LENGTHS to view + # arrays (extents are not otherwise recoverable from the ABI); it counts + # toward n_buffers here and nowhere else in the format. + nvariadic = spec.variadic ? ncore - length(spec.buffers) : 0 + nbuf = spec.variadic ? ncore + 1 : ncore + bufptrs = Ptr{Ptr{Cvoid}}( + _malloc!(root, AC.checked_mul(Int64(max(nbuf, 1)), Int64(sizeof(Ptr)))), + ) + for (i, b) in enumerate(d.buffers) + role = i <= length(spec.buffers) ? spec.buffers[i] : AC.DATA + bufferp = + if role == AC.OFFSETS && d.len == 0 && d.offset == 0 && AC.isempty_buffer(b) + # Core's canonical empty representation omits this otherwise + # unused allocation. C Data still exposes the Columnar + # length+1 offsets buffer, so root one terminal zero in the + # export aggregate without changing the Core array. + zerop = Ptr{UInt8}(_malloc!(root, spec.offsetwidth)) + for j = 1:spec.offsetwidth + unsafe_store!(zerop, UInt8(0), j) + end + Ptr{Cvoid}(zerop) + elseif AC.isempty_buffer(b) + # An absent validity bitmap, or any actual zero-byte buffer, is + # represented by a NULL pointer. + Ptr{Cvoid}(C_NULL) + else + Ptr{Cvoid}(AC.sliceptr(b)) + end + unsafe_store!(bufptrs, bufferp, i) + end + if spec.variadic + sizesp = + Ptr{Int64}(_malloc!(root, AC.checked_mul(Int64(max(nvariadic, 1)), Int64(8)))) + for k = 1:nvariadic + unsafe_store!(sizesp, d.buffers[length(spec.buffers) + k].len, k) + end + unsafe_store!(bufptrs, Ptr{Cvoid}(sizesp), nbuf) + end + nchildren = length(d.children) + canonical_children = Ptr{CArrowArray}[] + childptrs = Ptr{Ptr{CArrowArray}}(C_NULL) + if nchildren > 0 + childptrs = Ptr{Ptr{CArrowArray}}( + _malloc!(root, AC.checked_mul(Int64(nchildren), Int64(sizeof(Ptr)))), + ) + for (i, c) in enumerate(d.children) + child = _export_array!(root, c, release) + push!(canonical_children, child) + unsafe_store!(childptrs, child, i) + end + end + dict = + d.dictionary === nothing ? Ptr{CArrowArray}(C_NULL) : + _export_array!(root, d.dictionary, release) + control = _newcontrol!(root) + unsafe_store!( + p, + CArrowArray( + d.len, + nullcount(d), + d.offset, + nbuf, + nchildren, + bufptrs, + childptrs, + dict, + release, + control, + ), + ) + root.array_topology[control] = (canonical_children, dict) + return p +end + +function _build_c_data!(sp, skey, ap, akey, f::Field, d::ArrayData, arel, srel) + _newroot(Any[f]; result_slot=sp, key_slot=skey) do root + _export_schema!(root, f, srel) + end + _newroot(Any[d]; result_slot=ap, key_slot=akey) do root + _export_array!(root, d, arel) + end + return nothing +end + +""" + to_c_data(field, data) -> (Ptr{CArrowSchema}, Ptr{CArrowArray}) + +Export one column. The schema and array have separate sets of per-node +control blocks and separate Julia-side roots, as required by their +independent C Data lifetimes. Releasing either root recursively marks only +that structure tree released. Moved descendants defer aggregate cleanup. +The array root keeps the source ArrayData reachable until it is reaped; +that reachability is what keeps the exported buffer pointers valid. + +The column is validated through the semantic tier before publication — +the same tier the IPC writer applies. Content policy (`validate_full`: +UTF-8 well-formedness, the advisory nullability contract, canonical bits) +is the caller's opt-in, exactly as for IPC. +""" +function to_c_data(f::Field, d::ArrayData) + # Reject mismatched schema/data and malformed buffers before publishing + # either independently-owned C root (semantic composes structural). + validate_semantic(f, d) + arel = @cfunction(_release_array, Cvoid, (Ptr{CArrowArray},)) + srel = @cfunction(_release_schema, Cvoid, (Ptr{CArrowSchema},)) + sp = Ref{Ptr{CArrowSchema}}(C_NULL) + skey = Ref{Int64}(0) + ap = Ref{Ptr{CArrowArray}}(C_NULL) + akey = Ref{Int64}(0) + try + # The exact public method owns both output slots until its tuple return. + # A helper cannot lose a published pointer at its own return boundary: + # _newroot records each result in the caller's slot when it publishes. + _build_c_data!(sp, skey, ap, akey, f, d, arel, srel) + return sp[], ap[] + catch + # Schema and array are separate C lifetimes, but export is one API + # transaction. Neither has escaped on this path, so discard both. + _cleanup_export_slots!(sp, skey, ap, akey) + rethrow() + end +end + +function _free_export!(root::ExportedRoot) + # After a root is claimed, cleanup only drops Julia references and frees + # tracked mallocs. There is no fallible ownership transition to retry. + empty!(root.schema_topology) + empty!(root.array_topology) + while !isempty(root.mallocs) + Libc.free(pop!(root.mallocs)) + end + empty!(root.roots) + return nothing +end + +function _cleanup_registered_root!(key::Int64; require_released=true) + # Claim by removal: popping the root under the registry lock makes this + # cleanup naturally exclusive against concurrent reapers, and the frees + # below cannot throw, so a claimed root never needs re-publishing. + # `require_released=false` is only legal on paths where no C node has + # escaped to a consumer (build failures); a released consumer callback + # finds its root through this registry, so popping early would strand it. + root = lock(REGISTRY_LOCK) do + candidate = get(EXPORT_REGISTRY, key, nothing) + candidate === nothing && return nothing + require_released && candidate.remaining != 0 && return nothing + pop!(EXPORT_REGISTRY, key) + return candidate + end + root === nothing && return false + _free_export!(root) + return true +end + +""" + reap!() -> Int + +Find fully released exports: free every malloc they own and drop their +registry roots. Cleanup is this explicit scan — there is no background +reaper task; call it after consumers have released their structures. +""" +function reap!() + keys = lock(REGISTRY_LOCK) do + Int64[k for (k, root) in EXPORT_REGISTRY if root.remaining == 0] + end + reaped = 0 + for key in keys + reaped += _cleanup_registered_root!(key) + end + return reaped +end + +function _cleanup_private_root!(root::ExportedRoot, key::Int64) + registered = lock(REGISTRY_LOCK) do + get(EXPORT_REGISTRY, key, nothing) === root + end + if registered + _cleanup_registered_root!(key; require_released=false) + else + _free_export!(root) + end + return nothing +end + +function _cleanup_export_slots!(sp, skey, ap, akey) + # Clear raw pointer slots before any free. Stable registry keys remain + # valid cleanup tokens until their corresponding root is gone. + sp[] = C_NULL + ap[] = C_NULL + if akey[] != 0 + _cleanup_registered_root!(akey[]; require_released=false) + akey[] = 0 + end + if skey[] != 0 + _cleanup_registered_root!(skey[]; require_released=false) + skey[] = 0 + end + return nothing +end + +function _newroot(build, roots::Vector{Any}; result_slot=nothing, key_slot=nothing) + # key and root are single-assignment BEFORE the try: reassignment of a + # closure-captured local boxes it, which the trim verifier rejects. + # Nothing before the try owns native memory, so there is nothing to + # clean on those paths. + key = lock(REGISTRY_LOCK) do + NEXT_KEY[] = AC.checked_add(NEXT_KEY[], Int64(1)) + end + root = ExportedRoot( + roots, + Ptr{Cvoid}[], + key, + 0, + Dict{Ptr{Cvoid},Tuple{Vector{Ptr{CArrowSchema}},Ptr{CArrowSchema}}}(), + Dict{Ptr{Cvoid},Tuple{Vector{Ptr{CArrowArray}},Ptr{CArrowArray}}}(), + )::ExportedRoot + try + # The pointer cannot escape before `build` returns. Keep the root + # private until then: publishing it with `remaining == 0` would let a + # concurrent reaper free partial mallocs underneath the builder, + # before its first node control increments `remaining`. + result = build(root) + lock(REGISTRY_LOCK) do + EXPORT_REGISTRY[key] = root + key_slot === nothing || (key_slot[] = key) + result_slot === nothing || (result_slot[] = result) + end + return result + catch + # Export-failure cleanup: unregister (if published) and free. + result_slot === nothing || (result_slot[] = C_NULL) + key_slot === nothing || (key_slot[] = 0) + _cleanup_private_root!(root, key) + rethrow() + end +end + +# --------------------------------------------------------------------------- +# Import: C structs -> Core, one ForeignOwner per moved tree +# --------------------------------------------------------------------------- + +# ReleaseCell action: the producer release, run once the cell has revoked +# every region over the import. +function _release_owner_action(p::Ptr{Cvoid})::Cvoid + slot = unsafe_pointer_to_objref(p)::Base.RefValue{Any} + x = slot[] + x === nothing || _release_foreign_owner!(x::ForeignOwner, Libc.free) + return nothing +end + +""" +One owner for one MOVED ArrowArray tree. All `BufferSlice`s from the whole +tree — children and dictionary — use regions whose `root` is this object. +The tree therefore outlives every slice, and the C release callback runs +exactly once, from `release!` or from the GC finalizer, whichever comes +first. + +The malloc'd copy of the moved struct mirrors the C Data convention for its +own state: its release field is NULL (inert) until the move commits, and the +producer's real callback is stored only then (`_arm_foreign_owner!`). A +failure between construction and the move commit therefore frees just our +copy and never calls the producer — the source, whose release field is still +set, remains the owner. +""" +mutable struct ForeignOwner + # malloc'd copy: a stable native address for the producer's release callback + const arrayblock::Ptr{CArrowArray} + const producer_release::Ptr{Cvoid} # the moved struct's real callback + @atomic released::Bool # one swap picks the single releaser + # ONE revocation cell for every OwnerRegion built over this import: the + # producer's release frees the whole tree at once, so closing any + # imported buffer must revoke all of its siblings first (they share this + # lifetime). `release!` revokes through the cell, whose action runs the + # producer release — exactly-once against the GC-finalizer path. + const cell::AC.ReleaseCell + function ForeignOwner(arr::CArrowArray, registerfinalizer) + block = Libc.malloc(sizeof(CArrowArray)) + block == C_NULL && throw(OutOfMemoryError()) + p = Ptr{CArrowArray}(block) + # The cell must exist before the owner does, so its action receives a + # `Ref{Any}` slot that is filled in after `new`. `ReleaseCell.arg` + # roots the slot, the slot roots the owner, and + # `_release_owner_action` recovers the owner from that pointer. + slot = Ref{Any}(nothing) + cell = AC.ReleaseCell(@cfunction(_release_owner_action, Cvoid, (Ptr{Cvoid},)), slot) + o = try + unsafe_store!(p, arr) + _store_field!(p, Val(:release), Ptr{Cvoid}(C_NULL)) # inert until armed + new(p, arr.release, false, cell) + catch + # The native copy exists before the Julia owner does. If copy + # initialization or owner allocation fails, no finalizer can + # reclaim that copy for us. + Libc.free(block) + rethrow() + end + slot[] = o + try + registerfinalizer(release!, o) + catch + # The source still owns the producer resources. The copy is inert, + # so constructor cleanup frees only our malloc'd storage. If the + # registrar installed a finalizer before throwing, its later call + # observes released=true and is inert. + release!(o) + rethrow() + end + return o + end +end +ForeignOwner(arr::CArrowArray) = ForeignOwner(arr, finalizer) + +# The move commit: the source ArrowArray's release has been nulled, so this +# copy is now the sole owner of the producer's resources. Storing the real +# callback into the copy arms the release path; nothing between construction +# and this store can throw. +function _arm_foreign_owner!(o::ForeignOwner) + (@atomic o.released) && error("cannot arm a released foreign owner") + GC.@preserve o _store_field!(o.arrayblock, Val(:release), o.producer_release) + return nothing +end + +function _foreign_owner_armed(o::ForeignOwner) + GC.@preserve o begin + return unsafe_load(o.arrayblock).release != C_NULL + end +end + +function _release_moved_owner!(o::ForeignOwner) + # A failure may occur after the source move but before arming. Arm first + # so this release still runs the producer callback in that seam. + _foreign_owner_armed(o) || _arm_foreign_owner!(o) + release!(o) + return nothing +end + +""" + release!(owner::ForeignOwner) + +Deterministically release an imported C-data tree, in this order: revoke +every `OwnerRegion` built over the import through the shared cell (later +access throws `InvalidStateException`); run the producer's release callback +(if armed) on the malloc'd struct copy; check that the producer nulled the +copy's release field, as C Data requires; free the copy. Exactly once — one +atomic swap picks the single releaser between an explicit call and the GC +finalizer, and later calls return immediately. This is also the only +lifetime handle for an empty import that carries no region at all; +`ArrayData.owner` then holds it. A conformance failure throws, and reports +as a finalizer error on the finalizer path. +""" +function release!(o::ForeignOwner) + return release!(o.cell) +end + +function _release_foreign_owner!(o::ForeignOwner, deallocate!) + @atomicswap(o.released = true) && return nothing + GC.@preserve o begin + cb = unsafe_load(o.arrayblock).release + if cb != C_NULL + ccall(cb, Cvoid, (Ptr{CArrowArray},), o.arrayblock) + unsafe_load(o.arrayblock).release == C_NULL || ( + deallocate!(o.arrayblock); + error("C Data producer release did not mark the structure released") + ) + end + deallocate!(o.arrayblock) + end + return nothing +end + +"Load child `i`'s CArrowArray by value out of a parent's children table." +childat(a::CArrowArray, i::Int) = unsafe_load(unsafe_load(a.children, i)) +"Load buffer pointer `i` out of a CArrowArray's buffer table." +bufferptr(a::CArrowArray, i::Int) = unsafe_load(a.buffers, i) + +""" + from_c_data(schemaptr, arrayptr) -> (Field, ArrayData) + +Import a C-data column. The ArrowArray is moved: it is copied by value and its +source release is nulled so the producer side cannot double-free. The +ArrowSchema is parsed and then released in place. Buffer extents are computed +from length/offset/layout — DECLARED extents: the ABI cannot prove the +allocation sizes, so this is the trusted-in-process boundary, and validation +runs on the declared geometry. The imported column passes the semantic tier +(the same default as the IPC reader); `validate_full` on the returned pair is +the caller's opt-in for content policy. A failed import releases the moved +tree exactly once. +""" +from_c_data(sp::Ptr{CArrowSchema}, ap::Ptr{CArrowArray}) = _from_c_data(sp, ap) + +function _from_c_data( + sp::Ptr{CArrowSchema}, + ap::Ptr{CArrowArray}; + ownerfactory=ForeignOwner, +) + sp == C_NULL && throw(ArgumentError("ArrowSchema pointer is NULL")) + ap == C_NULL && throw(ArgumentError("ArrowArray pointer is NULL")) + sch = unsafe_load(sp) + arr = unsafe_load(ap) + (sch.release == C_NULL || arr.release == C_NULL) && + throw(ArgumentError("cannot import a released structure")) + owner = nothing + try + try + owner = ownerfactory(arr)::ForeignOwner + # MOVE: relinquish source ownership before arming the copied + # owner. The source release field is authoritative. + _store_field!(ap, Val(:release), Ptr{Cvoid}(C_NULL)) + _arm_foreign_owner!(owner) + _preflight_schema(sch) + f = _import_field(sch) + _preflight_array(f, arr) + d = _import_array(f, arr, owner) + validate_semantic(f, d) + return f, d + finally + # The schema lifetime is separate and must end on every path, + # including owner-construction failure. + _release_c_schema!(sp, sch) + end + catch + # Before the move, the caller's source remains the owner. After the + # move, this local copy must release exactly once even when finalizer + # registration or later validation failed. + owner !== nothing && + unsafe_load(ap).release == C_NULL && + _release_moved_owner!(owner) + rethrow() + end +end + +function _preflight_schema(sch::CArrowSchema, depth::Int=0) + depth <= 64 || throw(ValidationError("C schema nesting exceeds 64 levels")) + sch.release != C_NULL || throw(ValidationError("released C schema node")) + sch.format != C_NULL || throw(ValidationError("C schema format is NULL")) + sch.n_children >= 0 || throw(ValidationError("negative C schema child count")) + sch.n_children <= 1_000_000 || + throw(ValidationError("C schema child count exceeds import limit")) + sch.n_children == 0 || + sch.children != C_NULL || + throw(ValidationError("C schema child table is NULL")) + for i = 1:sch.n_children + childptr = unsafe_load(sch.children, i) + childptr != C_NULL || throw(ValidationError("C schema child $i is NULL")) + _preflight_schema(unsafe_load(childptr), depth + 1) + end + if sch.dictionary != C_NULL + _preflight_schema(unsafe_load(sch.dictionary), depth + 1) + end + return nothing +end + +function _preflight_array(f::Field, arr::CArrowArray, depth::Int=0) + depth <= 64 || throw(ValidationError("C array nesting exceeds 64 levels")) + arr.release != C_NULL || throw(ValidationError("released C array node")) + arr.length >= 0 || throw(ValidationError("negative C array length")) + arr.offset >= 0 || throw(ValidationError("negative C array offset")) + AC.checked_add(arr.offset, arr.length) + -1 <= arr.null_count <= arr.length || + throw(ValidationError("invalid C array null count $(arr.null_count)")) + arr.n_buffers >= 0 || throw(ValidationError("negative C array buffer count")) + arr.n_children >= 0 || throw(ValidationError("negative C array child count")) + arr.n_buffers == 0 || + arr.buffers != C_NULL || + throw(ValidationError("C array buffer table is NULL")) + arr.n_children == 0 || + arr.children != C_NULL || + throw(ValidationError("C array child table is NULL")) + + spec = AC.layoutspec_of(f.type) + expected_buffers = length(spec.buffers) + if spec.variadic + # validity + views + N variadic data buffers + the trailing int64 + # sizes buffer: at least the fixed pair plus the sizes buffer. + Int64(arr.n_buffers) >= expected_buffers + 1 || throw( + ValidationError( + "view layout $(typeof(f.type)) requires at least $(expected_buffers + 1) buffers, producer sent $(arr.n_buffers)", + ), + ) + else + Int64(arr.n_buffers) == expected_buffers || throw( + ValidationError( + "layout $(typeof(f.type)) declares $expected_buffers buffers, producer sent $(arr.n_buffers)", + ), + ) + end + expected_children = spec.childcount == -1 ? length(f.children) : spec.childcount + Int64(arr.n_children) == expected_children || throw( + ValidationError( + "layout $(typeof(f.type)) declares $expected_children children, producer sent $(arr.n_children)", + ), + ) + + for i = 1:arr.n_children + childptr = unsafe_load(arr.children, i) + childptr != C_NULL || throw(ValidationError("C array child $i is NULL")) + child = unsafe_load(childptr) + cf = f.children[i] + _preflight_array(cf, child, depth + 1) + end + if f.type isa DictionaryType + arr.dictionary != C_NULL || + throw(ValidationError("dictionary C array has no dictionary values")) + _preflight_array( + AC.dictvaluefield(f, f.type), + unsafe_load(arr.dictionary), + depth + 1, + ) + elseif arr.dictionary != C_NULL + throw(ValidationError("non-dictionary C array has dictionary values")) + end + return nothing +end + +function _release_c_schema!(sp::Ptr{CArrowSchema}, sch::CArrowSchema) + sch.release == C_NULL && return nothing + ccall(sch.release, Cvoid, (Ptr{CArrowSchema},), sp) + unsafe_load(sp).release == C_NULL || + error("C Data producer release did not mark the structure released") + return nothing +end + +function _release_c_array!(ap::Ptr{CArrowArray}, arr::CArrowArray) + arr.release == C_NULL && return nothing + ccall(arr.release, Cvoid, (Ptr{CArrowArray},), ap) + unsafe_load(ap).release == C_NULL || + error("C Data producer release did not mark the structure released") + return nothing +end + +# Longest C string a schema may carry. Format strings are tens of bytes; +# names and metadata keys are human-scale. The cap converts a missing NUL +# terminator from an unbounded memory scan into a clean refusal. +const CSTRING_SCAN_LIMIT = Int64(1) << 20 + +function _import_cstring(p::Ptr{UInt8}, what::AbstractString) + # The limit is enforced BEFORE every dereference: the scan window is + # exactly CSTRING_SCAN_LIMIT bytes, so the NUL must fall inside it + # (maximum payload is the limit minus one) and byte limit+1 is never + # touched — a guard page there must produce this refusal, not SIGBUS. + n = Int64(0) + while true + n >= CSTRING_SCAN_LIMIT && throw( + ValidationError( + "C Data $what has no NUL terminator within " * + "$(CSTRING_SCAN_LIMIT) bytes", + ), + ) + unsafe_load(p + n) == 0x00 && break + n += 1 + end + s = unsafe_string(p, n) + isvalid(s) || throw(ValidationError("C Data $what is not valid UTF-8")) + return s +end + +function _validate_schema_flags(sch::CArrowSchema, fmt::AbstractString) + sch.flags & ~ARROW_FLAG_ALL_SUPPORTED == 0 || + throw(ValidationError("C schema contains unsupported flag bits")) + (sch.flags & ARROW_FLAG_DICTIONARY_ORDERED == 0 || sch.dictionary != C_NULL) || + throw(ValidationError("ARROW_FLAG_DICTIONARY_ORDERED requires a dictionary schema")) + (sch.flags & ARROW_FLAG_MAP_KEYS_SORTED == 0 || fmt == "+m") || + throw(ValidationError("ARROW_FLAG_MAP_KEYS_SORTED requires a map schema")) + return nothing +end + +""" +Parse a C metadata blob. The pair count and the key/value lengths are +producer-declared, the same trust as every other C Data pointer, but a +negative length is rejected: it would wrap the walk. Metadata bytes are +taken by declared length and are not UTF-8-checked here: every import runs +`validate_semantic`, whose structural stage checks them before a `Field` +escapes, and the stream path re-checks per batch in `_validate_stream_field`. +""" +function _import_cmetadata(p::Ptr{UInt8}) + p == C_NULL && return nothing + n = unsafe_load(Ptr{Int32}(p)) + n < 0 && throw(ValidationError("C schema metadata declares a negative pair count")) + n == 0 && return nothing + off = Int64(4) + out = Pair{String,String}[] + for _ = 1:n + klen = unsafe_load(Ptr{Int32}(p + off)) + klen < 0 && + throw(ValidationError("C schema metadata declares a negative key length")) + k = unsafe_string(p + off + 4, klen) + off += 4 + Int64(klen) + vlen = unsafe_load(Ptr{Int32}(p + off)) + vlen < 0 && + throw(ValidationError("C schema metadata declares a negative value length")) + v = unsafe_string(p + off + 4, vlen) + off += 4 + Int64(vlen) + push!(out, k => v) + end + return out +end + +function _import_field(sch::CArrowSchema)::Field + fmt = _import_cstring(sch.format, "format") + _validate_schema_flags(sch, fmt) + name = sch.name == C_NULL ? "" : _import_cstring(sch.name, "field name") + nullable = (sch.flags & ARROW_FLAG_NULLABLE) != 0 + meta = _import_cmetadata(sch.metadata) + t = parseformat(fmt, sch.flags) + + # Check the schema shape before indexing any recursively-created child. + # Struct is the only mapped layout with field-declared arity. + spec = AC.layoutspec_of(t) + expected_children = spec.childcount + if expected_children >= 0 && sch.n_children != expected_children + throw( + ValidationError( + "C schema for $(typeof(t)) declares $(sch.n_children) children; expected $expected_children", + ), + ) + end + t isa UnionType && + length(t.typeids) != sch.n_children && + throw( + ValidationError( + "union format declares $(length(t.typeids)) type ids for $(sch.n_children) children", + ), + ) + + children = Field[] + for i = 1:sch.n_children + push!(children, _import_field(unsafe_load(unsafe_load(sch.children, i)))) + end + if sch.dictionary != C_NULL + vf = _import_field(unsafe_load(sch.dictionary)) + t isa IntType || + throw(ValidationError("dictionary index format must be an integer")) + isempty(children) || + throw(ValidationError("dictionary index schema must not have children")) + ordered = (sch.flags & ARROW_FLAG_DICTIONARY_ORDERED) != 0 + # The value node's metadata joins the wrapper's (wrapper pairs + # first; duplicate keys are legal): Core's one slot cannot express + # the two-node attribution, but no pair is lost. + vmeta = vf.metadata + dmeta = + meta === nothing ? + (vmeta === nothing ? nothing : collect(Pair{String,String}, vmeta)) : + (vmeta === nothing ? meta : vcat(meta, collect(Pair{String,String}, vmeta))) + # Branch on the metadata's presence: a Union-typed keyword makes + # the kwcall tuple imprecise, which trim cannot resolve. + dmeta === nothing && return Field( + name, + DictionaryType(t, vf.type, ordered); + nullable=nullable, + children=vf.children, + ) + return Field( + name, + DictionaryType(t, vf.type, ordered); + nullable=nullable, + metadata=dmeta, + children=vf.children, + ) + end + meta === nothing && return Field(name, t; nullable=nullable, children=children) + return Field(name, t; nullable=nullable, metadata=meta, children=children) +end + +""" +Compute each buffer's DECLARED byte extent from the layout registry and wrap +it as a slice over a foreign region rooted by `owner`. Offsets buffers are +sized first (len+1 entries) and then READ — inside their own declared bounds +— to size the data buffer they govern; that dependency order is exactly the +registry's buffer order, so the loop stays generic. +""" +function _import_array(f::Field, arr::CArrowArray, owner::ForeignOwner)::ArrayData + t = f.type + spec = AC.layoutspec_of(t) + total = AC.checked_add(arr.offset, arr.length) + buffers = BufferSlice[] + offsets_slice = nothing + for (i, role) in enumerate(spec.buffers) + p = bufferptr(arr, i) + nbytes = if role == AC.VALIDITY + p == C_NULL && + total > 0 && + arr.null_count != 0 && + throw(ValidationError("NULL validity buffer requires null_count == 0")) + p == C_NULL ? Int64(0) : AC.expected_validity_bytes(total) + elseif role == AC.OFFSETS + AC.checked_mul(AC.checked_add(total, Int64(1)), Int64(spec.offsetwidth)) + elseif role == AC.DATA + if spec.fixedwidth > 0 + AC.checked_mul(total, Int64(spec.fixedwidth)) + elseif spec.fixedwidth == -1 + AC.expected_validity_bytes(total) + else + # varbinary data: sized by the final offset, read from the + # offsets slice we just built (bounded by ITS declared size). + if offsets_slice === nothing || AC.isempty_buffer(offsets_slice) + Int64(0) + else + finaloffset = if spec.offsetwidth == 8 + AC.loadat(offsets_slice, Int64, AC.checked_mul(total, Int64(8))) + else + Int64(AC.loadat(offsets_slice, Int32, AC.checked_mul(total, Int64(4)))) + end + finaloffset >= 0 || + throw(ValidationError("negative final offset $finaloffset")) + finaloffset + end + end + elseif role == AC.TYPE_IDS + # One Int8 discriminator per union slot. + total + elseif role == AC.ELEMENT_OFFSETS || role == AC.SIZES + # Per-slot values (dense-union offsets; list-view offsets and + # sizes), not monotone ranges: exactly `total` entries, no +1. + AC.checked_mul(total, Int64(spec.offsetwidth)) + elseif role == AC.VIEWS + AC.checked_mul(total, Int64(16)) + else + throw(ValidationError("unsupported buffer role $role in C data import")) + end + if p == C_NULL + nbytes == 0 || + throw(ValidationError("NULL $role buffer with nonzero required size")) + push!(buffers, BufferSlice()) + else + region = OwnerRegion(Ptr{UInt8}(p), nbytes; root=owner, cell=owner.cell) + slice = BufferSlice(region, 0, nbytes) + role == AC.OFFSETS && (offsets_slice = slice) + push!(buffers, slice) + end + end + if spec.variadic + # The trailing int64 sizes buffer declares each variadic data + # buffer's extent — the one place the ABI carries a length for them. + nfixed = length(spec.buffers) + nvariadic = Int(arr.n_buffers) - nfixed - 1 + sizesp = Ptr{Int64}(bufferptr(arr, Int(arr.n_buffers))) + (nvariadic == 0 || sizesp != C_NULL) || throw( + ValidationError("view array with variadic buffers has a NULL sizes buffer"), + ) + for k = 1:nvariadic + len = unsafe_load(sizesp, k) + len >= 0 || throw(ValidationError("negative variadic buffer length $len")) + p = bufferptr(arr, nfixed + k) + if p == C_NULL + len == 0 || throw( + ValidationError("NULL variadic buffer with nonzero declared length"), + ) + push!(buffers, BufferSlice()) + else + region = OwnerRegion(Ptr{UInt8}(p), len; root=owner, cell=owner.cell) + push!(buffers, BufferSlice(region, 0, len)) + end + end + end + children = ArrayData[] + for i = 1:arr.n_children + # Preflight already rejects children on a dictionary node, so this + # branch is unreachable; it keeps the expression total for inference. + cf = t isa DictionaryType ? error("dictionary carries no children") : f.children[i] + push!(children, _import_array(cf, childat(arr, i), owner)) + end + dict = nothing + if arr.dictionary != C_NULL + t isa DictionaryType || + throw(ValidationError("dictionary array on a non-dictionary field")) + dict = _import_array(AC.dictvaluefield(f, t), unsafe_load(arr.dictionary), owner) + end + return AC._arraydata( + t, + arr.length, + buffers, + arr.offset, + children, + dict, + owner, + arr.null_count, + ) +end + +# --------------------------------------------------------------------------- +# C stream interface (ArrowArrayStream): batches over the same two mappings +# --------------------------------------------------------------------------- + +# Execution contract: stream callbacks call into Julia, so +# `get_schema`/`get_next`/`get_last_error`/`release` are legal ONLY from +# Julia-attached threads, and calls on one stream must not overlap (the C +# stream spec itself declares the structure not thread-safe). There is no +# marshaling to a Julia-owned worker for foreign-thread callers. + +""" + CArrowArrayStream + +The field-exact Julia representation of the Arrow C stream interface's +`ArrowArrayStream` ABI struct. +""" +struct CArrowArrayStream + get_schema::Ptr{Cvoid} # int (*)(ArrowArrayStream*, ArrowSchema* out) + get_next::Ptr{Cvoid} # int (*)(ArrowArrayStream*, ArrowArray* out) + get_last_error::Ptr{Cvoid} # const char* (*)(ArrowArrayStream*) + release::Ptr{Cvoid} # void (*)(ArrowArrayStream*) + private_data::Ptr{Cvoid} +end + +const EINVAL = Cint(Base.Libc.EINVAL) + +mutable struct ExportedStreamState + batchfield::Field # struct-typed: children are the schema + batches::Vector{AC.RecordBatch} + nextindex::Int + lasterror::Ptr{UInt8} # malloc'd NUL string; freed on replace/release +end + +const STREAM_REGISTRY = Dict{Int64,ExportedStreamState}() + +function _stream_state(sp::Ptr{CArrowArrayStream}) + stream = unsafe_load(sp) + stream.release == C_NULL && return nothing, Ptr{Cvoid}(C_NULL) + control = stream.private_data + control == C_NULL && return nothing, Ptr{Cvoid}(C_NULL) + key = unsafe_load(Ptr{Int64}(control + 8)) + state = lock(REGISTRY_LOCK) do + get(STREAM_REGISTRY, key, nothing) + end + return state, control +end + +function _set_stream_error!( + state::ExportedStreamState, + msg::AbstractString, + (allocate!)=Libc.malloc, + (deallocate!)=Libc.free, +) + # The prior pointer expires at the next stream operation even if building + # its replacement fails. Clear it first so malloc failure cannot report a + # stale error from an earlier operation. + old = state.lasterror + state.lasterror = Ptr{UInt8}(C_NULL) + try + old == C_NULL || deallocate!(old) + catch + # Error reporting is called from C callbacks and must never throw. + end + p = Ptr{UInt8}(C_NULL) + try + clean = replace(msg, '\0' => ' ') + bytes = codeunits(clean) + n = AC.checked_add(Int64(length(bytes)), Int64(1)) + p = Ptr{UInt8}(allocate!(n)) + p == C_NULL && return nothing + for (i, b) in enumerate(bytes) + unsafe_store!(p, b, i) + end + unsafe_store!(p, 0x00, length(bytes) + 1) + state.lasterror = p + catch + try + p == C_NULL || deallocate!(p) + catch + end + end + return nothing +end + +function _set_stream_exception!(state::ExportedStreamState, e) + try + _set_stream_error!(state, sprint(showerror, e)) + catch + # `_set_stream_error!` is itself best-effort, but keep the callback + # boundary closed if exception rendering fails before it is called. + _set_stream_error!(state, "stream callback failed") + end + return nothing +end + +function _publish_stream_result!(build, roots::Vector{Any}, result_slot, out, publish!) + key_slot = Ref{Int64}(0) + committed = false + try + _newroot(build, roots; result_slot=result_slot, key_slot=key_slot) + publish!(out, unsafe_load(result_slot[])) + committed = true + catch + # The result root became public inside Julia, but no usable C struct + # reached the consumer. Remove it immediately. Cleanup is best-effort + # here so it cannot replace the operation's original exception or + # cross the enclosing C callback boundary. + if !committed && key_slot[] != 0 + try + _cleanup_registered_root!(key_slot[]; require_released=false) + catch + end + key_slot[] = 0 + end + rethrow() + end + return nothing +end + +function _stream_get_schema_impl( + sp::Ptr{CArrowArrayStream}, + out::Ptr{CArrowSchema}, + publish!, +)::Cint + state = nothing + try + sp == C_NULL && return EINVAL + state, _ = _stream_state(sp) + state === nothing && return EINVAL + out == C_NULL && throw(ArgumentError("ArrowSchema output pointer is NULL")) + srel = @cfunction(_release_schema, Cvoid, (Ptr{CArrowSchema},)) + shell = Ref{Ptr{CArrowSchema}}(C_NULL) + _publish_stream_result!(Any[state.batchfield], shell, out, publish!) do root + _export_schema!(root, state.batchfield, srel) + end + return Cint(0) + catch e + state isa ExportedStreamState && _set_stream_exception!(state, e) + return EINVAL + end +end + +_stream_get_schema(sp::Ptr{CArrowArrayStream}, out::Ptr{CArrowSchema})::Cint = + _stream_get_schema_impl(sp, out, unsafe_store!) + +function _stream_get_next_impl( + sp::Ptr{CArrowArrayStream}, + out::Ptr{CArrowArray}, + publish!, +)::Cint + state = nothing + try + sp == C_NULL && return EINVAL + state, _ = _stream_state(sp) + state === nothing && return EINVAL + out == C_NULL && throw(ArgumentError("ArrowArray output pointer is NULL")) + if state.nextindex > length(state.batches) + # End of stream: a released (NULL-release) struct, per spec. + publish!( + out, + CArrowArray( + 0, + 0, + 0, + 0, + 0, + Ptr{Ptr{Cvoid}}(C_NULL), + Ptr{Ptr{CArrowArray}}(C_NULL), + Ptr{CArrowArray}(C_NULL), + Ptr{Cvoid}(C_NULL), + Ptr{Cvoid}(C_NULL), + ), + ) + return Cint(0) + end + b = state.batches[state.nextindex] + d = ArrayData( + StructType(), + b.nrows, + [BufferSlice()]; + children=collect(ArrayData, b.columns), + nullcount=0, + ) + validate_semantic(state.batchfield, d) + arel = @cfunction(_release_array, Cvoid, (Ptr{CArrowArray},)) + shell = Ref{Ptr{CArrowArray}}(C_NULL) + _publish_stream_result!(Any[d], shell, out, publish!) do root + _export_array!(root, d, arel) + end + state.nextindex += 1 + return Cint(0) + catch e + state isa ExportedStreamState && _set_stream_exception!(state, e) + return EINVAL + end +end + +_stream_get_next(sp::Ptr{CArrowArrayStream}, out::Ptr{CArrowArray})::Cint = + _stream_get_next_impl(sp, out, unsafe_store!) + +function _stream_get_last_error(sp::Ptr{CArrowArrayStream})::Ptr{UInt8} + try + sp == C_NULL && return Ptr{UInt8}(C_NULL) + state, _ = _stream_state(sp) + state === nothing && return Ptr{UInt8}(C_NULL) + return state.lasterror + catch + return Ptr{UInt8}(C_NULL) + end +end + +function _stream_release(sp::Ptr{CArrowArrayStream})::Cvoid + # Claim/commit with no error channel, like the node callbacks. Batch and + # schema roots already handed to the consumer keep their own lifetimes. + try + sp == C_NULL && return nothing + lock(REGISTRY_LOCK) do + stream = unsafe_load(sp) + stream.release == C_NULL && return nothing + control = stream.private_data + control == C_NULL && return nothing + key = unsafe_load(Ptr{Int64}(control + 8)) + state = get(STREAM_REGISTRY, key, nothing) + state === nothing && return nothing + pop!(STREAM_REGISTRY, key) + errorp = state.lasterror + state.lasterror = Ptr{UInt8}(C_NULL) + errorp == C_NULL || Libc.free(errorp) + _store_field!(sp, Val(:release), Ptr{Cvoid}(C_NULL)) + _store_field!(sp, Val(:private_data), Ptr{Cvoid}(C_NULL)) + Libc.free(control) + return nothing + end + catch + # A void C callback has no error channel. Never unwind into C. + end + return nothing +end + +# Core's structural schema invariants (endianness, UTF-8 names and metadata, +# valid descriptors) for the whole schema TREE. Stream schemas travel +# separately from any batch, so both stream directions apply this walk to the +# schema itself — a zero-batch stream never reaches batch validation. +function _validate_stream_schema(sch::Schema) + AC._validate_schema(sch) + for f in sch.fields + _validate_stream_field(f) + end + return nothing +end + +function _validate_stream_field(f::Field) + isvalid(f.name) || throw(ValidationError("field name is not valid UTF-8")) + AC._validate_metadata(f.metadata, "field") + AC._validate_descriptor_of(f.type) + for c in f.children + _validate_stream_field(c) + end + return nothing +end + +""" + export_stream!(sp::Ptr{CArrowArrayStream}, sch::Schema, batches) + +Fill a CALLER-owned ArrowArrayStream struct (the C stream convention: the +producer fills, the consumer owns the struct storage) streaming `batches` as +struct-typed arrays whose children are the schema's columns. The stream's +registry root keeps schema fields and batches reachable until `release`; +every `get_schema`/`get_next` result is its own export root with the same +lifecycle as `to_c_data` output. +""" +export_stream!( + sp::Ptr{CArrowArrayStream}, + sch::Schema, + batches::AbstractVector{AC.RecordBatch}, +) = _export_stream!(sp, sch, batches, Libc.malloc, Libc.free, unsafe_store!) + +function _export_stream!( + sp::Ptr{CArrowArrayStream}, + sch::Schema, + batches::AbstractVector{AC.RecordBatch}, + allocate!, + deallocate!, + publish!, +) + sp == C_NULL && throw(ArgumentError("ArrowArrayStream pointer is NULL")) + _validate_stream_schema(sch) + for b in batches + length(b.columns) == length(sch.fields) || + throw(ValidationError("stream batch column count does not match the schema")) + end + # The stream's struct-typed schema node carries the schema-level + # metadata (the C++/pyarrow convention for `schema.metadata`). + batchfield = Field( + "", + StructType(); + nullable=false, + metadata=sch.metadata, + children=collect(Field, sch.fields), + ) + state = ExportedStreamState( + batchfield, + collect(AC.RecordBatch, batches), + 1, + Ptr{UInt8}(C_NULL), + ) + get_schema = + @cfunction(_stream_get_schema, Cint, (Ptr{CArrowArrayStream}, Ptr{CArrowSchema})) + get_next = + @cfunction(_stream_get_next, Cint, (Ptr{CArrowArrayStream}, Ptr{CArrowArray})) + get_last_error = + @cfunction(_stream_get_last_error, Ptr{UInt8}, (Ptr{CArrowArrayStream},)) + release = @cfunction(_stream_release, Cvoid, (Ptr{CArrowArrayStream},)) + control = Ptr{Cvoid}(C_NULL) + key = Int64(0) + havekey = false + try + control = Ptr{Cvoid}(allocate!(CONTROL_BLOCK_BYTES)) + control == C_NULL && throw(OutOfMemoryError()) + key = lock(REGISTRY_LOCK) do + NEXT_KEY[] = AC.checked_add(NEXT_KEY[], Int64(1)) + end + havekey = true + # A stream reuses the node control-block layout, but only the key at + # offset 8 is ever read; the state byte at offset 0 exists so the two + # block kinds stay interchangeable. + unsafe_store!(Ptr{UInt8}(control), 0x00) + unsafe_store!(Ptr{Int64}(control + 8), key) + lock(REGISTRY_LOCK) do + STREAM_REGISTRY[key] = state + end + publish!( + sp, + CArrowArrayStream(get_schema, get_next, get_last_error, release, control), + ) + return sp + catch + if havekey + lock(REGISTRY_LOCK) do + get(STREAM_REGISTRY, key, nothing) === state && pop!(STREAM_REGISTRY, key) + end + end + errorp = state.lasterror + state.lasterror = Ptr{UInt8}(C_NULL) + errorp == C_NULL || deallocate!(errorp) + control == C_NULL || deallocate!(control) + rethrow() + end +end + +_stream_registry_count() = lock(REGISTRY_LOCK) do + length(STREAM_REGISTRY) +end + +# ---- import half ---------------------------------------------------------- + +""" +One owner for one MOVED ArrowArrayStream, mirroring ForeignOwner: a malloc'd +copy of the moved struct gives the producer's callbacks a stable address, one +atomic flag picks the single releaser between explicit `release!` and the GC +finalizer, and post-release calls are the spec's own undefined behavior. +""" +mutable struct StreamOwner + const block::Ptr{CArrowArrayStream} + const producer_release::Ptr{Cvoid} + @atomic released::Bool + function StreamOwner(stream::CArrowArrayStream, registerfinalizer) + block = Libc.malloc(sizeof(CArrowArrayStream)) + block == C_NULL && throw(OutOfMemoryError()) + p = Ptr{CArrowArrayStream}(block) + o = try + unsafe_store!(p, stream) + _store_field!(p, Val(:release), Ptr{Cvoid}(C_NULL)) # inert until moved + new(p, stream.release, false) + catch + Libc.free(block) + rethrow() + end + try + registerfinalizer(release!, o) + catch + # The source still owns the producer stream. Free only the inert + # copy; an already-installed finalizer observes released=true. + release!(o) + rethrow() + end + return o + end +end +StreamOwner(stream::CArrowArrayStream) = StreamOwner(stream, finalizer) + +function _stream_owner_armed(o::StreamOwner) + GC.@preserve o begin + return unsafe_load(o.block).release != C_NULL + end +end + +function _arm_stream_owner!(o::StreamOwner) + (@atomic o.released) && error("cannot arm a released stream owner") + GC.@preserve o _store_field!(o.block, Val(:release), o.producer_release) + return nothing +end + +function _release_moved_stream_owner!(o::StreamOwner) + _stream_owner_armed(o) || _arm_stream_owner!(o) + release!(o) + return nothing +end + +# Run the moved stream's producer release callback on the malloc'd struct +# copy, check that the producer nulled the copy's release field, and free the +# copy. Exactly one atomic swap picks the single releaser between explicit +# calls and the GC finalizer. +function release!(o::StreamOwner) + @atomicswap(o.released = true) && return nothing + GC.@preserve o begin + cb = unsafe_load(o.block).release + if cb != C_NULL + ccall(cb, Cvoid, (Ptr{CArrowArrayStream},), o.block) + unsafe_load(o.block).release == C_NULL || ( + Libc.free(o.block); + error("C stream producer release did not mark the structure released") + ) + end + Libc.free(o.block) + end + return nothing +end + +""" + ImportedStream + +Consumer side of a moved ArrowArrayStream: `schema(s)` is fixed at import, +`nextbatch!(s)` pulls one struct-typed batch (returning `nothing` at end of +stream), and `release!(s)` ends the producer's stream exactly once. Each +pulled batch owns its own ForeignOwner and outlives the stream if the caller +keeps it. Producer-reported failures surface as `ValidationError`s carrying +the producer's `get_last_error` text. One stream call at a time: an +overlapping `nextbatch!`/`release!` throws `ConcurrencyViolationError` +rather than entering the producer's callbacks concurrently. +""" +mutable struct ImportedStream <: AC.RecordBatchSource + const owner::StreamOwner + const batchfield::Field + const schema::Schema + done::Bool + @atomic pulling::Bool +end + +ImportedStream(owner, batchfield, schema, done) = + ImportedStream(owner, batchfield, schema, done, false) + +AC.schema(s::ImportedStream) = s.schema + +function _claimstreamcall!(s::ImportedStream) + _, ok = @atomicreplace s.pulling false => true + ok || throw( + Base.ConcurrencyViolationError( + "ImportedStream supports only one active stream call", + ), + ) + return nothing +end + +_finishstreamcall!(s::ImportedStream) = (@atomic :release s.pulling = false) + +""" + release!(s::ImportedStream) + +Run the producer's stream release callback exactly once (later calls and the +GC finalizer are no-ops); batches already pulled keep their own owners and +stay valid. +""" +function release!(s::ImportedStream) + _claimstreamcall!(s) + try + return release!(s.owner) + finally + _finishstreamcall!(s) + end +end + +function _stream_call_failed(o::StreamOwner, what::AbstractString) + msg = "C stream $what failed" + GC.@preserve o begin + cb = unsafe_load(o.block).get_last_error + if cb != C_NULL + p = ccall(cb, Ptr{UInt8}, (Ptr{CArrowArrayStream},), o.block) + p == C_NULL || (msg *= ": " * _import_cstring(p, "stream error")) + end + end + throw(ValidationError(msg)) +end + +""" + from_c_stream(sp::Ptr{CArrowArrayStream}) -> ImportedStream + +Move a producer's stream (copy the struct, null the source release) and read +its schema. The schema must be a struct-typed batch schema, per the C stream +convention; its fields become the imported `Schema`. +""" +function from_c_stream(sp::Ptr{CArrowArrayStream}) + sp == C_NULL && throw(ArgumentError("ArrowArrayStream pointer is NULL")) + stream = unsafe_load(sp) + stream.release == C_NULL && throw(ArgumentError("cannot import a released stream")) + ( + stream.get_schema == C_NULL || + stream.get_next == C_NULL || + stream.get_last_error == C_NULL + ) && throw(ArgumentError("C stream is missing required callbacks")) + owner = StreamOwner(stream) + moved = false + try + _store_field!(sp, Val(:release), Ptr{Cvoid}(C_NULL)) # the move commit + moved = true + _arm_stream_owner!(owner) + out = Ref( + CArrowSchema( + Ptr{UInt8}(C_NULL), + Ptr{UInt8}(C_NULL), + Ptr{UInt8}(C_NULL), + 0, + 0, + Ptr{Ptr{CArrowSchema}}(C_NULL), + Ptr{CArrowSchema}(C_NULL), + Ptr{Cvoid}(C_NULL), + Ptr{Cvoid}(C_NULL), + ), + ) + status = GC.@preserve owner out ccall( + unsafe_load(owner.block).get_schema, + Cint, + (Ptr{CArrowArrayStream}, Ptr{CArrowSchema}), + owner.block, + Base.unsafe_convert(Ptr{CArrowSchema}, out), + ) + status == 0 || _stream_call_failed(owner, "get_schema") + sch = out[] + batchfield = GC.@preserve out try + _preflight_schema(sch) + _import_field(sch) + finally + _release_c_schema!(Base.unsafe_convert(Ptr{CArrowSchema}, out), sch) + end + batchfield.type isa StructType || + throw(ValidationError("C stream schema must be a struct-typed batch schema")) + schema = Schema(collect(Field, batchfield.children); metadata=batchfield.metadata) + _validate_stream_schema(schema) + return ImportedStream(owner, batchfield, schema, false) + catch + moved ? _release_moved_stream_owner!(owner) : release!(owner) + rethrow() + end +end + +AC.nextbatch!(s::ImportedStream) = _nextbatch!(s, ForeignOwner) + +function _nextbatch!(s::ImportedStream, ownerfactory) + _claimstreamcall!(s) + try + # Fail closed on a released stream even when it already ended naturally: + # release terminates the consumer contract, not just the batch supply. + (@atomic s.owner.released) && + throw(ArgumentError("cannot pull from a released stream")) + s.done && return nothing + out = Ref( + CArrowArray( + 0, + 0, + 0, + 0, + 0, + Ptr{Ptr{Cvoid}}(C_NULL), + Ptr{Ptr{CArrowArray}}(C_NULL), + Ptr{CArrowArray}(C_NULL), + Ptr{Cvoid}(C_NULL), + Ptr{Cvoid}(C_NULL), + ), + ) + status = GC.@preserve s out ccall( + unsafe_load(s.owner.block).get_next, + Cint, + (Ptr{CArrowArrayStream}, Ptr{CArrowArray}), + s.owner.block, + Base.unsafe_convert(Ptr{CArrowArray}, out), + ) + status == 0 || _stream_call_failed(s.owner, "get_next") + arr = out[] + if arr.release == C_NULL + s.done = true + return nothing + end + # The producer filled consumer-owned storage. Build an inert destination + # owner first. If that construction fails, the live source slot still owns + # the result and must release it. Then null the source and arm the copy. + batchowner = try + ownerfactory(arr)::ForeignOwner + catch + GC.@preserve out _release_c_array!( + Base.unsafe_convert(Ptr{CArrowArray}, out), + arr, + ) + rethrow() + end + moved = false + d = try + GC.@preserve out _store_field!( + Base.unsafe_convert(Ptr{CArrowArray}, out), + :release, + Ptr{Cvoid}(C_NULL), + ) + moved = true + _arm_foreign_owner!(batchowner) + _preflight_array(s.batchfield, arr) + d0 = _import_array(s.batchfield, arr, batchowner) + validate_semantic(s.batchfield, d0) + d0 + catch + moved ? _release_moved_owner!(batchowner) : release!(batchowner) + rethrow() + end + return AC.RecordBatch(s.schema, collect(ArrayData, d.children), d.len) + finally + _finishstreamcall!(s) + end +end diff --git a/src/columnconstruction.jl b/src/columnconstruction.jl new file mode 100644 index 00000000..d2a7184e --- /dev/null +++ b/src/columnconstruction.jl @@ -0,0 +1,5592 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 construction: Julia columns -> Arrow Field + ArrayData. +# +# This file owns fresh inference, retained-schema reconstruction, recursive +# ArrowTypes lowering, dictionary pool policy, and field metadata. The write +# facade binds partitions and delegates each complete column through +# `_constructcolumn`. +# ============================================================================= + +"Copy a Field's metadata into stable, owned key-value pairs." +_fieldmetadata(f::AC.Field) = + f.metadata === nothing ? nothing : collect(Pair{String,String}, f.metadata) + +"One native (Field, ArrayData) column from a Julia vector, facade conversions included." +function _constructnativepart(name::String, v::AbstractVector; context=nothing) + T = Base.nonmissingtype(eltype(v)) + # With TimeZones loaded, a fresh ZonedDateTime column writes as a + # timezone-declared millisecond timestamp (ext/ArrowTimeZonesExt.jl) + # instead of reflecting the zone's whole transition table as a struct. + zoned = _zonednativepart(name, v, T) + zoned === nothing || return zoned + if eltype(v) === Missing + t = AC.NullType() + return AC.Field(name, t; nullable=true), + AC.ArrayData(t, length(v), AC.BufferSlice[]; nullcount=length(v)) + elseif T === Union{} + throw( + ArgumentError( + "column $name has bottom element type Union{} and cannot infer " * + "an Arrow type; give the empty column a declared element type", + ), + ) + elseif T <: Dates.Date + return _constructtemporalpart( + name, + v, + AC.DateType(AC.DAY), + x -> Int32(Dates.value(x) - _EPOCH_DAYS), + ) + elseif T <: Dates.DateTime + return _constructtemporalpart( + name, + v, + AC.TimestampType(AC.MILLISECOND, nothing), + x -> Int64(Dates.value(x) - Dates.UNIXEPOCH), + ) + elseif T <: Dates.Time + return _constructtemporalpart( + name, + v, + AC.TimeType(AC.NANOSECOND, 64), + x -> Int64(Dates.value(x)), + ) + elseif T <: Dates.Period && + T <: Union{Dates.Second,Dates.Millisecond,Dates.Microsecond,Dates.Nanosecond} + unit = + T <: Dates.Second ? AC.SECOND : + T <: Dates.Millisecond ? AC.MILLISECOND : + T <: Dates.Microsecond ? AC.MICROSECOND : AC.NANOSECOND + return _constructtemporalpart( + name, + v, + AC.DurationType(unit), + x -> Int64(Dates.value(x)), + ) + elseif T <: NamedTuple + any(ismissing, v) && throw( + ArgumentError( + "column $name has nullable NamedTuple rows that require the " * + "recursive column adapter", + ), + ) + isconcretetype(T) || throw( + ArgumentError( + "column $name has abstract NamedTuple element type; give it a " * + "concrete NamedTuple type with declared field names and types", + ), + ) + if fieldcount(T) == 0 + t = AC.StructType() + return AC.Field(name, t; nullable=false, children=AC.Field[]), + AC.ArrayData( + t, + length(v), + [AC.BufferSlice()]; + children=AC.ArrayData[], + nullcount=0, + ) + end + # Preserve the declared child types. A value-narrowing comprehension + # turns an empty child into `Any[]` and an all-missing nullable child + # into `Missing[]`, so neither can recover its Arrow descriptor. + cols = NamedTuple{fieldnames(T)}( + ntuple( + i -> collect(fieldtype(T, i), (getfield(x, i) for x in v)), + fieldcount(T), + ), + ) + return AC.fromjulia_struct(name, cols) + elseif T <: AbstractString && T != String + return AC.fromjulia(name, _missings_to(String, v)) + elseif T === Any || (T <: AbstractVector && eltype(T) === Any) + # Materialized facade columns are Any-eltype for composite layouts; + # one narrowing pass recovers list columns (inner vectors narrow + # element-wise, empties adopt the joined element type). + w = _narrowlists(v) + NW = Base.nonmissingtype(eltype(w)) + (NW === Any || (NW <: AbstractVector && eltype(NW) === Any)) && throw( + ArgumentError( + "column $name has element type Any and cannot be narrowed to " * + "a writable Arrow column; give it a concrete element type", + ), + ) + return _constructpart(name, w; context) + else + return AC.fromjulia(name, _plainvector(v)) + end +end + +"One (Field, ArrayData) column, including the ArrowTypes compatibility adapter." +function _constructpart( + name::String, + v::AbstractVector; + context=nothing, + narrowabstract::Bool=true, +) + context === nothing && (context = _WriterContext()) + if eltype(v) === Any + narrowed = _narrowlists(v) + eltype(narrowed) === Any && throw( + ArgumentError( + "column $name has element type Any and cannot be narrowed to a " * + "writable Arrow column; give it a concrete element type", + ), + ) + return _constructpart(name, narrowed; context) + end + if narrowabstract + runtimeparts = _narrowabstractwriterparts(name, AbstractVector[v], context) + if runtimeparts !== nothing + narrowed = only(runtimeparts) + # Julia can canonicalize a Union of concrete Tuple runtime types + # back to the original abstract declaration, so eltype may be + # unchanged. One narrowing pass is enough; recursing on the same + # declaration would loop. + eltype(narrowed) === eltype(v) || + return _constructpart(name, narrowed; context, narrowabstract=false) + end + end + return _arrowtypes_needs(v, context) ? _constructarrowtypespart(name, v; context) : + _constructnativepart(name, v; context) +end + +"Walk a bounded writer value graph, with optional short-circuit inspection." +function _walkwritercontainer!(visit, x, active::Base.IdSet{Any}, depth::Int) + x === missing && return false + depth < _MAX_WRITER_SCHEMA_DEPTH || throw( + ArgumentError( + "ArrowTypes value nesting exceeds the supported depth " * + "$_MAX_WRITER_SCHEMA_DEPTH", + ), + ) + visit(x) && return true + if x isa NamedTuple || x isa Tuple + for value in x + _walkwritercontainer!(visit, value, active, depth + 1) === true && return true + end + elseif x isa Pair + _walkwritercontainer!(visit, first(x), active, depth + 1) === true && return true + _walkwritercontainer!(visit, last(x), active, depth + 1) === true && return true + elseif x isa AbstractVector || x isa AbstractDict + x in active && + throw(ArgumentError("recursive ArrowTypes value container cannot be written")) + push!(active, x) + try + if x isa AbstractDict + for (key, value) in pairs(x) + _walkwritercontainer!(visit, key, active, depth + 1) === true && + return true + _walkwritercontainer!(visit, value, active, depth + 1) === true && + return true + end + else + for value in x + _walkwritercontainer!(visit, value, active, depth + 1) === true && + return true + end + end + finally + delete!(active, x) + end + end + return false +end + +"Reject recursive or excessively deep value-container graphs before narrowing." +function _preflightwritercontainers(v::AbstractVector) + active = Base.IdSet{Any}() + for value in v + _walkwritercontainer!(_ -> false, value, active, 0) + end + return nothing +end + +"Narrow an Any-eltype column, recovering list-of-T structure when present." +function _narrowlists(v::AbstractVector) + # Narrowing can turn an Any column into a concrete list type before the + # ArrowTypes value walk runs. Reject cycles and excessive nesting first, + # without resolving any user traits. + _preflightwritercontainers(v) + # map(identity, …) re-infers a narrowed element type from the values: + # inner maps narrow each row, the outer map narrows the column. + w = map(x -> x isa AbstractVector ? map(identity, x) : x, v) + w = map(identity, w) + NT = Base.nonmissingtype(eltype(w)) + NT <: AbstractVector || return w + # Join the inner element types (empties narrow to Union{} and would + # otherwise poison the join), then retype every inner vector. + E = Union{} + for x in w + x === missing && continue + isempty(x) && continue + E = typejoin(E, eltype(x)) + end + E === Union{} && (E = Any) + E === Any && return w + hasm = eltype(w) >: Missing + S = hasm ? Union{Missing,Vector{E}} : Vector{E} + return S[x === missing ? missing : convert(Vector{E}, x) for x in w] +end + +# An ArrowStrings column IS Utf8View memory: its payload vector is the views +# buffer and its byte buffers are the variadic data buffers — no copy, no +# String materialization. Nullability comes from the column's eltype. +# Both keywords are accepted for dispatch uniformity; a StringVector is +# already Arrow memory, so no context or narrowing applies. +function _constructpart( + name::String, + v::ArrowStrings.StringVector; + context=nothing, + narrowabstract::Bool=true, +) + return AC.fromviewentries(name, v.payloads, v.buffers; nullable=eltype(v) >: Missing) +end + +"Concrete Vector with an exact Union{Missing,T} or T eltype for fromjulia." +function _plainvector(v::AbstractVector) + T = eltype(v) + return v isa Vector{T} ? v : collect(T, v) +end + +_missings_to(::Type{S}, v) where {S} = + eltype(v) >: Missing ? Union{Missing,S}[x === missing ? missing : S(x) for x in v] : + S[S(x) for x in v] + +"Construct temporal ArrayData from nullable Int64 storage values." +function _temporaldata(t::AC.ArrowType, storage::AbstractVector) + _, d0 = AC.fromjulia("storage", storage) + # Narrow to Int32 for any 32-bit temporal descriptor (Date32 fresh, + # Time32 retained); every other temporal layout keeps the inferred Int64 + # buffer. + buffers = d0.buffers + if AC.primwidth(t) == 4 + narrow = Vector{Int32}(undef, length(storage)) + for (i, x) in enumerate(storage) + narrow[i] = x === missing ? Int32(0) : Int32(x) + end + buffers = AC.BufferSlice[d0.buffers[1], AC._databuffer(narrow)] + end + return AC._arraydata( + t, + d0.len, + buffers, + 0, + AC.ArrayData[], + nothing, + d0.owner, + AC.nullcount(d0), + ) +end + +"Construct one fresh temporal part and preserve its validity." +function _constructtemporalpart( + name::String, + v::AbstractVector, + t::AC.ArrowType, + tostorage::F, +) where {F} + storage = Union{Missing,Int64}[x === missing ? missing : Int64(tostorage(x)) for x in v] + return AC.Field(name, t; nullable=eltype(v) >: Missing), _temporaldata(t, storage) +end + +# ArrowTypes write adapter. ArrowCore intentionally stays conversion-free; +# this layer lowers logical Julia values before it asks Core to build buffers. + +"Stable metadata merge; `extra` wins except for protected extension labels." +function _mergemetapairs(base, extra; protectextension::Bool=false) + out = Pair{String,String}[] + base === nothing || append!(out, String(first(kv)) => String(last(kv)) for kv in base) + extra === nothing && return isempty(out) ? nothing : out + for kv in extra + key = String(first(kv)) + value = String(last(kv)) + i = findfirst(p -> first(p) == key, out) + if protectextension && + (key == _EXTENSION_NAME_KEY || key == _EXTENSION_METADATA_KEY) && + i !== nothing + continue + end + if i === nothing + push!(out, key => value) + else + out[i] = key => value + end + end + return isempty(out) ? nothing : out +end + +"Attach this logical type's extension label without disturbing child labels." +function _arrowtypeslogicalfield(context, f::AC.Field, T; nullable::Bool=f.nullable) + extension = _writerextension!(context, T) + extension === nothing && return f + metadata = _mergemetapairs(_fieldmetadata(f), extension) + return AC.Field( + f.name, + f.type; + nullable=nullable, + metadata=metadata, + children=collect(AC.Field, f.children), + ) +end + +_arrowtypesnativetype(T) = + T <: Union{ + Int8, + Int16, + Int32, + Int64, + UInt8, + UInt16, + UInt32, + UInt64, + Float16, + Float32, + Float64, + Bool, + AbstractString, + Dates.Date, + Dates.DateTime, + Dates.Time, + Dates.Period, + } || _zonedwritertype(T) + +# The TimeZones extension's write hooks. With the extension absent both +# answer "not mine", so every column keeps this file's ordinary routes. +function _zonedwritertype(T) + ext = Base.get_extension(@__MODULE__, :ArrowTimeZonesExt) + return ext !== nothing && ext.iszonedtype(T) +end + +function _zonednativepart(name::String, v::AbstractVector, T) + ext = Base.get_extension(@__MODULE__, :ArrowTimeZonesExt) + ext === nothing && return nothing + return ext.zonednativepart(name, v, T) +end + +function _arrowtypesneedstype(T, arrowtype, hasarrowname, depth::Int=0) + depth < _MAX_WRITER_SCHEMA_DEPTH || throw( + ArgumentError( + "ArrowTypes storage type exceeds the supported depth " * + "$_MAX_WRITER_SCHEMA_DEPTH at $T", + ), + ) + T === Union{} && return false + T === Missing && return false + T === Any && return false + if T isa Union + members = Base.uniontypes(T) + nonnull = count(!=(Missing), members) + nullable_struct = Missing in members && any(member -> member <: NamedTuple, members) + return nonnull > 1 || + nullable_struct || + any( + member -> + _arrowtypesneedstype(member, arrowtype, hasarrowname, depth + 1), + members, + ) + end + NT = Base.nonmissingtype(T) + NT !== T && return _arrowtypesneedstype(NT, arrowtype, hasarrowname, depth + 1) + arrowtype(T) !== T && return true + hasarrowname(T) && return true + T <: AbstractDict && return true + T <: Tuple && return true + if T <: NamedTuple + return any( + i -> _arrowtypesneedstype(fieldtype(T, i), arrowtype, hasarrowname, depth + 1), + 1:fieldcount(T), + ) + end + if T <: AbstractVector + return _arrowtypesneedstype(eltype(T), arrowtype, hasarrowname, depth + 1) + end + _arrowtypesnativetype(T) && return false + # Plain concrete structs are written by field reflection; native facade + # values such as Dates must not be routed through ArrowTypes' default + # StructKind. + return isconcretetype(T) && !isprimitivetype(T) +end + +_arrowtypesneedsvalue(x, context, depth::Int=0) = + _walkwritercontainer!( + value -> _writerneedstype!(context, typeof(value)), + x, + context.values, + depth, + ) === true + +function _arrowtypes_needs(v::AbstractVector, context) + _writerneedstype!(context, eltype(v)) && return true + # Abstract/Any containers do not expose nested logical types. Inspect only + # until the first ArrowTypes value is found; concrete columns take no scan. + return (eltype(v) === Any || !isconcretetype(eltype(v))) && + any(value -> _arrowtypesneedsvalue(value, context), v) +end + +_arrowtypesconcreteorunion(T) = + isconcretetype(T) || + (T isa Union && all(_arrowtypesconcreteorunion, Base.uniontypes(T))) + +const _WriterTypeRoutes = IdDict{Type,Int} + +function _arrowtypesstorageisspecified(context, T::Type, path::Vector{Type}=Type[]) + Base.@nospecialize T + _arrowtypesconcreteorunion(T) && return true + _arrowtypesnativetype(T) && return true + length(path) < _MAX_WRITER_SCHEMA_DEPTH || throw( + ArgumentError( + "ArrowTypes storage schema exceeds the supported depth " * + "$_MAX_WRITER_SCHEMA_DEPTH at $T", + ), + ) + T in path && throw( + ArgumentError( + "recursive ArrowTypes storage schema: " * + join((string(S) for S in (path..., T)), " -> "), + ), + ) + push!(path, T) + try + T isa Union && return all( + member -> _arrowtypesstorageisspecified(context, member, path), + Base.uniontypes(T), + ) + (T === Tuple || T === NamedTuple || (T <: Tuple && Base.isvatuple(T))) && + return false + T <: Union{NamedTuple,Tuple} && return all( + i -> _arrowtypesstorageisspecified(context, fieldtype(T, i), path), + 1:fieldcount(T), + ) + T <: AbstractVector && + return _arrowtypesstorageisspecified(context, eltype(T), path) + T <: AbstractDict && + return _arrowtypesstorageisspecified(context, keytype(T), path) && + _arrowtypesstorageisspecified(context, valtype(T), path) + storage = _writerstoragetype!(context, T) + storage === T && return false + return _arrowtypesstorageisspecified(context, storage, path) + finally + pop!(path) + end +end + +function _writerroute(routes::_WriterTypeRoutes, runtime::Type) + haskey(routes, runtime) && return routes[runtime] + route = -1 + for (declared, candidate) in routes + runtime <: declared || continue + (route == -1 || route == candidate) || throw( + ArgumentError( + "writer type $runtime matches conflicting retained-field routes " * + "$route and $candidate", + ), + ) + route = candidate + end + route >= 0 || throw(ArgumentError("writer type $runtime has no retained Field route")) + routes[runtime] = route + return route +end + +function _writerpromoteunion(T, S) + promoted = promote_type(T, S) + return _arrowtypesconcreteorunion(promoted) ? promoted : Union{T,S} +end + +const _WRITER_CONVERSION_ERRORS = Union{MethodError,InexactError,OverflowError,TypeError} + +function _writerconvert(::Type{T}, value) where {T} + value isa T && return value + if T isa Union + for member in Base.uniontypes(T) + try + return _writerconvert(member, value) + catch err + err isa _WRITER_CONVERSION_ERRORS || rethrow() + end + end + end + return convert(T, value) +end + +_writercolumnname(context, fallback::String) = + isempty(context.column) ? fallback : context.column + +function _writertoarrow(context, value, writertype::Type, fallback::String) + try + return ArrowTypes.toarrow(value) + catch err + err isa Union{InterruptException,OutOfMemoryError} && rethrow() + err isa _WRITER_CONVERSION_ERRORS || rethrow() + column = _writercolumnname(context, fallback) + throw( + ArgumentError( + "ArrowTypes.toarrow for writer type $writertype failed for " * + "column $column: $(sprint(showerror, err))", + ), + ) + end +end + +function _writerconvertlowered( + ::Type{T}, + value, + context, + writertype::Type, + fallback::String, +) where {T} + try + return _writerconvert(T, value) + catch err + err isa Union{InterruptException,OutOfMemoryError} && rethrow() + err isa _WRITER_CONVERSION_ERRORS || rethrow() + column = _writercolumnname(context, fallback) + throw( + ArgumentError( + "ArrowTypes.toarrow for writer type $writertype produced storage " * + "type $(typeof(value)), which cannot be represented as $T for " * + "column $column", + ), + ) + end +end + +function _writerconvertedcolumn( + ::Type{T}, + values, + context, + writertype::Type, + name::String, +) where {T} + out = Vector{T}(undef, length(values)) + for i in eachindex(values) + out[i] = _writerconvertlowered(T, values[i], context, writertype, name) + end + return out +end + +# Function barrier for the known-storage lowering loop: the storage type is +# computed at runtime, so specializing on it here keeps the per-element +# toarrow/convert path monomorphic instead of boxing every lowered value. +function _writerloweredcolumn( + ::Type{T}, + v::AbstractVector, + context, + writertype::Type, + name::String, +) where {T} + out = Vector{T}(undef, length(v)) + for (i, value) in enumerate(v) + lowered = _writertoarrow(context, value, writertype, name) + out[i] = _writerconvertlowered(T, lowered, context, writertype, name) + end + return out +end + +function _arrowtypesmappedcolumn(name::String, v::AbstractVector, declared, context) + S = eltype(v) + logical = Base.nonmissingtype(S) + storage = Missing <: S ? Union{Missing,declared} : declared + # Buffer builders index from 1; a non-one-based vector must go through + # the copy path. + S === storage && _arrowtypesconcreteorunion(S) && firstindex(v) == 1 && return v + + # An empty logical column has no runtime evidence. A recursively + # materializable declared storage shape is still authoritative; unsized + # Tuple and other inference-only declarations fall back to Arrow Null, as + # ArrowTypes.ToArrow does. + isempty(v) && + return _arrowtypesstorageisspecified(context, declared) ? storage[] : Missing[] + + if _arrowtypesconcreteorunion(storage) || + _arrowtypesstorageisspecified(context, declared) + return _writerloweredcolumn(storage, v, context, logical, name) + end + + lowered = Any[] + observed = Type[] + seen = Base.IdSet{Type}() + sawmissing = false + for value in v + mapped = _writertoarrow(context, value, logical, name) + push!(lowered, mapped) + if value === missing + sawmissing = true + continue + end + runtime = typeof(mapped) + runtime in seen && continue + _writerregisterstorageinference!(context, name, runtime) + push!(seen, runtime) + push!(observed, runtime) + end + if isempty(observed) + # A concrete logical declaration keeps ArrowTypes.ToArrow's + # all-missing fallback type. Abstract inference-only storage has no + # shape evidence, so the column is written as Arrow Null (what + # ToArrow returns: Missing[]). + fallbacktype = + _arrowtypesconcreteorunion(S) ? eltype(ArrowTypes.ToArrow(v)) : Missing + return _writerconvertedcolumn(fallbacktype, lowered, context, logical, name) + end + resulttype = reduce(_writerpromoteunion, observed) + sawmissing && (resulttype = Union{Missing,resulttype}) + return _writerconvertedcolumn(resulttype, lowered, context, logical, name) +end + +function _constructarrowtypespart( + name::String, + v::AbstractVector; + extension_shape::Bool=false, + context=nothing, +) + logical = Base.nonmissingtype(eltype(v)) + # `Base.nonmissingtype(Missing) === Union{}`. Handle a pure-null child + # before trait dispatch: bottom is a subtype of every container type and + # would otherwise enter an unrelated ArrowKind branch. + eltype(v) === Missing && return _constructnativepart(name, v; context) + logical === Union{} && return _constructnativepart(name, v; context) + context === nothing && (context = _WriterContext()) + return _withwriterschema(context, logical) do + _constructarrowtypespart_impl(name, v, logical; extension_shape, context) + end +end + +function _constructarrowtypespart_impl( + name::String, + v::AbstractVector, + logical; + extension_shape::Bool, + context, +) + # VersionNumber's legacy reflected representation contains variable-length + # tuples. String is an equivalent extension storage form: ArrowTypes' + # one-argument JuliaType registration and default VersionNumber(String) + # constructor read it in both Arrow 2.x and 3.x. + if logical === VersionNumber + mapped = + eltype(v) >: Missing ? + Union{Missing,String}[x === missing ? missing : string(x) for x in v] : + String[string(x) for x in v] + f, d = _constructnativepart(name, mapped; context) + return _arrowtypeslogicalfield(context, f, logical), d + end + declared = _writerstoragetype!(context, logical) + if Missing <: eltype(v) && + !(declared isa Union) && + _writerkind!(context, declared) isa ArrowTypes.NullKind + throw( + ArgumentError( + "nullable logical column $name maps its non-missing values to Arrow " * + "Null storage, so outer missing values cannot be distinguished; " * + "use DictEncode to preserve the two states", + ), + ) + end + if Missing <: eltype(v) && declared isa Union && Missing <: declared + throw( + ArgumentError( + "ArrowTypes logical type $logical uses Missing as a storage Union " * + "branch, so a nullable logical column cannot distinguish that " * + "storage value from an outer missing value", + ), + ) + end + if logical isa Union && declared === logical + # A declared Julia Union owns one Arrow child per declared member. + # Build those members independently so whole-column promotion cannot + # erase abstract-but-writable branches or same-storage logical types. + f, d = _arrowtypesunioncolumn(name, v; extension_shape, context) + return _arrowtypeslogicalfield(context, f, logical), d + end + # Abstract storage must be inferred from observed values. A recursively + # materializable declaration remains schema authority for empty columns, + # including every declared Union branch nested in a composite. + mapped = _arrowtypesmappedcolumn(name, v, declared, context) + ownshape = _writerextensionshape(context, logical, declared) + f, d = if declared isa Union && Missing <: declared + # Missing is a VALUE in this logical type's storage Union. Keep an + # explicit Null child; scalar null validity would erase the difference + # between that value and an outer nullable logical slot. + _arrowtypesunioncolumn( + name, + mapped; + extension_shape=extension_shape || ownshape, + context, + ) + else + _arrowtypesstoragecolumn( + name, + mapped; + extension_shape=extension_shape || ownshape, + context, + ) + end + # A logical type stored as Arrow Null is not nullable unless the column + # itself admits missing; otherwise the label's own values would read back + # as missing. + nullable = f.type isa AC.NullType && !(Missing <: eltype(v)) ? false : f.nullable + return _arrowtypeslogicalfield(context, f, logical; nullable=nullable), d +end + +function _arrowtypesstoragecolumn( + name::String, + v::AbstractVector; + extension_shape::Bool, + context, +) + eltype(v) === Missing && return _constructnativepart(name, v; context) + S = Base.nonmissingtype(eltype(v)) + S === Union{} && return _constructnativepart(name, v; context) + if S isa Union + return _arrowtypesunioncolumn(name, v; extension_shape, context) + end + _arrowtypesnativetype(S) && return _constructnativepart(name, v; context) + S <: NamedTuple && return _arrowtypesstructcolumn(name, v, S; extension_shape, context) + kind = _writerkind!(context, S) + if kind isa ArrowTypes.NullKind + nulls = S === Nothing ? fill(missing, length(v)) : v + return _constructnativepart(name, nulls; context) + end + if kind isa ArrowTypes.FixedSizeListKind + return _arrowtypesfixedlistcolumn(name, v, kind; extension_shape, context) + elseif kind isa ArrowTypes.MapKind + return _arrowtypesmapcolumn(name, v, S; extension_shape, context) + elseif kind isa ArrowTypes.ListKind + if S <: AbstractString + return _constructnativepart(name, v; context) + end + return _arrowtypeslistcolumn(name, v, S; extension_shape, context) + elseif kind isa ArrowTypes.StructKind + return _arrowtypesstructcolumn(name, v, S; extension_shape, context) + elseif kind isa Union{ArrowTypes.PrimitiveKind,ArrowTypes.BoolKind} + return _constructnativepart(name, v; context) + end + throw( + ArgumentError( + "unsupported ArrowTypes storage kind $(typeof(kind)) for column $name", + ), + ) +end + +# Arrow Union child limit; shares the ArrowTypes storage-Union limit. +const _MAX_WRITER_UNION_BRANCHES = _MAX_ARROWTYPE_UNION_BRANCHES +const _MAX_INFERRED_WRITER_TYPES = 8 + +_writertypevariants(T::Type) = + T === Union{} ? Type[] : T isa Union ? Type[Base.uniontypes(T)...] : Type[T] + +function _checkedwritervariants(owner::String, T::Type; allowempty::Bool=false) + variants = _writertypevariants(T) + length(variants) <= _MAX_WRITER_UNION_BRANCHES || throw( + ArgumentError("$owner has more than $_MAX_WRITER_UNION_BRANCHES Union branches"), + ) + allowempty || !isempty(variants) || throw(ArgumentError("$owner has no Union branches")) + return variants +end + +function _mergewritertypes(owner::String, left, right, limit::Int, kind::String) + types = Type[] + seen = Base.IdSet{Type}() + for source in (left, right), T in source + T in seen && continue + length(types) < limit || throw(ArgumentError("$owner has more than $limit $kind")) + push!(seen, T) + push!(types, T) + end + return types +end + +function _writeruniontype(types) + uniontype = Union{} + for T in types + uniontype = Union{uniontype,T} + end + return uniontype +end + +"Construct a fresh dense Arrow Union from one concrete Julia Union element type." +function _arrowtypesunioncolumn( + name::String, + v::AbstractVector; + extension_shape::Bool, + context, +) + variants = _checkedwritervariants("Arrow Union column $name", eltype(v)) + + childvalues = Any[T[] for T in variants] + typeids = Vector{Int8}(undef, length(v)) + offsets = Vector{Int32}(undef, length(v)) + for (i, x) in enumerate(v) + pos = findfirst(T -> x isa T, variants) + pos === nothing && throw( + ArgumentError( + "Arrow Union column $name contains $(typeof(x)), which is not in " * + "its declared element type $(eltype(v))", + ), + ) + child = childvalues[pos] + length(child) <= typemax(Int32) || + throw(ArgumentError("Arrow Union column $name exceeds the Int32 offset range")) + typeids[i] = Int8(pos - 1) + offsets[i] = Int32(length(child)) + push!(child, x) + end + + childfields = AC.Field[] + childdata = AC.ArrayData[] + for i in eachindex(variants) + field, data = _arrowtypeschildcolumn( + string(i - 1), + childvalues[i]; + extension_shape=extension_shape, + union_member=true, + context, + ) + push!(childfields, field) + push!(childdata, data) + end + ids = Int8.(0:(length(variants) - 1)) + t = AC.UnionType(AC.DenseMode, ids) + field = AC.Field(name, t; nullable=Missing in variants, children=childfields) + data = AC.ArrayData( + t, + length(v), + [AC._databuffer(typeids), AC._databuffer(offsets)]; + children=childdata, + nullcount=0, + ) + return field, data +end + +"Construct one child while keeping unmarked nested facade conversions disabled." +function _arrowtypeschildcolumn( + name::String, + v::AbstractVector; + extension_shape::Bool, + union_member::Bool=false, + context, +) + T = Base.nonmissingtype(eltype(v)) + union_member && + T === AbstractString && + return _constructnativepart(name, _missings_to(String, v); context) + if _arrowtypes_needs(v, context) || ( + T !== Any && ( + T <: NamedTuple || + T <: Tuple || + T <: AbstractDict || + (T <: AbstractVector && !(T <: AbstractString)) + ) + ) + return _constructarrowtypespart(name, v; extension_shape, context) + end + return extension_shape ? _constructnativepart(name, v; context) : + AC.fromjulia(name, _plainvector(v)) +end + +function _arrowtypeslistcolumn( + name::String, + v::AbstractVector, + S; + extension_shape::Bool, + context, +) + S <: AbstractVector || throw( + ArgumentError("ArrowTypes ListKind column $name must lower to AbstractVector rows"), + ) + E = eltype(S) + flat = E[] + present = Bool[x !== missing for x in v] + offsets = Vector{Int32}(undef, length(v) + 1) + offsets[1] = 0 + for (i, row) in enumerate(v) + if row !== missing + row isa AbstractVector || throw( + ArgumentError("ArrowTypes ListKind column $name contains a non-vector row"), + ) + length(flat) <= typemax(Int32) - length(row) || throw( + ArgumentError( + "ArrowTypes ListKind column $name exceeds the Int32 offset range", + ), + ) + append!(flat, row) + end + offsets[i + 1] = Int32(length(flat)) + end + childfield, childdata = _arrowtypeschildcolumn("item", flat; extension_shape, context) + t = AC.ListType(false) + field = AC.Field(name, t; nullable=Missing <: eltype(v), children=AC.Field[childfield]) + data = AC.ArrayData( + t, + length(v), + [AC._bitmapbuffer(present), AC._databuffer(offsets)]; + children=AC.ArrayData[childdata], + nullcount=count(!, present), + ) + return field, data +end + +struct _RetainedHidden end +const _RETAINED_HIDDEN = _RetainedHidden() + +abstract type _MaskedChildValues{T} <: AbstractVector{T} end + +""" +Lazy view of a fixed-size-list column's child slots; slots under a missing or +hidden row read as _RETAINED_HIDDEN. +""" +struct _MaskedFixedListValues{V<:AbstractVector} <: _MaskedChildValues{Any} + rows::V + width::Int + len::Int +end + +function _MaskedFixedListValues(rows::AbstractVector, width::Int) + width >= 0 || throw(ArgumentError("retained fixed-size-list width is negative")) + len64 = try + Base.checked_mul(Int64(length(rows)), Int64(width)) + catch err + err isa OverflowError || rethrow() + throw(ArgumentError("retained fixed-size-list child length is out of range")) + end + len64 <= typemax(Int) || + throw(ArgumentError("retained fixed-size-list child length is out of range")) + return _MaskedFixedListValues(rows, width, Int(len64)) +end + +Base.IndexStyle(::Type{<:_MaskedFixedListValues}) = IndexLinear() +Base.size(values::_MaskedFixedListValues) = (values.len,) +function Base.getindex(values::_MaskedFixedListValues, index::Int) + @boundscheck checkbounds(values, index) + rowoffset, childoffset = divrem(index - 1, values.width) + row = values.rows[firstindex(values.rows) + rowoffset] + row === missing && return _RETAINED_HIDDEN + row === _RETAINED_HIDDEN && return _RETAINED_HIDDEN + return row[firstindex(row) + childoffset] +end + +struct _FreshStructChildRows end +struct _RetainedStructChildRows end + +"Lazy values for one Struct child, including slots hidden by parent validity." +struct _MaskedStructChildValues{T,V<:AbstractVector,Mode} <: _MaskedChildValues{T} + rows::V + child::Int +end + +function _freshstructchildvalues(::Type{T}, rows::AbstractVector, child::Int) where {T} + return _MaskedStructChildValues{T,typeof(rows),_FreshStructChildRows}(rows, child) +end + +function _retainedstructchildvalues(rows::AbstractVector, child::Int) + return _MaskedStructChildValues{Any,typeof(rows),_RetainedStructChildRows}(rows, child) +end + +Base.IndexStyle(::Type{<:_MaskedStructChildValues}) = IndexLinear() +Base.size(values::_MaskedStructChildValues) = size(values.rows) + +function Base.getindex( + values::_MaskedStructChildValues{T,V,_FreshStructChildRows}, + index::Int, +) where {T,V} + @boundscheck checkbounds(values, index) + row = values.rows[firstindex(values.rows) + index - 1] + row === missing && return _RETAINED_HIDDEN + row === _RETAINED_HIDDEN && return _RETAINED_HIDDEN + return getfield(row, values.child) +end + +function Base.getindex( + values::_MaskedStructChildValues{T,V,_RetainedStructChildRows}, + index::Int, +) where {T,V} + @boundscheck checkbounds(values, index) + row = values.rows[firstindex(values.rows) + index - 1] + row === missing && return _RETAINED_HIDDEN + row === _RETAINED_HIDDEN && return _RETAINED_HIDDEN + return last(row[firstindex(row) + values.child - 1]) +end + +"A writer-side Union route retained until dense or sparse child construction." +struct _WriterRoutedUnion{T} + child::Int + value::T + writertype::Type +end +_WriterRoutedUnion(child::Int, value) = _WriterRoutedUnion(child, value, typeof(value)) + +struct _RegisteredWriterPlan + input::Type + routes::_WriterTypeRoutes +end + +struct _DeferredRegisteredWriter{T} + field::AC.Field + value::T + writetype::Type +end + +const _WriterUnionPlanCache = + IdDict{AC.Field,Dict{Tuple{Type,Bool},Tuple{Vector{Type},Vector{Int}}}} + +mutable struct _WriterContext + unions::_WriterUnionPlanCache + registered::IdDict{AC.Field,Dict{Tuple{Type,Type},_RegisteredWriterPlan}} + arrowtypes::_ArrowTypesContext + needs::IdDict{Type,Bool} + kinds::IdDict{Type,Any} + extensions::IdDict{Type,Any} + extensionshapes::IdDict{Type,Bool} + unionvariants::IdDict{Type,Vector{Type}} + unionbranches::Dict{Tuple{Type,Type},Int} + candidates::Dict{Tuple{Type,Bool},AC.Field} + inferred::Base.IdSet{Type} + physical::IdDict{AC.Field,Any} + schema::Vector{Type} + values::Base.IdSet{Any} + collecting::Bool + deferred::Bool + column::String +end + +function _WriterContext(column::AbstractString="") + arrowtypes = _ArrowTypesContext(; warn=false) + return _WriterContext( + _WriterUnionPlanCache(), + IdDict{AC.Field,Dict{Tuple{Type,Type},_RegisteredWriterPlan}}(), + arrowtypes, + IdDict{Type,Bool}(), + IdDict{Type,Any}(), + IdDict{Type,Any}(), + IdDict{Type,Bool}(), + IdDict{Type,Vector{Type}}(), + Dict{Tuple{Type,Type},Int}(), + Dict{Tuple{Type,Bool},AC.Field}(), + Base.IdSet{Type}(), + IdDict{AC.Field,Any}(), + Type[], + Base.IdSet{Any}(), + false, + false, + String(column), + ) +end + +_writerunionvariants!(context::_WriterContext, T::Type) = + get!(() -> _checkedwritervariants("writer type $T", T), context.unionvariants, T) + +function _writerunionbranch!( + context::_WriterContext, + uniontype::Type, + runtime::Type, + variants::Vector{Type}=_writerunionvariants!(context, uniontype), +) + return get!(context.unionbranches, (uniontype, runtime)) do + branch = findfirst(T -> runtime <: T, variants) + branch === nothing && throw( + ArgumentError("value type $runtime is not in declared writer type $uniontype"), + ) + branch + end +end + +const _MAX_WRITER_SCHEMA_DEPTH = 64 + +function _withwriterschema(f, context::_WriterContext, T::Type) + Base.@nospecialize T + if T in context.schema + path = join((string(S) for S in (context.schema..., T)), " -> ") + throw(ArgumentError("recursive ArrowTypes storage schema: $path")) + end + length(context.schema) < _MAX_WRITER_SCHEMA_DEPTH || throw( + ArgumentError( + "ArrowTypes storage schema exceeds the supported depth " * + "$_MAX_WRITER_SCHEMA_DEPTH at $T", + ), + ) + push!(context.schema, T) + try + return f() + finally + pop!(context.schema) + end +end + +function _writerstoragetype!(context::_WriterContext, T::Type) + Base.@nospecialize T + return _arrowtypesstoragetype!(context.arrowtypes, T) +end + +_writerneedstype!(context::_WriterContext, T::Type) = get!(context.needs, T) do + _arrowtypesneedstype( + T, + nested -> _writerstoragetype!(context, nested), + nested -> _writerextension!(context, nested) !== nothing, + 0, + ) +end + +function _writerkind!(context::_WriterContext, T::Type) + Base.@nospecialize T + return get!(() -> ArrowTypes.ArrowKind(T), context.kinds, T) +end + +function _writerextension!(context::_WriterContext, T::Type) + Base.@nospecialize T + return get!(context.extensions, T) do + ArrowTypes.hasarrowname(T) || return nothing + Pair{String,String}[ + _EXTENSION_NAME_KEY => String(ArrowTypes.arrowname(T)), + _EXTENSION_METADATA_KEY => String(ArrowTypes.arrowmetadata(T)), + ] + end +end + +function _writerextensionshape(context::_WriterContext, T::Type, storage) + Base.@nospecialize T + return get!(context.extensionshapes, T) do + kind = _writerkind!(context, T) + storage !== T || + _writerextension!(context, T) !== nothing || + (kind isa ArrowTypes.StructKind && !(T <: Union{NamedTuple,Tuple})) + end +end + +function _writercandidatefield!( + context::_WriterContext, + T::Type; + extension_shape::Bool=false, +) + Base.@nospecialize T + key = (T, extension_shape) + haskey(context.candidates, key) && return context.candidates[key] + return _withwriterschema(context, T) do + get!(context.candidates, key) do + if T === VersionNumber + field, _ = _constructnativepart("", String[]) + return _arrowtypeslogicalfield(context, field, T) + end + storage = _writerstoragetype!(context, T) + ownshape = _writerextensionshape(context, T, storage) + nestedshape = extension_shape || ownshape + if storage isa Union + variants = _checkedwritervariants("writer candidate type $T", storage) + children = AC.Field[] + for (i, variant) in enumerate(variants) + child = _writercandidatefield!( + context, + variant; + extension_shape=nestedshape, + ) + push!( + children, + AC.Field( + string(i - 1), + child.type; + nullable=child.nullable, + metadata=_fieldmetadata(child), + children=collect(AC.Field, child.children), + ), + ) + end + ids = Int8.(0:(length(children) - 1)) + uniontype = AC.UnionType(AC.DenseMode, ids) + field = AC.Field("", uniontype; nullable=Missing in variants, children) + return _arrowtypeslogicalfield(context, field, T; nullable=field.nullable) + end + field, _ = _arrowtypesstoragecolumn( + "", + storage[]; + extension_shape=nestedshape, + context, + ) + # Null-storage nullability rule: see _constructarrowtypespart_impl. + nullable = field.type isa AC.NullType && T !== Missing ? false : field.nullable + return _arrowtypeslogicalfield(context, field, T; nullable) + end + end +end + +"Give an inferred writer candidate its column name without rebuilding its shape." +function _namedwriterfield(name::String, candidate::AC.Field) + return AC.Field( + name, + candidate.type; + nullable=candidate.nullable, + metadata=_fieldmetadata(candidate), + children=collect(AC.Field, candidate.children), + ) +end + +_placeholderadd(a::Int, b::Int) = a > typemax(Int) - b ? typemax(Int) : a + b +_placeholdermul(a::Int, b::Int) = + a == 0 || b == 0 ? 0 : a > typemax(Int) ÷ b ? typemax(Int) : a * b + +# Two keywords control hidden-slot synthesis, and every function below shares +# them: `forcevalid` means an ancestor requires a NON-NULL value in this slot +# (a dense Union child, an REE value run); `inactive` means the whole subtree +# is unreachable, so validity is free and nothing must be synthesized valid. +# Hence needvalid = !inactive && (forcevalid || !f.nullable). +function _writercansynthesize(f::AC.Field; forcevalid::Bool=false, inactive::Bool=false) + needvalid = !inactive && (forcevalid || !f.nullable) + t = f.type + t isa AC.NullType && return !needvalid + if t isa AC.RunEndEncodedType + return length(f.children) == 2 && + _writercansynthesize(f.children[2]; forcevalid=needvalid, inactive) + elseif t isa AC.UnionType + return any( + child -> _writercansynthesize(child; forcevalid=needvalid, inactive), + f.children, + ) + end + # Other composite layouts remain non-null when their children are null. + # Dictionary construction likewise makes a valid index; its pool value is + # not the dictionary array's validity state. + return true +end + +function _writerplaceholdercost(f::AC.Field; forcevalid::Bool=false, inactive::Bool=false) + _writercansynthesize(f; forcevalid, inactive) || return typemax(Int) + needvalid = !inactive && (forcevalid || !f.nullable) + t = f.type + t isa AC.NullType && return 0 + t isa AC.BoolType && return 1 + t isa Union{ + AC.IntType, + AC.FloatType, + AC.DateType, + AC.TimeType, + AC.TimestampType, + AC.DurationType, + AC.FixedSizeBinaryType, + AC.DecimalType, + AC.IntervalType, + } && return AC.primwidth(t) + t isa Union{AC.Utf8Type,AC.BinaryType,AC.ListType} && return t.large ? 8 : 4 + t isa AC.MapType && return 4 + t isa AC.ViewType && return 16 + t isa AC.ListViewType && return t.large ? 16 : 8 + if t isa AC.FixedSizeListType + return _placeholdermul( + t.listsize, + _writerplaceholdercost(only(f.children); inactive), + ) + elseif t isa AC.StructType + return foldl( + (cost, child) -> _placeholderadd(cost, _writerplaceholdercost(child; inactive)), + f.children; + init=0, + ) + elseif t isa AC.RunEndEncodedType + return _placeholderadd( + 4, + _writerplaceholdercost(f.children[2]; forcevalid=needvalid, inactive), + ) + elseif t isa AC.UnionType + return minimum( + child -> _writerplaceholdercost(child; forcevalid=needvalid, inactive), + f.children; + init=typemax(Int), + ) + end + return typemax(Int) +end + +function _writerunionvariants(context::_WriterContext, f::AC.Field, writetype::Type) + f.type isa AC.UnionType || throw(ArgumentError("Field $(f.name) is not a Union")) + writetype isa Union || throw( + ArgumentError( + "writer type $writetype does not describe retained Union field $(f.name)", + ), + ) + variants = _writerunionvariants!(context, writetype) + length(variants) == length(f.children) || throw( + ArgumentError( + "writer type $writetype has $(length(variants)) Union branches, but " * + "retained field $(f.name) has $(length(f.children)) children", + ), + ) + return variants +end + +function _writerphysicalbasetype(context::_WriterContext, f::AC.Field) + return get!(context.physical, f) do + t = f.type + if t isa AC.DictionaryType + return _writerphysicalbasetype(context, AC.dictvaluefield(f, t)) + elseif t isa AC.RunEndEncodedType + return length(f.children) == 2 ? + _writerphysicalbasetype(context, f.children[2]) : Any + elseif t isa Union{AC.ListType,AC.ListViewType} + return length(f.children) == 1 ? + Vector{_writerphysicalbasetype(context, f.children[1])} : Vector{Any} + elseif t isa AC.FixedSizeListType + return length(f.children) == 1 ? + _boundedfixedliststoragetype( + t.listsize, + _writerphysicalbasetype(context, f.children[1]), + ) : NTuple{0,Any} + elseif t isa AC.StructType + names = _arrowtypesstructnames(f) + types = Type[_writerphysicalbasetype(context, child) for child in f.children] + return Core.apply_type(NamedTuple, names, Core.apply_type(Tuple, types...)) + elseif t isa AC.MapType + length(f.children) == 1 || return Dict{Any,Any} + entries = f.children[1] + length(entries.children) == 2 || return Dict{Any,Any} + return Dict{ + _writerphysicalbasetype(context, entries.children[1]), + _writerphysicalbasetype(context, entries.children[2]), + } + elseif t isa AC.UnionType + isempty(f.children) && return Union{} + T = _writerphysicalbasetype(context, f.children[1]) + for child in Iterators.drop(f.children, 1) + T = Union{T,_writerphysicalbasetype(context, child)} + end + return T + end + return _arrowtypesprimitivebasetype(f) + end +end + +_writerdatelike(t::AC.ArrowType) = + t isa AC.DateType || (t isa AC.TimestampType && t.unit in (AC.SECOND, AC.MILLISECOND)) + +_writersequence(t::AC.ArrowType) = + t isa Union{AC.ListType,AC.ListViewType,AC.FixedSizeListType} + +_writerplainstorageshape(f::AC.Field) = + !f.nullable && + _arrowtypesextension(f) === nothing && + all(_writerplainstorageshape, f.children) + +function _writeropaquestoragealias(retained::AC.Field, candidate::AC.Field) + isempty(retained.children) && !isempty(candidate.children) || return false + t = retained.type + supported = + t isa AC.BinaryType || + t isa AC.FixedSizeBinaryType || + (t isa AC.ViewType && !t.utf8) || + (t isa AC.DecimalType && t.bits > 64) || + (t isa AC.IntervalType && t.unit != AC.YEAR_MONTH) + supported || return false + return all(_writerplainstorageshape, candidate.children) +end + +"Physical descriptor distance; exact layouts win over compatible storage domains." +function _writerdescriptordistance( + context::_WriterContext, + retained::AC.Field, + candidate::AC.Field, +) + AC.typeequal(retained.type, candidate.type) && return 0 + if retained.type isa AC.UnionType && candidate.type isa AC.UnionType + return retained.type.mode == candidate.type.mode ? 1 : 3 + end + if retained.type isa AC.MapType || candidate.type isa AC.MapType + retained.type isa AC.MapType && candidate.type isa AC.MapType || return nothing + retained.type.keyssorted == candidate.type.keyssorted || return nothing + end + if retained.type isa AC.DurationType && candidate.type isa AC.DurationType + return 1 + end + if _writerdatelike(retained.type) && _writerdatelike(candidate.type) + return typeof(retained.type) === typeof(candidate.type) ? 1 : 3 + end + if _writersequence(retained.type) && _writersequence(candidate.type) + if retained.type isa AC.FixedSizeListType && + candidate.type isa AC.FixedSizeListType && + retained.type.listsize != candidate.type.listsize + return nothing + end + return typeof(retained.type) === typeof(candidate.type) ? 1 : 3 + end + retainedstorage = _writerphysicalbasetype(context, retained) + candidatestorage = _writerphysicalbasetype(context, candidate) + retainedstorage === Any && return nothing + candidatestorage === Any && return nothing + retainedstorage === candidatestorage || return nothing + return typeof(retained.type) === typeof(candidate.type) ? 1 : 4 +end + +"Physical writer identity distance for a retained logical branch." +function _writerfielddistance( + context::_WriterContext, + retained::AC.Field, + candidate::AC.Field; + checkname::Bool=true, + positionalstruct::Bool=false, +) + checkname && retained.name != candidate.name && return nothing + candidate.nullable && !retained.nullable && return nothing + _arrowtypesextension(retained) == _arrowtypesextension(candidate) || return nothing + if retained.type isa AC.RunEndEncodedType && !(candidate.type isa AC.RunEndEncodedType) + length(retained.children) == 2 || return nothing + childdistance = _writerfielddistance( + context, + retained.children[2], + _withoutownextension(candidate); + checkname=false, + positionalstruct, + ) + childdistance === nothing && return nothing + return 4 + childdistance + (retained.nullable == candidate.nullable ? 0 : 1) + end + typedistance = _writerdescriptordistance(context, retained, candidate) + typedistance === nothing && return nothing + if _writeropaquestoragealias(retained, candidate) + return typedistance + (retained.nullable == candidate.nullable ? 0 : 1) + end + length(retained.children) == length(candidate.children) || return nothing + distance = typedistance + (retained.nullable == candidate.nullable ? 0 : 1) + if retained.type isa AC.UnionType + matches = Vector{Pair{Int,Int}}(undef, length(candidate.children)) + for (candidateindex, candidatechild) in enumerate(candidate.children) + choices = Pair{Int,Int}[] + for (retainedindex, retainedchild) in enumerate(retained.children) + childdistance = _writerfielddistance( + context, + retainedchild, + candidatechild; + checkname=false, + positionalstruct=false, + ) + childdistance === nothing || push!(choices, childdistance => retainedindex) + end + isempty(choices) && return nothing + sort!(choices; by=choice -> (first(choice), last(choice))) + length(choices) > 1 && + first(choices[1]) == first(choices[2]) && + throw( + ArgumentError( + "writer Field has multiple equally exact children in retained " * + "field $(retained.name); logical routing would be ambiguous", + ), + ) + matches[candidateindex] = first(choices) + end + length(unique(last(match) for match in matches)) == length(matches) || + return nothing + return distance + sum(first, matches) + end + for i in eachindex(retained.children) + childdistance = _writerfielddistance( + context, + retained.children[i], + candidate.children[i]; + checkname=retained.type isa AC.StructType && !positionalstruct, + positionalstruct=retained.type isa AC.MapType, + ) + childdistance === nothing && return nothing + distance += childdistance + end + return distance +end + +function _writerchildcandidate(context::_WriterContext, f::AC.Field, candidate::AC.Field) + matches = Pair{Int,Int}[] + for (i, child) in enumerate(f.children) + distance = _writerfielddistance(context, child, candidate; checkname=false) + distance === nothing || push!(matches, distance => i) + end + isempty(matches) && return nothing + sort!(matches; by=pair -> (first(pair), last(pair))) + length(matches) > 1 && + first(matches[1]) == first(matches[2]) && + throw( + ArgumentError( + "writer Field has multiple equally exact children in retained field " * + "$(f.name); logical routing would be ambiguous", + ), + ) + return last(first(matches)) +end + +"Match one physical storage branch beneath a retained logical Union." +function _writerunionstoragechildcandidate( + context::_WriterContext, + f::AC.Field, + candidate::AC.Field, +) + child = _writerchildcandidate(context, f, candidate) + child === nothing || return child + + # An inferred dictionary can preserve the abstract logical label on each + # observed subtype branch. Those repeated labels describe the parent + # public-domain value, not distinct physical branches. Remove only labels that + # equal the parent label before matching the parent's storage Union. + parentextension = _arrowtypesextension(f) + parentextension === nothing && return nothing + storagechildren = AC.Field[ + _arrowtypesextension(branch) == parentextension ? _withoutownextension(branch) : + branch for branch in f.children + ] + return _writerchildcandidate( + context, + AC.Field( + f.name, + f.type; + nullable=f.nullable, + metadata=_fieldmetadata(f), + children=storagechildren, + ), + candidate, + ) +end + +function _writerunionplan!( + context::_WriterContext, + f::AC.Field, + storagetype::Type; + extension_shape::Bool=false, +) + bytype = get!(context.unions, f) do + Dict{Tuple{Type,Bool},Tuple{Vector{Type},Vector{Int}}}() + end + return get!(bytype, (storagetype, extension_shape)) do + variants = _writerunionvariants(context, f, storagetype) + children = Int[] + owned = Set{Int}() + for variant in variants + candidate = _writercandidatefield!(context, variant; extension_shape) + child = _writerunionstoragechildcandidate(context, f, candidate) + child === nothing && throw( + ArgumentError( + "storage Union branch $variant matches no child of retained " * + "field $(f.name)", + ), + ) + child in owned && throw( + ArgumentError( + "storage Union $storagetype maps more than one branch to child " * + "$child of retained field $(f.name)", + ), + ) + push!(owned, child) + push!(children, child) + end + # A branch must identify one retained child uniquely. Exact layouts + # win, but the retained descriptor may select a compatible storage form. + return variants, children + end +end + +_fieldcontainsunion(f::AC.Field) = + f.type isa AC.UnionType || any(_fieldcontainsunion, f.children) + +function _registeredoutermissing(context::_WriterContext, f::AC.Field, target, runtimes) + f.type isa AC.UnionType || return f.nullable + storage = _writerstoragetype!(context, target) + storage isa Union && return _arrowtypeslogicalnullable(f, target, storage) + + nullable = nothing + for T in runtimes + runtime_storage = _writerstoragetype!(context, T) + runtime_storage isa Union || continue + storagebranches = length(Base.uniontypes(runtime_storage)) + fieldbranches = length(f.children) + candidate = if fieldbranches == storagebranches + false + elseif fieldbranches == storagebranches + 1 + Missing <: runtime_storage && throw( + ArgumentError( + "registered writer type $T uses Missing as a storage Union " * + "branch, so retained field $(f.name) cannot add an " * + "indistinguishable outer missing branch", + ), + ) + true + else + throw( + ArgumentError( + "registered writer type $T has $storagebranches storage Union " * + "branches, but retained field $(f.name) has $fieldbranches children", + ), + ) + end + nullable === nothing || + nullable == candidate || + throw( + ArgumentError( + "runtime writer types disagree about outer nullability for retained " * + "field $(f.name)", + ), + ) + nullable = candidate + end + nullable === nothing || return nullable + nullchildren = findall( + child -> child.type isa AC.NullType && _arrowtypesextension(child) === nothing, + f.children, + ) + return length(nullchildren) == 1 +end + +function _writercheckabstractextensionidentity( + f::AC.Field, + target, + T::Type, + candidateextension, +) + isconcretetype(target) && return nothing + retainedextension = _arrowtypesextension(f) + candidateextension === nothing || + candidateextension == retainedextension || + throw( + ArgumentError( + "registered writer subtype $T declares extension " * + "$(repr(candidateextension)), which cannot replace abstract " * + "retained target $target with extension " * + "$(repr(retainedextension))", + ), + ) + return nothing +end + +_writercheckabstractextension(f::AC.Field, target, T::Type, candidate::AC.Field) = + _writercheckabstractextensionidentity(f, target, T, _arrowtypesextension(candidate)) + +function _registeredwriterroutes( + context::_WriterContext, + f::AC.Field, + runtimes, + outermissing::Bool, +) + routes = _WriterTypeRoutes() + _, target, _ = _arrowtypestarget(context.arrowtypes, f) + target === nothing && + throw(ArgumentError("retained field $(f.name) has no registered writer target")) + if !(f.type isa AC.UnionType) + for T in runtimes + candidate = _writercandidatefield!(context, T) + _writercheckabstractextension(f, target, T, candidate) + candidate = _withoutownextension(candidate) + retained = _withoutownextension(f) + _writerfielddistance(context, retained, candidate; checkname=false) === + nothing && throw( + ArgumentError( + "registered writer type $T does not match retained field " * + "$(f.name)", + ), + ) + routes[T] = 0 + end + outermissing && (routes[Missing] = 0) + return routes + end + + childowners = Dict{Int,Type}() + for T in runtimes + candidate = _writercandidatefield!(context, T) + _writercheckabstractextension(f, target, T, candidate) + storage = _writerstoragetype!(context, T) + if storage isa Union + plannedstorage = outermissing ? Union{Missing,storage} : storage + _writerunionplan!( + context, + f, + plannedstorage; + extension_shape=_writerextensionshape(context, T, storage), + ) + routes[T] = 0 + continue + end + child = _writerchildcandidate(context, f, candidate) + if child === nothing && _arrowtypesextension(candidate) !== nothing + # A direct registered Union keeps its logical label only on the + # parent. A dictionary can retain a concrete label on its value + # children. Prefer the exact child contract, then retry after + # removing the runtime subtype's outer label. + candidateextension = _arrowtypesextension(candidate) + if isconcretetype(target) || candidateextension == _arrowtypesextension(f) + child = _writerchildcandidate(context, f, _withoutownextension(candidate)) + end + end + child === nothing && throw( + ArgumentError( + "registered writer type $T matches no child of retained field $(f.name)", + ), + ) + owner = get(childowners, child, nothing) + owner === nothing || throw( + ArgumentError( + "registered writer types $owner and $T both map to child $child " * + "of retained field $(f.name)", + ), + ) + childowners[child] = T + routes[T] = child + end + if outermissing + nullchildren = findall( + child -> + child.type isa AC.NullType && _arrowtypesextension(child) === nothing, + f.children, + ) + length(nullchildren) == 1 || throw( + ArgumentError( + "retained registered Union field $(f.name) has no unique outer " * + "missing child", + ), + ) + routes[Missing] = only(nullchildren) + end + return routes +end + +"Append `T` once, under `limit`, or throw naming `owner` and `kind`." +function _writerpushtype!( + types::Vector{Type}, + seen::Base.IdSet{Type}, + T::Type, + owner::String, + limit::Int, + kind::String, +) + T in seen && return nothing + length(types) < limit || throw(ArgumentError("$owner has more than $limit $kind")) + push!(seen, T) + push!(types, T) + return nothing +end + +function _writerpushinferred!( + context::_WriterContext, + f::AC.Field, + types::Vector{Type}, + seen::Base.IdSet{Type}, + T::Type, +) + _writerregisterinferred!(context, f, T) + return _writerpushtype!( + types, + seen, + T, + "column $(f.name)", + _MAX_INFERRED_WRITER_TYPES, + "inferred registered writer runtime types", + ) +end + +function _writerregisterinferred!(context::_WriterContext, f::AC.Field, T::Type) + return _writerregisterinferredtype!( + context, + f.name, + T, + "inferred registered writer runtime types", + ) +end + +function _writerregisterinferredtype!( + context::_WriterContext, + name::String, + T::Type, + kind::String, +) + if !(T in context.inferred) + length(context.inferred) < _MAX_INFERRED_WRITER_TYPES || throw( + ArgumentError("column $name has more than $_MAX_INFERRED_WRITER_TYPES $kind"), + ) + push!(context.inferred, T) + end + return nothing +end + +function _writerregisterstorageinference!(context::_WriterContext, name::String, T::Type) + return _writerregisterinferredtype!( + context, + name, + T, + "inferred ArrowTypes storage types", + ) +end + +""" +Narrow one unresolved abstract declaration from whole-column runtime evidence. +Returns nothing when the declaration cannot be narrowed from runtime evidence. +""" +function _narrowabstractwriterparts(name::String, parts, context::_WriterContext) + logical = nothing + for part in parts + declared = Base.nonmissingtype(eltype(part)) + declared in (Union{}, Any) && continue + if !isconcretetype(declared) && !(declared isa Union) + logical === nothing ? (logical = declared) : + logical === declared || return nothing + end + end + logical === nothing && return nothing + _writerstoragetype!(context, logical) === logical || return nothing + _writerextension!(context, logical) === nothing || return nothing + + observed = Type[] + seen = Base.IdSet{Type}() + needsruntime = false + nullable = false + for part in parts + declared = Base.nonmissingtype(eltype(part)) + nullable |= Missing <: eltype(part) + declared in (Union{}, logical) || + (isconcretetype(declared) && declared <: logical) || + return nothing + for value in part + value === missing && continue + value isa logical || throw( + ArgumentError( + "column $name contains $(typeof(value)), which is outside its " * + "declared abstract element type $logical", + ), + ) + runtime = typeof(value) + runtime in seen && continue + push!(seen, runtime) + push!(observed, runtime) + needsruntime |= + _writerneedstype!(context, runtime) || + _writerextension!(context, runtime) !== nothing + end + end + isempty(observed) && return nothing + needsruntime || return nothing + for runtime in observed + _writerregisterinferredtype!( + context, + name, + runtime, + "inferred registered writer runtime types", + ) + end + storagetype = _writeruniontype(observed) + nullable && (storagetype = Union{Missing,storagetype}) + return AbstractVector[collect(storagetype, part) for part in parts] +end + +"One whole-column plan for an ArrowTypes mapping whose storage needs values." +struct _InferredStoragePlan + logical::Type + declared::Type + storagetype::Type + parts::Vector{AbstractVector} + nullable::Bool +end + +function _inferredstoragelogical(parts, target) + target === nothing || return target + logical = nothing + for part in parts + T = Base.nonmissingtype(eltype(part)) + T === Union{} && continue + isconcretetype(T) || return nothing + logical === nothing ? (logical = T) : logical === T || return nothing + end + return logical +end + +function _inferredstoragedeclaration(parts, context::_WriterContext; target=nothing) + logical = _inferredstoragelogical(parts, target) + logical === nothing && return nothing + logical === VersionNumber && return nothing + for part in parts + T = Base.nonmissingtype(eltype(part)) + T in (Union{}, logical) && continue + target !== nothing && + T === Any && + all(value -> value === missing || value isa logical, part) && + continue + return nothing + end + _writerneedstype!(context, logical) || return nothing + declared = _writerstoragetype!(context, logical) + _arrowtypesstorageisspecified(context, declared) && return nothing + return logical => declared +end + +""" +One whole-column storage plan for an ArrowTypes mapping with abstract storage, +or nothing when no plan applies. +""" +function _inferredstorageplan( + name::String, + parts, + context::_WriterContext; + target=nothing, + retained=nothing, +) + declaration = _inferredstoragedeclaration(parts, context; target) + declaration === nothing && return nothing + logical, declared = declaration + + lowered = Vector{Vector{Any}}(undef, length(parts)) + observed = Type[] + seen = Base.IdSet{Type}() + sawmissing = false + nullable = false + logicalseen = Base.IdSet{Type}() + for (partition, part) in enumerate(parts) + nullable |= Missing <: eltype(part) + values = Any[] + sizehint!(values, length(part)) + for value in part + if value === missing + push!(values, missing) + sawmissing = true + continue + end + mapped = _writertoarrow(context, value, logical, name) + mapped = _writerconvertlowered(declared, mapped, context, logical, name) + push!(values, mapped) + if retained !== nothing && !isconcretetype(logical) + runtime = typeof(value) + if !(runtime in logicalseen) + _writerregisterinferred!(context, retained, runtime) + _writercheckabstractextensionidentity( + retained, + logical, + runtime, + _writerextension!(context, runtime), + ) + push!(logicalseen, runtime) + end + end + runtime = typeof(mapped) + runtime in seen && continue + _writerregisterstorageinference!(context, name, runtime) + push!(seen, runtime) + push!(observed, runtime) + end + lowered[partition] = values + end + + if isempty(observed) + mappedparts = AbstractVector[fill(missing, length(values)) for values in lowered] + return _InferredStoragePlan(logical, declared, Missing, mappedparts, nullable) + end + nullable && + all(T -> _writerkind!(context, T) isa ArrowTypes.NullKind, observed) && + throw( + ArgumentError( + "nullable logical column $name maps its non-missing values to Arrow " * + "Null storage, so outer missing values cannot be distinguished; " * + "use DictEncode to preserve the two states", + ), + ) + storagetype = reduce(_writerpromoteunion, observed) + sawmissing && (storagetype = Union{Missing,storagetype}) + mappedparts = AbstractVector[ + _writerconvertedcolumn(storagetype, values, context, logical, name) for + values in lowered + ] + return _InferredStoragePlan(logical, declared, storagetype, mappedparts, nullable) +end + +function _constructinferredfreshpart( + name::String, + plan::_InferredStoragePlan, + values::AbstractVector, + context::_WriterContext, +) + ownshape = _writerextensionshape(context, plan.logical, plan.declared) + field, data = _arrowtypesstoragecolumn(name, values; extension_shape=ownshape, context) + # Null-storage nullability rule: see _constructarrowtypespart_impl. + nullable = field.type isa AC.NullType && !plan.nullable ? false : field.nullable + return _arrowtypeslogicalfield(context, field, plan.logical; nullable), data +end + +function _inferredstorageroutes( + context::_WriterContext, + field::AC.Field, + storagetype::Type; + extension_shape::Bool, +) + field.type isa AC.UnionType || return nothing + routes = _WriterTypeRoutes() + owners = Set{Int}() + for variant in _writertypevariants(storagetype) + candidate = _writercandidatefield!(context, variant; extension_shape) + child = _writerchildcandidate(context, field, candidate) + child === nothing && throw( + ArgumentError( + "inferred storage type $variant matches no child of retained " * + "field $(field.name)", + ), + ) + child in owners && throw( + ArgumentError( + "inferred storage types map ambiguously to child $child of " * + "retained field $(field.name)", + ), + ) + push!(owners, child) + routes[variant] = child + end + return routes +end + +function _constructinferredretainedpart( + field::AC.Field, + plan::_InferredStoragePlan, + values::AbstractVector, + context::_WriterContext, +) + isempty(values) && return _constructhiddenpart(field, 0, context) + ownshape = _writerextensionshape(context, plan.logical, plan.declared) + routes = + _inferredstorageroutes(context, field, plan.storagetype; extension_shape=ownshape) + return _constructwriterchild(field, plan.storagetype, values; routes, context) +end + +_writerfullydeclared(T::Type) = + isconcretetype(T) || (T isa Union && all(isconcretetype, _writertypevariants(T))) + +function _registeredwritercolumnplan(context::_WriterContext, f::AC.Field, valueparts) + _, target, _ = _arrowtypestarget(context.arrowtypes, f) + target === nothing && return nothing + declaredtypes = Type[] + declaredseen = Base.IdSet{Type}() + inferredtypes = Type[] + inferredseen = Base.IdSet{Type}() + schemaevidence = Type[] + schemaseen = Base.IdSet{Type}() + sawmissing = false + for values in valueparts + declared = Base.nonmissingtype(eltype(values)) + declared in (Union{}, Any) || + declared <: target || + throw( + ArgumentError( + "column $(f.name) has declared non-missing type $declared, " * + "but its retained ArrowTypes target is $target", + ), + ) + declaredvariants = _writertypevariants(declared) + fullydeclared = _writerfullydeclared(declared) + for value in values + value === _RETAINED_HIDDEN && continue + if value === missing + sawmissing = true + elseif value isa target + fullydeclared || _writerpushinferred!( + context, + f, + inferredtypes, + inferredseen, + typeof(value), + ) + else + throw( + ArgumentError( + "column $(f.name) holds $(typeof(value)) values, but its " * + "retained ArrowTypes target is $target", + ), + ) + end + end + if isconcretetype(declared) + _writerpushtype!( + declaredtypes, + declaredseen, + declared, + "column $(f.name)", + _MAX_WRITER_UNION_BRANCHES, + "declared registered writer types", + ) + elseif declared isa Union + for member in declaredvariants + isconcretetype(member) || continue + _writerpushtype!( + declaredtypes, + declaredseen, + member, + "column $(f.name)", + _MAX_WRITER_UNION_BRANCHES, + "declared registered writer types", + ) + end + elseif declared !== Union{} && _writerstoragetype!(context, declared) !== declared + # An abstract declaration with an explicit storage mapping is + # schema evidence even when no runtime value is visible. + _writerpushtype!( + schemaevidence, + schemaseen, + declared, + "column $(f.name)", + _MAX_WRITER_UNION_BRANCHES, + "registered writer schema-evidence types", + ) + end + end + plantypes = _mergewritertypes( + "column $(f.name)", + declaredtypes, + inferredtypes, + _MAX_WRITER_UNION_BRANCHES, + "registered writer planning types", + ) + plantypes = _mergewritertypes( + "column $(f.name)", + plantypes, + schemaevidence, + _MAX_WRITER_UNION_BRANCHES, + "registered writer planning types", + ) + outermissing = _registeredoutermissing(context, f, target, plantypes) + sawmissing && + !outermissing && + throw( + ArgumentError( + "column $(f.name) contains an outer missing value that its registered " * + "ArrowTypes field cannot represent", + ), + ) + input = outermissing ? Union{Missing,target} : target + return _RegisteredWriterPlan( + input, + _registeredwriterroutes(context, f, plantypes, outermissing), + ) +end + +_registeredwriterplan(context::_WriterContext, f::AC.Field, values::AbstractVector) = + _registeredwritercolumnplan(context, f, (values,)) + +"A storage-domain placeholder value tree for one hidden slot under a retained Field." +function _writerhiddenstorage(f::AC.Field; forcevalid::Bool=false, inactive::Bool=false) + t = f.type + if t isa AC.UnionType + needvalid = !inactive && (forcevalid || !f.nullable) + child = _writerhiddenunionchild(f; forcevalid=needvalid, inactive) + return _WriterRoutedUnion( + child, + _writerhiddenstorage(f.children[child]; forcevalid=needvalid, inactive), + ) + elseif t isa Union{AC.ListType,AC.ListViewType} + return Any[] + elseif t isa AC.FixedSizeListType + child = only(f.children) + return Any[_writerhiddenstorage(child; inactive) for _ = 1:(t.listsize)] + elseif t isa AC.StructType + return Pair{String,Any}[ + child.name => _writerhiddenstorage(child; inactive) for child in f.children + ] + elseif t isa AC.MapType + return Pair{Any,Any}[] + elseif t isa AC.RunEndEncodedType + needvalid = !inactive && (forcevalid || !f.nullable) + return _writerhiddenstorage(f.children[2]; forcevalid=needvalid, inactive) + elseif t isa AC.DictionaryType + # A dictionary index must point at something even in an inactive + # subtree, so a non-nullable dictionary always needs one pool entry + # (`_constructhiddenpart`'s dictionary branch shares this rule). In + # the value domain that entry is the VALUE field's placeholder, and + # it must be non-missing: value-level construction reads a missing + # value back as a null INDEX, which a needvalid slot cannot have. + needvalid = (!inactive && forcevalid) || !f.nullable + needvalid || return missing + return _writerhiddenstorage(AC.dictvaluefield(f, t); forcevalid=true) + end + return _retainedplaceholder(f; forcevalid=(!inactive && forcevalid)) +end + +function _writerhiddenunionchild( + f::AC.Field; + forcevalid::Bool=(!f.nullable), + inactive::Bool=false, +) + candidates = Int[ + i for i in eachindex(f.children) if + _writercansynthesize(f.children[i]; forcevalid, inactive) + ] + isempty(candidates) && throw( + ArgumentError( + "retained Union field $(f.name) has no child that can synthesize a " * + "$(forcevalid ? "non-null hidden" : "hidden") value", + ), + ) + _, position = findmin( + i -> _writerplaceholdercost(f.children[i]; forcevalid, inactive), + candidates, + ) + return candidates[position] +end + +# A Null-only child has exactly one public-domain value, so hidden and missing rows +# only contribute child LENGTH. Visible rows are still validated for exact +# width; no placeholder objects are allocated. +function _writercompactfixednull(f::AC.Field) + f.type isa AC.FixedSizeListType || return false + length(f.children) == 1 || return false + child = only(f.children) + # Recursive compaction can preserve nested validity without invoking any + # logical hooks only for an extension-free shape. A labelled immediate + # Null child is also safe because it has no descendant validity topology. + return _writerplainfixednullshape(f) || child.type isa AC.NullType +end + +function _writerplainfixednullshape(f::AC.Field) + _arrowtypesextension(f) === nothing || return false + f.type isa AC.NullType && return true + f.type isa AC.FixedSizeListType || return false + length(f.children) == 1 || return false + return _writerplainfixednullshape(only(f.children)) +end + +function _checkplainfixednullvalue(f::AC.Field, value) + if f.type isa AC.NullType + value === missing || throw( + ArgumentError( + "retained Null field $(f.name) received a non-null public-domain value", + ), + ) + # NullType is null by definition. Its Field.nullable flag is advisory + # metadata and cannot make a physical Null slot non-null. + return nothing + end + if value === missing + f.nullable || throw( + ArgumentError( + "retained field $(f.name) is non-nullable but received a null value", + ), + ) + return nothing + end + t = f.type::AC.FixedSizeListType + (value isa Tuple || value isa AbstractVector) || throw( + ArgumentError("retained fixed-size-list field $(f.name) requires a sequence value"), + ) + length(value) == t.listsize || throw( + ArgumentError( + "retained fixed-size-list field $(f.name) received $(length(value)) " * + "values; expected $(t.listsize)", + ), + ) + child = only(f.children) + for item in value + _checkplainfixednullvalue(child, item) + end + return nothing +end + +function _writerhiddenvalue(f::AC.Field, writetype::Type) + f.nullable && Missing <: writetype && !_fieldcontainsunion(f) && return missing + # Hidden parent slots have no logical Union branch. The retained Field is + # complete physical authority, so choose cheap storage placeholders without + # requiring a constructor for the logical writer type. + return _writerhiddenstorage(f; inactive=true) +end + +function _writerhiddenbytes(f::AC.Field, nbytes::Integer, what::String) + 0 <= nbytes <= typemax(Int) || throw( + ArgumentError( + "hidden $what for retained field $(f.name) needs an unsupported " * + "$nbytes-byte buffer", + ), + ) + nbytes == 0 && return AC.BufferSlice() + return AC._databuffer(zeros(UInt8, Int(nbytes))) +end + +function _writerhiddenmul(f::AC.Field, left::Integer, right::Integer, what::String) + try + return Base.checked_mul(Int64(left), Int64(right)) + catch err + err isa OverflowError || rethrow() + throw( + ArgumentError( + "hidden $what for retained field $(f.name) exceeds the supported " * + "length range", + ), + ) + end +end + +function _writerhiddenbuffers(f::AC.Field, n::Int; forcevalid::Bool=false) + spec = AC.layoutspec_of(f.type) + buffers = AC.BufferSlice[] + makenull = f.nullable && !forcevalid + nullcount = makenull ? n : 0 + for role in spec.buffers + if role == AC.VALIDITY + nbytes = makenull ? cld(n, 8) : 0 + push!(buffers, _writerhiddenbytes(f, nbytes, "validity")) + elseif role == AC.DATA + nbytes = + spec.fixedwidth == -1 ? cld(n, 8) : + spec.fixedwidth == 0 ? 0 : _writerhiddenmul(f, n, spec.fixedwidth, "data") + push!(buffers, _writerhiddenbytes(f, nbytes, "data")) + elseif role == AC.OFFSETS + nslots = Base.checked_add(Int64(n), Int64(1)) + nbytes = _writerhiddenmul(f, nslots, spec.offsetwidth, "offset") + push!(buffers, _writerhiddenbytes(f, nbytes, "offset")) + elseif role == AC.ELEMENT_OFFSETS || role == AC.SIZES + nbytes = _writerhiddenmul(f, n, spec.offsetwidth, "element-offset") + push!(buffers, _writerhiddenbytes(f, nbytes, "element-offset")) + elseif role == AC.VIEWS + nbytes = _writerhiddenmul(f, n, 16, "view") + push!(buffers, _writerhiddenbytes(f, nbytes, "view")) + else + throw( + ArgumentError( + "hidden construction for retained field $(f.name) needs a " * + "layout-specific $(role) buffer", + ), + ) + end + end + return buffers, nullcount +end + +"Build exact retained data for slots whose values are hidden by an ancestor." +function _constructhiddenpart( + f::AC.Field, + n::Int, + context::_WriterContext; + forcevalid::Bool=false, + inactive::Bool=false, +) + n >= 0 || throw(ArgumentError("negative hidden length for retained field $(f.name)")) + t = f.type + if t isa AC.NullType + !inactive && + forcevalid && + n > 0 && + throw( + ArgumentError( + "retained Null field $(f.name) cannot synthesize a non-null hidden value", + ), + ) + return _retainedfield(f), AC.ArrayData(t, n, AC.BufferSlice[]; nullcount=n) + elseif t isa AC.DictionaryType + valuefield = AC.dictvaluefield(f, t) + # Unlike the shared needvalid rule, `inactive` cannot waive validity + # here: a non-nullable dictionary's indices are all valid, so they + # must point at something — one pool entry — even in an inactive + # subtree. + needvalid = (!inactive && forcevalid) || !f.nullable + poollength = n == 0 || !needvalid ? 0 : 1 + rebuiltvaluefield, dictionary = + _constructhiddenpart(valuefield, poollength, context; inactive) + buffers, nullcount = + _writerhiddenbuffers(f, n; forcevalid=(!inactive && forcevalid)) + field = AC.Field( + f.name, + t; + nullable=f.nullable, + metadata=_fieldmetadata(f), + children=collect(AC.Field, rebuiltvaluefield.children), + ) + data = AC.ArrayData(t, n, buffers; dictionary, nullcount) + return field, data + elseif t isa Union{AC.ListType,AC.ListViewType,AC.MapType} + length(f.children) == 1 || throw( + ArgumentError( + "retained $(AC.descriptorname(t)) field $(f.name) needs one child", + ), + ) + childfield, childdata = _constructhiddenpart(f.children[1], 0, context; inactive) + buffers, nullcount = + _writerhiddenbuffers(f, n; forcevalid=(!inactive && forcevalid)) + data = AC.ArrayData(t, n, buffers; children=AC.ArrayData[childdata], nullcount) + return _retainedfield(f; children=AC.Field[childfield]), data + elseif t isa AC.FixedSizeListType + length(f.children) == 1 || + throw(ArgumentError("retained fixed-size-list field $(f.name) needs one child")) + childlength64 = _writerhiddenmul(f, n, t.listsize, "fixed-size-list child") + childlength64 <= typemax(Int) || throw( + ArgumentError( + "hidden fixed-size-list child for retained field $(f.name) exceeds " * + "the supported length range", + ), + ) + childfield, childdata = + _constructhiddenpart(f.children[1], Int(childlength64), context; inactive) + buffers, nullcount = + _writerhiddenbuffers(f, n; forcevalid=(!inactive && forcevalid)) + data = AC.ArrayData(t, n, buffers; children=AC.ArrayData[childdata], nullcount) + return _retainedfield(f; children=AC.Field[childfield]), data + elseif t isa AC.StructType + childfields = AC.Field[] + childdata = AC.ArrayData[] + for child in f.children + rebuiltfield, rebuiltdata = _constructhiddenpart(child, n, context; inactive) + push!(childfields, rebuiltfield) + push!(childdata, rebuiltdata) + end + buffers, nullcount = + _writerhiddenbuffers(f, n; forcevalid=(!inactive && forcevalid)) + data = AC.ArrayData(t, n, buffers; children=childdata, nullcount) + return _retainedfield(f; children=childfields), data + elseif t isa AC.RunEndEncodedType + length(f.children) == 2 || throw( + ArgumentError("retained run-end encoded field $(f.name) needs two children"), + ) + runfield, valuefield = f.children + if n == 0 + rebuiltrunfield, runenddata = _constructhiddenpart(runfield, 0, context) + rebuiltvaluefield, valuedata = _constructhiddenpart(valuefield, 0, context) + else + runfield.type isa AC.IntType || throw( + ArgumentError("retained REE field $(f.name) has a non-integer run end"), + ) + RT = AC.juliatype(runfield.type) + n <= typemax(RT) || throw( + ArgumentError( + "hidden REE length for retained field $(f.name) exceeds its " * + "run-end type", + ), + ) + rebuiltrunfield, runenddata = _constructpart(runfield, RT[RT(n)]; context) + rebuiltvaluefield, valuedata = _constructhiddenpart( + valuefield, + 1, + context; + forcevalid=(!inactive && (forcevalid || !f.nullable)), + inactive, + ) + end + data = AC.ArrayData( + t, + n, + AC.BufferSlice[]; + children=AC.ArrayData[runenddata, valuedata], + nullcount=0, + ) + return _retainedfield(f; children=AC.Field[rebuiltrunfield, rebuiltvaluefield]), + data + elseif t isa AC.UnionType + isempty(f.children) && + throw(ArgumentError("retained Union field $(f.name) has no children")) + if n == 0 + childfields = AC.Field[] + childdata = AC.ArrayData[] + for child in f.children + rebuiltfield, rebuiltdata = _constructhiddenpart(child, 0, context) + push!(childfields, rebuiltfield) + push!(childdata, rebuiltdata) + end + buffers = + t.mode == AC.SparseMode ? AC.BufferSlice[AC.BufferSlice()] : + AC.BufferSlice[AC.BufferSlice(), AC.BufferSlice()] + data = AC.ArrayData(t, 0, buffers; children=childdata, nullcount=0) + return _retainedfield(f; children=childfields), data + end + needvalid = !inactive && (forcevalid || !f.nullable) + chosen = _writerhiddenunionchild(f; forcevalid=needvalid, inactive) + childfields = AC.Field[] + childdata = AC.ArrayData[] + for (index, child) in enumerate(f.children) + childlength = t.mode == AC.SparseMode ? n : index == chosen && n > 0 ? 1 : 0 + rebuiltfield, rebuiltdata = _constructhiddenpart( + child, + childlength, + context; + forcevalid=needvalid && index == chosen, + inactive=inactive || index != chosen, + ) + push!(childfields, rebuiltfield) + push!(childdata, rebuiltdata) + end + typeids = n == 0 ? AC.BufferSlice() : AC._databuffer(fill(t.typeids[chosen], n)) + buffers = if t.mode == AC.SparseMode + AC.BufferSlice[typeids] + else + offsets = n == 0 ? AC.BufferSlice() : AC._databuffer(zeros(Int32, n)) + AC.BufferSlice[typeids, offsets] + end + data = AC.ArrayData(t, n, buffers; children=childdata, nullcount=0) + return _retainedfield(f; children=childfields), data + end + + isempty(f.children) || throw( + ArgumentError( + "retained leaf field $(f.name) unexpectedly has $(length(f.children)) children", + ), + ) + buffers, nullcount = _writerhiddenbuffers(f, n; forcevalid=(!inactive && forcevalid)) + return _retainedfield(f), AC.ArrayData(t, n, buffers; nullcount) +end + +"Build a fixed-size list whose immediate child has Null storage." +function _constructfixednullpart( + f::AC.Field, + present::AbstractVector{Bool}, + context::_WriterContext, +) + t = f.type::AC.FixedSizeListType + length(f.children) == 1 && only(f.children).type isa AC.NullType || throw( + ArgumentError( + "retained fixed-size-list field $(f.name) does not have an immediate Null child", + ), + ) + any(!, present) && + !f.nullable && + throw( + ArgumentError( + "column $(f.name) contains a null fixed-size-list value under a " * + "non-nullable field", + ), + ) + childlength64 = _writerhiddenmul(f, length(present), t.listsize, "fixed-list child") + childlength64 <= typemax(Int) || throw( + ArgumentError( + "retained fixed-size-list child $(f.name) exceeds the supported length range", + ), + ) + childfield, childdata = + _constructhiddenpart(only(f.children), Int(childlength64), context) + data = AC.ArrayData( + t, + length(present), + AC.BufferSlice[AC._bitmapbuffer(present)]; + children=AC.ArrayData[childdata], + nullcount=count(!, present), + ) + return _retainedfield(f; children=AC.Field[childfield]), data +end + +"Build exact extension-free nested fixed-list validity over a Null leaf." +function _constructplainfixednullpart( + f::AC.Field, + values::AbstractVector, + context::_WriterContext; + routed::Bool=false, +) + _writerplainfixednullshape(f) || throw( + ArgumentError( + "retained fixed-size-list field $(f.name) is not an extension-free Null-only shape", + ), + ) + slots = Pair{Int,Any}[] + sizehint!(slots, length(values)) + for (i, raw) in enumerate(values) + raw === _RETAINED_HIDDEN && continue + value = if routed + raw isa _WriterRoutedUnion || throw( + ArgumentError( + "retained sparse Union child $(f.name) lost its writer route", + ), + ) + raw.value + else + raw + end + _checkplainfixednullvalue(f, value) + push!(slots, i => value) + end + return _constructplainfixednullpart(f, length(values), slots, context) +end + +function _constructplainfixednullpart(f::AC.Field, n::Int, slots, context::_WriterContext) + t = f.type::AC.FixedSizeListType + child = only(f.children) + present = f.nullable ? falses(n) : nothing + sequences = Pair{Int,Any}[] + for slot in slots + index, value = first(slot), last(slot) + if value === missing + f.nullable || throw( + ArgumentError( + "retained field $(f.name) is non-nullable but received a null value", + ), + ) + continue + end + present === nothing || (present[index] = true) + push!(sequences, index => value) + end + + childlength64 = _writerhiddenmul(f, n, t.listsize, "fixed-list child") + childlength64 <= typemax(Int) || throw( + ArgumentError( + "retained fixed-size-list child $(f.name) exceeds the supported length range", + ), + ) + childlength = Int(childlength64) + childfield, childdata = if child.type isa AC.NullType + _constructhiddenpart(child, childlength, context) + else + childslots = ( + Base.checked_add( + Base.checked_mul(Int64(index - 1), Int64(t.listsize)), + Int64(childindex), + ) => item for (index, sequence) in sequences for + (childindex, item) in enumerate(sequence) + ) + _constructplainfixednullpart(child, childlength, childslots, context) + end + nullcount = present === nothing ? 0 : n - length(sequences) + buffers = + AC.BufferSlice[present === nothing ? AC.BufferSlice() : AC._bitmapbuffer(present),] + data = AC.ArrayData(t, n, buffers; children=AC.ArrayData[childdata], nullcount) + return _retainedfield(f; children=AC.Field[childfield]), data +end + +"Construct a sparse child whose visible fixed-list values have only Null storage." +function _constructcompactfixednull(f::AC.Field, values, context::_WriterContext) + _writercompactfixednull(f) || + throw(ArgumentError("retained fixed-size-list field $(f.name) is not Null-only")) + plain = _writerplainfixednullshape(f) + plain && return _constructplainfixednullpart(f, values, context; routed=true) + visible = _writerstoragevector(f) + present = Vector{Bool}(undef, length(values)) + for (i, value) in enumerate(values) + if value === _RETAINED_HIDDEN + present[i] = true + continue + end + value isa _WriterRoutedUnion || throw( + ArgumentError("retained sparse Union child $(f.name) lost its writer route"), + ) + present[i] = value.value !== missing + _writerpushstorage!(visible, value.value, f, value.writertype, context) + end + # Visible rows are validated then discarded; see _writercompactfixednull. + isempty(visible) || _constructwriterstorage(f, visible, context) + return _constructfixednullpart(f, present, context) +end + +function _arrowtypeswriterstoragevalue( + context::_WriterContext, + f::AC.Field, + value, + writetype::Type, +) + if f.type isa AC.UnionType + # A single logical type may lower to a physical Union. Keep that case + # distinct from a declared logical Union: the former must be lowered + # before routing, while the latter uses its logical branches to retain + # branch identity (including same-storage extension types). + logicaltype = Base.nonmissingtype(writetype) + logicalstorage = _writerstoragetype!(context, logicaltype) + lowers_to_union = isconcretetype(logicaltype) && logicalstorage isa Union + storagevalue = value + storagetype = writetype + if lowers_to_union + if value !== missing + T = typeof(value) + storagevalue = + T === VersionNumber ? string(value) : + _writerneedstype!(context, T) ? + _writertoarrow(context, value, T, f.name) : value + end + storagetype = + Missing <: writetype ? Union{Missing,logicalstorage} : logicalstorage + storagevalue = _writerconvertlowered( + storagetype, + storagevalue, + context, + logicaltype, + f.name, + ) + elseif !(writetype isa Union) + throw( + ArgumentError( + "writer type $writetype does not describe retained Union field " * + "$(f.name)", + ), + ) + end + variants, children = _writerunionplan!( + context, + f, + storagetype; + extension_shape=lowers_to_union && + _writerextensionshape(context, logicaltype, logicalstorage), + ) + branch = _writerunionbranch!(context, storagetype, typeof(storagevalue), variants) + child = children[branch] + return _WriterRoutedUnion( + child, + _arrowtypeswriterchildstoragevalue( + context, + f.children[child], + storagevalue, + variants[branch], + ), + typeof(value), + ) + end + if writetype isa Union + variants = _writerunionvariants!(context, writetype) + branch = _writerunionbranch!(context, writetype, typeof(value), variants) + writetype = variants[branch] + end + value === missing && return missing + T = typeof(value) + value = if T === VersionNumber + string(value) + elseif _writerneedstype!(context, T) + _writertoarrow(context, value, T, f.name) + else + value + end + storagewritetype = + T === VersionNumber ? String : _writerstoragetype!(context, writetype) + # Validate the complete lowered shape before recursive Field access. A + # malformed Struct/List/Map result must fail at the ArrowTypes seam instead + # of leaking a BoundsError or unrelated property/index exception. + value = _writerconvertlowered(storagewritetype, value, context, T, f.name) + t = f.type + t isa AC.NullType && value === nothing && return missing + if t isa AC.RunEndEncodedType + length(f.children) == 2 || throw( + ArgumentError("writer REE field $(f.name) needs run-end and value children"), + ) + return _arrowtypeswriterchildstoragevalue( + context, + f.children[2], + value, + storagewritetype, + ) + end + if t isa AC.DateType || + t isa AC.TimestampType || + t isa AC.TimeType || + t isa AC.DurationType + ok, storage = _facadetostorage(t, value) + ok && return storage + throw( + ArgumentError( + "registered writer value $(repr(value)) for field $(f.name) cannot " * + "be represented exactly by retained Arrow type $(repr(t))", + ), + ) + end + if t isa Union{AC.ListType,AC.ListViewType,AC.FixedSizeListType} + child = only(f.children) + childtype = eltype(storagewritetype) + return Any[ + _arrowtypeswriterchildstoragevalue(context, child, x, childtype) for x in value + ] + elseif t isa AC.StructType + fieldcount(storagewritetype) == length(f.children) || throw( + ArgumentError( + "writer type $storagewritetype does not match retained Struct field " * + "$(f.name)", + ), + ) + names = value isa NamedTuple ? _arrowtypesstructnames(f) : nothing + out = Pair{String,Any}[] + for (i, child) in enumerate(f.children) + childvalue = if value isa NamedTuple + getproperty(value, names[i]) + elseif value isa Tuple + getfield(value, i) + elseif value isa AbstractVector + kv = value[i] + kv isa Pair && String(first(kv)) == child.name || throw( + ArgumentError( + "writer Struct field $(f.name) has incompatible child $i", + ), + ) + last(kv) + else + getfield(value, i) + end + childtype = fieldtype(storagewritetype, i) + push!( + out, + child.name => _arrowtypeswriterchildstoragevalue( + context, + child, + childvalue, + childtype, + ), + ) + end + return out + elseif t isa AC.MapType + entries = only(f.children) + keyfield, valuefield = entries.children + keytype = Base.keytype(storagewritetype) + valuetype = Base.valtype(storagewritetype) + return Pair{Any,Any}[ + _arrowtypeswriterchildstoragevalue(context, keyfield, first(kv), keytype) => + _arrowtypeswriterchildstoragevalue( + context, + valuefield, + last(kv), + valuetype, + ) for kv in value + ] + end + return value +end + +function _registeredvalueplan!(context::_WriterContext, f::AC.Field, value, writetype::Type) + _, target, _ = _arrowtypestarget(context.arrowtypes, f) + target === nothing && return nothing + logical = Base.nonmissingtype(writetype) + (value === missing ? logical <: target : value isa target) || return nothing + bytype = get!(context.registered, f) do + Dict{Tuple{Type,Type},_RegisteredWriterPlan}() + end + runtime = typeof(value) + # Every value of a fully declared type has the same complete routing plan. + # Abstract declarations still need one plan per inferred runtime type. + evidence = _writerfullydeclared(logical) ? logical : runtime + return get!(bytype, (writetype, evidence)) do + values = Vector{writetype}(undef, 1) + values[1] = value + _registeredwriterplan(context, f, values)::_RegisteredWriterPlan + end +end + +function _deferredregisteredwriter( + context::_WriterContext, + f::AC.Field, + value, + writetype::Type, +) + _, target, _ = _arrowtypestarget(context.arrowtypes, f) + target === nothing && return nothing + logical = Base.nonmissingtype(writetype) + (value === missing ? logical <: target : value isa target) || return nothing + value === missing || + _writerfullydeclared(logical) || + _writerregisterinferred!(context, f, typeof(value)) + context.deferred = true + return _DeferredRegisteredWriter(f, value, writetype) +end + +function _writerconvertedvalue( + context::_WriterContext, + f::AC.Field, + value, + writetype::Type, + routes::Union{Nothing,_WriterTypeRoutes}, +) + routes === nothing && return _arrowtypeswriterstoragevalue(context, f, value, writetype) + runtime = typeof(value) + route = _writerroute(routes, runtime) + if route > 0 + route <= length(f.children) || + throw(ArgumentError("writer route $route is outside retained Field $(f.name)")) + return _WriterRoutedUnion( + route, + _arrowtypeswriterstoragevalue(context, f.children[route], value, runtime), + runtime, + ) + end + runtimewriter = + f.type isa AC.UnionType && Missing <: writetype && runtime !== Missing ? + Union{Missing,runtime} : runtime + return _arrowtypeswriterstoragevalue(context, f, value, runtimewriter) +end + +function _arrowtypeswriterchildstoragevalue( + context::_WriterContext, + f::AC.Field, + value, + writetype::Type, +) + if context.collecting + deferred = _deferredregisteredwriter(context, f, value, writetype) + deferred === nothing || return deferred + end + plan = _registeredvalueplan!(context, f, value, writetype) + plan === nothing && return _arrowtypeswriterstoragevalue(context, f, value, writetype) + return _writerconvertedvalue(context, f, value, plan.input, plan.routes) +end + +function _withwritercollection(f, context::_WriterContext) + collecting = context.collecting + deferred = context.deferred + context.collecting = true + context.deferred = false + try + value = f() + return value, context.deferred + finally + context.collecting = collecting + context.deferred = deferred + end +end + +function _writerresolvedeferredlevel( + context::_WriterContext, + value::_DeferredRegisteredWriter, +) + plan = _registeredvalueplan!(context, value.field, value.value, value.writetype) + plan === nothing && return _arrowtypeswriterstoragevalue( + context, + value.field, + value.value, + value.writetype, + ) + return _writerconvertedvalue(context, value.field, value.value, plan.input, plan.routes) +end + +function _writerresolvedeferredlevel(context::_WriterContext, value::_WriterRoutedUnion) + _writerhasdeferred(value.value) || return value + return _WriterRoutedUnion( + value.child, + _writerresolvedeferredlevel(context, value.value), + value.writertype, + ) +end + +function _writerresolvedeferredlevel(context::_WriterContext, value::Pair) + _writerhasdeferred(value) || return value + return _writerresolvedeferredlevel(context, first(value)) => + _writerresolvedeferredlevel(context, last(value)) +end + +function _writerresolvedeferredlevel(context::_WriterContext, values::AbstractVector) + _writerhasdeferred(values) || return values + for i in eachindex(values) + values[i] = _writerresolvedeferredlevel(context, values[i]) + end + return values +end + +_writerresolvedeferredlevel(::_WriterContext, value) = value + +_writerhasdeferred(::_DeferredRegisteredWriter) = true +_writerhasdeferred(value::_WriterRoutedUnion) = _writerhasdeferred(value.value) +_writerhasdeferred(value::Pair) = + _writerhasdeferred(first(value)) || _writerhasdeferred(last(value)) +_writerhasdeferred(values::AbstractVector) = any(_writerhasdeferred, values) +_writerhasdeferred(value) = false + +""" +Empty storage vector typed for one Field's physical values (Any for Unions or +unknown types). +""" +function _writerstoragevector(f::AC.Field) + T = _declaredeltype(f, false) + return f.type isa AC.UnionType || T === Any ? Any[] : T[] +end + +function _writerpushstorage!( + storage::AbstractVector, + value, + f::AC.Field, + writertype::Type, + context::_WriterContext, +) + stored = _writerconvertlowered(eltype(storage), value, context, writertype, f.name) + push!(storage, stored) + return nothing +end + +abstract type _MaskedLeafProjection end + +struct _NativeMaskedLeafProjection <: _MaskedLeafProjection + field::AC.Field + context::_WriterContext +end + +struct _RegisteredMaskedLeafProjection{R} <: _MaskedLeafProjection + field::AC.Field + writertype::Type + routes::R + context::_WriterContext +end + +function _maskedleafvalue( + projection::_NativeMaskedLeafProjection, + ::Type{S}, + value, +) where {S} + value === missing && return missing + return _writerconvertlowered(S, value, projection.context, S, projection.field.name) +end + +function _maskedleafvalue( + projection::_RegisteredMaskedLeafProjection, + ::Type{S}, + value, +) where {S} + if projection.routes === nothing && + projection.writertype === S && + !_writerneedstype!(projection.context, S) + value === missing && return missing + return _writerconvertlowered(S, value, projection.context, S, projection.field.name) + end + lowered = _writerconvertedvalue( + projection.context, + projection.field, + value, + projection.writertype, + projection.routes, + ) + lowered === missing && return missing + return _writerconvertlowered( + S, + lowered, + projection.context, + projection.writertype, + projection.field.name, + ) +end + +function _maskedleafnullerror(f::AC.Field) + throw( + ArgumentError( + "column $(f.name) holds missing values but its retained field is non-nullable", + ), + ) +end + +function _maskedleafvalidity(f::AC.Field, n::Int) + f.nullable || return nothing, 0 + bytes = zeros(UInt8, AC.expected_validity_bytes(Int64(n))) + return bytes, n +end + +function _maskedleafpresent!(validity, index::Int, nullcount::Int) + validity === nothing && return nullcount + AC._setbitmapbit!(validity, index) + return nullcount - 1 +end + +function _maskedleafvaliditybuffer(validity, nullcount::Int) + validity === nothing && return AC.BufferSlice() + nullcount == 0 && return AC.BufferSlice() + return AC._databuffer(validity) +end + +function _maskedleafarray( + f::AC.Field, + n::Int, + validity, + databuffer::AC.BufferSlice, + nullcount::Int, +) + buffers = AC.BufferSlice[_maskedleafvaliditybuffer(validity, nullcount), databuffer] + return f, AC.ArrayData(f.type, n, buffers; nullcount) +end + +function _constructmaskedboolleaf( + f::AC.Field, + values::V, + projection::_MaskedLeafProjection, +) where {V<:_MaskedChildValues} + n = length(values) + data = zeros(UInt8, AC.expected_validity_bytes(Int64(n))) + validity, nullcount = _maskedleafvalidity(f, n) + for index in eachindex(values) + raw = values[index] + raw === _RETAINED_HIDDEN && continue + stored = _maskedleafvalue(projection, Bool, raw) + if stored === missing + f.nullable || _maskedleafnullerror(f) + continue + end + nullcount = _maskedleafpresent!(validity, index, nullcount) + stored && AC._setbitmapbit!(data, index) + end + return _maskedleafarray(f, n, validity, AC._databuffer(data), nullcount) +end + +function _constructmaskedfixedleaf( + f::AC.Field, + values::V, + projection::_MaskedLeafProjection, + ::Type{S}, +) where {S,V<:_MaskedChildValues} + n = length(values) + placeholder = _writerhiddenstorage(f; forcevalid=true) + storedplaceholder = _writerconvertlowered(S, placeholder, projection.context, S, f.name) + data = fill(storedplaceholder, n) + validity, nullcount = _maskedleafvalidity(f, n) + for index in eachindex(values) + raw = values[index] + raw === _RETAINED_HIDDEN && continue + stored = _maskedleafvalue(projection, S, raw) + if stored === missing + f.nullable || _maskedleafnullerror(f) + continue + end + data[index] = stored + nullcount = _maskedleafpresent!(validity, index, nullcount) + end + return _maskedleafarray(f, n, validity, AC._databuffer(data), nullcount) +end + +function _constructmaskedvarbytesleaf( + f::AC.Field, + values::V, + projection::_MaskedLeafProjection, + ::Type{S}, +) where {S,V<:_MaskedChildValues} + t = f.type::Union{AC.Utf8Type,AC.BinaryType} + return _constructmaskedvarbytesleaf(f, values, projection, S, t.large ? Int64 : Int32) +end + +function _constructmaskedvarbytesleaf( + f::AC.Field, + values::V, + projection::_MaskedLeafProjection, + ::Type{S}, + ::Type{Offset}, +) where {S,Offset<:Union{Int32,Int64},V<:_MaskedChildValues} + t = f.type::Union{AC.Utf8Type,AC.BinaryType} + n = length(values) + offsets = Vector{Offset}(undef, n + 1) + offsets[1] = zero(Offset) + data = UInt8[] + validity, nullcount = _maskedleafvalidity(f, n) + for index in eachindex(values) + raw = values[index] + if raw !== _RETAINED_HIDDEN + stored = _maskedleafvalue(projection, S, raw) + if stored === missing + f.nullable || _maskedleafnullerror(f) + else + bytes = if t isa AC.Utf8Type + stored isa AbstractString || throw( + ArgumentError("column $(f.name) must contain string values"), + ) + codeunits(stored) + else + stored isa AbstractVector{UInt8} || throw( + ArgumentError("column $(f.name) must contain byte-vector values"), + ) + stored + end + length(data) <= typemax(Offset) - length(bytes) || throw( + ArgumentError( + "column $(f.name) data exceeds its retained offset width", + ), + ) + append!(data, bytes) + nullcount = _maskedleafpresent!(validity, index, nullcount) + end + end + offsets[index + 1] = Offset(length(data)) + end + databuffer = isempty(data) ? AC.BufferSlice() : AC._databuffer(data) + buffers = AC.BufferSlice[ + _maskedleafvaliditybuffer(validity, nullcount), + AC._databuffer(offsets), + databuffer, + ] + return f, AC.ArrayData(t, n, buffers; nullcount) +end + +"Bulk-fill mixed visible and parent-hidden composite slots in physical storage." +function _constructmaskedleaf(f::AC.Field, values, projection::_MaskedLeafProjection) + values isa _MaskedChildValues || return nothing + isempty(f.children) || return nothing + f.type isa AC.DictionaryType && return nothing + any(row -> row === missing || row === _RETAINED_HIDDEN, values.rows) || return nothing + T = _declaredeltype(f, false) + T === Any && return nothing + S = Base.nonmissingtype(T) + S === Union{} && return nothing + f.type isa AC.BoolType && + S === Bool && + return _constructmaskedboolleaf(f, values, projection) + f.type isa Union{AC.Utf8Type,AC.BinaryType} && + return _constructmaskedvarbytesleaf(f, values, projection, S) + spec = AC.layoutspec_of(f.type) + isbitstype(S) && spec.fixedwidth > 0 && sizeof(S) == spec.fixedwidth || return nothing + return _constructmaskedfixedleaf(f, values, projection, S) +end + +"Bulk-fill composite slots hidden by parent validity at a physical leaf." +function _constructmaskedwriterleaf( + f::AC.Field, + writetype::Type, + values, + routes, + context::_WriterContext, +) + # The masked-leaf fast path writes one physical child. A nonzero route + # means the value belongs under a Union child, which this path cannot + # express. + (routes === nothing || all(iszero, Base.values(routes))) || return nothing + projection = _RegisteredMaskedLeafProjection(f, writetype, routes, context) + masked = _constructmaskedleaf(f, values, projection) + masked === nothing || return masked + values isa _MaskedChildValues || return nothing + isempty(f.children) || return nothing + f.type isa AC.DictionaryType && return nothing + any(row -> row === missing || row === _RETAINED_HIDDEN, values.rows) || return nothing + storage = _writerstoragevector(f) + eltype(storage) === Any && return nothing + + hidden = _writerhiddenvalue(f, writetype) + storedhidden = + _writerconvertlowered(eltype(storage), hidden, context, writetype, f.name) + resize!(storage, length(values)) + fill!(storage, storedhidden) + for index in eachindex(values) + value = values[index] + value === _RETAINED_HIDDEN && continue + lowered = _writerconvertedvalue(context, f, value, writetype, routes) + storage[index] = + _writerconvertlowered(eltype(storage), lowered, context, writetype, f.name) + end + return _constructwriterstorage(f, storage, context) +end + +function _constructwriterchild( + f::AC.Field, + writetype::Type, + values; + routes::Union{Nothing,_WriterTypeRoutes}=nothing, + context::_WriterContext=_WriterContext(), +) + masked = _constructmaskedwriterleaf(f, writetype, values, routes, context) + masked === nothing || return masked + hashidden = any(value -> value === _RETAINED_HIDDEN, values) + hidden = hashidden ? _writerhiddenvalue(f, writetype) : nothing + prepared, deferred = _withwritercollection(context) do + out = Any[] + sizehint!(out, length(values)) + for value in values + push!( + out, + value === _RETAINED_HIDDEN ? hidden : + _writerconvertedvalue(context, f, value, writetype, routes), + ) + end + out + end + rounds = 0 + while deferred + rounds < _MAX_WRITER_SCHEMA_DEPTH || throw( + ArgumentError( + "registered writer values exceed the supported nested planning " * + "depth $_MAX_WRITER_SCHEMA_DEPTH for column $(f.name)", + ), + ) + rounds += 1 + prepared, deferred = _withwritercollection(context) do + Any[_writerresolvedeferredlevel(context, value) for value in prepared] + end + end + converted = _writerstoragevector(f) + sizehint!(converted, length(values)) + for value in prepared + _writerpushstorage!(converted, value, f, writetype, context) + end + return _constructwriterstorage(f, converted, context) +end + +"Construct values already lowered to one Field's storage domain." +function _constructwriterstorage( + f::AC.Field, + storage::AbstractVector, + context::_WriterContext, +) + t = f.type + if t isa AC.DateType || + t isa AC.TimestampType || + t isa AC.TimeType || + t isa AC.DurationType + values = Union{Missing,Int64}[x === missing ? missing : Int64(x) for x in storage] + return _rebuildtemporal(f, values) + end + storagefield = _writerstoragefield(f) + _, data = _constructpart(storagefield, storage; context) + return f, data +end + +"Remove logical labels while rebuilding values already in physical storage." +function _writerstoragefield(f::AC.Field) + own = _withoutownextension(f) + children = AC.Field[_writerstoragefield(child) for child in f.children] + return AC.Field( + own.name, + own.type; + nullable=own.nullable, + metadata=_fieldmetadata(own), + children, + ) +end + +function _arrowtypesfixedlistcolumn( + name::String, + v::AbstractVector, + kind; + extension_shape::Bool, + context, +) + N = ArrowTypes.getsize(kind) + E = ArrowTypes.gettype(kind) + present = Bool[x !== missing for x in v] + for row in v + row === missing || + length(row) == N || + throw( + ArgumentError( + "ArrowTypes fixed-list column $name expected $N values per row", + ), + ) + end + masked = if any(!, present) + emptychildfield, _ = _arrowtypeschildcolumn("item", E[]; extension_shape, context) + t = AC.FixedSizeListType(N) + emptyfield = AC.Field( + name, + t; + nullable=Missing <: eltype(v), + children=AC.Field[emptychildfield], + ) + if _writercompactfixednull(emptyfield) + # Only visible rows need validation; see _writercompactfixednull. + if _writerplainfixednullshape(emptyfield) + return _constructplainfixednullpart(emptyfield, v, context) + else + visible = E[] + nvisible = count(identity, present) + sizehint!(visible, Base.checked_mul(nvisible, N)) + for row in v + row === missing || append!(visible, row) + end + isempty(visible) || + _constructwriterchild(emptychildfield, E, visible; context) + end + return _constructfixednullpart(emptyfield, present, context) + end + values = _MaskedFixedListValues(v, N) + _constructwriterchild(emptychildfield, E, values; context) + else + nothing + end + if masked === nothing + flat = E[] + sizehint!(flat, Base.checked_mul(length(v), N)) + for row in v + append!(flat, row) + end + childfield, childdata = + _arrowtypeschildcolumn("item", flat; extension_shape, context) + else + childfield, childdata = masked + end + t = AC.FixedSizeListType(N) + field = AC.Field(name, t; nullable=Missing <: eltype(v), children=AC.Field[childfield]) + data = AC.ArrayData( + t, + length(v), + [AC._bitmapbuffer(present)]; + children=AC.ArrayData[childdata], + nullcount=count(!, present), + ) + return field, data +end + +function _arrowtypesstructcolumn( + name::String, + v::AbstractVector, + S; + extension_shape::Bool, + context, +) + _arrowtypesstorageisspecified(context, S) || throw( + ArgumentError( + "ArrowTypes StructKind column $name lowered to an underspecified type $S; " * + "give every struct field a writable storage type", + ), + ) + names = fieldnames(S) + nchildren = fieldcount(S) + present = Bool[x !== missing for x in v] + childfields = AC.Field[] + childdata = AC.ArrayData[] + hasnull = any(!, present) + for j = 1:nchildren + FT = fieldtype(S, j) + childname = string(names[j]) + # A null parent does not make its children nullable; the parent's + # validity bitmap masks their slots. Build hidden slots from the + # resolved child Field in the storage domain — no logical default + # constructors, no duplicated variable payloads. + cf, cd = if hasnull + emptychildfield, _ = + _arrowtypeschildcolumn(childname, FT[]; extension_shape, context) + values = _freshstructchildvalues(Any, v, j) + _constructwriterchild(emptychildfield, FT, values; context) + else + values = _freshstructchildvalues(FT, v, j) + _arrowtypeschildcolumn(childname, values; extension_shape, context) + end + push!(childfields, cf) + push!(childdata, cd) + end + t = AC.StructType() + field = AC.Field(name, t; nullable=Missing <: eltype(v), children=childfields) + data = AC.ArrayData( + t, + length(v), + [AC._bitmapbuffer(present)]; + children=childdata, + nullcount=count(!, present), + ) + return field, data +end + +function _arrowtypesmapcolumn( + name::String, + v::AbstractVector, + S; + extension_shape::Bool, + context, +) + S <: AbstractDict || throw( + ArgumentError("ArrowTypes MapKind column $name must lower to AbstractDict rows"), + ) + K = keytype(S) + V = valtype(S) + keys = K[] + values = V[] + present = Bool[x !== missing for x in v] + offsets = Vector{Int32}(undef, length(v) + 1) + offsets[1] = 0 + for (i, row) in enumerate(v) + if row !== missing + row isa AbstractDict || throw( + ArgumentError( + "ArrowTypes MapKind column $name contains a non-dictionary row", + ), + ) + length(keys) <= typemax(Int32) - length(row) || throw( + ArgumentError( + "ArrowTypes MapKind column $name exceeds the Int32 offset range", + ), + ) + for (key, value) in row + push!(keys, key) + push!(values, value) + end + end + offsets[i + 1] = Int32(length(keys)) + end + any(ismissing, keys) && throw(ArgumentError("Arrow Map keys cannot be missing")) + keyfield, keydata = _arrowtypeschildcolumn("key", keys; extension_shape, context) + keyfield.nullable && + throw(ArgumentError("Arrow Map keys must have a non-nullable type")) + _validatemapphysicalkeys(name, keyfield, keydata, offsets) + valuefield, valuedata = + _arrowtypeschildcolumn("value", values; extension_shape, context) + entriesfield = AC.Field( + "entries", + AC.StructType(); + nullable=false, + children=AC.Field[keyfield, valuefield], + ) + entriesdata = AC.ArrayData( + AC.StructType(), + length(keys), + [AC.BufferSlice()]; + children=AC.ArrayData[keydata, valuedata], + nullcount=0, + ) + t = AC.MapType(false) + field = + AC.Field(name, t; nullable=Missing <: eltype(v), children=AC.Field[entriesfield]) + data = AC.ArrayData( + t, + length(v), + [AC._bitmapbuffer(present), AC._databuffer(offsets)]; + children=AC.ArrayData[entriesdata], + nullcount=count(!, present), + ) + return field, data +end + +"Registered public-domain value Field and Julia target for a retained dictionary." +function _registereddictionaryfield(context::_WriterContext, f::AC.Field) + t = f.type::AC.DictionaryType + _, target, _ = _arrowtypestarget(context.arrowtypes, f) + target === nothing && return nothing, nothing + valuefield = _arrowtypesdictvaluefield(f, t; retainmetadata=true) + return valuefield, target +end + +"Recursive Field equality: name, type, nullability, ordered metadata, and children." +function _fieldcontractequal(a::AC.Field, b::AC.Field) + a.name == b.name || return false + AC.typeequal(a.type, b.type) || return false + a.nullable == b.nullable || return false + _fieldmetadata(a) == _fieldmetadata(b) || return false + length(a.children) == length(b.children) || return false + for i in eachindex(a.children) + _fieldcontractequal(a.children[i], b.children[i]) || return false + end + return true +end + +""" +Assemble a fresh Int32-indexed dictionary Field and its first batch from +already-built pool data. +""" +function _newdictfromdata( + name::String, + valuefield::AC.Field, + valuedata::AC.ArrayData, + indices::Vector; + nullable::Bool, +) + t = AC.DictionaryType(AC.IntType(32, true), valuefield.type, false) + field = AC.Field( + name, + t; + nullable, + metadata=_fieldmetadata(valuefield), + children=collect(AC.Field, valuefield.children), + ) + return field, _dictbatch(field, indices, valuedata) +end + +"Construct a fresh dictionary's candidate categories before index encoding." +function _constructnewdictpooldata( + name::String, + pool::Vector; + valuefield::Union{Nothing,AC.Field}=nothing, + writetype::Union{Nothing,Type}=nothing, + routes::Union{Nothing,_WriterTypeRoutes}=nothing, + context::_WriterContext=_WriterContext(), +) + builtfield, valuedata = + valuefield === nothing || writetype === nothing ? + _constructpart(name, pool; context) : + _constructwriterchild(valuefield, writetype, pool; routes, context) + if valuefield !== nothing + _fieldcontractequal(valuefield, builtfield) || throw( + ArgumentError( + "column $name dictionary pool maps to a different Arrow Field " * + "than its partition values; keep the declared value types and " * + "ArrowTypes metadata identical across the column", + ), + ) + # The inferred value Field stays the schema; the writer adapter only + # supplied matching fresh ArrayData, so this is not retained + # reconstruction. + builtfield = valuefield + end + return builtfield, valuedata +end + +# --- retained-schema rewrite (facade Table/Stream round-trips) -------------- + +"Storage integers for a public column under a RETAINED temporal descriptor." +function _retainedstorage( + t::Union{AC.DateType,AC.TimestampType,AC.TimeType,AC.DurationType}, + v::AbstractVector, + name::String, +) + return _retainedstorage(_facadetoken(t), t, v, name) +end + +@noinline function _retainedstorage( + token::Val{K}, + t::Union{AC.DateType,AC.TimestampType,AC.TimeType,AC.DurationType}, + v::AbstractVector, + name::String, +) where {K} + out = Union{Missing,Int64}[] + sizehint!(out, length(v)) + for x in v + if x === missing + push!(out, missing) + else + sv = _exactfacadevalue(token, x) + sv === nothing && throw( + ArgumentError( + "column $name holds $(typeof(x)) values that do not match " * + "its retained Arrow type $(repr(t)); the column was " * + "replaced with incompatible data", + ), + ) + push!(out, Int64(sv)) + end + end + return out +end + +"Enforce the first partition's complete recursive Field contract." +function _checkpartitionfield( + expected::AC.Field, + actual::AC.Field, + partition::Int, + column::Symbol; + path::String=String(column), +) + expected.name == actual.name || throw( + ArgumentError( + "partition $partition column $column has child name $(repr(actual.name)) " * + "at $path, but the first partition declared $(repr(expected.name))", + ), + ) + AC.typeequal(expected.type, actual.type) || throw( + ArgumentError( + "partition $partition column $column maps $path to Arrow type " * + "$(repr(actual.type)), but the first partition declared " * + "$(repr(expected.type)); make the column types agree across partitions", + ), + ) + expectedmeta = something(_fieldmetadata(expected), Pair{String,String}[]) + actualmeta = something(_fieldmetadata(actual), Pair{String,String}[]) + expectedmeta == actualmeta || throw( + ArgumentError( + "partition $partition column $column has different ordered metadata " * + "at $path; the first partition declared $(repr(expectedmeta)), but " * + "this partition declared $(repr(actualmeta))", + ), + ) + actual.nullable && + !expected.nullable && + throw( + ArgumentError( + "partition $partition column $column is nullable at $path, but " * + "the first partition declared it non-nullable; make the first " * + "partition's corresponding element type admit Missing to widen " * + "the schema", + ), + ) + length(expected.children) == length(actual.children) || throw( + ArgumentError( + "partition $partition column $column has $(length(actual.children)) " * + "children at $path, but the first partition declared " * + "$(length(expected.children))", + ), + ) + for i in eachindex(expected.children) + expectedchild = expected.children[i] + actualchild = actual.children[i] + childpath = "$path.$(expectedchild.name)[$i]" + _checkpartitionfield(expectedchild, actualchild, partition, column; path=childpath) + end + return nothing +end + +function _retainedfield(f::AC.Field; children=collect(AC.Field, f.children)) + return AC.Field( + f.name, + f.type; + nullable=f.nullable, + metadata=_fieldmetadata(f), + children=children, + ) +end + +function _retainedvalidity(f::AC.Field, v::AbstractVector) + present = Bool[x !== missing for x in v] + any(!, present) && + !f.nullable && + throw( + ArgumentError( + "column $(f.name) holds missing values but its retained field is non-nullable", + ), + ) + return present +end + +function _retainedvarbytes(f::AC.Field, v::AbstractVector) + t = f.type::Union{AC.Utf8Type,AC.BinaryType} + present = _retainedvalidity(f, v) + Offset = t.large ? Int64 : Int32 + offsets = Vector{Offset}(undef, length(v) + 1) + offsets[1] = zero(Offset) + data = UInt8[] + for (i, x) in enumerate(v) + if x !== missing + bytes = if t isa AC.Utf8Type + x isa AbstractString || + throw(ArgumentError("column $(f.name) must contain string values")) + codeunits(x) + else + x isa AbstractVector{UInt8} || throw( + ArgumentError("column $(f.name) must contain byte-vector values"), + ) + x + end + length(data) <= typemax(Offset) - length(bytes) || throw( + ArgumentError("column $(f.name) data exceeds its retained offset width"), + ) + append!(data, bytes) + end + offsets[i + 1] = Offset(length(data)) + end + databuf = isempty(data) ? AC.BufferSlice() : AC._databuffer(data) + d = AC.ArrayData( + t, + length(v), + [AC._bitmapbuffer(present), AC._databuffer(offsets), databuf]; + nullcount=count(!, present), + ) + return _retainedfield(f), d +end + +function _retainedview(f::AC.Field, v::AbstractVector) + t = f.type::AC.ViewType + present = _retainedvalidity(f, v) + payloads = Vector{ArrowStrings.ArrowStringPayload}(undef, length(v)) + data = UInt8[] + for (i, x) in enumerate(v) + if x === missing + payloads[i] = ArrowStrings.PAYLOAD_MISSING + continue + end + bytes = if t.utf8 + x isa AbstractString || + throw(ArgumentError("column $(f.name) must contain string values")) + codeunits(x) + else + x isa AbstractVector{UInt8} || + throw(ArgumentError("column $(f.name) must contain byte-vector values")) + x + end + n = length(bytes) + if n <= ArrowStrings.INLINE_MAX + payloads[i] = ArrowStrings.inline_payload(bytes, 1, n) + else + length(data) <= typemax(Int32) - n || throw( + ArgumentError("column $(f.name) view data exceeds the Int32 offset range"), + ) + off = length(data) + append!(data, bytes) + payloads[i] = ArrowStrings.view_payload(data, off + 1, n, 0, off) + end + end + buffers = AC.BufferSlice[AC._bitmapbuffer(present), AC._databuffer(payloads)] + isempty(data) || push!(buffers, AC._databuffer(data)) + d = AC.ArrayData(t, length(v), buffers; nullcount=count(!, present)) + return _retainedfield(f), d +end + +function _retainedfixedbytes(f::AC.Field, v::AbstractVector) + t = f.type::AC.FixedSizeBinaryType + present = _retainedvalidity(f, v) + data = zeros(UInt8, Base.checked_mul(length(v), t.nbytes)) + for (i, x) in enumerate(v) + x === missing && continue + x isa AbstractVector{UInt8} || + throw(ArgumentError("column $(f.name) must contain byte-vector values")) + length(x) == t.nbytes || throw( + ArgumentError( + "column $(f.name) fixed-size value $i has $(length(x)) bytes; " * + "expected $(t.nbytes)", + ), + ) + copyto!(data, (i - 1) * t.nbytes + 1, x, 1, t.nbytes) + end + d = AC.ArrayData( + t, + length(v), + [AC._bitmapbuffer(present), AC._databuffer(data)]; + nullcount=count(!, present), + ) + return _retainedfield(f), d +end + +function _retaineddecimal(f::AC.Field, v::AbstractVector) + t = f.type::AC.DecimalType + present = _retainedvalidity(f, v) + if t.bits == 32 || t.bits == 64 + T = t.bits == 32 ? Int32 : Int64 + values = Vector{T}(undef, length(v)) + for (i, x) in enumerate(v) + x === missing || + x isa T || + throw( + ArgumentError( + "column $(f.name) must contain $T decimal storage values", + ), + ) + values[i] = x === missing ? zero(T) : x + end + databuf = AC._databuffer(values) + else + width = AC.primwidth(t) + values = zeros(UInt8, Base.checked_mul(length(v), width)) + for (i, x) in enumerate(v) + x === missing && continue + x isa AbstractVector{UInt8} || throw( + ArgumentError("column $(f.name) must contain byte-vector decimal values"), + ) + length(x) == width || throw( + ArgumentError( + "column $(f.name) decimal value $i has $(length(x)) bytes; " * + "expected $width", + ), + ) + copyto!(values, (i - 1) * width + 1, x, 1, width) + end + databuf = AC._databuffer(values) + end + d = AC.ArrayData( + t, + length(v), + [AC._bitmapbuffer(present), databuf]; + nullcount=count(!, present), + ) + return _retainedfield(f), d +end + +function _retainedinterval(f::AC.Field, v::AbstractVector) + t = f.type::AC.IntervalType + present = _retainedvalidity(f, v) + if t.unit == AC.YEAR_MONTH + values = Int32[ + x === missing ? Int32(0) : + x isa Int32 ? x : + throw(ArgumentError("column $(f.name) must contain Int32 intervals")) for + x in v + ] + elseif t.unit == AC.DAY_TIME + T = NamedTuple{(:days, :millis),Tuple{Int32,Int32}} + values = T[ + x === missing ? T((0, 0)) : + x isa T ? x : + throw(ArgumentError("column $(f.name) must contain $T intervals")) for + x in v + ] + else + T = NamedTuple{(:months, :days, :nanos),Tuple{Int32,Int32,Int64}} + values = T[ + x === missing ? T((0, 0, 0)) : + x isa T ? x : + throw(ArgumentError("column $(f.name) must contain $T intervals")) for + x in v + ] + end + d = AC.ArrayData( + t, + length(v), + [AC._bitmapbuffer(present), AC._databuffer(values)]; + nullcount=count(!, present), + ) + return _retainedfield(f), d +end + +function _retainedtypedvalues(f::AC.Field, values; converted::Bool) + T = _declaredeltype(f, converted) + T === Any && return Any[x for x in values] + out = Vector{T}(undef, length(values)) + for (i, x) in enumerate(values) + x isa T || throw( + ArgumentError( + "column $(f.name) holds $(typeof(x)) values that do not match " * + "its retained Arrow type $(repr(f.type))", + ), + ) + out[i] = x + end + return out +end + +"A physical child value hidden by a null composite parent." +function _retainedplaceholder(f::AC.Field; forcevalid::Bool=false) + f.nullable && !forcevalid && return missing + t = f.type + if t isa AC.NullType + forcevalid && throw( + ArgumentError( + "retained Null field $(f.name) cannot synthesize a non-null placeholder", + ), + ) + return missing + end + t isa Union{ + AC.IntType, + AC.FloatType, + AC.DateType, + AC.TimeType, + AC.TimestampType, + AC.DurationType, + } && return zero(AC.juliatype(t)) + t isa AC.BoolType && return false + t isa AC.Utf8Type && return "" + t isa AC.BinaryType && return UInt8[] + t isa AC.FixedSizeBinaryType && return zeros(UInt8, t.nbytes) + t isa AC.ViewType && return t.utf8 ? "" : UInt8[] + if t isa AC.DecimalType + return t.bits == 32 ? Int32(0) : + t.bits == 64 ? Int64(0) : zeros(UInt8, AC.primwidth(t)) + end + if t isa AC.IntervalType + return t.unit == AC.YEAR_MONTH ? Int32(0) : + t.unit == AC.DAY_TIME ? (days=Int32(0), millis=Int32(0)) : + (months=Int32(0), days=Int32(0), nanos=Int64(0)) + end + t isa Union{AC.ListType,AC.ListViewType} && return Any[] + if t isa AC.FixedSizeListType + length(f.children) == 1 || + throw(ArgumentError("retained fixed-size list $(f.name) needs one child")) + return Any[_retainedplaceholder(f.children[1]) for _ = 1:(t.listsize)] + end + if t isa AC.StructType + return Pair{String,Any}[ + child.name => _retainedplaceholder(child) for child in f.children + ] + end + t isa AC.MapType && return Pair{Any,Any}[] + if t isa AC.RunEndEncodedType + length(f.children) == 2 || throw( + ArgumentError("retained run-end encoded field $(f.name) needs two children"), + ) + return _retainedplaceholder(f.children[2]; forcevalid=forcevalid || !f.nullable) + end + throw( + ArgumentError( + "cannot synthesize hidden child data for retained $(AC.descriptorname(t)) " * + "field $(f.name)", + ), + ) +end + +_retainedcontainer(t::AC.ArrowType) = + t isa Union{ + AC.ListType, + AC.ListViewType, + AC.FixedSizeListType, + AC.StructType, + AC.MapType, + AC.RunEndEncodedType, + } + +"Write values materialized inside a composite, where temporal values stay raw." +function _retainedchildcolumn(f::AC.Field, values, context::_WriterContext) + f.type isa AC.DictionaryType && throw( + ArgumentError( + "nested retained dictionary field $(f.name) cannot be reconstructed " * + "after facade materialization because its pool is unavailable", + ), + ) + all(value -> value === _RETAINED_HIDDEN, values) && + return _constructhiddenpart(f, length(values), context; inactive=true) + _, target, _ = _arrowtypestarget(context.arrowtypes, f) + inferredplan = + target === nothing ? nothing : + _inferredstorageplan(f.name, AbstractVector[values], context; target, retained=f) + inferredplan === nothing || return _constructinferredretainedpart( + f, + inferredplan, + only(inferredplan.parts), + context, + ) + registeredwriter = _registeredwriterplan(context, f, values) + registeredwriter === nothing || return _constructwriterchild( + f, + registeredwriter.input, + values; + routes=registeredwriter.routes, + context, + ) + routed = any(x -> x isa _WriterRoutedUnion, values) + allhiddenunion = f.type isa AC.UnionType && all(x -> x === _RETAINED_HIDDEN, values) + if routed || allhiddenunion + hidden = _writerhiddenstorage(f; inactive=true) + prepared = Any[x === _RETAINED_HIDDEN ? hidden : x for x in values] + return _constructpart(f, prepared; context) + end + if any(x -> x === _RETAINED_HIDDEN, values) + if _retainedcontainer(f.type) + # Preserve hidden state until recursive construction reaches a leaf. + return _constructpart(f, values; context) + end + masked = _constructmaskedleaf(f, values, _NativeMaskedLeafProjection(f, context)) + masked === nothing || return masked + hidden = _writerhiddenstorage(f; inactive=true) + storage = _writerstoragevector(f) + sizehint!(storage, length(values)) + for value in values + push!(storage, value === _RETAINED_HIDDEN ? hidden : value) + end + return _constructwriterstorage(f, storage, context) + end + _retainedcontainer(f.type) && return _constructpart(f, values; context) + v = _retainedtypedvalues(f, values; converted=false) + t = f.type + if t isa AC.DateType || + t isa AC.TimestampType || + t isa AC.TimeType || + t isa AC.DurationType + storage = Union{Missing,Int64}[x === missing ? missing : Int64(x) for x in v] + return _rebuildtemporal(f, storage) + end + return _constructpart(f, v; context) +end + +function _retainedlist(f::AC.Field, v::AbstractVector, context::_WriterContext) + t = f.type::Union{AC.ListType,AC.ListViewType,AC.FixedSizeListType} + length(f.children) == 1 || + throw(ArgumentError("column $(f.name) retained list descriptor needs one child")) + present = _retainedvalidity(f, v) + t isa AC.FixedSizeListType && + all(row -> row === missing || row === _RETAINED_HIDDEN, v) && + return _constructhiddenpart(f, length(v), context) + if t isa AC.FixedSizeListType && _writercompactfixednull(f) + # Hidden rows cost child length only; see _writercompactfixednull. + plain = _writerplainfixednullshape(f) + plain && return _constructplainfixednullpart(f, v, context) + visible = Any[] + for row in v + row === missing && continue + row === _RETAINED_HIDDEN && continue + row isa AbstractVector || throw( + ArgumentError( + "column $(f.name) retained fixed-size list rows must be vectors", + ), + ) + length(row) == t.listsize || throw( + ArgumentError( + "column $(f.name) fixed-size list row has $(length(row)) " * + "values; expected $(t.listsize)", + ), + ) + append!(visible, row) + end + isempty(visible) || _retainedchildcolumn(only(f.children), visible, context) + return _constructfixednullpart(f, present, context) + end + flat = Any[] + if t isa AC.ListType + Offset = t.large ? Int64 : Int32 + offsets = Vector{Offset}(undef, length(v) + 1) + offsets[1] = zero(Offset) + for (i, row) in enumerate(v) + if row !== missing && row !== _RETAINED_HIDDEN + row isa AbstractVector || throw( + ArgumentError("column $(f.name) retained list rows must be vectors"), + ) + length(flat) <= typemax(Offset) - length(row) || throw( + ArgumentError( + "column $(f.name) child data exceeds its retained offset width", + ), + ) + append!(flat, row) + end + offsets[i + 1] = Offset(length(flat)) + end + buffers = AC.BufferSlice[AC._bitmapbuffer(present), AC._databuffer(offsets)] + elseif t isa AC.ListViewType + Offset = t.large ? Int64 : Int32 + offsets = Vector{Offset}(undef, length(v)) + sizes = Vector{Offset}(undef, length(v)) + for (i, row) in enumerate(v) + if row === missing || row === _RETAINED_HIDDEN + offsets[i] = zero(Offset) + sizes[i] = zero(Offset) + continue + end + row isa AbstractVector || throw( + ArgumentError("column $(f.name) retained list-view rows must be vectors"), + ) + length(flat) <= typemax(Offset) - length(row) || throw( + ArgumentError( + "column $(f.name) child data exceeds its retained offset width", + ), + ) + offsets[i] = Offset(length(flat)) + sizes[i] = Offset(length(row)) + append!(flat, row) + end + buffers = AC.BufferSlice[ + AC._bitmapbuffer(present), + AC._databuffer(offsets), + AC._databuffer(sizes), + ] + else + for row in v + if row === missing || row === _RETAINED_HIDDEN + continue + end + row isa AbstractVector || throw( + ArgumentError( + "column $(f.name) retained fixed-size list rows must be vectors", + ), + ) + length(row) == t.listsize || throw( + ArgumentError( + "column $(f.name) fixed-size list row has $(length(row)) " * + "values; expected $(t.listsize)", + ), + ) + end + # Expose hidden child spans on demand. The child constructor still + # emits every descriptor-required physical slot, but it no longer + # allocates one Any reference per hidden slot before doing so. + flat = _MaskedFixedListValues(v, t.listsize) + buffers = AC.BufferSlice[AC._bitmapbuffer(present)] + end + childfield, childdata = _retainedchildcolumn(f.children[1], flat, context) + d = AC.ArrayData( + t, + length(v), + buffers; + children=AC.ArrayData[childdata], + nullcount=count(!, present), + ) + return _retainedfield(f; children=AC.Field[childfield]), d +end + +function _validateretainedstructrows(f::AC.Field, rows::AbstractVector) + nchildren = length(f.children) + for row in rows + row === missing && continue + row === _RETAINED_HIDDEN && continue + row isa AbstractVector || throw( + ArgumentError( + "column $(f.name) retained struct rows must be ordered Pair vectors", + ), + ) + length(row) == nchildren || throw( + ArgumentError( + "column $(f.name) retained struct row has $(length(row)) fields; " * + "expected $nchildren", + ), + ) + firstrowindex = firstindex(row) + for j = 1:nchildren + kv = row[firstrowindex + j - 1] + kv isa Pair || throw( + ArgumentError("column $(f.name) retained struct rows must contain Pairs"), + ) + first(kv) == f.children[j].name || throw( + ArgumentError( + "column $(f.name) retained struct child $j is named " * + "$(repr(first(kv))); expected $(repr(f.children[j].name))", + ), + ) + end + end + return nothing +end + +function _retainedstruct(f::AC.Field, v::AbstractVector, context::_WriterContext) + t = f.type::AC.StructType + present = _retainedvalidity(f, v) + nchildren = length(f.children) + _validateretainedstructrows(f, v) + children = AC.ArrayData[] + childfields = AC.Field[] + for j = 1:nchildren + values = _retainedstructchildvalues(v, j) + cf, cd = _retainedchildcolumn(f.children[j], values, context) + push!(childfields, cf) + push!(children, cd) + end + d = AC.ArrayData( + t, + length(v), + [AC._bitmapbuffer(present)]; + children=children, + nullcount=count(!, present), + ) + return _retainedfield(f; children=childfields), d +end + +function _writermapkeyisless(f::AC.Field, a, b) + try + result = isless(a, b) + result isa Bool || throw( + ArgumentError("Map key comparison for field $(f.name) did not return Bool"), + ) + return result + catch err + err isa Union{InterruptException,OutOfMemoryError} && rethrow() + err isa MethodError || rethrow() + throw( + ArgumentError( + "Map field $(f.name) declares sorted keys, but keys of types " * + "$(typeof(a)) and $(typeof(b)) cannot be ordered", + ), + ) + end +end + +function _validatemapphysicalkeys( + name::String, + keyfield::AC.Field, + keydata::AC.ArrayData, + offsets, +) + plan = _ArrowTypesRoutePlan() + for row = 1:(length(offsets) - 1) + lo = Int(offsets[row]) + 1 + hi = Int(offsets[row + 1]) + seen = Dict{_WriterStorageKey,Nothing}() + for index = lo:hi + value = _arrowtypesroutedvalue(keyfield, keydata, Int64(index), plan, true) + key = _WriterStorageKey(value) + haskey(seen, key) && throw( + ArgumentError( + "column $name row $row has duplicate physical Map key storage", + ), + ) + seen[key] = nothing + end + end + return nothing +end + +function _retainedmap(f::AC.Field, v::AbstractVector, context::_WriterContext) + t = f.type::AC.MapType + length(f.children) == 1 || throw( + ArgumentError("column $(f.name) retained map descriptor needs one entries child"), + ) + entries = f.children[1] + entries.type isa AC.StructType && length(entries.children) == 2 || + throw(ArgumentError("column $(f.name) retained map entries descriptor is invalid")) + present = _retainedvalidity(f, v) + offsets = Vector{Int32}(undef, length(v) + 1) + offsets[1] = 0 + keys = Any[] + values = Any[] + for (i, row) in enumerate(v) + if row !== missing && row !== _RETAINED_HIDDEN + row isa AbstractVector || throw( + ArgumentError("column $(f.name) retained map rows must be Pair vectors"), + ) + length(keys) <= typemax(Int32) - length(row) || throw( + ArgumentError("column $(f.name) map entries exceed the Int32 offset range"), + ) + for kv in row + kv isa Pair || throw( + ArgumentError("column $(f.name) retained map rows must contain Pairs"), + ) + key = first(kv) + key === missing && throw( + ArgumentError("column $(f.name) retained Map keys cannot be missing"), + ) + push!(keys, key) + push!(values, last(kv)) + end + end + offsets[i + 1] = Int32(length(keys)) + end + keyfield, keydata = _retainedchildcolumn(entries.children[1], keys, context) + valuefield, valuedata = _retainedchildcolumn(entries.children[2], values, context) + _validatemapphysicalkeys(f.name, keyfield, keydata, offsets) + if t.keyssorted + for row = 1:length(v) + lo = Int(offsets[row]) + 1 + hi = Int(offsets[row + 1]) + lo >= hi && continue + previous = AC.getvalue(keyfield, keydata, Int64(lo)) + for index = (lo + 1):hi + key = AC.getvalue(keyfield, keydata, Int64(index)) + _writermapkeyisless(f, key, previous) && throw( + ArgumentError( + "column $(f.name) row $row is not sorted by its physical " * + "key storage but its retained Map type declares sorted keys", + ), + ) + previous = key + end + end + end + entriesfield = _retainedfield(entries; children=AC.Field[keyfield, valuefield]) + entriesdata = AC.ArrayData( + entries.type, + length(keys), + [AC.BufferSlice()]; + children=AC.ArrayData[keydata, valuedata], + nullcount=0, + ) + d = AC.ArrayData( + t, + length(v), + [AC._bitmapbuffer(present), AC._databuffer(offsets)]; + children=[entriesdata], + nullcount=count(!, present), + ) + return _retainedfield(f; children=AC.Field[entriesfield]), d +end + +function _retainedree(f::AC.Field, v::AbstractVector, context::_WriterContext) + t = f.type::AC.RunEndEncodedType + length(f.children) == 2 || + throw(ArgumentError("column $(f.name) retained REE descriptor needs two children")) + runfield, valuefield = f.children + runtype = runfield.type + runtype isa AC.IntType || + throw(ArgumentError("column $(f.name) retained REE run ends must be integers")) + RT = AC.juliatype(runtype) + length(v) <= typemax(RT) || + throw(ArgumentError("column $(f.name) length exceeds its retained run-end type")) + runends = RT[] + runvalues = Any[] + for (i, x) in enumerate(v) + # Logical `isequal` may intentionally coarsen a custom type. Only the + # non-overridable `===` relation is safe for physical run coalescing. + if isempty(runvalues) || x !== last(runvalues) + push!(runvalues, x) + push!(runends, RT(i)) + else + runends[end] = RT(i) + end + end + rebuiltvaluefield, valuedata = _retainedchildcolumn(valuefield, runvalues, context) + physical = _writerstoredcolumn(rebuiltvaluefield, valuedata) + if length(physical) > 1 + selected = Int[1] + compactrunends = RT[runends[1]] + previous = _WriterStorageKey(physical[1]) + for i = 2:length(physical) + key = _WriterStorageKey(physical[i]) + if isequal(key, previous) + compactrunends[end] = runends[i] + else + push!(selected, i) + push!(compactrunends, runends[i]) + previous = key + end + end + if length(selected) < length(physical) + rebuiltvaluefield, valuedata = _rebuildwriterstorageselection( + rebuiltvaluefield, + physical, + selected, + context, + ) + runends = compactrunends + end + end + rebuiltrunfield, runenddata = _constructpart(runfield, runends; context) + d = AC.ArrayData( + t, + length(v), + AC.BufferSlice[]; + children=AC.ArrayData[runenddata, valuedata], + nullcount=0, + ) + return _retainedfield(f; children=AC.Field[rebuiltrunfield, rebuiltvaluefield]), d +end + +"Construct a retained dense or sparse Union from writer-side branch routes." +function _constructwriterunion(f::AC.Field, v::AbstractVector, context::_WriterContext) + t = f.type::AC.UnionType + length(t.typeids) == length(f.children) || throw( + ArgumentError("writer-routed Union field $(f.name) has invalid child metadata"), + ) + active = t.mode == AC.SparseMode ? falses(length(f.children)) : nothing + if active !== nothing + for x in v + x isa _WriterRoutedUnion || throw( + ArgumentError( + "column $(f.name) has a retained UnionType whose child type ids " * + "cannot be recovered from materialized facade values", + ), + ) + 1 <= x.child <= length(f.children) || throw( + ArgumentError("writer Union route $(x.child) is outside field $(f.name)"), + ) + active[x.child] = true + end + end + compact = + active === nothing ? nothing : + Bool[active[i] && _writercompactfixednull(f.children[i]) for i in eachindex(active)] + childvalues = AbstractVector[ + compact !== nothing && compact[i] ? Any[] : _writerstoragevector(f.children[i]) + for i in eachindex(f.children) + ] + t.mode == AC.SparseMode && foreach(values -> sizehint!(values, length(v)), childvalues) + typeids = Vector{Int8}(undef, length(v)) + offsets = t.mode == AC.DenseMode ? Vector{Int32}(undef, length(v)) : nothing + # Sparse Union: every active child needs a value in every slot. nothing + # marks children with no rows at all; _RETAINED_HIDDEN marks compacted + # Null-only children, whose hidden rows cost length only. + hidden = if active === nothing || isempty(v) + nothing + else + Any[ + !active[i] ? nothing : + compact[i] ? _RETAINED_HIDDEN : + _writerhiddenstorage(f.children[i]; inactive=true) for + i in eachindex(active) + ] + end + for (i, x) in enumerate(v) + x isa _WriterRoutedUnion || throw( + ArgumentError( + "column $(f.name) has a retained UnionType whose child type ids " * + "cannot be recovered from materialized facade values", + ), + ) + 1 <= x.child <= length(f.children) || + throw(ArgumentError("writer Union route $(x.child) is outside field $(f.name)")) + typeids[i] = t.typeids[x.child] + if t.mode == AC.DenseMode + child = childvalues[x.child] + length(child) <= typemax(Int32) || throw( + ArgumentError( + "writer Union field $(f.name) exceeds the Int32 offset range", + ), + ) + offsets[i] = Int32(length(child)) + _writerpushstorage!(child, x.value, f.children[x.child], x.writertype, context) + else + for childindex in eachindex(childvalues) + active[childindex] || continue + if compact[childindex] + push!( + childvalues[childindex], + childindex == x.child ? x : _RETAINED_HIDDEN, + ) + else + storagevalue = childindex == x.child ? x.value : hidden[childindex] + _writerpushstorage!( + childvalues[childindex], + storagevalue, + f.children[childindex], + x.writertype, + context, + ) + end + end + end + end + childfields = AC.Field[] + childdata = AC.ArrayData[] + for i in eachindex(f.children) + childfield = f.children[i] + rebuiltfield, rebuiltdata = if active !== nothing && !active[i] + _constructhiddenpart(childfield, length(v), context; inactive=true) + elseif compact !== nothing && compact[i] + _constructcompactfixednull(childfield, childvalues[i], context) + else + _constructwriterstorage(childfield, childvalues[i], context) + end + push!(childfields, rebuiltfield) + push!(childdata, rebuiltdata) + end + field = AC.Field( + f.name, + t; + nullable=f.nullable, + metadata=_fieldmetadata(f), + children=childfields, + ) + buffers = + t.mode == AC.DenseMode ? + AC.BufferSlice[AC._databuffer(typeids), AC._databuffer(offsets)] : + AC.BufferSlice[AC._databuffer(typeids)] + data = AC.ArrayData(t, length(v), buffers; children=childdata, nullcount=0) + return field, data +end + +"Construct one part under a retained Field: descriptor, nullability, metadata." +function _constructpart( + f::AC.Field, + v::AbstractVector; + context::_WriterContext=_WriterContext(), +) + _, target, _ = _arrowtypestarget(context.arrowtypes, f) + inferredplan = + target === nothing ? nothing : + _inferredstorageplan(f.name, AbstractVector[v], context; target, retained=f) + inferredplan === nothing || return _constructinferredretainedpart( + f, + inferredplan, + only(inferredplan.parts), + context, + ) + registeredwriter = _registeredwriterplan(context, f, v) + registeredwriter === nothing || return _constructwriterchild( + f, + registeredwriter.input, + v; + routes=registeredwriter.routes, + context, + ) + return _constructretainedpart(f, f.type, v, context) +end + +function _constructretainedpart( + f::AC.Field, + t::AC.NullType, + v::AbstractVector, + ::_WriterContext, +) + eltype(v) === Missing || throw( + ArgumentError( + "column $(f.name) no longer matches its retained Arrow NullType; " * + "give it a Missing element type", + ), + ) + d = AC.ArrayData(t, length(v), AC.BufferSlice[]; nullcount=length(v)) + return AC.Field(f.name, t; nullable=f.nullable, metadata=_fieldmetadata(f)), d +end + +function _checkretainedidentity(f::AC.Field, t::AC.ArrowType, v::AbstractVector) + # Identity FIRST, for every retained field with a known facade type: + # a replaced column is rejected on its declared element type before any + # value is read. + Fp = _facadebasetype(t) + if Fp !== Any + NT = Base.nonmissingtype(eltype(v)) + NT <: Fp || + NT === Union{} || + throw( + ArgumentError( + "column $(f.name) holds $(NT) values, but its retained Arrow " * + "type $(repr(t)) materializes as $(Fp); the column was " * + "replaced with incompatible data", + ), + ) + eltype(v) >: Missing && + !f.nullable && + throw( + ArgumentError( + "column $(f.name) may hold missing values but its retained " * + "field is non-nullable", + ), + ) + end + return Fp +end + +function _constructretainedpart( + f::AC.Field, + t::Union{AC.DateType,AC.TimestampType,AC.TimeType,AC.DurationType}, + v::AbstractVector, + ::_WriterContext, +) + Fp = _checkretainedidentity(f, t, v) + storage = + Fp === Int64 ? Union{Missing,Int64}[x === missing ? missing : Int64(x) for x in v] : + _retainedstorage(t, v, f.name) + return _rebuildtemporal(f, storage) +end + +function _constructretainedpart( + f::AC.Field, + t::Union{AC.Utf8Type,AC.BinaryType}, + v::AbstractVector, + ::_WriterContext, +) + _checkretainedidentity(f, t, v) + return _retainedvarbytes(f, v) +end + +function _constructretainedpart( + f::AC.Field, + t::AC.ViewType, + v::AbstractVector, + ::_WriterContext, +) + _checkretainedidentity(f, t, v) + return _retainedview(f, v) +end + +function _constructretainedpart( + f::AC.Field, + t::AC.FixedSizeBinaryType, + v::AbstractVector, + ::_WriterContext, +) + _checkretainedidentity(f, t, v) + return _retainedfixedbytes(f, v) +end + +function _constructretainedpart( + f::AC.Field, + t::AC.DecimalType, + v::AbstractVector, + ::_WriterContext, +) + _checkretainedidentity(f, t, v) + return _retaineddecimal(f, v) +end + +function _constructretainedpart( + f::AC.Field, + t::AC.IntervalType, + v::AbstractVector, + ::_WriterContext, +) + _checkretainedidentity(f, t, v) + return _retainedinterval(f, v) +end + +function _constructretainedpart( + f::AC.Field, + t::Union{AC.ListType,AC.ListViewType,AC.FixedSizeListType}, + v::AbstractVector, + context::_WriterContext, +) + _checkretainedidentity(f, t, v) + return _retainedlist(f, v, context) +end + +function _constructretainedpart( + f::AC.Field, + t::AC.StructType, + v::AbstractVector, + context::_WriterContext, +) + _checkretainedidentity(f, t, v) + return _retainedstruct(f, v, context) +end + +function _constructretainedpart( + f::AC.Field, + t::AC.MapType, + v::AbstractVector, + context::_WriterContext, +) + _checkretainedidentity(f, t, v) + return _retainedmap(f, v, context) +end + +function _constructretainedpart( + f::AC.Field, + t::AC.RunEndEncodedType, + v::AbstractVector, + context::_WriterContext, +) + _checkretainedidentity(f, t, v) + return _retainedree(f, v, context) +end + +function _constructretainedpart( + f::AC.Field, + t::AC.UnionType, + v::AbstractVector, + context::_WriterContext, +) + _checkretainedidentity(f, t, v) + all(x -> x isa _WriterRoutedUnion, v) && return _constructwriterunion(f, v, context) + throw( + ArgumentError( + "column $(f.name) has a retained UnionType whose child type ids " * + "cannot be recovered from materialized facade values", + ), + ) +end + +function _constructretainedpart( + f::AC.Field, + t::AC.ArrowType, + v::AbstractVector, + context::_WriterContext, +) + _checkretainedidentity(f, t, v) + # Build the column naturally, then impose the retained descriptor: types + # must match and nullability comes from the retained field. Missing values + # under a non-nullable field mean the column was replaced. For lists this + # applies recursively — child names, nullability, metadata, and each + # level's width are part of retained identity. + fn, dn = _constructpart(f.name, v; context) + AC.typeequal(fn.type, t) || throw( + ArgumentError( + "column $(f.name) no longer matches its retained Arrow type " * + "$(repr(t)); it now maps to $(repr(fn.type))", + ), + ) + fn.nullable && + !f.nullable && + AC.nullcount(dn) > 0 && + throw( + ArgumentError( + "column $(f.name) holds missing values but its retained field is " * + "non-nullable", + ), + ) + rebuilt = AC.Field( + f.name, + fn.type; + nullable=f.nullable, + metadata=_fieldmetadata(f), + children=collect(AC.Field, fn.children), + ) + return rebuilt, dn +end + +function _rebuildtemporal(f::AC.Field, storage) + t = f.type + nmissing = count(x -> x === missing, storage) + nmissing > 0 && + !f.nullable && + throw( + ArgumentError( + "column $(f.name) holds missing values but its retained field is " * + "non-nullable", + ), + ) + fld = AC.Field(f.name, t; nullable=f.nullable, metadata=_fieldmetadata(f)) + return fld, _temporaldata(t, storage) +end + +"Merge retained dictionary snapshots without changing the first pool's order." +struct _MergedDictionaryPool{V<:AbstractVector} + values::V + retainedprefix::Int +end + +function _mergeddictpool(poolhints; widen::Bool=false) + hints = Any[pool for pool in poolhints if pool !== nothing] + isempty(hints) && return nothing + values = _mergecategorypools(hints; widen) + return _MergedDictionaryPool(values, length(first(hints))) +end + +# Encoding identity belongs to the lowered Arrow storage domain. These +# wrappers provide structural hashing and equality without calling a logical +# value's overloadable `hash`, `isequal`, `==`, or `isless` methods. Dictionary +# pool dedup, dictionary compaction, and run-end compaction share this rule. +struct _WriterStorageKey{T} + value::T +end + +_writerstorageroute(value::_WriterRoutedUnion) = (value.child, value.value) +_writerstorageroute(value::_ArrowTypesRoutedUnion) = (value.child, value.value) +_writerstorageroute(value) = nothing + +const _WRITER_INTEGER_STORAGE = Union{Int8,Int16,Int32,Int64,UInt8,UInt16,UInt32,UInt64} + +function _writerstorageequal(left, right, depth::Int=0) + depth <= _MAX_WRITER_SCHEMA_DEPTH || + throw(ArgumentError("writer storage value exceeds the supported nesting depth")) + leftroute = _writerstorageroute(left) + rightroute = _writerstorageroute(right) + if leftroute !== nothing || rightroute !== nothing + leftroute === nothing && return false + rightroute === nothing && return false + return first(leftroute) == first(rightroute) && + _writerstorageequal(last(leftroute), last(rightroute), depth + 1) + end + (left === missing || right === missing) && return left === right + (left === nothing || right === nothing) && return left === right + if left isa AbstractString || right isa AbstractString + left isa AbstractString && right isa AbstractString || return false + ncodeunits(left) == ncodeunits(right) || return false + for i = 1:ncodeunits(left) + codeunit(left, i) === codeunit(right, i) || return false + end + return true + end + if left isa AbstractVector || right isa AbstractVector + left isa AbstractVector && right isa AbstractVector || return false + length(left) == length(right) || return false + for i = 1:length(left) + _writerstorageequal(left[i], right[i], depth + 1) || return false + end + return true + end + if left isa Pair || right isa Pair + left isa Pair && right isa Pair || return false + return _writerstorageequal(first(left), first(right), depth + 1) && + _writerstorageequal(last(left), last(right), depth + 1) + end + if left isa NamedTuple || right isa NamedTuple + left isa NamedTuple && right isa NamedTuple || return false + keys(left) === keys(right) || return false + return _writerstorageequal(Tuple(left), Tuple(right), depth + 1) + end + if left isa Tuple || right isa Tuple + left isa Tuple && right isa Tuple || return false + length(left) == length(right) || return false + for i in eachindex(left, right) + _writerstorageequal(left[i], right[i], depth + 1) || return false + end + return true + end + if left isa AbstractFloat || right isa AbstractFloat + # Compare floats by bits: pool compaction must keep 0.0 and -0.0 + # distinct and must treat identical NaN payloads as one category. + typeof(left) === typeof(right) || return false + left isa Float16 && return reinterpret(UInt16, left) === reinterpret(UInt16, right) + left isa Float32 && return reinterpret(UInt32, left) === reinterpret(UInt32, right) + left isa Float64 && return reinterpret(UInt64, left) === reinterpret(UInt64, right) + return false + end + if left isa _WRITER_INTEGER_STORAGE || + right isa _WRITER_INTEGER_STORAGE || + left isa Bool || + right isa Bool || + left isa Char || + right isa Char + return typeof(left) === typeof(right) && left === right + end + typeof(left) === typeof(right) || return false + # Mutable values have no structural identity we may trust; compare by + # reference and hash by objectid, so equal-but-distinct mutables stay + # separate categories. + Base.ismutabletype(typeof(left)) && return left === right + fieldcount(typeof(left)) == fieldcount(typeof(right)) || return false + for i = 1:fieldcount(typeof(left)) + _writerstorageequal(getfield(left, i), getfield(right, i), depth + 1) || + return false + end + return true +end + +function _writerstoragehash(value, seed::UInt, depth::Int=0) + depth <= _MAX_WRITER_SCHEMA_DEPTH || + throw(ArgumentError("writer storage value exceeds the supported nesting depth")) + route = _writerstorageroute(value) + if route !== nothing + h = hash(UInt8(0x01), seed) + h = hash(first(route), h) + return _writerstoragehash(last(route), h, depth + 1) + end + value === missing && return hash(UInt8(0x02), seed) + value === nothing && return hash(UInt8(0x03), seed) + if value isa AbstractString + h = hash(UInt8(0x04), seed) + h = hash(ncodeunits(value), h) + for byte in codeunits(value) + h = hash(byte, h) + end + return h + end + if value isa AbstractVector + h = hash(UInt8(0x05), seed) + h = hash(length(value), h) + for item in value + h = _writerstoragehash(item, h, depth + 1) + end + return h + end + if value isa Pair + h = hash(UInt8(0x06), seed) + h = _writerstoragehash(first(value), h, depth + 1) + return _writerstoragehash(last(value), h, depth + 1) + end + if value isa NamedTuple + h = hash(UInt8(0x07), seed) + h = hash(keys(value), h) + return _writerstoragehash(Tuple(value), h, depth + 1) + end + if value isa Tuple + h = hash(UInt8(0x08), seed) + h = hash(length(value), h) + for item in value + h = _writerstoragehash(item, h, depth + 1) + end + return h + end + if value isa AbstractFloat + # Hash floats by bits, matching _writerstorageequal. + h = hash(typeof(value), hash(UInt8(0x09), seed)) + bits = + value isa Float16 ? reinterpret(UInt16, value) : + value isa Float32 ? reinterpret(UInt32, value) : reinterpret(UInt64, value) + return hash(bits, h) + end + if value isa _WRITER_INTEGER_STORAGE || value isa Bool || value isa Char + return hash(value, hash(typeof(value), hash(UInt8(0x0a), seed))) + end + h = hash(typeof(value), hash(UInt8(0x0b), seed)) + # Mutables hash by objectid, matching _writerstorageequal's `===` rule. + Base.ismutabletype(typeof(value)) && return hash(objectid(value), h) + for i = 1:fieldcount(typeof(value)) + h = _writerstoragehash(getfield(value, i), h, depth + 1) + end + return h +end + +Base.isequal(left::_WriterStorageKey, right::_WriterStorageKey) = + _writerstorageequal(left.value, right.value) +Base.hash(key::_WriterStorageKey, seed::UInt) = _writerstoragehash(key.value, seed) + +"Declared common value type for a newly inferred dictionary pool." +function _dictvaluetype(vals) + T = Union{} + for v in vals + V = Base.nonmissingtype(eltype(v)) + V === Union{} || (T = typejoin(T, V)) + end + return T +end + +"One value pool: keep a retained prefix exactly and append new categories." +function _dictionarypool(vals, retainedpool; writetype=nothing) + if retainedpool === nothing + T = writetype === nothing ? _dictvaluetype(vals) : writetype + T === Union{} && throw( + ArgumentError( + "cannot infer a dictionary value type from empty or all-missing columns", + ), + ) + pool = T[] + else + pool = collect(retainedpool) + end + # Dedup with the non-overloadable structural key `_compactdictionarypool` + # also uses: public-domain values may overload `hash`/`isequal`, so only the + # `_WriterStorageKey` relation (bit-pattern floats, `===`/objectid for + # mutables) may pre-merge categories. Exact value-based dedup of the + # LOWERED categories still happens in `_compactdictionarypool`, in the + # storage domain; this pass only keeps the candidate pool at + # O(distinct-under-that-policy) instead of O(rows). + seen = Dict{_WriterStorageKey,Nothing}() + for x in pool + x === missing && continue + key = _WriterStorageKey(x) + haskey(seen, key) || (seen[key] = nothing) + end + for v in vals, x in v + x === missing && continue + key = _WriterStorageKey(x) + if !haskey(seen, key) + push!(pool, x) + seen[key] = nothing + end + end + return pool +end + +struct _DictionaryEvidence + field::AC.Field + writetype::Type + inferred::Bool + routes::_WriterTypeRoutes + observed::Vector{Type} +end + +function _directunionroutes(T::Type) + routes = _WriterTypeRoutes() + for (i, variant) in enumerate(_checkedwritervariants("writer route type $T", T)) + routes[variant] = i + end + return routes +end + +function _mergewriterroutes(left::_WriterTypeRoutes, right::_WriterTypeRoutes) + routes = copy(left) + for (T, route) in right + prior = get(routes, T, route) + prior == route || throw( + ArgumentError( + "writer type $T has conflicting dictionary routes $prior and $route", + ), + ) + routes[T] = route + end + return routes +end + +function _observeddictionarytypes(name::String, values) + types = Type[] + seen = Base.IdSet{Type}() + for value in values + value === missing && continue + T = typeof(value) + T in seen && continue + length(types) < _MAX_INFERRED_WRITER_TYPES || throw( + ArgumentError( + "dictionary column $name has more than " * + "$_MAX_INFERRED_WRITER_TYPES inferred runtime value types", + ), + ) + push!(seen, T) + push!(types, T) + end + return types +end + +"One Field inferred from the exact runtime writer types in abstract values." +function _observeddictionaryevidence( + name::String, + writetype::Type, + context::_WriterContext, + observed::Vector{Type}=_writertypevariants(writetype), +) + length(observed) <= _MAX_INFERRED_WRITER_TYPES || throw( + ArgumentError( + "dictionary column $name has more than " * + "$_MAX_INFERRED_WRITER_TYPES inferred runtime value types", + ), + ) + variants = _writertypevariants(writetype) + fields = AC.Field[] + for variant in variants + push!(fields, _namedwriterfield(name, _writercandidatefield!(context, variant))) + end + firstfield = first(fields) + if all(field -> _fieldcontractequal(firstfield, field), fields) + # Runtime types are observational evidence, not a declared Union. + # Collapse them when they share one exact Arrow Field. + routes = _WriterTypeRoutes() + for variant in variants + routes[variant] = 0 + end + return _DictionaryEvidence(firstfield, writetype, true, routes, observed) + end + field = _namedwriterfield(name, _writercandidatefield!(context, writetype)) + return _DictionaryEvidence( + field, + writetype, + true, + _directunionroutes(writetype), + observed, + ) +end + +"A fresh dictionary partition's non-null value Field, or no usable evidence." +function _dictionaryevidence( + name::String, + values::AbstractVector, + context::_WriterContext, + observed::Union{Nothing,Vector{Type}}=nothing, +) + declared = Base.nonmissingtype(eltype(values)) + if !isconcretetype(declared) && !(declared isa Union) + # Empty/all-missing abstract vectors cannot identify ArrowTypes metadata + # (for example, a parametric logical tag). Let a later concrete partition + # provide the schema instead of manufacturing an incomplete one. + writetype = _writeruniontype( + observed === nothing ? _observeddictionarytypes(name, values) : observed, + ) + writetype === Union{} && return nothing + return _observeddictionaryevidence(name, writetype, context) + end + # Field evidence is type-derived. Constructing row data here would lower + # every repeated category once for each row before the unique pool exists. + field, _ = _constructpart(name, declared[]; context) + routes = if declared isa Union + _directunionroutes(declared) + else + routes = _WriterTypeRoutes() + routes[declared] = 0 + routes + end + return _DictionaryEvidence(field, declared, false, routes, Type[]) +end + +"Common fresh dictionary value Field and writer type across all partitions." +function _dictionaryvaluefield( + name::String, + values, + column::Symbol, + context::_WriterContext, +) + allobserved = Type[] + observations = Vector{Union{Nothing,Vector{Type}}}(undef, length(values)) + for (partition, part) in enumerate(values) + declared = Base.nonmissingtype(eltype(part)) + if !isconcretetype(declared) && !(declared isa Union) + observations[partition] = _observeddictionarytypes(name, part) + allobserved = _mergewritertypes( + "dictionary column $name", + allobserved, + observations[partition], + _MAX_INFERRED_WRITER_TYPES, + "inferred runtime value types", + ) + else + observations[partition] = nothing + end + end + authority = nothing + for (partition, part) in enumerate(values) + evidence = _dictionaryevidence(name, part, context, observations[partition]) + evidence === nothing && continue + if authority === nothing + authority = evidence + else + observed = _mergewritertypes( + "dictionary column $name", + authority.observed, + evidence.observed, + _MAX_INFERRED_WRITER_TYPES, + "inferred runtime value types", + ) + variants = _mergewritertypes( + "dictionary column $name", + _writertypevariants(authority.writetype), + _writertypevariants(evidence.writetype), + _MAX_WRITER_UNION_BRANCHES, + "writer types", + ) + writetype = _writeruniontype(variants) + if authority.inferred && evidence.inferred + authority = _observeddictionaryevidence(name, writetype, context, observed) + continue + end + _checkpartitionfield(authority.field, evidence.field, partition, column) + authority = _DictionaryEvidence( + authority.field, + writetype, + false, + _mergewriterroutes(authority.routes, evidence.routes), + observed, + ) + end + end + return authority +end + +"First pool position for each non-null value; duplicate categories stay intact." +function _dictionarylookup(pool) + # Same `_WriterStorageKey` policy as `_dictionarypool`: no user-overloadable + # `hash`/`isequal` participates in encoding identity. + lookup = Dict{_WriterStorageKey,Int64}() + for (i, x) in enumerate(pool) + x === missing && continue + key = _WriterStorageKey(x) + haskey(lookup, key) || (lookup[key] = Int64(i - 1)) + end + return lookup +end + +function _dictionaryindices(v, lookup, missingindex=nothing) + return Union{Missing,Int64}[ + x === missing ? (missingindex === nothing ? missing : missingindex) : + lookup[_WriterStorageKey(x)] for x in v + ] +end + +_writerstoredvalue(value::_ArrowTypesRoutedNull) = missing +function _writerstoredvalue(value::_ArrowTypesRoutedUnion) + return _WriterRoutedUnion( + value.child, + _writerstoredvalue(value.value), + typeof(value.value), + ) +end +function _writerstoredvalue(value::Pair{A,B}) where {A,B} + firstvalue = _writerstoredvalue(first(value)) + lastvalue = _writerstoredvalue(last(value)) + return firstvalue isa A && lastvalue isa B ? Pair{A,B}(firstvalue, lastvalue) : + firstvalue => lastvalue +end +function _writerstoredvalue(value::AbstractVector{T}) where {T} + stored = Any[_writerstoredvalue(item) for item in value] + return T === Any || !all(item -> item isa T, stored) ? stored : collect(T, stored) +end +_writerstoredvalue(value) = value + +function _writerstoredcolumn(f::AC.Field, d::AC.ArrayData) + plan = _ArrowTypesRoutePlan() + return Any[_arrowtypesroutedvalue(f, d, Int64(i), plan, true) for i = 1:(d.len)] +end + +function _rebuildwriterstorageselection( + f::AC.Field, + physical::Vector{Any}, + selected::Vector{Int}, + context::_WriterContext, +) + storage = _writerstoragevector(f) + sizehint!(storage, length(selected)) + for i in selected + value = _writerstoredvalue(physical[i]) + _writerpushstorage!(storage, value, f, typeof(value), context) + end + return _constructwriterstorage(f, storage, context) +end + +"Compact candidate categories by exact lowered storage identity." +function _compactdictionarypool( + f::AC.Field, + d::AC.ArrayData, + retainedprefix::Int, + context::_WriterContext, +) + 0 <= retainedprefix <= d.len || + throw(ArgumentError("retained dictionary prefix is outside the candidate pool")) + physical = _writerstoredcolumn(f, d) + remap = Vector{Int64}(undef, length(physical)) + selected = Int[] + seen = Dict{_WriterStorageKey,Int64}() + for i in eachindex(physical) + key = _WriterStorageKey(physical[i]) + if i <= retainedprefix + index = Int64(length(selected)) + push!(selected, i) + remap[i] = index + haskey(seen, key) || (seen[key] = index) + continue + end + index = get(seen, key, Int64(-1)) + if index < 0 + index = Int64(length(selected)) + push!(selected, i) + seen[key] = index + end + remap[i] = index + end + length(selected) == length(physical) && return f, d, remap + + compactfield, compactdata = + _rebuildwriterstorageselection(f, physical, selected, context) + return compactfield, compactdata, remap +end + +function _remapdictionaryindices(indices, remap::Vector{Int64}) + return Union{Missing,Int64}[ + index === missing ? missing : remap[Int(index) + 1] for index in indices + ] +end + +"One shared-pool dictionary batch: identical pool OBJECT across batches." +function _dictbatch(fld::AC.Field, indices::Vector, pool_d::AC.ArrayData) + t = fld.type::AC.DictionaryType + IT = AC.juliatype(t.indextype) + present = [x !== missing for x in indices] + inds = IT[x === missing ? zero(IT) : IT(x) for x in indices] + nc = count(!, present) + d = AC.ArrayData( + t, + length(indices), + [AC._bitmapbuffer(present), AC._databuffer(inds)]; + dictionary=pool_d, + nullcount=nc, + ) + return d +end + +"Field + first-batch data for a dictionary column under a RETAINED type." +function _retaineddictfromdata( + rf::AC.Field, + valuefield::AC.Field, + valuedata::AC.ArrayData, + poollength::Int, + firstidx::Vector, + name::String, +) + t = rf.type::AC.DictionaryType + AC.typeequal(valuefield.type, t.valuetype) || throw( + ArgumentError( + "column $name pool maps to $(summary(valuefield.type)) but the retained " * + "dictionary value type is $(summary(t.valuetype))", + ), + ) + IT = AC.juliatype(t.indextype) + poollength - 1 <= typemax(IT) || throw( + ArgumentError( + "column $name pool of $poollength values exceeds the retained " * + "$(summary(t.indextype)) index range", + ), + ) + field = AC.Field( + name, + t; + nullable=rf.nullable, + metadata=_fieldmetadata(rf), + children=collect(AC.Field, valuefield.children), + ) + return field, _dictbatch(field, firstidx, valuedata) +end + +_withfieldmetadata(f::AC.Field, ::Nothing) = f +function _withfieldmetadata(f::AC.Field, metadata::Vector{Pair{String,String}}) + # Explicit application metadata augments retained/extension metadata. + # A generated ArrowTypes label is authoritative for the values that were + # lowered; accepting a conflicting explicit label would make them decode + # as a different logical type. + merged = _mergemetapairs(_fieldmetadata(f), metadata; protectextension=true) + return AC.Field( + f.name, + f.type; + nullable=f.nullable, + metadata=merged, + children=collect(AC.Field, f.children), + ) +end + +""" +Construct one logical column across all table partitions. + +This is the sole column-policy seam used by the write facade. It owns fresh +inference, retained reconstruction, ArrowTypes lowering, dictionary pooling, +partition agreement, and field metadata. The returned data vector is in the +same order as `parts`. + +`retained` is the Field a rewrite must reproduce; `poolhints` carries one +retained dictionary-pool snapshot (or nothing) per partition, enforced below; +`metadata` is explicit application metadata merged onto the final Field. +""" +function _constructcolumn( + name::Symbol, + parts::Vector{AbstractVector}; + retained::Union{Nothing,AC.Field}=nothing, + poolhints::Vector{Any}=Any[nothing for _ in parts], + metadata::Union{Nothing,Vector{Pair{String,String}}}=nothing, +) + isempty(parts) && throw(ArgumentError("column $name has no partitions")) + length(poolhints) == length(parts) || throw( + ArgumentError( + "column $name has $(length(poolhints)) dictionary-pool hints for " * + "$(length(parts)) partitions", + ), + ) + fieldname = String(name) + nparts = length(parts) + data = Vector{AC.ArrayData}(undef, nparts) + context = _WriterContext(fieldname) + dictintent = + (retained !== nothing && retained.type isa AC.DictionaryType) || + any(part isa DictEncode for part in parts) + + local field::AC.Field + if dictintent + values = AbstractVector[ + part isa DictEncode ? (part::DictEncode).data : part for part in parts + ] + retainedvaluefield = nothing + retainedtarget = nothing + if retained !== nothing + retainedvaluefield, retainedtarget = + _registereddictionaryfield(context, retained) + retainedtype = retained.type::AC.DictionaryType + retainedvaluefield === nothing && + retainedtype.valuetype isa AC.NullType && + retained.nullable && + _arrowtypesextension(retained) !== nothing && + throw( + ArgumentError( + "column $name has a nullable retained Dictionary with " * + "an unregistered extension; materialization cannot " * + "distinguish a valid null-pool index from an outer null index", + ), + ) + # Table materialization already lifts retained pool snapshots + # through the dictionary value Field. They stay in the same + # public domain as row values until the merged pool is lowered + # exactly once through `_constructpart` on that value Field below. + end + + inferreddeclaration = if retained === nothing + _inferredstoragedeclaration(values, context) + elseif retainedtarget === nothing + nothing + else + _inferredstoragedeclaration(values, context; target=retainedtarget) + end + freshevidence = + retained === nothing && inferreddeclaration === nothing ? + _dictionaryvaluefield(fieldname, values, name, context) : nothing + freshvaluefield = freshevidence === nothing ? nothing : freshevidence.field + freshwritetype = freshevidence === nothing ? nothing : freshevidence.writetype + freshroutes = freshevidence === nothing ? nothing : freshevidence.routes + widenpool = if inferreddeclaration !== nothing + true + elseif retained === nothing + freshvaluefield !== nothing && freshvaluefield.type isa AC.UnionType + else + (retained.type::AC.DictionaryType).valuetype isa AC.UnionType + end + mergedpool = _mergeddictpool(poolhints; widen=widenpool) + poolhint = mergedpool === nothing ? nothing : mergedpool.values + if retained !== nothing + (retained.type::AC.DictionaryType).valuetype isa AC.UnionType && + retainedvaluefield === nothing && + throw( + ArgumentError( + "column $name has a retained UnionType whose child type ids " * + "cannot be recovered from materialized facade values", + ), + ) + for valuespart in values + T = Base.nonmissingtype(eltype(valuespart)) + if retainedvaluefield === nothing + publictype = _facadebasetype(retained.type) + publictype !== Any && + !(T <: publictype) && + T !== Union{} && + throw( + ArgumentError( + "column $name holds $T values, but its retained " * + "dictionary materializes as $publictype; the column " * + "was replaced with incompatible data", + ), + ) + else + islogical = + T <: retainedtarget || ( + T === Any && + all(x -> x === missing || x isa retainedtarget, valuespart) + ) + islogical || throw( + ArgumentError( + "column $name holds $T values, but its retained " * + "dictionary materializes as $retainedtarget; the column " * + "was replaced with incompatible data", + ), + ) + end + # A non-nullable dictionary can materialize missing through a + # valid index into a null pool entry. The retained pool keeps + # that distinct from a missing index. + any(ismissing, valuespart) && + !retained.nullable && + (poolhint === nothing || !any(ismissing, poolhint)) && + throw( + ArgumentError( + "column $name may hold missing values but its retained " * + "dictionary field is non-nullable", + ), + ) + end + (retained.type::AC.DictionaryType).ordered && + poolhint === nothing && + throw( + ArgumentError( + "column $name has an ordered retained dictionary, but " * + "its original category pool is unavailable", + ), + ) + end + + poolwritetype = + inferreddeclaration === nothing ? freshwritetype : first(inferreddeclaration) + pool = _dictionarypool(values, poolhint; writetype=poolwritetype) + inferredplan = + inferreddeclaration === nothing ? nothing : + _inferredstorageplan( + fieldname, + AbstractVector[pool], + context; + target=first(inferreddeclaration), + retained=retainedvaluefield, + ) + lookup = _dictionarylookup(pool) + missingindex = + retained !== nothing && !retained.nullable ? findfirst(ismissing, pool) : + nothing + missingindex === nothing || (missingindex = Int64(missingindex - 1)) + candidateindices = Vector{Union{Missing,Int64}}[ + _dictionaryindices(part, lookup, missingindex) for part in values + ] + candidatevaluefield, candidatevaluedata = if inferredplan !== nothing + if retained === nothing + _constructinferredfreshpart( + fieldname, + inferredplan, + only(inferredplan.parts), + context, + ) + else + _constructinferredretainedpart( + retainedvaluefield, + inferredplan, + only(inferredplan.parts), + context, + ) + end + elseif retained === nothing + _constructnewdictpooldata( + fieldname, + pool; + valuefield=freshvaluefield, + writetype=freshwritetype, + routes=freshroutes, + context, + ) + else + retainedtype = retained.type::AC.DictionaryType + constructionfield = + retainedvaluefield === nothing ? + AC.dictvaluefield(retained, retainedtype) : retainedvaluefield + _constructpart(constructionfield, pool; context) + end + retainedprefix = mergedpool === nothing ? 0 : mergedpool.retainedprefix + valuefield, valuedata, remap = _compactdictionarypool( + candidatevaluefield, + candidatevaluedata, + retainedprefix, + context, + ) + indices = Vector{Union{Missing,Int64}}[ + _remapdictionaryindices(part, remap) for part in candidateindices + ] + if retained === nothing + valuedata.len - 1 <= typemax(Int32) || throw( + ArgumentError("column $name dictionary exceeds the Int32 index range"), + ) + field, data[1] = _newdictfromdata( + fieldname, + valuefield, + valuedata, + indices[1]; + nullable=any(v -> Missing <: eltype(v), values), + ) + else + field, data[1] = _retaineddictfromdata( + retained, + valuefield, + valuedata, + Int(valuedata.len), + indices[1], + fieldname, + ) + end + pooldata = data[1].dictionary::AC.ArrayData + for partition = 2:nparts + data[partition] = _dictbatch(field, indices[partition], pooldata) + end + else + narrowedparts = + retained === nothing ? _narrowabstractwriterparts(fieldname, parts, context) : + nothing + constructionparts = narrowedparts === nothing ? parts : narrowedparts + inferredplan = if retained === nothing + _inferredstorageplan(fieldname, constructionparts, context) + else + _, target, _ = _arrowtypestarget(context.arrowtypes, retained) + target === nothing ? nothing : + _inferredstorageplan(fieldname, constructionparts, context; target, retained) + end + registeredwriter = + retained === nothing || inferredplan !== nothing ? nothing : + _registeredwritercolumnplan(context, retained, constructionparts) + for partition = 1:nparts + partitionfield, data[partition] = if inferredplan !== nothing + if retained === nothing + _constructinferredfreshpart( + fieldname, + inferredplan, + inferredplan.parts[partition], + context, + ) + else + _constructinferredretainedpart( + retained, + inferredplan, + inferredplan.parts[partition], + context, + ) + end + elseif retained === nothing + # A successful column-scoped narrowing pass already saw every + # partition. Replanning each partition from its local values + # can split one declared nullable Tuple child into incompatible + # Int and Null Fields. + _constructpart( + fieldname, + constructionparts[partition]; + context, + narrowabstract=narrowedparts === nothing, + ) + elseif registeredwriter === nothing + _constructpart(retained, constructionparts[partition]; context) + else + _constructwriterchild( + retained, + registeredwriter.input, + constructionparts[partition]; + routes=registeredwriter.routes, + context, + ) + end + if partition == 1 + field = partitionfield + else + _checkpartitionfield(field, partitionfield, partition, name) + end + end + end + return _withfieldmetadata(field, metadata), data +end diff --git a/src/eltypes.jl b/src/eltypes.jl deleted file mode 100644 index 52dbb809..00000000 --- a/src/eltypes.jl +++ /dev/null @@ -1,578 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Given a flatbuffers metadata type definition (a Field instance from Schema.fbs), -translate to the appropriate Julia storage eltype -""" -function juliaeltype end - -finaljuliatype(T) = T -finaljuliatype(::Type{Missing}) = Missing -finaljuliatype(::Type{Union{T,Missing}}) where {T} = Union{Missing,finaljuliatype(T)} - -""" -Given a FlatBuffers.Builder and a Julia column or column eltype, -Write the field.type flatbuffer definition of the eltype -""" -function arrowtype end - -arrowtype(b, col::AbstractVector{T}) where {T} = arrowtype(b, maybemissing(T)) -arrowtype(b, col::DictEncoded) = arrowtype(b, col.encoding.data) -arrowtype(b, col::Compressed) = arrowtype(b, col.data) - -function juliaeltype(f::Meta.Field, ::Nothing, convert::Bool) - T = juliaeltype(f, convert) - return convert ? finaljuliatype(T) : T -end - -function juliaeltype(f::Meta.Field, meta::AbstractDict{String,String}, convert::Bool) - TT = juliaeltype(f, convert) - !convert && return TT - T = finaljuliatype(TT) - if haskey(meta, "ARROW:extension:name") - typename = meta["ARROW:extension:name"] - metadata = get(meta, "ARROW:extension:metadata", "") - JT = ArrowTypes.JuliaType(Val(Symbol(typename)), maybemissing(TT), metadata) - if JT !== nothing - return f.nullable ? Union{JT,Missing} : JT - else - @warn "unsupported ARROW:extension:name type: \"$typename\", arrow type = $TT" maxlog = - 1 _id = hash((:juliaeltype, typename, TT)) - end - end - return something(TT, T) -end - -function juliaeltype(f::Meta.Field, convert::Bool) - T = juliaeltype(f, f.type, convert) - return f.nullable ? Union{T,Missing} : T -end - -juliaeltype(f::Meta.Field, ::Meta.Null, convert) = Missing - -function arrowtype(b, ::Type{Missing}) - Meta.nullStart(b) - return Meta.Null, Meta.nullEnd(b), nothing -end - -function juliaeltype(f::Meta.Field, int::Meta.Int, convert) - if int.is_signed - if int.bitWidth == 8 - Int8 - elseif int.bitWidth == 16 - Int16 - elseif int.bitWidth == 32 - Int32 - elseif int.bitWidth == 64 - Int64 - elseif int.bitWidth == 128 - Int128 - else - error("$int is not valid arrow type metadata") - end - else - if int.bitWidth == 8 - UInt8 - elseif int.bitWidth == 16 - UInt16 - elseif int.bitWidth == 32 - UInt32 - elseif int.bitWidth == 64 - UInt64 - elseif int.bitWidth == 128 - UInt128 - else - error("$int is not valid arrow type metadata") - end - end -end - -function arrowtype(b, ::Type{T}) where {T<:Integer} - Meta.intStart(b) - Meta.intAddBitWidth(b, Int32(8 * sizeof(T))) - Meta.intAddIsSigned(b, T <: Signed) - return Meta.Int, Meta.intEnd(b), nothing -end - -# primitive types -function juliaeltype(f::Meta.Field, fp::Meta.FloatingPoint, convert) - if fp.precision == Meta.Precision.HALF - Float16 - elseif fp.precision == Meta.Precision.SINGLE - Float32 - elseif fp.precision == Meta.Precision.DOUBLE - Float64 - end -end - -function arrowtype(b, ::Type{T}) where {T<:AbstractFloat} - Meta.floatingPointStart(b) - Meta.floatingPointAddPrecision( - b, - T === Float16 ? Meta.Precision.HALF : - T === Float32 ? Meta.Precision.SINGLE : Meta.Precision.DOUBLE, - ) - return Meta.FloatingPoint, Meta.floatingPointEnd(b), nothing -end - -juliaeltype(f::Meta.Field, b::Union{Meta.Utf8,Meta.LargeUtf8,Meta.Utf8View}, convert) = - String - -datasizeof(x) = sizeof(x) -datasizeof(x::AbstractVector) = sum(datasizeof, x) - -juliaeltype( - f::Meta.Field, - b::Union{Meta.Binary,Meta.LargeBinary,Meta.BinaryView}, - convert, -) = Base.CodeUnits - -juliaeltype(f::Meta.Field, x::Meta.FixedSizeBinary, convert) = - NTuple{Int(x.byteWidth),UInt8} - -# arggh! -Base.write(io::IO, x::NTuple{N,T}) where {N,T} = sum(y -> Base.write(io, y), x) - -juliaeltype(f::Meta.Field, x::Meta.Bool, convert) = Bool - -function arrowtype(b, ::Type{Bool}) - Meta.boolStart(b) - return Meta.Bool, Meta.boolEnd(b), nothing -end - -struct Decimal{P,S,T} - value::T # only Int128 or Int256 -end - -Base.zero(::Type{Decimal{P,S,T}}) where {P,S,T} = Decimal{P,S,T}(T(0)) -==(a::Decimal{P,S,T}, b::Decimal{P,S,T}) where {P,S,T} = ==(a.value, b.value) -Base.isequal(a::Decimal{P,S,T}, b::Decimal{P,S,T}) where {P,S,T} = isequal(a.value, b.value) - -function juliaeltype(f::Meta.Field, x::Meta.Decimal, convert) - return Decimal{x.precision,x.scale,x.bitWidth == 256 ? Int256 : Int128} -end - -ArrowTypes.ArrowKind(::Type{<:Decimal}) = PrimitiveKind() - -function arrowtype(b, ::Type{Decimal{P,S,T}}) where {P,S,T} - Meta.decimalStart(b) - Meta.decimalAddPrecision(b, Int32(P)) - Meta.decimalAddScale(b, Int32(S)) - Meta.decimalAddBitWidth(b, Int32(T == Int256 ? 256 : 128)) - return Meta.Decimal, Meta.decimalEnd(b), nothing -end - -Base.write(io::IO, x::Decimal) = Base.write(io, x.value) - -abstract type ArrowTimeType end -Base.write(io::IO, x::ArrowTimeType) = Base.write(io, x.x) -ArrowTypes.ArrowKind(::Type{<:ArrowTimeType}) = PrimitiveKind() - -struct Date{U,T} <: ArrowTimeType - x::T -end - -const DATE = Date{Meta.DateUnit.DAY,Int32} -Base.zero(::Type{Date{U,T}}) where {U,T} = Date{U,T}(T(0)) -storagetype(::Type{Date{U,T}}) where {U,T} = T -bitwidth(x::Meta.DateUnit.T) = x == Meta.DateUnit.DAY ? Int32 : Int64 -Date{Meta.DateUnit.DAY}(days) = DATE(Int32(days)) -Date{Meta.DateUnit.MILLISECOND}(ms) = Date{Meta.DateUnit.MILLISECOND,Int64}(Int64(ms)) - -juliaeltype(f::Meta.Field, x::Meta.Date, convert) = Date{x.unit,bitwidth(x.unit)} -finaljuliatype(::Type{DATE}) = Dates.Date -Base.convert(::Type{Dates.Date}, x::DATE) = - Dates.Date(Dates.UTD(Int64(x.x + UNIX_EPOCH_DATE))) -finaljuliatype(::Type{Date{Meta.DateUnit.MILLISECOND,Int64}}) = Dates.DateTime -Base.convert(::Type{Dates.DateTime}, x::Date{Meta.DateUnit.MILLISECOND,Int64}) = - Dates.DateTime(Dates.UTM(Int64(x.x + UNIX_EPOCH_DATETIME))) - -function arrowtype(b, ::Type{Date{U,T}}) where {U,T} - Meta.dateStart(b) - Meta.dateAddUnit(b, U) - return Meta.Date, Meta.dateEnd(b), nothing -end - -const UNIX_EPOCH_DATE = Dates.value(Dates.Date(1970)) -Base.convert(::Type{DATE}, x::Dates.Date) = DATE(Int32(Dates.value(x) - UNIX_EPOCH_DATE)) - -const UNIX_EPOCH_DATETIME = Dates.value(Dates.DateTime(1970)) -Base.convert(::Type{Date{Meta.DateUnit.MILLISECOND,Int64}}, x::Dates.DateTime) = - Date{Meta.DateUnit.MILLISECOND,Int64}(Int64(Dates.value(x) - UNIX_EPOCH_DATETIME)) - -ArrowTypes.ArrowType(::Type{Dates.Date}) = DATE -ArrowTypes.toarrow(x::Dates.Date) = convert(DATE, x) -const DATE_SYMBOL = Symbol("JuliaLang.Date") -ArrowTypes.arrowname(::Type{Dates.Date}) = DATE_SYMBOL -ArrowTypes.JuliaType(::Val{DATE_SYMBOL}, S) = Dates.Date -ArrowTypes.fromarrow(::Type{Dates.Date}, x::DATE) = convert(Dates.Date, x) -ArrowTypes.default(::Type{Dates.Date}) = Dates.Date(1, 1, 1) - -struct Time{U,T} <: ArrowTimeType - x::T -end - -Base.zero(::Type{Time{U,T}}) where {U,T} = Time{U,T}(T(0)) -const TIME = Time{Meta.TimeUnit.NANOSECOND,Int64} - -bitwidth(x::Meta.TimeUnit.T) = - x == Meta.TimeUnit.SECOND || x == Meta.TimeUnit.MILLISECOND ? Int32 : Int64 -Time{U}(x) where {U<:Meta.TimeUnit.T} = Time{U,bitwidth(U)}(bitwidth(U)(x)) -storagetype(::Type{Time{U,T}}) where {U,T} = T -juliaeltype(f::Meta.Field, x::Meta.Time, convert) = Time{x.unit,bitwidth(x.unit)} -finaljuliatype(::Type{<:Time}) = Dates.Time -periodtype(U::Meta.TimeUnit.T) = - U === Meta.TimeUnit.SECOND ? Dates.Second : - U === Meta.TimeUnit.MILLISECOND ? Dates.Millisecond : - U === Meta.TimeUnit.MICROSECOND ? Dates.Microsecond : Dates.Nanosecond -Base.convert(::Type{Dates.Time}, x::Time{U,T}) where {U,T} = - Dates.Time(Dates.Nanosecond(Dates.tons(periodtype(U)(x.x)))) - -function arrowtype(b, ::Type{Time{U,T}}) where {U,T} - Meta.timeStart(b) - Meta.timeAddUnit(b, U) - Meta.timeAddBitWidth(b, Int32(8 * sizeof(T))) - return Meta.Time, Meta.timeEnd(b), nothing -end - -Base.convert(::Type{TIME}, x::Dates.Time) = TIME(Dates.value(x)) - -ArrowTypes.ArrowType(::Type{Dates.Time}) = TIME -ArrowTypes.toarrow(x::Dates.Time) = convert(TIME, x) -const TIME_SYMBOL = Symbol("JuliaLang.Time") -ArrowTypes.arrowname(::Type{Dates.Time}) = TIME_SYMBOL -ArrowTypes.JuliaType(::Val{TIME_SYMBOL}, S) = Dates.Time -ArrowTypes.fromarrow(::Type{Dates.Time}, x::Arrow.Time) = convert(Dates.Time, x) -ArrowTypes.default(::Type{Dates.Time}) = Dates.Time(1, 1, 1) - -struct Timestamp{U,TZ} <: ArrowTimeType - x::Int64 -end - -Base.zero(::Type{Timestamp{U,T}}) where {U,T} = Timestamp{U,T}(Int64(0)) - -function juliaeltype(f::Meta.Field, x::Meta.Timestamp, convert) - return Timestamp{x.unit,x.timezone === nothing ? nothing : Symbol(x.timezone)} -end - -const DATETIME = Timestamp{Meta.TimeUnit.MILLISECOND,nothing} - -finaljuliatype(::Type{Timestamp{U,TZ}}) where {U,TZ} = ZonedDateTime -finaljuliatype(::Type{Timestamp{U,nothing}}) where {U} = DateTime - -@noinline warntimestamp(U, T) = - @warn "automatically converting Arrow.Timestamp with precision = $U to `$T` which only supports millisecond precision; conversion may be lossy; to avoid converting, pass `Arrow.Table(source; convert=false)" maxlog = - 1 _id = hash((:warntimestamp, U, T)) - -function Base.convert(::Type{ZonedDateTime}, x::Timestamp{U,TZ}) where {U,TZ} - (U === Meta.TimeUnit.MICROSECOND || U == Meta.TimeUnit.NANOSECOND) && - warntimestamp(U, ZonedDateTime) - return ZonedDateTime( - Dates.DateTime( - Dates.UTM(Int64(Dates.toms(periodtype(U)(x.x)) + UNIX_EPOCH_DATETIME)), - ), - TimeZone(String(TZ)); - from_utc=true, - ) -end - -function Base.convert(::Type{DateTime}, x::Timestamp{U,nothing}) where {U} - (U === Meta.TimeUnit.MICROSECOND || U == Meta.TimeUnit.NANOSECOND) && - warntimestamp(U, DateTime) - return Dates.DateTime( - Dates.UTM(Int64(Dates.toms(periodtype(U)(x.x)) + UNIX_EPOCH_DATETIME)), - ) -end - -Base.convert(::Type{Timestamp{Meta.TimeUnit.MILLISECOND,TZ}}, x::ZonedDateTime) where {TZ} = - Timestamp{Meta.TimeUnit.MILLISECOND,TZ}( - Int64(Dates.value(DateTime(x, UTC)) - UNIX_EPOCH_DATETIME), - ) -Base.convert(::Type{Timestamp{Meta.TimeUnit.MILLISECOND,nothing}}, x::DateTime) = - Timestamp{Meta.TimeUnit.MILLISECOND,nothing}( - Int64(Dates.value(x) - UNIX_EPOCH_DATETIME), - ) - -function arrowtype(b, ::Type{Timestamp{U,TZ}}) where {U,TZ} - tz = TZ !== nothing ? FlatBuffers.createstring!(b, String(TZ)) : FlatBuffers.UOffsetT(0) - Meta.timestampStart(b) - Meta.timestampAddUnit(b, U) - Meta.timestampAddTimezone(b, tz) - return Meta.Timestamp, Meta.timestampEnd(b), nothing -end - -ArrowTypes.ArrowType(::Type{Dates.DateTime}) = DATETIME -ArrowTypes.toarrow(x::Dates.DateTime) = convert(DATETIME, x) -const DATETIME_SYMBOL = Symbol("JuliaLang.DateTime") -ArrowTypes.arrowname(::Type{Dates.DateTime}) = DATETIME_SYMBOL -ArrowTypes.JuliaType(::Val{DATETIME_SYMBOL}, S) = Dates.DateTime -ArrowTypes.fromarrow(::Type{Dates.DateTime}, x::Timestamp) = convert(Dates.DateTime, x) -ArrowTypes.fromarrow(::Type{Dates.DateTime}, x::Date{Meta.DateUnit.MILLISECOND,Int64}) = - convert(Dates.DateTime, x) -ArrowTypes.default(::Type{Dates.DateTime}) = Dates.DateTime(1, 1, 1, 1, 1, 1) - -ArrowTypes.ArrowType(::Type{ZonedDateTime}) = Timestamp -ArrowTypes.toarrow(x::ZonedDateTime) = - convert(Timestamp{Meta.TimeUnit.MILLISECOND,Symbol(x.timezone)}, x) -const ZONEDDATETIME_SYMBOL = Symbol("JuliaLang.ZonedDateTime-UTC") -ArrowTypes.arrowname(::Type{ZonedDateTime}) = ZONEDDATETIME_SYMBOL -ArrowTypes.JuliaType(::Val{ZONEDDATETIME_SYMBOL}, S) = ZonedDateTime -ArrowTypes.fromarrow(::Type{ZonedDateTime}, x::Timestamp) = convert(ZonedDateTime, x) -ArrowTypes.default(::Type{TimeZones.ZonedDateTime}) = - TimeZones.ZonedDateTime(1, 1, 1, 1, 1, 1, TimeZones.tz"UTC") - -# Backwards compatibility: older versions of Arrow saved ZonedDateTime's with this metdata: -const OLD_ZONEDDATETIME_SYMBOL = Symbol("JuliaLang.ZonedDateTime") -# and stored the local time instead of the UTC time. -struct LocalZonedDateTime end -ArrowTypes.JuliaType(::Val{OLD_ZONEDDATETIME_SYMBOL}, S) = LocalZonedDateTime -function ArrowTypes.fromarrow(::Type{LocalZonedDateTime}, x::Timestamp{U,TZ}) where {U,TZ} - (U === Meta.TimeUnit.MICROSECOND || U == Meta.TimeUnit.NANOSECOND) && - warntimestamp(U, ZonedDateTime) - return ZonedDateTime( - Dates.DateTime( - Dates.UTM(Int64(Dates.toms(periodtype(U)(x.x)) + UNIX_EPOCH_DATETIME)), - ), - TimeZone(String(TZ)), - ) -end - -""" - Arrow.ToTimestamp(x::AbstractVector{ZonedDateTime}) - -Wrapper array that provides a more efficient encoding of `ZonedDateTime` elements to the arrow format. In the arrow format, -timestamp columns with timezone information are encoded as the arrow equivalent of a Julia type parameter, meaning an entire column -_should_ have elements all with the same timezone. If a `ZonedDateTime` column is passed to `Arrow.write`, for correctness, it must -scan each element to check each timezone. `Arrow.ToTimestamp` provides a "bypass" of this process by encoding the timezone of the -first element of the `AbstractVector{ZonedDateTime}`, which in turn allows `Arrow.write` to avoid costly checking/conversion and -can encode the `ZonedDateTime` as `Arrow.Timestamp` directly. -""" -struct ToTimestamp{A,TZ} <: AbstractVector{Timestamp{Meta.TimeUnit.MILLISECOND,TZ}} - data::A # AbstractVector{ZonedDateTime} -end - -ToTimestamp(x::A) where {A<:AbstractVector{ZonedDateTime}} = - ToTimestamp{A,Symbol(x[1].timezone)}(x) -Base.IndexStyle(::Type{<:ToTimestamp}) = Base.IndexLinear() -Base.size(x::ToTimestamp) = (length(x.data),) -Base.eltype(::Type{ToTimestamp{A,TZ}}) where {A,TZ} = - Timestamp{Meta.TimeUnit.MILLISECOND,TZ} -Base.getindex(x::ToTimestamp{A,TZ}, i::Integer) where {A,TZ} = - convert(Timestamp{Meta.TimeUnit.MILLISECOND,TZ}, getindex(x.data, i)) - -struct Interval{U,T} <: ArrowTimeType - x::T -end - -Base.zero(::Type{Interval{U,T}}) where {U,T} = Interval{U,T}(T(0)) - -bitwidth(x::Meta.IntervalUnit.T) = x == Meta.IntervalUnit.YEAR_MONTH ? Int32 : Int64 -Interval{Meta.IntervalUnit.YEAR_MONTH}(x) = - Interval{Meta.IntervalUnit.YEAR_MONTH,Int32}(Int32(x)) -Interval{Meta.IntervalUnit.DAY_TIME}(x) = - Interval{Meta.IntervalUnit.DAY_TIME,Int64}(Int64(x)) - -function juliaeltype(f::Meta.Field, x::Meta.Interval, convert) - return Interval{x.unit,bitwidth(x.unit)} -end - -function arrowtype(b, ::Type{Interval{U,T}}) where {U,T} - Meta.intervalStart(b) - Meta.intervalAddUnit(b, U) - return Meta.Interval, Meta.intervalEnd(b), nothing -end - -struct Duration{U} <: ArrowTimeType - x::Int64 -end - -Base.zero(::Type{Duration{U}}) where {U} = Duration{U}(Int64(0)) - -function juliaeltype(f::Meta.Field, x::Meta.Duration, convert) - return Duration{x.unit} -end - -finaljuliatype(::Type{Duration{U}}) where {U} = periodtype(U) -Base.convert(::Type{P}, x::Duration{U}) where {P<:Dates.Period,U} = P(periodtype(U)(x.x)) - -function arrowtype(b, ::Type{Duration{U}}) where {U} - Meta.durationStart(b) - Meta.durationAddUnit(b, U) - return Meta.Duration, Meta.durationEnd(b), nothing -end - -arrowtype(b, ::Type{P}) where {P<:Dates.Period} = arrowtype(b, Duration{arrowperiodtype(P)}) - -arrowperiodtype(P) = Meta.TimeUnit.SECOND -arrowperiodtype(::Type{Dates.Millisecond}) = Meta.TimeUnit.MILLISECOND -arrowperiodtype(::Type{Dates.Microsecond}) = Meta.TimeUnit.MICROSECOND -arrowperiodtype(::Type{Dates.Nanosecond}) = Meta.TimeUnit.NANOSECOND - -Base.convert(::Type{Duration{U}}, x::Dates.Period) where {U} = - Duration{U}(Dates.value(periodtype(U)(x))) - -ArrowTypes.ArrowType(::Type{P}) where {P<:Dates.Period} = Duration{arrowperiodtype(P)} -ArrowTypes.toarrow(x::P) where {P<:Dates.Period} = convert(Duration{arrowperiodtype(P)}, x) -const PERIOD_SYMBOL = Symbol("JuliaLang.Dates.Period") -ArrowTypes.arrowname(::Type{P}) where {P<:Dates.Period} = PERIOD_SYMBOL -ArrowTypes.JuliaType(::Val{PERIOD_SYMBOL}, ::Type{Duration{U}}) where {U} = periodtype(U) -ArrowTypes.fromarrow(::Type{P}, x::Duration{U}) where {P<:Dates.Period,U} = convert(P, x) - -# nested types; call juliaeltype recursively on nested children -function juliaeltype( - f::Meta.Field, - list::Union{Meta.List,Meta.LargeList,Meta.ListView,Meta.LargeListView}, - convert, -) - return Vector{juliaeltype(f.children[1], buildmetadata(f.children[1]), convert)} -end - -# arrowtype will call fieldoffset recursively for children -function arrowtype(b, x::List{T,O,A}) where {T,O,A} - if liststringtype(x) - if T <: AbstractString || T <: Union{AbstractString,Missing} - if O == Int32 - Meta.utf8Start(b) - return Meta.Utf8, Meta.utf8End(b), nothing - else # if O == Int64 - Meta.largUtf8Start(b) - return Meta.LargeUtf8, Meta.largUtf8End(b), nothing - end - else # if Base.CodeUnits - if O == Int32 - Meta.binaryStart(b) - return Meta.Binary, Meta.binaryEnd(b), nothing - else # if O == Int64 - Meta.largeBinaryStart(b) - return Meta.LargeBinary, Meta.largeBinaryEnd(b), nothing - end - end - else - children = [fieldoffset(b, "", x.data)] - if O == Int32 - Meta.listStart(b) - return Meta.List, Meta.listEnd(b), children - else - Meta.largeListStart(b) - return Meta.LargeList, Meta.largeListEnd(b), children - end - end -end - -function juliaeltype(f::Meta.Field, list::Meta.FixedSizeList, convert) - type = juliaeltype(f.children[1], buildmetadata(f.children[1]), convert) - return NTuple{Int(list.listSize),type} -end - -function arrowtype(b, x::FixedSizeList{T,A}) where {T,A} - N = ArrowTypes.getsize( - ArrowTypes.ArrowKind(ArrowTypes.ArrowType(Base.nonmissingtype(T))), - ) - if eltype(A) == UInt8 - Meta.fixedSizeBinaryStart(b) - Meta.fixedSizeBinaryAddByteWidth(b, Int32(N)) - return Meta.FixedSizeBinary, Meta.fixedSizeBinaryEnd(b), nothing - else - children = [fieldoffset(b, "", x.data)] - Meta.fixedSizeListStart(b) - Meta.fixedSizeListAddListSize(b, Int32(N)) - return Meta.FixedSizeList, Meta.fixedSizeListEnd(b), children - end -end - -function juliaeltype(f::Meta.Field, map::Meta.Map, convert) - K = juliaeltype( - f.children[1].children[1], - buildmetadata(f.children[1].children[1]), - convert, - ) - V = juliaeltype( - f.children[1].children[2], - buildmetadata(f.children[1].children[2]), - convert, - ) - return Dict{K,V} -end - -function arrowtype(b, x::Map) - children = [fieldoffset(b, "entries", x.data)] - Meta.mapStart(b) - return Meta.Map, Meta.mapEnd(b), children -end - -struct KeyValue{K,V} - key::K - value::V -end -keyvalueK(::Type{KeyValue{K,V}}) where {K,V} = K -keyvalueV(::Type{KeyValue{K,V}}) where {K,V} = V -Base.length(kv::KeyValue) = 1 -Base.iterate(kv::KeyValue, st=1) = st === nothing ? nothing : (kv, nothing) -ArrowTypes.default(::Type{KeyValue{K,V}}) where {K,V} = KeyValue(default(K), default(V)) - -function arrowtype(b, ::Type{KeyValue{K,V}}) where {K,V} - children = [fieldoffset(b, "key", K), fieldoffset(b, "value", V)] - Meta.structStart(b) - return Meta.Struct, Meta.structEnd(b), children -end - -function juliaeltype(f::Meta.Field, list::Meta.Struct, convert) - names = Tuple(Symbol(x.name) for x in f.children) - types = Tuple(juliaeltype(x, buildmetadata(x), convert) for x in f.children) - return NamedTuple{names,Tuple{types...}} -end - -function arrowtype(b, x::Struct{T,S}) where {T,S} - names = fieldnames(Base.nonmissingtype(T)) - children = [fieldoffset(b, names[i], x.data[i]) for i = 1:length(names)] - Meta.structStart(b) - return Meta.Struct, Meta.structEnd(b), children -end - -# Unions -function UnionT(f::Meta.Field, convert) - typeids = f.type.typeIds === nothing ? nothing : Tuple(Int(x) for x in f.type.typeIds) - UT = UnionT{ - f.type.mode, - typeids, - Tuple{(juliaeltype(x, buildmetadata(x), convert) for x in f.children)...}, - } - return UT -end - -juliaeltype(f::Meta.Field, u::Meta.Union, convert) = - Union{(juliaeltype(x, buildmetadata(x), convert) for x in f.children)...} - -function arrowtype( - b, - x::Union{DenseUnion{S,UnionT{T,typeIds,U}},SparseUnion{S,UnionT{T,typeIds,U}}}, -) where {S,T,typeIds,U} - if typeIds !== nothing - Meta.unionStartTypeIdsVector(b, length(typeIds)) - for id in Iterators.reverse(typeIds) - FlatBuffers.prepend!(b, id) - end - TI = FlatBuffers.endvector!(b, length(typeIds)) - end - children = [fieldoffset(b, "", x.data[i]) for i = 1:fieldcount(U)] - Meta.unionStart(b) - Meta.unionAddMode(b, T) - if typeIds !== nothing - Meta.unionAddTypeIds(b, TI) - end - return Meta.Union, Meta.unionEnd(b), children -end diff --git a/src/ipc_read.jl b/src/ipc_read.jl new file mode 100644 index 00000000..42e9a987 --- /dev/null +++ b/src/ipc_read.jl @@ -0,0 +1,1452 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 IPC STREAM reader as a thin peer over ArrowCore. The file-format +# reader (readfile/ArrowFile) shares these primitives from ipc_write.jl. +# +# * Framing: checked spans, the generated FlatBuffers verifier before any +# generated getter runs, and explicit resource limits (`Limits` + +# `framemessages`) enforced before metadata-directed allocation. The +# message body is the decoding AUTHORITY: every wire buffer is first a +# checked subslice of its message-body slice, so corrupt metadata cannot +# alias the schema message, another batch, or anything else in the file. +# Positively compressed buffers are decoded into separate exact-sized +# owned regions. +# +# * Decoding: ONE generic recursive decoder (`decodefield`) walks nodes and +# buffers in the order `layoutspec` declares; variadic layouts carry +# their own bounded count handling. +# +# * IPC bookkeeping stays in the adapter: dictionary ids live in an +# adapter-side table; Core Fields carry `DictionaryType` object +# references and never see an id. +# +# * Metadata bindings and the verifier are GENERATED from the vendored +# apache/arrow format/*.fbs (tools/fbsgen.jl -> src/metadata/) over the +# vendored FlatBuffers runtime. +# ============================================================================= + +# --------------------------------------------------------------------------- +# Framing stage: resource limits before metadata-directed decode allocation +# --------------------------------------------------------------------------- + +""" + Arrow.Limits(; keyword arguments...) + +Resource limits for an Arrow IPC reader. Each bound is checked before the +work it governs can exceed it. + +| keyword | default | bounds | +|:----------------------------|:--------------|:-------------------------------| +| `max_metadata_bytes` | 16 MiB | metadata per message | +| `max_body_bytes` | 2 GiB | one message body | +| `max_buffer_bytes` | 2 GiB | one buffer | +| `max_total_allocated_bytes` | 256 MiB | cumulative allocation per read | +| `max_messages` | 1,000,000 | messages per stream or file | +| `max_metadata_objects` | 1,000,000 | flatbuffer objects per message | +| `max_nesting_depth` | 64 | schema and metadata nesting | +| `max_array_length` | 1,000,000,000 | array length | +| `max_concurrent_reads` | 64 | concurrent range reads | + +Pass a `Limits` value as the `limits` keyword to [`Arrow.Table`](@ref) or +[`Arrow.Stream`](@ref). +""" +Base.@kwdef struct Limits + max_metadata_bytes::Int64 = 16 * 1024 * 1024 + max_body_bytes::Int64 = Int64(2) * 1024 * 1024 * 1024 + max_buffer_bytes::Int64 = Int64(2) * 1024 * 1024 * 1024 + max_total_allocated_bytes::Int64 = 256 * 1024 * 1024 + max_messages::Int = 1_000_000 + max_metadata_objects::Int = 1_000_000 + max_nesting_depth::Int = 64 + max_array_length::Int64 = 1_000_000_000 + max_concurrent_reads::Int = 64 +end + +mutable struct AllocationBudget + left::Int64 + limit::Int64 + lock::ReentrantLock +end + +function AllocationBudget(limit::Integer) + value = Int64(limit) + value >= 0 || throw(ArgumentError("allocation budget must be nonnegative")) + return AllocationBudget(value, value, ReentrantLock()) +end + +""" + Arrow.AllocationLimitError + +Thrown when a reader exhausts `Limits.max_total_allocated_bytes`. This budget +covers package-controlled fetch, decode, scan, ArrowTypes, and facade +materialization allocations for one read. +""" +struct AllocationLimitError <: Exception + msg::String +end +Base.showerror(io::IO, e::AllocationLimitError) = print(io, e.msg) + +function _charge_unlocked!(budget::AllocationBudget, amount::Int64, what::AbstractString) + amount >= 0 || throw(ArgumentError("negative allocation charge")) + amount <= budget.left || + throw(AllocationLimitError("$what exceeds the reader allocation budget")) + budget.left -= amount + return nothing +end + +function _charge!(budget::AllocationBudget, amount::Int64, what::AbstractString) + lock(budget.lock) + try + return _charge_unlocked!(budget, amount, what) + finally + unlock(budget.lock) + end +end + +function _refund!(budget::AllocationBudget, amount::Int64) + amount >= 0 || throw(ArgumentError("negative allocation refund")) + lock(budget.lock) + try + refunded = Base.Checked.checked_add(budget.left, amount) + refunded <= budget.limit || + throw(ArgumentError("allocation refund exceeds the charged total")) + budget.left = refunded + return nothing + finally + unlock(budget.lock) + end +end + +function _remaining(budget::AllocationBudget) + lock(budget.lock) + try + return budget.left + finally + unlock(budget.lock) + end +end + +AC._charge_materialization!(budget::AllocationBudget, amount::Int64, what::AbstractString) = + _charge!(budget, amount, what) +AC._materialization_remaining(budget::AllocationBudget) = _remaining(budget) +AC._materialization_limit_exceeded!(::AllocationBudget, what::AbstractString) = + throw(AllocationLimitError("$what exceeds the reader allocation budget")) + +function _verify_ipc_metadata_budgeted( + metabytes::Vector{UInt8}, + limits::Limits, + budget::AllocationBudget, +) + lock(budget.lock) + try + result = verify_ipc_metadata(metabytes, limits, budget.left) + _charge_unlocked!(budget, result[4], "verified metadata expansion") + return result + finally + unlock(budget.lock) + end +end + +function _verify_footer_budgeted( + footerbytes::Vector{UInt8}, + limits::Limits, + budget::AllocationBudget, +) + lock(budget.lock) + try + result = verify_footer(footerbytes, limits, budget.left) + _charge_unlocked!(budget, result[5], "verified footer expansion") + return result + finally + unlock(budget.lock) + end +end + +@inline function _chargevector!( + budget::Union{Nothing,AllocationBudget}, + ::Type{T}, + n::Integer, + what::AbstractString, +) where {T} + budget === nothing || _charge!(budget, AC._materializedvectorbytes(T, n), what) + return nothing +end + +@inline function _chargebitvector!( + budget::Union{Nothing,AllocationBudget}, + n::Integer, + what::AbstractString, +) + budget === nothing || _charge!(budget, AC._materializedbitvectorbytes(n), what) + return nothing +end + +@inline function _chargeobject!( + budget::Union{Nothing,AllocationBudget}, + payload::Integer, + what::AbstractString, +) + budget === nothing || _charge!(budget, AC._materializedobjectbytes(payload), what) + return nothing +end + +@inline function _dictslotbytes(::Type{K}, ::Type{V}) where {K,V} + return AC.checked_add( + Int64(Base.elsize(Vector{K})), + AC.checked_add(Int64(Base.elsize(Vector{V})), Int64(16)), + ) +end + +const _MIN_EMPTY_DICT_SLOTS = Int64(16) + +"Reserve the supported-version minimum backing store for one empty dictionary." +function _chargeemptydict!( + budget::Union{Nothing,AllocationBudget}, + ::Type{K}, + ::Type{V}, + what::AbstractString, +) where {K,V} + budget === nothing && return nothing + payload = AC.checked_mul(_MIN_EMPTY_DICT_SLOTS, _dictslotbytes(K, V)) + # Julia 1.10 eagerly allocates minimum Dict backing arrays. Later Julia + # versions defer more of this work. Reserve the larger supported layout so + # cache construction has one stable reader-budget contract on every host. + _charge!(budget, AC.checked_add(Int64(512), payload), what) + return nothing +end + +function _chargedict!( + budget::Union{Nothing,AllocationBudget}, + ::Type{K}, + ::Type{V}, + n::Integer, + what::AbstractString, +) where {K,V} + budget === nothing && return nothing + n >= 0 || throw(ArgumentError("negative dictionary capacity")) + entries = Int64(n) + capacity = entries == 0 ? Int64(0) : AC.checked_mul(entries, Int64(2)) + payload = AC.checked_mul(capacity, _dictslotbytes(K, V)) + # Dict's requested hash-table capacity and Julia's backing-store capacity + # can each round upward. Reserve both layers conservatively. + bytes = AC.checked_add(Int64(512), AC.checked_mul(Int64(2), payload)) + _charge!(budget, bytes, what) + return nothing +end + +"Charge amortized growth for one new entry after its container reserve." +function _chargedictentry!( + budget::Union{Nothing,AllocationBudget}, + ::Type{K}, + ::Type{V}, + what::AbstractString, +) where {K,V} + budget === nothing && return nothing + # `_chargedict!` reserves two slots per requested entry and two backing + # layers. Charge that linear term once per new memo key. ArrowTypes memo + # owners reserve the supported-version empty container separately. + bytes = AC.checked_mul(Int64(4), _dictslotbytes(K, V)) + _charge!(budget, bytes, what) + return nothing +end + +"Memoize one value after reserving the new dictionary entry's allocation." +function _memoized!( + f::F, + cache::AbstractDict{K,V}, + key, + budget::Union{Nothing,AllocationBudget}, + what::AbstractString, +) where {F,K,V} + haskey(cache, key) && return cache[key] + _chargedictentry!(budget, K, V, what) + return get!(f, cache, key) +end + +struct FramedMessage + msg::Meta.Message # parsed flatbuffer metadata + body::BufferSlice # THE authority: buffers must subslice this + version::Int16 + header_type::UInt8 + features::Vector{Int64} # populated on schema messages +end + +# The encapsulated-message marker that precedes every metadata length. This +# reader requires it even on V4 metadata, so streams written before the +# marker existed are refused. +const CONTINUATION = 0xFFFFFFFF +const EXPERIMENTAL_COMPRESSION_KEY = "ARROW:experimental_compression" + +# --------------------------------------------------------------------------- +# FlatBuffers verification (generated walkers over a schema-blind runtime) +# --------------------------------------------------------------------------- + +# The shape verifier is GENERATED from the vendored format/*.fbs by +# tools/fbsgen.jl (src/metadata/Verifier.jl). Table geometry, scalar widths, +# enum domains, string and vector bounds, union dispatch, and the +# nesting/object/reserve budgets all derive from the schema, so binding drift +# cannot reach them. The wrappers below own only what the schema cannot +# express: the accepted metadata versions and message kinds, and the +# features/version coupling. +# The aliases below exist for test fixtures, which reuse the runtime's +# traversal primitives to LOCATE the bytes they corrupt. No src file uses them. +const _VTable = Meta.VTable +const _vtable = Meta._vtable +const _vfield = Meta._vfield +const _vref = Meta._vref +const _vvector = Meta._vvector +const _vrange = Meta._vrange +const _vu8 = Meta._vu8 +const _vu32 = Meta._vu32 +const _vi32 = Meta._vi32 +const _vi64 = Meta._vi64 + +_vfail(msg) = throw(ValidationError("invalid IPC FlatBuffer: $msg")) + +_verifyctx(limits::Limits, reserve_limit::Int64) = Meta.VerifyContext( + Int64(limits.max_metadata_objects), + limits.max_nesting_depth, + reserve_limit, +) + +# Translate the metadata module's verifier exceptions into the adapter's +# error vocabulary at the wrapper boundary. +function _verified(f::F) where {F} + try + return f() + catch e + e isa Meta.VerifyError && _vfail(e.msg) + e isa Meta.VerifyBudgetError && throw(AllocationLimitError(e.msg)) + rethrow() + end +end + +function _schemafeatures(sch::Meta.Schema, version::Int16) + fv = sch.features + features = fv === nothing ? Int64[] : Int64[Int64(x) for x in fv] + version == Int16(3) && + !isempty(features) && + _vfail("schema features require metadata V5") + return features +end + +function verify_ipc_metadata( + bytes::Vector{UInt8}, + limits::Limits, + reserve_limit::Int64=limits.max_total_allocated_bytes, +) + ctx = _verifyctx(limits, reserve_limit) + # STAGED root verification. The inline stage proves the table shell and + # every non-reference field, including the version. The adapter then + # gates the version, so an unsupported one rejects in constant time. Only + # after that does the reference stage walk the attacker-directed header + # graph. + t = _verified(() -> Meta.verifyrootstart_Message(bytes, ctx)) + msg = FB.getrootas(Meta.Message, bytes, 0) + version = Int16(Int64(msg.version)) + version in (Int16(3), Int16(4)) || + _vfail("unsupported metadata version $version (only V4/V5 are accepted)") + _verified(() -> Meta.verifyrootrest_Message(t, ctx)) + # The verifier proved header presence and rejected union members outside + # the generated schemas (the Tensor family), so this dispatch is total. + header = msg.header + header_type = + header isa Meta.Schema ? UInt8(1) : + header isa Meta.DictionaryBatch ? UInt8(2) : + header isa Meta.RecordBatch ? UInt8(3) : _vfail("unsupported message header tag") + features = header isa Meta.Schema ? _schemafeatures(header, version) : Int64[] + return version, header_type, features, ctx.reserved +end + +""" + framemessages(region, limits) -> Vector{FramedMessage} + +Walk the IPC stream framing (continuation marker, metadata length, metadata +flatbuffer, body), checking every declared length against the limits and the +region's real extent before metadata-directed decode allocation. A truncated +prefix, metadata block, or body throws. EOF exactly after a complete message +is the intentional missing-EOS boundary case and is accepted. The framing's +metadata length covers the padded metadata only; a Footer Block's +metaDataLength instead covers prefix + metadata. +""" +framemessages(region::OwnerRegion, limits::Limits=Limits()) = _framemessages( + region, + limits, + Base.ENDIAN_BOM, + AllocationBudget(limits.max_total_allocated_bytes), +) + +function _validatelimits(limits::Limits) + limits.max_metadata_bytes >= 0 || throw(ArgumentError("negative metadata limit")) + limits.max_body_bytes >= 0 || throw(ArgumentError("negative body limit")) + limits.max_buffer_bytes >= 0 || throw(ArgumentError("negative buffer limit")) + limits.max_total_allocated_bytes >= 0 || + throw(ArgumentError("negative allocation limit")) + limits.max_messages >= 0 || throw(ArgumentError("negative message limit")) + limits.max_metadata_objects >= 0 || + throw(ArgumentError("negative metadata-object limit")) + limits.max_nesting_depth >= 0 || throw(ArgumentError("negative nesting limit")) + limits.max_array_length >= 0 || throw(ArgumentError("negative array-length limit")) + limits.max_concurrent_reads >= 1 || + throw(ArgumentError("concurrent-read limit must be positive")) + return nothing +end + +function _framemessages( + region::OwnerRegion, + limits::Limits, + host_endian_bom::UInt32, + budget::AllocationBudget=AllocationBudget(limits.max_total_allocated_bytes), +) + # The generated FlatBuffers bindings use native-endian scalar + # loads. Reject an unsupported host before any generated getter sees the + # little-endian wire bytes. The explicit argument keeps this ordering + # testable on the supported little-endian CI host. + host_endian_bom == UInt32(0x04030201) || + throw(ValidationError("the IPC reader requires a little-endian host")) + _validatelimits(limits) + blob = BufferSlice(region, 0, region.len) + msgs = FramedMessage[] + pos = Int64(0) # 0-based byte position within the blob + while pos < blob.len + blob.len - pos >= 8 || throw(ValidationError("truncated IPC prefix at byte $pos")) + pos % 8 == 0 || throw(ValidationError("IPC message is not 8-byte aligned")) + cont = AC.loadat(blob, UInt32, pos) + cont == CONTINUATION || + throw(ValidationError("missing continuation marker at byte $pos")) + # The stream framing's metadata length covers the padded metadata + # only. A Footer Block.metaDataLength instead covers prefix + + # metadata (see block!). + metalen = Int64(AC.loadat(blob, Int32, AC.checked_add(pos, Int64(4)))) + # Continuation marker + zero length is the end-of-stream marker. + # Nothing may follow it. + if metalen == 0 + AC.checked_add(pos, Int64(8)) == blob.len || + throw(ValidationError("trailing bytes after IPC end-of-stream")) + return msgs + end + length(msgs) < limits.max_messages || + throw(ValidationError("message count exceeds limit")) + 0 < metalen <= limits.max_metadata_bytes || throw( + ValidationError( + "metadata length $metalen outside (0, $(limits.max_metadata_bytes)]", + ), + ) + metalen % 8 == 0 || + throw(ValidationError("metadata length $metalen is not 8-byte aligned")) + metastart = AC.checked_add(pos, Int64(8)) + bodyguess = AC.checked_add(metastart, metalen) + bodyguess <= blob.len || + throw(ValidationError("truncated metadata: need $metalen bytes at $pos")) + _chargevector!(budget, UInt8, metalen, "metadata allocation") + metabytes = AC.slicebytes(AC.subslice(blob, metastart, metalen)) + version, header_type, features, _ = + _verify_ipc_metadata_budgeted(metabytes, limits, budget) + # No generated getter runs before the verifier has bounded the full + # table/vector/string graph it may visit. + msg = FB.getrootas(Meta.Message, metabytes, 0) + bodylen = Int64(msg.bodyLength) + 0 <= bodylen <= limits.max_body_bytes || throw( + ValidationError("body length $bodylen outside [0, $(limits.max_body_bytes)]"), + ) + bodylen % 8 == 0 || + throw(ValidationError("body length $bodylen is not 8-byte aligned")) + bodystart = bodyguess + bodyend = AC.checked_add(bodystart, bodylen) + bodyend <= blob.len || + throw(ValidationError("truncated body: need $bodylen bytes at $bodystart")) + push!( + msgs, + FramedMessage( + msg, + AC.subslice(blob, bodystart, bodylen), + version, + header_type, + features, + ), + ) + pos = bodyend + end + return msgs +end + +# --------------------------------------------------------------------------- +# Metadata mapping: Meta.* type structs -> Core runtime descriptors +# --------------------------------------------------------------------------- + +# One value-level mapping table: one function per direction over runtime +# descriptors. Julia-type conversion is the facade's concern, not this +# adapter's. + +function coretype(t)::ArrowType + if t isa Meta.Int + IntType(Int(t.bitWidth), t.is_signed) + elseif t isa Meta.FloatingPoint + FloatType( + t.precision == Meta.Precision.HALF ? 16 : + t.precision == Meta.Precision.SINGLE ? 32 : 64, + ) + elseif t isa Meta.Bool + BoolType() + elseif t isa Meta.Utf8 + Utf8Type(false) + elseif t isa Meta.LargeUtf8 + Utf8Type(true) + elseif t isa Meta.Binary + BinaryType(false) + elseif t isa Meta.LargeBinary + BinaryType(true) + elseif t isa Meta.FixedSizeBinary + FixedSizeBinaryType(Int(something(t.byteWidth, Int32(0)))) + elseif t isa Meta.List + ListType(false) + elseif t isa Meta.LargeList + ListType(true) + elseif t isa Meta.FixedSizeList + FixedSizeListType(Int(t.listSize)) + elseif t isa Meta.Struct + StructType() + elseif t isa Meta.Map + MapType(something(t.keysSorted, false)) + elseif t isa Meta.Timestamp + timezone = t.timezone + TimestampType(timeunit(t.unit), timezone === nothing ? nothing : String(timezone)) + elseif t isa Meta.Date + DateType(t.unit == Meta.DateUnit.DAY ? AC.DAY : AC.MILLISECOND_DATE) + elseif t isa Meta.Time + TimeType(timeunit(t.unit), Int(t.bitWidth)) + elseif t isa Meta.Duration + DurationType(timeunit(t.unit)) + elseif t isa Meta.Decimal + DecimalType(Int(t.precision), Int(t.scale), Int(t.bitWidth)) + elseif t isa Meta.Interval + IntervalType( + t.unit == Meta.IntervalUnit.YEAR_MONTH ? AC.YEAR_MONTH : + t.unit == Meta.IntervalUnit.DAY_TIME ? AC.DAY_TIME : AC.MONTH_DAY_NANO, + ) + elseif t isa Meta.Utf8View + ViewType(true) + elseif t isa Meta.BinaryView + ViewType(false) + elseif t isa Meta.ListView + ListViewType(false) + elseif t isa Meta.LargeListView + ListViewType(true) + elseif t isa Meta.RunEndEncoded + RunEndEncodedType() + elseif t isa Meta.Null + NullType() + else + throw(ValidationError("IPC adapter does not map metadata type $(typeof(t))")) + end +end + +""" +Map one metadata type to a Core descriptor, with the built child Fields in +hand — Union is the one type whose descriptor (mode + type ids) spans the +type table AND the children vector, so it cannot go through `coretype`. +""" +function _coremetatype(mt, children::Vector{Field})::ArrowType + mt isa Meta.Union || return coretype(mt) + mode = mt.mode == Meta.UnionMode.Dense ? AC.DenseMode : AC.SparseMode + ids = mt.typeIds + nchildren = length(children) + # Type ids are Int8 in [0, 127] (checked below), so 128 distinct children + # is the ceiling. + nchildren <= 128 || throw(ValidationError("a union cannot have more than 128 children")) + if ids === nothing + return UnionType(mode, Int8[Int8(i) for i = 0:(nchildren - 1)]) + end + length(ids) == nchildren || + throw(ValidationError("union type-id count must equal child count")) + all(x -> 0 <= x <= 127, ids) || + throw(ValidationError("union type ids must be in [0, 127]")) + coreids = Int8[Int8(x) for x in ids] + length(unique(coreids)) == length(coreids) || + throw(ValidationError("union type ids must be unique")) + return UnionType(mode, coreids) +end + +timeunit(u) = + u == Meta.TimeUnit.SECOND ? AC.SECOND : + u == Meta.TimeUnit.MILLISECOND ? AC.MILLISECOND : + u == Meta.TimeUnit.MICROSECOND ? AC.MICROSECOND : AC.NANOSECOND + +function coremetadata(kvs) + kvs === nothing && return nothing + return Pair{String,String}[ + String(kv.key) => String(something(kv.value, "")) for kv in kvs + ] +end + +""" +Convert a metadata Field to a Core Field. Dictionary-encoded fields become +`DictionaryType` here; the IPC dictionary id is recorded in the adapter's +side table (`dictids`), NOT on the Core field — Core never learns about ids. +""" +function corefield( + f::Meta.Field, + dictids::Dict{Int64,Meta.Field}, + fielddictids::IdDict{Field,Int64}, +) + children = Field[ + corefield(c, dictids, fielddictids) for c in something(f.children, Meta.Field[]) + ] + t = _coremetatype(f.type, children) + if f.dictionary === nothing + return Field( + String(something(f.name, "")), + t, + f.nullable, + coremetadata(f.custom_metadata), + children, + ) + end + # Nested dictionary encoding (a dictionary field whose VALUE type has + # dictionary-encoded children) is spec-legal and present in the + # arrow-testing gold corpus (nested_dictionary: dict(list(dict(utf8)))). + # A dictionary batch's values decode through the same `decodefield` + # with the live pool table, so inner pools resolve as long as batches + # arrive in dependency order — which the IPC spec requires. + dictids[f.dictionary.id] = f + idxt = + f.dictionary.indexType === nothing ? IntType(32, true) : + coretype(f.dictionary.indexType)::IntType + cf = Field( + String(something(f.name, "")), + DictionaryType(idxt, t, f.dictionary.isOrdered), + f.nullable, + coremetadata(f.custom_metadata), + children, + ) + # Identity-keyed: safe for duplicate column names and nested dict fields + # (name matching would be neither). + fielddictids[cf] = f.dictionary.id + return cf +end + +function validatedictionaryids(fields, fielddictids::IdDict{Field,Int64}) + seen = Dict{Int64,Field}() + compatible(a::Field, b::Field; compare_name::Bool=false) = + (!compare_name || a.name == b.name) && + AC.typeequal(a.type, b.type) && + a.nullable == b.nullable && + # One id resolves ONE pool, so repeated ids must agree on the whole + # nested dictionary-id topology: compatible value schemas whose + # nested fields carry DIFFERENT wire ids would decode the second + # field through pools its schema never declared. + (!(a.type isa DictionaryType) || fielddictids[a] == fielddictids[b]) && + length(a.children) == length(b.children) && + all(compatible(x, y; compare_name=true) for (x, y) in zip(a.children, b.children)) + function walk(f::Field) + if f.type isa DictionaryType + id = fielddictids[f] + vf = AC.dictvaluefield(f, f.type) + if haskey(seen, id) + old = seen[id] + compatible(old, vf) || throw( + ValidationError( + "dictionary id $id is shared by incompatible value schemas", + ), + ) + else + seen[id] = vf + # A pool's value schema may itself hold dictionary-encoded + # fields (nested dictionary encoding); their ids resolve + # through this same table, so register them too. + walk(vf) + end + return + end + foreach(walk, f.children) + end + foreach(walk, fields) + return seen +end + +function validateschemafield(f::Field) + AC._validate_descriptor(f.type) + if f.type isa DictionaryType + validateschemafield(AC.dictvaluefield(f, f.type)) + return f + end + spec = layoutspec(f.type) + expected = spec.childcount == -1 ? length(f.children) : spec.childcount + length(f.children) == expected || throw( + ValidationError( + "$(typeof(f.type)) schema expects $expected children, got $(length(f.children))", + ), + ) + if f.type isa UnionType + length(f.type.typeids) == length(f.children) || + throw(ValidationError("union type-id count must equal child count")) + length(unique(f.type.typeids)) == length(f.type.typeids) || + throw(ValidationError("union type ids must be unique")) + all(>=(0), f.type.typeids) || + throw(ValidationError("union type ids must be in [0, 127]")) + elseif f.type isa MapType + entries = f.children[1] + entries.type isa StructType && + !entries.nullable && + length(entries.children) == 2 && + !entries.children[1].nullable || + throw(ValidationError("invalid map entries/key schema")) + elseif f.type isa RunEndEncodedType + length(f.children) == 2 || throw(ValidationError("REE requires two children")) + run, values = f.children + run.name == "run_ends" && + values.name == "values" && + run.type isa IntType && + run.type.signed && + run.type.bits in (16, 32, 64) && + !run.nullable || throw(ValidationError("invalid run-end encoded schema")) + end + foreach(validateschemafield, f.children) + return f +end + +# --------------------------------------------------------------------------- +# THE generic decoder: registry-driven node/buffer consumption +# --------------------------------------------------------------------------- + +# Node/buffer consumption order falls out of `layoutspec`: one field = one +# node (unless the layout says otherwise) + the registry's buffers in +# registry order + children in declared order. Nothing threads +# (nodeidx, bufferidx, varbufferidx) by hand per layout, so an off-by-one +# cannot silently shift every subsequent buffer; a mismatch is a thrown +# error at the *end* of the batch (leftover nodes/buffers), not corruption. + +# Buffer compression: one codec context per reader, reused across buffers +# and explicitly finalized when the reader is done — no global pools. +const CODEC_NONE = Int8(-1) +const CODEC_LZ4_FRAME = Int8(0) # Meta.CompressionType.LZ4_FRAME +const CODEC_ZSTD = Int8(1) # Meta.CompressionType.ZSTD + +mutable struct DecodeState + lz4::Ptr{CLZ4.LZ4F_dctx} + zstd::Ptr{ZSTD.ZSTD_DCtx} + budget::AllocationBudget +end + +DecodeState(budget::AllocationBudget) = + DecodeState(Ptr{CLZ4.LZ4F_dctx}(C_NULL), Ptr{ZSTD.ZSTD_DCtx}(C_NULL), budget) + +function _lz4ctx!(state::DecodeState) + state.lz4 != C_NULL && return state.lz4 + slot = Ref{Ptr{CLZ4.LZ4F_dctx}}(C_NULL) + CLZ4.LZ4F_createDecompressionContext(slot, CLZ4.LZ4F_getVersion()) + state.lz4 = slot[] + return state.lz4 +end + +function _zstdctx!(state::DecodeState) + state.zstd != C_NULL && return state.zstd + p = ZSTD.ZSTD_createDCtx() + p == C_NULL && throw(OutOfMemoryError()) + state.zstd = p + return p +end + +function Base.close(state::DecodeState) + lz4 = state.lz4 + state.lz4 = Ptr{CLZ4.LZ4F_dctx}(C_NULL) + try + lz4 == C_NULL || CLZ4.LZ4F_freeDecompressionContext(lz4) + finally + zstd = state.zstd + state.zstd = Ptr{ZSTD.ZSTD_DCtx}(C_NULL) + zstd == C_NULL || ZSTD.ZSTD_freeDCtx(zstd) + end + return nothing +end + +function _decode_lz4!( + state::DecodeState, + src::Ptr{UInt8}, + srclen::Int64, + out::Vector{UInt8}, + declared::Int64, +) + ctx = _lz4ctx!(state) + CLZ4.LZ4F_resetDecompressionContext(ctx) + inpos = Int64(0) + outpos = Int64(0) + while true + insize = Ref{Csize_t}(Csize_t(srclen - inpos)) + outsize = Ref{Csize_t}(Csize_t(declared - outpos)) + # A zero-capacity destination is valid. It lets the decoder consume + # an empty frame or the footer after the last output byte without a + # second allocation. + dst = outpos == declared ? Ptr{UInt8}(C_NULL) : pointer(out) + outpos + hint = CLZ4.LZ4F_decompress(ctx, dst, outsize, src + inpos, insize, C_NULL) + inpos += Int64(insize[]) + outpos += Int64(outsize[]) + if hint == 0 + inpos == srclen || throw( + ValidationError("LZ4 buffer contains trailing bytes or multiple frames"), + ) + outpos == declared || throw( + ValidationError( + "LZ4 output length $outpos does not match declared $declared", + ), + ) + return nothing + end + inpos < srclen || throw(ValidationError("truncated LZ4 frame")) + (insize[] != 0 || outsize[] != 0) || + throw(ValidationError("LZ4 output exceeds declared length $declared")) + end +end + +function _decode_zstd!( + state::DecodeState, + src::Ptr{UInt8}, + srclen::Int64, + out::Vector{UInt8}, + declared::Int64, +) + dst = declared == 0 ? Ptr{UInt8}(C_NULL) : pointer(out) + got = ZSTD.ZSTD_decompressDCtx( + _zstdctx!(state), + dst, + Csize_t(declared), + src, + Csize_t(srclen), + ) + if ZSTD.ZSTD_isError(got) != 0 + msg = unsafe_string(ZSTD.ZSTD_getErrorName(got)) + throw(ValidationError("ZSTD decompression failed: $msg")) + end + Int64(got) == declared || throw( + ValidationError( + "ZSTD output length $(Int64(got)) does not match declared $declared", + ), + ) + return nothing +end + +# `B` is the body representation: a contiguous `BufferSlice` for in-memory +# and mmapped messages, or a sparse body (scan.jl) whose fetched +# spans stand in for the contiguous message body. `_bodyslice` is the one +# seam between them; the parameter keeps the cursor concrete per use. +mutable struct DecodeCursor{B} + nodes::AbstractVector{Meta.FieldNode} + buffers::AbstractVector{Meta.Buffer} + body::B + max_buffer_bytes::Int64 + max_array_length::Int64 + nodeidx::Int + bufidx::Int + last_nonempty_end::Int64 + codec::Int8 # CODEC_NONE, or the batch's declared codec + state::Union{Nothing,DecodeState} + # One entry per view-typed field in depth-first schema order: how many + # variadic data buffers that field consumes (format 1.4). Non-view + # batches carry an empty vector; a leftover entry is a skew error. + variadics::AbstractVector{<:Integer} + varidx::Int +end + +"Resolve one declared buffer window against the message body." +_bodyslice(body::BufferSlice, offset::Int64, len::Int64) = AC.subslice(body, offset, len) + +DecodeCursor( + nodes, + buffers, + body, + limits::Limits; + codec::Int8=CODEC_NONE, + state::Union{Nothing,DecodeState}=nothing, + variadics=nothing, +) = DecodeCursor( + something(nodes, Meta.FieldNode[]), + something(buffers, Meta.Buffer[]), + body, + limits.max_buffer_bytes, + limits.max_array_length, + 1, + 1, + 0, + codec, + state, + something(variadics, Int64[]), + 1, +) + +""" + variadiccounts(rb::Meta.RecordBatch) -> Vector{Int64} + +The batch's `variadicBufferCounts` as a concrete `Vector{Int64}` (empty when +the slot is absent). The generated binding reads the spec's `[long]` at +8-byte width; this accessor exists so every site shares one normalized shape. +""" +variadiccounts(rb::Meta.RecordBatch) = + collect(Int64, something(rb.variadicBufferCounts, Int64[])) + +"One variadic-buffer count, in depth-first view-field order (format 1.4)." +function takevariadic!(c::DecodeCursor) + c.varidx <= length(c.variadics) || throw( + ValidationError( + "metadata declares fewer variadic buffer counts than the schema requires", + ), + ) + n = c.variadics[c.varidx] + c.varidx += 1 + 0 <= n <= length(c.buffers) || + throw(ValidationError("variadic buffer count $n outside [0, $(length(c.buffers))]")) + return Int(n) +end + +function takenode!(c::DecodeCursor) + c.nodeidx <= length(c.nodes) || throw( + ValidationError("metadata declares fewer field nodes than the schema requires"), + ) + n = c.nodes[c.nodeidx] + c.nodeidx += 1 + 0 <= n.length <= c.max_array_length || + throw(ValidationError("field-node length $(n.length) exceeds limit")) + 0 <= n.null_count <= n.length || + throw(ValidationError("invalid field-node null count $(n.null_count)")) + return n +end + +""" +Consume one buffer-table entry's METADATA: bounds, alignment, limits, and the +non-overlap/monotone invariants — every check that needs no body byte. +`takebuffer!` continues into the body subslice and decompression; +`skipbuffer!` stops here, so scan pushdown can skip a column's decode and its +body range. Ranged tail reads and coalescing may still over-read those bytes. +""" +function _buffermeta!(c::DecodeCursor) + c.bufidx <= length(c.buffers) || + throw(ValidationError("metadata declares fewer buffers than the schema requires")) + b = c.buffers[c.bufidx] + c.bufidx += 1 + offset = Int64(b.offset) + len = Int64(b.length) + offset >= 0 || throw(ValidationError("negative batch buffer offset $offset")) + offset % 8 == 0 || + throw(ValidationError("batch buffer offset $offset is not 8-byte aligned")) + 0 <= len <= c.max_buffer_bytes || + throw(ValidationError("batch buffer length $len exceeds limit")) + if len > 0 + offset >= c.last_nonempty_end || + throw(ValidationError("batch buffers overlap or move backwards")) + c.last_nonempty_end = try + AC.checked_add(offset, len) + catch e + e isa OverflowError || rethrow() + throw(ValidationError("batch buffer end overflows")) + end + end + return offset, len +end + +skipbuffer!(c::DecodeCursor) = (_buffermeta!(c); nothing) + +function takebuffer!(c::DecodeCursor) + offset, len = _buffermeta!(c) + # THE checked-subslice step: a buffer is only ever a window into this + # message's body span (or, for a sparse body, into a fetched span that + # was itself derived from this buffer table). Checked arithmetic turns a + # corrupt offset/length into a clean ValidationError. + wire = try + _bodyslice(c.body, offset, len) + catch e + e isa ArgumentError || e isa OverflowError || rethrow() + throw(ValidationError("batch buffer [$offset, $len] escapes its message body")) + end + (c.codec == CODEC_NONE || len == 0) && return wire + return _decompressbuffer!(c, wire) +end + +""" +Decode one compressed buffer: an Int64 uncompressed-length prefix, then the +payload. A prefix of -1 means the payload is stored raw and stays a view of +the wire mapping. Every declared size is bounded before allocation (the +prefix is attacker-controlled) and the decompressed size must match the +declaration exactly. A positively compressed buffer becomes its own +exact-sized owned region. +""" +function _decompressbuffer!(c::DecodeCursor, wire::BufferSlice) + wire.len >= 8 || throw( + ValidationError("compressed buffer of $(wire.len) bytes lacks its length prefix"), + ) + declared = AC.loadat(wire, Int64, Int64(0)) + declared == -1 && return AC.subslice(wire, 8, wire.len - 8) # stored raw + 0 <= declared <= c.max_buffer_bytes || throw( + ValidationError("declared decompressed length $declared exceeds the buffer limit"), + ) + declared <= typemax(Int) || + throw(ValidationError("declared decompressed buffer is not addressable")) + payloadlen = wire.len - 8 + payloadlen > 0 || throw(ValidationError("compressed buffer has an empty payload")) + state = c.state::DecodeState + _chargevector!(state.budget, UInt8, declared, "decompressed bytes") + committed = false + try + # This is the only output allocation. Its size was checked and + # charged before either native decoder sees the input frame. + out = Vector{UInt8}(undef, Int(declared)) + # The native decoders read through a raw pointer, so the wire + # region's root must stay reachable for the whole call (Core rule 2). + wireregion = wire.region::OwnerRegion + GC.@preserve out wireregion begin + src = AC.sliceptr(wire) + 8 + if c.codec == CODEC_LZ4_FRAME + _decode_lz4!(state, src, payloadlen, out, declared) + else + _decode_zstd!(state, src, payloadlen, out, declared) + end + end + result = BufferSlice(heapregion(out), 0, declared) + committed = true + return result + catch e + e isa Union{ValidationError,AllocationLimitError} && rethrow() + e isa OutOfMemoryError && rethrow() + e isa InterruptException && rethrow() + throw(ValidationError("buffer decompression failed: $(sprint(showerror, e))")) + finally + committed || _refund!(state.budget, declared) + end +end + +function finishcursor!(c::DecodeCursor) + c.nodeidx == length(c.nodes) + 1 || + throw(ValidationError("unconsumed field nodes: schema/batch mismatch")) + c.bufidx == length(c.buffers) + 1 || + throw(ValidationError("unconsumed buffers: schema/batch mismatch")) + c.varidx == length(c.variadics) + 1 || + throw(ValidationError("unconsumed variadic buffer counts: schema/batch mismatch")) + return nothing +end + +function missingdicts( + fields, + nodes, + dicts::Dict{Int64,ArrayData}, + fielddictids::IdDict{Field,Int64}, +) + ns = something(nodes, Meta.FieldNode[]) + idx = Ref(1) + missing = Set{Int64}() + function walk(f::Field) + idx[] <= length(ns) || throw( + ValidationError("metadata declares fewer field nodes than the schema requires"), + ) + node = ns[idx[]] + idx[] += 1 + if f.type isa DictionaryType + id = fielddictids[f] + if !haskey(dicts, id) + node.length >= 0 && node.null_count == node.length || throw( + ValidationError( + "record batch uses undefined dictionary id $id for a non-null slot", + ), + ) + push!(missing, id) + end + return + end + spec = layoutspec(f.type) + nchildren = spec.childcount == -1 ? length(f.children) : spec.childcount + for i = 1:nchildren + walk(f.children[i]) + end + end + foreach(walk, fields) + return missing +end + +""" + decodefield(field, cursor, dictionaries, fielddictids) -> ArrayData + +Generic over every mapped layout in this adapter, variadic (view) layouts +included. +Dictionary-encoded columns consume the INDEX layout's buffers (validity + +indices) and resolve their values through the adapter's dictionary table. +""" +function decodefield( + f::Field, + c::DecodeCursor, + dicts::Dict{Int64,ArrayData}, + fielddictids::IdDict{Field,Int64}, +) + t = f.type + node = takenode!(c) + spec = layoutspec(t) + buffers = BufferSlice[takebuffer!(c) for _ in spec.buffers] + if spec.variadic + # View layouts append their declared count of variadic data buffers + # after the fixed validity/views pair (format 1.4). + for _ = 1:takevariadic!(c) + push!(buffers, takebuffer!(c)) + end + end + for (role, buffer) in zip(spec.buffers, buffers) + # A zero-length array may omit its offsets buffer entirely; Core + # accepts that canonical empty form, and nanoarrow and C++ write it. + # A PARTIAL offsets buffer — nonempty but short of one slot — is + # still malformed framing. + if role == AC.OFFSETS && node.length == 0 && 0 < buffer.len < spec.offsetwidth + throw(ValidationError("IPC offsets buffer is shorter than one offset slot")) + end + end + children = ArrayData[] + if t isa DictionaryType + # Index buffers were just consumed; values come from the side table. + id = fielddictids[f] + haskey(dicts, id) || throw( + ValidationError( + "record batch references dictionary id $id before its dictionary batch", + ), + ) + return ArrayData( + t, + node.length, + buffers; + dictionary=dicts[id], + nullcount=node.null_count, + ) + end + nchildren = spec.childcount == -1 ? length(f.children) : spec.childcount + for i = 1:nchildren + push!(children, decodefield(f.children[i], c, dicts, fielddictids)) + end + if t isa UnionType && t.mode == AC.SparseMode + all(child -> child.len == node.length, children) || + throw(ValidationError("IPC sparse-union children must equal the union length")) + end + return ArrayData(t, node.length, buffers; children=children, nullcount=node.null_count) +end + +function validaterecordcolumns( + fields, + cols, + validated_dictionaries::AC._ValidatedDictionaries, +) + for (f, col) in zip(fields, cols) + # One IPC dictionary id can back many fields. Compatible value + # schemas were proved when the stream schema was built. Each immutable + # pool snapshot is fully certified at its DictionaryBatch, so record + # validation skips that pool tree. Index contracts still run + # independently for every field. + AC._validate_semantic(f, col, validated_dictionaries) + end + return validated_dictionaries +end + +function decoderecord( + fm::FramedMessage, + fields, + sch::Schema, + dicts::Dict{Int64,ArrayData}, + fielddictids::IdDict{Field,Int64}, + limits::Limits, + validated_dictionaries, + state::DecodeState, +) + header = fm.msg.header::Meta.RecordBatch + codec = _batchcodec(header.compression, fm.version) + rblen = something(header.length, Int64(0)) + 0 <= rblen <= limits.max_array_length || + throw(ValidationError("record batch length $rblen exceeds limit")) + cursor = DecodeCursor( + header.nodes, + header.buffers, + fm.body, + limits; + codec=codec, + state=state, + variadics=variadiccounts(header), + ) + cols = ArrayData[decodefield(f, cursor, dicts, fielddictids) for f in fields] + finishcursor!(cursor) + validaterecordcolumns(fields, cols, validated_dictionaries) + all(col -> col.len == rblen, cols) || + throw(ValidationError("RecordBatch length does not match top-level field nodes")) + return AC.RecordBatch(sch, cols, rblen, validated_dictionaries) +end + +""" +Arrow 0.17 V4 streams signaled buffer compression on the Message, before +RecordBatch.compression existed. Reject that legacy marker before treating +its length-prefixed compressed buffers as raw data. +""" +function rejectexperimentalcompression( + msg::Meta.Message, + version::Int16, + header_type::UInt8, +) + version == Int16(3) || return nothing # V4 + header_type in (UInt8(2), UInt8(3)) || return nothing + metadata = msg.custom_metadata + metadata === nothing && return nothing + any(kv -> kv.key == EXPERIMENTAL_COMPRESSION_KEY, metadata) && throw( + ValidationError( + "pre-1.0 experimental V4 IPC compression (the " * + "ARROW:experimental_compression metadata convention, superseded " * + "by V5 BodyCompression in 2020) is not supported", + ), + ) + return nothing +end +rejectexperimentalcompression(fm::FramedMessage) = + rejectexperimentalcompression(fm.msg, fm.version, fm.header_type) + +# --------------------------------------------------------------------------- +# Stream reader: RecordBatchSource over framed messages +# --------------------------------------------------------------------------- + +mutable struct IPCStream <: AC.RecordBatchSource + schema::Schema + corefields::AC.FrozenVector{Field} + batches::Vector{AC.RecordBatch} + region::OwnerRegion # one lifetime root for the wire bytes + nextindex::Int + @atomic pulling::Bool + fielddictids::IdDict{Field,Int64} # adapter-side id table (shared ids preserved) + budget::AllocationBudget # framing, decode, and facade materialization + limits::Limits # complete policy used when this handle opened +end + +mutable struct PendingRecord + fm::FramedMessage + dictionaries::Dict{Int64,ArrayData} + missing::Set{Int64} + slot::Int +end +AC.schema(s::IPCStream) = s.schema + +function AC.nextbatch!(s::IPCStream) + # One active pull at a time: the claim CAS rejects concurrent callers + # (fail closed, no duplicated or skipped batches) and is released on + # every exit path. + _, ok = @atomicreplace s.pulling false => true + ok || throw( + Base.ConcurrencyViolationError( + "IPCStream supports only one active nextbatch! call", + ), + ) + try + i = s.nextindex + i > length(s.batches) && return nothing + b = s.batches[i] + s.nextindex = i + 1 + return b + finally + @atomic :release s.pulling = false + end +end + +""" +Map a batch's declared BodyCompression to a codec id, enforcing the spec +subset this adapter supports: BUFFER-method LZ4_FRAME or ZSTD. +BodyCompression is V5-only; a V4 batch that declares it is refused. +""" +function _batchcodec(compression, version::Int16)::Int8 + compression === nothing && return CODEC_NONE + version == Int16(4) || throw(ValidationError("BodyCompression requires metadata V5")) + method = something(compression.method, Meta.BodyCompressionMethod.BUFFER) + method == Meta.BodyCompressionMethod.BUFFER || + throw(ValidationError("unsupported body-compression method $method")) + codec = something(compression.codec, Meta.CompressionType.LZ4_FRAME) + codec == Meta.CompressionType.LZ4_FRAME && return CODEC_LZ4_FRAME + codec == Meta.CompressionType.ZSTD && return CODEC_ZSTD + throw(ValidationError("unsupported compression codec $codec")) +end + +""" + readstream(bytes; limits=Limits()) -> IPCStream + +Decode a stream from a borrowed byte vector. Raw batch buffers remain +zero-copy views of `bytes`; positively compressed buffers become exact-sized +owned copies. The caller must not mutate or resize `bytes` until the returned +stream and all batches from it are unreachable (the facade's `Arrow.Table` +and `Arrow.Stream` own their backing storage and do not expose this borrow). +`IPCStream` is a single-owner cursor; overlapping `nextbatch!` calls throw +`ConcurrencyViolationError`. +""" +readstream(bytes::Vector{UInt8}; limits::Limits=Limits()) = + _readstream(bytes, limits, AllocationBudget(limits.max_total_allocated_bytes)) + +function _readstream(bytes::Vector{UInt8}, limits::Limits, budget::AllocationBudget) + region = heapregion(bytes) + msgs = _framemessages(region, limits, Base.ENDIAN_BOM, budget) + isempty(msgs) && throw(ValidationError("empty IPC stream")) + first(msgs).header_type == 1 || # Schema + throw(ValidationError("first IPC message must be a schema")) + msgs[1].msg.header isa Meta.Schema || + throw(ValidationError("first IPC message must be a schema")) + metaschema = msgs[1].msg.header + msgs[1].body.len == 0 || + throw(ValidationError("schema message must have an empty body")) + endian = something(metaschema.endianness, Meta.Endianness.Little) + endian == Meta.Endianness.Little || throw( + ValidationError("big-endian IPC is not supported (no endianness normalization)"), + ) + dictids = Dict{Int64,Meta.Field}() + fielddictids = IdDict{Field,Int64}() # adapter-side id table + fields = Field[ + corefield(f, dictids, fielddictids) for + f in something(metaschema.fields, Meta.Field[]) + ] + foreach(validateschemafield, fields) + dictvaluefields = validatedictionaryids(fields, fielddictids) + sch = Schema( + fields; + metadata=coremetadata(metaschema.custom_metadata), + endianness=AC.LittleEndian, + ) + dicts = Dict{Int64,ArrayData}() + # One codec context per reader, shared by every compressed batch in the + # stream and explicitly finalized on every exit path. + state = DecodeState(budget) + validated_dictionaries = AC._ValidatedDictionaries() + batchslots = Union{Nothing,AC.RecordBatch}[] + pending = PendingRecord[] + features = Set(msgs[1].features) + schemaversion = msgs[1].version + try + for fm in msgs[2:end] + fm.version == schemaversion || + throw(ValidationError("IPC metadata version changes within the stream")) + # Legacy V4 compression marker: see rejectexperimentalcompression. + rejectexperimentalcompression(fm) + header = fm.msg.header + if header isa Meta.DictionaryBatch + header.isDelta && + throw(ValidationError("delta dictionaries are not supported")) + rb = header.data + codec = _batchcodec(rb.compression, fm.version) + haskey(dictids, header.id) || + throw(ValidationError("dictionary batch has unknown id $(header.id)")) + replacement = haskey(dicts, header.id) + if replacement && !(1 in features) + throw( + ValidationError( + "dictionary replacement used without required schema feature", + ), + ) + end + # A dictionary batch's payload is a one-column record batch of + # the VALUE type, so the generic decoder handles it. Pool + # nullability is independent of the encoded index field. + haskey(dictvaluefields, header.id) || + throw(ValidationError("dictionary batch has unknown id $(header.id)")) + # Value fields were built once from the Core schema; reusing + # them avoids repeated metadata-string/container allocation on + # dictionary replacement messages. + vf = dictvaluefields[header.id] + rblen = something(rb.length, Int64(0)) + 0 <= rblen <= limits.max_array_length || + throw(ValidationError("dictionary batch length $rblen exceeds limit")) + cursor = DecodeCursor( + rb.nodes, + rb.buffers, + fm.body, + limits; + codec=codec, + state=state, + variadics=variadiccounts(rb), + ) + decoded = decodefield(vf, cursor, dicts, fielddictids) + finishcursor!(cursor) + decoded.len == rblen || throw( + ValidationError( + "dictionary RecordBatch length does not match its field node", + ), + ) + # Certify the entire immutable pool snapshot before publication. + # Later record validation may then skip every recursive stage for + # this exact identity. Replacements decode to a new identity and + # must earn their own certificate here. + validate_semantic(vf, decoded) + validated_dictionaries[decoded] = nothing + dicts[header.id] = decoded + + # The IPC spec permits an all-null dictionary column before its + # first DictionaryBatch. Resolve only the missing dictionary; + # preserve every dictionary snapshot already visible at the + # record's wire position. + if !replacement + stillpending = PendingRecord[] + for p in pending + if header.id in p.missing + p.dictionaries[header.id] = decoded + delete!(p.missing, header.id) + end + if isempty(p.missing) + batchslots[p.slot] = decoderecord( + p.fm, + fields, + sch, + p.dictionaries, + fielddictids, + limits, + validated_dictionaries, + state, + ) + else + push!(stillpending, p) + end + end + pending = stillpending + end + elseif header isa Meta.RecordBatch + missing = missingdicts(fields, header.nodes, dicts, fielddictids) + push!(batchslots, nothing) + slot = length(batchslots) + if isempty(missing) + batchslots[slot] = decoderecord( + fm, + fields, + sch, + dicts, + fielddictids, + limits, + validated_dictionaries, + state, + ) + else + push!(pending, PendingRecord(fm, copy(dicts), missing, slot)) + end + else + throw(ValidationError("unsupported IPC message header $(typeof(header))")) + end + end + isempty(pending) || throw( + ValidationError("stream ended before required dictionary batches arrived"), + ) + batches = AC.RecordBatch[b::AC.RecordBatch for b in batchslots] + return IPCStream( + sch, + AC.FrozenVector{Field}(fields), + batches, + region, + 1, + false, + fielddictids, + budget, + limits, + ) + finally + close(state) + end +end diff --git a/src/ipc_write.jl b/src/ipc_write.jl new file mode 100644 index 00000000..b218bff2 --- /dev/null +++ b/src/ipc_write.jl @@ -0,0 +1,1764 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 IPC writer, plus the file-format reader (`readfile`), over ArrowCore. +# +# * Batch encode is the inverse of decode, one implementation: ONE generic +# `encodefield!` walks the SAME `layoutspec` registry the decoder walks — +# node, registry buffers in registry order, children in declared order. +# There are no per-layout write methods to skew against the read side. +# +# * Dictionary state machine, replacement-on-change: each batch's pools +# are captured by identity. A dictionary batch is emitted before the +# first record batch that references its id and again ONLY when a later +# batch's pool for that id is a different snapshot. Replacement streams +# declare Feature.DICTIONARY_REPLACEMENT in the schema. +# +# * Compression at encode: per-buffer LZ4_FRAME/ZSTD with the Int64 +# uncompressed-length prefix, the `-1` stored-raw fallback when +# compression does not help, codec objects owned per writer and +# explicitly finalized. Compressed streams declare Feature.COMPRESSED_BODY; +# the reader accepts compressed batches with or without the declaration. +# +# * File format = stream framing + a Block index + a Footer: the writer +# isolates footer bookkeeping from generic message writing; `readfile` +# exposes the footer's record-batch index as a lazy random-access handle +# (`length`/`getindex`) over one borrowed or mmapped region. +# ============================================================================= + +# --------------------------------------------------------------------------- +# Encode-side codec state: per-writer objects, explicitly finalized (both +# codecs share TranscodingStreams' one streams API) +# --------------------------------------------------------------------------- + +mutable struct EncodeState + lz4::Union{Nothing,LZ4FrameCompressor} + zstd::Union{Nothing,ZstdCompressor} +end +EncodeState() = EncodeState(nothing, nothing) + +function _lz4c!(s::EncodeState) + if s.lz4 === nothing + c = LZ4FrameCompressor() + TS.initialize(c) + s.lz4 = c + end + return s.lz4::LZ4FrameCompressor +end + +function _zstdc!(s::EncodeState) + if s.zstd === nothing + c = ZstdCompressor() + TS.initialize(c) + s.zstd = c + end + return s.zstd::ZstdCompressor +end + +function Base.close(s::EncodeState) + lz4 = s.lz4 + s.lz4 = nothing + try + lz4 === nothing || TS.finalize(lz4) + finally + zstd = s.zstd + s.zstd = nothing + zstd === nothing || TS.finalize(zstd) + end + return nothing +end + +# --------------------------------------------------------------------------- +# Metadata building: Core descriptors -> Meta tables (inverse of `coretype`) +# --------------------------------------------------------------------------- + +_metatimeunit(u) = + u == AC.SECOND ? Meta.TimeUnit.SECOND : + u == AC.MILLISECOND ? Meta.TimeUnit.MILLISECOND : + u == AC.MICROSECOND ? Meta.TimeUnit.MICROSECOND : Meta.TimeUnit.NANOSECOND + +""" +Build the flatbuffer TYPE table for one Core descriptor. Returns +`(tag type, table offset)` for `fieldAddTypeType`/`fieldAddType`. The isa +ladder is the encode half of `coretype`; a descriptor outside the mapped set +is a clean writer refusal, mirroring the reader's refusal of unmapped tags. +""" +function metatype!(b::FB.Builder, t::ArrowType) + if t isa IntType + Meta.intStart(b) + Meta.intAddBitWidth(b, Int32(t.bits)) + Meta.intAddIsSigned(b, t.signed) + return Meta.Int, Meta.intEnd(b) + elseif t isa FloatType + Meta.floatingPointStart(b) + Meta.floatingPointAddPrecision( + b, + t.bits == 16 ? Meta.Precision.HALF : + t.bits == 32 ? Meta.Precision.SINGLE : Meta.Precision.DOUBLE, + ) + return Meta.FloatingPoint, Meta.floatingPointEnd(b) + elseif t isa BoolType + Meta.boolStart(b) + return Meta.Bool, Meta.boolEnd(b) + elseif t isa Utf8Type + if t.large + Meta.largeUtf8Start(b) + return Meta.LargeUtf8, Meta.largeUtf8End(b) + end + Meta.utf8Start(b) + return Meta.Utf8, Meta.utf8End(b) + elseif t isa BinaryType + if t.large + Meta.largeBinaryStart(b) + return Meta.LargeBinary, Meta.largeBinaryEnd(b) + end + Meta.binaryStart(b) + return Meta.Binary, Meta.binaryEnd(b) + elseif t isa FixedSizeBinaryType + Meta.fixedSizeBinaryStart(b) + Meta.fixedSizeBinaryAddByteWidth(b, Int32(t.nbytes)) + return Meta.FixedSizeBinary, Meta.fixedSizeBinaryEnd(b) + elseif t isa ListType + if t.large + Meta.largeListStart(b) + return Meta.LargeList, Meta.largeListEnd(b) + end + Meta.listStart(b) + return Meta.List, Meta.listEnd(b) + elseif t isa FixedSizeListType + Meta.fixedSizeListStart(b) + Meta.fixedSizeListAddListSize(b, Int32(t.listsize)) + return Meta.FixedSizeList, Meta.fixedSizeListEnd(b) + elseif t isa StructType + Meta.structStart(b) + return Meta.Struct, Meta.structEnd(b) + elseif t isa MapType + Meta.mapStart(b) + t.keyssorted && Meta.mapAddKeysSorted(b, true) + return Meta.Map, Meta.mapEnd(b) + elseif t isa DateType + Meta.dateStart(b) + Meta.dateAddUnit( + b, + t.unit == AC.DAY ? Meta.DateUnit.DAY : Meta.DateUnit.MILLISECOND, + ) + return Meta.Date, Meta.dateEnd(b) + elseif t isa TimeType + Meta.timeStart(b) + Meta.timeAddUnit(b, _metatimeunit(t.unit)) + Meta.timeAddBitWidth(b, Int32(t.bits)) + return Meta.Time, Meta.timeEnd(b) + elseif t isa TimestampType + tz = t.timezone === nothing ? FB.UOffsetT(0) : FB.createstring!(b, t.timezone) + Meta.timestampStart(b) + Meta.timestampAddUnit(b, _metatimeunit(t.unit)) + tz == 0 || Meta.timestampAddTimezone(b, tz) + return Meta.Timestamp, Meta.timestampEnd(b) + elseif t isa DurationType + Meta.durationStart(b) + Meta.durationAddUnit(b, _metatimeunit(t.unit)) + return Meta.Duration, Meta.durationEnd(b) + elseif t isa DecimalType + Meta.decimalStart(b) + Meta.decimalAddPrecision(b, Int32(t.precision)) + Meta.decimalAddScale(b, Int32(t.scale)) + Meta.decimalAddBitWidth(b, Int32(t.bits)) + return Meta.Decimal, Meta.decimalEnd(b) + elseif t isa IntervalType + Meta.intervalStart(b) + Meta.intervalAddUnit( + b, + t.unit == AC.YEAR_MONTH ? Meta.IntervalUnit.YEAR_MONTH : + t.unit == AC.DAY_TIME ? Meta.IntervalUnit.DAY_TIME : + Meta.IntervalUnit.MONTH_DAY_NANO, + ) + return Meta.Interval, Meta.intervalEnd(b) + elseif t isa UnionType + # The builder writes last-first, so every vector is prepended in + # reverse. A zero UOffsetT means "absent": the slot is left out and + # the getter returns the schema default. + Meta.unionStartTypeIdsVector(b, length(t.typeids)) + foreach(x -> FB.prepend!(b, Int32(x)), Iterators.reverse(t.typeids)) + idvec = FB.endvector!(b, length(t.typeids)) + Meta.unionStart(b) + Meta.unionAddMode( + b, + t.mode == AC.DenseMode ? Meta.UnionMode.Dense : Meta.UnionMode.Sparse, + ) + Meta.unionAddTypeIds(b, idvec) + return Meta.Union, Meta.unionEnd(b) + elseif t isa ViewType + if t.utf8 + Meta.utf8ViewStart(b) + return Meta.Utf8View, Meta.utf8ViewEnd(b) + end + Meta.binaryViewStart(b) + return Meta.BinaryView, Meta.binaryViewEnd(b) + elseif t isa ListViewType + if t.large + Meta.largeListViewStart(b) + return Meta.LargeListView, Meta.largeListViewEnd(b) + end + Meta.listViewStart(b) + return Meta.ListView, Meta.listViewEnd(b) + elseif t isa RunEndEncodedType + Meta.runEndEncodedStart(b) + return Meta.RunEndEncoded, Meta.runEndEncodedEnd(b) + elseif t isa NullType + Meta.nullStart(b) + return Meta.Null, Meta.nullEnd(b) + else + throw( + ValidationError( + "IPC writer does not map descriptor " * "$(AC.descriptorname(t))", + ), + ) + end +end + +function _metakeyvalues!(b::FB.Builder, metadata) + metadata === nothing && return FB.UOffsetT(0) + # Core metadata is an ordered Pair sequence. Preserve both that order and + # duplicate keys; callers that want canonical ordering can supply it. + pairs = collect(metadata) + kvs = FB.UOffsetT[] + for (k, v) in pairs + key = FB.createstring!(b, k) + val = FB.createstring!(b, v) + Meta.keyValueStart(b) + Meta.keyValueAddKey(b, key) + Meta.keyValueAddValue(b, val) + push!(kvs, Meta.keyValueEnd(b)) + end + FB.startvector!(b, 4, length(kvs), 4) + foreach(x -> FB.prependoffset!(b, x), Iterators.reverse(kvs)) + return FB.endvector!(b, length(kvs)) +end + +""" +Build the flatbuffer Field table for one Core Field (inverse of `corefield`). +A `DictionaryType` field writes its VALUE type into the type slots and its +index/id/ordering into a DictionaryEncoding table; the id comes from the +writer's side table — Core fields still never carry one. +""" +function metafield!(b::FB.Builder, f::Field, fielddictids::IdDict{Field,Int64}) + t = f.type + valuetype = t + dictoff = FB.UOffsetT(0) + if t isa DictionaryType + valuetype = t.valuetype + idxtag, idxoff = metatype!(b, t.indextype) + idxtag === Meta.Int || + throw(ValidationError("dictionary index type must be an integer")) + Meta.dictionaryEncodingStart(b) + Meta.dictionaryEncodingAddId(b, fielddictids[f]) + Meta.dictionaryEncodingAddIndexType(b, idxoff) + t.ordered && Meta.dictionaryEncodingAddIsOrdered(b, true) + dictoff = Meta.dictionaryEncodingEnd(b) + end + children = FB.UOffsetT[metafield!(b, c, fielddictids) for c in f.children] + Meta.fieldStartChildrenVector(b, length(children)) + foreach(x -> FB.prependoffset!(b, x), Iterators.reverse(children)) + childvec = FB.endvector!(b, length(children)) + kvvec = _metakeyvalues!(b, f.metadata) + name = FB.createstring!(b, f.name) + tag, typeoff = metatype!(b, valuetype) + Meta.fieldStart(b) + Meta.fieldAddName(b, name) + Meta.fieldAddNullable(b, f.nullable) + Meta.fieldAddTypeType(b, tag) + Meta.fieldAddType(b, typeoff) + dictoff == 0 || Meta.fieldAddDictionary(b, dictoff) + Meta.fieldAddChildren(b, childvec) + kvvec == 0 || Meta.fieldAddCustomMetadata(b, kvvec) + return Meta.fieldEnd(b) +end + +_pad8!(bytes::Vector{UInt8}) = append!(bytes, zeros(UInt8, mod(-length(bytes), 8))) + +""" +Finish the current builder content as one framed message: continuation +marker, padded metadata length, metadata, then the (already padded) body. +""" +function _finishmessage!(out::Vector{UInt8}, b::FB.Builder, msg, body::Vector{UInt8}) + FB.finish!(b, msg) + meta = collect(FB.finishedbytes(b)) + _pad8!(meta) + length(body) % 8 == 0 || throw(ArgumentError("message body must be padded")) + append!(out, reinterpret(UInt8, UInt32[CONTINUATION, UInt32(length(meta))])) + append!(out, meta) + append!(out, body) + return out +end + +function _metaschema!( + b::FB.Builder, + sch::Schema, + fielddictids::IdDict{Field,Int64}, + features::Vector{Int64}, +) + fields = FB.UOffsetT[metafield!(b, f, fielddictids) for f in sch.fields] + Meta.schemaStartFieldsVector(b, length(fields)) + foreach(x -> FB.prependoffset!(b, x), Iterators.reverse(fields)) + fieldvec = FB.endvector!(b, length(fields)) + kvvec = _metakeyvalues!(b, sch.metadata) + featurevec = FB.UOffsetT(0) + if !isempty(features) + FB.startvector!(b, 8, length(features), 8) + foreach(x -> FB.prepend!(b, x), Iterators.reverse(features)) + featurevec = FB.endvector!(b, length(features)) + end + Meta.schemaStart(b) + Meta.schemaAddEndianness(b, Meta.Endianness.Little) + Meta.schemaAddFields(b, fieldvec) + kvvec == 0 || Meta.schemaAddCustomMetadata(b, kvvec) + featurevec == 0 || Meta.schemaAddFeatures(b, featurevec) + return Meta.schemaEnd(b) +end + +function _schemamessage!( + out::Vector{UInt8}, + sch::Schema, + fielddictids::IdDict{Field,Int64}, + features::Vector{Int64}, +) + b = FB.Builder(1024) + schoff = _metaschema!(b, sch, fielddictids, features) + Meta.messageStart(b) + Meta.messageAddVersion(b, Meta.MetadataVersion.V5) + Meta.messageAddHeaderType(b, Meta.Schema) + Meta.messageAddHeader(b, schoff) + return _finishmessage!(out, b, Meta.messageEnd(b), UInt8[]) +end + +# --------------------------------------------------------------------------- +# THE generic encoder: registry-driven node/buffer emission +# --------------------------------------------------------------------------- + +mutable struct EncodeCursor + nodes::Vector{NTuple{2,Int64}} # (length, null_count), forward order + buffers::Vector{NTuple{2,Int64}} # (offset, length), forward order + body::Vector{UInt8} + codec::Int8 + state::Union{Nothing,EncodeState} + variadics::Vector{Int64} # per view field, depth-first order +end +EncodeCursor(codec::Int8, state::Union{Nothing,EncodeState}) = + EncodeCursor(NTuple{2,Int64}[], NTuple{2,Int64}[], UInt8[], codec, state, Int64[]) + +function _compressbytes(state::EncodeState, codec::Int8, raw::Vector{UInt8}) + codec == CODEC_LZ4_FRAME && return transcode(_lz4c!(state), raw) + return transcode(_zstdc!(state), raw) +end + +""" +Append one buffer to the message body: raw bytes for uncompressed batches; +for compressed batches, the spec's Int64 uncompressed-length prefix plus the +frame, falling back to `-1` + raw whenever compression does not shrink the +payload. Zero-length buffers write no body bytes in either mode. The recorded +Buffer length is the wire length; 8-byte alignment padding sits between +buffers and belongs to neither. +""" +function encodebuffer!(c::EncodeCursor, bytes::Vector{UInt8}) + offset = Int64(length(c.body)) + offset % 8 == 0 || throw(ArgumentError("encoder lost body alignment")) + if isempty(bytes) + push!(c.buffers, (offset, Int64(0))) + return nothing + end + if c.codec == CODEC_NONE + append!(c.body, bytes) + push!(c.buffers, (offset, Int64(length(bytes)))) + _pad8!(c.body) + return nothing + end + compressed = _compressbytes(c.state::EncodeState, c.codec, bytes) + if length(compressed) < length(bytes) + append!(c.body, reinterpret(UInt8, Int64[Int64(length(bytes))])) + append!(c.body, compressed) + push!(c.buffers, (offset, Int64(8 + length(compressed)))) + else + append!(c.body, reinterpret(UInt8, Int64[Int64(-1)])) + append!(c.body, bytes) + push!(c.buffers, (offset, Int64(8 + length(bytes)))) + end + _pad8!(c.body) + return nothing +end + +""" + encodefield!(cursor, f, d) + +The write half of the registry walk — the exact mirror of `decodefield`: +one node, then the layout's buffers in registry order, then children in +declared order. Dictionary-encoded fields emit their INDEX buffers here; +their pool travels in a dictionary batch. Buffer content is emitted from the +`ArrayData` slices verbatim: the encoder adds no per-layout interpretation, +so read and write cannot skew. +""" +function encodefield!(c::EncodeCursor, f::Field, d::ArrayData) + t = f.type + AC.typeequal(t, d.type) || + throw(ValidationError("column data type does not match its schema field")) + d.offset == 0 || throw( + ValidationError( + "IPC encode of sliced (nonzero-offset) arrays is not supported; materialize first", + ), + ) + push!(c.nodes, (d.len, AC.nullcount(d))) + spec = layoutspec(t) + if spec.variadic + length(d.buffers) >= length(spec.buffers) || + throw(ValidationError("column buffer count does not match its layout")) + else + length(d.buffers) == length(spec.buffers) || + throw(ValidationError("column buffer count does not match its layout")) + end + for (role, b) in zip(spec.buffers, d.buffers) + if role == AC.OFFSETS && d.len == 0 && b.len == 0 + # Core canonicalizes an empty offset array without allocating its + # otherwise-unused physical buffer. IPC still requires the one + # terminal zero offset (length + 1 entries). + encodebuffer!(c, zeros(UInt8, spec.offsetwidth)) + else + encodebuffer!(c, AC.slicebytes(b)) + end + end + if spec.variadic + # View layouts append their variadic data buffers after the fixed + # validity/views pair; the count travels in the header's + # variadicBufferCounts vector, depth-first (format 1.4). + push!(c.variadics, Int64(length(d.buffers) - length(spec.buffers))) + for b in Iterators.drop(d.buffers, length(spec.buffers)) + encodebuffer!(c, AC.slicebytes(b)) + end + end + t isa DictionaryType && return nothing + nchildren = spec.childcount == -1 ? length(f.children) : spec.childcount + length(d.children) == nchildren || + throw(ValidationError("column child count does not match its schema field")) + if t isa UnionType && t.mode == AC.SparseMode + all(child -> child.len == d.len, d.children) || + throw(ValidationError("IPC sparse-union children must equal the union length")) + end + for i = 1:nchildren + encodefield!(c, f.children[i], d.children[i]) + end + return nothing +end + +function _batchheader!(b::FB.Builder, c::EncodeCursor, nrows::Int64) + Meta.recordBatchStartNodesVector(b, length(c.nodes)) + for (len, nulls) in Iterators.reverse(c.nodes) + Meta.createFieldNode(b, len, nulls) + end + nodes = FB.endvector!(b, length(c.nodes)) + Meta.recordBatchStartBuffersVector(b, length(c.buffers)) + for (off, len) in Iterators.reverse(c.buffers) + Meta.createBuffer(b, off, len) + end + buffers = FB.endvector!(b, length(c.buffers)) + compression = FB.UOffsetT(0) + if c.codec != CODEC_NONE + Meta.bodyCompressionStart(b) + Meta.bodyCompressionAddCodec( + b, + c.codec == CODEC_LZ4_FRAME ? Meta.CompressionType.LZ4_FRAME : + Meta.CompressionType.ZSTD, + ) + compression = Meta.bodyCompressionEnd(b) + end + varvec = FB.UOffsetT(0) + if !isempty(c.variadics) + Meta.recordBatchStartVariadicBufferCountsVector(b, length(c.variadics)) + foreach(x -> FB.prepend!(b, x), Iterators.reverse(c.variadics)) + varvec = FB.endvector!(b, length(c.variadics)) + end + Meta.recordBatchStart(b) + Meta.recordBatchAddLength(b, nrows) + Meta.recordBatchAddNodes(b, nodes) + Meta.recordBatchAddBuffers(b, buffers) + compression == 0 || Meta.recordBatchAddCompression(b, compression) + varvec == 0 || Meta.recordBatchAddVariadicBufferCounts(b, varvec) + return Meta.recordBatchEnd(b) +end + +function _recordmessage!( + out::Vector{UInt8}, + batch::AC.RecordBatch, + fields, + codec::Int8, + state::Union{Nothing,EncodeState}, +) + c = EncodeCursor(codec, state) + for (f, col) in zip(fields, batch.columns) + encodefield!(c, f, col) + end + b = FB.Builder(1024) + rb = _batchheader!(b, c, batch.nrows) + Meta.messageStart(b) + Meta.messageAddVersion(b, Meta.MetadataVersion.V5) + Meta.messageAddHeaderType(b, Meta.RecordBatch) + Meta.messageAddHeader(b, rb) + Meta.messageAddBodyLength(b, Int64(length(c.body))) + return _finishmessage!(out, b, Meta.messageEnd(b), c.body) +end + +function _dictionarymessage!( + out::Vector{UInt8}, + id::Int64, + vf::Field, + pool::ArrayData, + codec::Int8, + state::Union{Nothing,EncodeState}, +) + c = EncodeCursor(codec, state) + encodefield!(c, vf, pool) + b = FB.Builder(1024) + rb = _batchheader!(b, c, pool.len) + Meta.dictionaryBatchStart(b) + Meta.dictionaryBatchAddId(b, id) + Meta.dictionaryBatchAddData(b, rb) + dictbatch = Meta.dictionaryBatchEnd(b) + Meta.messageStart(b) + Meta.messageAddVersion(b, Meta.MetadataVersion.V5) + Meta.messageAddHeaderType(b, Meta.DictionaryBatch) + Meta.messageAddHeader(b, dictbatch) + Meta.messageAddBodyLength(b, Int64(length(c.body))) + return _finishmessage!(out, b, Meta.messageEnd(b), c.body) +end + +# --------------------------------------------------------------------------- +# Stream writer driver +# --------------------------------------------------------------------------- + +const CODEC_NAMES = + Dict{Symbol,Int8}(:none => CODEC_NONE, :lz4 => CODEC_LZ4_FRAME, :zstd => CODEC_ZSTD) + +function _requirelittleendian(host_endian_bom::UInt32=Base.ENDIAN_BOM) + host_endian_bom == UInt32(0x04030201) || + throw(ValidationError("the IPC writer requires a little-endian host")) + return nothing +end + +""" +Assign one IPC dictionary id per dictionary-typed field, depth-first over the +schema — the writer-side half of the adapter id table (ids are adapter +bookkeeping; Core fields never carry them). +""" +function assigndictids(fields, given::IdDict{Field,Int64}=IdDict{Field,Int64}()) + # `given` lets a caller preserve ids from a source (a reader's table): two + # fields sharing one id then share one dictionary batch, exactly as the + # source did (the arrow-testing 4.0.0-shareddict integration case). Fresh + # ids fill the lowest unoccupied values so they never collide with given + # ones — including given ids at the top of the signed-long domain, where + # `max + 1` would wrap. + ids = IdDict{Field,Int64}(given) + seen = IdDict{Field,Nothing}() + used = Set{Int64}(values(ids)) + next = Ref(Int64(0)) + function freshid() + while next[] in used + next[] < typemax(Int64) || + throw(ValidationError("IPC dictionary id space is exhausted")) + next[] += 1 + end + push!(used, next[]) + return next[] + end + function walk(f::Field) + haskey(seen, f) && throw( + ValidationError( + "IPC writer schema reuses one Field object in multiple positions", + ), + ) + seen[f] = nothing + if f.type isa DictionaryType && !haskey(ids, f) + ids[f] = freshid() + end + foreach(walk, f.children) + end + foreach(walk, fields) + return ids +end + +""" +Collect `(field, pool)` pairs for every dictionary-typed field in one batch, +paired with the schema walk (nested dictionaries included). +""" +function dictionarypools(fields, cols) + pairs = Tuple{Field,ArrayData}[] + function walk(f::Field, d::ArrayData) + if f.type isa DictionaryType + d.dictionary === nothing && + throw(ValidationError("dictionary column carries no pool")) + # Post-order: pools nested INSIDE this pool's values are collected + # (and therefore emitted) before it — the dependency order the + # IPC spec requires for nested dictionary encoding. + walk(AC.dictvaluefield(f, f.type), d.dictionary) + push!(pairs, (f, d.dictionary)) + return + end + for (cf, cd) in zip(f.children, d.children) + walk(cf, cd) + end + end + for (f, d) in zip(fields, cols) + walk(f, d) + end + return pairs +end + +function _checkbatches(sch::Schema, batches) + for batch in batches + length(batch.columns) == length(sch.fields) || + throw(ValidationError("batch column count does not match the schema")) + for (f, bf) in zip(sch.fields, batch.schema.fields) + AC.typeequal(f.type, bf.type) && f.name == bf.name || + throw(ValidationError("batch schema does not match the stream schema")) + end + end + return nothing +end + +function _validatewriterschema(sch::Schema) + AC._validate_schema(sch) + function walk(f::Field) + isvalid(f.name) || throw(ValidationError("field name is not valid UTF-8")) + AC._validate_metadata(f.metadata, "field") + foreach(walk, f.children) + return nothing + end + foreach(walk, sch.fields) + foreach(validateschemafield, sch.fields) + return nothing +end + +function _validatewriterbatch!( + validated::AC._ValidatedDictionaries, + sch::Schema, + batch::AC.RecordBatch, + ids::IdDict{Field,Int64}, +) + # A shared immutable pool must satisfy every value-field contract + # through which the schema refers to it. Identity caching is safe only + # after those field-specific checks have run. + current = Dict{Int64,ArrayData}() + for (f, pool) in dictionarypools(sch.fields, batch.columns) + # One id names ONE pool within a record batch: every dictionary + # message precedes the record message on the wire, so an + # intra-batch pool change is not temporal replacement — it would + # silently retarget the earlier field to the later pool. + id = ids[f] + haskey(current, id) && + current[id] !== pool && + throw( + ValidationError( + "dictionary id $id carries two different pools in one record batch", + ), + ) + current[id] = pool + validate_semantic(AC.dictvaluefield(f, f.type::DictionaryType), pool) + validated[pool] = nothing + end + for (f, col) in zip(sch.fields, batch.columns) + AC._validate_semantic(f, col, validated) + end + return nothing +end + +function _validatewriterbatches(sch::Schema, batches, ids::IdDict{Field,Int64}) + validated = AC._ValidatedDictionaries() + for batch in batches + _validatewriterbatch!(validated, sch, batch, ids) + end + return nothing +end + +const FEATURE_DICTIONARY_REPLACEMENT = Int64(1) +const FEATURE_COMPRESSED_BODY = Int64(2) + +""" +Which features must the schema declare for these batches? Replacement is +detected by pool-identity change per id across the batch sequence +(replacement-on-change); compression declares COMPRESSED_BODY. +""" +function _streamfeatures(sch::Schema, batches, ids::IdDict{Field,Int64}, codec::Int8) + features = Int64[] + current = Dict{Int64,ArrayData}() + replacement = false + for batch in batches + for (f, pool) in dictionarypools(sch.fields, batch.columns) + id = ids[f] + old = get(current, id, nothing) + old === pool || old === nothing || (replacement = true) + current[id] = pool + end + end + replacement && push!(features, FEATURE_DICTIONARY_REPLACEMENT) + isempty(batches) || codec == CODEC_NONE || push!(features, FEATURE_COMPRESSED_BODY) + return features +end + +# --------------------------------------------------------------------------- +# Incremental writer core: begin/writebatch/finish over one IO +# --------------------------------------------------------------------------- + +""" +One in-progress IPC output. `beginwrite!` emits the schema message (and the +file-format preamble), `writebatch!` validates and publishes one record batch +(with any dictionary messages it needs), and `finishwrite!` emits the +end-of-stream marker (and the file-format Footer). Batch bytes are staged in +full and published to `io` once, so a validation or encoding failure never +publishes a partial message — but the SINK gains bytes batch by batch, and an +abandoned state leaves a torn stream or an unfooted file. + +The eager `writestream`/`writefile` drive this same state over an in-memory +sink, validating every batch up front (`validate=false` per batch) and +passing the features their whole batch sequence implies, so their output is +byte-identical to the pre-incremental writers. +""" +mutable struct IPCWriteState + const io::IO + const file::Bool + const schema::Schema + const fielddictids::IdDict{Field,Int64} + const ids::IdDict{Field,Int64} + const features::Vector{Int64} + const codec::Int8 + const state::Union{Nothing,EncodeState} + const validated::AC._ValidatedDictionaries + const current::Dict{Int64,ArrayData} # last pool per id, for replacement + const dictblocks::Vector{NTuple{3,Int64}} # file format Footer Blocks + const recordblocks::Vector{NTuple{3,Int64}} + written::Int64 # bytes published to `io` + finished::Bool +end + +# Batch sequences are unknown up front. A declared feature is a demand on +# readers (nanoarrow refuses any stream declaring one it does not support), +# so nothing is declared speculatively: pool replacement needs the caller's +# explicit opt-in, and a codec declares compressed bodies because every +# batch this state writes will use it. +function _incrementalfeatures( + ids::IdDict{Field,Int64}, + codec::Int8, + file::Bool, + dictreplacement::Bool, +) + features = Int64[] + dictreplacement && + !file && + !isempty(ids) && + push!(features, FEATURE_DICTIONARY_REPLACEMENT) + codec == CODEC_NONE || push!(features, FEATURE_COMPRESSED_BODY) + return features +end + +function _publish!(st::IPCWriteState, out::Vector{UInt8}) + Base.write(st.io, out) + st.written += length(out) + return nothing +end + +function beginwrite!( + io::IO, + sch::Schema; + file::Bool, + compress::Symbol=:none, + dictids::IdDict{Field,Int64}=IdDict{Field,Int64}(), + features::Union{Nothing,Vector{Int64}}=nothing, + dictreplacement::Bool=false, +) + _requirelittleendian() + haskey(CODEC_NAMES, compress) || + throw(ArgumentError("compress must be :none, :lz4, or :zstd")) + codec = CODEC_NAMES[compress] + _validatewriterschema(sch) + ids = assigndictids(sch.fields, dictids) + fielddictids = IdDict{Field,Int64}(ids) + validatedictionaryids(sch.fields, fielddictids) + feats = + features === nothing ? _incrementalfeatures(ids, codec, file, dictreplacement) : + features + st = IPCWriteState( + io, + file, + sch, + fielddictids, + ids, + feats, + codec, + codec == CODEC_NONE ? nothing : EncodeState(), + AC._ValidatedDictionaries(), + Dict{Int64,ArrayData}(), + NTuple{3,Int64}[], + NTuple{3,Int64}[], + Int64(0), + false, + ) + out = UInt8[] + if file + append!(out, FILE_MAGIC) + append!(out, zeros(UInt8, 2)) # pad to 8 before the first message + end + _schemamessage!(out, sch, fielddictids, feats) + _publish!(st, out) + return st +end + +# Record one file-format Block around `emit!`: the Block offset is global +# (bytes already published plus this staging buffer's position), and +# metaDataLength spans prefix + metadata (up to the body start). +function _fileblock!(emit!::F, st::IPCWriteState, out::Vector{UInt8}, blocks) where {F} + start = Int64(length(out)) + emit!() + total = Int64(length(out)) - start + metalen = Int64(8) + Int64(reinterpret(UInt32, out[(start + 5):(start + 8)])[1]) + push!(blocks, (st.written + start, metalen, total - metalen)) + return nothing +end + +"Order-sensitive pool content equality, in each pool's public value domain." +function _poolsequal(vf::Field, a::ArrayData, b::ArrayData) + a.len == b.len || return false + return isequal(AC.materialize(vf, a), AC.materialize(vf, b)) +end + +function writebatch!(st::IPCWriteState, batch::AC.RecordBatch; validate::Bool=true) + st.finished && throw(ArgumentError("this IPC writer is already finished")) + if validate + _checkbatches(st.schema, (batch,)) + _validatewriterbatch!(st.validated, st.schema, batch, st.ids) + end + out = UInt8[] + for (f, pool) in dictionarypools(st.schema.fields, batch.columns) + id = st.ids[f] + old = get(st.current, id, nothing) + old === pool && continue + if old !== nothing + vf = AC.dictvaluefield(f, f.type::DictionaryType) + if st.file + # The file format carries one dictionary batch per id. A new + # pool object with identical content references the emitted + # batch; anything else is replacement, which files cannot say. + _poolsequal(vf, old, pool) || throw( + ValidationError( + "the IPC file format carries one dictionary batch per " * + "id; dictionary id $id changed pools across batches — " * + "use the stream format for replacement", + ), + ) + st.current[id] = pool + continue + end + if !(FEATURE_DICTIONARY_REPLACEMENT in st.features) + # Emitting a replacement the schema message did not declare + # would publish a stream this package's own reader refuses. + _poolsequal(vf, old, pool) || throw( + ValidationError( + "dictionary id $id changed pools, but the stream's " * + "schema message does not declare the " * + "DictionaryReplacement feature", + ), + ) + st.current[id] = pool + continue + end + end + vf = AC.dictvaluefield(f, f.type::DictionaryType) + if st.file + _fileblock!(st, out, st.dictblocks) do + _dictionarymessage!(out, id, vf, pool, st.codec, st.state) + end + else + _dictionarymessage!(out, id, vf, pool, st.codec, st.state) + end + st.current[id] = pool + end + if st.file + _fileblock!(st, out, st.recordblocks) do + _recordmessage!(out, batch, st.schema.fields, st.codec, st.state) + end + else + _recordmessage!(out, batch, st.schema.fields, st.codec, st.state) + end + _publish!(st, out) + return nothing +end + +# The file-format trailer: Footer flatbuffer (schema again, then the two +# Block struct-vectors), the Int32 footer length, and the trailing magic. +function _filefooter!( + out::Vector{UInt8}, + sch::Schema, + fielddictids::IdDict{Field,Int64}, + features::Vector{Int64}, + dictblocks::Vector{NTuple{3,Int64}}, + recordblocks::Vector{NTuple{3,Int64}}, +) + b = FB.Builder(1024) + schoff = _metaschema!(b, sch, fielddictids, features) + Meta.footerStartDictionariesVector(b, length(dictblocks)) + for (off, metalen, bodylen) in Iterators.reverse(dictblocks) + Meta.createBlock(b, off, Int32(metalen), bodylen) + end + dictvec = FB.endvector!(b, length(dictblocks)) + Meta.footerStartRecordBatchesVector(b, length(recordblocks)) + for (off, metalen, bodylen) in Iterators.reverse(recordblocks) + Meta.createBlock(b, off, Int32(metalen), bodylen) + end + recordvec = FB.endvector!(b, length(recordblocks)) + Meta.footerStart(b) + Meta.footerAddVersion(b, Meta.MetadataVersion.V5) + Meta.footerAddSchema(b, schoff) + Meta.footerAddDictionaries(b, dictvec) + Meta.footerAddRecordBatches(b, recordvec) + FB.finish!(b, Meta.footerEnd(b)) + footer = collect(FB.finishedbytes(b)) + append!(out, footer) + append!(out, reinterpret(UInt8, Int32[Int32(length(footer))])) + append!(out, FILE_MAGIC) + return nothing +end + +function finishwrite!(st::IPCWriteState) + st.finished && return nothing + out = UInt8[] + append!(out, reinterpret(UInt8, UInt32[CONTINUATION, UInt32(0)])) + st.file && _filefooter!( + out, + st.schema, + st.fielddictids, + st.features, + st.dictblocks, + st.recordblocks, + ) + _publish!(st, out) + st.finished = true + st.state === nothing || close(st.state) + return nothing +end + +"Release the codec state without emitting anything; for error cleanup." +function abortwrite!(st::IPCWriteState) + st.finished && return nothing + st.finished = true + st.state === nothing || close(st.state) + return nothing +end + +""" +Resume a STREAM whose schema message is already on the wire: a state that +emits no preamble. `features` must be the features that schema message +declared, and `current` the last pool each dictionary id has emitted — +`writebatch!` then skips content-identical pools and gates replacement on +the declared features exactly as it does mid-stream. +""" +function resumestream!( + io::IO, + sch::Schema, + dictids::IdDict{Field,Int64}, + features::Vector{Int64}; + compress::Symbol=:none, + current::Dict{Int64,ArrayData}=Dict{Int64,ArrayData}(), +) + _requirelittleendian() + haskey(CODEC_NAMES, compress) || + throw(ArgumentError("compress must be :none, :lz4, or :zstd")) + codec = CODEC_NAMES[compress] + ids = assigndictids(sch.fields, dictids) + return IPCWriteState( + io, + false, + sch, + IdDict{Field,Int64}(ids), + ids, + features, + codec, + codec == CODEC_NONE ? nothing : EncodeState(), + AC._ValidatedDictionaries(), + current, + NTuple{3,Int64}[], + NTuple{3,Int64}[], + Int64(0), + false, + ) +end + +""" + writestream(sch, batches; compress=:none, dictids=IdDict{Field,Int64}()) -> Vector{UInt8} + writestream(stream::IPCStream; compress=:none) -> Vector{UInt8} + +Encode a complete IPC stream: schema message, dictionary batches emitted +before the first record batch that references them (and again on +pool-identity change), record batches, end-of-stream marker. Every column is +semantically validated before any of its bytes are emitted — the writer +refuses to publish data Core would refuse to read. `dictids` preserves ids +from a source table (the `IPCStream` form passes the reader's), so fields +sharing one id keep sharing one dictionary batch, exactly as the source did. +""" +function writestream( + sch::Schema, + batches::AbstractVector{AC.RecordBatch}; + compress::Symbol=:none, + dictids::IdDict{Field,Int64}=IdDict{Field,Int64}(), +) + _requirelittleendian() + haskey(CODEC_NAMES, compress) || + throw(ArgumentError("compress must be :none, :lz4, or :zstd")) + codec = CODEC_NAMES[compress] + _checkbatches(sch, batches) + _validatewriterschema(sch) + # `beginwrite!` recomputes the same ids from `dictids` (assignment is + # deterministic); this pass exists so the whole batch sequence is + # validated — and its EXACT features computed — before any bytes exist. + # A declared feature is a demand on readers, so an eager stream declares + # only what these batches actually use. + ids = assigndictids(sch.fields, dictids) + validatedictionaryids(sch.fields, IdDict{Field,Int64}(ids)) + _validatewriterbatches(sch, batches, ids) + io = IOBuffer() + st = beginwrite!( + io, + sch; + file=false, + compress=compress, + dictids=dictids, + features=_streamfeatures(sch, batches, ids, codec), + ) + try + for batch in batches + writebatch!(st, batch; validate=false) + end + finishwrite!(st) + finally + abortwrite!(st) + end + return take!(io) +end + +writestream(s::IPCStream; compress::Symbol=:none) = + writestream(s.schema, s.batches; compress=compress, dictids=s.fielddictids) + +# --------------------------------------------------------------------------- +# File format: magic + stream messages + Block index + Footer +# --------------------------------------------------------------------------- + +const FILE_MAGIC = b"ARROW1" + +""" + writefile(sch, batches; compress=:none, dictids=IdDict{Field,Int64}()) -> Vector{UInt8} + writefile(stream::IPCStream; compress=:none) -> Vector{UInt8} + +The file variant: leading magic, the same stream messages, an end-of-stream +marker, then the Footer with its dictionary and record-batch Block indexes, +the Int32 footer length, and the trailing magic. Footer bookkeeping is +isolated here; message writing is the stream code above. The file format +carries exactly one dictionary batch per id, so batches whose pools change +identity are a clean refusal (the stream format handles replacement). +`dictids` preserves ids from a source table (the `IPCStream` form passes the +reader's), so fields sharing one id keep sharing one dictionary batch, +exactly as the source did. +""" +function writefile( + sch::Schema, + batches::AbstractVector{AC.RecordBatch}; + compress::Symbol=:none, + dictids::IdDict{Field,Int64}=IdDict{Field,Int64}(), +) + _requirelittleendian() + haskey(CODEC_NAMES, compress) || + throw(ArgumentError("compress must be :none, :lz4, or :zstd")) + codec = CODEC_NAMES[compress] + _checkbatches(sch, batches) + _validatewriterschema(sch) + ids = assigndictids(sch.fields, dictids) + # CODEC_NONE so this refusal sees only DICTIONARY_REPLACEMENT. + # COMPRESSED_BODY is legal in a file and is declared through the features + # below. + isempty(_streamfeatures(sch, batches, ids, CODEC_NONE)) || throw( + ValidationError( + "the IPC file format carries one dictionary batch per id; " * + "a dictionary with changing pools requires the stream format", + ), + ) + # Same shared-id contract as the stream writer: compatible value schemas, + # one nested id topology, one pool per id within each batch. + validatedictionaryids(sch.fields, IdDict{Field,Int64}(ids)) + _validatewriterbatches(sch, batches, ids) + io = IOBuffer() + st = beginwrite!( + io, + sch; + file=true, + compress=compress, + dictids=dictids, + features=_streamfeatures(sch, batches, ids, codec), + ) + try + for batch in batches + writebatch!(st, batch; validate=false) + end + finishwrite!(st) + finally + abortwrite!(st) + end + return take!(io) +end + +writefile(s::IPCStream; compress::Symbol=:none) = + writefile(s.schema, s.batches; compress=compress, dictids=s.fielddictids) + +# --------------------------------------------------------------------------- +# File reader: footer verification + lazy random-access batch handle +# --------------------------------------------------------------------------- + +_blocktuples(blocks) = + blocks === nothing ? NTuple{3,Int64}[] : + NTuple{3,Int64}[ + (Int64(b.offset), Int64(b.metaDataLength), Int64(b.bodyLength)) for b in blocks + ] + +""" +Footer verification, same wrapper role as `verify_ipc_metadata`: the +generated walker bounds the whole table graph, then the verified getters +supply the Block indexes. +""" +function verify_footer( + bytes::Vector{UInt8}, + limits::Limits, + reserve_limit::Int64=limits.max_total_allocated_bytes, +) + ctx = _verifyctx(limits, reserve_limit) + # Staged like verify_ipc_metadata: version gates between the inline and + # reference stages, so an unsupported footer rejects in constant time. + t = _verified(() -> Meta.verifyrootstart_Footer(bytes, ctx)) + footer = FB.getrootas(Meta.Footer, bytes, 0) + version = Int16(Int64(footer.version)) + version in (Int16(3), Int16(4)) || + _vfail("unsupported footer version $version (only V4/V5 are accepted)") + _verified(() -> Meta.verifyrootrest_Footer(t, ctx)) + features = _schemafeatures(footer.schema::Meta.Schema, version) + return version, + features, + _blocktuples(footer.dictionaries), + _blocktuples(footer.recordBatches), + ctx.reserved +end + +function _metadataequal(a, b) + av = something(a, Meta.KeyValue[]) + bv = something(b, Meta.KeyValue[]) + length(av) == length(bv) || return false + for (x, y) in zip(av, bv) + x.key == y.key || return false + something(x.value, "") == something(y.value, "") || return false + end + return true +end + +function _fieldequal(a::Field, b::Field) + a.name == b.name && + a.nullable == b.nullable && + AC.typeequal(a.type, b.type) && + a.metadata == b.metadata && + length(a.children) == length(b.children) || return false + return all(_fieldequal(x, y) for (x, y) in zip(a.children, b.children)) +end + +function _schemaside(metaschema::Meta.Schema) + dictids = Dict{Int64,Meta.Field}() + fielddictids = IdDict{Field,Int64}() + fields = Field[ + corefield(f, dictids, fielddictids) for + f in something(metaschema.fields, Meta.Field[]) + ] + foreach(validateschemafield, fields) + valueschemas = validatedictionaryids(fields, fielddictids) + length(valueschemas) == length(dictids) || + throw(ValidationError("duplicate dictionary id in file schema")) + return fields +end + +function _fieldwireequal(a::Meta.Field, b::Meta.Field) + _metadataequal(a.custom_metadata, b.custom_metadata) || return false + adict, bdict = a.dictionary, b.dictionary + (adict === nothing) == (bdict === nothing) || return false + if adict !== nothing + adict.id == bdict.id || return false + end + achildren = something(a.children, Meta.Field[]) + bchildren = something(b.children, Meta.Field[]) + length(achildren) == length(bchildren) || return false + return all(_fieldwireequal(x, y) for (x, y) in zip(achildren, bchildren)) +end + +function _schemaequal(a::Meta.Schema, b::Meta.Schema) + something(a.endianness, Meta.Endianness.Little) == + something(b.endianness, Meta.Endianness.Little) || return false + ametafields = something(a.fields, Meta.Field[]) + bmetafields = something(b.fields, Meta.Field[]) + length(ametafields) == length(bmetafields) || return false + all(_fieldwireequal(x, y) for (x, y) in zip(ametafields, bmetafields)) || return false + afields = _schemaside(a) + bfields = _schemaside(b) + length(afields) == length(bfields) || return false + _metadataequal(a.custom_metadata, b.custom_metadata) || return false + all(_fieldequal(x, y) for (x, y) in zip(afields, bfields)) || return false + return true +end + +function _fileschema( + region::OwnerRegion, + footerstart::Int64, + limits::Limits, + budget::AllocationBudget, +) + blob = BufferSlice(region, 0, region.len) + AC.loadat(blob, UInt32, Int64(8)) == CONTINUATION || + throw(ValidationError("file data section does not start with an IPC message")) + declared = Int64(AC.loadat(blob, Int32, Int64(12))) + 0 < declared <= limits.max_metadata_bytes || + throw(ValidationError("file schema metadata length is outside the limit")) + declared % 8 == 0 || + throw(ValidationError("file schema metadata is not 8-byte aligned")) + metalen = AC.checked_add(Int64(8), declared) + fm = _blockmessage(region, (Int64(8), metalen, Int64(0)), footerstart, limits, budget) + fm.header_type == UInt8(1) && fm.msg.header isa Meta.Schema || # Schema + throw(ValidationError("file data section does not start with a schema")) + return fm, AC.checked_add(Int64(8), metalen) +end + +""" + ArrowFile + +The footer's record-batch index as a random-access handle: `length(file)` +batches, `file[i]` decodes batch `i` on demand. Open decodes only the schema +and the dictionary batches every record shares. Each `getindex` decodes from +the mapped bytes with its own allocation budget and codec contexts, and the +handle is read-only after open, so concurrent `getindex` calls are safe. The +region root (heap vector or Mmap array) is the only lifetime anchor, as in +Core. +""" +struct ArrowFile + region::OwnerRegion + schema::Schema + fields::AC.FrozenVector{Field} + fielddictids::IdDict{Field,Int64} + dictionaries::Dict{Int64,ArrayData} + validated::AC._ValidatedDictionaries + recordblocks::Vector{NTuple{3,Int64}} + dataend::Int64 + limits::Limits + schemaversion::Int16 +end + +Base.length(f::ArrowFile) = length(f.recordblocks) +AC.schema(f::ArrowFile) = f.schema + +""" +Range-check one Footer Block's declared extents against the data section and +return its `(offset, frameend)` span. +""" +function _blockextent(block::NTuple{3,Int64}, dataend::Int64) + offset, metalen, bodylen = block + # 16 is the smallest legal prefix + metadata: an 8-byte prefix plus at + # least one 8-byte-aligned metadata block. + (offset >= 0 && metalen >= 16 && bodylen >= 0) || + throw(ValidationError("footer block has invalid extents")) + offset % 8 == 0 || throw(ValidationError("footer block is not 8-byte aligned")) + metalen % 8 == 0 || + throw(ValidationError("footer block metadata length is not 8-byte aligned")) + bodylen % 8 == 0 || + throw(ValidationError("footer block body length is not 8-byte aligned")) + frameend = AC.checked_add(AC.checked_add(offset, metalen), bodylen) + frameend <= dataend || throw(ValidationError("footer block escapes the data section")) + return offset, frameend +end + +function _validateblockindex(dictblocks, recordblocks, dataend::Int64; datastart::Int64=0) + extents = Tuple{Int64,Int64}[] + indexedend = datastart + for block in Iterators.flatten((dictblocks, recordblocks)) + extent = _blockextent(block, dataend) + extent[1] >= datastart || + throw(ValidationError("footer block overlaps the file schema")) + push!(extents, extent) + indexedend = max(indexedend, extent[2]) + end + sort!(extents; by=first) + for i = 2:length(extents) + extents[i - 1][2] <= extents[i][1] || + throw(ValidationError("footer blocks overlap")) + end + return indexedend +end + +# The `_block*` traversal below mirrors the schema-blind vtable primitives of +# src/metadata/VerifierRuntime.jl (`_vtable`/`_vfield`/`_vref`/`_vvector`/ +# `_vrange`) over a `BufferSlice` source with block-specific diagnostics; the +# two must change in lockstep. +function _blockrange(b::BufferSlice, pos::Int64, len::Int64, what::AbstractString) + (pos >= 0 && len >= 0 && len <= b.len && pos <= b.len - len) || + throw(ValidationError("$what escapes block metadata")) + return nothing +end + +function _blockload(b::BufferSlice, ::Type{T}, pos::Int64, what::AbstractString) where {T} + _blockrange(b, pos, Int64(sizeof(T)), what) + return AC.loadat(b, T, pos) +end + +function _blockadd(a::Int64, b::Int64, what::AbstractString) + try + return AC.checked_add(a, b) + catch e + e isa OverflowError || rethrow() + throw(ValidationError("$what overflows")) + end +end + +function _blocksub(a::Int64, b::Int64, what::AbstractString) + try + return AC.checked_sub(a, b) + catch e + e isa OverflowError || rethrow() + throw(ValidationError("$what overflows")) + end +end + +function _blockmul(a::Int64, b::Int64, what::AbstractString) + try + return AC.checked_mul(a, b) + catch e + e isa OverflowError || rethrow() + throw(ValidationError("$what overflows")) + end +end + +struct _BlockTable + metadata::BufferSlice + pos::Int64 + vpos::Int64 + vlen::Int64 + olen::Int64 +end + +function _blocktable(metadata::BufferSlice, pos::Int64, what::AbstractString) + pos % 4 == 0 || throw(ValidationError("$what is misaligned")) + back = Int64(_blockload(metadata, Int32, pos, what)) + back != 0 || throw(ValidationError("$what has a zero vtable offset")) + vpos = _blocksub(pos, back, what) + vpos % 2 == 0 || throw(ValidationError("$what vtable is misaligned")) + vlen = Int64(_blockload(metadata, UInt16, vpos, what)) + olen = Int64(_blockload(metadata, UInt16, _blockadd(vpos, Int64(2), what), what)) + vlen >= 4 && iseven(vlen) || throw(ValidationError("invalid $what vtable length $vlen")) + olen >= 4 || throw(ValidationError("invalid $what object length $olen")) + _blockrange(metadata, vpos, vlen, what) + _blockrange(metadata, pos, olen, what) + return _BlockTable(metadata, pos, vpos, vlen, olen) +end + +function _blockfield( + t::_BlockTable, + slot::Int, + width::Int, + what::AbstractString; + required::Bool=false, +) + entryoff = Int64(4 + 2slot) + if entryoff > t.vlen - 2 + required && throw(ValidationError("required $what is absent")) + return nothing + end + entry = _blockadd(t.vpos, entryoff, what) + off = Int64(_blockload(t.metadata, UInt16, entry, what)) + if off == 0 + required && throw(ValidationError("required $what is absent")) + return nothing + end + (off >= 4 && width <= t.olen && off <= t.olen - width) || + throw(ValidationError("$what exceeds its table object")) + pos = _blockadd(t.pos, off, what) + width > 1 && pos % min(width, 8) != 0 && throw(ValidationError("$what is misaligned")) + _blockrange(t.metadata, pos, Int64(width), what) + return pos +end + +function _blockref(t::_BlockTable, slot::Int, what::AbstractString; required::Bool=false) + pos = _blockfield(t, slot, 4, what; required=required) + pos === nothing && return nothing + rel = Int64(_blockload(t.metadata, UInt32, pos, what)) + rel > 0 || throw(ValidationError("$what has a null or backward offset")) + target = _blockadd(pos, rel, what) + _blockrange(t.metadata, target, Int64(1), what) + return target +end + +function _blockvector(t::_BlockTable, slot::Int, elemsize::Int, what::AbstractString) + pos = _blockref(t, slot, what) + pos === nothing && return nothing + pos % 4 == 0 || throw(ValidationError("$what length is misaligned")) + n = Int64(_blockload(t.metadata, UInt32, pos, what)) + start = _blockadd(pos, Int64(4), what) + bytes = _blockmul(n, Int64(elemsize), what) + _blockrange(t.metadata, start, bytes, what) + n > 0 && + elemsize > 1 && + start % min(elemsize, 8) != 0 && + throw(ValidationError("$what data is misaligned")) + return start, n +end + +""" +Read the fixed Message/RecordBatch envelope and wire-buffer structs that bind +one Footer Block to its on-wire frame. `_blockmessage` still verifies the +full metadata graph lazily. This preflight runs first so that EOS +classification never trusts a forged Footer extent. +""" +function _blockmessagebatch(metadata::BufferSlice) + root = Int64(_blockload(metadata, UInt32, Int64(0), "message root")) + root >= 4 || throw(ValidationError("invalid block message root offset")) + msg = _blocktable(metadata, root, "block message table") + bodypos = _blockfield(msg, 3, 8, "message body-length slot") + bodylen = + bodypos === nothing ? Int64(0) : + _blockload(metadata, Int64, bodypos, "message body-length slot") + headerpos = _blockfield(msg, 1, 1, "message header type"; required=true) + headertype = _blockload(metadata, UInt8, headerpos, "message header type") + headerref = _blockref(msg, 2, "message header"; required=true) + header = _blocktable(metadata, headerref, "message header table") + batch = if headertype == UInt8(2) # DictionaryBatch.data + dataref = _blockref(header, 1, "dictionary batch data"; required=true) + _blocktable(metadata, dataref, "dictionary record-batch table") + elseif headertype == UInt8(3) # RecordBatch + header + else + throw(ValidationError("footer block has unsupported message header type $headertype")) + end + return bodylen, headertype, batch +end + +function _verifyblockbuffers(batch::_BlockTable, bodylen::Int64) + buffers = _blockvector(batch, 2, 16, "record-batch buffer vector") + buffers === nothing && return nothing + start, n = buffers + last_nonempty_end = Int64(0) + for i = Int64(0):(n - 1) + base = _blockadd( + start, + _blockmul(i, Int64(16), "record-batch buffer position"), + "record-batch buffer position", + ) + offset = _blockload(batch.metadata, Int64, base, "record-batch buffer offset") + len = _blockload( + batch.metadata, + Int64, + _blockadd(base, Int64(8), "record-batch buffer length"), + "record-batch buffer length", + ) + offset >= 0 || throw(ValidationError("negative batch buffer offset $offset")) + len >= 0 || throw(ValidationError("negative batch buffer length $len")) + offset % 8 == 0 || + throw(ValidationError("batch buffer offset $offset is not 8-byte aligned")) + bufferend = _blockadd(offset, len, "batch buffer end") + bufferend <= bodylen || + throw(ValidationError("batch buffer [$offset, $len] escapes its message body")) + if len > 0 + offset >= last_nonempty_end || + throw(ValidationError("batch buffers overlap or move backwards")) + last_nonempty_end = bufferend + end + end + return nothing +end + +function _verifyblockframe( + blob::BufferSlice, + block::NTuple{3,Int64}, + expectedheadertype::UInt8, +) + offset, metalen, bodylen = block + AC.loadat(blob, UInt32, offset) == CONTINUATION || + throw(ValidationError("footer block does not point at a message")) + declared = Int64(AC.loadat(blob, Int32, offset + 4)) + declared == metalen - 8 || + throw(ValidationError("footer block metadata length does not match the message")) + metadata = AC.subslice(blob, offset + 8, declared) + messagebodylen, headertype, batch = _blockmessagebatch(metadata) + messagebodylen == bodylen || + throw(ValidationError("footer block body length does not match the message")) + headertype == expectedheadertype || + throw(ValidationError("footer block has the wrong message header type")) + _verifyblockbuffers(batch, bodylen) + return nothing +end + +function _verifyblockframes( + blob::BufferSlice, + dictblocks, + recordblocks, + dataend::Int64; + datastart::Int64=0, +) + indexedend = _validateblockindex(dictblocks, recordblocks, dataend; datastart=datastart) + foreach(block -> _verifyblockframe(blob, block, UInt8(2)), dictblocks) + foreach(block -> _verifyblockframe(blob, block, UInt8(3)), recordblocks) + return indexedend +end + +""" +Frame and verify the single message a Block points at, against the block's +own declared extents and the enclosing region. +""" +function _blockmessage( + region::OwnerRegion, + block::NTuple{3,Int64}, + dataend::Int64, + limits::Limits, + budget::AllocationBudget, +) + offset, _ = _blockextent(block, dataend) + _, metalen, bodylen = block + blob = BufferSlice(region, 0, region.len) + AC.loadat(blob, UInt32, offset) == CONTINUATION || + throw(ValidationError("footer block does not point at a message")) + declared = Int64(AC.loadat(blob, Int32, offset + 4)) + declared == metalen - 8 || + throw(ValidationError("footer block metadata length does not match the message")) + 0 < declared <= limits.max_metadata_bytes || throw( + ValidationError( + "metadata length $declared outside (0, $(limits.max_metadata_bytes)]", + ), + ) + 0 <= bodylen <= limits.max_body_bytes || + throw(ValidationError("body length $bodylen outside [0, $(limits.max_body_bytes)]")) + _chargevector!(budget, UInt8, declared, "metadata allocation") + metabytes = AC.slicebytes(AC.subslice(blob, offset + 8, declared)) + version, header_type, features, _ = + _verify_ipc_metadata_budgeted(metabytes, limits, budget) + msg = FB.getrootas(Meta.Message, metabytes, 0) + Int64(msg.bodyLength) == bodylen || + throw(ValidationError("footer block body length does not match the message")) + return FramedMessage( + msg, + AC.subslice(blob, offset + metalen, bodylen), + version, + header_type, + features, + ) +end + +""" + readfile(bytes::Vector{UInt8}; limits=Limits()) -> ArrowFile + readfile(region::OwnerRegion; limits=Limits()) -> ArrowFile + +Open an IPC-format file: verify both magics, the footer length, the footer +flatbuffer, and the schema; eagerly decode the dictionary blocks (shared by +every record batch); expose record batches lazily through the Block index. +Pass `mmapregion(path)` to read a file through Core's mmap path. Duplicate +dictionary ids and delta dictionaries are format errors here — the file +format carries exactly one dictionary batch per id. +""" +readfile(bytes::Vector{UInt8}; limits::Limits=Limits()) = + _readfile(heapregion(bytes), limits, AllocationBudget(limits.max_total_allocated_bytes)) + +readfile(region::OwnerRegion; limits::Limits=Limits()) = + _readfile(region, limits, AllocationBudget(limits.max_total_allocated_bytes)) + +function _readfile(region::OwnerRegion, limits::Limits, budget::AllocationBudget) + _requirelittleendian() + _validatelimits(limits) + blob = BufferSlice(region, 0, region.len) + # leading magic + pad (8), smallest footer flatbuffer (8), footer length + # (4), trailing magic (6). + minlen = Int64(8 + 8 + 4 + 6) + region.len >= minlen || throw(ValidationError("file is too short to be an IPC file")) + for (i, byte) in enumerate(FILE_MAGIC) + AC.loadat(blob, UInt8, Int64(i - 1)) == byte || + throw(ValidationError("missing leading ARROW1 magic")) + AC.loadat(blob, UInt8, region.len - 6 + (i - 1)) == byte || + throw(ValidationError("missing trailing ARROW1 magic")) + end + # The trailer is