feat(type): Add precision to/from proto for IntervalDayType

* Note, this is a breaking change. Earlier consumer's of IntervalDayType on wire * precision was Microsecond (since not set). Now if Precision will be Seconds (if not set)
anshuldata · Oct 16, 2024 · 132dc66 · 132dc66
1 parent 002720e
commit 132dc66
Show file tree

Hide file tree

Showing 6 changed files with 34 additions and 7 deletions.
diff --git a/expr/literals.go b/expr/literals.go
@@ -812,9 +812,21 @@ func LiteralFromProto(l *proto.Expression_Literal) Literal {
 				Nullability:      nullability,
 			}}
 	case *proto.Expression_Literal_IntervalDayToSecond_:
+		precision := types.PrecisionMicroSeconds
+		switch lit.IntervalDayToSecond.PrecisionMode.(type) {
+		case *proto.Expression_Literal_IntervalDayToSecond_Precision:
+			var err error
+			precision, err = types.ProtoToTimePrecision(lit.IntervalDayToSecond.GetPrecision())
+			if err != nil {
+				return nil
+			}
+		default:
+			return nil
+		}
 		return &ProtoLiteral{
 			Value: lit.IntervalDayToSecond,
 			Type: &types.IntervalDayType{
+				Length:           precision.ToProtoVal(),
 				Nullability:      nullability,
 				TypeVariationRef: l.TypeVariationReference,
 			},

diff --git a/expr/proto_literals_test.go b/expr/proto_literals_test.go
@@ -36,6 +36,7 @@ func TestToProtoLiteral(t *testing.T) {
 }
 
 func TestLiteralFromProtoLiteral(t *testing.T) {
+	intDayToSecVal := &proto.Expression_Literal_IntervalDayToSecond{Days: 1, Seconds: 2, PrecisionMode: &proto.Expression_Literal_IntervalDayToSecond_Precision{Precision: 5}}
 	for _, tc := range []struct {
 		name             string
 		constructedProto *proto.Expression_Literal
@@ -49,6 +50,10 @@ func TestLiteralFromProtoLiteral(t *testing.T) {
 			&proto.Expression_Literal{LiteralType: &proto.Expression_Literal_PrecisionTimestampTz{PrecisionTimestampTz: &proto.Expression_Literal_PrecisionTimestamp{Precision: 9, Value: 12345678}}, Nullable: true},
 			&ProtoLiteral{Value: int64(12345678), Type: types.NewPrecisionTimestampTzType(types.PrecisionNanoSeconds).WithNullability(types.NullabilityNullable)},
 		},
+		{"IntervalDayType",
+			&proto.Expression_Literal{LiteralType: &proto.Expression_Literal_IntervalDayToSecond_{IntervalDayToSecond: intDayToSecVal}, Nullable: true},
+			&ProtoLiteral{Value: intDayToSecVal, Type: &types.IntervalDayType{Length: types.PrecisionEMinus5Seconds.ToProtoVal(), Nullability: types.NullabilityNullable}},
+		},
 		{"IntervalYearToMonthType",
 			&proto.Expression_Literal{LiteralType: &proto.Expression_Literal_IntervalYearToMonth_{IntervalYearToMonth: &proto.Expression_Literal_IntervalYearToMonth{Years: 1234, Months: 5}}, Nullable: true},
 			IntervalYearToMonthLiteral{Years: 1234, Months: 5, Nullability: types.NullabilityNullable},

diff --git a/extensions/extension_mgr.go b/extensions/extension_mgr.go
@@ -21,12 +21,8 @@ type AdvancedExtension = extensions.AdvancedExtension
 const SubstraitDefaultURIPrefix = "https://github.com/substrait-io/substrait/blob/main/extensions/"
 
 // DefaultCollection is loaded with the default Substrait extension
-// definitions.
-// Parser needs to enhanced until than below function files are not processed
-// 1. functions_arithmetic_decimal.yaml (need to parse complex return type)
-// 2. functions_geometry.yaml (need to parse geometry type)
-// 3. type_variations.yaml
-// 4. unknown.yaml
+// definitions. Not all files are currently parsable.
+// Parser needs to enhanced to support all files
 var DefaultCollection Collection
 
 func init() {

diff --git a/types/parser/type_parser_test.go b/types/parser/type_parser_test.go
@@ -80,7 +80,7 @@ func TestParserRetType(t *testing.T) {
 		shortName   string
 		expectedTyp types.Type
 	}{
-		{"interval_day?<1>", "interval_day?<1>", "iday", &types.IntervalDayType{}},
+		{"interval_day?<1>", "interval_day?<1>", "iday", &types.IntervalDayType{Length: 1, Nullability: types.NullabilityNullable}},
 	}
 
 	p, err := parser.New()
@@ -96,6 +96,7 @@ func TestParserRetType(t *testing.T) {
 				retType, err := d.Expr.(*parser.Type).RetType()
 				assert.NoError(t, err)
 				assert.Equal(t, reflect.TypeOf(td.expectedTyp), reflect.TypeOf(retType))
+				assert.True(t, td.expectedTyp.Equals(retType))
 			}
 		})
 	}

diff --git a/types/types.go b/types/types.go
@@ -253,9 +253,18 @@ func TypeFromProto(t *proto.Type) Type {
 			TypeVariationRef: t.IntervalYear.TypeVariationReference,
 		}
 	case *proto.Type_IntervalDay_:
+		var precision = PrecisionMicroSeconds
+		if t.IntervalDay.Precision != nil {
+			var err error
+			precision, err = ProtoToTimePrecision(*t.IntervalDay.Precision)
+			if err != nil {
+				panic(fmt.Sprintf("Invalid precision %v", err))
+			}
+		}
 		return &IntervalDayType{
 			Nullability:      t.IntervalDay.Nullability,
 			TypeVariationRef: t.IntervalDay.TypeVariationReference,
+			Length:           precision.ToProtoVal(),
 		}
 	case *proto.Type_TimestampTz:
 		return &TimestampTzType{
@@ -492,8 +501,10 @@ func TypeToProto(t Type) *proto.Type {
 				Nullability:            t.Nullability,
 				TypeVariationReference: t.TypeVariationRef}}}
 	case *IntervalDayType:
+		precision := t.Length
 		return &proto.Type{Kind: &proto.Type_IntervalDay_{
 			IntervalDay: &proto.Type_IntervalDay{
+				Precision:              &precision,
 				Nullability:            t.Nullability,
 				TypeVariationReference: t.TypeVariationRef}}}
 	case *UUIDType:

diff --git a/types/types_test.go b/types/types_test.go
@@ -82,6 +82,8 @@ func TestTypeRoundtrip(t *testing.T) {
 				&FixedCharType{Nullability: n, Length: 25},
 				&VarCharType{Nullability: n, Length: 35},
 				&FixedBinaryType{Nullability: n, Length: 45},
+				&IntervalDayType{Nullability: n, Length: 5},
+				&IntervalDayType{Nullability: n},
 				&DecimalType{Nullability: n, Precision: 34, Scale: 3},
 				&MapType{Nullability: n, Key: &Int8Type{}, Value: &Int16Type{Nullability: n}},
 				&ListType{Nullability: n, Type: &TimeType{Nullability: n}},