1use crate::SQLError;
10use uqa_core::{
11 memory::{Produced, ProductionControl, ProductionVec},
12 Value,
13};
14
15use crate::ast::ColumnType;
16
17#[derive(Debug, Clone)]
19pub enum SQLParam {
20 Scalar(Value),
21 TypedScalar {
23 value: Value,
24 ty: ColumnType,
25 },
26 #[doc(hidden)]
29 DeferredError {
30 ty: ColumnType,
31 error: crate::ast::DeferredSQLError,
32 },
33 Vector(Vec<f32>),
34 Tensor(Vec<Vec<f32>>),
35}
36
37impl SQLParam {
38 pub fn scalar(value: Value) -> Self {
39 Self::Scalar(value)
40 }
41
42 #[must_use]
43 pub fn typed_scalar(value: Value, ty: ColumnType) -> Self {
44 Self::TypedScalar { value, ty }
45 }
46
47 #[must_use]
49 pub fn scalar_value(&self) -> Option<&Value> {
50 match self {
51 Self::Scalar(value) | Self::TypedScalar { value, .. } => Some(value),
52 Self::Vector(_) | Self::Tensor(_) | Self::DeferredError { .. } => None,
53 }
54 }
55
56 #[must_use]
58 pub fn declared_scalar_type(&self) -> Option<&ColumnType> {
59 match self {
60 Self::TypedScalar { ty, .. } | Self::DeferredError { ty, .. } => Some(ty),
61 Self::Scalar(_) | Self::Vector(_) | Self::Tensor(_) => None,
62 }
63 }
64
65 pub fn to_value(&self) -> Result<Value, SQLError> {
67 self.to_value_with_control(&ProductionControl::uncontrolled())
68 .map(|value| value.into_uncontrolled().expect("ordinary parameter value"))
69 }
70
71 pub fn to_value_with_control(
73 &self,
74 control: &ProductionControl<'_>,
75 ) -> Result<Produced<Value>, SQLError> {
76 control.check()?;
77 match self {
78 Self::DeferredError { error, .. } => Err(error.clone().into()),
79 Self::Scalar(value) | Self::TypedScalar { value, .. } => Ok(control.copy_value(value)?),
80 Self::Vector(values) => vector_value(values, control),
81 Self::Tensor(vectors) => {
82 let mut output = ProductionVec::new(*control);
83 output.reserve(vectors.len())?;
84 for values in vectors {
85 output.push_produced(vector_value(values, control)?)?;
86 }
87 let (values, memory) = output.finish()?.into_parts();
88 Ok(control.finish(Value::List(values), memory)?)
89 }
90 }
91 }
92
93 pub fn vector(v: Vec<f32>) -> Self {
94 Self::Vector(v)
95 }
96
97 pub fn tensor(v: Vec<Vec<f32>>) -> Self {
98 Self::Tensor(v)
99 }
100}
101
102fn vector_value(
103 values: &[f32],
104 control: &ProductionControl<'_>,
105) -> Result<Produced<Value>, SQLError> {
106 let mut output = ProductionVec::new(*control);
107 output.reserve(values.len())?;
108 for value in values {
109 output.push_produced(
110 control.finish(Value::Float(f64::from(*value)), control.empty_reservation())?,
111 )?;
112 }
113 let (values, memory) = output.finish()?.into_parts();
114 Ok(control.finish(Value::List(values), memory)?)
115}
116
117#[cfg(test)]
118mod tests {
119 use super::*;
120
121 #[test]
122 fn typed_scalar_preserves_declared_type_without_changing_scalar_access() {
123 let value = Value::Int(7);
124 let typed = SQLParam::typed_scalar(value.clone(), ColumnType::SmallInteger);
125 assert_eq!(typed.scalar_value(), Some(&value));
126 assert_eq!(
127 typed.declared_scalar_type(),
128 Some(&ColumnType::SmallInteger)
129 );
130
131 let scalar = SQLParam::scalar(value.clone());
132 assert_eq!(scalar.scalar_value(), Some(&value));
133 assert_eq!(scalar.declared_scalar_type(), None);
134 }
135
136 #[test]
137 fn parameter_value_owners_keep_tensor_and_scalar_payloads_until_drop() {
138 use uqa_core::{memory::MemoryBudget, CancellationToken};
139 let budget = MemoryBudget::new(1 << 16);
140 let token = CancellationToken::new();
141 let control = ProductionControl::new(&budget, &token, &token);
142 for (parameter, expected) in [
143 (
144 SQLParam::typed_scalar(Value::Str("payload".repeat(32)), ColumnType::Text),
145 Value::Str("payload".repeat(32)),
146 ),
147 (
148 SQLParam::Vector(vec![1.0, 2.5]),
149 Value::List(vec![Value::Float(1.0), Value::Float(2.5)]),
150 ),
151 (
152 SQLParam::Tensor(vec![vec![1.0, 2.5], vec![]]),
153 Value::List(vec![
154 Value::List(vec![Value::Float(1.0), Value::Float(2.5)]),
155 Value::List(vec![]),
156 ]),
157 ),
158 ] {
159 let value = parameter.to_value_with_control(&control).unwrap();
160 assert_eq!(*value, expected);
161 assert!(value.reserved_bytes() > 0);
162 assert_eq!(budget.used(), value.reserved_bytes());
163 drop(value);
164 assert_eq!(budget.used(), 0);
165 assert_eq!(parameter.to_value().unwrap(), expected);
166 }
167 }
168
169 #[test]
170 fn parameter_production_releases_partial_output_on_quota_and_both_tokens() {
171 use uqa_core::{memory::MemoryBudget, CancellationToken};
172 let budget = MemoryBudget::new(256);
173 let original = CancellationToken::new();
174 let invoking = CancellationToken::new();
175 let control = ProductionControl::new(&budget, &original, &invoking);
176 let held = control.copy_text("held").unwrap();
177 let parameter = SQLParam::Tensor(vec![vec![1.0; 64], vec![2.0; 64]]);
178 assert_eq!(
179 parameter
180 .to_value_with_control(&control)
181 .unwrap_err()
182 .sqlstate(),
183 Some("53200")
184 );
185 assert_eq!(budget.used(), held.reserved_bytes());
186 for token in [&original, &invoking] {
187 token.cancel();
188 assert_eq!(
189 parameter
190 .to_value_with_control(&control)
191 .unwrap_err()
192 .sqlstate(),
193 Some("57014")
194 );
195 assert_eq!(budget.used(), held.reserved_bytes());
196 token.reset();
197 }
198 assert_eq!(
199 parameter.to_value().unwrap(),
200 Value::List(vec![
201 Value::List(vec![Value::Float(1.0); 64]),
202 Value::List(vec![Value::Float(2.0); 64])
203 ])
204 );
205 }
206}