From d170f872517b36efb9e6136cc415fa96566bbfa2 Mon Sep 17 00:00:00 2001 From: evenyag Date: Wed, 22 Jul 2026 16:01:34 +0800 Subject: [PATCH] fix(mito2): guard primary-key-only SST reads Signed-off-by: evenyag --- src/mito2/src/sst/parquet/file_range.rs | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/src/mito2/src/sst/parquet/file_range.rs b/src/mito2/src/sst/parquet/file_range.rs index f066cba6153..bbef934bc0a 100644 --- a/src/mito2/src/sst/parquet/file_range.rs +++ b/src/mito2/src/sst/parquet/file_range.rs @@ -30,7 +30,7 @@ use datatypes::schema::Schema; use mito_codec::row_converter::PrimaryKeyCodec; use parquet::arrow::arrow_reader::RowSelection; use parquet::file::metadata::ParquetMetaData; -use snafu::{OptionExt, ResultExt}; +use snafu::{OptionExt, ResultExt, ensure}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::RegionMetadataRef; use store_api::storage::{ColumnId, TimeSeriesRowSelector}; @@ -232,6 +232,11 @@ impl FileRange { } /// Creates a reader that returns only the encoded primary-key column. + /// + /// This returns the primary key as persisted in the SST and does not apply + /// [`FlatCompatBatch`]. Callers must not use it when the primary key needs + /// compatibility conversion and must interpret the key using the file's + /// metadata. The PK-only path currently supports sparse encoding only. pub(crate) async fn primary_key_reader( &self, fetch_metrics: Option<&ParquetFetchMetrics>, @@ -240,6 +245,18 @@ impl FileRange { return Ok(None); } + let file_encoding = self.context.read_format().metadata().primary_key_encoding; + ensure!( + file_encoding == PrimaryKeyEncoding::Sparse, + UnexpectedSnafu { + reason: format!( + "Primary-key-only SST reads require sparse encoding, file {} uses {} encoding", + self.file_handle().file_id(), + file_encoding, + ), + } + ); + let stream = self .context .reader_builder