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

arrow_flight/sql/metadata/
table_types.rs
