1use 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
34pub const PUFFIN_INDEX_TYPE_BLOOM_FILTER: &str = "bloom_filter";
36pub const PUFFIN_INDEX_TYPE_FULLTEXT_BLOOM: &str = "fulltext_bloom";
38pub const PUFFIN_INDEX_TYPE_FULLTEXT_TANTIVY: &str = "fulltext_tantivy";
40pub const PUFFIN_INDEX_TYPE_INVERTED: &str = "inverted";
42
43#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
45pub struct ManifestSstEntry {
46 pub table_dir: String,
48 pub region_id: RegionId,
50 pub table_id: TableId,
52 pub region_number: RegionNumber,
54 pub region_group: RegionGroup,
56 pub region_sequence: RegionSeq,
58 pub file_id: String,
60 pub index_version: u64,
62 pub level: u8,
64 pub file_path: String,
66 pub file_size: u64,
68 pub max_row_group_uncompressed_size: u64,
70 pub index_file_path: Option<String>,
72 pub index_file_size: Option<u64>,
74 pub num_rows: u64,
76 pub num_row_groups: u64,
78 pub num_series: Option<u64>,
80 pub min_ts: Timestamp,
82 pub max_ts: Timestamp,
84 pub sequence: Option<u64>,
86 pub partition_expr: Option<String>,
88 pub origin_region_id: RegionId,
90 pub node_id: Option<u64>,
92 pub visible: bool,
94 pub primary_key_min: Option<Bytes>,
96 pub primary_key_max: Option<Bytes>,
98}
99
100impl ManifestSstEntry {
101 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 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 pub fn reserved_table_name_for_inspection() -> &'static str {
217 "__inspect/__mito/__sst_manifest"
218 }
219
220 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#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
232pub struct StorageSstEntry {
233 pub file_path: String,
235 pub file_size: Option<u64>,
237 pub last_modified_ms: Option<Timestamp>,
239 pub node_id: Option<u64>,
241}
242
243impl StorageSstEntry {
244 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 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 pub fn reserved_table_name_for_inspection() -> &'static str {
286 "__inspect/__mito/__sst_storage"
287 }
288
289 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#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
301pub struct PuffinIndexMetaEntry {
302 pub table_dir: String,
304 pub index_file_path: String,
306 pub region_id: RegionId,
308 pub table_id: TableId,
310 pub region_number: RegionNumber,
312 pub region_group: RegionGroup,
314 pub region_sequence: RegionSeq,
316 pub file_id: String,
318 pub index_file_size: Option<u64>,
320 pub index_type: String,
322 pub target_type: String,
324 pub target_key: String,
326 pub target_json: String,
328 pub blob_size: u64,
330 pub meta_json: Option<String>,
332 pub node_id: Option<u64>,
334}
335
336impl PuffinIndexMetaEntry {
337 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 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 pub fn reserved_table_name_for_inspection() -> &'static str {
404 "__inspect/__mito/__puffin_index_meta"
405 }
406
407 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 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), max_ts: Timestamp::new_second(2), 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), max_ts: Timestamp::new_microsecond(2000), 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 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 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 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 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 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 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}