refactor: use sparse struct to parse

cutecutecat · cutecutecat · commit 5649ac8cf8a1 · 2024-06-17T16:13:24.000+08:00
Signed-off-by: cutecutecat &lt;junyuchen@tensorchord.ai&gt;
diff --git a/src/datatype/text_svecf32.rs b/src/datatype/text_svecf32.rs
@@ -1,39 +1,22 @@
 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};
 
 #[pgrx::pg_extern(immutable, strict, parallel_safe)]
-fn _vectors_svecf32_in(input: &CStr, _oid: Oid, typmod: i32) -> SVecf32Output {
+fn _vectors_svecf32_in(input: &CStr, _oid: Oid, _typmod: i32) -> SVecf32Output {
     use crate::utils::parse::parse_pgvector_svector;
-    let reserve = Typmod::parse_from_i32(typmod)
-        .unwrap()
-        .dims()
-        .map(|x| x.get())
-        .unwrap_or(0);
-    let v = parse_pgvector_svector(input.to_bytes(), reserve as usize, |s| {
-        s.parse::<F32>().ok()
-    });
+    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);
+            SVecf32Output::new(SVecf32Borrowed::new(dims as u32, &indexes, &values))
         }
     }
 }
diff --git a/src/utils/parse.rs b/src/utils/parse.rs
@@ -88,16 +88,16 @@ where
 #[inline(always)]
 pub fn parse_pgvector_svector<T: Zero + Clone, F>(
     input: &[u8],
-    reserve: usize,
     f: F,
-) -> Result<Vec<T>, ParseVectorError>
+) -> 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 dims: usize = 0;
     let left = 'a: {
         for position in 0..input.len() - 1 {
             match input[position] {
@@ -109,7 +109,6 @@ where
         return Err(ParseVectorError::BadParentheses { character: '{' });
     };
     let mut token: ArrayVec<u8, 48> = ArrayVec::new();
-    let mut capacity = reserve;
     let right = 'a: {
         for position in (1..input.len()).rev() {
             match input[position] {
@@ -121,7 +120,7 @@ where
                 b'/' => {
                     token.reverse();
                     let s = unsafe { std::str::from_utf8_unchecked(&token[..]) };
-                    capacity = s
+                    dims = s
                         .parse::<usize>()
                         .map_err(|_| ParseVectorError::BadParsing { position })?;
                 }
@@ -135,8 +134,9 @@ where
         }
         return Err(ParseVectorError::BadParentheses { character: '}' });
     };
-    let mut vector = vec![T::zero(); capacity];
-    let mut index: usize = 0;
+    let mut indexes = Vec::<u32>::new();
+    let mut values = Vec::<T>::new();
+    let mut index: u32 = 0;
     for position in left + 1..right {
         let c = input[position];
         match c {
@@ -153,7 +153,8 @@ where
                     // Safety: all bytes in `token` are ascii characters
                     let s = unsafe { std::str::from_utf8_unchecked(&token[1..]) };
                     let num = f(s).ok_or(ParseVectorError::BadParsing { position })?;
-                    vector[index] = num;
+                    indexes.push(index);
+                    values.push(num);
                     token.clear();
                 } else {
                     return Err(ParseVectorError::TooShortNumber { position });
@@ -164,7 +165,7 @@ where
                     // Safety: all bytes in `token` are ascii characters
                     let s = unsafe { std::str::from_utf8_unchecked(&token[1..]) };
                     index = s
-                        .parse::<usize>()
+                        .parse::<u32>()
                         .map_err(|_| ParseVectorError::BadParsing { position })?;
                     token.clear();
                 } else {
@@ -180,8 +181,9 @@ where
         // Safety: all bytes in `token` are ascii characters
         let s = unsafe { std::str::from_utf8_unchecked(&token[1..]) };
         let num = f(s).ok_or(ParseVectorError::BadParsing { position })?;
-        vector[index] = num;
+        indexes.push(index);
+        values.push(num);
         token.clear();
     }
-    Ok(vector)
+    Ok((indexes, values, dims))
 }
diff --git a/tests/sqllogictest/sparse.slt b/tests/sqllogictest/sparse.slt
@@ -40,12 +40,12 @@ DROP TABLE t;
 query I
 SELECT to_svector(5, '{1,2}', '{1,2}');
 ----
-{2:1, 3:2}/5
+{1:1, 2:2}/5
 
 query I
 SELECT to_svector(5, '{1,2}', '{1,1}') * to_svector(5, '{1,3}', '{2,2}');
 ----
-[0, 2, 0, 0, 0]
+{1:2}/5
 
 statement error Lengths of index and value are not matched.
 SELECT to_svector(5, '{1,2,3}', '{1,2}');
diff --git a/tests/sqllogictest/svector.slt b/tests/sqllogictest/svector.slt
@@ -6,7 +6,7 @@ CREATE TABLE t (id bigserial, val svector);
 
 statement ok
 INSERT INTO t (val)
-VALUES ('[1,2,3]'), ('[4,5,6]');
+VALUES ('{0:1, 1:2, 2:3}/3'), ('{0:4, 1:5, 2:6}/3');
 
 query I
 SELECT vector_dims(val) FROM t;
@@ -23,12 +23,12 @@ SELECT round(vector_norm(val)::numeric, 5) FROM t;
 query ?
 SELECT avg(val) FROM t;
 ----
-[2.5, 3.5, 4.5]
+{0:2.5, 1:3.5, 2:4.5}/3
 
 query ?
 SELECT sum(val) FROM t;
 ----
-[5, 7, 9]
+{0:5, 1:7, 2:9}/3
 
 statement ok
 CREATE TABLE test_vectors (id serial, data vector(1000));
@@ -46,35 +46,35 @@ SELECT count(*) FROM test_vectors;
 5000
 
 query R
-SELECT vector_norm('[3,4]'::svector);
+SELECT vector_norm('{0:3, 1:4}/2'::svector);
 ----
 5
 
 query I
-SELECT vector_dims(v) FROM unnest(ARRAY['[1,2]'::svector, '[3]']) v;
+SELECT vector_dims(v) FROM unnest(ARRAY['{0:1, 1:2}/2'::svector, '{0:3}/1'::svector]) v;
 ----
 2
 1
 
 query ?
-SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::svector, '[3,5,7]']) v;
+SELECT avg(v) FROM unnest(ARRAY['{0:1, 1:2, 2:3}/3'::svector, '{0:3, 1:5, 2:7}/3'::svector]) v;
 ----
-[2, 3.5, 5]
+{0:2, 1:3.5, 2:5}/3
 
 query ?
-SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::svector, '[-1,2,-3]']) v;
+SELECT avg(v) FROM unnest(ARRAY['{0:1, 1:2, 2:3}/3'::svector, '{0:-1, 1:2, 2:-3}/3'::svector]) v;
 ----
-[0, 2, 0]
+{1:2}/3
 
 query ?
-SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::svector, '[3,5,7]', NULL]) v;
+SELECT avg(v) FROM unnest(ARRAY['{0:1, 1:2, 2:3}/3'::svector, '{0:3, 1:5, 2:7}/3'::svector, NULL]) v;
 ----
-[2, 3.5, 5]
+{0:2, 1:3.5, 2:5}/3
 
 query ?
-SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::svector,NULL]) v;
+SELECT avg(v) FROM unnest(ARRAY['{0:1, 1:2, 2:3}/3'::svector,NULL]) v;
 ----
-[1, 2, 3]
+{0:1, 1:2, 2:3}/3
 
 query ?
 SELECT avg(v) FROM unnest(ARRAY[]::svector[]) v;
@@ -87,22 +87,22 @@ SELECT avg(v) FROM unnest(ARRAY[NULL]::svector[]) v;
 NULL
 
 query ?
-SELECT avg(v) FROM unnest(ARRAY['[3e38]'::svector, '[3e38]']) v;
+SELECT avg(v) FROM unnest(ARRAY['{0:3e38}/1'::svector, '{0:3e38}/1'::svector]) v;
 ----
-[inf]
+{0:inf}/1
 
 statement error differs in dimensions
-SELECT avg(v) FROM unnest(ARRAY['[1,2]'::svector, '[3]']) v;
+SELECT avg(v) FROM unnest(ARRAY['{0:1, 1:2}/2'::svector, '{0:3}/1'::svector]) v;
 
 query ?
 SELECT avg(v) FROM unnest(ARRAY[to_svector(5, '{0,1}', '{2,3}'), to_svector(5, '{0,2}', '{1,3}'), to_svector(5, '{3,4}', '{3,3}')]) v;
 ----
-[1, 1, 1, 1, 1]
+{0:1, 1:1, 2:1, 3:1, 4:1}/5
 
 query ?
 SELECT avg(v) FROM unnest(ARRAY[to_svector(32, '{0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15}', '{1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1}'), to_svector(32, '{1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16}', '{1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1}'), to_svector(32, '{2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17}', '{1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1}') ]) v;
 ----
-[0.33333334, 0.6666667, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0.6666667, 0.33333334, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
+{0:0.33333334, 1:0.6666667, 2:1, 3:1, 4:1, 5:1, 6:1, 7:1, 8:1, 9:1, 10:1, 11:1, 12:1, 13:1, 14:1, 15:1, 16:0.6666667, 17:0.33333334}/32
 
 # test avg(svector) get the same result as avg(vector)
 query ?
@@ -111,20 +111,20 @@ SELECT avg(data) = avg(data::svector)::vector FROM test_vectors;
 t
 
 query ?
-SELECT sum(v) FROM unnest(ARRAY['[1,2,3]'::svector, '[3,5,7]']) v;
+SELECT sum(v) FROM unnest(ARRAY['{0:1, 1:2, 2:3}/3'::svector, '{0:3, 1:5, 2:7}/3'::svector]) v;
 ----
-[4, 7, 10]
+{0:4, 1:7, 2:10}/3
 
 # test zero element
 query ?
-SELECT sum(v) FROM unnest(ARRAY['[1,2,3]'::svector, '[-1,2,-3]']) v;
+SELECT sum(v) FROM unnest(ARRAY['{0:1, 1:2, 2:3}/3'::svector, '{0:-1, 1:2, 2:-3}/3'::svector]) v;
 ----
-[0, 4, 0]
+{1:4}/3
 
 query ?
-SELECT sum(v) FROM unnest(ARRAY['[1,2,3]'::svector, '[3,5,7]', NULL]) v;
+SELECT sum(v) FROM unnest(ARRAY['{0:1, 1:2, 2:3}/3'::svector, '{0:3, 1:5, 2:7}/3'::svector, NULL]) v;
 ----
-[4, 7, 10]
+{0:4, 1:7, 2:10}/3
 
 query ?
 SELECT sum(v) FROM unnest(ARRAY[]::svector[]) v;
@@ -137,23 +137,23 @@ SELECT sum(v) FROM unnest(ARRAY[NULL]::svector[]) v;
 NULL
 
 statement error differs in dimensions
-SELECT sum(v) FROM unnest(ARRAY['[1,2]'::svector, '[3]']) v;
+SELECT sum(v) FROM unnest(ARRAY['{0:1, 1:2}/2'::svector, '{0:3}/1'::svector]) v;
 
 # should this return an error ? 
 query ?
-SELECT sum(v) FROM unnest(ARRAY['[3e38]'::svector, '[3e38]']) v;
+SELECT sum(v) FROM unnest(ARRAY['{0:3e38}/1'::svector, '{0:3e38}/1'::svector]) v;
 ----
-[inf]
+{0:inf}/1
 
 query ?
 SELECT sum(v) FROM unnest(ARRAY[to_svector(5, '{0,1}', '{1,2}'), to_svector(5, '{0,2}', '{1,2}'), to_svector(5, '{3,4}', '{3,3}')]) v;
 ----
-[2, 2, 2, 3, 3]
+{0:2, 1:2, 2:2, 3:3, 4:3}/5
 
 query ?
 SELECT sum(v) FROM unnest(ARRAY[to_svector(32, '{0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15}', '{1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1}'), to_svector(32, '{1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16}', '{1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1}'), to_svector(32, '{2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17}', '{1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1}') ]) v;
 ----
-[1, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 2, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
+{0:1, 1:2, 2:3, 3:3, 4:3, 5:3, 6:3, 7:3, 8:3, 9:3, 10:3, 11:3, 12:3, 13:3, 14:3, 15:3, 16:2, 17:1}/32
 
 # test sum(svector) get the same result as sum(vector)
 query ?
diff --git a/tests/sqllogictest/svector_subscript.slt b/tests/sqllogictest/svector_subscript.slt
@@ -2,83 +2,83 @@ statement ok
 SET search_path TO pg_temp, vectors;
 
 query I
-SELECT ('{2:1, 3:2, 4:3, 5:4, 6:5, 7:6, 8:7}/8'::svector)[3:6];
+SELECT ('{1:1, 2:2, 3:3, 4:4, 5:5, 6:6, 7:7}/8'::svector)[3:6];
 ----
-{1:3, 2:4, 3:5}/3
+{0:3, 1:4, 2:5}/3
 
 query I
-SELECT ('{2:1, 3:2, 4:3, 5:4, 6:5, 7:6, 8:7}/8'::svector)[:4];
+SELECT ('{1:1, 2:2, 3:3, 4:4, 5:5, 6:6, 7:7}/8'::svector)[:4];
 ----
-{2:1, 3:2, 4:3}/4
+{1:1, 2:2, 3:3}/4
 
 query I
-SELECT ('{2:1, 3:2, 4:3, 5:4, 6:5, 7:6, 8:7}/8'::svector)[5:];
+SELECT ('{1:1, 2:2, 3:3, 4:4, 5:5, 6:6, 7:7}/8'::svector)[5:];
 ----
-{1:5, 2:6, 3:7}/3
+{0:5, 1:6, 2:7}/3
 
 query I
-SELECT ('{2:1, 3:2, 4:3, 5:4, 6:5, 7:6, 8:7}/8'::svector)[1:8];
+SELECT ('{1:1, 2:2, 3:3, 4:4, 5:5, 6:6, 7:7}/8'::svector)[1:8];
 ----
-{1:1, 2:2, 3:3, 4:4, 5:5, 6:6, 7:7}/7
+{0:1, 1:2, 2:3, 3:4, 4:5, 5:6, 6:7}/7
 
 statement error type svector does only support one subscript
-SELECT ('{2:1, 3:2, 4:3, 5:4, 6:5, 7:6, 8:7}/8'::svector)[3:3][1:1];
+SELECT ('{1:1, 2:2, 3:3, 4:4, 5:5, 6:6, 7:7}/8'::svector)[3:3][1:1];
 
 statement error type svector does only support slice fetch
-SELECT ('{2:1, 3:2, 4:3, 5:4, 6:5, 7:6, 8:7}/8'::svector)[3];
+SELECT ('{1:1, 2:2, 3:3, 4:4, 5:5, 6:6, 7:7}/8'::svector)[3];
 
 query I
-SELECT ('{2:1, 3:2, 4:3, 5:4, 6:5, 7:6, 8:7}/8'::svector)[5:4];
+SELECT ('{1:1, 2:2, 3:3, 4:4, 5:5, 6:6, 7:7}/8'::svector)[5:4];
 ----
 NULL
 
 query I
-SELECT ('{2:1, 3:2, 4:3, 5:4, 6:5, 7:6, 8:7}/8'::svector)[9:];
+SELECT ('{1:1, 2:2, 3:3, 4:4, 5:5, 6:6, 7:7}/8'::svector)[9:];
 ----
 NULL
 
 query I
-SELECT ('{2:1, 3:2, 4:3, 5:4, 6:5, 7:6, 8:7}/8'::svector)[:0];
+SELECT ('{1:1, 2:2, 3:3, 4:4, 5:5, 6:6, 7:7}/8'::svector)[:0];
 ----
 NULL
 
 query I
-SELECT ('{2:1, 3:2, 4:3, 5:4, 6:5, 7:6, 8:7}/8'::svector)[:-1];
+SELECT ('{1:1, 2:2, 3:3, 4:4, 5:5, 6:6, 7:7}/8'::svector)[:-1];
 ----
 NULL
 
 query I
-SELECT ('{2:1, 3:2, 4:3, 5:4, 6:5, 7:6, 8:7}/8'::svector)[NULL:NULL];
+SELECT ('{1:1, 2:2, 3:3, 4:4, 5:5, 6:6, 7:7}/8'::svector)[NULL:NULL];
 ----
 NULL
 
 query I
-SELECT ('{2:1, 3:2, 4:3, 5:4, 6:5, 7:6, 8:7}/8'::svector)[NULL:8];
+SELECT ('{1:1, 2:2, 3:3, 4:4, 5:5, 6:6, 7:7}/8'::svector)[NULL:8];
 ----
 NULL
 
 query I
-SELECT ('{2:1, 3:2, 4:3, 5:4, 6:5, 7:6, 8:7}/8'::svector)[1:NULL];
+SELECT ('{1:1, 2:2, 3:3, 4:4, 5:5, 6:6, 7:7}/8'::svector)[1:NULL];
 ----
 NULL
 
 query I
-SELECT ('{2:1, 3:2, 4:3, 5:4, 6:5, 7:6, 8:7}/8'::svector)[NULL:];
+SELECT ('{1:1, 2:2, 3:3, 4:4, 5:5, 6:6, 7:7}/8'::svector)[NULL:];
 ----
 NULL
 
 query I
-SELECT ('{2:1, 3:2, 4:3, 5:4, 6:5, 7:6, 8:7}/8'::svector)[:NULL];
+SELECT ('{1:1, 2:2, 3:3, 4:4, 5:5, 6:6, 7:7}/8'::svector)[:NULL];
 ----
 NULL
 
 query I
-SELECT ('{3:2, 5:4, 8:7}/8'::svector)[3:7];
+SELECT ('{2:2, 4:4, 7:7}/8'::svector)[3:7];
 ----
-{2:4}/4
+{1:4}/4
 
 query I
-SELECT ('{3:2, 5:4, 8:7}/8'::svector)[5:7];
+SELECT ('{2:2, 4:4, 7:7}/8'::svector)[5:7];
 ----
 {}/2
 
