andygrove commented on code in PR #1485:
URL: https://github.com/apache/datafusion-comet/pull/1485#discussion_r1991982544
##########
spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocator.java:
##########
@@ -35,43 +35,47 @@
* store serialized rows. This class is simply an implementation of
`MemoryConsumer` that delegates
* memory allocation to the `TaskMemoryManager`. This requires that the
`TaskMemoryManager` is
* configured with `MemoryMode.OFF_HEAP`, i.e. it is using off-heap memory.
+ *
+ * <p>If the user does not enable off-heap memory then we want to use
+ * CometBoundedShuffleMemoryAllocator. The tests also need to default to using
this because off-heap
+ * is not enabled when running the Spark SQL tests. The
+ * COMET_COLUMNAR_SHUFFLE_BOUNDED_MEMORY_ALLOCATOR flag allows us to use
unified memory management
+ * in Comet tests (this does assume that off-heap is enabled).
*/
public final class CometShuffleMemoryAllocator extends
CometShuffleMemoryAllocatorTrait {
private static CometShuffleMemoryAllocatorTrait INSTANCE;
/**
* Returns the singleton instance of `CometShuffleMemoryAllocator`. This
method should be used
* instead of the constructor to ensure that only one instance of
`CometShuffleMemoryAllocator` is
- * created. For Spark tests, this returns `CometTestShuffleMemoryAllocator`
which is a test-only
- * allocator that should not be used in production.
+ * created. For Spark tests, this returns
`CometBoundedShuffleMemoryAllocator`.
*/
public static CometShuffleMemoryAllocatorTrait getInstance(
SparkConf conf, TaskMemoryManager taskMemoryManager, long pageSize) {
boolean isSparkTesting = Utils.isTesting();
- boolean useUnifiedMemAllocator =
- (boolean)
-
CometConf$.MODULE$.COMET_COLUMNAR_SHUFFLE_UNIFIED_MEMORY_ALLOCATOR_IN_TEST().get();
+ boolean useBoundedAllocator =
+ (boolean)
CometConf$.MODULE$.COMET_COLUMNAR_SHUFFLE_BOUNDED_MEMORY_ALLOCATOR().get();
- if (!useUnifiedMemAllocator) {
+ if (isSparkTesting || useBoundedAllocator) {
Review Comment:
Also, if we run in on-heap mode, this condition doesn't get triggered so we
end up using the unified memory allocator, which we don't want
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
