dbkit_core/
func.rs

1use crate::compile::CompiledSql;
2use crate::expr::{Expr, ExprNode, IntoExpr, NumericExprType, VectorBinaryOp};
3use crate::query::Select;
4use crate::PgVector;
5
6pub trait StringUnaryExpr {
7    type Output;
8}
9
10impl StringUnaryExpr for String {
11    type Output = String;
12}
13
14impl StringUnaryExpr for Option<String> {
15    type Output = Option<String>;
16}
17
18pub trait StringLengthExpr {
19    type Output;
20}
21
22impl StringLengthExpr for String {
23    type Output = i32;
24}
25
26impl StringLengthExpr for Option<String> {
27    type Output = Option<i32>;
28}
29
30fn unary_string_fn<T>(name: &'static str, arg: impl IntoExpr<T>) -> Expr<<T as StringUnaryExpr>::Output>
31where
32    T: StringUnaryExpr,
33{
34    let expr = arg.into_expr();
35    Expr::new(ExprNode::Func {
36        name,
37        args: vec![expr.node],
38    })
39}
40
41fn string_length_fn<T>(name: &'static str, arg: impl IntoExpr<T>) -> Expr<<T as StringLengthExpr>::Output>
42where
43    T: StringLengthExpr,
44{
45    let expr = arg.into_expr();
46    Expr::new(ExprNode::Func {
47        name,
48        args: vec![expr.node],
49    })
50}
51
52pub fn upper<T>(arg: impl IntoExpr<T>) -> Expr<<T as StringUnaryExpr>::Output>
53where
54    T: StringUnaryExpr,
55{
56    unary_string_fn("UPPER", arg)
57}
58
59pub fn trim<T>(arg: impl IntoExpr<T>) -> Expr<<T as StringUnaryExpr>::Output>
60where
61    T: StringUnaryExpr,
62{
63    unary_string_fn("TRIM", arg)
64}
65
66pub fn char_length<T>(arg: impl IntoExpr<T>) -> Expr<<T as StringLengthExpr>::Output>
67where
68    T: StringLengthExpr,
69{
70    string_length_fn("CHAR_LENGTH", arg)
71}
72
73pub fn count<T>(arg: impl IntoExpr<T>) -> Expr<i64> {
74    let expr = arg.into_expr();
75    Expr::new(ExprNode::Func {
76        name: "COUNT",
77        args: vec![expr.node],
78    })
79}
80
81pub fn sum<T>(arg: impl IntoExpr<T>) -> Expr<T> {
82    let expr = arg.into_expr();
83    Expr::new(ExprNode::Func {
84        name: "SUM",
85        args: vec![expr.node],
86    })
87}
88
89pub trait NullableAggregateOutput {
90    type Output;
91}
92
93macro_rules! impl_nullable_aggregate_output {
94    ($($ty:ty),+ $(,)?) => {
95        $(
96            impl NullableAggregateOutput for $ty {
97                type Output = Option<$ty>;
98            }
99
100            impl NullableAggregateOutput for Option<$ty> {
101                type Output = Option<$ty>;
102            }
103        )+
104    };
105}
106
107impl_nullable_aggregate_output!(
108    String,
109    i16,
110    i32,
111    i64,
112    f32,
113    f64,
114    uuid::Uuid,
115    chrono::NaiveDateTime,
116    chrono::DateTime<chrono::Utc>,
117    chrono::NaiveDate,
118    chrono::NaiveTime,
119    crate::PgInterval,
120);
121
122pub fn min<T>(arg: impl IntoExpr<T>) -> Expr<<T as NullableAggregateOutput>::Output>
123where
124    T: NullableAggregateOutput,
125{
126    let expr = arg.into_expr();
127    Expr::new(ExprNode::Func {
128        name: "MIN",
129        args: vec![expr.node],
130    })
131}
132
133pub fn max<T>(arg: impl IntoExpr<T>) -> Expr<<T as NullableAggregateOutput>::Output>
134where
135    T: NullableAggregateOutput,
136{
137    let expr = arg.into_expr();
138    Expr::new(ExprNode::Func {
139        name: "MAX",
140        args: vec![expr.node],
141    })
142}
143
144pub fn coalesce<T>(a: impl IntoExpr<T>, b: impl IntoExpr<T>) -> Expr<T> {
145    let left = a.into_expr();
146    let right = b.into_expr();
147    Expr::new(ExprNode::Func {
148        name: "COALESCE",
149        args: vec![left.node, right.node],
150    })
151}
152
153pub fn least<T>(a: impl IntoExpr<T>, b: impl IntoExpr<T>) -> Expr<T> {
154    let left = a.into_expr();
155    let right = b.into_expr();
156    Expr::new(ExprNode::Func {
157        name: "LEAST",
158        args: vec![left.node, right.node],
159    })
160}
161
162pub fn greatest<T>(a: impl IntoExpr<T>, b: impl IntoExpr<T>) -> Expr<T> {
163    let left = a.into_expr();
164    let right = b.into_expr();
165    Expr::new(ExprNode::Func {
166        name: "GREATEST",
167        args: vec![left.node, right.node],
168    })
169}
170
171pub fn power<B, E>(base: impl IntoExpr<B>, exponent: impl IntoExpr<E>) -> Expr<f64>
172where
173    B: NumericExprType,
174    E: NumericExprType,
175{
176    let base = base.into_expr();
177    let exponent = exponent.into_expr();
178    Expr::new(ExprNode::Func {
179        name: "POWER",
180        args: vec![base.node, exponent.node],
181    })
182}
183
184pub fn date_trunc<T>(part: impl IntoExpr<String>, value: impl IntoExpr<T>) -> Expr<T> {
185    let part = part.into_expr();
186    let value = value.into_expr();
187    Expr::new(ExprNode::Func {
188        name: "DATE_TRUNC",
189        args: vec![part.node, value.node],
190    })
191}
192
193fn exists_expr(subquery: CompiledSql) -> Expr<bool> {
194    Expr::new(ExprNode::Exists { subquery })
195}
196
197pub fn exists<Out, Loads, Lock, DistinctState, GroupState>(subquery: Select<Out, Loads, Lock, DistinctState, GroupState>) -> Expr<bool> {
198    exists_expr(subquery.compile_for_exists())
199}
200
201/// Marker trait for values that can participate in vector distance/similarity expressions.
202pub trait VectorExpr<const N: usize> {}
203
204impl<const N: usize> VectorExpr<N> for PgVector<N> {}
205impl<const N: usize> VectorExpr<N> for Option<PgVector<N>> {}
206
207fn vector_binary_fn<const N: usize, L, R>(name: &'static str, left: impl IntoExpr<L>, right: impl IntoExpr<R>) -> Expr<f32>
208where
209    L: VectorExpr<N>,
210    R: VectorExpr<N>,
211{
212    let left = left.into_expr();
213    let right = right.into_expr();
214    Expr::new(ExprNode::Func {
215        name,
216        args: vec![left.node, right.node],
217    })
218}
219
220fn vector_binary_operator<const N: usize, L, R>(op: VectorBinaryOp, left: impl IntoExpr<L>, right: impl IntoExpr<R>) -> Expr<f32>
221where
222    L: VectorExpr<N>,
223    R: VectorExpr<N>,
224{
225    let left = left.into_expr();
226    let right = right.into_expr();
227    Expr::new(ExprNode::VectorBinary {
228        left: Box::new(left.node),
229        op,
230        right: Box::new(right.node),
231    })
232}
233
234/// Euclidean (L2) distance using pgvector's `<->` operator.
235///
236/// Lower is more similar.
237///
238/// ANN note:
239/// - This form is operator-based and can use pgvector ivfflat/hnsw indexes for
240///   `ORDER BY ... LIMIT` nearest-neighbor queries.
241pub fn l2_distance<const N: usize, L, R>(left: impl IntoExpr<L>, right: impl IntoExpr<R>) -> Expr<f32>
242where
243    L: VectorExpr<N>,
244    R: VectorExpr<N>,
245{
246    vector_binary_operator::<N, L, R>(VectorBinaryOp::L2Distance, left, right)
247}
248
249/// Cosine distance using pgvector's `<=>` operator.
250///
251/// Lower is more similar.
252///
253/// ANN note:
254/// - This form is operator-based and can use pgvector ivfflat/hnsw indexes for
255///   `ORDER BY ... LIMIT` nearest-neighbor queries.
256pub fn cosine_distance<const N: usize, L, R>(left: impl IntoExpr<L>, right: impl IntoExpr<R>) -> Expr<f32>
257where
258    L: VectorExpr<N>,
259    R: VectorExpr<N>,
260{
261    vector_binary_operator::<N, L, R>(VectorBinaryOp::CosineDistance, left, right)
262}
263
264/// True inner product as a function expression (`INNER_PRODUCT(a, b)`).
265///
266/// Higher is more similar (for normalized embeddings, identical vectors are `1.0`).
267///
268/// ANN warning:
269/// - This is intentionally a function call to preserve true inner-product semantics,
270///   but function expressions are generally not pgvector ANN index-compatible for
271///   `ORDER BY ... LIMIT`.
272/// - For ANN-indexed retrieval, use [`inner_product_distance`] with `ORDER BY ASC`.
273pub fn inner_product<const N: usize, L, R>(left: impl IntoExpr<L>, right: impl IntoExpr<R>) -> Expr<f32>
274where
275    L: VectorExpr<N>,
276    R: VectorExpr<N>,
277{
278    vector_binary_fn::<N, L, R>("INNER_PRODUCT", left, right)
279}
280
281/// L1 (Manhattan) distance using pgvector's `<+>` operator.
282///
283/// Lower is more similar.
284///
285/// ANN note:
286/// - This form is operator-based and can use pgvector ivfflat/hnsw indexes for
287///   `ORDER BY ... LIMIT` nearest-neighbor queries.
288pub fn l1_distance<const N: usize, L, R>(left: impl IntoExpr<L>, right: impl IntoExpr<R>) -> Expr<f32>
289where
290    L: VectorExpr<N>,
291    R: VectorExpr<N>,
292{
293    vector_binary_operator::<N, L, R>(VectorBinaryOp::L1Distance, left, right)
294}
295
296/// Negative inner-product distance using pgvector's `<#>` operator.
297///
298/// Lower is more similar, so nearest-neighbor queries should use `ORDER BY ASC`.
299///
300/// ANN note:
301/// - This form is operator-based and can use pgvector ivfflat/hnsw indexes for
302///   `ORDER BY ... LIMIT` nearest-neighbor queries.
303/// - Thresholds are inverted relative to true inner product
304///   (for example `inner_product > 0.9` corresponds to
305///   `inner_product_distance < -0.9`).
306pub fn inner_product_distance<const N: usize, L, R>(left: impl IntoExpr<L>, right: impl IntoExpr<R>) -> Expr<f32>
307where
308    L: VectorExpr<N>,
309    R: VectorExpr<N>,
310{
311    vector_binary_operator::<N, L, R>(VectorBinaryOp::InnerProductDistance, left, right)
312}
dbkit_core/func.rs

dbkit_core/
func.rs
