diff --git a/ibis_bigquery/rewrites.py b/ibis_bigquery/rewrites.py index 9522de8..508a251 100644 --- a/ibis_bigquery/rewrites.py +++ b/ibis_bigquery/rewrites.py @@ -11,6 +11,11 @@ def bigquery_day_of_week_name(e): return arg.strftime("%A") +def bq_floor_divide(expr): + left, right = expr.op().args + return left.div(right).floor() + + def identical_to(expr): left, right = expr.op().args return (left.isnull() & right.isnull()) | (left == right) @@ -46,6 +51,7 @@ def bigquery_any_all_no_op(expr): REWRITES = { **sql_compiler.ExprTranslator._rewrites, ops.DayOfWeekName: bigquery_day_of_week_name, + ops.FloorDivide: bq_floor_divide, ops.IdenticalTo: identical_to, ops.Log2: log2, ops.Sum: bq_sum, diff --git a/tests/system/test_compiler.py b/tests/system/test_compiler.py index d12459d..7c78b9b 100644 --- a/tests/system/test_compiler.py +++ b/tests/system/test_compiler.py @@ -70,6 +70,15 @@ def test_ieee_divide(alltypes, project_id, dataset_id): assert result == expected +def test_floor_divide(alltypes, project_id, dataset_id): + expr = alltypes.double_col // 0 + result = expr.compile() + expected = f"""\ +SELECT CAST(FLOOR(IEEE_DIVIDE(`double_col`, 0)) AS INT64) AS `tmp` +FROM `{project_id}.{dataset_id}.functional_alltypes`""" + assert result == expected + + def test_identical_to(alltypes, project_id, dataset_id): t = alltypes pred = t.string_col.identical_to("a") & t.date_string_col.identical_to("b")