uqa-sql 0.4.0

PostgreSQL-compatible SQL compiler built on libpg_query
Documentation 
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

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

//! Vector conversion reserves destination buffers before evaluating their elements.

use super::{column_type_name, validate_vector_dimensions, ColumnType, SQLError, Value};
use crate::expr::{tensor_items, vector_element, vector_items};
use uqa_core::{
    memory::{Budgeted, BudgetedVec, MemoryBudget, MemoryError},
    QueryCancelled,
};

/// Preserve assignment conversion while charging tensor headers and float buffers to the caller's allowance. The result transfers its reservations together with its buffers to the retaining index owner.
pub fn index_vectors_for_type_budgeted(
    value: &Value,
    ty: &ColumnType,
    memory: &MemoryBudget,
    poll: &mut impl FnMut() -> Result<(), QueryCancelled>,
) -> Result<Budgeted<Vec<Vec<f32>>>, SQLError> {
    poll()?;
    if matches!(value, Value::Null) {
        return Ok(Budgeted::new(Vec::new(), memory.empty_reservation()));
    }
    let (values, dimensions) = match ty {
        ColumnType::Vector(dim) => (std::slice::from_ref(value), *dim),
        ColumnType::Tensor(dim) => (tensor_items(value)?, *dim),
        _ => {
            return Err(SQLError::TypeMismatch(format!(
                "{} is not vector-indexable",
                column_type_name(ty)
            )))
        }
    };
    // The nested buffers must drop before their payload reservation on every error.
    let mut pending = (BudgetedVec::new(memory), memory.empty_reservation());
    pending.0.reserve(values.len()).map_err(memory_error)?;
    for value in values {
        poll()?;
        let items = vector_items(value)?;
        let mut vector = BudgetedVec::new(memory);
        vector.reserve(items.len()).map_err(memory_error)?;
        for item in items {
            poll()?;
            vector.push(vector_element(item)?).map_err(memory_error)?;
        }
        let (vector, reservation) = vector.into_parts();
        pending.0.push(vector).map_err(memory_error)?;
        pending.1.absorb(reservation);
    }
    for vector in pending.0.iter() {
        poll()?;
        validate_vector_dimensions(dimensions, vector.len())?;
    }
    poll()?;
    let (vectors, mut reservation) = pending.0.into_parts();
    reservation.absorb(pending.1);
    Ok(Budgeted::new(vectors, reservation))
}

fn memory_error(error: MemoryError) -> SQLError {
    SQLError::Routine {
        sqlstate: "53200".into(),
        message: format!("convert vector index input: {error}"),
    }
}

#[cfg(test)]
mod tests;