Skip to content
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.

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

perf: calculate checksum using protobuf values #2848

Merged
merged 6 commits into from
Feb 9, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
import com.google.spanner.v1.ResultSetStats;

/** Forwarding implementation of ResultSet that forwards all calls to a delegate. */
public class ForwardingResultSet extends ForwardingStructReader implements ResultSet {
public class ForwardingResultSet extends ForwardingStructReader implements ProtobufResultSet {

private Supplier<ResultSet> delegate;

Expand Down Expand Up @@ -55,6 +55,22 @@ public boolean next() throws SpannerException {
return delegate.get().next();
}

@Override
public boolean canGetProtobufValue(int columnIndex) {
ResultSet resultSetDelegate = delegate.get();
return (resultSetDelegate instanceof ProtobufResultSet)
&& ((ProtobufResultSet) resultSetDelegate).canGetProtobufValue(columnIndex);
}

@Override
public com.google.protobuf.Value getProtobufValue(int columnIndex) {
ResultSet resultSetDelegate = delegate.get();
Preconditions.checkState(
resultSetDelegate instanceof ProtobufResultSet,
"The result set does not support protobuf values");
return ((ProtobufResultSet) resultSetDelegate).getProtobufValue(columnIndex);
}

@Override
public Struct getCurrentRowAsStruct() {
return delegate.get().getCurrentRowAsStruct();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import static com.google.common.base.Preconditions.checkState;

import com.google.common.annotations.VisibleForTesting;
import com.google.protobuf.Value;
import com.google.spanner.v1.PartialResultSet;
import com.google.spanner.v1.ResultSetMetadata;
import com.google.spanner.v1.ResultSetStats;
Expand All @@ -28,7 +29,7 @@
import javax.annotation.Nullable;

@VisibleForTesting
class GrpcResultSet extends AbstractResultSet<List<Object>> {
class GrpcResultSet extends AbstractResultSet<List<Object>> implements ProtobufResultSet {
private final GrpcValueIterator iterator;
private final Listener listener;
private final DecodeMode decodeMode;
Expand All @@ -49,6 +50,18 @@ class GrpcResultSet extends AbstractResultSet<List<Object>> {
this.decodeMode = decodeMode;
}

@Override
public boolean canGetProtobufValue(int columnIndex) {
return !closed && currRow != null && currRow.canGetProtoValue(columnIndex);
}

@Override
public Value getProtobufValue(int columnIndex) {
checkState(!closed, "ResultSet is closed");
checkState(currRow != null, "next() call required");
return currRow.getProtoValueInternal(columnIndex);
ankiaga marked this conversation as resolved.
Show resolved Hide resolved
}

@Override
protected GrpcStruct currRow() {
checkState(!closed, "ResultSet is closed");
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
import com.google.cloud.spanner.AbstractResultSet.Float64Array;
import com.google.cloud.spanner.AbstractResultSet.Int64Array;
import com.google.cloud.spanner.AbstractResultSet.LazyByteArray;
import com.google.cloud.spanner.Type.Code;
import com.google.cloud.spanner.Type.StructField;
import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;
Expand Down Expand Up @@ -459,10 +460,37 @@ protected Date getDateInternal(int columnIndex) {
return (Date) rowData.get(columnIndex);
}

private boolean isUnrecognizedType(int columnIndex) {
return type.getStructFields().get(columnIndex).getType().getCode() == Code.UNRECOGNIZED;
}

boolean canGetProtoValue(int columnIndex) {
return isUnrecognizedType(columnIndex)
|| (decodeMode == DecodeMode.LAZY_PER_ROW && !rowDecoded)
|| (decodeMode == DecodeMode.LAZY_PER_COL && !colDecoded.get(columnIndex));
}

protected com.google.protobuf.Value getProtoValueInternal(int columnIndex) {
checkProtoValueSupported(columnIndex);
return (com.google.protobuf.Value) rowData.get(columnIndex);
}

private void checkProtoValueSupported(int columnIndex) {
// Unrecognized types are returned as protobuf values.
if (isUnrecognizedType(columnIndex)) {
return;
}
Preconditions.checkState(
decodeMode != DecodeMode.DIRECT,
"Getting proto value is not supported when DecodeMode#DIRECT is used.");
Preconditions.checkState(
!(decodeMode == DecodeMode.LAZY_PER_ROW && rowDecoded),
"Getting proto value after the row has been decoded is not supported.");
Preconditions.checkState(
!(decodeMode == DecodeMode.LAZY_PER_COL && colDecoded.get(columnIndex)),
"Getting proto value after the column has been decoded is not supported.");
}

private void ensureDecoded(int columnIndex) {
if (decodeMode == DecodeMode.LAZY_PER_ROW && !rowDecoded) {
for (int i = 0; i < rowData.size(); i++) {
Expand Down Expand Up @@ -496,7 +524,8 @@ protected Value getValueInternal(int columnIndex) {
case INT64:
return Value.int64(isNull ? null : getLongInternal(columnIndex));
case ENUM:
return Value.protoEnum(getLongInternal(columnIndex), columnType.getProtoTypeFqn());
return Value.protoEnum(
isNull ? null : getLongInternal(columnIndex), columnType.getProtoTypeFqn());
case NUMERIC:
return Value.numeric(isNull ? null : getBigDecimalInternal(columnIndex));
case PG_NUMERIC:
Expand All @@ -512,7 +541,8 @@ protected Value getValueInternal(int columnIndex) {
case BYTES:
return Value.internalBytes(isNull ? null : getLazyBytesInternal(columnIndex));
case PROTO:
return Value.protoMessage(getBytesInternal(columnIndex), columnType.getProtoTypeFqn());
return Value.protoMessage(
isNull ? null : getBytesInternal(columnIndex), columnType.getProtoTypeFqn());
case TIMESTAMP:
return Value.timestamp(isNull ? null : getTimestampInternal(columnIndex));
case DATE:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -383,7 +383,6 @@ static final class DecodeOption extends InternalOption implements ReadAndQueryOp

DecodeOption(DecodeMode decodeMode) {
this.decodeMode = Preconditions.checkNotNull(decodeMode, "DecodeMode cannot be null");
;
}

@Override
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
/*
* Copyright 2024 Google LLC
*
* Licensed 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 com.google.cloud.spanner;

import com.google.api.core.InternalApi;
import com.google.protobuf.Value;

/** Interface for {@link ResultSet}s that can return a protobuf value. */
@InternalApi
public interface ProtobufResultSet extends ResultSet {

/** Returns true if the protobuf value for the given column is still available. */
boolean canGetProtobufValue(int columnIndex);

/**
* Returns the column value as a protobuf value.
*
* <p>This is an internal method not intended for external usage.
*
* <p>This method may only be called before the column value has been decoded to a plain Java
* object. This means that the {@link DecodeMode} that is used for the {@link ResultSet} must be
* one of {@link DecodeMode#LAZY_PER_ROW} and {@link DecodeMode#LAZY_PER_COL}, and that the
* corresponding {@link ResultSet#getValue(int)}, {@link ResultSet#getBoolean(int)}, ... method
* may not yet have been called for the column.
*/
@InternalApi
Value getProtobufValue(int columnIndex);
}
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ public ResultSet get() {
}
}

private static class PrePopulatedResultSet implements ResultSet {
private static class PrePopulatedResultSet implements ProtobufResultSet {
private final List<Struct> rows;
private final Type type;
private int index = -1;
Expand Down Expand Up @@ -137,6 +137,19 @@ public boolean next() throws SpannerException {
return ++index < rows.size();
}

@Override
public boolean canGetProtobufValue(int columnIndex) {
return !closed && index >= 0 && index < rows.size();
}

@Override
public com.google.protobuf.Value getProtobufValue(int columnIndex) {
Preconditions.checkState(!closed, "ResultSet is closed");
Preconditions.checkState(index >= 0, "Must be preceded by a next() call");
Preconditions.checkElementIndex(index, rows.size(), "All rows have been yielded");
return getValue(columnIndex).toProto();
}

@Override
public Struct getCurrentRowAsStruct() {
Preconditions.checkState(!closed, "ResultSet is closed");
Expand Down
Loading
Loading