fix: column dtypes in empty dfs... (#24)

Co-authored-by: Michał Sośnicki <michal.sosnicki@neptune.ai>

Author: michalsosn

src/neptune_query/internal/output_format.py

@@ -461,6 +461,14 @@ def _restore_path_column_names(
     Restores colum names in the DF based on the mapping.
     """
 
+    # No columns to rename, simply ensure the dtype of the path column changes from categorical int to str
+    if df.columns.empty:
+        if isinstance(df.columns, pd.MultiIndex):
+            df.columns = df.columns.set_levels(df.columns.get_level_values("path").astype(str), level="path")
+        else:
+            df.columns = df.columns.astype(str)
+        return df
+
     # We need to reverse the mapping to index -> column name
     if type_suffix:
         reverse_mapping = {index: f"{path}:{type_suffix}" for path, index in path_mapping.items()}

tests/unit/internal/test_output_format.py

@@ -930,9 +930,10 @@ def test_create_empty_metrics_dataframe(
 
     # With previews or timestamps, MultiIndex columns are returned
     if include_preview or timestamp_column_name:
-        expected_df.columns = pd.MultiIndex.from_tuples([], names=[None, None])
+        expected_df.columns = pd.MultiIndex.from_product([[], ["value"]], names=[None, None])
+        # the comparator seems not to delve into the exact column names on the 2nd level when the 1st level is empty..
 
-    pd.testing.assert_frame_equal(df, expected_df, check_column_type=False)
+    pd.testing.assert_frame_equal(df, expected_df)
 
 
 @pytest.mark.parametrize("timestamp_column_name", [None, "absolute"])
@@ -956,9 +957,9 @@ def test_create_empty_series_dataframe(timestamp_column_name: str):
     )
 
     if timestamp_column_name:
-        expected_df.columns = pd.MultiIndex.from_tuples([], names=[None, None])
+        expected_df.columns = pd.MultiIndex.from_product([[], ["value"]], names=[None, None])
 
-    pd.testing.assert_frame_equal(df, expected_df, check_index_type=False, check_column_type=False)
+    pd.testing.assert_frame_equal(df, expected_df)
 
 
 @pytest.mark.parametrize(