perf: decode BYTES columns lazily

BYTES columns are encoded as Base64 strings. Decoding these are relatively CPU-heavy,
especially for large values. Decoding them is not always necessary if the user only
needs the Base64 string. Also, the internally used Guava decoder is less efficient
than JDK implementations that are available from Java 8 and onwards. This change
therefore delays the decoding of BYTES columns until it is actually necessary, and
then uses the JDK implementation instead of the Guava version. The JDK implementation
in OpenJDK 17 uses approx 1/3 of the CPU cycles of the Guava version.

Author: olavloite

google-cloud-spanner/src/main/java/com/google/cloud/ByteArrayHelper.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2023 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.cloud;
+
+import com.google.api.core.InternalApi;
+import com.google.protobuf.ByteStringHelper;
+
+@InternalApi
+public class ByteArrayHelper {
+  @InternalApi
+  public static ByteArray wrap(byte[] bytes) {
+    return new ByteArray(ByteStringHelper.wrap(bytes));
+  }
+}

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java

@@ -26,6 +26,7 @@
 import com.google.api.client.util.ExponentialBackOff;
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.cloud.ByteArray;
+import com.google.cloud.ByteArrayHelper;
 import com.google.cloud.Date;
 import com.google.cloud.Timestamp;
 import com.google.cloud.spanner.Type.StructField;
@@ -55,6 +56,7 @@
 import java.math.BigDecimal;
 import java.util.AbstractList;
 import java.util.ArrayList;
+import java.util.Base64;
 import java.util.BitSet;
 import java.util.Collections;
 import java.util.Iterator;
@@ -353,6 +355,39 @@ private boolean isMergeable(KindCase kind) {
     }
   }
 
+  static class LazyByteArray implements Serializable {
+    private static final Base64.Decoder DECODER = Base64.getDecoder();
+    private final String base64String;
+    private final transient AbstractLazyInitializer<ByteArray> byteArray =
+        new AbstractLazyInitializer<ByteArray>() {
+          @Override
+          protected ByteArray initialize() {
+            return ByteArrayHelper.wrap(DECODER.decode(base64String));
+          }
+        };
+
+    LazyByteArray(String base64String) {
+      this.base64String = base64String;
+    }
+
+    ByteArray getByteArray() {
+      try {
+        return byteArray.get();
+      } catch (Throwable t) {
+        throw SpannerExceptionFactory.asSpannerException(t);
+      }
+    }
+
+    String getBase64String() {
+      return base64String;
+    }
+
+    @Override
+    public String toString() {
+      return getBase64String();
+    }
+  }
+
   static class GrpcStruct extends Struct implements Serializable {
     private final Type type;
     private final List<Object> rowData;
@@ -395,7 +430,12 @@ private Object writeReplace() {
             builder.set(fieldName).to(Value.pgJsonb((String) value));
             break;
           case BYTES:
-            builder.set(fieldName).to((ByteArray) value);
+            builder
+                .set(fieldName)
+                .to(
+                    value instanceof LazyByteArray
+                        ? ((LazyByteArray) value).getByteArray()
+                        : (ByteArray) value);
             break;
           case TIMESTAMP:
             builder.set(fieldName).to((Timestamp) value);
@@ -511,7 +551,7 @@ private static Object decodeValue(Type fieldType, com.google.protobuf.Value prot
           return proto.getStringValue();
         case BYTES:
           checkType(fieldType, proto, KindCase.STRING_VALUE);
-          return ByteArray.fromBase64(proto.getStringValue());
+          return new LazyByteArray(proto.getStringValue());
         case TIMESTAMP:
           checkType(fieldType, proto, KindCase.STRING_VALUE);
           return Timestamp.parseTimestamp(proto.getStringValue());
@@ -619,7 +659,8 @@ protected BigDecimal getBigDecimalInternal(int columnIndex) {
 
     @Override
     protected String getStringInternal(int columnIndex) {
-      return (String) rowData.get(columnIndex);
+      Object value = rowData.get(columnIndex);
+      return value == null ? null : value.toString();
     }
 
     @Override
@@ -634,7 +675,7 @@ protected String getPgJsonbInternal(int columnIndex) {
 
     @Override
     protected ByteArray getBytesInternal(int columnIndex) {
-      return (ByteArray) rowData.get(columnIndex);
+      return ((LazyByteArray) rowData.get(columnIndex)).getByteArray();
     }
 
     @Override
@@ -785,9 +826,10 @@ protected List<String> getPgJsonbListInternal(int columnIndex) {
     }
 
     @Override
-    @SuppressWarnings("unchecked") // We know ARRAY<BYTES> produces a List<ByteArray>.
+    @SuppressWarnings("unchecked") // We know ARRAY<BYTES> produces a List<LazyByteArray>.
     protected List<ByteArray> getBytesListInternal(int columnIndex) {
-      return Collections.unmodifiableList((List<ByteArray>) rowData.get(columnIndex));
+      return Lists.transform(
+          (List<LazyByteArray>) rowData.get(columnIndex), l -> l == null ? null : l.getByteArray());
     }
 
     @Override

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractStructReader.java

@@ -170,17 +170,20 @@ public BigDecimal getBigDecimal(String columnName) {
   public String getString(int columnIndex) {
     checkNonNullOfTypes(
         columnIndex,
-        Arrays.asList(Type.string(), Type.pgNumeric()),
+        Arrays.asList(Type.string(), Type.pgNumeric(), Type.bytes()),
         columnIndex,
-        "STRING, NUMERIC");
+        "STRING, NUMERIC, BYTES");
     return getStringInternal(columnIndex);
   }
 
   @Override
   public String getString(String columnName) {
     int columnIndex = getColumnIndex(columnName);
     checkNonNullOfTypes(
-        columnIndex, Arrays.asList(Type.string(), Type.pgNumeric()), columnName, "STRING, NUMERIC");
+        columnIndex,
+        Arrays.asList(Type.string(), Type.pgNumeric(), Type.bytes()),
+        columnName,
+        "STRING, NUMERIC, BYTES");
     return getStringInternal(columnIndex);
   }
 

google-cloud-spanner/src/main/java/com/google/protobuf/ByteStringHelper.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2023 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.protobuf;
+
+import com.google.api.core.InternalApi;
+
+@InternalApi
+public class ByteStringHelper {
+
+  @InternalApi
+  public static ByteString wrap(byte[] bytes) {
+    return ByteString.wrap(bytes);
+  }
+}

google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractStructReaderTypesTest.java

@@ -27,6 +27,7 @@
 import com.google.cloud.Date;
 import com.google.cloud.Timestamp;
 import com.google.common.base.Throwables;
+import com.google.common.collect.ImmutableList;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.math.BigDecimal;
@@ -215,7 +216,7 @@ public static Collection<Object[]> parameters() {
             "getStringInternal",
             "1.23",
             "getString",
-            Collections.singletonList("getValue")
+            ImmutableList.of("getValue", "getBigDecimal")
           },
           {
             Type.string(),
@@ -229,7 +230,7 @@ public static Collection<Object[]> parameters() {
             "getBytesInternal",
             ByteArray.copyFrom(new byte[] {0}),
             "getBytes",
-            Collections.singletonList("getValue")
+            ImmutableList.of("getValue", "getString")
           },
           {
             Type.json(),

google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java

@@ -37,6 +37,7 @@
 import com.google.api.core.ApiFutures;
 import com.google.api.gax.grpc.testing.LocalChannelProvider;
 import com.google.api.gax.retrying.RetrySettings;
+import com.google.cloud.ByteArray;
 import com.google.cloud.NoCredentials;
 import com.google.cloud.Timestamp;
 import com.google.cloud.spanner.AbstractResultSet.GrpcStreamIterator;
@@ -54,6 +55,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.SettableFuture;
 import com.google.protobuf.AbstractMessage;
+import com.google.protobuf.ListValue;
 import com.google.spanner.v1.CommitRequest;
 import com.google.spanner.v1.DeleteSessionRequest;
 import com.google.spanner.v1.ExecuteBatchDmlRequest;
@@ -75,8 +77,10 @@
 import io.grpc.inprocess.InProcessServerBuilder;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Base64;
 import java.util.Collections;
 import java.util.List;
+import java.util.Random;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
@@ -2381,4 +2385,51 @@ public void testAnalyzeUpdateStatement() {
     ExecuteSqlRequest request = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
     assertEquals(QueryMode.PLAN, request.getQueryMode());
   }
+
+  @Test
+  public void testByteArray() {
+    Random random = new Random();
+    byte[] bytes = new byte[random.nextInt(20_000)];
+    int numRows = 10;
+    List<ListValue> rows = new ArrayList<>(numRows);
+    for (int i = 0; i < numRows; i++) {
+      random.nextBytes(bytes);
+      rows.add(
+          ListValue.newBuilder()
+              .addValues(
+                  com.google.protobuf.Value.newBuilder()
+                      .setStringValue(Base64.getEncoder().encodeToString(bytes))
+                      .build())
+              .build());
+    }
+    Statement statement = Statement.of("select * from foo");
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            statement,
+            com.google.spanner.v1.ResultSet.newBuilder()
+                .setMetadata(
+                    ResultSetMetadata.newBuilder()
+                        .setRowType(
+                            StructType.newBuilder()
+                                .addFields(
+                                    Field.newBuilder()
+                                        .setType(Type.newBuilder().setCode(TypeCode.BYTES).build())
+                                        .setName("f1")
+                                        .build())
+                                .build())
+                        .build())
+                .addAllRows(rows)
+                .build()));
+    DatabaseClient client =
+        spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+    try (ResultSet resultSet = client.singleUse().executeQuery(statement)) {
+      while (resultSet.next()) {
+        String base64String = resultSet.getString(0);
+        ByteArray byteArray = resultSet.getBytes(0);
+        // Use the 'old' ByteArray.fromBase64(..) method that uses the Guava encoder to ensure that
+        // the two encoders (JDK and Guava) return the same values.
+        assertEquals(ByteArray.fromBase64(base64String), byteArray);
+      }
+    }
+  }
 }