sunchao commented on code in PR #5027: URL: https://github.com/apache/datafusion-comet/pull/5027#discussion_r4162550069
########## spark/src/test/scala/org/apache/comet/CometCodegenMemoryPressureSuite.scala: ########## @@ -0,0 +1,84 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet + +import org.apache.spark.SparkConf +import org.apache.spark.sql.{CometTestBase, Row} +import org.apache.spark.sql.comet.CometSortExec +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.functions.{col, regexp_replace} + +/** + * A codegen-dispatched expression feeding a native sort in a small off-heap pool. + * + * A JVM UDF's output is charged to the Spark task while the UDF holds it. Native operators + * reserve through a consumer whose `spill` returns 0, so they spill only when their own + * `try_grow` fails, and by then they have filled the pool. The next UDF allocation then asks for + * more than Spark has left. Refusing it would fail the task where the sort could have spilled, so + * the allocation is recorded instead. + * + * `greedy_unified` makes Spark's grant the only limit on the sort, so the sort fills the pool + * within a single task. DataFusion's default 10 MiB merge reservation leaves 6 MiB of the 16 MiB + * pool for buffered batches. The UDF allocates each output from the size of its input, several + * times what the sort reserves for the shrunken batch that comes out, so the pool runs short for + * the UDF (about 115,000 rows in) before it runs short for the sort (about 131,000 rows in). + */ +class CometCodegenMemoryPressureSuite + extends CometTestBase + with AdaptiveSparkPlanHelper + with CometCodegenAssertions { + + override protected def sparkConf: SparkConf = + super.sparkConf + .set("spark.memory.offHeap.size", "16m") + .set(CometConf.COMET_OFFHEAP_MEMORY_POOL_TYPE.key, "greedy_unified") + + test("a UDF feeding a native sort under memory pressure lets the sort spill") { + val numRows = 160000 + withTempPath { dir => + val path = dir.getCanonicalPath + // 100-character strings that the UDF shrinks to 20 characters. + spark + .range(numRows) Review Comment: [P2] Could you pass `numRows.toLong` to `range`? `numRows` is inferred as `Int`, while Spark’s `range` takes `Long`. The required Spark 3.5 strict-warnings build rejects this implicit widening, so the newly added suite fails test compilation and prevents the required CI gate from passing. An explicit conversion preserves the test’s behavior. Evidence: Exact-head CI job 110682965428 runs `./mvnw -B test-compile -Pspark-3.5 -Pstrict-warnings -DskipTests` and reports `CometCodegenMemoryPressureSuite.scala:59: implicit numeric widening`, followed by `scala-maven-plugin:4.9.6:testCompile` failure. Checkout logs identify a merge containing `1158716705151f3bda234f40730df273b3aa7009`. The exact checkout declares `val numRows = 160000` at line 54 and passes it without conversion at line 59. https://github.com/apache/datafusion-comet/actions/runs/36956997245/job/110682965428 -- 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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
