You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@druid.apache.org by ka...@apache.org on 2023/06/15 07:43:23 UTC
[druid] branch master updated: Limit results generated by SELECT queries in MSQ (#14370)
This is an automated email from the ASF dual-hosted git repository.
karan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new 4935f2470a Limit results generated by SELECT queries in MSQ (#14370)
4935f2470a is described below
commit 4935f2470a9303e228a274949103ac38f4612c45
Author: Laksh Singla <la...@gmail.com>
AuthorDate: Thu Jun 15 13:13:11 2023 +0530
Limit results generated by SELECT queries in MSQ (#14370)
* Limit select results in MSQ
* reduce number of files in test
* add truncated flag
* avoid materializing select results to list, use iterable instead
* javadocs
---
docs/api-reference/sql-ingestion-api.md | 2 +
.../org/apache/druid/msq/exec/ControllerImpl.java | 133 +++++++++++----------
.../java/org/apache/druid/msq/exec/Limits.java | 7 ++
.../msq/indexing/report/MSQResultsReport.java | 36 ++++--
.../org/apache/druid/msq/exec/MSQSelectTest.java | 63 ++++++++++
.../msq/indexing/report/MSQTaskReportTest.java | 13 +-
.../org/apache/druid/msq/test/MSQTestBase.java | 18 +--
.../druid/testing/utils/MsqTestQueryHelper.java | 7 +-
8 files changed, 180 insertions(+), 99 deletions(-)
diff --git a/docs/api-reference/sql-ingestion-api.md b/docs/api-reference/sql-ingestion-api.md
index a9cceb8d4d..43acca9b49 100644
--- a/docs/api-reference/sql-ingestion-api.md
+++ b/docs/api-reference/sql-ingestion-api.md
@@ -203,6 +203,8 @@ Keep the following in mind when using the task API to view reports:
- As an experimental feature, the MSQ task engine supports running SELECT queries. SELECT query results are written into
the `multiStageQuery.payload.results.results` task report key as an array of arrays. The behavior and result format of plain
SELECT queries (without INSERT or REPLACE) is subject to change.
+- `multiStageQuery.payload.results.resultsTruncated` denote whether the results of the report have been truncated to prevent
+the reports from blowing up
For an explanation of the fields in a report, see [Report response fields](#report-response-fields).
diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java
index bf16412b0e..a5ff643e8b 100644
--- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java
+++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java
@@ -206,6 +206,7 @@ import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
+import java.util.Iterator;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
@@ -1395,67 +1396,77 @@ public class ControllerImpl implements Controller
return Yielders.each(
Sequences.concat(
- StreamSupport.stream(queryKernel.getResultPartitionsForStage(finalStageId).spliterator(), false)
- .map(
- readablePartition -> {
- try {
- return new FrameChannelSequence(
- inputChannels.openChannel(
- new StagePartition(
- queryKernel.getStageDefinition(finalStageId).getId(),
- readablePartition.getPartitionNumber()
- )
- )
- );
- }
- catch (IOException e) {
- throw new RuntimeException(e);
- }
+ StreamSupport.stream(queryKernel.getResultPartitionsForStage(finalStageId).spliterator(), false)
+ .map(
+ readablePartition -> {
+ try {
+ return new FrameChannelSequence(
+ inputChannels.openChannel(
+ new StagePartition(
+ queryKernel.getStageDefinition(finalStageId).getId(),
+ readablePartition.getPartitionNumber()
+ )
+ )
+ );
+ }
+ catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ }
+ ).collect(Collectors.toList())
+ ).flatMap(
+ frame -> {
+ final Cursor cursor = FrameProcessors.makeCursor(
+ frame,
+ queryKernel.getStageDefinition(finalStageId).getFrameReader()
+ );
+
+ final ColumnSelectorFactory columnSelectorFactory = cursor.getColumnSelectorFactory();
+ final ColumnMappings columnMappings = task.getQuerySpec().getColumnMappings();
+ @SuppressWarnings("rawtypes")
+ final List<ColumnValueSelector> selectors =
+ columnMappings.getMappings()
+ .stream()
+ .map(
+ mapping ->
+ columnSelectorFactory.makeColumnValueSelector(mapping.getQueryColumn())
+ ).collect(Collectors.toList());
+
+ final List<SqlTypeName> sqlTypeNames = task.getSqlTypeNames();
+ Iterable<Object[]> retVal = () -> new Iterator<Object[]>()
+ {
+ @Override
+ public boolean hasNext()
+ {
+ return !cursor.isDone();
+ }
+
+ @Override
+ public Object[] next()
+ {
+ final Object[] row = new Object[columnMappings.size()];
+ for (int i = 0; i < row.length; i++) {
+ final Object value = selectors.get(i).getObject();
+ if (sqlTypeNames == null || task.getSqlResultsContext() == null) {
+ // SQL type unknown, or no SQL results context: pass-through as is.
+ row[i] = value;
+ } else {
+ row[i] = SqlResults.coerce(
+ context.jsonMapper(),
+ task.getSqlResultsContext(),
+ value,
+ sqlTypeNames.get(i)
+ );
}
- ).collect(Collectors.toList())
- ).flatMap(
- frame -> {
- final Cursor cursor = FrameProcessors.makeCursor(
- frame,
- queryKernel.getStageDefinition(finalStageId).getFrameReader()
- );
-
- final ColumnSelectorFactory columnSelectorFactory = cursor.getColumnSelectorFactory();
- final ColumnMappings columnMappings = task.getQuerySpec().getColumnMappings();
- @SuppressWarnings("rawtypes")
- final List<ColumnValueSelector> selectors =
- columnMappings.getMappings()
- .stream()
- .map(
- mapping ->
- columnSelectorFactory.makeColumnValueSelector(mapping.getQueryColumn())
- ).collect(Collectors.toList());
-
- final List<SqlTypeName> sqlTypeNames = task.getSqlTypeNames();
- final List<Object[]> retVal = new ArrayList<>();
- while (!cursor.isDone()) {
- final Object[] row = new Object[columnMappings.size()];
- for (int i = 0; i < row.length; i++) {
- final Object value = selectors.get(i).getObject();
- if (sqlTypeNames == null || task.getSqlResultsContext() == null) {
- // SQL type unknown, or no SQL results context: pass-through as is.
- row[i] = value;
- } else {
- row[i] = SqlResults.coerce(
- context.jsonMapper(),
- task.getSqlResultsContext(),
- value,
- sqlTypeNames.get(i)
- );
- }
- }
- retVal.add(row);
- cursor.advance();
- }
-
- return Sequences.simple(retVal);
- }
- ).withBaggage(resultReaderExec::shutdownNow)
+ }
+ cursor.advance();
+ return row;
+ }
+ };
+ return Sequences.simple(retVal);
+ }
+ )
+ .withBaggage(resultReaderExec::shutdownNow)
);
} else {
return null;
@@ -2026,7 +2037,7 @@ public class ControllerImpl implements Controller
);
}
- return new MSQResultsReport(mappedSignature.build(), sqlTypeNames, resultsYielder);
+ return new MSQResultsReport(mappedSignature.build(), sqlTypeNames, resultsYielder, null);
}
private static MSQStatusReport makeStatusReport(
diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/Limits.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/Limits.java
index f9a53ff219..c9e598c618 100644
--- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/Limits.java
+++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/Limits.java
@@ -90,4 +90,11 @@ public class Limits
* {@link ClusterStatisticsMergeMode#SEQUENTIAL} mode is chosen.
*/
public static final long MAX_WORKERS_FOR_PARALLEL_MERGE = 100;
+
+ /**
+ * Max number of rows in the query reports of the SELECT queries run by MSQ. This ensures that the reports donot blow
+ * up for queries operating on larger datasets. The full result of the select query should be available once the
+ * MSQ is able to run async queries
+ */
+ public static final long MAX_SELECT_RESULT_ROWS = 3_000;
}
diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/report/MSQResultsReport.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/report/MSQResultsReport.java
index c145ae77d5..56359f8cdf 100644
--- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/report/MSQResultsReport.java
+++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/report/MSQResultsReport.java
@@ -27,9 +27,11 @@ import org.apache.calcite.sql.type.SqlTypeName;
import org.apache.druid.java.util.common.guava.Sequences;
import org.apache.druid.java.util.common.guava.Yielder;
import org.apache.druid.java.util.common.guava.Yielders;
+import org.apache.druid.msq.exec.Limits;
import org.apache.druid.segment.column.ColumnType;
import javax.annotation.Nullable;
+import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
@@ -42,17 +44,30 @@ public class MSQResultsReport
private final List<ColumnAndType> signature;
@Nullable
private final List<SqlTypeName> sqlTypeNames;
- private final Yielder<Object[]> resultYielder;
+ private final List<Object[]> results;
+ private final boolean resultsTruncated;
public MSQResultsReport(
final List<ColumnAndType> signature,
@Nullable final List<SqlTypeName> sqlTypeNames,
- final Yielder<Object[]> resultYielder
+ Yielder<Object[]> resultYielder,
+ @Nullable Boolean resultsTruncated
)
{
this.signature = Preconditions.checkNotNull(signature, "signature");
this.sqlTypeNames = sqlTypeNames;
- this.resultYielder = Preconditions.checkNotNull(resultYielder, "resultYielder");
+ this.results = new ArrayList<>();
+ int rowCount = 0;
+ while (!resultYielder.isDone() && rowCount < Limits.MAX_SELECT_RESULT_ROWS) {
+ results.add(resultYielder.get());
+ resultYielder = resultYielder.next(null);
+ ++rowCount;
+ }
+ if (resultsTruncated != null) {
+ this.resultsTruncated = !resultYielder.isDone() || resultsTruncated;
+ } else {
+ this.resultsTruncated = !resultYielder.isDone();
+ }
}
/**
@@ -62,10 +77,11 @@ public class MSQResultsReport
static MSQResultsReport fromJson(
@JsonProperty("signature") final List<ColumnAndType> signature,
@JsonProperty("sqlTypeNames") @Nullable final List<SqlTypeName> sqlTypeNames,
- @JsonProperty("results") final List<Object[]> results
+ @JsonProperty("results") final List<Object[]> results,
+ @JsonProperty("resultsTruncated") final Boolean resultsTruncated
)
{
- return new MSQResultsReport(signature, sqlTypeNames, Yielders.each(Sequences.simple(results)));
+ return new MSQResultsReport(signature, sqlTypeNames, Yielders.each(Sequences.simple(results)), resultsTruncated);
}
@JsonProperty("signature")
@@ -83,9 +99,15 @@ public class MSQResultsReport
}
@JsonProperty("results")
- public Yielder<Object[]> getResultYielder()
+ public List<Object[]> getResults()
+ {
+ return results;
+ }
+
+ @JsonProperty("resultsTruncted")
+ public boolean isResultsTruncated()
{
- return resultYielder;
+ return resultsTruncated;
}
public static class ColumnAndType
diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java
index 0ac6550ea7..b09e3aa01f 100644
--- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java
@@ -22,13 +22,16 @@ package org.apache.druid.msq.exec;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import org.apache.druid.common.config.NullHandling;
+import org.apache.druid.data.input.impl.CsvInputFormat;
import org.apache.druid.data.input.impl.JsonInputFormat;
import org.apache.druid.data.input.impl.LocalInputSource;
import org.apache.druid.frame.util.DurableStorageUtils;
import org.apache.druid.java.util.common.DateTimes;
import org.apache.druid.java.util.common.ISE;
import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.common.granularity.Granularities;
+import org.apache.druid.math.expr.ExprEval;
import org.apache.druid.math.expr.ExprMacroTable;
import org.apache.druid.msq.indexing.MSQSpec;
import org.apache.druid.msq.indexing.MSQTuningConfig;
@@ -1784,6 +1787,66 @@ public class MSQSelectTest extends MSQTestBase
}
}
+ @Test
+ public void testSelectRowsGetTruncatedInReports() throws IOException
+ {
+ RowSignature dummyRowSignature = RowSignature.builder().add("timestamp", ColumnType.LONG).build();
+
+ final int numFiles = 200;
+
+ final File toRead = MSQTestFileUtils.getResourceAsTemporaryFile(temporaryFolder, this, "/wikipedia-sampled.json");
+ final String toReadFileNameAsJson = queryFramework().queryJsonMapper().writeValueAsString(toRead.getAbsolutePath());
+
+ String externalFiles = String.join(", ", Collections.nCopies(numFiles, toReadFileNameAsJson));
+
+ List<Object[]> result = new ArrayList<>();
+ for (int i = 0; i < Limits.MAX_SELECT_RESULT_ROWS; ++i) {
+ result.add(new Object[]{1});
+ }
+
+ testSelectQuery()
+ .setSql(StringUtils.format(
+ " SELECT 1 as \"timestamp\"\n"
+ + "FROM TABLE(\n"
+ + " EXTERN(\n"
+ + " '{ \"files\": [%s],\"type\":\"local\"}',\n"
+ + " '{\"type\": \"csv\", \"hasHeaderRow\": true}',\n"
+ + " '[{\"name\": \"timestamp\", \"type\": \"string\"}]'\n"
+ + " )\n"
+ + ")",
+ externalFiles
+ ))
+ .setExpectedRowSignature(dummyRowSignature)
+ .setExpectedMSQSpec(
+ MSQSpec
+ .builder()
+ .query(newScanQueryBuilder()
+ .dataSource(new ExternalDataSource(
+ new LocalInputSource(null, null, Collections.nCopies(numFiles, toRead)),
+ new CsvInputFormat(null, null, null, true, 0),
+ RowSignature.builder().add("timestamp", ColumnType.STRING).build()
+ ))
+ .intervals(querySegmentSpec(Filtration.eternity()))
+ .columns("v0")
+ .virtualColumns(new ExpressionVirtualColumn("v0", ExprEval.of(1L).toExpr(), ColumnType.LONG))
+ .context(defaultScanQueryContext(
+ context,
+ RowSignature.builder().add("v0", ColumnType.LONG).build()
+ ))
+ .build()
+ )
+ .columnMappings(new ColumnMappings(
+ ImmutableList.of(
+ new ColumnMapping("v0", "timestamp")
+ )
+ ))
+ .tuningConfig(MSQTuningConfig.defaultConfig())
+ .build())
+ .setQueryContext(context)
+ .setExpectedResultRows(result)
+ .verifyResults();
+ }
+
@Test
public void testMultiValueStringWithIncorrectType() throws IOException
{
diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/report/MSQTaskReportTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/report/MSQTaskReportTest.java
index aeca792024..6c111c7fd8 100644
--- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/report/MSQTaskReportTest.java
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/report/MSQTaskReportTest.java
@@ -31,7 +31,6 @@ import org.apache.druid.indexer.TaskState;
import org.apache.druid.indexing.common.SingleFileTaskReportFileWriter;
import org.apache.druid.indexing.common.TaskReport;
import org.apache.druid.java.util.common.guava.Sequences;
-import org.apache.druid.java.util.common.guava.Yielder;
import org.apache.druid.java.util.common.guava.Yielders;
import org.apache.druid.msq.counters.CounterSnapshotsTree;
import org.apache.druid.msq.guice.MSQIndexingModule;
@@ -51,7 +50,6 @@ import org.junit.rules.TemporaryFolder;
import java.io.File;
import java.util.ArrayDeque;
-import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
@@ -105,7 +103,8 @@ public class MSQTaskReportTest
new MSQResultsReport(
Collections.singletonList(new MSQResultsReport.ColumnAndType("s", ColumnType.STRING)),
ImmutableList.of(SqlTypeName.VARCHAR),
- Yielders.each(Sequences.simple(results))
+ Yielders.each(Sequences.simple(results)),
+ null
)
)
);
@@ -125,13 +124,7 @@ public class MSQTaskReportTest
Assert.assertEquals(report.getPayload().getStatus().getPendingTasks(), report2.getPayload().getStatus().getPendingTasks());
Assert.assertEquals(report.getPayload().getStages(), report2.getPayload().getStages());
- Yielder<Object[]> yielder = report2.getPayload().getResults().getResultYielder();
- final List<Object[]> results2 = new ArrayList<>();
-
- while (!yielder.isDone()) {
- results2.add(yielder.get());
- yielder = yielder.next(null);
- }
+ final List<Object[]> results2 = report2.getPayload().getResults().getResults();
Assert.assertEquals(results.size(), results2.size());
for (int i = 0; i < results.size(); i++) {
diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java
index e44d1974b8..de922451e4 100644
--- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java
@@ -69,7 +69,6 @@ import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.common.concurrent.Execs;
import org.apache.druid.java.util.common.granularity.Granularities;
import org.apache.druid.java.util.common.granularity.Granularity;
-import org.apache.druid.java.util.common.guava.Yielder;
import org.apache.druid.java.util.common.logger.Logger;
import org.apache.druid.math.expr.ExprMacroTable;
import org.apache.druid.metadata.input.InputSourceModule;
@@ -769,20 +768,7 @@ public class MSQTestBase extends BaseCalciteQueryTest
if (resultsReport == null) {
return null;
} else {
- Yielder<Object[]> yielder = resultsReport.getResultYielder();
- List<Object[]> rows = new ArrayList<>();
- while (!yielder.isDone()) {
- rows.add(yielder.get());
- yielder = yielder.next(null);
- }
- try {
- yielder.close();
- }
- catch (IOException e) {
- throw new ISE("Unable to get results from the report");
- }
-
- return rows;
+ return resultsReport.getResults();
}
}
@@ -1317,7 +1303,7 @@ public class MSQTestBase extends BaseCalciteQueryTest
{
Preconditions.checkArgument(expectedResultRows != null, "Result rows cannot be null");
Preconditions.checkArgument(expectedRowSignature != null, "Row signature cannot be null");
- Preconditions.checkArgument(expectedMSQSpec != null, "MultiStageQuery Query spec not ");
+ Preconditions.checkArgument(expectedMSQSpec != null, "MultiStageQuery Query spec cannot be null ");
Pair<MSQSpec, Pair<List<MSQResultsReport.ColumnAndType>, List<Object[]>>> specAndResults = runQueryWithResult();
if (specAndResults == null) { // A fault was expected and the assertion has been done in the runQueryWithResult
diff --git a/integration-tests/src/main/java/org/apache/druid/testing/utils/MsqTestQueryHelper.java b/integration-tests/src/main/java/org/apache/druid/testing/utils/MsqTestQueryHelper.java
index 7525cb9d87..5e1a40df07 100644
--- a/integration-tests/src/main/java/org/apache/druid/testing/utils/MsqTestQueryHelper.java
+++ b/integration-tests/src/main/java/org/apache/druid/testing/utils/MsqTestQueryHelper.java
@@ -31,7 +31,6 @@ import org.apache.druid.java.util.common.IAE;
import org.apache.druid.java.util.common.ISE;
import org.apache.druid.java.util.common.RetryUtils;
import org.apache.druid.java.util.common.StringUtils;
-import org.apache.druid.java.util.common.guava.Yielder;
import org.apache.druid.java.util.http.client.response.StatusResponseHolder;
import org.apache.druid.msq.indexing.report.MSQResultsReport;
import org.apache.druid.msq.indexing.report.MSQTaskReport;
@@ -201,17 +200,15 @@ public class MsqTestQueryHelper extends AbstractTestQueryHelper<MsqQueryWithResu
List<Map<String, Object>> actualResults = new ArrayList<>();
- Yielder<Object[]> yielder = resultsReport.getResultYielder();
+ List<Object[]> results = resultsReport.getResults();
List<MSQResultsReport.ColumnAndType> rowSignature = resultsReport.getSignature();
- while (!yielder.isDone()) {
- Object[] row = yielder.get();
+ for (Object[] row : results) {
Map<String, Object> rowWithFieldNames = new LinkedHashMap<>();
for (int i = 0; i < row.length; ++i) {
rowWithFieldNames.put(rowSignature.get(i).getName(), row[i]);
}
actualResults.add(rowWithFieldNames);
- yielder = yielder.next(null);
}
QueryResultVerifier.ResultVerificationObject resultsComparison = QueryResultVerifier.compareResults(
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org