arrow_flight/sql/metadata/
tables.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! [`GetTablesBuilder`] for building responses to [`CommandGetTables`] queries.
6//!
7//! [`CommandGetTables`]: crate::sql::CommandGetTables
8
9use std::sync::{Arc, LazyLock};
10
11use arrow_arith::boolean::{and, or};
12use arrow_array::builder::{BinaryBuilder, StringBuilder};
13use arrow_array::{ArrayRef, RecordBatch, StringArray};
14use arrow_ord::cmp::eq;
15use arrow_schema::{DataType, Field, Schema, SchemaRef};
16use arrow_select::{filter::filter_record_batch, take::take};
17use arrow_string::like::like;
18
19use super::lexsort_to_indices;
20use crate::error::*;
21use crate::sql::CommandGetTables;
22use crate::{IpcMessage, IpcWriteOptions, SchemaAsIpc};
23
24/// A builder for a [`CommandGetTables`] response.
25///
26/// Builds rows like this:
27///
28/// * catalog_name: utf8,
29/// * db_schema_name: utf8,
30/// * table_name: utf8 not null,
31/// * table_type: utf8 not null,
32/// * (optional) table_schema: bytes not null (schema of the table as described
33///   in Schema.fbs::Schema it is serialized as an IPC message.)
34pub struct GetTablesBuilder {
35    catalog_filter: Option<String>,
36    table_types_filter: Vec<String>,
37    // Optional filters to apply to schemas
38    db_schema_filter_pattern: Option<String>,
39    // Optional filters to apply to tables
40    table_name_filter_pattern: Option<String>,
41    // array builder for catalog names
42    catalog_name: StringBuilder,
43    // array builder for db schema names
44    db_schema_name: StringBuilder,
45    // array builder for tables names
46    table_name: StringBuilder,
47    // array builder for table types
48    table_type: StringBuilder,
49    // array builder for table schemas
50    table_schema: Option<BinaryBuilder>,
51}
52
53impl CommandGetTables {
54    /// Create a builder suitable for constructing a response
55    pub fn into_builder(self) -> GetTablesBuilder {
56        self.into()
57    }
58}
59
60impl From<CommandGetTables> for GetTablesBuilder {
61    fn from(value: CommandGetTables) -> Self {
62        Self::new(
63            value.catalog,
64            value.db_schema_filter_pattern,
65            value.table_name_filter_pattern,
66            value.table_types,
67            value.include_schema,
68        )
69    }
70}
71
72impl GetTablesBuilder {
73    /// Create a new instance of [`GetTablesBuilder`]
74    ///
75    /// # Parameters
76    ///
77    /// - `catalog`:  Specifies the Catalog to search for the tables.
78    ///   - An empty string retrieves those without a catalog.
79    ///   - If omitted the catalog name is not used to narrow the search.
80    /// - `db_schema_filter_pattern`: Specifies a filter pattern for schemas to search for.
81    ///   When no pattern is provided, the pattern will not be used to narrow the search.
82    ///   In the pattern string, two special characters can be used to denote matching rules:
83    ///     - "%" means to match any substring with 0 or more characters.
84    ///     - "_" means to match any one character.
85    /// - `table_name_filter_pattern`: Specifies a filter pattern for tables to search for.
86    ///   When no pattern is provided, all tables matching other filters are searched.
87    ///   In the pattern string, two special characters can be used to denote matching rules:
88    ///     - "%" means to match any substring with 0 or more characters.
89    ///     - "_" means to match any one character.
90    /// - `table_types`:  Specifies a filter of table types which must match.
91    ///   An empty Vec matches all table types.
92    /// - `include_schema`: Specifies if the Arrow schema should be returned for found tables.
93    ///
94    /// [`CommandGetTables`]: crate::sql::CommandGetTables
95    pub fn new(
96        catalog: Option<impl Into<String>>,
97        db_schema_filter_pattern: Option<impl Into<String>>,
98        table_name_filter_pattern: Option<impl Into<String>>,
99        table_types: impl IntoIterator<Item = impl Into<String>>,
100        include_schema: bool,
101    ) -> Self {
102        let table_schema = if include_schema {
103            Some(BinaryBuilder::new())
104        } else {
105            None
106        };
107        Self {
108            catalog_filter: catalog.map(|s| s.into()),
109            table_types_filter: table_types.into_iter().map(|tt| tt.into()).collect(),
110            db_schema_filter_pattern: db_schema_filter_pattern.map(|s| s.into()),
111            table_name_filter_pattern: table_name_filter_pattern.map(|t| t.into()),
112            catalog_name: StringBuilder::new(),
113            db_schema_name: StringBuilder::new(),
114            table_name: StringBuilder::new(),
115            table_type: StringBuilder::new(),
116            table_schema,
117        }
118    }
119
120    /// Append a row
121    pub fn append(
122        &mut self,
123        catalog_name: impl AsRef<str>,
124        schema_name: impl AsRef<str>,
125        table_name: impl AsRef<str>,
126        table_type: impl AsRef<str>,
127        table_schema: &Schema,
128    ) -> Result<()> {
129        self.catalog_name.append_value(catalog_name);
130        self.db_schema_name.append_value(schema_name);
131        self.table_name.append_value(table_name);
132        self.table_type.append_value(table_type);
133        if let Some(self_table_schema) = self.table_schema.as_mut() {
134            let options = IpcWriteOptions::default();
135            // encode the schema into the correct form
136            let message: std::result::Result<IpcMessage, _> =
137                SchemaAsIpc::new(table_schema, &options).try_into();
138            let IpcMessage(schema) = message?;
139            self_table_schema.append_value(schema);
140        }
141
142        Ok(())
143    }
144
145    /// builds a `RecordBatch` for `CommandGetTables`
146    pub fn build(self) -> Result<RecordBatch> {
147        let schema = self.schema();
148        let Self {
149            catalog_filter,
150            table_types_filter,
151            db_schema_filter_pattern,
152            table_name_filter_pattern,
153
154            mut catalog_name,
155            mut db_schema_name,
156            mut table_name,
157            mut table_type,
158            table_schema,
159        } = self;
160
161        // Make the arrays
162        let catalog_name = catalog_name.finish();
163        let db_schema_name = db_schema_name.finish();
164        let table_name = table_name.finish();
165        let table_type = table_type.finish();
166        let table_schema = table_schema.map(|mut table_schema| table_schema.finish());
167
168        // apply any filters, getting a BooleanArray that represents
169        // the rows that passed the filter
170        let mut filters = vec![];
171
172        if let Some(catalog_filter_name) = catalog_filter {
173            let scalar = StringArray::new_scalar(catalog_filter_name);
174            filters.push(eq(&catalog_name, &scalar)?);
175        }
176
177        let tt_filter = table_types_filter
178            .into_iter()
179            .map(|tt| eq(&table_type, &StringArray::new_scalar(tt)))
180            .collect::<std::result::Result<Vec<_>, _>>()?
181            .into_iter()
182            // We know the arrays are of same length as they are produced fromn the same root array
183            .reduce(|filter, arr| or(&filter, &arr).unwrap());
184        if let Some(filter) = tt_filter {
185            filters.push(filter);
186        }
187
188        if let Some(db_schema_filter_pattern) = db_schema_filter_pattern {
189            // use like kernel to get wildcard matching
190            let scalar = StringArray::new_scalar(db_schema_filter_pattern);
191            filters.push(like(&db_schema_name, &scalar)?)
192        }
193
194        if let Some(table_name_filter_pattern) = table_name_filter_pattern {
195            // use like kernel to get wildcard matching
196            let scalar = StringArray::new_scalar(table_name_filter_pattern);
197            filters.push(like(&table_name, &scalar)?)
198        }
199
200        let batch = if let Some(table_schema) = table_schema {
201            RecordBatch::try_new(
202                schema,
203                vec![
204                    Arc::new(catalog_name) as ArrayRef,
205                    Arc::new(db_schema_name) as ArrayRef,
206                    Arc::new(table_name) as ArrayRef,
207                    Arc::new(table_type) as ArrayRef,
208                    Arc::new(table_schema) as ArrayRef,
209                ],
210            )
211        } else {
212            RecordBatch::try_new(
213                // schema is different if table_schema is none
214                schema,
215                vec![
216                    Arc::new(catalog_name) as ArrayRef,
217                    Arc::new(db_schema_name) as ArrayRef,
218                    Arc::new(table_name) as ArrayRef,
219                    Arc::new(table_type) as ArrayRef,
220                ],
221            )
222        }?;
223
224        // `AND` any filters together
225        let mut total_filter = None;
226        while let Some(filter) = filters.pop() {
227            let new_filter = match total_filter {
228                Some(total_filter) => and(&total_filter, &filter)?,
229                None => filter,
230            };
231            total_filter = Some(new_filter);
232        }
233
234        // Apply the filters if needed
235        let filtered_batch = if let Some(total_filter) = total_filter {
236            filter_record_batch(&batch, &total_filter)?
237        } else {
238            batch
239        };
240
241        // Order filtered results by catalog_name, then db_schema_name, then table_name, then table_type
242        // https://github.com/apache/arrow/blob/130f9e981aa98c25de5f5bfe55185db270cec313/format/FlightSql.proto#LL1202C1-L1202C1
243        let sort_cols = filtered_batch.project(&[0, 1, 2, 3])?;
244        let indices = lexsort_to_indices(sort_cols.columns());
245        let columns = filtered_batch
246            .columns()
247            .iter()
248            .map(|c| take(c, &indices, None))
249            .collect::<std::result::Result<Vec<_>, _>>()?;
250
251        Ok(RecordBatch::try_new(filtered_batch.schema(), columns)?)
252    }
253
254    /// Return the schema of the RecordBatch that will be returned from [`CommandGetTables`]
255    ///
256    /// Note the schema differs based on the values of `include_schema`
257    ///
258    /// [`CommandGetTables`]: crate::sql::CommandGetTables
259    pub fn schema(&self) -> SchemaRef {
260        get_tables_schema(self.include_schema())
261    }
262
263    /// Should the "schema" column be included
264    pub fn include_schema(&self) -> bool {
265        self.table_schema.is_some()
266    }
267}
268
269fn get_tables_schema(include_schema: bool) -> SchemaRef {
270    if include_schema {
271        Arc::clone(&GET_TABLES_SCHEMA_WITH_TABLE_SCHEMA)
272    } else {
273        Arc::clone(&GET_TABLES_SCHEMA_WITHOUT_TABLE_SCHEMA)
274    }
275}
276
277/// The schema for GetTables without `table_schema` column
278static GET_TABLES_SCHEMA_WITHOUT_TABLE_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
279    Arc::new(Schema::new(vec![
280        Field::new("catalog_name", DataType::Utf8, true),
281        Field::new("db_schema_name", DataType::Utf8, true),
282        Field::new("table_name", DataType::Utf8, false),
283        Field::new("table_type", DataType::Utf8, false),
284    ]))
285});
286
287/// The schema for GetTables with `table_schema` column
288static GET_TABLES_SCHEMA_WITH_TABLE_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
289    Arc::new(Schema::new(vec![
290        Field::new("catalog_name", DataType::Utf8, true),
291        Field::new("db_schema_name", DataType::Utf8, true),
292        Field::new("table_name", DataType::Utf8, false),
293        Field::new("table_type", DataType::Utf8, false),
294        Field::new("table_schema", DataType::Binary, false),
295    ]))
296});
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301    use arrow_array::{StringArray, UInt32Array};
302
303    fn get_ref_batch() -> RecordBatch {
304        RecordBatch::try_new(
305            get_tables_schema(false),
306            vec![
307                Arc::new(StringArray::from(vec![
308                    "a_catalog",
309                    "a_catalog",
310                    "a_catalog",
311                    "a_catalog",
312                    "b_catalog",
313                    "b_catalog",
314                    "b_catalog",
315                    "b_catalog",
316                ])) as ArrayRef,
317                Arc::new(StringArray::from(vec![
318                    "a_schema", "a_schema", "b_schema", "b_schema", "a_schema", "a_schema",
319                    "b_schema", "b_schema",
320                ])) as ArrayRef,
321                Arc::new(StringArray::from(vec![
322                    "a_table", "b_table", "a_table", "b_table", "a_table", "a_table", "b_table",
323                    "b_table",
324                ])) as ArrayRef,
325                Arc::new(StringArray::from(vec![
326                    "TABLE", "TABLE", "TABLE", "TABLE", "TABLE", "VIEW", "TABLE", "VIEW",
327                ])) as ArrayRef,
328            ],
329        )
330        .unwrap()
331    }
332
333    fn get_ref_builder(
334        catalog: Option<&str>,
335        db_schema_filter_pattern: Option<&str>,
336        table_name_filter_pattern: Option<&str>,
337        table_types: Vec<&str>,
338        include_schema: bool,
339    ) -> GetTablesBuilder {
340        let dummy_schema = Schema::empty();
341        let tables = [
342            ("a_catalog", "a_schema", "a_table", "TABLE"),
343            ("a_catalog", "a_schema", "b_table", "TABLE"),
344            ("a_catalog", "b_schema", "a_table", "TABLE"),
345            ("a_catalog", "b_schema", "b_table", "TABLE"),
346            ("b_catalog", "a_schema", "a_table", "TABLE"),
347            ("b_catalog", "a_schema", "a_table", "VIEW"),
348            ("b_catalog", "b_schema", "b_table", "TABLE"),
349            ("b_catalog", "b_schema", "b_table", "VIEW"),
350        ];
351        let mut builder = GetTablesBuilder::new(
352            catalog,
353            db_schema_filter_pattern,
354            table_name_filter_pattern,
355            table_types,
356            include_schema,
357        );
358        for (catalog_name, schema_name, table_name, table_type) in tables {
359            builder
360                .append(
361                    catalog_name,
362                    schema_name,
363                    table_name,
364                    table_type,
365                    &dummy_schema,
366                )
367                .unwrap();
368        }
369        builder
370    }
371
372    #[test]
373    fn test_tables_are_filtered() {
374        let ref_batch = get_ref_batch();
375
376        let builder = get_ref_builder(None, None, None, Vec::new(), false);
377        let table_batch = builder.build().unwrap();
378        assert_eq!(table_batch, ref_batch);
379
380        let builder = get_ref_builder(None, Some("a%"), Some("a%"), Vec::new(), false);
381        let table_batch = builder.build().unwrap();
382        let indices = UInt32Array::from(vec![0, 4, 5]);
383        let ref_filtered = RecordBatch::try_new(
384            get_tables_schema(false),
385            ref_batch
386                .columns()
387                .iter()
388                .map(|c| take(c, &indices, None))
389                .collect::<std::result::Result<Vec<_>, _>>()
390                .unwrap(),
391        )
392        .unwrap();
393        assert_eq!(table_batch, ref_filtered);
394
395        let builder = get_ref_builder(Some("a_catalog"), None, None, Vec::new(), false);
396        let table_batch = builder.build().unwrap();
397        let indices = UInt32Array::from(vec![0, 1, 2, 3]);
398        let ref_filtered = RecordBatch::try_new(
399            get_tables_schema(false),
400            ref_batch
401                .columns()
402                .iter()
403                .map(|c| take(c, &indices, None))
404                .collect::<std::result::Result<Vec<_>, _>>()
405                .unwrap(),
406        )
407        .unwrap();
408        assert_eq!(table_batch, ref_filtered);
409
410        let builder = get_ref_builder(None, None, None, vec!["VIEW"], false);
411        let table_batch = builder.build().unwrap();
412        let indices = UInt32Array::from(vec![5, 7]);
413        let ref_filtered = RecordBatch::try_new(
414            get_tables_schema(false),
415            ref_batch
416                .columns()
417                .iter()
418                .map(|c| take(c, &indices, None))
419                .collect::<std::result::Result<Vec<_>, _>>()
420                .unwrap(),
421        )
422        .unwrap();
423        assert_eq!(table_batch, ref_filtered);
424    }
425
426    #[test]
427    fn test_tables_are_sorted() {
428        let ref_batch = get_ref_batch();
429        let dummy_schema = Schema::empty();
430
431        let tables = [
432            ("b_catalog", "a_schema", "a_table", "TABLE"),
433            ("b_catalog", "b_schema", "b_table", "TABLE"),
434            ("b_catalog", "b_schema", "b_table", "VIEW"),
435            ("b_catalog", "a_schema", "a_table", "VIEW"),
436            ("a_catalog", "a_schema", "a_table", "TABLE"),
437            ("a_catalog", "b_schema", "a_table", "TABLE"),
438            ("a_catalog", "b_schema", "b_table", "TABLE"),
439            ("a_catalog", "a_schema", "b_table", "TABLE"),
440        ];
441        let mut builder = GetTablesBuilder::new(
442            None::<String>,
443            None::<String>,
444            None::<String>,
445            None::<String>,
446            false,
447        );
448        for (catalog_name, schema_name, table_name, table_type) in tables {
449            builder
450                .append(
451                    catalog_name,
452                    schema_name,
453                    table_name,
454                    table_type,
455                    &dummy_schema,
456                )
457                .unwrap();
458        }
459        let table_batch = builder.build().unwrap();
460        assert_eq!(table_batch, ref_batch);
461    }
462}
arrow_flight/sql/metadata/tables.rs

arrow_flight/sql/metadata/
tables.rs
