Examples with ColumnSelection com.palantir.atlasdb.keyvalue.api.ColumnSelection used on opensource projects

Example 1 with ColumnSelection

use of com.palantir.atlasdb.keyvalue.api.ColumnSelection in project atlasdb by palantir.

the class CachingTransactionTest method testGetRows.

@Test
public void testGetRows() {
    final Set<byte[]> oneRow = ImmutableSortedSet.orderedBy(PtBytes.BYTES_COMPARATOR).add(ROW_BYTES).build();
    final ColumnSelection oneColumn = ColumnSelection.create(ImmutableList.of(COL_BYTES));
    final Set<byte[]> noRows = ImmutableSortedSet.orderedBy(PtBytes.BYTES_COMPARATOR).build();
    final SortedMap<byte[], RowResult<byte[]>> emptyResults = ImmutableSortedMap.<byte[], RowResult<byte[]>>orderedBy(PtBytes.BYTES_COMPARATOR).build();
    final RowResult<byte[]> rowResult = RowResult.of(Cell.create(ROW_BYTES, COL_BYTES), VALUE_BYTES);
    final SortedMap<byte[], RowResult<byte[]>> oneResult = ImmutableSortedMap.<byte[], RowResult<byte[]>>orderedBy(PtBytes.BYTES_COMPARATOR).put(ROW_BYTES, rowResult).build();
    mockery.checking(new Expectations() {

        {
            // row result is cached after first call, so second call requests no rows
            oneOf(txn).getRows(table, oneRow, oneColumn);
            will(returnValue(oneResult));
            oneOf(txn).getRows(table, noRows, oneColumn);
            will(returnValue(emptyResults));
        }
    });
    Assert.assertEquals(oneResult, ct.getRows(table, oneRow, oneColumn));
    Assert.assertEquals(oneResult, ct.getRows(table, oneRow, oneColumn));
    mockery.assertIsSatisfied();
}

Example 2 with ColumnSelection

use of com.palantir.atlasdb.keyvalue.api.ColumnSelection in project atlasdb by palantir.

the class CachingTransactionTest method testCacheEmptyGets.

@Test
public void testCacheEmptyGets() {
    final Set<byte[]> oneRow = ImmutableSortedSet.orderedBy(PtBytes.BYTES_COMPARATOR).add(ROW_BYTES).build();
    final ColumnSelection oneColumn = ColumnSelection.create(ImmutableList.of(COL_BYTES));
    final SortedMap<byte[], RowResult<byte[]>> emptyResults = ImmutableSortedMap.<byte[], RowResult<byte[]>>orderedBy(PtBytes.BYTES_COMPARATOR).build();
    mockery.checking(new Expectations() {

        {
            // the cache doesn't actually cache empty results in this case
            // this is probably an oversight, but this has been the behavior for a long time
            oneOf(txn).getRows(table, oneRow, oneColumn);
            will(returnValue(emptyResults));
            oneOf(txn).getRows(table, oneRow, oneColumn);
            will(returnValue(emptyResults));
        }
    });
    Assert.assertEquals(emptyResults, ct.getRows(table, oneRow, oneColumn));
    Assert.assertEquals(emptyResults, ct.getRows(table, oneRow, oneColumn));
    mockery.assertIsSatisfied();
}

Example 3 with ColumnSelection

use of com.palantir.atlasdb.keyvalue.api.ColumnSelection in project atlasdb by palantir.

the class CqlKeyValueService method getRangeWithPageCreator.

public <T> ClosableIterator<RowResult<T>> getRangeWithPageCreator(TableReference tableRef, RangeRequest rangeRequest, long timestamp, com.datastax.driver.core.ConsistencyLevel consistency, Supplier<ResultsExtractor<T>> resultsExtractor) {
    if (rangeRequest.isReverse()) {
        throw new UnsupportedOperationException();
    }
    if (rangeRequest.isEmptyRange()) {
        return ClosableIterators.wrap(ImmutableList.<RowResult<T>>of().iterator());
    }
    final int batchHint = rangeRequest.getBatchHint() == null ? 100 : rangeRequest.getBatchHint();
    final ColumnSelection selection = rangeRequest.getColumnNames().isEmpty() ? ColumnSelection.all() : ColumnSelection.create(rangeRequest.getColumnNames());
    final byte[] endExclusive = rangeRequest.getEndExclusive();
    final StringBuilder bindQuery = new StringBuilder();
    bindQuery.append("SELECT * FROM " + getFullTableName(tableRef) + " WHERE token(" + fieldNameProvider.row() + ") >= token(?) ");
    if (endExclusive.length > 0) {
        bindQuery.append("AND token(" + fieldNameProvider.row() + ") < token(?) ");
    }
    bindQuery.append("LIMIT " + batchHint);
    final String getLastRowQuery = "SELECT * FROM " + getFullTableName(tableRef) + " WHERE " + fieldNameProvider.row() + " = ?";
    return ClosableIterators.wrap(new AbstractPagingIterable<RowResult<T>, TokenBackedBasicResultsPage<RowResult<T>, byte[]>>() {

        @Override
        protected TokenBackedBasicResultsPage<RowResult<T>, byte[]> getFirstPage() throws Exception {
            return getPage(rangeRequest.getStartInclusive());
        }

        @Override
        protected TokenBackedBasicResultsPage<RowResult<T>, byte[]> getNextPage(TokenBackedBasicResultsPage<RowResult<T>, byte[]> previous) throws Exception {
            return getPage(previous.getTokenForNextPage());
        }

        TokenBackedBasicResultsPage<RowResult<T>, byte[]> getPage(final byte[] startKey) throws Exception {
            BoundStatement boundStatement = getPreparedStatement(tableRef, bindQuery.toString(), session).setConsistencyLevel(consistency).bind();
            boundStatement.setBytes(0, ByteBuffer.wrap(startKey));
            if (endExclusive.length > 0) {
                boundStatement.setBytes(1, ByteBuffer.wrap(endExclusive));
            }
            ResultSet resultSet = session.execute(boundStatement);
            List<Row> rows = Lists.newArrayList(resultSet.all());
            cqlKeyValueServices.logTracedQuery(bindQuery.toString(), resultSet, session, cqlStatementCache.normalQuery);
            byte[] maxRow = null;
            ResultsExtractor<T> extractor = resultsExtractor.get();
            for (Row row : rows) {
                byte[] rowName = getRowName(row);
                if (maxRow == null) {
                    maxRow = rowName;
                } else {
                    maxRow = PtBytes.BYTES_COMPARATOR.max(maxRow, rowName);
                }
            }
            if (maxRow == null) {
                return new SimpleTokenBackedResultsPage<>(endExclusive, ImmutableList.of(), false);
            }
            // get the rest of the last row
            BoundStatement boundLastRow = getPreparedStatement(tableRef, getLastRowQuery, session).bind();
            boundLastRow.setBytes(fieldNameProvider.row(), ByteBuffer.wrap(maxRow));
            try {
                resultSet = session.execute(boundLastRow);
            } catch (com.datastax.driver.core.exceptions.UnavailableException e) {
                throw new InsufficientConsistencyException("This operation requires all Cassandra" + " nodes to be up and available.", e);
            }
            rows.addAll(resultSet.all());
            cqlKeyValueServices.logTracedQuery(getLastRowQuery, resultSet, session, cqlStatementCache.normalQuery);
            for (Row row : rows) {
                extractor.internalExtractResult(timestamp, selection, getRowName(row), getColName(row), getValue(row), getTs(row));
            }
            SortedMap<byte[], SortedMap<byte[], T>> resultsByRow = Cells.breakCellsUpByRow(extractor.asMap());
            return ResultsExtractor.getRowResults(endExclusive, maxRow, resultsByRow);
        }
    }.iterator());
}

Example 4 with ColumnSelection

use of com.palantir.atlasdb.keyvalue.api.ColumnSelection in project atlasdb by palantir.

the class SchemaApiTestV2Table method getSmallRowRangeColumn1.

/**
 * Returns a mapping from all the row keys in a rangeRequest to their value at column Column1
 * (if that column exists for the row-key). As the Column1 values are all loaded in memory,
 * do not use for large amounts of data. The order of results is preserved in the map.
 */
public LinkedHashMap<String, Long> getSmallRowRangeColumn1(RangeRequest rangeRequest) {
    ColumnSelection colSelection = ColumnSelection.create(ImmutableList.of(PtBytes.toCachedBytes("c")));
    rangeRequest = rangeRequest.getBuilder().retainColumns(colSelection).build();
    Preconditions.checkArgument(rangeRequest.getColumnNames().size() <= 1, "Must not request columns other than Column1.");
    LinkedHashMap<String, Long> resultsMap = new LinkedHashMap<>();
    BatchingVisitableView.of(t.getRange(tableRef, rangeRequest)).immutableCopy().forEach(entry -> {
        SchemaApiTestTable.SchemaApiTestRowResult resultEntry = SchemaApiTestTable.SchemaApiTestRowResult.of(entry);
        resultsMap.put(resultEntry.getRowName().getComponent1(), resultEntry.getColumn1());
    });
    return resultsMap;
}

Example 5 with ColumnSelection

use of com.palantir.atlasdb.keyvalue.api.ColumnSelection in project atlasdb by palantir.

the class SchemaApiTestV2Table method getColumn2.

/**
 * Returns a mapping from the specified row keys to their value at column Column2.
 * As the Column2 values are all loaded in memory, do not use for large amounts of data.
 * If the column does not exist for a key, the entry will be omitted from the map.
 */
public Map<String, StringValue> getColumn2(Iterable<String> rowKeys) {
    ColumnSelection colSelection = ColumnSelection.create(ImmutableList.of(PtBytes.toCachedBytes("d")));
    List<SchemaApiTestTable.SchemaApiTestRow> rows = Lists.newArrayList(rowKeys).stream().map(SchemaApiTestTable.SchemaApiTestRow::of).collect(Collectors.toList());
    SortedMap<byte[], RowResult<byte[]>> results = t.getRows(tableRef, Persistables.persistAll(rows), colSelection);
    return results.values().stream().map(entry -> SchemaApiTestTable.SchemaApiTestRowResult.of(entry)).collect(Collectors.toMap(entry -> entry.getRowName().getComponent1(), SchemaApiTestTable.SchemaApiTestRowResult::getColumn2));
}