module Sequel::Plugins::Pgvector::DatasetMethods

def nearest_neighbors(column, value, distance:)
  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"
      "<->"
    when "taxicab"
      "<+>"
    when "hamming"
      "<~>"
    when "jaccard"
      "<%>"
    end
  raise ArgumentError, "Invalid distance: #{distance}" unless operator
  order = "#{quoted_column} #{operator} ?"
  neighbor_distance =
    if distance == "inner_product"
      "(#{order}) * -1"
    else
      order
    end
  select_append(Sequel.lit("#{neighbor_distance} AS neighbor_distance", value))
    .exclude(column => nil)
    .order(Sequel.lit(order, value))
end