Source code for flytekit.types.structured.structured_dataset

from __future__ import annotations

import collections
import importlib
import os
import re
import types
import typing
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Dict, Generator, Optional, Type, Union

import _datetime
import numpy as _np
import pandas
import pandas as pd
import pyarrow
import pyarrow as pa

if importlib.util.find_spec("pyspark") is not None:
    import pyspark
if importlib.util.find_spec("polars") is not None:
    import polars as pl
from dataclasses_json import config, dataclass_json
from marshmallow import fields
from typing_extensions import Annotated, TypeAlias, get_args, get_origin

from flytekit.core.context_manager import FlyteContext, FlyteContextManager
from flytekit.core.type_engine import TypeEngine, TypeTransformer
from flytekit.loggers import logger
from flytekit.models import literals
from flytekit.models import types as type_models
from flytekit.models.literals import Literal, Scalar, StructuredDatasetMetadata
from flytekit.models.types import LiteralType, SchemaType, StructuredDatasetType

T = typing.TypeVar("T")  # StructuredDataset type or a dataframe type
DF = typing.TypeVar("DF")  # Dataframe type

# Protocols
BIGQUERY = "bq"
S3 = "s3"
ABFS = "abfs"
GCS = "gs"
LOCAL = "/"

# For specifying the storage formats of StructuredDatasets. It's just a string, nothing fancy.
StructuredDatasetFormat: TypeAlias = str

# Storage formats
PARQUET: StructuredDatasetFormat = "parquet"


[docs]@dataclass_json @dataclass class StructuredDataset(object): """ This is the user facing StructuredDataset 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[os.PathLike] = field(default=None, metadata=config(mm_field=fields.String())) file_format: typing.Optional[str] = field(default=PARQUET, metadata=config(mm_field=fields.String())) DEFAULT_FILE_FORMAT = PARQUET
[docs] @classmethod def columns(cls) -> typing.Dict[str, typing.Type]: return {}
[docs] @classmethod 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, uri: Optional[str] = None, metadata: typing.Optional[literals.StructuredDatasetMetadata] = 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 # 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.StructuredDataset] = None # Not meant for users to set, will be set by an open() call self._dataframe_type = None @property def dataframe(self) -> Type[typing.Any]: return self._dataframe @property def metadata(self) -> Optional[StructuredDatasetMetadata]: return self._metadata @property def literal(self) -> Optional[literals.StructuredDataset]: return self._literal_sd
[docs] def open(self, dataframe_type: Type[DF]): self._dataframe_type = dataframe_type return self
[docs] def all(self) -> DF: if self._dataframe_type is None: raise ValueError("No dataframe type set. Use open() to set the local dataframe type you want to use.") ctx = FlyteContextManager.current_context() return flyte_dataset_transformer.open_as(ctx, self.literal, self._dataframe_type, self.metadata)
[docs] def iter(self) -> Generator[DF, None, None]: if self._dataframe_type is None: raise ValueError("No dataframe type set. Use open() to set the local dataframe type you want to use.") ctx = FlyteContextManager.current_context() return flyte_dataset_transformer.iter_as( ctx, self.literal, self._dataframe_type, updated_metadata=self.metadata )
def extract_cols_and_format( t: typing.Any, ) -> typing.Tuple[Type[T], Optional[typing.OrderedDict[str, Type]], Optional[str], Optional[pa.lib.Schema]]: """ 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), - pa.lib.Schema If more than one of any type of thing is found, an error will be raised. If no instances of a given type are found, then None will be returned. If we add more things, we should put all the returned items in a dataclass instead of just a tuple. :param t: The incoming type which may or may not be Annotated :return: Tuple representing the original type, optional OrderedDict of columns, optional str for the format, optional pyarrow Schema """ fmt = None ordered_dict_cols = None pa_schema = None if get_origin(t) is Annotated: base_type, *annotate_args = get_args(t) for aa in annotate_args: if isinstance(aa, StructuredDatasetFormat): if fmt is not None: raise ValueError(f"A format was already specified {fmt}, cannot use {aa}") fmt = aa elif isinstance(aa, collections.OrderedDict): if ordered_dict_cols is not None: raise ValueError(f"Column information was already found {ordered_dict_cols}, cannot use {aa}") ordered_dict_cols = aa elif isinstance(aa, pyarrow.Schema): if pa_schema is not None: raise ValueError(f"Arrow schema was already found {pa_schema}, cannot use {aa}") pa_schema = aa return base_type, ordered_dict_cols, fmt, pa_schema # We return None as the format instead of parquet or something because the transformer engine may find # a better default for the given dataframe type. return t, ordered_dict_cols, fmt, pa_schema
[docs]class StructuredDatasetEncoder(ABC): def __init__(self, python_type: Type[T], protocol: str, supported_format: Optional[str] = None): """ 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 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 :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 "s3" or "s3://". They are the same since the "://" will just be stripped by the constructor. :param supported_format: Arbitrary string representing the format. If not supplied then an empty string will be used. An empty string implies that the encoder works with any format. If the format being asked for does not exist, the transformer enginer will look for the "" endcoder instead and write a warning. """ self._python_type = python_type self._protocol = protocol.replace("://", "") self._supported_format = supported_format or "" @property def python_type(self) -> Type[T]: return self._python_type @property def protocol(self) -> str: return self._protocol @property def supported_format(self) -> str: return self._supported_format
[docs] @abstractmethod def encode( self, ctx: FlyteContext, structured_dataset: StructuredDataset, structured_dataset_type: StructuredDatasetType, ) -> literals.StructuredDataset: """ Even if the user code returns a plain dataframe instance, the dataset transformer engine will wrap the incoming dataframe with defaults set for that dataframe type. This simplifies this function's interface as a lot of data that could be specified by the user using the # TODO: Do we need to add a flag to indicate if it was wrapped by the transformer or by the user? :param ctx: :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 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 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. """ raise NotImplementedError
[docs]class StructuredDatasetDecoder(ABC): def __init__(self, python_type: Type[DF], protocol: str, supported_format: Optional[str] = None): """ 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 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 :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 "s3" or "s3://". They are the same since the "://" will just be stripped by the constructor. :param supported_format: Arbitrary string representing the format. If not supplied then an empty string will be used. An empty string implies that the decoder works with any format. If the format being asked for does not exist, the transformer enginer will look for the "" decoder instead and write a warning. """ self._python_type = python_type self._protocol = protocol.replace("://", "") self._supported_format = supported_format or "" @property def python_type(self) -> Type[DF]: return self._python_type @property def protocol(self) -> str: return self._protocol @property def supported_format(self) -> str: return self._supported_format
[docs] @abstractmethod def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, current_task_metadata: StructuredDatasetMetadata, ) -> Union[DF, Generator[DF, None, None]]: """ 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 ctx: A FlyteContext, useful in accessing the filesystem and other attributes :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 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. :return: This function can either return an instance of the dataframe that this decoder handles, or an iterator of those dataframes. """ raise NotImplementedError
def protocol_prefix(uri: str) -> str: g = re.search(r"([\w]+)://.*", uri) if g and g.groups(): return g.groups()[0] return LOCAL def convert_schema_type_to_structured_dataset_type( column_type: int, ) -> int: if column_type == SchemaType.SchemaColumn.SchemaColumnType.INTEGER: return type_models.SimpleType.INTEGER if column_type == SchemaType.SchemaColumn.SchemaColumnType.FLOAT: return type_models.SimpleType.FLOAT if column_type == SchemaType.SchemaColumn.SchemaColumnType.STRING: return type_models.SimpleType.STRING if column_type == SchemaType.SchemaColumn.SchemaColumnType.DATETIME: return type_models.SimpleType.DATETIME if column_type == SchemaType.SchemaColumn.SchemaColumnType.DURATION: return type_models.SimpleType.DURATION if column_type == SchemaType.SchemaColumn.SchemaColumnType.BOOLEAN: return type_models.SimpleType.BOOLEAN else: raise AssertionError(f"Unrecognized SchemaColumnType: {column_type}") class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]): """ 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. """ _SUPPORTED_TYPES: typing.Dict[Type, LiteralType] = { _np.int32: type_models.LiteralType(simple=type_models.SimpleType.INTEGER), _np.int64: type_models.LiteralType(simple=type_models.SimpleType.INTEGER), _np.uint32: type_models.LiteralType(simple=type_models.SimpleType.INTEGER), _np.uint64: type_models.LiteralType(simple=type_models.SimpleType.INTEGER), int: type_models.LiteralType(simple=type_models.SimpleType.INTEGER), _np.float32: type_models.LiteralType(simple=type_models.SimpleType.FLOAT), _np.float64: type_models.LiteralType(simple=type_models.SimpleType.FLOAT), float: type_models.LiteralType(simple=type_models.SimpleType.FLOAT), _np.bool_: type_models.LiteralType(simple=type_models.SimpleType.BOOLEAN), # type: ignore bool: type_models.LiteralType(simple=type_models.SimpleType.BOOLEAN), _np.datetime64: type_models.LiteralType(simple=type_models.SimpleType.DATETIME), _datetime.datetime: type_models.LiteralType(simple=type_models.SimpleType.DATETIME), _np.timedelta64: type_models.LiteralType(simple=type_models.SimpleType.DURATION), _datetime.timedelta: type_models.LiteralType(simple=type_models.SimpleType.DURATION), _np.string_: type_models.LiteralType(simple=type_models.SimpleType.STRING), _np.str_: type_models.LiteralType(simple=type_models.SimpleType.STRING), _np.object_: type_models.LiteralType(simple=type_models.SimpleType.STRING), str: type_models.LiteralType(simple=type_models.SimpleType.STRING), } ENCODERS: Dict[Type, Dict[str, Dict[str, StructuredDatasetEncoder]]] = {} DECODERS: Dict[Type, Dict[str, Dict[str, StructuredDatasetDecoder]]] = {} DEFAULT_PROTOCOLS: Dict[Type, str] = {} DEFAULT_FORMATS: Dict[Type, str] = {} Handlers = Union[StructuredDatasetEncoder, StructuredDatasetDecoder] @staticmethod def _finder(handler_map, df_type: Type, protocol: str, format: str): try: return handler_map[df_type][protocol][format] except KeyError: try: hh = handler_map[df_type][protocol][""] logger.info( f"Didn't find format specific handler {type(handler_map)} for protocol {protocol}" f" format {format}, using default instead." ) return hh except KeyError: ... raise ValueError(f"Failed to find a handler for {df_type}, protocol {protocol}, fmt {format}") @classmethod def get_encoder(cls, df_type: Type, protocol: str, format: str): return cls._finder(StructuredDatasetTransformerEngine.ENCODERS, df_type, protocol, format) @classmethod def get_decoder(cls, df_type: Type, protocol: str, format: str): return cls._finder(StructuredDatasetTransformerEngine.DECODERS, df_type, protocol, format) @classmethod def _handler_finder(cls, h: Handlers) -> Dict[str, Handlers]: # Maybe think about default dict in the future, but is typing as nice? if isinstance(h, StructuredDatasetEncoder): top_level = cls.ENCODERS elif isinstance(h, StructuredDatasetDecoder): top_level = cls.DECODERS else: raise TypeError(f"We don't support this type of handler {h}") if h.python_type not in top_level: top_level[h.python_type] = {} if h.protocol not in top_level[h.python_type]: top_level[h.python_type][h.protocol] = {} return top_level[h.python_type][h.protocol] def __init__(self): super().__init__("StructuredDataset Transformer", StructuredDataset) self._type_assertions_enabled = False # Instances of StructuredDataset opt-in to the ability of being cached. self._hash_overridable = True @classmethod def register(cls, h: Handlers, default_for_type: Optional[bool] = True, override: Optional[bool] = False): """ Call this with any handler to register it with this dataframe meta-transformer The string "://" should not be present in any handler's protocol so we don't check for it. """ lowest_level = cls._handler_finder(h) if h.supported_format in lowest_level and override is False: raise ValueError(f"Already registered a handler for {(h.python_type, h.protocol, h.supported_format)}") lowest_level[h.supported_format] = h logger.debug(f"Registered {h} as handler for {h.python_type}, protocol {h.protocol}, fmt {h.supported_format}") if default_for_type: # TODO: Add logging, think about better ux, maybe default False and warn if doesn't exist. cls.DEFAULT_FORMATS[h.python_type] = h.supported_format cls.DEFAULT_PROTOCOLS[h.python_type] = h.protocol # 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() TypeEngine.register_additional_type(engine, h.python_type, override=True) def assert_type(self, t: Type[StructuredDataset], v: typing.Any): return def to_literal( self, ctx: FlyteContext, python_val: Union[StructuredDataset, typing.Any], python_type: Union[Type[StructuredDataset], Type], expected: LiteralType, ) -> Literal: # Make a copy in case we need to hand off to encoders, since we can't be sure of mutations. # Check first to see if it's even an SD type. For backwards compatibility, we may be getting a FlyteSchema python_type, *attrs = extract_cols_and_format(python_type) # In case it's a FlyteSchema sdt = StructuredDatasetType(format=self.DEFAULT_FORMATS.get(python_type, None)) if expected and expected.structured_dataset_type: sdt = StructuredDatasetType( columns=expected.structured_dataset_type.columns, format=expected.structured_dataset_type.format, external_schema_type=expected.structured_dataset_type.external_schema_type, 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 issubclass(python_type, StructuredDataset) and isinstance(python_val, StructuredDataset): # 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 # # Ex. # def t1(dataset: Annotated[StructuredDataset, my_cols]) -> Annotated[StructuredDataset, my_cols]: # return dataset if python_val._literal_sd is not None: if python_val.dataframe is not None: raise ValueError( f"Shouldn't have specified both literal {python_val._literal_sd} and dataframe {python_val.dataframe}" ) return Literal(scalar=Scalar(structured_dataset=python_val._literal_sd)) # 2. A task returns a python StructuredDataset with a 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. # # Ex. # def t2(uri: str) -> Annotated[StructuredDataset, my_cols] # return StructuredDataset(uri=uri) if python_val.dataframe is None: if not python_val.uri: raise ValueError(f"If dataframe is not specified, then the uri should be specified. {python_val}") sd_model = literals.StructuredDataset( uri=python_val.uri, metadata=StructuredDatasetMetadata(structured_dataset_type=sdt), ) return Literal(scalar=Scalar(structured_dataset=sd_model)) # 3. This is the third and probably most common case. The python StructuredDataset 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) protocol = self._protocol_from_type_or_prefix(ctx, df_type, python_val.uri) return self.encode( ctx, python_val, df_type, protocol, sdt.format or typing.cast(StructuredDataset, python_val).DEFAULT_FILE_FORMAT, sdt, ) # Otherwise assume it's a dataframe instance. Wrap it with some defaults if python_type in self.DEFAULT_FORMATS: fmt = self.DEFAULT_FORMATS[python_type] else: logger.debug(f"No default format for type {python_type}, using system default.") fmt = StructuredDataset.DEFAULT_FILE_FORMAT protocol = self._protocol_from_type_or_prefix(ctx, python_type) meta = StructuredDatasetMetadata(structured_dataset_type=expected.structured_dataset_type if expected else None) sd = StructuredDataset(dataframe=python_val, metadata=meta) return self.encode(ctx, sd, python_type, protocol, fmt, sdt) def _protocol_from_type_or_prefix(self, ctx: FlyteContext, df_type: Type, uri: Optional[str] = None) -> str: """ Get the protocol from the default, if missing, then look it up from the uri if provided, if not then look up from the provided context's file access. """ if df_type in self.DEFAULT_PROTOCOLS: return self.DEFAULT_PROTOCOLS[df_type] else: protocol = protocol_prefix(uri or ctx.file_access.raw_output_prefix) logger.debug( f"No default protocol for type {df_type} found, using {protocol} from output prefix {ctx.file_access.raw_output_prefix}" ) return protocol def encode( self, ctx: FlyteContext, sd: StructuredDataset, df_type: Type, protocol: str, format: str, structured_literal_type: StructuredDatasetType, ) -> Literal: handler: StructuredDatasetEncoder handler = self.get_encoder(df_type, protocol, format) sd_model = handler.encode(ctx, sd, 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. if sd_model.metadata is None: sd_model._metadata = StructuredDatasetMetadata(structured_literal_type) if sd_model.metadata.structured_dataset_type is None: sd_model.metadata._structured_dataset_type = structured_literal_type # Always set the format here to the format of the handler. # Note that this will always be the same as the incoming format except for when the fallback handler # with a format of "" is used. sd_model.metadata._structured_dataset_type.format = handler.supported_format return Literal(scalar=Scalar(structured_dataset=sd_model)) def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T: """ 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, def t1() -> Annotated[StructuredDataset, kwtypes(col_a=int, col_b=float)]: ... def t2(in_a: Annotated[StructuredDataset, kwtypes(col_b=float)]): ... where t2(in_a=t1()), when t2 does in_a.open(pd.DataFrame).all(), it should get a DataFrame with only one column. +-----------------------------+-----------------------------------------+--------------------------------------+ | | StructuredDatasetType of the incoming Literal | +-----------------------------+-----------------------------------------+--------------------------------------+ | StructuredDatasetType | Has columns defined | [] columns or None | | of currently running task | | | +=============================+=========================================+======================================+ | Has columns | The StructuredDatasetType passed to the decoder will have the columns | | defined | as defined by the type annotation of the currently running task. | | | | | | Decoders **should** then subset the incoming data to the columns requested. | | | | +-----------------------------+-----------------------------------------+--------------------------------------+ | [] columns or None | StructuredDatasetType passed to decoder | StructuredDatasetType passed to the | | | will have the columns from the incoming | decoder will have an empty list of | | | Literal. This is the scenario where | columns. | | | the Literal returned by the running | | | | task will have more information than | | | | the running task's signature. | | +-----------------------------+-----------------------------------------+--------------------------------------+ """ # 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) # The literal that we get in might be an old FlyteSchema. # We'll continue to support this for the time being. There is some duplicated logic here but let's # keep it copy/pasted for clarity if lv.scalar.schema is not None: schema_columns = lv.scalar.schema.type.columns # See the repeated logic below for comments if column_dict is None or len(column_dict) == 0: final_dataset_columns = [] if schema_columns is not None and schema_columns != []: for c in schema_columns: final_dataset_columns.append( StructuredDatasetType.DatasetColumn( name=c.name, literal_type=LiteralType( simple=convert_schema_type_to_structured_dataset_type(c.type), ), ) ) # Dataframe will always be serialized to parquet file by FlyteSchema transformer new_sdt = StructuredDatasetType(columns=final_dataset_columns, format=PARQUET) else: final_dataset_columns = self._convert_ordered_dict_of_columns_to_list(column_dict) # Dataframe will always be serialized to parquet file by FlyteSchema transformer new_sdt = StructuredDatasetType(columns=final_dataset_columns, format=PARQUET) metad = literals.StructuredDatasetMetadata(structured_dataset_type=new_sdt) sd_literal = literals.StructuredDataset( uri=lv.scalar.schema.uri, metadata=metad, ) if issubclass(expected_python_type, StructuredDataset): sd = StructuredDataset(dataframe=None, metadata=metad) sd._literal_sd = sd_literal return sd else: return self.open_as(ctx, sd_literal, expected_python_type, metad) # Start handling for StructuredDataset 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 final_dataset_columns = [] # If the current running task's input does not have columns defined, or has an empty list of columns if column_dict is None or len(column_dict) == 0: # but if it does, then we just copy it over if incoming_columns is not None and incoming_columns != []: final_dataset_columns = incoming_columns.copy() # If the current running task's input does have columns defined else: final_dataset_columns = self._convert_ordered_dict_of_columns_to_list(column_dict) new_sdt = StructuredDatasetType( columns=final_dataset_columns, format=lv.scalar.structured_dataset.metadata.structured_dataset_type.format, external_schema_type=lv.scalar.structured_dataset.metadata.structured_dataset_type.external_schema_type, external_schema_bytes=lv.scalar.structured_dataset.metadata.structured_dataset_type.external_schema_bytes, ) metad = 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 # we should do the opening/downloading and whatever else it might entail right now. No iteration option here. return self.open_as(ctx, lv.scalar.structured_dataset, df_type=expected_python_type, updated_metadata=metad) def to_html(self, ctx: FlyteContext, 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 else: # Here we only render column information by default instead of opening the structured dataset. col = typing.cast(StructuredDataset, python_val).columns() df = pd.DataFrame(col, ["column type"]) return df.to_html() else: df = python_val if isinstance(df, pandas.DataFrame): return df.describe().to_html() elif isinstance(df, pa.Table): return df.to_string() elif isinstance(df, _np.ndarray): return pd.DataFrame(df).describe().to_html() elif importlib.util.find_spec("pyspark") is not None and isinstance(df, pyspark.sql.DataFrame): return pd.DataFrame(df.schema, columns=["StructField"]).to_html() elif importlib.util.find_spec("polars") is not None and isinstance(df, pl.DataFrame): describe_df = df.describe() return pd.DataFrame(describe_df.transpose(), columns=describe_df.columns).to_html(index=False) else: raise NotImplementedError("Conversion to html string should be implemented") def open_as( self, ctx: FlyteContext, sd: literals.StructuredDataset, df_type: Type[DF], updated_metadata: StructuredDatasetMetadata, ) -> DF: """ :param ctx: A FlyteContext, useful in accessing the filesystem and other attributes :param sd: :param df_type: :param updated_metadata: New metadata type, since it might be different from the metadata in the literal. :return: dataframe. It could be pandas dataframe or arrow table, etc. """ protocol = protocol_prefix(sd.uri) decoder = self.get_decoder(df_type, protocol, sd.metadata.structured_dataset_type.format) result = decoder.decode(ctx, sd, updated_metadata) if isinstance(result, types.GeneratorType): raise ValueError(f"Decoder {decoder} returned iterator {result} but whole value requested from {sd}") return result def iter_as( self, ctx: FlyteContext, sd: literals.StructuredDataset, df_type: Type[DF], updated_metadata: StructuredDatasetMetadata, ) -> Generator[DF, None, None]: protocol = protocol_prefix(sd.uri) decoder = self.DECODERS[df_type][protocol][sd.metadata.structured_dataset_type.format] result = decoder.decode(ctx, sd, updated_metadata) if not isinstance(result, types.GeneratorType): raise ValueError(f"Decoder {decoder} didn't return iterator {result} but should have from {sd}") return result def _get_dataset_column_literal_type(self, t: Type) -> type_models.LiteralType: if t in self._SUPPORTED_TYPES: return self._SUPPORTED_TYPES[t] if hasattr(t, "__origin__") and t.__origin__ == list: return type_models.LiteralType(collection_type=self._get_dataset_column_literal_type(t.__args__[0])) if hasattr(t, "__origin__") and t.__origin__ == dict: return type_models.LiteralType(map_value_type=self._get_dataset_column_literal_type(t.__args__[1])) raise AssertionError(f"type {t} is currently not supported by StructuredDataset") def _convert_ordered_dict_of_columns_to_list( self, column_map: typing.OrderedDict[str, Type] ) -> typing.List[StructuredDatasetType.DatasetColumn]: converted_cols: typing.List[StructuredDatasetType.DatasetColumn] = [] if column_map is None or len(column_map) == 0: return converted_cols for k, v in column_map.items(): lt = self._get_dataset_column_literal_type(v) converted_cols.append(StructuredDatasetType.DatasetColumn(name=k, literal_type=lt)) return converted_cols def _get_dataset_type(self, t: typing.Union[Type[StructuredDataset], typing.Any]) -> StructuredDatasetType: original_python_type, column_map, storage_format, pa_schema = extract_cols_and_format(t) # Get the column information converted_cols = self._convert_ordered_dict_of_columns_to_list(column_map) # Get the format default_format = ( original_python_type.DEFAULT_FILE_FORMAT if issubclass(original_python_type, StructuredDataset) else self.DEFAULT_FORMATS.get(original_python_type, PARQUET) ) fmt = storage_format or default_format return StructuredDatasetType( columns=converted_cols, format=fmt, external_schema_type="arrow" if pa_schema else None, 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]) -> 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. The other aspects of it - columns, external schema type, etc. can be read from associated metadata. :param t: The python dataframe type, which is mostly ignored. """ return LiteralType(structured_dataset_type=self._get_dataset_type(t)) def guess_python_type(self, literal_type: LiteralType) -> Type[T]: # 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.structured_dataset_type is not None: return StructuredDataset raise ValueError(f"StructuredDatasetTransformerEngine cannot reverse {literal_type}") flyte_dataset_transformer = StructuredDatasetTransformerEngine() TypeEngine.register(flyte_dataset_transformer)