Skip to content

Commit 5be5ed0

Browse files
ntkatholentkathole
committed
feat: Offline write batch for IcebergSource
Signed-off-by: ntkathole <[email protected]>
1 parent 5ee7ee3 commit 5be5ed0

3 files changed

Lines changed: 358 additions & 0 deletions

File tree

‎sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/tests/test_iceberg_offline_write.py‎

Lines changed: 273 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,273 @@
1+
"""Unit tests for IcebergSource offline_write_batch support.
2+
3+
Tests that SparkOfflineStore and DuckDB _write_data_source correctly
4+
route IcebergSource writes through PyIceberg's catalog-aware append/overwrite
5+
API instead of raw file writes.
6+
"""
7+
8+
from unittest.mock import MagicMock, patch
9+
10+
import pyarrow as pa
11+
import pytest
12+
13+
from feast.infra.data_sources.contrib.iceberg_catalog.iceberg_source import (
14+
IcebergSource,
15+
)
16+
17+
18+
@pytest.fixture
19+
def sample_arrow_table():
20+
"""A minimal Arrow table representing feature data."""
21+
return pa.table(
22+
{
23+
"entity_id": pa.array([1, 2, 3]),
24+
"feature_value": pa.array([10.0, 20.0, 30.0]),
25+
"event_timestamp": pa.array(
26+
["2026-01-01T00:00:00", "2026-01-02T00:00:00", "2026-01-03T00:00:00"]
27+
),
28+
}
29+
)
30+
31+
32+
@pytest.fixture
33+
def iceberg_source():
34+
"""An IcebergSource with REST catalog for testing."""
35+
return IcebergSource(
36+
endpoint="http://localhost:8181",
37+
warehouse="test_warehouse",
38+
namespace="test_ns",
39+
table="test_features",
40+
catalog_type="rest",
41+
catalog_name="feast_test",
42+
timestamp_field="event_timestamp",
43+
)
44+
45+
46+
@pytest.fixture
47+
def iceberg_source_hive():
48+
"""An IcebergSource with Hive catalog for testing."""
49+
return IcebergSource(
50+
catalog_type="hive",
51+
warehouse="hive_warehouse",
52+
namespace="hive_ns",
53+
table="hive_features",
54+
catalog_name="hive_catalog",
55+
catalog_properties={"uri": "thrift://metastore:9083"},
56+
timestamp_field="event_timestamp",
57+
)
58+
59+
60+
class TestSparkOfflineStoreIcebergWrite:
61+
"""Tests for _iceberg_offline_write_batch in the Spark offline store."""
62+
63+
@patch(
64+
"feast.infra.offline_stores.contrib.spark_offline_store.spark"
65+
"._load_pyiceberg_table"
66+
)
67+
def test_append_calls_pyiceberg(
68+
self, mock_load_table, iceberg_source, sample_arrow_table
69+
):
70+
"""Verify that IcebergSource write calls PyIceberg append()."""
71+
from feast.infra.offline_stores.contrib.spark_offline_store.spark import (
72+
_iceberg_offline_write_batch,
73+
)
74+
75+
mock_iceberg_table = MagicMock()
76+
mock_load_table.return_value = mock_iceberg_table
77+
78+
feature_view = MagicMock()
79+
feature_view.batch_source = iceberg_source
80+
81+
config = MagicMock()
82+
83+
_iceberg_offline_write_batch(config, feature_view, sample_arrow_table)
84+
85+
mock_load_table.assert_called_once_with(iceberg_source)
86+
mock_iceberg_table.append.assert_called_once_with(sample_arrow_table)
87+
88+
@patch(
89+
"feast.infra.offline_stores.contrib.spark_offline_store.spark"
90+
"._load_pyiceberg_table"
91+
)
92+
def test_passes_correct_data_source(
93+
self, mock_load_table, iceberg_source_hive, sample_arrow_table
94+
):
95+
"""Verify Hive catalog IcebergSource is passed to _load_pyiceberg_table."""
96+
from feast.infra.offline_stores.contrib.spark_offline_store.spark import (
97+
_iceberg_offline_write_batch,
98+
)
99+
100+
mock_iceberg_table = MagicMock()
101+
mock_load_table.return_value = mock_iceberg_table
102+
103+
feature_view = MagicMock()
104+
feature_view.batch_source = iceberg_source_hive
105+
106+
config = MagicMock()
107+
108+
_iceberg_offline_write_batch(config, feature_view, sample_arrow_table)
109+
110+
mock_load_table.assert_called_once_with(iceberg_source_hive)
111+
mock_iceberg_table.append.assert_called_once()
112+
113+
@patch(
114+
"feast.infra.offline_stores.contrib.spark_offline_store.spark"
115+
"._load_pyiceberg_table"
116+
)
117+
def test_append_error_propagates(
118+
self, mock_load_table, iceberg_source, sample_arrow_table
119+
):
120+
"""Verify that PyIceberg errors during append propagate."""
121+
from feast.infra.offline_stores.contrib.spark_offline_store.spark import (
122+
_iceberg_offline_write_batch,
123+
)
124+
125+
mock_iceberg_table = MagicMock()
126+
mock_iceberg_table.append.side_effect = RuntimeError("Catalog commit failed")
127+
mock_load_table.return_value = mock_iceberg_table
128+
129+
feature_view = MagicMock()
130+
feature_view.batch_source = iceberg_source
131+
132+
config = MagicMock()
133+
134+
with pytest.raises(RuntimeError, match="Catalog commit failed"):
135+
_iceberg_offline_write_batch(config, feature_view, sample_arrow_table)
136+
137+
@patch(
138+
"feast.infra.offline_stores.contrib.spark_offline_store.spark"
139+
"._load_pyiceberg_table"
140+
)
141+
@patch(
142+
"feast.infra.offline_stores.contrib.spark_offline_store.spark"
143+
".get_spark_session_or_start_new_with_repoconfig"
144+
)
145+
def test_offline_write_batch_dispatches_iceberg(
146+
self,
147+
mock_get_spark,
148+
mock_load_table,
149+
iceberg_source,
150+
sample_arrow_table,
151+
):
152+
"""End-to-end: SparkOfflineStore.offline_write_batch dispatches to
153+
_iceberg_offline_write_batch when batch_source is IcebergSource."""
154+
from feast.infra.offline_stores.contrib.spark_offline_store.spark import (
155+
SparkOfflineStore,
156+
SparkOfflineStoreConfig,
157+
)
158+
159+
mock_iceberg_table = MagicMock()
160+
mock_load_table.return_value = mock_iceberg_table
161+
162+
feature_view = MagicMock()
163+
feature_view.batch_source = iceberg_source
164+
165+
config = MagicMock()
166+
config.offline_store = SparkOfflineStoreConfig()
167+
168+
SparkOfflineStore.offline_write_batch(
169+
config, feature_view, sample_arrow_table, progress=None
170+
)
171+
172+
mock_load_table.assert_called_once_with(iceberg_source)
173+
mock_iceberg_table.append.assert_called_once_with(sample_arrow_table)
174+
mock_get_spark.assert_not_called()
175+
176+
def test_offline_write_batch_rejects_non_spark_non_iceberg(self):
177+
"""SparkOfflineStore.offline_write_batch raises for unknown source types."""
178+
from feast.infra.offline_stores.contrib.spark_offline_store.spark import (
179+
SparkOfflineStore,
180+
SparkOfflineStoreConfig,
181+
)
182+
183+
feature_view = MagicMock()
184+
feature_view.batch_source = MagicMock()
185+
feature_view.batch_source.__class__ = type("UnknownSource", (), {})
186+
187+
config = MagicMock()
188+
config.offline_store = SparkOfflineStoreConfig()
189+
190+
with pytest.raises(AssertionError):
191+
SparkOfflineStore.offline_write_batch(
192+
config,
193+
feature_view,
194+
pa.table({"a": [1]}),
195+
progress=None,
196+
)
197+
198+
199+
class TestDuckDBOfflineStoreIcebergWrite:
200+
"""Tests for _write_iceberg_data_source in the DuckDB offline store."""
201+
202+
@patch("feast.infra.offline_stores.duckdb._load_pyiceberg_table_for_write")
203+
def test_append_mode(self, mock_load_table, iceberg_source, sample_arrow_table):
204+
"""Verify append mode calls iceberg_table.append()."""
205+
from feast.infra.offline_stores.duckdb import _write_iceberg_data_source
206+
207+
mock_iceberg_table = MagicMock()
208+
mock_load_table.return_value = mock_iceberg_table
209+
210+
import ibis
211+
212+
ibis_table = ibis.memtable(sample_arrow_table)
213+
214+
_write_iceberg_data_source(ibis_table, iceberg_source, mode="append")
215+
216+
mock_load_table.assert_called_once_with(iceberg_source)
217+
mock_iceberg_table.append.assert_called_once()
218+
mock_iceberg_table.overwrite.assert_not_called()
219+
220+
@patch("feast.infra.offline_stores.duckdb._load_pyiceberg_table_for_write")
221+
def test_overwrite_mode(self, mock_load_table, iceberg_source, sample_arrow_table):
222+
"""Verify overwrite mode calls iceberg_table.overwrite()."""
223+
from feast.infra.offline_stores.duckdb import _write_iceberg_data_source
224+
225+
mock_iceberg_table = MagicMock()
226+
mock_load_table.return_value = mock_iceberg_table
227+
228+
import ibis
229+
230+
ibis_table = ibis.memtable(sample_arrow_table)
231+
232+
_write_iceberg_data_source(ibis_table, iceberg_source, mode="overwrite")
233+
234+
mock_load_table.assert_called_once_with(iceberg_source)
235+
mock_iceberg_table.overwrite.assert_called_once()
236+
mock_iceberg_table.append.assert_not_called()
237+
238+
@patch("feast.infra.offline_stores.duckdb._load_pyiceberg_table_for_write")
239+
def test_write_data_source_dispatches_iceberg(
240+
self, mock_load_table, iceberg_source, sample_arrow_table
241+
):
242+
"""_write_data_source dispatches to _write_iceberg_data_source for IcebergSource."""
243+
from feast.infra.offline_stores.duckdb import _write_data_source
244+
245+
mock_iceberg_table = MagicMock()
246+
mock_load_table.return_value = mock_iceberg_table
247+
248+
import ibis
249+
250+
ibis_table = ibis.memtable(sample_arrow_table)
251+
252+
_write_data_source(ibis_table, iceberg_source, repo_path="/tmp/repo")
253+
254+
mock_load_table.assert_called_once_with(iceberg_source)
255+
mock_iceberg_table.append.assert_called_once()
256+
257+
@patch("feast.infra.offline_stores.duckdb._load_pyiceberg_table_for_write")
258+
def test_error_propagates(
259+
self, mock_load_table, iceberg_source, sample_arrow_table
260+
):
261+
"""Verify that PyIceberg errors propagate through DuckDB path."""
262+
from feast.infra.offline_stores.duckdb import _write_iceberg_data_source
263+
264+
mock_iceberg_table = MagicMock()
265+
mock_iceberg_table.append.side_effect = RuntimeError("Schema mismatch")
266+
mock_load_table.return_value = mock_iceberg_table
267+
268+
import ibis
269+
270+
ibis_table = ibis.memtable(sample_arrow_table)
271+
272+
with pytest.raises(RuntimeError, match="Schema mismatch"):
273+
_write_iceberg_data_source(ibis_table, iceberg_source)

‎sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -361,7 +361,16 @@ def offline_write_batch(
361361
table: pyarrow.Table,
362362
progress: Optional[Callable[[int], Any]],
363363
):
364+
from feast.infra.data_sources.contrib.iceberg_catalog.iceberg_source import (
365+
IcebergSource,
366+
)
367+
364368
assert isinstance(config.offline_store, SparkOfflineStoreConfig)
369+
370+
if isinstance(feature_view.batch_source, IcebergSource):
371+
_iceberg_offline_write_batch(config, feature_view, table)
372+
return
373+
365374
assert isinstance(feature_view.batch_source, SparkSource)
366375

367376
pa_schema, column_names = offline_utils.get_pyarrow_schema_from_batch_source(
@@ -1275,6 +1284,28 @@ def _load_pyiceberg_table(data_source: "DataSource"):
12751284
return catalog.load_table(fqn)
12761285

12771286

1287+
def _iceberg_offline_write_batch(
1288+
config: "RepoConfig",
1289+
feature_view: "FeatureView",
1290+
table: "pyarrow.Table",
1291+
) -> None:
1292+
"""Write a PyArrow table to an IcebergSource via PyIceberg.
1293+
1294+
Uses PyIceberg's append() API which produces proper Iceberg snapshots,
1295+
handles schema evolution, and updates catalog metadata — making the
1296+
written data discoverable via any Iceberg-compatible engine.
1297+
"""
1298+
from feast.infra.data_sources.contrib.iceberg_catalog.iceberg_source import (
1299+
IcebergSource,
1300+
)
1301+
1302+
assert isinstance(feature_view.batch_source, IcebergSource)
1303+
data_source: IcebergSource = feature_view.batch_source
1304+
1305+
iceberg_table = _load_pyiceberg_table(data_source)
1306+
iceberg_table.append(table)
1307+
1308+
12781309
def _pull_from_iceberg_source(
12791310
spark_session: SparkSession,
12801311
data_source: "DataSource",

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

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,13 +126,67 @@ def _read_iceberg_catalog_source(data_source: "IcebergSource", repo_path: str) -
126126
return ibis.memtable(table.scan().to_arrow())
127127

128128

129+
def _load_pyiceberg_table_for_write(data_source: "IcebergSource"):
130+
"""Load a PyIceberg table object for write operations.
131+
132+
Reuses the same catalog resolution logic as ``_read_iceberg_catalog_source``
133+
but returns the raw PyIceberg table instead of an ibis memtable, so the
134+
caller can use ``table.append()`` or ``table.overwrite()`` for writes.
135+
"""
136+
fqn = f"{data_source.namespace}.{data_source.iceberg_table}"
137+
138+
if data_source.catalog_type != "rest":
139+
catalog_client = data_source.get_catalog_client()
140+
return catalog_client.load_table(fqn)
141+
142+
from pyiceberg.catalog import load_catalog
143+
144+
catalog_config: dict = {
145+
"type": "rest",
146+
"uri": data_source.endpoint,
147+
"warehouse": data_source.warehouse,
148+
}
149+
if data_source.token_env_var:
150+
token = os.environ.get(data_source.token_env_var)
151+
if token:
152+
catalog_config["token"] = token
153+
catalog = load_catalog(data_source.catalog_name, **catalog_config)
154+
return catalog.load_table(fqn)
155+
156+
157+
def _write_iceberg_data_source(
158+
table: Table,
159+
data_source: "IcebergSource",
160+
mode: str = "append",
161+
) -> None:
162+
"""Write an ibis table to an IcebergSource via PyIceberg.
163+
164+
Converts the ibis table to Arrow and uses PyIceberg's append/overwrite
165+
API to produce proper Iceberg snapshots with catalog metadata updates.
166+
"""
167+
iceberg_table = _load_pyiceberg_table_for_write(data_source)
168+
arrow_table = table.to_pyarrow()
169+
if mode == "overwrite":
170+
iceberg_table.overwrite(arrow_table)
171+
else:
172+
iceberg_table.append(arrow_table)
173+
174+
129175
def _write_data_source(
130176
table: Table,
131177
data_source: DataSource,
132178
repo_path: str,
133179
mode: str = "append",
134180
allow_overwrite: bool = False,
135181
):
182+
from feast.infra.data_sources.contrib.iceberg_catalog.iceberg_source import (
183+
IcebergSource,
184+
)
185+
186+
if isinstance(data_source, IcebergSource):
187+
_write_iceberg_data_source(table, data_source, mode)
188+
return
189+
136190
assert isinstance(data_source, FileSource)
137191

138192
file_options = data_source.file_options

0 commit comments

Comments
 (0)