lancelot 0.1.0

data/ext/lancelot/src/dataset.rs

@@ -0,0 +1,454 @@
+use magnus::{Error, Ruby, RHash, RArray, Symbol, TryConvert, function, method, RClass, Module, Object};
+use std::cell::RefCell;
+use std::sync::Arc;
+use tokio::runtime::Runtime;
+use lance::Dataset;
+use lance::index::vector::VectorIndexParams;
+use lance_index::{IndexType, DatasetIndexExt};
+use lance_index::scalar::{InvertedIndexParams, FullTextSearchQuery};
+use arrow_array::{RecordBatch, RecordBatchIterator, Float32Array};
+use futures::stream::TryStreamExt;
+
+use crate::schema::build_arrow_schema;
+use crate::conversion::{build_record_batch, convert_batch_to_ruby};
+
+#[magnus::wrap(class = "Lancelot::Dataset", free_immediately, size)]
+pub struct LancelotDataset {
+    dataset: RefCell<Option<Dataset>>,
+    runtime: RefCell<Runtime>,
+    path: String,
+}
+
+impl LancelotDataset {
+    pub fn new(path: String) -> Result<Self, Error> {
+        let runtime = Runtime::new()
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        
+        Ok(Self {
+            dataset: RefCell::new(None),
+            runtime: RefCell::new(runtime),
+            path,
+        })
+    }
+
+    pub fn path(&self) -> String {
+        self.path.clone()
+    }
+
+    pub fn create(&self, schema_hash: RHash) -> Result<(), Error> {
+        let schema = build_arrow_schema(schema_hash)?;
+        
+        let empty_batch = RecordBatch::new_empty(Arc::new(schema.clone()));
+        let batches = vec![empty_batch];
+        let reader = RecordBatchIterator::new(
+            batches.into_iter().map(Ok),
+            Arc::new(schema)
+        );
+        
+        let dataset = self.runtime.borrow_mut().block_on(async {
+            Dataset::write(
+                reader,
+                &self.path,
+                None,
+            )
+            .await
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))
+        })?;
+
+        self.dataset.replace(Some(dataset));
+        Ok(())
+    }
+
+    pub fn open(&self) -> Result<(), Error> {
+        let dataset = self.runtime.borrow_mut().block_on(async {
+            Dataset::open(&self.path)
+                .await
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))
+        })?;
+
+        self.dataset.replace(Some(dataset));
+        Ok(())
+    }
+
+    pub fn add_data(&self, data: RArray) -> Result<(), Error> {
+        let mut dataset = self.dataset.borrow_mut();
+        let dataset = dataset.as_mut()
+            .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Dataset not opened"))?;
+
+        // Check if data is empty
+        if data.len() == 0 {
+            return Ok(());  // Nothing to add
+        }
+
+        // Get the dataset's schema
+        let schema = self.runtime.borrow_mut().block_on(async {
+            dataset.schema()
+        });
+        
+        // Convert Lance schema to Arrow schema
+        let arrow_schema = schema.into();
+
+        let batch = build_record_batch(data, &arrow_schema)?;
+
+        let batches = vec![batch];
+        let reader = RecordBatchIterator::new(
+            batches.into_iter().map(Ok),
+            Arc::new(arrow_schema)
+        );
+        
+        self.runtime.borrow_mut().block_on(async move {
+            dataset.append(reader, None)
+                .await
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))
+        })?;
+
+        Ok(())
+    }
+
+    pub fn count_rows(&self) -> Result<i64, Error> {
+        let dataset = self.dataset.borrow();
+        let dataset = dataset.as_ref()
+            .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Dataset not opened"))?;
+
+        let count = self.runtime.borrow_mut().block_on(async {
+            dataset.count_rows(None)
+                .await
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))
+        })?;
+
+        Ok(count as i64)
+    }
+
+    pub fn schema(&self) -> Result<RHash, Error> {
+        let dataset = self.dataset.borrow();
+        let _dataset = dataset.as_ref()
+            .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Dataset not opened"))?;
+
+        let ruby = Ruby::get().unwrap();
+        let hash = ruby.hash_new();
+        
+        // TODO: Read actual schema from Lance dataset once we figure out the 0.31 API
+        // For now, return a hardcoded schema that matches what we support
+        hash.aset(Symbol::new("text"), "string")?;
+        hash.aset(Symbol::new("score"), "float32")?;
+
+        Ok(hash)
+    }
+
+    pub fn scan_all(&self) -> Result<RArray, Error> {
+        let dataset = self.dataset.borrow();
+        let dataset = dataset.as_ref()
+            .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Dataset not opened"))?;
+
+        let batches: Vec<RecordBatch> = self.runtime.borrow_mut().block_on(async {
+            let scanner = dataset.scan();
+            let stream = scanner
+                .try_into_stream()
+                .await
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            
+            stream
+                .try_collect::<Vec<_>>()
+                .await
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))
+        })?;
+
+        let ruby = Ruby::get().unwrap();
+        let result_array = ruby.ary_new();
+
+        for batch in batches {
+            let documents = convert_batch_to_ruby(&batch)?;
+            for doc in documents {
+                result_array.push(doc)?;
+            }
+        }
+
+        Ok(result_array)
+    }
+
+    pub fn scan_limit(&self, limit: i64) -> Result<RArray, Error> {
+        let dataset = self.dataset.borrow();
+        let dataset = dataset.as_ref()
+            .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Dataset not opened"))?;
+
+        let batches: Vec<RecordBatch> = self.runtime.borrow_mut().block_on(async {
+            let mut scanner = dataset.scan();
+            scanner.limit(Some(limit), None)
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            
+            let stream = scanner
+                .try_into_stream()
+                .await
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            
+            stream
+                .try_collect::<Vec<_>>()
+                .await
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))
+        })?;
+
+        let ruby = Ruby::get().unwrap();
+        let result_array = ruby.ary_new();
+
+        for batch in batches {
+            let documents = convert_batch_to_ruby(&batch)?;
+            for doc in documents {
+                result_array.push(doc)?;
+            }
+        }
+
+        Ok(result_array)
+    }
+
+    pub fn create_vector_index(&self, column: String) -> Result<(), Error> {
+        let mut dataset = self.dataset.borrow_mut();
+        let dataset = dataset.as_mut()
+            .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Dataset not opened"))?;
+
+        self.runtime.borrow_mut().block_on(async move {
+            // Get row count to determine optimal number of partitions
+            let num_rows = dataset.count_rows(None).await
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            
+            // Use fewer partitions for small datasets
+            let num_partitions = if num_rows < 256 {
+                std::cmp::max(1, (num_rows / 4) as usize)
+            } else {
+                256
+            };
+            
+            // Create IVF_FLAT vector index parameters
+            let params = VectorIndexParams::ivf_flat(num_partitions, lance_linalg::distance::MetricType::L2);
+            
+            dataset.create_index(
+                &[&column],
+                IndexType::Vector,
+                None,
+                &params,
+                true
+            )
+            .await
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))
+        })
+    }
+
+    pub fn vector_search(&self, column: String, query_vector: RArray, limit: i64) -> Result<RArray, Error> {
+        let dataset = self.dataset.borrow();
+        let dataset = dataset.as_ref()
+            .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Dataset not opened"))?;
+
+        // Convert Ruby array to Vec<f32>
+        let vector: Vec<f32> = query_vector
+            .into_iter()
+            .map(|v| f64::try_convert(v).map(|f| f as f32))
+            .collect::<Result<Vec<_>, _>>()?;
+
+        let batches: Vec<RecordBatch> = self.runtime.borrow_mut().block_on(async {
+            let mut scanner = dataset.scan();
+            
+            // Use nearest for vector search
+            scanner.nearest(&column, &Float32Array::from(vector), limit as usize)
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            
+            let stream = scanner
+                .try_into_stream()
+                .await
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            
+            stream
+                .try_collect::<Vec<_>>()
+                .await
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))
+        })?;
+
+        let ruby = Ruby::get().unwrap();
+        let result_array = ruby.ary_new();
+
+        for batch in batches {
+            let documents = convert_batch_to_ruby(&batch)?;
+            for doc in documents {
+                result_array.push(doc)?;
+            }
+        }
+
+        Ok(result_array)
+    }
+
+    pub fn create_text_index(&self, column: String) -> Result<(), Error> {
+        let mut dataset = self.dataset.borrow_mut();
+        let dataset = dataset.as_mut()
+            .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Dataset not opened"))?;
+
+        self.runtime.borrow_mut().block_on(async move {
+            // Create inverted index for full-text search
+            let params = InvertedIndexParams::default();
+            
+            dataset.create_index(
+                &[&column],
+                IndexType::Inverted,
+                None,
+                &params,
+                true
+            )
+            .await
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))
+        })
+    }
+
+    pub fn text_search(&self, column: String, query: String, limit: i64) -> Result<RArray, Error> {
+        let dataset = self.dataset.borrow();
+        let dataset = dataset.as_ref()
+            .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Dataset not opened"))?;
+
+        let batches: Vec<RecordBatch> = self.runtime.borrow_mut().block_on(async {
+            let mut scanner = dataset.scan();
+            
+            // Use full-text search with inverted index
+            let fts_query = FullTextSearchQuery::new(query)
+                .with_column(column)
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            
+            scanner.full_text_search(fts_query)
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            
+            // Apply limit
+            scanner.limit(Some(limit), None)
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            
+            let stream = scanner
+                .try_into_stream()
+                .await
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            
+            stream
+                .try_collect::<Vec<_>>()
+                .await
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))
+        })?;
+
+        let ruby = Ruby::get().unwrap();
+        let result_array = ruby.ary_new();
+
+        for batch in batches {
+            let documents = convert_batch_to_ruby(&batch)?;
+            for doc in documents {
+                result_array.push(doc)?;
+            }
+        }
+
+        Ok(result_array)
+    }
+
+    pub fn multi_column_text_search(&self, columns: RArray, query: String, limit: i64) -> Result<RArray, Error> {
+        let dataset = self.dataset.borrow();
+        let dataset = dataset.as_ref()
+            .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Dataset not opened"))?;
+
+        // Convert Ruby array of columns to Vec<String>
+        let columns: Vec<String> = columns
+            .into_iter()
+            .map(|v| String::try_convert(v))
+            .collect::<Result<Vec<_>, _>>()?;
+
+        let batches: Vec<RecordBatch> = self.runtime.borrow_mut().block_on(async {
+            let mut scanner = dataset.scan();
+            
+            // Create a full-text search query for multiple columns
+            let fts_query = FullTextSearchQuery::new(query)
+                .with_columns(&columns)
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            
+            scanner.full_text_search(fts_query)
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            
+            // Apply limit
+            scanner.limit(Some(limit), None)
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            
+            let stream = scanner
+                .try_into_stream()
+                .await
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            
+            stream
+                .try_collect::<Vec<_>>()
+                .await
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))
+        })?;
+
+        let ruby = Ruby::get().unwrap();
+        let result_array = ruby.ary_new();
+
+        for batch in batches {
+            let documents = convert_batch_to_ruby(&batch)?;
+            for doc in documents {
+                result_array.push(doc)?;
+            }
+        }
+
+        Ok(result_array)
+    }
+
+    pub fn filter_scan(&self, filter_expr: String, limit: Option<i64>) -> Result<RArray, Error> {
+        let dataset = self.dataset.borrow();
+        let dataset = dataset.as_ref()
+            .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Dataset not opened"))?;
+
+        let batches: Vec<RecordBatch> = self.runtime.borrow_mut().block_on(async {
+            let mut scanner = dataset.scan();
+            
+            // Apply SQL-like filter
+            scanner.filter(&filter_expr)
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            
+            // Apply limit if provided
+            if let Some(lim) = limit {
+                scanner.limit(Some(lim), None)
+                    .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            }
+            
+            let stream = scanner
+                .try_into_stream()
+                .await
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            
+            stream
+                .try_collect::<Vec<_>>()
+                .await
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))
+        })?;
+
+        let ruby = Ruby::get().unwrap();
+        let result_array = ruby.ary_new();
+
+        for batch in batches {
+            let documents = convert_batch_to_ruby(&batch)?;
+            for doc in documents {
+                result_array.push(doc)?;
+            }
+        }
+
+        Ok(result_array)
+    }
+}
+
+impl LancelotDataset {
+    pub fn bind(class: &RClass) -> Result<(), Error> {
+        class.define_singleton_method("new", function!(LancelotDataset::new, 1))?;
+        class.define_method("path", method!(LancelotDataset::path, 0))?;
+        class.define_method("create", method!(LancelotDataset::create, 1))?;
+        class.define_method("open", method!(LancelotDataset::open, 0))?;
+        class.define_method("add_data", method!(LancelotDataset::add_data, 1))?;
+        class.define_method("count_rows", method!(LancelotDataset::count_rows, 0))?;
+        class.define_method("schema", method!(LancelotDataset::schema, 0))?;
+        class.define_method("scan_all", method!(LancelotDataset::scan_all, 0))?;
+        class.define_method("scan_limit", method!(LancelotDataset::scan_limit, 1))?;
+        class.define_method("create_vector_index", method!(LancelotDataset::create_vector_index, 1))?;
+        class.define_method("create_text_index", method!(LancelotDataset::create_text_index, 1))?;
+        class.define_method("_rust_vector_search", method!(LancelotDataset::vector_search, 3))?;
+        class.define_method("_rust_text_search", method!(LancelotDataset::text_search, 3))?;
+        class.define_method("_rust_multi_column_text_search", method!(LancelotDataset::multi_column_text_search, 3))?;
+        class.define_method("filter_scan", method!(LancelotDataset::filter_scan, 2))?;
+        Ok(())
+    }
+}

data/ext/lancelot/src/lib.rs

@@ -0,0 +1,17 @@
+use magnus::{define_module, Error, Ruby, Module};
+
+mod dataset;
+mod schema;
+mod conversion;
+
+use dataset::LancelotDataset;
+
+#[magnus::init]
+fn init(ruby: &Ruby) -> Result<(), Error> {
+    let module = define_module("Lancelot")?;
+    
+    let dataset_class = module.define_class("Dataset", ruby.class_object())?;
+    LancelotDataset::bind(&dataset_class)?;
+    
+    Ok(())
+}

data/ext/lancelot/src/schema.rs

@@ -0,0 +1,50 @@
+use magnus::{Error, RHash, Symbol, Value, TryConvert, r_hash::ForEach, value::ReprValue};
+use arrow_schema::{DataType, Field, Schema as ArrowSchema};
+use std::sync::Arc;
+
+pub fn build_arrow_schema(schema_hash: RHash) -> Result<ArrowSchema, Error> {
+    let mut fields = Vec::new();
+
+    schema_hash.foreach(|key: Symbol, value: Value| {
+        let field_name = key.name()?.to_string();
+        
+        let data_type = if value.is_kind_of(magnus::class::hash()) {
+            let hash = RHash::from_value(value)
+                .ok_or_else(|| Error::new(magnus::exception::arg_error(), "Invalid hash value"))?;
+            let type_str: String = hash.fetch(Symbol::new("type"))?;
+            
+            match type_str.as_str() {
+                "vector" => {
+                    let dimension: i32 = hash.fetch(Symbol::new("dimension"))?;
+                    DataType::FixedSizeList(
+                        Arc::new(Field::new("item", DataType::Float32, true)),
+                        dimension,
+                    )
+                }
+                _ => return Err(Error::new(
+                    magnus::exception::arg_error(),
+                    format!("Unknown field type: {}", type_str)
+                ))
+            }
+        } else {
+            let type_str = String::try_convert(value)?;
+            match type_str.as_str() {
+                "string" => DataType::Utf8,
+                "float32" => DataType::Float32,
+                "float64" => DataType::Float64,
+                "int32" => DataType::Int32,
+                "int64" => DataType::Int64,
+                "boolean" => DataType::Boolean,
+                _ => return Err(Error::new(
+                    magnus::exception::arg_error(),
+                    format!("Unknown field type: {}", type_str)
+                ))
+            }
+        };
+
+        fields.push(Field::new(field_name, data_type, true));
+        Ok(ForEach::Continue)
+    })?;
+
+    Ok(ArrowSchema::new(fields))
+}

data/lib/lancelot/dataset.rb

@@ -0,0 +1,119 @@
+# frozen_string_literal: true
+
+module Lancelot
+  class Dataset
+    class << self
+      def create(path, schema:)
+        dataset = new(path)
+        dataset.create(normalize_schema(schema))
+        dataset
+      end
+
+      def open(path)
+        dataset = new(path)
+        dataset.open
+        dataset
+      end
+
+      private
+
+      def normalize_schema(schema)
+        schema.transform_values do |type|
+          case type
+          when Hash
+            type
+          when :string, "string"
+            "string"
+          when :float, :float32, "float", "float32"
+            "float32"
+          when :float64, "float64"
+            "float64"
+          when :int, :int32, "int", "int32"
+            "int32"
+          when :int64, "int64"
+            "int64"
+          when :bool, :boolean, "bool", "boolean"
+            "boolean"
+          else
+            raise ArgumentError, "Unknown type: #{type}"
+          end
+        end
+      end
+    end
+
+    def add_documents(documents)
+      add_data(documents.map { |doc| normalize_document(doc) })
+    end
+
+    def <<(document)
+      add_documents([document])
+      self
+    end
+
+    def size
+      count_rows
+    end
+
+    alias_method :count, :size
+    alias_method :length, :size
+
+    def all
+      scan_all
+    end
+
+    def first(n = nil)
+      if n.nil?
+        scan_limit(1).first
+      else
+        scan_limit(n)
+      end
+    end
+
+    def each(&block)
+      return enum_for(:each) unless block_given?
+      scan_all.each(&block)
+    end
+
+    include Enumerable
+
+    def vector_search(query_vector, column: "vector", limit: 10)
+      unless query_vector.is_a?(Array)
+        raise ArgumentError, "Query vector must be an array of numbers"
+      end
+      
+      _rust_vector_search(column.to_s, query_vector, limit)
+    end
+
+    def nearest_neighbors(vector, k: 10, column: "vector")
+      vector_search(vector, column: column, limit: k)
+    end
+
+    def text_search(query, column: nil, columns: nil, limit: 10)
+      unless query.is_a?(String)
+        raise ArgumentError, "Query must be a string"
+      end
+      
+      if column && columns
+        raise ArgumentError, "Cannot specify both column and columns"
+      elsif columns
+        # Multi-column search
+        columns = Array(columns).map(&:to_s)
+        _rust_multi_column_text_search(columns, query, limit)
+      else
+        # Single column search (default to "text" if not specified)
+        column ||= "text"
+        _rust_text_search(column.to_s, query, limit)
+      end
+    end
+
+    def where(filter_expression, limit: nil)
+      filter_scan(filter_expression.to_s, limit)
+    end
+
+    private
+
+    def normalize_document(doc)
+      doc.transform_keys(&:to_sym)
+    end
+  end
+end

data/lib/lancelot/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module Lancelot
+  VERSION = "0.1.0"
+end

data/lib/lancelot.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+require_relative "lancelot/version"
+require_relative "lancelot/lancelot"
+require_relative "lancelot/dataset"
+
+module Lancelot
+  class Error < StandardError; end
+end

data/sig/lancelot.rbs

@@ -0,0 +1,4 @@
+module Lancelot
+  VERSION: String
+  # See the writing guide of rbs: https://github.com/ruby/rbs#guides
+end