chore: clean up reader coerce in fragment.py (#3432)

lancedb · Feb 4, 2025 · 2295324 · 2295324
1 parent 42722fb
commit 2295324
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 23 deletions.
diff --git a/python/python/lance/fragment.py b/python/python/lance/fragment.py
@@ -24,8 +24,6 @@
 
 import pyarrow as pa
 
-from .dependencies import _check_for_pandas
-from .dependencies import pandas as pd
 from .lance import (
     DeletionFile as DeletionFile,
 )
@@ -257,7 +255,7 @@ def create_from_file(
     @staticmethod
     def create(
         dataset_uri: Union[str, Path],
-        data: Union[pa.Table, pa.RecordBatchReader],
+        data: ReaderLike,
         fragment_id: Optional[int] = None,
         schema: Optional[pa.Schema] = None,
         max_rows_per_group: int = 1024,
@@ -331,16 +329,7 @@ def create(
             else:
                 data_storage_version = "stable"
 
-        if _check_for_pandas(data) and isinstance(data, pd.DataFrame):
-            reader = pa.Table.from_pandas(data, schema=schema).to_reader()
-        elif isinstance(data, pa.Table):
-            reader = data.to_reader()
-        elif isinstance(data, pa.dataset.Scanner):
-            reader = data.to_reader()
-        elif isinstance(data, pa.RecordBatchReader):
-            reader = data
-        else:
-            raise TypeError(f"Unknown data_obj type {type(data)}")
+        reader = _coerce_reader(data, schema)
 
         if isinstance(dataset_uri, Path):
             dataset_uri = str(dataset_uri)
@@ -797,16 +786,7 @@ def write_fragments(
     """
     from .dataset import LanceDataset
 
-    if _check_for_pandas(data) and isinstance(data, pd.DataFrame):
-        reader = pa.Table.from_pandas(data, schema=schema).to_reader()
-    elif isinstance(data, pa.Table):
-        reader = data.to_reader()
-    elif isinstance(data, pa.dataset.Scanner):
-        reader = data.to_reader()
-    elif isinstance(data, pa.RecordBatchReader):
-        reader = data
-    else:
-        raise TypeError(f"Unknown data_obj type {type(data)}")
+    reader = _coerce_reader(data, schema)
 
     if isinstance(dataset_uri, Path):
         dataset_uri = str(dataset_uri)

diff --git a/python/python/lance/types.py b/python/python/lance/types.py
@@ -18,6 +18,7 @@
         pa.Table,
         pa.dataset.Dataset,
         pa.dataset.Scanner,
+        pa.RecordBatch,
         Iterable[RecordBatch],
         pa.RecordBatchReader,
     ]