This is an automated email from the ASF dual-hosted git repository.
lzljs3620320 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 105e50bd0 [spark] Support partition and bucket metadata column (#4180)
105e50bd0 is described below
commit 105e50bd0b7a707e355355789582632b83b5ed15
Author: Xiduo You <[email protected]>
AuthorDate: Thu Sep 19 10:03:02 2024 +0800
[spark] Support partition and bucket metadata column (#4180)
---
.../org/apache/paimon/spark/SparkTypeUtils.java | 10 +++++
.../org/apache/paimon/spark/PaimonBaseScan.scala | 4 +-
.../paimon/spark/PaimonPartitionReader.scala | 5 ++-
.../paimon/spark/PaimonRecordReaderIterator.scala | 26 ++++++++----
.../scala/org/apache/paimon/spark/SparkTable.scala | 8 +++-
.../paimon/spark/schema/PaimonMetadataColumn.scala | 33 ++++++++++-----
.../apache/paimon/spark/sql/PaimonQueryTest.scala | 47 +++++++++++++++++++++-
7 files changed, 110 insertions(+), 23 deletions(-)
diff --git
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkTypeUtils.java
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkTypeUtils.java
index 08fb2de32..bc5eed00e 100644
---
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkTypeUtils.java
+++
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkTypeUtils.java
@@ -19,6 +19,7 @@
package org.apache.paimon.spark;
import org.apache.paimon.spark.util.shim.TypeUtils;
+import org.apache.paimon.table.Table;
import org.apache.paimon.types.ArrayType;
import org.apache.paimon.types.BigIntType;
import org.apache.paimon.types.BinaryType;
@@ -59,6 +60,15 @@ public class SparkTypeUtils {
private SparkTypeUtils() {}
+ public static StructType toSparkPartitionType(Table table) {
+ int[] projections =
table.rowType().getFieldIndices(table.partitionKeys());
+ List<DataField> partitionTypes = new ArrayList<>();
+ for (int i : projections) {
+ partitionTypes.add(table.rowType().getFields().get(i));
+ }
+ return (StructType) SparkTypeUtils.fromPaimonType(new RowType(false,
partitionTypes));
+ }
+
public static StructType fromPaimonRowType(RowType type) {
return (StructType) fromPaimonType(type);
}
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonBaseScan.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonBaseScan.scala
index 2188879c7..cb43560dc 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonBaseScan.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonBaseScan.scala
@@ -118,7 +118,9 @@ abstract class PaimonBaseScan(
}
override def toBatch: Batch = {
- val metadataColumns = metadataFields.map(field =>
PaimonMetadataColumn.get(field.name))
+ val partitionType = SparkTypeUtils.toSparkPartitionType(table)
+ val metadataColumns =
+ metadataFields.map(field => PaimonMetadataColumn.get(field.name,
partitionType))
PaimonBatch(lazyInputPartitions, readBuilder, metadataColumns)
}
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonPartitionReader.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonPartitionReader.scala
index cea235120..fa9072df3 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonPartitionReader.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonPartitionReader.scala
@@ -87,8 +87,9 @@ case class PaimonPartitionReader(
private def readSplit(): PaimonRecordReaderIterator = {
if (splits.hasNext) {
- val reader = readFunc(splits.next())
- PaimonRecordReaderIterator(reader, metadataColumns)
+ val split = splits.next();
+ val reader = readFunc(split)
+ PaimonRecordReaderIterator(reader, metadataColumns, split)
} else {
null
}
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala
index ce880e670..6d6c1ca47 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala
@@ -22,6 +22,7 @@ import org.apache.paimon.data.{BinaryString, GenericRow,
InternalRow => PaimonIn
import org.apache.paimon.fs.Path
import org.apache.paimon.reader.{FileRecordIterator, RecordReader}
import org.apache.paimon.spark.schema.PaimonMetadataColumn
+import org.apache.paimon.table.source.{DataSplit, Split}
import org.apache.paimon.utils.CloseableIterator
import org.apache.spark.sql.PaimonUtils
@@ -30,7 +31,8 @@ import java.io.IOException
case class PaimonRecordReaderIterator(
reader: RecordReader[PaimonInternalRow],
- metadataColumns: Seq[PaimonMetadataColumn])
+ metadataColumns: Seq[PaimonMetadataColumn],
+ split: Split)
extends CloseableIterator[PaimonInternalRow] {
private var lastFilePath: Path = _
@@ -44,6 +46,15 @@ case class PaimonRecordReaderIterator(
GenericRow.of(Array.fill(metadataColumns.size)(null.asInstanceOf[AnyRef]):
_*)
private val joinedRow: JoinedRow = JoinedRow.join(null, metadataRow)
+ private def validateMetadataColumns(): Unit = {
+ if (needMetadata) {
+ if (!isFileRecordIterator || !split.isInstanceOf[DataSplit]) {
+ throw new RuntimeException(
+ "There need be FileRecordIterator when metadata columns are
required.")
+ }
+ }
+ }
+
override def hasNext: Boolean = {
if (currentIterator == null) {
false
@@ -98,10 +109,7 @@ case class PaimonRecordReaderIterator(
val dataRow = currentIterator.next()
if (dataRow != null) {
if (needMetadata) {
- if (!isFileRecordIterator) {
- throw new RuntimeException(
- "There need be FileRecoredIterator when metadata columns are
required.")
- }
+ validateMetadataColumns()
updateMetadataRow(currentIterator.asInstanceOf[FileRecordIterator[PaimonInternalRow]])
currentResult = joinedRow.replace(dataRow, metadataRow)
} else {
@@ -133,10 +141,14 @@ case class PaimonRecordReaderIterator(
metadataColumns.zipWithIndex.foreach {
case (metadataColumn, index) =>
metadataColumn.name match {
- case PaimonMetadataColumn.ROW_INDEX.name =>
+ case PaimonMetadataColumn.ROW_INDEX_COLUMN =>
metadataRow.setField(index, fileRecordIterator.returnedPosition())
- case PaimonMetadataColumn.FILE_PATH.name =>
+ case PaimonMetadataColumn.FILE_PATH_COLUMN =>
metadataRow.setField(index,
BinaryString.fromString(lastFilePath.toUri.toString))
+ case PaimonMetadataColumn.PARTITION_COLUMN =>
+ metadataRow.setField(index,
split.asInstanceOf[DataSplit].partition())
+ case PaimonMetadataColumn.BUCKET_COLUMN =>
+ metadataRow.setField(index, split.asInstanceOf[DataSplit].bucket())
}
}
}
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/SparkTable.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/SparkTable.scala
index dd24fc6b1..d06c125d2 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/SparkTable.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/SparkTable.scala
@@ -81,7 +81,13 @@ case class SparkTable(table: Table)
}
override def metadataColumns: Array[MetadataColumn] = {
- Array[MetadataColumn](PaimonMetadataColumn.FILE_PATH,
PaimonMetadataColumn.ROW_INDEX)
+ val partitionType = SparkTypeUtils.toSparkPartitionType(table)
+ Array[MetadataColumn](
+ PaimonMetadataColumn.FILE_PATH,
+ PaimonMetadataColumn.ROW_INDEX,
+ PaimonMetadataColumn.PARTITION(partitionType),
+ PaimonMetadataColumn.BUCKET
+ )
}
override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder
= {
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/schema/PaimonMetadataColumn.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/schema/PaimonMetadataColumn.scala
index b539a1351..991f6f9c9 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/schema/PaimonMetadataColumn.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/schema/PaimonMetadataColumn.scala
@@ -23,7 +23,7 @@ import org.apache.paimon.types.DataField
import org.apache.spark.sql.catalyst.expressions.AttributeReference
import org.apache.spark.sql.connector.catalog.MetadataColumn
-import org.apache.spark.sql.types.{DataType, LongType, StringType, StructField}
+import org.apache.spark.sql.types.{DataType, IntegerType, LongType,
StringType, StructField, StructType}
case class PaimonMetadataColumn(id: Int, override val name: String, override
val dataType: DataType)
extends MetadataColumn {
@@ -43,22 +43,35 @@ case class PaimonMetadataColumn(id: Int, override val name:
String, override val
object PaimonMetadataColumn {
- val FILE_PATH_COLUMN = "__paimon_file_path"
val ROW_INDEX_COLUMN = "__paimon_row_index"
- val SUPPORTED_METADATA_COLUMNS: Seq[String] = Seq(FILE_PATH_COLUMN,
ROW_INDEX_COLUMN)
+ val FILE_PATH_COLUMN = "__paimon_file_path"
+ val PARTITION_COLUMN = "__paimon_partition"
+ val BUCKET_COLUMN = "__paimon_bucket"
+ val SUPPORTED_METADATA_COLUMNS: Seq[String] = Seq(
+ ROW_INDEX_COLUMN,
+ FILE_PATH_COLUMN,
+ PARTITION_COLUMN,
+ BUCKET_COLUMN
+ )
val ROW_INDEX: PaimonMetadataColumn =
PaimonMetadataColumn(Int.MaxValue - 100, ROW_INDEX_COLUMN, LongType)
val FILE_PATH: PaimonMetadataColumn =
PaimonMetadataColumn(Int.MaxValue - 101, FILE_PATH_COLUMN, StringType)
+ val PARTITION: DataType => PaimonMetadataColumn = (dt: DataType) => {
+ PaimonMetadataColumn(Int.MaxValue - 102, PARTITION_COLUMN, dt)
+ }
+ val BUCKET: PaimonMetadataColumn =
+ PaimonMetadataColumn(Int.MaxValue - 103, BUCKET_COLUMN, IntegerType)
- def get(metadataColumn: String): PaimonMetadataColumn = {
- if (metadataColumn == FILE_PATH_COLUMN) {
- FILE_PATH
- } else if (metadataColumn == ROW_INDEX_COLUMN) {
- ROW_INDEX
- } else {
- throw new IllegalArgumentException(s"$metadataColumn metadata column is
not supported.")
+ def get(metadataColumn: String, partitionType: StructType):
PaimonMetadataColumn = {
+ metadataColumn match {
+ case ROW_INDEX_COLUMN => ROW_INDEX
+ case FILE_PATH_COLUMN => FILE_PATH
+ case PARTITION_COLUMN => PARTITION(partitionType)
+ case BUCKET_COLUMN => BUCKET
+ case _ =>
+ throw new IllegalArgumentException(s"$metadataColumn metadata column
is not supported.")
}
}
}
diff --git
a/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/sql/PaimonQueryTest.scala
b/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/sql/PaimonQueryTest.scala
index 702965701..c1814096f 100644
---
a/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/sql/PaimonQueryTest.scala
+++
b/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/sql/PaimonQueryTest.scala
@@ -29,6 +29,7 @@ import java.util
import scala.collection.JavaConverters._
class PaimonQueryTest extends PaimonSparkTestBase {
+ import testImplicits._
fileFormats.foreach {
fileFormat =>
@@ -66,13 +67,55 @@ class PaimonQueryTest extends PaimonSparkTestBase {
}
+ test("Query metadata columns for bucket") {
+ withTable("T") {
+ spark.sql(
+ """
+ |CREATE TABLE T (c1 INT, c2 STRING) TBLPROPERTIES
('bucket-key'='c1', 'bucket'='3')
+ |""".stripMargin)
+
+ spark.sql("""
+ |INSERT INTO T
+ |VALUES (1, 'x1'), (2, 'x3'), (3, 'x3'), (4, 'x4'), (5,
'x5'), (6, 'x6')
+ |""".stripMargin)
+
+ val res = spark.sql("""
+ |SELECT __paimon_partition, __paimon_bucket FROM T
+ |GROUP BY __paimon_partition, __paimon_bucket
+ |ORDER BY __paimon_partition, __paimon_bucket
+ |""".stripMargin)
+ checkAnswer(res, Row(Row(), 0) :: Row(Row(), 1) :: Row(Row(), 2) :: Nil)
+ }
+ }
+
+ test("Query metadata columns for partition") {
+ withTable("T") {
+ spark.sql("""
+ |CREATE TABLE T (c1 INT) PARTITIONED BY(p1 INT, p2 String)
+ |""".stripMargin)
+
+ spark.sql("""
+ |INSERT INTO T
+ |VALUES (1, 1, 'x1'), (2, 1, 'x2'), (3, 2, 'x1'), (4, 2,
'x2'), (5, 3, 'x3')
+ |""".stripMargin)
+
+ val res = spark.sql("""
+ |SELECT __paimon_partition, __paimon_bucket FROM T
+ |GROUP BY __paimon_partition, __paimon_bucket
+ |ORDER BY __paimon_partition, __paimon_bucket
+ |""".stripMargin)
+ checkAnswer(
+ res,
+ Row(Row(1, "x1"), 0) :: Row(Row(1, "x2"), 0) :: Row(Row(2, "x1"), 0) ::
+ Row(Row(2, "x2"), 0) :: Row(Row(3, "x3"), 0) :: Nil)
+ }
+ }
+
fileFormats.foreach {
fileFormat =>
bucketModes.foreach {
bucketMode =>
test(s"Query input_file_name(): file.format=$fileFormat,
bucket=$bucketMode") {
- val _spark: SparkSession = spark
- import _spark.implicits._
withTable("T") {
val bucketProp = if (bucketMode != -1) {
