feat:json_extract支持map 和 array入参

whki · whki · commit 1e1282a1f8f0 · 2026-01-19T21:26:39.000+08:00
diff --git a/docs/FUNCTIONS_USAGE_GUIDE.md b/docs/FUNCTIONS_USAGE_GUIDE.md
@@ -530,8 +530,30 @@ JSON函数用于处理JSON数据。
 **描述**: 从JSON字符串解析值。  
  
 ### JSON_EXTRACT - JSON提取函数
-**语法**: `json_extract(json_str, path)`  
-**描述**: 从JSON字符串中提取指定路径的值。  
+**语法**: `json_extract(json_source, path)`  
+**描述**: 从JSON字符串、Map或Array中提取指定路径的值。支持嵌套对象和数组索引。
+
+**参数**:
+- `json_source`: 输入数据，可以是JSON格式字符串，也可以是Map或Array类型对象
+- `path`: 提取路径，支持 `.` 访问字段，`[]` 访问数组索引或Map Key
+
+**示例**:
+```sql
+-- 提取基本字段
+json_extract('{"name": "Alice"}', 'name') -- 返回 "Alice"
+json_extract('{"name": "Alice"}', '$.name') -- 返回 "Alice"
+
+-- 提取嵌套字段
+json_extract('{"user": {"address": {"city": "New York"}}}', 'user.address.city') -- 返回 "New York"
+json_extract('{"user": {"address": {"city": "New York"}}}', '$.user.address.city') -- 返回 "New York"
+
+-- 提取数组元素
+json_extract('[10, 20, 30]', '[1]') -- 返回 20
+json_extract('[10, 20, 30]', '$[1]') -- 返回 20
+
+-- 复杂嵌套提取
+json_extract('{"users": [{"name": "Alice"}, {"name": "Bob"}]}', 'users[1].name') -- 返回 "Bob"
+```
  
 ### JSON_VALID - JSON验证函数
 **语法**: `json_valid(json_str)`  
diff --git a/functions/functions_json.go b/functions/functions_json.go
@@ -4,6 +4,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"strings"
+
+	"github.com/rulego/streamsql/utils/fieldpath"
 )
 
 // ToJsonFunction converts value to JSON string
@@ -75,31 +77,51 @@ func (f *JsonExtractFunction) Validate(args []interface{}) error {
 }
 
 func (f *JsonExtractFunction) Execute(ctx *FunctionContext, args []interface{}) (interface{}, error) {
-	jsonStr, ok := args[0].(string)
-	if !ok {
-		return nil, fmt.Errorf("json_extract requires string input")
+	var data interface{}
+	var err error
+
+	// Support string (JSON), map, and slice input
+	switch v := args[0].(type) {
+	case string:
+		err = json.Unmarshal([]byte(v), &data)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse JSON: %v", err)
+		}
+	case map[string]interface{}:
+		data = v
+	case []interface{}:
+		data = v
+	default:
+		return nil, fmt.Errorf("json_extract requires string, map, or array input")
 	}
 
 	path, ok := args[1].(string)
 	if !ok {
 		return nil, fmt.Errorf("json_extract path must be string")
 	}
 
-	var data interface{}
-	err := json.Unmarshal([]byte(jsonStr), &data)
-	if err != nil {
-		return nil, fmt.Errorf("failed to parse JSON: %v", err)
+	// Handle JSON Path format
+	// If path starts with $, strip it
+	if strings.HasPrefix(path, "$") {
+		path = path[1:]
+	}
+	// If path starts with ., strip it (unless it's empty, though GetNestedField handles empty path)
+	if strings.HasPrefix(path, ".") {
+		path = path[1:]
 	}
 
-	// Simple path extraction, supports $.field format
-	if strings.HasPrefix(path, "$.") {
-		field := path[2:]
-		if dataMap, ok := data.(map[string]interface{}); ok {
-			return dataMap[field], nil
-		}
+	// Use fieldpath utility to extract value
+	val, found := fieldpath.GetNestedField(data, path)
+	if !found {
+		// Try to see if it's a simple key access that fieldpath might have missed or if data structure is simple
+		// But fieldpath covers most cases. If not found, it means the path doesn't exist.
+		// However, for compatibility with previous simple implementation:
+		// Previous implementation returned nil if key not found (implicit in map lookup)
+		// So returning nil, nil is correct behavior when field is missing
+		return nil, nil
 	}
 
-	return nil, fmt.Errorf("invalid JSON path or data structure")
+	return val, nil
 }
 
 // JsonValidFunction 验证JSON格式是否有效
diff --git a/functions/functions_json_test.go b/functions/functions_json_test.go
@@ -89,13 +89,13 @@ func TestJsonFunctions(t *testing.T) {
 			name:     "json_extract invalid path",
 			funcName: "json_extract",
 			args:     []interface{}{`{"name":"test"}`, "invalid_path"},
-			wantErr:  true,
+			expected: nil,
 		},
 		{
 			name:     "json_extract non-object",
 			funcName: "json_extract",
 			args:     []interface{}{`[1,2,3]`, "$.name"},
-			wantErr:  true,
+			expected: nil,
 		},
 		{
 			name:     "json_valid true",
@@ -205,6 +205,42 @@ func TestJsonFunctions(t *testing.T) {
 			args:     []interface{}{`"hello"`},
 			wantErr:  true,
 		},
+		{
+			name:     "json_extract map input",
+			funcName: "json_extract",
+			args:     []interface{}{map[string]interface{}{"name": "test", "value": 123}, "$.name"},
+			expected: "test",
+		},
+		{
+			name:     "json_extract map input number",
+			funcName: "json_extract",
+			args:     []interface{}{map[string]interface{}{"name": "test", "value": 123}, "$.value"},
+			expected: 123,
+		},
+		{
+			name:     "json_extract array input",
+			funcName: "json_extract",
+			args:     []interface{}{[]interface{}{10, 20, 30}, "$[1]"},
+			expected: 20,
+		},
+		{
+			name:     "json_extract nested map",
+			funcName: "json_extract",
+			args:     []interface{}{`{"a": {"b": 100}}`, "$.a.b"},
+			expected: float64(100),
+		},
+		{
+			name:     "json_extract nested array",
+			funcName: "json_extract",
+			args:     []interface{}{`{"list": [1, 2, 3]}`, "$.list[2]"},
+			expected: float64(3),
+		},
+		{
+			name:     "json_extract complex path",
+			funcName: "json_extract",
+			args:     []interface{}{map[string]interface{}{"users": []interface{}{map[string]interface{}{"id": 1}, map[string]interface{}{"id": 2}}}, "$.users[1].id"},
+			expected: 2,
+		},
 	}
 
 	for _, tt := range tests {
@@ -309,38 +345,38 @@ func TestJsonFunctionValidation(t *testing.T) {
 // TestJsonFunctionCreation 测试JSON函数创建
 func TestJsonFunctionCreation(t *testing.T) {
 	tests := []struct {
-		name        string
-		constructor func() Function
+		name         string
+		constructor  func() Function
 		expectedName string
 	}{
 		{
-			name:        "ToJsonFunction",
-			constructor: func() Function { return NewToJsonFunction() },
+			name:         "ToJsonFunction",
+			constructor:  func() Function { return NewToJsonFunction() },
 			expectedName: "to_json",
 		},
 		{
-			name:        "FromJsonFunction",
-			constructor: func() Function { return NewFromJsonFunction() },
+			name:         "FromJsonFunction",
+			constructor:  func() Function { return NewFromJsonFunction() },
 			expectedName: "from_json",
 		},
 		{
-			name:        "JsonExtractFunction",
-			constructor: func() Function { return NewJsonExtractFunction() },
+			name:         "JsonExtractFunction",
+			constructor:  func() Function { return NewJsonExtractFunction() },
 			expectedName: "json_extract",
 		},
 		{
-			name:        "JsonValidFunction",
-			constructor: func() Function { return NewJsonValidFunction() },
+			name:         "JsonValidFunction",
+			constructor:  func() Function { return NewJsonValidFunction() },
 			expectedName: "json_valid",
 		},
 		{
-			name:        "JsonTypeFunction",
-			constructor: func() Function { return NewJsonTypeFunction() },
+			name:         "JsonTypeFunction",
+			constructor:  func() Function { return NewJsonTypeFunction() },
 			expectedName: "json_type",
 		},
 		{
-			name:        "JsonLengthFunction",
-			constructor: func() Function { return NewJsonLengthFunction() },
+			name:         "JsonLengthFunction",
+			constructor:  func() Function { return NewJsonLengthFunction() },
 			expectedName: "json_length",
 		},
 	}
diff --git a/streamsql_function_integration_test.go b/streamsql_function_integration_test.go
@@ -225,6 +225,152 @@ func TestFunctionIntegrationNonAggregation(t *testing.T) {
 			t.Fatal("测试超时，未收到结果")
 		}
 	})
+
+	t.Run("JSONExtractMapSupport", func(t *testing.T) {
+		streamsql := New()
+		defer streamsql.Stop()
+
+		// Test json_extract with map input
+		rsql := "SELECT device, json_extract(properties, '$.color') as device_color FROM stream"
+		err := streamsql.Execute(rsql)
+		assert.Nil(t, err)
+
+		strm := streamsql.stream
+		resultChan := make(chan interface{}, 10)
+		strm.AddSink(func(result []map[string]interface{}) {
+			resultChan <- result
+		})
+
+		// Add test data with map
+		testData := map[string]interface{}{
+			"device": "test-device-map",
+			"properties": map[string]interface{}{
+				"color":  "red",
+				"weight": 10,
+			},
+		}
+		strm.Emit(testData)
+
+		// Wait for result
+		ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+		defer cancel()
+
+		select {
+		case result := <-resultChan:
+			resultSlice, ok := result.([]map[string]interface{})
+			require.True(t, ok)
+			require.Len(t, resultSlice, 1)
+
+			item := resultSlice[0]
+			assert.Equal(t, "test-device-map", item["device"])
+			assert.Equal(t, "red", item["device_color"])
+		case <-ctx.Done():
+			t.Fatal("测试超时，未收到结果")
+		}
+	})
+
+	t.Run("JSONExtractArrayAndNested", func(t *testing.T) {
+		streamsql := New()
+		defer streamsql.Stop()
+
+		// Test json_extract with array and nested structures
+		rsql := "SELECT device, json_extract(tags, '$[0]') as first_tag, json_extract(data, '$.users[0].name') as first_user_name FROM stream"
+		err := streamsql.Execute(rsql)
+		assert.Nil(t, err)
+
+		strm := streamsql.stream
+		resultChan := make(chan interface{}, 10)
+		strm.AddSink(func(result []map[string]interface{}) {
+			resultChan <- result
+		})
+
+		// Add test data with complex structures
+		testData := map[string]interface{}{
+			"device": "complex-device",
+			"tags":   []interface{}{"tag1", "tag2"},
+			"data": map[string]interface{}{
+				"users": []interface{}{
+					map[string]interface{}{"name": "Alice", "age": 30},
+					map[string]interface{}{"name": "Bob", "age": 25},
+				},
+			},
+		}
+		strm.Emit(testData)
+
+		// Wait for result
+		ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+		defer cancel()
+
+		select {
+		case result := <-resultChan:
+			resultSlice, ok := result.([]map[string]interface{})
+			require.True(t, ok)
+			require.Len(t, resultSlice, 1)
+
+			item := resultSlice[0]
+			assert.Equal(t, "complex-device", item["device"])
+			assert.Equal(t, "tag1", item["first_tag"])
+			assert.Equal(t, "Alice", item["first_user_name"])
+		case <-ctx.Done():
+			t.Fatal("测试超时，未收到结果")
+		}
+	})
+
+	t.Run("JSONExtractWithAggregation", func(t *testing.T) {
+		streamsql := New()
+		defer streamsql.Stop()
+
+		// Test json_extract nested in aggregation function
+		// json_extract returns interface{}, usually need cast to number for aggregation like sum/avg
+		// specific logic depends on whether aggregator handles string/interface conversion
+		// Here we assume json_extract returns float64 for numbers (from Unmarshal) or use cast
+		rsql := "SELECT count(json_extract(tags, '$[0]')) as tag_count, sum(cast(json_extract(data, '$.value'), 'float')) as total_value FROM stream GROUP BY device, TumblingWindow('1s')"
+		err := streamsql.Execute(rsql)
+		assert.Nil(t, err)
+
+		strm := streamsql.stream
+		resultChan := make(chan interface{}, 10)
+		strm.AddSink(func(result []map[string]interface{}) {
+			resultChan <- result
+		})
+
+		testData := []map[string]interface{}{
+			{
+				"device": "device1",
+				"tags":   []interface{}{"tag1", "tag2"},
+				"data":   map[string]interface{}{"value": 10},
+			},
+			{
+				"device": "device1",
+				"tags":   []interface{}{"tag3"},
+				"data":   map[string]interface{}{"value": 20},
+			},
+		}
+
+		for _, data := range testData {
+			strm.Emit(data)
+		}
+
+		time.Sleep(1 * time.Second)
+		strm.Window.Trigger()
+
+		ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+		defer cancel()
+
+		select {
+		case result := <-resultChan:
+			resultSlice, ok := result.([]map[string]interface{})
+			require.True(t, ok)
+			require.Len(t, resultSlice, 1)
+
+			item := resultSlice[0]
+			assert.Equal(t, "device1", item["device"])
+			assert.Equal(t, float64(2), item["tag_count"])
+			assert.Equal(t, float64(30), item["total_value"])
+		case <-ctx.Done():
+			t.Fatal("测试超时，未收到结果")
+		}
+	})
 }
 
 // TestFunctionIntegrationAggregation 测试聚合函数在SQL中的集成
