Skip to content
This repository has been archived by the owner on May 2, 2024. It is now read-only.

Commit

Permalink
Introduce cassandra.skip-parition-check
Browse files Browse the repository at this point in the history 
NativeCassandraSession#getPartitions is very slow
if there are many partitions used in WHERE clause.
In our use, the partitions usually exist, so
introduce cassandra.skip-parition-check property
to reduce planning time.
  • Loading branch information
@abicky @okkez
abicky authored and okkez committed Dec 11, 2023
1 parent 8e29b59 commit f0997d3
Show file tree
Hide file tree
Showing 7 changed files with 303 additions and 5 deletions.
13 changes: 13 additions & 0 deletions presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraClientConfig.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@ public class CassandraClientConfig
private String truststorePassword;
private File keystorePath;
private String keystorePassword;
private boolean skipPartitionCheck;

@NotNull
@Size(min = 1)
Expand Down Expand Up @@ -469,11 +470,23 @@ public Optional<String> getTruststorePassword()
return Optional.ofNullable(truststorePassword);
}

public boolean isSkipPartitionCheck()
{
return skipPartitionCheck;
}

@Config("cassandra.tls.truststore-password")
@ConfigSecuritySensitive
public CassandraClientConfig setTruststorePassword(String truststorePassword)
{
this.truststorePassword = truststorePassword;
return this;
}

@Config("cassandra.skip-partition-check")
public CassandraClientConfig setSkipPartitionCheck(boolean skipPartitionCheck)
{
this.skipPartitionCheck = skipPartitionCheck;
return this;
}
}
3 changes: 2 additions & 1 deletion presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraClientModule.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,7 @@ public static CassandraSession createCassandraSession(
contactPoints.forEach(clusterBuilder::addContactPoint);
return clusterBuilder.build();
}),
config.getNoHostAvailableRetryTimeout());
config.getNoHostAvailableRetryTimeout(),
config.isSkipPartitionCheck());
}
}
16 changes: 16 additions & 0 deletions presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraType.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -355,6 +355,22 @@ public static String getColumnValueForCql(Row row, int position, CassandraType c
}
}

public static String getColumnValueForCql(Object object, CassandraType cassandraType)
{
switch (cassandraType) {
case ASCII:
case TEXT:
case VARCHAR:
return CassandraCqlUtils.quoteStringLiteral(((Slice) object).toStringUtf8());
case INT:
case BIGINT:
return object.toString();
default:
throw new IllegalStateException("Handling of type " + cassandraType
+ " is not implemented");
}
}

private static String objectToString(Object object, CassandraType elemType)
{
switch (elemType) {
Expand Down
94 changes: 93 additions & 1 deletion presto-cassandra/src/main/java/com/facebook/presto/cassandra/NativeCassandraSession.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,9 @@
import com.google.common.collect.Ordering;
import com.google.common.collect.Sets;
import com.google.common.util.concurrent.UncheckedExecutionException;
import com.facebook.airlift.json.JsonCodec;
import com.facebook.airlift.log.Logger;
import io.airlift.slice.Slice;
import io.airlift.units.Duration;

import java.nio.Buffer;
Expand All @@ -82,6 +85,7 @@
import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.Iterables.filter;
import static com.google.common.collect.Iterables.transform;
import static java.lang.Math.toIntExact;
import static java.lang.String.format;
import static java.util.Comparator.comparing;
import static java.util.Locale.ENGLISH;
Expand Down Expand Up @@ -116,13 +120,15 @@ public KeyspaceMetadata load(String key)
private final Cluster cluster;
private final Supplier<Session> session;
private final Duration noHostAvailableRetryTimeout;
private final boolean skipPartitionCheck;

public NativeCassandraSession(String connectorId, JsonCodec<List<ExtraColumnMetadata>> extraColumnMetadataCodec, Cluster cluster, Duration noHostAvailableRetryTimeout)
public NativeCassandraSession(String connectorId, JsonCodec<List<ExtraColumnMetadata>> extraColumnMetadataCodec, Cluster cluster, Duration noHostAvailableRetryTimeout, boolean skipPartitionCheck)
{
this.connectorId = requireNonNull(connectorId, "connectorId is null");
this.extraColumnMetadataCodec = requireNonNull(extraColumnMetadataCodec, "extraColumnMetadataCodec is null");
this.cluster = requireNonNull(cluster, "cluster is null");
this.noHostAvailableRetryTimeout = requireNonNull(noHostAvailableRetryTimeout, "noHostAvailableRetryTimeout is null");
this.skipPartitionCheck = skipPartitionCheck;
this.session = memoize(cluster::connect);
}

Expand Down Expand Up @@ -391,6 +397,10 @@ private CassandraColumnHandle buildColumnHandle(AbstractTableMetadata tableMetad
@Override
public List<CassandraPartition> getPartitions(CassandraTable table, List<Set<Object>> filterPrefixes)
{
if (skipPartitionCheck) {
return buildPartitionsFromFilterPrefixes(table, filterPrefixes);
}

List<CassandraColumnHandle> partitionKeyColumns = table.getPartitionKeyColumns();

if (filterPrefixes.size() != partitionKeyColumns.size()) {
Expand Down Expand Up @@ -458,6 +468,88 @@ public List<CassandraPartition> getPartitions(CassandraTable table, List<Set<Obj
return partitions.build();
}

private List<CassandraPartition> buildPartitionsFromFilterPrefixes(CassandraTable table, List<Set<Object>> filterPrefixes)
{
List<CassandraColumnHandle> partitionKeyColumns = table.getPartitionKeyColumns();

if (filterPrefixes.size() != partitionKeyColumns.size()) {
return ImmutableList.of(CassandraPartition.UNPARTITIONED);
}

ByteBuffer buffer = ByteBuffer.allocate(1000);
HashMap<ColumnHandle, NullableValue> map = new HashMap<>();
Set<String> uniquePartitionIds = new HashSet<>();
StringBuilder stringBuilder = new StringBuilder();

boolean isComposite = partitionKeyColumns.size() > 1;

ImmutableList.Builder<CassandraPartition> partitions = ImmutableList.builder();
for (List<Object> values : Sets.cartesianProduct(filterPrefixes)) {
buffer.clear();
map.clear();
stringBuilder.setLength(0);
for (int i = 0; i < partitionKeyColumns.size(); i++) {
Object value = values.get(i);
CassandraColumnHandle columnHandle = partitionKeyColumns.get(i);
CassandraType cassandraType = columnHandle.getCassandraType();

switch (cassandraType) {
case TEXT:
Slice slice = (Slice) value;
if (isComposite) {
buffer.putShort((short) slice.length());
buffer.put(slice.getBytes());
buffer.put((byte) 0);
}
else {
buffer.put(slice.getBytes());
}
break;
case INT:
int intValue = toIntExact((long) value);
if (isComposite) {
buffer.putShort((short) Integer.BYTES);
buffer.putInt(intValue);
buffer.put((byte) 0);
}
else {
buffer.putInt(intValue);
}
break;
case BIGINT:
if (isComposite) {
buffer.putShort((short) Long.BYTES);
buffer.putLong((long) value);
buffer.put((byte) 0);
}
else {
buffer.putLong((long) value);
}
break;
default:
throw new IllegalStateException("Handling of type " + cassandraType + " is not implemented");
}

map.put(columnHandle, NullableValue.of(cassandraType.getNativeType(), value));
if (i > 0) {
stringBuilder.append(" AND ");
}
stringBuilder.append(CassandraCqlUtils.validColumnName(columnHandle.getName()));
stringBuilder.append(" = ");
stringBuilder.append(CassandraType.getColumnValueForCql(value, cassandraType));
}
buffer.flip();
byte[] key = new byte[buffer.limit()];
buffer.get(key);
TupleDomain<ColumnHandle> tupleDomain = TupleDomain.fromFixedValues(map);
String partitionId = stringBuilder.toString();
if (uniquePartitionIds.add(partitionId)) {
partitions.add(new CassandraPartition(key, partitionId, tupleDomain, false));
}
}
return partitions.build();
}

@Override
public ResultSet execute(String cql, Object... values)
{
Expand Down
11 changes: 10 additions & 1 deletion presto-cassandra/src/test/java/com/facebook/presto/cassandra/EmbeddedCassandra.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ public final class EmbeddedCassandra
private static final Duration REFRESH_SIZE_ESTIMATES_TIMEOUT = new Duration(1, MINUTES);

private static CassandraSession session;
private static ReopeningCluster cluster;
private static boolean initialized;

private EmbeddedCassandra() {}
Expand Down Expand Up @@ -89,7 +90,8 @@ public static synchronized void start()
"EmbeddedCassandra",
JsonCodec.listJsonCodec(ExtraColumnMetadata.class),
cluster,
new Duration(1, MINUTES));
new Duration(1, MINUTES),
false);

try {
checkConnectivity(session);
Expand All @@ -101,6 +103,7 @@ public static synchronized void start()
}

EmbeddedCassandra.session = session;
EmbeddedCassandra.cluster = cluster;
initialized = true;
}

Expand Down Expand Up @@ -141,6 +144,12 @@ public static synchronized int getPort()
return PORT;
}

public static synchronized ReopeningCluster getCluster()
{
checkIsInitialized();
return cluster;
}

private static void checkIsInitialized()
{
checkState(initialized, "EmbeddedCassandra must be started with #start() method before retrieving the cluster retrieval");
Expand Down
7 changes: 5 additions & 2 deletions presto-cassandra/src/test/java/com/facebook/presto/cassandra/TestCassandraClientConfig.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,8 @@ public void testDefaults()
.setKeystorePassword(null)
.setTruststorePath(null)
.setTruststorePassword(null)
.setTlsEnabled(false));
.setTlsEnabled(false)
.setSkipPartitionCheck(false));
}

@Test
Expand Down Expand Up @@ -103,6 +104,7 @@ public void testExplicitPropertyMappings()
.put("cassandra.tls.keystore-password", "keystore-password")
.put("cassandra.tls.truststore-path", "/tmp/truststore")
.put("cassandra.tls.truststore-password", "truststore-password")
.put("cassandra.skip-partition-check", "true")
.build();

CassandraClientConfig expected = new CassandraClientConfig()
Expand Down Expand Up @@ -136,7 +138,8 @@ public void testExplicitPropertyMappings()
.setKeystorePath(new File("/tmp/keystore"))
.setKeystorePassword("keystore-password")
.setTruststorePath(new File("/tmp/truststore"))
.setTruststorePassword("truststore-password");
.setTruststorePassword("truststore-password")
.setSkipPartitionCheck(true);

ConfigAssertions.assertFullMapping(properties, expected);
}
Expand Down
Loading

0 comments on commit f0997d3

Please sign in to comment.