neilramaswamy commented on code in PR #45051: URL: https://github.com/apache/spark/pull/45051#discussion_r1518669207
########## sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TimerStateImpl.scala: ########## @@ -0,0 +1,224 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.streaming + +import java.io.Serializable +import java.nio.{ByteBuffer, ByteOrder} + +import org.apache.spark.internal.Logging +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.execution.streaming.state._ +import org.apache.spark.sql.streaming.TimeoutMode +import org.apache.spark.sql.types._ +import org.apache.spark.util.NextIterator + +/** + * Singleton utils class used primarily while interacting with TimerState + */ +object TimerStateUtils { + case class TimestampWithKey( + key: Any, + expiryTimestampMs: Long) extends Serializable + + val PROC_TIMERS_STATE_NAME = "_procTimers" + val EVENT_TIMERS_STATE_NAME = "_eventTimers" + val KEY_TO_TIMESTAMP_CF = "_keyToTimestamp" + val TIMESTAMP_TO_KEY_CF = "_timestampToKey" +} + +/** + * Class that provides the implementation for storing timers + * used within the `transformWithState` operator. + * @param store - state store to be used for storing timer data + * @param timeoutMode - mode of timeout (event time or processing time) + * @param keyExprEnc - encoder for key expression + * @tparam S - type of timer value + */ +class TimerStateImpl[S]( + store: StateStore, + timeoutMode: TimeoutMode, + keyExprEnc: ExpressionEncoder[Any]) extends Logging { + + private val EMPTY_ROW = + UnsafeProjection.create(Array[DataType](NullType)).apply(InternalRow.apply(null)) + + private val schemaForPrefixKey: StructType = new StructType() + .add("key", BinaryType) + + private val schemaForKeyRow: StructType = new StructType() + .add("key", BinaryType) + .add("expiryTimestampMs", LongType, nullable = false) + + private val keySchemaForSecIndex: StructType = new StructType() + .add("expiryTimestampMs", BinaryType, nullable = false) + .add("key", BinaryType) + + private val schemaForValueRow: StructType = + StructType(Array(StructField("__dummy__", NullType))) + + private val keySerializer = keyExprEnc.createSerializer() + + private val prefixKeyEncoder = UnsafeProjection.create(schemaForPrefixKey) + + private val keyEncoder = UnsafeProjection.create(schemaForKeyRow) + + private val secIndexKeyEncoder = UnsafeProjection.create(keySchemaForSecIndex) + + val timerCfName = if (timeoutMode == TimeoutMode.ProcessingTime) { + TimerStateUtils.PROC_TIMERS_STATE_NAME + } else { + TimerStateUtils.EVENT_TIMERS_STATE_NAME + } + + val keyToTsCFName = timerCfName + TimerStateUtils.KEY_TO_TIMESTAMP_CF + store.createColFamilyIfAbsent(keyToTsCFName, + schemaForKeyRow, numColsPrefixKey = 1, + schemaForValueRow, useMultipleValuesPerKey = false, + isInternal = true) + + val tsToKeyCFName = timerCfName + TimerStateUtils.TIMESTAMP_TO_KEY_CF + store.createColFamilyIfAbsent(tsToKeyCFName, + keySchemaForSecIndex, numColsPrefixKey = 0, + schemaForValueRow, useMultipleValuesPerKey = false, + isInternal = true) + + private def encodeKey(expiryTimestampMs: Long): UnsafeRow = { + val keyOption = ImplicitGroupingKeyTracker.getImplicitKeyOption + if (!keyOption.isDefined) { + throw StateStoreErrors.implicitKeyNotFound(keyToTsCFName) + } + + val keyByteArr = keySerializer.apply(keyOption.get).asInstanceOf[UnsafeRow].getBytes() + val keyRow = keyEncoder(InternalRow(keyByteArr, expiryTimestampMs)) + keyRow + } + + // We maintain a secondary index that inverts the ordering of the timestamp + // and grouping key and maintains the list of (expiry) timestamps in sorted order + // (using BIG_ENDIAN encoding) within RocksDB. + // This is because RocksDB uses byte-wise comparison using the default comparator to + // determine sorted order of keys. This is used to read expired timers at any given + // processing time/event time timestamp threshold by performing a range scan. + private def encodeSecIndexKey(expiryTimestampMs: Long): UnsafeRow = { + val keyOption = ImplicitGroupingKeyTracker.getImplicitKeyOption + if (!keyOption.isDefined) { + throw StateStoreErrors.implicitKeyNotFound(tsToKeyCFName) + } + + val keyByteArr = keySerializer.apply(keyOption.get).asInstanceOf[UnsafeRow].getBytes() + val bbuf = ByteBuffer.allocate(8) + bbuf.order(ByteOrder.BIG_ENDIAN) + bbuf.putLong(expiryTimestampMs) + val keyRow = secIndexKeyEncoder(InternalRow(bbuf.array(), keyByteArr)) Review Comment: If I understand correctly, `bbuf.array()` is an `Array[Byte]`, which means that Spark SQL will treat this as a non-mutable type, i.e. it will encode it as a reference in the `UnsafeRow` we get back from the `UnsafeProjection`. So, then, the big-endianness of the encoding we do up above isn't really at the _start_ of the key, but rather somewhere else later on. I printed out the bytes from `keyRow` here from a test suite, and I got back: `[0,1800000008,2000000020,cce6060000000000,0,1000000009,79656b5f74736574,31]` Is this behavior intentional? Would we instead want the `UnsafeRow` to look more like `[0, bbuf[0], bbuf[1], ...`. -- 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
