Pandas API addon: add_prefix method

hablapps · Nov 29, 2023 · ad3d65b · ad3d65b
1 parent 7de2b7c
commit ad3d65b
Showing 1 changed file with 24 additions and 0 deletions.
diff --git a/src/pykx/pandas_api/pandas_indexing.py b/src/pykx/pandas_api/pandas_indexing.py
@@ -328,6 +328,17 @@ def _rename_columns(tab, labels):
                  tab, labels)  # noqa
     else:
         return q('{c:cols x; c:@[c;c?key y;y]; c xcol x}', tab, labels)
+
+
+def _prefix_columns(tab, prefix):
+        if "Keyed" in str(type(tab)):
+            return q('''{
+                    c:cols value x;
+                    c: `$"hola" , /: string c;
+                    key[x]!c xcol value x}''',
+                    tab, prefix)  # noqa
+        else:
+            return q('{c:cols x; c: `$"hola" ,/: string c; c xcol x}', tab)
 
 
 class PandasIndexing:
@@ -453,6 +464,19 @@ def rename(self, labels=None, index=None, columns=None, axis=0,
                 t = _rename_columns(t, columns)
 
         return t
+
+    def add_prefix(self, prefix=None, axis=0):
+        t = self
+
+        if prefix:
+            if axis == 0:
+                t = _prefix_columns(t, prefix)
+            elif axis == 1:
+                t = _rename_columns(t, labels)
+            else:
+                raise ValueError(f'No axis named {axis}')
+
+        return t
 
     def sample(self, n=None, frac=None, replace=False, weights=None,
                random_state=None, axis=None, ignore_index=False):