diff --git a/weed/mq/schema/schema_builder.go b/weed/mq/schema/schema_builder.go index 449820ecb..19229c419 100644 --- a/weed/mq/schema/schema_builder.go +++ b/weed/mq/schema/schema_builder.go @@ -19,11 +19,11 @@ type RecordTypeBuilder struct { recordType *schema_pb.RecordType } -func NewRecordTypeBuilder() *RecordTypeBuilder { +func RecordTypeBegin() *RecordTypeBuilder { return &RecordTypeBuilder{recordType: &schema_pb.RecordType{}} } -func (rtb *RecordTypeBuilder) Build() *schema_pb.RecordType { +func (rtb *RecordTypeBuilder) RecordTypeEnd() *schema_pb.RecordType { // be consistent with parquet.node.go `func (g Group) Fields() []Field` sort.Slice(rtb.recordType.Fields, func(i, j int) bool { return rtb.recordType.Fields[i].Name < rtb.recordType.Fields[j].Name @@ -39,10 +39,10 @@ func (rtb *RecordTypeBuilder) SetField(name string, scalarType *schema_pb.Type) return rtb } -func (rtb *RecordTypeBuilder) SetRecordField(name string, recordTypeBuilder *RecordTypeBuilder) *RecordTypeBuilder { +func (rtb *RecordTypeBuilder) SetRecordField(name string, recordType *schema_pb.RecordType) *RecordTypeBuilder { rtb.recordType.Fields = append(rtb.recordType.Fields, &schema_pb.Field{ Name: name, - Type: &schema_pb.Type{Kind: &schema_pb.Type_RecordType{RecordType: recordTypeBuilder.Build()}}, + Type: &schema_pb.Type{Kind: &schema_pb.Type_RecordType{RecordType: recordType}}, }) return rtb } diff --git a/weed/mq/schema/struct_to_schema_test.go b/weed/mq/schema/struct_to_schema_test.go index 84fb48326..71f8215dc 100644 --- a/weed/mq/schema/struct_to_schema_test.go +++ b/weed/mq/schema/struct_to_schema_test.go @@ -30,10 +30,10 @@ func TestStructToSchema(t *testing.T) { Field2 string }{}, }, - want: NewRecordTypeBuilder(). - SetField("Field1", TypeInteger). - SetField("Field2", TypeString). - Build(), + want: RecordTypeBegin(). + SetField("Field1", TypeInteger). + SetField("Field2", TypeString). + RecordTypeEnd(), }, { name: "simple list", @@ -43,10 +43,10 @@ func TestStructToSchema(t *testing.T) { Field2 string }{}, }, - want: NewRecordTypeBuilder(). - SetField("Field1", ListOf(TypeInteger)). - SetField("Field2", TypeString). - Build(), + want: RecordTypeBegin(). + SetField("Field1", ListOf(TypeInteger)). + SetField("Field2", TypeString). + RecordTypeEnd(), }, { name: "simple []byte", @@ -55,9 +55,9 @@ func TestStructToSchema(t *testing.T) { Field2 []byte }{}, }, - want: NewRecordTypeBuilder(). - SetField("Field2", TypeBytes). - Build(), + want: RecordTypeBegin(). + SetField("Field2", TypeBytes). + RecordTypeEnd(), }, { name: "nested simpe structs", @@ -70,13 +70,15 @@ func TestStructToSchema(t *testing.T) { } }{}, }, - want: NewRecordTypeBuilder(). + want: RecordTypeBegin(). SetField("Field1", TypeInteger). - SetRecordField("Field2", NewRecordTypeBuilder(). - SetField("Field3", TypeString). - SetField("Field4", TypeInteger), + SetRecordField("Field2", + RecordTypeBegin(). + SetField("Field3", TypeString). + SetField("Field4", TypeInteger). + RecordTypeEnd(), ). - Build(), + RecordTypeEnd(), }, { name: "nested struct type", @@ -93,17 +95,19 @@ func TestStructToSchema(t *testing.T) { } }{}, }, - want: NewRecordTypeBuilder(). + want: RecordTypeBegin(). SetField("Field1", TypeInteger). - SetRecordField("Field2", NewRecordTypeBuilder(). + SetRecordField("Field2", RecordTypeBegin(). SetField("Field3", TypeString). SetField("Field4", ListOf(TypeInteger)). - SetRecordField("Field5", NewRecordTypeBuilder(). - SetField("Field6", TypeString). - SetField("Field7", TypeBytes), - ), + SetRecordField("Field5", + RecordTypeBegin(). + SetField("Field6", TypeString). + SetField("Field7", TypeBytes). + RecordTypeEnd(), + ).RecordTypeEnd(), ). - Build(), + RecordTypeEnd(), }, } diff --git a/weed/mq/schema/to_parquet_levels_test.go b/weed/mq/schema/to_parquet_levels_test.go index b7ed4dde1..65a90f956 100644 --- a/weed/mq/schema/to_parquet_levels_test.go +++ b/weed/mq/schema/to_parquet_levels_test.go @@ -18,16 +18,21 @@ func TestToParquetLevels(t *testing.T) { { name: "nested type", args: args{ - NewRecordTypeBuilder(). + RecordTypeBegin(). SetField("ID", TypeLong). SetField("CreatedAt", TypeLong). - SetRecordField("Person", NewRecordTypeBuilder(). - SetField("zName", TypeString). - SetField("emails", ListOf(TypeString))). + SetRecordField("Person", + RecordTypeBegin(). + SetField("zName", TypeString). + SetField("emails", ListOf(TypeString)). + RecordTypeEnd()). SetField("Company", TypeString). - SetRecordField("Address", NewRecordTypeBuilder(). - SetField("Street", TypeString). - SetField("City", TypeString)).Build(), + SetRecordField("Address", + RecordTypeBegin(). + SetField("Street", TypeString). + SetField("City", TypeString). + RecordTypeEnd()). + RecordTypeEnd(), }, want: &ParquetLevels{ startColumnIndex: 0, diff --git a/weed/mq/schema/write_parquet_test.go b/weed/mq/schema/write_parquet_test.go index 9402530a6..42a1d0916 100644 --- a/weed/mq/schema/write_parquet_test.go +++ b/weed/mq/schema/write_parquet_test.go @@ -12,16 +12,21 @@ import ( func TestWriteReadParquet(t *testing.T) { // create a schema_pb.RecordType - recordType := NewRecordTypeBuilder(). + recordType := RecordTypeBegin(). SetField("ID", TypeLong). SetField("CreatedAt", TypeLong). - SetRecordField("Person", NewRecordTypeBuilder(). - SetField("zName", TypeString). - SetField("emails", ListOf(TypeString))). + SetRecordField("Person", + RecordTypeBegin(). + SetField("zName", TypeString). + SetField("emails", ListOf(TypeString)). + RecordTypeEnd()). SetField("Company", TypeString). - SetRecordField("Address", NewRecordTypeBuilder(). - SetField("Street", TypeString). - SetField("City", TypeString)).Build() + SetRecordField("Address", + RecordTypeBegin(). + SetField("Street", TypeString). + SetField("City", TypeString). + RecordTypeEnd()). + RecordTypeEnd() fmt.Printf("RecordType: %v\n", recordType) // create a parquet schema