GV

GeezerVison

The World Wide Web, restored to 800 pixels of glory.

Home|About|Gaming cabinet|Archives|How it works|Starter sites

On this site

  • Sign in
  • ActionsCI/CD and automation
  • CodespacesCloud development environments
  • All Features
  • Documentation
  • Training
  • Certifications
  • Build with Copilot
  • Customize Copilot
  • Integrations
  • Trending repos
  • GitHub Sponsors
  • Developer articles
  • Enterprises

This view

  • Original page
  • New address

GeezerVison is showing github.com — fix(spark): S3/GCS PyArrow filesystem resolution for staging paths by abhijeet-dhumal · Pull Request #6442 · feast-dev/feast

fix(spark): S3/GCS PyArrow filesystem resolution for staging paths #6442

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.

Sign up for GitHub

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

Merged
ntkathole merged 3 commits into feast-dev:master from abhijeet-dhumal:fix/spark-staging-filesystem
Jun 3, 2026

Uh oh!

There was an error while loading. Please reload this page.
Merged

fix(spark): S3/GCS PyArrow filesystem resolution for staging paths #6442

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
28318b1
fix(spark): S3/GCS PyArrow filesystem for staging paths
abhijeet-dhumal May 27, 2026
bcb7340
style: ruff format spark.py
abhijeet-dhumal May 27, 2026
63adf09
fix(spark): address review feedback on staging filesystem PR
abhijeet-dhumal Jun 1, 2026
File filter
Filter by extension
Clear filters
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
Diff view

There are no files selected for viewing

51 changes: 40 additions & 11 deletions 51 sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py
Show comments
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -482,21 +482,50 @@ def _to_arrow_via_staging(self) -> pyarrow.Table:
if not parquet_paths:
return pyarrow.table({})

normalized_paths = self._normalize_staging_paths(parquet_paths)
dataset = ds.dataset(normalized_paths, format="parquet")
pa_fs, stripped_paths = self._resolve_staging_filesystem(parquet_paths)
dataset = ds.dataset(stripped_paths, format="parquet", filesystem=pa_fs)
return dataset.to_table()

def _normalize_staging_paths(self, paths: List[str]) -> List[str]:
"""Normalize staging paths for PyArrow datasets."""
def _resolve_staging_filesystem(
self, paths: List[str]
) -> Tuple[Optional[pyarrow.fs.FileSystem], List[str]]:
"""Return (pyarrow filesystem, prefix-stripped paths) for staging URIs."""
sample = paths[0]

import pyarrow.fs as pafs

if sample.startswith("s3://") or sample.startswith("s3a://"):
endpoint = os.environ.get("AWS_ENDPOINT_URL_S3") or os.environ.get(
"AWS_S3_ENDPOINT", ""
)
region = getattr(
self._config.offline_store, "region", None
) or os.environ.get("AWS_DEFAULT_REGION", "us-east-1")
kwargs: Dict[str, Any] = {"region": region}
if endpoint:
kwargs["endpoint_override"] = (
endpoint.rstrip("/")
.removeprefix("https://")
.removeprefix("http://")
)
kwargs["scheme"] = "https" if endpoint.startswith("https") else "http"
fs = pafs.S3FileSystem(**kwargs)
stripped = [p.removeprefix("s3a://").removeprefix("s3://") for p in paths]
return fs, stripped

if sample.startswith("gs://"):
fs = pafs.GcsFileSystem()
stripped = [p[len("gs://") :] for p in paths]
return fs, stripped

# Local paths
normalized = []
for path in paths:
if path.startswith("file://"):
normalized.append(path[len("file://") :])
elif "://" in path:
normalized.append(path)
for p in paths:
if p.startswith("file://"):
normalized.append(p[len("file://") :])
else:
normalized.append(path)
return normalized
normalized.append(p)
return None, normalized

def to_feast_df(
self,
Expand Down
134 changes: 134 additions & 0 deletions 134 ...ts/unit/infra/offline_stores/contrib/spark_offline_store/test_spark_staging_filesystem.py
Show comments
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,134 @@
"""
Unit tests for SparkRetrievalJob._resolve_staging_filesystem.

Verifies that the correct PyArrow filesystem and prefix-stripped paths
are returned for S3, S3A, GCS, file://, and plain local paths.
"""

from unittest.mock import MagicMock, patch

import pytest

from feast.infra.offline_stores.contrib.spark_offline_store.spark import (
SparkRetrievalJob,
)


@pytest.fixture()
def retrieval_job():
"""Minimal SparkRetrievalJob with a mock config that has no offline_store region."""
job = object.__new__(SparkRetrievalJob)
config = MagicMock()
config.offline_store.region = None
job._config = config
return job


class TestResolveS3Filesystem:
def test_s3_scheme_returns_s3_filesystem(self, retrieval_job):
with patch("pyarrow.fs.S3FileSystem") as mock_s3:
mock_s3.return_value = MagicMock(name="s3fs")
fs, paths = retrieval_job._resolve_staging_filesystem(
["s3://my-bucket/path/a.parquet", "s3://my-bucket/path/b.parquet"]
)
mock_s3.assert_called_once()
assert fs is mock_s3.return_value
assert paths == ["my-bucket/path/a.parquet", "my-bucket/path/b.parquet"]

def test_s3a_scheme_strips_prefix(self, retrieval_job):
with patch("pyarrow.fs.S3FileSystem") as mock_s3:
mock_s3.return_value = MagicMock(name="s3fs")
fs, paths = retrieval_job._resolve_staging_filesystem(
["s3a://bucket/dir/file.parquet"]
)
assert paths == ["bucket/dir/file.parquet"]

def test_s3_with_minio_endpoint(self, retrieval_job, monkeypatch):
monkeypatch.setenv("AWS_ENDPOINT_URL_S3", "http://minio.local:9000")
monkeypatch.setenv("AWS_DEFAULT_REGION", "us-east-1")
with patch("pyarrow.fs.S3FileSystem") as mock_s3:
mock_s3.return_value = MagicMock(name="s3fs")
retrieval_job._resolve_staging_filesystem(["s3://bucket/file.parquet"])
call_kwargs = mock_s3.call_args[1]
assert call_kwargs["endpoint_override"] == "minio.local:9000"
assert call_kwargs["scheme"] == "http"

def test_s3_with_https_endpoint(self, retrieval_job, monkeypatch):
monkeypatch.setenv("AWS_ENDPOINT_URL_S3", "https://s3.custom.corp")
with patch("pyarrow.fs.S3FileSystem") as mock_s3:
mock_s3.return_value = MagicMock(name="s3fs")
retrieval_job._resolve_staging_filesystem(["s3://bucket/file.parquet"])
call_kwargs = mock_s3.call_args[1]
assert call_kwargs["endpoint_override"] == "s3.custom.corp"
assert call_kwargs["scheme"] == "https"

def test_s3_falls_back_to_aws_s3_endpoint_env(self, retrieval_job, monkeypatch):
monkeypatch.delenv("AWS_ENDPOINT_URL_S3", raising=False)
monkeypatch.setenv("AWS_S3_ENDPOINT", "http://legacy-minio:9000")
with patch("pyarrow.fs.S3FileSystem") as mock_s3:
mock_s3.return_value = MagicMock(name="s3fs")
retrieval_job._resolve_staging_filesystem(["s3://bucket/file.parquet"])
call_kwargs = mock_s3.call_args[1]
assert "endpoint_override" in call_kwargs

def test_s3_no_endpoint_no_override(self, retrieval_job, monkeypatch):
monkeypatch.delenv("AWS_ENDPOINT_URL_S3", raising=False)
monkeypatch.delenv("AWS_S3_ENDPOINT", raising=False)
with patch("pyarrow.fs.S3FileSystem") as mock_s3:
mock_s3.return_value = MagicMock(name="s3fs")
retrieval_job._resolve_staging_filesystem(["s3://bucket/file.parquet"])
call_kwargs = mock_s3.call_args[1]
assert "endpoint_override" not in call_kwargs
assert "scheme" not in call_kwargs

def test_s3_region_from_offline_store_config(self, retrieval_job):
retrieval_job._config.offline_store.region = "eu-west-1"
with patch("pyarrow.fs.S3FileSystem") as mock_s3:
mock_s3.return_value = MagicMock(name="s3fs")
retrieval_job._resolve_staging_filesystem(["s3://bucket/file.parquet"])
call_kwargs = mock_s3.call_args[1]
assert call_kwargs["region"] == "eu-west-1"

def test_s3_region_fallback_to_env(self, retrieval_job, monkeypatch):
retrieval_job._config.offline_store.region = None
monkeypatch.setenv("AWS_DEFAULT_REGION", "ap-southeast-1")
with patch("pyarrow.fs.S3FileSystem") as mock_s3:
mock_s3.return_value = MagicMock(name="s3fs")
retrieval_job._resolve_staging_filesystem(["s3://bucket/file.parquet"])
call_kwargs = mock_s3.call_args[1]
assert call_kwargs["region"] == "ap-southeast-1"


class TestResolveGCSFilesystem:
def test_gs_scheme_returns_gcs_filesystem(self, retrieval_job):
with patch("pyarrow.fs.GcsFileSystem") as mock_gcs:
mock_gcs.return_value = MagicMock(name="gcsfs")
fs, paths = retrieval_job._resolve_staging_filesystem(
["gs://my-bucket/path/a.parquet", "gs://my-bucket/path/b.parquet"]
)
mock_gcs.assert_called_once()
assert fs is mock_gcs.return_value
assert paths == ["my-bucket/path/a.parquet", "my-bucket/path/b.parquet"]


class TestResolveLocalFilesystem:
def test_file_scheme_stripped(self, retrieval_job):
fs, paths = retrieval_job._resolve_staging_filesystem(
["file:///tmp/staging/a.parquet"]
)
assert fs is None
assert paths == ["/tmp/staging/a.parquet"]

def test_plain_local_path_unchanged(self, retrieval_job):
fs, paths = retrieval_job._resolve_staging_filesystem(
["/tmp/staging/a.parquet", "/tmp/staging/b.parquet"]
)
assert fs is None
assert paths == ["/tmp/staging/a.parquet", "/tmp/staging/b.parquet"]

def test_mixed_file_and_plain_paths(self, retrieval_job):
fs, paths = retrieval_job._resolve_staging_filesystem(
["file:///tmp/a.parquet", "/tmp/b.parquet"]
)
assert fs is None
assert paths == ["/tmp/a.parquet", "/tmp/b.parquet"]
Loading
Oops, something went wrong.

Uh oh!

There was an error while loading. Please reload this page.
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.
You are visitor 01402503 — GeezerVison does not store the pages you look up.
Best viewed with Netscape Navigator 4.7 or Internet Explorer 6 at 800×600. No cookies required.