This is an automated email from the ASF dual-hosted git repository.
gaborgsomogyi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new 23cec76515a [FLINK-39637][state] Add filter pushdown support for
savepoint connector
23cec76515a is described below
commit 23cec76515ad958502d5f5053fff07de3407bfd8
Author: Ilya Soin <[email protected]>
AuthorDate: Tue Jun 23 18:54:34 2026 +0300
[FLINK-39637][state] Add filter pushdown support for savepoint connector
---
.../apache/flink/state/api/SavepointReader.java | 32 +-
.../apache/flink/state/api/filter/BoundInfo.java | 45 ++
.../flink/state/api/filter/EmptyKeyFilter.java | 61 ++
.../flink/state/api/filter/ExactKeyFilter.java | 63 ++
.../flink/state/api/filter/RangeKeyFilter.java | 124 ++++
.../flink/state/api/filter/SavepointKeyFilter.java | 121 ++++
.../api/input/FilteringCloseableIterator.java | 71 ++
.../state/api/input/KeyedStateInputFormat.java | 77 +-
.../table/SavepointDataStreamScanProvider.java | 24 +-
.../state/table/SavepointDynamicTableSource.java | 40 +-
.../state/table/SavepointFilterTranslator.java | 371 ++++++++++
.../state/api/SavepointReaderKeyedStateITCase.java | 214 +++++-
.../api/input/FilteringCloseableIteratorTest.java | 121 ++++
.../state/api/input/KeyedStateInputFormatTest.java | 81 +++
.../table/SavepointDynamicTableSourceTest.java | 273 ++++++-
.../state/table/SavepointFilterTranslatorTest.java | 805 +++++++++++++++++++++
16 files changed, 2472 insertions(+), 51 deletions(-)
diff --git
a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/SavepointReader.java
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/SavepointReader.java
index 7c6bdf527d0..b556ee0739f 100644
---
a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/SavepointReader.java
+++
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/SavepointReader.java
@@ -18,6 +18,7 @@
package org.apache.flink.state.api;
+import org.apache.flink.annotation.Experimental;
import org.apache.flink.annotation.PublicEvolving;
import org.apache.flink.api.common.InvalidProgramException;
import org.apache.flink.api.common.functions.InvalidTypesException;
@@ -32,6 +33,7 @@ import org.apache.flink.runtime.checkpoint.OperatorState;
import org.apache.flink.runtime.checkpoint.metadata.CheckpointMetadata;
import org.apache.flink.runtime.state.StateBackend;
import org.apache.flink.runtime.state.VoidNamespace;
+import org.apache.flink.state.api.filter.SavepointKeyFilter;
import org.apache.flink.state.api.functions.KeyedStateReaderFunction;
import org.apache.flink.state.api.input.BroadcastStateInputFormat;
import org.apache.flink.state.api.input.KeyedStateInputFormat;
@@ -401,6 +403,33 @@ public class SavepointReader {
TypeInformation<K> keyTypeInfo,
TypeInformation<OUT> outTypeInfo)
throws IOException {
+ return readKeyedState(identifier, function, keyTypeInfo, outTypeInfo,
null);
+ }
+
+ /**
+ * Read keyed state from an operator in a {@code Savepoint}, optionally
pruning input splits and
+ * keys with a {@link SavepointKeyFilter}.
+ *
+ * @param identifier The identifier of the operator.
+ * @param function The {@link KeyedStateReaderFunction} that is called for
each key in state.
+ * @param keyTypeInfo The type information of the key in state.
+ * @param outTypeInfo The type information of the output of the transform
reader function.
+ * @param keyFilter Optional filter on the state key. When present, input
splits whose key
+ * groups cannot contain any matching key are skipped, and within each
split only matching
+ * keys are iterated.
+ * @param <K> The type of the key in state.
+ * @param <OUT> The output type of the transform function.
+ * @return A {@code DataStream} of objects read from keyed state.
+ * @throws IOException If the savepoint does not contain operator state
with the given uid.
+ */
+ @Experimental
+ public <K, OUT> DataStream<OUT> readKeyedState(
+ OperatorIdentifier identifier,
+ KeyedStateReaderFunction<K, OUT> function,
+ TypeInformation<K> keyTypeInfo,
+ TypeInformation<OUT> outTypeInfo,
+ @Nullable SavepointKeyFilter keyFilter)
+ throws IOException {
OperatorState operatorState = metadata.getOperatorState(identifier);
KeyedStateInputFormat<K, VoidNamespace, OUT> inputFormat =
@@ -409,7 +438,8 @@ public class SavepointReader {
stateBackend,
MutableConfig.of(env.getConfiguration()),
new KeyedStateReaderOperator<>(function, keyTypeInfo),
- env.getConfig());
+ env.getConfig(),
+ keyFilter);
return SourceBuilder.fromFormat(env, inputFormat, outTypeInfo);
}
diff --git
a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/filter/BoundInfo.java
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/filter/BoundInfo.java
new file mode 100644
index 00000000000..ddf3a24ef39
--- /dev/null
+++
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/filter/BoundInfo.java
@@ -0,0 +1,45 @@
+/*
+ * 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.flink.state.api.filter;
+
+import org.apache.flink.annotation.Experimental;
+
+import java.io.Serializable;
+
+/** Information about a bound in a range filter. */
+@Experimental
+public final class BoundInfo implements Serializable {
+ private static final long serialVersionUID = 4L;
+
+ private final Comparable<?> value;
+ private final boolean inclusive;
+
+ public BoundInfo(Comparable<?> value, boolean inclusive) {
+ this.value = value;
+ this.inclusive = inclusive;
+ }
+
+ public Comparable<?> getValue() {
+ return value;
+ }
+
+ public boolean isInclusive() {
+ return inclusive;
+ }
+}
diff --git
a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/filter/EmptyKeyFilter.java
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/filter/EmptyKeyFilter.java
new file mode 100644
index 00000000000..cca74d2dac2
--- /dev/null
+++
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/filter/EmptyKeyFilter.java
@@ -0,0 +1,61 @@
+/*
+ * 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.flink.state.api.filter;
+
+import java.util.Collections;
+import java.util.Set;
+
+/** A filter that rejects every key. */
+final class EmptyKeyFilter implements SavepointKeyFilter {
+
+ private static final long serialVersionUID = 1L;
+
+ static final EmptyKeyFilter INSTANCE = new EmptyKeyFilter();
+
+ private EmptyKeyFilter() {}
+
+ @Override
+ public boolean test(Object key) {
+ return false;
+ }
+
+ @Override
+ public boolean isEmpty() {
+ return true;
+ }
+
+ @Override
+ public Set<Object> getExactKeys() {
+ return Collections.emptySet();
+ }
+
+ @Override
+ public SavepointKeyFilter intersect(SavepointKeyFilter other) {
+ return this;
+ }
+
+ private Object readResolve() {
+ return INSTANCE;
+ }
+
+ @Override
+ public String toString() {
+ return "EmptyKeyFilter";
+ }
+}
diff --git
a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/filter/ExactKeyFilter.java
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/filter/ExactKeyFilter.java
new file mode 100644
index 00000000000..49b966c3f81
--- /dev/null
+++
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/filter/ExactKeyFilter.java
@@ -0,0 +1,63 @@
+/*
+ * 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.flink.state.api.filter;
+
+import java.util.HashSet;
+import java.util.Set;
+
+/** A filter that accepts a finite set of keys. */
+final class ExactKeyFilter implements SavepointKeyFilter {
+
+ private static final long serialVersionUID = 2L;
+
+ private final Set<Object> keys;
+
+ ExactKeyFilter(Set<Object> keys) {
+ this.keys = Set.copyOf(keys);
+ }
+
+ @Override
+ public boolean test(Object key) {
+ return keys.contains(key);
+ }
+
+ @Override
+ public Set<Object> getExactKeys() {
+ return keys;
+ }
+
+ @Override
+ public SavepointKeyFilter intersect(SavepointKeyFilter other) {
+ if (other.isEmpty()) {
+ return other;
+ }
+ final Set<Object> otherKeys = other.getExactKeys();
+ if (otherKeys != null) {
+ final Set<Object> intersection = new HashSet<>(keys);
+ intersection.retainAll(otherKeys);
+ return SavepointKeyFilter.exact(intersection);
+ }
+ return SavepointKeyFilter.filterKeys(keys, other);
+ }
+
+ @Override
+ public String toString() {
+ return "ExactKeyFilter" + keys;
+ }
+}
diff --git
a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/filter/RangeKeyFilter.java
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/filter/RangeKeyFilter.java
new file mode 100644
index 00000000000..3a2ee86af09
--- /dev/null
+++
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/filter/RangeKeyFilter.java
@@ -0,0 +1,124 @@
+/*
+ * 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.flink.state.api.filter;
+
+import javax.annotation.Nullable;
+
+import java.util.Set;
+
+/** A filter based on a comparable range. */
+final class RangeKeyFilter implements SavepointKeyFilter {
+
+ private static final long serialVersionUID = 3L;
+
+ @Nullable private final BoundInfo lower;
+ @Nullable private final BoundInfo upper;
+
+ RangeKeyFilter(@Nullable BoundInfo lower, @Nullable BoundInfo upper) {
+ this.lower = lower;
+ this.upper = upper;
+ }
+
+ @Override
+ public boolean test(Object key) {
+ if (lower != null) {
+ int cmp = compare(lower.getValue(), key);
+ if (cmp > 0 || (cmp == 0 && !lower.isInclusive())) {
+ return false;
+ }
+ }
+ if (upper != null) {
+ int cmp = compare(upper.getValue(), key);
+ if (cmp < 0 || (cmp == 0 && !upper.isInclusive())) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ @SuppressWarnings("unchecked")
+ private static int compare(Comparable<?> a, Object b) {
+ return ((Comparable<Object>) a).compareTo(b);
+ }
+
+ @Override
+ public BoundInfo getLowerBound() {
+ return lower;
+ }
+
+ @Override
+ public BoundInfo getUpperBound() {
+ return upper;
+ }
+
+ @Override
+ public SavepointKeyFilter intersect(SavepointKeyFilter other) {
+ if (other.isEmpty()) {
+ return other;
+ }
+ final Set<Object> otherExactKeys = other.getExactKeys();
+ if (otherExactKeys != null) {
+ return SavepointKeyFilter.filterKeys(otherExactKeys, this);
+ }
+ return intersectRange(other.getLowerBound(), other.getUpperBound());
+ }
+
+ private SavepointKeyFilter intersectRange(
+ @Nullable BoundInfo otherLower, @Nullable BoundInfo otherUpper) {
+ BoundInfo newLower = tighter(lower, otherLower, true);
+ BoundInfo newUpper = tighter(upper, otherUpper, false);
+
+ if (newLower != null && newUpper != null) {
+ int cmp = compare(newLower.getValue(), newUpper.getValue());
+ if (cmp > 0) {
+ return SavepointKeyFilter.empty();
+ }
+ if (cmp == 0 && (!newLower.isInclusive() ||
!newUpper.isInclusive())) {
+ return SavepointKeyFilter.empty();
+ }
+ }
+ return new RangeKeyFilter(newLower, newUpper);
+ }
+
+ @Nullable
+ private static BoundInfo tighter(
+ @Nullable BoundInfo a, @Nullable BoundInfo b, boolean
preferHigher) {
+ if (a == null) {
+ return b;
+ }
+ if (b == null) {
+ return a;
+ }
+ int c = compare(a.getValue(), b.getValue());
+ if (c == 0) {
+ return new BoundInfo(a.getValue(), a.isInclusive() &&
b.isInclusive());
+ }
+ boolean aWins = preferHigher ? c > 0 : c < 0;
+ return aWins ? a : b;
+ }
+
+ @Override
+ public String toString() {
+ String lowerStr =
+ lower == null ? "(-∞" : (lower.isInclusive() ? "[" : "(") +
lower.getValue();
+ String upperStr =
+ upper == null ? "+∞)" : upper.getValue() +
(upper.isInclusive() ? "]" : ")");
+ return "RangeKeyFilter" + lowerStr + ", " + upperStr;
+ }
+}
diff --git
a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/filter/SavepointKeyFilter.java
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/filter/SavepointKeyFilter.java
new file mode 100644
index 00000000000..b5a29db3043
--- /dev/null
+++
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/filter/SavepointKeyFilter.java
@@ -0,0 +1,121 @@
+/*
+ * 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.flink.state.api.filter;
+
+import org.apache.flink.annotation.Experimental;
+
+import javax.annotation.Nullable;
+
+import java.io.Serializable;
+import java.util.HashSet;
+import java.util.Set;
+
+/** Represents a key filter that can be pushed down into a savepoint scan. */
+@Experimental
+public interface SavepointKeyFilter extends Serializable {
+
+ /** Returns {@code true} if the given key passes this filter. */
+ boolean test(Object key);
+
+ /**
+ * Returns {@code true} if this filter rejects every key.
+ *
+ * <p>Used only while combining filters during push-down translation, not
during the scan.
+ */
+ default boolean isEmpty() {
+ return false;
+ }
+
+ /**
+ * Returns the finite set of keys this filter matches, or {@code null} if
the filter does not
+ * resolve to a finite key set.
+ */
+ @Nullable
+ default Set<Object> getExactKeys() {
+ return null;
+ }
+
+ /**
+ * Returns the lower bound of this filter's range, or {@code null} if the
filter does not define
+ * a lower bound.
+ *
+ * <p>Used only while combining filters during push-down translation, not
during the scan.
+ */
+ @Nullable
+ default BoundInfo getLowerBound() {
+ return null;
+ }
+
+ /**
+ * Returns the upper bound of this filter's range, or {@code null} if the
filter does not define
+ * an upper bound.
+ *
+ * <p>Used only while combining filters during push-down translation, not
during the scan.
+ */
+ @Nullable
+ default BoundInfo getUpperBound() {
+ return null;
+ }
+
+ /**
+ * Returns a filter that accepts a key if and only if both {@code this}
and {@code other} accept
+ * it.
+ *
+ * <p>Used only while combining filters during push-down translation, not
during the scan.
+ */
+ default SavepointKeyFilter intersect(SavepointKeyFilter other) {
+ throw new UnsupportedOperationException(
+ getClass().getSimpleName() + " does not support intersect()");
+ }
+
+ static SavepointKeyFilter filterKeys(Set<Object> keys, SavepointKeyFilter
predicate) {
+ final Set<Object> retained = new HashSet<>();
+ for (Object key : keys) {
+ if (predicate.test(key)) {
+ retained.add(key);
+ }
+ }
+ return exact(retained);
+ }
+
+ static SavepointKeyFilter exact(Set<Object> keys) {
+ if (keys.isEmpty()) {
+ return EmptyKeyFilter.INSTANCE;
+ }
+ return new ExactKeyFilter(keys);
+ }
+
+ static SavepointKeyFilter exact(Object value) {
+ return new ExactKeyFilter(Set.of(value));
+ }
+
+ static SavepointKeyFilter range(
+ @Nullable Comparable<?> lower,
+ boolean lowerInclusive,
+ @Nullable Comparable<?> upper,
+ boolean upperInclusive) {
+ BoundInfo lowerBoundInfo = lower != null ? new BoundInfo(lower,
lowerInclusive) : null;
+ BoundInfo upperBoundInfo = upper != null ? new BoundInfo(upper,
upperInclusive) : null;
+ return new RangeKeyFilter(lowerBoundInfo, upperBoundInfo);
+ }
+
+ static SavepointKeyFilter empty() {
+ return EmptyKeyFilter.INSTANCE;
+ }
+}
diff --git
a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/FilteringCloseableIterator.java
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/FilteringCloseableIterator.java
new file mode 100644
index 00000000000..25963e2e389
--- /dev/null
+++
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/FilteringCloseableIterator.java
@@ -0,0 +1,71 @@
+/*
+ * 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.flink.state.api.input;
+
+import org.apache.flink.util.CloseableIterator;
+
+import javax.annotation.Nullable;
+
+import java.util.NoSuchElementException;
+import java.util.function.Predicate;
+
+/** A {@link CloseableIterator} that only exposes elements matching a
predicate. */
+final class FilteringCloseableIterator<T> implements CloseableIterator<T> {
+
+ private final CloseableIterator<T> source;
+ private final Predicate<T> predicate;
+
+ @Nullable private T pending;
+ private boolean hasPending;
+
+ FilteringCloseableIterator(CloseableIterator<T> source, Predicate<T>
predicate) {
+ this.source = source;
+ this.predicate = predicate;
+ }
+
+ @Override
+ public boolean hasNext() {
+ while (!hasPending && source.hasNext()) {
+ final T candidate = source.next();
+ if (predicate.test(candidate)) {
+ pending = candidate;
+ hasPending = true;
+ }
+ }
+
+ return hasPending;
+ }
+
+ @Override
+ public T next() {
+ if (!hasNext()) {
+ throw new NoSuchElementException();
+ }
+
+ final T result = pending;
+ pending = null;
+ hasPending = false;
+ return result;
+ }
+
+ @Override
+ public void close() throws Exception {
+ source.close();
+ }
+}
diff --git
a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/KeyedStateInputFormat.java
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/KeyedStateInputFormat.java
index ae95f79048c..aba34a8c65d 100644
---
a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/KeyedStateInputFormat.java
+++
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/KeyedStateInputFormat.java
@@ -33,8 +33,10 @@ import
org.apache.flink.runtime.checkpoint.StateAssignmentOperation;
import org.apache.flink.runtime.state.AbstractKeyedStateBackend;
import org.apache.flink.runtime.state.DefaultKeyedStateStore;
import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
import org.apache.flink.runtime.state.KeyedStateHandle;
import org.apache.flink.runtime.state.StateBackend;
+import org.apache.flink.state.api.filter.SavepointKeyFilter;
import org.apache.flink.state.api.functions.KeyedStateReaderFunction;
import org.apache.flink.state.api.input.operator.StateReaderOperator;
import org.apache.flink.state.api.input.splits.KeyGroupRangeInputSplit;
@@ -54,8 +56,11 @@ import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.IOException;
+import java.util.ArrayList;
import java.util.Comparator;
+import java.util.HashSet;
import java.util.List;
+import java.util.Set;
/**
* Input format for reading partitioned state.
@@ -81,6 +86,8 @@ public class KeyedStateInputFormat<K, N, OUT>
private final SerializedValue<ExecutionConfig> serializedExecutionConfig;
+ @Nullable private final SavepointKeyFilter keyFilter;
+
private transient CloseableRegistry registry;
private transient BufferingCollector<OUT> out;
@@ -101,6 +108,27 @@ public class KeyedStateInputFormat<K, N, OUT>
StateReaderOperator<?, K, N, OUT> operator,
ExecutionConfig executionConfig)
throws IOException {
+ this(operatorState, stateBackend, configuration, operator,
executionConfig, null);
+ }
+
+ /**
+ * Creates an input format for reading partitioned state from an operator
in a savepoint.
+ *
+ * @param operatorState The state to be queried.
+ * @param stateBackend The state backed used to snapshot the operator.
+ * @param configuration The underlying Flink configuration used to
configure the state backend.
+ * @param keyFilter Optional filter on the state key. When present, splits
whose key groups
+ * cannot contain any matching key are skipped, and within each split
only matching keys are
+ * iterated.
+ */
+ public KeyedStateInputFormat(
+ OperatorState operatorState,
+ @Nullable StateBackend stateBackend,
+ Configuration configuration,
+ StateReaderOperator<?, K, N, OUT> operator,
+ ExecutionConfig executionConfig,
+ @Nullable SavepointKeyFilter keyFilter)
+ throws IOException {
Preconditions.checkNotNull(operatorState, "The operator state cannot
be null");
Preconditions.checkNotNull(configuration, "The configuration cannot be
null");
Preconditions.checkNotNull(operator, "The operator cannot be null");
@@ -114,6 +142,7 @@ public class KeyedStateInputFormat<K, N, OUT>
this.configuration = new Configuration(configuration);
this.operator = operator;
this.serializedExecutionConfig = new
SerializedValue<>(executionConfig);
+ this.keyFilter = keyFilter;
}
@Override
@@ -132,9 +161,15 @@ public class KeyedStateInputFormat<K, N, OUT>
@Override
public KeyGroupRangeInputSplit[] createInputSplits(int minNumSplits)
throws IOException {
final int maxParallelism = operatorState.getMaxParallelism();
-
final List<KeyGroupRange> keyGroups =
sortedKeyGroupRanges(minNumSplits, maxParallelism);
+ if (keyFilter != null) {
+ Set<Object> exactKeys = keyFilter.getExactKeys();
+ if (exactKeys != null) {
+ return pruneByExactKeys(keyGroups, exactKeys, maxParallelism);
+ }
+ }
+
return CollectionUtil.mapWithIndex(
keyGroups,
(keyGroupRange, index) ->
@@ -143,6 +178,28 @@ public class KeyedStateInputFormat<K, N, OUT>
.toArray(KeyGroupRangeInputSplit[]::new);
}
+ private KeyGroupRangeInputSplit[] pruneByExactKeys(
+ List<KeyGroupRange> keyGroups, Set<Object> exactKeys, int
maxParallelism) {
+ if (exactKeys.isEmpty()) {
+ return new KeyGroupRangeInputSplit[0];
+ }
+
+ final Set<Integer> targetKeyGroups = new HashSet<>();
+ for (Object key : exactKeys) {
+ targetKeyGroups.add(KeyGroupRangeAssignment.assignToKeyGroup(key,
maxParallelism));
+ }
+
+ final List<KeyGroupRangeInputSplit> prunedSplits = new ArrayList<>();
+ for (int i = 0; i < keyGroups.size(); i++) {
+ KeyGroupRange range = keyGroups.get(i);
+ if (rangeContainsAny(range, targetKeyGroups)) {
+ prunedSplits.add(
+ createKeyGroupRangeInputSplit(operatorState,
maxParallelism, range, i));
+ }
+ }
+ return prunedSplits.toArray(new KeyGroupRangeInputSplit[0]);
+ }
+
@Override
public void openInputFormat() {
out = new BufferingCollector<>();
@@ -188,7 +245,14 @@ public class KeyedStateInputFormat<K, N, OUT>
operator.setup(
runtimeContext::createSerializer, keyedStateBackend,
timeServiceManager, ctx);
operator.open();
- keysAndNamespaces = operator.getKeysAndNamespaces(ctx);
+ if (keyFilter != null) {
+ keysAndNamespaces =
+ new FilteringCloseableIterator<>(
+ operator.getKeysAndNamespaces(ctx),
+ keyAndNamespace ->
keyFilter.test(keyAndNamespace.f0));
+ } else {
+ keysAndNamespaces = operator.getKeysAndNamespaces(ctx);
+ }
} catch (Exception e) {
throw new IOException("Failed to restore timer state", e);
}
@@ -229,6 +293,15 @@ public class KeyedStateInputFormat<K, N, OUT>
return out.next();
}
+ private static boolean rangeContainsAny(KeyGroupRange range, Set<Integer>
keyGroups) {
+ for (int kg : keyGroups) {
+ if (range.contains(kg)) {
+ return true;
+ }
+ }
+ return false;
+ }
+
private static KeyGroupRangeInputSplit createKeyGroupRangeInputSplit(
OperatorState operatorState,
int maxParallelism,
diff --git
a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointDataStreamScanProvider.java
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointDataStreamScanProvider.java
index fd24b0f2448..35a9e04a7af 100644
---
a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointDataStreamScanProvider.java
+++
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointDataStreamScanProvider.java
@@ -30,6 +30,7 @@ import org.apache.flink.runtime.state.StateBackend;
import org.apache.flink.runtime.state.StateBackendLoader;
import org.apache.flink.state.api.OperatorIdentifier;
import org.apache.flink.state.api.SavepointReader;
+import org.apache.flink.state.api.filter.SavepointKeyFilter;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.table.connector.ProviderContext;
@@ -53,6 +54,7 @@ public class SavepointDataStreamScanProvider implements
DataStreamScanProvider {
private final TypeInformation keyTypeInfo;
private final Tuple2<Integer, List<StateValueColumnConfiguration>>
keyValueProjections;
private final RowType rowType;
+ @Nullable private final SavepointKeyFilter keyFilter;
public SavepointDataStreamScanProvider(
@Nullable final String stateBackendType,
@@ -61,12 +63,31 @@ public class SavepointDataStreamScanProvider implements
DataStreamScanProvider {
final TypeInformation keyTypeInfo,
final Tuple2<Integer, List<StateValueColumnConfiguration>>
keyValueProjections,
RowType rowType) {
+ this(
+ stateBackendType,
+ statePath,
+ operatorIdentifier,
+ keyTypeInfo,
+ keyValueProjections,
+ rowType,
+ null);
+ }
+
+ public SavepointDataStreamScanProvider(
+ @Nullable final String stateBackendType,
+ final String statePath,
+ final OperatorIdentifier operatorIdentifier,
+ final TypeInformation keyTypeInfo,
+ final Tuple2<Integer, List<StateValueColumnConfiguration>>
keyValueProjections,
+ RowType rowType,
+ @Nullable SavepointKeyFilter keyFilter) {
this.stateBackendType = stateBackendType;
this.statePath = statePath;
this.operatorIdentifier = operatorIdentifier;
this.keyTypeInfo = keyTypeInfo;
this.keyValueProjections = keyValueProjections;
this.rowType = rowType;
+ this.keyFilter = keyFilter;
}
@Override
@@ -133,7 +154,8 @@ public class SavepointDataStreamScanProvider implements
DataStreamScanProvider {
operatorIdentifier,
new KeyedStateReader(rowType, keyValueProjections),
keyTypeInfo,
- outTypeInfo);
+ outTypeInfo,
+ keyFilter);
} catch (Exception e) {
throw new RuntimeException(e);
}
diff --git
a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointDynamicTableSource.java
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointDynamicTableSource.java
index 0bb3bb36f9c..bc8e30c243e 100644
---
a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointDynamicTableSource.java
+++
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointDynamicTableSource.java
@@ -21,10 +21,14 @@ package org.apache.flink.state.table;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.state.api.OperatorIdentifier;
+import org.apache.flink.state.api.filter.SavepointKeyFilter;
import org.apache.flink.table.connector.ChangelogMode;
import org.apache.flink.table.connector.source.DynamicTableSource;
import org.apache.flink.table.connector.source.ScanTableSource;
+import
org.apache.flink.table.connector.source.abilities.SupportsFilterPushDown;
+import org.apache.flink.table.expressions.ResolvedExpression;
import org.apache.flink.table.types.logical.RowType;
+import org.apache.flink.table.types.utils.TypeConversions;
import javax.annotation.Nullable;
@@ -32,13 +36,14 @@ import java.util.List;
/** Savepoint dynamic source. */
@SuppressWarnings("rawtypes")
-public class SavepointDynamicTableSource implements ScanTableSource {
+public class SavepointDynamicTableSource implements ScanTableSource,
SupportsFilterPushDown {
@Nullable private final String stateBackendType;
private final String statePath;
private final OperatorIdentifier operatorIdentifier;
private final TypeInformation keyTypeInfo;
private final Tuple2<Integer, List<StateValueColumnConfiguration>>
keyValueProjections;
private final RowType rowType;
+ @Nullable private SavepointKeyFilter keyFilter;
public SavepointDynamicTableSource(
@Nullable final String stateBackendType,
@@ -57,13 +62,29 @@ public class SavepointDynamicTableSource implements
ScanTableSource {
@Override
public DynamicTableSource copy() {
- return new SavepointDynamicTableSource(
- stateBackendType,
- statePath,
- operatorIdentifier,
- keyTypeInfo,
- keyValueProjections,
- rowType);
+ SavepointDynamicTableSource copy =
+ new SavepointDynamicTableSource(
+ stateBackendType,
+ statePath,
+ operatorIdentifier,
+ keyTypeInfo,
+ keyValueProjections,
+ rowType);
+ copy.keyFilter = this.keyFilter;
+ return copy;
+ }
+
+ @Override
+ public Result applyFilters(List<ResolvedExpression> filters) {
+ final int keyColumnIndex = keyValueProjections.f0;
+ final SavepointFilterTranslator.Result result =
+ new SavepointFilterTranslator(
+ keyColumnIndex,
+ TypeConversions.fromLogicalToDataType(
+ rowType.getTypeAt(keyColumnIndex)))
+ .apply(filters);
+ keyFilter = result.keyFilter();
+ return Result.of(result.accepted(), result.remaining());
}
@Override
@@ -84,6 +105,7 @@ public class SavepointDynamicTableSource implements
ScanTableSource {
operatorIdentifier,
keyTypeInfo,
keyValueProjections,
- rowType);
+ rowType,
+ keyFilter);
}
}
diff --git
a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointFilterTranslator.java
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointFilterTranslator.java
new file mode 100644
index 00000000000..001973ed4ed
--- /dev/null
+++
b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointFilterTranslator.java
@@ -0,0 +1,371 @@
+/*
+ * 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.flink.state.table;
+
+import org.apache.flink.state.api.filter.SavepointKeyFilter;
+import org.apache.flink.table.expressions.CallExpression;
+import org.apache.flink.table.expressions.FieldReferenceExpression;
+import org.apache.flink.table.expressions.ResolvedExpression;
+import org.apache.flink.table.expressions.ValueLiteralExpression;
+import org.apache.flink.table.functions.BuiltInFunctionDefinitions;
+import org.apache.flink.table.functions.FunctionDefinition;
+import org.apache.flink.table.types.DataType;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.Nullable;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.BiFunction;
+
+/**
+ * Converts {@link ResolvedExpression} key filter predicates into {@link
SavepointKeyFilter}
+ * instances that can be used to prune key groups and key iterations during
savepoint reads.
+ */
+class SavepointFilterTranslator {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(SavepointFilterTranslator.class);
+
+ private static final Map<
+ FunctionDefinition,
+ BiFunction<SavepointFilterTranslator, CallExpression,
SavepointKeyFilter>>
+ FILTERS =
+ Map.of(
+ BuiltInFunctionDefinitions.EQUALS,
+ SavepointFilterTranslator::fromEquals,
+ BuiltInFunctionDefinitions.OR,
+ SavepointFilterTranslator::fromOr,
+ BuiltInFunctionDefinitions.AND,
+ SavepointFilterTranslator::fromAnd,
+ BuiltInFunctionDefinitions.BETWEEN,
+ SavepointFilterTranslator::fromBetween,
+ BuiltInFunctionDefinitions.GREATER_THAN,
+ SavepointFilterTranslator::fromGreaterThan,
+ BuiltInFunctionDefinitions.GREATER_THAN_OR_EQUAL,
+ SavepointFilterTranslator::fromGreaterThanOrEqual,
+ BuiltInFunctionDefinitions.LESS_THAN,
+ SavepointFilterTranslator::fromLessThan,
+ BuiltInFunctionDefinitions.LESS_THAN_OR_EQUAL,
+ SavepointFilterTranslator::fromLessThanOrEqual);
+
+ private final int keyColumnIndex;
+ private final DataType keyColumnType;
+
+ SavepointFilterTranslator(int keyColumnIndex, DataType keyColumnType) {
+ this.keyColumnIndex = keyColumnIndex;
+ this.keyColumnType = keyColumnType;
+ }
+
+ Result apply(List<ResolvedExpression> filters) {
+ final List<ResolvedExpression> accepted = new ArrayList<>();
+ final List<ResolvedExpression> remaining = new ArrayList<>();
+
+ SavepointKeyFilter keyFilter = null;
+ for (ResolvedExpression filter : filters) {
+ SavepointKeyFilter extracted = extractFilter(filter);
+ if (extracted == null) {
+ remaining.add(filter);
+ continue;
+ }
+
+ keyFilter = keyFilter == null ? extracted :
keyFilter.intersect(extracted);
+ accepted.add(filter);
+ }
+
+ return new Result(accepted, remaining, keyFilter);
+ }
+
+ @Nullable
+ private SavepointKeyFilter extractFilter(ResolvedExpression expr) {
+ final BiFunction<SavepointFilterTranslator, CallExpression,
SavepointKeyFilter> extractor =
+ expr instanceof CallExpression
+ ? FILTERS.get(((CallExpression)
expr).getFunctionDefinition())
+ : null;
+ if (extractor == null) {
+ LOG.debug(
+ "Unsupported predicate [{}] cannot be pushed into
savepoint key filter.", expr);
+ return null;
+ }
+ return extractor.apply(this, (CallExpression) expr);
+ }
+
+ //
-------------------------------------------------------------------------
+ // Equality
+ //
-------------------------------------------------------------------------
+
+ @Nullable
+ private SavepointKeyFilter fromEquals(CallExpression call) {
+ if (!isBinaryValid(call)) {
+ return null;
+ }
+ ResolvedExpression left = call.getResolvedChildren().get(0);
+ ResolvedExpression right = call.getResolvedChildren().get(1);
+
+ Object value = null;
+ if (isKeyField(left)) {
+ value = extractValue(right);
+ } else if (isKeyField(right)) {
+ value = extractValue(left);
+ }
+
+ if (value == null) {
+ return null;
+ }
+ return SavepointKeyFilter.exact(value);
+ }
+
+ @Nullable
+ private SavepointKeyFilter fromOr(CallExpression call) {
+ Set<Object> keys = new HashSet<>();
+ for (ResolvedExpression arg : call.getResolvedChildren()) {
+ SavepointKeyFilter sub = extractFilter(arg);
+ if (sub == null) {
+ return null;
+ }
+ Set<Object> subKeys = sub.getExactKeys();
+ // OR can only absorb finite key sets; a range branch cannot be
merged via union.
+ if (subKeys == null) {
+ return null;
+ }
+ keys.addAll(subKeys);
+ }
+ return SavepointKeyFilter.exact(keys);
+ }
+
+ //
-------------------------------------------------------------------------
+ // Range
+ //
-------------------------------------------------------------------------
+
+ @Nullable
+ private SavepointKeyFilter fromAnd(CallExpression call) {
+ SavepointKeyFilter merged = null;
+ for (ResolvedExpression arg : call.getResolvedChildren()) {
+ SavepointKeyFilter sub = extractFilter(arg);
+ // AND only absorbs range filters; exact (or null) children break
pushdown.
+ if (sub == null || sub.getExactKeys() != null) {
+ return null;
+ }
+ merged = (merged == null) ? sub : merged.intersect(sub);
+ if (merged.isEmpty()) {
+ return merged;
+ }
+ }
+ return merged;
+ }
+
+ @Nullable
+ private SavepointKeyFilter fromBetween(CallExpression call) {
+ List<ResolvedExpression> args = call.getResolvedChildren();
+ if (args.size() != 3) {
+ return null;
+ }
+ ResolvedExpression valueExpr = args.get(0);
+ ResolvedExpression lowerExpr = args.get(1);
+ ResolvedExpression upperExpr = args.get(2);
+
+ if (!isKeyField(valueExpr)) {
+ return null;
+ }
+
+ Object lower = extractValue(lowerExpr);
+ Object upper = extractValue(upperExpr);
+ if (lower == null || upper == null) {
+ return null;
+ }
+ if (!(lower instanceof Comparable) || !(upper instanceof Comparable)) {
+ LOG.debug(
+ "BETWEEN predicate on non-comparable key type {} cannot be
pushed into savepoint key filter.",
+ lower.getClass().getName());
+ return null;
+ }
+ return SavepointKeyFilter.range(
+ (Comparable<?>) lower, true,
+ (Comparable<?>) upper, true);
+ }
+
+ @Nullable
+ private SavepointKeyFilter fromGreaterThan(CallExpression call) {
+ return fromComparison(call, Comparison.GT);
+ }
+
+ @Nullable
+ private SavepointKeyFilter fromGreaterThanOrEqual(CallExpression call) {
+ return fromComparison(call, Comparison.GTE);
+ }
+
+ @Nullable
+ private SavepointKeyFilter fromLessThan(CallExpression call) {
+ return fromComparison(call, Comparison.LT);
+ }
+
+ @Nullable
+ private SavepointKeyFilter fromLessThanOrEqual(CallExpression call) {
+ return fromComparison(call, Comparison.LTE);
+ }
+
+ @Nullable
+ private SavepointKeyFilter fromComparison(CallExpression call, Comparison
cmp) {
+ if (!isBinaryValid(call)) {
+ return null;
+ }
+ ResolvedExpression left = call.getResolvedChildren().get(0);
+ ResolvedExpression right = call.getResolvedChildren().get(1);
+
+ final boolean keyOnLeft = isKeyField(left);
+ final boolean keyOnRight = isKeyField(right);
+ if (!keyOnLeft && !keyOnRight) {
+ return null;
+ }
+ Object bound = extractValue(keyOnLeft ? right : left);
+ if (bound == null) {
+ return null;
+ }
+ if (!(bound instanceof Comparable)) {
+ LOG.debug(
+ "Range predicate on non-comparable key type {} cannot be
pushed into savepoint key filter.",
+ bound.getClass().getName());
+ return null;
+ }
+ Comparable<?> b = (Comparable<?>) bound;
+ Comparison keyLeftCmp = keyOnLeft ? cmp : cmp.flip();
+ switch (keyLeftCmp) {
+ case GT:
+ return SavepointKeyFilter.range(b, false, null, true);
+ case GTE:
+ return SavepointKeyFilter.range(b, true, null, true);
+ case LT:
+ return SavepointKeyFilter.range(null, true, b, false);
+ case LTE:
+ return SavepointKeyFilter.range(null, true, b, true);
+ default:
+ throw new IllegalStateException("Unknown Comparison: " +
keyLeftCmp);
+ }
+ }
+
+ //
-------------------------------------------------------------------------
+ // Helpers
+ //
-------------------------------------------------------------------------
+
+ private static boolean isBinaryValid(CallExpression call) {
+ return call.getResolvedChildren().size() == 2;
+ }
+
+ private boolean isKeyField(ResolvedExpression expr) {
+ return expr instanceof FieldReferenceExpression
+ && ((FieldReferenceExpression) expr).getFieldIndex() ==
keyColumnIndex;
+ }
+
+ @Nullable
+ private Object extractValue(ResolvedExpression expr) {
+ if (!(expr instanceof ValueLiteralExpression)) {
+ LOG.debug("Refusing pushdown: predicate operand [{}] is not a
literal value.", expr);
+ return null;
+ }
+ ValueLiteralExpression literal = (ValueLiteralExpression) expr;
+ Class<?> literalClass =
literal.getOutputDataType().getConversionClass();
+ Object value = literal.getValueAs(literalClass).orElse(null);
+ if (value == null) {
+ LOG.debug(
+ "Refusing pushdown: literal {} of type {} cannot be read
as its conversion"
+ + " class {}.",
+ literal,
+ literal.getOutputDataType(),
+ literalClass.getName());
+ return null;
+ }
+ return widenToKeyType(value);
+ }
+
+ @Nullable
+ private Object widenToKeyType(Object value) {
+ Class<?> keyClass = keyColumnType.getConversionClass();
+ if (keyClass.isInstance(value)) {
+ return value;
+ }
+ if (value instanceof Number) {
+ if (keyClass == Long.class) {
+ return ((Number) value).longValue();
+ }
+ if (keyClass == Double.class) {
+ return ((Number) value).doubleValue();
+ }
+ }
+ LOG.debug(
+ "Refusing pushdown: literal value {} of type {} cannot be
widened to key type {}.",
+ value,
+ value.getClass().getName(),
+ keyColumnType);
+ return null;
+ }
+
+ static final class Result {
+ private final List<ResolvedExpression> accepted;
+ private final List<ResolvedExpression> remaining;
+ @Nullable private final SavepointKeyFilter keyFilter;
+
+ private Result(
+ List<ResolvedExpression> accepted,
+ List<ResolvedExpression> remaining,
+ @Nullable SavepointKeyFilter keyFilter) {
+ this.accepted = accepted;
+ this.remaining = remaining;
+ this.keyFilter = keyFilter;
+ }
+
+ List<ResolvedExpression> accepted() {
+ return accepted;
+ }
+
+ List<ResolvedExpression> remaining() {
+ return remaining;
+ }
+
+ @Nullable
+ SavepointKeyFilter keyFilter() {
+ return keyFilter;
+ }
+ }
+
+ private enum Comparison {
+ GT,
+ GTE,
+ LT,
+ LTE;
+
+ Comparison flip() {
+ switch (this) {
+ case GT:
+ return LT;
+ case GTE:
+ return LTE;
+ case LT:
+ return GT;
+ case LTE:
+ return GTE;
+ default:
+ throw new IllegalStateException("Unknown Comparison: " +
this);
+ }
+ }
+ }
+}
diff --git
a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/SavepointReaderKeyedStateITCase.java
b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/SavepointReaderKeyedStateITCase.java
index de3bdcd5f3a..5df10085c8a 100644
---
a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/SavepointReaderKeyedStateITCase.java
+++
b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/SavepointReaderKeyedStateITCase.java
@@ -18,6 +18,8 @@
package org.apache.flink.state.api;
+import org.apache.flink.api.common.JobExecutionResult;
+import org.apache.flink.api.common.accumulators.IntCounter;
import org.apache.flink.api.common.functions.OpenContext;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
@@ -25,23 +27,29 @@ import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.StateBackend;
+import org.apache.flink.state.api.filter.SavepointKeyFilter;
import org.apache.flink.state.api.functions.KeyedStateReaderFunction;
import org.apache.flink.state.api.utils.JobResultRetriever;
import org.apache.flink.state.api.utils.SavepointTestBase;
+import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.streaming.api.functions.sink.v2.DiscardingSink;
+import org.apache.flink.streaming.util.testing.CollectingSink;
import org.apache.flink.util.Collector;
import org.junit.Assert;
import org.junit.Test;
-import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Objects;
import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static org.assertj.core.api.Assertions.assertThat;
/** IT case for reading state. */
public abstract class SavepointReaderKeyedStateITCase<B extends StateBackend>
@@ -51,8 +59,10 @@ public abstract class SavepointReaderKeyedStateITCase<B
extends StateBackend>
private static ValueStateDescriptor<Integer> valueState =
new ValueStateDescriptor<>("value", Types.INT);
+ private static final String STATE_READS_ACCUMULATOR =
"stateReadsAccumulator";
+
private static final List<Pojo> elements =
- Arrays.asList(Pojo.of(1, 1), Pojo.of(2, 2), Pojo.of(3, 3));
+ IntStream.rangeClosed(1, 10).mapToObj(i -> Pojo.of(i,
i)).collect(Collectors.toList());
protected abstract Tuple2<Configuration, B> getStateBackendTuple();
@@ -63,13 +73,7 @@ public abstract class SavepointReaderKeyedStateITCase<B
extends StateBackend>
StreamExecutionEnvironment.getExecutionEnvironment(backendTuple.f0);
env.setParallelism(4);
- env.addSource(createSource(elements))
- .returns(Pojo.class)
- .rebalance()
- .keyBy(id -> id.key)
- .process(new KeyedStatefulOperator())
- .uid(uid)
- .sinkTo(new DiscardingSink<>());
+ applyStatefulPipeline(env);
String savepointPath = takeSavepoint(env);
@@ -85,6 +89,170 @@ public abstract class SavepointReaderKeyedStateITCase<B
extends StateBackend>
"Unexpected results from keyed state", expected, new
HashSet<>(results));
}
+ @Test
+ public void testReadKeyedStateWithExactFilter() throws Exception {
+ Tuple2<Configuration, B> backendTuple = getStateBackendTuple();
+ StreamExecutionEnvironment env =
+
StreamExecutionEnvironment.getExecutionEnvironment(backendTuple.f0);
+ env.setParallelism(4);
+
+ applyStatefulPipeline(env);
+
+ String savepointPath = takeSavepoint(env);
+
+ SavepointReader savepoint = SavepointReader.read(env, savepointPath,
backendTuple.f1);
+ CountingReadResult result =
+ readKeyedStateWithCountingReader(savepoint,
SavepointKeyFilter.exact(5));
+ // Only key=5 reaches the reader and reads state.
+ assertThat(result.values).containsExactly(5);
+ assertThat(result.counter).isEqualTo(1);
+ }
+
+ @Test
+ public void testReadKeyedStateWithMultiKeyExactFilter() throws Exception {
+ Tuple2<Configuration, B> backendTuple = getStateBackendTuple();
+ StreamExecutionEnvironment env =
+
StreamExecutionEnvironment.getExecutionEnvironment(backendTuple.f0);
+ env.setParallelism(4);
+
+ applyStatefulPipeline(env);
+
+ String savepointPath = takeSavepoint(env);
+
+ SavepointReader savepoint = SavepointReader.read(env, savepointPath,
backendTuple.f1);
+ CountingReadResult result =
+ readKeyedStateWithCountingReader(
+ savepoint, SavepointKeyFilter.exact(Set.of(3, 5, 7)));
+ assertThat(result.values).containsExactlyInAnyOrder(3, 5, 7);
+ assertThat(result.counter).isEqualTo(3);
+ }
+
+ @Test
+ public void testReadKeyedStateWithInclusiveRangeFilter() throws Exception {
+ Tuple2<Configuration, B> backendTuple = getStateBackendTuple();
+ StreamExecutionEnvironment env =
+
StreamExecutionEnvironment.getExecutionEnvironment(backendTuple.f0);
+ env.setParallelism(4);
+
+ applyStatefulPipeline(env);
+
+ String savepointPath = takeSavepoint(env);
+
+ SavepointReader savepoint = SavepointReader.read(env, savepointPath,
backendTuple.f1);
+ CountingReadResult result =
+ readKeyedStateWithCountingReader(
+ savepoint, SavepointKeyFilter.range(3, true, 6, true));
+ // [3, 6]: only the in-range keys reach the reader and read state.
+ assertThat(result.values).containsExactlyInAnyOrder(3, 4, 5, 6);
+ assertThat(result.counter).isEqualTo(4);
+ }
+
+ @Test
+ public void
testReadKeyedStateWithInclusiveLowerExclusiveUpperRangeFilter() throws
Exception {
+ Tuple2<Configuration, B> backendTuple = getStateBackendTuple();
+ StreamExecutionEnvironment env =
+
StreamExecutionEnvironment.getExecutionEnvironment(backendTuple.f0);
+ env.setParallelism(4);
+
+ applyStatefulPipeline(env);
+
+ String savepointPath = takeSavepoint(env);
+
+ SavepointReader savepoint = SavepointReader.read(env, savepointPath,
backendTuple.f1);
+ CountingReadResult result =
+ readKeyedStateWithCountingReader(
+ savepoint, SavepointKeyFilter.range(3, true, 6,
false));
+ // [3, 6): the exclusive upper bound drops key 6.
+ assertThat(result.values).containsExactlyInAnyOrder(3, 4, 5);
+ assertThat(result.counter).isEqualTo(3);
+ }
+
+ @Test
+ public void
testReadKeyedStateWithExclusiveLowerInclusiveUpperRangeFilter() throws
Exception {
+ Tuple2<Configuration, B> backendTuple = getStateBackendTuple();
+ StreamExecutionEnvironment env =
+
StreamExecutionEnvironment.getExecutionEnvironment(backendTuple.f0);
+ env.setParallelism(4);
+
+ applyStatefulPipeline(env);
+
+ String savepointPath = takeSavepoint(env);
+
+ SavepointReader savepoint = SavepointReader.read(env, savepointPath,
backendTuple.f1);
+ CountingReadResult result =
+ readKeyedStateWithCountingReader(
+ savepoint, SavepointKeyFilter.range(3, false, 6,
true));
+ // (3, 6]: the exclusive lower bound drops key 3.
+ assertThat(result.values).containsExactlyInAnyOrder(4, 5, 6);
+ assertThat(result.counter).isEqualTo(3);
+ }
+
+ @Test
+ public void testReadKeyedStateWithExclusiveRangeFilter() throws Exception {
+ Tuple2<Configuration, B> backendTuple = getStateBackendTuple();
+ StreamExecutionEnvironment env =
+
StreamExecutionEnvironment.getExecutionEnvironment(backendTuple.f0);
+ env.setParallelism(4);
+
+ applyStatefulPipeline(env);
+
+ String savepointPath = takeSavepoint(env);
+
+ SavepointReader savepoint = SavepointReader.read(env, savepointPath,
backendTuple.f1);
+ CountingReadResult result =
+ readKeyedStateWithCountingReader(
+ savepoint, SavepointKeyFilter.range(3, false, 6,
false));
+ // (3, 6): both bounds exclusive, dropping keys 3 and 6.
+ assertThat(result.values).containsExactlyInAnyOrder(4, 5);
+ assertThat(result.counter).isEqualTo(2);
+ }
+
+ @Test
+ public void testReadKeyedStateWithEmptyFilter() throws Exception {
+ Tuple2<Configuration, B> backendTuple = getStateBackendTuple();
+ StreamExecutionEnvironment env =
+
StreamExecutionEnvironment.getExecutionEnvironment(backendTuple.f0);
+ env.setParallelism(4);
+
+ applyStatefulPipeline(env);
+
+ String savepointPath = takeSavepoint(env);
+
+ SavepointReader savepoint = SavepointReader.read(env, savepointPath,
backendTuple.f1);
+ CountingReadResult result =
+ readKeyedStateWithCountingReader(savepoint,
SavepointKeyFilter.empty());
+ // No key reaches the reader, so no state is read.
+ assertThat(result.values).isEmpty();
+ assertThat(result.counter).isZero();
+ }
+
+ private void applyStatefulPipeline(StreamExecutionEnvironment env) {
+ env.addSource(createSource(elements))
+ .returns(Pojo.class)
+ .rebalance()
+ .keyBy(id -> id.key)
+ .process(new KeyedStatefulOperator())
+ .uid(uid)
+ .sinkTo(new DiscardingSink<>());
+ }
+
+ private CountingReadResult readKeyedStateWithCountingReader(
+ SavepointReader savepoint, SavepointKeyFilter keyFilter) throws
Exception {
+ DataStream<Integer> stateValues =
+ savepoint.readKeyedState(
+ OperatorIdentifier.forUid(uid),
+ new CountingReader(),
+ Types.INT,
+ Types.INT,
+ keyFilter);
+ CollectingSink<Integer> collectingSink = new CollectingSink<>();
+ stateValues.sinkTo(collectingSink);
+ JobExecutionResult jobResult =
stateValues.getExecutionEnvironment().execute();
+ Integer counter =
jobResult.getAccumulatorResult(STATE_READS_ACCUMULATOR);
+ return new CountingReadResult(
+ collectingSink.getRemainingOutput(), counter == null ? 0 :
counter);
+ }
+
private static class KeyedStatefulOperator extends
KeyedProcessFunction<Integer, Pojo, Void> {
private transient ValueState<Integer> state;
@@ -124,6 +292,34 @@ public abstract class SavepointReaderKeyedStateITCase<B
extends StateBackend>
}
}
+ private static class CountingReader extends
KeyedStateReaderFunction<Integer, Integer> {
+
+ private transient ValueState<Integer> state;
+ private transient IntCounter counter;
+
+ @Override
+ public void open(OpenContext openContext) {
+ state = getRuntimeContext().getState(valueState);
+ counter =
getRuntimeContext().getIntCounter(STATE_READS_ACCUMULATOR);
+ }
+
+ @Override
+ public void readKey(Integer key, Context ctx, Collector<Integer> out)
throws Exception {
+ counter.add(1);
+ out.collect(state.value());
+ }
+ }
+
+ private static class CountingReadResult {
+ private final List<Integer> values;
+ private final int counter;
+
+ private CountingReadResult(List<Integer> values, int counter) {
+ this.values = values;
+ this.counter = counter;
+ }
+ }
+
/** A simple pojo type. */
public static class Pojo {
public static Pojo of(Integer key, Integer state) {
diff --git
a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/FilteringCloseableIteratorTest.java
b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/FilteringCloseableIteratorTest.java
new file mode 100644
index 00000000000..600ff77839f
--- /dev/null
+++
b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/FilteringCloseableIteratorTest.java
@@ -0,0 +1,121 @@
+/*
+ * 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.flink.state.api.input;
+
+import org.apache.flink.util.CloseableIterator;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import java.util.NoSuchElementException;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+class FilteringCloseableIteratorTest {
+
+ @Test
+ void predicateMatchReturnsMatchingElements() throws Exception {
+ try (FilteringCloseableIterator<Integer> iterator =
+ new FilteringCloseableIterator<>(
+ CloseableIterator.fromList(Arrays.asList(1, 2, 3, 4),
value -> {}),
+ value -> value % 2 == 0)) {
+ assertThat(iterator.hasNext()).isTrue();
+ assertThat(iterator.next()).isEqualTo(2);
+ assertThat(iterator.hasNext()).isTrue();
+ assertThat(iterator.next()).isEqualTo(4);
+ assertThat(iterator.hasNext()).isFalse();
+ }
+ }
+
+ @Test
+ void predicateSkipReturnsNoElements() throws Exception {
+ try (FilteringCloseableIterator<Integer> iterator =
+ new FilteringCloseableIterator<>(
+ CloseableIterator.fromList(Arrays.asList(1, 2, 3),
value -> {}),
+ value -> value > 10)) {
+ assertThat(iterator.hasNext()).isFalse();
+
assertThatThrownBy(iterator::next).isInstanceOf(NoSuchElementException.class);
+ }
+ }
+
+ @Test
+ void emptySourceReturnsNoElements() throws Exception {
+ try (FilteringCloseableIterator<Integer> iterator =
+ new FilteringCloseableIterator<>(
+
CloseableIterator.fromList(Collections.<Integer>emptyList(), value -> {}),
+ value -> true)) {
+ assertThat(iterator.hasNext()).isFalse();
+
assertThatThrownBy(iterator::next).isInstanceOf(NoSuchElementException.class);
+ }
+ }
+
+ @Test
+ void closeDelegatesToSource() throws Exception {
+ final TestCloseableIterator source = new
TestCloseableIterator(Arrays.asList(1, 2, 3));
+ final FilteringCloseableIterator<Integer> iterator =
+ new FilteringCloseableIterator<>(source, value -> true);
+
+ iterator.close();
+
+ assertThat(source.closed).isTrue();
+ }
+
+ @Test
+ void hasNextMultipleCallsDoesNotAdvanceTwice() throws Exception {
+ try (FilteringCloseableIterator<Integer> iterator =
+ new FilteringCloseableIterator<>(
+ CloseableIterator.fromList(Arrays.asList(1, 2, 3),
value -> {}),
+ value -> value >= 2)) {
+ assertThat(iterator.hasNext()).isTrue();
+ assertThat(iterator.hasNext()).isTrue();
+ assertThat(iterator.next()).isEqualTo(2);
+ assertThat(iterator.hasNext()).isTrue();
+ assertThat(iterator.next()).isEqualTo(3);
+ assertThat(iterator.hasNext()).isFalse();
+ }
+ }
+
+ private static final class TestCloseableIterator implements
CloseableIterator<Integer> {
+ private final Iterator<Integer> iterator;
+ private boolean closed;
+
+ private TestCloseableIterator(List<Integer> values) {
+ this.iterator = values.iterator();
+ }
+
+ @Override
+ public boolean hasNext() {
+ return iterator.hasNext();
+ }
+
+ @Override
+ public Integer next() {
+ return iterator.next();
+ }
+
+ @Override
+ public void close() {
+ closed = true;
+ }
+ }
+}
diff --git
a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/KeyedStateInputFormatTest.java
b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/KeyedStateInputFormatTest.java
index 75018679b75..292eb61d2c4 100644
---
a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/KeyedStateInputFormatTest.java
+++
b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/KeyedStateInputFormatTest.java
@@ -32,6 +32,7 @@ import
org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
import org.apache.flink.runtime.jobgraph.OperatorID;
import org.apache.flink.runtime.state.VoidNamespace;
import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
+import org.apache.flink.state.api.filter.SavepointKeyFilter;
import org.apache.flink.state.api.functions.KeyedStateReaderFunction;
import org.apache.flink.state.api.input.operator.KeyedStateReaderOperator;
import org.apache.flink.state.api.input.splits.KeyGroupRangeInputSplit;
@@ -62,6 +63,7 @@ import java.util.Comparator;
import java.util.List;
import java.util.Set;
+import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
/** Tests for keyed state input format. */
@@ -116,6 +118,85 @@ class KeyedStateInputFormatTest {
splits.length);
}
+ @ParameterizedTest(name = "Enable async state = {0}")
+ @ValueSource(booleans = {false, true})
+ void testExactKeyFilterPrunesInputSplits(boolean asyncState) throws
Exception {
+ OperatorID operatorID = OperatorIDGenerator.fromUid("uid");
+
+ OperatorSubtaskState state =
+ createOperatorSubtaskState(createFlatMap(asyncState),
asyncState);
+ OperatorState operatorState = new OperatorState(null, null,
operatorID, 1, 128);
+ operatorState.putState(0, state);
+
+ KeyedStateInputFormat<?, ?, ?> format =
+ new KeyedStateInputFormat<>(
+ operatorState,
+ new HashMapStateBackend(),
+ new Configuration(),
+ new KeyedStateReaderOperator<>(new ReaderFunction(),
Types.INT),
+ new ExecutionConfig(),
+ SavepointKeyFilter.exact(5));
+ KeyGroupRangeInputSplit[] splits = format.createInputSplits(10);
+
+ assertThat(splits)
+ .as("Single-key exact filter maps to exactly one key group
range")
+ .hasSize(1);
+ }
+
+ @ParameterizedTest(name = "Enable async state = {0}")
+ @ValueSource(booleans = {false, true})
+ void testEmptyFilterProducesNoInputSplits(boolean asyncState) throws
Exception {
+ OperatorID operatorID = OperatorIDGenerator.fromUid("uid");
+
+ OperatorSubtaskState state =
+ createOperatorSubtaskState(createFlatMap(asyncState),
asyncState);
+ OperatorState operatorState = new OperatorState(null, null,
operatorID, 1, 128);
+ operatorState.putState(0, state);
+
+ KeyedStateInputFormat<?, ?, ?> format =
+ new KeyedStateInputFormat<>(
+ operatorState,
+ new HashMapStateBackend(),
+ new Configuration(),
+ new KeyedStateReaderOperator<>(new ReaderFunction(),
Types.INT),
+ new ExecutionConfig(),
+ SavepointKeyFilter.empty());
+ KeyGroupRangeInputSplit[] splits = format.createInputSplits(10);
+
+ assertThat(splits).isEmpty();
+ }
+
+ @ParameterizedTest(name = "Enable async state = {0}")
+ @ValueSource(booleans = {false, true})
+ void testRangeFilterDoesNotPruneInputSplits(boolean asyncState) throws
Exception {
+ OperatorID operatorID = OperatorIDGenerator.fromUid("uid");
+
+ OperatorSubtaskState state =
+ createOperatorSubtaskState(createFlatMap(asyncState),
asyncState);
+ OperatorState operatorState = new OperatorState(null, null,
operatorID, 1, 128);
+ operatorState.putState(0, state);
+
+ KeyedStateInputFormat<?, ?, ?> formatRange =
+ new KeyedStateInputFormat<>(
+ operatorState,
+ new HashMapStateBackend(),
+ new Configuration(),
+ new KeyedStateReaderOperator<>(new ReaderFunction(),
Types.INT),
+ new ExecutionConfig(),
+ SavepointKeyFilter.range(3, true, 7, true));
+ KeyedStateInputFormat<?, ?, ?> formatNoFilter =
+ new KeyedStateInputFormat<>(
+ operatorState,
+ new HashMapStateBackend(),
+ new Configuration(),
+ new KeyedStateReaderOperator<>(new ReaderFunction(),
Types.INT),
+ new ExecutionConfig());
+
+ assertThat(formatRange.createInputSplits(10))
+ .as("Range filters cannot prune key-group splits")
+ .hasSize(formatNoFilter.createInputSplits(10).length);
+ }
+
@ParameterizedTest(name = "Enable async state = {0}")
@ValueSource(booleans = {false, true})
void testReadState(boolean asyncState) throws Exception {
diff --git
a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/table/SavepointDynamicTableSourceTest.java
b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/table/SavepointDynamicTableSourceTest.java
index 5ddcc9fab64..e2a8ca574d0 100644
---
a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/table/SavepointDynamicTableSourceTest.java
+++
b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/table/SavepointDynamicTableSourceTest.java
@@ -31,6 +31,7 @@ import org.junit.jupiter.api.Test;
import java.util.List;
import java.util.Map;
import java.util.Set;
+import java.util.regex.Pattern;
import java.util.stream.Collectors;
import java.util.stream.LongStream;
@@ -39,6 +40,22 @@ import static org.assertj.core.api.Assertions.assertThat;
/** Unit tests for the savepoint SQL reader. */
class SavepointDynamicTableSourceTest {
+
+ private static final String STATE_TABLE_DDL =
+ "CREATE TABLE state_table (\n"
+ + " k bigint,\n"
+ + " KeyedPrimitiveValue bigint,\n"
+ + " KeyedPojoValue ROW<privateLong bigint, publicLong
bigint>,\n"
+ + " KeyedPrimitiveValueList ARRAY<bigint>,\n"
+ + " KeyedPrimitiveValueMap MAP<string, bigint>,\n"
+ + " PRIMARY KEY (k) NOT ENFORCED\n"
+ + ")\n"
+ + "with (\n"
+ + " 'connector' = 'savepoint',\n"
+ + " 'state.path' = 'src/test/resources/table-state',\n"
+ + " 'operator.uid' = 'keyed-state-process-uid'\n"
+ + ")";
+
@Test
@SuppressWarnings("unchecked")
public void testReadKeyedState() throws Exception {
@@ -47,21 +64,7 @@ class SavepointDynamicTableSourceTest {
StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment(config);
StreamTableEnvironment tEnv = StreamTableEnvironment.create(env);
- final String sql =
- "CREATE TABLE state_table (\n"
- + " k bigint,\n"
- + " KeyedPrimitiveValue bigint,\n"
- + " KeyedPojoValue ROW<privateLong bigint, publicLong
bigint>,\n"
- + " KeyedPrimitiveValueList ARRAY<bigint>,\n"
- + " KeyedPrimitiveValueMap MAP<string, bigint>,\n"
- + " PRIMARY KEY (k) NOT ENFORCED\n"
- + ")\n"
- + "with (\n"
- + " 'connector' = 'savepoint',\n"
- + " 'state.path' =
'src/test/resources/table-state',\n"
- + " 'operator.uid' = 'keyed-state-process-uid'\n"
- + ")";
- tEnv.executeSql(sql);
+ tEnv.executeSql(STATE_TABLE_DDL);
Table table = tEnv.sqlQuery("SELECT * FROM state_table");
List<Row> result = tEnv.toDataStream(table).executeAndCollect(100);
@@ -85,7 +88,7 @@ class SavepointDynamicTableSourceTest {
result.stream()
.map(r -> (Row) r.getField("KeyedPojoValue"))
.collect(Collectors.toSet());
- assertThat(pojoValues.size()).isEqualTo(1);
+ assertThat(pojoValues).hasSize(1);
Row pojoData = pojoValues.iterator().next();
assertThat(pojoData.getField("publicLong")).isEqualTo(1L);
assertThat(pojoData.getField("privateLong")).isEqualTo(1L);
@@ -100,10 +103,9 @@ class SavepointDynamicTableSourceTest {
(Long[])
r.getField("KeyedPrimitiveValueList")))
.flatMap(l -> Set.of(l).stream())
.collect(Collectors.toSet());
- assertThat(listValues.size()).isEqualTo(10);
- for (Tuple2<Long, Long[]> tuple2 : listValues) {
- assertThat(tuple2.f0).isEqualTo(tuple2.f1[0]);
- }
+ assertThat(listValues)
+ .hasSize(10)
+ .allSatisfy(tuple2 ->
assertThat(tuple2.f0).isEqualTo(tuple2.f1[0]));
// Check map state
Set<Tuple2<Long, Map<String, Long>>> mapValues =
@@ -116,12 +118,14 @@ class SavepointDynamicTableSourceTest {
r.getField("KeyedPrimitiveValueMap")))
.flatMap(l -> Set.of(l).stream())
.collect(Collectors.toSet());
- assertThat(mapValues.size()).isEqualTo(10);
- for (Tuple2<Long, Map<String, Long>> tuple2 : mapValues) {
- assertThat(tuple2.f1.size()).isEqualTo(1);
- String expectedKey = String.valueOf(tuple2.f0);
- assertThat(tuple2.f1.get(expectedKey)).isEqualTo(tuple2.f0);
- }
+ assertThat(mapValues)
+ .hasSize(10)
+ .allSatisfy(
+ tuple2 -> {
+ assertThat(tuple2.f1).hasSize(1);
+ String expectedKey = String.valueOf(tuple2.f0);
+
assertThat(tuple2.f1.get(expectedKey)).isEqualTo(tuple2.f0);
+ });
}
@Test
@@ -146,7 +150,7 @@ class SavepointDynamicTableSourceTest {
tEnv.executeSql(sql);
Table table = tEnv.sqlQuery("SELECT * FROM state_table");
List<Row> result = tEnv.toDataStream(table).executeAndCollect(100);
- assertThat(result.size()).isEqualTo(5);
+ assertThat(result).hasSize(5);
List<Long> keys =
result.stream().map(row -> (Long)
row.getField("k")).collect(Collectors.toList());
@@ -201,7 +205,7 @@ class SavepointDynamicTableSourceTest {
result.stream()
.map(r -> (Row) r.getField("KeyedSpecificAvroValue"))
.collect(Collectors.toSet());
- assertThat(specificAvroValues.size()).isEqualTo(1);
+ assertThat(specificAvroValues).hasSize(1);
Row avroData = specificAvroValues.iterator().next();
assertThat(avroData.getField("longData")).isEqualTo(1L);
@@ -209,8 +213,219 @@ class SavepointDynamicTableSourceTest {
result.stream()
.map(r -> (String) r.getField("KeyedGenericAvroValue"))
.collect(Collectors.toSet());
- assertThat(genericAvroValues.size()).isEqualTo(1);
+ assertThat(genericAvroValues).hasSize(1);
String avroGenericValue = genericAvroValues.iterator().next();
assertThat(avroGenericValue).isEqualTo("{\"longData\": 1}");
}
+
+ //
-------------------------------------------------------------------------
+ // Filter push-down tests
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void testFilterPushDownEqualityReturnsOnlyMatchingKey() throws Exception {
+ StreamTableEnvironment tEnv = createBatchTableEnv();
+ tEnv.executeSql(STATE_TABLE_DDL);
+
+ String sql = "SELECT k FROM state_table WHERE k = 5";
+
+ assertThat(hasPushedDownFilter(tEnv, sql)).isTrue();
+
+ List<Row> result =
tEnv.toDataStream(tEnv.sqlQuery(sql)).executeAndCollect(100);
+
+ assertThat(result).hasSize(1);
+ assertThat(result.get(0).getField("k")).isEqualTo(5L);
+ }
+
+ @Test
+ @SuppressWarnings("unchecked")
+ void testFilterPushDownEqualityReturnsCorrectResult() throws Exception {
+ StreamTableEnvironment tEnv = createBatchTableEnv();
+ tEnv.executeSql(STATE_TABLE_DDL);
+
+ String sql = "SELECT * FROM state_table WHERE k = 5";
+
+ assertThat(hasPushedDownFilter(tEnv, sql)).isTrue();
+
+ List<Row> result =
tEnv.toDataStream(tEnv.sqlQuery(sql)).executeAndCollect(100);
+
+ assertThat(result).hasSize(1);
+ Row row = result.get(0);
+ assertThat(row.getField("k")).isEqualTo(5L);
+ assertThat(row.getField("KeyedPrimitiveValue")).isEqualTo(1L);
+
+ Row pojo = (Row) row.getField("KeyedPojoValue");
+ assertThat(pojo.getField("privateLong")).isEqualTo(1L);
+ assertThat(pojo.getField("publicLong")).isEqualTo(1L);
+
+ Long[] list = (Long[]) row.getField("KeyedPrimitiveValueList");
+ assertThat(list).containsExactly(5L);
+
+ Map<String, Long> map = (Map<String, Long>)
row.getField("KeyedPrimitiveValueMap");
+ assertThat(map).containsExactlyEntriesOf(Map.of("5", 5L));
+ }
+
+ @Test
+ void testFilterPushDownRangeReturnsCorrectResult() throws Exception {
+ StreamTableEnvironment tEnv = createBatchTableEnv();
+ tEnv.executeSql(STATE_TABLE_DDL);
+
+ String sql = "SELECT k FROM state_table WHERE k >= 7 ORDER BY k";
+
+ assertThat(hasPushedDownFilter(tEnv, sql)).isTrue();
+
+ List<Row> result =
tEnv.toDataStream(tEnv.sqlQuery(sql)).executeAndCollect(100);
+
+ assertThat(result).hasSize(3);
+ assertThat(result.get(0).getField("k")).isEqualTo(7L);
+ assertThat(result.get(1).getField("k")).isEqualTo(8L);
+ assertThat(result.get(2).getField("k")).isEqualTo(9L);
+ }
+
+ @Test
+ void testFilterPushDownNonexistentKeyReturnsEmpty() throws Exception {
+ StreamTableEnvironment tEnv = createBatchTableEnv();
+ tEnv.executeSql(STATE_TABLE_DDL);
+
+ String sql = "SELECT k FROM state_table WHERE k = 999";
+
+ assertThat(hasPushedDownFilter(tEnv, sql)).isTrue();
+
+ List<Row> result =
tEnv.toDataStream(tEnv.sqlQuery(sql)).executeAndCollect(100);
+
+ assertThat(result).isEmpty();
+ }
+
+ @Test
+ void testFilterPushDownInListReturnsOnlyMatchingKeys() throws Exception {
+ StreamTableEnvironment tEnv = createBatchTableEnv();
+ tEnv.executeSql(STATE_TABLE_DDL);
+
+ String sql = "SELECT k FROM state_table WHERE k IN (3, 7) ORDER BY k";
+
+ assertThat(hasPushedDownFilter(tEnv, sql)).isTrue();
+
+ List<Row> result =
tEnv.toDataStream(tEnv.sqlQuery(sql)).executeAndCollect(100);
+
+ assertThat(result).hasSize(2);
+ assertThat(result.get(0).getField("k")).isEqualTo(3L);
+ assertThat(result.get(1).getField("k")).isEqualTo(7L);
+ }
+
+ @Test
+ @SuppressWarnings("unchecked")
+ void testFilterPushDownPartialPushDown() throws Exception {
+ // When the WHERE clause contains both a key filter and a non-key
filter,
+ // both must be applied correctly regardless of which is pushed into
the source.
+ StreamTableEnvironment tEnv = createBatchTableEnv();
+ tEnv.executeSql(STATE_TABLE_DDL);
+
+ String sql =
+ "SELECT k, KeyedPrimitiveValueMap FROM state_table"
+ + " WHERE k = 5 AND KeyedPrimitiveValueMap['5'] > 3";
+
+ assertThat(hasPushedDownFilter(tEnv, sql)).isTrue();
+
+ List<Row> result =
tEnv.toDataStream(tEnv.sqlQuery(sql)).executeAndCollect(100);
+
+ assertThat(result).hasSize(1);
+ assertThat(result.get(0).getField("k")).isEqualTo(5L);
+ Map<String, Long> map =
+ (Map<String, Long>)
result.get(0).getField("KeyedPrimitiveValueMap");
+ assertThat(map).containsEntry("5", 5L);
+ }
+
+ @Test
+ void testFilterPushDownBetweenReturnsCorrectResult() throws Exception {
+ StreamTableEnvironment tEnv = createBatchTableEnv();
+ tEnv.executeSql(STATE_TABLE_DDL);
+
+ String sql = "SELECT k FROM state_table WHERE k BETWEEN 3 AND 6 ORDER
BY k";
+
+ assertThat(hasPushedDownFilter(tEnv, sql)).isTrue();
+
+ List<Row> result =
tEnv.toDataStream(tEnv.sqlQuery(sql)).executeAndCollect(100);
+
+ assertThat(result).hasSize(4);
+ assertThat(result.get(0).getField("k")).isEqualTo(3L);
+ assertThat(result.get(1).getField("k")).isEqualTo(4L);
+ assertThat(result.get(2).getField("k")).isEqualTo(5L);
+ assertThat(result.get(3).getField("k")).isEqualTo(6L);
+ }
+
+ @Test
+ void testFilterPushDownLiteralOnLeftSide() throws Exception {
+ // verify that "5 = k" (literal on the left) works the same as "k = 5".
+ StreamTableEnvironment tEnv = createBatchTableEnv();
+ tEnv.executeSql(STATE_TABLE_DDL);
+
+ String sql = "SELECT k FROM state_table WHERE 5 = k";
+
+ assertThat(hasPushedDownFilter(tEnv, sql)).isTrue();
+
+ List<Row> result =
tEnv.toDataStream(tEnv.sqlQuery(sql)).executeAndCollect(100);
+
+ assertThat(result).hasSize(1);
+ assertThat(result.get(0).getField("k")).isEqualTo(5L);
+ }
+
+ @Test
+ @SuppressWarnings("unchecked")
+ void
testOrAcrossKeyAndNonKeyColumnIsNotPushedDownButReturnsCorrectResult() throws
Exception {
+ // OR involving a non-pushable column: correctness must be preserved
+ // regardless of whether the planner pushes the filter or not.
+ StreamTableEnvironment tEnv = createBatchTableEnv();
+ tEnv.executeSql(STATE_TABLE_DDL);
+
+ String sql =
+ "SELECT k, KeyedPrimitiveValueMap FROM state_table"
+ + " WHERE k = 5 OR KeyedPrimitiveValueMap['0'] = 0"
+ + " ORDER BY k";
+
+ assertThat(hasPushedDownFilter(tEnv, sql)).isFalse();
+
+ List<Row> result =
tEnv.toDataStream(tEnv.sqlQuery(sql)).executeAndCollect(100);
+
+ assertThat(result).hasSize(2);
+ assertThat(result.get(0).getField("k")).isEqualTo(0L);
+ Map<String, Long> map0 =
+ (Map<String, Long>)
result.get(0).getField("KeyedPrimitiveValueMap");
+ assertThat(map0).containsEntry("0", 0L);
+ assertThat(result.get(1).getField("k")).isEqualTo(5L);
+ }
+
+ @Test
+ void testUnsupportedFilterIsNotPushedDownButReturnsCorrectResult() throws
Exception {
+ StreamTableEnvironment tEnv = createBatchTableEnv();
+ tEnv.executeSql(STATE_TABLE_DDL);
+
+ String sql = "SELECT k FROM state_table WHERE k % 2 = 0 ORDER BY k";
+
+ assertThat(hasPushedDownFilter(tEnv, sql)).isFalse();
+
+ List<Row> result =
tEnv.toDataStream(tEnv.sqlQuery(sql)).executeAndCollect(100);
+
+ List<Long> keys =
+ result.stream().map(r -> (Long)
r.getField("k")).collect(Collectors.toList());
+ assertThat(keys).containsExactly(0L, 2L, 4L, 6L, 8L);
+ }
+
+ //
-------------------------------------------------------------------------
+ // Helpers
+ //
-------------------------------------------------------------------------
+
+ private static StreamTableEnvironment createBatchTableEnv() {
+ Configuration config = new Configuration();
+ config.set(RUNTIME_MODE, RuntimeExecutionMode.BATCH);
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment(config);
+ return StreamTableEnvironment.create(env);
+ }
+
+ private static final Pattern PUSHED_DOWN_FILTER =
+ Pattern.compile(
+ "TableSourceScan\\(table=\\[\\[default_catalog,
default_database, state_table, filter=\\[.+?]]]");
+
+ private static boolean hasPushedDownFilter(StreamTableEnvironment tEnv,
String sql) {
+ return PUSHED_DOWN_FILTER.matcher(tEnv.explainSql(sql)).find();
+ }
}
diff --git
a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/table/SavepointFilterTranslatorTest.java
b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/table/SavepointFilterTranslatorTest.java
new file mode 100644
index 00000000000..b2e4bc49cec
--- /dev/null
+++
b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/table/SavepointFilterTranslatorTest.java
@@ -0,0 +1,805 @@
+/*
+ * 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.flink.state.table;
+
+import org.apache.flink.state.api.filter.SavepointKeyFilter;
+import org.apache.flink.table.api.DataTypes;
+import org.apache.flink.table.expressions.CallExpression;
+import org.apache.flink.table.expressions.FieldReferenceExpression;
+import org.apache.flink.table.expressions.ResolvedExpression;
+import org.apache.flink.table.expressions.ValueLiteralExpression;
+import org.apache.flink.table.functions.BuiltInFunctionDefinitions;
+import org.apache.flink.table.types.DataType;
+
+import org.junit.jupiter.api.Test;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.math.BigDecimal;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+/** Unit tests for {@link SavepointFilterTranslator}. */
+class SavepointFilterTranslatorTest {
+
+ private static final int KEY_COL = 0;
+ private static final DataType LONG_KEY_TYPE = DataTypes.BIGINT().notNull();
+
+ //
-------------------------------------------------------------------------
+ // Exact key filter — fromEquals
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void equalsKeyOnLeft() {
+ SavepointKeyFilter filter = keyFilterOf(eq(longKeyRef(),
longLit(42L)));
+ assertNotNull(filter);
+ assertThat(filter.getExactKeys()).containsExactly(42L);
+ assertThat(filter.test(42L)).isTrue();
+ assertThat(filter.test(43L)).isFalse();
+ }
+
+ @Test
+ void equalsKeyOnRight() {
+ SavepointKeyFilter filter = keyFilterOf(eq(longLit(42L),
longKeyRef()));
+ assertNotNull(filter);
+ assertThat(filter.getExactKeys()).containsExactly(42L);
+ assertThat(filter.test(42L)).isTrue();
+ assertThat(filter.test(0L)).isFalse();
+ }
+
+ @Test
+ void equalsNeitherSideIsKeyColumn_returnsNull() {
+ SavepointKeyFilter filter = keyFilterOf(eq(otherRef(), longLit(42L)));
+ assertThat(filter).isNull();
+ }
+
+ @Test
+ void equalsNeitherSideIsLiteral_returnsNull() {
+ SavepointKeyFilter filter = keyFilterOf(eq(longKeyRef(), otherRef()));
+ assertThat(filter).isNull();
+ }
+
+ //
-------------------------------------------------------------------------
+ // Exact key filter — fromOr
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void orOfEqualsProducesMergedExactFilter() {
+ CallExpression expr =
+ or(
+ eq(longKeyRef(), longLit(1L)),
+ eq(longKeyRef(), longLit(2L)),
+ eq(longLit(3L), longKeyRef()));
+
+ SavepointKeyFilter filter = keyFilterOf(expr);
+ assertNotNull(filter);
+ assertThat(filter.getExactKeys()).containsExactlyInAnyOrder(1L, 2L,
3L);
+ assertThat(filter.test(4L)).isFalse();
+ }
+
+ @Test
+ void orWithNonPushableChild_returnsNull() {
+ // One branch is a range, which OR cannot absorb
+ CallExpression expr = or(eq(longKeyRef(), longLit(1L)),
gt(longKeyRef(), longLit(5L)));
+ assertThat(keyFilterOf(expr)).isNull();
+ }
+
+ //
-------------------------------------------------------------------------
+ // Range key filter — fromBetween
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void betweenProducesInclusiveRange() {
+ SavepointKeyFilter filter = keyFilterOf(between(longKeyRef(),
longLit(10L), longLit(20L)));
+
+ assertNotNull(filter);
+ assertThat(filter.getExactKeys()).isNull();
+ assertThat(filter.test(9L)).isFalse();
+ assertThat(filter.test(10L)).isTrue();
+ assertThat(filter.test(15L)).isTrue();
+ assertThat(filter.test(20L)).isTrue();
+ assertThat(filter.test(21L)).isFalse();
+ }
+
+ @Test
+ void betweenWithNonKeyField_returnsNull() {
+ SavepointKeyFilter filter = keyFilterOf(between(otherRef(),
longLit(1L), longLit(10L)));
+ assertThat(filter).isNull();
+ }
+
+ //
-------------------------------------------------------------------------
+ // Range key filter — comparison operators
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void greaterThanProducesExclusiveLowerBound() {
+ SavepointKeyFilter filter = keyFilterOf(gt(longKeyRef(), longLit(5L)));
+ assertNotNull(filter);
+ assertThat(filter.getExactKeys()).isNull();
+ assertThat(filter.test(5L)).isFalse();
+ assertThat(filter.test(6L)).isTrue();
+ }
+
+ @Test
+ void greaterThanOrEqualProducesInclusiveLowerBound() {
+ SavepointKeyFilter filter = keyFilterOf(gte(longKeyRef(),
longLit(5L)));
+ assertNotNull(filter);
+ assertThat(filter.getExactKeys()).isNull();
+ assertThat(filter.test(4L)).isFalse();
+ assertThat(filter.test(5L)).isTrue();
+ assertThat(filter.test(6L)).isTrue();
+ }
+
+ @Test
+ void lessThanProducesExclusiveUpperBound() {
+ SavepointKeyFilter filter = keyFilterOf(lt(longKeyRef(),
longLit(10L)));
+ assertNotNull(filter);
+ assertThat(filter.getExactKeys()).isNull();
+ assertThat(filter.test(9L)).isTrue();
+ assertThat(filter.test(10L)).isFalse();
+ }
+
+ @Test
+ void lessThanOrEqualProducesInclusiveUpperBound() {
+ SavepointKeyFilter filter = keyFilterOf(lte(longKeyRef(),
longLit(10L)));
+ assertNotNull(filter);
+ assertThat(filter.getExactKeys()).isNull();
+ assertThat(filter.test(10L)).isTrue();
+ assertThat(filter.test(11L)).isFalse();
+ }
+
+ @Test
+ void comparisonWithLiteralOnLeft_flipsDirection() {
+ // literal > key → key < literal → upper bound (exclusive)
+ SavepointKeyFilter filter = keyFilterOf(gt(longLit(10L),
longKeyRef()));
+ assertNotNull(filter);
+ assertThat(filter.getExactKeys()).isNull();
+ assertThat(filter.test(9L)).isTrue();
+ assertThat(filter.test(10L)).isFalse();
+ }
+
+ @Test
+ void comparisonWithLiteralOnLeft_lte_flipsDirection() {
+ // literal <= key → key >= literal → lower bound (inclusive)
+ SavepointKeyFilter filter = keyFilterOf(lte(longLit(5L),
longKeyRef()));
+ assertThat(filter.getExactKeys()).isNull();
+ assertThat(filter.test(4L)).isFalse();
+ assertThat(filter.test(5L)).isTrue();
+ }
+
+ //
-------------------------------------------------------------------------
+ // AND — range intersection
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void andOfTwoRangesProducesIntersection() {
+ // key >= 5 AND key <= 10
+ CallExpression expr = and(gte(longKeyRef(), longLit(5L)),
lte(longKeyRef(), longLit(10L)));
+ SavepointKeyFilter filter = keyFilterOf(expr);
+
+ assertNotNull(filter);
+ assertThat(filter.getExactKeys()).isNull();
+ assertThat(filter.test(4L)).isFalse();
+ assertThat(filter.test(5L)).isTrue();
+ assertThat(filter.test(10L)).isTrue();
+ assertThat(filter.test(11L)).isFalse();
+ }
+
+ @Test
+ void andWithProvablyEmptyIntersection_matchesNothing() {
+ // key > 10 AND key < 5 — disjoint
+ CallExpression expr = and(gt(longKeyRef(), longLit(10L)),
lt(longKeyRef(), longLit(5L)));
+ SavepointKeyFilter filter = keyFilterOf(expr);
+
+ assertNotNull(filter);
+ assertThat(filter.isEmpty()).isTrue();
+ }
+
+ @Test
+ void andWithExactKeyChildIsNotPushable() {
+ // AND requires all children to be range filters; exact filter breaks
pushdown
+ CallExpression expr = and(eq(longKeyRef(), longLit(5L)),
gt(longKeyRef(), longLit(3L)));
+ assertThat(keyFilterOf(expr)).isNull();
+ }
+
+ //
-------------------------------------------------------------------------
+ // Unsupported predicates
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void nonCallExpressionReturnsNull() {
+ assertThat(keyFilterOf(longKeyRef())).isNull();
+ }
+
+ @Test
+ void unrecognizedFunctionReturnsNull() {
+ CallExpression isNull =
+ CallExpression.permanent(
+ BuiltInFunctionDefinitions.IS_NULL,
+ Collections.singletonList(longKeyRef()),
+ DataTypes.BOOLEAN());
+ assertThat(keyFilterOf(isNull)).isNull();
+ }
+
+ //
-------------------------------------------------------------------------
+ // Range key filter behavior — string keys (natural comparison)
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void rangeFilterOnStringKey() {
+ SavepointKeyFilter filter =
+ keyFilterOf(between(stringKeyRef(), stringLit("beta"),
stringLit("delta")));
+
+ assertNotNull(filter);
+ assertThat(filter.test("alpha")).isFalse();
+ assertThat(filter.test("beta")).isTrue();
+ assertThat(filter.test("gamma")).isFalse(); // "gamma" > "delta"
lexicographically
+ assertThat(filter.test("delta")).isTrue();
+ assertThat(filter.test("epsilon")).isFalse();
+ }
+
+ @Test
+ void rangeFilterWithDoubleComparison() {
+ ValueLiteralExpression floatLower =
+ new ValueLiteralExpression(1.5f, DataTypes.FLOAT().notNull());
+ ValueLiteralExpression floatUpper =
+ new ValueLiteralExpression(3.5f, DataTypes.FLOAT().notNull());
+ FieldReferenceExpression floatKey =
+ new FieldReferenceExpression("key",
DataTypes.FLOAT().notNull(), 0, KEY_COL);
+
+ SavepointKeyFilter filter = keyFilterOf(between(floatKey, floatLower,
floatUpper));
+
+ assertNotNull(filter);
+ assertThat(filter.test(1.5f)).isTrue();
+ assertThat(filter.test(2.0f)).isTrue();
+ assertThat(filter.test(3.5f)).isTrue();
+ assertThat(filter.test(1.0f)).isFalse();
+ assertThat(filter.test(4.0f)).isFalse();
+ }
+
+ //
-------------------------------------------------------------------------
+ // Range intersection
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void intersectNarrowsBounds() {
+ // [5, ∞) ∩ (-∞, 10] = [5, 10]
+ SavepointKeyFilter lower = SavepointKeyFilter.range(5L, true, null,
true);
+ SavepointKeyFilter upper = SavepointKeyFilter.range(null, true, 10L,
true);
+ SavepointKeyFilter result = lower.intersect(upper);
+
+ assertThat(result.isEmpty()).isFalse();
+ assertThat(result.getExactKeys()).isNull();
+ assertThat(result.test(4L)).isFalse();
+ assertThat(result.test(5L)).isTrue();
+ assertThat(result.test(10L)).isTrue();
+ assertThat(result.test(11L)).isFalse();
+ }
+
+ @Test
+ void intersectDisjointRangesReturnsEmpty() {
+ // [10, ∞) ∩ (-∞, 5] — disjoint
+ SavepointKeyFilter a = SavepointKeyFilter.range(10L, true, null, true);
+ SavepointKeyFilter b = SavepointKeyFilter.range(null, true, 5L, true);
+ assertThat(a.intersect(b).isEmpty()).isTrue();
+ }
+
+ @Test
+ void intersectEqualBoundsInclusiveIsNonEmpty() {
+ // [7, ∞) ∩ (-∞, 7] = [7, 7]
+ SavepointKeyFilter a = SavepointKeyFilter.range(7L, true, null, true);
+ SavepointKeyFilter b = SavepointKeyFilter.range(null, true, 7L, true);
+ SavepointKeyFilter result = a.intersect(b);
+ assertThat(result.isEmpty()).isFalse();
+ assertThat(result.test(7L)).isTrue();
+ assertThat(result.test(6L)).isFalse();
+ assertThat(result.test(8L)).isFalse();
+ }
+
+ @Test
+ void intersectEqualBoundsOneExclusiveIsEmpty() {
+ // (7, ∞) ∩ (-∞, 7] — empty because lower is exclusive
+ SavepointKeyFilter a = SavepointKeyFilter.range(7L, false, null, true);
+ SavepointKeyFilter b = SavepointKeyFilter.range(null, true, 7L, true);
+ assertThat(a.intersect(b).isEmpty()).isTrue();
+ }
+
+ //
-------------------------------------------------------------------------
+ // SavepointFilters.apply — intersection handling
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void applyAccumulatesRangeAndRange() {
+ SavepointFilterTranslator.Result applied =
+ apply(
+ List.of(gte(longKeyRef(), longLit(3L)),
lte(longKeyRef(), longLit(8L))),
+ LONG_KEY_TYPE);
+ SavepointKeyFilter result = applied.keyFilter();
+ assertNotNull(result);
+
+ assertThat(applied.accepted()).hasSize(2);
+ assertThat(applied.remaining()).isEmpty();
+ assertThat(result.test(2L)).isFalse();
+ assertThat(result.test(3L)).isTrue();
+ assertThat(result.test(8L)).isTrue();
+ assertThat(result.test(9L)).isFalse();
+ }
+
+ @Test
+ void applyAccumulatesExactAndExact() {
+ SavepointFilterTranslator.Result applied =
+ apply(
+ List.of(
+ or(
+ eq(longKeyRef(), longLit(1L)),
+ eq(longKeyRef(), longLit(2L)),
+ eq(longKeyRef(), longLit(3L))),
+ or(
+ eq(longKeyRef(), longLit(2L)),
+ eq(longKeyRef(), longLit(3L)),
+ eq(longKeyRef(), longLit(4L)))),
+ LONG_KEY_TYPE);
+ SavepointKeyFilter result = applied.keyFilter();
+
+ assertNotNull(result);
+ assertThat(applied.accepted()).hasSize(2);
+ assertThat(applied.remaining()).isEmpty();
+ assertThat(result.getExactKeys()).containsExactlyInAnyOrder(2L, 3L);
+ }
+
+ @Test
+ void applyAccumulatesExactAndExactEmptyResult_matchesNothing() {
+ SavepointFilterTranslator.Result applied =
+ apply(
+ List.of(eq(longKeyRef(), longLit(1L)),
eq(longKeyRef(), longLit(2L))),
+ LONG_KEY_TYPE);
+ SavepointKeyFilter result = applied.keyFilter();
+
+ assertNotNull(result);
+ assertThat(applied.accepted()).hasSize(2);
+ assertThat(applied.remaining()).isEmpty();
+ assertThat(result.isEmpty()).isTrue();
+ }
+
+ @Test
+ void applyAccumulatesExactAndRange_keepsOnlyKeysInRange() {
+ SavepointFilterTranslator.Result applied =
+ apply(
+ List.of(
+ or(
+ eq(longKeyRef(), longLit(1L)),
+ eq(longKeyRef(), longLit(5L)),
+ eq(longKeyRef(), longLit(10L)),
+ eq(longKeyRef(), longLit(15L))),
+ between(longKeyRef(), longLit(4L),
longLit(12L))),
+ LONG_KEY_TYPE);
+ SavepointKeyFilter result = applied.keyFilter();
+
+ assertNotNull(result);
+ assertThat(applied.accepted()).hasSize(2);
+ assertThat(applied.remaining()).isEmpty();
+ assertThat(result.getExactKeys()).containsExactlyInAnyOrder(5L, 10L);
+ }
+
+ @Test
+ void applyAccumulatesRangeAndExact_keepsOnlyKeysInRange() {
+ SavepointFilterTranslator.Result applied =
+ apply(
+ List.of(
+ between(longKeyRef(), longLit(4L),
longLit(12L)),
+ or(
+ eq(longKeyRef(), longLit(1L)),
+ eq(longKeyRef(), longLit(5L)),
+ eq(longKeyRef(), longLit(10L)),
+ eq(longKeyRef(), longLit(15L)))),
+ LONG_KEY_TYPE);
+ SavepointKeyFilter result = applied.keyFilter();
+
+ assertNotNull(result);
+ assertThat(applied.accepted()).hasSize(2);
+ assertThat(applied.remaining()).isEmpty();
+ assertThat(result.getExactKeys()).containsExactlyInAnyOrder(5L, 10L);
+ }
+
+ //
-------------------------------------------------------------------------
+ // Empty key filter
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void emptyKeyFilter_rejectsEverything() {
+ SavepointKeyFilter empty = SavepointKeyFilter.empty();
+ assertThat(empty.isEmpty()).isTrue();
+ assertThat(empty.getExactKeys()).isEmpty();
+ assertThat(empty.test(42L)).isFalse();
+ assertThat(empty.test("hello")).isFalse();
+ }
+
+ @Test
+ void exactWithEmptySetReturnsEmptyKeyFilter() {
+ SavepointKeyFilter filter =
SavepointKeyFilter.exact(Collections.emptySet());
+ assertThat(filter.isEmpty()).isTrue();
+ }
+
+ @Test
+ void emptyKeyFilterSingletonPreservedAcrossSerialization() throws
Exception {
+ SavepointKeyFilter original = SavepointKeyFilter.empty();
+ ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ try (ObjectOutputStream oos = new ObjectOutputStream(baos)) {
+ oos.writeObject(original);
+ }
+ Object deserialized;
+ try (ObjectInputStream ois =
+ new ObjectInputStream(new
ByteArrayInputStream(baos.toByteArray()))) {
+ deserialized = ois.readObject();
+ }
+ assertThat(deserialized).isSameAs(SavepointKeyFilter.empty());
+ }
+
+ //
-------------------------------------------------------------------------
+ // Exact key filter — single-value factory
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void exactSingleValueFactory() {
+ SavepointKeyFilter filter = SavepointKeyFilter.exact(42L);
+ assertThat(filter.getExactKeys()).containsExactly(42L);
+ assertThat(filter.test(42L)).isTrue();
+ assertThat(filter.test(43L)).isFalse();
+ }
+
+ //
-------------------------------------------------------------------------
+ // AND with 3+ children
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void andOfThreeRangesProducesIntersection() {
+ // key >= 3 AND key <= 20 AND key < 10
+ CallExpression expr =
+ and(
+ gte(longKeyRef(), longLit(3L)),
+ lte(longKeyRef(), longLit(20L)),
+ lt(longKeyRef(), longLit(10L)));
+ SavepointKeyFilter filter = keyFilterOf(expr);
+
+ assertNotNull(filter);
+ assertThat(filter.getExactKeys()).isNull();
+ assertThat(filter.test(2L)).isFalse();
+ assertThat(filter.test(3L)).isTrue();
+ assertThat(filter.test(9L)).isTrue();
+ assertThat(filter.test(10L)).isFalse();
+ assertThat(filter.test(20L)).isFalse();
+ }
+
+ //
-------------------------------------------------------------------------
+ // OR edge cases
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void orWithSingleChild_returnsExactFilter() {
+ CallExpression expr = or(eq(longKeyRef(), longLit(7L)));
+ SavepointKeyFilter filter = keyFilterOf(expr);
+
+ assertNotNull(filter);
+ assertThat(filter.getExactKeys()).containsExactly(7L);
+ }
+
+ //
-------------------------------------------------------------------------
+ // Comparison with neither side being the key column
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void comparisonWithNeitherSideBeingKeyColumn_returnsNull() {
+ assertThat(keyFilterOf(gt(otherRef(), longLit(5L)))).isNull();
+ assertThat(keyFilterOf(lt(longLit(5L), otherRef()))).isNull();
+ }
+
+ //
-------------------------------------------------------------------------
+ // SavepointFilters.apply — empty filter handling
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void applyWithEmptyThenRange_returnsEmpty() {
+ SavepointFilterTranslator.Result applied =
+ apply(
+ List.of(
+ and(gt(longKeyRef(), longLit(10L)),
lt(longKeyRef(), longLit(5L))),
+ lte(longKeyRef(), longLit(10L))),
+ LONG_KEY_TYPE);
+ SavepointKeyFilter result = applied.keyFilter();
+
+ assertNotNull(result);
+ assertThat(applied.accepted()).hasSize(2);
+ assertThat(applied.remaining()).isEmpty();
+ assertThat(result.isEmpty()).isTrue();
+ }
+
+ @Test
+ void applyWithRangeThenEmpty_returnsEmpty() {
+ SavepointFilterTranslator.Result applied =
+ apply(
+ List.of(
+ lte(longKeyRef(), longLit(10L)),
+ and(gt(longKeyRef(), longLit(10L)),
lt(longKeyRef(), longLit(5L)))),
+ LONG_KEY_TYPE);
+ SavepointKeyFilter result = applied.keyFilter();
+
+ assertNotNull(result);
+ assertThat(applied.accepted()).hasSize(2);
+ assertThat(applied.remaining()).isEmpty();
+ assertThat(result.isEmpty()).isTrue();
+ }
+
+ @Test
+ void applyWithEmptyThenExact_returnsEmpty() {
+ SavepointFilterTranslator.Result applied =
+ apply(
+ List.of(
+ and(gt(longKeyRef(), longLit(10L)),
lt(longKeyRef(), longLit(5L))),
+ eq(longKeyRef(), longLit(1L))),
+ LONG_KEY_TYPE);
+ SavepointKeyFilter result = applied.keyFilter();
+
+ assertNotNull(result);
+ assertThat(applied.accepted()).hasSize(2);
+ assertThat(applied.remaining()).isEmpty();
+ assertThat(result.isEmpty()).isTrue();
+ }
+
+ @Test
+ void applyWithConflictingExactPredicates_returnsEmptyFilter() {
+ SavepointFilterTranslator.Result applied =
+ apply(
+ List.of(eq(longKeyRef(), longLit(1L)),
eq(longKeyRef(), longLit(2L))),
+ LONG_KEY_TYPE);
+ SavepointKeyFilter result = applied.keyFilter();
+
+ assertNotNull(result);
+ assertThat(applied.accepted()).hasSize(2);
+ assertThat(applied.remaining()).isEmpty();
+ assertThat(result.isEmpty()).isTrue();
+ }
+
+ //
-------------------------------------------------------------------------
+ // BETWEEN — wrong arg count
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void betweenWithWrongArgCount_returnsNull() {
+ CallExpression malformed =
+ CallExpression.permanent(
+ BuiltInFunctionDefinitions.BETWEEN,
+ Arrays.asList(longKeyRef(), longLit(1L)),
+ DataTypes.BOOLEAN());
+ assertThat(keyFilterOf(malformed)).isNull();
+ }
+
+ //
-------------------------------------------------------------------------
+ // BETWEEN — non-literal bounds
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void betweenWithNonLiteralBound_returnsNull() {
+ CallExpression expr =
+ CallExpression.permanent(
+ BuiltInFunctionDefinitions.BETWEEN,
+ Arrays.asList(longKeyRef(), otherRef(), longLit(10L)),
+ DataTypes.BOOLEAN());
+ assertThat(keyFilterOf(expr)).isNull();
+ }
+
+ //
-------------------------------------------------------------------------
+ // Comparison with non-literal value side
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void comparisonWithNonLiteralValue_returnsNull() {
+ assertThat(keyFilterOf(gt(longKeyRef(), otherRef()))).isNull();
+ }
+
+ //
-------------------------------------------------------------------------
+ // EQUALS — malformed (wrong arg count)
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void equalsWithWrongArgCount_returnsNull() {
+ CallExpression malformed =
+ CallExpression.permanent(
+ BuiltInFunctionDefinitions.EQUALS,
+ Collections.singletonList(longKeyRef()),
+ DataTypes.BOOLEAN());
+ assertThat(keyFilterOf(malformed)).isNull();
+ }
+
+ //
-------------------------------------------------------------------------
+ // Literal type widening to key type
+ //
-------------------------------------------------------------------------
+
+ @Test
+ void equalsWithIntLiteralIsWidenedToBigintKeyAndPushed() {
+ ValueLiteralExpression intLit = new ValueLiteralExpression(5,
DataTypes.INT().notNull());
+ SavepointKeyFilter filter = keyFilterOf(eq(longKeyRef(), intLit));
+
+ assertNotNull(filter);
+ assertThat(filter.getExactKeys()).containsExactly(5L);
+ assertThat(filter.test(5L)).isTrue();
+ assertThat(filter.test(6L)).isFalse();
+ }
+
+ @Test
+ void betweenWithIntLiteralBoundsIsWidenedToBigintKeyAndPushed() {
+ ValueLiteralExpression lower = new ValueLiteralExpression(1,
DataTypes.INT().notNull());
+ ValueLiteralExpression upper = new ValueLiteralExpression(10,
DataTypes.INT().notNull());
+ SavepointKeyFilter filter = keyFilterOf(between(longKeyRef(), lower,
upper));
+
+ assertNotNull(filter);
+ assertThat(filter.getExactKeys()).isNull();
+ assertThat(filter.test(0L)).isFalse();
+ assertThat(filter.test(1L)).isTrue();
+ assertThat(filter.test(10L)).isTrue();
+ assertThat(filter.test(11L)).isFalse();
+ }
+
+ @Test
+ void equalsWithIntLiteralIsWidenedToDoubleKeyAndPushed() {
+ FieldReferenceExpression doubleKey =
+ new FieldReferenceExpression("key",
DataTypes.DOUBLE().notNull(), 0, KEY_COL);
+ ValueLiteralExpression intLit = new ValueLiteralExpression(5,
DataTypes.INT().notNull());
+
+ SavepointKeyFilter filter = keyFilterOf(eq(doubleKey, intLit));
+
+ assertNotNull(filter);
+ assertThat(filter.getExactKeys()).containsExactly(5.0d);
+ assertThat(filter.test(5.0d)).isTrue();
+ assertThat(filter.test(6.0d)).isFalse();
+ }
+
+ @Test
+ void nonNumericLiteralAgainstNumericKeyIsNotWidenedAndNotPushed() {
+ assertThat(keyFilterOf(eq(longKeyRef(), stringLit("5")))).isNull();
+ }
+
+ @Test
+ void numericLiteralWithNonWidenableKeyTypeIsNotPushed() {
+ FieldReferenceExpression intKey =
+ new FieldReferenceExpression("key", DataTypes.INT().notNull(),
0, KEY_COL);
+ assertThat(keyFilterOf(eq(intKey, longLit(5L)))).isNull();
+ }
+
+ @Test
+ void decimalKeyEqualityIsPushedDownPreservingLiteralScale() {
+ FieldReferenceExpression decKey =
+ new FieldReferenceExpression("key", DataTypes.DECIMAL(10,
2).notNull(), 0, KEY_COL);
+ ValueLiteralExpression lit =
+ new ValueLiteralExpression(
+ new BigDecimal("5.00"), DataTypes.DECIMAL(10,
2).notNull());
+
+ SavepointKeyFilter filter = keyFilterOf(eq(decKey, lit));
+
+ assertNotNull(filter);
+ // Literal scale is preserved, so exact matching is scale sensitive:
5.00 matches, 5.0 not.
+ assertThat(filter.getExactKeys()).containsExactly(new
BigDecimal("5.00"));
+ assertThat(filter.test(new BigDecimal("5.00"))).isTrue();
+ assertThat(filter.test(new BigDecimal("5.0"))).isFalse();
+ }
+
+ //
-------------------------------------------------------------------------
+ // Expression helpers
+ //
-------------------------------------------------------------------------
+
+ private static SavepointKeyFilter keyFilterOf(ResolvedExpression expr) {
+ DataType keyType = findKeyType(expr);
+ return apply(Collections.singletonList(expr), keyType).keyFilter();
+ }
+
+ private static SavepointFilterTranslator.Result apply(
+ List<ResolvedExpression> filters, DataType keyType) {
+ return new SavepointFilterTranslator(KEY_COL, keyType).apply(filters);
+ }
+
+ private static DataType findKeyType(ResolvedExpression expr) {
+ if (expr instanceof FieldReferenceExpression
+ && ((FieldReferenceExpression) expr).getFieldIndex() ==
KEY_COL) {
+ return expr.getOutputDataType();
+ }
+ for (ResolvedExpression child : expr.getResolvedChildren()) {
+ DataType found = findKeyType(child);
+ if (found != null) {
+ return found;
+ }
+ }
+ return LONG_KEY_TYPE;
+ }
+
+ private static FieldReferenceExpression longKeyRef() {
+ return new FieldReferenceExpression("key",
DataTypes.BIGINT().notNull(), 0, KEY_COL);
+ }
+
+ private static FieldReferenceExpression stringKeyRef() {
+ return new FieldReferenceExpression("key",
DataTypes.STRING().notNull(), 0, KEY_COL);
+ }
+
+ private static FieldReferenceExpression otherRef() {
+ return new FieldReferenceExpression("val",
DataTypes.BIGINT().notNull(), 0, 1);
+ }
+
+ private static ValueLiteralExpression longLit(long value) {
+ return new ValueLiteralExpression(value, DataTypes.BIGINT().notNull());
+ }
+
+ private static ValueLiteralExpression stringLit(String value) {
+ return new ValueLiteralExpression(value, DataTypes.STRING().notNull());
+ }
+
+ private static CallExpression eq(ResolvedExpression left,
ResolvedExpression right) {
+ return CallExpression.permanent(
+ BuiltInFunctionDefinitions.EQUALS, Arrays.asList(left, right),
DataTypes.BOOLEAN());
+ }
+
+ private static CallExpression or(ResolvedExpression... args) {
+ return CallExpression.permanent(
+ BuiltInFunctionDefinitions.OR, Arrays.asList(args),
DataTypes.BOOLEAN());
+ }
+
+ private static CallExpression and(ResolvedExpression... args) {
+ return CallExpression.permanent(
+ BuiltInFunctionDefinitions.AND, Arrays.asList(args),
DataTypes.BOOLEAN());
+ }
+
+ private static CallExpression between(
+ ResolvedExpression value, ResolvedExpression lower,
ResolvedExpression upper) {
+ return CallExpression.permanent(
+ BuiltInFunctionDefinitions.BETWEEN,
+ Arrays.asList(value, lower, upper),
+ DataTypes.BOOLEAN());
+ }
+
+ private static CallExpression gt(ResolvedExpression left,
ResolvedExpression right) {
+ return CallExpression.permanent(
+ BuiltInFunctionDefinitions.GREATER_THAN,
+ Arrays.asList(left, right),
+ DataTypes.BOOLEAN());
+ }
+
+ private static CallExpression gte(ResolvedExpression left,
ResolvedExpression right) {
+ return CallExpression.permanent(
+ BuiltInFunctionDefinitions.GREATER_THAN_OR_EQUAL,
+ Arrays.asList(left, right),
+ DataTypes.BOOLEAN());
+ }
+
+ private static CallExpression lt(ResolvedExpression left,
ResolvedExpression right) {
+ return CallExpression.permanent(
+ BuiltInFunctionDefinitions.LESS_THAN,
+ Arrays.asList(left, right),
+ DataTypes.BOOLEAN());
+ }
+
+ private static CallExpression lte(ResolvedExpression left,
ResolvedExpression right) {
+ return CallExpression.permanent(
+ BuiltInFunctionDefinitions.LESS_THAN_OR_EQUAL,
+ Arrays.asList(left, right),
+ DataTypes.BOOLEAN());
+ }
+}
