Adding support for structured dataset (#369)

Signed-off-by: pmahindrakar-oss <prafulla.mahindrakar@gmail.com>
flyteorg · Feb 22, 2023 · e69fd54 · e69fd54
1 parent f3724b4
commit e69fd54
Show file tree

Hide file tree

Showing 4 changed files with 88 additions and 0 deletions.
diff --git a/clients/go/coreutils/extract_literal.go b/clients/go/coreutils/extract_literal.go
@@ -60,6 +60,8 @@ func ExtractFromLiteral(literal *core.Literal) (interface{}, error) {
 			return scalarValue.Schema.Uri, nil
 		case *core.Scalar_Generic:
 			return scalarValue.Generic, nil
+		case *core.Scalar_StructuredDataset:
+			return scalarValue.StructuredDataset.Uri, nil
 		default:
 			return nil, fmt.Errorf("unsupported literal scalar type %T", scalarValue)
 		}

diff --git a/clients/go/coreutils/extract_literal_test.go b/clients/go/coreutils/extract_literal_test.go
@@ -176,4 +176,27 @@ func TestFetchLiteral(t *testing.T) {
 			assert.Equal(t, val.Kind, extractedStructValue.Fields[key].Kind)
 		}
 	})
+
+	t.Run("Structured dataset", func(t *testing.T) {
+		literalVal := "s3://blah/blah/blah"
+		var dataSetColumns []*core.StructuredDatasetType_DatasetColumn
+		dataSetColumns = append(dataSetColumns, &core.StructuredDatasetType_DatasetColumn{
+			Name: "Price",
+			LiteralType: &core.LiteralType{
+				Type: &core.LiteralType_Simple{
+					Simple: core.SimpleType_FLOAT,
+				},
+			},
+		})
+		var literalType = &core.LiteralType{Type: &core.LiteralType_StructuredDatasetType{StructuredDatasetType: &core.StructuredDatasetType{
+			Columns: dataSetColumns,
+			Format:  "testFormat",
+		}}}
+
+		lit, err := MakeLiteralForType(literalType, literalVal)
+		assert.NoError(t, err)
+		extractedLiteralVal, err := ExtractFromLiteral(lit)
+		assert.NoError(t, err)
+		assert.Equal(t, literalVal, extractedLiteralVal)
+	})
 }
diff --git a/clients/go/coreutils/literals.go b/clients/go/coreutils/literals.go
@@ -436,6 +436,26 @@ func MakeLiteralForSchema(path storage.DataReference, columns []*core.SchemaType
 	}
 }
 
+func MakeLiteralForStructuredDataSet(path storage.DataReference, columns []*core.StructuredDatasetType_DatasetColumn, format string) *core.Literal {
+	return &core.Literal{
+		Value: &core.Literal_Scalar{
+			Scalar: &core.Scalar{
+				Value: &core.Scalar_StructuredDataset{
+					StructuredDataset: &core.StructuredDataset{
+						Uri: path.String(),
+						Metadata: &core.StructuredDatasetMetadata{
+							StructuredDatasetType: &core.StructuredDatasetType{
+								Columns: columns,
+								Format:  format,
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+}
+
 func MakeLiteralForBlob(path storage.DataReference, isDir bool, format string) *core.Literal {
 	dim := core.BlobType_SINGLE
 	if isDir {
@@ -538,6 +558,9 @@ func MakeLiteralForType(t *core.LiteralType, v interface{}) (*core.Literal, erro
 	case *core.LiteralType_Schema:
 		lv := MakeLiteralForSchema(storage.DataReference(fmt.Sprintf("%v", v)), newT.Schema.Columns)
 		return lv, nil
+	case *core.LiteralType_StructuredDatasetType:
+		lv := MakeLiteralForStructuredDataSet(storage.DataReference(fmt.Sprintf("%v", v)), newT.StructuredDatasetType.Columns, newT.StructuredDatasetType.Format)
+		return lv, nil
 
 	case *core.LiteralType_EnumType:
 		var newV string

diff --git a/clients/go/coreutils/literals_test.go b/clients/go/coreutils/literals_test.go
@@ -675,4 +675,44 @@ func TestMakeLiteralForType(t *testing.T) {
 		assert.NoError(t, err)
 		assert.Equal(t, "", l.GetScalar().GetPrimitive().GetStringValue())
 	})
+
+	t.Run("Structured Data Set", func(t *testing.T) {
+		var dataSetColumns []*core.StructuredDatasetType_DatasetColumn
+		dataSetColumns = append(dataSetColumns, &core.StructuredDatasetType_DatasetColumn{
+			Name: "Price",
+			LiteralType: &core.LiteralType{
+				Type: &core.LiteralType_Simple{
+					Simple: core.SimpleType_FLOAT,
+				},
+			},
+		})
+		var literalType = &core.LiteralType{Type: &core.LiteralType_StructuredDatasetType{StructuredDatasetType: &core.StructuredDatasetType{
+			Columns: dataSetColumns,
+			Format:  "testFormat",
+		}}}
+
+		expectedLV := &core.Literal{Value: &core.Literal_Scalar{Scalar: &core.Scalar{
+			Value: &core.Scalar_StructuredDataset{
+				StructuredDataset: &core.StructuredDataset{
+					Uri: "s3://blah/blah/blah",
+					Metadata: &core.StructuredDatasetMetadata{
+						StructuredDatasetType: &core.StructuredDatasetType{
+							Columns: dataSetColumns,
+							Format:  "testFormat",
+						},
+					},
+				},
+			},
+		}}}
+		lv, err := MakeLiteralForType(literalType, "s3://blah/blah/blah")
+		assert.NoError(t, err)
+
+		assert.Equal(t, expectedLV, lv)
+
+		expectedVal, err := ExtractFromLiteral(expectedLV)
+		assert.NoError(t, err)
+		actualVal, err := ExtractFromLiteral(lv)
+		assert.NoError(t, err)
+		assert.Equal(t, expectedVal, actualVal)
+	})
 }