[KYUUBI #5377] Spark engine query result save to file

[lsm1]

Why are the changes needed?

close #5377

How was this patch tested?

• Add some test cases that check the changes thoroughly including negative and positive cases if possible

• Add screenshots for manual tests if appropriate

• Run test locally before make a pull request

Was this patch authored or co-authored using generative AI tooling?

NO

[codecov-commenter]

Codecov Report

Attention: 99 lines in your changes are missing coverage. Please review.

Comparison is base (4463cc8) 61.51% compared to head (9d1a18c) 61.34%.
Report is 3 commits behind head on master.

Files	Patch %	Lines
...ubi/engine/spark/operation/FetchOrcStatement.scala	0.00%	57 Missing ⚠️
...uubi/engine/spark/operation/ExecuteStatement.scala	24.00%	18 Missing and 1 partial ⚠️
...rg/apache/kyuubi/engine/spark/SparkSQLEngine.scala	18.18%	8 Missing and 1 partial ⚠️
...g/apache/spark/sql/kyuubi/SparkDatasetHelper.scala	0.00%	8 Missing ⚠️
.../engine/spark/session/SparkSQLSessionManager.scala	0.00%	4 Missing and 1 partial ⚠️
...in/scala/org/apache/kyuubi/config/KyuubiConf.scala	93.75%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master    #5591      +/-   ##
============================================
- Coverage     61.51%   61.34%   -0.18%     
  Complexity       23       23              
============================================
  Files           608      609       +1     
  Lines         36091    36252     +161     
  Branches       4952     4993      +41     
============================================
+ Hits          22201    22237      +36     
- Misses        11506    11616     +110     
- Partials       2384     2399      +15     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[wForget]

Use resultDF instead of result. Also, is toDF(colName: _*) necessary?

[lsm1]

If the result has duplicate columns, we can not write it to file, so we rename all col name to avoid this case

spark.sql("select 1 as a,2 as a").write("/filepath")

org.apache.spark.sql.AnalysisException: Found duplicate column(s) when inserting into hdfs

[pan3793]

@lsm1 let's add such information to the comments

[cfmcgrady]

there is another known issue as I mentioned in the issue comment

directly call df.write will introduce an extra shuffle for the outermost limit, and hurt performance

I think we should also add this known issue to the comment and create a new ticket to track this issue.

[wForget]

iters(idx).hasNext) has been called in the hasNext method

[wForget]

shouldSaveResultToHdfs is inferred based on the execution plan and may not be accurate. Should we change it to sparkSave || shouldSaveResultToHdfs

[pan3793]

I prefer to keep it as-is, we need a configuration to disable this feature globally

[wForget]

It might be simpler for us to make changes in the executeStatement method, like:

change

result = spark.sql(statement)

to

if (saveResultToPath) {
  spark.sql(statement).write.format(format).save(resultPath)
  result = spark.read.load(resultPath)
} else {
  result = spark.sql(statement)
}

WDYT? cc @pan3793 @cxzl25

[cxzl25]

It might be simpler for us to make changes in the executeStatement method, like:

result = spark.read.load(resultPath)

This may lose the ordering of the query data, e.g. order by xx limit 100

[wForget]

This may lose the ordering of the query data, e.g. order by xx limit 100

I did a simple test and the results were as expected. (Test Env: Kyuubi 1.8.0 + Spark 3.5.0)

create table wangzhen_test_20231115_t1(id bigint, name string) stored as parquet;
insert into wangzhen_test_20231115_t1 values (1, 'a');
insert into wangzhen_test_20231115_t1 values (2, 'b');
insert into wangzhen_test_20231115_t1 values (3, 'c');

set kyuubi.operation.language=scala;

val df = spark.sql("select * from wangzhen_test_20231115_t1 order by id limit 2");
df.write.format("parquet").save("hfds://XXX/result.parquet");

spark.sql("set kyuubi.operation.language=sql");
select * from `parquet`.`hfds://XXX/result.parquet`;

== Physical Plan ==
Execute InsertIntoHadoopFsRelationCommand (5)
+- WriteFiles (4)
   +- TakeOrderedAndProject (3)
      +- * ColumnarToRow (2)
         +- Scan parquet spark_catalog.XXX.wangzhen_test_20231115_t1 (1)

[lsm1]

if (saveResultToPath) {
  spark.sql(statement).write.format(format).save(resultPath)
  result = spark.read.load(resultPath)
} else {
  result = spark.sql(statement)
}

we still call result.collect() later, if the result is too large,we can not avoid driver OOM

[wForget]

we still call result.collect() later, if the result is too large,we can not avoid driver OOM

Can we combine incremental collection mode?

[cxzl25]

I did a simple test and the results were as expected

Maybe you can test the scenario of generating multiple files

[wForget]

Maybe you can test the scenario of generating multiple files

Do you mean a large data set or multiple tasks?

[wForget]

Maybe you can test the scenario of generating multiple files

The output seems to be in order even when outputting multiple files.

set spark.sql.files.maxRecordsPerFile=10000;

set kyuubi.operation.language=scala;

spark.range(0, 1000000, 1, numPartitions = 10)
  .selectExpr("id", "cast(id as string) as name")
  .createOrReplaceTempView("wangzhen_test_20231115_tmp1")

val df = spark.sql("select * from wangzhen_test_20231115_tmp1 order by id limit 100000");
df.write.format("parquet").save("hdfs://XXX/result");

spark.sql("set kyuubi.operation.language=sql");
select * from `parquet`.`hdfs://XXX/result`;

[lsm1]

we still call result.collect() later, if the result is too large,we can not avoid driver OOM

Can we combine incremental collection mode?

When we use incremental collection mode, it may significantly impact performance.

[cxzl25]

The output seems to be in order even when outputting multiple files.

When Spark reads datasource, it will be sorted by file length, so there is no guarantee.

org.apache.spark.sql.execution.datasources.v2.FileScan#partitions

      partition.files.flatMap { file =>
        PartitionedFileUtil.splitFiles(
          file = file,
          isSplitable = isSplitable(file.getPath),
          maxSplitBytes = maxSplitBytes,
          partitionValues = partitionValues
        )
      }.toArray.sortBy(_.length)(implicitly[Ordering[Long]].reverse)

https://github.com/apache/spark/blob/4398bb5d37328e2f3594302d98f98803a379a2e9/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala#L146-L160

[cfmcgrady]

nit: only cleanup for the operation ExecuteStatement?

[lsm1]

There is no simple way to determine whether the session has executed ExecuteStatement

[turboFei]

Some question:

I wonder that, If the result is order needed, if we save the result into files and then read from when client fetching result, the result returned to users is not ordered as expected.

[pan3793]

@turboFei from the context, I think the implementation already reserves the global order, @cxzl25 could you please clarify it?

[lsm1]

Some question:

I wonder that, If the result is order needed, if we save the result into files and then read from when client fetching result, the result returned to users is not ordered as expected.

1. spark save ordered result to multiple part-X files in the filesystem, in order of the keys
   https://github.com/lsm1/spark/blob/master/core/src/main/scala/org/apache/spark/rdd/OrderedRDDFunctions.scala#L47

org.apache.spark.rdd.OrderedRDDFunctions#sortByKey
 /**
	* Sort the RDD by key, so that each partition contains a sorted range of the elements. Calling
	* `collect` or `save` on the resulting RDD will return or output an ordered list of records
	* (in the `save` case, they will be written to multiple `part-X` files in the filesystem, in
	* order of the keys).  
*/

2. When fetchOrcStatement read file, it will be sorted by file name, so it will return ordered result

[cfmcgrady]

thanks all. merging to master(v1.9.0)