Pickle error when trying to `append` to existing deltatable

[Thodorissio]

First of all I would like to thank you for you awesome contributions. During my development I came across the following issue.

Description

When trying to append to existing DeltaTable the following error occurs:

TypeError: ('Could not serialize object of type HighLevelGraph', '<ToPickle: HighLevelGraph with 3 layers.\n<dask.highlevelgraph.HighLevelGraph object at 0x1384dcc0ad0>\n 0. 1341334790528\n 1. finalize-02082eb4-e53c-4b1a-83dc-fb753d3f60dc\n 2. _commit-94f97f6c-675a-47c6-88a7-82b1a5234034\n>')

Reproducible Example

import pandas as pd
import dask.dataframe as dd
import dask_deltatable as ddt

from distributed import Client
from deltalake import DeltaTable

output_table = "./animals"


if __name__ == "__main__":
    client = Client()
    print(f"Dask Client: {client}")

    animals_df = pd.DataFrame(
        {
            "name": ["dog", "cat", "whale", "elephant"],
            "life_span": [13, 15, 90, 70],
        },
    )

    animals_ddf = dd.from_pandas(animals_df)
    animals_ddf["high_longevity"] = animals_ddf["life_span"] > 40
    ddt.to_deltalake(
        table_or_uri=output_table,
        df=animals_ddf,
        compute=True,
        mode="append",
    )

    delta_table = DeltaTable(output_table)
    delta_table_df = delta_table.to_pandas()
    print("Created DeltaTable:")
    print(delta_table_df)

    more_animals_df = pd.DataFrame(
        {
            "name": ["shark", "parrot"],
            "life_span": [20, 50],
        },
    )

    more_animals_ddf = dd.from_pandas(more_animals_df)
    more_animals_ddf["high_longevity"] = more_animals_ddf["life_span"] > 40
    ddt.to_deltalake(
        table_or_uri=output_table,
        df=more_animals_ddf,
        compute=True,
        mode="append",
    )

Stacktrace

Dask Client: <Client: 'tcp://127.0.0.1:60355' processes=4 threads=12, memory=31.90 GiB>
Created DeltaTable:
       name  life_span  high_longevity
0       dog         13           False
1       cat         15           False
2     whale         90            True
3  elephant         70            True
2024-12-16 13:58:42,783 - distributed.protocol.pickle - ERROR - Failed to serialize <ToPickle: HighLevelGraph with 3 layers.
<dask.highlevelgraph.HighLevelGraph object at 0x1e7f9b8b530>
 0. 2095838501424
 1. finalize-54be334c-3207-4a95-8908-1aac80f5edb6
 2. _commit-2c7fae99-a722-4a11-8b99-ca2120ebbb4d
>.
Traceback (most recent call last):
  File "C:\Users\thodo\miniconda3\envs\myenv\Lib\site-packages\distributed\protocol\pickle.py", line 60, in dumps
    result = pickle.dumps(x, **dump_kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: cannot pickle 'deltalake._internal.RawDeltaTable' object

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "C:\Users\thodo\miniconda3\envs\myenv\Lib\site-packages\distributed\protocol\pickle.py", line 65, in dumps
    pickler.dump(x)
TypeError: cannot pickle 'deltalake._internal.RawDeltaTable' object

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "C:\Users\thodo\miniconda3\envs\myenv\Lib\site-packages\distributed\protocol\pickle.py", line 77, in dumps
    result = cloudpickle.dumps(x, **dump_kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\thodo\miniconda3\envs\myenv\Lib\site-packages\cloudpickle\cloudpickle.py", line 1529, in dumps
    cp.dump(obj)
  File "C:\Users\thodo\miniconda3\envs\myenv\Lib\site-packages\cloudpickle\cloudpickle.py", line 1295, in dump
    return super().dump(obj)
           ^^^^^^^^^^^^^^^^^
TypeError: cannot pickle 'deltalake._internal.RawDeltaTable' object
Traceback (most recent call last):
  File "C:\Users\thodo\miniconda3\envs\myenv\Lib\site-packages\distributed\protocol\pickle.py", line 60, in dumps
    result = pickle.dumps(x, **dump_kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: cannot pickle 'deltalake._internal.RawDeltaTable' object

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "C:\Users\thodo\miniconda3\envs\myenv\Lib\site-packages\distributed\protocol\pickle.py", line 65, in dumps
    pickler.dump(x)
TypeError: cannot pickle 'deltalake._internal.RawDeltaTable' object

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "C:\Users\thodo\miniconda3\envs\myenv\Lib\site-packages\distributed\protocol\serialize.py", line 366, in serialize
    header, frames = dumps(x, context=context) if wants_context else dumps(x)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\thodo\miniconda3\envs\myenv\Lib\site-packages\distributed\protocol\serialize.py", line 78, in pickle_dumps
    frames[0] = pickle.dumps(
                ^^^^^^^^^^^^^
  File "C:\Users\thodo\miniconda3\envs\myenv\Lib\site-packages\distributed\protocol\pickle.py", line 77, in dumps
    result = cloudpickle.dumps(x, **dump_kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\thodo\miniconda3\envs\myenv\Lib\site-packages\cloudpickle\cloudpickle.py", line 1529, in dumps
    cp.dump(obj)
  File "C:\Users\thodo\miniconda3\envs\myenv\Lib\site-packages\cloudpickle\cloudpickle.py", line 1295, in dump
    return super().dump(obj)
           ^^^^^^^^^^^^^^^^^
TypeError: cannot pickle 'deltalake._internal.RawDeltaTable' object

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "C:\Users\thodo\Documents\libra\myenv\append_issue.py", line 45, in <module>
    ddt.to_deltalake(
  File "C:\Users\thodo\miniconda3\envs\myenv\Lib\site-packages\dask_deltatable\write.py", line 239, in to_deltalake
    result = result.compute()
             ^^^^^^^^^^^^^^^^
  File "C:\Users\thodo\miniconda3\envs\myenv\Lib\site-packages\dask\base.py", line 372, in compute
    (result,) = compute(self, traverse=False, **kwargs)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\thodo\miniconda3\envs\myenv\Lib\site-packages\dask\base.py", line 660, in compute
    results = schedule(dsk, keys, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\thodo\miniconda3\envs\myenv\Lib\site-packages\distributed\protocol\serialize.py", line 392, in serialize
    raise TypeError(msg, str_x) from exc
TypeError: ('Could not serialize object of type HighLevelGraph', '<ToPickle: HighLevelGraph with 3 layers.\n<dask.highlevelgraph.HighLevelGraph object at 0x1e7f9b8b530>\n 0. 2095838501424\n 1. finalize-54be334c-3207-4a95-8908-1aac80f5edb6\n 2. _commit-2c7fae99-a722-4a11-8b99-ca2120ebbb4d\n>')

Library Versions

dask==2024.11.2
dask-deltatable==0.3.3
deltalake==0.22.3
distributed==2024.11.2
pandas==2.2.3

[roflmaoqwertyqaz]

I ran into this issue as well, it seems the issue is with the deltalake library. In the DeltaTable class, the internal _table object is a RawDeltaTable, which is not picklable.

I worked around this issue by inheriting the DeltaTable class in deltalake and adding __getstate__ and __setstate__ methods to ensure the problematic object is not pickled:

import os
from pathlib import Path
from typing import Union, Dict, Optional

import deltalake
from deltalake._internal import RawDeltaTable
import cloudpickle

class PicklableDeltaTable(deltalake.DeltaTable):

    def __init__(
        self,
        table_uri: Union[str, Path, "os.PathLike[str]"],
        version: Optional[int] = None,
        storage_options: Optional[Dict[str, str]] = None,
        without_files: bool = False,
        log_buffer_size: Optional[int] = None,
    ):
        super().__init__(table_uri, version, storage_options, without_files, log_buffer_size)

        # Create class attributes for the constructor parameters that do not already have attributes
        self._table_uri = table_uri
        self._version = version
        self._without_files = without_files
        self._log_buffer_size = log_buffer_size

    def __getstate__(self):
        state = self.__dict__.copy()

        # drop the _table attribute, which is an unpicklable RawDeltaTable
        del state['_table']
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)

        # Recreate the RawDeltaTable from the saved constructor parameters and add it back as an attribute
        table = RawDeltaTable(
            str(self._table_uri),
            version=self._version,
            storage_options=self._storage_options,
            without_files=self._without_files,
            log_buffer_size=self._log_buffer_size,
        )
        self._table = table

This can be used as a drop in replacement for the input delta table object in to_deltalake, the only extra step is to explicitly pass in the PicklableDeltaTable, and to handle the case where the delta table does not yet exist:

out_path = 'foo/bar/baz'

# check if table exists, if it doesn't, use literal path instead of table object
if deltalake.DeltaTable.is_deltatable(out_path):
    out_table = PicklableDeltaTable(
        out_path,
        storage_options=storage_options
    )
else:
    out_table = out_path

ddt.to_deltalake(
    table_or_uri=out_table,
    df=your_dask_df,
    compute=True,
    mode="append",
)

This might be bit of a hack since it is manipulating DeltaTable internals, but I found it to be the most straightforward way to deal with this issue.

Hope this helps!

[Thodorissio]

Thanks a lot for this solution, it works nicely. I will use this until this case is properly handled by the library classes.

[jacobtomlinson]

This is great! @roflmaoqwertyqaz do you have any interest in making a PR to the upstream deltalake library to implement this properly there?

[HiromuHota]

Submitted a PR against deltalake delta-io/delta-rs#3501, using __reduce__ instead of __getstate__/__setstate__ because the former is good enough. Hope they take it.

[ion-elgreco]

I need some more context here, why do you serialize the whole DeltaTable object, instead of passing the configuration to the workers? If we would merge your pr @HiromuHota then one, could dump it's storage options in pickle format which isn't secure, and that feels like a big footgun for unaware users in terms of security.

[jacobtomlinson]

@ion-elgreco Dask passes objects between workers using cloudpickle. Therefore all objects needs to be pickle serializable. Being able to construct a DeltaTable on one machine and then share it between workers for them to all pull data in parallel is the desired functionality here.

one, could dump it's storage options in pickle format which isn't secure, and that feels like a big footgun for unaware users in terms of security

This feels a little like an overly cautious security practice and has broken functionality here, have you had this concern raised by users or external security researchers? If you enter credentials into an object and then dump that object with pickle it's not surprising that the credentials would be included in the dump.

[ion-elgreco]

@ion-elgreco Dask passes objects between workers using cloudpickle. Therefore all objects needs to be pickle serializable. Being able to construct a DeltaTable on one machine and then share it between workers for them to all pull data in parallel is the desired functionality here.

one, could dump it's storage options in pickle format which isn't secure, and that feels like a big footgun for unaware users in terms of security

This feels a little like an overly cautious security practice and has broken functionality here, have you had this concern raised by users or external security researchers? If you enter credentials into an object and then dump that object with pickle it's not surprising that the credentials would be included in the dump.

It's a concern I'm raising. I don't see why we need to export storage options as well, if you set the proper env vars in the workers it also works

[jacobtomlinson]

If the user is passing storage options via the Python API and not via environment variables then things will break for them. I agree that setting the storage options in the environment would be the security best practice here, but enforcing that as the only method seems heavy handed.

I wonder if a compromise would be to implement a flag where if DeltaTable(storage_options=None) then __reduce__ also sets storage_options to None. But if the user does pass storage options into the class constructor then it should include them in the dump? This way we avoid leaking credentials unintentionally, but still include them if they are passed explicitly. See delta-io/delta-rs#3501 (review)

[ion-elgreco]

If the user is passing storage options via the Python API and not via environment variables then things will break for them. I agree that setting the storage options in the environment would be the security best practice here, but enforcing that as the only method seems heavy handed.

I wonder if a compromise would be to implement a flag where if DeltaTable(storage_options=None) then __reduce__ also sets storage_options to None. But if the user does pass storage options into the class constructor then it should include them in the dump? This way we avoid leaking credentials unintentionally, but still include them if they are passed explicitly. See delta-io/delta-rs#3501 (review)

I don't want to bloat the api with a flag on DeltaTable to allow pickling storage_options

[jacobtomlinson]

@ion-elgreco I'm not suggesting exposing it via the public API. I'm suggesting keeping track privately and only pickling the storage options if they are passed as a kwarg. This should satisfy your security concerns while preserving expected functionality when passing credentials to the object via a kwarg and then pickling it.

[ion-elgreco]

The parsed storage options from env vars are never exposed into python, so that doesn't change anything

[HiromuHota]

@ion-elgreco / @jacobtomlinson Thanks for your response&discussion!
As a user, can I pass storage options as env vars to delta-rs or would you be open to allow it? If that's the case, we can pickle an DeltaTable object except for storage_options and pick up storage_options from env vars when unpickling.

The alternative you suggested (i.e., passing configuration to the workers) definitely works and this is exactly what I did as a workaround (see snorkel-ai@fdb2b93). However, this is a bit cumbersome because you have to pass up to 5 parameters (table_uri, version, storage_options, without_files, log_buffer_size) to the workers. As a user of delta-rs, I'd like to pass as few parameters as possible, a single parameter (object of DeltaTable) or two (DeltaTable and storage_options). Thoughts?

[jacobtomlinson]

The parsed storage options from env vars are never exposed into python, so that doesn't change anything

In that case I'm not sure I understand your security concern at all. As library developers I don't think we should be concerned about protecting users from pickling and dumping things they shouldn't. There are enough articles on the internet about why most folks shouldn't use pickle for serialisation to disk. This is especially true when the intentional blocking of pickle breaks other projects which use it legitimately for over the network serialisation like Dask does.

As a user, can I pass storage options as env vars to delta-rs

I think this highlights the risk here of creating a poor user experience. If users don't know this is possible they aren't going to know to configure this on all their Dask workers ahead of time.