From 6857751e7d99f6e37b1edbc8c9046502c6c4d826 Mon Sep 17 00:00:00 2001 From: Andrew Schechtman-Rook Date: Fri, 24 Nov 2023 15:47:53 -0500 Subject: [PATCH] fix(decompile): handle isin --- ibis/expr/decompile.py | 5 +++++ ibis/expr/tests/test_decompile.py | 3 +++ 2 files changed, 8 insertions(+) diff --git a/ibis/expr/decompile.py b/ibis/expr/decompile.py index ee406ecaa9dc..a6cbf56de716 100644 --- 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 index 7ffae7753ba6..23fe6a41645a 100644 --- 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,6 +52,7 @@ 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_), ], @@ -58,6 +60,7 @@ def test_decompile_invalid_type(): "top_with_highest_population", "overall_population_density", "population_density_per_country", + "eurasian_countries", "three", "nine", ],