Skip to main content

mito2/read/
read_columns.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::collections::BTreeMap;
16use std::hash::Hash;
17use std::mem;
18use std::sync::Arc;
19
20use datatypes::types::json_type::JsonNativeType;
21use store_api::storage::ColumnId;
22
23pub(crate) type JsonTargetTypes = Arc<BTreeMap<ColumnId, JsonNativeType>>;
24
25/// Logical columns to read from a region.
26///
27/// Read columns describe which logical root columns should be read from storage.
28/// JSON2 columns can carry query-time target types that are later translated to
29/// physical nested parquet paths by the parquet reader.
30#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
31pub struct ReadColumns {
32    pub col_ids: Vec<ColumnId>,
33    json_target_types: JsonTargetTypes,
34}
35
36impl ReadColumns {
37    /// Creates read columns from logical column ids.
38    ///
39    /// This preserves the input order and duplicate ids.
40    pub fn new<I>(col_ids: I) -> Self
41    where
42        I: IntoIterator<Item = ColumnId>,
43    {
44        Self {
45            col_ids: col_ids.into_iter().collect(),
46            json_target_types: Arc::default(),
47        }
48    }
49
50    /// Attaches query-time JSON2 projection types.
51    pub fn with_json_target_types(
52        mut self,
53        json_target_types: BTreeMap<ColumnId, JsonNativeType>,
54    ) -> Self {
55        self.json_target_types = Arc::new(json_target_types);
56        self
57    }
58
59    pub fn is_empty(&self) -> bool {
60        self.col_ids.is_empty()
61    }
62
63    pub fn column_ids_iter(&self) -> impl Iterator<Item = ColumnId> + '_ {
64        self.col_ids.iter().copied()
65    }
66
67    pub fn column_ids(&self) -> Vec<ColumnId> {
68        self.column_ids_iter().collect()
69    }
70
71    pub(crate) fn json_target_types(&self) -> &JsonTargetTypes {
72        &self.json_target_types
73    }
74
75    /// Returns the query-time JSON2 projection type for a column.
76    pub fn json_target_type(&self, column_id: ColumnId) -> Option<&JsonNativeType> {
77        self.json_target_types.get(&column_id)
78    }
79
80    pub fn estimated_size(&self) -> usize {
81        self.col_ids.capacity() * mem::size_of::<ColumnId>()
82            + self.col_ids.len() * mem::size_of::<ColumnId>()
83            + self.json_target_types.len() * (size_of::<ColumnId>() + size_of::<JsonNativeType>())
84    }
85}