Visitar URL original
fix(spark): S3/GCS PyArrow filesystem resolution for staging paths by abhijeet-dhumal · Pull Request #6442 · feast-dev/feast · GitHub
Skip to content

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 from
Jun 3, 2026
Merged

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

Changes from 1 commit
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

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
fix(spark): S3/GCS PyArrow filesystem for staging paths 
Signed-off-by: abhijeet-dhumal <abhijeetdhumal652@gmail.com>
  • Loading branch information
@abhijeet-dhumal @ntkathole
abhijeet-dhumal authored and ntkathole committed Jun 3, 2026
commit 28318b109378c6313a7ac256da7370df0f6fd307
53 changes: 42 additions & 11 deletions sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -482,21 +482,52 @@ 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]

if sample.startswith("s3://") or sample.startswith("s3a://"):
import pyarrow.fs as pafs

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("/").replace(
"https://", ""
).replace("http://", "")
kwargs["scheme"] = (
"https" if endpoint.startswith("https") else "http"
)
fs = pafs.S3FileSystem(**kwargs)
stripped = [p.replace("s3a://", "").replace("s3://", "") for p in paths]

@ntkathole ntkathole May 28, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

better use p.removeprefix("s3a://").removeprefix("s3://"), same for https/http and other paths

@abhijeet-dhumal abhijeet-dhumal Jun 1, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the catch! Updated to use removeprefix() for both the S3 path stripping and the endpoint scheme removal.

return fs, stripped

if sample.startswith("gs://"):
import pyarrow.fs as pafs

@ntkathole ntkathole May 28, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

multiple times same import

@abhijeet-dhumal abhijeet-dhumal Jun 1, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point.. Yes.. hoisted import pyarrow.fs as pafs above the S3/GCS branches so it's imported once.


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