Skip to content

Commit bc7cece

Browse files
add e2e tests.
Signed-off-by: my-vegetable-has-exploded <wy1109468038@gmail.com>
1 parent b4b51b9 commit bc7cece

File tree

6 files changed

+200
-1
lines changed

6 files changed

+200
-1
lines changed

src/datatype/veci8.rs

+17-1
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ impl Veci8 {
7272

7373
pub fn new_zeroed_in_postgres(len: usize) -> Veci8Output {
7474
unsafe {
75-
assert!(1 <= len && len <= 65535);
75+
assert!((1..=65535).contains(&len));
7676
let layout = Veci8::layout(len);
7777
let ptr = pgrx::pg_sys::palloc(layout.size()) as *mut Veci8;
7878
ptr.cast::<u8>().add(layout.size() - 8).write_bytes(0, 8);
@@ -555,3 +555,19 @@ fn _vectors_veci8_out(vector: Veci8Input<'_>) -> CString {
555555
buffer.push(']');
556556
CString::new(buffer).unwrap()
557557
}
558+
559+
#[pgrx::pg_extern(immutable, parallel_safe, strict)]
560+
fn _vectors_to_veci8(len: i32, alpha: f32, offset: f32, values: pgrx::Array<i32>) -> Veci8Output {
561+
check_value_dimensions(len as usize);
562+
if (len as usize) != values.len() {
563+
bad_literal("Lengths of values and len are not matched.");
564+
}
565+
if values.contains_nulls() {
566+
bad_literal("Index or value contains nulls.");
567+
}
568+
let values = values
569+
.iter()
570+
.map(|x| I8(x.unwrap() as i8))
571+
.collect::<Vec<_>>();
572+
Veci8::new_in_postgres(values.as_slice(), F32(alpha), F32(offset))
573+
}

src/sql/finalize.sql

+3
Original file line numberDiff line numberDiff line change
@@ -558,4 +558,7 @@ CREATE CAST (veci8 AS vector)
558558
CREATE CAST (vector AS veci8)
559559
WITH FUNCTION _vectors_cast_vecf32_to_veci8(vector, integer, boolean);
560560

561+
CREATE FUNCTION to_veci8("len" INT, "alpha" real, "offset" real, "values" INT[]) RETURNS veci8
562+
IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vectors_veci8_from_array_wrapper';
563+
561564
-- finalize end

tests/sqllogictest/int8.slt

+50
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
statement ok
2+
SET search_path TO pg_temp, vectors;
3+
4+
statement ok
5+
CREATE TABLE t (val veci8(3));
6+
7+
statement ok
8+
INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[]::vector::veci8 FROM generate_series(1, 1000);
9+
10+
statement ok
11+
CREATE INDEX ON t USING vectors (val veci8_l2_ops)
12+
WITH (options = "[indexing.hnsw]");
13+
14+
statement ok
15+
CREATE INDEX ON t USING vectors (val veci8_dot_ops)
16+
WITH (options = "[indexing.hnsw]");
17+
18+
statement ok
19+
CREATE INDEX ON t USING vectors (val veci8_cos_ops)
20+
WITH (options = "[indexing.ivf]");
21+
22+
23+
query I
24+
SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]'::veci8 limit 10) t2;
25+
----
26+
10
27+
28+
query I
29+
SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <=> '[0.5,0.5,0.5]'::veci8 limit 10) t2;
30+
----
31+
10
32+
33+
query I
34+
SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <#> '[0.5,0.5,0.5]'::veci8 limit 10) t2;
35+
----
36+
10
37+
38+
statement ok
39+
DROP TABLE t;
40+
41+
query I
42+
SELECT to_veci8(5, 1, 0, '{0,1,2,0,0}');
43+
----
44+
[0, 1, 2, 0, 0]
45+
46+
Lengths of values and len are not matched.
47+
SELECT to_veci8(5, 1, 0, '{0,1,2,0,0,0}');
48+
49+
Index or value contains nulls.
50+
SELECT to_veci8(5, 1, 0, '{0,1,2,NULL,0}');

tests/sqllogictest/veci8_binary.slt

+37
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
statement ok
2+
SET search_path TO pg_temp, vectors;
3+
4+
statement ok
5+
CREATE TABLE t (id bigserial, val veci8);
6+
7+
statement ok
8+
INSERT INTO t (val) SELECT NULL FROM generate_series(1, 1000);
9+
10+
statement ok
11+
INSERT INTO t (val) SELECT ARRAY[random()]::real[]::vector::veci8 FROM generate_series(1, 1000);
12+
13+
statement ok
14+
INSERT INTO t (val) SELECT ARRAY[random(), random()]::real[]::vector::veci8 FROM generate_series(1, 1000);
15+
16+
statement ok
17+
INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[]::vector::veci8 FROM generate_series(1, 1000);
18+
19+
statement ok
20+
COPY t TO '/tmp/data.bin' WITH (FORMAT binary);
21+
22+
statement ok
23+
CREATE TABLE t2 (id bigserial, val veci8);
24+
25+
statement ok
26+
COPY t2 FROM '/tmp/data.bin' WITH (FORMAT binary);
27+
28+
query I
29+
SELECT SUM(((t.val = t2.val) OR (t.val IS NULL and t2.val IS NULL))::int) FROM t FULL OUTER JOIN t2 ON t.id = t2.id;
30+
----
31+
4000
32+
33+
statement ok
34+
DROP TABLE t;
35+
36+
statement ok
37+
DROP TABLE t2;

tests/sqllogictest/veci8_storage.slt

+20
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
statement ok
2+
SET search_path TO pg_temp, vectors;
3+
4+
statement ok
5+
CREATE TABLE t (val veci8);
6+
7+
statement ok
8+
INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[]::vector::veci8 FROM generate_series(1, 1000);
9+
10+
statement ok
11+
ALTER TABLE t ALTER COLUMN val SET STORAGE PLAIN;
12+
13+
statement ok
14+
ALTER TABLE t ALTER COLUMN val SET STORAGE EXTERNAL;
15+
16+
statement ok
17+
ALTER TABLE t ALTER COLUMN val SET STORAGE EXTENDED;
18+
19+
statement ok
20+
ALTER TABLE t ALTER COLUMN val SET STORAGE MAIN;
+73
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
statement ok
2+
SET search_path TO pg_temp, vectors;
3+
4+
query I
5+
SELECT ('[0, 1, 2, 3, 4, 5, 6, 7]'::veci8)[3:6];
6+
----
7+
[3, 4, 5]
8+
9+
query I
10+
SELECT ('[0, 1, 2, 3, 4, 5, 6, 7]'::veci8)[:4];
11+
----
12+
[0, 1, 2, 3]
13+
14+
query I
15+
SELECT ('[0, 1, 2, 3, 4, 5, 6, 7]'::veci8)[5:];
16+
----
17+
[5, 6, 7]
18+
19+
query I
20+
SELECT ('[0, 1, 2, 3, 4, 5, 6, 7]'::veci8)[1:8];
21+
----
22+
[1, 2, 3, 4, 5, 6, 7]
23+
24+
statement error type veci8 does only support one subscript
25+
SELECT ('[0, 1, 2, 3, 4, 5, 6, 7]'::veci8)[3:3][1:1];
26+
27+
statement error type veci8 does only support slice fetch
28+
SELECT ('[0, 1, 2, 3, 4, 5, 6, 7]'::veci8)[3];
29+
30+
query I
31+
SELECT ('[0, 1, 2, 3, 4, 5, 6, 7]'::veci8)[5:4];
32+
----
33+
NULL
34+
35+
query I
36+
SELECT ('[0, 1, 2, 3, 4, 5, 6, 7]'::veci8)[9:];
37+
----
38+
NULL
39+
40+
query I
41+
SELECT ('[0, 1, 2, 3, 4, 5, 6, 7]'::veci8)[:0];
42+
----
43+
NULL
44+
45+
query I
46+
SELECT ('[0, 1, 2, 3, 4, 5, 6, 7]'::veci8)[:-1];
47+
----
48+
NULL
49+
50+
query I
51+
SELECT ('[0, 1, 2, 3, 4, 5, 6, 7]'::veci8)[NULL:NULL];
52+
----
53+
NULL
54+
55+
query I
56+
SELECT ('[0, 1, 2, 3, 4, 5, 6, 7]'::veci8)[NULL:8];
57+
----
58+
NULL
59+
60+
query I
61+
SELECT ('[0, 1, 2, 3, 4, 5, 6, 7]'::veci8)[1:NULL];
62+
----
63+
NULL
64+
65+
query I
66+
SELECT ('[0, 1, 2, 3, 4, 5, 6, 7]'::veci8)[NULL:];
67+
----
68+
NULL
69+
70+
query I
71+
SELECT ('[0, 1, 2, 3, 4, 5, 6, 7]'::veci8)[:NULL];
72+
----
73+
NULL

0 commit comments

Comments
 (0)