fix(decompile): handle isin

ibis-project · Nov 25, 2023 · 6857751 · 6857751
1 parent b20a790
commit 6857751
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/ibis/expr/decompile.py b/ibis/expr/decompile.py
@@ -286,6 +286,11 @@ def binary(op, left, right):
     return f"{left} {operator} {right}"
 
 
+@translate.register(ops.InValues)
+def isin(op, value, options):
+    return f"{value}.isin(({', '.join([str(option) for option in options])}))"
+
+
 class CodeContext:
     always_assign = (ops.ScalarParameter, ops.UnboundTable, ops.Aggregation)
     always_ignore = (ops.TableColumn, dt.Primitive, dt.Variadic, dt.Temporal)

diff --git a/ibis/expr/tests/test_decompile.py b/ibis/expr/tests/test_decompile.py
@@ -20,6 +20,7 @@
     name="countries",
 )
 asian_countries = countries.filter(countries.continent == "AS")
+eurasian_countries = countries.filter(countries.continent.isin(("AS", "EUR")))
 top_with_highest_population = asian_countries.order_by(
     asian_countries.population.desc()
 ).limit(10)
@@ -51,13 +52,15 @@ def test_decompile_invalid_type():
         (top_with_highest_population, top_with_highest_population),
         (overall_population_density, overall_population_density),
         (population_density_per_country, population_density_per_country),
+        (eurasian_countries, eurasian_countries),
         (three, 3),
         (nine, nine_),
     ],
     ids=[
         "top_with_highest_population",
         "overall_population_density",
         "population_density_per_country",
+        "eurasian_countries",
         "three",
         "nine",
     ],