private Stream<RowColumnRangeIterator> getBatchForRow(byte[] row, ColumnRangeSelection columnRange, int batchSize) { return kvs.getRowsColumnRange( tableRef, Collections.singleton(row), BatchColumnRangeSelection.create(columnRange, batchSize), SweepQueueUtils.READ_TS).values().stream(); }
public static BatchColumnRangeSelection valueOf(String serialized) { String[] split = deserializeRegex.split(serialized); byte[] startCol = PtBytes.decodeBase64(split[0]); byte[] endCol = PtBytes.decodeBase64(split[1]); int batchHint = Integer.parseInt(split[2]); return BatchColumnRangeSelection.create(startCol, endCol, batchHint); }
public static BatchColumnRangeSelection createPrefixRange(byte[] prefix, int batchSize) { byte[] startCol = Preconditions.checkNotNull(prefix, "prefix cannot be null").clone(); byte[] endCol = RangeRequests.createEndNameForPrefixScan(prefix); return BatchColumnRangeSelection.create(startCol, endCol, batchSize); }
public static RowColumnRangeIterator mergeGetRowsColumnRangeIntoSingleIterator(KeyValueService kvs, TableReference tableRef, Iterable<byte[]> rows, ColumnRangeSelection columnRangeSelection, int batchHint, long timestamp) { if (Iterables.isEmpty(rows)) { return new LocalRowColumnRangeIterator(Collections.emptyIterator()); } int columnBatchSize = Math.max(1, batchHint / Iterables.size(rows)); BatchColumnRangeSelection batchColumnRangeSelection = BatchColumnRangeSelection.create(columnRangeSelection, columnBatchSize); Map<byte[], RowColumnRangeIterator> rowsColumnRanges = kvs.getRowsColumnRange(tableRef, rows, batchColumnRangeSelection, timestamp); // Return results in the same order as the provided rows. Iterable<RowColumnRangeIterator> orderedRanges = Iterables.transform(rows, rowsColumnRanges::get); return new LocalRowColumnRangeIterator(Iterators.concat(orderedRanges.iterator())); } }
@Override public ClosableIterator<Map.Entry<Cell, Value>> getBatch(int batchSize, @Nullable byte[] lastToken) { byte[] startCol = columnRangeSelection.getStartCol(); if (lastToken != null) { startCol = RangeRequests.nextLexicographicName(lastToken); } BatchColumnRangeSelection newRange = BatchColumnRangeSelection.create(startCol, columnRangeSelection.getEndCol(), batchSize); Map<byte[], RowColumnRangeIterator> range = keyValueService.getRowsColumnRange( tableRef, ImmutableList.of(row), newRange, timestamp); if (range.isEmpty()) { return ClosableIterators.wrap(ImmutableList.<Map.Entry<Cell, Value>>of().iterator()); } return ClosableIterators.wrap(Iterables.getOnlyElement(range.values())); }
int getNextId() { RowColumnRangeIterator iterator = kvs.getRowsColumnRange( ID_TO_NAME, Collections.singleton(rowAsBytes), BatchColumnRangeSelection.create(null, null, 1), Long.MAX_VALUE).getOrDefault(rowAsBytes, emptyIterator()); if (!iterator.hasNext()) { return 1; } Map.Entry<Cell, Value> first = iterator.next(); RowResult<byte[]> rowResult = RowResult.of(first.getKey(), first.getValue().getContents()); SweepIdToNameRowResult deserializedRowResult = SweepIdToNameRowResult.of(rowResult); return Ints.checkedCast( Iterables.getOnlyElement(deserializedRowResult.getColumnValues()).getColumnName().getTableId()) + 1; }
private static BatchColumnRangeSelection nextLexicographicalRangeEnd( BatchColumnRangeSelection currentRange, byte[] rangeEnd) { if (rangeEnd.length != 0 && !RangeRequests.isTerminalRow(false, rangeEnd)) { return BatchColumnRangeSelection.create( currentRange.getStartCol(), RangeRequests.getNextStartRow(false, rangeEnd), currentRange.getBatchHint()); } else { return currentRange; } }
@Test public void testMultipleReadsToSameColumnRange() { String rowString = "row1"; byte[] row = PtBytes.toBytes(rowString); byte[] rowDifferentReference = PtBytes.toBytes(rowString); Transaction t1 = startTransaction(); Map<byte[], BatchingVisitable<Map.Entry<Cell, byte[]>>> columnRange = t1.getRowsColumnRange(TEST_TABLE, ImmutableList.of(row), BatchColumnRangeSelection.create(PtBytes.toBytes("col"), PtBytes.toBytes("col0"), 1)); columnRange.values().forEach(visitable -> visitable.batchAccept(10, t -> true)); Map<byte[], BatchingVisitable<Map.Entry<Cell, byte[]>>> columnRangeAgain = t1.getRowsColumnRange(TEST_TABLE, ImmutableList.of(rowDifferentReference), BatchColumnRangeSelection.create(PtBytes.toBytes("col"), PtBytes.toBytes("col0"), 1)); columnRangeAgain.values().forEach(visitable -> visitable.batchAccept(10, t -> true)); put(t1, "mutation to ensure", "conflict", "handling"); t1.commit(); }
@Test public void testMultipleReadsToSameColumnRangeAcrossRows() { byte[] row = PtBytes.toBytes("row"); byte[] differentRow = PtBytes.toBytes("differentRow"); Transaction transaction = startTransaction(); BatchColumnRangeSelection sameColumnRangeSelection = BatchColumnRangeSelection.create(PtBytes.toBytes("col"), PtBytes.toBytes("col0"), 1); Map<byte[], BatchingVisitable<Map.Entry<Cell, byte[]>>> columnRangeResultForRow = transaction.getRowsColumnRange(TEST_TABLE, ImmutableList.of(row), sameColumnRangeSelection); columnRangeResultForRow.values().forEach(visitable -> visitable.batchAccept(10, t -> true)); Map<byte[], BatchingVisitable<Map.Entry<Cell, byte[]>>> columnRangeResultForDifferentRow = transaction.getRowsColumnRange(TEST_TABLE, ImmutableList.of(differentRow), sameColumnRangeSelection); columnRangeResultForDifferentRow.values().forEach(visitable -> visitable.batchAccept(10, t -> true)); put(transaction, "mutation to ensure", "conflict", "handling"); transaction.commit(); }
@Test public void testEmptyColumnRangePagingTransaction() { byte[] row = PtBytes.toBytes("row1"); Transaction t = startTransaction(); Map<byte[], BatchingVisitable<Map.Entry<Cell, byte[]>>> columnRange = t.getRowsColumnRange(TEST_TABLE, ImmutableList.of(row), BatchColumnRangeSelection.create(PtBytes.EMPTY_BYTE_ARRAY, PtBytes.EMPTY_BYTE_ARRAY, 1)); List<Map.Entry<Cell, byte[]>> expected = ImmutableList.of(); verifyMatchingResult(expected, row, columnRange); put(t, "row1", "col1", "v1"); t.commit(); t = startTransaction(); delete(t, "row1", "col1"); columnRange = t.getRowsColumnRange(TEST_TABLE, ImmutableList.of(row), BatchColumnRangeSelection.create(PtBytes.EMPTY_BYTE_ARRAY, PtBytes.EMPTY_BYTE_ARRAY, 1)); verifyMatchingResult(expected, row, columnRange); t.commit(); t = startTransaction(); columnRange = t.getRowsColumnRange(TEST_TABLE, ImmutableList.of(row), BatchColumnRangeSelection.create(PtBytes.EMPTY_BYTE_ARRAY, PtBytes.EMPTY_BYTE_ARRAY, 1)); verifyMatchingResult(expected, row, columnRange); }
@Test public void canGetRowsColumnRange() { BatchColumnRangeSelection rangeSelection = BatchColumnRangeSelection.create(null, null, 1); Map<byte[], RowColumnRangeIterator> rowsColumnRange = getTestKvs() .getRowsColumnRange(TEST_TABLE, ImmutableList.of(FIRST_ROW), rangeSelection, Long.MAX_VALUE); assertThat(Iterables.getOnlyElement(rowsColumnRange.keySet())).isEqualTo(FIRST_ROW); assertThat(rowsColumnRange.get(FIRST_ROW)).containsExactlyElementsOf(expectedRowEntries); }
@Test public void getRowsColumnRangeBatch() throws Exception { RowColumnRangeIterator rowColumnIterator = mock(RowColumnRangeIterator.class); List<byte[]> rows = ImmutableList.of(ROW_NAME); Map<byte[], RowColumnRangeIterator> expectedResult = ImmutableMap.of(ROW_NAME, rowColumnIterator); BatchColumnRangeSelection range = BatchColumnRangeSelection.create(COL_NAME, COL_NAME, 2); when(delegate.getRowsColumnRange(TABLE_REF, rows, range, TIMESTAMP)).thenReturn(expectedResult); Map<byte[], RowColumnRangeIterator> result = kvs.getRowsColumnRange(TABLE_REF, rows, range, TIMESTAMP); assertThat(result, equalTo(expectedResult)); checkSpan("atlasdb-kvs.getRowsColumnRange({table}, 1 rows, ts 1)"); verify(delegate).getRowsColumnRange(TABLE_REF, rows, range, TIMESTAMP); verifyNoMoreInteractions(delegate); }
@Test public void testGetRowColumnRangeHistorical() { putTestDataForMultipleTimestamps(); Map<byte[], RowColumnRangeIterator> values = keyValueService.getRowsColumnRange(TEST_TABLE, ImmutableList.of(row(0)), BatchColumnRangeSelection.create(PtBytes.EMPTY_BYTE_ARRAY, PtBytes.EMPTY_BYTE_ARRAY, 1), TEST_TIMESTAMP + 2); assertEquals(1, values.size()); Map<Cell, Value> batchValues = getValuesForRow(values, row(0), 1); assertEquals(1, batchValues.size()); assertArrayEquals(val(0, 7), batchValues.get(TEST_CELL).getContents()); values = keyValueService.getRowsColumnRange(TEST_TABLE, ImmutableList.of(row(0)), BatchColumnRangeSelection.create(PtBytes.EMPTY_BYTE_ARRAY, PtBytes.EMPTY_BYTE_ARRAY, 1), TEST_TIMESTAMP + 1); assertEquals(1, values.size()); batchValues = getValuesForRow(values, row(0), 1); assertEquals(1, batchValues.size()); assertArrayEquals(val(0, 5), batchValues.get(TEST_CELL).getContents()); }
@Test public void notAllowBatchColumnRangeGets() { checkThrowsAndNoInteraction(() -> readTransaction.getRowsColumnRange( DUMMY_THOROUGH_TABLE, ImmutableList.of(EMPTY_BYTES), BatchColumnRangeSelection.create(EMPTY_BYTES, EMPTY_BYTES, 1)), IllegalStateException.class, "Cannot read"); }
@Test public void testColumnRangeReadWriteNoConflict() { byte[] row = PtBytes.toBytes("row1"); writeColumns(); Transaction t1 = startTransaction(); Map<byte[], BatchingVisitable<Map.Entry<Cell, byte[]>>> columnRange = t1.getRowsColumnRange(TEST_TABLE, ImmutableList.of(row), BatchColumnRangeSelection.create(PtBytes.EMPTY_BYTE_ARRAY, PtBytes.EMPTY_BYTE_ARRAY, 1)); // Serializable transaction records only the first column as read. Map.Entry<Cell, byte[]> read = BatchingVisitables.getFirst(Iterables.getOnlyElement(columnRange.values())); assertEquals(Cell.create(row, PtBytes.toBytes("col0")), read.getKey()); // Write to avoid the read only path. put(t1, "row1_1", "col0", "v0"); Transaction t2 = startTransaction(); put(t2, "row1", "col1", "v0_0"); t2.commit(); t1.commit(); }
@Test public void getRowsColumnRangeThrows() { BatchColumnRangeSelection rangeSelection = BatchColumnRangeSelection.create(null, null, 1); assertThrowsAtlasDbDependencyExceptionAndDoesNotChangeCassandraSchema(() -> getTestKvs() .getRowsColumnRange(TEST_TABLE, ImmutableList.of(FIRST_ROW), rangeSelection, Long.MAX_VALUE)); }
@Test public void testColumnRangeReadWriteEmptyRange() { byte[] row = PtBytes.toBytes("row1"); Transaction t1 = startTransaction(); Map<byte[], BatchingVisitable<Map.Entry<Cell, byte[]>>> columnRange = t1.getRowsColumnRange(TEST_TABLE, ImmutableList.of(row), BatchColumnRangeSelection.create(PtBytes.toBytes("col"), PtBytes.toBytes("col0"), 1)); assertNull(BatchingVisitables.getFirst(Iterables.getOnlyElement(columnRange.values()))); // Write to avoid the read only path. put(t1, "row1_1", "col0", "v0"); Transaction t2 = startTransaction(); put(t2, "row1", "col", "v0"); t2.commit(); try { t1.commit(); fail(); } catch (TransactionSerializableConflictException e) { // expected } }
@Test public void testColumnRangeReadWriteEmptyRangeUnread() { byte[] row = PtBytes.toBytes("row1"); Transaction t1 = startTransaction(); Map<byte[], BatchingVisitable<Map.Entry<Cell, byte[]>>> columnRange = t1.getRowsColumnRange(TEST_TABLE, ImmutableList.of(row), BatchColumnRangeSelection.create(PtBytes.toBytes("col"), PtBytes.toBytes("col0"), 1)); // Intentionally not reading anything from the result, so we shouldn't get a conflict. // Write to avoid the read only path. put(t1, "row1_1", "col0", "v0"); Transaction t2 = startTransaction(); put(t2, "row1", "col", "v0"); t2.commit(); t1.commit(); }
@Test public void testGetRowColumnRangeMultipleRows() { putTestDataForSingleTimestamp(); Map<byte[], RowColumnRangeIterator> values = keyValueService.getRowsColumnRange(TEST_TABLE, ImmutableList.of(row(1), row(0), row(2)), BatchColumnRangeSelection.create(PtBytes.EMPTY_BYTE_ARRAY, PtBytes.EMPTY_BYTE_ARRAY, 1), TEST_TIMESTAMP + 1); assertEquals(3, values.keySet().size()); Map<Cell, Value> row0Values = getValuesForRow(values, row(0), 2); assertArrayEquals(val(0, 0), row0Values.get(TEST_CELL).getContents()); assertArrayEquals(val(0, 1), row0Values.get(Cell.create(row(0), column(1))).getContents()); Map<Cell, Value> row1Values = getValuesForRow(values, row(1), 2); assertArrayEquals(val(1, 0), row1Values.get(Cell.create(row(1), column(0))).getContents()); assertArrayEquals(val(1, 2), row1Values.get(Cell.create(row(1), column(2))).getContents()); Map<Cell, Value> row2Values = getValuesForRow(values, row(2), 2); assertArrayEquals(val(2, 1), row2Values.get(Cell.create(row(2), column(1))).getContents()); assertArrayEquals(val(2, 2), row2Values.get(Cell.create(row(2), column(2))).getContents()); }
private RowsColumnRangeBatchRequest createRequest(int numTotalRows) { ColumnRangeSelection fullColumnRange = new ColumnRangeSelection(col(0), col(5)); ImmutableRowsColumnRangeBatchRequest.Builder request = ImmutableRowsColumnRangeBatchRequest.builder().columnRangeSelection(fullColumnRange); if (hasPartialFirstRow) { request.partialFirstRow(Maps.immutableEntry(row(0), BatchColumnRangeSelection.create(col(3), col(5), 10))); } int firstFullRowIndex = hasPartialFirstRow ? 2 : 1; int lastFullRowIndex = hasPartialLastRow ? numTotalRows - 1 : numTotalRows; for (int rowNum = firstFullRowIndex; rowNum <= lastFullRowIndex; rowNum++) { request.addRowsToLoadFully(row(rowNum)); } if (hasPartialLastRow) { request.partialLastRow( Maps.immutableEntry(row(numTotalRows), BatchColumnRangeSelection.create(fullColumnRange, 10))); } return request.build(); }