common_function/scalars/
primary_key.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::HashMap;
16use std::fmt::{self, Display};
17use std::sync::Arc;
18
19use datafusion_common::arrow::array::{
20    Array, ArrayRef, BinaryArray, BinaryViewArray, DictionaryArray, ListBuilder, StringBuilder,
21};
22use datafusion_common::arrow::datatypes::{DataType, Field};
23use datafusion_common::{DataFusionError, ScalarValue};
24use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility};
25use datatypes::arrow::datatypes::UInt32Type;
26use datatypes::value::Value;
27use mito_codec::row_converter::{
28    CompositeValues, PrimaryKeyCodec, SortField, build_primary_key_codec_with_fields,
29};
30use store_api::codec::PrimaryKeyEncoding;
31use store_api::metadata::RegionMetadata;
32use store_api::storage::ColumnId;
33use store_api::storage::consts::{PRIMARY_KEY_COLUMN_NAME, ReservedColumnId};
34
35use crate::function::{Function, extract_args};
36use crate::function_registry::FunctionRegistry;
37
38type NameValuePair = (String, Option<String>);
39
40#[derive(Clone, Debug)]
41pub(crate) struct DecodePrimaryKeyFunction {
42    signature: Signature,
43}
44
45const NAME: &str = "decode_primary_key";
46const NULL_VALUE_LITERAL: &str = "null";
47
48impl Default for DecodePrimaryKeyFunction {
49    fn default() -> Self {
50        Self {
51            signature: Signature::any(3, Volatility::Immutable),
52        }
53    }
54}
55
56impl DecodePrimaryKeyFunction {
57    pub fn register(registry: &FunctionRegistry) {
58        registry.register_scalar(Self::default());
59    }
60
61    fn return_data_type() -> DataType {
62        DataType::List(Arc::new(Field::new("item", DataType::Utf8, true)))
63    }
64}
65
66impl Function for DecodePrimaryKeyFunction {
67    fn name(&self) -> &str {
68        NAME
69    }
70
71    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
72        Ok(Self::return_data_type())
73    }
74
75    fn signature(&self) -> &Signature {
76        &self.signature
77    }
78
79    fn invoke_with_args(
80        &self,
81        args: ScalarFunctionArgs,
82    ) -> datafusion_common::Result<ColumnarValue> {
83        let [encoded, _, _] = extract_args(self.name(), &args)?;
84        let number_rows = args.number_rows;
85
86        let encoding = parse_encoding(&args.args[1])?;
87        let metadata = parse_region_metadata(&args.args[2])?;
88        let codec = build_codec(&metadata, encoding);
89        let name_lookup: HashMap<_, _> = metadata
90            .column_metadatas
91            .iter()
92            .map(|c| (c.column_id, c.column_schema.name.clone()))
93            .collect();
94
95        let decoded_rows = decode_primary_keys(encoded, number_rows, codec.as_ref(), &name_lookup)?;
96        let array = build_list_array(&decoded_rows)?;
97
98        Ok(ColumnarValue::Array(array))
99    }
100}
101
102impl Display for DecodePrimaryKeyFunction {
103    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
104        write!(f, "DECODE_PRIMARY_KEY")
105    }
106}
107
108fn parse_encoding(arg: &ColumnarValue) -> datafusion_common::Result<PrimaryKeyEncoding> {
109    let encoding = match arg {
110        ColumnarValue::Scalar(ScalarValue::Utf8(Some(v)))
111        | ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(v))) => v.as_str(),
112        ColumnarValue::Scalar(value) => {
113            return Err(DataFusionError::Execution(format!(
114                "encoding must be a string literal, got {value:?}"
115            )));
116        }
117        ColumnarValue::Array(_) => {
118            return Err(DataFusionError::Execution(
119                "encoding must be a scalar string".to_string(),
120            ));
121        }
122    };
123
124    match encoding.to_ascii_lowercase().as_str() {
125        "dense" => Ok(PrimaryKeyEncoding::Dense),
126        "sparse" => Ok(PrimaryKeyEncoding::Sparse),
127        _ => Err(DataFusionError::Execution(format!(
128            "unsupported primary key encoding: {encoding}"
129        ))),
130    }
131}
132
133fn build_codec(
134    metadata: &RegionMetadata,
135    encoding: PrimaryKeyEncoding,
136) -> Arc<dyn PrimaryKeyCodec> {
137    let fields = metadata.primary_key_columns().map(|c| {
138        (
139            c.column_id,
140            SortField::new(c.column_schema.data_type.clone()),
141        )
142    });
143    build_primary_key_codec_with_fields(encoding, fields)
144}
145
146fn parse_region_metadata(arg: &ColumnarValue) -> datafusion_common::Result<RegionMetadata> {
147    let json = match arg {
148        ColumnarValue::Scalar(ScalarValue::Utf8(Some(v)))
149        | ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(v))) => v.as_str(),
150        ColumnarValue::Scalar(value) => {
151            return Err(DataFusionError::Execution(format!(
152                "region metadata must be a string literal, got {value:?}"
153            )));
154        }
155        ColumnarValue::Array(_) => {
156            return Err(DataFusionError::Execution(
157                "region metadata must be a scalar string".to_string(),
158            ));
159        }
160    };
161
162    RegionMetadata::from_json(json)
163        .map_err(|e| DataFusionError::Execution(format!("failed to parse region metadata: {e:?}")))
164}
165
166fn decode_primary_keys(
167    encoded: ArrayRef,
168    number_rows: usize,
169    codec: &dyn PrimaryKeyCodec,
170    name_lookup: &HashMap<ColumnId, String>,
171) -> datafusion_common::Result<Vec<Vec<NameValuePair>>> {
172    if let Some(dict) = encoded
173        .as_any()
174        .downcast_ref::<DictionaryArray<UInt32Type>>()
175    {
176        decode_dictionary(dict, number_rows, codec, name_lookup)
177    } else if let Some(array) = encoded.as_any().downcast_ref::<BinaryArray>() {
178        decode_binary_array(array, codec, name_lookup)
179    } else if let Some(array) = encoded.as_any().downcast_ref::<BinaryViewArray>() {
180        decode_binary_view_array(array, codec, name_lookup)
181    } else {
182        Err(DataFusionError::Execution(format!(
183            "column {PRIMARY_KEY_COLUMN_NAME} must be binary or dictionary(binary) array"
184        )))
185    }
186}
187
188fn decode_dictionary(
189    dict: &DictionaryArray<UInt32Type>,
190    number_rows: usize,
191    codec: &dyn PrimaryKeyCodec,
192    name_lookup: &HashMap<ColumnId, String>,
193) -> datafusion_common::Result<Vec<Vec<NameValuePair>>> {
194    let values = dict
195        .values()
196        .as_any()
197        .downcast_ref::<BinaryArray>()
198        .ok_or_else(|| {
199            DataFusionError::Execution("primary key dictionary values are not binary".to_string())
200        })?;
201
202    let mut decoded_values = Vec::with_capacity(values.len());
203    for i in 0..values.len() {
204        let pk = values.value(i);
205        let pairs = decode_one(pk, codec, name_lookup)?;
206        decoded_values.push(pairs);
207    }
208
209    let mut rows = Vec::with_capacity(number_rows);
210    let keys = dict.keys();
211    let dict_indices = keys.values();
212    for &dict_index in dict_indices[..number_rows].iter() {
213        rows.push(decoded_values[dict_index as usize].clone());
214    }
215
216    Ok(rows)
217}
218
219fn decode_binary_array(
220    array: &BinaryArray,
221    codec: &dyn PrimaryKeyCodec,
222    name_lookup: &HashMap<ColumnId, String>,
223) -> datafusion_common::Result<Vec<Vec<NameValuePair>>> {
224    (0..array.len())
225        .map(|i| decode_one(array.value(i), codec, name_lookup))
226        .collect()
227}
228
229fn decode_binary_view_array(
230    array: &BinaryViewArray,
231    codec: &dyn PrimaryKeyCodec,
232    name_lookup: &HashMap<ColumnId, String>,
233) -> datafusion_common::Result<Vec<Vec<NameValuePair>>> {
234    (0..array.len())
235        .map(|i| decode_one(array.value(i), codec, name_lookup))
236        .collect()
237}
238
239fn decode_one(
240    pk: &[u8],
241    codec: &dyn PrimaryKeyCodec,
242    name_lookup: &HashMap<ColumnId, String>,
243) -> datafusion_common::Result<Vec<NameValuePair>> {
244    let decoded = codec
245        .decode(pk)
246        .map_err(|e| DataFusionError::Execution(format!("failed to decode primary key: {e}")))?;
247
248    Ok(match decoded {
249        CompositeValues::Dense(values) => values
250            .into_iter()
251            .map(|(column_id, value)| (column_name(column_id, name_lookup), value_to_string(value)))
252            .collect(),
253        CompositeValues::Sparse(values) => {
254            let mut values: Vec<_> = values
255                .iter()
256                .map(|(column_id, value)| {
257                    (
258                        *column_id,
259                        column_name(*column_id, name_lookup),
260                        value_to_string(value.clone()),
261                    )
262                })
263                .collect();
264            values.sort_by_key(|(column_id, _, _)| {
265                (ReservedColumnId::is_reserved(*column_id), *column_id)
266            });
267            values
268                .into_iter()
269                .map(|(_, name, value)| (name, value))
270                .collect()
271        }
272    })
273}
274
275fn column_name(column_id: ColumnId, name_lookup: &HashMap<ColumnId, String>) -> String {
276    if let Some(name) = name_lookup.get(&column_id) {
277        return name.clone();
278    }
279
280    if column_id == ReservedColumnId::table_id() {
281        return "__table_id".to_string();
282    }
283    if column_id == ReservedColumnId::tsid() {
284        return "__tsid".to_string();
285    }
286
287    column_id.to_string()
288}
289
290fn value_to_string(value: Value) -> Option<String> {
291    match value {
292        Value::Null => None,
293        _ => Some(value.to_string()),
294    }
295}
296
297fn build_list_array(rows: &[Vec<NameValuePair>]) -> datafusion_common::Result<ArrayRef> {
298    let mut builder = ListBuilder::new(StringBuilder::new());
299
300    for row in rows {
301        for (key, value) in row {
302            let value = value.as_deref().unwrap_or(NULL_VALUE_LITERAL);
303            builder.values().append_value(format!("{key} : {value}"));
304        }
305        builder.append(true);
306    }
307
308    Ok(Arc::new(builder.finish()))
309}
310
311#[cfg(test)]
312mod tests {
313    use api::v1::SemanticType;
314    use datafusion_common::ScalarValue;
315    use datatypes::arrow::array::builder::BinaryDictionaryBuilder;
316    use datatypes::arrow::array::{BinaryArray, ListArray, StringArray};
317    use datatypes::arrow::datatypes::UInt32Type;
318    use datatypes::prelude::ConcreteDataType;
319    use datatypes::schema::ColumnSchema;
320    use datatypes::value::Value;
321    use mito_codec::row_converter::{
322        DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField, SparsePrimaryKeyCodec,
323    };
324    use store_api::codec::PrimaryKeyEncoding;
325    use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
326    use store_api::storage::consts::ReservedColumnId;
327    use store_api::storage::{ColumnId, RegionId};
328
329    use super::*;
330
331    fn pk_field() -> Arc<Field> {
332        Arc::new(Field::new_dictionary(
333            PRIMARY_KEY_COLUMN_NAME,
334            DataType::UInt32,
335            DataType::Binary,
336            false,
337        ))
338    }
339
340    fn region_metadata_json(
341        columns: &[(ColumnId, &str, ConcreteDataType)],
342        encoding: PrimaryKeyEncoding,
343    ) -> String {
344        let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1));
345        builder.push_column_metadata(ColumnMetadata {
346            column_schema: ColumnSchema::new(
347                "ts",
348                ConcreteDataType::timestamp_millisecond_datatype(),
349                false,
350            ),
351            semantic_type: SemanticType::Timestamp,
352            column_id: 100,
353        });
354        builder.primary_key_encoding(encoding);
355        for (id, name, ty) in columns {
356            builder.push_column_metadata(ColumnMetadata {
357                column_schema: ColumnSchema::new(
358                    (*name).to_string(),
359                    ty.clone(),
360                    !ReservedColumnId::is_reserved(*id),
361                ),
362                semantic_type: SemanticType::Tag,
363                column_id: *id,
364            });
365        }
366        builder.primary_key(columns.iter().map(|(id, _, _)| *id).collect());
367
368        builder.build().unwrap().to_json().unwrap()
369    }
370
371    fn list_row(list: &ListArray, row_idx: usize) -> Vec<String> {
372        let values = list.value(row_idx);
373        let values = values.as_any().downcast_ref::<StringArray>().unwrap();
374        (0..values.len())
375            .map(|i| values.value(i).to_string())
376            .collect()
377    }
378
379    #[test]
380    fn test_decode_dense_primary_key() {
381        let columns = vec![
382            (0, "host", ConcreteDataType::string_datatype()),
383            (1, "core", ConcreteDataType::int64_datatype()),
384        ];
385        let metadata_json = region_metadata_json(&columns, PrimaryKeyEncoding::Dense);
386        let codec = DensePrimaryKeyCodec::with_fields(
387            columns
388                .iter()
389                .map(|(id, _, ty)| (*id, SortField::new(ty.clone())))
390                .collect(),
391        );
392
393        let rows = vec![
394            vec![Value::from("a"), Value::from(1_i64)],
395            vec![Value::from("b"), Value::from(2_i64)],
396            vec![Value::from("a"), Value::from(1_i64)],
397        ];
398
399        let mut builder = BinaryDictionaryBuilder::<UInt32Type>::new();
400        for row in &rows {
401            let encoded = codec.encode(row.iter().map(|v| v.as_value_ref())).unwrap();
402            builder.append(encoded.as_slice()).unwrap();
403        }
404        let dict_array: ArrayRef = Arc::new(builder.finish());
405
406        let args = ScalarFunctionArgs {
407            args: vec![
408                ColumnarValue::Array(dict_array),
409                ColumnarValue::Scalar(ScalarValue::Utf8(Some("dense".to_string()))),
410                ColumnarValue::Scalar(ScalarValue::Utf8(Some(metadata_json))),
411            ],
412            arg_fields: vec![
413                pk_field(),
414                Arc::new(Field::new("encoding", DataType::Utf8, false)),
415                Arc::new(Field::new("region_metadata", DataType::Utf8, false)),
416            ],
417            number_rows: 3,
418            return_field: Arc::new(Field::new(
419                "decoded",
420                DecodePrimaryKeyFunction::return_data_type(),
421                false,
422            )),
423            config_options: Default::default(),
424        };
425
426        let func = DecodePrimaryKeyFunction::default();
427        let result = func
428            .invoke_with_args(args)
429            .and_then(|v| v.to_array(3))
430            .unwrap();
431        let list = result.as_any().downcast_ref::<ListArray>().unwrap();
432
433        let expected = [
434            vec!["host : a".to_string(), "core : 1".to_string()],
435            vec!["host : b".to_string(), "core : 2".to_string()],
436            vec!["host : a".to_string(), "core : 1".to_string()],
437        ];
438
439        for (row_idx, expected_row) in expected.iter().enumerate() {
440            assert_eq!(*expected_row, list_row(list, row_idx));
441        }
442    }
443
444    #[test]
445    fn test_decode_sparse_primary_key() {
446        let columns = vec![
447            (
448                ReservedColumnId::table_id(),
449                "__table_id",
450                ConcreteDataType::uint32_datatype(),
451            ),
452            (
453                ReservedColumnId::tsid(),
454                "__tsid",
455                ConcreteDataType::uint64_datatype(),
456            ),
457            (10, "k0", ConcreteDataType::string_datatype()),
458            (11, "k1", ConcreteDataType::string_datatype()),
459        ];
460        let metadata_json = region_metadata_json(&columns, PrimaryKeyEncoding::Sparse);
461        let codec = SparsePrimaryKeyCodec::schemaless();
462
463        let rows = vec![
464            vec![
465                (ReservedColumnId::table_id(), Value::UInt32(1)),
466                (ReservedColumnId::tsid(), Value::UInt64(100)),
467                (10, Value::from("a")),
468                (11, Value::from("b")),
469            ],
470            vec![
471                (ReservedColumnId::table_id(), Value::UInt32(1)),
472                (ReservedColumnId::tsid(), Value::UInt64(200)),
473                (10, Value::from("c")),
474                (11, Value::from("d")),
475            ],
476        ];
477
478        let mut encoded_values = Vec::with_capacity(rows.len());
479        for row in &rows {
480            let mut buf = Vec::new();
481            codec.encode_values(row, &mut buf).unwrap();
482            encoded_values.push(buf);
483        }
484
485        let pk_array: ArrayRef = Arc::new(BinaryArray::from_iter_values(
486            encoded_values.iter().cloned(),
487        ));
488
489        let args = ScalarFunctionArgs {
490            args: vec![
491                ColumnarValue::Array(pk_array),
492                ColumnarValue::Scalar(ScalarValue::Utf8(Some("sparse".to_string()))),
493                ColumnarValue::Scalar(ScalarValue::Utf8(Some(metadata_json))),
494            ],
495            arg_fields: vec![
496                pk_field(),
497                Arc::new(Field::new("encoding", DataType::Utf8, false)),
498                Arc::new(Field::new("region_metadata", DataType::Utf8, false)),
499            ],
500            number_rows: rows.len(),
501            return_field: Arc::new(Field::new(
502                "decoded",
503                DecodePrimaryKeyFunction::return_data_type(),
504                false,
505            )),
506            config_options: Default::default(),
507        };
508
509        let func = DecodePrimaryKeyFunction::default();
510        let result = func
511            .invoke_with_args(args)
512            .and_then(|v| v.to_array(rows.len()))
513            .unwrap();
514        let list = result.as_any().downcast_ref::<ListArray>().unwrap();
515
516        let expected = [
517            vec![
518                "k0 : a".to_string(),
519                "k1 : b".to_string(),
520                "__tsid : 100".to_string(),
521                "__table_id : 1".to_string(),
522            ],
523            vec![
524                "k0 : c".to_string(),
525                "k1 : d".to_string(),
526                "__tsid : 200".to_string(),
527                "__table_id : 1".to_string(),
528            ],
529        ];
530
531        for (row_idx, expected_row) in expected.iter().enumerate() {
532            assert_eq!(*expected_row, list_row(list, row_idx));
533        }
534    }
535}
common_function/scalars/primary_key.rs

common_function/scalars/
primary_key.rs
