Repository: flyteorg/flyte (language: python)

Issue:
[Core feature] Support for Native JSON types in flyte
### Motivation: Why do you think this is important?

Currently flyte supports JSON through protobuf Struct. This causes a lot of problems for the users, especially when using ints as Structs only store numbers - which are double values.


### Goal: What should the final outcome look like, ideally?

Support a new FlyteIDL.LiteralType called LiteralType.JSON. So users can simply use any object that dumps to a JSON using binary / bytearray as the transport layer.

It should be possible to also have customer json serializers and deserializers for the data like JsonPickle etc

### Describe alternatives you've considered 

New literalType called LiteralType.JSON
This maps to a Literal of type binary (wire format)
Optionally the metadata / schema can be extracted and stored

###  Propose: Link/Inline OR Additional context 

_No response_

### Are you sure this issue hasn't been raised already?

- [X] Yes

### Have you read the Code of Conduct?

- [X] Yes


Discussion / hints:
To support native JSON types in Flyte, you propose introducing a new `LiteralType.JSON` in FlyteIDL, using binary/bytearray for transport, and allowing custom JSON serializers/deserializers. This approach aims to address issues with current JSON handling via protobuf Struct, particularly with integer representation.

Would you like more detail on how to implement this feature or any potential challenges? 

Sources:
- [GitHub Issue #4505](https://github.com/flyteorg/flyte/issues/4505)
- [GitHub Issue #1282](https://github.com/flyteorg/flyte/issues/1282)
- [Discuss Flyte](https://discuss.flyte.org/t/3651339/is-there-way-to-specify-a-custom-type-transformer-for-a-json/)

---
This is an AI-generated response and your feedback is appreciated! Please leave a 👍 if this is helpful and 👎 if it is not.

Mention @runllm to ask follow-up questions. See [runllm.com](https://runllm.com/) for more!
Related to https://github.com/flyteorg/flyte/issues/4505
#take
From https://github.com/flyteorg/flyte/issues/4505#issuecomment-2060451042, I see the follow Python types that can use `Literal.JSON`:

1. `dicts` -> `JsonPickle`
2. `dataclasses` -> Use `Mashumuro's` encoders/decorders or the object's `to_json` + `from_json` (These methods are defined with `DataClassJSONMixin`)

@kumare3 Are you considering other Python types that would use `Literal.JSON`?
I am saying wholesale migration to json for these unsupported types from struct
https://github.com/flyteorg/flyte/pull/5337

The PRs will be created as follows:
- Flyte
[Flyte][1][IDL] Binary IDL With MessagePack
[Flyte][2][Literal Type For Scalar] Binary IDL With MessagePack
[Flyte][3][Attribute Access] Binary IDL With MessagePack
[Flyte][4][FlyteCTL] Binary IDL With MessagePack
[Flyte][5][Compiler][Struct Literal Type using JSON SCHEMA] Binary IDL With MessagePack

- flytekit
[flytekit][1][SimpleTransformer] Binary IDL With MessagePack
[flytekit][2][untyped dict] Binary IDL With MessagePack 
[flytekit][3] [list, dict and nested cases] Binary IDL With MessagePack
[flytekit][4][pure dataclass and nested dataclass] Binary IDL With MessagePack
[flytekit][5][Attribute Access] Binary IDL With MessagePack
[flytekit][6][Flyte Types] Binary IDL With MessagePack

follow up:
for flytekit, we can suport `Dict[int, dataclass]` as input, but this need to also change `click_types.py`, which will be big enough to open a new PR.
follow up 2:
we should consider rewrite `get_literal_type` for `Dict[int, dataclasss]` attribute access.
> we should consider rewrite get_literal_type for Dict[int, dataclasss] attribute access

what is this now and what should it be?
> > we should consider rewrite get_literal_type for Dict[int, dataclasss] attribute access
> 
> what is this now and what should it be?

now: Dict[int, dataclasss] not supported, and Dict[int, dataclasss] not supported with attribute access
after: Dict[int, dataclasss] supported, and Dict[int, dataclasss] supported with attribute access
`get_literal_type` is the function to provide each field in the `Dict` for the propeller to access the attribute's type.

This case can be supported, but we don't have backward compatible issue now, since no users report about this before.
```python
# Status is an Enum class
@dataclass
class DC:
    grid: Dict[str, List[Optional[Union[int, str, float, bool, Status, InnerDC]]]] = field(default_factory=lambda: {
        'all_types': [InnerDC()],
    })
```

Mashumaro Issue: https://github.com/Fatal1ty/mashumaro/issues/252

Fix patch:
diff --git a/flytepropeller/pkg/controller/nodes/attr_path_resolver.go b/flytepropeller/pkg/controller/nodes/attr_path_resolver.go
index 192fa1956c..3b4e46ce50 100644
--- a/flytepropeller/pkg/controller/nodes/attr_path_resolver.go
+++ b/flytepropeller/pkg/controller/nodes/attr_path_resolver.go
@@ -46,8 +46,7 @@ func resolveAttrPathInPromise(ctx context.Context, datastore *storage.DataStore,
 			}
 			currVal = currVal.GetCollection().GetLiterals()[attr.GetIntValue()]
 			index++
-		// scalar is always the leaf, so we can break here
-		case *core.Literal_Scalar:
+		default:
 			break
 		}
 	}
@@ -107,9 +106,7 @@ func resolveAttrPathInPbStruct(nodeID string, st *structpb.Struct, bindAttrPath
 }
 
 // resolveAttrPathInBinary resolves the binary idl object (e.g. dataclass, pydantic basemodel) with attribute path
-func resolveAttrPathInBinary(nodeID string, binaryIDL *core.Binary, bindAttrPath []*core.PromiseAttribute) (*core.
-	Literal,
-	error) {
+func resolveAttrPathInBinary(nodeID string, binaryIDL *core.Binary, bindAttrPath []*core.PromiseAttribute) (*core.Literal, error) {
 
 	binaryBytes := binaryIDL.GetValue()
 	serializationFormat := binaryIDL.GetTag()
@@ -165,6 +162,28 @@ func resolveAttrPathInBinary(nodeID string, binaryIDL *core.Binary, bindAttrPath
 		}
 	}
 
+	// In arrayNodeHandler, the resolved value should be a literal collection.
+	// If the current value is already a collection, convert it to a literal collection.
+	// This conversion does not affect how Flytekit processes the resolved value.
+	if collection, ok := currVal.([]any); ok {
+		literals := make([]*core.Literal, len(collection))
+		for i, v := range collection {
+			resolvedBinaryBytes, err := msgpack.Marshal(v)
+			if err != nil {
+				return nil, err
+			}
+			literals[i] = constructResolvedBinary(resolvedBinaryBytes, serializationFormat)
+		}
+
+		return &core.Literal{
+			Value: &core.Literal_Collection{
+				Collection: &core.LiteralCollection{
+					Literals: literals,
+				},
+			},
+		}, nil
+	}
+
 	// Marshal the current value to MessagePack bytes
 	resolvedBinaryBytes, err := msgpack.Marshal(currVal)
 	if err != nil {


Test patch:
diff --git a/flytepropeller/pkg/controller/nodes/attr_path_resolver_test.go b/flytepropeller/pkg/controller/nodes/attr_path_resolver_test.go
index 1467fc0ea4..e8e28ac08f 100644
--- a/flytepropeller/pkg/controller/nodes/attr_path_resolver_test.go
+++ b/flytepropeller/pkg/controller/nodes/attr_path_resolver_test.go
@@ -10,13 +10,14 @@ import (
 	"github.com/stretchr/testify/assert"
 	"google.golang.org/protobuf/types/known/structpb"
 
+	"github.com/flyteorg/flyte/flyteidl/clients/go/coreutils"
 	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core"
 	"github.com/flyteorg/flyte/flytepropeller/pkg/controller/nodes/errors"
 )
 
-// FlyteFile and FlyteDirectory represented as map[interface{}]interface{}
-type FlyteFile map[interface{}]interface{}
-type FlyteDirectory map[interface{}]interface{}
+// FlyteFile and FlyteDirectory represented as map[any]any
+type FlyteFile map[any]any
+type FlyteDirectory map[any]any
 
 // InnerDC struct (equivalent to InnerDC dataclass in Python)
 type InnerDC struct {
@@ -73,7 +74,7 @@ func NewScalarLiteral(value string) *core.Literal {
 	}
 }
 
-func NewStructFromMap(m map[string]interface{}) *structpb.Struct {
+func NewStructFromMap(m map[string]any) *structpb.Struct {
 	st, _ := structpb.NewStruct(m)
 	return st
 }
@@ -135,7 +136,7 @@ func TestResolveAttrPathInStruct(t *testing.T) {
 				Value: &core.Literal_Scalar{
 					Scalar: &core.Scalar{
 						Value: &core.Scalar_Generic{
-							Generic: NewStructFromMap(map[string]interface{}{"foo": "bar"}),
+							Generic: NewStructFromMap(map[string]any{"foo": "bar"}),
 						},
 					},
 				},
@@ -157,8 +158,8 @@ func TestResolveAttrPathInStruct(t *testing.T) {
 					Scalar: &core.Scalar{
 						Value: &core.Scalar_Generic{
 							Generic: NewStructFromMap(
-								map[string]interface{}{
-									"foo": []interface{}{"bar1", "bar2"},
+								map[string]any{
+									"foo": []any{"bar1", "bar2"},
 								},
 							),
 						},
@@ -187,8 +188,8 @@ func TestResolveAttrPathInStruct(t *testing.T) {
 					Scalar: &core.Scalar{
 						Value: &core.Scalar_Generic{
 							Generic: NewStructFromMap(
-								map[string]interface{}{
-									"foo": []interface{}{[]interface{}{"bar1", "bar2"}},
+								map[string]any{
+									"foo": []any{[]any{"bar1", "bar2"}},
 								},
 							),
 						},
@@ -236,7 +237,7 @@ func TestResolveAttrPathInStruct(t *testing.T) {
 												Value: &core.Literal_Scalar{
 													Scalar: &core.Scalar{
 														Value: &core.Scalar_Generic{
-															Generic: NewStructFromMap(map[string]interface{}{"bar": "car"}),
+															Generic: NewStructFromMap(map[string]any{"bar": "car"}),
 														},
 													},
 												},
@@ -276,9 +277,9 @@ func TestResolveAttrPathInStruct(t *testing.T) {
 					Scalar: &core.Scalar{
 						Value: &core.Scalar_Generic{
 							Generic: NewStructFromMap(
-								map[string]interface{}{
-									"foo": map[string]interface{}{
-										"bar": map[string]interface{}{
+								map[string]any{
+									"foo": map[string]any{
+										"bar": map[string]any{
 											"baz": 42,
 										},
 									},
@@ -306,7 +307,7 @@ func TestResolveAttrPathInStruct(t *testing.T) {
 					Scalar: &core.Scalar{
 						Value: &core.Scalar_Generic{
 							Generic: NewStructFromMap(
-								map[string]interface{}{
+								map[string]any{
 									"baz": 42,
 								},
 							),
@@ -365,7 +366,7 @@ func TestResolveAttrPathInStruct(t *testing.T) {
 				Value: &core.Literal_Scalar{
 					Scalar: &core.Scalar{
 						Value: &core.Scalar_Generic{
-							Generic: NewStructFromMap(map[string]interface{}{"foo": "bar"}),
+							Generic: NewStructFromMap(map[string]any{"foo": "bar"}),
 						},
 					},
 				},
@@ -387,8 +388,8 @@ func TestResolveAttrPathInStruct(t *testing.T) {
 					Scalar: &core.Scalar{
 						Value: &core.Scalar_Generic{
 							Generic: NewStructFromMap(
-								map[string]interface{}{
-									"foo": []interface{}{"bar1", "bar2"},
+								map[string]any{
+									"foo": []any{"bar1", "bar2"},
 								},
 							),
 						},
@@ -495,11 +496,35 @@ func createNestedDC() DC {
 
 func TestResolveAttrPathInBinary(t *testing.T) {
 	// Helper function to convert a map to msgpack bytes and then to BinaryIDL
-	toMsgpackBytes := func(m interface{}) []byte {
+	toMsgpackBytes := func(m any) []byte {
 		msgpackBytes, err := msgpack.Marshal(m)
 		assert.NoError(t, err)
 		return msgpackBytes
 	}
+	toLiteralCollectionWithMsgpackBytes := func(collection []any) *core.Literal {
+		literals := make([]*core.Literal, len(collection))
+		for i, v := range collection {
+			resolvedBinaryBytes, _ := msgpack.Marshal(v)
+			literals[i] = constructResolvedBinary(resolvedBinaryBytes, coreutils.MESSAGEPACK)
+		}
+		return &core.Literal{
+			Value: &core.Literal_Collection{
+				Collection: &core.LiteralCollection{
+					Literals: literals,
+				},
+			},
+		}
+	}
+	fromLiteralCollectionWithMsgpackBytes := func(lv *core.Literal) []any {
+		literals := lv.GetCollection().GetLiterals()
+		collection := make([]any, len(literals))
+		for i, l := range literals {
+			var v any
+			_ = msgpack.Unmarshal(l.GetScalar().GetBinary().Value, &v)
+			collection[i] = v
+		}
+		return collection
+	}
 
 	flyteFile := FlyteFile{
 		"path": "s3://my-s3-bucket/example.txt",
@@ -630,18 +655,7 @@ func TestResolveAttrPathInBinary(t *testing.T) {
 					},
 				},
 			},
-			expected: &core.Literal{
-				Value: &core.Literal_Scalar{
-					Scalar: &core.Scalar{
-						Value: &core.Scalar_Binary{
-							Binary: &core.Binary{
-								Value: toMsgpackBytes([]int{0, 1, 2, -1, -2}),
-								Tag:   "msgpack",
-							},
-						},
-					},
-				},
-			},
+			expected: toLiteralCollectionWithMsgpackBytes([]any{0, 1, 2, -1, -2}),
 			hasError: false,
 		},
 		{
@@ -653,18 +667,7 @@ func TestResolveAttrPathInBinary(t *testing.T) {
 					},
 				},
 			},
-			expected: &core.Literal{
-				Value: &core.Literal_Scalar{
-					Scalar: &core.Scalar{
-						Value: &core.Scalar_Binary{
-							Binary: &core.Binary{
-								Value: toMsgpackBytes([]FlyteFile{{"path": "s3://my-s3-bucket/example.txt"}}),
-								Tag:   "msgpack",
-							},
-						},
-					},
-				},
-			},
+			expected: toLiteralCollectionWithMsgpackBytes([]any{flyteFile}),
 			hasError: false,
 		},
 		{
@@ -676,18 +679,7 @@ func TestResolveAttrPathInBinary(t *testing.T) {
 					},
 				},
 			},
-			expected: &core.Literal{
-				Value: &core.Literal_Scalar{
-					Scalar: &core.Scalar{
-						Value: &core.Scalar_Binary{
-							Binary: &core.Binary{
-								Value: toMsgpackBytes([][]int{{0}, {1}, {-1}}),
-								Tag:   "msgpack",
-							},
-						},
-					},
-				},
-			},
+			expected: toLiteralCollectionWithMsgpackBytes([]any{[]int{0}, []int{1}, []int{-1}}),
 			hasError: false,
 		},
 		{
@@ -699,18 +691,8 @@ func TestResolveAttrPathInBinary(t *testing.T) {
 					},
 				},
 			},
-			expected: &core.Literal{
-				Value: &core.Literal_Scalar{
-					Scalar: &core.Scalar{
-						Value: &core.Scalar_Binary{
-							Binary: &core.Binary{
-								Value: toMsgpackBytes([]map[int]bool{{0: false}, {1: true}, {-1: true}}),
-								Tag:   "msgpack",
-							},
-						},
-					},
-				},
-			},
+			expected: toLiteralCollectionWithMsgpackBytes([]any{map[int]bool{0: false}, map[int]bool{1: true},
+				map[int]bool{-1: true}}),
 			hasError: false,
 		},
 		{
@@ -1037,18 +1019,7 @@ func TestResolveAttrPathInBinary(t *testing.T) {
 					},
 				},
 			},
-			expected: &core.Literal{
-				Value: &core.Literal_Scalar{
-					Scalar: &core.Scalar{
-						Value: &core.Scalar_Binary{
-							Binary: &core.Binary{
-								Value: toMsgpackBytes([]int{0, 1, 2, -1, -2}),
-								Tag:   "msgpack",
-							},
-						},
-					},
-				},
-			},
+			expected: toLiteralCollectionWithMsgpackBytes([]any{0, 1, 2, -1, -2}),
 			hasError: false,
 		},
 		{
@@ -1065,18 +1036,7 @@ func TestResolveAttrPathInBinary(t *testing.T) {
 					},
 				},
 			},
-			expected: &core.Literal{
-				Value: &core.Literal_Scalar{
-					Scalar: &core.Scalar{
-						Value: &core.Scalar_Binary{
-							Binary: &core.Binary{
-								Value: toMsgpackBytes([]FlyteFile{flyteFile}),
-								Tag:   "msgpack",
-							},
-						},
-					},
-				},
-			},
+			expected: toLiteralCollectionWithMsgpackBytes([]any{flyteFile}),
 			hasError: false,
 		},
 		{
@@ -1093,18 +1053,7 @@ func TestResolveAttrPathInBinary(t *testing.T) {
 					},
 				},
 			},
-			expected: &core.Literal{
-				Value: &core.Literal_Scalar{
-					Scalar: &core.Scalar{
-						Value: &core.Scalar_Binary{
-							Binary: &core.Binary{
-								Value: toMsgpackBytes([][]int{{0}, {1}, {-1}}),
-								Tag:   "msgpack",
-							},
-						},
-					},
-				},
-			},
+			expected: toLiteralCollectionWithMsgpackBytes([]any{[]int{0}, []int{1}, []int{-1}}),
 			hasError: false,
 		},
 		{
@@ -1126,18 +1075,7 @@ func TestResolveAttrPathInBinary(t *testing.T) {
 					},
 				},
 			},
-			expected: &core.Literal{
-				Value: &core.Literal_Scalar{
-					Scalar: &core.Scalar{
-						Value: &core.Scalar_Binary{
-							Binary: &core.Binary{
-								Value: toMsgpackBytes([]int{0}),
-								Tag:   "msgpack",
-							},
-						},
-					},
-				},
-			},
+			expected: toLiteralCollectionWithMsgpackBytes([]any{0}),
 			hasError: false,
 		},
 		{
@@ -1192,18 +1130,11 @@ func TestResolveAttrPathInBinary(t *testing.T) {
 					},
 				},
 			},
-			expected: &core.Literal{
-				Value: &core.Literal_Scalar{
-					Scalar: &core.Scalar{
-						Value: &core.Scalar_Binary{
-							Binary: &core.Binary{
-								Value: toMsgpackBytes([]map[int]bool{{0: false}, {1: true}, {-1: true}}),
-								Tag:   "msgpack",
-							},
-						},
-					},
-				},
-			},
+			expected: toLiteralCollectionWithMsgpackBytes([]any{
+				map[int]bool{0: false},
+				map[int]bool{1: true},
+				map[int]bool{-1: true},
+			}),
 			hasError: false,
 		},
 		{
@@ -1422,10 +1353,10 @@ func TestResolveAttrPathInBinary(t *testing.T) {
 					Scalar: &core.Scalar{
 						Value: &core.Scalar_Binary{
 							Binary: &core.Binary{
-								Value: toMsgpackBytes(map[string]interface{}{
-									"foo": map[string]interface{}{
+								Value: toMsgpackBytes(map[string]any{
+									"foo": map[string]any{
 										"bar": int64(42),
-										"baz": map[string]interface{}{
+										"baz": map[string]any{
 											"qux":  3.14,
 											"quux": "str",
 										},
@@ -1465,8 +1396,8 @@ func TestResolveAttrPathInBinary(t *testing.T) {
 					Scalar: &core.Scalar{
 						Value: &core.Scalar_Binary{
 							Binary: &core.Binary{
-								Value: toMsgpackBytes(map[string]interface{}{
-									"foo": []interface{}{int64(42), 3.14, "str"},
+								Value: toMsgpackBytes(map[string]any{
+									"foo": []any{int64(42), 3.14, "str"},
 								}),
 								Tag: "msgpack",
 							},
@@ -1499,13 +1430,13 @@ func TestResolveAttrPathInBinary(t *testing.T) {
 			assert.Error(t, err, i)
 			assert.ErrorContains(t, err, errors.PromiseAttributeResolveError, i)
 		} else {
-			var expectedValue, actualValue interface{}
+			var expectedValue, actualValue any
 
-			// Helper function to unmarshal a Binary Literal into an interface{}
-			unmarshalBinaryLiteral := func(literal *core.Literal) (interface{}, error) {
+			// Helper function to unmarshal a Binary Literal into an any
+			unmarshalBinaryLiteral := func(literal *core.Literal) (any, error) {
 				if scalar, ok := literal.Value.(*core.Literal_Scalar); ok {
 					if binary, ok := scalar.Scalar.Value.(*core.Scalar_Binary); ok {
-						var value interface{}
+						var value any
 						err := msgpack.Unmarshal(binary.Binary.Value, &value)
 						return value, err
 					}
@@ -1513,16 +1444,22 @@ func TestResolveAttrPathInBinary(t *testing.T) {
 				return nil, fmt.Errorf("literal is not a Binary Scalar")
 			}
 
-			// Unmarshal the expected value
-			expectedValue, err := unmarshalBinaryLiteral(arg.expected)
-			if err != nil {
-				t.Fatalf("Failed to unmarshal expected value in test case %d: %v", i, err)
+			if arg.expected.GetCollection() != nil {
+				expectedValue = fromLiteralCollectionWithMsgpackBytes(arg.expected)
+			} else {
+				expectedValue, err = unmarshalBinaryLiteral(arg.expected)
+				if err != nil {
+					t.Fatalf("Failed to unmarshal expected value in test case %d: %v", i, err)
+				}
 			}
 
-			// Unmarshal the resolved value
-			actualValue, err = unmarshalBinaryLiteral(resolved)
-			if err != nil {
-				t.Fatalf("Failed to unmarshal resolved value in test case %d: %v", i, err)
+			if resolved.GetCollection() != nil {
+				actualValue = fromLiteralCollectionWithMsgpackBytes(resolved)
+			} else {
+				actualValue, err = unmarshalBinaryLiteral(resolved)
+				if err != nil {
+					t.Fatalf("Failed to unmarshal resolved value in test case %d: %v", i, err)
+				}
 			}
 
 			// Deeply compare the expected and actual values, ignoring map ordering