diff --git a/sdk/python/feast/infra/offline_stores/snowflake.py b/sdk/python/feast/infra/offline_stores/snowflake.py index 241627ba01..50f92164cc 100644 --- a/sdk/python/feast/infra/offline_stores/snowflake.py +++ b/sdk/python/feast/infra/offline_stores/snowflake.py @@ -19,7 +19,6 @@ import numpy as np import pandas as pd import pyarrow -import pyarrow as pa from pydantic import Field, StrictStr from pydantic.typing import Literal from pytz import utc @@ -410,7 +409,7 @@ def _to_df_internal(self) -> pd.DataFrame: return df - def _to_arrow_internal(self) -> pa.Table: + def _to_arrow_internal(self) -> pyarrow.Table: with self._query_generator() as query: pa_table = execute_snowflake_statement( @@ -423,7 +422,7 @@ def _to_arrow_internal(self) -> pa.Table: else: empty_result = execute_snowflake_statement(self.snowflake_conn, query) - return pa.Table.from_pandas( + return pyarrow.Table.from_pandas( pd.DataFrame(columns=[md.name for md in empty_result.description]) ) diff --git a/sdk/python/feast/type_map.py b/sdk/python/feast/type_map.py index 8a817358b3..2cb1c4fefb 100644 --- a/sdk/python/feast/type_map.py +++ b/sdk/python/feast/type_map.py @@ -15,6 +15,7 @@ from collections import defaultdict from datetime import datetime, timezone from typing import ( + TYPE_CHECKING, Any, Dict, Iterator, @@ -31,7 +32,6 @@ import numpy as np import pandas as pd -import pyarrow from google.protobuf.timestamp_pb2 import Timestamp from feast.protos.feast.types.Value_pb2 import ( @@ -46,6 +46,9 @@ from feast.protos.feast.types.Value_pb2 import Value as ProtoValue from feast.value_type import ListType, ValueType +if TYPE_CHECKING: + import pyarrow + # null timestamps get converted to -9223372036854775808 NULL_TIMESTAMP_INT_VALUE = np.datetime64("NaT").astype(int) @@ -554,7 +557,7 @@ def mssql_to_feast_value_type(mssql_type_as_str: str) -> ValueType: return type_map[mssql_type_as_str.lower()] -def pa_to_mssql_type(pa_type: pyarrow.DataType) -> str: +def pa_to_mssql_type(pa_type: "pyarrow.DataType") -> str: # PyArrow types: https://arrow.apache.org/docs/python/api/datatypes.html # MS Sql types: https://docs.microsoft.com/en-us/sql/t-sql/data-types/data-types-transact-sql?view=sql-server-ver16 pa_type_as_str = str(pa_type).lower()