micromegas_analytics/lakehouse/
list_view_sets_table_function.rs

1use super::view_factory::ViewFactory;
2use crate::lakehouse::catalog::list_view_sets;
3use async_trait::async_trait;
4use datafusion::arrow::array::{ArrayRef, BinaryArray, BooleanArray, StringArray};
5use datafusion::arrow::datatypes::DataType;
6use datafusion::arrow::datatypes::Field;
7use datafusion::arrow::datatypes::Schema;
8use datafusion::arrow::datatypes::SchemaRef;
9use datafusion::arrow::record_batch::RecordBatch;
10use datafusion::catalog::Session;
11use datafusion::catalog::TableFunctionArgs;
12use datafusion::catalog::TableFunctionImpl;
13use datafusion::catalog::TableProvider;
14use datafusion::datasource::TableType;
15use datafusion::datasource::memory::{DataSourceExec, MemorySourceConfig};
16use datafusion::error::DataFusionError;
17use datafusion::physical_plan::ExecutionPlan;
18use datafusion::prelude::Expr;
19use std::sync::Arc;
20
21/// A DataFusion `TableFunctionImpl` for listing view sets with their current schema information.
22#[derive(Debug)]
23pub struct ListViewSetsTableFunction {
24    view_factory: Arc<ViewFactory>,
25}
26
27impl ListViewSetsTableFunction {
28    pub fn new(view_factory: Arc<ViewFactory>) -> Self {
29        Self { view_factory }
30    }
31}
32
33impl TableFunctionImpl for ListViewSetsTableFunction {
34    fn call_with_args(
35        &self,
36        _args: TableFunctionArgs,
37    ) -> datafusion::error::Result<Arc<dyn TableProvider>> {
38        Ok(Arc::new(ListViewSetsTableProvider {
39            view_factory: self.view_factory.clone(),
40        }))
41    }
42}
43
44/// A DataFusion `TableProvider` for listing view sets with their current schema information.
45#[derive(Debug)]
46pub struct ListViewSetsTableProvider {
47    pub view_factory: Arc<ViewFactory>,
48}
49
50#[async_trait]
51impl TableProvider for ListViewSetsTableProvider {
52    fn schema(&self) -> SchemaRef {
53        Arc::new(Schema::new(vec![
54            Field::new("view_set_name", DataType::Utf8, false),
55            Field::new("current_schema_hash", DataType::Binary, false),
56            Field::new("schema", DataType::Utf8, false),
57            Field::new("has_view_maker", DataType::Boolean, false),
58            Field::new("global_instance_available", DataType::Boolean, false),
59        ]))
60    }
61
62    fn table_type(&self) -> TableType {
63        TableType::Temporary
64    }
65
66    async fn scan(
67        &self,
68        _state: &dyn Session,
69        projection: Option<&Vec<usize>>,
70        _filters: &[Expr],
71        limit: Option<usize>,
72    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
73        // Get current schema versions from the view factory
74        let schema_infos =
75            list_view_sets(&self.view_factory).map_err(|e| DataFusionError::External(e.into()))?;
76
77        // Apply limit early to avoid building unnecessary arrays.
78        // DataFusion trusts us to apply the limit - if we ignore it, too many rows
79        // will be returned to the client.
80        let limited_infos: &[_] = if let Some(n) = limit {
81            &schema_infos[..n.min(schema_infos.len())]
82        } else {
83            &schema_infos
84        };
85
86        // Convert to Arrow arrays
87        let view_set_names: Vec<String> = limited_infos
88            .iter()
89            .map(|info| info.view_set_name.clone())
90            .collect();
91        let schema_hashes: Vec<&[u8]> = limited_infos
92            .iter()
93            .map(|info| info.current_schema_hash.as_slice())
94            .collect();
95        let schemas: Vec<String> = limited_infos
96            .iter()
97            .map(|info| info.schema.clone())
98            .collect();
99        let has_view_makers: Vec<bool> = limited_infos
100            .iter()
101            .map(|info| info.has_view_maker)
102            .collect();
103        let global_instances: Vec<bool> = limited_infos
104            .iter()
105            .map(|info| info.global_instance_available)
106            .collect();
107
108        let view_set_name_array: ArrayRef = Arc::new(StringArray::from(view_set_names));
109        let schema_hash_array: ArrayRef = Arc::new(BinaryArray::from(schema_hashes));
110        let schema_array: ArrayRef = Arc::new(StringArray::from(schemas));
111        let has_view_maker_array: ArrayRef = Arc::new(BooleanArray::from(has_view_makers));
112        let global_instance_array: ArrayRef = Arc::new(BooleanArray::from(global_instances));
113
114        let columns = vec![
115            view_set_name_array,
116            schema_hash_array,
117            schema_array,
118            has_view_maker_array,
119            global_instance_array,
120        ];
121
122        let record_batch = RecordBatch::try_new(self.schema(), columns)
123            .map_err(|e| DataFusionError::External(e.into()))?;
124
125        let source = MemorySourceConfig::try_new(
126            &[vec![record_batch]],
127            self.schema(),
128            projection.map(|v| v.to_owned()),
129        )?;
130        Ok(DataSourceExec::from_data_source(source))
131    }
132}
micromegas_analytics/lakehouse/list_view_sets_table_function.rs

micromegas_analytics/lakehouse/
list_view_sets_table_function.rs
