fix(ingestion/nifi): Improve nifi lineage extraction performance #11490

skrydal · 2024-09-27T11:38:19Z

Checklist

The PR conforms to DataHub's Contributing Guideline (particularly Commit Message Format)
Links to related issues (if applicable)
Tests for the changes have been added/updated (if applicable)
Docs related to the changes have been added/updated (if applicable). If a new feature has been added a Usage Guide has been added for the same.
For any breaking change/potential downtime/deprecation/big changes an entry has been made in Updating DataHub

coderabbitai · 2024-09-27T11:38:26Z

Important

Review skipped

Auto reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media?

❤️ Share

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

Review comments: Directly reply to a review comment made by CodeRabbit. Example:
- I pushed a fix in commit <commit_id>, please review it.
- Generate unit testing code for this file.
- Open a follow-up GitHub issue for this discussion.
Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
- @coderabbitai generate unit testing code for this file.
- @coderabbitai modularize this function.
PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
- @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
- @coderabbitai read src/utils.ts and generate unit testing code.
- @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
- @coderabbitai help me debug CodeRabbit configuration file.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

@coderabbitai pause to pause the reviews on a PR.
@coderabbitai resume to resume the paused reviews.
@coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
@coderabbitai full review to do a full review from scratch and review all the files again.
@coderabbitai summary to regenerate the summary of the PR.
@coderabbitai resolve resolve all the CodeRabbit review comments.
@coderabbitai configuration to show the current CodeRabbit configuration for the repository.
@coderabbitai help to get help.

Other keywords and placeholders

Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (`.coderabbit.yaml`)

You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
Please see the configuration documentation for more information.
If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

Visit our Documentation for detailed information on how to use CodeRabbit.
Join our Discord Community to get help, request features, and share feedback.
Follow us on X/Twitter for updates and announcements.

skrydal · 2024-09-30T13:23:39Z

@hsheth2 please take a look again - I changed the PR completely, now it is improving the performance of the lineage extraction. It is ready to be merged if I get approvals.

mayurinehate · 2024-10-01T05:43:45Z

metadata-ingestion/src/datahub/ingestion/source/nifi.py

@@ -694,8 +781,11 @@ def fetch_provenance_events(

        if provenance_response.ok:
            provenance = provenance_response.json().get("provenance", {})
+            total = provenance.get("results", {}).get("total")


Would these counts be available in provenance results before verifying that fetching provenance results has finished - which is done below in while loop ?

That's a great catch - while it doesn't cause any exception it actually always returns 0 in such case for total and totalCount. Moved the retrieval of these values to back where they were.

mayurinehate · 2024-10-01T05:48:11Z

metadata-ingestion/src/datahub/ingestion/source/nifi.py

-    connections: List[Tuple[str, str]] = field(default_factory=list)
+    connections: BidirectionalComponentGraph = field(
+        default_factory=BidirectionalComponentGraph
+    )


Assuming, perf improvement is due to removing duplicate connection edges,
I wonder if changing connections: List[Tuple[str, str]] to connections: Set[Tuple[str, str]] and keeping everything else just as before give similar boost in perf ?

It would address duplicates problem but then we would need to still iterate over whole set to find matches for incoming/outgoing connections for an entity. The question is why would we do it this way then if we already have better implementation in place? Do you see any shortcoming of the current implementation, beside taking roughly twice the amount of memory (which I think is not much in first place, amount of connections, even for the biggest clusters, tends to be below 100k).

The implementation of BidirectionalComponentGraph is fair and is adapted as per operations we need. It looks good to me.

mayurinehate · 2024-10-01T09:25:15Z

metadata-ingestion/src/datahub/ingestion/source/nifi.py

+        self.outgoing: Dict[str, Set[str]] = defaultdict(set)
+        self.incoming: Dict[str, Set[str]] = defaultdict(set)
+        # this will not count duplicates/removal of non-existing connections correctly - it is only for quick check
+        self.connections_cnt = 0


Can we rename instance variables incoming, outgoing and connections_cnt to start with an underscore to clarify that they are not to be modified/accessed outside of class ?

Renamed the variables, as requested. Do we follow this pattern in other relevant classes as well? I haven't noticed many fields starting with a _.

mayurinehate · 2024-10-01T09:25:44Z

metadata-ingestion/src/datahub/ingestion/source/nifi.py

+        del self.outgoing[component]
+        del self.incoming[component]
+
+        self.connections_cnt -= deleted_connections_cnt


Suggested change

self.connections_cnt -= deleted_connections_cnt

self.connections_cnt += added_connections_cnt - deleted_connections_cnt

This would be wrong - the add_connection function already increments the counter.

Oh, can you add a comment mentioning so , or use remove_connection?

The count thing isn't so clear to me.

I added an explanatory comment, remove_connection shouldn't be used while removing a component (due to performance reasons)

mayurinehate · 2024-10-01T09:26:11Z

metadata-ingestion/src/datahub/ingestion/source/nifi.py

-    connections: List[Tuple[str, str]] = field(default_factory=list)
+    connections: BidirectionalComponentGraph = field(
+        default_factory=BidirectionalComponentGraph
+    )


The implementation of BidirectionalComponentGraph is fair and is adapted as per operations we need. It looks good to me.

skrydal added 2 commits September 27, 2024 13:34

Add verbosity to the nifi source

b8bc84b

Linting

6d102d1

skrydal changed the title ~~fix(ingestion/nifi) Increase nifi verbosity~~ fix(ingestion/nifi): Increase nifi verbosity Sep 27, 2024

github-actions bot added the ingestion PR or Issue related to the ingestion of metadata label Sep 27, 2024

vercel bot deployed to Preview September 27, 2024 11:56 View deployment

skrydal and others added 2 commits September 27, 2024 14:57

Merge branch 'datahub-project:master' into nifi_problem

042dd10

Linting

799b4e4

vercel bot deployed to Preview September 27, 2024 13:14 View deployment

Even more verbosity

1731b76

vercel bot deployed to Preview September 27, 2024 14:11 View deployment

skrydal added 2 commits September 27, 2024 17:10

print -> debug

d6cc882

Fixed outgoing connections logging

d951e02

vercel bot deployed to Preview September 27, 2024 15:31 View deployment

Fixing log

5c97171

hsheth2 approved these changes Sep 27, 2024

View reviewed changes

hsheth2 added the merge-pending-ci A PR that has passed review and should be merged once CI is green. label Sep 27, 2024

vercel bot deployed to Preview September 27, 2024 16:52 View deployment

skrydal removed the merge-pending-ci A PR that has passed review and should be merged once CI is green. label Sep 27, 2024

skrydal changed the title ~~fix(ingestion/nifi): Increase nifi verbosity~~ [DO NOT MERGE] fix(ingestion/nifi): Increase nifi verbosity Sep 27, 2024

skrydal added 2 commits September 28, 2024 23:42

More performant approach

b51ce00

Improved logging, fixed unit tests

d5a8b64

vercel bot deployed to Preview September 28, 2024 22:16 View deployment

skrydal changed the title ~~[DO NOT MERGE] fix(ingestion/nifi): Increase nifi verbosity~~ [DO NOT MERGE] fix(ingestion/nifi): Increase nifi lineage extraction performance Sep 29, 2024

Better connection counting

c567f7f

vercel bot deployed to Preview September 29, 2024 11:19 View deployment

Merge branch 'master' into nifi_problem

00577f2

vercel bot deployed to Preview September 29, 2024 18:58 View deployment

skrydal changed the title ~~[DO NOT MERGE] fix(ingestion/nifi): Increase nifi lineage extraction performance~~ fix(ingestion/nifi): Increase nifi lineage extraction performance Sep 30, 2024

Added some minor comments

b8d5e59

vercel bot deployed to Preview September 30, 2024 13:22 View deployment

skrydal requested a review from hsheth2 September 30, 2024 13:22

Merge branch 'master' into nifi_problem

23abb0a

vercel bot deployed to Preview September 30, 2024 13:48 View deployment

Merge branch 'master' into nifi_problem

d0e4b47

vercel bot deployed to Preview September 30, 2024 20:16 View deployment

mayurinehate reviewed Oct 1, 2024

View reviewed changes

skrydal changed the title ~~fix(ingestion/nifi): Increase nifi lineage extraction performance~~ fix(ingestion/nifi): Improve nifi lineage extraction performance Oct 1, 2024

Improvement of total counts logging

2c6a2e4

skrydal requested a review from mayurinehate October 1, 2024 08:34

vercel bot deployed to Preview October 1, 2024 08:51 View deployment

mayurinehate reviewed Oct 1, 2024

View reviewed changes

Adjusting fields names according to the PR comment

4389b33

skrydal requested a review from mayurinehate October 1, 2024 11:00

vercel bot deployed to Preview October 1, 2024 11:15 View deployment

mayurinehate approved these changes Oct 1, 2024

View reviewed changes

Added an explanatory comment

4f36ec2

vercel bot deployed to Preview October 1, 2024 12:43 View deployment

skrydal merged commit e1514d5 into datahub-project:master Oct 1, 2024
68 checks passed

skrydal deleted the nifi_problem branch October 1, 2024 19:51

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix(ingestion/nifi): Improve nifi lineage extraction performance #11490

fix(ingestion/nifi): Improve nifi lineage extraction performance #11490

skrydal commented Sep 27, 2024

coderabbitai bot commented Sep 27, 2024 •

edited

Loading

Review skipped

Chat

CodeRabbit Commands (Invoked using PR comments)

Other keywords and placeholders

CodeRabbit Configuration File (`.coderabbit.yaml`)

Documentation and Community

skrydal commented Sep 30, 2024

mayurinehate Oct 1, 2024

skrydal Oct 1, 2024

mayurinehate Oct 1, 2024

skrydal Oct 1, 2024

mayurinehate Oct 1, 2024

mayurinehate Oct 1, 2024

skrydal Oct 1, 2024

mayurinehate Oct 1, 2024

skrydal Oct 1, 2024

mayurinehate Oct 1, 2024

mayurinehate Oct 1, 2024 •

edited

Loading

skrydal Oct 1, 2024

mayurinehate Oct 1, 2024

	self.connections_cnt -= deleted_connections_cnt
	self.connections_cnt += added_connections_cnt - deleted_connections_cnt

fix(ingestion/nifi): Improve nifi lineage extraction performance #11490

fix(ingestion/nifi): Improve nifi lineage extraction performance #11490

Conversation

skrydal commented Sep 27, 2024

Checklist

coderabbitai bot commented Sep 27, 2024 • edited Loading

Review skipped

Chat

CodeRabbit Commands (Invoked using PR comments)

Other keywords and placeholders

CodeRabbit Configuration File (.coderabbit.yaml)

Documentation and Community

skrydal commented Sep 30, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mayurinehate Oct 1, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

coderabbitai bot commented Sep 27, 2024 •

edited

Loading

CodeRabbit Configuration File (`.coderabbit.yaml`)

mayurinehate Oct 1, 2024 •

edited

Loading