Skip to content

Commit 4771156

Browse files
Pranjulcr7Pranjulcr7
authored
perf: Avoid pandas copy in Spark reads (#6878)
* perf: Avoid pandas copy in Spark reads Signed-off-by: Pranjul Gupta <[email protected]> * fix: Complete Spark Arrow ingestion Signed-off-by: Pranjul Gupta <[email protected]> --------- Signed-off-by: Pranjul Gupta <[email protected]>
1 parent 2583c45 commit 4771156

4 files changed

Lines changed: 128 additions & 11 deletions

File tree

‎sdk/python/feast/infra/compute_engines/spark/nodes.py‎

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
import pandas as pd
77
from pyspark.sql import DataFrame, SparkSession, Window
88
from pyspark.sql import functions as F
9-
from pyspark.sql.pandas.types import from_arrow_schema
109
from pyspark.sql.types import (
1110
ArrayType,
1211
BinaryType,
@@ -192,13 +191,7 @@ def execute(self, context: ExecutionContext) -> DAGValue:
192191
spark_df = cast(SparkRetrievalJob, retrieval_job).to_spark_df()
193192
else:
194193
arrow_table = retrieval_job.to_arrow()
195-
if arrow_table.num_rows == 0:
196-
spark_schema = from_arrow_schema(arrow_table.schema)
197-
spark_df = self.spark_session.createDataFrame(
198-
self.spark_session.sparkContext.emptyRDD(), schema=spark_schema
199-
)
200-
else:
201-
spark_df = self.spark_session.createDataFrame(arrow_table.to_pandas())
194+
spark_df = self.spark_session.createDataFrame(arrow_table)
202195

203196
return DAGValue(
204197
data=spark_df,

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1256,7 +1256,7 @@ def _register_iceberg_source_as_temp_view(
12561256

12571257
iceberg_table = _load_pyiceberg_table(data_source)
12581258
arrow_table = iceberg_table.scan().to_arrow()
1259-
df = spark_session.createDataFrame(arrow_table.to_pandas())
1259+
df = spark_session.createDataFrame(arrow_table)
12601260
df.createOrReplaceTempView(view_name)
12611261

12621262

‎sdk/python/tests/component/spark/test_nodes.py‎

Lines changed: 89 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
1-
from datetime import datetime, timedelta
2-
from unittest.mock import MagicMock
1+
from datetime import datetime, timedelta, timezone
2+
from unittest.mock import MagicMock, patch
3+
4+
import pyarrow as pa
5+
import pytest
6+
from pyspark.sql.pandas.types import from_arrow_schema
37

48
from feast.aggregation import Aggregation
59
from feast.infra.compute_engines.dag.context import ColumnInfo, ExecutionContext
@@ -9,13 +13,96 @@
913
SparkAggregationNode,
1014
SparkDedupNode,
1115
SparkJoinNode,
16+
SparkReadNode,
1217
SparkTransformationNode,
1318
)
19+
from feast.infra.offline_stores.contrib.spark_offline_store.spark import (
20+
SparkRetrievalJob,
21+
)
1422
from tests.example_repos.example_feature_repo_with_bfvs import (
1523
driver,
1624
)
1725

1826

27+
@pytest.mark.parametrize("empty", [False, True])
28+
@pytest.mark.parametrize("arrow_enabled", [False, True])
29+
def test_spark_read_node_preserves_arrow_values_and_schema(
30+
spark_session, empty, arrow_enabled
31+
):
32+
schema = pa.schema(
33+
[
34+
pa.field("id", pa.int64(), nullable=False),
35+
pa.field("count", pa.int32()),
36+
pa.field("score", pa.float32()),
37+
pa.field("name", pa.string()),
38+
pa.field("active", pa.bool_()),
39+
pa.field("event_timestamp", pa.timestamp("us", tz="UTC")),
40+
]
41+
)
42+
rows = [
43+
{"id": 1, "count": 2, "score": 0.5, "name": "first", "active": True},
44+
{"id": 2, "count": None, "score": None, "name": None, "active": None},
45+
{"id": 3, "count": -1, "score": 1.5, "name": "last", "active": False},
46+
]
47+
for row in rows:
48+
row["event_timestamp"] = datetime(
49+
2026, 1, 1, 12, 0, 0, 123456, tzinfo=timezone.utc
50+
)
51+
table = pa.concat_tables(
52+
[pa.Table.from_pylist(rows[:1], schema), pa.Table.from_pylist(rows[1:], schema)]
53+
)
54+
if empty:
55+
table = table.slice(0, 0)
56+
rows = []
57+
job = MagicMock()
58+
job.to_arrow.return_value = table
59+
node = SparkReadNode(
60+
"read",
61+
MagicMock(),
62+
ColumnInfo(["id"], ["count"], "event_timestamp", None),
63+
spark_session,
64+
)
65+
config = {
66+
"spark.sql.execution.arrow.pyspark.enabled": str(arrow_enabled).lower(),
67+
"spark.sql.session.timeZone": "UTC",
68+
}
69+
previous = {key: spark_session.conf.get(key) for key in config}
70+
for key, value in config.items():
71+
spark_session.conf.set(key, value)
72+
try:
73+
with patch(
74+
"feast.infra.compute_engines.spark.nodes.create_offline_store_retrieval_job",
75+
return_value=job,
76+
):
77+
result = node.execute(MagicMock())
78+
assert result.format == DAGFormat.SPARK
79+
assert result.data.schema == from_arrow_schema(schema)
80+
actual_rows = result.data.orderBy("id").toArrow().to_pylist()
81+
assert actual_rows == rows
82+
finally:
83+
for key, value in previous.items():
84+
spark_session.conf.set(key, value)
85+
86+
87+
def test_spark_read_node_keeps_native_spark_dataframe(spark_session):
88+
dataframe = spark_session.range(3)
89+
job = MagicMock(spec=SparkRetrievalJob)
90+
job.to_spark_df.return_value = dataframe
91+
node = SparkReadNode(
92+
"read",
93+
MagicMock(),
94+
ColumnInfo(["id"], [], "event_timestamp", None),
95+
spark_session,
96+
)
97+
with patch(
98+
"feast.infra.compute_engines.spark.nodes.create_offline_store_retrieval_job",
99+
return_value=job,
100+
):
101+
result = node.execute(MagicMock())
102+
assert result.data is dataframe
103+
job.to_arrow.assert_not_called()
104+
105+
19106
def test_spark_transformation_node_executes_udf(spark_session):
20107
# Sample Spark input
21108
df = spark_session.createDataFrame(

‎sdk/python/tests/component/spark/test_spark.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,14 @@
44

55
import pandas as pd
66
import pyarrow as pa
7+
import pytest
8+
from pyspark.sql.pandas.types import from_arrow_schema
79

810
from feast.entity import Entity
911
from feast.feature_view import FeatureView, Field
12+
from feast.infra.data_sources.contrib.iceberg_catalog.iceberg_source import (
13+
IcebergSource,
14+
)
1015
from feast.infra.offline_stores.contrib.spark_offline_store import spark as spark_module
1116
from feast.infra.offline_stores.contrib.spark_offline_store.spark import (
1217
SparkOfflineStore,
@@ -21,6 +26,38 @@
2126
from feast.types import Float32, ValueType
2227

2328

29+
@pytest.mark.parametrize("empty", [False, True])
30+
def test_register_iceberg_source_preserves_arrow_values_and_schema(
31+
spark_session, empty
32+
):
33+
schema = pa.schema(
34+
[pa.field("id", pa.int64(), nullable=False), pa.field("count", pa.int32())]
35+
)
36+
rows = [{"id": 1, "count": 2}, {"id": 2, "count": None}]
37+
table = pa.concat_tables(
38+
[pa.Table.from_pylist(rows[:1], schema), pa.Table.from_pylist(rows[1:], schema)]
39+
)
40+
if empty:
41+
table = table.slice(0, 0)
42+
rows = []
43+
source = IcebergSource(warehouse="test", namespace="test", table="arrow_input")
44+
iceberg_table = MagicMock()
45+
iceberg_table.scan.return_value.to_arrow.return_value = table
46+
view_name = "iceberg_tmp_arrow_input"
47+
try:
48+
with patch.object(
49+
spark_module, "_load_pyiceberg_table", return_value=iceberg_table
50+
):
51+
spark_module._register_iceberg_source_as_temp_view(
52+
spark_session, source, MagicMock()
53+
)
54+
result = spark_session.table(view_name)
55+
assert result.schema == from_arrow_schema(schema)
56+
assert [row.asDict() for row in result.orderBy("id").collect()] == rows
57+
finally:
58+
spark_session.catalog.dropTempView(view_name)
59+
60+
2461
@patch(
2562
"feast.infra.offline_stores.contrib.spark_offline_store.spark.get_spark_session_or_start_new_with_repoconfig"
2663
)

0 commit comments

Comments
 (0)