fix: cast null arrays to the appropriate type when coercing to a table (

#3362) @eddyxu @westonpace @Jay-ju Don't ask me why :) just try testing `None` with `LanceDatasink`. See this comment: #3308 (comment) Ref issue: #3308 --------- Co-authored-by: Weston Pace <[email protected]>
lancedb · Jan 10, 2025 · 2142594 · 2142594
1 parent 69d4610
commit 2142594
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 18 deletions.
diff --git a/python/python/lance/ray/sink.py b/python/python/lance/ray/sink.py
@@ -46,23 +46,8 @@ def _pd_to_arrow(
         tbl.schema = tbl.schema.remove_metadata()
         return tbl
     elif isinstance(df, pa.Table):
-        fields = df.schema.names
-        new_columns = []
-        new_fields = []
-        for field in fields:
-            col = df[field]
-            new_field = pa.field(field, col.type)
-            if (
-                pa.types.is_null(col.type)
-                and schema.field_by_name(field).type == pa.string()
-            ):
-                new_field = pa.field(field, pa.string())
-                col = pa.compute.if_else(pa.compute.is_null(col), NONE_ARROW_STR, col)
-            new_columns.append(col)
-            new_fields.append(new_field)
-        new_schema = pa.schema(fields=new_fields)
-        new_table = pa.Table.from_arrays(new_columns, schema=new_schema)
-        return new_table
+        if schema is not None:
+            return df.cast(schema)
     return df
 
 
@@ -439,6 +424,7 @@ def write_lance(
     output_uri: str,
     *,
     schema: Optional[pa.Schema] = None,
+    mode: Literal["create", "append", "overwrite"] = "create",
     transform: Optional[
         Callable[[pa.Table], Union[pa.Table, Generator[None, pa.Table, None]]]
     ] = None,
@@ -485,7 +471,9 @@ def write_lance(
         ),
         batch_size=max_rows_per_file,
     ).write_datasink(
-        LanceCommitter(output_uri, schema=schema, storage_options=storage_options)
+        LanceCommitter(
+            output_uri, schema=schema, mode=mode, storage_options=storage_options
+        )
     )
 
 

diff --git a/python/python/tests/test_ray.py b/python/python/tests/test_ray.py
@@ -138,3 +138,26 @@ def f(row):
     assert len(pylist) == 10
     for item in pylist:
         assert item is None
+
+
+@pytest.mark.filterwarnings("ignore::DeprecationWarning")
+def test_ray_write_lance_none_str_datasink(tmp_path: Path):
+    def f(row):
+        return {
+            "id": row["id"],
+            "str": None,
+        }
+
+    schema = pa.schema([pa.field("id", pa.int64()), pa.field("str", pa.string())])
+
+    sink = LanceDatasink(tmp_path, schema=schema)
+    (ray.data.range(10).map(f).write_datasink(sink))
+    ds = lance.dataset(tmp_path)
+    ds.count_rows() == 10
+    assert ds.schema == schema
+
+    tbl = ds.to_table()
+    pylist = tbl["str"].to_pylist()
+    assert len(pylist) == 10
+    for item in pylist:
+        assert item is None