Protobuf schema to Spark schema conversion

protobuf_spark_schema_convertor.py

"""
MIT License

Copyright (c) [2022] [Rastislav Kisel]

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""

import google.protobuf.descriptor

from pyspark.sql.types import (
  StructType,
  StructField,
  StringType,
  IntegerType,
  LongType,
  DoubleType,
  FloatType,
  ArrayType,
  MapType,
  BooleanType,
)

class ProtobufToSparkSchemaConvertor:
  """
  ProtobufToSparkSchemaConvertor converts Protobuf schema to Spark schema.
  It does that by walking recursively through the protobuf definition and
  creating corresponding Spark type objects.
  """

  # List from: https://googleapis.dev/python/protobuf/latest/google/protobuf/descriptor.html#google.protobuf.descriptor.FieldDescriptor.TYPE_BOOL
  _primitive_types_map = {
    google.protobuf.descriptor.FieldDescriptor.TYPE_BOOL: lambda: BooleanType(),
    google.protobuf.descriptor.FieldDescriptor.TYPE_BYTES: lambda: StringType(),
    google.protobuf.descriptor.FieldDescriptor.TYPE_DOUBLE: lambda: DoubleType(),
    google.protobuf.descriptor.FieldDescriptor.TYPE_ENUM: lambda: LongType(),
    google.protobuf.descriptor.FieldDescriptor.TYPE_FIXED32: lambda: IntegerType(),
    google.protobuf.descriptor.FieldDescriptor.TYPE_FIXED64: lambda: LongType(),
    google.protobuf.descriptor.FieldDescriptor.TYPE_FLOAT: lambda: FloatType(),
    google.protobuf.descriptor.FieldDescriptor.TYPE_INT32: lambda: IntegerType(),
    google.protobuf.descriptor.FieldDescriptor.TYPE_INT64: lambda: LongType(),
    google.protobuf.descriptor.FieldDescriptor.TYPE_SFIXED32: lambda: IntegerType(),
    google.protobuf.descriptor.FieldDescriptor.TYPE_SFIXED64: lambda: LongType(),
    google.protobuf.descriptor.FieldDescriptor.TYPE_SINT32: lambda: IntegerType(),
    google.protobuf.descriptor.FieldDescriptor.TYPE_SINT64: lambda: LongType(),
    google.protobuf.descriptor.FieldDescriptor.TYPE_STRING: lambda: StringType(),
    google.protobuf.descriptor.FieldDescriptor.TYPE_UINT32: lambda: IntegerType(),
    google.protobuf.descriptor.FieldDescriptor.TYPE_UINT64: lambda: LongType(),
  }

  # List from: https://github.com/protocolbuffers/protobuf/blob/main/src/google/protobuf/wrappers.proto
  _wrapper_type_names_map = {
    "DoubleValue": lambda: DoubleType(),
    "FloatValue": lambda: FloatType(),
    "Int64Value": lambda: LongType(),
    "UInt64Value": lambda: LongType(),
    "Int32Value": lambda: IntegerType(),
    "UInt32Value": lambda: IntegerType(),
    "BoolValue": lambda: BooleanType(),
    "StringValue": lambda: StringType(),
    "BytesValue": lambda: StringType(),
  }

  def get_schema(self, descriptor: google.protobuf.descriptor.Descriptor):
    full_schema = []
    self._walk_protobuf_descriptor(descriptor, full_schema)
    return StructType(full_schema)

  def _is_int_type(self, field_descriptor: google.protobuf.descriptor.FieldDescriptor):
    return field_descriptor.type in [
      google.protobuf.descriptor.FieldDescriptor.TYPE_FIXED32,
      google.protobuf.descriptor.FieldDescriptor.TYPE_FIXED64,
      google.protobuf.descriptor.FieldDescriptor.TYPE_INT32,
      google.protobuf.descriptor.FieldDescriptor.TYPE_INT64,
      google.protobuf.descriptor.FieldDescriptor.TYPE_SFIXED32,
      google.protobuf.descriptor.FieldDescriptor.TYPE_SFIXED64,
      google.protobuf.descriptor.FieldDescriptor.TYPE_SINT32,
      google.protobuf.descriptor.FieldDescriptor.TYPE_SINT64,
      google.protobuf.descriptor.FieldDescriptor.TYPE_UINT32,
      google.protobuf.descriptor.FieldDescriptor.TYPE_UINT64,
    ]

  def _is_primitive_type(self, field_descriptor: google.protobuf.descriptor.FieldDescriptor):
    return field_descriptor.type in self._primitive_types_map

  def _is_map_type(self, field_descriptor: google.protobuf.descriptor.FieldDescriptor):
    return (
      field_descriptor.message_type
      and field_descriptor.message_type.has_options
      and field_descriptor.message_type.GetOptions().map_entry
    )

  def _is_proto_wrapper_type(self, field_descriptor: google.protobuf.descriptor.FieldDescriptor):
    return field_descriptor.message_type.name in self._wrapper_type_names_map

  def _handle_map(self, field: google.protobuf.descriptor.FieldDescriptor):
    # Maps hold key/value descriptors in "key" and "value" fields.
    key = [inner_field for inner_field in field.message_type.fields if inner_field.name == "key"][0]
    if self._is_int_type(key):
      # Protobuf serializes integers as strings in maps:
      # https://github.com/protocolbuffers/protobuf/issues/7769
      key_type = StringType()
    else:
      key_type = ProtobufToSparkSchemaConvertor._primitive_types_map[key.type]()

    value = [
      inner_field for inner_field in field.message_type.fields if inner_field.name == "value"
    ][0]
    # Handle primitive value type
    if self._is_primitive_type(value):
      value_type = self._primitive_types_map[value.type]()
    # Handle map value type - we handle it here because map doesn't have proper fields.
    elif self._is_map_type(value):
      value_type = self._handle_map(value)
    # Handle message value type
    else:
      schema_list = []
      self._walk_protobuf_descriptor(value.message_type, schema_list)
      value_type = StructType(schema_list)

    return MapType(key_type, value_type, True)

  def _walk_protobuf_descriptor(self, descriptor: google.protobuf.descriptor.Descriptor, schema):
    for field in descriptor.fields:
      if field.type == google.protobuf.descriptor.FieldDescriptor.TYPE_MESSAGE:
        # Handle timestamp: https://github.com/protocolbuffers/protobuf/blob/main/src/google/protobuf/timestamp.proto
        if field.message_type.name == "Timestamp":
          spark_field = StructField(field.name, StringType(), True)
        # Handle wrapper types: https://github.com/protocolbuffers/protobuf/blob/main/src/google/protobuf/wrappers.proto
        elif self._is_proto_wrapper_type(field):
          type_factory = self._wrapper_type_names_map[field.message_type.name]
          spark_field = StructField(field.name, type_factory(), True)
        # Handle map type
        elif self._is_map_type(field):
          map_type = self._handle_map(field)
          spark_field = StructField(field.name, map_type, True)
        else:
          # Handle message type
          schema_list = []
          self._walk_protobuf_descriptor(field.message_type, schema_list)
          spark_field = StructField(field.name, StructType(schema_list), True)
      else:
        # Handle primitive type
        type_factory = self._primitive_types_map[field.type]
        if not type_factory:
          raise ValueError(f"Missing primitive type for: {field.name} | {field.type}")

        spark_field = StructField(field.name, type_factory(), True)

      # Handle array type
      if (
        field.label == google.protobuf.descriptor.FieldDescriptor.LABEL_REPEATED
        and not self._is_map_type(field)
      ):
        spark_field = StructField(field.name, ArrayType(spark_field.dataType, True), True)

      schema.append(spark_field)