Skip to content
This repository has been archived by the owner on Mar 29, 2023. It is now read-only.

Commit

Permalink
feat: Add rewrite function for ops.FloorDivide (#85)
Browse files Browse the repository at this point in the history
* feat: Add floor divide rewrite function

* chore: release v1.0.0 with latest fixes (#100)

* chore: release v1.0.0 with latest fixes

* update changelog

* update pyarrow dependency

* chore: add author

* chore: fix classifier

* restore rewrite

Co-authored-by: Timothy Dijamco <timothy.dijamco@twosigma.com>
Co-authored-by: Tim Swast <swast@google.com>
  • Loading branch information
3 people authored Sep 16, 2022
1 parent b35cda9 commit a04a674
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 0 deletions.
6 changes: 6 additions & 0 deletions ibis_bigquery/rewrites.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand Down
9 changes: 9 additions & 0 deletions tests/system/test_compiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down

0 comments on commit a04a674

Please sign in to comment.