Reindex categorical codes from Polars

[Luca Bittarello (lbittarello)]

When loading Parquet data, I realised that Polars doesn't necessarily start categorical codes at zero, which causes the kernel to crash in, e.g., transpose_matvec_complex. This PR adds reindexing for safety.

[Jan Tilly (jtilly)]

Changelog?

Does this need a test (maybe Martin Stancsics (@stanmart) can contribute one)? Seems like a subtle bug that's easy to re-introduce.

[Martin Stancsics (stanmart)]

AFAICT _extract_codes_and_categories still relies on the fact that the pl.Series.get_categories() returns category labels such that their codes are ordered (increasing). I would expect that to be the case, but then I would have also expected category codes to have no gaps. Do you know of a way to make sure it is true (without comparing physical values and labels, which is probably quite expensive)? A cursory search told me polars considers the codes an implementation detail and do not really want to expose the mapping.

[Martin Stancsics (stanmart)]

Here's a mwe to reproduce the non-consecutive codes without loading/writing data:

import polars as pl

with pl.StringCache():
    s1 = pl.Series(["beagle", "poodle", "labrador"], dtype=pl.Categorical)
    s2 = pl.Series(["labrador", "boxer", "beagle"], dtype=pl.Categorical)

print("s1 categories: ", s1.cat.get_categories().to_numpy())
print("s1 codes: ", s1.to_physical().to_numpy())

print("s2 categories: ", s2.cat.get_categories().to_numpy())
print("s2 codes: ", s2.to_physical().to_numpy())

Output:

s1 categories:  ['beagle' 'poodle' 'labrador']
s1 codes:  [0 1 2]
s2 categories:  ['labrador' 'boxer' 'beagle']
s2 codes:  [2 3 0]

The good news is that get_categories() indeed seems to return the categories such that their codes are increasing, even when the observation with the lowest code comes last.

It seems that get_categories returns labels in order of their apparence, and not their category codes. I'll turn it into a test case.

[Martin Stancsics (stanmart)]

It will be an issue, though:

_extract_codes_and_categories(s2)

(array([1, 2, 0]), array(['labrador', 'boxer', 'beagle'], dtype=object))

If we reconstruct the categorical from this, it will be ['boxer', 'beagle', 'labrador'], which is incorrect.

[Martin Stancsics (stanmart)]

Need to handle polars.Series.get_categories() returning category labels such that their codes are not necessarily in an increasing order.

[Luca Bittarello (lbittarello)]

I implemented a different fix, which gets around the string cache by localising the series. It allows us to reconstitute both arrays in Martin's example. (My hunch is that Polars is using some sort of string caching when reading data from disk, which is why I thought that there might be gaps in the first place, but I did naively expect the categories to be in the right order. Thanks for catching that, Martin Stancsics (@stanmart).)

Changelog?

Polars support hasn't made it into any release yet. The existing log entry (under new features) should be enough.