Skip to content

QueryResults cleanup #25070

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
wendigo merged 13 commits into from
Feb 19, 2025
Merged

QueryResults cleanup #25070

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
9409e8a
Remove unused method
wendigo Feb 18, 2025
74223f1
Remove redundant this.
wendigo Feb 18, 2025
956e7a8
Rename variable to match meeaning
wendigo Feb 18, 2025
9834db8
Use OptionalLong for updated count
wendigo Feb 18, 2025
426f21c
Remove redundant Session from QueryResultRows
wendigo Feb 18, 2025
c8fccdc
Use field instead of calling method
wendigo Feb 18, 2025
137598b
Remove unused method
wendigo Feb 18, 2025
e8f66cb
Improve formatting by wrapping a new line
wendigo Feb 18, 2025
5d93725
Move null checks to a single place
wendigo Feb 18, 2025
a921696
Remove Nullable annotation
wendigo Feb 18, 2025
ef109dc
Move Optional check to a method
wendigo Feb 18, 2025
757b982
Move Optional null mapping to the callee
wendigo Feb 18, 2025
722aa96
Migrate getUpdateCount to OptionalLong
wendigo Feb 18, 2025
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
4 changes: 2 additions & 2 deletions client/trino-cli/src/main/java/io/trino/cli/Query.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -242,8 +242,8 @@ private void processInitialStatusUpdates(WarningsPrinter warningsPrinter)
private void renderUpdate(Terminal terminal, PrintStream out, QueryStatusInfo results, OutputFormat outputFormat, Optional<String> pager)
{
String status = results.getUpdateType();
if (results.getUpdateCount() != null) {
long count = results.getUpdateCount();
if (results.getUpdateCount().isPresent()) {
long count = results.getUpdateCount().getAsLong();
status += format(": %s row%s", count, (count != 1) ? "s" : "");
out.println(status);
}
Expand Down
3 changes: 2 additions & 1 deletion client/trino-cli/src/test/java/io/trino/cli/TestQueryRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
import java.util.Locale;
import java.util.Optional;
import java.util.OptionalDouble;
import java.util.OptionalLong;
import java.util.Properties;

import static com.google.common.io.ByteStreams.nullOutputStream;
Expand Down Expand Up @@ -146,7 +147,7 @@ static String createResults(MockWebServer server)
null,
ImmutableList.of(),
null,
null);
OptionalLong.empty());
return QUERY_RESULTS_CODEC.toJson(queryResults);
}

Expand Down
9 changes: 5 additions & 4 deletions client/trino-client/src/main/java/io/trino/client/QueryResults.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@

import java.net.URI;
import java.util.List;
import java.util.OptionalLong;

import static com.google.common.base.MoreObjects.firstNonNull;
import static com.google.common.base.MoreObjects.toStringHelper;
Expand All @@ -43,7 +44,7 @@ public class QueryResults
private final QueryError error;
private final List<Warning> warnings;
private final String updateType;
private final Long updateCount;
private final OptionalLong updateCount;

@JsonCreator
public QueryResults(
Expand All @@ -57,7 +58,7 @@ public QueryResults(
@JsonProperty("error") QueryError error,
@JsonProperty("warnings") List<Warning> warnings,
@JsonProperty("updateType") String updateType,
@JsonProperty("updateCount") Long updateCount)
@JsonProperty("updateCount") OptionalLong updateCount)
{
this.id = requireNonNull(id, "id is null");
this.infoUri = requireNonNull(infoUri, "infoUri is null");
Expand All @@ -70,7 +71,7 @@ public QueryResults(
this.error = error;
this.warnings = ImmutableList.copyOf(firstNonNull(warnings, ImmutableList.of()));
this.updateType = updateType;
this.updateCount = updateCount;
this.updateCount = requireNonNull(updateCount, "updateCount is null");
}

@JsonProperty
Expand Down Expand Up @@ -159,7 +160,7 @@ public String getUpdateType()
@Nullable
@JsonProperty
@Override
public Long getUpdateCount()
public OptionalLong getUpdateCount()
Copy link
Copy Markdown
Member

@findepi findepi Jan 23, 2026

Choose a reason for hiding this comment

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

can we remove @Nullable now?

Copy link
Copy Markdown
Contributor Author

@wendigo wendigo Jan 23, 2026

Choose a reason for hiding this comment

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

We can

Copy link
Copy Markdown
Contributor Author

@wendigo wendigo Jan 23, 2026

Choose a reason for hiding this comment

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

#27992

{
return updateCount;
}
Expand Down
3 changes: 2 additions & 1 deletion client/trino-client/src/main/java/io/trino/client/QueryStatusInfo.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

import java.net.URI;
import java.util.List;
import java.util.OptionalLong;

public interface QueryStatusInfo
{
Expand All @@ -36,5 +37,5 @@ public interface QueryStatusInfo

String getUpdateType();

Long getUpdateCount();
OptionalLong getUpdateCount();
}
3 changes: 2 additions & 1 deletion client/trino-client/src/test/java/io/trino/client/TestClientRedirect.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import java.time.ZoneId;
import java.util.Optional;
import java.util.OptionalDouble;
import java.util.OptionalLong;
import java.util.stream.IntStream;
import java.util.stream.Stream;

Expand Down Expand Up @@ -172,7 +173,7 @@ private String newQueryResults(MockWebServer server)
null,
ImmutableList.of(),
null,
null);
OptionalLong.empty());

return QUERY_RESULTS_CODEC.toJson(queryResults);
}
Expand Down
3 changes: 2 additions & 1 deletion client/trino-client/src/test/java/io/trino/client/TestResultRowsDecoder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
import java.util.Iterator;
import java.util.List;
import java.util.OptionalDouble;
import java.util.OptionalLong;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.IntStream;

Expand Down Expand Up @@ -274,7 +275,7 @@ private static QueryResults fromQueryData(QueryData queryData)
null,
ImmutableList.of(),
null,
0L);
OptionalLong.of(0L));
}

private static QueryResults fromSegments(Segment... segments)
Expand Down
3 changes: 2 additions & 1 deletion client/trino-client/src/test/java/io/trino/client/TestRetry.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import java.time.ZoneId;
import java.util.Optional;
import java.util.OptionalDouble;
import java.util.OptionalLong;
import java.util.stream.IntStream;
import java.util.stream.Stream;

Expand Down Expand Up @@ -147,7 +148,7 @@ private String newQueryResults(String state)
null,
ImmutableList.of(),
null,
null);
OptionalLong.empty());

return QUERY_RESULTS_CODEC.toJson(queryResults);
}
Expand Down
3 changes: 1 addition & 2 deletions client/trino-jdbc/src/main/java/io/trino/jdbc/TrinoStatement.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -274,8 +274,7 @@ final boolean internalExecute(String sql)

connection().updateSession(client);

Long updateCount = client.finalStatusInfo().getUpdateCount();
currentUpdateCount.set((updateCount != null) ? updateCount : 0);
currentUpdateCount.set(client.finalStatusInfo().getUpdateCount().orElse(0));
currentUpdateType.set(client.finalStatusInfo().getUpdateType());
warningsManager.addWarnings(client.finalStatusInfo().getWarnings());
return false;
Expand Down
3 changes: 2 additions & 1 deletion client/trino-jdbc/src/test/java/io/trino/jdbc/TestAsyncResultIterator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
import java.util.Map;
import java.util.Optional;
import java.util.OptionalDouble;
import java.util.OptionalLong;
import java.util.Set;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
Expand Down Expand Up @@ -365,7 +366,7 @@ public String getUpdateType()
}

@Override
public Long getUpdateCount()
public OptionalLong getUpdateCount()
{
throw new UnsupportedOperationException();
}
Expand Down
3 changes: 2 additions & 1 deletion client/trino-jdbc/src/test/java/io/trino/jdbc/TestProgressMonitor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@
import java.sql.Statement;
import java.util.List;
import java.util.OptionalDouble;
import java.util.OptionalLong;
import java.util.function.Consumer;

import static com.google.common.base.Preconditions.checkState;
Expand Down Expand Up @@ -103,7 +104,7 @@ private String newQueryResults(Integer partialCancelId, Integer nextUriId, List<
null,
ImmutableList.of(),
null,
null);
OptionalLong.empty());

return QUERY_RESULTS_CODEC.toJson(queryResults);
}
Expand Down
3 changes: 2 additions & 1 deletion core/trino-main/src/main/java/io/trino/dispatcher/QueuedStatementResource.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@
import java.net.URI;
import java.util.Optional;
import java.util.OptionalDouble;
import java.util.OptionalLong;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.Executor;
Expand Down Expand Up @@ -302,7 +303,7 @@ private static QueryResults createQueryResults(
queryError.orElse(null),
ImmutableList.of(),
null,
null);
OptionalLong.empty());
}

private static final class Query
Expand Down
4 changes: 1 addition & 3 deletions core/trino-main/src/main/java/io/trino/server/protocol/JsonBytesQueryDataProducer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,9 +39,7 @@ public QueryData produce(ExternalUriInfo uriInfo, Session session, QueryResultRo
return null;
}

List<OutputColumn> columns = rows.getOutputColumns()
.orElseThrow(() -> new IllegalStateException("Data present without columns"));

List<OutputColumn> columns = rows.getOutputColumns();
if (typeEncoders == null) {
typeEncoders = createTypeEncoders(session, columns);
sourcePageChannels = requireNonNull(columns, "columns is null").stream()
Expand Down
16 changes: 8 additions & 8 deletions core/trino-main/src/main/java/io/trino/server/protocol/Query.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,7 @@
import static io.trino.server.protocol.ProtocolUtil.createColumn;
import static io.trino.server.protocol.ProtocolUtil.toStatementStats;
import static io.trino.server.protocol.QueryInfoUrlFactory.getQueryInfoUri;
import static io.trino.server.protocol.QueryResultRows.empty;
import static io.trino.server.protocol.QueryResultRows.queryResultRowsBuilder;
import static io.trino.server.protocol.Slug.Context.EXECUTING_QUERY;
import static io.trino.spi.StandardErrorCode.SERIALIZATION_ERROR;
Expand Down Expand Up @@ -184,7 +185,7 @@ class Query
private Optional<Throwable> typeSerializationException = Optional.empty();

@GuardedBy("this")
private Long updateCount;
private OptionalLong updateCount = OptionalLong.empty();

public static Query create(
Session session,
Expand Down Expand Up @@ -442,13 +443,12 @@ private synchronized QueryResultsResponse getNextResult(long token, ExternalUriI
resultRows = removePagesFromExchange(queryInfo, targetResultSize.toBytes());
}
else {
resultRows = queryResultRowsBuilder(session).build();
resultRows = empty();
}

if ((queryInfo.updateType() != null) && (updateCount == null)) {
if ((queryInfo.updateType() != null) && updateCount.isEmpty()) {
// grab the update count for non-queries
Optional<Long> updatedRowsCount = resultRows.getUpdateCount();
updateCount = updatedRowsCount.orElse(null);
updateCount = resultRows.getUpdateCount();
}

if (isStarted && (queryInfo.outputStage().isEmpty() || exchangeDataSource.isFinished())) {
Expand Down Expand Up @@ -514,7 +514,7 @@ private synchronized QueryResultsResponse getNextResult(long token, ExternalUriI
getQueryInfoUri(queryInfoUrl, queryId, externalUriInfo),
partialCancelUri,
nextResultsUri,
resultRows.getColumns().orElse(null),
resultRows.getOptionalColumns(),
queryDataProducer.produce(externalUriInfo, session, resultRows, this::handleSerializationException),
toStatementStats(queryInfo),
toQueryError(queryInfo, typeSerializationException),
Expand Down Expand Up @@ -551,7 +551,7 @@ private synchronized QueryResultsResponse toResultsResponse(QueryResults queryRe
private synchronized QueryResultRows removePagesFromExchange(ResultQueryInfo queryInfo, long targetResultBytes)
{
if (!resultsConsumed && queryInfo.outputStage().isEmpty()) {
return queryResultRowsBuilder(session)
return queryResultRowsBuilder()
.withColumnsAndTypes(ImmutableList.of(), ImmutableList.of())
.build();
}
Expand All @@ -560,7 +560,7 @@ private synchronized QueryResultRows removePagesFromExchange(ResultQueryInfo que
// client while holding the lock because the query may transition to the finished state when the
// last page is removed. If another thread observes this state before the response is cached
// the pages will be lost.
QueryResultRows.Builder resultBuilder = queryResultRowsBuilder(session)
QueryResultRows.Builder resultBuilder = queryResultRowsBuilder()
.withColumnsAndTypes(columns, types);

try {
Expand Down
Loading