Skip to content

Commit

Permalink
DrillWindowQueryTest: use proper way to decide if the query is ordered (
Browse files Browse the repository at this point in the history
  • Loading branch information
kgyrtkirk authored and ycp2 committed Nov 17, 2023
1 parent 5ea5c45 commit 18536d9
Show file tree
Hide file tree
Showing 6 changed files with 43 additions and 11 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,7 @@ public void validate()
catch (SqlParseException e1) {
throw translateException(e1);
}
hook.captureSqlNode(root);
handler = createHandler(root);
handler.validate();
plannerContext.setResourceActions(handler.resourceActions());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@
import org.apache.calcite.rel.RelRoot;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.sql.SqlInsert;
import org.apache.calcite.sql.SqlNode;
import org.apache.druid.guice.annotations.UnstableApi;
import org.apache.druid.sql.calcite.rel.DruidRel;

/**
Expand All @@ -32,9 +34,13 @@
* none at the points where tests want to verify, except for the opportunity to
* capture the native query.
*/
@UnstableApi
public interface PlannerHook
{
void captureSql(String sql);
default void captureSqlNode(SqlNode node)
{
}
void captureQueryRel(RelRoot rootQueryRel);
void captureDruidRel(DruidRel<?> druidRel);
void captureBindableRel(BindableRel bindableRel);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1513,6 +1513,7 @@ public void verify(String sql, QueryResults queryResults)
}
catch (AssertionError e) {
log.info("sql: %s", sql);
log.info(resultsToString("Expected", expectedResults));
log.info(resultsToString("Actual", queryResults.results));
throw e;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@
import com.google.common.collect.Iterators;
import com.google.common.io.ByteStreams;
import com.google.inject.Injector;
import org.apache.calcite.sql.SqlNode;
import org.apache.calcite.sql2rel.SqlToRelConverter;
import org.apache.commons.io.FileUtils;
import org.apache.druid.common.config.NullHandling;
import org.apache.druid.data.input.InputRow;
Expand Down Expand Up @@ -52,6 +54,7 @@
import org.apache.druid.sql.calcite.NotYetSupported.Modes;
import org.apache.druid.sql.calcite.NotYetSupported.NotYetSupportedProcessor;
import org.apache.druid.sql.calcite.QueryTestRunner.QueryResults;
import org.apache.druid.sql.calcite.planner.PlannerCaptureHook;
import org.apache.druid.sql.calcite.planner.PlannerContext;
import org.apache.druid.sql.calcite.util.SpecificSegmentsQuerySegmentWalker;
import org.apache.druid.timeline.DataSegment;
Expand Down Expand Up @@ -363,7 +366,8 @@ public void verify(String sql, QueryResults queryResults)
List<Object[]> expectedResults = parseResults(currentRowSignature, expectedResultsText);
try {
Assert.assertEquals(StringUtils.format("result count: %s", sql), expectedResultsText.size(), results.size());
if (!isOrdered(sql)) {
if (!isOrdered(queryResults)) {
// in case the resultset is not ordered; order via the same comparator before comparision
results.sort(new ArrayRowCmp());
expectedResults.sort(new ArrayRowCmp());
}
Expand All @@ -377,12 +381,10 @@ public void verify(String sql, QueryResults queryResults)
}
}

private boolean isOrdered(String sql)
private boolean isOrdered(QueryResults queryResults)
{
// FIXME: SqlToRelConverter.isOrdered(null) would be better
sql = sql.toLowerCase(Locale.ENGLISH).replace('\n', ' ');
sql = sql.substring(sql.lastIndexOf(')'));
return sql.contains("order");
SqlNode sqlNode = ((PlannerCaptureHook) queryResults.capture).getSqlNode();
return SqlToRelConverter.isOrdered(sqlNode);
}
}

Expand Down Expand Up @@ -478,6 +480,7 @@ public void windowQueryTest()
.skipVectorize(true)
.queryContext(ImmutableMap.of(
PlannerContext.CTX_ENABLE_WINDOW_FNS, true,
PlannerCaptureHook.NEED_CAPTURE_HOOK, true,
QueryContexts.ENABLE_DEBUG, true)
)
.sql(testCase.getQueryString())
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -86,14 +86,14 @@ public interface QueryRunStepFactory
*/
public abstract static class QueryRunStep
{
private final QueryTestBuilder builder;
protected final QueryTestBuilder builder;

public QueryRunStep(final QueryTestBuilder builder)
{
this.builder = builder;
}

public QueryTestBuilder builder()
public final QueryTestBuilder builder()
{
return builder;
}
Expand Down Expand Up @@ -207,12 +207,10 @@ public void run()
public abstract static class BaseExecuteQuery extends QueryRunStep
{
protected final List<QueryResults> results = new ArrayList<>();
protected final boolean doCapture;

public BaseExecuteQuery(QueryTestBuilder builder)
{
super(builder);
doCapture = builder.expectedLogicalPlan != null;
}

public List<QueryResults> results()
Expand Down Expand Up @@ -278,7 +276,7 @@ public QueryResults runQuery(
)
{
try {
final PlannerCaptureHook capture = doCapture ? new PlannerCaptureHook() : null;
final PlannerCaptureHook capture = getCaptureHook();
final DirectStatement stmt = sqlStatementFactory.directStatement(query);
stmt.setHook(capture);
final Sequence<Object[]> results = stmt.execute().getResults();
Expand All @@ -300,6 +298,14 @@ public QueryResults runQuery(
}
}

private PlannerCaptureHook getCaptureHook()
{
if (builder.getQueryContext().containsKey(PlannerCaptureHook.NEED_CAPTURE_HOOK) || builder.expectedLogicalPlan != null) {
return new PlannerCaptureHook();
}
return null;
}

public static Pair<RowSignature, List<Object[]>> getResults(
final SqlStatementFactory sqlStatementFactory,
final SqlQueryPlus query
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,19 +23,29 @@
import org.apache.calcite.rel.RelRoot;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.sql.SqlInsert;
import org.apache.calcite.sql.SqlNode;
import org.apache.druid.sql.calcite.rel.DruidRel;

public class PlannerCaptureHook implements PlannerHook
{
public static final String NEED_CAPTURE_HOOK = "need_capture_hook";

private RelRoot relRoot;
private SqlInsert insertNode;
private SqlNode sqlNode;

@Override
public void captureSql(String sql)
{
// Not used at present. Add a field to capture this if you need it.
}

@Override
public void captureSqlNode(SqlNode node)
{
this.sqlNode = node;
}

@Override
public void captureQueryRel(RelRoot rootQueryRel)
{
Expand Down Expand Up @@ -75,4 +85,9 @@ public SqlInsert insertNode()
{
return insertNode;
}

public SqlNode getSqlNode()
{
return sqlNode;
}
}

0 comments on commit 18536d9

Please sign in to comment.