RowMerger implements StreamObserver<ReadRowsResponse> { @Override public final void onNext(ReadRowsResponse readRowsResponse) { if (complete) { onError(new IllegalStateException("Adding partialRow after completion")); return; } int rowsProcessed = 0; for (int i = 0; i < readRowsResponse.getChunksCount(); i++) { try { CellChunk chunk = readRowsResponse.getChunks(i); state.validateChunk(rowInProgress, lastCompletedRowKey, chunk); if (chunk.getResetRow()) { rowInProgress = null; state = RowMergerState.NewRow; continue; } if (state == RowMergerState.NewRow) { rowInProgress = new RowInProgress(); rowInProgress.updateCurrentKey(chunk); } else if (state == RowMergerState.RowInProgress) { rowInProgress.updateCurrentKey(chunk); } if (chunk.getValueSize() > 0) { rowInProgress.addPartialCellChunk(chunk); state = RowMergerState.CellInProgress; } else if (rowInProgress.hasChunkInProgess()) { rowInProgress.addPartialCellChunk(chunk); rowInProgress.completeMultiChunkCell(); state = RowMergerState.RowInProgress; } else { rowInProgress.addFullChunk(chunk); state = RowMergerState.RowInProgress; } if (chunk.getCommitRow()) { observer.onNext(rowInProgress.buildRow()); lastCompletedRowKey = rowInProgress.getRowKey(); state = RowMergerState.NewRow; rowInProgress = null; rowsProcessed++; } } catch (Throwable e) { onError(e); return; } } this.rowCountInLastMessage = rowsProcessed; } RowMerger(StreamObserver<FlatRow> observer); static List<FlatRow> toRows(Iterable<ReadRowsResponse> responses); void clearRowInProgress(); @Override final void onNext(ReadRowsResponse readRowsResponse); Integer getRowCountInLastMessage(); ByteString getLastCompletedRowKey(); @Override void onCompleted(); @Override void onError(Throwable e);  }
@Test public void testThreeCellRow() { CellChunk cellChunk1 = createCell("row_key1", "family", "qualifier", "value", 1, false); CellChunk cellChunk2 = createCell(null, null, "qualifier2", "value2", 2, false); CellChunk cellChunk3 = createCell(null, null, null, "value3", 1, true); RowMerger underTest = new RowMerger(observer); underTest.onNext(ReadRowsResponse.newBuilder().addChunks(cellChunk1).build()); underTest.onNext(ReadRowsResponse.newBuilder().addChunks(cellChunk2).build()); underTest.onNext(ReadRowsResponse.newBuilder().addChunks(cellChunk3).build()); FlatRow expected = FlatRow.newBuilder() .withRowKey(cellChunk1.getRowKey()) .addCell( cellChunk1.getFamilyName().getValue(), cellChunk1.getQualifier().getValue(), cellChunk1.getTimestampMicros(), cellChunk1.getValue()) .addCell( cellChunk1.getFamilyName().getValue(), cellChunk2.getQualifier().getValue(), cellChunk2.getTimestampMicros(), cellChunk2.getValue()) .addCell( cellChunk1.getFamilyName().getValue(), cellChunk2.getQualifier().getValue(), cellChunk3.getTimestampMicros(), cellChunk3.getValue()) .build(); verify(observer, times(1)).onNext(eq(expected)); }
@Test public void testMultiChunkCell() { CellChunk cellChunk1 = createCell("row_key1", "family", "qualifier", "v", 1, false) .toBuilder() .setValueSize("value".length()) .build(); CellChunk cellChunk2 = createCell(null, null, null, "alue", 1, true); RowMerger underTest = new RowMerger(observer); underTest.onNext(ReadRowsResponse.newBuilder().addChunks(cellChunk1).build()); underTest.onNext(ReadRowsResponse.newBuilder().addChunks(cellChunk2).build()); ByteString value = ByteString.copyFrom("value".getBytes()); FlatRow row = FlatRow.newBuilder() .withRowKey(cellChunk1.getRowKey()) .addCell( cellChunk1.getFamilyName().getValue(), cellChunk1.getQualifier().getValue(), cellChunk1.getTimestampMicros(), value) .build(); verify(observer, times(1)).onNext(eq(row)); }
@Test public void testSimpleReset() { CellChunk cellChunk1 = createCell("row_key1", "family", "qualifier", "value", 1, false); CellChunk cellChunk2 = createCell("row_key2", "family", "qualifier", "value", 1, true); RowMerger underTest = new RowMerger(observer); underTest.onNext(ReadRowsResponse.newBuilder().addChunks(cellChunk1).build()); underTest.onNext( ReadRowsResponse.newBuilder().addChunks(CellChunk.newBuilder().setResetRow(true)).build()); underTest.onNext(ReadRowsResponse.newBuilder().addChunks(cellChunk2).build()); verify(observer, times(0)).onNext(eq(toRow(cellChunk1))); verify(observer, times(1)).onNext(eq(toRow(cellChunk2))); }
@Test public void testVariousConditions() { RowMerger underTest = new RowMerger(observer); String rowKey = "row_key1"; String family1 = "family1"; String family2 = "family2"; String qualifier1 = "qualifier1"; String qualifier2 = "qualifier2"; String value1 = "value1"; String value2 = "value2"; String value3 = "value3"; String value4 = "value4"; String value5 = "value5"; ByteString rowKeyByteString = toByteString(rowKey); ByteString qualifier1ByteString = toByteString(qualifier1); ByteString qualifier2ByteString = toByteString(qualifier2); ByteString value1ByteString = toByteString(value1); ByteString value2ByteString = toByteString(value2); ByteString value3ByteString = toByteString(value3); ByteString value4ByteString = toByteString(value4); ByteString value5ByteString = toByteString(value5); final long ts1 = 54321L; final long ts2 = 12345L; final String label = "label"; underTest.onNext( ReadRowsResponse.newBuilder() .addAllChunks( Arrays.asList( createCell(rowKey, family1, qualifier1, value1, ts1, false), createCell(rowKey, family1, qualifier1, value1, ts1, false), createCell(rowKey, family1, qualifier1, value2, ts2, false), createCell(rowKey, family1, qualifier1, value2, ts2, false, label), createCell(rowKey, family1, qualifier2, value3, ts1, false), createCell(rowKey, family2, qualifier1, value4, ts1, false), createCell(rowKey, family2, qualifier2, value5, ts1, true))) .build()); verify(observer, times(1)) .onNext( eq( FlatRow.newBuilder() .withRowKey(rowKeyByteString) .addCell(family1, qualifier1ByteString, ts1, value1ByteString) .addCell(family1, qualifier1ByteString, ts2, value2ByteString) .addCell( family1, qualifier1ByteString, ts2, value2ByteString, Arrays.asList(label)) .addCell(family1, qualifier2ByteString, ts1, value3ByteString) .addCell(family2, qualifier1ByteString, ts1, value4ByteString) .addCell(family2, qualifier2ByteString, ts1, value5ByteString) .build())); }