[입 개발] Spark에서 Parquet 파일 Custom Schema 로 읽어들이기

최근(?)에 다음과 같은 에러를 많이 보았습니다.

scala> val df = spark.read.parquest("s3://bucket-name/path/")
org.apache.spark.sql.AnalysisException: Parquet type not supported: INT32 (UINT_8);
    at org.apache.spark.sql.execution.datasources.parquet.ParquetToSparkSchemaConverter.typeNotSupported$1(ParquetSchemaConverter.scala:101)
    at org.apache.spark.sql.execution.datasources.parquet.ParquetToSparkSchemaConverter.convertPrimitiveField(ParquetSchemaConverter.scala:137)
    at org.apache.spark.sql.execution.datasources.parquet.ParquetToSparkSchemaConverter.convertField(ParquetSchemaConverter.scala:89)
    at org.apache.spark.sql.execution.datasources.parquet.ParquetToSparkSchemaConverter$$anonfun$1.apply(ParquetSchemaConverter.scala:68)
    at org.apache.spark.sql.execution.datasources.parquet.ParquetToSparkSchemaConverter$$anonfun$1.apply(ParquetSchemaConverter.scala:65)
    at scala.collection.TraversableLike$$anonfun$map$1.apply(TraversableLike.scala:234)
    at 

에러를 보면 Parquet type not supported: INT32 (UINT_8) 로 해당 타입을 지원하지 않는 다는 뜻입니다.

결론 부터 말하자면, Spark에서는 기본적으로 Unsigned Integer를 지원하지 않습니다. 그래서 만약 다음과 같은 Schema 의 Parquet 파일이 있다면 해당 파일을 읽는데 위와 같은 에러를 내게 됩니다. (재밌는건 Athena(아마도 Presto) 에서는 Unsigned Integer도 잘 읽어드립니다.)

결론부터 말하자면, 이런 경우 spark.read.parquet 만 하신다면 value3를 읽는데, 실패하지만, value3 가 필요없다면 Custom Scheme 를 주면 읽을 수 있습니다. 물론, 처음부터 Unsigned Type을 안쓰면 되지 않느냐!!! 라는 멋진 의견을 주실 수 있는데, DB를 덤프해서 parquet으로 저장하다보면, 자신도 모르게 Unsigned Type이 함께 딸려들어갈 수 있습니다. 참고로 AWS DMS(Data Migration Service)는 DB의 내용을 자동으로 덤프해서 parquet으로 만들어주는 멋진 기능이 있는데, white list 기반이 아닌 black list 기반입니다. 즉, 등록한 컬럼들만 덤프하는 것이 아니라, 명시적으로 제외한 컬럼들만 제외되서 덤프되기 때문에, 자신이 모르게 얼마든지 새로운 컬럼이 추가될 수 있습니다. 먼저 정답부터 드린다면…

위와 같은 경우에 schema 메서드를 이용해서 아주 쉽게 읽을 수 있습니다. (다만 해당 컬럼은 포기하셔야…) schema는 두 가지 형태로 제공이 가능합니다. 두 가지 방법 다 잘 동작합니다. 자세한 것은 다음 페이지를 읽어봅시다.(https://spark.apache.org/docs/2.4.0/api/scala/index.html#org.apache.spark.sql.DataFrameReader)

spark.read.schema("value1 long, value2 STRING").parquet(path)

val customSchema = StructType(Array(
    StructField("value1", LongType, true),
    StructField("value2", StringType, true)
    ))

spark.read.schema(customSchema).parquet(path)

자 그럼 왜 이렇게 동작하는지를 알아보시죠. 그런데 사실 Parquet 은 원래부터, 특정 컬럼들만 읽어들이는 기능을 제공하고 있습니다.(http://engineering.vcnc.co.kr/2018/05/parquet-and-spark/)

여기를 보시면 https://parquet.apache.org/documentation/latest/ Parquet의 데이터 구조가 특정 컬럼들만 읽을 수 있는 Columnar 형식이라는 것을 알 수 있습니다.

4-byte magic number "PAR1" 
<Column 1 Chunk 1 + Column Metadata> 
<Column 2 Chunk 1 + Column Metadata> 
... 
<Column N Chunk 1 + Column Metadata> 
<Column 1 Chunk 2 + Column Metadata> 
<Column 2 Chunk 2 + Column Metadata> 
... 
<Column N Chunk 2 + Column Metadata>
... 
<Column 1 Chunk M + Column Metadata> 
<Column 2 Chunk M + Column Metadata> 
... 
<Column N Chunk M + Column Metadata> 
File Metadata 4-byte length in bytes of file metadata 
4-byte magic number "PAR1"

다시 처음으로 돌아가서 아까의 에러메시지를 살펴보시면, ParquetSchemaConverter.scala 라는 파일명이 나옵니다. 그럼 이제 해당 파일을 찾아봅시다. 해당 파일을 보면 대략 convertPrimitiveField 함수에서 에러가 난 것을 쉽게 찾을 수 있습니다.

  private def convertPrimitiveField(field: PrimitiveType): DataType = {
    val typeName = field.getPrimitiveTypeName
    val originalType = field.getOriginalType

    def typeString =
      if (originalType == null) s"$typeName" else s"$typeName ($originalType)"

    def typeNotSupported() =
      throw new AnalysisException(s"Parquet type not supported: $typeString")

    def typeNotImplemented() =
      throw new AnalysisException(s"Parquet type not yet supported: $typeString")

    def illegalType() =
      throw new AnalysisException(s"Illegal Parquet type: $typeString")

    // When maxPrecision = -1, we skip precision range check, and always respect the precision
    // specified in field.getDecimalMetadata.  This is useful when interpreting decimal types stored
    // as binaries with variable lengths.
    def makeDecimalType(maxPrecision: Int = -1): DecimalType = {
      val precision = field.getDecimalMetadata.getPrecision
      val scale = field.getDecimalMetadata.getScale

      ParquetSchemaConverter.checkConversionRequirement(
        maxPrecision == -1 || 1 <= precision && precision <= maxPrecision,
        s"Invalid decimal precision: $typeName cannot store $precision digits (max $maxPrecision)")

      DecimalType(precision, scale)
    }

    typeName match {
      case BOOLEAN => BooleanType

      case FLOAT => FloatType

      case DOUBLE => DoubleType

      case INT32 =>
        originalType match {
          case INT_8 => ByteType
          case INT_16 => ShortType
          case INT_32 | null => IntegerType
          case DATE => DateType
          case DECIMAL => makeDecimalType(Decimal.MAX_INT_DIGITS)
          case UINT_8 => typeNotSupported()
          case UINT_16 => typeNotSupported()
          case UINT_32 => typeNotSupported()
          case TIME_MILLIS => typeNotImplemented()
          case _ => illegalType()
        }

      case INT64 =>
        originalType match {
          case INT_64 | null => LongType
          case DECIMAL => makeDecimalType(Decimal.MAX_LONG_DIGITS)
          case UINT_64 => typeNotSupported()
          case TIMESTAMP_MICROS => TimestampType
          case TIMESTAMP_MILLIS => TimestampType
          case _ => illegalType()
        }

      case INT96 =>
        ParquetSchemaConverter.checkConversionRequirement(
          assumeInt96IsTimestamp,
          "INT96 is not supported unless it's interpreted as timestamp. " +
            s"Please try to set ${SQLConf.PARQUET_INT96_AS_TIMESTAMP.key} to true.")
        TimestampType

      case BINARY =>
        originalType match {
          case UTF8 | ENUM | JSON => StringType
          case null if assumeBinaryIsString => StringType
          case null => BinaryType
          case BSON => BinaryType
          case DECIMAL => makeDecimalType()
          case _ => illegalType()
        }

      case FIXED_LEN_BYTE_ARRAY =>
        originalType match {
          case DECIMAL => makeDecimalType(Decimal.maxPrecisionForBytes(field.getTypeLength))
          case INTERVAL => typeNotImplemented()
          case _ => illegalType()
        }

      case _ => illegalType()
    }
  }

 그럼 이제 저 함수를 어디서 부르는지 확인해 보면 될 것 같습니다. 그런데 convertPrimitiveField 함수는 convertField 함수에서 부르고 있습니다. 그리고 Parameter로 Type이 함께 넘어옵니다.

  def convertField(parquetType: Type): DataType = parquetType match {
    case t: PrimitiveType => convertPrimitiveField(t)
    case t: GroupType => convertGroupField(t.asGroupType())
  }

그리고 convertField 다시 convert 라는 함수에서 불려지고 있습니다. 코드를 보면 아시겠지만 파라매터로 GroupType 으로 Schema가 넘어오고 있고, map 후에 각각의 컬럼을 convertField로 처리한다는 것을 쉽게(?) 알 수 있습니다. 그럼 이제 convert를 부르는 곳을 확인해봅시다.

  private def convert(parquetSchema: GroupType): StructType = {
    val fields = parquetSchema.getFields.asScala.map { field =>
      field.getRepetition match {
        case OPTIONAL =>
          StructField(field.getName, convertField(field), nullable = true)

        case REQUIRED =>
          StructField(field.getName, convertField(field), nullable = false)

        case REPEATED =>
          // A repeated field that is neither contained by a `LIST`- or `MAP`-annotated group nor
          // annotated by `LIST` or `MAP` should be interpreted as a required list of required
          // elements where the element type is the type of the field.
          val arrayType = ArrayType(convertField(field), containsNull = false)
          StructField(field.getName, arrayType, nullable = false)
      }
    }

    StructType(fields.toSeq)
  }

그리고 다시 convert 함수는 readSchemaFromFooter라는 함수에서 넘겨주는 fileMetaData를 사용합니다.

  def readSchemaFromFooter(
      footer: Footer, converter: ParquetToSparkSchemaConverter): StructType = {
    val fileMetaData = footer.getParquetMetadata.getFileMetaData
    fileMetaData
      .getKeyValueMetaData
      .asScala.toMap
      .get(ParquetReadSupport.SPARK_METADATA_KEY)
      .flatMap(deserializeSchemaString)
      .getOrElse(converter.convert(fileMetaData.getSchema))
  }

이렇게 주룩주룩 고구마 줄기 처럼 다시 따라가기 전에 처음으로 돌아가서 젤 먼저 호출해서 schema 메서드를 살펴봅시다. 아래와 같이 schema 는 userSpecifiedSchema 로 저장됩니다. (이 이름을 잘 기억해 둡시다. 우리가 준 Custom Schema를 담고 있습니다.)

  def schema(schema: StructType): DataFrameReader = {
    this.userSpecifiedSchema = Option(schema)
    this
  }

  def schema(schemaString: String): DataFrameReader = {
    this.userSpecifiedSchema = Option(StructType.fromDDL(schemaString))
    this
  }

그리고 다시 parquet 함수는 load 함수를 부르고 이것은 다시 loadV1Source를 호출하게 됩니다. 뭔가 userSpecifiedSchema 값을 넘기고 있는 것을 볼 수 있습니다.

  private def loadV1Source(paths: String*) = {
    // Code path for data source v1.
    sparkSession.baseRelationToDataFrame(
      DataSource.apply(
        sparkSession,
        paths = paths,
        userSpecifiedSchema = userSpecifiedSchema,
        className = source,
        options = extraOptions.toMap).resolveRelation())
  }

여기서 다시 DataSource의 resolveRelation 를 호출하고 그 안에서, 다시 getOrInferFileFormatSchema 를 호출합니다. 그리고 이 안에서 다시userSpecifiedSchema 를 건드립니다. 결론적으로 dataSchema 는 userSpecifiedSchema 가 있으면 그걸 그대로 사용하고, 없으면 format 에 맞는 inferSchema를 하면서 해당 파일내의 Schema를 가져오게 되는데, 여기서 아까 말한 convert가 호출되게 됩니다.

    val dataSchema = userSpecifiedSchema.map { schema =>
      StructType(schema.filterNot(f => partitionSchema.exists(p => equality(p.name, f.name))))
    }.orElse {
      format.inferSchema(
        sparkSession,
        caseInsensitiveOptions,
        tempFileIndex.allFiles())
    }.getOrElse {
      throw new AnalysisException(
        s"Unable to infer schema for $format. It must be specified manually.")
    }

그냥 요약하면 schema 메서도를 쓰면, 내부적으로 Parquet의 전체 Schema를 읽지 않고, 주어진 Schema로 읽어오기 때문에 문제가 없다라고 보시면 될 것 같습니다. 이게 무슨 소리냐!!!

이 글은, 제가 Parquet에서 Unsigned Integer를 읽으면서 에러가 난다라고 하자, 아 그거 될텐데요 하면서 순식간에 Spark 코드를 찾아서 이 과정을 알려주신 옆자리 동료님께 바칩니다.