mito2/sst/index/
column.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 api::v1::SemanticType;
16use datatypes::arrow::array::Array;
17use datatypes::arrow::record_batch::RecordBatch;
18use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;
19
20use crate::sst::parquet::format::PrimaryKeyArray;
21
22/// Yields (first row, row count) for indexing a materialized column.
23/// Tags are constant within consecutive equal PK dictionary keys. Fields and
24/// timestamps must still be visited row by row. Distinct dictionary entries may
25/// contain equal PK bytes; keeping those runs separate is correct and avoids decoding.
26/// Inputs without a non-null PK dictionary fall back to visiting individual rows.
27pub(crate) fn column_index_rows(
28    batch: &RecordBatch,
29    semantic_type: SemanticType,
30) -> impl Iterator<Item = (usize, usize)> + '_ {
31    let keys = (semantic_type == SemanticType::Tag)
32        .then(|| batch.column_by_name(PRIMARY_KEY_COLUMN_NAME))
33        .flatten()
34        .and_then(|array| array.as_any().downcast_ref::<PrimaryKeyArray>())
35        .filter(|array| array.null_count() == 0 && array.values().null_count() == 0)
36        .map(|array| array.keys().values());
37    let mut row = 0;
38    std::iter::from_fn(move || {
39        if row == batch.num_rows() {
40            return None;
41        }
42        let start = row;
43        row += 1;
44        if let Some(keys) = keys {
45            while row < keys.len() && keys[row] == keys[start] {
46                row += 1;
47            }
48        }
49        Some((start, row - start))
50    })
51}
mito2/sst/index/column.rs

mito2/sst/index/
column.rs
