Avoid clearing cached_plan for every dataset (#286)

dask · Sep 12, 2023 · 753ffbf · 753ffbf
1 parent cd1dfaf
commit 753ffbf
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/dask_expr/io/parquet.py b/dask_expr/io/parquet.py
@@ -52,7 +52,7 @@
 
 _cached_dataset_info = {}
 _CACHED_DATASET_SIZE = 10
-# TODO: Allow _cached_plan to contain >1 item?
+_CACHED_PLAN_SIZE = 10
 _cached_plan = {}
 
 
@@ -65,6 +65,12 @@ def _control_cached_dataset_info(key):
         _cached_dataset_info.pop(key_to_pop)
 
 
+def _control_cached_plan(key):
+    if len(_cached_plan) > _CACHED_PLAN_SIZE and key not in _cached_plan:
+        key_to_pop = list(_cached_plan.keys())[0]
+        _cached_plan.pop(key_to_pop)
+
+
 @normalize_token.register(pa_ds.Dataset)
 def normalize_pa_ds(ds):
     return (ds.files, ds.schema)
@@ -602,7 +608,7 @@ def _plan(self):
                     common_kwargs,
                 )
 
-            _cached_plan.clear()
+            _control_cached_plan(dataset_token)
             _cached_plan[dataset_token] = {
                 "func": io_func,
                 "parts": parts,