Skip to content

[SPARK-48047][SQL] Reduce memory pressure of empty TreeNode tags #46285

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
n-young-db wants to merge 6 commits into from
Closed

[SPARK-48047][SQL] Reduce memory pressure of empty TreeNode tags #46285

Changes from 5 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
ad49a31
Reduce memory pressure of empty TreeNode tags
n-young-db Apr 29, 2024
70dfc02
comments
n-young-db Apr 29, 2024
dbf703d
comment
n-young-db Apr 29, 2024
866c028
comment
n-young-db Apr 29, 2024
6303972
comment
n-young-db Apr 29, 2024
09c93e4
isEmpty
n-young-db Apr 29, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 18 additions & 4 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,8 +78,16 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]]
/**
* A mutable map for holding auxiliary information of this tree node. It will be carried over
* when this node is copied via `makeCopy`, or transformed via `transformUp`/`transformDown`.
* We lazily evaluate the `tags` since the default size of a `mutable.Map` is nonzero. This
* will reduce unnecessary memory pressure.
*/
private val tags: mutable.Map[TreeNodeTag[_], Any] = mutable.Map.empty
private[this] var _tags: mutable.Map[TreeNodeTag[_], Any] = null
private def tags: mutable.Map[TreeNodeTag[_], Any] = {
if (_tags eq null) {
_tags = mutable.Map.empty
}
_tags
}

/**
* Default tree pattern [[BitSet] for a [[TreeNode]].
Expand Down Expand Up @@ -151,7 +159,7 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]]
// SPARK-32753: it only makes sense to copy tags to a new node
// but it's too expensive to detect other cases likes node removal
// so we make a compromise here to copy tags to node with no tags
if (tags.isEmpty) {
if ((_tags eq null) || _tags.isEmpty) {
tags ++= other.tags
Copy link
Contributor

@andylam-db andylam-db May 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since tags are def, doesn't this line mutate nothing?

Copy link
Contributor

@andylam-db andylam-db May 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, the map is probably pass-by reference.

}
}
Expand All @@ -161,11 +169,17 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]]
}

def getTagValue[T](tag: TreeNodeTag[T]): Option[T] = {
tags.get(tag).map(_.asInstanceOf[T])
if (_tags eq null) {
Copy link
Contributor

@nikhilsheoran-db nikhilsheoran-db Apr 29, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it makes sense to add a comment clarifying the intent of using both _tags and tags here -- "To avoid initializing the tags when getTagValue is called on a TreeNode without any tags set"

Copy link
Contributor Author

@n-young-db n-young-db Apr 29, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hm, I think the intent is clear enough as is

None
} else {
tags.get(tag).map(_.asInstanceOf[T])
}
}

def unsetTagValue[T](tag: TreeNodeTag[T]): Unit = {
tags -= tag
if (_tags ne null) {
tags -= tag
}
}

/**
Expand Down