Created
December 20, 2023 14:04
-
-
Save thomasaarholt/f206aa0bd254153f133a5cd7b9848dfc to your computer and use it in GitHub Desktop.
Serialize and deserialize polars datatypes to json
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import json | |
import polars as pl | |
def dtype_to_json(dtype: pl.DataType) -> str: | |
return json.dumps(str(dtype)) | |
def json_to_dtype(json_dtype_str: str) -> pl.DataType: | |
from polars.datatypes.classes import ( # noqa F401 | |
Array, | |
Binary, | |
Boolean, | |
Categorical, | |
Date, | |
Datetime, | |
Decimal, | |
Duration, | |
Enum, | |
Float32, | |
Float64, | |
Int8, | |
Int16, | |
Int32, | |
Int64, | |
List, | |
Null, | |
Object, | |
Struct, | |
Time, | |
UInt8, | |
UInt16, | |
UInt32, | |
UInt64, | |
Unknown, | |
Utf8, | |
) | |
return eval(json.loads(json_dtype_str)) | |
# tests | |
dtypes = [ | |
pl.Datetime(time_zone="UTC"), | |
pl.List(pl.Array(pl.Float32(), 3)), | |
pl.Float32(), | |
pl.Struct({"a": pl.Int8(), "b": pl.List(pl.Utf8())}), | |
] | |
for dtype in dtypes: | |
assert dtype == json_to_dtype(dtype_to_json(dtype)) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment