lancelot 0.1.0

data/examples/vector_search.rb

@@ -0,0 +1,102 @@
+#!/usr/bin/env ruby
+
+require "bundler/setup"
+require "lancelot"
+require "tmpdir"
+
+# Simple function to generate random embeddings
+def generate_embedding(text)
+  # In real applications, you'd use an actual embedding model
+  # This is just for demonstration
+  text.bytes.take(128).map { |b| b / 255.0 } + [0.0] * (128 - text.bytes.take(128).length)
+end
+
+# Create a temporary directory for our dataset
+Dir.mktmpdir do |dir|
+  dataset_path = File.join(dir, "vector_dataset")
+  
+  puts "Creating dataset with vector support at: #{dataset_path}"
+  
+  # Create a dataset with vector schema
+  dataset = Lancelot::Dataset.create(dataset_path, schema: {
+    text: :string,
+    score: :float32,
+    vector: { type: "vector", dimension: 128 }
+  })
+  
+  # Sample documents about programming languages
+  documents = [
+    { 
+      text: "Ruby is a dynamic, object-oriented programming language",
+      score: 0.95
+    },
+    { 
+      text: "Python is great for data science and machine learning",
+      score: 0.92
+    },
+    { 
+      text: "JavaScript runs in browsers and on servers with Node.js",
+      score: 0.88
+    },
+    { 
+      text: "Rust provides memory safety without garbage collection",
+      score: 0.91
+    },
+    { 
+      text: "Go makes concurrent programming easy with goroutines",
+      score: 0.89
+    },
+    { 
+      text: "Java is widely used in enterprise applications",
+      score: 0.85
+    },
+    { 
+      text: "C++ offers high performance and low-level control",
+      score: 0.90
+    },
+    { 
+      text: "TypeScript adds static typing to JavaScript",
+      score: 0.87
+    }
+  ]
+  
+  # Add embeddings to documents
+  documents_with_embeddings = documents.map do |doc|
+    doc.merge(vector: generate_embedding(doc[:text]))
+  end
+  
+  # Add documents to dataset
+  puts "Adding documents to dataset..."
+  dataset.add_documents(documents_with_embeddings)
+  puts "Added #{dataset.count} documents"
+  
+  # Create vector index
+  puts "\nCreating vector index..."
+  dataset.create_vector_index("vector")
+  
+  # Perform vector search
+  puts "\nSearching for documents similar to 'dynamic programming languages'..."
+  query_embedding = generate_embedding("dynamic programming languages")
+  
+  results = dataset.vector_search(query_embedding, column: "vector", limit: 3)
+  
+  puts "\nTop 3 most similar documents:"
+  results.each_with_index do |doc, i|
+    puts "#{i + 1}. #{doc[:text]} (score: #{doc[:score]})"
+    puts
+  end
+  
+  # Search with nearest_neighbors alias
+  puts "Searching for documents similar to 'memory safety and performance'..."
+  query_embedding2 = generate_embedding("memory safety and performance")
+  
+  similar = dataset.nearest_neighbors(query_embedding2, k: 3, column: "vector")
+  
+  puts "\nTop 3 nearest neighbors:"
+  similar.each_with_index do |doc, i|
+    puts "#{i + 1}. #{doc[:text]} (score: #{doc[:score]})"
+    puts
+  end
+end
+
+puts "\nDone!"

data/ext/lancelot/.gitignore

@@ -0,0 +1,10 @@
+# Rust
+target/
+Cargo.lock
+
+# Ruby extension build artifacts
+*.bundle
+*.so
+*.o
+*.a
+mkmf.log

data/ext/lancelot/Cargo.toml

@@ -0,0 +1,28 @@
+[package]
+name = "lancelot"
+version = "0.1.0"
+edition = "2021"
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+magnus = { version = "0.7", features = ["rb-sys"] }
+lance = { version = "0.31", default-features = false }
+lance-index = "0.31"
+lance-linalg = "0.31"
+tokio = { version = "1", features = ["rt-multi-thread", "macros"] }
+arrow = "55"
+arrow-array = "55"
+arrow-schema = "55"
+arrow-data = "55"
+futures = "0.3"
+thiserror = "2"
+serde = { version = "1", features = ["derive"] }
+serde_json = "1"
+
+[package.metadata]
+# Minimum supported Rust version
+rust-version = "1.86.0"
+
+[dev-dependencies]

data/ext/lancelot/extconf.rb

@@ -0,0 +1,4 @@
+require "mkmf"
+require "rb_sys/mkmf"
+
+create_rust_makefile("lancelot/lancelot")

data/ext/lancelot/src/conversion.rs

@@ -0,0 +1,243 @@
+use magnus::{Error, Ruby, RHash, RArray, Symbol, Value, TryConvert, value::ReprValue};
+use arrow_schema::{DataType, Schema as ArrowSchema};
+use arrow_array::{RecordBatch, StringArray, Float32Array, ArrayRef, Array, FixedSizeListArray};
+use std::collections::HashMap;
+use std::sync::Arc;
+
+pub fn build_record_batch(
+    data: RArray,
+    schema: &ArrowSchema,
+) -> Result<RecordBatch, Error> {
+    let mut columns: HashMap<String, Vec<Option<String>>> = HashMap::new();
+    let mut float_columns: HashMap<String, Vec<Option<f32>>> = HashMap::new();
+    let mut int_columns: HashMap<String, Vec<Option<i64>>> = HashMap::new();
+    let mut bool_columns: HashMap<String, Vec<Option<bool>>> = HashMap::new();
+    let mut vector_columns: HashMap<String, Vec<Option<Vec<f32>>>> = HashMap::new();
+    
+    for field in schema.fields() {
+        match field.data_type() {
+            DataType::Utf8 => {
+                columns.insert(field.name().to_string(), Vec::new());
+            }
+            DataType::Float32 => {
+                float_columns.insert(field.name().to_string(), Vec::new());
+            }
+            DataType::Int64 => {
+                int_columns.insert(field.name().to_string(), Vec::new());
+            }
+            DataType::Boolean => {
+                bool_columns.insert(field.name().to_string(), Vec::new());
+            }
+            DataType::FixedSizeList(_, _) => {
+                vector_columns.insert(field.name().to_string(), Vec::new());
+            }
+            _ => {}
+        }
+    }
+
+    for item in data.into_iter() {
+        let item = RHash::try_convert(item)?;
+        for field in schema.fields() {
+            let key = Symbol::new(field.name());
+            let value: Value = item.fetch(key)
+                .or_else(|_| {
+                    // Try with string key  
+                    item.fetch(field.name().as_str())
+                })?;
+            
+            match field.data_type() {
+                DataType::Utf8 => {
+                    if value.is_nil() {
+                        columns.get_mut(field.name()).unwrap().push(None);
+                    } else {
+                        let s = String::try_convert(value)?;
+                        columns.get_mut(field.name()).unwrap().push(Some(s));
+                    }
+                }
+                DataType::Float32 => {
+                    if value.is_nil() {
+                        float_columns.get_mut(field.name()).unwrap().push(None);
+                    } else {
+                        let f = f64::try_convert(value)?;
+                        float_columns.get_mut(field.name()).unwrap().push(Some(f as f32));
+                    }
+                }
+                DataType::Int64 => {
+                    if value.is_nil() {
+                        int_columns.get_mut(field.name()).unwrap().push(None);
+                    } else {
+                        let i = i64::try_convert(value)?;
+                        int_columns.get_mut(field.name()).unwrap().push(Some(i));
+                    }
+                }
+                DataType::Boolean => {
+                    if value.is_nil() {
+                        bool_columns.get_mut(field.name()).unwrap().push(None);
+                    } else {
+                        let b = bool::try_convert(value)?;
+                        bool_columns.get_mut(field.name()).unwrap().push(Some(b));
+                    }
+                }
+                DataType::FixedSizeList(_, _) => {
+                    if value.is_nil() {
+                        vector_columns.get_mut(field.name()).unwrap().push(None);
+                    } else {
+                        let arr = RArray::try_convert(value)?;
+                        let vec: Vec<f32> = arr.into_iter()
+                            .map(|v| f64::try_convert(v).map(|f| f as f32))
+                            .collect::<Result<Vec<_>, _>>()?;
+                        vector_columns.get_mut(field.name()).unwrap().push(Some(vec));
+                    }
+                }
+                _ => {}
+            }
+        }
+    }
+
+    let mut arrays: Vec<ArrayRef> = Vec::new();
+    
+    for field in schema.fields() {
+        let array: ArrayRef = match field.data_type() {
+            DataType::Utf8 => {
+                let values = columns.get(field.name()).unwrap();
+                Arc::new(StringArray::from(values.clone()))
+            }
+            DataType::Float32 => {
+                let values = float_columns.get(field.name()).unwrap();
+                Arc::new(Float32Array::from(values.clone()))
+            }
+            DataType::Int64 => {
+                let values = int_columns.get(field.name()).unwrap();
+                Arc::new(arrow_array::Int64Array::from(values.clone()))
+            }
+            DataType::Boolean => {
+                let values = bool_columns.get(field.name()).unwrap();
+                Arc::new(arrow_array::BooleanArray::from(values.clone()))
+            }
+            DataType::FixedSizeList(inner_field, list_size) => {
+                let values = vector_columns.get(field.name()).unwrap();
+                // Build flat array of all values
+                let mut flat_values = Vec::new();
+                for vec_opt in values {
+                    match vec_opt {
+                        Some(vec) => {
+                            if vec.len() != *list_size as usize {
+                                return Err(Error::new(
+                                    magnus::exception::arg_error(),
+                                    format!("Vector dimension mismatch. Expected {}, got {}", list_size, vec.len())
+                                ));
+                            }
+                            flat_values.extend(vec);
+                        }
+                        None => {
+                            // Add nulls for the entire vector
+                            flat_values.extend(vec![0.0f32; *list_size as usize]);
+                        }
+                    }
+                }
+                
+                let flat_array = Float32Array::from(flat_values);
+                Arc::new(FixedSizeListArray::new(
+                    inner_field.clone(),
+                    *list_size,
+                    Arc::new(flat_array),
+                    None
+                ))
+            }
+            _ => return Err(Error::new(
+                magnus::exception::runtime_error(),
+                format!("Unsupported data type: {:?}", field.data_type())
+            ))
+        };
+        
+        arrays.push(array);
+    }
+
+    RecordBatch::try_new(Arc::new(schema.clone()), arrays)
+        .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))
+}
+
+pub fn convert_batch_to_ruby(batch: &RecordBatch) -> Result<Vec<RHash>, Error> {
+    let ruby = Ruby::get().unwrap();
+    let mut documents = Vec::new();
+    
+    let num_rows = batch.num_rows();
+    let schema = batch.schema();
+    
+    for row_idx in 0..num_rows {
+        let doc = ruby.hash_new();
+        
+        for (col_idx, field) in schema.fields().iter().enumerate() {
+            let column = batch.column(col_idx);
+            let key = Symbol::new(field.name());
+            
+            match field.data_type() {
+                DataType::Utf8 => {
+                    let array = column.as_any().downcast_ref::<StringArray>()
+                        .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Failed to cast to StringArray"))?;
+                    
+                    if array.is_null(row_idx) {
+                        doc.aset(key, ruby.qnil())?;
+                    } else {
+                        doc.aset(key, array.value(row_idx))?;
+                    }
+                }
+                DataType::Float32 => {
+                    let array = column.as_any().downcast_ref::<Float32Array>()
+                        .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Failed to cast to Float32Array"))?;
+                    
+                    if array.is_null(row_idx) {
+                        doc.aset(key, ruby.qnil())?;
+                    } else {
+                        doc.aset(key, array.value(row_idx))?;
+                    }
+                }
+                DataType::Int64 => {
+                    let array = column.as_any().downcast_ref::<arrow_array::Int64Array>()
+                        .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Failed to cast to Int64Array"))?;
+                    
+                    if array.is_null(row_idx) {
+                        doc.aset(key, ruby.qnil())?;
+                    } else {
+                        doc.aset(key, array.value(row_idx))?;
+                    }
+                }
+                DataType::Boolean => {
+                    let array = column.as_any().downcast_ref::<arrow_array::BooleanArray>()
+                        .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Failed to cast to BooleanArray"))?;
+                    
+                    if array.is_null(row_idx) {
+                        doc.aset(key, ruby.qnil())?;
+                    } else {
+                        doc.aset(key, array.value(row_idx))?;
+                    }
+                }
+                DataType::FixedSizeList(_, list_size) => {
+                    let array = column.as_any().downcast_ref::<FixedSizeListArray>()
+                        .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Failed to cast to FixedSizeListArray"))?;
+                    
+                    if array.is_null(row_idx) {
+                        doc.aset(key, ruby.qnil())?;
+                    } else {
+                        let values = array.value(row_idx);
+                        let float_array = values.as_any().downcast_ref::<Float32Array>()
+                            .ok_or_else(|| Error::new(magnus::exception::runtime_error(), "Failed to cast vector values to Float32Array"))?;
+                        
+                        let ruby_array = ruby.ary_new();
+                        for i in 0..*list_size {
+                            ruby_array.push(float_array.value(i as usize))?;
+                        }
+                        doc.aset(key, ruby_array)?;
+                    }
+                }
+                _ => {
+                    // Skip unsupported types for now
+                }
+            }
+        }
+        
+        documents.push(doc);
+    }
+    
+    Ok(documents)
+}