diff --git a/sdks/python/apache_beam/metrics/metric_test.py b/sdks/python/apache_beam/metrics/metric_test.py index ae66200737b5..90c5a144ca2b 100644 --- a/sdks/python/apache_beam/metrics/metric_test.py +++ b/sdks/python/apache_beam/metrics/metric_test.py @@ -57,6 +57,29 @@ def test_basic_metric_name(self): self.assertEqual(key.metric.name, 'name1') self.assertEqual(key, MetricKey('step1', MetricName('namespace1', 'name1'))) + def test_metric_name_pickle_and_legacy_dict_state(self): + import pickle + name = MetricName('namespace1', 'name1', urn='urn1', labels={'k': 'v'}) + _ = hash(name) + _ = name.fast_name() + unpickled = pickle.loads(pickle.dumps(name)) + self.assertEqual(unpickled, name) + self.assertEqual(hash(unpickled), hash(name)) + self.assertEqual(unpickled.fast_name(), name.fast_name()) + + legacy = MetricName.__new__(MetricName) + legacy.__setstate__({ + 'namespace': 'namespace1', + 'name': 'name1', + 'urn': 'urn1', + 'labels': { + 'k': 'v' + }, + }) + self.assertEqual(legacy, name) + self.assertEqual(hash(legacy), hash(name)) + self.assertEqual(legacy.fast_name(), name.fast_name()) + class MetricResultsTest(unittest.TestCase): def test_metric_filter_namespace_matching(self): diff --git a/sdks/python/apache_beam/metrics/metricbase.py b/sdks/python/apache_beam/metrics/metricbase.py index addbf14153cf..d19a1f1c76c3 100644 --- a/sdks/python/apache_beam/metrics/metricbase.py +++ b/sdks/python/apache_beam/metrics/metricbase.py @@ -82,9 +82,11 @@ def __init__( self.name = name self.urn = urn self.labels = labels if labels else {} + self._cached_fast_name: Optional[str] = None + self._cached_hash: Optional[int] = None def __eq__(self, other): - return ( + return self is other or ( self.namespace == other.namespace and self.name == other.name and self.urn == other.urn and self.labels == other.labels) @@ -97,17 +99,38 @@ def __str__(self): self.namespace, self.name) def __hash__(self): - return hash((self.namespace, self.name, self.urn) + - tuple(self.labels.items())) + if self._cached_hash is None: + self._cached_hash = hash( + (self.namespace, self.name, self.urn) + + tuple(self.labels.items())) if self.labels else hash( + (self.namespace, self.name, self.urn)) + return self._cached_hash def fast_name(self): - name = self.name or '' - namespace = self.namespace or '' - urn = self.urn or '' - labels = '' - if self.labels: - labels = '_'.join(['%s=%s' % (k, v) for (k, v) in self.labels.items()]) - return '%d_%s%s%s%s' % (len(name), name, namespace, urn, labels) + if self._cached_fast_name is None: + name = self.name or '' + namespace = self.namespace or '' + urn = self.urn or '' + labels = '' + if self.labels: + labels = '_'.join(['%s=%s' % (k, v) for (k, v) in self.labels.items()]) + self._cached_fast_name = '%d_%s%s%s%s' % ( + len(name), name, namespace, urn, labels) + return self._cached_fast_name + + def __getstate__(self): + return (self.namespace, self.name, self.urn, self.labels) + + def __setstate__(self, state): + if isinstance(state, dict): + self.namespace = state['namespace'] + self.name = state['name'] + self.urn = state['urn'] + self.labels = state['labels'] + else: + self.namespace, self.name, self.urn, self.labels = state + self._cached_fast_name = None + self._cached_hash = None class Metric(object):