fix: don't use categoricals in output format (#63)

* fix: don't use categoricals in output format

* Remove categorical usage from fetch metrics buckets output format

Author: pitercl

src/neptune_query/internal/output_format.py

@@ -301,14 +301,13 @@ def generate_categorized_rows() -> Generator[Tuple, None, None]:
         np.fromiter(generate_categorized_rows(), dtype=types),
     )
 
-    experiment_dtype = pd.CategoricalDtype(categories=label_mapping)
-    df[index_column_name] = pd.Categorical.from_codes(df[index_column_name], dtype=experiment_dtype)
     if timestamp_column_name:
         df[timestamp_column_name] = pd.to_datetime(df[timestamp_column_name], unit="ms", origin="unix", utc=True)
 
-    df = _pivot_and_reindex_df(df, include_point_previews, index_column_name, timestamp_column_name)
+    df = _pivot_df(df, include_point_previews, index_column_name, timestamp_column_name)
+    df = _restore_labels_in_index(df, index_column_name, label_mapping)
     df = _restore_path_column_names(df, path_mapping, "float_series" if type_suffix_in_column_names else None)
-    df = _sort_indices(df)
+    df = _sort_index_and_columns(df, index_column_name)
 
     return df
 
@@ -390,14 +389,13 @@ def generate_categorized_rows() -> Generator[Tuple, None, None]:
         np.fromiter(generate_categorized_rows(), dtype=types),
     )
 
-    experiment_dtype = pd.CategoricalDtype(categories=label_mapping)
-    df[index_column_name] = pd.Categorical.from_codes(df[index_column_name], dtype=experiment_dtype)
     if timestamp_column_name:
         df[timestamp_column_name] = pd.to_datetime(df[timestamp_column_name], unit="ms", origin="unix", utc=True)
 
-    df = _pivot_and_reindex_df(df, False, index_column_name, timestamp_column_name)
+    df = _pivot_df(df, False, index_column_name, timestamp_column_name)
+    df = _restore_labels_in_index(df, index_column_name, label_mapping)
     df = _restore_path_column_names(df, path_mapping, None)
-    df = _sort_indices(df)
+    df = _sort_index_and_columns(df, index_column_name)
 
     return df
 
@@ -461,8 +459,6 @@ def generate_categorized_rows() -> Generator[Tuple, None, None]:
     df = pd.DataFrame(
         np.fromiter(generate_categorized_rows(), dtype=types),
     )
-    experiment_dtype = pd.CategoricalDtype(categories=label_mapping)
-    df[container_column_name] = pd.Categorical.from_codes(df[container_column_name], dtype=experiment_dtype)
 
     df["bucket"] = pd.IntervalIndex.from_arrays(df["from_x"], df["to_x"], closed="right")
     df = df.drop(columns=["from_x", "to_x"])
@@ -475,11 +471,8 @@ def generate_categorized_rows() -> Generator[Tuple, None, None]:
         dropna=False,
         sort=False,
     )
-    df.columns = df.columns.set_levels(
-        df.columns.get_level_values(container_column_name).unique().astype(str),
-        level=container_column_name,
-    )
 
+    df = _restore_labels_in_columns(df, container_column_name, label_mapping)
     df = _restore_path_column_names(df, path_mapping, None)
 
     # Clear out any columns that were not requested, but got added because of dropna=False
@@ -546,7 +539,7 @@ def _collapse_open_buckets(df: pd.DataFrame) -> pd.DataFrame:
     return df
 
 
-def _pivot_and_reindex_df(
+def _pivot_df(
     df: pd.DataFrame,
     include_point_previews: bool,
     index_column_name: str,
@@ -559,7 +552,7 @@ def _pivot_and_reindex_df(
         df[[index_column_name, "step"]]
         .astype(
             {
-                index_column_name: "category",
+                index_column_name: "uint32",
                 "step": "float64",
             }
         )
@@ -588,15 +581,31 @@ def _pivot_and_reindex_df(
         )
 
     # Include only observed (experiment, step) pairs
-    df = df.filter(observed_idx, axis="index")
+    return df.filter(observed_idx, axis="index")
 
-    # Replace categorical codes in `index_column_name` with strings
-    df.index = df.index.set_levels(
-        df.index.get_level_values(index_column_name).unique().astype(str),
-        level=index_column_name,
-    )
 
-    return df.sort_index(level=[index_column_name, "step"])
+def _restore_labels_in_index(
+    df: pd.DataFrame,
+    column_name: str,
+    label_mapping: list[str],
+) -> pd.DataFrame:
+    if df.index.empty:
+        df.index = df.index.set_levels(df.index.get_level_values(column_name).astype(str), level=column_name)
+        return df
+
+    return df.rename(index={i: label for i, label in enumerate(label_mapping)}, level=column_name)
+
+
+def _restore_labels_in_columns(
+    df: pd.DataFrame,
+    column_name: str,
+    label_mapping: list[str],
+) -> pd.DataFrame:
+    if df.index.empty:
+        df.columns = df.columns.set_levels(df.columns.get_level_values(column_name).astype(str), level=column_name)
+        return df
+
+    return df.rename(columns={i: label for i, label in enumerate(label_mapping)}, level=column_name)
 
 
 def _restore_path_column_names(
@@ -622,16 +631,18 @@ def _restore_path_column_names(
     return df.rename(columns=reverse_mapping)
 
 
-def _sort_indices(df: pd.DataFrame) -> pd.DataFrame:
+def _sort_index_and_columns(df: pd.DataFrame, index_column_name: str) -> pd.DataFrame:
     # MultiIndex DFs need to have column index order swapped: value/metric_name -> metric_name/value.
     # We also sort columns, but only after the original names have been restored.
     if isinstance(df.columns, pd.MultiIndex):
         df.columns.names = (None, None)
         df = df.swaplevel(axis="columns")
-        return df.sort_index(axis="columns", level=0)
+        df = df.sort_index(axis="columns", level=0)
     else:
         df.columns.name = None
-        return df.sort_index(axis="columns")
+        df = df.sort_index(axis="columns")
+
+    return df.sort_index(axis="index", level=[index_column_name, "step"])
 
 
 def create_files_dataframe(

tests/unit/internal/test_output_format.py

@@ -500,10 +500,13 @@ def _generate_bucket_metrics(
     }
 
 
+def _a_timestamp(seconds_delta) -> datetime:
+    return datetime(2023, 1, 1, 0, 0, 0, 0, timezone.utc) + timedelta(seconds=seconds_delta)
+
+
 def _generate_float_point_value(step: int, preview: bool) -> FloatPointValue:
-    timestamp = datetime(2023, 1, 1, 0, 0, 0, 0, timezone.utc) + timedelta(seconds=step)
     return (
-        timestamp.timestamp(),
+        _a_timestamp(seconds_delta=step).timestamp(),
         float(step),
         float(step) * 100,
         preview,
@@ -605,6 +608,84 @@ def test_create_metrics_dataframe_shape(include_preview):
     ), "The list of tuples should have the same number of rows as the DataFrame"
 
 
+def test_create_metrics_dataframe_from_exp_with_no_points():
+    df = create_metrics_dataframe(
+        # This input data produces a "hole" in our categorical mapping of experiment names to integers
+        metrics_data={
+            _generate_run_attribute_definition(1, 1): [_generate_float_point_value(1, False)],
+            _generate_run_attribute_definition(2, 2): [],
+            _generate_run_attribute_definition(3, 1): [_generate_float_point_value(2, False)],
+        },
+        sys_id_label_mapping={
+            SysId("sysid1"): "exp1",
+            SysId("sysid2"): "exp2",
+            SysId("sysid3"): "exp3",
+        },
+        include_point_previews=False,
+        type_suffix_in_column_names=False,
+        index_column_name="experiment",
+    )
+
+    expected_df = pd.DataFrame(
+        data={
+            "path1": [
+                100.0,
+                200.0,
+            ],
+        },
+        index=pd.MultiIndex.from_tuples(
+            tuples=[
+                ("exp1", 1.0),
+                ("exp3", 2.0),
+            ],
+            names=["experiment", "step"],
+        ),
+    )
+    pd.testing.assert_frame_equal(df, expected_df)
+
+
+def test_create_metrics_dataframe_from_exp_with_no_points_preview():
+    df = create_metrics_dataframe(
+        # This input data produces a "hole" in our categorical mapping of experiment names to integers
+        metrics_data={
+            _generate_run_attribute_definition(1, 1): [_generate_float_point_value(1, True)],
+            _generate_run_attribute_definition(2, 2): [],
+            _generate_run_attribute_definition(3, 1): [_generate_float_point_value(2, True)],
+        },
+        sys_id_label_mapping={
+            SysId("sysid1"): "exp1",
+            SysId("sysid2"): "exp2",
+            SysId("sysid3"): "exp3",
+        },
+        include_point_previews=True,
+        type_suffix_in_column_names=False,
+        index_column_name="experiment",
+    )
+
+    expected_df = pd.DataFrame(
+        data={
+            ("path1", "is_preview"): [
+                True,
+                True,
+            ],
+            ("path1", "preview_completion"): [
+                0.999,
+                0.998,
+            ],
+            ("path1", "value"): [100.0, 200.0],
+        },
+        index=pd.MultiIndex.from_tuples(
+            tuples=[
+                ("exp1", 1.0),
+                ("exp3", 2.0),
+            ],
+            names=["experiment", "step"],
+        ),
+    )
+    expected_df[("path1", "is_preview")] = expected_df[("path1", "is_preview")].astype("object")
+    pd.testing.assert_frame_equal(df, expected_df)
+
+
 @pytest.mark.parametrize("type_suffix_in_column_names", [True, False])
 @pytest.mark.parametrize("include_preview", [True, False])
 def test_create_metrics_dataframe_with_absolute_timestamp(type_suffix_in_column_names: bool, include_preview: bool):