GeezerVison is showing github.com — perf: Avoid pandas copy in Spark reads by Pranjulcr7 · Pull Request #6878 · feast-dev/feast
perf: Avoid pandas copy in Spark reads
#6878
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement . We’ll occasionally send you account related emails.
Already on GitHub?
Sign in
to your account
Uh oh!
There was an error while loading. Please reload this page.
Show file tree
Hide file tree
Changes from all commits
Commits
File filter
Filter by extension
.py
(4)
All 1 file type selected
Viewed files
Conversations
Failed to load comments.
Loading
Uh oh!
There was an error while loading. Please reload this page.
Jump to
Jump to file
Failed to load files.
Loading
Uh oh!
There was an error while loading. Please reload this page.
Diff view
Hide whitespace
Diff view
Hide whitespace
There are no files selected for viewing
Original file line number
Diff line number
Diff line change
Expand Up
@@ -6,7 +6,6 @@
import pandas as pd
from pyspark.sql import DataFrame, SparkSession, Window
from pyspark.sql import functions as F
from pyspark.sql.pandas.types import from_arrow_schema
from pyspark.sql.types import (
ArrayType,
BinaryType,
Expand Down
Expand Up
@@ -192,13 +191,7 @@ def execute(self, context: ExecutionContext) -> DAGValue:
spark_df = cast(SparkRetrievalJob, retrieval_job).to_spark_df()
else:
arrow_table = retrieval_job.to_arrow()
if arrow_table.num_rows == 0:
spark_schema = from_arrow_schema(arrow_table.schema)
spark_df = self.spark_session.createDataFrame(
self.spark_session.sparkContext.emptyRDD(), schema=spark_schema
)
else:
spark_df = self.spark_session.createDataFrame(arrow_table.to_pandas())
spark_df = self.spark_session.createDataFrame(arrow_table)
return DAGValue(
data=spark_df,
Expand Down
Original file line number
Diff line number
Diff line change
Expand Up
@@ -1256,7 +1256,7 @@ def _register_iceberg_source_as_temp_view(
iceberg_table = _load_pyiceberg_table(data_source)
arrow_table = iceberg_table.scan().to_arrow()
df = spark_session.createDataFrame(arrow_table.to_pandas() )
df = spark_session.createDataFrame(arrow_table)
df.createOrReplaceTempView(view_name)
Expand Down
Original file line number
Diff line number
Diff line change
@@ -1,5 +1,9 @@
from datetime import datetime, timedelta
from unittest.mock import MagicMock
from datetime import datetime, timedelta, timezone
from unittest.mock import MagicMock, patch
import pyarrow as pa
import pytest
from pyspark.sql.pandas.types import from_arrow_schema
from feast.aggregation import Aggregation
from feast.infra.compute_engines.dag.context import ColumnInfo, ExecutionContext
Expand All
@@ -9,13 +13,96 @@
SparkAggregationNode,
SparkDedupNode,
SparkJoinNode,
SparkReadNode,
SparkTransformationNode,
)
from feast.infra.offline_stores.contrib.spark_offline_store.spark import (
SparkRetrievalJob,
)
from tests.example_repos.example_feature_repo_with_bfvs import (
driver,
)
@pytest.mark.parametrize("empty", [False, True])
@pytest.mark.parametrize("arrow_enabled", [False, True])
def test_spark_read_node_preserves_arrow_values_and_schema(
spark_session, empty, arrow_enabled
):
schema = pa.schema(
[
pa.field("id", pa.int64(), nullable=False),
pa.field("count", pa.int32()),
pa.field("score", pa.float32()),
pa.field("name", pa.string()),
pa.field("active", pa.bool_()),
pa.field("event_timestamp", pa.timestamp("us", tz="UTC")),
]
)
rows = [
{"id": 1, "count": 2, "score": 0.5, "name": "first", "active": True},
{"id": 2, "count": None, "score": None, "name": None, "active": None},
{"id": 3, "count": -1, "score": 1.5, "name": "last", "active": False},
]
for row in rows:
row["event_timestamp"] = datetime(
2026, 1, 1, 12, 0, 0, 123456, tzinfo=timezone.utc
)
table = pa.concat_tables(
[pa.Table.from_pylist(rows[:1], schema), pa.Table.from_pylist(rows[1:], schema)]
)
if empty:
table = table.slice(0, 0)
rows = []
job = MagicMock()
job.to_arrow.return_value = table
node = SparkReadNode(
"read",
MagicMock(),
ColumnInfo(["id"], ["count"], "event_timestamp", None),
spark_session,
)
config = {
"spark.sql.execution.arrow.pyspark.enabled": str(arrow_enabled).lower(),
"spark.sql.session.timeZone": "UTC",
}
previous = {key: spark_session.conf.get(key) for key in config}
for key, value in config.items():
spark_session.conf.set(key, value)
try:
with patch(
"feast.infra.compute_engines.spark.nodes.create_offline_store_retrieval_job",
Comment thread
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more .
Spark collects TimestampType as naive datetime in the session timezone — .astimezone(timezone.utc) on a naive datetime assumes the system TZ is correct. Consider setting spark.sql.session.timeZone to UTC in this test for deterministic CI results.
Sorry, something went wrong.
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more .
Set the session timezone to UTC and compare Arrow rows, so the check avoids naive datetime conversion. Confirmed with TZ=UTC and TZ=Pacific/Honolulu in 4f1fcf9.
Sorry, something went wrong.
Uh oh!
There was an error while loading. Please reload this page.
return_value=job,
):
result = node.execute(MagicMock())
assert result.format == DAGFormat.SPARK
assert result.data.schema == from_arrow_schema(schema)
actual_rows = result.data.orderBy("id").toArrow().to_pylist()
assert actual_rows == rows
finally:
for key, value in previous.items():
spark_session.conf.set(key, value)
def test_spark_read_node_keeps_native_spark_dataframe(spark_session):
dataframe = spark_session.range(3)
job = MagicMock(spec=SparkRetrievalJob)
job.to_spark_df.return_value = dataframe
node = SparkReadNode(
"read",
MagicMock(),
ColumnInfo(["id"], [], "event_timestamp", None),
spark_session,
)
with patch(
"feast.infra.compute_engines.spark.nodes.create_offline_store_retrieval_job",
return_value=job,
):
result = node.execute(MagicMock())
assert result.data is dataframe
job.to_arrow.assert_not_called()
def test_spark_transformation_node_executes_udf(spark_session):
# Sample Spark input
df = spark_session.createDataFrame(
Expand Down
Original file line number
Diff line number
Diff line change
Expand Up
@@ -4,9 +4,14 @@
import pandas as pd
import pyarrow as pa
import pytest
from pyspark.sql.pandas.types import from_arrow_schema
from feast.entity import Entity
from feast.feature_view import FeatureView, Field
from feast.infra.data_sources.contrib.iceberg_catalog.iceberg_source import (
IcebergSource,
)
from feast.infra.offline_stores.contrib.spark_offline_store import spark as spark_module
from feast.infra.offline_stores.contrib.spark_offline_store.spark import (
SparkOfflineStore,
Expand All
@@ -21,6 +26,38 @@
from feast.types import Float32, ValueType
@pytest.mark.parametrize("empty", [False, True])
def test_register_iceberg_source_preserves_arrow_values_and_schema(
spark_session, empty
):
schema = pa.schema(
[pa.field("id", pa.int64(), nullable=False), pa.field("count", pa.int32())]
)
rows = [{"id": 1, "count": 2}, {"id": 2, "count": None}]
table = pa.concat_tables(
[pa.Table.from_pylist(rows[:1], schema), pa.Table.from_pylist(rows[1:], schema)]
)
if empty:
table = table.slice(0, 0)
rows = []
source = IcebergSource(warehouse="test", namespace="test", table="arrow_input")
iceberg_table = MagicMock()
iceberg_table.scan.return_value.to_arrow.return_value = table
view_name = "iceberg_tmp_arrow_input"
try:
with patch.object(
spark_module, "_load_pyiceberg_table", return_value=iceberg_table
):
spark_module._register_iceberg_source_as_temp_view(
spark_session, source, MagicMock()
)
result = spark_session.table(view_name)
assert result.schema == from_arrow_schema(schema)
assert [row.asDict() for row in result.orderBy("id").collect()] == rows
finally:
spark_session.catalog.dropTempView(view_name)
@patch(
"feast.infra.offline_stores.contrib.spark_offline_store.spark.get_spark_session_or_start_new_with_repoconfig"
)
Expand Down
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.