feat: add field_id for internal fields (#8284)

* feat: add field_id for internal fields

* test: address file size change after including field ids

* fix: address review comments
This commit is contained in:
Ning Sun
2026-06-12 11:10:24 +08:00
committed by GitHub
parent 827b904189
commit ba89ef40f4
4 changed files with 76 additions and 27 deletions

View File

@@ -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: "<file_id>", index_version: 0, level: 0, file_path: "test/11_0000000001/<file_id>.parquet", file_size: 2701, index_file_path: Some("test/11_0000000001/index/<file_id>.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: "<file_id>", index_version: 0, level: 0, file_path: "test/11_0000000002/<file_id>.parquet", file_size: 2701, index_file_path: Some("test/11_0000000002/index/<file_id>.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: "<file_id>", index_version: 0, level: 0, file_path: "test/22_0000000042/<file_id>.parquet", file_size: 2701, index_file_path: Some("test/22_0000000042/index/<file_id>.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: "<file_id>", index_version: 0, level: 0, file_path: "test/11_0000000001/<file_id>.parquet", file_size: 2951, index_file_path: Some("test/11_0000000001/index/<file_id>.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: "<file_id>", index_version: 0, level: 0, file_path: "test/11_0000000002/<file_id>.parquet", file_size: 2951, index_file_path: Some("test/11_0000000002/index/<file_id>.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: "<file_id>", index_version: 0, level: 0, file_path: "test/22_0000000042/<file_id>.parquet", file_size: 2951, index_file_path: Some("test/22_0000000042/index/<file_id>.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/<file_id>.parquet", file_size: None, last_modified_ms: None, node_id: None }
StorageSstEntry { file_path: "test/11_0000000001/index/<file_id>.puffin", file_size: None, last_modified_ms: None, node_id: None }
@@ -877,9 +877,9 @@ StorageSstEntry { file_path: "test/11_0000000002/index/<file_id>.puffin", file_s
StorageSstEntry { file_path: "test/22_0000000042/<file_id>.parquet", file_size: None, last_modified_ms: None, node_id: None }
StorageSstEntry { file_path: "test/22_0000000042/index/<file_id>.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: "<file_id>", index_version: 0, level: 0, file_path: "test/11_0000000001/<file_id>.parquet", file_size: 3099, index_file_path: Some("test/11_0000000001/index/<file_id>.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: "<file_id>", index_version: 0, level: 0, file_path: "test/11_0000000002/<file_id>.parquet", file_size: 3099, index_file_path: Some("test/11_0000000002/index/<file_id>.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: "<file_id>", index_version: 0, level: 0, file_path: "test/22_0000000042/<file_id>.parquet", file_size: 3099, index_file_path: Some("test/22_0000000042/index/<file_id>.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: "<file_id>", index_version: 0, level: 0, file_path: "test/11_0000000001/<file_id>.parquet", file_size: 3365, index_file_path: Some("test/11_0000000001/index/<file_id>.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: "<file_id>", index_version: 0, level: 0, file_path: "test/11_0000000002/<file_id>.parquet", file_size: 3365, index_file_path: Some("test/11_0000000002/index/<file_id>.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: "<file_id>", index_version: 0, level: 0, file_path: "test/22_0000000042/<file_id>.parquet", file_size: 3365, index_file_path: Some("test/22_0000000042/index/<file_id>.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/<file_id>.parquet", file_size: None, last_modified_ms: None, node_id: None }
StorageSstEntry { file_path: "test/11_0000000001/index/<file_id>.puffin", file_size: None, last_modified_ms: None, node_id: None }

View File

@@ -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()
}

View File

@@ -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"]));

View File

@@ -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()