[Enhancement] add key_metadata field in the iceberg distributed plan (StarRocks#45426)

Signed-off-by: stephen <[email protected]>

Author: stephen-shelby

be/src/runtime/metadata_result_writer.cpp

@@ -145,6 +145,7 @@ StatusOr<TFetchDataResultPtr> MetadataResultWriter::_process_chunk(Chunk* chunk)
 // 10 -> "file_sequence_number"
 // 11 -> "data_sequence_number"
 // 12 -> "column_stats"
+// 13 -> "key_metadata"
 Status MetadataResultWriter::_fill_iceberg_metadata(const Columns& columns, const Chunk* chunk,
                                                     TFetchDataResult* result) const {
     SCOPED_TIMER(_serialize_timer);
@@ -163,6 +164,7 @@ Status MetadataResultWriter::_fill_iceberg_metadata(const Columns& columns, cons
     auto file_sequence_number = down_cast<Int64Column*>(ColumnHelper::get_data_column(columns[10].get()));
     auto data_sequence_number = down_cast<Int64Column*>(ColumnHelper::get_data_column(columns[11].get()));
     auto iceberg_metrics = down_cast<BinaryColumn*>(ColumnHelper::get_data_column(columns[12].get()));
+    auto key_metadata = down_cast<BinaryColumn*>(ColumnHelper::get_data_column(columns[13].get()));
 
     std::vector<TMetadataEntry> meta_entries;
     int num_rows = chunk->num_rows();
@@ -212,6 +214,9 @@ Status MetadataResultWriter::_fill_iceberg_metadata(const Columns& columns, cons
         if (!columns[12]->is_null(i)) {
             iceberg_meta.__set_column_stats(iceberg_metrics->get_slice(i).to_string());
         }
+        if (!columns[13]->is_null(i)) {
+            iceberg_meta.__set_key_metadata(key_metadata->get_slice(i).to_string());
+        }
     }
 
     result->result_batch.rows.resize(num_rows);

fe/fe-core/src/main/java/com/starrocks/connector/metadata/iceberg/IcebergMetadataCollectJob.java

@@ -33,6 +33,7 @@ public class IcebergMetadataCollectJob extends MetadataCollectJob {
             ", file_sequence_number" + // BIGINT
             ", data_sequence_number " + // BIGINT
             ", column_stats " + // BINARY
+            ", key_metadata " + // BINARY
             "FROM `$catalogName`.`$dbName`.`$tableName$logical_iceberg_metadata` " +
             "FOR VERSION AS OF $snapshotId " +
             "WHERE $predicate'";

fe/fe-core/src/main/java/com/starrocks/connector/metadata/iceberg/LogicalIcebergMetadataTable.java

@@ -64,6 +64,7 @@ public static LogicalIcebergMetadataTable create(String catalogName, String orig
                         .column("file_sequence_number", ScalarType.createType(PrimitiveType.BIGINT))
                         .column("data_sequence_number", ScalarType.createType(PrimitiveType.BIGINT))
                         .column("column_stats", ScalarType.createType(PrimitiveType.VARBINARY))
+                        .column("key_metadata", ScalarType.createType(PrimitiveType.VARBINARY))
                         .build(),
                 originDb,
                 originTable,

fe/fe-core/src/main/java/org/apache/iceberg/MetadataParser.java

@@ -236,6 +236,9 @@ private ContentFile<?> parseThriftToIcebergDataFile(TMetadataEntry entry) {
         // build equality field id
         int[] equalityFieldIds = thrift.isSetEquality_ids() ? ArrayUtil.toIntArray(thrift.getEquality_ids()) : null;
 
+        // build key metadata
+        ByteBuffer keyMetadata = thrift.isSetKey_metadata() ? ByteBuffer.wrap(thrift.getKey_metadata()) : null;
+
         BaseFile<?> baseFile;
         // TODO(stephen): add keyMetadata field
         if (content == FileContent.DATA) {
@@ -246,7 +249,7 @@ private ContentFile<?> parseThriftToIcebergDataFile(TMetadataEntry entry) {
                     partitionData,
                     fileLength,
                     metrics,
-                    null,
+                    keyMetadata,
                     splitOffsets,
                     null);
         } else {
@@ -261,7 +264,7 @@ private ContentFile<?> parseThriftToIcebergDataFile(TMetadataEntry entry) {
                     equalityFieldIds,
                     sortId,
                     splitOffsets,
-                    null
+                    keyMetadata
             );
         }
 

fe/fe-core/src/test/java/com/starrocks/connector/iceberg/IcebergMetadataTest.java

@@ -1364,9 +1364,10 @@ org.apache.iceberg.Table getTable(String dbName, String tableName) throws StarRo
         MetadataCollectJob collectJob = new IcebergMetadataCollectJob("iceberg_catalog", "db", "table",
                 TResultSinkType.METADATA_ICEBERG, snapshotId, "");
         collectJob.init(starRocksAssert.getCtx().getSessionVariable());
-        String expectedSql = "SELECT content, file_path, file_format, spec_id, partition_data, record_count, " +
-                "file_size_in_bytes, split_offsets, sort_id, equality_ids, file_sequence_number, data_sequence_number , " +
-                "column_stats FROM `iceberg_catalog`.`db`.`table$logical_iceberg_metadata` FOR VERSION AS OF 1 WHERE 1=1'";
+        String expectedSql = "SELECT content, file_path, file_format, spec_id, partition_data, record_count," +
+                " file_size_in_bytes, split_offsets, sort_id, equality_ids, file_sequence_number," +
+                " data_sequence_number , column_stats , key_metadata FROM" +
+                " `iceberg_catalog`.`db`.`table$logical_iceberg_metadata` FOR VERSION AS OF 1 WHERE 1=1'";
         Assert.assertEquals(expectedSql, collectJob.getSql());
         Assert.assertNotNull(collectJob.getContext());
         Assert.assertTrue(collectJob.getContext().isMetadataContext());

gensrc/thrift/Data.thrift

@@ -142,6 +142,7 @@ struct TIcebergMetadata {
     11: optional i64 file_sequence_number
     12: optional i64 data_sequence_number
     13: optional binary column_stats;
+    14: optional binary key_metadata;
 }
 
 // Metadata data for metadata table

java-extensions/iceberg-metadata-reader/src/main/java/com/starrocks/connector/iceberg/IcebergMetadataScanner.java

@@ -35,7 +35,6 @@
 import org.apache.iceberg.expressions.Expression;
 import org.apache.iceberg.expressions.Expressions;
 import org.apache.iceberg.io.CloseableIterator;
-import org.apache.iceberg.util.ByteBuffers;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
@@ -46,6 +45,7 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 
+import static org.apache.iceberg.util.ByteBuffers.toByteArray;
 import static org.apache.iceberg.util.SerializationUtil.deserializeFromBase64;
 
 public class IcebergMetadataScanner extends ConnectorScanner {
@@ -243,6 +243,8 @@ private Object get(String columnName, ContentFile<?> file) {
                 return file.dataSequenceNumber();
             case "column_stats":
                 return getIcebergMetrics(file);
+            case "key_metadata":
+                return file.keyMetadata() == null ? null : toByteArray(file.keyMetadata());
             default:
                 throw new IllegalArgumentException("Unrecognized column name " + columnName);
         }
@@ -299,7 +301,7 @@ private Map<Integer, byte[]> convertByteBufferMap(Map<Integer, ByteBuffer> byteB
         return byteBufferMap.entrySet().stream()
                 .collect(Collectors.toMap(
                         Map.Entry::getKey,
-                        entry -> ByteBuffers.toByteArray(entry.getValue())));
+                        entry -> toByteArray(entry.getValue())));
     }
 
     private void parseRequiredTypes() {