<commit_before>package org.apache.prepbuddy.datacleansers.imputation;

import org.apache.prepbuddy.SparkTestCase;
import org.apache.prepbuddy.rdds.TransformableRDD;
import org.apache.spark.api.java.JavaRDD;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;

import static org.junit.Assert.assertTrue;

public class ImputationTest extends SparkTestCase {
    @Test
    public void shouldImputeTheValueWithTheMean() {
        JavaRDD<String> initialDataSet = javaSparkContext.parallelize(Arrays.asList(
                "07434677419,07371326239,Incoming,31,Wed Sep 15 19:17:44 +0100 2010",
                "07641036117,01666472054,Outgoing,20,Mon Feb 11 07:18:23 +0000 1980",
                "07641036117,07371326239,Incoming, ,Mon Feb 11 07:45:42 +0000 1980",
                "07641036117,07681546436,Missed,12,Mon Feb 11 08:04:42 +0000 1980"

        ));
        TransformableRDD initialRDD = new TransformableRDD(initialDataSet);
        TransformableRDD imputed = initialRDD.impute(3, new MeanSubstitution());
        List<String> listOfRecord = imputed.collect();

        String expected1 = "07641036117,07371326239,Incoming,15.75,Mon Feb 11 07:45:42 +0000 1980";
        assertTrue(listOfRecord.contains(expected1));
    }

    @Test
    public void shouldImputeTheValueWithTheMeanApprox() {
        JavaRDD<String> initialDataSet = javaSparkContext.parallelize(Arrays.asList(
                "07434677419,07371326239,Incoming,31,Wed Sep 15 19:17:44 +0100 2010",
                "07641036117,01666472054,Outgoing,20,Mon Feb 11 07:18:23 +0000 1980",
                "07641036117,07371326239,Incoming, ,Mon Feb 11 07:45:42 +0000 1980",
                "07641036117,07681546436,Missed,12,Mon Feb 11 08:04:42 +0000 1980"

        ));
        TransformableRDD initialRDD = new TransformableRDD(initialDataSet);
        TransformableRDD imputed = initialRDD.impute(3, new ApproxMeanSubstitution());
        List<String> listOfRecord = imputed.collect();

        String expected1 = "07641036117,07371326239,Incoming,15.75,Mon Feb 11 07:45:42 +0000 1980";
        assertTrue(listOfRecord.contains(expected1));
    }

    @Test
    public void shouldImputeTheValueWithTheMostOccurredValue() throws Exception {
        JavaRDD<String> initialDataSet = javaSparkContext.parallelize(Arrays.asList(
                "07434677419,07371326239,Incoming,31,Wed Sep 15 19:17:44 +0100 2010",
                "07641036117,01666472054,Outgoing,31,Mon Feb 11 07:18:23 +0000 1980",
                "07641036117,07371326239,Incoming, ,Mon Feb 11 07:45:42 +0000 1980",
                "07641036117,07681546436,Missed,12,Mon Feb 11 08:04:42 +0000 1980"

        ));
        TransformableRDD initialRDD = new TransformableRDD(initialDataSet);
        TransformableRDD imputed = initialRDD.impute(3, new ModeSubstitution());
        List<String> listOfRecord = imputed.collect();
        String expected1 = "07641036117,07371326239,Incoming,31,Mon Feb 11 07:45:42 +0000 1980";
        assertTrue(listOfRecord.contains(expected1));
    }

    @Test
    public void shouldImputeTheValueWithTheRegression() throws Exception {
        JavaRDD<String> initialDataSet = javaSparkContext.parallelize(Arrays.asList(
                "60,3.1", "61,3.6", "62,3.8", "63,4", "65,4.1", "64,"
        ));
        TransformableRDD initialRDD = new TransformableRDD(initialDataSet);
        TransformableRDD imputed = initialRDD.impute(1, new UnivariateLinearRegressionSubstitution(0));
        List<String> listOfRecord = imputed.collect();

        String expected1 = "64,4.06";
        assertTrue(listOfRecord.contains(expected1));
    }

    @Test
    public void shouldImputeTheValueWithTheRegressionTest2() throws Exception {
        JavaRDD<String> initialDataSet = javaSparkContext.parallelize(Arrays.asList(
                "3.4,5.67", "3.9,4.81", "2.6,4.93", "1.9,6.21", "2.2,6.83", "3.3,5.61", "1.7,5.45", "2.4,4.94", "2.8,5.73",
                "3.6,"
        ));
        TransformableRDD initialRDD = new TransformableRDD(initialDataSet);
        TransformableRDD imputed = initialRDD.impute(1, new UnivariateLinearRegressionSubstitution(0));
        List<String> listOfRecord = imputed.collect();

        String expected1 = "3.6,5.24";
        assertTrue(listOfRecord.contains(expected1));
    }

    @Test
    public void shouldImputeMissingValuesWithTheNaiveBayes() {
        JavaRDD<String> initialDataSet = javaSparkContext.parallelize(Arrays.asList(
                "known,new,long,home,skips",
                "unknown,new,short,work,reads",
                "unknown,follow Up,long,work,skips",
                "known,follow Up,long,home,skips",
                "known,new,short,home,reads",
                "known,follow Up,long,work,skips",
                "unknown,follow Up,short,work,skips",
                "unknown,new,short,work,reads",
                "known,follow Up,long,home,skips",
                "known,new,long,work,skips",
                "unknown,follow Up,short,home,skips",
                "known,new,long,work,skips",
                "known,follow Up,short,home,reads",
                "known,new,short,work,reads",
                "known,new,short,home,reads",
                "known,follow Up,short,work,reads",
                "known,new,short,home,reads",
                "unknown,new,short,work,reads",
                "unknown,new,long,work,",
                "unknown,follow Up,long,home,"
        ));
        TransformableRDD initialRDD = new TransformableRDD(initialDataSet);
        TransformableRDD imputed = initialRDD.impute(4, new NaiveBayesSubstitution(0, 1, 2, 3));
        List<String> listOfRecord = imputed.collect();

        String expected1 = "unknown,new,long,work,skips";
        assertTrue(listOfRecord.contains(expected1));
    }


    @Test
    public void shouldImputeMissingValuesWithTheNaiveBayesOnAnyDataset() {
        JavaRDD<String> initialDataSet = javaSparkContext.parallelize(Arrays.asList(
                "sunny,hot,high,false,N",
                "sunny,hot,high,true,N",
                "overcast,hot,high,false,P",
                "rain,mild,high,false,P",
                "rain,cool,normal,false,P",
                "rain,cool,normal,true,N",
                "overcast,cool,normal,true,P",
                "sunny,mild,high,false,N",
                "sunny,cool,normal,false,P",
                "rain,mild,normal,false,P",
                "sunny,mild,normal,true,P",
                "overcast,mild,high,true,P",
                "overcast,hot,normal,false,P",
                "rain,mild,high,true,N",
                "sunny,cool,high,false,"
        ));
        TransformableRDD initialRDD = new TransformableRDD(initialDataSet);
        TransformableRDD imputed = initialRDD.impute(4, new NaiveBayesSubstitution(0, 1, 2, 3));
        List<String> listOfRecord = imputed.collect();

        System.out.println("listOfRecord = " + listOfRecord);
        String expected1 = "sunny,cool,high,false,N";
        assertTrue(listOfRecord.contains(expected1));

    }
}<commit_msg>added one more expectation<commit_after>package org.apache.prepbuddy.datacleansers.imputation;

import org.apache.prepbuddy.SparkTestCase;
import org.apache.prepbuddy.rdds.TransformableRDD;
import org.apache.spark.api.java.JavaRDD;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;

import static org.junit.Assert.assertTrue;

public class ImputationTest extends SparkTestCase {
    @Test
    public void shouldImputeTheValueWithTheMean() {
        JavaRDD<String> initialDataSet = javaSparkContext.parallelize(Arrays.asList(
                "07434677419,07371326239,Incoming,31,Wed Sep 15 19:17:44 +0100 2010",
                "07641036117,01666472054,Outgoing,20,Mon Feb 11 07:18:23 +0000 1980",
                "07641036117,07371326239,Incoming, ,Mon Feb 11 07:45:42 +0000 1980",
                "07641036117,07681546436,Missed,12,Mon Feb 11 08:04:42 +0000 1980"

        ));
        TransformableRDD initialRDD = new TransformableRDD(initialDataSet);
        TransformableRDD imputed = initialRDD.impute(3, new MeanSubstitution());
        List<String> listOfRecord = imputed.collect();

        String expected1 = "07641036117,07371326239,Incoming,15.75,Mon Feb 11 07:45:42 +0000 1980";
        assertTrue(listOfRecord.contains(expected1));
    }

    @Test
    public void shouldImputeTheValueWithTheMeanApprox() {
        JavaRDD<String> initialDataSet = javaSparkContext.parallelize(Arrays.asList(
                "07434677419,07371326239,Incoming,31,Wed Sep 15 19:17:44 +0100 2010",
                "07641036117,01666472054,Outgoing,20,Mon Feb 11 07:18:23 +0000 1980",
                "07641036117,07371326239,Incoming, ,Mon Feb 11 07:45:42 +0000 1980",
                "07641036117,07681546436,Missed,12,Mon Feb 11 08:04:42 +0000 1980"

        ));
        TransformableRDD initialRDD = new TransformableRDD(initialDataSet);
        TransformableRDD imputed = initialRDD.impute(3, new ApproxMeanSubstitution());
        List<String> listOfRecord = imputed.collect();

        String expected1 = "07641036117,07371326239,Incoming,15.75,Mon Feb 11 07:45:42 +0000 1980";
        assertTrue(listOfRecord.contains(expected1));
    }

    @Test
    public void shouldImputeTheValueWithTheMostOccurredValue() throws Exception {
        JavaRDD<String> initialDataSet = javaSparkContext.parallelize(Arrays.asList(
                "07434677419,07371326239,Incoming,31,Wed Sep 15 19:17:44 +0100 2010",
                "07641036117,01666472054,Outgoing,31,Mon Feb 11 07:18:23 +0000 1980",
                "07641036117,07371326239,Incoming, ,Mon Feb 11 07:45:42 +0000 1980",
                "07641036117,07681546436,Missed,12,Mon Feb 11 08:04:42 +0000 1980"

        ));
        TransformableRDD initialRDD = new TransformableRDD(initialDataSet);
        TransformableRDD imputed = initialRDD.impute(3, new ModeSubstitution());
        List<String> listOfRecord = imputed.collect();
        String expected1 = "07641036117,07371326239,Incoming,31,Mon Feb 11 07:45:42 +0000 1980";
        assertTrue(listOfRecord.contains(expected1));
    }

    @Test
    public void shouldImputeTheValueWithTheRegression() throws Exception {
        JavaRDD<String> initialDataSet = javaSparkContext.parallelize(Arrays.asList(
                "60,3.1", "61,3.6", "62,3.8", "63,4", "65,4.1", "64,"
        ));
        TransformableRDD initialRDD = new TransformableRDD(initialDataSet);
        TransformableRDD imputed = initialRDD.impute(1, new UnivariateLinearRegressionSubstitution(0));
        List<String> listOfRecord = imputed.collect();

        String expected1 = "64,4.06";
        assertTrue(listOfRecord.contains(expected1));
    }

    @Test
    public void shouldImputeTheValueWithTheRegressionTest2() throws Exception {
        JavaRDD<String> initialDataSet = javaSparkContext.parallelize(Arrays.asList(
                "3.4,5.67", "3.9,4.81", "2.6,4.93", "1.9,6.21", "2.2,6.83", "3.3,5.61", "1.7,5.45", "2.4,4.94", "2.8,5.73",
                "3.6,"
        ));
        TransformableRDD initialRDD = new TransformableRDD(initialDataSet);
        TransformableRDD imputed = initialRDD.impute(1, new UnivariateLinearRegressionSubstitution(0));
        List<String> listOfRecord = imputed.collect();

        String expected1 = "3.6,5.24";
        assertTrue(listOfRecord.contains(expected1));
    }

    @Test
    public void shouldImputeMissingValuesWithTheNaiveBayes() {
        JavaRDD<String> initialDataSet = javaSparkContext.parallelize(Arrays.asList(
                "known,new,long,home,skips",
                "unknown,new,short,work,reads",
                "unknown,follow Up,long,work,skips",
                "known,follow Up,long,home,skips",
                "known,new,short,home,reads",
                "known,follow Up,long,work,skips",
                "unknown,follow Up,short,work,skips",
                "unknown,new,short,work,reads",
                "known,follow Up,long,home,skips",
                "known,new,long,work,skips",
                "unknown,follow Up,short,home,skips",
                "known,new,long,work,skips",
                "known,follow Up,short,home,reads",
                "known,new,short,work,reads",
                "known,new,short,home,reads",
                "known,follow Up,short,work,reads",
                "known,new,short,home,reads",
                "unknown,new,short,work,reads",
                "unknown,new,long,work,",
                "unknown,follow Up,long,home,"
        ));
        TransformableRDD initialRDD = new TransformableRDD(initialDataSet);
        TransformableRDD imputed = initialRDD.impute(4, new NaiveBayesSubstitution(0, 1, 2, 3));
        List<String> listOfRecord = imputed.collect();

        String expected1 = "unknown,new,long,work,skips";
        assertTrue(listOfRecord.contains(expected1));
    }


    @Test
    public void shouldImputeMissingValuesWithTheNaiveBayesOnAnyDataset() {
        JavaRDD<String> initialDataSet = javaSparkContext.parallelize(Arrays.asList(
                "sunny,hot,high,false,N",
                "sunny,hot,high,true,N",
                "overcast,hot,high,false,P",
                "rain,mild,high,false,P",
                "rain,cool,normal,false,P",
                "rain,cool,normal,true,N",
                "overcast,cool,normal,true,P",
                "sunny,mild,high,false,N",
                "sunny,cool,normal,false,P",
                "rain,mild,normal,false,P",
                "sunny,mild,normal,true,P",
                "overcast,mild,high,true,P",
                "overcast,hot,normal,false,P",
                "rain,mild,high,true,N",
                "sunny,cool,high,false,",
                "rain,hot,high,false,"
        ));
        TransformableRDD initialRDD = new TransformableRDD(initialDataSet);
        TransformableRDD imputed = initialRDD.impute(4, new NaiveBayesSubstitution(0, 1, 2, 3));
        List<String> listOfRecord = imputed.collect();

        System.out.println("listOfRecord = " + listOfRecord);
        String expected1 = "sunny,cool,high,false,N";
        String expected2 = "rain,hot,high,false,N";

        assertTrue(listOfRecord.contains(expected1));
        assertTrue(listOfRecord.contains(expected2));

    }
}
{'fraction_non_alphanumeric': 0.11978361669242658, 'fraction_numerical': 0.09962062666854012, 'mean_word_length': 2.9021381578947367, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 31, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2500784', 'n_tokens_mistral': 5566, 'n_tokens_neox': 4507, 'n_words': 702}