0.0.9 update upsert

Author: davidkhala

frame/davidkhala/data/frame/pandas/__init__.py

@@ -1,3 +1,23 @@
-def upsert(df, primary_key:str, record:dict):
-    df.loc[record[primary_key]] = record
-    return df
+import pandas as pd
+def upsert(df, *primary_keys: str, record: dict):
+    index_keys = df.index.names if isinstance(df.index, pd.MultiIndex) else [df.index.name]
+    condition = True
+
+    for key in primary_keys:
+        if key in df.columns:
+            condition &= (df[key] == record[key])
+        elif key in index_keys:
+            condition &= (df.index.get_level_values(key) == record[key])
+        else:
+            raise KeyError(f"'{key}' not found in either columns or index")
+
+    match_indices = df[condition].index
+
+    if not match_indices.empty:
+        for col, value in record.items():
+            if col in df.columns:
+                df.loc[match_indices, col] = value
+    else:
+        df = pd.concat([df, pd.DataFrame([record])], ignore_index=True)
+
+    return df

frame/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "davidkhala.data.frame"
-version = "0.0.8"
+version = "0.0.9"
 description = ""
 authors = [{ name = "David Liu", email = "[email protected]" }]
 requires-python = ">=3.10"

frame/tests/pandas_test.py

@@ -15,20 +15,40 @@ def test_duckDB(self):
         conn.register("airports", df_airports)
         conn.close()
 
+from davidkhala.data.frame.pandas import upsert
 class SyntaxTestCase(unittest.TestCase):
+
     def test_upsert(self):
-        from davidkhala.data.frame.pandas import upsert
-        prim_key ='id'
+
         df = pandas.DataFrame({
-            prim_key: [1, 2],
+            'id': [1, 2],
             'name': ['Alice', 'Bob'],
             'score': [85, 90]
-        }).set_index(prim_key)
+        })
+        prim_key ='id'
+        df = df.set_index(prim_key)
         new_record = {'name': 'Charlie', 'score': 95, prim_key:2}
-        upsert(df, prim_key, new_record)
+        upsert(df, prim_key, record = new_record)
 
         self.assertEqual(95, df.at[2, 'score'])
-
+    def test_upsert2(self):
+        df = pandas.DataFrame([
+            {'School': 'Oxford', 'Country': 'UK', 'Students': 1000},
+            {'School': 'Harvard', 'Country': 'US', 'Students': 1500}
+        ]).set_index(['School', 'Country'])
+        new_students = 2000
+        record = {'School': 'Oxford', 'Country': 'UK', 'Students': new_students}
+        df = upsert(df, 'School', 'Country', record=record)
+        self.assertEqual(new_students, df.loc[('Oxford', 'UK')].Students)
+    def test_upsert3(self):
+        df = pandas.DataFrame([
+            {'School': 'Oxford', 'Country': 'UK', 'Students': 1000},
+            {'School': 'Harvard', 'Country': 'US', 'Students': 1500}
+        ])
+        new_students = 2000
+        record = {'School': 'Oxford', 'Country': 'UK', 'Students': new_students}
+        df = upsert(df, 'School', 'Country', record=record)
+        self.assertEqual(new_students, df[(df['School'] == 'Oxford') & (df['Country'] == 'UK')].iloc[0].Students)