fix: Also preserve ordered categoricals for dask backend

[jtilly]

Part 2 of #361.

This also preserves ordered categoricals for the dask backend.

[fjetter]

FYI Dask's implementation around Categorical was historically really bad. I think that improved with the dataframe rewrite but I can't guarantee it. What I'm saying is: if you rely on categoricals in dask, make sure to test this with your code properly!

[fjetter]

Dask has its own alignment logic and it is stricter than the one proposed here. For example, the test below fails with

TypeError: to union ordered Categoricals, all categories must be the same

def test_read_ddf_differing_ordered_categorical_partition_schemas_fall_back_unordered(
    store_factory,
):
    priority_dtype_0 = pd.CategoricalDtype(
        categories=["low", "high"], ordered=True
    )
    priority_dtype_1 = pd.CategoricalDtype(
        categories=["medium", "low"], ordered=True
    )
    df_0 = pd.DataFrame(
        {
            "id": [0, 1],
            "priority": pd.Categorical(
                ["low", "high"], dtype=priority_dtype_0
            ),
        }
    )
    df_1 = pd.DataFrame(
        {
            "id": [2, 3],
            "priority": pd.Categorical(
                ["medium", "low"], dtype=priority_dtype_1
            ),
        }
    )

    store_dataframes_as_dataset(
        dfs=[df_0, df_1], dataset_uuid="dataset_uuid", store=store_factory
    )

    ddf = read_dataset_as_ddf(
        dataset_uuid="dataset_uuid",
        store=store_factory,
        categoricals=["priority"],
        table="table",
    )
    result = ddf.compute(scheduler="sync").sort_values("id").reset_index(drop=True)

    assert result["priority"].tolist() == ["low", "high", "medium", "low"]
    assert isinstance(result["priority"].dtype, pd.CategoricalDtype)
    assert not result["priority"].cat.ordered
    assert set(result["priority"].cat.categories) == {"high", "low", "medium"}
    assert not ddf._meta.dtypes["priority"].ordered

Generally, I would recommend to add a similar range of tests as done in the non-dask path.

The error pops up in the dask code that is merging partitions so this error would show up whenever dask internally performs a repartitioning or when collecting results. To fix this, we'd have to cast the read files already to the appropriate CategoricalDtype. That might require a bit of a larger refactoring

[fjetter]

I'm a bit surprised this wasn't necessary for the non-dask path

[jtilly]

• Eager path: align_categories is called on already-loaded partitions. Each partition arrives via pyarrow's to_pandas(), and we deliberately exclude the categories columns from _reset_dictionary_columns, so the parquet file's dictionary type (with ordered=True) is preserved into pandas. Eager therefore picks up ordered from the partition data itself, and PR fix: Preserve ordered categorials #361 only had to teach align_categories not to throw it away.

• Dask path: _get_dask_meta_for_dataset builds _meta from the dataset schema before any partition is read.