Retain user record partition keys (a8m#1)

aggregator.go

@@ -31,14 +31,16 @@ func (a *Aggregator) Count() int {
 
 // Put record using `data` and `partitionKey`. This method is thread-safe.
 func (a *Aggregator) Put(data []byte, partitionKey string) {
-	// For now, all records in the aggregated record will have
-	// the same partition key.
-	// later, we will add shard-mapper same as the KPL use.
-	// see: https://github.com/a8m/kinesis-producer/issues/1
-	if len(a.pkeys) == 0 {
-		a.pkeys = []string{partitionKey}
-		a.nbytes += len([]byte(partitionKey))
-	}
+	// a8m: For now, all records in the aggregated record will have
+	//      the same partition key.
+	//      later, we will add shard-mapper same as the KPL use.
+	//      see: https://github.com/a8m/kinesis-producer/issues/1
+	// jawang35: In this fork I'm allowing user records to retain
+	//           their individual partition keys with the
+	//           understanding that shard-mapping is not implemented
+	//           like it is in the KPL.
+	a.pkeys = append(a.pkeys, partitionKey)
+	a.nbytes += len([]byte(partitionKey))
 	keyIndex := uint64(len(a.pkeys) - 1)
 
 	a.nbytes += partitionKeyIndexSize

aggregator_test.go

@@ -22,7 +22,7 @@ func TestSizeAndCount(t *testing.T) {
 	for i := 0; i < n; i++ {
 		a.Put(data, pkey)
 	}
-	assert(t, a.Size() == 5*n+5+8*n, "size should equal to the data and the partition-key")
+	assert(t, a.Size() == 5*n+5*n+8*n, "size should equal to the data and the partition-keys")
 	assert(t, a.Count() == n, "count should be equal to the number of Put calls")
 }