com.palantir.atlasdb.keyvalue.cassandra.HostPartitioner Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of atlasdb-cassandra Show documentation
Show all versions of atlasdb-cassandra Show documentation
Palantir open source project
The newest version!
/*
* (c) Copyright 2018 Palantir Technologies Inc. All rights reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.palantir.atlasdb.keyvalue.cassandra;
import com.google.common.base.Function;
import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.ListMultimap;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimaps;
import com.palantir.atlasdb.keyvalue.api.Cell;
import com.palantir.atlasdb.keyvalue.cassandra.pool.CassandraServer;
import com.palantir.logsafe.Preconditions;
import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
final class HostPartitioner {
private HostPartitioner() {
// Static class
}
static Map> partitionMapByHost(
CassandraClientPool clientPool, Iterable> cells) {
Map>> partitionedByHost =
partitionByHost(clientPool, cells, entry -> entry.getKey().getRowName());
Map> cellsByHost = new HashMap<>();
for (Map.Entry>> hostAndCells : partitionedByHost.entrySet()) {
Map cellsForHost =
Maps.newHashMapWithExpectedSize(hostAndCells.getValue().size());
for (Map.Entry entry : hostAndCells.getValue()) {
cellsForHost.put(entry.getKey(), entry.getValue());
}
cellsByHost.put(hostAndCells.getKey(), cellsForHost);
}
return cellsByHost;
}
static Map> partitionByHost(
CassandraClientPool clientPool, Iterable iterable, Function keyExtractor) {
// Ensure that the same key goes to the same partition. This is important when writing multiple columns
// to the same row, since this is a normally a single write in cassandra, whereas splitting the columns
// into different requests results in multiple writes.
ListMultimap partitionedByKey = ArrayListMultimap.create();
for (V value : iterable) {
partitionedByKey.put(ByteBuffer.wrap(keyExtractor.apply(value)), value);
}
ListMultimap valuesByHost = ArrayListMultimap.create();
for (ByteBuffer key : partitionedByKey.keySet()) {
Preconditions.checkState(key.hasArray(), "Expected an array backed buffer");
Preconditions.checkState(key.arrayOffset() == 0, "Buffer array must have no offset");
Preconditions.checkState(key.limit() == key.array().length, "Array length must match the limit");
CassandraServer host = clientPool.getRandomServerForKey(key.array());
valuesByHost.putAll(host, partitionedByKey.get(key));
}
return Multimaps.asMap(valuesByHost);
}
}
| |