nemo_curator.stages.text.io.writer.lance

Module Contents

Classes

Name	Description
`LanceWriter`	Write `DocumentBatch` tables to Lance fragments and checkpoint the commit.

Functions

Name	Description
`_find_fragment_version`	Find the transaction that committed the fragment files.
`commit_lance_checkpoint`	Commit checkpointed fragments and return their exact Lance dataset version.

Data

_COMMITTED_MARKER

_RECORDS_DIR

API

class nemo_curator.stages.text.io.writer.lance.LanceWriter(
    path: str,
    commit_path: str,
    schema: pyarrow.Schema | None = None,
    write_kwargs: dict[str, typing.Any] = dict(),
    fields: list[str] | None = None,
    name: str = 'lance_writer',
    mode: typing.Literal['create', 'append', 'overwrite'] = 'create'
)

Dataclass

Bases: ProcessingStage[DocumentBatch, FileGroupTask]

Write DocumentBatch tables to Lance fragments and checkpoint the commit.

commit_path

str

fields

list[str] | None = None

mode

Literal['create', 'append', 'overwrite'] = 'create'

name

str = 'lance_writer'

path

str

schema

Schema | None = None

write_kwargs

dict[str, Any] = field(default_factory=dict)

nemo_curator.stages.text.io.writer.lance.LanceWriter.__post_init__() -> None

nemo_curator.stages.text.io.writer.lance.LanceWriter._output_table_and_schema(
    task: nemo_curator.tasks.DocumentBatch
) -> tuple[pyarrow.Table, pyarrow.Schema | None]

Select output columns and restore source Lance field types when available.

nemo_curator.stages.text.io.writer.lance.LanceWriter.inputs() -> tuple[list[str], list[str]]

nemo_curator.stages.text.io.writer.lance.LanceWriter.outputs() -> tuple[list[str], list[str]]

nemo_curator.stages.text.io.writer.lance.LanceWriter.process(
    task: nemo_curator.tasks.DocumentBatch
) -> nemo_curator.tasks.FileGroupTask

Write one batch as uncommitted fragments and persist their commit records.

nemo_curator.stages.text.io.writer.lance._find_fragment_version(
    dataset: lance.LanceDataset,
    fragments: list[lance.fragment.FragmentMetadata]
) -> int | None

Find the transaction that committed the fragment files.

nemo_curator.stages.text.io.writer.lance.commit_lance_checkpoint(
    dataset_path: str,
    checkpoint_path: str,
    dataset_storage_options: dict[str, typing.Any] | None = None,
    checkpoint_storage_options: dict[str, typing.Any] | None = None
) -> int

Commit checkpointed fragments and return their exact Lance dataset version.

nemo_curator.stages.text.io.writer.lance._COMMITTED_MARKER = '_COMMITTED'

nemo_curator.stages.text.io.writer.lance._RECORDS_DIR = 'records'