Skip to content

Fix schema alignment for buffered single-row Arrow tables - #8692

Open
YusefSyed wants to merge 1 commit into
huggingface:mainfrom
YusefSyed:codex/arrow-writer-row-schema
Open

YusefSyed wants to merge 1 commit into
huggingface:mainfrom
YusefSyed:codex/arrow-writer-row-schema

Conversation

@YusefSyed

Copy link
Copy Markdown

Buffered ArrowWriter.write_row calls concatenate the input tables before the usual schema cast. Compatible rows therefore fail when their columns arrive in a different order or a later row omits a column. This also breaks load_dataset("json", ...) for one-record JSON files, even though the same records work with immediate flushing or larger input tables.

For heterogeneous buffered schemas, align each row to the effective writer schema using the existing table_cast, then concatenate and write one batch. Homogeneous buffers retain their current path. The selected schema still governs missing-column null filling, casts and feature metadata; extra columns remain errors.

The regressions cover buffering, reordered/missing columns, declared and established schemas, semantic feature metadata, invalid casts, and JSON arrays/selected fields with inferred and explicit features.

Validation on macOS arm64, Python 3.11.15 and PyArrow 25.0.1:

  • New regression cases fail with the original method and pass with the fix.
  • pytest tests/test_arrow_writer.py tests/packaged_modules/test_json.py -q: 131 passed, 11 skipped (optional agent-trace dependency).
  • Focused tests/test_arrow_dataset.py map checks covering Arrow/pandas outputs, features, batching, column removal and multiprocessing: 12 passed.
  • make quality and git diff --check: passed.

AI assistance: OpenAI Codex generated the implementation, tests and this description, with GPT-6 Astra used for investigation. The listed checks were executed locally.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant