From 043b4d66855ec7dff742fcd4f0adc64b2422dc4d Mon Sep 17 00:00:00 2001
From: barak <barak.alon@airbnb.com>
Date: Sat, 20 Jul 2024 13:35:41 -0400
Subject: [PATCH] fix: convert range back to lookup (#4)

---
 odex/optimize.py        |  9 ++++++---
 odex/plan.py            | 10 ++++++++++
 tests/fixtures/e2e.yaml | 12 +++++++++++-
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/odex/optimize.py b/odex/optimize.py
index 405c21d..5e2ea98 100644
--- a/odex/optimize.py
+++ b/odex/optimize.py
@@ -131,12 +131,15 @@ def transform(self, plan: Plan, ctx: Context) -> Plan:
                     else:
                         new_range = combined
 
-                inputs.append(
-                    IndexRange(
+                if new_range.is_equality():
+                    assert isinstance(new_range.left, Bound)
+                    new_plan: Plan = IndexLookup(index=index, value=new_range.left.value)
+                else:
+                    new_plan = IndexRange(
                         index=index,
                         range=new_range,
                     )
-                )
+                inputs.append(new_plan)
 
             inputs.extend(others)
 
diff --git a/odex/plan.py b/odex/plan.py
index 3fa2413..4e14278 100644
--- a/odex/plan.py
+++ b/odex/plan.py
@@ -181,6 +181,16 @@ def combine(self, other: "Range[C]") -> "Optional[Range[C]]":
             right=right,
         )
 
+    def is_equality(self) -> bool:
+        """Determine if a range can be treated like an equality"""
+        return (
+            isinstance(self.left, Bound)
+            and isinstance(self.right, Bound)
+            and self.left.inclusive
+            and self.right.inclusive
+            and self.left == self.right
+        )
+
     def _combine_bounds(
         self,
         a: OptionalBound,
diff --git a/tests/fixtures/e2e.yaml b/tests/fixtures/e2e.yaml
index a3a0cd5..6f1664f 100644
--- a/tests/fixtures/e2e.yaml
+++ b/tests/fixtures/e2e.yaml
@@ -129,7 +129,7 @@ setups:
           - ScanFilter: a >= 3
         - ScanFilter: a <= 3
     optimized_plan: |-
-      IndexRange: 3 <= SortedDictIndex(a) <= 3
+      IndexLookup: SortedDictIndex(a) = 3
     result:
     - 2
   - title: Conflicting equalities
@@ -260,6 +260,16 @@ setups:
       IndexLookup: HashIndex(a) = 'foo'
     result:
     - 0
+  - title: Duplicate String equality
+    condition: a = 'foo' AND a = 'foo'
+    plan: |-
+      Intersect
+        - ScanFilter: a = 'foo'
+        - ScanFilter: a = 'foo'
+    optimized_plan: |-
+      IndexLookup: HashIndex(a) = 'foo'
+    result:
+      - 0
   - title: Conflicting equalities, string
     condition: a = 'foo' AND a = 'bar'
     plan: |-