Expose the supported types through the local type registry. (#122)

Author: EpsilonPrime

functions/registries.go

@@ -28,6 +28,9 @@ type LocalTypeRegistry interface {
 
 	//GetTypeClasses() []types.TypeClass // TODO
 
+	// GetSupportedTypes returns the types supported by this dialect.
+	GetSupportedTypes() map[string]types.Type
+
 	// IsTypeSupportedInTables checks whether a particular type is supported in tables.
 	// Some types (such as INTERVAL) may only be supported in literal contexts.
 	IsTypeSupportedInTables(typ types.Type) bool
@@ -46,6 +49,10 @@ type Dialect interface {
 	// the subset of types supported by this dialect. This will return an error if there are
 	// types declared in the dialect that are not available within the provided registry.
 	LocalizeTypeRegistry(registry TypeRegistry) (LocalTypeRegistry, error)
+
+	// GetLocalTypeRegistry returns the last created type registry using this dialect or constructs
+	// one using LocalizeTypeRegistry and a default type registry if one hasn't yet been made.
+	GetLocalTypeRegistry() (LocalTypeRegistry, error)
 }
 
 type FunctionName interface {

functions/types.go

@@ -214,6 +214,10 @@ func (t *localTypeRegistryImpl) GetLocalTypeFromSubstraitType(typ types.Type) (s
 	return "", substraitgo.ErrNotFound
 }
 
+func (t *localTypeRegistryImpl) GetSupportedTypes() map[string]types.Type {
+	return t.localNameToType
+}
+
 func (t *localTypeRegistryImpl) IsTypeSupportedInTables(typ types.Type) bool {
 	if ti, ok := t.typeInfoMap[typ.ShortString()]; ok {
 		return ti.supportedAsColumn

types/types_test.go

@@ -4,10 +4,13 @@ package types_test
 
 import (
 	"fmt"
+	"strings"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/substrait-io/substrait-go/v3/functions"
 	. "github.com/substrait-io/substrait-go/v3/types"
 	"github.com/substrait-io/substrait-go/v3/types/integer_parameters"
 )
@@ -497,3 +500,45 @@ func TestGetTimeValueByPrecision(t *testing.T) {
 		})
 	}
 }
+
+func TestGetSupportedTypes(t *testing.T) {
+	dialect, err := functions.LoadDialect("test_dialect",
+		strings.NewReader(`
+name: test_dialect
+type: sql
+dependencies:
+  arithmetic:
+    https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml
+supported_types:
+  fp64:
+    sql_type_name: float
+  bool:
+    sql_type_name: boolean
+  varchar:
+    sql_type_name: varchar
+  date:
+    sql_type_name: date
+  time:
+    sql_type_name: time
+  pts:
+    sql_type_name: timestamp
+  ptstz:
+    sql_type_name: timestamptz
+  dec:
+    sql_type_name: numeric
+scalar_functions:
+  - name: arithmetic.add
+    local_name: '+'
+    infix: true
+    required_options:
+      overflow: SILENT
+      rounding: TIE_TO_EVEN
+    supported_kernels:
+      - fp64_fp64
+`))
+	require.NoError(t, err)
+	typeRegistry, err := dialect.GetLocalTypeRegistry()
+	require.NoError(t, err)
+	st := typeRegistry.GetSupportedTypes()
+	assert.Len(t, st, 8)
+}