You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "fanyue-xia (via GitHub)" <gi...@apache.org> on 2024/04/10 00:23:03 UTC
[PR] [SPARK-47788] [Structured Streaming] Ensure the same hash partitioning for streaming stateful ops [spark]
fanyue-xia opened a new pull request, #45971:
URL: https://github.com/apache/spark/pull/45971
### What changes were proposed in this pull request?
It adds StablePartitioningSuite which uses golden files to ensure hash partitions will not be changed in the future for streaming stateful ops.
### Why are the changes needed?
To really make sure that any changes to hash function / partitioner in Spark doesn’t cause logical correctness issues in existing running streaming queries.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
Unit test
### Was this patch authored or co-authored using generative AI tooling?
No
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788][SS] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "fanyue-xia (via GitHub)" <gi...@apache.org>.
fanyue-xia commented on code in PR #45971:
URL: https://github.com/apache/spark/pull/45971#discussion_r1565133758
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,219 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ private val goldenFileURI =
+ this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ private val schemaFileName = "randomSchemas" // files for storing random input schemas
+ private val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ private val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ private val numSchemaTypes = 1
+ private val maxNumFields = 20
+
+ // Configs for shuffle
+ private val numRows = 10000
+ private val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach { schema =>
+ writer.write(schema.toDDL)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try {
+ source
+ .getLines()
+ .map { ddl =>
+ StructType.fromDDL(ddl)
+ }
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ // Save the total number of rows
+ os.writeInt(rows.length)
+ // Save all rows
+ rows.foreach { row =>
+ // Save the row's total number of bytes
+ val rowBytes = row.getBytes()
+ // Save the row's actual bytes
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)
+ }
+ // Save all partIds, which should be in the same order as rows
+ partIds.foreach { id =>
+ os.writeInt(id)
+ }
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map { _ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(
+ schema,
+ rand = new Random(rand.nextInt())
+ )
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+
+ // UnsafeProjection returns the same UnsafeRow instance intentionally, so
+ // unless doing deep copy, the hash partitions below will evaluate on the
+ // same row thus return same value.
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(rand: Random): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ RandomDataGenerator.randomNestedSchema(rand, maxNumFields, variableTypes)
+ }
+ }
+
+ private def getPartitionId(rows: Seq[UnsafeRow], hash: HashPartitioning): Seq[Int] = {
+ val partIdExpr = hash.partitionIdExpression
+ rows.map { row =>
+ partIdExpr.eval(row).asInstanceOf[Int]
+ }
+ }
+
+ test("SPARK-47788: Ensure the same hash function is used across batches.") {
Review Comment:
Thanks for the suggestion. "compatibility across Spark versions" makes more sense from the user's point of view. Updated.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788] [Structured Streaming] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "fanyue-xia (via GitHub)" <gi...@apache.org>.
fanyue-xia commented on PR #45971:
URL: https://github.com/apache/spark/pull/45971#issuecomment-2052704919
> > the seed might behave differently across runs/on different machines
>
> Ah I see, this indeed makes sense.
>
> In this case, I think we should fix the generator of rows. It's okay to sacrifice randomness of rows here. We can have a dedicated row generation function, depending on the input type, this function just return a fixed return (e.g. if input is int, just return 233, if input is byte, just return 0xdeadbeef)
>
> Giving up randomness of the row should still get the job done. The way the hash is computed is sth like `hash(field 1, hash (field 2, seed...)...)`, and this part hasn't been touched likely since the beginning.
>
> https://github.com/apache/spark/blob/6ee662c28ffb0deb70f08a971f9c1869288d39ba/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/hash.scala#L289-L298
>
> This should not be changed, even if it is changed, it makes less sense to be changed to a hash function that hash (1, 2) and (2, 1) to be the same bucket.
>
> Then as long as we increase the max number of field and the number of schemas (now this can be a fairly high number, also pay attention to the test run time), it should behave similar as having a large number of random generated rows.
Thanks for your suggestion. I preserved the randomness of the row since it feels more intuitive and we don't need to reason about the underlying hashing schemes for `StructType`. Changed the code to limit to just one nested schema and limit the `Array` size and `String` length when generating rows. Now the golden file is smaller than 10mb.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788] [Structured Streaming] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "fanyue-xia (via GitHub)" <gi...@apache.org>.
fanyue-xia commented on PR #45971:
URL: https://github.com/apache/spark/pull/45971#issuecomment-2052224756
> Thanks for the effort! This really requires some deep understanding of spark internals...
>
> There is still one important concern, that the golden file size is too big. I looked a bit, it seems that the largest golden file is ~7MB. We should find a way to limit the file size to < 10MB.
>
> One improvement I can see is that, here you are storing both the rows and partition ids, but we don't need to store rows.
>
> Instead, we store the random seed, and regenerate the random rows in the check. By doing this we only need to store the seed, the schemas, and for each schema:
>
> 1. partition ids, and
> 2. numRows
>
> Now golden file size should be much smaller.
>
> This means that we trust `RandomDataGenerator` to generate the same row for the same seed everytime. The code hasn't been touched like 10 years so I think this should be safe. Although unlikely, when people really need to touch that later, they'll notice this test failure.
I have concerned that the seed might behave differently across runs/on different machines. Talked to @HeartSaVioR about it, he mentioned that he isn't sure whether it is the case; it’s safer to just store the generated inputs; if there is any difference in any random generation, it’s going to be uneasy to find.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788][SS] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "fanyue-xia (via GitHub)" <gi...@apache.org>.
fanyue-xia commented on code in PR #45971:
URL: https://github.com/apache/spark/pull/45971#discussion_r1565102542
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,219 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ private val goldenFileURI =
+ this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ private val schemaFileName = "randomSchemas" // files for storing random input schemas
+ private val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ private val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ private val numSchemaTypes = 1
+ private val maxNumFields = 20
+
+ // Configs for shuffle
+ private val numRows = 10000
+ private val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach { schema =>
+ writer.write(schema.toDDL)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try {
+ source
+ .getLines()
+ .map { ddl =>
+ StructType.fromDDL(ddl)
+ }
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ // Save the total number of rows
+ os.writeInt(rows.length)
+ // Save all rows
+ rows.foreach { row =>
+ // Save the row's total number of bytes
+ val rowBytes = row.getBytes()
+ // Save the row's actual bytes
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)
+ }
+ // Save all partIds, which should be in the same order as rows
+ partIds.foreach { id =>
+ os.writeInt(id)
+ }
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map { _ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(
+ schema,
+ rand = new Random(rand.nextInt())
+ )
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+
+ // UnsafeProjection returns the same UnsafeRow instance intentionally, so
+ // unless doing deep copy, the hash partitions below will evaluate on the
+ // same row thus return same value.
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(rand: Random): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ RandomDataGenerator.randomNestedSchema(rand, maxNumFields, variableTypes)
+ }
+ }
+
+ private def getPartitionId(rows: Seq[UnsafeRow], hash: HashPartitioning): Seq[Int] = {
+ val partIdExpr = hash.partitionIdExpression
+ rows.map { row =>
+ partIdExpr.eval(row).asInstanceOf[Int]
+ }
+ }
+
+ test("SPARK-47788: Ensure the same hash function is used across batches.") {
+ val rowAndPartIdFile = new File(goldenFileURI.getPath, rowAndPartIdFilename)
+
+ if (regenerateGoldenFiles) {
+ // To limit the golden file size under 10Mb, please set the final val MAX_STR_LEN: Int = 100
Review Comment:
In the oss-spark, it doesn't take parameters for max string length and max array size, which differs databricks' internal version.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788] [Structured Streaming] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "fanyue-xia (via GitHub)" <gi...@apache.org>.
fanyue-xia commented on code in PR #45971:
URL: https://github.com/apache/spark/pull/45971#discussion_r1563367019
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,199 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.StreamTest
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ val goldenFileURI = this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ val schemaFileName = "randomSchemas" // files for storing random input schemas
+ val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ val numSchemaTypes = 5
+ val maxNumFields = 10
+
+ // Configs for shuffle
+ val numRows = 10000
+ val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach{ schema =>
+ writer.write(schema.json)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try{
+ source.getLines().map{json => DataType.fromJson(json).asInstanceOf[StructType]}
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally
+ source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ os.writeInt(rows.length)
+ rows.foreach { row =>
+ val rowBytes = row.getBytes()
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)}
+ partIds.foreach {id => os.writeInt(id)}
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map{_ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(schema, rand = new Random(rand.nextInt(RandomDataGenerator.MAX_STR_LEN)))
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+ unsafeRow.copy()
Review Comment:
Added comments.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788][SS] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "fanyue-xia (via GitHub)" <gi...@apache.org>.
fanyue-xia commented on code in PR #45971:
URL: https://github.com/apache/spark/pull/45971#discussion_r1565133340
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,219 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ private val goldenFileURI =
+ this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ private val schemaFileName = "randomSchemas" // files for storing random input schemas
+ private val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ private val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ private val numSchemaTypes = 1
+ private val maxNumFields = 20
+
+ // Configs for shuffle
+ private val numRows = 10000
+ private val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach { schema =>
+ writer.write(schema.toDDL)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try {
+ source
+ .getLines()
+ .map { ddl =>
+ StructType.fromDDL(ddl)
+ }
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ // Save the total number of rows
+ os.writeInt(rows.length)
+ // Save all rows
+ rows.foreach { row =>
+ // Save the row's total number of bytes
+ val rowBytes = row.getBytes()
+ // Save the row's actual bytes
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)
+ }
+ // Save all partIds, which should be in the same order as rows
+ partIds.foreach { id =>
+ os.writeInt(id)
+ }
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map { _ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(
+ schema,
+ rand = new Random(rand.nextInt())
+ )
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+
+ // UnsafeProjection returns the same UnsafeRow instance intentionally, so
+ // unless doing deep copy, the hash partitions below will evaluate on the
+ // same row thus return same value.
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(rand: Random): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ RandomDataGenerator.randomNestedSchema(rand, maxNumFields, variableTypes)
+ }
+ }
+
+ private def getPartitionId(rows: Seq[UnsafeRow], hash: HashPartitioning): Seq[Int] = {
+ val partIdExpr = hash.partitionIdExpression
+ rows.map { row =>
+ partIdExpr.eval(row).asInstanceOf[Int]
+ }
+ }
+
+ test("SPARK-47788: Ensure the same hash function is used across batches.") {
+ val rowAndPartIdFile = new File(goldenFileURI.getPath, rowAndPartIdFilename)
+
+ if (regenerateGoldenFiles) {
+ // To limit the golden file size under 10Mb, please set the final val MAX_STR_LEN: Int = 100
+ // and final val MAX_ARR_SIZE: Int = 4 in org.apache.spark.sql.RandomDataGenerator
+
+ val seed = Random.nextInt()
+ val random = new Random(seed)
+ logInfo(s"Get random inputs with seed $seed")
+
+ val schemas = getRandomSchemas(random)
+
+ val os = new DataOutputStream(
+ codec.compressedOutputStream(new FileOutputStream(rowAndPartIdFile))
+ )
+
+ saveSchemas(schemas)
+
+ schemas.foreach { schema =>
+ // Streaming stateful ops rely on this distribution to partition the data.
+ // Spark should make sure this class's partition dependency remain unchanged.
+ val hash = StatefulOpClusteredDistribution(
+ Seq(BoundReference(0, schema, nullable = true)),
+ numShufflePartitions
+ ).createPartitioning(numShufflePartitions)
+
+ assert(sql/catalyst/src/test/scala/org/apache/spark/sql/RandomDataGenerator.scala
+ hash.isInstanceOf[HashPartitioning],
+ "StatefulOpClusteredDistribution should " +
+ "rely on HashPartitioning to ensure partitions remain the same for streaming " +
+ "stateful operators."
+ )
+ val rows = getRandomRows(schema, numRows, random)
+ val partitions = getPartitionId(rows, hash.asInstanceOf[HashPartitioning])
+ saveRowsAndPartIds(rows, partitions, os)
+ }
+ os.close()
+ } else {
+ val schemas = readSchemas()
+ val is = new DataInputStream(
+ codec.compressedInputStream(new FileInputStream(rowAndPartIdFile))
+ )
+ schemas.foreach { schema =>
+ val hash = StatefulOpClusteredDistribution(
+ Seq(BoundReference(0, schema, nullable = true)),
+ numShufflePartitions
+ ).createPartitioning(numShufflePartitions)
+
+ assert(
+ hash.isInstanceOf[HashPartitioning],
+ "StatefulOpClusteredDistribution should " +
+ "rely on HashPartitioning to ensure partitions remain the same for streaming " +
+ "stateful operators."
+ )
+
+ val (rows, expectedPartitions) = readRowsAndPartIds(is)
+ val partitions = getPartitionId(rows, hash.asInstanceOf[HashPartitioning])
+ assert(
+ partitions === expectedPartitions,
Review Comment:
Expanded the assertions to include number of total mismatch rows and samples
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788][SS] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "fanyue-xia (via GitHub)" <gi...@apache.org>.
fanyue-xia commented on code in PR #45971:
URL: https://github.com/apache/spark/pull/45971#discussion_r1565133109
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,219 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ private val goldenFileURI =
+ this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ private val schemaFileName = "randomSchemas" // files for storing random input schemas
+ private val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ private val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
Review Comment:
Added.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788][SS] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "fanyue-xia (via GitHub)" <gi...@apache.org>.
fanyue-xia commented on code in PR #45971:
URL: https://github.com/apache/spark/pull/45971#discussion_r1565133340
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,219 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ private val goldenFileURI =
+ this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ private val schemaFileName = "randomSchemas" // files for storing random input schemas
+ private val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ private val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ private val numSchemaTypes = 1
+ private val maxNumFields = 20
+
+ // Configs for shuffle
+ private val numRows = 10000
+ private val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach { schema =>
+ writer.write(schema.toDDL)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try {
+ source
+ .getLines()
+ .map { ddl =>
+ StructType.fromDDL(ddl)
+ }
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ // Save the total number of rows
+ os.writeInt(rows.length)
+ // Save all rows
+ rows.foreach { row =>
+ // Save the row's total number of bytes
+ val rowBytes = row.getBytes()
+ // Save the row's actual bytes
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)
+ }
+ // Save all partIds, which should be in the same order as rows
+ partIds.foreach { id =>
+ os.writeInt(id)
+ }
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map { _ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(
+ schema,
+ rand = new Random(rand.nextInt())
+ )
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+
+ // UnsafeProjection returns the same UnsafeRow instance intentionally, so
+ // unless doing deep copy, the hash partitions below will evaluate on the
+ // same row thus return same value.
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(rand: Random): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ RandomDataGenerator.randomNestedSchema(rand, maxNumFields, variableTypes)
+ }
+ }
+
+ private def getPartitionId(rows: Seq[UnsafeRow], hash: HashPartitioning): Seq[Int] = {
+ val partIdExpr = hash.partitionIdExpression
+ rows.map { row =>
+ partIdExpr.eval(row).asInstanceOf[Int]
+ }
+ }
+
+ test("SPARK-47788: Ensure the same hash function is used across batches.") {
+ val rowAndPartIdFile = new File(goldenFileURI.getPath, rowAndPartIdFilename)
+
+ if (regenerateGoldenFiles) {
+ // To limit the golden file size under 10Mb, please set the final val MAX_STR_LEN: Int = 100
+ // and final val MAX_ARR_SIZE: Int = 4 in org.apache.spark.sql.RandomDataGenerator
+
+ val seed = Random.nextInt()
+ val random = new Random(seed)
+ logInfo(s"Get random inputs with seed $seed")
+
+ val schemas = getRandomSchemas(random)
+
+ val os = new DataOutputStream(
+ codec.compressedOutputStream(new FileOutputStream(rowAndPartIdFile))
+ )
+
+ saveSchemas(schemas)
+
+ schemas.foreach { schema =>
+ // Streaming stateful ops rely on this distribution to partition the data.
+ // Spark should make sure this class's partition dependency remain unchanged.
+ val hash = StatefulOpClusteredDistribution(
+ Seq(BoundReference(0, schema, nullable = true)),
+ numShufflePartitions
+ ).createPartitioning(numShufflePartitions)
+
+ assert(sql/catalyst/src/test/scala/org/apache/spark/sql/RandomDataGenerator.scala
+ hash.isInstanceOf[HashPartitioning],
+ "StatefulOpClusteredDistribution should " +
+ "rely on HashPartitioning to ensure partitions remain the same for streaming " +
+ "stateful operators."
+ )
+ val rows = getRandomRows(schema, numRows, random)
+ val partitions = getPartitionId(rows, hash.asInstanceOf[HashPartitioning])
+ saveRowsAndPartIds(rows, partitions, os)
+ }
+ os.close()
+ } else {
+ val schemas = readSchemas()
+ val is = new DataInputStream(
+ codec.compressedInputStream(new FileInputStream(rowAndPartIdFile))
+ )
+ schemas.foreach { schema =>
+ val hash = StatefulOpClusteredDistribution(
+ Seq(BoundReference(0, schema, nullable = true)),
+ numShufflePartitions
+ ).createPartitioning(numShufflePartitions)
+
+ assert(
+ hash.isInstanceOf[HashPartitioning],
+ "StatefulOpClusteredDistribution should " +
+ "rely on HashPartitioning to ensure partitions remain the same for streaming " +
+ "stateful operators."
+ )
+
+ val (rows, expectedPartitions) = readRowsAndPartIds(is)
+ val partitions = getPartitionId(rows, hash.asInstanceOf[HashPartitioning])
+ assert(
+ partitions === expectedPartitions,
Review Comment:
Expanded the assertions to include the number of total mismatch rows and samples. Also log the schema info at the top.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788][SS] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "fanyue-xia (via GitHub)" <gi...@apache.org>.
fanyue-xia commented on code in PR #45971:
URL: https://github.com/apache/spark/pull/45971#discussion_r1565102542
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,219 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ private val goldenFileURI =
+ this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ private val schemaFileName = "randomSchemas" // files for storing random input schemas
+ private val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ private val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ private val numSchemaTypes = 1
+ private val maxNumFields = 20
+
+ // Configs for shuffle
+ private val numRows = 10000
+ private val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach { schema =>
+ writer.write(schema.toDDL)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try {
+ source
+ .getLines()
+ .map { ddl =>
+ StructType.fromDDL(ddl)
+ }
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ // Save the total number of rows
+ os.writeInt(rows.length)
+ // Save all rows
+ rows.foreach { row =>
+ // Save the row's total number of bytes
+ val rowBytes = row.getBytes()
+ // Save the row's actual bytes
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)
+ }
+ // Save all partIds, which should be in the same order as rows
+ partIds.foreach { id =>
+ os.writeInt(id)
+ }
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map { _ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(
+ schema,
+ rand = new Random(rand.nextInt())
+ )
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+
+ // UnsafeProjection returns the same UnsafeRow instance intentionally, so
+ // unless doing deep copy, the hash partitions below will evaluate on the
+ // same row thus return same value.
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(rand: Random): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ RandomDataGenerator.randomNestedSchema(rand, maxNumFields, variableTypes)
+ }
+ }
+
+ private def getPartitionId(rows: Seq[UnsafeRow], hash: HashPartitioning): Seq[Int] = {
+ val partIdExpr = hash.partitionIdExpression
+ rows.map { row =>
+ partIdExpr.eval(row).asInstanceOf[Int]
+ }
+ }
+
+ test("SPARK-47788: Ensure the same hash function is used across batches.") {
+ val rowAndPartIdFile = new File(goldenFileURI.getPath, rowAndPartIdFilename)
+
+ if (regenerateGoldenFiles) {
+ // To limit the golden file size under 10Mb, please set the final val MAX_STR_LEN: Int = 100
Review Comment:
In the oss-spark, it doesn't take parameters for max string length and max array size, which differs databricks' internal version.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788] [Structured Streaming] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on code in PR #45971:
URL: https://github.com/apache/spark/pull/45971#discussion_r1562849582
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,199 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.StreamTest
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ val goldenFileURI = this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ val schemaFileName = "randomSchemas" // files for storing random input schemas
+ val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ val numSchemaTypes = 5
+ val maxNumFields = 10
+
+ // Configs for shuffle
+ val numRows = 10000
+ val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach{ schema =>
+ writer.write(schema.json)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try{
+ source.getLines().map{json => DataType.fromJson(json).asInstanceOf[StructType]}
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally
+ source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ os.writeInt(rows.length)
+ rows.foreach { row =>
+ val rowBytes = row.getBytes()
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)}
+ partIds.foreach {id => os.writeInt(id)}
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map{_ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(schema, rand = new Random(rand.nextInt(RandomDataGenerator.MAX_STR_LEN)))
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+ unsafeRow.copy()
Review Comment:
Just curious why need a copy here?
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,199 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.StreamTest
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ val goldenFileURI = this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ val schemaFileName = "randomSchemas" // files for storing random input schemas
+ val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ val numSchemaTypes = 5
+ val maxNumFields = 10
+
+ // Configs for shuffle
+ val numRows = 10000
+ val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach{ schema =>
+ writer.write(schema.json)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try{
+ source.getLines().map{json => DataType.fromJson(json).asInstanceOf[StructType]}
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally
+ source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ os.writeInt(rows.length)
+ rows.foreach { row =>
+ val rowBytes = row.getBytes()
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)}
+ partIds.foreach {id => os.writeInt(id)}
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map{_ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(schema, rand = new Random(rand.nextInt(RandomDataGenerator.MAX_STR_LEN)))
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
Review Comment:
Thanks for the useful comment!
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,199 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.StreamTest
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ val goldenFileURI = this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ val schemaFileName = "randomSchemas" // files for storing random input schemas
+ val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ val numSchemaTypes = 5
+ val maxNumFields = 10
+
+ // Configs for shuffle
+ val numRows = 10000
+ val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach{ schema =>
+ writer.write(schema.json)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try{
+ source.getLines().map{json => DataType.fromJson(json).asInstanceOf[StructType]}
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally
+ source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ os.writeInt(rows.length)
+ rows.foreach { row =>
+ val rowBytes = row.getBytes()
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)}
+ partIds.foreach {id => os.writeInt(id)}
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map{_ =>
Review Comment:
nit:
`(1 to numRows).map { _ =>`
spaces between `map`, `{` and `_`
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,199 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.StreamTest
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ val goldenFileURI = this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ val schemaFileName = "randomSchemas" // files for storing random input schemas
+ val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ val numSchemaTypes = 5
+ val maxNumFields = 10
+
+ // Configs for shuffle
+ val numRows = 10000
+ val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach{ schema =>
+ writer.write(schema.json)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try{
+ source.getLines().map{json => DataType.fromJson(json).asInstanceOf[StructType]}
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally
+ source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ os.writeInt(rows.length)
+ rows.foreach { row =>
+ val rowBytes = row.getBytes()
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)}
+ partIds.foreach {id => os.writeInt(id)}
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map{_ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
Review Comment:
nit:
space between `(` and `_`
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,199 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.StreamTest
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ val goldenFileURI = this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ val schemaFileName = "randomSchemas" // files for storing random input schemas
+ val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ val numSchemaTypes = 5
+ val maxNumFields = 10
+
+ // Configs for shuffle
+ val numRows = 10000
+ val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach{ schema =>
+ writer.write(schema.json)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try{
+ source.getLines().map{json => DataType.fromJson(json).asInstanceOf[StructType]}
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally
+ source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ os.writeInt(rows.length)
+ rows.foreach { row =>
+ val rowBytes = row.getBytes()
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)}
+ partIds.foreach {id => os.writeInt(id)}
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map{_ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(schema, rand = new Random(rand.nextInt(RandomDataGenerator.MAX_STR_LEN)))
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(numSchemaTypes: Int, rand: Random,
+ maxNumFields: Int): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ val numField = rand.nextInt(maxNumFields)
+ RandomDataGenerator.randomSchema(rand, Math.max(1, numField), variableTypes)
+ }
+ }
+
+ private def getPartitionId(rows: Seq[UnsafeRow], hash: HashPartitioning): Seq[Int] = {
+ val partIdExpr = hash.partitionIdExpression
+ rows.map{row =>
+ partIdExpr.eval(row).asInstanceOf[Int]
+ }
+ }
+
+ test("SPARK-47788: Ensure the same hash function is used across batches.") {
+ val rowAndPartIdFile = new File(goldenFileURI.getPath, rowAndPartIdFilename)
+
+ if (regenerateGoldenFiles) {
+ val seed = Random.nextInt()
+ val random = new Random(seed)
+ logInfo(s"Get random inputs with seed $seed")
+
+ val schemas = getRandomSchemas(numSchemaTypes, random, maxNumFields)
+
+ val os = new DataOutputStream(
+ codec.compressedOutputStream(new FileOutputStream(rowAndPartIdFile))
+ )
+
+ saveSchemas(schemas)
+ schemas.foreach { schema =>
+ // Streaming stateful ops rely on this distribution to partition the data.
+ // Spark should make sure this class's partition dependency remain unchanged.
+ val hash = StatefulOpClusteredDistribution(
+ Seq(BoundReference(0, schema, nullable = true)),
+ numShufflePartitions
+ ).createPartitioning(numShufflePartitions)
+
+ assert(
+ hash.isInstanceOf[HashPartitioning],
+ "StatefulOpClusteredDistribution should " +
+ "rely on HashPartitioning to ensure partitions remain the same for streaming " +
+ "stateful operators."
+ )
+ val rows = getRandomRows(schema, numRows, random)
+ val partitions = getPartitionId(rows, hash.asInstanceOf[HashPartitioning])
+ saveRowsAndPartIds(rows, partitions, os)
+
Review Comment:
nit remove this empty line
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,199 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.StreamTest
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ val goldenFileURI = this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ val schemaFileName = "randomSchemas" // files for storing random input schemas
+ val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ val numSchemaTypes = 5
+ val maxNumFields = 10
+
+ // Configs for shuffle
+ val numRows = 10000
+ val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach{ schema =>
+ writer.write(schema.json)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try{
+ source.getLines().map{json => DataType.fromJson(json).asInstanceOf[StructType]}
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally
+ source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ os.writeInt(rows.length)
+ rows.foreach { row =>
+ val rowBytes = row.getBytes()
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)}
+ partIds.foreach {id => os.writeInt(id)}
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map{_ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(schema, rand = new Random(rand.nextInt(RandomDataGenerator.MAX_STR_LEN)))
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(numSchemaTypes: Int, rand: Random,
+ maxNumFields: Int): Seq[StructType] = {
Review Comment:
indent
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,199 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.StreamTest
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ val goldenFileURI = this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ val schemaFileName = "randomSchemas" // files for storing random input schemas
+ val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ val numSchemaTypes = 5
+ val maxNumFields = 10
+
+ // Configs for shuffle
+ val numRows = 10000
+ val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach{ schema =>
+ writer.write(schema.json)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try{
+ source.getLines().map{json => DataType.fromJson(json).asInstanceOf[StructType]}
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally
+ source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ os.writeInt(rows.length)
+ rows.foreach { row =>
+ val rowBytes = row.getBytes()
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)}
+ partIds.foreach {id => os.writeInt(id)}
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map{_ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(schema, rand = new Random(rand.nextInt(RandomDataGenerator.MAX_STR_LEN)))
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(numSchemaTypes: Int, rand: Random,
+ maxNumFields: Int): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ val numField = rand.nextInt(maxNumFields)
+ RandomDataGenerator.randomSchema(rand, Math.max(1, numField), variableTypes)
Review Comment:
Should this be `Math.min`?
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,199 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.StreamTest
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ val goldenFileURI = this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ val schemaFileName = "randomSchemas" // files for storing random input schemas
+ val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ val numSchemaTypes = 5
+ val maxNumFields = 10
+
+ // Configs for shuffle
+ val numRows = 10000
+ val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach{ schema =>
+ writer.write(schema.json)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try{
+ source.getLines().map{json => DataType.fromJson(json).asInstanceOf[StructType]}
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally
+ source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ os.writeInt(rows.length)
+ rows.foreach { row =>
+ val rowBytes = row.getBytes()
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)}
+ partIds.foreach {id => os.writeInt(id)}
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map{_ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(schema, rand = new Random(rand.nextInt(RandomDataGenerator.MAX_STR_LEN)))
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(numSchemaTypes: Int, rand: Random,
+ maxNumFields: Int): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ val numField = rand.nextInt(maxNumFields)
+ RandomDataGenerator.randomSchema(rand, Math.max(1, numField), variableTypes)
+ }
+ }
+
+ private def getPartitionId(rows: Seq[UnsafeRow], hash: HashPartitioning): Seq[Int] = {
+ val partIdExpr = hash.partitionIdExpression
+ rows.map{row =>
Review Comment:
nit: `ros.map { row =>`
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,199 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.StreamTest
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ val goldenFileURI = this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ val schemaFileName = "randomSchemas" // files for storing random input schemas
+ val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ val numSchemaTypes = 5
+ val maxNumFields = 10
+
+ // Configs for shuffle
+ val numRows = 10000
+ val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach{ schema =>
+ writer.write(schema.json)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try{
+ source.getLines().map{json => DataType.fromJson(json).asInstanceOf[StructType]}
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally
+ source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ os.writeInt(rows.length)
+ rows.foreach { row =>
+ val rowBytes = row.getBytes()
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)}
+ partIds.foreach {id => os.writeInt(id)}
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map{_ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(schema, rand = new Random(rand.nextInt(RandomDataGenerator.MAX_STR_LEN)))
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(numSchemaTypes: Int, rand: Random,
+ maxNumFields: Int): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ val numField = rand.nextInt(maxNumFields)
+ RandomDataGenerator.randomSchema(rand, Math.max(1, numField), variableTypes)
+ }
+ }
+
+ private def getPartitionId(rows: Seq[UnsafeRow], hash: HashPartitioning): Seq[Int] = {
+ val partIdExpr = hash.partitionIdExpression
+ rows.map{row =>
+ partIdExpr.eval(row).asInstanceOf[Int]
+ }
+ }
+
+ test("SPARK-47788: Ensure the same hash function is used across batches.") {
+ val rowAndPartIdFile = new File(goldenFileURI.getPath, rowAndPartIdFilename)
+
+ if (regenerateGoldenFiles) {
+ val seed = Random.nextInt()
+ val random = new Random(seed)
+ logInfo(s"Get random inputs with seed $seed")
+
+ val schemas = getRandomSchemas(numSchemaTypes, random, maxNumFields)
+
+ val os = new DataOutputStream(
+ codec.compressedOutputStream(new FileOutputStream(rowAndPartIdFile))
+ )
+
+ saveSchemas(schemas)
+ schemas.foreach { schema =>
+ // Streaming stateful ops rely on this distribution to partition the data.
+ // Spark should make sure this class's partition dependency remain unchanged.
+ val hash = StatefulOpClusteredDistribution(
+ Seq(BoundReference(0, schema, nullable = true)),
+ numShufflePartitions
+ ).createPartitioning(numShufflePartitions)
+
+ assert(
+ hash.isInstanceOf[HashPartitioning],
+ "StatefulOpClusteredDistribution should " +
+ "rely on HashPartitioning to ensure partitions remain the same for streaming " +
+ "stateful operators."
+ )
+ val rows = getRandomRows(schema, numRows, random)
+ val partitions = getPartitionId(rows, hash.asInstanceOf[HashPartitioning])
+ saveRowsAndPartIds(rows, partitions, os)
Review Comment:
So in the binary golden file, the format is like (no new line, added for better understanding):
```
<num rows>
row1 part1
row2 part2
...
<num rows>
row1 part1
row2 part2
...
```
If so, could you add a comment here explaining this layout? Thanks in advance!
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,199 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.StreamTest
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ val goldenFileURI = this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ val schemaFileName = "randomSchemas" // files for storing random input schemas
+ val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ val numSchemaTypes = 5
+ val maxNumFields = 10
+
+ // Configs for shuffle
+ val numRows = 10000
+ val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach{ schema =>
Review Comment:
nit: `foreach { schema =>`
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788][SS] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR closed pull request #45971: [SPARK-47788][SS] Ensure the same hash partitioning for streaming stateful ops
URL: https://github.com/apache/spark/pull/45971
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788][SS] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #45971:
URL: https://github.com/apache/spark/pull/45971#issuecomment-2055535812
Thanks! Merging to master.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788][SS] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #45971:
URL: https://github.com/apache/spark/pull/45971#discussion_r1565034767
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,219 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ private val goldenFileURI =
+ this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ private val schemaFileName = "randomSchemas" // files for storing random input schemas
+ private val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ private val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ private val numSchemaTypes = 1
+ private val maxNumFields = 20
+
+ // Configs for shuffle
+ private val numRows = 10000
+ private val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach { schema =>
+ writer.write(schema.toDDL)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try {
+ source
+ .getLines()
+ .map { ddl =>
+ StructType.fromDDL(ddl)
+ }
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ // Save the total number of rows
+ os.writeInt(rows.length)
+ // Save all rows
+ rows.foreach { row =>
+ // Save the row's total number of bytes
+ val rowBytes = row.getBytes()
+ // Save the row's actual bytes
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)
+ }
+ // Save all partIds, which should be in the same order as rows
+ partIds.foreach { id =>
+ os.writeInt(id)
+ }
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map { _ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(
+ schema,
+ rand = new Random(rand.nextInt())
+ )
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+
+ // UnsafeProjection returns the same UnsafeRow instance intentionally, so
+ // unless doing deep copy, the hash partitions below will evaluate on the
+ // same row thus return same value.
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(rand: Random): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ RandomDataGenerator.randomNestedSchema(rand, maxNumFields, variableTypes)
+ }
+ }
+
+ private def getPartitionId(rows: Seq[UnsafeRow], hash: HashPartitioning): Seq[Int] = {
+ val partIdExpr = hash.partitionIdExpression
+ rows.map { row =>
+ partIdExpr.eval(row).asInstanceOf[Int]
+ }
+ }
+
+ test("SPARK-47788: Ensure the same hash function is used across batches.") {
+ val rowAndPartIdFile = new File(goldenFileURI.getPath, rowAndPartIdFilename)
+
+ if (regenerateGoldenFiles) {
+ // To limit the golden file size under 10Mb, please set the final val MAX_STR_LEN: Int = 100
Review Comment:
Do I understand correctly that RandomDataGenerator.forType receives the parameters for max string length and max array size? We don't need to let others to fix the code of RandomDataGenerator. Let's just do that by ourselves via providing the parameters.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788][SS] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "fanyue-xia (via GitHub)" <gi...@apache.org>.
fanyue-xia commented on code in PR #45971:
URL: https://github.com/apache/spark/pull/45971#discussion_r1565112083
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,219 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
Review Comment:
Good point. I changed it before in case of doing an end to end streaming query test, but now it can be deleted since the additional test is not necessary.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788] [Structured Streaming] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on code in PR #45971:
URL: https://github.com/apache/spark/pull/45971#discussion_r1562953554
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,199 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.StreamTest
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ val goldenFileURI = this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ val schemaFileName = "randomSchemas" // files for storing random input schemas
+ val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ val numSchemaTypes = 5
+ val maxNumFields = 10
+
+ // Configs for shuffle
+ val numRows = 10000
+ val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach{ schema =>
+ writer.write(schema.json)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try{
+ source.getLines().map{json => DataType.fromJson(json).asInstanceOf[StructType]}
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally
+ source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ os.writeInt(rows.length)
+ rows.foreach { row =>
+ val rowBytes = row.getBytes()
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)}
+ partIds.foreach {id => os.writeInt(id)}
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map{_ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(schema, rand = new Random(rand.nextInt(RandomDataGenerator.MAX_STR_LEN)))
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(numSchemaTypes: Int, rand: Random,
+ maxNumFields: Int): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ val numField = rand.nextInt(maxNumFields)
+ RandomDataGenerator.randomSchema(rand, Math.max(1, numField), variableTypes)
Review Comment:
oh you are right..
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788] [Structured Streaming] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on PR #45971:
URL: https://github.com/apache/spark/pull/45971#issuecomment-2052183431
Thanks for the effort! This really requires some deep understanding of spark internals...
There is still one important concern, that the golden file size is too big. I looked a bit, it seems that the largest golden file is ~7MB. We should find a way to limit the file size to < 10MB.
One improvement I can see is that, here you are storing both the rows and partition ids.
I think we don't need to store rows here. Instead, we store the random seed, and regenerate the random rows in the check.
By doing this we only need to store the seed, the schemas, and for each schema:
1. partition ids, and
2. numRows
Now golden file size should be much smaller.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788] [Structured Streaming] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "fanyue-xia (via GitHub)" <gi...@apache.org>.
fanyue-xia commented on code in PR #45971:
URL: https://github.com/apache/spark/pull/45971#discussion_r1563365767
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,199 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.StreamTest
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ val goldenFileURI = this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ val schemaFileName = "randomSchemas" // files for storing random input schemas
+ val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ val numSchemaTypes = 5
+ val maxNumFields = 10
+
+ // Configs for shuffle
+ val numRows = 10000
+ val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach{ schema =>
+ writer.write(schema.json)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try{
+ source.getLines().map{json => DataType.fromJson(json).asInstanceOf[StructType]}
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally
+ source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ os.writeInt(rows.length)
+ rows.foreach { row =>
+ val rowBytes = row.getBytes()
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)}
+ partIds.foreach {id => os.writeInt(id)}
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map{_ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(schema, rand = new Random(rand.nextInt(RandomDataGenerator.MAX_STR_LEN)))
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(numSchemaTypes: Int, rand: Random,
+ maxNumFields: Int): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ val numField = rand.nextInt(maxNumFields)
+ RandomDataGenerator.randomSchema(rand, Math.max(1, numField), variableTypes)
+ }
+ }
+
+ private def getPartitionId(rows: Seq[UnsafeRow], hash: HashPartitioning): Seq[Int] = {
+ val partIdExpr = hash.partitionIdExpression
+ rows.map{row =>
+ partIdExpr.eval(row).asInstanceOf[Int]
+ }
+ }
+
+ test("SPARK-47788: Ensure the same hash function is used across batches.") {
+ val rowAndPartIdFile = new File(goldenFileURI.getPath, rowAndPartIdFilename)
+
+ if (regenerateGoldenFiles) {
+ val seed = Random.nextInt()
+ val random = new Random(seed)
+ logInfo(s"Get random inputs with seed $seed")
+
+ val schemas = getRandomSchemas(numSchemaTypes, random, maxNumFields)
+
+ val os = new DataOutputStream(
+ codec.compressedOutputStream(new FileOutputStream(rowAndPartIdFile))
+ )
+
+ saveSchemas(schemas)
+ schemas.foreach { schema =>
+ // Streaming stateful ops rely on this distribution to partition the data.
+ // Spark should make sure this class's partition dependency remain unchanged.
+ val hash = StatefulOpClusteredDistribution(
+ Seq(BoundReference(0, schema, nullable = true)),
+ numShufflePartitions
+ ).createPartitioning(numShufflePartitions)
+
+ assert(
+ hash.isInstanceOf[HashPartitioning],
+ "StatefulOpClusteredDistribution should " +
+ "rely on HashPartitioning to ensure partitions remain the same for streaming " +
+ "stateful operators."
+ )
+ val rows = getRandomRows(schema, numRows, random)
+ val partitions = getPartitionId(rows, hash.asInstanceOf[HashPartitioning])
+ saveRowsAndPartIds(rows, partitions, os)
Review Comment:
The layout is not for human readable since the rows will be compressed (and without space and newline), so instead of visualizing the layout, added a few comments.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788][SS] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #45971:
URL: https://github.com/apache/spark/pull/45971#issuecomment-2055533095
CI failed only from Spark connect, and looks like known flakiness (or consistent failure):
```
interrupt tag *** FAILED *** (30 seconds, 197 milliseconds)
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788] [Structured Streaming] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "fanyue-xia (via GitHub)" <gi...@apache.org>.
fanyue-xia commented on code in PR #45971:
URL: https://github.com/apache/spark/pull/45971#discussion_r1563379719
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,199 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.StreamTest
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ val goldenFileURI = this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ val schemaFileName = "randomSchemas" // files for storing random input schemas
+ val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ val numSchemaTypes = 5
+ val maxNumFields = 10
+
+ // Configs for shuffle
+ val numRows = 10000
+ val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach{ schema =>
+ writer.write(schema.json)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try{
+ source.getLines().map{json => DataType.fromJson(json).asInstanceOf[StructType]}
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally
+ source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ os.writeInt(rows.length)
+ rows.foreach { row =>
+ val rowBytes = row.getBytes()
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)}
+ partIds.foreach {id => os.writeInt(id)}
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map{_ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
Review Comment:
I run the formatter and it doesn't seem to add the space.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788][SS] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #45971:
URL: https://github.com/apache/spark/pull/45971#discussion_r1565032444
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,219 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ private val goldenFileURI =
+ this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ private val schemaFileName = "randomSchemas" // files for storing random input schemas
+ private val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ private val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ private val numSchemaTypes = 1
+ private val maxNumFields = 20
+
+ // Configs for shuffle
+ private val numRows = 10000
+ private val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach { schema =>
+ writer.write(schema.toDDL)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try {
+ source
+ .getLines()
+ .map { ddl =>
+ StructType.fromDDL(ddl)
+ }
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ // Save the total number of rows
+ os.writeInt(rows.length)
+ // Save all rows
+ rows.foreach { row =>
+ // Save the row's total number of bytes
+ val rowBytes = row.getBytes()
+ // Save the row's actual bytes
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)
+ }
+ // Save all partIds, which should be in the same order as rows
+ partIds.foreach { id =>
+ os.writeInt(id)
+ }
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map { _ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(
+ schema,
+ rand = new Random(rand.nextInt())
+ )
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+
+ // UnsafeProjection returns the same UnsafeRow instance intentionally, so
+ // unless doing deep copy, the hash partitions below will evaluate on the
+ // same row thus return same value.
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(rand: Random): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ RandomDataGenerator.randomNestedSchema(rand, maxNumFields, variableTypes)
+ }
+ }
+
+ private def getPartitionId(rows: Seq[UnsafeRow], hash: HashPartitioning): Seq[Int] = {
+ val partIdExpr = hash.partitionIdExpression
+ rows.map { row =>
+ partIdExpr.eval(row).asInstanceOf[Int]
+ }
+ }
+
+ test("SPARK-47788: Ensure the same hash function is used across batches.") {
Review Comment:
Probably the better frame for this is, we ensure the same has function to be used across Spark versions. That's technically the same when the query upgrades the Spark version during the restart, but emphasizing "compatibility across Spark versions" is more direct to what we verify here.
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,219 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ private val goldenFileURI =
+ this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ private val schemaFileName = "randomSchemas" // files for storing random input schemas
+ private val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ private val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ private val numSchemaTypes = 1
+ private val maxNumFields = 20
+
+ // Configs for shuffle
+ private val numRows = 10000
+ private val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
Review Comment:
nit: empty line between val and def
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,219 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ private val goldenFileURI =
+ this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ private val schemaFileName = "randomSchemas" // files for storing random input schemas
+ private val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ private val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ private val numSchemaTypes = 1
+ private val maxNumFields = 20
+
+ // Configs for shuffle
+ private val numRows = 10000
+ private val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach { schema =>
+ writer.write(schema.toDDL)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try {
+ source
+ .getLines()
+ .map { ddl =>
+ StructType.fromDDL(ddl)
+ }
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ // Save the total number of rows
+ os.writeInt(rows.length)
+ // Save all rows
+ rows.foreach { row =>
+ // Save the row's total number of bytes
+ val rowBytes = row.getBytes()
+ // Save the row's actual bytes
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)
+ }
+ // Save all partIds, which should be in the same order as rows
+ partIds.foreach { id =>
+ os.writeInt(id)
+ }
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map { _ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(
+ schema,
+ rand = new Random(rand.nextInt())
+ )
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+
+ // UnsafeProjection returns the same UnsafeRow instance intentionally, so
+ // unless doing deep copy, the hash partitions below will evaluate on the
+ // same row thus return same value.
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(rand: Random): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ RandomDataGenerator.randomNestedSchema(rand, maxNumFields, variableTypes)
+ }
+ }
+
+ private def getPartitionId(rows: Seq[UnsafeRow], hash: HashPartitioning): Seq[Int] = {
+ val partIdExpr = hash.partitionIdExpression
+ rows.map { row =>
+ partIdExpr.eval(row).asInstanceOf[Int]
+ }
+ }
+
+ test("SPARK-47788: Ensure the same hash function is used across batches.") {
+ val rowAndPartIdFile = new File(goldenFileURI.getPath, rowAndPartIdFilename)
+
+ if (regenerateGoldenFiles) {
+ // To limit the golden file size under 10Mb, please set the final val MAX_STR_LEN: Int = 100
+ // and final val MAX_ARR_SIZE: Int = 4 in org.apache.spark.sql.RandomDataGenerator
+
+ val seed = Random.nextInt()
+ val random = new Random(seed)
+ logInfo(s"Get random inputs with seed $seed")
+
+ val schemas = getRandomSchemas(random)
+
+ val os = new DataOutputStream(
+ codec.compressedOutputStream(new FileOutputStream(rowAndPartIdFile))
+ )
+
+ saveSchemas(schemas)
+
+ schemas.foreach { schema =>
+ // Streaming stateful ops rely on this distribution to partition the data.
+ // Spark should make sure this class's partition dependency remain unchanged.
+ val hash = StatefulOpClusteredDistribution(
+ Seq(BoundReference(0, schema, nullable = true)),
+ numShufflePartitions
+ ).createPartitioning(numShufflePartitions)
+
+ assert(sql/catalyst/src/test/scala/org/apache/spark/sql/RandomDataGenerator.scala
Review Comment:
Looks like broken code?
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,219 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ private val goldenFileURI =
+ this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ private val schemaFileName = "randomSchemas" // files for storing random input schemas
+ private val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ private val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ private val numSchemaTypes = 1
+ private val maxNumFields = 20
+
+ // Configs for shuffle
+ private val numRows = 10000
+ private val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach { schema =>
+ writer.write(schema.toDDL)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try {
+ source
+ .getLines()
+ .map { ddl =>
+ StructType.fromDDL(ddl)
+ }
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ // Save the total number of rows
+ os.writeInt(rows.length)
+ // Save all rows
+ rows.foreach { row =>
+ // Save the row's total number of bytes
+ val rowBytes = row.getBytes()
+ // Save the row's actual bytes
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)
+ }
+ // Save all partIds, which should be in the same order as rows
+ partIds.foreach { id =>
+ os.writeInt(id)
+ }
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map { _ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(
+ schema,
+ rand = new Random(rand.nextInt())
+ )
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+
+ // UnsafeProjection returns the same UnsafeRow instance intentionally, so
+ // unless doing deep copy, the hash partitions below will evaluate on the
+ // same row thus return same value.
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(rand: Random): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ RandomDataGenerator.randomNestedSchema(rand, maxNumFields, variableTypes)
+ }
+ }
+
+ private def getPartitionId(rows: Seq[UnsafeRow], hash: HashPartitioning): Seq[Int] = {
+ val partIdExpr = hash.partitionIdExpression
+ rows.map { row =>
+ partIdExpr.eval(row).asInstanceOf[Int]
+ }
+ }
+
+ test("SPARK-47788: Ensure the same hash function is used across batches.") {
+ val rowAndPartIdFile = new File(goldenFileURI.getPath, rowAndPartIdFilename)
+
+ if (regenerateGoldenFiles) {
+ // To limit the golden file size under 10Mb, please set the final val MAX_STR_LEN: Int = 100
+ // and final val MAX_ARR_SIZE: Int = 4 in org.apache.spark.sql.RandomDataGenerator
+
+ val seed = Random.nextInt()
+ val random = new Random(seed)
Review Comment:
traditional approach for using random seed (except the case you create random instance too quickly - multiple in several milliseconds) is System.currentTimeMillis() as seed, which Random with no param will just do it.
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,219 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ private val goldenFileURI =
+ this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ private val schemaFileName = "randomSchemas" // files for storing random input schemas
+ private val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ private val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ private val numSchemaTypes = 1
+ private val maxNumFields = 20
+
+ // Configs for shuffle
+ private val numRows = 10000
+ private val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach { schema =>
+ writer.write(schema.toDDL)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try {
+ source
+ .getLines()
+ .map { ddl =>
+ StructType.fromDDL(ddl)
+ }
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ // Save the total number of rows
+ os.writeInt(rows.length)
+ // Save all rows
+ rows.foreach { row =>
+ // Save the row's total number of bytes
+ val rowBytes = row.getBytes()
+ // Save the row's actual bytes
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)
+ }
+ // Save all partIds, which should be in the same order as rows
+ partIds.foreach { id =>
+ os.writeInt(id)
+ }
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map { _ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(
+ schema,
+ rand = new Random(rand.nextInt())
+ )
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+
+ // UnsafeProjection returns the same UnsafeRow instance intentionally, so
+ // unless doing deep copy, the hash partitions below will evaluate on the
+ // same row thus return same value.
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(rand: Random): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ RandomDataGenerator.randomNestedSchema(rand, maxNumFields, variableTypes)
+ }
+ }
+
+ private def getPartitionId(rows: Seq[UnsafeRow], hash: HashPartitioning): Seq[Int] = {
+ val partIdExpr = hash.partitionIdExpression
+ rows.map { row =>
+ partIdExpr.eval(row).asInstanceOf[Int]
+ }
+ }
+
+ test("SPARK-47788: Ensure the same hash function is used across batches.") {
+ val rowAndPartIdFile = new File(goldenFileURI.getPath, rowAndPartIdFilename)
+
+ if (regenerateGoldenFiles) {
+ // To limit the golden file size under 10Mb, please set the final val MAX_STR_LEN: Int = 100
+ // and final val MAX_ARR_SIZE: Int = 4 in org.apache.spark.sql.RandomDataGenerator
+
+ val seed = Random.nextInt()
+ val random = new Random(seed)
+ logInfo(s"Get random inputs with seed $seed")
+
+ val schemas = getRandomSchemas(random)
+
+ val os = new DataOutputStream(
+ codec.compressedOutputStream(new FileOutputStream(rowAndPartIdFile))
+ )
+
+ saveSchemas(schemas)
+
+ schemas.foreach { schema =>
+ // Streaming stateful ops rely on this distribution to partition the data.
+ // Spark should make sure this class's partition dependency remain unchanged.
+ val hash = StatefulOpClusteredDistribution(
+ Seq(BoundReference(0, schema, nullable = true)),
+ numShufflePartitions
+ ).createPartitioning(numShufflePartitions)
+
+ assert(sql/catalyst/src/test/scala/org/apache/spark/sql/RandomDataGenerator.scala
+ hash.isInstanceOf[HashPartitioning],
+ "StatefulOpClusteredDistribution should " +
+ "rely on HashPartitioning to ensure partitions remain the same for streaming " +
+ "stateful operators."
+ )
+ val rows = getRandomRows(schema, numRows, random)
+ val partitions = getPartitionId(rows, hash.asInstanceOf[HashPartitioning])
+ saveRowsAndPartIds(rows, partitions, os)
+ }
+ os.close()
+ } else {
+ val schemas = readSchemas()
+ val is = new DataInputStream(
+ codec.compressedInputStream(new FileInputStream(rowAndPartIdFile))
+ )
+ schemas.foreach { schema =>
+ val hash = StatefulOpClusteredDistribution(
+ Seq(BoundReference(0, schema, nullable = true)),
+ numShufflePartitions
+ ).createPartitioning(numShufflePartitions)
+
+ assert(
+ hash.isInstanceOf[HashPartitioning],
+ "StatefulOpClusteredDistribution should " +
+ "rely on HashPartitioning to ensure partitions remain the same for streaming " +
+ "stateful operators."
+ )
+
+ val (rows, expectedPartitions) = readRowsAndPartIds(is)
+ val partitions = getPartitionId(rows, hash.asInstanceOf[HashPartitioning])
+ assert(
+ partitions === expectedPartitions,
Review Comment:
Please consider what will be logged when this assertion fails, and whether the information is sufficient to investigate the issue. If there are only a few rows being shown as mismatched partitions, what would be the better format of assertion error message to understand the issue?
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,219 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ private val goldenFileURI =
+ this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ private val schemaFileName = "randomSchemas" // files for storing random input schemas
+ private val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ private val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
Review Comment:
While I don't think we have to add all types (size restriction), let's add importantly used types. StructType with two TimestampTypes consists a type of time window, so `TimestampType` would be a good one to test. `LongType` is also a good one.
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,219 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ private val goldenFileURI =
+ this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ private val schemaFileName = "randomSchemas" // files for storing random input schemas
+ private val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ private val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ private val numSchemaTypes = 1
+ private val maxNumFields = 20
+
+ // Configs for shuffle
+ private val numRows = 10000
+ private val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach { schema =>
+ writer.write(schema.toDDL)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try {
+ source
+ .getLines()
+ .map { ddl =>
+ StructType.fromDDL(ddl)
+ }
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ // Save the total number of rows
+ os.writeInt(rows.length)
+ // Save all rows
+ rows.foreach { row =>
+ // Save the row's total number of bytes
+ val rowBytes = row.getBytes()
+ // Save the row's actual bytes
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)
+ }
+ // Save all partIds, which should be in the same order as rows
+ partIds.foreach { id =>
+ os.writeInt(id)
+ }
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map { _ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(
+ schema,
+ rand = new Random(rand.nextInt())
+ )
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+
+ // UnsafeProjection returns the same UnsafeRow instance intentionally, so
+ // unless doing deep copy, the hash partitions below will evaluate on the
+ // same row thus return same value.
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(rand: Random): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ RandomDataGenerator.randomNestedSchema(rand, maxNumFields, variableTypes)
+ }
+ }
+
+ private def getPartitionId(rows: Seq[UnsafeRow], hash: HashPartitioning): Seq[Int] = {
+ val partIdExpr = hash.partitionIdExpression
+ rows.map { row =>
+ partIdExpr.eval(row).asInstanceOf[Int]
+ }
+ }
+
+ test("SPARK-47788: Ensure the same hash function is used across batches.") {
+ val rowAndPartIdFile = new File(goldenFileURI.getPath, rowAndPartIdFilename)
+
+ if (regenerateGoldenFiles) {
+ // To limit the golden file size under 10Mb, please set the final val MAX_STR_LEN: Int = 100
Review Comment:
Do I understand correctly that RandomDataGenerator.forType receives the parameters for max string length and max array size? We don't need to let others to fix the code of RandomDataGenerator. Let's just do that by ourselves via providing the parameters.
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,219 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
Review Comment:
Where do we actually use this? We actually do not do actual shuffle, and we seem to pass the number of shuffle partition as parameter for all necessary places.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788][SS] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "fanyue-xia (via GitHub)" <gi...@apache.org>.
fanyue-xia commented on code in PR #45971:
URL: https://github.com/apache/spark/pull/45971#discussion_r1565133720
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,219 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ private val goldenFileURI =
+ this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ private val schemaFileName = "randomSchemas" // files for storing random input schemas
+ private val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ private val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ private val numSchemaTypes = 1
+ private val maxNumFields = 20
+
+ // Configs for shuffle
+ private val numRows = 10000
+ private val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach { schema =>
+ writer.write(schema.toDDL)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try {
+ source
+ .getLines()
+ .map { ddl =>
+ StructType.fromDDL(ddl)
+ }
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ // Save the total number of rows
+ os.writeInt(rows.length)
+ // Save all rows
+ rows.foreach { row =>
+ // Save the row's total number of bytes
+ val rowBytes = row.getBytes()
+ // Save the row's actual bytes
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)
+ }
+ // Save all partIds, which should be in the same order as rows
+ partIds.foreach { id =>
+ os.writeInt(id)
+ }
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map { _ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(
+ schema,
+ rand = new Random(rand.nextInt())
+ )
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+
+ // UnsafeProjection returns the same UnsafeRow instance intentionally, so
+ // unless doing deep copy, the hash partitions below will evaluate on the
+ // same row thus return same value.
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(rand: Random): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ RandomDataGenerator.randomNestedSchema(rand, maxNumFields, variableTypes)
+ }
+ }
+
+ private def getPartitionId(rows: Seq[UnsafeRow], hash: HashPartitioning): Seq[Int] = {
+ val partIdExpr = hash.partitionIdExpression
+ rows.map { row =>
+ partIdExpr.eval(row).asInstanceOf[Int]
+ }
+ }
+
+ test("SPARK-47788: Ensure the same hash function is used across batches.") {
+ val rowAndPartIdFile = new File(goldenFileURI.getPath, rowAndPartIdFilename)
+
+ if (regenerateGoldenFiles) {
+ // To limit the golden file size under 10Mb, please set the final val MAX_STR_LEN: Int = 100
+ // and final val MAX_ARR_SIZE: Int = 4 in org.apache.spark.sql.RandomDataGenerator
+
+ val seed = Random.nextInt()
+ val random = new Random(seed)
Review Comment:
Thanks for the suggestion.
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,219 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ private val goldenFileURI =
+ this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ private val schemaFileName = "randomSchemas" // files for storing random input schemas
+ private val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ private val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ private val numSchemaTypes = 1
+ private val maxNumFields = 20
+
+ // Configs for shuffle
+ private val numRows = 10000
+ private val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach { schema =>
+ writer.write(schema.toDDL)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try {
+ source
+ .getLines()
+ .map { ddl =>
+ StructType.fromDDL(ddl)
+ }
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ // Save the total number of rows
+ os.writeInt(rows.length)
+ // Save all rows
+ rows.foreach { row =>
+ // Save the row's total number of bytes
+ val rowBytes = row.getBytes()
+ // Save the row's actual bytes
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)
+ }
+ // Save all partIds, which should be in the same order as rows
+ partIds.foreach { id =>
+ os.writeInt(id)
+ }
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map { _ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(
+ schema,
+ rand = new Random(rand.nextInt())
+ )
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+
+ // UnsafeProjection returns the same UnsafeRow instance intentionally, so
+ // unless doing deep copy, the hash partitions below will evaluate on the
+ // same row thus return same value.
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(rand: Random): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ RandomDataGenerator.randomNestedSchema(rand, maxNumFields, variableTypes)
+ }
+ }
+
+ private def getPartitionId(rows: Seq[UnsafeRow], hash: HashPartitioning): Seq[Int] = {
+ val partIdExpr = hash.partitionIdExpression
+ rows.map { row =>
+ partIdExpr.eval(row).asInstanceOf[Int]
+ }
+ }
+
+ test("SPARK-47788: Ensure the same hash function is used across batches.") {
Review Comment:
Updated.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788] [Structured Streaming] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "fanyue-xia (via GitHub)" <gi...@apache.org>.
fanyue-xia commented on code in PR #45971:
URL: https://github.com/apache/spark/pull/45971#discussion_r1562948068
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryHashPartitionVerifySuite.scala:
##########
@@ -0,0 +1,199 @@
+/*
+ * 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.spark.sql.streaming
+
+import java.io.{BufferedWriter, DataInputStream, DataOutputStream, File, FileInputStream, FileOutputStream, FileWriter}
+
+import scala.io.Source
+import scala.util.Random
+
+import com.google.common.io.ByteStreams
+
+import org.apache.spark.SparkConf
+import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.StreamTest
+import org.apache.spark.sql.types.{BinaryType, DataType, DoubleType, FloatType, IntegerType, StringType, StructType}
+
+
+class StreamingQueryHashPartitionVerifySuite extends StreamTest {
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.SHUFFLE_PARTITIONS.key, numShufflePartitions.toString)
+
+ // Configs for golden file
+ val goldenFileURI = this.getClass.getResource("/structured-streaming/partition-tests/").toURI
+
+ val schemaFileName = "randomSchemas" // files for storing random input schemas
+ val rowAndPartIdFilename =
+ "rowsAndPartIds" // files for storing random input rows and resulting partition ids
+ val codec = CompressionCodec.createCodec(
+ sparkConf) // Used for compressing output to rowAndPartId file
+
+ // Configs for random schema generation
+ private val variableTypes = Seq(IntegerType, DoubleType, FloatType, BinaryType, StringType)
+ val numSchemaTypes = 5
+ val maxNumFields = 10
+
+ // Configs for shuffle
+ val numRows = 10000
+ val numShufflePartitions = 100
+ private def saveSchemas(schemas: Seq[StructType]) = {
+ val writer = new BufferedWriter(new FileWriter(new File(goldenFileURI.getPath, schemaFileName)))
+ schemas.foreach{ schema =>
+ writer.write(schema.json)
+ writer.newLine()
+ }
+ writer.close()
+ }
+
+ private def readSchemas(): Seq[StructType] = {
+ val source = Source.fromFile(new File(goldenFileURI.getPath, schemaFileName))
+ try{
+ source.getLines().map{json => DataType.fromJson(json).asInstanceOf[StructType]}
+ .toArray // Avoid Stream lazy materialization
+ .toSeq
+ } finally
+ source.close()
+ }
+
+ private def saveRowsAndPartIds(rows: Seq[UnsafeRow], partIds: Seq[Int], os: DataOutputStream) = {
+ os.writeInt(rows.length)
+ rows.foreach { row =>
+ val rowBytes = row.getBytes()
+ os.writeInt(rowBytes.size)
+ os.write(rowBytes)}
+ partIds.foreach {id => os.writeInt(id)}
+ }
+
+ private def readRowsAndPartIds(is: DataInputStream): (Seq[UnsafeRow], Seq[Int]) = {
+ val numRows = is.readInt()
+ val rows = (1 to numRows).map{_ =>
+ val rowSize = is.readInt()
+ val rowBuffer = new Array[Byte](rowSize)
+ ByteStreams.readFully(is, rowBuffer, 0, rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(rowBuffer, rowSize)
+ row
+ }
+ val partIds = (1 to numRows).map(_ => is.readInt()).toArray.toSeq
+ (rows, partIds)
+ }
+
+ private def getRandomRows(schema: StructType, numRows: Int, rand: Random): Seq[UnsafeRow] = {
+ val generator = RandomDataGenerator
+ .forType(schema, rand = new Random(rand.nextInt(RandomDataGenerator.MAX_STR_LEN)))
+ .get
+
+ // Create the converters needed to convert from external row to internal
+ // row and to UnsafeRows. Projection itself costs a lot on initialization
+ // (codegen and compile), so initialize it once.
+ val internalConverter = CatalystTypeConverters.createToCatalystConverter(schema)
+ val unsafeConverter = UnsafeProjection.create(Array(schema).asInstanceOf[Array[DataType]])
+
+ (1 to numRows).map { _ =>
+ val row = generator().asInstanceOf[Row]
+
+ val internalRow = new GenericInternalRow(1)
+ internalRow.update(0, internalConverter(row).asInstanceOf[GenericInternalRow])
+ val unsafeRow = unsafeConverter.apply(internalRow)
+ unsafeRow.copy()
+ }
+ }
+
+ private def getRandomSchemas(numSchemaTypes: Int, rand: Random,
+ maxNumFields: Int): Seq[StructType] = {
+ (1 to numSchemaTypes).map { _ =>
+ val numField = rand.nextInt(maxNumFields)
+ RandomDataGenerator.randomSchema(rand, Math.max(1, numField), variableTypes)
Review Comment:
This is to prevent the schema contains 0 field
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47788] [Structured Streaming] Ensure the same hash partitioning for streaming stateful ops [spark]
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on PR #45971:
URL: https://github.com/apache/spark/pull/45971#issuecomment-2052419964
> the seed might behave differently across runs/on different machines
Ah I see, this indeed makes sense.
In this case, I think we should fix the generator of rows. It's okay to sacrifice randomness of rows here.
We can have a dedicated row generation function, depending on the input type, this function just return a fixed return (e.g. if input is int, just return 233, if input is byte, just return 0xdeadbeef)
Giving up randomness of the row should still get the job done. The way the hash is computed is sth like `hash(field 1, hash (field 2, seed...)...)`, and this part hasn't been touched likely since the beginning.
https://github.com/apache/spark/blob/6ee662c28ffb0deb70f08a971f9c1869288d39ba/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/hash.scala#L289-L298
This should not be changed, even if it is changed, it makes less sense to be changed to a hash function that hash (1, 2) and (2, 1) to be the same bucket.
Then as long as we increase the max number of field and the number of schemas (now this can be a fairly high number, also pay attention to the test run time), it should behave similar as having a large number of random generated rows.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org