Add tests for CASE decomposition. (#15639)

I was looking into adding a rule to do this, and found that it was already
happening as part of Calcite's RexSimplify. So this patch simply adds some
tests to ensure that it continues to happen.
This commit is contained in:
Gian Merlino 2024-01-10 13:24:24 -08:00 committed by GitHub
parent 0b91cc4db2
commit ee77fa7fb3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 82 additions and 0 deletions

View File

@ -3209,6 +3209,88 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
); );
} }
@Test
public void testDecomposeCaseWhenThreeArg()
{
// Cannot vectorize due to virtual columns.
cannotVectorize();
testQuery(
"SELECT\n"
+ " dim1, dim2, CASE WHEN dim1 = 'abc' THEN dim1 ELSE dim2 END\n"
+ "FROM druid.foo\n"
+ "WHERE\n"
+ " CASE WHEN dim1 = 'abc' THEN dim1 ELSE dim2 END = 'abc'",
ImmutableList.of(
newScanQueryBuilder()
.dataSource(CalciteTests.DATASOURCE1)
.virtualColumns(
expressionVirtualColumn(
"v0",
"case_searched((\"dim1\" == 'abc'),\"dim1\",\"dim2\")",
ColumnType.STRING
)
)
.intervals(querySegmentSpec(Filtration.eternity()))
.filters(
or(
equality("dim1", "abc", ColumnType.STRING),
and(
equality("dim2", "abc", ColumnType.STRING),
NullHandling.sqlCompatible()
? not(istrue(equality("dim1", "abc", ColumnType.STRING)))
: not(equality("dim1", "abc", ColumnType.STRING))
)
)
)
.columns("dim1", "dim2", "v0")
.build()
),
ImmutableList.of(
new Object[]{"def", "abc", "abc"},
new Object[]{"abc", NullHandling.defaultStringValue(), "abc"}
)
);
}
@Test
public void testDecomposeCaseWhenTwoArg()
{
// Cannot vectorize due to virtual columns.
cannotVectorize();
testQuery(
"SELECT\n"
+ " dim1, dim2, CASE WHEN dim1 = 'def' THEN dim2 END\n"
+ "FROM druid.foo\n"
+ "WHERE\n"
+ " CASE WHEN dim1 = 'def' THEN dim2 END = 'abc'",
ImmutableList.of(
newScanQueryBuilder()
.dataSource(CalciteTests.DATASOURCE1)
.virtualColumns(
expressionVirtualColumn(
"v0",
"case_searched((\"dim1\" == 'def'),\"dim2\",null)",
ColumnType.STRING
)
)
.intervals(querySegmentSpec(Filtration.eternity()))
.filters(
and(
equality("dim1", "def", ColumnType.STRING),
equality("dim2", "abc", ColumnType.STRING)
)
)
.columns("dim1", "dim2", "v0")
.build()
),
ImmutableList.of(
new Object[]{"def", "abc", "abc"}
)
);
}
@Test @Test
public void testGroupByCaseWhenOfTripleAnd() public void testGroupByCaseWhenOfTripleAnd()
{ {