1use std::collections::BTreeMap;
18use std::sync::Arc;
19
20use api::v1::SemanticType;
21use common_base::readable_size::ReadableSize;
22use datatypes::json::JsonSettings;
23use datatypes::schema::ColumnEncoding;
24use parquet::basic::{Compression, Encoding, ZstdLevel};
25use parquet::file::metadata::ParquetMetaData;
26use parquet::file::properties::{DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT, WriterPropertiesBuilder};
27use parquet::schema::types::ColumnPath;
28use snafu::ResultExt;
29use store_api::metadata::RegionMetadataRef;
30use store_api::mito_engine_options::FloatFieldEncoding;
31use store_api::storage::consts::{OP_TYPE_COLUMN_NAME, SEQUENCE_COLUMN_NAME};
32use store_api::storage::{ColumnId, FileId};
33
34use crate::error::{CompressionOptionsSnafu, Result, WriteParquetSnafu};
35use crate::sst::DEFAULT_WRITE_BUFFER_SIZE;
36use crate::sst::file::FileTimeRange;
37use crate::sst::index::IndexOutput;
38use crate::sst::parquet::parallel_writer::EncodeRuntime;
39
40pub mod file_range;
41pub mod flat_format;
42pub mod format;
43pub(crate) mod helper;
44pub(crate) mod index_reader;
45pub(crate) mod index_writer;
46pub(crate) mod json_align;
47pub mod metadata;
48pub mod parallel_writer;
49pub mod prefilter;
50pub mod push_decoder;
51pub mod read_columns;
52pub mod reader;
53pub mod row_group;
54pub mod row_selection;
55pub(crate) mod stats;
56pub mod writer;
57
58pub const PARQUET_METADATA_KEY: &str = "greptime:metadata";
60
61pub(crate) const DEFAULT_READ_BATCH_SIZE: usize = 8 * 1024;
67
68pub(crate) type Json2RewriteTargets = Arc<BTreeMap<ColumnId, Json2TargetLayout>>;
70
71#[derive(Debug, Clone, PartialEq, Eq)]
73pub(crate) struct Json2TargetLayout {
74 pub(crate) extension_metadata: String,
76 pub(crate) target_layout: JsonSettings,
78}
79
80pub const DEFAULT_ROW_GROUP_SIZE: usize = 100 * 1024;
86
87pub(crate) const COLUMN_INDEX_TRUNCATE_LENGTH: Option<usize> =
97 parquet::file::properties::DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH;
98
99pub(crate) fn apply_internal_column_encoding(
101 builder: WriterPropertiesBuilder,
102 metadata: &RegionMetadataRef,
103) -> WriterPropertiesBuilder {
104 let ts_col = ColumnPath::new(vec![
105 metadata.time_index_column().column_schema.name.clone(),
106 ]);
107 let seq_col = ColumnPath::new(vec![SEQUENCE_COLUMN_NAME.to_string()]);
108 let op_type_col = ColumnPath::new(vec![OP_TYPE_COLUMN_NAME.to_string()]);
109 builder
110 .set_column_encoding(seq_col.clone(), Encoding::DELTA_BINARY_PACKED)
111 .set_column_dictionary_enabled(seq_col, false)
112 .set_column_encoding(ts_col.clone(), Encoding::DELTA_BINARY_PACKED)
113 .set_column_dictionary_enabled(ts_col, false)
114 .set_column_compression(op_type_col, Compression::UNCOMPRESSED)
115}
116
117pub(crate) fn apply_column_compression(
119 mut builder: WriterPropertiesBuilder,
120 metadata: &RegionMetadataRef,
121) -> Result<WriterPropertiesBuilder> {
122 for column in &metadata.column_metadatas {
123 let Some(options) =
124 column
125 .column_schema
126 .compression_options()
127 .context(CompressionOptionsSnafu {
128 column_name: &column.column_schema.name,
129 })?
130 else {
131 continue;
132 };
133 let path = ColumnPath::new(vec![column.column_schema.name.clone()]);
134 if let Some(encoding) = options.encoding {
135 let encoding = match encoding {
136 ColumnEncoding::Plain => Encoding::PLAIN,
137 ColumnEncoding::DeltaBinaryPacked => Encoding::DELTA_BINARY_PACKED,
138 };
139 builder = builder
140 .set_column_dictionary_enabled(path.clone(), false)
141 .set_column_encoding(path.clone(), encoding);
142 }
143 if let Some(level) = options.level {
144 let level = ZstdLevel::try_new(level).context(WriteParquetSnafu)?;
145 builder = builder.set_column_compression(path.clone(), Compression::ZSTD(level));
146 }
147 if let Some(page_size) = options.page_size {
148 let page_size = page_size.as_bytes() as usize;
149 builder = builder.set_column_data_page_size_limit(path.clone(), page_size);
150 if page_size < DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT {
154 builder = builder.set_column_dictionary_page_size_limit(path, page_size);
155 }
156 }
157 }
158 Ok(builder)
159}
160
161pub(crate) fn apply_float_field_encoding(
163 mut builder: WriterPropertiesBuilder,
164 metadata: &RegionMetadataRef,
165 encoding: FloatFieldEncoding,
166) -> WriterPropertiesBuilder {
167 if encoding == FloatFieldEncoding::ByteStreamSplit {
168 for column in &metadata.column_metadatas {
169 if column.semantic_type == SemanticType::Field
170 && column.column_schema.data_type.is_float()
171 {
172 let path = ColumnPath::new(vec![column.column_schema.name.clone()]);
173 builder = builder
174 .set_column_encoding(path.clone(), parquet::basic::Encoding::BYTE_STREAM_SPLIT)
175 .set_column_dictionary_enabled(path, false);
176 }
177 }
178 }
179 builder
180}
181
182#[derive(Debug, Clone)]
184pub struct WriteOptions {
185 pub write_buffer_size: ReadableSize,
187 pub row_group_size: usize,
189 pub max_file_size: Option<usize>,
193 pub float_field_encoding: FloatFieldEncoding,
195 pub encode_parallelism: usize,
198 pub encode_runtime: EncodeRuntime,
200}
201
202impl Default for WriteOptions {
203 fn default() -> Self {
204 WriteOptions {
205 write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE,
206 row_group_size: DEFAULT_ROW_GROUP_SIZE,
207 max_file_size: None,
208 float_field_encoding: FloatFieldEncoding::default(),
209 encode_parallelism: 1,
210 encode_runtime: EncodeRuntime::default(),
211 }
212 }
213}
214
215#[derive(Debug, Default)]
217pub struct SstInfo {
218 pub file_id: FileId,
220 pub time_range: FileTimeRange,
223 pub file_size: u64,
225 pub max_row_group_uncompressed_size: u64,
227 pub num_rows: usize,
229 pub num_row_groups: u64,
231 pub file_metadata: Option<Arc<ParquetMetaData>>,
233 pub index_metadata: IndexOutput,
235 pub num_series: u64,
237}
238
239#[cfg(test)]
240mod tests {
241 use std::collections::HashSet;
242 use std::sync::Arc;
243
244 use api::v1::{OpType, SemanticType};
245 use bytes::Bytes;
246 use common_base::readable_size::ReadableSize;
247 use common_function::function::FunctionRef;
248 use common_function::function_factory::ScalarFunctionFactory;
249 use common_function::scalars::matches::MatchesFunction;
250 use common_function::scalars::matches_term::MatchesTermFunction;
251 use common_time::Timestamp;
252 use datafusion_common::{Column, ScalarValue};
253 use datafusion_expr::expr::ScalarFunction;
254 use datafusion_expr::{BinaryExpr, Expr, Literal, Operator, col, lit};
255 use datatypes::arrow;
256 use datatypes::arrow::array::{
257 Array, ArrayRef, AsArray, BinaryDictionaryBuilder, Float32Array, Float64Array, Int32Array,
258 RecordBatch, StringArray, StringDictionaryBuilder, TimestampMillisecondArray, UInt8Array,
259 UInt64Array,
260 };
261 use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit, UInt32Type};
262 use datatypes::arrow::util::pretty::pretty_format_batches;
263 use datatypes::prelude::ConcreteDataType;
264 use datatypes::schema::{
265 CompressionOptions, FulltextAnalyzer, FulltextBackend, FulltextOptions,
266 };
267 use object_store::ObjectStore;
268 use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
269 use parquet::arrow::{ArrowWriter, AsyncArrowWriter};
270 use parquet::basic::{Compression, Encoding, ZstdLevel};
271 use parquet::file::metadata::{KeyValue, PageIndexPolicy};
272 use parquet::file::properties::WriterProperties;
273 use parquet::schema::types::ColumnPath;
274 use store_api::codec::PrimaryKeyEncoding;
275 use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder};
276 use store_api::mito_engine_options::FloatFieldEncoding;
277 use store_api::region_request::PathType;
278 use store_api::storage::{ColumnSchema, RegionId};
279 use table::predicate::Predicate;
280 use tokio_util::compat::FuturesAsyncWriteCompatExt;
281
282 use super::*;
283 use crate::access_layer::{FilePathProvider, Metrics, RegionFilePathFactory, WriteType};
284 use crate::cache::index::result_cache::PredicateKey;
285 use crate::cache::test_util::assert_parquet_metadata_equal;
286 use crate::cache::{CacheManager, CacheStrategy};
287 use crate::config::IndexConfig;
288 use crate::read::FlatSource;
289 use crate::region::options::{IndexOptions, InvertedIndexOptions};
290 use crate::sst::file::{FileHandle, FileMeta, RegionFileId, RegionIndexId};
291 use crate::sst::file_purger::NoopFilePurger;
292 use crate::sst::index::bloom_filter::applier::BloomFilterIndexApplierBuilder;
293 use crate::sst::index::fulltext_index::applier::builder::FulltextIndexApplierBuilder;
294 use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder;
295 use crate::sst::index::{IndexBuildType, Indexer, IndexerBuilder, IndexerBuilderImpl};
296 use crate::sst::parquet::flat_format::FlatWriteFormat;
297 use crate::sst::parquet::metadata::extract_primary_key_range;
298 use crate::sst::parquet::reader::{ParquetReader, ParquetReaderBuilder, ReaderMetrics};
299 use crate::sst::parquet::row_selection::RowGroupSelection;
300 use crate::sst::parquet::writer::ParquetWriter;
301 use crate::sst::{
302 DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, location, to_flat_sst_arrow_schema,
303 };
304 use crate::test_util::TestEnv;
305 use crate::test_util::sst_util::{
306 WriteChunkRecorder, build_test_binary_test_region_metadata,
307 new_flat_source_from_record_batches, new_primary_key, new_record_batch_by_range,
308 new_record_batch_with_custom_sequence, new_sparse_primary_key, sst_file_handle,
309 sst_file_handle_with_file_id, sst_region_metadata, sst_region_metadata_with_encoding,
310 };
311
312 const FILE_DIR: &str = "/";
313 const REGION_ID: RegionId = RegionId::new(0, 0);
314
315 #[test]
316 fn test_float_field_encoding_properties_and_roundtrip() {
317 let mut metadata_builder = RegionMetadataBuilder::new(REGION_ID);
318 metadata_builder
319 .push_column_metadata(ColumnMetadata {
320 column_schema: ColumnSchema::new("f32", ConcreteDataType::float32_datatype(), true),
321 semantic_type: SemanticType::Field,
322 column_id: 0,
323 })
324 .push_column_metadata(ColumnMetadata {
325 column_schema: ColumnSchema::new("f64", ConcreteDataType::float64_datatype(), true),
326 semantic_type: SemanticType::Field,
327 column_id: 1,
328 })
329 .push_column_metadata(ColumnMetadata {
330 column_schema: ColumnSchema::new("tag", ConcreteDataType::float32_datatype(), true),
331 semantic_type: SemanticType::Tag,
332 column_id: 2,
333 })
334 .push_column_metadata(ColumnMetadata {
335 column_schema: ColumnSchema::new("i32", ConcreteDataType::int32_datatype(), true),
336 semantic_type: SemanticType::Field,
337 column_id: 3,
338 })
339 .push_column_metadata(ColumnMetadata {
340 column_schema: ColumnSchema::new(
341 "ts",
342 ConcreteDataType::timestamp_millisecond_datatype(),
343 false,
344 ),
345 semantic_type: SemanticType::Timestamp,
346 column_id: 4,
347 });
348 metadata_builder.primary_key(vec![2]);
349 let metadata = Arc::new(metadata_builder.build().unwrap());
350
351 let f32_values = [
352 Some(1.5_f32),
353 Some(2.0),
354 Some(0.0),
355 Some(-0.0),
356 Some(f32::INFINITY),
357 Some(f32::from_bits(0x7fc0_1234)),
358 None,
359 ];
360 let f64_values = [
361 Some(1.5_f64),
362 Some(2.0),
363 Some(0.0),
364 Some(-0.0),
365 Some(f64::NEG_INFINITY),
366 Some(f64::from_bits(0x7ff8_0000_0000_1234)),
367 None,
368 ];
369 let schema = Arc::new(Schema::new(vec![
370 Field::new("f32", DataType::Float32, true),
371 Field::new("f64", DataType::Float64, true),
372 Field::new("tag", DataType::Float32, true),
373 Field::new("i32", DataType::Int32, true),
374 Field::new(
375 "ts",
376 DataType::Timestamp(TimeUnit::Millisecond, None),
377 false,
378 ),
379 ]));
380 let batch = RecordBatch::try_new(
381 schema.clone(),
382 vec![
383 Arc::new(Float32Array::from(f32_values.to_vec())) as ArrayRef,
384 Arc::new(Float64Array::from(f64_values.to_vec())) as ArrayRef,
385 Arc::new(Float32Array::from(vec![
386 Some(1.0),
387 None,
388 Some(-0.0),
389 Some(2.0),
390 Some(3.0),
391 Some(4.0),
392 None,
393 ])),
394 Arc::new(Int32Array::from(vec![
395 Some(1),
396 None,
397 Some(3),
398 Some(4),
399 Some(5),
400 Some(6),
401 None,
402 ])),
403 Arc::new(TimestampMillisecondArray::from_iter_values(0..7)),
404 ],
405 )
406 .unwrap();
407
408 let path = |name: &str| ColumnPath::new(vec![name.to_string()]);
409 let bss = apply_float_field_encoding(
410 WriterProperties::builder(),
411 &metadata,
412 FloatFieldEncoding::ByteStreamSplit,
413 );
414 assert_eq!(
415 Some(Encoding::BYTE_STREAM_SPLIT),
416 bss.clone().build().encoding(&path("f32"))
417 );
418 assert_eq!(
419 Some(Encoding::BYTE_STREAM_SPLIT),
420 bss.clone().build().encoding(&path("f64"))
421 );
422 assert!(!bss.clone().build().dictionary_enabled(&path("f32")));
423 assert!(!bss.clone().build().dictionary_enabled(&path("f64")));
424 assert_eq!(None, bss.clone().build().encoding(&path("i32")));
425 assert!(bss.clone().build().dictionary_enabled(&path("tag")));
426
427 let default = apply_float_field_encoding(
428 WriterProperties::builder().set_encoding(Encoding::PLAIN),
429 &metadata,
430 FloatFieldEncoding::Default,
431 )
432 .build();
433 assert_eq!(Some(Encoding::PLAIN), default.encoding(&path("f32")));
434 assert!(default.dictionary_enabled(&path("f32")));
435 assert!(default.dictionary_enabled(&path("f64")));
436
437 let mut bytes = Vec::new();
438 let mut writer = ArrowWriter::try_new(&mut bytes, schema, Some(bss.build())).unwrap();
439 writer.write(&batch).unwrap();
440 let footer = writer.finish().unwrap();
441 drop(writer);
442 for name in ["f32", "f64"] {
443 let column = footer.row_groups()[0]
444 .columns()
445 .iter()
446 .find(|column| column.column_path().string() == name)
447 .unwrap();
448 assert!(
449 column
450 .encodings()
451 .any(|encoding| encoding == Encoding::BYTE_STREAM_SPLIT)
452 );
453 }
454 let mut reader = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(bytes))
455 .unwrap()
456 .build()
457 .unwrap();
458 let actual = reader.next().unwrap().unwrap();
459 let actual_f32 = actual
460 .column(0)
461 .as_any()
462 .downcast_ref::<Float32Array>()
463 .unwrap();
464 let actual_f64 = actual
465 .column(1)
466 .as_any()
467 .downcast_ref::<Float64Array>()
468 .unwrap();
469 for (index, value) in f32_values.into_iter().enumerate() {
470 assert_eq!(value.is_none(), actual_f32.is_null(index));
471 if let Some(value) = value {
472 assert_eq!(value.to_bits(), actual_f32.value(index).to_bits());
473 }
474 }
475 for (index, value) in f64_values.into_iter().enumerate() {
476 assert_eq!(value.is_none(), actual_f64.is_null(index));
477 if let Some(value) = value {
478 assert_eq!(value.to_bits(), actual_f64.value(index).to_bits());
479 }
480 }
481 assert!(actual.column(2).is_null(1));
482 assert!(actual.column(3).is_null(1));
483 }
484
485 #[derive(Clone)]
486 struct FixedPathProvider {
487 region_file_id: RegionFileId,
488 }
489
490 impl FilePathProvider for FixedPathProvider {
491 fn build_index_file_path(&self, _file_id: RegionFileId) -> String {
492 location::index_file_path_legacy(FILE_DIR, self.region_file_id, PathType::Bare)
493 }
494
495 fn build_index_file_path_with_version(&self, index_id: RegionIndexId) -> String {
496 location::index_file_path(FILE_DIR, index_id, PathType::Bare)
497 }
498
499 fn build_sst_file_path(&self, _file_id: RegionFileId) -> String {
500 location::sst_file_path(FILE_DIR, self.region_file_id, PathType::Bare)
501 }
502 }
503
504 pub(crate) struct NoopIndexBuilder;
506
507 #[async_trait::async_trait]
508 impl IndexerBuilder for NoopIndexBuilder {
509 async fn build(
510 &self,
511 _file_id: RegionFileId,
512 _index_version: u64,
513 _row_group_size: Option<usize>,
514 ) -> Indexer {
515 Indexer::default()
516 }
517 }
518
519 #[tokio::test]
520 async fn test_parallel_encoding_writes_identical_file() {
521 let mut env = TestEnv::new().await;
522 let object_store = env.init_object_store_manager();
523 let metadata = Arc::new(sst_region_metadata());
524 let mut contents = Vec::new();
525 for encode_parallelism in [1, 3] {
526 let handle = sst_file_handle(0, 1000);
527 let file_path = FixedPathProvider {
528 region_file_id: handle.file_id(),
529 };
530 let path = file_path.build_sst_file_path(handle.file_id());
531 let source = new_flat_source_from_record_batches(vec![
532 new_record_batch_by_range(&["a", "d"], 0, 60),
533 new_record_batch_by_range(&["b", "f"], 0, 40),
534 new_record_batch_by_range(&["b", "h"], 100, 200),
535 ]);
536 let write_opts = WriteOptions {
537 row_group_size: 50,
538 encode_parallelism,
539 ..Default::default()
540 };
541
542 let mut metrics = Metrics::new(WriteType::Flush);
543 let mut writer = ParquetWriter::new_with_object_store(
544 object_store.clone(),
545 metadata.clone(),
546 IndexConfig::default(),
547 NoopIndexBuilder,
548 file_path,
549 &mut metrics,
550 )
551 .await;
552 let info = writer
553 .write_all_flat(source, None, &write_opts)
554 .await
555 .unwrap()
556 .remove(0);
557 assert_eq!(200, info.num_rows);
558 assert_eq!(4, info.num_row_groups);
559 contents.push(object_store.read(&path).await.unwrap().to_vec());
560 }
561 assert!(
562 contents[0] == contents[1],
563 "parallel encoding changed the SST"
564 );
565 }
566
567 #[tokio::test]
568 async fn test_write_read() {
569 let mut env = TestEnv::new().await;
570 let object_store = env.init_object_store_manager();
571 let handle = sst_file_handle(0, 1000);
572 let file_path = FixedPathProvider {
573 region_file_id: handle.file_id(),
574 };
575 let metadata = Arc::new(sst_region_metadata());
576 let source = new_flat_source_from_record_batches(vec![
577 new_record_batch_by_range(&["a", "d"], 0, 60),
578 new_record_batch_by_range(&["b", "f"], 0, 40),
579 new_record_batch_by_range(&["b", "h"], 100, 200),
580 ]);
581 let write_opts = WriteOptions {
583 row_group_size: 50,
584 ..Default::default()
585 };
586
587 let mut metrics = Metrics::new(WriteType::Flush);
588 let mut writer = ParquetWriter::new_with_object_store(
589 object_store.clone(),
590 metadata.clone(),
591 IndexConfig::default(),
592 NoopIndexBuilder,
593 file_path,
594 &mut metrics,
595 )
596 .await;
597
598 let info = writer
599 .write_all_flat_as_primary_key(source, None, &write_opts)
600 .await
601 .unwrap()
602 .remove(0);
603 assert_eq!(200, info.num_rows);
604 assert!(info.file_size > 0);
605 assert_eq!(
606 (
607 Timestamp::new_millisecond(0),
608 Timestamp::new_millisecond(199)
609 ),
610 info.time_range
611 );
612
613 let builder = ParquetReaderBuilder::new(
614 FILE_DIR.to_string(),
615 PathType::Bare,
616 handle.clone(),
617 object_store,
618 );
619 let mut reader = builder.build().await.unwrap().unwrap();
620 check_record_batch_reader_result(
621 &mut reader,
622 &[
623 new_record_batch_by_range(&["a", "d"], 0, 50),
624 new_record_batch_by_range(&["a", "d"], 50, 60),
625 new_record_batch_by_range(&["b", "f"], 0, 40),
626 new_record_batch_by_range(&["b", "h"], 100, 150),
627 new_record_batch_by_range(&["b", "h"], 150, 200),
628 ],
629 )
630 .await;
631 }
632
633 #[tokio::test]
634 async fn test_read_with_cache() {
635 let mut env = TestEnv::new().await;
636 let object_store = env.init_object_store_manager();
637 let handle = sst_file_handle(0, 1000);
638 let metadata = Arc::new(sst_region_metadata());
639 let source = new_flat_source_from_record_batches(vec![
640 new_record_batch_by_range(&["a", "d"], 0, 60),
641 new_record_batch_by_range(&["b", "f"], 0, 40),
642 new_record_batch_by_range(&["b", "h"], 100, 200),
643 ]);
644 let write_opts = WriteOptions {
646 row_group_size: 50,
647 ..Default::default()
648 };
649 let mut metrics = Metrics::new(WriteType::Flush);
651 let mut writer = ParquetWriter::new_with_object_store(
652 object_store.clone(),
653 metadata.clone(),
654 IndexConfig::default(),
655 NoopIndexBuilder,
656 FixedPathProvider {
657 region_file_id: handle.file_id(),
658 },
659 &mut metrics,
660 )
661 .await;
662
663 let sst_info = writer
664 .write_all_flat_as_primary_key(source, None, &write_opts)
665 .await
666 .unwrap()
667 .remove(0);
668
669 let cache = CacheStrategy::EnableAll(Arc::new(
671 CacheManager::builder()
672 .page_cache_size(64 * 1024 * 1024)
673 .build(),
674 ));
675 let builder = ParquetReaderBuilder::new(
676 FILE_DIR.to_string(),
677 PathType::Bare,
678 handle.clone(),
679 object_store,
680 )
681 .cache(cache.clone());
682 for _ in 0..3 {
683 let mut reader = builder.build().await.unwrap().unwrap();
684 check_record_batch_reader_result(
685 &mut reader,
686 &[
687 new_record_batch_by_range(&["a", "d"], 0, 50),
688 new_record_batch_by_range(&["a", "d"], 50, 60),
689 new_record_batch_by_range(&["b", "f"], 0, 40),
690 new_record_batch_by_range(&["b", "h"], 100, 150),
691 new_record_batch_by_range(&["b", "h"], 150, 200),
692 ],
693 )
694 .await;
695 }
696
697 let parquet_meta = sst_info.file_metadata.unwrap();
698 let get_ranges = |row_group_idx: usize| {
699 let row_group = parquet_meta.row_group(row_group_idx);
700 let mut ranges = Vec::with_capacity(row_group.num_columns());
701 for i in 0..row_group.num_columns() {
702 let (start, length) = row_group.column(i).byte_range();
703 ranges.push(start..start + length);
704 }
705
706 ranges
707 };
708
709 let pages = cache.parquet_page_cache().unwrap();
712 for i in 0..4 {
713 for range in get_ranges(i) {
714 for page in range.start / crate::cache::parquet_page::PAGE_SIZE
715 ..=(range.end - 1) / crate::cache::parquet_page::PAGE_SIZE
716 {
717 assert!(pages.contains(handle.file_id(), page).await);
718 }
719 }
720 }
721 let other =
722 crate::sst::file::RegionFileId::new(handle.file_id().region_id(), FileId::random());
723 assert!(!pages.contains(other, 0).await);
724 }
725
726 #[tokio::test]
727 async fn test_parquet_metadata_eq() {
728 let mut env = crate::test_util::TestEnv::new().await;
730 let object_store = env.init_object_store_manager();
731 let handle = sst_file_handle(0, 1000);
732 let metadata = Arc::new(sst_region_metadata());
733 let source = new_flat_source_from_record_batches(vec![
734 new_record_batch_by_range(&["a", "d"], 0, 60),
735 new_record_batch_by_range(&["b", "f"], 0, 40),
736 new_record_batch_by_range(&["b", "h"], 100, 200),
737 ]);
738 let write_opts = WriteOptions {
739 row_group_size: 50,
740 ..Default::default()
741 };
742
743 let mut metrics = Metrics::new(WriteType::Flush);
746 let mut writer = ParquetWriter::new_with_object_store(
747 object_store.clone(),
748 metadata.clone(),
749 IndexConfig::default(),
750 NoopIndexBuilder,
751 FixedPathProvider {
752 region_file_id: handle.file_id(),
753 },
754 &mut metrics,
755 )
756 .await;
757
758 let sst_info = writer
759 .write_all_flat_as_primary_key(source, None, &write_opts)
760 .await
761 .unwrap()
762 .remove(0);
763 let writer_metadata = sst_info.file_metadata.unwrap();
764
765 let builder = ParquetReaderBuilder::new(
767 FILE_DIR.to_string(),
768 PathType::Bare,
769 handle.clone(),
770 object_store,
771 )
772 .page_index_policy(PageIndexPolicy::Optional);
773 let reader = builder.build().await.unwrap().unwrap();
774 let reader_metadata = reader.parquet_metadata();
775 let cached_writer_metadata =
776 crate::cache::CachedSstMeta::try_new("test.sst", Arc::unwrap_or_clone(writer_metadata))
777 .unwrap()
778 .parquet_metadata();
779
780 assert_parquet_metadata_equal(cached_writer_metadata, reader_metadata);
781 }
782
783 #[tokio::test]
784 async fn test_read_with_tag_filter() {
785 let mut env = TestEnv::new().await;
786 let object_store = env.init_object_store_manager();
787 let handle = sst_file_handle(0, 1000);
788 let metadata = Arc::new(sst_region_metadata());
789 let source = new_flat_source_from_record_batches(vec![
790 new_record_batch_by_range(&["a", "d"], 0, 60),
791 new_record_batch_by_range(&["b", "f"], 0, 40),
792 new_record_batch_by_range(&["b", "h"], 100, 200),
793 ]);
794 let write_opts = WriteOptions {
796 row_group_size: 50,
797 ..Default::default()
798 };
799 let mut metrics = Metrics::new(WriteType::Flush);
801 let mut writer = ParquetWriter::new_with_object_store(
802 object_store.clone(),
803 metadata.clone(),
804 IndexConfig::default(),
805 NoopIndexBuilder,
806 FixedPathProvider {
807 region_file_id: handle.file_id(),
808 },
809 &mut metrics,
810 )
811 .await;
812 writer
813 .write_all_flat_as_primary_key(source, None, &write_opts)
814 .await
815 .unwrap()
816 .remove(0);
817
818 let predicate = Some(Predicate::new(vec![Expr::BinaryExpr(BinaryExpr {
820 left: Box::new(Expr::Column(Column::from_name("tag_0"))),
821 op: Operator::Eq,
822 right: Box::new("a".lit()),
823 })]));
824
825 let builder = ParquetReaderBuilder::new(
826 FILE_DIR.to_string(),
827 PathType::Bare,
828 handle.clone(),
829 object_store,
830 )
831 .predicate(predicate);
832 let mut reader = builder.build().await.unwrap().unwrap();
833 check_record_batch_reader_result(
834 &mut reader,
835 &[
836 new_record_batch_by_range(&["a", "d"], 0, 50),
837 new_record_batch_by_range(&["a", "d"], 50, 60),
838 ],
839 )
840 .await;
841 }
842
843 #[tokio::test]
844 async fn test_read_empty_batch() {
845 let mut env = TestEnv::new().await;
846 let object_store = env.init_object_store_manager();
847 let handle = sst_file_handle(0, 1000);
848 let metadata = Arc::new(sst_region_metadata());
849 let source = new_flat_source_from_record_batches(vec![
850 new_record_batch_by_range(&["a", "z"], 0, 0),
851 new_record_batch_by_range(&["a", "z"], 100, 100),
852 new_record_batch_by_range(&["a", "z"], 200, 230),
853 ]);
854 let write_opts = WriteOptions {
856 row_group_size: 50,
857 ..Default::default()
858 };
859 let mut metrics = Metrics::new(WriteType::Flush);
861 let mut writer = ParquetWriter::new_with_object_store(
862 object_store.clone(),
863 metadata.clone(),
864 IndexConfig::default(),
865 NoopIndexBuilder,
866 FixedPathProvider {
867 region_file_id: handle.file_id(),
868 },
869 &mut metrics,
870 )
871 .await;
872 writer
873 .write_all_flat_as_primary_key(source, None, &write_opts)
874 .await
875 .unwrap()
876 .remove(0);
877
878 let builder = ParquetReaderBuilder::new(
879 FILE_DIR.to_string(),
880 PathType::Bare,
881 handle.clone(),
882 object_store,
883 );
884 let mut reader = builder.build().await.unwrap().unwrap();
885 check_record_batch_reader_result(
886 &mut reader,
887 &[new_record_batch_by_range(&["a", "z"], 200, 230)],
888 )
889 .await;
890 }
891
892 #[tokio::test]
893 async fn test_read_with_field_filter() {
894 let mut env = TestEnv::new().await;
895 let object_store = env.init_object_store_manager();
896 let handle = sst_file_handle(0, 1000);
897 let metadata = Arc::new(sst_region_metadata());
898 let source = new_flat_source_from_record_batches(vec![
899 new_record_batch_by_range(&["a", "d"], 0, 60),
900 new_record_batch_by_range(&["b", "f"], 0, 40),
901 new_record_batch_by_range(&["b", "h"], 100, 200),
902 ]);
903 let write_opts = WriteOptions {
905 row_group_size: 50,
906 ..Default::default()
907 };
908 let mut metrics = Metrics::new(WriteType::Flush);
910 let mut writer = ParquetWriter::new_with_object_store(
911 object_store.clone(),
912 metadata.clone(),
913 IndexConfig::default(),
914 NoopIndexBuilder,
915 FixedPathProvider {
916 region_file_id: handle.file_id(),
917 },
918 &mut metrics,
919 )
920 .await;
921
922 writer
923 .write_all_flat_as_primary_key(source, None, &write_opts)
924 .await
925 .unwrap()
926 .remove(0);
927
928 let predicate = Some(Predicate::new(vec![Expr::BinaryExpr(BinaryExpr {
930 left: Box::new(Expr::Column(Column::from_name("field_0"))),
931 op: Operator::GtEq,
932 right: Box::new(150u64.lit()),
933 })]));
934
935 let builder = ParquetReaderBuilder::new(
936 FILE_DIR.to_string(),
937 PathType::Bare,
938 handle.clone(),
939 object_store,
940 )
941 .predicate(predicate);
942 let mut reader = builder.build().await.unwrap().unwrap();
943 check_record_batch_reader_result(
944 &mut reader,
945 &[new_record_batch_by_range(&["b", "h"], 150, 200)],
946 )
947 .await;
948 }
949
950 #[tokio::test]
951 async fn test_column_compression_options() {
952 let words: Vec<String> = (0..256).map(|i| format!("word{i:03}")).collect();
955 let mut seed = 42u64;
956 let values: Vec<Vec<u8>> = (0..2048)
957 .map(|_| {
958 (0..250)
959 .map(|_| {
960 seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
961 words[(seed >> 33) as usize % words.len()].as_str()
962 })
963 .collect::<Vec<_>>()
964 .join(" ")
965 .into_bytes()
966 })
967 .collect();
968
969 async fn write(
970 values: &[Vec<u8>],
971 compression: Option<CompressionOptions>,
972 ) -> (usize, usize, bool) {
973 let mut metadata = (*build_test_binary_test_region_metadata()).clone();
974 if let Some(options) = compression {
975 let column = &mut metadata.column_metadatas[1];
976 assert_eq!("field_0", column.column_schema.name);
977 column.column_schema = column
978 .column_schema
979 .clone()
980 .with_compression_options(options)
981 .unwrap();
982 }
983 let metadata = Arc::new(metadata);
984 let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default());
985 let num_rows = values.len();
986 let pk = new_primary_key(&["a"]);
987 let mut tags = StringDictionaryBuilder::<UInt32Type>::new();
988 let mut pks = BinaryDictionaryBuilder::<UInt32Type>::new();
989 for _ in 0..num_rows {
990 tags.append_value("a");
991 pks.append(&pk).unwrap();
992 }
993 let batch = RecordBatch::try_new(
994 flat_schema,
995 vec![
996 Arc::new(tags.finish()),
997 Arc::new(datatypes::arrow::array::BinaryArray::from_iter_values(
998 values.iter().map(|v| v.as_slice()),
999 )),
1000 Arc::new(TimestampMillisecondArray::from_iter_values(
1001 0..num_rows as i64,
1002 )),
1003 Arc::new(pks.finish()),
1004 Arc::new(UInt64Array::from_value(1000, num_rows)),
1005 Arc::new(UInt8Array::from_value(OpType::Put as u8, num_rows)),
1006 ],
1007 )
1008 .unwrap();
1009
1010 let mut env = TestEnv::new().await;
1011 let object_store = env.init_object_store_manager();
1012 let handle = sst_file_handle(0, 1000);
1013 let mut metrics = Metrics::new(WriteType::Flush);
1014 let mut writer = ParquetWriter::new_with_object_store(
1015 object_store.clone(),
1016 metadata,
1017 IndexConfig::default(),
1018 NoopIndexBuilder,
1019 FixedPathProvider {
1020 region_file_id: handle.file_id(),
1021 },
1022 &mut metrics,
1023 )
1024 .await;
1025 writer
1026 .write_all_flat_as_primary_key(
1027 new_flat_source_from_record_batches(vec![batch]),
1028 None,
1029 &WriteOptions::default(),
1030 )
1031 .await
1032 .unwrap();
1033
1034 let bytes = object_store
1035 .read(&handle.file_path(FILE_DIR, PathType::Bare))
1036 .await
1037 .unwrap()
1038 .to_bytes();
1039 let options = parquet::arrow::arrow_reader::ArrowReaderOptions::new()
1040 .with_page_index_policy(PageIndexPolicy::Required);
1041 let builder =
1042 ParquetRecordBatchReaderBuilder::try_new_with_options(bytes, options).unwrap();
1043 let schema = builder.metadata().file_metadata().schema_descr();
1044 let field = (0..schema.num_columns())
1045 .find(|i| schema.column(*i).name() == "field_0")
1046 .unwrap();
1047 let pages = builder.metadata().offset_index().unwrap()[0][field]
1048 .page_locations()
1049 .len();
1050 let chunk = builder.metadata().row_group(0).column(field);
1051 (
1052 chunk.compressed_size() as usize,
1053 pages,
1054 chunk.dictionary_page_offset().is_some(),
1055 )
1056 }
1057
1058 let (default_size, default_pages, default_dict) = write(&values, None).await;
1059 let (level_size, level_pages, _) = write(
1060 &values,
1061 Some(CompressionOptions {
1062 level: Some(19),
1063 ..Default::default()
1064 }),
1065 )
1066 .await;
1067 let (_, pages, _) = write(
1068 &values,
1069 Some(CompressionOptions {
1070 page_size: Some(ReadableSize::kb(64)),
1071 ..Default::default()
1072 }),
1073 )
1074 .await;
1075 let (_, _, plain_dict) = write(
1076 &values,
1077 Some(CompressionOptions {
1078 encoding: Some(ColumnEncoding::Plain),
1079 ..Default::default()
1080 }),
1081 )
1082 .await;
1083
1084 assert!(default_pages > 1, "{default_pages}");
1085 assert_eq!(default_pages, level_pages);
1086 assert!(level_size < default_size, "{level_size} >= {default_size}");
1087 assert!(pages > default_pages, "{pages} <= {default_pages}");
1088 assert!(default_dict);
1089 assert!(!plain_dict);
1090 }
1091
1092 #[tokio::test]
1093 async fn test_read_large_binary() {
1094 let mut env = TestEnv::new().await;
1095 let object_store = env.init_object_store_manager();
1096 let handle = sst_file_handle(0, 1000);
1097 let file_path = handle.file_path(FILE_DIR, PathType::Bare);
1098
1099 let write_opts = WriteOptions {
1100 row_group_size: 50,
1101 ..Default::default()
1102 };
1103
1104 let metadata = build_test_binary_test_region_metadata();
1105 let json = metadata.to_json().unwrap();
1106 let key_value_meta = KeyValue::new(PARQUET_METADATA_KEY.to_string(), json);
1107
1108 let props_builder = WriterProperties::builder()
1109 .set_key_value_metadata(Some(vec![key_value_meta]))
1110 .set_compression(Compression::ZSTD(ZstdLevel::default()))
1111 .set_encoding(Encoding::PLAIN)
1112 .set_max_row_group_row_count(Some(write_opts.row_group_size));
1113
1114 let writer_props = props_builder.build();
1115
1116 let write_format = FlatWriteFormat::new(metadata, &FlatSchemaOptions::default());
1117 let fields: Vec<_> = write_format
1118 .arrow_schema()
1119 .fields()
1120 .into_iter()
1121 .map(|field| {
1122 let data_type = field.data_type().clone();
1123 if data_type == DataType::Binary {
1124 Field::new(field.name(), DataType::LargeBinary, field.is_nullable())
1125 } else {
1126 Field::new(field.name(), data_type, field.is_nullable())
1127 }
1128 })
1129 .collect();
1130
1131 let arrow_schema = Arc::new(Schema::new(fields));
1132
1133 assert_eq!(
1135 &DataType::LargeBinary,
1136 arrow_schema.field_with_name("field_0").unwrap().data_type()
1137 );
1138 let mut writer = AsyncArrowWriter::try_new(
1139 object_store
1140 .writer_with(&file_path)
1141 .concurrent(DEFAULT_WRITE_CONCURRENCY)
1142 .await
1143 .map(|w| w.into_futures_async_write().compat_write())
1144 .unwrap(),
1145 arrow_schema.clone(),
1146 Some(writer_props),
1147 )
1148 .unwrap();
1149
1150 let batch = new_record_batch_with_binary(&["a"], 0, 60);
1151 let arrays: Vec<_> = batch
1152 .columns()
1153 .iter()
1154 .map(|array| {
1155 let data_type = array.data_type().clone();
1156 if data_type == DataType::Binary {
1157 arrow::compute::cast(array, &DataType::LargeBinary).unwrap()
1158 } else {
1159 array.clone()
1160 }
1161 })
1162 .collect();
1163 let result = RecordBatch::try_new(arrow_schema, arrays).unwrap();
1164
1165 writer.write(&result).await.unwrap();
1166 writer.close().await.unwrap();
1167
1168 let builder = ParquetReaderBuilder::new(
1169 FILE_DIR.to_string(),
1170 PathType::Bare,
1171 handle.clone(),
1172 object_store,
1173 );
1174 let mut reader = builder.build().await.unwrap().unwrap();
1175 check_record_batch_reader_result(
1176 &mut reader,
1177 &[
1178 new_record_batch_with_binary(&["a"], 0, 50),
1179 new_record_batch_with_binary(&["a"], 50, 60),
1180 ],
1181 )
1182 .await;
1183 }
1184
1185 #[tokio::test]
1186 async fn test_column_index_truncates_large_field_values() {
1187 let mut env = TestEnv::new().await;
1188 let object_store = env.init_object_store_manager();
1189 let handle = sst_file_handle(0, 1000);
1190 let file_path = FixedPathProvider {
1191 region_file_id: handle.file_id(),
1192 };
1193 let metadata = build_test_binary_test_region_metadata();
1194 let tag = "t".repeat(100);
1196 let values: Vec<Vec<u8>> = (0..64)
1197 .map(|i| format!("{i:08}").into_bytes().repeat(4 * 1024))
1198 .collect();
1199 let batch = new_record_batch_with_binary_values(&tag, &values);
1200 let mut metrics = Metrics::new(WriteType::Flush);
1201 let mut writer = ParquetWriter::new_with_object_store(
1202 object_store.clone(),
1203 metadata.clone(),
1204 IndexConfig::default(),
1205 NoopIndexBuilder,
1206 file_path,
1207 &mut metrics,
1208 )
1209 .await;
1210 writer
1211 .write_all_flat_as_primary_key(
1212 new_flat_source_from_record_batches(vec![batch]),
1213 None,
1214 &WriteOptions::default(),
1215 )
1216 .await
1217 .unwrap();
1218
1219 let path = handle.file_path(FILE_DIR, PathType::Bare);
1220 let bytes = object_store.read(&path).await.unwrap().to_bytes();
1221 let options = parquet::arrow::arrow_reader::ArrowReaderOptions::new()
1222 .with_page_index_policy(PageIndexPolicy::Required);
1223 let builder =
1224 ParquetRecordBatchReaderBuilder::try_new_with_options(bytes, options).unwrap();
1225 let parquet_meta = builder.metadata().clone();
1226 let schema = parquet_meta.file_metadata().schema_descr();
1227 let column = |name: &str| {
1228 (0..schema.num_columns())
1229 .find(|i| schema.column(*i).name() == name)
1230 .unwrap()
1231 };
1232
1233 let field_index = &parquet_meta.column_index().unwrap()[0][column("field_0")];
1234 let parquet::file::page_index::column_index::ColumnIndexMetaData::BYTE_ARRAY(field_index) =
1235 field_index
1236 else {
1237 panic!("unexpected column index: {field_index:?}");
1238 };
1239 let mut pages = 0;
1240 for (min, max) in field_index
1241 .min_values_iter()
1242 .zip(field_index.max_values_iter())
1243 {
1244 assert!(min.unwrap().len() <= 64);
1245 assert!(max.unwrap().len() <= 64);
1246 pages += 1;
1247 }
1248 assert!(pages > 1, "values should span several pages");
1249
1250 let pk_stats = parquet_meta
1252 .row_group(0)
1253 .column(column(store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME))
1254 .statistics()
1255 .unwrap();
1256 let pk = new_primary_key(&[&tag]);
1257 assert!(pk.len() > 64);
1258 assert_eq!(pk_stats.max_bytes_opt().unwrap(), pk.as_slice());
1259 }
1260
1261 #[rstest::rstest]
1262 #[tokio::test]
1263 async fn test_write_multiple_files(#[values(1024, 4096)] write_buffer_size: usize) {
1264 common_telemetry::init_default_ut_logging();
1265 let mut env = TestEnv::new().await;
1267 let chunks = WriteChunkRecorder::default();
1268 let object_store = env.init_object_store_manager().layer(chunks.layer());
1269 let metadata = Arc::new(sst_region_metadata());
1270 let batches = vec![
1271 new_record_batch_by_range(&["a", "a"], 0, 1000),
1272 new_record_batch_by_range(&["b", "b"], 0, 1000),
1273 new_record_batch_by_range(&["c", "c"], 0, 1000),
1274 new_record_batch_by_range(&["d", "d"], 100, 200),
1275 new_record_batch_by_range(&["d", "d"], 200, 300),
1276 new_record_batch_by_range(&["d", "d"], 300, 1000),
1277 new_record_batch_by_range(&["e", "e"], 0, 100),
1278 ];
1279 let total_rows: usize = batches.iter().map(|batch| batch.num_rows()).sum();
1280
1281 let source = new_flat_source_from_record_batches(batches);
1282 let write_opts = WriteOptions {
1283 write_buffer_size: ReadableSize(write_buffer_size as u64),
1284 row_group_size: 50,
1285 max_file_size: Some(1024 * 16),
1286 ..Default::default()
1287 };
1288
1289 let path_provider = RegionFilePathFactory {
1290 table_dir: "test".to_string(),
1291 path_type: PathType::Bare,
1292 };
1293 let mut metrics = Metrics::new(WriteType::Flush);
1294 let mut writer = ParquetWriter::new_with_object_store(
1295 object_store.clone(),
1296 metadata.clone(),
1297 IndexConfig::default(),
1298 NoopIndexBuilder,
1299 path_provider.clone(),
1300 &mut metrics,
1301 )
1302 .await;
1303
1304 let files = writer
1305 .write_all_flat_as_primary_key(source, None, &write_opts)
1306 .await
1307 .unwrap();
1308 assert_eq!(2, files.len());
1309
1310 assert_eq!(files.len(), chunks.num_files());
1312
1313 let mut rows_read = 0;
1314 for f in &files {
1315 assert!(f.file_size > write_buffer_size as u64);
1316 chunks.assert_chunks(
1317 &path_provider
1318 .build_sst_file_path(RegionFileId::new(metadata.region_id, f.file_id)),
1319 write_buffer_size,
1320 f.file_size as usize,
1321 );
1322 let file_handle = sst_file_handle_with_file_id(
1323 f.file_id,
1324 f.time_range.0.value(),
1325 f.time_range.1.value(),
1326 );
1327 let builder = ParquetReaderBuilder::new(
1328 "test".to_string(),
1329 PathType::Bare,
1330 file_handle,
1331 object_store.clone(),
1332 );
1333 let mut reader = builder.build().await.unwrap().unwrap();
1334 while let Some(batch) = reader.next_record_batch().await.unwrap() {
1335 rows_read += batch.num_rows();
1336 }
1337 }
1338 assert_eq!(total_rows, rows_read);
1339 }
1340
1341 #[tokio::test]
1342 async fn test_split_file_at_series_boundary_inside_batch() {
1343 let mut env = TestEnv::new().await;
1344 let object_store = env.init_object_store_manager();
1345 let metadata = Arc::new(sst_region_metadata());
1346 let first_batch_rows = (0..1000).map(|ts| ("a", "a", ts)).collect::<Vec<_>>();
1347 let second_batch_rows = (0..1000).map(|ts| ("b", "b", ts)).collect::<Vec<_>>();
1348 let mut third_batch_rows = (1000..2000).map(|ts| ("b", "b", ts)).collect::<Vec<_>>();
1349 third_batch_rows.extend((0..1000).map(|ts| ("c", "c", ts)));
1350 let batches = vec![
1351 new_record_batch_from_rows(&first_batch_rows),
1352 new_record_batch_from_rows(&second_batch_rows),
1353 new_record_batch_from_rows(&third_batch_rows),
1354 ];
1355 let total_rows = batches.iter().map(RecordBatch::num_rows).sum::<usize>();
1356 let source = new_flat_source_from_record_batches(batches);
1357 let write_opts = WriteOptions {
1358 row_group_size: 50,
1359 max_file_size: Some(1),
1360 ..Default::default()
1361 };
1362 let path_provider = RegionFilePathFactory {
1363 table_dir: "test_series_boundary".to_string(),
1364 path_type: PathType::Bare,
1365 };
1366 let mut metrics = Metrics::new(WriteType::Compaction);
1367 let mut writer = ParquetWriter::new_with_object_store(
1368 object_store,
1369 metadata.clone(),
1370 IndexConfig::default(),
1371 NoopIndexBuilder,
1372 path_provider,
1373 &mut metrics,
1374 )
1375 .await;
1376
1377 let files = writer
1378 .write_all_flat(source, None, &write_opts)
1379 .await
1380 .unwrap();
1381
1382 assert!(files.len() > 1);
1383 assert_eq!(
1384 total_rows,
1385 files.iter().map(|file| file.num_rows).sum::<usize>()
1386 );
1387 let primary_key_ranges = files
1388 .iter()
1389 .map(|file| {
1390 extract_primary_key_range(file.file_metadata.as_ref().unwrap(), metadata.as_ref())
1391 .unwrap()
1392 })
1393 .collect::<Vec<_>>();
1394 assert!(
1395 primary_key_ranges
1396 .windows(2)
1397 .all(|ranges| ranges[0].1 < ranges[1].0)
1398 );
1399 }
1400
1401 #[tokio::test]
1402 async fn test_oversized_single_series_stays_in_one_file() {
1403 let mut env = TestEnv::new().await;
1404 let object_store = env.init_object_store_manager();
1405 let metadata = Arc::new(sst_region_metadata());
1406 let first_batch_rows = (0..1000).map(|ts| ("a", "a", ts)).collect::<Vec<_>>();
1407 let second_batch_rows = (1000..2000).map(|ts| ("a", "a", ts)).collect::<Vec<_>>();
1408 let third_batch_rows = (2000..3000).map(|ts| ("a", "a", ts)).collect::<Vec<_>>();
1409 let batches = vec![
1410 new_record_batch_from_rows(&first_batch_rows),
1411 new_record_batch_from_rows(&second_batch_rows),
1412 new_record_batch_from_rows(&third_batch_rows),
1413 ];
1414 let total_rows = batches.iter().map(RecordBatch::num_rows).sum::<usize>();
1415 let source = new_flat_source_from_record_batches(batches);
1416 let write_opts = WriteOptions {
1417 row_group_size: 50,
1418 max_file_size: Some(1),
1419 ..Default::default()
1420 };
1421 let path_provider = RegionFilePathFactory {
1422 table_dir: "test_oversized_series".to_string(),
1423 path_type: PathType::Bare,
1424 };
1425 let mut metrics = Metrics::new(WriteType::Compaction);
1426 let mut writer = ParquetWriter::new_with_object_store(
1427 object_store,
1428 metadata,
1429 IndexConfig::default(),
1430 NoopIndexBuilder,
1431 path_provider,
1432 &mut metrics,
1433 )
1434 .await;
1435
1436 let files = writer
1437 .write_all_flat_as_primary_key(source, None, &write_opts)
1438 .await
1439 .unwrap();
1440
1441 assert_eq!(1, files.len());
1442 assert_eq!(total_rows, files[0].num_rows);
1443 assert!(files[0].file_size > write_opts.max_file_size.unwrap() as u64);
1444 }
1445
1446 #[tokio::test]
1447 async fn test_write_multiple_files_without_primary_key() {
1448 let mut env = TestEnv::new().await;
1449 let object_store = env.init_object_store_manager();
1450 let metadata = Arc::new(sst_region_metadata_without_primary_key());
1451 let batch_rows = 1000;
1452 let batches = vec![
1453 new_record_batch_without_primary_key(0, batch_rows),
1454 new_record_batch_without_primary_key(batch_rows, 2 * batch_rows),
1455 new_record_batch_without_primary_key(2 * batch_rows, 3 * batch_rows),
1456 ];
1457 let total_rows = batches.iter().map(RecordBatch::num_rows).sum::<usize>();
1458 let source = new_flat_source_from_record_batches(batches);
1459 let write_opts = WriteOptions {
1460 row_group_size: 50,
1461 max_file_size: Some(1),
1462 ..Default::default()
1463 };
1464 let path_provider = RegionFilePathFactory {
1465 table_dir: "test_no_primary_key".to_string(),
1466 path_type: PathType::Bare,
1467 };
1468 let mut metrics = Metrics::new(WriteType::Compaction);
1469 let mut writer = ParquetWriter::new_with_object_store(
1470 object_store,
1471 metadata,
1472 IndexConfig::default(),
1473 NoopIndexBuilder,
1474 path_provider,
1475 &mut metrics,
1476 )
1477 .await;
1478
1479 let files = writer
1480 .write_all_flat(source, None, &write_opts)
1481 .await
1482 .unwrap();
1483
1484 assert!(files.len() > 1);
1487 assert!(files.iter().all(|file| file.num_rows % batch_rows == 0));
1488 assert_eq!(
1489 total_rows,
1490 files.iter().map(|file| file.num_rows).sum::<usize>()
1491 );
1492 }
1493
1494 #[tokio::test]
1495 async fn test_write_read_with_index() {
1496 let mut env = TestEnv::new().await;
1497 let object_store = env.init_object_store_manager();
1498 let file_path = RegionFilePathFactory::new(FILE_DIR.to_string(), PathType::Bare);
1499 let metadata = Arc::new(sst_region_metadata());
1500 let row_group_size = 50;
1501
1502 let source = new_flat_source_from_record_batches(vec![
1503 new_record_batch_by_range(&["a", "d"], 0, 20),
1504 new_record_batch_by_range(&["b", "d"], 0, 20),
1505 new_record_batch_by_range(&["c", "d"], 0, 20),
1506 new_record_batch_by_range(&["c", "f"], 0, 40),
1507 new_record_batch_by_range(&["c", "h"], 100, 200),
1508 ]);
1509 let write_opts = WriteOptions {
1511 row_group_size,
1512 ..Default::default()
1513 };
1514
1515 let puffin_manager = env
1516 .get_puffin_manager()
1517 .build(object_store.clone(), file_path.clone());
1518 let intermediate_manager = env.get_intermediate_manager();
1519
1520 let indexer_builder = IndexerBuilderImpl {
1521 build_type: IndexBuildType::Flush,
1522 metadata: metadata.clone(),
1523 puffin_manager,
1524 write_cache_enabled: false,
1525 intermediate_manager,
1526 index_options: IndexOptions {
1527 inverted_index: InvertedIndexOptions {
1528 segment_row_count: 1,
1529 ..Default::default()
1530 },
1531 },
1532 inverted_index_config: Default::default(),
1533 fulltext_index_config: Default::default(),
1534 bloom_filter_index_config: Default::default(),
1535 };
1536
1537 let mut metrics = Metrics::new(WriteType::Flush);
1538 let mut writer = ParquetWriter::new_with_object_store(
1539 object_store.clone(),
1540 metadata.clone(),
1541 IndexConfig::default(),
1542 indexer_builder,
1543 file_path.clone(),
1544 &mut metrics,
1545 )
1546 .await;
1547
1548 let info = writer
1549 .write_all_flat_as_primary_key(source, None, &write_opts)
1550 .await
1551 .unwrap()
1552 .remove(0);
1553 assert_eq!(200, info.num_rows);
1554 assert!(info.file_size > 0);
1555 assert!(info.index_metadata.file_size > 0);
1556
1557 assert!(info.index_metadata.inverted_index.index_size > 0);
1558 assert_eq!(info.index_metadata.inverted_index.row_count, 200);
1559 assert_eq!(info.index_metadata.inverted_index.columns, vec![0]);
1560
1561 assert!(info.index_metadata.bloom_filter.index_size > 0);
1562 assert_eq!(info.index_metadata.bloom_filter.row_count, 200);
1563 assert_eq!(info.index_metadata.bloom_filter.columns, vec![1]);
1564
1565 assert_eq!(
1566 (
1567 Timestamp::new_millisecond(0),
1568 Timestamp::new_millisecond(199)
1569 ),
1570 info.time_range
1571 );
1572
1573 let handle = FileHandle::new(
1574 FileMeta {
1575 region_id: metadata.region_id,
1576 file_id: info.file_id,
1577 time_range: info.time_range,
1578 level: 0,
1579 file_size: info.file_size,
1580 max_row_group_uncompressed_size: info.max_row_group_uncompressed_size,
1581 available_indexes: info.index_metadata.build_available_indexes(),
1582 indexes: info.index_metadata.build_indexes(),
1583 index_file_size: info.index_metadata.file_size,
1584 index_version: 0,
1585 num_row_groups: info.num_row_groups,
1586 num_rows: info.num_rows as u64,
1587 sequence: None,
1588 partition_expr: match &metadata.partition_expr {
1589 Some(json_str) => partition::expr::PartitionExpr::from_json_str(json_str)
1590 .expect("partition expression should be valid JSON"),
1591 None => None,
1592 },
1593 num_series: 0,
1594 ..Default::default()
1595 },
1596 Arc::new(NoopFilePurger),
1597 );
1598
1599 let cache = Arc::new(
1600 CacheManager::builder()
1601 .index_result_cache_size(1024 * 1024)
1602 .index_metadata_size(1024 * 1024)
1603 .index_content_page_size(1024 * 1024)
1604 .index_content_size(1024 * 1024)
1605 .puffin_metadata_size(1024 * 1024)
1606 .build(),
1607 );
1608 let index_result_cache = cache.index_result_cache().unwrap();
1609
1610 let build_inverted_index_applier = |exprs: &[Expr]| {
1611 InvertedIndexApplierBuilder::new(
1612 FILE_DIR.to_string(),
1613 PathType::Bare,
1614 object_store.clone(),
1615 &metadata,
1616 HashSet::from_iter([0]),
1617 env.get_puffin_manager(),
1618 )
1619 .with_puffin_metadata_cache(cache.puffin_metadata_cache().cloned())
1620 .with_inverted_index_cache(cache.inverted_index_cache().cloned())
1621 .build(exprs)
1622 .unwrap()
1623 .map(Arc::new)
1624 };
1625
1626 let build_bloom_filter_applier = |exprs: &[Expr]| {
1627 BloomFilterIndexApplierBuilder::new(
1628 FILE_DIR.to_string(),
1629 PathType::Bare,
1630 object_store.clone(),
1631 &metadata,
1632 env.get_puffin_manager(),
1633 )
1634 .with_puffin_metadata_cache(cache.puffin_metadata_cache().cloned())
1635 .with_bloom_filter_index_cache(cache.bloom_filter_index_cache().cloned())
1636 .build(exprs)
1637 .unwrap()
1638 .map(Arc::new)
1639 };
1640
1641 let preds = vec![col("tag_0").eq(lit("b"))];
1658 let inverted_index_applier = build_inverted_index_applier(&preds);
1659 let bloom_filter_applier = build_bloom_filter_applier(&preds);
1660
1661 let builder = ParquetReaderBuilder::new(
1662 FILE_DIR.to_string(),
1663 PathType::Bare,
1664 handle.clone(),
1665 object_store.clone(),
1666 )
1667 .predicate(Some(Predicate::new(preds)))
1668 .inverted_index_appliers([inverted_index_applier.clone(), None])
1669 .bloom_filter_index_appliers([bloom_filter_applier.clone(), None])
1670 .cache(CacheStrategy::EnableAll(cache.clone()));
1671
1672 let mut metrics = ReaderMetrics::default();
1673 let (context, selection) = builder
1674 .build_reader_input(&mut metrics)
1675 .await
1676 .unwrap()
1677 .unwrap();
1678 let mut reader = ParquetReader::new(Arc::new(context), selection)
1679 .await
1680 .unwrap();
1681 check_record_batch_reader_result(
1682 &mut reader,
1683 &[new_record_batch_by_range(&["b", "d"], 0, 20)],
1684 )
1685 .await;
1686
1687 assert_eq!(metrics.filter_metrics.rg_total, 4);
1688 assert_eq!(metrics.filter_metrics.rg_minmax_filtered, 3);
1689 assert_eq!(metrics.filter_metrics.rg_inverted_filtered, 0);
1690 assert_eq!(metrics.filter_metrics.rows_inverted_filtered, 30);
1691 let plan = inverted_index_applier
1692 .as_ref()
1693 .unwrap()
1694 .plan_for_sst(&metadata)
1695 .unwrap()
1696 .unwrap();
1697 let cached = index_result_cache
1698 .get(&plan.predicate_key, handle.file_id().file_id())
1699 .unwrap();
1700 assert!(cached.contains_row_group(0));
1702 assert!(cached.contains_row_group(1));
1703 assert!(cached.contains_row_group(2));
1704 assert!(cached.contains_row_group(3));
1705
1706 let preds = vec![
1721 col("ts").gt_eq(lit(ScalarValue::TimestampMillisecond(Some(50), None))),
1722 col("ts").lt(lit(ScalarValue::TimestampMillisecond(Some(200), None))),
1723 col("tag_1").eq(lit("d")),
1724 ];
1725 let inverted_index_applier = build_inverted_index_applier(&preds);
1726 let bloom_filter_applier = build_bloom_filter_applier(&preds);
1727
1728 let builder = ParquetReaderBuilder::new(
1729 FILE_DIR.to_string(),
1730 PathType::Bare,
1731 handle.clone(),
1732 object_store.clone(),
1733 )
1734 .predicate(Some(Predicate::new(preds)))
1735 .inverted_index_appliers([inverted_index_applier.clone(), None])
1736 .bloom_filter_index_appliers([bloom_filter_applier.clone(), None])
1737 .cache(CacheStrategy::EnableAll(cache.clone()));
1738
1739 let mut metrics = ReaderMetrics::default();
1740 let read_input = builder.build_reader_input(&mut metrics).await.unwrap();
1741 assert!(read_input.is_none());
1742
1743 assert_eq!(metrics.filter_metrics.rg_total, 4);
1744 assert_eq!(metrics.filter_metrics.rg_minmax_filtered, 2);
1745 assert_eq!(metrics.filter_metrics.rg_bloom_filtered, 2);
1746 assert_eq!(metrics.filter_metrics.rows_bloom_filtered, 100);
1747 let bloom_predicates = bloom_filter_applier
1748 .as_ref()
1749 .unwrap()
1750 .compatible_predicate_for_sst(&metadata)
1751 .unwrap();
1752 let bloom_predicate_key = PredicateKey::new_bloom(bloom_predicates);
1753 let cached = index_result_cache
1754 .get(&bloom_predicate_key, handle.file_id().file_id())
1755 .unwrap();
1756 assert!(cached.contains_row_group(2));
1757 assert!(cached.contains_row_group(3));
1758 assert!(!cached.contains_row_group(0));
1759 assert!(!cached.contains_row_group(1));
1760
1761 let preds = vec![col("tag_1").eq(lit("d"))];
1782 let inverted_index_applier = build_inverted_index_applier(&preds);
1783 let bloom_filter_applier = build_bloom_filter_applier(&preds);
1784
1785 let builder = ParquetReaderBuilder::new(
1786 FILE_DIR.to_string(),
1787 PathType::Bare,
1788 handle.clone(),
1789 object_store.clone(),
1790 )
1791 .predicate(Some(Predicate::new(preds)))
1792 .inverted_index_appliers([inverted_index_applier.clone(), None])
1793 .bloom_filter_index_appliers([bloom_filter_applier.clone(), None])
1794 .cache(CacheStrategy::EnableAll(cache.clone()));
1795
1796 let mut metrics = ReaderMetrics::default();
1797 let (context, selection) = builder
1798 .build_reader_input(&mut metrics)
1799 .await
1800 .unwrap()
1801 .unwrap();
1802 let mut reader = ParquetReader::new(Arc::new(context), selection)
1803 .await
1804 .unwrap();
1805 check_record_batch_reader_result(
1806 &mut reader,
1807 &[
1808 new_record_batch_by_range(&["a", "d"], 0, 20),
1809 new_record_batch_by_range(&["b", "d"], 0, 20),
1810 new_record_batch_by_range(&["c", "d"], 0, 10),
1811 new_record_batch_by_range(&["c", "d"], 10, 20),
1812 ],
1813 )
1814 .await;
1815
1816 assert_eq!(metrics.filter_metrics.rg_total, 4);
1817 assert_eq!(metrics.filter_metrics.rg_minmax_filtered, 0);
1818 assert_eq!(metrics.filter_metrics.rg_bloom_filtered, 2);
1819 assert_eq!(metrics.filter_metrics.rows_bloom_filtered, 140);
1820 let bloom_predicates = bloom_filter_applier
1821 .as_ref()
1822 .unwrap()
1823 .compatible_predicate_for_sst(&metadata)
1824 .unwrap();
1825 let bloom_predicate_key = PredicateKey::new_bloom(bloom_predicates);
1826 let cached = index_result_cache
1827 .get(&bloom_predicate_key, handle.file_id().file_id())
1828 .unwrap();
1829 assert!(cached.contains_row_group(0));
1830 assert!(cached.contains_row_group(1));
1831 assert!(cached.contains_row_group(2));
1832 assert!(cached.contains_row_group(3));
1833 }
1834
1835 fn new_record_batch_with_binary_values(tag: &str, values: &[Vec<u8>]) -> RecordBatch {
1836 let metadata = build_test_binary_test_region_metadata();
1837 let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default());
1838 let num_rows = values.len();
1839 let mut tag_0_builder = StringDictionaryBuilder::<UInt32Type>::new();
1840 let mut pk_builder = BinaryDictionaryBuilder::<UInt32Type>::new();
1841 let pk = new_primary_key(&[tag]);
1842 for _ in 0..num_rows {
1843 tag_0_builder.append_value(tag);
1844 pk_builder.append(&pk).unwrap();
1845 }
1846 let columns: Vec<ArrayRef> = vec![
1847 Arc::new(tag_0_builder.finish()),
1848 Arc::new(datatypes::arrow::array::BinaryArray::from_iter_values(
1849 values.iter().map(|v| v.as_slice()),
1850 )),
1851 Arc::new(TimestampMillisecondArray::from_iter_values(
1852 0..num_rows as i64,
1853 )),
1854 Arc::new(pk_builder.finish()),
1855 Arc::new(UInt64Array::from_value(1000, num_rows)),
1856 Arc::new(UInt8Array::from_value(OpType::Put as u8, num_rows)),
1857 ];
1858 RecordBatch::try_new(flat_schema, columns).unwrap()
1859 }
1860
1861 fn new_record_batch_with_binary(tags: &[&str], start: usize, end: usize) -> RecordBatch {
1862 assert!(end >= start);
1863 let metadata = build_test_binary_test_region_metadata();
1864 let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default());
1865
1866 let num_rows = end - start;
1867 let mut columns = Vec::new();
1868
1869 let mut tag_0_builder = StringDictionaryBuilder::<UInt32Type>::new();
1870 for _ in 0..num_rows {
1871 tag_0_builder.append_value(tags[0]);
1872 }
1873 columns.push(Arc::new(tag_0_builder.finish()) as ArrayRef);
1874
1875 let values = (0..num_rows)
1876 .map(|_| "some data".as_bytes())
1877 .collect::<Vec<_>>();
1878 columns.push(
1879 Arc::new(datatypes::arrow::array::BinaryArray::from_iter_values(
1880 values,
1881 )) as ArrayRef,
1882 );
1883
1884 let timestamps: Vec<i64> = (start..end).map(|v| v as i64).collect();
1885 columns.push(Arc::new(TimestampMillisecondArray::from(timestamps)));
1886
1887 let pk = new_primary_key(tags);
1888 let mut pk_builder = BinaryDictionaryBuilder::<UInt32Type>::new();
1889 for _ in 0..num_rows {
1890 pk_builder.append(&pk).unwrap();
1891 }
1892 columns.push(Arc::new(pk_builder.finish()));
1893
1894 columns.push(Arc::new(UInt64Array::from_value(1000, num_rows)));
1895 columns.push(Arc::new(UInt8Array::from_value(
1896 OpType::Put as u8,
1897 num_rows,
1898 )));
1899
1900 RecordBatch::try_new(flat_schema, columns).unwrap()
1901 }
1902
1903 async fn check_record_batch_reader_result(
1904 reader: &mut ParquetReader,
1905 expected: &[RecordBatch],
1906 ) {
1907 let mut actual = Vec::new();
1908 while let Some(batch) = reader.next_record_batch().await.unwrap() {
1909 actual.push(batch);
1910 }
1911 assert_eq!(
1912 pretty_format_batches(expected).unwrap().to_string(),
1913 pretty_format_batches(&actual).unwrap().to_string()
1914 );
1915 assert!(reader.next_record_batch().await.unwrap().is_none());
1916 }
1917
1918 fn sst_region_metadata_without_primary_key() -> RegionMetadata {
1922 let mut builder = RegionMetadataBuilder::new(REGION_ID);
1923 builder
1924 .push_column_metadata(ColumnMetadata {
1925 column_schema: ColumnSchema::new(
1926 "field_0".to_string(),
1927 ConcreteDataType::uint64_datatype(),
1928 true,
1929 ),
1930 semantic_type: SemanticType::Field,
1931 column_id: 0,
1932 })
1933 .push_column_metadata(ColumnMetadata {
1934 column_schema: ColumnSchema::new(
1935 "ts".to_string(),
1936 ConcreteDataType::timestamp_millisecond_datatype(),
1937 false,
1938 ),
1939 semantic_type: SemanticType::Timestamp,
1940 column_id: 1,
1941 })
1942 .primary_key(vec![]);
1943 builder.build().unwrap()
1944 }
1945
1946 fn new_record_batch_without_primary_key(start: usize, end: usize) -> RecordBatch {
1948 assert!(end >= start);
1949 let metadata = Arc::new(sst_region_metadata_without_primary_key());
1950 let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default());
1951
1952 let num_rows = end - start;
1953 let mut pk_builder = BinaryDictionaryBuilder::<UInt32Type>::new();
1954 for _ in 0..num_rows {
1956 pk_builder.append([]).unwrap();
1957 }
1958
1959 RecordBatch::try_new(
1960 flat_schema,
1961 vec![
1962 Arc::new(UInt64Array::from_iter_values(start as u64..end as u64)) as ArrayRef,
1963 Arc::new(TimestampMillisecondArray::from_iter_values(
1964 start as i64..end as i64,
1965 )) as ArrayRef,
1966 Arc::new(pk_builder.finish()) as ArrayRef,
1967 Arc::new(UInt64Array::from_value(1000, num_rows)) as ArrayRef,
1968 Arc::new(UInt8Array::from_value(OpType::Put as u8, num_rows)) as ArrayRef,
1969 ],
1970 )
1971 .unwrap()
1972 }
1973
1974 fn new_record_batch_from_rows(rows: &[(&str, &str, i64)]) -> RecordBatch {
1975 let metadata = Arc::new(sst_region_metadata());
1976 let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default());
1977
1978 let mut tag_0_builder = StringDictionaryBuilder::<UInt32Type>::new();
1979 let mut tag_1_builder = StringDictionaryBuilder::<UInt32Type>::new();
1980 let mut pk_builder = BinaryDictionaryBuilder::<UInt32Type>::new();
1981 let mut field_values = Vec::with_capacity(rows.len());
1982 let mut timestamps = Vec::with_capacity(rows.len());
1983
1984 for (tag_0, tag_1, ts) in rows {
1985 tag_0_builder.append_value(*tag_0);
1986 tag_1_builder.append_value(*tag_1);
1987 pk_builder.append(new_primary_key(&[tag_0, tag_1])).unwrap();
1988 field_values.push(*ts as u64);
1989 timestamps.push(*ts);
1990 }
1991
1992 RecordBatch::try_new(
1993 flat_schema,
1994 vec![
1995 Arc::new(tag_0_builder.finish()) as ArrayRef,
1996 Arc::new(tag_1_builder.finish()) as ArrayRef,
1997 Arc::new(UInt64Array::from(field_values)) as ArrayRef,
1998 Arc::new(TimestampMillisecondArray::from(timestamps)) as ArrayRef,
1999 Arc::new(pk_builder.finish()) as ArrayRef,
2000 Arc::new(UInt64Array::from_value(1000, rows.len())) as ArrayRef,
2001 Arc::new(UInt8Array::from_value(OpType::Put as u8, rows.len())) as ArrayRef,
2002 ],
2003 )
2004 .unwrap()
2005 }
2006
2007 fn new_record_batch_by_range_sparse(
2010 tags: &[&str],
2011 start: usize,
2012 end: usize,
2013 metadata: &Arc<RegionMetadata>,
2014 ) -> RecordBatch {
2015 assert!(end >= start);
2016 let flat_schema = to_flat_sst_arrow_schema(
2017 metadata,
2018 &FlatSchemaOptions::from_encoding(PrimaryKeyEncoding::Sparse),
2019 );
2020
2021 let num_rows = end - start;
2022 let mut columns: Vec<ArrayRef> = Vec::new();
2023
2024 let field_values: Vec<u64> = (start..end).map(|v| v as u64).collect();
2028 columns.push(Arc::new(UInt64Array::from(field_values)) as ArrayRef);
2029
2030 let timestamps: Vec<i64> = (start..end).map(|v| v as i64).collect();
2032 columns.push(Arc::new(TimestampMillisecondArray::from(timestamps)) as ArrayRef);
2033
2034 let table_id = 1u32; let tsid = 100u64; let pk = new_sparse_primary_key(tags, metadata, table_id, tsid);
2038
2039 let mut pk_builder = BinaryDictionaryBuilder::<UInt32Type>::new();
2040 for _ in 0..num_rows {
2041 pk_builder.append(&pk).unwrap();
2042 }
2043 columns.push(Arc::new(pk_builder.finish()) as ArrayRef);
2044
2045 columns.push(Arc::new(UInt64Array::from_value(1000, num_rows)) as ArrayRef);
2047
2048 columns.push(Arc::new(UInt8Array::from_value(OpType::Put as u8, num_rows)) as ArrayRef);
2050
2051 RecordBatch::try_new(flat_schema, columns).unwrap()
2052 }
2053
2054 fn create_test_indexer_builder(
2056 env: &TestEnv,
2057 object_store: ObjectStore,
2058 file_path: RegionFilePathFactory,
2059 metadata: Arc<RegionMetadata>,
2060 ) -> IndexerBuilderImpl {
2061 let puffin_manager = env.get_puffin_manager().build(object_store, file_path);
2062 let intermediate_manager = env.get_intermediate_manager();
2063
2064 IndexerBuilderImpl {
2065 build_type: IndexBuildType::Flush,
2066 metadata,
2067 puffin_manager,
2068 write_cache_enabled: false,
2069 intermediate_manager,
2070 index_options: IndexOptions {
2071 inverted_index: InvertedIndexOptions {
2072 segment_row_count: 1,
2073 ..Default::default()
2074 },
2075 },
2076 inverted_index_config: Default::default(),
2077 fulltext_index_config: Default::default(),
2078 bloom_filter_index_config: Default::default(),
2079 }
2080 }
2081
2082 async fn write_flat_sst(
2084 object_store: ObjectStore,
2085 metadata: Arc<RegionMetadata>,
2086 indexer_builder: IndexerBuilderImpl,
2087 file_path: RegionFilePathFactory,
2088 flat_source: FlatSource,
2089 write_opts: &WriteOptions,
2090 ) -> SstInfo {
2091 let mut metrics = Metrics::new(WriteType::Flush);
2092 let mut writer = ParquetWriter::new_with_object_store(
2093 object_store,
2094 metadata,
2095 IndexConfig::default(),
2096 indexer_builder,
2097 file_path,
2098 &mut metrics,
2099 )
2100 .await;
2101
2102 writer
2103 .write_all_flat(flat_source, None, write_opts)
2104 .await
2105 .unwrap()
2106 .remove(0)
2107 }
2108
2109 fn create_file_handle_from_sst_info(
2111 info: &SstInfo,
2112 metadata: &Arc<RegionMetadata>,
2113 ) -> FileHandle {
2114 FileHandle::new(
2115 FileMeta {
2116 region_id: metadata.region_id,
2117 file_id: info.file_id,
2118 time_range: info.time_range,
2119 level: 0,
2120 file_size: info.file_size,
2121 max_row_group_uncompressed_size: info.max_row_group_uncompressed_size,
2122 available_indexes: info.index_metadata.build_available_indexes(),
2123 indexes: info.index_metadata.build_indexes(),
2124 index_file_size: info.index_metadata.file_size,
2125 index_version: 0,
2126 num_row_groups: info.num_row_groups,
2127 num_rows: info.num_rows as u64,
2128 sequence: None,
2129 partition_expr: match &metadata.partition_expr {
2130 Some(json_str) => partition::expr::PartitionExpr::from_json_str(json_str)
2131 .expect("partition expression should be valid JSON"),
2132 None => None,
2133 },
2134 num_series: 0,
2135 ..Default::default()
2136 },
2137 Arc::new(NoopFilePurger),
2138 )
2139 }
2140
2141 fn create_test_cache() -> Arc<CacheManager> {
2143 Arc::new(
2144 CacheManager::builder()
2145 .index_result_cache_size(1024 * 1024)
2146 .index_metadata_size(1024 * 1024)
2147 .index_content_page_size(1024 * 1024)
2148 .index_content_size(1024 * 1024)
2149 .puffin_metadata_size(1024 * 1024)
2150 .build(),
2151 )
2152 }
2153
2154 #[tokio::test]
2155 async fn test_write_flat_with_index() {
2156 let mut env = TestEnv::new().await;
2157 let object_store = env.init_object_store_manager();
2158 let file_path = RegionFilePathFactory::new(FILE_DIR.to_string(), PathType::Bare);
2159 let metadata = Arc::new(sst_region_metadata());
2160 let row_group_size = 50;
2161
2162 let flat_batches = vec![
2164 new_record_batch_by_range(&["a", "d"], 0, 20),
2165 new_record_batch_by_range(&["b", "d"], 0, 20),
2166 new_record_batch_by_range(&["c", "d"], 0, 20),
2167 new_record_batch_by_range(&["c", "f"], 0, 40),
2168 new_record_batch_by_range(&["c", "h"], 100, 200),
2169 ];
2170
2171 let flat_source = new_flat_source_from_record_batches(flat_batches);
2172
2173 let write_opts = WriteOptions {
2174 row_group_size,
2175 ..Default::default()
2176 };
2177
2178 let puffin_manager = env
2179 .get_puffin_manager()
2180 .build(object_store.clone(), file_path.clone());
2181 let intermediate_manager = env.get_intermediate_manager();
2182
2183 let indexer_builder = IndexerBuilderImpl {
2184 build_type: IndexBuildType::Flush,
2185 metadata: metadata.clone(),
2186 puffin_manager,
2187 write_cache_enabled: false,
2188 intermediate_manager,
2189 index_options: IndexOptions {
2190 inverted_index: InvertedIndexOptions {
2191 segment_row_count: 1,
2192 ..Default::default()
2193 },
2194 },
2195 inverted_index_config: Default::default(),
2196 fulltext_index_config: Default::default(),
2197 bloom_filter_index_config: Default::default(),
2198 };
2199
2200 let mut metrics = Metrics::new(WriteType::Flush);
2201 let mut writer = ParquetWriter::new_with_object_store(
2202 object_store.clone(),
2203 metadata.clone(),
2204 IndexConfig::default(),
2205 indexer_builder,
2206 file_path.clone(),
2207 &mut metrics,
2208 )
2209 .await;
2210
2211 let info = writer
2212 .write_all_flat(flat_source, None, &write_opts)
2213 .await
2214 .unwrap()
2215 .remove(0);
2216 assert_eq!(200, info.num_rows);
2217 assert!(info.file_size > 0);
2218 assert!(info.index_metadata.file_size > 0);
2219
2220 assert!(info.index_metadata.inverted_index.index_size > 0);
2221 assert_eq!(info.index_metadata.inverted_index.row_count, 200);
2222 assert_eq!(info.index_metadata.inverted_index.columns, vec![0]);
2223
2224 assert!(info.index_metadata.bloom_filter.index_size > 0);
2225 assert_eq!(info.index_metadata.bloom_filter.row_count, 200);
2226 assert_eq!(info.index_metadata.bloom_filter.columns, vec![1]);
2227
2228 assert_eq!(
2229 (
2230 Timestamp::new_millisecond(0),
2231 Timestamp::new_millisecond(199)
2232 ),
2233 info.time_range
2234 );
2235 }
2236
2237 #[tokio::test]
2238 async fn test_read_with_override_sequence() {
2239 test_read_with_override_sequence_with_format(false).await;
2240 test_read_with_override_sequence_with_format(true).await;
2241 }
2242
2243 async fn test_read_with_override_sequence_with_format(flat_format: bool) {
2244 let mut env = TestEnv::new().await;
2245 let object_store = env.init_object_store_manager();
2246 let metadata = Arc::new(sst_region_metadata());
2247
2248 async fn read_sequences(builder: ParquetReaderBuilder) -> Vec<u64> {
2249 let mut reader = builder.build().await.unwrap().unwrap();
2250 let mut sequences = Vec::new();
2251 while let Some(batch) = reader.next_record_batch().await.unwrap() {
2252 let sequence = batch
2253 .column(batch.num_columns() - 2)
2254 .as_primitive::<datatypes::arrow::datatypes::UInt64Type>();
2255 sequences.extend((0..sequence.len()).map(|idx| sequence.value(idx)));
2256 }
2257 sequences
2258 }
2259
2260 async fn write_sst(
2261 object_store: ObjectStore,
2262 metadata: Arc<RegionMetadata>,
2263 handle: FileHandle,
2264 flat_format: bool,
2265 sequence: u64,
2266 ) {
2267 let file_path = FixedPathProvider {
2268 region_file_id: handle.file_id(),
2269 };
2270 let source = new_flat_source_from_record_batches(vec![
2271 new_record_batch_with_custom_sequence(&["a", "d"], 0, 60, sequence),
2272 new_record_batch_with_custom_sequence(&["b", "f"], 0, 40, sequence),
2273 ]);
2274 let write_opts = WriteOptions {
2275 row_group_size: 50,
2276 ..Default::default()
2277 };
2278 let mut metrics = Metrics::new(WriteType::Flush);
2279 let mut writer = ParquetWriter::new_with_object_store(
2280 object_store,
2281 metadata,
2282 IndexConfig::default(),
2283 NoopIndexBuilder,
2284 file_path,
2285 &mut metrics,
2286 )
2287 .await;
2288 if flat_format {
2289 writer
2290 .write_all_flat(source, None, &write_opts)
2291 .await
2292 .unwrap();
2293 } else {
2294 writer
2295 .write_all_flat_as_primary_key(source, None, &write_opts)
2296 .await
2297 .unwrap();
2298 }
2299 }
2300
2301 fn handle_with_meta(
2302 handle: &FileHandle,
2303 sequence: Option<u64>,
2304 preserve_row_sequence: bool,
2305 ) -> FileHandle {
2306 let mut file_meta = handle.meta_ref().clone();
2307 file_meta.sequence = sequence.and_then(std::num::NonZeroU64::new);
2308 file_meta.preserve_row_sequence = preserve_row_sequence;
2309 FileHandle::new(file_meta, Arc::new(NoopFilePurger))
2310 }
2311
2312 let custom_sequence = 12345;
2313 let local_zero_handle = sst_file_handle(0, 1000);
2314 let local_nonzero_handle = sst_file_handle(0, 1000);
2315 write_sst(
2316 object_store.clone(),
2317 metadata.clone(),
2318 local_zero_handle.clone(),
2319 flat_format,
2320 0,
2321 )
2322 .await;
2323 write_sst(
2324 object_store.clone(),
2325 metadata.clone(),
2326 local_nonzero_handle.clone(),
2327 flat_format,
2328 7,
2329 )
2330 .await;
2331
2332 let local_zero_none = read_sequences(
2333 ParquetReaderBuilder::new(
2334 FILE_DIR.to_string(),
2335 PathType::Bare,
2336 local_zero_handle.clone(),
2337 object_store.clone(),
2338 )
2339 .expected_metadata(Some(metadata.clone())),
2340 )
2341 .await;
2342 assert!(local_zero_none.iter().all(|sequence| *sequence == 0));
2343
2344 let local_zero_override = read_sequences(
2346 ParquetReaderBuilder::new(
2347 FILE_DIR.to_string(),
2348 PathType::Bare,
2349 handle_with_meta(&local_zero_handle, Some(custom_sequence), false),
2350 object_store.clone(),
2351 )
2352 .expected_metadata(Some(metadata.clone())),
2353 )
2354 .await;
2355 assert!(
2356 local_zero_override
2357 .iter()
2358 .all(|sequence| *sequence == custom_sequence)
2359 );
2360
2361 let local_nonzero_override = read_sequences(
2363 ParquetReaderBuilder::new(
2364 FILE_DIR.to_string(),
2365 PathType::Bare,
2366 handle_with_meta(&local_nonzero_handle, Some(custom_sequence), false),
2367 object_store.clone(),
2368 )
2369 .expected_metadata(Some(metadata.clone())),
2370 )
2371 .await;
2372 assert!(local_nonzero_override.iter().all(|sequence| *sequence == 7));
2373
2374 let local_nonzero_none = read_sequences(
2375 ParquetReaderBuilder::new(
2376 FILE_DIR.to_string(),
2377 PathType::Bare,
2378 local_nonzero_handle.clone(),
2379 object_store.clone(),
2380 )
2381 .expected_metadata(Some(metadata.clone())),
2382 )
2383 .await;
2384 assert!(local_nonzero_none.iter().all(|sequence| *sequence == 7));
2385
2386 let local_trusted_nonzero = read_sequences(
2388 ParquetReaderBuilder::new(
2389 FILE_DIR.to_string(),
2390 PathType::Bare,
2391 handle_with_meta(&local_nonzero_handle, Some(custom_sequence), true),
2392 object_store.clone(),
2393 )
2394 .expected_metadata(Some(metadata.clone())),
2395 )
2396 .await;
2397 assert!(local_trusted_nonzero.iter().all(|sequence| *sequence == 7));
2398
2399 let local_trusted_zero = read_sequences(
2400 ParquetReaderBuilder::new(
2401 FILE_DIR.to_string(),
2402 PathType::Bare,
2403 handle_with_meta(&local_zero_handle, Some(custom_sequence), true),
2404 object_store.clone(),
2405 )
2406 .expected_metadata(Some(metadata.clone())),
2407 )
2408 .await;
2409 assert!(local_trusted_zero.iter().all(|sequence| *sequence == 0));
2410
2411 let mut target_metadata = (*metadata).clone();
2412 target_metadata.region_id = RegionId::new(0, 1);
2413 let target_metadata = Arc::new(target_metadata);
2414
2415 let foreign_marked = read_sequences(
2417 ParquetReaderBuilder::new(
2418 FILE_DIR.to_string(),
2419 PathType::Bare,
2420 handle_with_meta(&local_nonzero_handle, Some(custom_sequence), true),
2421 object_store.clone(),
2422 )
2423 .expected_metadata(Some(target_metadata.clone())),
2424 )
2425 .await;
2426 assert!(
2427 foreign_marked
2428 .iter()
2429 .all(|sequence| *sequence == custom_sequence)
2430 );
2431
2432 let foreign_unmarked = read_sequences(
2433 ParquetReaderBuilder::new(
2434 FILE_DIR.to_string(),
2435 PathType::Bare,
2436 handle_with_meta(&local_nonzero_handle, Some(custom_sequence), false),
2437 object_store.clone(),
2438 )
2439 .expected_metadata(Some(target_metadata.clone())),
2440 )
2441 .await;
2442 assert!(
2443 foreign_unmarked
2444 .iter()
2445 .all(|sequence| *sequence == custom_sequence)
2446 );
2447
2448 let foreign_none = read_sequences(
2450 ParquetReaderBuilder::new(
2451 FILE_DIR.to_string(),
2452 PathType::Bare,
2453 handle_with_meta(&local_nonzero_handle, None, false),
2454 object_store,
2455 )
2456 .expected_metadata(Some(target_metadata)),
2457 )
2458 .await;
2459 assert!(foreign_none.iter().all(|sequence| *sequence == 7));
2460 }
2461
2462 #[tokio::test]
2463 async fn test_write_flat_read_with_inverted_index() {
2464 let mut env = TestEnv::new().await;
2465 let object_store = env.init_object_store_manager();
2466 let file_path = RegionFilePathFactory::new(FILE_DIR.to_string(), PathType::Bare);
2467 let metadata = Arc::new(sst_region_metadata());
2468 let row_group_size = 100;
2469
2470 let flat_batches = vec![
2478 new_record_batch_by_range(&["a", "d"], 0, 50),
2479 new_record_batch_by_range(&["b", "d"], 50, 100),
2480 new_record_batch_by_range(&["c", "d"], 100, 150),
2481 new_record_batch_by_range(&["c", "f"], 150, 200),
2482 ];
2483
2484 let flat_source = new_flat_source_from_record_batches(flat_batches);
2485
2486 let write_opts = WriteOptions {
2487 row_group_size,
2488 ..Default::default()
2489 };
2490
2491 let indexer_builder = create_test_indexer_builder(
2492 &env,
2493 object_store.clone(),
2494 file_path.clone(),
2495 metadata.clone(),
2496 );
2497
2498 let info = write_flat_sst(
2499 object_store.clone(),
2500 metadata.clone(),
2501 indexer_builder,
2502 file_path.clone(),
2503 flat_source,
2504 &write_opts,
2505 )
2506 .await;
2507 assert_eq!(200, info.num_rows);
2508 assert!(info.file_size > 0);
2509 assert!(info.index_metadata.file_size > 0);
2510
2511 let handle = create_file_handle_from_sst_info(&info, &metadata);
2512
2513 let cache = create_test_cache();
2514
2515 let preds = vec![col("tag_0").eq(lit("b"))];
2518 let inverted_index_applier = InvertedIndexApplierBuilder::new(
2519 FILE_DIR.to_string(),
2520 PathType::Bare,
2521 object_store.clone(),
2522 &metadata,
2523 HashSet::from_iter([0]),
2524 env.get_puffin_manager(),
2525 )
2526 .with_puffin_metadata_cache(cache.puffin_metadata_cache().cloned())
2527 .with_inverted_index_cache(cache.inverted_index_cache().cloned())
2528 .build(&preds)
2529 .unwrap()
2530 .map(Arc::new);
2531
2532 let builder = ParquetReaderBuilder::new(
2533 FILE_DIR.to_string(),
2534 PathType::Bare,
2535 handle.clone(),
2536 object_store.clone(),
2537 )
2538 .predicate(Some(Predicate::new(preds)))
2539 .inverted_index_appliers([inverted_index_applier.clone(), None])
2540 .cache(CacheStrategy::EnableAll(cache.clone()));
2541
2542 let mut metrics = ReaderMetrics::default();
2543 let (_context, selection) = builder
2544 .build_reader_input(&mut metrics)
2545 .await
2546 .unwrap()
2547 .unwrap();
2548
2549 assert_eq!(selection.row_group_count(), 1);
2551 assert_eq!(50, selection.get(0).unwrap().row_count());
2552
2553 assert_eq!(metrics.filter_metrics.rg_total, 2);
2555 assert_eq!(metrics.filter_metrics.rg_minmax_filtered, 1);
2556 assert_eq!(metrics.filter_metrics.rg_inverted_filtered, 0);
2557 assert_eq!(metrics.filter_metrics.rows_inverted_filtered, 50);
2558 }
2559
2560 #[tokio::test]
2561 async fn test_write_flat_read_with_bloom_filter() {
2562 let mut env = TestEnv::new().await;
2563 let object_store = env.init_object_store_manager();
2564 let file_path = RegionFilePathFactory::new(FILE_DIR.to_string(), PathType::Bare);
2565 let metadata = Arc::new(sst_region_metadata());
2566 let row_group_size = 100;
2567
2568 let flat_batches = vec![
2576 new_record_batch_by_range(&["a", "d"], 0, 50),
2577 new_record_batch_by_range(&["b", "e"], 50, 100),
2578 new_record_batch_by_range(&["c", "d"], 100, 150),
2579 new_record_batch_by_range(&["c", "f"], 150, 200),
2580 ];
2581
2582 let flat_source = new_flat_source_from_record_batches(flat_batches);
2583
2584 let write_opts = WriteOptions {
2585 row_group_size,
2586 ..Default::default()
2587 };
2588
2589 let indexer_builder = create_test_indexer_builder(
2590 &env,
2591 object_store.clone(),
2592 file_path.clone(),
2593 metadata.clone(),
2594 );
2595
2596 let info = write_flat_sst(
2597 object_store.clone(),
2598 metadata.clone(),
2599 indexer_builder,
2600 file_path.clone(),
2601 flat_source,
2602 &write_opts,
2603 )
2604 .await;
2605 assert_eq!(200, info.num_rows);
2606 assert!(info.file_size > 0);
2607 assert!(info.index_metadata.file_size > 0);
2608
2609 let handle = create_file_handle_from_sst_info(&info, &metadata);
2610
2611 let cache = create_test_cache();
2612
2613 let preds = vec![
2616 col("ts").gt_eq(lit(ScalarValue::TimestampMillisecond(Some(50), None))),
2617 col("ts").lt(lit(ScalarValue::TimestampMillisecond(Some(200), None))),
2618 col("tag_1").eq(lit("d")),
2619 ];
2620 let bloom_filter_applier = BloomFilterIndexApplierBuilder::new(
2621 FILE_DIR.to_string(),
2622 PathType::Bare,
2623 object_store.clone(),
2624 &metadata,
2625 env.get_puffin_manager(),
2626 )
2627 .with_puffin_metadata_cache(cache.puffin_metadata_cache().cloned())
2628 .with_bloom_filter_index_cache(cache.bloom_filter_index_cache().cloned())
2629 .build(&preds)
2630 .unwrap()
2631 .map(Arc::new);
2632
2633 let builder = ParquetReaderBuilder::new(
2634 FILE_DIR.to_string(),
2635 PathType::Bare,
2636 handle.clone(),
2637 object_store.clone(),
2638 )
2639 .predicate(Some(Predicate::new(preds)))
2640 .bloom_filter_index_appliers([None, bloom_filter_applier.clone()])
2641 .cache(CacheStrategy::EnableAll(cache.clone()));
2642
2643 let mut metrics = ReaderMetrics::default();
2644 let (_context, selection) = builder
2645 .build_reader_input(&mut metrics)
2646 .await
2647 .unwrap()
2648 .unwrap();
2649
2650 assert_eq!(selection.row_group_count(), 2);
2652 assert_eq!(50, selection.get(0).unwrap().row_count());
2653 assert_eq!(50, selection.get(1).unwrap().row_count());
2654
2655 assert_eq!(metrics.filter_metrics.rg_total, 2);
2657 assert_eq!(metrics.filter_metrics.rg_minmax_filtered, 0);
2658 assert_eq!(metrics.filter_metrics.rg_bloom_filtered, 0);
2659 assert_eq!(metrics.filter_metrics.rows_bloom_filtered, 100);
2660 }
2661
2662 #[tokio::test]
2663 async fn test_reader_prefilter_with_outer_selection_and_trailing_filtered_rows() {
2664 let mut env = TestEnv::new().await;
2665 let object_store = env.init_object_store_manager();
2666 let file_path = RegionFilePathFactory::new(FILE_DIR.to_string(), PathType::Bare);
2667 let metadata = Arc::new(sst_region_metadata());
2668 let row_group_size = 10;
2669
2670 let flat_source = new_flat_source_from_record_batches(vec![
2671 new_record_batch_by_range(&["a", "d"], 0, 3),
2672 new_record_batch_by_range(&["b", "d"], 3, 10),
2673 ]);
2674 let write_opts = WriteOptions {
2675 row_group_size,
2676 ..Default::default()
2677 };
2678 let indexer_builder = create_test_indexer_builder(
2679 &env,
2680 object_store.clone(),
2681 file_path.clone(),
2682 metadata.clone(),
2683 );
2684 let info = write_flat_sst(
2685 object_store.clone(),
2686 metadata.clone(),
2687 indexer_builder,
2688 file_path,
2689 flat_source,
2690 &write_opts,
2691 )
2692 .await;
2693 let handle = create_file_handle_from_sst_info(&info, &metadata);
2694
2695 let builder =
2696 ParquetReaderBuilder::new(FILE_DIR.to_string(), PathType::Bare, handle, object_store)
2697 .predicate(Some(Predicate::new(vec![col("tag_0").eq(lit("a"))])));
2698
2699 let mut metrics = ReaderMetrics::default();
2700 let (context, _) = builder
2701 .build_reader_input(&mut metrics)
2702 .await
2703 .unwrap()
2704 .unwrap();
2705 let selection = RowGroupSelection::from_row_ranges(
2706 vec![(0, std::iter::once(0..6).collect())],
2707 row_group_size,
2708 );
2709
2710 let mut reader = ParquetReader::new(Arc::new(context), selection)
2711 .await
2712 .unwrap();
2713 check_record_batch_reader_result(
2714 &mut reader,
2715 &[new_record_batch_by_range(&["a", "d"], 0, 3)],
2716 )
2717 .await;
2718 }
2719
2720 #[tokio::test]
2721 async fn test_reader_prefilter_with_outer_selection_disjoint_matches_and_trailing_gap() {
2722 let mut env = TestEnv::new().await;
2723 let object_store = env.init_object_store_manager();
2724 let file_path = RegionFilePathFactory::new(FILE_DIR.to_string(), PathType::Bare);
2725 let metadata = Arc::new(sst_region_metadata());
2726 let row_group_size = 8;
2727
2728 let flat_source = new_flat_source_from_record_batches(vec![
2729 new_record_batch_by_range(&["a", "d"], 0, 2),
2730 new_record_batch_by_range(&["b", "d"], 2, 4),
2731 new_record_batch_by_range(&["a", "d"], 4, 6),
2732 new_record_batch_by_range(&["c", "d"], 6, 8),
2733 ]);
2734 let write_opts = WriteOptions {
2735 row_group_size,
2736 ..Default::default()
2737 };
2738 let indexer_builder = create_test_indexer_builder(
2739 &env,
2740 object_store.clone(),
2741 file_path.clone(),
2742 metadata.clone(),
2743 );
2744 let info = write_flat_sst(
2745 object_store.clone(),
2746 metadata.clone(),
2747 indexer_builder,
2748 file_path,
2749 flat_source,
2750 &write_opts,
2751 )
2752 .await;
2753 let handle = create_file_handle_from_sst_info(&info, &metadata);
2754
2755 let builder =
2756 ParquetReaderBuilder::new(FILE_DIR.to_string(), PathType::Bare, handle, object_store)
2757 .predicate(Some(Predicate::new(vec![col("tag_0").eq(lit("a"))])));
2758
2759 let mut metrics = ReaderMetrics::default();
2760 let (context, _) = builder
2761 .build_reader_input(&mut metrics)
2762 .await
2763 .unwrap()
2764 .unwrap();
2765 let selection = RowGroupSelection::from_row_ranges(
2766 vec![(0, std::iter::once(0..8).collect())],
2767 row_group_size,
2768 );
2769
2770 let mut reader = ParquetReader::new(Arc::new(context), selection)
2771 .await
2772 .unwrap();
2773 check_record_batch_reader_result(
2774 &mut reader,
2775 &[new_record_batch_from_rows(&[
2776 ("a", "d", 0),
2777 ("a", "d", 1),
2778 ("a", "d", 4),
2779 ("a", "d", 5),
2780 ])],
2781 )
2782 .await;
2783 }
2784
2785 #[tokio::test]
2786 async fn test_write_flat_read_with_inverted_index_sparse() {
2787 common_telemetry::init_default_ut_logging();
2788
2789 let mut env = TestEnv::new().await;
2790 let object_store = env.init_object_store_manager();
2791 let file_path = RegionFilePathFactory::new(FILE_DIR.to_string(), PathType::Bare);
2792 let metadata = Arc::new(sst_region_metadata_with_encoding(
2793 PrimaryKeyEncoding::Sparse,
2794 ));
2795 let row_group_size = 100;
2796
2797 let flat_batches = vec![
2805 new_record_batch_by_range_sparse(&["a", "d"], 0, 50, &metadata),
2806 new_record_batch_by_range_sparse(&["b", "d"], 50, 100, &metadata),
2807 new_record_batch_by_range_sparse(&["c", "d"], 100, 150, &metadata),
2808 new_record_batch_by_range_sparse(&["c", "f"], 150, 200, &metadata),
2809 ];
2810
2811 let flat_source = new_flat_source_from_record_batches(flat_batches);
2812
2813 let write_opts = WriteOptions {
2814 row_group_size,
2815 ..Default::default()
2816 };
2817
2818 let indexer_builder = create_test_indexer_builder(
2819 &env,
2820 object_store.clone(),
2821 file_path.clone(),
2822 metadata.clone(),
2823 );
2824
2825 let info = write_flat_sst(
2826 object_store.clone(),
2827 metadata.clone(),
2828 indexer_builder,
2829 file_path.clone(),
2830 flat_source,
2831 &write_opts,
2832 )
2833 .await;
2834 assert_eq!(200, info.num_rows);
2835 assert!(info.file_size > 0);
2836 assert!(info.index_metadata.file_size > 0);
2837
2838 let handle = create_file_handle_from_sst_info(&info, &metadata);
2839
2840 let cache = create_test_cache();
2841
2842 let preds = vec![col("tag_0").eq(lit("b"))];
2845 let inverted_index_applier = InvertedIndexApplierBuilder::new(
2846 FILE_DIR.to_string(),
2847 PathType::Bare,
2848 object_store.clone(),
2849 &metadata,
2850 HashSet::from_iter([0]),
2851 env.get_puffin_manager(),
2852 )
2853 .with_puffin_metadata_cache(cache.puffin_metadata_cache().cloned())
2854 .with_inverted_index_cache(cache.inverted_index_cache().cloned())
2855 .build(&preds)
2856 .unwrap()
2857 .map(Arc::new);
2858
2859 let builder = ParquetReaderBuilder::new(
2860 FILE_DIR.to_string(),
2861 PathType::Bare,
2862 handle.clone(),
2863 object_store.clone(),
2864 )
2865 .predicate(Some(Predicate::new(preds)))
2866 .inverted_index_appliers([inverted_index_applier.clone(), None])
2867 .cache(CacheStrategy::EnableAll(cache.clone()));
2868
2869 let mut metrics = ReaderMetrics::default();
2870 let (_context, selection) = builder
2871 .build_reader_input(&mut metrics)
2872 .await
2873 .unwrap()
2874 .unwrap();
2875
2876 assert_eq!(selection.row_group_count(), 1);
2878 assert_eq!(50, selection.get(0).unwrap().row_count());
2879
2880 assert_eq!(metrics.filter_metrics.rg_total, 2);
2884 assert_eq!(metrics.filter_metrics.rg_minmax_filtered, 0); assert_eq!(metrics.filter_metrics.rg_inverted_filtered, 1);
2886 assert_eq!(metrics.filter_metrics.rows_inverted_filtered, 150);
2887 }
2888
2889 #[tokio::test]
2890 async fn test_write_flat_read_with_bloom_filter_sparse() {
2891 let mut env = TestEnv::new().await;
2892 let object_store = env.init_object_store_manager();
2893 let file_path = RegionFilePathFactory::new(FILE_DIR.to_string(), PathType::Bare);
2894 let metadata = Arc::new(sst_region_metadata_with_encoding(
2895 PrimaryKeyEncoding::Sparse,
2896 ));
2897 let row_group_size = 100;
2898
2899 let flat_batches = vec![
2907 new_record_batch_by_range_sparse(&["a", "d"], 0, 50, &metadata),
2908 new_record_batch_by_range_sparse(&["b", "e"], 50, 100, &metadata),
2909 new_record_batch_by_range_sparse(&["c", "d"], 100, 150, &metadata),
2910 new_record_batch_by_range_sparse(&["c", "f"], 150, 200, &metadata),
2911 ];
2912
2913 let flat_source = new_flat_source_from_record_batches(flat_batches);
2914
2915 let write_opts = WriteOptions {
2916 row_group_size,
2917 ..Default::default()
2918 };
2919
2920 let indexer_builder = create_test_indexer_builder(
2921 &env,
2922 object_store.clone(),
2923 file_path.clone(),
2924 metadata.clone(),
2925 );
2926
2927 let info = write_flat_sst(
2928 object_store.clone(),
2929 metadata.clone(),
2930 indexer_builder,
2931 file_path.clone(),
2932 flat_source,
2933 &write_opts,
2934 )
2935 .await;
2936 assert_eq!(200, info.num_rows);
2937 assert!(info.file_size > 0);
2938 assert!(info.index_metadata.file_size > 0);
2939
2940 let handle = create_file_handle_from_sst_info(&info, &metadata);
2941
2942 let cache = create_test_cache();
2943
2944 let preds = vec![
2947 col("ts").gt_eq(lit(ScalarValue::TimestampMillisecond(Some(50), None))),
2948 col("ts").lt(lit(ScalarValue::TimestampMillisecond(Some(200), None))),
2949 col("tag_1").eq(lit("d")),
2950 ];
2951 let bloom_filter_applier = BloomFilterIndexApplierBuilder::new(
2952 FILE_DIR.to_string(),
2953 PathType::Bare,
2954 object_store.clone(),
2955 &metadata,
2956 env.get_puffin_manager(),
2957 )
2958 .with_puffin_metadata_cache(cache.puffin_metadata_cache().cloned())
2959 .with_bloom_filter_index_cache(cache.bloom_filter_index_cache().cloned())
2960 .build(&preds)
2961 .unwrap()
2962 .map(Arc::new);
2963
2964 let builder = ParquetReaderBuilder::new(
2965 FILE_DIR.to_string(),
2966 PathType::Bare,
2967 handle.clone(),
2968 object_store.clone(),
2969 )
2970 .predicate(Some(Predicate::new(preds)))
2971 .bloom_filter_index_appliers([None, bloom_filter_applier.clone()])
2972 .cache(CacheStrategy::EnableAll(cache.clone()));
2973
2974 let mut metrics = ReaderMetrics::default();
2975 let (_context, selection) = builder
2976 .build_reader_input(&mut metrics)
2977 .await
2978 .unwrap()
2979 .unwrap();
2980
2981 assert_eq!(selection.row_group_count(), 2);
2983 assert_eq!(50, selection.get(0).unwrap().row_count());
2984 assert_eq!(50, selection.get(1).unwrap().row_count());
2985
2986 assert_eq!(metrics.filter_metrics.rg_total, 2);
2988 assert_eq!(metrics.filter_metrics.rg_minmax_filtered, 0);
2989 assert_eq!(metrics.filter_metrics.rg_bloom_filtered, 0);
2990 assert_eq!(metrics.filter_metrics.rows_bloom_filtered, 100);
2991 }
2992
2993 fn fulltext_region_metadata() -> RegionMetadata {
2996 let mut builder = RegionMetadataBuilder::new(REGION_ID);
2997 builder
2998 .push_column_metadata(ColumnMetadata {
2999 column_schema: ColumnSchema::new(
3000 "tag_0".to_string(),
3001 ConcreteDataType::string_datatype(),
3002 true,
3003 ),
3004 semantic_type: SemanticType::Tag,
3005 column_id: 0,
3006 })
3007 .push_column_metadata(ColumnMetadata {
3008 column_schema: ColumnSchema::new(
3009 "text_bloom".to_string(),
3010 ConcreteDataType::string_datatype(),
3011 true,
3012 )
3013 .with_fulltext_options(FulltextOptions {
3014 enable: true,
3015 analyzer: FulltextAnalyzer::English,
3016 case_sensitive: false,
3017 backend: FulltextBackend::Bloom,
3018 granularity: 1,
3019 false_positive_rate_in_10000: 50,
3020 })
3021 .unwrap(),
3022 semantic_type: SemanticType::Field,
3023 column_id: 1,
3024 })
3025 .push_column_metadata(ColumnMetadata {
3026 column_schema: ColumnSchema::new(
3027 "text_tantivy".to_string(),
3028 ConcreteDataType::string_datatype(),
3029 true,
3030 )
3031 .with_fulltext_options(FulltextOptions {
3032 enable: true,
3033 analyzer: FulltextAnalyzer::English,
3034 case_sensitive: false,
3035 backend: FulltextBackend::Tantivy,
3036 granularity: 1,
3037 false_positive_rate_in_10000: 50,
3038 })
3039 .unwrap(),
3040 semantic_type: SemanticType::Field,
3041 column_id: 2,
3042 })
3043 .push_column_metadata(ColumnMetadata {
3044 column_schema: ColumnSchema::new(
3045 "field_0".to_string(),
3046 ConcreteDataType::uint64_datatype(),
3047 true,
3048 ),
3049 semantic_type: SemanticType::Field,
3050 column_id: 3,
3051 })
3052 .push_column_metadata(ColumnMetadata {
3053 column_schema: ColumnSchema::new(
3054 "ts".to_string(),
3055 ConcreteDataType::timestamp_millisecond_datatype(),
3056 false,
3057 ),
3058 semantic_type: SemanticType::Timestamp,
3059 column_id: 4,
3060 })
3061 .primary_key(vec![0]);
3062 builder.build().unwrap()
3063 }
3064
3065 fn new_fulltext_record_batch_by_range(
3067 tag: &str,
3068 text_bloom: &str,
3069 text_tantivy: &str,
3070 start: usize,
3071 end: usize,
3072 ) -> RecordBatch {
3073 assert!(end >= start);
3074 let metadata = Arc::new(fulltext_region_metadata());
3075 let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default());
3076
3077 let num_rows = end - start;
3078 let mut columns = Vec::new();
3079
3080 let mut tag_builder = StringDictionaryBuilder::<UInt32Type>::new();
3082 for _ in 0..num_rows {
3083 tag_builder.append_value(tag);
3084 }
3085 columns.push(Arc::new(tag_builder.finish()) as ArrayRef);
3086
3087 let text_bloom_values: Vec<_> = (0..num_rows).map(|_| text_bloom).collect();
3089 columns.push(Arc::new(StringArray::from(text_bloom_values)));
3090
3091 let text_tantivy_values: Vec<_> = (0..num_rows).map(|_| text_tantivy).collect();
3093 columns.push(Arc::new(StringArray::from(text_tantivy_values)));
3094
3095 let field_values: Vec<u64> = (start..end).map(|v| v as u64).collect();
3097 columns.push(Arc::new(UInt64Array::from(field_values)));
3098
3099 let timestamps: Vec<i64> = (start..end).map(|v| v as i64).collect();
3101 columns.push(Arc::new(TimestampMillisecondArray::from(timestamps)));
3102
3103 let pk = new_primary_key(&[tag]);
3105 let mut pk_builder = BinaryDictionaryBuilder::<UInt32Type>::new();
3106 for _ in 0..num_rows {
3107 pk_builder.append(&pk).unwrap();
3108 }
3109 columns.push(Arc::new(pk_builder.finish()));
3110
3111 columns.push(Arc::new(UInt64Array::from_value(1000, num_rows)));
3113
3114 columns.push(Arc::new(UInt8Array::from_value(
3116 OpType::Put as u8,
3117 num_rows,
3118 )));
3119
3120 RecordBatch::try_new(flat_schema, columns).unwrap()
3121 }
3122
3123 #[tokio::test]
3124 async fn test_write_flat_read_with_fulltext_index() {
3125 let mut env = TestEnv::new().await;
3126 let object_store = env.init_object_store_manager();
3127 let file_path = RegionFilePathFactory::new(FILE_DIR.to_string(), PathType::Bare);
3128 let metadata = Arc::new(fulltext_region_metadata());
3129 let row_group_size = 50;
3130
3131 let flat_batches = vec![
3137 new_fulltext_record_batch_by_range("a", "hello world", "quick brown fox", 0, 50),
3138 new_fulltext_record_batch_by_range("b", "hello world", "quick brown fox", 50, 100),
3139 new_fulltext_record_batch_by_range("c", "goodbye world", "lazy dog", 100, 150),
3140 new_fulltext_record_batch_by_range("d", "goodbye world", "lazy dog", 150, 200),
3141 ];
3142
3143 let flat_source = new_flat_source_from_record_batches(flat_batches);
3144
3145 let write_opts = WriteOptions {
3146 row_group_size,
3147 ..Default::default()
3148 };
3149
3150 let indexer_builder = create_test_indexer_builder(
3151 &env,
3152 object_store.clone(),
3153 file_path.clone(),
3154 metadata.clone(),
3155 );
3156
3157 let mut info = write_flat_sst(
3158 object_store.clone(),
3159 metadata.clone(),
3160 indexer_builder,
3161 file_path.clone(),
3162 flat_source,
3163 &write_opts,
3164 )
3165 .await;
3166 assert_eq!(200, info.num_rows);
3167 assert!(info.file_size > 0);
3168 assert!(info.index_metadata.file_size > 0);
3169
3170 assert!(info.index_metadata.fulltext_index.index_size > 0);
3172 assert_eq!(info.index_metadata.fulltext_index.row_count, 200);
3173 info.index_metadata.fulltext_index.columns.sort_unstable();
3175 assert_eq!(info.index_metadata.fulltext_index.columns, vec![1, 2]);
3176
3177 assert_eq!(
3178 (
3179 Timestamp::new_millisecond(0),
3180 Timestamp::new_millisecond(199)
3181 ),
3182 info.time_range
3183 );
3184
3185 let handle = create_file_handle_from_sst_info(&info, &metadata);
3186
3187 let cache = create_test_cache();
3188
3189 let matches_func = || {
3191 Arc::new(
3192 ScalarFunctionFactory::from(Arc::new(MatchesFunction::default()) as FunctionRef)
3193 .provide(Default::default()),
3194 )
3195 };
3196
3197 let matches_term_func = || {
3198 Arc::new(
3199 ScalarFunctionFactory::from(
3200 Arc::new(MatchesTermFunction::default()) as FunctionRef,
3201 )
3202 .provide(Default::default()),
3203 )
3204 };
3205
3206 let preds = vec![Expr::ScalarFunction(ScalarFunction {
3209 args: vec![col("text_bloom"), "hello".lit()],
3210 func: matches_term_func(),
3211 })];
3212
3213 let fulltext_applier = FulltextIndexApplierBuilder::new(
3214 FILE_DIR.to_string(),
3215 PathType::Bare,
3216 object_store.clone(),
3217 env.get_puffin_manager(),
3218 &metadata,
3219 )
3220 .with_puffin_metadata_cache(cache.puffin_metadata_cache().cloned())
3221 .with_bloom_filter_cache(cache.bloom_filter_index_cache().cloned())
3222 .build(&preds)
3223 .unwrap()
3224 .map(Arc::new);
3225
3226 let builder = ParquetReaderBuilder::new(
3227 FILE_DIR.to_string(),
3228 PathType::Bare,
3229 handle.clone(),
3230 object_store.clone(),
3231 )
3232 .predicate(Some(Predicate::new(preds)))
3233 .fulltext_index_appliers([None, fulltext_applier.clone()])
3234 .cache(CacheStrategy::EnableAll(cache.clone()));
3235
3236 let mut metrics = ReaderMetrics::default();
3237 let (_context, selection) = builder
3238 .build_reader_input(&mut metrics)
3239 .await
3240 .unwrap()
3241 .unwrap();
3242
3243 assert_eq!(selection.row_group_count(), 2);
3245 assert_eq!(50, selection.get(0).unwrap().row_count());
3246 assert_eq!(50, selection.get(1).unwrap().row_count());
3247
3248 assert_eq!(metrics.filter_metrics.rg_total, 4);
3250 assert_eq!(metrics.filter_metrics.rg_minmax_filtered, 0);
3251 assert_eq!(metrics.filter_metrics.rg_fulltext_filtered, 2);
3252 assert_eq!(metrics.filter_metrics.rows_fulltext_filtered, 100);
3253
3254 let preds = vec![Expr::ScalarFunction(ScalarFunction {
3257 args: vec![col("text_tantivy"), "lazy".lit()],
3258 func: matches_func(),
3259 })];
3260
3261 let fulltext_applier = FulltextIndexApplierBuilder::new(
3262 FILE_DIR.to_string(),
3263 PathType::Bare,
3264 object_store.clone(),
3265 env.get_puffin_manager(),
3266 &metadata,
3267 )
3268 .with_puffin_metadata_cache(cache.puffin_metadata_cache().cloned())
3269 .with_bloom_filter_cache(cache.bloom_filter_index_cache().cloned())
3270 .build(&preds)
3271 .unwrap()
3272 .map(Arc::new);
3273
3274 let builder = ParquetReaderBuilder::new(
3275 FILE_DIR.to_string(),
3276 PathType::Bare,
3277 handle.clone(),
3278 object_store.clone(),
3279 )
3280 .predicate(Some(Predicate::new(preds)))
3281 .fulltext_index_appliers([None, fulltext_applier.clone()])
3282 .cache(CacheStrategy::EnableAll(cache.clone()));
3283
3284 let mut metrics = ReaderMetrics::default();
3285 let (_context, selection) = builder
3286 .build_reader_input(&mut metrics)
3287 .await
3288 .unwrap()
3289 .unwrap();
3290
3291 assert_eq!(selection.row_group_count(), 2);
3293 assert_eq!(50, selection.get(2).unwrap().row_count());
3294 assert_eq!(50, selection.get(3).unwrap().row_count());
3295
3296 assert_eq!(metrics.filter_metrics.rg_total, 4);
3298 assert_eq!(metrics.filter_metrics.rg_minmax_filtered, 0);
3299 assert_eq!(metrics.filter_metrics.rg_fulltext_filtered, 2);
3300 assert_eq!(metrics.filter_metrics.rows_fulltext_filtered, 100);
3301 }
3302}