Skip to content

Commit

Permalink
feat: add support for ops.Log
Browse files Browse the repository at this point in the history
if `base` isn't specified we default to the value of `e`
  • Loading branch information
gforsyth authored and cpcloud committed Nov 2, 2022
1 parent 21f677e commit 226074e
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 0 deletions.
1 change: 1 addition & 0 deletions ibis_substrait/compiler/mapping.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
"Less": "lt",
"LessEqual": "lte",
"Ln": "ln",
"Log": "logb",
"Log2": "log2",
"Log10": "log10",
"Lowercase": "lower",
Expand Down
24 changes: 24 additions & 0 deletions ibis_substrait/compiler/translate.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,12 @@
import decimal
import functools
import itertools
import math
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
Expand Down Expand Up @@ -1054,3 +1056,25 @@ def _extractdatefield(
scalar_func.arguments.add(enum=stalg.FunctionArgument.Enum(specified=span))
scalar_func.arguments.extend(arguments)
return stalg.Expression(scalar_function=scalar_func)


@translate.register(ops.Log)
def _log(
op: ops.Log,
expr: ir.TableExpr,
compiler: SubstraitCompiler,
**kwargs: Any,
) -> stalg.Expression:
arg = stalg.FunctionArgument(value=translate(op.arg, compiler, **kwargs))
base = stalg.FunctionArgument(
value=translate(
op.base if op.base is not None else ibis.literal(math.e), compiler, **kwargs
)
)

scalar_func = stalg.Expression.ScalarFunction(
function_reference=compiler.function_id(expr),
output_type=translate(expr.type()),
arguments=[arg, base],
)
return stalg.Expression(scalar_function=scalar_func)

0 comments on commit 226074e

Please sign in to comment.