risingwave_frontend/catalog/
index_catalog.rs

1// Copyright 2022 RisingWave Labs
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::{BTreeMap, HashMap, HashSet};
16use std::ops::Deref;
17use std::sync::Arc;
18
19use educe::Educe;
20use itertools::Itertools;
21use risingwave_common::catalog::{ColumnDesc, CreateType, Field, IndexId, Schema};
22use risingwave_common::session_config::SessionConfig;
23use risingwave_common::util::epoch::Epoch;
24use risingwave_common::util::sort_util::ColumnOrder;
25use risingwave_pb::catalog::{
26    PbCreateType, PbIndex, PbIndexColumnProperties, PbVectorIndexInfo, vector_index_info,
27};
28
29use crate::catalog::table_catalog::TableType;
30use crate::catalog::{OwnedByUserCatalog, TableCatalog};
31use crate::expr::{ExprDisplay, ExprImpl, ExprRewriter as _, FunctionCall};
32use crate::user::UserId;
33
34#[derive(Clone, Debug, Educe)]
35#[educe(PartialEq, Eq, Hash)]
36pub struct TableIndex {
37    pub name: String,
38    /// The properties of the index columns.
39    /// <https://www.postgresql.org/docs/current/functions-info.html#FUNCTIONS-INFO-INDEX-COLUMN-PROPS>
40    pub index_column_properties: Vec<PbIndexColumnProperties>,
41
42    pub index_table: Arc<TableCatalog>,
43
44    pub primary_table: Arc<TableCatalog>,
45
46    pub primary_to_secondary_mapping: BTreeMap<usize, usize>,
47
48    pub secondary_to_primary_mapping: BTreeMap<usize, usize>,
49
50    /// Map function call from the primary table to the index table.
51    /// Use `HashMap` instead of `BTreeMap`, because `FunctionCall` can't be used as the key for
52    /// `BTreeMap`. BTW, the trait `std::hash::Hash` is not implemented for
53    /// `HashMap<function_call::FunctionCall, usize>`, so we need to ignore it. It will not
54    /// affect the correctness, since it can be derived by `index_item`.
55    #[educe(PartialEq(ignore))]
56    #[educe(Hash(ignore))]
57    pub function_mapping: HashMap<FunctionCall, usize>,
58
59    pub index_columns_len: u32,
60}
61
62#[derive(Clone, Debug, PartialEq, Eq, Educe)]
63#[educe(Hash)]
64pub struct VectorIndex {
65    pub index_table: Arc<TableCatalog>,
66    pub vector_expr: ExprImpl,
67    #[educe(Hash(ignore))]
68    pub primary_to_included_info_column_mapping: HashMap<usize, usize>,
69    /// Positions of the primary table's primary key columns in the info columns.
70    ///
71    /// `None` if some primary key column is not stored in the index, e.g. the order key of an
72    /// `ORDER BY` materialized view left out of `INCLUDE`. Such an index can serve only queries it
73    /// fully covers, since it cannot locate rows of the primary table.
74    pub primary_key_idx_in_info_columns: Option<Vec<usize>>,
75    pub included_info_columns: Vec<usize>,
76    pub vector_index_info: PbVectorIndexInfo,
77}
78
79impl VectorIndex {
80    pub fn info_column_desc(&self) -> Vec<ColumnDesc> {
81        self.index_table.columns[1..=self.included_info_columns.len()]
82            .iter()
83            .map(|col| col.column_desc.clone())
84            .collect()
85    }
86
87    pub fn resolve_hnsw_ef_search(&self, config: &SessionConfig) -> Option<usize> {
88        match self.vector_index_info.config.as_ref().unwrap() {
89            vector_index_info::Config::Flat(_) => None,
90            vector_index_info::Config::HnswFlat(_) => Some(config.batch_hnsw_ef_search()),
91        }
92    }
93}
94
95#[derive(Clone, Debug, Hash, PartialEq, Eq)]
96pub enum IndexType {
97    Table(Arc<TableIndex>),
98    Vector(Arc<VectorIndex>),
99}
100
101#[derive(Clone, Debug, Hash, PartialEq, Eq)]
102pub struct IndexCatalog {
103    pub id: IndexId,
104
105    pub name: String,
106
107    /// Only `InputRef` and `FuncCall` type index is supported Now.
108    /// The index of `InputRef` is the column index of the primary table.
109    /// The `index_item` size is equal to the index table columns size
110    /// The input args of `FuncCall` is also the column index of the primary table.
111    pub index_item: Vec<ExprImpl>,
112
113    pub index_type: IndexType,
114
115    pub primary_table: Arc<TableCatalog>,
116
117    /// Indicate whether to create index in background or foreground.
118    pub create_type: CreateType,
119
120    pub created_at_epoch: Option<Epoch>,
121
122    pub initialized_at_epoch: Option<Epoch>,
123
124    pub created_at_cluster_version: Option<String>,
125
126    pub initialized_at_cluster_version: Option<String>,
127}
128
129impl IndexCatalog {
130    pub fn build_from(
131        index_prost: &PbIndex,
132        index_table: &Arc<TableCatalog>,
133        primary_table: &Arc<TableCatalog>,
134    ) -> Self {
135        let index_item: Vec<ExprImpl> = index_prost
136            .index_item
137            .iter()
138            .map(|expr| ExprImpl::from_expr_proto(expr).unwrap())
139            .map(|expr| item_rewriter::CompositeCastEliminator.rewrite_expr(expr))
140            .collect();
141
142        let create_type = index_prost
143            .get_create_type()
144            .unwrap_or(PbCreateType::Foreground);
145
146        let index_type = match index_table.table_type {
147            TableType::Index => {
148                let primary_to_secondary_mapping: BTreeMap<usize, usize> = index_item
149                    .iter()
150                    .enumerate()
151                    .filter_map(|(i, expr)| match expr {
152                        ExprImpl::InputRef(input_ref) => Some((input_ref.index, i)),
153                        ExprImpl::FunctionCall(_) => None,
154                        _ => unreachable!(),
155                    })
156                    .collect();
157
158                let secondary_to_primary_mapping = BTreeMap::from_iter(
159                    primary_to_secondary_mapping
160                        .clone()
161                        .into_iter()
162                        .map(|(x, y)| (y, x)),
163                );
164
165                let function_mapping: HashMap<FunctionCall, usize> = index_item
166                    .iter()
167                    .enumerate()
168                    .filter_map(|(i, expr)| match expr {
169                        ExprImpl::InputRef(_) => None,
170                        ExprImpl::FunctionCall(func) => Some((func.deref().clone(), i)),
171                        _ => unreachable!(),
172                    })
173                    .collect();
174                IndexType::Table(Arc::new(TableIndex {
175                    name: index_prost.name.clone(),
176                    index_column_properties: index_prost.index_column_properties.clone(),
177                    index_columns_len: index_prost.index_columns_len,
178                    index_table: index_table.clone(),
179                    primary_table: primary_table.clone(),
180                    primary_to_secondary_mapping,
181                    secondary_to_primary_mapping,
182                    function_mapping,
183                }))
184            }
185            TableType::VectorIndex => {
186                assert_eq!(index_prost.index_columns_len, 1);
187                let included_info_columns = index_item[1..].iter().map(|item| {
188                    let ExprImpl::InputRef(input) = item else {
189                        panic!("vector index included columns must be from direct input column, but got: {:?}", item);
190                    };
191                    input.index
192                }).collect_vec();
193                let primary_to_included_info_column_mapping: HashMap<_, _> = included_info_columns
194                    .iter()
195                    .enumerate()
196                    .map(|(included_info_column_idx, primary_column_idx)| {
197                        (*primary_column_idx, included_info_column_idx)
198                    })
199                    .collect();
200                let primary_key_idx_in_info_columns = primary_table
201                    .pk()
202                    .iter()
203                    .map(|order| {
204                        primary_to_included_info_column_mapping
205                            .get(&order.column_index)
206                            .copied()
207                    })
208                    .collect();
209                IndexType::Vector(Arc::new(VectorIndex {
210                    index_table: index_table.clone(),
211                    vector_expr: index_item[0].clone(),
212                    primary_to_included_info_column_mapping,
213                    primary_key_idx_in_info_columns,
214                    included_info_columns,
215                    vector_index_info: index_table
216                        .vector_index_info
217                        .expect("should exist for vector index"),
218                }))
219            }
220            TableType::Table | TableType::MaterializedView | TableType::Internal => {
221                unreachable!()
222            }
223        };
224
225        IndexCatalog {
226            id: index_prost.id,
227            name: index_prost.name.clone(),
228            index_item,
229            index_type,
230            primary_table: primary_table.clone(),
231            created_at_epoch: index_prost.created_at_epoch.map(Epoch::from),
232            initialized_at_epoch: index_prost.initialized_at_epoch.map(Epoch::from),
233            created_at_cluster_version: index_prost.created_at_cluster_version.clone(),
234            initialized_at_cluster_version: index_prost.initialized_at_cluster_version.clone(),
235            create_type: CreateType::from_proto(create_type),
236        }
237    }
238}
239
240impl TableIndex {
241    /// Map the primary key of the primary table to the columns of the index table.
242    ///
243    /// Returns `None` if some primary key column is not stored in the index table. The index
244    /// table always contains the *stream key* of the primary table, but the primary key (storage
245    /// order key) can be a superset of it, e.g. for a `TopN` materialized view whose primary key
246    /// is `[order_key..., stream_key...]`.
247    pub fn primary_table_pk_ref_to_index_table(&self) -> Option<Vec<ColumnOrder>> {
248        let mapping = self.primary_to_secondary_mapping();
249
250        self.primary_table
251            .pk
252            .iter()
253            .map(|x| {
254                mapping
255                    .get(&x.column_index)
256                    .map(|idx| ColumnOrder::new(*idx, x.order_type))
257            })
258            .collect()
259    }
260
261    pub fn primary_table_distribute_key_ref_to_index_table(&self) -> Vec<usize> {
262        let mapping = self.primary_to_secondary_mapping();
263
264        self.primary_table
265            .distribution_key
266            .iter()
267            .map(|x| *mapping.get(x).unwrap())
268            .collect_vec()
269    }
270
271    pub fn full_covering(&self) -> bool {
272        self.index_table.columns.len() == self.primary_table.columns.len()
273    }
274
275    /// A mapping maps the column index of the secondary index to the column index of the primary
276    /// table.
277    pub fn secondary_to_primary_mapping(&self) -> &BTreeMap<usize, usize> {
278        &self.secondary_to_primary_mapping
279    }
280
281    /// A mapping maps the column index of the primary table to the column index of the secondary
282    /// index.
283    pub fn primary_to_secondary_mapping(&self) -> &BTreeMap<usize, usize> {
284        &self.primary_to_secondary_mapping
285    }
286
287    pub fn function_mapping(&self) -> &HashMap<FunctionCall, usize> {
288        &self.function_mapping
289    }
290}
291
292impl IndexCatalog {
293    pub fn index_table(&self) -> &Arc<TableCatalog> {
294        match &self.index_type {
295            IndexType::Table(index) => &index.index_table,
296            IndexType::Vector(index) => &index.index_table,
297        }
298    }
299
300    /// Get the column properties of the index column.
301    pub fn get_column_properties(&self, column_idx: usize) -> Option<PbIndexColumnProperties> {
302        match &self.index_type {
303            IndexType::Table(index) => index.index_column_properties.get(column_idx).cloned(),
304            IndexType::Vector { .. } => {
305                if column_idx == 0 {
306                    // return with the default value defined in [https://www.postgresql.org/docs/current/sql-createindex.html]
307                    Some(PbIndexColumnProperties {
308                        is_desc: false,
309                        nulls_first: false,
310                    })
311                } else {
312                    None
313                }
314            }
315        }
316    }
317
318    pub fn get_column_def(&self, column_idx: usize) -> Option<String> {
319        let col = self.index_table().columns.get(column_idx)?;
320        if col.is_hidden {
321            return None;
322        }
323        let expr_display = ExprDisplay {
324            expr: &self.index_item[column_idx],
325            input_schema: &Schema::new(
326                self.primary_table
327                    .columns
328                    .iter()
329                    .map(|col| Field::from(&col.column_desc))
330                    .collect_vec(),
331            ),
332        };
333
334        // TODO(Kexiang): Currently, extra info like ":Int32" introduced by `ExprDisplay` is kept for simplity.
335        // We'd better remove it in the future.
336        Some(expr_display.to_string())
337    }
338
339    pub fn display(&self) -> IndexDisplay {
340        let index_table = self.index_table().clone();
341        let index_columns_with_ordering = index_table
342            .pk
343            .iter()
344            .filter(|x| !index_table.columns[x.column_index].is_hidden)
345            .map(|x| {
346                let index_column_name = index_table.columns[x.column_index].name().to_owned();
347                format!("{} {}", index_column_name, x.order_type)
348            })
349            .collect_vec();
350
351        let pk_column_index_set = index_table
352            .pk
353            .iter()
354            .map(|x| x.column_index)
355            .collect::<HashSet<_>>();
356
357        let include_columns = index_table
358            .columns
359            .iter()
360            .enumerate()
361            .filter(|(i, _)| !pk_column_index_set.contains(i))
362            .filter(|(_, x)| !x.is_hidden)
363            .map(|(_, x)| x.name().to_owned())
364            .collect_vec();
365
366        let distributed_by_columns = index_table
367            .distribution_key
368            .iter()
369            .map(|&x| index_table.columns[x].name().to_owned())
370            .collect_vec();
371
372        IndexDisplay {
373            index_columns_with_ordering,
374            include_columns,
375            distributed_by_columns,
376        }
377    }
378
379    pub fn is_created(&self) -> bool {
380        self.index_table().is_created()
381    }
382}
383
384pub struct IndexDisplay {
385    pub index_columns_with_ordering: Vec<String>,
386    pub include_columns: Vec<String>,
387    pub distributed_by_columns: Vec<String>,
388}
389
390impl OwnedByUserCatalog for IndexCatalog {
391    fn owner(&self) -> UserId {
392        self.index_table().owner
393    }
394}
395
396mod item_rewriter {
397    use risingwave_pb::expr::expr_node;
398
399    use crate::expr::{Expr, ExprImpl, ExprRewriter, FunctionCall};
400
401    /// Rewrite the expression of index item to eliminate `CompositeCast`, if any. This is needed
402    /// if the type of a column was changed and there's functional index on it.
403    ///
404    /// # Example
405    ///
406    /// Imagine there's a table created with `CREATE TABLE t (v struct<a int, b int>)`.
407    /// Then we create an index on it with `CREATE INDEX idx ON t ((v).a)`, which will create an
408    /// index item `Field(InputRef(0), 0)`.
409    ///
410    /// If we alter the column with `ALTER TABLE t ALTER COLUMN v TYPE struct<x varchar, a int>`,
411    /// the meta service will wrap the `InputRef(0)` with a `CompositeCast` to maintain the correct
412    /// return type. The index item will now become `Field(CompositeCast(InputRef(0)), 0)`.
413    ///
414    /// `CompositeCast` is for internal use only, and cannot be constructed or executed. To allow
415    /// this functional index to work and be matched with user queries, we need to eliminate it
416    /// here. By comparing the input and output types of `CompositeCast` and matching the field id,
417    /// we can find the real `Field` index and rewrite it to `Field(InputRef(0), 1)`.
418    ///
419    /// Note that if the field is dropped, we will leave the index item as is. This makes the index
420    /// item invalid, and it will never be matched and used.
421    pub struct CompositeCastEliminator;
422
423    impl ExprRewriter for CompositeCastEliminator {
424        fn rewrite_function_call(&mut self, func_call: FunctionCall) -> ExprImpl {
425            let (func_type, inputs, ret) = func_call.decompose();
426
427            // Flatten consecutive `CompositeCast`.
428            if func_type == expr_node::Type::CompositeCast {
429                let child = inputs[0].clone();
430
431                if let Some(child) = child.as_function_call()
432                    && child.func_type() == expr_node::Type::CompositeCast
433                {
434                    let new_child = child.inputs()[0].clone();
435
436                    // If the type already matches, no need to wrap again.
437                    // Recursively eliminate more composite cast by calling rewrite again.
438                    if new_child.return_type() == ret {
439                        return self.rewrite_expr(new_child);
440                    } else {
441                        let new_composite_cast =
442                            FunctionCall::new_unchecked(func_type, vec![new_child], ret);
443                        return self.rewrite_function_call(new_composite_cast);
444                    }
445                }
446            }
447            // Rewrite `Field(CompositeCast(x), y)` to `Field(x, y')`.
448            // TODO: also support rewriting `ArrayAccess` and `MapAccess`.
449            else if func_type == expr_node::Type::Field {
450                let child = inputs[0].clone();
451
452                if let Some(child) = child.as_function_call()
453                    && child.func_type() == expr_node::Type::CompositeCast
454                {
455                    let index = (inputs[1].clone().into_literal().unwrap())
456                        .get_data()
457                        .clone()
458                        .unwrap()
459                        .into_int32();
460
461                    let struct_type = child.return_type().into_struct();
462                    let field_id = struct_type
463                        .id_at(index as usize)
464                        .expect("ids should be set");
465
466                    // Unwrap the composite cast.
467                    let new_child = child.inputs()[0].clone();
468                    let new_struct_type = new_child.return_type().into_struct();
469
470                    let Some(new_index) = new_struct_type
471                        .ids()
472                        .expect("ids should be set")
473                        .position(|x| x == field_id)
474                    else {
475                        // Previously we have index on this field, but now it's dropped.
476                        // As a result, this entire index item becomes invalid.
477                        // Simply leave it as is. Users cannot construct a `CompositeCast` (which is
478                        // not user-facing), thus this index item will never be matched and used.
479                        return FunctionCall::new_unchecked(func_type, inputs, ret).into();
480                    };
481                    let new_index = ExprImpl::literal_int(new_index as i32);
482
483                    let new_inputs = vec![new_child, new_index];
484                    let new_field_call = FunctionCall::new_unchecked(func_type, new_inputs, ret);
485
486                    // Recursively eliminate more composite cast by calling rewrite again.
487                    return self.rewrite_function_call(new_field_call);
488                }
489            }
490
491            let inputs = inputs
492                .into_iter()
493                .map(|expr| self.rewrite_expr(expr))
494                .collect();
495            FunctionCall::new_unchecked(func_type, inputs, ret).into()
496        }
497    }
498}
risingwave_frontend/catalog/index_catalog.rs

risingwave_frontend/catalog/
index_catalog.rs
