Get notebook asset working correctly

catalyst-cooperative · Jan 15, 2025 · 4474239 · 4474239
1 parent ce8b109
commit 4474239
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 29 deletions.
diff --git a/src/mozilla_sec_eia/models/sec_eia_record_linkage/notebooks/splink-sec-eia.ipynb b/src/mozilla_sec_eia/models/sec_eia_record_linkage/notebooks/splink-sec-eia.ipynb
@@ -15,9 +15,7 @@
    "id": "e136eeca-133b-4f61-8e37-5e1ea8c99683",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "import dagstermill"
-   ]
+   "source": []
   },
   {
    "cell_type": "code",
@@ -3742,39 +3740,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 62,
+   "execution_count": 66,
    "id": "7385664a-4a40-44e8-b7bf-917a623be158",
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Index(['record_id', 'company_name', 'street_address', 'filename', 'phone_number', 'central_index_key', 'city', 'company_name_raw', 'date_of_name_change', 'film_number', 'fiscal_year_end', 'form_type', 'former_conformed_name', 'irs_number', 'organization_name', 'sec_act', 'sec_file_number', 'standard_industrial_classification', 'state', 'state_of_incorporation', 'street_address_2', 'zip_code', 'report_date', 'report_year', 'location_of_inc', 'company_name_no_legal', 'company_name_mphone', 'files_10k', 'sec_company_id', 'utility_id_eia'], dtype='object')"
-      ]
-     },
-     "execution_count": 62,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "out_df.columns"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "7c9f50d5-f612-4887-b3ef-1ba810973d1a",
-   "metadata": {},
    "outputs": [],
    "source": [
-    "dagstermill.yield_result(out_df, output_name=\"sec_10k_filers_matched_df\")"
+    "# Hacky write directly to GCS\n",
+    "out_df.to_parquet(\"gs://sec10k-outputs/v2/core_sec_10k__filers.parquet\")"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9bfccfa0-17eb-4373-8a89-b578ff00349e",
+   "id": "2a1aa1c1-c497-439a-a66f-71d308772fce",
    "metadata": {},
    "outputs": [],
    "source": []

diff --git a/src/mozilla_sec_eia/models/sec_eia_record_linkage/transform_sec_input.py b/src/mozilla_sec_eia/models/sec_eia_record_linkage/transform_sec_input.py
@@ -354,13 +354,13 @@ def transformed_basic_10k(
 
 @asset(
     ins={
-        "sec_10k_filers_matched_df": AssetIn("core_sec_10k__filers"),
         "clean_ex21_df": AssetIn("transformed_ex21_subsidiary_table"),
         "clean_eia_df": AssetIn("core_eia__parents_and_subsidiaries"),
     },
+    deps=["core_sec_10k__filers"],
+    io_manager_key="pandas_parquet_io_manager",
 )
 def out_sec_10k__parents_and_subsidiaries(
-    sec_10k_filers_matched_df: pd.DataFrame,
     clean_ex21_df: pd.DataFrame,
     clean_eia_df: pd.DataFrame,
 ) -> pd.DataFrame:
@@ -370,6 +370,9 @@ def out_sec_10k__parents_and_subsidiaries(
     filing companies. Create an sec_company_id for subsidiaries
     that aren't linked to a CIK.
     """
+    sec_10k_filers_matched_df = pd.read_parquet(
+        "gs://sec10k-outputs/v2/core_sec_10k__filers.parquet"
+    )
     ex21_df_with_cik = match_ex21_subsidiaries_to_filer_company(
         basic10k_df=sec_10k_filers_matched_df, ex21_df=clean_ex21_df
     )