Mirror
/
seaweedfs
mirror of https://github.com/seaweedfs/seaweedfs.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity
Browse Source

works with single level of list

mq-subscribe
chrislu 8 months ago
parent
d7e5f6b2a5
commit
d88c1872ac
4 changed files with 71 additions and 60 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 7
      weed/mq/schema/to_parquet_schema.go
  2. 27
      weed/mq/schema/to_parquet_value.go
  3. 66
      weed/mq/schema/to_schema_value.go
  4. 31
      weed/mq/schema/write_parquet_test.go

7
weed/mq/schema/to_parquet_schema.go
View File

@ -21,6 +21,7 @@ func toParquetFieldType(fieldType *schema_pb.Type) (dataType parquet.Node, err e
switch fieldType.Kind.(type) {
case *schema_pb.Type_ScalarType:
dataType, err = toParquetFieldTypeScalar(fieldType.GetScalarType())
dataType = parquet.Optional(dataType)
case *schema_pb.Type_RecordType:
dataType, err = toParquetFieldTypeRecord(fieldType.GetRecordType())
case *schema_pb.Type_ListType:
@ -29,6 +30,7 @@ func toParquetFieldType(fieldType *schema_pb.Type) (dataType parquet.Node, err e
return nil, fmt.Errorf("unknown field type: %T", fieldType.Kind)
}
return dataType, err
}
@ -37,7 +39,7 @@ func toParquetFieldTypeList(listType *schema_pb.ListType) (parquet.Node, error)
if err != nil {
return nil, err
}
return parquet.List(elementType), nil
return parquet.Repeated(elementType), nil
}
func toParquetFieldTypeScalar(scalarType schema_pb.ScalarType) (parquet.Node, error) {
@ -67,9 +69,6 @@ func toParquetFieldTypeRecord(recordType *schema_pb.RecordType) (parquet.Node, e
if err != nil {
return nil, err
}
if !field.IsRequired {
parquetFieldType = parquet.Optional(parquetFieldType)
}
recordNode[field.Name] = parquetFieldType
}
return recordNode, nil

27
weed/mq/schema/to_parquet_value.go
View File

@ -6,28 +6,33 @@ import (
"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
)
func rowBuilderVisit(rowBuilder *parquet.RowBuilder, fieldType *schema_pb.Type, fieldValue *schema_pb.Value, columnIndex int) error {
func rowBuilderVisit(rowBuilder *parquet.RowBuilder, fieldType *schema_pb.Type, fieldValue *schema_pb.Value, columnIndex int) (endIndex int, err error) {
switch fieldType.Kind.(type) {
case *schema_pb.Type_ScalarType:
parquetValue, err := toParquetValue(fieldValue)
endIndex = columnIndex+1
var parquetValue parquet.Value
parquetValue, err = toParquetValue(fieldValue)
if err != nil {
return err
return
}
rowBuilder.Add(columnIndex, parquetValue)
// fmt.Printf("rowBuilder.Add %d %v\n", columnIndex, parquetValue)
case *schema_pb.Type_ListType:
rowBuilder.Next(columnIndex)
// fmt.Printf("rowBuilder.Next %d\n", columnIndex)
elementType := fieldType.GetListType().ElementType
for _, value := range fieldValue.GetListValue().Values {
if err := rowBuilderVisit(rowBuilder, elementType, value, columnIndex); err != nil {
return err
if endIndex, err = rowBuilderVisit(rowBuilder, elementType, value, columnIndex); err != nil {
return
}
}
rowBuilder.Next(columnIndex)
}
return nil
return
}
func AddRecordValue(rowBuilder *parquet.RowBuilder, recordType *schema_pb.RecordType, recordValue *schema_pb.RecordValue) error {
visitor := func(fieldType *schema_pb.Type, fieldValue *schema_pb.Value, index int) error {
visitor := func(fieldType *schema_pb.Type, fieldValue *schema_pb.Value, index int) (endIndex int, err error) {
return rowBuilderVisit(rowBuilder, fieldType, fieldValue, index)
}
fieldType := &schema_pb.Type{Kind: &schema_pb.Type_RecordType{RecordType: recordType}}
@ -38,7 +43,7 @@ func AddRecordValue(rowBuilder *parquet.RowBuilder, recordType *schema_pb.Record
// typeValueVisitor is a function that is called for each value in a schema_pb.Value
// Find the column index.
// intended to be used in RowBuilder.Add(columnIndex, value)
type typeValueVisitor func(fieldType *schema_pb.Type, fieldValue *schema_pb.Value, index int) error
type typeValueVisitor func(fieldType *schema_pb.Type, fieldValue *schema_pb.Value, index int) (endIndex int, err error)
func visitValue(fieldType *schema_pb.Type, fieldValue *schema_pb.Value, visitor typeValueVisitor) (err error) {
_, err = doVisitValue(fieldType, fieldValue, 0, visitor)
@ -50,9 +55,9 @@ func visitValue(fieldType *schema_pb.Type, fieldValue *schema_pb.Value, visitor
func doVisitValue(fieldType *schema_pb.Type, fieldValue *schema_pb.Value, columnIndex int, visitor typeValueVisitor) (endIndex int, err error) {
switch fieldType.Kind.(type) {
case *schema_pb.Type_ScalarType:
return columnIndex+1, visitor(fieldType, fieldValue, columnIndex)
return visitor(fieldType, fieldValue, columnIndex)
case *schema_pb.Type_ListType:
return columnIndex+1, visitor(fieldType, fieldValue, columnIndex)
return visitor(fieldType, fieldValue, columnIndex)
case *schema_pb.Type_RecordType:
for _, field := range fieldType.GetRecordType().Fields {
fieldValue, found := fieldValue.GetRecordValue().Fields[field.Name]

66
weed/mq/schema/to_schema_value.go
View File

@ -8,78 +8,76 @@ import (
func ToRecordValue(recordType *schema_pb.RecordType, row parquet.Row) (*schema_pb.RecordValue, error) {
values := []parquet.Value(row)
recordValue, _, err := toRecordValue(recordType, values, 0)
recordValue, _, _, err := toRecordValue(recordType, values, 0, 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) {
func ToValue(t *schema_pb.Type, values []parquet.Value, valueIndex, columnIndex int) (value *schema_pb.Value, endValueIndex, endColumnIndex int, err error) {
switch t.Kind.(type) {
case *schema_pb.Type_ScalarType:
return toScalarValue(t.GetScalarType(), values, columnIndex)
value, err = toScalarValue(t.GetScalarType(), values, valueIndex, columnIndex)
return value, valueIndex + 1, columnIndex + 1, err
case *schema_pb.Type_ListType:
return toListValue(t.GetListType(), values, columnIndex)
return toListValue(t.GetListType(), values, valueIndex, columnIndex)
case *schema_pb.Type_RecordType:
return toRecordValue(t.GetRecordType(), values, columnIndex)
return toRecordValue(t.GetRecordType(), values, valueIndex, columnIndex)
}
return nil, 0, fmt.Errorf("unsupported type: %v", t)
return nil, 0, 0, fmt.Errorf("unsupported type: %v", t)
}
func toRecordValue(recordType *schema_pb.RecordType, values []parquet.Value, columnIndex int) (*schema_pb.Value, int, error) {
func toRecordValue(recordType *schema_pb.RecordType, values []parquet.Value, valueIndex, columnIndex int) (*schema_pb.Value, int, 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)
fieldValue, endValueIndex, endColumnIndex, err := ToValue(field.Type, values, valueIndex, columnIndex)
if err != nil {
return nil, 0, err
return nil, 0, 0, err
}
if endIndex == columnIndex {
continue
}
columnIndex = endIndex
columnIndex = endColumnIndex
valueIndex = endValueIndex
recordValue.Fields[field.Name] = fieldValue
}
return &schema_pb.Value{Kind: &schema_pb.Value_RecordValue{RecordValue: &recordValue}}, columnIndex, nil
return &schema_pb.Value{Kind: &schema_pb.Value_RecordValue{RecordValue: &recordValue}}, valueIndex, columnIndex, nil
}
func toListValue(listType *schema_pb.ListType, values []parquet.Value, index int) (listValue *schema_pb.Value, endIndex int, err error) {
func toListValue(listType *schema_pb.ListType, values []parquet.Value, valueIndex, columnIndex int) (listValue *schema_pb.Value, endValueIndex, endColumnIndex 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 {
for ;valueIndex < len(values); {
if values[valueIndex].Column() != columnIndex {
break
}
value, valueIndex, endColumnIndex, err = ToValue(listType.ElementType, values, valueIndex, columnIndex)
if err != nil {
return nil, 0,0, err
}
listValues = append(listValues, value)
i = endIndex
}
return &schema_pb.Value{Kind: &schema_pb.Value_ListValue{ListValue: &schema_pb.ListValue{Values: listValues}}}, endIndex, nil
return &schema_pb.Value{Kind: &schema_pb.Value_ListValue{ListValue: &schema_pb.ListValue{Values: listValues}}}, valueIndex, endColumnIndex, nil
}
func toScalarValue(scalarType schema_pb.ScalarType, values []parquet.Value, columnIndex int) (*schema_pb.Value, int, error) {
value := values[columnIndex]
func toScalarValue(scalarType schema_pb.ScalarType, values []parquet.Value, valueIndex, columnIndex int) (*schema_pb.Value, error) {
value := values[valueIndex]
if value.Column() != columnIndex {
return nil, columnIndex, nil
return nil, nil
}
switch scalarType {
case schema_pb.ScalarType_BOOLEAN:
return &schema_pb.Value{Kind: &schema_pb.Value_BoolValue{BoolValue: value.Boolean()}}, columnIndex + 1, nil
return &schema_pb.Value{Kind: &schema_pb.Value_BoolValue{BoolValue: value.Boolean()}}, nil
case schema_pb.ScalarType_INTEGER:
return &schema_pb.Value{Kind: &schema_pb.Value_Int32Value{Int32Value: value.Int32()}}, columnIndex + 1, nil
return &schema_pb.Value{Kind: &schema_pb.Value_Int32Value{Int32Value: value.Int32()}}, nil
case schema_pb.ScalarType_LONG:
return &schema_pb.Value{Kind: &schema_pb.Value_Int64Value{Int64Value: value.Int64()}}, columnIndex + 1, nil
return &schema_pb.Value{Kind: &schema_pb.Value_Int64Value{Int64Value: value.Int64()}}, nil
case schema_pb.ScalarType_FLOAT:
return &schema_pb.Value{Kind: &schema_pb.Value_FloatValue{FloatValue: value.Float()}}, columnIndex + 1, nil
return &schema_pb.Value{Kind: &schema_pb.Value_FloatValue{FloatValue: value.Float()}}, nil
case schema_pb.ScalarType_DOUBLE:
return &schema_pb.Value{Kind: &schema_pb.Value_DoubleValue{DoubleValue: value.Double()}}, columnIndex + 1, nil
return &schema_pb.Value{Kind: &schema_pb.Value_DoubleValue{DoubleValue: value.Double()}}, nil
case schema_pb.ScalarType_BYTES:
return &schema_pb.Value{Kind: &schema_pb.Value_BytesValue{BytesValue: value.ByteArray()}}, columnIndex + 1, nil
return &schema_pb.Value{Kind: &schema_pb.Value_BytesValue{BytesValue: value.ByteArray()}}, nil
case schema_pb.ScalarType_STRING:
return &schema_pb.Value{Kind: &schema_pb.Value_StringValue{StringValue: string(value.ByteArray())}}, columnIndex + 1, nil
return &schema_pb.Value{Kind: &schema_pb.Value_StringValue{StringValue: string(value.ByteArray())}}, nil
}
return nil, columnIndex, fmt.Errorf("unsupported scalar type: %v", scalarType)
return nil, fmt.Errorf("unsupported scalar type: %v", scalarType)
}

31
weed/mq/schema/write_parquet_test.go
View File

@ -27,7 +27,8 @@ func TestWriteParquet(t *testing.T) {
t.Fatalf("ToParquetSchema failed: %v", err)
}
fmt.Printf("ParquetSchema: %v\n", parquetSchema)
parquet.PrintSchema(os.Stdout, "example", parquetSchema)
fmt.Printf("Go Type: %+v\n", parquetSchema.GoType())
filename := "example.parquet"
@ -50,21 +51,29 @@ func testWritingParquetFile(t *testing.T, filename string, parquetSchema *parque
defer file.Close()
writer := parquet.NewWriter(file, parquetSchema, parquet.Compression(&zstd.Codec{Level: zstd.SpeedDefault}))
rowBuilder := parquet.NewRowBuilder(parquetSchema)
for i := 0; i < 128; i++ {
for i := 0; i < 128*1024; i++ {
rowBuilder.Reset()
// generate random data
AddRecordValue(rowBuilder, recordType, NewRecordValueBuilder().
AddLongValue("ID", int64(1+i)).
AddLongValue("CreatedAt", 2*int64(i)).
recordValue := NewRecordValueBuilder().
AddLongValue("ID", 1+int64(i)).
AddLongValue("CreatedAt", 2+2*int64(i)).
AddRecordValue("Person", NewRecordValueBuilder().
AddStringValue("zName", fmt.Sprintf("john_%d", i)).
AddStringListValue("emails",
fmt.Sprintf("john_%d@y.com", i),
fmt.Sprintf("john_%d@g.com", i),
fmt.Sprintf("john_%d@t.com", i))).
AddStringValue("Company", fmt.Sprintf("company_%d", i)).Build())
fmt.Sprintf("john_%d@a.com", i),
fmt.Sprintf("john_%d@b.com", i),
fmt.Sprintf("john_%d@c.com", i),
fmt.Sprintf("john_%d@d.com", i),
fmt.Sprintf("john_%d@e.com", i))).
AddStringValue("Company", fmt.Sprintf("company_%d", i)).Build()
AddRecordValue(rowBuilder, recordType, recordValue)
// fmt.Printf("RecordValue: %v\n", recordValue)
row := rowBuilder.Row()
// fmt.Printf("Row: %+v\n", row)
if err != nil {
t.Fatalf("rowBuilder.Build failed: %v", err)
}
@ -98,11 +107,11 @@ func testReadingParquetFile(t *testing.T, filename string, parquetSchema *parque
for i := 0; i < rowCount; i++ {
row := rows[i]
// convert parquet row to schema_pb.RecordValue
recordValue, err := ToRecordValue(recordType, row)
_, err := ToRecordValue(recordType, row)
if err != nil {
t.Fatalf("ToRecordValue failed: %v", err)
}
fmt.Printf("RecordValue: %v\n", recordValue)
// fmt.Printf("RecordValue: %v\n", recordValue)
}
total += rowCount
}

Write Preview
Loading…
Cancel
Save