diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index 0cc122573e..8af1773a64 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -866,9 +866,9 @@ async fn test_cache_null_primary_key_with_format(flat_format: bool) { #[tokio::test] async fn test_list_ssts() { test_list_ssts_with_format(false, r#" -ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000001/.parquet", file_size: 2701, index_file_path: Some("test/11_0000000001/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } -ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000002/.parquet", file_size: 2701, index_file_path: Some("test/11_0000000002/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } -ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test/22_0000000042/.parquet", file_size: 2701, index_file_path: Some("test/22_0000000042/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"# , +ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000001/.parquet", file_size: 2951, index_file_path: Some("test/11_0000000001/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } +ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000002/.parquet", file_size: 2951, index_file_path: Some("test/11_0000000002/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } +ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test/22_0000000042/.parquet", file_size: 2951, index_file_path: Some("test/22_0000000042/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"# , r#" StorageSstEntry { file_path: "test/11_0000000001/.parquet", file_size: None, last_modified_ms: None, node_id: None } StorageSstEntry { file_path: "test/11_0000000001/index/.puffin", file_size: None, last_modified_ms: None, node_id: None } @@ -877,9 +877,9 @@ StorageSstEntry { file_path: "test/11_0000000002/index/.puffin", file_s StorageSstEntry { file_path: "test/22_0000000042/.parquet", file_size: None, last_modified_ms: None, node_id: None } StorageSstEntry { file_path: "test/22_0000000042/index/.puffin", file_size: None, last_modified_ms: None, node_id: None }"#).await; test_list_ssts_with_format(true, r#" -ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000001/.parquet", file_size: 3099, index_file_path: Some("test/11_0000000001/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } -ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000002/.parquet", file_size: 3099, index_file_path: Some("test/11_0000000002/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } -ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test/22_0000000042/.parquet", file_size: 3099, index_file_path: Some("test/22_0000000042/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"#, +ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000001/.parquet", file_size: 3365, index_file_path: Some("test/11_0000000001/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } +ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000002/.parquet", file_size: 3365, index_file_path: Some("test/11_0000000002/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } +ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test/22_0000000042/.parquet", file_size: 3365, index_file_path: Some("test/22_0000000042/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"#, r#" StorageSstEntry { file_path: "test/11_0000000001/.parquet", file_size: None, last_modified_ms: None, node_id: None } StorageSstEntry { file_path: "test/11_0000000001/index/.puffin", file_size: None, last_modified_ms: None, node_id: None } diff --git a/src/mito2/src/sst.rs b/src/mito2/src/sst.rs index d8edb76de6..20c9383f1f 100644 --- a/src/mito2/src/sst.rs +++ b/src/mito2/src/sst.rs @@ -72,6 +72,17 @@ pub fn with_field_id(mut field: Field, column_id: u32) -> Field { field } +/// Parquet field ID base for internal columns (__primary_key, __sequence, __op_type). +/// Uses bit 30 to distinguish from user column IDs and fit in positive i32 range. +pub(crate) const INTERNAL_PARQUET_FIELD_ID_BASE: u32 = 1 << 30; + +/// Parquet field ID for the __primary_key column. +pub(crate) const PRIMARY_KEY_PARQUET_FIELD_ID: u32 = INTERNAL_PARQUET_FIELD_ID_BASE; +/// Parquet field ID for the __sequence column. +pub(crate) const SEQUENCE_PARQUET_FIELD_ID: u32 = INTERNAL_PARQUET_FIELD_ID_BASE + 1; +/// Parquet field ID for the __op_type column. +pub(crate) const OP_TYPE_PARQUET_FIELD_ID: u32 = INTERNAL_PARQUET_FIELD_ID_BASE + 2; + /// Gets the arrow schema to store in parquet. pub fn to_sst_arrow_schema(metadata: &RegionMetadata) -> SchemaRef { let fields = Fields::from_iter( @@ -255,18 +266,23 @@ pub(crate) fn tag_maybe_to_dictionary_field( pub(crate) fn internal_fields() -> [FieldRef; 3] { // Internal columns are always not null. [ - Arc::new(Field::new_dictionary( - PRIMARY_KEY_COLUMN_NAME, - ArrowDataType::UInt32, - ArrowDataType::Binary, - false, + Arc::new(with_field_id( + Field::new_dictionary( + PRIMARY_KEY_COLUMN_NAME, + ArrowDataType::UInt32, + ArrowDataType::Binary, + false, + ), + PRIMARY_KEY_PARQUET_FIELD_ID, )), - Arc::new(Field::new( - SEQUENCE_COLUMN_NAME, - ArrowDataType::UInt64, - false, + Arc::new(with_field_id( + Field::new(SEQUENCE_COLUMN_NAME, ArrowDataType::UInt64, false), + SEQUENCE_PARQUET_FIELD_ID, + )), + Arc::new(with_field_id( + Field::new(OP_TYPE_COLUMN_NAME, ArrowDataType::UInt8, false), + OP_TYPE_PARQUET_FIELD_ID, )), - Arc::new(Field::new(OP_TYPE_COLUMN_NAME, ArrowDataType::UInt8, false)), ] } @@ -277,11 +293,19 @@ pub(crate) fn override_pk_field_to_binary(schema: &SchemaRef) -> SchemaRef { .iter() .map(|field| { if field.name() == PRIMARY_KEY_COLUMN_NAME { - Arc::new(Field::new( + let mut new_field = Field::new( PRIMARY_KEY_COLUMN_NAME, ArrowDataType::Binary, field.is_nullable(), - )) + ); + // Preserve the field_id metadata so parquet readers that require + // all columns to carry a field_id don't fail. + if let Some(field_id) = field.metadata().get(PARQUET_FIELD_ID_KEY) { + new_field + .metadata_mut() + .insert(PARQUET_FIELD_ID_KEY.to_string(), field_id.clone()); + } + Arc::new(new_field) } else { field.clone() } diff --git a/src/mito2/src/sst/parquet/flat_format.rs b/src/mito2/src/sst/parquet/flat_format.rs index 39d8691a13..7e84d5ba61 100644 --- a/src/mito2/src/sst/parquet/flat_format.rs +++ b/src/mito2/src/sst/parquet/flat_format.rs @@ -932,12 +932,13 @@ mod tests { use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder}; use store_api::storage::RegionId; + use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; use super::*; use crate::read::read_columns::ReadColumns; use crate::sst::{ - FlatSchemaOptions, flat_sst_arrow_schema_column_num, override_pk_field_to_binary, - to_flat_sst_arrow_schema, + FlatSchemaOptions, PARQUET_FIELD_ID_KEY, PRIMARY_KEY_PARQUET_FIELD_ID, + flat_sst_arrow_schema_column_num, override_pk_field_to_binary, to_flat_sst_arrow_schema, }; /// Builds a `RegionMetadata` with the given number of tags and fields. @@ -1040,6 +1041,17 @@ mod tests { let output_schema = read_format.output_arrow_schema().unwrap(); let binary_schema = override_pk_field_to_binary(&output_schema); + // The __primary_key field must preserve its field_id metadata after + // being converted from dictionary to plain binary. + let pk_field = binary_schema + .field_with_name(PRIMARY_KEY_COLUMN_NAME) + .unwrap(); + assert_eq!( + pk_field.metadata().get(PARQUET_FIELD_ID_KEY), + Some(&PRIMARY_KEY_PARQUET_FIELD_ID.to_string()), + "__primary_key field must retain its PARQUET:field_id after override_pk_field_to_binary" + ); + // Repeat the second pk to verify identity keys (no dedup). let tag_keys = UInt32Array::from(vec![0u32, 1, 1]); let tag_values = Arc::new(StringArray::from(vec!["t0", "t1"])); diff --git a/src/mito2/src/sst/parquet/format.rs b/src/mito2/src/sst/parquet/format.rs index bb29143b69..a6338c1a1c 100644 --- a/src/mito2/src/sst/parquet/format.rs +++ b/src/mito2/src/sst/parquet/format.rs @@ -844,7 +844,10 @@ mod tests { use crate::sst::parquet::flat_format::{ FlatReadFormat, FlatWriteFormat, sequence_column_index, sst_column_id_indices, }; - use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema, with_field_id}; + use crate::sst::{ + FlatSchemaOptions, OP_TYPE_PARQUET_FIELD_ID, PRIMARY_KEY_PARQUET_FIELD_ID, + SEQUENCE_PARQUET_FIELD_ID, to_flat_sst_arrow_schema, with_field_id, + }; const TEST_SEQUENCE: u64 = 1; const TEST_OP_TYPE: u8 = OpType::Put as u8; @@ -910,10 +913,20 @@ mod tests { Box::new(ArrowDataType::Binary), ), false, - None, + Some(PRIMARY_KEY_PARQUET_FIELD_ID), + ), + make_field( + "__sequence", + ArrowDataType::UInt64, + false, + Some(SEQUENCE_PARQUET_FIELD_ID), + ), + make_field( + "__op_type", + ArrowDataType::UInt8, + false, + Some(OP_TYPE_PARQUET_FIELD_ID), ), - make_field("__sequence", ArrowDataType::UInt64, false, None), - make_field("__op_type", ArrowDataType::UInt8, false, None), ]; Arc::new(Schema::new(fields)) } @@ -1223,9 +1236,9 @@ mod tests { Some(4), Some(2), Some(5), - None, - None, - None, + Some(PRIMARY_KEY_PARQUET_FIELD_ID), + Some(SEQUENCE_PARQUET_FIELD_ID), + Some(OP_TYPE_PARQUET_FIELD_ID), ]; let fields: Vec<_> = build_test_flat_sst_schema() .fields()