Skip to content

Commit 318f831

Browse files
yassramyassram
committed
test(isthmus): additional tests for substrait to Calcite with Sort
1 parent 3338960 commit 318f831

File tree

1 file changed

+205
-0
lines changed

1 file changed

+205
-0
lines changed

‎isthmus/src/test/java/io/substrait/isthmus/ComplexSortTest.java

Lines changed: 205 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,205 @@
1+
package io.substrait.isthmus;
2+
3+
import static io.substrait.isthmus.SqlConverterBase.EXTENSION_COLLECTION;
4+
import static org.junit.jupiter.api.Assertions.assertEquals;
5+
6+
import io.substrait.dsl.SubstraitBuilder;
7+
import io.substrait.expression.Expression;
8+
import io.substrait.relation.Rel;
9+
import io.substrait.type.TypeCreator;
10+
import java.io.PrintWriter;
11+
import java.io.StringWriter;
12+
import java.util.List;
13+
import org.apache.calcite.rel.RelNode;
14+
import org.apache.calcite.rel.externalize.RelWriterImpl;
15+
import org.apache.calcite.sql.SqlExplainLevel;
16+
import org.apache.calcite.util.Pair;
17+
import org.checkerframework.checker.nullness.qual.Nullable;
18+
import org.junit.jupiter.api.Test;
19+
20+
public class ComplexSortTest extends PlanTestBase {
21+
22+
final TypeCreator R = TypeCreator.of(false);
23+
SubstraitBuilder b = new SubstraitBuilder(extensions);
24+
25+
final SubstraitToCalcite substraitToCalcite =
26+
new SubstraitToCalcite(EXTENSION_COLLECTION, typeFactory);
27+
28+
public static class CollationRelWriter extends RelWriterImpl {
29+
public CollationRelWriter(StringWriter sw) {
30+
super(new PrintWriter(sw), SqlExplainLevel.EXPPLAN_ATTRIBUTES, false);
31+
}
32+
33+
@Override
34+
protected void explain_(RelNode rel, List<Pair<String, @Nullable Object>> values) {
35+
var collation = rel.getTraitSet().getCollation();
36+
if (!collation.isDefault()) {
37+
StringBuilder s = new StringBuilder();
38+
spacer.spaces(s);
39+
s.append("Collation: ").append(collation.toString());
40+
pw.println(s);
41+
}
42+
super.explain_(rel, values);
43+
}
44+
}
45+
46+
@Test
47+
void handleInputReferenceSort() {
48+
// CREATE TABLE example (a VARCHAR)
49+
// SELECT a FROM example ORDER BY a
50+
51+
Rel rel =
52+
b.project(
53+
input -> b.fieldReferences(input, 0),
54+
b.remap(1),
55+
b.sort(
56+
input ->
57+
List.of(
58+
b.sortField(
59+
b.fieldReference(input, 0), Expression.SortDirection.ASC_NULLS_LAST)),
60+
b.namedScan(List.of("example"), List.of("a"), List.of(R.STRING))));
61+
62+
String expected =
63+
"""
64+
Collation: [0]
65+
LogicalSort(sort0=[$0], dir0=[ASC])
66+
LogicalTableScan(table=[[example]])
67+
""";
68+
69+
RelNode relReturned = substraitToCalcite.convert(rel);
70+
var sw = new StringWriter();
71+
relReturned.explain(new CollationRelWriter(sw));
72+
assertEquals(expected, sw.toString());
73+
}
74+
75+
@Test
76+
void handleCastExpressionSort() {
77+
// CREATE TABLE example (a VARCHAR)
78+
// SELECT a FROM example ORDER BY a::INT
79+
80+
Rel rel =
81+
b.project(
82+
input -> b.fieldReferences(input, 0),
83+
b.remap(1),
84+
b.sort(
85+
input ->
86+
List.of(
87+
b.sortField(
88+
b.cast(b.fieldReference(input, 0), R.I32),
89+
Expression.SortDirection.ASC_NULLS_LAST)),
90+
b.namedScan(List.of("example"), List.of("a"), List.of(R.STRING))));
91+
92+
String expected =
93+
"""
94+
LogicalProject(a0=[$0])
95+
Collation: [1]
96+
LogicalSort(sort0=[$1], dir0=[ASC])
97+
LogicalProject(a=[$0], a0=[CAST($0):INTEGER NOT NULL])
98+
LogicalTableScan(table=[[example]])
99+
""";
100+
101+
RelNode relReturned = substraitToCalcite.convert(rel);
102+
var sw = new StringWriter();
103+
relReturned.explain(new CollationRelWriter(sw));
104+
assertEquals(expected, sw.toString());
105+
}
106+
107+
@Test
108+
void handleCastProjectAndSortWithSortDirection() {
109+
// CREATE TABLE example (a VARCHAR)
110+
// SELECT a::INT FROM example ORDER BY a::INT DESC NULLS LAST
111+
112+
Rel rel =
113+
b.project(
114+
input -> List.of(b.cast(b.fieldReference(input, 0), R.I32)),
115+
b.remap(1),
116+
b.sort(
117+
input ->
118+
List.of(
119+
b.sortField(
120+
b.cast(b.fieldReference(input, 0), R.I32),
121+
Expression.SortDirection.DESC_NULLS_LAST)),
122+
b.namedScan(List.of("example"), List.of("a"), List.of(R.STRING))));
123+
124+
String expected =
125+
"""
126+
LogicalProject(a0=[CAST($0):INTEGER NOT NULL])
127+
Collation: [1 DESC-nulls-last]
128+
LogicalSort(sort0=[$1], dir0=[DESC-nulls-last])
129+
LogicalProject(a=[$0], a0=[CAST($0):INTEGER NOT NULL])
130+
LogicalTableScan(table=[[example]])
131+
""";
132+
133+
RelNode relReturned = substraitToCalcite.convert(rel);
134+
var sw = new StringWriter();
135+
relReturned.explain(new CollationRelWriter(sw));
136+
assertEquals(expected, sw.toString());
137+
}
138+
139+
@Test
140+
void handleCastSortToOriginalType() {
141+
// CREATE TABLE example (a VARCHAR)
142+
// SELECT a FROM example ORDER BY a::VARCHAR
143+
144+
Rel rel =
145+
b.project(
146+
input -> List.of(b.fieldReference(input, 0)),
147+
b.remap(1),
148+
b.sort(
149+
input ->
150+
List.of(
151+
b.sortField(
152+
b.cast(b.fieldReference(input, 0), R.STRING),
153+
Expression.SortDirection.DESC_NULLS_LAST)),
154+
b.namedScan(List.of("example"), List.of("a"), List.of(R.STRING))));
155+
156+
String expected =
157+
"""
158+
LogicalProject(a0=[$0])
159+
Collation: [1 DESC-nulls-last]
160+
LogicalSort(sort0=[$1], dir0=[DESC-nulls-last])
161+
LogicalProject(a=[$0], a0=[$0])
162+
LogicalTableScan(table=[[example]])
163+
""";
164+
165+
RelNode relReturned = substraitToCalcite.convert(rel);
166+
var sw = new StringWriter();
167+
relReturned.explain(new CollationRelWriter(sw));
168+
assertEquals(expected, sw.toString());
169+
}
170+
171+
@Test
172+
void handleComplex2ExpressionSort() {
173+
// CREATE TABLE example (a VARCHAR, b INT)
174+
// SELECT b, a FROM example ORDER BY a::INT > 0 DESC, -b + 42 ASC NULLS LAST
175+
176+
Rel rel =
177+
b.project(
178+
input -> List.of(b.fieldReference(input, 0), b.fieldReference(input, 1)),
179+
b.remap(2, 3),
180+
b.sort(
181+
input ->
182+
List.of(
183+
b.sortField(
184+
b.cast(b.fieldReference(input, 0), R.I32),
185+
Expression.SortDirection.DESC_NULLS_FIRST),
186+
b.sortField(
187+
b.add(b.negate(b.cast(b.fieldReference(input, 0), R.I32)), b.i32(42)),
188+
Expression.SortDirection.ASC_NULLS_LAST)),
189+
b.namedScan(List.of("example"), List.of("a", "b"), List.of(R.STRING, R.I64))));
190+
191+
String expected =
192+
"""
193+
LogicalProject(a0=[$0], b0=[$1])
194+
Collation: [2 DESC, 3]
195+
LogicalSort(sort0=[$2], sort1=[$3], dir0=[DESC], dir1=[ASC])
196+
LogicalProject(a=[$0], b=[$1], a0=[CAST($0):INTEGER NOT NULL], $f3=[+(-(CAST($0):INTEGER NOT NULL), 42)])
197+
LogicalTableScan(table=[[example]])
198+
""";
199+
200+
RelNode relReturned = substraitToCalcite.convert(rel);
201+
var sw = new StringWriter();
202+
relReturned.explain(new CollationRelWriter(sw));
203+
assertEquals(expected, sw.toString());
204+
}
205+
}

0 commit comments

Comments
 (0)