flow/repr/
relation.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 datafusion_common::DFSchema;
16use datatypes::data_type::DataType;
17use datatypes::prelude::ConcreteDataType;
18use serde::{Deserialize, Serialize};
19use snafu::{ResultExt, ensure};
20
21use crate::error::{DatafusionSnafu, InternalSnafu, InvalidQuerySnafu, Result};
22
23/// a set of column indices that are "keys" for the collection.
24#[derive(Default, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize, Hash)]
25pub struct Key {
26    /// indicate whose column form key
27    pub column_indices: Vec<usize>,
28}
29
30impl Key {
31    /// create a new Key
32    pub fn new() -> Self {
33        Default::default()
34    }
35
36    /// create a new Key from a vector of column indices
37    pub fn from(mut column_indices: Vec<usize>) -> Self {
38        column_indices.sort_unstable();
39        Self { column_indices }
40    }
41
42    /// Add a column to Key
43    pub fn add_col(&mut self, col: usize) {
44        self.column_indices.push(col);
45    }
46
47    /// Remove a column from Key
48    pub fn remove_col(&mut self, col: usize) {
49        self.column_indices.retain(|&r| r != col);
50    }
51
52    /// get all columns in Key
53    pub fn get(&self) -> &Vec<usize> {
54        &self.column_indices
55    }
56
57    /// True if Key is empty
58    pub fn is_empty(&self) -> bool {
59        self.column_indices.is_empty()
60    }
61
62    /// True if all columns in self are also in other
63    pub fn subset_of(&self, other: &Key) -> bool {
64        self.column_indices
65            .iter()
66            .all(|c| other.column_indices.contains(c))
67    }
68}
69
70/// The type of a relation.
71#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize, Hash)]
72pub struct RelationType {
73    /// The type for each column, in order.
74    pub column_types: Vec<ColumnType>,
75    /// Sets of indices that are "keys" for the collection.
76    ///
77    /// Each element in this list is a set of column indices, each with the
78    /// property that the collection contains at most one record with each
79    /// distinct set of values for each column. Alternately, for a specific set
80    /// of values assigned to the these columns there is at most one record.
81    ///
82    /// A collection can contain multiple sets of keys, although it is common to
83    /// have either zero or one sets of key indices.
84    pub keys: Vec<Key>,
85    /// optionally indicate the column that is TIME INDEX
86    pub time_index: Option<usize>,
87    /// mark all the columns that are added automatically by flow, but are not present in original sql
88    pub auto_columns: Vec<usize>,
89}
90
91impl RelationType {
92    pub fn with_autos(mut self, auto_cols: &[usize]) -> Self {
93        self.auto_columns = auto_cols.to_vec();
94        self
95    }
96
97    /// Constructs a `RelationType` representing the relation with no columns and
98    /// no keys.
99    pub fn empty() -> Self {
100        RelationType::new(vec![])
101    }
102
103    /// Constructs a new `RelationType` from specified column types.
104    ///
105    /// The `RelationType` will have no keys.
106    pub fn new(column_types: Vec<ColumnType>) -> Self {
107        RelationType {
108            column_types,
109            keys: Vec::new(),
110            time_index: None,
111            auto_columns: vec![],
112        }
113    }
114
115    /// Adds a new key for the relation. Also sorts the key indices.
116    ///
117    /// will ignore empty key
118    pub fn with_key(mut self, mut indices: Vec<usize>) -> Self {
119        if indices.is_empty() {
120            return self;
121        }
122        indices.sort_unstable();
123        let key = Key::from(indices);
124        if !self.keys.contains(&key) {
125            self.keys.push(key);
126        }
127        self
128    }
129
130    /// Adds new keys for the relation. Also sorts the key indices.
131    ///
132    /// will ignore empty keys
133    pub fn with_keys(mut self, keys: Vec<Vec<usize>>) -> Self {
134        for key in keys {
135            self = self.with_key(key)
136        }
137        self
138    }
139
140    /// will also remove time index from keys if it's in keys
141    pub fn with_time_index(mut self, time_index: Option<usize>) -> Self {
142        self.time_index = time_index;
143        for key in &mut self.keys {
144            key.remove_col(time_index.unwrap_or(usize::MAX));
145        }
146        // remove empty keys
147        self.keys.retain(|key| !key.is_empty());
148        self
149    }
150
151    /// Computes the number of columns in the relation.
152    pub fn arity(&self) -> usize {
153        self.column_types.len()
154    }
155
156    /// Gets the index of the columns used when creating a default index.
157    pub fn default_key(&self) -> Vec<usize> {
158        if let Some(key) = self.keys.first() {
159            if key.is_empty() {
160                (0..self.column_types.len()).collect()
161            } else {
162                key.get().clone()
163            }
164        } else {
165            (0..self.column_types.len()).collect()
166        }
167    }
168
169    /// True if any collection described by `self` could safely be described by `other`.
170    ///
171    /// In practice this means checking that the scalar types match exactly, and that the
172    /// nullability of `self` is at least as strict as `other`, and that all keys of `other`
173    /// contain some key of `self` (as a set of key columns is less strict than any subset).
174    pub fn subtypes(&self, other: &RelationType) -> bool {
175        if self.column_types.len() != other.column_types.len() {
176            return false;
177        }
178
179        for (col1, col2) in self.column_types.iter().zip(other.column_types.iter()) {
180            if col1.nullable && !col2.nullable {
181                return false;
182            }
183            if col1.scalar_type != col2.scalar_type {
184                return false;
185            }
186        }
187
188        let all_keys = other
189            .keys
190            .iter()
191            .all(|key1| self.keys.iter().any(|key2| key1.subset_of(key2)));
192        if !all_keys {
193            return false;
194        }
195
196        true
197    }
198
199    /// Return relation describe with column names
200    pub fn into_named(self, names: Vec<Option<ColumnName>>) -> RelationDesc {
201        RelationDesc { typ: self, names }
202    }
203
204    /// Return relation describe without column names
205    pub fn into_unnamed(self) -> RelationDesc {
206        RelationDesc {
207            names: vec![None; self.column_types.len()],
208            typ: self,
209        }
210    }
211}
212
213/// The type of a `Value`
214///
215/// [`ColumnType`] bundles information about the scalar type of a datum (e.g.,
216/// Int32 or String) with its nullability.
217///
218/// To construct a column type, either initialize the struct directly, or
219/// use the [`ScalarType::nullable`] method.
220#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize, Hash)]
221pub struct ColumnType {
222    /// The underlying scalar type (e.g., Int32 or String) of this column.
223    pub scalar_type: ConcreteDataType,
224    /// Whether this datum can be null.
225    #[serde(default = "return_true")]
226    pub nullable: bool,
227}
228
229impl ColumnType {
230    /// Constructs a new `ColumnType` from a scalar type and a nullability flag.
231    pub fn new(scalar_type: ConcreteDataType, nullable: bool) -> Self {
232        ColumnType {
233            scalar_type,
234            nullable,
235        }
236    }
237
238    /// Constructs a new `ColumnType` from a scalar type, with nullability set to
239    /// ***true***
240    pub fn new_nullable(scalar_type: ConcreteDataType) -> Self {
241        ColumnType {
242            scalar_type,
243            nullable: true,
244        }
245    }
246
247    /// Returns the scalar type of this column.
248    pub fn scalar_type(&self) -> &ConcreteDataType {
249        &self.scalar_type
250    }
251
252    /// Returns true if this column can be null.
253    pub fn nullable(&self) -> bool {
254        self.nullable
255    }
256}
257
258/// This method exists solely for the purpose of making ColumnType nullable by
259/// default in unit tests. The default value of a bool is false, and the only
260/// way to make an object take on any other value by default is to pass it a
261/// function that returns the desired default value. See
262/// <https://github.com/serde-rs/serde/issues/1030>
263#[inline(always)]
264fn return_true() -> bool {
265    true
266}
267
268/// A description of the shape of a relation.
269///
270/// It bundles a [`RelationType`] with the name of each column in the relation.
271/// Individual column names are optional.
272#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize, Hash)]
273pub struct RelationDesc {
274    pub typ: RelationType,
275    pub names: Vec<Option<ColumnName>>,
276}
277
278impl RelationDesc {
279    pub fn len(&self) -> Result<usize> {
280        ensure!(
281            self.typ.column_types.len() == self.names.len(),
282            InternalSnafu {
283                reason: "Expect typ and names field to be of same length"
284            }
285        );
286        Ok(self.names.len())
287    }
288
289    pub fn to_df_schema(&self) -> Result<DFSchema> {
290        let fields: Vec<_> = self
291            .iter()
292            .enumerate()
293            .map(|(i, (name, typ))| {
294                let name = name.clone().unwrap_or(format!("Col_{i}"));
295                let nullable = typ.nullable;
296                let data_type = typ.scalar_type.clone().as_arrow_type();
297                arrow_schema::Field::new(name, data_type, nullable)
298            })
299            .collect();
300        let arrow_schema = arrow_schema::Schema::new(fields);
301
302        DFSchema::try_from(arrow_schema.clone()).with_context(|_e| DatafusionSnafu {
303            context: format!("Error when converting to DFSchema: {:?}", arrow_schema),
304        })
305    }
306}
307
308impl RelationDesc {
309    /// Constructs a new `RelationDesc` that represents the empty relation
310    /// with no columns and no keys.
311    pub fn empty() -> Self {
312        RelationDesc {
313            typ: RelationType::empty(),
314            names: vec![],
315        }
316    }
317
318    /// Constructs a new `RelationDesc` from a `RelationType` and an iterator
319    /// over column names.
320    ///
321    pub fn try_new<I, N>(typ: RelationType, names: I) -> Result<Self>
322    where
323        I: IntoIterator<Item = N>,
324        N: Into<Option<ColumnName>>,
325    {
326        let names: Vec<_> = names.into_iter().map(|name| name.into()).collect();
327        ensure!(
328            typ.arity() == names.len(),
329            InvalidQuerySnafu {
330                reason: format!(
331                    "Length mismatch between RelationType {:?} and column names {:?}",
332                    typ.column_types, names
333                )
334            }
335        );
336        Ok(RelationDesc { typ, names })
337    }
338
339    /// Constructs a new `RelationDesc` from a `RelationType` and an iterator
340    /// over column names.
341    ///
342    /// # Panics
343    ///
344    /// Panics if the arity of the `RelationType` is not equal to the number of
345    /// items in `names`.
346    pub fn new_unchecked<I, N>(typ: RelationType, names: I) -> Self
347    where
348        I: IntoIterator<Item = N>,
349        N: Into<Option<ColumnName>>,
350    {
351        let names: Vec<_> = names.into_iter().map(|name| name.into()).collect();
352        assert_eq!(typ.arity(), names.len());
353        RelationDesc { typ, names }
354    }
355
356    pub fn from_names_and_types<I, T, N>(iter: I) -> Self
357    where
358        I: IntoIterator<Item = (N, T)>,
359        T: Into<ColumnType>,
360        N: Into<Option<ColumnName>>,
361    {
362        let (names, types): (Vec<_>, Vec<_>) = iter.into_iter().unzip();
363        let types = types.into_iter().map(Into::into).collect();
364        let typ = RelationType::new(types);
365        Self::new_unchecked(typ, names)
366    }
367    /// Concatenates a `RelationDesc` onto the end of this `RelationDesc`.
368    pub fn concat(mut self, other: Self) -> Self {
369        let self_len = self.typ.column_types.len();
370        self.names.extend(other.names);
371        self.typ.column_types.extend(other.typ.column_types);
372        for k in other.typ.keys {
373            let k = k
374                .column_indices
375                .into_iter()
376                .map(|idx| idx + self_len)
377                .collect();
378            self = self.with_key(k);
379        }
380        self
381    }
382
383    /// Appends a column with the specified name and type.
384    pub fn with_column<N>(mut self, name: N, column_type: ColumnType) -> Self
385    where
386        N: Into<Option<ColumnName>>,
387    {
388        self.typ.column_types.push(column_type);
389        self.names.push(name.into());
390        self
391    }
392
393    /// Adds a new key for the relation.
394    pub fn with_key(mut self, indices: Vec<usize>) -> Self {
395        self.typ = self.typ.with_key(indices);
396        self
397    }
398
399    /// Builds a new relation description with the column names replaced with
400    /// new names.
401    ///
402    pub fn try_with_names<I, N>(self, names: I) -> Result<Self>
403    where
404        I: IntoIterator<Item = N>,
405        N: Into<Option<ColumnName>>,
406    {
407        Self::try_new(self.typ, names)
408    }
409
410    /// Computes the number of columns in the relation.
411    pub fn arity(&self) -> usize {
412        self.typ.arity()
413    }
414
415    /// Returns the relation type underlying this relation description.
416    pub fn typ(&self) -> &RelationType {
417        &self.typ
418    }
419
420    /// Returns an iterator over the columns in this relation.
421    pub fn iter(&self) -> impl Iterator<Item = (&Option<ColumnName>, &ColumnType)> {
422        self.iter_names().zip(self.iter_types())
423    }
424
425    /// Returns an iterator over the types of the columns in this relation.
426    pub fn iter_types(&self) -> impl Iterator<Item = &ColumnType> {
427        self.typ.column_types.iter()
428    }
429
430    /// Returns an iterator over the names of the columns in this relation.
431    pub fn iter_names(&self) -> impl Iterator<Item = &Option<ColumnName>> {
432        self.names.iter()
433    }
434
435    /// Finds a column by name.
436    ///
437    /// Returns the index and type of the column named `name`. If no column with
438    /// the specified name exists, returns `None`. If multiple columns have the
439    /// specified name, the leftmost column is returned.
440    pub fn get_by_name(&self, name: &ColumnName) -> Option<(usize, &ColumnType)> {
441        self.iter_names()
442            .position(|n| n.as_ref() == Some(name))
443            .map(|i| (i, &self.typ.column_types[i]))
444    }
445
446    /// Gets the name of the `i`th column.
447    ///
448    /// # Panics
449    ///
450    /// Panics if `i` is not a valid column index.
451    pub fn get_name(&self, i: usize) -> &Option<ColumnName> {
452        &self.names[i]
453    }
454}
455
456/// The name of a column in a [`RelationDesc`].
457pub type ColumnName = String;
flow/repr/relation.rs

flow/repr/
relation.rs
