Skip to content

Commit 9b227d7

Browse files
authored
fix: Fix timestamp consistency in push api (feast-dev#3614)
Signed-off-by: Jiwon Park <bakjeeone@hotmail.com>
1 parent bfb26c3 commit 9b227d7

File tree

3 files changed

+9
-6
lines changed

3 files changed

+9
-6
lines changed

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -363,7 +363,7 @@ def offline_write_batch(
363363
assert isinstance(feature_view.batch_source, BigQuerySource)
364364

365365
pa_schema, column_names = offline_utils.get_pyarrow_schema_from_batch_source(
366-
config, feature_view.batch_source
366+
config, feature_view.batch_source, timestamp_unit="ns"
367367
)
368368
if column_names != table.column_names:
369369
raise ValueError(

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

+3-2
Original file line numberDiff line numberDiff line change
@@ -232,7 +232,7 @@ def get_offline_store_from_config(offline_store_config: Any) -> OfflineStore:
232232

233233

234234
def get_pyarrow_schema_from_batch_source(
235-
config: RepoConfig, batch_source: DataSource
235+
config: RepoConfig, batch_source: DataSource, timestamp_unit: str = "us"
236236
) -> Tuple[pa.Schema, List[str]]:
237237
"""Returns the pyarrow schema and column names for the given batch source."""
238238
column_names_and_types = batch_source.get_table_column_names_and_types(config)
@@ -244,7 +244,8 @@ def get_pyarrow_schema_from_batch_source(
244244
(
245245
column_name,
246246
feast_value_type_to_pa(
247-
batch_source.source_datatype_to_feast_value_type()(column_type)
247+
batch_source.source_datatype_to_feast_value_type()(column_type),
248+
timestamp_unit=timestamp_unit,
248249
),
249250
)
250251
)

sdk/python/feast/type_map.py

+5-3
Original file line numberDiff line numberDiff line change
@@ -844,7 +844,9 @@ def pg_type_to_feast_value_type(type_str: str) -> ValueType:
844844
return value
845845

846846

847-
def feast_value_type_to_pa(feast_type: ValueType) -> "pyarrow.DataType":
847+
def feast_value_type_to_pa(
848+
feast_type: ValueType, timestamp_unit: str = "us"
849+
) -> "pyarrow.DataType":
848850
import pyarrow
849851

850852
type_map = {
@@ -855,15 +857,15 @@ def feast_value_type_to_pa(feast_type: ValueType) -> "pyarrow.DataType":
855857
ValueType.STRING: pyarrow.string(),
856858
ValueType.BYTES: pyarrow.binary(),
857859
ValueType.BOOL: pyarrow.bool_(),
858-
ValueType.UNIX_TIMESTAMP: pyarrow.timestamp("us"),
860+
ValueType.UNIX_TIMESTAMP: pyarrow.timestamp(timestamp_unit),
859861
ValueType.INT32_LIST: pyarrow.list_(pyarrow.int32()),
860862
ValueType.INT64_LIST: pyarrow.list_(pyarrow.int64()),
861863
ValueType.DOUBLE_LIST: pyarrow.list_(pyarrow.float64()),
862864
ValueType.FLOAT_LIST: pyarrow.list_(pyarrow.float32()),
863865
ValueType.STRING_LIST: pyarrow.list_(pyarrow.string()),
864866
ValueType.BYTES_LIST: pyarrow.list_(pyarrow.binary()),
865867
ValueType.BOOL_LIST: pyarrow.list_(pyarrow.bool_()),
866-
ValueType.UNIX_TIMESTAMP_LIST: pyarrow.list_(pyarrow.timestamp("us")),
868+
ValueType.UNIX_TIMESTAMP_LIST: pyarrow.list_(pyarrow.timestamp(timestamp_unit)),
867869
ValueType.NULL: pyarrow.null(),
868870
}
869871
return type_map[feast_type]

0 commit comments

Comments
 (0)