save

Author: levkk

pgdog/src/backend/pool/connection/buffer.rs

@@ -4,7 +4,7 @@ use std::{cmp::Ordering, collections::VecDeque};
 
 use crate::{
     frontend::router::parser::{Aggregate, OrderBy},
-    net::messages::{DataRow, FromBytes, Message, Protocol, RowDescription, ToBytes},
+    net::messages::{DataRow, FromBytes, Message, Protocol, RowDescription, ToBytes, Vector},
 };
 
 use super::Aggregates;
@@ -34,33 +34,60 @@ impl Buffer {
 
     /// Sort the buffer.
     pub(super) fn sort(&mut self, columns: &[OrderBy], rd: &RowDescription) {
-        // Calculate column indecies once, since
-        // fetching indecies by name is O(n).
+        // Calculate column indicies once, since
+        // fetching indicies by name is O(n).
         let mut cols = vec![];
         for column in columns {
-            if let Some(index) = column.index() {
-                cols.push(Some((index, column.asc())));
-            } else if let Some(name) = column.name() {
-                if let Some(index) = rd.field_index(name) {
-                    cols.push(Some((index, column.asc())));
-                } else {
-                    cols.push(None);
+            match column {
+                OrderBy::Asc(_) => cols.push(column.clone()),
+                OrderBy::AscColumn(name) => {
+                    if let Some(index) = rd.field_index(name) {
+                        cols.push(OrderBy::Asc(index + 1));
+                    }
+                }
+                OrderBy::Desc(_) => cols.push(column.clone()),
+                OrderBy::DescColumn(name) => {
+                    if let Some(index) = rd.field_index(name) {
+                        cols.push(OrderBy::Desc(index + 1));
+                    }
+                }
+                OrderBy::AscVectorL2(_, _) => cols.push(column.clone()),
+                OrderBy::AscVectorL2Column(name, vector) => {
+                    if let Some(index) = rd.field_index(name) {
+                        cols.push(OrderBy::AscVectorL2(index + 1, vector.clone()));
+                    }
                 }
-            } else {
-                cols.push(None);
             };
         }
 
         // Sort rows.
         let order_by = move |a: &DataRow, b: &DataRow| -> Ordering {
-            for col in cols.iter().flatten() {
-                let (index, asc) = col;
-                let left = a.get_column(*index, rd);
-                let right = b.get_column(*index, rd);
+            for col in cols.iter() {
+                let index = col.index();
+                let asc = col.asc();
+                let index = if let Some(index) = index {
+                    index
+                } else {
+                    continue;
+                };
+                let left = a.get_column(index, rd);
+                let right = b.get_column(index, rd);
 
                 let ordering = match (left, right) {
                     (Ok(Some(left)), Ok(Some(right))) => {
-                        if *asc {
+                        if let OrderBy::AscVectorL2(_, vector) = col {
+                            let left: Option<Vector> = left.value.try_into().ok();
+                            let right: Option<Vector> = right.value.try_into().ok();
+
+                            if let (Some(left), Some(right)) = (left, right) {
+                                let left = left.distance_l2(vector);
+                                let right = right.distance_l2(vector);
+
+                                left.partial_cmp(&right)
+                            } else {
+                                Some(Ordering::Equal)
+                            }
+                        } else if asc {
                             left.value.partial_cmp(&right.value)
                         } else {
                             right.value.partial_cmp(&left.value)

pgdog/src/frontend/router/parser/order_by.rs

@@ -8,15 +8,19 @@ pub enum OrderBy {
     Desc(usize),
     AscColumn(String),
     DescColumn(String),
-    AscVectorL2(String, Vector),
+    AscVectorL2Column(String, Vector),
+    AscVectorL2(usize, Vector),
 }
 
 impl OrderBy {
     /// ORDER BY x ASC
     pub fn asc(&self) -> bool {
         matches!(
             self,
-            OrderBy::Asc(_) | OrderBy::AscColumn(_) | OrderBy::AscVectorL2(_, _)
+            OrderBy::Asc(_)
+                | OrderBy::AscColumn(_)
+                | OrderBy::AscVectorL2Column(_, _)
+                | OrderBy::AscVectorL2(_, _)
         )
     }
 
@@ -25,6 +29,7 @@ impl OrderBy {
         match self {
             OrderBy::Asc(column) => Some(*column - 1),
             OrderBy::Desc(column) => Some(*column - 1),
+            OrderBy::AscVectorL2(column, _) => Some(*column - 1),
             _ => None,
         }
     }
@@ -34,15 +39,15 @@ impl OrderBy {
         match self {
             OrderBy::AscColumn(ref name) => Some(name.as_str()),
             OrderBy::DescColumn(ref name) => Some(name.as_str()),
-            OrderBy::AscVectorL2(ref name, _) => Some(name.as_str()),
+            OrderBy::AscVectorL2Column(ref name, _) => Some(name.as_str()),
             _ => None,
         }
     }
 
     /// ORDER BY clause contains a vector.
     pub fn vector(&self) -> Option<&Vector> {
         match self {
-            OrderBy::AscVectorL2(_, vector) => Some(&vector),
+            OrderBy::AscVectorL2Column(_, vector) => Some(&vector),
             _ => None,
         }
     }

pgdog/src/frontend/router/parser/query.rs

@@ -345,7 +345,7 @@ impl QueryParser {
 
                                                 if let Some(vector) = vector {
                                                     if let Some(column) = column {
-                                                        order_by.push(OrderBy::AscVectorL2(
+                                                        order_by.push(OrderBy::AscVectorL2Column(
                                                             column, vector,
                                                         ));
                                                     }

pgdog/src/net/messages/data_types/mod.rs

@@ -45,6 +45,8 @@ pub enum Datum {
     Uuid(Uuid),
     /// NUMERIC, REAL, DOUBLE PRECISION.
     Numeric(Numeric),
+    /// Vector
+    Vector(Vector),
     /// NULL.
     Null,
 }
@@ -63,6 +65,7 @@ impl ToDataRowColumn for Datum {
             TimestampTz(tz) => tz.to_data_row_column(),
             Uuid(uuid) => uuid.to_data_row_column(),
             Numeric(num) => num.to_data_row_column(),
+            Vector(vector) => vector.to_data_row_column(),
             Null => Data::null(),
         }
     }
@@ -104,6 +107,7 @@ impl Datum {
             DataType::Uuid => Ok(Datum::Uuid(Uuid::decode(bytes, encoding)?)),
             DataType::Timestamp => Ok(Datum::Timestamp(Timestamp::decode(bytes, encoding)?)),
             DataType::TimestampTz => Ok(Datum::TimestampTz(TimestampTz::decode(bytes, encoding)?)),
+            DataType::Vector => Ok(Datum::Vector(Vector::decode(bytes, encoding)?)),
             _ => Ok(Datum::Null),
         }
     }
@@ -130,4 +134,5 @@ pub enum DataType {
     Numeric,
     Other(i32),
     Uuid,
+    Vector,
 }

pgdog/src/net/messages/data_types/numeric.rs

@@ -65,7 +65,7 @@ impl Add for Numeric {
 
 impl Ord for Numeric {
     fn cmp(&self, other: &Self) -> std::cmp::Ordering {
-        match self.partial_cmp(other) {
+        match self.data.partial_cmp(&other.data) {
             Some(ordering) => ordering,
             None => {
                 if self.data.is_nan() || other.data.is_nan() {

pgdog/src/net/messages/data_types/vector.rs

@@ -1,6 +1,9 @@
 use crate::{
     frontend::router::sharding::vector::Distance,
-    net::{messages::Format, Error},
+    net::{
+        messages::{Format, ToDataRowColumn},
+        Error,
+    },
 };
 use bytes::Bytes;
 use serde::{
@@ -10,7 +13,7 @@ use serde::{
 };
 use std::{ops::Deref, str::from_utf8};
 
-use super::{FromDataType, Numeric};
+use super::{Datum, FromDataType, Numeric};
 
 #[derive(Debug, Clone, PartialEq, PartialOrd, Ord, Eq, Hash)]
 #[repr(C)]
@@ -51,6 +54,12 @@ impl FromDataType for Vector {
     }
 }
 
+impl ToDataRowColumn for Vector {
+    fn to_data_row_column(&self) -> crate::net::messages::data_row::Data {
+        Bytes::from(self.encode(Format::Text).unwrap()).into()
+    }
+}
+
 impl Vector {
     /// Length of the vector.
     pub fn len(&self) -> usize {
@@ -100,6 +109,23 @@ impl TryFrom<&str> for Vector {
     }
 }
 
+impl Into<Datum> for Vector {
+    fn into(self) -> Datum {
+        Datum::Vector(self)
+    }
+}
+
+impl TryFrom<Datum> for Vector {
+    type Error = Error;
+
+    fn try_from(value: Datum) -> Result<Self, Self::Error> {
+        match value {
+            Datum::Vector(vector) => Ok(vector),
+            _ => Err(Error::UnexpectedPayload),
+        }
+    }
+}
+
 struct VectorVisitor;
 
 impl<'de> Visitor<'de> for VectorVisitor {

pgdog/src/net/messages/row_description.rs

@@ -92,6 +92,7 @@ impl Field {
             1184 => DataType::TimestampTz,
             1186 => DataType::Interval,
             2950 => DataType::Uuid,
+            27231 => DataType::Vector, // TODO: this is actually variable.
             _ => DataType::Other(self.type_oid),
         }
     }