backup

Author: davidkhala

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

@@ -1,23 +1,39 @@
 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 = pd.Series(True, index=df.index)
-    if df.columns.empty and df.index.empty:
-        return pd.DataFrame([record])
 
-    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")
+def upsert(df: pd.DataFrame, *primary_keys: str, record: dict) -> pd.DataFrame:
+    # 判断是否已有索引结构
+    index_names = df.index.names if isinstance(df.index, pd.MultiIndex) else [df.index.name]
+    has_index = all(k in index_names for k in primary_keys if k is not None)
 
-    match_indices = df[condition].index
+    # 如果非空，尝试匹配并更新
+    if not df.empty:
+        condition = pd.Series(True, index=df.index)
+        for key in primary_keys:
+            if key in df.columns:
+                condition &= (df[key] == record[key])
+            elif has_index:
+                condition &= (df.index.get_level_values(key) == record[key])
+            else:
+                raise KeyError(f"'{key}' not found in columns or index")
 
-    if not match_indices.empty:
-        for col, value in record.items():
-            if col in df.columns:
-                df.loc[match_indices, col] = value
-        return df
-    return pd.concat([df, pd.DataFrame([record])])
+        if condition.any():
+            for col, val in record.items():
+                if col in df.columns:
+                    df.loc[condition, col] = val
+            return df
+
+    # 插入新行
+    index_vals = tuple(record[k] for k in primary_keys)
+    data = {k: v for k, v in record.items() if not has_index or k not in primary_keys}
+
+    new_row = pd.DataFrame([data], columns=df.columns if not df.columns.empty else None)
+
+    if has_index:
+        new_row.index = (
+            pd.Index([index_vals[0]], name=primary_keys[0])
+            if len(primary_keys) == 1
+            else pd.MultiIndex.from_tuples([index_vals], names=primary_keys)
+        )
+        return pd.concat([df, new_row])
+    else:
+        return pd.concat([df, new_row], ignore_index=True)

frame/tests/pandas_test.py

@@ -5,8 +5,8 @@
 
 class AirlineTestCase(unittest.TestCase):
     def test_import(self):
-        df_airports = pandas.read_csv("fixtures/airports.csv")
-        df_airlines = pandas.read_csv("fixtures/airlines.csv")
+        pandas.read_csv("fixtures/airports.csv")
+        pandas.read_csv("fixtures/airlines.csv")
 
     def test_duckDB(self):
         conn = duckdb.connect()
@@ -15,31 +15,37 @@ 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):
 
-        df = pandas.DataFrame({
-            'id': [1, 2],
-            'name': ['Alice', 'Bob'],
-            'score': [85, 90]
-        })
-        prim_key ='id'
-        df = df.set_index(prim_key)
-        new_record = {'name': 'Charlie', 'score': 95, prim_key:2}
-        upsert(df, prim_key, record = new_record)
+class SyntaxTestCase(unittest.TestCase):
+    single_index_df = pandas.DataFrame({
+        'id': [1, 2],
+        'name': ['Alice', 'Bob'],
+        'score': [85, 90]
+    }).set_index('id')
 
+    def test_upsert(self):
+        df = self.single_index_df
+        prim_key = df.index.name
+        new_record = {'name': 'Charlie', 'score': 95, prim_key: 2}
+        df = upsert(df, prim_key, record=new_record)
         self.assertEqual(95, df.at[2, 'score'])
+        self.assertEqual(2, len(df.columns))
+
     def test_upsert_empty(self):
         df = pandas.DataFrame()
         prim_key = 'id'
         new_record = {'name': 'Charlie', 'score': 95, prim_key: 2}
-        upsert(df, prim_key, record=new_record)
-
+        df = upsert(df, prim_key, record=new_record)
+        self.assertEqual(3, len(df.columns))
+        self.assertEqual(1, len(df))
 
         df = pandas.DataFrame(columns=[prim_key]).set_index(prim_key)
-        upsert(df, prim_key, record=new_record)
+        df = upsert(df, prim_key, record=new_record)
+        self.assertEqual(2, len(df.columns))
+        self.assertEqual(1, len(df))
 
     def test_upsert2(self):
         df = pandas.DataFrame([
@@ -51,6 +57,7 @@ def test_upsert2(self):
         df = upsert(df, 'School', 'Country', record=record)
         self.assertEqual(new_students, df.loc[('Oxford', 'UK')].Students)
         print(df)
+
     def test_upsert3(self):
         df = pandas.DataFrame([
             {'School': 'Oxford', 'Country': 'UK', 'Students': 1000},
@@ -62,6 +69,5 @@ def test_upsert3(self):
         self.assertEqual(new_students, df[(df['School'] == 'Oxford') & (df['Country'] == 'UK')].iloc[0].Students)
 
 
-
 if __name__ == '__main__':
     unittest.main()