Expose distribute_indexes_ruby

src/distribute_indexes.rs

@@ -118,4 +118,16 @@ mod tests {
             vec![0, 3, 4, 1, 6, 5, 2, 7, 8, 10, 9, 12, 11, 13, 14]
         );
     }
+
+    #[test]
+    fn test_distribute2() {
+        let data = vec!["Picture", "Post", "Video", "Video", "Picture", "Post", "Picture", "Picture", "Video"];
+        let result = distribute_ids(&data, 3);
+        let result_with_labels = result
+            .iter()
+            .map(|idx| data[*idx])
+            .collect::<Vec<_>>();
+        assert_eq!(result_with_labels, vec!["Picture", "Post", "Video", "Picture", "Post", "Video", "Picture", "Video", "Picture"]);
+        assert_eq!(result, vec![0, 1, 2, 4, 5, 3, 6, 8, 7]);
+    }
 }

src/lib.rs

@@ -7,4 +7,5 @@ mod ruby_ext;
 
 pub use distribute_csv::distribute as distribute_csv;
 pub use distribute_indexes::distribute;
+pub use distribute_indexes::distribute_ids;
 pub use distributing_iterator::DistributingIterator;

src/ruby_ext.rs

@@ -1,19 +1,22 @@
-use magnus::{error::Result, function, exception, Error};
+use magnus::{error::Result, exception, function, Error, Value};
 
-use crate::distribute_csv;
+use crate::{distribute_ids, distribute_csv};
 
 fn distribute_csv_ruby(data: String, field: String, spread: u64) -> Result<String> {
     match distribute_csv(&data, &field, spread) {
         Ok(result) => Ok(result),
-        Err(e) => {
-            Err(Error::new(exception::standard_error(), format!("{:?}", e)))
-        }
+        Err(e) => Err(Error::new(exception::standard_error(), format!("{:?}", e))),
     }
 }
 
+fn distribute_indexes_ruby(data: Vec<String>, spread: usize) -> Result<Vec<usize>> {
+    Ok(distribute_ids(&data, spread))
+}
+
 #[magnus::init]
 fn init(ruby: &magnus::Ruby) -> Result<()> {
     let module = ruby.define_module("DistributingIterator")?;
-    module.define_module_function("distribute", function!(distribute_csv_ruby, 3))?;
+    module.define_module_function("distribute_csv", function!(distribute_csv_ruby, 3))?;
+    module.define_module_function("distribute_indexes", function!(distribute_indexes_ruby, 2))?;
     Ok(())
 }