Class: Neighbor::Redis::VectorSet

Inherits:
Object
  • Object
show all
Defined in:
lib/neighbor/redis/vector_set.rb

Constant Summary collapse

NO_DEFAULT = 
Object.new

Instance Method Summary collapse

Constructor Details

#initialize(name, m: nil, ef_construction: nil, ef_search: nil, epsilon: nil, quantization: nil, reduce: nil, id_type: "string") ⇒ VectorSet

Returns a new instance of VectorSet.



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
43
44
45
46
47
48
49
50
51
52
 
# File 'lib/neighbor/redis/vector_set.rb', line 6

def initialize(
  name,
  m: nil,
  ef_construction: nil,
  ef_search: nil,
  epsilon: nil,
  quantization: nil,
  reduce: nil,
  id_type: "string"
)
  name = name.to_str
  if name.include?(":")
    raise ArgumentError, "invalid name"
  end

  @name = name
  @m = m&.to_i
  @ef_construction = ef_construction&.to_i
  @ef_search = ef_search&.to_i
  @epsilon = epsilon&.to_f

  @quant_type =
    case quantization&.to_s
    when nil
      "NOQUANT"
    when "binary"
      "BIN"
    when "int8"
      "Q8"
    else
      raise ArgumentError, "invalid quantization"
    end

  case id_type.to_s
  when "string", "integer"
    @int_ids = id_type == "integer"
  else
    raise ArgumentError, "invalid id_type"
  end

  @reduce_args = []
  @reduce_args.push("REDUCE", reduce.to_i) if reduce

  @add_args = []
  @add_args.push("M", @m) if @m
  @add_args.push("EF", @ef_construction) if @ef_construction
end

Instance Method Details

#add(id, vector, metadata: nil) ⇒ Object 



73
74
75
 
# File 'lib/neighbor/redis/vector_set.rb', line 73

def add(id, vector, metadata: nil)
  add_all([id], [vector], metadata: metadata ? [metadata] : nil)[0]
end

#add_all(ids, vectors, metadata: nil) ⇒ Object

Raises:

  • (ArgumentError) 


77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
 
# File 'lib/neighbor/redis/vector_set.rb', line 77

def add_all(ids, vectors, metadata: nil)
  # perform checks first to reduce chance of non-atomic updates
  ids = ids.to_a.map { |v| item_id(v) }
  vectors = vectors.to_a
  metadata = metadata.to_a if metadata

  raise ArgumentError, "different sizes" if ids.size != vectors.size

  if vectors.size > 1
    dimensions = vectors.first.size
    unless vectors.all? { |v| v.size == dimensions }
      raise ArgumentError, "different dimensions"
    end
  end

  if metadata
    raise ArgumentError, "different sizes" if metadata.size != ids.size
  end

  result =
    client.pipelined do |pipeline|
      ids.zip(vectors).each_with_index do |(id, vector), i|
        attributes = metadata[i] if metadata
        attribute_args = []
        attribute_args.push("SETATTR", JSON.generate(attributes)) if attributes
        pipeline.call("VADD", key, *@reduce_args, "FP32", to_binary(vector), id, @quant_type, *attribute_args, *@add_args)
      end
    end
  result.map { |v| bool_result(v) }
end

#count ⇒ Object 



69
70
71
 
# File 'lib/neighbor/redis/vector_set.rb', line 69

def count
  run_command("VCARD", key)
end

#dimensions ⇒ Object 



62
63
64
65
66
67
 
# File 'lib/neighbor/redis/vector_set.rb', line 62

def dimensions
  run_command("VDIM", key)
rescue => e
  raise e unless e.message.include?("key does not exist")
  nil
end

#drop ⇒ Object 



198
199
200
 
# File 'lib/neighbor/redis/vector_set.rb', line 198

def drop
  bool_result(run_command("DEL", key))
end

#exists? ⇒ Boolean

Returns:

  • (Boolean) 


54
55
56
 
# File 'lib/neighbor/redis/vector_set.rb', line 54

def exists?
  !run_command("VINFO", key).nil?
end

#find(id) ⇒ Object 



134
135
136
137
138
 
# File 'lib/neighbor/redis/vector_set.rb', line 134

def find(id)
  id = item_id(id)

  run_command("VEMB", key, id)&.map(&:to_f)
end

#info ⇒ Object 



58
59
60
 
# File 'lib/neighbor/redis/vector_set.rb', line 58

def info
  hash_result(run_command("VINFO", key))
end

#links(id) ⇒ Object 



181
182
183
184
185
186
187
188
189
 
# File 'lib/neighbor/redis/vector_set.rb', line 181

def links(id)
  id = item_id(id)

  run_command("VLINKS", key, id, "WITHSCORES")&.map do |links|
    hash_result(links).map do |k, v|
      search_result(k, v)
    end
  end
end

#member?(id) ⇒ Boolean Also known as: include?

Returns:

  • (Boolean) 


108
109
110
111
112
 
# File 'lib/neighbor/redis/vector_set.rb', line 108

def member?(id)
  id = item_id(id)

  bool_result(run_command("VISMEMBER", key, id))
end

#metadata(id) ⇒ Object 



140
141
142
143
144
145
 
# File 'lib/neighbor/redis/vector_set.rb', line 140

def metadata(id)
  id = item_id(id)

  a = run_command("VGETATTR", key, id)
  a ? JSON.parse(a) : nil
end

#remove(id) ⇒ Object 



115
116
117
118
119
 
# File 'lib/neighbor/redis/vector_set.rb', line 115

def remove(id)
  id = item_id(id)

  bool_result(run_command("VREM", key, id))
end

#remove_all(ids) ⇒ Object 



121
122
123
124
125
126
127
128
129
130
131
132
 
# File 'lib/neighbor/redis/vector_set.rb', line 121

def remove_all(ids)
  # perform checks first to reduce chance of non-atomic updates
  ids = ids.to_a.map { |v| item_id(v) }

  result =
    client.pipelined do |pipeline|
      ids.each do |id|
        pipeline.call("VREM", key, id)
      end
    end
  result.map { |v| bool_result(v) }
end

#remove_metadata(id) ⇒ Object 



153
154
155
156
157
 
# File 'lib/neighbor/redis/vector_set.rb', line 153

def remove_metadata(id)
  id = item_id(id)

  bool_result(run_command("VSETATTR", key, id, ""))
end

#sample(n = NO_DEFAULT) ⇒ Object 



191
192
193
194
195
196
 
# File 'lib/neighbor/redis/vector_set.rb', line 191

def sample(n = NO_DEFAULT)
  count = n == NO_DEFAULT ? 1 : n.to_i

  result = run_command("VRANDMEMBER", key, count).map { |v| item_id(v) }
  n == NO_DEFAULT ? result.first : result
end

#search(vector, count: 5, with_metadata: false, ef_search: nil, exact: false, _filter: nil, _ef_filter: nil) ⇒ Object 



159
160
161
162
163
164
165
 
# File 'lib/neighbor/redis/vector_set.rb', line 159

def search(vector, count: 5, with_metadata: false, ef_search: nil, exact: false, _filter: nil, _ef_filter: nil)
  count = count.to_i

  search_command(["FP32", to_binary(vector)], count:, with_metadata:, ef_search:, exact:, _filter:, _ef_filter:).map do |k, v|
    search_result(k, v, with_metadata:)
  end
end

#search_id(id, count: 5, with_metadata: false, ef_search: nil, exact: false, _filter: nil, _ef_filter: nil) ⇒ Object Also known as: nearest 



167
168
169
170
171
172
173
174
175
176
177
178
 
# File 'lib/neighbor/redis/vector_set.rb', line 167

def search_id(id, count: 5, with_metadata: false, ef_search: nil, exact: false, _filter: nil, _ef_filter: nil)
  id = item_id(id)
  count = count.to_i

  result =
    search_command(["ELE", id], count: count + 1, with_metadata:, ef_search:, exact:, _filter:, _ef_filter:).filter_map do |k, v|
      if k != id.to_s
        search_result(k, v, with_metadata:)
      end
    end
  result.first(count)
end

#set_metadata(id, metadata) ⇒ Object 



147
148
149
150
151
 
# File 'lib/neighbor/redis/vector_set.rb', line 147

def set_metadata(id, metadata)
  id = item_id(id)

  bool_result(run_command("VSETATTR", key, id, JSON.generate(metadata)))
end