zvec-ruby 0.1.0

data/test/test_helper.rb

@@ -0,0 +1,170 @@
+require "minitest/autorun"
+require "minitest/pride"
+require "fileutils"
+require "tmpdir"
+
+# The C++ extension requires zvec to be installed. For pure-Ruby layer tests,
+# we stub the Ext module when the native extension is unavailable.
+begin
+  require "zvec"
+  NATIVE_EXT_AVAILABLE = true
+rescue LoadError
+  NATIVE_EXT_AVAILABLE = false
+
+  # Minimal stubs so pure-Ruby logic can be tested without the compiled extension.
+  module Zvec
+    class Error < StandardError; end
+
+    module Ext
+      # Stub enums as simple modules with constants
+      module DataType
+        UNDEFINED = :undefined
+        BINARY = :binary; STRING = :string; BOOL = :bool
+        INT32 = :int32; INT64 = :int64; UINT32 = :uint32; UINT64 = :uint64
+        FLOAT = :float; DOUBLE = :double
+        VECTOR_FP32 = :vector_fp32; VECTOR_FP64 = :vector_fp64
+        VECTOR_FP16 = :vector_fp16; VECTOR_INT8 = :vector_int8
+        SPARSE_VECTOR_FP32 = :sparse_vector_fp32; SPARSE_VECTOR_FP16 = :sparse_vector_fp16
+        ARRAY_STRING = :array_string; ARRAY_INT32 = :array_int32
+        ARRAY_INT64 = :array_int64; ARRAY_FLOAT = :array_float
+        ARRAY_DOUBLE = :array_double; ARRAY_BOOL = :array_bool
+      end
+
+      module MetricType
+        UNDEFINED = :undefined; L2 = :l2; IP = :ip; COSINE = :cosine; MIPSL2 = :mipsl2
+      end
+
+      module IndexType
+        UNDEFINED = :undefined; HNSW = :hnsw; IVF = :ivf; FLAT = :flat; INVERT = :invert
+      end
+
+      module QuantizeType
+        UNDEFINED = :undefined; FP16 = :fp16; INT8 = :int8; INT4 = :int4
+      end
+
+      class Doc
+        attr_accessor :pk, :score
+        def initialize; @pk = ""; @score = 0.0; @fields = {}; end
+        def field_names; @fields.keys; end
+        def has?(f); @fields.key?(f); end
+        def has_value?(f); @fields.key?(f) && !@fields[f].nil?; end
+        def is_empty; @fields.empty?; end
+        alias_method :empty?, :is_empty
+        def set_null(f); @fields[f] = nil; end
+        def set_string(f, v); @fields[f] = v; end
+        def set_bool(f, v); @fields[f] = v; end
+        def set_int32(f, v); @fields[f] = v; end
+        def set_int64(f, v); @fields[f] = v; end
+        def set_uint32(f, v); @fields[f] = v; end
+        def set_uint64(f, v); @fields[f] = v; end
+        def set_float(f, v); @fields[f] = v; end
+        def set_double(f, v); @fields[f] = v; end
+        def set_float_vector(f, v); @fields[f] = v; end
+        def set_double_vector(f, v); @fields[f] = v; end
+        def set_string_array(f, v); @fields[f] = v; end
+        def get_string(f); @fields[f].is_a?(String) ? @fields[f] : nil; end
+        def get_bool(f); [true, false].include?(@fields[f]) ? @fields[f] : nil; end
+        def get_int32(f); @fields[f].is_a?(Integer) ? @fields[f] : nil; end
+        def get_int64(f); @fields[f].is_a?(Integer) ? @fields[f] : nil; end
+        def get_float(f); @fields[f].is_a?(Float) ? @fields[f] : nil; end
+        def get_double(f); @fields[f].is_a?(Float) ? @fields[f] : nil; end
+        def get_float_vector(f); @fields[f].is_a?(Array) && @fields[f].first.is_a?(Float) ? @fields[f] : nil; end
+        def get_double_vector(f); get_float_vector(f); end
+        def get_string_array(f); @fields[f].is_a?(Array) && @fields[f].first.is_a?(String) ? @fields[f] : nil; end
+        def to_s; "[pk:#{@pk}, score:#{@score}, fields:#{@fields.size}]"; end
+      end
+
+      class FieldSchema
+        attr_reader :name, :data_type
+        attr_accessor :dimension, :nullable
+        def initialize(name, data_type); @name = name; @data_type = data_type; @dimension = 0; @nullable = false; end
+        def set_index_params(_); end
+        def vector_field?; [:vector_fp32, :vector_fp64, :vector_fp16, :vector_int8].include?(@data_type); end
+        def to_s; "FieldSchema(#{@name}, #{@data_type})"; end
+      end
+
+      class CollectionSchema
+        attr_reader :name
+        def initialize(name); @name = name; @fields = {}; end
+        def add_field(fs); @fields[fs.name] = fs; end
+        def has_field?(n); @fields.key?(n); end
+        def field_names; @fields.keys; end
+        def all_field_names; @fields.keys; end
+        alias_method :field_names, :all_field_names
+        def fields; @fields.values; end
+        def vector_fields; @fields.values.select(&:vector_field?); end
+        def forward_fields; @fields.values.reject(&:vector_field?); end
+        def to_s; "CollectionSchema(#{@name})"; end
+      end
+
+      class HnswIndexParams
+        def initialize(metric, m: 16, ef_construction: 200, quantize_type: nil); end
+      end
+
+      class FlatIndexParams
+        def initialize(metric, quantize_type: nil); end
+      end
+
+      class IVFIndexParams
+        def initialize(metric, n_list: 1024, n_iters: 10, use_soar: false, quantize_type: nil); end
+      end
+
+      class InvertIndexParams
+        def initialize(enable_range_optimization: true, enable_extended_wildcard: false); end
+      end
+
+      class HnswQueryParams
+        attr_reader :ef
+        def initialize(ef: 200); @ef = ef; end
+      end
+
+      class CollectionOptions
+        attr_accessor :read_only, :enable_mmap, :max_buffer_size
+        def initialize; @read_only = false; @enable_mmap = true; @max_buffer_size = 64 * 1024 * 1024; end
+        alias_method :read_only?, :read_only
+        alias_method :enable_mmap?, :enable_mmap
+      end
+
+      class VectorQuery
+        attr_accessor :topk, :field_name, :filter, :include_vector
+        def initialize; @topk = 10; @field_name = ""; @filter = ""; @include_vector = false; end
+        def set_query_vector(arr); @query_vector = arr; end
+        def set_output_fields(f); @output_fields = f; end
+        def set_query_params(p); @query_params = p; end
+        alias_method :include_vector?, :include_vector
+      end
+
+      class CollectionStats
+        attr_accessor :doc_count
+        def initialize; @doc_count = 0; end
+        def index_completeness; {}; end
+        def to_s; "CollectionStats(doc_count=#{@doc_count})"; end
+      end
+
+      class Status
+        def initialize(ok = true, msg = ""); @ok = ok; @msg = msg; end
+        def ok?; @ok; end
+        def message; @msg; end
+        def to_s; @ok ? "OK" : @msg; end
+      end
+    end
+
+    require_relative "../lib/zvec/version"
+    require_relative "../lib/zvec/data_types"
+    require_relative "../lib/zvec/schema"
+    require_relative "../lib/zvec/doc"
+    require_relative "../lib/zvec/query"
+
+    include DataTypes
+  end
+end
+
+# Helper to create a temporary directory that is cleaned up after the test
+module TempDirHelper
+  def with_temp_dir(prefix = "zvec_test")
+    dir = Dir.mktmpdir(prefix)
+    yield dir
+  ensure
+    FileUtils.rm_rf(dir) if dir && Dir.exist?(dir)
+  end
+end

data/test/test_query.rb

@@ -0,0 +1,64 @@
+require_relative "test_helper"
+
+class TestVectorQuery < Minitest::Test
+  def test_initialize_basic
+    q = Zvec::VectorQuery.new(
+      field_name: "embedding",
+      vector: [1.0, 2.0, 3.0]
+    )
+    assert_kind_of Zvec::VectorQuery, q
+    assert_kind_of Zvec::Ext::VectorQuery, q.ext_query
+  end
+
+  def test_field_name
+    q = Zvec::VectorQuery.new(field_name: "vec", vector: [1.0])
+    assert_equal "vec", q.ext_query.field_name
+  end
+
+  def test_topk_default
+    q = Zvec::VectorQuery.new(field_name: "vec", vector: [1.0])
+    assert_equal 10, q.ext_query.topk
+  end
+
+  def test_topk_custom
+    q = Zvec::VectorQuery.new(field_name: "vec", vector: [1.0], topk: 5)
+    assert_equal 5, q.ext_query.topk
+  end
+
+  def test_filter
+    q = Zvec::VectorQuery.new(
+      field_name: "vec", vector: [1.0],
+      filter: "year > 2024"
+    )
+    assert_equal "year > 2024", q.ext_query.filter
+  end
+
+  def test_include_vector_default_false
+    q = Zvec::VectorQuery.new(field_name: "vec", vector: [1.0])
+    refute q.ext_query.include_vector?
+  end
+
+  def test_include_vector_true
+    q = Zvec::VectorQuery.new(
+      field_name: "vec", vector: [1.0],
+      include_vector: true
+    )
+    assert q.ext_query.include_vector?
+  end
+
+  def test_symbol_field_name
+    q = Zvec::VectorQuery.new(field_name: :embedding, vector: [1.0])
+    assert_equal "embedding", q.ext_query.field_name
+  end
+
+  def test_integer_vector_converted_to_float
+    q = Zvec::VectorQuery.new(field_name: "vec", vector: [1, 2, 3])
+    # Should not raise — integers get .to_f
+    assert_kind_of Zvec::VectorQuery, q
+  end
+
+  def test_ext_query_accessible
+    q = Zvec::VectorQuery.new(field_name: "vec", vector: [1.0])
+    refute_nil q.ext_query
+  end
+end

data/test/test_ruby_llm_store.rb

@@ -0,0 +1,166 @@
+require_relative "test_helper"
+
+# Tests for the RubyLLM::Store integration.
+# Requires native extension.
+class TestRubyLLMStore < Minitest::Test
+  include TempDirHelper
+
+  def setup
+    skip "Native extension not available" unless NATIVE_EXT_AVAILABLE
+    require "zvec/ruby_llm"
+  end
+
+  def test_initialize_creates_collection
+    with_temp_dir do |dir|
+      path = File.join(dir, "store")
+      store = Zvec::RubyLLM::Store.new(path, dimension: 4)
+      assert_kind_of Zvec::Collection, store.collection
+      assert_equal 4, store.dimension
+      store.collection.destroy
+    end
+  end
+
+  def test_initialize_with_metric
+    with_temp_dir do |dir|
+      [:cosine, :l2, :ip].each_with_index do |metric, i|
+        path = File.join(dir, "store_#{i}")
+        store = Zvec::RubyLLM::Store.new(path, dimension: 4, metric: metric)
+        assert_kind_of Zvec::Collection, store.collection
+        store.collection.destroy
+      end
+    end
+  end
+
+  def test_initialize_invalid_metric
+    with_temp_dir do |dir|
+      path = File.join(dir, "store")
+      assert_raises(ArgumentError) do
+        Zvec::RubyLLM::Store.new(path, dimension: 4, metric: :invalid)
+      end
+    end
+  end
+
+  def test_add_and_count
+    with_temp_dir do |dir|
+      path = File.join(dir, "store")
+      store = Zvec::RubyLLM::Store.new(path, dimension: 4)
+
+      store.add("doc1", embedding: [0.1, 0.2, 0.3, 0.4], content: "Hello")
+      assert_equal 1, store.count
+
+      store.add("doc2", embedding: [0.4, 0.3, 0.2, 0.1], content: "World")
+      assert_equal 2, store.count
+
+      store.collection.destroy
+    end
+  end
+
+  def test_add_many
+    with_temp_dir do |dir|
+      path = File.join(dir, "store")
+      store = Zvec::RubyLLM::Store.new(path, dimension: 4)
+
+      docs = [
+        { id: "a", embedding: [0.1, 0.2, 0.3, 0.4], content: "A" },
+        { id: "b", embedding: [0.4, 0.3, 0.2, 0.1], content: "B" },
+        { id: "c", embedding: [0.2, 0.4, 0.1, 0.3], content: "C" },
+      ]
+      store.add_many(docs)
+      assert_equal 3, store.count
+
+      store.collection.destroy
+    end
+  end
+
+  def test_search
+    with_temp_dir do |dir|
+      path = File.join(dir, "store")
+      store = Zvec::RubyLLM::Store.new(path, dimension: 4)
+
+      store.add("near", embedding: [1.0, 0.0, 0.0, 0.0], content: "Near")
+      store.add("far", embedding: [0.0, 0.0, 0.0, 1.0], content: "Far")
+
+      results = store.search([1.0, 0.0, 0.0, 0.0], top_k: 2)
+      assert_kind_of Array, results
+      assert_equal 2, results.size
+
+      # Each result is a hash
+      first = results.first
+      assert_kind_of Hash, first
+      assert first.key?(:id)
+      assert first.key?(:score)
+      assert first.key?(:content)
+      assert first.key?(:metadata)
+
+      # Nearest should be "near"
+      assert_equal "near", first[:id]
+
+      store.collection.destroy
+    end
+  end
+
+  def test_search_returns_content
+    with_temp_dir do |dir|
+      path = File.join(dir, "store")
+      store = Zvec::RubyLLM::Store.new(path, dimension: 4)
+
+      store.add("d1", embedding: [1.0, 0.0, 0.0, 0.0], content: "Ruby rocks")
+      results = store.search([1.0, 0.0, 0.0, 0.0], top_k: 1)
+      assert_equal "Ruby rocks", results.first[:content]
+
+      store.collection.destroy
+    end
+  end
+
+  def test_delete
+    with_temp_dir do |dir|
+      path = File.join(dir, "store")
+      store = Zvec::RubyLLM::Store.new(path, dimension: 4)
+
+      store.add("del1", embedding: [0.1, 0.2, 0.3, 0.4])
+      store.add("del2", embedding: [0.4, 0.3, 0.2, 0.1])
+      assert_equal 2, store.count
+
+      store.delete("del1")
+      assert_equal 1, store.count
+
+      store.collection.destroy
+    end
+  end
+
+  def test_flush
+    with_temp_dir do |dir|
+      path = File.join(dir, "store")
+      store = Zvec::RubyLLM::Store.new(path, dimension: 4)
+      store.add("f1", embedding: [0.1, 0.2, 0.3, 0.4])
+      store.flush
+      # Should not raise
+      store.collection.destroy
+    end
+  end
+
+  def test_custom_vector_field
+    with_temp_dir do |dir|
+      path = File.join(dir, "store")
+      store = Zvec::RubyLLM::Store.new(path, dimension: 4, vector_field: "my_vec")
+      store.add("v1", embedding: [0.1, 0.2, 0.3, 0.4])
+      assert_equal 1, store.count
+      store.collection.destroy
+    end
+  end
+
+  def test_custom_content_field
+    with_temp_dir do |dir|
+      path = File.join(dir, "store")
+      store = Zvec::RubyLLM::Store.new(path, dimension: 4, content_field: "text")
+      store.add("c1", embedding: [0.1, 0.2, 0.3, 0.4], content: "Custom")
+      results = store.search([0.1, 0.2, 0.3, 0.4], top_k: 1)
+      assert_equal "Custom", results.first[:content]
+      store.collection.destroy
+    end
+  end
+
+  def test_reopen_existing_store
+    skip "Cannot reopen without explicit close (collection locks on open)"
+  end
+end

data/test/test_schema.rb

@@ -0,0 +1,133 @@
+require_relative "test_helper"
+
+class TestSchema < Minitest::Test
+  def test_initialize_with_name
+    schema = Zvec::Schema.new("test_schema")
+    assert_equal "test_schema", schema.name
+  end
+
+  def test_initialize_with_block
+    schema = Zvec::Schema.new("test") do
+      string "title"
+      int32 "count"
+    end
+
+    assert schema.has_field?("title")
+    assert schema.has_field?("count")
+  end
+
+  def test_string_field
+    schema = Zvec::Schema.new("test") { string "name" }
+    assert schema.has_field?("name")
+    assert_equal Zvec::DataTypes::STRING, schema.field_type("name")
+  end
+
+  def test_int32_field
+    schema = Zvec::Schema.new("test") { int32 "count" }
+    assert schema.has_field?("count")
+    assert_equal Zvec::DataTypes::INT32, schema.field_type("count")
+  end
+
+  def test_int64_field
+    schema = Zvec::Schema.new("test") { int64 "big_count" }
+    assert schema.has_field?("big_count")
+    assert_equal Zvec::DataTypes::INT64, schema.field_type("big_count")
+  end
+
+  def test_float_field
+    schema = Zvec::Schema.new("test") { float "score" }
+    assert schema.has_field?("score")
+    assert_equal Zvec::DataTypes::FLOAT, schema.field_type("score")
+  end
+
+  def test_double_field
+    schema = Zvec::Schema.new("test") { double "precise" }
+    assert schema.has_field?("precise")
+    assert_equal Zvec::DataTypes::DOUBLE, schema.field_type("precise")
+  end
+
+  def test_bool_field
+    schema = Zvec::Schema.new("test") { bool "active" }
+    assert schema.has_field?("active")
+    assert_equal Zvec::DataTypes::BOOL, schema.field_type("active")
+  end
+
+  def test_vector_field
+    schema = Zvec::Schema.new("test") do
+      vector "embedding", dimension: 128
+    end
+    assert schema.has_field?("embedding")
+    assert_equal Zvec::DataTypes::VECTOR_FP32, schema.field_type("embedding")
+  end
+
+  def test_vector_field_custom_type
+    schema = Zvec::Schema.new("test") do
+      vector "embedding", dimension: 128, type: Zvec::DataTypes::VECTOR_FP64
+    end
+    assert_equal Zvec::DataTypes::VECTOR_FP64, schema.field_type("embedding")
+  end
+
+  def test_field_with_explicit_type
+    schema = Zvec::Schema.new("test")
+    schema.field("tags", Zvec::DataTypes::ARRAY_STRING)
+    assert schema.has_field?("tags")
+    assert_equal Zvec::DataTypes::ARRAY_STRING, schema.field_type("tags")
+  end
+
+  def test_field_names
+    schema = Zvec::Schema.new("test") do
+      string "a"
+      int32 "b"
+      vector "c", dimension: 4
+    end
+    names = schema.field_names
+    assert_includes names, "a"
+    assert_includes names, "b"
+    assert_includes names, "c"
+    assert_equal 3, names.size
+  end
+
+  def test_has_field_returns_false_for_missing
+    schema = Zvec::Schema.new("test")
+    refute schema.has_field?("nonexistent")
+  end
+
+  def test_field_type_returns_nil_for_missing
+    schema = Zvec::Schema.new("test")
+    assert_nil schema.field_type("nonexistent")
+  end
+
+  def test_symbol_field_name
+    schema = Zvec::Schema.new("test") { string "name" }
+    # field_type converts to string internally
+    assert_equal Zvec::DataTypes::STRING, schema.field_type(:name)
+  end
+
+  def test_field_returns_self_for_chaining
+    schema = Zvec::Schema.new("test")
+    result = schema.field("a", Zvec::DataTypes::STRING)
+    assert_same schema, result
+  end
+
+  def test_to_s_returns_string
+    schema = Zvec::Schema.new("test") { string "x" }
+    assert_kind_of String, schema.to_s
+  end
+
+  def test_ext_schema_accessible
+    schema = Zvec::Schema.new("test")
+    assert_kind_of Zvec::Ext::CollectionSchema, schema.ext_schema
+  end
+
+  def test_multiple_fields_in_block
+    schema = Zvec::Schema.new("docs") do
+      string "title"
+      string "body", nullable: true
+      int32 "year"
+      float "rating"
+      bool "published"
+      vector "embedding", dimension: 384
+    end
+    assert_equal 6, schema.field_names.size
+  end
+end

data/test/test_version.rb

@@ -0,0 +1,19 @@
+require_relative "test_helper"
+
+class TestVersion < Minitest::Test
+  def test_version_is_defined
+    refute_nil Zvec::VERSION
+  end
+
+  def test_version_is_a_string
+    assert_kind_of String, Zvec::VERSION
+  end
+
+  def test_version_format
+    assert_match(/\A\d+\.\d+\.\d+/, Zvec::VERSION)
+  end
+
+  def test_version_value
+    assert_equal "0.1.0", Zvec::VERSION
+  end
+end

data/zvec.gemspec

@@ -0,0 +1,43 @@
+require_relative "lib/zvec/version"
+
+Gem::Specification.new do |spec|
+  spec.name          = "zvec-ruby"
+  spec.version       = Zvec::VERSION
+  spec.authors       = ["Johannes Dwi Cahyo"]
+  spec.summary       = "Ruby bindings for zvec vector database"
+  spec.description   = "Ruby gem wrapping the zvec C++ vector database (https://github.com/alibaba/zvec) " \
+                        "using Rice. Provides Collection, Doc, Schema, and VectorQuery classes for " \
+                        "high-performance vector similarity search, plus integrations for ruby_llm and ActiveRecord."
+  spec.homepage      = "https://github.com/johannesdwicahyo/zvec-ruby"
+  spec.license       = "Apache-2.0"
+
+  spec.required_ruby_version = ">= 3.1.0"
+
+  spec.files = Dir[
+    "lib/**/*.rb",
+    "ext/**/*.{rb,cpp,h,hpp}",
+    "examples/**/*.rb",
+    "test/**/*.rb",
+    "README.md",
+    "LICENSE",
+    "CHANGELOG.md",
+    "Rakefile",
+    "zvec.gemspec"
+  ]
+
+  spec.metadata = {
+    "homepage_uri"    => spec.homepage,
+    "source_code_uri" => spec.homepage,
+    "changelog_uri"   => "#{spec.homepage}/blob/main/CHANGELOG.md",
+    "bug_tracker_uri" => "#{spec.homepage}/issues",
+  }
+
+  spec.extensions    = ["ext/zvec/extconf.rb"]
+  spec.require_paths = ["lib"]
+
+  spec.add_dependency "rice", ">= 4.0"
+
+  spec.add_development_dependency "rake-compiler", "~> 1.2"
+  spec.add_development_dependency "rake", "~> 13.0"
+  spec.add_development_dependency "minitest", "~> 5.0"
+end