Skip to content

Commit

Permalink
fix: dispatch underlying op for aliases and decimal literal
Browse files Browse the repository at this point in the history
`ops.Alias` was added in Ibis recently and the dispatcher was defaulting
to `ops.ValueOp` when it received one -- instead this added in a
specific Alias dispatch rule to pull out the underlying op.

Need to add proper decimal handling, but for now this at least keeps
things consistent between 2.1.1 and 3.x
  • Loading branch information
gforsyth authored and cpcloud committed May 11, 2022
1 parent 1bbf4aa commit 66e93e1
Show file tree
Hide file tree
Showing 4 changed files with 25 additions and 3 deletions.
21 changes: 21 additions & 0 deletions ibis_substrait/compiler/translate.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,17 +8,20 @@
import collections
import collections.abc
import datetime
import decimal
import functools
import itertools
import operator
import uuid
from typing import Any, Mapping, MutableMapping, Sequence, TypeVar

import ibis
import ibis.expr.datatypes as dt
import ibis.expr.operations as ops
import ibis.expr.schema as sch
import ibis.expr.types as ir
from ibis import util
from packaging import version

from ..proto.substrait import algebra_pb2 as stalg
from ..proto.substrait import type_pb2 as stt
Expand Down Expand Up @@ -264,6 +267,11 @@ def _literal_float64(_: dt.Float64, value: float) -> stalg.Expression.Literal:
return stalg.Expression.Literal(fp64=value)


@translate_literal.register
def _literal_decimal(_: dt.Decimal, value: decimal.Decimal) -> stalg.Expression.Literal:
raise NotImplementedError


@translate_literal.register
def _literal_string(_: dt.String, value: str) -> stalg.Expression.Literal:
return stalg.Expression.Literal(string=value)
Expand Down Expand Up @@ -441,6 +449,19 @@ def _translate_window_bounds(
return translate_preceding(*preceding), translate_following(*following)


if version.parse(ibis.__version__) >= version.parse("3.0.0"):

@translate.register(ops.Alias)
def alias_op(
op: ops.Alias,
expr: ir.ValueExpr,
compiler: SubstraitCompiler,
**kwargs: Any,
) -> stalg.Expression:
# For an alias, dispatch on the underlying argument
return translate(op.arg.op(), op.arg, compiler, **kwargs)


@translate.register(ops.ValueOp)
def value_op(
op: ops.ValueOp,
Expand Down
2 changes: 1 addition & 1 deletion ibis_substrait/tests/compiler/test_literal.py
Original file line number Diff line number Diff line change
Expand Up @@ -303,7 +303,7 @@ def test_literal(compiler, expr, ir):


@pytest.mark.xfail(
raises=ibis.common.exceptions.IbisTypeError,
raises=(ibis.common.exceptions.IbisTypeError, NotImplementedError),
reason="Ibis doesn't allow decimal values through validation",
)
def test_decimal_literal(compiler):
Expand Down
4 changes: 2 additions & 2 deletions poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ python = ">=3.8,<3.11"
ibis-framework = ">=2,<4"
protobuf = "^3.19.4"
platformdirs = "<2.5.2"
packaging = "^21.3"

[tool.poetry.dev-dependencies]
black = "^21.9b0"
Expand Down

0 comments on commit 66e93e1

Please sign in to comment.