Skip to content

Add page information to SqlStatementResource API #14512

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cryptoe merged 5 commits into from
Jul 3, 2023
Merged

Add page information to SqlStatementResource API #14512

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
f94fe1e
Add page information to druid statement resource
adarshsanjeev Jun 30, 2023
5580e9a
Fix tests
adarshsanjeev Jul 1, 2023
13e889c
Fix builds
adarshsanjeev Jul 1, 2023
349cd47
Address review comments
adarshsanjeev Jul 3, 2023
d04ffdd
Fix builds
adarshsanjeev Jul 3, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
103 changes: 103 additions & 0 deletions ...core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/PageInformation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/

package org.apache.druid.msq.sql.entity;


import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;

import javax.annotation.Nullable;
import java.util.Objects;

/**
* Contains information about a single page in the results.
*/
public class PageInformation
{
@Nullable
private final Long numRows;
@Nullable
private final Long sizeInBytes;
private final long id;

@JsonCreator
public PageInformation(
@JsonProperty("numRows") @Nullable Long numRows,
@JsonProperty("sizeInBytes") @Nullable Long sizeInBytes,
@JsonProperty("id") long id
)
{
this.numRows = numRows;
this.sizeInBytes = sizeInBytes;
this.id = id;
}

@JsonProperty
@Nullable
public Long getNumRows()
{
return numRows;
}

@JsonProperty
@Nullable
public Long getSizeInBytes()
{
return sizeInBytes;
}

@JsonProperty
public long getId()
{
return id;
}

@Override
public boolean equals(Object o)
{
if (this == o) {
return true;
}
if (o == null || getClass() != o.getClass()) {
return false;
}
PageInformation that = (PageInformation) o;
return id == that.id && Objects.equals(numRows, that.numRows) && Objects.equals(
sizeInBytes,
that.sizeInBytes
);
}

@Override
public int hashCode()
{
return Objects.hash(numRows, sizeInBytes, id);
}

@Override
public String toString()
{
return "PageInformation{" +
"numRows=" + numRows +
", sizeInBytes=" + sizeInBytes +
", id=" + id +
'}';
}
}
64 changes: 36 additions & 28 deletions ...multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/ResultSetInformation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,50 +32,50 @@ public class ResultSetInformation
{

@Nullable
private final Long numRows;
private final Long numTotalRows;
@Nullable
private final Long sizeInBytes;

private final Long totalSizeInBytes;
@Nullable
private final ResultFormat resultFormat;

@Nullable
private final List<Object> records;

@Nullable
private final String dataSource;
@Nullable
private final List<PageInformation> pages;

@JsonCreator
public ResultSetInformation(
@JsonProperty("numTotalRows") @Nullable Long numTotalRows,
@JsonProperty("totalSizeInBytes") @Nullable Long totalSizeInBytes,
@JsonProperty("resultFormat") @Nullable ResultFormat resultFormat,
@JsonProperty("numRows") @Nullable Long numRows,
Comment thread
adarshsanjeev marked this conversation as resolved.
@JsonProperty("sizeInBytes") @Nullable Long sizeInBytes,
@JsonProperty("dataSource") @Nullable String dataSource,
@JsonProperty("sampleRecords") @Nullable
List<Object> records
@JsonProperty("sampleRecords") @Nullable List<Object> records,
@JsonProperty("pages") @Nullable List<PageInformation> pages
)
{
this.numRows = numRows;
this.sizeInBytes = sizeInBytes;
this.numTotalRows = numTotalRows;
this.totalSizeInBytes = totalSizeInBytes;
this.resultFormat = resultFormat;
this.dataSource = dataSource;
this.records = records;
this.pages = pages;
}

@Nullable
@JsonProperty
@Nullable
@JsonInclude(JsonInclude.Include.NON_NULL)
public Long getNumRows()
public Long getNumTotalRows()
{
return numRows;
return numTotalRows;
}

@Nullable
@JsonProperty
@Nullable
@JsonInclude(JsonInclude.Include.NON_NULL)
public Long getSizeInBytes()
public Long getTotalSizeInBytes()
{
return sizeInBytes;
return totalSizeInBytes;
}

@JsonProperty
Expand All @@ -94,14 +94,21 @@ public String getDataSource()
return dataSource;
}

@Nullable
@JsonProperty("sampleRecords")
@Nullable
@JsonInclude(JsonInclude.Include.NON_NULL)
public List<Object> getRecords()
{
return records;
}

@JsonProperty
@Nullable
@JsonInclude(JsonInclude.Include.NON_NULL)
public List<PageInformation> getPages()
{
return pages;
}

@Override
public boolean equals(Object o)
Expand All @@ -113,30 +120,31 @@ public boolean equals(Object o)
return false;
}
ResultSetInformation that = (ResultSetInformation) o;
return Objects.equals(numRows, that.numRows)
&& Objects.equals(sizeInBytes, that.sizeInBytes)
&& resultFormat == that.resultFormat
&& Objects.equals(records, that.records)
&& Objects.equals(dataSource, that.dataSource);
return Objects.equals(numTotalRows, that.numTotalRows) && Objects.equals(
totalSizeInBytes,
that.totalSizeInBytes
) && resultFormat == that.resultFormat && Objects.equals(records, that.records) && Objects.equals(
dataSource,
that.dataSource
) && Objects.equals(pages, that.pages);
}

@Override
public int hashCode()
{
return Objects.hash(numRows, sizeInBytes, resultFormat, records, dataSource);
return Objects.hash(numTotalRows, totalSizeInBytes, resultFormat, records, dataSource, pages);
}

@Override
public String toString()
{
return "ResultSetInformation{" +
"totalRows=" + numRows +
", totalSize=" + sizeInBytes +
"numTotalRows=" + numTotalRows +
", totalSizeInBytes=" + totalSizeInBytes +
", resultFormat=" + resultFormat +
", records=" + records +
", dataSource='" + dataSource + '\'' +
", pages=" + pages +
'}';
}

}

52 changes: 30 additions & 22 deletions ...ti-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@


import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.io.CountingOutputStream;
import com.google.common.util.concurrent.ListenableFuture;
Expand All @@ -46,6 +47,7 @@
import org.apache.druid.msq.sql.MSQTaskSqlEngine;
import org.apache.druid.msq.sql.SqlStatementState;
import org.apache.druid.msq.sql.entity.ColumnNameAndTypes;
import org.apache.druid.msq.sql.entity.PageInformation;
import org.apache.druid.msq.sql.entity.ResultSetInformation;
import org.apache.druid.msq.sql.entity.SqlStatementResult;
import org.apache.druid.msq.util.SqlStatementResourceHelper;
Expand Down Expand Up @@ -268,8 +270,7 @@ public Response doGetStatus(
@Produces(MediaType.APPLICATION_JSON)
public Response doGetResults(
@PathParam("id") final String queryId,
@QueryParam("offset") Long offset,
@QueryParam("numRows") Long numberOfRows,
@QueryParam("page") Long page,
@Context final HttpServletRequest req
)
{
Expand All @@ -284,28 +285,16 @@ public Response doGetResults(
}
final AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(req);

if (offset != null && offset < 0) {
if (page != null && page < 0) {
return buildNonOkResponse(
DruidException.forPersona(DruidException.Persona.USER)
.ofCategory(DruidException.Category.INVALID_INPUT)
.build(
"offset cannot be negative. Please pass a positive number."
)
);
}
if (numberOfRows != null && numberOfRows < 0) {
return buildNonOkResponse(
DruidException.forPersona(DruidException.Persona.USER)
.ofCategory(DruidException.Category.INVALID_INPUT)
.build(
"numRows cannot be negative. Please pass a positive number."
"Page cannot be negative. Please pass a positive number."
)
);
}

Copy link
Copy Markdown
Contributor

@cryptoe cryptoe Jul 2, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the page does not exist its also an error.
For eg: pageId=999 is passed but we only have 2 pages in the result, its also a invalid input.

Copy link
Copy Markdown
Contributor Author

@adarshsanjeev adarshsanjeev Jul 3, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have added this check below as "if page > 0" for the report since we don't really have a list of pages to check the index for yet and manually creating a list of a single page to check this doesn't make sense

final long start = offset == null ? 0 : offset;
final long last = SqlStatementResourceHelper.getLastIndex(numberOfRows, start);

TaskStatusResponse taskResponse = contactOverlord(overlordClient.taskStatus(queryId));
if (taskResponse == null) {
return Response.status(Response.Status.NOT_FOUND).build();
Expand Down Expand Up @@ -343,8 +332,21 @@ public Response doGetResults(
if (!signature.isPresent()) {
return Response.ok().build();
}
Optional<List<Object>> results = SqlStatementResourceHelper.getResults(SqlStatementResourceHelper.getPayload(
contactOverlord(overlordClient.taskReportAsMap(queryId))));

if (page != null && page > 0) {
// Results from task report are only present as one page.
return buildNonOkResponse(
DruidException.forPersona(DruidException.Persona.USER)
.ofCategory(DruidException.Category.INVALID_INPUT)
.build("Page number is out of range of the results.")
);
}

Optional<List<Object>> results = SqlStatementResourceHelper.getResults(
SqlStatementResourceHelper.getPayload(
contactOverlord(overlordClient.taskReportAsMap(queryId))
)
);

return Response.ok((StreamingOutput) outputStream -> {
CountingOutputStream os = new CountingOutputStream(outputStream);
Expand All @@ -353,7 +355,7 @@ public Response doGetResults(
List<ColumnNameAndTypes> rowSignature = signature.get();
writer.writeResponseStart();

for (long k = start; k < Math.min(last, results.get().size()); k++) {
for (long k = 0; k < results.get().size(); k++) {
writer.writeRowStart();
for (int i = 0; i < rowSignature.size(); i++) {
writer.writeRowField(
Expand Down Expand Up @@ -575,13 +577,19 @@ private Optional<ResultSetInformation> getSampleResults(
isSelectQuery
);
return Optional.of(new ResultSetInformation(
null,
// since the rows can be sampled, get the number of rows from counters
rowsAndSize.orElse(new Pair<>(null, null)).lhs,
rowsAndSize.orElse(new Pair<>(null, null)).rhs,
null,
dataSource,
// only populate sample results in case a select query is successful
isSelectQuery ? SqlStatementResourceHelper.getResults(payload).orElse(null) : null
isSelectQuery ? SqlStatementResourceHelper.getResults(payload).orElse(null) : null,
ImmutableList.of(
new PageInformation(
rowsAndSize.orElse(new Pair<>(null, null)).lhs,
Copy link
Copy Markdown
Contributor

@cryptoe cryptoe Jul 2, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should be populate from the counters.
We need to rework

SqlStatementResourceHelper.getRowsAndSizeFromPayload(
          payload,
          isSelectQuery
      );

so its returns pageInformation.
The next set of changes that I am making then would not touch this piece of code.

Copy link
Copy Markdown
Contributor Author

@adarshsanjeev adarshsanjeev Jul 3, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This might need to be handled at least partially in the next change, as counters don't really mean anything for results if the destination is the task report.

rowsAndSize.orElse(new Pair<>(null, null)).rhs,
0
)
)
));
} else {
return Optional.empty();
Expand Down
6 changes: 4 additions & 2 deletions ...i-stage-query/src/test/java/org/apache/druid/msq/sql/SqlMsqStatementResourcePostTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import org.apache.druid.msq.indexing.error.InsertCannotBeEmptyFault;
import org.apache.druid.msq.indexing.error.MSQException;
import org.apache.druid.msq.sql.entity.ColumnNameAndTypes;
import org.apache.druid.msq.sql.entity.PageInformation;
import org.apache.druid.msq.sql.entity.ResultSetInformation;
import org.apache.druid.msq.sql.entity.SqlStatementResult;
import org.apache.druid.msq.sql.resources.SqlStatementResource;
Expand Down Expand Up @@ -112,17 +113,18 @@ public void testMSQSelectQueryTest() throws IOException
),
MSQTestOverlordServiceClient.DURATION,
new ResultSetInformation(
null,
6L,
316L,
null,
MSQControllerTask.DUMMY_DATASOURCE_FOR_SELECT,
objectMapper.readValue(
objectMapper.writeValueAsString(
results),
new TypeReference<List<Object>>()
{
}
)
),
ImmutableList.of(new PageInformation(6L, 316L, 0))
),
null
);
Expand Down
Loading