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

arrow_flight/sql/metadata/
table_types.rs
