Skip to content

Commit 45750bf

Browse files
nielspardonnielspardon
authored
fix(isthmus): fix argument order for trim function mapping (#714)
The Calcite/SQL `trim` function has the following syntax: ``` TRIM(BOTH characters FROM input) TRIM(LEADING characters FROM input) TRIM(TRAILING characters FROM input) ``` https://calcite.apache.org/docs/reference.html#character-string-operators-and-functions While the Substrait `trim`, `ltrim`, `rtrim` functions are defined as: ``` trim(input, characters) ltrim(input, characters) rtrim(input, characters) ``` https://github.com/substrait-io/substrait/blob/ee109ad48f149f2e3706d51b597febe86b79b225/extensions/functions_string.yaml#L1299-L1363 This PR fixes the argument order for the trim function mapping in isthmus by swapping the `input` and `characters` arguments during conversion. Signed-off-by: Niels Pardon <par@zurich.ibm.com>
1 parent 5462443 commit 45750bf

2 files changed

Lines changed: 153 additions & 7 deletions

File tree

‎isthmus/src/main/java/io/substrait/isthmus/expression/TrimFunctionMapper.java‎

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -23,12 +23,13 @@
2323
/**
2424
* Custom mapping for the Calcite TRIM function to various Substrait functions. The first TRIM
2525
* operand indicates the Substrait function to which it should be mapped. The first operand is then
26-
* omitted from the arguments supplied to the Substrait function.
26+
* omitted from the arguments supplied to the Substrait function. The second and third operands are
27+
* swapped in order.
2728
*
2829
* <ul>
29-
* <li>TRIM('BOTH', characters, string) -> trim(characters, string)
30-
* <li>TRIM('LEADING', characters, string) -> ltrim(characters, string)
31-
* <li>TRIM('TRAILING', .characters, string) -> rtrim(characters, string)
30+
* <li>TRIM('BOTH', characters, string) -> trim(string, characters)
31+
* <li>TRIM('LEADING', characters, string) -> ltrim(string, characters)
32+
* <li>TRIM('TRAILING', .characters, string) -> rtrim(string, characters)
3233
* </ul>
3334
*/
3435
final class TrimFunctionMapper implements ScalarFunctionMapper {
@@ -99,8 +100,11 @@ public Optional<SubstraitFunctionMapping> toSubstrait(final RexCall call) {
99100
}
100101

101102
String name = trim.substraitName();
102-
List<RexNode> operands =
103-
call.getOperands().stream().skip(1).collect(Collectors.toUnmodifiableList());
103+
List<RexNode> operands = call.getOperands().stream().skip(1).collect(Collectors.toList());
104+
105+
// Substrait expects (string, characters) while Calcite has (characters, string)
106+
Collections.swap(operands, 0, 1);
107+
104108
return new SubstraitFunctionMapping(name, operands, functions);
105109
});
106110
}
@@ -129,7 +133,9 @@ public Optional<List<FunctionArg>> getExpressionArguments(
129133
.map(EnumArg::of)
130134
.map(
131135
trimTypeArg -> {
132-
LinkedList args = new LinkedList<>(expression.arguments());
136+
LinkedList<FunctionArg> args = new LinkedList<>(expression.arguments());
137+
// Substrait expects (string, characters) while Calcite has (characters, string)
138+
Collections.swap(args, 0, 1);
133139
args.addFirst(trimTypeArg);
134140
return args;
135141
});

‎isthmus/src/test/java/io/substrait/isthmus/expression/TrimFunctionMapperTest.java‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
package io.substrait.isthmus.expression;
2+
3+
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
5+
import io.substrait.expression.EnumArg;
6+
import io.substrait.expression.Expression.ScalarFunctionInvocation;
7+
import io.substrait.expression.Expression.StrLiteral;
8+
import io.substrait.expression.FunctionArg;
9+
import io.substrait.extension.DefaultExtensionCatalog;
10+
import io.substrait.isthmus.PlanTestBase;
11+
import io.substrait.type.TypeCreator;
12+
import java.util.List;
13+
import java.util.Optional;
14+
import org.apache.calcite.rex.RexBuilder;
15+
import org.apache.calcite.rex.RexCall;
16+
import org.apache.calcite.rex.RexNode;
17+
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
18+
import org.apache.calcite.sql.fun.SqlTrimFunction;
19+
import org.junit.jupiter.api.Test;
20+
21+
class TrimFunctionMapperTest extends PlanTestBase {
22+
final TrimFunctionMapper trimFunctionMapper =
23+
new TrimFunctionMapper(DefaultExtensionCatalog.DEFAULT_COLLECTION.scalarFunctions());
24+
25+
final RexBuilder rexBuilder = builder.getRexBuilder();
26+
27+
@Test
28+
void calciteTrimBothArgumentOrder() {
29+
final RexNode characters = rexBuilder.makeLiteral(" ");
30+
final RexNode input = rexBuilder.makeLiteral(" whitespace ");
31+
final RexNode trimBothRexCall =
32+
rexBuilder.makeCall(
33+
SqlStdOperatorTable.TRIM,
34+
List.of(rexBuilder.makeFlag(SqlTrimFunction.Flag.BOTH), characters, input));
35+
36+
Optional<SubstraitFunctionMapping> substraitCall =
37+
trimFunctionMapper.toSubstrait((RexCall) trimBothRexCall);
38+
39+
assertEquals("trim", substraitCall.get().substraitName());
40+
// operands should be swapped now
41+
assertEquals(input, substraitCall.get().operands().get(0));
42+
assertEquals(characters, substraitCall.get().operands().get(1));
43+
}
44+
45+
@Test
46+
void calciteTrimLeadingArgumentOrder() {
47+
final RexNode characters = rexBuilder.makeLiteral(" ");
48+
final RexNode input = rexBuilder.makeLiteral(" whitespace ");
49+
final RexNode trimBothRexCall =
50+
rexBuilder.makeCall(
51+
SqlStdOperatorTable.TRIM,
52+
List.of(rexBuilder.makeFlag(SqlTrimFunction.Flag.LEADING), characters, input));
53+
54+
Optional<SubstraitFunctionMapping> substraitCall =
55+
trimFunctionMapper.toSubstrait((RexCall) trimBothRexCall);
56+
57+
assertEquals("ltrim", substraitCall.get().substraitName());
58+
// operands should be swapped now
59+
assertEquals(input, substraitCall.get().operands().get(0));
60+
assertEquals(characters, substraitCall.get().operands().get(1));
61+
}
62+
63+
@Test
64+
void calciteTrimTrailingArgumentOrder() {
65+
final RexNode characters = rexBuilder.makeLiteral(" ");
66+
final RexNode input = rexBuilder.makeLiteral(" whitespace ");
67+
final RexNode trimBothRexCall =
68+
rexBuilder.makeCall(
69+
SqlStdOperatorTable.TRIM,
70+
List.of(rexBuilder.makeFlag(SqlTrimFunction.Flag.TRAILING), characters, input));
71+
72+
Optional<SubstraitFunctionMapping> substraitCall =
73+
trimFunctionMapper.toSubstrait((RexCall) trimBothRexCall);
74+
75+
assertEquals("rtrim", substraitCall.get().substraitName());
76+
// operands should be swapped now
77+
assertEquals(input, substraitCall.get().operands().get(0));
78+
assertEquals(characters, substraitCall.get().operands().get(1));
79+
}
80+
81+
@Test
82+
void substraitTrimArgumentOrder() {
83+
final StrLiteral characters = sb.str(" ");
84+
final StrLiteral input = sb.str(" whitespace ");
85+
ScalarFunctionInvocation trimFn =
86+
sb.scalarFn(
87+
DefaultExtensionCatalog.FUNCTIONS_STRING,
88+
"trim:str_str",
89+
TypeCreator.REQUIRED.STRING,
90+
input,
91+
characters);
92+
93+
Optional<List<FunctionArg>> arguments = trimFunctionMapper.getExpressionArguments(trimFn);
94+
95+
assertEquals(EnumArg.of(SqlTrimFunction.Flag.BOTH.name()), arguments.get().get(0));
96+
// arguments should be swapped now
97+
assertEquals(characters, arguments.get().get(1));
98+
assertEquals(input, arguments.get().get(2));
99+
}
100+
101+
@Test
102+
void substraitLtrimArgumentOrder() {
103+
final StrLiteral characters = sb.str(" ");
104+
final StrLiteral input = sb.str(" whitespace ");
105+
ScalarFunctionInvocation trimFn =
106+
sb.scalarFn(
107+
DefaultExtensionCatalog.FUNCTIONS_STRING,
108+
"ltrim:str_str",
109+
TypeCreator.REQUIRED.STRING,
110+
input,
111+
characters);
112+
113+
Optional<List<FunctionArg>> arguments = trimFunctionMapper.getExpressionArguments(trimFn);
114+
115+
assertEquals(EnumArg.of(SqlTrimFunction.Flag.LEADING.name()), arguments.get().get(0));
116+
// arguments should be swapped now
117+
assertEquals(characters, arguments.get().get(1));
118+
assertEquals(input, arguments.get().get(2));
119+
}
120+
121+
@Test
122+
void substraitRtrimArgumentOrder() {
123+
final StrLiteral characters = sb.str(" ");
124+
final StrLiteral input = sb.str(" whitespace ");
125+
ScalarFunctionInvocation trimFn =
126+
sb.scalarFn(
127+
DefaultExtensionCatalog.FUNCTIONS_STRING,
128+
"rtrim:str_str",
129+
TypeCreator.REQUIRED.STRING,
130+
input,
131+
characters);
132+
133+
Optional<List<FunctionArg>> arguments = trimFunctionMapper.getExpressionArguments(trimFn);
134+
135+
assertEquals(EnumArg.of(SqlTrimFunction.Flag.TRAILING.name()), arguments.get().get(0));
136+
// arguments should be swapped now
137+
assertEquals(characters, arguments.get().get(1));
138+
assertEquals(input, arguments.get().get(2));
139+
}
140+
}

0 commit comments

Comments
 (0)