Does Flyte support the use of Polars LazyFrames? I keep getting a warrning regrading it, see below:

[f6198790a618a4396a3a-n0-0] terminated with exit code (247). Reason [OOMKilled]. Message: 

{"asctime": "2024-05-01 20:01:11,141", "name": "flytekit", "levelname": "WARNING", "message": "Unsupported Type <built-in function any> found, Flyte will default to use PickleFile as the transport. Pickle can only be used to send objects between the exact same version of Python, and we strongly recommend to use python type that flyte support."}

thoughts?

looks like it’s not: https://github.com/flyteorg/flytekit/blob/master/plugins/flytekit-polars/flytekitplugins/polars/sd_transformers.py#L87 would be fairly straightforward to register a new

PolarsDataFrameToParquetEncodingHandler

and

ParquetToPolarsDataFrameDecodingHandler

for lazyframes

Im pretty new to flyte, im probably even more "green" when it comes to th interworking's. So you will have bare with me here 🙂 .

if you look at that

sd_transformers

file, you’ll see those classes defined, which operate on

pl.DataFrame

here, here, here, and here.

class PolarsLazyFrameToParquetEncodingHandler(StructuredDatasetEncoder): ...

class ParquetToPolarsLazyFrameDecodingHandler(StructuredDatasetDecoder):

You could define your own classes (or contribute it back into the polars plugin 🙂) and then register it

StructuredDatasetTransformerEngine.register(PolarsLazyFrameToParquetEncodingHandler())
StructuredDatasetTransformerEngine.register(ParquetToPolarsLazyFrameDecodingHandler())

so what something like this:

class PolarsLazyFrameRenderer:
    """
    Renders Polars LazyFrame's statistics as an HTML table after collecting the data.
    """

    def to_html(self, lf: pl.LazyFrame) -> str:
        assert isinstance(lf, pl.LazyFrame)
        df = lf.collect()  # Convert LazyFrame to DataFrame
        describe_df = df.describe()
        return pd.DataFrame(describe_df.transpose(), columns=describe_df.columns).to_html(index=False)


class PolarsLazyFrameToParquetEncodingHandler(StructuredDatasetEncoder):
    def __init__(self):
        super().__init__(pl.LazyFrame, None, PARQUET)

    def encode(
        self,
        ctx: FlyteContext,
        structured_dataset: StructuredDataset,
        structured_dataset_type: StructuredDatasetType,
    ) -> literals.StructuredDataset:
        lf = typing.cast(pl.LazyFrame, structured_dataset.dataframe)
        df = lf.collect()  # Collect data to a DataFrame before writing

        output_bytes = io.BytesIO()
        df.write_parquet(output_bytes)  # Polars now only supports write_parquet

        if structured_dataset.uri is not None:
            fs = ctx.file_access.get_filesystem_for_path(path=structured_dataset.uri)
            with fs.open(structured_dataset.uri, "wb") as s:
                s.write(output_bytes.getvalue())
            output_uri = structured_dataset.uri
        else:
            remote_fn = "00000"  # Default filename
            output_uri = ctx.file_access.put_raw_data(output_bytes.getvalue(), file_name=remote_fn)
        return literals.StructuredDataset(uri=output_uri, metadata=StructuredDatasetMetadata(structured_dataset_type))


class ParquetToPolarsLazyFrameDecodingHandler(StructuredDatasetDecoder):
    def __init__(self):
        super().__init__(pl.LazyFrame, None, PARQUET)

    def decode(
        self,
        ctx: FlyteContext,
        flyte_value: literals.StructuredDataset,
        current_task_metadata: StructuredDatasetMetadata,
    ) -> pl.LazyFrame:
        uri = flyte_value.uri

        kwargs = get_fsspec_storage_options(
            protocol=fsspec_utils.get_protocol(uri),
            data_config=ctx.file_access.data_config,
        )
        # Return LazyFrame directly from read_parquet
        return pl.scan_parquet(uri, storage_options=kwargs)


StructuredDatasetTransformerEngine.register(PolarsLazyFrameToParquetEncodingHandler())
StructuredDatasetTransformerEngine.register(ParquetToPolarsLazyFrameDecodingHandler())
StructuredDatasetTransformerEngine.register_renderer(pl.LazyFrame, PolarsLazyFrameRenderer())

yeah, that looks about right

one thing im confused of, is that StructredDataset type.

as long as you

collect

at serialization time it should be fine

intresting..

yeah you can test this by having these classes and registration call inside the same file you have the tasks/workflows. you could have it in a separate file, as long as you import it in the workflow file.

Cool, let me give it a go

it needs to be in the top-level scope of the module

# lazyframe encoder/decoder class defs
# registration
...


@workflow
def wf(): ...

yeah, so its a bit odd, i get this now:

{
  "asctime": "2024-05-01 20:54:00,858",
  "name": "flytekit",
  "levelname": "WARNING",
  "message": "Unsupported Type <built-in function any> found, Flyte will default to use PickleFile as the transport. Pickle can only be used to send objects between the exact same version of Python, and we strongly recommend to use python type that flyte support."
}

weird, this is the

any

built-in function, not the

typing.Any

type

well, i think i got that cleared

cool stuff

@Brad could you make an issue for this? [flyte-core] and if you have capacity a contribution would be 🔥. Basically the code you have here is a good start and can be added to this file. You can follow this contribution guide if you get stuck.

⭐ Create a new Flyte Core Feature issue: https://github.com/flyteorg/flyte/issues/new?assignees=&labels=enhancement%2Cuntriaged&template=feature_request.yaml&title=%5BCore+feature%5D+

will do!