Make unique index sorted

src/fondant/components/load_from_csv/src/main.py

@@ -72,14 +72,15 @@ def _set_unique_index(dataframe: pd.DataFrame, partition_info=None):
                 Function that sets a unique index
                 based on the partition and row number.
                 """
+                n_digits = len(str(len(dataframe)))
+
                 dataframe["id"] = 1
                 dataframe["id"] = (
                     str(partition_info["number"])
                     + "_"
-                    + (dataframe.id.cumsum()).astype(str)
+                    + dataframe.id.cumsum().astype(str).str.zfill(n_digits)
                 )
-                dataframe.index = dataframe.pop("id")
-                return dataframe
+                return dataframe.set_index("id", drop=True)
 
             def _get_meta_df() -> pd.DataFrame:
                 meta_dict = {"id": pd.Series(dtype="object")}

src/fondant/components/load_from_hf_hub/src/main.py

@@ -81,14 +81,15 @@ def set_df_index(self, dask_df: dd.DataFrame) -> dd.DataFrame:
 
             def _set_unique_index(dataframe: pd.DataFrame, partition_info=None):
                 """Function that sets a unique index based on the partition and row number."""
+                n_digits = len(str(len(dataframe)))
+
                 dataframe["id"] = 1
                 dataframe["id"] = (
                     str(partition_info["number"])
                     + "_"
-                    + (dataframe.id.cumsum()).astype(str)
+                    + dataframe.id.cumsum().astype(str).str.zfill(n_digits)
                 )
-                dataframe.index = dataframe.pop("id")
-                return dataframe
+                return dataframe.set_index("id", drop=True)
 
             def _get_meta_df() -> pd.DataFrame:
                 meta_dict = {"id": pd.Series(dtype="object")}

src/fondant/components/load_from_parquet/src/main.py

@@ -68,14 +68,15 @@ def set_df_index(self, dask_df: dd.DataFrame) -> dd.DataFrame:
 
             def _set_unique_index(dataframe: pd.DataFrame, partition_info=None):
                 """Function that sets a unique index based on the partition and row number."""
+                n_digits = len(str(len(dataframe)))
+
                 dataframe["id"] = 1
                 dataframe["id"] = (
                     str(partition_info["number"])
                     + "_"
-                    + (dataframe.id.cumsum()).astype(str)
+                    + dataframe.id.cumsum().astype(str).str.zfill(n_digits)
                 )
-                dataframe.index = dataframe.pop("id")
-                return dataframe
+                return dataframe.set_index("id", drop=True)
 
             def _get_meta_df() -> pd.DataFrame:
                 meta_dict = {"id": pd.Series(dtype="object")}

src/fondant/components/load_from_pdf/src/main.py

@@ -47,14 +47,15 @@ def set_df_index(self, dask_df: dd.DataFrame) -> dd.DataFrame:
 
             def _set_unique_index(dataframe: pd.DataFrame, partition_info=None):
                 """Function that sets a unique index based on the partition and row number."""
+                n_digits = len(str(len(dataframe)))
+
                 dataframe["id"] = 1
                 dataframe["id"] = (
                     str(partition_info["number"])
                     + "_"
-                    + (dataframe.id.cumsum()).astype(str)
+                    + dataframe.id.cumsum().astype(str).str.zfill(n_digits)
                 )
-                dataframe.index = dataframe.pop("id")
-                return dataframe
+                return dataframe.set_index("id", drop=True)
 
             def _get_meta_df() -> pd.DataFrame:
                 meta_dict = {"id": pd.Series(dtype="object")}