Module: Sequel::Plugins::Pgvector::DatasetMethods

Defined in:
lib/sequel/plugins/pgvector.rb

Instance Method Summary collapse

Instance Method Details

#nearest_neighbors(column, value, distance:, threshold: nil) ⇒ Object

Raises:

  • (ArgumentError) 


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
 
# File 'lib/sequel/plugins/pgvector.rb', line 13

def nearest_neighbors(column, value, distance:, threshold: nil)
  value = ::Pgvector.encode(value) unless value.is_a?(String)
  quoted_column = quote_identifier(column)
  distance = distance.to_s

  operator =
    case distance
    when "inner_product"
      "<#>"
    when "cosine"
      "<=>"
    when "euclidean"
      "<->"
    end

  raise ArgumentError, "Invalid distance: #{distance}" unless operator

  order = "#{quoted_column} #{operator} ?"

  neighbor_distance =
    if distance == "inner_product"
      "(#{order}) * -1"
    else
      order
    end


  query = select_append(Sequel.lit("#{neighbor_distance} AS neighbor_distance", value))
                .exclude(column => nil)
                .order(Sequel.lit(order, value))
    
  # Apply the WHERE condition only if threshold is provided
  query = query.where(Sequel.lit("(#{neighbor_distance}) < ?", value, threshold)) if threshold
      
  query
end