refactor(data): extract shared utilities for class balancing and CSV writing

vxa8502 · vxa8502 · commit f38047b7bcaa · 2026-01-28T11:32:21.000-06:00
diff --git a/src/main/java/sentiment/data/DataPreparer.java b/src/main/java/sentiment/data/DataPreparer.java
@@ -4,14 +4,12 @@
 import org.slf4j.LoggerFactory;
 import sentiment.evaluation.StratifiedDataSplitter;
 
-import java.io.FileWriter;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.time.Instant;
 import java.util.*;
-import java.util.stream.Collectors;
 
 /**
  * Prepares and manages immutable data splits for training.
@@ -165,8 +163,10 @@ private SplitManifest createSplits(String domain, String rawDataPath, int maxSam
         Path trainPath = processedDir.resolve("train.csv");
         Path testPath = processedDir.resolve("test.csv");
 
-        saveSplitToCsv(split.train, trainPath);
-        saveSplitToCsv(split.test, testPath);
+        SimpleDatasetLoader.writeCsv(split.train, trainPath);
+        logger.info("Saved {} samples to {}", split.train.size(), trainPath);
+        SimpleDatasetLoader.writeCsv(split.test, testPath);
+        logger.info("Saved {} samples to {}", split.test.size(), testPath);
 
         // Step 4: Create and save manifest
         logger.info("Step 4/4: Creating manifest with checksums...");
@@ -221,74 +221,8 @@ private List<Dataset> loadAndBalance(String rawDataPath, int maxSamples, int see
         List<Dataset> shuffled = new ArrayList<>(allData);
         Collections.shuffle(shuffled, new Random(seed));
 
-        // Balance classes by undersampling majority
-        return balanceClasses(shuffled, seed);
-    }
-
-    /**
-     * Balance classes by undersampling the majority class.
-     */
-    private List<Dataset> balanceClasses(List<Dataset> data, int seed) {
-        List<Dataset> positive = data.stream()
-                .filter(d -> d.getSentiment() == Dataset.SentimentLabel.POSITIVE)
-                .collect(Collectors.toList());
-        List<Dataset> negative = data.stream()
-                .filter(d -> d.getSentiment() == Dataset.SentimentLabel.NEGATIVE)
-                .collect(Collectors.toList());
-
-        int minoritySize = Math.min(positive.size(), negative.size());
-
-        if (minoritySize == 0) {
-            logger.warn("One class has zero samples - cannot balance");
-            return data;
-        }
-
-        // Check if already balanced (within 5% tolerance)
-        double ratio = (double) Math.min(positive.size(), negative.size()) /
-                       Math.max(positive.size(), negative.size());
-        if (ratio >= 0.95) {
-            logger.info("Classes already balanced (ratio: {}) - no undersampling needed",
-                    String.format("%.2f", ratio));
-            return data;
-        }
-
-        // Undersample majority class
-        List<Dataset> balancedPositive = positive.subList(0, minoritySize);
-        List<Dataset> balancedNegative = negative.subList(0, minoritySize);
-
-        List<Dataset> balanced = new ArrayList<>(minoritySize * 2);
-        balanced.addAll(balancedPositive);
-        balanced.addAll(balancedNegative);
-        Collections.shuffle(balanced, new Random(seed));
-
-        logger.info("Balanced classes: {} -> {} samples (undersampled {} from majority)",
-                data.size(), balanced.size(), data.size() - balanced.size());
-
-        return balanced;
-    }
-
-    /**
-     * Save a data split to CSV file.
-     */
-    private void saveSplitToCsv(List<Dataset> data, Path filePath) throws IOException {
-        try (FileWriter writer = new FileWriter(filePath.toFile())) {
-            writer.write("review,sentiment\n");
-
-            for (Dataset sample : data) {
-                String sentiment = sample.getSentiment().name().toLowerCase();
-                String text = escapeCsv(sample.getText());
-                writer.write("\"" + text + "\"," + sentiment + "\n");
-            }
-        }
-        logger.info("Saved {} samples to {}", data.size(), filePath);
-    }
-
-    /**
-     * Escape CSV special characters in text.
-     */
-    private String escapeCsv(String text) {
-        if (text == null) return "";
-        return text.replace("\"", "\"\"");
+        // Balance classes by undersampling majority (uses shared utility)
+        return DatasetUtils.balanceByUndersampling(shuffled, seed);
     }
 
     // ===== CLI Entry Point =====
diff --git a/src/main/java/sentiment/data/DatasetUtils.java b/src/main/java/sentiment/data/DatasetUtils.java
@@ -0,0 +1,89 @@
+package sentiment.data;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Random;
+import java.util.stream.Collectors;
+
+/**
+ * Utility methods for dataset manipulation.
+ * Provides shared operations used across data preparation and training pipelines.
+ */
+public final class DatasetUtils {
+
+    private static final Logger logger = LoggerFactory.getLogger(DatasetUtils.class);
+
+    /**
+     * Tolerance threshold for considering classes "already balanced".
+     * If minority/majority ratio >= this value, no undersampling is performed.
+     */
+    private static final double BALANCE_TOLERANCE = 0.95;
+
+    private DatasetUtils() {
+        // Utility class - prevent instantiation
+    }
+
+    /**
+     * Balance classes by undersampling the majority class to match the minority class.
+     * This ensures fair model comparison across datasets with different class distributions.
+     *
+     * The method:
+     * 1. Separates data by sentiment class
+     * 2. If already balanced (within 5% tolerance), returns original data
+     * 3. Otherwise, undersamples majority class to match minority class size
+     * 4. Shuffles the result to mix classes
+     *
+     * @param data input dataset (may be imbalanced)
+     * @param seed random seed for reproducible shuffling
+     * @return balanced dataset with equal positive/negative samples
+     */
+    public static List<Dataset> balanceByUndersampling(List<Dataset> data, long seed) {
+        if (data == null || data.isEmpty()) {
+            return data;
+        }
+
+        // Separate by class
+        List<Dataset> positive = data.stream()
+                .filter(d -> d.getSentiment() == Dataset.SentimentLabel.POSITIVE)
+                .collect(Collectors.toList());
+        List<Dataset> negative = data.stream()
+                .filter(d -> d.getSentiment() == Dataset.SentimentLabel.NEGATIVE)
+                .collect(Collectors.toList());
+
+        // Find minority class size
+        int minoritySize = Math.min(positive.size(), negative.size());
+
+        if (minoritySize == 0) {
+            logger.warn("One class has zero samples - cannot balance. Returning original data.");
+            return data;
+        }
+
+        // Check if already balanced (within tolerance)
+        double ratio = (double) minoritySize / Math.max(positive.size(), negative.size());
+        if (ratio >= BALANCE_TOLERANCE) {
+            logger.info("Classes already balanced (ratio: {}) - no undersampling needed",
+                    String.format("%.2f", ratio));
+            return data;
+        }
+
+        // Undersample majority class
+        List<Dataset> balancedPositive = positive.subList(0, minoritySize);
+        List<Dataset> balancedNegative = negative.subList(0, minoritySize);
+
+        // Combine and shuffle to mix classes
+        List<Dataset> balanced = new ArrayList<>(minoritySize * 2);
+        balanced.addAll(balancedPositive);
+        balanced.addAll(balancedNegative);
+        Collections.shuffle(balanced, new Random(seed));
+
+        logger.info("Balanced classes: {} -> {} samples (undersampled {} from majority class)",
+                data.size(), balanced.size(), data.size() - balanced.size());
+        logger.info("Final distribution: positive={}, negative={}", minoritySize, minoritySize);
+
+        return balanced;
+    }
+}
diff --git a/src/main/java/sentiment/data/SimpleDatasetLoader.java b/src/main/java/sentiment/data/SimpleDatasetLoader.java
@@ -9,6 +9,7 @@
 import sentiment.util.ValidationUtils;
 
 import java.io.BufferedReader;
+import java.io.BufferedWriter;
 import java.io.FileReader;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
@@ -397,6 +398,44 @@ private List<Dataset> loadPlaintext(String filePath) throws DataLoadingException
         );
     }
 
+    // CSV WRITING
+
+    /**
+     * Write datasets to CSV file in the standard format used by this loader.
+     * Format: "review,sentiment" with quoted text field and lowercase sentiment.
+     *
+     * This is the inverse of loadCsv() - keeping read/write together ensures format consistency.
+     *
+     * @param data list of datasets to write
+     * @param path destination file path
+     * @throws IOException if file cannot be written
+     */
+    public static void writeCsv(List<Dataset> data, Path path) throws IOException {
+        try (BufferedWriter writer = Files.newBufferedWriter(path, StandardCharsets.UTF_8)) {
+            // Use explicit \n for cross-platform consistency (matches original format)
+            writer.write("review,sentiment\n");
+
+            for (Dataset sample : data) {
+                writer.write('"');
+                writer.write(escapeCsvField(sample.getText()));
+                writer.write("\",");
+                writer.write(sample.getSentiment().name().toLowerCase());
+                writer.write('\n');
+            }
+        }
+    }
+
+    /**
+     * Escape special characters in CSV field (quotes and newlines).
+     */
+    private static String escapeCsvField(String text) {
+        if (text == null) {
+            return "";
+        }
+        // Double quotes to escape them in CSV
+        return text.replace("\"", "\"\"");
+    }
+
     //  UTILITIES
 
     /**
diff --git a/src/main/java/sentiment/training/ModelTrainer.java b/src/main/java/sentiment/training/ModelTrainer.java
@@ -7,6 +7,7 @@
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Component;
 import sentiment.data.Dataset;
+import sentiment.data.DatasetUtils;
 import sentiment.data.SimpleDatasetLoader;
 import sentiment.data.DatasetLoadResult;
 import sentiment.data.SplitManifest;
@@ -21,13 +22,11 @@
 import sentiment.config.FeatureExtractionProperties;
 import weka.core.Instances;
 
-import java.io.FileWriter;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.*;
-import java.util.stream.Collectors;
 
 /**
  * Service for training and persisting sentiment analysis models.
@@ -327,59 +326,8 @@ private List<Dataset> loadTrainingData(String dataPath, int maxSamples) throws E
 
         // Balance classes by undersampling majority class to match minority class
         // This ensures fair comparison across datasets (IMDB 50/50, Amazon 50/50, Yelp 75/25 -> 50/50)
-        List<Dataset> balanced = balanceClasses(shuffled);
-
-        return balanced;
-    }
-
-    /**
-     * Balances classes by undersampling the majority class to match the minority class.
-     * This ensures fair model comparison across datasets with different class distributions.
-     *
-     * @param data shuffled dataset (may be imbalanced)
-     * @return balanced dataset with equal positive/negative samples
-     */
-    private List<Dataset> balanceClasses(List<Dataset> data) {
-        // Separate by class
-        List<Dataset> positive = data.stream()
-                .filter(d -> d.getSentiment() == Dataset.SentimentLabel.POSITIVE)
-                .collect(Collectors.toList());
-        List<Dataset> negative = data.stream()
-                .filter(d -> d.getSentiment() == Dataset.SentimentLabel.NEGATIVE)
-                .collect(Collectors.toList());
-
-        // Find minority class size
-        int minoritySize = Math.min(positive.size(), negative.size());
-
-        if (minoritySize == 0) {
-            logger.warn("One class has zero samples - cannot balance. Returning original data.");
-            return data;
-        }
-
-        // Check if already balanced (within 5% tolerance)
-        double ratio = (double) Math.min(positive.size(), negative.size()) /
-                       Math.max(positive.size(), negative.size());
-        if (ratio >= 0.95) {
-            logger.info("Classes already balanced (ratio: {}) - no undersampling needed",
-                    String.format("%.2f", ratio));
-            return data;
-        }
-
-        // Undersample majority class
-        List<Dataset> balancedPositive = positive.subList(0, minoritySize);
-        List<Dataset> balancedNegative = negative.subList(0, minoritySize);
-
-        // Combine and shuffle again to mix classes
-        List<Dataset> balanced = new ArrayList<>(minoritySize * 2);
-        balanced.addAll(balancedPositive);
-        balanced.addAll(balancedNegative);
-        Collections.shuffle(balanced, new Random(42));
-
-        logger.info("Balanced classes: {} -> {} samples (undersampled {} from majority class)",
-                data.size(), balanced.size(), data.size() - balanced.size());
-        logger.info("Final distribution: positive={}, negative={}", minoritySize, minoritySize);
-
-        return balanced;
+        // Uses shared utility to ensure consistent balancing across training and data preparation
+        return DatasetUtils.balanceByUndersampling(shuffled, 42L);
     }
 
     /**
@@ -675,28 +623,19 @@ private void saveDataSplits(List<Dataset> trainData, List<Dataset> testData, Str
         }
 
         // Save both splits (backward compatibility path - no manifest)
+        // Uses shared CSV writer to ensure format consistency with data preparation pipeline
         logger.warn("Saving splits without manifest (legacy mode). Run ./scripts/prepare_data.sh for proper setup.");
-        saveSplitToFile(trainData, processedDir.resolve("train.csv"), "train");
-        saveSplitToFile(testData, processedDir.resolve("test.csv"), "test");
-    }
-
-    /**
-     * Save a data split to CSV file.
-     */
-    private void saveSplitToFile(List<Dataset> data, Path filePath, String splitName) {
-        try (FileWriter writer = new FileWriter(filePath.toFile())) {
-            writer.write("review,sentiment\n");
-
-            for (Dataset sample : data) {
-                String sentiment = sample.getSentiment().name().toLowerCase();
-                String text = escapeCsv(sample.getText());
-                writer.write("\"" + text + "\"," + sentiment + "\n");
-            }
-
-            logger.info("Saved {} split ({} samples) to: {}", splitName, data.size(), filePath);
-
-        } catch (Exception e) {
-            logger.error("ARTIFACT_SAVE_FAILED: {} split - {}", splitName, e.getMessage());
+        try {
+            SimpleDatasetLoader.writeCsv(trainData, processedDir.resolve("train.csv"));
+            logger.info("Saved train split ({} samples) to: {}", trainData.size(), processedDir.resolve("train.csv"));
+        } catch (IOException e) {
+            logger.error("ARTIFACT_SAVE_FAILED: train split - {}", e.getMessage());
+        }
+        try {
+            SimpleDatasetLoader.writeCsv(testData, processedDir.resolve("test.csv"));
+            logger.info("Saved test split ({} samples) to: {}", testData.size(), processedDir.resolve("test.csv"));
+        } catch (IOException e) {
+            logger.error("ARTIFACT_SAVE_FAILED: test split - {}", e.getMessage());
         }
     }
 
@@ -722,15 +661,6 @@ private String inferDatasetName(String dataPath) {
         return dotIndex > 0 ? filename.substring(0, dotIndex) : filename;
     }
 
-    /**
-     * Escape CSV special characters in text.
-     */
-    private String escapeCsv(String text) {
-        if (text == null) return "";
-        // Escape quotes by doubling them
-        return text.replace("\"", "\"\"");
-    }
-
     private void analyzeAndPrintFeatureImportance(SentimentClassifier classifier,
                                                    StratifiedDataSplitter.DataSplit split,
                                                    int topFeaturesCount,
diff --git a/src/test/java/sentiment/data/DatasetUtilsTest.java b/src/test/java/sentiment/data/DatasetUtilsTest.java
diff --git a/src/test/java/sentiment/data/SimpleDatasetLoaderTest.java b/src/test/java/sentiment/data/SimpleDatasetLoaderTest.java
