chrislu
8 months ago
6 changed files with 317 additions and 19 deletions
-
5weed/mq/schema/schema_builder.go
-
4weed/mq/schema/to_parquet_schema.go
-
29weed/mq/schema/to_parquet_value.go
-
85weed/mq/schema/to_schema_value.go
-
113weed/mq/schema/value_builder.go
-
90weed/mq/schema/write_parquet_test.go
@ -0,0 +1,85 @@ |
|||||
|
package schema |
||||
|
|
||||
|
import ( |
||||
|
"fmt" |
||||
|
"github.com/parquet-go/parquet-go" |
||||
|
"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb" |
||||
|
) |
||||
|
|
||||
|
func ToRecordValue(recordType *schema_pb.RecordType, row parquet.Row) (*schema_pb.RecordValue, error) { |
||||
|
values := []parquet.Value(row) |
||||
|
recordValue, _, err := toRecordValue(recordType, values, 0) |
||||
|
if err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
return recordValue.GetRecordValue(), nil |
||||
|
} |
||||
|
|
||||
|
func ToValue(t *schema_pb.Type, values []parquet.Value, columnIndex int) (value *schema_pb.Value, endIndex int, err error) { |
||||
|
switch t.Kind.(type) { |
||||
|
case *schema_pb.Type_ScalarType: |
||||
|
return toScalarValue(t.GetScalarType(), values, columnIndex) |
||||
|
case *schema_pb.Type_ListType: |
||||
|
return toListValue(t.GetListType(), values, columnIndex) |
||||
|
case *schema_pb.Type_RecordType: |
||||
|
return toRecordValue(t.GetRecordType(), values, columnIndex) |
||||
|
} |
||||
|
return nil, 0, fmt.Errorf("unsupported type: %v", t) |
||||
|
} |
||||
|
|
||||
|
func toRecordValue(recordType *schema_pb.RecordType, values []parquet.Value, columnIndex int) (*schema_pb.Value, int, error) { |
||||
|
recordValue := schema_pb.RecordValue{Fields: make(map[string]*schema_pb.Value)} |
||||
|
for _, field := range recordType.Fields { |
||||
|
fieldValue, endIndex, err := ToValue(field.Type, values, columnIndex) |
||||
|
if err != nil { |
||||
|
return nil, 0, err |
||||
|
} |
||||
|
if endIndex == columnIndex { |
||||
|
continue |
||||
|
} |
||||
|
columnIndex = endIndex |
||||
|
recordValue.Fields[field.Name] = fieldValue |
||||
|
} |
||||
|
return &schema_pb.Value{Kind: &schema_pb.Value_RecordValue{RecordValue: &recordValue}}, columnIndex, nil |
||||
|
} |
||||
|
|
||||
|
func toListValue(listType *schema_pb.ListType, values []parquet.Value, index int) (listValue *schema_pb.Value, endIndex int, err error) { |
||||
|
listValues := make([]*schema_pb.Value, 0) |
||||
|
var value *schema_pb.Value |
||||
|
for i := index; i < len(values); { |
||||
|
value, endIndex, err = ToValue(listType.ElementType, values, i) |
||||
|
if err != nil { |
||||
|
return nil, 0, err |
||||
|
} |
||||
|
if endIndex == i { |
||||
|
break |
||||
|
} |
||||
|
listValues = append(listValues, value) |
||||
|
i = endIndex |
||||
|
} |
||||
|
return &schema_pb.Value{Kind: &schema_pb.Value_ListValue{ListValue: &schema_pb.ListValue{Values: listValues}}}, endIndex, nil |
||||
|
} |
||||
|
|
||||
|
func toScalarValue(scalarType schema_pb.ScalarType, values []parquet.Value, columnIndex int) (*schema_pb.Value, int, error) { |
||||
|
value := values[columnIndex] |
||||
|
if value.Column() != columnIndex { |
||||
|
return nil, columnIndex, nil |
||||
|
} |
||||
|
switch scalarType { |
||||
|
case schema_pb.ScalarType_BOOLEAN: |
||||
|
return &schema_pb.Value{Kind: &schema_pb.Value_BoolValue{BoolValue: value.Boolean()}}, columnIndex + 1, nil |
||||
|
case schema_pb.ScalarType_INTEGER: |
||||
|
return &schema_pb.Value{Kind: &schema_pb.Value_Int32Value{Int32Value: value.Int32()}}, columnIndex + 1, nil |
||||
|
case schema_pb.ScalarType_LONG: |
||||
|
return &schema_pb.Value{Kind: &schema_pb.Value_Int64Value{Int64Value: value.Int64()}}, columnIndex + 1, nil |
||||
|
case schema_pb.ScalarType_FLOAT: |
||||
|
return &schema_pb.Value{Kind: &schema_pb.Value_FloatValue{FloatValue: value.Float()}}, columnIndex + 1, nil |
||||
|
case schema_pb.ScalarType_DOUBLE: |
||||
|
return &schema_pb.Value{Kind: &schema_pb.Value_DoubleValue{DoubleValue: value.Double()}}, columnIndex + 1, nil |
||||
|
case schema_pb.ScalarType_BYTES: |
||||
|
return &schema_pb.Value{Kind: &schema_pb.Value_BytesValue{BytesValue: value.ByteArray()}}, columnIndex + 1, nil |
||||
|
case schema_pb.ScalarType_STRING: |
||||
|
return &schema_pb.Value{Kind: &schema_pb.Value_StringValue{StringValue: string(value.ByteArray())}}, columnIndex + 1, nil |
||||
|
} |
||||
|
return nil, columnIndex, fmt.Errorf("unsupported scalar type: %v", scalarType) |
||||
|
} |
@ -0,0 +1,113 @@ |
|||||
|
package schema |
||||
|
|
||||
|
import "github.com/seaweedfs/seaweedfs/weed/pb/schema_pb" |
||||
|
|
||||
|
// RecordValueBuilder helps in constructing RecordValue protobuf messages
|
||||
|
type RecordValueBuilder struct { |
||||
|
recordValue *schema_pb.RecordValue |
||||
|
} |
||||
|
|
||||
|
// NewRecordValueBuilder creates a new RecordValueBuilder instance
|
||||
|
func NewRecordValueBuilder() *RecordValueBuilder { |
||||
|
return &RecordValueBuilder{recordValue: &schema_pb.RecordValue{Fields: make(map[string]*schema_pb.Value)}} |
||||
|
} |
||||
|
|
||||
|
// Build returns the constructed RecordValue message
|
||||
|
func (rvb *RecordValueBuilder) Build() *schema_pb.RecordValue { |
||||
|
return rvb.recordValue |
||||
|
} |
||||
|
|
||||
|
func (rvb *RecordValueBuilder) AddBoolValue(key string, value bool) *RecordValueBuilder { |
||||
|
rvb.recordValue.Fields[key] = &schema_pb.Value{Kind: &schema_pb.Value_BoolValue{BoolValue: value}} |
||||
|
return rvb |
||||
|
} |
||||
|
func (rvb *RecordValueBuilder) AddIntValue(key string, value int32) *RecordValueBuilder { |
||||
|
rvb.recordValue.Fields[key] = &schema_pb.Value{Kind: &schema_pb.Value_Int32Value{Int32Value: value}} |
||||
|
return rvb |
||||
|
} |
||||
|
func (rvb *RecordValueBuilder) AddLongValue(key string, value int64) *RecordValueBuilder { |
||||
|
rvb.recordValue.Fields[key] = &schema_pb.Value{Kind: &schema_pb.Value_Int64Value{Int64Value: value}} |
||||
|
return rvb |
||||
|
} |
||||
|
func (rvb *RecordValueBuilder) AddFloatValue(key string, value float32) *RecordValueBuilder { |
||||
|
rvb.recordValue.Fields[key] = &schema_pb.Value{Kind: &schema_pb.Value_FloatValue{FloatValue: value}} |
||||
|
return rvb |
||||
|
} |
||||
|
func (rvb *RecordValueBuilder) AddDoubleValue(key string, value float64) *RecordValueBuilder { |
||||
|
rvb.recordValue.Fields[key] = &schema_pb.Value{Kind: &schema_pb.Value_DoubleValue{DoubleValue: value}} |
||||
|
return rvb |
||||
|
} |
||||
|
func (rvb *RecordValueBuilder) AddBytesValue(key string, value []byte) *RecordValueBuilder { |
||||
|
rvb.recordValue.Fields[key] = &schema_pb.Value{Kind: &schema_pb.Value_BytesValue{BytesValue: value}} |
||||
|
return rvb |
||||
|
} |
||||
|
func (rvb *RecordValueBuilder) AddStringValue(key string, value string) *RecordValueBuilder { |
||||
|
rvb.recordValue.Fields[key] = &schema_pb.Value{Kind: &schema_pb.Value_StringValue{StringValue: value}} |
||||
|
return rvb |
||||
|
} |
||||
|
func (rvb *RecordValueBuilder) AddRecordValue(key string, value *RecordValueBuilder) *RecordValueBuilder { |
||||
|
rvb.recordValue.Fields[key] = &schema_pb.Value{Kind: &schema_pb.Value_RecordValue{RecordValue: value.Build()}} |
||||
|
return rvb |
||||
|
} |
||||
|
|
||||
|
func (rvb *RecordValueBuilder) addListValue(key string, values []*schema_pb.Value) *RecordValueBuilder { |
||||
|
rvb.recordValue.Fields[key] = &schema_pb.Value{Kind: &schema_pb.Value_ListValue{ListValue: &schema_pb.ListValue{Values: values}}} |
||||
|
return rvb |
||||
|
} |
||||
|
|
||||
|
func (rvb *RecordValueBuilder) AddBoolListValue(key string, values ...bool) *RecordValueBuilder { |
||||
|
var listValues []*schema_pb.Value |
||||
|
for _, v := range values { |
||||
|
listValues = append(listValues, &schema_pb.Value{Kind: &schema_pb.Value_BoolValue{BoolValue: v}}) |
||||
|
} |
||||
|
return rvb.addListValue(key, listValues) |
||||
|
} |
||||
|
func (rvb *RecordValueBuilder) AddIntListValue(key string, values ...int32) *RecordValueBuilder { |
||||
|
var listValues []*schema_pb.Value |
||||
|
for _, v := range values { |
||||
|
listValues = append(listValues, &schema_pb.Value{Kind: &schema_pb.Value_Int32Value{Int32Value: v}}) |
||||
|
} |
||||
|
return rvb.addListValue(key, listValues) |
||||
|
} |
||||
|
func (rvb *RecordValueBuilder) AddLongListValue(key string, values ...int64) *RecordValueBuilder { |
||||
|
var listValues []*schema_pb.Value |
||||
|
for _, v := range values { |
||||
|
listValues = append(listValues, &schema_pb.Value{Kind: &schema_pb.Value_Int64Value{Int64Value: v}}) |
||||
|
} |
||||
|
return rvb.addListValue(key, listValues) |
||||
|
} |
||||
|
func (rvb *RecordValueBuilder) AddFloatListValue(key string, values ...float32) *RecordValueBuilder { |
||||
|
var listValues []*schema_pb.Value |
||||
|
for _, v := range values { |
||||
|
listValues = append(listValues, &schema_pb.Value{Kind: &schema_pb.Value_FloatValue{FloatValue: v}}) |
||||
|
} |
||||
|
return rvb.addListValue(key, listValues) |
||||
|
} |
||||
|
func (rvb *RecordValueBuilder) AddDoubleListValue(key string, values ...float64) *RecordValueBuilder { |
||||
|
var listValues []*schema_pb.Value |
||||
|
for _, v := range values { |
||||
|
listValues = append(listValues, &schema_pb.Value{Kind: &schema_pb.Value_DoubleValue{DoubleValue: v}}) |
||||
|
} |
||||
|
return rvb.addListValue(key, listValues) |
||||
|
} |
||||
|
func (rvb *RecordValueBuilder) AddBytesListValue(key string, values ...[]byte) *RecordValueBuilder { |
||||
|
var listValues []*schema_pb.Value |
||||
|
for _, v := range values { |
||||
|
listValues = append(listValues, &schema_pb.Value{Kind: &schema_pb.Value_BytesValue{BytesValue: v}}) |
||||
|
} |
||||
|
return rvb.addListValue(key, listValues) |
||||
|
} |
||||
|
func (rvb *RecordValueBuilder) AddStringListValue(key string, values ...string) *RecordValueBuilder { |
||||
|
var listValues []*schema_pb.Value |
||||
|
for _, v := range values { |
||||
|
listValues = append(listValues, &schema_pb.Value{Kind: &schema_pb.Value_StringValue{StringValue: v}}) |
||||
|
} |
||||
|
return rvb.addListValue(key, listValues) |
||||
|
} |
||||
|
func (rvb *RecordValueBuilder) AddRecordListValue(key string, values ...*RecordValueBuilder) *RecordValueBuilder { |
||||
|
var listValues []*schema_pb.Value |
||||
|
for _, v := range values { |
||||
|
listValues = append(listValues, &schema_pb.Value{Kind: &schema_pb.Value_RecordValue{RecordValue: v.Build()}}) |
||||
|
} |
||||
|
return rvb.addListValue(key, listValues) |
||||
|
} |
Write
Preview
Loading…
Cancel
Save
Reference in new issue