Implement casting for Vector to HalfVec

Author: A2va

src/diesel_ext/expression_methods.rs

@@ -2,6 +2,9 @@ use diesel::expression::{AsExpression, Expression};
 use diesel::pg::Pg;
 use diesel::sql_types::{Double, SqlType};
 
+#[cfg(feature = "halfvec")]
+use crate::diesel_ext::halfvec::HalfVecCast;
+
 diesel::infix_operator!(L2Distance, " <-> ", Double, backend: Pg);
 diesel::infix_operator!(MaxInnerProduct, " <#> ", Double, backend: Pg);
 diesel::infix_operator!(CosineDistance, " <=> ", Double, backend: Pg);
@@ -10,6 +13,11 @@ diesel::infix_operator!(HammingDistance, " <~> ", Double, backend: Pg);
 diesel::infix_operator!(JaccardDistance, " <%> ", Double, backend: Pg);
 
 pub trait VectorExpressionMethods: Expression + Sized {
+    #[cfg(feature = "halfvec")]
+    fn cast_to_halfvec(self, dim: usize) -> HalfVecCast<Self> {
+        HalfVecCast::new(self, dim)
+    }
+
     fn l2_distance<T>(self, other: T) -> L2Distance<Self, T::Expression>
     where
         Self::SqlType: SqlType,

src/diesel_ext/halfvec.rs

@@ -1,8 +1,10 @@
 use diesel::deserialize::{self, FromSql};
+use diesel::expression::{AppearsOnTable, Expression, SelectableExpression, ValidGrouping};
 use diesel::pg::{Pg, PgValue};
-use diesel::query_builder::QueryId;
+use diesel::query_builder::{AstPass, QueryFragment, QueryId};
 use diesel::serialize::{self, IsNull, Output, ToSql};
 use diesel::sql_types::SqlType;
+use diesel::QueryResult;
 use std::convert::TryFrom;
 use std::io::Write;
 
@@ -32,6 +34,42 @@ impl FromSql<HalfVectorType, Pg> for HalfVector {
     }
 }
 
+#[derive(Debug, Clone, Copy, QueryId, ValidGrouping)]
+pub struct HalfVecCast<Expr> {
+    pub expr: Expr,
+    pub dim: usize,
+}
+
+impl<Expr> HalfVecCast<Expr> {
+    pub fn new(expr: Expr, dim: usize) -> Self {
+        Self { expr, dim }
+    }
+}
+
+impl<Expr> Expression for HalfVecCast<Expr>
+where
+    Expr: Expression,
+{
+    type SqlType = HalfVectorType;
+}
+
+impl<Expr> QueryFragment<Pg> for HalfVecCast<Expr>
+where
+    Expr: QueryFragment<Pg>,
+{
+    fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, Pg>) -> QueryResult<()> {
+        out.push_sql("(");
+        self.expr.walk_ast(out.reborrow())?;
+        out.push_sql(")::halfvec(");
+        out.push_sql(&self.dim.to_string());
+        out.push_sql(")");
+        Ok(())
+    }
+}
+
+impl<Expr, QS> AppearsOnTable<QS> for HalfVecCast<Expr> where Expr: AppearsOnTable<QS> {}
+impl<Expr, QS> SelectableExpression<QS> for HalfVecCast<Expr> where Expr: SelectableExpression<QS> {}
+
 #[cfg(test)]
 mod tests {
     use crate::{HalfVector, VectorExpressionMethods};