Skip to content

Commit 1299e4e

Browse files
haoxu0ntkathole
authored andcommitted
perf: Scan Lance sources lazily in the DuckDB offline store
Register the LanceDataset with DuckDB instead of materializing it with to_arrow(), so the query's column projection is pushed into the Lance scanner. Retrievals that select scalar features no longer read the embedding columns stored next to them. Signed-off-by: haoxu0 <[email protected]>
1 parent c9b9d72 commit 1299e4e

3 files changed

Lines changed: 110 additions & 1 deletion

File tree

‎docs/reference/data-sources/table-formats.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,10 @@ lance_format = LanceFormat(
3131
)
3232
```
3333

34+
The DuckDB offline store reads a `LanceSource` as a lazy scan, so only the
35+
columns a query selects are read from the dataset. Materializing a few scalar
36+
features does not load the embedding columns stored beside them.
37+
3438
### Apache Iceberg
3539

3640
[Apache Iceberg](https://iceberg.apache.org/) is an open table format designed for huge analytic datasets. It provides:

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

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from __future__ import annotations
22

3+
import hashlib
34
import json
45
import os
56
from datetime import date, datetime, timezone
@@ -193,8 +194,24 @@ def _read_lance_source(data_source: "LanceSource") -> Table:
193194
is honoured before any data crosses into ibis. The declared-schema contract
194195
is enforced separately, by ``validate_lance_source_schema``, because the
195196
reader callback is handed a ``DataSource`` without its feature view.
197+
198+
The dataset is registered with DuckDB as a lazy scan rather than read into
199+
memory. A ``LanceDataset`` is a ``pyarrow.dataset.Dataset``, so DuckDB pushes
200+
the query's column projection into the Lance scanner: a retrieval that needs
201+
a few scalar features never reads the embedding columns stored beside them.
202+
203+
The view is named after the dataset's location and resolved version. An
204+
opened dataset is a snapshot, so the name always denotes the same data and
205+
re-registering it is harmless, while a job built before a later write keeps
206+
reading the version it was built against.
196207
"""
197-
return ibis.memtable(data_source.to_arrow())
208+
dataset = data_source.get_dataset()
209+
digest = hashlib.sha256(f"{dataset.uri}@{dataset.version}".encode()).hexdigest()
210+
view_name = f"__feast_lance_{digest[:16]}"
211+
212+
backend = ibis.get_backend()
213+
backend.con.register(view_name, dataset)
214+
return backend.table(view_name)
198215

199216

200217
def _validate_lance_feature_views(feature_views: List[FeatureView]) -> None:

‎sdk/python/tests/unit/infra/data_sources/test_lance_source.py‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -592,6 +592,94 @@ def test_pull_latest_from_a_lance_source(tmp_path):
592592
assert job.to_arrow().num_rows == 3
593593

594594

595+
def _driver_stats_with_embeddings(n_rows=3, width=8):
596+
stats = _driver_stats(n_rows=n_rows)
597+
embeddings = pa.FixedSizeListArray.from_arrays(
598+
pa.array([0.5] * (n_rows * width), pa.float32()), width
599+
)
600+
return stats.append_column("embedding", embeddings)
601+
602+
603+
@pytest.fixture
604+
def scanned_columns(monkeypatch):
605+
"""Record the column projection of every scan Lance is asked to run."""
606+
scans = []
607+
original = lance.LanceDataset.scanner
608+
609+
def recording_scanner(self, *args, **kwargs):
610+
scans.append(kwargs.get("columns"))
611+
return original(self, *args, **kwargs)
612+
613+
monkeypatch.setattr(lance.LanceDataset, "scanner", recording_scanner)
614+
return scans
615+
616+
617+
def test_the_duckdb_reader_is_lazy(tmp_path, scanned_columns):
618+
uri = str(tmp_path / "stats.lance")
619+
_write(uri, _driver_stats())
620+
table = _read_data_source(
621+
LanceSource(uri=uri, timestamp_field="event_timestamp"), str(tmp_path)
622+
)
623+
assert scanned_columns == []
624+
assert table.to_pyarrow().num_rows == 3
625+
assert scanned_columns
626+
627+
628+
def test_pull_latest_does_not_scan_columns_it_does_not_select(
629+
tmp_path, scanned_columns
630+
):
631+
uri = str(tmp_path / "stats.lance")
632+
_write(uri, _driver_stats_with_embeddings())
633+
job = DuckDBOfflineStore.pull_latest_from_table_or_query(
634+
config=_repo_config(tmp_path),
635+
data_source=LanceSource(uri=uri, timestamp_field="event_timestamp"),
636+
join_key_columns=["driver_id"],
637+
feature_name_columns=["conv_rate"],
638+
timestamp_field="event_timestamp",
639+
created_timestamp_column=None,
640+
start_date=datetime(2026, 1, 1),
641+
end_date=datetime(2026, 1, 2),
642+
)
643+
result = job.to_arrow()
644+
assert result.num_rows == 3
645+
assert "embedding" not in result.column_names
646+
assert scanned_columns
647+
assert all(
648+
columns is not None and "embedding" not in columns
649+
for columns in scanned_columns
650+
)
651+
652+
653+
def test_historical_retrieval_does_not_scan_columns_it_does_not_select(
654+
tmp_path, scanned_columns
655+
):
656+
uri = str(tmp_path / "stats.lance")
657+
_write(uri, _driver_stats_with_embeddings())
658+
source = LanceSource(uri=uri, timestamp_field="event_timestamp")
659+
df = _historical(tmp_path, source, entity_hour=13)
660+
# driver 1000 is the first row of the fixture.
661+
assert df["conv_rate"].iloc[0] == pytest.approx(0.1)
662+
assert scanned_columns
663+
assert all(
664+
columns is not None and "embedding" not in columns
665+
for columns in scanned_columns
666+
)
667+
668+
669+
def test_a_job_reads_the_version_it_was_built_against(tmp_path):
670+
"""A lazy scan must not pick up a write that lands before it executes."""
671+
uri = str(tmp_path / "stats.lance")
672+
_write(uri, _driver_stats(n_rows=3))
673+
source = LanceSource(uri=uri, timestamp_field="event_timestamp")
674+
before = _read_data_source(source, str(tmp_path))
675+
676+
_write(uri, _driver_stats(n_rows=2, base_hour=12), mode="append")
677+
after = _read_data_source(source, str(tmp_path))
678+
679+
assert before.to_pyarrow().num_rows == 3
680+
assert after.to_pyarrow().num_rows == 5
681+
682+
595683
def _pit_fixture(tmp_path, namespace=None):
596684
"""One driver with two feature rows, an hour apart.
597685

0 commit comments

Comments
 (0)