[core] Parse primary-key vector index source metadata (_SOURCE_META) (#515)

Author: JunRuiLee

crates/paimon/src/spec/avro/index_manifest_entry_decode.rs

@@ -18,7 +18,8 @@
 use super::cursor::AvroCursor;
 use super::decode::{neg_count_to_usize, AvroRecordDecode};
 use super::decode_helpers::{
-    normalize_partition, read_bytes_field, read_int_field, read_long_field, read_string_field,
+    extract_record_schema, normalize_partition, read_bytes_field, read_int_field, read_long_field,
+    read_string_field,
 };
 use super::schema::{skip_nullable_field, WriterSchema};
 use crate::spec::index_manifest::IndexManifestEntry;
@@ -65,7 +66,8 @@ impl AvroRecordDecode for IndexManifestEntry {
                     deletion_vectors_ranges = decode_nullable_dv_ranges(cursor, field.nullable)?;
                 }
                 "_GLOBAL_INDEX" => {
-                    global_index_meta = decode_nullable_global_index(cursor, field.nullable)?;
+                    global_index_meta =
+                        decode_nullable_global_index(cursor, field.nullable, &field.schema)?;
                 }
                 _ => skip_nullable_field(cursor, &field.schema, field.nullable)?,
             }
@@ -145,6 +147,7 @@ fn decode_nullable_dv_ranges(
 fn decode_nullable_global_index(
     cursor: &mut AvroCursor,
     nullable: bool,
+    schema: &super::schema::FieldSchema,
 ) -> crate::Result<Option<GlobalIndexMeta>> {
     if nullable {
         let idx = cursor.read_union_index()?;
@@ -192,11 +195,29 @@ fn decode_nullable_global_index(
         }
     };
 
+    // _SOURCE_META: nullable bytes — only present in >= #8549 writer schemas.
+    // Guard on the writer's nested field list so a legacy 5-field _GLOBAL_INDEX
+    // record does not misalign the cursor into the next record.
+    let has_source_meta = extract_record_schema(schema)
+        .map(|s| s.fields.iter().any(|f| f.name == "_SOURCE_META"))
+        .unwrap_or(false);
+    let source_meta = if has_source_meta {
+        let u_idx = cursor.read_union_index()?;
+        if u_idx == 0 {
+            None
+        } else {
+            Some(cursor.read_bytes()?.to_vec())
+        }
+    } else {
+        None
+    };
+
     Ok(Some(GlobalIndexMeta {
         row_range_start,
         row_range_end,
         index_field_id,
         extra_field_ids,
         index_meta,
+        source_meta,
     }))
 }

crates/paimon/src/spec/index_file_meta.rs

@@ -46,6 +46,9 @@ pub struct GlobalIndexMeta {
 
     #[serde(default, rename = "_INDEX_META", with = "serde_bytes")]
     pub index_meta: Option<Vec<u8>>,
+
+    #[serde(default, rename = "_SOURCE_META", with = "serde_bytes")]
+    pub source_meta: Option<Vec<u8>>,
 }
 
 /// Metadata of index file.

crates/paimon/src/spec/index_manifest.rs

@@ -71,7 +71,8 @@ pub const INDEX_MANIFEST_ENTRY_SCHEMA: &str = r#"{
                     {"name": "_ROW_RANGE_END", "type": "long"},
                     {"name": "_INDEX_FIELD_ID", "type": "int"},
                     {"name": "_EXTRA_FIELD_IDS", "type": ["null", {"type": "array", "items": "int"}], "default": null},
-                    {"name": "_INDEX_META", "type": ["null", "bytes"], "default": null}
+                    {"name": "_INDEX_META", "type": ["null", "bytes"], "default": null},
+                    {"name": "_SOURCE_META", "type": ["null", "bytes"], "default": null}
                 ]
             }]
         }
@@ -173,7 +174,7 @@ mod tests {
     use indexmap::IndexMap;
 
     use super::*;
-    use crate::spec::DeletionVectorMeta;
+    use crate::spec::{DeletionVectorMeta, GlobalIndexMeta};
 
     #[test]
     fn test_read_index_manifest_file() {
@@ -276,4 +277,144 @@ mod tests {
         };
         assert_eq!(sample, decoded);
     }
+
+    fn global_index_entry(source_meta: Option<Vec<u8>>) -> IndexManifestEntry {
+        IndexManifestEntry {
+            version: 1,
+            kind: FileKind::Add,
+            partition: vec![0, 1, 0, 2, 0, 3, 0, 4, 0, 5, 0, 6],
+            bucket: 0,
+            index_file: IndexFileMeta {
+                index_type: "GLOBAL_INDEX".into(),
+                file_name: "gi-1".into(),
+                file_size: 42,
+                row_count: 7,
+                deletion_vectors_ranges: None,
+                global_index_meta: Some(GlobalIndexMeta {
+                    row_range_start: 10,
+                    row_range_end: 20,
+                    index_field_id: 3,
+                    extra_field_ids: Some(vec![4, 5]),
+                    index_meta: Some(vec![9, 8, 7]),
+                    source_meta,
+                }),
+            },
+        }
+    }
+
+    #[test]
+    fn source_meta_round_trips_through_index_manifest() {
+        // New-format writer schema carries _SOURCE_META; a Some(..) value must round-trip.
+        let entry = global_index_entry(Some(vec![1, 2, 3]));
+        let bytes = crate::spec::to_avro_bytes_with_compression(
+            INDEX_MANIFEST_ENTRY_SCHEMA,
+            std::slice::from_ref(&entry),
+            crate::spec::DEFAULT_AVRO_COMPRESSION,
+        )
+        .unwrap();
+        let decoded = IndexManifest::read_from_bytes(&bytes).unwrap();
+        assert_eq!(decoded[0], entry);
+        assert_eq!(
+            decoded[0]
+                .index_file
+                .global_index_meta
+                .as_ref()
+                .unwrap()
+                .source_meta,
+            Some(vec![1, 2, 3])
+        );
+
+        // A None source_meta must also round-trip as None.
+        let entry_none = global_index_entry(None);
+        let bytes_none = crate::spec::to_avro_bytes_with_compression(
+            INDEX_MANIFEST_ENTRY_SCHEMA,
+            std::slice::from_ref(&entry_none),
+            crate::spec::DEFAULT_AVRO_COMPRESSION,
+        )
+        .unwrap();
+        let decoded_none = IndexManifest::read_from_bytes(&bytes_none).unwrap();
+        assert_eq!(decoded_none[0], entry_none);
+        assert_eq!(
+            decoded_none[0]
+                .index_file
+                .global_index_meta
+                .as_ref()
+                .unwrap()
+                .source_meta,
+            None
+        );
+    }
+
+    #[test]
+    fn legacy_five_field_global_index_decodes_without_source_meta() {
+        // 5-field _GLOBAL_INDEX schema (pre-#8549): no _SOURCE_META. Identical to
+        // INDEX_MANIFEST_ENTRY_SCHEMA with the trailing _SOURCE_META line removed.
+        const LEGACY_SCHEMA: &str = r#"{
+    "type": "record",
+    "name": "org.apache.paimon.avro.generated.record",
+    "fields": [
+        {"name": "_VERSION", "type": "int"},
+        {"name": "_KIND", "type": "int"},
+        {"name": "_PARTITION", "type": "bytes"},
+        {"name": "_BUCKET", "type": "int"},
+        {"name": "_INDEX_TYPE", "type": "string"},
+        {"name": "_FILE_NAME", "type": "string"},
+        {"name": "_FILE_SIZE", "type": "long"},
+        {"name": "_ROW_COUNT", "type": "long"},
+        {
+            "default": null,
+            "name": "_DELETIONS_VECTORS_RANGES",
+            "type": ["null", {
+                "type": "array",
+                "items": ["null", {
+                    "type": "record",
+                    "name": "org.apache.paimon.avro.generated.record__DELETIONS_VECTORS_RANGES",
+                    "fields": [
+                        {"name": "f0", "type": "string"},
+                        {"name": "f1", "type": "int"},
+                        {"name": "f2", "type": "int"},
+                        {"name": "_CARDINALITY", "type": ["null", "long"], "default": null}
+                    ]
+                }]
+            }]
+        },
+        {
+            "default": null,
+            "name": "_GLOBAL_INDEX",
+            "type": ["null", {
+                "type": "record",
+                "name": "org.apache.paimon.avro.generated.record__GLOBAL_INDEX",
+                "fields": [
+                    {"name": "_ROW_RANGE_START", "type": "long"},
+                    {"name": "_ROW_RANGE_END", "type": "long"},
+                    {"name": "_INDEX_FIELD_ID", "type": "int"},
+                    {"name": "_EXTRA_FIELD_IDS", "type": ["null", {"type": "array", "items": "int"}], "default": null},
+                    {"name": "_INDEX_META", "type": ["null", "bytes"], "default": null}
+                ]
+            }]
+        }
+    ]
+}"#;
+
+        // Written by a pre-#8549 writer (5-field record, source_meta absent).
+        let entry = global_index_entry(None);
+        let bytes = crate::spec::to_avro_bytes_with_compression(
+            LEGACY_SCHEMA,
+            std::slice::from_ref(&entry),
+            crate::spec::DEFAULT_AVRO_COMPRESSION,
+        )
+        .unwrap();
+        // Decoding with the current 6-field reader must not misalign the stream.
+        let decoded = IndexManifest::read_from_bytes(&bytes).unwrap();
+        assert_eq!(decoded[0], entry);
+        assert_eq!(
+            decoded[0]
+                .index_file
+                .global_index_meta
+                .as_ref()
+                .unwrap()
+                .source_meta,
+            None
+        );
+    }
 }

crates/paimon/src/spec/mod.rs

@@ -64,6 +64,9 @@ pub use manifest_file_meta::*;
 mod index_file_meta;
 pub use index_file_meta::*;
 
+mod pk_vector_source;
+pub use pk_vector_source::*;
+
 mod index_manifest;
 pub use index_manifest::{IndexManifest, IndexManifestEntry};
 mod manifest;