BatchUpdatesCollector.java
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you 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 org.apache.cassandra.cql3.statements;
import java.nio.ByteBuffer;
import java.util.*;
import com.google.common.collect.HashMultiset;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Maps;
import org.apache.cassandra.db.commitlog.CommitLogSegment;
import org.apache.cassandra.db.virtual.VirtualMutation;
import org.apache.cassandra.net.MessagingService;
import org.apache.cassandra.schema.TableId;
import org.apache.cassandra.schema.TableMetadata;
import org.apache.cassandra.db.*;
import org.apache.cassandra.db.partitions.PartitionUpdate;
import static org.apache.cassandra.utils.MonotonicClock.Global.approxTime;
/**
* Utility class to collect updates.
*
* <p>In a batch statement we don't want to recreate mutations every time as this is particularly inefficient when
* applying multiple batch to the same partition (see #6737). </p>
*
*/
final class BatchUpdatesCollector implements UpdatesCollector
{
/**
* The columns that will be updated for each table (keyed by the table ID).
*/
private final Map<TableId, RegularAndStaticColumns> updatedColumns;
/**
* The number of updated rows per table and key.
*/
private final Map<TableId, HashMultiset<ByteBuffer>> perPartitionKeyCounts;
/**
* The mutations per keyspace.
*
* optimised for the common single-keyspace case
*
* Key is keyspace name, then we have an IMutationBuilder for each touched partition key in that keyspace
*
* MutationBuilder holds a PartitionUpdate.Builder
*/
private final Map<String, Map<ByteBuffer, IMutationBuilder>> mutationBuilders = Maps.newHashMapWithExpectedSize(1);
BatchUpdatesCollector(Map<TableId, RegularAndStaticColumns> updatedColumns, Map<TableId, HashMultiset<ByteBuffer>> perPartitionKeyCounts)
{
super();
this.updatedColumns = updatedColumns;
this.perPartitionKeyCounts = perPartitionKeyCounts;
}
/**
* Gets the <code>PartitionUpdate.Builder</code> for the specified column family and key. If the builder does not
* exist it will be created.
*
* @param metadata the column family meta data
* @param dk the partition key
* @param consistency the consistency level
* @return the <code>PartitionUpdate.Builder</code> for the specified column family and key
*/
public PartitionUpdate.Builder getPartitionUpdateBuilder(TableMetadata metadata, DecoratedKey dk, ConsistencyLevel consistency)
{
IMutationBuilder mut = getMutationBuilder(metadata, dk, consistency);
PartitionUpdate.Builder upd = mut.get(metadata.id);
if (upd == null)
{
RegularAndStaticColumns columns = updatedColumns.get(metadata.id);
assert columns != null;
upd = new PartitionUpdate.Builder(metadata, dk, columns, perPartitionKeyCounts.get(metadata.id).count(dk.getKey()));
mut.add(upd);
}
return upd;
}
private IMutationBuilder getMutationBuilder(TableMetadata metadata, DecoratedKey dk, ConsistencyLevel consistency)
{
Map<ByteBuffer, IMutationBuilder> ksMap = keyspaceMap(metadata.keyspace);
IMutationBuilder mutationBuilder = ksMap.get(dk.getKey());
if (mutationBuilder == null)
{
mutationBuilder = makeMutationBuilder(metadata, dk, consistency);
ksMap.put(dk.getKey(), mutationBuilder);
}
return mutationBuilder;
}
private IMutationBuilder makeMutationBuilder(TableMetadata metadata, DecoratedKey partitionKey, ConsistencyLevel cl)
{
if (metadata.isVirtual())
{
return new VirtualMutationBuilder(metadata.keyspace, partitionKey);
}
else
{
MutationBuilder builder = new MutationBuilder(metadata.keyspace, partitionKey, 1);
return metadata.isCounter() ? new CounterMutationBuilder(builder, cl) : builder;
}
}
/**
* Returns a collection containing all the mutations.
* @return a collection containing all the mutations.
*/
public List<IMutation> toMutations()
{
List<IMutation> ms = new ArrayList<>();
for (Map<ByteBuffer, IMutationBuilder> ksMap : mutationBuilders.values())
{
for (IMutationBuilder builder : ksMap.values())
{
IMutation mutation = builder.build();
mutation.validateIndexedColumns();
mutation.validateSize(MessagingService.current_version, CommitLogSegment.ENTRY_OVERHEAD_SIZE);
ms.add(mutation);
}
}
return ms;
}
/**
* Returns the key-mutation mappings for the specified keyspace.
*
* @param ksName the keyspace name
* @return the key-mutation mappings for the specified keyspace.
*/
private Map<ByteBuffer, IMutationBuilder> keyspaceMap(String ksName)
{
Map<ByteBuffer, IMutationBuilder> ksMap = mutationBuilders.get(ksName);
if (ksMap == null)
{
ksMap = Maps.newHashMapWithExpectedSize(1);
mutationBuilders.put(ksName, ksMap);
}
return ksMap;
}
private interface IMutationBuilder
{
/**
* Add a new PartitionUpdate builder to this mutation builder
* @param builder the builder to add
* @return this
*/
IMutationBuilder add(PartitionUpdate.Builder builder);
/**
* Build the immutable mutation
*/
IMutation build();
/**
* Get the builder for the given tableId
*/
PartitionUpdate.Builder get(TableId tableId);
}
private static class MutationBuilder implements IMutationBuilder
{
private final Map<TableId, PartitionUpdate.Builder> modifications;
private final DecoratedKey key;
private final String keyspaceName;
private final long createdAt = approxTime.now();
private MutationBuilder(String keyspaceName, DecoratedKey key, int initialSize)
{
this.keyspaceName = keyspaceName;
this.key = key;
this.modifications = Maps.newHashMapWithExpectedSize(initialSize);
}
public MutationBuilder add(PartitionUpdate.Builder updateBuilder)
{
assert updateBuilder != null;
assert updateBuilder.partitionKey().getPartitioner() == key.getPartitioner();
PartitionUpdate.Builder prev = modifications.put(updateBuilder.metadata().id, updateBuilder);
if (prev != null)
// developer error
throw new IllegalArgumentException("Table " + updateBuilder.metadata().name + " already has modifications in this mutation: " + prev);
return this;
}
public Mutation build()
{
ImmutableMap.Builder<TableId, PartitionUpdate> updates = new ImmutableMap.Builder<>();
for (Map.Entry<TableId, PartitionUpdate.Builder> updateEntry : modifications.entrySet())
{
PartitionUpdate update = updateEntry.getValue().build();
updates.put(updateEntry.getKey(), update);
}
return new Mutation(keyspaceName, key, updates.build(), createdAt);
}
public PartitionUpdate.Builder get(TableId tableId)
{
return modifications.get(tableId);
}
public DecoratedKey key()
{
return key;
}
public boolean isEmpty()
{
return modifications.isEmpty();
}
public String getKeyspaceName()
{
return keyspaceName;
}
}
private static class CounterMutationBuilder implements IMutationBuilder
{
private final MutationBuilder mutationBuilder;
private final ConsistencyLevel cl;
private CounterMutationBuilder(MutationBuilder mutationBuilder, ConsistencyLevel cl)
{
this.mutationBuilder = mutationBuilder;
this.cl = cl;
}
public IMutationBuilder add(PartitionUpdate.Builder builder)
{
return mutationBuilder.add(builder);
}
public IMutation build()
{
return new CounterMutation(mutationBuilder.build(), cl);
}
public PartitionUpdate.Builder get(TableId id)
{
return mutationBuilder.get(id);
}
}
private static class VirtualMutationBuilder implements IMutationBuilder
{
private final String keyspaceName;
private final DecoratedKey partitionKey;
private final HashMap<TableId, PartitionUpdate.Builder> modifications = new HashMap<>();
private VirtualMutationBuilder(String keyspaceName, DecoratedKey partitionKey)
{
this.keyspaceName = keyspaceName;
this.partitionKey = partitionKey;
}
@Override
public VirtualMutationBuilder add(PartitionUpdate.Builder builder)
{
PartitionUpdate.Builder prev = modifications.put(builder.metadata().id, builder);
if (null != prev)
throw new IllegalStateException();
return this;
}
@Override
public VirtualMutation build()
{
ImmutableMap.Builder<TableId, PartitionUpdate> updates = new ImmutableMap.Builder<>();
modifications.forEach((tableId, updateBuilder) -> updates.put(tableId, updateBuilder.build()));
return new VirtualMutation(keyspaceName, partitionKey, updates.build());
}
@Override
public PartitionUpdate.Builder get(TableId tableId)
{
return modifications.get(tableId);
}
}
}