[SPARK-34817][SQL] Read parquet unsigned types that stored as int32 physical type in parquet

sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedColumnReader.java

@@ -565,6 +565,10 @@ private void readIntBatch(int rowId, int num, WritableColumnVector column) throw
         canReadAsIntDecimal(column.dataType())) {
       defColumn.readIntegers(
           num, column, rowId, maxDefLevel, (VectorizedValuesReader) dataColumn);
+    } else if (column.dataType() == DataTypes.LongType) {
+      //  We use LongType to handle UINT32
+      defColumn.readIntegersAsUnsigned(
+          num, column, rowId, maxDefLevel, (VectorizedValuesReader) dataColumn);
     } else if (column.dataType() == DataTypes.ByteType) {
       defColumn.readBytes(
           num, column, rowId, maxDefLevel, (VectorizedValuesReader) dataColumn);

sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedPlainValuesReader.java

@@ -83,6 +83,15 @@ public final void readIntegers(int total, WritableColumnVector c, int rowId) {
     }
   }
 
+  @Override
+  public final void readIntegersAsUnsigned(int total, WritableColumnVector c, int rowId) {
+    int requiredBytes = total * 4;
+    ByteBuffer buffer = getBuffer(requiredBytes);
+    for (int i = 0; i < total; i += 1) {
+      c.putLong(rowId + i, Integer.toUnsignedLong(buffer.getInt()));
+    }
+  }
+
   // A fork of `readIntegers` to rebase the date values. For performance reasons, this method
   // iterates the values twice: check if we need to rebase first, then go to the optimized branch
   // if rebase is not needed.

sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedRleValuesReader.java

@@ -203,6 +203,41 @@ public void readIntegers(
     }
   }
 
+  // A fork of `readIntegers`, reading the signed integers as unsigned in long type
+  public void readIntegersAsUnsigned(
+    int total,
+    WritableColumnVector c,
+    int rowId,
+    int level,
+    VectorizedValuesReader data) throws IOException {
+    int left = total;
+    while (left > 0) {
+      if (this.currentCount == 0) this.readNextGroup();
+      int n = Math.min(left, this.currentCount);
+      switch (mode) {
+        case RLE:
+          if (currentValue == level) {
+            data.readIntegersAsUnsigned(n, c, rowId);
+          } else {
+            c.putNulls(rowId, n);
+          }
+          break;
+        case PACKED:
+          for (int i = 0; i < n; ++i) {
+            if (currentBuffer[currentBufferIdx++] == level) {
+              c.putLong(rowId + i, Integer.toUnsignedLong(data.readInteger()));
+            } else {
+              c.putNull(rowId + i);
+            }
+          }
+          break;
+      }
+      rowId += n;
+      left -= n;
+      currentCount -= n;
+    }
+  }
+
   // A fork of `readIntegers`, which rebases the date int value (days) before filling
   // the Spark column vector.
   public void readIntegersWithRebase(
@@ -602,6 +637,11 @@ public void readIntegers(int total, WritableColumnVector c, int rowId) {
     }
   }
 
+  @Override
+  public void readIntegersAsUnsigned(int total, WritableColumnVector c, int rowId) {
+    throw new UnsupportedOperationException("only readInts is valid.");
+  }
+
   @Override
   public void readIntegersWithRebase(
       int total, WritableColumnVector c, int rowId, boolean failIfRebase) {

sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedValuesReader.java

@@ -41,6 +41,7 @@ public interface VectorizedValuesReader {
   void readBytes(int total, WritableColumnVector c, int rowId);
   void readIntegers(int total, WritableColumnVector c, int rowId);
   void readIntegersWithRebase(int total, WritableColumnVector c, int rowId, boolean failIfRebase);
+  void readIntegersAsUnsigned(int total, WritableColumnVector c, int rowId);
   void readLongs(int total, WritableColumnVector c, int rowId);
   void readLongsWithRebase(int total, WritableColumnVector c, int rowId, boolean failIfRebase);
   void readFloats(int total, WritableColumnVector c, int rowId);

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowConverter.scala

@@ -252,6 +252,11 @@ private[parquet] class ParquetRowConverter(
       updater: ParentContainerUpdater): Converter with HasParentContainerUpdater = {
 
     catalystType match {
+      case LongType if parquetType.getOriginalType == OriginalType.UINT_32 =>
+        new ParquetPrimitiveConverter(updater) {
+          override def addInt(value: Int): Unit =
+            updater.setLong(Integer.toUnsignedLong(value))
+        }
       case BooleanType | IntegerType | LongType | FloatType | DoubleType | BinaryType =>
         new ParquetPrimitiveConverter(updater)
 

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaConverter.scala

@@ -130,13 +130,11 @@ class ParquetToSparkSchemaConverter(
       case INT32 =>
         originalType match {
           case INT_8 => ByteType
-          case INT_16 => ShortType
-          case INT_32 | null => IntegerType
+          case INT_16 | UINT_8 => ShortType
+          case INT_32 | UINT_16 | 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 UINT_32 => LongType
           case TIME_MILLIS => typeNotImplemented()
           case _ => illegalType()
         }

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala

@@ -313,18 +313,20 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSparkSession
   test("SPARK-10113 Support for unsigned Parquet logical types") {
     val parquetSchema = MessageTypeParser.parseMessageType(
       """message root {
-        |  required int32 c(UINT_32);
+        |  required INT32 a(UINT_8);
+        |  required INT32 b(UINT_16);
+        |  required INT32 c(UINT_32);
         |}
       """.stripMargin)
 
+    val expectedSparkTypes = Seq(ShortType, IntegerType, LongType)
+
     withTempPath { location =>
       val path = new Path(location.getCanonicalPath)
       val conf = spark.sessionState.newHadoopConf()
       writeMetadata(parquetSchema, path, conf)
-      val errorMessage = intercept[Throwable] {
-        spark.read.parquet(path.toString).printSchema()
-      }.toString
-      assert(errorMessage.contains("Parquet type not supported"))
+      val sparkTypes = spark.read.parquet(path.toString).schema.map(_.dataType)
+      assert(sparkTypes === expectedSparkTypes)
     }
   }
 
@@ -381,9 +383,25 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSparkSession
     checkCompressionCodec(CompressionCodecName.SNAPPY)
   }
 
+  private def createParquetWriter(schema: MessageType, path: Path): ParquetWriter[Group] = {
+    val testWriteSupport = new TestGroupWriteSupport(schema)
+    /**
+     * Provide a builder for constructing a parquet writer - after PARQUET-248 directly
+     * constructing the writer is deprecated and should be done through a builder. The default
+     * builders include Avro - but for raw Parquet writing we must create our own builder.
+     */
+    class ParquetWriterBuilder() extends
+      ParquetWriter.Builder[Group, ParquetWriterBuilder](path) {
+      override def getWriteSupport(conf: Configuration) = testWriteSupport
+
+      override def self() = this
+    }
+    new ParquetWriterBuilder().build()
+  }
+
   test("read raw Parquet file") {
     def makeRawParquetFile(path: Path): Unit = {
-      val schema = MessageTypeParser.parseMessageType(
+      val schemaStr =
         """
           |message root {
           |  required boolean _1;
@@ -392,22 +410,11 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSparkSession
           |  required float   _4;
           |  required double  _5;
           |}
-        """.stripMargin)
-
-      val testWriteSupport = new TestGroupWriteSupport(schema)
-      /**
-       * Provide a builder for constructing a parquet writer - after PARQUET-248 directly
-       * constructing the writer is deprecated and should be done through a builder. The default
-       * builders include Avro - but for raw Parquet writing we must create our own builder.
-       */
-      class ParquetWriterBuilder() extends
-          ParquetWriter.Builder[Group, ParquetWriterBuilder](path) {
-        override def getWriteSupport(conf: Configuration) = testWriteSupport
-
-        override def self() = this
-      }
+        """.stripMargin
+      val schema = MessageTypeParser.parseMessageType(schemaStr)
+
 
-      val writer = new ParquetWriterBuilder().build()
+      val writer = createParquetWriter(schema, path)
 
       (0 until 10).foreach { i =>
         val record = new SimpleGroup(schema)
@@ -432,6 +439,42 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSparkSession
     }
   }
 
+  test("SPARK-34817: Read UINT_8/UINT_16/UINT_32 from parquet") {
+    def makeRawParquetFile(path: Path): Unit = {
+      val schemaStr =
+        """message root {
+          |  required INT32 a(UINT_8);
+          |  required INT32 b(UINT_16);
+          |  required INT32 c(UINT_32);
+          |}
+        """.stripMargin
+      val schema = MessageTypeParser.parseMessageType(schemaStr)
+
+
+      val writer = createParquetWriter(schema, path)
+
+      (0 until 10).foreach { i =>
+        val record = new SimpleGroup(schema)
+        record.add(0, i + Byte.MaxValue)
+        record.add(1, i + Short.MaxValue)
+        record.add(2, i + Int.MaxValue)
+        writer.write(record)
+      }
+      writer.close()
+    }
+
+    withTempDir { dir =>
+      val path = new Path(dir.toURI.toString, "part-r-0.parquet")
+      makeRawParquetFile(path)
+      readParquetFile(path.toString) { df =>
+        checkAnswer(df, (0 until 10).map { i =>
+          Row(i + Byte.MaxValue.toShort, i + Short.MaxValue.toInt, i + Int.MaxValue.toLong)
+          Row(i + Byte.MaxValue, i + Short.MaxValue, i + Int.MaxValue.toLong)
+
+        })
+      }
+    }
+  }
   test("write metadata") {
     val hadoopConf = spark.sessionState.newHadoopConf()
     withTempPath { file =>