feat: enable handling of URNs alongside URIs

errors.go

@@ -5,13 +5,16 @@ package substraitgo
 import "errors"
 
 var (
-	ErrNotImplemented    = errors.New("not implemented")
-	ErrInvalidType       = errors.New("invalid type")
-	ErrInvalidExpr       = errors.New("invalid expression")
-	ErrNotFound          = errors.New("not found")
-	ErrKeyExists         = errors.New("key already exists")
-	ErrInvalidRel        = errors.New("invalid relation")
-	ErrInvalidArg        = errors.New("invalid argument")
-	ErrInvalidInputCount = errors.New("invalid input count")
-	ErrInvalidDialect    = errors.New("invalid dialect")
+	ErrNotImplemented            = errors.New("not implemented")
+	ErrInvalidType               = errors.New("invalid type")
+	ErrInvalidExpr               = errors.New("invalid expression")
+	ErrNotFound                  = errors.New("not found")
+	ErrKeyExists                 = errors.New("key already exists")
+	ErrInvalidRel                = errors.New("invalid relation")
+	ErrInvalidArg                = errors.New("invalid argument")
+	ErrInvalidInputCount         = errors.New("invalid input count")
+	ErrInvalidDialect            = errors.New("invalid dialect")
+	ErrInvalidSimpleExtention    = errors.New("invalid simple extension")
+	ErrInvalidPlan               = errors.New("invalid plan")
+	ErrExtensionURINotResolvable = errors.New("extension URI not resolvable")
 )

expr/binding_test.go

@@ -14,33 +14,33 @@ import (
 var (
 	extReg    = NewEmptyExtensionRegistry(extensions.GetDefaultCollectionWithNoError())
 	uPointRef = extReg.GetTypeAnchor(extensions.ID{
-		URI:  extensions.SubstraitDefaultURIPrefix + "extension_types.yaml",
+		URN:  extensions.SubstraitDefaultURNPrefix + "extension_types",
 		Name: "point",
 	})
 
 	subID = extensions.ID{
-		URI:  extensions.SubstraitDefaultURIPrefix + "functions_arithmetic.yaml",
+		URN:  extensions.SubstraitDefaultURNPrefix + "functions_arithmetic",
 		Name: "subtract"}
 	addID = extensions.ID{
-		URI:  extensions.SubstraitDefaultURIPrefix + "functions_arithmetic.yaml",
+		URN:  extensions.SubstraitDefaultURNPrefix + "functions_arithmetic",
 		Name: "add"}
 	indexInID = extensions.ID{
-		URI:  extensions.SubstraitDefaultURIPrefix + "functions_set.yaml",
+		URN:  extensions.SubstraitDefaultURNPrefix + "functions_set",
 		Name: "index_in"}
 	rankID = extensions.ID{
-		URI:  extensions.SubstraitDefaultURIPrefix + "functions_arithmetic.yaml",
+		URN:  extensions.SubstraitDefaultURNPrefix + "functions_arithmetic",
 		Name: "rank"}
 	firstValueID = extensions.ID{
-		URI:  extensions.SubstraitDefaultURIPrefix + "functions_arithmetic.yaml",
+		URN:  extensions.SubstraitDefaultURNPrefix + "functions_arithmetic",
 		Name: "first_value"}
 	extractID = extensions.ID{
-		URI:  extensions.SubstraitDefaultURIPrefix + "functions_datetime.yaml",
+		URN:  extensions.SubstraitDefaultURNPrefix + "functions_datetime",
 		Name: "extract"}
 	ntileID = extensions.ID{
-		URI:  extensions.SubstraitDefaultURIPrefix + "functions_arithmetic.yaml",
+		URN:  extensions.SubstraitDefaultURNPrefix + "functions_arithmetic",
 		Name: "ntile"}
 	sumID = extensions.ID{
-		URI:  extensions.SubstraitDefaultURIPrefix + "functions_arithmetic.yaml",
+		URN:  extensions.SubstraitDefaultURNPrefix + "functions_arithmetic",
 		Name: "sum"}
 
 	boringSchema = types.NamedStruct{

expr/builder_test.go

@@ -66,7 +66,7 @@ func TestExprBuilder(t *testing.T) {
 			b.ScalarFunc(subID).Args(b.RootRef(expr.NewStructFieldRef(3)),
 				b.Wrap(expr.NewLiteral(int32(3), false))), ""},
 		{"window func", "",
-			b.WindowFunc(rankID), "invalid expression: non-decomposable window or agg function '{https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml rank}' must use InitialToResult phase"},
+			b.WindowFunc(rankID), "invalid expression: non-decomposable window or agg function '{extension:io.substrait:functions_arithmetic rank}' must use InitialToResult phase"},
 		{"window func", "rank(; phase: AGGREGATION_PHASE_INITIAL_TO_RESULT, invocation: AGGREGATION_INVOCATION_UNSPECIFIED) => i64?",
 			b.WindowFunc(rankID).Phase(types.AggPhaseInitialToResult), ""},
 		{"window func",
@@ -118,6 +118,7 @@ func TestExprBuilder(t *testing.T) {
 func TestCustomTypesInFunctionOutput(t *testing.T) {
 	custom := `%YAML 1.2
 ---
+urn: extension:test:custom
 types:
   - name: custom_type1
   - name: custom_type2
@@ -159,9 +160,9 @@ window_functions:
 
 	planBuilder := plan.NewBuilder(&collection)
 
-	customType1 := planBuilder.UserDefinedType("custom", "custom_type1")
-	customType2 := planBuilder.UserDefinedType("custom", "custom_type2")
-	customType3 := planBuilder.UserDefinedType("custom", "custom_type3")
+	customType1 := planBuilder.UserDefinedType("extension:test:custom", "custom_type1")
+	customType2 := planBuilder.UserDefinedType("extension:test:custom", "custom_type2")
+	customType3 := planBuilder.UserDefinedType("extension:test:custom", "custom_type3")
 
 	anyVal, err := anypb.New(expr.NewPrimitiveLiteral("foo", false).ToProto())
 	require.NoError(t, err)
@@ -173,7 +174,7 @@ window_functions:
 
 	// check scalar function
 	scalar, err := planBuilder.GetExprBuilder().ScalarFunc(extensions.ID{
-		URI:  "custom",
+		URN:  "extension:test:custom",
 		Name: "custom_function",
 	}).Args(
 		customLiteral,
@@ -188,7 +189,7 @@ window_functions:
 
 	// check aggregate function
 	aggr, err := planBuilder.GetExprBuilder().AggFunc(extensions.ID{
-		URI:  "custom",
+		URN:  "extension:test:custom",
 		Name: "custom_aggr",
 	}).Args(
 		customLiteral,
@@ -202,7 +203,7 @@ window_functions:
 
 	// check window function
 	window, err := planBuilder.GetExprBuilder().WindowFunc(extensions.ID{
-		URI:  "custom",
+		URN:  "extension:test:custom",
 		Name: "custom_window",
 	}).Args(
 		customLiteral,

expr/expression.go

@@ -373,7 +373,7 @@ type Expression interface {
 	//     // it's a pre-order traversal
 	//     if f, ok := e.(*ScalarFunction); ok {
 	//       return &ScalarFunction{
-	//         ID: ExtID{URI: "some other uri", Name: "some other func"},
+	//         ID: ExtID{URN: "some other urn", Name: "some other func"},
 	//         Args: f.Args,
 	//         Options: f.Options,
 	//         OutputType: f.OutputType,
@@ -1565,14 +1565,19 @@ type Extended struct {
 	BaseSchema       types.NamedStruct
 	AdvancedExts     *extensions.AdvancedExtension
 	ExpectedTypeURLs []string
+
+	reg ExtensionRegistry
 }
 
 func ExtendedFromProto(ex *proto.ExtendedExpression, c *extensions.Collection) (*Extended, error) {
+	extSet, err := extensions.GetExtensionSet(ex, c)
+	if err != nil {
+		return nil, err
+	}
 	var (
-		base   = types.NewNamedStructFromProto(ex.BaseSchema)
-		extSet = extensions.GetExtensionSet(ex)
-		reg    = NewExtensionRegistry(extSet, c)
-		refs   = make([]ExpressionReference, len(ex.ReferredExpr))
+		base = types.NewNamedStructFromProto(ex.BaseSchema)
+		reg  = NewExtensionRegistry(extSet, c)
+		refs = make([]ExpressionReference, len(ex.ReferredExpr))
 	)
 
 	for i, r := range ex.ReferredExpr {
@@ -1602,18 +1607,20 @@ func ExtendedFromProto(ex *proto.ExtendedExpression, c *extensions.Collection) (
 		BaseSchema:       base,
 		AdvancedExts:     ex.AdvancedExtensions,
 		ExpectedTypeURLs: ex.ExpectedTypeUrls,
+		reg:              reg,
 	}, nil
 }
 
 func (ex *Extended) ToProto() *proto.ExtendedExpression {
-	uris, decls := ex.Extensions.ToProto()
+	urns, uris, decls := ex.Extensions.ToProto(ex.reg.Collection())
 	refs := make([]*proto.ExpressionReference, len(ex.ReferredExpr))
 	for i, ref := range ex.ReferredExpr {
 		refs[i] = ref.ToProto()
 	}
 
 	return &proto.ExtendedExpression{
 		Version:            ex.Version,
+		ExtensionUrns:      urns,
 		ExtensionUris:      uris,
 		Extensions:         decls,
 		BaseSchema:         ex.BaseSchema.ToProto(),