MONGOID-5918 Support BSON::Vector field type (#6162)

Declaring a field with type: BSON::Vector now stores the value as a BSON
binary of the vector subtype and reads it back as a BSON::Vector,
preserving dtype and padding.

The extension is guarded so Mongoid still loads on bson-ruby < 5.1, which
predates BSON::Vector. Querying by an exact vector value is not supported
because BSON::Vector subclasses Array; use $vectorSearch for similarity
queries.

Author: comandeo-mongo

lib/mongoid/extensions.rb

@@ -26,3 +26,4 @@
 require 'mongoid/extensions/time'
 require 'mongoid/extensions/time_with_zone'
 require 'mongoid/extensions/true_class'
+require 'mongoid/extensions/vector'

lib/mongoid/extensions/vector.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+module Mongoid
+  module Extensions
+    # Adds type-casting behavior to BSON::Vector class so that a field
+    # declared with +type: BSON::Vector+ is stored as a BSON binary of the
+    # vector subtype and read back as a BSON::Vector.
+    #
+    # Querying by an exact vector value (e.g. +where(embedding: vector)+) is
+    # not supported: BSON::Vector subclasses Array, so the criteria selector
+    # treats it as a list of elements rather than a scalar. Use Atlas Vector
+    # Search ($vectorSearch) for similarity queries on vector fields.
+    module Vector
+      # Turn the object from the ruby type we deal with to a Mongo friendly
+      # type.
+      #
+      # @example Mongoize the object.
+      #   object.mongoize
+      #
+      # @return [ BSON::Binary | nil ] The object as a vector binary.
+      def mongoize
+        BSON::Vector.mongoize(self)
+      end
+
+      module ClassMethods
+        # Mongoize an object of any type to how it's stored in the db.
+        #
+        # @example Mongoize the object.
+        #   BSON::Vector.mongoize(vector)
+        #
+        # @param [ Object ] object The object to Mongoize.
+        #
+        # @return [ BSON::Binary | nil ] A vector binary or nil.
+        def mongoize(object)
+          case object
+          when BSON::Vector then BSON::Binary.from_vector(object)
+          when BSON::Binary then object
+          end
+        end
+
+        # Convert the object from its mongo friendly ruby type back to a
+        # BSON::Vector.
+        #
+        # @example Demongoize the object.
+        #   BSON::Vector.demongoize(binary)
+        #
+        # @param [ Object ] object The object to demongoize.
+        #
+        # @return [ BSON::Vector | nil ] The vector or nil.
+        def demongoize(object)
+          case object
+          when BSON::Binary then (object.type == :vector) ? object.as_vector : nil
+          when BSON::Vector then object
+          end
+        end
+      end
+    end
+  end
+end
+
+if defined?(BSON::Vector)
+  BSON::Vector.include Mongoid::Extensions::Vector
+  BSON::Vector.extend(Mongoid::Extensions::Vector::ClassMethods)
+end

spec/mongoid/extensions/vector_spec.rb

@@ -0,0 +1,175 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+# BSON::Vector only exists in bson-ruby >= 5.1, so the model and all examples
+# are guarded by min_bson_version. The describe argument is a string (not the
+# constant) so this file loads cleanly on older bson versions.
+if defined?(BSON::Vector)
+  class VectorEmbedding
+    include Mongoid::Document
+
+    field :embedding, type: BSON::Vector
+  end
+end
+
+describe 'BSON::Vector field type' do
+  min_bson_version '5.1'
+
+  shared_examples 'a vector round-trip' do
+    let(:vector) do
+      BSON::Vector.new(values, dtype, padding)
+    end
+
+    let(:binary) do
+      BSON::Binary.from_vector(vector)
+    end
+
+    describe '.mongoize' do
+      it 'returns the matching vector binary' do
+        expect(BSON::Vector.mongoize(vector)).to eq(binary)
+        expect(BSON::Vector.mongoize(vector).type).to eq(:vector)
+      end
+    end
+
+    describe '.demongoize' do
+      let(:demongoized) do
+        BSON::Vector.demongoize(binary)
+      end
+
+      it 'preserves the values, dtype and padding' do
+        expect(demongoized).to be_a(BSON::Vector)
+        expect(demongoized.to_a).to eq(values)
+        expect(demongoized.dtype).to eq(dtype)
+        expect(demongoized.padding).to eq(padding)
+      end
+    end
+
+    context 'when used as a field type' do
+      let!(:model) do
+        VectorEmbedding.create!(embedding: vector)
+      end
+
+      let(:reloaded) do
+        VectorEmbedding.find(model._id)
+      end
+
+      it 'stores the value as a vector binary' do
+        stored = VectorEmbedding.collection.find(_id: model._id).first['embedding']
+        expect(stored).to be_a(BSON::Binary)
+        expect(stored.type).to eq(:vector)
+      end
+
+      it 'reads the value back as an equivalent BSON::Vector' do
+        expect(reloaded.embedding).to be_a(BSON::Vector)
+        expect(reloaded.embedding.to_a).to eq(values)
+        expect(reloaded.embedding.dtype).to eq(dtype)
+        expect(reloaded.embedding.padding).to eq(padding)
+      end
+    end
+  end
+
+  context 'with an int8 vector' do
+    let(:values) { [ 1, 2, 3 ] }
+    let(:dtype) { :int8 }
+    let(:padding) { 0 }
+
+    it_behaves_like 'a vector round-trip'
+  end
+
+  context 'with a float32 vector' do
+    let(:values) { [ 1.5, -2.0, 0.25 ] }
+    let(:dtype) { :float32 }
+    let(:padding) { 0 }
+
+    it_behaves_like 'a vector round-trip'
+  end
+
+  context 'with a packed_bit vector' do
+    let(:values) { [ 255, 0, 128 ] }
+    let(:dtype) { :packed_bit }
+    let(:padding) { 3 }
+
+    it_behaves_like 'a vector round-trip'
+  end
+
+  describe '#mongoize' do
+    let(:vector) do
+      BSON::Vector.new([ 1, 2, 3 ], :int8)
+    end
+
+    it 'delegates to the class method and returns a vector binary' do
+      expect(vector.mongoize).to eq(BSON::Binary.from_vector(vector))
+      expect(vector.mongoize.type).to eq(:vector)
+    end
+
+    # Regression guard: BSON::Vector < ::Array, and Array has its own Mongoid
+    # extension. Without our override the value would be mongoized as a plain
+    # array, losing the vector subtype.
+    it 'does not mongoize via the Array extension' do
+      expect(vector.mongoize).to be_a(BSON::Binary)
+      expect(vector.mongoize).not_to be_a(Array)
+    end
+  end
+
+  describe '.mongoize' do
+    let(:vector) do
+      BSON::Vector.new([ 1, 2, 3 ], :int8)
+    end
+
+    context 'when given a vector BSON::Binary' do
+      it 'returns it unchanged' do
+        binary = BSON::Binary.from_vector(vector)
+        expect(BSON::Vector.mongoize(binary)).to eq(binary)
+      end
+    end
+
+    context 'when given nil' do
+      it 'returns nil' do
+        expect(BSON::Vector.mongoize(nil)).to be_nil
+      end
+    end
+
+    context 'when given a plain array' do
+      it 'returns nil' do
+        expect(BSON::Vector.mongoize([ 1, 2, 3 ])).to be_nil
+      end
+    end
+
+    context 'when given an uncastable type' do
+      it 'returns nil' do
+        expect(BSON::Vector.mongoize(true)).to be_nil
+      end
+    end
+  end
+
+  describe '.demongoize' do
+    let(:vector) do
+      BSON::Vector.new([ 1, 2, 3 ], :int8)
+    end
+
+    context 'when given a BSON::Vector' do
+      it 'returns it unchanged' do
+        expect(BSON::Vector.demongoize(vector)).to eq(vector)
+      end
+    end
+
+    context 'when given a non-vector BSON::Binary' do
+      it 'returns nil' do
+        expect(BSON::Vector.demongoize(BSON::Binary.new('x', :generic))).to be_nil
+      end
+    end
+
+    context 'when given nil' do
+      it 'returns nil' do
+        expect(BSON::Vector.demongoize(nil)).to be_nil
+      end
+    end
+
+    context 'when given an uncastable type' do
+      it 'returns nil' do
+        expect(BSON::Vector.demongoize(true)).to be_nil
+      end
+    end
+  end
+end