ggjh-159 commented on issue #12805:
URL: https://github.com/apache/gluten/issues/12805#issuecomment-5326337239
# Velox Stateful Path Key Hash Collision Fix
## 1. Problem
The Velox `experimental/stateful/` path (streaming aggregation / join /
rank) stores the hash value as the key identity inside state. Under hash
collision, distinct keys are silently merged, corrupting data for stateful
operators.
**Concrete collision**: BIGINT keys `391` and `32728` both produce `fmix64 %
INT_MAX = 250707955`. The two keys are recognized as the same group inside
state, and accumulated values pollute each other.
**Collision is inevitable**: callers use numPartitions=INT_MAX, and by the
birthday paradox only ~2^16 distinct keys are needed for a 50% collision
probability. In other words, collisions are not "an extremely rare edge case" —
they are **guaranteed to fire in production**.
## 2. Root Cause
```
KeySelector::partition(input)
-> HashPartitionFunction::partition(...)
-> VectorHasher::hash() // compute 64-bit hash
-> hash % numPartitions // yields uint32_t
<- std::map<int64_t, RowVectorPtr> // <- this int64_t is the partition
id (hash-derived)
GroupWindowAggregator::advance():
for (auto& [key, data] : keyToData) {
windowState_->value(key, window); // <- key is a hash-derived value,
used as primary key for state lookup
windowState_->update(key, window, acc);
}
```
Along this entire path **there is no original key field comparison
whatsoever**. `KeySelector` treats hash as identity, diverging from the
semantics of Velox's native `HashAggregation` — native `HashTable::groupProbe`
still does precise key comparison via `RowContainer::equals` after a hash hit.
## 3. Fix Mechanism (Core)
The core is to **replace V1 `KeySelector::partition` with Velox native
`HashTable<false>::groupProbe`**: hash is used only for bucket positioning, and
after a hash hit, `RowContainer::equals` does precise key comparison. This path
is identical to the batch HashAggregation path, which is industry-proven.
### 3.1 KV separation
The key side (used by HashTable probe) and the value side (each State's
storage) are physically separated — this holds uniformly for all state types:
- **key side**: key RowContainer, `[user_key columns...]`, pure user key,
zero metadata fields (Velox native RowContainer, no modification)
- **value side**: each State subclass decides its own data form — AccState's
value is a RowContainer row (layout `[acc columns...]`); ValueState / ListState
/ MapState hold V / list / map respectively; none contain ns
**Why KV separation is mandatory**: in the window scenario, namespace
(window ID) is dynamically generated — all ns cannot be known before access, so
HashTable cannot do a joint (key, ns) probe. The namespace dimension must be
stripped from the HashTable and maintained by the backend — the Heap backend
carries it in StateTable's `(K, N) → State` mapping (StateTable is
Heap-specific); the RocksDB backend has no StateTable and encodes ns directly
into its composite key `[key-group][key][ns]`.
### 3.2 Three-layer architecture
| Component | Responsibility | Implementation |
|---|---|---|
| **RowContainer** (key side only; the value side is not necessarily a
RowContainer — its form is decided by each State type, e.g. AccState uses a
RowContainer, separate from the key side) | memory management | Velox native,
zero modification |
| **HashTable** (inside KeySelector utility) | grouping: find key row by
user key | Velox native, zero modification |
| **KeyedStateBackendV2\<K\>** | persistence: (K, N) → State mapping,
snapshot serialization | V2 new |
### 3.3 StateKey / Namespace abstraction
Abstract base classes (virtual + template coexist; in template context the
compiler devirtualizes to static calls):
```cpp
class StateKey {
public:
virtual bool equals(const StateKey& other) const = 0;
virtual uint64_t hash() const = 0;
virtual uint32_t keyGroup() const = 0; // derived: hash % maxParallelism
};
class Namespace {
public:
virtual bool equals(const Namespace& other) const = 0;
virtual uint64_t hash() const = 0; // participates in StateMap's
composite (K, N) hash
};
```
Subclasses (instances are immutable; hash / keyGroup cached at construction):
- `RowContainerStateKey : public StateKey` — holds `const char* row_` (view
semantics, pointing into key RowContainer)
- `VoidNamespace : public Namespace` — for non-window operators
(GroupAggregation, etc.)
- `RowContainerNamespace : public Namespace` — for window operators
(GroupWindowAggregator)
### 3.4 State subclasses (templated on K/N)
```cpp
class State { /* abstract */ };
template <typename K, typename N> class AccState : public State {
public:
// batch: resolve N keys' value row pointers in one call; keys may contain
// duplicates (per-row, feeds addRawInput directly) or be deduplicated
(distinct iteration)
void rows(folly::Range<const K*> keys, const N& ns, char** outRows);
char* row(const K& key, const N& ns);
RowContainer* valueRows();
const std::vector<std::unique_ptr<Aggregate>>& aggregates() const;
};
template <typename K, typename N, typename V> class ValueState : public
State;
template <typename K, typename N, typename T> class ListState : public State;
template <typename K, typename N, typename UK, typename UV> class MapState :
public State;
```
**Key design points**:
- **API uses explicit `(K, N, x)` parameters**, not Flink's implicit
`currentKey + currentNamespace`: Flink is single-row streaming where
`setCurrentKey` is cheap; Velox is batch processing where one advance()
processes N input rows with potentially different keys, so calling
`setCurrentKey` per access is expensive and unnatural
- **State handle is per-descriptor** (one instance per aggregate), not
per-(K,N): one ValueState / AccState instance serves all keys, same granularity
as Flink
- **K/N are both compile-time fixed on the operator side**: operator classes
hardcode K/N; on the hot path the compiler devirtualizes to static calls
### 3.5 StateDescriptor subclassing
Mirrors the Flink StateDescriptor pattern: abstract base holds name,
subclasses hold the value description fields:
```cpp
class StateDescriptor { name_ };
// velox-specific acc value description: row layout + accumulation semantics
class AccStateDescriptor : public StateDescriptor {
RowContainer* valueRows_; // describes how
acc value is stored
std::vector<std::unique_ptr<Aggregate>> aggregates_; // describes how
acc value is accumulated (native exec::Aggregate)
};
// Flink-style generic state descriptors
template <typename V> class ValueStateDescriptor : public StateDescriptor;
template <typename T> class ListStateDescriptor : public StateDescriptor;
template <typename UK, typename UV> class MapStateDescriptor : public
StateDescriptor;
```
### 3.6 KeyedStateBackendV2
K is the key type and N is the namespace type. The backend class is
templated on K only; N is a template parameter of the getOrCreate methods
(mirroring Flink `AbstractKeyedStateBackend<K>`):
```cpp
template <typename K>
class KeyedStateBackendV2 {
public:
template <typename N>
AccState<K, N>& getOrCreateAccState(const N& ns, const AccStateDescriptor&
desc);
template <typename N, typename V>
ValueState<K, N, V>& getOrCreateValueState(const N& ns, const
ValueStateDescriptor<V>& desc);
// ... ListState / MapState likewise
void snapshot(int64_t checkpointId);
void restore(...);
private:
StateKeySerializer<K>* keySerializer_;
std::unordered_map<std::string, std::unique_ptr<StateTableBase>>
stateTables_;
};
```
### 3.7 StateTable (bucketed by key-group)
```cpp
template <typename K, typename N, typename S>
class StateTable : public StateTableBase {
// Layer 1: vector, bucketed by key-group (size = maxParallelism)
// Layer 2: StateMap, a single-level map keyed by composite (K, N) — one
hash lookup
std::vector<StateMap<K, N, S>> buckets_;
};
```
Snapshot streams out bucket by bucket — no need to recompute key-group
(bucket index IS the key-group). Aligns with Flink `StateTable<K,N,S>`
structure (`StateMap[]` + StateMap internal `(K,N)→S` composite-key
single-level map).
### 3.8 KeySelector utility
```cpp
class KeySelector {
public:
void probe(const RowVectorPtr& input); //
internally calls HashTable::groupProbe
// Usage 1: per-row (main aggregation-accumulate path)
folly::Range<const RowContainerStateKey*> keys() const; // N
entries, 1:1 with input rows
// Usage 2: distinct (join build / rank / sorted agg, per-group processing)
folly::Range<const RowContainerStateKey*> distinctKeys() const; // D
deduplicated
const std::vector<SelectivityVector>& groupRows() const; // input
rows covered by key g
folly::Range<const vector_size_t> newGroups() const; //
first-occurrence row numbers of new groups
RowContainer* keyRowContainer() const;
private:
std::unique_ptr<HashTable<false>> hashTable_;
exec::HashLookup lookup_; // private: hit/hash details never exposed;
StateKey construction internal to KeySelector
};
```
**Only StateKey is exposed to operators**: hit/hash details (HashLookup)
stay fully private; `RowContainerStateKey(row, hash, maxParallelism)`
construction happens only inside KeySelector — operators touch nothing but the
StateKey type from probe to accumulation. Deduplication and the key→row mapping
(distinct usage) are also done inside KeySelector, not by operators.
### 3.9 Removing StreamKeyedOperator
The StreamKeyedOperator intermediate base class is deleted: its grouping is
inefficient and cannot use Velox's batch processing capability. Operators
directly extend StatefulOperator, compose a KeySelector + state handle, and
implement their own advance().
### 3.10 KeyGroup algorithm
```
hash = VectorHasher::hash(user_key) // 64-bit
key_group = hash % maxParallelism
```
All four places (shuffle / probe / snapshot / restore) reuse the same
algorithm naturally. Introducing `maxParallelism` decouples from parallelism
(Flink core design): rescale does not change maxParallelism → key-group
boundaries stay stable.
## 4. Operator integration
### 4.1 Class hierarchy
```
StatefulOperator (only base class)
└─ concrete stateful operators (direct inheritance, no intermediate keyed
base)
GroupAggregation / GroupWindowAggregator / WindowJoin /
AppendOnlyTopNRanker / ...
├─ composes KeySelector (does groupProbe)
└─ composes AccState<K, N> / ValueState<K, N, V> / ... handle
StreamOperatorStateHandler
└─ holds KeyedStateBackendV2<K>*
StateBackend (Task-level)
└─ createKeyedStateBackend<K>(...) → unique_ptr<KeyedStateBackendV2<K>>
```
### 4.2 Initialization flow example
```cpp
void GroupAggregation::initializeStateBackend(StateBackend* stateBackend) {
StatefulOperator::initializeStateBackend(stateBackend);
keyedStateBackend_ =
stateBackend->createKeyedStateBackend<RowContainerStateKey>(...);
}
void GroupAggregation::initializeState() {
StatefulOperator::initializeState();
valueRows_ = std::make_unique<RowContainer>(accTypes_, ..., pool_);
auto desc = AccStateDescriptor("agg", valueRows_.get(),
std::move(aggregates_), pool_);
accState_ =
&keyedStateBackend_->getOrCreateAccState(VoidNamespace::instance(), desc);
}
```
1. initializeStateBackend: the Task-level StateBackend is handed down to the
operator, which creates the KeyedStateBackendV2 instance with K specified
explicitly as the template argument (the method is currently non-virtual and
must become virtual so operators can override it to specify K).
2. Build the value RowContainer: owned by the operator; the acc row layout
is defined by the operator.
3. Build the AccStateDescriptor: holds the value description (acc row layout
+ accumulation semantics).
4. Get the state handle: the backend's getOrCreateAccState is called once
(namespace and descriptor as parameters); the returned handle is stored as an
operator member, and advance uses that member directly without re-fetching.
<img width="846" height="348" alt="Image"
src="https://github.com/user-attachments/assets/60f677c9-f4d3-4d11-8bed-b408d8cadfa8";
/>
### 4.3 advance flow example
```cpp
void GroupAggregation::advance(RowVectorPtr input) {
keySelector_.probe(input);
outRows_.resize(input->size());
accState_->rows(keySelector_.keys(), VoidNamespace::instance(),
outRows_.data());
for (auto& agg : accState_->aggregates())
agg->addRawInput(outRows_.data(), rows, ...);
}
```
1. Grouping: KeySelector internally uses the HashTable for probe-based
grouping, producing a hits sequence 1:1 with input rows for StateKey
construction;
2. State lookup: AccState queries per StateKey row by row; on miss it
automatically creates the acc row, initializes it via the Aggregate's
initializeNewGroups, and puts it into the state backend;
3. Vectorized accumulation: each Aggregate calls addRawInput; outRows is 1:1
with input rows and passed straight through as velox's groups parameter.
<img width="1654" height="502" alt="Image"
src="https://github.com/user-attachments/assets/b2097d08-59c4-429f-8f71-a42e1b42fd05";
/>
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
