Skip to main content

store_api/
sst_entry.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
15use std::sync::Arc;
16
17use bytes::Bytes;
18use common_recordbatch::DfRecordBatch;
19use common_time::Timestamp;
20use common_time::timestamp::TimeUnit;
21use datafusion_common::DataFusionError;
22use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, LogicalTableSource};
23use datatypes::arrow::array::{
24    ArrayRef, BinaryArray, BooleanArray, TimestampMillisecondArray, TimestampNanosecondArray,
25    UInt8Array, UInt32Array, UInt64Array,
26};
27use datatypes::arrow::error::ArrowError;
28use datatypes::arrow_array::StringArray;
29use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
30use serde::{Deserialize, Serialize};
31
32use crate::storage::{RegionGroup, RegionId, RegionNumber, RegionSeq, ScanRequest, TableId};
33
34/// Puffin index type for bloom filter indexes.
35pub const PUFFIN_INDEX_TYPE_BLOOM_FILTER: &str = "bloom_filter";
36/// Puffin index type for bloom-backed full-text indexes.
37pub const PUFFIN_INDEX_TYPE_FULLTEXT_BLOOM: &str = "fulltext_bloom";
38/// Puffin index type for Tantivy-backed full-text indexes.
39pub const PUFFIN_INDEX_TYPE_FULLTEXT_TANTIVY: &str = "fulltext_tantivy";
40/// Puffin index type for inverted indexes.
41pub const PUFFIN_INDEX_TYPE_INVERTED: &str = "inverted";
42
43/// An entry describing a SST file known by the engine's manifest.
44#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
45pub struct ManifestSstEntry {
46    /// The table directory this file belongs to.
47    pub table_dir: String,
48    /// The region id of region that refers to the file.
49    pub region_id: RegionId,
50    /// The table id this file belongs to.
51    pub table_id: TableId,
52    /// The region number this file belongs to.
53    pub region_number: RegionNumber,
54    /// The region group this file belongs to.
55    pub region_group: RegionGroup,
56    /// The region sequence this file belongs to.
57    pub region_sequence: RegionSeq,
58    /// Engine-specific file identifier (string form).
59    pub file_id: String,
60    /// Index version, increment when the index file is rebuilt.
61    pub index_version: u64,
62    /// SST level.
63    pub level: u8,
64    /// Full path of the SST file in object store.
65    pub file_path: String,
66    /// File size in bytes.
67    pub file_size: u64,
68    /// Maximum uncompressed row group size in bytes. Zero means unknown.
69    pub max_row_group_uncompressed_size: u64,
70    /// Full path of the index file in object store.
71    pub index_file_path: Option<String>,
72    /// File size of the index file in object store.
73    pub index_file_size: Option<u64>,
74    /// Number of rows in the SST.
75    pub num_rows: u64,
76    /// Number of row groups in the SST.
77    pub num_row_groups: u64,
78    /// Number of series in the SST.
79    pub num_series: Option<u64>,
80    /// Min timestamp.
81    pub min_ts: Timestamp,
82    /// Max timestamp.
83    pub max_ts: Timestamp,
84    /// The sequence number associated with this file.
85    pub sequence: Option<u64>,
86    /// Human-readable partition expression associated with this file.
87    pub partition_expr: Option<String>,
88    /// The region id of region that creates the file.
89    pub origin_region_id: RegionId,
90    /// The node id fetched from the manifest.
91    pub node_id: Option<u64>,
92    /// Whether this file is visible in current version.
93    pub visible: bool,
94    /// Minimum encoded primary key in the SST.
95    pub primary_key_min: Option<Bytes>,
96    /// Maximum encoded primary key in the SST.
97    pub primary_key_max: Option<Bytes>,
98}
99
100impl ManifestSstEntry {
101    /// Returns the schema of the manifest sst entry.
102    pub fn schema() -> SchemaRef {
103        use datatypes::prelude::ConcreteDataType as Ty;
104        Arc::new(Schema::new(vec![
105            ColumnSchema::new("table_dir", Ty::string_datatype(), false),
106            ColumnSchema::new("region_id", Ty::uint64_datatype(), false),
107            ColumnSchema::new("table_id", Ty::uint32_datatype(), false),
108            ColumnSchema::new("region_number", Ty::uint32_datatype(), false),
109            ColumnSchema::new("region_group", Ty::uint8_datatype(), false),
110            ColumnSchema::new("region_sequence", Ty::uint32_datatype(), false),
111            ColumnSchema::new("file_id", Ty::string_datatype(), false),
112            ColumnSchema::new("index_version", Ty::uint64_datatype(), false),
113            ColumnSchema::new("level", Ty::uint8_datatype(), false),
114            ColumnSchema::new("file_path", Ty::string_datatype(), false),
115            ColumnSchema::new("file_size", Ty::uint64_datatype(), false),
116            ColumnSchema::new("index_file_path", Ty::string_datatype(), true),
117            ColumnSchema::new("index_file_size", Ty::uint64_datatype(), true),
118            ColumnSchema::new("num_rows", Ty::uint64_datatype(), false),
119            ColumnSchema::new("num_row_groups", Ty::uint64_datatype(), false),
120            ColumnSchema::new("num_series", Ty::uint64_datatype(), true),
121            ColumnSchema::new("min_ts", Ty::timestamp_nanosecond_datatype(), true),
122            ColumnSchema::new("max_ts", Ty::timestamp_nanosecond_datatype(), true),
123            ColumnSchema::new("sequence", Ty::uint64_datatype(), true),
124            ColumnSchema::new("origin_region_id", Ty::uint64_datatype(), false),
125            ColumnSchema::new("node_id", Ty::uint64_datatype(), true),
126            ColumnSchema::new("visible", Ty::boolean_datatype(), false),
127            ColumnSchema::new("primary_key_min", Ty::binary_datatype(), true),
128            ColumnSchema::new("primary_key_max", Ty::binary_datatype(), true),
129            ColumnSchema::new(
130                "max_row_group_uncompressed_size",
131                Ty::uint64_datatype(),
132                false,
133            ),
134            ColumnSchema::new("partition_expr", Ty::string_datatype(), true),
135        ]))
136    }
137
138    /// Converts a list of manifest sst entries to a record batch.
139    pub fn to_record_batch(entries: &[Self]) -> std::result::Result<DfRecordBatch, ArrowError> {
140        let schema = Self::schema();
141        let table_dirs = entries.iter().map(|e| e.table_dir.as_str());
142        let region_ids = entries.iter().map(|e| e.region_id.as_u64());
143        let table_ids = entries.iter().map(|e| e.table_id);
144        let region_numbers = entries.iter().map(|e| e.region_number);
145        let region_groups = entries.iter().map(|e| e.region_group);
146        let region_sequences = entries.iter().map(|e| e.region_sequence);
147        let file_ids = entries.iter().map(|e| e.file_id.as_str());
148        let index_versions = entries.iter().map(|e| e.index_version);
149        let levels = entries.iter().map(|e| e.level);
150        let file_paths = entries.iter().map(|e| e.file_path.as_str());
151        let file_sizes = entries.iter().map(|e| e.file_size);
152        let max_row_group_uncompressed_sizes =
153            entries.iter().map(|e| e.max_row_group_uncompressed_size);
154        let index_file_paths = entries.iter().map(|e| e.index_file_path.as_ref());
155        let index_file_sizes = entries.iter().map(|e| e.index_file_size);
156        let num_rows = entries.iter().map(|e| e.num_rows);
157        let num_row_groups = entries.iter().map(|e| e.num_row_groups);
158        let num_series = entries.iter().map(|e| e.num_series);
159        let min_ts = entries.iter().map(|e| {
160            e.min_ts
161                .convert_to(TimeUnit::Nanosecond)
162                .map(|ts| ts.value())
163        });
164        let max_ts = entries.iter().map(|e| {
165            e.max_ts
166                .convert_to(TimeUnit::Nanosecond)
167                .map(|ts| ts.value())
168        });
169        let sequences = entries.iter().map(|e| e.sequence);
170        let partition_exprs = entries.iter().map(|e| e.partition_expr.as_ref());
171        let origin_region_ids = entries.iter().map(|e| e.origin_region_id.as_u64());
172        let node_ids = entries.iter().map(|e| e.node_id);
173        let visible_flags = entries.iter().map(|e| Some(e.visible));
174        let primary_key_min = entries.iter().map(|e| e.primary_key_min.as_deref());
175        let primary_key_max = entries.iter().map(|e| e.primary_key_max.as_deref());
176
177        let columns: Vec<ArrayRef> = vec![
178            Arc::new(StringArray::from_iter_values(table_dirs)),
179            Arc::new(UInt64Array::from_iter_values(region_ids)),
180            Arc::new(UInt32Array::from_iter_values(table_ids)),
181            Arc::new(UInt32Array::from_iter_values(region_numbers)),
182            Arc::new(UInt8Array::from_iter_values(region_groups)),
183            Arc::new(UInt32Array::from_iter_values(region_sequences)),
184            Arc::new(StringArray::from_iter_values(file_ids)),
185            Arc::new(UInt64Array::from_iter(index_versions)),
186            Arc::new(UInt8Array::from_iter_values(levels)),
187            Arc::new(StringArray::from_iter_values(file_paths)),
188            Arc::new(UInt64Array::from_iter_values(file_sizes)),
189            Arc::new(StringArray::from_iter(index_file_paths)),
190            Arc::new(UInt64Array::from_iter(index_file_sizes)),
191            Arc::new(UInt64Array::from_iter_values(num_rows)),
192            Arc::new(UInt64Array::from_iter_values(num_row_groups)),
193            Arc::new(UInt64Array::from_iter(num_series)),
194            Arc::new(TimestampNanosecondArray::from_iter(min_ts)),
195            Arc::new(TimestampNanosecondArray::from_iter(max_ts)),
196            Arc::new(UInt64Array::from_iter(sequences)),
197            Arc::new(UInt64Array::from_iter_values(origin_region_ids)),
198            Arc::new(UInt64Array::from_iter(node_ids)),
199            Arc::new(BooleanArray::from_iter(visible_flags)),
200            Arc::new(BinaryArray::from_iter(primary_key_min)),
201            Arc::new(BinaryArray::from_iter(primary_key_max)),
202            Arc::new(UInt64Array::from_iter_values(
203                max_row_group_uncompressed_sizes,
204            )),
205            Arc::new(StringArray::from_iter(partition_exprs)),
206        ];
207
208        DfRecordBatch::try_new(schema.arrow_schema().clone(), columns)
209    }
210
211    /// Reserved internal inspect table name.
212    ///
213    /// This table name is used only for building logical plans on the
214    /// frontend -> datanode path. It is not user-visible and cannot be
215    /// referenced by user queries.
216    pub fn reserved_table_name_for_inspection() -> &'static str {
217        "__inspect/__mito/__sst_manifest"
218    }
219
220    /// Builds a logical plan for scanning the manifest sst entries.
221    pub fn build_plan(scan_request: ScanRequest) -> Result<LogicalPlan, DataFusionError> {
222        build_plan_helper(
223            scan_request,
224            Self::reserved_table_name_for_inspection(),
225            Self::schema(),
226        )
227    }
228}
229
230/// An entry describing a SST file listed from storage layer directly.
231#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
232pub struct StorageSstEntry {
233    /// Full path of the SST file in object store.
234    pub file_path: String,
235    /// File size in bytes.
236    pub file_size: Option<u64>,
237    /// Last modified time in milliseconds since epoch, if available from storage.
238    pub last_modified_ms: Option<Timestamp>,
239    /// The node id fetched from the manifest.
240    pub node_id: Option<u64>,
241}
242
243impl StorageSstEntry {
244    /// Returns the schema of the storage sst entry.
245    pub fn schema() -> SchemaRef {
246        use datatypes::prelude::ConcreteDataType as Ty;
247        Arc::new(Schema::new(vec![
248            ColumnSchema::new("file_path", Ty::string_datatype(), false),
249            ColumnSchema::new("file_size", Ty::uint64_datatype(), true),
250            ColumnSchema::new(
251                "last_modified_ms",
252                Ty::timestamp_millisecond_datatype(),
253                true,
254            ),
255            ColumnSchema::new("node_id", Ty::uint64_datatype(), true),
256        ]))
257    }
258
259    /// Converts a list of storage sst entries to a record batch.
260    pub fn to_record_batch(entries: &[Self]) -> std::result::Result<DfRecordBatch, ArrowError> {
261        let schema = Self::schema();
262        let file_paths = entries.iter().map(|e| e.file_path.as_str());
263        let file_sizes = entries.iter().map(|e| e.file_size);
264        let last_modified_ms = entries.iter().map(|e| {
265            e.last_modified_ms
266                .and_then(|ts| ts.convert_to(TimeUnit::Millisecond).map(|ts| ts.value()))
267        });
268        let node_ids = entries.iter().map(|e| e.node_id);
269
270        let columns: Vec<ArrayRef> = vec![
271            Arc::new(StringArray::from_iter_values(file_paths)),
272            Arc::new(UInt64Array::from_iter(file_sizes)),
273            Arc::new(TimestampMillisecondArray::from_iter(last_modified_ms)),
274            Arc::new(UInt64Array::from_iter(node_ids)),
275        ];
276
277        DfRecordBatch::try_new(schema.arrow_schema().clone(), columns)
278    }
279
280    /// Reserved internal inspect table name.
281    ///
282    /// This table name is used only for building logical plans on the
283    /// frontend -> datanode path. It is not user-visible and cannot be
284    /// referenced by user queries.
285    pub fn reserved_table_name_for_inspection() -> &'static str {
286        "__inspect/__mito/__sst_storage"
287    }
288
289    /// Builds a logical plan for scanning the storage sst entries.
290    pub fn build_plan(scan_request: ScanRequest) -> Result<LogicalPlan, DataFusionError> {
291        build_plan_helper(
292            scan_request,
293            Self::reserved_table_name_for_inspection(),
294            Self::schema(),
295        )
296    }
297}
298
299/// An entry describing puffin index metadata for inspection.
300#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
301pub struct PuffinIndexMetaEntry {
302    /// The table directory this index belongs to.
303    pub table_dir: String,
304    /// The full path of the index file in object store.
305    pub index_file_path: String,
306    /// The region id referencing the index file.
307    pub region_id: RegionId,
308    /// The table id referencing the index file.
309    pub table_id: TableId,
310    /// The region number referencing the index file.
311    pub region_number: RegionNumber,
312    /// The region group referencing the index file.
313    pub region_group: RegionGroup,
314    /// The region sequence referencing the index file.
315    pub region_sequence: RegionSeq,
316    /// Engine-specific file identifier (string form).
317    pub file_id: String,
318    /// Size of the index file in object store (if available).
319    pub index_file_size: Option<u64>,
320    /// Logical index type (`bloom_filter`, `fulltext_bloom`, `fulltext_tantivy`, `inverted`).
321    pub index_type: String,
322    /// Target type (`column`, ...).
323    pub target_type: String,
324    /// Encoded target key string.
325    pub target_key: String,
326    /// Structured JSON describing the target.
327    pub target_json: String,
328    /// Size of the blob storing this target.
329    pub blob_size: u64,
330    /// Structured JSON describing index-specific metadata (if available).
331    pub meta_json: Option<String>,
332    /// Node id associated with the index file (if known).
333    pub node_id: Option<u64>,
334}
335
336impl PuffinIndexMetaEntry {
337    /// Returns the schema describing puffin index metadata entries.
338    pub fn schema() -> SchemaRef {
339        use datatypes::prelude::ConcreteDataType as Ty;
340        Arc::new(Schema::new(vec![
341            ColumnSchema::new("table_dir", Ty::string_datatype(), false),
342            ColumnSchema::new("index_file_path", Ty::string_datatype(), false),
343            ColumnSchema::new("region_id", Ty::uint64_datatype(), false),
344            ColumnSchema::new("table_id", Ty::uint32_datatype(), false),
345            ColumnSchema::new("region_number", Ty::uint32_datatype(), false),
346            ColumnSchema::new("region_group", Ty::uint8_datatype(), false),
347            ColumnSchema::new("region_sequence", Ty::uint32_datatype(), false),
348            ColumnSchema::new("file_id", Ty::string_datatype(), false),
349            ColumnSchema::new("index_file_size", Ty::uint64_datatype(), true),
350            ColumnSchema::new("index_type", Ty::string_datatype(), false),
351            ColumnSchema::new("target_type", Ty::string_datatype(), false),
352            ColumnSchema::new("target_key", Ty::string_datatype(), false),
353            ColumnSchema::new("target_json", Ty::string_datatype(), false),
354            ColumnSchema::new("blob_size", Ty::uint64_datatype(), false),
355            ColumnSchema::new("meta_json", Ty::string_datatype(), true),
356            ColumnSchema::new("node_id", Ty::uint64_datatype(), true),
357        ]))
358    }
359
360    /// Converts a list of puffin index metadata entries to a record batch.
361    pub fn to_record_batch(entries: &[Self]) -> std::result::Result<DfRecordBatch, ArrowError> {
362        let schema = Self::schema();
363        let table_dirs = entries.iter().map(|e| e.table_dir.as_str());
364        let index_file_paths = entries.iter().map(|e| e.index_file_path.as_str());
365        let region_ids = entries.iter().map(|e| e.region_id.as_u64());
366        let table_ids = entries.iter().map(|e| e.table_id);
367        let region_numbers = entries.iter().map(|e| e.region_number);
368        let region_groups = entries.iter().map(|e| e.region_group);
369        let region_sequences = entries.iter().map(|e| e.region_sequence);
370        let file_ids = entries.iter().map(|e| e.file_id.as_str());
371        let index_file_sizes = entries.iter().map(|e| e.index_file_size);
372        let index_types = entries.iter().map(|e| e.index_type.as_str());
373        let target_types = entries.iter().map(|e| e.target_type.as_str());
374        let target_keys = entries.iter().map(|e| e.target_key.as_str());
375        let target_jsons = entries.iter().map(|e| e.target_json.as_str());
376        let blob_sizes = entries.iter().map(|e| e.blob_size);
377        let meta_jsons = entries.iter().map(|e| e.meta_json.as_deref());
378        let node_ids = entries.iter().map(|e| e.node_id);
379
380        let columns: Vec<ArrayRef> = vec![
381            Arc::new(StringArray::from_iter_values(table_dirs)),
382            Arc::new(StringArray::from_iter_values(index_file_paths)),
383            Arc::new(UInt64Array::from_iter_values(region_ids)),
384            Arc::new(UInt32Array::from_iter_values(table_ids)),
385            Arc::new(UInt32Array::from_iter_values(region_numbers)),
386            Arc::new(UInt8Array::from_iter_values(region_groups)),
387            Arc::new(UInt32Array::from_iter_values(region_sequences)),
388            Arc::new(StringArray::from_iter_values(file_ids)),
389            Arc::new(UInt64Array::from_iter(index_file_sizes)),
390            Arc::new(StringArray::from_iter_values(index_types)),
391            Arc::new(StringArray::from_iter_values(target_types)),
392            Arc::new(StringArray::from_iter_values(target_keys)),
393            Arc::new(StringArray::from_iter_values(target_jsons)),
394            Arc::new(UInt64Array::from_iter_values(blob_sizes)),
395            Arc::new(StringArray::from_iter(meta_jsons)),
396            Arc::new(UInt64Array::from_iter(node_ids)),
397        ];
398
399        DfRecordBatch::try_new(schema.arrow_schema().clone(), columns)
400    }
401
402    /// Reserved internal inspect table name for puffin index metadata.
403    pub fn reserved_table_name_for_inspection() -> &'static str {
404        "__inspect/__mito/__puffin_index_meta"
405    }
406
407    /// Builds a logical plan for scanning puffin index metadata entries.
408    pub fn build_plan(scan_request: ScanRequest) -> Result<LogicalPlan, DataFusionError> {
409        build_plan_helper(
410            scan_request,
411            Self::reserved_table_name_for_inspection(),
412            Self::schema(),
413        )
414    }
415}
416
417fn build_plan_helper(
418    scan_request: ScanRequest,
419    table_name: &str,
420    schema: SchemaRef,
421) -> Result<LogicalPlan, DataFusionError> {
422    let table_source = LogicalTableSource::new(schema.arrow_schema().clone());
423
424    let projection = scan_request.projection;
425    let mut builder = LogicalPlanBuilder::scan(table_name, Arc::new(table_source), projection)?;
426
427    for filter in scan_request.filters {
428        builder = builder.filter(filter)?;
429    }
430
431    if let Some(limit) = scan_request.limit {
432        builder = builder.limit(0, Some(limit))?;
433    }
434
435    builder.build()
436}
437
438#[cfg(test)]
439mod tests {
440    use datafusion_common::TableReference;
441    use datafusion_expr::{LogicalPlan, Operator, binary_expr, col, lit};
442    use datatypes::arrow::array::{
443        Array, BinaryArray, TimestampMillisecondArray, TimestampNanosecondArray, UInt8Array,
444        UInt32Array, UInt64Array,
445    };
446    use datatypes::arrow_array::StringArray;
447
448    use super::*;
449
450    #[test]
451    fn test_sst_entry_manifest_to_record_batch() {
452        // Prepare entries
453        let table_id1: TableId = 1;
454        let region_group1: RegionGroup = 2;
455        let region_seq1: RegionSeq = 3;
456        let region_number1: RegionNumber = ((region_group1 as u32) << 24) | region_seq1;
457        let region_id1 = RegionId::with_group_and_seq(table_id1, region_group1, region_seq1);
458
459        let table_id2: TableId = 5;
460        let region_group2: RegionGroup = 1;
461        let region_seq2: RegionSeq = 42;
462        let region_number2: RegionNumber = ((region_group2 as u32) << 24) | region_seq2;
463        let region_id2 = RegionId::with_group_and_seq(table_id2, region_group2, region_seq2);
464
465        let entries = vec![
466            ManifestSstEntry {
467                table_dir: "tdir1".to_string(),
468                region_id: region_id1,
469                table_id: table_id1,
470                region_number: region_number1,
471                region_group: region_group1,
472                region_sequence: region_seq1,
473                file_id: "f1".to_string(),
474                index_version: 0,
475                level: 1,
476                file_path: "/p1".to_string(),
477                file_size: 100,
478                max_row_group_uncompressed_size: 80,
479                index_file_path: None,
480                index_file_size: None,
481                num_rows: 10,
482                num_row_groups: 2,
483                num_series: Some(5),
484                min_ts: Timestamp::new_millisecond(1000), // 1s -> 1_000_000_000ns
485                max_ts: Timestamp::new_second(2),         // 2s -> 2_000_000_000ns
486                sequence: None,
487                partition_expr: Some("a < 10".to_string()),
488                origin_region_id: region_id1,
489                node_id: Some(1),
490                visible: false,
491                primary_key_min: Some(Bytes::from_static(b"aaa")),
492                primary_key_max: Some(Bytes::from_static(b"zzz")),
493            },
494            ManifestSstEntry {
495                table_dir: "tdir2".to_string(),
496                region_id: region_id2,
497                table_id: table_id2,
498                region_number: region_number2,
499                region_group: region_group2,
500                region_sequence: region_seq2,
501                file_id: "f2".to_string(),
502                index_version: 1,
503                level: 3,
504                file_path: "/p2".to_string(),
505                file_size: 200,
506                max_row_group_uncompressed_size: 160,
507                index_file_path: Some("idx".to_string()),
508                index_file_size: Some(11),
509                num_rows: 20,
510                num_row_groups: 4,
511                num_series: None,
512                min_ts: Timestamp::new_nanosecond(5),     // 5ns
513                max_ts: Timestamp::new_microsecond(2000), // 2ms -> 2_000_000ns
514                sequence: Some(9),
515                partition_expr: None,
516                origin_region_id: region_id2,
517                node_id: None,
518                visible: true,
519                primary_key_min: None,
520                primary_key_max: None,
521            },
522        ];
523
524        let schema = ManifestSstEntry::schema();
525        let batch = ManifestSstEntry::to_record_batch(&entries).unwrap();
526
527        // Schema checks
528        assert_eq!(schema.arrow_schema().fields().len(), batch.num_columns());
529        assert_eq!(2, batch.num_rows());
530        let expected_columns = [
531            "table_dir",
532            "region_id",
533            "table_id",
534            "region_number",
535            "region_group",
536            "region_sequence",
537            "file_id",
538            "index_version",
539            "level",
540            "file_path",
541            "file_size",
542            "index_file_path",
543            "index_file_size",
544            "num_rows",
545            "num_row_groups",
546            "num_series",
547            "min_ts",
548            "max_ts",
549            "sequence",
550            "origin_region_id",
551            "node_id",
552            "visible",
553            "primary_key_min",
554            "primary_key_max",
555            "max_row_group_uncompressed_size",
556            "partition_expr",
557        ];
558        assert_eq!(
559            expected_columns,
560            schema
561                .arrow_schema()
562                .fields()
563                .iter()
564                .map(|field| field.name().as_str())
565                .collect::<Vec<_>>()
566                .as_slice()
567        );
568        for (i, f) in schema.arrow_schema().fields().iter().enumerate() {
569            assert_eq!(f.name(), batch.schema().field(i).name());
570            assert_eq!(f.is_nullable(), batch.schema().field(i).is_nullable());
571            assert_eq!(f.data_type(), batch.schema().field(i).data_type());
572        }
573
574        // Column asserts
575        let table_dirs = batch
576            .column(0)
577            .as_any()
578            .downcast_ref::<StringArray>()
579            .unwrap();
580        assert_eq!("tdir1", table_dirs.value(0));
581        assert_eq!("tdir2", table_dirs.value(1));
582
583        let region_ids = batch
584            .column(1)
585            .as_any()
586            .downcast_ref::<UInt64Array>()
587            .unwrap();
588        assert_eq!(region_id1.as_u64(), region_ids.value(0));
589        assert_eq!(region_id2.as_u64(), region_ids.value(1));
590
591        let table_ids = batch
592            .column(2)
593            .as_any()
594            .downcast_ref::<UInt32Array>()
595            .unwrap();
596        assert_eq!(table_id1, table_ids.value(0));
597        assert_eq!(table_id2, table_ids.value(1));
598
599        let region_numbers = batch
600            .column(3)
601            .as_any()
602            .downcast_ref::<UInt32Array>()
603            .unwrap();
604        assert_eq!(region_number1, region_numbers.value(0));
605        assert_eq!(region_number2, region_numbers.value(1));
606
607        let region_groups = batch
608            .column(4)
609            .as_any()
610            .downcast_ref::<UInt8Array>()
611            .unwrap();
612        assert_eq!(region_group1, region_groups.value(0));
613        assert_eq!(region_group2, region_groups.value(1));
614
615        let region_sequences = batch
616            .column(5)
617            .as_any()
618            .downcast_ref::<UInt32Array>()
619            .unwrap();
620        assert_eq!(region_seq1, region_sequences.value(0));
621        assert_eq!(region_seq2, region_sequences.value(1));
622
623        let file_ids = batch
624            .column(6)
625            .as_any()
626            .downcast_ref::<StringArray>()
627            .unwrap();
628        assert_eq!("f1", file_ids.value(0));
629        assert_eq!("f2", file_ids.value(1));
630
631        let index_versions = batch
632            .column(7)
633            .as_any()
634            .downcast_ref::<UInt64Array>()
635            .unwrap();
636        assert_eq!(0, index_versions.value(0));
637        assert_eq!(1, index_versions.value(1));
638
639        let levels = batch
640            .column(8)
641            .as_any()
642            .downcast_ref::<UInt8Array>()
643            .unwrap();
644        assert_eq!(1, levels.value(0));
645        assert_eq!(3, levels.value(1));
646
647        let file_paths = batch
648            .column(9)
649            .as_any()
650            .downcast_ref::<StringArray>()
651            .unwrap();
652        assert_eq!("/p1", file_paths.value(0));
653        assert_eq!("/p2", file_paths.value(1));
654
655        let file_sizes = batch
656            .column(10)
657            .as_any()
658            .downcast_ref::<UInt64Array>()
659            .unwrap();
660        assert_eq!(100, file_sizes.value(0));
661        assert_eq!(200, file_sizes.value(1));
662
663        let index_file_paths = batch
664            .column(11)
665            .as_any()
666            .downcast_ref::<StringArray>()
667            .unwrap();
668        assert!(index_file_paths.is_null(0));
669        assert_eq!("idx", index_file_paths.value(1));
670
671        let index_file_sizes = batch
672            .column(12)
673            .as_any()
674            .downcast_ref::<UInt64Array>()
675            .unwrap();
676        assert!(index_file_sizes.is_null(0));
677        assert_eq!(11, index_file_sizes.value(1));
678
679        let num_rows = batch
680            .column(13)
681            .as_any()
682            .downcast_ref::<UInt64Array>()
683            .unwrap();
684        assert_eq!(10, num_rows.value(0));
685        assert_eq!(20, num_rows.value(1));
686
687        let num_row_groups = batch
688            .column(14)
689            .as_any()
690            .downcast_ref::<UInt64Array>()
691            .unwrap();
692        assert_eq!(2, num_row_groups.value(0));
693        assert_eq!(4, num_row_groups.value(1));
694
695        let num_series = batch
696            .column(15)
697            .as_any()
698            .downcast_ref::<UInt64Array>()
699            .unwrap();
700        assert_eq!(5, num_series.value(0));
701        assert!(num_series.is_null(1));
702
703        let min_ts = batch
704            .column(16)
705            .as_any()
706            .downcast_ref::<TimestampNanosecondArray>()
707            .unwrap();
708        assert_eq!(1_000_000_000, min_ts.value(0));
709        assert_eq!(5, min_ts.value(1));
710
711        let max_ts = batch
712            .column(17)
713            .as_any()
714            .downcast_ref::<TimestampNanosecondArray>()
715            .unwrap();
716        assert_eq!(2_000_000_000, max_ts.value(0));
717        assert_eq!(2_000_000, max_ts.value(1));
718
719        let sequences = batch
720            .column(18)
721            .as_any()
722            .downcast_ref::<UInt64Array>()
723            .unwrap();
724        assert!(sequences.is_null(0));
725        assert_eq!(9, sequences.value(1));
726
727        let origin_region_ids = batch
728            .column(19)
729            .as_any()
730            .downcast_ref::<UInt64Array>()
731            .unwrap();
732        assert_eq!(region_id1.as_u64(), origin_region_ids.value(0));
733        assert_eq!(region_id2.as_u64(), origin_region_ids.value(1));
734
735        let node_ids = batch
736            .column(20)
737            .as_any()
738            .downcast_ref::<UInt64Array>()
739            .unwrap();
740        assert_eq!(1, node_ids.value(0));
741        assert!(node_ids.is_null(1));
742
743        let visible = batch
744            .column(21)
745            .as_any()
746            .downcast_ref::<BooleanArray>()
747            .unwrap();
748        assert!(!visible.value(0));
749        assert!(visible.value(1));
750
751        let primary_key_min = batch
752            .column(22)
753            .as_any()
754            .downcast_ref::<BinaryArray>()
755            .unwrap();
756        assert_eq!(b"aaa", primary_key_min.value(0));
757        assert!(primary_key_min.is_null(1));
758
759        let primary_key_max = batch
760            .column(23)
761            .as_any()
762            .downcast_ref::<BinaryArray>()
763            .unwrap();
764        assert_eq!(b"zzz", primary_key_max.value(0));
765        assert!(primary_key_max.is_null(1));
766
767        let max_row_group_uncompressed_sizes = batch
768            .column(24)
769            .as_any()
770            .downcast_ref::<UInt64Array>()
771            .unwrap();
772        assert_eq!(80, max_row_group_uncompressed_sizes.value(0));
773        assert_eq!(160, max_row_group_uncompressed_sizes.value(1));
774
775        let partition_exprs = batch
776            .column(25)
777            .as_any()
778            .downcast_ref::<StringArray>()
779            .unwrap();
780        assert_eq!("a < 10", partition_exprs.value(0));
781        assert!(partition_exprs.is_null(1));
782    }
783
784    #[test]
785    fn test_sst_entry_storage_to_record_batch() {
786        let entries = vec![
787            StorageSstEntry {
788                file_path: "/s1".to_string(),
789                file_size: None,
790                last_modified_ms: None,
791                node_id: Some(1),
792            },
793            StorageSstEntry {
794                file_path: "/s2".to_string(),
795                file_size: Some(123),
796                last_modified_ms: Some(Timestamp::new_millisecond(456)),
797                node_id: None,
798            },
799        ];
800
801        let schema = StorageSstEntry::schema();
802        let batch = StorageSstEntry::to_record_batch(&entries).unwrap();
803
804        assert_eq!(schema.arrow_schema().fields().len(), batch.num_columns());
805        assert_eq!(2, batch.num_rows());
806
807        let file_paths = batch
808            .column(0)
809            .as_any()
810            .downcast_ref::<StringArray>()
811            .unwrap();
812        assert_eq!("/s1", file_paths.value(0));
813        assert_eq!("/s2", file_paths.value(1));
814
815        let file_sizes = batch
816            .column(1)
817            .as_any()
818            .downcast_ref::<UInt64Array>()
819            .unwrap();
820        assert!(file_sizes.is_null(0));
821        assert_eq!(123, file_sizes.value(1));
822
823        let last_modified = batch
824            .column(2)
825            .as_any()
826            .downcast_ref::<TimestampMillisecondArray>()
827            .unwrap();
828        assert!(last_modified.is_null(0));
829        assert_eq!(456, last_modified.value(1));
830
831        let node_ids = batch
832            .column(3)
833            .as_any()
834            .downcast_ref::<UInt64Array>()
835            .unwrap();
836        assert_eq!(1, node_ids.value(0));
837        assert!(node_ids.is_null(1));
838    }
839
840    #[test]
841    fn test_puffin_index_meta_to_record_batch() {
842        let entries = vec![
843            PuffinIndexMetaEntry {
844                table_dir: "table1".to_string(),
845                index_file_path: "index1".to_string(),
846                region_id: RegionId::with_group_and_seq(10, 0, 20),
847                table_id: 10,
848                region_number: 20,
849                region_group: 0,
850                region_sequence: 20,
851                file_id: "file1".to_string(),
852                index_file_size: Some(1024),
853                index_type: "bloom_filter".to_string(),
854                target_type: "column".to_string(),
855                target_key: "1".to_string(),
856                target_json: "{\"column\":1}".to_string(),
857                blob_size: 256,
858                meta_json: Some("{\"bloom\":{}}".to_string()),
859                node_id: Some(42),
860            },
861            PuffinIndexMetaEntry {
862                table_dir: "table2".to_string(),
863                index_file_path: "index2".to_string(),
864                region_id: RegionId::with_group_and_seq(11, 0, 21),
865                table_id: 11,
866                region_number: 21,
867                region_group: 0,
868                region_sequence: 21,
869                file_id: "file2".to_string(),
870                index_file_size: None,
871                index_type: "inverted".to_string(),
872                target_type: "unknown".to_string(),
873                target_key: "legacy".to_string(),
874                target_json: "{}".to_string(),
875                blob_size: 0,
876                meta_json: None,
877                node_id: None,
878            },
879        ];
880
881        let schema = PuffinIndexMetaEntry::schema();
882        let batch = PuffinIndexMetaEntry::to_record_batch(&entries).unwrap();
883
884        assert_eq!(schema.arrow_schema().fields().len(), batch.num_columns());
885        assert_eq!(2, batch.num_rows());
886
887        let table_dirs = batch
888            .column(0)
889            .as_any()
890            .downcast_ref::<StringArray>()
891            .unwrap();
892        assert_eq!("table1", table_dirs.value(0));
893        assert_eq!("table2", table_dirs.value(1));
894
895        let index_file_paths = batch
896            .column(1)
897            .as_any()
898            .downcast_ref::<StringArray>()
899            .unwrap();
900        assert_eq!("index1", index_file_paths.value(0));
901        assert_eq!("index2", index_file_paths.value(1));
902
903        let region_ids = batch
904            .column(2)
905            .as_any()
906            .downcast_ref::<UInt64Array>()
907            .unwrap();
908        assert_eq!(
909            RegionId::with_group_and_seq(10, 0, 20).as_u64(),
910            region_ids.value(0)
911        );
912        assert_eq!(
913            RegionId::with_group_and_seq(11, 0, 21).as_u64(),
914            region_ids.value(1)
915        );
916
917        let table_ids = batch
918            .column(3)
919            .as_any()
920            .downcast_ref::<UInt32Array>()
921            .unwrap();
922        assert_eq!(10, table_ids.value(0));
923        assert_eq!(11, table_ids.value(1));
924
925        let region_numbers = batch
926            .column(4)
927            .as_any()
928            .downcast_ref::<UInt32Array>()
929            .unwrap();
930        assert_eq!(20, region_numbers.value(0));
931        assert_eq!(21, region_numbers.value(1));
932
933        let region_groups = batch
934            .column(5)
935            .as_any()
936            .downcast_ref::<UInt8Array>()
937            .unwrap();
938        assert_eq!(0, region_groups.value(0));
939        assert_eq!(0, region_groups.value(1));
940
941        let region_sequences = batch
942            .column(6)
943            .as_any()
944            .downcast_ref::<UInt32Array>()
945            .unwrap();
946        assert_eq!(20, region_sequences.value(0));
947        assert_eq!(21, region_sequences.value(1));
948
949        let file_ids = batch
950            .column(7)
951            .as_any()
952            .downcast_ref::<StringArray>()
953            .unwrap();
954        assert_eq!("file1", file_ids.value(0));
955        assert_eq!("file2", file_ids.value(1));
956
957        let index_file_sizes = batch
958            .column(8)
959            .as_any()
960            .downcast_ref::<UInt64Array>()
961            .unwrap();
962        assert_eq!(1024, index_file_sizes.value(0));
963        assert!(index_file_sizes.is_null(1));
964
965        let index_types = batch
966            .column(9)
967            .as_any()
968            .downcast_ref::<StringArray>()
969            .unwrap();
970        assert_eq!("bloom_filter", index_types.value(0));
971        assert_eq!("inverted", index_types.value(1));
972
973        let target_types = batch
974            .column(10)
975            .as_any()
976            .downcast_ref::<StringArray>()
977            .unwrap();
978        assert_eq!("column", target_types.value(0));
979        assert_eq!("unknown", target_types.value(1));
980
981        let target_keys = batch
982            .column(11)
983            .as_any()
984            .downcast_ref::<StringArray>()
985            .unwrap();
986        assert_eq!("1", target_keys.value(0));
987        assert_eq!("legacy", target_keys.value(1));
988
989        let target_json = batch
990            .column(12)
991            .as_any()
992            .downcast_ref::<StringArray>()
993            .unwrap();
994        assert_eq!("{\"column\":1}", target_json.value(0));
995        assert_eq!("{}", target_json.value(1));
996
997        let blob_sizes = batch
998            .column(13)
999            .as_any()
1000            .downcast_ref::<UInt64Array>()
1001            .unwrap();
1002        assert_eq!(256, blob_sizes.value(0));
1003        assert_eq!(0, blob_sizes.value(1));
1004
1005        let meta_jsons = batch
1006            .column(14)
1007            .as_any()
1008            .downcast_ref::<StringArray>()
1009            .unwrap();
1010        assert_eq!("{\"bloom\":{}}", meta_jsons.value(0));
1011        assert!(meta_jsons.is_null(1));
1012
1013        let node_ids = batch
1014            .column(15)
1015            .as_any()
1016            .downcast_ref::<UInt64Array>()
1017            .unwrap();
1018        assert_eq!(42, node_ids.value(0));
1019        assert!(node_ids.is_null(1));
1020    }
1021
1022    #[test]
1023    fn test_manifest_build_plan() {
1024        // Note: filter must reference a column in the projected schema
1025        let projection = Some(vec![0, 1, 2]);
1026        let request = ScanRequest {
1027            projection,
1028            filters: vec![binary_expr(col("table_id"), Operator::Gt, lit(0))],
1029            limit: Some(5),
1030            ..Default::default()
1031        };
1032
1033        let plan = ManifestSstEntry::build_plan(request).unwrap();
1034
1035        // Expect plan to be Filter -> Limit -> TableScan or Filter+Limit wrapped.
1036        // We'll pattern match to reach TableScan and verify key fields.
1037        let (scan, has_filter, has_limit) = extract_scan(&plan);
1038
1039        assert!(has_filter);
1040        assert!(has_limit);
1041        assert_eq!(
1042            scan.table_name,
1043            TableReference::bare(ManifestSstEntry::reserved_table_name_for_inspection())
1044        );
1045        assert_eq!(scan.projection, Some(vec![0, 1, 2]));
1046
1047        // projected schema should match projection
1048        let fields = scan.projected_schema.fields();
1049        assert_eq!(fields.len(), 3);
1050        assert_eq!(fields[0].name(), "table_dir");
1051        assert_eq!(fields[1].name(), "region_id");
1052        assert_eq!(fields[2].name(), "table_id");
1053    }
1054
1055    #[test]
1056    fn test_storage_build_plan() {
1057        let projection = Some(vec![0, 2]);
1058        let request = ScanRequest {
1059            projection,
1060            filters: vec![binary_expr(col("file_path"), Operator::Eq, lit("/a"))],
1061            limit: Some(1),
1062            ..Default::default()
1063        };
1064
1065        let plan = StorageSstEntry::build_plan(request).unwrap();
1066        let (scan, has_filter, has_limit) = extract_scan(&plan);
1067        assert!(has_filter);
1068        assert!(has_limit);
1069        assert_eq!(
1070            scan.table_name,
1071            TableReference::bare(StorageSstEntry::reserved_table_name_for_inspection())
1072        );
1073        assert_eq!(scan.projection, Some(vec![0, 2]));
1074
1075        let fields = scan.projected_schema.fields();
1076        assert_eq!(fields.len(), 2);
1077        assert_eq!(fields[0].name(), "file_path");
1078        assert_eq!(fields[1].name(), "last_modified_ms");
1079    }
1080
1081    // Helper to reach TableScan and detect presence of Filter/Limit in plan
1082    fn extract_scan(plan: &LogicalPlan) -> (&datafusion_expr::logical_plan::TableScan, bool, bool) {
1083        use datafusion_expr::logical_plan::Limit;
1084
1085        match plan {
1086            LogicalPlan::Filter(f) => {
1087                let (scan, _, has_limit) = extract_scan(&f.input);
1088                (scan, true, has_limit)
1089            }
1090            LogicalPlan::Limit(Limit { input, .. }) => {
1091                let (scan, has_filter, _) = extract_scan(input);
1092                (scan, has_filter, true)
1093            }
1094            LogicalPlan::TableScan(scan) => (scan, false, false),
1095            other => panic!("unexpected plan: {other:?}"),
1096        }
1097    }
1098}