Adding support for Threshold, Limit, and Order Arguments

[sebscholl]

This pull request adds 3 keyword arguments to the nearest_neighbor method. They are:

order

movie = Movie.find_by(name: "Star Wars (1977)")
# Order all results by the neighbor_distance column in descending order
movie.nearest_neighbors(:factors, distance: "inner_product", order: { neighbor_distance: :desc })

limit

movie = Movie.find_by(name: "Star Wars (1977)")
# Limit the results to 3 records
movie.nearest_neighbors(:factors, distance: "inner_product", limit: 3)

threshold

movie = Movie.find_by(name: "Star Wars (1977)")
# Only return records where the neighbor_distance is greater than or equal to 0.9
movie.nearest_neighbors(:factors, distance: "inner_product", threshold: { gte: 0.9 })

Multiple Options

All options can be used at the same time or separately.

movie = Movie.find_by(name: "Star Wars (1977)")

# Only return 5 records where the neighbor_distance is greater than or equal to 0.9 in descending order
movie.nearest_neighbors(
  :factors,
  distance: "inner_product", 
  limit: 5,
  threshold: { gte: 0.9 },
  order: { neighbor_distance: :desc }
)

These options manipulate the SQL statement generated by ActiveRecord. All original test suits are intact and passing, and the new tests were written with the new options.

[ankane]

Hi @sebscholl, thanks for the PR.

1. nearest_neighbors currently returns a relation, so you can limit with limit(n) or first(n).
2. Results are currently ordered by distance. However, if you have a default scope on the model, that'll take precedence.
3. For thresholds, you can use where("(embedding <#> ?) * -1 > ?", vector, 0.9) or filter in memory with select { |v| v.neighbor_distance > 0.9 }. I may add an option for this at some point, but want to think more about the design.

[sebscholl]

Makes sense. Do you believe it would be helpful to add this info to the docs (e.g, where("(embedding <#> ?) * -1 > ?", vector, 0.9)) or prefer to sit tight until you feel you have more clarity on the design? Lmk, and I can make an update if it would help.

[gvkhna]

@ankane I believe in the case of using class method like Movie.nearest_neighbor(embedding, my_gen_embedding, ...) the ordering is not set by distance. Instead I'm getting ORDER BY "text_nodes"."id" ASC LIMIT $1 on these queries. I'm encountering this exact problem and so i'll implement the query manually as a workaround for now.

P.S. I set Movie.unscoped {} but still am getting ORDER BY ID, AFAIK there is no way to set to order by distance with the gem.

P.P.S I set .order(Arel.sql("neighbor_distance DESC")) but it didn't actually apply that to the query instead still ordering by ID.

[vestedpr-dev]

Regarding thresholds, if others are working on this, here is some relevant code I came up with:

def filter_by_within_distance(scope)
        return scope unless @params[:within_distance] && @params[:distance_type]

        distance_type = @params[:distance_type].to_sym

        # Determine the correct operator based on distance type
        operator = case distance_type
                   when :euclidean
                     "<->"
                   when :cosine
                     "<=>"
                   when :inner_product
                     "<#>"
                   else
                     raise ArgumentError, "Unsupported distance type: #{@params[:distance_type]}"
                   end

        condition_pattern = if distance_type == :inner_product
                              # Negative inner product
                              "((#{@params[:search_vector_column]} #{operator} '[?]') * -1) < ?"
                            else
                              "(#{@params[:search_vector_column]} #{operator} '[?]') < ?"
                            end

        scope.where(condition_pattern, @query_vector, @params[:within_distance])
      end

Some nuances to take note of:

• When used within a Rails order clause, I needed to wrap the vector in single quotes and the square bracket to end up with valid SQL: '[?]'
• This is elementary but not obvious at first to those new to pgvector: depending on the distance type you're using for the ordering, you need to use the same comparison operator (<->, <=>, <#>) in your limit statement.

[allcentury]

@vestedpr-dev thanks for writing that out! I wasn't able to get this line to work in rails 8, pg 17:

"((#{@params[:search_vector_column]} #{operator} '[?]') * -1) < ?"

My actual command:

m.nearest_neighbors(:embeddings, distance: "cosine").where("embeddings <=> '[?]' <= ?", m.embeddings, 1 - 0.9).first(3)

error:

exec_params': ERROR:  invalid input syntax for type vector: "[$2, $3, $4]" (PG::InvalidTextRepresentation)
LINE 1: ...nts"."embeddings" IS NOT NULL AND (embeddings <=> '[$2, $3, ...

I tried various ways of casting and finally got this way to work, but this feels very inefficient:

    # Convert embeddings array to a parameterized string safely
    embeddings_array = self.embeddings.map { |v| ActiveRecord::Base.sanitize_sql(v) }.join(', ')

    # Safely bind parameters to avoid SQL injection
    MyModel
      .where.not(id: self.id)
      .where.not(embeddings: nil)
      .where(Arel.sql("embeddings <=> ARRAY[#{embeddings_array}]::vector <= ?"), 1 - threshold)
      .order(Arel.sql("embeddings <=> ARRAY[#{embeddings_array}]::vector"))
      .limit(3)

[Petercopter]

This seems to work for me using Postgres.

Model
  .nearest_neighbors(:embedding, embedding, distance: 'cosine')
  .where(
    "(embedding <=> ARRAY[#{embedding.join(', ')}]::vector) < ?",
    0.2
  )

I wasn't able to get the '[?]' syntax working either.

[jayelkaake]

bump - because this looks great! We should get it reviewed and merged.