arrow_flight/sql/metadata/
tables.rs

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

arrow_flight/sql/metadata/
tables.rs
