fix: change to a generator

Author: rustyconover

README.md

@@ -66,11 +66,11 @@ files = [
 planner = Planner(files)
 
 # Filter files based on SQL expressions
-matching_files = planner.get_matching_files("sales_amount > 40000 AND region = 'US'")
+matching_files = set(planner.get_matching_files("sales_amount > 40000 AND region = 'US'"))
 print(matching_files)  # {'data_2023_q1.parquet', 'data_2023_q2.parquet'}
 
 # More complex queries
-matching_files = planner.get_matching_files("region IN ('EU', 'UK')")
+matching_files = set(planner.get_matching_files("region IN ('EU', 'UK')"))
 print(matching_files)  # {'data_2023_q2.parquet'}
 ```
 

src/query_farm_sql_scan_planning/planner.py

@@ -1,6 +1,6 @@
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Generator
 import duckdb
 import pyarrow as pa
 import sqlglot
@@ -476,7 +476,9 @@ def _evaluate_sql_node(
                     f"Supported types: Connector, Predicate, Not, Boolean, Case, Null"
                 )
 
-    def get_matching_files(self, exp: sqlglot.expressions.Expression | str) -> set[str]:
+    def get_matching_files(
+        self, exp: sqlglot.expressions.Expression | str, *, dialect: str = "duckdb"
+    ) -> Generator[str, None, None]:
         """
         Get a set of files that match the given SQL expression.
         Args:
@@ -487,24 +489,19 @@ def get_matching_files(self, exp: sqlglot.expressions.Expression | str) -> set[s
         """
         if isinstance(exp, str):
             # Parse the expression if it is a string.
-            expression = sqlglot.parse_one(exp, dialect="duckdb")
+            expression = sqlglot.parse_one(exp, dialect=dialect)
         else:
             expression = exp
 
-        assert isinstance(expression, sqlglot.expressions.Expression), (
-            f"Expected a sqlglot expression, got {type(expression)}"
-        )
+        if not isinstance(expression, sqlglot.expressions.Expression):
+            raise ValueError(f"Expected a sqlglot expression, got {type(expression)}")
 
         # Simplify the parsed expression, move all of the literals to the right side
         expression = sqlglot.optimizer.simplify.simplify(expression)
 
-        matching_files = set()
-
         for filename, file_info in self.files:
             eval_result = self._evaluate_sql_node(expression, file_info)
             if eval_result is None or eval_result is True:
                 # If the expression evaluates to True or cannot be evaluated, add the file
                 # to the result set since the caller will be able to filter the rows further.
-                matching_files.add(filename)
-
-        return matching_files
+                yield filename

src/query_farm_sql_scan_planning/test_planner.py

@@ -305,7 +305,9 @@ def test_scan_planning(
     filter_obj = Planner(sample_files)
 
     # Apply the filter
-    result = filter_obj.get_matching_files(sqlglot.parse_one(clause, dialect="duckdb"))
+    result = set(
+        filter_obj.get_matching_files(sqlglot.parse_one(clause, dialect="duckdb"))
+    )
 
     # Check if files were filtered as expected
     if result != expected_files: