This is an automated email from the ASF dual-hosted git repository.
taiyang-li pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 630d5ace8a [CORE] optimize driver heap memory by lazily store metadata
columns and fix mismatch result when input_file_name is only in filter
condition (#11899)
630d5ace8a is described below
commit 630d5ace8af5287e49c0942835450198babc121c
Author: 李扬 <[email protected]>
AuthorDate: Wed Jul 8 20:56:25 2026 +0800
[CORE] optimize driver heap memory by lazily store metadata columns and fix
mismatch result when input_file_name is only in filter condition (#11899)
* avoid driver oom caused by unnecessary metadata columns in splitinfo
* support input_file_name only in filter
* fix code style
---
.../gluten/backendsapi/clickhouse/CHIteratorApi.scala | 15 +++++++++++----
.../gluten/backendsapi/velox/VeloxIteratorApi.scala | 18 ++++++++++++++----
.../gluten/execution/BasicScanExecTransformer.scala | 14 +++++++++++++-
.../columnar/PushDownInputFileExpression.scala | 8 +++++++-
.../apache/spark/sql/GlutenColumnExpressionSuite.scala | 15 +++++++++++++++
.../apache/spark/sql/GlutenColumnExpressionSuite.scala | 15 +++++++++++++++
.../apache/spark/sql/GlutenColumnExpressionSuite.scala | 15 +++++++++++++++
.../apache/spark/sql/GlutenColumnExpressionSuite.scala | 15 +++++++++++++++
.../scala/org/apache/gluten/sql/shims/SparkShims.scala | 5 +++--
9 files changed, 108 insertions(+), 12 deletions(-)
diff --git
a/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHIteratorApi.scala
b/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHIteratorApi.scala
index 8e20a72a65..22e7a75aed 100644
---
a/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHIteratorApi.scala
+++
b/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHIteratorApi.scala
@@ -161,7 +161,10 @@ class CHIteratorApi extends IteratorApi with Logging with
LogLevelUtil {
val fileSizes = new JArrayList[JLong]()
val modificationTimes = new JArrayList[JLong]()
val partitionColumns = new JArrayList[JMap[String, String]]
- val metadataColumns = new JArrayList[JMap[String, String]]
+ val needMetadataColumns = metadataColumnNames != null &&
metadataColumnNames.nonEmpty
+ val emptyMetadataColumn: JMap[String, String] =
+ java.util.Collections.emptyMap[String, String]()
+ val metadataColumns = new JArrayList[JMap[String,
String]](f.files.length)
val otherMetadataColumns = new JArrayList[JMap[String, Object]]
val dateFormatter = DateFormatter()
val timestampFormatter =
TimestampFormatter.getFractionFormatter(ZoneOffset.UTC)
@@ -171,9 +174,13 @@ class CHIteratorApi extends IteratorApi with Logging with
LogLevelUtil {
starts.add(JLong.valueOf(file.start))
lengths.add(JLong.valueOf(file.length))
val metadataColumn =
- SparkShimLoader.getSparkShims
- .generateMetadataColumns(file, metadataColumnNames)
- .asJava
+ if (needMetadataColumns) {
+ SparkShimLoader.getSparkShims
+ .generateMetadataColumns(file, metadataColumnNames)
+ .asJava
+ } else {
+ emptyMetadataColumn
+ }
metadataColumns.add(metadataColumn)
val partitionColumn = new JHashMap[String, String]()
for (i <- 0 until file.partitionValues.numFields) {
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxIteratorApi.scala
b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxIteratorApi.scala
index 9ad1e8dc8e..f81b78c71f 100644
---
a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxIteratorApi.scala
+++
b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxIteratorApi.scala
@@ -91,9 +91,19 @@ class VeloxIteratorApi extends IteratorApi with Logging {
.unzip
val partitionColumns = getPartitionColumns(partitionSchema, partitionFiles)
- val metadataColumns = partitionFiles
- .map(
- f => SparkShimLoader.getSparkShims.generateMetadataColumns(f,
metadataColumnNames).asJava)
+ val needMetadataColumns = metadataColumnNames != null &&
metadataColumnNames.nonEmpty
+ val emptyMetadataColumn: java.util.Map[String, String] =
+ java.util.Collections.emptyMap[String, String]()
+ val metadataColumns: java.util.List[java.util.Map[String, String]] =
+ if (needMetadataColumns) {
+ partitionFiles
+ .map(
+ f =>
+ SparkShimLoader.getSparkShims.generateMetadataColumns(f,
metadataColumnNames).asJava)
+ .asJava
+ } else {
+ java.util.Collections.nCopies(partitionFiles.size, emptyMetadataColumn)
+ }
val otherMetadataColumns = partitionFiles
.map(f =>
SparkShimLoader.getSparkShims.getOtherConstantMetadataColumnValues(f))
@@ -106,7 +116,7 @@ class VeloxIteratorApi extends IteratorApi with Logging {
fileSizes.asJava,
modificationTimes.asJava,
partitionColumns.map(_.asJava).asJava,
- metadataColumns.asJava,
+ metadataColumns,
fileFormat,
locations.toList.asJava,
mapAsJavaMap(properties),
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/BasicScanExecTransformer.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/BasicScanExecTransformer.scala
index 6d265100d0..c765ed333d 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/BasicScanExecTransformer.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/BasicScanExecTransformer.scala
@@ -115,6 +115,18 @@ trait BasicScanExecTransformer extends
LeafTransformSupport with BaseDataSource
case _ => Seq(partition)
}
+ val metadataFromSpark = getMetadataColumns().map(_.name)
+
+ val inputFileRelatedMetadataKeys = Seq(
+ InputFileName().prettyName,
+ InputFileBlockStart().prettyName,
+ InputFileBlockLength().prettyName)
+
+ val neededInputFileRelatedMetadataKeys =
+ inputFileRelatedMetadataKeys.filter(k => output.exists(_.name == k))
+
+ val metadataColumnNames = (metadataFromSpark ++
neededInputFileRelatedMetadataKeys).distinct
+
BackendsApiManager.getIteratorApiInstance
.genSplitInfo(
partition.index,
@@ -122,7 +134,7 @@ trait BasicScanExecTransformer extends LeafTransformSupport
with BaseDataSource
getPartitionSchema,
getDataSchema,
readFileFormat,
- getMetadataColumns().map(_.name),
+ metadataColumnNames,
getProperties)
}
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/PushDownInputFileExpression.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/PushDownInputFileExpression.scala
index 8e3a2d53ad..d50c10e7b4 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/PushDownInputFileExpression.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/PushDownInputFileExpression.scala
@@ -21,7 +21,7 @@ import
org.apache.gluten.execution.{BatchScanExecTransformerBase, FileSourceScan
import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute,
AttributeReference, Expression, InputFileBlockLength, InputFileBlockStart,
InputFileName, NamedExpression}
import org.apache.spark.sql.catalyst.optimizer.CollapseProjectShim
import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.execution.{DeserializeToObjectExec,
FileSourceScanExec, LeafExecNode, ProjectExec, SerializeFromObjectExec,
SparkPlan, UnionExec}
+import org.apache.spark.sql.execution.{DeserializeToObjectExec,
FileSourceScanExec, FilterExec, LeafExecNode, ProjectExec,
SerializeFromObjectExec, SparkPlan, UnionExec}
import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
import org.apache.spark.sql.hive.HiveTableScanExecTransformer
@@ -96,6 +96,12 @@ object PushDownInputFileExpression {
}
val newChild = addMetadataCol(child, replacedExprs)
ProjectExec(newProjectList, newChild)
+ case f @ FilterExec(condition, child)
+ if containsInputFileRelatedExpr(condition) &&
hasInputFileRelatedSource(child) =>
+ val replacedExprs = mutable.Map[String, Alias]()
+ val newCondition = rewriteExpr(condition, replacedExprs)
+ val newChild = addMetadataCol(child, replacedExprs)
+ ProjectExec(f.output, FilterExec(newCondition, newChild))
}
private def addMetadataCol(
diff --git
a/gluten-ut/spark33/src/test/scala/org/apache/spark/sql/GlutenColumnExpressionSuite.scala
b/gluten-ut/spark33/src/test/scala/org/apache/spark/sql/GlutenColumnExpressionSuite.scala
index f4f85825c8..f525bab61a 100644
---
a/gluten-ut/spark33/src/test/scala/org/apache/spark/sql/GlutenColumnExpressionSuite.scala
+++
b/gluten-ut/spark33/src/test/scala/org/apache/spark/sql/GlutenColumnExpressionSuite.scala
@@ -97,4 +97,19 @@ class GlutenColumnExpressionSuite extends
ColumnExpressionSuite with GlutenSQLTe
}
}
}
+
+ test("scan filter references input_file_name but project does not select
it") {
+ withTempPath {
+ dir =>
+ val data = sparkContext.parallelize(0 to 10).toDF("id")
+ data.write.parquet(dir.getCanonicalPath)
+
+ // Filter references input_file_name(), but the project only selects
`id`.
+ val q = spark.read
+ .parquet(dir.getCanonicalPath)
+ .filter(input_file_name().contains("parquet"))
+ .select($"id")
+ checkAnswer(q, (0 to 10).map(Row(_)))
+ }
+ }
}
diff --git
a/gluten-ut/spark35/src/test/scala/org/apache/spark/sql/GlutenColumnExpressionSuite.scala
b/gluten-ut/spark35/src/test/scala/org/apache/spark/sql/GlutenColumnExpressionSuite.scala
index f4f85825c8..f525bab61a 100644
---
a/gluten-ut/spark35/src/test/scala/org/apache/spark/sql/GlutenColumnExpressionSuite.scala
+++
b/gluten-ut/spark35/src/test/scala/org/apache/spark/sql/GlutenColumnExpressionSuite.scala
@@ -97,4 +97,19 @@ class GlutenColumnExpressionSuite extends
ColumnExpressionSuite with GlutenSQLTe
}
}
}
+
+ test("scan filter references input_file_name but project does not select
it") {
+ withTempPath {
+ dir =>
+ val data = sparkContext.parallelize(0 to 10).toDF("id")
+ data.write.parquet(dir.getCanonicalPath)
+
+ // Filter references input_file_name(), but the project only selects
`id`.
+ val q = spark.read
+ .parquet(dir.getCanonicalPath)
+ .filter(input_file_name().contains("parquet"))
+ .select($"id")
+ checkAnswer(q, (0 to 10).map(Row(_)))
+ }
+ }
}
diff --git
a/gluten-ut/spark40/src/test/scala/org/apache/spark/sql/GlutenColumnExpressionSuite.scala
b/gluten-ut/spark40/src/test/scala/org/apache/spark/sql/GlutenColumnExpressionSuite.scala
index 6ee95026fd..ad43946768 100644
---
a/gluten-ut/spark40/src/test/scala/org/apache/spark/sql/GlutenColumnExpressionSuite.scala
+++
b/gluten-ut/spark40/src/test/scala/org/apache/spark/sql/GlutenColumnExpressionSuite.scala
@@ -100,4 +100,19 @@ class GlutenColumnExpressionSuite extends
ColumnExpressionSuite with GlutenSQLTe
}
}
}
+
+ test("scan filter references input_file_name but project does not select
it") {
+ withTempPath {
+ dir =>
+ val data = sparkContext.parallelize(0 to 10).toDF("id")
+ data.write.parquet(dir.getCanonicalPath)
+
+ // Filter references input_file_name(), but the project only selects
`id`.
+ val q = spark.read
+ .parquet(dir.getCanonicalPath)
+ .filter(input_file_name().contains("parquet"))
+ .select($"id")
+ checkAnswer(q, (0 to 10).map(Row(_)))
+ }
+ }
}
diff --git
a/gluten-ut/spark41/src/test/scala/org/apache/spark/sql/GlutenColumnExpressionSuite.scala
b/gluten-ut/spark41/src/test/scala/org/apache/spark/sql/GlutenColumnExpressionSuite.scala
index 6ee95026fd..ad43946768 100644
---
a/gluten-ut/spark41/src/test/scala/org/apache/spark/sql/GlutenColumnExpressionSuite.scala
+++
b/gluten-ut/spark41/src/test/scala/org/apache/spark/sql/GlutenColumnExpressionSuite.scala
@@ -100,4 +100,19 @@ class GlutenColumnExpressionSuite extends
ColumnExpressionSuite with GlutenSQLTe
}
}
}
+
+ test("scan filter references input_file_name but project does not select
it") {
+ withTempPath {
+ dir =>
+ val data = sparkContext.parallelize(0 to 10).toDF("id")
+ data.write.parquet(dir.getCanonicalPath)
+
+ // Filter references input_file_name(), but the project only selects
`id`.
+ val q = spark.read
+ .parquet(dir.getCanonicalPath)
+ .filter(input_file_name().contains("parquet"))
+ .select($"id")
+ checkAnswer(q, (0 to 10).map(Row(_)))
+ }
+ }
}
diff --git
a/shims/common/src/main/scala/org/apache/gluten/sql/shims/SparkShims.scala
b/shims/common/src/main/scala/org/apache/gluten/sql/shims/SparkShims.scala
index c7b8f89c87..166d7fcce5 100644
--- a/shims/common/src/main/scala/org/apache/gluten/sql/shims/SparkShims.scala
+++ b/shims/common/src/main/scala/org/apache/gluten/sql/shims/SparkShims.scala
@@ -171,11 +171,12 @@ trait SparkShims {
def generateMetadataColumns(
file: PartitionedFile,
metadataColumnNames: Seq[String] = Seq.empty): Map[String, String] = {
- Map(
+ val requested = metadataColumnNames.toSet
+ Seq(
InputFileName().prettyName -> file.filePath.toString,
InputFileBlockStart().prettyName -> file.start.toString,
InputFileBlockLength().prettyName -> file.length.toString
- )
+ ).collect { case (name, value) if requested.contains(name) => name ->
value }.toMap
}
// For compatibility with Spark-3.5.
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
