From da72500c93e1ff1107a24f1e67eb453a443e4d82 Mon Sep 17 00:00:00 2001 From: Daniel Sola Date: Mon, 29 Sep 2025 12:34:58 -0700 Subject: [PATCH] add timestamp transformer --- flytekit/core/type_engine.py | 6 +++++- flytekit/types/schema/types_pandas.py | 24 +++++++++++++++++++++--- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 6b3760056e..b547d497ab 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -1363,7 +1363,11 @@ def lazy_import_transformers(cls): from flytekit.extras import pydantic_transformer # noqa: F401 if is_imported("pandas"): try: - from flytekit.types.schema.types_pandas import PandasSchemaReader, PandasSchemaWriter # noqa: F401 + from flytekit.types.schema.types_pandas import ( # noqa: F401 + PandasSchemaReader, + PandasSchemaWriter, + PandasTimestampTransformer, + ) except ValueError: logger.debug("Transformer for pandas is already registered.") if is_imported("numpy"): diff --git a/flytekit/types/schema/types_pandas.py b/flytekit/types/schema/types_pandas.py index bff3572cfe..83db243a3d 100644 --- a/flytekit/types/schema/types_pandas.py +++ b/flytekit/types/schema/types_pandas.py @@ -5,9 +5,9 @@ import pandas from flytekit import FlyteContext -from flytekit.core.type_engine import AsyncTypeTransformer, T, TypeEngine -from flytekit.models.literals import Literal, Scalar, Schema -from flytekit.models.types import LiteralType, SchemaType +from flytekit.core.type_engine import AsyncTypeTransformer, SimpleTransformer, T, TypeEngine +from flytekit.models.literals import Literal, Primitive, Scalar, Schema +from flytekit.models.types import LiteralType, SchemaType, SimpleType from flytekit.types.schema import LocalIOSchemaReader, LocalIOSchemaWriter, SchemaEngine, SchemaFormat, SchemaHandler @@ -122,7 +122,25 @@ def to_html(self, ctx: FlyteContext, python_val: pandas.DataFrame, expected_pyth return python_val.describe().to_html() +class PandasTimestampTransformer(SimpleTransformer[pandas.Timestamp]): + """ + Type transformer for pandas.Timestamp objects. + + Converts pandas.Timestamp to/from Flyte DATETIME literals. + """ + + def __init__(self): + super().__init__( + name="pandas.Timestamp", + t=pandas.Timestamp, + lt=LiteralType(simple=SimpleType.DATETIME), + to_literal_transformer=lambda x: Literal(scalar=Scalar(primitive=Primitive(datetime=x.to_pydatetime()))), + from_literal_transformer=lambda x: pandas.Timestamp(x.scalar.primitive.datetime), + ) + + SchemaEngine.register_handler( SchemaHandler("pandas-dataframe-schema", pandas.DataFrame, PandasSchemaReader, PandasSchemaWriter) ) TypeEngine.register(PandasDataFrameTransformer()) +TypeEngine.register(PandasTimestampTransformer())