fix(datasink): correct write-result payload types and guard missing URIs

lance_ray/datasink.py

@@ -12,7 +12,7 @@
 from lance_namespace import DescribeTableRequest
 from ray.data import DataContext
 from ray.data._internal.util import _check_import
-from ray.data.datasource.datasink import Datasink
+from ray.data.datasource.datasink import Datasink, WriteResult
 
 from .fragment import prepare_fragment_write_options, write_fragment
 from .utils import (
@@ -23,16 +23,25 @@
 )
 
 if TYPE_CHECKING:
+    from lance.fragment import FragmentMetadata
+    from lance_namespace import LanceNamespace
+
     import pandas as pd
 
+#: What a Lance write task hands back to the commit phase: one
+#: ``(pickled fragment, pickled schema)`` pair per written fragment.
+WriteReturn = list[tuple[bytes, bytes]]
+
 
 def _declare_table_with_fallback(
-    namespace, table_id: list[str]
-) -> tuple[str, Optional[dict[str, str]]]:
+    namespace: "LanceNamespace", table_id: list[str]
+) -> tuple[Optional[str], Optional[dict[str, str]]]:
     """Declare a table using declare_table, falling back to create_empty_table.
 
     Returns:
-        Tuple of (uri, storage_options)
+        Tuple of (uri, storage_options). The namespace protocol marks the
+        location as optional, so ``uri`` may be ``None``; the datasink raises
+        when it later needs a location it does not have.
     """
     try:
         from lance_namespace import DeclareTableRequest
@@ -49,7 +58,7 @@ def _declare_table_with_fallback(
         return create_response.location, create_response.storage_options
 
 
-class _BaseLanceDatasink(Datasink):
+class _BaseLanceDatasink(Datasink[WriteReturn]):
     """Base class for Lance Datasink."""
 
     def __init__(
@@ -84,6 +93,8 @@ def __init__(
         # Construct namespace from impl and properties (cached per worker)
         namespace = get_or_create_namespace(namespace_impl, namespace_properties)
 
+        self.table_id: Optional[list[str]]
+        self.uri: Optional[str]
         if namespace is not None and table_id is not None:
             self.table_id = table_id
 
@@ -137,21 +148,32 @@ def namespace_kwargs(self) -> dict[str, Any]:
             self._namespace_impl, self._namespace_properties, self.table_id
         )
 
+    @property
+    def dataset_uri(self) -> str:
+        """The resolved dataset URI, which every write/commit step requires."""
+        if self.uri is None:
+            raise ValueError(
+                "No dataset location is available. Provide 'uri', or "
+                "'namespace_impl' + 'table_id' pointing at a table whose "
+                "namespace reports a location."
+            )
+        return self.uri
+
     @property
     def supports_distributed_writes(self) -> bool:
         return True
 
-    def on_write_start(self, schema: Optional[pa.Schema] = None):
+    def on_write_start(self, schema: Optional[pa.Schema] = None) -> None:
         _check_import(self, module="lance", package="pylance")
 
         import lance
 
         if self.mode == "append":
-            base_store_params_kwargs = {}
+            base_store_params_kwargs: dict[str, Any] = {}
             if self.base_store_params:
                 base_store_params_kwargs = {"base_store_params": self.base_store_params}
             ds = lance.LanceDataset(
-                self.uri,
+                self.dataset_uri,
                 storage_options=self.storage_options,
                 **self.namespace_kwargs,
                 **base_store_params_kwargs,
@@ -162,22 +184,27 @@ def on_write_start(self, schema: Optional[pa.Schema] = None):
 
     def on_write_complete(
         self,
-        write_result: list[list[tuple[str, str]]],
-    ):
+        write_result: WriteResult[WriteReturn] | list[WriteReturn],
+    ) -> None:
         import warnings
 
         import lance
 
-        write_results = write_result
-        if not write_results:
+        if not write_result:
             warnings.warn(
                 "write_results is empty.",
                 DeprecationWarning,
                 stacklevel=2,
             )
             return
-        if hasattr(write_results, "write_returns"):
-            write_results = write_results.write_returns  # type: ignore
+
+        # Older Ray releases pass the raw per-task return values, newer ones
+        # pass an aggregated ``WriteResult``; support both.
+        write_results: list[WriteReturn]
+        if hasattr(write_result, "write_returns"):
+            write_results = write_result.write_returns
+        else:
+            write_results = write_result
 
         if len(write_results) == 0:
             warnings.warn(
@@ -187,8 +214,8 @@ def on_write_complete(
             )
             return
 
-        fragments = []
-        schema = None
+        fragments: list[FragmentMetadata] = []
+        schema: Optional[pa.Schema] = None
         for batch in write_results:
             for fragment_str, schema_str in batch:
                 fragment = pickle.loads(fragment_str)
@@ -198,7 +225,7 @@ def on_write_complete(
         # Skip commit when there are no fragments.
         if not schema:
             return
-        op = None
+        op: Optional[lance.LanceOperation.BaseOperation] = None
         if self.mode in {"create", "overwrite"}:
             op = lance.LanceOperation.Overwrite(
                 schema,
@@ -212,11 +239,11 @@ def on_write_complete(
         elif self.mode == "append":
             op = lance.LanceOperation.Append(fragments)
         if op:
-            base_store_params_kwargs = {}
+            base_store_params_kwargs: dict[str, Any] = {}
             if self.base_store_params:
                 base_store_params_kwargs = {"base_store_params": self.base_store_params}
             lance.LanceDataset.commit(
-                self.uri,
+                self.dataset_uri,
                 op,
                 read_version=self.read_version,
                 storage_options=self.storage_options,
@@ -369,10 +396,10 @@ def write(
         self,
         blocks: Iterable[Union[pa.Table, "pd.DataFrame"]],
         ctx: Any,
-    ):
+    ) -> WriteReturn:
         fragments_and_schema = write_fragment(
             blocks,
-            self.uri,
+            self.dataset_uri,
             schema=self.schema,
             max_rows_per_file=self.max_rows_per_file,
             data_storage_version=self.data_storage_version,
@@ -412,9 +439,9 @@ def write(
         self,
         blocks: Iterable[Union[pa.Table, "pd.DataFrame"]],
         _ctx: Any,
-    ):
+    ) -> WriteReturn:
         """Passthrough the fragments to commit phase"""
-        v = []
+        v: WriteReturn = []
         for block in blocks:
             # If block is empty, skip to get "fragment" and "schema" filed
             if len(block) == 0: