Bigtable: make row & cell ordering more consistent.

google-cloud-clients/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/DefaultRowAdapter.java

@@ -18,6 +18,8 @@
 import com.google.common.collect.ImmutableList;
 import com.google.protobuf.ByteString;
 import java.util.List;
+import java.util.Objects;
+import java.util.TreeMap;
 
 /**
  * Default implementation of a {@link RowAdapter} that uses {@link Row}s to represent logical rows.
@@ -44,13 +46,21 @@ public ByteString getKey(Row row) {
   /** {@inheritDoc} */
   public class DefaultRowBuilder implements RowBuilder<Row> {
     private ByteString currentKey;
-    private ImmutableList.Builder<RowCell> cells;
+    private TreeMap<String, ImmutableList.Builder<RowCell>> cellsByFamily;
+    private ImmutableList.Builder<RowCell> currentFamilyCells;
+    private String previousFamily;
+    private int totalCellCount;
+
     private String family;
     private ByteString qualifier;
     private List<String> labels;
     private long timestamp;
     private ByteString value;
 
+    public DefaultRowBuilder() {
+      reset();
+    }
+
     /** {@inheritDoc} */
     @Override
     public Row createScanMarkerRow(ByteString key) {
@@ -61,7 +71,6 @@ public Row createScanMarkerRow(ByteString key) {
     @Override
     public void startRow(ByteString key) {
       currentKey = key;
-      cells = ImmutableList.builder();
     }
 
     /** {@inheritDoc} */
@@ -84,20 +93,52 @@ public void cellValue(ByteString value) {
     /** {@inheritDoc} */
     @Override
     public void finishCell() {
-      cells.add(RowCell.create(family, qualifier, timestamp, labels, value));
+      if (!Objects.equals(family, previousFamily)) {
+        previousFamily = family;
+        currentFamilyCells = ImmutableList.builder();
+        cellsByFamily.put(family, currentFamilyCells);
+      }
+
+      RowCell rowCell = RowCell.create(family, qualifier, timestamp, labels, value);
+      currentFamilyCells.add(rowCell);
+      totalCellCount++;
     }
 
     /** {@inheritDoc} */
     @Override
     public Row finishRow() {
-      return Row.create(currentKey, cells.build());
+      final ImmutableList<RowCell> sortedCells;
+
+      // Optimization: If there are no cells, then just return the static empty list.
+      if (cellsByFamily.size() == 0) {
+        sortedCells = ImmutableList.of();
+      } else if (cellsByFamily.size() == 1) {
+        // Optimization: If there is a single family, avoid copies and return that one list.
+        sortedCells = currentFamilyCells.build();
+      } else {
+        // Normal path: concatenate the cells order by family.
+        ImmutableList.Builder<RowCell> sortedCellsBuilder =
+            ImmutableList.builderWithExpectedSize(totalCellCount);
+
+        for (ImmutableList.Builder<RowCell> familyCells : cellsByFamily.values()) {
+          sortedCellsBuilder.addAll(familyCells.build());
+        }
+        sortedCells = sortedCellsBuilder.build();
+      }
+
+      return Row.create(currentKey, sortedCells);
     }
 
     /** {@inheritDoc} */
     @Override
     public void reset() {
       currentKey = null;
-      cells = null;
+
+      cellsByFamily = new TreeMap<>();
+      currentFamilyCells = null;
+      previousFamily = null;
+      totalCellCount = 0;
+
       family = null;
       qualifier = null;
       labels = null;

google-cloud-clients/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Row.java

@@ -18,15 +18,38 @@
 import com.google.api.core.InternalApi;
 import com.google.api.core.InternalExtensionOnly;
 import com.google.auto.value.AutoValue;
+import com.google.cloud.bigtable.data.v2.internal.ByteStringComparator;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
 import com.google.protobuf.ByteString;
 import java.io.Serializable;
+import java.util.Comparator;
 import java.util.List;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 
-/** Default representation of a logical row. */
+/**
+ * Default representation of a logical row.
+ *
+ * <p>The cells contained within, will be sorted by the native order. Please see
+ * {@link RowCell#compareByNative()} for details.
+ */
 @InternalExtensionOnly
 @AutoValue
-public abstract class Row implements Comparable<Row>, Serializable {
+public abstract class Row implements Serializable {
+  /**
+   * Returns a comparator that compares two Row objects by comparing the result of {@link
+   * #getKey()}} for each.
+   */
+  public static Comparator<Row> compareByKey() {
+    return new Comparator<Row>() {
+      @Override
+      public int compare(Row r1, Row r2) {
+        return ByteStringComparator.INSTANCE.compare(r1.getKey(), r2.getKey());
+      }
+    };
+  }
+
   /** Creates a new instance of the {@link Row}. */
   @InternalApi
   public static Row create(ByteString key, List<RowCell> cells) {
@@ -38,30 +61,110 @@ public static Row create(ByteString key, List<RowCell> cells) {
   public abstract ByteString getKey();
 
   /**
-   * Returns the list of cells. The cells will be clustered by their family and sorted by their
-   * qualifier.
+   * Returns a sorted list of cells. The cells will be sorted natively.
+   *
+   * @see RowCell#compareByNative() For details about the ordering.
    */
   public abstract List<RowCell> getCells();
 
-  /** Lexicographically compares this row's key to another row's key. */
-  @Override
-  public int compareTo(@Nonnull Row row) {
-    int sizeA = getKey().size();
-    int sizeB = row.getKey().size();
-    int size = Math.min(sizeA, sizeB);
-
-    for (int i = 0; i < size; i++) {
-      int byteA = getKey().byteAt(i) & 0xff;
-      int byteB = row.getKey().byteAt(i) & 0xff;
-      if (byteA == byteB) {
-        continue;
+  /**
+   * Returns a sublist of the cells that belong to the specified family.
+   *
+   * @see RowCell#compareByNative() For details about the ordering.
+   */
+  public List<RowCell> getCells(@Nonnull String family) {
+    Preconditions.checkNotNull(family, "family");
+
+    int start = getFirst(family, null);
+    if (start < 0) {
+      return ImmutableList.of();
+    }
+
+    int end = getLast(family, null, start);
+
+    return getCells().subList(start, end + 1);
+  }
+
+  /**
+   * Returns a sublist of the cells that belong to the specified family and qualifier.
+   *
+   * @see RowCell#compareByNative() For details about the ordering.
+   */
+  public List<RowCell> getCells(@Nonnull String family, @Nonnull String qualifier) {
+    Preconditions.checkNotNull(family, "family");
+    Preconditions.checkNotNull(qualifier, "qualifier");
+
+    return getCells(family, ByteString.copyFromUtf8(qualifier));
+  }
+
+  /**
+   * Returns a sublist of the cells that belong to the specified family and qualifier.
+   *
+   * @see RowCell#compareByNative() For details about the ordering.
+   */
+  public List<RowCell> getCells(@Nonnull String family, @Nonnull ByteString qualifier) {
+    Preconditions.checkNotNull(family, "family");
+    Preconditions.checkNotNull(qualifier, "qualifier");
+
+    int start = getFirst(family, qualifier);
+    if (start < 0) {
+      return ImmutableList.of();
+    }
+
+    int end = getLast(family, qualifier, start);
+
+    return getCells().subList(start, end + 1);
+  }
+
+  private int getFirst(@Nonnull String family, @Nullable ByteString qualifier) {
+    int low = 0;
+    int high = getCells().size();
+    int index = -1;
+
+    while (low < high) {
+      int mid = (high + low) / 2;
+      RowCell midCell = getCells().get(mid);
+
+      int c = midCell.getFamily().compareTo(family);
+      if (c == 0 && qualifier != null) {
+        c = ByteStringComparator.INSTANCE.compare(midCell.getQualifier(), qualifier);
+      }
+
+      if (c < 0) {
+        low = mid + 1;
+      } else if (c == 0) {
+        index = mid;
+        high = mid;
       } else {
-        return byteA < byteB ? -1 : 1;
+        high = mid;
       }
     }
-    if (sizeA == sizeB) {
-      return 0;
+    return index;
+  }
+
+  private int getLast(@Nonnull String family, @Nullable ByteString qualifier, int startIndex) {
+    int low = startIndex;
+    int high = getCells().size();
+    int index = -1;
+
+    while (low < high) {
+      int mid = (high + low) / 2;
+      RowCell midCell = getCells().get(mid);
+
+      int c = midCell.getFamily().compareTo(family);
+      if (c == 0 && qualifier != null) {
+        c = ByteStringComparator.INSTANCE.compare(midCell.getQualifier(), qualifier);
+      }
+
+      if (c < 0) {
+        low = mid + 1;
+      } else if (c == 0) {
+        index = mid;
+        low = mid + 1;
+      } else {
+        high = mid;
+      }
     }
-    return sizeA < sizeB ? -1 : 1;
+    return index;
   }
 }

google-cloud-clients/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/RowCell.java

@@ -18,15 +18,38 @@
 import com.google.api.core.InternalApi;
 import com.google.api.core.InternalExtensionOnly;
 import com.google.auto.value.AutoValue;
+import com.google.cloud.bigtable.data.v2.internal.ByteStringComparator;
+import com.google.common.collect.ComparisonChain;
 import com.google.protobuf.ByteString;
 import java.io.Serializable;
+import java.util.Comparator;
 import java.util.List;
 import javax.annotation.Nonnull;
 
 /** Default representation of a cell in a {@link Row}. */
 @InternalExtensionOnly
 @AutoValue
 public abstract class RowCell implements Serializable {
+  /**
+   * A comparator that compares the cells by Bigtable native ordering.
+   *
+   * <p>family lexicographically, then by qualifier
+   * lexicographically and finally by timestamp in reverse chronological order. Labels and values
+   * are not included in the comparison.
+   */
+  public static Comparator<RowCell> compareByNative() {
+    return new Comparator<RowCell>() {
+      @Override
+      public int compare(RowCell c1, RowCell c2) {
+        return ComparisonChain.start()
+            .compare(c1.getFamily(), c2.getFamily())
+            .compare(c1.getQualifier(), c2.getQualifier(), ByteStringComparator.INSTANCE)
+            .compare(c2.getTimestamp(), c1.getTimestamp())
+            .result();
+      }
+    };
+  }
+
   /** Creates a new instance of the {@link RowCell}. */
   @InternalApi
   public static RowCell create(

google-cloud-clients/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/ReadModifyWriteRowCallable.java

@@ -26,22 +26,25 @@
 import com.google.bigtable.v2.ReadModifyWriteRowRequest;
 import com.google.bigtable.v2.ReadModifyWriteRowResponse;
 import com.google.cloud.bigtable.data.v2.internal.RequestContext;
+import com.google.cloud.bigtable.data.v2.models.DefaultRowAdapter;
 import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
 import com.google.cloud.bigtable.data.v2.models.Row;
-import com.google.cloud.bigtable.data.v2.models.RowCell;
-import com.google.common.collect.ImmutableList;
+import com.google.cloud.bigtable.data.v2.models.RowAdapter;
+import com.google.cloud.bigtable.data.v2.models.RowAdapter.RowBuilder;
 import com.google.common.util.concurrent.MoreExecutors;
 
 /** Simple wrapper for ReadModifyWriteRow to wrap the request and response protobufs. */
 class ReadModifyWriteRowCallable extends UnaryCallable<ReadModifyWriteRow, Row> {
   private final UnaryCallable<ReadModifyWriteRowRequest, ReadModifyWriteRowResponse> inner;
   private final RequestContext requestContext;
+  private final RowAdapter<Row> rowAdapter;
 
   ReadModifyWriteRowCallable(
       UnaryCallable<ReadModifyWriteRowRequest, ReadModifyWriteRowResponse> inner,
       RequestContext requestContext) {
     this.inner = inner;
     this.requestContext = requestContext;
+    this.rowAdapter = new DefaultRowAdapter();
   }
 
   @Override
@@ -61,21 +64,25 @@ public Row apply(ReadModifyWriteRowResponse readModifyWriteRowResponse) {
   }
 
   private Row convertResponse(ReadModifyWriteRowResponse response) {
-    ImmutableList.Builder<RowCell> cells = ImmutableList.builder();
+    RowBuilder<Row> rowBuilder = rowAdapter.createRowBuilder();
+    rowBuilder.startRow(response.getRow().getKey());
 
     for (Family family : response.getRow().getFamiliesList()) {
       for (Column column : family.getColumnsList()) {
         for (Cell cell : column.getCellsList()) {
-          cells.add(
-              RowCell.create(
-                  family.getName(),
-                  column.getQualifier(),
-                  cell.getTimestampMicros(),
-                  cell.getLabelsList(),
-                  cell.getValue()));
+          rowBuilder.startCell(
+              family.getName(),
+              column.getQualifier(),
+              cell.getTimestampMicros(),
+              cell.getLabelsList(),
+              cell.getValue().size());
+
+          rowBuilder.cellValue(cell.getValue());
+
+          rowBuilder.finishCell();
         }
       }
     }
-    return Row.create(response.getRow().getKey(), cells.build());
+    return rowBuilder.finishRow();
   }
 }