def extract_output(self, accumulator): feature_names_to_value_counts = dict() feature_names_to_weighted_value_counts = dict() for feature_name, value_counts in accumulator.items(): if value_counts.unweighted_counts: feature_value_counts = [ top_k_stats_generator.FeatureValueCount(key, value) for key, value in value_counts.unweighted_counts.items() ] feature_names_to_value_counts[ feature_name] = feature_value_counts if value_counts.weighted_counts: weighted_feature_value_counts = [ top_k_stats_generator.FeatureValueCount(key, value) for key, value in value_counts.weighted_counts.items() ] feature_names_to_weighted_value_counts[ feature_name] = weighted_feature_value_counts return _make_dataset_feature_stats_proto_with_multiple_features( feature_names_to_value_counts, feature_names_to_weighted_value_counts, self._categorical_features, self._num_top_values, self._num_rank_histogram_buckets)
def test_make_feature_stats_proto_with_topk_stats_weighted(self): expected_result = text_format.Parse( """ name: 'fa' type: STRING string_stats { weighted_string_stats { top_values { value: 'a' frequency: 4 } top_values { value: 'c' frequency: 3 } top_values { value: 'd' frequency: 2 } rank_histogram { buckets { low_rank: 0 high_rank: 0 label: "a" sample_count: 4.0 } buckets { low_rank: 1 high_rank: 1 label: "c" sample_count: 3.0 } } } }""", statistics_pb2.FeatureNameStatistics()) value_counts = [('a', 4), ('c', 3), ('d', 2), ('b', 2)] top_k_value_count_list = [ top_k_stats_generator.FeatureValueCount(value_count[0], value_count[1]) for value_count in value_counts ] result = top_k_stats_generator.make_feature_stats_proto_with_topk_stats( 'fa', top_k_value_count_list, False, True, 3, 2) compare.assertProtoEqual(self, result, expected_result)