From fc78d5337beae381063ada893a976383c4a8d0eb Mon Sep 17 00:00:00 2001 From: Paul Nathan Date: Sun, 5 Oct 2025 11:42:02 -0700 Subject: [PATCH 1/2] Create test --- test | 1 + 1 file changed, 1 insertion(+) create mode 100644 test diff --git a/test b/test new file mode 100644 index 0000000..ed9b086 --- /dev/null +++ b/test @@ -0,0 +1 @@ +test the file \ No newline at end of file From 285c07581248a633fce65d77b62efd390ae03c5f Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Mon, 6 Oct 2025 16:15:40 +0000 Subject: [PATCH 2/2] feat(rust): Implement Collection for multi-source data management This commit adds a `Collection` struct that provides a unified view over multiple data sources, backed by the `BufferPool` for efficient caching. A `DataSource` struct is introduced to wrap raw data and define the `stride` for paging. The `Collection` is initialized with a list of these sources, chunks the data into pages, and loads them into a `FramePool`. The `Collection` implements a `get()` method for indexed access and the `Iterator` trait for seamless iteration, with the `BufferPool` handling data loading and caching transparently. Comprehensive integration tests are included to validate the new functionality. --- rust/src/collection.rs | 183 ++++++++++++++++++++++ rust/src/lib.rs | 1 + rust/tests/collection_integration_test.rs | 104 ++++++++++++ 3 files changed, 288 insertions(+) create mode 100644 rust/src/collection.rs create mode 100644 rust/tests/collection_integration_test.rs diff --git a/rust/src/collection.rs b/rust/src/collection.rs new file mode 100644 index 0000000..c78498b --- /dev/null +++ b/rust/src/collection.rs @@ -0,0 +1,183 @@ +use crate::bufferpool::{bottom_evictor, BufferPool}; +use crate::framepool::FramePool; +use std::cell::RefCell; +use std::sync::Arc; + +/// A DataSource encapsulates a vector of data items and a stride. +/// The stride determines how many items are packed into a single page +/// in the buffer pool. +pub struct DataSource +where + T: Clone, +{ + pub data: Vec>, + pub stride: usize, +} + +impl DataSource +where + T: Clone, +{ + /// Creates a new DataSource. + pub fn new(data: Vec>, stride: usize) -> Self { + DataSource { data, stride } + } +} + +#[derive(Debug)] +struct SourceInfo { + start_page: u64, + stride: usize, + item_count: usize, + page_count: u64, +} + +/// A Collection provides a unified, iterable, and indexable view over +/// multiple data sources, backed by a buffer pool for efficient caching. +pub struct Collection<'a, T> +where + T: Clone, +{ + buffer_pool: RefCell>>>, + source_info: Vec, + total_items: usize, +} + +impl<'a, T> Collection<'a, T> +where + T: Clone, +{ + /// Creates a new Collection. + /// + /// It takes a vector of `DataSource`s, populates the provided `FramePool`, + /// and sets up a `BufferPool` to manage access. + pub fn new( + sources: Vec>, + pool: &'a mut dyn FramePool>>, + buffer_pool_size: u64, + ) -> Result { + let mut source_info = Vec::new(); + let mut total_items = 0; + let mut page_counter = 0; + + for source in sources { + let item_count = source.data.len(); + let stride = source.stride; + if stride == 0 { + // A stride of 0 is invalid. + continue; + } + let page_count = ((item_count + stride - 1) / stride) as u64; + + let info = SourceInfo { + start_page: page_counter, + stride, + item_count, + page_count, + }; + source_info.push(info); + + for (i, chunk) in source.data.chunks(stride).enumerate() { + let current_page_id = page_counter + i as u64; + pool.put_frame(current_page_id, Arc::new(chunk.to_vec()))?; + } + + page_counter += page_count; + total_items += item_count; + } + + pool.resize(page_counter)?; + + let buffer_pool = BufferPool::new(buffer_pool_size as usize, pool, bottom_evictor); + + Ok(Self { + buffer_pool: RefCell::new(buffer_pool), + source_info, + total_items, + }) + } + + /// Returns the total number of items in the collection. + pub fn len(&self) -> usize { + self.total_items + } + + /// Retrieves an item by its global index. + /// + /// This method provides indexed access to the collection's data. It handles + /// mapping the global index to the correct page and offset, and uses the + /// buffer pool to fetch the page if it's not already in memory. + pub fn get(&self, index: usize) -> Option> { + if index >= self.total_items { + return None; + } + + let mut items_seen = 0; + for info in &self.source_info { + if index < items_seen + info.item_count { + let local_index = index - items_seen; + let page_offset = local_index / info.stride; + let item_offset = local_index % info.stride; + let page_id = info.start_page + page_offset as u64; + + let mut bp = self.buffer_pool.borrow_mut(); + if let Some(page_frame) = bp.get_page(page_id) { + let page_data = page_frame.data(); // Arc>> + return Some(page_data[item_offset].clone()); + } else { + return None; + } + } + items_seen += info.item_count; + } + + None + } + + /// Returns an iterator over the items in the collection. + pub fn iter(&self) -> CollectionIterator<'_, 'a, T> { + CollectionIterator { + collection: self, + current_index: 0, + } + } +} + +/// An iterator over the items in a `Collection`. +pub struct CollectionIterator<'iter, 'collection, T> +where + T: Clone, +{ + collection: &'iter Collection<'collection, T>, + current_index: usize, +} + +impl<'iter, 'collection, T> Iterator for CollectionIterator<'iter, 'collection, T> +where + T: Clone, +{ + type Item = Arc; + + fn next(&mut self) -> Option { + if self.current_index < self.collection.len() { + let item = self.collection.get(self.current_index); + self.current_index += 1; + item + } else { + None + } + } +} + +/// Allows `for item in &collection` syntax. +impl<'a, 'b, T> IntoIterator for &'b Collection<'a, T> +where + T: Clone, +{ + type Item = Arc; + type IntoIter = CollectionIterator<'b, 'a, T>; + + fn into_iter(self) -> Self::IntoIter { + self.iter() + } +} \ No newline at end of file diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 4e202aa..4b2561c 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -169,3 +169,4 @@ pub mod bufferpool; pub mod framepool; pub mod unique_stack; +pub mod collection; diff --git a/rust/tests/collection_integration_test.rs b/rust/tests/collection_integration_test.rs new file mode 100644 index 0000000..cbbadc1 --- /dev/null +++ b/rust/tests/collection_integration_test.rs @@ -0,0 +1,104 @@ +use bufferpool::collection::{Collection, DataSource}; +use bufferpool::framepool::{FramePool, MemPool}; +use std::sync::Arc; + +#[test] +fn test_collection_initialization_and_len() { + let data1: Vec> = (0..10) + .map(|i| Arc::new(format!("data1_{}", i))) + .collect(); + let data2: Vec> = (0..5) + .map(|i| Arc::new(format!("data2_{}", i))) + .collect(); + + let source1 = DataSource::new(data1, 3); + let source2 = DataSource::new(data2, 2); + + let mut pool = MemPool::new(); + let collection = Collection::new(vec![source1, source2], &mut pool, 2).unwrap(); + + assert_eq!(collection.len(), 15); +} + +#[test] +fn test_collection_get() { + let data1: Vec> = (0..10) + .map(|i| Arc::new(format!("data1_{}", i))) + .collect(); + let data2: Vec> = (0..5) + .map(|i| Arc::new(format!("data2_{}", i))) + .collect(); + + let source1 = DataSource::new(data1, 3); // 4 pages + let source2 = DataSource::new(data2, 2); // 3 pages + + let mut pool = MemPool::new(); + let collection = Collection::new(vec![source1, source2], &mut pool, 2).unwrap(); + + // Test indexing within the first data source + assert_eq!(*collection.get(0).unwrap(), "data1_0"); + assert_eq!(*collection.get(3).unwrap(), "data1_3"); // Crosses a page boundary + assert_eq!(*collection.get(9).unwrap(), "data1_9"); + + // Test indexing into the second data source + assert_eq!(*collection.get(10).unwrap(), "data2_0"); + assert_eq!(*collection.get(12).unwrap(), "data2_2"); // Crosses a page boundary + assert_eq!(*collection.get(14).unwrap(), "data2_4"); + + // Test out of bounds + assert!(collection.get(15).is_none()); +} + +#[test] +fn test_collection_iteration() { + let data1: Vec> = (0..10) + .map(|i| Arc::new(format!("data1_{}", i))) + .collect(); + let data2: Vec> = (0..5) + .map(|i| Arc::new(format!("data2_{}", i))) + .collect(); + + let source1 = DataSource::new(data1.clone(), 3); + let source2 = DataSource::new(data2.clone(), 2); + + let mut pool = MemPool::new(); + // Use a small buffer pool to force eviction + let collection = Collection::new(vec![source1, source2], &mut pool, 2).unwrap(); + + let mut expected_data = Vec::new(); + expected_data.extend(data1); + expected_data.extend(data2); + + let collected_data: Vec> = collection.iter().collect(); + + assert_eq!(collected_data.len(), 15); + assert_eq!(collected_data, expected_data); +} + +#[test] +fn test_collection_into_iter() { + let data1: Vec> = (0..8) + .map(|i| Arc::new(format!("d1_{}", i))) + .collect(); + let data2: Vec> = (0..6) + .map(|i| Arc::new(format!("d2_{}", i))) + .collect(); + + let source1 = DataSource::new(data1.clone(), 4); + let source2 = DataSource::new(data2.clone(), 3); + + let mut pool = MemPool::new(); + let collection = Collection::new(vec![source1, source2], &mut pool, 2).unwrap(); + + let mut expected_data = Vec::new(); + expected_data.extend(data1); + expected_data.extend(data2); + + // Use `into_iter()` which is called by `for ... in &collection` + let mut collected_data = vec![]; + for item in &collection { + collected_data.push(item); + } + + assert_eq!(collected_data, expected_data); +} \ No newline at end of file