Refactor code in TableSegment

Author: erezsh

data_diff/table_segment.py

@@ -14,6 +14,12 @@
 
 RECOMMENDED_CHECKSUM_DURATION = 20
 
+def _make_range(col, min_value=None, max_value=None):
+    if min_value is not None:
+        yield min_value <= col
+    if max_value is not None:
+        yield col < max_value
+
 
 @dataclass
 class TableSegment:
@@ -78,28 +84,25 @@ def with_schema(self) -> "TableSegment":
 
         return self._with_raw_schema(self.database.query_table_schema(self.table_path))
 
-    def _make_key_range(self):
-        if self.min_key is not None:
-            assert len(self.key_columns) == 1
-            (k,) = self.key_columns
-            yield self.min_key <= this[k]
-        if self.max_key is not None:
+    def _make_where(self):
+        # min_key <= x < max_key
+        if self.min_key is not None and self.max_key is not None:
             assert len(self.key_columns) == 1
             (k,) = self.key_columns
-            yield this[k] < self.max_key
+            yield from _make_range(this[k], self.min_key, self.max_key)
+
+        # min_update <= x < max_update
+        yield from _make_range(this[self.update_column], self.min_update, self.max_update)
 
-    def _make_update_range(self):
-        if self.min_update is not None:
-            yield self.min_update <= this[self.update_column]
-        if self.max_update is not None:
-            yield this[self.update_column] < self.max_update
+        # user-defined where
+        yield self.where or SKIP
 
     @property
     def source_table(self):
         return table(*self.table_path, schema=self._schema)
 
     def make_select(self):
-        return self.source_table.where(*self._make_key_range(), *self._make_update_range(), self.where or SKIP)
+        return self.source_table.where(*self._make_where())
 
     def get_values(self) -> list:
         "Download all the relevant values of the segment from the database"