From 3f0e4efb51946623b0871daf91b662c580db398f Mon Sep 17 00:00:00 2001 From: jiangxt2 Date: Wed, 19 Aug 2026 22:54:52 +0800 Subject: [PATCH] feat(inference): normalize Ray tensor embeddings for Lance Signed-off-by: jiangxt2 --- docs/how-to/inference.md | 19 +- docs/reference/support-matrix.md | 2 +- src/tributo/inference/contracts.py | 6 +- src/tributo/integrations/sinks/lance.py | 134 +++++++- tests/inference/test_lance_sink.py | 315 ++++++++++++++++++ tests/integrations/test_lance_vector_index.py | 82 ++++- 6 files changed, 519 insertions(+), 39 deletions(-) diff --git a/docs/how-to/inference.md b/docs/how-to/inference.md index e76bff0..645c77f 100644 --- a/docs/how-to/inference.md +++ b/docs/how-to/inference.md @@ -134,11 +134,20 @@ contract is approved. `output_format="lance"` selects a generic Lance ResultSink. Declare `output_vector_columns` when fixed-size floating-point vector validation is -required. The sink builds a credential-safe `WriteRequest` and selects the -stable `tributo.ray.lance` Binding. That Binding currently delegates all -data-plane work to `lance_ray.write_lance(stream=False)` and records the schema -fingerprint in the sink receipt; it does not probe a post-write dataset version -or build an ANN index. A Lance table may contain ordinary columns as well as +required. For declared columns, the sink accepts an exact `FixedSizeList`, a +fixed-shape Ray V1/V2 tensor, or a one-dimensional Arrow fixed-shape tensor. +The declared dimension and dtype must match exactly; the sink does not flatten +higher-rank tensors or perform dtype conversion. It normalizes accepted tensor +columns to Lance `FixedSizeList` columns in the existing distributed Arrow +validation step. Only declared vector columns are converted; if another column +in a worker batch differs from the driver target schema, the sink fails instead +of coercing that column. + +The sink builds a credential-safe `WriteRequest` and selects the stable +`tributo.ray.lance` Binding. That Binding delegates the physical write to +`lance_ray.write_lance(stream=False)`. The sink receipt records the normalized +schema fingerprint; the sink does not probe a post-write dataset version or +build an ANN index. A Lance table may contain ordinary columns as well as user-produced vectors. Direct `LanceResultSinkRequest` construction defaults to provider-native `create`. diff --git a/docs/reference/support-matrix.md b/docs/reference/support-matrix.md index a50cdb2..e59ad47 100644 --- a/docs/reference/support-matrix.md +++ b/docs/reference/support-matrix.md @@ -49,7 +49,7 @@ compatible profile, while the generated `Validated profiles` column remains | Doris reads | Adapter only | Ray routes use locked `ray-doris==1.0`; Daft routes use locked `daft-doris==1.0`; real-database Conformance is still required and tablet planning remains provider/binding-owned | | ORC and Hive external-table reads | Not implemented | Locked Ray/Daft versions expose no validated public reader | | Third-party ingestion Provider/Binding SPI | Implemented | Installed packages use `tributo.ingestion_providers` plus `tributo.ingestion_bindings`; bad plugins are isolated, duplicate routes never replace built-ins, and Binding selection can constrain filesystem, catalog, and storage format | -| Lance output | Implemented as a generic ResultSink path | User Predictor owns vector semantics; the sink does not pool, normalize, or automatically invoke the separate vector-index workflow | +| Lance output | Implemented as a generic ResultSink path | Declared fixed-shape Ray/Arrow tensor columns are strictly validated and normalized to Lance `FixedSizeList`; the user Predictor still owns vector semantics, and the sink does not pool, mathematically normalize embedding values, change dtype, or automatically invoke the separate vector-index workflow | | Native bounded writes | Basic local/S3 native round-trips are implemented through `WriteGateway` for Ray/Daft Parquet, CSV, Iceberg, and Lance. Ray Lance delegates to locked `lance-ray==0.5.0`/PyLance 9; Daft delegates to `DataFrame.write_lance`. Mode support comes from the selected native Binding capability. Existing-target `CREATE`, missing-target `APPEND`, schema evolution, and empty writes remain provider-owned, are not Tributo guarantees, and are outside the current Gate | Tributo owns control-plane validation only; Ray Data, Daft, or an official native integration owns data-plane writes | | Custom Hugging Face Predictor | User-provided Ray Data/Jobs extension point | Tokenization, task semantics, output interpretation, pooling, normalization, and metadata remain user-owned | | Database inference sinks | Extension point | No built-in ClickHouse or Doris sink | diff --git a/src/tributo/inference/contracts.py b/src/tributo/inference/contracts.py index ba492d2..318ccdb 100644 --- a/src/tributo/inference/contracts.py +++ b/src/tributo/inference/contracts.py @@ -296,8 +296,10 @@ class LanceVectorColumnSpec(_FrozenContract): class LanceResultSinkRequest(_FrozenContract): """Configuration for explicit Lance result materialization. - Tributo validates only the Arrow schema contract declared here. It does - not infer model task semantics, pooling, normalization, or vector metadata. + Tributo validates the Arrow schema contract declared here and normalizes + supported fixed-shape Ray or Arrow tensor columns to Lance fixed-size list + columns. It does not infer model task semantics, pooling, normalization, + dtype conversion, or vector metadata. Direct requests default to provider-native ``create``. The legacy ``InferenceConfig`` pipeline separately retains its historical ``overwrite`` default for compatibility. Tributo does not add stricter diff --git a/src/tributo/integrations/sinks/lance.py b/src/tributo/integrations/sinks/lance.py index f2e5699..a9455df 100644 --- a/src/tributo/integrations/sinks/lance.py +++ b/src/tributo/integrations/sinks/lance.py @@ -30,6 +30,13 @@ logger = logging.getLogger(__name__) +_RAY_FIXED_SHAPE_TENSOR_EXTENSION_NAMES = frozenset( + { + "ray.data.arrow_tensor", + "ray.data.arrow_tensor_v2", + } +) + class _StorageProfileResolverLike(Protocol): def resolve(self, profile: str | None) -> StorageProfile: ... @@ -40,12 +47,13 @@ class LanceResultSink: """Write a Ray Dataset through the shared native write Gateway. The sink is deliberately explicit: it always writes Lance, regardless of - whether the output contains a vector column. Vector schema checks apply - only to columns declared by ``request.vector_columns``; model semantics - remain the responsibility of the caller's Predictor. The Gateway selects - the stable Ray Lance Binding; the selected provider owns all data-plane and - save-mode behavior. All callers use this same boundary; the sink does not - expose a format-specific compatibility facade. + whether the output contains a vector column. Vector schema checks and + supported fixed-shape tensor normalization apply only to columns declared + by ``request.vector_columns``; model semantics remain the responsibility + of the caller's Predictor. The Gateway selects the stable Ray Lance + Binding; the selected provider owns physical writes and save-mode behavior. + All callers use this same boundary; the sink does not expose a format-specific + compatibility facade. """ api_version: ClassVar[int] = 1 @@ -64,7 +72,7 @@ def write( run_id: str, plan_digest: str, ) -> ResultSinkReceipt: - """Validate the declared Arrow schema and materialize a Lance table.""" + """Normalize the declared Arrow vectors and materialize a Lance table.""" from tributo.inference.contracts import ( LanceResultSinkRequest, ResultSinkReceipt, @@ -75,14 +83,18 @@ def write( f"Lance result sink cannot write {request.sink_id!r}" ) arrow_schema = _arrow_schema(dataset.schema()) - _validate_vector_schema(arrow_schema, request) + target_schema = _canonical_vector_schema(arrow_schema, request) runtime_s3 = _runtime_s3( self._storage_resolver, request.storage_profile, request.uri ) try: if request.vector_columns: dataset = dataset.map_batches( - partial(validate_vector_batch, request=request), + partial( + _normalize_vector_batch, + request=request, + target_schema=target_schema, + ), batch_format="pyarrow", ) options: dict[str, Any] = { @@ -125,7 +137,7 @@ def write( source_error_type or type(exc).__name__ ) from None - fingerprint = schema_fingerprint(arrow_schema) + fingerprint = schema_fingerprint(target_schema) result_id = _result_id( run_id=run_id, plan_digest=plan_digest, @@ -172,29 +184,115 @@ def _runtime_s3( return profile -def _validate_vector_schema(schema: pa.Schema, request: LanceResultSinkRequest) -> None: +def _canonical_vector_schema( + schema: pa.Schema, request: LanceResultSinkRequest +) -> pa.Schema: + fields = list(schema) for spec in request.vector_columns: if spec.name not in schema.names: raise ResultWriteError( f"Lance vector column {spec.name!r} is missing from the output schema" ) + field_index = schema.get_field_index(spec.name) field = schema.field(spec.name) data_type = field.type - if not pa.types.is_fixed_size_list(data_type): + vector_type = _fixed_vector_type(data_type) + if vector_type is None: raise ResultWriteError( - f"Lance vector column {spec.name!r} must use fixed_size_list" + f"Lance vector column {spec.name!r} has unsupported type {data_type}; " + "expected fixed_size_list or a supported fixed-shape tensor type" ) - if data_type.list_size != spec.dimension: + shape, value_type = vector_type + if shape != (spec.dimension,): raise ResultWriteError( - f"Lance vector column {spec.name!r} has dimension " - f"{data_type.list_size}, expected {spec.dimension}" + f"Lance vector column {spec.name!r} has shape {shape}, expected " + f"one-dimensional vectors of dimension {spec.dimension}" ) expected_type = getattr(pa, spec.dtype)() - if data_type.value_type != expected_type: + if value_type != expected_type: raise ResultWriteError( f"Lance vector column {spec.name!r} has dtype " - f"{data_type.value_type}, expected {spec.dtype}" + f"{value_type}, expected {spec.dtype}" + ) + if not pa.types.is_fixed_size_list(data_type): + fields[field_index] = pa.field( + field.name, + pa.list_(expected_type, spec.dimension), + nullable=field.nullable, + metadata=field.metadata, + ) + return pa.schema(fields, metadata=schema.metadata) + + +def _fixed_vector_type( + data_type: pa.DataType, +) -> tuple[tuple[int, ...], pa.DataType] | None: + if pa.types.is_fixed_size_list(data_type): + return (data_type.list_size,), data_type.value_type + if isinstance(data_type, pa.FixedShapeTensorType): + return tuple(data_type.shape), data_type.value_type + if not isinstance(data_type, pa.ExtensionType): + return None + if data_type.extension_name not in _RAY_FIXED_SHAPE_TENSOR_EXTENSION_NAMES: + return None + shape = getattr(data_type, "shape", None) + value_type = getattr(data_type, "value_type", None) + if shape is None or not isinstance(value_type, pa.DataType): + return None + try: + normalized_shape = tuple(int(dimension) for dimension in shape) + except (TypeError, ValueError): + return None + return normalized_shape, value_type + + +def _normalize_vector_batch( + batch: pa.Table, + *, + request: LanceResultSinkRequest, + target_schema: pa.Schema, +) -> pa.Table: + batch_target_schema = _canonical_vector_schema(batch.schema, request) + if not batch_target_schema.equals(target_schema, check_metadata=True): + raise ResultWriteError( + "Lance vector batch schema does not match the driver target schema: " + f"{_schema_mismatch_reason(batch_target_schema, target_schema)}" + ) + try: + normalized = ( + batch + if batch.schema.equals(target_schema, check_metadata=True) + else batch.cast(batch_target_schema) + ) + except (pa.ArrowException, ValueError) as exc: + from tributo.inference._credential_safety import safe_exception_summary + + raise ResultWriteError( + "Lance vector batch cannot be normalized to the declared schema: " + f"{type(exc).__name__}: {safe_exception_summary(exc)}" + ) from None + return validate_vector_batch(normalized, request) + + +def _schema_mismatch_reason(actual: pa.Schema, expected: pa.Schema) -> str: + if actual.names != expected.names: + return f"field names or order are {actual.names!r}, expected {expected.names!r}" + for actual_field, expected_field in zip(actual, expected, strict=True): + if actual_field.type != expected_field.type: + return ( + f"field {actual_field.name!r} has type {actual_field.type}, " + f"expected {expected_field.type}" + ) + if actual_field.nullable != expected_field.nullable: + return ( + f"field {actual_field.name!r} has nullable={actual_field.nullable}, " + f"expected nullable={expected_field.nullable}" ) + if actual_field.metadata != expected_field.metadata: + return f"field {actual_field.name!r} metadata differs" + if actual.metadata != expected.metadata: + return "schema metadata differs" + return "schemas differ" def validate_vector_batch(batch: pa.Table, request: LanceResultSinkRequest) -> pa.Table: diff --git a/tests/inference/test_lance_sink.py b/tests/inference/test_lance_sink.py index cbbed96..b16c377 100644 --- a/tests/inference/test_lance_sink.py +++ b/tests/inference/test_lance_sink.py @@ -2,17 +2,22 @@ from __future__ import annotations +import subprocess +import sys from typing import Any from unittest.mock import patch +import numpy as np import pyarrow as pa import pytest +from ray.air.util.tensor_extensions.arrow import ArrowTensorType, ArrowTensorTypeV2 from tests.inference.conformance.test_result_sink_contract import ( assert_result_sink_conformance, ) from tributo._common.storage_profiles import StorageProfile from tributo.data.base import WriteMode +from tributo.data.refs import schema_fingerprint from tributo.data.writing.contracts import WriteBindingError, WriteCapabilityError from tributo.exceptions import ResultMaterializationError, ResultWriteError from tributo.inference.contracts import ( @@ -21,6 +26,8 @@ ) from tributo.integrations.sinks.lance import ( LanceResultSink, + _canonical_vector_schema, + _normalize_vector_batch, validate_vector_batch, ) @@ -92,6 +99,22 @@ def test_lance_sink_runs_result_sink_conformance() -> None: assert dataset.map_calls +def test_lance_sink_module_imports_in_fresh_interpreter() -> None: + result = subprocess.run( + [ + sys.executable, + "-c", + "from tributo.integrations.sinks.lance import LanceResultSink", + ], + check=False, + capture_output=True, + text=True, + timeout=30, + ) + + assert result.returncode == 0, result.stderr + + def test_lance_sink_is_explicit_and_returns_schema_metadata() -> None: dataset = _Dataset(_vector_schema()) request = _request(mode="overwrite", data_storage_version="2.1") @@ -113,12 +136,304 @@ def test_lance_sink_is_explicit_and_returns_schema_metadata() -> None: assert len(receipt.metadata["schema_fingerprint"]) == 64 +def test_lance_sink_fingerprints_normalized_ray_tensor_schema() -> None: + source_schema = pa.schema( + [ + pa.field("id", pa.int64(), nullable=False), + pa.field( + "vector", + ArrowTensorTypeV2((3,), pa.float32()), + metadata={b"role": b"embedding"}, + ), + ], + metadata={b"dataset": b"inference"}, + ) + dataset = _Dataset(source_schema) + with patch( + "tributo.integrations.sinks.lance.default_write_gateway" + ) as default_gateway: + receipt = LanceResultSink().write( + dataset, _request(), run_id="run-1", plan_digest="a" * 64 + ) + + target_schema = pa.schema( + [ + pa.field("id", pa.int64(), nullable=False), + pa.field( + "vector", + pa.list_(pa.float32(), 3), + metadata={b"role": b"embedding"}, + ), + ], + metadata={b"dataset": b"inference"}, + ) + assert receipt.metadata["schema_fingerprint"] == schema_fingerprint(target_schema) + _, handle = default_gateway.return_value.execute.call_args.args + assert handle.dataset is dataset + assert dataset.map_calls[0][1]["batch_format"] == "pyarrow" + + +def test_lance_sink_without_declared_vectors_preserves_existing_path() -> None: + schema = pa.schema( + [ + pa.field("id", pa.int64()), + pa.field("payload", pa.list_(pa.float32())), + ] + ) + dataset = _Dataset(schema) + with patch( + "tributo.integrations.sinks.lance.default_write_gateway" + ) as default_gateway: + receipt = LanceResultSink().write( + dataset, + _request(vector_columns=()), + run_id="run-1", + plan_digest="a" * 64, + ) + + assert dataset.map_calls == [] + assert receipt.metadata["schema_fingerprint"] == schema_fingerprint(schema) + _, handle = default_gateway.return_value.execute.call_args.args + assert handle.dataset is dataset + + +@pytest.mark.parametrize( + "data_type", + [ + ArrowTensorType((3,), pa.float32()), + ArrowTensorTypeV2((3,), pa.float32()), + pa.fixed_shape_tensor(pa.float32(), (3,)), + pa.list_(pa.float32(), 3), + ], +) +def test_vector_schema_accepts_supported_fixed_shape_types( + data_type: pa.DataType, +) -> None: + source_schema = pa.schema( + [ + pa.field("id", pa.int64(), nullable=False), + pa.field( + "vector", + data_type, + nullable=False, + metadata={b"role": b"embedding"}, + ), + ], + metadata={b"dataset": b"inference"}, + ) + + target_schema = _canonical_vector_schema(source_schema, _request()) + + assert target_schema.names == source_schema.names + vector_field = target_schema.field("vector") + assert vector_field.type == pa.list_(pa.float32(), 3) + assert vector_field.nullable is False + assert vector_field.metadata == {b"role": b"embedding"} + assert target_schema.metadata == {b"dataset": b"inference"} + + +@pytest.mark.parametrize("representation", ["ray-v1", "ray-v2", "arrow-native"]) +def test_vector_batch_normalizes_fixed_shape_tensor_and_preserves_values( + representation: str, +) -> None: + values = np.asarray( + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + dtype=np.float32, + ) + if representation == "arrow-native": + tensor = pa.FixedShapeTensorArray.from_numpy_ndarray(values) + else: + tensor_type = ( + ArrowTensorType((3,), pa.float32()) + if representation == "ray-v1" + else ArrowTensorTypeV2((3,), pa.float32()) + ) + tensor = pa.ExtensionArray.from_storage( + tensor_type, + pa.array(values.tolist(), type=tensor_type.storage_type), + ) + batch = pa.table( + { + "id": pa.array([1, 2], type=pa.int64()), + "vector": tensor, + } + ) + target_schema = _canonical_vector_schema(batch.schema, _request()) + + normalized = _normalize_vector_batch( + batch, + request=_request(), + target_schema=target_schema, + ) + + assert normalized.schema == target_schema + assert normalized.column("vector").to_pylist() == [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + ] + + +def test_vector_batch_rejects_source_dtype_drift_before_cast() -> None: + target_schema = _vector_schema() + batch = pa.table( + { + "id": pa.array([1], type=pa.int64()), + "vector": pa.array( + [[1.0, 2.0, 3.0]], + type=pa.list_(pa.float64(), 3), + ), + } + ) + + with pytest.raises(ResultWriteError, match="dtype"): + _normalize_vector_batch( + batch, + request=_request(), + target_schema=target_schema, + ) + + +def test_vector_batch_rejects_source_rank_drift_before_cast() -> None: + tensor_type = ArrowTensorTypeV2((1, 3), pa.float32()) + batch = pa.table( + { + "id": pa.array([], type=pa.int64()), + "vector": pa.ExtensionArray.from_storage( + tensor_type, + pa.array([], type=tensor_type.storage_type), + ), + } + ) + + with pytest.raises(ResultWriteError, match="shape"): + _normalize_vector_batch( + batch, + request=_request(), + target_schema=_vector_schema(), + ) + + +@pytest.mark.parametrize( + ("batch_schema", "message"), + [ + ( + pa.schema( + [ + pa.field("id", pa.int32(), nullable=False), + pa.field("vector", pa.list_(pa.float32(), 3)), + ] + ), + "'id'.*int32.*int64", + ), + ( + pa.schema( + [ + pa.field("id", pa.int64(), nullable=True), + pa.field("vector", pa.list_(pa.float32(), 3)), + ] + ), + "'id'.*nullable=True.*nullable=False", + ), + ( + pa.schema( + [ + pa.field( + "id", pa.int64(), nullable=False, metadata={b"role": b"key"} + ), + pa.field("vector", pa.list_(pa.float32(), 3)), + ] + ), + "'id' metadata differs", + ), + ( + pa.schema( + [ + pa.field("id", pa.int64(), nullable=False), + pa.field("vector", pa.list_(pa.float32(), 3)), + ], + metadata={b"dataset": b"worker"}, + ), + "schema metadata differs", + ), + ( + pa.schema( + [ + pa.field("vector", pa.list_(pa.float32(), 3)), + pa.field("id", pa.int64(), nullable=False), + ] + ), + "field names or order", + ), + ], +) +def test_vector_batch_rejects_non_vector_schema_drift_before_cast( + batch_schema: pa.Schema, message: str +) -> None: + values = { + "id": pa.array([1], type=batch_schema.field("id").type), + "vector": pa.array( + [[1.0, 2.0, 3.0]], + type=pa.list_(pa.float32(), 3), + ), + } + batch = pa.Table.from_arrays( + [values[field.name] for field in batch_schema], + schema=batch_schema, + ) + + with pytest.raises(ResultWriteError, match=message): + _normalize_vector_batch( + batch, + request=_request(), + target_schema=_vector_schema(), + ) + + +def test_vector_batch_cast_failure_retains_only_safe_detail() -> None: + source_schema = pa.schema( + [ + pa.field("id", pa.int64(), nullable=False), + pa.field("vector", ArrowTensorTypeV2((3,), pa.float32())), + ] + ) + target_schema = _canonical_vector_schema(source_schema, _request()) + + class _FailingCastBatch: + schema = source_schema + + def cast(self, schema: pa.Schema) -> pa.Table: + del schema + raise pa.ArrowInvalid("field vector failed; token=secret-value") + + batch: Any = _FailingCastBatch() + with pytest.raises(ResultWriteError) as exc_info: + _normalize_vector_batch( + batch, + request=_request(), + target_schema=target_schema, + ) + + message = str(exc_info.value) + assert "ArrowInvalid: field vector failed" in message + assert "token=" in message + assert "secret-value" not in message + assert exc_info.value.__cause__ is None + + @pytest.mark.parametrize( ("schema", "message"), [ (_vector_schema(dimension=2), "dimension"), (pa.schema([pa.field("vector", pa.list_(pa.float32()))]), "fixed_size_list"), (_vector_schema(dtype=pa.float64()), "dtype"), + ( + pa.schema([pa.field("vector", ArrowTensorTypeV2((1, 3), pa.float32()))]), + "shape", + ), + ( + pa.schema([pa.field("vector", ArrowTensorTypeV2((3,), pa.float64()))]), + "dtype", + ), (pa.schema([pa.field("id", pa.int64())]), "missing"), ], ) diff --git a/tests/integrations/test_lance_vector_index.py b/tests/integrations/test_lance_vector_index.py index 0b203fe..1c3a788 100644 --- a/tests/integrations/test_lance_vector_index.py +++ b/tests/integrations/test_lance_vector_index.py @@ -36,6 +36,11 @@ import pyarrow.parquet as pq import ray +from tributo.inference.contracts import ( + LanceResultSinkRequest, + LanceVectorColumnSpec, +) +from tributo.integrations.sinks.lance import LanceResultSink from tributo.job import TributoClient from tributo.vector_index.contracts import ( CoverageStatus, @@ -70,6 +75,23 @@ _PROFILE = "vector_it" +class _EmbeddingPredictor: + """Return a conventional two-dimensional NumPy embedding batch.""" + + def __call__(self, batch: pa.Table) -> dict[str, np.ndarray]: + row_ids = np.asarray(batch.column("id").to_pylist(), dtype=np.int64) + return { + "id": row_ids, + "group": np.asarray( + ["even" if row_id % 2 == 0 else "odd" for row_id in row_ids] + ), + "vector": np.asarray( + [_vector_for(int(row_id)) for row_id in row_ids], + dtype=np.float32, + ), + } + + def _configure_cluster() -> None: ray.init(address="auto", ignore_reinit_error=True) deadline = time.monotonic() + 60 @@ -164,20 +186,54 @@ def _table(row_ids: list[int], *, duplicate_query: bool = False) -> pa.Table: def _write_initial_dataset(uri: str, storage_options: dict[str, str]) -> list[int]: - all_ids: list[int] = [] - for fragment in range(_INITIAL_FRAGMENTS): - start = fragment * _ROWS_PER_FRAGMENT - row_ids = list(range(start, start + _ROWS_PER_FRAGMENT)) - all_ids.extend(row_ids) - lance.write_dataset( - _table(row_ids), - uri, - mode="overwrite" if fragment == 0 else "append", - storage_options=storage_options, - ) + row_ids = list(range(_INITIAL_FRAGMENTS * _ROWS_PER_FRAGMENT)) + source = ray.data.from_items( + [{"id": row_id} for row_id in row_ids], + override_num_blocks=_INITIAL_FRAGMENTS, + ) + embeddings = source.map_batches( + _EmbeddingPredictor, + batch_format="pyarrow", + batch_size=128, + ) + inferred_schema = embeddings.schema() + arrow_schema = getattr(inferred_schema, "base_schema", inferred_schema) + assert isinstance(arrow_schema, pa.Schema) + ray_vector_type = arrow_schema.field("vector").type + assert getattr(ray_vector_type, "extension_name", None) in { + "ray.data.arrow_tensor", + "ray.data.arrow_tensor_v2", + } + + LanceResultSink().write( + embeddings, + LanceResultSinkRequest( + uri=uri, + storage_profile=_PROFILE, + mode="create", + min_rows_per_file=_ROWS_PER_FRAGMENT, + max_rows_per_file=_ROWS_PER_FRAGMENT, + vector_columns=( + LanceVectorColumnSpec(name="vector", dimension=_DIMENSION), + ), + ), + run_id="distributed-embedding-it", + plan_digest="e" * 64, + ) + dataset = lance.LanceDataset(uri, storage_options=storage_options) - assert len(dataset.get_fragments()) == _INITIAL_FRAGMENTS - return all_ids + assert dataset.schema.field("vector").type == pa.list_(pa.float32(), _DIMENSION) + persisted = dataset.to_table(columns=["id", "group", "vector"]).sort_by("id") + assert persisted.num_rows == len(row_ids) + assert [int(value) for value in persisted["id"].to_pylist()] == row_ids + np.testing.assert_allclose( + np.asarray(persisted["vector"].to_pylist(), dtype=np.float32), + np.asarray([_vector_for(row_id) for row_id in row_ids], dtype=np.float32), + ) + # Ray may split each input block into multiple batches, and Lance-Ray may + # materialize those batches as separate fragments. + assert len(dataset.get_fragments()) >= _INITIAL_FRAGMENTS + return row_ids def _build_with_concurrent_append(