necessary changes to allow graph execution in dataloader (#152)

* necessary changes to allow graph execution in dataloader

* added test for wildcard, removed comments, refactored ternary

Co-authored-by: Karl Higley <kmhigley@gmail.com>

Author: jperez999

merlin/dag/executors.py

@@ -120,7 +120,8 @@ def _build_input_data(self, node, transformable, capture_dtypes=False):
         else:
             # If there are no parents, this is an input node,
             # so pull columns directly from root data
-            input_data = transformable[node_input_cols + list(addl_input_cols)]
+            addl_input_cols = list(addl_input_cols) if addl_input_cols else []
+            input_data = transformable[node_input_cols + addl_input_cols]
 
         return input_data
 
@@ -161,6 +162,10 @@ def _transform_data(self, node, input_data, capture_dtypes=False):
 
                 if is_list:
                     col_dtype = list_val_dtype(col_series)
+                if hasattr(col_dtype, "as_numpy_dtype"):
+                    col_dtype = col_dtype.as_numpy_dtype()
+                elif hasattr(col_series, "numpy"):
+                    col_dtype = col_series[0].cpu().numpy().dtype
 
                 output_data_schema = output_col_schema.with_dtype(
                     col_dtype, is_list=is_list, is_ragged=is_list

merlin/dag/ops/selection.py

@@ -108,4 +108,6 @@ def compute_output_schema(
             The schemas of the columns produced by this operator
         """
         selector = col_selector or self.selector
+        if selector.all:
+            selector = ColumnSelector(input_schema.column_names)
         return super().compute_output_schema(input_schema, selector, prev_output_schema)

tests/unit/dag/ops/test_selection.py

@@ -49,3 +49,14 @@ def test_selection_output_schema(df):
     result_schema = op.compute_output_schema(schema, ColumnSelector())
 
     assert result_schema.column_names == ["x", "y"]
+
+
+@pytest.mark.parametrize("engine", ["parquet"])
+def test_selection_wildcard_output_schema(df):
+    selector = ColumnSelector("*")
+    schema = Schema([ColumnSchema(col) for col in df.columns])
+    op = SelectionOp(selector)
+
+    result_schema = op.compute_output_schema(schema, ColumnSelector())
+
+    assert result_schema.column_names == schema.column_names