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

import org.apache.prepbuddy.groupingops.TextFacets;
import org.apache.prepbuddy.rdds.TransformableRDD;
import org.apache.prepbuddy.utils.*;
import scala.Tuple2;

import java.io.Serializable;
import java.util.List;

public class NaiveBayesClassifier implements Serializable {
    private PivotTable<Probability> probs;
    private int missingDataColumn;
    private int[] otherColumns;
    private List<String> permissibleValues;

    public NaiveBayesClassifier(int... columnIndexes) {
        this.otherColumns = columnIndexes;
    }

    public NaiveBayesClassifier(int missingDataColumn, int[] otherColumns) {
        this.missingDataColumn = missingDataColumn;
        this.otherColumns = otherColumns;
    }

    public void train(TransformableRDD trainingSet) {
        TextFacets facets = trainingSet.listFacets(missingDataColumn);
        permissibleValues = facets.cardinalValues();
        List<Tuple2<String, Integer>> rowKeys = facets.rdd().collect();


        PivotTable<Integer> frequencyTable = trainingSet.pivotByCount(missingDataColumn, otherColumns);
        final long totalRows = trainingSet.count();
        probs = (PivotTable<Probability>) frequencyTable.transform(new TransformationFunction<Integer, Probability>() {
            public Probability transform(Integer input) {
                Probability probability = new Probability((double) input.intValue() / totalRows);
                return probability;
            }

            @Override
            public Probability defaultValue() {
                return new Probability(0);
            }
        });<commit_msg>Refactoring<commit_after>package org.apache.prepbuddy.datacleansers.imputation;

import org.apache.prepbuddy.groupingops.TextFacets;
import org.apache.prepbuddy.rdds.TransformableRDD;
import org.apache.prepbuddy.utils.*;
import scala.Tuple2;

import java.io.Serializable;
import java.util.List;

public class NaiveBayesClassifier implements Serializable {
    private PivotTable<Probability> probs;
    private int missingDataColumn;
    private int[] otherColumns;
    private List<String> permissibleValues;

    public NaiveBayesClassifier(int... columnIndexes) {
        this.otherColumns = columnIndexes;
    }

    public NaiveBayesClassifier(int missingDataColumn, int[] otherColumns) {
        this.missingDataColumn = missingDataColumn;
        this.otherColumns = otherColumns;
    }

    public void train(TransformableRDD trainingSet) {
        TextFacets facets = trainingSet.listFacets(missingDataColumn);
        permissibleValues = facets.cardinalValues();
        List<Tuple2<String, Integer>> rowKeys = facets.rdd().collect();


        PivotTable<Integer> frequencyTable = trainingSet.pivotByCount(missingDataColumn, otherColumns);
        final long totalRows = trainingSet.count();
        probs = (PivotTable<Probability>) frequencyTable.transform(new TransformationFunction<Integer, Probability>() {
            public Probability transform(Integer input) {
                double value = (double) input.intValue() / totalRows;
                Probability probability = new Probability(value);
                return probability;
            }

            @Override
            public Probability defaultValue() {
                return new Probability(0);
            }
        });
{'fraction_non_alphanumeric': 0.08590047393364929, 'fraction_numerical': 0.0017772511848341231, 'mean_word_length': 3.312899106002554, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, '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': '5835722', 'n_tokens_mistral': 930, 'n_tokens_neox': 878, 'n_words': 183}