arrow_flight/sql/metadata/
table_types.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//! [`GetTableTypesBuilder`] for building responses to [`CommandGetTableTypes`] queries.
19//!
20//! [`CommandGetTableTypes`]: crate::sql::CommandGetTableTypes
21
22use std::sync::Arc;
23
24use arrow_array::{builder::StringBuilder, ArrayRef, RecordBatch};
25use arrow_schema::{DataType, Field, Schema, SchemaRef};
26use arrow_select::take::take;
27use once_cell::sync::Lazy;
28
29use crate::error::*;
30use crate::sql::CommandGetTableTypes;
31
32use super::lexsort_to_indices;
33
34/// A builder for a [`CommandGetTableTypes`] response.
35///
36/// Builds rows like this:
37///
38/// * table_type: utf8,
39#[derive(Default)]
40pub struct GetTableTypesBuilder {
41    // array builder for table types
42    table_type: StringBuilder,
43}
44
45impl CommandGetTableTypes {
46    /// Create a builder suitable for constructing a response
47    pub fn into_builder(self) -> GetTableTypesBuilder {
48        self.into()
49    }
50}
51
52impl From<CommandGetTableTypes> for GetTableTypesBuilder {
53    fn from(_value: CommandGetTableTypes) -> Self {
54        Self::new()
55    }
56}
57
58impl GetTableTypesBuilder {
59    /// Create a new instance of [`GetTableTypesBuilder`]
60    pub fn new() -> Self {
61        Self {
62            table_type: StringBuilder::new(),
63        }
64    }
65
66    /// Append a row
67    pub fn append(&mut self, table_type: impl AsRef<str>) {
68        self.table_type.append_value(table_type);
69    }
70
71    /// builds a `RecordBatch` with the correct schema for a `CommandGetTableTypes` response
72    pub fn build(self) -> Result<RecordBatch> {
73        let schema = self.schema();
74        let Self { mut table_type } = self;
75
76        // Make the arrays
77        let table_type = table_type.finish();
78
79        let batch = RecordBatch::try_new(schema, vec![Arc::new(table_type) as ArrayRef])?;
80
81        // Order filtered results by table_type
82        let indices = lexsort_to_indices(batch.columns());
83        let columns = batch
84            .columns()
85            .iter()
86            .map(|c| take(c, &indices, None))
87            .collect::<std::result::Result<Vec<_>, _>>()?;
88
89        Ok(RecordBatch::try_new(batch.schema(), columns)?)
90    }
91
92    /// Return the schema of the RecordBatch that will be returned
93    /// from [`CommandGetTableTypes`]
94    pub fn schema(&self) -> SchemaRef {
95        get_table_types_schema()
96    }
97}
98
99fn get_table_types_schema() -> SchemaRef {
100    Arc::clone(&GET_TABLE_TYPES_SCHEMA)
101}
102
103/// The schema for [`CommandGetTableTypes`].
104static GET_TABLE_TYPES_SCHEMA: Lazy<SchemaRef> = Lazy::new(|| {
105    Arc::new(Schema::new(vec![Field::new(
106        "table_type",
107        DataType::Utf8,
108        false,
109    )]))
110});
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115    use arrow_array::StringArray;
116
117    fn get_ref_batch() -> RecordBatch {
118        RecordBatch::try_new(
119            get_table_types_schema(),
120            vec![Arc::new(StringArray::from(vec![
121                "a_table_type",
122                "b_table_type",
123                "c_table_type",
124                "d_table_type",
125            ])) as ArrayRef],
126        )
127        .unwrap()
128    }
129
130    #[test]
131    fn test_table_types_are_sorted() {
132        let ref_batch = get_ref_batch();
133
134        let mut builder = GetTableTypesBuilder::new();
135        builder.append("b_table_type");
136        builder.append("a_table_type");
137        builder.append("d_table_type");
138        builder.append("c_table_type");
139        let schema_batch = builder.build().unwrap();
140
141        assert_eq!(schema_batch, ref_batch)
142    }
143
144    #[test]
145    fn test_builder_from_query() {
146        let ref_batch = get_ref_batch();
147        let query = CommandGetTableTypes {};
148
149        let mut builder = query.into_builder();
150        builder.append("a_table_type");
151        builder.append("b_table_type");
152        builder.append("c_table_type");
153        builder.append("d_table_type");
154        let schema_batch = builder.build().unwrap();
155
156        assert_eq!(schema_batch, ref_batch)
157    }
158}
arrow_flight/sql/metadata/table_types.rs

arrow_flight/sql/metadata/
table_types.rs
