You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "zifeif2 (via GitHub)" <gi...@apache.org> on 2025/11/18 01:05:00 UTC
[PR] [WIP] [SPARK-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
zifeif2 opened a new pull request, #53104:
URL: https://github.com/apache/spark/pull/53104
<!--
Thanks for sending a pull request! Here are some tips for you:
1. If this is your first time, please read our contributor guidelines: https://spark.apache.org/contributing.html
2. Ensure you have added or run the appropriate tests for your PR: https://spark.apache.org/developer-tools.html
3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][SPARK-XXXX] Your PR title ...'.
4. Be sure to keep the PR description updated to reflect all changes.
5. Please write your PR title to summarize what this PR proposes.
6. If possible, provide a concise example to reproduce the issue for a faster review.
7. If you want to add a new configuration, please read the guideline first for naming configurations in
'core/src/main/scala/org/apache/spark/internal/config/ConfigEntry.scala'.
8. If you want to add or modify an error type or message, please read the guideline first in
'common/utils/src/main/resources/error/README.md'.
-->
## WIP : will add test on HDFSStateStore later
### What changes were proposed in this pull request?
<!--
Please clarify what changes you are proposing. The purpose of this section is to outline the changes and how this PR fixes the issue.
If possible, please consider writing useful notes for better and faster reviews in your PR. See the examples below.
1. If you refactor some codes with changing classes, showing the class hierarchy will help reviewers.
2. If you fix some SQL features, you can provide some references of other DBMSes.
3. If there is design documentation, please add the link.
4. If there is a discussion in the mailing list, please add the link.
-->
Introducing a new StatePartitionReader - StatePartitionReaderAllColumnFamilies to support offline repartition.
StatePartitionReaderAllColumnFamilies is invoked when user specify option `readAllColumnFamilies` to true.
We have the StateDataSource Reader, which allows customers to read the rows in an operator state store using the DataFrame API, just like they read a normal table. But it currently only supports reading one column family in the state store at a time.
We would introduce a change to allow reading all the state rows in all the column families, so that we can repartition them at once. This would allow us to read the entire state store, repartition the rows, and then save the new repartition state rows to the cloud. This also has a perf impact, since we don’t have to read each column family separately. We would read the state based on the last committed batch version.
Since each column family can have a different schema, the DataFrame we will return will treat the key and value row as bytes -
- partition_key (string)
- key_bytes (binary)
- value_bytes (binary)
- column_family_name (string)
### Why are the changes needed?
<!--
Please clarify why the changes are needed. For instance,
1. If you propose a new API, clarify the use case for a new API.
2. If you fix a bug, you can clarify why it is a bug.
-->
See above
### Does this PR introduce _any_ user-facing change?
<!--
Note that it means *any* user-facing change including all aspects such as new features, bug fixes, or other behavior changes. Documentation-only updates are not considered user-facing changes.
If yes, please clarify the previous behavior and the change this PR proposes - provide the console output, description and/or an example to show the behavior difference if possible.
If possible, please also clarify if this is a user-facing change compared to the released Spark versions or within the unreleased branches such as master.
If no, write 'No'.
-->
No
### How was this patch tested?
<!--
If tests were added, say they were added here. Please make sure to add some test cases that check the changes thoroughly including negative and positive cases if possible.
If it was tested in a way different from regular unit tests, please clarify how you tested step by step, ideally copy and paste-able, so that other reviewers can test and check, and descendants can verify in the future.
If tests were not added, please describe why they were not added and/or why it was difficult to add.
If benchmark tests were added, please run the benchmarks in GitHub Actions for the consistent environment, and the instructions could accord to: https://spark.apache.org/developer-tools.html#github-workflow-benchmarks.
-->
See unit test. It not only verify the schema, but also validate the data are serialized to bytes correctly by comparing them against the normal queried data frame
### Was this patch authored or co-authored using generative AI tooling?
<!--
If generative AI tooling has been used in the process of authoring this patch, please include the
phrase: 'Generated-by: ' followed by the name of the tool and its version.
If no, write 'No'.
Please refer to the [ASF Generative Tooling Guidance](https://www.apache.org/legal/generative-tooling.html) for details.
-->
Yes. haiku, sonnet.
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "zifeif2 (via GitHub)" <gi...@apache.org>.
zifeif2 commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2557243370
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionAllColumnFamiliesReaderSuite.scala:
##########
@@ -0,0 +1,449 @@
+/*
+ * 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.execution.datasources.v2.state
+
+import java.nio.ByteOrder
+import java.util.Arrays
+
+import org.apache.spark.SparkUnsupportedOperationException
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.UnsafeProjection
+import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.{HDFSBackedStateStoreProvider, RocksDBStateStoreProvider}
+import org.apache.spark.sql.functions.{count, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.OutputMode
+import org.apache.spark.sql.types.{IntegerType, LongType, NullType, StructField, StructType}
+
+/**
+ * Note: This extends StateDataSourceTestBase to access
+ * helper methods like runDropDuplicatesQuery without inheriting all predefined tests.
+ */
+class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase {
+
+ import testImplicits._
+
+ override def beforeAll(): Unit = {
+ super.beforeAll()
+ spark.conf.set(SQLConf.STATE_STORE_PROVIDER_CLASS.key,
+ classOf[RocksDBStateStoreProvider].getName)
+ }
+
+ private def getNormalReadDf(checkpointDir: String): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .load()
+ .selectExpr("partition_id", "key", "value")
+ }
+
+ private def getBytesReadDf(checkpointDir: String): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .option(StateSourceOptions.INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES, "true")
+ .load()
+ .selectExpr("partition_key", "key_bytes", "value_bytes", "column_family_name")
+ }
+
+ /**
+ * Validates the schema and column families of the bytes read DataFrame.
+ */
+ private def validateBytesReadDfSchema(df: DataFrame): Unit = {
+ // Verify schema
+ val schema = df.schema
+ assert(schema.fieldNames === Array(
+ "partition_key", "key_bytes", "value_bytes", "column_family_name"))
+ assert(schema("partition_key").dataType.typeName === "struct")
+ assert(schema("key_bytes").dataType.typeName === "binary")
+ assert(schema("value_bytes").dataType.typeName === "binary")
+ assert(schema("column_family_name").dataType.typeName === "string")
+ }
+
+ /**
+ * Compares normal read data with bytes read data for a specific column family.
+ * Converts normal rows to bytes then compares with bytes read.
+ */
+ private def compareNormalAndBytesData(
+ normalDf: Array[Row],
+ bytesDf: Array[Row],
+ columnFamily: String,
+ keySchema: StructType,
+ valueSchema: StructType): Unit = {
+
+ // Filter bytes data for the specified column family and extract raw bytes directly
+ val filteredBytesData = bytesDf.filter { row =>
+ row.getString(3) == columnFamily
+ }
+
+ // Verify same number of rows
+ assert(filteredBytesData.length == normalDf.length,
+ s"Row count mismatch for column family '$columnFamily': " +
+ s"normal read has ${normalDf.length} rows, " +
+ s"bytes read has ${filteredBytesData.length} rows")
+
+ // Create projections to convert Row to UnsafeRow bytes
+ val keyProjection = UnsafeProjection.create(keySchema)
+ val valueProjection = UnsafeProjection.create(valueSchema)
+
+ // Create converters to convert external Row types to internal Catalyst types
+ val keyConverter = CatalystTypeConverters.createToCatalystConverter(keySchema)
+ val valueConverter = CatalystTypeConverters.createToCatalystConverter(valueSchema)
+
+ // Convert normal data to bytes
+ val normalAsBytes = normalDf.map { row =>
+ val key = row.getStruct(1)
+ val value = if (row.isNullAt(2)) null else row.getStruct(2)
+
+ // Convert key to InternalRow, then to UnsafeRow, then get bytes
+ val keyInternalRow = keyConverter(key).asInstanceOf[InternalRow]
+ val keyUnsafeRow = keyProjection(keyInternalRow)
+ // IMPORTANT: Must clone the bytes array since getBytes() returns a reference
+ // that may be overwritten by subsequent UnsafeRow operations
+ val keyBytes = keyUnsafeRow.getBytes.clone()
+
+ // Convert value to bytes
+ val valueBytes = if (value == null) {
+ Array.empty[Byte]
+ } else {
+ val valueInternalRow = valueConverter(value).asInstanceOf[InternalRow]
+ val valueUnsafeRow = valueProjection(valueInternalRow)
+ // IMPORTANT: Must clone the bytes array since getBytes() returns a reference
+ // that may be overwritten by subsequent UnsafeRow operations
+ valueUnsafeRow.getBytes.clone()
+ }
+
+ (keyBytes, valueBytes)
+ }
+
+ // Extract raw bytes from bytes read data (no deserialization/reserialization)
+ val bytesAsBytes = filteredBytesData.map { row =>
+ val keyBytes = row.getAs[Array[Byte]](1)
+ val valueBytes = row.getAs[Array[Byte]](2)
+ (keyBytes, valueBytes)
+ }
+
+ // Sort both for comparison (since Set equality doesn't work well with byte arrays)
+ val normalSorted = normalAsBytes.sortBy(x => (x._1.mkString(","), x._2.mkString(",")))
+ val bytesSorted = bytesAsBytes.sortBy(x => (x._1.mkString(","), x._2.mkString(",")))
+
+ assert(normalSorted.length == bytesSorted.length,
+ s"Size mismatch: normal has ${normalSorted.length}, bytes has ${bytesSorted.length}")
+
+ // Compare each pair
+ normalSorted.zip(bytesSorted).zipWithIndex.foreach {
+ case (((normalKey, normalValue), (bytesKey, bytesValue)), idx) =>
+ assert(Arrays.equals(normalKey, bytesKey),
+ s"Key mismatch at index $idx:\n" +
+ s" Normal: ${normalKey.mkString("[", ",", "]")}\n" +
+ s" Bytes: ${bytesKey.mkString("[", ",", "]")}")
+ assert(Arrays.equals(normalValue, bytesValue),
+ s"Value mismatch at index $idx:\n" +
+ s" Normal: ${normalValue.mkString("[", ",", "]")}\n" +
+ s" Bytes: ${bytesValue.mkString("[", ",", "]")}")
+ }
+ }
+
+ // Run all tests with both changelog checkpointing enabled and disabled
+ Seq(true, false).foreach { changelogCheckpointingEnabled =>
+ val testSuffix = if (changelogCheckpointingEnabled) {
+ "with changelog checkpointing"
+ } else {
+ "without changelog checkpointing"
+ }
+
+ def testWithChangelogConfig(testName: String)(testFun: => Unit): Unit = {
+ test(s"$testName ($testSuffix)") {
+ withSQLConf(
+ "spark.sql.streaming.stateStore.rocksdb.changelogCheckpointing.enabled" ->
+ changelogCheckpointingEnabled.toString) {
+ testFun
+ }
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: simple aggregation state ver 1") {
+ withSQLConf(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key -> "1") {
+ withTempDir { tempDir =>
+ runLargeDataStreamingAggregationQuery(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(StructField("groupKey", IntegerType, nullable = false)))
+ // State version 1 includes key columns in the value
+ val valueSchema = StructType(Array(
+ StructField("groupKey", IntegerType, nullable = false),
+ StructField("count", LongType, nullable = false),
+ StructField("sum", LongType, nullable = false),
+ StructField("max", IntegerType, nullable = false),
+ StructField("min", IntegerType, nullable = false)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath)
+
+ validateBytesReadDfSchema(bytesDf)
+ compareNormalAndBytesData(normalData, bytesDf.collect(), "default", keySchema, valueSchema)
+ }
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: simple aggregation state ver 2") {
+ withSQLConf(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key -> "2") {
+ withTempDir { tempDir =>
+ runLargeDataStreamingAggregationQuery(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(StructField("groupKey", IntegerType, nullable = false)))
+ val valueSchema = StructType(Array(
+ StructField("count", LongType, nullable = false),
+ StructField("sum", LongType, nullable = false),
+ StructField("max", IntegerType, nullable = false),
+ StructField("min", IntegerType, nullable = false)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath).collect()
+
+ compareNormalAndBytesData(normalData, bytesDf, "default", keySchema, valueSchema)
+ }
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: composite key aggregation state ver 1") {
+ withSQLConf(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key -> "1") {
+ withTempDir { tempDir =>
+ runCompositeKeyStreamingAggregationQuery(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(
+ StructField("groupKey", IntegerType, nullable = false),
+ StructField("fruit", org.apache.spark.sql.types.StringType, nullable = true)
+ ))
+ // State version 1 includes key columns in the value
+ val valueSchema = StructType(Array(
+ StructField("groupKey", IntegerType, nullable = false),
+ StructField("fruit", org.apache.spark.sql.types.StringType, nullable = true),
+ StructField("count", LongType, nullable = false),
+ StructField("sum", LongType, nullable = false),
+ StructField("max", IntegerType, nullable = false),
+ StructField("min", IntegerType, nullable = false)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath).collect()
+
+ compareNormalAndBytesData(normalData, bytesDf, "default", keySchema, valueSchema)
+ }
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: composite key aggregation state ver 2") {
+ withSQLConf(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key -> "2") {
+ withTempDir { tempDir =>
+ runCompositeKeyStreamingAggregationQuery(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(
+ StructField("groupKey", IntegerType, nullable = false),
+ StructField("fruit", org.apache.spark.sql.types.StringType, nullable = true)
+ ))
+ val valueSchema = StructType(Array(
+ StructField("count", LongType, nullable = false),
+ StructField("sum", LongType, nullable = false),
+ StructField("max", IntegerType, nullable = false),
+ StructField("min", IntegerType, nullable = false)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath).collect()
+
+ compareNormalAndBytesData(normalData, bytesDf, "default", keySchema, valueSchema)
+ }
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: dropDuplicates validation") {
+ withTempDir { tempDir =>
+ runDropDuplicatesQuery(tempDir.getAbsolutePath)
Review Comment:
Will fix indentation!
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "zifeif2 (via GitHub)" <gi...@apache.org>.
zifeif2 commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2557223014
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionAllColumnFamiliesReaderSuite.scala:
##########
@@ -0,0 +1,449 @@
+/*
+ * 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.execution.datasources.v2.state
+
+import java.nio.ByteOrder
+import java.util.Arrays
+
+import org.apache.spark.SparkUnsupportedOperationException
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.UnsafeProjection
+import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.{HDFSBackedStateStoreProvider, RocksDBStateStoreProvider}
+import org.apache.spark.sql.functions.{count, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.OutputMode
+import org.apache.spark.sql.types.{IntegerType, LongType, NullType, StructField, StructType}
+
+/**
+ * Note: This extends StateDataSourceTestBase to access
+ * helper methods like runDropDuplicatesQuery without inheriting all predefined tests.
+ */
+class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase {
+
+ import testImplicits._
+
+ override def beforeAll(): Unit = {
+ super.beforeAll()
+ spark.conf.set(SQLConf.STATE_STORE_PROVIDER_CLASS.key,
+ classOf[RocksDBStateStoreProvider].getName)
+ }
+
+ private def getNormalReadDf(checkpointDir: String): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .load()
+ .selectExpr("partition_id", "key", "value")
+ }
+
+ private def getBytesReadDf(checkpointDir: String): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .option(StateSourceOptions.INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES, "true")
+ .load()
+ .selectExpr("partition_key", "key_bytes", "value_bytes", "column_family_name")
+ }
+
+ /**
+ * Validates the schema and column families of the bytes read DataFrame.
+ */
+ private def validateBytesReadDfSchema(df: DataFrame): Unit = {
+ // Verify schema
+ val schema = df.schema
+ assert(schema.fieldNames === Array(
+ "partition_key", "key_bytes", "value_bytes", "column_family_name"))
+ assert(schema("partition_key").dataType.typeName === "struct")
+ assert(schema("key_bytes").dataType.typeName === "binary")
+ assert(schema("value_bytes").dataType.typeName === "binary")
+ assert(schema("column_family_name").dataType.typeName === "string")
+ }
+
+ /**
+ * Compares normal read data with bytes read data for a specific column family.
+ * Converts normal rows to bytes then compares with bytes read.
+ */
+ private def compareNormalAndBytesData(
+ normalDf: Array[Row],
+ bytesDf: Array[Row],
+ columnFamily: String,
+ keySchema: StructType,
+ valueSchema: StructType): Unit = {
+
+ // Filter bytes data for the specified column family and extract raw bytes directly
+ val filteredBytesData = bytesDf.filter { row =>
+ row.getString(3) == columnFamily
+ }
+
+ // Verify same number of rows
+ assert(filteredBytesData.length == normalDf.length,
+ s"Row count mismatch for column family '$columnFamily': " +
+ s"normal read has ${normalDf.length} rows, " +
+ s"bytes read has ${filteredBytesData.length} rows")
+
+ // Create projections to convert Row to UnsafeRow bytes
+ val keyProjection = UnsafeProjection.create(keySchema)
+ val valueProjection = UnsafeProjection.create(valueSchema)
+
+ // Create converters to convert external Row types to internal Catalyst types
+ val keyConverter = CatalystTypeConverters.createToCatalystConverter(keySchema)
+ val valueConverter = CatalystTypeConverters.createToCatalystConverter(valueSchema)
+
+ // Convert normal data to bytes
+ val normalAsBytes = normalDf.map { row =>
+ val key = row.getStruct(1)
+ val value = if (row.isNullAt(2)) null else row.getStruct(2)
+
+ // Convert key to InternalRow, then to UnsafeRow, then get bytes
+ val keyInternalRow = keyConverter(key).asInstanceOf[InternalRow]
+ val keyUnsafeRow = keyProjection(keyInternalRow)
+ // IMPORTANT: Must clone the bytes array since getBytes() returns a reference
+ // that may be overwritten by subsequent UnsafeRow operations
+ val keyBytes = keyUnsafeRow.getBytes.clone()
+
+ // Convert value to bytes
+ val valueBytes = if (value == null) {
+ Array.empty[Byte]
+ } else {
+ val valueInternalRow = valueConverter(value).asInstanceOf[InternalRow]
+ val valueUnsafeRow = valueProjection(valueInternalRow)
+ // IMPORTANT: Must clone the bytes array since getBytes() returns a reference
+ // that may be overwritten by subsequent UnsafeRow operations
+ valueUnsafeRow.getBytes.clone()
+ }
+
+ (keyBytes, valueBytes)
+ }
+
+ // Extract raw bytes from bytes read data (no deserialization/reserialization)
+ val bytesAsBytes = filteredBytesData.map { row =>
+ val keyBytes = row.getAs[Array[Byte]](1)
+ val valueBytes = row.getAs[Array[Byte]](2)
+ (keyBytes, valueBytes)
+ }
+
+ // Sort both for comparison (since Set equality doesn't work well with byte arrays)
+ val normalSorted = normalAsBytes.sortBy(x => (x._1.mkString(","), x._2.mkString(",")))
+ val bytesSorted = bytesAsBytes.sortBy(x => (x._1.mkString(","), x._2.mkString(",")))
+
+ assert(normalSorted.length == bytesSorted.length,
+ s"Size mismatch: normal has ${normalSorted.length}, bytes has ${bytesSorted.length}")
+
+ // Compare each pair
+ normalSorted.zip(bytesSorted).zipWithIndex.foreach {
+ case (((normalKey, normalValue), (bytesKey, bytesValue)), idx) =>
+ assert(Arrays.equals(normalKey, bytesKey),
+ s"Key mismatch at index $idx:\n" +
+ s" Normal: ${normalKey.mkString("[", ",", "]")}\n" +
+ s" Bytes: ${bytesKey.mkString("[", ",", "]")}")
+ assert(Arrays.equals(normalValue, bytesValue),
+ s"Value mismatch at index $idx:\n" +
+ s" Normal: ${normalValue.mkString("[", ",", "]")}\n" +
+ s" Bytes: ${bytesValue.mkString("[", ",", "]")}")
+ }
+ }
+
+ // Run all tests with both changelog checkpointing enabled and disabled
+ Seq(true, false).foreach { changelogCheckpointingEnabled =>
+ val testSuffix = if (changelogCheckpointingEnabled) {
+ "with changelog checkpointing"
+ } else {
+ "without changelog checkpointing"
+ }
+
+ def testWithChangelogConfig(testName: String)(testFun: => Unit): Unit = {
+ test(s"$testName ($testSuffix)") {
+ withSQLConf(
+ "spark.sql.streaming.stateStore.rocksdb.changelogCheckpointing.enabled" ->
+ changelogCheckpointingEnabled.toString) {
+ testFun
+ }
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: simple aggregation state ver 1") {
+ withSQLConf(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key -> "1") {
+ withTempDir { tempDir =>
+ runLargeDataStreamingAggregationQuery(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(StructField("groupKey", IntegerType, nullable = false)))
+ // State version 1 includes key columns in the value
+ val valueSchema = StructType(Array(
+ StructField("groupKey", IntegerType, nullable = false),
+ StructField("count", LongType, nullable = false),
+ StructField("sum", LongType, nullable = false),
+ StructField("max", IntegerType, nullable = false),
+ StructField("min", IntegerType, nullable = false)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath)
+
+ validateBytesReadDfSchema(bytesDf)
+ compareNormalAndBytesData(normalData, bytesDf.collect(), "default", keySchema, valueSchema)
+ }
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: simple aggregation state ver 2") {
+ withSQLConf(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key -> "2") {
+ withTempDir { tempDir =>
+ runLargeDataStreamingAggregationQuery(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(StructField("groupKey", IntegerType, nullable = false)))
+ val valueSchema = StructType(Array(
+ StructField("count", LongType, nullable = false),
+ StructField("sum", LongType, nullable = false),
+ StructField("max", IntegerType, nullable = false),
+ StructField("min", IntegerType, nullable = false)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath).collect()
+
+ compareNormalAndBytesData(normalData, bytesDf, "default", keySchema, valueSchema)
+ }
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: composite key aggregation state ver 1") {
+ withSQLConf(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key -> "1") {
+ withTempDir { tempDir =>
+ runCompositeKeyStreamingAggregationQuery(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(
+ StructField("groupKey", IntegerType, nullable = false),
+ StructField("fruit", org.apache.spark.sql.types.StringType, nullable = true)
+ ))
+ // State version 1 includes key columns in the value
+ val valueSchema = StructType(Array(
+ StructField("groupKey", IntegerType, nullable = false),
+ StructField("fruit", org.apache.spark.sql.types.StringType, nullable = true),
+ StructField("count", LongType, nullable = false),
+ StructField("sum", LongType, nullable = false),
+ StructField("max", IntegerType, nullable = false),
+ StructField("min", IntegerType, nullable = false)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath).collect()
+
+ compareNormalAndBytesData(normalData, bytesDf, "default", keySchema, valueSchema)
+ }
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: composite key aggregation state ver 2") {
+ withSQLConf(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key -> "2") {
+ withTempDir { tempDir =>
+ runCompositeKeyStreamingAggregationQuery(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(
+ StructField("groupKey", IntegerType, nullable = false),
+ StructField("fruit", org.apache.spark.sql.types.StringType, nullable = true)
+ ))
+ val valueSchema = StructType(Array(
+ StructField("count", LongType, nullable = false),
+ StructField("sum", LongType, nullable = false),
+ StructField("max", IntegerType, nullable = false),
+ StructField("min", IntegerType, nullable = false)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath).collect()
+
+ compareNormalAndBytesData(normalData, bytesDf, "default", keySchema, valueSchema)
+ }
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: dropDuplicates validation") {
+ withTempDir { tempDir =>
+ runDropDuplicatesQuery(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(
+ StructField("value", IntegerType, nullable = false),
+ StructField("eventTime", org.apache.spark.sql.types.TimestampType)
+ ))
+ val valueSchema = StructType(Array(
+ StructField("__dummy__", NullType, nullable = true)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath).collect()
+
+ compareNormalAndBytesData(normalData, bytesDf, "default", keySchema, valueSchema)
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: dropDuplicates with column specified") {
+ withTempDir { tempDir =>
+ runDropDuplicatesQueryWithColumnSpecified(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(
+ StructField("col1", org.apache.spark.sql.types.StringType, nullable = true)
+ ))
+ val valueSchema = StructType(Array(
+ StructField("__dummy__", NullType, nullable = true)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath).collect()
+
+ compareNormalAndBytesData(normalData, bytesDf, "default", keySchema, valueSchema)
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: dropDuplicatesWithinWatermark") {
+ withTempDir { tempDir =>
+ runDropDuplicatesWithinWatermarkQuery(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(
+ StructField("_1", org.apache.spark.sql.types.StringType, nullable = true)
+ ))
+ val valueSchema = StructType(Array(
+ StructField("expiresAtMicros", LongType, nullable = false)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath).collect()
+
+ compareNormalAndBytesData(normalData, bytesDf, "default", keySchema, valueSchema)
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: session window aggregation") {
+ withTempDir { tempDir =>
+ runSessionWindowAggregationQuery(tempDir.getAbsolutePath)
Review Comment:
will fix indentation
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "zifeif2 (via GitHub)" <gi...@apache.org>.
zifeif2 commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2579588370
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OfflineStateRepartitionErrors.scala:
##########
@@ -201,3 +207,11 @@ class StateRepartitionUnsupportedOffsetSeqVersionError(
checkpointLocation,
subClass = "UNSUPPORTED_OFFSET_SEQ_VERSION",
messageParameters = Map("version" -> version.toString))
+
+class StateRepartitionUnsupportedProviderError(
Review Comment:
<img width="554" height="311" alt="Screenshot 2025-12-01 at 9 57 46 PM" src="https://github.com/user-attachments/assets/a60cb2de-15cf-44bc-8a54-9fda4ed7271f" />
I asked Claude and still couldn't figure out what's the issue with the indentation 😅 , could you specify? Thanks!
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2582762607
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala:
##########
@@ -407,6 +418,7 @@ object StateSourceOptions extends DataSourceOptions {
val STATE_VAR_NAME = newOption("stateVarName")
val READ_REGISTERED_TIMERS = newOption("readRegisteredTimers")
val FLATTEN_COLLECTION_TYPES = newOption("flattenCollectionTypes")
+ val INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES = newOption("internalOnlyReadAllColumnFamilies")
Review Comment:
Why do we need to expose an option at the state data source ? We only need to expose an interface at the partition reader level 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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "micheal-o (via GitHub)" <gi...@apache.org>.
micheal-o commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2548006075
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala:
##########
@@ -65,6 +65,13 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
val sourceOptions = StateSourceOptions.modifySourceOptions(hadoopConf,
StateSourceOptions.apply(session, hadoopConf, properties))
val stateConf = buildStateStoreConf(sourceOptions.resolvedCpLocation, sourceOptions.batchId)
+ if (sourceOptions.internalOnlyReadAllColumnFamilies
+ && !stateConf.providerClass.contains("RocksDB")) {
Review Comment:
nit: Just compare the class. Look at how we do this in other places
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala:
##########
@@ -237,6 +240,44 @@ class StatePartitionReader(
}
}
+/**
+ * An implementation of [[StatePartitionReaderBase]] for reading all column families
+ * in binary format. This reader returns raw key and value bytes along with column family names.
Review Comment:
nit: Add that "We are returning key/value bytes because each column family can have different schema"
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala:
##########
@@ -76,6 +88,22 @@ object SchemaUtil {
row
}
+ /**
+ * Creates a unified row from raw key and value bytes.
+ * This is an alias for unifyStateRowPairAsBytes that takes individual byte arrays
+ * instead of a tuple for better readability.
+ */
+ def unifyStateRowPairAsRawBytes(
+ pair: (UnsafeRow, UnsafeRow),
+ colFamilyName: String): InternalRow = {
+ val row = new GenericInternalRow(6)
Review Comment:
This is 6, but you're only updating 4?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala:
##########
@@ -237,6 +240,44 @@ class StatePartitionReader(
}
}
+/**
+ * An implementation of [[StatePartitionReaderBase]] for reading all column families
+ * in binary format. This reader returns raw key and value bytes along with column family names.
+ */
+class StatePartitionReaderAllColumnFamilies(
Review Comment:
nit: `StatePartitionAllColumnFamiliesReader`
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReaderAllColumnFamiliesSuite.scala:
##########
@@ -0,0 +1,244 @@
+/*
+ * 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.execution.datasources.v2.state
+
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.expressions.{GenericRowWithSchema, UnsafeRow}
+import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.{HDFSBackedStateStoreProvider, RocksDBStateStoreProvider}
+import org.apache.spark.sql.functions.{count, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.OutputMode
+import org.apache.spark.sql.types.{IntegerType, LongType, StructField, StructType}
+import org.apache.spark.tags.SlowSQLTest
+import org.apache.spark.unsafe.Platform
+
+/**
+ * Test suite to verify StatePartitionReaderAllColumnFamilies functionality.
+ */
+@SlowSQLTest
+class StatePartitionReaderAllColumnFamiliesSuite extends StateDataSourceTestBase {
+
+ import testImplicits._
+
+ /**
+ * Returns a set of (partitionId, key, value) tuples from a normal state read.
+ */
+ private def getNormalReadData(checkpointDir: String): DataFrame = {
Review Comment:
nit: `getNormalReadDf`
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReaderAllColumnFamiliesSuite.scala:
##########
@@ -0,0 +1,244 @@
+/*
+ * 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.execution.datasources.v2.state
+
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.expressions.{GenericRowWithSchema, UnsafeRow}
+import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.{HDFSBackedStateStoreProvider, RocksDBStateStoreProvider}
+import org.apache.spark.sql.functions.{count, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.OutputMode
+import org.apache.spark.sql.types.{IntegerType, LongType, StructField, StructType}
+import org.apache.spark.tags.SlowSQLTest
+import org.apache.spark.unsafe.Platform
+
+/**
+ * Test suite to verify StatePartitionReaderAllColumnFamilies functionality.
+ */
+@SlowSQLTest
+class StatePartitionReaderAllColumnFamiliesSuite extends StateDataSourceTestBase {
+
+ import testImplicits._
+
+ /**
+ * Returns a set of (partitionId, key, value) tuples from a normal state read.
+ */
+ private def getNormalReadData(checkpointDir: String): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .load()
+ .selectExpr("partition_id", "key", "value")
+ }
+
+ /**
+ * Returns a DataFrame with raw bytes mode (INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES = true).
+ */
+ private def getBytesReadDf(checkpointDir: String): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .option(StateSourceOptions.INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES, "true")
+ .load()
+ }
+
+ /**
+ * Validates the schema and column families of the bytes read DataFrame.
+ */
+ private def validateBytesReadSchema(df: DataFrame): Unit = {
+ // Verify schema
+ val schema = df.schema
+ assert(schema.fieldNames === Array(
+ "partition_key", "key_bytes", "value_bytes", "column_family_name", "value", "key"))
+ assert(schema("partition_key").dataType.typeName === "struct")
+ assert(schema("key_bytes").dataType.typeName === "binary")
+ assert(schema("value_bytes").dataType.typeName === "binary")
+ assert(schema("column_family_name").dataType.typeName === "string")
+ }
+
+ private def parseBytesReadData(df: Array[Row], keyLength: Int, valueLength: Int)
+ : Set[(GenericRowWithSchema, UnsafeRow, UnsafeRow, String)] = {
+ df.map { row =>
+ val partitionKey = row.getAs[GenericRowWithSchema](0)
+ val keyBytes = row.getAs[Array[Byte]](1)
+ val valueBytes = row.getAs[Array[Byte]](2)
+ val columnFamily = row.getString(3)
+
+ // Deserialize key bytes to UnsafeRow
+ val keyRow = new UnsafeRow(keyLength)
+ keyRow.pointTo(
+ keyBytes,
+ Platform.BYTE_ARRAY_OFFSET,
+ keyBytes.length)
+
+ // Deserialize value bytes to UnsafeRow
+ val valueRow = new UnsafeRow(valueLength)
+ valueRow.pointTo(
+ valueBytes,
+ Platform.BYTE_ARRAY_OFFSET,
+ valueBytes.length)
+ (partitionKey, keyRow.copy(), valueRow.copy(), columnFamily)
+ }
+ .toSet
+ }
+
+ /**
+ * Compares normal read data with bytes read data for a specific column family.
+ */
+ private def compareNormalAndBytesData(
+ normalReadDf: DataFrame,
+ bytesReadDf: DataFrame,
+ columnFamily: String,
+ keySchema: StructType,
+ valueSchema: StructType): Unit = {
+ // Verify data
+ val bytesDf = bytesReadDf
+ .selectExpr("partition_key", "key_bytes", "value_bytes", "column_family_name")
+ .collect()
+ assert(bytesDf.length == 10,
+ s"Expected 10 rows but got: ${bytesDf.length}")
+
+ // Filter bytes data for the specified column family
+ val bytesData = parseBytesReadData(bytesDf, keySchema.length, valueSchema.length)
+ val filteredBytesData = bytesData.filter(_._4 == columnFamily)
+
+ // Apply the projection
+ // Convert to comparable format (extract field values)
+ val normalSet = normalReadDf.collect().map { row =>
+ val key = row.getStruct(1)
+ val value = row.getStruct(2)
+ val keyFields = (0 until key.length).map(i => key.get(i))
+ val valueFields = (0 until value.length).map(i => value.get(i))
+ (keyFields, valueFields)
+ }.toSet
+
+ val bytesSet = filteredBytesData.map { case (_, keyRow, valueRow, _) =>
+ val keyFields = (0 until keySchema.length).map(i =>
+ keyRow.get(i, keySchema(i).dataType))
+ val valueFields = (0 until valueSchema.length).map(i =>
+ valueRow.get(i, valueSchema(i).dataType))
+ (keyFields, valueFields)
+ }
+ // Verify same number of rows
+ assert(filteredBytesData.size == normalSet.size,
+ s"Row count mismatch for column family '$columnFamily': " +
+ s"normal read has ${filteredBytesData.size} rows, bytes read has ${normalSet.size} rows")
Review Comment:
wrong, switch the vals?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReaderAllColumnFamiliesSuite.scala:
##########
@@ -0,0 +1,244 @@
+/*
+ * 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.execution.datasources.v2.state
+
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.expressions.{GenericRowWithSchema, UnsafeRow}
+import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.{HDFSBackedStateStoreProvider, RocksDBStateStoreProvider}
+import org.apache.spark.sql.functions.{count, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.OutputMode
+import org.apache.spark.sql.types.{IntegerType, LongType, StructField, StructType}
+import org.apache.spark.tags.SlowSQLTest
+import org.apache.spark.unsafe.Platform
+
+/**
+ * Test suite to verify StatePartitionReaderAllColumnFamilies functionality.
+ */
+@SlowSQLTest
+class StatePartitionReaderAllColumnFamiliesSuite extends StateDataSourceTestBase {
+
+ import testImplicits._
+
+ /**
+ * Returns a set of (partitionId, key, value) tuples from a normal state read.
+ */
+ private def getNormalReadData(checkpointDir: String): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .load()
+ .selectExpr("partition_id", "key", "value")
+ }
+
+ /**
+ * Returns a DataFrame with raw bytes mode (INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES = true).
+ */
+ private def getBytesReadDf(checkpointDir: String): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .option(StateSourceOptions.INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES, "true")
+ .load()
+ }
+
+ /**
+ * Validates the schema and column families of the bytes read DataFrame.
+ */
+ private def validateBytesReadSchema(df: DataFrame): Unit = {
+ // Verify schema
+ val schema = df.schema
+ assert(schema.fieldNames === Array(
+ "partition_key", "key_bytes", "value_bytes", "column_family_name", "value", "key"))
+ assert(schema("partition_key").dataType.typeName === "struct")
+ assert(schema("key_bytes").dataType.typeName === "binary")
+ assert(schema("value_bytes").dataType.typeName === "binary")
+ assert(schema("column_family_name").dataType.typeName === "string")
+ }
+
+ private def parseBytesReadData(df: Array[Row], keyLength: Int, valueLength: Int)
+ : Set[(GenericRowWithSchema, UnsafeRow, UnsafeRow, String)] = {
+ df.map { row =>
+ val partitionKey = row.getAs[GenericRowWithSchema](0)
+ val keyBytes = row.getAs[Array[Byte]](1)
+ val valueBytes = row.getAs[Array[Byte]](2)
+ val columnFamily = row.getString(3)
+
+ // Deserialize key bytes to UnsafeRow
+ val keyRow = new UnsafeRow(keyLength)
+ keyRow.pointTo(
+ keyBytes,
+ Platform.BYTE_ARRAY_OFFSET,
+ keyBytes.length)
+
+ // Deserialize value bytes to UnsafeRow
+ val valueRow = new UnsafeRow(valueLength)
+ valueRow.pointTo(
+ valueBytes,
+ Platform.BYTE_ARRAY_OFFSET,
+ valueBytes.length)
+ (partitionKey, keyRow.copy(), valueRow.copy(), columnFamily)
+ }
+ .toSet
+ }
+
+ /**
+ * Compares normal read data with bytes read data for a specific column family.
+ */
+ private def compareNormalAndBytesData(
+ normalReadDf: DataFrame,
+ bytesReadDf: DataFrame,
+ columnFamily: String,
+ keySchema: StructType,
+ valueSchema: StructType): Unit = {
+ // Verify data
+ val bytesDf = bytesReadDf
+ .selectExpr("partition_key", "key_bytes", "value_bytes", "column_family_name")
+ .collect()
+ assert(bytesDf.length == 10,
+ s"Expected 10 rows but got: ${bytesDf.length}")
+
+ // Filter bytes data for the specified column family
+ val bytesData = parseBytesReadData(bytesDf, keySchema.length, valueSchema.length)
+ val filteredBytesData = bytesData.filter(_._4 == columnFamily)
+
+ // Apply the projection
+ // Convert to comparable format (extract field values)
+ val normalSet = normalReadDf.collect().map { row =>
+ val key = row.getStruct(1)
+ val value = row.getStruct(2)
+ val keyFields = (0 until key.length).map(i => key.get(i))
+ val valueFields = (0 until value.length).map(i => value.get(i))
+ (keyFields, valueFields)
+ }.toSet
+
+ val bytesSet = filteredBytesData.map { case (_, keyRow, valueRow, _) =>
+ val keyFields = (0 until keySchema.length).map(i =>
+ keyRow.get(i, keySchema(i).dataType))
+ val valueFields = (0 until valueSchema.length).map(i =>
+ valueRow.get(i, valueSchema(i).dataType))
+ (keyFields, valueFields)
+ }
+ // Verify same number of rows
+ assert(filteredBytesData.size == normalSet.size,
+ s"Row count mismatch for column family '$columnFamily': " +
+ s"normal read has ${filteredBytesData.size} rows, bytes read has ${normalSet.size} rows")
+
+ assert(normalSet == bytesSet)
+ }
+
+ test(s"read all column families with simple operator") {
Review Comment:
You need some more test cases to cover all operators that use single CF e.g. agg with composite key, DropDuplicate, DropDuplicateWithinWatermark, SessionWindow, FMGWS, etc. See `StateDataSourceReadSuite` it already has a bunch of query functions you can reuse and example test cases.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala:
##########
@@ -371,6 +378,7 @@ case class StateSourceOptions(
stateVarName: Option[String],
readRegisteredTimers: Boolean,
flattenCollectionTypes: Boolean,
+ internalOnlyReadAllColumnFamilies: Boolean,
Review Comment:
nit: default to false?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala:
##########
@@ -65,6 +65,13 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
val sourceOptions = StateSourceOptions.modifySourceOptions(hadoopConf,
StateSourceOptions.apply(session, hadoopConf, properties))
val stateConf = buildStateStoreConf(sourceOptions.resolvedCpLocation, sourceOptions.batchId)
+ if (sourceOptions.internalOnlyReadAllColumnFamilies
+ && !stateConf.providerClass.contains("RocksDB")) {
+ throw StateDataSourceErrors.invalidOptionValue(
+ StateSourceOptions.INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES,
Review Comment:
This is an internal only option, so we shouldn't expose it in user error message right?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala:
##########
@@ -60,6 +61,17 @@ object SchemaUtil {
.add("key", keySchema)
.add("value", valueSchema)
.add("partition_id", IntegerType)
+ } else if (sourceOptions.internalOnlyReadAllColumnFamilies) {
+ new StructType()
+ // todo: change this to some more specific type after we
+ // can extract partition key from keySchema
+ .add("partition_key", keySchema)
+ .add("key_bytes", BinaryType)
+ .add("value_bytes", BinaryType)
+ .add("column_family_name", StringType)
+ // need key and value schema so that state store can encode data
+ .add("value", valueSchema)
Review Comment:
why do we need the `key` and `value` row here?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala:
##########
@@ -76,6 +88,22 @@ object SchemaUtil {
row
}
+ /**
+ * Creates a unified row from raw key and value bytes.
+ * This is an alias for unifyStateRowPairAsBytes that takes individual byte arrays
Review Comment:
not sure what this line means?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala:
##########
@@ -237,6 +240,44 @@ class StatePartitionReader(
}
}
+/**
+ * An implementation of [[StatePartitionReaderBase]] for reading all column families
+ * in binary format. This reader returns raw key and value bytes along with column family names.
+ */
+class StatePartitionReaderAllColumnFamilies(
+ storeConf: StateStoreConf,
+ hadoopConf: SerializableConfiguration,
+ partition: StateStoreInputPartition,
+ schema: StructType,
+ keyStateEncoderSpec: KeyStateEncoderSpec)
+ extends StatePartitionReaderBase(storeConf, hadoopConf, partition, schema,
+ keyStateEncoderSpec, None, None, None, None) {
+
+ private lazy val store: ReadStateStore = {
+ assert(getStartStoreUniqueId == getEndStoreUniqueId,
+ "Start and end store unique IDs must be the same when reading all column families")
+ provider.getReadStore(
+ partition.sourceOptions.batchId + 1,
+ getStartStoreUniqueId
+ )
+ }
+
+ override lazy val iter: Iterator[InternalRow] = {
+ // Single store with column families (join v3, transformWithState, or simple operators)
Review Comment:
what does this comment mean?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReaderAllColumnFamiliesSuite.scala:
##########
@@ -0,0 +1,244 @@
+/*
+ * 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.execution.datasources.v2.state
+
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.expressions.{GenericRowWithSchema, UnsafeRow}
+import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.{HDFSBackedStateStoreProvider, RocksDBStateStoreProvider}
+import org.apache.spark.sql.functions.{count, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.OutputMode
+import org.apache.spark.sql.types.{IntegerType, LongType, StructField, StructType}
+import org.apache.spark.tags.SlowSQLTest
+import org.apache.spark.unsafe.Platform
+
+/**
+ * Test suite to verify StatePartitionReaderAllColumnFamilies functionality.
Review Comment:
ditto
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala:
##########
@@ -231,7 +259,11 @@ object SchemaUtil {
"user_map_key" -> classOf[StructType],
"user_map_value" -> classOf[StructType],
"expiration_timestamp_ms" -> classOf[LongType],
- "partition_id" -> classOf[IntegerType])
+ "partition_id" -> classOf[IntegerType],
+ "partition_key" -> classOf[StructType],
+ "key_bytes"->classOf[BinaryType],
+ "value_bytes"->classOf[BinaryType],
+ "column_family_name"->classOf[StringType])
Review Comment:
nit: fix space between `->` for `key_bytes`, `value_bytes` and `column_family_name`
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala:
##########
@@ -76,6 +88,22 @@ object SchemaUtil {
row
}
+ /**
+ * Creates a unified row from raw key and value bytes.
+ * This is an alias for unifyStateRowPairAsBytes that takes individual byte arrays
+ * instead of a tuple for better readability.
+ */
+ def unifyStateRowPairAsRawBytes(
+ pair: (UnsafeRow, UnsafeRow),
+ colFamilyName: String): InternalRow = {
+ val row = new GenericInternalRow(6)
+ row.update(0, pair._1)
Review Comment:
Add todo for setting the actual partition key
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala:
##########
@@ -76,6 +88,22 @@ object SchemaUtil {
row
}
+ /**
+ * Creates a unified row from raw key and value bytes.
+ * This is an alias for unifyStateRowPairAsBytes that takes individual byte arrays
+ * instead of a tuple for better readability.
+ */
+ def unifyStateRowPairAsRawBytes(
+ pair: (UnsafeRow, UnsafeRow),
Review Comment:
nit: fix indentation
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReaderAllColumnFamiliesSuite.scala:
##########
@@ -0,0 +1,244 @@
+/*
+ * 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.execution.datasources.v2.state
+
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.expressions.{GenericRowWithSchema, UnsafeRow}
+import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.{HDFSBackedStateStoreProvider, RocksDBStateStoreProvider}
+import org.apache.spark.sql.functions.{count, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.OutputMode
+import org.apache.spark.sql.types.{IntegerType, LongType, StructField, StructType}
+import org.apache.spark.tags.SlowSQLTest
+import org.apache.spark.unsafe.Platform
+
+/**
+ * Test suite to verify StatePartitionReaderAllColumnFamilies functionality.
+ */
+@SlowSQLTest
+class StatePartitionReaderAllColumnFamiliesSuite extends StateDataSourceTestBase {
Review Comment:
ditto
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReaderAllColumnFamiliesSuite.scala:
##########
@@ -0,0 +1,244 @@
+/*
Review Comment:
`StatePartitionAllColumnFamiliesReaderSuite.scala`
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReaderAllColumnFamiliesSuite.scala:
##########
@@ -0,0 +1,244 @@
+/*
+ * 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.execution.datasources.v2.state
+
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.expressions.{GenericRowWithSchema, UnsafeRow}
+import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.{HDFSBackedStateStoreProvider, RocksDBStateStoreProvider}
+import org.apache.spark.sql.functions.{count, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.OutputMode
+import org.apache.spark.sql.types.{IntegerType, LongType, StructField, StructType}
+import org.apache.spark.tags.SlowSQLTest
+import org.apache.spark.unsafe.Platform
+
+/**
+ * Test suite to verify StatePartitionReaderAllColumnFamilies functionality.
+ */
+@SlowSQLTest
Review Comment:
why is this marked as slow?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReaderAllColumnFamiliesSuite.scala:
##########
@@ -0,0 +1,244 @@
+/*
+ * 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.execution.datasources.v2.state
+
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.expressions.{GenericRowWithSchema, UnsafeRow}
+import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.{HDFSBackedStateStoreProvider, RocksDBStateStoreProvider}
+import org.apache.spark.sql.functions.{count, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.OutputMode
+import org.apache.spark.sql.types.{IntegerType, LongType, StructField, StructType}
+import org.apache.spark.tags.SlowSQLTest
+import org.apache.spark.unsafe.Platform
+
+/**
+ * Test suite to verify StatePartitionReaderAllColumnFamilies functionality.
+ */
+@SlowSQLTest
+class StatePartitionReaderAllColumnFamiliesSuite extends StateDataSourceTestBase {
+
+ import testImplicits._
+
+ /**
+ * Returns a set of (partitionId, key, value) tuples from a normal state read.
+ */
+ private def getNormalReadData(checkpointDir: String): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .load()
+ .selectExpr("partition_id", "key", "value")
+ }
+
+ /**
+ * Returns a DataFrame with raw bytes mode (INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES = true).
+ */
+ private def getBytesReadDf(checkpointDir: String): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .option(StateSourceOptions.INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES, "true")
+ .load()
+ }
+
+ /**
+ * Validates the schema and column families of the bytes read DataFrame.
+ */
+ private def validateBytesReadSchema(df: DataFrame): Unit = {
+ // Verify schema
+ val schema = df.schema
+ assert(schema.fieldNames === Array(
+ "partition_key", "key_bytes", "value_bytes", "column_family_name", "value", "key"))
+ assert(schema("partition_key").dataType.typeName === "struct")
+ assert(schema("key_bytes").dataType.typeName === "binary")
+ assert(schema("value_bytes").dataType.typeName === "binary")
+ assert(schema("column_family_name").dataType.typeName === "string")
+ }
+
+ private def parseBytesReadData(df: Array[Row], keyLength: Int, valueLength: Int)
Review Comment:
nit: `convertBytesToRow`?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReaderAllColumnFamiliesSuite.scala:
##########
@@ -0,0 +1,244 @@
+/*
+ * 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.execution.datasources.v2.state
+
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.expressions.{GenericRowWithSchema, UnsafeRow}
+import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.{HDFSBackedStateStoreProvider, RocksDBStateStoreProvider}
+import org.apache.spark.sql.functions.{count, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.OutputMode
+import org.apache.spark.sql.types.{IntegerType, LongType, StructField, StructType}
+import org.apache.spark.tags.SlowSQLTest
+import org.apache.spark.unsafe.Platform
+
+/**
+ * Test suite to verify StatePartitionReaderAllColumnFamilies functionality.
+ */
+@SlowSQLTest
+class StatePartitionReaderAllColumnFamiliesSuite extends StateDataSourceTestBase {
Review Comment:
Lets make the test cases in this suite run with both checkpoint v1 and v2
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReaderAllColumnFamiliesSuite.scala:
##########
@@ -0,0 +1,244 @@
+/*
+ * 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.execution.datasources.v2.state
+
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.expressions.{GenericRowWithSchema, UnsafeRow}
+import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.{HDFSBackedStateStoreProvider, RocksDBStateStoreProvider}
+import org.apache.spark.sql.functions.{count, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.OutputMode
+import org.apache.spark.sql.types.{IntegerType, LongType, StructField, StructType}
+import org.apache.spark.tags.SlowSQLTest
+import org.apache.spark.unsafe.Platform
+
+/**
+ * Test suite to verify StatePartitionReaderAllColumnFamilies functionality.
+ */
+@SlowSQLTest
+class StatePartitionReaderAllColumnFamiliesSuite extends StateDataSourceTestBase {
+
+ import testImplicits._
+
+ /**
+ * Returns a set of (partitionId, key, value) tuples from a normal state read.
+ */
+ private def getNormalReadData(checkpointDir: String): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .load()
+ .selectExpr("partition_id", "key", "value")
+ }
+
+ /**
+ * Returns a DataFrame with raw bytes mode (INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES = true).
+ */
+ private def getBytesReadDf(checkpointDir: String): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .option(StateSourceOptions.INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES, "true")
+ .load()
+ }
+
+ /**
+ * Validates the schema and column families of the bytes read DataFrame.
+ */
+ private def validateBytesReadSchema(df: DataFrame): Unit = {
+ // Verify schema
+ val schema = df.schema
+ assert(schema.fieldNames === Array(
+ "partition_key", "key_bytes", "value_bytes", "column_family_name", "value", "key"))
+ assert(schema("partition_key").dataType.typeName === "struct")
+ assert(schema("key_bytes").dataType.typeName === "binary")
+ assert(schema("value_bytes").dataType.typeName === "binary")
+ assert(schema("column_family_name").dataType.typeName === "string")
+ }
+
+ private def parseBytesReadData(df: Array[Row], keyLength: Int, valueLength: Int)
+ : Set[(GenericRowWithSchema, UnsafeRow, UnsafeRow, String)] = {
+ df.map { row =>
+ val partitionKey = row.getAs[GenericRowWithSchema](0)
+ val keyBytes = row.getAs[Array[Byte]](1)
+ val valueBytes = row.getAs[Array[Byte]](2)
+ val columnFamily = row.getString(3)
+
+ // Deserialize key bytes to UnsafeRow
+ val keyRow = new UnsafeRow(keyLength)
+ keyRow.pointTo(
+ keyBytes,
+ Platform.BYTE_ARRAY_OFFSET,
+ keyBytes.length)
+
+ // Deserialize value bytes to UnsafeRow
+ val valueRow = new UnsafeRow(valueLength)
+ valueRow.pointTo(
+ valueBytes,
+ Platform.BYTE_ARRAY_OFFSET,
+ valueBytes.length)
+ (partitionKey, keyRow.copy(), valueRow.copy(), columnFamily)
+ }
+ .toSet
+ }
+
+ /**
+ * Compares normal read data with bytes read data for a specific column family.
+ */
+ private def compareNormalAndBytesData(
+ normalReadDf: DataFrame,
+ bytesReadDf: DataFrame,
+ columnFamily: String,
+ keySchema: StructType,
+ valueSchema: StructType): Unit = {
+ // Verify data
+ val bytesDf = bytesReadDf
+ .selectExpr("partition_key", "key_bytes", "value_bytes", "column_family_name")
+ .collect()
+ assert(bytesDf.length == 10,
+ s"Expected 10 rows but got: ${bytesDf.length}")
+
+ // Filter bytes data for the specified column family
+ val bytesData = parseBytesReadData(bytesDf, keySchema.length, valueSchema.length)
+ val filteredBytesData = bytesData.filter(_._4 == columnFamily)
+
+ // Apply the projection
+ // Convert to comparable format (extract field values)
+ val normalSet = normalReadDf.collect().map { row =>
+ val key = row.getStruct(1)
+ val value = row.getStruct(2)
+ val keyFields = (0 until key.length).map(i => key.get(i))
+ val valueFields = (0 until value.length).map(i => value.get(i))
+ (keyFields, valueFields)
+ }.toSet
+
+ val bytesSet = filteredBytesData.map { case (_, keyRow, valueRow, _) =>
+ val keyFields = (0 until keySchema.length).map(i =>
+ keyRow.get(i, keySchema(i).dataType))
+ val valueFields = (0 until valueSchema.length).map(i =>
+ valueRow.get(i, valueSchema(i).dataType))
+ (keyFields, valueFields)
+ }
+ // Verify same number of rows
+ assert(filteredBytesData.size == normalSet.size,
+ s"Row count mismatch for column family '$columnFamily': " +
+ s"normal read has ${filteredBytesData.size} rows, bytes read has ${normalSet.size} rows")
+
+ assert(normalSet == bytesSet)
+ }
+
+ test(s"read all column families with simple operator") {
+ withTempDir { tempDir =>
+ withSQLConf(
+ SQLConf.STATE_STORE_PROVIDER_CLASS.key -> classOf[RocksDBStateStoreProvider].getName,
Review Comment:
should we just set this at the class level, so that all test cases will use RocksDB. Except when a test case overrides the conf.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReaderAllColumnFamiliesSuite.scala:
##########
@@ -0,0 +1,244 @@
+/*
+ * 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.execution.datasources.v2.state
+
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.expressions.{GenericRowWithSchema, UnsafeRow}
+import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.{HDFSBackedStateStoreProvider, RocksDBStateStoreProvider}
+import org.apache.spark.sql.functions.{count, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.OutputMode
+import org.apache.spark.sql.types.{IntegerType, LongType, StructField, StructType}
+import org.apache.spark.tags.SlowSQLTest
+import org.apache.spark.unsafe.Platform
+
+/**
+ * Test suite to verify StatePartitionReaderAllColumnFamilies functionality.
+ */
+@SlowSQLTest
+class StatePartitionReaderAllColumnFamiliesSuite extends StateDataSourceTestBase {
+
+ import testImplicits._
+
+ /**
+ * Returns a set of (partitionId, key, value) tuples from a normal state read.
+ */
+ private def getNormalReadData(checkpointDir: String): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .load()
+ .selectExpr("partition_id", "key", "value")
+ }
+
+ /**
+ * Returns a DataFrame with raw bytes mode (INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES = true).
+ */
+ private def getBytesReadDf(checkpointDir: String): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .option(StateSourceOptions.INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES, "true")
+ .load()
+ }
+
+ /**
+ * Validates the schema and column families of the bytes read DataFrame.
+ */
+ private def validateBytesReadSchema(df: DataFrame): Unit = {
+ // Verify schema
+ val schema = df.schema
+ assert(schema.fieldNames === Array(
+ "partition_key", "key_bytes", "value_bytes", "column_family_name", "value", "key"))
+ assert(schema("partition_key").dataType.typeName === "struct")
+ assert(schema("key_bytes").dataType.typeName === "binary")
+ assert(schema("value_bytes").dataType.typeName === "binary")
+ assert(schema("column_family_name").dataType.typeName === "string")
+ }
+
+ private def parseBytesReadData(df: Array[Row], keyLength: Int, valueLength: Int)
+ : Set[(GenericRowWithSchema, UnsafeRow, UnsafeRow, String)] = {
+ df.map { row =>
+ val partitionKey = row.getAs[GenericRowWithSchema](0)
+ val keyBytes = row.getAs[Array[Byte]](1)
+ val valueBytes = row.getAs[Array[Byte]](2)
+ val columnFamily = row.getString(3)
+
+ // Deserialize key bytes to UnsafeRow
+ val keyRow = new UnsafeRow(keyLength)
+ keyRow.pointTo(
+ keyBytes,
+ Platform.BYTE_ARRAY_OFFSET,
+ keyBytes.length)
+
+ // Deserialize value bytes to UnsafeRow
+ val valueRow = new UnsafeRow(valueLength)
+ valueRow.pointTo(
+ valueBytes,
+ Platform.BYTE_ARRAY_OFFSET,
+ valueBytes.length)
+ (partitionKey, keyRow.copy(), valueRow.copy(), columnFamily)
+ }
+ .toSet
+ }
+
+ /**
+ * Compares normal read data with bytes read data for a specific column family.
+ */
+ private def compareNormalAndBytesData(
+ normalReadDf: DataFrame,
+ bytesReadDf: DataFrame,
+ columnFamily: String,
+ keySchema: StructType,
+ valueSchema: StructType): Unit = {
+ // Verify data
+ val bytesDf = bytesReadDf
+ .selectExpr("partition_key", "key_bytes", "value_bytes", "column_family_name")
+ .collect()
+ assert(bytesDf.length == 10,
+ s"Expected 10 rows but got: ${bytesDf.length}")
+
+ // Filter bytes data for the specified column family
+ val bytesData = parseBytesReadData(bytesDf, keySchema.length, valueSchema.length)
+ val filteredBytesData = bytesData.filter(_._4 == columnFamily)
+
+ // Apply the projection
+ // Convert to comparable format (extract field values)
+ val normalSet = normalReadDf.collect().map { row =>
+ val key = row.getStruct(1)
+ val value = row.getStruct(2)
+ val keyFields = (0 until key.length).map(i => key.get(i))
+ val valueFields = (0 until value.length).map(i => value.get(i))
+ (keyFields, valueFields)
+ }.toSet
+
+ val bytesSet = filteredBytesData.map { case (_, keyRow, valueRow, _) =>
+ val keyFields = (0 until keySchema.length).map(i =>
+ keyRow.get(i, keySchema(i).dataType))
+ val valueFields = (0 until valueSchema.length).map(i =>
+ valueRow.get(i, valueSchema(i).dataType))
+ (keyFields, valueFields)
+ }
+ // Verify same number of rows
+ assert(filteredBytesData.size == normalSet.size,
Review Comment:
are you comparing the correct vals here?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala:
##########
@@ -60,6 +61,17 @@ object SchemaUtil {
.add("key", keySchema)
.add("value", valueSchema)
.add("partition_id", IntegerType)
+ } else if (sourceOptions.internalOnlyReadAllColumnFamilies) {
+ new StructType()
+ // todo: change this to some more specific type after we
Review Comment:
Add this ticket number `SPARK-54443` to the todo
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "zifeif2 (via GitHub)" <gi...@apache.org>.
zifeif2 commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2550736299
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala:
##########
@@ -60,6 +61,17 @@ object SchemaUtil {
.add("key", keySchema)
.add("value", valueSchema)
.add("partition_id", IntegerType)
+ } else if (sourceOptions.internalOnlyReadAllColumnFamilies) {
+ new StructType()
+ // todo: change this to some more specific type after we
+ // can extract partition key from keySchema
+ .add("partition_key", keySchema)
+ .add("key_bytes", BinaryType)
+ .add("value_bytes", BinaryType)
+ .add("column_family_name", StringType)
+ // need key and value schema so that state store can encode data
+ .add("value", valueSchema)
Review Comment:
I thought RocksDBStateStoreProvider will need valueSchema and keySchema to pass `validateStateRowFormat` after getting raw bytes from RocksDB. But looks like we can bypass the validation by setting `formatValidationEnabled` and `formatValidationCheckValue` in StateStoreConf to false, so I am going to do so in the next version
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala:
##########
@@ -237,6 +240,44 @@ class StatePartitionReader(
}
}
+/**
+ * An implementation of [[StatePartitionReaderBase]] for reading all column families
+ * in binary format. This reader returns raw key and value bytes along with column family names.
+ */
+class StatePartitionReaderAllColumnFamilies(
+ storeConf: StateStoreConf,
+ hadoopConf: SerializableConfiguration,
+ partition: StateStoreInputPartition,
+ schema: StructType,
+ keyStateEncoderSpec: KeyStateEncoderSpec)
+ extends StatePartitionReaderBase(storeConf, hadoopConf, partition, schema,
+ keyStateEncoderSpec, None, None, None, None) {
+
+ private lazy val store: ReadStateStore = {
+ assert(getStartStoreUniqueId == getEndStoreUniqueId,
+ "Start and end store unique IDs must be the same when reading all column families")
+ provider.getReadStore(
+ partition.sourceOptions.batchId + 1,
+ getStartStoreUniqueId
+ )
+ }
+
+ override lazy val iter: Iterator[InternalRow] = {
+ // Single store with column families (join v3, transformWithState, or simple operators)
Review Comment:
It's not supposed to be there. I'll remove it
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala:
##########
@@ -65,6 +65,13 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
val sourceOptions = StateSourceOptions.modifySourceOptions(hadoopConf,
StateSourceOptions.apply(session, hadoopConf, properties))
val stateConf = buildStateStoreConf(sourceOptions.resolvedCpLocation, sourceOptions.batchId)
+ if (sourceOptions.internalOnlyReadAllColumnFamilies
+ && !stateConf.providerClass.contains("RocksDB")) {
+ throw StateDataSourceErrors.invalidOptionValue(
+ StateSourceOptions.INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES,
Review Comment:
true! Added a new error class in OfflineStateRepartitionErrors in next version
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReaderAllColumnFamiliesSuite.scala:
##########
@@ -0,0 +1,244 @@
+/*
+ * 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.execution.datasources.v2.state
+
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.expressions.{GenericRowWithSchema, UnsafeRow}
+import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.{HDFSBackedStateStoreProvider, RocksDBStateStoreProvider}
+import org.apache.spark.sql.functions.{count, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.OutputMode
+import org.apache.spark.sql.types.{IntegerType, LongType, StructField, StructType}
+import org.apache.spark.tags.SlowSQLTest
+import org.apache.spark.unsafe.Platform
+
+/**
+ * Test suite to verify StatePartitionReaderAllColumnFamilies functionality.
+ */
+@SlowSQLTest
+class StatePartitionReaderAllColumnFamiliesSuite extends StateDataSourceTestBase {
+
+ import testImplicits._
+
+ /**
+ * Returns a set of (partitionId, key, value) tuples from a normal state read.
+ */
+ private def getNormalReadData(checkpointDir: String): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .load()
+ .selectExpr("partition_id", "key", "value")
+ }
+
+ /**
+ * Returns a DataFrame with raw bytes mode (INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES = true).
+ */
+ private def getBytesReadDf(checkpointDir: String): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .option(StateSourceOptions.INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES, "true")
+ .load()
+ }
+
+ /**
+ * Validates the schema and column families of the bytes read DataFrame.
+ */
+ private def validateBytesReadSchema(df: DataFrame): Unit = {
+ // Verify schema
+ val schema = df.schema
+ assert(schema.fieldNames === Array(
+ "partition_key", "key_bytes", "value_bytes", "column_family_name", "value", "key"))
+ assert(schema("partition_key").dataType.typeName === "struct")
+ assert(schema("key_bytes").dataType.typeName === "binary")
+ assert(schema("value_bytes").dataType.typeName === "binary")
+ assert(schema("column_family_name").dataType.typeName === "string")
+ }
+
+ private def parseBytesReadData(df: Array[Row], keyLength: Int, valueLength: Int)
+ : Set[(GenericRowWithSchema, UnsafeRow, UnsafeRow, String)] = {
+ df.map { row =>
+ val partitionKey = row.getAs[GenericRowWithSchema](0)
+ val keyBytes = row.getAs[Array[Byte]](1)
+ val valueBytes = row.getAs[Array[Byte]](2)
+ val columnFamily = row.getString(3)
+
+ // Deserialize key bytes to UnsafeRow
+ val keyRow = new UnsafeRow(keyLength)
+ keyRow.pointTo(
+ keyBytes,
+ Platform.BYTE_ARRAY_OFFSET,
+ keyBytes.length)
+
+ // Deserialize value bytes to UnsafeRow
+ val valueRow = new UnsafeRow(valueLength)
+ valueRow.pointTo(
+ valueBytes,
+ Platform.BYTE_ARRAY_OFFSET,
+ valueBytes.length)
+ (partitionKey, keyRow.copy(), valueRow.copy(), columnFamily)
+ }
+ .toSet
+ }
+
+ /**
+ * Compares normal read data with bytes read data for a specific column family.
+ */
+ private def compareNormalAndBytesData(
+ normalReadDf: DataFrame,
+ bytesReadDf: DataFrame,
+ columnFamily: String,
+ keySchema: StructType,
+ valueSchema: StructType): Unit = {
+ // Verify data
+ val bytesDf = bytesReadDf
+ .selectExpr("partition_key", "key_bytes", "value_bytes", "column_family_name")
+ .collect()
+ assert(bytesDf.length == 10,
+ s"Expected 10 rows but got: ${bytesDf.length}")
+
+ // Filter bytes data for the specified column family
+ val bytesData = parseBytesReadData(bytesDf, keySchema.length, valueSchema.length)
+ val filteredBytesData = bytesData.filter(_._4 == columnFamily)
+
+ // Apply the projection
+ // Convert to comparable format (extract field values)
+ val normalSet = normalReadDf.collect().map { row =>
+ val key = row.getStruct(1)
+ val value = row.getStruct(2)
+ val keyFields = (0 until key.length).map(i => key.get(i))
+ val valueFields = (0 until value.length).map(i => value.get(i))
+ (keyFields, valueFields)
+ }.toSet
+
+ val bytesSet = filteredBytesData.map { case (_, keyRow, valueRow, _) =>
+ val keyFields = (0 until keySchema.length).map(i =>
+ keyRow.get(i, keySchema(i).dataType))
+ val valueFields = (0 until valueSchema.length).map(i =>
+ valueRow.get(i, valueSchema(i).dataType))
+ (keyFields, valueFields)
+ }
+ // Verify same number of rows
+ assert(filteredBytesData.size == normalSet.size,
Review Comment:
Refactored in the next version so that it's more obvious to compare the "number of rows from rawBytesData that belongs to a column family" against "normal dataframe"
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala:
##########
@@ -76,6 +88,22 @@ object SchemaUtil {
row
}
+ /**
+ * Creates a unified row from raw key and value bytes.
+ * This is an alias for unifyStateRowPairAsBytes that takes individual byte arrays
+ * instead of a tuple for better readability.
+ */
+ def unifyStateRowPairAsRawBytes(
+ pair: (UnsafeRow, UnsafeRow),
+ colFamilyName: String): InternalRow = {
+ val row = new GenericInternalRow(6)
Review Comment:
Because previously the schema has 6 fields but we only need 4 (as we didn't need key and value for repartitioning). Next version removes key and value columns from schema, so I'll change this to 4 instead.
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2582793268
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala:
##########
@@ -407,6 +418,7 @@ object StateSourceOptions extends DataSourceOptions {
val STATE_VAR_NAME = newOption("stateVarName")
val READ_REGISTERED_TIMERS = newOption("readRegisteredTimers")
val FLATTEN_COLLECTION_TYPES = newOption("flattenCollectionTypes")
+ val INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES = newOption("internalOnlyReadAllColumnFamilies")
Review Comment:
Lets reword this
`_readAllColumnFamilies` and mention that options beginning with `_` are for internal use only and should not be supported externally
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "micheal-o (via GitHub)" <gi...@apache.org>.
micheal-o commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2539756121
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala:
##########
@@ -65,28 +65,38 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
val sourceOptions = StateSourceOptions.modifySourceOptions(hadoopConf,
StateSourceOptions.apply(session, hadoopConf, properties))
val stateConf = buildStateStoreConf(sourceOptions.resolvedCpLocation, sourceOptions.batchId)
- val stateStoreReaderInfo: StateStoreReaderInfo = getStoreMetadataAndRunChecks(
- sourceOptions)
-
- // The key state encoder spec should be available for all operators except stream-stream joins
- val keyStateEncoderSpec = if (stateStoreReaderInfo.keyStateEncoderSpecOpt.isDefined) {
- stateStoreReaderInfo.keyStateEncoderSpecOpt.get
+ if (sourceOptions.readAllColumnFamilies) {
+ // For readAllColumnFamilies mode, we don't need specific metadata
+ val keyStateEncoderSpec = NoPrefixKeyStateEncoderSpec(new StructType())
+ new StateTable(session, schema, sourceOptions, stateConf, keyStateEncoderSpec,
Review Comment:
We need to know the metadata/schema for the default CF right. Since this is only supporting default CF for now.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala:
##########
@@ -65,28 +65,38 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
val sourceOptions = StateSourceOptions.modifySourceOptions(hadoopConf,
StateSourceOptions.apply(session, hadoopConf, properties))
val stateConf = buildStateStoreConf(sourceOptions.resolvedCpLocation, sourceOptions.batchId)
- val stateStoreReaderInfo: StateStoreReaderInfo = getStoreMetadataAndRunChecks(
- sourceOptions)
-
- // The key state encoder spec should be available for all operators except stream-stream joins
- val keyStateEncoderSpec = if (stateStoreReaderInfo.keyStateEncoderSpecOpt.isDefined) {
- stateStoreReaderInfo.keyStateEncoderSpecOpt.get
+ if (sourceOptions.readAllColumnFamilies) {
+ // For readAllColumnFamilies mode, we don't need specific metadata
+ val keyStateEncoderSpec = NoPrefixKeyStateEncoderSpec(new StructType())
Review Comment:
Lets get the encoder from the schema.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala:
##########
@@ -84,13 +88,17 @@ abstract class StatePartitionReaderBase(
protected val keySchema = {
if (SchemaUtil.checkVariableType(stateVariableInfoOpt, StateVariableType.MapState)) {
SchemaUtil.getCompositeKeySchema(schema, partition.sourceOptions)
+ } else if (partition.sourceOptions.readAllColumnFamilies) {
+ new StructType().add("keyBytes", BinaryType, nullable = false)
} else {
SchemaUtil.getSchemaAsDataType(schema, "key").asInstanceOf[StructType]
}
}
protected val valueSchema = if (stateVariableInfoOpt.isDefined) {
schemaForValueRow
+ } else if (partition.sourceOptions.readAllColumnFamilies) {
+ new StructType().add("valueBytes", BinaryType, nullable = false)
Review Comment:
ditto
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala:
##########
@@ -406,6 +418,7 @@ object StateSourceOptions extends DataSourceOptions {
val STATE_VAR_NAME = newOption("stateVarName")
val READ_REGISTERED_TIMERS = newOption("readRegisteredTimers")
val FLATTEN_COLLECTION_TYPES = newOption("flattenCollectionTypes")
+ val READ_ALL_COLUMN_FAMILIES = newOption("readAllColumnFamilies")
Review Comment:
make this `internalOnlyReadAllColumnFamilies`
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala:
##########
@@ -84,13 +88,17 @@ abstract class StatePartitionReaderBase(
protected val keySchema = {
if (SchemaUtil.checkVariableType(stateVariableInfoOpt, StateVariableType.MapState)) {
SchemaUtil.getCompositeKeySchema(schema, partition.sourceOptions)
+ } else if (partition.sourceOptions.readAllColumnFamilies) {
+ new StructType().add("keyBytes", BinaryType, nullable = false)
Review Comment:
This is wrong, this is meant to be the actual key schema. Because this is the schema that would be used to load the store
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala:
##########
@@ -237,6 +245,91 @@ class StatePartitionReader(
}
}
+/**
+ * An implementation of [[StatePartitionReaderBase]] for reading all column families
+ * in binary format. This reader returns raw key and value bytes along with column family names.
+ */
+class StatePartitionReaderAllColumnFamilies(
+ storeConf: StateStoreConf,
+ hadoopConf: SerializableConfiguration,
+ partition: StateStoreInputPartition,
+ schema: StructType)
+ extends StatePartitionReaderBase(storeConf, hadoopConf, partition, schema,
+ NoPrefixKeyStateEncoderSpec(new StructType()), None, None, None, None) {
Review Comment:
ditto. Use the actual key encoder of the default CF. We need that when creating the provider
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala:
##########
@@ -75,7 +75,7 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
private val providerName = "HDFSBackedStateStoreProvider"
class HDFSBackedReadStateStore(val version: Long, map: HDFSBackedStateStoreMap)
- extends ReadStateStore {
+ extends ReadStateStore with SupportsRawBytesRead {
Review Comment:
We should throw exception if trying to use this reader with HDFS checkpoint
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala:
##########
@@ -419,6 +420,21 @@ private[sql] class RocksDBStateStoreProvider
}
}
+ override def rawIterator(colFamilyName: String): Iterator[(Array[Byte], Array[Byte])] = {
Review Comment:
no need to implement this rawIterator, since we need the keyRow too. The normal iterator is fine
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala:
##########
@@ -60,6 +61,12 @@ object SchemaUtil {
.add("key", keySchema)
.add("value", valueSchema)
.add("partition_id", IntegerType)
+ } else if (sourceOptions.readAllColumnFamilies) {
+ new StructType()
+ .add("partition_id", IntegerType)
Review Comment:
This is `partition_key`. And its type will be the actual schema of the partition key. For now use the actual key schema as the partition key schema. And when you write the row, write the key row as the partition_key for now.
I'm working on the change for specifying the partition key schema and extracting the partition key from the key row.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala:
##########
@@ -237,6 +245,91 @@ class StatePartitionReader(
}
}
+/**
+ * An implementation of [[StatePartitionReaderBase]] for reading all column families
+ * in binary format. This reader returns raw key and value bytes along with column family names.
+ */
+class StatePartitionReaderAllColumnFamilies(
+ storeConf: StateStoreConf,
+ hadoopConf: SerializableConfiguration,
+ partition: StateStoreInputPartition,
+ schema: StructType)
+ extends StatePartitionReaderBase(storeConf, hadoopConf, partition, schema,
+ NoPrefixKeyStateEncoderSpec(new StructType()), None, None, None, None) {
+
+ val allStateStoreMetadata = {
+ new StateMetadataPartitionReader(
+ partition.sourceOptions.resolvedCpLocation,
+ new SerializableConfiguration(hadoopConf.value),
+ partition.sourceOptions.batchId).stateMetadata.toArray
+ }
+
+ private lazy val store: ReadStateStore = {
+ assert(getStartStoreUniqueId == getEndStoreUniqueId,
+ "Start and end store unique IDs must be the same when reading all column families")
+ provider.getReadStore(
+ partition.sourceOptions.batchId + 1,
+ getStartStoreUniqueId
+ )
+ }
+
+ val colFamilyNames: Seq[String] = {
+ // todo: Support operator with multiple column family names in next PR
+ Seq[String]()
+ }
+
+ override protected lazy val provider: StateStoreProvider = {
+ val stateStoreId = StateStoreId(partition.sourceOptions.stateCheckpointLocation.toString,
+ partition.sourceOptions.operatorId, partition.partition, partition.sourceOptions.storeName)
+ val stateStoreProviderId = StateStoreProviderId(stateStoreId, partition.queryId)
+
+ // Disable format validation when reading raw bytes.
+ // We use binary schemas (keyBytes/valueBytes) which don't match the actual schema
+ // of the stored data. Validation would fail in HDFSBackedStateStoreProvider when
+ // loading data from disk, so we disable it for raw bytes mode.
+ val modifiedStoreConf = storeConf.withFormatValidationDisabled()
+
+ val keyStateEncoderSpec = NoPrefixKeyStateEncoderSpec(keySchema)
+ val provider = StateStoreProvider.createAndInit(
+ stateStoreProviderId, keySchema, valueSchema, keyStateEncoderSpec,
+ useColumnFamilies = colFamilyNames.nonEmpty, modifiedStoreConf, hadoopConf.value, false, None)
+
+ provider
+ }
+
+ override lazy val iter: Iterator[InternalRow] = {
+ // Single store with column families (join v3, transformWithState, or simple operators)
+ require(store.isInstanceOf[SupportsRawBytesRead],
+ s"State store ${store.getClass.getName} does not support raw bytes reading")
+
+ val rawStore = store.asInstanceOf[SupportsRawBytesRead]
+ if (colFamilyNames.isEmpty) {
+ rawStore
+ .rawIterator()
Review Comment:
no need to implement raw bytes iterator. Lets use the normal iterator. And when you get the UnsafeRow, you can call `row.getBytes` to get its bytes
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala:
##########
@@ -76,6 +83,24 @@ object SchemaUtil {
row
}
+ /**
+ * Creates a unified row from raw key and value bytes.
+ * This is an alias for unifyStateRowPairAsBytes that takes individual byte arrays
+ * instead of a tuple for better readability.
+ */
+ def unifyStateRowPairAsRawBytes(
+ partition: Int,
+ keyBytes: Array[Byte],
+ valueBytes: Array[Byte],
+ colFamilyName: String): InternalRow = {
+ val row = new GenericInternalRow(4)
+ row.update(0, partition)
Review Comment:
I explained this above, write the key row (not key bytes) for now as the partition key.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala:
##########
@@ -237,6 +245,91 @@ class StatePartitionReader(
}
}
+/**
+ * An implementation of [[StatePartitionReaderBase]] for reading all column families
+ * in binary format. This reader returns raw key and value bytes along with column family names.
+ */
+class StatePartitionReaderAllColumnFamilies(
+ storeConf: StateStoreConf,
+ hadoopConf: SerializableConfiguration,
+ partition: StateStoreInputPartition,
+ schema: StructType)
+ extends StatePartitionReaderBase(storeConf, hadoopConf, partition, schema,
+ NoPrefixKeyStateEncoderSpec(new StructType()), None, None, None, None) {
+
+ val allStateStoreMetadata = {
+ new StateMetadataPartitionReader(
+ partition.sourceOptions.resolvedCpLocation,
+ new SerializableConfiguration(hadoopConf.value),
+ partition.sourceOptions.batchId).stateMetadata.toArray
+ }
+
+ private lazy val store: ReadStateStore = {
+ assert(getStartStoreUniqueId == getEndStoreUniqueId,
+ "Start and end store unique IDs must be the same when reading all column families")
+ provider.getReadStore(
+ partition.sourceOptions.batchId + 1,
+ getStartStoreUniqueId
+ )
+ }
+
+ val colFamilyNames: Seq[String] = {
+ // todo: Support operator with multiple column family names in next PR
Review Comment:
add Spark ticket number for the todo
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala:
##########
@@ -75,7 +75,7 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with
private val providerName = "HDFSBackedStateStoreProvider"
class HDFSBackedReadStateStore(val version: Long, map: HDFSBackedStateStoreMap)
- extends ReadStateStore {
+ extends ReadStateStore with SupportsRawBytesRead {
Review Comment:
We are not supporting HDFS state store for repartitioning
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreConf.scala:
##########
@@ -163,6 +163,26 @@ class StateStoreConf(
*/
val sqlConfs: Map[String, String] =
sqlConf.getAllConfs.filter(_._1.startsWith("spark.sql.streaming.stateStore."))
+
+ /**
+ * Creates a copy of this StateStoreConf with format validation disabled.
+ * This is useful when reading raw bytes where the schema used (binary) doesn't match
+ * the actual stored data schema.
+ */
+ def withFormatValidationDisabled(): StateStoreConf = {
Review Comment:
Why?
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2582801732
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala:
##########
@@ -492,6 +505,29 @@ object StateSourceOptions extends DataSourceOptions {
val readChangeFeed = Option(options.get(READ_CHANGE_FEED)).exists(_.toBoolean)
+ val internalOnlyReadAllColumnFamilies = try {
+ Option(options.get(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES)).exists(_.toBoolean)
+ } catch {
+ case _: IllegalArgumentException =>
+ throw StateDataSourceErrors.invalidOptionValue(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES,
+ "Boolean value is expected")
+ }
+
+ if (internalOnlyReadAllColumnFamilies && stateVarName.isDefined) {
Review Comment:
We don't need to do individual if conditions.
Lets do something like this
````
if (_readAllColumnFamilies) {
if (stateVarName.isDefined) {
... throw error
}
// cover other conditions
// add a log here to indicate that the flag is enabled
}
```
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2582801732
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala:
##########
@@ -492,6 +505,29 @@ object StateSourceOptions extends DataSourceOptions {
val readChangeFeed = Option(options.get(READ_CHANGE_FEED)).exists(_.toBoolean)
+ val internalOnlyReadAllColumnFamilies = try {
+ Option(options.get(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES)).exists(_.toBoolean)
+ } catch {
+ case _: IllegalArgumentException =>
+ throw StateDataSourceErrors.invalidOptionValue(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES,
+ "Boolean value is expected")
+ }
+
+ if (internalOnlyReadAllColumnFamilies && stateVarName.isDefined) {
Review Comment:
We don't need to do individual if conditions.
Lets do something like this
```
if (_readAllColumnFamilies) {
if (stateVarName.isDefined) {
... throw error
}
// cover other conditions
// add a log here to indicate that the flag is enabled
}
```
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "zifeif2 (via GitHub)" <gi...@apache.org>.
zifeif2 commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2582799523
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala:
##########
@@ -407,6 +418,7 @@ object StateSourceOptions extends DataSourceOptions {
val STATE_VAR_NAME = newOption("stateVarName")
val READ_REGISTERED_TIMERS = newOption("readRegisteredTimers")
val FLATTEN_COLLECTION_TYPES = newOption("flattenCollectionTypes")
+ val INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES = newOption("internalOnlyReadAllColumnFamilies")
Review Comment:
Could you explain a bit more regarding to exposing an interface at the partition reader level?
I am thinking about the interface of using StatePartitionAllColumnFamiliesReader as the following
```
spark.read
.format("statestore")
.option(StateSourceOptions.PATH, tempDir.getAbsolutePath)
.option(StateSourceOptions.INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES, "true")
.load()
.collect()
```
Therefore, we need to expose this option in StateDataSource.
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "micheal-o (via GitHub)" <gi...@apache.org>.
micheal-o commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2579600707
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OfflineStateRepartitionErrors.scala:
##########
@@ -201,3 +207,11 @@ class StateRepartitionUnsupportedOffsetSeqVersionError(
checkpointLocation,
subClass = "UNSUPPORTED_OFFSET_SEQ_VERSION",
messageParameters = Map("version" -> version.toString))
+
+class StateRepartitionUnsupportedProviderError(
Review Comment:
@zifeif2 probably my eyes. You can resolve it.
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "micheal-o (via GitHub)" <gi...@apache.org>.
micheal-o commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2579262405
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala:
##########
@@ -492,6 +504,29 @@ object StateSourceOptions extends DataSourceOptions {
val readChangeFeed = Option(options.get(READ_CHANGE_FEED)).exists(_.toBoolean)
+ val internalOnlyReadAllColumnFamilies = try {
+ Option(options.get(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES)).exists(_.toBoolean)
+ } catch {
+ case _: IllegalArgumentException =>
+ throw StateDataSourceErrors.invalidOptionValue(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES,
+ "Boolean value is expected")
+ }
+
+ if (internalOnlyReadAllColumnFamilies && stateVarName.isDefined) {
+ throw StateDataSourceErrors.conflictOptions(
+ Seq(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES, STATE_VAR_NAME))
+ }
+
+ if (internalOnlyReadAllColumnFamilies && joinSide != JoinSideValues.none) {
Review Comment:
You forgot to address this comment: https://github.com/apache/spark/pull/53104/files#r2567180585
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OfflineStateRepartitionErrors.scala:
##########
@@ -85,6 +85,12 @@ object OfflineStateRepartitionErrors {
version: Int): StateRepartitionInvalidCheckpointError = {
new StateRepartitionUnsupportedOffsetSeqVersionError(checkpointLocation, version)
}
+
+ def unsupportedStateStoreProviderError(
+ checkpointLocation: String,
+ providerClass: String): StateRepartitionUnsupportedProviderError = {
Review Comment:
nit: return type `StateRepartitionInvalidCheckpointError`
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OfflineStateRepartitionErrors.scala:
##########
@@ -201,3 +207,11 @@ class StateRepartitionUnsupportedOffsetSeqVersionError(
checkpointLocation,
subClass = "UNSUPPORTED_OFFSET_SEQ_VERSION",
messageParameters = Map("version" -> version.toString))
+
+class StateRepartitionUnsupportedProviderError(
Review Comment:
still haven't fixed this indentation
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "zifeif2 (via GitHub)" <gi...@apache.org>.
zifeif2 commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2542776947
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala:
##########
@@ -60,6 +61,17 @@ object SchemaUtil {
.add("key", keySchema)
.add("value", valueSchema)
.add("partition_id", IntegerType)
+ } else if (sourceOptions.internalOnlyReadAllColumnFamilies) {
+ new StructType()
+ // todo: change this to some more specific type after we
Review Comment:
@micheal-o would like to check where the partition key extraction logic will happen? Thanks!
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "zifeif2 (via GitHub)" <gi...@apache.org>.
zifeif2 commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2540300182
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala:
##########
@@ -237,6 +245,91 @@ class StatePartitionReader(
}
}
+/**
+ * An implementation of [[StatePartitionReaderBase]] for reading all column families
+ * in binary format. This reader returns raw key and value bytes along with column family names.
+ */
+class StatePartitionReaderAllColumnFamilies(
+ storeConf: StateStoreConf,
+ hadoopConf: SerializableConfiguration,
+ partition: StateStoreInputPartition,
+ schema: StructType)
+ extends StatePartitionReaderBase(storeConf, hadoopConf, partition, schema,
+ NoPrefixKeyStateEncoderSpec(new StructType()), None, None, None, None) {
+
+ val allStateStoreMetadata = {
+ new StateMetadataPartitionReader(
+ partition.sourceOptions.resolvedCpLocation,
+ new SerializableConfiguration(hadoopConf.value),
+ partition.sourceOptions.batchId).stateMetadata.toArray
+ }
+
+ private lazy val store: ReadStateStore = {
+ assert(getStartStoreUniqueId == getEndStoreUniqueId,
+ "Start and end store unique IDs must be the same when reading all column families")
+ provider.getReadStore(
+ partition.sourceOptions.batchId + 1,
+ getStartStoreUniqueId
+ )
+ }
+
+ val colFamilyNames: Seq[String] = {
+ // todo: Support operator with multiple column family names in next PR
Review Comment:
Spark jira is taking forever to load today... will add the ticket number in the next 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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db closed pull request #53104: [SPARK-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily
URL: https://github.com/apache/spark/pull/53104
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "zifeif2 (via GitHub)" <gi...@apache.org>.
zifeif2 commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2557223014
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionAllColumnFamiliesReaderSuite.scala:
##########
@@ -0,0 +1,449 @@
+/*
+ * 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.execution.datasources.v2.state
+
+import java.nio.ByteOrder
+import java.util.Arrays
+
+import org.apache.spark.SparkUnsupportedOperationException
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.UnsafeProjection
+import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.{HDFSBackedStateStoreProvider, RocksDBStateStoreProvider}
+import org.apache.spark.sql.functions.{count, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.OutputMode
+import org.apache.spark.sql.types.{IntegerType, LongType, NullType, StructField, StructType}
+
+/**
+ * Note: This extends StateDataSourceTestBase to access
+ * helper methods like runDropDuplicatesQuery without inheriting all predefined tests.
+ */
+class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase {
+
+ import testImplicits._
+
+ override def beforeAll(): Unit = {
+ super.beforeAll()
+ spark.conf.set(SQLConf.STATE_STORE_PROVIDER_CLASS.key,
+ classOf[RocksDBStateStoreProvider].getName)
+ }
+
+ private def getNormalReadDf(checkpointDir: String): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .load()
+ .selectExpr("partition_id", "key", "value")
+ }
+
+ private def getBytesReadDf(checkpointDir: String): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .option(StateSourceOptions.INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES, "true")
+ .load()
+ .selectExpr("partition_key", "key_bytes", "value_bytes", "column_family_name")
+ }
+
+ /**
+ * Validates the schema and column families of the bytes read DataFrame.
+ */
+ private def validateBytesReadDfSchema(df: DataFrame): Unit = {
+ // Verify schema
+ val schema = df.schema
+ assert(schema.fieldNames === Array(
+ "partition_key", "key_bytes", "value_bytes", "column_family_name"))
+ assert(schema("partition_key").dataType.typeName === "struct")
+ assert(schema("key_bytes").dataType.typeName === "binary")
+ assert(schema("value_bytes").dataType.typeName === "binary")
+ assert(schema("column_family_name").dataType.typeName === "string")
+ }
+
+ /**
+ * Compares normal read data with bytes read data for a specific column family.
+ * Converts normal rows to bytes then compares with bytes read.
+ */
+ private def compareNormalAndBytesData(
+ normalDf: Array[Row],
+ bytesDf: Array[Row],
+ columnFamily: String,
+ keySchema: StructType,
+ valueSchema: StructType): Unit = {
+
+ // Filter bytes data for the specified column family and extract raw bytes directly
+ val filteredBytesData = bytesDf.filter { row =>
+ row.getString(3) == columnFamily
+ }
+
+ // Verify same number of rows
+ assert(filteredBytesData.length == normalDf.length,
+ s"Row count mismatch for column family '$columnFamily': " +
+ s"normal read has ${normalDf.length} rows, " +
+ s"bytes read has ${filteredBytesData.length} rows")
+
+ // Create projections to convert Row to UnsafeRow bytes
+ val keyProjection = UnsafeProjection.create(keySchema)
+ val valueProjection = UnsafeProjection.create(valueSchema)
+
+ // Create converters to convert external Row types to internal Catalyst types
+ val keyConverter = CatalystTypeConverters.createToCatalystConverter(keySchema)
+ val valueConverter = CatalystTypeConverters.createToCatalystConverter(valueSchema)
+
+ // Convert normal data to bytes
+ val normalAsBytes = normalDf.map { row =>
+ val key = row.getStruct(1)
+ val value = if (row.isNullAt(2)) null else row.getStruct(2)
+
+ // Convert key to InternalRow, then to UnsafeRow, then get bytes
+ val keyInternalRow = keyConverter(key).asInstanceOf[InternalRow]
+ val keyUnsafeRow = keyProjection(keyInternalRow)
+ // IMPORTANT: Must clone the bytes array since getBytes() returns a reference
+ // that may be overwritten by subsequent UnsafeRow operations
+ val keyBytes = keyUnsafeRow.getBytes.clone()
+
+ // Convert value to bytes
+ val valueBytes = if (value == null) {
+ Array.empty[Byte]
+ } else {
+ val valueInternalRow = valueConverter(value).asInstanceOf[InternalRow]
+ val valueUnsafeRow = valueProjection(valueInternalRow)
+ // IMPORTANT: Must clone the bytes array since getBytes() returns a reference
+ // that may be overwritten by subsequent UnsafeRow operations
+ valueUnsafeRow.getBytes.clone()
+ }
+
+ (keyBytes, valueBytes)
+ }
+
+ // Extract raw bytes from bytes read data (no deserialization/reserialization)
+ val bytesAsBytes = filteredBytesData.map { row =>
+ val keyBytes = row.getAs[Array[Byte]](1)
+ val valueBytes = row.getAs[Array[Byte]](2)
+ (keyBytes, valueBytes)
+ }
+
+ // Sort both for comparison (since Set equality doesn't work well with byte arrays)
+ val normalSorted = normalAsBytes.sortBy(x => (x._1.mkString(","), x._2.mkString(",")))
+ val bytesSorted = bytesAsBytes.sortBy(x => (x._1.mkString(","), x._2.mkString(",")))
+
+ assert(normalSorted.length == bytesSorted.length,
+ s"Size mismatch: normal has ${normalSorted.length}, bytes has ${bytesSorted.length}")
+
+ // Compare each pair
+ normalSorted.zip(bytesSorted).zipWithIndex.foreach {
+ case (((normalKey, normalValue), (bytesKey, bytesValue)), idx) =>
+ assert(Arrays.equals(normalKey, bytesKey),
+ s"Key mismatch at index $idx:\n" +
+ s" Normal: ${normalKey.mkString("[", ",", "]")}\n" +
+ s" Bytes: ${bytesKey.mkString("[", ",", "]")}")
+ assert(Arrays.equals(normalValue, bytesValue),
+ s"Value mismatch at index $idx:\n" +
+ s" Normal: ${normalValue.mkString("[", ",", "]")}\n" +
+ s" Bytes: ${bytesValue.mkString("[", ",", "]")}")
+ }
+ }
+
+ // Run all tests with both changelog checkpointing enabled and disabled
+ Seq(true, false).foreach { changelogCheckpointingEnabled =>
+ val testSuffix = if (changelogCheckpointingEnabled) {
+ "with changelog checkpointing"
+ } else {
+ "without changelog checkpointing"
+ }
+
+ def testWithChangelogConfig(testName: String)(testFun: => Unit): Unit = {
+ test(s"$testName ($testSuffix)") {
+ withSQLConf(
+ "spark.sql.streaming.stateStore.rocksdb.changelogCheckpointing.enabled" ->
+ changelogCheckpointingEnabled.toString) {
+ testFun
+ }
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: simple aggregation state ver 1") {
+ withSQLConf(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key -> "1") {
+ withTempDir { tempDir =>
+ runLargeDataStreamingAggregationQuery(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(StructField("groupKey", IntegerType, nullable = false)))
+ // State version 1 includes key columns in the value
+ val valueSchema = StructType(Array(
+ StructField("groupKey", IntegerType, nullable = false),
+ StructField("count", LongType, nullable = false),
+ StructField("sum", LongType, nullable = false),
+ StructField("max", IntegerType, nullable = false),
+ StructField("min", IntegerType, nullable = false)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath)
+
+ validateBytesReadDfSchema(bytesDf)
+ compareNormalAndBytesData(normalData, bytesDf.collect(), "default", keySchema, valueSchema)
+ }
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: simple aggregation state ver 2") {
+ withSQLConf(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key -> "2") {
+ withTempDir { tempDir =>
+ runLargeDataStreamingAggregationQuery(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(StructField("groupKey", IntegerType, nullable = false)))
+ val valueSchema = StructType(Array(
+ StructField("count", LongType, nullable = false),
+ StructField("sum", LongType, nullable = false),
+ StructField("max", IntegerType, nullable = false),
+ StructField("min", IntegerType, nullable = false)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath).collect()
+
+ compareNormalAndBytesData(normalData, bytesDf, "default", keySchema, valueSchema)
+ }
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: composite key aggregation state ver 1") {
+ withSQLConf(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key -> "1") {
+ withTempDir { tempDir =>
+ runCompositeKeyStreamingAggregationQuery(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(
+ StructField("groupKey", IntegerType, nullable = false),
+ StructField("fruit", org.apache.spark.sql.types.StringType, nullable = true)
+ ))
+ // State version 1 includes key columns in the value
+ val valueSchema = StructType(Array(
+ StructField("groupKey", IntegerType, nullable = false),
+ StructField("fruit", org.apache.spark.sql.types.StringType, nullable = true),
+ StructField("count", LongType, nullable = false),
+ StructField("sum", LongType, nullable = false),
+ StructField("max", IntegerType, nullable = false),
+ StructField("min", IntegerType, nullable = false)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath).collect()
+
+ compareNormalAndBytesData(normalData, bytesDf, "default", keySchema, valueSchema)
+ }
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: composite key aggregation state ver 2") {
+ withSQLConf(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key -> "2") {
+ withTempDir { tempDir =>
+ runCompositeKeyStreamingAggregationQuery(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(
+ StructField("groupKey", IntegerType, nullable = false),
+ StructField("fruit", org.apache.spark.sql.types.StringType, nullable = true)
+ ))
+ val valueSchema = StructType(Array(
+ StructField("count", LongType, nullable = false),
+ StructField("sum", LongType, nullable = false),
+ StructField("max", IntegerType, nullable = false),
+ StructField("min", IntegerType, nullable = false)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath).collect()
+
+ compareNormalAndBytesData(normalData, bytesDf, "default", keySchema, valueSchema)
+ }
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: dropDuplicates validation") {
+ withTempDir { tempDir =>
+ runDropDuplicatesQuery(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(
+ StructField("value", IntegerType, nullable = false),
+ StructField("eventTime", org.apache.spark.sql.types.TimestampType)
+ ))
+ val valueSchema = StructType(Array(
+ StructField("__dummy__", NullType, nullable = true)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath).collect()
+
+ compareNormalAndBytesData(normalData, bytesDf, "default", keySchema, valueSchema)
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: dropDuplicates with column specified") {
+ withTempDir { tempDir =>
+ runDropDuplicatesQueryWithColumnSpecified(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(
+ StructField("col1", org.apache.spark.sql.types.StringType, nullable = true)
+ ))
+ val valueSchema = StructType(Array(
+ StructField("__dummy__", NullType, nullable = true)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath).collect()
+
+ compareNormalAndBytesData(normalData, bytesDf, "default", keySchema, valueSchema)
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: dropDuplicatesWithinWatermark") {
+ withTempDir { tempDir =>
+ runDropDuplicatesWithinWatermarkQuery(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(
+ StructField("_1", org.apache.spark.sql.types.StringType, nullable = true)
+ ))
+ val valueSchema = StructType(Array(
+ StructField("expiresAtMicros", LongType, nullable = false)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath).collect()
+
+ compareNormalAndBytesData(normalData, bytesDf, "default", keySchema, valueSchema)
+ }
+ }
+
+ testWithChangelogConfig("all-column-families: session window aggregation") {
+ withTempDir { tempDir =>
+ runSessionWindowAggregationQuery(tempDir.getAbsolutePath)
Review Comment:
will fix indentation
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "zifeif2 (via GitHub)" <gi...@apache.org>.
zifeif2 commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2557236283
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala:
##########
@@ -50,7 +50,16 @@ class StatePartitionReaderFactory(
override def createReader(partition: InputPartition): PartitionReader[InternalRow] = {
val stateStoreInputPartition = partition.asInstanceOf[StateStoreInputPartition]
if (stateStoreInputPartition.sourceOptions.internalOnlyReadAllColumnFamilies) {
- new StatePartitionReaderAllColumnFamilies(storeConf, hadoopConf,
+ // Disable format validation because the schema returned by
+ // StatePartitionAllColumnFamiliesReader does not contain the corresponding
+ // keySchema or valueSchema.
+ // It's safe to do so we also don't expect the caller of StatePartitionAllColumnFamiliesReader
+ // to extract specific fields out of the returning row.
+ val modifiedStoreConf = storeConf.withExtraOptions(Map(
+ StateStoreConf.FORMAT_VALIDATION_ENABLED_CONFIG -> "false",
Review Comment:
this version disable format validation so that we don't need the keySchema and valueSchema
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "micheal-o (via GitHub)" <gi...@apache.org>.
micheal-o commented on PR #53104:
URL: https://github.com/apache/spark/pull/53104#issuecomment-3573993761
Also CI is failing due to linter error for your changes. PTAL
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "zifeif2 (via GitHub)" <gi...@apache.org>.
zifeif2 commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2540249132
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreConf.scala:
##########
@@ -163,6 +163,26 @@ class StateStoreConf(
*/
val sqlConfs: Map[String, String] =
sqlConf.getAllConfs.filter(_._1.startsWith("spark.sql.streaming.stateStore."))
+
+ /**
+ * Creates a copy of this StateStoreConf with format validation disabled.
+ * This is useful when reading raw bytes where the schema used (binary) doesn't match
+ * the actual stored data schema.
+ */
+ def withFormatValidationDisabled(): StateStoreConf = {
Review Comment:
Was written for HDFS but since we don't support it, I'll remove it
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "micheal-o (via GitHub)" <gi...@apache.org>.
micheal-o commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2582784549
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala:
##########
@@ -407,6 +418,7 @@ object StateSourceOptions extends DataSourceOptions {
val STATE_VAR_NAME = newOption("stateVarName")
val READ_REGISTERED_TIMERS = newOption("readRegisteredTimers")
val FLATTEN_COLLECTION_TYPES = newOption("flattenCollectionTypes")
+ val INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES = newOption("internalOnlyReadAllColumnFamilies")
Review Comment:
@anishshri-db it allows us to reuse alot of the existing state data source code for impl and testing. So we don't have to duplicate code. It also allows dataframe transformation on the data easier.
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "micheal-o (via GitHub)" <gi...@apache.org>.
micheal-o commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2567181951
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala:
##########
@@ -78,19 +83,39 @@ abstract class StatePartitionReaderBase(
extends PartitionReader[InternalRow] with Logging {
// Used primarily as a placeholder for the value schema in the context of
// state variables used within the transformWithState operator.
- private val schemaForValueRow: StructType =
+ // Also used as a placeholder for both key and value schema for
+ // StatePartitionAllColumnFamiliesReader
+ private val placeholderSchema: StructType =
StructType(Array(StructField("__dummy__", NullType)))
+ private val colFamilyToSchema : mutable.HashMap[String, StateStoreColFamilySchema] = {
+ val stateStoreId = StateStoreId(
+ partition.sourceOptions.stateCheckpointLocation.toString,
+ partition.sourceOptions.operatorId,
+ StateStore.PARTITION_ID_TO_CHECK_SCHEMA,
+ partition.sourceOptions.storeName)
+ val stateStoreProviderId = StateStoreProviderId(stateStoreId, partition.queryId)
+ val manager = new StateSchemaCompatibilityChecker(stateStoreProviderId, hadoopConf.value)
Review Comment:
Why duplicate the code from `getStoreMetadataAndRunChecks` here? We already read the schema file in the StateDataSource i.e. in `inferSchema` and `getStoreMetadataAndRunChecks` methods, and we pass it in here via `stateStoreColFamilySchemaOpt` param. So we can rely on stateStoreColFamilySchemaOpt here.
Also, since this PR is for single CF, then lets focus on single CF for now. In the subsequent PR, we can then pass in schema for all CFs.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OfflineStateRepartitionErrors.scala:
##########
@@ -201,3 +208,10 @@ class StateRepartitionUnsupportedOffsetSeqVersionError(
checkpointLocation,
subClass = "UNSUPPORTED_OFFSET_SEQ_VERSION",
messageParameters = Map("version" -> version.toString))
+
+class StateRepartitionUnsupportedProviderError(
+ checkpointLocation: String,
+ provider: String) extends StateRepartitionInvalidCheckpointError(
Review Comment:
Please reformat. See how the others are formatted. Please make sure you run scalafmt before submitting for review ([doc](https://spark.apache.org/developer-tools.html#:~:text=the%20style%20guide.-,Formatting%20code,-To%20format%20Scala))
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionAllColumnFamiliesReaderSuite.scala:
##########
@@ -0,0 +1,523 @@
+/*
+ * 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.execution.datasources.v2.state
+
+import java.nio.ByteOrder
+import java.util.Arrays
+
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.UnsafeProjection
+import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.{HDFSBackedStateStoreProvider, RocksDBStateStoreProvider, StateRepartitionUnsupportedProviderError, StateStore}
+import org.apache.spark.sql.functions.{count, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.OutputMode
+import org.apache.spark.sql.types.{BooleanType, IntegerType, LongType, NullType, StructField, StructType, TimestampType}
+
+/**
+ * Note: This extends StateDataSourceTestBase to access
+ * helper methods like runDropDuplicatesQuery without inheriting all predefined tests.
+ */
+class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase {
+
+ import testImplicits._
+
+ override def beforeAll(): Unit = {
+ super.beforeAll()
+ spark.conf.set(SQLConf.STATE_STORE_PROVIDER_CLASS.key,
+ classOf[RocksDBStateStoreProvider].getName)
+ }
+
+ private def getNormalReadDf(
+ checkpointDir: String,
+ storeName: Option[String] = Option.empty[String]): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .option(StateSourceOptions.STORE_NAME, storeName.orNull)
+ .load()
+ .selectExpr("partition_id", "key", "value")
+ }
+
+ private def getBytesReadDf(
+ checkpointDir: String,
+ storeName: Option[String] = Option.empty[String]): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .option(StateSourceOptions.INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES, "true")
+ .option(StateSourceOptions.STORE_NAME, storeName.orNull)
+ .load()
+ .selectExpr("partition_key", "key_bytes", "value_bytes", "column_family_name")
+ }
+
+ /**
+ * Validates the schema and column families of the bytes read DataFrame.
+ */
+ private def validateBytesReadDfSchema(df: DataFrame): Unit = {
+ // Verify schema
+ val schema = df.schema
+ assert(schema.fieldNames === Array(
+ "partition_key", "key_bytes", "value_bytes", "column_family_name"))
+ assert(schema("partition_key").dataType.typeName === "struct")
+ assert(schema("key_bytes").dataType.typeName === "binary")
+ assert(schema("value_bytes").dataType.typeName === "binary")
+ assert(schema("column_family_name").dataType.typeName === "string")
+ }
+
+ /**
+ * Compares normal read data with bytes read data for a specific column family.
+ * Converts normal rows to bytes then compares with bytes read.
+ */
+ private def compareNormalAndBytesData(
+ normalDf: Array[Row],
+ bytesDf: Array[Row],
+ columnFamily: String,
+ keySchema: StructType,
+ valueSchema: StructType): Unit = {
+
+ // Filter bytes data for the specified column family and extract raw bytes directly
+ val filteredBytesData = bytesDf.filter { row =>
+ row.getString(3) == columnFamily
+ }
+
+ // Verify same number of rows
+ assert(filteredBytesData.length == normalDf.length,
+ s"Row count mismatch for column family '$columnFamily': " +
+ s"normal read has ${normalDf.length} rows, " +
+ s"bytes read has ${filteredBytesData.length} rows")
+
+ // Create projections to convert Row to UnsafeRow bytes
+ val keyProjection = UnsafeProjection.create(keySchema)
+ val valueProjection = UnsafeProjection.create(valueSchema)
+
+ // Create converters to convert external Row types to internal Catalyst types
+ val keyConverter = CatalystTypeConverters.createToCatalystConverter(keySchema)
+ val valueConverter = CatalystTypeConverters.createToCatalystConverter(valueSchema)
+
+ // Convert normal data to bytes
+ val normalAsBytes = normalDf.toSeq.map { row =>
+ val key = row.getStruct(1)
+ val value = if (row.isNullAt(2)) null else row.getStruct(2)
+
+ // Convert key to InternalRow, then to UnsafeRow, then get bytes
+ val keyInternalRow = keyConverter(key).asInstanceOf[InternalRow]
+ val keyUnsafeRow = keyProjection(keyInternalRow)
+ // IMPORTANT: Must clone the bytes array since getBytes() returns a reference
+ // that may be overwritten by subsequent UnsafeRow operations
+ val keyBytes = keyUnsafeRow.getBytes.clone()
+
+ // Convert value to bytes
+ val valueBytes = if (value == null) {
+ Array.empty[Byte]
+ } else {
+ val valueInternalRow = valueConverter(value).asInstanceOf[InternalRow]
+ val valueUnsafeRow = valueProjection(valueInternalRow)
+ // IMPORTANT: Must clone the bytes array since getBytes() returns a reference
+ // that may be overwritten by subsequent UnsafeRow operations
+ valueUnsafeRow.getBytes.clone()
+ }
+
+ (keyBytes, valueBytes)
+ }
+
+ // Extract raw bytes from bytes read data (no deserialization/reserialization)
+ val bytesAsBytes = filteredBytesData.map { row =>
+ val keyBytes = row.getAs[Array[Byte]](1)
+ val valueBytes = row.getAs[Array[Byte]](2)
+ (keyBytes, valueBytes)
+ }
+
+ // Sort both for comparison (since Set equality doesn't work well with byte arrays)
+ val normalSorted = normalAsBytes.sortBy(x => (x._1.mkString(","), x._2.mkString(",")))
+ val bytesSorted = bytesAsBytes.sortBy(x => (x._1.mkString(","), x._2.mkString(",")))
+
+ assert(normalSorted.length == bytesSorted.length,
+ s"Size mismatch: normal has ${normalSorted.length}, bytes has ${bytesSorted.length}")
+
+ // Compare each pair
+ normalSorted.zip(bytesSorted).zipWithIndex.foreach {
+ case (((normalKey, normalValue), (bytesKey, bytesValue)), idx) =>
+ assert(Arrays.equals(normalKey, bytesKey),
+ s"Key mismatch at index $idx:\n" +
+ s" Normal: ${normalKey.mkString("[", ",", "]")}\n" +
+ s" Bytes: ${bytesKey.mkString("[", ",", "]")}")
+ assert(Arrays.equals(normalValue, bytesValue),
+ s"Value mismatch at index $idx:\n" +
+ s" Normal: ${normalValue.mkString("[", ",", "]")}\n" +
+ s" Bytes: ${bytesValue.mkString("[", ",", "]")}")
+ }
+ }
+
+ // Run all tests with both changelog checkpointing enabled and disabled
+ Seq(true, false).foreach { changelogCheckpointingEnabled =>
+ val testSuffix = if (changelogCheckpointingEnabled) {
+ "with changelog checkpointing"
+ } else {
+ "without changelog checkpointing"
+ }
+
+ def testWithChangelogConfig(testName: String)(testFun: => Unit): Unit = {
+ test(s"$testName ($testSuffix)") {
+ withSQLConf(
+ "spark.sql.streaming.stateStore.rocksdb.changelogCheckpointing.enabled" ->
+ changelogCheckpointingEnabled.toString) {
+ testFun
+ }
+ }
+ }
+
+ testWithChangelogConfig("SPARK-54388: simple aggregation state ver 1") {
+ withSQLConf(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key -> "1") {
+ withTempDir { tempDir =>
+ runLargeDataStreamingAggregationQuery(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(StructField("groupKey", IntegerType, nullable = false)))
+ // State version 1 includes key columns in the value
+ val valueSchema = StructType(Array(
+ StructField("groupKey", IntegerType, nullable = false),
+ StructField("count", LongType, nullable = false),
+ StructField("sum", LongType, nullable = false),
+ StructField("max", IntegerType, nullable = false),
+ StructField("min", IntegerType, nullable = false)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath)
+
+ validateBytesReadDfSchema(bytesDf)
+ compareNormalAndBytesData(normalData, bytesDf.collect(), "default", keySchema, valueSchema)
+ }
+ }
+ }
+
+ testWithChangelogConfig("SPARK-54388: simple aggregation state ver 2") {
+ withSQLConf(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key -> "2") {
+ withTempDir { tempDir =>
+ runLargeDataStreamingAggregationQuery(tempDir.getAbsolutePath)
+
+ val keySchema = StructType(Array(StructField("groupKey", IntegerType, nullable = false)))
+ val valueSchema = StructType(Array(
+ StructField("count", LongType, nullable = false),
+ StructField("sum", LongType, nullable = false),
+ StructField("max", IntegerType, nullable = false),
+ StructField("min", IntegerType, nullable = false)
+ ))
+
+ val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
+ val bytesDf = getBytesReadDf(tempDir.getAbsolutePath).collect()
+
+ compareNormalAndBytesData(normalData, bytesDf, "default", keySchema, valueSchema)
Review Comment:
We are not validating the bytesDf i.e. `validateBytesReadDfSchema`. Here and the other test cases below
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala:
##########
@@ -492,6 +503,29 @@ object StateSourceOptions extends DataSourceOptions {
val readChangeFeed = Option(options.get(READ_CHANGE_FEED)).exists(_.toBoolean)
+ val internalOnlyReadAllColumnFamilies = try {
+ Option(options.get(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES)).exists(_.toBoolean)
+ } catch {
+ case _: IllegalArgumentException =>
+ throw StateDataSourceErrors.invalidOptionValue(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES,
+ "Boolean value is expected")
+ }
+
+ if (internalOnlyReadAllColumnFamilies && stateVarName.isDefined) {
+ throw StateDataSourceErrors.conflictOptions(
+ Seq(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES, STATE_VAR_NAME))
+ }
+
+ if (internalOnlyReadAllColumnFamilies && joinSide != JoinSideValues.none) {
Review Comment:
nit: add comment that for this option, we use `storeName` to specify the join store instead
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionAllColumnFamiliesReaderSuite.scala:
##########
@@ -0,0 +1,523 @@
+/*
+ * 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.execution.datasources.v2.state
+
+import java.nio.ByteOrder
+import java.util.Arrays
+
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.UnsafeProjection
+import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.{HDFSBackedStateStoreProvider, RocksDBStateStoreProvider, StateRepartitionUnsupportedProviderError, StateStore}
+import org.apache.spark.sql.functions.{count, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.OutputMode
+import org.apache.spark.sql.types.{BooleanType, IntegerType, LongType, NullType, StructField, StructType, TimestampType}
+
+/**
+ * Note: This extends StateDataSourceTestBase to access
+ * helper methods like runDropDuplicatesQuery without inheriting all predefined tests.
+ */
+class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase {
+
+ import testImplicits._
+
+ override def beforeAll(): Unit = {
+ super.beforeAll()
+ spark.conf.set(SQLConf.STATE_STORE_PROVIDER_CLASS.key,
+ classOf[RocksDBStateStoreProvider].getName)
+ }
+
+ private def getNormalReadDf(
+ checkpointDir: String,
+ storeName: Option[String] = Option.empty[String]): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .option(StateSourceOptions.STORE_NAME, storeName.orNull)
+ .load()
+ .selectExpr("partition_id", "key", "value")
+ }
+
+ private def getBytesReadDf(
+ checkpointDir: String,
+ storeName: Option[String] = Option.empty[String]): DataFrame = {
+ spark.read
+ .format("statestore")
+ .option(StateSourceOptions.PATH, checkpointDir)
+ .option(StateSourceOptions.INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES, "true")
+ .option(StateSourceOptions.STORE_NAME, storeName.orNull)
+ .load()
+ .selectExpr("partition_key", "key_bytes", "value_bytes", "column_family_name")
Review Comment:
nit: select is not needed right
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/OfflineStateRepartitionErrors.scala:
##########
@@ -201,3 +208,10 @@ class StateRepartitionUnsupportedOffsetSeqVersionError(
checkpointLocation,
subClass = "UNSUPPORTED_OFFSET_SEQ_VERSION",
messageParameters = Map("version" -> version.toString))
+
+class StateRepartitionUnsupportedProviderError(
Review Comment:
please fix indentation
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "micheal-o (via GitHub)" <gi...@apache.org>.
micheal-o commented on PR #53104:
URL: https://github.com/apache/spark/pull/53104#issuecomment-3584221268
@zifeif2 please also fix the CI failure.
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2582763965
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala:
##########
@@ -492,6 +505,29 @@ object StateSourceOptions extends DataSourceOptions {
val readChangeFeed = Option(options.get(READ_CHANGE_FEED)).exists(_.toBoolean)
+ val internalOnlyReadAllColumnFamilies = try {
+ Option(options.get(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES)).exists(_.toBoolean)
+ } catch {
+ case _: IllegalArgumentException =>
+ throw StateDataSourceErrors.invalidOptionValue(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES,
+ "Boolean value is expected")
+ }
+
+ if (internalOnlyReadAllColumnFamilies && stateVarName.isDefined) {
+ throw StateDataSourceErrors.conflictOptions(
+ Seq(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES, STATE_VAR_NAME))
+ }
+
+ if (internalOnlyReadAllColumnFamilies && joinSide != JoinSideValues.none) {
+ throw StateDataSourceErrors.conflictOptions(
+ Seq(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES, JOIN_SIDE))
+ }
+
+ if (internalOnlyReadAllColumnFamilies && readChangeFeed) {
Review Comment:
Prob don't need this if we dont expose an option ?
--
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-54388][SS] Introduce StatePartitionReader that scan raw bytes for Single ColFamily [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #53104:
URL: https://github.com/apache/spark/pull/53104#discussion_r2582806450
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala:
##########
@@ -492,6 +505,29 @@ object StateSourceOptions extends DataSourceOptions {
val readChangeFeed = Option(options.get(READ_CHANGE_FEED)).exists(_.toBoolean)
+ val internalOnlyReadAllColumnFamilies = try {
+ Option(options.get(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES)).exists(_.toBoolean)
+ } catch {
+ case _: IllegalArgumentException =>
+ throw StateDataSourceErrors.invalidOptionValue(INTERNAL_ONLY_READ_ALL_COLUMN_FAMILIES,
+ "Boolean value is expected")
+ }
+
+ if (internalOnlyReadAllColumnFamilies && stateVarName.isDefined) {
Review Comment:
Lets also note here that this path should only be invoked by internal callers - for eg - repartitioning
--
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