burn-dataset 0.22.0-pre.2

use crate::Dataset;
use std::error::Error;
use std::{marker::PhantomData, sync::Arc};

/// Only use a fraction of an existing dataset lazily.
#[derive(new, Clone)]
pub struct PartialDataset<D, I> {
    dataset: D,
    start_index: usize,
    end_index: usize,
    input: PhantomData<I>,
}

impl<D, I> PartialDataset<D, I>
where
    D: Dataset<I>,
{
    /// Splits a dataset into multiple partial datasets.
    pub fn split(dataset: D, num: usize) -> Vec<PartialDataset<Arc<D>, I>> {
        let dataset = Arc::new(dataset); // cheap cloning.

        let mut current = 0;
        let mut datasets = Vec::with_capacity(num);

        let batch_size = dataset.len() / num;

        for i in 0..num {
            let start = current;
            let mut end = current + batch_size;

            if i == (num - 1) {
                end = dataset.len();
            }

            let dataset = PartialDataset::new(dataset.clone(), start, end);

            current += batch_size;
            datasets.push(dataset);
        }

        datasets
    }

    /// Splits a dataset by distributing complete chunks/batches across multiple partial datasets.
    pub fn split_chunks(
        dataset: D,
        num: usize,
        batch_size: usize,
    ) -> Vec<PartialDataset<Arc<D>, I>> {
        let dataset = Arc::new(dataset); // cheap cloning.
        let total_items = dataset.len();

        // Total number of complete batches
        let total_batches = total_items.div_ceil(batch_size);
        let batches_per_split = total_batches / num;
        let extra_batches = total_batches % num;

        let mut datasets = Vec::with_capacity(num);
        let mut current_batch = 0;

        for i in 0..num {
            // Extra batches distributed across first splits
            let split_batches = if i < extra_batches {
                batches_per_split + 1
            } else {
                batches_per_split
            };

            let start_batch = current_batch;
            let end_batch = start_batch + split_batches;

            let start_index = start_batch * batch_size;
            let end_index = core::cmp::min(end_batch * batch_size, total_items);

            if start_index < total_items {
                datasets.push(PartialDataset::new(dataset.clone(), start_index, end_index));
            }

            current_batch = end_batch;
        }

        datasets
    }
}

impl<D, I, E> Dataset<I, E> for PartialDataset<D, I>
where
    D: Dataset<I, E>,
    I: Clone + Send + Sync,
    E: Error + Send + Sync + 'static,
{
    fn get(&self, index: usize) -> Result<I, E> {
        let index = index + self.start_index;
        assert!(
            index >= self.start_index && index < self.end_index,
            "Index out of bounds for PartialDataset: {} >= {}",
            index,
            self.end_index
        );
        self.dataset.get(index)
    }

    fn get_many(&self, indexes: Vec<usize>) -> Result<Vec<I>, E> {
        let translated = indexes
            .into_iter()
            .map(|index| {
                let index = index + self.start_index;
                assert!(
                    index >= self.start_index && index < self.end_index,
                    "Index out of bounds for PartialDataset: {} >= {}",
                    index,
                    self.end_index
                );
                index
            })
            .collect();
        self.dataset.get_many(translated)
    }

    fn len(&self) -> usize {
        usize::min(self.end_index - self.start_index, self.dataset.len())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::FakeDataset;
    use std::collections::HashSet;

    #[test]
    fn test_start_from_beginning() {
        let dataset_original = FakeDataset::<String>::new(27);
        let mut items_original_1 = HashSet::new();
        let mut items_original_2 = HashSet::new();
        let mut items_partial = HashSet::new();
        dataset_original
            .iter()
            .map(Result::unwrap)
            .enumerate()
            .for_each(|(i, item)| {
                match i >= 10 {
                    true => items_original_2.insert(item),
                    false => items_original_1.insert(item),
                };
            });

        let dataset_partial = PartialDataset::new(dataset_original, 0, 10);

        for item in dataset_partial.iter().map(Result::unwrap) {
            items_partial.insert(item);
        }

        assert_eq!(dataset_partial.len(), 10);
        assert_eq!(items_original_1, items_partial);
        for item in items_original_2 {
            assert!(!items_partial.contains(&item));
        }
    }

    #[test]
    fn test_start_inside() {
        let dataset_original = FakeDataset::<String>::new(27);
        let mut items_original_1 = HashSet::new();
        let mut items_original_2 = HashSet::new();
        let mut items_partial = HashSet::new();

        dataset_original
            .iter()
            .map(Result::unwrap)
            .enumerate()
            .for_each(|(i, item)| {
                match !(10..20).contains(&i) {
                    true => items_original_2.insert(item),
                    false => items_original_1.insert(item),
                };
            });

        let dataset_partial = PartialDataset::new(dataset_original, 10, 20);
        for item in dataset_partial.iter().map(Result::unwrap) {
            items_partial.insert(item);
        }

        assert_eq!(dataset_partial.len(), 10);
        assert_eq!(items_original_1, items_partial);
        for item in items_original_2 {
            assert!(!items_partial.contains(&item));
        }
    }

    #[test]
    fn test_get_many() {
        let dataset_original = FakeDataset::<String>::new(27);
        let source_items: Vec<String> = dataset_original.iter().map(Result::unwrap).collect();

        let dataset_partial = PartialDataset::new(dataset_original, 10, 20);

        // Local indices, out of order and with a duplicate.
        let requested = vec![5, 0, 3, 0];
        let expected: Vec<String> = requested
            .iter()
            .map(|&i| source_items[10 + i].clone())
            .collect();

        let items = dataset_partial.get_many(requested).unwrap();

        assert_eq!(items, expected);
    }

    #[test]
    #[should_panic(expected = "Index out of bounds for PartialDataset: 20 >= 20")]
    fn test_get_many_out_of_bounds() {
        let dataset_original = FakeDataset::<String>::new(27);
        let dataset_partial = PartialDataset::new(dataset_original, 10, 20);

        let _ = dataset_partial.get_many(vec![0, 10]);
    }

    #[test]
    fn test_split_contains_all_items_without_duplicates() {
        let dataset_original = FakeDataset::<String>::new(27);
        let mut items_original = Vec::new();
        let mut items_partial = Vec::new();
        for item in dataset_original.iter().map(Result::unwrap) {
            items_original.push(item);
        }

        let dataset_partials = PartialDataset::split(dataset_original, 4);
        let expected_len = [6, 6, 6, 9];

        for (i, dataset) in dataset_partials.iter().enumerate() {
            assert_eq!(dataset.len(), expected_len[i]);
            for item in dataset.iter().map(Result::unwrap) {
                items_partial.push(item);
            }
        }

        assert_eq!(items_original, items_partial);
    }

    #[test]
    fn test_split_chunks_contains_all_items_without_duplicates() {
        let dataset_original = FakeDataset::<String>::new(27);
        let mut items_original = Vec::new();
        let mut items_partial = Vec::new();
        for item in dataset_original.iter().map(Result::unwrap) {
            items_original.push(item);
        }

        let dataset_partials = PartialDataset::split_chunks(dataset_original, 4, 5);
        // [(2 * 5), (2 * 5), 5, 2] -> 5 complete chunks + 1 incomplete with 2 remaining items
        // OTOH, `split(dataset, 4)` would yield [6, 6, 6, 9] -> 4 incomplete chunks + 4 incomplete with [1, 1, 1, 4]
        let expected_len = [10, 10, 5, 2];

        for (i, dataset) in dataset_partials.iter().enumerate() {
            assert_eq!(dataset.len(), expected_len[i]);
            for item in dataset.iter().map(Result::unwrap) {
                items_partial.push(item);
            }
        }

        assert_eq!(items_original, items_partial);
    }
}