chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,427 @@
|
||||
"""
|
||||
TurbopufferDatasink - Ray Data datasink for Turbopuffer vector database
|
||||
|
||||
Implementation following the pattern of MongoDatasink and Daft's Turbopuffer sink.
|
||||
|
||||
This is based on [Turbopuffer Write API](https://turbopuffer.com/docs/write)
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Iterable, Literal, Optional, Union
|
||||
|
||||
import pyarrow as pa
|
||||
import pyarrow.compute as pc
|
||||
|
||||
from ray._common.retry import call_with_retry
|
||||
from ray.data._internal.execution.interfaces import TaskContext
|
||||
from ray.data._internal.planner.exchange.sort_task_spec import SortKey
|
||||
from ray.data._internal.util import _check_import
|
||||
from ray.data.block import Block, BlockAccessor
|
||||
from ray.data.datasource.datasink import Datasink
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import turbopuffer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Reserved column names for Turbopuffer
|
||||
_ID_COLUMN = "id"
|
||||
_VECTOR_COLUMN = "vector"
|
||||
TURBOPUFFER_API_KEY_ENV_VAR = "TURBOPUFFER_API_KEY"
|
||||
|
||||
|
||||
class TurbopufferDatasink(Datasink):
|
||||
"""Turbopuffer Ray Datasink.
|
||||
|
||||
A Ray :class:`~ray.data.datasource.Datasink` for writing data into the
|
||||
`Turbopuffer <https://turbopuffer.com/>`_ vector database.
|
||||
|
||||
Supports two modes of operation:
|
||||
|
||||
* **Single namespace** -- provide ``namespace`` to write all rows into one
|
||||
Turbopuffer namespace.
|
||||
* **Multi-namespace** -- provide ``namespace_column`` to route each row to
|
||||
the namespace whose name is stored in that column. The column is
|
||||
automatically dropped before the data is sent to Turbopuffer.
|
||||
|
||||
Exactly one of ``namespace`` or ``namespace_column`` must be supplied.
|
||||
|
||||
Args:
|
||||
namespace: Name of the Turbopuffer namespace to write into.
|
||||
Mutually exclusive with ``namespace_column``.
|
||||
namespace_column: Name of a column whose values determine the
|
||||
target namespace for each row. Rows are grouped by this column
|
||||
and each group is written to its corresponding namespace. The
|
||||
column is removed from the data before writing. Mutually
|
||||
exclusive with ``namespace``.
|
||||
region: Turbopuffer region identifier (for example,
|
||||
``"gcp-us-central1"``). Mutually exclusive with ``base_url``.
|
||||
Exactly one of ``region`` or ``base_url`` must be supplied.
|
||||
base_url: Base URL for the Turbopuffer API (for example,
|
||||
``"https://gcp-us-central1.turbopuffer.com"``). Mutually
|
||||
exclusive with ``region``. Exactly one of ``region`` or
|
||||
``base_url`` must be supplied.
|
||||
api_key: Turbopuffer API key. If omitted, the value is read from the
|
||||
``TURBOPUFFER_API_KEY`` environment variable.
|
||||
schema: Optional Turbopuffer schema definition to pass along with
|
||||
writes. If provided, it is forwarded as the ``schema`` argument
|
||||
to ``namespace.write``.
|
||||
id_column: Name of the column to treat as the document identifier.
|
||||
Rows with null IDs are dropped before writing. Defaults to ``"id"``.
|
||||
vector_column: Name of the column containing embedding vectors.
|
||||
If this differs from ``"vector"``, it is renamed to ``"vector"``
|
||||
before writing. Defaults to ``"vector"``.
|
||||
batch_size: Maximum number of rows to include in a single Turbopuffer
|
||||
write call (logical row batching; subject to Turbopuffer's
|
||||
256MiB request-size limit). Defaults to ``10000``.
|
||||
distance_metric: Distance metric for the namespace. Passed to
|
||||
``namespace.write`` as the ``distance_metric`` argument.
|
||||
Defaults to ``"cosine_distance"``.
|
||||
concurrency: Unused; Ray Data controls write parallelism via
|
||||
:meth:`~ray.data.Dataset.write_datasink` ``concurrency``.
|
||||
|
||||
Examples:
|
||||
Write to a single namespace using a region:
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
import ray
|
||||
from ray.data._internal.datasource.turbopuffer_datasink import (
|
||||
TurbopufferDatasink,
|
||||
)
|
||||
|
||||
ds = ray.data.range(100)
|
||||
ds = ds.map_batches(lambda batch: {"id": batch["id"], "vector": ...})
|
||||
|
||||
ds.write_datasink(
|
||||
TurbopufferDatasink(
|
||||
namespace="my-namespace",
|
||||
api_key="<YOUR_API_KEY>",
|
||||
region="gcp-us-central1",
|
||||
)
|
||||
)
|
||||
|
||||
Write using a base URL instead of a region:
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
ds.write_datasink(
|
||||
TurbopufferDatasink(
|
||||
namespace="my-namespace",
|
||||
api_key="<YOUR_API_KEY>",
|
||||
base_url="https://gcp-us-central1.turbopuffer.com",
|
||||
)
|
||||
)
|
||||
|
||||
Write to multiple namespaces driven by a column:
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
ds.write_datasink(
|
||||
TurbopufferDatasink(
|
||||
namespace_column="tenant",
|
||||
api_key="<YOUR_API_KEY>",
|
||||
region="gcp-us-central1",
|
||||
)
|
||||
)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
namespace: Optional[str] = None,
|
||||
*,
|
||||
namespace_column: Optional[str] = None,
|
||||
region: Optional[str] = None,
|
||||
base_url: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
schema: Optional[dict] = None,
|
||||
id_column: str = "id",
|
||||
vector_column: str = "vector",
|
||||
batch_size: int = 10000,
|
||||
distance_metric: Literal[
|
||||
"cosine_distance", "euclidean_distance"
|
||||
] = "cosine_distance",
|
||||
concurrency: Optional[int] = None,
|
||||
):
|
||||
_check_import(self, module="turbopuffer", package="turbopuffer")
|
||||
|
||||
# Validate namespace / namespace_column mutual exclusivity.
|
||||
if namespace and namespace_column:
|
||||
raise ValueError(
|
||||
"Specify exactly one of 'namespace' or 'namespace_column', " "not both."
|
||||
)
|
||||
if not namespace and not namespace_column:
|
||||
raise ValueError(
|
||||
"Either 'namespace' or 'namespace_column' must be provided."
|
||||
)
|
||||
|
||||
# Validate region / base_url mutual exclusivity.
|
||||
if region is not None and base_url is not None:
|
||||
raise ValueError("Specify exactly one of 'region' or 'base_url', not both.")
|
||||
if region is None and base_url is None:
|
||||
raise ValueError("Either 'region' or 'base_url' must be provided.")
|
||||
|
||||
# Store configuration
|
||||
self.namespace = namespace
|
||||
self.namespace_column = namespace_column
|
||||
self.api_key = api_key or os.getenv(TURBOPUFFER_API_KEY_ENV_VAR)
|
||||
self.region = region
|
||||
self.base_url = base_url
|
||||
self.schema = schema
|
||||
self.id_column = id_column
|
||||
self.vector_column = vector_column
|
||||
self.batch_size = batch_size
|
||||
self.distance_metric = distance_metric
|
||||
|
||||
# Validate column configuration
|
||||
if self.id_column == self.vector_column:
|
||||
raise ValueError(
|
||||
"id_column and vector_column refer to the same column "
|
||||
f"'{self.id_column}'. They must be distinct."
|
||||
)
|
||||
|
||||
if self.namespace_column and self.namespace_column in (
|
||||
self.id_column,
|
||||
self.vector_column,
|
||||
):
|
||||
raise ValueError(
|
||||
f"namespace_column '{self.namespace_column}' must not be the "
|
||||
f"same as id_column ('{self.id_column}') or vector_column "
|
||||
f"('{self.vector_column}')."
|
||||
)
|
||||
|
||||
# Validate API key
|
||||
if not self.api_key:
|
||||
raise ValueError(
|
||||
"API key is required. Provide via api_key parameter or "
|
||||
"TURBOPUFFER_API_KEY environment variable"
|
||||
)
|
||||
|
||||
# Initialize client
|
||||
self._client = None
|
||||
|
||||
def __getstate__(self) -> dict:
|
||||
"""Exclude `_client` during pickling."""
|
||||
state = self.__dict__.copy()
|
||||
state.pop("_client", None)
|
||||
return state
|
||||
|
||||
def __setstate__(self, state: dict) -> None:
|
||||
self.__dict__.update(state)
|
||||
self._client = None
|
||||
|
||||
def _get_client(self):
|
||||
"""Lazy initialize Turbopuffer client."""
|
||||
if self._client is None:
|
||||
import turbopuffer
|
||||
|
||||
kwargs = {"api_key": self.api_key}
|
||||
if self.region is not None:
|
||||
kwargs["region"] = self.region
|
||||
else:
|
||||
kwargs["base_url"] = self.base_url
|
||||
self._client = turbopuffer.Turbopuffer(**kwargs)
|
||||
return self._client
|
||||
|
||||
def write(
|
||||
self,
|
||||
blocks: Iterable[Block],
|
||||
ctx: TaskContext,
|
||||
) -> None:
|
||||
"""
|
||||
Write blocks to Turbopuffer in a streaming fashion.
|
||||
|
||||
For memory efficiency, blocks are processed one at a time rather than
|
||||
concatenating all blocks into a single large table. This follows the
|
||||
pattern used by ClickHouseDatasink.
|
||||
|
||||
Each block is prepared (columns renamed, null IDs filtered), then
|
||||
written in batches of ``batch_size``.
|
||||
|
||||
When ``namespace_column`` is set, each block is grouped by the
|
||||
namespace column and each group is written to its corresponding
|
||||
Turbopuffer namespace.
|
||||
"""
|
||||
client = self._get_client()
|
||||
|
||||
for block in blocks:
|
||||
accessor = BlockAccessor.for_block(block)
|
||||
table = accessor.to_arrow()
|
||||
|
||||
if table.num_rows == 0:
|
||||
continue
|
||||
|
||||
if self.namespace_column:
|
||||
# Multi-namespace: group by namespace column, write to each.
|
||||
self._write_multi_namespace(client, table)
|
||||
else:
|
||||
# Single namespace.
|
||||
table = self._prepare_arrow_table(table)
|
||||
|
||||
if table.num_rows == 0:
|
||||
continue
|
||||
|
||||
ns = client.namespace(self.namespace)
|
||||
for batch in table.to_batches(max_chunksize=self.batch_size):
|
||||
self._write_batch_with_retry(ns, batch, self.namespace)
|
||||
|
||||
def _rename_column_if_needed(
|
||||
self,
|
||||
table: pa.Table,
|
||||
source_column: str,
|
||||
target_column: str,
|
||||
column_type: str,
|
||||
) -> pa.Table:
|
||||
"""
|
||||
Rename a column in the table if it differs from the target name.
|
||||
|
||||
Args:
|
||||
table: The Arrow table to modify.
|
||||
source_column: The current column name in the table.
|
||||
target_column: The required column name for Turbopuffer.
|
||||
column_type: Human-readable type for error messages (e.g., "ID", "Vector").
|
||||
|
||||
Returns:
|
||||
The table with the column renamed, or the original table if no rename needed.
|
||||
|
||||
Raises:
|
||||
ValueError: If source column is missing or target column already exists.
|
||||
"""
|
||||
if source_column not in table.column_names:
|
||||
raise ValueError(
|
||||
f"{column_type} column '{source_column}' not found in table"
|
||||
)
|
||||
|
||||
# No rename needed if source and target are the same
|
||||
if source_column == target_column:
|
||||
return table
|
||||
|
||||
if target_column in table.column_names:
|
||||
raise ValueError(
|
||||
f"Table already has a '{target_column}' column; cannot also rename "
|
||||
f"'{source_column}' to '{target_column}'. Please disambiguate your schema."
|
||||
)
|
||||
|
||||
return BlockAccessor.for_block(table).rename_columns(
|
||||
{source_column: target_column}
|
||||
)
|
||||
|
||||
def _prepare_arrow_table(self, table: pa.Table) -> pa.Table:
|
||||
"""
|
||||
Prepare Arrow table for Turbopuffer write.
|
||||
|
||||
1. Rename ID column to "id" if needed
|
||||
2. Rename vector column to "vector" if needed
|
||||
3. Filter out rows with null IDs
|
||||
"""
|
||||
table = self._rename_column_if_needed(table, self.id_column, _ID_COLUMN, "ID")
|
||||
table = self._rename_column_if_needed(
|
||||
table, self.vector_column, _VECTOR_COLUMN, "Vector"
|
||||
)
|
||||
|
||||
# Filter out rows with null IDs
|
||||
if _ID_COLUMN in table.column_names:
|
||||
table = table.filter(pc.is_valid(table.column(_ID_COLUMN)))
|
||||
|
||||
return table
|
||||
|
||||
def _write_multi_namespace(
|
||||
self, client: "turbopuffer.Turbopuffer", table: pa.Table
|
||||
) -> None:
|
||||
"""Group rows by ``namespace_column`` and write each group to its namespace.
|
||||
|
||||
Uses :meth:`BlockAccessor._iter_groups_sorted` for efficient
|
||||
zero-copy slicing by group.
|
||||
"""
|
||||
group_col_name = self.namespace_column
|
||||
|
||||
if group_col_name not in table.column_names:
|
||||
raise ValueError(
|
||||
f"Namespace column '{group_col_name}' not found in table. "
|
||||
f"Available columns: {table.column_names}"
|
||||
)
|
||||
|
||||
# Reject null namespace values early -- we cannot route them.
|
||||
ns_col = table.column(group_col_name)
|
||||
if pc.any(pc.is_null(ns_col)).as_py():
|
||||
raise ValueError(
|
||||
f"Namespace column '{group_col_name}' contains null values; "
|
||||
"fill or drop them before writing with namespace_column."
|
||||
)
|
||||
|
||||
# Sort by the namespace column so _iter_groups_sorted can yield
|
||||
# contiguous zero-copy slices for each unique namespace value.
|
||||
sort_key = SortKey(key=group_col_name, descending=False)
|
||||
block_accessor = BlockAccessor.for_block(
|
||||
BlockAccessor.for_block(table).sort(sort_key)
|
||||
)
|
||||
|
||||
for (namespace_name,), group_table in block_accessor._iter_groups_sorted(
|
||||
sort_key
|
||||
):
|
||||
# Drop the namespace column -- it is routing metadata, not data.
|
||||
group_table = group_table.drop(group_col_name)
|
||||
|
||||
# Prepare (rename id/vector columns, filter null IDs).
|
||||
group_table = self._prepare_arrow_table(group_table)
|
||||
if group_table.num_rows == 0:
|
||||
continue
|
||||
|
||||
ns = client.namespace(namespace_name)
|
||||
for batch in group_table.to_batches(max_chunksize=self.batch_size):
|
||||
self._write_batch_with_retry(ns, batch, namespace_name)
|
||||
|
||||
def _transform_to_turbopuffer_format(
|
||||
self, table: Union[pa.Table, pa.RecordBatch]
|
||||
) -> dict:
|
||||
if _ID_COLUMN not in table.column_names:
|
||||
raise ValueError(f"Table must have '{_ID_COLUMN}' column")
|
||||
|
||||
# Cast 16-byte binary ID column to native UUID type for Turbopuffer performance.
|
||||
# Native UUIDs are 16 bytes vs 36 bytes for string-encoded UUIDs.
|
||||
# See: https://turbopuffer.com/docs/performance
|
||||
id_col = table.column(_ID_COLUMN)
|
||||
if pa.types.is_fixed_size_binary(id_col.type) and id_col.type.byte_width == 16:
|
||||
# Cast fixed_size_binary(16) to uuid type
|
||||
uuid_col = id_col.cast(pa.uuid())
|
||||
table = table.set_column(
|
||||
table.schema.get_field_index(_ID_COLUMN), _ID_COLUMN, uuid_col
|
||||
)
|
||||
|
||||
# to_pydict() on UuidArray automatically returns uuid.UUID objects
|
||||
return table.to_pydict()
|
||||
|
||||
def _write_batch_with_retry(
|
||||
self,
|
||||
namespace: "turbopuffer.Namespace",
|
||||
batch: pa.Table,
|
||||
namespace_name: Optional[str] = None,
|
||||
):
|
||||
"""Write a single batch with exponential backoff retry.
|
||||
|
||||
Args:
|
||||
namespace: The Turbopuffer namespace object to write to.
|
||||
batch: Arrow table or record-batch to write.
|
||||
namespace_name: Human-readable namespace name for log messages.
|
||||
Falls back to ``self.namespace`` when not provided.
|
||||
"""
|
||||
ns_label = namespace_name or self.namespace
|
||||
try:
|
||||
batch_data = self._transform_to_turbopuffer_format(batch)
|
||||
call_with_retry(
|
||||
lambda: namespace.write(
|
||||
upsert_columns=batch_data,
|
||||
schema=self.schema,
|
||||
distance_metric=self.distance_metric,
|
||||
),
|
||||
description=f"write batch to namespace '{ns_label}'",
|
||||
max_attempts=5,
|
||||
max_backoff_s=32,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Write failed for namespace '{ns_label}' after retries: {e}")
|
||||
raise
|
||||
Reference in New Issue
Block a user