@@ -844,7 +844,9 @@ def pg_type_to_feast_value_type(type_str: str) -> ValueType:
844
844
return value
845
845
846
846
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" :
848
850
import pyarrow
849
851
850
852
type_map = {
@@ -855,15 +857,15 @@ def feast_value_type_to_pa(feast_type: ValueType) -> "pyarrow.DataType":
855
857
ValueType .STRING : pyarrow .string (),
856
858
ValueType .BYTES : pyarrow .binary (),
857
859
ValueType .BOOL : pyarrow .bool_ (),
858
- ValueType .UNIX_TIMESTAMP : pyarrow .timestamp ("us" ),
860
+ ValueType .UNIX_TIMESTAMP : pyarrow .timestamp (timestamp_unit ),
859
861
ValueType .INT32_LIST : pyarrow .list_ (pyarrow .int32 ()),
860
862
ValueType .INT64_LIST : pyarrow .list_ (pyarrow .int64 ()),
861
863
ValueType .DOUBLE_LIST : pyarrow .list_ (pyarrow .float64 ()),
862
864
ValueType .FLOAT_LIST : pyarrow .list_ (pyarrow .float32 ()),
863
865
ValueType .STRING_LIST : pyarrow .list_ (pyarrow .string ()),
864
866
ValueType .BYTES_LIST : pyarrow .list_ (pyarrow .binary ()),
865
867
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 )),
867
869
ValueType .NULL : pyarrow .null (),
868
870
}
869
871
return type_map [feast_type ]
0 commit comments