ml6team / fondant / 6010460310

91%

push

github

Preserve divisions when writing and reading (#391)

Apparently we are currently not preserving the divisions of the data
when writing and reading to parquet. We are preserving the partitions
when reading initially, but since divisions are no longer known, the
partitions will be shuffled when we merge the index and subsets
together.

Let's create a dataframe with known partitions and write it to parquet:
```python
ddf = dd.from_dict({"data": list(range(5))}, npartitions=5)
ddf.known_divisions
Out: True
ddf.divisions
Out: (0, 1, 2, 3, 4, 4)

ddf.to_parquet("/tmp/data")
```

If we read it back, we see that the partitions are preserved, but the
division information is not.
```python
index_df = dd.read_parquet("/tmp/data")
index_df = index_df.drop("data", axis=1)

index_df.map_partitions(len).compute()
Out: 
0    1
1    1
2    1
3    1
4    1
dtype: int64

index_df.known_divisions
Out: False
index_df.divisions
Out: (None, None, None, None, None, None)
```

We read it back as a subset as well and double check that we get the
same behavior, we do:
```python
subset_df = dd.read_parquet("/tmp/data")

subset_df.map_partitions(len).compute()
Out: 
0    1
1    1
2    1
3    1
4    1
dtype: int64

subset_df.known_divisions
Out: False
subset_df.divisions
Out: (None, None, None, None, None, None)
```

If we now merge these 2 dataframes together, we see that the partitions
are shuffled, because Dask thinks the divisions are unknown.
```python
dataframe = dd.merge(index_df, subset_df, left_index=True, right_index=True, how="left")
dataframe.map_partitions(len).compute()
Out: 
0    3
1    1
2    1
3    0
4    0
dtype: int64
```

Running this again with the changes in this PR will preserve the
division information and the partitions after merging.

3 of 3 new or added lines in 1 file covered. (100.0%)

1120 of 1206 relevant lines covered (92.87%)

3.71 hits per line