fix so databricks pipes doesnt fail io manager (#47)

dagster-io · Oct 24, 2023 · 59e5f4b · 59e5f4b
1 parent 4f9abe9
commit 59e5f4b
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/hooli_data_eng/assets/forecasting/__init__.py b/hooli_data_eng/assets/forecasting/__init__.py
@@ -13,6 +13,7 @@
     Int,
     Config,
     AssetExecutionContext,
+    MaterializeResult
 )
 from dagstermill import define_dagstermill_asset
 from dagster._utils import file_relative_path
@@ -177,11 +178,11 @@ def big_orders(context, predicted_orders: pd.DataFrame):
 # or use that upstream Snowflake table, it is used here for illustrative purposes
 @asset(
         deps=[predicted_orders],
-        compute_kind="databricks"
+        compute_kind="databricks",
 )
 def databricks_asset(
     context: AssetExecutionContext, pipes_client: PipesDatabricksClient
-):
+) -> MaterializeResult:
     # cluster config
     cluster_config = {
         "num_workers": 1,
@@ -216,4 +217,4 @@ def databricks_asset(
         task=task,
         context=context,
         extras=extras,
-    ).get_results()
+    ).get_materialize_result()