Add timing context manager
@@ -2,6 +2,7 @@ import json
import time
import weakref
from collections import OrderedDict
+from contextlib import contextmanager
import numpy as np
@@ -22,6 +23,18 @@ except ImportError:
DataFrame = None
+@contextmanager
+def timing(key):
+ if not use_stats():
+ yield
+ else:
+ start = time.time()
+ try:
+ yield
+ finally:
+ increment_stat(key, time.time() - start)
+
+
def mr_dense_result_shape(ranges, base_shape=None):
if base_shape is not None:
assert len(ranges) == len(base_shape), "internal error: mismatched shapes"
@@ -132,7 +145,7 @@ class DataFrameIndexer(MultiRangeIndexer):
def __getitem__(self, idx):
check_dataframe_deps()
- idx_start = time.time()
+ with timing("py.__getitem__time"):
array = self.array
# we need to use a Query in order to get coords for a dense array
@@ -141,8 +154,7 @@ class DataFrameIndexer(MultiRangeIndexer):
idx, array, query, use_arrow=self.use_arrow, preload_metadata=True
)
- pd_start = time.time()
-
+ with timing("py.buffer_conversion_time"):
if isinstance(result, dict):
df = DataFrame.from_dict(result)
elif isinstance(result, PyQuery):
@@ -154,13 +166,9 @@ class DataFrameIndexer(MultiRangeIndexer):
raise TypeError(f"Unhandled result type {type(result)}")
if isinstance(df, DataFrame):
+ with timing("py.pandas_index_update_time"):
df = _update_df_from_meta(df, array.meta, query.index_col)
- if use_stats():
- end = time.time()
- increment_stat("py.buffer_conversion_time", end - pd_start)
- increment_stat("py.__getitem__time", end - idx_start)
-
return df
@@ -252,8 +260,6 @@ def _get_pyquery_results(pyquery, schema):
def _update_df_from_meta(df, array_meta, index_col=True):
- start = time.time()
-
col_dtypes = {}
if "__pandas_attribute_repr" in array_meta:
col_dtypes.update(json.loads(array_meta["__pandas_attribute_repr"]))
@@ -289,9 +295,6 @@ def _update_df_from_meta(df, array_meta, index_col=True):
if df.index.name != index_name:
df.index.rename(index_name, inplace=True)
- if use_stats():
- increment_stat("py.pandas_index_update_time", time.time() - start)
-
return df
