flyte 0.0.1b0__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/io/{structured_dataset/structured_dataset.py → _dataframe/dataframe.py}

@@ -1,20 +1,17 @@
 from __future__ import annotations
 
 import _datetime
-import asyncio
 import collections
 import types
 import typing
 from abc import ABC, abstractmethod
-from dataclasses import dataclass, field, is_dataclass
+from dataclasses import is_dataclass
 from typing import Any, ClassVar, Coroutine, Dict, Generic, List, Optional, Type, Union
 
-import msgpack
-from flyteidl.core import literals_pb2, types_pb2
+from flyteidl2.core import literals_pb2, types_pb2
 from fsspec.utils import get_protocol
-from mashumaro.mixins.json import DataClassJSONMixin
 from mashumaro.types import SerializableType
-from pydantic import model_serializer, model_validator
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_serializer, model_validator
 from typing_extensions import Annotated, TypeAlias, get_args, get_origin
 
 import flyte.storage as storage
@@ -35,58 +32,66 @@ else:
     pd = lazy_module("pandas")
     pa = lazy_module("pyarrow")
 
-T = typing.TypeVar("T")  # StructuredDataset type or a dataframe type
+T = typing.TypeVar("T")  # DataFrame type or a dataframe type
 DF = typing.TypeVar("DF")  # Dataframe type
 
-# For specifying the storage formats of StructuredDatasets. It's just a string, nothing fancy.
-StructuredDatasetFormat: TypeAlias = str
+# For specifying the storage formats of DataFrames. It's just a string, nothing fancy.
+DataFrameFormat: TypeAlias = str
 
 # Storage formats
-PARQUET: StructuredDatasetFormat = "parquet"
-CSV: StructuredDatasetFormat = "csv"
-GENERIC_FORMAT: StructuredDatasetFormat = ""
+PARQUET: DataFrameFormat = "parquet"
+CSV: DataFrameFormat = "csv"
+GENERIC_FORMAT: DataFrameFormat = ""
 GENERIC_PROTOCOL: str = "generic protocol"
 
 
-@dataclass
-class StructuredDataset(SerializableType, DataClassJSONMixin):
+class DataFrame(BaseModel, SerializableType):
     """
-    This is the user facing StructuredDataset class. Please don't confuse it with the literals.StructuredDataset
+    This is the user facing DataFrame class. Please don't confuse it with the literals.StructuredDataset
     class (that is just a model, a Python class representation of the protobuf).
     """
 
-    uri: typing.Optional[str] = field(default=None)
-    file_format: typing.Optional[str] = field(default=GENERIC_FORMAT)
+    uri: typing.Optional[str] = Field(default=None)
+    format: typing.Optional[str] = Field(default=GENERIC_FORMAT)
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    # Private attributes that are not part of the Pydantic model schema
+    _raw_df: typing.Optional[typing.Any] = PrivateAttr(default=None)
+    _metadata: typing.Optional[literals_pb2.StructuredDatasetMetadata] = PrivateAttr(default=None)
+    _literal_sd: Optional[literals_pb2.StructuredDataset] = PrivateAttr(default=None)
+    _dataframe_type: Optional[Type[Any]] = PrivateAttr(default=None)
+    _already_uploaded: bool = PrivateAttr(default=False)
 
     # loop manager is working better than synchronicity for some reason, was getting an error but may be an easy fix
     def _serialize(self) -> Dict[str, Optional[str]]:
         # dataclass case
         lt = TypeEngine.to_literal_type(type(self))
-        engine = StructuredDatasetTransformerEngine()
+        engine = DataFrameTransformerEngine()
         lv = loop_manager.run_sync(engine.to_literal, self, type(self), lt)
-        sd = StructuredDataset(uri=lv.scalar.structured_dataset.uri)
-        sd.file_format = lv.scalar.structured_dataset.metadata.structured_dataset_type.format
+        sd = DataFrame(uri=lv.scalar.structured_dataset.uri)
+        sd.format = lv.scalar.structured_dataset.metadata.structured_dataset_type.format
         return {
             "uri": sd.uri,
-            "file_format": sd.file_format,
+            "format": sd.format,
         }
 
     @classmethod
-    def _deserialize(cls, value) -> "StructuredDataset":
+    def _deserialize(cls, value) -> DataFrame:
         uri = value.get("uri", None)
-        file_format = value.get("file_format", None)
+        format_val = value.get("format", None)
 
         if uri is None:
-            raise ValueError("StructuredDataset's uri and file format should not be None")
+            raise ValueError("DataFrame's uri and file format should not be None")
 
-        engine = StructuredDatasetTransformerEngine()
+        engine = DataFrameTransformerEngine()
         return loop_manager.run_sync(
             engine.to_python_value,
             literals_pb2.Literal(
                 scalar=literals_pb2.Scalar(
                     structured_dataset=literals_pb2.StructuredDataset(
                         metadata=literals_pb2.StructuredDatasetMetadata(
-                            structured_dataset_type=types_pb2.StructuredDatasetType(format=file_format)
+                            structured_dataset_type=types_pb2.StructuredDatasetType(format=format_val)
                         ),
                         uri=uri,
                     )
@@ -96,28 +101,28 @@ class StructuredDataset(SerializableType, DataClassJSONMixin):
         )
 
     @model_serializer
-    def serialize_structured_dataset(self) -> Dict[str, Optional[str]]:
+    def serialize_dataframe(self) -> Dict[str, Optional[str]]:
         lt = TypeEngine.to_literal_type(type(self))
-        sde = StructuredDatasetTransformerEngine()
+        sde = DataFrameTransformerEngine()
         lv = loop_manager.run_sync(sde.to_literal, self, type(self), lt)
         return {
             "uri": lv.scalar.structured_dataset.uri,
-            "file_format": lv.scalar.structured_dataset.metadata.structured_dataset_type.format,
+            "format": lv.scalar.structured_dataset.metadata.structured_dataset_type.format,
         }
 
     @model_validator(mode="after")
-    def deserialize_structured_dataset(self, info) -> StructuredDataset:
+    def deserialize_dataframe(self, info) -> DataFrame:
         if info.context is None or info.context.get("deserialize") is not True:
             return self
 
-        engine = StructuredDatasetTransformerEngine()
+        engine = DataFrameTransformerEngine()
         return loop_manager.run_sync(
             engine.to_python_value,
             literals_pb2.Literal(
                 scalar=literals_pb2.Scalar(
                     structured_dataset=literals_pb2.StructuredDataset(
                         metadata=literals_pb2.StructuredDatasetMetadata(
-                            structured_dataset_type=types_pb2.StructuredDatasetType(format=self.file_format)
+                            structured_dataset_type=types_pb2.StructuredDatasetType(format=self.format)
                         ),
                         uri=self.uri,
                     )
@@ -134,30 +139,46 @@ class StructuredDataset(SerializableType, DataClassJSONMixin):
     def column_names(cls) -> typing.List[str]:
         return [k for k, v in cls.columns().items()]
 
-    def __init__(
-        self,
-        dataframe: typing.Optional[typing.Any] = None,
+    @classmethod
+    def from_df(
+        cls,
+        val: typing.Optional[typing.Any] = None,
         uri: typing.Optional[str] = None,
-        metadata: typing.Optional[literals_pb2.StructuredDatasetMetadata] = None,
+    ) -> DataFrame:
+        """
+        Wrapper to create a DataFrame from a dataframe.
+        The reason this is implemented as a wrapper instead of a full translation invoking
+        the type engine and the encoders is because there's too much information in the type
+        signature of the task that we don't want the user to have to replicate.
+        """
+        instance = cls(uri=uri)
+        instance._raw_df = val
+        return instance
+
+    @classmethod
+    def from_existing_remote(
+        cls,
+        remote_path: str,
+        format: typing.Optional[str] = None,
         **kwargs,
-    ):
-        self._dataframe = dataframe
-        # Make these fields public, so that the dataclass transformer can set a value for it
-        # https://github.com/flyteorg/flytekit/blob/bcc8541bd6227b532f8462563fe8aac902242b21/flytekit/core/type_engine.py#L298
-        self.uri = uri
-        # When dataclass_json runs from_json, we need to set it here, otherwise the format will be empty string
-        self.file_format = kwargs["file_format"] if "file_format" in kwargs else GENERIC_FORMAT
-        # This is a special attribute that indicates if the data was either downloaded or uploaded
-        self._metadata = metadata
-        # This is not for users to set, the transformer will set this.
-        self._literal_sd: Optional[literals_pb2.StructuredDataset] = None
-        # Not meant for users to set, will be set by an open() call
-        self._dataframe_type: Optional[DF] = None  # type: ignore
-        self._already_uploaded = False
+    ) -> "DataFrame":
+        """
+        Create a DataFrame reference from an existing remote dataframe.
+
+        Args:
+            remote_path: The remote path to the existing dataframe
+            format: Format of the stored dataframe
+
+        Example:
+            ```python
+            df = DataFrame.from_existing_remote("s3://bucket/data.parquet", format="parquet")
+            ```
+        """
+        return cls(uri=remote_path, format=format or GENERIC_FORMAT, **kwargs)
 
     @property
-    def dataframe(self) -> Optional[DF]:
-        return self._dataframe
+    def val(self) -> Optional[DF]:
+        return self._raw_df
 
     @property
     def metadata(self) -> Optional[literals_pb2.StructuredDatasetMetadata]:
@@ -168,18 +189,18 @@ class StructuredDataset(SerializableType, DataClassJSONMixin):
         return self._literal_sd
 
     def open(self, dataframe_type: Type[DF]):
-        from flyte.io.structured_dataset import lazy_import_structured_dataset_handler
-
         """
         Load the handler if needed. For the use case like:
         @task
-        def t1(sd: StructuredDataset):
+        def t1(df: DataFrame):
           import pandas as pd
-          sd.open(pd.DataFrame).all()
+          df.open(pd.DataFrame).all()
 
-        pandas is imported inside the task, so pandas handler won't be loaded during deserialization in type engine.
+        pandas is imported inside the task, so panda handler won't be loaded during deserialization in type engine.
         """
-        lazy_import_structured_dataset_handler()
+        from flyte.io._dataframe import lazy_import_dataframe_handler
+
+        lazy_import_dataframe_handler()
         self._dataframe_type = dataframe_type
         return self
 
@@ -187,22 +208,22 @@ class StructuredDataset(SerializableType, DataClassJSONMixin):
         if self._dataframe_type is None:
             raise ValueError("No dataframe type set. Use open() to set the local dataframe type you want to use.")
 
-        if self.uri is not None and self.dataframe is None:
-            expected = TypeEngine.to_literal_type(StructuredDataset)
+        if self.uri is not None and self.val is None:
+            expected = TypeEngine.to_literal_type(DataFrame)
             await self._set_literal(expected)
 
         return await flyte_dataset_transformer.open_as(self.literal, self._dataframe_type, self.metadata)
 
     async def _set_literal(self, expected: types_pb2.LiteralType) -> None:
         """
-        Explicitly set the StructuredDataset Literal to handle the following cases:
+        Explicitly set the DataFrame Literal to handle the following cases:
 
-        1. Read a dataframe from a StructuredDataset with an uri, for example:
+        1. Read the content from a DataFrame with an uri, for example:
 
         @task
-        def return_sd() -> StructuredDataset:
-            sd = StructuredDataset(uri="s3://my-s3-bucket/s3_flyte_dir/df.parquet", file_format="parquet")
-            df = sd.open(pd.DataFrame).all()
+        def return_df() -> DataFrame:
+            df = DataFrame(uri="s3://my-s3-bucket/s3_flyte_dir/df.parquet", format="parquet")
+            df = df.open(pd.DataFrame).all()
             return df
 
         For details, please refer to this issue: https://github.com/flyteorg/flyte/issues/5954.
@@ -212,14 +233,14 @@ class StructuredDataset(SerializableType, DataClassJSONMixin):
 
         For details, please refer to this issue: https://github.com/flyteorg/flyte/issues/5956.
         """
-        to_literal = await flyte_dataset_transformer.to_literal(self, StructuredDataset, expected)
+        to_literal = await flyte_dataset_transformer.to_literal(self, DataFrame, expected)
         self._literal_sd = to_literal.scalar.structured_dataset
         if self.metadata is None:
             self._metadata = self._literal_sd.metadata
 
     async def set_literal(self, expected: types_pb2.LiteralType) -> None:
         """
-        A public wrapper method to set the StructuredDataset Literal.
+        A public wrapper method to set the DataFrame Literal.
 
         This method provides external access to the internal _set_literal method.
         """
@@ -244,6 +265,9 @@ def flatten_dict(sub_dict: dict, parent_key: str = "") -> typing.Dict:
             fields = getattr(value, "__dataclass_fields__")
             d = {k: v.type for k, v in fields.items()}
             result.update(flatten_dict(sub_dict=d, parent_key=current_key))
+        elif hasattr(value, "model_fields"):  # Pydantic model
+            d = {k: v.annotation for k, v in value.model_fields.items()}
+            result.update(flatten_dict(sub_dict=d, parent_key=current_key))
         else:
             result[current_key] = value
     return result
@@ -256,7 +280,7 @@ def extract_cols_and_format(
     Helper function, just used to iterate through Annotations and extract out the following information:
       - base type, if not Annotated, it will just be the type that was passed in.
       - column information, as a collections.OrderedDict,
-      - the storage format, as a ``StructuredDatasetFormat`` (str),
+      - the storage format, as a ``DataFrameFormat`` (str),
       - pa.lib.Schema
 
     If more than one of any type of thing is found, an error will be raised.
@@ -286,7 +310,7 @@ def extract_cols_and_format(
                 d = collections.OrderedDict()
                 d.update(aa)
                 ordered_dict_cols = d
-            elif isinstance(aa, StructuredDatasetFormat):
+            elif isinstance(aa, DataFrameFormat):
                 if fmt != "":
                     raise ValueError(f"A format was already specified {fmt}, cannot use {aa}")
                 fmt = aa
@@ -305,7 +329,7 @@ def extract_cols_and_format(
     return t, ordered_dict_cols, fmt, pa_schema
 
 
-class StructuredDatasetEncoder(ABC, Generic[T]):
+class DataFrameEncoder(ABC, Generic[T]):
     def __init__(
         self,
         python_type: Type[T],
@@ -314,10 +338,10 @@ class StructuredDatasetEncoder(ABC, Generic[T]):
     ):
         """
         Extend this abstract class, implement the encode function, and register your concrete class with the
-        StructuredDatasetTransformerEngine class in order for the core flytekit type engine to handle
+        DataFrameTransformerEngine class in order for the core flytekit type engine to handle
         dataframe libraries. This is the encoding interface, meaning it is used when there is a Python value that the
         flytekit type engine is trying to convert into a Flyte Literal. For the other way, see
-        the StructuredDatasetEncoder
+        the DataFrameEncoder
 
         :param python_type: The dataframe class in question that you want to register this encoder with
         :param protocol: A prefix representing the storage driver (e.g. 's3, 'gs', 'bq', etc.). You can use either
@@ -347,7 +371,7 @@ class StructuredDatasetEncoder(ABC, Generic[T]):
     @abstractmethod
     async def encode(
         self,
-        structured_dataset: StructuredDataset,
+        dataframe: DataFrame,
         structured_dataset_type: types_pb2.StructuredDatasetType,
     ) -> literals_pb2.StructuredDataset:
         """
@@ -357,20 +381,20 @@ class StructuredDatasetEncoder(ABC, Generic[T]):
         the
         # TODO: Do we need to add a flag to indicate if it was wrapped by the transformer or by the user?
 
-        :param structured_dataset: This is a StructuredDataset wrapper object. See more info above.
-        :param structured_dataset_type: This the StructuredDatasetType, as found in the LiteralType of the interface
+        :param dataframe: This is a DataFrame wrapper object. See more info above.
+        :param structured_dataset_type: This the DataFrameType, as found in the LiteralType of the interface
           of the task that invoked this encoding call. It is passed along to encoders so that authors of encoders
-          can include it in the returned literals.StructuredDataset. See the IDL for more information on why this
+          can include it in the returned literals.DataFrame. See the IDL for more information on why this
           literal in particular carries the type information along with it. If the encoder doesn't supply it, it will
           also be filled in after the encoder runs by the transformer engine.
-        :return: This function should return a StructuredDataset literal object. Do not confuse this with the
-          StructuredDataset wrapper class used as input to this function - that is the user facing Python class.
-          This function needs to return the IDL StructuredDataset.
+        :return: This function should return a DataFrame literal object. Do not confuse this with the
+          DataFrame wrapper class used as input to this function - that is the user facing Python class.
+          This function needs to return the IDL DataFrame.
         """
         raise NotImplementedError
 
 
-class StructuredDatasetDecoder(ABC, Generic[DF]):
+class DataFrameDecoder(ABC, Generic[DF]):
     def __init__(
         self,
         python_type: Type[DF],
@@ -380,9 +404,9 @@ class StructuredDatasetDecoder(ABC, Generic[DF]):
     ):
         """
         Extend this abstract class, implement the decode function, and register your concrete class with the
-        StructuredDatasetTransformerEngine class in order for the core flytekit type engine to handle
+        DataFrameTransformerEngine class in order for the core flytekit type engine to handle
         dataframe libraries. This is the decoder interface, meaning it is used when there is a Flyte Literal value,
-        and we have to get a Python value out of it. For the other way, see the StructuredDatasetEncoder
+        and we have to get a Python value out of it. For the other way, see the DataFrameEncoder
 
         :param python_type: The dataframe class in question that you want to register this decoder with
         :param protocol: A prefix representing the storage driver (e.g. 's3, 'gs', 'bq', etc.). You can use either
@@ -419,8 +443,8 @@ class StructuredDatasetDecoder(ABC, Generic[DF]):
         This is code that will be called by the dataset transformer engine to ultimately translate from a Flyte Literal
         value into a Python instance.
 
-        :param flyte_value: This will be a Flyte IDL StructuredDataset Literal - do not confuse this with the
-          StructuredDataset class defined also in this module.
+        :param flyte_value: This will be a Flyte IDL DataFrame Literal - do not confuse this with the
+          DataFrame class defined also in this module.
         :param current_task_metadata: Metadata object containing the type (and columns if any) for the currently
            executing task. This type may have more or less information than the type information bundled
            inside the incoming flyte_value.
@@ -459,19 +483,19 @@ def get_supported_types():
 class DuplicateHandlerError(ValueError): ...
 
 
-class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
+class DataFrameTransformerEngine(TypeTransformer[DataFrame]):
     """
     Think of this transformer as a higher-level meta transformer that is used for all the dataframe types.
     If you are bringing a custom data frame type, or any data frame type, to flytekit, instead of
     registering with the main type engine, you should register with this transformer instead.
     """
 
-    ENCODERS: ClassVar[Dict[Type, Dict[str, Dict[str, StructuredDatasetEncoder]]]] = {}
-    DECODERS: ClassVar[Dict[Type, Dict[str, Dict[str, StructuredDatasetDecoder]]]] = {}
+    ENCODERS: ClassVar[Dict[Type, Dict[str, Dict[str, DataFrameEncoder]]]] = {}
+    DECODERS: ClassVar[Dict[Type, Dict[str, Dict[str, DataFrameDecoder]]]] = {}
     DEFAULT_PROTOCOLS: ClassVar[Dict[Type, str]] = {}
     DEFAULT_FORMATS: ClassVar[Dict[Type, str]] = {}
 
-    Handlers = Union[StructuredDatasetEncoder, StructuredDatasetDecoder]
+    Handlers = Union[DataFrameEncoder, DataFrameDecoder]
     Renderers: ClassVar[Dict[Type, Renderable]] = {}
 
     @classmethod
@@ -527,17 +551,17 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
 
     @classmethod
     def get_encoder(cls, df_type: Type, protocol: str, format: str):
-        return cls._finder(StructuredDatasetTransformerEngine.ENCODERS, df_type, protocol, format)
+        return cls._finder(DataFrameTransformerEngine.ENCODERS, df_type, protocol, format)
 
     @classmethod
-    def get_decoder(cls, df_type: Type, protocol: str, format: str) -> StructuredDatasetDecoder:
-        return cls._finder(StructuredDatasetTransformerEngine.DECODERS, df_type, protocol, format)
+    def get_decoder(cls, df_type: Type, protocol: str, format: str) -> DataFrameDecoder:
+        return cls._finder(DataFrameTransformerEngine.DECODERS, df_type, protocol, format)
 
     @classmethod
     def _handler_finder(cls, h: Handlers, protocol: str) -> Dict[str, Handlers]:
-        if isinstance(h, StructuredDatasetEncoder):
+        if isinstance(h, DataFrameEncoder):
             top_level = cls.ENCODERS
-        elif isinstance(h, StructuredDatasetDecoder):
+        elif isinstance(h, DataFrameDecoder):
             top_level = cls.DECODERS  # type: ignore
         else:
             raise TypeError(f"We don't support this type of handler {h}")
@@ -548,7 +572,7 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
         return top_level[h.python_type][protocol]  # type: ignore
 
     def __init__(self):
-        super().__init__("StructuredDataset Transformer", StructuredDataset)
+        super().__init__("DataFrame Transformer", DataFrame)
         self._type_assertions_enabled = False
 
     @classmethod
@@ -568,7 +592,7 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
         Call this with any Encoder or Decoder to register it with the flytekit type system. If your handler does not
         specify a protocol (e.g. s3, gs, etc.) field, then
 
-        :param h: The StructuredDatasetEncoder or StructuredDatasetDecoder you wish to register with this transformer.
+        :param h: The DataFrameEncoder or DataFrameDecoder you wish to register with this transformer.
         :param default_for_type: If set, when a user returns from a task an instance of the dataframe the handler
           handles, e.g. ``return pd.DataFrame(...)``, not wrapped around the ``StructuredDataset`` object, we will
           use this handler's protocol and format as the default, effectively saying that this handler will be called.
@@ -582,7 +606,7 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
         :param default_storage_for_type: Same as above but only for the storage format. Error if already set,
           unless override is specified.
         """
-        if not (isinstance(h, StructuredDatasetEncoder) or isinstance(h, StructuredDatasetDecoder)):
+        if not (isinstance(h, DataFrameEncoder) or isinstance(h, DataFrameDecoder)):
             raise TypeError(f"We don't support this type of handler {h}")
 
         if h.protocol is None:
@@ -623,17 +647,21 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
                 f"Already registered a handler for {(h.python_type, protocol, h.supported_format)}"
             )
         lowest_level[h.supported_format] = h
-        logger.debug(f"Registered {h} as handler for {h.python_type}, protocol {protocol}, fmt {h.supported_format}")
+        logger.debug(
+            f"Registered {h.__class__.__name__} as handler for {h.python_type.__class__.__name__},"
+            f" protocol {protocol}, fmt {h.supported_format}"
+        )
 
         if (default_format_for_type or default_for_type) and h.supported_format != GENERIC_FORMAT:
             if h.python_type in cls.DEFAULT_FORMATS and not override:
                 if cls.DEFAULT_FORMATS[h.python_type] != h.supported_format:
                     logger.info(
-                        f"Not using handler {h} with format {h.supported_format}"
-                        f" as default for {h.python_type}, {cls.DEFAULT_FORMATS[h.python_type]} already specified."
+                        f"Not using handler {h.__class__.__name__} with format {h.supported_format}"
+                        f" as default for {h.python_type.__class__.__name__},"
+                        f" {cls.DEFAULT_FORMATS[h.python_type]} already specified."
                     )
             else:
-                logger.debug(f"Use {type(h).__name__} as default handler for {h.python_type}.")
+                logger.debug(f"Use {type(h).__name__} as default handler for {h.python_type.__class__.__name__}.")
                 cls.DEFAULT_FORMATS[h.python_type] = h.supported_format
         if default_storage_for_type or default_for_type:
             if h.protocol in cls.DEFAULT_PROTOCOLS and not override:
@@ -648,27 +676,27 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
         # Register with the type engine as well
         # The semantics as of now are such that it doesn't matter which order these transformers are loaded in, as
         # long as the older Pandas/FlyteSchema transformer do not also specify the override
-        engine = StructuredDatasetTransformerEngine()
+        engine = DataFrameTransformerEngine()
         TypeEngine.register_additional_type(engine, h.python_type, override=True)
 
-    def assert_type(self, t: Type[StructuredDataset], v: typing.Any):
+    def assert_type(self, t: Type[DataFrame], v: typing.Any):
         return
 
     async def to_literal(
         self,
-        python_val: Union[StructuredDataset, typing.Any],
-        python_type: Union[Type[StructuredDataset], Type],
+        python_val: Union[DataFrame, typing.Any],
+        python_type: Union[Type[DataFrame], Type],
         expected: types_pb2.LiteralType,
     ) -> literals_pb2.Literal:
         # Make a copy in case we need to hand off to encoders, since we can't be sure of mutations.
-        python_type, *attrs = extract_cols_and_format(python_type)
+        python_type, *_attrs = extract_cols_and_format(python_type)
         sdt = types_pb2.StructuredDatasetType(format=self.DEFAULT_FORMATS.get(python_type, GENERIC_FORMAT))
 
-        if issubclass(python_type, StructuredDataset) and not isinstance(python_val, StructuredDataset):
+        if issubclass(python_type, DataFrame) and not isinstance(python_val, DataFrame):
             # Catch a common mistake
             raise TypeTransformerFailedError(
-                f"Expected a StructuredDataset instance, but got {type(python_val)} instead."
-                f" Did you forget to wrap your dataframe in a StructuredDataset instance?"
+                f"Expected a DataFrame instance, but got {type(python_val)} instead."
+                f" Did you forget to wrap your dataframe in a DataFrame instance?"
             )
 
         if expected and expected.structured_dataset_type:
@@ -679,60 +707,60 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
                 external_schema_bytes=expected.structured_dataset_type.external_schema_bytes,
             )
 
-        # If the type signature has the StructuredDataset class, it will, or at least should, also be a
-        # StructuredDataset instance.
-        if isinstance(python_val, StructuredDataset):
+        # If the type signature has the DataFrame class, it will, or at least should, also be a
+        # DataFrame instance.
+        if isinstance(python_val, DataFrame):
             # There are three cases that we need to take care of here.
 
-            # 1. A task returns a StructuredDataset that was just a passthrough input. If this happens
-            # then return the original literals.StructuredDataset without invoking any encoder
+            # 1. A task returns a DataFrame that was just a passthrough input. If this happens
+            # then return the original literals.DataFrame without invoking any encoder
             #
             # Ex.
-            #   def t1(dataset: Annotated[StructuredDataset, my_cols]) -> Annotated[StructuredDataset, my_cols]:
+            #   def t1(dataset: Annotated[DataFrame, my_cols]) -> Annotated[DataFrame, my_cols]:
             #       return dataset
             if python_val._literal_sd is not None:
                 if python_val._already_uploaded:
                     return literals_pb2.Literal(scalar=literals_pb2.Scalar(structured_dataset=python_val._literal_sd))
-                if python_val.dataframe is not None:
+                if python_val.val is not None:
                     raise ValueError(
-                        f"Shouldn't have specified both literal {python_val._literal_sd}"
-                        f" and dataframe {python_val.dataframe}"
+                        f"Shouldn't have specified both literal {python_val._literal_sd} and dataframe {python_val.val}"
                     )
                 return literals_pb2.Literal(scalar=literals_pb2.Scalar(structured_dataset=python_val._literal_sd))
 
-            # 2. A task returns a python StructuredDataset with an uri.
-            # Note: this case is also what happens we start a local execution of a task with a python StructuredDataset.
-            #  It gets converted into a literal first, then back into a python StructuredDataset.
+            # 2. A task returns a python DataFrame with an uri.
+            # Note: this case is also what happens we start a local execution of a task with a python DataFrame.
+            #  It gets converted into a literal first, then back into a python DataFrame.
             #
             # Ex.
-            #   def t2(uri: str) -> Annotated[StructuredDataset, my_cols]
-            #       return StructuredDataset(uri=uri)
-            if python_val.dataframe is None:
+            #   def t2(uri: str) -> Annotated[DataFrame, my_cols]
+            #       return DataFrame(uri=uri)
+            if python_val.val is None:
                 uri = python_val.uri
-                file_format = python_val.file_format
+                format_val = python_val.format
 
                 # Check the user-specified uri
                 if not uri:
                     raise ValueError(f"If dataframe is not specified, then the uri should be specified. {python_val}")
                 if not storage.is_remote(uri):
-                    uri = await storage.put(uri)
+                    uri = await storage.put(uri, recursive=True)
 
-                # Check the user-specified file_format
-                # When users specify file_format for a StructuredDataset, the file_format should be retained
+                # Check the user-specified format
+                # When users specify format for a DataFrame, the format should be retained
                 # conditionally. For details, please refer to https://github.com/flyteorg/flyte/issues/6096.
                 # Following illustrates why we can't always copy the user-specified file_format over:
                 #
                 # @task
-                # def modify_format(sd: Annotated[StructuredDataset, {}, "task-format"]) -> StructuredDataset:
-                #     return sd
+                # def modify_format(df: Annotated[DataFrame, {}, "task-format"]) -> DataFrame:
+                #     return df
                 #
-                # sd = StructuredDataset(uri="s3://my-s3-bucket/df.parquet", file_format="user-format")
-                # sd2 = modify_format(sd=sd)
+                # df = DataFrame(uri="s3://my-s3-bucket/df.parquet", format="user-format")
+                # df2 = modify_format(df=df)
                 #
-                # In this case, we expect sd2.file_format to be task-format (as shown in Annotated), not user-format.
-                # If we directly copy the user-specified file_format over, the type hint information will be missing.
-                if sdt.format == GENERIC_FORMAT and file_format != GENERIC_FORMAT:
-                    sdt.format = file_format
+                # In this case, we expect the df2.format to be task-format (as shown in Annotated),
+                # not user-format. If we directly copy the user-specified format over,
+                # the type hint information will be missing.
+                if sdt.format == GENERIC_FORMAT and format_val != GENERIC_FORMAT:
+                    sdt.format = format_val
 
                 sd_model = literals_pb2.StructuredDataset(
                     uri=uri,
@@ -740,9 +768,9 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
                 )
                 return literals_pb2.Literal(scalar=literals_pb2.Scalar(structured_dataset=sd_model))
 
-            # 3. This is the third and probably most common case. The python StructuredDataset object wraps a dataframe
+            # 3. This is the third and probably most common case. The python DataFrame object wraps a dataframe
             # that we will need to invoke an encoder for. Figure out which encoder to call and invoke it.
-            df_type = type(python_val.dataframe)
+            df_type = type(python_val.val)
             protocol = self._protocol_from_type_or_prefix(df_type, python_val.uri)
 
             return await self.encode(
@@ -760,8 +788,9 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
             structured_dataset_type=expected.structured_dataset_type if expected else None
         )
 
-        sd = StructuredDataset(dataframe=python_val, metadata=meta)
-        return await self.encode(sd, python_type, protocol, fmt, sdt)
+        fdf = DataFrame.from_df(val=python_val)
+        fdf._metadata = meta
+        return await self.encode(fdf, python_type, protocol, fmt, sdt)
 
     def _protocol_from_type_or_prefix(self, df_type: Type, uri: Optional[str] = None) -> str:
         """
@@ -782,16 +811,16 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
 
     async def encode(
         self,
-        sd: StructuredDataset,
+        df: DataFrame,
         df_type: Type,
         protocol: str,
         format: str,
         structured_literal_type: types_pb2.StructuredDatasetType,
     ) -> literals_pb2.Literal:
-        handler: StructuredDatasetEncoder
+        handler: DataFrameEncoder
         handler = self.get_encoder(df_type, protocol, format)
 
-        sd_model = await handler.encode(sd, structured_literal_type)
+        sd_model = await handler.encode(df, structured_literal_type)
         # This block is here in case the encoder did not set the type information in the metadata. Since this literal
         # is special in that it carries around the type itself, we want to make sure the type info therein is at
         # least as good as the type of the interface.
@@ -807,75 +836,16 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
         lit = literals_pb2.Literal(scalar=literals_pb2.Scalar(structured_dataset=sd_model))
 
         # Because the handler.encode may have uploaded something, and because the sd may end up living inside a
-        # dataclass, we need to modify any uploaded flyte:// urls here.
-        modify_literal_uris(lit)  # todo: verify that this can be removed.
-        sd._literal_sd = sd_model
-        sd._already_uploaded = True
+        # dataclass, we need to modify any uploaded flyte:// urls here. Needed here even though the Type engine
+        # already does this because the DataframeTransformerEngine may be called directly.
+        modify_literal_uris(lit)
+        df._literal_sd = sd_model
+        df._already_uploaded = True
         return lit
 
-    # pr: han-ru: can this be removed if we make StructuredDataset a pydantic model?
-    def dict_to_structured_dataset(
-        self, dict_obj: typing.Dict[str, str], expected_python_type: Type[T] | StructuredDataset
-    ) -> T | StructuredDataset:
-        uri = dict_obj.get("uri", None)
-        file_format = dict_obj.get("file_format", None)
-
-        if uri is None:
-            raise ValueError("StructuredDataset's uri and file format should not be None")
-
-        # Instead of using python native StructuredDataset, we need to build a literals.StructuredDataset
-        # The reason is that _literal_sd of python sd is accessed when task output LiteralMap is
-        # converted back to flyteidl. Hence, _literal_sd must have to_flyte_idl method
-        # See https://github.com/flyteorg/flytekit/blob/f938661ff8413219d1bea77f6914a58c302d5c6c/flytekit/bin/entrypoint.py#L326
-        # For details, please refer to this issue: https://github.com/flyteorg/flyte/issues/5956.
-        sdt = types_pb2.StructuredDatasetType(format=file_format)
-        metad = literals_pb2.StructuredDatasetMetadata(structured_dataset_type=sdt)
-        sd_literal = literals_pb2.StructuredDataset(uri=uri, metadata=metad)
-
-        return asyncio.run(
-            StructuredDatasetTransformerEngine().to_python_value(
-                literals_pb2.Literal(scalar=literals_pb2.Scalar(structured_dataset=sd_literal)),
-                expected_python_type,
-            )
-        )
-
-    def from_binary_idl(
-        self, binary_idl_object: literals_pb2.Binary, expected_python_type: Type[T] | StructuredDataset
-    ) -> T | StructuredDataset:
-        """
-        If the input is from flytekit, the Life Cycle will be as follows:
-
-        Life Cycle:
-        binary IDL                 -> resolved binary         -> bytes                   -> expected Python object
-        (flytekit customized          (propeller processing)     (flytekit binary IDL)      (flytekit customized
-        serialization)                                                                       deserialization)
-
-        Example Code:
-        @dataclass
-        class DC:
-            sd: StructuredDataset
-
-        @workflow
-        def wf(dc: DC):
-            t_sd(dc.sd)
-
-        Note:
-        - The deserialization is the same as put a structured dataset in a dataclass,
-          which will deserialize by the mashumaro's API.
-
-        Related PR:
-        - Title: Override Dataclass Serialization/Deserialization Behavior for FlyteTypes via Mashumaro
-        - Link: https://github.com/flyteorg/flytekit/pull/2554
-        """
-        if binary_idl_object.tag == MESSAGEPACK:
-            python_val = msgpack.loads(binary_idl_object.value)
-            return self.dict_to_structured_dataset(dict_obj=python_val, expected_python_type=expected_python_type)
-        else:
-            raise TypeTransformerFailedError(f"Unsupported binary format: `{binary_idl_object.tag}`")
-
     async def to_python_value(
-        self, lv: literals_pb2.Literal, expected_python_type: Type[T] | StructuredDataset
-    ) -> T | StructuredDataset:
+        self, lv: literals_pb2.Literal, expected_python_type: Type[T] | DataFrame
+    ) -> T | DataFrame:
         """
         The only tricky thing with converting a Literal (say the output of an earlier task), to a Python value at
         the start of a task execution, is the column subsetting behavior. For example, if you have,
@@ -906,14 +876,13 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
         |                             | the running task's signature.           |                                      |
         +-----------------------------+-----------------------------------------+--------------------------------------+
         """
-        # Handle dataclass attribute access
         if lv.HasField("scalar") and lv.scalar.HasField("binary"):
-            return self.from_binary_idl(lv.scalar.binary, expected_python_type)
+            raise TypeTransformerFailedError("Attribute access unsupported.")
 
         # Detect annotations and extract out all the relevant information that the user might supply
-        expected_python_type, column_dict, storage_fmt, pa_schema = extract_cols_and_format(expected_python_type)
+        expected_python_type, column_dict, _storage_fmt, _pa_schema = extract_cols_and_format(expected_python_type)
 
-        # Start handling for StructuredDataset scalars, first look at the columns
+        # Start handling for DataFrame scalars, first look at the columns
         incoming_columns = lv.scalar.structured_dataset.metadata.structured_dataset_type.columns
 
         # If the incoming literal, also doesn't have columns, then we just have an empty list, so initialize here
@@ -935,30 +904,27 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
         )
         metad = literals_pb2.StructuredDatasetMetadata(structured_dataset_type=new_sdt)
 
-        # A StructuredDataset type, for example
-        #   t1(input_a: StructuredDataset)  # or
-        #   t1(input_a: Annotated[StructuredDataset, my_cols])
-        if issubclass(expected_python_type, StructuredDataset):
-            sd = expected_python_type(
-                dataframe=None,
-                # Note here that the type being passed in
-                metadata=metad,
-            )
-            sd._literal_sd = lv.scalar.structured_dataset
-            sd.file_format = metad.structured_dataset_type.format
-            return sd
-
-        # If the requested type was not a StructuredDataset, then it means it was a plain dataframe type, which means
+        # A DataFrame type, for example
+        #   t1(input_a: DataFrame)  # or
+        #   t1(input_a: Annotated[DataFrame, my_cols])
+        if issubclass(expected_python_type, DataFrame):
+            fdf = DataFrame(format=metad.structured_dataset_type.format, uri=lv.scalar.structured_dataset.uri)
+            fdf._already_uploaded = True
+            fdf._literal_sd = lv.scalar.structured_dataset
+            fdf._metadata = metad
+            return fdf
+
+        # If the requested type was not a flyte.DataFrame, then it means it was a raw dataframe type, which means
         # we should do the opening/downloading and whatever else it might entail right now. No iteration option here.
         return await self.open_as(lv.scalar.structured_dataset, df_type=expected_python_type, updated_metadata=metad)
 
     def to_html(self, python_val: typing.Any, expected_python_type: Type[T]) -> str:
-        if isinstance(python_val, StructuredDataset):
-            if python_val.dataframe is not None:
-                df = python_val.dataframe
+        if isinstance(python_val, DataFrame):
+            if python_val.val is not None:
+                df = python_val.val
             else:
                 # Here we only render column information by default instead of opening the structured dataset.
-                col = typing.cast(StructuredDataset, python_val).columns()
+                col = typing.cast(DataFrame, python_val).columns()
                 dataframe = pd.DataFrame(col, ["column type"])
                 return dataframe.to_html()  # type: ignore
         else:
@@ -1004,11 +970,12 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
     def _get_dataset_column_literal_type(self, t: Type) -> types_pb2.LiteralType:
         if t in get_supported_types():
             return get_supported_types()[t]
-        if hasattr(t, "__origin__") and t.__origin__ is list:
+        origin = getattr(t, "__origin__", None)
+        if origin is list:
             return types_pb2.LiteralType(collection_type=self._get_dataset_column_literal_type(t.__args__[0]))
-        if hasattr(t, "__origin__") and t.__origin__ is dict:
+        if origin is dict:
             return types_pb2.LiteralType(map_value_type=self._get_dataset_column_literal_type(t.__args__[1]))
-        raise AssertionError(f"type {t} is currently not supported by StructuredDataset")
+        raise AssertionError(f"type {t} is currently not supported by DataFrame")
 
     def _convert_ordered_dict_of_columns_to_list(
         self, column_map: typing.Optional[typing.OrderedDict[str, Type]]
@@ -1022,10 +989,8 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
             converted_cols.append(types_pb2.StructuredDatasetType.DatasetColumn(name=k, literal_type=lt))
         return converted_cols
 
-    def _get_dataset_type(
-        self, t: typing.Union[Type[StructuredDataset], typing.Any]
-    ) -> types_pb2.StructuredDatasetType:
-        original_python_type, column_map, storage_format, pa_schema = extract_cols_and_format(t)  # type: ignore
+    def _get_dataset_type(self, t: typing.Union[Type[DataFrame], typing.Any]) -> types_pb2.StructuredDatasetType:
+        _original_python_type, column_map, storage_format, pa_schema = extract_cols_and_format(t)  # type: ignore
 
         # Get the column information
         converted_cols: typing.List[types_pb2.StructuredDatasetType.DatasetColumn] = (
@@ -1039,7 +1004,7 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
             external_schema_bytes=typing.cast(pa.lib.Schema, pa_schema).to_string().encode() if pa_schema else None,
         )
 
-    def get_literal_type(self, t: typing.Union[Type[StructuredDataset], typing.Any]) -> types_pb2.LiteralType:
+    def get_literal_type(self, t: typing.Union[Type[DataFrame], typing.Any]) -> types_pb2.LiteralType:
         """
         Provide a concrete implementation so that writers of custom dataframe handlers since there's nothing that
         special about the literal type. Any dataframe type will always be associated with the structured dataset type.
@@ -1049,13 +1014,13 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
         """
         return types_pb2.LiteralType(structured_dataset_type=self._get_dataset_type(t))
 
-    def guess_python_type(self, literal_type: types_pb2.LiteralType) -> Type[StructuredDataset]:
+    def guess_python_type(self, literal_type: types_pb2.LiteralType) -> Type[DataFrame]:
         # todo: technically we should return the dataframe type specified in the constructor, but to do that,
         #   we'd have to store that, which we don't do today. See possibly #1363
         if literal_type.HasField("structured_dataset_type"):
-            return StructuredDataset
-        raise ValueError(f"StructuredDatasetTransformerEngine cannot reverse {literal_type}")
+            return DataFrame
+        raise ValueError(f"DataFrameTransformerEngine cannot reverse {literal_type}")
 
 
-flyte_dataset_transformer = StructuredDatasetTransformerEngine()
+flyte_dataset_transformer = DataFrameTransformerEngine()
 TypeEngine.register(flyte_dataset_transformer)