fix: add support for simple casts

Author: rustyconover

README.md

@@ -31,18 +31,19 @@ rye add query-farm-sql-scan-planning
 
 ```python
 from query_farm_sql_scan_planning import Planner, RangeFieldInfo, SetFieldInfo
+import pyarrow as pa
 
 # Define file metadata
 files = [
     (
         "data_2023_q1.parquet",
         {
-            "sales_amount": RangeFieldInfo[int](
-                min_value=100, max_value=50000,
+            "sales_amount": RangeFieldInfo(
+                min_value=pa.scalar(100), max_value=pa.scalar(50000),
                 has_nulls=False, has_non_nulls=True
             ),
             "region": SetFieldInfo[str](
-                values={"US", "CA", "MX"},
+                values={pa.scalar("US"), pa.scalar("CA"), pa.scalar("MX")},
                 has_nulls=False, has_non_nulls=True
             ),
         }
@@ -51,11 +52,11 @@ files = [
         "data_2023_q2.parquet",
         {
             "sales_amount": RangeFieldInfo[int](
-                min_value=200, max_value=75000,
+                min_value=pa.scalar(200), max_value=pa.scalar(75000),
                 has_nulls=False, has_non_nulls=True
             ),
             "region": SetFieldInfo[str](
-                values={"US", "EU", "UK"},
+                values={pa.scalar("US"), pa.scalar("EU"), pa.scalar("UK")},
                 has_nulls=False, has_non_nulls=True
             ),
         }
@@ -81,9 +82,9 @@ print(matching_files)  # {'data_2023_q2.parquet'}
 For fields with known minimum and maximum values:
 
 ```python
-RangeFieldInfo[int](
-    min_value=0,
-    max_value=100,
+RangeFieldInfo(
+    min_value=pa.scalar(0),
+    max_value=pa.scalar(100),
     has_nulls=False,      # Whether the field contains NULL values
     has_non_nulls=True    # Whether the field contains non-NULL values
 )
@@ -94,8 +95,8 @@ RangeFieldInfo[int](
 For fields with a known set of possible values (useful for categorical data):
 
 ```python
-SetFieldInfo[str](
-    values={"apple", "banana", "cherry"},
+SetFieldInfo(
+    values={pa.scalar("apple"), pa.scalar("banana"), pa.scalar("cherry")},
     has_nulls=False,
     has_non_nulls=True
 )

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "query-farm-sql-scan-planning"
-version = "0.1.4"
+version = "0.1.5"
 description = "A Python library for intelligent file filtering using SQL expressions and metadata-based scan planning. This library enables efficient data lake query optimization by determining which files need to be scanned based on their statistical metadata."
 authors = [
     { name = "Rusty Conover", email = "rusty@conover.me" }

src/query_farm_sql_scan_planning/planner.py

@@ -3,9 +3,10 @@
 from typing import Any, Generator
 import duckdb
 import pyarrow as pa
+import pyarrow.compute as pc
 import sqlglot
 import sqlglot.expressions
-import sqlglot.optimizer.simplify
+import sqlglot.optimizer
 
 
 @dataclass
@@ -139,9 +140,20 @@ def _eval_predicate(
         if isinstance(node, sqlglot.expressions.Is):
             return self._evaluate_node_is(node, file_info)
 
-        # Handle comparison operations
-        if not isinstance(node.left, sqlglot.expressions.Column):
+        # So if the left is just a cast of an column ref we can handle that
+        # because the right hand side will also be that way.
+
+        left_is_column = isinstance(node.left, sqlglot.expressions.Column)
+        left_is_cast = isinstance(node.left, sqlglot.expressions.Cast)
+        need_left_value_cast = False
+        if left_is_cast and isinstance(node.left.this, sqlglot.expressions.Column):
+            # If the left side is a cast of a column, we can treat it as a column reference.
+            need_left_value_cast = True
+            left_column_name = node.left.this.this.this
+        elif not left_is_column:
             return None
+        else:
+            left_column_name = node.left.this.this
 
         if node.right.find(sqlglot.expressions.Column) is not None:
             # Can't evaluate this since it has a right hand column ref, ideally
@@ -165,14 +177,7 @@ def _eval_predicate(
         if type(right_val) is pa.Int32Scalar and right_val.as_py() is None:
             right_val = pa.scalar(None, type=pa.null())
 
-        left_val = node.left
-        assert isinstance(left_val, sqlglot.expressions.Column), (
-            f"Expected a column on left side of {node}, got {left_val}"
-        )
-        assert isinstance(left_val.this, sqlglot.expressions.Identifier), (
-            f"Expected an identifier on left side of {node}, got {left_val.this}"
-        )
-        referenced_field_name = left_val.this.this
+        referenced_field_name = left_column_name
 
         field_info = file_info.get(referenced_field_name)
 
@@ -181,6 +186,22 @@ def _eval_predicate(
         if field_info is None:
             return None
 
+        if need_left_value_cast:
+            if not isinstance(field_info, RangeFieldInfo):
+                # If we need a value cast but the field info is not a range,
+                # we can't evaluate this expression.
+                return None
+            field_info = RangeFieldInfo(
+                has_nulls=field_info.has_nulls,
+                has_non_nulls=field_info.has_non_nulls,
+                min_value=pc.cast(field_info.min_value, right_val.type)
+                if field_info.min_value is not None
+                else None,
+                max_value=pc.cast(field_info.max_value, right_val.type)
+                if field_info.max_value is not None
+                else None,
+            )
+
         if isinstance(field_info, SetFieldInfo):
             match type(node):
                 case sqlglot.expressions.EQ:

src/query_farm_sql_scan_planning/test_planner.py

@@ -165,8 +165,12 @@ def sample_files() -> list[tuple[str, FileFieldInfo]]:
 @pytest.mark.parametrize(
     "clause, expected_files",
     [
+        ("t1::date >= '2030-01-01'", set()),
+        ("t1::date = '2023-08-01'", {"file1"}),
         ("t1 = DATE '2023-08-01'", {"file1"}),
         ("t1 > DATE '2023-08-01'", {"file1"}),
+        ("t1::date > '2023-08-01'", {"file1"}),
+        ("cast(t1 as timestamp) > TIMESTAMP '2023-08-01'", {"file1"}),
         ("t1 <> DATE '2023-08-01'", {"file1"}),
         ("t1 <> DATE '2023-08-01' - interval '6 days'", {"file1"}),
         # This isn't possible, to evaluate, we need to check for additional
@@ -186,6 +190,7 @@ def sample_files() -> list[tuple[str, FileFieldInfo]]:
         ),
         ("'apple' in (d1)", ALL_FILES),  # could be improved.
         ("v1 < 100 and d1 = 'apple'", {"file1"}),
+        ("v1::uhugeint * 5 > 400", ALL_FILES),
         ("v1 > 500 and v1 < 600", {"file4", "file5"}),
         ("v1 != 500 and v1 < 400", {"file1", "file2", "file3"}),
         ("v1 >= 300 and v1 <= 500", {"file2", "file3", "file4", "file7"}),