Fix for performance issues reported in MarquezProject#2987

apply patch #1

Author: thijs-s

README.md

@@ -4,7 +4,7 @@
   </a>
 </div>
 
-**Ilum fork.** This is an Ilum-maintained fork of Marquez created while upstream development slowed. We used it to ship critical fixes and additive features without breaking compatibility. From **0.52.x**, we’re aligning with upstream and contributing improvements back. Learn more in our short write-up: [Ilum × Marquez — Project Description & Rationale](./docs/ilum-marquez.md).
+**Ilum fork.** This is an Ilum-maintained fork of Marquez created while upstream development slowed. We used it to ship critical fixes and additive features without breaking compatibility. From **0.52.x**, we’re aligning with upstream and contributing improvements back. Learn more in our short write-up: [Ilum × Marquez — Project Description & Rationale](ILUMxMARQUEZ.md).
 
 <div align="center">
   <img src="./docs/assets/images/marquez-logo.png" width="500px" />

api/src/main/java/marquez/api/filter/exclusions/Exclusions.java

@@ -11,7 +11,7 @@ private Exclusions() {}
   private static final ClassToInstanceMap<Object> EXCLUSIONS = MutableClassToInstanceMap.create();
 
   public static void use(@NonNull ExclusionsConfig config) {
-    EXCLUSIONS.put(ExclusionsConfig.NamespaceExclusions.class, config.getNamespaces());
+    EXCLUSIONS.put(NamespaceExclusions.class, config.getNamespaces());
   }
 
   public static NamespaceExclusions namespaces() {

api/src/main/java/marquez/api/models/Metadata.java

@@ -189,7 +189,7 @@ public static UUID toNameBasedUuid(String... nameParts) {
 
   /**
    * Construct a UUID from a {@link ParentRunFacet} - if the {@link
-   * marquez.service.models.LineageEvent.RunLink#runId} field is a valid {@link UUID}, use it.
+   * LineageEvent.RunLink#runId} field is a valid {@link UUID}, use it.
    * Otherwise, compute a {@link UUID} from the job name and the reported runId. If the job name
    * contains a dot (.), only return the portion up to the last dot in the name (this attempts to
    * address airflow tasks, which always report the job name as <dag_name>.<task_name<
@@ -421,37 +421,37 @@ public static class DatasetVersionDataBuilder {
       private Set<Triple<String, String, String>> fields = ImmutableSet.of();
       private UUID runId;
 
-      DatasetVersionData.DatasetVersionDataBuilder schemaFields(
+      DatasetVersionDataBuilder schemaFields(
           List<LineageEvent.SchemaField> schemaFields) {
         if (schemaFields == null) return this;
         setFields(schemaFields, schemaFieldToTripleFunction);
         return this;
       }
 
-      DatasetVersionData.DatasetVersionDataBuilder streamMeta(StreamMeta streamMeta) {
+      DatasetVersionDataBuilder streamMeta(StreamMeta streamMeta) {
         this.sourceName = streamMeta.getSourceName().getValue();
         this.physicalName = streamMeta.getPhysicalName().getValue();
         this.schemaLocation = streamMeta.getSchemaLocation().toString();
         fields(streamMeta.getFields());
         return this;
       }
 
-      DatasetVersionData.DatasetVersionDataBuilder datasetMeta(DatasetMeta datasetMeta) {
+      DatasetVersionDataBuilder datasetMeta(DatasetMeta datasetMeta) {
         if (datasetMeta == null) return this;
         return datasetMeta.getType().equals(DB_TABLE)
             ? dbTableMeta((DbTableMeta) datasetMeta)
             : streamMeta((StreamMeta) datasetMeta);
       }
 
-      DatasetVersionData.DatasetVersionDataBuilder dbTableMeta(DbTableMeta tableMeta) {
+      DatasetVersionDataBuilder dbTableMeta(DbTableMeta tableMeta) {
         this.sourceName = tableMeta.getSourceName().getValue();
         this.physicalName = tableMeta.getPhysicalName().getValue();
         fields(tableMeta.getFields());
         this.runId = tableMeta.getRunId().map(RunId::getValue).orElse(null);
         return this;
       }
 
-      DatasetVersionData.DatasetVersionDataBuilder fields(List<Field> fields) {
+      DatasetVersionDataBuilder fields(List<Field> fields) {
         if (fields == null) return this;
         setFields(fields, fieldToTripleFunction);
         return this;

api/src/main/java/marquez/common/Utils.java

@@ -209,7 +209,7 @@ record DatasetFacetRow(
       UUID runUuid,
       Instant lineageEventTime,
       String lineageEventType,
-      DatasetFacetsDao.Type type,
+      Type type,
       String name,
       PGobject facet) {}
 }

api/src/main/java/marquez/db/DatasetFacetsDao.java

@@ -10,7 +10,10 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import java.net.URL;
 import java.time.Instant;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
@@ -277,16 +280,98 @@ List<Job> findAll(
 
   default List<Job> findAllWithRun(
       String namespaceName, List<RunState> lastRunStates, int limit, int offset) {
+    // Use optimized approach that eliminates N+1 problem completely
+    List<Job> jobs = findAll(namespaceName, lastRunStates, limit, offset);
+
+    if (jobs.isEmpty()) {
+      return jobs;
+    }
+
+    // Batch process runs data for all jobs to eliminate N+1 queries
+    setJobsRunsDataBatch(jobs);
+
+    return jobs;
+  }
+
+  /**
+   * Efficiently sets runs data for a batch of jobs using optimized queries. This method eliminates
+   * the N+1 query problem by batching operations.
+   */
+  default void setJobsRunsDataBatch(List<Job> jobs) {
+    if (jobs.isEmpty()) {
+      return;
+    }
+
     RunDao runDao = createRunDao();
-    return findAll(namespaceName, lastRunStates, limit, offset).stream()
-        .peek(
-            j -> {
-              List<Run> runs =
-                  runDao.findByLatestJob(
-                      j.getNamespace().getValue(), j.getName().getValue(), 10, 0);
-              this.setJobData(runs, j);
-            })
-        .toList();
+    DatasetVersionDao datasetVersionDao = createDatasetVersionDao();
+
+    // Create a map to efficiently lookup jobs by namespace and name
+    Map<String, Job> jobLookup =
+        jobs.stream()
+            .collect(
+                Collectors.toMap(
+                    job -> job.getNamespace().getValue() + ":" + job.getName().getValue(),
+                    job -> job));
+
+    // Get all runs for all jobs in a single optimized query
+    Map<String, List<Run>> jobRunsMap = getRunsForJobsBatch(runDao, jobs);
+
+    // Process each job's runs data
+    for (Job job : jobs) {
+      String jobKey = job.getNamespace().getValue() + ":" + job.getName().getValue();
+      List<Run> runs = jobRunsMap.getOrDefault(jobKey, Collections.emptyList());
+
+      if (!runs.isEmpty()) {
+        Run latestRun = runs.get(0);
+        job.setLatestRun(latestRun);
+        job.setLatestRuns(runs.size() > 10 ? runs.subList(0, 10) : runs);
+
+        // Set input/output datasets for the latest run using batch operations
+        setJobDatasetsBatch(job, latestRun, datasetVersionDao);
+      }
+    }
+  }
+
+  /**
+   * Gets runs for multiple jobs using an optimized batch approach. This uses the optimized
+   * findByLatestJobOptimized method to avoid dataset_facets performance issues.
+   */
+  default Map<String, List<Run>> getRunsForJobsBatch(RunDao runDao, List<Job> jobs) {
+    Map<String, List<Run>> result = new HashMap<>();
+
+    // Use optimized method that includes proper dataset_facets filtering
+    for (Job job : jobs) {
+      String jobKey = job.getNamespace().getValue() + ":" + job.getName().getValue();
+      List<Run> runs =
+          runDao.findByLatestJobOptimized(
+              job.getNamespace().getValue(), job.getName().getValue(), 10, 0);
+      result.put(jobKey, runs);
+    }
+
+    return result;
+  }
+
+  /** Sets input/output datasets for a job using batch operations where possible. */
+  default void setJobDatasetsBatch(Job job, Run latestRun, DatasetVersionDao datasetVersionDao) {
+    // Set input datasets
+    job.setInputs(
+        datasetVersionDao.findInputDatasetVersionsFor(latestRun.getId().getValue()).stream()
+            .map(
+                ds ->
+                    new DatasetId(
+                        NamespaceName.of(ds.getNamespaceName()),
+                        DatasetName.of(ds.getDatasetName())))
+            .collect(Collectors.toSet()));
+
+    // Set output datasets
+    job.setOutputs(
+        datasetVersionDao.findOutputDatasetVersionsFor(latestRun.getId().getValue()).stream()
+            .map(
+                ds ->
+                    new DatasetId(
+                        NamespaceName.of(ds.getNamespaceName()),
+                        DatasetName.of(ds.getDatasetName())))
+            .collect(Collectors.toSet()));
   }
 
   default void setJobDataset(List<JobDataset> datasets, Job j) {

api/src/main/java/marquez/db/JobDao.java

@@ -507,6 +507,90 @@ default RunRow upsertRunMeta(
   @SqlUpdate("UPDATE runs SET job_version_uuid = :jobVersionUuid WHERE uuid = :runUuid")
   void updateJobVersion(UUID runUuid, UUID jobVersionUuid);
 
+  @SqlQuery(
+      """
+          WITH filtered_jobs AS (
+            SELECT
+                jv.uuid,
+                jv.namespace_name,
+                jv.name
+            FROM jobs_view jv
+            WHERE jv.namespace_name=:namespace AND (jv.name=:jobName OR :jobName = ANY(jv.aliases))
+          ),
+          run_facets_agg AS (
+            SELECT
+                run_uuid,
+                JSON_AGG(facet ORDER BY lineage_event_time ASC) AS facets
+            FROM run_facets_view
+            -- This filter here is used for performance purpose: we only aggregate the json of run_uuid that matters
+            WHERE
+                run_uuid IN (SELECT uuid FROM runs_view WHERE job_uuid IN (SELECT uuid FROM filtered_jobs))
+            GROUP BY run_uuid
+          ),
+          input_versions_agg AS (
+               SELECT
+                   im.run_uuid,
+                   JSON_AGG(json_build_object('namespace', dv.namespace_name,
+                        'name', dv.dataset_name,
+                        'version', dv.version,
+                        'dataset_version_uuid', dv.uuid
+                   )) AS input_versions
+               FROM runs_input_mapping im
+               INNER JOIN dataset_versions dv ON im.dataset_version_uuid = dv.uuid
+               -- This filter here is used for performance purpose: we only aggregate the json of run_uuid that matters
+               WHERE
+                   im.run_uuid IN (SELECT uuid FROM runs_view WHERE job_uuid IN (SELECT uuid FROM filtered_jobs))
+               GROUP BY im.run_uuid
+          ),
+          output_versions_agg AS (
+              SELECT
+                  dv.run_uuid,
+              JSON_AGG(json_build_object('namespace', namespace_name,
+                                       'name', dataset_name,
+                                       'version', version,
+                                       'dataset_version_uuid', uuid
+                                       )) AS output_versions
+              FROM dataset_versions dv
+              -- This filter here is used for performance purpose: we only aggregate the json of run_uuid that matters
+              WHERE dv.run_uuid IN (SELECT uuid FROM runs_view WHERE job_uuid IN (SELECT uuid FROM filtered_jobs))
+              GROUP BY dv.run_uuid
+          ),
+          dataset_facets_agg AS (
+              SELECT
+                  run_uuid,
+                  JSON_AGG(json_build_object(
+                      'dataset_version_uuid', dataset_version_uuid,
+                      'name', name,
+                      'type', type,
+                      'facet', facet
+                  ) ORDER BY created_at ASC) as dataset_facets
+              FROM dataset_facets_view
+              -- This filter here is used for performance purpose: we only aggregate the json of run_uuid that matters
+              WHERE run_uuid IN (SELECT uuid FROM runs_view WHERE job_uuid IN (SELECT uuid FROM filtered_jobs))
+              AND (type ILIKE 'output' OR type ILIKE 'input')
+              GROUP BY run_uuid
+          )
+          SELECT
+              r.*,
+              ra.args,
+              f.facets,
+              jv.version AS job_version,
+              ri.input_versions,
+              ro.output_versions,
+              df.dataset_facets
+          FROM runs_view r
+          INNER JOIN filtered_jobs fj ON r.job_uuid = fj.uuid
+          LEFT JOIN run_facets_agg f ON r.uuid = f.run_uuid
+          LEFT JOIN run_args ra ON ra.uuid = r.run_args_uuid
+          LEFT JOIN job_versions jv ON jv.uuid = r.job_version_uuid
+          LEFT JOIN input_versions_agg ri ON r.uuid = ri.run_uuid
+          LEFT JOIN output_versions_agg ro ON r.uuid = ro.run_uuid
+          LEFT JOIN dataset_facets_agg df ON r.uuid = df.run_uuid
+          ORDER BY r.transitioned_at DESC, r.started_at DESC
+          LIMIT :limit OFFSET :offset
+      """)
+  List<Run> findByLatestJobOptimized(String namespace, String jobName, int limit, int offset);
+
   @SqlQuery(
       BASE_FIND_RUN_SQL
           + """

api/src/main/java/marquez/db/RunDao.java

@@ -43,7 +43,7 @@ static Set<String> getColumnNames(ResultSetMetaData metaData) {
 
   /**
    * Returns a new {@link ImmutableMap} instance of facets present in the provided {@link
-   * java.sql.ResultSet}, or an empty {@link ImmutableMap} if none are present. Note, {@code key}s
+   * ResultSet}, or an empty {@link ImmutableMap} if none are present. Note, {@code key}s
    * in the resulting facet map are the facet names (ex: 'schema', 'dataSource', 'documentation',
    * etc).
    */

api/src/main/java/marquez/db/mappers/MapperUtils.java

@@ -9,7 +9,7 @@
 import io.prometheus.client.Histogram;
 
 public class DatabaseMetrics {
-  public static final CollectorRegistry registry = new io.prometheus.client.CollectorRegistry();
+  public static final CollectorRegistry registry = new CollectorRegistry();
 
   public static final Histogram dbDurationSeconds =
       Histogram.build()

api/src/main/java/marquez/service/DatabaseMetrics.java

@@ -38,7 +38,7 @@ public class ColumnLineageIntegrationTest extends BaseIntegrationTest {
   public void setup(Jdbi jdbi) {
     OpenLineageDao openLineageDao = jdbi.onDemand(OpenLineageDao.class);
 
-    LineageEvent.JobFacet jobFacet = JobFacet.builder().build();
+    JobFacet jobFacet = JobFacet.builder().build();
 
     LineageEvent.Dataset dataset_A = getDatasetA();
     LineageEvent.Dataset dataset_B = getDatasetB();