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 arrow_schema::extension::ExtensionType;
17use common_function::scalars::json::json_get::{JsonGetWithType, json_object_path};
18use datafusion_expr::Expr;
19use datatypes::arrow::array::{Array, ArrayRef, StructArray};
20use datatypes::arrow::record_batch::RecordBatch;
21use datatypes::data_type::{ConcreteDataType, DataType};
22use datatypes::extension::json::Json2ExtensionType;
23use datatypes::json::JsonTypeHint;
24use datatypes::schema::FulltextOptions;
25use datatypes::vectors::{Helper, VectorRef};
26use index::target::IndexTarget;
27use snafu::{OptionExt, ResultExt};
28use store_api::metadata::RegionMetadata;
29use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;
30
31use crate::error::{ConvertVectorSnafu, IndexOptionsSnafu, InvalidRecordBatchSnafu, Result};
32use crate::sst::parquet::format::PrimaryKeyArray;
33
34/// Extracts leaf values for a JSON type hint at the given path.
35///
36/// Returns an error if the path does not exist, traverses a non-object,
37/// or the leaf array cannot be converted to a vector.
38///
39/// # Details
40///
41/// The caller guarantees that the materialized leaf type matches the type hint.
42///
43/// JSON2 builders propagate ancestor nulls to child arrays. SST reads and schema
44/// adaptation preserve this invariant, so no ancestor null bitmap merging is needed.
45pub(crate) fn extract_json_type_hint_values(root: &ArrayRef, path: &[String]) -> Result<VectorRef> {
46    let mut array = root.clone();
47    for segment in path {
48        let object = array
49            .as_any()
50            .downcast_ref::<StructArray>()
51            .with_context(|| InvalidRecordBatchSnafu {
52                reason: format!("JSON index path {path:?} traverses a non-object"),
53            })?;
54        array = object
55            .column_by_name(segment)
56            .with_context(|| InvalidRecordBatchSnafu {
57                reason: format!("JSON index path {path:?} has no materialized field {segment}"),
58            })?
59            .clone();
60    }
61    Helper::try_into_vector(array).context(ConvertVectorSnafu)
62}
63
64/// Yields (first row, row count) for indexing a materialized column.
65/// Tags are constant within consecutive equal PK dictionary keys. Fields and
66/// timestamps must still be visited row by row. Distinct dictionary entries may
67/// contain equal PK bytes; keeping those runs separate is correct and avoids decoding.
68/// Inputs without a non-null PK dictionary fall back to visiting individual rows.
69pub(crate) fn column_index_rows(
70    batch: &RecordBatch,
71    semantic_type: SemanticType,
72) -> impl Iterator<Item = (usize, usize)> + '_ {
73    let keys = (semantic_type == SemanticType::Tag)
74        .then(|| batch.column_by_name(PRIMARY_KEY_COLUMN_NAME))
75        .flatten()
76        .and_then(|array| array.as_any().downcast_ref::<PrimaryKeyArray>())
77        .filter(|array| array.null_count() == 0 && array.values().null_count() == 0)
78        .map(|array| array.keys().values());
79    let mut row = 0;
80    std::iter::from_fn(move || {
81        if row == batch.num_rows() {
82            return None;
83        }
84        let start = row;
85        row += 1;
86        if let Some(keys) = keys {
87            while row < keys.len() && keys[row] == keys[start] {
88                row += 1;
89            }
90        }
91        Some((start, row - start))
92    })
93}
94
95/// Resolves a typed JSON access when its hint enables the requested index.
96pub(crate) fn json_index_target(
97    metadata: &RegionMetadata,
98    expr: &Expr,
99    enabled: impl Fn(&JsonTypeHint) -> bool,
100) -> Result<Option<(IndexTarget, ConcreteDataType)>> {
101    let Expr::ScalarFunction(func) = expr else {
102        return Ok(None);
103    };
104
105    if func.func.name() != JsonGetWithType::NAME || !(2..=3).contains(&func.args.len()) {
106        return Ok(None);
107    }
108
109    let Expr::Column(column) = &func.args[0] else {
110        return Ok(None);
111    };
112
113    let Some(col) = metadata.column_by_name(&column.name) else {
114        return Ok(None);
115    };
116
117    if !col.column_schema.data_type.is_json2() {
118        return Ok(None);
119    }
120
121    let Some(path_literal) = func.args[1]
122        .as_literal()
123        .and_then(|v| v.try_as_str())
124        .flatten()
125    else {
126        return Ok(None);
127    };
128
129    let Some(path) = json_object_path(path_literal).ok().flatten() else {
130        return Ok(None);
131    };
132
133    let Ok(Some(extension)) = col.column_schema.extension_type::<Json2ExtensionType>() else {
134        return Ok(None);
135    };
136
137    let Some(hint) = extension
138        .metadata()
139        .json_settings()
140        .type_hints()
141        .iter()
142        .find(|hint| hint.path == path && enabled(hint))
143    else {
144        return Ok(None);
145    };
146
147    if let Some(type_arg) = func.args.get(2) {
148        let Expr::Literal(value, _) = type_arg else {
149            return Ok(None);
150        };
151        if !value.is_null() {
152            return Ok(None);
153        }
154        // A type conversion may change comparison semantics and invalidate index pruning.
155        let result_type = value.data_type();
156        if result_type != hint.data_type.as_arrow_type()
157            && !(result_type.is_string() && hint.data_type.is_string())
158        {
159            return Ok(None);
160        }
161    }
162
163    if !(hint.data_type.is_numeric() || hint.data_type.is_boolean() || hint.data_type.is_string()) {
164        return Ok(None);
165    }
166    Ok(IndexTarget::new_json_path(col.column_id, path)
167        .ok()
168        .map(|target| (target, hint.data_type.clone())))
169}
170
171/// Resolves a skipping target's encoding type from an SST's own hints.
172pub(crate) fn skipping_target_data_type(
173    metadata: &RegionMetadata,
174    target: &IndexTarget,
175) -> Option<ConcreteDataType> {
176    let path = match target {
177        IndexTarget::ColumnId(id) => {
178            return metadata
179                .column_by_id(*id)
180                .map(|col| col.column_schema.data_type.clone());
181        }
182        IndexTarget::JsonPath(path) => path,
183    };
184    let column = metadata.column_by_id(path.column_id())?;
185    let extension = column
186        .column_schema
187        .extension_type::<Json2ExtensionType>()
188        .ok()??;
189    extension
190        .metadata()
191        .json_settings()
192        .type_hints()
193        .iter()
194        .find(|hint| hint.path == path.path() && hint.skipping_index.is_some())
195        .map(|hint| hint.data_type.clone())
196}
197
198/// Resolves enabled fulltext options using a column or string hint's metadata.
199pub(crate) fn fulltext_target_options(
200    metadata: &RegionMetadata,
201    target: &IndexTarget,
202) -> Result<Option<FulltextOptions>> {
203    let path = match target {
204        IndexTarget::ColumnId(id) => {
205            let Some(column) = metadata.column_by_id(*id) else {
206                return Ok(None);
207            };
208            if !column.column_schema.data_type.is_string() {
209                return Ok(None);
210            }
211            return column
212                .column_schema
213                .fulltext_options()
214                .with_context(|_| IndexOptionsSnafu {
215                    column_name: &column.column_schema.name,
216                })
217                .map(|options| options.filter(|options| options.enable));
218        }
219        IndexTarget::JsonPath(path) => path,
220    };
221    let Some(column) = metadata.column_by_id(path.column_id()) else {
222        return Ok(None);
223    };
224    if !column.column_schema.data_type.is_json2() {
225        return Ok(None);
226    }
227    let Some(extension) = column
228        .column_schema
229        .extension_type::<Json2ExtensionType>()
230        .with_context(|_| IndexOptionsSnafu {
231            column_name: &column.column_schema.name,
232        })?
233    else {
234        return Ok(None);
235    };
236    Ok(extension
237        .metadata()
238        .json_settings()
239        .type_hints()
240        .iter()
241        .find(|hint| hint.path == path.path() && hint.data_type.is_string())
242        .and_then(|hint| hint.fulltext_index.clone())
243        .filter(|options| options.enable))
244}
245
246#[cfg(test)]
247mod tests {
248    use std::sync::Arc;
249
250    use datatypes::arrow::array::{Int64Array, NullArray};
251    use datatypes::arrow::buffer::NullBuffer;
252    use datatypes::arrow::datatypes::Field;
253    use datatypes::schema::{ColumnSchema, FulltextOptions};
254    use datatypes::value::Value;
255    use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
256    use store_api::storage::RegionId;
257
258    use super::*;
259
260    #[test]
261    fn test_fulltext_target_options_requires_string_column() {
262        let options = FulltextOptions {
263            enable: true,
264            ..Default::default()
265        };
266        for data_type in [
267            ConcreteDataType::string_datatype(),
268            ConcreteDataType::int64_datatype(),
269        ] {
270            let expected = data_type.is_string().then_some(options.clone());
271            let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1));
272            builder.push_column_metadata(ColumnMetadata {
273                column_schema: ColumnSchema::new("value", data_type, true)
274                    .with_fulltext_options(options.clone())
275                    .unwrap(),
276                column_id: 1,
277                semantic_type: SemanticType::Field,
278            });
279            builder.push_column_metadata(ColumnMetadata {
280                column_schema: ColumnSchema::new(
281                    "ts",
282                    ConcreteDataType::timestamp_millisecond_datatype(),
283                    false,
284                ),
285                column_id: 2,
286                semantic_type: SemanticType::Timestamp,
287            });
288            let metadata = builder.build().unwrap();
289            assert_eq!(
290                fulltext_target_options(&metadata, &IndexTarget::ColumnId(1)).unwrap(),
291                expected
292            );
293        }
294    }
295
296    #[test]
297    fn test_extract_json_type_hint_values_null_type() {
298        let nulls: ArrayRef = Arc::new(NullArray::new(2));
299        let root: ArrayRef = Arc::new(StructArray::new(
300            vec![Field::new("a", nulls.data_type().clone(), true)].into(),
301            vec![nulls],
302            None,
303        ));
304        // A null leaf can be extracted, but cannot be traversed as an object.
305        assert!(extract_json_type_hint_values(&root, &["a".into(), "b".into()]).is_err());
306        let leaf = extract_json_type_hint_values(&root, &["a".into()]).unwrap();
307        assert_eq!(leaf.len(), 2);
308        assert_eq!(leaf.get(0), Value::Null);
309        assert_eq!(leaf.get(1), Value::Null);
310
311        let values: ArrayRef = Arc::new(Int64Array::from(vec![None, None]));
312        let root: ArrayRef = Arc::new(StructArray::new(
313            vec![Field::new("a", values.data_type().clone(), true)].into(),
314            vec![values],
315            None,
316        ));
317        let leaf = extract_json_type_hint_values(&root, &["a".into()]).unwrap();
318        assert_eq!(leaf.len(), 2);
319        assert_eq!(leaf.get(0), Value::Null);
320        assert_eq!(leaf.get(1), Value::Null);
321    }
322
323    #[test]
324    fn test_extract_json_type_hint_values_nulls_and_slices() {
325        let child = Arc::new(Int64Array::from(vec![Some(11), None, None, None, Some(55)]));
326        let nested = StructArray::new(
327            vec![Field::new("a.b", child.data_type().clone(), true)].into(),
328            vec![child],
329            Some(NullBuffer::from(vec![true, false, false, true, true])),
330        );
331        let root: ArrayRef = Arc::new(StructArray::new(
332            vec![Field::new("nested", nested.data_type().clone(), true)].into(),
333            vec![Arc::new(nested)],
334            Some(NullBuffer::from(vec![true, true, false, true, true])),
335        ));
336        let path = vec!["nested".into(), "a.b".into()];
337        let leaf = extract_json_type_hint_values(&root, &path).unwrap();
338        assert_eq!(
339            (0..leaf.len()).map(|i| leaf.get(i)).collect::<Vec<_>>(),
340            vec![
341                Value::Int64(11),
342                Value::Null,
343                Value::Null,
344                Value::Null,
345                Value::Int64(55)
346            ]
347        );
348        let slice = extract_json_type_hint_values(&root.slice(1, 4), &path).unwrap();
349        assert_eq!(
350            (0..slice.len()).map(|i| slice.get(i)).collect::<Vec<_>>(),
351            vec![Value::Null, Value::Null, Value::Null, Value::Int64(55)]
352        );
353        assert!(extract_json_type_hint_values(&root, &["missing".into()]).is_err());
354        assert!(
355            extract_json_type_hint_values(&root, &["nested".into(), "a.b".into(), "child".into()])
356                .is_err()
357        );
358    }
359}
mito2/sst/index/column.rs

mito2/sst/index/
column.rs
