FilesExpand file tree

master

/

example.rb

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

42 lines (34 loc) · 1.18 KB

master

/

example.rb

Copy path

Top

File metadata and controls

• Code
• Blame

42 lines (34 loc) · 1.18 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

require "numo/narray/alt"

require "pg"

require "pgvector"

# generate random data

rows = 1000000

dimensions = 128

embeddings = Numo::SFloat.new(rows, dimensions).rand

# enable extension

conn = PG.connect(dbname: "pgvector_example")

conn.exec("CREATE EXTENSION IF NOT EXISTS vector")

# create table

conn.exec("DROP TABLE IF EXISTS items")

conn.exec("CREATE TABLE items (id bigserial, embedding vector(#{dimensions}))")

# load data

puts "Loading #{embeddings.shape[0]} rows"

coder = PG::BinaryEncoder::CopyRow.new

conn.copy_data("COPY items (embedding) FROM STDIN WITH (FORMAT BINARY)", coder) do

embeddings.each_over_axis(0).with_index do |embedding, i|

conn.put_copy_data([Pgvector::Vector.new(embedding).to_binary])

# show progress

putc "." if i % 10000 == 0

end

end

puts "\nSuccess!"

# create any indexes *after* loading initial data (skipping for this example)

create_index = false

if create_index

puts "Creating index"

conn.exec("SET maintenance_work_mem = '8GB'")

conn.exec("SET max_parallel_maintenance_workers = 7")

conn.exec("CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops)")

end

# update planner statistics for good measure

conn.exec("ANALYZE items")