store_api/
inspect.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 common_recordbatch::DfRecordBatch;
18use datafusion_common::DataFusionError;
19use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, LogicalTableSource};
20use datatypes::arrow::array::ArrayRef;
21use datatypes::arrow::datatypes::SchemaRef as ArrowSchemaRef;
22use datatypes::arrow::error::ArrowError;
23use datatypes::arrow::record_batch::RecordBatchOptions;
24use datatypes::schema::SchemaRef;
25
26use crate::storage::ScanRequest;
27
28/// Builds projected inspection batches while sharing arrays for repeated columns.
29pub(crate) struct ProjectedBatchBuilder {
30    schema: ArrowSchemaRef,
31    duplicates: Vec<Option<usize>>,
32}
33
34impl ProjectedBatchBuilder {
35    /// Prepares duplicate-column mappings once per scan.
36    pub(crate) fn new(schema: ArrowSchemaRef) -> Self {
37        let duplicates = schema
38            .fields()
39            .iter()
40            .enumerate()
41            .map(|(index, field)| {
42                schema.fields()[..index]
43                    .iter()
44                    .position(|previous| previous == field)
45            })
46            .collect();
47        Self { schema, duplicates }
48    }
49
50    /// Builds each distinct projected column once and preserves empty-projection row counts.
51    pub(crate) fn build(
52        &self,
53        row_count: usize,
54        mut build_column: impl FnMut(&str) -> Result<ArrayRef, ArrowError>,
55    ) -> Result<DfRecordBatch, ArrowError> {
56        let mut columns: Vec<ArrayRef> = Vec::with_capacity(self.schema.fields().len());
57        for (index, field) in self.schema.fields().iter().enumerate() {
58            let array = match self.duplicates[index] {
59                Some(previous) => columns[previous].clone(),
60                None => build_column(field.name())?,
61            };
62            columns.push(array);
63        }
64        DfRecordBatch::try_new_with_options(
65            self.schema.clone(),
66            columns,
67            &RecordBatchOptions::new().with_row_count(Some(row_count)),
68        )
69    }
70}
71
72/// Builds an inspection scan with filtering before projection and limit.
73pub(crate) fn build_plan(
74    scan_request: ScanRequest,
75    table_name: &str,
76    schema: SchemaRef,
77) -> Result<LogicalPlan, DataFusionError> {
78    let table_source = LogicalTableSource::new(schema.arrow_schema().clone());
79
80    // Filters may reference columns omitted from the requested output. Apply
81    // them to the full schema and let the optimizer push the minimal projection.
82    let projection = scan_request.projection;
83    let mut builder = LogicalPlanBuilder::scan(table_name, Arc::new(table_source), None)?;
84
85    for filter in scan_request.filters {
86        builder = builder.filter(filter)?;
87    }
88
89    if let Some(projection) = projection {
90        let projected = schema.arrow_schema().project(&projection)?;
91        builder = builder.project(
92            projected
93                .fields()
94                .iter()
95                .map(|field| datafusion_expr::col(field.name())),
96        )?;
97    }
98
99    if let Some(limit) = scan_request.limit {
100        builder = builder.limit(0, Some(limit))?;
101    }
102
103    builder.build()
104}
store_api/inspect.rs

store_api/
inspect.rs
