src/java/org/apache/cassandra/cql3/statements/BatchUpdatesCollector.java - cassandra - Git at Google

 /*
  * 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);
         }
     }
 }
	/*
	* 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);
	}
	}
	}