common_grpc_expr/
insert.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::helper;
16use api::v1::column::Values;
17use common_base::BitVec;
18use datatypes::data_type::{ConcreteDataType, DataType};
19use datatypes::prelude::VectorRef;
20use snafu::{ensure, ResultExt};
21
22use crate::error::{CreateVectorSnafu, Result, UnexpectedValuesLengthSnafu};
23
24pub(crate) fn add_values_to_builder(
25    data_type: ConcreteDataType,
26    values: Values,
27    row_count: usize,
28    null_mask: Vec<u8>,
29) -> Result<VectorRef> {
30    if null_mask.is_empty() {
31        Ok(helper::pb_values_to_vector_ref(&data_type, values))
32    } else {
33        let builder = &mut data_type.create_mutable_vector(row_count);
34        let values = helper::pb_values_to_values(&data_type, values);
35        let null_mask = BitVec::from_vec(null_mask);
36        ensure!(
37            null_mask.count_ones() + values.len() == row_count,
38            UnexpectedValuesLengthSnafu {
39                reason: "If null_mask is not empty, the sum of the number of nulls and the length of values must be equal to row_count."
40            }
41        );
42
43        let mut idx_of_values = 0;
44        for idx in 0..row_count {
45            match is_null(&null_mask, idx) {
46                Some(true) => builder.push_null(),
47                _ => {
48                    builder
49                        .try_push_value_ref(values[idx_of_values].as_value_ref())
50                        .context(CreateVectorSnafu)?;
51                    idx_of_values += 1
52                }
53            }
54        }
55        Ok(builder.to_vector())
56    }
57}
58
59fn is_null(null_mask: &BitVec, idx: usize) -> Option<bool> {
60    null_mask.get(idx).as_deref().copied()
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66
67    #[test]
68    fn test_is_null() {
69        let null_mask = BitVec::from_slice(&[0b0000_0001, 0b0000_1000]);
70
71        assert_eq!(Some(true), is_null(&null_mask, 0));
72        assert_eq!(Some(false), is_null(&null_mask, 1));
73        assert_eq!(Some(false), is_null(&null_mask, 10));
74        assert_eq!(Some(true), is_null(&null_mask, 11));
75        assert_eq!(Some(false), is_null(&null_mask, 12));
76
77        assert_eq!(None, is_null(&null_mask, 16));
78        assert_eq!(None, is_null(&null_mask, 99));
79    }
80}
common_grpc_expr/insert.rs

common_grpc_expr/
insert.rs
