uqa-engine 0.3.6

Engine: schema-aware table store, catalog restore, transactions
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Compose scalar execution services from current session and catalog state.
use crate::{Engine, TableState};
use std::{collections::BTreeMap, sync::Arc};
use uqa_core::DocId;
use uqa_execution::query::{
    model_training::{ModelTrainingContext, TrainedModels, TrainingTable, TrainingTables},
    scalar_functions::{ScalarFunctionContext, ScalarSession},
    scalar_projection::AnalyzerRevisions,
};
use uqa_ml::DeepModel;
use uqa_sql::SQLError;
use uqa_storage::{document_store::Document, StorageBackendResult};
impl ScalarSession for Engine {
    fn backend_process_id(&self) -> i32 {
        self.backend_process_id()
    }
    fn notify(&self, channel: &str, payload: &str) -> Result<(), SQLError> {
        self.notify(channel, payload)
    }
    fn notification_queue_usage(&self) -> Result<f64, SQLError> {
        self.notification_queue_usage()
    }
}
impl AnalyzerRevisions for Engine {
    fn analyzer_revision(&self, name: &str) -> Result<Arc<uqa_analysis::CompiledAnalyzer>, String> {
        self.resolve_analyzer_revision(name)
    }
}

struct TrainingTableState(Arc<TableState>);

impl TrainingTable for TrainingTableState {
    fn doc_ids(&self) -> StorageBackendResult<Vec<DocId>> {
        self.0.document_store.read().doc_ids()
    }
}

impl TrainingTables for Engine {
    fn training_table(
        &self,
        name: &str,
    ) -> StorageBackendResult<Option<Box<dyn TrainingTable + '_>>> {
        self.try_table(name).map(|table| {
            table.map(|state| Box::new(TrainingTableState(state)) as Box<dyn TrainingTable>)
        })
    }

    fn training_documents(
        &self,
        table: &str,
        doc_ids: &[DocId],
        projection: &[String],
    ) -> Result<BTreeMap<DocId, Document>, SQLError> {
        self.get_documents_with_materialized_projection(table, doc_ids, projection)
    }
}

impl TrainedModels for Engine {
    fn save_model(&self, name: &str, model: &DeepModel) -> Result<(), SQLError> {
        self.save_model(name, model)
    }
}

impl Engine {
    pub(crate) fn model_training_context(&self) -> ModelTrainingContext<'_> {
        ModelTrainingContext {
            tables: self,
            models: self,
        }
    }

    pub(crate) fn scalar_function_context(&self) -> ScalarFunctionContext<'_> {
        ScalarFunctionContext {
            catalog: self.catalog_execution(),
            sequences: self.sequence_introspection_context(),
            names: self,
            roles: self,
            database: self.database_privilege_inquiry(),
            schemas: self.schema_privilege_inquiry(),
            sequence_privileges: self.sequence_privilege_inquiry(),
            tables: self.table_privilege_context(),
            session: self,
            graphs: self,
            models: self.model_training_context(),
            analyzers: self,
            runtime: self.query_runtime_view(),
        }
    }
}