fix(airflow): Grouping of in memory nodes

[ankatiyar]

Description

Development notes

Checklist

• [ ] Opened this PR as a 'Draft Pull Request' if it is work-in-progress
• [ ] Updated the documentation to reflect the code changes
• [ ] Added a description of this change in the relevant RELEASE.md file
• [ ] Added tests to cover my changes