dask / dask/dask

Preserving divisions when reading/loading dataframes with structs containing multiple fields

Open
#11,021 1 comment 0 reactions 0 assignees View on GitHub
dataframe io needs attention
Dominant language
Python
Stars
13.9k
Forks
2k
PR merge metrics
No merged PRs in 30d

Description

Hi,

We're facing issue preserving divisions when writing pyarrow structs containing multiple fields, this is better illustrated in the following example:

```python
data = {
'id': ["01", "02", "03", "04", "05", "06", "07", "08", "09", "10"],
'A': [11, 12, 13, 14, 15, 16, 17, 18, 19, 20],

}
df = pd.DataFrame(data)

# Create a Dask DataFrame
ddf = dd.from_pandas(df, npartitions=3) # Set the number of partitions

# Sort the index
ddf = ddf.set_index('id', drop=True, sorted=True)

def add_nested_struct_multiple_fields(rows):

example_dict = {
"foo": "foo",
"bar": "bar",
}

return [example_dict, example_dict]

def transform_multiple_fields(dataframe):

dataframe["nested_struct"] = dataframe.apply(add_nested_struct_multiple_fields, axis=1)

return dataframe

nested_struct_multiple_fields_schema = pa.list_(
pa.struct(
[
pa.field("foo", pa.string()),
pa.field("bar", pa.string()),
],
)
)

def create_meta_dict(nested_struct_schema):
meta_dict = {"id": pd.Series(dtype="object")}
meta_dict["A"] = pd.Series(dtype=pd.ArrowDtype(pa.int32()))
meta_dict["nested_struct"] = pd.Series(dtype=pd.ArrowDtype(nested_struct_schema))
meta_df = pd.DataFrame(meta_dict).set_index("id")

return meta_df

schema_nested_struct_multiple_fields = {
"A":pa.int32(),
"nested_struct":nested_struct_multiple_fields_schema
}

ddf_multiple_fields = ddf.map_partitions(
transform_multiple_fields,
meta=nested_struct_multiple_fields_meta,
)

ddf_multiple_fields.known_divisions
-> True # Divisions are still preserved before writing

dd.to_parquet(ddf_multiple_fields, "multiple_fields", schema=schema_nested_struct_multiple_fields)

dd.read_parquet(path="multiple_fields", index="id", calculate_divisions=True).known_divisions
-> False # Division are not known after reading the saved dataframe again
```

Here, we are writing two new fields ("foo", "bar"), when reading the parquet file again, the divisions are not preserved. Oddly enough this does not seem to occur when only one field is added to the struct. See [this notebook](https://colab.research.google.com/drive/1-aAeaU9yZYT7Qguh9EeNO5Gg8Y7A4nHn?usp=sharing) for more clarity on the issue where we compare both scenarios.

Could you please provide clarity on the issue? Thanks in advance!

**Environment**:

- Dask version: 2024.2.1
- Python version: 3.11
- Operating System: Linux
- Install method (conda, pip, source): pip

Contributor guide

Open the contributing guide

Research direction

Reproduce the supplied example with dd.to_parquet and dd.read_parquet, comparing the one-field and multiple-field struct cases. Trace how calculate_divisions derives divisions after loading, and consider the issue complete when the multiple-field example preserves known divisions after reading the saved dataframe.

Written by the indexing model from the issue text.

Assessment

Tech stack
pandas, python
Domain
data-engineering
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.