Fix output formatting logic

Author: pitercl

src/neptune_query/internal/output_format.py

@@ -406,14 +406,28 @@ def _pivot_and_reindex_df(
     index_column_name: str,
     timestamp_column_name: Optional[str],
 ) -> pd.DataFrame:
+    # Holds all existing (experiment, step) pairs
+    # This is needed because pivot_table will add rows for all combinations of (experiment, step)
+    # even if they don't exist in the original data, filling the rows with NaNs.
+    observed_idx = pd.MultiIndex.from_frame(
+        df[[index_column_name, "step"]]
+        .astype(
+            {
+                index_column_name: "category",
+                "step": "float64",
+            }
+        )
+        .drop_duplicates()
+    )
+
     if df.empty and timestamp_column_name:
         # Handle empty DataFrame case to avoid pandas dtype errors
         df[timestamp_column_name] = pd.Series(dtype="datetime64[ns]")
 
     if include_point_previews or timestamp_column_name:
         # if there are multiple value columns, don't specify them and rely on pandas to create the column multi-index
         df = df.pivot_table(
-            index=[index_column_name, "step"], columns="path", aggfunc="first", observed=True, dropna=False
+            index=[index_column_name, "step"], columns="path", aggfunc="first", observed=True, dropna=False, sort=False
         )
     else:
         # when there's only "value", define values explicitly, to make pandas generate a flat index
@@ -424,14 +438,19 @@ def _pivot_and_reindex_df(
             aggfunc="first",
             observed=True,
             dropna=False,
+            sort=False,
         )
 
-    df = df.reset_index()
-    df[index_column_name] = df[index_column_name].astype(str)
-    df = df.sort_values(by=[index_column_name, "step"], ignore_index=True)
-    df = df.set_index([index_column_name, "step"])
+    # Include only observed (experiment, step) pairs
+    df = df.reindex(index=observed_idx)
 
-    return df
+    # 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_path_column_names(

tests/unit/internal/test_output_format.py

@@ -911,8 +911,6 @@ def test_create_metrics_dataframe_random_order():
 def test_create_empty_metrics_dataframe(
     type_suffix_in_column_names: bool, include_preview: bool, timestamp_column_name: str
 ):
-    # Given empty dataframe
-
     # When
     df = create_metrics_dataframe(
         metrics_data={},
@@ -924,21 +922,17 @@ def test_create_empty_metrics_dataframe(
     )
 
     # Then
+    expected_df = (
+        pd.DataFrame(data={"experiment": [], "step": []})
+        .astype(dtype={"experiment": "object", "step": "float64"})
+        .set_index(["experiment", "step"])
+    )
+
+    # With previews or timestamps, MultiIndex columns are returned
     if include_preview or timestamp_column_name:
-        expected_df = pd.DataFrame(
-            index=pd.MultiIndex.from_tuples([], names=["experiment", "step"]),
-            columns=pd.MultiIndex.from_tuples([], names=["path", "metric"]),  # Create empty MultiIndex for columns
-        )
-        expected_df.columns.names = None, None
-    else:
-        expected_df = pd.DataFrame(
-            {
-                "experiment": [],
-                "step": [],
-            }
-        ).set_index(["experiment", "step"])
+        expected_df.columns = pd.MultiIndex.from_tuples([], names=[None, None])
 
-    pd.testing.assert_frame_equal(df, expected_df, check_index_type=False)
+    pd.testing.assert_frame_equal(df, expected_df, check_column_type=False)
 
 
 @pytest.mark.parametrize("timestamp_column_name", [None, "absolute"])
@@ -955,21 +949,16 @@ def test_create_empty_series_dataframe(timestamp_column_name: str):
     )
 
     # Then
+    expected_df = (
+        pd.DataFrame(data={"experiment": [], "step": []})
+        .astype(dtype={"experiment": "object", "step": "float64"})
+        .set_index(["experiment", "step"])
+    )
+
     if timestamp_column_name:
-        expected_df = pd.DataFrame(
-            index=pd.MultiIndex.from_tuples([], names=["experiment", "step"]),
-            columns=pd.MultiIndex.from_tuples([], names=["path", "metric"]),  # Create empty MultiIndex for columns
-        )
-        expected_df.columns.names = None, None
-    else:
-        expected_df = pd.DataFrame(
-            {
-                "experiment": [],
-                "step": [],
-            }
-        ).set_index(["experiment", "step"])
+        expected_df.columns = pd.MultiIndex.from_tuples([], names=[None, None])
 
-    pd.testing.assert_frame_equal(df, expected_df, check_index_type=False)
+    pd.testing.assert_frame_equal(df, expected_df, check_index_type=False, check_column_type=False)
 
 
 @pytest.mark.parametrize(