feat(python): add Arrow string batch tokenization - #2400
Open
granthamtaylor wants to merge 1 commit into
Open
Conversation
Accept Arrow string and large_string arrays through the C Data Interface without materializing Python strings. Preserve character offsets and support configurable null handling with Literal typing. Add FFI lifetime and parity tests plus a Python benchmark matrix for conversion, encoding, throughput, and process peak memory.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Arrow callers currently need to materialize Python strings before calling
encode_batch. This addsTokenizer.encode_batch_arrow, which accepts UTF-8stringandlarge_stringarrays through__arrow_c_array__and returns the existingEncodingobjects. It supports slices and empty arrays and preserves tokenizer configuration and character offsets.The Python binding uses
arrow-arrayandarrow-schemafor C Data Interface ownership and validation. Imported buffers remain alive throughout synchronous parallel encoding, and release callbacks run after reattaching to Python. Non-null text is checked for valid UTF-8; null payload bytes are left uninterpreted. PyArrow is an explicit test dependency, with no new Python runtime dependency.Null values raise an error by default. The keyword-only
null_handling="empty"option encodes them as empty strings, andnull_handling="skip"omits them while preserving retained row order. Generated stubs expose these choices asLiteral["error", "empty", "skip"], with runtime validation retained. Inputs are flat string arrays with offsets aligned for their native type; nested and chunked arrays, sequence pairs, and pre-tokenized inputs are outside this method's scope.Includes API documentation, tests for encoding parity and capsule ownership, and a standalone benchmark. The benchmark compares list conversion and Arrow input in fresh processes, reports timing and absolute process peak RSS separately, and retains raw samples and environment details. Tokenization and output construction can still allocate memory; no general performance improvement is claimed from the synthetic benchmark.
Python throughput measurements cover 64 configurations: batch sizes 1, 16, 256 and 4,096; body lengths 32, 256, 2,048 and 8,192 Unicode characters; both Arrow string widths; and one or four Rayon threads. Each configuration has five paired repetitions, for 640 fresh worker processes.
For 32-character bodies, single-item calls ran at 0.89–0.91x the list baseline. With four threads, batches of 16–4,096 ran at 1.07–1.13x. One-thread cases with 2,048- or 8,192-character bodies were near parity (0.998–1.024x). For example, at 4,096 rows × 32 characters with four threads,
stringthroughput was 551,025 documents/s for the list path and 601,467 documents/s for Arrow, with a median paired ratio of 1.07x. These measurements show workload-dependent gains and small-call overhead.Measured on an Apple M4 Pro (14 cores, 48 GiB), macOS 26.5.2 ARM64, Python 3.13.4, PyArrow 25.0.1, NumPy 2.5.3 and pandas 3.0.5. Both paths used the same release build of this Arrow patch on base
6cfd9d385ca0ed91c10b49f0ce97d02cfde1b607, with Rust 1.98.1 and the default Python binding features. This was a shared development machine without CPU affinity.The fixture uses
WordLevelwithWhitespace, repeating"hello world arrow café 東京 "to the requested length. The four body lengths occupy 37, 302, 2,438 and 9,767 UTF-8 bytes respectively. All rows in a batch have the same body; there are no nulls, padding or truncation. This is a synthetic measurement of repeated calls and does not establish performance for BPE, other text distributions or complete applications.Each fresh worker warms up one full batch, then measures calls until cumulative timed duration reaches at least 0.2 seconds. Timings include
array.to_pylist()plusencode_batchfor the list path, andencode_batch_arrowfor the Arrow path, including Python dispatch, Arrow import/validation, tokenization and returnedEncodingconstruction. Fixture setup, warmup, correctness checks, previous input/output destruction and loop bookkeeping are outside the timers. Workers run sequentially; list/Arrow order alternates across paired repetitions and cases.The following 48 matrices cover all 12 recorded metrics, with List / Arrow values for each of the 64 configurations. Timings and throughput summarize per-worker mean calls; RSS, counts and cumulative duration summarize their corresponding process or batch observations. The setup RSS and workload/call-count matrices provide context, not independent measures of performance.
Medians are calculated independently for each metric: median conversion plus median encoding need not equal median total, although each raw sample's components sum to its total. Arrow's explicit Python conversion time is zero; its native import and validation are included in encoding time. Peak RSS includes imports, fixtures, full-batch warmup and outputs. For example, the 4,096 × 8,192-character
stringcase with four threads reached 1,029.0 MiB for the list path and 830.6 MiB for Arrow. These are whole-process high-water marks, not allocation deltas.Conversion time (ms/call)
Conversion time (ms/call)
Every cell: List / Arrow, median across five workers. Lower is better for this measured stage.
List measures to_pylist(); Arrow is exactly 0 because it does not run this conversion. Arrow import and validation costs are included in encoding time, not zero overall overhead.
Timing medians are independent; their components need not sum after aggregation.
string · 1 Rayon thread
string · 4 Rayon threads
large_string · 1 Rayon thread
large_string · 4 Rayon threads
List = array.to_pylist() + encode_batch(); Arrow = encode_batch_arrow(). Native import/validation belongs to Arrow encoding time. Full precision is retained in the saved CSV and JSON results.
Encoding time (ms/call)
Encoding time (ms/call)
Every cell: List / Arrow, median across five workers. Lower is better for the defined encoding stage.
List excludes its explicit conversion; Arrow includes native import and validation. Use total time to compare the complete measured calls.
Timing medians are independent; their components need not sum after aggregation.
string · 1 Rayon thread
string · 4 Rayon threads
large_string · 1 Rayon thread
large_string · 4 Rayon threads
List = array.to_pylist() + encode_batch(); Arrow = encode_batch_arrow(). Native import/validation belongs to Arrow encoding time. Full precision is retained in the saved CSV and JSON results.
Total call time (ms/call)
Total call time (ms/call)
Every cell: List / Arrow, median across five workers. Lower is better.
Includes explicit conversion where applicable and the call returning Encoding objects.
Timing medians are independent; their components need not sum after aggregation.
string · 1 Rayon thread
string · 4 Rayon threads
large_string · 1 Rayon thread
large_string · 4 Rayon threads
List = array.to_pylist() + encode_batch(); Arrow = encode_batch_arrow(). Native import/validation belongs to Arrow encoding time. Full precision is retained in the saved CSV and JSON results.
Document throughput (documents/s)
Document throughput (documents/s)
Every cell: List / Arrow, median across five workers. Higher is better.
Derived from each worker's mean total call time, then summarized across workers.
Timing medians are independent; their components need not sum after aggregation.
string · 1 Rayon thread
string · 4 Rayon threads
large_string · 1 Rayon thread
large_string · 4 Rayon threads
List = array.to_pylist() + encode_batch(); Arrow = encode_batch_arrow(). Native import/validation belongs to Arrow encoding time. Full precision is retained in the saved CSV and JSON results.
UTF-8 throughput (MiB/s)
UTF-8 throughput (MiB/s)
Every cell: List / Arrow, median across five workers. Higher is better.
Uses actual UTF-8 input bytes and total call time; 1 MiB = 1,048,576 bytes.
Timing medians are independent; their components need not sum after aggregation.
string · 1 Rayon thread
string · 4 Rayon threads
large_string · 1 Rayon thread
large_string · 4 Rayon threads
List = array.to_pylist() + encode_batch(); Arrow = encode_batch_arrow(). Native import/validation belongs to Arrow encoding time. Full precision is retained in the saved CSV and JSON results.
Output token throughput (tokens/s)
Output token throughput (tokens/s)
Every cell: List / Arrow, median across five workers. Higher is better.
Uses returned output token counts and total call time.
Timing medians are independent; their components need not sum after aggregation.
string · 1 Rayon thread
string · 4 Rayon threads
large_string · 1 Rayon thread
large_string · 4 Rayon threads
List = array.to_pylist() + encode_batch(); Arrow = encode_batch_arrow(). Native import/validation belongs to Arrow encoding time. Full precision is retained in the saved CSV and JSON results.
Process peak resident memory (MiB)
Process peak resident memory (MiB)
Every cell: List / Arrow, median across five workers. Lower absolute process peak is preferable.
Median process high-water mark through encoding, including imports, fixture setup, warmup and outputs. This is not isolated tokenizer allocation; do not subtract process high-water marks.
Timing medians are independent; their components need not sum after aggregation.
string · 1 Rayon thread
string · 4 Rayon threads
large_string · 1 Rayon thread
large_string · 4 Rayon threads
List = array.to_pylist() + encode_batch(); Arrow = encode_batch_arrow(). Native import/validation belongs to Arrow encoding time. Full precision is retained in the saved CSV and JSON results.
Process peak memory before measurement (MiB)
Process peak memory before measurement (MiB)
Every cell: List / Arrow, median across five workers. Context diagnostic; no performance ranking.
Median high-water mark after fixture setup and full-batch warmup, before timed calls. Subtracting this from the later peak does not measure operation allocations.
Timing medians are independent; their components need not sum after aggregation.
string · 1 Rayon thread
string · 4 Rayon threads
large_string · 1 Rayon thread
large_string · 4 Rayon threads
List = array.to_pylist() + encode_batch(); Arrow = encode_batch_arrow(). Native import/validation belongs to Arrow encoding time. Full precision is retained in the saved CSV and JSON results.
Timed calls per worker (calls)
Timed calls per worker (calls)
Every cell: List / Arrow, median across five workers. Context diagnostic; no performance ranking.
Median completed calls across workers. The adaptive loop runs until its time target is reached.
Timing medians are independent; their components need not sum after aggregation.
string · 1 Rayon thread
string · 4 Rayon threads
large_string · 1 Rayon thread
large_string · 4 Rayon threads
List = array.to_pylist() + encode_batch(); Arrow = encode_batch_arrow(). Native import/validation belongs to Arrow encoding time. Full precision is retained in the saved CSV and JSON results.
Cumulative timed duration per worker (ms)
Cumulative timed duration per worker (ms)
Every cell: List / Arrow, median across five workers. Context diagnostic; no performance ranking.
Median sum of timed calls per worker, rather than per-call time or full worker wall time. At least 200 ms was requested; one long call can overshoot that target.
Timing medians are independent; their components need not sum after aggregation.
string · 1 Rayon thread
string · 4 Rayon threads
large_string · 1 Rayon thread
large_string · 4 Rayon threads
List = array.to_pylist() + encode_batch(); Arrow = encode_batch_arrow(). Native import/validation belongs to Arrow encoding time. Full precision is retained in the saved CSV and JSON results.
UTF-8 payload per batch (bytes)
UTF-8 payload per batch (bytes)
Every cell: List / Arrow, median across five workers. Context diagnostic; no performance ranking.
Actual input bytes per batch; expected to match between methods. Body-length column labels remain characters.
Timing medians are independent; their components need not sum after aggregation.
string · 1 Rayon thread
string · 4 Rayon threads
large_string · 1 Rayon thread
large_string · 4 Rayon threads
List = array.to_pylist() + encode_batch(); Arrow = encode_batch_arrow(). Native import/validation belongs to Arrow encoding time. Full precision is retained in the saved CSV and JSON results.
Output tokens per batch (tokens)
Output tokens per batch (tokens)
Every cell: List / Arrow, median across five workers. Context diagnostic; no performance ranking.
Returned tokens per batch; expected to match between methods, independent of worker iteration count.
Timing medians are independent; their components need not sum after aggregation.
string · 1 Rayon thread
string · 4 Rayon threads
large_string · 1 Rayon thread
large_string · 4 Rayon threads
List = array.to_pylist() + encode_batch(); Arrow = encode_batch_arrow(). Native import/validation belongs to Arrow encoding time. Full precision is retained in the saved CSV and JSON results.
Both input paths, UTF-8 throughput, and variability
Values are medians with Q1–Q3 ranges across five fresh workers, or five paired ratios. Quartiles use inclusive linear interpolation. These ranges are descriptive, not confidence intervals; eight of the 64 paired ranges include 1x. All cells are retained, including slowdowns. MiB/s uses actual UTF-8 bytes divided by 1,048,576.
Reproduce from the repository root after installing PyArrow and a release build (
maturin develop --releasefrombindings/python):--format markdownemits all metric matrices directly. Select a subset with, for example,--metrics pylist_ms encode_ms total_ms peak_rss_mib. JSON retains conversion/encoding time, documents/s, UTF-8 MiB/s, output tokens/s, iteration counts, individual samples and environment details. Process peak RSS includes imports, fixtures and warmup; it is not an allocation measurement for the encoding call alone.Validation on macOS ARM64:
cargo fmt --checkandcargo clippy --all-targets --all-features -- -D warningspassed.make check-stylepassed in a temporary source copy. The final Arrow method matches generated output, includingLiteraland its import. Ruff and ty checks passed; caller checks accept all three choices and reject an invalid value.cargo audit -D warningspassed with the repository's two existing advisory exclusions.The full regression and interpreter compatibility checks preceded the final typing-only update. After that update, the wheel was rebuilt and all 151 Arrow tests, stub generation, Clippy, and static typing checks passed again. Linux, Windows, and the remaining wheel targets still require upstream CI.
Related: #1415.