Skip to content

Commit 5986998

Browse files
committed
Add type annotations
Signed-off-by: Willem Pienaar <git@willem.co>
1 parent eda2764 commit 5986998

File tree

2 files changed

+10
-7
lines changed

2 files changed

+10
-7
lines changed

sdk/python/feast/infra/offline_stores/file.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,7 @@ def evaluate_historical_retrieval():
112112
)
113113

114114
# Read offline parquet data in pyarrow format.
115-
filesystem, path = FileSource.prepare_path(
115+
filesystem, path = FileSource.create_filesystem_and_path(
116116
feature_view.batch_source.path,
117117
feature_view.batch_source.file_options.s3_endpoint_override,
118118
)
@@ -242,7 +242,7 @@ def pull_latest_from_table_or_query(
242242

243243
# Create lazy function that is only called from the RetrievalJob object
244244
def evaluate_offline_job():
245-
filesystem, path = FileSource.prepare_path(
245+
filesystem, path = FileSource.create_filesystem_and_path(
246246
data_source.path, data_source.file_options.s3_endpoint_override
247247
)
248248
source_df = pd.read_parquet(path, filesystem=filesystem)

sdk/python/feast/infra/offline_stores/file_source.py

+8-5
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from typing import Callable, Dict, Iterable, Optional, Tuple
22

3-
from pyarrow import fs
3+
from pyarrow._fs import FileSystem
4+
from pyarrow._s3fs import S3FileSystem
45
from pyarrow.parquet import ParquetFile
56

67
from feast import type_map
@@ -138,7 +139,7 @@ def source_datatype_to_feast_value_type() -> Callable[[str], ValueType]:
138139
def get_table_column_names_and_types(
139140
self, config: RepoConfig
140141
) -> Iterable[Tuple[str, str]]:
141-
filesystem, path = FileSource.prepare_path(
142+
filesystem, path = FileSource.create_filesystem_and_path(
142143
self.path, self._file_options.s3_endpoint_override
143144
)
144145
schema = ParquetFile(
@@ -147,12 +148,14 @@ def get_table_column_names_and_types(
147148
return zip(schema.names, map(str, schema.types))
148149

149150
@staticmethod
150-
def prepare_path(path: str, s3_endpoint_override: str):
151+
def create_filesystem_and_path(
152+
path: str, s3_endpoint_override: str
153+
) -> Tuple[Optional[FileSystem], str]:
151154
if path.startswith("s3://"):
152-
s3 = fs.S3FileSystem(
155+
s3fs = S3FileSystem(
153156
endpoint_override=s3_endpoint_override if s3_endpoint_override else None
154157
)
155-
return s3, path.replace("s3://", "")
158+
return s3fs, path.replace("s3://", "")
156159
else:
157160
return None, path
158161

0 commit comments

Comments
 (0)