perf: use PyArrow fast path for JSONL reads - #2325
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
🌿 Preview your docs: https://nvidia-preview-codex-pr2303-jsonl-arrow.docs.buildwithfern.com/nemo/curator Here are the markdown pages you've updated: |
9f31d09 to
b391189
Compare
praateekmahajan
changed the title
2155137 to
01423d0
Compare
|
@claude review |
| DEFAULT_PYARROW_BLOCK_SIZE = 8 * 1024 * 1024 | ||
| DEFAULT_PYARROW_MAX_BLOCK_SIZE = 256 * 1024 * 1024 |
There was a problem hiding this comment.
Can you explain what this do and how you chose them?
There was a problem hiding this comment.
TLDR is, pyarrow json reads one "block" (default of 1mb) of bytes at time. If a singular row is larger than that then it might error out. To my knowledge, pd.read_json follows a similar code path (with engine=pyarrow) but doesn't expose blocksize.
We do 8mb as default (assuming 8mb is enough for 1 row), however if we have pdf/image bytes in our jsonl row then it can explode to more, which is why we go upto 256 mb (to avoid the runtime error)
| except pa.ArrowInvalid as error: | ||
| if "straddling object" not in str(error) or block_size >= max_block_size: | ||
| raise | ||
| block_size = min(block_size * 2, max_block_size) |
There was a problem hiding this comment.
In what case would we hit this? Could we end up in a loop where we try reading larger and larger block sizes when we run into this error?
There was a problem hiding this comment.
Imagine a row, which is <8 mb, then we'll be able to read at first try. If not it'll double that, and goto 16mb and try reading again.. It'll stop at 256mb
| read_options = paj.ReadOptions(block_size=block_size, use_threads=False) | ||
| if not is_remote_url(file_path) and not storage_options and compression == "infer": | ||
| return paj.read_json(file_path, read_options=read_options) | ||
| with fsspec.open( |
There was a problem hiding this comment.
Pyarrow parquet supports remote IO, does jsonl as well?
There was a problem hiding this comment.
| def _assign_ids_func(self, filepath: str | list[str], df: pd.DataFrame) -> pd.DataFrame: | ||
| @staticmethod | ||
| def _id_generator_key(task: ReaderTask) -> str | list[str]: | ||
| # TODO(NMCUR-315): Use the deterministic task ID for FileGroupTask as well. |
There was a problem hiding this comment.
Nit but can we not reference internal Linear issues as TODOs?
Signed-off-by: Praateek <praateekm@gmail.com>
Signed-off-by: Praateek <praateekm@gmail.com>
Signed-off-by: Praateek <praateekm@gmail.com>
Signed-off-by: Praateek <praateekm@gmail.com>
Signed-off-by: Praateek <praateekm@gmail.com>
Signed-off-by: Praateek <praateekm@gmail.com>
Signed-off-by: Praateek <praateekm@gmail.com>
Signed-off-by: Praateek <praateekm@gmail.com>
Signed-off-by: Praateek <praateekm@gmail.com>
Signed-off-by: Praateek <praateekm@gmail.com>
3 participants
What changed
This PR makes direct PyArrow parsing the default for JSONL files and keeps the result as a
pa.Tableuntil a stage actually needs pandas.Reader-side
_curator_dedup_idgeneration and assignment now work with both Arrow tables and pandas DataFrames. As a result,LanceReadernow supports_generate_idsand_assign_idswithout converting its Arrow output to pandas first.When pandas is needed,
DocumentBatch.to_pandas()handles the conversion and preserves Arrow-backed string columns.Choosing an engine
JSONL reads use direct PyArrow parsing by default. Callers can select
engine="pandas"when they need pandas-specific parsing or type inference.The two engines do not always infer identical types. For example, pandas may interpret an ISO
created_atvalue as a timezone-aware datetime while PyArrow retains it as a string. Mixed-type columns are another situation where the pandas engine may be the better fit.The PyArrow path continues to support remote files through
fsspec. It normally reads with an 8 MiB block. If PyArrow reports that a single JSONL record straddles the block boundary, the reader retries with progressively larger blocks so that it can read the complete row. The 256 MiB maximum is a safety ceiling for unusually large individual records, such as rows containing a base64-encoded image or PDF; it is not the default amount read for every row.Lance and reader IDs
Because ID handling lives in
BaseReader, Arrow-backed Lance reads now support_generate_idsand_assign_idstoo.Lance tasks use a stable identity based on the dataset path, version, and fragment IDs. File readers retain their existing path-based registry keys for backward compatibility; NMCUR-315 tracks moving them to deterministic task IDs as well.
Compatibility
This work is based directly on
main, including RAPIDS 26.08 and pandas 3 from #2303. PyArrow remains at 19.0.1 under the current dependency constraints.Testing
The tests cover:
fsspecinputsDocumentBatchArrow-string conversionBenchmarks
The table compares the latest completed
mainnightly (nightly-2026_08_20__06_39_17_UTC) with this PR (pr-2325-2026_08_21__00_25_05_UTC-519ad6a9). Reader time is the meanjsonl_reader_process_timeper task. Lower is better.domain_classification_raydatadomain_classification_xennaembedding_generation_raydataembedding_generation_xennadedup_removal_raydatadedup_removal_xennascore_filter_raydatascore_filter_xennafasttext_filter_raydatafasttext_filter_xennamodifier_raydatamodifier_xennandd_dynamo_dp8ndd_ray_serve_dp8Across all 14 entries, the median JSONL reader task is 75.2% faster and median E2E runtime improves by 2.0%. Looking only at the 12 entries where JSONL reading is not the dominant cost, median E2E runtime improves by 0.8%.