From 4e4e31df22b27668fb9c13db30b598e7606ec1d4 Mon Sep 17 00:00:00 2001
From: Jaroslaw Bojar <jarek.bojar@gmail.com>
Date: Sat, 1 Mar 2014 01:59:27 +0100
Subject: [PATCH] Better support for virtual nodes (many ranges on same
 cassandra instance)

---
 .../apache/cassandra/hadoop/cql3/CqlRecordWriter.java | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/src/java/org/apache/cassandra/hadoop/cql3/CqlRecordWriter.java b/src/java/org/apache/cassandra/hadoop/cql3/CqlRecordWriter.java
index d08e08a8f0bb..5d9b1bdd75f7 100644
--- a/src/java/org/apache/cassandra/hadoop/cql3/CqlRecordWriter.java
+++ b/src/java/org/apache/cassandra/hadoop/cql3/CqlRecordWriter.java
@@ -67,7 +67,7 @@ final class CqlRecordWriter extends AbstractColumnFamilyRecordWriter<Map<String,
     private static final Logger logger = LoggerFactory.getLogger(CqlRecordWriter.class);
 
     // handles for clients for each range running in the threadpool
-    private final Map<Range, RangeClient> clients;
+    private final Map<List<InetAddress>, RangeClient> clients;
 
     // host to prepared statement id mappings
     private ConcurrentHashMap<Cassandra.Client, Integer> preparedStatements = new ConcurrentHashMap<Cassandra.Client, Integer>();
@@ -100,7 +100,7 @@ final class CqlRecordWriter extends AbstractColumnFamilyRecordWriter<Map<String,
     CqlRecordWriter(Configuration conf)
     {
         super(conf);
-        this.clients = new HashMap<Range, RangeClient>();
+        this.clients = new HashMap<List<InetAddress>, RangeClient>();
 
         try
         {
@@ -165,13 +165,14 @@ public void write(Map<String, ByteBuffer> keyColumns, List<ByteBuffer> values) t
         Range<Token> range = ringCache.getRange(getPartitionKey(keyColumns));
 
         // get the client for the given range, or create a new one
-        RangeClient client = clients.get(range);
+        final List<InetAddress> endpoints = ringCache.getEndpoint(range);
+        RangeClient client = clients.get(endpoints);
         if (client == null)
         {
             // haven't seen keys for this range: create new client
-            client = new RangeClient(ringCache.getEndpoint(range));
+            client = new RangeClient(endpoints);
             client.start();
-            clients.put(range, client);
+            clients.put(endpoints, client);
         }
 
         // add primary key columns to the bind variables