HeartSaVioR commented on code in PR #45051: URL: https://github.com/apache/spark/pull/45051#discussion_r1518698125
########## 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: Yeah, good point. I see two caveats about this: https://github.com/apache/spark/blob/master/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java#L44-L58 1. Here two columns are both variable-length types - the first column will have the same relative offset and same length, so the value in reference section will all be same. The problem is in second column. The relative offset is expected to the same, but the length of second column can differ, which changes the value in reference section. It's great to have a test if we ever make the key in test as variable-length, but we should test the key with various lengths. 2. (Looks like it should not happen here but worth noting) null bits come first, which will change the orderness in overall if any of column can ever have the value of null and non-null. That said, either we have to deal with single binary column and use its own serde format, or require state store to support orderness of prefix against some column(s), like we did for prefix scan. -- 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
