Skip to content

Commit

Permalink
Support FLOAT32 type in Spanner (#30893)
Browse files Browse the repository at this point in the history
  • Loading branch information
@arawind
arawind authored Apr 12, 2024
1 parent 0d5d567 commit 325c0b5
Show file tree
Hide file tree
Showing 7 changed files with 67 additions and 27 deletions.
4 changes: 4 additions & 0 deletions ...loud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,8 @@ private static long estimatePrimitiveValue(Value v) {
switch (v.getType().getCode()) {
case BOOL:
return 1;
case FLOAT32:
return 4;
case INT64:
case FLOAT64:
case ENUM:
Expand Down Expand Up @@ -142,6 +144,8 @@ private static long estimateArrayValue(Value v) {
switch (v.getType().getArrayElementType().getCode()) {
case BOOL:
return v.getBoolArray().size();
case FLOAT32:
return 4L * v.getFloat32Array().size();
case INT64:
case ENUM:
return 8L * v.getInt64Array().size();
Expand Down
9 changes: 3 additions & 6 deletions ...google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationUtils.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -207,12 +207,7 @@ private static void setBeamValueToMutation(
mutationBuilder.set(columnName).to(row.getInt64(columnName));
break;
case FLOAT:
@Nullable Float floatValue = row.getFloat(columnName);
if (floatValue == null) {
mutationBuilder.set(columnName).to(((Double) null));
} else {
mutationBuilder.set(columnName).to(floatValue);
}
mutationBuilder.set(columnName).to(row.getFloat(columnName));
break;
case DOUBLE:
mutationBuilder.set(columnName).to(row.getDouble(columnName));
Expand Down Expand Up @@ -311,6 +306,8 @@ private static void addIterableToMutationBuilder(
mutationBuilder.set(column).toInt64Array((Iterable<Long>) ((Object) iterable));
break;
case FLOAT:
mutationBuilder.set(column).toFloat32Array((Iterable<Float>) ((Object) iterable));
break;
case DOUBLE:
mutationBuilder.set(column).toFloat64Array((Iterable<Double>) ((Object) iterable));
break;
Expand Down
6 changes: 6 additions & 0 deletions ...google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchema.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,9 @@ private static Type parseSpannerType(String spannerType, Dialect dialect) {
if ("INT64".equals(spannerType)) {
return Type.int64();
}
if ("FLOAT32".equals(spannerType)) {
return Type.float32();
}
if ("FLOAT64".equals(spannerType)) {
return Type.float64();
}
Expand Down Expand Up @@ -227,6 +230,9 @@ private static Type parseSpannerType(String spannerType, Dialect dialect) {
if ("BIGINT".equals(spannerType)) {
return Type.int64();
}
if ("REAL".equals(spannerType)) {
return Type.float32();
}
if ("DOUBLE PRECISION".equals(spannerType)) {
return Type.float64();
}
Expand Down
16 changes: 9 additions & 7 deletions ...o/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/StructUtils.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,12 +76,7 @@ public static Struct beamRowToStruct(Row row) {
addIterableToStructBuilder(structBuilder, row.getIterable(column), field);
break;
case FLOAT:
@Nullable Float floatValue = row.getFloat(column);
if (floatValue == null) {
structBuilder.set(column).to((Double) null);
} else {
structBuilder.set(column).to(floatValue);
}
structBuilder.set(column).to(row.getFloat(column));
break;
case DOUBLE:
structBuilder.set(column).to(row.getDouble(column));
Expand Down Expand Up @@ -187,8 +182,9 @@ private static Type simpleBeamTypeToSpannerType(Schema.FieldType beamType) {
case INT16:
return Type.int64();
case DOUBLE:
case FLOAT:
return Type.float64();
case FLOAT:
return Type.float32();
case DECIMAL:
return Type.numeric();
case STRING:
Expand Down Expand Up @@ -242,6 +238,8 @@ private static void addIterableToStructBuilder(
structBuilder.set(column).toInt64Array((Iterable<Long>) ((Object) iterable));
break;
case FLOAT:
structBuilder.set(column).toFloat32Array((Iterable<Float>) ((Object) iterable));
break;
case DOUBLE:
structBuilder.set(column).toFloat64Array((Iterable<Double>) ((Object) iterable));
break;
Expand Down Expand Up @@ -306,6 +304,8 @@ private static void addIterableToStructBuilder(
return DateTime.parse(struct.getDate(column).toString());
case INT64:
return struct.getLong(column);
case FLOAT32:
return struct.getFloat(column);
case FLOAT64:
return struct.getDouble(column);
case NUMERIC:
Expand Down Expand Up @@ -352,6 +352,8 @@ private static void addIterableToStructBuilder(
.collect(toList());
case INT64:
return struct.getLongList(column);
case FLOAT32:
return struct.getFloatList(column);
case FLOAT64:
return struct.getDoubleList(column);
case STRING:
Expand Down
16 changes: 14 additions & 2 deletions ...-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimatorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ public class MutationSizeEstimatorTest {
@Test
public void primitives() throws Exception {
Mutation int64 = Mutation.newInsertOrUpdateBuilder("test").set("one").to(1).build();
Mutation float32 = Mutation.newInsertOrUpdateBuilder("test").set("one").to(1.3f).build();
Mutation float64 = Mutation.newInsertOrUpdateBuilder("test").set("one").to(2.9).build();
Mutation bool = Mutation.newInsertOrUpdateBuilder("test").set("one").to(false).build();
Mutation numeric =
Expand Down Expand Up @@ -72,6 +73,7 @@ public void primitives() throws Exception {
.build();

assertThat(MutationSizeEstimator.sizeOf(int64), is(8L));
assertThat(MutationSizeEstimator.sizeOf(float32), is(4L));
assertThat(MutationSizeEstimator.sizeOf(float64), is(8L));
assertThat(MutationSizeEstimator.sizeOf(bool), is(1L));
assertThat(MutationSizeEstimator.sizeOf(numeric), is(30L));
Expand All @@ -89,6 +91,11 @@ public void primitiveArrays() throws Exception {
.set("one")
.toInt64Array(new long[] {1L, 2L, 3L})
.build();
Mutation float32 =
Mutation.newInsertOrUpdateBuilder("test")
.set("one")
.toFloat32Array(new float[] {1.0f, 2.0f})
.build();
Mutation float64 =
Mutation.newInsertOrUpdateBuilder("test")
.set("one")
Expand Down Expand Up @@ -160,6 +167,7 @@ public void primitiveArrays() throws Exception {
"customer.app.TestMessage")
.build();
assertThat(MutationSizeEstimator.sizeOf(int64), is(24L));
assertThat(MutationSizeEstimator.sizeOf(float32), is(8L));
assertThat(MutationSizeEstimator.sizeOf(float64), is(16L));
assertThat(MutationSizeEstimator.sizeOf(bool), is(4L));
assertThat(MutationSizeEstimator.sizeOf(numeric), is(153L));
Expand All @@ -180,6 +188,8 @@ public void nullPrimitiveArrays() throws Exception {
.set("one")
.toProtoEnumArray(null, "customer.app.TestEnum")
.build();
Mutation float32 =
Mutation.newInsertOrUpdateBuilder("test").set("one").toFloat32Array((float[]) null).build();
Mutation float64 =
Mutation.newInsertOrUpdateBuilder("test")
.set("one")
Expand All @@ -202,6 +212,7 @@ public void nullPrimitiveArrays() throws Exception {
Mutation.newInsertOrUpdateBuilder("test").set("one").toPgJsonbArray(null).build();

assertThat(MutationSizeEstimator.sizeOf(int64), is(0L));
assertThat(MutationSizeEstimator.sizeOf(float32), is(0L));
assertThat(MutationSizeEstimator.sizeOf(float64), is(0L));
assertThat(MutationSizeEstimator.sizeOf(bool), is(0L));
assertThat(MutationSizeEstimator.sizeOf(numeric), is(0L));
Expand Down Expand Up @@ -384,12 +395,13 @@ public void dates() throws Exception {
@Test
public void group() throws Exception {
Mutation int64 = Mutation.newInsertOrUpdateBuilder("test").set("one").to(1).build();
Mutation float32 = Mutation.newInsertOrUpdateBuilder("test").set("one").to(1.3f).build();
Mutation float64 = Mutation.newInsertOrUpdateBuilder("test").set("one").to(2.9).build();
Mutation bool = Mutation.newInsertOrUpdateBuilder("test").set("one").to(false).build();

MutationGroup group = MutationGroup.create(int64, float64, bool);
MutationGroup group = MutationGroup.create(int64, float32, float64, bool);

assertThat(MutationSizeEstimator.sizeOf(group), is(17L));
assertThat(MutationSizeEstimator.sizeOf(group), is(21L));
}

@Test
Expand Down
22 changes: 15 additions & 7 deletions ...le-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/MutationUtilsTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ public class MutationUtilsTest {
private static final Schema WRITE_ROW_SCHEMA =
Schema.builder()
.addNullableField("f_int64", Schema.FieldType.INT64)
.addNullableField("f_float32", Schema.FieldType.FLOAT)
.addNullableField("f_float64", Schema.FieldType.DOUBLE)
.addNullableField("f_string", Schema.FieldType.STRING)
.addNullableField("f_bytes", Schema.FieldType.BYTES)
Expand All @@ -56,6 +57,7 @@ public class MutationUtilsTest {
.addNullableField("f_struct", Schema.FieldType.row(EMPTY_SCHEMA))
.addNullableField("f_struct_int64", Schema.FieldType.row(INT64_SCHEMA))
.addNullableField("f_array", Schema.FieldType.array(Schema.FieldType.INT64))
.addNullableField("f_float_array", Schema.FieldType.array(Schema.FieldType.FLOAT))
.addNullableField("f_double_array", Schema.FieldType.array(Schema.FieldType.DOUBLE))
.addNullableField("f_decimal_array", Schema.FieldType.array(Schema.FieldType.DECIMAL))
.addNullableField("f_boolean_array", Schema.FieldType.array(Schema.FieldType.BOOLEAN))
Expand All @@ -66,7 +68,6 @@ public class MutationUtilsTest {
"f_struct_array", Schema.FieldType.array(Schema.FieldType.row(INT64_SCHEMA)))
.addNullableField("f_int16", Schema.FieldType.INT16)
.addNullableField("f_int32", Schema.FieldType.INT32)
.addNullableField("f_float", Schema.FieldType.FLOAT)
.addNullableField("f_decimal", Schema.FieldType.DECIMAL)
.addNullableField("f_byte", Schema.FieldType.BYTE)
.addNullableField("f_iterable", Schema.FieldType.iterable(Schema.FieldType.INT64))
Expand All @@ -75,6 +76,7 @@ public class MutationUtilsTest {
private static final Row WRITE_ROW =
Row.withSchema(WRITE_ROW_SCHEMA)
.withFieldValue("f_int64", 1L)
.withFieldValue("f_float32", 2.1f)
.withFieldValue("f_float64", 1.1)
.withFieldValue("f_string", "donald_duck")
.withFieldValue("f_bytes", "some_bytes".getBytes(UTF_8))
Expand All @@ -83,6 +85,7 @@ public class MutationUtilsTest {
.withFieldValue("f_struct", EMPTY_ROW)
.withFieldValue("f_struct_int64", INT64_ROW)
.withFieldValue("f_array", ImmutableList.of(2L, 3L))
.withFieldValue("f_float_array", ImmutableList.of(3.0f, 4.0f))
.withFieldValue("f_double_array", ImmutableList.of(1., 2.))
.withFieldValue(
"f_decimal_array",
Expand All @@ -101,7 +104,6 @@ public class MutationUtilsTest {
.withFieldValue("f_struct_array", ImmutableList.of(INT64_ROW, INT64_ROW))
.withFieldValue("f_int16", (short) 2)
.withFieldValue("f_int32", 0x7fffffff)
.withFieldValue("f_float", 0.0f)
.withFieldValue("f_decimal", BigDecimal.valueOf(Long.MIN_VALUE))
.withFieldValue("f_byte", Byte.parseByte("127"))
.withFieldValue("f_iterable", ImmutableList.of(2L, 3L))
Expand All @@ -110,6 +112,7 @@ public class MutationUtilsTest {
private static final Schema WRITE_ROW_SCHEMA_NULLS =
Schema.builder()
.addNullableField("f_int64", Schema.FieldType.INT64)
.addNullableField("f_float32", Schema.FieldType.FLOAT)
.addNullableField("f_float64", Schema.FieldType.DOUBLE)
.addNullableField("f_string", Schema.FieldType.STRING)
.addNullableField("f_bytes", Schema.FieldType.BYTES)
Expand All @@ -134,34 +137,35 @@ public class MutationUtilsTest {
.addValue(null)
.addValue(null)
.addValue(null)
.addValue(null)
.build();

private static final Schema KEY_SCHEMA =
Schema.builder()
.addNullableField("f_int64", Schema.FieldType.INT64)
.addNullableField("f_float32", Schema.FieldType.FLOAT)
.addNullableField("f_float64", Schema.FieldType.DOUBLE)
.addNullableField("f_string", Schema.FieldType.STRING)
.addNullableField("f_bytes", Schema.FieldType.BYTES)
.addNullableField("f_date_time", Schema.FieldType.DATETIME)
.addNullableField("f_bool", Schema.FieldType.BOOLEAN)
.addNullableField("f_int16", Schema.FieldType.INT16)
.addNullableField("f_int32", Schema.FieldType.INT32)
.addNullableField("f_float", Schema.FieldType.FLOAT)
.addNullableField("f_decimal", Schema.FieldType.DECIMAL)
.addNullableField("f_byte", Schema.FieldType.BYTE)
.build();

private static final Row KEY_ROW =
Row.withSchema(KEY_SCHEMA)
.withFieldValue("f_int64", 1L)
.withFieldValue("f_float32", 2.1f)
.withFieldValue("f_float64", 1.1)
.withFieldValue("f_string", "donald_duck")
.withFieldValue("f_bytes", "some_bytes".getBytes(UTF_8))
.withFieldValue("f_date_time", DateTime.parse("2077-10-15T00:00:00+00:00"))
.withFieldValue("f_bool", false)
.withFieldValue("f_int16", (short) 2)
.withFieldValue("f_int32", 0x7fffffff)
.withFieldValue("f_float", 0.0f)
.withFieldValue("f_decimal", BigDecimal.valueOf(Long.MIN_VALUE))
.withFieldValue("f_byte", Byte.parseByte("127"))
.build();
Expand Down Expand Up @@ -263,14 +267,14 @@ private static Mutation createDeleteMutation() {
Key key =
Key.newBuilder()
.append(1L)
.append(2.1f)
.append(1.1)
.append("donald_duck")
.append(ByteArray.copyFrom("some_bytes".getBytes(UTF_8)))
.append(Timestamp.parseTimestamp("2077-10-15T00:00:00"))
.append(false)
.append((short) 2)
.append(0x7fffffff)
.append(0.0f)
.append(BigDecimal.valueOf(Long.MIN_VALUE))
.append(Byte.parseByte("127"))
.build();
Expand All @@ -295,6 +299,8 @@ private static Mutation createMutation(Mutation.Op operation) {
return builder
.set("f_int64")
.to(1L)
.set("f_float32")
.to(2.1f)
.set("f_float64")
.to(1.1)
.set("f_string")
Expand All @@ -311,6 +317,8 @@ private static Mutation createMutation(Mutation.Op operation) {
.to(Struct.newBuilder().set("int64").to(3L).build())
.set("f_array")
.toInt64Array(ImmutableList.of(2L, 3L))
.set("f_float_array")
.toFloat32Array(ImmutableList.of(3.0f, 4.0f))
.set("f_double_array")
.toFloat64Array(ImmutableList.of(1., 2.))
.set("f_decimal_array")
Expand Down Expand Up @@ -339,8 +347,6 @@ private static Mutation createMutation(Mutation.Op operation) {
.to((short) 2)
.set("f_int32")
.to(0x7fffffff)
.set("f_float")
.to(0.0f)
.set("f_decimal")
.to(BigDecimal.valueOf(Long.MIN_VALUE))
.set("f_byte")
Expand All @@ -355,6 +361,8 @@ private static Mutation createMutationNulls(Mutation.Op operation) {
return builder
.set("f_int64")
.to((Long) null)
.set("f_float32")
.to((Float) null)
.set("f_float64")
.to((Double) null)
.set("f_string")
Expand Down
Loading

0 comments on commit 325c0b5

Please sign in to comment.