feat: overhaul many things

Author: NickCrews

mismo/_counts_table.py

@@ -7,7 +7,7 @@
 from ibis import _
 from ibis.expr import types as ir
 
-from mismo.types._table_wrapper import TableWrapper
+from mismo.types._wrapper import TableWrapper
 
 if TYPE_CHECKING:
     import altair as alt

mismo/_util.py

@@ -89,13 +89,28 @@ def cases(
         return builder.else_(else_).end()
 
 
-@overload
-def bind(t: ibis.Table, ref: Any, /) -> tuple[ibis.Value, ...]: ...
-@overload
-def bind(t: ibis.Deferred, ref: Any, /) -> tuple[ibis.Deferred]: ...
+IntoValue = str | int | ibis.Deferred | ibis.Value | Callable[[ibis.Table], ibis.Value]
 
 
-def bind(t: ibis.Deferred | ibis.Table, ref: Any) -> tuple[ibis.Value, ...]:
+@overload
+def bind(
+    t: ibis.Table,
+    ref: IntoValue | Iterable[IntoValue] | Mapping[str, IntoValue],
+    /,
+) -> tuple[ibis.Value, ...]: ...
+@overload
+def bind(
+    t: ibis.Deferred,
+    ref: IntoValue | Iterable[IntoValue] | Mapping[str, IntoValue],
+    /,
+) -> tuple[ibis.Deferred]: ...
+
+
+def bind(
+    t: ibis.Deferred | ibis.Table,
+    ref: IntoValue | Iterable[IntoValue] | Mapping[str, IntoValue],
+    /,
+) -> tuple[ibis.Value, ...]:
     """Reference into a table to get Columns and Scalars.
 
     ibis._.bind(ref) does not work because it returns another Deferred.
@@ -110,49 +125,21 @@ def bind(t: ibis.Deferred | ibis.Table, ref: Any) -> tuple[ibis.Value, ...]:
 
 
 @overload
-def bind_one(t: ibis.Table, ref: Any, /) -> ibis.Value: ...
+def bind_one(t: ibis.Table, ref: IntoValue, /) -> ibis.Value: ...
 @overload
-def bind_one(t: ibis.Deferred, ref: Any, /) -> ibis.Deferred: ...
+def bind_one(t: ibis.Deferred, ref: IntoValue, /) -> ibis.Deferred: ...
 
 
-def bind_one(t: ibis.Deferred | ibis.Table, ref: Any) -> ibis.Value | ibis.Deferred:
+def bind_one(
+    t: ibis.Deferred | ibis.Table, ref: IntoValue, /
+) -> ibis.Value | ibis.Deferred:
     """Like bind(), but ensure that exactly one value is returned."""
     vals = bind(t, ref)
     if len(vals) != 1:
         raise ValueError(f"Expected 1 value, got {len(vals)} from {ref}")
     return vals[0]
 
 
-def get_column(
-    t: ir.Table, ref: Any, *, on_many: Literal["error", "struct"] = "error"
-) -> ir.Column:
-    """Get a column from a table using some sort of reference to the column.
-
-    ref can be a string, a Deferred, a callable, an ibis selector, etc.
-
-    Parameters
-    ----------
-    t :
-        The table
-    ref :
-        The reference to the column
-    on_many :
-        What to do if ref returns multiple columns. If "error", raise an error.
-        If "struct", return a StructColumn containing all the columns.
-    """
-    cols = bind(t, ref)
-    if isinstance(t, ibis.Deferred):
-        # This is by definition a single column
-        return cols[0]
-    if len(cols) != 1:
-        if on_many == "error":
-            raise ValueError(f"Expected 1 column, got {len(cols)}")
-        if on_many == "struct":
-            return ibis.struct({c.get_name(): c for c in cols})
-        raise ValueError(f"on_many must be 'error' or 'struct'. Got {on_many}")
-    return cols[0]
-
-
 def ensure_ibis(
     val: Any, type: str | dt.DataType | None = None
 ) -> ir.Value | ibis.Deferred:

mismo/arrays/_array.py

@@ -110,8 +110,8 @@ def array_filter_isin_other(
         The table with a new column named following `result_format` with the
         filtered array.
     """  # noqa E501
-    array_col = _util.get_column(t, array)
-    t = t.mutate(__array=array_col, __id=ibis.row_number())
+    array_val = _util.bind_one(t, array)
+    t = t.mutate(__array=array_val, __id=ibis.row_number())
     temp = t.select("__id", __unnested=_.__array.unnest())
     # When we re-.collect() items below, the order matters,
     # but the .filter() can mess up the order, so we need to
@@ -131,7 +131,7 @@ def array_filter_isin_other(
             [], _.__filtered
         )
     ).drop("__array")
-    result_name = result_format.format(name=array_col.get_name())
+    result_name = result_format.format(name=array_val.get_name())
     return re_joined.rename({result_name: "__filtered"})
 
 

mismo/lib/email/_core.py

@@ -5,7 +5,7 @@
 import ibis
 from ibis.expr import types as ir
 
-from mismo._util import cases, get_column
+from mismo._util import bind_one, cases
 from mismo.arrays import array_combinations, array_min
 from mismo.compare import MatchLevel
 from mismo.linker import UnnestLinker
@@ -164,7 +164,7 @@ def __init__(
     def prepare_for_fast_linking(self, t: ir.Table) -> ir.Table:
         """Add a column with the parsed and normalized email addresses."""
         return t.mutate(
-            get_column(t, self.column)
+            bind_one(t, self.column)
             .map(
                 lambda email: ParsedEmail(
                     clean_email(email, normalize=True)

mismo/lib/geo/_latlon.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
+from collections.abc import Mapping
 import dataclasses
 import math
-from typing import Callable
+from typing import Callable, Literal, TypedDict
 
 import ibis
-from ibis import Deferred
 from ibis.expr import types as ir
 
 from mismo import _util, linkage, types
@@ -57,6 +57,61 @@ def haversine(theta: ir.FloatingValue) -> ir.FloatingValue:
     return (R_earth * 2) * a.sqrt().asin()
 
 
+Coordinates = Mapping[Literal["lat", "lon"], ir.FloatingColumn]
+CoordinatesMapping = Mapping[Literal["lat", "lon"], _util.IntoValue]
+
+
+class CoordinatesDict(TypedDict):
+    lat: ir.FloatingColumn
+    lon: ir.FloatingColumn
+
+
+IntoCoordinates = (
+    str
+    | ibis.Deferred
+    | Coordinates
+    | CoordinatesMapping
+    | Callable[[ir.Table], "IntoCoordinates"]
+)
+
+
+def default_resolver() -> CoordinatesMapping:
+    return {"lat": "lat", "lon": "lon"}
+
+
+def get_coordinate_pair(
+    table: ir.Table,
+    mapping: IntoCoordinates,
+) -> CoordinatesDict:
+    """Get the coordinates from a table.
+
+    Parameters
+    ----------
+    table
+        The table to get the coordinates from.
+    mapping
+        A mapping of column names to use for the coordinates,
+        or a function that takes a table and returns the coordinate pair.
+
+    Returns
+    -------
+    coordinates
+        The coordinates.
+    """
+    if callable(mapping):
+        called = mapping(table)
+        return get_coordinate_pair(table, called)
+
+    if isinstance(mapping, str) or isinstance(mapping, ibis.Deferred):
+        coords_col = _util.bind_one(table, mapping)
+        lat: ir.FloatingColumn = coords_col["lat"]  # ty:ignore[not-subscriptable]
+        lon: ir.FloatingColumn = coords_col["lon"]  # ty:ignore[not-subscriptable]
+    else:
+        lat: ir.FloatingColumn = _util.bind_one(table, mapping["lat"])  # ty:ignore[invalid-assignment]
+        lon: ir.FloatingColumn = _util.bind_one(table, mapping["lon"])  # ty:ignore[invalid-assignment]
+    return CoordinatesDict(lat=lat, lon=lon)
+
+
 @dataclasses.dataclass(frozen=True)
 class CoordinateLinker:
     """Links two locations together if they are within a certain distance.
@@ -108,9 +163,8 @@ class CoordinateLinker:
     ... )
     >>> linker = CoordinateLinker(
     ...     distance_km=1,
-    ...     left_coord="latlon",
-    ...     right_lat="latitude",
-    ...     right_lon="longitude",
+    ...     left_resolver="latlon",
+    ...     right_resolver={"lat": "latitude", "lon": "longitude"},
     ... )
     >>> linker(left, right).links
     ┏━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━┓
@@ -130,127 +184,66 @@ class CoordinateLinker:
     """
     The (approx) max distance in kilometers that two coords will be blocked together.
     """
-    coord: str | Deferred | Callable[[ir.Table], ir.StructColumn] | None = None
-    """The column in both tables containing the `struct<lat: float, lon: float>` coordinates."""  # noqa: E501
-    lat: str | Deferred | Callable[[ir.Table], ir.FloatingColumn] | None = None
-    """The column in both tables containing the latitude coordinates."""
-    lon: str | Deferred | Callable[[ir.Table], ir.FloatingColumn] | None = None
-    """The column in both tables containing the longitude coordinates."""
-    left_coord: str | Deferred | Callable[[ir.Table], ir.StructColumn] | None = None
-    """The column in the left tables containing the `struct<lat: float, lon: float>` coordinates."""  # noqa: E501
-    right_coord: str | Deferred | Callable[[ir.Table], ir.StructColumn] | None = None
-    """The column in the right tables containing the `struct<lat: float, lon: float>` coordinates."""  # noqa: E501
-    left_lat: str | Deferred | Callable[[ir.Table], ir.FloatingColumn] | None = None
-    """The column in the left tables containing the latitude coordinates."""
-    left_lon: str | Deferred | Callable[[ir.Table], ir.FloatingColumn] | None = None
-    """The column in the left tables containing the longitude coordinates."""
-    right_lat: str | Deferred | Callable[[ir.Table], ir.FloatingColumn] | None = None
-    """The column in the right tables containing the latitude coordinates."""
-    right_lon: str | Deferred | Callable[[ir.Table], ir.FloatingColumn] | None = None
-    """The column in the right tables containing the longitude coordinates."""
     max_pairs: int | None = None
     """The maximum number of pairs that any single block of coordinates can contain.
     
     eg if you have 1000 records all with the same coordinates, this would
     naively result in ~(1000 * 1000) / 2 = 500_000 pairs.
     If we set max_pairs to less than this, this group of records will be skipped.
     """
-
-    def __post_init__(self):
-        ok_subsets = [
-            {"coord"},
-            {"left_coord", "right_coord"},
-            {"left_coord", "right_lat", "right_lon"},
-            {"left_lat", "left_lon", "right_coord"},
-            {"left_lat", "left_lon", "right_lat", "right_lon"},
-            {"lat", "lon"},
-            {"lat", "right_lat", "right_lon"},
-            {"left_lat", "left_lon", "lon"},
-        ]
-        options = [
-            "coord",
-            "left_coord",
-            "right_coord",
-            "lat",
-            "lon",
-            "left_lat",
-            "left_lon",
-            "right_lat",
-            "right_lon",
-        ]
-        present = {k for k in options if getattr(self, k) is not None}
-        if present not in ok_subsets:
-            ok_subsets_str = "\n".join("- " + str(s) for s in ok_subsets)
-            raise ValueError(
-                "You must specify exactly one of the following subsets of options:\n"
-                + ok_subsets_str
-                + f"\nYou provided:\n{present}"
-            )
-
-    def _left_coord(
-        self, left: ir.Table
-    ) -> tuple[ir.FloatingColumn, ir.FloatingColumn]:
-        if self.coord is not None:
-            left_coord = _util.get_column(left, self.coord, on_many="struct")
-            return (left_coord.lat, left_coord.lon)
-        if self.left_coord is not None:
-            left_coord = _util.get_column(left, self.left_coord, on_many="struct")
-            left_lat = left_coord.lat
-            left_lon = left_coord.lon
-        if self.lat is not None:
-            left_lat = _util.get_column(left, self.lat)
-        if self.lon is not None:
-            left_lon = _util.get_column(left, self.lon)
-        if self.left_lat is not None:
-            left_lat = _util.get_column(left, self.left_lat)
-        if self.left_lon is not None:
-            left_lon = _util.get_column(left, self.left_lon)
-        return left_lat, left_lon
-
-    def _right_coord(
-        self, right: ir.Table
-    ) -> tuple[ir.FloatingColumn, ir.FloatingColumn]:
-        if self.coord is not None:
-            right_coord = _util.get_column(right, self.coord, on_many="struct")
-            return (right_coord.lat, right_coord.lon)
-        if self.right_coord is not None:
-            right_coord = _util.get_column(right, self.right_coord, on_many="struct")
-            right_lat = right_coord.lat
-            right_lon = right_coord.lon
-        if self.lat is not None:
-            right_lat = _util.get_column(right, self.lat)
-        if self.lon is not None:
-            right_lon = _util.get_column(right, self.lon)
-        if self.right_lat is not None:
-            right_lat = _util.get_column(right, self.right_lat)
-        if self.right_lon is not None:
-            right_lon = _util.get_column(right, self.right_lon)
-        return right_lat, right_lon
-
-    def _join_keys(
-        self, left: ir.Table | ibis.Deferred, right: ir.Table | ibis.Deferred
-    ) -> tuple[tuple[ir.Column, ir.Column], tuple[ir.Column, ir.Column]]:
-        left_lat, left_lon = self._left_coord(left)
-        right_lat, right_lon = self._right_coord(right)
+    left_resolver: IntoCoordinates = dataclasses.field(default_factory=default_resolver)
+    """A specification of how to get the lat and lon values from the left table.
+    
+    Can be:
+    - A `str` or `ibis.Deferred`, which are assumed to point to a column
+      containing a struct with `lat` and `lon` fields.
+    - A `Mapping` of `{"lat": ..., "lon": ...}` where the values are
+      column names or `ibis.Deferred` expressions.
+    - A callable that takes a table and returns one of the above.
+    """
+    right_resolver: IntoCoordinates = dataclasses.field(
+        default_factory=default_resolver
+    )
+    """See `left_resolver`, but for the right table."""
+
+    def hash_coord(
+        self, coord: CoordinatesDict, /
+    ) -> tuple[ir.IntegerValue, ir.IntegerValue]:
+        lat, lon = coord["lat"], coord["lon"]
         # We have to use a grid size of ~3x the precision to avoid
         # two points falling right on either side of a grid cell boundary
         grid_size = self.distance_km * 3
-        left_lat_key, left_lon_key = _bin_lat_lon(left_lat, left_lon, grid_size)
-        right_lat_key, right_lon_key = _bin_lat_lon(right_lat, right_lon, grid_size)
-        return (
-            left_lat_key.name("lat_binned"),
-            right_lat_key.name("lat_binned"),
-        ), (
-            left_lon_key.name("lon_binned"),
-            right_lon_key.name("lon_binned"),
-        )
+        lat_key, lon_key = _bin_lat_lon(lat, lon, grid_size)
+        return lat_key, lon_key
+
+    def hash_left(self, t: ir.Table, /) -> tuple[ir.IntegerValue, ir.IntegerValue]:
+        coords = get_coordinate_pair(t, self.left_resolver)
+        return self.hash_coord(coords)
+
+    def hash_right(self, t: ir.Table, /) -> tuple[ir.IntegerValue, ir.IntegerValue]:
+        coords = get_coordinate_pair(t, self.right_resolver)
+        return self.hash_coord(coords)
 
     @property
     def _key_linker(self) -> KeyLinker:
         import mismo
 
-        lat_key, lon_key = self._join_keys(ibis._, ibis._)
-        return mismo.KeyLinker([lat_key, lon_key], max_pairs=self.max_pairs)
+        def lat_key_left(t: ir.Table) -> ir.IntegerValue:
+            return self.hash_left(t)[0]
+
+        def lon_key_left(t: ir.Table) -> ir.IntegerValue:
+            return self.hash_left(t)[1]
+
+        def lat_key_right(t: ir.Table) -> ir.IntegerValue:
+            return self.hash_right(t)[0]
+
+        def lon_key_right(t: ir.Table) -> ir.IntegerValue:
+            return self.hash_right(t)[1]
+
+        return mismo.KeyLinker(
+            [(lat_key_left, lat_key_right), (lon_key_left, lon_key_right)],
+            max_pairs=self.max_pairs,
+        )
 
     def __join_condition__(self, left: ir.Table, right: ir.Table) -> ir.BooleanValue:
         return self._key_linker.__join_condition__(left, right)