diff --git a/clients/go/coreutils/extract_literal.go b/clients/go/coreutils/extract_literal.go index 57c517945..afc7fd2a0 100644 --- 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 index 73ae7a7cf..32d392322 100644 --- 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 index 103a44cf4..4ad84f181 100644 --- 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 index 7b156e81e..4b16478ad 100644 --- 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) + }) }