Add parquet layer to attempt predicate pushdown, update table_scan logic

Author: ayushdg

dask_sql/physical/rel/logical/table_scan.py

@@ -3,10 +3,13 @@
 from functools import reduce
 from typing import TYPE_CHECKING
 
+from dask.utils_test import hlg_layer
+
 from dask_sql.datacontainer import DataContainer
 from dask_sql.physical.rel.base import BaseRelPlugin
 from dask_sql.physical.rel.logical.filter import filter_or_scalar
 from dask_sql.physical.rex import RexConverter
+from dask_sql.physical.utils.filter import attempt_predicate_pushdown
 
 if TYPE_CHECKING:
     import dask_sql
@@ -81,25 +84,45 @@ def _apply_filters(self, table_scan, rel, dc, context):
         conjunctive_dnf_filters = table_scan.getDNFFilters().filtered_exprs
         non_dnf_filters = table_scan.getDNFFilters().io_unfilterable_exprs
         # All filters here are applied in conjunction (&)
-        if non_dnf_filters or conjunctive_dnf_filters:
-            df_condition = reduce(
-                operator.and_,
-                [
-                    RexConverter.convert(rel, rex, dc, context=context)
-                    for rex in non_dnf_filters
-                ],
-            )
-            df = filter_or_scalar(
-                df, df_condition, conjunctive_filters=conjunctive_dnf_filters
-            )
-            if conjunctive_dnf_filters:
+        if conjunctive_dnf_filters:
+            if non_dnf_filters:
                 df_condition = reduce(
                     operator.and_,
                     [
                         RexConverter.convert(rel, rex, dc, context=context)
-                        for rex in all_filters
+                        for rex in non_dnf_filters
                     ],
                 )
-                df = filter_or_scalar(df, df_condition)
+                df = filter_or_scalar(
+                    df, df_condition, conjunctive_filters=conjunctive_dnf_filters
+                )
+            else:
+                df = attempt_predicate_pushdown(
+                    df, conjunctive_filters=conjunctive_dnf_filters
+                )
+
+            df_condition = reduce(
+                operator.and_,
+                [
+                    RexConverter.convert(
+                        rel, rex, DataContainer(df, cc), context=context
+                    )
+                    for rex in all_filters
+                ],
+            )
+            df = filter_or_scalar(df, df_condition)
+        elif all_filters:
+            df_condition = reduce(
+                operator.and_,
+                [
+                    RexConverter.convert(rel, rex, dc, context=context)
+                    for rex in all_filters
+                ],
+            )
+            df = filter_or_scalar(df, df_condition)
+        try:
+            logger.debug(hlg_layer(df.dask, "read-parquet").creation_info)
+        except KeyError:
+            pass
 
         return DataContainer(df, cc)

dask_sql/physical/utils/filter.py

@@ -83,9 +83,13 @@ def attempt_predicate_pushdown(
     try:
         filters = dsk.layers[name]._dnf_filter_expression(dsk)
         if not isinstance(filters, frozenset):
-            # No filters encountered
-            return ddf
-        filters = filters.to_list_tuple()
+            if conjunctive_filters or disjunctive_filters:
+                filters = []
+            else:
+                # No filters encountered
+                return ddf
+        else:
+            filters = filters.to_list_tuple()
     except ValueError:
         # DNF dispatching failed for 1+ layers
         logger.warning(
@@ -97,8 +101,11 @@ def attempt_predicate_pushdown(
     # Expand the filter set with provided conjunctive and disjunctive filters
     filters.extend([f] for f in disjunctive_filters or [])
     # Add conjunctive filters to each disjunctive filter
-    for f in filters:
-        f.extend(conjunctive_filters or [])
+    if filters:
+        for f in filters:
+            f.extend(conjunctive_filters or [])
+    else:
+        filters = [conjunctive_filters or []]
     # Regenerate collection with filtered IO layer
     try:
         return dsk.layers[name]._regenerate_collection(
@@ -275,6 +282,8 @@ def _dnf_filter_expression(self, dsk):
             func = _blockwise_getitem_dnf
         elif op == dd._Frame.fillna:
             func = _blockwise_fillna_dnf
+        elif op == dd.io.read_parquet:
+            return []
         else:
             raise ValueError(f"No DNF expression for {op}")