Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat(bigquery/storage/managedwriter/adapt): support packed field option #6312

Merged
merged 5 commits into from
Jul 14, 2022
Merged

feat(bigquery/storage/managedwriter/adapt): support packed field option #6312

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
312a649
feat(bigquery/storage/managedwriter/adapt): support packed field option
shollyman Jul 7, 2022
183e223
Merge branch 'main' into packed-repeated
shollyman Jul 8, 2022
38eddac
widen type list eligible for packed encoding
shollyman Jul 8, 2022
7d70dbb
cleanup comments (including unrelated)
shollyman Jul 8, 2022
ab927cf
Merge branch 'main' into packed-repeated
shollyman Jul 14, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 42 additions & 8 deletions bigquery/storage/managedwriter/adapt/protoconversion.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,27 @@ var bqTypeToFieldTypeMap = map[storagepb.TableFieldSchema_Type]descriptorpb.Fiel
storagepb.TableFieldSchema_TIMESTAMP: descriptorpb.FieldDescriptorProto_TYPE_INT64,
}

// Primitive types which can leverage packed encoding when repeated/arrays.
//
// Note: many/most of these aren't used when doing schema to proto conversion, but
// are included for completeness.
var packedTypes = []descriptorpb.FieldDescriptorProto_Type{
descriptorpb.FieldDescriptorProto_TYPE_INT32,
descriptorpb.FieldDescriptorProto_TYPE_INT64,
descriptorpb.FieldDescriptorProto_TYPE_UINT32,
descriptorpb.FieldDescriptorProto_TYPE_UINT64,
descriptorpb.FieldDescriptorProto_TYPE_SINT32,
descriptorpb.FieldDescriptorProto_TYPE_SINT64,
descriptorpb.FieldDescriptorProto_TYPE_FIXED32,
descriptorpb.FieldDescriptorProto_TYPE_FIXED64,
descriptorpb.FieldDescriptorProto_TYPE_SFIXED32,
descriptorpb.FieldDescriptorProto_TYPE_SFIXED64,
descriptorpb.FieldDescriptorProto_TYPE_FLOAT,
descriptorpb.FieldDescriptorProto_TYPE_DOUBLE,
descriptorpb.FieldDescriptorProto_TYPE_BOOL,
descriptorpb.FieldDescriptorProto_TYPE_ENUM,
}

// For TableFieldSchema OPTIONAL mode, we use the wrapper types to allow for the
// proper representation of NULL values, as proto3 semantics would just use default value.
var bqTypeToWrapperMap = map[storagepb.TableFieldSchema_Type]string{
Expand All @@ -85,7 +106,7 @@ var wellKnownTypesWrapperName = "google/protobuf/wrappers.proto"

// dependencyCache is used to reduce the number of unique messages we generate by caching based on the tableschema.
//
// keys are based on the base64-encoded serialized tableschema value.
// Keys are based on the base64-encoded serialized tableschema value.
type dependencyCache map[string]protoreflect.MessageDescriptor

func (dm dependencyCache) get(schema *storagepb.TableSchema) protoreflect.MessageDescriptor {
Expand Down Expand Up @@ -143,7 +164,7 @@ func StorageSchemaToProto3Descriptor(inSchema *storagepb.TableSchema, scope stri
return storageSchemaToDescriptorInternal(inSchema, scope, &dc, true)
}

// internal implementation of the conversion code.
// Internal implementation of the conversion code.
func storageSchemaToDescriptorInternal(inSchema *storagepb.TableSchema, scope string, cache *dependencyCache, useProto3 bool) (protoreflect.MessageDescriptor, error) {
if inSchema == nil {
return nil, newConversionError(scope, fmt.Errorf("no input schema was provided"))
Expand All @@ -170,11 +191,11 @@ func storageSchemaToDescriptorInternal(inSchema *storagepb.TableSchema, scope st
break
}
}
// if dep is missing, add to current dependencies
// If dep is missing, add to current dependencies.
if !haveDep {
deps = append(deps, foundDesc.ParentFile())
}
// construct field descriptor for the message
// Construct field descriptor for the message.
fdp, err := tableFieldSchemaToFieldDescriptorProto(f, fNumber, string(foundDesc.FullName()), useProto3)
if err != nil {
return nil, newConversionError(scope, fmt.Errorf("couldn't convert field to FieldDescriptorProto: %w", err))
Expand Down Expand Up @@ -277,12 +298,25 @@ func tableFieldSchemaToFieldDescriptorProto(field *storagepb.TableFieldSchema, i

// For (REQUIRED||REPEATED) fields for proto3, or all cases for proto2, we can use the expected scalar types.
if field.GetMode() != storagepb.TableFieldSchema_NULLABLE || !useProto3 {
return &descriptorpb.FieldDescriptorProto{
outType := bqTypeToFieldTypeMap[field.GetType()]
fdp := &descriptorpb.FieldDescriptorProto{
Name: proto.String(name),
Number: proto.Int32(idx),
Type: bqTypeToFieldTypeMap[field.GetType()].Enum(),
Type: outType.Enum(),
Label: convertModeToLabel(field.GetMode(), useProto3),
}, nil
}
// Special case: proto2 repeated fields may benefit from using packed annotation.
if field.GetMode() == storagepb.TableFieldSchema_REPEATED && !useProto3 {
for _, v := range packedTypes {
if outType == v {
fdp.Options = &descriptorpb.FieldOptions{
Packed: proto.Bool(true),
}
break
}
}
}
return fdp, nil
}
// For NULLABLE proto3 fields, use a wrapper type.
return &descriptorpb.FieldDescriptorProto{
Expand Down Expand Up @@ -445,7 +479,7 @@ func normalizeName(in string) string {
return strings.Replace(in, ".", "_", -1)
}

// these types don't get normalized into the fully-contained structure.
// These types don't get normalized into the fully-contained structure.
var normalizationSkipList = []string{
/*
TODO: when backend supports resolving well known types, this list should be enabled.
Expand Down
29 changes: 29 additions & 0 deletions bigquery/storage/managedwriter/adapt/protoconversion_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -384,6 +384,35 @@ func TestSchemaToProtoConversion(t *testing.T) {
},
},
},
{
description: "repeated w/packed",
bq: &storagepb.TableSchema{
Fields: []*storagepb.TableFieldSchema{
{Name: "name", Type: storagepb.TableFieldSchema_STRING, Mode: storagepb.TableFieldSchema_NULLABLE},
{Name: "some_lengths", Type: storagepb.TableFieldSchema_INT64, Mode: storagepb.TableFieldSchema_REPEATED},
{Name: "nicknames", Type: storagepb.TableFieldSchema_STRING, Mode: storagepb.TableFieldSchema_REPEATED},
}},
wantProto2: &descriptorpb.DescriptorProto{
Name: proto.String("root"),
Field: []*descriptorpb.FieldDescriptorProto{
{
Name: proto.String("name"),
Number: proto.Int32(1),
Type: descriptorpb.FieldDescriptorProto_TYPE_STRING.Enum(),
Label: descriptorpb.FieldDescriptorProto_LABEL_OPTIONAL.Enum()},
{
Name: proto.String("some_lengths"),
Number: proto.Int32(2),
Type: descriptorpb.FieldDescriptorProto_TYPE_INT64.Enum(),
Label: descriptorpb.FieldDescriptorProto_LABEL_REPEATED.Enum(),
Options: &descriptorpb.FieldOptions{
Packed: proto.Bool(true),
},
},
{Name: proto.String("nicknames"), Number: proto.Int32(3), Type: descriptorpb.FieldDescriptorProto_TYPE_STRING.Enum(), Label: descriptorpb.FieldDescriptorProto_LABEL_REPEATED.Enum()},
},
},
},
}
for _, tc := range testCases {
// Proto2
Expand Down