mito2/read/
projection.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
15//! Projection helpers shared by flat projection code.
16
17use std::cmp::Ordering;
18
19use common_recordbatch::error::DataTypesSnafu;
20use datatypes::prelude::{ConcreteDataType, DataType};
21use datatypes::value::Value;
22use datatypes::vectors::{Helper, VectorRef};
23use snafu::{OptionExt, ResultExt};
24use store_api::metadata::RegionMetadataRef;
25use store_api::storage::ColumnId;
26
27use crate::cache::CacheStrategy;
28use crate::error::{InvalidRequestSnafu, Result};
29
30/// Only cache vector when its length `<=` this value.
31pub(crate) const MAX_VECTOR_LENGTH_TO_CACHE: usize = 16384;
32
33pub(crate) fn read_column_ids_from_projection(
34    metadata: &RegionMetadataRef,
35    projection: &[usize],
36) -> Result<Vec<ColumnId>> {
37    let mut column_ids = Vec::with_capacity(projection.len().max(1));
38    if projection.is_empty() {
39        column_ids.push(metadata.time_index_column().column_id);
40        return Ok(column_ids);
41    }
42
43    for idx in projection {
44        let column = metadata
45            .column_metadatas
46            .get(*idx)
47            .with_context(|| InvalidRequestSnafu {
48                region_id: metadata.region_id,
49                reason: format!("projection index {} is out of bound", idx),
50            })?;
51        column_ids.push(column.column_id);
52    }
53    Ok(column_ids)
54}
55
56/// Gets a vector with repeated values from specific cache or creates a new one.
57pub(crate) fn repeated_vector_with_cache(
58    data_type: &ConcreteDataType,
59    value: &Value,
60    num_rows: usize,
61    cache_strategy: &CacheStrategy,
62) -> common_recordbatch::error::Result<VectorRef> {
63    if let Some(vector) = cache_strategy.get_repeated_vector(data_type, value) {
64        match vector.len().cmp(&num_rows) {
65            Ordering::Less => (),
66            Ordering::Equal => return Ok(vector),
67            Ordering::Greater => return Ok(vector.slice(0, num_rows)),
68        }
69    }
70
71    let vector = new_repeated_vector(data_type, value, num_rows)?;
72    if vector.len() <= MAX_VECTOR_LENGTH_TO_CACHE {
73        cache_strategy.put_repeated_vector(value.clone(), vector.clone());
74    }
75
76    Ok(vector)
77}
78
79/// Returns a vector with repeated values.
80pub(crate) fn new_repeated_vector(
81    data_type: &ConcreteDataType,
82    value: &Value,
83    num_rows: usize,
84) -> common_recordbatch::error::Result<VectorRef> {
85    if let Ok(vector) = value
86        .try_to_scalar_value(data_type)
87        .and_then(|scalar| Helper::try_from_scalar_value(scalar, num_rows, Some(data_type)))
88    {
89        return Ok(vector);
90    }
91
92    // Preserve extension types that cannot safely round-trip through ScalarValue.
93    let mut mutable_vector = data_type.create_mutable_vector(num_rows);
94    for _ in 0..num_rows {
95        mutable_vector
96            .try_push_value_ref(&value.as_value_ref())
97            .context(DataTypesSnafu)?;
98    }
99    Ok(mutable_vector.to_vector())
100}
101
102#[cfg(test)]
103mod tests {
104    use std::sync::Arc;
105
106    use api::v1::OpType;
107    use common_recordbatch::RecordBatch;
108    use datatypes::arrow::array::{Int64Array, TimestampMillisecondArray, UInt8Array, UInt64Array};
109    use datatypes::arrow::datatypes::Field;
110    use datatypes::arrow::util::pretty;
111    use datatypes::value::ValueRef;
112    use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField};
113    use mito_codec::test_util::TestRegionMetadataBuilder;
114    use store_api::storage::consts::{
115        OP_TYPE_COLUMN_NAME, PRIMARY_KEY_COLUMN_NAME, SEQUENCE_COLUMN_NAME,
116    };
117
118    use super::*;
119    use crate::read::flat_projection::FlatProjectionMapper;
120    use crate::read::read_columns::ReadColumns;
121
122    #[test]
123    fn test_repeated_struct_null_fields_and_json() {
124        use datatypes::types::{StructField, StructType};
125        use datatypes::value::StructValue;
126
127        let inner_type = StructType::from([StructField::new(
128            "x",
129            ConcreteDataType::int32_datatype(),
130            true,
131        )]);
132        let inner = Value::Struct(StructValue::new(vec![Value::Null], inner_type));
133        let json = datatypes::json::JsonSettings::default()
134            .encode(serde_json::json!({"answer": 42}))
135            .unwrap();
136        let nested_type = StructType::from([StructField::new("nested", inner.data_type(), true)]);
137        let json_type = StructType::from([StructField::new("json", json.data_type(), true)]);
138        let values = [
139            inner.clone(),
140            Value::Struct(StructValue::new(vec![inner], nested_type)),
141            Value::Struct(StructValue::new(vec![], StructType::default())),
142            Value::Struct(StructValue::new(vec![json], json_type)),
143        ];
144        for value in values {
145            let data_type = value.data_type();
146            // JSON children are read back as their underlying struct values.
147            let expected = serde_json::Value::try_from(value.clone()).unwrap();
148            for num_rows in [0, 1, 3] {
149                let vector = new_repeated_vector(&data_type, &value, num_rows).unwrap();
150                assert_eq!(data_type, vector.data_type());
151                assert_eq!(
152                    data_type.as_arrow_type(),
153                    *vector.to_arrow_array().data_type()
154                );
155                assert_eq!(num_rows, vector.len());
156                assert_eq!(0, vector.null_count());
157                for row in 0..num_rows {
158                    assert_eq!(
159                        expected,
160                        serde_json::Value::try_from(vector.get(row)).unwrap()
161                    );
162                }
163            }
164        }
165    }
166
167    fn print_record_batch(record_batch: RecordBatch) -> String {
168        pretty::pretty_format_batches(&[record_batch.into_df_record_batch()])
169            .unwrap()
170            .to_string()
171    }
172
173    fn new_flat_batch(
174        ts_start: Option<i64>,
175        idx_tags: &[(usize, i64)],
176        idx_fields: &[(usize, i64)],
177        num_rows: usize,
178    ) -> datatypes::arrow::record_batch::RecordBatch {
179        let mut columns = Vec::with_capacity(1 + idx_tags.len() + idx_fields.len() + 3);
180        let mut fields = Vec::with_capacity(1 + idx_tags.len() + idx_fields.len() + 3);
181
182        for (i, tag) in idx_tags {
183            let array = Arc::new(Int64Array::from_iter_values(std::iter::repeat_n(
184                *tag, num_rows,
185            ))) as _;
186            columns.push(array);
187            fields.push(Field::new(
188                format!("k{i}"),
189                datatypes::arrow::datatypes::DataType::Int64,
190                true,
191            ));
192        }
193
194        for (i, field) in idx_fields {
195            let array = Arc::new(Int64Array::from_iter_values(std::iter::repeat_n(
196                *field, num_rows,
197            ))) as _;
198            columns.push(array);
199            fields.push(Field::new(
200                format!("v{i}"),
201                datatypes::arrow::datatypes::DataType::Int64,
202                true,
203            ));
204        }
205
206        if let Some(ts_start) = ts_start {
207            let timestamps = Arc::new(TimestampMillisecondArray::from_iter_values(
208                (0..num_rows).map(|i| ts_start + i as i64 * 1000),
209            )) as _;
210            columns.push(timestamps);
211            fields.push(Field::new(
212                "ts",
213                datatypes::arrow::datatypes::DataType::Timestamp(
214                    datatypes::arrow::datatypes::TimeUnit::Millisecond,
215                    None,
216                ),
217                true,
218            ));
219        }
220
221        let converter = DensePrimaryKeyCodec::with_fields(
222            (0..idx_tags.len())
223                .map(|idx| {
224                    (
225                        idx as u32,
226                        SortField::new(ConcreteDataType::int64_datatype()),
227                    )
228                })
229                .collect(),
230        );
231        let encoded_pk = converter
232            .encode(idx_tags.iter().map(|(_, v)| ValueRef::Int64(*v)))
233            .unwrap();
234
235        let pk_values: Vec<&[u8]> = std::iter::repeat_n(encoded_pk.as_slice(), num_rows).collect();
236        let keys = datatypes::arrow::array::UInt32Array::from_iter(0..num_rows as u32);
237        let values = Arc::new(datatypes::arrow::array::BinaryArray::from_vec(pk_values));
238        let pk_array =
239            Arc::new(datatypes::arrow::array::DictionaryArray::try_new(keys, values).unwrap()) as _;
240        columns.push(pk_array);
241        fields.push(Field::new_dictionary(
242            PRIMARY_KEY_COLUMN_NAME,
243            datatypes::arrow::datatypes::DataType::UInt32,
244            datatypes::arrow::datatypes::DataType::Binary,
245            false,
246        ));
247
248        columns.push(Arc::new(UInt64Array::from_iter_values(0..num_rows as u64)) as _);
249        fields.push(Field::new(
250            SEQUENCE_COLUMN_NAME,
251            datatypes::arrow::datatypes::DataType::UInt64,
252            false,
253        ));
254
255        columns.push(Arc::new(UInt8Array::from_iter_values(
256            (0..num_rows).map(|_| OpType::Put as u8),
257        )) as _);
258        fields.push(Field::new(
259            OP_TYPE_COLUMN_NAME,
260            datatypes::arrow::datatypes::DataType::UInt8,
261            false,
262        ));
263
264        let schema = Arc::new(datatypes::arrow::datatypes::Schema::new(fields));
265
266        datatypes::arrow::record_batch::RecordBatch::try_new(schema, columns).unwrap()
267    }
268
269    #[test]
270    fn test_flat_projection_mapper_all() {
271        let metadata = Arc::new(
272            TestRegionMetadataBuilder::default()
273                .num_tags(2)
274                .num_fields(2)
275                .build(),
276        );
277        let cache = CacheStrategy::Disabled;
278        let mapper = FlatProjectionMapper::all(&metadata).unwrap();
279        assert_eq!(
280            &[0, 1, 2, 3, 4],
281            mapper.read_columns().column_ids().as_slice()
282        );
283        assert_eq!(
284            [
285                (1, ConcreteDataType::int64_datatype()),
286                (2, ConcreteDataType::int64_datatype()),
287                (3, ConcreteDataType::int64_datatype()),
288                (4, ConcreteDataType::int64_datatype()),
289                (0, ConcreteDataType::timestamp_millisecond_datatype())
290            ],
291            mapper.batch_schema()
292        );
293
294        let batch = new_flat_batch(Some(0), &[(1, 1), (2, 2)], &[(3, 3), (4, 4)], 3);
295        let record_batch = mapper.convert(&batch, &cache).unwrap();
296        let expect = "\
297+---------------------+----+----+----+----+
298| ts                  | k0 | k1 | v0 | v1 |
299+---------------------+----+----+----+----+
300| 1970-01-01T00:00:00 | 1  | 2  | 3  | 4  |
301| 1970-01-01T00:00:01 | 1  | 2  | 3  | 4  |
302| 1970-01-01T00:00:02 | 1  | 2  | 3  | 4  |
303+---------------------+----+----+----+----+";
304        assert_eq!(expect, print_record_batch(record_batch));
305    }
306
307    #[test]
308    fn test_flat_projection_mapper_with_projection() {
309        let metadata = Arc::new(
310            TestRegionMetadataBuilder::default()
311                .num_tags(2)
312                .num_fields(2)
313                .build(),
314        );
315        let cache = CacheStrategy::Disabled;
316        let mapper = FlatProjectionMapper::new(&metadata, [4, 1]).unwrap();
317        assert_eq!(&[4, 1], mapper.read_columns().column_ids().as_slice());
318        assert_eq!(
319            [
320                (1, ConcreteDataType::int64_datatype()),
321                (4, ConcreteDataType::int64_datatype()),
322                (0, ConcreteDataType::timestamp_millisecond_datatype())
323            ],
324            mapper.batch_schema()
325        );
326
327        let batch = new_flat_batch(None, &[(1, 1)], &[(4, 4)], 3);
328        let record_batch = mapper.convert(&batch, &cache).unwrap();
329        let expect = "\
330+----+----+
331| v1 | k0 |
332+----+----+
333| 4  | 1  |
334| 4  | 1  |
335| 4  | 1  |
336+----+----+";
337        assert_eq!(expect, print_record_batch(record_batch));
338    }
339
340    #[test]
341    fn test_flat_projection_mapper_read_superset() {
342        let metadata = Arc::new(
343            TestRegionMetadataBuilder::default()
344                .num_tags(2)
345                .num_fields(2)
346                .build(),
347        );
348        let cache = CacheStrategy::Disabled;
349        let mapper = FlatProjectionMapper::new_with_read_columns(
350            &metadata,
351            vec![4, 1],
352            ReadColumns::new([4, 1, 3]),
353        )
354        .unwrap();
355        assert_eq!(&[4, 1, 3], mapper.read_columns().column_ids().as_slice());
356
357        let batch = new_flat_batch(None, &[(1, 1)], &[(3, 3), (4, 4)], 3);
358        let record_batch = mapper.convert(&batch, &cache).unwrap();
359        let expect = "\
360+----+----+
361| v1 | k0 |
362+----+----+
363| 4  | 1  |
364| 4  | 1  |
365| 4  | 1  |
366+----+----+";
367        assert_eq!(expect, print_record_batch(record_batch));
368    }
369
370    #[test]
371    fn test_flat_projection_mapper_empty_projection() {
372        let metadata = Arc::new(
373            TestRegionMetadataBuilder::default()
374                .num_tags(2)
375                .num_fields(2)
376                .build(),
377        );
378        let cache = CacheStrategy::Disabled;
379        let mapper = FlatProjectionMapper::new(&metadata, []).unwrap();
380        assert_eq!(&[0], mapper.read_columns().column_ids().as_slice());
381        assert!(mapper.output_schema().is_empty());
382        assert_eq!(
383            [(0, ConcreteDataType::timestamp_millisecond_datatype())],
384            mapper.batch_schema()
385        );
386
387        let batch = new_flat_batch(Some(0), &[], &[], 3);
388        let record_batch = mapper.convert(&batch, &cache).unwrap();
389        assert_eq!(3, record_batch.num_rows());
390        assert_eq!(0, record_batch.num_columns());
391        assert!(record_batch.schema.is_empty());
392    }
393}
mito2/read/projection.rs

mito2/read/
projection.rs
