Merge pull request #5 from drobison00/devin-dfp-fixes

Devin dfp fixes

morpheus/utils/nvt/mutate.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
 import functools
 import os
 import typing
@@ -139,7 +140,19 @@ def transform(self, col_selector: ColumnSelector, df: DataFrameType) -> DataFram
             Transformed dataframe.
         """
 
-        return self._func(col_selector, df)
+        df = self._func(col_selector, df)
+
+        # If our dataframe doesn't contain the expected output columns, even after processing, we add dummy columns.
+        # This could occur if our JSON data doesn't always contain columns we expect to be expanded.
+        df_cols_set = set(df.columns)
+        new_cols = {
+            col[0]: np.zeros(df.shape[0], dtype=col[1])
+            for col in self._output_columns if col[0] not in df_cols_set
+        }
+
+        df = df.assign(**new_cols)
+
+        return df
 
     def column_mapping(self, col_selector: ColumnSelector) -> typing.Dict[str, str]:
         """

tests/test_downloader.py

@@ -40,7 +40,7 @@ def dask_distributed(fail_missing: bool):
 @pytest.fixture(autouse=True, scope='session')
 def dask_cuda(fail_missing: bool):
     """
-    Mark tests requiring dat_cuda
+    Mark tests requiring dask_cuda
     """
     yield import_or_skip("dask_cuda", reason="Downloader requires dask_cuda", fail_missing=fail_missing)
 
@@ -101,7 +101,7 @@ def test_get_dask_cluster(mock_dask_cluster: mock.MagicMock,
     assert downloader.get_dask_cluster() is mock_dask_cluster
 
     mock_dask_config.set.assert_called_once()
-    mock_dask_cluster.assert_called_once_with(start=True, processes=use_processes)
+    mock_dask_cluster.assert_called_once_with()
 
 
 @mock.patch('dask.config')