com.google.zetasql.PreparedExpression Maven / Gradle / Ivy
Show all versions of zetasql-client Show documentation
/*
* Copyright 2019 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.zetasql;
import static com.google.common.base.Verify.verify;
import com.google.common.base.Ascii;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Queues;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.SettableFuture;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.zetasql.LocalService.EvaluateRequest;
import com.google.zetasql.LocalService.EvaluateRequestBatch;
import com.google.zetasql.LocalService.EvaluateResponse;
import com.google.zetasql.LocalService.EvaluateResponseBatch;
import com.google.zetasql.LocalService.PrepareRequest;
import com.google.zetasql.LocalService.PrepareResponse;
import com.google.zetasql.LocalService.PreparedState;
import com.google.zetasql.LocalService.UnprepareRequest;
import io.grpc.Channel;
import io.grpc.ManagedChannel;
import io.grpc.StatusRuntimeException;
import io.grpc.stub.ClientCallStreamObserver;
import io.grpc.stub.ClientResponseObserver;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Optional;
import java.util.Queue;
import java.util.concurrent.atomic.AtomicLong;
/**
* ZetaSQL expression evaluation using Service RPC.
*
* When evaluating an expression, callers can provide
*
*
* - A set of expression columns - column names usable in the expression.
*
- A set of parameters - parameters usable in the expression as \@param.
*
*
* Columns / Parameters are passed as String:Value Maps.
*
*
A prepared expression will create server side state, which should be released by calling
* close() when the expression is no longer used. Note that the close() method is called in
* finalize(), but you should not rely on that because the Java garbage collector has no idea how
* much memory is used on C++ side and GC may happen later than necessary.
*
*
Read the unit tests for examples.
*
*
This class is not thread-safe. External synchronization is needed when it is shared by
* multiple threads.
*/
public class PreparedExpression implements AutoCloseable {
private final String sql;
private boolean prepared = false;
private boolean closed = false;
private Type outputType;
private long preparedId;
private FileDescriptorSetsBuilder fileDescriptorSetsBuilder;
private TypeFactory factory = TypeFactory.nonUniqueNames();
private AnalyzerOptions options;
// expected are the fields we expect the user to provide to execute, but not all are required.
private Map expectedColumns;
private Map expectedParameters;
// referenced is the subset of expected that the analyzer determined are actually required.
private List referencedColumns;
private List referencedParameters;
public PreparedExpression(String sql) {
this.sql = sql;
}
/**
* Prepare the expression with given options. Throwing SqlException if there's an error (not
* necessarily network/server failure). See unit tests for examples how options can be set.
*
* @param options
*/
public void prepare(AnalyzerOptions options) {
prepareInternal(options, Optional.empty());
}
public void prepareWithCatalog(AnalyzerOptions options, SimpleCatalog catalog) {
Preconditions.checkNotNull(catalog);
prepareInternal(options, Optional.of(catalog));
}
private void prepareInternal(AnalyzerOptions options, Optional maybeCatalog) {
Preconditions.checkState(!prepared);
Preconditions.checkState(!closed);
this.options = options;
fileDescriptorSetsBuilder = new FileDescriptorSetsBuilder();
fileDescriptorSetsBuilder.addAllFileDescriptors(BuiltinDescriptorPool.getInstance());
PrepareRequest.Builder request = PrepareRequest.newBuilder();
request.setSql(sql);
request.setOptions(options.serialize(fileDescriptorSetsBuilder));
Map catalogRegisteredDescriptorPoolIds = Collections.emptyMap();
if (maybeCatalog.isPresent()) {
SimpleCatalog catalog = maybeCatalog.get();
if (catalog.isRegistered()) {
catalogRegisteredDescriptorPoolIds = catalog.getRegisteredDescriptorPoolIds();
for (DescriptorPool pool : catalogRegisteredDescriptorPoolIds.keySet()) {
fileDescriptorSetsBuilder.addAllFileDescriptors(pool);
}
request.setRegisteredCatalogId(catalog.getRegisteredId());
} else {
request.setSimpleCatalog(catalog.serialize(fileDescriptorSetsBuilder));
}
} else {
// A catalog must be provided, so we request a default initialized catalog with all of
// the builtin functions.
request
.getSimpleCatalogBuilder()
.getBuiltinFunctionOptionsBuilder()
.setLanguageOptions(request.getOptions().getLanguageOptions());
}
request.setDescriptorPoolList(
DescriptorPoolSerializer.createDescriptorPoolListWithRegisteredIds(
fileDescriptorSetsBuilder, catalogRegisteredDescriptorPoolIds));
PrepareResponse resp;
try {
resp = Client.getStub().prepare(request.build());
} catch (StatusRuntimeException e) {
throw new SqlException(e);
}
setPrepared(resp.getPrepared());
}
private void setPrepared(PreparedState resp) {
preparedId = resp.getPreparedExpressionId();
outputType =
factory.deserialize(resp.getOutputType(), fileDescriptorSetsBuilder.getDescriptorPools());
expectedColumns = toLower(options.getExpressionColumns());
expectedParameters = toLower(options.getQueryParameters());
referencedColumns = resp.getReferencedColumnsList();
referencedParameters = resp.getReferencedParametersList();
prepared = true;
}
/** Get the output type of this expression. */
public Type getOutputType() {
Preconditions.checkState(prepared);
Preconditions.checkState(!closed);
return outputType;
}
/**
* Get the list of column names referenced in this expression.
*
* The columns will be returned in lower case, as column expressions are case-insensitive when
* evaluated. This can be used for efficiency, the list of columns returned from this method are
* the minimal set that must be provided to execute().
*/
public List getReferencedColumns() {
Preconditions.checkState(prepared);
Preconditions.checkState(!closed);
return referencedColumns;
}
/**
* Get the list of parameters referenced in this expression.
*
* This is similar to getReferencedColumns(), but for parameters.
*/
public List getReferencedParameters() {
Preconditions.checkState(prepared);
Preconditions.checkState(!closed);
return referencedParameters;
}
/**
* Evaluate the sql expression via Service RPC.
*
* @return The evaluation result.
*/
@CanIgnoreReturnValue // TODO: consider removing this?
public Value execute() {
return execute(Collections.emptyMap(), Collections.emptyMap());
}
/**
* Evaluate the sql expression via Service RPC.
*
* @param columns Map of column name:value pairs used in the sql expression.
* @param parameters Map of parameter name:value pairs.
* @return The evaluation result.
*/
@CanIgnoreReturnValue // TODO: consider removing this?
public Value execute(Map columns, Map parameters) {
return execute(buildRequest(columns, parameters));
}
private Value execute(EvaluateRequest request) {
Preconditions.checkState(!closed);
final EvaluateResponse resp;
try {
resp = Client.getStub().evaluate(request);
} catch (StatusRuntimeException e) {
throw new SqlException(e);
}
if (!prepared) {
setPrepared(resp.getPrepared());
}
return Value.deserialize(outputType, resp.getValue());
}
private EvaluateRequest buildRequest(Map columns, Map parameters) {
Preconditions.checkNotNull(columns);
Preconditions.checkNotNull(parameters);
EvaluateRequest.Builder request = EvaluateRequest.newBuilder();
if (prepared) {
request.setPreparedExpressionId(preparedId);
final Map normalizedColumns =
normalizeParameters(columns, expectedColumns, "column");
final Map normalizedParameters =
normalizeParameters(parameters, expectedParameters, "query");
for (String column : referencedColumns) {
Value value = normalizedColumns.get(column);
if (value == null) {
throw new SqlException("Incomplete column parameters " + column);
}
request.addColumns(serializeParameter(column, value));
}
for (String param : referencedParameters) {
Value value = normalizedParameters.get(param);
if (value == null) {
throw new SqlException("Incomplete query parameters " + param);
}
request.addParams(serializeParameter(param, value));
}
// TODO: Remove once new descriptor pool is the default
// Force adding this field, as empty to ensure we use the 'new' local_service
// descriptor pool codepath. It is not really needed, as all of the descriptor
// pools are already synchronized via prepare.
request.getDescriptorPoolListBuilder();
} else {
request.setSql(sql);
fileDescriptorSetsBuilder = new FileDescriptorSetsBuilder();
fileDescriptorSetsBuilder.addAllFileDescriptors(BuiltinDescriptorPool.getInstance());
ImmutableMap registeredDescriptorPoolIds = ImmutableMap.of();
options = new AnalyzerOptions();
for (Entry column : columns.entrySet()) {
options.addExpressionColumn(column.getKey(), column.getValue().getType());
request.addColumns(serializeParameter(column.getKey(), column.getValue()));
}
for (Entry param : parameters.entrySet()) {
options.addQueryParameter(param.getKey(), param.getValue().getType());
request.addParams(serializeParameter(param.getKey(), param.getValue()));
}
request.setOptions(options.serialize(fileDescriptorSetsBuilder));
request.setDescriptorPoolList(
DescriptorPoolSerializer.createDescriptorPoolListWithRegisteredIds(
fileDescriptorSetsBuilder, registeredDescriptorPoolIds));
}
return request.build();
}
private static Map normalizeParameters(
Map parameters, Map expected, String kind) {
HashMap lower = new HashMap<>();
for (Map.Entry entry : parameters.entrySet()) {
String name = Ascii.toLowerCase(entry.getKey());
if (!expected.containsKey(name)) {
throw new SqlException("Unexpected " + kind + " parameter '" + name + "'");
}
Value value = entry.getValue();
Type type = expected.get(name);
if (!type.equals(value.getType())) {
throw new SqlException(
"Expected " + kind + " parameter '" + name + "' to be of type " + type);
}
if (lower.putIfAbsent(name, value) != null) {
throw new SqlException("Duplicate expression " + kind + " name '" + name + "'");
}
}
return lower;
}
private static ImmutableMap toLower(Map parameters) {
ImmutableMap.Builder b = new ImmutableMap.Builder<>();
for (Map.Entry entry : parameters.entrySet()) {
b.put(Ascii.toLowerCase(entry.getKey()), entry.getValue());
}
return b.buildOrThrow();
}
/** Opens a handle for streaming execution */
public Stream stream() {
Preconditions.checkState(prepared);
Preconditions.checkState(!closed);
return new Stream();
}
/**
* ZetaSQL streaming expression evaluation.
*
* Streaming evaluation requests are buffered until there is either no outstanding work or an
* optimal number of pending requests (tuned by benchmarks) to send as a single batch. This
* minimizes the average latency, but increases deviation in request latency.
*
*
This class is not thread-safe. External synchronization is needed when it is shared by
* multiple threads.
*/
public final class Stream implements AutoCloseable {
private final Queue> pending;
private final Queue batch;
private final AtomicLong batchCount; // eventually consistent
private final AtomicLong batchSerializedSize; // eventually consistent
private final AtomicLong outstandingCount;
private final Channel channel;
private final ClientCallStreamObserver requestObserver;
private Stream() {
pending = Queues.newConcurrentLinkedQueue();
batch = Queues.newConcurrentLinkedQueue();
batchCount = new AtomicLong();
batchSerializedSize = new AtomicLong();
outstandingCount = new AtomicLong();
channel = ClientChannelProvider.loadChannel();
ZetaSqlLocalServiceGrpc.ZetaSqlLocalServiceStub stub =
ZetaSqlLocalServiceGrpc.newStub(channel);
requestObserver =
(ClientCallStreamObserver)
stub.evaluateStream(new ResponseObserver());
}
/**
* Evaluate the sql expression via streaming RPC.
*
* This method is not thread-safe. External synchronization is needed when it is shared by
* multiple threads.
*
* @param columns Map of column name:value pairs used in the sql expression.
* @param parameters Map of parameter name:value pairs.
* @return The evaluation result.
*/
public ListenableFuture execute(
Map columns, Map parameters) {
final SettableFuture f = SettableFuture.create();
final EvaluateRequest r = buildRequest(columns, parameters);
final int size = r.getSerializedSize();
/* critical block: This block is not thread safe, as documented in the
* PreparedExpression class comment, for maximal performance. The order
* of these operations is critical to the correctness of this code.
*/
pending.add(f);
batch.add(r);
batchCount.getAndIncrement();
batchSerializedSize.getAndAdd(size);
// end critical block
maybeFlush();
return f;
}
private void maybeFlush() {
/* batchCount and batchSerializedSize are eventually consistent hints.
* They may be too low if called concurrently with execute, which will
* be mitigated when execute calls this method. They may be too high
* if called concurrently with flush, which will result in an unneeded
* call to flush.
*/
final long batchCount = this.batchCount.get();
if (batchCount <= 0) {
return;
}
if (outstandingCount.get() != 0
&& batchSerializedSize.get() < 65536
&& (batchCount < 16 || !requestObserver.isReady())) {
return;
}
flush();
}
/** Flush buffered requests to RPC. */
public synchronized void flush() {
try {
final EvaluateRequestBatch.Builder b = EvaluateRequestBatch.newBuilder();
for (EvaluateRequest r; (r = batch.poll()) != null; ) {
batchSerializedSize.getAndAdd(-r.getSerializedSize());
b.addRequest(r);
}
if (b.getRequestCount() > 0) {
batchCount.getAndAdd(-b.getRequestCount());
outstandingCount.getAndAdd(b.getRequestCount());
requestObserver.onNext(b.build());
}
} catch (RuntimeException e) {
requestObserver.onError(e);
throw e;
}
}
private class ResponseObserver
implements ClientResponseObserver {
@Override
public void onNext(EvaluateResponseBatch respb) {
long remain = outstandingCount.addAndGet(-respb.getResponseCount());
verify(remain >= 0);
for (EvaluateResponse resp : respb.getResponseList()) {
final SettableFuture f = pending.remove();
final Value v;
try {
v = Value.deserialize(outputType, resp.getValue());
} catch (RuntimeException e) {
f.setException(e);
continue;
}
f.set(v);
}
if (remain == 0) {
/* If no outstanding requests remain, flush the current batch. This
* avoids a hang if the caller stops calling evaluate but doesn't
* call flush and limits the latency of a request to 2T (where T is
* the round-trip latency of a single request). This may result in
* suboptimal behavior for a time if the request rate slows to
* between T and 2T.
*/
flush();
}
}
private void setException(Throwable t) {
for (SettableFuture f; (f = pending.poll()) != null; ) {
f.setException(t);
}
}
@Override
public void onError(Throwable t) {
if (t instanceof StatusRuntimeException) {
setException(new SqlException((StatusRuntimeException) t));
} else {
setException(t);
}
}
@Override
public void onCompleted() {
setException(new RuntimeException("Stream closed"));
}
@Override
public void beforeStart(final ClientCallStreamObserver requestStream) {
requestStream.setOnReadyHandler(this::onReady);
}
private void onReady() {
maybeFlush();
}
}
@Override
public void close() {
synchronized (this) {
requestObserver.onCompleted();
}
if (channel instanceof ManagedChannel) {
((ManagedChannel) channel).shutdown();
}
}
@Override
protected void finalize() throws Throwable {
super.finalize();
close();
}
}
/** Release the server side state for this prepared expression. */
@Override
public void close() {
if (prepared && !closed) {
try {
Client.getStub()
.unprepare(UnprepareRequest.newBuilder().setPreparedExpressionId(preparedId).build());
} catch (StatusRuntimeException e) {
// ignore
}
prepared = false;
closed = true;
}
}
@Override
protected void finalize() throws Throwable {
super.finalize();
close();
}
private static EvaluateRequest.Parameter serializeParameter(String name, Value value) {
return EvaluateRequest.Parameter.newBuilder()
.setName(name)
.setValue(value.serialize())
.build();
}
}