Switch from_flat 'index' keyword to 'on'

src/nested_dask/core.py

@@ -287,7 +287,7 @@ def from_map(
         return NestedFrame.from_dask_dataframe(nf)
 
     @classmethod
-    def from_flat(cls, df, base_columns, nested_columns=None, index=None, name="nested"):
+    def from_flat(cls, df, base_columns, nested_columns=None, on=None, name="nested"):
         """Creates a NestedFrame with base and nested columns from a flat
         dataframe.
 
@@ -303,7 +303,7 @@ def from_flat(cls, df, base_columns, nested_columns=None, index=None, name="nest
             in the list will attempt to be packed into a single nested column
             with the name provided in `nested_name`. If None, is defined as all
             columns not in `base_columns`.
-        index: str, or None
+        on: str or None
             The name of a column to use as the new index. Typically, the index
             should have a unique value per row for base columns, and should
             repeat for nested columns. For example, a dataframe with two
@@ -323,15 +323,15 @@ def from_flat(cls, df, base_columns, nested_columns=None, index=None, name="nest
         meta = npd.NestedFrame(df[base_columns]._meta)
 
         if nested_columns is None:
-            nested_columns = [col for col in df.columns if (col not in base_columns) and col != index]
+            nested_columns = [col for col in df.columns if (col not in base_columns) and col != on]
 
         if len(nested_columns) > 0:
             nested_meta = pack(df[nested_columns]._meta, name)
             meta = meta.join(nested_meta)
 
         return df.map_partitions(
             lambda x: npd.NestedFrame.from_flat(
-                df=x, base_columns=base_columns, nested_columns=nested_columns, index=index, name=name
+                df=x, base_columns=base_columns, nested_columns=nested_columns, on=on, name=name
             ),
             meta=meta,
         )

tests/nested_dask/test_nestedframe.py

@@ -1,14 +1,15 @@
 import dask
 import dask.dataframe as dd
-import nested_dask as nd
 import nested_pandas as npd
 import numpy as np
 import pandas as pd
 import pyarrow as pa
 import pytest
-from nested_dask.datasets import generate_data
 from nested_pandas.series.dtype import NestedDtype
 
+import nested_dask as nd
+from nested_dask.datasets import generate_data
+
 dask.config.set({"dataframe.convert-string": False})
 
 
@@ -173,7 +174,7 @@ def test_from_flat():
     assert len(ndf_comp) == 2
 
     # Check using an index
-    ndf = nd.NestedFrame.from_flat(nf, base_columns=["b"], index="a")
+    ndf = nd.NestedFrame.from_flat(nf, base_columns=["b"], on="a")
     assert list(ndf.columns) == ["b", "nested"]
     assert list(ndf["nested"].nest.fields) == ["c", "d"]
     ndf_comp = ndf.compute()