version up

GetMapValue refactoring $ go test -bench=. -count=5 goos: darwin goarch: amd64 pkg: github.com/yandex/pandora/lib/mp cpu: Intel(R) Core(TM) i7-9750H CPU @ 2.60GHz BenchmarkGetMapValue-12 45223 24348 ns/op BenchmarkGetMapValue-12 51969 23631 ns/op BenchmarkGetMapValue-12 52303 24771 ns/op BenchmarkGetMapValue-12 52562 23441 ns/op BenchmarkGetMapValue-12 49144 23131 ns/op BenchmarkGetMapValue2-12 49866 23739 ns/op BenchmarkGetMapValue2-12 47876 23441 ns/op BenchmarkGetMapValue2-12 49641 24082 ns/op BenchmarkGetMapValue2-12 50052 23706 ns/op BenchmarkGetMapValue2-12 51109 23714 ns/op PASS ok github.com/yandex/pandora/lib/mp 15.501s
yandex · Nov 24, 2023 · bbfbd6e · bbfbd6e
1 parent 4a858e2
commit bbfbd6e
Show file tree

Hide file tree

Showing 3 changed files with 394 additions and 81 deletions.
diff --git a/cli/cli.go b/cli/cli.go
@@ -25,7 +25,7 @@ import (
 	"go.uber.org/zap/zapcore"
 )
 
-const Version = "0.5.16"
+const Version = "0.5.17"
 const defaultConfigFile = "load"
 const stdinConfigSelector = "-"
 

diff --git a/lib/mp/map.go b/lib/mp/map.go
@@ -2,6 +2,7 @@ package mp
 
 import (
 	"fmt"
+	"reflect"
 	"strconv"
 	"strings"
 )
@@ -15,11 +16,10 @@ func (e *ErrSegmentNotFound) Error() string {
 	return fmt.Sprintf("segment %s not found in path %s", e.segment, e.path)
 }
 
-func GetMapValue(input map[string]any, path string, iter Iterator) (any, error) {
+func GetMapValue(current map[string]any, path string, iter Iterator) (any, error) {
 	var curSegment strings.Builder
 	segments := strings.Split(path, ".")
 
-	currentData := input
 	for i, segment := range segments {
 		segment = strings.TrimSpace(segment)
 		curSegment.WriteByte('.')
@@ -29,80 +29,98 @@ func GetMapValue(input map[string]any, path string, iter Iterator) (any, error)
 			indexStr := strings.ToLower(strings.TrimSpace(segment[openBraceIdx+1 : len(segment)-1]))
 
 			segment = segment[:openBraceIdx]
-			value, exists := currentData[segment]
-			if !exists {
+			pathVal, ok := current[segment]
+			if !ok {
 				return nil, &ErrSegmentNotFound{path: path, segment: segment}
 			}
-
-			mval, isMval := value.([]map[string]string)
-			if isMval {
-				index, err := calcIndex(indexStr, curSegment.String(), len(mval), iter)
-				if err != nil {
-					return nil, fmt.Errorf("failed to calc index: %w", err)
-				}
-				vval := mval[index]
-				currentData = make(map[string]any, len(vval))
-				for k, v := range vval {
-					currentData[k] = v
-				}
-				continue
+			sliceElement, err := extractFromSlice(pathVal, indexStr, curSegment.String(), iter)
+			if err != nil {
+				return nil, fmt.Errorf("cant extract value path=`%s`,segment=`%s`,err=%w", segment, path, err)
 			}
-
-			mapSlice, isMapSlice := value.([]map[string]any)
-			if !isMapSlice {
-				anySlice, isAnySlice := value.([]any)
-				if isAnySlice {
-					index, err := calcIndex(indexStr, curSegment.String(), len(anySlice), iter)
-					if err != nil {
-						return nil, fmt.Errorf("failed to calc index: %w", err)
-					}
-					if i != len(segments)-1 {
-						return nil, fmt.Errorf("not last segment %s in path %s", segment, path)
-					}
-					return anySlice[index], nil
-				}
-				stringSlice, isStringSlice := value.([]string)
-				if isStringSlice {
-					index, err := calcIndex(indexStr, curSegment.String(), len(stringSlice), iter)
-					if err != nil {
-						return nil, fmt.Errorf("failed to calc index: %w", err)
-					}
-					if i != len(segments)-1 {
-						return nil, fmt.Errorf("not last segment %s in path %s", segment, path)
-					}
-					return stringSlice[index], nil
+			current, ok = sliceElement.(map[string]any)
+			if !ok {
+				if i != len(segments)-1 {
+					return nil, fmt.Errorf("not last segment %s in path %s", segment, path)
 				}
-				return nil, fmt.Errorf("invalid type of segment %s in path %s", segment, path)
+				return sliceElement, nil
 			}
-
-			index, err := calcIndex(indexStr, curSegment.String(), len(mapSlice), iter)
-			if err != nil {
-				return nil, fmt.Errorf("failed to calc index: %w", err)
-			}
-			currentData = mapSlice[index]
 		} else {
-			value, exists := currentData[segment]
-			if !exists {
+			pathVal, ok := current[segment]
+			if !ok {
 				return nil, &ErrSegmentNotFound{path: path, segment: segment}
 			}
-			var ok bool
-			currentData, ok = value.(map[string]any)
+			current, ok = pathVal.(map[string]any)
 			if !ok {
 				if i != len(segments)-1 {
 					return nil, fmt.Errorf("not last segment %s in path %s", segment, path)
 				}
-				return value, nil
+				return pathVal, nil
 			}
 		}
 	}
 
-	return currentData, nil
+	return current, nil
+}
+
+func extractFromSlice(curValue any, indexStr string, curSegment string, iter Iterator) (result any, err error) {
+	validTypes := []reflect.Type{
+		reflect.TypeOf([]map[string]string{}),
+		reflect.TypeOf([]map[string]any{}),
+		reflect.TypeOf([]any{}),
+		reflect.TypeOf([]string{}),
+		reflect.TypeOf([]int{}),
+		reflect.TypeOf([]int64{}),
+		reflect.TypeOf([]float64{}),
+	}
+
+	var valueLen int
+	var valueFound bool
+	for _, valueType := range validTypes {
+		if reflect.TypeOf(curValue) == valueType {
+			valueLen = reflect.ValueOf(curValue).Len()
+			valueFound = true
+			break
+		}
+	}
+
+	if !valueFound {
+		return nil, fmt.Errorf("invalid type of value `%+v`, %T", curValue, curValue)
+	}
+
+	index, err := calcIndex(indexStr, curSegment, valueLen, iter)
+	if err != nil {
+		return nil, fmt.Errorf("failed to calc index for %T; err: %w", curValue, err)
+	}
+
+	switch v := curValue.(type) {
+	case []map[string]string:
+		currentData := make(map[string]any, len(v[index]))
+		for k, val := range v[index] {
+			currentData[k] = val
+		}
+		return currentData, nil
+	case []map[string]any:
+		return v[index], nil
+	case []any:
+		return v[index], nil
+	case []string:
+		return v[index], nil
+	case []int:
+		return v[index], nil
+	case []int64:
+		return v[index], nil
+	case []float64:
+		return v[index], nil
+	}
+
+	// This line should never be reached, as we've covered all valid types above
+	return nil, fmt.Errorf("invalid type of value `%+v`, %T", curValue, curValue)
 }
 
 func calcIndex(indexStr string, segment string, length int, iter Iterator) (int, error) {
 	index, err := strconv.Atoi(indexStr)
 	if err != nil && indexStr != "next" && indexStr != "rand" && indexStr != "last" {
-		return 0, fmt.Errorf("invalid index: %s", indexStr)
+		return 0, fmt.Errorf("index should be integer or one of [next, rand, last], but got `%s`", indexStr)
 	}
 	if indexStr != "next" && indexStr != "rand" && indexStr != "last" {
 		if index >= 0 && index < length {