wehub-resource-sync
461 lines
16 KiB
Python
461 lines
16 KiB
Python
import pickle
|
|
from itertools import chain
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
Callable,
|
|
Dict,
|
|
Iterable,
|
|
Iterator,
|
|
List,
|
|
Optional,
|
|
Tuple,
|
|
Union,
|
|
)
|
|
|
|
import pyarrow as pa
|
|
|
|
from ray._common.retry import call_with_retry
|
|
from ray.data._internal.arrow_ops.transform_pyarrow import (
|
|
reorder_columns_by_schema,
|
|
)
|
|
from ray.data._internal.datasource.lance_utils import (
|
|
create_storage_options_provider,
|
|
get_or_create_namespace,
|
|
)
|
|
from ray.data._internal.savemode import SaveMode
|
|
from ray.data._internal.util import _check_import, unify_schemas_with_validation
|
|
from ray.data.block import Block, BlockAccessor
|
|
from ray.data.context import DataContext
|
|
from ray.data.datasource.datasink import Datasink
|
|
|
|
if TYPE_CHECKING:
|
|
import pandas as pd
|
|
from lance import LanceDataset
|
|
from lance.fragment import FragmentMetadata
|
|
|
|
|
|
_WRITE_LANCE_FRAGMENTS_DESCRIPTION = "write lance fragments"
|
|
|
|
|
|
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 _make_stream_factory(
|
|
stream: Iterable[Block], replayable: bool
|
|
) -> Tuple[Optional[Callable[[], Iterator[Block]]], Optional[Block]]:
|
|
"""Return a reusable stream factory and the first block, or (None, None)."""
|
|
if replayable:
|
|
blocks = list(stream)
|
|
if not blocks:
|
|
return None, None
|
|
|
|
def stream_factory() -> Iterator[Block]:
|
|
return iter(blocks)
|
|
|
|
return stream_factory, blocks[0]
|
|
|
|
stream_iter = iter(stream)
|
|
first = next(stream_iter, None)
|
|
if first is None:
|
|
return None, None
|
|
|
|
def stream_factory() -> Iterator[Block]:
|
|
return chain([first], stream_iter)
|
|
|
|
return stream_factory, first
|
|
|
|
|
|
def _write_fragment(
|
|
stream: Iterable[Block],
|
|
uri: str,
|
|
*,
|
|
schema: Optional["pa.Schema"] = None,
|
|
max_rows_per_file: int = 64 * 1024 * 1024,
|
|
max_bytes_per_file: Optional[int] = None,
|
|
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
|
|
|
|
if retry_params is None:
|
|
retry_params = {
|
|
"description": _WRITE_LANCE_FRAGMENTS_DESCRIPTION,
|
|
"match": [],
|
|
"max_attempts": 1,
|
|
"max_backoff_s": 0,
|
|
}
|
|
|
|
max_attempts = retry_params.get("max_attempts", 1)
|
|
stream_factory, first = _make_stream_factory(stream, replayable=max_attempts > 1)
|
|
if stream_factory is None or first is None:
|
|
return []
|
|
|
|
if schema is None:
|
|
if isinstance(first, pd.DataFrame):
|
|
schema = pa.Schema.from_pandas(first).remove_metadata()
|
|
else:
|
|
schema = first.schema
|
|
if len(schema.names) == 0:
|
|
# Empty table.
|
|
schema = None
|
|
|
|
def record_batch_converter(block_stream):
|
|
for block in block_stream:
|
|
tbl = BlockAccessor.for_block(block).to_arrow()
|
|
# `RecordBatchReader.from_batches(schema, ...)` is positional.
|
|
if schema is not None:
|
|
tbl = reorder_columns_by_schema(tbl, schema)
|
|
yield from tbl.to_batches()
|
|
|
|
max_bytes_per_file = (
|
|
DEFAULT_MAX_BYTES_PER_FILE if max_bytes_per_file is None else max_bytes_per_file
|
|
)
|
|
|
|
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(stream_factory())
|
|
)
|
|
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]
|
|
|
|
|
|
class _BaseLanceDatasink(Datasink):
|
|
"""Base class for Lance Datasink."""
|
|
|
|
def __init__(
|
|
self,
|
|
uri: Optional[str] = None,
|
|
schema: Optional[pa.Schema] = None,
|
|
mode: SaveMode = SaveMode.CREATE,
|
|
storage_options: Optional[Dict[str, Any]] = None,
|
|
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)
|
|
|
|
if mode not in {SaveMode.CREATE, SaveMode.APPEND, SaveMode.OVERWRITE}:
|
|
raise ValueError(
|
|
f"Unsupported Lance write mode: {mode!r}. "
|
|
"Supported modes are SaveMode.CREATE, SaveMode.APPEND, and SaveMode.OVERWRITE."
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
self.table_id = table_id
|
|
if mode != SaveMode.CREATE:
|
|
raise ValueError(
|
|
"Namespace writes currently only support mode='create'. "
|
|
"Use mode='create' for now."
|
|
)
|
|
|
|
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 = True
|
|
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 = 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:
|
|
return True
|
|
|
|
def _open_dataset(self) -> "LanceDataset":
|
|
"""Open the Lance dataset at ``self.uri``.
|
|
|
|
Raises whatever Lance raises if the dataset can't be opened (missing,
|
|
bad ``storage_options``, etc.). Opening natively honors
|
|
``storage_options``/``storage_options_provider``.
|
|
"""
|
|
import lance
|
|
|
|
return lance.LanceDataset(
|
|
self.uri,
|
|
storage_options=self.storage_options,
|
|
storage_options_provider=self.storage_options_provider,
|
|
)
|
|
|
|
def _dataset_exists(self) -> bool:
|
|
"""Whether a Lance dataset already exists at ``self.uri``.
|
|
|
|
A *successful open* is the authoritative existence signal, and it honors
|
|
``storage_options`` for every backend. We intentionally do not try to
|
|
classify failures (e.g. by matching Lance error strings, which drift
|
|
across versions): if the dataset can't be opened we report it as
|
|
not-existing and let the subsequent write surface any real error (such
|
|
as invalid ``storage_options``) with Lance's own message.
|
|
"""
|
|
try:
|
|
self._open_dataset()
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
def on_write_start(self, schema: Optional["pa.Schema"] = None) -> None:
|
|
_check_import(self, module="lance", package="pylance")
|
|
|
|
if self.mode == SaveMode.CREATE:
|
|
# CREATE must not clobber an existing dataset. Users who want to
|
|
# replace existing data should use SaveMode.OVERWRITE. Namespace
|
|
# writes manage table creation separately (the table location is
|
|
# declared/created up front), so skip the check in that case.
|
|
if self.table_id is None and self._dataset_exists():
|
|
raise ValueError(
|
|
f"Dataset at {self.uri} already exists. "
|
|
"Use mode=SaveMode.OVERWRITE to replace it, or "
|
|
"mode=SaveMode.APPEND to add to it."
|
|
)
|
|
elif self.mode == SaveMode.APPEND:
|
|
# APPEND needs the existing dataset's version/schema. Let Lance
|
|
# raise its own error (e.g. dataset not found) if it can't open.
|
|
ds = self._open_dataset()
|
|
self.read_version = ds.version
|
|
if self.schema is None:
|
|
self.schema = ds.schema
|
|
|
|
def on_write_complete(
|
|
self,
|
|
write_results: List[List[Tuple[str, str]]],
|
|
):
|
|
import warnings
|
|
|
|
import lance
|
|
|
|
if not write_results:
|
|
warnings.warn(
|
|
"write_results is empty.",
|
|
DeprecationWarning,
|
|
)
|
|
return
|
|
if hasattr(write_results, "write_returns"):
|
|
write_results = write_results.write_returns
|
|
|
|
if len(write_results) == 0:
|
|
warnings.warn(
|
|
"write results is empty. please check ray version or internal error",
|
|
DeprecationWarning,
|
|
)
|
|
return
|
|
|
|
fragments = []
|
|
schemas = []
|
|
for batch in write_results:
|
|
for fragment_str, schema_str in batch:
|
|
fragment = pickle.loads(fragment_str)
|
|
fragments.append(fragment)
|
|
schema = pickle.loads(schema_str)
|
|
if schema is not None:
|
|
schemas.append(schema)
|
|
# Skip commit when there are no fragments/schemas to commit.
|
|
if not schemas:
|
|
return
|
|
|
|
unified_schema = unify_schemas_with_validation(schemas)
|
|
if unified_schema is None:
|
|
return
|
|
|
|
if self.mode in {SaveMode.CREATE, SaveMode.OVERWRITE}:
|
|
op = lance.LanceOperation.Overwrite(unified_schema, fragments)
|
|
elif self.mode == SaveMode.APPEND:
|
|
op = lance.LanceOperation.Append(fragments)
|
|
lance.LanceDataset.commit(
|
|
self.uri,
|
|
op,
|
|
read_version=self.read_version,
|
|
storage_options=self.storage_options,
|
|
storage_options_provider=self.storage_options_provider,
|
|
)
|
|
|
|
|
|
class LanceDatasink(_BaseLanceDatasink):
|
|
"""Lance Ray Datasink.
|
|
|
|
Write a Ray dataset to lance.
|
|
|
|
If we expect to write larger-than-memory files,
|
|
we can use `LanceFragmentWriter` and `LanceCommitter`.
|
|
|
|
Args:
|
|
uri: The base URI of the dataset.
|
|
schema: The schema of the dataset.
|
|
mode: The write mode. Default is SaveMode.CREATE. Choices are
|
|
SaveMode.CREATE, SaveMode.APPEND, SaveMode.OVERWRITE. Namespace-backed
|
|
writes currently support only SaveMode.CREATE.
|
|
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.
|
|
When namespace params are provided, only SaveMode.CREATE is currently
|
|
supported.
|
|
*args: Additional positional arguments forwarded to the base class.
|
|
**kwargs: Additional keyword arguments forwarded to the base class.
|
|
"""
|
|
|
|
NAME = "Lance"
|
|
|
|
def __init__(
|
|
self,
|
|
uri: Optional[str] = None,
|
|
schema: Optional[pa.Schema] = None,
|
|
mode: SaveMode = SaveMode.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,
|
|
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,
|
|
)
|
|
|
|
self.min_rows_per_file = min_rows_per_file
|
|
self.max_rows_per_file = max_rows_per_file
|
|
self.data_storage_version = data_storage_version
|
|
# if mode is append, read_version is read from existing dataset.
|
|
self.read_version: Optional[int] = None
|
|
|
|
data_context = DataContext.get_current()
|
|
lance_config = data_context.lance_config
|
|
match = []
|
|
match.extend(lance_config.write_fragments_errors_to_retry)
|
|
match.extend(data_context.retried_io_errors)
|
|
self._retry_params = {
|
|
"description": _WRITE_LANCE_FRAGMENTS_DESCRIPTION,
|
|
"match": match,
|
|
"max_attempts": lance_config.write_fragments_max_attempts,
|
|
"max_backoff_s": lance_config.write_fragments_retry_max_backoff_s,
|
|
}
|
|
|
|
@property
|
|
def min_rows_per_write(self) -> int:
|
|
return self.min_rows_per_file
|
|
|
|
def get_name(self) -> str:
|
|
return self.NAME
|
|
|
|
def write(
|
|
self,
|
|
blocks: Iterable[Union[pa.Table, "pd.DataFrame"]],
|
|
_ctx,
|
|
):
|
|
fragments_and_schema = _write_fragment(
|
|
blocks,
|
|
self.uri,
|
|
schema=self.schema,
|
|
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))
|
|
for fragment, schema in fragments_and_schema
|
|
]
|