Make paging work for SELECT queries with GROUP BY

Patch by Alex Petrov; reviewed by Caleb Rackliffe and Andrés de la Peña for CASSANDRA-16427
This commit is contained in:
Alex Petrov 2021-02-04 09:40:39 +01:00
parent 9fa2c28dc2
commit 4c3b426120
3 changed files with 114 additions and 61 deletions

View File

@ -32,6 +32,7 @@ import org.apache.cassandra.thrift.CqlResultType;
import org.apache.cassandra.thrift.CqlRow;
import org.apache.cassandra.utils.ByteBufferUtil;
import org.apache.cassandra.service.pager.PagingState;
import com.google.common.annotations.VisibleForTesting;
public class ResultSet
{
@ -286,6 +287,12 @@ public class ResultSet
names.add(name);
}
@VisibleForTesting
public PagingState getPagingState()
{
return pagingState;
}
public void setHasMorePages(PagingState pagingState)
{
this.pagingState = pagingState;

View File

@ -21,17 +21,23 @@ package org.apache.cassandra.distributed.impl;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.Future;
import com.google.common.collect.Iterators;
import com.datastax.driver.core.ProtocolVersion;
import org.apache.cassandra.config.CFMetaData;
import org.apache.cassandra.cql3.CQLStatement;
import org.apache.cassandra.cql3.QueryOptions;
import org.apache.cassandra.cql3.QueryProcessor;
import org.apache.cassandra.cql3.UntypedResultSet;
import org.apache.cassandra.cql3.statements.SelectStatement;
import org.apache.cassandra.db.Keyspace;
import org.apache.cassandra.distributed.api.ConsistencyLevel;
import org.apache.cassandra.distributed.api.ICoordinator;
import org.apache.cassandra.distributed.api.IInstance;
@ -122,49 +128,64 @@ public class Coordinator implements ICoordinator
throw new IllegalArgumentException("Page size should be strictly positive but was " + pageSize);
return instance.sync(() -> {
ClientState clientState = makeFakeClientState();
ConsistencyLevel consistencyLevel = ConsistencyLevel.valueOf(consistencyLevelOrigin.name());
CQLStatement prepared = QueryProcessor.getStatement(query, ClientState.forInternalCalls()).statement;
CQLStatement prepared = QueryProcessor.getStatement(query, clientState).statement;
List<ByteBuffer> boundBBValues = new ArrayList<>();
for (Object boundValue : boundValues)
{
boundBBValues.add(ByteBufferUtil.objectToBytes(boundValue));
}
prepared.validate(QueryState.forInternalCalls().getClientState());
prepared.validate(clientState);
assert prepared instanceof SelectStatement : "Only SELECT statements can be executed with paging";
ClientState clientState = QueryState.forInternalCalls().getClientState();
SelectStatement selectStatement = (SelectStatement) prepared;
QueryOptions queryOptions = QueryOptions.create(toCassandraCL(consistencyLevel),
boundBBValues,
false,
pageSize,
null,
null,
Server.CURRENT_VERSION);
Pageable pageable = selectStatement.getPageableCommand(queryOptions);
// Usually pager fetches a single page (see SelectStatement#execute). We need to iterate over all
// of the results lazily.
QueryPager pager = QueryPagers.pager(pageable, toCassandraCL(consistencyLevel), clientState, null);
Iterator<Object[]> iter = RowUtil.toObjects(selectStatement.getResultMetadata().names,
UntypedResultSet.create(selectStatement,
pager,
pageSize).iterator());
QueryState queryState = new QueryState(clientState);
QueryOptions initialOptions = QueryOptions.create(toCassandraCL(consistencyLevel),
boundBBValues,
false,
pageSize,
null,
null,
Server.CURRENT_VERSION);
ResultMessage.Rows initialRows = selectStatement.execute(queryState, initialOptions);
Iterator<Object[]> iter = new Iterator<Object[]>() {
ResultMessage.Rows rows = selectStatement.execute(queryState, initialOptions);
Iterator<Object[]> iter = RowUtil.toIter(rows);
// We have to make sure iterator is not running on main thread.
Iterator<Object[]> it = new Iterator<Object[]>() {
public boolean hasNext()
{
return instance.sync(() -> iter.hasNext()).call();
if (iter.hasNext())
return true;
if (rows.result.metadata.getPagingState() == null)
return false;
QueryOptions nextOptions = QueryOptions.create(toCassandraCL(consistencyLevel),
boundBBValues,
true,
pageSize,
rows.result.metadata.getPagingState(),
null,
Server.CURRENT_VERSION);
rows = selectStatement.execute(queryState, nextOptions);
iter = Iterators.forArray(RowUtil.toObjects(initialRows.result.metadata.names, rows.result.rows));
return hasNext();
}
public Object[] next()
{
return instance.sync(() -> iter.next()).call();
return iter.next();
}
};
return QueryResults.fromObjectArrayIterator(RowUtil.getColumnNames(selectStatement.getResultMetadata().names), it);
return QueryResults.fromObjectArrayIterator(RowUtil.getColumnNames(initialRows.result.metadata.names), iter);
}).call();
}

View File

@ -19,10 +19,13 @@
package org.apache.cassandra.distributed.impl;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import com.google.common.collect.Iterators;
import com.google.common.collect.Lists;
import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.Row;
@ -56,44 +59,7 @@ public class RowUtil
public static Object[][] toObjects(ResultMessage.Rows rows)
{
Object[][] result = new Object[rows.result.rows.size()][];
List<ColumnSpecification> specs = rows.result.metadata.names;
for (int i = 0; i < rows.result.rows.size(); i++)
{
List<ByteBuffer> row = rows.result.rows.get(i);
result[i] = new Object[row.size()];
for (int j = 0; j < row.size(); j++)
{
ByteBuffer bb = row.get(j);
if (bb != null)
result[i][j] = specs.get(j).type.getSerializer().deserialize(bb);
}
}
return result;
}
public static Iterator<Object[]> toObjects(UntypedResultSet rs)
{
return toObjects(rs.metadata(), rs.iterator());
}
public static Iterator<Object[]> toObjects(List<ColumnSpecification> columnSpecs, Iterator<UntypedResultSet.Row> rs)
{
return Iterators.transform(rs,
(row) -> {
Object[] objectRow = new Object[columnSpecs.size()];
for (int i = 0; i < columnSpecs.size(); i++)
{
ColumnSpecification columnSpec = columnSpecs.get(i);
ByteBuffer bb = row.getBytes(columnSpec.name.toString());
if (bb != null)
objectRow[i] = columnSpec.type.getSerializer().deserialize(bb);
}
return objectRow;
});
return toObjects(rows.result.metadata.names, rows.result.rows);
}
public static Iterator<Object[]> toObjects(ResultSet rs)
@ -109,5 +75,64 @@ public class RowUtil
});
}
public static Object[][] toObjects(List<ColumnSpecification> specs, List<List<ByteBuffer>> rows)
{
Object[][] result = new Object[rows.size()][];
for (int i = 0; i < rows.size(); i++)
{
List<ByteBuffer> row = rows.get(i);
result[i] = new Object[row.size()];
for (int j = 0; j < row.size(); j++)
{
ByteBuffer bb = row.get(j);
if (bb != null)
result[i][j] = specs.get(j).type.getSerializer().deserialize(bb);
}
}
return result;
}
public static Iterator<Object[]> toIter(UntypedResultSet rs)
{
return toIter(rs.metadata(), rs.iterator());
}
public static Iterator<Object[]> toIter(ResultMessage.Rows rows)
{
return toIterInternal(rows.result.metadata.names, rows.result.rows);
}
public static Iterator<Object[]> toIter(List<ColumnSpecification> columnSpecs, Iterator<UntypedResultSet.Row> rs)
{
Iterator<List<ByteBuffer>> iter = Iterators.transform(rs,
(row) -> {
List<ByteBuffer> bbs = new ArrayList<>(columnSpecs.size());
for (int i = 0; i < columnSpecs.size(); i++)
{
ColumnSpecification columnSpec = columnSpecs.get(i);
bbs.add(row.getBytes(columnSpec.name.toString()));
}
return bbs;
});
return toIterInternal(columnSpecs, Lists.newArrayList(iter));
}
private static Iterator<Object[]> toIterInternal(List<ColumnSpecification> columnSpecs, List<List<ByteBuffer>> rs)
{
return Iterators.transform(rs.iterator(),
(row) -> {
Object[] objectRow = new Object[columnSpecs.size()];
for (int i = 0; i < columnSpecs.size(); i++)
{
ColumnSpecification columnSpec = columnSpecs.get(i);
ByteBuffer bb = row.get(i);
if (bb != null)
objectRow[i] = columnSpec.type.getSerializer().deserialize(bb);
}
return objectRow;
});
}
}