Getting factors and numeric columns

Author: wtsimple

data_preprocessor.py

@@ -1,4 +1,19 @@
+from typing import List
+
+
 class DataPreprocessor(object):
     def __init__(self, train_df=None, test_df=None):
         self.train_df = train_df
         self.test_df = test_df
+
+
+    def get_factors(self):
+        return self._get_cols_by_types(types=['string'])
+
+
+    def get_numeric_columns(self):
+        return self._get_cols_by_types(types=['double', 'int'])
+
+
+    def _get_cols_by_types(self, types: List[str] = None):
+        return [col for col, data_type in self.train_df.dtypes if data_type in types]

settings.py

@@ -1,16 +1,22 @@
+# Adult data set configs
+ADULT_COLUMN_NAMES = ["age",
+                      "workclass",
+                      "fnlwgt",
+                      "education",
+                      "education_num",
+                      "marital_status",
+                      "occupation",
+                      "relationship",
+                      "race",
+                      "sex",
+                      "capital_gain",
+                      "capital_loss",
+                      "hours_per_week",
+                      "native_country",
+                      "income"]
+
+ADULT_TRAIN_DATA = "data/adult.data"
+ADULT_TEST_DATA = "data/adult.test"
+
+# Test configurations
 TEST_DATA_PATH = "data/data_example_for_tests.csv"
-COLUMN_NAMES = ["age",
-                "workclass",
-                "fnlwgt",
-                "education",
-                "education_num",
-                "marital_status",
-                "occupation",
-                "relationship",
-                "race",
-                "sex",
-                "capital_gain",
-                "capital_loss",
-                "hours_per_week",
-                "native_country",
-                "income"]

tests/test_data_loader.py

@@ -2,7 +2,7 @@
 from pyspark.sql import DataFrame
 
 from data_loader import DataLoader
-from settings import TEST_DATA_PATH, COLUMN_NAMES
+from settings import TEST_DATA_PATH, ADULT_COLUMN_NAMES
 
 
 @pytest.fixture
@@ -11,9 +11,9 @@ def data_loader():
 
 
 def test_data_loader_loads_data_frame(data_loader):
-    df = data_loader.load_relative(path=TEST_DATA_PATH, columns=COLUMN_NAMES)
+    df = data_loader.load_relative(path=TEST_DATA_PATH, columns=ADULT_COLUMN_NAMES)
     assert isinstance(df, DataFrame)
-    assert df.columns == COLUMN_NAMES
+    assert df.columns == ADULT_COLUMN_NAMES
     # Check some values from the first row
     first_row = df.first()
     assert first_row.income == ' <=50K'

tests/test_data_preprocessor.py

@@ -5,14 +5,33 @@
 
 from data_loader import DataLoader
 from data_preprocessor import DataPreprocessor
-from settings import TEST_DATA_PATH, COLUMN_NAMES
+from settings import TEST_DATA_PATH, ADULT_COLUMN_NAMES
 
 
 @pytest.fixture
 def preprocessor():
-    df = DataLoader().load_relative(path=TEST_DATA_PATH, columns=COLUMN_NAMES)
+    df = DataLoader().load_relative(path=TEST_DATA_PATH, columns=ADULT_COLUMN_NAMES)
     return DataPreprocessor(train_df=df, test_df=df)
 
 
+def test_preprocessor_get_factors(preprocessor):
+    # Example 1
+    factors_example = ['workclass', 'education', 'marital_status', 'occupation', 'relationship', 'race', 'sex',
+                       'native_country', 'income']
+    assert preprocessor.get_factors() == factors_example
+    # Example 2 (rename last column)
+    preprocessor.train_df = preprocessor.train_df.withColumnRenamed('income', 'income2')
+    factors_example[-1] = "income2"
+    assert preprocessor.get_factors() == factors_example
+
+
+def test_preprocessor_get_numeric_columns(preprocessor):
+    numeric_cols = ['age', 'fnlwgt', 'education_num', 'capital_gain', 'capital_loss', 'hours_per_week']
+    assert preprocessor.get_numeric_columns() == numeric_cols
+    preprocessor.train_df = preprocessor.train_df.withColumnRenamed('capital_loss', 'capital_super_loss')
+    numeric_cols[-2] = 'capital_super_loss'
+    assert preprocessor.get_numeric_columns() == numeric_cols
+
+
 def test_data_preprocessor_explore_factors(preprocessor):
-    pass
+    pass  # preprocessor.explore_factors()