Document Lance public APIs

Author: VibhuJawa

nemo_curator/stages/text/io/lance_commit.py

@@ -31,6 +31,7 @@ def commit_lance_checkpoint(
     storage_options: dict[str, Any] | None = None,
     checkpoint_storage_options: dict[str, Any] | None = None,
 ) -> int:
+    """Commit records written by ``LanceWriter`` and return the Lance version."""
     import lance
     from lance.schema import json_to_schema
     from lance_ray import LanceFragmentCommitter
@@ -70,6 +71,7 @@ def commit_lance_annotation_checkpoint(
     storage_options: dict[str, Any] | None = None,
     checkpoint_storage_options: dict[str, Any] | None = None,
 ) -> int:
+    """Commit records written by ``LanceAnnotationWriter`` and return the Lance version."""
     import lance
 
     records, committed_version = read_lance_checkpoint(
@@ -103,9 +105,5 @@ def commit_lance_annotation_checkpoint(
         read_version=read_version,
         storage_options=storage_options,
     ).version
-    write_lance_checkpoint_marker(
-        commit_path,
-        version,
-        checkpoint_storage_options,
-    )
+    write_lance_checkpoint_marker(commit_path, version, checkpoint_storage_options)
     return version

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

@@ -30,17 +30,8 @@
 
 
 def _read_dataset_kwargs(read_kwargs: dict[str, Any], version: int | None = None) -> dict[str, Any]:
-    return {
-        **dict(read_kwargs.get("dataset_options") or {}),
-        **{
-            key: value
-            for key, value in (
-                ("storage_options", read_kwargs.get("storage_options")),
-                ("version", read_kwargs.get("version", version)),
-            )
-            if value is not None
-        },
-    }
+    options = {"storage_options": read_kwargs.get("storage_options"), "version": read_kwargs.get("version", version)}
+    return {**dict(read_kwargs.get("dataset_options") or {}), **{k: v for k, v in options.items() if v is not None}}
 
 
 def _scanner_kwargs(read_kwargs: dict[str, Any], fields: list[str] | None) -> dict[str, Any]:
@@ -207,6 +198,7 @@ def process(self, task: LanceReadTask) -> DocumentBatch | None:
 
 @dataclass
 class LanceReader(CompositeStage[EmptyTask, DocumentBatch]):
+    """Read a Lance dataset into Curator ``DocumentBatch`` objects by fragment."""
     path: str
     fragments_per_partition: int = 32
     fields: list[str] | None = None

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

@@ -56,6 +56,7 @@ def _schema_for_table(schema: pa.Schema, table: pa.Table) -> pa.Schema:
 
 @dataclass
 class LanceWriter(ProcessingStage[DocumentBatch, FileGroupTask]):
+    """Write ``DocumentBatch`` tables to Lance fragments and checkpoint the commit."""
     path: str
     commit_path: str
     schema: pa.Schema | None = None
@@ -125,6 +126,7 @@ def process(self, task: DocumentBatch) -> FileGroupTask:
 
 @dataclass
 class LanceAnnotationWriter(ProcessingStage[DocumentBatch, FileGroupTask]):
+    """Update existing Lance rows using metadata columns emitted by ``LanceReader``."""
     path: str
     commit_path: str
     schema: pa.Schema
@@ -151,6 +153,7 @@ def outputs(self) -> tuple[list[str], list[str]]:
         return ["data"], []
 
     def prepare(self) -> int:
+        """Create or validate annotation columns and pin the Lance version for the run."""
         import lance
 
         dataset = lance.dataset(self.path, storage_options=self.storage_options)
@@ -208,14 +211,8 @@ def process(self, task: DocumentBatch) -> FileGroupTask:
             msg = f"Lance annotation update table is missing required columns: {missing}"
             raise ValueError(msg)
         version = self._update_version()
-        dataset = lance.dataset(
-            self.path,
-            **{
-                key: value
-                for key, value in {"storage_options": self.storage_options, "version": version}.items()
-                if value is not None
-            },
-        )
+        options = {"storage_options": self.storage_options, "version": version}
+        dataset = lance.dataset(self.path, **{k: v for k, v in options.items() if v is not None})
 
         record_paths = []
         fragment_ids = sorted(int(value) for value in pc.unique(table[LANCE_FRAGID_COLUMN].combine_chunks()).to_pylist())