[SPARK-54446][SQL][ML][CONNECT] FPGrowth supports local filesystem with Arrow file format #53232
Conversation
zhengruifeng
requested review from
HyukjinKwon,
WeichenXu123,
cloud-fan,
holdenk and
viirya
and removed request for
WeichenXu123
|
This PR is another attempt to save ml models containing dataframes to driver's local fs. |
| fileWriter.start() | ||
| while (batchBytesIter.hasNext) { | ||
| val batchBytes = batchBytesIter.next() | ||
| val batch = ArrowConverters.loadBatch(batchBytes, allocator) |
There was a problem hiding this comment.
The batch: ArrowRecordBatch doesn't extends Serializable, so still use the Array[Byte] as the underlying data in the PR.
sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowFileReadWrite.scala
Outdated
Show resolved
Hide resolved
sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowFileReadWrite.scala
Show resolved
Hide resolved
| def saveDataFrame(path: String, df: DataFrame): Unit = { | ||
| if (localSavingModeState.get()) { | ||
| val filePath = Paths.get(path) | ||
| Files.createDirectories(filePath.getParent) | ||
|
|
||
| df match { | ||
| case d: org.apache.spark.sql.classic.DataFrame => | ||
| ArrowFileReadWrite.save(d, path) | ||
| case _ => throw new UnsupportedOperationException("Unsupported dataframe type") | ||
| } | ||
| } else { | ||
| df.write.parquet(path) | ||
| } | ||
| } | ||
|
|
||
| def loadDataFrame(path: String, spark: SparkSession): DataFrame = { | ||
| if (localSavingModeState.get()) { | ||
| spark match { | ||
| case s: org.apache.spark.sql.classic.SparkSession => | ||
| ArrowFileReadWrite.load(s, path) | ||
| case _ => throw new UnsupportedOperationException("Unsupported session type") | ||
| } | ||
| } else { | ||
| spark.read.parquet(path) | ||
| } | ||
| } |
There was a problem hiding this comment.
So if we have localSavingModeState set to true this will write out an arrow file which is not stable format wise. It does look like localSavingModeState is only set to true in internal methods in Scala. Looking in the PySpark docstrings I see we tell people to use this API so I remain -0.9.
There was a problem hiding this comment.
hi @holdenk , as @WeichenXu123 explained #53150 (comment), this is a runtime temporary file in spark connect server side, and will be cleaned after session close.
So I think we don't have to use a stable format here.
There was a problem hiding this comment.
localSavingModeState is also used internally, (only Spark driver code can set the flag) . Where does the doc string mentioned it ? we should remove it from doc and mark localSavingModeState as private field
There was a problem hiding this comment.
Hmm, even it is just a temporary session file, is there any reason not to use Parquet but Arrow file format?
There was a problem hiding this comment.
we can read/write parquet with arrow, but it requires a new dependency
<dependency>
<groupId>org.apache.parquet</groupId>
<artifactId>parquet-arrow</artifactId>
</dependency>
otherwise, I am not sure whether we have utils to read/write parquet.
There was a problem hiding this comment.
At the end we need the in-memory data to be in arrow format, so using arrow file is more efficient.
viirya
left a comment
viirya
left a comment
There was a problem hiding this comment.
Wonder why choosing Arrow file format now instead of Parquet?
Due to the process of batch -> bytes -> batch -> bytes (when writing to file), it doesn't look like an efficient way.
| val rdd = df.toArrowBatchRdd(maxRecordsPerBatch, "UTC", true, false) | ||
| val arrowSchema = ArrowUtils.toArrowSchema(df.schema, "UTC", true, false) | ||
| val writer = new SparkArrowFileWriter(arrowSchema, path) | ||
| writer.write(rdd.toLocalIterator) |
There was a problem hiding this comment.
Instead, can we call toLocalIterator on original DataFrame's rdd and write rows to Arrow batches locally? Then we don't need to have the redundant Bytes?
There was a problem hiding this comment.
we can make best of the ArrowConverters utils, if we use the Bytes
sql/core/src/main/scala/org/apache/spark/sql/classic/Dataset.scala
Outdated
Show resolved
Hide resolved
sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowFileReadWrite.scala
Outdated
Show resolved
Hide resolved
zhengruifeng
changed the title
Can we have a shared util to produce RDD of arrow batches? Then we can either turn it to RDD of bytes, or write it to local files. |
This is actually already reusing a lot of existing utiles at Basically Below code is for |
zhengruifeng
changed the title
6 participants
What changes were proposed in this pull request?
FPGrowth supports local filesystem
Why are the changes needed?
to make FPGrowth work with local filesystem
Does this PR introduce any user-facing change?
yes, FPGrowth will work when local saving mode is one
How was this patch tested?
updated tests
Was this patch authored or co-authored using generative AI tooling?
no