You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "ericm-db (via GitHub)" <gi...@apache.org> on 2024/11/23 21:36:54 UTC
[PR] [WIP] init [spark]
ericm-db opened a new pull request, #48944:
URL: https://github.com/apache/spark/pull/48944
<!--
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'.
-->
### 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.
-->
### 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.
-->
### Does this PR introduce _any_ user-facing change?
<!--
Note that it means *any* user-facing change including all aspects such as the documentation fix.
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'.
-->
### 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.
-->
### 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.
-->
--
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-50017][SQL] Support Avro encoding for TransformWithState operator [spark]
Posted by "brkyvz (via GitHub)" <gi...@apache.org>.
brkyvz commented on code in PR #48944:
URL: https://github.com/apache/spark/pull/48944#discussion_r1887799446
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala:
##########
@@ -60,15 +60,92 @@ sealed trait RocksDBValueStateEncoder {
* by the callers. The metadata in each row does not need to be written as Avro or UnsafeRow,
* but the actual data provided by the caller does.
*/
+/** Interface for encoding and decoding state store data between UnsafeRow and raw bytes.
+ *
+ * @note All encode methods expect non-null input rows. Handling of null values is left to the
+ * implementing classes.
+ */
trait DataEncoder {
+ /** Encodes a complete key row into bytes. Used as the primary key for state lookups.
+ *
+ * @param row An UnsafeRow containing all key columns as defined in the keySchema
+ * @return Serialized byte array representation of the key
+ */
def encodeKey(row: UnsafeRow): Array[Byte]
+
+ /** Encodes the non-prefix portion of a key row. Used with prefix scan and
Review Comment:
Text should start at the second line, not the first in scaladoc
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala:
##########
@@ -356,6 +376,14 @@ case class PrefixKeyScanStateEncoderSpec(
if (numColsPrefixKey == 0 || numColsPrefixKey >= keySchema.length) {
throw StateStoreErrors.incorrectNumOrderingColsForPrefixScan(numColsPrefixKey.toString)
}
+ override def toEncoder(
Review Comment:
nit: new line please
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala:
##########
@@ -789,44 +940,58 @@ abstract class RocksDBKeyStateEncoderBase(
}
}
+/** Factory object for creating state encoders used by RocksDB state store.
Review Comment:
move text to next line
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala:
##########
@@ -402,6 +479,80 @@ class AvroStateEncoder(
private lazy val remainingKeyAvroProjection = UnsafeProjection.create(remainingKeySchema)
+
Review Comment:
uber nit: delete extra empty lines
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala:
##########
@@ -402,6 +479,80 @@ class AvroStateEncoder(
private lazy val remainingKeyAvroProjection = UnsafeProjection.create(remainingKeySchema)
+
+
+ private def getAvroSerializer(schema: StructType): AvroSerializer = {
+ val avroType = SchemaConverters.toAvroType(schema)
+ new AvroSerializer(schema, avroType, nullable = false)
+ }
+
+ private def getAvroDeserializer(schema: StructType): AvroDeserializer = {
+ val avroType = SchemaConverters.toAvroType(schema)
+ val avroOptions = AvroOptions(Map.empty)
+ new AvroDeserializer(avroType, schema,
+ avroOptions.datetimeRebaseModeInRead, avroOptions.useStableIdForUnionType,
+ avroOptions.stableIdPrefixForUnionType, avroOptions.recursiveFieldMaxDepth)
+ }
+
+ /**
+ * Creates an AvroEncoder that handles both key and value serialization/deserialization.
+ * This method sets up the complete encoding infrastructure needed for state store operations.
+ *
+ * The encoder handles different key encoding specifications:
+ * - NoPrefixKeyStateEncoderSpec: Simple key encoding without prefix
+ * - PrefixKeyScanStateEncoderSpec: Keys with prefix for efficient scanning
+ * - RangeKeyScanStateEncoderSpec: Keys with ordering requirements for range scans
+ *
+ * For prefix scan cases, it also creates separate encoders for the suffix portion of keys.
+ *
+ * @param keyStateEncoderSpec Specification for how to encode keys
+ * @param valueSchema Schema for the values to be encoded
+ * @return An AvroEncoder containing all necessary serializers and deserializers
+ */
+ private def createAvroEnc(
+ keyStateEncoderSpec: KeyStateEncoderSpec,
+ valueSchema: StructType
+ ): AvroEncoder = {
Review Comment:
please fix the indent
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala:
##########
@@ -325,6 +325,18 @@ sealed trait KeyStateEncoderSpec {
def keySchema: StructType
def jsonValue: JValue
def json: String = compact(render(jsonValue))
+
+ /** Creates a RocksDBKeyStateEncoder for this specification.
Review Comment:
ditto on scaladoc
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala:
##########
@@ -60,15 +60,92 @@ sealed trait RocksDBValueStateEncoder {
* by the callers. The metadata in each row does not need to be written as Avro or UnsafeRow,
* but the actual data provided by the caller does.
*/
+/** Interface for encoding and decoding state store data between UnsafeRow and raw bytes.
+ *
+ * @note All encode methods expect non-null input rows. Handling of null values is left to the
+ * implementing classes.
+ */
Review Comment:
is this in the right place?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceTransformWithStateSuite.scala:
##########
@@ -126,8 +126,7 @@ class SessionGroupsStatefulProcessorWithTTL extends
* Test suite to verify integration of state data source reader with the transformWithState operator
*/
class StateDataSourceTransformWithStateSuite extends StateStoreMetricsTest
- with AlsoTestWithEncodingTypes
- with AlsoTestWithChangelogCheckpointingEnabled {
Review Comment:
did you lose the `AlsoTestWithChangelogCheckpointingEnabled`? Does `AlsoTestWithRocksDBFeatures` contain that?
--
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-50017][SQL] Support Avro encoding for TransformWithState operator [spark]
Posted by "brkyvz (via GitHub)" <gi...@apache.org>.
brkyvz commented on code in PR #48944:
URL: https://github.com/apache/spark/pull/48944#discussion_r1884513211
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala:
##########
@@ -76,17 +76,16 @@ private[sql] class RocksDBStateStoreProvider
isInternal: Boolean = false): Unit = {
verifyColFamilyCreationOrDeletion("create_col_family", colFamilyName, isInternal)
val newColFamilyId = rocksDB.createColFamilyIfAbsent(colFamilyName)
- // Create cache key using store ID to avoid collisions
- val avroEncCacheKey = s"${getRunId(hadoopConf)}_${stateStoreId.operatorId}_" +
- s"${stateStoreId.partitionId}_$colFamilyName"
+ val dataEncoderCacheKey = s"${getRunId(hadoopConf)}_${stateStoreId.operatorId}_" +
+ s"${stateStoreId.partitionId}_${colFamilyName}"
- val avroEnc = getAvroEnc(
- stateStoreEncoding, avroEncCacheKey, keyStateEncoderSpec, valueSchema)
+ val dataEncoder = getDataEncoder(
+ stateStoreEncoding, dataEncoderCacheKey, keyStateEncoderSpec, valueSchema)
keyValueEncoderMap.putIfAbsent(colFamilyName,
- (RocksDBStateEncoder.getKeyEncoder(keyStateEncoderSpec, useColumnFamilies,
- Some(newColFamilyId), avroEnc), RocksDBStateEncoder.getValueEncoder(valueSchema,
- useMultipleValuesPerKey, avroEnc)))
+ (RocksDBStateEncoder.getKeyEncoder(dataEncoder, keyStateEncoderSpec, useColumnFamilies,
+ Some(newColFamilyId)), RocksDBStateEncoder.getValueEncoder(dataEncoder, valueSchema,
+ useMultipleValuesPerKey)))
Review Comment:
nit: can you make these one parameter per line please? maybe define them outside of the putIfAbsent first
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala:
##########
@@ -51,93 +51,56 @@ sealed trait RocksDBValueStateEncoder {
def decodeValues(valueBytes: Array[Byte]): Iterator[UnsafeRow]
}
-abstract class RocksDBKeyStateEncoderBase(
- useColumnFamilies: Boolean,
- virtualColFamilyId: Option[Short] = None) extends RocksDBKeyStateEncoder {
- def offsetForColFamilyPrefix: Int =
- if (useColumnFamilies) VIRTUAL_COL_FAMILY_PREFIX_BYTES else 0
-
- val out = new ByteArrayOutputStream
- /**
- * Get Byte Array for the virtual column family id that is used as prefix for
- * key state rows.
- */
- override def getColumnFamilyIdBytes(): Array[Byte] = {
- assert(useColumnFamilies, "Cannot return virtual Column Family Id Bytes" +
- " because multiple Column is not supported for this encoder")
- val encodedBytes = new Array[Byte](VIRTUAL_COL_FAMILY_PREFIX_BYTES)
- Platform.putShort(encodedBytes, Platform.BYTE_ARRAY_OFFSET, virtualColFamilyId.get)
- encodedBytes
- }
-
- /**
- * Encode and put column family Id as a prefix to a pre-allocated byte array.
- *
- * @param numBytes - size of byte array to be created for storing key row (without
- * column family prefix)
- * @return Array[Byte] for an array byte to put encoded key bytes
- * Int for a starting offset to put the encoded key bytes
- */
- protected def encodeColumnFamilyPrefix(numBytes: Int): (Array[Byte], Int) = {
- val encodedBytes = new Array[Byte](numBytes + offsetForColFamilyPrefix)
- var offset = Platform.BYTE_ARRAY_OFFSET
- if (useColumnFamilies) {
- Platform.putShort(encodedBytes, Platform.BYTE_ARRAY_OFFSET, virtualColFamilyId.get)
- offset = Platform.BYTE_ARRAY_OFFSET + offsetForColFamilyPrefix
- }
- (encodedBytes, offset)
- }
+/**
+ * The DataEncoder can encode UnsafeRows into raw bytes in two ways:
+ * - Using the direct byte layout of the UnsafeRow
+ * - Converting the UnsafeRow into an Avro row, and encoding that
+ * In both of these cases, the raw bytes that are written into RockDB have
+ * headers, footers and other metadata, but they also have data that is provided
+ * by the callers. The metadata in each row does not need to be written as Avro or UnsafeRow,
+ * but the actual data provided by the caller does.
+ */
+trait DataEncoder {
+ def encodeKey(row: UnsafeRow): Array[Byte]
+ def encodeRemainingKey(row: UnsafeRow): Array[Byte]
+ def encodePrefixKeyForRangeScan(row: UnsafeRow): Array[Byte]
+ def encodeValue(row: UnsafeRow): Array[Byte]
- /**
- * Get starting offset for decoding an encoded key byte array.
- */
- protected def decodeKeyStartOffset: Int = {
- if (useColumnFamilies) {
- Platform.BYTE_ARRAY_OFFSET + VIRTUAL_COL_FAMILY_PREFIX_BYTES
- } else Platform.BYTE_ARRAY_OFFSET
- }
+ def decodeKey(bytes: Array[Byte]): UnsafeRow
+ def decodeRemainingKey(bytes: Array[Byte]): UnsafeRow
+ def decodePrefixKeyForRangeScan(bytes: Array[Byte]): UnsafeRow
+ def decodeValue(bytes: Array[Byte]): UnsafeRow
Review Comment:
ditto
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala:
##########
@@ -615,29 +613,56 @@ object RocksDBStateStoreProvider {
private val AVRO_ENCODER_LIFETIME_HOURS = 1L
// Add the cache at companion object level so it persists across provider instances
- private val avroEncoderMap: NonFateSharingCache[String, AvroEncoder] =
+ private val dataEncoderCache: NonFateSharingCache[String, RocksDBDataEncoder] =
NonFateSharingCache(
maximumSize = MAX_AVRO_ENCODERS_IN_CACHE,
expireAfterAccessTime = AVRO_ENCODER_LIFETIME_HOURS,
expireAfterAccessTimeUnit = TimeUnit.HOURS
)
- def getAvroEnc(
+ /**
+ * Creates and returns a data encoder for the state store based on the specified encoding type.
+ * This method handles caching of encoders to improve performance by reusing encoder instances
+ * when possible.
+ *
+ * The method supports two encoding types:
+ * - Avro: Uses Apache Avro for serialization with schema evolution support
+ * - UnsafeRow: Uses Spark's internal row format for optimal performance
+ *
+ * @param stateStoreEncoding The encoding type to use ("avro" or "unsaferow")
+ * @param encoderCacheKey A unique key for caching the encoder instance, typically combining
+ * query ID, operator ID, partition ID, and column family name
+ * @param keyStateEncoderSpec Specification for how to encode keys, including schema and any
+ * prefix/range scan requirements
+ * @param valueSchema The schema for the values to be encoded
+ * @return A RocksDBDataEncoder instance configured for the specified encoding type
+ */
+ def getDataEncoder(
stateStoreEncoding: String,
- avroEncCacheKey: String,
+ encoderCacheKey: String,
keyStateEncoderSpec: KeyStateEncoderSpec,
- valueSchema: StructType): Option[AvroEncoder] = {
+ valueSchema: StructType): RocksDBDataEncoder = {
stateStoreEncoding match {
- case "avro" => Some(
- RocksDBStateStoreProvider.avroEncoderMap.get(
- avroEncCacheKey,
- new java.util.concurrent.Callable[AvroEncoder] {
- override def call(): AvroEncoder = createAvroEnc(keyStateEncoderSpec, valueSchema)
+ case "avro" =>
+ RocksDBStateStoreProvider.dataEncoderCache.get(
+ encoderCacheKey,
+ new java.util.concurrent.Callable[AvroStateEncoder] {
+ override def call(): AvroStateEncoder = {
+ val avroEncoder = createAvroEnc(keyStateEncoderSpec, valueSchema)
+ new AvroStateEncoder(keyStateEncoderSpec, valueSchema, avroEncoder)
+ }
+ }
+ )
+ case "unsaferow" =>
+ RocksDBStateStoreProvider.dataEncoderCache.get(
+ encoderCacheKey,
+ new java.util.concurrent.Callable[UnsafeRowDataEncoder] {
+ override def call(): UnsafeRowDataEncoder = {
+ new UnsafeRowDataEncoder(keyStateEncoderSpec, valueSchema)
+ }
Review Comment:
nit: maybe change the order of the match-case with the cache.get?
```scala
assert(Set("avro", "unsaferow").contains(stateStoreEncoding))
RocksDBStateStoreProvider.dataEncoderCache.get(
encoderCacheKey,
new java.util.concurrent.Callable[DataEncoder] { () =>
if (stateStoreEncoding == "avro") {
...
} else {
...
}
}
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala:
##########
@@ -800,156 +583,509 @@ class RangeKeyScanStateEncoder(
})
record.put(fieldIdx, ByteBuffer.wrap(markerBuffer.array()))
- val valueBuffer = ByteBuffer.allocate(8)
- valueBuffer.order(ByteOrder.BIG_ENDIAN)
- if ((rawBits & doubleSignBitMask) != 0) {
- val updatedVal = rawBits ^ doubleFlipBitMask
- valueBuffer.putDouble(longBitsToDouble(updatedVal))
- } else {
- valueBuffer.putDouble(doubleVal)
- }
- record.put(fieldIdx + 1, ByteBuffer.wrap(valueBuffer.array()))
+ val valueBuffer = ByteBuffer.allocate(8)
+ valueBuffer.order(ByteOrder.BIG_ENDIAN)
+ if ((rawBits & doubleSignBitMask) != 0) {
+ val updatedVal = rawBits ^ doubleFlipBitMask
+ valueBuffer.putDouble(longBitsToDouble(updatedVal))
+ } else {
+ valueBuffer.putDouble(doubleVal)
+ }
+ record.put(fieldIdx + 1, ByteBuffer.wrap(valueBuffer.array()))
+
+ case _ => throw new UnsupportedOperationException(
+ s"Range scan encoding not supported for data type: ${field.dataType}")
+ }
+ }
+ fieldIdx += 2
+ }
+
+ out.reset()
+ val writer = new GenericDatumWriter[GenericRecord](rangeScanAvroType)
+ val encoder = EncoderFactory.get().binaryEncoder(out, null)
+ writer.write(record, encoder)
+ encoder.flush()
+ out.toByteArray
+ }
+
+ override def encodeValue(row: UnsafeRow): Array[Byte] =
+ encodeUnsafeRowToAvro(row, avroEncoder.valueSerializer, valueAvroType, out)
+
+ override def decodeKey(bytes: Array[Byte]): UnsafeRow = {
+ keyStateEncoderSpec match {
+ case NoPrefixKeyStateEncoderSpec(_) =>
+ decodeFromAvroToUnsafeRow(bytes, avroEncoder.keyDeserializer, keyAvroType, keyProj)
+ case PrefixKeyScanStateEncoderSpec(_, _) =>
+ decodeFromAvroToUnsafeRow(
+ bytes, avroEncoder.keyDeserializer, prefixKeyAvroType, prefixKeyProj)
+ case _ => throw unsupportedOperationForKeyStateEncoder("decodeKey")
+ }
+ }
+
+
+ override def decodeRemainingKey(bytes: Array[Byte]): UnsafeRow = {
+ keyStateEncoderSpec match {
+ case PrefixKeyScanStateEncoderSpec(_, _) =>
+ decodeFromAvroToUnsafeRow(bytes,
+ avroEncoder.suffixKeyDeserializer.get, remainingKeyAvroType, remainingKeyAvroProjection)
+ case RangeKeyScanStateEncoderSpec(_, _) =>
+ decodeFromAvroToUnsafeRow(
+ bytes, avroEncoder.keyDeserializer, remainingKeyAvroType, remainingKeyAvroProjection)
+ case _ => throw unsupportedOperationForKeyStateEncoder("decodeRemainingKey")
+ }
+ }
+
+ /**
+ * Decodes an Avro-encoded byte array back into an UnsafeRow for range scan operations.
+ *
+ * This method reverses the encoding process performed by encodePrefixKeyForRangeScan:
+ * - Reads the marker byte to determine null status or sign
+ * - Reconstructs the original values from big-endian format
+ * - Handles special cases for floating point numbers by reversing bit manipulations
+ *
+ * The decoding process preserves the original data types and values, including:
+ * - Null values marked by nullValMarker
+ * - Sign information for numeric types
+ * - Proper restoration of negative floating point values
+ *
+ * @param bytes The Avro-encoded byte array to decode
+ * @param avroType The Avro schema defining the structure for decoding
+ * @return UnsafeRow containing the decoded data
+ * @throws UnsupportedOperationException if a field's data type is not supported for range
+ * scan decoding
+ */
+ override def decodePrefixKeyForRangeScan(bytes: Array[Byte]): UnsafeRow = {
+ val reader = new GenericDatumReader[GenericRecord](rangeScanAvroType)
+ val decoder = DecoderFactory.get().binaryDecoder(bytes, 0, bytes.length, null)
+ val record = reader.read(null, decoder)
+
+ val rowWriter = new UnsafeRowWriter(rangeScanKeyFieldsWithOrdinal.length)
+ rowWriter.resetRowWriter()
+
+ var fieldIdx = 0
+ rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
+ val field = fieldWithOrdinal._1
+
+ val markerBytes = record.get(fieldIdx).asInstanceOf[ByteBuffer].array()
+ val markerBuf = ByteBuffer.wrap(markerBytes)
+ markerBuf.order(ByteOrder.BIG_ENDIAN)
+ val marker = markerBuf.get()
+
+ if (marker == nullValMarker) {
+ rowWriter.setNullAt(idx)
+ } else {
+ field.dataType match {
+ case BooleanType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ rowWriter.write(idx, bytes(0) == 1)
+
+ case ByteType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.get())
+
+ case ShortType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.getShort())
+
+ case IntegerType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.getInt())
+
+ case LongType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.getLong())
+
+ case FloatType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ if (marker == negativeValMarker) {
+ val floatVal = valueBuf.getFloat
+ val updatedVal = floatToRawIntBits(floatVal) ^ floatFlipBitMask
+ rowWriter.write(idx, intBitsToFloat(updatedVal))
+ } else {
+ rowWriter.write(idx, valueBuf.getFloat())
+ }
+
+ case DoubleType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ if (marker == negativeValMarker) {
+ val doubleVal = valueBuf.getDouble
+ val updatedVal = doubleToRawLongBits(doubleVal) ^ doubleFlipBitMask
+ rowWriter.write(idx, longBitsToDouble(updatedVal))
+ } else {
+ rowWriter.write(idx, valueBuf.getDouble())
+ }
+
+ case _ => throw new UnsupportedOperationException(
+ s"Range scan decoding not supported for data type: ${field.dataType}")
+ }
+ }
+ fieldIdx += 2
+ }
+
+ rowWriter.getRow()
+ }
+
+ override def decodeValue(bytes: Array[Byte]): UnsafeRow =
+ decodeFromAvroToUnsafeRow(
+ bytes, avroEncoder.valueDeserializer, valueAvroType, valueProj)
+}
+
+abstract class RocksDBKeyStateEncoderBase(
+ useColumnFamilies: Boolean,
+ virtualColFamilyId: Option[Short] = None) extends RocksDBKeyStateEncoder {
+ def offsetForColFamilyPrefix: Int =
+ if (useColumnFamilies) VIRTUAL_COL_FAMILY_PREFIX_BYTES else 0
+
+ val out = new ByteArrayOutputStream
+ /**
+ * Get Byte Array for the virtual column family id that is used as prefix for
+ * key state rows.
+ */
+ override def getColumnFamilyIdBytes(): Array[Byte] = {
+ assert(useColumnFamilies, "Cannot return virtual Column Family Id Bytes" +
+ " because multiple Column is not supported for this encoder")
+ val encodedBytes = new Array[Byte](VIRTUAL_COL_FAMILY_PREFIX_BYTES)
+ Platform.putShort(encodedBytes, Platform.BYTE_ARRAY_OFFSET, virtualColFamilyId.get)
+ encodedBytes
+ }
+
+ /**
+ * Encode and put column family Id as a prefix to a pre-allocated byte array.
+ *
+ * @param numBytes - size of byte array to be created for storing key row (without
+ * column family prefix)
+ * @return Array[Byte] for an array byte to put encoded key bytes
+ * Int for a starting offset to put the encoded key bytes
+ */
+ protected def encodeColumnFamilyPrefix(numBytes: Int): (Array[Byte], Int) = {
+ val encodedBytes = new Array[Byte](numBytes + offsetForColFamilyPrefix)
+ var offset = Platform.BYTE_ARRAY_OFFSET
+ if (useColumnFamilies) {
+ Platform.putShort(encodedBytes, Platform.BYTE_ARRAY_OFFSET, virtualColFamilyId.get)
+ offset = Platform.BYTE_ARRAY_OFFSET + offsetForColFamilyPrefix
+ }
+ (encodedBytes, offset)
+ }
+
+ /**
+ * Get starting offset for decoding an encoded key byte array.
+ */
+ protected def decodeKeyStartOffset: Int = {
+ if (useColumnFamilies) {
+ Platform.BYTE_ARRAY_OFFSET + VIRTUAL_COL_FAMILY_PREFIX_BYTES
+ } else Platform.BYTE_ARRAY_OFFSET
+ }
+}
+
+object RocksDBStateEncoder extends Logging {
+ def getKeyEncoder(
+ dataEncoder: RocksDBDataEncoder,
+ keyStateEncoderSpec: KeyStateEncoderSpec,
+ useColumnFamilies: Boolean,
+ virtualColFamilyId: Option[Short] = None,
+ avroEnc: Option[AvroEncoder] = None): RocksDBKeyStateEncoder = {
+ // Return the key state encoder based on the requested type
+ keyStateEncoderSpec match {
+ case NoPrefixKeyStateEncoderSpec(keySchema) =>
+ new NoPrefixKeyStateEncoder(dataEncoder, keySchema, useColumnFamilies, virtualColFamilyId)
+
+ case PrefixKeyScanStateEncoderSpec(keySchema, numColsPrefixKey) =>
+ new PrefixKeyScanStateEncoder(dataEncoder, keySchema, numColsPrefixKey,
+ useColumnFamilies, virtualColFamilyId)
+
+ case RangeKeyScanStateEncoderSpec(keySchema, orderingOrdinals) =>
+ new RangeKeyScanStateEncoder(dataEncoder, keySchema, orderingOrdinals,
+ useColumnFamilies, virtualColFamilyId)
+
+ case _ =>
+ throw new IllegalArgumentException(s"Unsupported key state encoder spec: " +
+ s"$keyStateEncoderSpec")
+ }
+ }
+
+ def getValueEncoder(
+ dataEncoder: RocksDBDataEncoder,
+ valueSchema: StructType,
+ useMultipleValuesPerKey: Boolean,
+ avroEnc: Option[AvroEncoder] = None): RocksDBValueStateEncoder = {
Review Comment:
ditto
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala:
##########
@@ -51,93 +51,56 @@ sealed trait RocksDBValueStateEncoder {
def decodeValues(valueBytes: Array[Byte]): Iterator[UnsafeRow]
}
-abstract class RocksDBKeyStateEncoderBase(
- useColumnFamilies: Boolean,
- virtualColFamilyId: Option[Short] = None) extends RocksDBKeyStateEncoder {
- def offsetForColFamilyPrefix: Int =
- if (useColumnFamilies) VIRTUAL_COL_FAMILY_PREFIX_BYTES else 0
-
- val out = new ByteArrayOutputStream
- /**
- * Get Byte Array for the virtual column family id that is used as prefix for
- * key state rows.
- */
- override def getColumnFamilyIdBytes(): Array[Byte] = {
- assert(useColumnFamilies, "Cannot return virtual Column Family Id Bytes" +
- " because multiple Column is not supported for this encoder")
- val encodedBytes = new Array[Byte](VIRTUAL_COL_FAMILY_PREFIX_BYTES)
- Platform.putShort(encodedBytes, Platform.BYTE_ARRAY_OFFSET, virtualColFamilyId.get)
- encodedBytes
- }
-
- /**
- * Encode and put column family Id as a prefix to a pre-allocated byte array.
- *
- * @param numBytes - size of byte array to be created for storing key row (without
- * column family prefix)
- * @return Array[Byte] for an array byte to put encoded key bytes
- * Int for a starting offset to put the encoded key bytes
- */
- protected def encodeColumnFamilyPrefix(numBytes: Int): (Array[Byte], Int) = {
- val encodedBytes = new Array[Byte](numBytes + offsetForColFamilyPrefix)
- var offset = Platform.BYTE_ARRAY_OFFSET
- if (useColumnFamilies) {
- Platform.putShort(encodedBytes, Platform.BYTE_ARRAY_OFFSET, virtualColFamilyId.get)
- offset = Platform.BYTE_ARRAY_OFFSET + offsetForColFamilyPrefix
- }
- (encodedBytes, offset)
- }
+/**
+ * The DataEncoder can encode UnsafeRows into raw bytes in two ways:
+ * - Using the direct byte layout of the UnsafeRow
+ * - Converting the UnsafeRow into an Avro row, and encoding that
+ * In both of these cases, the raw bytes that are written into RockDB have
+ * headers, footers and other metadata, but they also have data that is provided
+ * by the callers. The metadata in each row does not need to be written as Avro or UnsafeRow,
+ * but the actual data provided by the caller does.
+ */
+trait DataEncoder {
+ def encodeKey(row: UnsafeRow): Array[Byte]
+ def encodeRemainingKey(row: UnsafeRow): Array[Byte]
+ def encodePrefixKeyForRangeScan(row: UnsafeRow): Array[Byte]
+ def encodeValue(row: UnsafeRow): Array[Byte]
Review Comment:
can you add docs for these actually? What additional information would be needed for these?
For example, when would `encodeRemainingKey` be used? What schema information needs to be passed in for this?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala:
##########
@@ -800,156 +583,509 @@ class RangeKeyScanStateEncoder(
})
record.put(fieldIdx, ByteBuffer.wrap(markerBuffer.array()))
- val valueBuffer = ByteBuffer.allocate(8)
- valueBuffer.order(ByteOrder.BIG_ENDIAN)
- if ((rawBits & doubleSignBitMask) != 0) {
- val updatedVal = rawBits ^ doubleFlipBitMask
- valueBuffer.putDouble(longBitsToDouble(updatedVal))
- } else {
- valueBuffer.putDouble(doubleVal)
- }
- record.put(fieldIdx + 1, ByteBuffer.wrap(valueBuffer.array()))
+ val valueBuffer = ByteBuffer.allocate(8)
+ valueBuffer.order(ByteOrder.BIG_ENDIAN)
+ if ((rawBits & doubleSignBitMask) != 0) {
+ val updatedVal = rawBits ^ doubleFlipBitMask
+ valueBuffer.putDouble(longBitsToDouble(updatedVal))
+ } else {
+ valueBuffer.putDouble(doubleVal)
+ }
+ record.put(fieldIdx + 1, ByteBuffer.wrap(valueBuffer.array()))
+
+ case _ => throw new UnsupportedOperationException(
+ s"Range scan encoding not supported for data type: ${field.dataType}")
+ }
+ }
+ fieldIdx += 2
+ }
+
+ out.reset()
+ val writer = new GenericDatumWriter[GenericRecord](rangeScanAvroType)
+ val encoder = EncoderFactory.get().binaryEncoder(out, null)
+ writer.write(record, encoder)
+ encoder.flush()
+ out.toByteArray
+ }
+
+ override def encodeValue(row: UnsafeRow): Array[Byte] =
+ encodeUnsafeRowToAvro(row, avroEncoder.valueSerializer, valueAvroType, out)
+
+ override def decodeKey(bytes: Array[Byte]): UnsafeRow = {
+ keyStateEncoderSpec match {
+ case NoPrefixKeyStateEncoderSpec(_) =>
+ decodeFromAvroToUnsafeRow(bytes, avroEncoder.keyDeserializer, keyAvroType, keyProj)
+ case PrefixKeyScanStateEncoderSpec(_, _) =>
+ decodeFromAvroToUnsafeRow(
+ bytes, avroEncoder.keyDeserializer, prefixKeyAvroType, prefixKeyProj)
+ case _ => throw unsupportedOperationForKeyStateEncoder("decodeKey")
+ }
+ }
+
+
+ override def decodeRemainingKey(bytes: Array[Byte]): UnsafeRow = {
+ keyStateEncoderSpec match {
+ case PrefixKeyScanStateEncoderSpec(_, _) =>
+ decodeFromAvroToUnsafeRow(bytes,
+ avroEncoder.suffixKeyDeserializer.get, remainingKeyAvroType, remainingKeyAvroProjection)
+ case RangeKeyScanStateEncoderSpec(_, _) =>
+ decodeFromAvroToUnsafeRow(
+ bytes, avroEncoder.keyDeserializer, remainingKeyAvroType, remainingKeyAvroProjection)
+ case _ => throw unsupportedOperationForKeyStateEncoder("decodeRemainingKey")
+ }
+ }
+
+ /**
+ * Decodes an Avro-encoded byte array back into an UnsafeRow for range scan operations.
+ *
+ * This method reverses the encoding process performed by encodePrefixKeyForRangeScan:
+ * - Reads the marker byte to determine null status or sign
+ * - Reconstructs the original values from big-endian format
+ * - Handles special cases for floating point numbers by reversing bit manipulations
+ *
+ * The decoding process preserves the original data types and values, including:
+ * - Null values marked by nullValMarker
+ * - Sign information for numeric types
+ * - Proper restoration of negative floating point values
+ *
+ * @param bytes The Avro-encoded byte array to decode
+ * @param avroType The Avro schema defining the structure for decoding
+ * @return UnsafeRow containing the decoded data
+ * @throws UnsupportedOperationException if a field's data type is not supported for range
+ * scan decoding
+ */
+ override def decodePrefixKeyForRangeScan(bytes: Array[Byte]): UnsafeRow = {
+ val reader = new GenericDatumReader[GenericRecord](rangeScanAvroType)
+ val decoder = DecoderFactory.get().binaryDecoder(bytes, 0, bytes.length, null)
+ val record = reader.read(null, decoder)
+
+ val rowWriter = new UnsafeRowWriter(rangeScanKeyFieldsWithOrdinal.length)
+ rowWriter.resetRowWriter()
+
+ var fieldIdx = 0
+ rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
+ val field = fieldWithOrdinal._1
+
+ val markerBytes = record.get(fieldIdx).asInstanceOf[ByteBuffer].array()
+ val markerBuf = ByteBuffer.wrap(markerBytes)
+ markerBuf.order(ByteOrder.BIG_ENDIAN)
+ val marker = markerBuf.get()
+
+ if (marker == nullValMarker) {
+ rowWriter.setNullAt(idx)
+ } else {
+ field.dataType match {
+ case BooleanType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ rowWriter.write(idx, bytes(0) == 1)
+
+ case ByteType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.get())
+
+ case ShortType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.getShort())
+
+ case IntegerType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.getInt())
+
+ case LongType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.getLong())
+
+ case FloatType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ if (marker == negativeValMarker) {
+ val floatVal = valueBuf.getFloat
+ val updatedVal = floatToRawIntBits(floatVal) ^ floatFlipBitMask
+ rowWriter.write(idx, intBitsToFloat(updatedVal))
+ } else {
+ rowWriter.write(idx, valueBuf.getFloat())
+ }
+
+ case DoubleType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ if (marker == negativeValMarker) {
+ val doubleVal = valueBuf.getDouble
+ val updatedVal = doubleToRawLongBits(doubleVal) ^ doubleFlipBitMask
+ rowWriter.write(idx, longBitsToDouble(updatedVal))
+ } else {
+ rowWriter.write(idx, valueBuf.getDouble())
+ }
+
+ case _ => throw new UnsupportedOperationException(
+ s"Range scan decoding not supported for data type: ${field.dataType}")
+ }
+ }
+ fieldIdx += 2
+ }
+
+ rowWriter.getRow()
+ }
+
+ override def decodeValue(bytes: Array[Byte]): UnsafeRow =
+ decodeFromAvroToUnsafeRow(
+ bytes, avroEncoder.valueDeserializer, valueAvroType, valueProj)
+}
+
+abstract class RocksDBKeyStateEncoderBase(
+ useColumnFamilies: Boolean,
+ virtualColFamilyId: Option[Short] = None) extends RocksDBKeyStateEncoder {
+ def offsetForColFamilyPrefix: Int =
+ if (useColumnFamilies) VIRTUAL_COL_FAMILY_PREFIX_BYTES else 0
+
+ val out = new ByteArrayOutputStream
+ /**
+ * Get Byte Array for the virtual column family id that is used as prefix for
+ * key state rows.
+ */
+ override def getColumnFamilyIdBytes(): Array[Byte] = {
+ assert(useColumnFamilies, "Cannot return virtual Column Family Id Bytes" +
+ " because multiple Column is not supported for this encoder")
+ val encodedBytes = new Array[Byte](VIRTUAL_COL_FAMILY_PREFIX_BYTES)
+ Platform.putShort(encodedBytes, Platform.BYTE_ARRAY_OFFSET, virtualColFamilyId.get)
+ encodedBytes
+ }
+
+ /**
+ * Encode and put column family Id as a prefix to a pre-allocated byte array.
+ *
+ * @param numBytes - size of byte array to be created for storing key row (without
+ * column family prefix)
+ * @return Array[Byte] for an array byte to put encoded key bytes
+ * Int for a starting offset to put the encoded key bytes
+ */
+ protected def encodeColumnFamilyPrefix(numBytes: Int): (Array[Byte], Int) = {
+ val encodedBytes = new Array[Byte](numBytes + offsetForColFamilyPrefix)
+ var offset = Platform.BYTE_ARRAY_OFFSET
+ if (useColumnFamilies) {
+ Platform.putShort(encodedBytes, Platform.BYTE_ARRAY_OFFSET, virtualColFamilyId.get)
+ offset = Platform.BYTE_ARRAY_OFFSET + offsetForColFamilyPrefix
+ }
+ (encodedBytes, offset)
+ }
+
+ /**
+ * Get starting offset for decoding an encoded key byte array.
+ */
+ protected def decodeKeyStartOffset: Int = {
+ if (useColumnFamilies) {
+ Platform.BYTE_ARRAY_OFFSET + VIRTUAL_COL_FAMILY_PREFIX_BYTES
+ } else Platform.BYTE_ARRAY_OFFSET
+ }
+}
+
+object RocksDBStateEncoder extends Logging {
Review Comment:
Can you add scaladoc for all the methods below?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala:
##########
@@ -800,156 +583,509 @@ class RangeKeyScanStateEncoder(
})
record.put(fieldIdx, ByteBuffer.wrap(markerBuffer.array()))
- val valueBuffer = ByteBuffer.allocate(8)
- valueBuffer.order(ByteOrder.BIG_ENDIAN)
- if ((rawBits & doubleSignBitMask) != 0) {
- val updatedVal = rawBits ^ doubleFlipBitMask
- valueBuffer.putDouble(longBitsToDouble(updatedVal))
- } else {
- valueBuffer.putDouble(doubleVal)
- }
- record.put(fieldIdx + 1, ByteBuffer.wrap(valueBuffer.array()))
+ val valueBuffer = ByteBuffer.allocate(8)
+ valueBuffer.order(ByteOrder.BIG_ENDIAN)
+ if ((rawBits & doubleSignBitMask) != 0) {
+ val updatedVal = rawBits ^ doubleFlipBitMask
+ valueBuffer.putDouble(longBitsToDouble(updatedVal))
+ } else {
+ valueBuffer.putDouble(doubleVal)
+ }
+ record.put(fieldIdx + 1, ByteBuffer.wrap(valueBuffer.array()))
+
+ case _ => throw new UnsupportedOperationException(
+ s"Range scan encoding not supported for data type: ${field.dataType}")
+ }
+ }
+ fieldIdx += 2
+ }
+
+ out.reset()
+ val writer = new GenericDatumWriter[GenericRecord](rangeScanAvroType)
+ val encoder = EncoderFactory.get().binaryEncoder(out, null)
+ writer.write(record, encoder)
+ encoder.flush()
+ out.toByteArray
+ }
+
+ override def encodeValue(row: UnsafeRow): Array[Byte] =
+ encodeUnsafeRowToAvro(row, avroEncoder.valueSerializer, valueAvroType, out)
+
+ override def decodeKey(bytes: Array[Byte]): UnsafeRow = {
+ keyStateEncoderSpec match {
+ case NoPrefixKeyStateEncoderSpec(_) =>
+ decodeFromAvroToUnsafeRow(bytes, avroEncoder.keyDeserializer, keyAvroType, keyProj)
+ case PrefixKeyScanStateEncoderSpec(_, _) =>
+ decodeFromAvroToUnsafeRow(
+ bytes, avroEncoder.keyDeserializer, prefixKeyAvroType, prefixKeyProj)
+ case _ => throw unsupportedOperationForKeyStateEncoder("decodeKey")
+ }
+ }
+
+
+ override def decodeRemainingKey(bytes: Array[Byte]): UnsafeRow = {
+ keyStateEncoderSpec match {
+ case PrefixKeyScanStateEncoderSpec(_, _) =>
+ decodeFromAvroToUnsafeRow(bytes,
+ avroEncoder.suffixKeyDeserializer.get, remainingKeyAvroType, remainingKeyAvroProjection)
+ case RangeKeyScanStateEncoderSpec(_, _) =>
+ decodeFromAvroToUnsafeRow(
+ bytes, avroEncoder.keyDeserializer, remainingKeyAvroType, remainingKeyAvroProjection)
+ case _ => throw unsupportedOperationForKeyStateEncoder("decodeRemainingKey")
+ }
+ }
+
+ /**
+ * Decodes an Avro-encoded byte array back into an UnsafeRow for range scan operations.
+ *
+ * This method reverses the encoding process performed by encodePrefixKeyForRangeScan:
+ * - Reads the marker byte to determine null status or sign
+ * - Reconstructs the original values from big-endian format
+ * - Handles special cases for floating point numbers by reversing bit manipulations
+ *
+ * The decoding process preserves the original data types and values, including:
+ * - Null values marked by nullValMarker
+ * - Sign information for numeric types
+ * - Proper restoration of negative floating point values
+ *
+ * @param bytes The Avro-encoded byte array to decode
+ * @param avroType The Avro schema defining the structure for decoding
+ * @return UnsafeRow containing the decoded data
+ * @throws UnsupportedOperationException if a field's data type is not supported for range
+ * scan decoding
+ */
+ override def decodePrefixKeyForRangeScan(bytes: Array[Byte]): UnsafeRow = {
+ val reader = new GenericDatumReader[GenericRecord](rangeScanAvroType)
+ val decoder = DecoderFactory.get().binaryDecoder(bytes, 0, bytes.length, null)
+ val record = reader.read(null, decoder)
+
+ val rowWriter = new UnsafeRowWriter(rangeScanKeyFieldsWithOrdinal.length)
+ rowWriter.resetRowWriter()
+
+ var fieldIdx = 0
+ rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
+ val field = fieldWithOrdinal._1
+
+ val markerBytes = record.get(fieldIdx).asInstanceOf[ByteBuffer].array()
+ val markerBuf = ByteBuffer.wrap(markerBytes)
+ markerBuf.order(ByteOrder.BIG_ENDIAN)
+ val marker = markerBuf.get()
+
+ if (marker == nullValMarker) {
+ rowWriter.setNullAt(idx)
+ } else {
+ field.dataType match {
+ case BooleanType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ rowWriter.write(idx, bytes(0) == 1)
+
+ case ByteType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.get())
+
+ case ShortType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.getShort())
+
+ case IntegerType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.getInt())
+
+ case LongType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.getLong())
+
+ case FloatType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ if (marker == negativeValMarker) {
+ val floatVal = valueBuf.getFloat
+ val updatedVal = floatToRawIntBits(floatVal) ^ floatFlipBitMask
+ rowWriter.write(idx, intBitsToFloat(updatedVal))
+ } else {
+ rowWriter.write(idx, valueBuf.getFloat())
+ }
+
+ case DoubleType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ if (marker == negativeValMarker) {
+ val doubleVal = valueBuf.getDouble
+ val updatedVal = doubleToRawLongBits(doubleVal) ^ doubleFlipBitMask
+ rowWriter.write(idx, longBitsToDouble(updatedVal))
+ } else {
+ rowWriter.write(idx, valueBuf.getDouble())
+ }
+
+ case _ => throw new UnsupportedOperationException(
+ s"Range scan decoding not supported for data type: ${field.dataType}")
+ }
+ }
+ fieldIdx += 2
+ }
+
+ rowWriter.getRow()
+ }
+
+ override def decodeValue(bytes: Array[Byte]): UnsafeRow =
+ decodeFromAvroToUnsafeRow(
+ bytes, avroEncoder.valueDeserializer, valueAvroType, valueProj)
+}
+
+abstract class RocksDBKeyStateEncoderBase(
+ useColumnFamilies: Boolean,
+ virtualColFamilyId: Option[Short] = None) extends RocksDBKeyStateEncoder {
+ def offsetForColFamilyPrefix: Int =
+ if (useColumnFamilies) VIRTUAL_COL_FAMILY_PREFIX_BYTES else 0
+
+ val out = new ByteArrayOutputStream
+ /**
+ * Get Byte Array for the virtual column family id that is used as prefix for
+ * key state rows.
+ */
+ override def getColumnFamilyIdBytes(): Array[Byte] = {
+ assert(useColumnFamilies, "Cannot return virtual Column Family Id Bytes" +
+ " because multiple Column is not supported for this encoder")
+ val encodedBytes = new Array[Byte](VIRTUAL_COL_FAMILY_PREFIX_BYTES)
+ Platform.putShort(encodedBytes, Platform.BYTE_ARRAY_OFFSET, virtualColFamilyId.get)
+ encodedBytes
+ }
+
+ /**
+ * Encode and put column family Id as a prefix to a pre-allocated byte array.
+ *
+ * @param numBytes - size of byte array to be created for storing key row (without
+ * column family prefix)
+ * @return Array[Byte] for an array byte to put encoded key bytes
+ * Int for a starting offset to put the encoded key bytes
+ */
+ protected def encodeColumnFamilyPrefix(numBytes: Int): (Array[Byte], Int) = {
+ val encodedBytes = new Array[Byte](numBytes + offsetForColFamilyPrefix)
+ var offset = Platform.BYTE_ARRAY_OFFSET
+ if (useColumnFamilies) {
+ Platform.putShort(encodedBytes, Platform.BYTE_ARRAY_OFFSET, virtualColFamilyId.get)
+ offset = Platform.BYTE_ARRAY_OFFSET + offsetForColFamilyPrefix
+ }
+ (encodedBytes, offset)
+ }
+
+ /**
+ * Get starting offset for decoding an encoded key byte array.
+ */
+ protected def decodeKeyStartOffset: Int = {
+ if (useColumnFamilies) {
+ Platform.BYTE_ARRAY_OFFSET + VIRTUAL_COL_FAMILY_PREFIX_BYTES
+ } else Platform.BYTE_ARRAY_OFFSET
+ }
+}
+
+object RocksDBStateEncoder extends Logging {
+ def getKeyEncoder(
+ dataEncoder: RocksDBDataEncoder,
+ keyStateEncoderSpec: KeyStateEncoderSpec,
+ useColumnFamilies: Boolean,
+ virtualColFamilyId: Option[Short] = None,
+ avroEnc: Option[AvroEncoder] = None): RocksDBKeyStateEncoder = {
Review Comment:
you can remove this line now?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala:
##########
@@ -800,156 +583,509 @@ class RangeKeyScanStateEncoder(
})
record.put(fieldIdx, ByteBuffer.wrap(markerBuffer.array()))
- val valueBuffer = ByteBuffer.allocate(8)
- valueBuffer.order(ByteOrder.BIG_ENDIAN)
- if ((rawBits & doubleSignBitMask) != 0) {
- val updatedVal = rawBits ^ doubleFlipBitMask
- valueBuffer.putDouble(longBitsToDouble(updatedVal))
- } else {
- valueBuffer.putDouble(doubleVal)
- }
- record.put(fieldIdx + 1, ByteBuffer.wrap(valueBuffer.array()))
+ val valueBuffer = ByteBuffer.allocate(8)
+ valueBuffer.order(ByteOrder.BIG_ENDIAN)
+ if ((rawBits & doubleSignBitMask) != 0) {
+ val updatedVal = rawBits ^ doubleFlipBitMask
+ valueBuffer.putDouble(longBitsToDouble(updatedVal))
+ } else {
+ valueBuffer.putDouble(doubleVal)
+ }
+ record.put(fieldIdx + 1, ByteBuffer.wrap(valueBuffer.array()))
+
+ case _ => throw new UnsupportedOperationException(
+ s"Range scan encoding not supported for data type: ${field.dataType}")
+ }
+ }
+ fieldIdx += 2
+ }
+
+ out.reset()
+ val writer = new GenericDatumWriter[GenericRecord](rangeScanAvroType)
+ val encoder = EncoderFactory.get().binaryEncoder(out, null)
+ writer.write(record, encoder)
+ encoder.flush()
+ out.toByteArray
+ }
+
+ override def encodeValue(row: UnsafeRow): Array[Byte] =
+ encodeUnsafeRowToAvro(row, avroEncoder.valueSerializer, valueAvroType, out)
+
+ override def decodeKey(bytes: Array[Byte]): UnsafeRow = {
+ keyStateEncoderSpec match {
+ case NoPrefixKeyStateEncoderSpec(_) =>
+ decodeFromAvroToUnsafeRow(bytes, avroEncoder.keyDeserializer, keyAvroType, keyProj)
+ case PrefixKeyScanStateEncoderSpec(_, _) =>
+ decodeFromAvroToUnsafeRow(
+ bytes, avroEncoder.keyDeserializer, prefixKeyAvroType, prefixKeyProj)
+ case _ => throw unsupportedOperationForKeyStateEncoder("decodeKey")
+ }
+ }
+
+
+ override def decodeRemainingKey(bytes: Array[Byte]): UnsafeRow = {
+ keyStateEncoderSpec match {
+ case PrefixKeyScanStateEncoderSpec(_, _) =>
+ decodeFromAvroToUnsafeRow(bytes,
+ avroEncoder.suffixKeyDeserializer.get, remainingKeyAvroType, remainingKeyAvroProjection)
+ case RangeKeyScanStateEncoderSpec(_, _) =>
+ decodeFromAvroToUnsafeRow(
+ bytes, avroEncoder.keyDeserializer, remainingKeyAvroType, remainingKeyAvroProjection)
+ case _ => throw unsupportedOperationForKeyStateEncoder("decodeRemainingKey")
+ }
+ }
+
+ /**
+ * Decodes an Avro-encoded byte array back into an UnsafeRow for range scan operations.
+ *
+ * This method reverses the encoding process performed by encodePrefixKeyForRangeScan:
+ * - Reads the marker byte to determine null status or sign
+ * - Reconstructs the original values from big-endian format
+ * - Handles special cases for floating point numbers by reversing bit manipulations
+ *
+ * The decoding process preserves the original data types and values, including:
+ * - Null values marked by nullValMarker
+ * - Sign information for numeric types
+ * - Proper restoration of negative floating point values
+ *
+ * @param bytes The Avro-encoded byte array to decode
+ * @param avroType The Avro schema defining the structure for decoding
+ * @return UnsafeRow containing the decoded data
+ * @throws UnsupportedOperationException if a field's data type is not supported for range
+ * scan decoding
+ */
+ override def decodePrefixKeyForRangeScan(bytes: Array[Byte]): UnsafeRow = {
+ val reader = new GenericDatumReader[GenericRecord](rangeScanAvroType)
+ val decoder = DecoderFactory.get().binaryDecoder(bytes, 0, bytes.length, null)
+ val record = reader.read(null, decoder)
+
+ val rowWriter = new UnsafeRowWriter(rangeScanKeyFieldsWithOrdinal.length)
+ rowWriter.resetRowWriter()
+
+ var fieldIdx = 0
+ rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
+ val field = fieldWithOrdinal._1
+
+ val markerBytes = record.get(fieldIdx).asInstanceOf[ByteBuffer].array()
+ val markerBuf = ByteBuffer.wrap(markerBytes)
+ markerBuf.order(ByteOrder.BIG_ENDIAN)
+ val marker = markerBuf.get()
+
+ if (marker == nullValMarker) {
+ rowWriter.setNullAt(idx)
+ } else {
+ field.dataType match {
+ case BooleanType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ rowWriter.write(idx, bytes(0) == 1)
+
+ case ByteType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.get())
+
+ case ShortType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.getShort())
+
+ case IntegerType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.getInt())
+
+ case LongType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.getLong())
+
+ case FloatType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ if (marker == negativeValMarker) {
+ val floatVal = valueBuf.getFloat
+ val updatedVal = floatToRawIntBits(floatVal) ^ floatFlipBitMask
+ rowWriter.write(idx, intBitsToFloat(updatedVal))
+ } else {
+ rowWriter.write(idx, valueBuf.getFloat())
+ }
+
+ case DoubleType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ if (marker == negativeValMarker) {
+ val doubleVal = valueBuf.getDouble
+ val updatedVal = doubleToRawLongBits(doubleVal) ^ doubleFlipBitMask
+ rowWriter.write(idx, longBitsToDouble(updatedVal))
+ } else {
+ rowWriter.write(idx, valueBuf.getDouble())
+ }
+
+ case _ => throw new UnsupportedOperationException(
+ s"Range scan decoding not supported for data type: ${field.dataType}")
+ }
+ }
+ fieldIdx += 2
+ }
+
+ rowWriter.getRow()
+ }
+
+ override def decodeValue(bytes: Array[Byte]): UnsafeRow =
+ decodeFromAvroToUnsafeRow(
+ bytes, avroEncoder.valueDeserializer, valueAvroType, valueProj)
+}
+
+abstract class RocksDBKeyStateEncoderBase(
+ useColumnFamilies: Boolean,
+ virtualColFamilyId: Option[Short] = None) extends RocksDBKeyStateEncoder {
+ def offsetForColFamilyPrefix: Int =
+ if (useColumnFamilies) VIRTUAL_COL_FAMILY_PREFIX_BYTES else 0
+
+ val out = new ByteArrayOutputStream
+ /**
+ * Get Byte Array for the virtual column family id that is used as prefix for
+ * key state rows.
+ */
+ override def getColumnFamilyIdBytes(): Array[Byte] = {
+ assert(useColumnFamilies, "Cannot return virtual Column Family Id Bytes" +
+ " because multiple Column is not supported for this encoder")
+ val encodedBytes = new Array[Byte](VIRTUAL_COL_FAMILY_PREFIX_BYTES)
+ Platform.putShort(encodedBytes, Platform.BYTE_ARRAY_OFFSET, virtualColFamilyId.get)
+ encodedBytes
+ }
+
+ /**
+ * Encode and put column family Id as a prefix to a pre-allocated byte array.
+ *
+ * @param numBytes - size of byte array to be created for storing key row (without
+ * column family prefix)
+ * @return Array[Byte] for an array byte to put encoded key bytes
+ * Int for a starting offset to put the encoded key bytes
+ */
+ protected def encodeColumnFamilyPrefix(numBytes: Int): (Array[Byte], Int) = {
+ val encodedBytes = new Array[Byte](numBytes + offsetForColFamilyPrefix)
+ var offset = Platform.BYTE_ARRAY_OFFSET
+ if (useColumnFamilies) {
+ Platform.putShort(encodedBytes, Platform.BYTE_ARRAY_OFFSET, virtualColFamilyId.get)
+ offset = Platform.BYTE_ARRAY_OFFSET + offsetForColFamilyPrefix
+ }
+ (encodedBytes, offset)
+ }
+
+ /**
+ * Get starting offset for decoding an encoded key byte array.
+ */
+ protected def decodeKeyStartOffset: Int = {
+ if (useColumnFamilies) {
+ Platform.BYTE_ARRAY_OFFSET + VIRTUAL_COL_FAMILY_PREFIX_BYTES
+ } else Platform.BYTE_ARRAY_OFFSET
+ }
+}
+
+object RocksDBStateEncoder extends Logging {
+ def getKeyEncoder(
+ dataEncoder: RocksDBDataEncoder,
+ keyStateEncoderSpec: KeyStateEncoderSpec,
+ useColumnFamilies: Boolean,
+ virtualColFamilyId: Option[Short] = None,
+ avroEnc: Option[AvroEncoder] = None): RocksDBKeyStateEncoder = {
+ // Return the key state encoder based on the requested type
+ keyStateEncoderSpec match {
+ case NoPrefixKeyStateEncoderSpec(keySchema) =>
+ new NoPrefixKeyStateEncoder(dataEncoder, keySchema, useColumnFamilies, virtualColFamilyId)
+
+ case PrefixKeyScanStateEncoderSpec(keySchema, numColsPrefixKey) =>
+ new PrefixKeyScanStateEncoder(dataEncoder, keySchema, numColsPrefixKey,
+ useColumnFamilies, virtualColFamilyId)
+
+ case RangeKeyScanStateEncoderSpec(keySchema, orderingOrdinals) =>
+ new RangeKeyScanStateEncoder(dataEncoder, keySchema, orderingOrdinals,
+ useColumnFamilies, virtualColFamilyId)
+
+ case _ =>
+ throw new IllegalArgumentException(s"Unsupported key state encoder spec: " +
+ s"$keyStateEncoderSpec")
+ }
+ }
Review Comment:
should the encoderSpec's expose a method instead:
```
def toEncoder(dataEncoder: DataEncoder, useColumnFamilies, virtualColFamilyId): RocksDBKeyStateEncoderBase
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala:
##########
@@ -800,156 +583,509 @@ class RangeKeyScanStateEncoder(
})
record.put(fieldIdx, ByteBuffer.wrap(markerBuffer.array()))
- val valueBuffer = ByteBuffer.allocate(8)
- valueBuffer.order(ByteOrder.BIG_ENDIAN)
- if ((rawBits & doubleSignBitMask) != 0) {
- val updatedVal = rawBits ^ doubleFlipBitMask
- valueBuffer.putDouble(longBitsToDouble(updatedVal))
- } else {
- valueBuffer.putDouble(doubleVal)
- }
- record.put(fieldIdx + 1, ByteBuffer.wrap(valueBuffer.array()))
+ val valueBuffer = ByteBuffer.allocate(8)
+ valueBuffer.order(ByteOrder.BIG_ENDIAN)
+ if ((rawBits & doubleSignBitMask) != 0) {
+ val updatedVal = rawBits ^ doubleFlipBitMask
+ valueBuffer.putDouble(longBitsToDouble(updatedVal))
+ } else {
+ valueBuffer.putDouble(doubleVal)
+ }
+ record.put(fieldIdx + 1, ByteBuffer.wrap(valueBuffer.array()))
+
+ case _ => throw new UnsupportedOperationException(
+ s"Range scan encoding not supported for data type: ${field.dataType}")
+ }
+ }
+ fieldIdx += 2
+ }
+
+ out.reset()
+ val writer = new GenericDatumWriter[GenericRecord](rangeScanAvroType)
+ val encoder = EncoderFactory.get().binaryEncoder(out, null)
+ writer.write(record, encoder)
+ encoder.flush()
+ out.toByteArray
+ }
+
+ override def encodeValue(row: UnsafeRow): Array[Byte] =
+ encodeUnsafeRowToAvro(row, avroEncoder.valueSerializer, valueAvroType, out)
+
+ override def decodeKey(bytes: Array[Byte]): UnsafeRow = {
+ keyStateEncoderSpec match {
+ case NoPrefixKeyStateEncoderSpec(_) =>
+ decodeFromAvroToUnsafeRow(bytes, avroEncoder.keyDeserializer, keyAvroType, keyProj)
+ case PrefixKeyScanStateEncoderSpec(_, _) =>
+ decodeFromAvroToUnsafeRow(
+ bytes, avroEncoder.keyDeserializer, prefixKeyAvroType, prefixKeyProj)
+ case _ => throw unsupportedOperationForKeyStateEncoder("decodeKey")
+ }
+ }
+
+
+ override def decodeRemainingKey(bytes: Array[Byte]): UnsafeRow = {
+ keyStateEncoderSpec match {
+ case PrefixKeyScanStateEncoderSpec(_, _) =>
+ decodeFromAvroToUnsafeRow(bytes,
+ avroEncoder.suffixKeyDeserializer.get, remainingKeyAvroType, remainingKeyAvroProjection)
+ case RangeKeyScanStateEncoderSpec(_, _) =>
+ decodeFromAvroToUnsafeRow(
+ bytes, avroEncoder.keyDeserializer, remainingKeyAvroType, remainingKeyAvroProjection)
+ case _ => throw unsupportedOperationForKeyStateEncoder("decodeRemainingKey")
+ }
+ }
+
+ /**
+ * Decodes an Avro-encoded byte array back into an UnsafeRow for range scan operations.
+ *
+ * This method reverses the encoding process performed by encodePrefixKeyForRangeScan:
+ * - Reads the marker byte to determine null status or sign
+ * - Reconstructs the original values from big-endian format
+ * - Handles special cases for floating point numbers by reversing bit manipulations
+ *
+ * The decoding process preserves the original data types and values, including:
+ * - Null values marked by nullValMarker
+ * - Sign information for numeric types
+ * - Proper restoration of negative floating point values
+ *
+ * @param bytes The Avro-encoded byte array to decode
+ * @param avroType The Avro schema defining the structure for decoding
+ * @return UnsafeRow containing the decoded data
+ * @throws UnsupportedOperationException if a field's data type is not supported for range
+ * scan decoding
+ */
+ override def decodePrefixKeyForRangeScan(bytes: Array[Byte]): UnsafeRow = {
+ val reader = new GenericDatumReader[GenericRecord](rangeScanAvroType)
+ val decoder = DecoderFactory.get().binaryDecoder(bytes, 0, bytes.length, null)
+ val record = reader.read(null, decoder)
+
+ val rowWriter = new UnsafeRowWriter(rangeScanKeyFieldsWithOrdinal.length)
+ rowWriter.resetRowWriter()
+
+ var fieldIdx = 0
+ rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
+ val field = fieldWithOrdinal._1
+
+ val markerBytes = record.get(fieldIdx).asInstanceOf[ByteBuffer].array()
+ val markerBuf = ByteBuffer.wrap(markerBytes)
+ markerBuf.order(ByteOrder.BIG_ENDIAN)
+ val marker = markerBuf.get()
+
+ if (marker == nullValMarker) {
+ rowWriter.setNullAt(idx)
+ } else {
+ field.dataType match {
+ case BooleanType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ rowWriter.write(idx, bytes(0) == 1)
+
+ case ByteType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.get())
+
+ case ShortType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.getShort())
+
+ case IntegerType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.getInt())
+
+ case LongType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.getLong())
+
+ case FloatType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ if (marker == negativeValMarker) {
+ val floatVal = valueBuf.getFloat
+ val updatedVal = floatToRawIntBits(floatVal) ^ floatFlipBitMask
+ rowWriter.write(idx, intBitsToFloat(updatedVal))
+ } else {
+ rowWriter.write(idx, valueBuf.getFloat())
+ }
+
+ case DoubleType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ if (marker == negativeValMarker) {
+ val doubleVal = valueBuf.getDouble
+ val updatedVal = doubleToRawLongBits(doubleVal) ^ doubleFlipBitMask
+ rowWriter.write(idx, longBitsToDouble(updatedVal))
+ } else {
+ rowWriter.write(idx, valueBuf.getDouble())
+ }
+
+ case _ => throw new UnsupportedOperationException(
+ s"Range scan decoding not supported for data type: ${field.dataType}")
+ }
+ }
+ fieldIdx += 2
+ }
+
+ rowWriter.getRow()
+ }
+
+ override def decodeValue(bytes: Array[Byte]): UnsafeRow =
+ decodeFromAvroToUnsafeRow(
+ bytes, avroEncoder.valueDeserializer, valueAvroType, valueProj)
+}
+
+abstract class RocksDBKeyStateEncoderBase(
+ useColumnFamilies: Boolean,
+ virtualColFamilyId: Option[Short] = None) extends RocksDBKeyStateEncoder {
+ def offsetForColFamilyPrefix: Int =
+ if (useColumnFamilies) VIRTUAL_COL_FAMILY_PREFIX_BYTES else 0
+
+ val out = new ByteArrayOutputStream
+ /**
+ * Get Byte Array for the virtual column family id that is used as prefix for
+ * key state rows.
+ */
+ override def getColumnFamilyIdBytes(): Array[Byte] = {
+ assert(useColumnFamilies, "Cannot return virtual Column Family Id Bytes" +
+ " because multiple Column is not supported for this encoder")
+ val encodedBytes = new Array[Byte](VIRTUAL_COL_FAMILY_PREFIX_BYTES)
+ Platform.putShort(encodedBytes, Platform.BYTE_ARRAY_OFFSET, virtualColFamilyId.get)
+ encodedBytes
+ }
+
+ /**
+ * Encode and put column family Id as a prefix to a pre-allocated byte array.
+ *
+ * @param numBytes - size of byte array to be created for storing key row (without
+ * column family prefix)
+ * @return Array[Byte] for an array byte to put encoded key bytes
+ * Int for a starting offset to put the encoded key bytes
+ */
+ protected def encodeColumnFamilyPrefix(numBytes: Int): (Array[Byte], Int) = {
+ val encodedBytes = new Array[Byte](numBytes + offsetForColFamilyPrefix)
+ var offset = Platform.BYTE_ARRAY_OFFSET
+ if (useColumnFamilies) {
+ Platform.putShort(encodedBytes, Platform.BYTE_ARRAY_OFFSET, virtualColFamilyId.get)
+ offset = Platform.BYTE_ARRAY_OFFSET + offsetForColFamilyPrefix
+ }
+ (encodedBytes, offset)
+ }
+
+ /**
+ * Get starting offset for decoding an encoded key byte array.
+ */
+ protected def decodeKeyStartOffset: Int = {
+ if (useColumnFamilies) {
+ Platform.BYTE_ARRAY_OFFSET + VIRTUAL_COL_FAMILY_PREFIX_BYTES
+ } else Platform.BYTE_ARRAY_OFFSET
+ }
+}
+
+object RocksDBStateEncoder extends Logging {
+ def getKeyEncoder(
+ dataEncoder: RocksDBDataEncoder,
+ keyStateEncoderSpec: KeyStateEncoderSpec,
+ useColumnFamilies: Boolean,
+ virtualColFamilyId: Option[Short] = None,
+ avroEnc: Option[AvroEncoder] = None): RocksDBKeyStateEncoder = {
+ // Return the key state encoder based on the requested type
+ keyStateEncoderSpec match {
+ case NoPrefixKeyStateEncoderSpec(keySchema) =>
+ new NoPrefixKeyStateEncoder(dataEncoder, keySchema, useColumnFamilies, virtualColFamilyId)
+
+ case PrefixKeyScanStateEncoderSpec(keySchema, numColsPrefixKey) =>
+ new PrefixKeyScanStateEncoder(dataEncoder, keySchema, numColsPrefixKey,
+ useColumnFamilies, virtualColFamilyId)
+
+ case RangeKeyScanStateEncoderSpec(keySchema, orderingOrdinals) =>
+ new RangeKeyScanStateEncoder(dataEncoder, keySchema, orderingOrdinals,
+ useColumnFamilies, virtualColFamilyId)
+
+ case _ =>
+ throw new IllegalArgumentException(s"Unsupported key state encoder spec: " +
+ s"$keyStateEncoderSpec")
+ }
+ }
+
+ def getValueEncoder(
+ dataEncoder: RocksDBDataEncoder,
+ valueSchema: StructType,
+ useMultipleValuesPerKey: Boolean,
+ avroEnc: Option[AvroEncoder] = None): RocksDBValueStateEncoder = {
+ if (useMultipleValuesPerKey) {
+ new MultiValuedStateEncoder(dataEncoder, valueSchema)
+ } else {
+ new SingleValueStateEncoder(dataEncoder, valueSchema)
+ }
+ }
+
+ def getColumnFamilyIdBytes(virtualColFamilyId: Short): Array[Byte] = {
+ val encodedBytes = new Array[Byte](VIRTUAL_COL_FAMILY_PREFIX_BYTES)
+ Platform.putShort(encodedBytes, Platform.BYTE_ARRAY_OFFSET, virtualColFamilyId)
+ encodedBytes
+ }
+}
+
+/**
+ * RocksDB Key Encoder for UnsafeRow that supports prefix scan
+ *
+ * @param dataEncoder - the encoder that handles actual encoding/decoding of data
+ * @param keySchema - schema of the key to be encoded
+ * @param numColsPrefixKey - number of columns to be used for prefix key
+ * @param useColumnFamilies - if column family is enabled for this encoder
+ * @param avroEnc - if Avro encoding is specified for this StateEncoder, this encoder will
+ * be defined
+ */
+class PrefixKeyScanStateEncoder(
+ dataEncoder: RocksDBDataEncoder,
+ keySchema: StructType,
+ numColsPrefixKey: Int,
+ useColumnFamilies: Boolean = false,
+ virtualColFamilyId: Option[Short] = None)
+ extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) with Logging {
+
+ private val prefixKeyFieldsWithIdx: Seq[(StructField, Int)] = {
+ keySchema.zipWithIndex.take(numColsPrefixKey)
+ }
+
+ private val remainingKeyFieldsWithIdx: Seq[(StructField, Int)] = {
+ keySchema.zipWithIndex.drop(numColsPrefixKey)
+ }
+
+ private val prefixKeyProjection: UnsafeProjection = {
+ val refs = prefixKeyFieldsWithIdx.map(x => BoundReference(x._2, x._1.dataType, x._1.nullable))
+ UnsafeProjection.create(refs)
+ }
+
+ private val remainingKeyProjection: UnsafeProjection = {
+ val refs = remainingKeyFieldsWithIdx.map(x =>
+ BoundReference(x._2, x._1.dataType, x._1.nullable))
+ UnsafeProjection.create(refs)
+ }
+
+ // Prefix Key schema and projection definitions used by the Avro Serializers
+ // and Deserializers
+ private val prefixKeySchema = StructType(keySchema.take(numColsPrefixKey))
+ private lazy val prefixKeyAvroType = SchemaConverters.toAvroType(prefixKeySchema)
+ private val prefixKeyProj = UnsafeProjection.create(prefixKeySchema)
+
+ // Remaining Key schema and projection definitions used by the Avro Serializers
+ // and Deserializers
+ private val remainingKeySchema = StructType(keySchema.drop(numColsPrefixKey))
+ private lazy val remainingKeyAvroType = SchemaConverters.toAvroType(remainingKeySchema)
+ private val remainingKeyProj = UnsafeProjection.create(remainingKeySchema)
Review Comment:
some of these need to be cleaned up. They're unused and no longer needed
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateTTLTest.scala:
##########
@@ -41,7 +41,7 @@ case class OutputEvent(
* Test suite base for TransformWithState with TTL support.
*/
abstract class TransformWithStateTTLTest
- extends StreamTest with AlsoTestWithChangelogCheckpointingEnabled
+ extends StreamTest
Review Comment:
are these changes merge conflicts?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala:
##########
@@ -664,12 +689,29 @@ object RocksDBStateStoreProvider {
avroOptions.stableIdPrefixForUnionType, avroOptions.recursiveFieldMaxDepth)
}
+ /**
Review Comment:
IMHO all of this logic should live inside the AvroDataEncoder
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala:
##########
@@ -387,17 +386,16 @@ private[sql] class RocksDBStateStoreProvider
defaultColFamilyId = Some(rocksDB.createColFamilyIfAbsent(StateStore.DEFAULT_COL_FAMILY_NAME))
}
- val colFamilyName = StateStore.DEFAULT_COL_FAMILY_NAME
- // Create cache key using store ID to avoid collisions
- val avroEncCacheKey = s"${getRunId(hadoopConf)}_${stateStoreId.operatorId}_" +
- s"${stateStoreId.partitionId}_$colFamilyName"
- val avroEnc = getAvroEnc(
- stateStoreEncoding, avroEncCacheKey, keyStateEncoderSpec, valueSchema)
+ val dataEncoderCacheKey = s"${getRunId(hadoopConf)}_${stateStoreId.operatorId}_" +
+ s"${stateStoreId.partitionId}_${StateStore.DEFAULT_COL_FAMILY_NAME}"
+
+ val dataEncoder = getDataEncoder(
+ stateStoreEncoding, dataEncoderCacheKey, keyStateEncoderSpec, valueSchema)
keyValueEncoderMap.putIfAbsent(StateStore.DEFAULT_COL_FAMILY_NAME,
- (RocksDBStateEncoder.getKeyEncoder(keyStateEncoderSpec,
- useColumnFamilies, defaultColFamilyId, avroEnc),
- RocksDBStateEncoder.getValueEncoder(valueSchema, useMultipleValuesPerKey, avroEnc)))
+ (RocksDBStateEncoder.getKeyEncoder(dataEncoder, keyStateEncoderSpec,
+ useColumnFamilies, defaultColFamilyId),
+ RocksDBStateEncoder.getValueEncoder(dataEncoder, valueSchema, useMultipleValuesPerKey)))
Review Comment:
ditto
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateTTLTest.scala:
##########
@@ -41,7 +41,7 @@ case class OutputEvent(
* Test suite base for TransformWithState with TTL support.
*/
abstract class TransformWithStateTTLTest
- extends StreamTest with AlsoTestWithChangelogCheckpointingEnabled
+ extends StreamTest
Review Comment:
I think the merge accidentally removed this one
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala:
##########
@@ -800,156 +583,509 @@ class RangeKeyScanStateEncoder(
})
record.put(fieldIdx, ByteBuffer.wrap(markerBuffer.array()))
- val valueBuffer = ByteBuffer.allocate(8)
- valueBuffer.order(ByteOrder.BIG_ENDIAN)
- if ((rawBits & doubleSignBitMask) != 0) {
- val updatedVal = rawBits ^ doubleFlipBitMask
- valueBuffer.putDouble(longBitsToDouble(updatedVal))
- } else {
- valueBuffer.putDouble(doubleVal)
- }
- record.put(fieldIdx + 1, ByteBuffer.wrap(valueBuffer.array()))
+ val valueBuffer = ByteBuffer.allocate(8)
+ valueBuffer.order(ByteOrder.BIG_ENDIAN)
+ if ((rawBits & doubleSignBitMask) != 0) {
+ val updatedVal = rawBits ^ doubleFlipBitMask
+ valueBuffer.putDouble(longBitsToDouble(updatedVal))
+ } else {
+ valueBuffer.putDouble(doubleVal)
+ }
+ record.put(fieldIdx + 1, ByteBuffer.wrap(valueBuffer.array()))
+
+ case _ => throw new UnsupportedOperationException(
+ s"Range scan encoding not supported for data type: ${field.dataType}")
+ }
+ }
+ fieldIdx += 2
+ }
+
+ out.reset()
+ val writer = new GenericDatumWriter[GenericRecord](rangeScanAvroType)
+ val encoder = EncoderFactory.get().binaryEncoder(out, null)
+ writer.write(record, encoder)
+ encoder.flush()
+ out.toByteArray
+ }
+
+ override def encodeValue(row: UnsafeRow): Array[Byte] =
+ encodeUnsafeRowToAvro(row, avroEncoder.valueSerializer, valueAvroType, out)
+
+ override def decodeKey(bytes: Array[Byte]): UnsafeRow = {
+ keyStateEncoderSpec match {
+ case NoPrefixKeyStateEncoderSpec(_) =>
+ decodeFromAvroToUnsafeRow(bytes, avroEncoder.keyDeserializer, keyAvroType, keyProj)
+ case PrefixKeyScanStateEncoderSpec(_, _) =>
+ decodeFromAvroToUnsafeRow(
+ bytes, avroEncoder.keyDeserializer, prefixKeyAvroType, prefixKeyProj)
+ case _ => throw unsupportedOperationForKeyStateEncoder("decodeKey")
+ }
+ }
+
+
+ override def decodeRemainingKey(bytes: Array[Byte]): UnsafeRow = {
+ keyStateEncoderSpec match {
+ case PrefixKeyScanStateEncoderSpec(_, _) =>
+ decodeFromAvroToUnsafeRow(bytes,
+ avroEncoder.suffixKeyDeserializer.get, remainingKeyAvroType, remainingKeyAvroProjection)
+ case RangeKeyScanStateEncoderSpec(_, _) =>
+ decodeFromAvroToUnsafeRow(
+ bytes, avroEncoder.keyDeserializer, remainingKeyAvroType, remainingKeyAvroProjection)
+ case _ => throw unsupportedOperationForKeyStateEncoder("decodeRemainingKey")
+ }
+ }
+
+ /**
+ * Decodes an Avro-encoded byte array back into an UnsafeRow for range scan operations.
+ *
+ * This method reverses the encoding process performed by encodePrefixKeyForRangeScan:
+ * - Reads the marker byte to determine null status or sign
+ * - Reconstructs the original values from big-endian format
+ * - Handles special cases for floating point numbers by reversing bit manipulations
+ *
+ * The decoding process preserves the original data types and values, including:
+ * - Null values marked by nullValMarker
+ * - Sign information for numeric types
+ * - Proper restoration of negative floating point values
+ *
+ * @param bytes The Avro-encoded byte array to decode
+ * @param avroType The Avro schema defining the structure for decoding
+ * @return UnsafeRow containing the decoded data
+ * @throws UnsupportedOperationException if a field's data type is not supported for range
+ * scan decoding
+ */
+ override def decodePrefixKeyForRangeScan(bytes: Array[Byte]): UnsafeRow = {
+ val reader = new GenericDatumReader[GenericRecord](rangeScanAvroType)
+ val decoder = DecoderFactory.get().binaryDecoder(bytes, 0, bytes.length, null)
+ val record = reader.read(null, decoder)
+
+ val rowWriter = new UnsafeRowWriter(rangeScanKeyFieldsWithOrdinal.length)
+ rowWriter.resetRowWriter()
+
+ var fieldIdx = 0
+ rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
+ val field = fieldWithOrdinal._1
+
+ val markerBytes = record.get(fieldIdx).asInstanceOf[ByteBuffer].array()
+ val markerBuf = ByteBuffer.wrap(markerBytes)
+ markerBuf.order(ByteOrder.BIG_ENDIAN)
+ val marker = markerBuf.get()
+
+ if (marker == nullValMarker) {
+ rowWriter.setNullAt(idx)
+ } else {
+ field.dataType match {
+ case BooleanType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ rowWriter.write(idx, bytes(0) == 1)
+
+ case ByteType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.get())
+
+ case ShortType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.getShort())
+
+ case IntegerType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.getInt())
+
+ case LongType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ rowWriter.write(idx, valueBuf.getLong())
+
+ case FloatType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ if (marker == negativeValMarker) {
+ val floatVal = valueBuf.getFloat
+ val updatedVal = floatToRawIntBits(floatVal) ^ floatFlipBitMask
+ rowWriter.write(idx, intBitsToFloat(updatedVal))
+ } else {
+ rowWriter.write(idx, valueBuf.getFloat())
+ }
+
+ case DoubleType =>
+ val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
+ val valueBuf = ByteBuffer.wrap(bytes)
+ valueBuf.order(ByteOrder.BIG_ENDIAN)
+ if (marker == negativeValMarker) {
+ val doubleVal = valueBuf.getDouble
+ val updatedVal = doubleToRawLongBits(doubleVal) ^ doubleFlipBitMask
+ rowWriter.write(idx, longBitsToDouble(updatedVal))
+ } else {
+ rowWriter.write(idx, valueBuf.getDouble())
+ }
+
+ case _ => throw new UnsupportedOperationException(
+ s"Range scan decoding not supported for data type: ${field.dataType}")
+ }
+ }
+ fieldIdx += 2
+ }
+
+ rowWriter.getRow()
+ }
+
+ override def decodeValue(bytes: Array[Byte]): UnsafeRow =
+ decodeFromAvroToUnsafeRow(
+ bytes, avroEncoder.valueDeserializer, valueAvroType, valueProj)
+}
+
+abstract class RocksDBKeyStateEncoderBase(
+ useColumnFamilies: Boolean,
+ virtualColFamilyId: Option[Short] = None) extends RocksDBKeyStateEncoder {
+ def offsetForColFamilyPrefix: Int =
+ if (useColumnFamilies) VIRTUAL_COL_FAMILY_PREFIX_BYTES else 0
+
+ val out = new ByteArrayOutputStream
+ /**
+ * Get Byte Array for the virtual column family id that is used as prefix for
+ * key state rows.
+ */
+ override def getColumnFamilyIdBytes(): Array[Byte] = {
+ assert(useColumnFamilies, "Cannot return virtual Column Family Id Bytes" +
+ " because multiple Column is not supported for this encoder")
+ val encodedBytes = new Array[Byte](VIRTUAL_COL_FAMILY_PREFIX_BYTES)
+ Platform.putShort(encodedBytes, Platform.BYTE_ARRAY_OFFSET, virtualColFamilyId.get)
+ encodedBytes
+ }
+
+ /**
+ * Encode and put column family Id as a prefix to a pre-allocated byte array.
+ *
+ * @param numBytes - size of byte array to be created for storing key row (without
+ * column family prefix)
+ * @return Array[Byte] for an array byte to put encoded key bytes
+ * Int for a starting offset to put the encoded key bytes
+ */
+ protected def encodeColumnFamilyPrefix(numBytes: Int): (Array[Byte], Int) = {
+ val encodedBytes = new Array[Byte](numBytes + offsetForColFamilyPrefix)
+ var offset = Platform.BYTE_ARRAY_OFFSET
+ if (useColumnFamilies) {
+ Platform.putShort(encodedBytes, Platform.BYTE_ARRAY_OFFSET, virtualColFamilyId.get)
+ offset = Platform.BYTE_ARRAY_OFFSET + offsetForColFamilyPrefix
+ }
+ (encodedBytes, offset)
+ }
+
+ /**
+ * Get starting offset for decoding an encoded key byte array.
+ */
+ protected def decodeKeyStartOffset: Int = {
+ if (useColumnFamilies) {
+ Platform.BYTE_ARRAY_OFFSET + VIRTUAL_COL_FAMILY_PREFIX_BYTES
+ } else Platform.BYTE_ARRAY_OFFSET
+ }
+}
+
+object RocksDBStateEncoder extends Logging {
+ def getKeyEncoder(
+ dataEncoder: RocksDBDataEncoder,
+ keyStateEncoderSpec: KeyStateEncoderSpec,
+ useColumnFamilies: Boolean,
+ virtualColFamilyId: Option[Short] = None,
+ avroEnc: Option[AvroEncoder] = None): RocksDBKeyStateEncoder = {
+ // Return the key state encoder based on the requested type
+ keyStateEncoderSpec match {
+ case NoPrefixKeyStateEncoderSpec(keySchema) =>
+ new NoPrefixKeyStateEncoder(dataEncoder, keySchema, useColumnFamilies, virtualColFamilyId)
+
+ case PrefixKeyScanStateEncoderSpec(keySchema, numColsPrefixKey) =>
+ new PrefixKeyScanStateEncoder(dataEncoder, keySchema, numColsPrefixKey,
+ useColumnFamilies, virtualColFamilyId)
+
+ case RangeKeyScanStateEncoderSpec(keySchema, orderingOrdinals) =>
+ new RangeKeyScanStateEncoder(dataEncoder, keySchema, orderingOrdinals,
+ useColumnFamilies, virtualColFamilyId)
+
+ case _ =>
+ throw new IllegalArgumentException(s"Unsupported key state encoder spec: " +
+ s"$keyStateEncoderSpec")
+ }
+ }
+
+ def getValueEncoder(
+ dataEncoder: RocksDBDataEncoder,
+ valueSchema: StructType,
+ useMultipleValuesPerKey: Boolean,
+ avroEnc: Option[AvroEncoder] = None): RocksDBValueStateEncoder = {
+ if (useMultipleValuesPerKey) {
+ new MultiValuedStateEncoder(dataEncoder, valueSchema)
+ } else {
+ new SingleValueStateEncoder(dataEncoder, valueSchema)
+ }
+ }
+
+ def getColumnFamilyIdBytes(virtualColFamilyId: Short): Array[Byte] = {
+ val encodedBytes = new Array[Byte](VIRTUAL_COL_FAMILY_PREFIX_BYTES)
+ Platform.putShort(encodedBytes, Platform.BYTE_ARRAY_OFFSET, virtualColFamilyId)
+ encodedBytes
+ }
+}
+
+/**
+ * RocksDB Key Encoder for UnsafeRow that supports prefix scan
+ *
+ * @param dataEncoder - the encoder that handles actual encoding/decoding of data
+ * @param keySchema - schema of the key to be encoded
+ * @param numColsPrefixKey - number of columns to be used for prefix key
+ * @param useColumnFamilies - if column family is enabled for this encoder
+ * @param avroEnc - if Avro encoding is specified for this StateEncoder, this encoder will
+ * be defined
+ */
+class PrefixKeyScanStateEncoder(
+ dataEncoder: RocksDBDataEncoder,
+ keySchema: StructType,
+ numColsPrefixKey: Int,
+ useColumnFamilies: Boolean = false,
+ virtualColFamilyId: Option[Short] = None)
+ extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) with Logging {
+
+ private val prefixKeyFieldsWithIdx: Seq[(StructField, Int)] = {
+ keySchema.zipWithIndex.take(numColsPrefixKey)
+ }
+
+ private val remainingKeyFieldsWithIdx: Seq[(StructField, Int)] = {
+ keySchema.zipWithIndex.drop(numColsPrefixKey)
+ }
+
+ private val prefixKeyProjection: UnsafeProjection = {
+ val refs = prefixKeyFieldsWithIdx.map(x => BoundReference(x._2, x._1.dataType, x._1.nullable))
+ UnsafeProjection.create(refs)
+ }
+
+ private val remainingKeyProjection: UnsafeProjection = {
+ val refs = remainingKeyFieldsWithIdx.map(x =>
+ BoundReference(x._2, x._1.dataType, x._1.nullable))
+ UnsafeProjection.create(refs)
+ }
+
+ // Prefix Key schema and projection definitions used by the Avro Serializers
+ // and Deserializers
+ private val prefixKeySchema = StructType(keySchema.take(numColsPrefixKey))
+ private lazy val prefixKeyAvroType = SchemaConverters.toAvroType(prefixKeySchema)
+ private val prefixKeyProj = UnsafeProjection.create(prefixKeySchema)
+
+ // Remaining Key schema and projection definitions used by the Avro Serializers
+ // and Deserializers
+ private val remainingKeySchema = StructType(keySchema.drop(numColsPrefixKey))
+ private lazy val remainingKeyAvroType = SchemaConverters.toAvroType(remainingKeySchema)
+ private val remainingKeyProj = UnsafeProjection.create(remainingKeySchema)
+
+ // This is quite simple to do - just bind sequentially, as we don't change the order.
+ private val restoreKeyProjection: UnsafeProjection = UnsafeProjection.create(keySchema)
+
+ // Reusable objects
+ private val joinedRowOnKey = new JoinedRow()
+
+ override def encodeKey(row: UnsafeRow): Array[Byte] = {
+ val prefixKeyEncoded = dataEncoder.encodeKey(extractPrefixKey(row))
+ val remainingEncoded = dataEncoder.encodeRemainingKey(remainingKeyProjection(row))
+
+ val (encodedBytes, startingOffset) = encodeColumnFamilyPrefix(
+ prefixKeyEncoded.length + remainingEncoded.length + 4
+ )
+
+ Platform.putInt(encodedBytes, startingOffset, prefixKeyEncoded.length)
+ Platform.copyMemory(prefixKeyEncoded, Platform.BYTE_ARRAY_OFFSET,
+ encodedBytes, startingOffset + 4, prefixKeyEncoded.length)
+ // NOTE: We don't put the length of remainingEncoded as we can calculate later
+ // on deserialization.
+ Platform.copyMemory(remainingEncoded, Platform.BYTE_ARRAY_OFFSET,
+ encodedBytes, startingOffset + 4 + prefixKeyEncoded.length,
+ remainingEncoded.length)
+
+ encodedBytes
+ }
+
+ override def decodeKey(keyBytes: Array[Byte]): UnsafeRow = {
+ val prefixKeyEncodedLen = Platform.getInt(keyBytes, decodeKeyStartOffset)
+ val prefixKeyEncoded = new Array[Byte](prefixKeyEncodedLen)
+ Platform.copyMemory(keyBytes, decodeKeyStartOffset + 4,
+ prefixKeyEncoded, Platform.BYTE_ARRAY_OFFSET, prefixKeyEncodedLen)
+
+ // Here we calculate the remainingKeyEncodedLen leveraging the length of keyBytes
+ val remainingKeyEncodedLen = keyBytes.length - 4 - prefixKeyEncodedLen -
+ offsetForColFamilyPrefix
+
+ val remainingKeyEncoded = new Array[Byte](remainingKeyEncodedLen)
+ Platform.copyMemory(keyBytes, decodeKeyStartOffset + 4 + prefixKeyEncodedLen,
+ remainingKeyEncoded, Platform.BYTE_ARRAY_OFFSET, remainingKeyEncodedLen)
+
+ val prefixKeyDecoded = dataEncoder.decodeKey(
+ prefixKeyEncoded)
+ val remainingKeyDecoded = dataEncoder.decodeRemainingKey(remainingKeyEncoded)
+
+ restoreKeyProjection(joinedRowOnKey.withLeft(prefixKeyDecoded).withRight(remainingKeyDecoded))
+ }
+
+ private def extractPrefixKey(key: UnsafeRow): UnsafeRow = {
+ prefixKeyProjection(key)
+ }
+
+ override def encodePrefixKey(prefixKey: UnsafeRow): Array[Byte] = {
+ val prefixKeyEncoded = dataEncoder.encodeKey(prefixKey)
+ val (prefix, startingOffset) = encodeColumnFamilyPrefix(
+ prefixKeyEncoded.length + 4
+ )
+
+ Platform.putInt(prefix, startingOffset, prefixKeyEncoded.length)
+ Platform.copyMemory(prefixKeyEncoded, Platform.BYTE_ARRAY_OFFSET, prefix,
+ startingOffset + 4, prefixKeyEncoded.length)
+ prefix
+ }
+
+ override def supportPrefixKeyScan: Boolean = true
+}
+
+/**
+ * RocksDB Key Encoder for UnsafeRow that supports range scan for fixed size fields
+ *
+ * To encode a row for range scan, we first project the orderingOrdinals from the oridinal
+ * UnsafeRow into another UnsafeRow; we then rewrite that new UnsafeRow's fields in BIG_ENDIAN
+ * to allow for scanning keys in sorted order using the byte-wise comparison method that
+ * RocksDB uses.
+ *
+ * Then, for the rest of the fields, we project those into another UnsafeRow.
+ * We then effectively join these two UnsafeRows together, and finally take those bytes
+ * to get the resulting row.
+ *
+ * We cannot support variable sized fields in the range scan because the UnsafeRow format
+ * stores variable sized fields as offset and length pointers to the actual values,
+ * thereby changing the required ordering.
+ *
+ * Note that we also support "null" values being passed for these fixed size fields. We prepend
+ * a single byte to indicate whether the column value is null or not. We cannot change the
+ * nullability on the UnsafeRow itself as the expected ordering would change if non-first
+ * columns are marked as null. If the first col is null, those entries will appear last in
+ * the iterator. If non-first columns are null, ordering based on the previous columns will
+ * still be honored. For rows with null column values, ordering for subsequent columns
+ * will also be maintained within those set of rows. We use the same byte to also encode whether
+ * the value is negative or not. For negative float/double values, we flip all the bits to ensure
+ * the right lexicographical ordering. For the rationale around this, please check the link
+ * here: https://en.wikipedia.org/wiki/IEEE_754#Design_rationale
+ *
+ * @param dataEncoder - the encoder that handles the actual encoding/decoding of data
+ * @param keySchema - schema of the key to be encoded
+ * @param orderingOrdinals - the ordinals for which the range scan is constructed
+ * @param useColumnFamilies - if column family is enabled for this encoder
+ * @param avroEnc - if Avro encoding is specified for this StateEncoder, this encoder will
+ * be defined
+ */
+class RangeKeyScanStateEncoder(
+ dataEncoder: RocksDBDataEncoder,
+ keySchema: StructType,
+ orderingOrdinals: Seq[Int],
+ useColumnFamilies: Boolean = false,
+ virtualColFamilyId: Option[Short] = None)
+ extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) with Logging {
- case _ => throw new UnsupportedOperationException(
- s"Range scan encoding not supported for data type: ${field.dataType}")
- }
- }
- fieldIdx += 2
+ private val rangeScanKeyFieldsWithOrdinal: Seq[(StructField, Int)] = {
+ orderingOrdinals.map { ordinal =>
+ val field = keySchema(ordinal)
+ (field, ordinal)
}
-
- out.reset()
- val writer = new GenericDatumWriter[GenericRecord](rangeScanAvroType)
- val encoder = EncoderFactory.get().binaryEncoder(out, null)
- writer.write(record, encoder)
- encoder.flush()
- out.toByteArray
}
- /**
- * Decodes an Avro-encoded byte array back into an UnsafeRow for range scan operations.
- *
- * This method reverses the encoding process performed by encodePrefixKeyForRangeScan:
- * - Reads the marker byte to determine null status or sign
- * - Reconstructs the original values from big-endian format
- * - Handles special cases for floating point numbers by reversing bit manipulations
- *
- * The decoding process preserves the original data types and values, including:
- * - Null values marked by nullValMarker
- * - Sign information for numeric types
- * - Proper restoration of negative floating point values
- *
- * @param bytes The Avro-encoded byte array to decode
- * @param avroType The Avro schema defining the structure for decoding
- * @return UnsafeRow containing the decoded data
- * @throws UnsupportedOperationException if a field's data type is not supported for range
- * scan decoding
- */
- def decodePrefixKeyForRangeScan(
- bytes: Array[Byte],
- avroType: Schema): UnsafeRow = {
+ private def isFixedSize(dataType: DataType): Boolean = dataType match {
+ case _: ByteType | _: BooleanType | _: ShortType | _: IntegerType | _: LongType |
+ _: FloatType | _: DoubleType => true
+ case _ => false
+ }
- val reader = new GenericDatumReader[GenericRecord](avroType)
- val decoder = DecoderFactory.get().binaryDecoder(bytes, 0, bytes.length, null)
- val record = reader.read(null, decoder)
+ // verify that only fixed sized columns are used for ordering
+ rangeScanKeyFieldsWithOrdinal.foreach { case (field, ordinal) =>
+ if (!isFixedSize(field.dataType)) {
+ // NullType is technically fixed size, but not supported for ordering
+ if (field.dataType == NullType) {
+ throw StateStoreErrors.nullTypeOrderingColsNotSupported(field.name, ordinal.toString)
+ } else {
+ throw StateStoreErrors.variableSizeOrderingColsNotSupported(field.name, ordinal.toString)
+ }
+ }
+ }
- val rowWriter = new UnsafeRowWriter(rangeScanKeyFieldsWithOrdinal.length)
- rowWriter.resetRowWriter()
+ private val remainingKeyFieldsWithOrdinal: Seq[(StructField, Int)] = {
+ 0.to(keySchema.length - 1).diff(orderingOrdinals).map { ordinal =>
+ val field = keySchema(ordinal)
+ (field, ordinal)
+ }
+ }
- var fieldIdx = 0
- rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
- val field = fieldWithOrdinal._1
+ private val rangeScanKeyProjection: UnsafeProjection = {
+ val refs = rangeScanKeyFieldsWithOrdinal.map(x =>
+ BoundReference(x._2, x._1.dataType, x._1.nullable))
+ UnsafeProjection.create(refs)
+ }
- val markerBytes = record.get(fieldIdx).asInstanceOf[ByteBuffer].array()
- val markerBuf = ByteBuffer.wrap(markerBytes)
- markerBuf.order(ByteOrder.BIG_ENDIAN)
- val marker = markerBuf.get()
+ private val remainingKeyProjection: UnsafeProjection = {
+ val refs = remainingKeyFieldsWithOrdinal.map(x =>
+ BoundReference(x._2, x._1.dataType, x._1.nullable))
+ UnsafeProjection.create(refs)
+ }
- if (marker == nullValMarker) {
- rowWriter.setNullAt(idx)
+ // The original schema that we might get could be:
+ // [foo, bar, baz, buzz]
+ // We might order by bar and buzz, leading to:
+ // [bar, buzz, foo, baz]
+ // We need to create a projection that sends, for example, the buzz at index 1 to index
+ // 3. Thus, for every record in the original schema, we compute where it would be in
+ // the joined row and created a projection based on that.
+ private val restoreKeyProjection: UnsafeProjection = {
+ val refs = keySchema.zipWithIndex.map { case (field, originalOrdinal) =>
+ val ordinalInJoinedRow = if (orderingOrdinals.contains(originalOrdinal)) {
+ orderingOrdinals.indexOf(originalOrdinal)
} else {
- field.dataType match {
- case BooleanType =>
- val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
- rowWriter.write(idx, bytes(0) == 1)
+ orderingOrdinals.length +
+ remainingKeyFieldsWithOrdinal.indexWhere(_._2 == originalOrdinal)
+ }
- case ByteType =>
- val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
- val valueBuf = ByteBuffer.wrap(bytes)
- valueBuf.order(ByteOrder.BIG_ENDIAN)
- rowWriter.write(idx, valueBuf.get())
+ BoundReference(ordinalInJoinedRow, field.dataType, field.nullable)
+ }
+ UnsafeProjection.create(refs)
+ }
- case ShortType =>
- val bytes = record.get(fieldIdx + 1).asInstanceOf[ByteBuffer].array()
- val valueBuf = ByteBuffer.wrap(bytes)
- valueBuf.order(ByteOrder.BIG_ENDIAN)
- rowWriter.write(idx, valueBuf.getShort())
+ private val rangeScanAvroSchema = StateStoreColumnFamilySchemaUtils.convertForRangeScan(
+ StructType(rangeScanKeyFieldsWithOrdinal.map(_._1).toArray))
Review Comment:
these are also no longer used. Please clean them up
--
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-50017][SQL] Support Avro encoding for TransformWithState operator [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on code in PR #48944:
URL: https://github.com/apache/spark/pull/48944#discussion_r1887868528
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceTransformWithStateSuite.scala:
##########
@@ -126,8 +126,7 @@ class SessionGroupsStatefulProcessorWithTTL extends
* Test suite to verify integration of state data source reader with the transformWithState operator
*/
class StateDataSourceTransformWithStateSuite extends StateStoreMetricsTest
- with AlsoTestWithEncodingTypes
- with AlsoTestWithChangelogCheckpointingEnabled {
Review Comment:
Yeah, AlsoTestWithRocksDBFeatures does test that - not sure why this is showing up in the diff.
--
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-50599][SQL] Create the DataEncoder trait that allows for Avro and UnsafeRow encoding [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR closed pull request #48944: [SPARK-50599][SQL] Create the DataEncoder trait that allows for Avro and UnsafeRow encoding
URL: https://github.com/apache/spark/pull/48944
--
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-50599][SQL] Create the DataEncoder trait that allows for Avro and UnsafeRow encoding [spark]
Posted by "brkyvz (via GitHub)" <gi...@apache.org>.
brkyvz commented on code in PR #48944:
URL: https://github.com/apache/spark/pull/48944#discussion_r1889211051
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala:
##########
@@ -51,93 +51,138 @@ sealed trait RocksDBValueStateEncoder {
def decodeValues(valueBytes: Array[Byte]): Iterator[UnsafeRow]
}
-abstract class RocksDBKeyStateEncoderBase(
- useColumnFamilies: Boolean,
- virtualColFamilyId: Option[Short] = None) extends RocksDBKeyStateEncoder {
- def offsetForColFamilyPrefix: Int =
- if (useColumnFamilies) VIRTUAL_COL_FAMILY_PREFIX_BYTES else 0
+/**
+ * The DataEncoder can encode UnsafeRows into raw bytes in two ways:
+ * - Using the direct byte layout of the UnsafeRow
+ * - Converting the UnsafeRow into an Avro row, and encoding that
+ * In both of these cases, the raw bytes that are written into RockDB have
+ * headers, footers and other metadata, but they also have data that is provided
+ * by the callers. The metadata in each row does not need to be written as Avro or UnsafeRow,
+ * but the actual data provided by the caller does.
+ * The classes that use this trait require specialized partial encoding which makes them much
+ * easier to cache and use, which is why each DataEncoder deals with multiple schemas.
+ */
+trait DataEncoder {
Review Comment:
Awesome scaladoc throughout the class. Thank you!
--
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-50599][SQL] Create the DataEncoder trait that allows for Avro and UnsafeRow encoding [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #48944:
URL: https://github.com/apache/spark/pull/48944#issuecomment-2550521437
The comment link does not show the comment, which I assume the comment is outdated via new commit.
I'm merging on behalf of @brkyvz . Thanks! Merging to master.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-50017][SQL] Support Avro encoding for TransformWithState operator [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on code in PR #48944:
URL: https://github.com/apache/spark/pull/48944#discussion_r1887871951
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala:
##########
@@ -219,469 +222,402 @@ object RocksDBStateEncoder extends Logging {
}
}
-/**
- * RocksDB Key Encoder for UnsafeRow that supports prefix scan
- *
- * @param keySchema - schema of the key to be encoded
- * @param numColsPrefixKey - number of columns to be used for prefix key
- * @param useColumnFamilies - if column family is enabled for this encoder
- * @param avroEnc - if Avro encoding is specified for this StateEncoder, this encoder will
- * be defined
- */
-class PrefixKeyScanStateEncoder(
- keySchema: StructType,
- numColsPrefixKey: Int,
- useColumnFamilies: Boolean = false,
- virtualColFamilyId: Option[Short] = None,
- avroEnc: Option[AvroEncoder] = None)
- extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) {
-
- import RocksDBStateEncoder._
-
- private val usingAvroEncoding = avroEnc.isDefined
- private val prefixKeyFieldsWithIdx: Seq[(StructField, Int)] = {
- keySchema.zipWithIndex.take(numColsPrefixKey)
- }
+class UnsafeRowDataEncoder(
+ keyStateEncoderSpec: KeyStateEncoderSpec,
+ valueSchema: StructType) extends RocksDBDataEncoder(keyStateEncoderSpec, valueSchema) {
- private val remainingKeyFieldsWithIdx: Seq[(StructField, Int)] = {
- keySchema.zipWithIndex.drop(numColsPrefixKey)
+ override def encodeKey(row: UnsafeRow): Array[Byte] = {
+ encodeUnsafeRow(row)
}
- private val prefixKeyProjection: UnsafeProjection = {
- val refs = prefixKeyFieldsWithIdx.map(x => BoundReference(x._2, x._1.dataType, x._1.nullable))
- UnsafeProjection.create(refs)
+ override def encodeRemainingKey(row: UnsafeRow): Array[Byte] = {
+ encodeUnsafeRow(row)
}
- private val remainingKeyProjection: UnsafeProjection = {
- val refs = remainingKeyFieldsWithIdx.map(x =>
- BoundReference(x._2, x._1.dataType, x._1.nullable))
- UnsafeProjection.create(refs)
- }
+ override def encodePrefixKeyForRangeScan(row: UnsafeRow): Array[Byte] = {
+ assert(keyStateEncoderSpec.isInstanceOf[RangeKeyScanStateEncoderSpec])
+ val rsk = keyStateEncoderSpec.asInstanceOf[RangeKeyScanStateEncoderSpec]
+ val rangeScanKeyFieldsWithOrdinal = rsk.orderingOrdinals.map { ordinal =>
+ val field = rsk.keySchema(ordinal)
+ (field, ordinal)
+ }
+ val writer = new UnsafeRowWriter(rsk.orderingOrdinals.length)
+ writer.resetRowWriter()
+ rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
+ val field = fieldWithOrdinal._1
+ val value = row.get(idx, field.dataType)
+ // Note that we cannot allocate a smaller buffer here even if the value is null
+ // because the effective byte array is considered variable size and needs to have
+ // the same size across all rows for the ordering to work as expected.
+ val bbuf = ByteBuffer.allocate(field.dataType.defaultSize + 1)
+ bbuf.order(ByteOrder.BIG_ENDIAN)
+ if (value == null) {
+ bbuf.put(nullValMarker)
+ writer.write(idx, bbuf.array())
+ } else {
+ field.dataType match {
+ case BooleanType =>
+ case ByteType =>
+ val byteVal = value.asInstanceOf[Byte]
+ val signCol = if (byteVal < 0) {
+ negativeValMarker
+ } else {
+ positiveValMarker
+ }
+ bbuf.put(signCol)
+ bbuf.put(byteVal)
+ writer.write(idx, bbuf.array())
- // Prefix Key schema and projection definitions used by the Avro Serializers
- // and Deserializers
- private val prefixKeySchema = StructType(keySchema.take(numColsPrefixKey))
- private lazy val prefixKeyAvroType = SchemaConverters.toAvroType(prefixKeySchema)
- private val prefixKeyProj = UnsafeProjection.create(prefixKeySchema)
+ case ShortType =>
+ val shortVal = value.asInstanceOf[Short]
+ val signCol = if (shortVal < 0) {
+ negativeValMarker
+ } else {
+ positiveValMarker
+ }
+ bbuf.put(signCol)
+ bbuf.putShort(shortVal)
+ writer.write(idx, bbuf.array())
- // Remaining Key schema and projection definitions used by the Avro Serializers
- // and Deserializers
- private val remainingKeySchema = StructType(keySchema.drop(numColsPrefixKey))
- private lazy val remainingKeyAvroType = SchemaConverters.toAvroType(remainingKeySchema)
- private val remainingKeyProj = UnsafeProjection.create(remainingKeySchema)
+ case IntegerType =>
+ val intVal = value.asInstanceOf[Int]
+ val signCol = if (intVal < 0) {
+ negativeValMarker
+ } else {
+ positiveValMarker
+ }
+ bbuf.put(signCol)
+ bbuf.putInt(intVal)
+ writer.write(idx, bbuf.array())
- // This is quite simple to do - just bind sequentially, as we don't change the order.
- private val restoreKeyProjection: UnsafeProjection = UnsafeProjection.create(keySchema)
+ case LongType =>
+ val longVal = value.asInstanceOf[Long]
+ val signCol = if (longVal < 0) {
+ negativeValMarker
+ } else {
+ positiveValMarker
+ }
+ bbuf.put(signCol)
+ bbuf.putLong(longVal)
+ writer.write(idx, bbuf.array())
- // Reusable objects
- private val joinedRowOnKey = new JoinedRow()
+ case FloatType =>
+ val floatVal = value.asInstanceOf[Float]
+ val rawBits = floatToRawIntBits(floatVal)
+ // perform sign comparison using bit manipulation to ensure NaN values are handled
+ // correctly
+ if ((rawBits & floatSignBitMask) != 0) {
+ // for negative values, we need to flip all the bits to ensure correct ordering
+ val updatedVal = rawBits ^ floatFlipBitMask
+ bbuf.put(negativeValMarker)
+ // convert the bits back to float
+ bbuf.putFloat(intBitsToFloat(updatedVal))
+ } else {
+ bbuf.put(positiveValMarker)
+ bbuf.putFloat(floatVal)
+ }
+ writer.write(idx, bbuf.array())
- override def encodeKey(row: UnsafeRow): Array[Byte] = {
- val (prefixKeyEncoded, remainingEncoded) = if (usingAvroEncoding) {
- (
- encodeUnsafeRowToAvro(
- extractPrefixKey(row),
- avroEnc.get.keySerializer,
- prefixKeyAvroType,
- out
- ),
- encodeUnsafeRowToAvro(
- remainingKeyProjection(row),
- avroEnc.get.suffixKeySerializer.get,
- remainingKeyAvroType,
- out
- )
- )
- } else {
- (encodeUnsafeRow(extractPrefixKey(row)), encodeUnsafeRow(remainingKeyProjection(row)))
+ case DoubleType =>
+ val doubleVal = value.asInstanceOf[Double]
+ val rawBits = doubleToRawLongBits(doubleVal)
+ // perform sign comparison using bit manipulation to ensure NaN values are handled
+ // correctly
+ if ((rawBits & doubleSignBitMask) != 0) {
+ // for negative values, we need to flip all the bits to ensure correct ordering
+ val updatedVal = rawBits ^ doubleFlipBitMask
+ bbuf.put(negativeValMarker)
+ // convert the bits back to double
+ bbuf.putDouble(longBitsToDouble(updatedVal))
+ } else {
+ bbuf.put(positiveValMarker)
+ bbuf.putDouble(doubleVal)
+ }
+ writer.write(idx, bbuf.array())
+ }
+ }
}
- val (encodedBytes, startingOffset) = encodeColumnFamilyPrefix(
- prefixKeyEncoded.length + remainingEncoded.length + 4
- )
+ encodeUnsafeRow(writer.getRow())
+ }
- Platform.putInt(encodedBytes, startingOffset, prefixKeyEncoded.length)
- Platform.copyMemory(prefixKeyEncoded, Platform.BYTE_ARRAY_OFFSET,
- encodedBytes, startingOffset + 4, prefixKeyEncoded.length)
- // NOTE: We don't put the length of remainingEncoded as we can calculate later
- // on deserialization.
- Platform.copyMemory(remainingEncoded, Platform.BYTE_ARRAY_OFFSET,
- encodedBytes, startingOffset + 4 + prefixKeyEncoded.length,
- remainingEncoded.length)
-
- encodedBytes
- }
-
- override def decodeKey(keyBytes: Array[Byte]): UnsafeRow = {
- val prefixKeyEncodedLen = Platform.getInt(keyBytes, decodeKeyStartOffset)
- val prefixKeyEncoded = new Array[Byte](prefixKeyEncodedLen)
- Platform.copyMemory(keyBytes, decodeKeyStartOffset + 4,
- prefixKeyEncoded, Platform.BYTE_ARRAY_OFFSET, prefixKeyEncodedLen)
+ override def encodeValue(row: UnsafeRow): Array[Byte] = encodeUnsafeRow(row)
- // Here we calculate the remainingKeyEncodedLen leveraging the length of keyBytes
- val remainingKeyEncodedLen = keyBytes.length - 4 - prefixKeyEncodedLen -
- offsetForColFamilyPrefix
-
- val remainingKeyEncoded = new Array[Byte](remainingKeyEncodedLen)
- Platform.copyMemory(keyBytes, decodeKeyStartOffset + 4 + prefixKeyEncodedLen,
- remainingKeyEncoded, Platform.BYTE_ARRAY_OFFSET, remainingKeyEncodedLen)
-
- val (prefixKeyDecoded, remainingKeyDecoded) = if (usingAvroEncoding) {
- (
- decodeFromAvroToUnsafeRow(
- prefixKeyEncoded,
- avroEnc.get.keyDeserializer,
- prefixKeyAvroType,
- prefixKeyProj
- ),
- decodeFromAvroToUnsafeRow(
- remainingKeyEncoded,
- avroEnc.get.suffixKeyDeserializer.get,
- remainingKeyAvroType,
- remainingKeyProj
- )
- )
- } else {
- (decodeToUnsafeRow(prefixKeyEncoded, numFields = numColsPrefixKey),
- decodeToUnsafeRow(remainingKeyEncoded, numFields = keySchema.length - numColsPrefixKey))
+ override def decodeKey(bytes: Array[Byte]): UnsafeRow = {
+ keyStateEncoderSpec match {
+ case NoPrefixKeyStateEncoderSpec(_) =>
+ decodeToUnsafeRow(bytes, reusedKeyRow)
+ case PrefixKeyScanStateEncoderSpec(_, numColsPrefixKey) =>
+ decodeToUnsafeRow(bytes, numFields = numColsPrefixKey)
+ case _ => throw unsupportedOperationForKeyStateEncoder("decodeKey")
}
-
- restoreKeyProjection(joinedRowOnKey.withLeft(prefixKeyDecoded).withRight(remainingKeyDecoded))
}
- private def extractPrefixKey(key: UnsafeRow): UnsafeRow = {
- prefixKeyProjection(key)
+ override def decodeRemainingKey(bytes: Array[Byte]): UnsafeRow = {
+ keyStateEncoderSpec match {
+ case PrefixKeyScanStateEncoderSpec(_, numColsPrefixKey) =>
+ decodeToUnsafeRow(bytes, numFields = numColsPrefixKey)
+ case RangeKeyScanStateEncoderSpec(_, orderingOrdinals) =>
+ decodeToUnsafeRow(bytes, keySchema.length - orderingOrdinals.length)
+ case _ => throw unsupportedOperationForKeyStateEncoder("decodeRemainingKey")
+ }
}
- override def encodePrefixKey(prefixKey: UnsafeRow): Array[Byte] = {
- val prefixKeyEncoded = if (usingAvroEncoding) {
- encodeUnsafeRowToAvro(prefixKey, avroEnc.get.keySerializer, prefixKeyAvroType, out)
- } else {
- encodeUnsafeRow(prefixKey)
+ override def decodePrefixKeyForRangeScan(bytes: Array[Byte]): UnsafeRow = {
+ assert(keyStateEncoderSpec.isInstanceOf[RangeKeyScanStateEncoderSpec])
+ val rsk = keyStateEncoderSpec.asInstanceOf[RangeKeyScanStateEncoderSpec]
+ val writer = new UnsafeRowWriter(rsk.orderingOrdinals.length)
+ val rangeScanKeyFieldsWithOrdinal = rsk.orderingOrdinals.map { ordinal =>
+ val field = rsk.keySchema(ordinal)
+ (field, ordinal)
}
- val (prefix, startingOffset) = encodeColumnFamilyPrefix(
- prefixKeyEncoded.length + 4
- )
-
- Platform.putInt(prefix, startingOffset, prefixKeyEncoded.length)
- Platform.copyMemory(prefixKeyEncoded, Platform.BYTE_ARRAY_OFFSET, prefix,
- startingOffset + 4, prefixKeyEncoded.length)
- prefix
- }
+ writer.resetRowWriter()
+ val row = decodeToUnsafeRow(bytes, numFields = rsk.orderingOrdinals.length)
+ rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
+ val field = fieldWithOrdinal._1
- override def supportPrefixKeyScan: Boolean = true
-}
+ val value = row.getBinary(idx)
+ val bbuf = ByteBuffer.wrap(value.asInstanceOf[Array[Byte]])
+ bbuf.order(ByteOrder.BIG_ENDIAN)
+ val isNullOrSignCol = bbuf.get()
+ if (isNullOrSignCol == nullValMarker) {
+ // set the column to null and skip reading the next byte(s)
+ writer.setNullAt(idx)
+ } else {
+ field.dataType match {
+ case BooleanType =>
+ case ByteType =>
+ writer.write(idx, bbuf.get)
-/**
- * RocksDB Key Encoder for UnsafeRow that supports range scan for fixed size fields
- *
- * To encode a row for range scan, we first project the orderingOrdinals from the oridinal
- * UnsafeRow into another UnsafeRow; we then rewrite that new UnsafeRow's fields in BIG_ENDIAN
- * to allow for scanning keys in sorted order using the byte-wise comparison method that
- * RocksDB uses.
- *
- * Then, for the rest of the fields, we project those into another UnsafeRow.
- * We then effectively join these two UnsafeRows together, and finally take those bytes
- * to get the resulting row.
- *
- * We cannot support variable sized fields in the range scan because the UnsafeRow format
- * stores variable sized fields as offset and length pointers to the actual values,
- * thereby changing the required ordering.
- *
- * Note that we also support "null" values being passed for these fixed size fields. We prepend
- * a single byte to indicate whether the column value is null or not. We cannot change the
- * nullability on the UnsafeRow itself as the expected ordering would change if non-first
- * columns are marked as null. If the first col is null, those entries will appear last in
- * the iterator. If non-first columns are null, ordering based on the previous columns will
- * still be honored. For rows with null column values, ordering for subsequent columns
- * will also be maintained within those set of rows. We use the same byte to also encode whether
- * the value is negative or not. For negative float/double values, we flip all the bits to ensure
- * the right lexicographical ordering. For the rationale around this, please check the link
- * here: https://en.wikipedia.org/wiki/IEEE_754#Design_rationale
- *
- * @param keySchema - schema of the key to be encoded
- * @param orderingOrdinals - the ordinals for which the range scan is constructed
- * @param useColumnFamilies - if column family is enabled for this encoder
- * @param avroEnc - if Avro encoding is specified for this StateEncoder, this encoder will
- * be defined
- */
-class RangeKeyScanStateEncoder(
- keySchema: StructType,
- orderingOrdinals: Seq[Int],
- useColumnFamilies: Boolean = false,
- virtualColFamilyId: Option[Short] = None,
- avroEnc: Option[AvroEncoder] = None)
- extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) with Logging {
+ case ShortType =>
+ writer.write(idx, bbuf.getShort)
- import RocksDBStateEncoder._
+ case IntegerType =>
+ writer.write(idx, bbuf.getInt)
- private val rangeScanKeyFieldsWithOrdinal: Seq[(StructField, Int)] = {
- orderingOrdinals.map { ordinal =>
- val field = keySchema(ordinal)
- (field, ordinal)
- }
- }
+ case LongType =>
+ writer.write(idx, bbuf.getLong)
- private def isFixedSize(dataType: DataType): Boolean = dataType match {
- case _: ByteType | _: BooleanType | _: ShortType | _: IntegerType | _: LongType |
- _: FloatType | _: DoubleType => true
- case _ => false
- }
+ case FloatType =>
+ if (isNullOrSignCol == negativeValMarker) {
+ // if the number is negative, get the raw binary bits for the float
+ // and flip the bits back
+ val updatedVal = floatToRawIntBits(bbuf.getFloat) ^ floatFlipBitMask
+ writer.write(idx, intBitsToFloat(updatedVal))
+ } else {
+ writer.write(idx, bbuf.getFloat)
+ }
- // verify that only fixed sized columns are used for ordering
- rangeScanKeyFieldsWithOrdinal.foreach { case (field, ordinal) =>
- if (!isFixedSize(field.dataType)) {
- // NullType is technically fixed size, but not supported for ordering
- if (field.dataType == NullType) {
- throw StateStoreErrors.nullTypeOrderingColsNotSupported(field.name, ordinal.toString)
- } else {
- throw StateStoreErrors.variableSizeOrderingColsNotSupported(field.name, ordinal.toString)
+ case DoubleType =>
+ if (isNullOrSignCol == negativeValMarker) {
+ // if the number is negative, get the raw binary bits for the double
+ // and flip the bits back
+ val updatedVal = doubleToRawLongBits(bbuf.getDouble) ^ doubleFlipBitMask
+ writer.write(idx, longBitsToDouble(updatedVal))
+ } else {
+ writer.write(idx, bbuf.getDouble)
+ }
+ }
}
}
+ writer.getRow()
}
- private val remainingKeyFieldsWithOrdinal: Seq[(StructField, Int)] = {
- 0.to(keySchema.length - 1).diff(orderingOrdinals).map { ordinal =>
- val field = keySchema(ordinal)
- (field, ordinal)
- }
- }
+ override def decodeValue(bytes: Array[Byte]): UnsafeRow = decodeToUnsafeRow(bytes, reusedValueRow)
+}
- private val rangeScanKeyProjection: UnsafeProjection = {
- val refs = rangeScanKeyFieldsWithOrdinal.map(x =>
- BoundReference(x._2, x._1.dataType, x._1.nullable))
- UnsafeProjection.create(refs)
- }
+class AvroStateEncoder(
+ keyStateEncoderSpec: KeyStateEncoderSpec,
+ valueSchema: StructType) extends RocksDBDataEncoder(keyStateEncoderSpec, valueSchema)
+ with Logging {
- private val remainingKeyProjection: UnsafeProjection = {
- val refs = remainingKeyFieldsWithOrdinal.map(x =>
- BoundReference(x._2, x._1.dataType, x._1.nullable))
- UnsafeProjection.create(refs)
- }
+ private val avroEncoder = createAvroEnc(keyStateEncoderSpec, valueSchema)
+ // Avro schema used by the avro encoders
+ private lazy val keyAvroType: Schema = SchemaConverters.toAvroType(keySchema)
+ private lazy val keyProj = UnsafeProjection.create(keySchema)
- // The original schema that we might get could be:
- // [foo, bar, baz, buzz]
- // We might order by bar and buzz, leading to:
- // [bar, buzz, foo, baz]
- // We need to create a projection that sends, for example, the buzz at index 1 to index
- // 3. Thus, for every record in the original schema, we compute where it would be in
- // the joined row and created a projection based on that.
- private val restoreKeyProjection: UnsafeProjection = {
- val refs = keySchema.zipWithIndex.map { case (field, originalOrdinal) =>
- val ordinalInJoinedRow = if (orderingOrdinals.contains(originalOrdinal)) {
- orderingOrdinals.indexOf(originalOrdinal)
- } else {
- orderingOrdinals.length +
- remainingKeyFieldsWithOrdinal.indexWhere(_._2 == originalOrdinal)
- }
+ private lazy val valueAvroType: Schema = SchemaConverters.toAvroType(valueSchema)
+ private lazy val valueProj = UnsafeProjection.create(valueSchema)
- BoundReference(ordinalInJoinedRow, field.dataType, field.nullable)
- }
- UnsafeProjection.create(refs)
+ // Prefix Key schema and projection definitions used by the Avro Serializers
+ // and Deserializers
+ private lazy val prefixKeySchema = keyStateEncoderSpec match {
+ case PrefixKeyScanStateEncoderSpec(keySchema, numColsPrefixKey) =>
+ StructType(keySchema.take (numColsPrefixKey))
+ case _ => throw unsupportedOperationForKeyStateEncoder("prefixKeySchema")
+ }
+ private lazy val prefixKeyAvroType = SchemaConverters.toAvroType(prefixKeySchema)
+ private lazy val prefixKeyProj = UnsafeProjection.create(prefixKeySchema)
+
+ // Range Key schema nd projection definitions used by the Avro Serializers and
+ // Deserializers
+ private lazy val rangeScanKeyFieldsWithOrdinal = keyStateEncoderSpec match {
+ case RangeKeyScanStateEncoderSpec(keySchema, orderingOrdinals) =>
+ orderingOrdinals.map { ordinal =>
+ val field = keySchema(ordinal)
+ (field, ordinal)
+ }
+ case _ =>
+ throw unsupportedOperationForKeyStateEncoder("rangeScanKey")
}
- private val rangeScanAvroSchema = StateStoreColumnFamilySchemaUtils.convertForRangeScan(
+ private lazy val rangeScanAvroSchema = StateStoreColumnFamilySchemaUtils.convertForRangeScan(
StructType(rangeScanKeyFieldsWithOrdinal.map(_._1).toArray))
private lazy val rangeScanAvroType = SchemaConverters.toAvroType(rangeScanAvroSchema)
- private val rangeScanAvroProjection = UnsafeProjection.create(rangeScanAvroSchema)
+ private lazy val rangeScanAvroProjection = UnsafeProjection.create(rangeScanAvroSchema)
- // Existing remainder key schema stuff
- private val remainingKeySchema = StructType(
- 0.to(keySchema.length - 1).diff(orderingOrdinals).map(keySchema(_))
- )
+ // Existing remainder key schema definitions
+ // Remaining Key schema and projection definitions used by the Avro Serializers
+ // and Deserializers
+ private lazy val remainingKeySchema = keyStateEncoderSpec match {
+ case PrefixKeyScanStateEncoderSpec(keySchema, numColsPrefixKey) =>
+ StructType(keySchema.drop(numColsPrefixKey))
+ case RangeKeyScanStateEncoderSpec(keySchema, orderingOrdinals) =>
+ StructType(0.until(keySchema.length).diff(orderingOrdinals).map(keySchema(_)))
+ case _ => throw unsupportedOperationForKeyStateEncoder("remainingKeySchema")
+ }
private lazy val remainingKeyAvroType = SchemaConverters.toAvroType(remainingKeySchema)
- private val remainingKeyAvroProjection = UnsafeProjection.create(remainingKeySchema)
+ private lazy val remainingKeyAvroProjection = UnsafeProjection.create(remainingKeySchema)
- // Reusable objects
- private val joinedRowOnKey = new JoinedRow()
-
- private def extractPrefixKey(key: UnsafeRow): UnsafeRow = {
- rangeScanKeyProjection(key)
+ private def getAvroSerializer(schema: StructType): AvroSerializer = {
+ val avroType = SchemaConverters.toAvroType(schema)
+ new AvroSerializer(schema, avroType, nullable = false)
}
- // bit masks used for checking sign or flipping all bits for negative float/double values
- private val floatFlipBitMask = 0xFFFFFFFF
- private val floatSignBitMask = 0x80000000
+ private def getAvroDeserializer(schema: StructType): AvroDeserializer = {
+ val avroType = SchemaConverters.toAvroType(schema)
+ val avroOptions = AvroOptions(Map.empty)
+ new AvroDeserializer(avroType, schema,
+ avroOptions.datetimeRebaseModeInRead, avroOptions.useStableIdForUnionType,
+ avroOptions.stableIdPrefixForUnionType, avroOptions.recursiveFieldMaxDepth)
+ }
- private val doubleFlipBitMask = 0xFFFFFFFFFFFFFFFFL
- private val doubleSignBitMask = 0x8000000000000000L
+ /**
+ * Creates an AvroEncoder that handles both key and value serialization/deserialization.
+ * This method sets up the complete encoding infrastructure needed for state store operations.
+ *
+ * The encoder handles different key encoding specifications:
+ * - NoPrefixKeyStateEncoderSpec: Simple key encoding without prefix
+ * - PrefixKeyScanStateEncoderSpec: Keys with prefix for efficient scanning
+ * - RangeKeyScanStateEncoderSpec: Keys with ordering requirements for range scans
+ *
+ * For prefix scan cases, it also creates separate encoders for the suffix portion of keys.
+ *
+ * @param keyStateEncoderSpec Specification for how to encode keys
+ * @param valueSchema Schema for the values to be encoded
+ * @return An AvroEncoder containing all necessary serializers and deserializers
+ */
+ private def createAvroEnc(
+ keyStateEncoderSpec: KeyStateEncoderSpec,
+ valueSchema: StructType): AvroEncoder = {
+ val valueSerializer = getAvroSerializer(valueSchema)
+ val valueDeserializer = getAvroDeserializer(valueSchema)
+
+ // Get key schema based on encoder spec type
+ val keySchema = keyStateEncoderSpec match {
+ case NoPrefixKeyStateEncoderSpec(schema) =>
+ schema
+ case PrefixKeyScanStateEncoderSpec(schema, numColsPrefixKey) =>
+ StructType(schema.take(numColsPrefixKey))
+ case RangeKeyScanStateEncoderSpec(schema, orderingOrdinals) =>
+ val remainingSchema = {
+ 0.until(schema.length).diff(orderingOrdinals).map { ordinal =>
+ schema(ordinal)
+ }
+ }
+ StructType(remainingSchema)
+ }
- // Byte markers used to identify whether the value is null, negative or positive
- // To ensure sorted ordering, we use the lowest byte value for negative numbers followed by
- // positive numbers and then null values.
- private val negativeValMarker: Byte = 0x00.toByte
- private val positiveValMarker: Byte = 0x01.toByte
- private val nullValMarker: Byte = 0x02.toByte
-
- // Rewrite the unsafe row by replacing fixed size fields with BIG_ENDIAN encoding
- // using byte arrays.
- // To handle "null" values, we prepend a byte to the byte array indicating whether the value
- // is null or not. If the value is null, we write the null byte followed by zero bytes.
- // If the value is not null, we write the null byte followed by the value.
- // Note that setting null for the index on the unsafeRow is not feasible as it would change
- // the sorting order on iteration.
- // Also note that the same byte is used to indicate whether the value is negative or not.
- private def encodePrefixKeyForRangeScan(row: UnsafeRow): UnsafeRow = {
- val writer = new UnsafeRowWriter(orderingOrdinals.length)
- writer.resetRowWriter()
- rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
- val field = fieldWithOrdinal._1
- val value = row.get(idx, field.dataType)
- // Note that we cannot allocate a smaller buffer here even if the value is null
- // because the effective byte array is considered variable size and needs to have
- // the same size across all rows for the ordering to work as expected.
- val bbuf = ByteBuffer.allocate(field.dataType.defaultSize + 1)
- bbuf.order(ByteOrder.BIG_ENDIAN)
- if (value == null) {
- bbuf.put(nullValMarker)
- writer.write(idx, bbuf.array())
- } else {
- field.dataType match {
- case BooleanType =>
- case ByteType =>
- val byteVal = value.asInstanceOf[Byte]
- val signCol = if (byteVal < 0) {
- negativeValMarker
- } else {
- positiveValMarker
- }
- bbuf.put(signCol)
- bbuf.put(byteVal)
- writer.write(idx, bbuf.array())
+ // Handle suffix key schema for prefix scan case
+ val suffixKeySchema = keyStateEncoderSpec match {
+ case PrefixKeyScanStateEncoderSpec(schema, numColsPrefixKey) =>
+ Some(StructType(schema.drop(numColsPrefixKey)))
+ case _ =>
+ None
+ }
- case ShortType =>
- val shortVal = value.asInstanceOf[Short]
- val signCol = if (shortVal < 0) {
- negativeValMarker
- } else {
- positiveValMarker
- }
- bbuf.put(signCol)
- bbuf.putShort(shortVal)
- writer.write(idx, bbuf.array())
+ val keySerializer = getAvroSerializer(keySchema)
+ val keyDeserializer = getAvroDeserializer(keySchema)
+
+ // Create the AvroEncoder with all components
+ AvroEncoder(
+ keySerializer,
+ keyDeserializer,
+ valueSerializer,
+ valueDeserializer,
+ suffixKeySchema.map(getAvroSerializer),
+ suffixKeySchema.map(getAvroDeserializer)
+ )
+ }
- case IntegerType =>
- val intVal = value.asInstanceOf[Int]
- val signCol = if (intVal < 0) {
- negativeValMarker
- } else {
- positiveValMarker
- }
- bbuf.put(signCol)
- bbuf.putInt(intVal)
- writer.write(idx, bbuf.array())
-
- case LongType =>
- val longVal = value.asInstanceOf[Long]
- val signCol = if (longVal < 0) {
- negativeValMarker
- } else {
- positiveValMarker
- }
- bbuf.put(signCol)
- bbuf.putLong(longVal)
- writer.write(idx, bbuf.array())
-
- case FloatType =>
- val floatVal = value.asInstanceOf[Float]
- val rawBits = floatToRawIntBits(floatVal)
- // perform sign comparison using bit manipulation to ensure NaN values are handled
- // correctly
- if ((rawBits & floatSignBitMask) != 0) {
- // for negative values, we need to flip all the bits to ensure correct ordering
- val updatedVal = rawBits ^ floatFlipBitMask
- bbuf.put(negativeValMarker)
- // convert the bits back to float
- bbuf.putFloat(intBitsToFloat(updatedVal))
- } else {
- bbuf.put(positiveValMarker)
- bbuf.putFloat(floatVal)
- }
- writer.write(idx, bbuf.array())
+ /**
+ * This method takes an UnsafeRow, and serializes to a byte array using Avro encoding.
+ */
+ def encodeUnsafeRowToAvro(
+ row: UnsafeRow,
+ avroSerializer: AvroSerializer,
+ valueAvroType: Schema,
+ out: ByteArrayOutputStream): Array[Byte] = {
+ // InternalRow -> Avro.GenericDataRecord
+ val avroData =
+ avroSerializer.serialize(row)
+ out.reset()
+ val encoder = EncoderFactory.get().directBinaryEncoder(out, null)
+ val writer = new GenericDatumWriter[Any](
+ valueAvroType) // Defining Avro writer for this struct type
+ writer.write(avroData, encoder) // Avro.GenericDataRecord -> byte array
+ encoder.flush()
+ out.toByteArray
+ }
- case DoubleType =>
- val doubleVal = value.asInstanceOf[Double]
- val rawBits = doubleToRawLongBits(doubleVal)
- // perform sign comparison using bit manipulation to ensure NaN values are handled
- // correctly
- if ((rawBits & doubleSignBitMask) != 0) {
- // for negative values, we need to flip all the bits to ensure correct ordering
- val updatedVal = rawBits ^ doubleFlipBitMask
- bbuf.put(negativeValMarker)
- // convert the bits back to double
- bbuf.putDouble(longBitsToDouble(updatedVal))
- } else {
- bbuf.put(positiveValMarker)
- bbuf.putDouble(doubleVal)
- }
- writer.write(idx, bbuf.array())
- }
- }
+ /**
+ * This method takes a byte array written using Avro encoding, and
+ * deserializes to an UnsafeRow using the Avro deserializer
+ */
+ def decodeFromAvroToUnsafeRow(
+ valueBytes: Array[Byte],
+ avroDeserializer: AvroDeserializer,
+ valueAvroType: Schema,
+ valueProj: UnsafeProjection): UnsafeRow = {
+ if (valueBytes != null) {
+ val reader = new GenericDatumReader[Any](valueAvroType)
+ val decoder = DecoderFactory.get().binaryDecoder(
+ valueBytes, 0, valueBytes.length, null)
+ // bytes -> Avro.GenericDataRecord
+ val genericData = reader.read(null, decoder)
+ // Avro.GenericDataRecord -> InternalRow
+ val internalRow = avroDeserializer.deserialize(
+ genericData).orNull.asInstanceOf[InternalRow]
+ // InternalRow -> UnsafeRow
+ valueProj.apply(internalRow)
+ } else {
+ null
}
- writer.getRow()
}
- // Rewrite the unsafe row by converting back from BIG_ENDIAN byte arrays to the
- // original data types.
- // For decode, we extract the byte array from the UnsafeRow, and then read the first byte
- // to determine if the value is null or not. If the value is null, we set the ordinal on
- // the UnsafeRow to null. If the value is not null, we read the rest of the bytes to get the
- // actual value.
- // For negative float/double values, we need to flip all the bits back to get the original value.
- private def decodePrefixKeyForRangeScan(row: UnsafeRow): UnsafeRow = {
- val writer = new UnsafeRowWriter(orderingOrdinals.length)
- writer.resetRowWriter()
- rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
- val field = fieldWithOrdinal._1
-
- val value = row.getBinary(idx)
- val bbuf = ByteBuffer.wrap(value.asInstanceOf[Array[Byte]])
- bbuf.order(ByteOrder.BIG_ENDIAN)
- val isNullOrSignCol = bbuf.get()
- if (isNullOrSignCol == nullValMarker) {
- // set the column to null and skip reading the next byte(s)
- writer.setNullAt(idx)
- } else {
- field.dataType match {
- case BooleanType =>
- case ByteType =>
- writer.write(idx, bbuf.get)
-
- case ShortType =>
- writer.write(idx, bbuf.getShort)
-
- case IntegerType =>
- writer.write(idx, bbuf.getInt)
-
- case LongType =>
- writer.write(idx, bbuf.getLong)
+ private val out = new ByteArrayOutputStream
- case FloatType =>
- if (isNullOrSignCol == negativeValMarker) {
- // if the number is negative, get the raw binary bits for the float
- // and flip the bits back
- val updatedVal = floatToRawIntBits(bbuf.getFloat) ^ floatFlipBitMask
- writer.write(idx, intBitsToFloat(updatedVal))
- } else {
- writer.write(idx, bbuf.getFloat)
- }
+ override def encodeKey(row: UnsafeRow): Array[Byte] = {
+ keyStateEncoderSpec match {
+ case NoPrefixKeyStateEncoderSpec(_) =>
+ encodeUnsafeRowToAvro(row, avroEncoder.keySerializer, keyAvroType, out)
+ case PrefixKeyScanStateEncoderSpec(_, _) =>
+ encodeUnsafeRowToAvro(row, avroEncoder.keySerializer, prefixKeyAvroType, out)
+ case _ => throw unsupportedOperationForKeyStateEncoder("encodeKey")
+ }
+ }
- case DoubleType =>
- if (isNullOrSignCol == negativeValMarker) {
- // if the number is negative, get the raw binary bits for the double
- // and flip the bits back
- val updatedVal = doubleToRawLongBits(bbuf.getDouble) ^ doubleFlipBitMask
- writer.write(idx, longBitsToDouble(updatedVal))
- } else {
- writer.write(idx, bbuf.getDouble)
- }
- }
- }
+ override def encodeRemainingKey(row: UnsafeRow): Array[Byte] = {
+ keyStateEncoderSpec match {
+ case PrefixKeyScanStateEncoderSpec(_, _) =>
+ encodeUnsafeRowToAvro(row, avroEncoder.suffixKeySerializer.get, remainingKeyAvroType, out)
+ case RangeKeyScanStateEncoderSpec(_, _) =>
+ encodeUnsafeRowToAvro(row, avroEncoder.keySerializer, remainingKeyAvroType, out)
Review Comment:
No, this is right - RangeKeyScan encoder doesn't use Avro serialization for both parts of the ke
--
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-50017][SQL] Support Avro encoding for TransformWithState operator [spark]
Posted by "brkyvz (via GitHub)" <gi...@apache.org>.
brkyvz commented on code in PR #48944:
URL: https://github.com/apache/spark/pull/48944#discussion_r1887870989
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateEncoder.scala:
##########
@@ -219,469 +222,402 @@ object RocksDBStateEncoder extends Logging {
}
}
-/**
- * RocksDB Key Encoder for UnsafeRow that supports prefix scan
- *
- * @param keySchema - schema of the key to be encoded
- * @param numColsPrefixKey - number of columns to be used for prefix key
- * @param useColumnFamilies - if column family is enabled for this encoder
- * @param avroEnc - if Avro encoding is specified for this StateEncoder, this encoder will
- * be defined
- */
-class PrefixKeyScanStateEncoder(
- keySchema: StructType,
- numColsPrefixKey: Int,
- useColumnFamilies: Boolean = false,
- virtualColFamilyId: Option[Short] = None,
- avroEnc: Option[AvroEncoder] = None)
- extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) {
-
- import RocksDBStateEncoder._
-
- private val usingAvroEncoding = avroEnc.isDefined
- private val prefixKeyFieldsWithIdx: Seq[(StructField, Int)] = {
- keySchema.zipWithIndex.take(numColsPrefixKey)
- }
+class UnsafeRowDataEncoder(
+ keyStateEncoderSpec: KeyStateEncoderSpec,
+ valueSchema: StructType) extends RocksDBDataEncoder(keyStateEncoderSpec, valueSchema) {
- private val remainingKeyFieldsWithIdx: Seq[(StructField, Int)] = {
- keySchema.zipWithIndex.drop(numColsPrefixKey)
+ override def encodeKey(row: UnsafeRow): Array[Byte] = {
+ encodeUnsafeRow(row)
}
- private val prefixKeyProjection: UnsafeProjection = {
- val refs = prefixKeyFieldsWithIdx.map(x => BoundReference(x._2, x._1.dataType, x._1.nullable))
- UnsafeProjection.create(refs)
+ override def encodeRemainingKey(row: UnsafeRow): Array[Byte] = {
+ encodeUnsafeRow(row)
}
- private val remainingKeyProjection: UnsafeProjection = {
- val refs = remainingKeyFieldsWithIdx.map(x =>
- BoundReference(x._2, x._1.dataType, x._1.nullable))
- UnsafeProjection.create(refs)
- }
+ override def encodePrefixKeyForRangeScan(row: UnsafeRow): Array[Byte] = {
+ assert(keyStateEncoderSpec.isInstanceOf[RangeKeyScanStateEncoderSpec])
+ val rsk = keyStateEncoderSpec.asInstanceOf[RangeKeyScanStateEncoderSpec]
+ val rangeScanKeyFieldsWithOrdinal = rsk.orderingOrdinals.map { ordinal =>
+ val field = rsk.keySchema(ordinal)
+ (field, ordinal)
+ }
+ val writer = new UnsafeRowWriter(rsk.orderingOrdinals.length)
+ writer.resetRowWriter()
+ rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
+ val field = fieldWithOrdinal._1
+ val value = row.get(idx, field.dataType)
+ // Note that we cannot allocate a smaller buffer here even if the value is null
+ // because the effective byte array is considered variable size and needs to have
+ // the same size across all rows for the ordering to work as expected.
+ val bbuf = ByteBuffer.allocate(field.dataType.defaultSize + 1)
+ bbuf.order(ByteOrder.BIG_ENDIAN)
+ if (value == null) {
+ bbuf.put(nullValMarker)
+ writer.write(idx, bbuf.array())
+ } else {
+ field.dataType match {
+ case BooleanType =>
+ case ByteType =>
+ val byteVal = value.asInstanceOf[Byte]
+ val signCol = if (byteVal < 0) {
+ negativeValMarker
+ } else {
+ positiveValMarker
+ }
+ bbuf.put(signCol)
+ bbuf.put(byteVal)
+ writer.write(idx, bbuf.array())
- // Prefix Key schema and projection definitions used by the Avro Serializers
- // and Deserializers
- private val prefixKeySchema = StructType(keySchema.take(numColsPrefixKey))
- private lazy val prefixKeyAvroType = SchemaConverters.toAvroType(prefixKeySchema)
- private val prefixKeyProj = UnsafeProjection.create(prefixKeySchema)
+ case ShortType =>
+ val shortVal = value.asInstanceOf[Short]
+ val signCol = if (shortVal < 0) {
+ negativeValMarker
+ } else {
+ positiveValMarker
+ }
+ bbuf.put(signCol)
+ bbuf.putShort(shortVal)
+ writer.write(idx, bbuf.array())
- // Remaining Key schema and projection definitions used by the Avro Serializers
- // and Deserializers
- private val remainingKeySchema = StructType(keySchema.drop(numColsPrefixKey))
- private lazy val remainingKeyAvroType = SchemaConverters.toAvroType(remainingKeySchema)
- private val remainingKeyProj = UnsafeProjection.create(remainingKeySchema)
+ case IntegerType =>
+ val intVal = value.asInstanceOf[Int]
+ val signCol = if (intVal < 0) {
+ negativeValMarker
+ } else {
+ positiveValMarker
+ }
+ bbuf.put(signCol)
+ bbuf.putInt(intVal)
+ writer.write(idx, bbuf.array())
- // This is quite simple to do - just bind sequentially, as we don't change the order.
- private val restoreKeyProjection: UnsafeProjection = UnsafeProjection.create(keySchema)
+ case LongType =>
+ val longVal = value.asInstanceOf[Long]
+ val signCol = if (longVal < 0) {
+ negativeValMarker
+ } else {
+ positiveValMarker
+ }
+ bbuf.put(signCol)
+ bbuf.putLong(longVal)
+ writer.write(idx, bbuf.array())
- // Reusable objects
- private val joinedRowOnKey = new JoinedRow()
+ case FloatType =>
+ val floatVal = value.asInstanceOf[Float]
+ val rawBits = floatToRawIntBits(floatVal)
+ // perform sign comparison using bit manipulation to ensure NaN values are handled
+ // correctly
+ if ((rawBits & floatSignBitMask) != 0) {
+ // for negative values, we need to flip all the bits to ensure correct ordering
+ val updatedVal = rawBits ^ floatFlipBitMask
+ bbuf.put(negativeValMarker)
+ // convert the bits back to float
+ bbuf.putFloat(intBitsToFloat(updatedVal))
+ } else {
+ bbuf.put(positiveValMarker)
+ bbuf.putFloat(floatVal)
+ }
+ writer.write(idx, bbuf.array())
- override def encodeKey(row: UnsafeRow): Array[Byte] = {
- val (prefixKeyEncoded, remainingEncoded) = if (usingAvroEncoding) {
- (
- encodeUnsafeRowToAvro(
- extractPrefixKey(row),
- avroEnc.get.keySerializer,
- prefixKeyAvroType,
- out
- ),
- encodeUnsafeRowToAvro(
- remainingKeyProjection(row),
- avroEnc.get.suffixKeySerializer.get,
- remainingKeyAvroType,
- out
- )
- )
- } else {
- (encodeUnsafeRow(extractPrefixKey(row)), encodeUnsafeRow(remainingKeyProjection(row)))
+ case DoubleType =>
+ val doubleVal = value.asInstanceOf[Double]
+ val rawBits = doubleToRawLongBits(doubleVal)
+ // perform sign comparison using bit manipulation to ensure NaN values are handled
+ // correctly
+ if ((rawBits & doubleSignBitMask) != 0) {
+ // for negative values, we need to flip all the bits to ensure correct ordering
+ val updatedVal = rawBits ^ doubleFlipBitMask
+ bbuf.put(negativeValMarker)
+ // convert the bits back to double
+ bbuf.putDouble(longBitsToDouble(updatedVal))
+ } else {
+ bbuf.put(positiveValMarker)
+ bbuf.putDouble(doubleVal)
+ }
+ writer.write(idx, bbuf.array())
+ }
+ }
}
- val (encodedBytes, startingOffset) = encodeColumnFamilyPrefix(
- prefixKeyEncoded.length + remainingEncoded.length + 4
- )
+ encodeUnsafeRow(writer.getRow())
+ }
- Platform.putInt(encodedBytes, startingOffset, prefixKeyEncoded.length)
- Platform.copyMemory(prefixKeyEncoded, Platform.BYTE_ARRAY_OFFSET,
- encodedBytes, startingOffset + 4, prefixKeyEncoded.length)
- // NOTE: We don't put the length of remainingEncoded as we can calculate later
- // on deserialization.
- Platform.copyMemory(remainingEncoded, Platform.BYTE_ARRAY_OFFSET,
- encodedBytes, startingOffset + 4 + prefixKeyEncoded.length,
- remainingEncoded.length)
-
- encodedBytes
- }
-
- override def decodeKey(keyBytes: Array[Byte]): UnsafeRow = {
- val prefixKeyEncodedLen = Platform.getInt(keyBytes, decodeKeyStartOffset)
- val prefixKeyEncoded = new Array[Byte](prefixKeyEncodedLen)
- Platform.copyMemory(keyBytes, decodeKeyStartOffset + 4,
- prefixKeyEncoded, Platform.BYTE_ARRAY_OFFSET, prefixKeyEncodedLen)
+ override def encodeValue(row: UnsafeRow): Array[Byte] = encodeUnsafeRow(row)
- // Here we calculate the remainingKeyEncodedLen leveraging the length of keyBytes
- val remainingKeyEncodedLen = keyBytes.length - 4 - prefixKeyEncodedLen -
- offsetForColFamilyPrefix
-
- val remainingKeyEncoded = new Array[Byte](remainingKeyEncodedLen)
- Platform.copyMemory(keyBytes, decodeKeyStartOffset + 4 + prefixKeyEncodedLen,
- remainingKeyEncoded, Platform.BYTE_ARRAY_OFFSET, remainingKeyEncodedLen)
-
- val (prefixKeyDecoded, remainingKeyDecoded) = if (usingAvroEncoding) {
- (
- decodeFromAvroToUnsafeRow(
- prefixKeyEncoded,
- avroEnc.get.keyDeserializer,
- prefixKeyAvroType,
- prefixKeyProj
- ),
- decodeFromAvroToUnsafeRow(
- remainingKeyEncoded,
- avroEnc.get.suffixKeyDeserializer.get,
- remainingKeyAvroType,
- remainingKeyProj
- )
- )
- } else {
- (decodeToUnsafeRow(prefixKeyEncoded, numFields = numColsPrefixKey),
- decodeToUnsafeRow(remainingKeyEncoded, numFields = keySchema.length - numColsPrefixKey))
+ override def decodeKey(bytes: Array[Byte]): UnsafeRow = {
+ keyStateEncoderSpec match {
+ case NoPrefixKeyStateEncoderSpec(_) =>
+ decodeToUnsafeRow(bytes, reusedKeyRow)
+ case PrefixKeyScanStateEncoderSpec(_, numColsPrefixKey) =>
+ decodeToUnsafeRow(bytes, numFields = numColsPrefixKey)
+ case _ => throw unsupportedOperationForKeyStateEncoder("decodeKey")
}
-
- restoreKeyProjection(joinedRowOnKey.withLeft(prefixKeyDecoded).withRight(remainingKeyDecoded))
}
- private def extractPrefixKey(key: UnsafeRow): UnsafeRow = {
- prefixKeyProjection(key)
+ override def decodeRemainingKey(bytes: Array[Byte]): UnsafeRow = {
+ keyStateEncoderSpec match {
+ case PrefixKeyScanStateEncoderSpec(_, numColsPrefixKey) =>
+ decodeToUnsafeRow(bytes, numFields = numColsPrefixKey)
+ case RangeKeyScanStateEncoderSpec(_, orderingOrdinals) =>
+ decodeToUnsafeRow(bytes, keySchema.length - orderingOrdinals.length)
+ case _ => throw unsupportedOperationForKeyStateEncoder("decodeRemainingKey")
+ }
}
- override def encodePrefixKey(prefixKey: UnsafeRow): Array[Byte] = {
- val prefixKeyEncoded = if (usingAvroEncoding) {
- encodeUnsafeRowToAvro(prefixKey, avroEnc.get.keySerializer, prefixKeyAvroType, out)
- } else {
- encodeUnsafeRow(prefixKey)
+ override def decodePrefixKeyForRangeScan(bytes: Array[Byte]): UnsafeRow = {
+ assert(keyStateEncoderSpec.isInstanceOf[RangeKeyScanStateEncoderSpec])
+ val rsk = keyStateEncoderSpec.asInstanceOf[RangeKeyScanStateEncoderSpec]
+ val writer = new UnsafeRowWriter(rsk.orderingOrdinals.length)
+ val rangeScanKeyFieldsWithOrdinal = rsk.orderingOrdinals.map { ordinal =>
+ val field = rsk.keySchema(ordinal)
+ (field, ordinal)
}
- val (prefix, startingOffset) = encodeColumnFamilyPrefix(
- prefixKeyEncoded.length + 4
- )
-
- Platform.putInt(prefix, startingOffset, prefixKeyEncoded.length)
- Platform.copyMemory(prefixKeyEncoded, Platform.BYTE_ARRAY_OFFSET, prefix,
- startingOffset + 4, prefixKeyEncoded.length)
- prefix
- }
+ writer.resetRowWriter()
+ val row = decodeToUnsafeRow(bytes, numFields = rsk.orderingOrdinals.length)
+ rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
+ val field = fieldWithOrdinal._1
- override def supportPrefixKeyScan: Boolean = true
-}
+ val value = row.getBinary(idx)
+ val bbuf = ByteBuffer.wrap(value.asInstanceOf[Array[Byte]])
+ bbuf.order(ByteOrder.BIG_ENDIAN)
+ val isNullOrSignCol = bbuf.get()
+ if (isNullOrSignCol == nullValMarker) {
+ // set the column to null and skip reading the next byte(s)
+ writer.setNullAt(idx)
+ } else {
+ field.dataType match {
+ case BooleanType =>
+ case ByteType =>
+ writer.write(idx, bbuf.get)
-/**
- * RocksDB Key Encoder for UnsafeRow that supports range scan for fixed size fields
- *
- * To encode a row for range scan, we first project the orderingOrdinals from the oridinal
- * UnsafeRow into another UnsafeRow; we then rewrite that new UnsafeRow's fields in BIG_ENDIAN
- * to allow for scanning keys in sorted order using the byte-wise comparison method that
- * RocksDB uses.
- *
- * Then, for the rest of the fields, we project those into another UnsafeRow.
- * We then effectively join these two UnsafeRows together, and finally take those bytes
- * to get the resulting row.
- *
- * We cannot support variable sized fields in the range scan because the UnsafeRow format
- * stores variable sized fields as offset and length pointers to the actual values,
- * thereby changing the required ordering.
- *
- * Note that we also support "null" values being passed for these fixed size fields. We prepend
- * a single byte to indicate whether the column value is null or not. We cannot change the
- * nullability on the UnsafeRow itself as the expected ordering would change if non-first
- * columns are marked as null. If the first col is null, those entries will appear last in
- * the iterator. If non-first columns are null, ordering based on the previous columns will
- * still be honored. For rows with null column values, ordering for subsequent columns
- * will also be maintained within those set of rows. We use the same byte to also encode whether
- * the value is negative or not. For negative float/double values, we flip all the bits to ensure
- * the right lexicographical ordering. For the rationale around this, please check the link
- * here: https://en.wikipedia.org/wiki/IEEE_754#Design_rationale
- *
- * @param keySchema - schema of the key to be encoded
- * @param orderingOrdinals - the ordinals for which the range scan is constructed
- * @param useColumnFamilies - if column family is enabled for this encoder
- * @param avroEnc - if Avro encoding is specified for this StateEncoder, this encoder will
- * be defined
- */
-class RangeKeyScanStateEncoder(
- keySchema: StructType,
- orderingOrdinals: Seq[Int],
- useColumnFamilies: Boolean = false,
- virtualColFamilyId: Option[Short] = None,
- avroEnc: Option[AvroEncoder] = None)
- extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) with Logging {
+ case ShortType =>
+ writer.write(idx, bbuf.getShort)
- import RocksDBStateEncoder._
+ case IntegerType =>
+ writer.write(idx, bbuf.getInt)
- private val rangeScanKeyFieldsWithOrdinal: Seq[(StructField, Int)] = {
- orderingOrdinals.map { ordinal =>
- val field = keySchema(ordinal)
- (field, ordinal)
- }
- }
+ case LongType =>
+ writer.write(idx, bbuf.getLong)
- private def isFixedSize(dataType: DataType): Boolean = dataType match {
- case _: ByteType | _: BooleanType | _: ShortType | _: IntegerType | _: LongType |
- _: FloatType | _: DoubleType => true
- case _ => false
- }
+ case FloatType =>
+ if (isNullOrSignCol == negativeValMarker) {
+ // if the number is negative, get the raw binary bits for the float
+ // and flip the bits back
+ val updatedVal = floatToRawIntBits(bbuf.getFloat) ^ floatFlipBitMask
+ writer.write(idx, intBitsToFloat(updatedVal))
+ } else {
+ writer.write(idx, bbuf.getFloat)
+ }
- // verify that only fixed sized columns are used for ordering
- rangeScanKeyFieldsWithOrdinal.foreach { case (field, ordinal) =>
- if (!isFixedSize(field.dataType)) {
- // NullType is technically fixed size, but not supported for ordering
- if (field.dataType == NullType) {
- throw StateStoreErrors.nullTypeOrderingColsNotSupported(field.name, ordinal.toString)
- } else {
- throw StateStoreErrors.variableSizeOrderingColsNotSupported(field.name, ordinal.toString)
+ case DoubleType =>
+ if (isNullOrSignCol == negativeValMarker) {
+ // if the number is negative, get the raw binary bits for the double
+ // and flip the bits back
+ val updatedVal = doubleToRawLongBits(bbuf.getDouble) ^ doubleFlipBitMask
+ writer.write(idx, longBitsToDouble(updatedVal))
+ } else {
+ writer.write(idx, bbuf.getDouble)
+ }
+ }
}
}
+ writer.getRow()
}
- private val remainingKeyFieldsWithOrdinal: Seq[(StructField, Int)] = {
- 0.to(keySchema.length - 1).diff(orderingOrdinals).map { ordinal =>
- val field = keySchema(ordinal)
- (field, ordinal)
- }
- }
+ override def decodeValue(bytes: Array[Byte]): UnsafeRow = decodeToUnsafeRow(bytes, reusedValueRow)
+}
- private val rangeScanKeyProjection: UnsafeProjection = {
- val refs = rangeScanKeyFieldsWithOrdinal.map(x =>
- BoundReference(x._2, x._1.dataType, x._1.nullable))
- UnsafeProjection.create(refs)
- }
+class AvroStateEncoder(
+ keyStateEncoderSpec: KeyStateEncoderSpec,
+ valueSchema: StructType) extends RocksDBDataEncoder(keyStateEncoderSpec, valueSchema)
+ with Logging {
- private val remainingKeyProjection: UnsafeProjection = {
- val refs = remainingKeyFieldsWithOrdinal.map(x =>
- BoundReference(x._2, x._1.dataType, x._1.nullable))
- UnsafeProjection.create(refs)
- }
+ private val avroEncoder = createAvroEnc(keyStateEncoderSpec, valueSchema)
+ // Avro schema used by the avro encoders
+ private lazy val keyAvroType: Schema = SchemaConverters.toAvroType(keySchema)
+ private lazy val keyProj = UnsafeProjection.create(keySchema)
- // The original schema that we might get could be:
- // [foo, bar, baz, buzz]
- // We might order by bar and buzz, leading to:
- // [bar, buzz, foo, baz]
- // We need to create a projection that sends, for example, the buzz at index 1 to index
- // 3. Thus, for every record in the original schema, we compute where it would be in
- // the joined row and created a projection based on that.
- private val restoreKeyProjection: UnsafeProjection = {
- val refs = keySchema.zipWithIndex.map { case (field, originalOrdinal) =>
- val ordinalInJoinedRow = if (orderingOrdinals.contains(originalOrdinal)) {
- orderingOrdinals.indexOf(originalOrdinal)
- } else {
- orderingOrdinals.length +
- remainingKeyFieldsWithOrdinal.indexWhere(_._2 == originalOrdinal)
- }
+ private lazy val valueAvroType: Schema = SchemaConverters.toAvroType(valueSchema)
+ private lazy val valueProj = UnsafeProjection.create(valueSchema)
- BoundReference(ordinalInJoinedRow, field.dataType, field.nullable)
- }
- UnsafeProjection.create(refs)
+ // Prefix Key schema and projection definitions used by the Avro Serializers
+ // and Deserializers
+ private lazy val prefixKeySchema = keyStateEncoderSpec match {
+ case PrefixKeyScanStateEncoderSpec(keySchema, numColsPrefixKey) =>
+ StructType(keySchema.take (numColsPrefixKey))
+ case _ => throw unsupportedOperationForKeyStateEncoder("prefixKeySchema")
+ }
+ private lazy val prefixKeyAvroType = SchemaConverters.toAvroType(prefixKeySchema)
+ private lazy val prefixKeyProj = UnsafeProjection.create(prefixKeySchema)
+
+ // Range Key schema nd projection definitions used by the Avro Serializers and
+ // Deserializers
+ private lazy val rangeScanKeyFieldsWithOrdinal = keyStateEncoderSpec match {
+ case RangeKeyScanStateEncoderSpec(keySchema, orderingOrdinals) =>
+ orderingOrdinals.map { ordinal =>
+ val field = keySchema(ordinal)
+ (field, ordinal)
+ }
+ case _ =>
+ throw unsupportedOperationForKeyStateEncoder("rangeScanKey")
}
- private val rangeScanAvroSchema = StateStoreColumnFamilySchemaUtils.convertForRangeScan(
+ private lazy val rangeScanAvroSchema = StateStoreColumnFamilySchemaUtils.convertForRangeScan(
StructType(rangeScanKeyFieldsWithOrdinal.map(_._1).toArray))
private lazy val rangeScanAvroType = SchemaConverters.toAvroType(rangeScanAvroSchema)
- private val rangeScanAvroProjection = UnsafeProjection.create(rangeScanAvroSchema)
+ private lazy val rangeScanAvroProjection = UnsafeProjection.create(rangeScanAvroSchema)
- // Existing remainder key schema stuff
- private val remainingKeySchema = StructType(
- 0.to(keySchema.length - 1).diff(orderingOrdinals).map(keySchema(_))
- )
+ // Existing remainder key schema definitions
+ // Remaining Key schema and projection definitions used by the Avro Serializers
+ // and Deserializers
+ private lazy val remainingKeySchema = keyStateEncoderSpec match {
+ case PrefixKeyScanStateEncoderSpec(keySchema, numColsPrefixKey) =>
+ StructType(keySchema.drop(numColsPrefixKey))
+ case RangeKeyScanStateEncoderSpec(keySchema, orderingOrdinals) =>
+ StructType(0.until(keySchema.length).diff(orderingOrdinals).map(keySchema(_)))
+ case _ => throw unsupportedOperationForKeyStateEncoder("remainingKeySchema")
+ }
private lazy val remainingKeyAvroType = SchemaConverters.toAvroType(remainingKeySchema)
- private val remainingKeyAvroProjection = UnsafeProjection.create(remainingKeySchema)
+ private lazy val remainingKeyAvroProjection = UnsafeProjection.create(remainingKeySchema)
- // Reusable objects
- private val joinedRowOnKey = new JoinedRow()
-
- private def extractPrefixKey(key: UnsafeRow): UnsafeRow = {
- rangeScanKeyProjection(key)
+ private def getAvroSerializer(schema: StructType): AvroSerializer = {
+ val avroType = SchemaConverters.toAvroType(schema)
+ new AvroSerializer(schema, avroType, nullable = false)
}
- // bit masks used for checking sign or flipping all bits for negative float/double values
- private val floatFlipBitMask = 0xFFFFFFFF
- private val floatSignBitMask = 0x80000000
+ private def getAvroDeserializer(schema: StructType): AvroDeserializer = {
+ val avroType = SchemaConverters.toAvroType(schema)
+ val avroOptions = AvroOptions(Map.empty)
+ new AvroDeserializer(avroType, schema,
+ avroOptions.datetimeRebaseModeInRead, avroOptions.useStableIdForUnionType,
+ avroOptions.stableIdPrefixForUnionType, avroOptions.recursiveFieldMaxDepth)
+ }
- private val doubleFlipBitMask = 0xFFFFFFFFFFFFFFFFL
- private val doubleSignBitMask = 0x8000000000000000L
+ /**
+ * Creates an AvroEncoder that handles both key and value serialization/deserialization.
+ * This method sets up the complete encoding infrastructure needed for state store operations.
+ *
+ * The encoder handles different key encoding specifications:
+ * - NoPrefixKeyStateEncoderSpec: Simple key encoding without prefix
+ * - PrefixKeyScanStateEncoderSpec: Keys with prefix for efficient scanning
+ * - RangeKeyScanStateEncoderSpec: Keys with ordering requirements for range scans
+ *
+ * For prefix scan cases, it also creates separate encoders for the suffix portion of keys.
+ *
+ * @param keyStateEncoderSpec Specification for how to encode keys
+ * @param valueSchema Schema for the values to be encoded
+ * @return An AvroEncoder containing all necessary serializers and deserializers
+ */
+ private def createAvroEnc(
+ keyStateEncoderSpec: KeyStateEncoderSpec,
+ valueSchema: StructType): AvroEncoder = {
+ val valueSerializer = getAvroSerializer(valueSchema)
+ val valueDeserializer = getAvroDeserializer(valueSchema)
+
+ // Get key schema based on encoder spec type
+ val keySchema = keyStateEncoderSpec match {
+ case NoPrefixKeyStateEncoderSpec(schema) =>
+ schema
+ case PrefixKeyScanStateEncoderSpec(schema, numColsPrefixKey) =>
+ StructType(schema.take(numColsPrefixKey))
+ case RangeKeyScanStateEncoderSpec(schema, orderingOrdinals) =>
+ val remainingSchema = {
+ 0.until(schema.length).diff(orderingOrdinals).map { ordinal =>
+ schema(ordinal)
+ }
+ }
+ StructType(remainingSchema)
+ }
- // Byte markers used to identify whether the value is null, negative or positive
- // To ensure sorted ordering, we use the lowest byte value for negative numbers followed by
- // positive numbers and then null values.
- private val negativeValMarker: Byte = 0x00.toByte
- private val positiveValMarker: Byte = 0x01.toByte
- private val nullValMarker: Byte = 0x02.toByte
-
- // Rewrite the unsafe row by replacing fixed size fields with BIG_ENDIAN encoding
- // using byte arrays.
- // To handle "null" values, we prepend a byte to the byte array indicating whether the value
- // is null or not. If the value is null, we write the null byte followed by zero bytes.
- // If the value is not null, we write the null byte followed by the value.
- // Note that setting null for the index on the unsafeRow is not feasible as it would change
- // the sorting order on iteration.
- // Also note that the same byte is used to indicate whether the value is negative or not.
- private def encodePrefixKeyForRangeScan(row: UnsafeRow): UnsafeRow = {
- val writer = new UnsafeRowWriter(orderingOrdinals.length)
- writer.resetRowWriter()
- rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
- val field = fieldWithOrdinal._1
- val value = row.get(idx, field.dataType)
- // Note that we cannot allocate a smaller buffer here even if the value is null
- // because the effective byte array is considered variable size and needs to have
- // the same size across all rows for the ordering to work as expected.
- val bbuf = ByteBuffer.allocate(field.dataType.defaultSize + 1)
- bbuf.order(ByteOrder.BIG_ENDIAN)
- if (value == null) {
- bbuf.put(nullValMarker)
- writer.write(idx, bbuf.array())
- } else {
- field.dataType match {
- case BooleanType =>
- case ByteType =>
- val byteVal = value.asInstanceOf[Byte]
- val signCol = if (byteVal < 0) {
- negativeValMarker
- } else {
- positiveValMarker
- }
- bbuf.put(signCol)
- bbuf.put(byteVal)
- writer.write(idx, bbuf.array())
+ // Handle suffix key schema for prefix scan case
+ val suffixKeySchema = keyStateEncoderSpec match {
+ case PrefixKeyScanStateEncoderSpec(schema, numColsPrefixKey) =>
+ Some(StructType(schema.drop(numColsPrefixKey)))
+ case _ =>
+ None
+ }
- case ShortType =>
- val shortVal = value.asInstanceOf[Short]
- val signCol = if (shortVal < 0) {
- negativeValMarker
- } else {
- positiveValMarker
- }
- bbuf.put(signCol)
- bbuf.putShort(shortVal)
- writer.write(idx, bbuf.array())
+ val keySerializer = getAvroSerializer(keySchema)
+ val keyDeserializer = getAvroDeserializer(keySchema)
+
+ // Create the AvroEncoder with all components
+ AvroEncoder(
+ keySerializer,
+ keyDeserializer,
+ valueSerializer,
+ valueDeserializer,
+ suffixKeySchema.map(getAvroSerializer),
+ suffixKeySchema.map(getAvroDeserializer)
+ )
+ }
- case IntegerType =>
- val intVal = value.asInstanceOf[Int]
- val signCol = if (intVal < 0) {
- negativeValMarker
- } else {
- positiveValMarker
- }
- bbuf.put(signCol)
- bbuf.putInt(intVal)
- writer.write(idx, bbuf.array())
-
- case LongType =>
- val longVal = value.asInstanceOf[Long]
- val signCol = if (longVal < 0) {
- negativeValMarker
- } else {
- positiveValMarker
- }
- bbuf.put(signCol)
- bbuf.putLong(longVal)
- writer.write(idx, bbuf.array())
-
- case FloatType =>
- val floatVal = value.asInstanceOf[Float]
- val rawBits = floatToRawIntBits(floatVal)
- // perform sign comparison using bit manipulation to ensure NaN values are handled
- // correctly
- if ((rawBits & floatSignBitMask) != 0) {
- // for negative values, we need to flip all the bits to ensure correct ordering
- val updatedVal = rawBits ^ floatFlipBitMask
- bbuf.put(negativeValMarker)
- // convert the bits back to float
- bbuf.putFloat(intBitsToFloat(updatedVal))
- } else {
- bbuf.put(positiveValMarker)
- bbuf.putFloat(floatVal)
- }
- writer.write(idx, bbuf.array())
+ /**
+ * This method takes an UnsafeRow, and serializes to a byte array using Avro encoding.
+ */
+ def encodeUnsafeRowToAvro(
+ row: UnsafeRow,
+ avroSerializer: AvroSerializer,
+ valueAvroType: Schema,
+ out: ByteArrayOutputStream): Array[Byte] = {
+ // InternalRow -> Avro.GenericDataRecord
+ val avroData =
+ avroSerializer.serialize(row)
+ out.reset()
+ val encoder = EncoderFactory.get().directBinaryEncoder(out, null)
+ val writer = new GenericDatumWriter[Any](
+ valueAvroType) // Defining Avro writer for this struct type
+ writer.write(avroData, encoder) // Avro.GenericDataRecord -> byte array
+ encoder.flush()
+ out.toByteArray
+ }
- case DoubleType =>
- val doubleVal = value.asInstanceOf[Double]
- val rawBits = doubleToRawLongBits(doubleVal)
- // perform sign comparison using bit manipulation to ensure NaN values are handled
- // correctly
- if ((rawBits & doubleSignBitMask) != 0) {
- // for negative values, we need to flip all the bits to ensure correct ordering
- val updatedVal = rawBits ^ doubleFlipBitMask
- bbuf.put(negativeValMarker)
- // convert the bits back to double
- bbuf.putDouble(longBitsToDouble(updatedVal))
- } else {
- bbuf.put(positiveValMarker)
- bbuf.putDouble(doubleVal)
- }
- writer.write(idx, bbuf.array())
- }
- }
+ /**
+ * This method takes a byte array written using Avro encoding, and
+ * deserializes to an UnsafeRow using the Avro deserializer
+ */
+ def decodeFromAvroToUnsafeRow(
+ valueBytes: Array[Byte],
+ avroDeserializer: AvroDeserializer,
+ valueAvroType: Schema,
+ valueProj: UnsafeProjection): UnsafeRow = {
+ if (valueBytes != null) {
+ val reader = new GenericDatumReader[Any](valueAvroType)
+ val decoder = DecoderFactory.get().binaryDecoder(
+ valueBytes, 0, valueBytes.length, null)
+ // bytes -> Avro.GenericDataRecord
+ val genericData = reader.read(null, decoder)
+ // Avro.GenericDataRecord -> InternalRow
+ val internalRow = avroDeserializer.deserialize(
+ genericData).orNull.asInstanceOf[InternalRow]
+ // InternalRow -> UnsafeRow
+ valueProj.apply(internalRow)
+ } else {
+ null
}
- writer.getRow()
}
- // Rewrite the unsafe row by converting back from BIG_ENDIAN byte arrays to the
- // original data types.
- // For decode, we extract the byte array from the UnsafeRow, and then read the first byte
- // to determine if the value is null or not. If the value is null, we set the ordinal on
- // the UnsafeRow to null. If the value is not null, we read the rest of the bytes to get the
- // actual value.
- // For negative float/double values, we need to flip all the bits back to get the original value.
- private def decodePrefixKeyForRangeScan(row: UnsafeRow): UnsafeRow = {
- val writer = new UnsafeRowWriter(orderingOrdinals.length)
- writer.resetRowWriter()
- rangeScanKeyFieldsWithOrdinal.zipWithIndex.foreach { case (fieldWithOrdinal, idx) =>
- val field = fieldWithOrdinal._1
-
- val value = row.getBinary(idx)
- val bbuf = ByteBuffer.wrap(value.asInstanceOf[Array[Byte]])
- bbuf.order(ByteOrder.BIG_ENDIAN)
- val isNullOrSignCol = bbuf.get()
- if (isNullOrSignCol == nullValMarker) {
- // set the column to null and skip reading the next byte(s)
- writer.setNullAt(idx)
- } else {
- field.dataType match {
- case BooleanType =>
- case ByteType =>
- writer.write(idx, bbuf.get)
-
- case ShortType =>
- writer.write(idx, bbuf.getShort)
-
- case IntegerType =>
- writer.write(idx, bbuf.getInt)
-
- case LongType =>
- writer.write(idx, bbuf.getLong)
+ private val out = new ByteArrayOutputStream
- case FloatType =>
- if (isNullOrSignCol == negativeValMarker) {
- // if the number is negative, get the raw binary bits for the float
- // and flip the bits back
- val updatedVal = floatToRawIntBits(bbuf.getFloat) ^ floatFlipBitMask
- writer.write(idx, intBitsToFloat(updatedVal))
- } else {
- writer.write(idx, bbuf.getFloat)
- }
+ override def encodeKey(row: UnsafeRow): Array[Byte] = {
+ keyStateEncoderSpec match {
+ case NoPrefixKeyStateEncoderSpec(_) =>
+ encodeUnsafeRowToAvro(row, avroEncoder.keySerializer, keyAvroType, out)
+ case PrefixKeyScanStateEncoderSpec(_, _) =>
+ encodeUnsafeRowToAvro(row, avroEncoder.keySerializer, prefixKeyAvroType, out)
+ case _ => throw unsupportedOperationForKeyStateEncoder("encodeKey")
+ }
+ }
- case DoubleType =>
- if (isNullOrSignCol == negativeValMarker) {
- // if the number is negative, get the raw binary bits for the double
- // and flip the bits back
- val updatedVal = doubleToRawLongBits(bbuf.getDouble) ^ doubleFlipBitMask
- writer.write(idx, longBitsToDouble(updatedVal))
- } else {
- writer.write(idx, bbuf.getDouble)
- }
- }
- }
+ override def encodeRemainingKey(row: UnsafeRow): Array[Byte] = {
+ keyStateEncoderSpec match {
+ case PrefixKeyScanStateEncoderSpec(_, _) =>
+ encodeUnsafeRowToAvro(row, avroEncoder.suffixKeySerializer.get, remainingKeyAvroType, out)
+ case RangeKeyScanStateEncoderSpec(_, _) =>
+ encodeUnsafeRowToAvro(row, avroEncoder.keySerializer, remainingKeyAvroType, out)
Review Comment:
does this need to be suffixKeySerializer?
--
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