feat: link crates referencing metric_doc macro at doc gen runtime

Author: aaron-ang

datafusion/core/build.rs

@@ -0,0 +1,133 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::env;
+use std::fs;
+use std::io::{Read, Write};
+use std::path::Path;
+
+fn main() {
+    let out_dir = env::var("OUT_DIR").unwrap();
+    let dest_path = Path::new(&out_dir).join("link_metrics.rs");
+    let mut f = fs::File::create(&dest_path).unwrap();
+
+    // Identify crates that use #[metric_doc]
+    // We scan sibling directories in `datafusion/`
+    let core_crate_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
+    let datafusion_dir = Path::new(&core_crate_dir).parent().unwrap();
+
+    let mut crates_with_metrics = Vec::new();
+
+    // Read Cargo.toml to check dependencies
+    let cargo_toml_path = Path::new(&core_crate_dir).join("Cargo.toml");
+    let mut cargo_toml_content = String::new();
+    if let Ok(mut file) = fs::File::open(&cargo_toml_path) {
+        file.read_to_string(&mut cargo_toml_content).unwrap();
+    }
+
+    if let Ok(entries) = fs::read_dir(datafusion_dir) {
+        for entry in entries.flatten() {
+            let path = entry.path();
+            if path.is_dir() {
+                // Check if it's a crate (has Cargo.toml)
+                if path.join("Cargo.toml").exists() {
+                    let dir_name = path.file_name().unwrap().to_str().unwrap();
+                    let crate_name = format!("datafusion-{}", dir_name);
+
+                    // Skip self (core) and macros (definition only)
+                    if dir_name == "core" || dir_name == "macros" {
+                        continue;
+                    }
+
+                    // Skip if not a dependency in Cargo.toml
+                    // This is a rough check
+                    if !cargo_toml_content.contains(&format!("{} =", crate_name))
+                        && !cargo_toml_content.contains(&format!("\"{}\"", crate_name))
+                    {
+                        continue;
+                    }
+
+                    if has_metric_doc(&path) {
+                        crates_with_metrics.push(crate_name);
+                    }
+                }
+            }
+        }
+    }
+
+    // Sort for deterministic output
+    crates_with_metrics.sort();
+
+    writeln!(
+        f,
+        "/// Automatically generated by build.rs to link crates with metrics"
+    )
+    .unwrap();
+    writeln!(f, "pub fn link_metrics() {{").unwrap();
+    for krate in crates_with_metrics {
+        let krate_snake = krate.replace("-", "_");
+
+        let is_optional = cargo_toml_content
+            .lines()
+            .any(|line| line.contains(&krate) && line.contains("optional = true"));
+
+        if is_optional {
+            writeln!(f, "    #[cfg(feature = \"{}\")]", krate_snake).unwrap();
+        }
+
+        writeln!(f, "    {{").unwrap();
+        writeln!(f, "        // Force link {}", krate).unwrap();
+        writeln!(f, "        #[allow(unused_imports)]").unwrap();
+        writeln!(f, "        use {} as _;", krate_snake).unwrap();
+        writeln!(f, "    }}").unwrap();
+    }
+    writeln!(f, "}}").unwrap();
+
+    println!("cargo:rerun-if-changed=build.rs");
+    println!("cargo:rerun-if-changed=Cargo.toml");
+}
+
+fn has_metric_doc(dir: &Path) -> bool {
+    if let Ok(entries) = fs::read_dir(dir) {
+        for entry in entries.flatten() {
+            let path = entry.path();
+            if path.is_dir() {
+                // Skip target, .git, etc
+                if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
+                    if name == "target" || name.starts_with('.') {
+                        continue;
+                    }
+                }
+                if has_metric_doc(&path) {
+                    return true;
+                }
+            } else if let Some(ext) = path.extension() {
+                if ext == "rs" {
+                    if let Ok(mut file) = fs::File::open(&path) {
+                        let mut content = String::new();
+                        if file.read_to_string(&mut content).is_ok() {
+                            if content.contains("#[metric_doc") {
+                                return true;
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+    false
+}

datafusion/core/src/bin/print_metric_docs.rs

@@ -23,10 +23,10 @@ use std::collections::HashSet;
 use datafusion_doc::metric_doc_sections::{
     ExecDoc, MetricDoc, MetricDocPosition, exec_docs, metric_docs,
 };
-use datafusion_execution as _; // Link metrics defined in execution crate.
-use datafusion_physical_plan as _; // Link metrics and execs defined in physical plan.
 
 fn main() -> std::io::Result<()> {
+    datafusion::doc::link_metrics();
+
     let mut content = String::new();
     let mut metrics: Vec<&MetricDoc> = metric_docs().collect();
     metrics.sort_by(|a, b| a.name.cmp(b.name));

datafusion/core/src/lib.rs

@@ -900,6 +900,11 @@ pub mod test;
 mod schema_equivalence;
 pub mod test_util;
 
+/// Documentation generation helpers
+pub mod doc {
+    include!(concat!(env!("OUT_DIR"), "/link_metrics.rs"));
+}
+
 #[cfg(doctest)]
 doc_comment::doctest!("../../../README.md", readme_example_test);