Add materialzed layer to the list of supported layers for fast limit …

…codepath
dask-contrib · Oct 5, 2022 · f34536f · f34536f
1 parent 607eccf
commit f34536f
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/dask_sql/physical/rel/logical/limit.py b/dask_sql/physical/rel/logical/limit.py
@@ -3,6 +3,7 @@
 import dask.dataframe as dd
 from dask import config as dask_config
 from dask.blockwise import Blockwise
+from dask.highlevelgraph import MaterializedLayer
 from dask.layers import DataFrameIOLayer
 
 from dask_sql.datacontainer import DataContainer
@@ -59,7 +60,9 @@ def _apply_limit(self, df: dd.DataFrame, limit: int, offset: int) -> dd.DataFram
                 dask_config.get("sql.limit.check-first-partition")
                 and all(
                     [
-                        isinstance(layer, (DataFrameIOLayer, Blockwise))
+                        isinstance(
+                            layer, (DataFrameIOLayer, Blockwise, MaterializedLayer)
+                        )
                         for layer in df.dask.layers.values()
                     ]
                 )