diff --git a/src/mito2/src/series_index/searcher.rs b/src/mito2/src/series_index/searcher.rs index 9b88a5b826..b5b45e2ed6 100644 --- a/src/mito2/src/series_index/searcher.rs +++ b/src/mito2/src/series_index/searcher.rs @@ -12,23 +12,23 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::SemanticType; use async_stream::try_stream; use common_recordbatch::filter::SimpleFilterEvaluator; use common_time::range::TimestampRange; +use common_time::timestamp::TimeUnit; use datafusion_expr::{Expr, col, lit}; use datatypes::arrow::array::{ArrayRef, UInt32Array, UInt64Array}; use datatypes::arrow::buffer::BooleanBuffer; use datatypes::arrow::datatypes::{DataType, SchemaRef}; +use datatypes::arrow::record_batch::RecordBatch; +use datatypes::value::timestamp_to_scalar_value; use futures::TryStreamExt; use object_store::ObjectStore; use snafu::{OptionExt, ResultExt, ensure}; use store_api::metadata::RegionMetadataRef; use table::predicate::Predicate; -use crate::error::{ - InvalidMetaSnafu, InvalidRecordBatchSnafu, RecordBatchSnafu, Result, UnexpectedSnafu, -}; +use crate::error::{InvalidRecordBatchSnafu, RecordBatchSnafu, Result, UnexpectedSnafu}; use crate::series_index::{ MAX_TS_COLUMN, METRIC_SERIES_ID_BATCH_SIZE, MIN_TS_COLUMN, MetricSeriesId, MetricSeriesIdStream, ROW_COUNT_COLUMN, TABLE_ID_COLUMN, TSID_COLUMN, series_index_schema, @@ -38,54 +38,68 @@ use crate::sst::parquet::prefilter::simple_tag_filters; /// Searches a series-index file for metric series matching query predicates. pub struct SeriesIndexSearcher { - object_store: ObjectStore, - filters: Vec<(Expr, SimpleFilterEvaluator)>, - empty_time_range: bool, + /// The file to search, or `None` when an empty time range rules out + /// every series without opening the file. + reader: Option, + /// Row-group pruning predicate built from all of the file's filters. + pruning_predicate: Predicate, + filters: Vec, } impl SeriesIndexSearcher { - /// Creates a searcher reusable across series-index files of `metadata`. - pub fn try_new( + /// Creates a searcher for the series-index file of `metadata` at `path`. + /// Predicates are built from the unit recorded in the file's schema, so a + /// file written before a time index unit widening keeps being interpreted + /// in its own unit. + pub async fn try_new( metadata: RegionMetadataRef, object_store: ObjectStore, + path: &str, predicate: Option<&Predicate>, time_range: Option, ) -> Result { // Keep search-time metadata validation identical to the writer. series_index_schema(&metadata)?; + if time_range.as_ref().is_some_and(TimestampRange::is_empty) { + return Ok(Self { + reader: None, + pruning_predicate: Predicate::new(Vec::new()), + filters: Vec::new(), + }); + } + + let reader = ParquetIndexReader::open(object_store, path).await?; + let unit = validate_index_schema(reader.schema())?; let mut filters = simple_tag_filters(&metadata, None, predicate); - let (empty_time_range, time_exprs) = time_range_filters(&metadata, time_range)?; - for expr in time_exprs { + for expr in time_range_exprs(unit, time_range.as_ref()) { let filter = SimpleFilterEvaluator::try_new(&expr).context(UnexpectedSnafu { reason: "failed to build an internal series-index time filter", })?; filters.push((expr, filter)); } + // An older index file may not contain tags added by schema evolution. + // Ignore filters on those tags to preserve a conservative candidate set. + let (pruning_predicate, filters) = filters_for_schema(reader.schema(), &filters); Ok(Self { - object_store, + reader: Some(reader), + pruning_predicate, filters, - empty_time_range, }) } - /// Searches `path` and returns sorted batches of matching metric-series IDs. - pub async fn search(&self, path: &str) -> Result { - if self.empty_time_range { + /// Searches the index file and returns sorted batches of matching + /// metric-series IDs. + pub fn search(&self) -> Result { + let Some(reader) = self.reader.as_ref() else { return Ok(Box::pin(futures::stream::empty())); - } - - let reader = ParquetIndexReader::open(self.object_store.clone(), path).await?; - validate_index_schema(reader.schema())?; - - // An older index file may not contain tags added by schema evolution. - // Ignore filters on those tags to preserve a conservative candidate set. - let (pruning_predicate, filters) = self.filters_for_schema(reader.schema()); - let mut projection_columns = Vec::with_capacity(filters.len() + 2); + }; + let mut projection_columns = Vec::with_capacity(self.filters.len() + 2); projection_columns.extend([TABLE_ID_COLUMN, TSID_COLUMN]); - projection_columns.extend(filters.iter().map(SimpleFilterEvaluator::column_name)); - let mut batches = reader.read(&pruning_predicate, &projection_columns)?; + projection_columns.extend(self.filters.iter().map(SimpleFilterEvaluator::column_name)); + let mut batches = reader.read(&self.pruning_predicate, &projection_columns)?; + let filters = self.filters.clone(); Ok(Box::pin(try_stream! { let mut last_series = None; @@ -137,50 +151,28 @@ impl SeriesIndexSearcher { } })) } - - fn filters_for_schema(&self, schema: &SchemaRef) -> (Predicate, Vec) { - let (exprs, filters): (Vec<_>, Vec<_>) = self - .filters - .iter() - .filter(|(_, filter)| schema.field_with_name(filter.column_name()).is_ok()) - .cloned() - .unzip(); - (Predicate::new(exprs), filters) - } } -// Builds `__series_min_ts`/`__series_max_ts` predicates in the unit of the -// given (region) metadata. NOTE: the searcher is constructed once per region -// while the raw i64 bounds were written in each file's unit; files written -// before/after a time index unit widening would need a per-file unit before -// this comparison is safe. See the note on `timestamp_values` in the writer. -fn time_range_filters( - metadata: &RegionMetadataRef, - time_range: Option, -) -> Result<(bool, Vec)> { - let Some(time_range) = time_range else { - return Ok((false, Vec::new())); - }; - if time_range.is_empty() { - return Ok((true, Vec::new())); - } +fn filters_for_schema( + schema: &SchemaRef, + filters: &[(Expr, SimpleFilterEvaluator)], +) -> (Predicate, Vec) { + let (exprs, filters): (Vec<_>, Vec<_>) = filters + .iter() + .filter(|(_, filter)| schema.field_with_name(filter.column_name()).is_ok()) + .cloned() + .unzip(); + (Predicate::new(exprs), filters) +} - let time_index = metadata.time_index_column(); - ensure!( - time_index.semantic_type == SemanticType::Timestamp, - InvalidMetaSnafu { - reason: "series index metadata has no timestamp time index", - } - ); - let timestamp_type = - time_index - .column_schema - .data_type - .as_timestamp() - .context(InvalidMetaSnafu { - reason: "series index time index is not a timestamp", - })?; - let unit = timestamp_type.unit(); +// Builds `__series_min_ts`/`__series_max_ts` predicates in `unit`, the unit +// recorded in the file being searched: its raw i64 bounds were written in that +// unit even if the region's time index has since been widened. The searcher +// handles empty ranges itself, so `time_range`, if present, is non-empty. +fn time_range_exprs(unit: TimeUnit, time_range: Option<&TimestampRange>) -> Vec { + let Some(time_range) = time_range else { + return Vec::new(); + }; let mut exprs = Vec::with_capacity(2); // A series overlaps [start, end) only if its maximum is at least start. // Round start up so a series ending before an unaligned start is pruned. @@ -188,20 +180,24 @@ fn time_range_filters( .start() .and_then(|start| start.convert_to_ceil(unit)) { - exprs.push(col(MAX_TS_COLUMN).gt_eq(lit(start.value()))); + exprs.push( + col(MAX_TS_COLUMN).gt_eq(lit(timestamp_to_scalar_value(unit, Some(start.value())))), + ); } // A series overlaps [start, end) only if its minimum is less than end. // Round the exclusive end up to avoid pruning the containing unit interval. if let Some(end) = time_range.end().and_then(|end| end.convert_to_ceil(unit)) { - exprs.push(col(MIN_TS_COLUMN).lt(lit(end.value()))); + exprs.push(col(MIN_TS_COLUMN).lt(lit(timestamp_to_scalar_value(unit, Some(end.value()))))); } - Ok((false, exprs)) + exprs } -fn validate_index_schema(schema: &SchemaRef) -> Result<()> { +/// Validates an index file's schema and returns the time unit of its +/// `__series_min_ts`/`__series_max_ts` columns. They are native +/// `Timestamp(unit)` columns, so the unit is part of the datatype and each +/// file is interpreted in the unit it was written with. +fn validate_index_schema(schema: &SchemaRef) -> Result { for (name, data_type) in [ - (MIN_TS_COLUMN, DataType::Int64), - (MAX_TS_COLUMN, DataType::Int64), (ROW_COUNT_COLUMN, DataType::UInt64), (TABLE_ID_COLUMN, DataType::UInt32), (TSID_COLUMN, DataType::UInt64), @@ -222,13 +218,43 @@ fn validate_index_schema(schema: &SchemaRef) -> Result<()> { } ); } - Ok(()) + let unit = |name: &str| { + let field = schema + .field_with_name(name) + .ok() + .context(InvalidRecordBatchSnafu { + reason: format!("series index is missing internal column {name}"), + })?; + ensure!( + !field.is_nullable(), + InvalidRecordBatchSnafu { + reason: format!("series index column {name} must be non-nullable"), + } + ); + match field.data_type() { + DataType::Timestamp(unit, _) => Ok(unit.into()), + data_type => InvalidRecordBatchSnafu { + reason: format!( + "series index column {name} must be a Timestamp, got {data_type:?}" + ), + } + .fail(), + } + }; + let min_unit = unit(MIN_TS_COLUMN)?; + let max_unit = unit(MAX_TS_COLUMN)?; + ensure!( + min_unit == max_unit, + InvalidRecordBatchSnafu { + reason: format!( + "series index columns {MIN_TS_COLUMN} and {MAX_TS_COLUMN} have different time units" + ), + } + ); + Ok(min_unit) } -fn column<'a>( - batch: &'a datatypes::arrow::record_batch::RecordBatch, - name: &str, -) -> Result<&'a ArrayRef> { +fn column<'a>(batch: &'a RecordBatch, name: &str) -> Result<&'a ArrayRef> { let index = batch .schema() .index_of(name) @@ -243,9 +269,13 @@ fn column<'a>( mod tests { use std::sync::Arc; + use api::v1::SemanticType; use datafusion_expr::{col, lit}; - use datatypes::arrow::array::{BinaryArray, TimestampMillisecondArray, UInt8Array}; - use datatypes::arrow::datatypes::{Field, Schema}; + use datatypes::arrow::array::{ + BinaryArray, TimestampMicrosecondArray, TimestampMillisecondArray, + TimestampNanosecondArray, TimestampSecondArray, UInt8Array, + }; + use datatypes::arrow::datatypes::{Field, Schema, TimeUnit as ArrowTimeUnit}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; @@ -262,13 +292,27 @@ mod tests { ObjectStore::new(Memory::default()).unwrap() } - fn flat_batch(primary_keys: &[Vec], timestamps: &[i64]) -> RecordBatch { + fn flat_batch_with_time_unit( + primary_keys: &[Vec], + timestamps: &[i64], + unit: ArrowTimeUnit, + ) -> RecordBatch { + let ts_column = match unit { + ArrowTimeUnit::Second => { + Arc::new(TimestampSecondArray::from(timestamps.to_vec())) as ArrayRef + } + ArrowTimeUnit::Millisecond => { + Arc::new(TimestampMillisecondArray::from(timestamps.to_vec())) as ArrayRef + } + ArrowTimeUnit::Microsecond => { + Arc::new(TimestampMicrosecondArray::from(timestamps.to_vec())) as ArrayRef + } + ArrowTimeUnit::Nanosecond => { + Arc::new(TimestampNanosecondArray::from(timestamps.to_vec())) as ArrayRef + } + }; let schema = Arc::new(Schema::new(vec![ - Field::new( - "ts", - DataType::Timestamp(datatypes::arrow::datatypes::TimeUnit::Millisecond, None), - false, - ), + Field::new("ts", DataType::Timestamp(unit, None), false), Field::new("__primary_key", DataType::Binary, false), Field::new("__sequence", DataType::UInt64, false), Field::new("__op_type", DataType::UInt8, false), @@ -276,7 +320,7 @@ mod tests { RecordBatch::try_new( schema, vec![ - Arc::new(TimestampMillisecondArray::from(timestamps.to_vec())), + ts_column, Arc::new(BinaryArray::from_iter_values( primary_keys.iter().map(Vec::as_slice), )), @@ -293,6 +337,25 @@ mod tests { path: &str, rows: &[(u32, u64, &str, &str, i64)], row_group_size: usize, + ) { + write_index_with_time_unit( + metadata, + object_store, + path, + rows, + row_group_size, + ArrowTimeUnit::Millisecond, + ) + .await + } + + async fn write_index_with_time_unit( + metadata: RegionMetadataRef, + object_store: ObjectStore, + path: &str, + rows: &[(u32, u64, &str, &str, i64)], + row_group_size: usize, + unit: ArrowTimeUnit, ) { let primary_keys = rows .iter() @@ -311,7 +374,7 @@ mod tests { .await .unwrap(); writer - .write(&flat_batch(&primary_keys, ×tamps)) + .write(&flat_batch_with_time_unit(&primary_keys, ×tamps, unit)) .await .unwrap(); writer.finish().await.unwrap(); @@ -361,11 +424,13 @@ mod tests { let searcher = SeriesIndexSearcher::try_new( metadata.clone(), object_store.clone(), + path, Some(&predicate), Some(time_range), ) + .await .unwrap(); - let ids = collect_ids(searcher.search(path).await.unwrap()).await; + let ids = collect_ids(searcher.search().unwrap()).await; assert_eq!( ids, vec![MetricSeriesId { @@ -384,11 +449,13 @@ mod tests { let searcher = SeriesIndexSearcher::try_new( metadata.clone(), object_store.clone(), + path, None, Some(time_range), ) + .await .unwrap(); - let ids = collect_ids(searcher.search(path).await.unwrap()).await; + let ids = collect_ids(searcher.search().unwrap()).await; assert_eq!( ids, vec![MetricSeriesId { @@ -405,8 +472,10 @@ mod tests { ) .unwrap(); let searcher = - SeriesIndexSearcher::try_new(metadata, object_store, None, Some(time_range)).unwrap(); - let ids = collect_ids(searcher.search(path).await.unwrap()).await; + SeriesIndexSearcher::try_new(metadata, object_store, path, None, Some(time_range)) + .await + .unwrap(); + let ids = collect_ids(searcher.search().unwrap()).await; assert_eq!( ids, vec![MetricSeriesId { @@ -449,10 +518,16 @@ mod tests { let predicate = Predicate::new(vec![col("tag_0").eq(lit("a")), col("tag_2").eq(lit("new"))]); - let searcher = - SeriesIndexSearcher::try_new(current_metadata, object_store, Some(&predicate), None) - .unwrap(); - let ids = collect_ids(searcher.search(path).await.unwrap()).await; + let searcher = SeriesIndexSearcher::try_new( + current_metadata, + object_store, + path, + Some(&predicate), + None, + ) + .await + .unwrap(); + let ids = collect_ids(searcher.search().unwrap()).await; assert_eq!( ids, vec![ @@ -468,6 +543,215 @@ mod tests { ); } + #[tokio::test] + async fn search_uses_file_time_unit_after_time_index_widen() { + // The index is written while the region's time index is milliseconds. + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let object_store = object_store(); + let path = "widen.parquet"; + write_index( + metadata.clone(), + object_store.clone(), + path, + &[ + (1, 10, "a", "x", 10), + (1, 20, "b", "x", 20), + (1, 30, "c", "x", 30), + ], + 2, + ) + .await; + + // The region's time index is then widened to microseconds. + let mut widened = (*metadata).clone(); + for column in &mut widened.column_metadatas { + if column.column_schema.name == "ts" { + column.column_schema.data_type = ConcreteDataType::timestamp_microsecond_datatype(); + } + } + let widened = Arc::new(widened); + + // A query range of [10.001ms, 25ms) expressed in microseconds must be + // compared against the file's millisecond bounds (ceil: [11ms, + // 25ms)): only the 20ms series intersects. Interpreting the file in + // microseconds instead would match nothing. + let time_range = TimestampRange::new( + common_time::Timestamp::new_microsecond(10_001), + common_time::Timestamp::new_microsecond(25_000), + ) + .unwrap(); + let searcher = SeriesIndexSearcher::try_new( + widened, + object_store.clone(), + path, + None, + Some(time_range), + ) + .await + .unwrap(); + let ids = collect_ids(searcher.search().unwrap()).await; + assert_eq!( + ids, + vec![MetricSeriesId { + table_id: 1, + tsid: 20 + }] + ); + + // A millisecond-unit range behaves identically to the pre-widen + // searcher: [20ms, 30ms) keeps only the 20ms series. + let time_range = TimestampRange::new( + common_time::Timestamp::new_millisecond(20), + common_time::Timestamp::new_millisecond(30), + ) + .unwrap(); + let searcher = + SeriesIndexSearcher::try_new(metadata, object_store, path, None, Some(time_range)) + .await + .unwrap(); + let ids = collect_ids(searcher.search().unwrap()).await; + assert_eq!( + ids, + vec![MetricSeriesId { + table_id: 1, + tsid: 20 + }] + ); + } + + #[tokio::test] + async fn search_reads_each_file_in_its_recorded_unit() { + // The first file is written while the region's time index is + // milliseconds; its series sit at 10ms and 20ms. + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let object_store = object_store(); + write_index( + metadata.clone(), + object_store.clone(), + "old_ms.parquet", + &[(1, 10, "a", "x", 10), (1, 20, "b", "x", 20)], + 2, + ) + .await; + + // The region's time index is then widened to microseconds and a + // second file is written; its series sit at 15_000µs and 15µs. + let mut widened = (*metadata).clone(); + for column in &mut widened.column_metadatas { + if column.column_schema.name == "ts" { + column.column_schema.data_type = ConcreteDataType::timestamp_microsecond_datatype(); + } + } + let widened = Arc::new(widened); + write_index_with_time_unit( + widened.clone(), + object_store.clone(), + "new_us.parquet", + &[(1, 30, "c", "x", 15_000), (1, 40, "d", "x", 15)], + 2, + ArrowTimeUnit::Microsecond, + ) + .await; + + // Each file is searched by its own searcher, which interprets it in + // the unit recorded in its schema. For [10.5ms, 25ms): + // - the millisecond file (ceil: [11ms, 25ms)) keeps only the 20ms + // series; reading it in microseconds would match nothing; + // - the microsecond file keeps only the 15_000µs series; reading it + // in milliseconds would also keep the 15µs series. + let time_range = TimestampRange::new( + common_time::Timestamp::new_microsecond(10_500), + common_time::Timestamp::new_microsecond(25_000), + ) + .unwrap(); + let searcher = SeriesIndexSearcher::try_new( + widened.clone(), + object_store.clone(), + "old_ms.parquet", + None, + Some(time_range), + ) + .await + .unwrap(); + let ids = collect_ids(searcher.search().unwrap()).await; + assert_eq!( + ids, + vec![MetricSeriesId { + table_id: 1, + tsid: 20 + }] + ); + let searcher = SeriesIndexSearcher::try_new( + widened, + object_store, + "new_us.parquet", + None, + Some(time_range), + ) + .await + .unwrap(); + let ids = collect_ids(searcher.search().unwrap()).await; + assert_eq!( + ids, + vec![MetricSeriesId { + table_id: 1, + tsid: 30 + }] + ); + } + + fn ts_field(name: &str, unit: Option) -> Field { + let data_type = match unit { + Some(unit) => DataType::Timestamp(unit, None), + None => DataType::Int64, + }; + Field::new(name, data_type, false) + } + + fn index_file_schema( + min_unit: Option, + max_unit: Option, + ) -> SchemaRef { + Arc::new(Schema::new(vec![ + ts_field(MIN_TS_COLUMN, min_unit), + ts_field(MAX_TS_COLUMN, max_unit), + Field::new(ROW_COUNT_COLUMN, DataType::UInt64, false), + Field::new(TABLE_ID_COLUMN, DataType::UInt32, false), + Field::new(TSID_COLUMN, DataType::UInt64, false), + ])) + } + + #[test] + fn validate_index_schema_rejects_unusable_columns() { + // Min/max columns that are not Timestamps are rejected. + let err = validate_index_schema(&index_file_schema(None, None)) + .unwrap_err() + .to_string(); + assert!(err.contains("must be a Timestamp, got Int64"), "{err}"); + + // The min and max columns must agree on the unit. + let err = validate_index_schema(&index_file_schema( + Some(ArrowTimeUnit::Millisecond), + Some(ArrowTimeUnit::Microsecond), + )) + .unwrap_err() + .to_string(); + assert!(err.contains("have different time units"), "{err}"); + + assert_eq!( + TimeUnit::Nanosecond, + validate_index_schema(&index_file_schema( + Some(ArrowTimeUnit::Nanosecond), + Some(ArrowTimeUnit::Nanosecond), + )) + .unwrap() + ); + } + #[tokio::test] async fn search_streams_fixed_size_batches() { let metadata = Arc::new(sst_region_metadata_with_encoding( @@ -486,10 +770,12 @@ mod tests { ) .await; - let searcher = SeriesIndexSearcher::try_new(metadata, object_store, None, None).unwrap(); + let searcher = + SeriesIndexSearcher::try_new(metadata, object_store, "batching.parquet", None, None) + .await + .unwrap(); let batches = searcher - .search("batching.parquet") - .await + .search() .unwrap() .try_collect::>() .await @@ -534,31 +820,48 @@ mod tests { ) .await; - let predicate = Predicate::new(vec![col("tag_0").eq(lit("m"))]); - let searcher = SeriesIndexSearcher::try_new( - metadata.clone(), - object_store.clone(), - Some(&predicate), - None, - ) - .unwrap(); let reader = ParquetIndexReader::open(object_store.clone(), path) .await .unwrap(); - let (pruning_predicate, _) = searcher.filters_for_schema(reader.schema()); + let unit = validate_index_schema(reader.schema()).unwrap(); + + // Tag predicates prune row groups through the tag-column statistics. + let predicate = Predicate::new(vec![col("tag_0").eq(lit("m"))]); + let mut filters = simple_tag_filters(&metadata, None, Some(&predicate)); + let (pruning_predicate, _) = filters_for_schema(reader.schema(), &filters); assert_eq!(reader.row_groups_to_read(&pruning_predicate), vec![1]); + // Time-range predicates prune row groups in the file's own unit: + // [2ms, 4ms) keeps only the row group holding the 2ms and 3ms series. + let time_range = TimestampRange::new( + common_time::Timestamp::new_millisecond(2), + common_time::Timestamp::new_millisecond(4), + ) + .unwrap(); + for expr in time_range_exprs(unit, Some(&time_range)) { + let filter = SimpleFilterEvaluator::try_new(&expr) + .context(UnexpectedSnafu { + reason: "failed to build an internal series-index time filter", + }) + .unwrap(); + filters.push((expr, filter)); + } + let (pruning_predicate, _) = filters_for_schema(reader.schema(), &filters); + assert_eq!(reader.row_groups_to_read(&pruning_predicate), vec![1]); + + // An empty time range yields no series without opening the file. let empty = SeriesIndexSearcher::try_new( metadata, object_store, + "does-not-need-to-exist.parquet", None, Some(TimestampRange::empty()), ) + .await .unwrap(); assert!( empty - .search("does-not-need-to-exist.parquet") - .await + .search() .unwrap() .try_collect::>() .await diff --git a/src/mito2/src/series_index/writer.rs b/src/mito2/src/series_index/writer.rs index db8e105ed4..ec1ec6df48 100644 --- a/src/mito2/src/series_index/writer.rs +++ b/src/mito2/src/series_index/writer.rs @@ -16,9 +16,11 @@ use std::cmp::Ordering; use std::sync::Arc; use std::time::{Duration, Instant}; +use common_time::timestamp::TimeUnit; use datatypes::arrow::array::{ - Array, ArrayRef, BinaryArray, DictionaryArray, Int64Array, StringArray, UInt32Array, - UInt64Array, + Array, ArrayRef, BinaryArray, DictionaryArray, Int64Array, StringArray, + TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, + TimestampSecondArray, UInt32Array, UInt64Array, }; use datatypes::arrow::datatypes::{DataType, Field, Schema, SchemaRef, UInt32Type}; use datatypes::arrow::record_batch::RecordBatch; @@ -112,6 +114,9 @@ pub struct SeriesIndexWriter { codec: Arc, tag_columns: Vec<(ColumnId, String)>, schema: SchemaRef, + /// Unit of the min/max ts Timestamp columns, i.e. the region's time + /// index unit at writer creation. + time_unit: TimeUnit, writer: ParquetIndexWriter, current_primary_key: Option>, current_row: Option, @@ -139,6 +144,7 @@ impl SeriesIndexWriter { reason: "series index row group size must be greater than zero", } ); + let time_unit = time_index_unit(&metadata)?; let schema = series_index_schema(&metadata)?; let tag_columns = tag_columns(&metadata); let writer = ParquetIndexWriter::try_new( @@ -156,6 +162,7 @@ impl SeriesIndexWriter { codec, tag_columns, schema, + time_unit, writer, current_primary_key: None, current_row: None, @@ -239,7 +246,7 @@ impl SeriesIndexWriter { let pk_idx = primary_key_column_index(batch.num_columns()); let ts_idx = time_index_column_index(batch.num_columns()); let primary_keys = batch.column(pk_idx); - let timestamps = timestamp_values(batch.column(ts_idx))?; + let timestamps = timestamp_values(batch.column(ts_idx), self.time_unit)?; ensure!( primary_keys.len() == batch.num_rows() && timestamps.len() == batch.num_rows(), InvalidRecordBatchSnafu { @@ -394,7 +401,7 @@ impl SeriesIndexWriter { if self.buffered_rows.is_empty() { return Ok(()); } - let batch = rows_to_batch(&self.schema, &self.buffered_rows)?; + let batch = rows_to_batch(&self.schema, &self.buffered_rows, self.time_unit)?; let start = Instant::now(); let result = self.writer.write(&batch).await; self.metrics.write_elapsed += start.elapsed(); @@ -432,9 +439,13 @@ impl SeriesIndexWriter { /// Returns the Arrow schema of a series index. pub fn series_index_schema(metadata: &RegionMetadataRef) -> Result { validate_metadata(metadata)?; + // Native Timestamp columns carry the time index unit in their datatype, + // so each file is interpreted in the unit it was written with even after + // the region's time index unit has been widened. + let ts_type = DataType::Timestamp(time_index_unit(metadata)?.into(), None); let mut fields = vec![ - Field::new(MIN_TS_COLUMN, DataType::Int64, false), - Field::new(MAX_TS_COLUMN, DataType::Int64, false), + Field::new(MIN_TS_COLUMN, ts_type.clone(), false), + Field::new(MAX_TS_COLUMN, ts_type, false), Field::new(ROW_COUNT_COLUMN, DataType::UInt64, false), Field::new(TABLE_ID_COLUMN, DataType::UInt32, false), Field::new(TSID_COLUMN, DataType::UInt64, false), @@ -447,6 +458,21 @@ pub fn series_index_schema(metadata: &RegionMetadataRef) -> Result { Ok(Arc::new(Schema::new(fields))) } +/// Returns the unit of the region's timestamp time index; the writer stamps +/// it into index files so a searcher interprets each file in the unit it was +/// written with (the region's time index unit may have been widened since). +fn time_index_unit(metadata: &RegionMetadataRef) -> Result { + Ok(metadata + .time_index_column() + .column_schema + .data_type + .as_timestamp() + .context(InvalidMetaSnafu { + reason: "series index requires a timestamp time index", + })? + .unit()) +} + fn validate_metadata(metadata: &RegionMetadataRef) -> Result<()> { for column in &metadata.column_metadatas { ensure!( @@ -522,32 +548,35 @@ fn is_reserved_column(column_id: ColumnId) -> bool { column_id == ReservedColumnId::table_id() || column_id == ReservedColumnId::tsid() } -/// Extracts raw i64 timestamp values for the `__series_min_ts`/`__series_max_ts` -/// columns. NOTE: the unit is dropped; the values must always be interpreted in -/// the unit of the region metadata the index is written with. Once this index -/// is wired into scans, files written before/after a time index unit widening -/// would carry mixed units — the index schema or the searcher must record the -/// unit per file, or the index must be rebuilt on such an alter. -fn timestamp_values(array: &ArrayRef) -> Result { - let timestamps = if let Some(array) = array.as_any().downcast_ref::() { - array.clone() - } else { - timestamp_array_to_primitive(array) - .map(|(array, _)| array) - .with_context(|| InvalidRecordBatchSnafu { - reason: format!( - "series index requires an Int64 or timestamp time index, got {:?}", - array.data_type() - ), - })? - }; +/// Extracts the time index column's raw values in the writer's `unit`. +/// A timestamp array carrying a different unit is an invariant violation: +/// the alter path flushes memtables before widening the region's time index +/// unit, so a writer never receives batches in the region's previous unit. +/// Reject the mismatch rather than reinterpreting or rescaling the values. +fn timestamp_values(array: &ArrayRef, unit: TimeUnit) -> Result { ensure!( - timestamps.null_count() == 0, + array.null_count() == 0, InvalidRecordBatchSnafu { reason: "series index input contains null timestamps", } ); - Ok(timestamps) + let (values, array_unit) = + timestamp_array_to_primitive(array).with_context(|| InvalidRecordBatchSnafu { + reason: format!( + "series index requires a timestamp time index column, got {:?}", + array.data_type() + ), + })?; + let array_unit: TimeUnit = array_unit.into(); + ensure!( + array_unit == unit, + InvalidRecordBatchSnafu { + reason: format!( + "series index input time index unit {array_unit:?} does not match the index unit {unit:?}" + ), + } + ); + Ok(values) } // TODO(yingwen): Bench and optimize the performance if this is costly. @@ -606,14 +635,27 @@ fn decode_primary_key( }) } -fn rows_to_batch(schema: &SchemaRef, rows: &[SeriesIndexRow]) -> Result { +/// Builds a timestamp array in `unit` from raw i64 values. +fn ts_array(timestamps: impl Iterator, unit: TimeUnit) -> ArrayRef { + match unit { + TimeUnit::Second => Arc::new(TimestampSecondArray::from_iter_values(timestamps)), + TimeUnit::Millisecond => Arc::new(TimestampMillisecondArray::from_iter_values(timestamps)), + TimeUnit::Microsecond => Arc::new(TimestampMicrosecondArray::from_iter_values(timestamps)), + TimeUnit::Nanosecond => Arc::new(TimestampNanosecondArray::from_iter_values(timestamps)), + } +} + +/// Builds a batch in the index `schema` from aggregated rows. The rows carry +/// raw timestamps in `unit` (guaranteed by `timestamp_values`), which +/// `series_index_schema` stamps into the min/max ts columns. +fn rows_to_batch( + schema: &SchemaRef, + rows: &[SeriesIndexRow], + unit: TimeUnit, +) -> Result { let mut arrays: Vec = vec![ - Arc::new(Int64Array::from_iter_values( - rows.iter().map(|row| row.min_ts), - )), - Arc::new(Int64Array::from_iter_values( - rows.iter().map(|row| row.max_ts), - )), + ts_array(rows.iter().map(|row| row.min_ts), unit), + ts_array(rows.iter().map(|row| row.max_ts), unit), Arc::new(UInt64Array::from_iter_values( rows.iter().map(|row| row.row_count), )), @@ -636,11 +678,11 @@ fn rows_to_batch(schema: &SchemaRef, rows: &[SeriesIndexRow]) -> Result = vec![ - Arc::new(Int64Array::from(vec![1, 2])), - Arc::new(TimestampSecondArray::from(vec![1, 2])), - Arc::new(TimestampMillisecondArray::from(vec![1, 2])), - Arc::new(TimestampMicrosecondArray::from(vec![1, 2])), - Arc::new(TimestampNanosecondArray::from(vec![1, 2])), - ]; - for array in arrays { - assert_eq!(timestamp_values(&array).unwrap().values().as_ref(), &[1, 2]); - } + // Values already in the writer's unit pass through unchanged. + let array: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![1, 2])); + assert_eq!( + timestamp_values(&array, TimeUnit::Millisecond).unwrap(), + Int64Array::from(vec![1, 2]) + ); + // A timestamp array carrying a different unit is rejected instead of + // being silently reinterpreted or rescaled. + let array: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![1, 2])); + let error = timestamp_values(&array, TimeUnit::Microsecond).unwrap_err(); + assert!( + error.to_string().contains("does not match the index unit"), + "{error}" + ); + + // Plain Int64 columns carry no unit to validate against and are + // rejected. + let int64: ArrayRef = Arc::new(Int64Array::from(vec![1, 2])); + let error = timestamp_values(&int64, TimeUnit::Millisecond).unwrap_err(); + assert!( + error + .to_string() + .contains("requires a timestamp time index column"), + "{error}" + ); let nulls: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![Some(1), None])); - let error = timestamp_values(&nulls).unwrap_err(); + let error = timestamp_values(&nulls, TimeUnit::Millisecond).unwrap_err(); assert!(error.to_string().contains("null timestamps"), "{error}"); let unsupported: ArrayRef = Arc::new(UInt8Array::from(vec![1, 2])); - let error = timestamp_values(&unsupported).unwrap_err(); + let error = timestamp_values(&unsupported, TimeUnit::Millisecond).unwrap_err(); assert!( - error.to_string().contains("requires an Int64 or timestamp"), + error + .to_string() + .contains("requires a timestamp time index column"), "{error}" ); } @@ -869,17 +937,17 @@ mod tests { batch .column(0) .as_any() - .downcast_ref::() + .downcast_ref::() .unwrap(), - &Int64Array::from(vec![70, 200]) + &TimestampMillisecondArray::from(vec![70, 200]) ); assert_eq!( batch .column(1) .as_any() - .downcast_ref::() + .downcast_ref::() .unwrap(), - &Int64Array::from(vec![130, 230]) + &TimestampMillisecondArray::from(vec![130, 230]) ); assert_eq!( batch diff --git a/src/mito2/src/sst/parquet/index_reader.rs b/src/mito2/src/sst/parquet/index_reader.rs index bd894f207a..0cf221986c 100644 --- a/src/mito2/src/sst/parquet/index_reader.rs +++ b/src/mito2/src/sst/parquet/index_reader.rs @@ -185,6 +185,15 @@ impl IndexRowGroupPruningStats<'_> { fn column_values(&self, column: &Column, is_min: bool) -> Option { let column_index = self.schema.index_of(&column.name).ok()?; let data_type = self.schema.field(column_index).data_type(); - column_values_by_type(self.row_groups, data_type, column_index, is_min) + let values = column_values_by_type(self.row_groups, data_type, column_index, is_min)?; + if values.data_type() == data_type { + Some(values) + } else { + // Parquet timestamp statistics surface as raw Int64; reinterpret + // them in the column's Timestamp type so pruning compares + // like-typed values. A cast failure yields no stats and keeps the + // row group (conservative). + datatypes::arrow::compute::cast(&values, data_type).ok() + } } }