com.google.zetasql.TableValuedFunction Maven / Gradle / Ivy
/*
* Copyright 2019 ZetaSQL Authors
*
* 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 java.nio.charset.StandardCharsets.UTF_8;
import com.google.common.base.Joiner;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.zetasql.FunctionProtos.TVFRelationColumnProto;
import com.google.zetasql.FunctionProtos.TVFRelationProto;
import com.google.zetasql.FunctionProtos.TableValuedFunctionProto;
import com.google.zetasql.ZetaSQLFunctions.FunctionEnums;
import com.google.zetasql.ZetaSQLFunctions.FunctionEnums.TableValuedFunctionType;
import java.io.Serializable;
/**
* This interface describes a table-valued function (TVF) available in a query engine.
*
* More information in zetasql/public/table_valued_function.h
*/
public abstract class TableValuedFunction implements Serializable {
private final ImmutableList namePath;
private final FunctionSignature signature;
private final ImmutableList columns;
/**
* Constructs a new TVF object with the given name and argument signature.
*
* Each TVF may accept value or relation arguments. The signature specifies whether each
* argument should be a value or a relation. For a value argument, the signature may specify a
* concrete Type or a (possibly templated) SignatureArgumentKind. For relation arguments, the
* signature should use ARG_TYPE_RELATION, and any relation will be accepted as an argument.
*/
public TableValuedFunction(ImmutableList namePath, FunctionSignature signature) {
this(namePath, signature, ImmutableList.of());
}
public TableValuedFunction(
ImmutableList namePath,
FunctionSignature signature,
ImmutableList columns) {
this.namePath = namePath;
this.signature = signature;
this.columns = columns;
}
/**
* Deserializes a table-valued function from a protocol buffer.
*/
public static TableValuedFunction deserialize(
TableValuedFunctionProto proto,
ImmutableList pools,
TypeFactory typeFactory) {
switch (proto.getType()) {
case FIXED_OUTPUT_SCHEMA_TVF:
return FixedOutputSchemaTVF.deserialize(proto, pools, typeFactory);
case FORWARD_INPUT_SCHEMA_TO_OUTPUT_SCHEMA_TVF:
return ForwardInputSchemaToOutputSchemaTVF.deserialize(proto, pools, typeFactory);
case TEMPLATED_SQL_TVF:
return TemplatedSQLTVF.deserialize(proto, pools, typeFactory);
case FORWARD_INPUT_SCHEMA_TO_OUTPUT_SCHEMA_WITH_APPENDED_COLUMNS:
return ForwardInputSchemaToOutputSchemaWithAppendedColumnTVF.deserialize(
proto, pools, typeFactory);
default:
StringBuilder builder = new StringBuilder();
for (String name : proto.getNamePathList()) {
builder.append(name);
}
throw new IllegalArgumentException(
"Serialization is not implemented yet for table-valued function: "
+ builder.toString());
}
}
/**
* Serializes this table-valued function to a protocol buffer. Each TVF only supports a single
* signature for now, which we represent with a FunctionSignatureProto within 'proto'.
*/
public TableValuedFunctionProto serialize(FileDescriptorSetsBuilder fileDescriptorSetsBuilder) {
TableValuedFunctionProto.Builder builder =
TableValuedFunctionProto.newBuilder()
.addAllNamePath(namePath)
.setSignature(signature.serialize(fileDescriptorSetsBuilder));
builder.setType(getType());
switch (getType()) {
case FIXED_OUTPUT_SCHEMA_TVF:
case FORWARD_INPUT_SCHEMA_TO_OUTPUT_SCHEMA_TVF:
break;
case FORWARD_INPUT_SCHEMA_TO_OUTPUT_SCHEMA_WITH_APPENDED_COLUMNS:
TVFRelationProto.Builder relationBuilder = TVFRelationProto.newBuilder();
for (TVFRelation.Column column : columns) {
relationBuilder.addColumn(
TVFRelationColumnProto.newBuilder()
.setName(column.getName())
.setType(column.getType().serialize())
.build());
}
builder.setCustomContext(new String(relationBuilder.build().toByteArray(), UTF_8));
break;
case TEMPLATED_SQL_TVF:
for (String name : ((TemplatedSQLTVF) this).argumentNames) {
builder.addArgumentName(name);
}
builder.setParseResumeLocation(((TemplatedSQLTVF) this).parseResumeLocation.serialize());
break;
default:
throw new IllegalArgumentException(
"Serialization is not implemented yet for table-valued function: " + getFullName());
}
return builder.build();
}
abstract FunctionEnums.TableValuedFunctionType getType();
public String getName() {
return namePath.get(namePath.size() - 1);
}
public ImmutableList getNamePath() {
return namePath;
}
public String getFullName() {
return getFullName(true);
}
public String getFullName(boolean includeGroup) {
return Joiner.on('.').join(namePath);
}
@Override
public String toString() {
return "TableValuedFunction";
}
public boolean isDefaultValue() {
return true;
}
/** A TVF that always returns a relation with the same fixed output schema. */
public static class FixedOutputSchemaTVF extends TableValuedFunction {
private final TVFRelation outputSchema;
public FixedOutputSchemaTVF(
ImmutableList namePath, FunctionSignature signature, TVFRelation outputSchema) {
super(namePath, signature);
this.outputSchema = outputSchema;
}
public TVFRelation getOutputSchema() {
return outputSchema;
}
@Override
public FunctionEnums.TableValuedFunctionType getType() {
return FunctionEnums.TableValuedFunctionType.FIXED_OUTPUT_SCHEMA_TVF;
}
/**
* Deserializes this table-valued function from a protocol buffer.
*/
public static FixedOutputSchemaTVF deserialize(
TableValuedFunctionProto proto,
ImmutableList pools,
TypeFactory typeFactory) {
ImmutableList namePath = ImmutableList.copyOf(proto.getNamePathList());
FunctionSignature signature = FunctionSignature.deserialize(proto.getSignature(), pools);
boolean hasRelationInputSchema =
proto.getSignature().getReturnType().getOptions().hasRelationInputSchema();
Preconditions.checkArgument(hasRelationInputSchema, proto);
Preconditions.checkArgument(
proto.getType() == FunctionEnums.TableValuedFunctionType.FIXED_OUTPUT_SCHEMA_TVF, proto);
return new FixedOutputSchemaTVF(
namePath,
signature,
TVFRelation.deserialize(
proto.getSignature().getReturnType().getOptions().getRelationInputSchema(),
pools,
typeFactory));
}
}
/**
* This represents a TVF that accepts a relation for its first argument. The TVF returns a
* relation with the same output schema as this input relation.
*/
public static class ForwardInputSchemaToOutputSchemaTVF extends TableValuedFunction {
public ForwardInputSchemaToOutputSchemaTVF(
ImmutableList namePath, FunctionSignature signature) {
super(namePath, signature);
}
@Override
public FunctionEnums.TableValuedFunctionType getType() {
return FunctionEnums.TableValuedFunctionType.FORWARD_INPUT_SCHEMA_TO_OUTPUT_SCHEMA_TVF;
}
/**
* Deserializes this table-valued function from a protocol buffer.
*/
public static ForwardInputSchemaToOutputSchemaTVF deserialize(
TableValuedFunctionProto proto,
ImmutableList pools,
TypeFactory typeFactory) {
Preconditions.checkArgument(proto.getType()
== FunctionEnums.TableValuedFunctionType.FORWARD_INPUT_SCHEMA_TO_OUTPUT_SCHEMA_TVF,
proto);
ImmutableList namePath = ImmutableList.copyOf(proto.getNamePathList());
FunctionSignature signature = FunctionSignature.deserialize(proto.getSignature(), pools);
return new ForwardInputSchemaToOutputSchemaTVF(namePath, signature);
}
}
/**
* This represents a templated function with a SQL body.
*
* The purpose of this class is to help support statements of the form
* "CREATE FUNCTION () AS ", where the
* may have templated types like "ANY TYPE". In this case, ZetaSQL cannot
* resolve the function expression right away and must defer this work until
* later when the function is called with concrete argument types.
* A TVF that always returns a relation with the same fixed output schema.
*/
public static class TemplatedSQLTVF extends TableValuedFunction {
private final ImmutableList argumentNames;
private final ParseResumeLocation parseResumeLocation;
public TemplatedSQLTVF(
ImmutableList namePath, FunctionSignature signature,
ImmutableList argumentNames, ParseResumeLocation parseResumeLocation) {
super(namePath, signature);
this.argumentNames = argumentNames;
this.parseResumeLocation = parseResumeLocation;
}
public ImmutableList getArgumentNames() {
return argumentNames;
}
public String getSqlBody() {
return this.parseResumeLocation.getInput().substring(
this.parseResumeLocation.getBytePosition());
}
@Override
public FunctionEnums.TableValuedFunctionType getType() {
return FunctionEnums.TableValuedFunctionType.TEMPLATED_SQL_TVF;
}
/**
* Deserializes this table-valued function from a protocol buffer.
*/
public static TemplatedSQLTVF deserialize(
TableValuedFunctionProto proto,
ImmutableList pools,
TypeFactory typeFactory) {
Preconditions.checkArgument(
proto.getType() == FunctionEnums.TableValuedFunctionType.TEMPLATED_SQL_TVF, proto);
ImmutableList namePath = ImmutableList.copyOf(proto.getNamePathList());
FunctionSignature signature = FunctionSignature.deserialize(proto.getSignature(), pools);
ImmutableList.Builder builder = new ImmutableList.Builder();
for (String name : proto.getArgumentNameList()) {
builder.add(name);
}
return new TemplatedSQLTVF(
namePath, signature, builder.build(),
new ParseResumeLocation(proto.getParseResumeLocation()));
}
}
/**
* This represents a TVF that accepts a relation for its first argument. The TVF returns a
* relation with a schema that is constructed by copying the schema of input relation and
* appending extra columns to the schema.
*/
public static class ForwardInputSchemaToOutputSchemaWithAppendedColumnTVF
extends TableValuedFunction {
public ForwardInputSchemaToOutputSchemaWithAppendedColumnTVF(
ImmutableList namePath,
FunctionSignature signature,
ImmutableList columns) {
super(namePath, signature, columns);
}
public static ForwardInputSchemaToOutputSchemaWithAppendedColumnTVF deserialize(
TableValuedFunctionProto proto,
ImmutableList pools,
TypeFactory typeFactory) {
Preconditions.checkArgument(
proto.getType()
== TableValuedFunctionType
.FORWARD_INPUT_SCHEMA_TO_OUTPUT_SCHEMA_WITH_APPENDED_COLUMNS,
proto);
ImmutableList namePath = ImmutableList.copyOf(proto.getNamePathList());
FunctionSignature signature = FunctionSignature.deserialize(proto.getSignature(), pools);
ImmutableList.Builder builder = ImmutableList.builder();
if (proto.hasCustomContext()) {
try {
TVFRelationProto relationProto =
TVFRelationProto.parseFrom(proto.getCustomContextBytes());
for (TVFRelationColumnProto columnProto : relationProto.getColumnList()) {
Type type = typeFactory.deserialize(columnProto.getType(), pools);
builder.add(TVFRelation.Column.create(columnProto.getName(), type));
}
} catch (InvalidProtocolBufferException e) {
throw new IllegalArgumentException(
"Failed to deserialize TVFRelationProto from custom_context in "
+ proto.getNamePath(proto.getNamePathCount() - 1),
e);
}
}
return new ForwardInputSchemaToOutputSchemaWithAppendedColumnTVF(
namePath, signature, builder.build());
}
@Override
TableValuedFunctionType getType() {
return TableValuedFunctionType.FORWARD_INPUT_SCHEMA_TO_OUTPUT_SCHEMA_WITH_APPENDED_COLUMNS;
}
}
}