Merge pull request #5 from dapper91/compare-feature

- sorting custom comparator feature implemented.

Author: dapper91

src/main.rs

@@ -63,7 +63,15 @@ fn main() {
         }
     };
 
-    let sorted_stream = match sorter.sort(input_stream.lines()) {
+    let compare = |a: &String, b: &String| {
+        if order == Order::Asc {
+            a.cmp(&b)
+        } else {
+            a.cmp(&b).reverse()
+        }
+    };
+
+    let sorted_stream = match sorter.sort_by(input_stream.lines(), compare) {
         Ok(sorted_stream) => sorted_stream,
         Err(err) => {
             log::error!("data sorting error: {}", err);
@@ -115,7 +123,7 @@ impl std::str::FromStr for LogLevel {
     }
 }
 
-#[derive(Copy, Clone, clap::ArgEnum)]
+#[derive(Copy, Clone, PartialEq, clap::ArgEnum)]
 enum Order {
     Asc,
     Desc,

src/merger.rs

@@ -1,36 +1,88 @@
 //! Binary heap merger.
 
+use std::cmp::Ordering;
 use std::collections::BinaryHeap;
 use std::error::Error;
 
+/// Value wrapper binding custom compare function to a value.
+struct OrderedWrapper<T, F>
+where
+    F: Fn(&T, &T) -> Ordering,
+{
+    value: T,
+    compare: F,
+}
+
+impl<T, F> OrderedWrapper<T, F>
+where
+    F: Fn(&T, &T) -> Ordering,
+{
+    fn wrap(value: T, compare: F) -> Self {
+        OrderedWrapper { value, compare }
+    }
+
+    fn unwrap(self) -> T {
+        self.value
+    }
+}
+
+impl<T, F> PartialEq for OrderedWrapper<T, F>
+where
+    F: Fn(&T, &T) -> Ordering,
+{
+    fn eq(&self, other: &Self) -> bool {
+        self.cmp(other) == Ordering::Equal
+    }
+}
+
+impl<T, F> Eq for OrderedWrapper<T, F> where F: Fn(&T, &T) -> Ordering {}
+
+impl<T, F> PartialOrd for OrderedWrapper<T, F>
+where
+    F: Fn(&T, &T) -> Ordering,
+{
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+impl<T, F> Ord for OrderedWrapper<T, F>
+where
+    F: Fn(&T, &T) -> Ordering,
+{
+    fn cmp(&self, other: &Self) -> Ordering {
+        (self.compare)(&self.value, &other.value)
+    }
+}
+
 /// Binary heap merger implementation.
 /// Merges multiple sorted inputs into a single sorted output.
 /// Time complexity is *m* \* log(*n*) in worst case where *m* is the number of items,
 /// *n* is the number of chunks (inputs).
-pub struct BinaryHeapMerger<T, E, C>
+pub struct BinaryHeapMerger<T, E, F, C>
 where
-    T: Ord,
     E: Error,
+    F: Fn(&T, &T) -> Ordering,
     C: IntoIterator<Item = Result<T, E>>,
 {
     // binary heap is max-heap by default so we reverse it to convert it to min-heap
-    items: BinaryHeap<(std::cmp::Reverse<T>, usize)>,
+    items: BinaryHeap<(std::cmp::Reverse<OrderedWrapper<T, F>>, usize)>,
     chunks: Vec<C::IntoIter>,
     initiated: bool,
+    compare: F,
 }
 
-impl<T, E, C> BinaryHeapMerger<T, E, C>
+impl<T, E, F, C> BinaryHeapMerger<T, E, F, C>
 where
-    T: Ord,
     E: Error,
+    F: Fn(&T, &T) -> Ordering,
     C: IntoIterator<Item = Result<T, E>>,
 {
     /// Creates an instance of a binary heap merger using chunks as inputs.
     /// Chunk items should be sorted in ascending order otherwise the result is undefined.
     ///
     /// # Arguments
     /// * `chunks` - Chunks to be merged in a single sorted one
-    pub fn new<I>(chunks: I) -> Self
+    pub fn new<I>(chunks: I, compare: F) -> Self
     where
         I: IntoIterator<Item = C>,
     {
@@ -40,15 +92,16 @@ where
         return BinaryHeapMerger {
             chunks,
             items,
+            compare,
             initiated: false,
         };
     }
 }
 
-impl<T, E, C> Iterator for BinaryHeapMerger<T, E, C>
+impl<T, E, F, C> Iterator for BinaryHeapMerger<T, E, F, C>
 where
-    T: Ord,
     E: Error,
+    F: Fn(&T, &T) -> Ordering + Copy,
     C: IntoIterator<Item = Result<T, E>>,
 {
     type Item = Result<T, E>;
@@ -59,7 +112,9 @@ where
             for (idx, chunk) in self.chunks.iter_mut().enumerate() {
                 if let Some(item) = chunk.next() {
                     match item {
-                        Ok(item) => self.items.push((std::cmp::Reverse(item), idx)),
+                        Ok(item) => self
+                            .items
+                            .push((std::cmp::Reverse(OrderedWrapper::wrap(item, self.compare)), idx)),
                         Err(err) => return Some(Err(err)),
                     }
                 }
@@ -70,12 +125,14 @@ where
         let (result, idx) = self.items.pop()?;
         if let Some(item) = self.chunks[idx].next() {
             match item {
-                Ok(item) => self.items.push((std::cmp::Reverse(item), idx)),
+                Ok(item) => self
+                    .items
+                    .push((std::cmp::Reverse(OrderedWrapper::wrap(item, self.compare)), idx)),
                 Err(err) => return Some(Err(err)),
             }
         }
 
-        return Some(Ok(result.0));
+        return Some(Ok(result.0.unwrap()));
     }
 }
 
@@ -131,7 +188,7 @@ mod test {
         #[case] chunks: Vec<Vec<Result<i32, io::Error>>>,
         #[case] expected_result: Vec<Result<i32, io::Error>>,
     ) {
-        let merger = BinaryHeapMerger::new(chunks);
+        let merger = BinaryHeapMerger::new(chunks, i32::cmp);
         let actual_result = merger.collect();
         assert!(
             compare_vectors_of_result::<_, io::Error>(&actual_result, &expected_result),

src/sort.rs

@@ -1,6 +1,7 @@
 //! External sorter.
 
 use log;
+use std::cmp::Ordering;
 use std::error::Error;
 use std::fmt;
 use std::fmt::{Debug, Display};
@@ -64,7 +65,7 @@ impl<S: Error, D: Error, I: Error> Display for SortError<S, D, I> {
 #[derive(Clone)]
 pub struct ExternalSorterBuilder<T, E, B = LimitedBufferBuilder, C = RmpExternalChunk<T>>
 where
-    T: Ord + Send,
+    T: Send,
     E: Error,
     B: ChunkBufferBuilder<T>,
     C: ExternalChunk<T>,
@@ -88,7 +89,7 @@ where
 
 impl<T, E, B, C> ExternalSorterBuilder<T, E, B, C>
 where
-    T: Ord + Send,
+    T: Send,
     E: Error,
     B: ChunkBufferBuilder<T>,
     C: ExternalChunk<T>,
@@ -137,7 +138,7 @@ where
 
 impl<T, E, B, C> Default for ExternalSorterBuilder<T, E, B, C>
 where
-    T: Ord + Send,
+    T: Send,
     E: Error,
     B: ChunkBufferBuilder<T>,
     C: ExternalChunk<T>,
@@ -158,7 +159,7 @@ where
 /// External sorter.
 pub struct ExternalSorter<T, E, B = LimitedBufferBuilder, C = RmpExternalChunk<T>>
 where
-    T: Ord + Send,
+    T: Send,
     E: Error,
     B: ChunkBufferBuilder<T>,
     C: ExternalChunk<T>,
@@ -182,7 +183,7 @@ where
 
 impl<T, E, B, C> ExternalSorter<T, E, B, C>
 where
-    T: Ord + Send,
+    T: Send,
     E: Error,
     B: ChunkBufferBuilder<T>,
     C: ExternalChunk<T>,
@@ -246,17 +247,42 @@ where
         return Ok(tmp_dir);
     }
 
-    /// Sorts data from input using external sort algorithm.
+    /// Sorts data from the input.
     /// Returns an iterator that can be used to get sorted data stream.
+    ///
+    /// # Arguments
+    /// * `input` - Input stream data to be fetched from
     pub fn sort<I>(
         &self,
         input: I,
     ) -> Result<
-        BinaryHeapMerger<T, C::DeserializationError, C>,
+        BinaryHeapMerger<T, C::DeserializationError, impl Fn(&T, &T) -> Ordering + Copy, C>,
+        SortError<C::SerializationError, C::DeserializationError, E>,
+    >
+    where
+        T: Ord,
+        I: IntoIterator<Item = Result<T, E>>,
+    {
+        self.sort_by(input, T::cmp)
+    }
+
+    /// Sorts data from the input using a custom compare function.
+    /// Returns an iterator that can be used to get sorted data stream.
+    ///
+    /// # Arguments
+    /// * `input` - Input stream data to be fetched from
+    /// * `compare` - Function be be used to compare items
+    pub fn sort_by<I, F>(
+        &self,
+        input: I,
+        compare: F,
+    ) -> Result<
+        BinaryHeapMerger<T, C::DeserializationError, F, C>,
         SortError<C::SerializationError, C::DeserializationError, E>,
     >
     where
         I: IntoIterator<Item = Result<T, E>>,
+        F: Fn(&T, &T) -> Ordering + Sync + Send + Copy,
     {
         let mut chunk_buf = self.buffer_builder.build();
         let mut external_chunks = Vec::new();
@@ -268,34 +294,39 @@ where
             }
 
             if chunk_buf.is_full() {
-                external_chunks.push(self.create_chunk(chunk_buf)?);
+                external_chunks.push(self.create_chunk(chunk_buf, compare)?);
                 chunk_buf = self.buffer_builder.build();
             }
         }
 
         if chunk_buf.len() > 0 {
-            external_chunks.push(self.create_chunk(chunk_buf)?);
+            external_chunks.push(self.create_chunk(chunk_buf, compare)?);
         }
 
         log::debug!("external sort preparation done");
 
-        return Ok(BinaryHeapMerger::new(external_chunks));
+        return Ok(BinaryHeapMerger::new(external_chunks, compare));
     }
 
-    fn create_chunk(
+    fn create_chunk<F>(
         &self,
-        mut chunk: impl ChunkBuffer<T>,
-    ) -> Result<C, SortError<C::SerializationError, C::DeserializationError, E>> {
+        mut buffer: impl ChunkBuffer<T>,
+        compare: F,
+    ) -> Result<C, SortError<C::SerializationError, C::DeserializationError, E>>
+    where
+        F: Fn(&T, &T) -> Ordering + Sync + Send,
+    {
         log::debug!("sorting chunk data ...");
         self.thread_pool.install(|| {
-            chunk.par_sort();
+            buffer.par_sort_by(compare);
         });
 
         log::debug!("saving chunk data");
-        let external_chunk = ExternalChunk::build(&self.tmp_dir, chunk, self.rw_buf_size).map_err(|err| match err {
-            ExternalChunkError::IO(err) => SortError::IO(err),
-            ExternalChunkError::SerializationError(err) => SortError::SerializationError(err),
-        })?;
+        let external_chunk =
+            ExternalChunk::build(&self.tmp_dir, buffer, self.rw_buf_size).map_err(|err| match err {
+                ExternalChunkError::IO(err) => SortError::IO(err),
+                ExternalChunkError::SerializationError(err) => SortError::SerializationError(err),
+            })?;
 
         return Ok(external_chunk);
     }
@@ -307,11 +338,14 @@ mod test {
     use std::path::Path;
 
     use rand::seq::SliceRandom;
+    use rstest::*;
 
     use super::{ExternalSorter, ExternalSorterBuilder, LimitedBufferBuilder};
 
-    #[test]
-    fn test_external_sorter() {
+    #[rstest]
+    #[case(false)]
+    #[case(true)]
+    fn test_external_sorter(#[case] reversed: bool) {
         let input_sorted = 0..100;
 
         let mut input: Vec<Result<i32, io::Error>> = Vec::from_iter(input_sorted.clone().map(|item| Ok(item)));
@@ -324,11 +358,21 @@ mod test {
             .build()
             .unwrap();
 
-        let result = sorter.sort(input).unwrap();
+        let compare = if reversed {
+            |a: &i32, b: &i32| a.cmp(b).reverse()
+        } else {
+            |a: &i32, b: &i32| a.cmp(b)
+        };
+
+        let result = sorter.sort_by(input, compare).unwrap();
 
         let actual_result: Result<Vec<i32>, _> = result.collect();
         let actual_result = actual_result.unwrap();
-        let expected_result = Vec::from_iter(input_sorted.clone());
+        let expected_result = if reversed {
+            Vec::from_iter(input_sorted.clone().rev())
+        } else {
+            Vec::from_iter(input_sorted.clone())
+        };
 
         assert_eq!(actual_result, expected_result)
     }