refactor: update various modules to use new wrapper and improve code

Update _util.py with code simplifications
Update arrays, linker, sets, and tf modules to use new wrapper
Add assert_equal utility function to tests

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/linker/_lsh.py

@@ -6,7 +6,7 @@
 from ibis import _
 from ibis.expr import types as ir
 
-from mismo._util import get_column
+from mismo._util import bind_one
 from mismo.arrays import array_choice
 from mismo.linker import _common
 
@@ -79,8 +79,8 @@ def __init__(
 
     def __call__(self, left: ir.Table, right: ir.Table) -> ir.Table:
         """Block two tables using Minhash LSH."""
-        left_terms = get_column(left, self.terms_column)
-        right_terms = get_column(right, self.terms_column)
+        left_terms = bind_one(left, self.terms_column)
+        right_terms = bind_one(right, self.terms_column)
         keys_name = self.keys_column.format(terms_column=left_terms.get_name())
         left = left.mutate(
             minhash_lsh_keys(

mismo/sets/_tfidf.py

@@ -166,7 +166,7 @@ def add_array_value_counts(
     │ NULL                         │ NULL                             │
     └──────────────────────────────┴──────────────────────────────────┘
     """  # noqa: E501
-    t = t.mutate(__terms=_util.get_column(t, column))
+    t = t.mutate(__terms=_util.bind_one(t, column))
     normalized = (
         t.select("__terms")
         .distinct()

mismo/tests/util.py

@@ -4,6 +4,7 @@
 
 import ibis
 from ibis.expr import types as ir
+from ibis.tests.util import assert_equal as ibis_assert_equal
 import pandas as pd
 import pytest
 
@@ -47,6 +48,10 @@ def assert_tables_equal(
     assert left_records == right_records
 
 
+def assert_equal(left: ibis.Value, right: ibis.Value) -> None:
+    ibis_assert_equal(left, right)
+
+
 def make_record_approx(record: dict) -> dict:
     return {k: make_float_comparable(v) for k, v in record.items()}
 

mismo/tf/_tf.py

@@ -9,7 +9,7 @@
 from ibis.expr import types as ir
 
 from mismo import _util
-from mismo.types._table_wrapper import TableWrapper
+from mismo.types._wrapper import TableWrapper
 
 T = TypeVar("T", bound=ir.Column)
 
@@ -52,24 +52,24 @@ def add_frequencies(
         self,
         table: ibis.Table,
         *,
-        column: str | ibis.Deferred | ibis.Column | None = None,
+        column: str | ibis.Deferred | ibis.Value | None = None,
         name_as: str | None = None,
-        default: Literal["1/N"] | int | float = "1/N",
+        default: Literal["1/N"] | int | float | ibis.Scalar = "1/N",
     ) -> ibis.Table:
         """Add frequency columns to the given table."""
         if name_as is None:
             name_as = f"frequency_{self.name}"
         if column is None:
             column = self.name
 
+        default_resolved: ibis.Scalar
         if default == "1/N" or default == "1/n":
             n_total = table.count().as_scalar()
-            default = (1 / n_total).cast("float64")
+            default_resolved = (1 / n_total).cast("float64")
         elif isinstance(default, ibis.Scalar):
-            default = default.cast("float64")
+            default_resolved = default.cast("float64")  # ty:ignore[invalid-assignment]
         else:
-            default = ibis.literal(default, "float64")
-
+            default_resolved = ibis.literal(default, "float64")
         table_column = _util.bind_one(table, column)
 
         unique_name = _util.unique_name("join_key")
@@ -82,7 +82,7 @@ def add_frequencies(
             .as_table()
             .distinct()
             .anti_join(stats_raw, unique_name)
-            .mutate(default.name(name_as))
+            .mutate(default_resolved.name(name_as))
         )
         stats = ibis.union(stats_raw, filler)
         assert stats.columns == (unique_name, name_as)
@@ -93,6 +93,8 @@ def add_frequencies(
 
 
 class TermFrequencyModel:
+    columns: dict[str, ir.Column]
+
     def __init__(
         self,
         columns: ibis.Table
@@ -102,7 +104,7 @@ def __init__(
         /,
     ):
         if isinstance(columns, Mapping):
-            self.columns = columns
+            self.columns = dict(columns)
         elif isinstance(columns, ibis.Table):
             self.columns = {col: columns[col] for col in columns.columns}
         elif isinstance(columns, ibis.Column):