Skip to main content

mito2/sst/
parquet.rs

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