[Data] - Port over changes from lance-ray into Ray Data

Signed-off-by: yaommen <myanstu@163.com>

Author: myandpr

ci/lint/pydoclint-baseline.txt

@@ -1002,10 +1002,6 @@ python/ray/data/_internal/block_batching/util.py
     DOC402: Function `finalize_batches` has "yield" statements, but the docstring does not have a "Yields" section
     DOC404: Function `finalize_batches` yield type(s) in docstring not consistent with the return annotation. Return annotation exists, but docstring "yields" section does not exist or has 0 type(s).
 --------------------
-python/ray/data/_internal/datasource/lance_datasink.py
-    DOC101: Method `LanceDatasink.__init__`: Docstring contains fewer arguments than in function signature.
-    DOC103: Method `LanceDatasink.__init__`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [**kwargs: , *args: , max_rows_per_file: int, min_rows_per_file: int, mode: Literal['create', 'append', 'overwrite'], schema: Optional[pa.Schema], storage_options: Optional[Dict[str, Any]], uri: str]. Arguments in the docstring but not in the function signature: [max_rows_per_file : , min_rows_per_file : , mode : , schema : , storage_options : , uri : ].
---------------------
 python/ray/data/_internal/datasource/sql_datasource.py
     DOC101: Method `SQLDatasource.supports_sharding`: Docstring contains fewer arguments than in function signature.
     DOC103: Method `SQLDatasource.supports_sharding`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [parallelism: int].

python/ray/data/_internal/datasource/lance_datasink.py

@@ -1,5 +1,4 @@
 import pickle
-from itertools import chain
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -14,15 +13,40 @@
 
 import pyarrow as pa
 
+from ray._common.retry import call_with_retry
+from ray.data._internal.datasource.lance_utils import (
+    create_storage_options_provider,
+    get_or_create_namespace,
+)
 from ray.data._internal.util import _check_import
 from ray.data.block import BlockAccessor
+from ray.data.context import DataContext
 from ray.data.datasource.datasink import Datasink
 
 if TYPE_CHECKING:
     import pandas as pd
     from lance.fragment import FragmentMetadata
 
 
+def _declare_table_with_fallback(
+    namespace, table_id: List[str]
+) -> Tuple[str, Optional[Dict[str, str]]]:
+    """Declare a table using declare_table, falling back to create_empty_table."""
+    try:
+        from lance_namespace import DeclareTableRequest
+
+        declare_request = DeclareTableRequest(id=table_id, location=None)
+        declare_response = namespace.declare_table(declare_request)
+        return declare_response.location, declare_response.storage_options
+    except (AttributeError, NotImplementedError):
+        # Fallback for older namespace implementations without declare_table.
+        from lance_namespace import CreateEmptyTableRequest
+
+        create_request = CreateEmptyTableRequest(id=table_id)
+        create_response = namespace.create_empty_table(create_request)
+        return create_response.location, create_response.storage_options
+
+
 def _write_fragment(
     stream: Iterable[Union["pa.Table", "pd.DataFrame"]],
     uri: str,
@@ -33,12 +57,20 @@ def _write_fragment(
     max_rows_per_group: int = 1024,  # Only useful for v1 writer.
     data_storage_version: Optional[str] = None,
     storage_options: Optional[Dict[str, Any]] = None,
+    namespace_impl: Optional[str] = None,
+    namespace_properties: Optional[Dict[str, str]] = None,
+    table_id: Optional[List[str]] = None,
+    retry_params: Optional[Dict[str, Any]] = None,
 ) -> List[Tuple["FragmentMetadata", "pa.Schema"]]:
     import pandas as pd
     from lance.fragment import DEFAULT_MAX_BYTES_PER_FILE, write_fragments
 
+    stream = list(stream)
+    if not stream:
+        return []
+
     if schema is None:
-        first = next(stream)
+        first = stream[0]
         if isinstance(first, pd.DataFrame):
             schema = pa.Schema.from_pandas(first).remove_metadata()
         else:
@@ -47,8 +79,6 @@ def _write_fragment(
             # Empty table.
             schema = None
 
-        stream = chain([first], stream)
-
     def record_batch_converter():
         for block in stream:
             tbl = BlockAccessor.for_block(block).to_arrow()
@@ -58,17 +88,35 @@ def record_batch_converter():
         DEFAULT_MAX_BYTES_PER_FILE if max_bytes_per_file is None else max_bytes_per_file
     )
 
-    reader = pa.RecordBatchReader.from_batches(schema, record_batch_converter())
-    fragments = write_fragments(
-        reader,
-        uri,
-        schema=schema,
-        max_rows_per_file=max_rows_per_file,
-        max_rows_per_group=max_rows_per_group,
-        max_bytes_per_file=max_bytes_per_file,
-        data_storage_version=data_storage_version,
-        storage_options=storage_options,
+    if retry_params is None:
+        retry_params = {
+            "description": "write lance fragments",
+            "match": [],
+            "max_attempts": 1,
+            "max_backoff_s": 0,
+        }
+
+    storage_options_provider = create_storage_options_provider(
+        namespace_impl,
+        namespace_properties,
+        table_id,
     )
+
+    def _write_once():
+        reader = pa.RecordBatchReader.from_batches(schema, record_batch_converter())
+        return write_fragments(
+            reader,
+            uri,
+            schema=schema,
+            max_rows_per_file=max_rows_per_file,
+            max_rows_per_group=max_rows_per_group,
+            max_bytes_per_file=max_bytes_per_file,
+            data_storage_version=data_storage_version,
+            storage_options=storage_options,
+            storage_options_provider=storage_options_provider,
+        )
+
+    fragments = call_with_retry(_write_once, **retry_params)
     return [(fragment, schema) for fragment in fragments]
 
 
@@ -77,21 +125,97 @@ class _BaseLanceDatasink(Datasink):
 
     def __init__(
         self,
-        uri: str,
+        uri: Optional[str] = None,
         schema: Optional[pa.Schema] = None,
         mode: Literal["create", "append", "overwrite"] = "create",
         storage_options: Optional[Dict[str, Any]] = None,
-        *args,
-        **kwargs,
+        table_id: Optional[List[str]] = None,
+        namespace_impl: Optional[str] = None,
+        namespace_properties: Optional[Dict[str, str]] = None,
+        *args: Any,
+        **kwargs: Any,
     ):
         super().__init__(*args, **kwargs)
 
-        self.uri = uri
+        merged_storage_options: Dict[str, Any] = {}
+        if storage_options:
+            merged_storage_options.update(storage_options)
+
+        self._namespace_impl = namespace_impl
+        self._namespace_properties = namespace_properties
+
+        namespace = get_or_create_namespace(namespace_impl, namespace_properties)
+
+        if namespace is not None and table_id is not None:
+            if uri is not None:
+                import warnings
+
+                warnings.warn(
+                    "The 'uri' argument is ignored when namespace parameters are "
+                    "provided. The resolved namespace location will be used instead.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+            from lance_namespace import DescribeTableRequest
+            from lance_namespace.errors import TableNotFoundError
+
+            self.table_id = table_id
+            has_namespace_storage_options = True
+
+            if mode == "append":
+                describe_request = DescribeTableRequest(id=table_id)
+                describe_response = namespace.describe_table(describe_request)
+                self.uri = describe_response.location
+                if describe_response.storage_options:
+                    merged_storage_options.update(describe_response.storage_options)
+            elif mode == "overwrite":
+                try:
+                    describe_request = DescribeTableRequest(id=table_id)
+                    describe_response = namespace.describe_table(describe_request)
+                    self.uri = describe_response.location
+                    if describe_response.storage_options:
+                        merged_storage_options.update(describe_response.storage_options)
+                except TableNotFoundError:
+                    uri, ns_storage_options = _declare_table_with_fallback(
+                        namespace, table_id
+                    )
+                    self.uri = uri
+                    if ns_storage_options:
+                        merged_storage_options.update(ns_storage_options)
+            else:
+                uri, ns_storage_options = _declare_table_with_fallback(
+                    namespace, table_id
+                )
+                self.uri = uri
+                if ns_storage_options:
+                    merged_storage_options.update(ns_storage_options)
+
+            self._has_namespace_storage_options = has_namespace_storage_options
+        else:
+            self.table_id = None
+            if uri is None:
+                raise ValueError(
+                    "Must provide either 'uri' or ('namespace_impl' and 'table_id')."
+                )
+            self.uri = uri
+            self._has_namespace_storage_options = False
+
         self.schema = schema
         self.mode = mode
 
         self.read_version: Optional[int] = None
-        self.storage_options = storage_options
+        self.storage_options = merged_storage_options
+
+    @property
+    def storage_options_provider(self):
+        """Lazily create storage options provider using namespace_impl/properties."""
+        if not self._has_namespace_storage_options:
+            return None
+        return create_storage_options_provider(
+            self._namespace_impl,
+            self._namespace_properties,
+            self.table_id,
+        )
 
     @property
     def supports_distributed_writes(self) -> bool:
@@ -103,7 +227,11 @@ def on_write_start(self, schema: Optional["pa.Schema"] = None) -> None:
         import lance
 
         if self.mode == "append":
-            ds = lance.LanceDataset(self.uri, storage_options=self.storage_options)
+            ds = lance.LanceDataset(
+                self.uri,
+                storage_options=self.storage_options,
+                storage_options_provider=self.storage_options_provider,
+            )
             self.read_version = ds.version
             if self.schema is None:
                 self.schema = ds.schema
@@ -152,6 +280,7 @@ def on_write_complete(
             op,
             read_version=self.read_version,
             storage_options=self.storage_options,
+            storage_options_provider=self.storage_options_provider,
         )
 
 
@@ -164,44 +293,52 @@ class LanceDatasink(_BaseLanceDatasink):
     we can use `LanceFragmentWriter` and `LanceCommitter`.
 
     Args:
-        uri : the base URI of the dataset.
-        schema : pyarrow.Schema, optional.
-            The schema of the dataset.
-        mode : str, optional
-            The write mode. Default is 'append'.
-            Choices are 'append', 'create', 'overwrite'.
-        min_rows_per_file : int, optional
-            The minimum number of rows per file. Default is 1024 * 1024.
-        max_rows_per_file : int, optional
-            The maximum number of rows per file. Default is 64 * 1024 * 1024.
-        data_storage_version: optional, str, default None
-            The version of the data storage format to use. Newer versions are more
-            efficient but require newer versions of lance to read.  The default is
-            "legacy" which will use the legacy v1 version.  See the user guide
-            for more details.
-        storage_options : Dict[str, Any], optional
-            The storage options for the writer. Default is None.
+        uri: The base URI of the dataset.
+        schema: The schema of the dataset.
+        mode: The write mode. Default is 'append'. Choices are 'append', 'create', 'overwrite'.
+        min_rows_per_file: The minimum number of rows per file. Default is 1024 * 1024.
+        max_rows_per_file: The maximum number of rows per file. Default is 64 * 1024 * 1024.
+        data_storage_version: The version of the data storage format to use. Newer versions are more
+            efficient but require newer versions of lance to read. The default is "legacy",
+            which will use the legacy v1 version. See the user guide for more details.
+        storage_options: The storage options for the writer. Default is None.
+        table_id: The table identifier as a list of strings, used with namespace params.
+        namespace_impl: The namespace implementation type (e.g., "rest", "dir").
+            Used together with namespace_properties and table_id for credentials vending.
+        namespace_properties: Properties for connecting to the namespace.
+            Used together with namespace_impl and table_id for credentials vending.
+        *args: Additional positional arguments forwarded to the base class.
+        **kwargs: Additional keyword arguments forwarded to the base class.
     """
 
     NAME = "Lance"
+    WRITE_FRAGMENTS_ERRORS_TO_RETRY = ["LanceError(IO)"]
+    WRITE_FRAGMENTS_MAX_ATTEMPTS = 10
+    WRITE_FRAGMENTS_RETRY_MAX_BACKOFF_SECONDS = 32
 
     def __init__(
         self,
-        uri: str,
+        uri: Optional[str] = None,
         schema: Optional[pa.Schema] = None,
         mode: Literal["create", "append", "overwrite"] = "create",
         min_rows_per_file: int = 1024 * 1024,
         max_rows_per_file: int = 64 * 1024 * 1024,
         data_storage_version: Optional[str] = None,
         storage_options: Optional[Dict[str, Any]] = None,
-        *args,
-        **kwargs,
+        table_id: Optional[List[str]] = None,
+        namespace_impl: Optional[str] = None,
+        namespace_properties: Optional[Dict[str, str]] = None,
+        *args: Any,
+        **kwargs: Any,
     ):
         super().__init__(
             uri,
             schema=schema,
             mode=mode,
             storage_options=storage_options,
+            table_id=table_id,
+            namespace_impl=namespace_impl,
+            namespace_properties=namespace_properties,
             *args,
             **kwargs,
         )
@@ -212,6 +349,16 @@ def __init__(
         # if mode is append, read_version is read from existing dataset.
         self.read_version: Optional[int] = None
 
+        match = []
+        match.extend(self.WRITE_FRAGMENTS_ERRORS_TO_RETRY)
+        match.extend(DataContext.get_current().retried_io_errors)
+        self._retry_params = {
+            "description": "write lance fragments",
+            "match": match,
+            "max_attempts": self.WRITE_FRAGMENTS_MAX_ATTEMPTS,
+            "max_backoff_s": self.WRITE_FRAGMENTS_RETRY_MAX_BACKOFF_SECONDS,
+        }
+
     @property
     def min_rows_per_write(self) -> int:
         return self.min_rows_per_file
@@ -231,6 +378,10 @@ def write(
             max_rows_per_file=self.max_rows_per_file,
             data_storage_version=self.data_storage_version,
             storage_options=self.storage_options,
+            namespace_impl=self._namespace_impl,
+            namespace_properties=self._namespace_properties,
+            table_id=self.table_id,
+            retry_params=self._retry_params,
         )
         return [
             (pickle.dumps(fragment), pickle.dumps(schema))