record type begin and end
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user