Computing the length of `DaskDataFrame` is slow

[LucaMarconato]

len(points) is slow because it is needs to compute the value of a Delayed object. This is a problem when printing a sdata object because the length needs to be computed: print(sdata) takes 5.5 seconds (40 million points among 19 dask dataframes).

Even when we load the data from .parquet, the DaskDataFrame object doesn't know the length (which is stored in .parquet).
This is because for dask the dataframe could have been filtered, so the length is treated as unknown.

I have opened an issue in dask, I am not sure how proceed dask/dask#9973.

This is the dask computational graph for an object loaded from disk, one idea could be to look at the data associated to read-parquet layer, find the parquet file and read it with pyarrow.parquet.read_table, which is ultrafast. But it is difficult to distinguish between objects whose rows may have been filtered and whose who didn't. The computational graph shown here is just for an object that has been read from disk (I think the schema tried to convert things to categorical).

points.dask
Out[3]: 
HighLevelGraph with 7 layers.
<dask.highlevelgraph.HighLevelGraph object at 0x7fde0a68bfd0>
 0. read-parquet-2f971e340a47781eb685a87a213bf826
 1. getitem-cdfb73fd637e9e0ee5a292183add9458
 2. cat-set_categories-56f20ea96556105e025a60a300d9ea96
 3. assign-182714d359d11968871ea582587d16ec
 4. getitem-97ec6823c9a58fc7a55c932bcef626c3
 5. cat-set_categories-908e5810fd0c7d7bb0d1b4525d366fd7
 6. assign-298db01be23dc397178dae1c8ef8aa3e

[LucaMarconato]

This is instead the computational graph for an object where the rows have been subset:

v[v['target'] == 0].dask
Out[7]: 
HighLevelGraph with 10 layers.
<dask.highlevelgraph.HighLevelGraph object at 0x7fdde943ff70>
 0. read-parquet-2f971e340a47781eb685a87a213bf826
 1. getitem-cdfb73fd637e9e0ee5a292183add9458
 2. cat-set_categories-56f20ea96556105e025a60a300d9ea96
 3. assign-182714d359d11968871ea582587d16ec
 4. getitem-97ec6823c9a58fc7a55c932bcef626c3
 5. cat-set_categories-908e5810fd0c7d7bb0d1b4525d366fd7
 6. assign-298db01be23dc397178dae1c8ef8aa3e
 7. getitem-e9166b7cae06f6a56b73efba4de9d7c8
 8. eq-ad8298a68d04fdd261df01cf994c8107
 9. getitem-a1121884f71ac60f6dd789f94ccd0211

From this I would not be able to distinguish objects whose length may have been modified and objects that are expected having the same length specified in the parquet file.

[LucaMarconato]

This is the code in the PointsModel.validate() that adds the layers 1-6, if I remove it I can be sure that the object has never been touched and I can compute the length with pyarrow.

for c in data.columns:
 # this is not strictly a validation since we are explicitly importing the categories
 # but it is a convenient way to ensure that the categories are known. It also just changes the state of the
 # series, so it is not a big deal.
 if is_categorical_dtype(data[c]):
 if not data[c].cat.known:
 try:
 data[c] = data[c].cat.set_categories(data[c].head(1).cat.categories)
 except ValueError:
 logger.info(f"Column `{c}` contains unknown categories. Consider casting it.")

In fact now the computational graph is

<dask.highlevelgraph.HighLevelGraph object at 0x7f77e2c42380>
 0. read-parquet-2f971e340a47781eb685a87a213bf826

And this code gives the length very fast (from 5.5 seconds to 0.9 seconds with the cosmx 40 million points dataset).

if len(v.dask.layers) == 1:
 name, layer = v.dask.layers.items().__iter__().__next__()
 if 'read-parquet' in name:
 t = layer.creation_info['args']
 assert isinstance(t, tuple)
 assert len(t) == 1
 parquet_file = t[0]
 from pyarrow.parquet import read_table
 table = read_table(parquet_file)
 length = len(table)
 else:
 length = len(v)
else:
 length = len(v)

Feels a bit hacky, but gonna keep it like this for now. I pushed the code in this pr: #141