data loader seems to be working OK

Author: wtsimple

data/data_example_for_tests.csv

@@ -0,0 +1,10 @@
+39, State-gov, 77516, Bachelors, 13, Never-married, Adm-clerical, Not-in-family, White, Male, 2174, 0, 40, United-States, <=50K
+50, Self-emp-not-inc, 83311, Bachelors, 13, Married-civ-spouse, Exec-managerial, Husband, White, Male, 0, 0, 13, United-States, <=50K
+38, Private, 215646, HS-grad, 9, Divorced, Handlers-cleaners, Not-in-family, White, Male, 0, 0, 40, United-States, <=50K
+53, Private, 234721, 11th, 7, Married-civ-spouse, Handlers-cleaners, Husband, Black, Male, 0, 0, 40, United-States, <=50K
+28, Private, 338409, Bachelors, 13, Married-civ-spouse, Prof-specialty, Wife, Black, Female, 0, 0, 40, Cuba, <=50K
+37, Private, 284582, Masters, 14, Married-civ-spouse, Exec-managerial, Wife, White, Female, 0, 0, 40, United-States, <=50K
+49, Private, 160187, 9th, 5, Married-spouse-absent, Other-service, Not-in-family, Black, Female, 0, 0, 16, Jamaica, <=50K
+52, Self-emp-not-inc, 209642, HS-grad, 9, Married-civ-spouse, Exec-managerial, Husband, White, Male, 0, 0, 45, United-States, >50K
+31, Private, 45781, Masters, 14, Never-married, Prof-specialty, Not-in-family, White, Female, 14084, 0, 50, United-States, >50K
+42, Private, 159449, Bachelors, 13, Married-civ-spouse, Exec-managerial, Husband, White, Male, 5178, 0, 40, United-States, >50K

data_loader.py

@@ -0,0 +1,32 @@
+import inspect
+import os
+
+from spark_launcher import SparkLauncher
+
+
+class DataLoader(object):
+    def __init__(self):
+        self.spark = SparkLauncher()
+
+
+    def load_relative(self, path='', columns=None):
+        absolute_path = self._get_absolute_path(path)
+        df = self.spark.session.read.csv(absolute_path, header=False, inferSchema=True)
+        if columns:
+            df = self._rename_columns(columns, df)
+        return df
+
+
+    @staticmethod
+    def _rename_columns(columns, df):
+        for new_col, old_col in zip(columns, df.columns):
+            df = df.withColumnRenamed(old_col, new_col)
+        return df
+
+
+    def _get_absolute_path(self, relative_path=''):
+        current_file = inspect.getfile(self.__class__)
+        directory = os.path.dirname(current_file)
+        file_name = os.path.join(
+            directory, relative_path)
+        return file_name

tests/test_data_loader.py

@@ -0,0 +1,35 @@
+import pytest
+from pyspark.sql import DataFrame
+
+from data_loader import DataLoader
+
+COLUMN_NAMES = ["age",
+                "workclass",
+                "fnlwgt",
+                "education",
+                "education_num",
+                "marital_status",
+                "occupation",
+                "relationship",
+                "race",
+                "sex",
+                "capital_gain",
+                "capital_loss",
+                "hours_per_week",
+                "native_country",
+                "income"]
+
+
+@pytest.fixture
+def data_loader():
+    return DataLoader()
+
+
+def test_data_loader_loads_data(data_loader):
+    df = data_loader.load_relative(path="data/data_example_for_tests.csv", columns=COLUMN_NAMES)
+    assert isinstance(df, DataFrame)
+    assert df.columns == COLUMN_NAMES
+    # Check some values from the first row
+    first_row = df.first()
+    assert first_row.income == ' <=50K'
+    assert first_row.age == 39

todo_list.md

@@ -1,5 +1,5 @@
 - [x] start and configure Spark
-- [ ] load data from the files
+- [ ] load_relative data from the files
 - [ ] prepare data for classification models
     - [ ] string encode the factor columns
     - [ ] one-hot encode the numerically encoded factor columns