Custom serializer not being called

[j-carson]

I have defined a custom serializer as follows

from io import BytesIO

import pandas as pd
from prefect import flow, task
from prefect.serializers import Serializer
from typing import Literal

class ParquetSerializer(Serializer):
    type: Literal["parquet"] = "parquet"

    def dumps(self, obj: pd.DataFrame) -> bytes:
        if not isinstance(obj, pd.DataFrame):
            raise ValueError("Hey, this is for dataframes")
            
        fp = BytesIO()
        obj.to_parquet(path=fp, engine="pyarrow")
        fp.seek(0)
        return fp.read()

    def loads(self, blob: bytes) -> pd.DataFrame:
        fp = BytesIO(blob)
        return pd.read_parquet(fp, engine="pyarrow")


@task(
    persist_result=True,
    result_serializer=ParquetSerializer(),
    result_storage_key="test_parquet.parquet",
)
def save_these_results():
    return pd.DataFrame(
        {
            "a": [1, 2, 3],
            "b": [4, 5, 6],
        }
    )


@flow
async def test_parquet_serializer():
    save_these_results()


if __name__ == "__main__":
    test_parquet_serializer.serve(name="Serialize-Test")

However, it is failing with a stack trace that seems to indicate that it is trying to use a json encoder instead of the serializer I wrote

  File ".../lib/python3.10/site-packages/prefect/futures.py", line 237, in _result
    return await final_state.result(raise_on_failure=raise_on_failure, fetch=True)
  File ".../lib/python3.10/site-packages/prefect/states.py", line 91, in _get_state_result
    raise await get_state_exception(state)
  File ".../lib/python3.10/site-packages/prefect/task_runners.py", line 231, in submit
    result = await call()
  File ".../lib/python3.10/site-packages/prefect/engine.py", line 1762, in begin_task_run
    state = await orchestrate_task_run(
  File ".../lib/python3.10/site-packages/prefect/engine.py", line 2064, in orchestrate_task_run
    terminal_state = await return_value_to_state(
  File ".../lib/python3.10/site-packages/prefect/states.py", line 305, in return_value_to_state
    return Completed(data=await result_factory.create_result(data))
  File ".../lib/python3.10/site-packages/prefect/results.py", line 377, in create_result
    return await PersistedResult.create(
  File ".../lib/python3.10/site-packages/prefect/results.py", line 574, in create
    await storage_block.write_path(key, content=blob.to_bytes())
  File ".../lib/python3.10/site-packages/prefect/results.py", line 615, in to_bytes
    return self.json().encode()
  File "pydantic/main.py", line 504, in pydantic.main.BaseModel.json
  File "/usr/lib/python3.10/json/__init__.py", line 238, in dumps
    **kw).encode(obj)
  File "/usr/lib/python3.10/json/encoder.py", line 199, in encode
    chunks = self.iterencode(o, _one_shot=True)
  File "/usr/lib/python3.10/json/encoder.py", line 257, in iterencode
    return _iterencode(o, 0)
  File "pydantic/json.py", line 88, in pydantic.json.pydantic_encoder
  File "pydantic/json.py", line 45, in pydantic.json.lambda
UnicodeDecodeError: 'utf-8' codec can't decode byte 0xe4 in position 115: invalid continuation byte
15:04:21.891 | ERROR   | Flow run 'slick-yak' - Finished in state Failed("Flow run encountered an exception. UnicodeDecodeError: 'utf-8' codec can't decode byte 0xe4 in position 115: invalid continuation byte")

[tehunter]

I ran into something similar. It turns out it was using the custom serializer to get the parquet byte string - then it tries to write it to a JSON along with some task information (think an object like '{"task_id": "123", "result": "<parquet byte blob>"}'). The byte string was causing some problems, it had to be encoded in base64.

class ParquetSerializer(Serializer):
    type: Literal["parquet"] = "parquet"

    def dumps(self, obj: Any) -> bytes:
        import base64
        assert isinstance(obj, pd.DataFrame)
        buffer = io.BytesIO()
        obj.to_parquet(buffer)
        return base64.encodebytes(buffer.getvalue())

    def loads(self, blob: bytes) -> Any:
        import base64
        buffer = io.BytesIO(base64.decodebytes(blob))
        return pd.read_parquet(buffer)