diff --git a/.ipynb_checkpoints/main-checkpoint.ipynb b/.ipynb_checkpoints/main-checkpoint.ipynb
index 4bce2ae..c9c683c 100644
--- a/.ipynb_checkpoints/main-checkpoint.ipynb
+++ b/.ipynb_checkpoints/main-checkpoint.ipynb
@@ -105,7 +105,7 @@
     "\n",
     "query = \"\"\"\n",
     "        SELECT \n",
-    "        t.transaction_id,\n",
+    "        t.transaction_key,\n",
     "        t.transaction_date,\n",
     "        t.transaction_amount,\n",
     "        t.transaction_desc,\n",
@@ -162,8 +162,10 @@
     "                                                    acc_data=\"account_data_v1\",\n",
     "                                                    alert_data=\"alert_data_v1\")\n",
     "                                       )\n",
+    "        \"Focal_id\", \"Cash_deposit_total\", \"Cash_deposit_count\",\n",
+    "                \"Segment\", \"Risk\", \"SAR_FLAG\"\n",
     "        cols = [\n",
-    "                \"transaction_id\",\n",
+    "                \"Focal_id\",\n",
     "                \"transaction_date\",\n",
     "                \"transaction_amount\",\n",
     "                \"transaction_desc\",\n",
@@ -176,9 +178,9 @@
     "                \"customer_name\",\n",
     "                \"date_of_birth\",\n",
     "                \"nationality\",\n",
-    "                \"business_segment\",\n",
-    "                \"risk_level\",\n",
-    "                \"sar_flag\"\n",
+    "                \"Segment\",\n",
+    "                \"Risk\", \n",
+    "                \"SAR_FLAG\"\n",
     "            ]\n",
     "        df = pd.DataFrame(row_list, columns = cols)\n",
     "        return df"
diff --git a/main.ipynb b/main.ipynb
index 4bce2ae..c9c683c 100644
--- a/main.ipynb
+++ b/main.ipynb
@@ -105,7 +105,7 @@
     "\n",
     "query = \"\"\"\n",
     "        SELECT \n",
-    "        t.transaction_id,\n",
+    "        t.transaction_key,\n",
     "        t.transaction_date,\n",
     "        t.transaction_amount,\n",
     "        t.transaction_desc,\n",
@@ -162,8 +162,10 @@
     "                                                    acc_data=\"account_data_v1\",\n",
     "                                                    alert_data=\"alert_data_v1\")\n",
     "                                       )\n",
+    "        \"Focal_id\", \"Cash_deposit_total\", \"Cash_deposit_count\",\n",
+    "                \"Segment\", \"Risk\", \"SAR_FLAG\"\n",
     "        cols = [\n",
-    "                \"transaction_id\",\n",
+    "                \"Focal_id\",\n",
     "                \"transaction_date\",\n",
     "                \"transaction_amount\",\n",
     "                \"transaction_desc\",\n",
@@ -176,9 +178,9 @@
     "                \"customer_name\",\n",
     "                \"date_of_birth\",\n",
     "                \"nationality\",\n",
-    "                \"business_segment\",\n",
-    "                \"risk_level\",\n",
-    "                \"sar_flag\"\n",
+    "                \"Segment\",\n",
+    "                \"Risk\", \n",
+    "                \"SAR_FLAG\"\n",
     "            ]\n",
     "        df = pd.DataFrame(row_list, columns = cols)\n",
     "        return df"
diff --git a/main.py b/main.py
index 81b6a29..140c1cf 100644
--- a/main.py
+++ b/main.py
@@ -94,7 +94,7 @@ import pandas as pd
 
 query = """
         SELECT 
-        t.transaction_id,
+        t.transaction_key,
         t.transaction_date,
         t.transaction_amount,
         t.transaction_desc,
@@ -151,8 +151,10 @@ class Scenario:
                                                     acc_data="account_data_v1",
                                                     alert_data="alert_data_v1")
                                        )
+        "Focal_id", "Cash_deposit_total", "Cash_deposit_count",
+                "Segment", "Risk", "SAR_FLAG"
         cols = [
-                "transaction_id",
+                "Focal_id",
                 "transaction_date",
                 "transaction_amount",
                 "transaction_desc",
@@ -165,9 +167,9 @@ class Scenario:
                 "customer_name",
                 "date_of_birth",
                 "nationality",
-                "business_segment",
-                "risk_level",
-                "sar_flag"
+                "Segment",
+                "Risk", 
+                "SAR_FLAG"
             ]
         df = pd.DataFrame(row_list, columns = cols)
         return df