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

feat: Add support for new functions #2287

Merged
merged 13 commits into from
Jul 24, 2024
Original file line number Diff line number Diff line change
Expand Up @@ -31,10 +31,15 @@ public abstract class Type {
private Type() {}

/**
* This type is a marker type that allows types to be used as the input to the SUM aggregate
* function.
* These types are marker types that allow types to be used as the input to aggregate function.
*/
public abstract static class SumAggregateInput extends Type {}
public interface SumAggregateInput {}

public interface MinAggregateInput {}

public interface MaxAggregateInput {}

public interface HllAggregateInput {}

abstract com.google.bigtable.admin.v2.Type toProto();

Expand Down Expand Up @@ -87,13 +92,48 @@ public static Aggregate int64Sum() {
}

/** Creates an Aggregate type with a SUM aggregator and specified input type. */
public static Aggregate sum(SumAggregateInput inputType) {
public static <Input extends Type & SumAggregateInput> Aggregate sum(Input inputType) {
return Aggregate.create(inputType, Aggregate.Aggregator.Sum.create());
}

@Deprecated
public static Aggregate sum(SumAggregateInput inputType) {
return Aggregate.create((Type) inputType, Aggregate.Aggregator.Sum.create());
}

/** Creates an Aggregate type with a MIN aggregator and Int64 input type. */
public static Aggregate int64Min() {
return min(bigEndianInt64());
}

/** Creates an Aggregate type with a MIN aggregator and specified input type. */
public static <Input extends Type & MinAggregateInput> Aggregate min(Input inputType) {
return Aggregate.create(inputType, Aggregate.Aggregator.Min.create());
}

/** Creates an Aggregate type with a MAX aggregator and Int64 input type. */
public static Aggregate int64Max() {
return max(bigEndianInt64());
}

/** Creates an Aggregate type with a MAX aggregator and specified input type. */
public static <Input extends Type & MaxAggregateInput> Aggregate max(Input inputType) {
return Aggregate.create(inputType, Aggregate.Aggregator.Max.create());
}

/** Creates an Aggregate type with a HLL aggregator and Int64 input type. */
public static Aggregate rawBytesHll() {
Copy link
Contributor

Choose a reason for hiding this comment

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

shouldn't this be called int64Hll?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Whoops! D= fixed!

return hll(rawBytes());
}

/** Creates an Aggregate type with a HLL aggregator and specified input type. */
public static <Input extends Type & HllAggregateInput> Aggregate hll(Input inputType) {
return Aggregate.create(inputType, Aggregate.Aggregator.Hll.create());
}

/** Represents a string of bytes with a specific encoding. */
@AutoValue
public abstract static class Bytes extends Type {
public abstract static class Bytes extends Type implements HllAggregateInput {
Copy link
Contributor

Choose a reason for hiding this comment

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

same, we only support int64 inputs right now so we shouldn't have this here.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

fixed

public static Bytes create(Encoding encoding) {
return new AutoValue_Type_Bytes(encoding);
}
Expand Down Expand Up @@ -151,7 +191,8 @@ com.google.bigtable.admin.v2.Type.Bytes.Encoding toProto() {

/** Represents a 64-bit integer with a specific encoding. */
@AutoValue
public abstract static class Int64 extends SumAggregateInput {
public abstract static class Int64 extends Type
implements SumAggregateInput, MinAggregateInput, MaxAggregateInput {
public static Int64 create(Encoding encoding) {
return new AutoValue_Type_Int64(encoding);
}
Expand Down Expand Up @@ -250,6 +291,44 @@ void buildTo(com.google.bigtable.admin.v2.Type.Aggregate.Builder builder) {
}
}

@AutoValue
public abstract static class Min extends Aggregator {
public static Min create() {
return new AutoValue_Type_Aggregate_Aggregator_Min();
}

@Override
void buildTo(com.google.bigtable.admin.v2.Type.Aggregate.Builder builder) {
builder.setMin(com.google.bigtable.admin.v2.Type.Aggregate.Min.getDefaultInstance());
}
}

@AutoValue
public abstract static class Max extends Aggregator {
public static Max create() {
return new AutoValue_Type_Aggregate_Aggregator_Max();
}

@Override
void buildTo(com.google.bigtable.admin.v2.Type.Aggregate.Builder builder) {
builder.setMax(com.google.bigtable.admin.v2.Type.Aggregate.Max.getDefaultInstance());
}
}

@AutoValue
public abstract static class Hll extends Aggregator {
public static Hll create() {
return new AutoValue_Type_Aggregate_Aggregator_Hll();
}

@Override
void buildTo(com.google.bigtable.admin.v2.Type.Aggregate.Builder builder) {
builder.setHllppUniqueCount(
com.google.bigtable.admin.v2.Type.Aggregate.HyperLogLogPlusPlusUniqueCount
.getDefaultInstance());
}
}

abstract void buildTo(com.google.bigtable.admin.v2.Type.Aggregate.Builder builder);
}

Expand All @@ -271,6 +350,15 @@ static Aggregate fromProto(com.google.bigtable.admin.v2.Type.Aggregate source) {
case SUM:
aggregator = Aggregator.Sum.create();
break;
case MIN:
aggregator = Aggregator.Min.create();
break;
case MAX:
aggregator = Aggregator.Max.create();
break;
case HLLPP_UNIQUE_COUNT:
aggregator = Aggregator.Hll.create();
break;
case AGGREGATOR_NOT_SET:
throw new UnsupportedOperationException();
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -256,6 +256,24 @@ public void testCreateTable() {
ColumnFamily.newBuilder()
.setGcRule(GcRule.getDefaultInstance())
.setValueType(TypeProtos.intSumType())
.build())
.putColumnFamilies(
"cf2",
ColumnFamily.newBuilder()
.setGcRule(GcRule.getDefaultInstance())
.setValueType(TypeProtos.intMinType())
.build())
.putColumnFamilies(
"cf3",
ColumnFamily.newBuilder()
.setGcRule(GcRule.getDefaultInstance())
.setValueType(TypeProtos.intMaxType())
.build())
.putColumnFamilies(
"cf4",
ColumnFamily.newBuilder()
.setGcRule(GcRule.getDefaultInstance())
.setValueType(TypeProtos.bytesHllType())
.build()))
.build();

Expand All @@ -267,7 +285,12 @@ public void testCreateTable() {

// Execute
Table result =
adminClient.createTable(CreateTableRequest.of(TABLE_ID).addFamily("cf1", Type.int64Sum()));
adminClient.createTable(
CreateTableRequest.of(TABLE_ID)
.addFamily("cf1", Type.int64Sum())
.addFamily("cf2", Type.int64Min())
.addFamily("cf3", Type.int64Max())
.addFamily("cf4", Type.rawBytesHll()));

// Verify
assertThat(result).isEqualTo(Table.fromProto(expectedResponse));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,10 @@ public static com.google.bigtable.admin.v2.Type.Bytes bytesType() {
.build();
}

public static com.google.bigtable.admin.v2.Type rawBytesType() {
return com.google.bigtable.admin.v2.Type.newBuilder().setBytesType(bytesType()).build();
}

public static com.google.bigtable.admin.v2.Type int64Type() {
return com.google.bigtable.admin.v2.Type.newBuilder()
.setInt64Type(
Expand All @@ -48,4 +52,33 @@ public static com.google.bigtable.admin.v2.Type intSumType() {
.setSum(com.google.bigtable.admin.v2.Type.Aggregate.Sum.getDefaultInstance()))
.build();
}

public static com.google.bigtable.admin.v2.Type intMinType() {
return com.google.bigtable.admin.v2.Type.newBuilder()
.setAggregateType(
com.google.bigtable.admin.v2.Type.Aggregate.newBuilder()
.setInputType(TypeProtos.int64Type())
.setMin(com.google.bigtable.admin.v2.Type.Aggregate.Min.getDefaultInstance()))
.build();
}

public static com.google.bigtable.admin.v2.Type intMaxType() {
return com.google.bigtable.admin.v2.Type.newBuilder()
.setAggregateType(
com.google.bigtable.admin.v2.Type.Aggregate.newBuilder()
.setInputType(TypeProtos.int64Type())
.setMax(com.google.bigtable.admin.v2.Type.Aggregate.Max.getDefaultInstance()))
.build();
}

public static com.google.bigtable.admin.v2.Type bytesHllType() {
return com.google.bigtable.admin.v2.Type.newBuilder()
.setAggregateType(
com.google.bigtable.admin.v2.Type.Aggregate.newBuilder()
.setInputType(TypeProtos.rawBytesType())
.setHllppUniqueCount(
com.google.bigtable.admin.v2.Type.Aggregate.HyperLogLogPlusPlusUniqueCount
.getDefaultInstance()))
.build();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,12 @@ public void modifyFamilies() {
.addFamily("cf3")
.addFamily("cf4", Type.int64Sum())
.addFamily("cf5", GCRules.GCRULES.maxVersions(1), Type.int64Sum())
.addFamily("cf6", Type.int64Min())
.addFamily("cf7", GCRules.GCRULES.maxVersions(1), Type.int64Min())
.addFamily("cf8", Type.int64Max())
.addFamily("cf9", GCRules.GCRULES.maxVersions(1), Type.int64Max())
.addFamily("cf10", Type.rawBytesHll())
.addFamily("cf11", GCRules.GCRULES.maxVersions(1), Type.rawBytesHll())
.updateFamily("cf1", GCRules.GCRULES.maxVersions(5))
.dropFamily("cf3")
.toProto(PROJECT_ID, INSTANCE_ID);
Expand Down Expand Up @@ -119,6 +125,48 @@ public void modifyFamilies() {
com.google.bigtable.admin.v2.ColumnFamily.newBuilder()
.setGcRule(GCRules.GCRULES.maxVersions(1).toProto())
.setValueType(Type.int64Sum().toProto())))
.addModifications(
com.google.bigtable.admin.v2.ModifyColumnFamiliesRequest.Modification.newBuilder()
.setId("cf6")
.setCreate(
com.google.bigtable.admin.v2.ColumnFamily.newBuilder()
.setGcRule(GcRule.getDefaultInstance())
.setValueType(Type.int64Min().toProto())))
.addModifications(
com.google.bigtable.admin.v2.ModifyColumnFamiliesRequest.Modification.newBuilder()
.setId("cf7")
.setCreate(
com.google.bigtable.admin.v2.ColumnFamily.newBuilder()
.setGcRule(GCRules.GCRULES.maxVersions(1).toProto())
.setValueType(Type.int64Min().toProto())))
.addModifications(
com.google.bigtable.admin.v2.ModifyColumnFamiliesRequest.Modification.newBuilder()
.setId("cf8")
.setCreate(
com.google.bigtable.admin.v2.ColumnFamily.newBuilder()
.setGcRule(GcRule.getDefaultInstance())
.setValueType(Type.int64Max().toProto())))
.addModifications(
com.google.bigtable.admin.v2.ModifyColumnFamiliesRequest.Modification.newBuilder()
.setId("cf9")
.setCreate(
com.google.bigtable.admin.v2.ColumnFamily.newBuilder()
.setGcRule(GCRules.GCRULES.maxVersions(1).toProto())
.setValueType(Type.int64Max().toProto())))
.addModifications(
com.google.bigtable.admin.v2.ModifyColumnFamiliesRequest.Modification.newBuilder()
.setId("cf10")
.setCreate(
com.google.bigtable.admin.v2.ColumnFamily.newBuilder()
.setGcRule(GcRule.getDefaultInstance())
.setValueType(Type.rawBytesHll().toProto())))
.addModifications(
com.google.bigtable.admin.v2.ModifyColumnFamiliesRequest.Modification.newBuilder()
.setId("cf11")
.setCreate(
com.google.bigtable.admin.v2.ColumnFamily.newBuilder()
.setGcRule(GCRules.GCRULES.maxVersions(1).toProto())
.setValueType(Type.rawBytesHll().toProto())))
.addModifications(
com.google.bigtable.admin.v2.ModifyColumnFamiliesRequest.Modification.newBuilder()
.setId("cf1")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -77,4 +77,61 @@ public void intSumFromProtoToProto() {
assertThat(Type.fromProto(proto)).isEqualTo(Type.int64Sum());
assertThat(Type.fromProto(proto).toProto()).isEqualTo(proto);
}

@Test
public void int64Min() {
Type type = Type.int64Min();
assertThat(type.toProto()).isEqualTo(TypeProtos.intMinType());
}

@Test
public void min() {
Type type = Type.min(Type.bigEndianInt64());
assertThat(type.toProto()).isEqualTo(TypeProtos.intMinType());
}

@Test
public void intMinFromProtoToProto() {
com.google.bigtable.admin.v2.Type proto = TypeProtos.intMinType();
assertThat(Type.fromProto(proto)).isEqualTo(Type.int64Min());
assertThat(Type.fromProto(proto).toProto()).isEqualTo(proto);
}

@Test
public void int64Max() {
Type type = Type.int64Max();
assertThat(type.toProto()).isEqualTo(TypeProtos.intMaxType());
}

@Test
public void max() {
Type type = Type.max(Type.bigEndianInt64());
assertThat(type.toProto()).isEqualTo(TypeProtos.intMaxType());
}

@Test
public void intMaxFromProtoToProto() {
com.google.bigtable.admin.v2.Type proto = TypeProtos.intMaxType();
assertThat(Type.fromProto(proto)).isEqualTo(Type.int64Max());
assertThat(Type.fromProto(proto).toProto()).isEqualTo(proto);
}

@Test
public void bytesHll() {
Type type = Type.rawBytesHll();
assertThat(type.toProto()).isEqualTo(TypeProtos.bytesHllType());
}

@Test
public void hll() {
Type type = Type.hll(Type.rawBytes());
assertThat(type.toProto()).isEqualTo(TypeProtos.bytesHllType());
}

@Test
public void bytesHllFromProtoToProto() {
com.google.bigtable.admin.v2.Type proto = TypeProtos.bytesHllType();
assertThat(Type.fromProto(proto)).isEqualTo(Type.rawBytesHll());
assertThat(Type.fromProto(proto).toProto()).isEqualTo(proto);
}
}
Loading