diff --git a/.github/workflows/TagBot.yml b/.github/workflows/TagBot.yml index ad381d1f..6f39f1db 100644 --- a/.github/workflows/TagBot.yml +++ b/.github/workflows/TagBot.yml @@ -25,9 +25,16 @@ 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 diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0a1e9530..ad0144b3 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -39,7 +39,7 @@ jobs: release: name: Verify release - ${{ matrix.os }} runs-on: ${{ matrix.os }} - timeout-minutes: 30 + timeout-minutes: 45 strategy: fail-fast: false matrix: @@ -57,38 +57,52 @@ jobs: run: | sudo apt update sudo apt install -y -V subversion + # Without an installed Julia, verify_rc.sh asks api.github.com for the + # latest release without a token, which a rate-limited runner refuses. + - uses: julia-actions/setup-julia@v3 + if: matrix.os == 'ubuntu-latest' + with: + version: '1' - name: Create 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 key: ${{ runner.os }}-release-${{ hashFiles('**/Project.toml') }} restore-keys: | ${{ runner.os }}-release- + # The release candidate's suite, with one generated property case in four + # (the test job runs them all). - name: Verify + env: + ARROW_PROPERTY_CASES: 3/4 run: | version=$(grep -o '^version = ".*"' "Project.toml" | \ sed -e 's/^version = "//g' \ -e 's/"$//g') VERIFY_DEFAULT=0 dev/release/verify_rc.sh ${version} 100 test: - name: ${{ matrix.pkg.name }} - Julia ${{ matrix.version }} - ${{ matrix.runner || matrix.os }} - ${{ matrix.nthreads }} threads - ${{ github.event_name }} + # Coverage makes Julia compile the package, and each child process the tests + # start, without its precompiled code, so it runs on one job: ubuntu, Julia 1, + # 2 threads. That job also runs every generated property case; the others run + # one in four, chosen by their matrix position (ARROW_PROPERTY_CASES). + name: Arrow.jl - Julia ${{ matrix.version }} - ${{ matrix.runner || matrix.os }} - ${{ matrix.nthreads }} threads - ${{ github.event_name }} runs-on: ${{ matrix.runner || matrix.os }} timeout-minutes: 45 + permissions: + contents: read + id-token: write 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: @@ -98,7 +112,7 @@ jobs: - 1 - 2 include: - # Julia 1.9's libopenlibm requires executable-stack support rejected by newer glibc. + # Older Julia binaries need executable-stack support from this runner. - version: 'min' os: ubuntu-latest runner: ubuntu-24.04 @@ -107,6 +121,7 @@ jobs: - uses: julia-actions/setup-julia@v3 with: version: ${{ matrix.version }} + project: '.' - uses: actions/cache@v6 env: cache-name: cache-artifacts @@ -117,24 +132,75 @@ 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' + # Julia 1.10 needs the local ArrowTypes source developed explicitly. + - name: Dev local subpackages shell: julia --project=. {0} run: | using Pkg - Pkg.develop(PackageSpec(path="src/ArrowTypes")) + Pkg.develop([PackageSpec(path="src/ArrowTypes")]) + - uses: julia-actions/julia-buildpkg@v1.6 - uses: julia-actions/julia-runtest@v1 + with: + coverage: ${{ matrix.os == 'ubuntu-latest' && matrix.version == '1' && matrix.nthreads == 2 }} env: JULIA_NUM_THREADS: ${{ matrix.nthreads }} + ARROW_PROPERTY_CASES: ${{ (matrix.os == 'ubuntu-latest' && matrix.version == '1' && matrix.nthreads == 2) && 'all' || format('{0}/4', strategy.job-index) }} + - uses: julia-actions/julia-processcoverage@v1 + if: matrix.os == 'ubuntu-latest' && matrix.version == '1' && matrix.nthreads == 2 + with: + directories: src,ext + - uses: codecov/codecov-action@v7 + if: matrix.os == 'ubuntu-latest' && matrix.version == '1' && matrix.nthreads == 2 + with: + files: lcov.info + fail_ci_if_error: true + use_oidc: ${{ github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository }} + # 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-24.04 + timeout-minutes: 20 + permissions: + contents: read + id-token: write + strategy: + fail-fast: false + matrix: + pkg: + - name: ArrowTypes.jl + dir: './src/ArrowTypes' + 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 + fail_ci_if_error: true + use_oidc: ${{ github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository }} test_monorepo: name: Monorepo dev - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }} runs-on: ${{ matrix.os }} @@ -171,19 +237,105 @@ jobs: Pkg.develop([PackageSpec(path="."), PackageSpec(path="src/ArrowTypes")]) - name: Run monorepo tests continue-on-error: false + env: + ARROW_PROPERTY_CASES: 2/4 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.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' + 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/ArrowTypes")]) + - uses: julia-actions/julia-buildpkg@v1.6 + with: + project: ${{ matrix.pkg.dir }} + # Nothing uploads this job's coverage, so it runs without it. + - uses: julia-actions/julia-runtest@v1 + env: + JULIA_NUM_THREADS: 2 + ARROW_PROPERTY_CASES: 1/4 + with: + project: ${{ matrix.pkg.dir }} + coverage: false + trim: + name: ArrowCore and DataStrings 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 DataStrings workload + run: julia --startup-file=no --history-file=no test/trim_compile_tests.jl docs: name: Documentation runs-on: ubuntu-latest - timeout-minutes: 5 + timeout-minutes: 15 steps: - uses: actions/checkout@v7 - uses: julia-actions/setup-julia@v3 with: version: '1' - - uses: julia-actions/julia-buildpkg@latest - - uses: julia-actions/julia-docdeploy@latest + - name: Prepare Arrow.jl development dependencies for the docs + shell: julia --project=docs {0} + run: | + using Pkg + Pkg.develop([PackageSpec(path="."), 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 }} @@ -194,11 +346,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..d4fd3375 100644 --- a/.github/workflows/ci_nightly.yml +++ b/.github/workflows/ci_nightly.yml @@ -25,6 +25,9 @@ jobs: name: ${{ matrix.pkg.name }} - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ matrix.nthreads }} threads - ${{ github.event_name }} runs-on: ${{ matrix.os }} timeout-minutes: 45 + permissions: + contents: read + id-token: write strategy: fail-fast: false matrix: @@ -64,10 +67,21 @@ 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 + fail_ci_if_error: true + use_oidc: true test_monorepo: name: Monorepo dev - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }} runs-on: ${{ matrix.os }} 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..d9aaba8b --- /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/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..29de11c0 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,134 @@ + + +# Changelog + +## 3.0.0 - Unreleased + +Arrow.jl 3.0 is a new implementation of the reader, the writer, and the C data +interface. The main names stay: `Arrow.Table`, `Arrow.Stream`, `Arrow.write`, +`Arrow.Writer`, `Arrow.append`, `Arrow.tobuffer`, `Arrow.DictEncode`, and +`Arrow.getmetadata`. The [migration guide](docs/src/migration.md) shows how to update +code, and what Arrow 2.x can read of files Arrow 3.0 writes. + +### Breaking + +- Requires Julia 1.10, ArrowTypes.jl 2, Tables.jl 1.14, DataStrings.jl 1, + DataDecimals.jl 1, and Durations.jl 1.4. TimeZones.jl and BitIntegers.jl are no + longer dependencies. +- Value types: strings read as `DataStrings.DataString` and binary as `DataBytes` + (`stringtype=String` reads `String`s); timestamps as `Durations.Timestamp{P}` at the + column's unit, no longer truncated to `DateTime`; zoned timestamps as + `Durations.ZonedTimestamp{P,Z}` instead of `ZonedDateTime`; decimals as + `DataDecimals.Decimal`; intervals as `Durations.Duration`. +- Nested values: list rows are `Arrow.ListValue` views; map rows are `Arrow.MapValue`s, + an `AbstractDict` in stored order, instead of `Dict`s (copies, as in `DataFrame(table)`, + hold `Dict`s); fixed-size lists read as tuples only when they have at most 64 values + and no nulls. +- Columns are read-only; assigning into one is an error. Copy with `copy(table)`, + `DataFrame(table)`, or `Arrow.read(source, DataFrame)`. +- Removed: the `convert` keyword, `Arrow.Table(bytes, pos, len)`, reading a vector of + sources as one table or stream, `Arrow.ToTimestamp`, and the `Arrow.ToArrow` + binding. Compressor objects are no longer accepted by `compress`. +- Unknown keywords are an `ArgumentError`. The 2.x writer keywords `alignment`, + `dictencodenested`, `denseunions`, and `maxdepth` are ignored with a warning. +- `Arrow.Writer` takes its schema from the first table, with every column nullable + unless `nullable=false`; later tables, and tables given to `Arrow.append`, must + convert to the schema exactly. +- `dictencode=false` writes `PooledArray` and `CategoricalArray` columns as plain values. +- Writing a string that is not valid UTF-8 is an error. +- New column encodings that Arrow 2.x reads differently or not at all: `UUID` as the + canonical `arrow.uuid` extension; `Int128` and `UInt128` as 16-byte fixed-size binary + labelled `JuliaLang.Int128`/`UInt128`; `Minute`, `Hour`, `Day`, and `Week` as + durations in seconds and `Month` and `Year` as year-month intervals, labelled; + `VersionNumber` as labelled strings; `Durations.Duration` as month-day-nano intervals. +- A `Symbol` column interns at most `Limits.max_new_symbols` new symbols per read. +- Columns that share a name read as `a`, `a_1`, `a_2`, as in CSV.jl. +- Rejected input: big-endian IPC, metadata versions before V4, bytes after a stream's + end marker in `Arrow.Table` (read several streams with one `Arrow.Stream` each), and + a time of day outside [0, 24 h), which 2.x wrapped. + +### Added + +- `select`, `drop`, `filter`, `limit`, `offset`, and `scan` (a `Tables.Scan`) on + `Arrow.Table`, `Arrow.Stream`, and `Arrow.read`, with the filter functions `Arrow.col`, + `Arrow.colcmp`, `Arrow.colin`, and `Arrow.isnull`. +- `Arrow.read(source, sink)` reads into a sink with owned columns. +- `==`, `isequal`, and `hash` on `Arrow.Table`, comparing names and columns, and + `DataAPI.nrow` and `DataAPI.ncol`. +- `Arrow.AbstractArrowSource` (`Arrow.sourcelength`, `Arrow.readrange`, + `Arrow.concurrentreads`) reads only the byte ranges a read needs, and a CloudStore.jl + extension reads S3 and Azure objects that way. +- `Arrow.Stream` reads a stream-format `IO` one record batch at a time as it arrives, + and gives file-format batches by index (`stream[i]`), with `length`, `DataAPI.nrow`, + `Tables.schema`, and DataAPI metadata without reading data. +- `ntasks` on reads. Reads and writes give the same values and bytes for every `ntasks`. +- `Arrow.release!`, `copy(table)`, `Arrow.Limits`, `Arrow.ValidationError`, and + `Arrow.AllocationLimitError`. +- Writer keywords `compresslevel`, `batchsize`, `dictdeltas` (delta dictionaries), + `stringview` (string and binary views), and, on `Arrow.Writer`, `nullable` and + `dictreplacement` (replacement dictionaries); `dictencode` takes column names. +- Reading list views, run-end encoding, month-day-nano intervals, and the canonical + `arrow.uuid`, `arrow.bool8`, and `arrow.fixed_shape_tensor` extensions. Rewriting a + table keeps these types. +- The C data and C stream interfaces: `Arrow.to_c_data(table)` and + `Arrow.export_stream!(streamptr, table)` export any Tables.jl table, built as + `Arrow.write` builds it: plain integer and float `Vector`s and an `Arrow.Table`'s + integer, float, and boolean buffers cross without copying, and other columns are built + into new buffers. `Arrow.Table(schemaptr, arrayptr)`, + `Arrow.Table(streamptr)`, and `Arrow.Stream(streamptr)` import a record batch or a + stream with the column types an IPC read gives, reading fixed-width values in the + producer's memory. Memory handed to a consumer is freed once the consumer releases + it, with nothing to call. +- Rows of a dictionary column, `t.c[rows]` for any vector of integer indices or a `Bool` + mask, are a `PooledArray` with the column's categories, as `copy(t.c)` is. +- A TimeZones.jl extension that writes `ZonedDateTime` columns, and a CategoricalArrays + extension that writes ordered `CategoricalArray`s as ordered dictionaries and reports + `isordered` for ordered dictionary columns. +- A FileIO.jl extension: `load("x.arrow")` reads an `Arrow.Table` and + `save("x.arrow", table)` writes one. +- A DataFrames.jl extension that precompiles reading into and writing from a + `DataFrame`, so the first one in a session takes about 0.1 s rather than 0.3 s. + +### Changed + +- On Windows, a path is read into memory unless `mmap=true`, since Windows cannot + delete a mapped file. +- String columns are `StringVector`s of 16-byte values over the source's own bytes + (the mapped file, the byte vector, or the decompressed buffer), built once when the + column is first taken; 2.x allocated a new `String` on every access. `copy(table)` + and `Arrow.read` give string columns that own their bytes. +- `Arrow.write` writes partition by partition, so a partitioned source such as + `CSV.Chunks` writes in about one partition's memory, and validates each record batch + before writing it. +- Writing to an existing path writes a new file and renames it over the old one: a + failed write leaves the old file, and tables that map the old file keep working. +- Decimals write as decimal128, or decimal256 above 38 digits, whatever the storage of + the Julia value. +- `metadata` and `colmetadata` default to the table's DataAPI metadata. +- A table with several partitions writes every column as nullable. +- A buffer that compression would not shrink is written uncompressed. + +### Fixed + +- Polars' IPC files read; 2.x read them as an empty table. +- Decimal32 and Decimal64 columns read their values; 2.x read them as 128-bit values. +- Timestamps finer than a millisecond read exactly. +- Corrupt or hostile input fails with a `ValidationError` or `AllocationLimitError` + instead of crashing, hanging, or allocating without bound. 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..f71f2503 100644 --- a/Project.toml +++ b/Project.toml @@ -17,38 +17,70 @@ name = "Arrow" uuid = "69666777-d1a9-59fb-9406-91d4454c9d45" authors = ["quinnj "] -version = "2.8.1" +version = "3.0.0-DEV" [deps] +DataDecimals = "3e2245cb-6932-498d-a7cf-39d39de8bde3" +Durations = "41ac09f5-0a95-4810-aa18-c88bf48ee130" +DataStrings = "48204bd6-5611-42ea-b167-fc1d713f9be2" ArrowTypes = "31f734f8-188a-4ce0-8406-c8a06bd891cd" -BitIntegers = "c3b6d118-76ef-56ca-8cc7-ebb389d030a1" 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" +PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" SentinelArrays = "91c51154-3ec4-41a3-a24f-3f23e20d615c" -StringViews = "354b36f9-a18e-4713-926e-db85100087ba" Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" -TimeZones = "f269a46b-ccf7-5d73-abea-4c690281aa53" TranscodingStreams = "3bb67fe8-82b1-5028-8e26-92a6c54297fa" UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4" +# ArrowTypes is developed in this repository. DataStrings resolves from General. +[sources] +ArrowTypes = {path = "src/ArrowTypes"} + +[weakdeps] +DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" +FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549" +FilePathsBase = "48062228-2e41-5def-b9a4-89aafe57970f" +Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b" +CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" +CloudStore = "3365d9ee-d53b-4a56-812d-5344d5b716d7" +TimeZones = "f269a46b-ccf7-5d73-abea-4c690281aa53" + +[extensions] +ArrowDataFramesExt = "DataFrames" +ArrowFileIOExt = "FileIO" +ArrowFilePathsBaseExt = "FilePathsBase" +ArrowSerializationExt = "Serialization" +ArrowCategoricalArraysExt = "CategoricalArrays" +ArrowCloudStoreExt = "CloudStore" +ArrowTimeZonesExt = "TimeZones" + [compat] -ArrowTypes = "1.1,2" -BitIntegers = "0.2, 0.3" +DataFrames = "1" +FileIO = "1.16" +FilePathsBase = "0.9" +Serialization = "1.10" +DataDecimals = "1" +Durations = "1.4" +DataStrings = "1" +ArrowTypes = "2" +CategoricalArrays = "0.10, 1" +CloudStore = "1.6" CodecLz4 = "0.4" -CodecZstd = "0.7, 0.8" -ConcurrentUtilities = "2" -DataAPI = "1" +DataAPI = "1.12" +Dates = "1.10" +CodecZstd = "0.8" EnumX = "1" -PooledArrays = "0.5, 1.0" -SentinelArrays = "1" -StringViews = "1, 2" -Tables = "1.1" +Mmap = "1.10" +PooledArrays = "1.4" +PrecompileTools = "1" +SentinelArrays = "1.4" +Tables = "1.14" TimeZones = "1" TranscodingStreams = "0.9.12, 0.10, 0.11" -julia = "1.9" +UUIDs = "1.10" +julia = "1.10" diff --git a/README.md b/README.md index 98bc9fd9..7c8450f9 100644 --- a/README.md +++ b/README.md @@ -17,55 +17,80 @@ 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 reads and writes the [Apache Arrow](https://arrow.apache.org) IPC file and +stream formats in pure Julia. A file reads as Tables.jl columns that view its memory, and +any Tables.jl table writes. It also implements the Arrow C data and C stream interfaces. -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; 3.0 is not registered yet. Arrow 3.0 is a +> breaking release: read the [migration guide](docs/src/migration.md) before you upgrade +> from 2.x. A checkout uses the in-repository `src/ArrowTypes` package: Julia 1.11+ +> resolves it 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: - ```julia -julia> using Pkg; Pkg.add("Arrow") +import Pkg +Pkg.add("Arrow") ``` -## Local Development +## Quick start + +```julia +using Arrow, DataFrames + +Arrow.write("data.arrow", (id=[1, 2, 3], name=["Ada", "Ben", missing])) + +table = Arrow.Table("data.arrow") # columns view the file +table.name # ["Ada", "Ben", missing] +df = Arrow.read("data.arrow", DataFrame) # a DataFrame with its own columns -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: +subset = Arrow.Table("data.arrow"; select=[:id], filter=Arrow.col(:id) > 1) -```sh -julia --project -e 'using Pkg; Pkg.develop(path="src/ArrowTypes")' +for batch in Arrow.Stream("data.arrow") # one record batch at a time + # ... +end ``` -## Format Support +The [documentation](https://arrow.apache.org/julia/) covers reading, writing, streams, +the type mapping, limits and errors, the C data interface, and migrating from 2.x. The +[changelog](CHANGELOG.md) summarizes what changed in 3.0. -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 +## Development -It currently doesn't include support for: - * Tensors or sparse tensors - * Flight RPC - * C data interface +In a checkout of this branch, prepare the local subpackages, then run the +tests: -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 +```julia +import Pkg +Pkg.activate(".") +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/README.md b/bench/README.md new file mode 100644 index 00000000..357d5353 --- /dev/null +++ b/bench/README.md @@ -0,0 +1,70 @@ + + +# IPC throughput benchmarks + +Run the package and PyArrow against deterministic primitive, nullable, string, +list, and dictionary columns: + +```sh +julia --project=. -t 1 bench/run.jl /tmp/arrow-bench-1 +julia --project=. -t 8 bench/run.jl /tmp/arrow-bench-8 +``` + +The driver uses a host `python3` with PyArrow installed. Set `ARROW_BENCH_PYTHON` +to a different Python executable if needed. If host PyArrow is unavailable, it +uses the image built by `julia conformance/run.jl`, or reports that the Python leg +was skipped. A Docker fallback is useful for checking the harness; its timings +include different filesystem and CPU scheduling costs from the host Julia run. + +Both readers use the **same Julia-written file**, memory mapped, and request all +columns. Arrow.jl validates buffers and constructs its DataStrings columns on +first access. Numeric and nested columns retain views. PyArrow's `read_all` +returns its native Arrow arrays, usually without further construction. These +measure the cost of making columns available, not iterating every value or +copying them into mutable Julia/Python objects. The driver prints this distinction +with every report. Both legs check the complete logical values outside the timing. + +Writes start from equivalent in-memory data, use batches of at most 1,048,576 +rows, and write uncompressed IPC files. Dictionary construction is timed in both +implementations. Each repetition writes a fresh destination; deletion of the +previous output happens before the timer. This does not measure the additional +cost of Arrow.jl's atomic replacement of an existing file. + +Each operation warms up once, then reports the median of three repetitions. +The Julia process and PyArrow CPU pool use the driver's thread count. Filesystem +caches are warm. File size divided by elapsed time gives the displayed MB/s; +writer output sizes can differ because physical layout choices differ. + +Use a quiet machine, alternate thread counts, and repeat runs before drawing +conclusions from small timing differences. This harness does not measure peak +memory, first-use compilation, compressed IO, or predicate pushdown. It is a +small throughput check, not evidence that every workload is faster than another +implementation. + +For a quick correctness check of the harness, scale every workload down: + +```sh +ARROW_BENCH_SCALE=0.001 julia --project=. bench/run.jl /tmp/arrow-bench-smoke +``` + +`ARROW_BENCH_SCALE` must be positive and finite. Keep it at the same value for +all compared runs; the default is 1 (10M primitive rows, 2M strings, 1M lists, +2M dictionary rows). The work directory retains both implementations' files and +JSONL results for inspection. diff --git a/bench/bench_pyarrow.py b/bench/bench_pyarrow.py new file mode 100644 index 00000000..08390dfc --- /dev/null +++ b/bench/bench_pyarrow.py @@ -0,0 +1,135 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# PyArrow serialize/deserialize timing over the same logical workloads. +# Both readers use the Julia-written files. read_all returns Arrow arrays; +# neither leg times scalar iteration or conversion to Python objects. +# Usage: bench/run.jl chooses a local Python or the conformance image. + +import sys, time, os, math +import pyarrow as pa +import pyarrow.ipc as ipc + +SCALE = float(os.environ.get("ARROW_BENCH_SCALE", "1")) +if not math.isfinite(SCALE) or SCALE <= 0: + raise ValueError("ARROW_BENCH_SCALE must be positive and finite") +ROWS_PRIMITIVE = max(1, round(10_000_000 * SCALE)) +ROWS_STRINGS = max(1, round(2_000_000 * SCALE)) +ROWS_LISTS = max(1, round(1_000_000 * SCALE)) +ROWS_DICT = max(1, round(2_000_000 * SCALE)) +BATCHSIZE = 1_048_576 +THREADS = int(sys.argv[2]) if len(sys.argv) > 2 else 1 +pa.set_cpu_count(THREADS) + + +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 + # both 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, setup=lambda: None): + setup() + f() + ts = [] + for _ in range(runs): + setup() + 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, max_chunksize=BATCHSIZE) + + def fresh(): + if os.path.exists(path): + os.remove(path) + + twrite = bench(write, setup=fresh) + size = os.path.getsize(path) + inputpath = os.path.join(outdir, "rewrite-%s.arrow" % name) + + def read(): + with pa.memory_map(inputpath) as src: + return ipc.open_file( + src, options=ipc.IpcReadOptions(use_threads=THREADS > 1) + ).read_all() + + tread = bench(read) + # Cast dictionary columns back to their plain-string input domain. + if not read().cast(tbl.schema).equals(tbl): + raise AssertionError("round-trip mismatch in %s" % name) + for op, secs in (("write", twrite), ("read", tread)): + nbytes = size if op == "write" else os.path.getsize(inputpath) + print('{"impl":"pyarrow","workload":"%s","op":"%s",' + '"seconds":%r,"bytes":%d}' % (name, op, secs, nbytes)) + + +main(sys.argv[1]) diff --git a/bench/bench_rewrite.jl b/bench/bench_rewrite.jl new file mode 100644 index 00000000..3d747d02 --- /dev/null +++ b/bench/bench_rewrite.jl @@ -0,0 +1,60 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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") + # Fresh destinations keep atomic replacement of existing paths out of + # the measurement. Removing the preceding output is untimed. + twrite = bench_time( + () -> Arrow.write(path, tbl; batchsize=BENCH_BATCHSIZE); + setup=() -> rm(path; force=true), + ) + sz = filesize(path) + # Request every column: string payload construction is deferred until + # first access. Numeric and nested columns still contain views. + tread = bench_time() do + t = Arrow.Table(path) + for nm in Tables.columnnames(t) + length(Tables.getcolumn(t, nm)) + end + Arrow.release!(t) + end + # Check every value outside the timing, including dictionary decoding. + t = Arrow.Table(path) + for nm in Tables.columnnames(tbl) + isequal(Tables.getcolumn(t, nm), Tables.getcolumn(tbl, nm)) || + error("round-trip mismatch in $name.$nm") + end + Arrow.release!(t) + 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..23056109 --- /dev/null +++ b/bench/run.jl @@ -0,0 +1,158 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 separate processes on the host when PyArrow is installed. +# The conformance image is a fallback. Results print seconds and throughput. +# +# Both readers open the Julia-written files and return native columns. +# Arrow.jl also validates buffers and builds DataStrings payloads; PyArrow +# generally wraps buffers. Neither times scalar iteration or mutable copies. + +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}[] + threads = Threads.nthreads() + scale = get(ENV, "ARROW_BENCH_SCALE", "1") + println("Julia $(VERSION), $threads threads; ARROW_BENCH_SCALE=$scale") + + rewriteout = joinpath(workdir, "rewrite.jsonl") + rewritescript = joinpath(here, "bench_rewrite.jl") + _runleg( + `$(Base.julia_cmd()) --startup-file=no --threads=$threads --project=$repo $rewritescript $workdir`, + rewriteout, + ) + push!(legs, ("rewrite", rewriteout)) + + pyout = joinpath(workdir, "pyarrow.jsonl") + python = get(ENV, "ARROW_BENCH_PYTHON", something(Sys.which("python3"), "python3")) + havepython = try + success(pipeline(`$python -c "import pyarrow"`; stdout=devnull, stderr=devnull)) + catch + false + end + !havepython && + haskey(ENV, "ARROW_BENCH_PYTHON") && + error("ARROW_BENCH_PYTHON must name a Python executable with pyarrow installed") + # A container is useful for smoke checks, but its IO and CPU scheduling + # differ from the host Julia process. Prefer a local Python for timings. + havedocker = + !havepython && + Sys.which("docker") !== nothing && + try + success( + pipeline( + `docker image inspect arrow-julia-conformance:latest`; + stdout=devnull, + stderr=devnull, + ), + ) + catch + false + end + if havepython + version = readchomp(`$python -c "import pyarrow; print(pyarrow.__version__)"`) + println("PyArrow $version, $threads threads, host Python: $python") + _runleg(`$python $(joinpath(here, "bench_pyarrow.py")) $workdir $threads`, pyout) + push!(legs, ("pyarrow", pyout)) + elseif havedocker + image = "arrow-julia-conformance:latest" + python = "/opt/pyarrow/bin/python" + println( + "PyArrow runs in Docker; host/container timing ratios are not directly comparable.", + ) + _runleg( + `docker run --rm -e ARROW_BENCH_SCALE=$scale -v $workdir:/bench -v $here:/src $image $python /src/bench_pyarrow.py /bench $threads`, + pyout, + ) + push!(legs, ("pyarrow", pyout)) + else + println( + "(pyarrow leg skipped: neither host PyArrow nor the conformance Docker image is 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("Reads use the same memory-mapped files and request every column.") + println( + "Arrow.jl validates buffers and builds string payloads; PyArrow generally wraps buffers.", + ) + println( + "Neither read times scalar iteration or mutable copies. Writes use fresh destinations.", + ) + 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..04b7409f --- /dev/null +++ b/bench/workloads.jl @@ -0,0 +1,82 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Shared workload definitions for the serialize/deserialize benchmarks. +# Deterministic arithmetic data — every implementation +# builds the same logical tables. The driver reads the Julia-written files +# in both implementations; writers may choose different physical layouts. + +const BENCH_SCALE = parse(Float64, get(ENV, "ARROW_BENCH_SCALE", "1")) +isfinite(BENCH_SCALE) && BENCH_SCALE > 0 || + error("ARROW_BENCH_SCALE must be positive and finite") +const BENCH_ROWS_PRIMITIVE = max(1, round(Int, 10_000_000 * BENCH_SCALE)) +const BENCH_ROWS_STRINGS = max(1, round(Int, 2_000_000 * BENCH_SCALE)) +const BENCH_ROWS_LISTS = max(1, round(Int, 1_000_000 * BENCH_SCALE)) +const BENCH_ROWS_DICT = max(1, round(Int, 2_000_000 * BENCH_SCALE)) +const BENCH_BATCHSIZE = 1_048_576 + +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, setup=() -> nothing) where {F} + setup() + f() + times = Float64[] + for _ = 1:runs + setup() + 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..e8daef95 100644 --- a/codecov.yaml +++ b/codecov.yaml @@ -18,5 +18,5 @@ --- codecov: notify: - # Wait for all "test" matrix jobs - after_n_builds: 6 + # One CI run uploads 16 Arrow reports and one ArrowTypes report. + after_n_builds: 17 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..c94a7eea --- /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" +DataStrings = "48204bd6-5611-42ea-b167-fc1d713f9be2" +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] +DataStrings = "=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..95690b1c --- /dev/null +++ b/conformance/cdata_oracle.jl @@ -0,0 +1,481 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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!() + Arrow._reap!() + leaked = length(Arrow.EXPORTS) + 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..8249f25a --- /dev/null +++ b/conformance/oracle.jl @@ -0,0 +1,336 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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++. +# +# 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 TARGETED_ORACLE_VARIANTS = Tuple( + (family, codec) for family in COMPRESSED_ORACLE_FAMILIES for codec in (:lz4, :zstd) +) + +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 + 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 + +""" +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) + 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 + push!(verdicts, documentcheck(case.label, check, gold) 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..afbddbc2 --- /dev/null +++ b/conformance/run.jl @@ -0,0 +1,127 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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/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..cdcc01a4 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,8 @@ If ArrowTypes is also registered, we also need to post a comment on the commit w @JuliaRegistrator register subdir=src/ArrowTypes ``` +DataStrings is registered separately in [JuliaData/DataStrings.jl](https://github.com/JuliaData/DataStrings.jl). + ### 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..d339ba6c 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -17,7 +17,26 @@ [deps] Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45" +ArrowTypes = "31f734f8-188a-4ce0-8406-c8a06bd891cd" +CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" +DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a" +DataDecimals = "3e2245cb-6932-498d-a7cf-39d39de8bde3" +DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" +DataStrings = "48204bd6-5611-42ea-b167-fc1d713f9be2" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +Durations = "41ac09f5-0a95-4810-aa18-c88bf48ee130" +PooledArrays = "2dfb63ee-cc39-5dd5-95bd-886bf059d720" +Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" [compat] -Documenter = "1" +ArrowTypes = "2" +CategoricalArrays = "0.10, 1" +DataAPI = "1" +DataDecimals = "1" +DataFrames = "1" +DataStrings = "1" +Documenter = "1.17" +Durations = "1.4" +PooledArrays = "1" +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..b7a15911 --- /dev/null +++ b/docs/dev/DESIGN-incremental-writer.md @@ -0,0 +1,209 @@ + + +# Design: the Arrow 3.x incremental writer + +Status: implemented, and the implementation differs from this proposal: +`Arrow.write` also streams (one partition in memory at a time), the Writer +declares every field nullable by default and widens later tables losslessly, +and dictionaries grow by delta batches in both formats. The manual +(`docs/src/manual.md`) and the docstrings describe the current behavior. +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 emission (the reader accepts deltas from other producers) 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-nested-public-values.md b/docs/dev/DESIGN-nested-public-values.md new file mode 100644 index 00000000..93b1c625 --- /dev/null +++ b/docs/dev/DESIGN-nested-public-values.md @@ -0,0 +1,146 @@ + + +# Nested public values + +## The rule + +The public value domain is uniform at every depth. A facade-converted leaf +— Date32/Date64, Time, Timestamp (zone-naive and zone-declared), Duration, +supported Decimal, Interval — materializes as the same public Julia value +whether it is a top-level column, a dictionary pool, a run-end-encoded +values child, a union branch, a struct field, or a list/map element. The +same holds in reverse: the writer accepts those public values at every +depth, fresh and retained. + +This replaces the pre-3.0-release rule where only the top level and +dictionary encoding converted, and nested temporal leaves stayed raw +storage integers inside `Vector{Pair{String,Any}}` / `Vector{Any}` rows +(the `converted::Bool` threading through `_declaredeltype`, and +`_rawdeclaredbasetype`). That asymmetry — and the write-side refusal of +nested public values that mirrored it — is removed. + +## Layering + +* **ArrowCore stays the storage authority, untouched.** `AC._value`, + `AC.materialize`, and `AC.fromjulia` keep the storage domain: raw counts + for temporals, raw coefficients for decimals. They serve the C interface + and conformance. Core + keeps zero knowledge of Dates/Durations/DataDecimals public types and + remains trim-compileable on its own. + +* **One leaf-conversion authority in the facade.** A single mapping per + convertible leaf descriptor supplies the triple (public element type, + storage→public scalar, public→storage scalar). `_facadebasetype` (types), + `_publicleafvalue` (storage→public; the per-element body `_postconvert` + loops), and `_exactfacadescalar`/`_facadetoken` (public→storage) are that + authority. Every reader, writer, and scan-planner rule reads it; no + second copy of any conversion exists. + +* **The facade owns a recursive public row builder.** For composite and + wrapper roots the facade builds rows itself (`_publicvalue`), mirroring + Core's `_value` shapes — `Vector{Pair{String,Any}}` struct rows, + `Vector{Any}` list rows, `Pair{Any,Any}` map rows — but converting leaf + values through the authority AS the row is built. Union branches convert + at build time, while the type id is in hand (a built row cannot recover + its branch, so post-hoc conversion is impossible; this is why conversion + moved into row construction). Top-level scalar columns keep the typed + bulk path plus one whole-column conversion pass; run-end-encoded and + dictionary roots keep their flat expansion plus one whole-column pass + over the values domain (the expansion output is governed by one + descriptor, so the one-pass rewrite stays valid and fast). + +* **Declared types have no `converted` flag.** `_declaredeltype(f)` walks + the field tree and every leaf answers with its public type. Empty, + all-missing, and populated columns of one field agree by construction. + +* **The writer is symmetric.** Fresh composite columns accept public leaf + values (a `NamedTuple` field of `DateTime` or `Timestamp{P}` lowers + through the authority instead of refusing), and retained write-back + lowers nested public values through the same public-to-storage authority + the top level uses (`_retainedstorage`/`_facadetostorage` over + `_exactfacadevalue`). A value that cannot lower exactly refuses, + exactly as at the top level. + +* **ArrowTypes-labeled subtrees are a different domain on purpose.** An + extension-labeled field routes through the ArrowTypes contract + (`ArrowType(T)`-declared storage views); registered mappings written + against that interface keep working unchanged. The boundary is the + extension label, as before. + +* **The allocation budget covers the dynamic rows' overhead, not just their + payload.** The per-row builders run dynamically typed, so each value + leaves transient boxes beside the payload Core's budgeted `getvalue` + reserves. The facade charges a fixed measured-worst-case reserve at each + seam: per converted leaf (facade and raw-domain restoration), per + shared-decimal `_postconvert` element (the one type-unstable `_mapcol` + body — the stable bodies charge nothing extra), per routed-walker leaf + (Core's transient estimate object), per value lifted through a + registered `fromarrow` hook (fixed reserves split three ways: childless + scalar storage, width-charged container storage, and union or + transparent-wrapper storage), per struct-hook child (the `applicable` + probe and the variadic hook call re-splat and re-box every child + value) plus the row's inline payload size (reconstruction copies isbits + storage wholesale at each seam, so struct charges scale with + `sizeof` of the constructed row — heap children count as pointer + slots), per lifted string's byte length (the default ArrowTypes pointer + adapter copies it), and per nested restoration call + (`_chargerestoredelements!`, the dynamic keyword-call tuple, with an + extra share on the raw-domain struct tail). The other + half of the model is staying allocation-free where no reserve exists: + null slots return `missing` before the budgeted extraction (NullType + first — its layout has no validity buffer), every per-value cache + lookup (`_arrowtypestarget`, `_hasarrowtypesextension`, the eltype, + dictionary-field and storage-type memos) returns hits without + constructing the memoization thunk, the wide-decimal scalar conversion + runs unboxed behind `_shareddecimalscalar`'s type barrier, restoration + loops hand child Fields to dynamic calls through the per-Field + `_arrowtypesboxedchildren` cache instead of re-boxing them per row, the + struct and map branches read their storage base type from the target + and eltype caches instead of rebuilding it per row (an unknown label + keeps the direct computation: its cached slot holds the bounded + fallback shape), and the row-lifting loop boxes its `Field` once behind + `_boxonce`, which the optimizer cannot fold away. The reserves were + sized by measuring + flat/union/REE/struct/list roots × timestamp/decimal/date leaves × + unlabeled/unknown-label/registered-label shapes; a nestedvalues testset + pins per-value charge floors for each seam and (on current Julia) that a + warmed read allocates no more than it charges. + +## Trim contract + +Two tiers, by what the caller knows statically: + +* **Dynamic schema** (runtime-discovered): the recursive builder is + dynamic dispatch over `Any`-shaped rows and is NOT trim-guaranteed. +* **Static schema**: two surfaces. Core's typed `getvalue(::Type{T}, ...)` + and `materialize(::Type{T}, ...)` claims stay the storage-domain + monomorphic tier and remain part of the trim workload, unchanged. The + facade's typed select override (`:col => NamedTuple{...}`, recursively, + including `Vector{...}` targets) converts each value with `convert`, as + `Tables.scan` does. It is a convenience over the dynamic rows and is not + trim-audited; making the whole facade trim-compileable stays future work. + +The typed tiers share no code with the dynamic row shapes, so the dynamic +`Pair{String,Any}` containers cost them nothing. + +## Scans + +Scan filters evaluate on the public values a full read returns, composite +columns included: plain Julia `==` over public rows is the semantics. diff --git a/docs/dev/DESIGN-scan-ranges-trim.md b/docs/dev/DESIGN-scan-ranges-trim.md new file mode 100644 index 00000000..44f877b5 --- /dev/null +++ b/docs/dev/DESIGN-scan-ranges-trim.md @@ -0,0 +1,323 @@ + + +# 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=…)`, the `select`/`drop`/`filter`/`limit`/ +`offset` keywords, and `Arrow.Stream`); §4 lists what is and is not built. +The 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 3 separates that design +intent from what the trim harness actually compiles. + +--- + +## 1. Tables.Scan support + +`Tables.Scan` (Tables.jl 1.14) is a plain-data scan request: select, rename, +and type items, a closed predicate algebra (`Cmp`/`In`/`IsNull`/`StrPred`/ +`And`/`Or`/`Not`, with `OpNode` as the growth channel), and `limit`/`offset`. +Arrow resolves each request once against the source schema with +`Tables.resolve` (`_ScanPlan`) and executes the bound request itself. Key +contract points this design leans on: + +- a pushed result must equal the generic executor's over the same scan, row + for row (the differential battery pins this for every source kind); +- `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. + +### One planner, one executor + +Every read goes through one planner and one executor, whatever the source: +a memory-mapped file, a byte vector, an `IO` read whole, a byte-range source +such as a cloud object, or stream-format input. A whole read is the scan that +selects every column. + + plan(scan, fields) = resolve once against the schema names → + check literal kinds against their columns (`_checkliterals`) → + give zone-free time literals a UTC column's zone (`_zonedbound`) + + _scanfile(src, footer, plan) = + one pass over every Block's metadata (a region serves it in place; a + byte-range source fetches it in coalesced ranges that also cover the + start of each body) → + row counts → the batch window when there is no filter → + the decode set's dictionaries, in footer order → + a decode plan per batch in the window → _runscan + + _scanstream(stream) = + frame the messages, apply dictionaries in stream order, and plan each + record batch against the dictionaries in effect for it → _runscan + + _runscan(records, plan) = + no filter and no window: decode the decode set of every batch on up to + `ntasks` tasks and build the columns a full read builds; + otherwise, per batch: decode the filter columns, evaluate the filter on + their public values (`Tables.filtermask`) in chunks of rows, then decode + the other selected columns only if a row qualifies, and keep copies of + the qualifying rows; without `limit` or `offset` batches run on up to + `ntasks` tasks, with a window they run in order and stop once it is full + +`Arrow.Stream` runs the same executor over one batch at a time. + +A region (a mapped file or a byte vector) is its own span provider: one span +over the whole region serves every range the planner names, without a copy. +A byte-range source fetches each round's ranges through `readrange`, +coalesced, and a fetched range is referenced in place by the columns that +read it, unless a select's columns fill less than half of it: then their +bytes are copied out, so a few columns do not keep a whole fetched range +alive. The planner is the same for both; only the fetch differs. + +### What Arrow pushes, by axis + +| Axis | Mechanism | +|---|---| +| `select` | only the decode set (selected ∪ filter columns) is decoded; `skipfield!` advances the node and buffer cursor past every other field with each buffer-table invariant checked, without slicing, decompressing, or validating its body; over a byte-range source no body range is planned for it | +| `limit`/`offset` | without a filter, `RecordBatch.length` is wire metadata: batches outside the window are never decoded or fetched; with a filter, the window composes over the qualifying rows, in order, and the scan stops once it is full | +| `filter` | evaluated per batch over the public values of the filter columns, the values a full read returns; the other selected columns are decoded (and over a byte-range source, fetched) only for batches with a qualifying row | +| `types` (`ref => T`) | an elementwise `convert` after the scan, as the generic executor applies it. Two overrides read stored values for their own output instead: a time of day as a `Dates.Period`, which also reads times outside a day, and a `JuliaLang.Symbol` column as strings, which creates no `Symbol`. Another output of the same column, and a filter on it, read its public values | + +The filter is evaluated by the generic evaluator over the batch's public +values, so pushdown and the executor share one three-valued semantics by +construction (a `missing` predicate excludes the row), and every predicate a +user can write pushes down: no literal needs a stored-value form. + +Element types do not depend on the source kind or on a filter; container +types do. A read that keeps every row of the batches it reads builds the +columns a full read builds: views of the source, `ChainedVector`s across +batches, and string columns over the source's bytes. A read with a filter or +a window keeps copies of its rows: `Vector`s for flat and nested values, +string and binary columns that hold just their values' bytes, and dictionary +columns that keep their pool, so a filtered flat or string column does not +keep a batch's decompressed bytes or the mapping alive for a few of its rows. +Kept nested rows still read their batch's child buffers. An empty result +column has the type a read of no batches gives it. + +Properties the composition holds: + +- **Positional filter references resolve to source names** before + evaluation (`Tables.resolve` normalizes them). +- **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. The rows of all batches together must fit + Tables' `Int` row count, whatever the window keeps. +- **One scan has one allocation budget,** shared by footer work, metadata, + fetched ranges, decompression, and column building. +- **Windows saturate.** `offset`/`limit` compose with `min` arithmetic over + row counts, so `offset + limit` never overflows. +- **Zero-field sources** read their row counts from metadata; a filter that + reads no column evaluates once per batch, and no per-row mask is allocated + from an untrusted row count. +- **A scan checks only what it reads.** Every Block's metadata is verified, + but the body of a dictionary outside the decode set, or of a record batch + outside the window, is never read, so damage there does not fail the scan. + +--- + +## 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 and Block indexes. The leading magic is not fetched: the Footer + is the sole authority. +2. **Block metadata fetches**: the metadata of every dictionary and record + batch Block, each with the start of its body (all of a body up to four + `gapbytes`, else its first `gapbytes`), coalesced across nearby spans. A + scan with a row window reads metadata alone. RecordBatch row counts are + here, not in the Footer, so the window is chosen after this pass. +3. **Window and buffer-range plan**: row counts choose the batch window when + there is no filter; the decode set then maps to a buffer index set + (subtree-inclusive; dictionary Blocks for the decode set's dictionaries, + including those a dictionary's values use) → byte ranges → **coalesce** + ranges with gaps below `gapbytes` (default 256 KiB — a gap fetch is usually + cheaper than a request round-trip). +4. **Body fetches**: the decode set's dictionary bodies and the record bodies + the executor reads first (every decode-set column without a filter; the + filter columns with one) in one round. A filter then fetches the other + selected columns of the batches that keep rows in one more round; a + filtered `limit` fetches batch by batch (a body of at most four `gapbytes` + with all its decode-set columns in one round) and stops when it is full, + freeing each batch's bytes once it is read. Each coalesced range lands in its own + owned heap region; decode resolves each declared buffer `(offset, len)` to + its containing fetched range and subslices it: *every buffer must fall + inside a fetched range that was itself derived from the verified buffer + table*, the message-body authority invariant with sparse backing. + +Request-count model (what matters against cloud latency): `1` tail (plus one +exact Footer follow-up when the tail is too small) + `⌈metadata spans after +coalescing⌉` + `⌈coalesced body ranges⌉`. A read without a filter makes +three sequential rounds, a filtered read four. The source can issue a round's +ranges concurrently. For a 40-column file reading 3 columns of every batch, +this moves roughly `3/40` of the body bytes plus metadata. Coalescing is an +explicit over-read policy: a requested span may cross unneeded bytes when the +configured gap permits it. + +`Arrow.Stream` over a byte-range file source reads the footer and the +dictionaries when it opens. A batch that decodes every column, or whose body is +at most four `gapbytes`, reads its whole block, and the blocks after it while +they fit 8 MiB, with one window fetched ahead. A batch that decodes some +columns of a larger body reads its metadata with the start of its body, then +the ranges of its decode set it lacks. + +### 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. + +- `Arrow.Table(src; …)` builds the internal `SourceFile(src; limits, + tailbytes, gapbytes)` handle and reads through the planner; the source's + length is read once, at handle construction, and the tail once per handle. + `Arrow.Table(src)` without a scan, or over a stream-format object, reads + the object whole in 8 MiB ranges (one request beyond the cached tail for a + source that reads one range at a time). +- 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`, `test/ranged_scan_tests.jl`), and the CloudStore + extension end to end against a local S3-compatible server + (`test/cloudstore_tests.jl`). +- Explicitly out of scope, documented: caching/prefetch policy beyond + coalescing and `Arrow.Stream`'s read-ahead, 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`: +every dictionary and record block's metadata is checked. 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. 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 DataStrings workload. It does **not** yet compile +a scan-and-materialize app, so §1–§2 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. + +--- + +## 4. Status + +Implemented (`src/scan.jl`, `src/table.jl`, `src/source.jl`): + +- **One planner and one executor** for every source, as decided on Aug 19 + 2026: memory-mapped files, byte vectors, an `IO` read whole, byte-range + sources, and stream-format input. A region serves the planner's ranges in + place, a byte-range source fetches them in coalesced rounds, and + `Arrow.Stream` runs the executor over one batch at a time. `skipfield!` + skips what the scan does not decode; `_ScanPlan` resolves a request once; + `_runscan` reads filter columns first, evaluates filters on public values in + chunks of rows on up to `ntasks` tasks, and decodes the other selected + columns only for batches that keep rows. +- **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). + +Not implemented: an HTTP transport extension (the source contract is the +extension point; the CloudStore extension is the model) and the +scan-and-materialize trim harness. `Arrow.Stream` decodes every dictionary +batch of a file when it opens, including those of columns it does not +select. + +Possible future work: per-batch statistics (row counts, null counts, minimum +and maximum values) in IPC footers, to prune batches a filter cannot match +before their metadata is fetched. The Arrow statistics schema standardizes +the value layout but not a placement in IPC files; without a placement other +implementations read, a private key would help only Julia readers of files +Julia wrote, and wrong statistics silently drop rows. diff --git a/docs/dev/core-README.md b/docs/dev/core-README.md new file mode 100644 index 00000000..5212f674 --- /dev/null +++ b/docs/dev/core-README.md @@ -0,0 +1,234 @@ + + +# The Arrow.jl engine + +Arrow.jl separates the Arrow memory layout from its transports and Julia table +interface. `ArrowCore` owns descriptors, buffers, and validation. IPC and C data +adapters consume those values. The table facade presents Julia columns and +Tables.jl operations; the writer converts those columns back into Arrow buffers. +`ArrowCore` is an internal module, not a separately versioned package. + +The user guides describe the supported [types](../src/types.md), +[read operations](../src/reading.md), [writes](../src/writing.md), +[streams](../src/streams.md), and [limits](../src/limits.md). This document explains +where those contracts are implemented and the constraints on changing them. + +## Code map + +| File | Responsibility | +|---|---| +| `src/ArrowCore.jl` | Ownership regions, checked buffer slices, type descriptors, fields, schemas, array layouts, validation, and low-level builders/accessors | +| `src/FlatBuffers/`, `src/metadata/` | Vendored FlatBuffers runtime and generated metadata bindings/verifier; `tools/fbsgen.jl` regenerates the latter from the vendored schemas | +| `src/ipc_read.jl` | Framing, limits, metadata decoding, layout decoding, decompression, the stream reader, and the file reader | +| `src/ipc_dictionary.jl` | Dictionary snapshots, replacement and delta updates, shared IDs, and dictionary dependencies | +| `src/ipc_write.jl` | IPC encoding, batch emission, file footers, compression, and incremental writer state | +| `src/cdata.jl`, `src/cstream.jl` | C data and C stream import/export of Core arrays, release callbacks, and foreign ownership | +| `src/cdata_table.jl` | `Arrow.Table`, `Arrow.Stream`, and Tables.jl tables across the C data and C stream interfaces | +| `src/source.jl` | Byte-range sources and read scheduling; `ext/ArrowCloudStoreExt.jl` supplies the CloudStore integration | +| `src/columns.jl` | Read-only buffer views, typed nested rows, dictionary/run-end columns, and deferred string construction | +| `src/table.jl` | `Table`/`Stream`, source opening, Tables.jl and DataAPI interfaces, copying, and public value conversion | +| `src/scan.jl` | Scan binding, the planner and executor every read uses, filters, row windows, and sparse ranged reads | +| `src/arrowtypes.jl`, `src/sharedvalues.jl` | ArrowTypes mappings and the shared string, decimal, and temporal value domains | +| `src/columnconstruction.jl` | Native and custom column construction, retained schemas, nested layouts, and dictionary pooling | +| `src/write.jl` | Public write options, partition iteration, `Writer`, append, and destination ownership | +| `src/precompile.jl` | Representative public reads and writes used to prepare the package image | +| `src/ArrowTypes/` | The separately registered ArrowTypes interface package; it is developed with Arrow in this repository | + +The only exported compatibility binding is `ArrowTypes`. Public entry points are +listed in `src/Arrow.jl`; implementation helpers are not part of that API. + +## Buffers and lifetime + +An `ArrayData` describes a column through checked `BufferSlice`s, child arrays, +an optional dictionary, and a logical offset and length. Each slice retains an +`OwnerRegion`: a pointer, extent, alignment, GC root, and shared release cell. +The layout registry describes buffer roles and child shapes. IPC dictionary IDs +stay in the adapter; Core dictionaries describe their values and index layout. + +Low-level access checks bounds and the region's release cell. Releasing a foreign +owner revokes its regions and invokes the producer's callback exactly once. +An export keeps its roots until the consumer has released every C node of it; +the release callbacks only count the nodes off, and later exports or garbage +collections free the export. Re-exporting imported data roots the foreign owner +the same way. + +The public table facade has a different lifetime contract. Its numeric and nested +columns keep their backing storage alive and read without a per-element release +check. `release!(table)` prevents later column retrieval; `release!(stream)` also +prevents iteration. Columns already handed out remain valid. Facade memory maps +are therefore unmapped by GC after the table and its columns become unreachable, +not eagerly by `release!`. A low-level region can use eager release where its +caller controls every access. Keep these two contracts separate. + +Input byte vectors and borrowed Julia vectors must not be resized or changed +while a read or low-level view uses them. External truncation of a mapped file +can cause a process-level bus error; the facade cannot guard against it. The +reader's `mmap=false` option avoids that dependency. Replacing an existing string +path through `Arrow.write` uses a temporary file and rename, so existing mappings +continue to see the old file. + +C pointers carry declared extents, not verifiable allocation sizes. The C data +producer must keep those allocations alive and unchanged until release. Exported +callbacks run Julia code: call them from Julia-attached threads, and serialize +calls on one C stream. The interfaces do not supply a foreign-thread trampoline. + +## Reading and validation + +IPC framing and the generated FlatBuffers verifier check sizes and metadata +geometry before metadata-directed allocations. Every batch buffer belongs to a +checked slice of that message's body. Compressed buffers decode into bounded, +exact-sized outputs; declared lengths and codec consumption are checked. Both +LZ4_FRAME and ZSTD are supported, including the stored-raw buffer sentinel. + +Structural validation checks layout geometry. Semantic validation checks the +conditions needed to read those layouts safely. `validate_full` adds stricter +content checks such as UTF-8, decimal precision, time-of-day range, advisory +nullability, and canonical bitmap padding. These are distinct contracts: do not +turn an optional content check into an unconditional reader refusal without +checking interoperability and the public type conversion that follows it. + +File reads use the footer's schema and batch index, including files whose data +section has no leading schema message. Stream reads consume messages in order. +Dictionary replacements and deltas publish new pools; earlier batches retain +their own snapshots. File indexes and ranged reads must preserve the same +values when dictionary updates and record batches occur in different orders. + +`Arrow.Table` exposes read-only views for primitive, Boolean, dictionary, +run-end, list, map, and struct layouts. Multi-batch columns use +`SentinelArrays.ChainedVector`. Native scalar conversions preserve the shared +public domains at every depth: Durations timestamps, DataDecimals decimals, +and the supported date/time/interval values. Custom ArrowTypes mappings may +require materialization. + +String and binary columns use DataStrings. The reader reserves their construction +budget up front and builds their payloads on first column access. Owned +decompression buffers can be retained; other input bytes are copied into +DataStrings buffers. Opening a table alone does not pay all of that work. +`stringtype=String` instead constructs Julia strings, including nested +strings. A benchmark that claims to read every column must retrieve every column +inside its timing. Neither retrieval nor `length(column)` visits every value in +a numeric or nested view. + +`Arrow.Limits` bounds metadata, array sizes, concurrent reads, new Symbol values, +and allocated output. The allocation allowance depends on input bytes, with a +minimum of 256 MiB and the caller's `max_total_allocated_bytes` ceiling. Compressed +input receives allowance for its supported decompressed data. Each stream batch +has its own budget, so a long stream is not charged against one cumulative cap. +The accounting bounds package-controlled work; it does not measure arbitrary +user-defined conversion hooks or the entire process heap. + +## Scans and ranged sources + +`Tables.Scan` is available from registered Tables.jl 1.14 or later; no source +revision is required. The facade accepts it directly and translates its simpler +read keywords into the same plan. Binding and column selection happen once per +request; filters run on the public values a read returns. + +Every read goes through one planner and one executor (`scan.jl`), whatever the +source: a region in memory (a mapped file, a byte vector, an `IO` read whole), a +byte-range source, or stream-format input. A scan decodes the selected columns +and those its filter reads. Without a filter, offsets and limits select batches +by their row counts. With a filter, the filter columns are read first, the other +selected columns only for batches with a qualifying row, and qualifying rows are +counted before applying the result window. Kept rows are copied; a read that +keeps every row builds a full read's views. Type overrides convert after the +scan. Repeated top-level names receive unique suffixes before selection and +metadata lookup. + +`SourceFile` uses a byte-range source's footer and index to request needed metadata +and buffers. Adjacent ranges may be coalesced under `gapbytes`; concurrency is +bounded by the source and reader limits. `Arrow.Stream` reads one batch at a time +through the same executor, with read-ahead for whole batches. Do not infer a bound +on bytes fetched from a bound on rows returned: filtering can still require +reading buffers for rows that are later rejected. + +Schema and field metadata remain ordered pairs so duplicate metadata keys and +their ordering can survive a retained-schema rewrite. + +## Writing and concurrency + +`write` and `Writer` share the incremental IPC machinery. They consume one +partition at a time and divide it into record batches. `batchsize` caps rows; +overflow of a 32-bit string/list offset splits a batch further. The source +partition itself remains live while its batches are written. Reused construction +buffers and dictionary pools also contribute to peak memory, so a row-count cap +is not a byte quota on arbitrary input iterators. + +`ntasks` bounds work on independent columns and compression buffers. Publication +and dictionary state remain ordered. Output must be independent of the task +count, and a single caller owns a `Writer`; overlapping writes are unsupported. +Readers similarly parallelize suitable column construction while preserving +batch order and ownership. + +The first input establishes an incremental writer's schema. Later inputs use +compatible conversions and the declared nullability policy. Dictionary pools can +grow: file output normally emits the complete dictionary at close, while stream +output emits an updated dictionary before the next dependent batch. Delta and +replacement modes have explicit options and consumer compatibility constraints +in the [stream guide](../src/streams.md). `append` resumes an existing IPC stream's +schema and final dictionary state; it does not append to the random-access file +format. + +Fresh Julia columns and retained-schema rewrites use the same recursive column +constructor. Retained descriptors remain authoritative where values and views +preserve enough information to reproduce them. Unsupported reconstruction must +fail instead of silently relabeling values. `dictencode=false` deliberately writes +pooled columns as plain values, while `DictEncode` remains an explicit encoding +request. Automatic dictionaries use DataAPI ref pools; optional CategoricalArrays +methods preserve ordered categories. These hooks avoid making DataFrames or CSV +runtime dependencies. + +## Verification and compilation + +On Julia 1.10, develop ArrowTypes explicitly because that Julia version does not +read `[sources]` entries: + +```sh +julia --project=. -e 'using Pkg; Pkg.develop(path="src/ArrowTypes"); Pkg.instantiate()' +julia --project=. -e 'using Pkg; Pkg.test()' +julia --project=. test/fuzz.jl --cases 16 --mutations 64 +julia --startup-file=no test/trim_compile_tests.jl +julia --startup-file=no conformance/run.jl +julia --project=. bench/run.jl +julia tools/fbsgen.jl src/metadata/fbs src/metadata +``` + +The package suite includes Core, public interfaces, optional integrations, +frozen Arrow 2.x fixtures, corruption cases, and four adapter acceptance suites. +The deterministic fuzzer records seeds and replay artifacts. Conformance compares +IPC with the pinned Apache corpus and PyArrow/nanoarrow, and exchanges actual +C data/C stream pointers with PyArrow in one process. `conformance/run.jl` prepares +a Docker image and its Julia/Python environments; direct tool and corpus versions +are pinned in the conformance files, while transitive dependencies and the base +image are not fully locked. + +The trim gate exercises a specific low-level Core/C data workload and requires +zero verifier warnings or errors plus a runnable binary. It does not establish +that arbitrary table schemas, custom mappings, or the complete facade compile +under `--trim=safe`. Closed descriptor dispatch, concrete containers, and runtime +C-function pointers keep that workload compilable. Never store raw `@cfunction` +pointers in module-level constants: precompilation cannot preserve their runtime +addresses. + +The public precompile workload and invalidation regression test cover a separate +concern: loading other data packages should not force Arrow's common operations +to compile again. Abstract extension calls use bounded specialization and +inference barriers where required. Preserve these boundaries when adding public +methods, and measure first use separately from warmed throughput. diff --git a/docs/make.jl b/docs/make.jl index b26c44bc..0d3c31db 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -22,16 +22,27 @@ 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", - "API Reference" => "reference.md", + "Reading" => "reading.md", + "Writing" => "writing.md", + "Streams and IPC" => "streams.md", + "Type mapping" => "types.md", + "Limits and errors" => "limits.md", + "C data interface" => "cdata.md", + "Migrating from Arrow.jl 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/cdata.md b/docs/src/cdata.md new file mode 100644 index 00000000..bbf0faa9 --- /dev/null +++ b/docs/src/cdata.md @@ -0,0 +1,161 @@ +```@raw html + +``` + +# C data interface + +The Arrow [C data interface](https://arrow.apache.org/docs/format/CDataInterface.html) +and [C stream interface](https://arrow.apache.org/docs/format/CStreamInterface.html) +hand tables to another Arrow implementation in the same process, such as pyarrow, Polars, +or DuckDB through PythonCall, as pointers to Arrow buffers instead of IPC bytes. The +producer fills C structs (`ArrowSchema`, `ArrowArray`, `ArrowArrayStream`); the consumer +moves them out and calls their `release` callbacks when it is done. + +[`Arrow.export_stream!`](@ref) exports any Tables.jl table as a stream of record batches +and [`Arrow.to_c_data`](@ref) as one record batch; [`Arrow.Table`](@ref) and +[`Arrow.Stream`](@ref) import from pointers to the structs. An exported table is built +as [`Arrow.write`](@ref) builds it, so the other implementation sees what it would read +from the IPC bytes, and an imported table has the column types an IPC read gives. + +With PythonCall, a `DataFrame` goes to pyarrow and a pyarrow table comes back like this: + +```julia +using Arrow, PythonCall, DataFrames +pa = pyimport("pyarrow") + +# Julia -> pyarrow: any Tables.jl table, as a stream of record batches +df = DataFrame(id=1:3, name=["a", "b", missing]) +stream = Ref{Arrow.CArrowArrayStream}() +pytable = GC.@preserve stream begin + p = Base.unsafe_convert(Ptr{Arrow.CArrowArrayStream}, stream) + Arrow.export_stream!(p, df) + pa.RecordBatchReader._import_from_c(UInt(p)).read_all() +end + +# pyarrow -> Julia: pyarrow fills the struct and Arrow.Table moves it out +stream = Ref{Arrow.CArrowArrayStream}() +t = GC.@preserve stream begin + p = Base.unsafe_convert(Ptr{Arrow.CArrowArrayStream}, stream) + pytable.to_reader()._export_to_c(UInt(p)) + Arrow.Table(p) +end +t.name # the strings and the missing value, as Arrow.Table(bytes) reads them +``` + +`read_all()` reads every batch on the calling Julia thread. Polars and DuckDB take the +pyarrow table it returns: `polars.from_arrow(pytable)`, or DuckDB's +`connection.register("name", pytable)`. Give DuckDB the table, not the reader: DuckDB +reads a reader on its own threads, which can deadlock with Julia's garbage collector +(see [Threads](@ref cdata-threads)). The `Ref`, `GC.@preserve`, and `unsafe_convert` +lines are the C ABI: the caller owns the struct's memory while the other side fills or +empties it. + +One record batch crosses as two structs that Arrow.jl allocates, and comes back from +structs the caller allocates: + +```julia +schemaptr, arrayptr = Arrow.to_c_data(df) +batch = pa.RecordBatch._import_from_c(UInt(arrayptr), UInt(schemaptr)) + +aref = Ref{Arrow.CArrowArray}() +sref = Ref{Arrow.CArrowSchema}() +t = GC.@preserve aref sref begin + ap = Base.unsafe_convert(Ptr{Arrow.CArrowArray}, aref) + sp = Base.unsafe_convert(Ptr{Arrow.CArrowSchema}, sref) + batch._export_to_c(UInt(ap), UInt(sp)) + Arrow.Table(sp, ap) # a struct array is a table; any other array is one column +end +``` + +[`Arrow.Stream(streamptr)`](@ref Arrow.Stream) pulls one record batch per iteration +instead of reading the whole stream, for a DuckDB or Polars result larger than memory. + +## What crosses, and what is copied + +- Exports use [`Arrow.write`](@ref)'s column construction and keywords: `dictencode`, + `stringview`, `largelists`, `metadata`, and `colmetadata` mean what they mean there, + and the consumer sees the same schema and values as from the IPC bytes. A `Vector` of + `Int8` to `Int64`, `UInt8` to `UInt64`, or `Float16` to `Float64` crosses without + copying, and so do an [`Arrow.Table`](@ref) column's integer, float, and boolean + values and validity bitmaps, in struct fields and list values too. Everything else is + built into new Arrow buffers, as a write builds it: `Bool` and `missing` values in + Julia vectors (a `Vector{Union{Missing,T}}` is copied even when it holds no + `missing`), other numbers (`Int128`, `Complex`, `Rational`, and decimals, except a + `Vector` of 128-bit decimals), dates, times, timestamps, durations, fixed-size binary + and UUID values, string and list offsets, a list's own validity bitmap, fixed-size + list values, and dictionary-encoded columns. String bytes may be copied. + `export_stream!` sends one record batch per `Tables.partitions` partition (an + `Arrow.Stream` is a partition per batch) and splits a batch only when its string data + outgrows 32-bit offsets; `to_c_data` refuses a table that needs more than one batch. +- Imports have the column types an IPC read gives: strings as `DataString`, timestamps + as `Timestamp` or `ZonedTimestamp`, decimals as `Decimal`, nested values as the same + views, with `stringtype`, `ntasks`, and `limits` as for `Arrow.Table`. Numbers, dates, + and other fixed-width values read the producer's memory in place; string and binary + values are copied into Julia-owned buffers. A C stream read whole chains its batches + into one table. `copy(table)` copies all of it into Julia-owned memory. +- Exported and imported data is validated as an IPC read is. + +## Memory + +An export keeps its source reachable until the consumer has released it: the table's +columns, or the vectors it was built from. The structs Arrow.jl allocated are then freed +by later exports or garbage collections, with nothing to call. The consumer reads +exported vectors in place, so modifying one changes what it reads; resizing one is safe +from Julia 1.11 on, where the exported storage stays alive, and frees the consumer's +bytes on Julia 1.10. + +An imported table's columns read the producer's memory, and the producer's `release` +callback runs once the table and its columns are garbage collected, as a mapped file is +unmapped. [`Arrow.release!`](@ref)`(table)` only refuses further column access. On an +[`Arrow.Stream`](@ref) over a C stream, `Arrow.release!` of the stream or of a table it +yielded releases the producer's stream at once and refuses further column access, as +for any stream; columns already taken stay valid. `deepcopy` of imported data shares +its memory and its release with the original; use `copy`. + +## [Threads](@id cdata-threads) + +The callbacks Arrow.jl exports are Julia functions. A consumer that calls them from its +own threads, as DuckDB's and Arrow C++'s thread pools do, makes Julia adopt those +threads. Such a consumer can deadlock with Julia's garbage collector while the Julia +thread that called it waits for it: a collection waits for that Julia thread, which +waits for the consumer, whose threads wait inside the callbacks for the collection. +DuckDB's scan of a `RecordBatchReader` works this way, so read the stream with +`read_all()` on the calling thread first and give DuckDB the pyarrow table. + +On Julia 1.12 and 1.13 a process in which such a consumer thread ran a callback hangs +at exit, when the consumer's library joins its threads (a Julia bug; 1.10 and 1.11 exit +normally). Closing the consumer, as with DuckDB's `connection.close()`, does not avoid +it, because the threads belong to the library's thread pool. Until Julia fixes it, +`ccall(:_exit, Cvoid, (Cint,), 0)` ends such a process at once; it skips exit hooks, so +flush output and close files first. + +A producer's `release` callback for imported data runs on the thread that imported +it, or on thread 1 when another library's thread imported it inside a callback: when a +garbage collection on another thread finds the data unreachable, that thread runs the +callback at its next yield. pyarrow's release of memory it shares with NumPy takes +Python's GIL, which PythonCall's main thread holds, so data imported on the main thread +is released there without waiting on the GIL. + +## Low-level columns + +[`Arrow.to_c_data`](@ref)`(field, data)` exports one column given as an `Arrow.Field` +and an `Arrow.ArrayData`, and [`Arrow.from_c_data`](@ref) imports one, returning that +pair. These are the forms the conformance suite and a custom producer or consumer use; +tables need neither. diff --git a/docs/src/index.md b/docs/src/index.md index 4f900590..bdb6154a 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -21,11 +21,42 @@ # Arrow.jl -```@contents -Pages = ["manual.md", "reference.md"] -Depth = 3 +Arrow.jl reads and writes the [Apache Arrow](https://arrow.apache.org) IPC file and +stream formats in pure Julia. A file reads as Tables.jl columns that view its memory, and +any Tables.jl table writes. The API is namespaced: `Arrow.Table`, `Arrow.write`, and the +rest are called with the `Arrow.` prefix, and only `ArrowTypes` is exported. + +## Quick start + +```@example quickstart +using Arrow, DataFrames + +path = joinpath(mktempdir(), "example.arrow") +Arrow.write(path, (id=[1, 2, 3], name=["Ada", "Ben", missing])) + +table = Arrow.Table(path) +table.name ``` -```@docs -Arrow +```@example quickstart +df = DataFrame(table) ``` + +[Reading](@ref) covers tables, columns, filters, threads, and sources; +[Writing](@ref) covers the writer's options, dictionaries, and incremental writes; +[Streams and IPC](@ref) covers record batches, pipes, and sockets; [Type mapping](@ref) +lists the column types in both directions; [Limits and errors](@ref) covers resource +limits, errors, and what is not supported. Upgrading from Arrow.jl 2.x? Read +[Migrating from Arrow.jl 2.x](@ref) first. The [API reference](@ref "API reference") +lists the public types and functions. + +## Format + +Arrow.jl implements version 1.5 of the Arrow columnar format: the IPC file and stream +formats, LZ4 frame and Zstandard compression, dictionary encoding with delta and +replacement dictionaries, extension types, and the [C data interface](@ref "C data +interface"). It reads every data type of that version, and writes all of them except +dictionary encoding inside a nested column. A few types have no Julia type of their own, +such as run-end encoding, list views, sparse unions, and date64; they are written when a +table read from Arrow is written again ([Rewriting tables](@ref)). It needs no Arrow C++ +library. diff --git a/docs/src/limits.md b/docs/src/limits.md new file mode 100644 index 00000000..e213b363 --- /dev/null +++ b/docs/src/limits.md @@ -0,0 +1,123 @@ +```@raw html + +``` + +# Limits and errors + +## Errors + +Reads and writes throw one of two package errors, or one of these Base errors: + +- [`Arrow.ValidationError`](@ref): the input is not valid Arrow data, such as a truncated + file, inconsistent metadata, an offset or dictionary index out of range, or bytes that + are not Arrow at all; it uses something this package does not read ([Not + supported](@ref)); or it declares a size past one of the [Resource limits](@ref). The + message names the problem, and for a limit the field to raise. +- [`Arrow.AllocationLimitError`](@ref): a read needed more memory than its allocation + budget allows ([Resource limits](@ref)); the message names the `Arrow.Limits` fields + to raise. +- `ArgumentError`: an argument that cannot work, such as an unknown keyword, a `select` + that matches no column, a filter literal of another kind than its column, or, when + writing, a value with no Arrow type, invalid UTF-8, or a table that does not fit an + `Arrow.Writer`'s schema. The message names the column. +- `InvalidStateException`: a table or stream used after [`Arrow.release!`](@ref + Arrow.release!(::Arrow.Table)). +- `SystemError` and `Base.IOError`: the file system, such as a missing file or a + read-only one. + +```@example limits +using Arrow + +bytes = take!(Arrow.tobuffer((id=collect(1:10_000),))) +try + Arrow.Table(bytes[1:(end - 100)]) +catch err + showerror(stdout, err) +end +``` + +## Resource limits + +Pass an [`Arrow.Limits`](@ref) as the `limits` keyword of `Arrow.Table` or `Arrow.Stream`: + +| field | default | caps | +|:--|:--|:--| +| `max_metadata_bytes` | 16 MiB | the metadata of one message | +| `max_body_bytes` | none | the body of one message | +| `max_buffer_bytes` | none | one buffer, compressed or not | +| `max_total_allocated_bytes` | physical memory | the memory one read, or one batch of a stream, allocates | +| `max_allocation_ratio` | 64 | that memory per byte of input | +| `max_messages` | 1,000,000 | messages in one file or table | +| `max_metadata_objects` | 1,000,000 | FlatBuffers objects in one message | +| `max_nesting_depth` | 64 | nesting of types and metadata | +| `max_array_length` | 1,000,000,000 | the length of one array | +| `max_concurrent_reads` | 64 | ranges an `Arrow.AbstractArrowSource` reads at once | +| `max_new_symbols` | 1,000,000 | new `Symbol` values one read interns | + +A read may allocate `max_allocation_ratio` bytes per byte of its input, at least 256 MiB +and at most `max_total_allocated_bytes`; each batch of an `Arrow.Stream` gets its own +allowance from its own bytes. A compressed buffer counts at its decompressed size, so +ordinary files of any size and compression ratio read with the defaults. Lengths that a +file declares without the bytes to back them fail before memory is reserved for them. +Data that expands much further than its bytes, such as long runs or a heavily repeated +dictionary value, can need a larger `max_allocation_ratio`. + +```@example limits +try + Arrow.Table(bytes; limits=Arrow.Limits(max_array_length=1_000)) +catch err + showerror(stdout, err) +end +``` + +Symbols last for the rest of the process, so a read interns at most `max_new_symbols` +new `Symbol` values from `Symbol` columns, and refuses a column that needs more before +interning any. Such a column can be read as strings: `Arrow.Table(path; select=[:id => +String])`. A filter on that column still compares `Symbol`s, so it interns them. + +Column names and time-zone names become `Symbol`s too. A read refuses, with a +`ValidationError`, a name longer than 4096 bytes, a schema that would intern more than +65,536 new column names or 1 MiB of them, and more than 1,024 new time zones. + +## Validation + +Every record batch is checked before it is used: the metadata by a FlatBuffers +verifier, then each buffer's size against the schema, offsets, dictionary indices, union +type ids, and run ends. Corrupt or hostile input therefore fails with a +`ValidationError` or an `AllocationLimitError` rather than crashing, hanging, or +allocating without bound; the package's fuzz tests check this. Three things the format +leaves to the writer are not checked, as other implementations do not check them: that +strings read are valid UTF-8, that a field declared non-nullable has no nulls (such a +column reads as `Union{Missing,T}`), and the padding bytes. + +Writes validate every record batch before its bytes reach the sink. + +## Not supported + +These fail with a `ValidationError`, or an `ArgumentError` when writing: + +- Big-endian IPC data. +- Metadata versions before V4, and the pre-1.0 experimental compression that predates + the format's body compression. +- Appending to a file-format file; append to the stream format or use an `Arrow.Writer`. +- Writing a `BigInt`, an `@enum`, a multidimensional array, or another type with no Arrow + type: convert it, or map it with ArrowTypes ([Custom types](@ref)). +- Writing dictionary encoding inside a nested column (it reads), and rewriting, as it + was read, a union with two children of the same type ([Rewriting tables](@ref)). diff --git a/docs/src/manual.md b/docs/src/manual.md deleted file mode 100644 index 5a3330fd..00000000 --- a/docs/src/manual.md +++ /dev/null @@ -1,247 +0,0 @@ -```@raw html - -``` - -# 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. - -## 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` - -When a custom `IO` subtype is returned (`I`) then the following methods also need to be defined: - -- `Base.read(io::I, ::Type{UInt8})` or `Base.read(io::I)` -- `Base.write(io::I, x)` - -## Reading arrow data - -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: - -```julia -using Arrow - -table = Arrow.Table("data.arrow") -``` - -### `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.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. - -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). - -So, what can you do with an `Arrow.Table` full of data? Quite a bit actually! - -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: - -* `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 full list of Julia packages leveraging the Tables.jl inteface can be found [here](https://github.com/JuliaData/Tables.jl/blob/master/INTEGRATIONS.md). - -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)` - -### Arrow types - -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: - -* `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` - - -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. - -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. - -#### Custom types - -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: - -```julia -using Arrow - -struct Person - id::Int - name::String -end - -# 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 - -table = (col1=[Person(1, "Bob"), Person(2, "Jane")],) -io = IOBuffer() -Arrow.write(io, table) -seekstart(io) -table2 = Arrow.Table(io) -``` - -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`. - -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. - -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: - -```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 -``` - -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. - -Let's run through one more complex example, just for fun and to really see how far the system can be pushed: - -```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]} -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) -``` - -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`. - -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. - -### `Arrow.Stream` - -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`. - -### Custom application metadata - -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. - -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). - -To attach custom schema/column metadata to Arrow tables at serialization time, see the `metadata` and `colmetadata` keyword arguments to [`Arrow.write`](@ref). - -## Writing arrow data - -Ok, so that's a pretty good rundown of *reading* arrow data, but how do you *produce* arrow data? Enter `Arrow.write`. - -### `Arrow.write` - -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. - -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 - -And these are just a few examples of the numerous [integrations](https://github.com/JuliaData/Tables.jl/blob/master/INTEGRATIONS.md). - -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. - -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: - -```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) - -# 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) -``` - -### `Arrow.Writer` - -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. - -### Multithreaded writing - -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. - -### Compression - -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. diff --git a/docs/src/migration.md b/docs/src/migration.md new file mode 100644 index 00000000..abb88531 --- /dev/null +++ b/docs/src/migration.md @@ -0,0 +1,162 @@ +```@raw html + +``` + +# Migrating from Arrow.jl 2.x + +Arrow.jl 3.0 is a new implementation with the same main names: `Arrow.Table`, +`Arrow.Stream`, `Arrow.write`, `Arrow.Writer`, `Arrow.append`, `Arrow.tobuffer`, +`Arrow.DictEncode`, and `Arrow.getmetadata`. It reads Arrow 2.x files using the Julia +types described below. Code that reads a table into a `DataFrame` and writes it back +usually runs unchanged. Files containing `ZonedDateTime` written before Arrow 2.4.2 +need rewriting first; see [Reading and writing alongside Arrow.jl 2.x](@ref). +Other changes include a few removed arguments and, for teams that run both versions, +column types that Arrow 2.x cannot read correctly. + +## Before you upgrade + +Arrow 3.0 needs Julia 1.10 or later and uses ArrowTypes.jl 2, Tables.jl 1.14, +DataStrings.jl, DataDecimals.jl, and Durations.jl. It no longer depends on TimeZones.jl +or BitIntegers.jl. + +Some registered packages still declare compatibility with Arrow 2 only, so adding Arrow 3 +can make the package manager hold them, or downgrade them, to older versions; until its +next release this includes CategoricalArrays, through its Arrow extension. After +upgrading, check `Pkg.status()` for packages that went backwards. + +## Calls + +| Arrow 2.x | Arrow 3.0 | +|:--|:--| +| `Arrow.Table(path; convert=false)` | removed: values always convert to Julia types ([Values](@ref)) | +| `Arrow.Table(bytes, pos, len)` | `Arrow.Table(bytes[pos:(pos + len - 1)])` | +| `Arrow.Table(path)`, `Arrow.write(path, table)` with a path object, such as a FilePathsBase.jl path or an AWSS3.jl `S3Path` | unchanged: any object that `open(path, "r")` (to read) or `open(path, "w")` (to write) accepts works | +| `Arrow.Table([path1, path2])`, `Arrow.Stream([...])` | read each source: `reduce(vcat, DataFrame.(Arrow.Table.(paths)))`, or `for path in paths, batch in Arrow.Stream(path)` | +| `table.col[i] = x` | columns are read-only: `df = DataFrame(table); df.col[i] = x` | +| `table.col isa Arrow.List` (or `Map`, `Struct`, `DenseUnion`, ...) | the column types are now `Arrow.ListVector`, `MapVector`, `StructVector`, and plain `Vector`s for unions; better, test the element type | +| `Arrow.getmetadata(table)` | unchanged; `DataAPI.metadata(table)` also works | +| `Arrow.getmetadata(table.col)` | `DataAPI.colmetadata(table, :col)` | +| `Arrow.write(...; compress=LZ4FrameCompressor())` | `compress=:lz4` (or `:zstd`), with `compresslevel` | +| `Arrow.write(...; alignment, dictencodenested, denseunions, maxdepth)` | accepted with a warning and ignored | +| `Arrow.write(...; ntasks=0)` | `ntasks=1` | +| `Arrow.ToTimestamp(x)` | removed: write `Durations.ZonedTimestamp` values, or `ZonedDateTime`s with TimeZones loaded | +| `open(Arrow.Writer, path) do writer ... end` | unchanged, or `Arrow.Writer(path) do writer ... end` | + +A keyword Arrow 3.0 does not know is an `ArgumentError` that lists the valid ones and +suggests the name to use for pyarrow's, Polars', and Arrow 2.x's names. `release!` is new +and public, and like every name but `ArrowTypes` it is called with the `Arrow.` prefix. + +## Values + +Columns are still views of the file, and most element types are the same as in 2.x. +These changed: + +| Arrow type | Arrow 2.x | Arrow 3.0 | To get the 2.x value | +|:--|:--|:--|:--| +| Utf8 | `String` | `DataStrings.DataString` | `String(x)`; or read with `stringtype=String` | +| Binary | `Base.CodeUnits` | `DataStrings.DataBytes` | `Vector{UInt8}(x)` | +| Timestamp | `DateTime`, truncated to milliseconds | `Durations.Timestamp{unit}`, exact | `DateTime(x)`; or select the column as `:col => DateTime` | +| Timestamp with a zone | `TimeZones.ZonedDateTime` | `Durations.ZonedTimestamp{unit,zone}` | `ZonedDateTime(x)`, with TimeZones loaded | +| Decimal | `Arrow.Decimal{P,S,Int128}` | `DataDecimals.Decimal{P,S,T}` | | +| Interval | `Arrow.Interval` | `Durations.Duration` | | +| List | a `SubArray` view | `Arrow.ListValue{T}` view | `copy(x)` | +| FixedSizeList with nullable values | a `Tuple` | `Arrow.ListValue{T}` | `Tuple(x)` | +| Map | `Dict{K,V}` | `Arrow.MapValue{K,V}`, an `AbstractDict` in stored order | `Dict(x)`; `Dict{String,V}(x)` for `String` keys | +| Struct | `NamedTuple` | `NamedTuple`, with `DataString` fields for strings | | + +A `DataString` compares and hashes equal to the `String` with the same characters, and a +`Timestamp` to the `DateTime` at the same time, so lookups, joins, and comparisons +against `String` and `DateTime` values keep working. Code that dispatches on, or asserts, +`String`, `DateTime`, or `Dict` element types needs the conversions above: + +```@example migration +using Arrow, Dates + +path = joinpath(mktempdir(), "values.arrow") +Arrow.write(path, (name=["Ada"], at=[DateTime(2025, 1, 1, 12)], attrs=[Dict("k" => 1)])) +table = Arrow.Table(path) +(table.name[1] == "Ada", String(table.name[1]), DateTime(table.at[1]), + Dict{String,Int}(table.attrs[1])) +``` + +[Type mapping](@ref) lists every type. + +Ordered dictionary columns keep their order in `DataAPI.levels`, and with +CategoricalArrays loaded `isordered(column)` reports it; see [Dictionary +columns](@ref) for rebuilding an ordered `CategoricalArray`. + +`Symbol` columns intern at most `Arrow.Limits().max_new_symbols` new symbols per read +([Resource limits](@ref)). + +## Writing changes + +- A path still gets the file format and an `IO` the stream format; `file=` overrides. +- `Arrow.Writer` takes its schema from the first table, with every column nullable + unless `nullable=false`. Later tables, and tables given to `Arrow.append`, must have + the same columns with values that convert exactly (an `Int32` column fits an `Int64` + one), or the write is an error naming the column. Arrow 2.x did not check them. +- `dictencode=false` writes `PooledArray` and `CategoricalArray` columns as plain values; + 2.x dictionary-encoded them regardless. `dictencode` also takes column names. +- Strings must be valid UTF-8; write other bytes as `Vector{UInt8}`. +- `UUID` columns write as the canonical `arrow.uuid` extension, which pyarrow and Polars + read as UUIDs. Files from 2.x, which used a `JuliaLang.UUID` label, still read as + `UUID`. +- `Int128` and `UInt128` write as 16-byte fixed-size binary labelled `JuliaLang.Int128` + and `JuliaLang.UInt128`. +- A vector of `Pair`s writes as an Arrow map, where 2.x wrote a list of structs. +- `metadata` and `colmetadata` default to the table's DataAPI metadata, so a + `DataFrame`'s metadata is written. +- New keywords: `compresslevel`, `batchsize`, `dictdeltas`, `stringview`, and, for + `Arrow.Writer`, `nullable` and `dictreplacement` ([Writing](@ref)). + +## Packages that define mappings + +ArrowTypes.jl 2 keeps the 2.x interface. Define `ArrowTypes.ArrowType` and +`ArrowTypes.toarrow` to write a type as a supported one, and `ArrowTypes.arrowname`, +`ArrowTypes.JuliaType`, and `ArrowTypes.fromarrow` to read it back ([Custom +types](@ref)). Files written with 2.x mappings read back through the same definitions. +Overriding `ArrowTypes.ArrowKind` alone no longer chooses how a type is stored. A +package that defines mappings should depend on ArrowTypes.jl directly rather than reach +it through Arrow. + +## Reading and writing alongside Arrow.jl 2.x + +Arrow 3.0 reads Arrow 2.x files as the types listed under [Values](@ref), with the +older `ZonedDateTime` exception described below. The other direction has gaps: +Arrow 2.x reads most files Arrow 3.0 writes, but not all. While some readers still +run Arrow 2.x, avoid these columns or write them in the form the last column suggests: + +| Written by Arrow 3.0 | Arrow 2.x reads | For 2.x readers, write | +|:--|:--|:--| +| `UUID` | `NTuple{16,UInt8}`, with a warning | `string.(uuids)` | +| `Int128`, `UInt128` | `NTuple{16,UInt8}`, with a warning | `string.(x)` | +| `Minute`, `Hour`, `Day`, `Week` | an `Arrow.Duration` of seconds, with a warning | `Dates.Second.(x)` | +| `Month`, `Year` | an `Arrow.Interval` of months, with a warning | `Dates.value.(x)` | +| `VersionNumber` | fails when the column is read | `string.(versions)` | +| `Durations.Duration` | cannot open the file | separate columns, or `string.(x)` | +| A dictionary column that gained values across partitions or `Arrow.Writer` tables, in the file format | with several threads, can fail with a `KeyError` | `dictdeltas=true`, or the stream format | +| Decimal32 or Decimal64, kept from a file another tool wrote | wrong values, without an error | `DataFrame(table)` first, which writes Decimal128 | + +Arrow 2.x reads the other columns as its own types, as listed under [Values](@ref): +timestamps as `DateTime`, strings as `String`, maps as `Dict`s. + +Arrow.jl releases before 2.4.2 stored `ZonedDateTime` columns as local times, labelled +`JuliaLang.ZonedDateTime`. Arrow 3.0 warns about that label and reads the stored values +as UTC instants, so rewrite such files with Arrow 2.x before reading them with 3.0. diff --git a/docs/src/reading.md b/docs/src/reading.md new file mode 100644 index 00000000..20678b78 --- /dev/null +++ b/docs/src/reading.md @@ -0,0 +1,264 @@ +```@raw html + +``` + +# Reading + +## Read a table + +[`Arrow.Table(source)`](@ref Arrow.Table) reads Arrow IPC data, in the file or the stream +format, into columns. A table is a `Tables.AbstractColumns`: `table.fare`, +`table[:fare]`, `table["fare"]`, `table[2]`, and `Tables.getcolumn(table, :fare)` return a +column, `names(table)` lists them, `DataAPI.nrow(table)` and `DataAPI.ncol(table)` (the +`nrow` and `ncol` DataFrames exports) count rows and columns, and any Tables.jl sink +accepts the table. + +```@example reading +using Arrow, DataFrames, Dates, Tables + +dir = mktempdir() +path = joinpath(dir, "trips.arrow") +Arrow.write(path, (id=1:6, fare=[7.5, 12.0, 3.25, 40.0, 9.5, 21.0], + zone=["A", "B", "A", "C", "B", "A"], day=Date(2025, 1, 1):Day(1):Date(2025, 1, 6))) + +table = Arrow.Table(path) +``` + +A source is a path, an `IO`, a byte vector, or an [`Arrow.AbstractArrowSource`](@ref) +([Sources](@ref)); any other object that `open(source, "r")` accepts, such as a +FilePathsBase.jl path, reads as the file or `IO` it opens. With FileIO.jl loaded, +`load("trips.arrow")` reads a table too. Column types are listed under +[Type mapping](@ref). Columns that share a name read as `a`, `a_1`, `a_2`, and so on, as +CSV.jl names them. Two tables are `==` when their column names and columns are. + +## Columns are read-only views + +Numbers, booleans, dates, times, timestamps, decimals, and the indices of dictionary +columns are read in place: such a column is an [`Arrow.ArrowVector`](@ref) that views the +mapped file, the byte vector, or the decompressed buffer it was read from. Opening a +table costs time per column and record batch, not per byte. A column of a file with +several record batches is a `SentinelArrays.ChainedVector` of one view per batch. + +```@example reading +table.fare +``` + +Strings and binary values read as a `DataStrings.StringVector` of `DataString`s and a +`BytesVector` of `DataBytes`, built the first time the table hands the column out: 16 +bytes per value, over the source's own bytes. A string longer than 12 bytes is a view +into the mapped file, the byte vector, or the decompressed buffer, so a `DataString` +taken from the column keeps those bytes alive until it is garbage collected; +`String(x)` copies it out. `stringtype=String` reads strings as `String`s, nested ones +included, at the cost of one allocation per value. Lists, maps, and structs read as +typed rows ([Nested values](@ref)). + +Columns are read-only: + +```@example reading +try + table.fare[1] = 0.0 +catch err + err +end +``` + +A column keeps the memory it reads alive for as long as the column is reachable, even +after its table is gone. + +## Copy into Julia vectors + +`copy(column)` and `copy(table)` copy into owned, mutable columns: a `Vector` per column, +a `PooledArray` for a dictionary column, and a `StringVector` for strings. A copied +table's string columns hold their own bytes, so nothing of the source stays +referenced. `DataFrame(table)` copies once, as `DataFrame` always does, but its string +columns still share the source's bytes (DataStrings' `copy` keeps a column's buffers); +`DataFrame(table; copycols=false)` keeps the views. [`Arrow.read(source, +sink)`](@ref Arrow.read) reads and hands fully owned copies to a sink in one step: + +```@example reading +df = Arrow.read(path, DataFrame) +``` + +`copy(table)` and `Arrow.read` also detach nested Arrow views. Lists become vectors, +maps become `Dict`s (a repeated key keeps its last value, as a lookup on the map finds +it), and structs and tuples keep their shape with owned child values. A copied table +therefore keeps no dependency on the input's buffers. Nested values that were already +ordinary Julia arrays follow Julia's usual shallow-copy behavior. + +## Memory mapping and lifetime + +A path is memory mapped by default, so opening a file larger than memory is cheap and +only the pages a read touches are loaded. `mmap=false` reads the file into memory +instead. The mapping lasts until the table and every column taken from it are garbage +collected. + +`Arrow.write` to a mapped path writes a new file and renames it over the old one, so +columns of the old file keep reading the old contents. On Linux and macOS, deleting or +renaming a mapped file is safe too. + +!!! warning "Files that other programs rewrite in place" + If another program truncates and rewrites a file in place (`open(path, "w")`, + `CSV.write`, pyarrow, Polars, a shell redirect) while a column of a mapped table of + that file is alive, the next read of that column crashes Julia with a bus error + (SIGBUS). No check can prevent it: the operating system removes the pages under the + mapping. Read such files with `mmap=false`. + +!!! note "Windows" + Windows cannot delete a file while it is mapped, so there a path is read into memory + unless `mmap=true` is passed. With `mmap=true`, deleting the file fails until the + table and its columns are garbage collected; `Arrow.write` can still replace it. + +[`Arrow.release!(table)`](@ref Arrow.release!(::Arrow.Table)) ends a read: the table +refuses further column access. A column taken before the release keeps working, and its +memory is freed when it is garbage collected. + +```@example reading +t = Arrow.Table(path) +fares = t.fare +Arrow.release!(t) +sum(fares) +``` + +## Choose columns and rows + +`select` and `drop` choose columns (names, indices, a regex, a vector of them, or a +function `(i, name) -> Bool` as in CSV.jl), `filter` keeps the rows where a predicate is +`true`, `limit` caps the row count, and `offset` skips rows first. Only the selected and +filter columns are read, a filter's other columns only for the record batches that keep +rows, and a `limit` stops reading once it has its rows. + +```@example reading +DataFrame(Arrow.Table(path; select=[:id, :fare], filter=Arrow.col(:fare) > 10, limit=2)) +``` + +Predicates are the `Tables.col` family, which Arrow provides under its own names, so a +filter needs no other package: + +- [`Arrow.col`](@ref) with `<`, `<=`, `>`, and `>=` against numbers, strings, and + characters; +- [`Arrow.colcmp`](@ref) for equality and for other values, such as dates and + timestamps; +- [`Arrow.colin`](@ref) for membership and [`Arrow.isnull`](@ref) for missing values; +- `&`, `|`, and `!` to combine them. + +A row whose predicate is `missing` is dropped. + +```@example reading +keep = Arrow.colcmp(==, Arrow.col(:zone), "A") & + Arrow.colcmp(>=, Arrow.col(:day), Date(2025, 1, 2)) +DataFrame(Arrow.Table(path; filter=keep)) +``` + +A literal must be of its column's kind: comparing a string column with a number is an +`ArgumentError` that names the column and both kinds of value. A `Date`, `DateTime`, or `Durations.Timestamp` +compared with a UTC timestamp column (pandas' time-zone-aware data, DuckDB's +TIMESTAMPTZ) is that time in UTC. A column in another zone needs a +`Durations.ZonedTimestamp` literal, and the error says how to write one. + +`select` also takes `name => newname` to rename a column and `name => Type` to convert +it, such as `:day => DateTime`; `scan` takes a `Tables.Scan` in place of all these +keywords. + +```@example reading +DataFrame(Arrow.Table(path; select=[:id, :fare => :amount], offset=1, limit=3)) +``` + +A read that keeps every row reads its columns as a full read does: as views. A read with +`filter`, `limit`, or `offset` copies the rows it keeps: numbers and other values into +`Vector`s, strings and binary values into columns that hold just those values' bytes, +and dictionary columns with their categories. These copies keep nothing else of the +source alive, but kept list, map, and struct rows still read their record batch's +buffers, which stay in memory (or keep the file mapped) as long as the rows do. A +column's element type does not depend on the filter or on the kind of source. + +## Threads + +A read uses up to `ntasks` tasks (default `Threads.nthreads()`): record batches decode +and decompress in parallel, columns build in parallel, a long string column builds in +chunks of rows, and a filter is evaluated on several record batches at once, and in +chunks of rows of a large one. `ntasks=1` reads on the calling task. The result is the same for every +`ntasks`. Start Julia with threads (`julia -t 8`) for the tasks to run at once. + +## Metadata + +`DataAPI.metadata(table)` gives the schema's key-value metadata and +`DataAPI.colmetadata(table, :fare)` a column's. `DataFrame(table)` and `Arrow.read(path, +DataFrame)` keep both, and writing the table again keeps them too. +[`Arrow.getmetadata(table)`](@ref Arrow.getmetadata) is the Arrow 2.x form. + +```@example reading +using DataAPI +meta = joinpath(dir, "meta.arrow") +Arrow.write(meta, (fare=[7.5],); metadata=["source" => "meter"], + colmetadata=Dict(:fare => ["unit" => "USD"])) +t = Arrow.Table(meta) +(DataAPI.metadata(t), DataAPI.colmetadata(t, :fare)) +``` + +## Sources + +- A **path** is memory mapped ([Memory mapping and lifetime](@ref)). +- A **byte vector** is read in place: do not modify or resize it while its table's + columns are in use. +- An **`IO`** is read to its end into memory by `Arrow.Table`; [`Arrow.Stream`](@ref) + reads a stream-format `IO` one record batch at a time ([Streams and IPC](@ref)). The + `IO` is not closed. +- An [`Arrow.AbstractArrowSource`](@ref) supplies bytes by range, for storage the others + do not cover. It implements [`Arrow.sourcelength`](@ref) and [`Arrow.readrange`](@ref), + and [`Arrow.concurrentreads`](@ref) when ranges may be read on several tasks at once. + A read with `select`, `drop`, `filter`, `limit`, `offset`, or `scan` reads the + file-format footer from the end of the source, then plans only the byte ranges of the + record batches and columns it needs, whatever the filter compares, and fetches them in + a few coalesced ranges. A record batch of at most 1 MiB is fetched whole with its + metadata, so a file of many small batches is fetched almost whole. A `select`'s columns + view the fetched bytes in place, except that bytes filling less than half of a fetched + range are copied out of it, so a few columns do not keep the whole object in memory. A + read without them fetches the whole source in 8 MiB ranges. + +```@example reading +struct CountingSource <: Arrow.AbstractArrowSource + bytes::Vector{UInt8} + fetched::Base.RefValue{Int} +end +Arrow.sourcelength(src::CountingSource) = length(src.bytes) +function Arrow.readrange(src::CountingSource, offset, count) + src.fetched[] += count + return src.bytes[(offset + 1):(offset + count)] +end + +big = joinpath(dir, "big.arrow") +Arrow.write(big, (id=1:100_000, fare=rand(100_000), note=fill("no comment", 100_000))) +src = CountingSource(read(big), Ref(0)) +fares = Arrow.Table(src; select=[:fare]).fare +(bytes_fetched=src.fetched[], file_bytes=filesize(big)) +``` + +With [CloudStore.jl](https://github.com/JuliaServices/CloudStore.jl) loaded, a +`CloudStore.Object` (S3 or Azure Blob Storage) reads by ranged requests the same way: + +```julia +using Arrow, CloudStore + +object = CloudStore.Object(bucket, "events/2025-01.arrow"; credentials) +table = Arrow.Table(object; select=[:id, :amount]) +``` + +`Serialization.serialize` of a table saves its values and schema; the deserialized table +owns its columns. diff --git a/docs/src/reference.md b/docs/src/reference.md index 0409d76f..86786268 100644 --- a/docs/src/reference.md +++ b/docs/src/reference.md @@ -19,17 +19,85 @@ --> ``` -# API Reference +# API reference -```@autodocs -Modules = [Arrow] -Order = [:type, :function] +```@docs +Arrow ``` -## Internals: `Arrow.FlatBuffers` +## Read -The `FlatBuffers` module is not part of Arrow.jl's public API, and these functions may change without notice. +```@docs +Arrow.Table +Arrow.read +Arrow.Stream +Arrow.release!(::Arrow.Table) +Base.copy(::Arrow.Table) +Arrow.getmetadata +Arrow.Limits +``` + +## Filter + +```@docs +Arrow.col +Arrow.colcmp +Arrow.colin +Arrow.isnull +``` + +## Write + +```@docs +Arrow.write +Arrow.write(::Arrow.Writer, ::Any) +Arrow.Writer +Arrow.append +Arrow.tobuffer +Arrow.DictEncode +``` + +## Column and value types + +```@docs +Arrow.ArrowVector +Arrow.Primitive +Arrow.BoolVector +Arrow.DictEncoded +Arrow.RunEndEncoded +Arrow.ListVector +Arrow.ListValue +Arrow.MapVector +Arrow.MapValue +Arrow.StructVector +``` + +## Byte-range sources + +```@docs +Arrow.AbstractArrowSource +Arrow.sourcelength +Arrow.readrange +Arrow.concurrentreads +``` + +## Error types + +```@docs +Arrow.ValidationError +Arrow.AllocationLimitError +``` + +## C data and C stream interfaces -```@autodocs -Modules = [Arrow.FlatBuffers] +```@docs +Arrow.Table(::Ptr{Arrow.CArrowSchema}, ::Ptr{Arrow.CArrowArray}) +Arrow.Table(::Ptr{Arrow.CArrowArrayStream}) +Arrow.Stream(::Ptr{Arrow.CArrowArrayStream}) +Arrow.to_c_data +Arrow.export_stream! +Arrow.from_c_data +Arrow.CArrowSchema +Arrow.CArrowArray +Arrow.CArrowArrayStream ``` diff --git a/docs/src/streams.md b/docs/src/streams.md new file mode 100644 index 00000000..0eb1f6dc --- /dev/null +++ b/docs/src/streams.md @@ -0,0 +1,137 @@ +```@raw html + +``` + +# Streams and IPC + +Arrow has two IPC formats. The **file** format starts with the `ARROW1` magic and ends +with a footer that indexes the record batches, so a reader can open any batch directly; +it suits files on disk. The **stream** format is a schema message, then dictionaries and +record batches, then an end marker; a reader can process it as it arrives, so it suits +pipes, sockets, and HTTP bodies. `Arrow.Table` and `Arrow.Stream` read both and detect +which one they are given. `Arrow.write` writes the file format to a path and the stream +format to an `IO`. + +## Read record batches + +[`Arrow.Stream(source)`](@ref Arrow.Stream) reads a source one record batch at a time; +each iteration gives an `Arrow.Table` of one batch. It takes the keywords of +`Arrow.Table`: `select`, `drop`, and `filter` apply to every batch, and `limit` and +`offset` count rows across batches. A batch decodes only the selected and filter +columns. Over an [`Arrow.AbstractArrowSource`](@ref) it fetches only their bytes, except +that batches of at most 1 MiB are fetched whole, several to a request. A loop +that processes and drops each batch holds one batch at a time, plus the dictionaries, +so a source larger than memory reads in bounded memory. + +```@example streams +using Arrow, DataAPI, DataFrames, Tables + +path = joinpath(mktempdir(), "batches.arrow") +Arrow.write(path, Tables.partitioner([(id=[1, 2], x=[0.5, 1.5]), (id=[3], x=[2.5])])) + +stream = Arrow.Stream(path) +[sum(batch.x) for batch in stream] +``` + +A stream is its own `Tables.partitions`, so `Arrow.write(out, Arrow.Stream(path))` writes +one record batch per batch read. `Tables.schema(stream)`, `DataAPI.metadata(stream)`, and +`DataAPI.colmetadata(stream, :x)` read no data. For the file format, `length(stream)` is +the number of batches, `stream[i]` reads batch `i`, and `DataAPI.nrow(stream)` counts rows +from the batches' metadata: + +```@example streams +(length(stream), DataAPI.nrow(stream), DataFrame(stream[2])) +``` + +A stream-format source has no index: it is read in order, its batch count is known only +once it has been read, and an `Arrow.Stream` over an `IO` can be iterated once. + +## Read a stream as it arrives + +Over a stream-format `IO`, `Arrow.Stream` reads each record batch when the iteration +reaches it, so a consumer processes batches while the producer is still writing. +[`Arrow.Writer(io)`](@ref Arrow.Writer) is the producer's side: each +`Arrow.write(writer, table)` writes that table's batches to the `IO` before it returns, +and `close(writer)` writes the end marker and leaves the `IO` open. Here the two ends of +a TCP connection are in one process: + +```@example streams +using Sockets + +port, server = listenany(ip"127.0.0.1", 40000) +producer = @async begin + socket = accept(server) + Arrow.Writer(socket) do writer + for i in 1:3 + Arrow.write(writer, (batch=[i, i], value=[10i, 10i + 1])) + end + end + close(socket) +end + +client = connect(ip"127.0.0.1", port) +received = [collect(batch.value) for batch in Arrow.Stream(client)] +wait(producer); close(client); close(server) +received +``` + +An `IO` that buffers its writes, such as a file's `IOStream`, may hold a batch until +`flush(io)`. [`Arrow.append`](@ref) adds batches to the end of an existing stream in a +file or a seekable `IO`. + +`Arrow.Table(io)` reads an `IO` to its end and refuses bytes after a stream's end marker. +To read several streams written one after another to one `IO`, call `Arrow.Stream(io)` +once per stream: each call reads one stream and stops after its end marker. + +```@example streams +io = IOBuffer() +Arrow.write(io, (a=[1],)) +Arrow.write(io, (a=[2, 3],)) +seekstart(io) +[[collect(batch.a) for batch in Arrow.Stream(io)] for _ in 1:2] +``` + +## Bytes in memory + +[`Arrow.tobuffer(table)`](@ref Arrow.tobuffer) writes a table in the stream format to a +new `IOBuffer`, positioned at its start, and `Arrow.Table(bytes)` reads a byte vector in +place: + +```@example streams +bytes = take!(Arrow.tobuffer((a=[1, 2],))) +Arrow.Table(bytes).a +``` + +## Other Arrow implementations + +Other Arrow implementations read both formats: pyarrow reads streams with +`pyarrow.ipc.open_stream` and files with `pyarrow.ipc.open_file`, Polars with +`polars.read_ipc_stream` and `polars.read_ipc`. Arrow.jl reads what they write: + +```python +import pyarrow as pa, pyarrow.ipc as ipc, sys + +table = pa.table({"id": [1, 2, 3]}) +with ipc.new_stream(sys.stdout.buffer, table.schema) as writer: + writer.write_table(table) # julia: for batch in Arrow.Stream(stdin) ... end +``` + +To exchange columns with another implementation in the same process without IPC bytes, +use the [C data interface](@ref "C data interface"). diff --git a/docs/src/types.md b/docs/src/types.md new file mode 100644 index 00000000..9b19adf9 --- /dev/null +++ b/docs/src/types.md @@ -0,0 +1,295 @@ +```@raw html + +``` + +# Type mapping + +## Arrow to Julia + +A column's element type comes from its Arrow type alone, never from its values, so an +empty or all-missing column has the same type as a full one. A nullable field's elements +are `Union{Missing,T}`. + +| Arrow type | Julia element type | +|:--|:--| +| Int8-Int64, UInt8-UInt64, Float16-Float64, Bool | the same type | +| Utf8, LargeUtf8, Utf8View | `DataStrings.DataString` (the column is a `StringVector`); `String` with `stringtype=String` | +| Binary, LargeBinary, BinaryView | `DataStrings.DataBytes` (the column is a `BytesVector`) | +| FixedSizeBinary(N) | `NTuple{N,UInt8}`; `DataBytes` when N > 256 | +| Date32, Date64 | `Dates.Date`, `Dates.DateTime` | +| Time32, Time64 | `Dates.Time` | +| Timestamp(unit) | `Durations.Timestamp{unit}` | +| Timestamp(unit, zone) | `Durations.ZonedTimestamp{unit,Symbol(zone)}` | +| Duration(unit) | `Dates.Second`, `Millisecond`, `Microsecond`, or `Nanosecond` | +| Interval | `Durations.Duration` | +| Decimal32/64/128/256(P, S) | `DataDecimals.Decimal{P,S,T}`, `T` the 32- to 256-bit integer of the type's width | +| List, LargeList, ListView, LargeListView | [`Arrow.ListValue{T}`](@ref Arrow.ListValue) | +| FixedSizeList(N) | `NTuple{N,T}`, or `ListValue{T}` when N > 64 or values may be null | +| Struct | a `NamedTuple` | +| Map | [`Arrow.MapValue{K,V}`](@ref Arrow.MapValue) | +| Dense and sparse Union | a `Union` of the children's types | +| Dictionary, run-end encoded | the values' type, in an [`Arrow.DictEncoded`](@ref) or [`Arrow.RunEndEncoded`](@ref) column | +| Null | `Missing` | + +Extension types read as their Julia type: canonical `arrow.uuid` as `UUIDs.UUID`, +`arrow.bool8` as `Bool`, `arrow.fixed_shape_tensor` as `Array{T,N}` rows, and the labels +Arrow.jl writes ([Julia to Arrow](@ref)) as the Julia type that wrote them. `arrow.json` +and `arrow.opaque` read as their storage, strings or bytes. Another extension name reads +as its storage too, with one warning per name per process, unless a package defines a +mapping for it ([Custom types](@ref)). + +A decimal with a negative scale, or a scale above its precision, reads as it is stored: +integers for Decimal32 and Decimal64, little-endian bytes for Decimal128 and Decimal256. +A decimal value with more digits than its precision reads as stored, as pyarrow reads +it. A time of day outside [0, 24 h) has no `Dates.Time` value and fails the read +with an error naming the column. + +## Julia to Arrow + +Each column's element type decides its Arrow type, at any depth. Values read back as the +same type unless the table says otherwise. + +| Julia element type | Arrow type | Reads back as | +|:--|:--|:--| +| `Int8`-`Int64`, `UInt8`-`UInt64`, `Float16`-`Float64`, `Bool` | the same type | the same type | +| `AbstractString` (`String`, `SubString`, `DataString`) | Utf8 | `DataString` | +| InlineStrings' `String1`-`String255` | Utf8, labelled by InlineStrings | the same type | +| `Symbol`, `Char` | Utf8, UInt32; labelled | `Symbol`, `Char` | +| `Vector{UInt8}`, `DataBytes` | Binary | `DataBytes` | +| `NTuple{N,UInt8}` | FixedSizeBinary(N) | `NTuple{N,UInt8}` | +| `Int128`, `UInt128` | FixedSizeBinary(16), labelled `JuliaLang.Int128`/`UInt128` | the same type | +| `UUID` | `arrow.uuid` | `UUID` | +| `Date`, `Time` | Date32, Time64(nanosecond) | `Date`, `Time` | +| `DateTime` | Timestamp(millisecond) | `Durations.Timestamp{Millisecond}` | +| `Durations.Timestamp{P}` | Timestamp(P) | `Durations.Timestamp{P}` | +| `Durations.ZonedTimestamp{P,Z}` | Timestamp(P, zone `Z`) | `Durations.ZonedTimestamp{P,Z}` | +| `TimeZones.ZonedDateTime` (with TimeZones loaded), one zone per column | Timestamp(millisecond, zone) | `ZonedTimestamp{Millisecond,Z}` | +| `Nanosecond`-`Second` | Duration of that unit | the same type | +| `Minute`, `Hour`, `Day`, `Week` | Duration(second), labelled | the same type | +| `Month`, `Year` | Interval(year-month), labelled | the same type | +| `Durations.Duration` | Interval(month-day-nano) | `Durations.Duration` | +| `DataDecimals.Decimal{P,S}` | Decimal128(P, S), or Decimal256 when P > 38 | `Decimal{P,S}`, stored in `Int128` or `Int256` | +| `PooledArray`, `CategoricalArray` | Dictionary of the pool's values; ordered if the array is | an `Arrow.DictEncoded` column | +| `AbstractVector{T}` | List | `ListValue{T}` | +| `NTuple{N,T}` | FixedSizeList(N) | `NTuple{N,T}` | +| other `Tuple`s | Struct with fields `1`, `2`, ..., labelled | the same type | +| `NamedTuple` | Struct | `NamedTuple` | +| `AbstractDict{K,V}`, `AbstractVector{<:Pair}` | Map | `MapValue{K,V}` | +| `Union{Missing,T}` | the type of `T`, nullable | `Union{Missing,T}` | +| `Union{T1,T2,...}` | dense Union, or one type that every member writes | a `Union` of the members' types as read | +| `VersionNumber`, `Nothing` | Utf8, Null; labelled | the same type | +| `Missing` only | Null | `Missing` | + +"Labelled" means the column carries an `ARROW:extension:name` such as +`JuliaLang.Symbol`, so Arrow.jl reads the Julia type back; other readers see the +storage type. SentinelArrays' `ChainedVector` and CSV.jl's columns write without being +collected first. + +A column whose element type is abstract (`Any`, `Real`, a `NamedTuple` with open field +types) is written as the types its values have: one type writes as that type, and up to +8 types as a dense Union of them. When all the types in such a column, or all the members +of a `Union` element type, write the same Arrow type, label included, the column is one +column of that type and reads back as it: `Union{String,SubString{String}}` writes one +Utf8 column, and two structs whose ArrowTypes.jl mappings give an `Int32` with no label +(or with one label they share) write one Int32 column. Two structs without mappings +that have the same fields are the exception: they write a dense Union. A column of +`Timestamp`s, `ZonedTimestamp`s, or `Decimal`s that mixes units, zones, or precisions is +an error, since one Arrow column has one of each (`Durations.astimezone` gives a column +one zone). A type with no Arrow type, such as `BigInt`, an `@enum`, or a `Matrix`, is an +error that names the column; map it with [Custom types](@ref), or convert it first. + +## Timestamps + +Timestamps read as `Durations.Timestamp{P}` and `Durations.ZonedTimestamp{P,Z}` at the +column's unit: 8-byte values that hold the stored instant exactly, at every unit, with no +time-zone database. A `Timestamp` compares and hashes equal to the `DateTime` at the +same time, so a `DateTime` column reads back equal to what was written, and +`DateTime(ts)` converts (truncating below a millisecond). + +A `ZonedTimestamp{P,Z}` holds a UTC instant and its zone `Z`. Instants compare with each +other across zones, but never equal a `DateTime`, which has no zone; `DateTime(zt)` gives +the local time in `Z`. UTC and fixed offsets such as `+05:30` need nothing else; named +zones such as `America/New_York` need TimeZones.jl loaded for local times and for +`ZonedDateTime(zt)`. + +```@example types +using Arrow, Dates, Durations + +io = IOBuffer() +Arrow.write(io, (at=[DateTime(2025, 1, 1, 12)], + utc=[Durations.ZonedTimestamp{Millisecond,:UTC}(2025, 1, 1, 12)])) +table = Arrow.Table(take!(io)) +table.at +``` + +```@example types +table.at[1] == DateTime(2025, 1, 1, 12) +``` + +```@example types +table.utc +``` + +To read a timestamp column as `DateTime`s, give the type in `select`: +`Arrow.Table(path; select=[:id, :at => DateTime])`. + +## Decimals + +Decimals read as `DataDecimals.Decimal{P,S,T}` whose storage integer `T` matches the +Arrow type's width: `Int32` for Decimal32, `Int64` for Decimal64, `Int128` for +Decimal128, and `DataDecimals.Int256` for Decimal256. They write as Decimal128 up to 38 +digits and Decimal256 above, as pyarrow, Polars, and DuckDB write them, whatever the +storage of the Julia value. So `dec"1.25"`, a `Decimal{3,2,Int32}`, reads back as the +equal `Decimal{3,2,Int128}`. Rewriting a column read from a Decimal32 or Decimal64 field +keeps that width. + +## Strings and bytes + +A `DataStrings.DataString` is a 16-byte value: strings of up to 12 bytes are stored in +the value, longer ones refer to the column's bytes. It is an `AbstractString` that +compares and hashes equal to the `String` with the same characters; `String(s)` makes +one. A `StringVector` writes as Utf8 in one pass over its values, or as Utf8View with +`stringview=true`. DuckDB 1.5's `read_arrow` and other older readers refuse the view +types, so they are written only on request. + +## Nested values + +Lists, maps, and structs read as typed rows, in [`Arrow.ListVector`](@ref), +[`Arrow.MapVector`](@ref), and [`Arrow.StructVector`](@ref) columns: + +- A list row is an [`Arrow.ListValue{T}`](@ref Arrow.ListValue), a read-only + `AbstractVector{T}` view of the column's values; `copy(row)` gives a `Vector{T}`. +- A struct row is a `NamedTuple` of its fields, built when it is read. A struct with + more than 64 fields, or with repeated field names, reads as a `Vector` of + `name => value` pairs instead. +- A map row is an [`Arrow.MapValue{K,V}`](@ref Arrow.MapValue), an `AbstractDict{K,V}` + over the map's entries in stored order: `row[key]`, `get`, and `haskey` look keys up, + and iteration gives `key => value` pairs, duplicate keys included (a lookup finds the + last). `Dict(row)` builds a `Dict`. +- A union value is a value of one of the children's types. + +`T` includes `Missing` only when the schema allows null elements, and a null list, +struct, or map reads as `missing`. Rows stay valid after their table is gone. Copies of a +nested column (`copy`, `DataFrame(table)`) are `Vector`s of plain rows, and the writer +takes plain rows back: vectors, named tuples, dictionaries, and vectors of pairs. + +```@example types +io = IOBuffer() +Arrow.write(io, (tags=[["a", "b"], String[]], point=[(x=1.0, y=2.0), (x=3.0, y=4.0)], + attrs=[Dict("k" => 1), Dict("j" => 2, "k" => 3)])) +nested = Arrow.Table(take!(io)) +nested.attrs +``` + +```@example types +(nested.point[2].y, nested.attrs[2]["k"], copy(nested.tags[1])) +``` + +Each level of a typed row nests one more type, and Julia compiles deep ones slowly. A +column nested more than 8 levels deep therefore reads with untyped rows: lists as +`Vector{Any}` and structs as `name => value` pairs. + +## Dictionary columns + +A dictionary column reads as an [`Arrow.DictEncoded`](@ref) of its values' type. +`DataAPI.refarray`, `DataAPI.refpool`, and `DataAPI.levels` expose the encoding without +decoding it; `levels` follows the dictionary's order. `copy(column)` gives a +`PooledArray`. With CategoricalArrays.jl loaded, `isordered(column)` tells whether the +dictionary is ordered, and `categorical` rebuilds a `CategoricalArray` with its order: + +```@example types +using CategoricalArrays, DataAPI + +sizes = categorical(["small", "large", "medium"]; levels=["small", "medium", "large"], + ordered=true) +io = IOBuffer() +Arrow.write(io, (size=sizes,)) +column = Arrow.Table(take!(io)).size +cat = categorical(column; levels=DataAPI.levels(column), ordered=isordered(column)) +(levels(cat), cat[1] < cat[2]) +``` + +## Custom types + +[ArrowTypes.jl](https://github.com/apache/arrow-julia/tree/main/src/ArrowTypes) maps +Julia types to Arrow types. A package that owns a type depends on ArrowTypes and defines +`ArrowTypes.ArrowType` and `ArrowTypes.toarrow` to write it as a supported type; to read +it back, it also names an extension with `ArrowTypes.arrowname` and defines +`ArrowTypes.JuliaType` and `ArrowTypes.fromarrow`: + +```@example types +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) + +io = IOBuffer() +Arrow.write(io, (id=AccountID.(1:3),)) +Arrow.Table(take!(io)).id +``` + +The mapping applies at any depth: in lists, structs, maps, unions, and dictionaries. +An `ArrowType` that is a `Union` with `Missing` as a member is refused, since a stored +`missing` could not be told from a null value; declare the storage without `Missing` and +write the column's missing values as such. +`ArrowTypes.arrowmetadata` adds a string the reader passes to `JuliaType`. Without a +mapping, a struct with supported field types writes as an Arrow struct and reads back as +a `NamedTuple`. In a process that has not loaded the mapping, the column reads as its +storage, with one warning. + +## Rewriting tables + +Writing an `Arrow.Table` or `Arrow.Stream` keeps its schema: each column's Arrow type, +with its units, offset widths, union mode, dictionary index type and order, run-end +encoding, and list views, along with nullability and metadata. Copying it into Julia +vectors first (`DataFrame(table)`, `copy(table)`) keeps only what the Julia types say, +so some types change on the way back while the values stay equal: + +| Read from | Written from Julia vectors as | +|:--|:--| +| LargeUtf8, Utf8View; LargeBinary, BinaryView | Utf8; Binary | +| LargeList, ListView, FixedSizeList with nullable values | List | +| Date64 | Timestamp(millisecond) | +| Time32, Time64(microsecond) | Time64(nanosecond) | +| Decimal32, Decimal64 | Decimal128 | +| Dictionary with Int16-Int64 indices | Dictionary with the narrowest indices | +| Ordered Dictionary | Dictionary, not ordered: a `PooledArray` has no order ([Dictionary columns](@ref) shows how to keep it) | +| Run-end encoded | plain values | +| Sparse or dense Union | dense Union with a Null child | +| `arrow.bool8`, `arrow.json` | Bool, Utf8 | +| Interval(year-month or day-time) | Interval(month-day-nano) | + +A rewrite keeps every type except Int(128) and UInt(128), which Polars and Arrow 2.x +write and pyarrow cannot read: they rewrite as FixedSizeBinary(16) labelled +`JuliaLang.Int128` or `JuliaLang.UInt128`, as a fresh Julia `Int128` column writes, which +Arrow.jl reads back as the same values and Polars reads as binary. + +A union with two children of the same type cannot be rewritten as it was, because its +values do not say which child they came from; writing it is an error that suggests +`Arrow.write(path, DataFrame(table))`. A dictionary column nested inside another column +is refused the same way. diff --git a/docs/src/writing.md b/docs/src/writing.md new file mode 100644 index 00000000..e307c914 --- /dev/null +++ b/docs/src/writing.md @@ -0,0 +1,228 @@ +```@raw html + +``` + +# Writing + +## Write a table + +[`Arrow.write(sink, table)`](@ref Arrow.write) writes any Tables.jl source (a +`DataFrame`, a `NamedTuple` of vectors, a `CSV.File`, an `Arrow.Table`) to a path or an +`IO` and returns the sink. FilePathsBase.jl local paths (`SystemPath`) work like string +paths, including atomic replacement and `Arrow.append`. Any other object that +`open(sink, "w")` accepts is written in place. `Arrow.write(sink)` returns a +function, so a table can be piped into it, as in +`CSV.File("trips.csv") |> Arrow.write("trips.arrow")`. With FileIO.jl loaded, +`save("trips.arrow", table)` writes a table as well. + +```@example writing +using Arrow, Tables + +dir = mktempdir() +path = joinpath(dir, "trips.arrow") +table = (id=[1, 2, 3], fare=[7.5, 12.0, missing]) +Arrow.write(path, table) +table |> Arrow.write(path) +Arrow.write(path, table; compress=:zstd) +Arrow.Table(path).fare +``` + +A path gets the IPC **file** format: the `ARROW1` magic, the record batches, and a footer +that indexes them, so readers can open any batch directly. An `IO` gets the **stream** +format, which has no footer and can be read as it arrives ([Streams and IPC](@ref)). +`file=true` or `file=false` overrides the choice. [`Arrow.tobuffer(table)`](@ref +Arrow.tobuffer) writes the stream format to a new `IOBuffer`. + +```@example writing +(String(read(path, 6)), String(take!(Arrow.write(IOBuffer(), table; file=true))[1:6])) +``` + +Writing to an existing path writes a temporary file in the same directory and renames it +over the old one at the end. A failed write leaves the old file unchanged, a table that +maps the old file keeps reading it, the file keeps its permissions, a symlink is written +through, and a read-only file is refused. An `IO` is left open. + +The keywords: + +| keyword | default | | +|:--|:--|:--| +| `file` | `true` for a path, `false` for an `IO` | the file format, or the stream format | +| `compress` | `nothing` | `:lz4` (LZ4 frame) or `:zstd` | +| `compresslevel` | `nothing` | the codec's level: LZ4 up to 12 (default 0), Zstandard up to 22 (default 3) | +| `batchsize` | `1_048_576` | rows per record batch | +| `dictencode` | `nothing` | `true`, or a column name or names: dictionary-encode those columns too; `false`: write pooled columns as values | +| `dictdeltas` | `false` | write the categories a dictionary gains as delta dictionaries | +| `largelists` | `false` | 64-bit offsets: LargeUtf8, LargeBinary, LargeList | +| `stringview` | `false` | write strings and bytes as Utf8View and BinaryView | +| `metadata`, `colmetadata` | the table's | schema and column key-value metadata | +| `ntasks` | `Threads.nthreads()` | tasks that build columns and compress buffers | + +[`Arrow.Writer`](@ref) also takes `nullable` and `dictreplacement`. A misspelled keyword is +an `ArgumentError` that lists the valid ones; the Arrow 2.x keywords `alignment`, +`dictencodenested`, `denseunions`, and `maxdepth` are accepted with a warning and ignored. + +## Column types + +Each column's element type decides its Arrow type, at any depth: a `Vector{Date}` inside a +list writes the type a `Date` column writes. [Type mapping](@ref) lists every type in both +directions. + +Strings must be valid UTF-8, as the format requires: a string that is not fails the write +with an error naming its column and row. Write bytes that are not text as `Vector{UInt8}` +values, which write as Binary. + +Utf8, Binary, and List store 32-bit offsets, which address 2 GiB of data in one record +batch. A batch whose strings, bytes, or list values would pass that is cut into smaller +batches, so the declared types never change. `largelists=true` declares 64-bit offsets +instead; a single value larger than 2 GiB needs it. + +## Batches and partitions + +Each `Tables.partitions` partition of the source becomes one or more record batches of at +most `batchsize` rows; a partition with no rows adds none, so a table with no rows writes +its schema alone. `Arrow.write(path, Arrow.Stream(src))` keeps the source's record +batches. Partitions are read, encoded, and written one at a time, so +`Arrow.write("big.arrow", CSV.Chunks("big.csv"))` holds about one chunk in memory. + +The first partition fixes the schema. With several partitions, or a source whose +partition count is unknown, every field is declared nullable, since a later partition +may hold missing values the first did not (pyarrow and Polars write partitioned data the +same way), so such a table reads back with +`Union{Missing,T}` columns. Later partitions must have the same columns in the same +order, with values that convert exactly (an `Int32` partition fits an `Int64` column); +anything else is an error naming the column. + +## Compression + +`compress=:lz4` or `:zstd` compresses each buffer of each record batch and dictionary on +its own, as the format defines; a buffer that does not shrink is stored uncompressed. +pyarrow and Polars read both codecs; DuckDB 1.5's `read_arrow` reads Zstandard but not +LZ4. Zstandard usually writes smaller files; LZ4 is faster to write and read. + +## Dictionary encoding + +A dictionary-encoded column stores each distinct value once and an integer index per row. +Low-cardinality strings, such as a region or a status, are the usual case. A column is +dictionary encoded when + +- it is wrapped in [`Arrow.DictEncode`](@ref); +- it has a `DataAPI.refpool`, as a `PooledArray`, a `CategoricalArray`, and a dictionary + column read from Arrow do: its pool becomes the dictionary, in pool order and with + unused entries, and an ordered `CategoricalArray` writes an ordered dictionary; or +- `dictencode` names it, or is `true`. + +`dictencode=false` writes pooled and categorical columns, and dictionary columns read +from Arrow, as plain values, for readers that refuse dictionaries such as DuckDB 1.5's +`read_arrow`; an `Arrow.DictEncode` column stays encoded. A pooled column of `UUID`s, or +one `dictencode` names, is written unencoded, because pyarrow cannot read a dictionary of +`arrow.uuid` values; an `Arrow.DictEncode` of `UUID`s is an error that names the column. +A dictionary of `UUID`s read from Arrow keeps its encoding unless `dictencode=false`. + +```@example writing +using PooledArrays + +Arrow.write(path, (region=Arrow.DictEncode(["east", "west", "east"]), + status=PooledArray(["open", "closed", "open"]), id=[1, 2, 3])) +t = Arrow.Table(path) +t.region +``` + +The indices are the narrowest signed integers that hold the dictionary (`Int8` up to 128 +values) for a known one-partition table, and `Int32` for several or an unknown number of +partitions and for an `Arrow.Writer`, since later tables may add values. Reading gives a +[`Arrow.DictEncoded`](@ref) column: `DataAPI.refarray`, `DataAPI.refpool`, and +`DataAPI.levels` expose the encoding, and `copy` gives a `PooledArray`. + +### Dictionaries that grow + +When a later partition, or a later table given to an `Arrow.Writer`, has new values, the +column keeps one dictionary: the new values follow the earlier ones. How the growth is +written, and who reads it: + +| output | default | `dictdeltas=true` | +|:--|:--|:--| +| file | the whole dictionary once, after the record batches | a delta dictionary with the new values | +| stream | the whole dictionary again, before the batch that needs it | a delta dictionary with the new values | +| read by | pyarrow, Polars | pyarrow; not Polars 1.43 | + +Arrow 2.x reads all of these, but with several threads it can fail on the default file +output with a `KeyError`, decoding a record batch before the dictionary written after it; +write such files with `dictdeltas=true` when Arrow 2.x readers need them. DuckDB 1.5's +`read_arrow` refuses dictionaries in every form. + +A stream [`Arrow.Writer`](@ref) created with `dictreplacement=true` instead replaces each +dictionary with the next table's own values. The schema then declares the +DictionaryReplacement feature, which some readers (nanoarrow) refuse. + +## Write in parts + +[`Arrow.Writer`](@ref) writes one output from several tables, such as batches that do +not fit in memory together. Each `Arrow.write(writer, table)` encodes that table's record +batches and hands them to the sink before it returns; `close` finishes the output, and the +function form closes it. If the function throws, the writer is abandoned instead, so an +existing file is kept, unless the sink is an `IO` that already received a table: then the +function form closes it, and the `IO` holds the tables written before the error. A writer +to a path writes nothing to that path until `close`: an existing file is replaced then, +from a temporary file beside it, and a new file holds no complete output before. + +```@example writing +Arrow.Writer(path) do writer + Arrow.write(writer, (id=[1, 2], name=["a", "b"])) + Arrow.write(writer, (id=Int32[3], name=[missing])) +end +Arrow.Table(path).name +``` + +The first table fixes the column names, order, and types. Every column is declared +nullable unless the writer was created with `nullable=false`, so later tables may hold +missing values. Later tables convert to the first table's types exactly: an `Int32` +column fits an `Int64` one, but a `Float64` column does not. A table that does not conform is +an error naming the column and both types; the batches written before it stay in the +output and the writer stays open. A writer closed before it receives a table writes an +output with no columns. A writer that is never closed leaves an `IO` without its footer or +end marker, and leaves a path as it was, beside a stray temporary file (`jl_` and random +letters) or a new file without its footer or end marker. + +[`Arrow.append(path, table)`](@ref Arrow.append) adds a table's record batches to an +existing stream-format file (or a seekable `IO`), under the same rules; a missing or empty +file starts a new stream. The file format cannot be appended to, since its footer would +have to be rewritten: produce it with an `Arrow.Writer`, or rewrite it with +`Arrow.write`. + +## Metadata + +`metadata` sets the schema's key-value pairs and `colmetadata` each column's, as a `Dict` +or a vector of pairs; `colmetadata` maps column names, as `Symbol`s or strings, to them. +Values that are not strings are written as they print. Both default to the +table's DataAPI metadata, so a `DataFrame`'s `metadata` and `colmetadata` are written; +keys under the reserved `ARROW:` prefix are not. Writing an `Arrow.Table` or +`Arrow.Stream` keeps its schema and column metadata. + +## Memory and threads + +A write holds the partition it is writing, the record batch it is building and that +batch's encoded buffers, and, in the file format, each dictionary that grows until it is +written at the end. Columns that are already Arrow layouts, such as the views of an +`Arrow.Table`, are written from their own memory without a copy. + +The columns of a record batch are built and its buffers compressed on up to `ntasks` +tasks (default `Threads.nthreads()`); with compression, a batch compresses while the next +one builds. `ntasks=1` writes on the calling task. The bytes written are the same for +every `ntasks`. diff --git a/ext/ArrowCategoricalArraysExt.jl b/ext/ArrowCategoricalArraysExt.jl new file mode 100644 index 00000000..a313f15d --- /dev/null +++ b/ext/ArrowCategoricalArraysExt.jl @@ -0,0 +1,29 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# An ordered CategoricalArray writes an ordered Arrow dictionary, and +# `isordered` reads a dictionary column's flag. +module ArrowCategoricalArraysExt + +import Arrow +import CategoricalArrays + +Arrow._isordered(x::CategoricalArrays.CategoricalArray) = CategoricalArrays.isordered(x) +CategoricalArrays.isordered(x::Arrow.DictEncoded) = Arrow._isordered(x) +CategoricalArrays.isordered(x::Arrow.ChainedVector{T,<:Arrow.DictEncoded}) where {T} = + Arrow._isordered(x) + +end 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/ArrowDataFramesExt.jl b/ext/ArrowDataFramesExt.jl new file mode 100644 index 00000000..d296dce8 --- /dev/null +++ b/ext/ArrowDataFramesExt.jl @@ -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. + +# Precompiles the reads and writes of a DataFrame, so the first one in a session +# does not compile them. It adds no methods. +module ArrowDataFramesExt + +import Arrow +import DataFrames +import Dates +import PooledArrays +import PrecompileTools + +PrecompileTools.@setup_workload begin + table = ( + id=Int64[1, 2, 3], + x=Float64[1.5, 2.5, 3.5], + y=Union{Missing,Float64}[1.5, missing, 3.5], + name=["a", "a longer string than twelve bytes", "c"], + note=Union{Missing,String}["a", missing, "c"], + day=[Dates.Date(2025, 1, 1), Dates.Date(2025, 1, 2), Dates.Date(2025, 1, 3)], + at=[ + Dates.DateTime(2025, 1, 1), + Dates.DateTime(2025, 1, 2), + Dates.DateTime(2025, 1, 3), + ], + flag=[true, false, true], + pool=PooledArrays.PooledArray(["u", "v", "u"]), + ) + PrecompileTools.@compile_workload begin + bytes = take!(Arrow.write(IOBuffer(), table; file=true)) + arrowtable = Arrow.Table(bytes) + df = DataFrames.DataFrame(arrowtable) + DataFrames.DataFrame(arrowtable; copycols=false) + Arrow.read(bytes, DataFrames.DataFrame) + take!(Arrow.write(IOBuffer(), df; file=true)) + take!(Arrow.write(IOBuffer(), DataFrames.DataFrame(table); file=true)) + end +end + +end # module ArrowDataFramesExt diff --git a/ext/ArrowFileIOExt.jl b/ext/ArrowFileIOExt.jl new file mode 100644 index 00000000..0bf99732 --- /dev/null +++ b/ext/ArrowFileIOExt.jl @@ -0,0 +1,38 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# FileIO.jl registers `.arrow` files (magic `ARROW1`) as `format"Arrow"` and +# hands them to this package: `load(path)` reads an `Arrow.Table` and +# `save(path, table)` writes any Tables.jl table, with `Arrow.Table`'s and +# `Arrow.write`'s keywords. +module ArrowFileIOExt + +import Arrow +using FileIO: File, Stream, @format_str, filename, stream + +Arrow.fileio_load(f::File{format"Arrow"}; kwargs...) = Arrow.Table(filename(f); kwargs...) +Arrow.fileio_load(s::Stream{format"Arrow"}; kwargs...) = Arrow.Table(stream(s); kwargs...) + +function Arrow.fileio_save(f::File{format"Arrow"}, table; kwargs...) + Arrow.write(filename(f), table; kwargs...) + return nothing +end +function Arrow.fileio_save(s::Stream{format"Arrow"}, table; kwargs...) + Arrow.write(stream(s), table; kwargs...) + return nothing +end + +end # module ArrowFileIOExt diff --git a/ext/ArrowFilePathsBaseExt.jl b/ext/ArrowFilePathsBaseExt.jl new file mode 100644 index 00000000..26ce7102 --- /dev/null +++ b/ext/ArrowFilePathsBaseExt.jl @@ -0,0 +1,33 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Local paths share the string-path source and atomic destination behavior. +module ArrowFilePathsBaseExt + +import Arrow +using FilePathsBase: SystemPath + +Arrow.Table(path::SystemPath; kwargs...) = Arrow.Table(string(path); kwargs...) +Arrow.Stream(path::SystemPath; kwargs...) = Arrow.Stream(string(path); kwargs...) +Arrow._opensink(path::SystemPath) = Arrow._openpath(string(path)) + +function Arrow.append(path::SystemPath, table; kwargs...) + Arrow.append(string(path), table; kwargs...) + return path +end +Arrow.append(path::SystemPath; kwargs...) = table -> Arrow.append(path, table; kwargs...) + +end diff --git a/ext/ArrowSerializationExt.jl b/ext/ArrowSerializationExt.jl new file mode 100644 index 00000000..3f8d749f --- /dev/null +++ b/ext/ArrowSerializationExt.jl @@ -0,0 +1,85 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +module ArrowSerializationExt + +import Arrow +import Serialization +const AC = Arrow.ArrowCore +const S = Serialization + +# Write only the physical values. ForeignOwner, OwnerRegion, ReleaseCell, +# raw addresses, and semantic certificates must never cross this boundary. +function S.serialize(s::S.AbstractSerializer, d::AC.ArrayData) + S.serialize_cycle_header(s, d) && return nothing + S.serialize(s, d.type) + S.serialize(s, d.len) + S.serialize(s, d.offset) + S.serialize(s, [AC.slicebytes(b) for b in d.buffers]) + S.serialize(s, d.children) + S.serialize(s, d.dictionary) + S.serialize(s, @atomic(:monotonic, d.nullcount)) + return nothing +end + +function S.deserialize(s::S.AbstractSerializer, ::Type{AC.ArrayData}) + type = S.deserialize(s) + len = S.deserialize(s) + offset = S.deserialize(s) + buffers = S.deserialize(s) + children = S.deserialize(s) + dictionary = S.deserialize(s) + nullcount = S.deserialize(s) + d = AC.ArrayData( + type, + len, + [AC._databuffer(b) for b in buffers]; + offset=offset, + children=children, + dictionary=dictionary, + nullcount=nullcount, + ) + # ArrayData's immutable topology cannot contain a cycle. Register after + # its children are read to preserve repeated references to the same node. + S.deserialize_cycle(s, d) + return d +end + +function S.serialize(s::S.AbstractSerializer, t::Arrow.Table) + S.serialize_type(s, Arrow.Table) + S.serialize(s, getfield(t, :names)) + # Views serialize as the owned vectors they copy to, not as their buffers. + S.serialize( + s, + AbstractVector[Arrow._ownedcolumn(Arrow._force(c)) for c in getfield(t, :columns)], + ) + for name in (:schema, :nrows, :retainedpools) + S.serialize(s, getfield(t, name)) + end + return nothing +end + +function S.deserialize(s::S.AbstractSerializer, ::Type{Arrow.Table}) + names = S.deserialize(s) + columns = S.deserialize(s) + schema = S.deserialize(s) + nrows = S.deserialize(s) + pools = S.deserialize(s) + return Arrow._table(names, columns, schema, AC.OwnerRegion[], nrows, pools) +end + +end # module ArrowSerializationExt diff --git a/ext/ArrowTimeZonesExt.jl b/ext/ArrowTimeZonesExt.jl new file mode 100644 index 00000000..fdac9c45 --- /dev/null +++ b/ext/ArrowTimeZonesExt.jl @@ -0,0 +1,91 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +TimeZones.jl conveniences for timezone-aware timestamp columns. Reading +never needs this extension: a zone-declared timestamp column materializes as +`Durations.ZonedTimestamp` at every unit (Durations' own TimeZones extension +adds named-zone rules and `ZonedDateTime` conversions to those values). +Loading TimeZones.jl adds two things here: + + * Writing: a fresh `TimeZones.ZonedDateTime` column writes as a + timezone-declared millisecond timestamp, matching Arrow 2.x, instead of + falling through to reflected-struct lowering. + * Retained writes: a `ZonedDateTime` lowers exactly to a zone-declared + column's storage domain (the UTC instant in the column unit). +""" +module ArrowTimeZonesExt + +import Arrow +import Dates +import TimeZones + +const AC = Arrow.ArrowCore + +# Lower one ZonedDateTime to a zone-declared timestamp's storage domain +# exactly: the UTC instant in the descriptor's unit, or `nothing` when the +# value is not a ZonedDateTime or the unit cannot hold the instant exactly. +# The declared zone deliberately plays no role: storage is the UTC instant +# and zoned comparisons use only it, so 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 && return ms % 1000 == 0 ? ms ÷ 1000 : nothing + scale = unit == AC.MICROSECOND ? Int64(1_000) : Int64(1_000_000) + wide = Int128(ms) * Int128(scale) + return Int128(typemin(Int64)) <= wide <= Int128(typemax(Int64)) ? Int64(wide) : nothing +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..3f71647e 100644 --- a/src/Arrow.jl +++ b/src/Arrow.jl @@ -14,131 +14,162 @@ # See the License for the specific language governing permissions and # limitations under the License. +# Layering (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`, `cstream.jl`: the C data and C stream interfaces over Core, +# import and export, with lifecycle accounting; `cdata_table.jl` puts +# `Arrow.Table`, `Arrow.Stream`, and any Tables.jl table across them. +# - `scan.jl`: `Tables.Scan` pushdown and byte-range reads. +# - `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. """ - Arrow.jl - -A pure Julia implementation of the [apache arrow](https://arrow.apache.org/) memory format specification. - -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 + Arrow -It currently doesn't include support for: - * Tensors or sparse tensors - * Flight RPC - * C data interface - -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 - -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). +Read and write the Apache Arrow IPC file and stream formats. [`Arrow.Table`](@ref) reads +a path, an `IO`, a byte vector, or an [`Arrow.AbstractArrowSource`](@ref) as Tables.jl +columns; [`Arrow.Stream`](@ref) reads one record batch at a time; [`Arrow.write`](@ref) +and [`Arrow.Writer`](@ref) write any Tables.jl table. The names are public but not +exported, apart from `ArrowTypes`: call them as `Arrow.Table` and so on. """ module Arrow -using Base.Iterators -using Mmap +using Tables +import DataAPI +import DataStrings, DataDecimals, Durations +import PooledArrays +using SentinelArrays: ChainedVector, ChainedVectorIndex +import ArrowTypes import Dates -using DataAPI, - Tables, - SentinelArrays, - PooledArrays, - CodecLz4, - CodecZstd, - TimeZones, - BitIntegers, - ConcurrentUtilities, - StringViews - -export ArrowTypes - -using Base: @propagate_inbounds -import Base: == - -const FILE_FORMAT_MAGIC_BYTES = b"ARROW1" -const CONTINUATION_INDICATOR_BYTES = 0xffffffff - -# vendored flatbuffers code for now -include("FlatBuffers/FlatBuffers.jl") -using .FlatBuffers - -include("metadata/Flatbuf.jl") -using .Flatbuf -const Meta = Flatbuf +import UUIDs +import CodecLz4 +import CodecZstd +import TranscodingStreams +using CodecLz4: LZ4FrameCompressor + +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 -using ArrowTypes -include("utils.jl") -include("arraytypes/arraytypes.jl") -include("eltypes.jl") +include("ArrowCore.jl") +using .ArrowCore +import .ArrowCore: release! +const AC = ArrowCore + +include("ipc_read.jl") +include("ipc_dictionary.jl") +include("ipc_write.jl") +include("cdata.jl") +include("cstream.jl") +include("source.jl") +include("arrowtypes.jl") +include("columns.jl") + +# 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") -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 +include("cdata_table.jl") +include("sharedvalues.jl") -function init_zstd_decompressor() - zstd = ZstdDecompressor() - CodecZstd.TranscodingStreams.initialize(zstd) - return Lockable(zstd) -end +""" + Arrow.ValidationError -function init_lz4_frame_compressor() - lz4 = LZ4FrameCompressor(; compressionlevel=4) - CodecLz4.TranscodingStreams.initialize(lz4) - return Lockable(lz4) +Thrown when input is not valid Arrow data, or uses something this package does not +read: corrupt, truncated, or inconsistent metadata or buffers, bytes that are not Arrow +at all, or big-endian data. The message names the problem. +""" +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, + :read, + :release!, + :col, + :colcmp, + :colin, + :isnull, + :ArrowVector, + :Primitive, + :BoolVector, + :ListVector, + :ListValue, + :MapVector, + :MapValue, + :StructVector, + :DictEncoded, + :RunEndEncoded, + :write, + :Writer, + :append, + :tobuffer, + :getmetadata, + :DictEncode, + :AbstractArrowSource, + :sourcelength, + :readrange, + :concurrentreads, + :CArrowSchema, + :CArrowArray, + :CArrowArrayStream, + :to_c_data, + :from_c_data, + :export_stream!, + :Limits, + :AllocationLimitError, + :ValidationError, + ), + ) end -function init_lz4_frame_decompressor() - lz4 = LZ4FrameDecompressor() - CodecLz4.TranscodingStreams.initialize(lz4) - return Lockable(lz4) -end +# 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 -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 -end +# FileIO.jl's `load("x.arrow")` and `save("x.arrow", table)` call these; the +# FileIO extension defines their methods. +function fileio_load end +function fileio_save end + +# Precompile the common write/read/scan shapes last, once every layer above +# is defined. +include("precompile.jl") -end # module Arrow +end # module Arrow diff --git a/src/ArrowCore.jl b/src/ArrowCore.jl new file mode 100644 index 00000000..5ff8c16c --- /dev/null +++ b/src/ArrowCore.jl @@ -0,0 +1,4012 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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) + +# 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. (Comments, not docstrings: `?Arrow.release!` lists only the +# public methods.) +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 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 +# (for hosts where a GC-timed unmap is not enough, such as deleting a mapped +# file on Windows); 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 + +# A deep copy shares the region and its cell: the bytes belong to the root and +# are freed through the cell once, so a copy with its own cell would read them +# after they are released and release them again. +Base.deepcopy_internal(r::OwnerRegion, ::IdDict) = r +Base.deepcopy_internal(c::ReleaseCell, ::IdDict) = c + +""" + heapregion(v::Vector{T}) -> OwnerRegion + +Borrow a Julia array as a region (zero-copy). The array's storage is the `root`, +so the bytes stay valid while the region is reachable even if `v` is resized +afterwards. On Julia 1.10 the array itself is the root, and resizing it +reallocates the bytes the region points at. +""" +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=_storageroot(v)) +end + +# `_storageroot(v)`: the object that owns a vector's bytes, and the one Mmap +# registers its unmap finalizer on — the array's backing `Memory` from Julia +# 1.11 (`finalize(arr)` is then a no-op), the array itself before. A release +# cell targets that object so `release!` truly unmaps now. +@static if VERSION >= v"1.11" + _storageroot(v::Vector) = v.ref.mem + function _release_mmap(p::Ptr{Cvoid})::Cvoid + finalize(unsafe_pointer_to_objref(p)::Memory{UInt8}) + return nothing + end +else + _storageroot(v::Vector) = v + 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. With +# `eager=false`, `release!` only revokes the region and the mapping is unmapped +# when its root is garbage collected: the form for memory that unowned views +# still read after a release. +function _mmapregion(io::IO, label; eager::Bool=true) + arr = Mmap.mmap(io, Vector{UInt8}) + isempty(arr) && throw(ArgumentError("cannot map empty file: $label")) + cell = + eager ? + ReleaseCell(@cfunction(_release_mmap, Cvoid, (Ptr{Cvoid},)), _storageroot(arr)) : + ReleaseCell() + return OwnerRegion(Ptr{UInt8}(pointer(arr)), length(arr); root=arr, cell=cell) +end + +"Unmap a mapped region's memory now. Nothing may read the region afterward." +_unmap!(r::OwnerRegion) = finalize(_storageroot(r.root::Vector{UInt8})) + +""" + 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 per the spec; 128 is Polars' Int128/UInt128 extension + 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 + +""" + copy(data::ArrayData) -> ArrayData + +Detach all buffer slices, children, and dictionary values into Julia-owned +byte arrays. Preserve the layout and element offset. The result has no +foreign owner or release callback and remains usable after the source is +released. Keep its `Field` to materialize named values with `materialize`. + +The source must remain valid and unchanged during the copy. This copies +whole declared buffer slices, which can include values outside a logical +array slice. Semantic validation is not inherited. `deepcopy` does not +detach: the copy shares the source's buffers and their release. +""" +function Base.copy(d::ArrayData) + return ArrayData( + d.type, + d.len, + [_databuffer(slicebytes(b)) for b in d.buffers]; + offset=d.offset, + children=map(copy, d.children), + dictionary=d.dictionary === nothing ? nothing : copy(d.dictionary), + nullcount=@atomic(:monotonic, d.nullcount), + ) +end + +# 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) + return d.len - _countsetbits(v, d.offset, d.len) +end + +"Set bits among the `len` bitmap bits that start at 0-based bit `start`." +function _countsetbits(v::BufferSlice, start::Int64, len::Int64) + len > 0 || return Int64(0) + start >= 0 || throw(BoundsError(v, start)) + stop = checked_add(start, len) # exclusive + firstbyte, lastbyte = start >> 3, (stop - 1) >> 3 + nbytes = lastbyte - firstbyte + 1 + # One bounds check covers every byte the bits touch, as getbit's would. + (nbytes <= v.len && firstbyte <= v.len - nbytes) || throw(BoundsError(v, lastbyte)) + n = Int64(0) + GC.@preserve v begin + p = sliceptr(v) + firstbyte + for k = 0:(nbytes - 1) + byte = unsafe_load(p + k) + k == 0 && (byte &= 0xff << (start & 7)) + k == nbytes - 1 && (byte &= 0xff >> (7 - ((stop - 1) & 7))) + n += count_ones(byte) + end + end + return n +end + +# --------------------------------------------------------------------------- +# §5 Staged validation +# --------------------------------------------------------------------------- + +struct ValidationError <: Exception + msg::String +end +Base.showerror(io::IO, e::ValidationError) = print(io, "Arrow.ValidationError: ", e.msg) + +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, 128) || + throw(ValidationError("integer bit width must be 8, 16, 32, 64, or 128")) +_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) + t.indextype.bits <= 64 || + throw(ValidationError("dictionary index bit width must be 8, 16, 32, or 64")) + # 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_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", + ), + ) + 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")) + 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 + _load_offset(offs, wide, d.offset) >= 0 || + throw(ValidationError("negative first offset")) + last = _monotonicoffsets(offs, wide, d.offset, d.len) + last <= databytes || throw( + ValidationError("final offset $last exceeds data extent $databytes"), + ) + end + end + if t isa DictionaryType + dictlen = length(d.dictionary) + data = rolebuffer(d, DATA) + w = primwidth(t.indextype) + if !_dictindicesinrange(d, t.indextype, dictlen) + 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 + 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 + +""" +Whether every valid slot's dictionary index lies in `[0, dictlen)`: one +bounds check for the index window, then a branch-free pass. `false` sends +the caller to the per-slot loop that names the offending index. +""" +function _dictindicesinrange(d::ArrayData, it::IntType, dictlen::Int64) + d.len == 0 && return true + w = Int64(primwidth(it)) + src = subslice(rolebuffer(d, DATA), checked_mul(d.offset, w), checked_mul(d.len, w)) + validity = validitybuffer(d) + ok = GC.@preserve d begin + p = sliceptr(src) + if it.signed + w == 1 ? _indicesinrange(Ptr{Int8}(p), d, validity, dictlen) : + w == 2 ? _indicesinrange(Ptr{Int16}(p), d, validity, dictlen) : + w == 4 ? _indicesinrange(Ptr{Int32}(p), d, validity, dictlen) : + _indicesinrange(Ptr{Int64}(p), d, validity, dictlen) + else + w == 1 ? _indicesinrange(Ptr{UInt8}(p), d, validity, dictlen) : + w == 2 ? _indicesinrange(Ptr{UInt16}(p), d, validity, dictlen) : + w == 4 ? _indicesinrange(Ptr{UInt32}(p), d, validity, dictlen) : + _indicesinrange(Ptr{UInt64}(p), d, validity, dictlen) + end + end + return ok +end + +function _indicesinrange( + p::Ptr{T}, + d::ArrayData, + validity::BufferSlice, + dictlen::Int64, +) where {T} + ok = true + if isempty_buffer(validity) + for i = 1:(d.len) + # A UInt64 above typemax(Int64) reinterprets negative and fails. + x = + T === UInt64 ? reinterpret(Int64, unsafe_load(p, i)) : + Int64(unsafe_load(p, i)) + ok &= (0 <= x) & (x < dictlen) + end + else + for i = 1:(d.len) + x = + T === UInt64 ? reinterpret(Int64, unsafe_load(p, i)) : + Int64(unsafe_load(p, i)) + ok &= !getbit(validity, _slotindex0(d, Int64(i))) | ((0 <= x) & (x < dictlen)) + end + end + return ok +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) + d.len == 0 && return nothing + views = rolebuffer(d, VIEWS) + # One bounds check for every entry, then raw loads per row. + entries = subslice(views, _viewbase(d, Int64(1)), checked_mul(d.len, Int64(16))) + validity = validitybuffer(d) + GC.@preserve d begin + p = sliceptr(entries) + for i = 1:(d.len) + (isempty_buffer(validity) || getbit(validity, _slotindex0(d, Int64(i)))) || + continue + e = p + 16 * (i - 1) + len = unsafe_load(Ptr{Int32}(e)) + len >= 0 || throw(ValidationError("negative view length $len")) + len <= VIEW_INLINE_MAX && continue + bufidx = unsafe_load(Ptr{Int32}(e + 8)) + off = Int64(unsafe_load(Ptr{Int32}(e + 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. + off + Int64(len) <= data.len || throw( + ValidationError( + "view range [$off, $(off + len)) escapes data buffer $bufidx", + ), + ) + unsafe_load(Ptr{UInt32}(e + 4)) == + unsafe_load(Ptr{UInt32}(sliceptr(data) + off)) || + 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 : t.bits == 128 ? Int128 : Int64 + ) : + ( + t.bits == 8 ? UInt8 : + t.bits == 16 ? UInt16 : t.bits == 32 ? UInt32 : t.bits == 128 ? UInt128 : 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 + +# 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)) + t.bits == 128 && return loadat(b, Int128, _slotbyteoff(d, i, 16)) + 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)) + t.bits == 128 && return loadat(b, UInt128, _slotbyteoff(d, i, 16)) + 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 + +""" +The last of the `n + 1` offsets that start at 0-based entry `start0` of +`offs`, after checking that they never decrease. One bounds check covers the +window, and the check runs branch-free so it vectorizes. +""" +function _monotonicoffsets(offs::BufferSlice, wide::Bool, start0::Int64, n::Int64)::Int64 + w = wide ? Int64(8) : Int64(4) + first = checked_mul(start0, w) + span = checked_mul(checked_add(n, Int64(1)), w) + (first >= 0 && span <= offs.len && first <= offs.len - span) || + throw(BoundsError(offs, checked_add(start0, n))) + last = GC.@preserve offs begin + p = sliceptr(offs) + first + wide ? _monotonic(Ptr{Int64}(p), n) : _monotonic(Ptr{Int32}(p), n) + end + return last +end + +function _monotonic(p::Ptr{T}, n::Int64)::Int64 where {T} + ok = true + for i = 1:n + ok &= unsafe_load(p, i) <= unsafe_load(p, i + 1) + end + if !ok + for i = 1:n + unsafe_load(p, i) <= unsafe_load(p, i + 1) || + throw(ValidationError("offsets not monotonically non-decreasing at $i")) + end + end + return Int64(unsafe_load(p, n + 1)) +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") + # Fixed-width values charge nothing beyond the column under a typed + # claim (only `Any` boxes them), so they skip the per-value walk. + fixedwidth = + d.type isa Union{ + IntType, + FloatType, + TimestampType, + DurationType, + DateType, + TimeType, + IntervalType, + BoolType, + NullType, + } + if !fixedwidth + for i = 1:(d.len) + _prechargevalue_of!(T, d.type, f, d, Int64(i), estimate) + end + 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)) + # 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: one byte copy builds the + # public vector. + if !(Missing <: T) + 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 + isempty_buffer(v) || _countsetbits(v, d.offset, n) == n || _typednullrefuse(f) + return vals + end + # A Missing-admitting vector has no contiguous byte view to copy into, so + # load each value straight from the buffer (unsafe_load is unaligned-safe) + # instead of staging a Vector{E} copy first. + out = Vector{T}(undef, n) + if n > 0 + GC.@preserve d begin + p = Ptr{E}(sliceptr(src)) + @inbounds for i = 1:n + out[i] = unsafe_load(p, i) + end + end + end + isempty_buffer(v) || _punchnulls!(out, v, d.offset, n) + return out +end + +# Set out[i] = missing for each cleared bit among the `n` validity bits that +# start at 0-based bit `start`, under one bounds check for the bitmap bytes. +function _punchnulls!(out::Vector, v::BufferSlice, start::Int64, n::Int64) + n > 0 || return out + start >= 0 || throw(BoundsError(v, start)) + stop = checked_add(start, n) + firstbyte, lastbyte = start >> 3, (stop - 1) >> 3 + nbytes = lastbyte - firstbyte + 1 + (nbytes <= v.len && firstbyte <= v.len - nbytes) || throw(BoundsError(v, lastbyte)) + GC.@preserve v begin + p = sliceptr(v) + for i = 1:n + bit = start + (i - 1) + (unsafe_load(p + (bit >> 3)) >> (bit & 7)) & 0x01 == 0x00 && (out[i] = missing) + end + 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) + # One pass fills the values and the validity bitmap. + n = length(v) + vals = Vector{S}(undef, n) + bits = zeros(UInt8, expected_validity_bytes(Int64(n))) + nc = 0 + for i = 1:n + x = v[i] + if x === missing + nc += 1 + vals[i] = zero(S) + else + vals[i] = S(x) + _setbitmapbit!(bits, i) + end + end + validity = nc == 0 ? BufferSlice() : BufferSlice(heapregion(bits), 0, length(bits)) + data = _databuffer(vals) + # 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 DataStrings' `DataString` 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 +"Copy a record batch with detached column buffers and the same immutable schema." +Base.copy(b::RecordBatch) = RecordBatch(b.schema, map(copy, b.columns), b.nrows) + +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/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..ada72de8 --- /dev/null +++ b/src/arrowtypes.jl @@ -0,0 +1,1756 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT 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) + # Checked as a `String`: other packages add methods to the iteration of + # other string types. + name isa String || return _existingjlsymbol(String(name)::String) + 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 + +# ArrowTypes.jl hooks take FixedSizeBinary storage as `NTuple{N,UInt8}` (for +# example `fromarrow(::Type{UUID}, ::NTuple{16,UInt8})`). Wider descriptors keep +# their bytes as a vector so an untrusted width cannot force a huge tuple type. +_arrowtypesfixedbinarytype(t::AC.FixedSizeBinaryType) = + t.nbytes <= _MAX_ARROWTYPE_EXACT_ARITY ? NTuple{t.nbytes,UInt8} : Vector{UInt8} + +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 + +# --------------------------------------------------------------------------- +# Mappings the facade owns, consulted before ArrowTypes.jl's hooks: canonical +# Arrow extension types, and Julia types ArrowTypes.jl gives no Arrow +# storage. `UUID` writes canonical `arrow.uuid` (big-endian bytes); the +# little-endian `JuliaLang.UUID` layout older writers used still reads and +# rewrites through ArrowTypes.jl's own UUID hooks. +# --------------------------------------------------------------------------- + +const _UUID_EXTENSION = "arrow.uuid" +const _FACADE_PERIOD_NAMES = ( + (Dates.Week, "JuliaLang.Dates.Week"), + (Dates.Day, "JuliaLang.Dates.Day"), + (Dates.Hour, "JuliaLang.Dates.Hour"), + (Dates.Minute, "JuliaLang.Dates.Minute"), + (Dates.Month, "JuliaLang.Dates.Month"), + (Dates.Year, "JuliaLang.Dates.Year"), +) +# Canonical extensions whose storage is already their public value. +_quietextension(name::String) = name == "arrow.json" || name == "arrow.opaque" + +function _facadearrowname(T::Type) + T === UUIDs.UUID && return _UUID_EXTENSION + T === Int128 && return "JuliaLang.Int128" + T === UInt128 && return "JuliaLang.UInt128" + for (P, name) in _FACADE_PERIOD_NAMES + T === P && return name + end + return nothing +end + +_facadestoragetype(T::Type) = T === Int128 || T === UInt128 ? NTuple{16,UInt8} : nothing + +# The Julia type a facade-owned label reads as, given its ArrowTypes storage +# type, or nothing. +function _facadejuliatype(name::String, storage) + bytes16 = storage === NTuple{16,UInt8} + name == _UUID_EXTENSION && return bytes16 ? UUIDs.UUID : nothing + name == "JuliaLang.Int128" && return bytes16 ? Int128 : nothing + name == "JuliaLang.UInt128" && return bytes16 ? UInt128 : nothing + name == "arrow.bool8" && return storage === Int8 ? Bool : nothing + for (P, label) in _FACADE_PERIOD_NAMES + name == label || continue + P <: Union{Dates.Month,Dates.Year} && return storage === Int32 ? P : nothing + return storage isa Type && storage <: Dates.FixedPeriod ? P : nothing + end + return nothing +end + +# Storage for a facade-owned value. `field` is the Field being rewritten, if +# any: a retained `JuliaLang.UUID` field keeps its little-endian layout. +function _facadetoarrow(x::UUIDs.UUID, field) + legacy = + field !== nothing && _fieldmetavalue(field, _EXTENSION_NAME_KEY) == "JuliaLang.UUID" + legacy && return ArrowTypes.toarrow(x) + return ArrowTypes._cast(NTuple{16,UInt8}, hton(x.value)) +end +_facadetoarrow(x::Union{Int128,UInt128}, field) = ArrowTypes._cast(NTuple{16,UInt8}, x) + +# The public value of facade-owned storage under `f`'s label, or nothing when +# the label is not one the facade owns (ArrowTypes.jl's hooks then apply). +function _facadefromarrow(T::Type, f::AC.Field, storage) + name = _fieldmetavalue(f, _EXTENSION_NAME_KEY) + (name === nothing || _facadejuliatype(name, typeof(storage)) !== T) && return nothing + # `storage`'s type is known at run time; see `_uninferred`. + T === UUIDs.UUID && + return UUIDs.UUID(ntoh(_uninferred(ArrowTypes._cast, UInt128, storage)::UInt128)) + T === Bool && return storage != 0 + T <: Union{Int128,UInt128} && return _uninferred(ArrowTypes._cast, T, storage) + T <: Union{Dates.Month,Dates.Year} && return convert(T, Dates.Month(storage)) + return convert(T, storage) +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 _arrowtypesfixedbinarytype(t) + 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, +) + # Measured as a `String`: other packages add methods to the iteration of + # other string types. + name isa String || return _extensionwarningname(String(name)::String, budget) + 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 + +# Labels that Arrow.jl 2.x (and Durations.jl's Arrow 2.x extension) wrote on +# temporal columns whose storage Arrow 3.0 reads natively into the same +# instants and values. They are expected in 2.x files, so they do not warn. +# The pre-2.4.2 `JuliaLang.ZonedDateTime` label is not one of them: that +# storage holds local wall time, not a UTC instant. +function _arrow2temporallabel(name::String, t::AC.ArrowType) + t isa AC.DictionaryType && (t = t.valuetype) + naive = t isa AC.TimestampType && _timestampzone(t) === nothing + zoned = t isa AC.TimestampType && _timestampzone(t) !== nothing + name == "JuliaLang.Date" && return t isa AC.DateType + name == "JuliaLang.DateTime" && + return naive || (t isa AC.DateType && t.unit == AC.MILLISECOND) + name == "JuliaLang.Time" && return t isa AC.TimeType + name == "JuliaLang.Dates.Period" && return t isa AC.DurationType + name == "JuliaLang.ZonedDateTime-UTC" && return zoned + name == "JuliaLang.Durations.Timestamp" && return naive + name == "JuliaLang.Durations.ZonedTimestamp" && return zoned + return false +end + +# Unsupported extension names already reported, process-wide, so reading +# many files that carry the same label warns once. +const _WARNED_EXTENSIONS = Set{String}() +const _WARNED_EXTENSIONS_LOCK = ReentrantLock() + +function _warnunsupportedextension(ctx, name::String, f::AC.Field) + ctx.warn || return nothing + _arrow2temporallabel(name, f.type) && return nothing + displayname = _extensionwarningname(name, ctx.budget) + report = lock(_WARNED_EXTENSIONS_LOCK) do + displayname in _WARNED_EXTENSIONS && return :seen + length(_WARNED_EXTENSIONS) > _MAX_EXTENSION_WARNINGS && return :seen + push!(_WARNED_EXTENSIONS, displayname) + return length(_WARNED_EXTENSIONS) > _MAX_EXTENSION_WARNINGS ? :suppress : :warn + end + report === :seen && return nothing + if report === :suppress + @warn "additional unsupported Arrow extension warnings were suppressed" + return nothing + end + 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) + # The routed walker asks per value: a cache hit must not construct the + # memoization thunk, or one heap closure rides beside every row. + cached = get(plan.extensions, f, nothing) + cached === nothing || return cached + return _memoized!( + plan.extensions, + f, + plan.budget, + "ArrowTypes extension-presence cache", + ) do + _arrowtypesextension(f) === nothing || return true + return _uninferred( + any, + child -> _hasarrowtypesextension(child, plan), + f.children, + )::Bool + end +end + +function _routedictionaryfield!( + plan::_ArrowTypesRoutePlan, + f::AC.Field, + t::AC.DictionaryType, +) + # Per dictionary value in the routed walker: hits must not allocate. + cached = get(plan.dictionaryfields, f, nothing) + cached === nothing || return cached + 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 _uninferred( + any, + child -> _needsarrowtypesroute(child, plan), + f.children, + )::Bool + end +end + +function _arrowtypesroutedvalue( + f::AC.Field, + d::AC.ArrayData, + i::Int64, + plan::_ArrowTypesRoutePlan, + routeallunions::Bool=false, + budget::Union{Nothing,AllocationBudget}=nothing, + underlabel::Bool=false, +) + # A label-free subtree that is NOT below any label lands in PUBLIC rows, + # so build it with the facade's public row builder. Everything at or + # below a label — registered or not — stays in the ArrowTypes STORAGE + # domain and keeps walking, so leaves reach the raw tail and unions + # carry routed markers; restoration then either feeds a registered + # target's fromarrow or converts the raw values itself when the label + # resolves to no target. + (routeallunions || underlabel || _hasarrowtypesextension(f, plan)) || + return _publicvalue(f, d, i, budget, true) + under = underlabel || _arrowtypesextension(f) !== nothing + 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, + under, + ) + # 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, + under, + ) + 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, + under, + ) + 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, + under, + ) + 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, + under, + ) + 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, + under, + ) + 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, + under, + ) + 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, + under, + ), + _arrowtypesroutedvalue( + valuefield, + valuedata, + entryindex, + plan, + routeallunions, + budget, + under, + ), + ) + end + return out + end + # A null slot must not run the budgeted extraction: its transient + # charging machinery would allocate per missing value with nothing to + # reserve it against. The Null layout has no validity buffer (asking + # for it throws) and every slot is missing. + t isa AC.NullType && return missing + AC.isvalid_at(d, i) || return missing + budget === nothing && return AC.getvalue(f, d, i) + # Core's budgeted getvalue reserves the materialized value but its own + # transient estimate object rides beside every leaf; flat wrapper + # columns have no container reserve to absorb it, so charge it here. + _chargeobject!(budget, 24, "ArrowTypes routed leaf estimate") + return 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; + force::Bool=false, +) + # Consult the route cache even when forced: the plan memoizes and + # budget-charges one route entry per schema Field either way. + needed = _needsarrowtypesroute(f, plan) + (force || needed) || 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} + boxedchildren::Dict{AC.Field,Vector{Any}} + routeplan::Union{Nothing,_ArrowTypesRoutePlan} + warn::Bool + newsymbols::Int # Symbols this read created (see `_liftsymbol!`) + 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, + AC.Field, + Vector{Any}, + "ArrowTypes boxed-children cache container", + ) + 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}(), + Dict{AC.Field,Vector{Any}}(), + nothing, + warn, + 0, + 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 _uninferred( + any, + child -> _hasarrowtypesextension(child, ctx), + f.children, + )::Bool + 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 + storage = _facadestoragetype(T) + storage === nothing || return storage + return _uninferred(ArrowTypes.ArrowType, 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 + # Row lifting asks per union value: a cache hit must not construct the + # memoization thunk. + cached = get(ctx.storage, T, nothing) + cached === nothing || return cached + 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) + # Row lifting asks per value: a cache hit must not construct the + # memoization thunk. + cached = get(ctx.targets, f, nothing) + cached === nothing || return cached + 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 + if isempty(f.children) + storage = _arrowtypesstoragebasetype(ctx, f) + facade = _facadejuliatype(name, storage) + facade === nothing || return (true, facade, storage) + _quietextension(name) && return (true, nothing, storage) + end + 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) + # Row restoration asks per value: a cache hit must not construct the + # memoization thunk. + cached = get(ctx.logicaltypes, f, nothing) + cached === nothing || return cached + 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) + # Row restoration asks per value: a cache hit must not construct the + # memoization thunk. + cached = get(ctx.publictypes, f, nothing) + cached === nothing || return cached + 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 + # An unmarked or unknown-label leaf materializes in the native + # public domain regardless of any storage fallback the label lookup + # produced: no registered mapping consumes these values. The native + # declared rule already carries nullability, and a Null leaf keeps + # its intrinsic Missing domain. + return _declaredeltype(f, true) + 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 _uninferred(map, identity, values) +end +function _typedvalues(::Type{T}, values, budget=nothing) where {T} + _chargevector!(budget, T, length(values), "ArrowTypes typed column") + return _collectas(T, 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 + t isa AC.TimeType && return _publictime(t.unit, x) + 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 + # Per dictionary value during row lifting: hits must not allocate. + cached = get(cache, f, nothing) + cached === nothing || return cached + 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, + rawdomain::Bool=false, +) + # Indexing instead of destructuring: per-value callers must not box + # iteration-state tuples beside every row. + target = _arrowtypestarget(ctx, f)[2] + # 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) + # A label with no target left its whole subtree in the raw storage + # domain (see the routed walker); its values convert here on the way + # into public rows. + return _arrowtypesstoragevalue( + ctx, + f, + x; + extension_shape=extension_shape, + rawdomain=rawdomain || _arrowtypesextension(f) !== nothing, + ) +end + +function _arrowtypesnestedeltype( + ctx::_ArrowTypesContext, + f::AC.Field; + extension_shape::Bool, +) + target = _arrowtypestarget(ctx, f)[2] + if target !== nothing || extension_shape + return _arrowtypeslogicaleltype(ctx, f) + end + return _arrowtypespubliceltype(ctx, f) +end + +""" +Reserve the dynamic per-element restoration overhead for `n` nested calls: +each one is a dynamic keyword invocation that leaves a small argument tuple +and scalar boxes behind (32 bytes covers the measured worst case). +""" +function _chargerestoredelements!( + budget::Union{Nothing,AllocationBudget}, + n::Integer, + what::AbstractString, +) + budget === nothing && return nothing + _charge!(budget, AC.checked_mul(Int64(n), Int64(32)), what) + return nothing +end + +""" +One heap-boxed reference per child Field. `Field` is an inline immutable, so +handing `f.children[i]` to a dynamic per-row call re-boxes it beside every +value; restoration loops index this cached vector instead. +""" +function _arrowtypesboxedchildren(ctx::_ArrowTypesContext, f::AC.Field) + cached = get(ctx.boxedchildren, f, nothing) + cached === nothing || return cached + return _memoized!( + ctx.boxedchildren, + f, + ctx.budget, + "ArrowTypes boxed-children cache", + ) do + _chargevector!( + ctx.budget, + AC.Field, + length(f.children), + "ArrowTypes boxed children", + ) + Any[child for child in f.children] + end +end + +"Convert children, preserving the reader's unmarked row containers unless requested." +function _arrowtypesstoragevalue( + ctx::_ArrowTypesContext, + f::AC.Field, + x; + extension_shape::Bool, + rawdomain::Bool=false, +) + t = f.type + x === missing && return missing + t isa AC.DictionaryType && return _arrowtypesstoragevalue( + ctx, + _arrowtypesdictvaluefield(ctx, f, t), + x; + extension_shape=extension_shape, + rawdomain=rawdomain, + ) + if t isa AC.RunEndEncodedType + length(f.children) == 2 || return x + return _arrowtypesnestedvalue( + ctx, + _arrowtypesboxedchildren(ctx, f)[2], + x; + extension_shape=extension_shape, + rawdomain=rawdomain, + ) + end + if t isa Union{AC.ListType,AC.ListViewType,AC.FixedSizeListType} + length(f.children) == 1 || return x + child = f.children[1] + childbox = _arrowtypesboxedchildren(ctx, f)[1] + _chargevector!(ctx.budget, Any, length(x), "ArrowTypes converted list value") + _chargerestoredelements!(ctx.budget, length(x), "ArrowTypes restored list elements") + vals = Any[ + _arrowtypesnestedvalue( + ctx, + childbox, + y; + extension_shape=extension_shape, + rawdomain=rawdomain, + ) for y in x + ] + if t isa AC.FixedSizeListType && extension_shape + _chargevector!(ctx.budget, Any, length(vals), "ArrowTypes tuple storage") + out = Tuple(vals) + # The tuple construction and the dynamic hook call each copy + # inline (isbits) child payloads wholesale; heap children + # contribute pointer slots only, which sizeof reflects. + _chargeobject!( + ctx.budget, + AC.checked_mul(Int64(2), Int64(sizeof(typeof(out)))), + "ArrowTypes tuple payload", + ) + return out + 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") + _chargerestoredelements!( + ctx.budget, + length(f.children), + "ArrowTypes restored struct children", + ) + boxedchildren = _arrowtypesboxedchildren(ctx, f) + 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, + boxedchildren[i], + last(kv); + extension_shape=extension_shape, + rawdomain=rawdomain, + ) + end + if extension_shape + # Rebuilding the storage base type per row walks every child and + # re-applies the NamedTuple constructor. A registered Field's + # target cache carries the exact type; an unlabeled Field reads + # it through the logical-eltype memo. An unknown label keeps the + # direct computation: its cached slot holds the bounded fallback + # shape, not the exact one. + tgt = _arrowtypestarget(ctx, f) + NT = + tgt[2] !== nothing ? tgt[3] : + tgt[1] === false ? Base.nonmissingtype(_arrowtypeslogicaleltype(ctx, f)) : + _arrowtypesstoragebasetype(ctx, f) + if NT <: NamedTuple + _chargevector!( + ctx.budget, + Any, + length(vals), + "ArrowTypes NamedTuple storage", + ) + # The transient tuple and the constructed row each copy the + # row's inline (isbits) payload; heap-backed children only + # contribute pointer slots, which sizeof reflects. + NT isa DataType && + isconcretetype(NT) && + _chargeobject!( + ctx.budget, + AC.checked_mul(Int64(2), Int64(sizeof(NT))), + "ArrowTypes NamedTuple payload", + ) + 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", + ) + # The raw-domain path (a label with no resolvable target) rebuilds + # every child pair with freshly boxed scalars; reserve that beside + # the pair vector. + _chargerestoredelements!( + ctx.budget, + AC.checked_mul(Int64(3), Int64(length(vals))), + "ArrowTypes raw struct rows", + ) + 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 + entryboxes = _arrowtypesboxedchildren(ctx, entries) + kfbox, vfbox = entryboxes[1], entryboxes[2] + _chargevector!( + ctx.budget, + Pair{Any,Any}, + length(x), + "ArrowTypes converted map value", + ) + _chargerestoredelements!( + ctx.budget, + AC.checked_mul(Int64(2), Int64(length(x))), + "ArrowTypes restored map entries", + ) + vals = Pair{Any,Any}[ + _arrowtypesnestedvalue( + ctx, + kfbox, + first(kv); + extension_shape=extension_shape, + rawdomain=rawdomain, + ) => _arrowtypesnestedvalue( + ctx, + vfbox, + last(kv); + extension_shape=extension_shape, + rawdomain=rawdomain, + ) for kv in x + ] + if extension_shape + # Same caching rule as the struct branch: registered Fields use + # the exact cached storage type, unlabeled Fields the memoized + # logical eltype, unknown labels the direct computation. + tgt = _arrowtypestarget(ctx, f) + D = + tgt[2] !== nothing ? tgt[3] : + tgt[1] === false ? Base.nonmissingtype(_arrowtypeslogicaleltype(ctx, f)) : + _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, + _arrowtypesboxedchildren(ctx, f)[x.child], + x.value; + extension_shape=extension_shape, + rawdomain=rawdomain, + ) + end + boxedchildren = _arrowtypesboxedchildren(ctx, f) + for (k, child) in enumerate(f.children) + T = _arrowtypesnestedeltype(ctx, child; extension_shape=extension_shape) + x isa T && return _arrowtypesnestedvalue( + ctx, + boxedchildren[k], + x; + extension_shape=extension_shape, + rawdomain=rawdomain, + ) + end + return x + end + # Only a marked parent's ArrowTypes storage shape asks for ArrowTypes + # scalars here — that domain is the registered mappings' contract. A + # leaf outside any label already arrived public from the routed walker; + # a leaf inside a label-without-target subtree arrived as raw storage + # and converts here, landing in public rows like every unmarked value. + if extension_shape + if t isa AC.FixedSizeBinaryType && _arrowtypesfixedbinarytype(t) !== Vector{UInt8} + # The tuple is built here and copied again into the dynamic hook call. + _chargeobject!(ctx.budget, 2 * t.nbytes, "ArrowTypes fixed-size binary tuple") + return NTuple{t.nbytes,UInt8}(x) + end + return _arrowtypesscalar(t, x) + end + if rawdomain && _isconvertibleleaf(t) + # Same reserve as the facade's converted leaf: the public box plus + # its transient conversion temporaries. + _chargeobject!(ctx.budget, 96, "ArrowTypes converted leaf") + return _publicleafscalar(t, x) + end + return 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) + value = _facadefromarrow(T, f, storage) + value === nothing || return value + t = f.type + if _arrowtypesutf8storage(ctx, f) && storage isa AbstractString + T === Symbol && return _liftsymbol!(ctx, storage, f.name) + bytes = _utf8units(storage) + GC.@preserve storage bytes begin + return _uninferred(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 + # Reconstruction copies the row's inline (isbits) payload at each + # seam — the transient tuple, the splatted arguments, and an isbits + # target's returned box; heap-backed children contribute pointer + # slots only, which sizeof reflects. + storage isa Union{NamedTuple,Tuple} && _chargeobject!( + ctx.budget, + AC.checked_mul(Int64(3), Int64(sizeof(typeof(storage)))), + "ArrowTypes struct hook payload", + ) + if T <: NamedTuple || T <: Tuple + # The transient tuple and constructed row beside the payload. + _chargerestoredelements!( + ctx.budget, + AC.checked_mul(Int64(2), Int64(length(storage))), + "ArrowTypes struct hook tuple", + ) + return T(Tuple(storage)) + end + names = _arrowtypesstructnames(ctx, f) + # The `applicable` probe and the variadic hook call each re-splat + # every child value, re-boxing scalars whose types Julia does not + # cache; charge that width-dependent dispatch work per child. + ctx.budget === nothing || _charge!( + ctx.budget, + AC.checked_mul(Int64(length(storage)), Int64(160)), + "ArrowTypes struct hook dispatch", + ) + 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) + target = _arrowtypestarget(ctx, f)[2] + if target === nothing + return _arrowtypesstoragevalue( + ctx, + f, + x; + extension_shape=false, + rawdomain=_arrowtypesextension(f) !== nothing, + ) + 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) + # Reserve the lifted value plus the dynamic-dispatch temporaries around + # the fromarrow hook. Childless (flat scalar) storage leaves about 140 + # bytes of boxes per lifted value; container storage's per-element + # restoration and hook-dispatch reserves already carry its width, so + # only a small fixed remainder rides per row; union storage and the + # transparent wrappers (run-end, dictionary) lift one scalar with no + # per-element reserve and keep the larger fixed one. + _chargeobject!( + ctx.budget, + isempty(f.children) ? 152 : + f.type isa + Union{AC.StructType,AC.ListType,AC.ListViewType,AC.FixedSizeListType,AC.MapType} ? + 96 : 216, + "ArrowTypes lifted value", + ) + # The default ArrowTypes pointer adapter copies a string's bytes again + # (`unsafe_string`) before the hook sees them: that copy grows with the + # value, so reserve it by length (plus the copy's measured header and + # size-class slop) beside the fixed overhead. + storage isa AbstractString && _chargeobject!( + ctx.budget, + AC.checked_add(Int64(ncodeunits(storage)), Int64(256)), + "ArrowTypes lifted string copy", + ) + return _arrowtypesfromarrow(ctx, target, f, storage) +end + +# Interning a Symbol permanently allocates process-global state, so one read +# (one stream batch) creates at most `Limits.max_new_symbols` of them, and +# charges their bytes to its allocation budget; existing Symbols cost nothing. +_maxnewsymbols(ctx::_ArrowTypesContext) = + ctx.budget === nothing ? _DEFAULT_MAX_NEW_SYMBOLS : ctx.budget.maxnewsymbols + +function _liftsymbol!(ctx::_ArrowTypesContext, s::AbstractString, name::AbstractString) + # As a `String`: see `_existingjlsymbol`. + s isa String || return _liftsymbol!(ctx, String(s)::String, name) + symbol = _existingjlsymbol(s) + symbol === nothing || return symbol + occursin('\0', s) && _symbolnulerror(ctx, s, name) + # Top-level columns were counted before conversion (`_checknewsymbols`). + ctx.newsymbols < _maxnewsymbols(ctx) || _symbollimiterror(ctx, name, false) + ctx.budget === nothing || + _charge!(ctx.budget, Int64(ncodeunits(s)), "interned Symbol values") + ctx.newsymbols += 1 + return Symbol(s) +end + +@noinline function _symbolnulerror(ctx::_ArrowTypesContext, s, name) + displayname = _extensionwarningname(s, ctx.budget) + throw( + AC.ValidationError( + "field $name holds the JuliaLang.Symbol value $(repr(displayname)), " * + "which contains a NUL byte and cannot be a Symbol", + ), + ) +end + +# `column`: `name` is a top-level column, which `select` can read as strings. +@noinline function _symbollimiterror(ctx::_ArrowTypesContext, name, column::Bool) + limit = _maxnewsymbols(ctx) + # `repr` of a Symbol runs Base code that other packages invalidate. + sym = _uninferred(repr, Symbol(name))::String + hint = column ? ", or read the column as strings with select=[$sym => String]" : "" + throw( + AC.ValidationError( + "$(column ? "column" : "field") $name holds more than $limit Symbol values " * + "this process has not created yet, the most one read may create (Julia " * + "never frees a Symbol). Raise the limit with " * + "limits=Arrow.Limits(max_new_symbols=...)$hint", + ), + ) +end + +# A JuliaLang.Symbol column's new names, counted (deduplicated) before any +# is created, so a refused read leaves no Symbols behind. +function _checknewsymbols(ctx::_ArrowTypesContext, f::AC.Field, col::AbstractVector) + limit = _maxnewsymbols(ctx) - ctx.newsymbols + new = Set{String}() + charged = Int64(0) + for x in col + x isa AbstractString || continue + _existingjlsymbol(x) === nothing || continue + occursin('\0', x) && _symbolnulerror(ctx, x, f.name) + x in new && continue + length(new) < limit || _symbollimiterror(ctx, f.name, true) + # The set's entry and its copy of the name, returned below. + bytes = Int64(ncodeunits(x) + 64) + ctx.budget === nothing || _charge!(ctx.budget, bytes, "new Symbol name count") + charged += bytes + push!(new, String(x)) + end + ctx.budget === nothing || _refund!(ctx.budget, charged) + return nothing +end + +"Heap-box an inline immutable once; @noinline keeps the box from folding away." +@noinline _boxonce(@nospecialize(x)) = x + +"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 && + isempty(f.children) + # No hook or routed provenance is available to consume, and the + # column is flat (its raw values sit under one leaf descriptor, via + # transparent wrappers at most): convert the whole column at once + # without building a per-row ArrowTypes workspace. A composite + # column with an unconsumed label falls through to the per-row + # lifting below, which converts its raw-domain leaves in place. + return _publiccolumn(f, _postconvertfield(f, col, ctx.budget), ctx.budget) + end + target === Symbol && _checknewsymbols(ctx, f, col) + _chargevector!(ctx.budget, Any, length(col), "ArrowTypes lifting workspace") + # `f` is an immutable struct held inline in this frame: box it once, + # behind a call boundary the optimizer cannot fold away, so the dynamic + # per-row call does not re-box it beside every value. + fbox = _boxonce(f) + values = Any[_arrowtypesvalue(ctx, fbox, 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 + # The lifted rows: routed markers consumed, and raw-domain leaves + # under the unconsumed label already converted to their public + # values — no whole-column conversion may run again. + return _publiccolumn(f, values, 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..1e730b40 --- /dev/null +++ b/src/cdata.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. + +# ============================================================================= +# 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: every C node of one export (a schema tree, an array tree, or a +# stream) carries in `private_data` a pointer to one malloc'd count of the +# nodes the consumer has not released. A release callback releases the +# children and dictionary still in place, marks the node released, and +# counts it off. It takes no lock, allocates nothing, and touches no Julia +# object, because consumers call it from their own threads and from inside +# GC finalizers. The Julia side keeps each export's roots (the columns, +# whose regions keep the buffers alive) and mallocs in `EXPORTS` until the +# count reaches zero; later exports or garbage collections free them. +# `to_c_data`'s two top-level structs, the only memory Arrow.jl allocates +# that a consumer can read after releasing it, are freed by the first reap +# at least 10 ms after the rest. +# +# * Import: the moved ArrowArray becomes ONE ForeignOwner shared by every +# child/dictionary BufferSlice (a single release for the whole tree). +# 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, and every region shares one `ReleaseCell`, so +# `release!` on the owner or on any region revokes all siblings (later +# access is a clean `InvalidStateException`) and then runs the producer's +# release exactly once. +# +# * Streams: `ArrowArrayStream` maps in both directions. An exported stream +# is one export whose `get_schema` and `get_next` results are exports of +# their own; an imported stream is a moved struct whose batches import as +# arrays do. +# ============================================================================= + +# --------------------------------------------------------------------------- +# ABI structs (field-exact per https://arrow.apache.org/docs/format/CDataInterface.html) +# --------------------------------------------------------------------------- + +""" + Arrow.CArrowSchema + +The C data interface's `ArrowSchema` struct, field for field. +""" +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 + +""" + Arrow.CArrowArray + +The C data interface's `ArrowArray` struct, field for field. +""" +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 + +# 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) + +# Call the release callback of the C struct at `p`, if it still has one. +function _release_c!(p::Ptr{T}) where {T} + r = unsafe_load(p).release + r == C_NULL || ccall(r, Cvoid, (Ptr{T},), p) + return nothing +end + +# Release a producer's struct and check that its callback marked the struct +# released, as the interface requires. +function _release_producer!(p::Ptr{T}) where {T} + _release_c!(p) + unsafe_load(p).release == C_NULL || + error("C Data producer release did not mark the structure released") + return nothing +end + +# --------------------------------------------------------------------------- +# 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" + +function formatstring(t::IntType) + t.bits == 128 && throw( + ArgumentError( + "the C data interface has no 128-bit integer format; " * + "convert the column to a decimal or fixed-size binary first", + ), + ) + return ( + t.signed ? Dict(8 => "c", 16 => "s", 32 => "i", 64 => "l") : + Dict(8 => "C", 16 => "S", 32 => "I", 64 => "L") + )[t.bits] +end +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, one live-node count per export +# --------------------------------------------------------------------------- + +# One export's Julia roots and native allocations. `live` is a malloc'd count +# of the C nodes the consumer has not released; every node's `private_data` +# points at it, so a release callback needs no Julia object. The 8 bytes after +# the count hold an exported stream's last error message. `top` holds +# `to_c_data`'s two top-level structs (C_NULL for a stream's exports, which fill +# structs the consumer allocated); see `_RETIRED`. +struct Export + roots::Vector{Any} + mallocs::Vector{Ptr{Cvoid}} + live::Ptr{Int64} + top::Ptr{Cvoid} +end + +const EXPORTS = Dict{Ptr{Int64},Export}() +const REGISTRY_LOCK = ReentrantLock() +# Exports left to make before one reaps. A reap scans every registered export, so it +# sets this to the number still unreleased (at least 64), and reaping costs constant +# time per export however many exports a consumer holds. +const _EXPORTS_UNTIL_REAP = Ref(64) +# The most exports registered since `EXPORTS` last shrank. Deleting keeps the slots a +# past peak grew, and a reap walks every slot, so a reap that leaves far fewer shrinks it. +const _EXPORTS_PEAK = Ref(0) +# The `top` blocks of freed exports, oldest first, with the `time_ns()` of the reap that +# freed them. A consumer that releases a top-level struct in place may read it after the +# release callback returns (Arrow C++ checks that its `release` is NULL), and a preempted +# consumer can take a while to get there, so a reap frees a block only once it has been +# retired for `_TOP_FREE_AFTER_NS`. +const _RETIRED = Tuple{UInt64,Ptr{Cvoid}}[] +const _TOP_FREE_AFTER_NS = UInt64(10_000_000) + +# A new export of `roots`; `topsize` bytes for top-level structs it allocates itself. +function _newexport(roots::Vector{Any}, topsize::Integer=0) + reap = lock(REGISTRY_LOCK) do + return (_EXPORTS_UNTIL_REAP[] -= 1) <= 0 + end + reap && _reap!() + live = Ptr{Int64}(Libc.calloc(2, 8)) + live == C_NULL && throw(OutOfMemoryError()) + top = topsize > 0 ? Libc.malloc(topsize) : C_NULL + if topsize > 0 && top == C_NULL + Libc.free(live) + throw(OutOfMemoryError()) + end + return Export(roots, Ptr{Cvoid}[live], live, top) +end + +function _malloc!(e::Export, n::Integer) + p = Libc.malloc(max(Int(n), 1)) + p == C_NULL && throw(OutOfMemoryError()) + push!(e.mallocs, p) + return p +end + +function _free!(e::Export) + message = unsafe_load(Ptr{Ptr{UInt8}}(e.live + 8)) + message == C_NULL || Libc.free(message) + foreach(Libc.free, e.mallocs) + empty!(e.mallocs) + empty!(e.roots) + return nothing +end + +# An object whose finalizer reaps and then registers itself again, so a +# consumer's releases free memory at later garbage collections even when +# nothing exports again. Created once, by the first export. +mutable struct _ReapHook end +const _REAP_HOOK_ARMED = Ref(false) +# Set when Julia starts exiting; see `_finalizeowner`. +const _EXITING = Ref(false) +function _setexiting() + _EXITING[] = true + return nothing +end + +function __init__() + # No export outlives the process that made it: drop what a precompile + # workload left (its pointers come back NULL), and re-arm the reap hook, + # whose finalizer is not saved. + empty!(EXPORTS) + empty!(_RETIRED) + _EXPORTS_UNTIL_REAP[] = 64 + _EXPORTS_PEAK[] = 0 + _REAP_HOOK_ARMED[] = false + atexit(_setexiting) + return nothing +end + +function _reaphook(h::_ReapHook) + finalizer(_reaphook, h) + trylock(REGISTRY_LOCK) || return nothing + try + _reap!() + finally + unlock(REGISTRY_LOCK) + end + return nothing +end + +function _register!(e::Export) + lock(REGISTRY_LOCK) do + EXPORTS[e.live] = e + if !_REAP_HOOK_ARMED[] + _REAP_HOOK_ARMED[] = true + finalizer(_reaphook, _ReapHook()) + end + end + return nothing +end + +# Free every export whose C nodes are all released, retire its `top` block, free the +# blocks retired long enough ago, and return how many exports it freed. Runs from +# `_newexport` (see `_EXPORTS_UNTIL_REAP`) and from `_ReapHook`, so no caller needs to. +function _reap!() + done = Export[] + expired = Ptr{Cvoid}[] + lock(REGISTRY_LOCK) do + peak = _EXPORTS_PEAK[] = max(_EXPORTS_PEAK[], length(EXPORTS)) + for e in values(EXPORTS) + unsafe_load(e.live, :sequentially_consistent) == 0 && push!(done, e) + end + # Read after the scan, so a block's stamp is later than the release the scan saw. + now = time_ns() + for e in done + delete!(EXPORTS, e.live) + e.top == C_NULL || push!(_RETIRED, (now, e.top)) + end + while !isempty(_RETIRED) && now - first(_RETIRED)[1] >= _TOP_FREE_AFTER_NS + push!(expired, popfirst!(_RETIRED)[2]) + end + # A drained queue gives back the capacity a burst grew. + isempty(_RETIRED) && sizehint!(_RETIRED, 0) + n = length(EXPORTS) + if peak > 4 * max(64, n) + sizehint!(EXPORTS, n) + _EXPORTS_PEAK[] = n + end + _EXPORTS_UNTIL_REAP[] = max(64, n) + return nothing + end + foreach(Libc.free, expired) + foreach(_free!, done) + return length(done) +end + +# Release callbacks do only native work: release the children and dictionary +# still in place (a moved one has `release == NULL`), mark the node released, +# and count it off its export. No lock, no allocation, no Julia object, and +# nothing that can throw: consumers call them from their own threads, and +# PythonCall runs pyarrow's releases from GC finalizers. +function _release_array(a::Ptr{CArrowArray})::Cvoid + arr = unsafe_load(a) + arr.release == C_NULL && return nothing + for i = 1:arr.n_children + _release_c!(unsafe_load(arr.children, i)) + end + arr.dictionary == C_NULL || _release_c!(arr.dictionary) + _store_field!(a, Val(:release), Ptr{Cvoid}(C_NULL)) + Base.unsafe_modify!(Ptr{Int64}(arr.private_data), -, 1, :sequentially_consistent) + return nothing +end + +function _release_schema(s::Ptr{CArrowSchema})::Cvoid + sch = unsafe_load(s) + sch.release == C_NULL && return nothing + for i = 1:sch.n_children + _release_c!(unsafe_load(sch.children, i)) + end + sch.dictionary == C_NULL || _release_c!(sch.dictionary) + _store_field!(s, Val(:release), Ptr{Cvoid}(C_NULL)) + Base.unsafe_modify!(Ptr{Int64}(sch.private_data), -, 1, :sequentially_consistent) + return nothing +end + +# 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!( + e::Export, + metadata::Union{Nothing,AC.FrozenVector{Pair{String,String}}}, +)::Ptr{UInt8} + (metadata === nothing || isempty(metadata)) && return Ptr{UInt8}(C_NULL) + buf = UInt8[] + append!(buf, reinterpret(UInt8, Int32[Int32(length(metadata))])) + 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!(e, length(buf))) + GC.@preserve buf unsafe_copyto!(p, pointer(buf), length(buf)) + return p +end + +function _cstring!(e::Export, s::String) + 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!(e, n + 1)) + GC.@preserve s unsafe_copyto!(p, pointer(s), n) + unsafe_store!(p, 0x00, n + 1) + return p +end + +# A dictionary Field's extension pairs name its VALUE type, as in IPC. For a +# canonical extension (the format reserves the `arrow.` prefix for them), export +# them on the dependent value node, as the C++ bridge does for an extension type +# it knows: on the wrapper they would declare dictionary storage, which PyArrow +# refuses. Other extension names stay on the wrapper with the field's metadata, +# where a consumer that does not know them keeps them as field metadata (PyArrow +# drops unknown extension pairs on a dependent node). +# (Literal keys: the trim gate includes this file without arrowtypes.jl.) +_isextensionpair(kv) = + first(kv) == "ARROW:extension:name" || first(kv) == "ARROW:extension:metadata" + +_iscanonicalextension(kv) = + first(kv) == "ARROW:extension:name" && startswith(last(kv), "arrow.") + +function _splitdictmetadata(metadata::Union{Nothing,AC.FrozenVector{Pair{String,String}}}) + (metadata === nothing || !any(_iscanonicalextension, metadata)) && + return metadata, nothing + wrapper = Pair{String,String}[kv for kv in metadata if !_isextensionpair(kv)] + values = Pair{String,String}[kv for kv in metadata if _isextensionpair(kv)] + return (isempty(wrapper) ? nothing : AC.FrozenVector{Pair{String,String}}(wrapper)), + AC.FrozenVector{Pair{String,String}}(values) +end + +# Export the schema tree of `f` into `at`, a struct the caller allocated, or into +# a malloc'd one. The node is stored last, after its children, so a failed +# export leaves `at` untouched. +function _export_schema!( + e::Export, + f::Field, + release::Ptr{Cvoid}; + at::Ptr{CArrowSchema}=Ptr{CArrowSchema}(C_NULL), + metadata::Union{Nothing,AC.FrozenVector{Pair{String,String}}}=f.metadata, +)::Ptr{CArrowSchema} + p = at == C_NULL ? Ptr{CArrowSchema}(_malloc!(e, sizeof(CArrowSchema))) : at + childfields = f.type isa DictionaryType ? Field[] : f.children + childptrs = Ptr{Ptr{CArrowSchema}}(C_NULL) + if !isempty(childfields) + childptrs = Ptr{Ptr{CArrowSchema}}(_malloc!(e, length(childfields) * sizeof(Ptr))) + for (i, cf) in enumerate(childfields) + unsafe_store!(childptrs, _export_schema!(e, cf, release), i) + end + end + dict = Ptr{CArrowSchema}(C_NULL) + wrappermetadata = metadata + if f.type isa DictionaryType + wrappermetadata, valuemetadata = _splitdictmetadata(metadata) + dict = _export_schema!( + e, + AC.dictvaluefield(f, f.type), + release; + metadata=valuemetadata, + ) + 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) + unsafe_store!(e.live, unsafe_load(e.live) + 1) + unsafe_store!( + p, + CArrowSchema( + _cstring!(e, formatstring_of(f.type)), + _cstring!(e, f.name), + # Field metadata rides the OUTER node for every field, dictionary + # wrappers included; PyArrow imports only the wrapper's pairs. + _cmetadata!(e, wrappermetadata), + flags, + length(childfields), + childptrs, + dict, + release, + Ptr{Cvoid}(e.live), + ), + ) + return p +end + +function _export_array!( + e::Export, + d::ArrayData, + release::Ptr{Cvoid}; + at::Ptr{CArrowArray}=Ptr{CArrowArray}(C_NULL), +)::Ptr{CArrowArray} + p = at == C_NULL ? Ptr{CArrowArray}(_malloc!(e, sizeof(CArrowArray))) : at + 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!(e, max(nbuf, 1) * 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 omits the offsets of an empty array; C Data still + # expects the length+1 offsets, so export one zero. + Libc.memset(_malloc!(e, spec.offsetwidth), 0, spec.offsetwidth) + elseif AC.isempty_buffer(b) + # An absent validity bitmap, or any zero-byte buffer, is NULL. + Ptr{Cvoid}(C_NULL) + else + Ptr{Cvoid}(AC.sliceptr(b)) + end + unsafe_store!(bufptrs, bufferp, i) + end + if spec.variadic + sizesp = Ptr{Int64}(_malloc!(e, max(nvariadic, 1) * 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 + childptrs = Ptr{Ptr{CArrowArray}}(C_NULL) + if !isempty(d.children) + childptrs = Ptr{Ptr{CArrowArray}}(_malloc!(e, length(d.children) * sizeof(Ptr))) + for (i, c) in enumerate(d.children) + unsafe_store!(childptrs, _export_array!(e, c, release), i) + end + end + dict = + d.dictionary === nothing ? Ptr{CArrowArray}(C_NULL) : + _export_array!(e, d.dictionary, release) + unsafe_store!(e.live, unsafe_load(e.live) + 1) + unsafe_store!( + p, + CArrowArray( + d.len, + nullcount(d), + d.offset, + nbuf, + length(d.children), + bufptrs, + childptrs, + dict, + release, + Ptr{Cvoid}(e.live), + ), + ) + return p +end + +""" + Arrow.to_c_data(field, data) -> (Ptr{CArrowSchema}, Ptr{CArrowArray}) + +The low-level form of [`Arrow.to_c_data`](@ref): export one column, an +`Arrow.Field` and its `Arrow.ArrayData`, as an `ArrowSchema` and an `ArrowArray` +that Arrow.jl allocates, pointing at `data`'s buffers without copying them. The +column is validated before it is exported. The consumer moves or releases the +two structs; `data` stays reachable until it has released the array, and what +Arrow.jl allocated is freed by later exports or garbage collections. +""" +function to_c_data(f::Field, d::ArrayData) + validate_semantic(f, d) + e = _newexport(Any[f, d], sizeof(CArrowSchema) + sizeof(CArrowArray)) + try + sp = _export_schema!( + e, + f, + @cfunction(_release_schema, Cvoid, (Ptr{CArrowSchema},)); + at=Ptr{CArrowSchema}(e.top), + ) + ap = _export_array!( + e, + d, + @cfunction(_release_array, Cvoid, (Ptr{CArrowArray},)); + at=Ptr{CArrowArray}(e.top + sizeof(CArrowSchema)), + ) + _register!(e) + return sp, ap + catch + _free!(e) + Libc.free(e.top) + rethrow() + end +end + +# --------------------------------------------------------------------------- +# Import: C structs -> Core, one ForeignOwner per moved tree +# --------------------------------------------------------------------------- + +# One owner for one MOVED ArrowArray tree. Every `BufferSlice` of the tree +# (children and dictionary) has a region whose `root` is this object and whose +# cell is `cell`, so the tree outlives every slice, `release!` on the owner or +# on any region revokes all of them, and the producer's release callback runs +# exactly once: from `release!` or from the GC finalizer, whichever comes first. +mutable struct ForeignOwner + # A malloc'd copy of the moved struct: the stable address its callback takes. + const block::Ptr{CArrowArray} + const thread::Int # where a collected owner releases (`_releasethread`) + cell::AC.ReleaseCell + ForeignOwner(block::Ptr{CArrowArray}) = new(block, _releasethread()) +end + +# The thread whose task releases an import made on this thread once it is collected: +# this thread, or thread 1 for a thread Julia adopted from a foreign library, which +# runs no Julia tasks. +_releasethread() = Threads.threadpool() === :foreign ? 1 : Threads.threadid() + +# The finalizer of an imported array or stream owner. A producer's release can +# wait on a lock the importing thread holds (pyarrow's release of NumPy memory +# takes Python's GIL, which PythonCall's main thread keeps), and a thread waiting +# inside a finalizer stalls every garbage collection. So a finalizer on another +# thread hands the release to a task on the owner's `thread` (the importing thread, +# or thread 1), which runs it at its next yield. At exit, a release that would be +# handed to another thread is skipped: Julia runs the finalizers of live objects then, +# and the process is ending. +function _finalizeowner(o) + Threads.threadid() == o.thread && return release!(o) + _EXITING[] && return nothing + t = Task(() -> release!(o)) + t.sticky = true + ccall(:jl_set_task_tid, Cint, (Any, Cint), t, o.thread - 1) + schedule(errormonitor(t)) + return nothing +end + +function _release_owner_action(p::Ptr{Cvoid})::Cvoid + o = unsafe_pointer_to_objref(p)::ForeignOwner + GC.@preserve o try + _release_producer!(o.block) + finally + Libc.free(o.block) + end + return nothing +end + +# Release an imported tree now: later access to its buffers throws an +# `InvalidStateException`, and the producer's release callback runs, once. A +# producer whose callback does not mark the struct released makes this throw. +release!(o::ForeignOwner) = release!(o.cell) + +# A deep copy of imported data shares its owner, so the producer's release runs +# once, for the original and every copy. +Base.deepcopy_internal(o::ForeignOwner, ::IdDict) = o + +# Move the array at `ap`: copy its struct, mark the source released, and give +# the copy an owner whose `release!` or finalizer runs the producer's release. +function _moveowner(ap::Ptr{CArrowArray}) + block = Ptr{CArrowArray}(Libc.malloc(sizeof(CArrowArray))) + block == C_NULL && throw(OutOfMemoryError()) + unsafe_store!(block, unsafe_load(ap)) + _store_field!(ap, Val(:release), Ptr{Cvoid}(C_NULL)) + try + o = ForeignOwner(block) + o.cell = AC.ReleaseCell(@cfunction(_release_owner_action, Cvoid, (Ptr{Cvoid},)), o) + return finalizer(_finalizeowner, o) + catch + try + _release_producer!(block) + finally + Libc.free(block) + end + rethrow() + end +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) + +""" + Arrow.from_c_data(schemaptr, arrayptr) -> (Arrow.Field, Arrow.ArrayData) + +The low-level form of [`Arrow.Table(schemaptr, arrayptr)`](@ref Arrow.Table): import +one column. The array is moved, as the interface requires: Arrow.jl takes over its +`release` callback and marks the struct at `arrayptr` released. The schema is read and +then released. The imported buffers stay valid while the returned data is reachable; +the producer's release callback runs when it no longer is. + +Buffer sizes follow from the declared lengths and offsets, which the producer is +trusted to get right, as the interface cannot check them; the column is then validated +as an IPC read is. A failed import releases the array. +""" +function from_c_data(sp::Ptr{CArrowSchema}, ap::Ptr{CArrowArray}) + sp == C_NULL && throw(ArgumentError("ArrowSchema pointer is NULL")) + ap == C_NULL && throw(ArgumentError("ArrowArray pointer is NULL")) + sch = unsafe_load(sp) + (sch.release == C_NULL || unsafe_load(ap).release == C_NULL) && + throw(ArgumentError("cannot import a released structure")) + owner = _moveowner(ap) + f = try + _preflight_schema(sch) + _import_field(sch) + catch + release!(owner) + rethrow() + finally + _release_producer!(sp) + end + return f, _importarray(f, owner) +end + +# Import the moved array `owner` holds as a column of `f`; a failed import +# releases it. +function _importarray(f::Field, owner::ForeignOwner) + try + arr = unsafe_load(owner.block) + _preflight_array(f, arr) + d = _import_array(f, arr, owner) + validate_semantic(f, d) + return d + catch + release!(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")) + arr.length <= typemax(Int64) - arr.offset || throw( + ValidationError( + "C array length $(arr.length) plus offset $(arr.offset) overflows Int64", + ), + ) + -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 + +# 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 +opaque binary key/value data, copied by declared length without UTF-8 +validation. Embedded NUL bytes are preserved. +""" +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 + # Hostile extents overflow as ValidationErrors, as in Core's own + # structural checks. + total < typemax(Int64) || + throw(ValidationError("required offsets buffer bytes overflow Int64")) + AC._required_extent( + total + Int64(1), + Int64(spec.offsetwidth), + "offsets buffer bytes", + ) + elseif role == AC.DATA + if spec.fixedwidth > 0 + AC._required_extent(total, Int64(spec.fixedwidth), "data buffer bytes") + 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._required_extent(total, Int64(spec.offsetwidth), "per-slot buffer bytes") + elseif role == AC.VIEWS + AC._required_extent(total, Int64(16), "views buffer bytes") + 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 diff --git a/src/cdata_table.jl b/src/cdata_table.jl new file mode 100644 index 00000000..6da691ef --- /dev/null +++ b/src/cdata_table.jl @@ -0,0 +1,274 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 table-level C data interface: `Arrow.Table` and `Arrow.Stream` over +# imported C structs, and `Arrow.to_c_data` and `Arrow.export_stream!` over any +# Tables.jl table. Imports build columns the way an IPC read does; exports +# build record batches the way `Arrow.write` does. + +# --- import ----------------------------------------------------------------- + +# The read keywords, checked before anything is moved out of a C struct, so a +# bad keyword cannot strand an import. +function _creadargs(; + stringtype::Type=DataStrings.DataString, + ntasks::Integer=Threads.nthreads(), + limits::Limits=Limits(), + unknown..., +) + isempty(unknown) || throw( + ArgumentError( + "a C data import takes the keywords stringtype, ntasks, and limits, not " * + join(keys(unknown), ", "), + ), + ) + _checkstringtype(stringtype) + return stringtype, _checkntasks(ntasks), limits +end + +# The table of `batches` with schema `sch`, built by the scan executor as an IPC +# read builds one. View columns read the producer's memory without checking a flag, +# and the Table contract lets columns outlive `release!`, so `release!(table)` only +# refuses further column access; the producer's release runs when the last column is +# garbage collected, as a mapped file is unmapped. +function _ctable( + sch::AC.Schema, + batches::Vector{AC.RecordBatch}, + stringtype, + ntasks, + limits, +) + budget = _readbudget(limits, sum(_databytes, batches; init=Int64(0))) + fields = collect(AC.Field, sch.fields) + plan = _ScanPlan(Tables.Scan(), fields, budget) + ctx = _ScanContext(_ViewContext(stringtype, budget), _ArrowTypesContext(budget=budget)) + t, _ = _scantable( + _DecodedRecord[_DecodedRecord(b) for b in batches], + _importdecoder(fields, limits), + plan, + nothing, + ctx, + _uniqueschema(sch), + AC.OwnerRegion[AC.OwnerRegion(Ptr{UInt8}(C_NULL), 0)], + budget; + ntasks, + lazy=true, + ) + return t +end + +# A decoder for imported batches, which arrive decoded and validated. +_importdecoder(fields, limits) = + _Decoder(fields, IdDict{Field,Int64}(), AC._ValidatedDictionaries(), limits) + +# The bytes of a batch's buffers, which size the allowance its read charges. +_databytes(b::AC.RecordBatch) = sum(_databytes, b.columns; init=Int64(0)) +_databytes(d::AC.ArrayData) = + sum(x -> Int64(x.len), d.buffers; init=Int64(0)) + + sum(_databytes, d.children; init=Int64(0)) + +# Release the batches of a failed import; a batch without columns has nothing to hold. +function _releasebatches!(batches) + for b in batches + isempty(b.columns) || release!(b.columns[1].owner::ForeignOwner) + end + return nothing +end + +# A struct array is a record batch of its fields; any other array is a +# one-column batch named after its schema. +function _recordbatch(f::AC.Field, d::AC.ArrayData) + if f.type isa AC.StructType + sch = AC.Schema(collect(AC.Field, f.children); metadata=f.metadata) + return sch, _structbatch(sch, d) + end + sch = AC.Schema(AC.Field[f]) + return sch, AC.RecordBatch(sch, AC.ArrayData[d], d.len) +end + +""" + Arrow.Table(schemaptr::Ptr{Arrow.CArrowSchema}, arrayptr::Ptr{Arrow.CArrowArray}; + stringtype=DataString, ntasks=Threads.nthreads(), limits=Arrow.Limits()) + +Import one array through the C data interface as a table. A struct array (a record +batch) becomes a table of its fields; any other array becomes a one-column table named +after the schema. The array is moved, as the interface requires, and the schema is read +and released. The columns have the types an IPC read of the same data has. Numbers, +dates, and other fixed-width values are read in the producer's memory, and string and +binary values are copied. `Arrow.release!(table)` refuses further column access, and the +producer's `release` callback runs once the table and its columns are garbage collected. +`stringtype`, `ntasks`, and `limits` mean what they mean for `Arrow.Table`. +""" +function Table(sp::Ptr{CArrowSchema}, ap::Ptr{CArrowArray}; kw...) + stringtype, ntasks, limits = _creadargs(; kw...) + f, d = from_c_data(sp, ap) + try + sch, b = _recordbatch(f, d) + return _ctable(sch, AC.RecordBatch[b], stringtype, ntasks, limits) + catch + release!(d.owner::ForeignOwner) + rethrow() + end +end + +""" + Arrow.Table(streamptr::Ptr{Arrow.CArrowArrayStream}; + stringtype=DataString, ntasks=Threads.nthreads(), limits=Arrow.Limits()) + +Import a whole C stream as a table. The stream is moved, as the interface requires, +every record batch is read, and the stream is released. The columns chain the batches +as an IPC read's do; fixed-width values stay in the producer's memory, which its +`release` callbacks give back once the table and its columns are garbage collected. +[`Arrow.Stream(streamptr)`](@ref Arrow.Stream) reads one batch at a time instead. +""" +function Table(p::Ptr{CArrowArrayStream}; kw...) + stringtype, ntasks, limits = _creadargs(; kw...) + s = from_c_stream(p; limits) + batches = AC.RecordBatch[] + try + while (b = AC.nextbatch!(s)) !== nothing + push!(batches, b) + end + return _ctable(s.schema, batches, stringtype, ntasks, limits) + catch + _releasebatches!(batches) + rethrow() + finally + release!(s) + end +end + +# The cell action of an `Arrow.Stream` over a C stream: release the producer's stream. +function _release_imported_stream(p::Ptr{Cvoid})::Cvoid + release!(unsafe_pointer_to_objref(p)::ImportedStream) + return nothing +end + +""" + Arrow.Stream(streamptr::Ptr{Arrow.CArrowArrayStream}; + stringtype=DataString, ntasks=Threads.nthreads(), limits=Arrow.Limits()) + +Read a C stream one record batch at a time: each iteration pulls the next batch from the +producer and gives it as an [`Arrow.Table`](@ref), built as +[`Arrow.Table(streamptr)`](@ref Arrow.Table) builds its columns. The stream is moved, as the interface requires. It is iterated once, +from one task at a time, and its length is unknown until it is read. As for any +`Arrow.Stream`, [`Arrow.release!`](@ref) of the stream or of a table it yielded +releases the producer's stream and refuses further column access; columns already +taken stay valid. +""" +function Stream(p::Ptr{CArrowArrayStream}; kw...) + stringtype, ntasks, limits = _creadargs(; kw...) + src = from_c_stream(p; limits) + budget = _readbudget(limits, 0) + closed = AC.ReleaseCell(@cfunction(_release_imported_stream, Cvoid, (Ptr{Cvoid},)), src) + fields = _batchfields(src) + return Stream( + src, + AC.OwnerRegion[AC.OwnerRegion(Ptr{UInt8}(C_NULL), 0; cell=closed)], + budget, + closed, + _ArrowTypesContext(budget=budget), + ReentrantLock(), + _ViewContext(stringtype, budget), + ntasks, + nothing, + _ScanPlan(Tables.Scan(), fields, budget), + trues(length(fields)), + ) +end + +# The imported stream as an `Arrow.Stream` source: its batches arrive as the +# iteration pulls them (`_streambatch`), so their number is unknown until then. +_configuredlimits(s::ImportedStream) = s.limits +_tableschema(s::ImportedStream) = _uniqueschema(s.schema) +_batchfields(s::ImportedStream) = s.schema.fields +_decoder(s::ImportedStream) = _importdecoder(s.schema.fields, s.limits) +_nbatches(::ImportedStream) = throw( + ArgumentError( + "the number of batches of a C stream is unknown until it is read; " * + "iterate the Arrow.Stream", + ), +) + +# --- export ----------------------------------------------------------------- + +# The schema and record batches of `tbl`: a one-shot IPC stream write to nowhere +# that keeps each batch it builds, so construction, dictionary state, and refusals +# are the writer's own and a C consumer sees what pyarrow reads from the IPC bytes. +# The kept buffers are fresh allocations (`exported` turns off the writer's reusable +# pool), because the consumer holds on to them. One record batch per partition, unless +# `batchsize` splits it or a batch's string data outgrows 32-bit offsets. A first +# partition with no rows still fixes the schema; its empty batch is the only one with +# no rows. +function _exportbatches(tbl; batchsize::Integer=typemax(Int), kw...) + for k in (:file, :compress, :compresslevel, :dictdeltas, :dictreplacement) + haskey(kw, k) && throw(ArgumentError("$k does not apply to the C data interface")) + end + batches = AC.RecordBatch[] + w = _writer(devnull; oneshot=true, batchsize, exported=batches, kw...) + _writeall!(w, tbl) + return w.schema::AC.Schema, batches +end + +""" + Arrow.to_c_data(table; kw...) -> (Ptr{Arrow.CArrowSchema}, Ptr{Arrow.CArrowArray}) + +Export a Tables.jl table as one record batch through the C data interface: an +`ArrowSchema` and an `ArrowArray` that Arrow.jl allocates, whose struct array has one +child per column. The columns are built as [`Arrow.write`](@ref) builds them, with the +same keywords (`dictencode`, `stringview`, `largelists`, `metadata`, `colmetadata`), so a +consumer reads what it would read from the IPC bytes. A `Vector` of `Int8` to `Int64`, +`UInt8` to `UInt64`, or `Float16` to `Float64` crosses without copying, as do an +[`Arrow.Table`](@ref) column's integer, float, and boolean values and validity bitmaps; +what crosses without copying must not be modified or resized while the consumer reads +it. Everything else, `Bool` and `missing` in Julia vectors, decimals, dates and +timestamps, string and list offsets, a list's own validity bitmap, and fixed-size list +values included, is built into new buffers, and string bytes may be copied. The consumer +moves or releases the two structs; the table's memory stays reachable until it has +released the array, and what Arrow.jl allocated is freed by later exports or garbage +collections. A table that needs more than one record batch (several partitions +with rows, or string data past 2 GiB) is an `ArgumentError`; export it with +[`Arrow.export_stream!`](@ref). +""" +function to_c_data(tbl; kw...) + sch, batches = _exportbatches(tbl; kw...) + # An empty first partition only fixed the schema. + length(batches) > 1 && filter!(b -> b.nrows > 0, batches) + length(batches) == 1 || throw( + ArgumentError( + "the table needs $(length(batches)) record batches; export it with " * + "Arrow.export_stream!", + ), + ) + return to_c_data(_batchfield(sch), _structarray(only(batches))) +end + +""" + Arrow.export_stream!(streamptr::Ptr{Arrow.CArrowArrayStream}, table; kw...) + +Fill the `ArrowArrayStream` the caller allocated at `streamptr` with the record batches +of a Tables.jl table: one per `Tables.partitions` partition (an [`Arrow.Stream`](@ref) +is a partition per batch read), split only when `batchsize` is given or a batch's string +data outgrows 32-bit offsets. The batches are built as [`Arrow.write`](@ref) builds them, +with the same keywords, before this returns; the consumer pulls them through the stream +and releases it. Buffers cross without copying where `Arrow.write` would not copy them. +Memory stays reachable until the consumer has released the stream and every batch it +pulled. +""" +function export_stream!(sp::Ptr{CArrowArrayStream}, tbl; kw...) + sch, batches = _exportbatches(tbl; kw...) + return export_stream!(sp, sch, AC.RecordBatch[b for b in batches if b.nrows > 0]) +end diff --git a/src/columnconstruction.jl b/src/columnconstruction.jl new file mode 100644 index 00000000..2bb03fb1 --- /dev/null +++ b/src/columnconstruction.jl @@ -0,0 +1,4398 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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) + +# The Arrow unit of a Timestamp/ZonedTimestamp element type. Only the four +# built-in resolutions have an Arrow timestamp unit; Durations 1.3+ permits +# package-defined periods, which have no Arrow representation. +function _instantunit(T::Type) + P = _instantelementperiod(T) + P === Dates.Second && return AC.SECOND + P === Dates.Millisecond && return AC.MILLISECOND + P === Dates.Microsecond && return AC.MICROSECOND + P === Dates.Nanosecond && return AC.NANOSECOND + throw( + ArgumentError( + "timestamp period $P has no Arrow timestamp unit; use Second, " * + "Millisecond, Microsecond, or Nanosecond resolution", + ), + ) +end +_instantelementperiod(::Type{Durations.Timestamp{P}}) where {P} = P +_instantelementperiod(::Type{<:Durations.ZonedTimestamp{P}}) where {P} = P + +"One native (Field, ArrayData) column from a Julia vector, facade conversions included." +# The column-level dispatchers (`_constructnativepart`, `_constructpart`) do +# not specialize on the vector type: the builders they call do, so each new +# column type compiles its builder and not the whole dispatch. Each keyword +# method calls its body through `_uninferred`, so a method another package adds +# under a body makes only that body compile again, not every column type's +# keyword method. +Base.@nospecializeinfer _constructnativepart( + name::String, + @nospecialize(v::AbstractVector); + context=nothing, +) = _uninferred(_nativepart, name, v, context)::Tuple{AC.Field,AC.ArrayData} + +Base.@nospecializeinfer function _nativepart( + name::String, + @nospecialize(v::AbstractVector), + context, +) + T = Base.nonmissingtype(_eltypeof(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. + if _zonedwritertype(T) + zoned = _uninferred(_zonednativepart, name, v, T) + zoned === nothing || return zoned + end + if _eltypeof(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 $(_errname(context, name)) has no values to infer an Arrow " * + "type from; give the empty column a concrete element type " * + "(Int64[] or String[], say)", + ), + ) + elseif T <: Union{DataDecimals.Decimal,Durations.Duration} + return _uninferred(_constructsharedpart, name, v, T) + elseif T <: Durations.Timestamp + isconcretetype(T) || throw( + ArgumentError( + "column $(_errname(context, name)) has abstract Timestamp element " * + "type $T; give the column one concrete resolution", + ), + ) + return _uninferred( + _constructtemporalpart, + name, + v, + AC.TimestampType(_instantunit(T), nothing), + x -> Int64(Dates.value(x)), + ) + elseif T <: Durations.ZonedTimestamp + isconcretetype(T) || throw( + ArgumentError( + "column $(_errname(context, name)) has abstract ZonedTimestamp " * + "element type $T; an Arrow column declares one time zone and one " * + "resolution — " * + "convert the values with Durations.astimezone first", + ), + ) + return _uninferred( + _constructtemporalpart, + name, + v, + AC.TimestampType(_instantunit(T), Durations.zonename(T)::String), + x -> Int64(Dates.value(x)), + ) + elseif T <: Dates.Date + return _uninferred( + _constructtemporalpart, + name, + v, + AC.DateType(AC.DAY), + x -> Int32(Dates.value(x) - _EPOCH_DAYS), + ) + elseif T <: Dates.DateTime + return _uninferred( + _constructtemporalpart, + name, + v, + AC.TimestampType(AC.MILLISECOND, nothing), + x -> Int64(Dates.value(x) - Dates.UNIXEPOCH), + ) + elseif T <: Dates.Time + return _uninferred( + _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 _uninferred( + _constructtemporalpart, + name, + v, + AC.DurationType(unit), + x -> Int64(Dates.value(x)), + ) + elseif T <: Union{Dates.Week,Dates.Day,Dates.Hour,Dates.Minute} + return _uninferred( + _constructtemporalpart, + name, + v, + AC.DurationType(AC.SECOND), + _periodseconds, + ) + elseif T <: Union{Dates.Month,Dates.Year} + return _uninferred( + _constructtemporalpart, + name, + v, + AC.IntervalType(AC.YEAR_MONTH), + _periodmonths, + ) + elseif T <: NamedTuple + _uninferred(any, ismissing, v)::Bool && throw( + ArgumentError( + "column $name has nullable NamedTuple rows that require the " * + "recursive column adapter", + ), + ) + isconcretetype(T) || throw( + ArgumentError( + "column $(_errname(context, name)) has abstract NamedTuple element " * + "type; give it a concrete NamedTuple type with declared field names " * + "and types", + ), + ) + # Only the empty NamedTuple reaches here: `_nestedkind` takes every + # other NamedTuple column. + return _structpart(name, v, context) + elseif T <: AbstractString + _stringview(context) && return _uninferred( + _retainedview, + AC.Field(name, AC.ViewType(true); nullable=Missing <: _eltypeof(v)), + v, + ) + return _uninferred(_varbytespart, name, v, AC.Utf8Type(_largelists(context))) + elseif T <: AbstractVector{UInt8} + _stringview(context) && return _uninferred( + _retainedview, + AC.Field(name, AC.ViewType(false); nullable=Missing <: _eltypeof(v)), + v, + ) + return _uninferred(_varbytespart, name, v, AC.BinaryType(_largelists(context))) + 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 = _uninferred(_narrowlists, v, _errname(context, name)) + NW = Base.nonmissingtype(_eltypeof(w)) + (NW === Any || (NW <: AbstractVector && eltype(NW) === Any)) && throw( + ArgumentError( + "column $(_errname(context, 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) + elseif T <: AbstractVector + context === nothing && (context = _WriterContext()) + return _uninferred(_listpart, name, v, context) + elseif T === Bool + return _uninferred(_boolpart, name, v) + elseif T <: _FixedPrimitive + t = AC.arrowtype_for(T) + if _eltypeof(v) === T + data = _uninferred(_fixedbuffer, v) + data === nothing && return _uninferred(AC.fromjulia, name, _plainvector(v)) + return AC.Field(name, t; nullable=false), + AC.ArrayData(t, length(v), [AC.BufferSlice(), data]; nullcount=0) + end + vals, validity, nc = _uninferred(_nullablevalues, T, v, zero(T))::Tuple + return AC.Field(name, t; nullable=true), + AC.ArrayData(t, length(v), [validity, AC._databuffer(vals)]; nullcount=nc) + else + # The vector's type is known at run time; see `_uninferred`. + return _uninferred(AC.fromjulia, name, _plainvector(v)) + end +end + +const _FixedPrimitive = + Union{Int8,Int16,Int32,Int64,UInt8,UInt16,UInt32,UInt64,Float16,Float32,Float64} + +"A Bool column: values and validity bit-packed in one pass." +function _boolpart(name::String, v::AbstractVector) + n = length(v) + i0 = firstindex(v) - 1 + values = Vector{UInt8}(undef, cld(n, 8)) + bits = Vector{UInt8}(undef, cld(n, 8)) + valid = 0 + @inbounds for b in eachindex(bits) + base = (b - 1) << 3 + byte = 0x00 + valuebyte = 0x00 + for j = 0:(min(8, n - base) - 1) + x = v[i0 + base + j + 1] + present = x !== missing + byte |= UInt8(present) << (j % UInt) + valuebyte |= UInt8(present && x::Bool) << (j % UInt) + end + bits[b] = byte + values[b] = valuebyte + valid += count_ones(byte) + end + nc = n - valid + t = AC.BoolType() + validity = nc == 0 ? AC.BufferSlice() : AC._databuffer(bits) + return AC.Field(name, t; nullable=Missing <: _eltypeof(v)), + AC.ArrayData(t, n, [validity, AC._databuffer(values)]; nullcount=nc) +end + +""" +The values of a non-null fixed-width column as one buffer: a `Vector`, or a +contiguous view of one (a record batch's rows), is borrowed without a copy; +nothing for any other vector or a nullable element type. +""" +function _fixedbuffer(v::AbstractVector{T}) where {T} + isbitstype(T) || return nothing + v isa Vector{T} && return AC._databuffer(v) + v isa SubArray{T,1,Vector{T},<:Tuple{AbstractUnitRange{Int}},true} || return nothing + offset = (first(parentindices(v)[1]) - 1) * sizeof(T) + return AC.BufferSlice(AC.heapregion(parent(v)), offset, length(v) * sizeof(T)) +end + +"One (Field, ArrayData) column, including the ArrowTypes compatibility adapter." +Base.@nospecializeinfer _constructpart( + name::String, + @nospecialize(v::AbstractVector); + context=nothing, + narrowabstract::Bool=true, +) = _uninferred(_partof, name, v, context, narrowabstract)::Tuple{AC.Field,AC.ArrayData} + +Base.@nospecializeinfer function _partof( + name::String, + @nospecialize(v::AbstractVector), + context, + narrowabstract::Bool, +) + context === nothing && (context = _WriterContext()) + if _eltypeof(v) === Any + narrowed = _uninferred(_narrowlists, v, _errname(context, name)) + _eltypeof(narrowed) === Any && throw( + ArgumentError( + "column $(_errname(context, 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 + nested = _uninferred(_nestedpart, name, v, context) + nested === nothing || return nested + if narrowabstract + narrowed = _uninferred(_narrowabstractcolumn, name, v, context) + if narrowed !== nothing + # 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. + _eltypeof(narrowed) === _eltypeof(v) || + return _constructpart(name, narrowed; context, narrowabstract=false) + end + end + # The builders compile for the column's own type (see `_uninferred`). + return _uninferred(_arrowtypes_needs, v, context)::Bool ? + _uninferred(_constructarrowtypespart, name, v; context) : + _constructnativepart(name, v; context) +end + +# --- plain nested columns ------------------------------------------------------- +# +# Lists (rows are AbstractVectors), fixed-size lists (NTuples of one element +# type), structs (NamedTuples), and maps (AbstractDicts or AbstractVectors of +# Pairs) build here in a few typed passes, at any depth and with missing at +# any level. Their children build as columns do, so leaves keep every +# mapping. A container with an ArrowTypes.jl mapping or label of its own takes +# the general route. + +# How an error names a column: a nested child also names its top-level column. +_errname(context, name) = + context === nothing || isempty(context.column) || context.column == name ? name : + "$(context.column) (child $name)" + +"The layout `T` writes as (`:list`, `:fixedlist`, `:struct`, `:map`), or nothing." +function _nestedkind(context, T::Type) + Base.@nospecialize T + T isa DataType && T <: Union{AbstractVector,AbstractDict,Tuple,NamedTuple} || + return nothing + T <: AbstractVector{UInt8} && return nothing + _writerstoragetype!(context, T) === T && _writerextension!(context, T) === nothing || + return nothing + if T <: NamedTuple + return fieldcount(T) > 0 ? :struct : nothing + elseif T <: Tuple + # NTuple{N,UInt8} is fixed_size_binary; tuples of mixed types are structs. + Base.isvatuple(T) && return nothing + types = fieldtypes(T) + (isempty(types) || !allequal(types) || types[1] === UInt8) && return nothing + return :fixedlist + end + T <: AbstractDict && return :map + E = eltype(T) + return E !== Union{} && E <: Pair ? :map : :list +end + +""" +`v` as a plain nested column, or nothing when it is not one. Nothing too for +a struct or fixed-size list with missing rows whose non-nullable children +cannot hide placeholder slots (see `_hidevalidity`). +""" +function _nestedpart(name::String, @nospecialize(v::AbstractVector), context) + kind = _nestedkind(context, Base.nonmissingtype(_eltypeof(v))) + kind === nothing && return nothing + # Dispatching on the kind compiles only the matching builder for `v`'s type. + return _nestedbuild(Val(kind), name, v, context) +end +_nestedbuild(::Val{:list}, name, @nospecialize(v), context) = _listpart(name, v, context) +_nestedbuild(::Val{:map}, name, @nospecialize(v), context) = _mappart(name, v, context) +_nestedbuild(::Val{:struct}, name, @nospecialize(v), context) = + _structpart(name, v, context) +_nestedbuild(::Val{:fixedlist}, name, @nospecialize(v), context) = + _fixedlistpart(name, v, context) + +""" +List offsets (row `i` holds values `offsets[i]+1:offsets[i+1]` of the +child), validity, and null count of `v`'s rows, in one pass. +""" +function _listoffsets(name::String, v::AbstractVector, ::Type{O}, pool=nothing) where {O} + n = length(v) + i0 = firstindex(v) - 1 + offsets = _takevector(pool, O, n + 1) + offsets[1] = 0 + bits = _takevector(pool, UInt8, cld(n, 8)) + pos = 0 + valid = 0 + @inbounds for b in eachindex(bits) + base = (b - 1) << 3 + byte = 0x00 + for j = 0:(min(8, n - base) - 1) + i = base + j + 1 + x = v[i0 + i] + present = x !== missing + pos += present ? length(x) : 0 + pos <= typemax(O) || throw(_OffsetOverflow(name, 0, i)) + offsets[i + 1] = pos + byte |= UInt8(present) << (j % UInt) + end + bits[b] = byte + valid += count_ones(byte) + end + nc = n - valid + validity = + nc == 0 ? AC.BufferSlice() : AC.BufferSlice(AC.heapregion(bits), 0, length(bits)) + return offsets, validity, nc +end + +# A list or map column read from Arrow keeps its offsets, rebased to start at +# zero: a null row may span child values, which stay in the child. +function _listoffsets( + name::String, + @nospecialize(v::Union{ListVector,MapVector}), + ::Type{O}, + pool=nothing, +) where {O} + n = length(v) + offsets = _takevector(pool, O, n + 1) + # A read column's offsets may view its batch (Int32 or Int64): kept alive + # through `v`. + GC.@preserve v _rebaseoffsets!(offsets, v.offsets, name) + validity, nc = _storedvalidity(v.validity, v.bitoffset, n, pool) + return offsets, validity, nc +end + +function _rebaseoffsets!(offsets::AbstractVector{O}, raw, name::String) where {O} + lo = raw[1] + @inbounds for i in eachindex(offsets) + x = raw[i] - lo + x <= typemax(O) || throw(_OffsetOverflow(name, 0, i - 1)) + offsets[i] = x + end + return offsets +end + +""" +The validity of `n` rows a read column stores as `bits` from `bitoffset` (no +bits: no nulls), as a write buffer, and the null count. +""" +function _storedvalidity(bits::Vector{UInt8}, bitoffset::Int, n::Int, pool=nothing) + isempty(bits) && return AC.BufferSlice(), 0 + out = _takevector(pool, UInt8, cld(n, 8)) + valid = 0 + @inbounds for b in eachindex(out) + byte = 0x00 + for j = 0:(min(8, n - ((b - 1) << 3)) - 1) + byte |= UInt8(_bitset(bits, bitoffset, ((b - 1) << 3) + j + 1)) << j + end + out[b] = byte + valid += count_ones(byte) + end + nc = n - valid + return nc == 0 ? AC.BufferSlice() : AC.BufferSlice(AC.heapregion(out), 0, length(out)), + nc +end + +""" +Elements `lo+1:hi` of a read column's child values without copying them: a +view, or for a nested child the same storage over those rows. +""" +_childslice(v::AbstractVector, lo::Int, hi::Int) = + lo == 0 && hi == length(v) ? v : view(v, (lo + 1):hi) +function _childslice(@nospecialize(v::StructVector), lo::Int, hi::Int) + lo == 0 && hi == length(v) && return v + return typeof(v)( + v.children, + v.offset + lo, + hi - lo, + v.validity, + v.bitoffset + lo, + v.owner, + ) +end +function _childslice(@nospecialize(v::ListVector), lo::Int, hi::Int) + lo == 0 && hi == length(v) && return v + return typeof(v)( + GC.@preserve(v, v.offsets[(lo + 1):(hi + 1)]), + v.values, + v.validity, + v.bitoffset + lo, + v.owner, + ) +end +function _childslice(@nospecialize(v::MapVector), lo::Int, hi::Int) + lo == 0 && hi == length(v) && return v + offsets = GC.@preserve v v.offsets[(lo + 1):(hi + 1)] + return typeof(v)(offsets, v.keys, v.values, v.validity, v.bitoffset + lo, v.owner) +end + +# A row whose length differs from the first pass's would write past its slots. +_changedrows(name) = ArgumentError("column $name changed while it was being written") + +# A list column read from Arrow already holds its rows' values in order, as a +# child of element type `E`; its values are that child, sliced. +_storedlist(@nospecialize(v), E) = v isa ListVector && _eltypeof(v.values) === E +function _storedlistvalues(name::String, @nospecialize(v::ListVector), offsets) + lo, hi = GC.@preserve v Int(v.offsets[1]), Int(v.offsets[end]) + hi - lo == offsets[end] || throw(_changedrows(name)) + return _childslice(v.values, lo, hi) +end + +"The values of `v`'s rows back to back, as `_listoffsets` counted them." +function _listvalues(name::String, ::Type{E}, v::AbstractVector, offsets) where {E} + _storedlist(v, E) && return _storedlistvalues(name, v, offsets) + flat = Vector{E}(undef, Int(offsets[end])) + i0 = firstindex(v) - 1 + @inbounds for i = 1:(length(offsets) - 1) + x = v[i0 + i] + x === missing && continue + lo = Int(offsets[i]) + m = Int(offsets[i + 1]) - lo + length(x) == m || throw(_changedrows(name)) + copyto!(flat, lo + 1, x, firstindex(x), m) + end + return flat +end + +# A nested builder's `spec` is a fresh column's name, or the retained Field it +# rebuilds, whose descriptor, child names, nullability, and metadata it keeps. +_specname(spec::String) = spec +_specname(f::AC.Field) = f.name +_childspec(::String, k::Int, name::String) = name +_childspec(f::AC.Field, k::Int, ::String) = f.children[k] + +# The Field a nested builder returns over its built children. +_nestedfield(name::String, t, v, children) = + AC.Field(name, t; nullable=Missing <: _eltypeof(v), children) +_nestedfield(f::AC.Field, t, v, children) = _retainedfield(f; children) + +# Null rows under a retained field that cannot be null mean the column was +# replaced. +_checknulls(::String, nc) = nothing +_checknulls(f::AC.Field, nc) = nc == 0 || f.nullable || throw(_nonnullableerror(f.name)) + +# A retained column's rows are of its layout's kind `T`. +function _checkrows(spec, v, ::Type{T}, what::String) where {T} + spec isa AC.Field || return nothing + Base.nonmissingtype(_eltypeof(v)) <: T && return nothing + for x in v + x === missing || + x isa T || + throw(ArgumentError("column $(spec.name) retained rows must be $what")) + end + return nothing +end + +function _listpart(spec, @nospecialize(v::AbstractVector), context) + name = _errname(context, _specname(spec)) + t = spec isa AC.Field ? spec.type : AC.ListType(_largelists(context)) + _checkrows(spec, v, AbstractVector, "vectors") + offsets, validity, nc = _listoffsets(name, v, t.large ? Int64 : Int32, _pool(context)) + _checknulls(spec, nc) + E = eltype(Base.nonmissingtype(_eltypeof(v))) + values = + _storedlist(v, E) ? _storedlistvalues(name, v, offsets) : + _listvalues(name, E, v, offsets) + cf, cd = _constructpart(_childspec(spec, 1, "item"), values; context) + # A list view's rows are consecutive spans of its child: starts and sizes. + buffers = + t isa AC.ListViewType ? + [validity, AC._databuffer(offsets[1:(end - 1)]), AC._databuffer(diff(offsets))] : + [validity, AC._databuffer(offsets)] + d = AC.ArrayData(t, length(v), buffers; children=AC.ArrayData[cd], nullcount=nc) + return _nestedfield(spec, t, v, AC.Field[cf]), d +end + +""" +Child `K` of each row of `rows` (a struct field, a map entry's key or value) +as a lazy vector of `T`; a missing row reads as missing. +""" +struct _ChildValues{T,K,V<:AbstractVector} <: AbstractVector{T} + rows::V +end +_ChildValues{T,K}(rows::V) where {T,K,V} = _ChildValues{T,K,V}(rows) +Base.size(c::_ChildValues) = (length(c.rows),) +Base.@propagate_inbounds function Base.getindex(c::_ChildValues{T,K}, i::Int) where {T,K} + x = c.rows[firstindex(c.rows) + i - 1] + return (x === missing ? missing : getfield(x, K))::T +end + +# Element `j` of a fixed-size list column's child: value `(j-1) % N + 1` of +# row `(j-1) ÷ N + 1`. +struct _TupleValues{T,N,V<:AbstractVector} <: AbstractVector{T} + rows::V +end +_TupleValues{T,N}(rows::V) where {T,N,V} = _TupleValues{T,N,V}(rows) +Base.size(c::_TupleValues{T,N}) where {T,N} = (N * length(c.rows),) +Base.@propagate_inbounds function Base.getindex(c::_TupleValues{T,N}, j::Int) where {T,N} + r, k = divrem(j - 1, N) + x = c.rows[firstindex(c.rows) + r] + return (x === missing ? missing : x[k + 1])::T +end + +""" +Whether a non-nullable child of type `T` can sit under missing parent rows: +it is built nullable, with missing in those slots, and `_hidevalidity` then +drops its validity, so the slots hold placeholders (zeros, empty strings, +empty lists). That needs a layout with a validity buffer that a nullable +build of `T` keeps: leaves without an ArrowTypes.jl mapping and the plain +nested kinds. +""" +_hideable(context, T::Type) = + T !== Missing && + T !== Union{} && + !(T isa Union) && + ( + _nestedkind(context, T) !== nothing || + !_writerneedstype!(context, T) || + _hideablestorage(context, T) + ) + +# A mapped type whose declared storage is one concrete type (not Null, not a +# storage union, not inferred from the values) builds with a validity buffer. +function _hideablestorage(context, T::Type) + T === VersionNumber && return true + S = _writerstoragetype!(context, T) + S === T && return false + S isa Union && return false + # Null storage (`Nothing`'s is `Missing`) has no validity to hide. + (S === Nothing || S === Missing) && return false + isconcretetype(S) || return false + return _writerneedstype!(context, S) ? _hideablestorage(context, S) : true +end + +""" +Data built for `f` declared nullable (or for a nullable fresh child) as `f` +declares it: where `f` cannot be null, the validity that marked only hidden +slots is dropped, through a union's children and run-end values too. Hidden +slots are never read, so they keep whatever placeholder the build wrote. +""" +function _hidevalidity(f::AC.Field, built::AC.Field, d::AC.ArrayData) + t = d.type + fields = collect(AC.Field, built.children) + datas = collect(AC.ArrayData, d.children) + buffers = collect(AC.BufferSlice, d.buffers) + nc = AC.nullcount(d) + if t isa Union{AC.UnionType,AC.RunEndEncodedType} + for i in (t isa AC.UnionType ? eachindex(datas) : 2:2) + fields[i], datas[i] = _hidevalidity(f.children[i], fields[i], datas[i]) + end + elseif !f.nullable && !(t isa AC.NullType) + buffers[1] = AC.BufferSlice() + nc = 0 + end + field = AC.Field( + built.name, + built.type; + nullable=f.nullable, + metadata=_fieldmetadata(built), + children=fields, + ) + data = AC.ArrayData( + t, + d.len, + buffers; + offset=d.offset, + children=datas, + dictionary=d.dictionary, + owner=d.owner, + nullcount=nc, + ) + return field, data +end + +""" +Whether field `f` holds a null. A Null field with an extension label (such as +`JuliaLang.Nothing`) holds a value. A union holds nulls in the child +`_unionnullchild` names, and a run-end column in its values: neither has +validity of its own. +""" +function _holdsnull(f::AC.Field) + t = f.type + t isa AC.UnionType && return _unionnullchild(f) !== nothing + t isa AC.RunEndEncodedType && return _holdsnull(last(f.children)) + t isa AC.NullType && return _arrowtypesextension(f) === nothing + return f.nullable +end + +""" +The child of union field `f` that takes a missing value, or nothing: a Null +child without an extension label, else, in a union without a label of its +own, a child of another type that holds a null. Every child of a labelled +union holds its storage values. +""" +function _unionnullchild(f::AC.Field) + k = findfirst(c -> c.type isa AC.NullType && _holdsnull(c), f.children) + k === nothing && + _arrowtypesextension(f) === nothing && + (k = findfirst(c -> !(c.type isa AC.NullType) && _holdsnull(c), f.children)) + return k +end + +# Missing values under a union or Null field that holds no null. +function _nonullerror(f::AC.Field) + f.type isa AC.NullType && return ArgumentError( + "column $(f.name) holds missing values, but it writes as Arrow Null labelled " * + "$(first(_arrowtypesextension(f))), whose slots read as that label's value", + ) + return ArgumentError( + "column $(f.name) holds missing values, but no child of the union it writes " * + "as holds a null: a missing value needs a Null child without an extension " * + "label or, in a union without a label of its own, a nullable child", + ) +end + +""" +Retained child `f` over `values`, `hidden` of which sit under null parent +rows and read as missing. A child that holds nulls takes them there. In a +child that cannot, any other missing value means the column was replaced, +and the hidden slots take a stand-in: a dense union points each at the slot +of the visible row before it; a sparse union, run-end, or Null column, which +has no validity buffer, takes the visible value before it; any other child is +built nullable and its validity then dropped (`_hidevalidity`). +""" +function _hiddenchild(f::AC.Field, values, hidden::Int, context) + hidden == 0 && return _constructpart(f, values; context) + hidden == length(values) && return _constructhiddenpart(f, hidden, context) + _holdsnull(f) && return _constructpart(f, values; context) + count(ismissing, values) == hidden || throw(_nonnullableerror(f.name)) + if f.type isa AC.UnionType && f.type.mode == AC.DenseMode + return _uninferred( + _hiddendenseunion, + f, + values, + context, + )::Tuple{AC.Field,AC.ArrayData} + end + # ponytail: a sparse union still copies a value into each hidden slot + # (sparse unions come only from other writers' schemas); a child's empty + # value there would cost nothing. + if f.type isa Union{AC.UnionType,AC.RunEndEncodedType,AC.NullType} + return _constructpart(f, _uninferred(_standins, values)::AbstractVector; context) + end + relaxed = AC.Field( + f.name, + f.type; + nullable=true, + metadata=_fieldmetadata(f), + children=collect(AC.Field, f.children), + ) + return _hidevalidity(f, _constructpart(relaxed, values; context)...) +end + +# `values` with each missing value replaced by the value before it (the first +# value, before any), so a run-end column's hidden slots extend a run. +function _standins(values) + x = values[findfirst(!ismissing, values)] + return [y === missing ? x : (x = y) for y in values] +end + +# A dense union over its visible values, whose hidden rows share the slot of +# the visible row before them (the first one's, before any): no value is +# copied for them, and each child's offsets stay non-decreasing. +function _hiddendenseunion(f::AC.Field, values, context) + uf, ud = _constructpart(f, [x for x in values if x !== missing]; context) + src = max.(cumsum(x !== missing for x in values), 1) + ids = reinterpret(Int8, AC.slicebytes(ud.buffers[1]))[src] + offsets = reinterpret(Int32, AC.slicebytes(ud.buffers[2]))[src] + buffers = AC.BufferSlice[AC._databuffer(ids), AC._databuffer(offsets)] + d = AC.ArrayData(ud.type, length(values), buffers; children=ud.children, nullcount=0) + return uf, d +end + +_nonnullableerror(name) = ArgumentError( + "column $name holds missing values but its retained field is non-nullable", +) + +""" +Child `spec` of nested rows `v`: `V{T}` over `v` gives its values, `hidden` of +which sit under null rows and read as missing. A retained child builds as +`_hiddenchild` says. A fresh child that cannot be null is built nullable and +its validity dropped when its layout keeps one (`_hideable`), else built under +the field an empty `T` column writes as, as a retained child is. +""" +function _rowchild(V::Type, v, spec, T::Type, hidden::Int, context) + hidden == 0 && return _constructpart(spec, V{T}(v); context) + spec isa AC.Field && return _hiddenchild(spec, V{Union{Missing,T}}(v), hidden, context) + Missing <: T && return _constructpart(spec, V{T}(v); context) + if _hideable(context, T) + f, d = _constructpart(spec, V{Union{Missing,T}}(v); context) + declared = AC.Field( + f.name, + f.type; + nullable=false, + metadata=_fieldmetadata(f), + children=f.children, + ) + return _hidevalidity(declared, f, d) + end + f, _ = _constructpart(spec, T[]; context) + return _hiddenchild(f, V{Union{Missing,T}}(v), hidden, context) +end + +function _structpart(spec, v::AbstractVector, context) + R = Base.nonmissingtype(_eltypeof(v)) + present = Bool[x !== missing for x in v] + nc = count(!, present) + _checknulls(spec, nc) + # Retained rows are typed NamedTuples with the field's names, or rows + # materialized under labels (ordered `Pair` vectors). + typed = R <: NamedTuple && isconcretetype(R) || !(spec isa AC.Field) + if spec isa AC.Field + typed ? _checkrownames(spec, R) : _validateretainedstructrows(spec, v) + end + fields = AC.Field[] + children = AC.ArrayData[] + for k = 1:(spec isa AC.Field ? length(spec.children) : fieldcount(R)) + f, d = + typed ? + _rowchild( + _ChildValues{C,k} where {C}, + v, + _childspec(spec, k, string(fieldname(R, k))), + fieldtype(R, k), + nc, + context, + ) : _hiddenchild(spec.children[k], _RowFields(v, k), nc, context) + push!(fields, f) + push!(children, d) + end + t = spec isa AC.Field ? spec.type : AC.StructType() + d = AC.ArrayData( + t, + length(v), + [nc == 0 ? AC.BufferSlice() : AC._bitmapbuffer(present)]; + children, + nullcount=nc, + ) + return _nestedfield(spec, t, v, fields), d +end + +# A struct column read from Arrow writes its children as stored, slots under +# null rows included, so no child needs hiding. +function _structpart(spec, @nospecialize(v::StructVector), context) + R = Base.nonmissingtype(_eltypeof(v)) + spec isa AC.Field && _checkrownames(spec, R) + n = length(v) + validity, nc = _storedvalidity(v.validity, v.bitoffset, n, _pool(context)) + _checknulls(spec, nc) + fields = AC.Field[] + children = AC.ArrayData[] + for k = 1:fieldcount(R) + child = _childslice(v.children[k], v.offset, v.offset + n) + f, d = _constructpart(_childspec(spec, k, string(fieldname(R, k))), child; context) + push!(fields, f) + push!(children, d) + end + t = spec isa AC.Field ? spec.type : AC.StructType() + d = AC.ArrayData(t, n, [validity]; children, nullcount=nc) + return _nestedfield(spec, t, v, fields), d +end + +# A typed struct row's field names must be the retained field's children's. +function _checkrownames(f::AC.Field, R::Type) + names = fieldnames(R) + length(names) == length(f.children) && + all(k -> string(names[k]) == f.children[k].name, eachindex(names)) || throw( + ArgumentError( + "column $(f.name) retained struct row has fields $names; expected " * + "$(Tuple(c.name for c in f.children))", + ), + ) + return nothing +end + +""" +A fixed-size list column of `N` values per row: `NTuple`s or vectors. A +retained one's rows may be `Array`s of its tensor shape. +""" +function _fixedlistpart(spec, v::AbstractVector, context, N::Int=0) + shape = spec isa AC.Field ? _tensorshape(spec) : nothing + shape === nothing || (v = _tensorstorage(v, shape)) + R = Base.nonmissingtype(_eltypeof(v)) + t = + spec isa AC.Field ? spec.type::AC.FixedSizeListType : + AC.FixedSizeListType(N == 0 ? fieldcount(R) : N) + N = t.listsize + # (Not `R <: NTuple{N,Any}`: that type has N parameters.) + if !(R isa DataType && R <: Tuple && !Base.isvatuple(R) && fieldcount(R) == N) + _checkrows(spec, v, Union{AbstractVector,Tuple}, "vectors") + for x in v + x === missing || + length(x) == N || + throw( + ArgumentError( + "column $(_specname(spec)) fixed-size list row has " * + "$(length(x)) values; expected $N", + ), + ) + end + end + present = Bool[x !== missing for x in v] + nc = count(!, present) + _checknulls(spec, nc) + V = _TupleValues{C,N} where {C} + cf, cd = _rowchild(V, v, _childspec(spec, 1, "item"), eltype(R), nc * N, context) + d = AC.ArrayData( + t, + length(v), + [nc == 0 ? AC.BufferSlice() : AC._bitmapbuffer(present)]; + children=AC.ArrayData[cd], + nullcount=nc, + ) + return _nestedfield(spec, t, v, AC.Field[cf]), d +end + +_maptypes(::Type{T}) where {T<:AbstractDict} = (keytype(T), valtype(T)) +_maptypes(::Type{T}) where {T<:AbstractVector} = + eltype(T) isa DataType && eltype(T) <: Pair ? fieldtypes(eltype(T)) : (Any, Any) +_maptypes(::Type) = (Any, Any) + +"The entries of `v`'s map rows back to back, as `_listoffsets` counted them." +function _mapentries(name::String, ::Type{P}, v::AbstractVector, offsets) where {P} + R = Base.nonmissingtype(_eltypeof(v)) + R <: AbstractVector && eltype(R) === P && return _listvalues(name, P, v, offsets) + entries = Vector{P}(undef, Int(offsets[end])) + i0 = firstindex(v) - 1 + @inbounds for i = 1:(length(offsets) - 1) + x = v[i0 + i] + x === missing && continue + pos = Int(offsets[i]) + length(x) == Int(offsets[i + 1]) - pos || throw(_changedrows(name)) + # A dictionary iterates its entries as key => value pairs too. + for p in x + p isa Pair || throw(ArgumentError("column $name map rows must hold Pairs")) + entries[pos += 1] = P(first(p), last(p)) + end + end + return entries +end + +# Child `k` (key or value) of a map's entries. +_entryspec(spec::String, k::Int, name::String) = name +_entryspec(f::AC.Field, k::Int, ::String) = only(f.children).children[k] + +function _mappart(spec, v::AbstractVector, context) + name = _errname(context, _specname(spec)) + _checkrows(spec, v, Union{AbstractVector,AbstractDict}, "Pair vectors or dictionaries") + R = Base.nonmissingtype(_eltypeof(v)) + K, V = _maptypes(R) + offsets, validity, nc = _listoffsets(name, v, Int32, _pool(context)) + _checknulls(spec, nc) + entries = _mapentries(name, Pair{K,V}, v, offsets) + KK = Base.nonmissingtype(K) + Missing <: K && + any(p -> first(p) === missing, entries) && + throw( + ArgumentError( + "column $name has a missing map key; Arrow map keys cannot be null, " * + "so drop those entries or replace the key", + ), + ) + kf, kd = + _constructpart(_entryspec(spec, 1, "key"), _ChildValues{KK,1}(entries); context) + vf, vd = + _constructpart(_entryspec(spec, 2, "value"), _ChildValues{V,2}(entries); context) + # Distinct keys that an ArrowTypes mapping stores alike would read back as + # repeated keys of one row. Keys stored as themselves keep a row's + # repeated keys, which Arrow maps (and `Arrow.MapValue` rows) allow. + keyspec = _entryspec(spec, 1, "key") + ( + R <: AbstractDict && _writerneedstype!(context, KK) || + keyspec isa AC.Field && _arrowtypesextension(keyspec) !== nothing + ) && _validatemapphysicalkeys(name, kf, kd, offsets) + return _mapcolumn(spec, v, offsets, validity, nc, length(entries), kf, kd, vf, vd) +end + +# A map column read from Arrow writes its offsets, keys, and values as stored. +function _mappart(spec, @nospecialize(v::MapVector), context) + name = _errname(context, _specname(spec)) + offsets, validity, nc = _listoffsets(name, v, Int32, _pool(context)) + _checknulls(spec, nc) + lo, hi = GC.@preserve v Int(v.offsets[1]), Int(v.offsets[end]) + kf, kd = + _constructpart(_entryspec(spec, 1, "key"), _childslice(v.keys, lo, hi); context) + vf, vd = + _constructpart(_entryspec(spec, 2, "value"), _childslice(v.values, lo, hi); context) + return _mapcolumn(spec, v, offsets, validity, nc, hi - lo, kf, kd, vf, vd) +end + +function _mapcolumn(spec, @nospecialize(v), offsets, validity, nc, nentries, kf, kd, vf, vd) + kf.nullable && throw( + ArgumentError( + "column $(_specname(spec)) has map keys that write as a nullable Arrow " * + "type; Arrow map keys cannot be null", + ), + ) + t = spec isa AC.Field ? spec.type::AC.MapType : AC.MapType(false) + t.keyssorted && _checksortedkeys(spec, kf, kd, offsets) + entriestype = AC.StructType() + children = AC.Field[kf, vf] + ef = + spec isa AC.Field ? _retainedfield(only(spec.children); children) : + AC.Field("entries", entriestype; nullable=false, children) + ed = AC.ArrayData( + entriestype, + nentries, + [AC.BufferSlice()]; + children=AC.ArrayData[kd, vd], + nullcount=0, + ) + d = AC.ArrayData( + t, + length(v), + [validity, AC._databuffer(offsets)]; + children=AC.ArrayData[ed], + nullcount=nc, + ) + return _nestedfield(spec, t, v, AC.Field[ef]), d +end + +# A retained map that declares sorted keys gets rows sorted by stored key. +function _checksortedkeys(f::AC.Field, kf::AC.Field, kd::AC.ArrayData, offsets) + for row = 1:(length(offsets) - 1) + lo = Int(offsets[row]) + 1 + hi = Int(offsets[row + 1]) + lo >= hi && continue + previous = AC.getvalue(kf, kd, Int64(lo)) + for index = (lo + 1):hi + key = AC.getvalue(kf, kd, 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 + return nothing +end + +""" +A retained Union column built from its values (public values, or lowered +storage values under a label), or nothing when some value fits none of the +union's children or several. A missing value goes to the child that holds +nulls (`_unionnullchild`), or is refused. Any other value goes to the child +its runtime type routes to (`_unionroute`); children of one layout cannot be +told apart and do not route. Each child then builds from its values as a +composite's child does, so values that lower to storage (dates, extension +types) lower there. +""" +function _constructfacadeunion(f::AC.Field, v::AbstractVector, context) + t = f.type::AC.UnionType + children = f.children + n = length(v) + length(t.typeids) == length(children) && n <= typemax(Int32) || return nothing + # Lowered values of a labelled union: member i of the storage Union is + # child i, as a fresh write orders the children (a nullable one's outer + # Null child among them). + S = Base.nonmissingtype(_eltypeof(v)) + members = + !context.lowered || !(S isa Union) ? Type[] : + Base.uniontypes( + length(children) > length(Base.uniontypes(S)) ? Union{Missing,S} : S, + ) + nullchild = _unionnullchild(f) + routes = IdDict{Type,Int}() + route = Vector{Int}(undef, n) + for (i, x) in enumerate(v) + k = if x === missing + nullchild === nothing && throw(_nonullerror(f)) + nullchild + else + get!(routes, typeof(x)) do + _unionroute(children, x, members, context) + end + end + k == 0 && return nothing + route[i] = k + end + sparse = t.mode == AC.SparseMode + typeids = Vector{Int8}(undef, n) + offsets = sparse ? nothing : Vector{Int32}(undef, n) + childvalues = [Any[] for _ in children] + for (i, x) in enumerate(v) + k = route[i] + typeids[i] = t.typeids[k] + if sparse + # A child's slots of other rows are hidden. + for j in eachindex(children) + push!(childvalues[j], j == k ? x : missing) + end + elseif x === missing && !isempty(childvalues[k]) && last(childvalues[k]) === missing + # Consecutive missing rows share one null slot. + offsets[i] = length(childvalues[k]) - 1 + else + offsets[i] = length(childvalues[k]) + push!(childvalues[k], x) + end + end + hidden = + sparse ? [count(!=(k), route) for k in eachindex(children)] : + zeros(Int, length(children)) + fields = AC.Field[] + datas = AC.ArrayData[] + for (k, c) in enumerate(children) + cf, cd = _hiddenchild(c, childvalues[k], hidden[k], context) + push!(fields, cf) + push!(datas, cd) + end + field = AC.Field( + f.name, + t; + nullable=f.nullable, + metadata=_fieldmetadata(f), + children=fields, + ) + buffers = + sparse ? AC.BufferSlice[AC._databuffer(typeids)] : + AC.BufferSlice[AC._databuffer(typeids), AC._databuffer(offsets)] + return field, AC.ArrayData(t, n, buffers; children=datas, nullcount=0) +end + +""" +The child of a union that takes value `x`: its storage member's child by +position when that child has the layout of a fresh column of `x`'s type, +labels included; else the one child with that layout; else, for a lowered +value, the one child whose storage domain holds it; else the one child of the +same layout family (`_samefamily`); else 0. +""" +function _unionroute(children, x, members, context) + probe = first(_constructpart("probe", typeof(x)[x]; context)) + i = findfirst(M -> x isa M, members) + i !== nothing && + i <= length(children) && + _unionfits(children[i], probe; exact=true) && + return i + matches = findall(c -> _unionfits(c, probe; exact=true), children) + # A lowered temporal value is its storage integer or DateTime. + isempty(matches) && + context.lowered && + ( + matches = findall( + c -> isempty(c.children) && x isa _arrowtypesprimitivebasetype(c), + children, + ) + ) + isempty(matches) && + (matches = findall(c -> _unionfits(c, probe; exact=false), children)) + return length(matches) == 1 ? only(matches) : 0 +end + +# A child takes a value whose fresh field has the child's layout (`exact`), or +# the same layout at another width (a largelists or stringview source, +# pyarrow's large_string), or, when the child is a union itself, that one of +# its children takes. Exact matches are tried first: an int64 child and a +# timestamp child both store Int64. +_unionfits(c::AC.Field, probe::AC.Field; exact::Bool) = + (exact ? _asretained(c, probe) !== nothing : _samefamily(c, probe)) || + c.type isa AC.UnionType && any(cc -> _unionfits(cc, probe; exact), c.children) + +# Fields of one layout family, label for label and child for child: the same +# type, or utf8 or binary at any width (offsets or views), a list, list view, +# or fixed-size list at any width, a timestamp without a zone (or Date64) at +# any unit, a duration or time at any unit, a decimal of one precision and +# scale at any width. +function _samefamily(f::AC.Field, built::AC.Field) + family = _layoutfamily(f.type) + ( + AC.typeequal(f.type, built.type) || + family !== nothing && family == _layoutfamily(built.type) + ) && + (f.nullable || !built.nullable) && + _extensionpairs(f) == _extensionpairs(built) && + length(f.children) == length(built.children) || return false + return all(p -> _samefamily(p[1], p[2]), zip(f.children, built.children)) +end +_layoutfamily(t::AC.ArrowType) = + t isa AC.Utf8Type || (t isa AC.ViewType && t.utf8) ? :utf8 : + t isa AC.BinaryType || t isa AC.ViewType ? :binary : + t isa Union{AC.ListType,AC.ListViewType,AC.FixedSizeListType} ? :list : + (t isa AC.DateType && t.unit != AC.DAY) || + (t isa AC.TimestampType && t.timezone === nothing) ? :timestamp : + t isa AC.DurationType ? :duration : + t isa AC.TimeType ? :time : + t isa AC.DecimalType ? (:decimal, t.precision, t.scale) : nothing + +""" +A freshly built field as the retained field `f` declares it, or nothing when +their layouts differ. Names, nullability, and metadata come from `f`; a fresh +field may be non-nullable where `f` is nullable. +""" +function _asretained(f::AC.Field, built::AC.Field) + AC.typeequal(f.type, built.type) && + (f.nullable || !built.nullable) && + _extensionpairs(f) == _extensionpairs(built) && + length(f.children) == length(built.children) || return nothing + children = AC.Field[] + for (fc, bc) in zip(f.children, built.children) + c = _asretained(fc, bc) + c === nothing && return nothing + push!(children, c) + end + return AC.Field( + f.name, + f.type; + nullable=f.nullable, + metadata=_fieldmetadata(f), + children, + ) +end + +# A field's extension name and metadata: its logical type. +_extensionpairs(f::AC.Field) = + f.metadata === nothing ? Pair{String,String}[] : + Pair{String,String}[ + kv for kv in f.metadata if + first(kv) == _EXTENSION_NAME_KEY || first(kv) == _EXTENSION_METADATA_KEY + ] + +"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 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 + +# An abstract element type no native builder writes as one column (`Real`, +# `Any`): such values are written as a dense Union of their runtime types. +_abstractnonnative(T) = + T !== Union{} && + !isconcretetype(T) && + !(T isa Union) && + !(T <: Union{AbstractString,AbstractVector,NamedTuple}) && + !_singletypefamily(T) + +# Families one Arrow column holds at a single concrete type (one resolution, +# zone, or precision): their abstract columns narrow only when every value +# has the same type, and refuse otherwise. +_singletypefamily(T) = + T <: Union{Durations.Timestamp,Durations.ZonedTimestamp,DataDecimals.Decimal} + +"The Union of the runtime types of column `name`'s values, `Missing` included." +function _observedtype(v::AbstractVector, name="") + types = Type[] + seen = Base.IdSet{Type}() + hasmissing = false + for x in v + if x === missing + hasmissing = true + continue + end + _pushobserved!(types, seen, typeof(x), name) + end + U = _writeruniontype(types) + return hasmissing ? Union{Missing,U} : U +end + +"Narrow an Any-eltype column, recovering list-of-T structure when present." +function _narrowlists(v::AbstractVector, name="") + # 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(_eltypeof(w)) + _abstractnonnative(NT) && return _collectas(_observedtype(w, name), w) + NT !== Union{} && NT <: AbstractVector || return w + # Join the inner element types (empties narrow to Union{} and would + # otherwise poison the join), then retype every inner vector. With only + # empty rows, their declared element types are all there is. + E = Union{} + for x in w + x === missing && continue + isempty(x) && continue + E = typejoin(E, eltype(x)) + end + if E === Union{} + for x in w + x === missing || (E = typejoin(E, eltype(x))) + end + end + E === Union{} && (E = Any) + # Rows of different element types: the elements' runtime types. + if _abstractnonnative(E) + values = collect(Iterators.flatten(x for x in w if x !== missing)) + isempty(values) || (E = _observedtype(values, name)) + end + E === Any && return w + hasm = _eltypeof(w) >: Missing + S = hasm ? Union{Missing,Vector{E}} : Vector{E} + return S[x === missing ? missing : convert(Vector{E}, x) for x in w] +end + +# A column read as a view over decoded Arrow buffers (one record batch's +# `ArrayData`) writes those buffers as they are when the writer would declare +# the same type for the column's element type: no value is converted. +function _constructpart( + name::String, + v::Union{Primitive,BoolVector}; + context=nothing, + narrowabstract::Bool=true, +) + d = v.owner + # Fixed-size binary rows (`NTuple{N,UInt8}`) and 128-bit integers (stored + # as fixed-size binary) have no native construction; the general path + # writes them. + if d isa AC.ArrayData && + d.offset == 0 && + d.len == length(v) && + !(Base.nonmissingtype(_eltypeof(v)) <: Union{Tuple,Int128,UInt128}) + f, _ = _constructnativepart(name, _eltypeof(v)[]; context) + AC.typeequal(f.type, d.type) && return f, _viewbuffers(d) + end + return invoke( + _constructpart, + Tuple{String,AbstractVector}, + name, + v; + context, + narrowabstract, + ) +end + +# A fixed-width or Bool batch's buffers cut to its length, so no byte past the +# last value (padding of the source) is written. +function _viewbuffers(d::AC.ArrayData) + n = d.len + databytes = d.type isa AC.BoolType ? cld(n, 8) : n * AC.primwidth(d.type) + data = AC.subslice(AC.rolebuffer(d, AC.DATA), 0, databytes) + nc = AC.nullcount(d) + validity = nc == 0 ? AC.BufferSlice() : AC.subslice(AC.validitybuffer(d), 0, cld(n, 8)) + return AC.ArrayData(d.type, n, [validity, data]; nullcount=nc) +end + +const _PayloadColumn = Union{DataStrings.StringVector,DataStrings.BytesVector} +# A record batch's rows of a DataStrings column. +const _PayloadRows = SubArray{<:Any,1,<:_PayloadColumn,<:Tuple{UnitRange{Int}}} +_payloadparent(v::_PayloadColumn) = v +_payloadparent(v::_PayloadRows) = parent(v) +_payloadentries(v::_PayloadColumn) = v.payloads +_payloadentries(v::_PayloadRows) = view(parent(v).payloads, only(parentindices(v))) + +# A DataStrings column IS Utf8View (BinaryView) memory: its payload vector is +# the views buffer and its byte buffers are the variadic data buffers. With +# `stringview=true` it is written as is, without a copy (a row range copies +# its view entries); otherwise its bytes are packed into utf8 (binary) +# offsets and data in one pass over the payloads. Nullability comes from the +# column's eltype. Both keywords are accepted for dispatch uniformity; no +# narrowing applies. +function _constructpart( + name::String, + v::Union{_PayloadColumn,_PayloadRows}; + context=nothing, + narrowabstract::Bool=true, +) + p = _payloadparent(v) + utf8 = p isa DataStrings.StringVector + nullable = Missing <: _eltypeof(v) + if _stringview(context) + # View entries are the payloads themselves, so a row range copies them. + payloads = v isa _PayloadRows ? collect(_payloadentries(v)) : p.payloads + # DataStrings keeps a private append arena, which may be empty. + lastbuffer = something(findlast(!isempty, p.buffers), 0) + f, d = AC.fromviewentries(name, payloads, p.buffers[1:lastbuffer]; nullable) + utf8 && return f, d + t = AC.ViewType(false) + return AC.Field(name, t; nullable), + AC.ArrayData(t, d.len, d.buffers; nullcount=AC.nullcount(d)) + end + large = _largelists(context) + t = utf8 ? AC.Utf8Type(large) : AC.BinaryType(large) + return _packedpayloadpart( + name, + _payloadentries(v), + p.buffers, + t, + large ? Int64 : Int32, + nullable, + _pool(context), + ) +end + +""" +utf8/binary offsets and data from DataStrings payloads. One pass measures the +values and copies them while every value so far is inline; a batch with an +out-of-line value is a slice of that value's buffer when every value's bytes +sit there in row order (as a read of utf8 data stores them), else one copy. +""" +function _packedpayloadpart( + name::String, + payloads, + buffers, + t::AC.ArrowType, + ::Type{Offset}, + nullable::Bool, + pool=nothing, +) where {Offset} + n = length(payloads) + offsets = _takevector(pool, Offset, n + 1) + offsets[1] = 0 + bits = _takevector(pool, UInt8, cld(n, 8)) + # Inline values are at most 12 bytes, so this holds every value until the + # first out-of-line one; a batch that starts out of line goes without. + firstlong = + n > 0 && DataStrings.payloadlength(@inbounds payloads[1]) > DataStrings.INLINE_MAX + inline = _takevector(pool, UInt8, firstlong ? 16 : 12n + 16) + # From the first out-of-line value on, the bytes are a slice of its buffer + # (`runbuf`, from `runstart`) when every out-of-line value sits there at + # its row's position and every inline value matches the bytes there, as in + # a column read from Arrow. This pass checks the values after the first + # out-of-line one; the ones copied to `inline` before it are compared once + # after. + runbuf = 0 + runstart = 0 + runok = true + prefix = 0 + run = Ptr{UInt8}(0) + room = 0 + mismatch = UInt64(0) + pos = 0 + valid = 0 + GC.@preserve payloads inline buffers begin + words = Ptr{UInt64}(pointer(payloads)) + out = pointer(inline) + @inbounds for b in eachindex(bits) + base = (b - 1) << 3 + byte = 0x00 + for j = 0:(min(8, n - base) - 1) + i = base + j + 1 + a = unsafe_load(words, 2i - 1) + w = unsafe_load(words, 2i) + len = Int(reinterpret(Int32, a % UInt32)) + present = len >= 0 + long = len > DataStrings.INLINE_MAX + # The payload's last 12 bytes, an inline value's bytes, as words. + lo = (a >> 32) | (w << 32) + hi = (w >> 32) % UInt32 + if runbuf == 0 + if long + runbuf = Int(reinterpret(Int32, w % UInt32)) + 1 + runstart = Int(reinterpret(Int32, (w >> 32) % UInt32)) - pos + prefix = pos + if 1 <= runbuf <= length(buffers) && runstart >= 0 + run = pointer(buffers[runbuf]) + runstart + room = length(buffers[runbuf]) - runstart + else + runok = false + end + else + unsafe_store!(Ptr{UInt64}(out + pos), lo) + unsafe_store!(Ptr{UInt32}(out + pos + 8), hi) + end + else + k = Int(reinterpret(Int32, w % UInt32)) + 1 + at = Int(reinterpret(Int32, (w >> 32) % UInt32)) - pos + # Without branches: the values are a random mix of kinds. + runok &= !long | ((k == runbuf) & (at == runstart)) + if pos + 12 <= room + m = clamp(len, 0, DataStrings.INLINE_MAX) + 1 + lomask = _INLINELOMASKS[m] + himask = _INLINEHIMASKS[m] + diff = + ((unsafe_load(Ptr{UInt64}(run + pos)) ⊻ lo) & lomask) | + ((unsafe_load(Ptr{UInt32}(run + pos + 8)) ⊻ hi) & himask) + mismatch |= ifelse(long, UInt64(0), diff) + elseif !long & (len > 0) + # Within 12 bytes of the buffer's end: byte by byte. + entry = Ptr{UInt8}(words) + 16 * (i - 1) + 4 + runok &= pos + len <= room && _memequal(run + pos, entry, len) + end + end + pos += ifelse(present, len, 0) + pos <= typemax(Offset) || throw(_OffsetOverflow(name, 0, i)) + offsets[i + 1] = pos % Offset + byte |= UInt8(present) << (j % UInt) + end + bits[b] = byte + valid += count_ones(byte) + end + runok &= (runbuf != 0) & (mismatch == 0) & (pos <= room) + runok && (runok = _memequal(pointer(inline), run, prefix)) + end + nc = n - valid + validity = + nc == 0 ? AC.BufferSlice() : AC.BufferSlice(AC.heapregion(bits), 0, length(bits)) + if runbuf == 0 + # The bytes are a slice: the vector keeps its length, which a reused + # vector's pool reads as its capacity. + databuf = + pos == 0 ? AC.BufferSlice() : AC.BufferSlice(AC.heapregion(inline), 0, pos) + d = AC.ArrayData(t, n, [validity, AC._databuffer(offsets), databuf]; nullcount=nc) + return AC.Field(name, t; nullable), _constructedvalid(d) + end + if runok + databuf = AC.BufferSlice(AC.heapregion(buffers[runbuf]), runstart, pos) + d = AC.ArrayData(t, n, [validity, AC._databuffer(offsets), databuf]; nullcount=nc) + return AC.Field(name, t; nullable), _constructedvalid(d) + end + data = _takevector(pool, UInt8, pos + 32) + GC.@preserve payloads buffers data begin + out = pointer(data) + words = Ptr{UInt64}(pointer(payloads)) + @inbounds for i = 1:n + a = unsafe_load(words, 2i - 1) + w = unsafe_load(words, 2i) + len = Int(reinterpret(Int32, a % UInt32)) + len <= 0 && continue + if len <= DataStrings.INLINE_MAX + # The inline bytes are the payload's last 12; `data` keeps 32 + # bytes of slack, so whole words are stored. + unsafe_store!(Ptr{UInt64}(out), (a >> 32) | (w << 32)) + unsafe_store!(Ptr{UInt32}(out + 8), (w >> 32) % UInt32) + else + buf = buffers[reinterpret(Int32, w % UInt32) + 1] + off = Int(reinterpret(Int32, (w >> 32) % UInt32)) + src = pointer(buf) + off + if len <= 32 && off + 32 <= length(buf) + for k = 0:8:24 + unsafe_store!( + Ptr{UInt64}(out + k), + unsafe_load(Ptr{UInt64}(src + k)), + ) + end + else + unsafe_copyto!(out, src, len) + end + end + out += len + end + end + databuf = pos == 0 ? AC.BufferSlice() : AC.BufferSlice(AC.heapregion(data), 0, pos) + d = AC.ArrayData(t, n, [validity, AC._databuffer(offsets), databuf]; nullcount=nc) + return AC.Field(name, t; nullable), _constructedvalid(d) +end + +# The masks of an inline value's first `len` bytes (0 to 12) in the two words +# `_packedpayloadpart` compares it as, at index `len + 1`. +const _INLINELOMASKS = ntuple(k -> (UInt64(1) << (8 * min(k - 1, 8))) - UInt64(1), 13) +const _INLINEHIMASKS = ntuple(k -> (UInt32(1) << (8 * clamp(k - 9, 0, 4))) - UInt32(1), 13) + +_memequal(p::Ptr{UInt8}, q::Ptr{UInt8}, n::Int) = + ccall(:memcmp, Cint, (Ptr{UInt8}, Ptr{UInt8}, Csize_t), p, q, n) == 0 + +_stringview(context) = context !== nothing && context.stringview +_pool(context) = context === nothing ? nothing : context.pool +_largelists(context) = context !== nothing && context.largelists + +""" +A 32-bit-offset column's data passes the Int32 range at `row` (0 when the row +is not known, as in nested data) of partition `part` (0 until the partition +loop records it). The writer cuts its record batches smaller and retries. +""" +struct _OffsetOverflow <: Exception + column::String + part::Int + row::Int +end +Base.showerror(io::IO, e::_OffsetOverflow) = + print(io, "column $(e.column) exceeds the Int32 offset range of one record batch") + +""" +A string of `column` (a child as "a (child b.c)") is not valid UTF-8, at `row` +(0 when no row refers to it) of part `part` of the values being built. The +batch and partition loops report it as an `ArgumentError` naming the row of +the table being written. +""" +struct _InvalidUTF8 <: Exception + column::String + part::Int + row::Int +end +Base.showerror(io::IO, e::_InvalidUTF8) = print(io, _invalidutf8message(e.column, e.row, 0)) + +function _invalidutf8message(column::String, row::Int, partition::Int) + at = row == 0 ? "" : partition == 0 ? " row $row" : " row $row of partition $partition" + return "column $column$at holds a string that is not valid UTF-8; Arrow strings " * + "must be UTF-8 (Polars and DuckDB refuse files with others), so write " * + "binary data as a Vector{UInt8} column" +end + +"Throw `_InvalidUTF8` for the first string of a built column that is not UTF-8." +function _checkutf8(name::Symbol, f::AC.Field, d::AC.ArrayData) + bad = _badutf8(f, d) + bad === nothing && return nothing + row, path = bad + column = isempty(path) ? String(name) : "$name (child $path)" + throw(_InvalidUTF8(column, 1, row)) +end + +""" +The first row of `d` holding a string, at any depth, that is not valid UTF-8, +and the path of child names to it, or nothing. Every non-null string slot is +checked, as Arrow's full validation does; row 0 means no row refers to the +slot (a dictionary value no row uses, say). +""" +function _badutf8(f::AC.Field, d::AC.ArrayData) + t = d.type + if t isa AC.Utf8Type + i = _badutf8offsets(d, t.large ? Int64 : Int32) + return i == 0 ? nothing : (i, "") + elseif t isa AC.ViewType + t.utf8 || return nothing + i = _badutf8views(d) + return i == 0 ? nothing : (i, "") + elseif t isa AC.DictionaryType + pool = d.dictionary + pool === nothing && return nothing + bad = _badutf8(AC.dictvaluefield(f, t), pool) + bad === nothing && return nothing + return (bad[1] == 0 ? 0 : _dictrow(d, t, bad[1]), bad[2]) + end + for (j, (cf, cd)) in enumerate(zip(f.children, d.children)) + bad = _badutf8(cf, cd) + bad === nothing && continue + path = isempty(bad[2]) ? cf.name : cf.name * "." * bad[2] + return (bad[1] == 0 ? 0 : _parentrow(d, j, bad[1]), path) + end + return nothing +end + +function _isascii(bytes::Vector{UInt8}) + acc = UInt64(0) + words = length(bytes) >> 3 + GC.@preserve bytes begin + p = Ptr{UInt64}(pointer(bytes)) + for i = 1:words + acc |= unsafe_load(p, i) + end + end + for i = (8words + 1):length(bytes) + acc |= @inbounds bytes[i] + end + return iszero(acc & 0x8080808080808080) +end + +# A utf8 column's first non-null row that is not UTF-8, or 0. ASCII data, or +# data valid as a whole whose rows all start on a character, is valid; only +# other data is checked row by row. `T` is the offset type. +function _badutf8offsets(d::AC.ArrayData, ::Type{T}) where {T} + n = Int(d.len) + n == 0 && return 0 + offs = AC.rolebuffer(d, AC.OFFSETS) + data = AC.rolebuffer(d, AC.DATA) + (d.offset + n + 1) * sizeof(T) <= offs.len || return 0 + GC.@preserve offs data begin + op = Ptr{T}(AC.sliceptr(offs)) + d.offset * sizeof(T) + lo = Int(unsafe_load(op, 1)) + hi = Int(unsafe_load(op, n + 1)) + 0 <= lo < hi <= data.len || return 0 + bytes = unsafe_wrap(Array, AC.sliceptr(data) + lo, hi - lo) + _isascii(bytes) && return 0 + if isvalid(String, bytes) + # A row starts on a character unless its first byte continues one; + # a row at the end (empty) starts anywhere. + starts = true + m = length(bytes) + for i = 2:n + o = Int(unsafe_load(op, i)) - lo + lead = @inbounds bytes[clamp(o + 1, 1, m)] + starts &= (o >= m) | (lead & 0xc0 != 0x80) + end + starts && return 0 + end + for i = 1:n + AC.isvalid_at(d, i) || continue + a = Int(unsafe_load(op, i)) - lo + b = Int(unsafe_load(op, i + 1)) - lo + isvalid(String, view(bytes, (a + 1):b)) || return i + end + end + return 0 +end + +# A utf8_view column's first non-null row that is not UTF-8, or 0. +function _badutf8views(d::AC.ArrayData) + n = Int(d.len) + views = AC.rolebuffer(d, AC.VIEWS) + 16 * (d.offset + n) <= views.len || return 0 + GC.@preserve d begin + entries = unsafe_wrap(Array, AC.sliceptr(views), views.len) + buffers = Vector{UInt8}[ + unsafe_wrap(Array, AC.sliceptr(b), b.len) for b in @view(d.buffers[3:end]) + ] + longascii = all(_isascii, buffers) + for i = 1:n + AC.isvalid_at(d, i) || continue + base = 16 * (d.offset + i - 1) + len = Int(unsafe_load(Ptr{Int32}(pointer(entries, base + 1)))) + bytes = if len <= AC.VIEW_INLINE_MAX + view(entries, (base + 5):(base + 4 + len)) + else + longascii && continue + k = Int(unsafe_load(Ptr{Int32}(pointer(entries, base + 9)))) + at = Int(unsafe_load(Ptr{Int32}(pointer(entries, base + 13)))) + view(buffers[k + 1], (at + 1):(at + len)) + end + isvalid(String, bytes) || return i + end + end + return 0 +end + +# The first row of a dictionary column whose index refers to pool value `k`. +function _dictrow(d::AC.ArrayData, t::AC.DictionaryType, k::Int) + w = AC.primwidth(t.indextype) + indices = AC.rolebuffer(d, AC.DATA) + for i = 1:Int(d.len) + AC.isvalid_at(d, i) || continue + AC._load_int(indices, t.indextype, AC._slotbyteoff(d, Int64(i), w)) + 1 == k && + return i + end + return 0 +end + +# The first row of `d` whose values include element `k` of its child `j`. +function _parentrow(d::AC.ArrayData, j::Int, k::Int) + t = d.type + for i = 1:Int(d.len) + k in _childspan(t, d, j, Int64(i)) && return i + end + return 0 +end + +function _childspan(t::AC.ArrowType, d::AC.ArrayData, j::Int, i::Int64) + if t isa Union{AC.ListType,AC.MapType} + lo, hi = AC._offsets_at(d, i, t isa AC.ListType && t.large) + return (lo + 1):hi + elseif t isa AC.ListViewType + off, sz = AC._listview_range(t, d, i) + return (off + 1):(off + sz) + elseif t isa AC.FixedSizeListType + base = AC._slotindex0(d, i) * t.listsize + return (base + 1):(base + t.listsize) + elseif t isa AC.StructType + return (d.offset + i):(d.offset + i) + elseif t isa AC.UnionType + tid = AC.loadat(AC.rolebuffer(d, AC.TYPE_IDS), Int8, AC._slotindex0(d, i)) + tid == t.typeids[j] || return 1:0 + t.mode == AC.SparseMode && return (d.offset + i):(d.offset + i) + at = + AC.loadat(AC.rolebuffer(d, AC.ELEMENT_OFFSETS), Int32, AC._slotbyteoff(d, i, 4)) + return (at + 1):(at + 1) + elseif t isa AC.RunEndEncodedType + j == 2 || return 1:0 + r = AC._ree_runindex(d, i) + return r:r + end + return 1:0 +end + +# Duration seconds and year-month interval months of the larger Dates periods. +_periodseconds(x::Dates.Week) = Base.checked_mul(Dates.value(x), Int64(604_800)) +_periodseconds(x::Dates.Day) = Base.checked_mul(Dates.value(x), Int64(86_400)) +_periodseconds(x::Dates.Hour) = Base.checked_mul(Dates.value(x), Int64(3_600)) +_periodseconds(x::Dates.Minute) = Base.checked_mul(Dates.value(x), Int64(60)) +_periodmonths(x::Dates.Month) = Dates.value(x) +_periodmonths(x::Dates.Year) = Base.checked_mul(Dates.value(x), Int64(12)) + +_nbytes(x::AbstractString) = ncodeunits(x) +_nbytes(x::AbstractVector{UInt8}) = length(x) +function _copybytes!( + data::Vector{UInt8}, + pos::Int, + x::Union{String,SubString{String}}, + n::Int, +) + GC.@preserve data x unsafe_copyto!(pointer(data, pos + 1), pointer(x), n) + return nothing +end +# A DataStrings value's bytes are inline in its 16-byte payload (after the +# 4-byte length) or at `payloadpos` in its buffer. +_copybytes!(data::Vector{UInt8}, pos::Int, x::DataStrings.DataString, n::Int) = + _copypayloadbytes!(data, pos, x, n) +_copybytes!(data::Vector{UInt8}, pos::Int, x::DataStrings.DataBytes, n::Int) = + _copypayloadbytes!(data, pos, x, n) +function _copypayloadbytes!(data::Vector{UInt8}, pos::Int, x, n::Int) + if n <= DataStrings.INLINE_MAX + r = Ref(x.p) + GC.@preserve r data begin + src = Ptr{UInt8}(Base.unsafe_convert(Ptr{DataStrings.StringPayload}, r)) + 4 + unsafe_copyto!(pointer(data, pos + 1), src, n) + end + else + GC.@preserve x data begin + src = pointer(x.data, DataStrings.payloadpos(x.p)) + unsafe_copyto!(pointer(data, pos + 1), src, n) + end + end + return nothing +end +_copybytes!(data::Vector{UInt8}, pos::Int, x::AbstractVector{UInt8}, n::Int) = + (copyto!(data, pos + 1, x, firstindex(x), n); nothing) +function _copybytes!(data::Vector{UInt8}, pos::Int, x::AbstractString, n::Int) + for (j, b) in enumerate(codeunits(x)) + @inbounds data[pos + j] = b + end + return nothing +end + +""" +A utf8 or binary column (`t`). One pass measures the values (offsets, +validity, and the offset-range check), a second copies their bytes into a +data buffer of the exact size. +""" +function _varbytespart(name::String, v::AbstractVector, t::AC.ArrowType) + Offset = t.large ? Int64 : Int32 + return _varbytespart(name, v, t, Offset) +end +function _varbytespart(name::String, v::AbstractVector, t, ::Type{Offset}) where {Offset} + n = length(v) + i0 = firstindex(v) - 1 + offsets = Vector{Offset}(undef, n + 1) + offsets[1] = 0 + bits = zeros(UInt8, cld(n, 8)) + pos = 0 + nc = 0 + @inbounds for i = 1:n + x = v[i0 + i] + if x === missing + nc += 1 + else + m = _nbytes(x)::Int + pos <= typemax(Offset) - m || throw(_OffsetOverflow(name, 0, i)) + pos += m + bits[((i - 1) >> 3) + 1] |= UInt8(1) << ((i - 1) & 7) + end + offsets[i + 1] = pos + end + data = Vector{UInt8}(undef, pos) + @inbounds for i = 1:n + x = v[i0 + i] + start = Int(offsets[i]) + m = Int(offsets[i + 1]) - start + # Every data byte is written exactly once, even if a lazy column + # yields different values on the second pass. + (x === missing ? m == 0 : _nbytes(x) == m) || + throw(ArgumentError("column $name changed while it was being written")) + x === missing || _copybytes!(data, start, x, m) + end + validity = + nc == 0 ? AC.BufferSlice() : AC.BufferSlice(AC.heapregion(bits), 0, length(bits)) + databuf = pos == 0 ? AC.BufferSlice() : AC._databuffer(data) + d = AC.ArrayData(t, n, [validity, AC._databuffer(offsets), databuf]; nullcount=nc) + return AC.Field(name, t; nullable=Missing <: _eltypeof(v)), _constructedvalid(d) +end + +# Offsets built here never decrease and end at the data length, and the null +# count is exact: the data passes Core's semantic checks by construction. +function _constructedvalid(d::AC.ArrayData) + @atomic :monotonic d.semachecked = true + return d +end + +"Concrete Vector with an exact Union{Missing,T} or T eltype for fromjulia." +function _plainvector(v::AbstractVector) + T = _eltypeof(v) + return v isa Vector{T} ? v : _collectas(T, v) +end + +# `eltype` of a vector the caller does not specialize on: Base's `eltype` of an +# abstract vector type has methods other packages add to (DataFrames does). +_eltypeof(::AbstractVector{T}) where {T} = T + +# A string's UTF-8 bytes. `codeunits` of an `AbstractString` has methods other +# packages add to (LaTeXStrings does); the bytes of a `String` do not. +_utf8units(x::AbstractString) = codeunits((x isa String ? x : String(x))::String) + +# `collect(T, v)` for an element type known only at run time. Other packages +# add `collect(::Type{...}, v)` methods (InlineStrings does), and writer code +# inferred against Base's one method would recompile once they load. +_collectas(::Type{T}, v::AbstractVector) where {T} = copyto!(Vector{T}(undef, length(v)), v) + +_missings_to(::Type{S}, v) where {S} = + _eltypeof(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 = _uninferred(AC.fromjulia, "storage", storage)::Tuple{AC.Field,AC.ArrayData} + # 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)::Int) + 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} + S = AC.primwidth(t) == 4 ? Int32 : Int64 + vals, validity, nc = _nullablevalues(S, v, zero(S), x -> S(tostorage(x))) + d = AC.ArrayData(t, length(v), [validity, AC._databuffer(vals)]; nullcount=nc) + return AC.Field(name, t; nullable=_eltypeof(v) >: Missing), d +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 => String(last(kv))::String for kv in base) + extra === nothing && return isempty(out) ? nothing : out + for kv in extra + key = String(first(kv))::String + value = String(last(kv))::String + 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, + DataDecimals.Decimal, + Durations.Duration, + Durations.Timestamp, + Durations.ZonedTimestamp, + } || _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) + T = _eltypeof(v) + _writerneedstype!(context, T) && return true + # Abstract/Any containers do not expose nested logical types. Inspect only + # until the first ArrowTypes value is found; concrete columns, and Unions + # of leaf types such as Union{Missing,Int64}, take no scan: the check + # above already decided every value they can hold. + return !isconcretetype(T) && + !_writerleafunion(T) && + any(value -> _arrowtypesneedsvalue(value, context), v) +end + +# A Union of concrete types the value walk does not descend into. +_writerleafunion(T) = + T isa Union && all(Base.uniontypes(T)) do S + isconcretetype(S) && + !(S <: Union{Tuple,NamedTuple,Pair,AbstractVector,AbstractDict}) + end + +_arrowtypesconcreteorunion(T) = + isconcretetype(T) || + (T isa Union && all(_arrowtypesconcreteorunion, Base.uniontypes(T))) + +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: " * + _uninferred(join, (string(S) for S in (path..., T)), " -> ")::String, + ), + ) + 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 _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} + +# Base's `convert(Missing, x)` raises a plain error; report a conversion error. +_writerconvert(::Type{Missing}, value) = + value === missing ? missing : throw(MethodError(convert, (Missing, value))) + +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, field=nothing) + value isa Union{UUIDs.UUID,Int128,UInt128} && return _facadetoarrow(value, field) + 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: $(_uninferred(sprint, showerror, err)::String)", + ), + ) + 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, + field=nothing, +) where {T} + S = Base.nonmissingtype(T) + # Primitive storage is read once by the builder, which lowers each value + # as it reads it: no lowered copy of the column. + S !== Union{} && + S <: Union{_FixedPrimitive,Bool} && + return _LoweredValues{T}(v) do x + lowered = _writertoarrow(context, x, writertype, name, field) + return _writerconvertlowered(S, lowered, context, writertype, name) + end + out = Vector{T}(undef, length(v)) + for (i, value) in enumerate(v) + lowered = _writertoarrow(context, value, writertype, name, field) + out[i] = _writerconvertlowered(T, lowered, context, writertype, name) + end + return out +end + +function _arrowtypesmappedcolumn(name::String, v::AbstractVector, declared, context) + S = _eltypeof(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 + _pushobserved!(observed, seen, typeof(mapped), name) + 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) ? + _eltypeof(_uninferred(ArrowTypes.ToArrow, v)::AbstractVector) : 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; context=nothing) + logical = Base.nonmissingtype(_eltypeof(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. + _eltypeof(v) === Missing && return _constructnativepart(name, v; context) + logical === Union{} && return _constructnativepart(name, v; context) + context === nothing && (context = _WriterContext()) + return _withwriterschema(context, logical) do + # `v`'s type is known at run time; see `_uninferred`. + _uninferred(_constructarrowtypespart_impl, name, v, logical; context) + end +end + +function _constructarrowtypespart_impl(name::String, v::AbstractVector, logical; context) + # VersionNumber's legacy reflected representation contains variable-length + # tuples, so Arrow 3.0 writes String extension storage instead. ArrowTypes' + # one-argument JuliaType registration and the VersionNumber(String) + # constructor read it back. Arrow 2.x readers cannot: they index the + # column with the Struct layout ArrowTypes declares for VersionNumber. + if logical === VersionNumber + mapped = + _eltypeof(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 <: _eltypeof(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 + _checkstoragemissing(logical, declared) + if logical isa Union && declared === logical + # A Julia Union writes one Arrow child per member, except members that + # all write one field (several string types, two structs stored as + # Int32 without a label): they write one column of it. + members = Base.uniontypes(logical) + merged = if all(isconcretetype, members) + try + fields = + AC.Field[first(_constructpart(name, M[]; context)) for M in members] + if all( + f -> _conformancemismatch(fields[1], f; exact=true) === nothing, + fields, + ) + f = fields[1] + f = AC.Field( + f.name, + f.type; + nullable=Missing <: _eltypeof(v), + metadata=_fieldmetadata(f), + children=f.children, + ) + _constructpart(f, v; context) + end + catch err + # A member whose empty column has no type to infer, or rows the + # one field's builder does not take (plain structs), write as + # a union. + err isa ArgumentError || rethrow() + nothing + end + end + merged === nothing || return merged + f, d = _arrowtypesunioncolumn(name, v; 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) + f, d = _arrowtypesstoragecolumn(name, mapped, context) + # 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 <: _eltypeof(v)) ? false : f.nullable + return _arrowtypeslogicalfield(context, f, logical; nullable=nullable), d +end + +# Missing marks a null slot; as a value of a storage Union it could not be told +# from one. +function _checkstoragemissing(logical, declared) + declared isa Union && + Missing <: declared && + throw( + ArgumentError( + "ArrowTypes.ArrowType($logical) is $declared, with Missing as a Union " * + "member, which Arrow.jl does not write; declare the storage without " * + "Missing and write the column with missing values", + ), + ) + return nothing +end + +"A column of lowered storage values, built by the layout its storage type's kind names." +function _arrowtypesstoragecolumn(name::String, v::AbstractVector, context) + _eltypeof(v) === Missing && return _constructnativepart(name, v; context) + S = Base.nonmissingtype(_eltypeof(v)) + S === Union{} && return _constructnativepart(name, v; context) + S isa Union && return _uninferred(_arrowtypesunioncolumn, name, v; context) + _arrowtypesnativetype(S) && return _constructnativepart(name, v; context) + kind = S <: NamedTuple ? ArrowTypes.StructKind() : _writerkind!(context, S) + if kind isa ArrowTypes.NullKind + nulls = S === Nothing ? fill(missing, length(v)) : v + return _constructnativepart(name, nulls; context) + elseif kind isa ArrowTypes.FixedSizeListKind + N = ArrowTypes.getsize(kind) + ArrowTypes.gettype(kind) === UInt8 && return _fixedbinarypart(name, v, N) + return _uninferred(_fixedlistpart, name, v, context, N) + elseif kind isa ArrowTypes.MapKind + S <: AbstractDict || throw( + ArgumentError( + "ArrowTypes MapKind column $name must lower to AbstractDict rows", + ), + ) + return _uninferred(_mappart, name, v, context) + elseif kind isa ArrowTypes.ListKind + S <: Union{AbstractString,AbstractVector{UInt8}} && + return _constructnativepart(name, v; context) + S <: AbstractVector || throw( + ArgumentError( + "column $(_errname(context, name)) holds $S values, which are not " * + "vectors; an Arrow list holds a vector per row (vec(x) flattens an array)", + ), + ) + return _uninferred(_listpart, name, v, context) + elseif kind isa ArrowTypes.StructKind + _arrowtypesstorageisspecified(context, S) || throw( + ArgumentError( + "ArrowTypes StructKind column $name lowered to an underspecified type " * + "$S; give every struct field a writable storage type", + ), + ) + return _uninferred(_structpart, name, v, 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 _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; context) + variants = _checkedwritervariants("Arrow Union column $name", _eltypeof(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 $(_eltypeof(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) + # An abstract string member writes its values as Strings. + values = + variants[i] === AbstractString ? _missings_to(String, childvalues[i]) : + childvalues[i] + field, data = _constructpart(string(i - 1), values; 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 + +""" +Field `k` of each row of a retained struct column: a typed read's row is a +`NamedTuple`, a materialized one an ordered `Pair` vector; a missing row reads +as missing. +""" +struct _RowFields{V<:AbstractVector} <: AbstractVector{Any} + rows::V + k::Int +end +Base.size(c::_RowFields) = (length(c.rows),) +function Base.getindex(c::_RowFields, i::Int) + x = c.rows[firstindex(c.rows) + i - 1] + x === missing && return missing + # Inferring a NamedTuple field's exact type makes Julia 1.10 with code + # coverage fail ("fatal error in type inference (type bound)"). + x isa NamedTuple && return _uninferred(getfield, x, c.k) + return last(x[firstindex(x) + c.k - 1]) +end + +mutable struct _WriterContext + arrowtypes::_ArrowTypesContext + needs::IdDict{Type,Bool} + kinds::IdDict{Type,Any} + extensions::IdDict{Type,Any} + schema::Vector{Type} + values::Base.IdSet{Any} + # True while the values being rebuilt are already in their storage domain + # (see `_constructlowered`); false for facade-materialized public rows. + lowered::Bool + column::String + # Fresh strings and bytes write as utf8/binary views instead of offsets. + stringview::Bool + # Fresh strings, bytes, and lists declare 64-bit offsets. + largelists::Bool + # Where a writer's rebuilt buffers come from (`_BufferPool`), or nothing. + pool::Any +end + +function _WriterContext( + column::AbstractString=""; + stringview::Bool=false, + largelists::Bool=false, + pool=nothing, +) + arrowtypes = _ArrowTypesContext(; warn=false) + return _WriterContext( + arrowtypes, + IdDict{Type,Bool}(), + IdDict{Type,Any}(), + IdDict{Type,Any}(), + Type[], + Base.IdSet{Any}(), + false, + String(column), + stringview, + largelists, + pool, + ) +end + +const _MAX_WRITER_SCHEMA_DEPTH = 64 + +function _withwriterschema(f, context::_WriterContext, T::Type) + Base.@nospecialize T + if T in context.schema + path = + _uninferred(join, (string(S) for S in (context.schema..., T)), " -> ")::String + 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 + +function _writerneedstype!(context::_WriterContext, T::Type) + Base.@nospecialize T + haskey(context.needs, T) && return context.needs[T] + return context.needs[T] = _arrowtypesneedstype( + T, + nested -> _writerstoragetype!(context, nested), + nested -> _writerextension!(context, nested) !== nothing, + 0, + ) +end + +function _writerkind!(context::_WriterContext, T::Type) + Base.@nospecialize T + haskey(context.kinds, T) && return context.kinds[T] + return context.kinds[T] = ArrowTypes.ArrowKind(T) +end + +function _writerextension!(context::_WriterContext, T::Type) + Base.@nospecialize T + haskey(context.extensions, T) && return context.extensions[T] + return context.extensions[T] = _writerextensionlabel(T) +end + +function _writerextensionlabel(T::Type) + Base.@nospecialize T + # A Union has the label its members share, if they all have one: a label + # that only matches the Union as a whole (tuple types' JuliaLang.Tuple) + # belongs to none of the children it writes as. + U = Base.nonmissingtype(T) + if U isa Union + labels = Any[_writerextensionlabel(M) for M in Base.uniontypes(U)] + return allequal(labels) ? first(labels) : nothing + end + name = _facadearrowname(T) + name === nothing || return Pair{String,String}[ + _EXTENSION_NAME_KEY => name, + _EXTENSION_METADATA_KEY => "", + ] + ArrowTypes.hasarrowname(T) || return nothing + # Asserted: the label methods are inferred for any `T`, and a `String` of an + # unknown value would convert through methods other packages add. + return Pair{String,String}[ + _EXTENSION_NAME_KEY => String(ArrowTypes.arrowname(T))::String, + _EXTENSION_METADATA_KEY => String(ArrowTypes.arrowmetadata(T))::String, + ] +end + +function _checkabstractlabel(f::AC.Field, target, T::Type, context::_WriterContext) + isconcretetype(target) && return nothing + label = _writerextension!(context, T) + label === nothing && return nothing + retained = _arrowtypesextension(f) + (last(label[1]), last(label[2])) == retained || throw( + ArgumentError( + "registered writer subtype $T declares extension $(repr(label)), which " * + "cannot replace abstract retained target $target with extension " * + "$(repr(retained))", + ), + ) + return nothing +end + +""" +Narrow one unresolved abstract declaration from whole-column runtime evidence. +Returns nothing when the declaration cannot be narrowed from runtime evidence. +""" +function _narrowabstractcolumn(name::String, v::AbstractVector, context::_WriterContext) + logical = Base.nonmissingtype(_eltypeof(v)) + logical in (Union{}, Any) && return nothing + (isconcretetype(logical) || logical isa Union) && return nothing + _writerstoragetype!(context, logical) === logical || return nothing + _writerextension!(context, logical) === nothing || return nothing + # Abstract row types of a plain nested layout (NTuples of nullable values, + # `AbstractVector{Int}`) write as that layout without runtime narrowing. + _nestedkind(context, logical) === nothing || return nothing + + observed = Type[] + seen = Base.IdSet{Type}() + nullable = Missing <: _eltypeof(v) + for value in v + value === missing && continue + value isa logical || throw( + ArgumentError( + "column $name contains $(typeof(value)), which is outside its " * + "declared abstract element type $logical", + ), + ) + _pushobserved!(observed, seen, typeof(value), name) + end + isempty(observed) && return nothing + needsruntime = any( + T -> _writerneedstype!(context, T) || _writerextension!(context, T) !== nothing, + observed, + ) + if !needsruntime && _singletypefamily(logical) + length(observed) == 1 || return nothing + storagetype = nullable ? Union{Missing,only(observed)} : only(observed) + return _collectas(storagetype, v) + end + if !needsruntime && _abstractnonnative(logical) + # `Real[1, 2.5]`: one runtime type writes as that type, several as a + # dense Union of them. + storagetype = _writeruniontype(observed) + nullable && (storagetype = Union{Missing,storagetype}) + return _collectas(storagetype, v) + end + if !needsruntime + # Rows of an abstract NamedTuple column (some fields missing in some + # rows, say) share one concrete row type once widened field by field. + logical <: NamedTuple || return nothing + rowtype = _widenednamedtuple(observed) + rowtype === nothing && return nothing + nullable && (rowtype = Union{Missing,rowtype}) + return _collectas(rowtype, v) + end + storagetype = _writeruniontype(observed) + nullable && (storagetype = Union{Missing,storagetype}) + return _collectas(storagetype, v) +end + +"The field-wise Union of same-named NamedTuple row types, or nothing." +function _widenednamedtuple(observed) + names = fieldnames(first(observed)) + all(T -> T <: NamedTuple && fieldnames(T) === names, observed) || return nothing + types = _uninferred( + ntuple, + i -> Union{(fieldtype(T, i) for T in observed)...}, + length(names), + )::Tuple + return NamedTuple{names,Tuple{types...}} +end + +# ---- registered rewrites: lower a labelled column once ---------------------- + +"Append runtime type `T` to `types` once; more than 8 is an error naming the column." +function _pushobserved!(types::Vector{Type}, seen::Base.IdSet{Type}, T::Type, name::String) + T in seen && return nothing + length(types) < _MAX_INFERRED_WRITER_TYPES || throw( + ArgumentError( + "column $name holds values of more than $_MAX_INFERRED_WRITER_TYPES types " * + "($(_uninferred(join, types, ", ")::String), $T, ...); give it a concrete " * + "element type", + ), + ) + push!(seen, T) + push!(types, T) + return nothing +end + +""" +Whether a column's values lower to storage before `f`'s builder runs: a +labelled field's do, and so do values with mappings of their own, except +under a union or an encoding, whose children lower their own values. +""" +_lowersvalues(f::AC.Field, target, v, context) = + target !== nothing || + !(f.type isa Union{AC.UnionType,AC.RunEndEncodedType,AC.DictionaryType}) && + _mappedvalues(v, context) + +""" +Whether `v`'s values lower through an `ArrowTypes.toarrow` mapping of their +own: the element type's, or a runtime type's in an abstract column. +""" +function _mappedvalues(v::AbstractVector, context::_WriterContext) + T = Base.nonmissingtype(_eltypeof(v)) + T === Union{} && return false + # A type stored as a null (`Nothing`, whose storage is `Missing`) is not + # lowered: only its label tells its values from missing ones. + mapped(S) = + S !== VersionNumber && + !(S <: Union{Int128,UInt128}) && + !(_writerstoragetype!(context, S) in (S, Missing, Nothing)) + isconcretetype(T) && return mapped(T) + T isa Union && + all(isconcretetype, Base.uniontypes(T)) && + return any(mapped, Base.uniontypes(T)) + return any(x -> x !== missing && mapped(typeof(x)), v) +end + +""" +The values of a column under the retained field `f` lowered once to their +storage: `ArrowTypes.toarrow` on each value, converted to the storage type +`f`'s target declares (`Missing` kept), or, when that storage is abstract or +`f` has no target, to the promoted type of the values' lowered types. A +labelled child inside a lowered value keeps its logical value, which lowers +when that child builds. +""" +function _lowercolumn(f::AC.Field, target, v::AbstractVector, context::_WriterContext) + name = f.name + target === nothing && return _inferlowered(f, nothing, v, context) + if target === VersionNumber + S = any(ismissing, v) ? Union{Missing,String} : String + return S[x === missing ? missing : string(x) for x in v] + end + declared = _writerstoragetype!(context, target) + _checkstoragemissing(target, declared) + _arrowtypesconcreteorunion(declared) || + _arrowtypesstorageisspecified(context, declared) || + return _inferlowered(f, target, v, context) + if !isconcretetype(target) + seen = Base.IdSet{Type}() + for x in v + x === missing && continue + T = typeof(x) + T in seen && continue + push!(seen, T) + _checkabstractlabel(f, target, T, context) + end + end + storage = any(x -> x === missing, v) ? Union{Missing,declared} : declared + return _writerloweredcolumn(storage, v, context, target, name, f) +end + +"`values` lowered by `lower` one at a time as they are read; missing stays missing." +struct _LoweredValues{T,V<:AbstractVector,F} <: AbstractVector{T} + lower::F + values::V +end +_LoweredValues{T}(lower::F, values::V) where {T,V,F} = _LoweredValues{T,V,F}(lower, values) +Base.size(l::_LoweredValues) = (length(l.values),) +Base.@propagate_inbounds function Base.getindex(l::_LoweredValues{T}, i::Int) where {T} + x = l.values[firstindex(l.values) + i - 1] + return (x === missing ? missing : l.lower(x))::T +end + +# Every value of a column under a labelled field is its target's; a concrete +# element type that is not the target says the column was replaced. +function _checktarget(f::AC.Field, target, v::AbstractVector) + T = Base.nonmissingtype(_eltypeof(v)) + (T === Union{} || T <: target) && return nothing + i = + isconcretetype(T) ? firstindex(v) : + findfirst(x -> x !== missing && !(x isa target), v) + i === nothing && return nothing + S = isconcretetype(T) ? T : typeof(v[i]) + throw( + ArgumentError( + "column $(f.name) holds $S values, but its retained ArrowTypes target is " * + "$target", + ), + ) +end + +""" +`v` lowered value by value through each value's own mapping, as a column of +the promoted type of the lowered types (a Union of them when they do not +promote to one concrete type). Under an abstract target, each runtime type +must carry the target's label. +""" +function _inferlowered(f::AC.Field, target, v::AbstractVector, context::_WriterContext) + name = f.name + values = Vector{Any}(undef, length(v)) + observed = Type[] + seen = Base.IdSet{Type}() + runtimes = Base.IdSet{Type}() + hasmissing = false + for (i, x) in enumerate(v) + if x === missing + values[i] = missing + hasmissing = true + continue + end + T = typeof(x) + if !(T in runtimes) + push!(runtimes, T) + target === nothing || _checkabstractlabel(f, target, T, context) + end + mapped = T === VersionNumber ? string(x) : _writertoarrow(context, x, T, name, f) + # Without a label, a type stored as a null (`Nothing`) would read back as + # missing values: its values stay themselves, for the field's builder to + # refuse. + mapped isa Union{Missing,Nothing} && + target === nothing && + _writerstoragetype!(context, T) in (Missing, Nothing) && + (mapped = x) + values[i] = mapped + _pushobserved!(observed, seen, typeof(mapped), name) + end + isempty(observed) && return fill(missing, length(v)) + storagetype = reduce(_writerpromoteunion, observed) + hasmissing && (storagetype = Union{Missing,storagetype}) + return _writerconvertedcolumn( + storagetype, + values, + context, + something(target, Any), + name, + ) +end + +""" +Build `f` without its label over `storage`, values in its storage domain +(`context.lowered`), and give the result `f`'s label and nullability back. +""" +function _constructlowered(f::AC.Field, storage::AbstractVector, context::_WriterContext) + lowered = context.lowered + context.lowered = true + built, data = try + _constructpart(_withoutownextension(f), storage; context) + finally + context.lowered = lowered + end + return _retainedfield(f; children=collect(AC.Field, built.children)), data +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 + +# Zeroed buffers for `n` slots of `f`'s layout, all null when `f` is nullable. +function _writerhiddenbuffers(f::AC.Field, n::Int) + spec = AC.layoutspec_of(f.type) + buffers = AC.BufferSlice[] + for role in spec.buffers + nbytes = + role == AC.VALIDITY ? (f.nullable ? cld(n, 8) : 0) : + role == AC.DATA ? + ( + spec.fixedwidth == -1 ? cld(n, 8) : + _writerhiddenmul(f, n, spec.fixedwidth, "data") + ) : + role == AC.OFFSETS ? _writerhiddenmul(f, n + 1, spec.offsetwidth, "offset") : + role == AC.ELEMENT_OFFSETS || role == AC.SIZES ? + _writerhiddenmul(f, n, spec.offsetwidth, "element-offset") : + role == AC.VIEWS ? _writerhiddenmul(f, n, 16, "view") : + throw( + ArgumentError( + "hidden construction for retained field $(f.name) needs a " * + "layout-specific $(role) buffer", + ), + ) + push!(buffers, _writerhiddenbytes(f, nbytes, "$role")) + end + return buffers, f.nullable ? n : 0 +end + +""" +`n` slots of retained field `f` that nothing reads, from zeroed buffers and no +values: nulls where `f` is nullable, zero placeholders where it is not (an +all-missing column of a later table, a child of null rows only). +""" +function _constructhiddenpart(f::AC.Field, n::Int, context::_WriterContext) + n >= 0 || throw(ArgumentError("negative hidden length for retained field $(f.name)")) + t = f.type + t isa AC.NullType && + return _retainedfield(f), AC.ArrayData(t, n, AC.BufferSlice[]; nullcount=n) + fields = AC.Field[] + datas = AC.ArrayData[] + function child!(c, m) + cf, cd = _constructhiddenpart(c, m, context) + push!(fields, cf) + push!(datas, cd) + return nothing + end + dictionary = nothing + if t isa AC.DictionaryType + # A non-nullable dictionary's valid indices point at one pool entry. + poollength = n == 0 || f.nullable ? 0 : 1 + _, dictionary = _constructhiddenpart(AC.dictvaluefield(f, t), poollength, context) + append!(fields, f.children) + elseif t isa Union{AC.ListType,AC.ListViewType,AC.MapType} + child!(only(f.children), 0) + elseif t isa AC.FixedSizeListType + child!(only(f.children), Int(_writerhiddenmul(f, n, t.listsize, "child"))) + elseif t isa AC.StructType + foreach(c -> child!(c, n), f.children) + elseif t isa AC.RunEndEncodedType + # One run of one value. + RT = AC.juliatype(f.children[1].type) + n <= typemax(RT) || throw( + ArgumentError("column $(f.name) length exceeds its retained run-end type"), + ) + rf, rd = _constructpart(f.children[1], n == 0 ? RT[] : RT[RT(n)]; context) + push!(fields, rf) + push!(datas, rd) + child!(f.children[2], min(n, 1)) + data = AC.ArrayData(t, n, AC.BufferSlice[]; children=datas, nullcount=0) + return _retainedfield(f; children=fields), data + elseif t isa AC.UnionType + # Every slot takes slot 0 of the child that holds nulls, else (slots + # nothing reads) of the first child. + k = something(_unionnullchild(f), 1) + sparse = t.mode == AC.SparseMode + for (j, c) in enumerate(f.children) + child!(c, sparse ? n : j == k ? min(n, 1) : 0) + end + typeids = n == 0 ? AC.BufferSlice() : AC._databuffer(fill(t.typeids[k], n)) + offsets = n == 0 ? AC.BufferSlice() : AC._databuffer(zeros(Int32, n)) + buffers = sparse ? AC.BufferSlice[typeids] : AC.BufferSlice[typeids, offsets] + data = AC.ArrayData(t, n, buffers; children=datas, nullcount=0) + return _retainedfield(f; children=fields), data + end + buffers, nullcount = _writerhiddenbuffers(f, n) + data = AC.ArrayData(t, n, buffers; children=datas, dictionary, nullcount) + return _retainedfield(f; children=fields), data +end + +# Registered mappings exchange ArrowTypes-domain temporals (`_arrowtypesscalar`): +# micro/nanosecond timestamps as raw Int64 and second/millisecond timestamps as +# zone-free DateTime. Unmarked fresh composites bring public values to the same +# seam, so both domains lower here; anything else has no exact storage. +function _arrowtypestemporalstorage(t::AC.ArrowType, value) + ok, storage = _facadetostorage(t, value) + ok && return storage + t isa AC.TimestampType || return nothing + t.unit in (AC.MICROSECOND, AC.NANOSECOND) && value isa Int64 && return value + return value isa Dates.DateTime ? _timestampstorage(t.unit, value) : nothing +end + +"A fixed_size_binary(N) column from `NTuple{N,UInt8}` rows." +function _fixedbinarypart(name::String, v::AbstractVector, N::Int) + T = NTuple{N,UInt8} + t = AC.FixedSizeBinaryType(N) + field = AC.Field(name, t; nullable=Missing <: _eltypeof(v)) + if _eltypeof(v) === T + data = something(_fixedbuffer(v), AC._databuffer(collect(T, v))) + return field, AC.ArrayData(t, length(v), [AC.BufferSlice(), data]; nullcount=0) + end + rows, validity, nc = _nullablevalues(T, v, ntuple(_ -> 0x00, N)) + return field, AC.ArrayData(t, length(v), [validity, AC._databuffer(rows)]; nullcount=nc) +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." +_fieldcontractequal(a::AC.Field, b::AC.Field) = + _conformancemismatch(a, b; exact=true) === nothing + +""" +Assemble a fresh Int32-indexed dictionary Field and its first batch from +already-built pool data. +""" +# The narrowest signed index type for `n` categories (indices 0 to n - 1). +_dictindextype(n::Integer) = + n <= 128 ? AC.IntType(8, true) : + n <= 32_768 ? AC.IntType(16, true) : + n - 1 <= typemax(Int32) ? AC.IntType(32, true) : AC.IntType(64, true) + +# Whether a column's categories are ordered: a read dictionary column's flag, +# a CategoricalArray's (the CategoricalArrays extension), else false. +_isordered(::AbstractVector) = false +_isordered(x::SubArray) = _isordered(parent(x)) +_isordered(x::DictEncode) = _isordered(x.data) +_isordered(x::DictEncoded) = x.ordered +_isordered(x::ChainedVector{T,<:DictEncoded}) where {T} = + !isempty(x.arrays) && all(_isordered, x.arrays) + +function _newdictfromdata( + name::String, + valuefield::AC.Field, + valuedata::AC.ArrayData, + indices::AbstractVector; + nullable::Bool, + indextype::AC.IntType=AC.IntType(32, true), + ordered::Bool=false, +) + t = AC.DictionaryType(indextype, valuefield.type, ordered) + field = AC.Field( + name, + t; + nullable, + metadata=_fieldmetadata(valuefield), + children=collect(AC.Field, valuefield.children), + ) + return field, _dictbatch(field, indices, valuedata) +end + +# --- retained-schema rewrite (facade Table/Stream round-trips) -------------- + +"Storage integers for a column of public values 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 + +# The same for lowered values, in the ArrowTypes domain too: raw integers for +# microsecond and nanosecond timestamps, `DateTime`s for the coarser ones. +function _loweredstorage(t::AC.ArrowType, v::AbstractVector, name::String) + out = Union{Missing,Int64}[] + sizehint!(out, length(v)) + for x in v + if x === missing + push!(out, missing) + else + sv = _arrowtypestemporalstorage(t, x) + sv === nothing && throw( + ArgumentError( + "column $name holds the value $(repr(x)), which cannot be " * + "represented exactly by its retained Arrow type $(repr(t))", + ), + ) + push!(out, Int64(sv)) + end + end + return out +end + +# The reader takes arrow.bool8 as Bool only on int8 storage; on other +# storage the label is kept and the values are written as they read. +_isbool8(f::AC.Field) = + f.type == AC.IntType(8, true) && + _fieldmetavalue(f, _EXTENSION_NAME_KEY) == "arrow.bool8" +_bool8storage(v::AbstractVector) = + Missing <: _eltypeof(v) ? Union{Missing,Int8}[x isa Bool ? Int8(x) : x for x in v] : + Int8[x isa Bool ? Int8(x) : x for x in v] + +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] + f.nullable || all(present) || throw(_nonnullableerror(f.name)) + return present +end + +# The writer's utf8/binary builders, after checking the values against the +# retained field: strings for utf8, byte vectors for binary. +function _retainedvarbytes(f::AC.Field, v::AbstractVector, pool=nothing) + t = f.type::Union{AC.Utf8Type,AC.BinaryType} + utf8 = t isa AC.Utf8Type + Offset = t.large ? Int64 : Int32 + _, d = + if v isa Union{_PayloadColumn,_PayloadRows} && + _payloadparent(v) isa (utf8 ? DataStrings.StringVector : DataStrings.BytesVector) + _packedpayloadpart( + f.name, + _payloadentries(v), + _payloadparent(v).buffers, + t, + Offset, + Missing <: _eltypeof(v), + pool, + ) + else + V = utf8 ? AbstractString : AbstractVector{UInt8} + Base.nonmissingtype(_eltypeof(v)) <: V || + all(x -> x === missing || x isa V, v) || + throw( + ArgumentError( + "column $(f.name) must contain " * + (utf8 ? "string" : "byte-vector") * + " values", + ), + ) + _varbytespart(f.name, v, t, Offset) + end + AC.nullcount(d) == 0 || + f.nullable || + throw( + ArgumentError( + "column $(f.name) holds missing values but its retained field is non-nullable", + ), + ) + return _retainedfield(f), d +end + +function _retainedview(f::AC.Field, v::AbstractVector) + t = f.type::AC.ViewType + present = _retainedvalidity(f, v) + payloads = Vector{DataStrings.StringPayload}(undef, length(v)) + data = UInt8[] + for (i, x) in enumerate(v) + if x === missing + # An empty view: the validity bit marks the null, and the slot + # stays valid when a parent hides this child's validity. + payloads[i] = DataStrings.StringPayload(UInt64(0), UInt64(0)) + continue + end + bytes = if t.utf8 + x isa AbstractString || + throw(ArgumentError("column $(f.name) must contain string values")) + _utf8units(x) + else + x isa AbstractVector{UInt8} || + throw(ArgumentError("column $(f.name) must contain byte-vector values")) + x + end + n = length(bytes) + if n <= DataStrings.INLINE_MAX + payloads[i] = DataStrings.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] = DataStrings.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 + if Base.nonmissingtype(_eltypeof(v)) === NTuple{t.nbytes,UInt8} + _, d = _fixedbinarypart(f.name, v, t.nbytes) + AC.nullcount(d) > 0 && + !f.nullable && + throw( + ArgumentError( + "column $(f.name) holds missing values but its retained field is non-nullable", + ), + ) + return _retainedfield(f), d + end + 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 Union{AbstractVector{UInt8},NTuple{<:Any,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)", + ), + ) + for k = 1:(t.nbytes) + data[(i - 1) * t.nbytes + k] = x[k] + end + end + d = AC.ArrayData( + t, + length(v), + [AC._bitmapbuffer(present), AC._databuffer(data)]; + nullcount=count(!, present), + ) + return _retainedfield(f), d +end + +"Values (`fill` at missing slots), validity buffer, and null count of a column." +function _nullablevalues(::Type{S}, v::AbstractVector, fill::S, f::F=identity) where {S,F} + n = length(v) + i0 = firstindex(v) - 1 + vals = Vector{S}(undef, n) + bits = Vector{UInt8}(undef, cld(n, 8)) + valid = 0 + # Each validity byte is built in a register from 8 values, branch-free; + # whole bytes take a fixed 8-step loop the compiler unrolls. + @inbounds for b in eachindex(bits) + base = (b - 1) << 3 + byte = 0x00 + if base + 8 <= n + for j = 0:7 + x = v[i0 + base + j + 1] + present = x !== missing + vals[base + j + 1] = present ? f(x) : fill + byte |= UInt8(present) << (j % UInt) + end + else + for j = 0:(n - base - 1) + x = v[i0 + base + j + 1] + present = x !== missing + vals[base + j + 1] = present ? f(x) : fill + byte |= UInt8(present) << (j % UInt) + end + end + bits[b] = byte + valid += count_ones(byte) + end + nc = n - valid + validity = + nc == 0 ? AC.BufferSlice() : AC.BufferSlice(AC.heapregion(bits), 0, length(bits)) + return vals, validity, nc +end + +# A `Decimal{P,S,T}` is its storage integer, so a column of them already is +# the Arrow coefficient buffer: a `Vector` is wrapped without copying and any +# other vector is collected once. Values write as stored, as Arrow C++ does; +# the precision check is `validate_full`'s. +function _decimaldata(f::AC.Field, v::AbstractVector, ::Type{D}) where {D} + t = f.type + if _eltypeof(v) === D + data = something(_fixedbuffer(v), AC._databuffer(collect(D, v))) + return AC.ArrayData(t, length(v), [AC.BufferSlice(), data]; nullcount=0) + end + vals, validity, nc = _nullablevalues(D, v, zero(D)) + nc > 0 && + !f.nullable && + throw( + ArgumentError( + "column $(f.name) holds missing values but its retained field is non-nullable", + ), + ) + return AC.ArrayData(t, length(v), [validity, AC._databuffer(vals)]; nullcount=nc) +end + +# `v`'s decimals as `H`s, whose storage is `T`. +_widendecimals(::Type{H}, ::Type{T}, v) where {H,T} = + map(x -> x === missing ? missing : reinterpret(H, T(DataDecimals.unscaled(x))), v) + +function _retaineddecimal(f::AC.Field, v::AbstractVector) + t = f.type::AC.DecimalType + D = Base.nonmissingtype(_eltypeof(v)) + if _shareddecimal(t) + H = _decimalhost(t) + D === H && return _retainedfield(f), _decimaldata(f, v, D) + # Another storage width at the field's scale (`dec"1.25"` stores an + # Int32): widen every coefficient in one pass. + if D <: DataDecimals.Decimal && isconcretetype(D) && D.parameters[2] == t.scale + w = _widendecimals(H, _decimalstorage(t), v) + return _retainedfield(f), _decimaldata(f, w, H) + end + end + v = [x isa DataDecimals.AbstractDecimal ? _shareddecimalraw(t, x) : x for x in v] + 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 + v = [x isa _IntervalValue ? _sharedintervalraw(t, x) : x for x in v] + 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 _validateretainedstructrows(f::AC.Field, rows::AbstractVector) + nchildren = length(f.children) + for row in rows + row === missing && continue + if row isa NamedTuple + # A typed read's row: the field names must be the children's. + ( + length(row) == nchildren && + all(j -> String(keys(row)[j]) == f.children[j].name, 1:nchildren) + ) || throw( + ArgumentError( + "column $(f.name) retained struct row has fields $(keys(row)); " * + "expected $(Tuple(c.name for c in f.children))", + ), + ) + continue + end + 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 _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, +) + for row = 1:(length(offsets) - 1) + seen = Set{Any}() + for index = (Int(offsets[row]) + 1):Int(offsets[row + 1]) + key = _storedkey(keyfield, keydata, Int64(index)) + key in seen && throw( + ArgumentError( + "column $name row $row has duplicate physical Map key storage", + ), + ) + push!(seen, key) + end + end + return nothing +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")) + if v isa RunEndEncoded + # A run-end column read from Arrow keeps its runs: those over its rows, + # each cut to them. + lo = searchsortedfirst(v.runends, v.offset + 1) + hi = searchsortedfirst(v.runends, v.offset + v.len) + runends = RT[RT(min(v.runends[r], v.offset + v.len) - v.offset) for r = lo:hi] + runvalues = v.len == 0 ? v.values[1:0] : v.values[lo:hi] + else + 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 run coalescing. + if isempty(runvalues) || x !== last(runvalues) + push!(runvalues, x) + push!(runends, RT(i)) + else + runends[end] = RT(i) + end + end + end + rebuiltvaluefield, valuedata = _constructpart(valuefield, runvalues; context) + # Neighbouring runs whose values are stored alike are one run. + selected = Int[] + ends = RT[] + previous = nothing + for i = 1:(valuedata.len) + key = _storedkey(rebuiltvaluefield, valuedata, Int64(i)) + if i > 1 && isequal(key, previous) + ends[end] = runends[i] + else + push!(selected, i) + push!(ends, runends[i]) + end + previous = key + end + if length(selected) < valuedata.len + valuedata = _takerows(rebuiltvaluefield, valuedata, selected) + runends = ends + 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 one part under a retained Field: descriptor, nullability, metadata." +Base.@nospecializeinfer _constructpart( + f::AC.Field, + @nospecialize(v::AbstractVector); + context::_WriterContext=_WriterContext(), +) = _uninferred(_retainedpartof, f, v, context)::Tuple{AC.Field,AC.ArrayData} + +Base.@nospecializeinfer function _retainedpartof( + f::AC.Field, + @nospecialize(v::AbstractVector), + context::_WriterContext, +) + # No value to build: an empty column is its field's empty data. + isempty(v) && return _constructhiddenpart(f, 0, context) + # Values of an abstract element type (a flattened list's, a struct field's) + # narrow to the types they hold. + _arrowtypesconcreteorunion(_eltypeof(v)) || (v = map(identity, v)) + # Values with no type of their own (an all-missing child of `Any` values) + # are the field's nulls, whatever its layout. + _holdsnull(f) && + _eltypeof(v) === Missing && + return _constructhiddenpart(f, length(v), context) + _, target, _ = _arrowtypestarget(context.arrowtypes, f) + if _lowersvalues(f, target, v, context) + target === nothing || _checktarget(f, target, v) + # Lowered and built without its label, a labelled union or Null field + # would take missing values that only an unlabelled one keeps (and + # `nothing` lowers to `missing`). + target !== nothing && + f.type isa Union{AC.UnionType,AC.NullType} && + !_holdsnull(f) && + any(ismissing, v) && + throw(_nonullerror(f)) + # No value to lower: an all-missing column is its field's nulls. + _holdsnull(f) && + all(ismissing, v) && + return _constructhiddenpart(f, length(v), context) + return _constructlowered(f, _lowercolumn(f, target, v, context), context) + end + return _constructretainedpart(f, v, context) +end + +# Writer-side element types a retained descriptor's builder converts exactly, +# beyond its public type: a written `DateTime` (or another +# date-time) column under a timestamp, any `AbstractString` under a string. +_retainedwritertype(::AC.TimestampType, ::Type{T}) where {T} = T <: Dates.AbstractDateTime +_retainedwritertype(::AC.Utf8Type, ::Type{T}) where {T} = T <: AbstractString +_retainedwritertype(t::AC.ViewType, ::Type{T}) where {T} = + t.utf8 ? T <: AbstractString : T <: AbstractVector{UInt8} +_retainedwritertype(::AC.BinaryType, ::Type{T}) where {T} = T <: AbstractVector{UInt8} +_retainedwritertype(t::AC.FixedSizeBinaryType, ::Type{T}) where {T} = + T <: NTuple{t.nbytes,UInt8} || T <: AbstractVector{UInt8} +_retainedwritertype(::AC.ArrowType, ::Type) = false + +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(_eltypeof(v)) + NT <: Fp || + NT === Union{} || + _retainedwritertype(t, NT) || + 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", + ), + ) + _eltypeof(v) >: Missing && + !f.nullable && + throw( + ArgumentError( + "column $(f.name) may hold missing values but its retained " * + "field is non-nullable", + ), + ) + end + return Fp +end + +# The builder of a retained layout other than Null, Union, and the primitive +# types, or nothing. +_retainedbuilder(t::AC.ArrowType) = + t isa Union{AC.DateType,AC.TimestampType,AC.TimeType,AC.DurationType} ? + _retainedtemporal : + t isa Union{AC.Utf8Type,AC.BinaryType} ? _retainedvarbytes : + t isa AC.ViewType ? _retainedview : + t isa AC.FixedSizeBinaryType ? _retainedfixedbytes : + t isa AC.DecimalType ? _retaineddecimal : + t isa AC.IntervalType ? _retainedinterval : + t isa Union{AC.ListType,AC.ListViewType} ? _listpart : + t isa AC.FixedSizeListType ? _fixedlistpart : + t isa AC.StructType ? _structpart : + t isa AC.MapType ? _mappart : t isa AC.RunEndEncodedType ? _retainedree : nothing + +function _retainedtemporal(f::AC.Field, v::AbstractVector, context::_WriterContext) + t = f.type + context.lowered && return _rebuildtemporal(f, _loweredstorage(t, v, f.name)) + _checkretainedidentity(f, t, v) + return _rebuildtemporal(f, _retainedstorage(t, v, f.name)) +end +_retainedvarbytes(f::AC.Field, v::AbstractVector, context::_WriterContext) = + _retainedvarbytes(f, v, context.pool) +_retainedview(f::AC.Field, v::AbstractVector, ::_WriterContext) = _retainedview(f, v) +_retainedfixedbytes(f::AC.Field, v::AbstractVector, ::_WriterContext) = + _retainedfixedbytes(f, v) +_retaineddecimal(f::AC.Field, v::AbstractVector, ::_WriterContext) = _retaineddecimal(f, v) +_retainedinterval(f::AC.Field, v::AbstractVector, ::_WriterContext) = + _retainedinterval(f, v) + +"One part under a retained field: the identity check, then the type's builder." +function _constructretainedpart(f::AC.Field, v::AbstractVector, context::_WriterContext) + t = f.type + t isa AC.DictionaryType && throw( + ArgumentError( + "nested retained dictionary field $(f.name) cannot be reconstructed " * + "after facade materialization because its pool is unavailable", + ), + ) + if t isa AC.NullType + # `nothing` is Null storage only where a label lowered a value to it. + all(x -> x === missing || (x === nothing && context.lowered), v) || 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 + t isa Union{AC.DateType,AC.TimestampType,AC.TimeType,AC.DurationType} || + _checkretainedidentity(f, t, v) + # Each layout's builder compiles for the columns that reach it, not for + # every column type this method sees (see `_uninferred`). + builder = _retainedbuilder(t) + builder === nothing || + return _uninferred(builder, f, v, context)::Tuple{AC.Field,AC.ArrayData} + if t isa AC.UnionType + built = _uninferred(_constructfacadeunion, f, v, context) + built === nothing || return built + throw( + ArgumentError( + "column $(f.name) is an Arrow union, and some of its values fit none or " * + "several of the union's children (two children of one type, say), so " * + "which child each value belongs to is unknown; write the values as a " * + "fresh column instead (Arrow.write(path, DataFrame(table)) infers the " * + "column's type from them) or drop the column", + ), + ) + end + # 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. + fn, dn = _constructpart(f.name, v; context) + # Polars' and Arrow 2.x's 128-bit Int type, which pyarrow cannot read, + # rewrites as a fresh Int128 column writes: fixed-size binary plus the + # JuliaLang.(U)Int128 label. + wide = t isa AC.IntType && t.bits == 128 && fn.type isa AC.FixedSizeBinaryType + wide || + 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", + ), + ) + metadata = _fieldmetadata(f) + if wide + metadata = Pair{String,String}[ + p for p in something(metadata, Pair{String,String}[]) if + !startswith(first(p), "ARROW:extension:") + ] + append!(metadata, something(_fieldmetadata(fn), Pair{String,String}[])) + end + rebuilt = AC.Field( + f.name, + fn.type; + nullable=f.nullable, + metadata=metadata, + 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 + +""" +A dictionary pool hint whose first `prefix` categories are kept exactly (a +pool already written or retained from an Arrow source); the categories after +them are candidates, deduplicated by storage value like a column's own. +""" +struct _PoolHint{V<:AbstractVector} + values::V + prefix::Int +end +_hintvalues(hint::_PoolHint) = hint.values +_hintvalues(hint::AbstractVector) = hint +_hintprefix(hint::_PoolHint) = hint.prefix +_hintprefix(hint::AbstractVector) = length(hint) + +# Dictionary category identity of a stored value: its type and value, a +# float's by its bits (NaN payloads, and 0.0 and -0.0, are distinct +# categories). +_poolkey(x::AbstractFloat) = (typeof(x), reinterpret(Unsigned, x)) +_poolkey(x) = (typeof(x), x) + +# A candidate category before its pool is built: the value itself (`===`, +# which no `isequal` method widens), or the contents of an array of plain +# values (a list row). Building the pool then merges categories stored alike. +struct _Same + x::Any +end +Base.isequal(a::_Same, b::_Same) = a.x === b.x +Base.hash(a::_Same, h::UInt) = hash(objectid(a.x), h) +_candidatekey(x) = + x isa Array && eltype(x) <: Union{Missing,Base.BitInteger,Base.IEEEFloat,Bool,String} ? + _poolkey(x) : _Same(x) + +"The category (`_poolkey`) of value `i` of a built column, with its union child." +function _storedkey(f::AC.Field, d::AC.ArrayData, i::Int64) + d.type isa AC.UnionType || return _poolkey(AC.getvalue(f, d, i)) + cf, cd, ci = AC._union_child(f, d, i) + return (findfirst(c -> c === cf, f.children), _storedkey(cf, cd, ci)) +end + +""" +The first position of each distinct stored category of a built pool +(`selected`) and every position's 0-based index among them (`remap`). The +first `prefix` categories (a pool already written or retained) are all kept, +in order. +""" +function _compactpool(f::AC.Field, d::AC.ArrayData, prefix::Int) + selected = Int[] + remap = Vector{Int64}(undef, d.len) + seen = Dict{Any,Int64}() + for i = 1:(d.len) + key = _storedkey(f, d, Int64(i)) + index = i <= prefix ? Int64(-1) : get(seen, key, Int64(-1)) + if index < 0 + index = Int64(length(selected)) + push!(selected, i) + haskey(seen, key) || (seen[key] = index) + end + remap[i] = index + end + return selected, remap +end + +"Rows `rows` of built column `d`, as one column." +_takerows(f::AC.Field, d::AC.ArrayData, rows) = _dictionaryconcat( + f, + _DictionaryRange[_DictionaryRange(d, Int64(i - 1), Int64(1)) for i in rows], + Limits(), + nothing, +) + +""" +An element type that holds both `T` and `S` values: `T` when it already +does, else their common supertype, with `missing` if either has it. +""" +function _widenedtype(T::Type, S::Type) + S <: T && return T + W = typejoin(Base.nonmissingtype(T), Base.nonmissingtype(S)) + return Missing <: T || Missing <: S ? Union{Missing,W} : W +end + +"One value pool: keep a retained prefix exactly and append new categories." +function _dictionarypool(v::AbstractVector, retainedpool; writetype=nothing) + if retainedpool === nothing + T = writetype === nothing ? Base.nonmissingtype(_eltypeof(v)) : writetype + T === Union{} && throw( + ArgumentError( + "cannot infer a dictionary value type from empty or all-missing columns", + ), + ) + pool = T[] + else + # New categories may be other types than the retained ones (writer + # values after a public pool read back, DataStrings after Strings): + # widen the value type instead of converting. + pool = _collectas( + _widenedtype(_eltypeof(retainedpool), Base.nonmissingtype(_eltypeof(v))), + retainedpool, + ) + end + # Candidates dedupe by `_candidatekey`, so the pool holds about one value + # per category, not one per row. + seen = Dict{Any,Nothing}() + for x in pool + x === missing || (seen[_candidatekey(x)] = nothing) + end + _addcategories!(pool, seen, v) + return pool +end + +# Function barrier: `pool` and each part of `vals` are only known at run time. +function _addcategories!(pool, seen, v) + for x in v + x === missing && continue + key = _candidatekey(x) + if !haskey(seen, key) + push!(pool, x) + seen[key] = nothing + end + end + return nothing +end + +"First pool position for each non-null value; duplicate categories stay intact." +function _dictionarylookup(pool) + lookup = Dict{Any,Int64}() + for (i, x) in enumerate(pool) + x === missing && continue + key = _candidatekey(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[_candidatekey(x)] for x in v + ] +end + +function _remapdictionaryindices(indices, remap::Vector{Int64}) + return Union{Missing,Int64}[ + index === missing ? missing : remap[Int(index) + 1] for index in indices + ] +end + +""" +The 0-based codes of a dictionary column read from Arrow (or a range of its +rows) when its pool is the first `prefix` categories of `hint`, else nothing. +""" +function _dictcodes(v, hint, prefix::Int) + hint === nothing && return nothing + parentcol, rows = v isa SubArray ? (parent(v), only(parentindices(v))) : (v, nothing) + parentcol isa DictEncoded && (rows === nothing || rows isa UnitRange{Int}) || + return nothing + pool = parentcol.pool + length(pool) <= prefix && all(i -> isequal(pool[i], hint[i]), eachindex(pool)) || + return nothing + return rows === nothing ? parentcol.indices : view(parentcol.indices, rows) +end + +"One shared-pool dictionary batch: identical pool OBJECT across batches." +function _dictbatch(fld::AC.Field, indices::AbstractVector, pool_d::AC.ArrayData) + t = fld.type::AC.DictionaryType + present = [x !== missing for x in indices] + # The index type is a run-time property of the field. + inds = _uninferred(_dictindexbuffer, AC.juliatype(t.indextype), indices)::AC.BufferSlice + nc = count(!, present) + d = AC.ArrayData( + t, + length(indices), + [AC._bitmapbuffer(present), inds]; + dictionary=pool_d, + nullcount=nc, + ) + return d +end + +_dictindexbuffer(::Type{IT}, indices) where {IT} = + AC._databuffer(IT[x === missing ? zero(IT) : IT(x) for x in indices]) + +"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::AbstractVector, + 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 + +# A column with a DataAPI ref pool (a PooledArray or CategoricalArray, say) is +# written dictionary-encoded, as Arrow 2.x did: its pool, in pool order and +# with unused entries, starts the Arrow dictionary. `refpools=false` writes +# such a column as its values instead (a later table's column whose schema +# field is not a dictionary). +_hasrefpool(x::AbstractVector) = + DataAPI.refpool(x) !== nothing && + DataAPI.refarray(x) !== x && + Base.nonmissingtype(_eltypeof(x)) !== Union{} +_hasrefpool(::Any) = false + +# Whether a column holds UUIDs (an abstract one, only UUIDs and missing +# values), which a pooled column or `dictencode` writes unencoded: pyarrow +# cannot read a dictionary of arrow.uuid values. +function _uuidvalues(x::AbstractVector) + T = Base.nonmissingtype(_eltypeof(x)) + T === UUIDs.UUID && return true + isconcretetype(T) && return false + # The scan compiles for the column's own type (see `_uninferred`). + return _uninferred(_onlyuuids, x)::Bool +end +_uuidvalues(::Any) = false +_onlyuuids(x) = + all(y -> y === missing || y isa UUIDs.UUID, x) && any(y -> y isa UUIDs.UUID, x) + +""" +The values and category pool of a ref-pool column. Pool values that are +wrappers without an ArrowTypes.jl mapping (CategoricalArrays' values) unwrap +through `DataAPI.unwrap`; a mapped element type keeps its values so its +ArrowTypes.jl lowering applies. +""" +function _refpoolvalues(x::AbstractVector) + T = Base.nonmissingtype(_eltypeof(x)) + mapped = ArrowTypes.ArrowType(T) !== T || ArrowTypes.arrowname(T) !== Symbol() + pool = DataAPI.refpool(x) + unwrap = + !mapped && + _uninferred(any, v -> v !== missing && DataAPI.unwrap(v) !== v, pool)::Bool + values = unwrap ? [DataAPI.unwrap(v) for v in x] : x + hint = Base.nonmissingtype(_eltypeof(values))[] + # Iterate: CategoricalArrays' ref pool is indexed from zero. + for v in pool + v === missing || push!(hint, unwrap ? DataAPI.unwrap(v) : v) + end + return values, hint +end + +""" +Construct one logical column of a record batch. + +This is the sole column-policy seam used by the write facade. It owns fresh +inference, retained reconstruction, ArrowTypes lowering, dictionary pooling, +field metadata, and the check that strings are UTF-8 (`_InvalidUTF8`). + +`retained` is the Field a rewrite must reproduce; `poolhint` carries the +retained dictionary-pool snapshot (or nothing), enforced below; `metadata` is +explicit application metadata merged onto the final Field. +""" +function _constructcolumn(name::Symbol, v::AbstractVector; kwargs...) + try + field, data = _constructcolumnparts(name, v; kwargs...) + _checkutf8(name, field, data) + return field, data + catch err + prefix = "fromjulia: unsupported element type " + msg = err isa ArgumentError ? err.msg : nothing + msg isa String && startswith(msg, prefix) || rethrow() + T = _reflectedstructtype(v) + T === nothing && throw( + ArgumentError( + "column $name holds $(chopprefix(msg::String, prefix)) values, which have " * + "no Arrow type; convert them first, or define ArrowTypes.ArrowType and " * + "ArrowTypes.toarrow for their type", + ), + ) + throw( + ArgumentError( + "column $name holds $T values, which have no ArrowTypes.jl mapping, so " * + "Arrow writes their fields as a Struct; a field failed with " * + "\"$msg\". Define ArrowTypes.ArrowType and ArrowTypes.toarrow " * + "for $T, or convert the column first (DataAPI.unwrap.(column) for " * + "wrapped values such as CategoricalValue)", + ), + ) + end +end + +# The element type of a column written by reflecting a plain struct's fields +# (ArrowTypes.jl's default StructKind mapping), or nothing. +function _reflectedstructtype(v) + T = Base.nonmissingtype(_eltypeof(v)) + isconcretetype(T) && isstructtype(T) && !(T <: Union{NamedTuple,Tuple}) || + return nothing + ArrowTypes.ArrowType(T) === T && ArrowTypes.arrowname(T) === Symbol() || return nothing + return ArrowTypes.ArrowKind(T) isa ArrowTypes.StructKind ? T : nothing +end + +function _constructcolumnparts( + name::Symbol, + v::AbstractVector; + retained::Union{Nothing,AC.Field}=nothing, + poolhint=nothing, + metadata::Union{Nothing,Vector{Pair{String,String}}}=nothing, + refpools::Bool=true, + stringview::Bool=false, + largelists::Bool=false, + growable::Bool=false, + pool=nothing, +) + if retained !== nothing && _isbool8(retained) + # arrow.bool8 reads as Bool; its storage is Int8 0 or 1. + _, int8data = _constructcolumnparts( + name, + _uninferred(_bool8storage, v); + retained=_withoutownextension(retained), + metadata, + ) + return _withfieldmetadata(_retainedfield(retained), metadata), int8data + end + fieldname = String(name) + dictintent = + (retained !== nothing && retained.type isa AC.DictionaryType) || + v isa DictEncode || + (refpools && retained === nothing && _hasrefpool(v) && !_uuidvalues(v)) + # A dictionary's pool outlives its batch, so it never takes reused buffers. + context = + _WriterContext(fieldname; stringview, largelists, pool=dictintent ? nothing : pool) + # A ref-pool column written as its values still unwraps them. + !dictintent && _hasrefpool(v) && (v = first(_uninferred(_refpoolvalues, v)::Tuple)) + # Built by calls this function does not infer (see `_uninferred`). + field, data = + dictintent ? + _uninferred( + _dictcolumnparts, + name, + fieldname, + v, + poolhint, + retained, + growable, + context, + )::Tuple{AC.Field,AC.ArrayData} : + _uninferred( + _plaincolumnparts, + name, + fieldname, + v, + retained, + context, + )::Tuple{AC.Field,AC.ArrayData} + return _withfieldmetadata(field, metadata), data +end + +"A dictionary column of `_constructcolumnparts`: its field and data." +function _dictcolumnparts(name, fieldname, part, poolhint, retained, growable, context) + values = part isa DictEncode ? (part::DictEncode).data : part + if _hasrefpool(values) + values, hint = _uninferred(_refpoolvalues, values)::Tuple + poolhint === nothing && (poolhint = hint) + end + 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 + + # The hint's first categories (a pool already written or retained from an + # Arrow source) are kept exactly; the categories after them dedupe. + retainedprefix = poolhint === nothing ? 0 : _hintprefix(poolhint) + poolhint === nothing || (poolhint = _hintvalues(poolhint)) + # A dictionary column read from Arrow whose pool starts the hint keeps its + # indices; others index their values into the pool. + codes = _dictcodes(values, poolhint, retainedprefix) + 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", + ), + ) + T = Base.nonmissingtype(_eltypeof(values)) + if retainedvaluefield === nothing + publictype = _facadebasetype(retained.type) + # Any string writes to a utf8 pool. + publictype === String && (publictype = AbstractString) + publictype !== Any && + !(T <: publictype) && + !_retainedwritertype((retained.type::AC.DictionaryType).valuetype, T) && + 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, values)) + 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. + !retained.nullable && + any(ismissing, something(codes, values)) && + (codes !== nothing || poolhint === nothing || !any(ismissing, poolhint)) && + throw( + ArgumentError( + "column $name may hold missing values but its retained " * + "dictionary field is non-nullable", + ), + ) + (retained.type::AC.DictionaryType).ordered && + poolhint === nothing && + throw( + ArgumentError( + "column $name has an ordered retained dictionary, but " * + "its original category pool is unavailable", + ), + ) + end + + if codes === nothing + pool = _dictionarypool(values, poolhint) + lookup = _dictionarylookup(pool) + missingindex = + retained !== nothing && !retained.nullable ? findfirst(ismissing, pool) : + nothing + missingindex === nothing || (missingindex = Int64(missingindex - 1)) + candidateindices = _dictionaryindices(values, lookup, missingindex) + else + pool = poolhint + candidateindices = codes + end + candidatevaluefield, candidatevaluedata = if retained === nothing + f, d = _constructpart(fieldname, pool; context) + # A dictionary field's extension label is, to pyarrow, the label of the + # dictionary type itself, which arrow.uuid's storage rules out: pyarrow + # 25 refuses such a file, its own output included. + _fieldmetavalue(f, _EXTENSION_NAME_KEY) == _UUID_EXTENSION && throw( + ArgumentError( + "column $name holds UUIDs, which Arrow.jl does not dictionary-encode: " * + "pyarrow cannot read a dictionary of arrow.uuid values. Write the " * + "column without Arrow.DictEncode", + ), + ) + f, d + else + retainedtype = retained.type::AC.DictionaryType + constructionfield = + retainedvaluefield === nothing ? AC.dictvaluefield(retained, retainedtype) : + retainedvaluefield + _constructpart(constructionfield, pool; context) + end + # Candidates stored alike are one category. + valuefield = candidatevaluefield + selected, remap = _compactpool(valuefield, candidatevaluedata, retainedprefix) + valuedata = + length(selected) == candidatevaluedata.len ? candidatevaluedata : + _takerows(valuefield, candidatevaluedata, selected) + indices = + length(selected) == length(remap) ? candidateindices : + _remapdictionaryindices(candidateindices, remap) + # `indices`' type is known at run time; see `_uninferred`. + retained === nothing && return _uninferred( + _newdictfromdata, + fieldname, + valuefield, + valuedata, + indices; + nullable=Missing <: _eltypeof(values), + # The narrowest index type the categories need; a pool that later + # tables may extend takes at least Int32, as pyarrow writes. + indextype=_dictindextype(growable ? max(valuedata.len, 1 << 16) : valuedata.len), + # ArrowCategoricalArraysExt adds an `_isordered` method. + ordered=_uninferred(_isordered, part)::Bool, + )::Tuple{AC.Field,AC.ArrayData} + return _uninferred( + _retaineddictfromdata, + retained, + valuefield, + valuedata, + Int(valuedata.len), + indices, + fieldname, + )::Tuple{AC.Field,AC.ArrayData} +end + +"A non-dictionary column of `_constructcolumnparts`: its field and data." +function _plaincolumnparts(name, fieldname, part, retained, context) + narrowed = + retained === nothing ? _narrowabstractcolumn(fieldname, part, context) : nothing + try + # A successful narrowing pass already saw every value; replanning + # from part of them can split one declared nullable Tuple child into + # incompatible Int and Null Fields. + retained === nothing || return _constructpart(retained, part; context) + return _constructpart( + fieldname, + something(narrowed, part); + context, + narrowabstract=(narrowed === nothing), + ) + catch err + (err isa _OffsetOverflow && err.part == 0) || rethrow() + e = err::_OffsetOverflow + # The overflow row is a row of this column only at its top level. + throw(_OffsetOverflow(e.column, 1, e.column == fieldname ? e.row : 0)) + end +end diff --git a/src/columns.jl b/src/columns.jl new file mode 100644 index 00000000..2fd855c0 --- /dev/null +++ b/src/columns.jl @@ -0,0 +1,2592 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# ============================================================================= +# Read-side columns for flat layouts: zero-copy views over decoded Arrow +# buffers, DataStrings string columns, and dictionary vectors. +# +# A view wraps its buffers as unowned Julia vectors (so loops vectorize like +# `Vector` loops) and holds the decoded `ArrayData`, whose buffer slices keep +# the backing memory (a memory map, the input bytes, or a decompressed +# buffer) reachable. Every access preserves the view, so memory a view can +# read is memory that is alive. Views check no release flag: the facade never +# unmaps a source eagerly (see `release!(::Table)`). +# ============================================================================= + +""" + Arrow.ArrowVector{T} <: AbstractVector{T} + +Supertype of the read-only columns `Arrow.Table` and `Arrow.Stream` return +for flat, list, map, and struct layouts. A flat column reads its Arrow buffers in +place (a memory-mapped file, the input bytes, or a decompressed buffer) and +keeps that memory alive while the column is reachable, independent of its +table. `copy(column)` returns an owned, mutable vector: a `Vector`, or a +`PooledArray` for a dictionary-encoded column. +""" +abstract type ArrowVector{T} <: AbstractVector{T} end + +Base.IndexStyle(::Type{<:ArrowVector}) = IndexLinear() + +@noinline _readonlycolumn() = error( + "Arrow columns are read-only views of the data they were read from; " * + "copy(column), copy(table), or DataFrame(table) returns a mutable copy", +) +Base.setindex!(::ArrowVector, v, ::Int) = _readonlycolumn() +# `push!` and `append!` resize first. +Base.resize!(::ArrowVector, ::Integer) = _readonlycolumn() +# Containers that know only a column's type (a `ChainedVector` whose batches +# are all empty) allocate through this; a new vector is a mutable `Array`. +Base.similar(::Type{<:ArrowVector{T}}, dims::Dims) where {T} = Array{T}(undef, dims) + +# --- buffer wrapping ------------------------------------------------------------ + +""" +A read-only vector of `S` values in memory that need not be aligned for `S`, +read with unaligned loads. Arrow IPC buffers are aligned to 8 bytes (pyarrow +writes no more), so 16- and 32-byte values (`Int128`, 128- and 256-bit +decimals) start at a 16-byte boundary only half of the time. `owner` keeps +the memory alive. +""" +struct _UnalignedVector{S} <: AbstractVector{S} + ptr::Ptr{S} + len::Int + owner::Any +end +Base.size(v::_UnalignedVector) = (v.len,) +Base.IndexStyle(::Type{<:_UnalignedVector}) = IndexLinear() +Base.@propagate_inbounds function Base.getindex(v::_UnalignedVector, i::Int) + @boundscheck checkbounds(v, i) + owner = v.owner + # `unsafe_load` loads with alignment 1. + return GC.@preserve owner unsafe_load(v.ptr, i) +end +function Base.copyto!( + dest::Vector{S}, + doffs::Integer, + src::_UnalignedVector{S}, + soffs::Integer, + n::Integer, +) where {S} + n == 0 && return dest + ( + n > 0 && + checkbounds(Bool, dest, doffs:(doffs + n - 1)) && + checkbounds(Bool, src, soffs:(soffs + n - 1)) + ) || throw(BoundsError(dest, doffs:(doffs + n - 1))) + owner = src.owner + GC.@preserve dest owner unsafe_copyto!( + Ptr{UInt8}(pointer(dest, doffs)), + Ptr{UInt8}(src.ptr + (soffs - 1) * sizeof(S)), + n * sizeof(S), + ) + return dest +end + +Base.copy(v::_UnalignedVector{S}) where {S} = + copyto!(Vector{S}(undef, v.len), 1, v, 1, v.len) + +# How a column stores its `S` values: in place with unaligned loads when `S` +# needs more than the 8-byte alignment Arrow buffers have, else a `Vector`. +_buffertype(::Type{S}) where {S} = + Base.datatype_alignment(S) > 8 ? _UnalignedVector{S} : Vector{S} + +# An unowned vector over `n` elements of `S` starting `offset` elements into +# `b` (a `_buffertype(S)`). Callers keep `b`'s region reachable for as long +# as the vector is used. A start misaligned for `S` (possible only for +# foreign or hand-built buffers) gets an owned, aligned `Vector` copy. +function _wrapbuffer( + ::Type{S}, + b::AC.BufferSlice, + offset::Int64, + n::Int64, + budget=nothing, +) where {S} + isbitstype(S) || throw(ArgumentError("a wrapped buffer needs an isbits element type")) + unaligned = _buffertype(S) === _UnalignedVector{S} + n == 0 && return unaligned ? _UnalignedVector{S}(Ptr{S}(C_NULL), 0, nothing) : S[] + w = Int64(sizeof(S)) + src = AC.subslice(b, AC.checked_mul(offset, w), AC.checked_mul(n, w)) + p = AC.sliceptr(src) + unaligned && return _UnalignedVector{S}(Ptr{S}(p), Int(n), src) + if UInt(p) % Base.datatype_alignment(S) != 0 + _chargeexact!(budget, S, n, "aligned buffer copy") + out = Vector{S}(undef, Int(n)) + GC.@preserve out src begin + # `S` is known at run time; see `_uninferred`. + unsafe_copyto!(Ptr{UInt8}(_uninferred(pointer, out)::Ptr), p, Int(src.len)) + end + return out + end + return unsafe_wrap(Array, Ptr{S}(p), Int(n)) +end + +# The validity bitmap bytes covering `d`'s slots and the bit position of slot +# 1 in them; an empty vector when no slot is null. +function _validitybits(d::AC.ArrayData) + v = AC.validitybuffer(d) + (AC.isempty_buffer(v) || AC.nullcount(d) == 0 || d.len == 0) && return UInt8[], 0 + first = d.offset >> 3 + nbytes = ((d.offset + d.len - 1) >> 3) - first + 1 + return _wrapbuffer(UInt8, v, first, nbytes), Int(d.offset & 7) +end + +@inline function _bitset(bits::Vector{UInt8}, bitoffset::Int, i::Int) + k = bitoffset + i - 1 + return @inbounds((bits[(k >> 3) + 1] >> (k & 7)) & 0x01) == 0x01 +end + +@inline _isnullslot(validity::Vector{UInt8}, bitoffset::Int, i::Int) = + !isempty(validity) && !_bitset(validity, bitoffset, i) + +# --- storage conversions -------------------------------------------------------- + +# `C` in `Primitive{T,S,C}` names how one stored `S` becomes a `T`; +# `Nothing` means the layouts are identical. +@inline _fromstorage(::Type{Nothing}, x) = x +struct _DaysToDate end +@inline _fromstorage(::Type{_DaysToDate}, x::Int32) = + Dates.Date(Dates.UTD(Int64(x) + _EPOCH_DAYS)) +struct _MillisToDateTime end +@inline _fromstorage(::Type{_MillisToDateTime}, x::Int64) = + Dates.DateTime(Dates.UTM(x + Dates.UNIXEPOCH)) +"Time of day stored in ticks of `N` nanoseconds (range-checked when built)." +struct _TimeOfDay{N} end +@inline _fromstorage(::Type{_TimeOfDay{N}}, x) where {N} = + Dates.Time(Dates.Nanosecond(Int64(x) * N)) +struct _YearMonth end +@inline _fromstorage(::Type{_YearMonth}, x::Int32) = Durations.Duration(x, 0, 0) +struct _DayTime end +@inline _fromstorage(::Type{_DayTime}, x::NTuple{2,Int32}) = + Durations.Duration(0, x[1], Int64(x[2]) * 1_000_000) +"Stored bytes as a byte vector (the raw read of a decimal no `Decimal` type holds)." +struct _StoredBytes end +@inline _fromstorage(::Type{_StoredBytes}, x::NTuple{N,UInt8}) where {N} = collect(x) + +# The element type each conversion produces. +_storagetarget(::Type{_DaysToDate}) = Dates.Date +_storagetarget(::Type{_MillisToDateTime}) = Dates.DateTime +_storagetarget(::Type{<:_TimeOfDay}) = Dates.Time +_storagetarget(::Type{_YearMonth}) = Durations.Duration +_storagetarget(::Type{_DayTime}) = Durations.Duration +_storagetarget(::Type{_StoredBytes}) = Vector{UInt8} + +# --- fixed width ------------------------------------------------------------------ + +""" + Arrow.Primitive{T,S,C,D} <: Arrow.ArrowVector{T} + +A fixed-width column: element `i` is the `i`th stored `S` value converted to +`T` (the same bits when the layouts match), or `missing` where the validity +bitmap marks a null. `D` is how the `S` values are held: a `Vector{S}`, or +for 16- and 32-byte values, a view of their buffer read with unaligned loads. +""" +struct Primitive{T,S,C,D<:AbstractVector{S}} <: ArrowVector{T} + data::D + validity::Vector{UInt8} + bitoffset::Int + owner::Any +end +Primitive{T,S,C}(data::D, validity, bitoffset, owner) where {T,S,C,D<:AbstractVector{S}} = + Primitive{T,S,C,D}(data, validity, bitoffset, owner) + +Base.size(v::Primitive) = size(v.data) + +Base.@propagate_inbounds function Base.getindex(v::Primitive{T,S,C}, i::Int) where {T,S,C} + @boundscheck checkbounds(v, i) + x = GC.@preserve v @inbounds(v.data[i]) + if Missing <: T + null = GC.@preserve v _isnullslot(v.validity, v.bitoffset, i) + null && return missing + end + return _fromstorage(C, x) +end + +function _primitive( + ::Type{T}, + ::Type{S}, + ::Type{C}, + d::AC.ArrayData, + budget=nothing, +) where {T,S,C} + data = _wrapbuffer(S, AC.rolebuffer(d, AC.DATA), d.offset, d.len, budget) + validity, bitoffset = _validitybits(d) + return Primitive{T,S,C}(data, validity, bitoffset, d) +end + +# One pass over the buffer, and its validity bitmap where some slot is null: +# the stored values are the values. +function Base.copyto!( + dest::Vector{T}, + doffs::Integer, + src::Primitive{T,S,Nothing}, + soffs::Integer, + n::Integer, +) where {T,S} + n == 0 && return dest + ( + n > 0 && + checkbounds(Bool, dest, doffs:(doffs + n - 1)) && + checkbounds(Bool, src, soffs:(soffs + n - 1)) + ) || throw(BoundsError(dest, doffs:(doffs + n - 1))) + GC.@preserve src begin + data, bits, off = src.data, src.validity, src.bitoffset + if isempty(bits) + copyto!(dest, doffs, data, soffs, n) + else + # Validity is present only when some slot is null, so `T` has `Missing`. + @inbounds for i = 0:(n - 1) + dest[doffs + i] = _bitset(bits, off, soffs + i) ? data[soffs + i] : missing + end + end + end + return dest +end + +Base.copy(v::Primitive{T,S,Nothing}) where {T,S} = + copyto!(Vector{T}(undef, length(v)), 1, v, 1, length(v)) + +# Base's generic `map`, `filter`, and `collect` do not vectorize through the +# per-access preserve. A view whose values are its stored layout (no nulls +# can occur in a non-`Missing` element type) runs them on its buffer. +Base.map(f, v::Primitive{T,T,Nothing}) where {T} = GC.@preserve v map(f, v.data) +Base.filter(f, v::Primitive{T,T,Nothing}) where {T} = GC.@preserve v filter(f, v.data) +Base.collect(v::Primitive{T,T,Nothing}) where {T} = GC.@preserve v copy(v.data) + +# Reductions over a view with no nulls run on its buffer, so they vectorize +# as they do over a `Vector`. +const _BufferPrimitive = Primitive{T,S,Nothing} where {T,S} +function Base.mapreduce(f, op, A::_BufferPrimitive; kw...) + isempty(A.validity) && return GC.@preserve A mapreduce(f, op, A.data; kw...) + return invoke(mapreduce, Tuple{Any,Any,AbstractArray}, f, op, A; kw...) +end +function Base.mapreduce(f, op, itr::Base.SkipMissing{<:_BufferPrimitive}) + A = itr.x + isempty(A.validity) && return GC.@preserve A mapreduce(f, op, A.data) + _maskablesum(f, op, A) && return GC.@preserve A _maskedsum(A) + return invoke(mapreduce, Tuple{Any,Any,Base.SkipMissing{<:AbstractArray}}, f, op, itr) +end +function Base.mapreduce( + f, + op, + itr::Base.SkipMissing{<:ChainedVector{T,<:_BufferPrimitive}}, +) where {T} + chunks = itr.x.arrays + if !isempty(chunks) && + (all(c -> isempty(c.validity), chunks) || all(c -> _maskablesum(f, op, c), chunks)) + return reduce(op, (mapreduce(f, op, skipmissing(c)) for c in chunks)) + end + return invoke(mapreduce, Tuple{Any,Any,Base.SkipMissing{<:AbstractArray}}, f, op, itr) +end + +# `sum(skipmissing(col))` over numbers with nulls: a branch-free pairwise sum +# in which null slots add the neutral element (`-0.0` keeps the sign of an +# all-`-0.0` float sum). +_maskablesum(f, op, ::Primitive) = false +_maskablesum( + ::typeof(identity), + ::Union{typeof(+),typeof(Base.add_sum)}, + ::Primitive{T,<:Union{Base.BitInteger,Base.IEEEFloat},Nothing}, +) where {T} = true + +function _maskedsum(A::Primitive{T,S}) where {T,S} + R = typeof(Base.add_sum(zero(S), zero(S))) + neutral = S <: AbstractFloat ? -zero(R) : zero(R) + isempty(A.data) && return neutral + return _maskedsum(A.data, A.validity, A.bitoffset, 1, length(A.data), neutral) +end +function _maskedsum( + data::AbstractVector, + bits::Vector{UInt8}, + off::Int, + lo::Int, + hi::Int, + neutral::R, +) where {R} + if hi - lo < 1024 + s = neutral + @inbounds @simd for i = lo:hi + k = off + i - 1 + valid = (bits[(k >> 3) + 1] >> (k & 7)) & 0x01 == 0x01 + s += ifelse(valid, R(data[i]), neutral) + end + return s + end + mid = (lo + hi) >>> 1 + return _maskedsum(data, bits, off, lo, mid, neutral) + + _maskedsum(data, bits, off, mid + 1, hi, neutral) +end + +# --- booleans --------------------------------------------------------------------- + +""" + Arrow.BoolVector{T} <: Arrow.ArrowVector{T} + +A boolean column read from its bit-packed values and validity bitmaps. +""" +struct BoolVector{T} <: ArrowVector{T} + values::Vector{UInt8} + valueoffset::Int + validity::Vector{UInt8} + bitoffset::Int + len::Int + owner::Any +end + +Base.size(v::BoolVector) = (v.len,) + +Base.@propagate_inbounds function Base.getindex(v::BoolVector{T}, i::Int) where {T} + @boundscheck checkbounds(v, i) + if Missing <: T + null = GC.@preserve v _isnullslot(v.validity, v.bitoffset, i) + null && return missing + end + return GC.@preserve v _bitset(v.values, v.valueoffset, i) +end + +function _boolvector(::Type{T}, d::AC.ArrayData) where {T} + validity, bitoffset = _validitybits(d) + d.len == 0 && return BoolVector{T}(UInt8[], 0, validity, bitoffset, 0, d) + first = d.offset >> 3 + nbytes = ((d.offset + d.len - 1) >> 3) - first + 1 + values = _wrapbuffer(UInt8, AC.rolebuffer(d, AC.DATA), first, nbytes) + return BoolVector{T}(values, Int(d.offset & 7), validity, bitoffset, Int(d.len), d) +end + +# Non-null counts come from the validity bitmap. +Base.count( + ::ComposedFunction{typeof(!),typeof(ismissing)}, + A::Union{Primitive,BoolVector}, +) = length(A) - _nullslots(A) +function _nullslots(A::Union{Primitive,BoolVector}) + n = length(A) + (isempty(A.validity) || n == 0) && return 0 + bits, off = A.validity, A.bitoffset + stop = off + n - 1 + last = stop >> 3 + 1 + valid = 0 + GC.@preserve A begin + @inbounds @simd for k = 1:last + valid += count_ones(bits[k]) + end + # Drop the bits before slot 1 and after slot n. + valid -= count_ones(bits[1] & ~(0xff << off)) + valid -= count_ones(bits[last] & ~(0xff >> (7 - (stop & 7)))) + end + return n - valid +end + +# --- run-end encoding ------------------------------------------------------------- + +""" + Arrow.RunEndEncoded{T,V} <: Arrow.ArrowVector{T} + +A run-end encoded column: element `i` is the value of the run that covers +it, found by binary search over the run ends. +""" +struct RunEndEncoded{T,V<:AbstractVector} <: ArrowVector{T} + runends::Vector{Int64} + values::V + offset::Int + len::Int +end + +Base.size(v::RunEndEncoded) = (v.len,) + +Base.@propagate_inbounds function Base.getindex(v::RunEndEncoded{T}, i::Int) where {T} + @boundscheck checkbounds(v, i) + # Validation proved the run ends ascending and covering every slot. + run = searchsortedfirst(v.runends, v.offset + i) + return @inbounds(v.values[run])::T +end + +# --- dictionary encoding ---------------------------------------------------------- + +""" + Arrow.DictEncoded{T} <: Arrow.ArrowVector{T} + +A dictionary-encoded column: each row is an index into a pool of values. +`DataAPI.refarray`, `DataAPI.refpool`, and `DataAPI.levels` (in dictionary order) +expose the encoding without decoding it, and `copy` returns a `PooledArray`. With +CategoricalArrays loaded, `isordered` tells whether the dictionary is ordered. +""" +struct DictEncoded{T,R,P<:AbstractVector,I<:AbstractVector} <: ArrowVector{T} + indices::I + pool::P + ordered::Bool +end + +Base.size(v::DictEncoded) = size(v.indices) + +Base.@propagate_inbounds function Base.getindex(v::DictEncoded{T}, i::Int) where {T} + @boundscheck checkbounds(v, i) + code = @inbounds v.indices[i] + code === missing && return missing + # Validation proved every valid index in [0, length(pool)). + return @inbounds(v.pool[Int(code) + 1])::T +end + +""" +The `DataAPI.refarray` of a dictionary column: the one-based pool position +of each value, with 0 for a null slot. +""" +struct DictRefs{R,I<:AbstractVector} <: AbstractVector{R} + indices::I +end +Base.size(v::DictRefs) = size(v.indices) +Base.IndexStyle(::Type{<:DictRefs}) = IndexLinear() +Base.@propagate_inbounds function Base.getindex(v::DictRefs{R}, i::Int) where {R} + code = v.indices[i] + return code === missing ? zero(R) : R(code) + one(R) +end + +""" +The `DataAPI.refpool` of a dictionary column whose codes can be null: index +0 holds `missing`, then the pool values in order. +""" +struct DictRefPool{T,P<:AbstractVector} <: AbstractVector{T} + pool::P +end +Base.size(v::DictRefPool) = (length(v.pool) + 1,) +Base.axes(v::DictRefPool) = (Base.IdentityUnitRange(0:length(v.pool)),) +Base.IndexStyle(::Type{<:DictRefPool}) = IndexLinear() +Base.@propagate_inbounds Base.getindex(v::DictRefPool{T}, i::Int) where {T} = + i == 0 ? missing : convert(T, v.pool[i]) + +# Codes widen by one step so `code + 1` never overflows. +_reftype(::Type{R}) where {R<:Union{Int8,UInt8,Int16,UInt16}} = Int32 +_reftype(::Type{R}) where {R<:Union{Int32,UInt32,Int64,UInt64}} = Int64 + +DataAPI.refarray(v::DictEncoded{T,R}) where {T,R} = + DictRefs{_reftype(R),typeof(v.indices)}(v.indices) +DataAPI.refpool(v::DictEncoded{T}) where {T} = + Missing <: T ? DictRefPool{T,typeof(v.pool)}(v.pool) : v.pool +DataAPI.levels(v::DictEncoded) = unique(skipmissing(v.pool)) + +# Chunks that share one pool keep a column-wide encoding. +_sharedpool(x::ChainedVector{T,<:DictEncoded}) where {T} = + all(c -> c.pool === x.arrays[1].pool, x.arrays) +# Categories in dictionary order, as one batch's are: those of later batches' +# dictionaries (deltas, replacements) follow the first's. +function DataAPI.levels(x::ChainedVector{T,<:DictEncoded}) where {T} + isempty(x.arrays) && return Base.nonmissingtype(T)[] + _sharedpool(x) && return DataAPI.levels(x.arrays[1]) + return unique!(reduce(vcat, [DataAPI.levels(c) for c in x.arrays])) +end +DataAPI.refarray(x::ChainedVector{T,<:DictEncoded}) where {T} = + isempty(x.arrays) || !_sharedpool(x) ? x : + ChainedVector([DataAPI.refarray(c) for c in x.arrays]) +DataAPI.refpool(x::ChainedVector{T,<:DictEncoded}) where {T} = + isempty(x.arrays) || !_sharedpool(x) ? nothing : DataAPI.refpool(x.arrays[1]) + +Base.copy(v::DictEncoded{T}) where {T} = _pooledcopy(T, (v,)) +# The batches of a column copy into one `PooledArray`, whether they share a +# pool, extend it (delta dictionaries), or replace it. +Base.copy(c::ChainedVector{T,<:DictEncoded}) where {T} = _pooledcopy(T, c.arrays) +# Rows of a dictionary column, as of a `PooledArray`, keep the encoding: any +# vector of integer indices or a `Bool` mask. +function Base.getindex(v::DictEncoded{T}, rows::AbstractVector{<:Integer}) where {T} + @boundscheck checkbounds(v, rows) + return _pooledcopy(T, (v,), to_indices(v, (rows,))[1]) +end +Base.getindex(c::ChainedVector{<:Any,<:DictEncoded}, rows::AbstractVector{<:Integer}) = + _dictrows(c, rows) +# As specific as SentinelArrays' methods for these index vectors, so not ambiguous. +Base.getindex(c::ChainedVector{<:Any,<:DictEncoded}, rows::AbstractVector{Int}) = + _dictrows(c, rows) +Base.getindex(c::ChainedVector{<:Any,<:DictEncoded}, rows::UnitRange{Int}) = + _dictrows(c, rows) +Base.getindex( + c::ChainedVector{<:Any,<:DictEncoded}, + rows::AbstractVector{<:ChainedVectorIndex}, +) = _dictrows(c, Int[Int(i) for i in rows]) +function _dictrows(c::ChainedVector{T}, rows) where {T} + checkbounds(c, rows) + return _pooledcopy(T, c.arrays, to_indices(c, (rows,))[1]) +end + +""" +Rows `rows` (all of them by default) of the dictionary column whose batches +are `chunks`, as a `PooledArray`. A batch's pool joins the result's when its +first row is copied (a copy of a column with no rows takes its batch's +pool): its distinct values, in pool order; duplicate and null entries remap, +so the PooledArray's invariants hold. +""" +function _pooledcopy(::Type{T}, chunks, rows=nothing) where {T} + R = _ownedtype(T) + invpool = Dict{R,UInt32}() + pool = R[] + ref(value) = get!(invpool, value) do + push!(pool, value) + UInt32(length(pool)) + end + remaps = IdDict{Any,Vector{UInt32}}() + remapof(v) = get!(() -> UInt32[ref(_ownedpoolvalue(x)) for x in v.pool], remaps, v.pool) + ends = cumsum(Int[length(v) for v in chunks]) + n = rows === nothing ? (isempty(ends) ? 0 : ends[end]) : length(rows) + refs = Vector{UInt32}(undef, n) + c = 1 + remap = UInt32[] + lastpool = nothing + for (k, i) in enumerate(rows === nothing ? (1:n) : rows) + # Ascending rows advance through the batches; others search. + if !(c <= length(ends) && (c == 1 ? 0 : ends[c - 1]) < i <= ends[c]) + c = searchsortedfirst(ends, i) + end + v = chunks[c] + if v.pool !== lastpool + remap = remapof(v) + lastpool = v.pool + end + code = @inbounds v.indices[i - (c == 1 ? 0 : ends[c - 1])] + refs[k] = code === missing ? ref(missing) : @inbounds remap[Int(code) + 1] + end + # A copy of an empty batch keeps its pool (a chained column has none empty). + rows === nothing && n == 0 && foreach(remapof, chunks) + return PooledArrays.PooledArray(PooledArrays.RefArray(refs), invpool, pool) +end + +# --- strings and binary ----------------------------------------------------------- + +const _StringPayload = DataStrings.StringPayload + +# Byte runs at most this long fit DataStrings' Int32 view words. +const _MAX_VIEW_BUFFER = Int64(typemax(Int32)) + +# Rows per chunk when a string column builds on several tasks. +const _STRING_CHUNK_ROWS = 1 << 16 + +_isstringlayout(t::AC.ArrowType) = + t isa Union{AC.Utf8Type,AC.BinaryType,AC.ViewType} || _widebinary(t) + +""" +Rows `first:last` of string batch `d`, whose output slots follow index `at`. +An offsets layout (`utf8`/`binary`, normal or large) copies the rows' bytes +into data buffer `buf` (0-based), where byte `a` of the batch's data buffer +sits at position `a - lo`. For a view layout `buf` is the batch's first data +buffer, and a chunk with `variadic > 0` copies that buffer of the batch +instead of encoding rows. +""" +struct _StringChunk + d::AC.ArrayData + at::Int + first::Int + last::Int + buf::Int + lo::Int64 + variadic::Int + copybytes::Bool # false: buffer `buf` already holds the batch's bytes (`_adoptbytes`) +end + +""" +A string or binary column under construction: the payloads and data buffers +of a `StringVector`/`BytesVector` of type `V`, or the values of a +`Vector{String}` (`V === String`). `_stringbuild` plans and charges it, +`_allocate!` allocates it; its chunks are independent, so several tasks can +fill them (`_buildchunk!`); `_finishbuild` returns the column. +""" +mutable struct _StringBuild{V,O} + n::Int # values + out::O # payloads or values, `n` long once allocated + buffers::Vector{Vector{UInt8}} # data buffers: kept bytes, or copies once allocated + sizes::Vector{Int64} # the bytes `_allocate!` gives buffer k (-1: kept) + bases::Vector{Int64} # where a kept buffer's batch bytes start in it + chunks::Vector{_StringChunk} +end + +function _allocate!(b::_StringBuild{V,O}) where {V,O} + b.out = O(undef, b.n) + for (k, len) in enumerate(b.sizes) + len >= 0 && (b.buffers[k] = Vector{UInt8}(undef, Int(len))) + end + return b +end + +# Call `f` with the `n + 1` offsets of offsets-layout batch `d`. A fixed-size +# binary batch's value `i` starts `(offset + i - 1) * width` bytes into its +# data buffer. +function _withoffsets(f, d::AC.ArrayData, budget=nothing) + t = d.type + n = Int64(d.len) + if t isa AC.FixedSizeBinaryType + w = Int64(t.nbytes) + o = Int64(d.offset) + return f((o * w):w:((o + n) * w)) + end + offs = AC.rolebuffer(d, AC.OFFSETS) + return t.large ? f(_wrapbuffer(Int64, offs, d.offset, n + 1, budget)) : + f(_wrapbuffer(Int32, offs, d.offset, n + 1, budget)) +end + +function _rowchunks!( + chunks, + d, + at::Int, + first::Int, + last::Int, + buf::Int, + lo::Int64, + rows::Int, + copybytes::Bool=true, +) + a = first + while a <= last + b = last - a >= rows ? a + rows - 1 : last + push!(chunks, _StringChunk(d, at, a, b, buf, lo, 0, copybytes)) + a = b + 1 + end + return chunks +end + +# The bytes behind slice `b` that a column can reference instead of copying, +# and where `b` starts in them: the decompressed buffer this read owns, or, +# with `source`, the source's own bytes (a mapped file, a byte vector, a +# fetched range) while `b` starts within the Int32 offsets a DataStrings view +# addresses. A `DataString` holds its buffer, so a value taken from such a +# column keeps the source bytes alive; `String(x)` copies it out. +function _adoptbytes(b::AC.BufferSlice, source::Bool) + r = b.region + r === nothing && return nothing + root = r.root + root isa _Decompressed && b.offset == 0 && return root.bytes, Int64(0) + source || return nothing + # From Julia 1.11 a byte vector's region roots its `Memory`; the column + # gets a vector of its own over that memory. + @static if VERSION >= v"1.11" + if root isa Memory{UInt8} + base = Int64(r.ptr - pointer(root)) + b.offset + base <= _MAX_VIEW_BUFFER || return nothing + return Base.wrap(Array, root, length(root)), base + end + end + root isa Vector{UInt8} && b.offset <= _MAX_VIEW_BUFFER && return root, b.offset + return nothing +end + +# Cut offsets batch `d` into runs of rows whose bytes span at most 2 GiB, one +# data buffer each, and the runs into chunks. +function _offsetruns!( + chunks, + buffers, + sizes, + bases, + d, + offsets::AbstractVector{O}, + at::Int, + budget, + rows::Int, + source::Bool, +) where {O} + n = Int(d.len) + # Bytes the column can reference in place need no copy. + adopted = _adoptbytes(AC.rolebuffer(d, AC.DATA), source) + if adopted !== nothing + own, base = adopted + if base + Int64(offsets[n + 1]) <= _MAX_VIEW_BUFFER + push!(buffers, own) + push!(sizes, -1) + push!(bases, base) + return _rowchunks!(chunks, d, at, 1, n, length(buffers) - 1, -base, rows, false) + end + end + i = 1 + while i <= n + lo = Int64(offsets[i]) + j = i + if Int64(offsets[n + 1]) - lo <= _MAX_VIEW_BUFFER + j = n + 1 + else + while j <= n && Int64(offsets[j + 1]) - lo <= _MAX_VIEW_BUFFER + j += 1 + end + end + j == i && throw( + ArgumentError( + "a string longer than 2 GiB does not fit a DataString; " * + "read the column with stringtype=String", + ), + ) + hi = Int64(offsets[j]) + _chargeexact!(budget, UInt8, hi - lo, "string column data") + push!(buffers, UInt8[]) + push!(sizes, hi - lo) + push!(bases, 0) + _rowchunks!(chunks, d, at, i, j - 1, length(buffers) - 1, lo, rows) + i = j + end + return chunks +end + +# The bytes the values of string batch `d` reference: its data span, or its +# variadic buffers. +_stringbytes(d::AC.ArrayData) = + d.type isa AC.ViewType ? + sum(k -> d.buffers[k].len, 3:length(d.buffers); init=Int64(0)) : + d.len == 0 ? Int64(0) : _withoffsets(o -> Int64(o[end]) - Int64(o[1]), d) + +""" +Plan the column of string or binary batches `datas` with element type `T`: +charge its output and data buffers (`_allocate!` allocates them), and cut its +rows into chunks of at most `rows`. A data buffer is the bytes the batch +already has in memory, referenced in place (`_adoptbytes`; with `source`, the +source's own bytes too), or else one copy of a run of the referenced bytes; +a run stays under 2 GiB so view offsets fit DataStrings' Int32 words. +""" +function _stringbuild( + ::Type{T}, + datas, + budget, + rows::Int=typemax(Int); + source::Bool=true, +) where {T} + B = Base.nonmissingtype(T) + n = sum(d -> Int(d.len), datas; init=0) + chunks = _StringChunk[] + buffers = Vector{UInt8}[] + sizes = Int64[] + bases = Int64[] + at = 0 + if B === String + _chargeexact!(budget, T, n, "string column") + bytes = Int64(0) + for d in datas + _rowchunks!(chunks, d, at, 1, Int(d.len), 0, Int64(0), rows) + bytes = AC.checked_add(bytes, _stringbytes(d)) + at += Int(d.len) + end + # One String object per value: its bytes plus a header and rounding. + budget === nothing || _charge!( + budget, + AC.checked_add(AC.checked_mul(Int64(n), Int64(32)), bytes), + "string values", + ) + return _StringBuild{String,Vector{T}}(n, T[], buffers, sizes, bases, chunks) + end + V = + B === DataStrings.DataString ? DataStrings.StringVector{T} : + DataStrings.BytesVector{T} + _chargeexact!(budget, _StringPayload, n, "string column payloads") + for d in datas + m = Int(d.len) + if d.type isa AC.ViewType + bufbase = length(buffers) # the batch's buffer 3 is column buffer `bufbase` + for k = 3:length(d.buffers) + len = d.buffers[k].len + adopted = _adoptbytes(d.buffers[k], source) + if adopted !== nothing && adopted[2] + len <= _MAX_VIEW_BUFFER + push!(buffers, adopted[1]) + push!(sizes, -1) + push!(bases, adopted[2]) + continue + end + _chargeexact!(budget, UInt8, len, "string column data") + push!(buffers, UInt8[]) + push!(sizes, len) + push!(bases, 0) + push!(chunks, _StringChunk(d, at, 1, 0, bufbase, Int64(0), k, true)) + end + _rowchunks!(chunks, d, at, 1, m, bufbase, Int64(0), rows) + elseif m > 0 + _withoffsets( + o -> _offsetruns!( + chunks, + buffers, + sizes, + bases, + d, + o, + at, + budget, + rows, + source, + ), + d, + budget, + ) + end + at += m + end + return _StringBuild{V,Vector{_StringPayload}}( + n, + _StringPayload[], + buffers, + sizes, + bases, + chunks, + ) +end + +_finishbuild(b::_StringBuild{String}) = b.out +_finishbuild(b::_StringBuild{V}) where {V} = V(b.out, b.buffers, Val(:trusted)) + +""" +A string or binary column of a whole-table read that builds when it is first +taken from its table (`_force`): the read charges its allowance up front, and +the payloads are built and the bytes copied on first access, on up to +`ntasks` tasks. One task builds; tasks that ask meanwhile wait for it. The +table hands out only the built column, so loops run over a `StringVector` +(`BytesVector`, `Vector{String}`) as for an eager read. +""" +mutable struct _PendingStrings{T} <: AbstractVector{T} + const build::_StringBuild + const ntasks::Int + const lock::ReentrantLock + @atomic column::Union{Nothing,AbstractVector} +end +_PendingStrings{T}(b::_StringBuild, ntasks::Int) where {T} = + _PendingStrings{T}(b, ntasks, ReentrantLock(), nothing) + +Base.size(p::_PendingStrings) = (p.build.n,) +Base.getindex(p::_PendingStrings, i::Int) = _force(p)[i] + +# A table column as it is handed out: a pending string column, built. +_force(c::AbstractVector) = c isa _PendingStrings ? _forcestrings(c) : c +function _forcestrings(p::_PendingStrings) + c = @atomic :acquire p.column + c === nothing || return c + return @lock p.lock begin + c = @atomic :acquire p.column + if c === nothing + b = _allocate!(p.build) + # `b`'s type is known at run time; see `_uninferred`. + build! = (k, _) -> (_buildchunk!(b, k); nothing) + _uninferred(_foreach, build!, length(b.chunks), p.ntasks) + c = _finishbuild(b) + @atomic :release p.column = c + # The chunks hold the batches' offsets, decompressed ones included. + empty!(b.chunks) + end + c + end +end + +"Fill chunk `k` of a string column under construction." +function _buildchunk!( + b::_StringBuild{V,Vector{_StringPayload}}, + k::Int, +) where {V<:AbstractVector} + c = b.chunks[k] + d = c.d + if d.type isa AC.ViewType + c.variadic > 0 && + return _copyvariadic!(b.buffers[c.buf + c.variadic - 2], d, c.variadic) + return _viewrows!(b.out, d, c, b.bases) + end + buf = b.buffers[c.buf + 1] + return _withoffsets(o -> _offsetrows!(b.out, buf, d, o, c), d) +end +function _buildchunk!(b::_StringBuild{String}, k::Int) + c = b.chunks[k] + d = c.d + d.type isa AC.ViewType && return _viewstrings!(b.out, d, c) + return _withoffsets(o -> _offsetstrings!(b.out, d, o, c), d) +end + +function _copyvariadic!(dst::Vector{UInt8}, d::AC.ArrayData, k::Int) + src = d.buffers[k] + GC.@preserve dst d unsafe_copyto!(pointer(dst), AC.sliceptr(src), length(dst)) + return nothing +end + +# Copy a chunk's bytes into its data buffer and encode its payloads. +function _offsetrows!( + payloads, + buf::Vector{UInt8}, + d::AC.ArrayData, + offsets::AbstractVector{O}, + c, +) where {O} + a = Int64(offsets[c.first]) + z = Int64(offsets[c.last + 1]) + # The chunk reads each of these bytes once (see `_willneed`). + GC.@preserve d offsets begin + z > a && _willneed(AC.sliceptr(AC.rolebuffer(d, AC.DATA)) + a, Int(z - a)) + offsets isa Vector && _willneed( + Ptr{UInt8}(pointer(offsets, c.first)), + (c.last - c.first + 2) * sizeof(O), + ) + end + if c.copybytes && z > a + src = AC.subslice(AC.rolebuffer(d, AC.DATA), a, z - a) + GC.@preserve buf d unsafe_copyto!( + pointer(buf, Int(a - c.lo) + 1), + AC.sliceptr(src), + Int(z - a), + ) + end + validity, bitoffset = _validitybits(d) + GC.@preserve d _encodepayloads!( + payloads, + c.at, + buf, + offsets, + c.lo, + c.first, + c.last, + Int(z - c.lo), + validity, + bitoffset, + c.buf, + ) + return nothing +end + +@inline _load8(p::Ptr{UInt8}, i::Int) = ltoh(unsafe_load(Ptr{UInt64}(p + i))) + +# The payload of the `len` bytes at `p + start`, which may read 12 bytes from +# `start` on: the length and up to four content bytes, then either the next +# eight content bytes (inline, zero-padded) or the buffer index and offset. +@inline function _payload(p::Ptr{UInt8}, start::Int, len::UInt64, viewword::UInt64) + w0 = _load8(p, start) + w1 = _load8(p, start + 4) + n4 = min(len, UInt64(4)) + a = (len % UInt32) | ((w0 & ((UInt64(1) << (8 * n4)) - UInt64(1))) << 32) + n8 = min(max(len, UInt64(4)) - UInt64(4), UInt64(8)) + inline = n8 == 8 ? w1 : w1 & ((UInt64(1) << (8 * n8)) - UInt64(1)) + view = viewword | (UInt64(start % UInt32) << 32) + return _StringPayload(a, len <= DataStrings.INLINE_MAX ? inline : view) +end + +""" +Encode rows `first:last` of an offsets-based string batch as payloads over +`buf`, the copy of the batch's bytes from offset `lo` on, reading only its +first `limit` bytes. Rows whose 12-byte window ends within `limit` encode +branch-free; the rows after them (inline, since their bytes end by `limit`) +byte by byte. +""" +function _encodepayloads!( + payloads::Vector{_StringPayload}, + at::Int, + buf::Vector{UInt8}, + offsets::AbstractVector{O}, + lo::Int64, + first::Int, + last::Int, + limit::Int, + validity::Vector{UInt8}, + bitoffset::Int, + bufidx::Int=0, +) where {O} + viewword = UInt64(bufidx % UInt32) + # Offsets ascend, so the rows with a whole window form a prefix. + fast = last + while fast >= first && Int64(offsets[fast]) - lo > limit - 12 + fast -= 1 + end + GC.@preserve buf begin + p = pointer(buf) + if isempty(validity) + @inbounds for k = first:fast + o = Int64(offsets[k]) + payloads[at + k] = + _payload(p, Int(o - lo), (Int64(offsets[k + 1]) - o) % UInt64, viewword) + end + else + @inbounds for k = first:fast + o = Int64(offsets[k]) + x = _payload(p, Int(o - lo), (Int64(offsets[k + 1]) - o) % UInt64, viewword) + payloads[at + k] = + _bitset(validity, bitoffset, k) ? x : DataStrings.PAYLOAD_MISSING + end + end + @inbounds for k = (fast + 1):last + start = Int(Int64(offsets[k]) - lo) + len = Int(Int64(offsets[k + 1]) - Int64(offsets[k])) + x = DataStrings.inline_payload(buf, start + 1, len) + payloads[at + k] = + _isnullslot(validity, bitoffset, k) ? DataStrings.PAYLOAD_MISSING : x + end + end + return payloads +end + +# Zero the bytes after an inline payload's content, as DataStrings requires. +@inline function _canonicalinline(p::_StringPayload) + len = Int(DataStrings.payloadlength(p)) + len > DataStrings.INLINE_MAX && return p + a = len >= 4 ? p.a : p.a & ((UInt64(1) << (32 + 8 * len)) - 1) + nb = max(len - 4, 0) + b = nb >= 8 ? p.b : p.b & ((UInt64(1) << (8 * nb)) - 1) + return _StringPayload(a, b) +end + +# The rows of a `utf8_view`/`binary_view` chunk: the view entries copied as +# payloads (the layouts are identical), re-pointed at the batch's buffers +# within the column's list, and at where each buffer's bytes start (`bases`). +function _viewrows!( + payloads::Vector{_StringPayload}, + d::AC.ArrayData, + c::_StringChunk, + bases::Vector{Int64}, +) + views = _wrapbuffer(_StringPayload, AC.rolebuffer(d, AC.VIEWS), d.offset, Int64(d.len)) + validity, bitoffset = _validitybits(d) + GC.@preserve d for k = c.first:c.last + if _isnullslot(validity, bitoffset, k) + payloads[c.at + k] = DataStrings.PAYLOAD_MISSING + continue + end + p = _canonicalinline(views[k]) + if DataStrings.payloadlength(p) > DataStrings.INLINE_MAX + idx = DataStrings.payloadbufidx(p) + c.buf + off = Int64(DataStrings.payloadoffset(p)) + bases[idx + 1] + p = _StringPayload(p.a, UInt64(idx % UInt32) | (UInt64(off % UInt32) << 32)) + end + payloads[c.at + k] = p + end + return nothing +end + +const _StringColumn = Union{DataStrings.StringVector,DataStrings.BytesVector} + +""" +The string column `V` of `payloads`, owning its bytes. A payload is inline, +missing, or the length and first four bytes of a value whose bytes +`bytes(i)` points at; those values' bytes are copied, in column order, into +one data buffer (several, each under 2 GiB, for more bytes) that the +payloads are re-pointed into. +""" +function _ownstrings(::Type{V}, payloads::Vector{_StringPayload}, bytes) where {V} + sizes = Int64[] + total = Int64(0) + for p in payloads + len = Int64(DataStrings.payloadlength(p)) + len > DataStrings.INLINE_MAX || continue + if total + len > _MAX_VIEW_BUFFER + push!(sizes, total) + total = Int64(0) + end + total += len + end + push!(sizes, total) + data = Vector{UInt8}[Vector{UInt8}(undef, Int(size)) for size in sizes] + k = 1 + pos = Int64(0) + @inbounds for i in eachindex(payloads) + p = payloads[i] + len = Int64(DataStrings.payloadlength(p)) + len > DataStrings.INLINE_MAX || continue + if pos + len > sizes[k] + k += 1 + pos = Int64(0) + end + dst = data[k] + GC.@preserve dst unsafe_copyto!(pointer(dst, Int(pos) + 1), bytes(i), Int(len)) + payloads[i] = + _StringPayload(p.a, UInt64((k - 1) % UInt32) | (UInt64(pos % UInt32) << 32)) + pos += len + end + return V(payloads, data, Val(:trusted)) +end + +# Where a view payload's bytes start in string column `v`. +_viewptr(v::_StringColumn, p::_StringPayload) = + pointer(v.buffers[DataStrings.payloadbufidx(p) + 1], DataStrings.payloadpos(p)) + +""" +A copy of a string or binary column that owns its bytes: the same values over +data buffers holding just the bytes they reference. A column read over a +mapped file, a byte vector, or a fetched range references those bytes in +place, so this is what detaches a copy from them. +""" +_compactstrings(v::_StringColumn) = _rowsof(v, eachindex(v)) + +# Rows `rows` of a column. A string column's rows own just their bytes, so a +# filtered read does not keep the rest of the batch's bytes alive. +_rowsof(v::AbstractVector, rows) = v[rows] +function _rowsof(v::V, rows) where {V<:_StringColumn} + payloads = v.payloads + return GC.@preserve v _ownstrings( + V, + payloads[rows], + i -> _viewptr(v, payloads[rows[i]]), + ) +end +# A dictionary column's rows keep its pool. +function _rowsof(v::DictEncoded{T,R,P}, rows) where {T,R,P} + indices = v.indices[rows] + return DictEncoded{T,R,P,typeof(indices)}(indices, v.pool, v.ordered) +end + +# The payload of the `len` bytes at `p`: inline when they fit, else their +# length and first four bytes (`_ownstrings` adds where they are stored). +@inline function _rowpayload(p::Ptr{UInt8}, len::Int) + a = UInt64(len % UInt32) + len > DataStrings.INLINE_MAX && + return _StringPayload(a | (UInt64(ltoh(unsafe_load(Ptr{UInt32}(p)))) << 32), 0) + b = zero(UInt64) + for i = 1:min(len, 4) + a |= UInt64(unsafe_load(p, i)) << (32 + 8 * (i - 1)) + end + for i = 5:len + b |= UInt64(unsafe_load(p, i)) << (8 * (i - 5)) + end + return _StringPayload(a, b) +end + +""" +Rows `rows` of string or binary batch `d` as a column of element type `T` +that owns its bytes: a `StringVector` or `BytesVector` whose data buffer +holds just those values' bytes, or `String`s. +""" +function _rowsof(::Type{T}, d::AC.ArrayData, rows) where {T} + validity, bitoffset = _validitybits(d) + isnull(i) = _isnullslot(validity, bitoffset, i) + if d.type isa AC.ViewType + views = + _wrapbuffer(_StringPayload, AC.rolebuffer(d, AC.VIEWS), d.offset, Int64(d.len)) + viewvalue(i) = + let p = views[i], len = Int(DataStrings.payloadlength(p)) + len <= DataStrings.INLINE_MAX && + return Ptr{UInt8}(pointer(views, i)) + 4, len + data = d.buffers[DataStrings.payloadbufidx(p) + 3] + return AC.sliceptr(AC.subslice(data, DataStrings.payloadoffset(p), len)), + len + end + return GC.@preserve d views _gatherstrings(T, rows, isnull, viewvalue) + end + return _withoffsets(d) do offsets + base = AC.sliceptr(AC.rolebuffer(d, AC.DATA)) + function offsetvalue(i) + o = Int64(offsets[i]) + len = Int64(offsets[i + 1]) - o + len <= _MAX_VIEW_BUFFER || throw( + ArgumentError( + "a string longer than 2 GiB does not fit a DataString; " * + "read the column with stringtype=String", + ), + ) + return base + o, Int(len) + end + return GC.@preserve d _gatherstrings(T, rows, isnull, offsetvalue) + end +end + +# The values `value(i)` (a pointer and a length) of rows `rows`, or `missing` +# where `isnull(i)`, as an owned column of element type `T`. +function _gatherstrings(::Type{T}, rows, isnull, value) where {T} + B = Base.nonmissingtype(T) + if B === String + out = Vector{T}(undef, length(rows)) + for (k, i) in enumerate(rows) + out[k] = isnull(i) ? missing : unsafe_string(value(i)...) + end + return out + end + V = + B === DataStrings.DataString ? DataStrings.StringVector{T} : + DataStrings.BytesVector{T} + payloads = Vector{_StringPayload}(undef, length(rows)) + for (k, i) in enumerate(rows) + payloads[k] = isnull(i) ? DataStrings.PAYLOAD_MISSING : _rowpayload(value(i)...) + end + return _ownstrings(V, payloads, k -> first(value(rows[k]))) +end + +# The string columns `parts`, one per batch, joined into one of element type +# `T` over all their buffers: a part's view payloads re-point into the joined +# buffer list. +function _joinstrings(::Type{T}, parts) where {T} + V = + Base.nonmissingtype(T) === DataStrings.DataString ? DataStrings.StringVector{T} : + DataStrings.BytesVector{T} + buffers = Vector{UInt8}[] + out = Vector{_StringPayload}(undef, sum(length, parts; init=0)) + at = 0 + for part in parts + base = length(buffers) + append!(buffers, part.buffers) + _repoint!(out, at, part.payloads, base) + at += length(part) + end + return V(out, buffers, Val(:trusted)) +end + +# Copy `payloads` into `out` after index `at`, view payloads re-pointed `base` +# buffers further into the joined buffer list. +function _repoint!(out, at::Int, payloads::Vector{_StringPayload}, base::Int) + @inbounds for i in eachindex(payloads) + p = payloads[i] + if base > 0 && DataStrings.payloadlength(p) > DataStrings.INLINE_MAX + p = _StringPayload( + p.a, + UInt64((DataStrings.payloadbufidx(p) + base) % UInt32) | + (p.b & 0xffffffff00000000), + ) + end + out[at + i] = p + end + return out +end + +# A dictionary pool value owning its bytes: a `DataString` or `DataBytes` view +# into the pool column's buffer, which may be the source's bytes, gets its own. +_ownedpoolvalue(x) = _ownedvalue(x) +function _ownedpoolvalue(x::Union{DataStrings.DataString,DataStrings.DataBytes}) + len = Int(DataStrings.payloadlength(x.p)) + len <= DataStrings.INLINE_MAX && return x + data = Vector{UInt8}(undef, len) + GC.@preserve x data unsafe_copyto!( + pointer(data), + pointer(x.data, DataStrings.payloadpos(x.p)), + len, + ) + return typeof(x)(_StringPayload(x.p.a, UInt64(0)), data) +end + +# `String`s of an offsets-based chunk, read from the batch's data. +function _offsetstrings!( + out, + d::AC.ArrayData, + offsets::AbstractVector{O}, + c::_StringChunk, +) where {O} + a = Int64(offsets[c.first]) + z = Int64(offsets[c.last + 1]) + src = AC.subslice(AC.rolebuffer(d, AC.DATA), a, z - a) + validity, bitoffset = _validitybits(d) + GC.@preserve d begin + p = AC.sliceptr(src) - a + for k = c.first:c.last + o = Int64(offsets[k]) + out[c.at + k] = + _isnullslot(validity, bitoffset, k) ? missing : + unsafe_string(p + o, Int(Int64(offsets[k + 1]) - o)) + end + end + return nothing +end + +# `String`s of a view chunk: inline bytes from the payload, others from the +# batch's variadic buffers. +function _viewstrings!(out, d::AC.ArrayData, c::_StringChunk) + views = _wrapbuffer(_StringPayload, AC.rolebuffer(d, AC.VIEWS), d.offset, Int64(d.len)) + validity, bitoffset = _validitybits(d) + scratch = Vector{UInt8}(undef, 16) + GC.@preserve d scratch for k = c.first:c.last + if _isnullslot(validity, bitoffset, k) + out[c.at + k] = missing + continue + end + p = views[k] + len = Int(DataStrings.payloadlength(p)) + out[c.at + k] = if len <= DataStrings.INLINE_MAX + q = pointer(scratch) + unsafe_store!(Ptr{UInt64}(q), htol((p.a >> 32) | (p.b << 32))) + unsafe_store!(Ptr{UInt64}(q + 8), htol(p.b >> 32)) + unsafe_string(q, len) + else + data = d.buffers[DataStrings.payloadbufidx(p) + 3] + unsafe_string( + AC.sliceptr(AC.subslice(data, DataStrings.payloadoffset(p), len)), + len, + ) + end + end + return nothing +end + +# --- multi-batch columns ----------------------------------------------------------- + +# Broadcasting over a column of several batches runs batch by batch, each +# over its own view (as fast as over a `Vector`), instead of locating the +# batch of every element. +struct _ChainedStyle <: Broadcast.AbstractArrayStyle{1} end +_ChainedStyle(::Val{1}) = _ChainedStyle() +_ChainedStyle(::Val{N}) where {N} = Broadcast.DefaultArrayStyle{N}() +Base.BroadcastStyle(::Type{<:ChainedVector{<:Any,<:ArrowVector}}) = _ChainedStyle() +Base.similar(::Broadcast.Broadcasted{_ChainedStyle}, ::Type{T}, dims) where {T} = + similar(Array{T}, dims) +Base.similar(::Broadcast.Broadcasted{_ChainedStyle}, ::Type{Bool}, dims) = + similar(BitArray, dims) + +function Base.copyto!(dest::AbstractVector, bc::Broadcast.Broadcasted{_ChainedStyle}) + axes(dest) == axes(bc) || Broadcast.throwdm(axes(dest), axes(bc)) + c = _drivingchain(bc) + c === nothing && return copyto!(dest, convert(Broadcast.Broadcasted{Nothing}, bc)) + for r in _batchranges(c) + part = Broadcast.instantiate(_batchargs(bc, r)) + if dest isa BitVector + # Bits pack fastest into a BitVector of their own. + bits = copyto!(BitVector(undef, length(r)), part) + copyto!(dest, first(r), bits, 1, length(r)) + else + copyto!(view(dest, r), part) + end + end + return dest +end + +# Without a concrete element type, Base narrows the result to the values it +# computes (`x .+ 1` over a nullable column with no nulls gives a +# `Vector{Int}`): each batch narrows on its own, and the batches join under +# the same widening (`promote_typejoin`). +function Base.copy(bc::Broadcast.Broadcasted{_ChainedStyle}) + c = _drivingchain(bc) + c === nothing && return copy(_elementwise(bc)) + E = Broadcast.combine_eltypes(bc.f, bc.args) + isconcretetype(E) && return copyto!(similar(bc, E), bc) + ranges = _batchranges(c) + parts = [copy(Broadcast.instantiate(_batchargs(bc, r))) for r in ranges] + isempty(parts) && return copy(_elementwise(bc)) + T = mapreduce(eltype, Base.promote_typejoin, parts) + dest = similar(bc, T) + for (r, part) in zip(ranges, parts) + copyto!(dest, first(r), part, 1, length(part)) + end + return dest +end + +_elementwise(bc::Broadcast.Broadcasted) = + Broadcast.Broadcasted{Broadcast.DefaultArrayStyle{1}}(bc.f, bc.args, bc.axes) + +# The column the broadcast is cut by (its first multi-batch view argument), or +# `nothing` when some argument cannot be cut into rows (a tuple, an array of +# other dimensions). +_drivingchain(bc::Broadcast.Broadcasted) = + all(_cuttable, bc.args) ? _firstchain(bc) : nothing + +# The row ranges of `c`'s batches, empty batches left out. +function _batchranges(c::ChainedVector) + ranges = UnitRange{Int}[] + lo = 1 + for a in c.arrays + n = length(a) + n > 0 && push!(ranges, lo:(lo + n - 1)) + lo += n + end + return ranges +end + +# Arguments a broadcast can be cut into rows by: vectors, scalars, and +# broadcasts of them (not tuples or other arrays). +_cuttable(a::Broadcast.Broadcasted) = all(_cuttable, a.args) +_cuttable(a::AbstractVector) = axes(a, 1) isa Base.OneTo +_cuttable(a) = axes(a) == () + +_firstchain(a::ChainedVector{<:Any,<:ArrowVector}) = a +function _firstchain(a::Broadcast.Broadcasted) + for x in a.args + c = _firstchain(x) + c === nothing || return c + end + return nothing +end +_firstchain(a) = nothing + +# `bc` restricted to rows `r`: a column whose batch is exactly `r` gives that +# batch, other vectors a view, and scalars and length-1 vectors stay whole. +_batchargs(bc::Broadcast.Broadcasted, r) = + Broadcast.Broadcasted(bc.f, map(a -> _batcharg(a, r), bc.args)) +_batcharg(a::Broadcast.Broadcasted, r) = _batchargs(a, r) +function _batcharg(a::ChainedVector{<:Any,<:ArrowVector}, r) + k = searchsortedfirst(a.inds, last(r)) + k <= length(a.inds) && + a.inds[k] == last(r) && + (k == 1 ? 0 : a.inds[k - 1]) == first(r) - 1 && + return a.arrays[k] + return view(a, r) +end +_batcharg(a::AbstractVector, r) = length(a) == 1 ? a : view(a, r) +_batcharg(a, r) = a + +# --- column assembly -------------------------------------------------------------- + +# Shared decimals read as `DataDecimals.Decimal{P,S,T}`, which has the storage +# layout of every Arrow decimal width. +_isbulkdecimal(f::AC.Field) = f.type isa AC.DecimalType && _shareddecimal(f.type) +"One batch of a decimal column, read as stored." +_decimalcolumn(f::AC.Field, d::AC.ArrayData, budget=nothing) = + _viewcolumn(f, AC.ArrayData[d], _ViewContext(DataStrings.DataString, budget)) + +# One read's view settings, its allocation budget, and its cache of +# dictionary pools (built once per pool snapshot; columns building on several +# tasks share it under the lock). +struct _ViewContext + stringtype::Type + budget::Union{Nothing,AllocationBudget} + pools::IdDict{AC.ArrayData,AbstractVector} + poolslock::ReentrantLock +end +_ViewContext(stringtype::Type=DataStrings.DataString, budget=nothing) = + _ViewContext(stringtype, budget, IdDict{AC.ArrayData,AbstractVector}(), ReentrantLock()) +# The same settings and pool cache under another allowance (`Arrow.Stream` batches). +_ViewContext(ctx::_ViewContext, budget) = + _ViewContext(ctx.stringtype, budget, ctx.pools, ctx.poolslock) + +# The view column of a dictionary's values, built once per pool snapshot. The +# build runs outside the lock; a concurrent duplicate is dropped. +function _dictpool(ctx::_ViewContext, vf::AC.Field, dict::AC.ArrayData) + pool = @lock ctx.poolslock get(ctx.pools, dict, nothing) + pool === nothing || return pool + built = _viewcolumn(vf, AC.ArrayData[dict], ctx) + return @lock ctx.poolslock get!(ctx.pools, dict, built) +end + +function _markpools!(live, data::AC.ArrayData, cached) + dict = data.dictionary + if dict !== nothing + haskey(cached, dict) && (live[dict] = nothing) + _markpools!(live, dict, cached) + end + foreach(child -> _markpools!(live, child, cached), data.children) + return nothing +end + +# A stream reuses the pool columns of the dictionaries still in effect +# (`dictionaries`), but only returned tables may keep older ones alive. +function _prunepools!(ctx::_ViewContext, dictionaries, budget) + lock(ctx.poolslock) do + isempty(ctx.pools) && return nothing + _chargedict!(budget, AC.ArrayData, Nothing, length(ctx.pools), "live stream pools") + live = IdDict{AC.ArrayData,Nothing}() + for dictionary in dictionaries + haskey(ctx.pools, dictionary) && (live[dictionary] = nothing) + _markpools!(live, dictionary, ctx.pools) + end + filter!(entry -> haskey(live, first(entry)), ctx.pools) + return nothing + end + return nothing +end + +# A view object's own size (struct, wrapped vector headers), charged per batch. +const _VIEW_OBJECT_BYTES = Int64(256) + +function _checkstringtype(stringtype) + stringtype === DataStrings.DataString || + stringtype === String || + throw(ArgumentError("stringtype must be DataString or String, got $stringtype")) + return stringtype +end + +""" +Whether `f` reads through this file's columns: a flat layout (as a view) or +a list, fixed-size list, map, struct, or union of such layouts (as typed +rows). An ArrowTypes extension label keeps the materializing path unless it +maps to no Julia type (see `_storagelabel`) or is a fixed-shape tensor. List +views, structs with duplicate or more than 64 field names, and layouts +nested more than 8 deep also keep the materializing path. +""" +function _viewable(f::AC.Field) + if _hasarrowtypesextension(f) + _tensorshape(f) === nothing || return true + _storagelabel(f) || return false + end + return _viewablelayout(f) +end + +# Whether `f` carries an extension label that reads as its storage: a quiet +# canonical label (`arrow.json`, `arrow.opaque`), or a label the facade does +# not own and no ArrowTypes.jl `JuliaType` method names, which warns once +# per process. Other labels, and labels below the top level, keep the +# materializing path, which resolves ArrowTypes.jl hooks once per field. +function _storagelabel(f::AC.Field) + ext = _arrowtypesextension(f) + ext === nothing && return false + any(_hasarrowtypesextension, f.children) && return false + name = first(ext) + _quietextension(name) && return true + _facadelabel(name) && return false + # A name that is not a Symbol yet cannot be registered; looking it up + # does not intern it. + sym = _existingjlsymbol(name) + sym === nothing || !_juliatypenames(sym) || return false + _warnunsupportedextension((warn=true, budget=nothing), name, f) + return true +end + +# Whether an `ArrowTypes.JuliaType` method other than its catch-all +# fallbacks can take `Val(sym)` as its first argument. +_juliatypenames(sym::Symbol) = any(methods(ArrowTypes.JuliaType)) do m + sig = Base.unwrap_unionall(m.sig) + length(sig.parameters) >= 2 || return false + arg = Base.rewrap_unionall(sig.parameters[2], m.sig) + return arg !== Any && Val{sym} <: arg +end + +_facadelabel(name::String) = + name in (_UUID_EXTENSION, "JuliaLang.Int128", "JuliaLang.UInt128", "arrow.bool8") || + any(p -> last(p) == name, _FACADE_PERIOD_NAMES) +function _viewablelayout(f::AC.Field, depth::Int=0) + t = f.type + t isa Union{ + AC.IntType, + AC.FloatType, + AC.BoolType, + AC.DateType, + AC.TimeType, + AC.TimestampType, + AC.DurationType, + AC.IntervalType, + AC.FixedSizeBinaryType, + AC.Utf8Type, + AC.BinaryType, + AC.ViewType, + AC.NullType, + } && return true + t isa AC.DecimalType && return true + t isa AC.DictionaryType && return _viewablelayout(AC.dictvaluefield(f, t), depth) + t isa AC.RunEndEncodedType && + return length(f.children) == 2 && _viewablelayout(f.children[2], depth) + # Row types nest one type per level, and Julia compiles deep ones slowly. + t isa _NestedLayout && depth == _MAX_ROW_DEPTH && return false + viewable(c) = _viewablelayout(c, depth + 1) + t isa Union{AC.ListType,AC.FixedSizeListType} && + return length(f.children) == 1 && viewable(f.children[1]) + # A list view gathers its rows' values, which it does for flat values. + t isa AC.ListViewType && return length(f.children) == 1 && + !_nestedvalue(f.children[1]) && + viewable(f.children[1]) + if t isa AC.MapType + length(f.children) == 1 || return false + entries = f.children[1] + return entries.type isa AC.StructType && + length(entries.children) == 2 && + all(viewable, entries.children) + end + t isa AC.StructType && return length(f.children) <= _MAX_ROW_FIELDS && + allunique(String[c.name for c in f.children]) && + all(viewable, f.children) + t isa AC.UnionType && return !isempty(f.children) && all(viewable, f.children) + return false +end + +# Fixed-size binary values wider than this read as `DataBytes`: code over +# an `NTuple{N,UInt8}` compiles in time that grows faster than `N` (minutes +# at 16 KiB), and a hostile schema can declare any width. +const _MAX_TUPLE_BYTES = 256 +_widebinary(t::AC.ArrowType) = t isa AC.FixedSizeBinaryType && t.nbytes > _MAX_TUPLE_BYTES + +# The stored element type and storage conversion of a fixed-width layout. +function _fixedstorage(t::AC.ArrowType) + t isa AC.IntType && return AC.juliatype(t), Nothing + t isa AC.FloatType && return AC.juliatype(t), Nothing + if t isa AC.DecimalType + _shareddecimal(t) && return _decimalhost(t), Nothing + # Other decimals read as stored: integers, or bytes at 128 and 256 bits. + t.bits <= 64 && return (t.bits == 32 ? Int32 : Int64), Nothing + return NTuple{t.bits ÷ 8,UInt8}, _StoredBytes + end + t isa AC.TimestampType && return _facadebasetype(t), Nothing + t isa AC.DurationType && return _facadebasetype(t), Nothing + t isa AC.FixedSizeBinaryType && return NTuple{t.nbytes,UInt8}, Nothing + t isa AC.DateType && + return t.unit == AC.DAY ? (Int32, _DaysToDate) : (Int64, _MillisToDateTime) + if t isa AC.TimeType + S = t.bits == 32 ? Int32 : Int64 + return S, _TimeOfDay{_NANOS_PER_DAY ÷ _timeticksperday(t.unit)} + end + if t isa AC.IntervalType + t.unit == AC.YEAR_MONTH && return Int32, _YearMonth + t.unit == AC.DAY_TIME && return NTuple{2,Int32}, _DayTime + return Durations.Duration, Nothing + end + throw(ArgumentError("no fixed-width storage for $(AC.descriptorname(t))")) +end + +_timeticksperday(unit) = + unit == AC.SECOND ? Int64(86_400) : + unit == AC.MILLISECOND ? Int64(86_400_000) : + unit == AC.MICROSECOND ? Int64(86_400_000_000) : _NANOS_PER_DAY + +""" +The public element type of a column, before nullability. Nested rows take +their children's element types over `datas`, the batches being read, so a +child whose data holds undeclared nulls widens the row type of every batch. +""" +function _viewbasetype(f::AC.Field, ctx::_ViewContext, datas=AC.ArrayData[]) + t = f.type + t isa AC.NullType && return Missing + t isa AC.BoolType && return Bool + t isa Union{AC.Utf8Type,AC.ViewType} && + (!(t isa AC.ViewType) || t.utf8) && + return ctx.stringtype === String ? String : DataStrings.DataString + (t isa Union{AC.BinaryType,AC.ViewType} || _widebinary(t)) && + return DataStrings.DataBytes + t isa AC.DictionaryType && return _viewbasetype(AC.dictvaluefield(f, t), ctx) + t isa AC.RunEndEncodedType && return _viewbasetype(f.children[2], ctx) + shape = _tensorshape(f) + shape === nothing || return _tensortype(f, ctx, datas, shape) + t isa _NestedLayout && return _rowtype(f, ctx, datas) + S, C = _fixedstorage(t) + return C === Nothing ? S : _storagetarget(C) +end + +# Physical nulls the column's values can hold, through dictionary pools and +# run-end values. +function _viewnulls(f::AC.Field, d::AC.ArrayData) + t = f.type + t isa AC.NullType && return d.len > 0 + if t isa AC.RunEndEncodedType + return _viewnulls(f.children[2], d.children[2]) + end + if t isa AC.DictionaryType + AC.nullcount(d) > 0 && return true + dict = d.dictionary::AC.ArrayData + return _viewnulls(AC.dictvaluefield(f, t), dict) + end + return AC.nullcount(d) > 0 +end + +# The declared type decides `Missing` (`_declaredeltype` is the authority); +# data that holds nulls a non-nullable declaration did not admit widens the +# column, as on every other read path. +function _vieweltype(f::AC.Field, datas, ctx::_ViewContext) + B = _viewbasetype(f, ctx, datas) + B === Missing && return Missing + # A union has no validity of its own: its children's types carry nulls. + f.type isa AC.UnionType && return B + nullable = Missing <: _declaredeltype(f, true) || any(d -> _viewnulls(f, d), datas) + return nullable ? Union{Missing,B} : B +end + +function _checktimes(f::AC.Field, t::AC.TimeType, d::AC.ArrayData, budget) + perday = _timeticksperday(t.unit) + S = t.bits == 32 ? Int32 : Int64 + col = _primitive(Union{Missing,S}, S, Nothing, d, budget) + for x in col + (x === missing || 0 <= x < perday) && continue + P = nameof(_unitperiod(t.unit)) + throw( + ValidationError( + "column $(repr(f.name)): time-of-day value $x is outside [0, $perday) " * + "in its unit; read the column as durations since midnight with " * + "`scan=Tables.Scan(select=($(_uninferred(repr, Symbol(f.name))::String) " * + "=> Dates.$P,))`", + ), + ) + end + return nothing +end + +# The period of one tick of a time-of-day unit. +_unitperiod(unit) = + unit == AC.SECOND ? Dates.Second : + unit == AC.MILLISECOND ? Dates.Millisecond : + unit == AC.MICROSECOND ? Dates.Microsecond : Dates.Nanosecond + +# Time-of-day storage `raw` read as `P` periods since midnight: the scan +# override that reads times a `Dates.Time` cannot hold. +function _timeperiods(t::AC.TimeType, raw, ::Type{P}) where {P} + unit = _unitperiod(t.unit) + out = Vector{Missing<:eltype(raw) ? Union{Missing,P} : P}(undef, length(raw)) + for i in eachindex(out, raw) + x = raw[i] + out[i] = x === missing ? missing : convert(P, unit(x)) + end + return out +end + +# One batch of a view column. Nested layouts skip `_batchview`, whose body +# for every flat layout would otherwise compile once per row type. +_viewbatch(T::Type, f::AC.Field, d::AC.ArrayData, ctx::_ViewContext) = + f.type isa _NestedLayout ? _nestedbatch(T, f, d, ctx) : _batchview(T, f, d, ctx) + +""" +One batch of a fixed-width, boolean, dictionary, or run-end view column. +`T` is only passed on, so this body compiles once, not once per element type +(every time zone and decimal scale is a type of its own). +""" +function _batchview(T::Type, f::AC.Field, d::AC.ArrayData, ctx::_ViewContext) + t = f.type + budget = ctx.budget + _chargeobject!(budget, _VIEW_OBJECT_BYTES, "column view") + t isa AC.BoolType && return _boolvector(T, d) + t isa AC.DictionaryType && return _dictview(T, f, t, d, ctx) + t isa AC.RunEndEncodedType && return _runendview(T, f, d, ctx) + t isa AC.TimeType && _checktimes(f, t, d, budget) + S, C = _fixedstorage(t) + return _primitive(T, S, C, d, budget) +end + +function _dictview( + ::Type{T}, + f::AC.Field, + t::AC.DictionaryType, + d::AC.ArrayData, + ctx::_ViewContext, +) where {T} + pool = _dictpool(ctx, AC.dictvaluefield(f, t), d.dictionary::AC.ArrayData) + R = AC.juliatype(t.indextype) + indices = _primitive(Missing <: T ? Union{Missing,R} : R, R, Nothing, d, ctx.budget) + return DictEncoded{T,R,typeof(pool),typeof(indices)}(indices, pool, t.ordered) +end + +function _runendview(::Type{T}, f::AC.Field, d::AC.ArrayData, ctx::_ViewContext) where {T} + runs, values = d.children[1], d.children[2] + R = AC.juliatype(runs.type::AC.IntType) + _chargeexact!(ctx.budget, Int64, runs.len, "run ends") + runends = _collectas(Int64, _primitive(R, R, Nothing, runs, ctx.budget)) + vcol = _viewcolumn(f.children[2], [values], ctx, T; source=false) + return RunEndEncoded{T,typeof(vcol)}(runends, vcol, Int(d.offset), Int(d.len)) +end + +""" +The view column of field `f` over its decoded batches `datas`: one view for +one batch, a `ChainedVector` of per-batch views for several, and one +`StringVector`/`BytesVector` (or `Vector{String}`) across all batches of a +string or binary column. `T` overrides the element type (run-end values). +A string column references the source's bytes in place unless `source` is +false: the values of nested rows and run-end columns copy them, so a copy of +a row (`copy(table)`, `DataFrame(table)`) owns its strings without copying +each one. +""" +function _viewcolumn( + f::AC.Field, + datas::AbstractVector{AC.ArrayData}, + ctx::_ViewContext, + T::Type=_vieweltype(f, datas, ctx); + source::Bool=true, +) + t = f.type + t isa AC.NullType && return fill(missing, sum(d -> Int(d.len), datas; init=0)) + if _isstringlayout(t) + b = _allocate!(_stringbuild(T, datas, ctx.budget; source)) + for k in eachindex(b.chunks) + _buildchunk!(b, k) + end + return _finishbuild(b) + end + isempty(datas) && return T[] + first = _viewbatch(T, f, datas[1], ctx) + length(datas) == 1 && return first + # `first`'s type is known at run time; see `_uninferred`. + return _uninferred(_chainviews, first, f, datas, ctx, T) +end + +# A concretely typed chunk vector (a comprehension over `_batchview` would +# infer its element type at run time, per column). Batches whose views differ +# in type (dictionary pools that differ in nullability) chain abstractly. +function _chainviews(first::V, f::AC.Field, datas, ctx::_ViewContext, T::Type) where {V} + _chargevector!(ctx.budget, V, length(datas), "column views") + _chargevector!(ctx.budget, Int, length(datas), "chunked column index") + cols = Vector{V}(undef, length(datas)) + cols[1] = first + for k = 2:length(datas) + col = _viewbatch(T, f, datas[k], ctx) + # The abstract chain compiles apart (see `_uninferred`). + col isa V || return _uninferred(_mixedchainviews, first, f, datas, ctx, T) + cols[k] = col + end + return ChainedVector(cols) +end + +_mixedchainviews(first, f::AC.Field, datas, ctx::_ViewContext, T::Type) = ChainedVector( + AbstractVector{T}[ + k == 1 ? first : _viewbatch(T, f, d, ctx) for (k, d) in enumerate(datas) + ], +) + +# --- nested rows ---------------------------------------------------------------- +# +# Lists read as `ListValue` views, maps as `MapValue` views, fixed-size lists of up +# to 64 non-null values as `NTuple`s, structs as `NamedTuple`s built on access, +# and unions as a `Union` of their children's element types. A column's +# element type alone fixes how its child values are stored (`_columntype`): +# strings in a `StringVector`, binary in a `BytesVector`, lists in a +# `ListVector`, maps in a `MapVector`, structs in a `StructVector`, other +# fixed-width values in a `Primitive{E,S,Nothing}` (`S` the non-missing `E`), +# other values in a `Vector`. Columns keep their children untyped and read +# them back through that type, so a column type has only its element type as +# a parameter, however deep the schema, and access is type-stable. A +# fixed-width child whose stored layout is its values views the source, as a +# top-level column does, and keeps it alive; other children are owned copies. + +const _NestedLayout = Union{ + AC.ListType, + AC.ListViewType, + AC.FixedSizeListType, + AC.MapType, + AC.StructType, + AC.UnionType, +} +const _MAX_ROW_FIELDS = 64 +# Each level nests a `Union{Missing,...}` in the row type, and the first read's +# compile time grows steeply with it: on Julia 1.13, lists 10 deep take about +# 1.5 s, 17 deep about 50 s, and structs 32 deep over 8 minutes. +const _MAX_ROW_DEPTH = 8 +const _MAX_TUPLE_ROW = 64 + +""" + Arrow.ListValue{T} <: AbstractVector{T} + +A row of a list column, or of a fixed-size list column longer than 64 values +or with nullable values: a read-only view of consecutive child values of its +column. It stays valid after its table is gone. `copy(row)` gives a +`Vector{T}`. +""" +struct ListValue{T} <: AbstractVector{T} + values::Any # the column's child values, a `_columntype(T)` + offset::Int # child values before the row's first + len::Int +end + +""" + Arrow.MapValue{K,V} <: AbstractDict{K,V} + +A row of a map column: a read-only dictionary that iterates the map's +`key => value` entries in stored order, duplicate keys included, which `length` +counts. `row[key]`, `get`, and `haskey` find the last entry with that key, the +one `Dict(row)` keeps; `keys(row)` and `values(row)` are vectors in stored +order. Each lookup scans the row, so for many lookups in a large map +`Dict(row)` is faster. `collect(row)` gives the `Vector{Pair{K,V}}` of entries +and `copy(row)` a `Dict{K,V}`. It stays valid after its table is gone. +""" +struct MapValue{K,V} <: AbstractDict{K,V} + keys::Any # the column's keys and values, `_columntype(K)`, `_columntype(V)` + values::Any + offset::Int + len::Int +end + +# List and map offsets: the batch's `Int32` (`Int64` for the large layouts) +# offsets, or ones the read computed. +const _Offsets = Union{Vector{Int32},Vector{Int64}} + +""" + Arrow.ListVector{T} <: Arrow.ArrowVector{T} + +A list column: row `i` is the [`Arrow.ListValue`](@ref) of child values +`offsets[i]+1:offsets[i+1]`, or `missing` where the validity bitmap marks a +null. +""" +struct ListVector{T} <: ArrowVector{T} + offsets::_Offsets # the batch's own offsets where it can (see `_listoffsets`) + values::Any # child values, a `_columntype(eltype(row))` + validity::Vector{UInt8} + bitoffset::Int + owner::Any +end + +""" + Arrow.MapVector{T} <: Arrow.ArrowVector{T} + +A map column: row `i` is the [`Arrow.MapValue`](@ref) of entries +`offsets[i]+1:offsets[i+1]`, or `missing` where the validity bitmap marks a +null. +""" +struct MapVector{T} <: ArrowVector{T} + offsets::_Offsets + keys::Any + values::Any + validity::Vector{UInt8} + bitoffset::Int + owner::Any +end + +""" + Arrow.StructVector{T} <: Arrow.ArrowVector{T} + +A struct column: row `i` is the `NamedTuple` of each field's value `i`, +built when it is read, or `missing` where the validity bitmap marks a null. +""" +struct StructVector{T} <: ArrowVector{T} + children::Vector{AbstractVector} # field `k`'s values, a `_columntype(fieldtype(row, k))` + offset::Int # child values before row 1 + len::Int + validity::Vector{UInt8} + bitoffset::Int + owner::Any +end + +_columntype(::Type{T}) where {T} = _columntype(T, Base.nonmissingtype(T)) +_columntype(::Type{Missing}) = Vector{Missing} +_columntype(::Type{T}, ::Type{S}) where {T,S} = + isbitstype(S) && S !== Bool ? Primitive{T,S,Nothing,_buffertype(S)} : Vector{T} +_columntype(::Type{T}, ::Type{DataStrings.DataString}) where {T} = + DataStrings.StringVector{T} +_columntype(::Type{T}, ::Type{DataStrings.DataBytes}) where {T} = DataStrings.BytesVector{T} +# A `Union` of row types (a union column's values) is stored as a `Vector`. +_columntype(::Type{T}, ::Type{R}) where {T,R<:ListValue} = + R isa DataType ? ListVector{T} : Vector{T} +_columntype(::Type{T}, ::Type{R}) where {T,R<:MapValue} = + R isa DataType ? MapVector{T} : Vector{T} +_columntype(::Type{T}, ::Type{R}) where {T,R<:NamedTuple} = + R isa DataType ? StructVector{T} : Vector{T} + +_childdatas(datas, k::Int) = AC.ArrayData[d.children[k] for d in datas] + +# The row type of nested field `f` over batches `datas`: each child's element +# type, so a child whose data holds undeclared nulls widens every batch's rows. +function _rowtype(f::AC.Field, ctx::_ViewContext, datas) + t = f.type + if t isa AC.StructType + names = Tuple(Symbol(c.name) for c in f.children) + types = + (_vieweltype(c, _childdatas(datas, k), ctx) for (k, c) in enumerate(f.children)) + return NamedTuple{names,Tuple{types...}} + end + t isa AC.UnionType && return Union{ + (_vieweltype(c, _childdatas(datas, k), ctx) for (k, c) in enumerate(f.children))..., + } + if t isa AC.MapType + entries = f.children[1] + entrydatas = _childdatas(datas, 1) + K = _vieweltype(entries.children[1], _childdatas(entrydatas, 1), ctx) + V = _vieweltype(entries.children[2], _childdatas(entrydatas, 2), ctx) + return MapValue{K,V} + end + E = _vieweltype(f.children[1], _childdatas(datas, 1), ctx) + # A tuple of a non-concrete `E` (nullable values) is an abstract type: + # every row would be boxed, with up to 2^N runtime types. + t isa AC.FixedSizeListType && + t.listsize <= _MAX_TUPLE_ROW && + isconcretetype(E) && + return NTuple{t.listsize,E} + return ListValue{E} +end + +# Whether field `f`'s values are nested rows (through a dictionary's values +# or run-end encoding too). +_nestedvalue(f::AC.Field) = + f.type isa AC.DictionaryType ? _nestedvalue(AC.dictvaluefield(f, f.type)) : + f.type isa AC.RunEndEncodedType ? _nestedvalue(f.children[2]) : f.type isa _NestedLayout + +# Element `j` of `values`, a column of `T` values stored as its `_columntype`. +Base.@propagate_inbounds _childvalue(::Type{T}, values, j::Int) where {T} = + (values::_columntype(T))[j] + +Base.IndexStyle(::Type{<:ListValue}) = IndexLinear() +Base.size(r::ListValue) = (r.len,) +Base.size(v::Union{ListVector,MapVector}) = (length(v.offsets) - 1,) +Base.size(v::StructVector) = (v.len,) + +Base.@propagate_inbounds function Base.getindex(r::ListValue{T}, i::Int) where {T} + @boundscheck checkbounds(r, i) + return @inbounds _childvalue(T, r.values, r.offset + i) +end + +# `push!`, `append!`, `merge!`, and the like go through `resize!`, +# `setindex!`, and `delete!`. +@noinline _readonlyvalue(x) = throw( + ArgumentError( + "an Arrow.$(nameof(typeof(x))) is read-only; copy(x) gives a " * + "$(x isa MapValue ? "Dict" : "Vector") that can change", + ), +) +Base.setindex!(x::ListValue, v, ::Int) = _readonlyvalue(x) +Base.resize!(x::ListValue, ::Integer) = _readonlyvalue(x) +Base.setindex!(x::MapValue, v, k) = _readonlyvalue(x) +Base.delete!(x::MapValue, k) = _readonlyvalue(x) +Base.empty!(x::MapValue) = _readonlyvalue(x) + +Base.length(r::MapValue) = r.len +function Base.iterate(r::MapValue{K,V}, i::Int=1) where {K,V} + i > r.len && return nothing + j = r.offset + i + return Pair{K,V}(_childvalue(K, r.keys, j), _childvalue(V, r.values, j)), i + 1 +end + +# The child position of the last entry whose key is `key`, or `nothing`. +function _mapentry(r::MapValue{K}, key) where {K} + for j = (r.offset + r.len):-1:(r.offset + 1) + isequal(_childvalue(K, r.keys, j), key) && return j + end + return nothing +end +Base.haskey(r::MapValue, key) = _mapentry(r, key) !== nothing +function Base.get(r::MapValue{K,V}, key, default) where {K,V} + j = _mapentry(r, key) + return j === nothing ? default : _childvalue(V, r.values, j) +end +function Base.get(default::Base.Callable, r::MapValue{K,V}, key) where {K,V} + j = _mapentry(r, key) + return j === nothing ? default() : _childvalue(V, r.values, j) +end +function Base.getindex(r::MapValue{K,V}, key) where {K,V} + j = _mapentry(r, key) + j === nothing && throw(KeyError(key)) + return _childvalue(V, r.values, j) +end +Base.keys(r::MapValue{K}) where {K} = ListValue{K}(r.keys, r.offset, r.len) +Base.values(r::MapValue{K,V}) where {K,V} = ListValue{V}(r.values, r.offset, r.len) + +# Inside a container or under `:compact` (DataFrame cells), list and map values +# print as the plain Julia values they hold, without element types: +# `[1, missing]`, `Dict("a" => 1)`. +_compactshow(io::IO) = get(io, :compact, false) === true || haskey(io, :typeinfo) +function Base.show(io::IO, r::ListValue) + _compactshow(io) || return invoke(show, Tuple{IO,AbstractArray}, io, r) + # A `:typeinfo` of the row's own type leaves out the element type. + return invoke(show, Tuple{IO,AbstractArray}, IOContext(io, :typeinfo => typeof(r)), r) +end +function Base.show(io::IO, r::MapValue) + _compactshow(io) || return invoke(show, Tuple{IO,AbstractDict}, io, r) + entries = IOContext(io, :typeinfo => eltype(r)) + limit = get(io, :limit, false)::Bool + print(io, "Dict(") + for (count, pair) in enumerate(r) + count > 1 && print(io, ", ") + show(entries, pair) + limit && count == 10 < length(r) && (print(io, ", …"); break) + end + print(io, ')') + return nothing +end + +Base.@propagate_inbounds function Base.getindex(v::ListVector{T}, i::Int) where {T} + @boundscheck checkbounds(v, i) + if Missing <: T + null = GC.@preserve v _isnullslot(v.validity, v.bitoffset, i) + null && return missing + end + lo, hi = GC.@preserve v @inbounds Int(v.offsets[i]), Int(v.offsets[i + 1]) + return Base.nonmissingtype(T)(v.values, lo, hi - lo) +end + +Base.@propagate_inbounds function Base.getindex(v::MapVector{T}, i::Int) where {T} + @boundscheck checkbounds(v, i) + if Missing <: T + null = GC.@preserve v _isnullslot(v.validity, v.bitoffset, i) + null && return missing + end + lo, hi = GC.@preserve v @inbounds Int(v.offsets[i]), Int(v.offsets[i + 1]) + return Base.nonmissingtype(T)(v.keys, v.values, lo, hi - lo) +end + +Base.@propagate_inbounds function Base.getindex(v::StructVector{T}, i::Int) where {T} + @boundscheck checkbounds(v, i) + if Missing <: T + null = GC.@preserve v _isnullslot(v.validity, v.bitoffset, i) + null && return missing + end + return _structrow(Base.nonmissingtype(T), v.children, v.offset + i) +end + +# Generated so each field reads through its own child type, without +# allocating. The generator loops: a comprehension's closure would compile +# once per row type. +@generated function _structrow( + ::Type{R}, + children::Vector{AbstractVector}, + j::Int, +) where {R<:NamedTuple} + fields = Expr[] + for k = 1:fieldcount(R) + push!(fields, :(@inbounds _childvalue($(fieldtype(R, k)), children[$k], j))) + end + return :(R(($(fields...),))) +end + +# Owned column copies replace nested views too: a struct or list can contain +# another list whose child values still point into the input's buffers. +_ownedtype(::Type{T}) where {T} = + T isa Union ? Union{(_ownedtype(S) for S in Base.uniontypes(T))...} : T +_ownedtype(::Type{ListValue{T}}) where {T} = Vector{_ownedtype(T)} +_ownedtype(::Type{MapValue{K,V}}) where {K,V} = Dict{_ownedtype(K),_ownedtype(V)} +_ownedtype(::Type{NamedTuple{N,T}}) where {N,T} = NamedTuple{N,_ownedtype(T)} +_ownedtype(::Type{Pair{K,V}}) where {K,V} = Pair{_ownedtype(K),_ownedtype(V)} +_ownedtype(::Type{T}) where {T<:Tuple} = + T isa Union ? Union{(_ownedtype(S) for S in Base.uniontypes(T))...} : + (T === Union{} || T isa UnionAll || Base.isvatuple(T)) ? T : + Tuple{map(_ownedtype, fieldtypes(T))...} + +# Unknown element types need a value check; concrete scalar types cannot +# contain any of the nested views this module creates. +_mayborrow(::Type{T}) where {T} = + T isa Union ? any(_mayborrow, Base.uniontypes(T)) : !isconcretetype(T) +_mayborrow(::Type{<:ListValue}) = true +_mayborrow(::Type{<:MapValue}) = true +_mayborrow(::Type{T}) where {T<:Union{NamedTuple,Tuple,Pair}} = + !isconcretetype(T) || any(_mayborrow, fieldtypes(T)) +_mayborrow(::Type{T}) where {T<:AbstractVector} = + !isconcretetype(T) || _mayborrow(eltype(T)) +_mayborrow(::Type{T}) where {T<:AbstractDict} = + !isconcretetype(T) || _mayborrow(keytype(T)) || _mayborrow(valtype(T)) + +_hasborrowed(x) = false +_hasborrowed(::Union{ListValue,MapValue}) = true +_hasborrowed(x::Union{NamedTuple,Tuple,Pair}) = any(_hasborrowed, x) +_hasborrowed(x::Union{AbstractVector,AbstractDict}) = + _mayborrow(eltype(x)) && any(_hasborrowed, x) + +_needsowned(c::AbstractVector{T}) where {T} = + _ownedtype(T) !== T || (_mayborrow(T) && any(_hasborrowed, c)) + +_ownedvalue(x) = x +_ownedvalue(x::ListValue) = _copyrows(eltype(x), x) +# A `Dict`, as Arrow 2.x read maps: a repeated key keeps its last value, the +# one the view's lookup finds. +_ownedvalue(x::MapValue{K,V}) where {K,V} = + Dict{_ownedtype(K),_ownedtype(V)}(_ownedvalue(k) => _ownedvalue(v) for (k, v) in x) +_ownedvalue(x::Union{NamedTuple,Tuple}) = map(_ownedvalue, x) +_ownedvalue(x::Pair) = _ownedvalue(first(x)) => _ownedvalue(last(x)) +_ownedvalue(x::AbstractVector) = _needsowned(x) ? _copyrows(eltype(x), x) : x +_ownedvalue(x::AbstractDict) = _hasborrowed(x) ? Dict(_ownedvalue(p) for p in x) : x + +function _copyrows(::Type{T}, v) where {T} + R = _ownedtype(T) + # An abstract container may change its element type when a hidden view + # becomes a vector. Let map infer or widen it from the converted values. + R === T && _mayborrow(T) && !isempty(v) && return map(_ownedvalue, v) + out = Vector{R}(undef, length(v)) + for (i, x) in enumerate(v) + out[i] = _ownedvalue(x) + end + return out +end +Base.copy(v::ArrowVector{T}) where {T} = _copyrows(T, v) +Base.copy(c::ChainedVector{T,<:ArrowVector}) where {T} = + _ownedtype(T) === T ? invoke(copy, Tuple{ChainedVector}, c) : _copyrows(T, c) + +# The values of child `d` (field `f`) as their `_columntype(E)`: the column a +# read builds, copied into owned storage when it is a view of the source +# (rows view their column's children, and outlive their table). +function _childcolumn( + f::AC.Field, + d::AC.ArrayData, + ctx::_ViewContext, + @nospecialize(E::Type) +) + col = _viewcolumn(f, AC.ArrayData[d], ctx, E; source=false) + return _storedas(_columntype(E), col, ctx.budget) +end +_storedas(::Type{C}, col, budget) where {C} = col isa C ? col : _storeas(C, col, budget) + +function _storeas(::Type{Vector{E}}, col, budget) where {E} + _chargeexact!(budget, E, length(col), "nested child values") + return copyto!(Vector{E}(undef, length(col)), 1, col, 1, length(col)) +end +# Converted values (dates, dictionary or run-end encoded values, ...) stored +# as `S`s, with a validity bitmap when some slot is null; null slots hold zero +# bytes. +function _storeas(::Type{Primitive{T,S,Nothing,D}}, col, budget) where {T,S,D} + n = length(col) + _chargeexact!(budget, S, n, "nested child values") + data = Vector{S}(undef, n) + GC.@preserve data Libc.memset(pointer(data), 0, sizeof(data)) + validity = UInt8[] + for (i, x) in enumerate(col) + if x === missing + validity = _setnull(validity, i, n, budget) + else + @inbounds data[i] = x + end + end + values = D === Vector{S} ? data : _UnalignedVector{S}(pointer(data), n, data) + return Primitive{T,S,Nothing,D}(values, validity, 0, nothing) +end + +# A validity bitmap for `n` slots with slot `i` null; an empty `validity` (no +# nulls yet) becomes an all-valid bitmap first. +function _setnull(validity::Vector{UInt8}, i::Int, n::Int, budget) + if isempty(validity) + _chargeexact!(budget, UInt8, cld(n, 8), "nested child validity") + validity = fill(0xff, cld(n, 8)) + end + @inbounds validity[((i - 1) >> 3) + 1] &= ~(0x01 << ((i - 1) & 7)) + return validity +end + +# Converted list, map, and struct values (run-end encoded or dictionary +# children) gather their rows into a column of their own. +function _storeas(::Type{ListVector{T}}, col, budget) where {T} + E = eltype(Base.nonmissingtype(T)) + n = length(col) + _chargeexact!(budget, Int, n + 1, "list offsets") + offsets = zeros(Int, n + 1) + values, validity = E[], UInt8[] + for (i, row) in enumerate(col) + if row === missing + validity = _setnull(validity, i, n, budget) + else + append!(values, row) + end + offsets[i + 1] = length(values) + end + _chargeexact!(budget, E, length(values), "nested child values") + return ListVector{T}( + offsets, + _storedas(_columntype(E), values, budget), + validity, + 0, + nothing, + ) +end +function _storeas(::Type{MapVector{T}}, col, budget) where {T} + K, V = fieldtypes(eltype(Base.nonmissingtype(T))) + n = length(col) + _chargeexact!(budget, Int, n + 1, "map offsets") + offsets = zeros(Int, n + 1) + keys, values, validity = K[], V[], UInt8[] + for (i, row) in enumerate(col) + if row === missing + validity = _setnull(validity, i, n, budget) + else + for (k, v) in row + push!(keys, k) + push!(values, v) + end + end + offsets[i + 1] = length(keys) + end + _chargeexact!(budget, Pair{K,V}, length(keys), "nested child values") + keys = _storedas(_columntype(K), keys, budget) + return MapVector{T}( + offsets, + keys, + _storedas(_columntype(V), values, budget), + validity, + 0, + nothing, + ) +end +function _storeas(::Type{StructVector{T}}, col, budget) where {T} + R = Base.nonmissingtype(T) + n = length(col) + validity = UInt8[] + for (i, row) in enumerate(col) + row === missing && (validity = _setnull(validity, i, n, budget)) + end + # A null row's field slots are never read; they repeat a valid row's + # values (none when every row is null). + k = findfirst(!ismissing, col) + rows = k === nothing ? R[] : R[row === missing ? col[k] : row for row in col] + _chargevector!(budget, AbstractVector, fieldcount(R), "struct children") + children = Vector{AbstractVector}(undef, fieldcount(R)) + for j = 1:fieldcount(R) + F = fieldtype(R, j) + _chargeexact!(budget, F, length(rows), "nested child values") + children[j] = _storedas(_columntype(F), F[getfield(r, j) for r in rows], budget) + end + return StructVector{T}(children, 0, n, validity, 0, nothing) +end +_storeas(::Type{DataStrings.StringVector{E}}, col, budget) where {E} = + DataStrings.StringVector{E}(collect(col)) +function _storeas(::Type{DataStrings.BytesVector{E}}, col, budget) where {E} + payloads = Vector{_StringPayload}(undef, length(col)) + buffers = Vector{UInt8}[] + for (i, x) in enumerate(col) + if x === missing + payloads[i] = DataStrings.PAYLOAD_MISSING + continue + end + bytes = collect(UInt8, x) + n = length(bytes) + payloads[i] = if n <= DataStrings.INLINE_MAX + DataStrings.inline_payload(bytes, 1, n) + else + push!(buffers, bytes) + DataStrings.view_payload(bytes, 1, n, length(buffers) - 1, 0) + end + end + return DataStrings.BytesVector{E}(payloads, buffers, Val(:trusted)) +end + +# The `n + 1` offsets of list or map batch `d`, plus `base` (a map's entries +# struct offset): a view of the batch's buffer when `base` is 0 (copied only if +# misaligned), else shifted copies. +function _listoffsets(d::AC.ArrayData, base::Int, budget) + n = Int64(d.len) + n == 0 && return Int[base] + buf = AC.rolebuffer(d, AC.OFFSETS) + raw = + AC.layoutspec_of(d.type).offsetwidth == 8 ? + _wrapbuffer(Int64, buf, d.offset, n + 1, budget) : + _wrapbuffer(Int32, buf, d.offset, n + 1, budget) + # The view stays valid while the column (whose owner is `d`) is reachable. + base == 0 && return raw + _chargeexact!(budget, Int, n + 1, "list offsets") + return GC.@preserve d raw .+ base +end + +# Not specialized on the row type, so it compiles once for every nested +# column; the per-row loops it calls are function barriers specialized on +# the row and column types. +function _nestedbatch( + @nospecialize(T::Type), + f::AC.Field, + d::AC.ArrayData, + ctx::_ViewContext, +) + t = f.type + R = Base.nonmissingtype(T) + budget = ctx.budget + _chargeobject!(budget, _VIEW_OBJECT_BYTES, "column view") + t isa AC.UnionType && return _unionrows(T, f, d, ctx) + validity, bitoffset = _validitybits(d) + if t isa AC.StructType + children = Vector{AbstractVector}(undef, length(f.children)) + for (k, c) in enumerate(f.children) + children[k] = _childcolumn(c, d.children[k], ctx, fieldtype(R, k)) + end + return StructVector{T}(children, Int(d.offset), Int(d.len), validity, bitoffset, d) + end + if t isa AC.MapType + entries, ed = f.children[1], d.children[1] + K, V = fieldtypes(eltype(R)) + keys = _childcolumn(entries.children[1], ed.children[1], ctx, K) + vals = _childcolumn(entries.children[2], ed.children[2], ctx, V) + offsets = _listoffsets(d, Int(ed.offset), budget) + return MapVector{T}(offsets, keys, vals, validity, bitoffset, d) + end + R <: Array && return _tensorrows(T, f, d, ctx, _tensorshape(f)::Tuple) + E = eltype(R) + if t isa AC.FixedSizeListType && R <: Tuple + # `NTuple{0,E}` is `Tuple{}`, whose rows read no child values. + N = fieldcount(R) + N == 0 && return _tuplerows(T, Val(0), nothing, d, budget) + child = _viewcolumn(f.children[1], AC.ArrayData[d.children[1]], ctx, E) + return _tuplerows(T, Val(N), child, d, budget) + end + values = _childcolumn(f.children[1], d.children[1], ctx, E) + t isa AC.ListViewType && + return _listviewvector(T, d, values, validity, bitoffset, budget) + offsets = if t isa AC.FixedSizeListType + w = t.listsize + _chargeexact!(budget, Int, d.len + 1, "fixed-size list offsets") + Int[(Int(d.offset) + i) * w for i = 0:(d.len)] + else + _listoffsets(d, 0, budget) + end + return ListVector{T}(offsets, values, validity, bitoffset, d) +end + +# A list view batch: rows may overlap or come in any order, so the child +# values of each row are gathered in row order (their own storage keeps +# flat values only; see `_viewablelayout`). +function _listviewvector( + @nospecialize(T::Type), + d::AC.ArrayData, + values, + validity, + bitoffset, + budget, +) + n = Int64(d.len) + _chargeexact!(budget, Int, n + 1, "list view offsets") + offsets = zeros(Int, n + 1) + n == 0 && + return ListVector{T}(offsets, _gathered(values, Int[]), validity, bitoffset, d) + wide = AC.layoutspec_of(d.type).offsetwidth == 8 + starts, sizes = GC.@preserve d ( + wide ? + ( + Int.(_wrapbuffer(Int64, AC.rolebuffer(d, AC.ELEMENT_OFFSETS), d.offset, n)), + Int.(_wrapbuffer(Int64, AC.rolebuffer(d, AC.SIZES), d.offset, n)), + ) : + ( + Int.(_wrapbuffer(Int32, AC.rolebuffer(d, AC.ELEMENT_OFFSETS), d.offset, n)), + Int.(_wrapbuffer(Int32, AC.rolebuffer(d, AC.SIZES), d.offset, n)), + ) + ) + for i = 1:n + len = _isnullslot(validity, bitoffset, i) ? 0 : sizes[i] + offsets[i + 1] = offsets[i] + len + end + _chargeexact!(budget, Int, offsets[end], "list view value index") + take = Vector{Int}(undef, offsets[end]) + for i = 1:n + for k = 1:(offsets[i + 1] - offsets[i]) + take[offsets[i] + k] = starts[i] + k + end + end + return ListVector{T}(offsets, _gathered(values, take), validity, bitoffset, d) +end + +_gathered(values::Vector, take) = values[take] +_gathered(v::Primitive{T,S,Nothing}, take) where {T,S} = + _storeas(typeof(v), view(v, take), nothing) +_gathered(v::DataStrings.StringVector{E}, take) where {E} = + DataStrings.StringVector{E}(v.payloads[take], v.buffers, Val(:trusted)) +_gathered(v::DataStrings.BytesVector{E}, take) where {E} = + DataStrings.BytesVector{E}(v.payloads[take], v.buffers, Val(:trusted)) + +function _tuplerows(::Type{T}, ::Val{N}, child, d::AC.ArrayData, budget) where {T,N} + n = Int(d.len) + _chargeexact!(budget, T, n, "fixed-size list rows") + out = Vector{T}(undef, n) + validity, bitoffset = _validitybits(d) + GC.@preserve validity child for i = 1:n + if Missing <: T && _isnullslot(validity, bitoffset, i) + out[i] = missing + else + base = (Int(d.offset) + i - 1) * N + out[i] = ntuple(j -> child[base + j], Val(N)) + end + end + return out +end + +# The value at child index `j` of the `k`th column, unrolled over the tuple. +_pickchild(cols::Tuple, k::Int, j::Int) = + k == 1 ? first(cols)[j] : _pickchild(Base.tail(cols), k - 1, j) +_pickchild(::Tuple{}, ::Int, ::Int) = throw(ValidationError("union child out of range")) + +function _unionrows(@nospecialize(T::Type), f::AC.Field, d::AC.ArrayData, ctx::_ViewContext) + t = f.type::AC.UnionType + budget = ctx.budget + cols = Tuple( + Any[ + _viewcolumn( + c, + AC.ArrayData[d.children[k]], + ctx, + _vieweltype(c, AC.ArrayData[d.children[k]], ctx), + ) for (k, c) in enumerate(f.children) + ], + ) + n = Int(d.len) + _chargeexact!(budget, T, n, "union rows") + out = Vector{T}(undef, n) + ids = _wrapbuffer(Int8, AC.rolebuffer(d, AC.TYPE_IDS), d.offset, n, budget) + dense = t.mode == AC.DenseMode + offsets = + dense ? + _wrapbuffer(Int32, AC.rolebuffer(d, AC.ELEMENT_OFFSETS), d.offset, n, budget) : + Int32[] + GC.@preserve ids offsets d _fillunion!( + out, + cols, + ids, + offsets, + t.typeids, + dense, + Int(d.offset), + ) + return out +end + +# Specialized on the children's column types (a function barrier). +function _fillunion!(out, cols::Tuple, ids, offsets, typeids, dense::Bool, base::Int) + for i in eachindex(out) + k = findfirst(==(ids[i]), typeids) + k === nothing && + throw(ValidationError("union type id $(ids[i]) not in declared domain")) + j = dense ? Int(offsets[i]) + 1 : base + i + out[i] = _pickchild(cols, k, j) + end + return out +end + +# --- fixed-shape tensors ---------------------------------------------------------- +# +# The canonical `arrow.fixed_shape_tensor` extension stores each tensor as a +# fixed-size list of its elements in row-major order over its physical +# `shape`; `permutation` says which physical dimension each logical one is. +# Rows read as `Array{T,N}`s of the logical shape. + +const _TENSOR_EXTENSION = "arrow.fixed_shape_tensor" + +# `(shape, permutation)` of a tensor field (1-based permutation), or nothing +# when `f` is not one this reader can read: a fixed-size list of a flat, +# unlabelled child with a well-formed `shape`. +function _tensorshape(f::AC.Field) + f.type isa AC.FixedSizeListType || return nothing + ext = _arrowtypesextension(f) + (ext === nothing || ext[1] != _TENSOR_EXTENSION) && return nothing + (length(f.children) == 1 && !_hasarrowtypesextension(f.children[1])) || return nothing + child = f.children[1] + (_viewablelayout(child) && !_nestedvalue(child)) || return nothing + shape = _jsonints(ext[2], "shape") + shape === nothing && return nothing + perm = _jsonints(ext[2], "permutation") + perm = perm === nothing ? collect(1:length(shape)) : perm .+ 1 + ( + all(>=(0), shape) && + length(perm) == length(shape) && + isperm(perm) && + prod(shape; init=1) == f.type.listsize + ) || return nothing + return (Tuple(shape), Tuple(perm)) +end + +# The integers of JSON array `key` in `json`, or nothing. +function _jsonints(json::AbstractString, key::String) + m = match(Regex("\"" * key * "\"\\s*:\\s*\\[([^\\]]*)\\]"), json) + m === nothing && return nothing + values = Int[] + for part in split(m[1], ','; keepempty=false) + x = tryparse(Int, strip(part)) + x === nothing && return nothing + push!(values, x) + end + return values +end + +# Tensor rows are `Array{E,N}`: `E` is the element type, without `Missing` +# unless an element of a non-null tensor is null. +function _tensortype(f::AC.Field, ctx::_ViewContext, datas, shape) + child = f.children[1] + B = _viewbasetype(child, ctx, _childdatas(datas, 1)) + w = (f.type::AC.FixedSizeListType).listsize + nulls = any(d -> _tensornulls(child, d, w), datas) + return Array{nulls ? Union{Missing,B} : B,length(first(shape))} +end + +# Whether a non-null row of tensor batch `d` holds a null element. +function _tensornulls(child::AC.Field, d::AC.ArrayData, w::Int) + c = d.children[1] + _viewnulls(child, c) || return false + # Null elements under a null tensor are never read. + rows, rowoff = _validitybits(d) + elems, elemoff = _validitybits(c) + isempty(elems) && return true # dictionary or run-end values with nulls + GC.@preserve d for i = 1:Int(d.len) + _isnullslot(rows, rowoff, i) && continue + base = (Int(d.offset) + i - 1) * w + any(j -> _isnullslot(elems, elemoff, base + j), 1:w) && return true + end + return false +end + +function _tensorrows(@nospecialize(T::Type), f::AC.Field, d::AC.ArrayData, ctx, shapeperm) + shape, perm = shapeperm + R = Base.nonmissingtype(T) + child = _viewcolumn(f.children[1], AC.ArrayData[d.children[1]], ctx, eltype(R)) + return _uninferred(_tensorrows, T, R, child, d, ctx.budget, shape, perm) +end + +function _tensorrows( + ::Type{T}, + ::Type{R}, + child, + d::AC.ArrayData, + budget, + shape, + perm, +) where {T,R} + n = Int(d.len) + w = prod(shape; init=1) + N = length(shape) + _chargeexact!(budget, T, n, "tensor rows") + _chargeexact!(budget, eltype(R), n * w, "tensor values") + out = Vector{T}(undef, n) + validity, bitoffset = _validitybits(d) + for i = 1:n + if Missing <: T && _isnullslot(validity, bitoffset, i) + out[i] = missing + continue + end + base = (Int(d.offset) + i - 1) * w + flat = eltype(R)[child[base + j] for j = 1:w] + # Row-major storage: reverse the dimensions to reshape, then back. + physical = permutedims(reshape(flat, reverse(shape)), N:-1:1) + out[i] = perm == Tuple(1:N) ? physical : permutedims(physical, perm) + end + return out +end + +# Tensor rows back to their row-major storage values, for a rewrite. +function _tensorstorage(v::AbstractVector, (shape, perm)) + N = length(shape) + return [ + row isa AbstractArray && ndims(row) == N ? + vec( + permutedims(perm == Tuple(1:N) ? row : permutedims(row, invperm(perm)), N:-1:1), + ) : row for row in v + ] +end diff --git a/src/cstream.jl b/src/cstream.jl new file mode 100644 index 00000000..d7a40be7 --- /dev/null +++ b/src/cstream.jl @@ -0,0 +1,429 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 stream interface (ArrowArrayStream) over the C data adapter in +# cdata.jl: an exported stream is one export whose `get_schema` and `get_next` +# results are exports of their own, and an imported stream is a moved struct +# whose batches import as arrays do. Kept apart from cdata.jl so the trim gate +# can compile the array interface without the IPC reader's `Limits`. + +""" + Arrow.CArrowArrayStream + +The C stream interface's `ArrowArrayStream` struct, field for field. +""" +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) + +# An all-zero struct: a released one (`release == NULL`), which also ends a stream. +function _zeroed(::Type{T}) where {T} + r = Ref{T}() + GC.@preserve r Libc.memset(Base.unsafe_convert(Ptr{T}, r), 0, sizeof(T)) + return r +end + +# The batches an exported stream yields, pulled one `get_next` at a time. +mutable struct ExportedStreamState + const batchfield::Field # struct-typed: children are the schema + const batches::Vector{AC.RecordBatch} + nextindex::Int +end + +# A record batch crosses as a struct array: the struct field's children are +# the schema's fields and it carries the schema-level metadata (the +# C++/pyarrow convention for `schema.metadata`). +_batchfield(sch::Schema) = Field( + "", + StructType(); + nullable=false, + metadata=sch.metadata, + children=collect(Field, sch.fields), +) +_structarray(b::AC.RecordBatch) = ArrayData( + StructType(), + b.nrows, + [BufferSlice()]; + children=collect(ArrayData, b.columns), + nullcount=0, +) + +# The exported stream's state, or `nothing` once released or reaped. +function _stream_state(sp::Ptr{CArrowArrayStream}) + stream = unsafe_load(sp) + (stream.release == C_NULL || stream.private_data == C_NULL) && return nothing + e = lock(REGISTRY_LOCK) do + get(EXPORTS, Ptr{Int64}(stream.private_data), nothing) + end + e === nothing && return nothing + return e.roots[1]::ExportedStreamState +end + +# Store `err` as the stream's last error, for `get_last_error`, in the 8 bytes +# after the export's live count. Best effort: a C callback cannot throw. +function _streamerror!(sp::Ptr{CArrowArrayStream}, err) + try + control = unsafe_load(sp).private_data + control == C_NULL && return nothing + slot = Ptr{Ptr{UInt8}}(control + 8) + old = unsafe_load(slot) + unsafe_store!(slot, Ptr{UInt8}(C_NULL)) + old == C_NULL || Libc.free(old) + message = replace(sprint(showerror, err), '\0' => ' ') + n = ncodeunits(message) + p = Ptr{UInt8}(Libc.malloc(n + 1)) + p == C_NULL && return nothing + GC.@preserve message unsafe_copyto!(p, pointer(message), n) + unsafe_store!(p, 0x00, n + 1) + unsafe_store!(slot, p) + catch + end + return nothing +end + +function _stream_get_schema(sp::Ptr{CArrowArrayStream}, out::Ptr{CArrowSchema})::Cint + sp == C_NULL && return EINVAL + try + state = _stream_state(sp) + state === nothing && return EINVAL + out == C_NULL && throw(ArgumentError("ArrowSchema output pointer is NULL")) + e = _newexport(Any[state.batchfield]) + try + _export_schema!( + e, + state.batchfield, + @cfunction(_release_schema, Cvoid, (Ptr{CArrowSchema},)); + at=out, + ) + _register!(e) + catch + _free!(e) + rethrow() + end + return Cint(0) + catch err + _streamerror!(sp, err) + return EINVAL + end +end + +function _stream_get_next(sp::Ptr{CArrowArrayStream}, out::Ptr{CArrowArray})::Cint + sp == C_NULL && return EINVAL + try + state = _stream_state(sp) + state === nothing && return EINVAL + out == C_NULL && throw(ArgumentError("ArrowArray output pointer is NULL")) + i = state.nextindex + if i > length(state.batches) + unsafe_store!(out, _zeroed(CArrowArray)[]) + return Cint(0) + end + d = _structarray(state.batches[i]) + validate_semantic(state.batchfield, d) + e = _newexport(Any[d]) + try + _export_array!( + e, + d, + @cfunction(_release_array, Cvoid, (Ptr{CArrowArray},)); + at=out, + ) + _register!(e) + catch + _free!(e) + rethrow() + end + state.nextindex = i + 1 + return Cint(0) + catch err + _streamerror!(sp, err) + return EINVAL + end +end + +function _stream_get_last_error(sp::Ptr{CArrowArrayStream})::Ptr{UInt8} + sp == C_NULL && return Ptr{UInt8}(C_NULL) + stream = unsafe_load(sp) + (stream.release == C_NULL || stream.private_data == C_NULL) && return Ptr{UInt8}(C_NULL) + return unsafe_load(Ptr{Ptr{UInt8}}(stream.private_data + 8)) +end + +function _stream_release(sp::Ptr{CArrowArrayStream})::Cvoid + sp == C_NULL && return nothing + stream = unsafe_load(sp) + stream.release == C_NULL && return nothing + _store_field!(sp, Val(:release), Ptr{Cvoid}(C_NULL)) + stream.private_data == C_NULL || + Base.unsafe_modify!(Ptr{Int64}(stream.private_data), -, 1, :sequentially_consistent) + return nothing +end + +# Core's structural schema invariants (endianness, UTF-8 names, +# 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_descriptor_of(f.type) + for c in f.children + _validate_stream_field(c) + end + return nothing +end + +""" + Arrow.export_stream!(streamptr::Ptr{CArrowArrayStream}, schema::Arrow.Schema, batches) + +The low-level form of [`Arrow.export_stream!`](@ref): fill the `ArrowArrayStream` the +caller allocated at `streamptr` with a stream that yields `batches`, a vector of +`Arrow.RecordBatch` values with `schema`. Each batch the consumer pulls is a struct +array of the batch's columns, exported as [`Arrow.to_c_data`](@ref) exports a column. +The batches stay reachable until the consumer releases the stream. +""" +function export_stream!( + sp::Ptr{CArrowArrayStream}, + sch::Schema, + batches::AbstractVector{AC.RecordBatch}, +) + 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 + state = ExportedStreamState(_batchfield(sch), collect(AC.RecordBatch, batches), 1) + e = _newexport(Any[state]) + unsafe_store!(e.live, 1) # the stream itself; its release counts it off + unsafe_store!( + sp, + CArrowArrayStream( + @cfunction( + _stream_get_schema, + Cint, + (Ptr{CArrowArrayStream}, Ptr{CArrowSchema}) + ), + @cfunction(_stream_get_next, Cint, (Ptr{CArrowArrayStream}, Ptr{CArrowArray})), + @cfunction(_stream_get_last_error, Ptr{UInt8}, (Ptr{CArrowArrayStream},)), + @cfunction(_stream_release, Cvoid, (Ptr{CArrowArrayStream},)), + Ptr{Cvoid}(e.live), + ), + ) + _register!(e) + return sp +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, +# and one atomic swap picks the single releaser between `release!` and the GC +# finalizer. +mutable struct StreamOwner + const block::Ptr{CArrowArrayStream} + const thread::Int # where a collected owner releases (`_releasethread`) + @atomic released::Bool +end + +function _movestream(sp::Ptr{CArrowArrayStream}) + block = Ptr{CArrowArrayStream}(Libc.malloc(sizeof(CArrowArrayStream))) + block == C_NULL && throw(OutOfMemoryError()) + unsafe_store!(block, unsafe_load(sp)) + _store_field!(sp, Val(:release), Ptr{Cvoid}(C_NULL)) + try + return finalizer(_finalizeowner, StreamOwner(block, _releasethread(), false)) + catch + try + _release_producer!(block) + finally + Libc.free(block) + end + rethrow() + end +end + +Base.deepcopy_internal(o::StreamOwner, ::IdDict) = o + +function release!(o::StreamOwner) + @atomicswap(o.released = true) && return nothing + GC.@preserve o try + _release_producer!(o.block) + finally + Libc.free(o.block) + end + return nothing +end + +# A stream imported with `from_c_stream`. `nextbatch!` returns its next +# `RecordBatch`, or `nothing` at the end, and `release!` releases the +# producer's stream; batches already read stay valid. An error the producer +# reports is a `ValidationError` with the producer's message. One task uses a +# stream at a time: overlapping calls throw a `ConcurrencyViolationError`. +mutable struct ImportedStream <: AC.RecordBatchSource + const owner::StreamOwner + const batchfield::Field + const schema::Schema + const limits::Limits + done::Bool + @atomic pulling::Bool +end + +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) + +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 + +# Import a C stream: take over the stream, marking the struct at `sp` released, +# and read its schema, whose fields are the batches' columns. `limits` bounds +# what the facade allocates for each batch. +function from_c_stream(sp::Ptr{CArrowArrayStream}; limits::Limits=Limits()) + 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 = _movestream(sp) + try + out = _zeroed(CArrowSchema) + 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_producer!(Base.unsafe_convert(Ptr{CArrowSchema}, out)) + 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, limits, false, false) + catch + release!(owner) + rethrow() + end +end + +function AC.nextbatch!(s::ImportedStream) + _claimstreamcall!(s) + try + # Fail closed on a released stream even when it already ended: release + # ends 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 = _zeroed(CArrowArray) + 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") + if out[].release == C_NULL + s.done = true + return nothing + end + owner = GC.@preserve out _moveowner(Base.unsafe_convert(Ptr{CArrowArray}, out)) + d = _importarray(s.batchfield, owner) + try + return _structbatch(s.schema, d) + catch + release!(owner) + rethrow() + end + finally + _finishstreamcall!(s) + end +end + +# A record batch has no row validity, so null struct rows are refused rather +# than exposing the child values under them (Arrow C++ refuses them too). +# Row i of a struct array is child row `offset + i`, and children may run past +# `offset + length`; each column is that visible child window. +function _structbatch(sch::Schema, d::ArrayData) + AC.nullcount(d) == 0 || throw( + ValidationError("a struct array with null rows has no record batch or table form"), + ) + columns = ArrayData[ + d.offset == 0 && c.len == d.len ? c : + AC._arraydata( + c.type, + d.len, + c.buffers, + c.offset + d.offset, + c.children, + c.dictionary, + c.owner, + -1, + ) for c in d.children + ] + return AC.RecordBatch(sch, columns, d.len) +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_dictionary.jl b/src/ipc_dictionary.jl new file mode 100644 index 00000000..9aac2061 --- /dev/null +++ b/src/ipc_dictionary.jl @@ -0,0 +1,430 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Dictionary updates retain immutable snapshots. Concatenate physical storage, +# never materialized Julia values: Union routes, nested dictionary indices, +# logical descriptors, and null payloads must not pass through type inference. + +mutable struct _DictionaryRange + const data::ArrayData + const first::Int64 # zero-based logical position within data + len::Int64 +end + +function _dictionaryvector(::Type{T}, n::Integer, limits::Limits, budget) where {T} + bytes = AC.checked_mul(Int64(n), Int64(sizeof(T))) + bytes <= limits.max_buffer_bytes || _limiterror( + "combined dictionary buffer of $bytes bytes", + "max_buffer_bytes", + limits.max_buffer_bytes, + ) + _chargevector!(budget, T, n, "dictionary concatenation") + return zeros(T, n) +end + +function _dictionaryrange!(ranges, data, first, len, budget) + len == 0 && return ranges + if !isempty(ranges) + previous = last(ranges) + if previous.data === data && previous.first + previous.len == first + previous.len = AC.checked_add(previous.len, len) + return ranges + end + end + # Charge capacity growth before push!, including the range object. + _chargeobject!(budget, 2 * sizeof(_DictionaryRange), "dictionary child ranges") + push!(ranges, _DictionaryRange(data, first, len)) + return ranges +end + +function _dictionaryput!( + bytes::Vector{UInt8}, + offset::Integer, + value::Integer, + width::Int; + signed=true, +) + bits = 8 * width - signed + maximum = bits >= 63 ? typemax(Int64) : (Int64(1) << bits) - 1 + 0 <= value <= maximum || + throw(ValidationError("combined dictionary offset exceeds its integer width")) + for k = 0:(width - 1) + bytes[offset + k + 1] = UInt8((value >> (8 * k)) & 0xff) + end +end + +function _dictionarycopy!(out, dest, src::BufferSlice, first, len) + len == 0 && return + 0 <= dest <= length(out) && 0 <= len <= length(out) - dest || + throw(ValidationError("dictionary copy escapes destination buffer")) + # subslice checks the source before exposing bytes. + slice = AC.subslice(src, first, len) + GC.@preserve out slice begin + unsafe_copyto!(pointer(out, dest + 1), AC.sliceptr(slice), len) + end +end + +function _dictionarylength(ranges, limits) + n = Int64(0) + for r in ranges + n = AC.checked_add(n, r.len) + end + n <= limits.max_array_length || _limiterror( + "combined dictionary length $n", + "max_array_length", + limits.max_array_length, + ) + return n +end + +# Also used recursively to compact selected child ranges. Every range points +# into a pool that has already passed full semantic validation. +function _dictionaryconcat(f::Field, ranges, limits::Limits, budget) + n = _dictionarylength(ranges, limits) + t = f.type + spec = AC.layoutspec_of(t) + _chargeobject!(budget, sizeof(ArrayData), "combined dictionary array") + _chargevector!(budget, BufferSlice, length(spec.buffers), "dictionary buffers") + _chargevector!(budget, ArrayData, length(f.children), "dictionary children") + buffers = BufferSlice[] + children = ArrayData[] + dictionary = nothing + nc = Int64(0) + if !isempty(spec.buffers) && spec.buffers[1] == AC.VALIDITY + validity = _dictionaryvector(UInt8, cld(n, 8), limits, budget) + k = 0 + for r in ranges, i = (r.first + 1):(r.first + r.len) + if AC.isvalid_at(r.data, i) + validity[(k >> 3) + 1] |= UInt8(1 << (k & 7)) + else + nc += 1 + end + k += 1 + end + push!(buffers, nc == 0 ? BufferSlice() : AC._databuffer(validity)) + end + if t isa NullType + nc = n + elseif t isa Union{Utf8Type,BinaryType,ListType,ListViewType,MapType} + width = spec.offsetwidth + isview = t isa ListViewType + isbinary = t isa Union{Utf8Type,BinaryType} + offsets = _dictionaryvector( + UInt8, + AC.checked_mul(n + !isview, Int64(width)), + limits, + budget, + ) + sizes = + isview ? + _dictionaryvector(UInt8, AC.checked_mul(n, Int64(width)), limits, budget) : + UInt8[] + childranges = _DictionaryRange[] + total = Int64(0) + k = Int64(0) + # Compute lengths and child selection before allocating the payload. + for r in ranges, i = (r.first + 1):(r.first + r.len) + _dictionaryput!(offsets, k * width, total, width) + if AC.isvalid_at(r.data, i) + lo, hi = + isview ? AC._listview_range(t, r.data, i) : + AC._offsets_at(r.data, i, width == 8) + len = isview ? hi : hi - lo + isview && _dictionaryput!(sizes, k * width, len, width) + total = AC.checked_add(total, len) + if !isbinary + total <= limits.max_array_length || _limiterror( + "combined dictionary child length $total", + "max_array_length", + limits.max_array_length, + ) + _dictionaryrange!(childranges, r.data.children[1], lo, len, budget) + end + end + k += 1 + end + !isview && _dictionaryput!(offsets, n * width, total, width) + push!(buffers, AC._databuffer(offsets)) + if isbinary + payload = _dictionaryvector(UInt8, total, limits, budget) + pos = Int64(0) + for r in ranges, i = (r.first + 1):(r.first + r.len) + AC.isvalid_at(r.data, i) || continue + lo, hi = AC._offsets_at(r.data, i, width == 8) + _dictionarycopy!(payload, pos, AC.rolebuffer(r.data, AC.DATA), lo, hi - lo) + pos += hi - lo + end + push!(buffers, AC._databuffer(payload)) + else + isview && push!(buffers, AC._databuffer(sizes)) + push!(children, _dictionaryconcat(f.children[1], childranges, limits, budget)) + end + elseif t isa Union{StructType,FixedSizeListType} || + (t isa UnionType && t.mode == AC.SparseMode) + if t isa UnionType + ids = _dictionaryvector(UInt8, n, limits, budget) + k = 0 + for r in ranges + _dictionarycopy!( + ids, + k, + AC.rolebuffer(r.data, AC.TYPE_IDS), + r.data.offset + r.first, + r.len, + ) + k += r.len + end + push!(buffers, AC._databuffer(ids)) + end + width = t isa FixedSizeListType ? Int64(t.listsize) : Int64(1) + for j in eachindex(f.children) + childranges = _DictionaryRange[] + for r in ranges + _dictionaryrange!( + childranges, + r.data.children[j], + AC.checked_mul(r.data.offset + r.first, width), + AC.checked_mul(r.len, width), + budget, + ) + end + push!(children, _dictionaryconcat(f.children[j], childranges, limits, budget)) + end + elseif t isa UnionType + ids = _dictionaryvector(UInt8, n, limits, budget) + offsets = _dictionaryvector(UInt8, AC.checked_mul(n, Int64(4)), limits, budget) + _chargevector!( + budget, + Vector{_DictionaryRange}, + length(f.children), + "dictionary union routes", + ) + routes = [_DictionaryRange[] for _ in f.children] + counts = _dictionaryvector(Int64, length(f.children), limits, budget) + k = Int64(0) + for r in ranges, i = (r.first + 1):(r.first + r.len) + tid = AC.loadat(AC.rolebuffer(r.data, AC.TYPE_IDS), Int8, r.data.offset + i - 1) + j = findfirst(==(tid), t.typeids)::Int + off = Int64( + AC.loadat( + AC.rolebuffer(r.data, AC.ELEMENT_OFFSETS), + Int32, + (r.data.offset + i - 1) * 4, + ), + ) + ids[k + 1] = reinterpret(UInt8, tid) + _dictionaryput!(offsets, k * 4, counts[j], 4) + _dictionaryrange!(routes[j], r.data.children[j], off, Int64(1), budget) + counts[j] += 1 + k += 1 + end + append!(buffers, (AC._databuffer(ids), AC._databuffer(offsets))) + for j in eachindex(f.children) + push!(children, _dictionaryconcat(f.children[j], routes[j], limits, budget)) + end + elseif t isa RunEndEncodedType + # Intersect physical runs with each requested logical range. This + # handles sliced/overlong last runs without expanding one value per row. + runranges = _DictionaryRange[] + _chargevector!(budget, Int64, 0, "dictionary run ends") + ends = Int64[] + total = Int64(0) + for r in ranges + pos = r.first + 1 + stop = r.first + r.len + while pos <= stop + run = AC._ree_runindex(r.data, pos) + runend = + Int64(AC.getvalue(f.children[1], r.data.children[1], run)) - + r.data.offset + len = min(runend, stop) - pos + 1 + total = AC.checked_add(total, len) + _chargeobject!(budget, 2 * sizeof(Int64), "dictionary run ends") + push!(ends, total) + _dictionaryrange!(runranges, r.data.children[2], run - 1, Int64(1), budget) + pos += len + end + end + _chargeobject!(budget, sizeof(ArrayData), "dictionary run-end array") + _chargevector!(budget, BufferSlice, 2, "dictionary run-end buffers") + rt = f.children[1].type::IntType + width = AC.primwidth(rt) + data = _dictionaryvector( + UInt8, + AC.checked_mul(Int64(length(ends)), Int64(width)), + limits, + budget, + ) + for (i, value) in enumerate(ends) + _dictionaryput!(data, (i - 1) * width, value, width) + end + push!( + children, + ArrayData(rt, length(ends), [BufferSlice(), AC._databuffer(data)]; nullcount=0), + ) + push!(children, _dictionaryconcat(f.children[2], runranges, limits, budget)) + elseif t isa ViewType + views = _dictionaryvector(UInt8, AC.checked_mul(n, Int64(16)), limits, budget) + # Keep long-value buffers borrowed; only descriptors need rebasing. + _chargedict!(budget, ArrayData, Int64, length(ranges), "dictionary view buffers") + bases = IdDict{ArrayData,Int64}() + databuffers = BufferSlice[] + k = Int64(0) + for r in ranges + base = get(bases, r.data, Int64(-1)) + if base < 0 + base = Int64(length(databuffers)) + bases[r.data] = base + _chargevector!( + budget, + BufferSlice, + 2 * (length(r.data.buffers) - 2), + "dictionary view buffers", + ) + append!(databuffers, r.data.buffers[3:end]) + end + for i = (r.first + 1):(r.first + r.len) + if AC.isvalid_at(r.data, i) + src = AC.rolebuffer(r.data, AC.VIEWS) + off = (r.data.offset + i - 1) * 16 + _dictionarycopy!(views, k * 16, src, off, Int64(16)) + if AC.loadat(src, Int32, off) > AC.VIEW_INLINE_MAX + idx = AC.loadat(src, Int32, off + 8) + _dictionaryput!( + views, + k * 16 + 8, + AC.checked_add(base, Int64(idx)), + 4, + ) + end + end + k += 1 + end + end + push!(buffers, AC._databuffer(views)) + append!(buffers, databuffers) + elseif t isa DictionaryType + # Nested dictionary values can refer to different snapshots. Concatenate + # unique pools and rebase their indices, preserving their value routes. + _chargedict!(budget, ArrayData, Int64, length(ranges), "nested dictionary pools") + bases = IdDict{ArrayData,Int64}() + poolranges = _DictionaryRange[] + total = Int64(0) + for r in ranges + pool = r.data.dictionary::ArrayData + if !haskey(bases, pool) + bases[pool] = total + _dictionaryrange!(poolranges, pool, Int64(0), pool.len, budget) + total = AC.checked_add(total, pool.len) + end + end + dictionary = + length(poolranges) == 1 ? poolranges[1].data : + _dictionaryconcat(AC.dictvaluefield(f, t), poolranges, limits, budget) + width = AC.primwidth(t.indextype) + data = _dictionaryvector(UInt8, AC.checked_mul(n, Int64(width)), limits, budget) + k = Int64(0) + for r in ranges, i = (r.first + 1):(r.first + r.len) + if AC.isvalid_at(r.data, i) + idx = AC._load_int( + AC.rolebuffer(r.data, AC.DATA), + t.indextype, + (r.data.offset + i - 1) * width, + ) + value = AC.checked_add(bases[r.data.dictionary], Int64(idx)) + _dictionaryput!(data, k * width, value, width; signed=t.indextype.signed) + end + k += 1 + end + push!(buffers, AC._databuffer(data)) + elseif t isa BoolType + data = _dictionaryvector(UInt8, cld(n, 8), limits, budget) + k = Int64(0) + for r in ranges, i = (r.first + 1):(r.first + r.len) + if AC.getbit(AC.rolebuffer(r.data, AC.DATA), r.data.offset + i - 1) + data[(k >> 3) + 1] |= UInt8(1 << (k & 7)) + end + k += 1 + end + push!(buffers, AC._databuffer(data)) + else + width = spec.fixedwidth + width >= 0 || throw(ValidationError("unsupported dictionary value layout")) + data = _dictionaryvector(UInt8, AC.checked_mul(n, Int64(width)), limits, budget) + k = Int64(0) + for r in ranges + len = AC.checked_mul(r.len, Int64(width)) + _dictionarycopy!( + data, + k, + AC.rolebuffer(r.data, AC.DATA), + AC.checked_mul(r.data.offset + r.first, Int64(width)), + len, + ) + k += len + end + push!(buffers, AC._databuffer(data)) + end + return ArrayData(t, n, buffers; children, dictionary, nullcount=nc) +end + +function _dictionarytransition(id, isdelta, present; file=false) + isdelta && + !present && + throw(ValidationError("delta dictionary id $id has no base dictionary")) + file && + !isdelta && + present && + throw(ValidationError("dictionary replacement is forbidden in the IPC file format")) + return nothing +end + +function _updatedictionary!( + dicts, + validated, + id, + isdelta, + vf, + decoded, + limits, + budget; + file=false, +) + _dictionarytransition(id, isdelta, haskey(dicts, id); file) + validate_semantic(vf, decoded) + if isdelta + previous = dicts[id] + if decoded.len == 0 + return previous + end + _chargevector!(budget, _DictionaryRange, 2, "dictionary delta ranges") + _chargeobject!(budget, 2 * sizeof(_DictionaryRange), "dictionary delta ranges") + decoded = _dictionaryconcat( + vf, + [ + _DictionaryRange(previous, 0, previous.len), + _DictionaryRange(decoded, 0, decoded.len), + ], + limits, + budget, + ) + validate_semantic(vf, decoded) + end + haskey(dicts, id) && delete!(validated, dicts[id]) + validated[decoded] = nothing + dicts[id] = decoded + return decoded +end diff --git a/src/ipc_read.jl b/src/ipc_read.jl new file mode 100644 index 00000000..6de62e8f --- /dev/null +++ b/src/ipc_read.jl @@ -0,0 +1,2754 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 and, at the end of this file, the file-format reader +# (readfile/ArrowFile), as thin peers over ArrowCore. +# +# * 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 +# --------------------------------------------------------------------------- + +const _DEFAULT_MAX_NEW_SYMBOLS = 1_000_000 + +""" + 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` | none | one message body | +| `max_buffer_bytes` | none | one buffer, compressed or not | +| `max_total_allocated_bytes` | physical memory | allocation per read or stream batch | +| `max_allocation_ratio` | 64 | allocation per input byte | +| `max_messages` | 1,000,000 | messages per file or `Arrow.Table` | +| `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 | +| `max_new_symbols` | 1,000,000 | new `Symbol` values per read | + +A read may allocate `max_allocation_ratio` bytes per byte of its input (the +file, buffer, IO, or object), at least 256 MiB, and at most +`max_total_allocated_bytes`. A compressed buffer counts at its decompressed +size, up to the size its column needs, so a compressed file reads with the +allowance it would have uncompressed. Each batch an [`Arrow.Stream`](@ref) +yields gets its own allowance from its own bytes. Ordinary columns cost a +small multiple of their stored bytes, so files of any size and compression +ratio read with the defaults; lengths a small malformed file declares without +the bytes to back them (run-end encoded, view, null-child, or decompressed +lengths beyond what the columns need) fail before they are allocated. Raise +`max_allocation_ratio` for trusted data that expands further, such as long runs +or heavily repeated dictionary values. + +A `Symbol` column (Julia's `JuliaLang.Symbol` extension) creates the `Symbol`s +this process has not seen yet, and Julia never frees them, so a read (or one +stream batch) creates at most `max_new_symbols`; a column needing more is +refused before any of its values is converted. `select=[:col => String]` reads +such a column as strings instead. + +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 = typemax(Int64) + max_buffer_bytes::Int64 = typemax(Int64) + max_total_allocated_bytes::Int64 = _physicalmemory() + max_allocation_ratio::Int64 = 64 + 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 + max_new_symbols::Int = _DEFAULT_MAX_NEW_SYMBOLS +end + +_physicalmemory() = Int64(min(Sys.total_memory(), UInt64(typemax(Int64)))) + +# `left` is charged and refunded with compare-and-swap, so decoding tasks +# share a budget without a lock. `limit` grows as decompressed input arrives +# (`ratio` bytes per byte, never past `cap`). +# A read's budget also carries its `Limits.max_new_symbols`. +mutable struct AllocationBudget + @atomic left::Int64 + @atomic limit::Int64 + const ratio::Int64 + const cap::Int64 + const maxnewsymbols::Int +end + +function AllocationBudget( + limit::Integer; + ratio::Integer=0, + cap::Integer=limit, + maxnewsymbols::Integer=_DEFAULT_MAX_NEW_SYMBOLS, +) + value = Int64(limit) + value >= 0 || throw(ArgumentError("allocation budget must be nonnegative")) + return AllocationBudget( + value, + value, + Int64(ratio), + max(value, Int64(cap)), + maxnewsymbols, + ) +end + +# Below this, a read's budget does not shrink with its input: small files keep +# the allowance every input had before the budget scaled with input size. +const _MIN_INPUT_BUDGET = Int64(256) * 1024 * 1024 + +_saturatingmul(a::Int64, b::Int64) = + (p=Base.Checked.mul_with_overflow(a, b); p[2] ? typemax(Int64) : p[1]) + +"The allocation budget for one read of `nbytes` input bytes (see [`Limits`](@ref))." +function _readlimit(limits::Limits, nbytes::Integer) + perinput = _saturatingmul(limits.max_allocation_ratio, Int64(nbytes)) + return min(limits.max_total_allocated_bytes, max(_MIN_INPUT_BUDGET, perinput)) +end + +_readbudget(limits::Limits, nbytes::Integer) = AllocationBudget( + _readlimit(limits, nbytes); + ratio=limits.max_allocation_ratio, + cap=limits.max_total_allocated_bytes, + maxnewsymbols=limits.max_new_symbols, +) + +# Move `limit` to `new` and shift `left` by the same amount (not below 0). +function _setlimit!(budget::AllocationBudget, old::Int64, new::Int64) + _, ok = @atomicreplace budget.limit old => new + ok || return false + left = @atomic budget.left + while true + left, done = @atomicreplace budget.left left => max(Int64(0), left + (new - old)) + done && return true + end +end + +"Lower an open read's budget to its input's allowance once the input size is known." +function _boundbudget!(budget::AllocationBudget, limits::Limits, nbytes::Integer) + limit = _readlimit(limits, nbytes) + while true + old = @atomic budget.limit + (limit >= old || _setlimit!(budget, old, limit)) && return budget + end +end + +"Raise the allowance as if `extra` more input bytes had been read, up to the cap." +function _grow!(budget::AllocationBudget, extra::Int64) + (budget.ratio > 0 && extra > 0) || return nothing + add = _saturatingmul(budget.ratio, extra) + while true + old = @atomic budget.limit + new = min(budget.cap, old > typemax(Int64) - add ? typemax(Int64) : old + add) + (new <= old || _setlimit!(budget, old, new)) && return nothing + end +end + +""" + Arrow.AllocationLimitError + +Thrown when a read needs more memory than its allocation budget: `max_allocation_ratio` +bytes per byte of input, at least 256 MiB and at most `max_total_allocated_bytes` +([`Arrow.Limits`](@ref)). The message names the fields to raise for trusted input. +""" +struct AllocationLimitError <: Exception + msg::String +end + +function Base.showerror(io::IO, e::AllocationLimitError) + print( + io, + e.msg, + " (Limits.max_allocation_ratio bytes per input byte, at most ", + "Limits.max_total_allocated_bytes; to read trusted input that needs more, pass ", + "`limits=Arrow.Limits(max_allocation_ratio=..., max_total_allocated_bytes=...)`)", + ) + return nothing +end + +function _charge!(budget::AllocationBudget, amount::Int64, what::AbstractString) + amount >= 0 || throw(ArgumentError("negative allocation charge")) + left = @atomic budget.left + while true + amount <= left || + throw(AllocationLimitError("$what exceeds the reader allocation budget")) + left, ok = @atomicreplace budget.left left => left - amount + ok && return nothing + end +end + +function _refund!(budget::AllocationBudget, amount::Int64) + amount >= 0 || throw(ArgumentError("negative allocation refund")) + left = @atomic budget.left + while true + refunded = Base.Checked.checked_add(left, amount) + refunded <= (@atomic budget.limit) || + throw(ArgumentError("allocation refund exceeds the charged total")) + left, ok = @atomicreplace budget.left left => refunded + ok && return nothing + end +end + +_remaining(budget::AllocationBudget) = @atomic budget.left + +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, +) + result = verify_ipc_metadata(metabytes, limits, _remaining(budget)) + _charge!(budget, result[4], "verified metadata expansion") + return result +end + +function _verify_footer_budgeted( + footerbytes::Vector{UInt8}, + limits::Limits, + budget::AllocationBudget, +) + result = verify_footer(footerbytes, limits, _remaining(budget)) + _charge!(budget, result[5], "verified footer expansion") + return result +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 + +# A vector allocated once at its final length (`Vector{T}(undef, n)`). +# Beyond Julia's small-object size classes its storage is the payload plus +# a header, so large columns are charged what they cost. +@inline function _exactbytes(::Type{T}, n::Integer) where {T} + payload = AC.checked_mul(Int64(n), Int64(Base.elsize(Vector{T}))) + Base.isbitsunion(T) && (payload = AC.checked_add(payload, Int64(n))) + return payload <= 4096 ? AC._materializedvectorbytes(T, n) : + AC.checked_add(payload, Int64(128)) +end + +@inline function _chargeexact!( + budget::Union{Nothing,AllocationBudget}, + ::Type{T}, + n::Integer, + what::AbstractString, +) where {T} + budget === nothing || _charge!(budget, _exactbytes(T, 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.VerifyLimitError && ( + e.ceiling === :depth ? + _limiterror("metadata nesting", "max_nesting_depth", e.limit) : + _limiterror("metadata object count", "max_metadata_objects", e.limit) + ) + 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, _readbudget(limits, region.len)) + +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_allocation_ratio >= 0 || throw(ArgumentError("negative allocation ratio")) + 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")) + limits.max_new_symbols >= 0 || throw(ArgumentError("negative new-Symbol limit")) + return nothing +end + +function _framemessages( + region::OwnerRegion, + limits::Limits, + host_endian_bom::UInt32, + budget::AllocationBudget=_readbudget(limits, region.len), +) + # 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 true + fm, pos = _framemessage(blob, pos, limits, budget, length(msgs)) + fm === nothing && return msgs + push!(msgs, fm) + end +end + +# The error for a message prefix without the continuation marker: at the +# start of the input, the input is not Arrow IPC data (or is a stream from +# before Arrow 0.15, whose messages start with their metadata length). +# An IPC stream of `n` bytes that holds no schema message. +_emptystream(n) = + n == 0 ? "the input is empty: Arrow IPC data has at least a schema message" : + "empty IPC stream: it ends before its schema message" + +function _nomarker(pos::Int64, word::UInt32) + pos == 0 || return "missing continuation marker at byte $pos" + msg = + "not Arrow IPC data: the input starts with neither the file magic " * + "ARROW1 nor an IPC stream message" + # "PAR1", read as a little-endian word. + word == 0x31524150 && return msg * " (it starts with PAR1, as a Parquet file does)" + word & 0xffff == 0x8b1f && + return msg * " (it starts with 1f 8b, as gzip data does; decompress it first)" + 0 < (word % Int32) <= 1 << 24 || return msg + return msg * + " (a stream written before Arrow 0.15 starts like this; " * + "rewrite it with a newer Arrow writer)" +end + +""" +Frame the message at 0-based `pos` of `blob`, the `nread + 1`th of its stream: +returns it and the position after it, or `nothing` at the end-of-stream +marker or the end of the bytes. +""" +function _framemessage( + blob::BufferSlice, + pos::Int64, + limits::Limits, + budget::AllocationBudget, + nread::Int, +) + pos == blob.len && return nothing, pos + blob.len - pos >= 8 || throw( + ValidationError( + pos == 0 ? "not Arrow IPC data: the input is $(blob.len) bytes" : + "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(_nomarker(pos, cont))) + # 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; for IPC streams written " * + "one after another, an Arrow.Stream of an IO reads one stream and " * + "leaves the IO at the next", + ), + ) + return nothing, blob.len + end + _checkmessagecount(nread, limits) + _checkmetalen(metalen, limits) + metastart = AC.checked_add(pos, Int64(8)) + bodystart = AC.checked_add(metastart, metalen) + bodystart <= 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)) + msg, version, header_type, features, bodylen = + _verifiedmessage(metabytes, limits, budget) + bodyend = AC.checked_add(bodystart, bodylen) + bodyend <= blob.len || + throw(ValidationError("truncated body: need $bodylen bytes at $bodystart")) + fm = FramedMessage( + msg, + AC.subslice(blob, bodystart, bodylen), + version, + header_type, + features, + ) + return fm, bodyend +end + +# `what` (a declared size) passes `Limits` field `field`, whose value is `limit`. +@noinline _limiterror(what::AbstractString, field::String, limit) = throw( + ValidationError( + "$what exceeds Arrow.Limits $field ($limit); raise it with " * + "limits=Arrow.Limits($field=...)", + ), +) + +# A footer block's declared metadata and body lengths. +function _checkblocklengths(metalen::Int64, bodylen::Int64, limits::Limits) + metalen > 0 || throw(ValidationError("metadata length $metalen is not positive")) + metalen <= limits.max_metadata_bytes || _limiterror( + "metadata length $metalen", + "max_metadata_bytes", + limits.max_metadata_bytes, + ) + bodylen >= 0 || throw(ValidationError("body length $bodylen is negative")) + bodylen <= limits.max_body_bytes || + _limiterror("body length $bodylen", "max_body_bytes", limits.max_body_bytes) + return nothing +end + +_checkmessagecount(nread::Int, limits::Limits) = + nread < limits.max_messages || + _limiterror("message count", "max_messages", limits.max_messages) + +function _checkmetalen(metalen::Int64, limits::Limits) + metalen > 0 || throw(ValidationError("metadata length $metalen is not positive")) + metalen <= limits.max_metadata_bytes || _limiterror( + "metadata length $metalen", + "max_metadata_bytes", + limits.max_metadata_bytes, + ) + metalen % 8 == 0 || + throw(ValidationError("metadata length $metalen is not 8-byte aligned")) + return nothing +end + +# Verify a message's metadata, then read it and its body length. +function _verifiedmessage( + metabytes::Vector{UInt8}, + limits::Limits, + budget::AllocationBudget, +) + 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) + bodylen >= 0 || throw(ValidationError("body length $bodylen is negative")) + bodylen <= limits.max_body_bytes || + _limiterror("body length $bodylen", "max_body_bytes", limits.max_body_bytes) + bodylen % 8 == 0 || throw(ValidationError("body length $bodylen is not 8-byte aligned")) + return msg, version, header_type, features, bodylen +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 = AC.layoutspec_of(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 + typeids = (f.type::UnionType).typeids + length(typeids) == length(f.children) || + throw(ValidationError("union type-id count must equal child count")) + length(unique(typeids)) == length(typeids) || + throw(ValidationError("union type ids must be unique")) + all(>=(0), 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 + +# --- task parallelism ------------------------------------------------------- + +function _checkntasks(ntasks) + (ntasks isa Integer && ntasks >= 1) || + throw(ArgumentError("ntasks must be a positive integer, got $(repr(ntasks))")) + return Int(ntasks) +end + +# The next item `_foreach` hands out and the smallest item that failed. +mutable struct _Claims + @atomic next::Int + @atomic failed::Int +end + +""" +Call `f(i, w)` for each `i` in `1:n` on at most `ntasks` tasks, the caller +being one of them; `w` in `1:min(n, ntasks)` names the task, so `f` can use +per-task scratch state. Items are claimed in increasing order. When items +throw, items after the smallest failing one are skipped and its exception is +rethrown once every task has stopped, so errors do not depend on `ntasks`. +""" +function _foreach(f::F, n::Int, ntasks::Int) where {F} + nt = min(n, ntasks) + if nt <= 1 + for i = 1:n + f(i, 1) + end + return nothing + end + claims = _Claims(0, typemax(Int)) + errors = Vector{Any}(nothing, n) + # The tasks reach `f` (and whatever it captures, such as a read's decoded + # batches) only through `job`, which is cleared once they are done: a + # finished task can stay reachable after this call (an idle thread keeps + # the last task it ran), and must not keep that data alive. + job = Ref{Union{Nothing,F}}(f) + function work(w) + g = job[]::F + while true + i = @atomic claims.next += 1 + (i > n || i > @atomic(claims.failed)) && return nothing + try + g(i, w) + catch e + errors[i] = (e, catch_backtrace()) + old = @atomic claims.failed + while i < old + old, ok = @atomicreplace claims.failed old => i + ok && break + end + end + end + end + tasks = Task[errormonitor(Threads.@spawn work(w)) for w = 2:nt] + try + work(1) + finally + foreach(wait, tasks) + job[] = nothing + end + failed = @atomic claims.failed + failed == typemax(Int) && return nothing + e, bt = errors[failed] + # The caller sees the worker's exception, with the worker's backtrace as + # its cause. + try + throw(CapturedException(e, bt)) + catch + throw(e) + end +end + +# Decode into `out[1:declared]` and return the bytes produced. With +# `exact=false` the frame may hold more (decoding stops once `out` is full) +# or less; the caller checks the count. +function _decode_lz4!( + state::DecodeState, + src::Ptr{UInt8}, + srclen::Int64, + out::Vector{UInt8}, + declared::Int64, + exact::Bool=true, +) + 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[]) + exact || outpos < declared || return outpos + if hint == 0 + exact || return outpos + 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 outpos + end + inpos < srclen || + (exact ? throw(ValidationError("truncated LZ4 frame")) : return outpos) + (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 Int64(got) +end + +# Decode at most `length(out)` bytes of a zstd buffer and return how many +# were produced; the rest of the input is not read. +function _decode_zstd_prefix!( + state::DecodeState, + src::Ptr{UInt8}, + srclen::Int64, + out::Vector{UInt8}, +) + dctx = _zstdctx!(state) + ZSTD.ZSTD_DCtx_reset(dctx, ZSTD.ZSTD_reset_session_only) + input = ZSTD.ZSTD_inBuffer(Ptr{Cvoid}(src), Csize_t(srclen), Csize_t(0)) + output = ZSTD.ZSTD_outBuffer(Ptr{Cvoid}(pointer(out)), Csize_t(length(out)), Csize_t(0)) + while output.pos < length(out) + before = (input.pos, output.pos) + r = ZSTD.ZSTD_decompressStream(dctx, output, input) + ZSTD.ZSTD_isError(r) == 0 || throw( + ValidationError( + "ZSTD decompression failed: $(unsafe_string(ZSTD.ZSTD_getErrorName(r)))", + ), + ) + (r == 0 || (input.pos, output.pos) == before) && break + end + return Int64(output.pos) +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. +# Concretely typed vectors: node and buffer reads run once per field per batch. +mutable struct DecodeCursor{ + B, + N<:AbstractVector{Meta.FieldNode}, + U<:AbstractVector{Meta.Buffer}, + V<:AbstractVector{<:Integer}, +} + nodes::N + buffers::U + 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::V + 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 + n.length >= 0 || throw(ValidationError("field-node length $(n.length) is negative")) + n.length <= c.max_array_length || + _limiterror("field-node length $(n.length)", "max_array_length", c.max_array_length) + 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")) + len >= 0 || throw(ValidationError("batch buffer length $len is negative")) + len <= c.max_buffer_bytes || + _limiterror("batch buffer length $len", "max_buffer_bytes", c.max_buffer_bytes) + if len > 0 + offset >= c.last_nonempty_end || + throw(ValidationError("batch buffers overlap or move backwards")) + c.last_nonempty_end = _checkedadd(offset, len, "batch buffer end") + end + return offset, len +end + +# Checked arithmetic on metadata values: an overflow is a `ValidationError` +# naming `what`. +function _checkedadd(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 _checkedmul(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 + +skipbuffer!(c::DecodeCursor) = (_buffermeta!(c); nothing) + +# `needed` is the byte count the buffer's column reads (-1 when unknown); it +# bounds how far decompressed bytes extend the read's allowance. +function takebuffer!(c::DecodeCursor, needed::Int64=Int64(-1)) + 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, needed) +end + +# Advise the kernel to map now the pages holding `n` bytes at `p`, which a +# reader is about to read once each: one request instead of a fault per page +# of a memory-mapped file (Unix). A range under 64 KiB spans a few pages at +# most, and advising it costs more than it saves. +function _willneed(p::Ptr{UInt8}, n::Int) + @static if Sys.isunix() + n >= 1 << 16 || return nothing + lo = UInt(p) & ~(UInt(AC.Mmap.PAGESIZE) - 1) + len = UInt(p) + UInt(n) - lo + ccall( + :madvise, + Cint, + (Ptr{Cvoid}, Csize_t, Cint), + Ptr{Cvoid}(lo), + len, + AC.Mmap.MADV_WILLNEED, + ) + end + return nothing +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. The declared size is attacker-controlled: it must agree with +the size the frame headers record, if any, and is charged to the read's budget +before allocation. A buffer that declares more than its column reads +(`needed`, from the field node; a writer may store a sliced column's whole +buffer) decodes only the column's bytes; otherwise the decompressed size must +match the declaration exactly. A positively compressed buffer becomes its own +exact-sized owned region. + +Decompressed bytes count as input to the budget as far as the column reads +them, so a highly compressible buffer is funded while a claim beyond what the +column reads is not. +""" +function _decompressbuffer!(c::DecodeCursor, wire::BufferSlice, needed::Int64=Int64(-1)) + 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 + declared >= 0 || + throw(ValidationError("declared decompressed length $declared is negative")) + declared <= c.max_buffer_bytes || _limiterror( + "declared decompressed length $declared", + "max_buffer_bytes", + c.max_buffer_bytes, + ) + declared <= typemax(Int) >> 1 || + 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 + wireregion = wire.region::OwnerRegion + recorded = GC.@preserve wireregion _framecontentsize( + c.codec, + AC.sliceptr(wire) + 8, + payloadlen, + ) + (recorded < 0 || recorded == declared) || throw( + ValidationError( + "declared decompressed length $declared does not match the $recorded bytes " * + "its frame records", + ), + ) + # A buffer longer than its column reads (a writer may store a sliced + # column's whole buffer) decodes only the column's bytes. + exact = needed < 0 || declared <= _padded64(needed) + outlen = exact ? declared : _padded64(needed) + # No payload decodes to more than its codec's largest expansion: LZ4 + # turns one byte into at most 255, a 4-byte zstd RLE block into at most + # 128 KiB. A larger claim is refused before it is allocated. + expansion = c.codec == CODEC_LZ4_FRAME ? Int64(255) : Int64(32768) + outlen <= AC.checked_mul(expansion, payloadlen) || throw( + ValidationError( + "declared decompressed length $outlen exceeds what $payloadlen " * + "compressed bytes can hold", + ), + ) + needed >= 0 && _grow!(state.budget, outlen - payloadlen) + charged = _exactbytes(UInt8, outlen) + _charge!(state.budget, charged, "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(outlen)) + # The native decoders read through a raw pointer, so the wire + # region's root must stay reachable for the whole call (Core rule 2). + got = GC.@preserve out wireregion begin + src = AC.sliceptr(wire) + 8 + _willneed(src, Int(payloadlen)) + if c.codec == CODEC_LZ4_FRAME + _decode_lz4!(state, src, payloadlen, out, outlen, exact) + elseif exact + _decode_zstd!(state, src, payloadlen, out, outlen) + else + _decode_zstd_prefix!(state, src, payloadlen, out) + end + end + exact || + got >= needed || + throw( + ValidationError( + "compressed buffer holds $got bytes; its column needs $needed", + ), + ) + region = OwnerRegion(Ptr{UInt8}(pointer(out)), sizeof(out); root=_Decompressed(out)) + result = BufferSlice(region, 0, got) + 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, charged) + end +end + +_padded64(n::Int64) = n > typemax(Int64) - 63 ? typemax(Int64) : (n + 63) & ~Int64(63) + +# The root of a decompressed buffer's region: bytes this read allocated and no +# other buffer shares, so a column may keep them instead of a copy (string +# data, `_adoptbytes`). +struct _Decompressed + bytes::Vector{UInt8} +end + +# The decompressed size the frame headers record: zstd frames usually do (a +# buffer may hold several), LZ4 frames optionally. -1 when a header does not +# say; a malformed header is left for the decoder to reject. +function _framecontentsize(codec::Int8, src::Ptr{UInt8}, len::Int64)::Int64 + if codec == CODEC_ZSTD + total = Int64(0) + pos = Int64(0) + while pos < len + n = ZSTD.ZSTD_getFrameContentSize(src + pos, Csize_t(len - pos)) + n >= ZSTD.ZSTD_CONTENTSIZE_ERROR && return Int64(-1) + n > typemax(Int64) - total && return Int64(-1) + total += Int64(n) + framelen = ZSTD.ZSTD_findFrameCompressedSize(src + pos, Csize_t(len - pos)) + (ZSTD.ZSTD_isError(framelen) == 0 && framelen > 0) || return Int64(-1) + pos += Int64(framelen) + end + return total + end + # LZ4 frame: magic, FLG (bit 3: content size present), BD, then the size. + ( + len >= 14 && + unsafe_load(src, 1) == 0x04 && + unsafe_load(src, 2) == 0x22 && + unsafe_load(src, 3) == 0x4d && + unsafe_load(src, 4) == 0x18 + ) || return Int64(-1) + unsafe_load(src, 5) & 0x08 == 0 && return Int64(-1) + n = UInt64(0) + for k = 14:-1:7 + n = (n << 8) | unsafe_load(src, k) + end + return n > typemax(Int64) ? Int64(-1) : Int64(n) +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 = AC.layoutspec_of(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. +""" +# The bytes a column of `n` slots reads from a buffer in `role`, from its +# field node alone (-1 when that needs bytes not decoded yet). Variable-width +# data ends at the last offset, which precedes it in `previous`. +function _neededbytes(role::AC.BufferRole, spec::AC.LayoutSpec, n::Int64, previous) + role == AC.VALIDITY && return cld(n, 8) + role == AC.TYPE_IDS && return n + role == AC.VIEWS && return _saturatingmul(n, Int64(16)) + role == AC.OFFSETS && + return n == 0 ? Int64(0) : _saturatingmul(n + 1, Int64(spec.offsetwidth)) + (role == AC.ELEMENT_OFFSETS || role == AC.SIZES) && + return _saturatingmul(n, Int64(spec.offsetwidth)) + role == AC.DATA || return Int64(-1) + spec.fixedwidth > 0 && return _saturatingmul(n, Int64(spec.fixedwidth)) + spec.fixedwidth == -1 && return cld(n, 8) + # Byte-addressed data follows its offsets (validity, offsets, data). + (spec.offsetwidth > 0 && length(previous) == 2) || return Int64(-1) + n == 0 && return Int64(0) + offsets = previous[2] + w = spec.offsetwidth + offsets.len >= (n + 1) * w || return Int64(-1) + last = + w == 8 ? AC.loadat(offsets, Int64, n * 8) : Int64(AC.loadat(offsets, Int32, n * 4)) + return last >= 0 ? last : Int64(-1) +end + +# The end of the last value each variadic data buffer of a view column holds, +# from its views: what a compressed data buffer's column reads. `nothing` for +# an uncompressed batch, whose buffers are not decompressed. +function _variadicends(c::DecodeCursor, views::BufferSlice, n::Int64, nvariadic::Int) + c.codec == CODEC_NONE && return nothing + _chargevector!((c.state::DecodeState).budget, Int64, nvariadic, "variadic buffer ends") + ends = fill(Int64(-1), nvariadic) + (nvariadic == 0 || views.len < 16n) && return ends + fill!(ends, 0) + for i = 0:(n - 1) + len = AC.loadat(views, Int32, 16i) + len > 12 || continue + k = AC.loadat(views, Int32, 16i + 8) + off = AC.loadat(views, Int32, 16i + 12) + (0 <= k < nvariadic && off >= 0) || continue + ends[k + 1] = max(ends[k + 1], Int64(off) + len) + end + return ends +end + +function decodefield( + f::Field, + c::DecodeCursor, + dicts::Dict{Int64,ArrayData}, + fielddictids::IdDict{Field,Int64}, +) + t = f.type + node = takenode!(c) + spec = AC.layoutspec_of(t) + n = Int64(node.length) + buffers = BufferSlice[] + for role in spec.buffers + push!(buffers, takebuffer!(c, _neededbytes(role, spec, n, buffers))) + end + if spec.variadic + # View layouts append their declared count of variadic data buffers + # after the fixed validity/views pair (format 1.4). + nvariadic = takevariadic!(c) + ends = _variadicends(c, buffers[end], n, nvariadic) + for k = 1:nvariadic + push!(buffers, takebuffer!(c, ends === nothing ? Int64(-1) : ends[k])) + 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 + +""" +The sequential half of decoding a record batch: its codec, its length, and +where each top-level field's nodes and buffers start. A metadata-only walk +(`skipfield!`) checks the whole node and buffer table first; the fields then +decode independently (`_decodeplanned`), in any order or concurrently. +""" +struct _RecordPlan{B} + header::Meta.RecordBatch + body::B + codec::Int8 + variadics::Vector{Int64} + rblen::Int64 + # Cursor (node, buffer, variadic count, last buffer end) at each field. + starts::Vector{NTuple{4,Int64}} +end + +function _recordplan(header::Meta.RecordBatch, body, version::Int16, fields, limits::Limits) + codec = _batchcodec(header.compression, version) + rblen = something(header.length, Int64(0)) + rblen >= 0 || throw(ValidationError("record batch length $rblen is negative")) + rblen <= limits.max_array_length || _limiterror( + "record batch length $rblen", + "max_array_length", + limits.max_array_length, + ) + variadics = variadiccounts(header) + c = DecodeCursor( + header.nodes, + header.buffers, + body, + limits; + codec=codec, + variadics=variadics, + ) + starts = Vector{NTuple{4,Int64}}(undef, length(fields)) + for (j, f) in enumerate(fields) + starts[j] = (c.nodeidx, c.bufidx, c.varidx, c.last_nonempty_end) + skipfield!(f, c) + end + finishcursor!(c) + return _RecordPlan(header, body, codec, variadics, rblen, starts) +end + +"Decode and validate field `j` of a planned record batch." +function _decodeplanned( + p::_RecordPlan, + j::Int, + f::Field, + dicts::Dict{Int64,ArrayData}, + fielddictids::IdDict{Field,Int64}, + limits::Limits, + validated_dictionaries, + state::DecodeState, +) + c = DecodeCursor( + p.header.nodes, + p.header.buffers, + p.body, + limits; + codec=p.codec, + state=state, + variadics=p.variadics, + ) + c.nodeidx, c.bufidx, c.varidx, c.last_nonempty_end = p.starts[j] + col = decodefield(f, c, dicts, fielddictids) + # One IPC dictionary id can back many fields. Each immutable pool + # snapshot was certified at its DictionaryBatch, so semantic validation + # skips that pool tree; index contracts still run for every field. + AC._validate_semantic(f, col, validated_dictionaries) + col.len == p.rblen || + throw(ValidationError("RecordBatch length does not match top-level field nodes")) + return col +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) + dictionaries::Dict{Int64,ArrayData} # final pools, including updates after the last record + budget::AllocationBudget # framing, decode, and facade materialization + limits::Limits # complete policy used when this handle opened +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, _readbudget(limits, length(bytes))) + +""" +The schema of an IPC stream from its first message: the Core schema and +fields, the dictionary ids, and each dictionary's value field. +""" +function _streamschema(fm::FramedMessage) + fm.header_type == 1 || # Schema + throw(ValidationError("first IPC message must be a schema")) + fm.msg.header isa Meta.Schema || + throw(ValidationError("first IPC message must be a schema")) + metaschema = fm.msg.header + fm.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, + ) + return sch, fields, dictids, fielddictids, dictvaluefields +end + +""" +The decoding state of an IPC stream after its schema: the dictionaries as of +the last message applied and the stream's metadata version. +""" +struct _StreamState + schema::Schema + fields::Vector{Field} + dictids::Dict{Int64,Meta.Field} + fielddictids::IdDict{Field,Int64} + dictvaluefields::Dict{Int64,Field} + dicts::Dict{Int64,ArrayData} + validated::AC._ValidatedDictionaries + version::Int16 + limits::Limits +end + +function _StreamState(fm::FramedMessage, limits::Limits) + sch, fields, dictids, fielddictids, dictvaluefields = _streamschema(fm) + return _StreamState( + sch, + fields, + dictids, + fielddictids, + dictvaluefields, + Dict{Int64,ArrayData}(), + AC._ValidatedDictionaries(), + fm.version, + limits, + ) +end + +# A message after the schema: its version, and the legacy V4 compression +# marker (see rejectexperimentalcompression). +function _checkstreammessage(st::_StreamState, fm::FramedMessage) + fm.version == st.version || + throw(ValidationError("IPC metadata version changes within the stream")) + rejectexperimentalcompression(fm) + return nothing +end + +""" +The values of dictionary batch `rb` (its RecordBatch, under message version +`version`), decoded from `body` as field `vf`. +""" +function _decodedictionary( + rb::Meta.RecordBatch, + body, + version::Int16, + vf::Field, + dicts, + fielddictids, + limits::Limits, + state::DecodeState, +) + rblen = something(rb.length, Int64(0)) + rblen >= 0 || throw(ValidationError("dictionary batch length $rblen is negative")) + rblen <= limits.max_array_length || _limiterror( + "dictionary batch length $rblen", + "max_array_length", + limits.max_array_length, + ) + cursor = DecodeCursor( + rb.nodes, + rb.buffers, + body, + limits; + codec=_batchcodec(rb.compression, version), + 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"), + ) + return decoded +end + +""" +Decode dictionary batch `fm` and apply it to the stream's dictionaries (a +first pool, a delta, or a replacement). Returns its id, the pool now in +effect, and whether the id already had one. +""" +function _applydictionary!( + st::_StreamState, + fm::FramedMessage, + budget::AllocationBudget, + state::DecodeState, +) + header = fm.msg.header::Meta.DictionaryBatch + rb = header.data + haskey(st.dictids, header.id) || + throw(ValidationError("dictionary batch has unknown id $(header.id)")) + replacement = haskey(st.dicts, header.id) + _dictionarytransition(header.id, header.isDelta, replacement) + # 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(st.dictvaluefields, header.id) || + throw(ValidationError("dictionary batch has unknown id $(header.id)")) + vf = st.dictvaluefields[header.id] + decoded = _decodedictionary( + rb, + fm.body, + fm.version, + vf, + st.dicts, + st.fielddictids, + st.limits, + state, + ) + decoded = _updatedictionary!( + st.dicts, + st.validated, + header.id, + header.isDelta, + vf, + decoded, + st.limits, + budget, + ) + return header.id, decoded, replacement +end + +""" +A record batch of a stream, planned against the dictionaries in effect at its +position. The IPC format lets a batch whose dictionary column is all null +arrive before that dictionary; such a batch waits (`missing` is not empty) +until the dictionary arrives, and then decodes with it. +""" +mutable struct PendingRecord + plan::_RecordPlan{BufferSlice} + dictionaries::Dict{Int64,ArrayData} + missing::Set{Int64} + budget::AllocationBudget +end + +function _pendingrecord(st::_StreamState, fm::FramedMessage, budget::AllocationBudget) + header = fm.msg.header::Meta.RecordBatch + missing = missingdicts(st.fields, header.nodes, st.dicts, st.fielddictids) + plan = _recordplan(header, fm.body, fm.version, st.fields, st.limits) + return PendingRecord(plan, copy(st.dicts), missing, budget) +end + +# A dictionary that arrived resolves the records waiting for it. +function _resolvepending!(queue, id::Int64, pool::ArrayData) + for p in queue + p isa PendingRecord && id in p.missing || continue + p.dictionaries[id] = pool + delete!(p.missing, id) + end + return queue +end + +""" +Decode the fields of planned records `ps` (every one resolved) on at most +`ntasks` tasks, one (record, field) pair at a time. +""" +function _decoderecords(st::_StreamState, ps::Vector{PendingRecord}, ntasks::Int) + isempty(ps) && return AC.RecordBatch[] + nf = length(st.fields) + cols = [Vector{ArrayData}(undef, nf) for _ in ps] + # One codec state per task, charging the budget of the record at hand. + states = [DecodeState(first(ps).budget) for _ = 1:min(ntasks, length(ps) * nf)] + try + _foreach(length(ps) * nf, ntasks) do item, w + k, j = divrem(item - 1, nf) .+ (1, 1) + p = ps[k] + state = states[w] + state.budget = p.budget + cols[k][j] = _decodeplanned( + p.plan, + j, + st.fields[j], + p.dictionaries, + st.fielddictids, + st.limits, + st.validated, + state, + ) + return nothing + end + finally + foreach(close, states) + end + return AC.RecordBatch[ + AC.RecordBatch(st.schema, cols[k], p.plan.rblen, st.validated) for + (k, p) in enumerate(ps) + ] +end + +""" +The framed stream in `region`: its decoding state after every message, and its +record batches planned against the dictionaries in effect at each. +""" +function _streamrecords(region::OwnerRegion, limits::Limits, budget::AllocationBudget) + msgs = _framemessages(region, limits, Base.ENDIAN_BOM, budget) + isempty(msgs) && throw(ValidationError(_emptystream(region.len))) + st = _StreamState(msgs[1], limits) + # One codec context per reader for the dictionary batches, which apply in + # order; record batches are planned in order and decoded together after. + state = DecodeState(budget) + records = PendingRecord[] + try + for fm in msgs[2:end] + _checkstreammessage(st, fm) + header = fm.msg.header + if header isa Meta.DictionaryBatch + id, pool, replaced = _applydictionary!(st, fm, budget, state) + # Waiting records take an id's first pool only; a later + # replacement is for the records after it. + replaced || _resolvepending!(records, id, pool) + elseif header isa Meta.RecordBatch + push!(records, _pendingrecord(st, fm, budget)) + else + throw(ValidationError("unsupported IPC message header $(typeof(header))")) + end + end + finally + close(state) + end + all(p -> isempty(p.missing), records) || + throw(ValidationError("stream ended before required dictionary batches arrived")) + return st, records +end + +function _readstream( + bytes::Vector{UInt8}, + limits::Limits, + budget::AllocationBudget, + ntasks::Int=1, +) + region = heapregion(bytes) + st, records = _streamrecords(region, limits, budget) + return IPCStream( + st.schema, + AC.FrozenVector{Field}(st.fields), + _decoderecords(st, records, ntasks), + region, + 1, + false, + st.fielddictids, + st.dicts, + budget, + limits, + ) +end + +# --- incremental streams ---------------------------------------------------- + +""" +The messages of an IPC stream, framed one at a time from a byte region (a +byte vector or a memory map) as the reader reaches them. +""" +mutable struct _RegionMessages + const blob::BufferSlice + pos::Int64 +end +_RegionMessages(region::OwnerRegion) = + _RegionMessages(BufferSlice(region, 0, region.len), 0) + +function _nextmessage!(m::_RegionMessages, limits::Limits, budget::AllocationBudget) + fm, m.pos = _framemessage(m.blob, m.pos, limits, budget, 0) + return fm +end + +""" +The messages of an IPC stream read from an `IO` one at a time: each message's +metadata and body are read when the reader reaches them, so a pipe or socket +streams with one message in memory. `head` holds bytes already read from the +`IO` (while telling the formats apart). +""" +mutable struct _IOMessages{I<:IO} + const io::I + head::Vector{UInt8} + pos::Int64 +end + +# Read `n` bytes into `out` from index `at` on; false when the stream ends first. +function _readinto!(m::_IOMessages, out::Vector{UInt8}, at::Int, n::Int) + k = min(length(m.head), n) + copyto!(out, at, m.head, 1, k) + deleteat!(m.head, 1:k) + if k < n + try + GC.@preserve out unsafe_read(m.io, pointer(out, at + k), UInt(n - k)) + catch e + e isa EOFError || rethrow() + return false + end + end + m.pos += n + return true +end + +function _readbytes!(m::_IOMessages, n::Int64) + out = Vector{UInt8}(undef, Int(n)) + return _readinto!(m, out, 1, Int(n)) ? out : nothing +end + +# A body is read in pieces of growing size, so a declared length the stream +# does not back costs only the bytes that arrive. Those bytes are the input +# itself, so they are not charged to the reader's allowance. +const _IO_BODY_PIECE = 64 * 1024 * 1024 + +function _readbody!(m::_IOMessages, bodylen::Int64) + body = UInt8[] + while length(body) < bodylen + have = length(body) + n = Int(min(bodylen - have, max(_IO_BODY_PIECE, have))) + resize!(body, have + n) + _readinto!(m, body, have + 1, n) || + throw(ValidationError("truncated body: need $bodylen bytes at byte $(m.pos)")) + end + return body +end + +function _nextmessage!(m::_IOMessages, limits::Limits, budget::AllocationBudget) + at = m.pos + isempty(m.head) && eof(m.io) && return nothing # the stream ended between messages + prefix = _readbytes!(m, Int64(8)) + prefix === nothing && throw(ValidationError("truncated IPC prefix at byte $at")) + word = reinterpret(UInt32, prefix)[1] + word == CONTINUATION || throw(ValidationError(_nomarker(at, word))) + metalen = Int64(ltoh(reinterpret(Int32, prefix)[2])) + metalen == 0 && return nothing # end-of-stream marker + _checkmetalen(metalen, limits) + _chargevector!(budget, UInt8, metalen, "metadata allocation") + metabytes = _readbytes!(m, metalen) + metabytes === nothing && + throw(ValidationError("truncated metadata: need $metalen bytes at $at")) + msg, version, header_type, features, bodylen = + _verifiedmessage(metabytes, limits, budget) + body = _readbody!(m, bodylen) + return FramedMessage( + msg, + BufferSlice(heapregion(body), 0, length(body)), + version, + header_type, + features, + ) +end + +# The budget a stream message frames under; its batch then decodes under one +# sized from the message's own bytes. +_framebudget(limits::Limits) = _readbudget(limits, 0) + +""" +An IPC stream read incrementally: messages are read and decoded as record +batches are requested (`_nextrecord!`), each under its own allowance, so a +stream of any length reads with one batch in memory. A stream over a byte +region can start over; one over an `IO` reads once. +""" +mutable struct _StreamReader{M} + const messages::M + const start::Int64 # position after the schema + # The lifetime authority of the stream and its batches: the byte region, + # or for an `IO` (each body is its own region) an empty one. + const region::OwnerRegion + st::_StreamState + missing::Set{Int64} # dictionaries promised by all-null batches + done::Bool + failed::Bool # a read threw; the stream stops there + const limits::Limits + const lock::ReentrantLock +end + +function _StreamReader(messages, limits::Limits) + _validatelimits(limits) + fm = _nextmessage!(messages, limits, _framebudget(limits)) + fm === nothing && throw(ValidationError(_emptystream(messages.pos))) + st = _StreamState(fm, limits) + region = messages isa _RegionMessages ? messages.blob.region : heapregion(UInt8[]) + return _StreamReader( + messages, + messages.pos, + region, + st, + Set{Int64}(), + false, + false, + limits, + ReentrantLock(), + ) +end + +# Start over from the first batch (a byte region or a byte-range source). +function _restart!(r::_StreamReader) + # A reader that has not read past its schema is at the start already. + r.messages.pos == r.start && isempty(r.missing) && !r.done && !r.failed && return r + m = r.messages + if m isa _IOMessages{_SourceIO} + empty!(m.head) + m.io.pos = r.start + elseif !(m isa _RegionMessages) + throw( + ArgumentError( + "an Arrow.Stream over an IO reads its batches once; " * + "read the IO into a byte vector to iterate it again", + ), + ) + end + m.pos = r.start + st = r.st + r.st = _StreamState( + st.schema, + st.fields, + st.dictids, + st.fielddictids, + st.dictvaluefields, + Dict{Int64,ArrayData}(), + AC._ValidatedDictionaries(), + st.version, + st.limits, + ) + empty!(r.missing) + r.done = false + r.failed = false + return r +end + +""" +The next record batch of the stream, planned against the dictionaries in +effect for it, and the allowance it decodes under; `nothing` at its end. +Messages are read until the next batch in stream order is decodable. +""" +function _nextrecord!(r::_StreamReader) + r.failed && throw( + ArgumentError( + "this Arrow.Stream stopped at an error; open a new one to read the source again", + ), + ) + try + return _readrecord!(r) + catch + r.failed = true + rethrow() + end +end + +function _readrecord!(r::_StreamReader) + while true + if r.done + isempty(r.missing) && return nothing + throw( + ValidationError("stream ended before required dictionary batches arrived"), + ) + end + fm = _nextmessage!(r.messages, r.limits, _framebudget(r.limits)) + if fm === nothing + r.done = true + continue + end + _checkstreammessage(r.st, fm) + budget = _readbudget(r.limits, fm.body.len) + header = fm.msg.header + if header isa Meta.DictionaryBatch + state = DecodeState(budget) + try + id, _, _ = _applydictionary!(r.st, fm, budget, state) + delete!(r.missing, id) + finally + close(state) + end + elseif header isa Meta.RecordBatch + p = _pendingrecord(r.st, fm, budget) + # An all-null batch needs no pool values. Yield it now, even when + # its dictionary has not arrived, without retaining later batches. + union!(r.missing, p.missing) + for id in p.missing + p.dictionaries[id] = _dictionaryconcat( + r.st.dictvaluefields[id], + _DictionaryRange[], + r.limits, + budget, + ) + end + return p, budget + else + throw(ValidationError("unsupported IPC message header $(typeof(header))")) + end + end +end + +"The number of record batches of a region stream, from its framing alone." +function _countrecords(r::_StreamReader{_RegionMessages}) + m = _RegionMessages(r.messages.blob, r.start) + n = 0 + while true + fm = _nextmessage!(m, r.limits, _framebudget(r.limits)) + fm === nothing && return n + fm.msg.header isa Meta.RecordBatch && (n += 1) + end +end + +# --------------------------------------------------------------------------- +# 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 + +""" + 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")) + # Every extent is nonnegative, so these comparisons cannot overflow. + ( + offset <= dataend && + metalen <= dataend - offset && + bodylen <= dataend - offset - metalen + ) || throw(ValidationError("footer block escapes the data section")) + frameend = offset + metalen + bodylen + 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 + +""" +Verify the message footer Block `block` points at, from `meta`, the Block's +prefix and metadata bytes: the continuation marker, the metadata length the +Block declares, the verified metadata, and the body length the Block +declares. Returns the message, its metadata version, header type, and +features. +""" +function _blockmeta( + meta::BufferSlice, + block::NTuple{3,Int64}, + limits::Limits, + budget::AllocationBudget, +) + _, metalen, bodylen = block + AC.loadat(meta, UInt32, Int64(0)) == CONTINUATION || + throw(ValidationError("footer block does not point at a message")) + declared = Int64(AC.loadat(meta, Int32, Int64(4))) + declared == metalen - 8 || + throw(ValidationError("footer block metadata length does not match the message")) + _checkblocklengths(declared, bodylen, limits) + _chargevector!(budget, UInt8, declared, "metadata allocation") + metabytes = AC.slicebytes(AC.subslice(meta, 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 msg, version, header_type, features +end + +"The message footer Block `block` points at in `region`, framed and verified." +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) + msg, version, header_type, features = + _blockmeta(AC.subslice(blob, offset, metalen), block, limits, budget) + 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 its schema, and the frame of every message the footer +indexes; eagerly decode the dictionary blocks (shared by every record batch); +expose record batches lazily through the Block index. The schema comes from +the footer; the leading Schema message is not read. +Pass `mmapregion(path)` to read a file through Core's mmap path. The file +format permits one base dictionary per id, followed by deltas. +Full dictionary replacement remains a format error. + +`Arrow.Table(handle)` and `Arrow.Stream(handle)` copy a region that can be +released eagerly, so their columns remain valid after release. Pass the path +directly to those constructors for memory-mapped views without this copy. +""" +readfile(bytes::Vector{UInt8}; limits::Limits=Limits()) = + _readfile(heapregion(bytes), limits, _readbudget(limits, length(bytes))) + +readfile(region::OwnerRegion; limits::Limits=Limits()) = + _readfile(region, limits, _readbudget(limits, region.len)) + +# C++ WriteFileFooter uses its current metadata version even when the +# writer emits V4 messages, and Arrow.jl before 2.6 wrote a V4 footer over V5 +# messages. The footer must not select the message layout. +function _filemetadataversion(footer::Int16, message::Int16, expected::Int16=message) + message == expected || + throw(ValidationError("IPC metadata version changes within the file")) + ( + footer == message || + (footer in (Int16(3), Int16(4)) && message in (Int16(3), Int16(4))) + ) || throw(ValidationError("IPC message and footer metadata versions differ")) + return message +end + +""" +The verified footer of the IPC file in `src`, a region or a byte-range source: +its schema, dictionary ids, metadata version, and Block index. The Footer is +authoritative, as in Arrow C++: the schema comes from it and every message is +reached through its Blocks, so the leading Schema message is never read. +(Polars writes that message without its stream framing.) Blocks start after +the padded leading magic (`ARROW1` plus two bytes), end before the Footer, +and do not overlap; each message is verified against its Block when it is +read. +""" +function _readfooter(src, limits::Limits, budget::AllocationBudget) + _requirelittleendian() + _validatelimits(limits) + # leading magic + pad (8), smallest footer flatbuffer (8), footer length + # (4), trailing magic (6). + src.len >= Int64(8 + 8 + 4 + 6) || + throw(ValidationError("file is too short to be an IPC file")) + # The trailer is