Skip to content

Commit

Permalink
feat: Add support for additional types
Browse files Browse the repository at this point in the history
  • Loading branch information
ron-gal committed Jul 23, 2024
1 parent 509a0c6 commit e235a6e
Show file tree
Hide file tree
Showing 5 changed files with 242 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -31,11 +31,16 @@ 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 abstract static class MinAggregateInput extends Type {}

public abstract static class MaxAggregateInput extends Type {}

public abstract static class HllAggregateInput extends Type {}

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

static Type fromProto(com.google.bigtable.admin.v2.Type source) {
Expand Down Expand Up @@ -91,6 +96,36 @@ public static Aggregate sum(SumAggregateInput inputType) {
return Aggregate.create(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 Aggregate min(MinAggregateInput 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 Aggregate max(MaxAggregateInput 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() {
return hll(rawBytes());
}

/** Creates an Aggregate type with a HLL aggregator and specified input type. */
public static Aggregate hll(HllAggregateInput 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 {
Expand Down Expand Up @@ -250,6 +285,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.setHll(
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 +344,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 HLL:
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.rawBytesHll())
.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 @@ -48,4 +48,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.int64Type())
.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);
}
}

0 comments on commit e235a6e

Please sign in to comment.