wehub-resource-sync
428 lines
16 KiB
Python
428 lines
16 KiB
Python
"""
|
|
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
|