Experimental const-T-list

src/pydiverse/transform/core/ops/logical.py

@@ -16,6 +16,7 @@
     "Xor",
     "RXor",
     "Invert",
+    "Isin",
 ]
 
 
@@ -104,3 +105,10 @@ class Invert(ElementWise, Unary):
     signatures = [
         "bool -> bool",
     ]
+
+
+class Isin(Comparison):
+    name = "isin"
+    signatures = [
+        "T, const-T-list -> bool",
+    ]

src/pydiverse/transform/core/ops/registry.py

@@ -323,6 +323,8 @@ def parse_cstypes(cst: str):
         for type_ in (*base_args_types, base_rtype):
             if type_.startswith("const-"):
                 type_ = type_[6:]  # only look at type after dash
+            if type_.endswith("-list"):
+                type_ = type_[:-5]  # only look at type before dash
             if not type_.isalnum():
                 raise ValueError(f"Invalid type '{type_}'. Types must be alphanumeric.")
 

src/pydiverse/transform/core/table_impl.py

@@ -249,6 +249,23 @@ def const_func(*args, **kwargs):
                 return TypedValue(literal_func, "float", const_value=const_func)
             if isinstance(expr, str):
                 return TypedValue(literal_func, "str", const_value=const_func)
+            if isinstance(expr, list):
+
+                def assert_func(*args, **kwargs):
+                    raise ValueError(
+                        "List type is currently only supported for constants, "
+                        f"not: {expr}"
+                    )
+
+                if all(isinstance(exp, bool) for exp in expr):
+                    return TypedValue(assert_func, "bool-list", const_value=const_func)
+                if all(isinstance(exp, int) for exp in expr):
+                    return TypedValue(assert_func, "int-list", const_value=const_func)
+                if all(isinstance(exp, float) for exp in expr):
+                    return TypedValue(assert_func, "float-list", const_value=const_func)
+                if all(isinstance(exp, str) for exp in expr):
+                    return TypedValue(assert_func, "str-list", const_value=const_func)
+            raise ValueError(f"Unknown expression type: {expr}")
 
     class AlignedExpressionEvaluator(Generic[AlignedT], DelegatingTranslator[AlignedT]):
         """

src/pydiverse/transform/eager/pandas_table.py

@@ -685,3 +685,10 @@ def _rank(x: pd.Series):
     @op.auto(variant="transform")
     def _rank(x: pd.Series):
         return x.transform("rank", method="min")
+
+
+with PandasTableImpl.op(ops.Isin()) as op:
+
+    @op.auto
+    def _isin(x: pd.Series, *args):
+        return x.isin(args)

src/pydiverse/transform/lazy/sql_table/sql_table.py

@@ -913,3 +913,10 @@ def _rank(x):
         # row_number() is like rank(method="first")
         # rank() is like method="min"
         return sa.func.RANK(), ImplicitArrange(sql.expression.ClauseList(x))
+
+
+with SQLTableImpl.op(ops.Isin()) as op:
+
+    @op.auto
+    def _isin(x, *args):
+        return x.in_(args)

tests/test_backend_equivalence.py

@@ -588,8 +588,10 @@ def test_noop(self, df2_x, df2_y):
 
     @tables(["df2"])
     def test_simple_filter(self, df2_x, df2_y):
-        assert_result_equal(df2_x, df2_y, lambda t: t >> filter(t.col1 == 2))
-        assert_result_equal(df2_x, df2_y, lambda t: t >> filter(t.col1 != 2))
+        # assert_result_equal(df2_x, df2_y, lambda t: t >> filter(t.col1 == 2))
+        # assert_result_equal(df2_x, df2_y, lambda t: t >> filter(t.col1 != 2))
+        # assert_result_equal(df2_x, df2_y, lambda t: t >> filter(t.col1.isin(2)))
+        assert_result_equal(df2_x, df2_y, lambda t: t >> filter(t.col1.isin([2, 3])))
 
     @tables(["df2"])
     def test_chained_filters(self, df2_x, df2_y):