added batch aggregator operations

Author: htemelski-redis

internal/routing/aggregator.go

@@ -25,6 +25,8 @@ type ResponseAggregator interface {
 	// AddWithKey processes a single shard response for a specific key (used by keyed aggregators).
 	AddWithKey(key string, result interface{}, err error) error
 
+	BatchAdd(map[string]interface{}, error) error
+
 	// Result returns the final aggregated result and any error.
 	Result() (interface{}, error)
 }
@@ -93,6 +95,14 @@ func (a *AllSucceededAggregator) Add(result interface{}, err error) error {
 	return nil
 }
 
+func (a *AllSucceededAggregator) BatchAdd(results map[string]interface{}, err error) error {
+	for _, res := range results {
+		a.Add(res, err)
+	}
+
+	return nil
+}
+
 func (a *AllSucceededAggregator) Result() (interface{}, error) {
 	var err error
 	res, e := a.res.Load(), a.err.Load()
@@ -127,6 +137,14 @@ func (a *OneSucceededAggregator) Add(result interface{}, err error) error {
 	return nil
 }
 
+func (a *OneSucceededAggregator) BatchAdd(results map[string]interface{}, err error) error {
+	for _, res := range results {
+		a.Add(res, err)
+	}
+
+	return nil
+}
+
 func (a *OneSucceededAggregator) AddWithKey(key string, result interface{}, err error) error {
 	return a.Add(result, err)
 }
@@ -162,6 +180,14 @@ func (a *AggSumAggregator) Add(result interface{}, err error) error {
 	return nil
 }
 
+func (a *AggSumAggregator) BatchAdd(results map[string]interface{}, err error) error {
+	for _, res := range results {
+		a.Add(res, err)
+	}
+
+	return nil
+}
+
 func (a *AggSumAggregator) AddWithKey(key string, result interface{}, err error) error {
 	return a.Add(result, err)
 }
@@ -198,6 +224,14 @@ func (a *AggMinAggregator) Add(result interface{}, err error) error {
 	return nil
 }
 
+func (a *AggMinAggregator) BatchAdd(results map[string]interface{}, err error) error {
+	for _, res := range results {
+		a.Add(res, err)
+	}
+
+	return nil
+}
+
 func (a *AggMinAggregator) AddWithKey(key string, result interface{}, err error) error {
 	return a.Add(result, err)
 }
@@ -238,6 +272,14 @@ func (a *AggMaxAggregator) Add(result interface{}, err error) error {
 	return nil
 }
 
+func (a *AggMaxAggregator) BatchAdd(results map[string]interface{}, err error) error {
+	for _, res := range results {
+		a.Add(res, err)
+	}
+
+	return nil
+}
+
 func (a *AggMaxAggregator) AddWithKey(key string, result interface{}, err error) error {
 	return a.Add(result, err)
 }
@@ -285,6 +327,14 @@ func (a *AggLogicalAndAggregator) Add(result interface{}, err error) error {
 	return nil
 }
 
+func (a *AggLogicalAndAggregator) BatchAdd(results map[string]interface{}, err error) error {
+	for _, res := range results {
+		a.Add(res, err)
+	}
+
+	return nil
+}
+
 func (a *AggLogicalAndAggregator) AddWithKey(key string, result interface{}, err error) error {
 	return a.Add(result, err)
 }
@@ -331,6 +381,14 @@ func (a *AggLogicalOrAggregator) Add(result interface{}, err error) error {
 	return nil
 }
 
+func (a *AggLogicalOrAggregator) BatchAdd(results map[string]interface{}, err error) error {
+	for _, res := range results {
+		a.Add(res, err)
+	}
+
+	return nil
+}
+
 func (a *AggLogicalOrAggregator) AddWithKey(key string, result interface{}, err error) error {
 	return a.Add(result, err)
 }
@@ -391,10 +449,7 @@ type DefaultKeylessAggregator struct {
 	firstErr error
 }
 
-func (a *DefaultKeylessAggregator) Add(result interface{}, err error) error {
-	a.mu.Lock()
-	defer a.mu.Unlock()
-
+func (a *DefaultKeylessAggregator) add(result interface{}, err error) error {
 	if err != nil && a.firstErr == nil {
 		a.firstErr = err
 		return nil
@@ -405,6 +460,21 @@ func (a *DefaultKeylessAggregator) Add(result interface{}, err error) error {
 	return nil
 }
 
+func (a *DefaultKeylessAggregator) Add(result interface{}, err error) error {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+
+	return a.add(result, err)
+}
+
+func (a *DefaultKeylessAggregator) BatchAdd(results map[string]interface{}, err error) error {
+	for _, res := range results {
+		a.add(res, err)
+	}
+
+	return nil
+}
+
 func (a *DefaultKeylessAggregator) AddWithKey(key string, result interface{}, err error) error {
 	return a.Add(result, err)
 }
@@ -434,10 +504,7 @@ func NewDefaultKeyedAggregator(keyOrder []string) *DefaultKeyedAggregator {
 	}
 }
 
-func (a *DefaultKeyedAggregator) Add(result interface{}, err error) error {
-	a.mu.Lock()
-	defer a.mu.Unlock()
-
+func (a *DefaultKeyedAggregator) add(result interface{}, err error) error {
 	if err != nil && a.firstErr == nil {
 		a.firstErr = err
 		return nil
@@ -449,10 +516,22 @@ func (a *DefaultKeyedAggregator) Add(result interface{}, err error) error {
 	return nil
 }
 
-func (a *DefaultKeyedAggregator) AddWithKey(key string, result interface{}, err error) error {
+func (a *DefaultKeyedAggregator) Add(result interface{}, err error) error {
 	a.mu.Lock()
 	defer a.mu.Unlock()
 
+	return a.add(result, err)
+}
+
+func (a *DefaultKeyedAggregator) BatchAdd(results map[string]interface{}, err error) error {
+	for _, res := range results {
+		a.add(res, err)
+	}
+
+	return nil
+}
+
+func (a *DefaultKeyedAggregator) addWithKey(key string, result interface{}, err error) error {
 	if err != nil && a.firstErr == nil {
 		a.firstErr = err
 		return nil
@@ -463,6 +542,26 @@ func (a *DefaultKeyedAggregator) AddWithKey(key string, result interface{}, err
 	return nil
 }
 
+func (a *DefaultKeyedAggregator) AddWithKey(key string, result interface{}, err error) error {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+
+	a.addWithKey(key, result, err)
+	return nil
+}
+
+func (a *DefaultKeyedAggregator) BatchAddWithKeyOrder(results map[string]interface{}, keyOrder []string) error {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+
+	a.keyOrder = keyOrder
+	for key, val := range results {
+		_ = a.addWithKey(key, val, nil)
+	}
+
+	return nil
+}
+
 func (a *DefaultKeyedAggregator) SetKeyOrder(keyOrder []string) {
 	a.mu.Lock()
 	defer a.mu.Unlock()
@@ -504,12 +603,24 @@ type SpecialAggregator struct {
 	errors         []error
 }
 
+func (a *SpecialAggregator) add(result interface{}, err error) error {
+	a.results = append(a.results, result)
+	a.errors = append(a.errors, err)
+	return nil
+}
+
 func (a *SpecialAggregator) Add(result interface{}, err error) error {
 	a.mu.Lock()
 	defer a.mu.Unlock()
 
-	a.results = append(a.results, result)
-	a.errors = append(a.errors, err)
+	return a.add(result, err)
+}
+
+func (a *SpecialAggregator) BatchAdd(results map[string]interface{}, err error) error {
+	for _, res := range results {
+		a.add(res, err)
+	}
+
 	return nil
 }
 

osscluster_router.go

@@ -393,21 +393,17 @@ func (c *ClusterClient) aggregateKeyedValues(cmd Cmder, keyedResults map[string]
 	aggregator := c.createAggregator(policy, cmd, true)
 
 	// Set key order for keyed aggregators
-	if keyedAgg, ok := aggregator.(*routing.DefaultKeyedAggregator); ok {
-		keyedAgg.SetKeyOrder(keyOrder)
+	var keyedAgg *routing.DefaultKeyedAggregator
+	var isKeyedAgg bool
+	var err error
+	if keyedAgg, isKeyedAgg = aggregator.(*routing.DefaultKeyedAggregator); isKeyedAgg {
+		err = keyedAgg.BatchAddWithKeyOrder(keyedResults, keyOrder)
+	} else {
+		err = aggregator.BatchAdd(keyedResults, nil)
 	}
 
-	// Add results with keys
-	for key, value := range keyedResults {
-		if keyedAgg, ok := aggregator.(*routing.DefaultKeyedAggregator); ok {
-			if err := keyedAgg.AddWithKey(key, value, nil); err != nil {
-				return err
-			}
-		} else {
-			if err := aggregator.Add(value, nil); err != nil {
-				return err
-			}
-		}
+	if err != nil {
+		return err
 	}
 
 	return c.finishAggregation(cmd, aggregator)