Skip to content

Commit beca56c

Browse files
gforsythcpcloud
andcommitted
fix(extract): substrait function argument ordering
According to the spec, you cannot have the function option argument specified after a "regular" argument. Update ibis_substrait/compiler/translate.py Co-authored-by: Phillip Cloud <417981+cpcloud@users.noreply.github.com>
1 parent d70e953 commit beca56c

File tree

2 files changed

+10
-8
lines changed

2 files changed

+10
-8
lines changed

ibis_substrait/compiler/mapping.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,6 @@
3434
v: getattr(ops, k) for k, v in IBIS_SUBSTRAIT_OP_MAPPING.items()
3535
}
3636
# override when reversing many-to-one mappings
37-
SUBSTRAIT_IBIS_OP_MAPPING["extract"] = lambda table, span: getattr(
37+
SUBSTRAIT_IBIS_OP_MAPPING["extract"] = lambda span, table: getattr(
3838
ops, f"Extract{span.capitalize()}"
3939
)(table)

ibis_substrait/compiler/translate.py

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -968,16 +968,18 @@ def _extractdatefield(
968968
compiler: SubstraitCompiler,
969969
**kwargs: Any,
970970
) -> stalg.Expression:
971+
# e.g. "ExtractYear" -> "YEAR"
972+
span = type(op).__name__[len("Extract") :].upper()
973+
args = (
974+
translate(arg, compiler, **kwargs)
975+
for arg in op.args
976+
if isinstance(arg, ir.Expr)
977+
)
978+
971979
scalar_func = stalg.Expression.ScalarFunction(
972980
function_reference=compiler.function_id(expr),
973981
output_type=translate(expr.type()),
974-
args=[
975-
translate(arg, compiler, **kwargs)
976-
for arg in op.args
977-
if isinstance(arg, ir.Expr)
978-
],
979982
)
980-
# e.g. "ExtractYear" -> "YEAR"
981-
span = type(op).__name__.lstrip("Extract").upper()
982983
scalar_func.args.add(enum=stalg.Expression.Enum(specified=span))
984+
scalar_func.args.extend(args)
983985
return stalg.Expression(scalar_function=scalar_func)

0 commit comments

Comments
 (0)