perf: use PyArrow fast path for JSONL reads

fern/versions/main/pages/curate-text/load-data/read-existing.mdx

@@ -119,6 +119,41 @@ Both `JsonlReader` and `ParquetReader` support these configuration options:
 | `fields` | list[str] \| None | Column names to read (column selection) | None (all columns) |
 | `read_kwargs` | dict[str, Any] \| None | Extra arguments for the underlying reader | None |
 
+### JSONL Engine Selection
+
+`JsonlReader` parses JSONL directly into a PyArrow table by default. Choose the engine through
+`read_kwargs` when your dataset requires different inference behavior.
+
+| `read_kwargs["engine"]` | Output backing type | Use when |
+| --- | --- | --- |
+| Omitted or `"pyarrow_direct"` | `pyarrow.Table` | Recommended for throughput and Arrow-native processing |
+| `"pandas"` | `pandas.DataFrame` | You need pandas-specific options or inference, such as `convert_dates` or mixed-type object columns |
+| Another pandas engine, such as `"ujson"` or `"pyarrow"` | `pandas.DataFrame` | You explicitly need an engine supported by `pandas.read_json` |
+
+The direct engine does not silently fall back to pandas. If PyArrow cannot infer a consistent
+schema—for example, when one JSON field contains both numbers and strings—select the pandas
+engine explicitly:
+
+```python
+reader = JsonlReader(
+    file_paths="/path/to/data",
+    read_kwargs={
+        "engine": "pandas",
+        "convert_dates": ["created_at"],
+    },
+)
+```
+
+The direct engine accepts these additional options:
+
+| Option | Description | Default |
+| --- | --- | --- |
+| `pyarrow_block_size` | Initial number of bytes PyArrow processes per parser block | 8 MiB |
+| `pyarrow_max_block_size` | Maximum parser block used when retrying an unusually large JSON object | 256 MiB |
+
+The direct reader processes 8 MiB chunks and retries with larger chunks, up to 256 MiB, for rows
+containing large payloads such as base64-encoded images or PDFs.
+
 ### Parquet-Specific Features
 
 `ParquetReader` provides these optimizations:

nemo_curator/stages/text/classifiers/utils.py

@@ -36,10 +36,10 @@ def outputs(self) -> tuple[list[str], list[str]]:
         return ["data"], [ATTENTION_MASK_FIELD, SEQ_ORDER_FIELD]
 
     def process(self, batch: DocumentBatch) -> DocumentBatch:
-        if SEQ_ORDER_FIELD in batch.data.columns:
+        if SEQ_ORDER_FIELD in batch.get_columns():
             return batch
 
-        output = batch.data.copy()
+        output = batch.to_pandas()
 
         # Add column to preserve original order
         output[SEQ_ORDER_FIELD] = np.arange(len(output))

nemo_curator/stages/text/io/reader/base.py

@@ -90,19 +90,19 @@ def process(self, task: ReaderTask) -> DocumentBatch:
         return self._document_batch(task, output)
 
     def _document_batch(self, task: ReaderTask, output: ReaderOutput) -> DocumentBatch:
-        result = output.data
-        # Apply IDs only for Pandas DataFrames
-        if isinstance(result, pd.DataFrame):
-            if self._generate_ids:
-                result = self._generate_ids_func(task.data, result)
-            elif self._assign_ids:
-                result = self._assign_ids_func(task.data, result)
-
-        return DocumentBatch(
+        batch = DocumentBatch(
             dataset_name=task.dataset_name,
-            data=result,
+            data=output.data,
             _metadata=output.metadata if output.metadata is not None else task._metadata,
         )
+        if self._generate_ids or self._assign_ids:
+            batch_key = self._id_generator_key(task)
+            if self._generate_ids:
+                self._generate_ids_func(batch_key, batch)
+            else:
+                self._assign_ids_func(batch_key, batch)
+
+        return batch
 
     def _validate_result(self, task: ReaderTask, result: ReaderData) -> None:
         if self.allow_empty:
@@ -125,26 +125,41 @@ def read_task(
         raise NotImplementedError
 
     # ID helpers ----------------------------------------------------------------
-    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.
+        # Keep returning file paths for backward compatibility until existing ID registries are migrated.
+        if isinstance(task, FileGroupTask):
+            return task.data
+        return task.get_deterministic_id()
+
+    @staticmethod
+    def _append_ids(batch: DocumentBatch, start_id: int, count: int) -> None:
+        from nemo_curator.stages.deduplication.id_generator import CURATOR_DEDUP_ID_STR
+
+        ids = np.arange(start_id, start_id + count)
+        if isinstance(batch.data, pd.DataFrame):
+            batch.data[CURATOR_DEDUP_ID_STR] = ids
+        else:
+            batch.data = batch.data.append_column(CURATOR_DEDUP_ID_STR, pa.array(ids, type=pa.int64()))
+
+    def _assign_ids_func(self, batch_key: str | list[str], batch: DocumentBatch) -> None:
         from nemo_curator.stages.deduplication.id_generator import CURATOR_DEDUP_ID_STR
 
-        if CURATOR_DEDUP_ID_STR not in df.columns:
-            min_id, max_id = ray.get(self.id_generator.get_batch_range.remote(filepath, None))
-            df[CURATOR_DEDUP_ID_STR] = np.arange(min_id, max_id + 1)
+        if CURATOR_DEDUP_ID_STR not in batch.get_columns():
+            min_id, max_id = ray.get(self.id_generator.get_batch_range.remote(batch_key, None))
+            self._append_ids(batch, min_id, max_id - min_id + 1)
         else:
-            logger.warning(f"Column {CURATOR_DEDUP_ID_STR} already exists in {filepath}, not re-assigning IDs")
-        return df
+            logger.warning(f"Column {CURATOR_DEDUP_ID_STR} already exists in {batch_key}, not re-assigning IDs")
 
-    def _generate_ids_func(self, filepath: str | list[str], df: pd.DataFrame) -> pd.DataFrame:
+    def _generate_ids_func(self, batch_key: str | list[str], batch: DocumentBatch) -> None:
         from nemo_curator.stages.deduplication.id_generator import CURATOR_DEDUP_ID_STR
 
-        if CURATOR_DEDUP_ID_STR not in df.columns:
-            num_rows = len(df)
-            min_id = ray.get(self.id_generator.register_batch.remote(filepath, num_rows))
-            df[CURATOR_DEDUP_ID_STR] = np.arange(min_id, min_id + num_rows)
+        if CURATOR_DEDUP_ID_STR not in batch.get_columns():
+            min_id = ray.get(self.id_generator.register_batch.remote(batch_key, batch.num_items))
+            self._append_ids(batch, min_id, batch.num_items)
         else:
-            logger.warning(f"Column {CURATOR_DEDUP_ID_STR} already exists in {filepath}, not generating new IDs")
-        return df
+            logger.warning(f"Column {CURATOR_DEDUP_ID_STR} already exists in {batch_key}, not generating new IDs")
 
     def ray_stage_spec(self) -> dict[str, Any]:
         return {RayStageSpecKeys.IS_ACTOR_STAGE: self._generate_ids or self._assign_ids}

nemo_curator/stages/text/io/reader/jsonl.py

@@ -15,16 +15,146 @@
 from dataclasses import dataclass, field
 from typing import Any, Literal
 
+import fsspec
 import pandas as pd
+import pyarrow as pa
+import pyarrow.json as paj
 from loguru import logger
 
 from nemo_curator.stages.base import CompositeStage
 from nemo_curator.stages.file_partitioning import FilePartitioningStage
 from nemo_curator.tasks import DocumentBatch, EmptyTask
+from nemo_curator.utils.client_utils import is_remote_url
 from nemo_curator.utils.file_utils import FILETYPE_TO_DEFAULT_EXTENSIONS, pandas_select_columns
 
 from .base import BaseFileReader
 
+PANDAS_ENGINE = "pandas"
+PYARROW_DIRECT_ENGINE = "pyarrow_direct"
+# Read 8 MiB chunks to accommodate most rows, but retry up to 256 MiB for
+# rows containing large binary payloads such as base64-encoded images or PDFs.
+DEFAULT_PYARROW_BLOCK_SIZE = 8 * 1024 * 1024
+DEFAULT_PYARROW_MAX_BLOCK_SIZE = 256 * 1024 * 1024
+
+
+def _validate_jsonl_read_kwargs(read_kwargs: dict[str, Any] | None) -> None:
+    if read_kwargs is not None and read_kwargs.get("lines", True) is False:
+        msg = "JsonlReader only supports lines=True"
+        raise RuntimeError(msg)
+
+
+def _pyarrow_select_columns(table: pa.Table, fields: list[str] | None, file_path: str) -> pa.Table | None:
+    if fields is None:
+        return table
+
+    existing_fields = [column for column in fields if column in table.column_names]
+    missing_fields = [column for column in fields if column not in table.column_names]
+    if missing_fields:
+        logger.warning(f"Columns {missing_fields} not found in {file_path}")
+    if existing_fields:
+        return table.select(existing_fields)
+
+    logger.error(f"None of the requested columns found in {file_path}")
+    return None
+
+
+def _read_jsonl_file_with_pyarrow(
+    file_path: str,
+    block_size: int,
+    max_block_size: int,
+    storage_options: dict[str, Any],
+    compression: str | None,
+) -> pa.Table:
+    """Read one JSONL file, growing the parser block for an oversized record.
+
+    PyArrow reports ``straddling object`` when a JSON object is too large for
+    its current parsing window. Each retry doubles the block size up to
+    ``max_block_size``; the error is re-raised at the ceiling, so this loop is
+    bounded. Remote paths and custom storage options are opened through
+    ``fsspec`` and passed to PyArrow as a file-like stream.
+    """
+    while True:
+        try:
+            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(
+                file_path,
+                mode="rb",
+                compression=compression,
+                **storage_options,
+            ) as stream:
+                return paj.read_json(stream, read_options=read_options)
+        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)
+
+
+def _read_jsonl_with_pyarrow(
+    paths: list[str],
+    read_kwargs: dict[str, Any],
+    fields: list[str] | None,
+) -> pa.Table:
+    """Read JSONL paths directly with PyArrow."""
+    read_kwargs = dict(read_kwargs)
+    read_kwargs.pop("engine", None)
+    read_kwargs.pop("lines", None)
+
+    block_size = read_kwargs.pop("pyarrow_block_size", DEFAULT_PYARROW_BLOCK_SIZE)
+    max_block_size = read_kwargs.pop("pyarrow_max_block_size", DEFAULT_PYARROW_MAX_BLOCK_SIZE)
+    storage_options = read_kwargs.pop("storage_options", {}) or {}
+    compression = read_kwargs.pop("compression", "infer")
+    if block_size <= 0 or max_block_size < block_size:
+        msg = "pyarrow block sizes must be positive and max block size must be at least the initial size"
+        raise ValueError(msg)
+    if read_kwargs:
+        unsupported = ", ".join(sorted(read_kwargs))
+        msg = f"Unsupported read_kwargs for engine={PYARROW_DIRECT_ENGINE!r}: {unsupported}"
+        raise TypeError(msg)
+
+    tables = []
+    for file_path in paths:
+        table = _read_jsonl_file_with_pyarrow(
+            file_path,
+            block_size,
+            max_block_size,
+            storage_options,
+            compression,
+        )
+        table = _pyarrow_select_columns(table, fields, file_path)
+        if table is not None:
+            tables.append(table)
+    if not tables:
+        msg = f"No data read from files in task {paths} with direct PyArrow JSONL reader"
+        logger.error(msg)
+        raise ValueError(msg)
+
+    return pa.concat_tables(tables, promote_options="permissive")
+
+
+def _read_jsonl_with_pandas(
+    paths: list[str],
+    read_kwargs: dict[str, Any],
+    fields: list[str] | None,
+) -> pd.DataFrame:
+    read_kwargs = dict(read_kwargs)
+    if read_kwargs.get("engine") == PANDAS_ENGINE:
+        read_kwargs.pop("engine")
+    read_kwargs["lines"] = True
+
+    dfs = []
+    for file_path in paths:
+        df = pd.read_json(file_path, **read_kwargs)
+        if fields is not None:
+            df = pandas_select_columns(df, fields, file_path)
+        dfs.append(df)
+    if not dfs:
+        msg = f"No data read from files in task {paths} with read_kwargs {read_kwargs} in JSONL reader"
+        logger.error(msg)
+        raise ValueError(msg)
+    return pd.concat(dfs, ignore_index=True)
+
 
 @dataclass
 class JsonlReaderStage(BaseFileReader):
@@ -35,7 +165,9 @@ class JsonlReaderStage(BaseFileReader):
 
     Args:
         fields (list[str], optional): If specified, only read these fields (columns). Defaults to None.
-        read_kwargs (dict[str, Any], optional): Keyword arguments for the reader. Defaults to {}.
+        read_kwargs (dict[str, Any], optional): Reader options. ``engine="pyarrow_direct"``
+            uses the direct PyArrow reader; all other engines, including ``"pyarrow"``,
+            are passed to ``pd.read_json``. Defaults to {}.
         _generate_ids (bool): Whether to generate monotonically increasing IDs across all files.
             This uses IdGenerator actor, which needs to be instantiated before using this stage.
             This can be slow, so it is recommended to use AddId stage instead, unless monotonically increasing IDs
@@ -48,36 +180,23 @@ class JsonlReaderStage(BaseFileReader):
 
     name: str = "jsonl_reader"
 
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        _validate_jsonl_read_kwargs(self.read_kwargs)
+
     def read_data(
         self,
         paths: list[str],
         read_kwargs: dict[str, Any] | None = None,
         fields: list[str] | None = None,
-    ) -> pd.DataFrame:
-        """Read JSONL files using Pandas."""
+    ) -> pd.DataFrame | pa.Table:
+        """Read JSONL files using the selected engine."""
 
-        # Normalize read_kwargs to a dict to avoid TypeError when None
-        # Work on a copy to avoid mutating caller's dict
         read_kwargs = {} if read_kwargs is None else dict(read_kwargs)
-        # Default to lines=True if not specified
-        if "lines" in read_kwargs and read_kwargs["lines"] is False:
-            msg = "lines=False is not supported for JSONL reader"
-            raise ValueError(msg)
-        else:
-            read_kwargs["lines"] = True
-
-        dfs = []
-        for file_path in paths:
-            df = pd.read_json(file_path, **read_kwargs)
-            if fields is not None:
-                df = pandas_select_columns(df, fields, file_path)
-            dfs.append(df)
-        # Concatenate all dataframes
-        if not dfs:
-            msg = f"No data read from files in task {paths} with read_kwargs {read_kwargs} in JSONL reader"
-            logger.error(msg)
-            raise ValueError(msg)
-        return pd.concat(dfs, ignore_index=True)
+        engine = read_kwargs.get("engine", PYARROW_DIRECT_ENGINE)
+        if engine == PYARROW_DIRECT_ENGINE:
+            return _read_jsonl_with_pyarrow(paths, read_kwargs, fields)
+        return _read_jsonl_with_pandas(paths, read_kwargs, fields)
 
 
 @dataclass
@@ -87,6 +206,19 @@ class JsonlReader(CompositeStage[EmptyTask, DocumentBatch]):
     This high-level stage decomposes into:
     1. FilePartitioningStage - partitions files into groups
     2. JsonlReaderStage - reads file groups into DocumentBatches
+
+    Args:
+        file_paths: File paths, directories, or glob patterns to read.
+        files_per_partition: Number of files grouped into each reader task.
+            When set, this takes precedence over ``blocksize``.
+        blocksize: Target storage size for each file-group task.
+        fields: Optional columns to retain.
+        read_kwargs: Options passed to ``pd.read_json``, or to the direct
+            PyArrow reader when ``engine="pyarrow_direct"``.
+        task_type: Output task modality. Only ``"document"`` is supported.
+        file_extensions: File extensions considered during partitioning.
+        _generate_ids: Generate stable, monotonically increasing document IDs.
+        _assign_ids: Assign IDs previously registered for the same reader task.
     """
 
     file_paths: str | list[str]
@@ -103,6 +235,7 @@ class JsonlReader(CompositeStage[EmptyTask, DocumentBatch]):
     def __post_init__(self):
         """Initialize parent class after dataclass initialization."""
         super().__init__()
+        _validate_jsonl_read_kwargs(self.read_kwargs)
         if self.read_kwargs is not None:
             self.storage_options = self.read_kwargs.get("storage_options", {})
 

nemo_curator/stages/text/io/reader/lance.py

@@ -142,6 +142,10 @@ def outputs(self) -> tuple[list[str], list[str]]:
         output_fields = list(columns or [])
         if self.include_lance_metadata:
             output_fields.extend([LANCE_ROWID_COLUMN, LANCE_ROWADDR_COLUMN, LANCE_FRAGID_COLUMN])
+        if self._generate_ids or self._assign_ids:
+            from nemo_curator.stages.deduplication.id_generator import CURATOR_DEDUP_ID_STR
+
+            output_fields.append(CURATOR_DEDUP_ID_STR)
         return ["data"], output_fields
 
     def _scanner_kwargs(self, read_kwargs: dict[str, Any], fields: list[str] | None) -> dict[str, Any]: