feat(json): Add JSON.SET fpha support (#3797)

* add json.set fpha support

* fix json set test

Author: ndyakov

json.go

@@ -35,6 +35,7 @@ type JSONCmdable interface {
 	JSONObjLen(ctx context.Context, key, path string) *IntPointerSliceCmd
 	JSONSet(ctx context.Context, key, path string, value interface{}) *StatusCmd
 	JSONSetMode(ctx context.Context, key, path string, value interface{}, mode string) *StatusCmd
+	JSONSetWithArgs(ctx context.Context, key, path string, value interface{}, options *JSONSetArgsOptions) *StatusCmd
 	JSONStrAppend(ctx context.Context, key, path, value string) *IntPointerSliceCmd
 	JSONStrLen(ctx context.Context, key, path string) *IntPointerSliceCmd
 	JSONToggle(ctx context.Context, key, path string) *IntPointerSliceCmd
@@ -57,6 +58,25 @@ type JSONArrTrimArgs struct {
 	Stop  *int
 }
 
+// FPHAType is the floating-point type used for storing FP homogeneous arrays
+// in JSON.SET (Redis 8.8+).
+type FPHAType string
+
+const (
+	FPHATypeBF16 FPHAType = "BF16"
+	FPHATypeFP16 FPHAType = "FP16"
+	FPHATypeFP32 FPHAType = "FP32"
+	FPHATypeFP64 FPHAType = "FP64"
+)
+
+// JSONSetArgsOptions are the optional arguments for JSONSetWithArgs.
+// Mode is "NX" or "XX" (case-insensitive). FPHA, when set, forces Redis to
+// store all FP homogeneous arrays using the specified floating-point type.
+type JSONSetArgsOptions struct {
+	Mode string
+	FPHA FPHAType
+}
+
 type JSONCmd struct {
 	baseCmd
 	val      string
@@ -584,6 +604,15 @@ func (c cmdable) JSONSet(ctx context.Context, key, path string, value interface{
 // the argument is a string or []byte when we assume that it can be passed directly as JSON.
 // For more information, see https://redis.io/commands/json.set
 func (c cmdable) JSONSetMode(ctx context.Context, key, path string, value interface{}, mode string) *StatusCmd {
+	return c.JSONSetWithArgs(ctx, key, path, value, &JSONSetArgsOptions{Mode: mode})
+}
+
+// JSONSetWithArgs sets the JSON value at the given path in the given key with optional arguments
+// for setting mode (NX/XX) and the FPHA (Floating-Point Homogeneous Array) type used for storing
+// FP arrays. The value must be something that can be marshaled to JSON (using encoding/JSON) unless
+// the argument is a string or []byte when we assume that it can be passed directly as JSON.
+// For more information, see https://redis.io/commands/json.set
+func (c cmdable) JSONSetWithArgs(ctx context.Context, key, path string, value interface{}, options *JSONSetArgsOptions) *StatusCmd {
 	var bytes []byte
 	var err error
 	switch v := value.(type) {
@@ -595,13 +624,17 @@ func (c cmdable) JSONSetMode(ctx context.Context, key, path string, value interf
 		bytes, err = json.Marshal(v)
 	}
 	args := []interface{}{"JSON.SET", key, path, util.BytesToString(bytes)}
-	if mode != "" {
-		switch strings.ToUpper(mode) {
-		case "XX", "NX":
-			args = append(args, strings.ToUpper(mode))
-
-		default:
-			panic("redis: JSON.SET mode must be NX or XX")
+	if options != nil {
+		if options.Mode != "" {
+			switch strings.ToUpper(options.Mode) {
+			case "XX", "NX":
+				args = append(args, strings.ToUpper(options.Mode))
+			default:
+				panic("redis: JSON.SET mode must be NX or XX")
+			}
+		}
+		if options.FPHA != "" {
+			args = append(args, "FPHA", string(options.FPHA))
 		}
 	}
 	cmd := NewStatusCmd(ctx, args...)

json_test.go

@@ -262,6 +262,55 @@ var _ = Describe("JSON Commands", Label("json"), func() {
 				Expect(cmd.Val()).To(Equal("OK"))
 			})
 
+			It("should JSONSetWithArgs with FPHA", Label("json.set", "json"), func() {
+				SkipBeforeRedisVersion(8.8, "FPHA argument requires Redis 8.8+")
+
+				fpArray := `[1.1, 2.2, 3.3, 4.4]`
+				for _, fpha := range []redis.FPHAType{
+					redis.FPHATypeBF16,
+					redis.FPHATypeFP16,
+					redis.FPHATypeFP32,
+					redis.FPHATypeFP64,
+				} {
+					key := "fpha_" + string(fpha)
+					cmd := client.JSONSetWithArgs(ctx, key, "$", fpArray, &redis.JSONSetArgsOptions{FPHA: fpha})
+					Expect(cmd.Err()).NotTo(HaveOccurred())
+					Expect(cmd.Val()).To(Equal("OK"))
+
+					res, err := client.JSONGet(ctx, key, "$").Result()
+					Expect(err).NotTo(HaveOccurred())
+					Expect(res).NotTo(BeEmpty())
+				}
+			})
+
+			It("should JSONSetWithArgs with FPHA and NX/XX", Label("json.set", "json"), func() {
+				SkipBeforeRedisVersion(8.8, "FPHA argument requires Redis 8.8+")
+
+				key := "fpha_nx"
+				fpArray := `[1.5, 2.5, 3.5]`
+
+				// NX + FPHA succeeds when key does not exist
+				res, err := client.JSONSetWithArgs(ctx, key, "$", fpArray, &redis.JSONSetArgsOptions{
+					Mode: "NX", FPHA: redis.FPHATypeFP32,
+				}).Result()
+				Expect(err).NotTo(HaveOccurred())
+				Expect(res).To(Equal("OK"))
+
+				// NX + FPHA returns redis.Nil when key already exists
+				cmd := client.JSONSetWithArgs(ctx, key, "$", `[4.5, 5.5]`, &redis.JSONSetArgsOptions{
+					Mode: "NX", FPHA: redis.FPHATypeFP32,
+				})
+				Expect(cmd.Err()).To(Equal(redis.Nil))
+				Expect(cmd.Val()).To(Equal(""))
+
+				// XX + FPHA succeeds when key exists
+				res, err = client.JSONSetWithArgs(ctx, key, "$", `[4.5, 5.5]`, &redis.JSONSetArgsOptions{
+					Mode: "XX", FPHA: redis.FPHATypeFP32,
+				}).Result()
+				Expect(err).NotTo(HaveOccurred())
+				Expect(res).To(Equal("OK"))
+			})
+
 			It("should JSONGet", Label("json.get", "json", "NonRedisEnterprise"), func() {
 				res, err := client.JSONSet(ctx, "get3", "$", `{"a": 1, "b": 2}`).Result()
 				Expect(err).NotTo(HaveOccurred())