Implement pluggable Lineage in Java SDK

[shnapz]

Addresses #36790: "[Feature Request]: Make lineage tracking pluggable"

Changes

• Created org.apache.beam.sdk.lineage.LineageBase - Plugin interface with single add() method
• Refactored org.apache.beam.sdk.metrics.Lineage - Hardcoded metrics → Delegation to LineageBase plugins
• Created org.apache.beam.sdk.lineage.LineageRegistrar - Plugin discovery interface (ServiceLoader)
• Extracted org.apache.beam.sdk.metrics.MetricsLineage - Default metrics-based implementation (implements LineageBase)
• Added Plugin initialization in FileSystems.setDefaultPipelineOptions()

Architecture

Before (master): Lineage was a concrete class hardcoded to use Beam metrics:

public class Lineage {
  private static final Lineage SOURCES = new Lineage(Type.SOURCE);
  private static final Lineage SINKS = new Lineage(Type.SINK);
  private final Metric metric;  // Hardcoded to Beam metrics

  private Lineage(Type type) {
    this.metric = Metrics.stringSet(LINEAGE_NAMESPACE, type.toString());
  }

  public void add(Iterable<String> segments) {
    ((StringSet) metric).add(String.join("", segments));  // Always metrics
  }
}

After (this PR): Clean separation via composition pattern:

// Plugin contract (simple interface)
public interface LineageBase {
  void add(Iterable<String> rollupSegments);
}

// Public API (final facade delegating to plugin)
public final class Lineage {
  private final LineageBase delegate;  // Plugin implementation

  private Lineage(LineageBase delegate) {
    this.delegate = delegate;
  }

  // Delegates to plugin
  public void add(Iterable<String> segments) {
    delegate.add(segments);
  }

  // Convenience overloads
  public void add(String system, Iterable<String> segments) { ... }
  public void add(String system, String subtype, ...) { ... }

  // Static utilities (unchanged)
  public static Lineage getSources() { ... }
  public static Lineage getSinks() { ... }
  public static String wrapSegment(String value) { ... }
  public static Set<String> query(MetricResults results, Type type) { ... }
}

// Default implementation (backward compatible)
public class MetricsLineage implements LineageBase {
  private final Metric metric;

  @Override
  public void add(Iterable<String> segments) {
    ((BoundedTrie) metric).add(segments);
  }
}

Plugin Selection: ServiceLoader discovery, first match wins, fallback to MetricsLineage.

Backward Compatibility: ✅ All existing code works unchanged (24+ call sites, static utilities, enums).

Why Pluggable Lineage?

1. Runner Fragmentation

Metrics-based lineage is scattered across runners with inconsistent support:

• Dataflow: Real-time metrics export to Cloud Monitoring
• Flink: Batch-only aggregation (no streaming support yet)
• Spark/Direct: Varying levels of support

Impact: Multi-runner organizations must consolidate lineage from different metrics backends, each with different APIs and formats.

Plugin Solution: Single implementation works consistently across all runners.

2. Enterprise Integration

Organizations with existing lineage infrastructure need:

• Direct API integration (Atlan, Collibra, Marquez, DataHub, OpenLineage)
• Custom metadata enrichment not in metrics subsystem

Example: Flyte workflow executing a Beam pipeline needs to tag lineage with Flyte execution ID and cost allocation. This context exists in the orchestrator, not in Beam workers' metrics.

3. Standard Formats

OpenLineage is the industry standard. Plugin enables direct emission vs. export metrics → parse → transform → send.

Initialization

Lineage.setDefaultPipelineOptions(options) is called from FileSystems.setDefaultPipelineOptions() (same pattern as Metrics).

Rationale: FileSystems.setDefaultPipelineOptions() is called at 48+ locations covering all execution scenarios (pipeline construction, worker startup, deserialization).

Known Limitation: Follows existing FileSystems pattern despite known issues (#18430). Architectural improvements would address all subsystems together.

Thread Safety

Uses AtomicReference with compareAndSet loop (same pattern as FileSystems/Metrics):

• AtomicReference<KV<Long, Integer>> tracks PipelineOptions identity
• AtomicReference<Lineage> for SOURCES/SINKS instances

Example: OpenLineage Plugin

For demonstration only (OpenLineage integration out of scope)

// 1. Plugin options
public interface OpenLineageOptions extends PipelineOptions {
  @Description("OpenLineage endpoint URL")
  String getOpenLineageUrl();
  void setOpenLineageUrl(String url);

  @Description("Enable OpenLineage plugin")
  @Default.Boolean(false)
  Boolean getEnableOpenLineage();
  void setEnableOpenLineage(Boolean enable);
}

// 2. Implement LineageBase
class OpenLineageReporter implements LineageBase {
  private final String endpoint;
  private final Lineage.LineageDirection direction;

  @Override
  public void add(Iterable<String> rollupSegments) {
    String fqn = String.join("", rollupSegments);
    // POST to OpenLineage API with workflow context
    sendToOpenLineage(endpoint, direction, fqn);
  }
}

// 3. Register via ServiceLoader
@AutoService(LineageRegistrar.class)
public class OpenLineageRegistrar implements LineageRegistrar {
  @Override
  public LineageBase fromOptions(PipelineOptions options, Lineage.LineageDirection direction) {
    OpenLineageOptions opts = options.as(OpenLineageOptions.class);
    if (opts.getEnableOpenLineage()) {
      return new OpenLineageReporter(opts.getOpenLineageUrl(), direction);
    }
    return null; // Fall back to MetricsLineage
  }
}

// 4. Usage
PipelineOptions options = PipelineOptionsFactory.create();
options.as(OpenLineageOptions.class).setEnableOpenLineage(true);
options.as(OpenLineageOptions.class).setOpenLineageUrl("https://lineage-api.example.com");
Pipeline p = Pipeline.create(options);

[shnapz]

irrelevant change to fix flaky tests

[shnapz]

irrelevant change to fix flaky tests

[github-actions]

Assigning reviewers:

R: @kennknowles for label java.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 56.86%. Comparing base (298f309) to head (7bcab65).
⚠️ Report is 12 commits behind head on master.

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #36781      +/-   ##
============================================
- Coverage     56.87%   56.86%   -0.01%     
+ Complexity     3417     3414       -3     
============================================
  Files          1178     1178              
  Lines        187492   187495       +3     
  Branches       3590     3590              
============================================
- Hits         106628   106620       -8     
- Misses        77472    77480       +8     
- Partials       3392     3395       +3     

Flag	Coverage Δ
java	71.92% <ø> (-0.04%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[github-actions]

Reminder, please take a look at this pr: @kennknowles

[shnapz]

[PreCommit Java / beam_PreCommit_Java](https://github.com/apache/beam/actions/runs/20927101528/job/60127906908?pr=36781) has apparently unrelated grpc client failure. Before I rebased on latest master it was green

[github-actions]

Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment assign to next reviewer:

R: @Abacn for label java.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

[shnapz]

@Abacn or @kennknowles would you mind taking a look? The change is low risk and introduces some important flexibility

[github-actions]

Reminder, please take a look at this pr: @Abacn

[github-actions]

Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment assign to next reviewer:

R: @kennknowles for label java.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

[github-actions]

Reminder, please take a look at this pr: @kennknowles

[shnapz]

FYI
We are using Lineage as a facade around LineageBase, so we don't expose the latter and we avoid any (breaking) changes in other parts of Beam

[shnapz]

@kennknowles Thanks for the feedback! I've updated the PR description with a detailed design section addressing your questions.

Selection Model: Single active plugin (first match wins via ServiceLoader), fallback to MetricsLineage. This matches Beam's existing patterns (similar to PipelineOptions registrars).

I've added a "Why Pluggable Lineage?" section to the PR description with four key use cases:

• Runner Fragmentation. Multi-runner orgs need unified lineage.
• Enterprise Integration. Direct integration with existing lineage systems + workflow-level context unavailable to worker metrics.
• Standard Formats. Direct OpenLineage emission vs. export -> parse -> transform -> send.

let me know if you have any questions!

[kennknowles]

R: @rohitsinha54

[github-actions]

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

[shnapz]

@rohitsinha54 could you take a look please?

[rohitsinha54]

Sorry for the delay in response. I have not had a chance to look into this even though I want to due to other priorities and commitments. I do really appreciate the contribution and believe it is will be very helpful to many Beam users. To not have this blocked on me any further I requested @Abacn to help with review. Thank you @Abacn.

[shnapz]

@Abacn could you please take a look?

[Abacn]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a new plugin mechanism for lineage implementations in Apache Beam. It refactors the existing Lineage class to act as a facade, delegating to LineageBase implementations discovered via LineageRegistrar. A default MetricsLineage implementation is provided, and the FileSystems class is updated for proper initialization. Comprehensive tests for the new plugin system are also included. A review comment suggests refactoring the constructor logic in MetricsLineage.java for improved readability.

[gemini-code-assist]

The constructor logic can be simplified for better readability. The type variable is only used in the else block but is calculated unconditionally, and the if block re-evaluates direction. A clearer approach is to determine the metric type within each branch of the if statement.

  public MetricsLineage(Lineage.LineageDirection direction) {
    if (MetricsFlag.lineageRollupEnabled()) {
      Lineage.Type metricType =
          (direction == Lineage.LineageDirection.SOURCE)
              ? Lineage.Type.SOURCEV2
              : Lineage.Type.SINKV2;
      this.metric = Metrics.boundedTrie(Lineage.LINEAGE_NAMESPACE, metricType.toString());
    } else {
      Lineage.Type metricType =
          (direction == Lineage.LineageDirection.SOURCE) ? Lineage.Type.SOURCE : Lineage.Type.SINK;
      this.metric = Metrics.stringSet(Lineage.LINEAGE_NAMESPACE, metricType.toString());
    }
  }

[Abacn]

There was an assumption (and coupling) between lineageRollupEnabled<->BoundedTrie. With common interface we can now avoid unchecked cast by implementing CounterMetricsLineage and TrieMetricsLineage separately.

[Abacn]

This work around is needed because AutoService gets loaded whenever it's in class path. I'm thinking about whether we should use approach similar to --runner selection instead of FileSystemRegistrar. Using a AutoService backed registrar to discover supported filesystems make sense as they won't conflict with each other. But pluggable lineage does---when there is registrar found, the default path (metrics based lineage) is disabled.

[Abacn]

If multiple registrars in classpath, the activated one is the alphabetically smallest one (for Tree Set). This could be confusing. Consider using a pipeline option to select a Lineage implementation.