anishshri-db commented on code in PR #45038: URL: https://github.com/apache/spark/pull/45038#discussion_r1487240502
########## sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImpl.scala: ########## @@ -52,14 +58,12 @@ class ValueStateImpl[S]( val keyByteArr = toRow .apply(keyOption.get).asInstanceOf[UnsafeRow].getBytes() - val schemaForKeyRow: StructType = new StructType().add("key", BinaryType) val keyEncoder = UnsafeProjection.create(schemaForKeyRow) Review Comment: Done ########## sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala: ########## @@ -92,10 +108,17 @@ private[sql] class RocksDBStateStoreProvider override def prefixScan(prefixKey: UnsafeRow, colFamilyName: String): Iterator[UnsafeRowPair] = { - require(encoder.supportPrefixKeyScan, "Prefix scan requires setting prefix key!") - - val prefix = encoder.encodePrefixKey(prefixKey) - rocksDB.prefixScan(prefix, colFamilyName).map(kv => encoder.decode(kv)) + val kvEncoder = keyValueEncoderMap.get(colFamilyName) + require(kvEncoder._1.supportPrefixKeyScan, + "Prefix scan requires setting prefix key!") + + val prefix = kvEncoder._1.encodePrefixKey(prefixKey) + rocksDB.prefixScan(prefix, colFamilyName).map { kv => + val rowPair = new UnsafeRowPair() Review Comment: Done ########## sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala: ########## @@ -69,19 +79,25 @@ private[sql] class RocksDBStateStoreProvider verify(state == UPDATING, "Cannot put after already committed or aborted") verify(key != null, "Key cannot be null") require(value != null, "Cannot put a null value") - rocksDB.put(encoder.encodeKey(key), encoder.encodeValue(value), colFamilyName) + val kvEncoder = keyValueEncoderMap.get(colFamilyName) + rocksDB.put(kvEncoder._1.encodeKey(key), + kvEncoder._2.encodeValue(value), colFamilyName) } override def remove(key: UnsafeRow, colFamilyName: String): Unit = { verify(state == UPDATING, "Cannot remove after already committed or aborted") verify(key != null, "Key cannot be null") - rocksDB.remove(encoder.encodeKey(key), colFamilyName) + val kvEncoder = keyValueEncoderMap.get(colFamilyName) + rocksDB.remove(kvEncoder._1.encodeKey(key), colFamilyName) } override def iterator(colFamilyName: String): Iterator[UnsafeRowPair] = { + val kvEncoder = keyValueEncoderMap.get(colFamilyName) rocksDB.iterator(colFamilyName).map { kv => - val rowPair = encoder.decode(kv) - if (!isValidated && rowPair.value != null) { + val rowPair = new UnsafeRowPair() Review Comment: Done -- 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-unsubscr...@spark.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org