coalesce on unnest row mismatch fix (#15019)

* coalesce on unnest row mismatch fix

* new example with coalesce over unnest with nested array columns

* New example with change in order which triggers the nvl

* new test plan update for useDefault=true
This commit is contained in:
Soumyava 2023-10-02 17:26:50 -07:00 committed by GitHub
parent f1edd671fb
commit 261f54dc04
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 93 additions and 0 deletions

View File

@ -174,6 +174,7 @@ public class DruidCorrelateUnnestRule extends RelOptRule
)
);
relBuilder.convert(correlate.getRowType(), false);
final RelNode build = relBuilder.build();
call.transformTo(build);
}

View File

@ -6308,4 +6308,96 @@ public class CalciteNestedDataQueryTest extends BaseCalciteQueryTest
);
}
@Test
public void testCoalesceOnNestedColumns()
{
// jo.unnest is first entry in coalesce
// so Calcite removes the coalesce to be used here
testQuery(
"select coalesce(c,long) as col "
+ " from druid.all_auto, unnest(json_value(arrayNestedLong, '$[1]' returning bigint array)) as u(c) ",
ImmutableList.of(
Druids.newScanQueryBuilder()
.dataSource(UnnestDataSource.create(
new TableDataSource(DATA_SOURCE_ALL),
new NestedFieldVirtualColumn("arrayNestedLong", "$[1]", "j0.unnest", ColumnType.LONG_ARRAY),
null
))
.intervals(querySegmentSpec(Filtration.eternity()))
.columns("j0.unnest")
.resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST)
.legacy(false)
.context(QUERY_CONTEXT_DEFAULT)
.build()
),
ImmutableList.of(
new Object[]{null},
new Object[]{3L},
new Object[]{4L},
new Object[]{3L},
new Object[]{4L},
new Object[]{1L},
new Object[]{2L},
new Object[]{null}
),
RowSignature.builder()
.add("col", ColumnType.LONG)
.build()
);
}
@Test
public void testCoalesceOnNestedColumnsLater()
{
// the first column in coalesce comes from the table
// so a virtual expression is present for the coalesce
testQuery(
"select coalesce(long,c) as col "
+ " from druid.all_auto, unnest(json_value(arrayNestedLong, '$[1]' returning bigint array)) as u(c) ",
useDefault ?
ImmutableList.of(
Druids.newScanQueryBuilder()
.dataSource(UnnestDataSource.create(
new TableDataSource(DATA_SOURCE_ALL),
new NestedFieldVirtualColumn("arrayNestedLong", "$[1]", "j0.unnest", ColumnType.LONG_ARRAY),
null
))
.intervals(querySegmentSpec(Filtration.eternity()))
.columns("long")
.resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST)
.legacy(false)
.context(QUERY_CONTEXT_DEFAULT)
.build()
) :
ImmutableList.of(
Druids.newScanQueryBuilder()
.dataSource(UnnestDataSource.create(
new TableDataSource(DATA_SOURCE_ALL),
new NestedFieldVirtualColumn("arrayNestedLong", "$[1]", "j0.unnest", ColumnType.LONG_ARRAY),
null
))
.intervals(querySegmentSpec(Filtration.eternity()))
.virtualColumns(expressionVirtualColumn("v0", "nvl(\"long\",\"j0.unnest\")", ColumnType.LONG))
.columns("v0")
.resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST)
.legacy(false)
.context(QUERY_CONTEXT_DEFAULT)
.build()
),
ImmutableList.of(
new Object[]{2L},
new Object[]{1L},
new Object[]{1L},
new Object[]{4L},
new Object[]{4L},
new Object[]{5L},
new Object[]{5L},
new Object[]{5L}
),
RowSignature.builder()
.add("col", ColumnType.LONG)
.build()
);
}
}