lancedb-ray 0.1.0__py3-none-any.whl

lancedb_ray/__init__.py

@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+"""lancedb-ray: Ray Data integration for LanceDB and LanceDB Enterprise.
+
+Read and write LanceDB tables as Ray Datasets, using the most parallel strategy
+each backend supports:
+
+- **Local / OSS** tables are backed by a Lance dataset, so reads run one task
+  per fragment group and appends write fragments in parallel that the driver
+  commits as a single atomic transaction.
+- **Cloud / Enterprise** (``db://``) tables are a remote service with no
+  fragment access, so reads shard the row space across tasks against a pinned
+  table version and writes fan out batched requests.
+
+Example:
+    >>> import lancedb_ray as ldbr
+    >>> ds = ldbr.read_lancedb("my_table", uri="/data/lancedb")  # doctest: +SKIP
+    >>> ldbr.write_lancedb(ds, "copy", uri="/data/lancedb", mode="create")  # doctest: +SKIP
+"""
+
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _dist_version
+
+from ._plan import OffsetRange
+from ._retry import RetryPolicy
+from .connection import LanceDBConnectionSpec
+from .datasink import LanceDBDatasink, WriteStats
+from .datasource import LanceDBDatasource
+from .io import read_lancedb, write_lancedb
+
+# Read from the installed distribution rather than hard-coded here, so the
+# version release-please writes into pyproject.toml is the only one there is.
+try:
+    __version__ = _dist_version("lancedb-ray")
+except PackageNotFoundError:  # pragma: no cover - only in an uninstalled tree
+    __version__ = "0.0.0.dev0"
+
+__all__ = [
+    "LanceDBConnectionSpec",
+    "LanceDBDatasink",
+    "LanceDBDatasource",
+    "OffsetRange",
+    "RetryPolicy",
+    "WriteStats",
+    "__version__",
+    "read_lancedb",
+    "write_lancedb",
+]

lancedb_ray/_plan.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pure planning helpers for distributed LanceDB reads and writes.
+
+Nothing in this module touches Ray, the network, or the filesystem. Keeping the
+planning arithmetic free of I/O is what makes it cheap to test exhaustively --
+the interesting edge cases (empty tables, parallelism exceeding row counts,
+uneven remainders) are all covered by ``tests/test_plan.py``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from typing import NamedTuple
+
+__all__ = [
+    "OffsetRange",
+    "chunk_offsets",
+    "plan_offset_shards",
+    "rows_within_byte_budget",
+    "split_arrow_table",
+]
+
+
+class OffsetRange(NamedTuple):
+    """A half-open ``[start, end)`` range of row offsets within a table."""
+
+    start: int
+    end: int
+
+    @property
+    def num_rows(self) -> int:
+        return self.end - self.start
+
+
+def plan_offset_shards(num_rows: int, parallelism: int) -> list[OffsetRange]:
+    """Split ``[0, num_rows)`` into at most ``parallelism`` contiguous shards.
+
+    Rows are distributed as evenly as possible: the first ``num_rows %
+    parallelism`` shards receive one extra row. Empty shards are never
+    produced, so the result is empty when ``num_rows`` is zero and has
+    ``num_rows`` single-row entries when ``parallelism`` exceeds ``num_rows``.
+
+    Args:
+        num_rows: Total number of rows in the table. Must not be negative.
+        parallelism: Desired number of shards. Must be positive.
+
+    Returns:
+        Contiguous, non-overlapping ranges covering ``[0, num_rows)`` in order.
+    """
+    if num_rows < 0:
+        raise ValueError(f"num_rows must not be negative, got {num_rows}")
+    if parallelism <= 0:
+        raise ValueError(f"parallelism must be positive, got {parallelism}")
+
+    if num_rows == 0:
+        return []
+
+    num_shards = min(parallelism, num_rows)
+    base, remainder = divmod(num_rows, num_shards)
+
+    shards: list[OffsetRange] = []
+    start = 0
+    for i in range(num_shards):
+        size = base + (1 if i < remainder else 0)
+        shards.append(OffsetRange(start, start + size))
+        start += size
+
+    # The loop above is exact arithmetic; this guards against future edits.
+    assert start == num_rows, f"shards covered {start} of {num_rows} rows"
+    return shards
+
+
+def chunk_offsets(offsets: OffsetRange, batch_size: int) -> Iterator[list[int]]:
+    """Materialise an offset range as batches of explicit row offsets.
+
+    ``Table.take_offsets`` requires an explicit list of offsets rather than a
+    range. For a large shard that list can be enormous, so it is generated
+    lazily inside the worker -- only the two integers of the ``OffsetRange``
+    ever cross the wire -- and yielded in ``batch_size`` chunks so no single
+    request carries the whole shard.
+
+    Args:
+        offsets: The half-open range to expand.
+        batch_size: Maximum number of offsets per yielded chunk. Must be positive.
+
+    Yields:
+        Lists of row offsets, each of length ``batch_size`` except possibly the last.
+    """
+    if batch_size <= 0:
+        raise ValueError(f"batch_size must be positive, got {batch_size}")
+
+    for start in range(offsets.start, offsets.end, batch_size):
+        end = min(start + batch_size, offsets.end)
+        yield list(range(start, end))
+
+
+def split_arrow_table(num_rows: int, max_rows: int) -> list[OffsetRange]:
+    """Split ``num_rows`` into consecutive slices of at most ``max_rows``.
+
+    Used by the write path to break an accumulated batch into request-sized
+    pieces before handing them to LanceDB.
+    """
+    if num_rows < 0:
+        raise ValueError(f"num_rows must not be negative, got {num_rows}")
+    if max_rows <= 0:
+        raise ValueError(f"max_rows must be positive, got {max_rows}")
+
+    return [
+        OffsetRange(start, min(start + max_rows, num_rows))
+        for start in range(0, num_rows, max_rows)
+    ]
+
+
+def rows_within_byte_budget(num_rows: int, nbytes: int, max_bytes: int) -> int:
+    """Largest slice of a batch whose share of ``nbytes`` fits in ``max_bytes``.
+
+    A row count does not bound memory. One row of a 1536-dimension float32
+    embedding is 6KB and one row carrying a JPEG is unbounded, so a schema wide
+    enough exhausts a worker long before any row limit is reached. Sizing a
+    slice by the batch's own average row width is what makes the ceiling
+    actually about bytes.
+
+    The estimate is uniform across rows, so a batch with a few outlying rows can
+    still overshoot -- this is a ceiling on the common case, not a hard bound.
+
+    Args:
+        num_rows: Rows in the batch being split.
+        nbytes: In-memory size of that batch.
+        max_bytes: Byte budget for one slice. Must be positive.
+
+    Returns:
+        A row count in ``[1, num_rows]``, or ``0`` when the batch is empty.
+        Never zero for a non-empty batch: a single row that exceeds the budget
+        on its own still has to be written, and refusing it would deadlock the
+        write rather than bound it.
+    """
+    if num_rows < 0:
+        raise ValueError(f"num_rows must not be negative, got {num_rows}")
+    if nbytes < 0:
+        raise ValueError(f"nbytes must not be negative, got {nbytes}")
+    if max_bytes <= 0:
+        raise ValueError(f"max_bytes must be positive, got {max_bytes}")
+
+    if num_rows == 0 or nbytes <= max_bytes:
+        return num_rows
+
+    # Integer arithmetic throughout: at these magnitudes a float ratio can
+    # round a slice up past the budget it was supposed to enforce.
+    return max(1, (max_bytes * num_rows) // nbytes)

lancedb_ray/_retry.py

@@ -0,0 +1,196 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Retry helpers for transient LanceDB failures.
+
+Remote (Cloud/Enterprise) calls go over HTTP and fail transiently; local writes
+contend on the dataset commit lock. Both are worth retrying with backoff, and
+neither should be retried forever.
+"""
+
+from __future__ import annotations
+
+import logging
+import random
+import re
+import time
+from collections.abc import Callable
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["RetryPolicy", "call_with_retry", "is_commit_conflict", "is_transient"]
+
+# Substrings identifying errors that are worth another attempt. Matched against
+# the lowercased ``str()`` of the exception because the LanceDB Python SDK
+# surfaces server-side failures as generic exception types with descriptive
+# messages rather than a dedicated exception hierarchy.
+_TRANSIENT_MARKERS = (
+    "timeout",
+    "timed out",
+    "connection reset",
+    "connection aborted",
+    "connection refused",
+    "broken pipe",
+    "temporarily unavailable",
+    "service unavailable",
+    "too many requests",
+    "internal server error",
+    "bad gateway",
+    "gateway timeout",
+)
+
+
+def _status_pattern(*codes: str) -> re.Pattern[str]:
+    """Match an HTTP status code as a whole number, not as digits anywhere.
+
+    A bare substring test reads ``429`` out of "dimension 1429" and ``503`` out
+    of "row 8503", classifying a deterministic schema error as retryable. Word
+    boundaries keep a code from matching inside a longer number, which is how
+    row counts, dimensions and byte sizes reach these messages.
+    """
+    return re.compile(rf"\b(?:{'|'.join(codes)})\b")
+
+
+_TRANSIENT_STATUS = _status_pattern("429", "502", "503", "504")
+
+_COMMIT_CONFLICT_MARKERS = (
+    "commit conflict",
+    "concurrent",
+    "version already exists",
+    "retryable commit",
+    "commit was rejected",
+)
+
+
+def is_transient(error: BaseException) -> bool:
+    """Return whether ``error`` looks like a retryable transient failure."""
+    message = str(error).lower()
+    if any(marker in message for marker in _TRANSIENT_MARKERS):
+        return True
+    return _TRANSIENT_STATUS.search(message) is not None
+
+
+#: Failures that mean the request never reached the service, or was rejected
+#: without being applied. Re-sending these cannot duplicate anything.
+_NOT_APPLIED_MARKERS = (
+    "connection refused",
+    "temporarily unavailable",
+    "service unavailable",
+    "too many requests",
+    "name or service not known",
+    "nodename nor servname",
+    "failed to resolve",
+)
+
+_NOT_APPLIED_STATUS = _status_pattern("429", "503")
+
+
+def is_definitely_not_applied(error: BaseException) -> bool:
+    """Whether ``error`` proves the write never took effect.
+
+    A read timeout or a dropped connection is ambiguous: the service may have
+    committed and only the response was lost. Re-sending a non-idempotent
+    append in that case silently duplicates rows, so appends retry only on
+    this narrower class.
+    """
+    message = str(error).lower()
+    if any(marker in message for marker in _NOT_APPLIED_MARKERS):
+        return True
+    return _NOT_APPLIED_STATUS.search(message) is not None
+
+
+#: Arrow-rs refuses to import a buffer whose pointer is not aligned for its
+#: scalar type. Ray hands out zero-copy views into its object store, and a
+#: view can land unaligned for a type that needs more than 8 bytes, such as
+#: decimal128. The import fails before anything is written.
+_ALIGNMENT_MARKERS = ("is not aligned with the specified scalar type",)
+
+
+def is_arrow_alignment_error(error: BaseException) -> bool:
+    """Whether ``error`` is arrow-rs rejecting an unaligned FFI buffer.
+
+    Recoverable by copying the batch into freshly allocated memory, and safe
+    to retry: the import fails before any data is committed.
+    """
+    message = str(error)
+    return any(marker in message for marker in _ALIGNMENT_MARKERS)
+
+
+def is_commit_conflict(error: BaseException) -> bool:
+    """Return whether ``error`` looks like a losing race on a dataset commit."""
+    message = str(error).lower()
+    return any(marker in message for marker in _COMMIT_CONFLICT_MARKERS)
+
+
+class RetryPolicy:
+    """Exponential backoff with full jitter.
+
+    Args:
+        max_attempts: Total attempts including the first. ``1`` disables retrying.
+        initial_backoff_s: Backoff before the second attempt.
+        max_backoff_s: Ceiling on the backoff between attempts.
+        predicate: Returns whether a given exception should be retried.
+    """
+
+    def __init__(
+        self,
+        *,
+        max_attempts: int = 5,
+        initial_backoff_s: float = 0.5,
+        max_backoff_s: float = 32.0,
+        predicate: Callable[[BaseException], bool] = is_transient,
+    ) -> None:
+        if max_attempts < 1:
+            raise ValueError(f"max_attempts must be at least 1, got {max_attempts}")
+        self.max_attempts = max_attempts
+        self.initial_backoff_s = initial_backoff_s
+        self.max_backoff_s = max_backoff_s
+        self.predicate = predicate
+
+    def backoff_for(self, attempt: int) -> float:
+        """Return the sleep, in seconds, before attempt number ``attempt`` (1-based)."""
+        uncapped = self.initial_backoff_s * (2 ** (attempt - 1))
+        return random.uniform(0.0, min(uncapped, self.max_backoff_s))
+
+
+def call_with_retry[T](
+    fn: Callable[[], T],
+    policy: RetryPolicy,
+    *,
+    description: str,
+    sleep: Callable[[float], None] = time.sleep,
+) -> T:
+    """Call ``fn``, retrying while ``policy.predicate`` accepts the raised error.
+
+    Args:
+        fn: Zero-argument callable to invoke.
+        policy: Attempt count, backoff schedule and retry predicate.
+        description: Human-readable operation name used in log messages.
+        sleep: Injectable sleep, so tests do not spend real time backing off.
+
+    Returns:
+        Whatever ``fn`` returns on its first successful attempt.
+
+    Raises:
+        BaseException: The final error, once attempts are exhausted or the
+            predicate rejects it.
+    """
+    last_error: BaseException | None = None
+    for attempt in range(1, policy.max_attempts + 1):
+        try:
+            return fn()
+        except Exception as error:  # noqa: BLE001 - re-raised below
+            last_error = error
+            if attempt == policy.max_attempts or not policy.predicate(error):
+                raise
+            backoff = policy.backoff_for(attempt)
+            logger.warning(
+                "%s failed (attempt %d/%d): %s. Retrying in %.2fs.",
+                description,
+                attempt,
+                policy.max_attempts,
+                error,
+                backoff,
+            )
+            sleep(backoff)
+
+    # Unreachable: the loop either returns or raises.
+    raise AssertionError(f"retry loop exited without result: {last_error}")