com.google.zetasql.SimpleModel 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 com.google.common.base.Ascii;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.zetasql.SimpleModelProtos.SimpleModelProto;
import com.google.zetasql.SimpleTableProtos.SimpleColumnProto;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicLong;
/** SimpleModel is a concrete implementation of the Model interface. */
public final class SimpleModel implements Model {
private final String name;
private final long id;
private final Map inputs = new HashMap<>();
private final Map outputs = new HashMap<>();
private static final AtomicLong nextModelId = new AtomicLong(0);
/** Ensures the next Model id is greater that the provided id */
private static void updateNextIdIfNotGreaterThan(long id) {
nextModelId.getAndAccumulate(
id, (previousValue, minimumValue) -> previousValue <= id ? id + 1 : previousValue);
}
/**
* A pair of a name and a type.
*
* Used to represent inputs and outputs when constructing SimpleModels
*/
public static class NameAndType {
private final String name;
private final Type type;
public NameAndType(String name, Type type) {
this.name = name;
this.type = type;
}
public String getName() {
return this.name;
}
public Type getType() {
return this.type;
}
private SimpleColumn toSimpleColumn(String modelName) {
return new SimpleColumn(modelName, name, type, false, false, false);
}
private static NameAndType deserialize(
SimpleColumnProto proto,
ImmutableList pools,
TypeFactory typeFactory) {
Type type = typeFactory.deserialize(proto.getType(), pools);
return new NameAndType(proto.getName(), type);
}
}
SimpleModel(String name, long id) {
this.name = name;
this.id = id;
updateNextIdIfNotGreaterThan(id);
}
/** Build a SimpleModel with the specified name, initually without any inputs or outputs. */
public SimpleModel(String name) {
this(name, nextModelId.getAndIncrement());
}
SimpleModel(String name, long id, List inputs, List outputs) {
this(name, id);
inputs.forEach(this::addInput);
outputs.forEach(this::addOutput);
}
/**
* Builds a model with the specified name, inputs and outputs.
*
* Fails if there are duplicate column names in either the inputs or the outputs.
*/
public SimpleModel(String name, List inputs, List outputs) {
this(name, nextModelId.getAndIncrement(), inputs, outputs);
}
@Override
public String getName() {
return this.name;
}
@Override
public String getFullName() {
return this.name;
}
@Override
public long getId() {
return this.id;
}
@Override
public ImmutableList getInputs() {
return ImmutableList.copyOf(this.inputs.values());
}
@Override
public ImmutableList getOutputs() {
return ImmutableList.copyOf(this.outputs.values());
}
/**
* Adds an input to this model.
*
* Fails if there's an existing input with the same name.
*/
public void addInput(NameAndType input) {
String inputNameLowercase = Ascii.toLowerCase(input.getName());
Preconditions.checkArgument(
!this.inputs.containsKey(inputNameLowercase),
"Duplicate input name %s for model %s",
input.getName(),
this.name);
this.inputs.put(inputNameLowercase, input.toSimpleColumn(this.name));
}
/**
* Adds an output to this model.
*
*
Fails if there's an existing output with the same name.
*/
public void addOutput(NameAndType output) {
String outputNameLowercase = Ascii.toLowerCase(output.getName());
Preconditions.checkArgument(
!this.outputs.containsKey(outputNameLowercase),
"Duplicate output name %s for model %s",
output.getName(),
this.name);
this.outputs.put(outputNameLowercase, output.toSimpleColumn(this.name));
}
@Override
public boolean equals(Object other) {
if (!(other instanceof SimpleModel)) {
return false;
}
SimpleModel otherAsModel = (SimpleModel) other;
return Objects.equals(this.name, otherAsModel.name)
&& Objects.equals(this.getInputs(), otherAsModel.getInputs())
&& Objects.equals(this.getOutputs(), otherAsModel.getOutputs())
&& this.id == otherAsModel.id;
}
@Override
public int hashCode() {
return Objects.hash(name, inputs, outputs, id);
}
public SimpleModelProto serialize(FileDescriptorSetsBuilder fileDescriptorSetsBuilder) {
SimpleModelProto.Builder builder = SimpleModelProto.newBuilder();
builder.setName(this.name);
builder.setId(this.id);
for (SimpleColumn input : this.inputs.values()) {
SimpleColumnProto.Builder inputBuilder = SimpleColumnProto.newBuilder();
inputBuilder.setName(input.getName());
input.getType().serialize(inputBuilder.getTypeBuilder(), fileDescriptorSetsBuilder);
builder.addInput(inputBuilder.build());
}
for (SimpleColumn output : this.outputs.values()) {
SimpleColumnProto.Builder outputBuilder = SimpleColumnProto.newBuilder();
outputBuilder.setName(output.getName());
output.getType().serialize(outputBuilder.getTypeBuilder(), fileDescriptorSetsBuilder);
builder.addOutput(outputBuilder.build());
}
return builder.build();
}
public static SimpleModel deserialize(
SimpleModelProto proto,
ImmutableList pools,
TypeFactory typeFactory) {
String name = proto.getName();
long id = proto.getId();
ImmutableList inputs =
proto.getInputList().stream()
.map(input -> NameAndType.deserialize(input, pools, typeFactory))
.collect(ImmutableList.toImmutableList());
ImmutableList outputs =
proto.getOutputList().stream()
.map(output -> NameAndType.deserialize(output, pools, typeFactory))
.collect(ImmutableList.toImmutableList());
return new SimpleModel(name, id, inputs, outputs);
}
}