Emit only interned dimension keys in the metric bundle output
diff --git a/protos/perfetto/trace_summary/v2_metric.proto b/protos/perfetto/trace_summary/v2_metric.proto
index 481c63c..ea336ba 100644
--- a/protos/perfetto/trace_summary/v2_metric.proto
+++ b/protos/perfetto/trace_summary/v2_metric.proto
@@ -334,18 +334,23 @@
   // Specification for a interned dimension associated with this metric.
   // For a given dimension, there can only be one interned dimension attached to
   // it.
+  // Interned dimensions are useful for providing additional information about
+  // a dimension without adding that information as a dimension to every metric
+  // row. This is particularly useful when you have columns associated with a
+  // dimension that would otherwise blow up the proto size due to repetition. By
+  // using interned dimensions, you can emit only the necessary keys, reducing
+  // data redundancy.
   //
-  // Interned dimensions are useful when you have a bunch of columns
-  // associated with a single dimension, and repeating them many times would
-  // blow up the proto size. Using interned dimensions, you can emit
-  // only an ID in the key and then "expand" to match, reducing data
-  // redundancy.
-  // For example, if a metric is dimensioned by package name, we can use
-  // interned dimension to provide APK version code for each package name
-  // without adding version code as a dimension to every row.
+  // For example, if a metric is dimensioned by package name, you can use an
+  // interned dimension to provide the APK version code for each package without
+  // adding the version code as a dimension to every row.
   //
-  // Note: The query for an interned dimension must return a unique row per
-  // dimension value, as only one row per input dimension is allowed.
+  // Note: To optimize the output size, only the keys that are actually present
+  // in the main metric query's output will be included in the interned
+  // dimension bundles.
+  //
+  // The query for an interned dimension must return a unique row per dimension
+  // value, as only one row per input dimension is allowed.
   //
   // Example:
   // ```
diff --git a/src/trace_processor/trace_summary/summary.cc b/src/trace_processor/trace_summary/summary.cc
index bcb8e77..d78d724 100644
--- a/src/trace_processor/trace_summary/summary.cc
+++ b/src/trace_processor/trace_summary/summary.cc
@@ -67,6 +67,23 @@
 using PerfettoSqlStructuredQuery = protos::pbzero::PerfettoSqlStructuredQuery;
 using InternedDimensionSpec = TraceMetricV2Spec::InternedDimensionSpec;
 
+uint64_t HashOf(const SqlValue& val) {
+  base::FnvHasher hasher;
+  hasher.Update(val.type);
+  if (val.is_null()) {
+  } else if (val.type == SqlValue::kLong) {
+    hasher.Update(val.long_value);
+  } else if (val.type == SqlValue::kDouble) {
+    hasher.Update(val.double_value);
+  } else if (val.type == SqlValue::kString) {
+    hasher.Update(val.string_value);
+  } else {
+    PERFETTO_FATAL("Unsupported SqlValue type %d for hashing",
+                   static_cast<int>(val.type));
+  }
+  return hasher.digest();
+}
+
 struct Metric {
   std::string id;
   std::string query;
@@ -526,26 +543,12 @@
   return base::OkStatus();
 }
 
-base::Status HashKeyAndInsert(const SqlValue& key_val,
-                              base::StringView key_column_name,
-                              base::FlatHashMap<uint64_t, bool>& seen_keys) {
-  base::FnvHasher key_hasher;
-  key_hasher.Update(key_val.type);
-  if (key_val.is_null()) {
-  } else if (key_val.type == SqlValue::kLong) {
-    key_hasher.Update(key_val.long_value);
-  } else if (key_val.type == SqlValue::kDouble) {
-    key_hasher.Update(key_val.double_value);
-  } else if (key_val.type == SqlValue::kString) {
-    key_hasher.Update(key_val.string_value);
-  } else {
-    return base::ErrStatus(
-        "Unsupported key type %d for interned dimension bundle with key "
-        "column '%.*s'",
-        static_cast<int>(key_val.type),
-        static_cast<int>(key_column_name.size()), key_column_name.data());
-  }
-  if (!seen_keys.Insert(key_hasher.digest(), true).second) {
+// Renamed from HashKeyAndInsert.
+base::Status CheckForDuplicateAndInsert(
+    uint64_t hash,
+    base::StringView key_column_name,
+    base::FlatHashMap<uint64_t, bool>& seen_keys) {
+  if (!seen_keys.Insert(hash, true).second) {
     return base::ErrStatus(
         "Duplicate key found in interned dimension bundle with key column "
         "'%.*s'",
@@ -558,6 +561,8 @@
     TraceProcessor* processor,
     const TraceMetricV2Spec::Decoder& spec,
     const std::vector<std::string>& interned_dimension_queries,
+    const base::FlatHashMap<std::string, base::FlatHashMap<uint64_t, bool>>&
+        interned_dim_keys_in_metric_bundle,
     TraceMetricV2Bundle* bundle) {
   RETURN_IF_ERROR(ValidateInternedDimensionSpecs(spec));
   size_t interned_dimension_idx = 0;
@@ -566,6 +571,8 @@
     InternedDimensionSpec::ColumnSpec::Decoder key_col_spec(
         ms.key_column_spec());
     base::StringView key_column_name = key_col_spec.name();
+    const auto& allowed_keys =
+        *interned_dim_keys_in_metric_bundle.Find(key_column_name.ToStdString());
     auto* interned_dimension_bundle = bundle->add_interned_dimension_bundles();
 
     const std::string& sql =
@@ -613,7 +620,15 @@
     base::FlatHashMap<uint64_t, bool> seen_keys;
     while (query_it.Next()) {
       const auto& key_val = query_it.Get(column_infos[0].first);
-      RETURN_IF_ERROR(HashKeyAndInsert(key_val, key_column_name, seen_keys));
+      uint64_t hash = HashOf(key_val);
+
+      // If the key was not in the metric bundle, we don't need to output
+      // its interned data.
+      if (!allowed_keys.Find(hash)) {
+        continue;
+      }
+      RETURN_IF_ERROR(
+          CheckForDuplicateAndInsert(hash, key_column_name, seen_keys));
       auto* row = interned_dimension_bundle->add_interned_dimension_rows();
       RETURN_IF_ERROR(WriteInternedDimensionValue(
           query_it.Get(column_infos[0].first), column_infos[0].second,
@@ -658,6 +673,18 @@
     const Metric* first = value.front();
     TraceMetricV2Spec::Decoder first_spec(first->spec);
 
+    base::FlatHashMap<std::string, base::FlatHashMap<uint64_t, bool>>
+        interned_dim_keys_in_metric_bundle;
+    for (auto ms_it = first_spec.interned_dimension_specs(); ms_it; ++ms_it) {
+      InternedDimensionSpec::Decoder ms_decoder(*ms_it);
+      interned_dim_keys_in_metric_bundle.Insert(
+          InternedDimensionSpec::ColumnSpec::Decoder(
+              ms_decoder.key_column_spec())
+              .name()
+              .ToStdString(),
+          base::FlatHashMap<uint64_t, bool>());
+    }
+
     auto query_it = processor->ExecuteQuery(first->query);
     if (!query_it.Status().ok()) {
       return base::ErrStatus(
@@ -734,6 +761,10 @@
       for (const auto& dim : dimensions_with_index) {
         RETURN_IF_ERROR(WriteDimension(dim, bundle_id, query_it,
                                        row->add_dimension(), &hasher));
+        if (auto* key_set = interned_dim_keys_in_metric_bundle.Find(dim.name)) {
+          const auto& key_val = query_it.Get(dim.index);
+          key_set->Insert(HashOf(key_val), true);
+        }
       }
       uint64_t hash = hasher.digest();
       if (is_unique_dimensions && !seen_dimensions.Insert(hash, true).second) {
@@ -772,7 +803,8 @@
     RETURN_IF_ERROR(query_it.Status());
     if (!first->interned_dimension_queries.empty()) {
       RETURN_IF_ERROR(WriteInternedDimensionBundles(
-          processor, first_spec, first->interned_dimension_queries, bundle));
+          processor, first_spec, first->interned_dimension_queries,
+          interned_dim_keys_in_metric_bundle, bundle));
     }
   }
   return base::OkStatus();
diff --git a/src/trace_processor/trace_summary/summary_integrationtest.cc b/src/trace_processor/trace_summary/summary_integrationtest.cc
index 2eb186f..2cef769 100644
--- a/src/trace_processor/trace_summary/summary_integrationtest.cc
+++ b/src/trace_processor/trace_summary/summary_integrationtest.cc
@@ -907,6 +907,129 @@
   )-"));
 }
 
+TEST_F(TraceSummaryTest, InternedDimensionBundleUnusedKeysDropped) {
+  ASSERT_OK_AND_ASSIGN(auto output, RunSummarize(R"(
+    metric_template_spec {
+      id_prefix: "my_metric"
+      value_columns: "dur"
+      dimensions_specs { name: "dim_1" type: INT64 }
+      dimensions_specs { name: "dim_2" type: INT64 }
+      query {
+        sql {
+          sql: "SELECT 123 as dim_1, 123 as dim_2, 750.0 as dur UNION ALL SELECT 456 as dim_1, 789 as dim_2, 850.0 as dur"
+          column_names: "dim_1"
+          column_names: "dim_2"
+          column_names: "dur"
+        }
+      }
+      interned_dimension_specs {
+        key_column_spec { name: "dim_1" type: INT64 }
+        data_column_specs { name: "version_1" type: INT64 }
+        query {
+          sql {
+            sql: "SELECT 123 as dim_1, 100 as version_1 UNION ALL SELECT 456 as dim_1, 200 as version_1 UNION ALL SELECT 789 as dim_1, 300 as version_1"
+          }
+        }
+      }
+      interned_dimension_specs {
+        key_column_spec { name: "dim_2" type: INT64 }
+        data_column_specs { name: "version_2" type: INT64 }
+        query {
+          sql {
+            sql: "SELECT 123 as dim_2, 1000 as version_2 UNION ALL SELECT 456 as dim_2, 2000 as version_2 UNION ALL SELECT 789 as dim_2, 3000 as version_2"
+          }
+        }
+      }
+    }
+  )"));
+  EXPECT_THAT(output, EqualsIgnoringWhitespace(R"-(
+    metric_bundles {
+      bundle_id: "my_metric"
+      specs {
+        id: "my_metric_dur"
+        value: "dur"
+        dimensions_specs {
+          name: "dim_1"
+          type: INT64
+        }
+        dimensions_specs {
+          name: "dim_2"
+          type: INT64
+        }
+        query {
+          sql {
+            sql: "SELECT 123 as dim_1, 123 as dim_2, 750.0 as dur UNION ALL SELECT 456 as dim_1, 789 as dim_2, 850.0 as dur"
+            column_names: "dim_1"
+            column_names: "dim_2"
+            column_names: "dur"
+          }
+        }
+        bundle_id: "my_metric"
+        interned_dimension_specs {
+          key_column_spec {
+            name: "dim_1"
+            type: INT64
+          }
+          data_column_specs {
+            name: "version_1"
+            type: INT64
+          }
+          query {
+            sql {
+              sql: "SELECT 123 as dim_1, 100 as version_1 UNION ALL SELECT 456 as dim_1, 200 as version_1 UNION ALL SELECT 789 as dim_1, 300 as version_1"
+            }
+          }
+        }
+        interned_dimension_specs {
+          key_column_spec {
+            name: "dim_2"
+            type: INT64
+          }
+          data_column_specs {
+            name: "version_2"
+            type: INT64
+          }
+          query {
+            sql {
+              sql: "SELECT 123 as dim_2, 1000 as version_2 UNION ALL SELECT 456 as dim_2, 2000 as version_2 UNION ALL SELECT 789 as dim_2, 3000 as version_2"
+            }
+          }
+        }
+      }
+      row {
+        dimension { int64_value: 123 }
+        dimension { int64_value: 123 }
+        values { double_value: 750.000000 }
+      }
+      row {
+        dimension { int64_value: 456 }
+        dimension { int64_value: 789 }
+        values { double_value: 850.000000 }
+      }
+      interned_dimension_bundles {
+        interned_dimension_rows {
+          key_dimension_value { int64_value: 123 }
+          interned_dimension_values { int64_value: 100 }
+        }
+        interned_dimension_rows {
+          key_dimension_value { int64_value: 456 }
+          interned_dimension_values { int64_value: 200 }
+        }
+      }
+      interned_dimension_bundles {
+        interned_dimension_rows {
+          key_dimension_value { int64_value: 123 }
+          interned_dimension_values { int64_value: 1000 }
+        }
+        interned_dimension_rows {
+          key_dimension_value { int64_value: 789 }
+          interned_dimension_values { int64_value: 3000 }
+        }
+      }
+    }
+  )-"));
+}
+
 TEST_F(TraceSummaryTest, TemplateSpecWithValueColumnsAndSpecsError) {
   base::StatusOr<std::string> status_or_output = RunSummarize(R"(
     metric_template_spec {