Computing the length of DaskDataFrame is slow #140

Activity

1. LucaMarconato commented on Feb 16, 2023

   

   LucaMarconato

   on Feb 16, 2023

   Last edited by LucaMarconato

   MemberAuthor

   More actions

   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.
2. LucaMarconato commented on Feb 16, 2023

   

   LucaMarconato

   on Feb 16, 2023

   Last edited by LucaMarconato

   MemberAuthor

   More actions

   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
3. LucaMarconato

   closed this as completedon Mar 15, 2023