Visitar URL original
feat: Add a non-JVM write path for Lance data sources (#6945) · feast-dev/feast@f754323 · GitHub
Skip to content

Commit f754323

Browse files
authored
feat: Add a non-JVM write path for Lance data sources (#6945)
* feat: Add a non-JVM write path for Lance data sources #6943 made a `LanceSource` readable without Spark. This makes one writable, through the same two addressing modes: a uri, or a `namespace_client` plus `table_id` resolved through a Lance namespace, so a catalog-addressed dataset is committed through its catalog rather than behind its back. `_write_data_source` gains a `LanceSource` branch, following the `IcebergSource` branch already there, which makes the path reachable from `DuckDBOfflineStore.offline_write_batch`. The `isinstance` narrowing the read path open-coded is now a shared `_as_lance_source` helper used by both, rather than a second copy of the optional-import guard. Three things are settled before Lance is called. A pinned source is refused. A Lance commit always produces a new version, so `write_dataset` has no `version` argument at all; a write through a source pinned to version 1 or to a tag would succeed and then be invisible through the very source that performed it. Measured: after appending to a dataset tagged `prod` at version 1, the unpinned source reads 5 rows and the pinned source still reads 3. An absent dataset is created rather than appended to, because Lance has no create-or-append mode and rejects `create` on an existing dataset. An existing dataset's vector widths are compared against the incoming data. Lance rejects an `append` whose schema disagrees, but it accepts an `overwrite` that replaces a 8-wide embedding column with a 32-wide one, silently rewriting the declared shape and leaving a version history whose vectors are not mutually comparable, with every index built on the old width invalidated. No legitimate schema evolution changes an embedding's dimension, so this is refused. The guard is narrow on purpose: only vector widths are policed, and other type changes remain Lance's business. The declared `vector_length` is checked at the store entry point rather than in the writer callback, which is handed a `DataSource` and so cannot see the feature view. It reuses `_validate_vector_field_lengths` from #6909 rather than adding a second implementation. Without it, creating a dataset at a width other than the declared one would succeed and fail only on the next read. 19 tests added, all using the `dir` namespace implementation for the catalog-based cases so no server is needed. Signed-off-by: hao-xu5 <hxu44@apple.com> * fix: Correct what a Lance namespace does during a write The docstring claimed a catalog-addressed write is "committed through its namespace". That is only true of a create. Instrumenting the namespace client shows how far it is actually involved: create -> declare_table append -> describe_table only open -> describe_table, namespace_id So a namespace records that a table exists and where it lives, and does not track its versions: the version an append produces is never reported back to it. Worth stating precisely, because it is the reason the pin contract has to be enforced in `assert_writable` rather than left to the catalog -- no catalog is going to reject a write on a pin's behalf. Adds a test that records the namespace calls, so the claim is pinned by a test rather than asserted in prose. Signed-off-by: hao-xu5 <hxu44@apple.com> * fix: Accept an optional source in _as_lance_source offline_write_batch passes FeatureView.batch_source, which is optional, so the narrow DataSource annotation made mypy reject the call. The body already handled None, since isinstance(None, LanceSource) is False, so only the signature changes. Signed-off-by: hao-xu5 <hxu44@apple.com> --------- Signed-off-by: hao-xu5 <hxu44@apple.com>
1 parent c005dc3 commit f754323

3 files changed

Lines changed: 512 additions & 11 deletions

File tree

‎sdk/python/feast/infra/data_sources/contrib/lance/lance_source.py‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -289,6 +289,70 @@ def to_arrow(
289289
_validate_vector_field_lengths(table, feature_view)
290290
return table
291291

292+
# --------------------------------------------------------------- writing
293+
294+
def get_write_target(self) -> Dict[str, Any]:
295+
"""Keyword arguments addressing this source for ``lance.write_dataset``.
296+
297+
Deliberately omits ``version``. A Lance commit always produces a new
298+
version, so there is no argument that writes *into* an existing one;
299+
``write_dataset`` has no ``version`` parameter at all. Callers must
300+
reject a pinned source before reaching a write, which
301+
``assert_writable`` does.
302+
"""
303+
kwargs: Dict[str, Any] = {}
304+
if self.storage_options:
305+
kwargs["storage_options"] = dict(self.storage_options)
306+
307+
if self.is_catalog_based:
308+
kwargs["namespace_client"] = self.get_namespace_client()
309+
kwargs["table_id"] = self.table_id
310+
else:
311+
kwargs["uri"] = self.uri
312+
return kwargs
313+
314+
def assert_writable(self) -> None:
315+
"""Refuse to write through a pinned source.
316+
317+
A write commits a new version, which a source pinned to an older version
318+
or to a tag does not read. The write would therefore succeed and then be
319+
invisible through the very source that performed it, so the only honest
320+
outcome is to refuse up front.
321+
"""
322+
if self.pin is None:
323+
return
324+
raise ValueError(
325+
f"LanceSource '{self.name}' is pinned to {self.pin_description} and "
326+
f"cannot be written through. A Lance write commits a new version, so "
327+
f"it would not be visible through this pin. Write through an unpinned "
328+
f"LanceSource addressing the same dataset, then move the tag or re-pin "
329+
f"to the version the write produced."
330+
)
331+
332+
def get_existing_schema(self) -> Optional["pyarrow.Schema"]:
333+
"""Schema of the dataset this source addresses, or ``None`` if absent.
334+
335+
Used to choose between creating and appending, and to compare shapes
336+
before handing data to Lance. A failure to open is reported as absent
337+
rather than raised, because the only distinction the caller draws is
338+
whether there is an existing shape to honour.
339+
"""
340+
import lance
341+
342+
kwargs = self.get_write_target()
343+
try:
344+
if self.is_catalog_based:
345+
return lance.dataset(
346+
namespace_client=kwargs["namespace_client"],
347+
table_id=kwargs["table_id"],
348+
storage_options=kwargs.get("storage_options"),
349+
).schema
350+
return lance.dataset(
351+
uri=kwargs["uri"], storage_options=kwargs.get("storage_options")
352+
).schema
353+
except Exception:
354+
return None
355+
292356
# ----------------------------------------------------------- DataSource
293357

294358
def source_type(self) -> DataSourceProto.SourceType.ValueType:
@@ -469,3 +533,60 @@ def validate_lance_source_schema(feature_view: "FeatureView") -> None:
469533
from feast.utils import _validate_vector_field_lengths
470534

471535
_validate_vector_field_lengths(schema.empty_table(), feature_view)
536+
537+
538+
def _vector_widths(schema: "pyarrow.Schema") -> Dict[str, int]:
539+
"""Width of every fixed-width vector column in an Arrow schema."""
540+
import pyarrow
541+
542+
widths: Dict[str, int] = {}
543+
for field in schema:
544+
if pyarrow.types.is_fixed_size_list(field.type):
545+
widths[field.name] = field.type.list_size
546+
return widths
547+
548+
549+
def validate_lance_write_shape(
550+
table: "pyarrow.Table",
551+
source: LanceSource,
552+
existing_schema: Optional["pyarrow.Schema"],
553+
) -> None:
554+
"""Refuse a write that would change an existing vector column's width.
555+
556+
Lance rejects an ``append`` whose schema disagrees with the dataset, but it
557+
accepts an ``overwrite`` that replaces a vector column with one of a
558+
different width, silently rewriting the declared shape and leaving a version
559+
history whose vectors are not mutually comparable. Every ANN index and every
560+
downstream consumer built against the old width is invalidated by such a
561+
write, and no legitimate schema evolution changes an embedding's dimension,
562+
so this is reported as an error rather than carried out.
563+
564+
Only vector widths are policed here. Other kinds of type change remain
565+
Lance's business, because for those an ``overwrite`` is a defensible
566+
evolution.
567+
"""
568+
if existing_schema is None:
569+
return
570+
571+
existing = _vector_widths(existing_schema)
572+
incoming = _vector_widths(table.schema)
573+
574+
changed = {
575+
name: (existing[name], incoming[name])
576+
for name in existing.keys() & incoming.keys()
577+
if existing[name] != incoming[name]
578+
}
579+
if not changed:
580+
return
581+
582+
detail = ", ".join(
583+
f"'{name}' is {was} in the dataset but {now} in the data being written"
584+
for name, (was, now) in sorted(changed.items())
585+
)
586+
raise ValueError(
587+
f"Write to Lance dataset '{source.name}' would change the width of "
588+
f"vector column(s): {detail}. A width change invalidates every index and "
589+
f"consumer built against the current width, and leaves a version history "
590+
f"whose vectors cannot be compared with each other. Write the re-embedded "
591+
f"data to a new dataset instead of overwriting this one."
592+
)

‎sdk/python/feast/infra/offline_stores/duckdb.py‎

Lines changed: 86 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,25 @@
5353
from feast.repo_config import FeastConfigBaseModel, RepoConfig
5454

5555

56+
def _as_lance_source(data_source: Optional[DataSource]) -> Optional["LanceSource"]:
57+
"""Narrow a DataSource to a LanceSource, or ``None`` if it is not one.
58+
59+
``None`` is accepted so that a caller holding an optional source, such as
60+
``FeatureView.batch_source``, can ask the question without narrowing first.
61+
62+
The import is guarded because the Lance extra is optional: a repo that does
63+
not install it must still be able to read and write every other source
64+
type. Only the ``isinstance`` check is guarded, never the caller's
65+
subsequent work, so an ImportError raised by a Lance read or write itself
66+
surfaces instead of being mistaken for "not a Lance source".
67+
"""
68+
try:
69+
from feast.infra.data_sources.contrib.lance.lance_source import LanceSource
70+
except ImportError:
71+
return None
72+
return data_source if isinstance(data_source, LanceSource) else None
73+
74+
5675
def _read_data_source(data_source: DataSource, repo_path: str) -> Table:
5776
from feast.infra.data_sources.contrib.iceberg_catalog.iceberg_source import (
5877
IcebergSource,
@@ -61,17 +80,7 @@ def _read_data_source(data_source: DataSource, repo_path: str) -> Table:
6180
if isinstance(data_source, IcebergSource):
6281
return _read_iceberg_catalog_source(data_source, repo_path)
6382

64-
lance_source: Optional["LanceSource"] = None
65-
try:
66-
from feast.infra.data_sources.contrib.lance.lance_source import LanceSource
67-
68-
if isinstance(data_source, LanceSource):
69-
lance_source = data_source
70-
except ImportError:
71-
pass
72-
# Deliberately outside the try: an ImportError raised by the read itself (a
73-
# missing pylance, say) must surface rather than fall through to the
74-
# FileSource assert below and be reported as the wrong source type.
83+
lance_source = _as_lance_source(data_source)
7584
if lance_source is not None:
7685
return _read_lance_source(lance_source)
7786

@@ -241,6 +250,58 @@ def _validate_lance_feature_views(feature_views: List[FeatureView]) -> None:
241250
validate_lance_source_schema(feature_view)
242251

243252

253+
def _write_lance_data_source(
254+
table: Table,
255+
data_source: "LanceSource",
256+
mode: str = "append",
257+
allow_overwrite: bool = False,
258+
) -> None:
259+
"""Write an ibis table to a LanceSource without Spark.
260+
261+
Mirrors ``_read_lance_source``: the source resolves either a uri or a
262+
``namespace_client`` + ``table_id``, and the same resolution is reused for
263+
the write, so a catalog-addressed dataset is written to the location its
264+
namespace resolves rather than to one assembled here.
265+
266+
How far the namespace is involved depends on the operation. Creating a table
267+
declares it through the namespace; appending only describes the table to
268+
resolve its location, and the version the append produces is not reported
269+
back. A Lance namespace therefore records that a table exists and where it
270+
lives, and does not track its versions. That is why a pin is a Lance-level
271+
concern that ``assert_writable`` has to enforce here: no catalog is going to
272+
reject the write on the pin's behalf.
273+
274+
Three things are settled before Lance is called. A pinned source is refused,
275+
because a commit produces a new version that the pin would not read. An
276+
absent dataset is created rather than appended to, since Lance has no
277+
create-or-append mode. An existing dataset's vector widths are compared
278+
against the incoming data, which closes the one case Lance itself permits
279+
silently.
280+
"""
281+
import lance
282+
283+
from feast.infra.data_sources.contrib.lance.lance_source import (
284+
validate_lance_write_shape,
285+
)
286+
287+
data_source.assert_writable()
288+
289+
arrow_table = table.to_pyarrow()
290+
existing_schema = data_source.get_existing_schema()
291+
292+
if existing_schema is None:
293+
lance_mode = "create"
294+
else:
295+
if mode == "overwrite" and not allow_overwrite:
296+
raise SavedDatasetLocationAlreadyExists(
297+
location=data_source.uri or ".".join(data_source.table_id)
298+
)
299+
validate_lance_write_shape(arrow_table, data_source, existing_schema)
300+
lance_mode = mode
301+
302+
lance.write_dataset(arrow_table, mode=lance_mode, **data_source.get_write_target())
303+
304+
244305
def _write_data_source(
245306
table: Table,
246307
data_source: DataSource,
@@ -256,6 +317,11 @@ def _write_data_source(
256317
_write_iceberg_data_source(table, data_source, mode)
257318
return
258319

320+
lance_source = _as_lance_source(data_source)
321+
if lance_source is not None:
322+
_write_lance_data_source(table, lance_source, mode, allow_overwrite)
323+
return
324+
259325
assert isinstance(data_source, FileSource)
260326

261327
file_options = data_source.file_options
@@ -735,6 +801,15 @@ def offline_write_batch(
735801
table: pyarrow.Table,
736802
progress: Optional[Callable[[int], Any]],
737803
):
804+
# Checked here rather than in the writer callback, which is handed a
805+
# DataSource and so cannot see the declared schema. Creating a Lance
806+
# dataset whose vectors are a different width from the declared
807+
# vector_length would otherwise succeed and fail only on the next read.
808+
if _as_lance_source(feature_view.batch_source) is not None:
809+
from feast.utils import _validate_vector_field_lengths
810+
811+
_validate_vector_field_lengths(table, feature_view)
812+
738813
offline_write_batch_ibis(
739814
config=config,
740815
feature_view=feature_view,

0 commit comments

Comments
 (0)