feat: new text representation for sparse vector

src/datatype/text_svecf32.rs

@@ -1,55 +1,49 @@
 use super::memory_svecf32::SVecf32Output;
 use crate::datatype::memory_svecf32::SVecf32Input;
-use crate::datatype::typmod::Typmod;
 use crate::error::*;
 use base::scalar::*;
 use base::vector::*;
-use num_traits::Zero;
 use pgrx::pg_sys::Oid;
 use std::ffi::{CStr, CString};
+use std::fmt::Write;
 
 #[pgrx::pg_extern(immutable, strict, parallel_safe)]
-fn _vectors_svecf32_in(input: &CStr, _oid: Oid, typmod: i32) -> SVecf32Output {
-    use crate::utils::parse::parse_vector;
-    let reserve = Typmod::parse_from_i32(typmod)
-        .unwrap()
-        .dims()
-        .map(|x| x.get())
-        .unwrap_or(0);
-    let v = parse_vector(input.to_bytes(), reserve as usize, |s| {
-        s.parse::<F32>().ok()
-    });
+fn _vectors_svecf32_in(input: &CStr, _oid: Oid, _typmod: i32) -> SVecf32Output {
+    use crate::utils::parse::{parse_pgvector_svector, svector_filter_nonzero};
+    let v = parse_pgvector_svector(input.to_bytes(), |s| s.parse::<F32>().ok());
     match v {
         Err(e) => {
             bad_literal(&e.to_string());
         }
-        Ok(vector) => {
-            check_value_dims_1048575(vector.len());
-            let mut indexes = Vec::<u32>::new();
-            let mut values = Vec::<F32>::new();
-            for (i, &x) in vector.iter().enumerate() {
-                if !x.is_zero() {
-                    indexes.push(i as u32);
-                    values.push(x);
-                }
-            }
-            SVecf32Output::new(SVecf32Borrowed::new(vector.len() as u32, &indexes, &values))
+        Ok((indexes, values, dims)) => {
+            check_value_dims_1048575(dims);
+            check_index_in_bound(&indexes, dims);
+            let (non_zero_indexes, non_zero_values) = svector_filter_nonzero(&indexes, &values);
+            SVecf32Output::new(SVecf32Borrowed::new(
+                dims as u32,
+                &non_zero_indexes,
+                &non_zero_values,
+            ))
         }
     }
 }
 
 #[pgrx::pg_extern(immutable, strict, parallel_safe)]
 fn _vectors_svecf32_out(vector: SVecf32Input<'_>) -> CString {
+    let dims = vector.for_borrow().dims();
     let mut buffer = String::new();
-    buffer.push('[');
-    let vec = vector.for_borrow().to_vec();
-    let mut iter = vec.iter();
-    if let Some(x) = iter.next() {
-        buffer.push_str(format!("{}", x).as_str());
-    }
-    for x in iter {
-        buffer.push_str(format!(", {}", x).as_str());
+    buffer.push('{');
+    let svec = vector.for_borrow();
+    let mut need_splitter = false;
+    for (&index, &value) in svec.indexes().iter().zip(svec.values().iter()) {
+        match need_splitter {
+            false => {
+                write!(buffer, "{}:{}", index, value).unwrap();
+                need_splitter = true;
+            }
+            true => write!(buffer, ", {}:{}", index, value).unwrap(),
+        }
     }
-    buffer.push(']');
+    write!(buffer, "}}/{}", dims).unwrap();
     CString::new(buffer).unwrap()
 }

src/error.rs

@@ -69,6 +69,20 @@ ADVICE: Check if dimensions of the vector are among 1 and 1_048_575."
     NonZeroU32::new(dims as u32).unwrap()
 }
 
+pub fn check_index_in_bound(indexes: &[u32], dims: usize) -> NonZeroU32 {
+    let mut last: u32 = 0;
+    for (i, index) in indexes.iter().enumerate() {
+        if i > 0 && last == *index {
+            error!("Indexes need to be unique, but there are more than one same index {index}")
+        }
+        if *index >= dims as u32 {
+            error!("Index out of bounds: the dim is {dims} but the index is {index}");
+        }
+        last = *index;
+    }
+    NonZeroU32::new(dims as u32).unwrap()
+}
+
 pub fn bad_literal(hint: &str) -> ! {
     error!(
         "\

src/utils/parse.rs

@@ -1,6 +1,7 @@
+use num_traits::Zero;
 use thiserror::Error;
 
-#[derive(Debug, Error)]
+#[derive(Debug, Error, PartialEq)]
 pub enum ParseVectorError {
     #[error("The input string is empty.")]
     EmptyString {},
@@ -83,3 +84,286 @@ where
     }
     Ok(vector)
 }
+
+#[derive(PartialEq, Debug)]
+enum ParseState {
+    Start,
+    LeftBracket,
+    Index,
+    Value,
+    Splitter,
+    Comma,
+    Length,
+}
+
+#[inline(always)]
+pub fn svector_filter_nonzero<T: Zero + Clone + PartialEq>(
+    indexes: &[u32],
+    values: &[T],
+) -> (Vec<u32>, Vec<T>) {
+    let non_zero_indexes: Vec<u32> = indexes
+        .iter()
+        .enumerate()
+        .filter(|(i, _)| values.get(*i).unwrap() != &T::zero())
+        .map(|(_, x)| *x)
+        .collect();
+    let non_zero_values: Vec<T> = indexes
+        .iter()
+        .enumerate()
+        .filter(|(i, _)| values.get(*i).unwrap() != &T::zero())
+        .map(|(i, _)| values.get(i).unwrap().clone())
+        .collect();
+    (non_zero_indexes, non_zero_values)
+}
+
+#[inline(always)]
+pub fn parse_pgvector_svector<T: Zero + Clone, F>(
+    input: &[u8],
+    f: F,
+) -> Result<(Vec<u32>, Vec<T>, usize), ParseVectorError>
+where
+    F: Fn(&str) -> Option<T>,
+{
+    use arrayvec::ArrayVec;
+    if input.is_empty() {
+        return Err(ParseVectorError::EmptyString {});
+    }
+    let mut token: ArrayVec<u8, 48> = ArrayVec::new();
+    let mut indexes = Vec::<u32>::new();
+    let mut values = Vec::<T>::new();
+
+    let mut state = ParseState::Start;
+    for (position, char) in input.iter().enumerate() {
+        let c = *char;
+        match (&state, c) {
+            (_, b' ') => {}
+            (ParseState::Start, b'{') => {
+                state = ParseState::LeftBracket;
+            }
+            (
+                ParseState::LeftBracket | ParseState::Index | ParseState::Comma,
+                b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'.' | b'+' | b'-',
+            ) => {
+                if token.is_empty() {
+                    token.push(b'$');
+                }
+                if token.try_push(c).is_err() {
+                    return Err(ParseVectorError::TooLongNumber { position });
+                }
+                state = ParseState::Index;
+            }
+            (ParseState::LeftBracket | ParseState::Comma, b'}') => {
+                state = ParseState::Splitter;
+            }
+            (ParseState::Index, b':') => {
+                if token.is_empty() {
+                    return Err(ParseVectorError::TooShortNumber { position });
+                }
+                let s = unsafe { std::str::from_utf8_unchecked(&token[1..]) };
+                let index = s
+                    .parse::<u32>()
+                    .map_err(|_| ParseVectorError::BadParsing { position })?;
+                indexes.push(index);
+                token.clear();
+                state = ParseState::Value;
+            }
+            (ParseState::Value, b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'.' | b'+' | b'-') => {
+                if token.is_empty() {
+                    token.push(b'$');
+                }
+                if token.try_push(c).is_err() {
+                    return Err(ParseVectorError::TooLongNumber { position });
+                }
+            }
+            (ParseState::Value, b',') => {
+                if token.is_empty() {
+                    return Err(ParseVectorError::TooShortNumber { position });
+                }
+                let s = unsafe { std::str::from_utf8_unchecked(&token[1..]) };
+                let num = f(s).ok_or(ParseVectorError::BadParsing { position })?;
+                values.push(num);
+                token.clear();
+                state = ParseState::Comma;
+            }
+            (ParseState::Value, b'}') => {
+                if token.is_empty() {
+                    return Err(ParseVectorError::TooShortNumber { position });
+                }
+                let s = unsafe { std::str::from_utf8_unchecked(&token[1..]) };
+                let num = f(s).ok_or(ParseVectorError::BadParsing { position })?;
+                values.push(num);
+                token.clear();
+                state = ParseState::Splitter;
+            }
+            (ParseState::Splitter, b'/') => {
+                state = ParseState::Length;
+            }
+            (ParseState::Length, b'0'..=b'9') => {
+                if token.is_empty() {
+                    token.push(b'$');
+                }
+                if token.try_push(c).is_err() {
+                    return Err(ParseVectorError::TooLongNumber { position });
+                }
+            }
+            (_, _) => {
+                return Err(ParseVectorError::BadCharacter { position });
+            }
+        }
+    }
+    if state != ParseState::Length {
+        return Err(ParseVectorError::BadParsing {
+            position: input.len(),
+        });
+    }
+    if token.is_empty() {
+        return Err(ParseVectorError::TooShortNumber {
+            position: input.len(),
+        });
+    }
+    let s = unsafe { std::str::from_utf8_unchecked(&token[1..]) };
+    let dims = s
+        .parse::<usize>()
+        .map_err(|_| ParseVectorError::BadParsing {
+            position: input.len(),
+        })?;
+
+    let mut indices = (0..indexes.len()).collect::<Vec<_>>();
+    indices.sort_by_key(|&i| &indexes[i]);
+    let sorted_values: Vec<T> = indices
+        .iter()
+        .map(|i| values.get(*i).unwrap().clone())
+        .collect();
+    indexes.sort();
+
+    Ok((indexes, sorted_values, dims))
+}
+
+#[cfg(test)]
+mod tests {
+    use base::scalar::F32;
+
+    use super::*;
+
+    #[test]
+    fn test_svector_parse_accept() {
+        let exprs: Vec<(&str, (Vec<u32>, Vec<F32>, usize))> = vec![
+            ("{}/1", (vec![], vec![], 1)),
+            ("{0:1}/1", (vec![0], vec![F32(1.0)], 1)),
+            (
+                "{0:1, 1:-2, }/2",
+                (vec![0, 1], vec![F32(1.0), F32(-2.0)], 2),
+            ),
+            ("{0:1, 1:1.5}/2", (vec![0, 1], vec![F32(1.0), F32(1.5)], 2)),
+            (
+                "{0:+3, 2:-4.1}/3",
+                (vec![0, 2], vec![F32(3.0), F32(-4.1)], 3),
+            ),
+            (
+                "{0:0, 1:0, 2:0}/3",
+                (vec![0, 1, 2], vec![F32(0.0), F32(0.0), F32(0.0)], 3),
+            ),
+            (
+                "{3:3, 2:2, 1:1, 0:0}/4",
+                (
+                    vec![0, 1, 2, 3],
+                    vec![F32(0.0), F32(1.0), F32(2.0), F32(3.0)],
+                    4,
+                ),
+            ),
+        ];
+        for (e, parsed) in exprs {
+            let ret = parse_pgvector_svector(e.as_bytes(), |s| s.parse::<F32>().ok());
+            assert!(ret.is_ok(), "at expr {:?}: {:?}", e, ret);
+            assert_eq!(ret.unwrap(), parsed, "parsed at expr {:?}", e);
+        }
+    }
+
+    #[test]
+    fn test_svector_parse_reject() {
+        let exprs: Vec<(&str, ParseVectorError)> = vec![
+            ("{", ParseVectorError::BadParsing { position: 1 }),
+            ("}", ParseVectorError::BadCharacter { position: 0 }),
+            ("{:", ParseVectorError::BadCharacter { position: 1 }),
+            (":}", ParseVectorError::BadCharacter { position: 0 }),
+            (
+                "{0:1, 1:2, 2:3}",
+                ParseVectorError::BadParsing { position: 15 },
+            ),
+            (
+                "{0:1, 1:2, 2:3",
+                ParseVectorError::BadParsing { position: 14 },
+            ),
+            (
+                "{0:1, 1:2}/",
+                ParseVectorError::TooShortNumber { position: 11 },
+            ),
+            ("{0}/5", ParseVectorError::BadCharacter { position: 2 }),
+            ("{0:}/5", ParseVectorError::TooShortNumber { position: 3 }),
+            ("{:0}/5", ParseVectorError::BadCharacter { position: 1 }),
+            (
+                "{0:, 1:2}/5",
+                ParseVectorError::TooShortNumber { position: 3 },
+            ),
+            ("{0:1, 1}/5", ParseVectorError::BadCharacter { position: 7 }),
+            ("/2", ParseVectorError::BadCharacter { position: 0 }),
+            ("{}/1/2", ParseVectorError::BadCharacter { position: 4 }),
+            (
+                "{0:1, 1:2}/4/2",
+                ParseVectorError::BadCharacter { position: 12 },
+            ),
+            ("{}/-4", ParseVectorError::BadCharacter { position: 3 }),
+            (
+                "{1,2,3,4}/5",
+                ParseVectorError::BadCharacter { position: 2 },
+            ),
+        ];
+        for (e, err) in exprs {
+            let ret = parse_pgvector_svector(e.as_bytes(), |s| s.parse::<F32>().ok());
+            assert!(ret.is_err(), "at expr {:?}: {:?}", e, ret);
+            assert_eq!(ret.unwrap_err(), err, "parsed at expr {:?}", e);
+        }
+    }
+
+    #[test]
+    fn test_svector_parse_filter() {
+        let exprs: Vec<(&str, (Vec<u32>, Vec<F32>, usize), (Vec<u32>, Vec<F32>))> = vec![
+            ("{}/0", (vec![], vec![], 0), (vec![], vec![])),
+            ("{}/1919810", (vec![], vec![], 1919810), (vec![], vec![])),
+            (
+                "{0:1, 0:2}/1",
+                (vec![0, 0], vec![F32(1.0), F32(2.0)], 1),
+                (vec![0, 0], vec![F32(1.0), F32(2.0)]),
+            ),
+            (
+                "{0:1, 1:1.5}/1",
+                (vec![0, 1], vec![F32(1.0), F32(1.5)], 1),
+                (vec![0, 1], vec![F32(1.0), F32(1.5)]),
+            ),
+            (
+                "{0:0, 1:0, 2:0}/2",
+                (vec![0, 1, 2], vec![F32(0.0), F32(0.0), F32(0.0)], 2),
+                (vec![], vec![]),
+            ),
+            (
+                "{2:0, 1:0}/2",
+                (vec![1, 2], vec![F32(0.0), F32(0.0)], 2),
+                (vec![], vec![]),
+            ),
+            (
+                "{2:0, 1:0, }/2",
+                (vec![1, 2], vec![F32(0.0), F32(0.0)], 2),
+                (vec![], vec![]),
+            ),
+        ];
+        for (e, parsed, filtered) in exprs {
+            let ret = parse_pgvector_svector(e.as_bytes(), |s| s.parse::<F32>().ok());
+            assert!(ret.is_ok(), "at expr {:?}: {:?}", e, ret);
+            assert_eq!(ret.unwrap(), parsed, "parsed at expr {:?}", e);
+
+            let (indexes, values, _) = parsed;
+            let nonzero = svector_filter_nonzero(&indexes, &values);
+            assert_eq!(nonzero, filtered, "filtered at expr {:?}", e);
+        }
+    }
+}