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