diff --git a/docs/content/docs/libs/state_processor_api.md b/docs/content/docs/libs/state_processor_api.md index 0b0472b7f39e40..336502ab3f2173 100644 --- a/docs/content/docs/libs/state_processor_api.md +++ b/docs/content/docs/libs/state_processor_api.md @@ -749,17 +749,13 @@ The following predicates on the key column can be pushed down: | Option | Required | Default | Type | Description | |----------------------------------|----------|---------|----------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| | fields.#.state-name | optional | (none) | String | Overrides the state name which must be used for state reading. This can be useful when the state name contains characters which are not compliant with SQL column names. | -| fields.#.state-type | optional | (none) | Enum Possible values: list, map, value | Defines the state type which must be used for state reading, including value, list and map. When it's not provided then it tries to infer from the SQL type (ARRAY=list, MAP=map, all others=value). | -| fields.#.key-class | optional | (none) | String | Defines the format class scheme for decoding map key data (for ex. java.lang.Long). Either key-class or key-type-factory can be specified. When none of them are provided then the format class scheme tries to infer from the SQL type (only primitive types supported). | -| fields.#.key-type-factory | optional | (none) | String | Defines the type information factory for decoding map key data. Either key-class or key-type-factory can be specified. When none of them are provided then the format class scheme tries to infer from the SQL type (only primitive types supported). | -| fields.#.value-class | optional | (none) | String | Defines the format class scheme for decoding value data (for ex. java.lang.Long). Either value-class or value-info-factory can be specified. When none of them are provided then the format class scheme tries to infer from the SQL type (only primitive types supported). | -| fields.#.value-type-factory | optional | (none) | String | Defines the type information factory for decoding value data. Either value-class or value-type-factory can be specified. When none of them are provided then the format class scheme tries to infer from the SQL type (only primitive types supported). | ### Default Data Type Mapping -The state SQL connector infers the data type for primitive types when `fields.#.value-class` and `fields.#.key-class` -are not defined. The following table shows the `Flink SQL type` -> `Java type` default mapping. If the mapping is not calculated properly -then it can be overridden with the two mentioned config parameters on a per-column basis. +The state SQL connector automatically infers each column's data type from the savepoint's serializer +snapshot metadata (this covers primitive, Avro, Row and POJO types). When no serializer snapshot is +available for a column, it falls back to inferring a primitive Java type directly from the SQL type, +using the mapping below. | Flink SQL type | Java type | |-------------------------|-------------------------------------------------------------------------| diff --git a/flink-core/src/main/java/org/apache/flink/api/common/typeutils/CustomRestoreSerializerFactory.java b/flink-core/src/main/java/org/apache/flink/api/common/typeutils/CustomRestoreSerializerFactory.java new file mode 100644 index 00000000000000..3ad8c0870f0044 --- /dev/null +++ b/flink-core/src/main/java/org/apache/flink/api/common/typeutils/CustomRestoreSerializerFactory.java @@ -0,0 +1,118 @@ +/* + * 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.api.common.typeutils; + +import org.apache.flink.annotation.Internal; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.function.Function; + +/** + * Thread-scoped holder for a factory that builds a fallback {@link TypeSerializer} when a {@link + * TypeSerializerSnapshot} restores itself but a class it depends on is not on the classpath — for + * example a POJO's declared type, or an Avro record's specific/reflect runtime type. + * + *

This exists solely to support the Flink State Processing API's ability to read state whose + * original classes are not on the classpath (e.g. converting savepoint state into table rows + * without the user's job JAR). It must never be set by, or otherwise affect, regular job restores: + * a {@code TypeSerializerSnapshot} only ever consults this factory after it has already determined + * — independently of this class — that the class it needs is genuinely missing, and only when a + * factory has actually been registered. With no factory registered (the default for every job that + * is not using the State Processing API), behavior is unchanged from before this hook existed: the + * snapshot fails fast with a {@code ClassNotFoundException} or equivalent. + * + *

A {@code ThreadLocal} is used because {@link + * CompositeTypeSerializerSnapshot#restoreSerializer()} eagerly restores all of its nested + * serializers and offers no hook for substituting one of them, so a POJO or Avro type nested + * arbitrarily deep inside a composite snapshot (e.g. a list or map serializer snapshot) cannot be + * reached otherwise. + */ +@Internal +public final class CustomRestoreSerializerFactory { + + private static final Logger LOG = LoggerFactory.getLogger(CustomRestoreSerializerFactory.class); + + private static final ThreadLocal, TypeSerializer>> + FACTORY = new ThreadLocal<>(); + + private CustomRestoreSerializerFactory() {} + + /** Registers the fallback factory for the current thread. */ + public static void set(Function, TypeSerializer> factory) { + FACTORY.set(factory); + } + + /** Returns the fallback factory registered via {@link #set}, or {@code null} if none. */ + public static Function, TypeSerializer> get() { + return FACTORY.get(); + } + + /** Clears the fallback factory registered for the current thread. */ + public static void remove() { + FACTORY.remove(); + } + + /** + * Resolves {@code className} via {@code classLoader}, or returns {@code null} if it cannot be + * found and a fallback factory is registered for the current thread. + * + * @throws NoClassDefFoundError if the class cannot be found and no fallback factory is + * registered. + */ + @SuppressWarnings("unchecked") + public static Class resolveOrNull(String className, ClassLoader classLoader) { + try { + return (Class) Class.forName(className, false, classLoader); + } catch (ClassNotFoundException e) { + if (get() == null) { + throw missingClass(className, e); + } + LOG.debug( + "Class '{}' not found on classpath; a CustomRestoreSerializerFactory is" + + " registered to read the data without it.", + className); + return null; + } + } + + /** + * Builds the fallback serializer for a {@code snapshot} whose runtime class, {@code + * missingClassName}, could not be loaded, using the factory registered via {@link #set}. + * + * @throws NoClassDefFoundError if no factory is registered for the current thread. + */ + @SuppressWarnings("unchecked") + public static TypeSerializer restoreFallbackSerializer( + TypeSerializerSnapshot snapshot, String missingClassName) { + Function, TypeSerializer> fallback = get(); + if (fallback == null) { + throw missingClass(missingClassName, new ClassNotFoundException(missingClassName)); + } + return (TypeSerializer) fallback.apply(snapshot); + } + + private static NoClassDefFoundError missingClass( + String className, ClassNotFoundException cause) { + NoClassDefFoundError error = new NoClassDefFoundError(className); + error.initCause(cause); + return error; + } +} diff --git a/flink-core/src/main/java/org/apache/flink/api/common/typeutils/base/EnumSerializer.java b/flink-core/src/main/java/org/apache/flink/api/common/typeutils/base/EnumSerializer.java index ca7d1fc56326a9..ee5b5898ab2ca9 100644 --- a/flink-core/src/main/java/org/apache/flink/api/common/typeutils/base/EnumSerializer.java +++ b/flink-core/src/main/java/org/apache/flink/api/common/typeutils/base/EnumSerializer.java @@ -20,12 +20,12 @@ import org.apache.flink.annotation.Internal; import org.apache.flink.annotation.VisibleForTesting; +import org.apache.flink.api.common.typeutils.CustomRestoreSerializerFactory; import org.apache.flink.api.common.typeutils.TypeSerializer; import org.apache.flink.api.common.typeutils.TypeSerializerSchemaCompatibility; import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; import org.apache.flink.core.memory.DataInputView; import org.apache.flink.core.memory.DataOutputView; -import org.apache.flink.util.InstantiationUtil; import java.io.IOException; import java.io.ObjectInputStream; @@ -184,6 +184,8 @@ public static final class EnumSerializerSnapshot> private T[] enums; private Class enumClass; + private String enumClassName; + private String[] enumNames; @SuppressWarnings("unused") public EnumSerializerSnapshot() { @@ -213,16 +215,23 @@ public void writeSnapshot(DataOutputView out) throws IOException { @Override public void readSnapshot(int readVersion, DataInputView in, ClassLoader userCodeClassLoader) throws IOException { - enumClass = InstantiationUtil.resolveClassByName(in, userCodeClassLoader); + final String className = in.readUTF(); + enumClass = + CustomRestoreSerializerFactory.resolveOrNull(className, userCodeClassLoader); + enumClassName = className; int numEnumConstants = in.readInt(); - @SuppressWarnings("unchecked") - T[] previousEnums = (T[]) Array.newInstance(enumClass, numEnumConstants); + T[] previousEnums = + enumClass != null ? (T[]) Array.newInstance(enumClass, numEnumConstants) : null; + String[] names = new String[numEnumConstants]; for (int i = 0; i < numEnumConstants; i++) { - String enumName = in.readUTF(); + names[i] = in.readUTF(); + if (previousEnums == null) { + continue; + } try { - previousEnums[i] = Enum.valueOf(enumClass, enumName); + previousEnums[i] = Enum.valueOf(enumClass, names[i]); } catch (IllegalArgumentException e) { throw new IllegalStateException( "Could not create a restore serializer for enum " @@ -230,14 +239,25 @@ public void readSnapshot(int readVersion, DataInputView in, ClassLoader userCode + ". Probably because an enum value was removed."); } } + enumNames = names; + if (enumClass == null) { + return; + } this.enums = previousEnums; } + /** Returns the enum constant names in this snapshot, ordered by their wire ordinal. */ + public String[] getEnumNames() { + return enumNames; + } + @Override public TypeSerializer restoreSerializer() { - checkState(enumClass != null, "Enum class can not be null."); - + if (enumClass == null) { + return CustomRestoreSerializerFactory.restoreFallbackSerializer( + this, enumClassName); + } return new EnumSerializer<>(enumClass, enums); } diff --git a/flink-core/src/main/java/org/apache/flink/api/java/typeutils/runtime/PojoSerializer.java b/flink-core/src/main/java/org/apache/flink/api/java/typeutils/runtime/PojoSerializer.java index 9d7e4254e861eb..19e31c9f211d3c 100644 --- a/flink-core/src/main/java/org/apache/flink/api/java/typeutils/runtime/PojoSerializer.java +++ b/flink-core/src/main/java/org/apache/flink/api/java/typeutils/runtime/PojoSerializer.java @@ -48,10 +48,10 @@ public final class PojoSerializer extends TypeSerializer { // Flags for the header - private static final byte IS_NULL = 1; - private static final byte NO_SUBCLASS = 2; - private static final byte IS_SUBCLASS = 4; - private static final byte IS_TAGGED_SUBCLASS = 8; + public static final byte IS_NULL = 1; + public static final byte NO_SUBCLASS = 2; + public static final byte IS_SUBCLASS = 4; + public static final byte IS_TAGGED_SUBCLASS = 8; private static final long serialVersionUID = 1L; diff --git a/flink-core/src/main/java/org/apache/flink/api/java/typeutils/runtime/PojoSerializerSnapshot.java b/flink-core/src/main/java/org/apache/flink/api/java/typeutils/runtime/PojoSerializerSnapshot.java index cb8deb5b9e7cc8..cdfe1b9b12ebae 100644 --- a/flink-core/src/main/java/org/apache/flink/api/java/typeutils/runtime/PojoSerializerSnapshot.java +++ b/flink-core/src/main/java/org/apache/flink/api/java/typeutils/runtime/PojoSerializerSnapshot.java @@ -23,6 +23,7 @@ import org.apache.flink.api.common.serialization.SerializerConfigImpl; import org.apache.flink.api.common.typeutils.CompositeTypeSerializerUtil; import org.apache.flink.api.common.typeutils.CompositeTypeSerializerUtil.IntermediateCompatibilityResult; +import org.apache.flink.api.common.typeutils.CustomRestoreSerializerFactory; import org.apache.flink.api.common.typeutils.TypeSerializer; import org.apache.flink.api.common.typeutils.TypeSerializerSchemaCompatibility; import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; @@ -34,10 +35,12 @@ import java.io.IOException; import java.lang.reflect.Field; +import java.util.AbstractMap.SimpleEntry; import java.util.ArrayList; import java.util.Arrays; import java.util.Iterator; import java.util.LinkedHashMap; +import java.util.List; import java.util.Map; import java.util.Set; @@ -143,6 +146,11 @@ public void readSnapshot(int readVersion, DataInputView in, ClassLoader userCode @Override @SuppressWarnings("unchecked") public TypeSerializer restoreSerializer() { + if (snapshotData.getPojoClass() == null) { + return CustomRestoreSerializerFactory.restoreFallbackSerializer( + this, snapshotData.getPojoClassName()); + } + final int numFields = snapshotData.getFieldSerializerSnapshots().size(); final ArrayList restoredFields = new ArrayList<>(numFields); @@ -257,6 +265,56 @@ public TypeSerializerSchemaCompatibility resolveSchemaCompatibility( return TypeSerializerSchemaCompatibility.compatibleAsIs(); } + // --------------------------------------------------------------------------------------------- + // Schema extraction support + // --------------------------------------------------------------------------------------------- + + /** + * Returns {@code true} if the POJO class could be loaded from the classloader that read this + * snapshot. When {@code false}, {@link #restoreSerializer()} delegates to the serializer + * supplied via {@link CustomRestoreSerializerFactory} instead of building a {@link + * PojoSerializer}. + */ + @Internal + public boolean isPojoClassAvailable() { + return snapshotData.getPojoClass() != null; + } + + /** + * Returns the POJO class name as stored in the snapshot. Available even when the class cannot + * be loaded. + */ + @Internal + public String getPojoClassName() { + return snapshotData.getPojoClassName(); + } + + /** + * Returns an ordered list of (field name, field serializer snapshot) pairs. Field names are + * always present; snapshot values may be {@code null} when the field snapshot could not be + * read. + */ + @Internal + public List>> getFieldSnapshotEntries() { + List>> result = new ArrayList<>(); + snapshotData + .getFieldSerializerSnapshots() + .forEach( + (fieldName, field, fieldSnapshot) -> + result.add(new SimpleEntry<>(fieldName, fieldSnapshot))); + return result; + } + + /** + * Returns the registered subclass serializer snapshots in tag order (tag 0, 1, 2, …). Values + * may be {@code null} if a subclass snapshot was not readable. + */ + @Internal + public List> getRegisteredSubclassSnapshotsOrdered() { + return new ArrayList<>( + snapshotData.getRegisteredSubclassSerializerSnapshots().unwrapOptionals().values()); + } + // --------------------------------------------------------------------------------------------- // Utility methods // --------------------------------------------------------------------------------------------- diff --git a/flink-core/src/main/java/org/apache/flink/api/java/typeutils/runtime/PojoSerializerSnapshotData.java b/flink-core/src/main/java/org/apache/flink/api/java/typeutils/runtime/PojoSerializerSnapshotData.java index 6b2bd112ad14f7..ceaaa3158e4693 100644 --- a/flink-core/src/main/java/org/apache/flink/api/java/typeutils/runtime/PojoSerializerSnapshotData.java +++ b/flink-core/src/main/java/org/apache/flink/api/java/typeutils/runtime/PojoSerializerSnapshotData.java @@ -19,12 +19,12 @@ package org.apache.flink.api.java.typeutils.runtime; import org.apache.flink.annotation.Internal; +import org.apache.flink.api.common.typeutils.CustomRestoreSerializerFactory; import org.apache.flink.api.common.typeutils.TypeSerializer; import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; import org.apache.flink.core.memory.DataInputView; import org.apache.flink.core.memory.DataOutputView; import org.apache.flink.util.CollectionUtil; -import org.apache.flink.util.InstantiationUtil; import org.apache.flink.util.LinkedOptionalMap; import org.apache.flink.util.function.BiConsumerWithException; import org.apache.flink.util.function.BiFunctionWithException; @@ -32,6 +32,8 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import javax.annotation.Nullable; + import java.io.IOException; import java.lang.reflect.Field; import java.util.LinkedHashMap; @@ -113,6 +115,7 @@ static PojoSerializerSnapshotData createFrom( return new PojoSerializerSnapshotData<>( pojoClass, + pojoClass.getName(), fieldSerializerSnapshots, optionalMapOf(registeredSubclassSerializerSnapshots, Class::getName), optionalMapOf(nonRegisteredSubclassSerializerSnapshots, Class::getName)); @@ -153,12 +156,14 @@ static PojoSerializerSnapshotData createFrom( return new PojoSerializerSnapshotData<>( pojoClass, + pojoClass.getName(), fieldSerializerSnapshots, optionalMapOf(existingRegisteredSubclassSerializerSnapshots, Class::getName), optionalMapOf(existingNonRegisteredSubclassSerializerSnapshots, Class::getName)); } - private Class pojoClass; + @Nullable private Class pojoClass; + private String pojoClassName; private LinkedOptionalMap> fieldSerializerSnapshots; private LinkedOptionalMap, TypeSerializerSnapshot> registeredSubclassSerializerSnapshots; @@ -166,14 +171,16 @@ static PojoSerializerSnapshotData createFrom( nonRegisteredSubclassSerializerSnapshots; private PojoSerializerSnapshotData( - Class typeClass, + @Nullable Class typeClass, + String pojoClassName, LinkedOptionalMap> fieldSerializerSnapshots, LinkedOptionalMap, TypeSerializerSnapshot> registeredSubclassSerializerSnapshots, LinkedOptionalMap, TypeSerializerSnapshot> nonRegisteredSubclassSerializerSnapshots) { - this.pojoClass = checkNotNull(typeClass); + this.pojoClass = typeClass; + this.pojoClassName = checkNotNull(pojoClassName); this.fieldSerializerSnapshots = checkNotNull(fieldSerializerSnapshots); this.registeredSubclassSerializerSnapshots = checkNotNull(registeredSubclassSerializerSnapshots); @@ -186,7 +193,7 @@ private PojoSerializerSnapshotData( // --------------------------------------------------------------------------------------------- void writeSnapshotData(DataOutputView out) throws IOException { - out.writeUTF(pojoClass.getName()); + out.writeUTF(pojoClassName); writeOptionalMap( out, fieldSerializerSnapshots, @@ -206,7 +213,14 @@ void writeSnapshotData(DataOutputView out) throws IOException { private static PojoSerializerSnapshotData readSnapshotData( DataInputView in, ClassLoader userCodeClassLoader) throws IOException { - Class pojoClass = InstantiationUtil.resolveClassByName(in, userCodeClassLoader); + final String pojoClassName = in.readUTF(); + Class pojoClass = + CustomRestoreSerializerFactory.resolveOrNull(pojoClassName, userCodeClassLoader); + if (pojoClass == null) { + LOG.debug( + "POJO class '{}' not found on classpath; schema can still be read from field snapshots.", + pojoClassName); + } LinkedOptionalMap> fieldSerializerSnapshots = readOptionalMap( @@ -226,6 +240,7 @@ private static PojoSerializerSnapshotData readSnapshotData( return new PojoSerializerSnapshotData<>( pojoClass, + pojoClassName, fieldSerializerSnapshots, registeredSubclassSerializerSnapshots, nonRegisteredSubclassSerializerSnapshots); @@ -235,10 +250,15 @@ private static PojoSerializerSnapshotData readSnapshotData( // Snapshot data accessors // --------------------------------------------------------------------------------------------- + @Nullable Class getPojoClass() { return pojoClass; } + String getPojoClassName() { + return pojoClassName; + } + LinkedOptionalMap> getFieldSerializerSnapshots() { return fieldSerializerSnapshots; } diff --git a/flink-core/src/test/java/org/apache/flink/api/common/typeutils/base/EnumSerializerSnapshotMissingClassTest.java b/flink-core/src/test/java/org/apache/flink/api/common/typeutils/base/EnumSerializerSnapshotMissingClassTest.java new file mode 100644 index 00000000000000..a485c7cf85b354 --- /dev/null +++ b/flink-core/src/test/java/org/apache/flink/api/common/typeutils/base/EnumSerializerSnapshotMissingClassTest.java @@ -0,0 +1,104 @@ +/* + * 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.api.common.typeutils.base; + +import org.apache.flink.api.common.typeutils.CustomRestoreSerializerFactory; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; +import org.apache.flink.api.common.typeutils.base.EnumSerializer.EnumSerializerSnapshot; +import org.apache.flink.core.memory.DataInputDeserializer; +import org.apache.flink.core.memory.DataOutputSerializer; + +import org.junit.jupiter.api.Test; + +import java.io.IOException; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Tests for the lenient enum class-loading in {@link EnumSerializerSnapshot}: with a {@link + * CustomRestoreSerializerFactory} registered, an {@link EnumSerializerSnapshot} must be readable + * even when the enum class is not on the classpath, with the constant names (in wire-ordinal order) + * remaining accessible. + */ +class EnumSerializerSnapshotMissingClassTest { + + enum SomeEnum { + FOO, + BAR, + BAZ + } + + @Test + void testReadWithClassAbsent() throws IOException { + EnumSerializerSnapshot read; + CustomRestoreSerializerFactory.set( + snapshot -> { + throw new UnsupportedOperationException("not exercised in this test"); + }); + try { + read = roundtripSnapshot(writeSnapshot(), withoutEnumClassLoader()); + } finally { + CustomRestoreSerializerFactory.remove(); + } + + assertThat(read.getEnumNames()).containsExactly("FOO", "BAR", "BAZ"); + } + + /** + * Regular job restores (i.e. without a {@link CustomRestoreSerializerFactory} registered, as is + * always the case outside of the State Processing API) must still fail fast when the enum class + * is genuinely missing. + */ + @Test + void testReadFailsFastWithoutFallbackFactory() { + assertThatThrownBy(() -> roundtripSnapshot(writeSnapshot(), withoutEnumClassLoader())) + .isInstanceOf(NoClassDefFoundError.class) + .hasCauseInstanceOf(ClassNotFoundException.class); + } + + /** Hides the enum class from the classloader used to read the snapshot back. */ + private ClassLoader withoutEnumClassLoader() { + return new ClassLoader(getClass().getClassLoader()) { + @Override + protected Class loadClass(String name, boolean resolve) + throws ClassNotFoundException { + if (name.contains(SomeEnum.class.getSimpleName())) { + throw new ClassNotFoundException(name); + } + return super.loadClass(name, resolve); + } + }; + } + + private static EnumSerializerSnapshot writeSnapshot() { + return new EnumSerializer<>(SomeEnum.class).snapshotConfiguration(); + } + + @SuppressWarnings("unchecked") + private static EnumSerializerSnapshot roundtripSnapshot( + EnumSerializerSnapshot snapshot, ClassLoader classLoader) throws IOException { + DataOutputSerializer out = new DataOutputSerializer(256); + TypeSerializerSnapshot.writeVersionedSnapshot(out, snapshot); + + DataInputDeserializer in = new DataInputDeserializer(out.getSharedBuffer()); + return (EnumSerializerSnapshot) + TypeSerializerSnapshot.readVersionedSnapshot(in, classLoader); + } +} diff --git a/flink-core/src/test/java/org/apache/flink/api/java/typeutils/runtime/PojoSerializerSnapshotLenientReadTest.java b/flink-core/src/test/java/org/apache/flink/api/java/typeutils/runtime/PojoSerializerSnapshotLenientReadTest.java new file mode 100644 index 00000000000000..6f5fedb5bae511 --- /dev/null +++ b/flink-core/src/test/java/org/apache/flink/api/java/typeutils/runtime/PojoSerializerSnapshotLenientReadTest.java @@ -0,0 +1,173 @@ +/* + * 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.api.java.typeutils.runtime; + +import org.apache.flink.api.common.serialization.SerializerConfigImpl; +import org.apache.flink.api.common.typeutils.CustomRestoreSerializerFactory; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; +import org.apache.flink.api.common.typeutils.base.IntSerializer; +import org.apache.flink.api.common.typeutils.base.LongSerializer; +import org.apache.flink.api.common.typeutils.base.StringSerializer; +import org.apache.flink.api.java.typeutils.TypeExtractor; +import org.apache.flink.core.memory.DataInputDeserializer; +import org.apache.flink.core.memory.DataOutputSerializer; + +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.AbstractMap.SimpleEntry; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Tests for the lenient POJO class-loading in {@link PojoSerializerSnapshotData}: with a {@link + * CustomRestoreSerializerFactory} registered, a {@link PojoSerializerSnapshot} must be readable + * even when the POJO class is not on the classpath, with the class name, the field names, and the + * field serializer snapshots all remaining accessible. + */ +class PojoSerializerSnapshotLenientReadTest { + + private static final Map> EXPECTED_FIELD_SNAPSHOTS = + Map.of( + "name", StringSerializer.StringSerializerSnapshot.class, + "age", IntSerializer.IntSerializerSnapshot.class, + "score", LongSerializer.LongSerializerSnapshot.class); + + /** POJO present while the snapshot is written, hidden from the classloader while it is read. */ + public static class SomePojo { + public String name; + public int age; + public long score; + } + + @Test + void testReadWithClassPresent() throws IOException { + PojoSerializerSnapshot read = + roundtripSnapshot(writeSnapshot(), getClass().getClassLoader()); + + assertThat(read.isPojoClassAvailable()).isTrue(); + assertThat(read.getPojoClassName()).isEqualTo(SomePojo.class.getName()); + assertFieldSnapshots(read); + } + + @Test + void testReadWithClassAbsent() throws IOException { + PojoSerializerSnapshot read; + CustomRestoreSerializerFactory.set( + snapshot -> { + throw new UnsupportedOperationException("not exercised in this test"); + }); + try { + read = roundtripSnapshot(writeSnapshot(), withoutPojoClassLoader()); + } finally { + CustomRestoreSerializerFactory.remove(); + } + + assertThat(read.isPojoClassAvailable()).isFalse(); + assertThat(read.getPojoClassName()).isEqualTo(SomePojo.class.getName()); + // Field names stay available because the key name is written before the framed value. + assertFieldSnapshots(read); + } + + /** + * Regular job restores (i.e. without a {@link CustomRestoreSerializerFactory} registered, as is + * always the case outside of the State Processing API) must still fail fast when the POJO class + * is genuinely missing, exactly as before lenient reading was introduced. + */ + @Test + void testReadFailsFastWithoutFallbackFactory() { + assertThatThrownBy(() -> roundtripSnapshot(writeSnapshot(), withoutPojoClassLoader())) + .isInstanceOf(NoClassDefFoundError.class) + .hasCauseInstanceOf(ClassNotFoundException.class); + } + + /** + * A {@link CustomRestoreSerializerFactory} only needs to be registered for the duration it is + * actually relied on: reading here succeeds with one present, but restoring a working + * serializer must still fail once it has been removed again. + */ + @Test + void testRestoreSerializerFailsWithoutFallbackFactory() throws IOException { + PojoSerializerSnapshot read; + CustomRestoreSerializerFactory.set( + snapshot -> { + throw new UnsupportedOperationException("not exercised in this test"); + }); + try { + read = roundtripSnapshot(writeSnapshot(), withoutPojoClassLoader()); + } finally { + CustomRestoreSerializerFactory.remove(); + } + + assertThat(read.isPojoClassAvailable()).isFalse(); + assertThatThrownBy(read::restoreSerializer) + .isInstanceOf(NoClassDefFoundError.class) + .hasCauseInstanceOf(ClassNotFoundException.class); + } + + /** Hides both the POJO class itself and the declaring class of each of its fields. */ + private ClassLoader withoutPojoClassLoader() { + return new ClassLoader(getClass().getClassLoader()) { + @Override + protected Class loadClass(String name, boolean resolve) + throws ClassNotFoundException { + if (name.contains(SomePojo.class.getSimpleName())) { + throw new ClassNotFoundException(name); + } + return super.loadClass(name, resolve); + } + }; + } + + private static void assertFieldSnapshots(PojoSerializerSnapshot snapshot) { + List>> entries = + snapshot.getFieldSnapshotEntries(); + + assertThat(entries).hasSize(EXPECTED_FIELD_SNAPSHOTS.size()); + for (SimpleEntry> entry : entries) { + Class expectedType = EXPECTED_FIELD_SNAPSHOTS.get(entry.getKey()); + assertThat(expectedType).as("unexpected field '%s'", entry.getKey()).isNotNull(); + assertThat(entry.getValue()) + .as("snapshot of field '%s'", entry.getKey()) + .isExactlyInstanceOf(expectedType); + } + } + + @SuppressWarnings("unchecked") + private static PojoSerializerSnapshot writeSnapshot() { + return (PojoSerializerSnapshot) + TypeExtractor.createTypeInfo(SomePojo.class) + .createSerializer(new SerializerConfigImpl()) + .snapshotConfiguration(); + } + + @SuppressWarnings("unchecked") + private static PojoSerializerSnapshot roundtripSnapshot( + PojoSerializerSnapshot snapshot, ClassLoader classLoader) throws IOException { + DataOutputSerializer out = new DataOutputSerializer(256); + TypeSerializerSnapshot.writeVersionedSnapshot(out, snapshot); + + DataInputDeserializer in = new DataInputDeserializer(out.getSharedBuffer()); + return (PojoSerializerSnapshot) + TypeSerializerSnapshot.readVersionedSnapshot(in, classLoader); + } +} diff --git a/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/typeutils/AvroSerializerSnapshot.java b/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/typeutils/AvroSerializerSnapshot.java index 92580e71b2906d..4a097ed7e5d06e 100644 --- a/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/typeutils/AvroSerializerSnapshot.java +++ b/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/typeutils/AvroSerializerSnapshot.java @@ -19,6 +19,7 @@ package org.apache.flink.formats.avro.typeutils; import org.apache.flink.annotation.VisibleForTesting; +import org.apache.flink.api.common.typeutils.CustomRestoreSerializerFactory; import org.apache.flink.api.common.typeutils.TypeSerializer; import org.apache.flink.api.common.typeutils.TypeSerializerSchemaCompatibility; import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; @@ -33,7 +34,7 @@ import org.apache.avro.specific.SpecificData; import org.apache.avro.specific.SpecificRecord; -import javax.annotation.Nonnull; +import javax.annotation.Nullable; import java.io.IOException; import java.util.Objects; @@ -49,7 +50,8 @@ * @param The data type that the originating serializer of this configuration serializes. */ public class AvroSerializerSnapshot implements TypeSerializerSnapshot { - private Class runtimeType; + @Nullable private Class runtimeType; + private String runtimeTypeName; private Schema schema; private Schema runtimeSchema; @@ -61,6 +63,7 @@ public AvroSerializerSnapshot() { AvroSerializerSnapshot(Schema schema, Class runtimeType) { this.schema = schema; this.runtimeType = runtimeType; + this.runtimeTypeName = runtimeType.getName(); } @Override @@ -106,7 +109,11 @@ public void readSnapshot(int readVersion, DataInputView in, ClassLoader userCode private void readV1(DataInputView in, ClassLoader userCodeClassLoader) throws IOException { final String previousSchemaDefinition = in.readUTF(); this.schema = parseAvroSchema(previousSchemaDefinition); - this.runtimeType = findClassOrFallbackToGeneric(userCodeClassLoader, schema.getFullName()); + this.runtimeTypeName = schema.getFullName(); + // V1 snapshots predate CustomRestoreSerializerFactory support: preserve their original + // behavior of falling back to GenericRecord rather than failing when the runtime type is + // missing from the classpath. + this.runtimeType = findClassOrFallbackToGeneric(userCodeClassLoader, runtimeTypeName); this.runtimeSchema = tryExtractAvroSchema(userCodeClassLoader, runtimeType); } @@ -114,18 +121,22 @@ private void readV2(DataInputView in, ClassLoader userCodeClassLoader) throws IO final String previousRuntimeTypeName = in.readUTF(); final String previousSchemaDefinition = in.readUTF(); - this.runtimeType = findClassOrThrow(userCodeClassLoader, previousRuntimeTypeName); + this.runtimeTypeName = previousRuntimeTypeName; + this.runtimeType = tryFindClass(userCodeClassLoader, runtimeTypeName); this.schema = parseAvroSchema(previousSchemaDefinition); - this.runtimeSchema = tryExtractAvroSchema(userCodeClassLoader, runtimeType); + this.runtimeSchema = + runtimeType == null ? null : tryExtractAvroSchema(userCodeClassLoader, runtimeType); } private void readV3(DataInputView in, ClassLoader userCodeClassLoader) throws IOException { final String previousRuntimeTypeName = readString(in); final String previousSchemaDefinition = readString(in); - this.runtimeType = findClassOrThrow(userCodeClassLoader, previousRuntimeTypeName); + this.runtimeTypeName = previousRuntimeTypeName; + this.runtimeType = tryFindClass(userCodeClassLoader, runtimeTypeName); this.schema = parseAvroSchema(previousSchemaDefinition); - this.runtimeSchema = tryExtractAvroSchema(userCodeClassLoader, runtimeType); + this.runtimeSchema = + runtimeType == null ? null : tryExtractAvroSchema(userCodeClassLoader, runtimeType); } @Override @@ -141,9 +152,12 @@ public TypeSerializerSchemaCompatibility resolveSchemaCompatibility( @Override public TypeSerializer restoreSerializer() { - checkNotNull(runtimeType); checkNotNull(schema); + if (runtimeType == null) { + return CustomRestoreSerializerFactory.restoreFallbackSerializer(this, runtimeTypeName); + } + if (runtimeSchema != null) { return new AvroSerializer<>( runtimeType, @@ -157,6 +171,11 @@ public TypeSerializer restoreSerializer() { } } + /** Returns the Avro writer schema stored in this snapshot. */ + public Schema getSchema() { + return schema; + } + // ------------------------------------------------------------------------------------------------------------ // Helpers // ------------------------------------------------------------------------------------------------------------ @@ -225,32 +244,16 @@ private static Schema tryExtractAvroSchema(ClassLoader cl, Class runtimeType) return d.getSchema(runtimeType); } - @SuppressWarnings("unchecked") - @Nonnull - private static Class findClassOrThrow( - ClassLoader userCodeClassLoader, String className) { - try { - Class runtimeTarget = Class.forName(className, false, userCodeClassLoader); - return (Class) runtimeTarget; - } catch (ClassNotFoundException e) { - throw new IllegalStateException( - "" - + "Unable to find the class '" - + className - + "' which is used to deserialize " - + "the elements of this serializer. " - + "Were the class was moved or renamed?", - e); - } + @Nullable + private static Class tryFindClass(ClassLoader userCodeClassLoader, String className) { + return CustomRestoreSerializerFactory.resolveOrNull(className, userCodeClassLoader); } @SuppressWarnings("unchecked") - @Nonnull private static Class findClassOrFallbackToGeneric( ClassLoader userCodeClassLoader, String className) { try { - Class runtimeTarget = Class.forName(className, false, userCodeClassLoader); - return (Class) runtimeTarget; + return (Class) Class.forName(className, false, userCodeClassLoader); } catch (ClassNotFoundException e) { return (Class) GenericRecord.class; } diff --git a/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/typeutils/AvroSerializerSnapshotTest.java b/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/typeutils/AvroSerializerSnapshotTest.java index c0838e9b842f9c..ba7b0231d806d1 100644 --- a/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/typeutils/AvroSerializerSnapshotTest.java +++ b/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/typeutils/AvroSerializerSnapshotTest.java @@ -18,6 +18,7 @@ package org.apache.flink.formats.avro.typeutils; +import org.apache.flink.api.common.typeutils.CustomRestoreSerializerFactory; import org.apache.flink.api.common.typeutils.TypeSerializer; import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; import org.apache.flink.api.common.typeutils.TypeSerializerSnapshotSerializationUtil; @@ -45,6 +46,7 @@ import static org.apache.flink.api.common.typeutils.TypeSerializerConditions.isCompatibleAsIs; import static org.apache.flink.api.common.typeutils.TypeSerializerConditions.isIncompatible; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Test {@link AvroSerializerSnapshot}. */ class AvroSerializerSnapshotTest { @@ -269,6 +271,75 @@ void restorePastSnapshots() throws IOException { } } + /** + * V1 snapshots predate {@code CustomRestoreSerializerFactory} support and must keep their + * original behavior of falling back to {@link GenericRecord} rather than failing when the + * runtime type is missing from the classpath (regression test for a fallback that was + * accidentally dropped while adding lenient reading for the State Processing API). + */ + @Test + void restoringV1SnapshotWithMissingRuntimeTypeFallsBackToGenericRecord() throws IOException { + DataOutputSerializer out = new DataOutputSerializer(256); + out.writeUTF(Address.getClassSchema().toString(false)); + + AvroSerializerSnapshot

restored = new AvroSerializerSnapshot<>(); + ClassLoader withoutAddress = classLoaderHiding(Address.class.getSimpleName()); + restored.readSnapshot(1, new DataInputDeserializer(out.getCopyOfBuffer()), withoutAddress); + + @SuppressWarnings("unchecked") + AvroSerializer
serializer = (AvroSerializer
) restored.restoreSerializer(); + assertThat(serializer.getType()).isEqualTo(GenericRecord.class); + } + + /** + * Regular job restores (i.e. without a {@link CustomRestoreSerializerFactory} registered, as is + * always the case outside of the State Processing API) must still fail fast when a V2/V3 + * snapshot's runtime type is genuinely missing. + */ + @Test + void restoringV3SnapshotWithMissingRuntimeTypeFailsFastWithoutFallbackFactory() + throws IOException { + DataOutputSerializer out = new DataOutputSerializer(256); + new AvroSerializerSnapshot<>(Address.getClassSchema(), Address.class).writeSnapshot(out); + ClassLoader withoutAddress = classLoaderHiding(Address.class.getSimpleName()); + + assertThatThrownBy( + () -> + new AvroSerializerSnapshot
() + .readSnapshot( + 3, + new DataInputDeserializer(out.getCopyOfBuffer()), + withoutAddress)) + .isInstanceOf(NoClassDefFoundError.class) + .hasCauseInstanceOf(ClassNotFoundException.class); + } + + /** + * With a {@link CustomRestoreSerializerFactory} registered, a V2/V3 snapshot must be readable + * even when the runtime type is missing from the classpath. + */ + @Test + void restoringV3SnapshotWithMissingRuntimeTypeReadsLenientlyWithFallbackFactory() + throws IOException { + DataOutputSerializer out = new DataOutputSerializer(256); + new AvroSerializerSnapshot<>(Address.getClassSchema(), Address.class).writeSnapshot(out); + ClassLoader withoutAddress = classLoaderHiding(Address.class.getSimpleName()); + + AvroSerializerSnapshot
restored = new AvroSerializerSnapshot<>(); + CustomRestoreSerializerFactory.set( + snapshot -> { + throw new UnsupportedOperationException("not exercised in this test"); + }); + try { + restored.readSnapshot( + 3, new DataInputDeserializer(out.getCopyOfBuffer()), withoutAddress); + } finally { + CustomRestoreSerializerFactory.remove(); + } + + assertThat(restored.getSchema()).isEqualTo(Address.getClassSchema()); + } + /** * Creates a new serializer snapshot for the current version. Use this before bumping the * snapshot version and also add the version (before bumping) to {@link #PAST_VERSIONS}. @@ -332,6 +403,20 @@ private static T deserialize(TypeSerializer serializer, ByteBuffer serial return serializer.deserialize(in); } + /** Returns a class loader that fails to find any class whose simple name is {@code hidden}. */ + private static ClassLoader classLoaderHiding(String hidden) { + return new ClassLoader(AvroSerializerSnapshotTest.class.getClassLoader()) { + @Override + protected Class loadClass(String name, boolean resolve) + throws ClassNotFoundException { + if (name.contains(hidden)) { + throw new ClassNotFoundException(name); + } + return super.loadClass(name, resolve); + } + }; + } + // --------------------------------------------------------------------------------------------------------------- // Test classes // --------------------------------------------------------------------------------------------------------------- diff --git a/flink-libraries/flink-state-processing-api/pom.xml b/flink-libraries/flink-state-processing-api/pom.xml index eb027221a0c4f1..d8cc8d1654d0d9 100644 --- a/flink-libraries/flink-state-processing-api/pom.xml +++ b/flink-libraries/flink-state-processing-api/pom.xml @@ -68,6 +68,13 @@ under the License. provided + + org.apache.flink + flink-table-type-utils + ${project.version} + provided + + @@ -89,7 +96,7 @@ under the License. org.apache.flink flink-avro ${project.version} - test + true diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/StateTableUtils.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/StateTableUtils.java new file mode 100644 index 00000000000000..b6a4eacba30b34 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/StateTableUtils.java @@ -0,0 +1,641 @@ +/* + * 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; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.api.common.state.StateDescriptor; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; +import org.apache.flink.runtime.checkpoint.OperatorState; +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState; +import org.apache.flink.runtime.checkpoint.metadata.CheckpointMetadata; +import org.apache.flink.runtime.state.IncrementalKeyedStateHandle; +import org.apache.flink.runtime.state.KeyGroupsSavepointStateHandle; +import org.apache.flink.runtime.state.KeyGroupsStateHandle; +import org.apache.flink.runtime.state.KeyedStateHandle; +import org.apache.flink.runtime.state.StateBackendLoader; +import org.apache.flink.runtime.state.VoidNamespaceSerializer; +import org.apache.flink.runtime.state.changelog.ChangelogStateBackendHandle; +import org.apache.flink.state.api.schema.KeyedStateSchemaInfo; +import org.apache.flink.state.api.schema.SerializerSnapshotToLogicalTypeConverter; +import org.apache.flink.state.api.schema.StateSchemaExtractor; +import org.apache.flink.state.api.schema.StateSchemaInfo; +import org.apache.flink.state.table.SavepointConnectorOptions; +import org.apache.flink.streaming.api.operators.InternalTimeServiceManagerImpl; +import org.apache.flink.streaming.runtime.operators.windowing.WindowOperator; +import org.apache.flink.table.api.Schema; +import org.apache.flink.table.catalog.CatalogTable; +import org.apache.flink.table.factories.FactoryUtil; +import org.apache.flink.table.types.DataType; +import org.apache.flink.table.types.logical.ArrayType; +import org.apache.flink.table.types.logical.BigIntType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.MapType; +import org.apache.flink.table.types.logical.VarBinaryType; +import org.apache.flink.table.types.utils.LogicalTypeDataTypeConverter; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * High-level utility for inspecting and reading keyed state from a checkpoint / savepoint without + * requiring user POJO classes on the classpath. + */ +@Internal +public final class StateTableUtils { + + private static final Logger LOG = LoggerFactory.getLogger(StateTableUtils.class); + + private StateTableUtils() {} + + /** + * Returns the {@link OperatorIdentifier}s of all operators present in the given checkpoint + * metadata that have at least one non-internal keyed state. + * + * @param metadata the checkpoint metadata to inspect + * @return list of operator identifiers; never null, may be empty + */ + public static List getOperatorIdentifiers(CheckpointMetadata metadata) { + return metadata.getOperatorStates().stream() + .filter(StateTableUtils::hasNonInternalKeyedState) + .map( + op -> + op.getOperatorUid() + .map(OperatorIdentifier::forUid) + .orElseGet( + () -> + OperatorIdentifier.forUidHash( + op.getOperatorID().toHexString()))) + .collect(Collectors.toList()); + } + + private static boolean hasNonInternalKeyedState(OperatorState op) { + try { + List schemas = StateSchemaExtractor.extractSchema(op); + ClassifiedStates classified = classifyStates(op.getOperatorID().toHexString(), schemas); + return !classified.voidNamespaceStates.isEmpty() + || !classified.windowNamespaceStates.isEmpty(); + } catch (Exception e) { + LOG.error( + "Could not extract state schema for operator '{}': {}. Excluding from catalog.", + op.getOperatorID(), + e.getMessage()); + return false; + } + } + + /** + * Returns the names of all keyed states registered by the given operator. + * + * @param metadata the checkpoint metadata to inspect + * @param operatorId identifies the operator + * @param classLoader the class loader used when reading serializer snapshots + * @return list of state names; never null, may be empty + * @throws IOException if the state header cannot be read + */ + public static List getKeyedStates( + CheckpointMetadata metadata, OperatorIdentifier operatorId) throws IOException { + OperatorState opState = findOperatorState(metadata, operatorId); + List schemaInfos = StateSchemaExtractor.extractSchema(opState); + ClassifiedStates classified = classifyStates(operatorId.toString(), schemaInfos); + return classified.voidNamespaceStates.stream() + .map(info -> info.stateName) + .collect(Collectors.toList()); + } + + /** + * Returns the {@link KeyedStateSchemaInfo} for the plain per-key (void-namespace) states of the + * given operator — the ones exposed by the {@code _keyed}/{@code _keyed_flat} tables. + * + *

Schema extraction is lenient: POJO field names and types are derived from the serializer + * snapshot and do not require the user POJO class to be on the classpath. + * + * @param metadata the checkpoint metadata to inspect + * @param operatorId identifies the operator + * @return schema information covering the key type and all registered state entries + * @throws IOException if the state header cannot be read + */ + public static KeyedStateSchemaInfo getKeyedStateSchema( + CheckpointMetadata metadata, OperatorIdentifier operatorId) throws IOException { + OperatorState opState = findOperatorState(metadata, operatorId); + List schemas = StateSchemaExtractor.extractSchema(opState); + ClassifiedStates classified = classifyStates(operatorId.toString(), schemas); + return buildKeyedStateSchemaInfo(schemas, classified.voidNamespaceStates, null); + } + + private static KeyedStateSchemaInfo buildKeyedStateSchemaInfo( + List allSchemas, + List statesToInclude, + @Nullable LogicalType windowLogicalType) { + LogicalType keyType = + allSchemas.isEmpty() + ? new VarBinaryType(true, VarBinaryType.MAX_LENGTH) + : SerializerSnapshotToLogicalTypeConverter.convert( + allSchemas.get(0).keySnapshot); + + LinkedHashMap stateSchemas = + new LinkedHashMap<>(); + for (StateSchemaInfo info : statesToInclude) { + SavepointConnectorOptions.StateType stateType; + if (info.stateKind == StateDescriptor.Type.LIST) { + stateType = SavepointConnectorOptions.StateType.LIST; + } else if (info.stateKind == StateDescriptor.Type.MAP) { + stateType = SavepointConnectorOptions.StateType.MAP; + } else { + stateType = SavepointConnectorOptions.StateType.VALUE; + } + + try { + LogicalType logicalType = + SerializerSnapshotToLogicalTypeConverter.convert(info.valueSnapshot); + stateSchemas.put( + info.stateName, + new KeyedStateSchemaInfo.StateEntryInfo( + stateType, logicalType, windowLogicalType)); + } catch (Exception e) { + logSchemaExtractionFailure("", info.stateName, info.valueSnapshot, e); + } + } + + return new KeyedStateSchemaInfo(keyType, stateSchemas); + } + + /** + * Logs that a single state's schema could not be extracted and will therefore be excluded from + * the table schema, shared by {@link #buildKeyedStateSchemaInfo}. + * + * @param label a prefix inserted before "state" in the log message (e.g. {@code "non-keyed "} + * or {@code ""}), distinguishing which caller excluded the state + */ + private static void logSchemaExtractionFailure( + String label, + String stateName, + @Nullable TypeSerializerSnapshot valueSnapshot, + Exception e) { + LOG.error( + "Cannot extract schema for {}state '{}' (serializer type: {}): {}. " + + "This state will be excluded from the table schema. " + + "Use explicit connector options to include it.", + label, + stateName, + valueSnapshot == null ? "null" : valueSnapshot.getClass().getSimpleName(), + e.getMessage()); + } + + /** + * Builds a {@link CatalogTable} representing all keyed states of an operator. + * + *

The resulting table has one column named {@code "state_key"} for the key and one column + * per keyed state. The connector options are pre-populated so the table can be registered + * directly in a {@link org.apache.flink.table.catalog.CatalogManager}. + * + *

When the state backend that produced the operator's keyed state can be unambiguously + * determined from the checkpoint metadata, {@link SavepointConnectorOptions#STATE_BACKEND_TYPE} + * is pre-populated as well, so callers don't need to specify it themselves. + * + * @param metadata the checkpoint metadata the operator belongs to + * @param schemaInfo the schema information returned by {@link #getKeyedStateSchema} + * @param statePath the path to the savepoint / checkpoint + * @param operatorIdentifier identifies the operator whose state to read + * @return a {@link CatalogTable} ready for registration + */ + public static CatalogTable getStateCatalogTable( + CheckpointMetadata metadata, + KeyedStateSchemaInfo schemaInfo, + String statePath, + OperatorIdentifier operatorIdentifier) { + return buildKeyedCatalogTable(metadata, schemaInfo, statePath, operatorIdentifier, null); + } + + /** + * Builds a {@link CatalogTable} representing all keyed states of an operator, or, when {@code + * windowType} is non-null, all namespaced (e.g. window-scoped) states of an operator. + */ + private static CatalogTable buildKeyedCatalogTable( + CheckpointMetadata metadata, + KeyedStateSchemaInfo schemaInfo, + String statePath, + OperatorIdentifier operatorIdentifier, + @Nullable LogicalType windowType) { + + Schema.Builder schemaBuilder = Schema.newBuilder(); + schemaBuilder.column( + "state_key", LogicalTypeDataTypeConverter.toDataType(schemaInfo.keyType).notNull()); + if (windowType != null) { + schemaBuilder.column( + "state_window", LogicalTypeDataTypeConverter.toDataType(windowType).notNull()); + } + + for (Map.Entry entry : + schemaInfo.stateSchemas.entrySet()) { + schemaBuilder.column(entry.getKey(), stateValueColumnDataType(entry.getValue())); + } + if (windowType == null) { + schemaBuilder.primaryKeyNamed("PK_state_key", "state_key"); + } + Schema schema = schemaBuilder.build(); + + Map options = buildBaseConnectorOptions(statePath, operatorIdentifier); + options.put( + SavepointConnectorOptions.STATE_READER_MODE.key(), + (windowType == null + ? SavepointConnectorOptions.StateReaderMode.KEYED + : SavepointConnectorOptions.StateReaderMode.WINDOWED) + .toString()); + withStateBackendType(options, metadata, operatorIdentifier); + + return CatalogTable.newBuilder().schema(schema).options(options).build(); + } + + /** + * Builds a {@link CatalogTable} exposing a single keyed LIST or MAP state flattened into one + * row per list element / map entry, rather than one row per key. + * + *

The resulting table has 3 columns, with a composite primary key on {@code state_key} and + * the sub-key column (the {@code state_key} value repeats across rows belonging to the same + * key, but the pair uniquely identifies a row). The third column has a fixed name — not the + * state's own name, to avoid collisions with other (reserved) column names: + * + *

    + *
  • LIST: {@code (state_key, list_index, list_value)}, primary key {@code (state_key, + * list_index)} + *
  • MAP: {@code (state_key, map_key, map_value)}, primary key {@code (state_key, map_key)} + *
+ * + * @param metadata the checkpoint metadata the operator belongs to + * @param schemaInfo the schema information returned by {@link #getKeyedStateSchema} + * @param stateName the name of the LIST or MAP state to flatten + * @param statePath the path to the savepoint / checkpoint + * @param operatorIdentifier identifies the operator whose state to read + * @return a {@link CatalogTable} ready for registration + */ + public static CatalogTable getFlattenedStateCatalogTable( + CheckpointMetadata metadata, + KeyedStateSchemaInfo schemaInfo, + String stateName, + String statePath, + OperatorIdentifier operatorIdentifier) { + return buildFlattenedKeyedCatalogTable( + metadata, schemaInfo, stateName, statePath, operatorIdentifier, false); + } + + /** + * Builds a {@link CatalogTable} exposing a single LIST or MAP state flattened into one row per + * list element / map entry, either plain-keyed ({@code windowed == false}, see {@link + * #getFlattenedStateCatalogTable}) or namespaced ({@code windowed == true}). + */ + private static CatalogTable buildFlattenedKeyedCatalogTable( + CheckpointMetadata metadata, + KeyedStateSchemaInfo schemaInfo, + String stateName, + String statePath, + OperatorIdentifier operatorIdentifier, + boolean windowed) { + + KeyedStateSchemaInfo.StateEntryInfo entryInfo = schemaInfo.stateSchemas.get(stateName); + if (entryInfo == null) { + throw new IllegalArgumentException( + "State '" + + stateName + + "' not found for operator '" + + operatorIdentifier + + "'."); + } + if (entryInfo.stateType != SavepointConnectorOptions.StateType.LIST + && entryInfo.stateType != SavepointConnectorOptions.StateType.MAP) { + throw new IllegalArgumentException( + "Flattened state tables are only supported for LIST and MAP states, but '" + + stateName + + "' is " + + entryInfo.stateType + + "."); + } + if (windowed && entryInfo.windowLogicalType == null) { + throw new IllegalArgumentException( + "State '" + + stateName + + "' is not a namespaced state for operator '" + + operatorIdentifier + + "'."); + } + + Schema.Builder schemaBuilder = Schema.newBuilder(); + schemaBuilder.column( + "state_key", LogicalTypeDataTypeConverter.toDataType(schemaInfo.keyType).notNull()); + if (windowed) { + schemaBuilder.column( + "state_window", + LogicalTypeDataTypeConverter.toDataType(entryInfo.windowLogicalType).notNull()); + } + + String subKeyColumnName = addFlattenedValueColumns(schemaBuilder, entryInfo); + if (!windowed) { + schemaBuilder.primaryKeyNamed( + "PK_state_key_" + subKeyColumnName, "state_key", subKeyColumnName); + } + Schema schema = schemaBuilder.build(); + + Map options = buildBaseConnectorOptions(statePath, operatorIdentifier); + options.put( + SavepointConnectorOptions.STATE_READER_MODE.key(), + (windowed + ? SavepointConnectorOptions.StateReaderMode.WINDOWED_FLAT + : SavepointConnectorOptions.StateReaderMode.KEYED_FLAT) + .toString()); + options.put(SavepointConnectorOptions.FLATTENED_STATE_NAME.key(), stateName); + withStateBackendType(options, metadata, operatorIdentifier); + + return CatalogTable.newBuilder().schema(schema).options(options).build(); + } + + /** + * Adds the LIST- or MAP-shaped sub-key and value columns (e.g. {@code (list_index, list_value)} + * or {@code (map_key, map_value)}) for a flattened state table, and returns the sub-key + * column's name. + */ + private static String addFlattenedValueColumns( + Schema.Builder schemaBuilder, KeyedStateSchemaInfo.StateEntryInfo entryInfo) { + LogicalType valueType; + String subKeyColumnName; + String valueColumnName; + if (entryInfo.stateType == SavepointConnectorOptions.StateType.LIST) { + valueType = ((ArrayType) entryInfo.logicalType).getElementType(); + subKeyColumnName = "list_index"; + valueColumnName = "list_value"; + schemaBuilder.column( + subKeyColumnName, + LogicalTypeDataTypeConverter.toDataType(new BigIntType(false))); + } else { + MapType mapType = (MapType) entryInfo.logicalType; + valueType = mapType.getValueType(); + subKeyColumnName = "map_key"; + valueColumnName = "map_value"; + schemaBuilder.column( + subKeyColumnName, + LogicalTypeDataTypeConverter.toDataType(mapType.getKeyType()).notNull()); + } + schemaBuilder.column(valueColumnName, LogicalTypeDataTypeConverter.toDataType(valueType)); + return subKeyColumnName; + } + + // ------------------------------------------------------------------------- + // Private helpers + // ------------------------------------------------------------------------- + + /** + * Resolves the SQL column {@link org.apache.flink.table.types.DataType} for a single state's + * value column, forcing it nullable for VALUE-shaped state: unlike LIST/MAP (which always have + * a value, possibly empty), a {@code ValueState}/{@code ReducingState}/{@code AggregatingState} + * can legitimately hold no value (e.g. never written, or cleared by a trigger such as {@code + * CountTrigger}), in which case a read returns {@code null}. + */ + private static DataType stateValueColumnDataType( + KeyedStateSchemaInfo.StateEntryInfo entryInfo) { + DataType dataType = LogicalTypeDataTypeConverter.toDataType(entryInfo.logicalType); + return entryInfo.stateType == SavepointConnectorOptions.StateType.VALUE + ? dataType.nullable() + : dataType; + } + + private static OperatorState findOperatorState( + CheckpointMetadata metadata, OperatorIdentifier operatorId) { + for (OperatorState op : metadata.getOperatorStates()) { + if (op.getOperatorID().equals(operatorId.getOperatorId())) { + return op; + } + } + throw new IllegalArgumentException( + "Operator '" + operatorId + "' not found in checkpoint metadata."); + } + + /** + * Returns the base connector options ({@link FactoryUtil#CONNECTOR}, {@link + * SavepointConnectorOptions#STATE_PATH}, and the operator identifier option) shared by every + * savepoint-backed {@link CatalogTable}. + */ + private static Map buildBaseConnectorOptions( + String statePath, OperatorIdentifier operatorIdentifier) { + Map options = new HashMap<>(); + options.put(FactoryUtil.CONNECTOR.key(), "savepoint"); + options.put(SavepointConnectorOptions.STATE_PATH.key(), statePath); + operatorIdentifier + .getUid() + .ifPresentOrElse( + uid -> options.put(SavepointConnectorOptions.OPERATOR_UID.key(), uid), + () -> + options.put( + SavepointConnectorOptions.OPERATOR_UID_HASH.key(), + operatorIdentifier.getOperatorId().toHexString())); + return options; + } + + /** + * Adds {@link SavepointConnectorOptions#STATE_BACKEND_TYPE} to {@code options} when it can be + * unambiguously determined from the checkpoint metadata. Only meaningful for keyed state + * tables: non-keyed (list/union/broadcast) state isn't stored in a state backend, so callers + * for those table kinds must not call this. + */ + private static void withStateBackendType( + Map options, + CheckpointMetadata metadata, + OperatorIdentifier operatorIdentifier) { + OperatorState opState = findOperatorState(metadata, operatorIdentifier); + detectStateBackendType(opState) + .ifPresent( + type -> + options.put( + SavepointConnectorOptions.STATE_BACKEND_TYPE.key(), type)); + } + + /** + * Attempts to determine the state backend (shortcut name, see {@link + * StateBackendLoader#HASHMAP_STATE_BACKEND_NAME} / {@link + * StateBackendLoader#ROCKSDB_STATE_BACKEND_NAME}) that produced the operator's keyed state, by + * inspecting the concrete {@link KeyedStateHandle} subtype found in the checkpoint metadata: + * heap/HashMap backends produce {@link KeyGroupsStateHandle}, RocksDB/ForSt backends produce + * {@link IncrementalKeyedStateHandle}. + * + *

Canonical-format savepoints rewrite keyed state into the backend-agnostic {@link + * KeyGroupsSavepointStateHandle}, in which case the originating backend can no longer be + * determined from the handle alone; an empty result is returned rather than guessing. + */ + static Optional detectStateBackendType(OperatorState opState) { + Set detectedTypes = new HashSet<>(); + for (OperatorSubtaskState subtaskState : opState.getStates()) { + collectStateBackendTypes(subtaskState.getManagedKeyedState(), detectedTypes); + collectStateBackendTypes(subtaskState.getRawKeyedState(), detectedTypes); + } + if (detectedTypes.size() != 1) { + if (detectedTypes.size() > 1) { + LOG.warn( + "Operator '{}' has keyed state handles from multiple state backends {}; " + + "not setting '{}'.", + opState.getOperatorID(), + detectedTypes, + SavepointConnectorOptions.STATE_BACKEND_TYPE.key()); + } + return Optional.empty(); + } + return Optional.of(detectedTypes.iterator().next()); + } + + private static void collectStateBackendTypes( + Iterable handles, Set detectedTypes) { + for (KeyedStateHandle handle : handles) { + if (handle instanceof ChangelogStateBackendHandle) { + collectStateBackendTypes( + ((ChangelogStateBackendHandle) handle).getMaterializedStateHandles(), + detectedTypes); + } else if (handle instanceof IncrementalKeyedStateHandle) { + detectedTypes.add(StateBackendLoader.ROCKSDB_STATE_BACKEND_NAME); + } else if (handle instanceof KeyGroupsSavepointStateHandle) { + // Canonical-format savepoints rewrite keyed state into a backend-agnostic + // format; the originating backend can no longer be told apart from the handle. + } else if (handle instanceof KeyGroupsStateHandle) { + detectedTypes.add(StateBackendLoader.HASHMAP_STATE_BACKEND_NAME); + } else { + LOG.warn("Unknown handle type '{}'.", handle.getClass().getSimpleName()); + } + } + } + + /** Returns {@code true} for Flink-internal states that are not user-registered states. */ + private static boolean isInternalState(String stateName) { + return stateName.startsWith(InternalTimeServiceManagerImpl.TIMER_STATE_PREFIX + "/") + || stateName.equals(WindowOperator.MERGING_WINDOW_SET_STATE_NAME); + } + + /** + * Returns {@code true} if a state is plain per-key state (registered with {@code + * VoidNamespace}) and {@code false} if it is scoped by some other namespace (e.g. a window). + * + *

A missing namespace snapshot (e.g. from an older savepoint format) is treated as void, + * matching pre-existing behavior. + */ + private static boolean isVoidNamespace(TypeSerializerSnapshot namespaceSnapshot) { + return namespaceSnapshot == null + || namespaceSnapshot + instanceof VoidNamespaceSerializer.VoidNamespaceSerializerSnapshot; + } + + /** + * The result of {@link #classifyStates}: user-registered states of an operator, partitioned + * into plain per-key (void-namespace) states and namespaced (e.g. window-scoped) states. + */ + private static final class ClassifiedStates { + final List voidNamespaceStates; + final List windowNamespaceStates; + @Nullable final LogicalType windowLogicalType; + + ClassifiedStates( + List voidNamespaceStates, + List windowNamespaceStates, + @Nullable LogicalType windowLogicalType) { + this.voidNamespaceStates = voidNamespaceStates; + this.windowNamespaceStates = windowNamespaceStates; + this.windowLogicalType = windowLogicalType; + } + } + + /** + * Partitions the user-registered states of a single operator into plain per-key + * (void-namespace) states and namespaced states, resolving the namespaced states' shared {@link + * LogicalType} along the way. + * + *

An operator may register states under more than one distinct namespace type only + * via hand-rolled state access (no built-in windowing API does this); when that happens, the + * first namespace type whose schema can be determined is kept, and every other group is + * excluded with a logged warning. + * + * @param operatorLabel a human-readable operator identifier, used only for log messages + */ + private static ClassifiedStates classifyStates( + String operatorLabel, List schemas) { + List voidStates = new ArrayList<>(); + Map> namespacedGroups = new LinkedHashMap<>(); + for (StateSchemaInfo info : schemas) { + if (isInternalState(info.stateName)) { + continue; + } + if (isVoidNamespace(info.namespaceSnapshot)) { + voidStates.add(info); + } else { + namespacedGroups + .computeIfAbsent( + info.namespaceSnapshot.getClass().getName(), k -> new ArrayList<>()) + .add(info); + } + } + + List chosenGroup = Collections.emptyList(); + LogicalType chosenNamespaceType = null; + for (Map.Entry> entry : namespacedGroups.entrySet()) { + if (chosenNamespaceType != null) { + logExcludedNamespaceGroup( + operatorLabel, + entry.getKey(), + entry.getValue(), + "an operator has states registered under more than one namespace type"); + continue; + } + + TypeSerializerSnapshot representative = entry.getValue().get(0).namespaceSnapshot; + try { + chosenNamespaceType = + SerializerSnapshotToLogicalTypeConverter.convert(representative); + } catch (UnsupportedOperationException e) { + logExcludedNamespaceGroup( + operatorLabel, + entry.getKey(), + entry.getValue(), + "cannot extract schema for this namespace type: " + e.getMessage()); + continue; + } + chosenGroup = entry.getValue(); + } + + return new ClassifiedStates(voidStates, chosenGroup, chosenNamespaceType); + } + + private static void logExcludedNamespaceGroup( + String operatorLabel, + String namespaceClassName, + List excluded, + String reason) { + LOG.warn( + "Excluding namespace type '{}' on operator '{}' from the catalog: {}. States: {}.", + namespaceClassName, + operatorLabel, + reason, + excluded.stream().map(i -> i.stateName).collect(Collectors.toList())); + } +} 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 00bb5c262e1931..3850265a861829 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 @@ -24,7 +24,8 @@ import org.apache.flink.api.common.io.DefaultInputSplitAssigner; import org.apache.flink.api.common.io.RichInputFormat; import org.apache.flink.api.common.io.statistics.BaseStatistics; -import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.api.common.typeutils.CustomRestoreSerializerFactory; +import org.apache.flink.api.java.tuple.Tuple3; import org.apache.flink.configuration.Configuration; import org.apache.flink.core.fs.CloseableRegistry; import org.apache.flink.core.io.InputSplitAssigner; @@ -38,6 +39,7 @@ 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.deserializer.MissingClassSerializerFactory; import org.apache.flink.state.api.input.operator.StateReaderOperator; import org.apache.flink.state.api.input.splits.KeyGroupRangeInputSplit; import org.apache.flink.state.api.runtime.SavepointRuntimeContext; @@ -92,7 +94,7 @@ public class KeyedStateInputFormat private transient BufferingCollector out; - private transient CloseableIterator> keysAndNamespaces; + private transient CloseableIterator> keysAndNamespaces; /** * Creates an input format for reading partitioned state from an operator in a savepoint. @@ -211,15 +213,10 @@ public void open(KeyGroupRangeInputSplit split) throws IOException { registry = new CloseableRegistry(); RuntimeContext runtimeContext = getRuntimeContext(); - ExecutionConfig executionConfig; - try { - executionConfig = - serializedExecutionConfig.deserializeValue( - runtimeContext.getUserCodeClassLoader()); - } catch (ClassNotFoundException e) { - throw new RuntimeException("Could not deserialize ExecutionConfig.", e); - } - final StreamOperatorStateContext context = + ExecutionConfig executionConfig = + deserialize(serializedExecutionConfig, runtimeContext.getUserCodeClassLoader()); + + StreamOperatorContextBuilder builder = new StreamOperatorContextBuilder( runtimeContext, configuration, @@ -229,19 +226,29 @@ public void open(KeyGroupRangeInputSplit split) throws IOException { stateBackend, executionConfig) .withMaxParallelism(split.getNumKeyGroups()) - .withKey(operator, runtimeContext.createSerializer(operator.getKeyType())) - .build(LOG); + .withKey(operator, runtimeContext.createSerializer(operator.getKeyType())); - AbstractKeyedStateBackend keyedStateBackend = - (AbstractKeyedStateBackend) context.keyedStateBackend(); + // Deserialize any POJO/Avro state whose class is missing from the classpath into + // RowData/GenericRecord instead of failing the restore. PojoSerializerSnapshot and + // AvroSerializerSnapshot only consult this factory once they've already determined that the + // class they need is genuinely missing, so registering it unconditionally is safe and has + // no effect on states whose classes are present. + CustomRestoreSerializerFactory.set(MissingClassSerializerFactory::create); - final DefaultKeyedStateStore keyedStateStore = - new DefaultKeyedStateStore(keyedStateBackend, runtimeContext::createSerializer); - SavepointRuntimeContext ctx = new SavepointRuntimeContext(runtimeContext, keyedStateStore); - - InternalTimeServiceManager timeServiceManager = - (InternalTimeServiceManager) context.internalTimerServiceManager(); try { + final StreamOperatorStateContext context = builder.build(LOG); + + AbstractKeyedStateBackend keyedStateBackend = + (AbstractKeyedStateBackend) context.keyedStateBackend(); + + final DefaultKeyedStateStore keyedStateStore = + new DefaultKeyedStateStore(keyedStateBackend, runtimeContext::createSerializer); + SavepointRuntimeContext ctx = + new SavepointRuntimeContext(runtimeContext, keyedStateStore); + + InternalTimeServiceManager timeServiceManager = + (InternalTimeServiceManager) context.internalTimerServiceManager(); + operator.setup( runtimeContext::createSerializer, keyedStateBackend, timeServiceManager, ctx); operator.open(); @@ -280,8 +287,8 @@ public OUT nextRecord(OUT reuse) throws IOException { return out.next(); } - final Tuple2 keyAndNamespace = keysAndNamespaces.next(); - operator.setCurrentKey(keyAndNamespace.f0); + final Tuple3 keyAndNamespace = keysAndNamespaces.next(); + operator.setCurrentKeyAndKeyGroup(keyAndNamespace.f0, keyAndNamespace.f2); try { operator.processElement(keyAndNamespace.f0, keyAndNamespace.f1, out); @@ -325,4 +332,14 @@ private static List sortedKeyGroupRanges(int minNumSplits, int ma keyGroups.sort(Comparator.comparing(KeyGroupRange::getStartKeyGroup)); return keyGroups; } + + static ExecutionConfig deserialize( + SerializedValue serializedExecutionConfig, ClassLoader classLoader) + throws IOException { + try { + return serializedExecutionConfig.deserializeValue(classLoader); + } catch (ClassNotFoundException e) { + throw new RuntimeException("Could not deserialize ExecutionConfig.", e); + } + } } diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/MultiStateKeyIterator.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/MultiStateKeyIterator.java index 38d1d3505a2577..7d6953eac51bbe 100644 --- a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/MultiStateKeyIterator.java +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/MultiStateKeyIterator.java @@ -20,6 +20,7 @@ import org.apache.flink.annotation.Internal; import org.apache.flink.api.common.state.StateDescriptor; +import org.apache.flink.api.java.tuple.Tuple2; import org.apache.flink.core.fs.CloseableRegistry; import org.apache.flink.runtime.state.KeyedStateBackend; import org.apache.flink.runtime.state.VoidNamespace; @@ -36,13 +37,20 @@ /** * An iterator for reading all keys in a state backend across multiple partitioned states. * + *

Uses {@link KeyedStateBackend#getKeysAndKeyGroups} rather than {@link + * KeyedStateBackend#getKeys(List, Object)}, because the State Processing API can read state whose + * key class is missing from the classpath, deserializing the key into a substitute representation + * instead of failing (see {@code CustomRestoreSerializerFactory}). Such a substituted key's {@code + * hashCode()} is not guaranteed to reproduce the original, so the key-group it was physically + * written under must be carried alongside it rather than recomputed from the key. + * * @param Type of the key by which state is keyed. */ @Internal -public final class MultiStateKeyIterator implements CloseableIterator { +public final class MultiStateKeyIterator implements CloseableIterator> { private final List> descriptors; - private final Iterator iterator; + private final Iterator> iterator; private final CloseableRegistry registry; @@ -52,8 +60,8 @@ public MultiStateKeyIterator( this.descriptors = Preconditions.checkNotNull(descriptors); Preconditions.checkNotNull(backend); registry = new CloseableRegistry(); - Stream stream = - backend.getKeys( + Stream> stream = + backend.getKeysAndKeyGroups( this.descriptors.stream() .map(StateDescriptor::getName) .collect(Collectors.toList()), @@ -72,7 +80,7 @@ public boolean hasNext() { } @Override - public K next() { + public Tuple2 next() { if (!hasNext()) { throw new NoSuchElementException(); } else { diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/OperatorStateInputFormat.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/OperatorStateInputFormat.java index 978160dc0e069a..65745a1b16ab55 100644 --- a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/OperatorStateInputFormat.java +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/OperatorStateInputFormat.java @@ -173,14 +173,9 @@ public void open(OperatorStateInputSplit split) throws IOException { registry = new CloseableRegistry(); RuntimeContext runtimeContext = getRuntimeContext(); - ExecutionConfig executionConfig; - try { - executionConfig = - serializedExecutionConfig.deserializeValue( - runtimeContext.getUserCodeClassLoader()); - } catch (ClassNotFoundException e) { - throw new RuntimeException("Could not deserialize ExecutionConfig.", e); - } + ExecutionConfig executionConfig = + KeyedStateInputFormat.deserialize( + serializedExecutionConfig, runtimeContext.getUserCodeClassLoader()); final StreamOperatorStateContext context = new StreamOperatorContextBuilder( runtimeContext, diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/deserializer/EnumNameDeserializer.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/deserializer/EnumNameDeserializer.java new file mode 100644 index 00000000000000..5105b1d65a048e --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/deserializer/EnumNameDeserializer.java @@ -0,0 +1,145 @@ +/* + * 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.deserializer; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; +import org.apache.flink.api.common.typeutils.base.EnumSerializer; +import org.apache.flink.core.memory.DataInputView; +import org.apache.flink.core.memory.DataOutputView; + +import java.io.IOException; +import java.util.Arrays; + +/** + * A {@link TypeSerializer} that reads the enum ordinal written by {@link EnumSerializer} and + * produces the enum constant's {@code name()} as a plain {@link String}, without the user enum + * class being on the classpath. + * + *

{@link EnumSerializer} writes a maintained ordinal (see {@link + * EnumSerializer.EnumSerializerSnapshot#getEnumNames()}) rather than {@link Enum#ordinal()}, so the + * name lookup here uses the exact same array the original serializer would have used to resolve + * that ordinal back to a constant. + */ +@Internal +public final class EnumNameDeserializer extends TypeSerializer { + + private static final long serialVersionUID = 1L; + + private final String[] enumNames; + + public static EnumNameDeserializer create(EnumSerializer.EnumSerializerSnapshot snapshot) { + return new EnumNameDeserializer(snapshot.getEnumNames()); + } + + EnumNameDeserializer(String[] enumNames) { + this.enumNames = enumNames; + } + + @Override + public String deserialize(DataInputView source) throws IOException { + int ordinal = source.readInt(); + if (ordinal < 0 || ordinal >= enumNames.length) { + throw new IOException( + "Unknown enum ordinal " + + ordinal + + " (have " + + enumNames.length + + " known constants). The savepoint may have been written with a" + + " different enum definition."); + } + return enumNames[ordinal]; + } + + @Override + public String deserialize(String reuse, DataInputView source) throws IOException { + return deserialize(source); + } + + // ------------------------------------------------------------------------- + // TypeSerializer boilerplate — copy/snapshot operations not needed for + // schema-extraction use cases but required by the interface. + // ------------------------------------------------------------------------- + + @Override + public boolean isImmutableType() { + return true; + } + + @Override + public TypeSerializer duplicate() { + return this; + } + + @Override + public String createInstance() { + return enumNames.length > 0 ? enumNames[0] : null; + } + + @Override + public String copy(String from) { + return from; + } + + @Override + public String copy(String from, String reuse) { + return from; + } + + @Override + public int getLength() { + return 4; + } + + @Override + public void serialize(String record, DataOutputView target) throws IOException { + throw new UnsupportedOperationException( + "EnumNameDeserializer is read-only; serialization is not supported."); + } + + @Override + public void copy(DataInputView source, DataOutputView target) throws IOException { + throw new UnsupportedOperationException( + "EnumNameDeserializer is read-only; copy is not supported."); + } + + @Override + public boolean equals(Object obj) { + if (this == obj) { + return true; + } + if (!(obj instanceof EnumNameDeserializer)) { + return false; + } + return Arrays.equals(enumNames, ((EnumNameDeserializer) obj).enumNames); + } + + @Override + public int hashCode() { + return Arrays.hashCode(enumNames); + } + + @Override + public TypeSerializerSnapshot snapshotConfiguration() { + throw new UnsupportedOperationException( + "EnumNameDeserializer is only ever used directly within a single read; it is never" + + " re-snapshotted."); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/deserializer/InternalTypeConverter.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/deserializer/InternalTypeConverter.java new file mode 100644 index 00000000000000..2d726f3de26a6f --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/deserializer/InternalTypeConverter.java @@ -0,0 +1,293 @@ +/* + * 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.deserializer; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.table.data.DecimalData; +import org.apache.flink.table.data.GenericArrayData; +import org.apache.flink.table.data.GenericMapData; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.StringData; +import org.apache.flink.table.data.TimestampData; +import org.apache.flink.table.types.logical.ArrayType; +import org.apache.flink.table.types.logical.DecimalType; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.MapType; +import org.apache.flink.table.types.logical.MultisetType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.types.Row; + +import javax.annotation.Nullable; + +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.sql.Timestamp; +import java.time.Instant; +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.StreamSupport; + +/** + * Converts external Java objects (as produced by DataStream serializers) to Flink table internal + * types as expected by {@link GenericRowData}. + * + *

Conversion rules: + * + *

    + *
  • {@link String} → {@link StringData} + *
  • {@link BigDecimal} → {@link DecimalData} (precision/scale from {@link DecimalType}) + *
  • {@link ByteBuffer} or {@code byte[]} → {@link DecimalData} (unscaled bytes) + *
  • {@link ByteBuffer} → {@code byte[]} for BINARY/VARBINARY + *
  • {@link java.sql.Date}, {@link LocalDate} → {@code int} (days since epoch) + *
  • {@link Timestamp}, {@link Instant}, {@link LocalDateTime} → {@link TimestampData} + *
  • {@link List}, arrays, {@link Iterable} → {@link GenericArrayData} (elements recursively + * converted) + *
  • {@link Map} or {@link Iterable} of {@link Map.Entry} → {@link GenericMapData} (keys/values + * recursively converted) + *
  • {@link Row} → {@link GenericRowData} (fields recursively converted) + *
  • {@link RowData} subtypes → passed through unchanged + *
  • Primitives (boxed) → passed through unchanged + *
+ */ +@Internal +public final class InternalTypeConverter { + + private InternalTypeConverter() {} + + /** + * Converts {@code value} to the Flink table internal representation dictated by {@code type}. + * + * @param value the raw Java object; may be null + * @param type the target logical type; used to drive nested conversions + * @return the converted value, or null if value is null + */ + @Nullable + public static Object toInternal(@Nullable Object value, LogicalType type) { + if (value == null) { + return null; + } + + switch (type.getTypeRoot()) { + case CHAR: + case VARCHAR: + if (value instanceof StringData) { + return value; + } + return StringData.fromString(value.toString()); + + case BOOLEAN: + case TINYINT: + case SMALLINT: + case INTEGER: + case BIGINT: + case FLOAT: + case DOUBLE: + case TIME_WITHOUT_TIME_ZONE: + case INTERVAL_YEAR_MONTH: + case INTERVAL_DAY_TIME: + return value; + + case DECIMAL: + if (value instanceof DecimalData) { + return value; + } + if (value instanceof BigDecimal) { + DecimalType dt = (DecimalType) type; + return DecimalData.fromBigDecimal( + (BigDecimal) value, dt.getPrecision(), dt.getScale()); + } + if (value instanceof ByteBuffer) { + DecimalType dt = (DecimalType) type; + return DecimalData.fromUnscaledBytes( + toByteArray((ByteBuffer) value), dt.getPrecision(), dt.getScale()); + } + if (value instanceof byte[]) { + DecimalType dt = (DecimalType) type; + return DecimalData.fromUnscaledBytes( + (byte[]) value, dt.getPrecision(), dt.getScale()); + } + return value; + + case DATE: + if (value instanceof Integer) { + return value; + } + if (value instanceof java.sql.Date) { + return (int) ((java.sql.Date) value).toLocalDate().toEpochDay(); + } + if (value instanceof LocalDate) { + return (int) ((LocalDate) value).toEpochDay(); + } + return value; + + case TIMESTAMP_WITHOUT_TIME_ZONE: + case TIMESTAMP_WITH_TIME_ZONE: + case TIMESTAMP_WITH_LOCAL_TIME_ZONE: + if (value instanceof TimestampData) { + return value; + } + if (value instanceof Timestamp) { + return TimestampData.fromTimestamp((Timestamp) value); + } + if (value instanceof Instant) { + return TimestampData.fromInstant((Instant) value); + } + if (value instanceof LocalDateTime) { + return TimestampData.fromLocalDateTime((LocalDateTime) value); + } + return value; + + case BINARY: + case VARBINARY: + if (value instanceof ByteBuffer) { + return toByteArray((ByteBuffer) value); + } + return value; + + case NULL: + return null; + + case ROW: + case STRUCTURED_TYPE: + if (value instanceof GenericRowData) { + return value; + } + if (value instanceof Row) { + return rowToGenericRowData((Row) value, (RowType) type); + } + return value; + + case ARRAY: + if (value instanceof GenericArrayData) { + return value; + } + ArrayType at = (ArrayType) type; + if (value instanceof Object[]) { + return objectArrayToArrayData((Object[]) value, at.getElementType()); + } + if (value instanceof Iterable) { + return iterableToArrayData((Iterable) value, at.getElementType()); + } + return value; + + case MAP: + if (value instanceof GenericMapData) { + return value; + } + MapType mt = (MapType) type; + if (value instanceof Map) { + return mapToMapData((Map) value, mt.getKeyType(), mt.getValueType()); + } + if (value instanceof Iterable) { + return mapEntryIterableToMapData( + (Iterable) value, mt.getKeyType(), mt.getValueType()); + } + return value; + + case MULTISET: + // MultisetType is not a MapType: it has only an element type, represented + // internally as Map (element -> multiplicity). + if (value instanceof GenericMapData) { + return value; + } + LogicalType elementType = ((MultisetType) type).getElementType(); + if (value instanceof Map) { + return mapToMapData((Map) value, elementType, new IntType()); + } + if (value instanceof Iterable) { + return mapEntryIterableToMapData( + (Iterable) value, elementType, new IntType()); + } + return value; + + default: + throw new UnsupportedOperationException( + "Cannot convert value of type '" + + value.getClass().getName() + + "' to internal representation for LogicalTypeRoot " + + type.getTypeRoot() + + "."); + } + } + + private static byte[] toByteArray(ByteBuffer bb) { + byte[] bytes = new byte[bb.remaining()]; + bb.get(bytes); + return bytes; + } + + private static GenericRowData rowToGenericRowData(Row row, RowType rowType) { + List fields = rowType.getFields(); + GenericRowData out = new GenericRowData(row.getArity()); + out.setRowKind(row.getKind()); + for (int i = 0; i < row.getArity(); i++) { + LogicalType fieldType = i < fields.size() ? fields.get(i).getType() : null; + Object rawField = row.getField(i); + out.setField(i, fieldType != null ? toInternal(rawField, fieldType) : rawField); + } + return out; + } + + private static GenericArrayData objectArrayToArrayData(Object[] src, LogicalType elementType) { + Object[] arr = new Object[src.length]; + for (int i = 0; i < src.length; i++) { + arr[i] = toInternal(src[i], elementType); + } + return new GenericArrayData(arr); + } + + private static GenericArrayData iterableToArrayData( + Iterable iterable, LogicalType elementType) { + return new GenericArrayData( + StreamSupport.stream(iterable.spliterator(), false) + .map(v -> toInternal(v, elementType)) + .toArray()); + } + + private static GenericMapData mapToMapData( + Map map, LogicalType keyType, LogicalType valueType) { + LinkedHashMap converted = new LinkedHashMap<>(map.size()); + for (Map.Entry entry : map.entrySet()) { + converted.put( + toInternal(entry.getKey(), keyType), toInternal(entry.getValue(), valueType)); + } + return new GenericMapData(converted); + } + + private static GenericMapData mapEntryIterableToMapData( + Iterable iterable, LogicalType keyType, LogicalType valueType) { + LinkedHashMap converted = new LinkedHashMap<>(); + for (Object element : iterable) { + if (!(element instanceof Map.Entry)) { + throw new UnsupportedOperationException( + "Map conversion supports only Iterable but received: " + + iterable.getClass().getName()); + } + Map.Entry entry = (Map.Entry) element; + converted.put( + toInternal(entry.getKey(), keyType), toInternal(entry.getValue(), valueType)); + } + return new GenericMapData(converted); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/deserializer/MissingClassSerializerFactory.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/deserializer/MissingClassSerializerFactory.java new file mode 100644 index 00000000000000..7a4c99a8687386 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/deserializer/MissingClassSerializerFactory.java @@ -0,0 +1,76 @@ +/* + * 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.deserializer; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.api.common.typeutils.CustomRestoreSerializerFactory; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; +import org.apache.flink.api.common.typeutils.base.EnumSerializer; +import org.apache.flink.api.java.typeutils.runtime.PojoSerializerSnapshot; +import org.apache.flink.state.api.schema.AvroStateUtils; + +/** + * The {@link CustomRestoreSerializerFactory} implementation registered by the State Processing API + * while reading state whose original classes are not on the classpath. + * + *

Dispatches on the kind of {@link TypeSerializerSnapshot} that could not resolve its class: + * + *

    + *
  • {@link PojoSerializerSnapshot} → a {@link PojoToRowDataDeserializer} that reads the POJO + * binary format directly into {@code RowData}. + *
  • {@link EnumSerializer.EnumSerializerSnapshot} → an {@link EnumNameDeserializer} that maps + * the serialized ordinal back to its constant name. + *
  • {@code AvroSerializerSnapshot} → an {@code AvroSerializer} reading into {@code + * GenericRecord}, using the schema embedded in the snapshot. + *
+ * + *

See {@link CustomRestoreSerializerFactory} for why this must only ever be registered by the + * State Processing API's own read path, never for regular job restores. + * + *

flink-avro is an optional dependency of this module (see the module {@code pom.xml}), so this + * class must never mention an Avro type directly: doing so - even in an {@code instanceof} branch + * that is never taken - would make the JVM try to resolve that type the moment this method is + * reached for *any* unrecognized snapshot, throwing {@code NoClassDefFoundError} for callers who + * never use Avro at all. The Avro-specific branch is instead selected by class name and delegated + * to {@link AvroStateUtils}, which is only ever loaded once that name comparison has already + * confirmed Avro is genuinely on the classpath. + */ +@Internal +public final class MissingClassSerializerFactory { + + private MissingClassSerializerFactory() {} + + public static TypeSerializer create(TypeSerializerSnapshot snapshot) { + if (snapshot instanceof PojoSerializerSnapshot) { + return PojoToRowDataDeserializer.create((PojoSerializerSnapshot) snapshot); + } + if (snapshot instanceof EnumSerializer.EnumSerializerSnapshot) { + return EnumNameDeserializer.create((EnumSerializer.EnumSerializerSnapshot) snapshot); + } + if (AvroStateUtils.AVRO_SERIALIZER_SNAPSHOT_CLASS_NAME.equals( + snapshot.getClass().getName())) { + return AvroStateUtils.createFallbackSerializer(snapshot); + } + throw new UnsupportedOperationException( + "No fallback serializer available for snapshot of type '" + + snapshot.getClass().getName() + + "'."); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/deserializer/PojoDeserializerCompatibilitySnapshot.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/deserializer/PojoDeserializerCompatibilitySnapshot.java new file mode 100644 index 00000000000000..eba41a1797c3d8 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/deserializer/PojoDeserializerCompatibilitySnapshot.java @@ -0,0 +1,89 @@ +/* + * 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.deserializer; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.api.common.typeutils.TypeSerializerSchemaCompatibility; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; +import org.apache.flink.api.java.typeutils.runtime.PojoSerializerSnapshot; +import org.apache.flink.core.memory.DataInputView; +import org.apache.flink.core.memory.DataOutputView; + +import javax.annotation.Nullable; + +/** + * A {@link TypeSerializerSnapshot} for deserializers that can read POJO binary data without the + * user POJO class being on the classpath, such as {@link PojoToRowDataDeserializer}. It declares + * itself {@link TypeSerializerSchemaCompatibility#compatibleAsIs() compatible as-is} with any + * stored {@link PojoSerializerSnapshot}. + * + *

The snapshot only ever exists in memory, wrapping the live deserializer it was created from: + * composite compatibility checks (e.g. {@code + * CompositeTypeSerializerSnapshot#resolveOuterSchemaCompatibility}) restore the "new" side of a + * composite serializer even when nested-level compatibility already short-circuited to {@code + * compatibleAsIs()}, so {@link #restoreSerializer()} hands back the wrapped instance rather than + * reconstructing one from persisted bytes. + */ +@Internal +public final class PojoDeserializerCompatibilitySnapshot implements TypeSerializerSnapshot { + + @Nullable private final TypeSerializer restoredSerializer; + + /** Constructor for reading the snapshot; see {@link #restoreSerializer()}. */ + public PojoDeserializerCompatibilitySnapshot() { + this(null); + } + + public PojoDeserializerCompatibilitySnapshot(TypeSerializer restoredSerializer) { + this.restoredSerializer = restoredSerializer; + } + + @Override + public int getCurrentVersion() { + return 1; + } + + @Override + public TypeSerializerSchemaCompatibility resolveSchemaCompatibility( + TypeSerializerSnapshot oldSerializerSnapshot) { + if (oldSerializerSnapshot instanceof PojoSerializerSnapshot + || oldSerializerSnapshot instanceof PojoDeserializerCompatibilitySnapshot) { + return TypeSerializerSchemaCompatibility.compatibleAsIs(); + } + return TypeSerializerSchemaCompatibility.incompatible(); + } + + @Override + public TypeSerializer restoreSerializer() { + if (restoredSerializer == null) { + throw new UnsupportedOperationException( + "PojoDeserializerCompatibilitySnapshot cannot reconstruct the deserializer on " + + "its own. Use PojoSerializerSnapshot.restoreSerializer() or " + + "PojoToRowDataDeserializer.create()."); + } + return restoredSerializer; + } + + @Override + public void writeSnapshot(DataOutputView out) {} + + @Override + public void readSnapshot(int readVersion, DataInputView in, ClassLoader userCodeClassLoader) {} +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/deserializer/PojoToRowDataDeserializer.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/deserializer/PojoToRowDataDeserializer.java new file mode 100644 index 00000000000000..390f79bcb5bf73 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/deserializer/PojoToRowDataDeserializer.java @@ -0,0 +1,314 @@ +/* + * 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.deserializer; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; +import org.apache.flink.api.java.typeutils.runtime.PojoSerializer; +import org.apache.flink.api.java.typeutils.runtime.PojoSerializerSnapshot; +import org.apache.flink.core.memory.DataInputView; +import org.apache.flink.core.memory.DataOutputView; +import org.apache.flink.state.api.schema.SerializerSnapshotToLogicalTypeConverter; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.types.logical.LogicalType; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.util.AbstractMap; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +/** + * A {@link TypeSerializer} that reads the POJO binary format written by {@link + * org.apache.flink.api.java.typeutils.runtime.PojoSerializer} and produces {@link GenericRowData}. + * + *

This deserializer does not require the user POJO class to be on the classpath. It + * mirrors the exact binary protocol of {@code PojoSerializer}: + * + * 

{@code
+ * 1 byte: flags (bitmask)
+ *   0x01 IS_NULL            → value is null, return null
+ *   0x02 NO_SUBCLASS        → exact POJO class: read numFields × (isNull boolean + field bytes)
+ *   0x08 IS_TAGGED_SUBCLASS → 1 byte subclass tag; delegate to registered subclass deserializer
+ *   0x04 IS_SUBCLASS        → UTF class name (must be read); Kryo not supported → throws IOException
+ * }
+ * + *

Use {@link #create(PojoSerializerSnapshot)} to build an instance from a savepoint snapshot. + */ +@Internal +public final class PojoToRowDataDeserializer extends TypeSerializer { + + private static final long serialVersionUID = 1L; + + private static final Logger LOG = LoggerFactory.getLogger(PojoToRowDataDeserializer.class); + + private final int numFields; + private final TypeSerializer[] fieldDeserializers; + private final LogicalType[] fieldTypes; + private final String[] fieldNames; + private final List registeredSubclassDeserializers; + + /** + * Builds a {@link PojoToRowDataDeserializer} from a {@link PojoSerializerSnapshot}. + * + *

For each field: + * + *

    + *
  • If the field snapshot is itself a {@link PojoSerializerSnapshot}, this method recurses + * to build a nested {@link PojoToRowDataDeserializer}. + *
  • For all other field types, the field's original serializer is restored via {@link + * TypeSerializerSnapshot#restoreSerializer()}. + *
+ * + *

Registered subclasses are handled by building a deserializer for each registered subclass + * snapshot in order (matching the tag index used in the binary format). + * + * @throws IllegalStateException if a required field serializer snapshot is absent + */ + public static PojoToRowDataDeserializer create(PojoSerializerSnapshot snapshot) { + List>> fieldEntries = + snapshot.getFieldSnapshotEntries(); + + List> fieldDeserializerList = new ArrayList<>(fieldEntries.size()); + List fieldTypeList = new ArrayList<>(fieldEntries.size()); + List fieldNameList = new ArrayList<>(fieldEntries.size()); + + for (AbstractMap.SimpleEntry> entry : fieldEntries) { + String fieldName = entry.getKey(); + TypeSerializerSnapshot fieldSnapshot = entry.getValue(); + + if (fieldSnapshot == null) { + throw new IllegalStateException( + "Cannot build deserializer for field '" + + fieldName + + "': its serializer snapshot was not readable from the savepoint. " + + "This field cannot be deserialized without the original snapshot."); + } + + TypeSerializer fieldDeserializer; + if (fieldSnapshot instanceof PojoSerializerSnapshot) { + fieldDeserializer = create((PojoSerializerSnapshot) fieldSnapshot); + } else { + fieldDeserializer = fieldSnapshot.restoreSerializer(); + } + + fieldDeserializerList.add(fieldDeserializer); + fieldTypeList.add(SerializerSnapshotToLogicalTypeConverter.convert(fieldSnapshot)); + fieldNameList.add(fieldName); + } + + List> subSnapshots = + snapshot.getRegisteredSubclassSnapshotsOrdered(); + List subDeserializers = new ArrayList<>(subSnapshots.size()); + for (TypeSerializerSnapshot subSnap : subSnapshots) { + subDeserializers.add( + subSnap instanceof PojoSerializerSnapshot + ? create((PojoSerializerSnapshot) subSnap) + : null); + } + + return new PojoToRowDataDeserializer( + fieldDeserializerList.toArray(new TypeSerializer[0]), + fieldTypeList.toArray(new LogicalType[0]), + fieldNameList.toArray(new String[0]), + subDeserializers); + } + + PojoToRowDataDeserializer( + TypeSerializer[] fieldDeserializers, + LogicalType[] fieldTypes, + String[] fieldNames, + List registeredSubclassDeserializers) { + this.numFields = fieldDeserializers.length; + this.fieldDeserializers = fieldDeserializers; + this.fieldTypes = fieldTypes; + this.fieldNames = fieldNames; + this.registeredSubclassDeserializers = registeredSubclassDeserializers; + } + + @Override + public RowData deserialize(DataInputView source) throws IOException { + int flags = source.readByte() & 0xFF; + + if ((flags & PojoSerializer.IS_NULL) != 0) { + return null; + } + + if ((flags & PojoSerializer.NO_SUBCLASS) != 0) { + return readFields(source); + } + + if ((flags & PojoSerializer.IS_TAGGED_SUBCLASS) != 0) { + int tag = source.readByte() & 0xFF; + if (tag < registeredSubclassDeserializers.size()) { + PojoToRowDataDeserializer subDeserializer = + registeredSubclassDeserializers.get(tag); + if (subDeserializer == null) { + // Either the subclass's own snapshot was unreadable, or the subclass is not + // itself a POJO (e.g. it falls back to Kryo) — either way we have no way to + // decode its bytes, and, like the IS_SUBCLASS/Kryo case below, its length is + // unknown so the bytes cannot even be skipped. + throw new IOException( + "Cannot deserialize registered POJO subclass at tag " + + tag + + ": its serializer snapshot is missing or is not a POJO " + + "serializer (e.g. it uses Kryo), which requires the class on " + + "the classpath."); + } + return subDeserializer.deserialize(source); + } + throw new IOException( + "Unknown registered subclass tag " + + tag + + " (have " + + registeredSubclassDeserializers.size() + + " registered). The savepoint may have been written with more subclasses registered."); + } + + if ((flags & PojoSerializer.IS_SUBCLASS) != 0) { + String className = source.readUTF(); + throw new IOException( + "Cannot deserialize POJO subclass '" + + className + + "': the subclass uses Kryo serialization, which requires the class on the" + + " classpath. Kryo-encoded bytes have unknown length and cannot be skipped." + + " Register the subclass or add the JAR to the classpath."); + } + + throw new IOException("Unrecognised POJO flags byte: 0x" + Integer.toHexString(flags)); + } + + @Override + public RowData deserialize(RowData reuse, DataInputView source) throws IOException { + return deserialize(source); + } + + private GenericRowData readFields(DataInputView source) throws IOException { + GenericRowData row = new GenericRowData(numFields); + for (int i = 0; i < numFields; i++) { + boolean isNull = source.readBoolean(); + if (isNull) { + row.setField(i, null); + continue; + } + // Unlike the conversion step below, a deserialize() failure here means the stream + // position for this and every subsequent field/row is now unknown: continuing to read + // would silently cascade garbage into later rows. Fail loudly instead, mirroring + // PojoSerializer.deserialize(), which never catches per-field failures either. + Object raw = fieldDeserializers[i].deserialize(source); + try { + row.setField(i, InternalTypeConverter.toInternal(raw, fieldTypes[i])); + } catch (Exception e) { + // Bytes were already consumed correctly, so the stream is still aligned for + // subsequent fields/rows; only this field's value could not be mapped to its + // table-internal representation. Safe to null just this field and continue. + LOG.warn( + "Failed to convert field '{}' (index {}) value: {}. Setting field to null.", + fieldNames[i], + i, + e.getMessage()); + row.setField(i, null); + } + } + return row; + } + + // ------------------------------------------------------------------------- + // TypeSerializer boilerplate — copy/snapshot operations not needed for + // schema-extraction use cases but required by the interface. + // ------------------------------------------------------------------------- + + @Override + public boolean isImmutableType() { + return false; + } + + @Override + public TypeSerializer duplicate() { + return this; + } + + @Override + public RowData createInstance() { + return new GenericRowData(numFields); + } + + @Override + public RowData copy(RowData from) { + return from; + } + + @Override + public RowData copy(RowData from, RowData reuse) { + return from; + } + + @Override + public int getLength() { + return -1; + } + + @Override + public void serialize(RowData record, DataOutputView target) throws IOException { + throw new UnsupportedOperationException( + "PojoToRowDataDeserializer is read-only; serialization is not supported."); + } + + @Override + public void copy(DataInputView source, DataOutputView target) throws IOException { + throw new UnsupportedOperationException( + "PojoToRowDataDeserializer is read-only; copy is not supported."); + } + + @Override + public boolean equals(Object obj) { + if (this == obj) { + return true; + } + if (!(obj instanceof PojoToRowDataDeserializer)) { + return false; + } + PojoToRowDataDeserializer other = (PojoToRowDataDeserializer) obj; + return numFields == other.numFields + && Arrays.equals(fieldDeserializers, other.fieldDeserializers) + && Arrays.equals(fieldTypes, other.fieldTypes) + && Arrays.equals(fieldNames, other.fieldNames) + && registeredSubclassDeserializers.equals(other.registeredSubclassDeserializers); + } + + @Override + public int hashCode() { + int result = numFields; + result = 31 * result + Arrays.hashCode(fieldDeserializers); + result = 31 * result + Arrays.hashCode(fieldTypes); + result = 31 * result + Arrays.hashCode(fieldNames); + result = 31 * result + registeredSubclassDeserializers.hashCode(); + return result; + } + + @Override + public TypeSerializerSnapshot snapshotConfiguration() { + return new PojoDeserializerCompatibilitySnapshot<>(this); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/operator/KeyedStateReaderOperator.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/operator/KeyedStateReaderOperator.java index 20ab18b3f22b34..ba831f0c52fdd2 100644 --- a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/operator/KeyedStateReaderOperator.java +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/operator/KeyedStateReaderOperator.java @@ -19,28 +19,20 @@ package org.apache.flink.state.api.input.operator; import org.apache.flink.annotation.Internal; -import org.apache.flink.api.common.state.ListState; -import org.apache.flink.api.common.state.ListStateDescriptor; import org.apache.flink.api.common.state.StateDescriptor; import org.apache.flink.api.common.typeinfo.TypeInformation; -import org.apache.flink.api.common.typeinfo.Types; -import org.apache.flink.api.common.typeutils.base.StringSerializer; import org.apache.flink.api.java.tuple.Tuple2; -import org.apache.flink.runtime.state.KeyedStateBackend; +import org.apache.flink.api.java.tuple.Tuple3; import org.apache.flink.runtime.state.VoidNamespace; import org.apache.flink.runtime.state.VoidNamespaceSerializer; import org.apache.flink.state.api.functions.KeyedStateReaderFunction; import org.apache.flink.state.api.input.MultiStateKeyIterator; import org.apache.flink.state.api.runtime.SavepointRuntimeContext; -import org.apache.flink.streaming.api.operators.InternalTimerService; import org.apache.flink.util.CloseableIterator; import org.apache.flink.util.Collector; -import java.util.Collections; import java.util.List; import java.util.Set; -import java.util.stream.Collectors; -import java.util.stream.StreamSupport; /** * A {@link StateReaderOperator} for executing a {@link KeyedStateReaderFunction}. @@ -54,7 +46,7 @@ public class KeyedStateReaderOperator private static final String USER_TIMERS_NAME = "user-timers"; - private transient Context context; + private transient Context context; public KeyedStateReaderOperator( KeyedStateReaderFunction function, TypeInformation keyType) { @@ -65,9 +57,12 @@ public KeyedStateReaderOperator( public void open() throws Exception { super.open(); - InternalTimerService timerService = - getInternalTimerService(USER_TIMERS_NAME); - context = new Context<>(getKeyedStateBackend(), timerService); + TimerRegistration timerRegistration = + registerTimers( + getInternalTimerService(USER_TIMERS_NAME), + USER_TIMERS_NAME, + VoidNamespace.INSTANCE::equals); + context = new Context(timerRegistration); } @Override @@ -77,7 +72,7 @@ public void processElement(KEY key, VoidNamespace namespace, Collector out) } @Override - public CloseableIterator> getKeysAndNamespaces( + public CloseableIterator> getKeysAndNamespaces( SavepointRuntimeContext ctx) throws Exception { ctx.disableStateRegistration(); List> stateDescriptors = ctx.getStateDescriptors(); @@ -86,74 +81,31 @@ public CloseableIterator> getKeysAndNamespaces( return new NamespaceDecorator<>(keys); } - private static class Context implements KeyedStateReaderFunction.Context { - - private static final String EVENT_TIMER_STATE = "event-time-timers"; - - private static final String PROC_TIMER_STATE = "proc-time-timers"; - - ListState eventTimers; - - ListState procTimers; - - private Context( - KeyedStateBackend keyedStateBackend, - InternalTimerService timerService) - throws Exception { - eventTimers = - keyedStateBackend.getPartitionedState( - USER_TIMERS_NAME, - StringSerializer.INSTANCE, - new ListStateDescriptor<>(EVENT_TIMER_STATE, Types.LONG)); - - timerService.forEachEventTimeTimer( - (namespace, timer) -> { - if (namespace.equals(VoidNamespace.INSTANCE)) { - eventTimers.add(timer); - } - }); - - procTimers = - keyedStateBackend.getPartitionedState( - USER_TIMERS_NAME, - StringSerializer.INSTANCE, - new ListStateDescriptor<>(PROC_TIMER_STATE, Types.LONG)); - - timerService.forEachProcessingTimeTimer( - (namespace, timer) -> { - if (namespace.equals(VoidNamespace.INSTANCE)) { - procTimers.add(timer); - } - }); + private static class Context implements KeyedStateReaderFunction.Context { + + private final TimerRegistration timerRegistration; + + private Context(TimerRegistration timerRegistration) { + this.timerRegistration = timerRegistration; } @Override public Set registeredEventTimeTimers() throws Exception { - Iterable timers = eventTimers.get(); - if (timers == null) { - return Collections.emptySet(); - } - - return StreamSupport.stream(timers.spliterator(), false).collect(Collectors.toSet()); + return timerRegistration.registeredEventTimeTimers(); } @Override public Set registeredProcessingTimeTimers() throws Exception { - Iterable timers = procTimers.get(); - if (timers == null) { - return Collections.emptySet(); - } - - return StreamSupport.stream(timers.spliterator(), false).collect(Collectors.toSet()); + return timerRegistration.registeredProcessingTimeTimers(); } } private static class NamespaceDecorator - implements CloseableIterator> { + implements CloseableIterator> { - private final CloseableIterator keys; + private final CloseableIterator> keys; - private NamespaceDecorator(CloseableIterator keys) { + private NamespaceDecorator(CloseableIterator> keys) { this.keys = keys; } @@ -163,9 +115,9 @@ public boolean hasNext() { } @Override - public Tuple2 next() { - KEY key = keys.next(); - return Tuple2.of(key, VoidNamespace.INSTANCE); + public Tuple3 next() { + Tuple2 keyAndKeyGroup = keys.next(); + return Tuple3.of(keyAndKeyGroup.f0, VoidNamespace.INSTANCE, keyAndKeyGroup.f1); } @Override diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/operator/StateReaderOperator.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/operator/StateReaderOperator.java index 2fbb635daa0f13..68cc4e3a3373a9 100644 --- a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/operator/StateReaderOperator.java +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/operator/StateReaderOperator.java @@ -23,9 +23,13 @@ import org.apache.flink.api.common.functions.Function; import org.apache.flink.api.common.functions.SerializerFactory; import org.apache.flink.api.common.functions.util.FunctionUtils; +import org.apache.flink.api.common.state.ListState; +import org.apache.flink.api.common.state.ListStateDescriptor; import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.common.typeinfo.Types; import org.apache.flink.api.common.typeutils.TypeSerializer; -import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.api.common.typeutils.base.StringSerializer; +import org.apache.flink.api.java.tuple.Tuple3; import org.apache.flink.runtime.state.KeyedStateBackend; import org.apache.flink.state.api.runtime.SavepointRuntimeContext; import org.apache.flink.state.api.runtime.VoidTriggerable; @@ -37,6 +41,11 @@ import org.apache.flink.util.Preconditions; import java.io.Serializable; +import java.util.Collections; +import java.util.Set; +import java.util.function.Predicate; +import java.util.stream.Collectors; +import java.util.stream.StreamSupport; /** * Base class for executing functions that read keyed state. @@ -52,6 +61,17 @@ public abstract class StateReaderOperator private static final long serialVersionUID = 1L; + /** + * Sentinel for {@link #setCurrentKeyAndKeyGroup} meaning the key-group is unknown, in which + * case it is derived from {@code key.hashCode()} instead. Used by backend lookup APIs that do + * not expose the physically stored key-group. + */ + public static final int UNKNOWN_KEY_GROUP = -1; + + private static final String EVENT_TIMER_STATE = "event-time-timers"; + + private static final String PROC_TIMER_STATE = "proc-time-timers"; + protected final F function; private final TypeInformation keyType; @@ -80,7 +100,7 @@ protected StateReaderOperator( public abstract void processElement(KEY key, N namespace, Collector out) throws Exception; - public abstract CloseableIterator> getKeysAndNamespaces( + public abstract CloseableIterator> getKeysAndNamespaces( SavepointRuntimeContext ctx) throws Exception; public final void setup( @@ -102,6 +122,77 @@ protected final InternalTimerService getInternalTimerService(String name) { name, keySerializer, namespaceSerializer, VoidTriggerable.instance()); } + /** + * Snapshots the timers currently registered in {@code timerService} into keyed list state under + * {@code timerStateName} and returns a {@link TimerRegistration} exposing them. + * + *

Only timers whose namespace matches {@code namespaceFilter} are snapshotted: namespaced + * readers (e.g. window state) and non-namespaced ones (plain keyed state, registered under + * {@link org.apache.flink.runtime.state.VoidNamespace}) differ in which timers belong to the + * state being read. + */ + protected final TimerRegistration registerTimers( + InternalTimerService timerService, + String timerStateName, + Predicate namespaceFilter) + throws Exception { + ListState eventTimers = + keyedStateBackend.getPartitionedState( + timerStateName, + StringSerializer.INSTANCE, + new ListStateDescriptor<>(EVENT_TIMER_STATE, Types.LONG)); + + timerService.forEachEventTimeTimer( + (namespace, timer) -> { + if (namespaceFilter.test(namespace)) { + eventTimers.add(timer); + } + }); + + ListState procTimers = + keyedStateBackend.getPartitionedState( + timerStateName, + StringSerializer.INSTANCE, + new ListStateDescriptor<>(PROC_TIMER_STATE, Types.LONG)); + + timerService.forEachProcessingTimeTimer( + (namespace, timer) -> { + if (namespaceFilter.test(namespace)) { + procTimers.add(timer); + } + }); + + return new TimerRegistration(eventTimers, procTimers); + } + + /** Read-only view over the timers snapshotted by {@link #registerTimers}. */ + protected static final class TimerRegistration { + + private final ListState eventTimers; + + private final ListState procTimers; + + private TimerRegistration(ListState eventTimers, ListState procTimers) { + this.eventTimers = eventTimers; + this.procTimers = procTimers; + } + + public Set registeredEventTimeTimers() throws Exception { + return toSet(eventTimers.get()); + } + + public Set registeredProcessingTimeTimers() throws Exception { + return toSet(procTimers.get()); + } + + private static Set toSet(Iterable timers) { + if (timers == null) { + return Collections.emptySet(); + } + return StreamSupport.stream(timers.spliterator(), false).collect(Collectors.toSet()); + } + } + public void open() throws Exception { FunctionUtils.openFunction(function, DefaultOpenContext.INSTANCE); } @@ -132,6 +223,16 @@ public final void setCurrentKey(Object key) { keyedStateBackend.setCurrentKey((KEY) key); } + /** Restores the reading context for the given key. See {@link #UNKNOWN_KEY_GROUP}. */ + @SuppressWarnings("unchecked") + public final void setCurrentKeyAndKeyGroup(Object key, int keyGroup) { + if (keyGroup == UNKNOWN_KEY_GROUP) { + keyedStateBackend.setCurrentKey((KEY) key); + } else { + keyedStateBackend.setCurrentKeyAndKeyGroup((KEY) key, keyGroup); + } + } + @Override public final Object getCurrentKey() { return keyedStateBackend.getCurrentKey(); diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/operator/WindowReaderOperator.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/operator/WindowReaderOperator.java index 2f09cacb79ce2f..262ff31eb2bc58 100644 --- a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/operator/WindowReaderOperator.java +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/operator/WindowReaderOperator.java @@ -32,16 +32,13 @@ import org.apache.flink.api.common.state.State; import org.apache.flink.api.common.state.StateDescriptor; import org.apache.flink.api.common.typeinfo.TypeInformation; -import org.apache.flink.api.common.typeinfo.Types; import org.apache.flink.api.common.typeutils.TypeSerializer; -import org.apache.flink.api.common.typeutils.base.StringSerializer; -import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.api.java.tuple.Tuple3; import org.apache.flink.runtime.state.DefaultKeyedStateStore; import org.apache.flink.runtime.state.KeyedStateBackend; import org.apache.flink.state.api.functions.WindowReaderFunction; import org.apache.flink.state.api.input.operator.window.WindowContents; import org.apache.flink.state.api.runtime.SavepointRuntimeContext; -import org.apache.flink.streaming.api.operators.InternalTimerService; import org.apache.flink.streaming.api.windowing.windows.Window; import org.apache.flink.streaming.runtime.streamrecord.StreamElementSerializer; import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; @@ -49,13 +46,10 @@ import org.apache.flink.util.Collector; import org.apache.flink.util.Preconditions; -import java.util.Collections; import java.util.Iterator; import java.util.List; import java.util.Set; -import java.util.stream.Collectors; import java.util.stream.Stream; -import java.util.stream.StreamSupport; /** * A {@link StateReaderOperator} for reading {@code WindowOperator} state. @@ -163,7 +157,12 @@ private WindowReaderOperator( public void open() throws Exception { super.open(); - ctx = new Context(getKeyedStateBackend(), getInternalTimerService(WINDOW_TIMER_NAME)); + ctx = + new Context( + registerTimers( + getInternalTimerService(WINDOW_TIMER_NAME), + WINDOW_TIMER_NAME, + namespace -> true)); } @Override @@ -176,57 +175,32 @@ public void processElement(KEY key, W namespace, Collector out) throws Exce } @Override - public CloseableIterator> getKeysAndNamespaces(SavepointRuntimeContext ctx) - throws Exception { - Stream> keysAndWindows = - getKeyedStateBackend().getKeysAndNamespaces(descriptor.getName()); + public CloseableIterator> getKeysAndNamespaces( + SavepointRuntimeContext ctx) throws Exception { + // getKeysAndNamespaces(String) does not expose the physically stored key-group. + Stream> keysAndWindows = + getKeyedStateBackend() + .getKeysAndNamespaces(descriptor.getName()) + .map(t -> Tuple3.of(t.f0, t.f1, UNKNOWN_KEY_GROUP)); return new IteratorWithRemove<>(keysAndWindows); } private class Context implements WindowReaderFunction.Context { - private static final String EVENT_TIMER_STATE = "event-time-timers"; - - private static final String PROC_TIMER_STATE = "proc-time-timers"; - W window; final PerWindowKeyedStateStore perWindowKeyedStateStore; final DefaultKeyedStateStore keyedStateStore; - ListState eventTimers; + final TimerRegistration timerRegistration; - ListState procTimers; - - private Context( - KeyedStateBackend keyedStateBackend, InternalTimerService timerService) - throws Exception { + private Context(TimerRegistration timerRegistration) { + KeyedStateBackend keyedStateBackend = getKeyedStateBackend(); keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, getSerializerFactory()); perWindowKeyedStateStore = new PerWindowKeyedStateStore(keyedStateBackend); - - eventTimers = - keyedStateBackend.getPartitionedState( - WINDOW_TIMER_NAME, - StringSerializer.INSTANCE, - new ListStateDescriptor<>(EVENT_TIMER_STATE, Types.LONG)); - - timerService.forEachEventTimeTimer( - (namespace, timer) -> { - eventTimers.add(timer); - }); - - procTimers = - keyedStateBackend.getPartitionedState( - WINDOW_TIMER_NAME, - StringSerializer.INSTANCE, - new ListStateDescriptor<>(PROC_TIMER_STATE, Types.LONG)); - - timerService.forEachProcessingTimeTimer( - (namespace, timer) -> { - procTimers.add(timer); - }); + this.timerRegistration = timerRegistration; } @Override @@ -257,22 +231,12 @@ public KeyedStateStore globalState() { @Override public Set registeredEventTimeTimers() throws Exception { - Iterable timers = eventTimers.get(); - if (timers == null) { - return Collections.emptySet(); - } - - return StreamSupport.stream(timers.spliterator(), false).collect(Collectors.toSet()); + return timerRegistration.registeredEventTimeTimers(); } @Override public Set registeredProcessingTimeTimers() throws Exception { - Iterable timers = procTimers.get(); - if (timers == null) { - return Collections.emptySet(); - } - - return StreamSupport.stream(timers.spliterator(), false).collect(Collectors.toSet()); + return timerRegistration.registeredProcessingTimeTimers(); } } diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/runtime/SavepointLoader.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/runtime/SavepointLoader.java index e4ea285c0e876b..b3ec1a0ce3cc19 100644 --- a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/runtime/SavepointLoader.java +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/runtime/SavepointLoader.java @@ -19,12 +19,15 @@ package org.apache.flink.state.api.runtime; import org.apache.flink.annotation.Internal; +import org.apache.flink.api.common.typeutils.CustomRestoreSerializerFactory; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; import org.apache.flink.core.fs.FSDataInputStream; import org.apache.flink.core.memory.DataInputViewStreamWrapper; import org.apache.flink.runtime.checkpoint.Checkpoints; import org.apache.flink.runtime.checkpoint.OperatorState; import org.apache.flink.runtime.checkpoint.metadata.CheckpointMetadata; import org.apache.flink.runtime.state.CompletedCheckpointStorageLocation; +import org.apache.flink.runtime.state.IncrementalKeyedStateHandle; import org.apache.flink.runtime.state.KeyGroupsStateHandle; import org.apache.flink.runtime.state.KeyedBackendSerializationProxy; import org.apache.flink.runtime.state.KeyedStateHandle; @@ -32,6 +35,9 @@ import org.apache.flink.runtime.state.filesystem.AbstractFsCheckpointStorageAccess; import org.apache.flink.runtime.state.metainfo.StateMetaInfoSnapshot; import org.apache.flink.state.api.OperatorIdentifier; +import org.apache.flink.state.api.input.deserializer.MissingClassSerializerFactory; + +import javax.annotation.Nullable; import java.io.DataInputStream; import java.io.IOException; @@ -44,6 +50,29 @@ public final class SavepointLoader { private SavepointLoader() {} + /** + * Operator-level metadata loaded in a single I/O pass: the per-state serializer snapshots and + * the backend key serializer snapshot. + */ + public static final class OperatorStateMetadata { + + /** Per-state serializer snapshots, keyed by state name. */ + public final Map stateSnapshots; + + /** + * The key serializer snapshot shared by all states in the operator's keyed backend, or + * {@code null} if none is available. + */ + @Nullable public final TypeSerializerSnapshot keySerializerSnapshot; + + OperatorStateMetadata( + Map stateSnapshots, + @Nullable TypeSerializerSnapshot keySerializerSnapshot) { + this.stateSnapshots = stateSnapshots; + this.keySerializerSnapshot = keySerializerSnapshot; + } + } + /** * Takes the given string (representing a pointer to a checkpoint) and resolves it to a file * status for the checkpoint's metadata file. @@ -78,6 +107,20 @@ public static CheckpointMetadata loadSavepointMetadata(String savepointPath) */ public static Map loadOperatorStateMetadata( String savepointPath, OperatorIdentifier operatorIdentifier) throws IOException { + return loadOperatorMetadata(savepointPath, operatorIdentifier).stateSnapshots; + } + + /** + * Loads both the per-state serializer snapshots and the backend key serializer snapshot for an + * operator in a single I/O operation. + * + * @param savepointPath Path to the savepoint directory + * @param operatorIdentifier Operator UID or hash + * @return combined operator metadata + * @throws IOException If reading fails + */ + public static OperatorStateMetadata loadOperatorMetadata( + String savepointPath, OperatorIdentifier operatorIdentifier) throws IOException { CheckpointMetadata checkpointMetadata = loadSavepointMetadata(savepointPath); @@ -107,16 +150,26 @@ public static Map loadOperatorStateMetadata( + operatorIdentifier)); KeyedBackendSerializationProxy proxy = readSerializationProxy(keyedStateHandle); - return proxy.getStateMetaInfoSnapshots().stream() - .collect(Collectors.toMap(StateMetaInfoSnapshot::getName, Function.identity())); + Map stateSnapshots = + proxy.getStateMetaInfoSnapshots().stream() + .collect( + Collectors.toMap( + StateMetaInfoSnapshot::getName, Function.identity())); + return new OperatorStateMetadata(stateSnapshots, proxy.getKeySerializerSnapshot()); } private static KeyedBackendSerializationProxy readSerializationProxy( KeyedStateHandle stateHandle) throws IOException { + // KeyGroupsStateHandle (heap/HashMapStateBackend) is itself a StreamStateHandle whose + // stream starts with the metadata header. IncrementalKeyedStateHandle (RocksDB) instead + // keeps the metadata in a separate handle, exposed via getMetaDataStateHandle(). StreamStateHandle streamStateHandle; if (stateHandle instanceof KeyGroupsStateHandle) { streamStateHandle = ((KeyGroupsStateHandle) stateHandle).getDelegateStateHandle(); + } else if (stateHandle instanceof IncrementalKeyedStateHandle) { + streamStateHandle = + ((IncrementalKeyedStateHandle) stateHandle).getMetaDataStateHandle(); } else { throw new IllegalArgumentException( "Unsupported KeyedStateHandle type: " + stateHandle.getClass()); @@ -128,6 +181,7 @@ private static KeyedBackendSerializationProxy readSerializationProxy( KeyedBackendSerializationProxy proxy = new KeyedBackendSerializationProxy<>( Thread.currentThread().getContextClassLoader()); + CustomRestoreSerializerFactory.set(MissingClassSerializerFactory::create); proxy.read(inputView); return proxy; diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/AvroStateUtils.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/AvroStateUtils.java new file mode 100644 index 00000000000000..df5de4c2fc9163 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/AvroStateUtils.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.schema; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; +import org.apache.flink.formats.avro.typeutils.AvroSchemaConverter; +import org.apache.flink.formats.avro.typeutils.AvroSerializer; +import org.apache.flink.formats.avro.typeutils.AvroSerializerSnapshot; +import org.apache.flink.table.types.logical.LogicalType; + +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericRecord; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +/** + * All Avro-specific logic used by the State Processing API's schema-based table access, gathered in + * one place so that the classes which dispatch to it never need to mention an Avro type themselves. + * + *

flink-avro is an optional dependency of this module (see the module {@code pom.xml}). A class + * that mentions an Avro type directly - even in an {@code instanceof} branch that is never taken - + * makes the JVM try to resolve that type the moment the check is reached for *any* value, throwing + * {@code NoClassDefFoundError} for callers who never use Avro at all. Callers must therefore select + * the Avro-specific branch by class/interface name (see {@link + * #AVRO_SERIALIZER_SNAPSHOT_CLASS_NAME} and {@link #isGenericRecord}) and only then delegate here; + * this class is consequently only ever loaded once that name comparison has already confirmed Avro + * is genuinely on the classpath. + */ +@Internal +public final class AvroStateUtils { + + /** Fully-qualified name of {@code AvroSerializerSnapshot}, for dispatch by class name. */ + public static final String AVRO_SERIALIZER_SNAPSHOT_CLASS_NAME = + "org.apache.flink.formats.avro.typeutils.AvroSerializerSnapshot"; + + private static final String GENERIC_RECORD_CLASS_NAME = "org.apache.avro.generic.GenericRecord"; + + /** + * Memoizes {@link #isGenericRecord}: it walks a class's full interface hierarchy, so callers + * that check the same class repeatedly (e.g. once per field of every row of the same type) + * would otherwise repeat that walk every time. A plain static map is safe here since the answer + * is a pure function of the {@code Class} object - it never changes for a given class, and is + * shared happily across every {@link AvroStateUtils} caller and instance. + */ + private static final Map, Boolean> GENERIC_RECORD_CACHE = new ConcurrentHashMap<>(); + + private AvroStateUtils() {} + + /** + * Builds the fallback serializer for an {@code AvroSerializerSnapshot} with a missing class. + */ + public static TypeSerializer createFallbackSerializer(TypeSerializerSnapshot snapshot) { + Schema schema = ((AvroSerializerSnapshot) snapshot).getSchema(); + return new AvroSerializer<>(GenericRecord.class, schema); + } + + /** + * Converts an {@code AvroSerializerSnapshot}'s embedded writer schema into a {@link + * LogicalType}. + */ + public static LogicalType convertToLogicalType(TypeSerializerSnapshot snapshot) { + // getSchema() returns the writer schema embedded in the snapshot — always present + // regardless of whether the specific record class is on the classpath. + AvroSerializerSnapshot avroSnapshot = (AvroSerializerSnapshot) snapshot; + return AvroSchemaConverter.convertToDataType(avroSnapshot.getSchema().toString()) + .getLogicalType(); + } + + /** + * Returns {@code true} if {@code clazz}, or any class/interface in its hierarchy, is named + * {@code org.apache.avro.generic.GenericRecord}. Safe to call even when {@code GenericRecord} + * itself is not on the classpath: {@code clazz} could only have been loaded and instantiated if + * all interfaces it declares were already resolved, so walking {@link Class#getInterfaces()} + * never triggers a fresh classload of {@code GenericRecord}. + */ + public static boolean isGenericRecord(Class clazz) { + return GENERIC_RECORD_CACHE.computeIfAbsent(clazz, AvroStateUtils::computeIsGenericRecord); + } + + /** + * Does the actual interface-hierarchy walk for {@link #isGenericRecord}. Recurses into itself + * rather than back into {@link #isGenericRecord}: {@code ConcurrentHashMap.computeIfAbsent} + * forbids its mapping function from calling back into the same map - even for a different key - + * and will throw {@code IllegalStateException("Recursive update")} if it does. + */ + private static boolean computeIsGenericRecord(Class clazz) { + for (Class current = clazz; current != null; current = current.getSuperclass()) { + for (Class iface : current.getInterfaces()) { + if (iface.getName().equals(GENERIC_RECORD_CLASS_NAME) + || computeIsGenericRecord(iface)) { + return true; + } + } + } + return false; + } + + /** Reads a field from an Avro {@code GenericRecord}. */ + public static Object getGenericRecordField(Object record, String fieldName) { + return ((GenericRecord) record).get(fieldName); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/KeyedStateSchemaInfo.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/KeyedStateSchemaInfo.java new file mode 100644 index 00000000000000..014a89dc018668 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/KeyedStateSchemaInfo.java @@ -0,0 +1,90 @@ +/* + * 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.schema; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.state.table.SavepointConnectorOptions; +import org.apache.flink.table.types.logical.LogicalType; + +import javax.annotation.Nullable; + +import java.util.LinkedHashMap; + +/** + * Schema information for all keyed states of a single operator, extracted from a savepoint without + * requiring user POJO classes on the classpath. + */ +@Internal +public final class KeyedStateSchemaInfo { + + /** The logical type of the keyed state backend key (e.g., BigIntType for Long keys). */ + public final LogicalType keyType; + + /** + * Ordered map of registered state names to their entry information. Ordered by the registration + * order found in the savepoint. + */ + public final LinkedHashMap stateSchemas; + + public KeyedStateSchemaInfo( + LogicalType keyType, LinkedHashMap stateSchemas) { + this.keyType = keyType; + this.stateSchemas = stateSchemas; + } + + /** Schema information for one keyed state entry. */ + public static final class StateEntryInfo { + + /** VALUE, LIST, or MAP. */ + public final SavepointConnectorOptions.StateType stateType; + + /** + * The SQL column logical type. + * + *

    + *
  • VALUE<Long>: BigIntType + *
  • VALUE<POJO>: RowType (field names + types from the serializer snapshot) + *
  • LIST<Long>: ArrayType(BigIntType) + *
  • MAP<Long,Long>: MapType(BigIntType, BigIntType) + *
+ */ + public final LogicalType logicalType; + + /** + * The resolved {@link LogicalType} of the state's namespace, or {@code null} if the state + * is plain per-key state (registered under {@code VoidNamespace}). Non-null means the state + * is scoped by some other namespace (e.g. a window), resolved generically by {@link + * SerializerSnapshotToLogicalTypeConverter} rather than as a fixed {@code + * TimeWindow}-shaped type. + * + *

Named "window" rather than "namespace" because this is the user-facing, + * post-conversion form; see {@link StateSchemaInfo} for the raw/resolved naming convention. + */ + @Nullable public final LogicalType windowLogicalType; + + public StateEntryInfo( + SavepointConnectorOptions.StateType stateType, + LogicalType logicalType, + @Nullable LogicalType windowLogicalType) { + this.stateType = stateType; + this.logicalType = logicalType; + this.windowLogicalType = windowLogicalType; + } + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/SerializerSnapshotToLogicalTypeConverter.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/SerializerSnapshotToLogicalTypeConverter.java new file mode 100644 index 00000000000000..b689cbe9656a72 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/SerializerSnapshotToLogicalTypeConverter.java @@ -0,0 +1,230 @@ +/* + * 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.schema; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.api.common.typeutils.CompositeTypeSerializerSnapshot; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; +import org.apache.flink.api.common.typeutils.base.BooleanSerializer; +import org.apache.flink.api.common.typeutils.base.ByteSerializer; +import org.apache.flink.api.common.typeutils.base.CharSerializer; +import org.apache.flink.api.common.typeutils.base.DoubleSerializer; +import org.apache.flink.api.common.typeutils.base.EnumSerializer; +import org.apache.flink.api.common.typeutils.base.FloatSerializer; +import org.apache.flink.api.common.typeutils.base.IntSerializer; +import org.apache.flink.api.common.typeutils.base.ListSerializerSnapshot; +import org.apache.flink.api.common.typeutils.base.LongSerializer; +import org.apache.flink.api.common.typeutils.base.MapSerializerSnapshot; +import org.apache.flink.api.common.typeutils.base.ShortSerializer; +import org.apache.flink.api.common.typeutils.base.StringSerializer; +import org.apache.flink.api.java.typeutils.runtime.NullableSerializer.NullableSerializerSnapshot; +import org.apache.flink.api.java.typeutils.runtime.PojoSerializerSnapshot; +import org.apache.flink.api.java.typeutils.runtime.TupleSerializerSnapshot; +import org.apache.flink.streaming.api.windowing.windows.GlobalWindow; +import org.apache.flink.streaming.api.windowing.windows.TimeWindow; +import org.apache.flink.table.runtime.typeutils.RowDataSerializer.RowDataSerializerSnapshot; +import org.apache.flink.table.types.logical.ArrayType; +import org.apache.flink.table.types.logical.BigIntType; +import org.apache.flink.table.types.logical.BooleanType; +import org.apache.flink.table.types.logical.CharType; +import org.apache.flink.table.types.logical.DoubleType; +import org.apache.flink.table.types.logical.FloatType; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.MapType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.table.types.logical.SmallIntType; +import org.apache.flink.table.types.logical.TimestampType; +import org.apache.flink.table.types.logical.TinyIntType; +import org.apache.flink.table.types.logical.VarBinaryType; +import org.apache.flink.table.types.logical.VarCharType; + +import java.util.AbstractMap; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * Converts a {@link TypeSerializerSnapshot} tree into a Flink SQL {@link LogicalType}. + * + *

This is a pure function — no I/O, no class loading beyond what the snapshot itself already + * contains. Field names are extracted from {@link PojoSerializerSnapshot} entries; all other + * composite types use positional names {@code f0, f1, …}. + * + *

flink-avro is an optional dependency of this module (see the module {@code pom.xml}), so this + * class must never mention an Avro type directly: doing so - even in an {@code instanceof} branch - + * would make the JVM try to resolve that type the moment this method is reached for *any* + * unrecognized snapshot (e.g. a plain {@code ListSerializerSnapshot}), throwing {@code + * NoClassDefFoundError} for callers who never use Avro at all. The Avro-specific branch is instead + * selected by class name and delegated to {@link AvroStateUtils}, which is only ever loaded once + * that name comparison has already confirmed Avro is genuinely on the classpath. + */ +@Internal +public final class SerializerSnapshotToLogicalTypeConverter { + + /** Type used for snapshots we cannot describe any further, e.g. a missing nested snapshot. */ + private static final LogicalType OPAQUE_TYPE = + new VarBinaryType(true, VarBinaryType.MAX_LENGTH); + + private static final TypeSerializerSnapshot[] NO_NESTED_SNAPSHOTS = + new TypeSerializerSnapshot[0]; + + /** + * Snapshot classes that always map to the same {@link LogicalType}. Matched by exact class + * because every one of them is final; {@link LogicalType} instances are immutable and therefore + * safe to share. + * + *

The window entries are not namespace-specific: they fire identically if a window type is + * ever used as an ordinary value rather than as a namespace. + */ + private static final Map, LogicalType> FIXED_TYPES = createFixedTypes(); + + private SerializerSnapshotToLogicalTypeConverter() {} + + private static Map, LogicalType> createFixedTypes() { + Map, LogicalType> types = new HashMap<>(); + types.put(IntSerializer.IntSerializerSnapshot.class, new IntType(false)); + types.put(LongSerializer.LongSerializerSnapshot.class, new BigIntType(false)); + types.put(FloatSerializer.FloatSerializerSnapshot.class, new FloatType(false)); + types.put(DoubleSerializer.DoubleSerializerSnapshot.class, new DoubleType(false)); + types.put(BooleanSerializer.BooleanSerializerSnapshot.class, new BooleanType(false)); + types.put(ByteSerializer.ByteSerializerSnapshot.class, new TinyIntType(false)); + types.put(ShortSerializer.ShortSerializerSnapshot.class, new SmallIntType(false)); + types.put(CharSerializer.CharSerializerSnapshot.class, new CharType(false, 1)); + types.put( + StringSerializer.StringSerializerSnapshot.class, + new VarCharType(true, VarCharType.MAX_LENGTH)); + types.put( + TimeWindow.Serializer.TimeWindowSerializerSnapshot.class, + new RowType( + false, + List.of( + new RowType.RowField("window_start", new TimestampType(false, 3)), + new RowType.RowField("window_end", new TimestampType(false, 3))))); + types.put( + GlobalWindow.Serializer.GlobalWindowSerializerSnapshot.class, + new RowType(false, List.of())); + return types; + } + + /** + * Converts the given snapshot to a {@link LogicalType}. + * + * @param snapshot the serializer snapshot to convert, may be null + * @return a {@link LogicalType} corresponding to the snapshot type + * @throws UnsupportedOperationException if the snapshot type is not supported for schema-based + * table access + */ + public static LogicalType convert(TypeSerializerSnapshot snapshot) { + if (snapshot == null) { + return OPAQUE_TYPE; + } + + LogicalType fixedType = FIXED_TYPES.get(snapshot.getClass()); + if (fixedType != null) { + return fixedType; + } + + if (snapshot instanceof PojoSerializerSnapshot) { + return convertPojo((PojoSerializerSnapshot) snapshot); + } + if (snapshot instanceof EnumSerializer.EnumSerializerSnapshot) { + return new VarCharType(true, VarCharType.MAX_LENGTH); + } + String avroSnapshotClassName = AvroStateUtils.AVRO_SERIALIZER_SNAPSHOT_CLASS_NAME; + if (avroSnapshotClassName.equals(snapshot.getClass().getName())) { + return AvroStateUtils.convertToLogicalType(snapshot); + } + if (snapshot instanceof ListSerializerSnapshot) { + return new ArrayType(true, convertNested(snapshot, 0)); + } + if (snapshot instanceof MapSerializerSnapshot) { + return new MapType(true, convertNested(snapshot, 0), convertNested(snapshot, 1)); + } + if (snapshot instanceof NullableSerializerSnapshot) { + return convertNested(snapshot, 0).copy(true); + } + if (snapshot instanceof TupleSerializerSnapshot) { + return convertTuple((TupleSerializerSnapshot) snapshot); + } + if (snapshot instanceof RowDataSerializerSnapshot) { + return convertRowData((RowDataSerializerSnapshot) snapshot); + } + + throw new UnsupportedOperationException( + "Cannot extract schema for TypeSerializerSnapshot of type '" + + snapshot.getClass().getName() + + "'. This serializer type is not supported for schema-based table" + + " access."); + } + + private static LogicalType convertPojo(PojoSerializerSnapshot snapshot) { + List>> fieldEntries = + snapshot.getFieldSnapshotEntries(); + List fields = new ArrayList<>(fieldEntries.size()); + for (AbstractMap.SimpleEntry> entry : fieldEntries) { + fields.add(new RowType.RowField(entry.getKey(), convert(entry.getValue()))); + } + return new RowType(true, fields); + } + + private static LogicalType convertTuple(TupleSerializerSnapshot snapshot) { + TypeSerializerSnapshot[] nested = nestedSnapshots(snapshot); + List fields = new ArrayList<>(nested.length); + for (int i = 0; i < nested.length; i++) { + fields.add(new RowType.RowField("f" + i, convert(nested[i]))); + } + return new RowType(true, fields); + } + + /** + * Converts a {@link RowDataSerializerSnapshot} from its {@link + * RowDataSerializerSnapshot#getTypes()} instead of walking its nested field snapshots: a {@link + * LogicalType} is already a complete, self-describing schema (including nested field names), + * unlike the POJO/Avro case where the snapshot tree is the only source of field names. + */ + private static LogicalType convertRowData(RowDataSerializerSnapshot snapshot) { + LogicalType[] types = snapshot.getTypes(); + String[] fieldNames = snapshot.getFieldNames(); + List fields = new ArrayList<>(types.length); + for (int i = 0; i < types.length; i++) { + String fieldName = fieldNames != null ? fieldNames[i] : "f" + i; + fields.add(new RowType.RowField(fieldName, types[i])); + } + return new RowType(true, fields); + } + + private static LogicalType convertNested(TypeSerializerSnapshot snapshot, int index) { + TypeSerializerSnapshot[] nested = nestedSnapshots(snapshot); + return convert(index < nested.length ? nested[index] : null); + } + + private static TypeSerializerSnapshot[] nestedSnapshots(TypeSerializerSnapshot snapshot) { + if (snapshot instanceof CompositeTypeSerializerSnapshot) { + TypeSerializerSnapshot[] nested = + ((CompositeTypeSerializerSnapshot) snapshot) + .getNestedSerializerSnapshots(); + if (nested != null) { + return nested; + } + } + return NO_NESTED_SNAPSHOTS; + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/StateSchemaExtractor.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/StateSchemaExtractor.java new file mode 100644 index 00000000000000..c58c7bf77370b6 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/StateSchemaExtractor.java @@ -0,0 +1,141 @@ +/* + * 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.schema; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.api.common.state.StateDescriptor; +import org.apache.flink.api.common.typeutils.CustomRestoreSerializerFactory; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; +import org.apache.flink.core.memory.DataInputView; +import org.apache.flink.core.memory.DataInputViewStreamWrapper; +import org.apache.flink.runtime.checkpoint.OperatorState; +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState; +import org.apache.flink.runtime.state.IncrementalKeyedStateHandle; +import org.apache.flink.runtime.state.KeyedBackendSerializationProxy; +import org.apache.flink.runtime.state.KeyedStateHandle; +import org.apache.flink.runtime.state.StreamStateHandle; +import org.apache.flink.runtime.state.metainfo.StateMetaInfoSnapshot; +import org.apache.flink.runtime.state.metainfo.StateMetaInfoSnapshot.CommonOptionsKeys; +import org.apache.flink.runtime.state.metainfo.StateMetaInfoSnapshot.CommonSerializerKeys; +import org.apache.flink.state.api.input.deserializer.MissingClassSerializerFactory; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +/** + * Utility for extracting {@link StateSchemaInfo} from a savepoint without instantiating the full + * state backend or requiring user POJO classes on the classpath. + * + *

It reads the {@link KeyedBackendSerializationProxy} header that every heap/RocksDB keyed state + * file starts with. + */ +@Internal +public final class StateSchemaExtractor { + + private StateSchemaExtractor() {} + + /** + * Reads state schema information from the first available keyed state handle in the given + * operator state. + * + *

Returns an empty list rather than throwing when no keyed state handle is found: an + * operator may register only non-keyed (list/union/broadcast) state, in which case it has no + * keyed state to describe. + * + *

The metadata header lives in different places depending on the state backend: heap ({@code + * HashMapStateBackend}) savepoints hand back a {@code KeyGroupsStateHandle}, which is itself a + * {@link StreamStateHandle} starting with the header; RocksDB (incremental or full native) + * snapshots hand back an {@link IncrementalKeyedStateHandle}, whose own data stream starts with + * the SST payload instead, so the header must be read from {@link + * IncrementalKeyedStateHandle#getMetaDataStateHandle()}. + * + * @param operatorState the operator state from a loaded savepoint / checkpoint metadata + * @return list of schema info, one entry per registered state; never null, may be empty + * @throws IOException if the state header cannot be read + */ + public static List extractSchema(OperatorState operatorState) + throws IOException { + + for (OperatorSubtaskState subtask : operatorState.getSubtaskStates().values()) { + for (KeyedStateHandle handle : subtask.getManagedKeyedState()) { + StreamStateHandle metadataHandle = null; + if (handle instanceof IncrementalKeyedStateHandle) { + metadataHandle = + ((IncrementalKeyedStateHandle) handle).getMetaDataStateHandle(); + } else if (handle instanceof StreamStateHandle) { + metadataHandle = (StreamStateHandle) handle; + } + if (metadataHandle != null) { + try (java.io.InputStream stream = metadataHandle.openInputStream()) { + return extractSchema(new DataInputViewStreamWrapper(stream)); + } + } + } + } + return Collections.emptyList(); + } + + /** + * Package-private overload that accepts a {@link DataInputView} directly. Allows unit tests to + * inject pre-built byte arrays without a real filesystem. + */ + static List extractSchema(DataInputView in) throws IOException { + ClassLoader classLoader = Thread.currentThread().getContextClassLoader(); + KeyedBackendSerializationProxy proxy = new KeyedBackendSerializationProxy<>(classLoader); + CustomRestoreSerializerFactory.set(MissingClassSerializerFactory::create); + proxy.read(in); + + TypeSerializerSnapshot keySnapshot = proxy.getKeySerializerSnapshot(); + List result = new ArrayList<>(); + + for (StateMetaInfoSnapshot meta : proxy.getStateMetaInfoSnapshots()) { + String kindStr = meta.getOption(CommonOptionsKeys.KEYED_STATE_TYPE); + StateDescriptor.Type stateKind; + try { + stateKind = StateDescriptor.Type.valueOf(kindStr); + } catch (IllegalArgumentException | NullPointerException e) { + stateKind = StateDescriptor.Type.UNKNOWN; + } + + TypeSerializerSnapshot valueSnapshot = + meta.getTypeSerializerSnapshot(CommonSerializerKeys.VALUE_SERIALIZER); + TypeSerializerSnapshot mapKeySnapshot = + meta.getTypeSerializerSnapshot(CommonSerializerKeys.USER_KEY_SERIALIZER); + TypeSerializerSnapshot namespaceSnapshot = + meta.getTypeSerializerSnapshot(CommonSerializerKeys.NAMESPACE_SERIALIZER); + + if (valueSnapshot == null) { + continue; + } + + result.add( + new StateSchemaInfo( + meta.getName(), + stateKind, + keySnapshot, + valueSnapshot, + mapKeySnapshot, + namespaceSnapshot)); + } + + return result; + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/StateSchemaInfo.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/StateSchemaInfo.java new file mode 100644 index 00000000000000..9031996dc8b3e5 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/StateSchemaInfo.java @@ -0,0 +1,82 @@ +/* + * 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.schema; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.api.common.state.StateDescriptor; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; + +import javax.annotation.Nullable; + +/** + * Carries the schema information extracted from a single keyed state entry in a savepoint, without + * requiring user POJO classes on the classpath. + * + *

Naming convention used throughout this package and {@code state.table}: "namespace" names the + * raw concept as Flink's runtime state backend knows it ({@link #namespaceSnapshot}), "window" the + * same concept once resolved to a user-facing table/column ({@link + * KeyedStateSchemaInfo.StateEntryInfo#windowLogicalType}). + */ +@Internal +public final class StateSchemaInfo { + + /** Name of the state as registered by the operator. */ + public final String stateName; + + /** The kind of state (VALUE, LIST, MAP, etc.). */ + public final StateDescriptor.Type stateKind; + + /** Serializer snapshot for the key type. */ + public final TypeSerializerSnapshot keySnapshot; + + /** + * Serializer snapshot for the state value type. For MAP state this is the value type; use + * {@link #mapKeySnapshot} for the map key type. + */ + public final TypeSerializerSnapshot valueSnapshot; + + /** + * Serializer snapshot for the map key type. Non-null only for {@link StateDescriptor.Type#MAP} + * state. + */ + @Nullable public final TypeSerializerSnapshot mapKeySnapshot; + + /** + * Serializer snapshot for the state's namespace. Plain per-key state (the only kind the + * savepoint/checkpoint table connector can read) is registered with {@code VoidNamespace}; a + * different namespace (e.g. a window) means the state is scoped per-window rather than per-key, + * and cannot be exposed as a flat keyed table. + */ + @Nullable public final TypeSerializerSnapshot namespaceSnapshot; + + public StateSchemaInfo( + String stateName, + StateDescriptor.Type stateKind, + TypeSerializerSnapshot keySnapshot, + TypeSerializerSnapshot valueSnapshot, + @Nullable TypeSerializerSnapshot mapKeySnapshot, + @Nullable TypeSerializerSnapshot namespaceSnapshot) { + this.stateName = stateName; + this.stateKind = stateKind; + this.keySnapshot = keySnapshot; + this.valueSnapshot = valueSnapshot; + this.mapKeySnapshot = mapKeySnapshot; + this.namespaceSnapshot = namespaceSnapshot; + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/SnapshotDiscovery.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/SnapshotDiscovery.java new file mode 100644 index 00000000000000..cc73ce36c3ef4d --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/SnapshotDiscovery.java @@ -0,0 +1,337 @@ +/* + * 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.catalog; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.core.fs.FileStatus; +import org.apache.flink.core.fs.Path; +import org.apache.flink.util.StringUtils; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.time.Instant; +import java.time.ZoneOffset; +import java.time.format.DateTimeFormatter; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; + +/** + * Discovers Flink checkpoints and savepoints within a set of labelled directories. + * + *

Each configured directory is associated with a user-chosen label. By default, database names + * are derived as {@code label/creationTs/relative-path}, where {@code creationTs} is the + * modification time of the snapshot's {@code _metadata} file formatted as {@code + * yyyy-MM-dd'T'HH:mm:ssX} (e.g. {@code 2026-07-22T10:30:45Z}) and {@code relative-path} is the + * verbatim path from the configured directory to the snapshot directory (e.g. {@code + * my-app/2026-07-22T10:30:45Z/savepoint-acce1cedsad} or {@code + * my-app/2026-07-22T10:30:45Z/a1b2c3d4.../chk-3}). The {@code creationTs} segment can be disabled + * via {@code dbNameIncludeTs}, in which case names fall back to {@code label/relative-path}. + */ +@Internal +class SnapshotDiscovery { + + private static final Logger LOG = LoggerFactory.getLogger(SnapshotDiscovery.class); + + private static final String METADATA_FILE_NAME = "_metadata"; + + private static final DateTimeFormatter CREATION_TS_FORMATTER = + DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ssX").withZone(ZoneOffset.UTC); + + private final Map labelToDir; + private final int listingParallelism; + private final boolean dbNameIncludeTs; + + private ExecutorService listingExecutor; + + SnapshotDiscovery( + Map labelToDirPath, int listingParallelism, boolean dbNameIncludeTs) { + this.labelToDir = validateAndConvert(labelToDirPath); + this.listingParallelism = listingParallelism; + this.dbNameIncludeTs = dbNameIncludeTs; + } + + void start() { + listingExecutor = createListingExecutor(listingParallelism); + } + + void stop() { + if (listingExecutor != null) { + listingExecutor.shutdownNow(); + listingExecutor = null; + } + } + + // ------------------------------------------------------------------------- + // Public API + // ------------------------------------------------------------------------- + + /** + * Full BFS scan of all configured directories, listing directories at the same depth + * concurrently (up to {@code listingParallelism} at a time, with automatic backpressure via + * {@link ThreadPoolExecutor.CallerRunsPolicy}). Returns one database name per discovered {@code + * _metadata} file. + * + * @throws IOException if every configured directory fails to scan + */ + List list() throws IOException { + List result = new ArrayList<>(); + Set seen = new LinkedHashSet<>(); + IOException err = null; + boolean allFailed = true; + + for (Map.Entry entry : labelToDir.entrySet()) { + String label = entry.getKey(); + Path dir = entry.getValue(); + try { + for (FileStatus metadataFileStatus : findMetadataFileStatuses(dir)) { + String dbName = buildDatabaseName(label, dir, metadataFileStatus); + if (!seen.add(dbName)) { + LOG.warn("Duplicate database name '{}'. Skipping.", dbName); + continue; + } + result.add(dbName); + } + allFailed = false; + } catch (IOException e) { + LOG.warn("Failed to scan {}: {}", dir, e.getMessage(), e); + err = e; + } + } + + if (allFailed) { + throw new IOException("All configured directories failed to scan", err); + } + return result; + } + + /** + * Reverses {@link #buildDatabaseName} to recover the label and the verbatim relative path from + * {@code dbName}, then verifies the snapshot's {@code _metadata} file with a single {@code + * getFileStatus} call. Returns the snapshot directory path if it exists. + */ + Optional find(String dbName) { + if (StringUtils.isNullOrWhitespaceOnly(dbName)) { + return Optional.empty(); + } + int labelSlash = dbName.indexOf('/'); + if (labelSlash == 0) { + return Optional.empty(); + } + + String label; + String relativePath; + if (dbNameIncludeTs) { + // A name with no '/' has no room for the mandatory creationTs segment. + if (labelSlash < 0) { + return Optional.empty(); + } + label = dbName.substring(0, labelSlash); + String afterLabel = dbName.substring(labelSlash + 1); + int tsSlash = afterLabel.indexOf('/'); + relativePath = tsSlash < 0 ? "" : afterLabel.substring(tsSlash + 1); + } else { + // A name with no '/' means the configured directory itself is the snapshot. + label = labelSlash < 0 ? dbName : dbName.substring(0, labelSlash); + relativePath = labelSlash < 0 ? "" : dbName.substring(labelSlash + 1); + } + + Path dir = labelToDir.get(label); + if (dir == null) { + return Optional.empty(); + } + + Path metadataFile = + relativePath.isEmpty() + ? new Path(dir, METADATA_FILE_NAME) + : new Path(dir, relativePath + "/" + METADATA_FILE_NAME); + try { + dir.getFileSystem().getFileStatus(metadataFile); + return Optional.of(metadataFile.getParent().toString()); + } catch (IOException e) { + return Optional.empty(); + } + } + + // ------------------------------------------------------------------------- + // Full BFS scan + // ------------------------------------------------------------------------- + + private List findMetadataFileStatuses(Path directory) throws IOException { + List metadataFiles = new ArrayList<>(); + List currentLevel = Collections.singletonList(directory); + boolean allFailed = true; + IOException err = null; + + while (!currentLevel.isEmpty()) { + List> futures = new ArrayList<>(currentLevel.size()); + for (Path dir : currentLevel) { + futures.add(listingExecutor.submit(() -> listDirectory(dir))); + } + + List nextLevel = new ArrayList<>(); + for (int i = 0; i < futures.size(); i++) { + FileStatus[] statuses = null; + try { + statuses = getResult(futures.get(i)); + } catch (IOException e) { + LOG.warn("Failed to list {}: {}", currentLevel.get(i), e.getMessage()); + err = e; + } + if (statuses == null) { + continue; + } + allFailed = false; + for (FileStatus status : statuses) { + if (status.isDir()) { + nextLevel.add(status.getPath()); + } else if (METADATA_FILE_NAME.equals(status.getPath().getName())) { + metadataFiles.add(status); + } + } + } + currentLevel = nextLevel; + } + + if (allFailed) { + throw new IOException("All directory listings failed under: " + directory, err); + } + return metadataFiles; + } + + private FileStatus[] listDirectory(Path dir) throws IOException { + FileStatus[] result = dir.getFileSystem().listStatus(dir); + if (result == null) { + throw new IOException( + "Cannot list directory (path does not exist or is not a directory): " + dir); + } + return result; + } + + private static FileStatus[] getResult(Future future) throws IOException { + try { + return future.get(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return null; + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof IOException) { + throw (IOException) cause; + } + throw new IOException("Directory listing failed", cause); + } + } + + // ------------------------------------------------------------------------- + // Database name derivation + // ------------------------------------------------------------------------- + + private String buildDatabaseName(String label, Path configuredDir, FileStatus metadataFile) { + Path snapshotDir = metadataFile.getPath().getParent(); + String configuredPath = configuredDir.toUri().getPath(); + String snapshotPath = snapshotDir.toUri().getPath(); + + String relative = snapshotPath.substring(configuredPath.length()); + if (relative.startsWith("/")) { + relative = relative.substring(1); + } + + StringBuilder dbName = new StringBuilder(label); + if (dbNameIncludeTs) { + Instant creationTs = Instant.ofEpochMilli(metadataFile.getModificationTime()); + dbName.append('/').append(CREATION_TS_FORMATTER.format(creationTs)); + } + if (!relative.isEmpty()) { + dbName.append('/').append(relative); + } + return dbName.toString(); + } + + // ------------------------------------------------------------------------- + // Construction-time helpers + // ------------------------------------------------------------------------- + + /** + * Converts the configured label → directory paths, rejecting empty configurations, directories + * assigned to more than one label, and directories nested inside one another (which would + * discover the same snapshots under multiple labels). + */ + private static Map validateAndConvert(Map labelToDirPath) { + if (labelToDirPath.isEmpty()) { + throw new IllegalArgumentException( + "At least one directory must be configured via 'directory.{label}' options."); + } + + Map result = new LinkedHashMap<>(); + Set normalizedDirs = new LinkedHashSet<>(); + for (Map.Entry entry : labelToDirPath.entrySet()) { + Path dir = new Path(entry.getValue()); + if (!normalizedDirs.add(dir.toUri().getPath())) { + throw new IllegalArgumentException( + String.format( + "Directory '%s' is assigned to more than one label.", + entry.getValue())); + } + result.put(entry.getKey(), dir); + } + + for (String dir : normalizedDirs) { + String prefix = dir.endsWith("/") ? dir : dir + "/"; + for (String other : normalizedDirs) { + if (!other.equals(dir) && other.startsWith(prefix)) { + throw new IllegalArgumentException( + String.format( + "Directory '%s' is an ancestor of '%s'. Providing both would " + + "discover the same snapshots under multiple labels.", + dir, other)); + } + } + } + return result; + } + + private static ExecutorService createListingExecutor(int parallelism) { + return new ThreadPoolExecutor( + parallelism, + parallelism, + 0L, + TimeUnit.MILLISECONDS, + new LinkedBlockingQueue<>(parallelism), + r -> { + Thread t = new Thread(r, "state-catalog-listing"); + t.setDaemon(true); + return t; + }, + new ThreadPoolExecutor.CallerRunsPolicy()); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/StateCatalog.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/StateCatalog.java new file mode 100644 index 00000000000000..66c6b174859543 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/StateCatalog.java @@ -0,0 +1,702 @@ +/* + * 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.catalog; + +import org.apache.flink.annotation.PublicEvolving; +import org.apache.flink.runtime.checkpoint.metadata.CheckpointMetadata; +import org.apache.flink.state.api.OperatorIdentifier; +import org.apache.flink.state.api.StateTableUtils; +import org.apache.flink.state.api.runtime.SavepointLoader; +import org.apache.flink.state.api.schema.KeyedStateSchemaInfo; +import org.apache.flink.state.table.SavepointConnectorOptions.StateReaderMode; +import org.apache.flink.state.table.SavepointConnectorOptions.StateType; +import org.apache.flink.table.api.DataTypes; +import org.apache.flink.table.api.Schema; +import org.apache.flink.table.catalog.AbstractCatalog; +import org.apache.flink.table.catalog.CatalogBaseTable; +import org.apache.flink.table.catalog.CatalogDatabase; +import org.apache.flink.table.catalog.CatalogDatabaseImpl; +import org.apache.flink.table.catalog.CatalogFunction; +import org.apache.flink.table.catalog.CatalogPartition; +import org.apache.flink.table.catalog.CatalogPartitionSpec; +import org.apache.flink.table.catalog.CatalogView; +import org.apache.flink.table.catalog.ObjectPath; +import org.apache.flink.table.catalog.exceptions.CatalogException; +import org.apache.flink.table.catalog.exceptions.DatabaseAlreadyExistException; +import org.apache.flink.table.catalog.exceptions.DatabaseNotEmptyException; +import org.apache.flink.table.catalog.exceptions.DatabaseNotExistException; +import org.apache.flink.table.catalog.exceptions.FunctionAlreadyExistException; +import org.apache.flink.table.catalog.exceptions.FunctionNotExistException; +import org.apache.flink.table.catalog.exceptions.PartitionAlreadyExistsException; +import org.apache.flink.table.catalog.exceptions.PartitionNotExistException; +import org.apache.flink.table.catalog.exceptions.PartitionSpecInvalidException; +import org.apache.flink.table.catalog.exceptions.TableAlreadyExistException; +import org.apache.flink.table.catalog.exceptions.TableNotExistException; +import org.apache.flink.table.catalog.exceptions.TableNotPartitionedException; +import org.apache.flink.table.catalog.stats.CatalogColumnStatistics; +import org.apache.flink.table.catalog.stats.CatalogTableStatistics; +import org.apache.flink.table.expressions.Expression; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +/** + * A read-only Flink SQL catalog that discovers checkpoints and savepoints from a configured set of + * directories and exposes their metadata as queryable SQL databases and views. + * + *

The catalog maps Flink's three-level hierarchy as follows: + * + *

    + *
  • Catalog: the name given at registration time (e.g. {@code "state"}) + *
  • Database: one entry per discovered snapshot (e.g. {@code "app1_savepoint-acce1cedsad"}) + *
  • Table: a single view named {@code "metadata"} per database, backed by the {@code + * savepoint_metadata} function from {@code StateModule} + *
+ * + *

Database names preserve hyphens from the original directory names. Backtick quoting is + * required in SQL for identifiers containing hyphens: + * + * 

{@code
+ * USE CATALOG state;
+ * USE `app1_savepoint-acce1cedsad`;
+ * SELECT * FROM metadata;
+ * }
+ * + *

{@code StateModule} must be loaded before querying any {@code metadata} view: + * + * 

{@code
+ * tableEnv.loadModule("state", StateModule.INSTANCE);
+ * }
+ * + *

Each catalog operation fetches state on demand. {@link #listDatabases()} performs a full + * directory scan; all other operations perform a single file check on the specific snapshot path + * reconstructed from the database name. There is no background polling and no shared cache. + * + *

All write operations throw {@link UnsupportedOperationException}. + */ +@PublicEvolving +public class StateCatalog extends AbstractCatalog { + + private static final Logger LOG = LoggerFactory.getLogger(StateCatalog.class); + + public static final String METADATA_TABLE = "metadata"; + public static final String OPERATOR_UID_PREFIX = "uid_"; + public static final String OPERATOR_ID_PREFIX = "id_"; + public static final String OPERATOR_TABLE_SUFFIX = "_keyed"; + public static final String FLAT_STATE_TABLE_SUFFIX = "_keyed_flat"; + + private static final CatalogDatabase EMPTY_DATABASE = + new CatalogDatabaseImpl(Collections.emptyMap(), ""); + + private final SnapshotDiscovery discovery; + + public StateCatalog(String name, Map labelsToDirs) { + this(name, labelsToDirs, StateCatalogOptions.LISTING_PARALLELISM.defaultValue()); + } + + public StateCatalog(String name, Map labelsToDirs, int listingParallelism) { + this( + name, + labelsToDirs, + listingParallelism, + StateCatalogOptions.DB_NAME_INCLUDE_TS.defaultValue()); + } + + public StateCatalog( + String name, + Map labelsToDirs, + int listingParallelism, + boolean dbNameIncludeTs) { + super(name, "default"); + this.discovery = new SnapshotDiscovery(labelsToDirs, listingParallelism, dbNameIncludeTs); + } + + @Override + @Nullable + public String getDefaultDatabase() { + return null; + } + + // ------------------------------------------------------------------------- + // Lifecycle + // ------------------------------------------------------------------------- + + @Override + public void open() throws CatalogException { + discovery.start(); + listDatabases(); + } + + @Override + public void close() throws CatalogException { + discovery.stop(); + } + + // ------------------------------------------------------------------------- + // Databases + // ------------------------------------------------------------------------- + + @Override + public List listDatabases() throws CatalogException { + try { + return discovery.list(); + } catch (IOException e) { + LOG.warn("Failed to list databases in catalog '{}'", getName(), e); + return Collections.emptyList(); + } + } + + @Override + public CatalogDatabase getDatabase(String databaseName) + throws DatabaseNotExistException, CatalogException { + if (discovery.find(databaseName).isEmpty()) { + throw new DatabaseNotExistException(getName(), databaseName); + } + return EMPTY_DATABASE; + } + + @Override + public boolean databaseExists(String databaseName) throws CatalogException { + return discovery.find(databaseName).isPresent(); + } + + @Override + public void createDatabase(String name, CatalogDatabase database, boolean ignoreIfExists) + throws DatabaseAlreadyExistException, CatalogException { + throw new UnsupportedOperationException("StateCatalog is read-only."); + } + + @Override + public void dropDatabase(String name, boolean ignoreIfNotExists, boolean cascade) + throws DatabaseNotExistException, DatabaseNotEmptyException, CatalogException { + throw new UnsupportedOperationException("StateCatalog is read-only."); + } + + @Override + public void alterDatabase(String name, CatalogDatabase newDatabase, boolean ignoreIfNotExists) + throws DatabaseNotExistException, CatalogException { + throw new UnsupportedOperationException("StateCatalog is read-only."); + } + + // ------------------------------------------------------------------------- + // Tables and views + // ------------------------------------------------------------------------- + + @Override + public List listTables(String databaseName) + throws DatabaseNotExistException, CatalogException { + Optional snapshotPath = discovery.find(databaseName); + if (snapshotPath.isEmpty()) { + throw new DatabaseNotExistException(getName(), databaseName); + } + List tables = new ArrayList<>(); + tables.add(METADATA_TABLE); + Set seen = new LinkedHashSet<>(); + try { + CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(snapshotPath.get()); + for (OperatorIdentifier opId : StateTableUtils.getOperatorIdentifiers(metadata)) { + for (ResolvedTable candidate : candidateTablesForOperator(metadata, opId)) { + String name = + tableName( + candidate.operatorIdentifier, + candidate.kind, + candidate.stateName); + // Two distinct (operator, state) combinations can legitimately derive the + // same table name, since operator UIDs and state names may themselves + // contain underscores (see #tableName). Skip and warn rather than exposing + // the same name twice, mirroring how SnapshotDiscovery#list handles + // colliding database names. + if (!seen.add(name)) { + LOG.warn( + "Table name '{}' is ambiguous between multiple operators/states " + + "in database '{}' and only the first one found is " + + "exposed. Consider renaming the colliding operator " + + "UID(s) or state name(s).", + name, + databaseName); + continue; + } + tables.add(name); + } + } + } catch (IOException e) { + throw new CatalogException( + "Failed to load checkpoint metadata for database '" + databaseName + "'", e); + } + return tables; + } + + @Override + public List listViews(String databaseName) + throws DatabaseNotExistException, CatalogException { + if (discovery.find(databaseName).isEmpty()) { + throw new DatabaseNotExistException(getName(), databaseName); + } + return Collections.singletonList(METADATA_TABLE); + } + + @Override + public CatalogBaseTable getTable(ObjectPath tablePath) + throws TableNotExistException, CatalogException { + Optional snapshotPath = discovery.find(tablePath.getDatabaseName()); + if (snapshotPath.isEmpty()) { + throw new TableNotExistException(getName(), tablePath); + } + String tableName = tablePath.getObjectName(); + if (METADATA_TABLE.equals(tableName)) { + return buildMetadataView(snapshotPath.get()); + } + try { + CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(snapshotPath.get()); + ResolvedTable resolved = + resolveTable(metadata, tableName) + .orElseThrow(() -> new TableNotExistException(getName(), tablePath)); + switch (resolved.kind) { + case KEYED: + { + KeyedStateSchemaInfo schemaInfo = + StateTableUtils.getKeyedStateSchema( + metadata, resolved.operatorIdentifier); + return StateTableUtils.getStateCatalogTable( + metadata, + schemaInfo, + snapshotPath.get(), + resolved.operatorIdentifier); + } + case KEYED_FLAT: + { + KeyedStateSchemaInfo schemaInfo = + StateTableUtils.getKeyedStateSchema( + metadata, resolved.operatorIdentifier); + return StateTableUtils.getFlattenedStateCatalogTable( + metadata, + schemaInfo, + resolved.stateName, + snapshotPath.get(), + resolved.operatorIdentifier); + } + default: + throw new IllegalStateException("Unhandled table kind " + resolved.kind); + } + } catch (IOException e) { + throw new CatalogException( + "Failed to load state schema for table '" + tablePath + "'", e); + } + } + + @Override + public boolean tableExists(ObjectPath tablePath) throws CatalogException { + Optional snapshotPath = discovery.find(tablePath.getDatabaseName()); + if (snapshotPath.isEmpty()) { + return false; + } + String tableName = tablePath.getObjectName(); + if (METADATA_TABLE.equals(tableName)) { + return true; + } + if (!tableName.startsWith(OPERATOR_UID_PREFIX) + && !tableName.startsWith(OPERATOR_ID_PREFIX)) { + return false; + } + try { + CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(snapshotPath.get()); + return resolveTable(metadata, tableName).isPresent(); + } catch (IOException e) { + LOG.warn( + "Failed to load checkpoint metadata while checking existence of table '{}'", + tablePath, + e); + return false; + } + } + + @Override + public void createTable(ObjectPath tablePath, CatalogBaseTable table, boolean ignoreIfExists) + throws TableAlreadyExistException, DatabaseNotExistException, CatalogException { + throw new UnsupportedOperationException("StateCatalog is read-only."); + } + + @Override + public void alterTable( + ObjectPath tablePath, CatalogBaseTable newTable, boolean ignoreIfNotExists) + throws TableNotExistException, CatalogException { + throw new UnsupportedOperationException("StateCatalog is read-only."); + } + + @Override + public void dropTable(ObjectPath tablePath, boolean ignoreIfNotExists) + throws TableNotExistException, CatalogException { + throw new UnsupportedOperationException("StateCatalog is read-only."); + } + + @Override + public void renameTable(ObjectPath tablePath, String newTableName, boolean ignoreIfNotExists) + throws TableAlreadyExistException, TableNotExistException, CatalogException { + throw new UnsupportedOperationException("StateCatalog is read-only."); + } + + // ------------------------------------------------------------------------- + // Partitions (not supported) + // ------------------------------------------------------------------------- + + @Override + public List listPartitions(ObjectPath tablePath) + throws TableNotExistException, TableNotPartitionedException, CatalogException { + if (!tableExists(tablePath)) { + throw new TableNotExistException(getName(), tablePath); + } + throw new TableNotPartitionedException(getName(), tablePath); + } + + @Override + public List listPartitions( + ObjectPath tablePath, CatalogPartitionSpec partitionSpec) + throws TableNotExistException, TableNotPartitionedException, CatalogException { + return listPartitions(tablePath); + } + + @Override + public List listPartitionsByFilter( + ObjectPath tablePath, List filters) + throws TableNotExistException, TableNotPartitionedException, CatalogException { + return listPartitions(tablePath); + } + + @Override + public CatalogPartition getPartition(ObjectPath tablePath, CatalogPartitionSpec partitionSpec) + throws PartitionNotExistException, CatalogException { + throw new PartitionNotExistException(getName(), tablePath, partitionSpec); + } + + @Override + public boolean partitionExists(ObjectPath tablePath, CatalogPartitionSpec partitionSpec) + throws CatalogException { + return false; + } + + @Override + public void createPartition( + ObjectPath tablePath, + CatalogPartitionSpec partitionSpec, + CatalogPartition partition, + boolean ignoreIfExists) + throws TableNotExistException, + TableNotPartitionedException, + PartitionSpecInvalidException, + PartitionAlreadyExistsException, + CatalogException { + throw new UnsupportedOperationException("StateCatalog is read-only."); + } + + @Override + public void dropPartition( + ObjectPath tablePath, CatalogPartitionSpec partitionSpec, boolean ignoreIfNotExists) + throws PartitionNotExistException, CatalogException { + throw new UnsupportedOperationException("StateCatalog is read-only."); + } + + @Override + public void alterPartition( + ObjectPath tablePath, + CatalogPartitionSpec partitionSpec, + CatalogPartition newPartition, + boolean ignoreIfNotExists) + throws PartitionNotExistException, CatalogException { + throw new UnsupportedOperationException("StateCatalog is read-only."); + } + + // ------------------------------------------------------------------------- + // Functions (not supported) + // ------------------------------------------------------------------------- + + @Override + public List listFunctions(String dbName) + throws DatabaseNotExistException, CatalogException { + return Collections.emptyList(); + } + + @Override + public CatalogFunction getFunction(ObjectPath functionPath) + throws FunctionNotExistException, CatalogException { + throw new FunctionNotExistException(getName(), functionPath); + } + + @Override + public boolean functionExists(ObjectPath functionPath) throws CatalogException { + return false; + } + + @Override + public void createFunction( + ObjectPath functionPath, CatalogFunction function, boolean ignoreIfExists) + throws FunctionAlreadyExistException, DatabaseNotExistException, CatalogException { + throw new UnsupportedOperationException("StateCatalog is read-only."); + } + + @Override + public void alterFunction( + ObjectPath functionPath, CatalogFunction newFunction, boolean ignoreIfNotExists) + throws FunctionNotExistException, CatalogException { + throw new UnsupportedOperationException("StateCatalog is read-only."); + } + + @Override + public void dropFunction(ObjectPath functionPath, boolean ignoreIfNotExists) + throws FunctionNotExistException, CatalogException { + throw new UnsupportedOperationException("StateCatalog is read-only."); + } + + // ------------------------------------------------------------------------- + // Statistics (read-only stubs) + // ------------------------------------------------------------------------- + + @Override + public CatalogTableStatistics getTableStatistics(ObjectPath tablePath) + throws TableNotExistException, CatalogException { + if (!tableExists(tablePath)) { + throw new TableNotExistException(getName(), tablePath); + } + return CatalogTableStatistics.UNKNOWN; + } + + @Override + public CatalogColumnStatistics getTableColumnStatistics(ObjectPath tablePath) + throws TableNotExistException, CatalogException { + if (!tableExists(tablePath)) { + throw new TableNotExistException(getName(), tablePath); + } + return CatalogColumnStatistics.UNKNOWN; + } + + @Override + public CatalogTableStatistics getPartitionStatistics( + ObjectPath tablePath, CatalogPartitionSpec partitionSpec) + throws PartitionNotExistException, CatalogException { + throw new PartitionNotExistException(getName(), tablePath, partitionSpec); + } + + @Override + public CatalogColumnStatistics getPartitionColumnStatistics( + ObjectPath tablePath, CatalogPartitionSpec partitionSpec) + throws PartitionNotExistException, CatalogException { + throw new PartitionNotExistException(getName(), tablePath, partitionSpec); + } + + @Override + public void alterTableStatistics( + ObjectPath tablePath, CatalogTableStatistics tableStatistics, boolean ignoreIfNotExists) + throws TableNotExistException, CatalogException { + throw new UnsupportedOperationException("StateCatalog is read-only."); + } + + @Override + public void alterTableColumnStatistics( + ObjectPath tablePath, + CatalogColumnStatistics columnStatistics, + boolean ignoreIfNotExists) + throws TableNotExistException, CatalogException { + throw new UnsupportedOperationException("StateCatalog is read-only."); + } + + @Override + public void alterPartitionStatistics( + ObjectPath tablePath, + CatalogPartitionSpec partitionSpec, + CatalogTableStatistics partitionStatistics, + boolean ignoreIfNotExists) + throws PartitionNotExistException, CatalogException { + throw new UnsupportedOperationException("StateCatalog is read-only."); + } + + @Override + public void alterPartitionColumnStatistics( + ObjectPath tablePath, + CatalogPartitionSpec partitionSpec, + CatalogColumnStatistics columnStatistics, + boolean ignoreIfNotExists) + throws PartitionNotExistException, CatalogException { + throw new UnsupportedOperationException("StateCatalog is read-only."); + } + + // ------------------------------------------------------------------------- + // View construction + // ------------------------------------------------------------------------- + + private static CatalogView buildMetadataView(String snapshotPath) { + String escapedPath = snapshotPath.replace("'", "''"); + String query = String.format("SELECT * FROM TABLE(savepoint_metadata('%s'))", escapedPath); + // Once the upstream StateCatalog PR is merged and OUTPUT_DATA_TYPE is available in + // SavepointMetadataTableFunction, replace the schema definition below with: + // Schema.newBuilder() + // .fromRowDataType(SavepointMetadataTableFunction.OUTPUT_DATA_TYPE) + // .build() + + Schema schema = + Schema.newBuilder() + .fromRowDataType( + DataTypes.ROW( + DataTypes.FIELD( + "checkpoint-id", DataTypes.BIGINT().notNull()), + DataTypes.FIELD("operator-name", DataTypes.STRING()), + DataTypes.FIELD("operator-uid", DataTypes.STRING()), + DataTypes.FIELD( + "operator-uid-hash", DataTypes.STRING().notNull()), + DataTypes.FIELD( + "operator-parallelism", DataTypes.INT().notNull()), + DataTypes.FIELD( + "operator-max-parallelism", + DataTypes.INT().notNull()), + DataTypes.FIELD( + "operator-subtask-state-count", + DataTypes.INT().notNull()), + DataTypes.FIELD( + "operator-coordinator-state-size-in-bytes", + DataTypes.BIGINT().notNull()), + DataTypes.FIELD( + "operator-total-size-in-bytes", + DataTypes.BIGINT().notNull()))) + .build(); + return CatalogView.of( + schema, + "Operator metadata for snapshot at " + snapshotPath, + query, + query, + Collections.emptyMap()); + } + + // ------------------------------------------------------------------------- + // Operator table helpers + // ------------------------------------------------------------------------- + + /** + * Table name for a {@code kind} of operator state, optionally scoped to one flattened/non-keyed + * state (see {@link #OPERATOR_TABLE_SUFFIX}/{@link #FLAT_STATE_TABLE_SUFFIX}). + * + *

{@code stateName} must be {@code null} for {@link StateReaderMode#KEYED}/{@link + * StateReaderMode#WINDOWED} (the general keyed/namespaced table, one per operator) and non-null + * for every other kind (a table scoped to one flattened LIST/MAP state, or one non-keyed state + * — the state name alone disambiguates the table since keyed/non-keyed state names are unique + * within an operator). + */ + private static final Map TABLE_SUFFIXES = + Map.of( + StateReaderMode.KEYED, OPERATOR_TABLE_SUFFIX, + StateReaderMode.KEYED_FLAT, FLAT_STATE_TABLE_SUFFIX); + + static String tableName( + OperatorIdentifier opId, StateReaderMode kind, @Nullable String stateName) { + String base = + opId.getUid() + .map(uid -> OPERATOR_UID_PREFIX + uid) + .orElseGet(() -> OPERATOR_ID_PREFIX + opId.getOperatorId().toHexString()); + String suffix = TABLE_SUFFIXES.get(kind); + if (suffix == null) { + throw new IllegalArgumentException("Unknown state reader mode: " + kind); + } + return stateName == null ? base + suffix : base + "_" + stateName + suffix; + } + + /** + * Identifies which table a table name refers to: the operator, the table shape ({@link + * StateReaderMode}), and — for flattened tables — the name of the flattened state. + */ + private static final class ResolvedTable { + final OperatorIdentifier operatorIdentifier; + final StateReaderMode kind; + @Nullable final String stateName; + + ResolvedTable(OperatorIdentifier operatorIdentifier, StateReaderMode kind) { + this(operatorIdentifier, kind, null); + } + + ResolvedTable( + OperatorIdentifier operatorIdentifier, + StateReaderMode kind, + @Nullable String stateName) { + this.operatorIdentifier = operatorIdentifier; + this.kind = kind; + this.stateName = stateName; + } + } + + /** + * Resolves a table name to the operator (and, for flattened tables, the state) it refers to, by + * generating candidate names for every operator/state in the checkpoint and matching against + * {@code tableName}. Names cannot be parsed directly since operator UIDs and state names may + * themselves contain underscores. + */ + private static Optional resolveTable( + CheckpointMetadata metadata, String tableName) { + for (OperatorIdentifier opId : StateTableUtils.getOperatorIdentifiers(metadata)) { + List candidates; + try { + candidates = candidateTablesForOperator(metadata, opId); + } catch (IOException e) { + LOG.warn("Failed to load state schema for operator '{}'. Skipping.", opId, e); + continue; + } + for (ResolvedTable candidate : candidates) { + if (tableName(candidate.operatorIdentifier, candidate.kind, candidate.stateName) + .equals(tableName)) { + return Optional.of(candidate); + } + } + } + return Optional.empty(); + } + + /** + * Enumerates every table that {@code opId} contributes: the general keyed table (if any plain + * per-key state is registered), plus one flattened table per LIST/MAP keyed state. + * + *

Shared by {@link #listTables} (which collects names for every candidate) and {@link + * #resolveTable} (which matches candidates against a target name), so that adding a new state + * kind only requires updating this one traversal. + */ + private static List candidateTablesForOperator( + CheckpointMetadata metadata, OperatorIdentifier opId) throws IOException { + List candidates = new ArrayList<>(); + + KeyedStateSchemaInfo schemaInfo = StateTableUtils.getKeyedStateSchema(metadata, opId); + if (!schemaInfo.stateSchemas.isEmpty()) { + candidates.add(new ResolvedTable(opId, StateReaderMode.KEYED)); + for (Map.Entry entry : + schemaInfo.stateSchemas.entrySet()) { + StateType stateType = entry.getValue().stateType; + if (stateType == StateType.LIST || stateType == StateType.MAP) { + candidates.add( + new ResolvedTable(opId, StateReaderMode.KEYED_FLAT, entry.getKey())); + } + } + } + + return candidates; + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/StateCatalogFactory.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/StateCatalogFactory.java new file mode 100644 index 00000000000000..9f18a8ae9bc19c --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/StateCatalogFactory.java @@ -0,0 +1,86 @@ +/* + * 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.catalog; + +import org.apache.flink.annotation.PublicEvolving; +import org.apache.flink.configuration.ConfigOption; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.table.catalog.Catalog; +import org.apache.flink.table.factories.CatalogFactory; + +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Set; + +/** + * Factory for creating {@link StateCatalog} instances via SQL DDL or programmatically. + * + *

Directories are configured with {@code directory.{label}} options: + * + * 

{@code
+ * CREATE CATALOG state WITH (
+ *     'type'              = 'state',
+ *     'directory.my-app'  = '/checkpoints/app1',
+ *     'directory.staging' = '/savepoints/staging'
+ * );
+ * }
+ */ +@PublicEvolving +public class StateCatalogFactory implements CatalogFactory { + + public static final String IDENTIFIER = "state"; + + @Override + public String factoryIdentifier() { + return IDENTIFIER; + } + + @Override + public Set> requiredOptions() { + return Collections.emptySet(); + } + + @Override + public Set> optionalOptions() { + return Set.of( + StateCatalogOptions.LISTING_PARALLELISM, StateCatalogOptions.DB_NAME_INCLUDE_TS); + } + + @Override + public Catalog createCatalog(Context context) { + Map options = context.getOptions(); + + Map labelsToDirs = new LinkedHashMap<>(); + for (Map.Entry entry : options.entrySet()) { + if (entry.getKey().startsWith(StateCatalogOptions.DIRECTORY_PREFIX)) { + String label = + entry.getKey().substring(StateCatalogOptions.DIRECTORY_PREFIX.length()); + labelsToDirs.put(label, entry.getValue()); + } + } + + Configuration configuration = Configuration.fromMap(options); + return new StateCatalog( + context.getName(), + labelsToDirs, + configuration.get(StateCatalogOptions.LISTING_PARALLELISM), + configuration.get(StateCatalogOptions.DB_NAME_INCLUDE_TS)); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/StateCatalogOptions.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/StateCatalogOptions.java new file mode 100644 index 00000000000000..6b0e7f10a672f6 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/StateCatalogOptions.java @@ -0,0 +1,59 @@ +/* + * 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.catalog; + +import org.apache.flink.annotation.PublicEvolving; +import org.apache.flink.configuration.ConfigOption; +import org.apache.flink.configuration.ConfigOptions; + +/** Configuration options for {@link StateCatalog}. */ +@PublicEvolving +public class StateCatalogOptions { + + /** + * Prefix for directory options. Each option of the form {@code directory.{label}} maps a + * human-readable label to a filesystem path. The label becomes the first segment of every + * database name derived from that directory (e.g. {@code my-app/savepoint-abc}). + */ + public static final String DIRECTORY_PREFIX = "directory."; + + public static final ConfigOption LISTING_PARALLELISM = + ConfigOptions.key("listing-parallelism") + .intType() + .defaultValue(10) + .withDescription( + "Maximum number of concurrent directory listing requests issued " + + "during a scan. Directories at the same depth are listed " + + "in parallel. Increase for high-latency remote filesystems " + + "(e.g. S3); decrease to reduce load on the filesystem."); + + public static final ConfigOption DB_NAME_INCLUDE_TS = + ConfigOptions.key("db-name.include-ts") + .booleanType() + .defaultValue(true) + .withDescription( + "Whether derived database names include the snapshot's creation " + + "timestamp as a segment, i.e. label/creationTs/relativePath " + + "instead of label/relativePath. The timestamp is the " + + "modification time of the snapshot's _metadata file, " + + "formatted with yyyy-MM-dd'T'HH:mm:ssX (e.g. " + + "2026-07-22T10:30:45Z)."); + + private StateCatalogOptions() {} +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/AbstractMultiColumnScanProvider.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/AbstractMultiColumnScanProvider.java new file mode 100644 index 00000000000000..73772a2c3ec21c --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/AbstractMultiColumnScanProvider.java @@ -0,0 +1,83 @@ +/* + * 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.annotation.Internal; +import org.apache.flink.api.common.state.StateDescriptor; +import org.apache.flink.state.api.OperatorIdentifier; +import org.apache.flink.state.api.filter.SavepointKeyFilter; +import org.apache.flink.state.api.schema.StateSchemaInfo; +import org.apache.flink.table.types.logical.RowType; + +import javax.annotation.Nullable; + +import java.util.List; +import java.util.Map; +import java.util.function.Supplier; + +/** + * Base for scan providers whose mapping registers an arbitrary number of value columns, each as its + * own state (see {@link MultiColumnStateMapping}): {@link SavepointDataStreamScanProvider} and + * {@link WindowSavepointDataStreamScanProvider}. + */ +@Internal +abstract class AbstractMultiColumnScanProvider + extends AbstractSavepointDataStreamScanProvider { + + protected AbstractMultiColumnScanProvider( + @Nullable final String stateBackendType, + final String statePath, + final OperatorIdentifier operatorIdentifier, + final Supplier mappingSupplier, + final RowType rowType, + @Nullable final SavepointKeyFilter keyFilter) { + super(stateBackendType, statePath, operatorIdentifier, mappingSupplier, rowType, keyFilter); + } + + /** + * Builds descriptors for ALL original value columns so that key(-and-namespace) enumeration + * works even when a projection removes all value columns from the output (e.g. {@code SELECT k + * FROM t WHERE k = 5}). + */ + @Override + protected final void prepareStateDescriptors(M mapping) { + List columns = mapping.getAllValueColumns(); + Map fallbackSchemas = + loadFallbackSchemas( + columns.stream() + .anyMatch( + c -> + isSerializerMissing( + c.getStateType(), + c.getMapKeyTypeSerializer(), + c.getValueTypeSerializer()))); + + for (StateValueColumnConfiguration columnConfig : columns) { + StateDescriptor descriptor = + buildStateDescriptor( + columnConfig.getStateName(), + columnConfig.getStateType(), + columnConfig.getActualStateKind(), + columnConfig.getMapKeyTypeSerializer(), + columnConfig.getValueTypeSerializer(), + fallbackSchemas); + columnConfig.setStateDescriptor(descriptor); + } + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/AbstractSavepointDataStreamScanProvider.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/AbstractSavepointDataStreamScanProvider.java new file mode 100644 index 00000000000000..0a780fd1a744a6 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/AbstractSavepointDataStreamScanProvider.java @@ -0,0 +1,295 @@ +/* + * 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.annotation.Internal; +import org.apache.flink.api.common.functions.AggregateFunction; +import org.apache.flink.api.common.functions.ReduceFunction; +import org.apache.flink.api.common.state.AggregatingStateDescriptor; +import org.apache.flink.api.common.state.ListStateDescriptor; +import org.apache.flink.api.common.state.MapStateDescriptor; +import org.apache.flink.api.common.state.ReducingStateDescriptor; +import org.apache.flink.api.common.state.StateDescriptor; +import org.apache.flink.api.common.state.ValueStateDescriptor; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.StateBackendOptions; +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.state.api.functions.KeyedStateReaderFunction; +import org.apache.flink.state.api.schema.StateSchemaInfo; +import org.apache.flink.streaming.api.datastream.DataStream; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.table.connector.ProviderContext; +import org.apache.flink.table.connector.source.DataStreamScanProvider; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.runtime.typeutils.InternalTypeInfo; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.util.StringUtils; + +import javax.annotation.Nullable; + +import java.util.Map; +import java.util.function.Supplier; + +/** + * Shared scan-time logic for the keyed state scan providers: opening the {@link SavepointReader} + * against the configured state backend, resolving the lazy mapping and building its state + * descriptors. Subclasses supply the mapping-specific descriptor setup and {@link SavepointReader} + * call. + */ +@Internal +@SuppressWarnings("rawtypes") +abstract class AbstractSavepointDataStreamScanProvider + implements DataStreamScanProvider { + + @Nullable protected final String stateBackendType; + protected final String statePath; + protected final OperatorIdentifier operatorIdentifier; + private final Supplier mappingSupplier; + protected final RowType rowType; + @Nullable protected final SavepointKeyFilter keyFilter; + + protected AbstractSavepointDataStreamScanProvider( + @Nullable final String stateBackendType, + final String statePath, + final OperatorIdentifier operatorIdentifier, + final Supplier mappingSupplier, + final RowType rowType, + @Nullable final SavepointKeyFilter keyFilter) { + this.stateBackendType = stateBackendType; + this.statePath = statePath; + this.operatorIdentifier = operatorIdentifier; + this.mappingSupplier = mappingSupplier; + this.rowType = rowType; + this.keyFilter = keyFilter; + } + + @Override + public boolean isBounded() { + return true; + } + + @Override + public DataStream produceDataStream( + ProviderContext providerContext, StreamExecutionEnvironment execEnv) { + try { + SavepointReader savepointReader = + createSavepointReader( + stateBackendType, statePath, execEnv, getClass().getClassLoader()); + + // Resolve the lazy mapping at scan time (class loading deferred from planning). + M mapping = mappingSupplier.get(); + prepareStateDescriptors(mapping); + + return readState(savepointReader, mapping); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + /** + * Builds the {@link StateDescriptor}(s) needed for this scan and registers them onto {@code + * mapping}, resolving fallback schemas from the savepoint header where no explicit serializer + * is available. + */ + protected abstract void prepareStateDescriptors(M mapping); + + /** Drives the actual {@link SavepointReader} call for this scan. */ + protected abstract DataStream readState(SavepointReader savepointReader, M mapping) + throws Exception; + + /** + * Reads plain ({@code VoidNamespace}) keyed state. Namespaced (window) scan providers instead + * drive {@link SavepointReader#readWindowKeyedState} from their own {@link #readState}. + */ + @SuppressWarnings("unchecked") + protected final DataStream readVoidNamespaceKeyedState( + SavepointReader savepointReader, + M mapping, + KeyedStateReaderFunction readerFunction) + throws Exception { + return savepointReader.readKeyedState( + operatorIdentifier, + readerFunction, + (TypeInformation) mapping.getKeyTypeInfo(), + InternalTypeInfo.of(rowType), + keyFilter); + } + + /** + * Loads the configured (or overridden) {@link StateBackend} and opens a {@link SavepointReader} + * against it. Also used by the non-keyed (operator state) scan providers, which cannot extend + * this keyed-specific abstraction. + */ + static SavepointReader createSavepointReader( + @Nullable String stateBackendType, + String statePath, + StreamExecutionEnvironment execEnv, + ClassLoader classLoader) + throws Exception { + Configuration configuration = Configuration.fromMap(execEnv.getConfiguration().toMap()); + if (!StringUtils.isNullOrWhitespaceOnly(stateBackendType)) { + configuration.set(StateBackendOptions.STATE_BACKEND, stateBackendType); + } + StateBackend stateBackend = + StateBackendLoader.loadStateBackendFromConfig(configuration, classLoader, null); + return SavepointReader.read(execEnv, statePath, stateBackend); + } + + /** + * Reads the savepoint header and returns a map of state name → {@link StateSchemaInfo}, or an + * empty map if no serializer is missing (avoiding the header read entirely). + */ + protected final Map loadFallbackSchemas(boolean anyMissingSerializer) { + return SavepointFallbackSchemaLoader.loadFallbackSchemas( + statePath, operatorIdentifier, anyMissingSerializer); + } + + /** + * Whether the serializer(s) needed to build a state descriptor could not be resolved from the + * preloaded savepoint metadata, in which case they must be restored from the savepoint header. + */ + static boolean isSerializerMissing( + SavepointConnectorOptions.StateType stateType, + @Nullable TypeSerializer mapKeyTypeSerializer, + @Nullable TypeSerializer valueTypeSerializer) { + return valueTypeSerializer == null + || (stateType == SavepointConnectorOptions.StateType.MAP + && mapKeyTypeSerializer == null); + } + + /** + * Builds the {@link StateDescriptor} for a single VALUE/LIST/MAP state, restoring the original + * serializer from {@code fallbackSchemas} for any serializer that could not be resolved from + * the preloaded savepoint metadata. + * + *

For the coarse {@code VALUE} shape, {@code actualStateKind} distinguishes what the state + * was actually registered as: a {@code .reduce()}/{@code .aggregate()} window function's + * window-contents state is registered as {@code REDUCING}/{@code AGGREGATING} rather than plain + * {@code VALUE}, and the state backend rejects descriptor lookups whose {@link + * StateDescriptor.Type} doesn't exactly match. In those cases a matching {@link + * ReducingStateDescriptor}/{@link AggregatingStateDescriptor} is built instead, using a + * reduce/aggregate function that is never invoked since callers only ever read this state. + */ + @SuppressWarnings("unchecked") + protected static StateDescriptor buildStateDescriptor( + String name, + SavepointConnectorOptions.StateType stateType, + StateDescriptor.Type actualStateKind, + @Nullable TypeSerializer mapKeyTypeSerializer, + @Nullable TypeSerializer valueTypeSerializer, + Map fallbackSchemas) { + switch (stateType) { + case VALUE: + TypeSerializer valueSerializer = + resolveValueSerializer(name, valueTypeSerializer, fallbackSchemas); + switch (actualStateKind) { + case REDUCING: + return new ReducingStateDescriptor<>( + name, new NoOpReduceFunction(), valueSerializer); + case AGGREGATING: + return new AggregatingStateDescriptor<>( + name, new IdentityAggregateFunction(), valueSerializer); + default: + return new ValueStateDescriptor<>(name, valueSerializer); + } + + case LIST: + return new ListStateDescriptor<>( + name, resolveValueSerializer(name, valueTypeSerializer, fallbackSchemas)); + + case MAP: + if (valueTypeSerializer != null && mapKeyTypeSerializer != null) { + return new MapStateDescriptor<>( + name, mapKeyTypeSerializer, valueTypeSerializer); + } + StateSchemaInfo schema = + SavepointFallbackSchemaLoader.getSchema(name, fallbackSchemas); + return new MapStateDescriptor<>( + name, + SavepointFallbackSchemaLoader.buildFallbackSerializer( + schema.mapKeySnapshot), + SavepointFallbackSchemaLoader.buildFallbackSerializer( + schema.valueSnapshot)); + + default: + throw new UnsupportedOperationException("Unsupported state type: " + stateType); + } + } + + private static TypeSerializer resolveValueSerializer( + String name, + @Nullable TypeSerializer valueTypeSerializer, + Map fallbackSchemas) { + if (valueTypeSerializer != null) { + return valueTypeSerializer; + } + return SavepointFallbackSchemaLoader.buildFallbackSerializer( + SavepointFallbackSchemaLoader.getSchema(name, fallbackSchemas).valueSnapshot); + } + + /** + * Never-invoked {@link ReduceFunction} used to build a {@link ReducingStateDescriptor} for + * read-only access: callers only call {@code ReducingState.get()}, which never merges. + */ + private static final class NoOpReduceFunction implements ReduceFunction { + @Override + public Object reduce(Object value1, Object value2) { + throw new UnsupportedOperationException( + "This reduce function only supports read-only state access and should never be invoked."); + } + } + + /** + * {@link AggregateFunction} used to build an {@link AggregatingStateDescriptor} for read-only + * access, with the accumulator type used as both {@code ACC} and {@code OUT} so that {@code + * AggregatingState.get()} returns the raw accumulator (matching the SQL column, which reflects + * the accumulator's serializer). + */ + private static final class IdentityAggregateFunction + implements AggregateFunction { + @Override + public Object createAccumulator() { + throw new UnsupportedOperationException( + "This aggregate function only supports read-only state access and should never be invoked."); + } + + @Override + public Object add(Object value, Object accumulator) { + throw new UnsupportedOperationException( + "This aggregate function only supports read-only state access and should never be invoked."); + } + + @Override + public Object getResult(Object accumulator) { + return accumulator; + } + + @Override + public Object merge(Object a, Object b) { + throw new UnsupportedOperationException( + "This aggregate function only supports read-only state access and should never be invoked."); + } + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/AbstractSavepointDynamicTableSource.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/AbstractSavepointDynamicTableSource.java new file mode 100644 index 00000000000000..45ab0ffb5b6061 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/AbstractSavepointDynamicTableSource.java @@ -0,0 +1,138 @@ +/* + * 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.annotation.Internal; +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.DataStreamScanProvider; +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 javax.annotation.Nullable; + +import java.util.List; +import java.util.function.Supplier; + +/** + * Shared planning-time state and behaviour for {@link SavepointDynamicTableSource} and {@link + * FlattenedSavepointDynamicTableSource}: the common constructor fields, key-column filter push-down + * (via {@link SavepointFilterTranslator}) and the fixed insert-only changelog mode. + * + *

The scan itself is delegated to the {@link DataStreamScanProvider} supplied by {@link + * SavepointDynamicTableSourceFactory} as a constructor reference (e.g. {@code + * SavepointDataStreamScanProvider::new}), so a single table-source class serves every keyed state + * table kind without a subclass per kind. + */ +@Internal +abstract class AbstractSavepointDynamicTableSource + implements ScanTableSource, SupportsFilterPushDown { + + /** Builds the {@link DataStreamScanProvider} for a given set of scan-time arguments. */ + interface ScanProviderFactory { + DataStreamScanProvider create( + @Nullable String stateBackendType, + String statePath, + OperatorIdentifier operatorIdentifier, + Supplier mappingSupplier, + RowType rowType, + @Nullable SavepointKeyFilter keyFilter); + } + + @Nullable protected final String stateBackendType; + protected final String statePath; + protected final OperatorIdentifier operatorIdentifier; + protected final String summaryString; + protected final ScanProviderFactory scanProviderFactory; + + protected Supplier mappingSupplier; + protected RowType rowType; + + /** + * Index of the (single) key column in {@link #rowType}. Tracked eagerly so filter push-down can + * reference it during planning without resolving the lazy mapping, and updated by projection + * push-down where supported. + */ + protected int keyColumnIndex; + + @Nullable protected SavepointKeyFilter keyFilter; + + protected AbstractSavepointDynamicTableSource( + @Nullable final String stateBackendType, + final String statePath, + final OperatorIdentifier operatorIdentifier, + final int keyColumnIndex, + final Supplier mappingSupplier, + final RowType rowType, + final String summaryString, + final ScanProviderFactory scanProviderFactory) { + this.stateBackendType = stateBackendType; + this.statePath = statePath; + this.operatorIdentifier = operatorIdentifier; + this.keyColumnIndex = keyColumnIndex; + this.mappingSupplier = mappingSupplier; + this.rowType = rowType; + this.summaryString = summaryString; + this.scanProviderFactory = scanProviderFactory; + } + + @Override + public Result applyFilters(List filters) { + return SavepointFilterTranslator.applyKeyColumnFilters( + keyColumnIndex, rowType, filters, kf -> this.keyFilter = kf); + } + + @Override + public final DynamicTableSource copy() { + AbstractSavepointDynamicTableSource copy = newInstance(); + copy.keyFilter = this.keyFilter; + return copy; + } + + /** + * Creates a fresh instance carrying the same constructor state as this one (used by {@link + * #copy()}, which separately copies the mutable {@link #keyFilter}). + */ + protected abstract AbstractSavepointDynamicTableSource newInstance(); + + @Override + public ChangelogMode getChangelogMode() { + return ChangelogMode.insertOnly(); + } + + @Override + public String asSummaryString() { + return summaryString; + } + + @Override + public ScanRuntimeProvider getScanRuntimeProvider(ScanContext scanContext) { + return scanProviderFactory.create( + stateBackendType, + statePath, + operatorIdentifier, + mappingSupplier, + rowType, + keyFilter); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/AbstractSingleColumnScanProvider.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/AbstractSingleColumnScanProvider.java new file mode 100644 index 00000000000000..590324b206fa43 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/AbstractSingleColumnScanProvider.java @@ -0,0 +1,73 @@ +/* + * 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.annotation.Internal; +import org.apache.flink.api.common.state.StateDescriptor; +import org.apache.flink.state.api.OperatorIdentifier; +import org.apache.flink.state.api.filter.SavepointKeyFilter; +import org.apache.flink.state.api.schema.StateSchemaInfo; +import org.apache.flink.table.types.logical.RowType; + +import javax.annotation.Nullable; + +import java.util.Map; +import java.util.function.Supplier; + +/** + * Base for scan providers whose mapping describes exactly one flattened LIST/MAP state via a single + * fixed descriptor (see {@link SingleColumnStateMapping}): {@link + * FlattenedSavepointDataStreamScanProvider} and {@link + * WindowFlattenedSavepointDataStreamScanProvider}. + */ +@Internal +abstract class AbstractSingleColumnScanProvider + extends AbstractSavepointDataStreamScanProvider { + + protected AbstractSingleColumnScanProvider( + @Nullable final String stateBackendType, + final String statePath, + final OperatorIdentifier operatorIdentifier, + final Supplier mappingSupplier, + final RowType rowType, + @Nullable final SavepointKeyFilter keyFilter) { + super(stateBackendType, statePath, operatorIdentifier, mappingSupplier, rowType, keyFilter); + } + + @Override + @SuppressWarnings("rawtypes") + protected final void prepareStateDescriptors(M mapping) { + Map fallbackSchemas = + loadFallbackSchemas( + isSerializerMissing( + mapping.getStateType(), + mapping.getMapKeyTypeSerializer(), + mapping.getValueTypeSerializer())); + + StateDescriptor descriptor = + buildStateDescriptor( + mapping.getStateName(), + mapping.getStateType(), + StateDescriptor.Type.UNKNOWN, + mapping.getMapKeyTypeSerializer(), + mapping.getValueTypeSerializer(), + fallbackSchemas); + mapping.setStateDescriptor(descriptor); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/FlattenedKeyedStateReader.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/FlattenedKeyedStateReader.java new file mode 100644 index 00000000000000..11592e48502de1 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/FlattenedKeyedStateReader.java @@ -0,0 +1,148 @@ +/* + * 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.api.common.functions.OpenContext; +import org.apache.flink.api.common.state.ListState; +import org.apache.flink.api.common.state.ListStateDescriptor; +import org.apache.flink.api.common.state.MapState; +import org.apache.flink.api.common.state.MapStateDescriptor; +import org.apache.flink.api.common.state.State; +import org.apache.flink.state.api.functions.KeyedStateReaderFunction; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.types.RowKind; +import org.apache.flink.util.Collector; + +import java.util.Map; + +/** + * Reads a single flattened keyed list/map state, emitting one row per list element / map entry + * instead of one row per key: {@code (state_key, index, value)} for LIST, {@code (state_key, + * map_key, value)} for MAP. + * + *

Shares value-conversion logic ({@link StateValueConverter}) with {@link KeyedStateReader}. + */ +@SuppressWarnings({"rawtypes", "unchecked"}) +public class FlattenedKeyedStateReader extends KeyedStateReaderFunction { + + private final RowType rowType; + private final FlattenedStateTableMapping mapping; + private final StateValueConverter converter = new StateValueConverter(); + + private transient State state; + + public FlattenedKeyedStateReader(RowType rowType, FlattenedStateTableMapping mapping) { + this.rowType = rowType; + this.mapping = mapping; + } + + @Override + public void open(OpenContext openContext) throws Exception { + switch (mapping.getStateType()) { + case LIST: + state = + getRuntimeContext() + .getListState((ListStateDescriptor) mapping.getStateDescriptor()); + break; + + case MAP: + state = + getRuntimeContext() + .getMapState((MapStateDescriptor) mapping.getStateDescriptor()); + break; + + default: + throw new UnsupportedOperationException( + "Unsupported flattened state type: " + mapping.getStateType()); + } + } + + @Override + public void close() { + state = null; + } + + @Override + public void readKey(Object key, Context context, Collector out) throws Exception { + LogicalType keyLogicalType = + rowType.getFields() + .get(FlattenedStateTableMapping.STATE_KEY_COLUMN_INDEX) + .getType(); + Object convertedKey = converter.getValue(keyLogicalType, key); + + switch (mapping.getStateType()) { + case LIST: + readList(convertedKey, out); + break; + + case MAP: + readMap(convertedKey, out); + break; + + default: + throw new UnsupportedOperationException( + "Unsupported flattened state type: " + mapping.getStateType()); + } + } + + private void readList(Object convertedKey, Collector out) throws Exception { + LogicalType valueLogicalType = + rowType.getFields().get(FlattenedStateTableMapping.VALUE_COLUMN_INDEX).getType(); + LogicalType indexLogicalType = + rowType.getFields().get(FlattenedStateTableMapping.SUB_KEY_COLUMN_INDEX).getType(); + + Iterable values = (Iterable) ((ListState) state).get(); + converter.writeListRows( + values, + () -> { + GenericRowData row = new GenericRowData(RowKind.INSERT, 3); + row.setField(FlattenedStateTableMapping.STATE_KEY_COLUMN_INDEX, convertedKey); + return row; + }, + FlattenedStateTableMapping.SUB_KEY_COLUMN_INDEX, + FlattenedStateTableMapping.VALUE_COLUMN_INDEX, + indexLogicalType, + valueLogicalType, + out); + } + + private void readMap(Object convertedKey, Collector out) throws Exception { + LogicalType valueLogicalType = + rowType.getFields().get(FlattenedStateTableMapping.VALUE_COLUMN_INDEX).getType(); + LogicalType mapKeyLogicalType = + rowType.getFields().get(FlattenedStateTableMapping.SUB_KEY_COLUMN_INDEX).getType(); + + Iterable> entries = ((MapState) state).entries(); + converter.writeMapRows( + entries, + () -> { + GenericRowData row = new GenericRowData(RowKind.INSERT, 3); + row.setField(FlattenedStateTableMapping.STATE_KEY_COLUMN_INDEX, convertedKey); + return row; + }, + FlattenedStateTableMapping.SUB_KEY_COLUMN_INDEX, + FlattenedStateTableMapping.VALUE_COLUMN_INDEX, + mapKeyLogicalType, + valueLogicalType, + out); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/FlattenedSavepointDataStreamScanProvider.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/FlattenedSavepointDataStreamScanProvider.java new file mode 100644 index 00000000000000..650f64d93eb12b --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/FlattenedSavepointDataStreamScanProvider.java @@ -0,0 +1,56 @@ +/* + * 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.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.table.data.RowData; +import org.apache.flink.table.types.logical.RowType; + +import javax.annotation.Nullable; + +import java.util.function.Supplier; + +/** + * Savepoint data stream scan provider for a single flattened keyed LIST/MAP state, emitting one row + * per list element / map entry (see {@link FlattenedKeyedStateReader}). + */ +@SuppressWarnings("rawtypes") +public class FlattenedSavepointDataStreamScanProvider + extends AbstractSingleColumnScanProvider { + + public FlattenedSavepointDataStreamScanProvider( + @Nullable final String stateBackendType, + final String statePath, + final OperatorIdentifier operatorIdentifier, + final Supplier mappingSupplier, + final RowType rowType, + @Nullable final SavepointKeyFilter keyFilter) { + super(stateBackendType, statePath, operatorIdentifier, mappingSupplier, rowType, keyFilter); + } + + @Override + protected DataStream readState( + SavepointReader savepointReader, FlattenedStateTableMapping mapping) throws Exception { + return readVoidNamespaceKeyedState( + savepointReader, mapping, new FlattenedKeyedStateReader(rowType, mapping)); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/FlattenedSavepointDynamicTableSource.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/FlattenedSavepointDynamicTableSource.java new file mode 100644 index 00000000000000..437e29de846b72 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/FlattenedSavepointDynamicTableSource.java @@ -0,0 +1,76 @@ +/* + * 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.OperatorIdentifier; +import org.apache.flink.state.api.filter.SavepointKeyFilter; +import org.apache.flink.table.types.logical.RowType; + +import javax.annotation.Nullable; + +import java.util.function.Supplier; + +/** + * Dynamic source for a table exposing a single flattened LIST/MAP state, i.e. every table kind + * whose mapping describes exactly one such state via a single fixed descriptor (see {@link + * SingleColumnStateMapping}): the plain keyed variant ({@link FlattenedStateTableMapping}, 3-column + * schema) and the namespaced (e.g. window-scoped) variant ({@link + * WindowFlattenedStateTableMapping}, 4-column schema). + * + *

Unlike {@link SavepointDynamicTableSource}, projection push-down is not supported: the schema + * is always exactly {@code (state_key[, state_window], index/map_key, value)}. Filter push-down on + * {@code state_key} is supported (via {@link SavepointKeyFilter}), pruning key groups/keys even + * though {@code state_key} is only part of the composite primary key. + */ +public class FlattenedSavepointDynamicTableSource + extends AbstractSavepointDynamicTableSource { + + public FlattenedSavepointDynamicTableSource( + @Nullable final String stateBackendType, + final String statePath, + final OperatorIdentifier operatorIdentifier, + final int keyColumnIndex, + final Supplier mappingSupplier, + final RowType rowType, + final String summaryString, + final ScanProviderFactory scanProviderFactory) { + super( + stateBackendType, + statePath, + operatorIdentifier, + keyColumnIndex, + mappingSupplier, + rowType, + summaryString, + scanProviderFactory); + } + + @Override + protected AbstractSavepointDynamicTableSource newInstance() { + return new FlattenedSavepointDynamicTableSource<>( + stateBackendType, + statePath, + operatorIdentifier, + keyColumnIndex, + mappingSupplier, + rowType, + summaryString, + scanProviderFactory); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/FlattenedStateTableMapping.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/FlattenedStateTableMapping.java new file mode 100644 index 00000000000000..553ddeddbd4241 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/FlattenedStateTableMapping.java @@ -0,0 +1,234 @@ +/* + * 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.annotation.Internal; +import org.apache.flink.api.common.serialization.SerializerConfig; +import org.apache.flink.api.common.state.StateDescriptor; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.state.api.OperatorIdentifier; +import org.apache.flink.table.api.ValidationException; +import org.apache.flink.table.catalog.Column; +import org.apache.flink.table.catalog.ResolvedCatalogTable; +import org.apache.flink.table.catalog.ResolvedSchema; +import org.apache.flink.table.catalog.UniqueConstraint; +import org.apache.flink.table.types.DataType; +import org.apache.flink.table.types.logical.LogicalTypeRoot; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.util.Preconditions; + +import javax.annotation.Nullable; + +import java.io.Serializable; +import java.util.List; + +/** + * Maps the fixed 3-column schema of a flattened keyed list/map state table: + * + *

    + *
  • LIST: {@code (state_key, list_index, list_value)} + *
  • MAP: {@code (state_key, map_key, map_value)} + *
+ * + *

The third column has a fixed name ({@code list_value}/{@code map_value}) rather than being + * named after the flattened state itself, to avoid collisions with other (reserved) column names; + * the true state name is instead resolved from {@link + * SavepointConnectorOptions#FLATTENED_STATE_NAME}. + * + *

A flattened table always exposes exactly one keyed state and emits one row per list element / + * map entry (as opposed to one row per key), so unlike {@link StateTableMapping} there is no + * per-column projection bookkeeping: column indices in the (fixed) output row are always {@code + * 0=state_key}, {@code 1=list_index/map_key}, {@code 2=list_value/map_value}. + */ +@Internal +public class FlattenedStateTableMapping implements Serializable, SingleColumnStateMapping { + + private static final long serialVersionUID = 1L; + + public static final int STATE_KEY_COLUMN_INDEX = 0; + public static final int SUB_KEY_COLUMN_INDEX = 1; + public static final int VALUE_COLUMN_INDEX = 2; + + private final String stateName; + private final SavepointConnectorOptions.StateType stateType; + private final TypeInformation keyTypeInfo; + @Nullable private final TypeSerializer mapKeyTypeSerializer; + private final TypeSerializer valueTypeSerializer; + @Nullable private StateDescriptor stateDescriptor; + + public FlattenedStateTableMapping( + String stateName, + SavepointConnectorOptions.StateType stateType, + TypeInformation keyTypeInfo, + @Nullable TypeSerializer mapKeyTypeSerializer, + TypeSerializer valueTypeSerializer) { + Preconditions.checkArgument( + stateType == SavepointConnectorOptions.StateType.LIST + || stateType == SavepointConnectorOptions.StateType.MAP, + "Flattened state tables only support LIST and MAP states, got: " + stateType); + this.stateName = stateName; + this.stateType = stateType; + this.keyTypeInfo = keyTypeInfo; + this.mapKeyTypeSerializer = mapKeyTypeSerializer; + this.valueTypeSerializer = valueTypeSerializer; + } + + @Override + public String getStateName() { + return stateName; + } + + @Override + public SavepointConnectorOptions.StateType getStateType() { + return stateType; + } + + @Override + public TypeInformation getKeyTypeInfo() { + return keyTypeInfo; + } + + @Override + @Nullable + public TypeSerializer getMapKeyTypeSerializer() { + return mapKeyTypeSerializer; + } + + @Override + public TypeSerializer getValueTypeSerializer() { + return valueTypeSerializer; + } + + @Override + @SuppressWarnings("rawtypes") + public void setStateDescriptor(StateDescriptor stateDescriptor) { + this.stateDescriptor = stateDescriptor; + } + + @Nullable + @SuppressWarnings("rawtypes") + public StateDescriptor getStateDescriptor() { + return stateDescriptor; + } + + // ------------------------------------------------------------------------- + // Factory + // ------------------------------------------------------------------------- + + /** + * Validates that the table schema matches the fixed 3-column flattened layout with a composite + * primary key on {@code (state_key, list_index/map_key)}, and returns the state type (LIST or + * MAP), inferred from whether the second column is named {@code list_index} or {@code map_key}. + * This is a purely structural check; it performs no I/O or class loading. + */ + public static SavepointConnectorOptions.StateType validateFlattenedSchema( + ResolvedCatalogTable catalogTable) { + ResolvedSchema schema = catalogTable.getResolvedSchema(); + List columns = schema.getColumns(); + if (columns.size() != 3) { + throw new ValidationException( + "Flattened keyed state tables must have exactly 3 columns " + + "(state_key, list_index/map_key, list_value/map_value), but found " + + columns.size() + + "."); + } + DataType physicalDataType = schema.toPhysicalRowDataType(); + Preconditions.checkArgument( + physicalDataType.getLogicalType().is(LogicalTypeRoot.ROW), + "Row data type expected."); + + String stateKeyColumnName = columns.get(STATE_KEY_COLUMN_INDEX).getName(); + String subKeyColumnName = columns.get(SUB_KEY_COLUMN_INDEX).getName(); + String valueColumnName = columns.get(VALUE_COLUMN_INDEX).getName(); + SavepointConnectorOptions.StateType stateType = + TableMappingSupport.inferFlattenedStateTypeAndValidateValueColumn( + "Flattened keyed state tables", + "second", + "third", + subKeyColumnName, + valueColumnName); + + List expectedKeyColumns = List.of(stateKeyColumnName, subKeyColumnName); + List primaryKeyColumns = + schema.getPrimaryKey().map(UniqueConstraint::getColumns).orElse(List.of()); + if (!primaryKeyColumns.equals(expectedKeyColumns)) { + throw new ValidationException( + "Flattened keyed state tables must declare a composite primary key on (" + + stateKeyColumnName + + ", " + + subKeyColumnName + + "), but found: " + + (primaryKeyColumns.isEmpty() ? "none" : primaryKeyColumns) + + "."); + } + + return stateType; + } + + /** + * Builds a complete {@link FlattenedStateTableMapping}, loading operator state metadata from + * the savepoint and resolving serializers and key type from it. + * + *

Assumes {@link #validateFlattenedSchema} has already been called for this table. + * + *

This performs I/O (savepoint metadata loading); callers should invoke it lazily, deferred + * to scan time, to keep planning free of savepoint access. + * + * @param catalogTable the resolved table whose schema drives the mapping + * @param stateName the name of the flattened LIST/MAP state, resolved from {@link + * SavepointConnectorOptions#FLATTENED_STATE_NAME} + * @param statePath path to the savepoint containing the operator state metadata + * @param operatorIdentifier identifies the operator whose state metadata is loaded + * @param serializerConfig serializer config used when creating serializers from resolved types + * @param stateType {@code LIST} or {@code MAP} + */ + public static FlattenedStateTableMapping from( + ResolvedCatalogTable catalogTable, + String stateName, + String statePath, + OperatorIdentifier operatorIdentifier, + SerializerConfig serializerConfig, + SavepointConnectorOptions.StateType stateType) { + + SavepointTypeInfoResolver typeResolver = + TableMappingSupport.createTypeResolver( + statePath, operatorIdentifier, serializerConfig); + + DataType physicalDataType = catalogTable.getResolvedSchema().toPhysicalRowDataType(); + RowType rowType = (RowType) physicalDataType.getLogicalType(); + + TableMappingSupport.FlattenedSerializers serializers = + TableMappingSupport.resolveFlattenedSerializers( + rowType, + typeResolver, + stateName, + stateType, + STATE_KEY_COLUMN_INDEX, + SUB_KEY_COLUMN_INDEX, + VALUE_COLUMN_INDEX); + + return new FlattenedStateTableMapping( + stateName, + stateType, + serializers.keyTypeInfo, + serializers.mapKeyTypeSerializer, + serializers.valueTypeSerializer); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/KeyedStateReader.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/KeyedStateReader.java index d3c88f06d43daf..975949ca385ac3 100644 --- a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/KeyedStateReader.java +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/KeyedStateReader.java @@ -19,80 +19,62 @@ package org.apache.flink.state.table; import org.apache.flink.api.common.functions.OpenContext; +import org.apache.flink.api.common.state.AggregatingStateDescriptor; import org.apache.flink.api.common.state.ListState; import org.apache.flink.api.common.state.ListStateDescriptor; import org.apache.flink.api.common.state.MapState; import org.apache.flink.api.common.state.MapStateDescriptor; +import org.apache.flink.api.common.state.ReducingStateDescriptor; import org.apache.flink.api.common.state.State; -import org.apache.flink.api.common.state.ValueState; +import org.apache.flink.api.common.state.StateDescriptor; import org.apache.flink.api.common.state.ValueStateDescriptor; -import org.apache.flink.api.java.tuple.Tuple2; import org.apache.flink.state.api.functions.KeyedStateReaderFunction; -import org.apache.flink.table.data.DecimalData; -import org.apache.flink.table.data.GenericArrayData; -import org.apache.flink.table.data.GenericMapData; import org.apache.flink.table.data.GenericRowData; import org.apache.flink.table.data.RowData; -import org.apache.flink.table.data.StringData; -import org.apache.flink.table.types.logical.ArrayType; -import org.apache.flink.table.types.logical.DecimalType; import org.apache.flink.table.types.logical.LogicalType; -import org.apache.flink.table.types.logical.MapType; import org.apache.flink.table.types.logical.RowType; import org.apache.flink.types.RowKind; import org.apache.flink.util.Collector; -import org.apache.flink.shaded.guava33.com.google.common.cache.Cache; -import org.apache.flink.shaded.guava33.com.google.common.cache.CacheBuilder; - -import java.lang.reflect.Field; -import java.lang.reflect.InvocationTargetException; -import java.lang.reflect.Method; -import java.math.BigDecimal; -import java.nio.ByteBuffer; import java.util.HashMap; -import java.util.Iterator; import java.util.List; import java.util.Map; -import java.util.concurrent.ExecutionException; -import java.util.stream.StreamSupport; /** Keyed state reader function for value, list and map state types. */ @SuppressWarnings({"rawtypes", "unchecked"}) public class KeyedStateReader extends KeyedStateReaderFunction { - private static final long CACHE_MAX_SIZE = 64000L; - private final Tuple2> keyValueProjections; + private final StateTableMapping mapping; private final RowType rowType; - private final Map states = new HashMap<>(); - private final Cache, Field> classFieldCache; - private final Cache, Method> classMethodCache; - - public KeyedStateReader( - final RowType rowType, - final Tuple2> keyValueProjections) { - this.keyValueProjections = keyValueProjections; + private final StateValueConverter converter = new StateValueConverter(); + + /** + * States keyed by state name. Populated from {@link StateTableMapping#getAllValueColumns()} so + * that ALL original states are registered with the Flink runtime (required for key enumeration + * in {@code KeyedStateReaderOperator.getKeysAndNamespaces}), even when some columns have been + * projected out. + */ + private final Map states = new HashMap<>(); + + public KeyedStateReader(final RowType rowType, final StateTableMapping mapping) { + this.mapping = mapping; this.rowType = rowType; - this.classMethodCache = CacheBuilder.newBuilder().maximumSize(CACHE_MAX_SIZE).build(); - this.classFieldCache = CacheBuilder.newBuilder().maximumSize(CACHE_MAX_SIZE).build(); } @Override public void open(OpenContext openContext) throws Exception { - for (StateValueColumnConfiguration columnConfig : keyValueProjections.f1) { + // Register ALL original value columns so that key enumeration always works, + // even when a projection has removed some (or all) value columns from the output. + for (StateValueColumnConfiguration columnConfig : mapping.getAllValueColumns()) { switch (columnConfig.getStateType()) { case VALUE: states.put( - columnConfig.getColumnIndex(), - getRuntimeContext() - .getState( - (ValueStateDescriptor) - columnConfig.getStateDescriptor())); + columnConfig.getStateName(), getOrCreateValueLikeState(columnConfig)); break; case LIST: states.put( - columnConfig.getColumnIndex(), + columnConfig.getStateName(), getRuntimeContext() .getListState( (ListStateDescriptor) @@ -101,7 +83,7 @@ public void open(OpenContext openContext) throws Exception { case MAP: states.put( - columnConfig.getColumnIndex(), + columnConfig.getStateName(), getRuntimeContext() .getMapState( (MapStateDescriptor) @@ -115,6 +97,27 @@ public void open(OpenContext openContext) throws Exception { } } + /** + * Registers the VALUE-shaped state for {@code columnConfig}, using the state-getter matching + * its {@link StateValueColumnConfiguration#getActualStateKind()} (e.g. {@code + * getReducingState}/{@code getAggregatingState} for a {@code .reduce()}/{@code .aggregate()} + * window function's window-contents state, which is registered as {@code REDUCING}/{@code + * AGGREGATING} rather than plain {@code VALUE}). + */ + private State getOrCreateValueLikeState(StateValueColumnConfiguration columnConfig) + throws Exception { + StateDescriptor descriptor = columnConfig.getStateDescriptor(); + switch (columnConfig.getActualStateKind()) { + case REDUCING: + return getRuntimeContext().getReducingState((ReducingStateDescriptor) descriptor); + case AGGREGATING: + return getRuntimeContext() + .getAggregatingState((AggregatingStateDescriptor) descriptor); + default: + return getRuntimeContext().getState((ValueStateDescriptor) descriptor); + } + } + @Override public void close() { states.clear(); @@ -122,43 +125,42 @@ public void close() { @Override public void readKey(Object key, Context context, Collector out) throws Exception { - GenericRowData row = new GenericRowData(RowKind.INSERT, 1 + keyValueProjections.f1.size()); + GenericRowData row = new GenericRowData(RowKind.INSERT, rowType.getFieldCount()); List fields = rowType.getFields(); - // Fill column from key - int columnIndex = keyValueProjections.f0; - LogicalType keyLogicalType = fields.get(columnIndex).getType(); - row.setField(columnIndex, getValue(keyLogicalType, key)); + int columnIndex = mapping.getKeyColumnIndex(); + if (columnIndex >= 0) { + LogicalType keyLogicalType = fields.get(columnIndex).getType(); + row.setField(columnIndex, converter.getValue(keyLogicalType, key)); + } - // Fill columns from values - for (StateValueColumnConfiguration columnConfig : keyValueProjections.f1) { + // Only write the projected value columns to the output row. + for (StateValueColumnConfiguration columnConfig : mapping.getValueColumns()) { LogicalType valueLogicalType = fields.get(columnConfig.getColumnIndex()).getType(); + State state = states.get(columnConfig.getStateName()); switch (columnConfig.getStateType()) { case VALUE: row.setField( columnConfig.getColumnIndex(), - getValue( + converter.getValue( valueLogicalType, - ((ValueState) states.get(columnConfig.getColumnIndex())) - .value())); + StateValueConverter.readValueLikeState( + state, columnConfig.getActualStateKind()))); break; case LIST: row.setField( columnConfig.getColumnIndex(), - getValue( + converter.getValue( valueLogicalType, - ((ListState) states.get(columnConfig.getColumnIndex())).get())); + StateValueConverter.readListLikeState((ListState) state))); break; case MAP: row.setField( columnConfig.getColumnIndex(), - getValue( - valueLogicalType, - ((MapState) states.get(columnConfig.getColumnIndex())) - .entries())); + converter.getValue(valueLogicalType, ((MapState) state).entries())); break; default: @@ -169,209 +171,4 @@ public void readKey(Object key, Context context, Collector out) throws out.collect(row); } - - private Object getValue(LogicalType logicalType, Object object) { - if (object == null) { - return null; - } - switch (logicalType.getTypeRoot()) { - case CHAR: // String - case VARCHAR: // String - return StringData.fromString(object.toString()); - - case BOOLEAN: // Boolean - return object; - - case BINARY: // byte[] - case VARBINARY: // ByteBuffer, byte[] - return convertToBytes(object); - - case DECIMAL: // BigDecimal, ByteBuffer, byte[] - return convertToDecimal(object, logicalType); - - case TINYINT: // Byte - case SMALLINT: // Short - case INTEGER: // Integer - case BIGINT: // Long - case FLOAT: // Float - case DOUBLE: // Double - case DATE: // Integer - return object; - - case INTERVAL_YEAR_MONTH: // Long - case INTERVAL_DAY_TIME: // Long - return object; - - case ARRAY: - return convertToArray(object, logicalType); - - case MAP: - return convertToMap(object, logicalType); - - case ROW: - return convertToRow(object, logicalType); - - case NULL: - return null; - - case MULTISET: - case TIME_WITHOUT_TIME_ZONE: - case TIMESTAMP_WITHOUT_TIME_ZONE: - case TIMESTAMP_WITH_TIME_ZONE: - case TIMESTAMP_WITH_LOCAL_TIME_ZONE: - case DISTINCT_TYPE: - case STRUCTURED_TYPE: - case RAW: - case SYMBOL: - case UNRESOLVED: - case DESCRIPTOR: - default: - throw new UnsupportedOperationException("Unsupported type: " + logicalType); - } - } - - private Object getObjectField(Object object, RowType.RowField rowField) { - String rowFieldName = rowField.getName(); - - Class objectClass = object.getClass(); - Object objectField; - try { - Field field = - classFieldCache.get( - Tuple2.of(objectClass, rowFieldName), - () -> objectClass.getField(rowFieldName)); - objectField = field.get(object); - } catch (ExecutionException e1) { - Method method = getMethod(objectClass, rowFieldName); - try { - objectField = method.invoke(object); - } catch (IllegalAccessException | InvocationTargetException e2) { - throw new RuntimeException(e2); - } - } catch (IllegalAccessException e) { - throw new UnsupportedOperationException( - "Cannot access field by either public member or getter function: " - + rowField.getName()); - } - - return objectField; - } - - private Method getMethod(Class objectClass, String rowFieldName) { - String upperRowFieldName = - rowFieldName.substring(0, 1).toUpperCase() + rowFieldName.substring(1); - try { - String methodName = "get" + upperRowFieldName; - return classMethodCache.get( - Tuple2.of(objectClass, methodName), () -> objectClass.getMethod(methodName)); - } catch (ExecutionException e1) { - try { - String methodName = "is" + upperRowFieldName; - return classMethodCache.get( - Tuple2.of(objectClass, methodName), - () -> objectClass.getMethod(methodName)); - } catch (ExecutionException e2) { - throw new RuntimeException(e2); - } - } - } - - private static DecimalData convertToDecimal(Object object, LogicalType logicalType) { - DecimalType decimalType = (DecimalType) logicalType; - - final int precision = decimalType.getPrecision(); - final int scale = decimalType.getScale(); - if (object instanceof BigDecimal) { - return DecimalData.fromBigDecimal((BigDecimal) object, precision, scale); - } else if (object instanceof ByteBuffer) { - ByteBuffer byteBuffer = (ByteBuffer) object; - final byte[] bytes = new byte[byteBuffer.remaining()]; - byteBuffer.get(bytes); - return DecimalData.fromUnscaledBytes(bytes, precision, scale); - } else if (object instanceof byte[]) { - final byte[] bytes = (byte[]) object; - return DecimalData.fromUnscaledBytes(bytes, precision, scale); - } else { - throw new UnsupportedOperationException( - "Decimal conversion supports only BigDecimal, ByteBuffer and byte[] but received: " - + object.getClass().getName()); - } - } - - private byte[] convertToBytes(Object object) { - if (object instanceof ByteBuffer) { - ByteBuffer byteBuffer = (ByteBuffer) object; - byte[] bytes = new byte[byteBuffer.remaining()]; - byteBuffer.get(bytes); - return bytes; - } else if (object instanceof byte[]) { - return (byte[]) object; - } else { - throw new UnsupportedOperationException( - "Byte array conversion supports only ByteBuffer and byte[] but received: " - + object.getClass().getName()); - } - } - - private GenericArrayData convertToArray(Object object, LogicalType logicalType) { - LogicalType elementLogicalType = ((ArrayType) logicalType).getElementType(); - - if (object instanceof Iterable) { - Iterable iterable = (Iterable) object; - return new GenericArrayData( - StreamSupport.stream(iterable.spliterator(), false) - .map(v -> getValue(elementLogicalType, v)) - .toArray()); - } else { - throw new UnsupportedOperationException( - "Array conversion supports only Iterable but received: " - + object.getClass().getName()); - } - } - - private GenericMapData convertToMap(Object object, LogicalType logicalType) { - MapType mapType = (MapType) logicalType; - LogicalType keyLogicalType = mapType.getKeyType(); - LogicalType valueLogicalType = mapType.getValueType(); - - if (object instanceof Iterable) { - Iterable iterable = (Iterable) object; - Iterator iterator = iterable.iterator(); - Map result = new HashMap<>(); - boolean typeChecked = false; - while (iterator.hasNext()) { - Object e = iterator.next(); - // The boolean check here is for performance tuning because instanceof is slow, and - // it's enough to check the type only once. - if (!typeChecked && !(e instanceof Map.Entry)) { - throw new UnsupportedOperationException( - "Map conversion supports only Iterable but received: " - + object.getClass().getName()); - } else { - typeChecked = true; - } - Map.Entry entry = (Map.Entry) e; - result.put( - getValue(keyLogicalType, entry.getKey()), - getValue(valueLogicalType, entry.getValue())); - } - return new GenericMapData(result); - } else { - throw new UnsupportedOperationException( - "Map conversion supports only Iterable but received: " - + object.getClass().getName()); - } - } - - private GenericRowData convertToRow(Object object, LogicalType logicalType) { - RowType rowType = (RowType) logicalType; - GenericRowData result = new GenericRowData(RowKind.INSERT, rowType.getFieldCount()); - List fields = rowType.getFields(); - for (int i = 0; i < rowType.getFieldCount(); i++) { - RowType.RowField subRowField = fields.get(i); - Object subObject = getObjectField(object, subRowField); - result.setField(i, getValue(subRowField.getType(), subObject)); - } - return result; - } } diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/MultiColumnStateMapping.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/MultiColumnStateMapping.java new file mode 100644 index 00000000000000..766aebed8c5042 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/MultiColumnStateMapping.java @@ -0,0 +1,46 @@ +/* + * 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.annotation.Internal; + +import java.util.List; + +/** + * Mixed into mapping classes that back a table with an arbitrary number of value columns, each + * registered as its own state, driven from {@link #getAllValueColumns()}. Implemented by {@link + * StateTableMapping} and {@link WindowStateTableMapping}. + */ +@Internal +interface MultiColumnStateMapping extends SavepointStateMapping { + + /** + * Full original value columns required for state-descriptor registration / key(-and-namespace) + * enumeration, preserved across projections. Never empty when the source table has at least one + * state column. + */ + List getAllValueColumns(); + + /** + * Creates a new mapping of the same concrete type with column indices remapped to the projected + * output. Declared here (rather than per concrete class) so that {@link + * SavepointDynamicTableSource} can apply projection push-down generically. + */ + MultiColumnStateMapping project(int[][] projectedFields); +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointConnectorOptions.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointConnectorOptions.java index daa6022071c790..f584d8ab612015 100644 --- a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointConnectorOptions.java +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointConnectorOptions.java @@ -32,13 +32,6 @@ public class SavepointConnectorOptions { public static final String FIELDS = "fields"; public static final String STATE_NAME = "state-name"; - public static final String STATE_TYPE = "state-type"; - public static final String DEPRECATED_MAP_KEY_FORMAT = "map-key-format"; - public static final String KEY_CLASS = "key-class"; - public static final String DEPRECATED_VALUE_FORMAT = "value-format"; - public static final String VALUE_CLASS = "value-class"; - public static final String KEY_TYPE_FACTORY = "key-type-factory"; - public static final String VALUE_TYPE_FACTORY = "value-type-factory"; /** Value state types. */ public enum StateType { @@ -47,6 +40,51 @@ public enum StateType { MAP } + /** Determines how a savepoint table's schema and rows are derived from keyed state. */ + public enum StateReaderMode { + + /** One row per key, one column per keyed state (the general keyed-state table). */ + KEYED("keyed"), + + /** + * Exposes a single keyed LIST/MAP state flattened into one row per list element / map + * entry, instead of one row per key. + */ + KEYED_FLAT("keyed-flat"), + + /** + * One row per (key, namespace), one column per VALUE-shaped namespaced state (e.g. a window + * operator's window-contents accumulator or window-registered value state). + */ + WINDOWED("windowed"), + + /** + * Exposes a single namespaced LIST/MAP state flattened into one row per list element / map + * entry, instead of one row per (key, namespace). + */ + WINDOWED_FLAT("windowed-flat"), + + /** One row per element of an operator {@code ListState} (no key/namespace concept). */ + LIST("list"), + + /** One row per element of an operator {@code UnionState} (no key/namespace concept). */ + UNION("union"), + + /** One row per entry of an operator {@code BroadcastState} (no key/namespace concept). */ + BROADCAST("broadcast"); + + private final String value; + + StateReaderMode(String value) { + this.value = value; + } + + @Override + public String toString() { + return value; + } + } + // -------------------------------------------------------------------------------------------- // Common options // -------------------------------------------------------------------------------------------- @@ -94,6 +132,33 @@ public enum StateType { .withDescription( "Defines the operator UID hash which must be used for state reading (Can't be used together with UID)."); + /** + * Determines whether the table exposes the general keyed-state schema (one row per key, one + * column per keyed state) or the flattened schema for a single LIST/MAP state (one row per list + * element / map entry). Set automatically by {@code StateCatalog}; not intended for manual use. + * In flattened mode, the state type (LIST or MAP) is not configured separately: it is inferred + * from whether the table's second column is named {@code list_index} (LIST) or {@code map_key} + * (MAP), and the state name is inferred from the name of the third column. + */ + public static final ConfigOption STATE_READER_MODE = + ConfigOptions.key("state.reader.mode") + .enumType(StateReaderMode.class) + .defaultValue(StateReaderMode.KEYED) + .withDescription( + Description.builder() + .text( + "Determines whether the table exposes the general keyed-state schema " + + "(%s, the default) or the flattened schema for a single LIST/MAP " + + "state (%s), which exposes one row per list element / map entry " + + "instead of one row per key, or one of the namespaced-state " + + "equivalents (%s, %s), which expose state registered under a " + + "non-void namespace (e.g. window-scoped state).", + code(StateReaderMode.KEYED.toString()), + code(StateReaderMode.KEYED_FLAT.toString()), + code(StateReaderMode.WINDOWED.toString()), + code(StateReaderMode.WINDOWED_FLAT.toString())) + .build()); + // -------------------------------------------------------------------------------------------- // Value options // -------------------------------------------------------------------------------------------- @@ -106,73 +171,23 @@ public enum StateType { .withDescription( "Defines the state name which must be used for state reading."); - /** Placeholder {@link ConfigOption}. Not used for retrieving values. */ - public static final ConfigOption STATE_TYPE_PLACEHOLDER = - ConfigOptions.key(String.format("%s.#.%s", FIELDS, STATE_TYPE)) - .enumType(StateType.class) - .noDefaultValue() - .withDescription( - Description.builder() - .text( - "Defines the state type which must be used for state reading, including %s, %s and %s. " - + "When it's not provided then it tries to be inferred from the SQL type (ARRAY=list, MAP=map, all others=value).", - code(StateType.VALUE.toString()), - code(StateType.LIST.toString()), - code(StateType.MAP.toString())) - .build()); - - /** Placeholder {@link ConfigOption}. Not used for retrieving values. */ - public static final ConfigOption KEY_CLASS_PLACEHOLDER = - ConfigOptions.key(String.format("%s.#.%s", FIELDS, KEY_CLASS)) + /** + * Explicitly identifies the single LIST/MAP/UNION/BROADCAST state exposed by a table whose + * value column(s) no longer carry the state's name themselves — either because the value is + * flattened directly into top-level columns (one column per field for a structured value, or a + * single value column named {@code list_value}/{@code map_value} for a scalar one, used by + * {@link #STATE_READER_MODE} {@code KEYED_FLAT}, {@code WINDOWED_FLAT}, {@code LIST}, and + * {@code UNION}), or because the value column has a fixed name ({@code map_value}, used by + * {@code BROADCAST}). Naming the value column after the state itself risked colliding with a + * table's other (reserved) column names, e.g. a state literally named {@code map_key}. Set + * automatically by {@code StateCatalog}; not intended for manual use. + */ + public static final ConfigOption FLATTENED_STATE_NAME = + ConfigOptions.key(STATE_NAME) .stringType() .noDefaultValue() - .withDeprecatedKeys(String.format("%s.#.%s", FIELDS, DEPRECATED_MAP_KEY_FORMAT)) .withDescription( - "Defines the format class scheme for decoding map key data. " - + "When it's not provided then it tries to be inferred from the SQL type (only primitive types supported)."); - - /** Placeholder {@link ConfigOption}. Not used for retrieving values. */ - public static final ConfigOption KEY_TYPE_INFO_FACTORY_PLACEHOLDER = - ConfigOptions.key(String.format("%s.#.%s", FIELDS, KEY_TYPE_FACTORY)) - .stringType() - .noDefaultValue() - .withDescription( - Description.builder() - .text( - "Defines the type information factory for decoding map key data. " - + "Either %s or %s can be specified. " - + "When none of them are provided then the format class scheme tries to be inferred from the SQL type (only primitive types supported).", - code(KEY_CLASS), code(KEY_TYPE_FACTORY)) - .build()); - - /** Placeholder {@link ConfigOption}. Not used for retrieving values. */ - public static final ConfigOption VALUE_CLASS_PLACEHOLDER = - ConfigOptions.key(String.format("%s.#.%s", FIELDS, VALUE_CLASS)) - .stringType() - .noDefaultValue() - .withDeprecatedKeys(String.format("%s.#.%s", FIELDS, DEPRECATED_VALUE_FORMAT)) - .withDescription( - Description.builder() - .text( - "Defines the format class scheme for decoding value data. " - + "Either %s or %s can be specified. " - + "When none of them are provided then format class scheme tries to be inferred from the SQL type (only primitive types supported).", - code(VALUE_CLASS), code(VALUE_TYPE_FACTORY)) - .build()); - - /** Placeholder {@link ConfigOption}. Not used for retrieving values. */ - public static final ConfigOption VALUE_TYPE_INFO_FACTORY_PLACEHOLDER = - ConfigOptions.key(String.format("%s.#.%s", FIELDS, VALUE_TYPE_FACTORY)) - .stringType() - .noDefaultValue() - .withDescription( - Description.builder() - .text( - "Defines the type information factory for decoding value data. " - + "Either %s or %s can be specified. " - + "When none of them are provided then the format class scheme tries to be inferred from the SQL type (only primitive types supported).", - code(VALUE_CLASS), code(VALUE_TYPE_FACTORY)) - .build()); + "Defines the name of the single LIST/MAP/UNION/BROADCAST state exposed by this table."); private SavepointConnectorOptions() {} } diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointConnectorOptionsUtil.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointConnectorOptionsUtil.java index bdaba67411b797..fcff8f302327ba 100644 --- a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointConnectorOptionsUtil.java +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointConnectorOptionsUtil.java @@ -35,9 +35,21 @@ public static OperatorIdentifier getOperatorIdentifier(ReadableConfig options) { final Optional operatorUid = options.getOptional(OPERATOR_UID); final Optional operatorUidHash = options.getOptional(OPERATOR_UID_HASH); - if (operatorUid.isPresent() == operatorUidHash.isPresent()) { + if (operatorUid.isPresent() && operatorUidHash.isPresent()) { throw new IllegalArgumentException( - "Either operator uid or operator uid hash must be specified."); + "Options '" + + OPERATOR_UID.key() + + "' and '" + + OPERATOR_UID_HASH.key() + + "' cannot be specified together."); + } + if (operatorUid.isEmpty() && operatorUidHash.isEmpty()) { + throw new IllegalArgumentException( + "Either '" + + OPERATOR_UID.key() + + "' or '" + + OPERATOR_UID_HASH.key() + + "' must be specified."); } return operatorUid 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 35a9e04a7af8c5..c9ec617111c5d1 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 @@ -18,146 +18,39 @@ package org.apache.flink.state.table; -import org.apache.flink.api.common.state.ListStateDescriptor; -import org.apache.flink.api.common.state.MapStateDescriptor; -import org.apache.flink.api.common.state.ValueStateDescriptor; -import org.apache.flink.api.common.typeinfo.TypeInformation; -import org.apache.flink.api.common.typeutils.TypeSerializer; -import org.apache.flink.api.java.tuple.Tuple2; -import org.apache.flink.configuration.Configuration; -import org.apache.flink.configuration.StateBackendOptions; -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; -import org.apache.flink.table.connector.source.DataStreamScanProvider; import org.apache.flink.table.data.RowData; -import org.apache.flink.table.runtime.typeutils.InternalTypeInfo; import org.apache.flink.table.types.logical.RowType; -import org.apache.flink.util.StringUtils; import javax.annotation.Nullable; -import javax.naming.ConfigurationException; -import java.util.List; +import java.util.function.Supplier; -/** Savepoint data stream scan provider. */ +/** + * Savepoint data stream scan provider for the general keyed state table, emitting one row per key + * (see {@link KeyedStateReader}). + */ @SuppressWarnings("rawtypes") -public class SavepointDataStreamScanProvider implements DataStreamScanProvider { - @Nullable private final String stateBackendType; - private final String statePath; - private final OperatorIdentifier operatorIdentifier; - private final TypeInformation keyTypeInfo; - private final Tuple2> keyValueProjections; - private final RowType rowType; - @Nullable private final SavepointKeyFilter keyFilter; +public class SavepointDataStreamScanProvider + extends AbstractMultiColumnScanProvider { public SavepointDataStreamScanProvider( @Nullable final String stateBackendType, final String statePath, final OperatorIdentifier operatorIdentifier, - final TypeInformation keyTypeInfo, - final Tuple2> 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> keyValueProjections, + final Supplier mappingSupplier, 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 - public boolean isBounded() { - return true; + super(stateBackendType, statePath, operatorIdentifier, mappingSupplier, rowType, keyFilter); } @Override - @SuppressWarnings({"rawtypes", "unchecked"}) - public DataStream produceDataStream( - ProviderContext providerContext, StreamExecutionEnvironment execEnv) { - try { - Configuration configuration = Configuration.fromMap(execEnv.getConfiguration().toMap()); - if (!StringUtils.isNullOrWhitespaceOnly(stateBackendType)) { - configuration.set(StateBackendOptions.STATE_BACKEND, stateBackendType); - } - StateBackend stateBackend = - StateBackendLoader.loadStateBackendFromConfig( - configuration, getClass().getClassLoader(), null); - - SavepointReader savepointReader = - SavepointReader.read(execEnv, statePath, stateBackend); - - // Get value state descriptors - for (StateValueColumnConfiguration columnConfig : keyValueProjections.f1) { - TypeSerializer valueTypeSerializer = columnConfig.getValueTypeSerializer(); - - switch (columnConfig.getStateType()) { - case VALUE: - columnConfig.setStateDescriptor( - new ValueStateDescriptor<>( - columnConfig.getStateName(), valueTypeSerializer)); - break; - - case LIST: - columnConfig.setStateDescriptor( - new ListStateDescriptor<>( - columnConfig.getStateName(), valueTypeSerializer)); - break; - - case MAP: - TypeSerializer mapKeyTypeSerializer = - columnConfig.getMapKeyTypeSerializer(); - if (mapKeyTypeSerializer == null) { - throw new ConfigurationException( - "Map key type serializer is required for map state"); - } - columnConfig.setStateDescriptor( - new MapStateDescriptor<>( - columnConfig.getStateName(), - mapKeyTypeSerializer, - valueTypeSerializer)); - break; - - default: - throw new UnsupportedOperationException( - "Unsupported state type: " + columnConfig.getStateType()); - } - } - - TypeInformation outTypeInfo = InternalTypeInfo.of(rowType); - - return savepointReader.readKeyedState( - operatorIdentifier, - new KeyedStateReader(rowType, keyValueProjections), - keyTypeInfo, - outTypeInfo, - keyFilter); - } catch (Exception e) { - throw new RuntimeException(e); - } + protected DataStream readState( + SavepointReader savepointReader, StateTableMapping mapping) throws Exception { + return readVoidNamespaceKeyedState( + savepointReader, mapping, new KeyedStateReader(rowType, mapping)); } } 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 bc8e30c243ebd4..4df2d94e925d7d 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 @@ -18,94 +18,70 @@ 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.connector.source.abilities.SupportsProjectionPushDown; +import org.apache.flink.table.types.DataType; import org.apache.flink.table.types.logical.RowType; -import org.apache.flink.table.types.utils.TypeConversions; import javax.annotation.Nullable; -import java.util.List; +import java.util.function.Supplier; -/** Savepoint dynamic source. */ -@SuppressWarnings("rawtypes") -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> keyValueProjections; - private final RowType rowType; - @Nullable private SavepointKeyFilter keyFilter; +/** + * Dynamic source for the general keyed/namespaced state tables, i.e. every table kind whose mapping + * registers an arbitrary number of value columns (see {@link MultiColumnStateMapping}): the plain + * keyed-state table ({@link StateTableMapping}) and the namespaced (e.g. window-scoped) keyed-state + * table ({@link WindowStateTableMapping}). + */ +public class SavepointDynamicTableSource + extends AbstractSavepointDynamicTableSource implements SupportsProjectionPushDown { public SavepointDynamicTableSource( @Nullable final String stateBackendType, final String statePath, final OperatorIdentifier operatorIdentifier, - final TypeInformation keyTypeInfo, - final Tuple2> keyValueProjections, - RowType rowType) { - this.stateBackendType = stateBackendType; - this.statePath = statePath; - this.operatorIdentifier = operatorIdentifier; - this.keyValueProjections = keyValueProjections; - this.keyTypeInfo = keyTypeInfo; - this.rowType = rowType; - } - - @Override - public DynamicTableSource copy() { - SavepointDynamicTableSource copy = - new SavepointDynamicTableSource( - stateBackendType, - statePath, - operatorIdentifier, - keyTypeInfo, - keyValueProjections, - rowType); - copy.keyFilter = this.keyFilter; - return copy; - } - - @Override - public Result applyFilters(List 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()); + final int keyColumnIndex, + final Supplier mappingSupplier, + final RowType rowType, + final String summaryString, + final ScanProviderFactory scanProviderFactory) { + super( + stateBackendType, + statePath, + operatorIdentifier, + keyColumnIndex, + mappingSupplier, + rowType, + summaryString, + scanProviderFactory); } @Override - public String asSummaryString() { - return "Savepoint Table Source"; + protected AbstractSavepointDynamicTableSource newInstance() { + return new SavepointDynamicTableSource<>( + stateBackendType, + statePath, + operatorIdentifier, + keyColumnIndex, + mappingSupplier, + rowType, + summaryString, + scanProviderFactory); } @Override - public ChangelogMode getChangelogMode() { - return ChangelogMode.insertOnly(); + public boolean supportsNestedProjection() { + return false; } @Override - public ScanRuntimeProvider getScanRuntimeProvider(ScanContext scanContext) { - return new SavepointDataStreamScanProvider( - stateBackendType, - statePath, - operatorIdentifier, - keyTypeInfo, - keyValueProjections, - rowType, - keyFilter); + @SuppressWarnings("unchecked") + public void applyProjection(int[][] projectedFields, DataType producedDataType) { + this.rowType = (RowType) producedDataType.getLogicalType(); + this.keyColumnIndex = + TableMappingSupport.remapColumnIndex(projectedFields, this.keyColumnIndex); + // Compose the projection lazily so the mapping is still resolved at scan time. + final Supplier prev = this.mappingSupplier; + this.mappingSupplier = () -> (M) prev.get().project(projectedFields); } } diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointDynamicTableSourceFactory.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointDynamicTableSourceFactory.java index 85220c4e5362aa..b3a9c17d1a7f6a 100644 --- a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointDynamicTableSourceFactory.java +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointDynamicTableSourceFactory.java @@ -20,65 +20,31 @@ import org.apache.flink.api.common.serialization.SerializerConfig; import org.apache.flink.api.common.serialization.SerializerConfigImpl; -import org.apache.flink.api.common.typeinfo.TypeInformation; -import org.apache.flink.api.common.typeutils.TypeSerializer; -import org.apache.flink.api.java.tuple.Tuple2; import org.apache.flink.configuration.ConfigOption; -import org.apache.flink.configuration.ConfigOptions; import org.apache.flink.configuration.Configuration; -import org.apache.flink.runtime.state.metainfo.StateMetaInfoSnapshot; import org.apache.flink.state.api.OperatorIdentifier; -import org.apache.flink.state.api.runtime.SavepointLoader; -import org.apache.flink.table.api.ValidationException; -import org.apache.flink.table.catalog.ResolvedCatalogTable; -import org.apache.flink.table.catalog.ResolvedSchema; import org.apache.flink.table.connector.source.DynamicTableSource; +import org.apache.flink.table.factories.DynamicTableFactory.Context; import org.apache.flink.table.factories.DynamicTableSourceFactory; import org.apache.flink.table.factories.FactoryUtil; -import org.apache.flink.table.types.DataType; -import org.apache.flink.table.types.logical.LogicalType; -import org.apache.flink.table.types.logical.LogicalTypeRoot; import org.apache.flink.table.types.logical.RowType; -import org.apache.flink.table.types.logical.utils.LogicalTypeChecks; -import org.apache.flink.util.Preconditions; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.util.Arrays; import java.util.HashSet; -import java.util.List; -import java.util.Map; import java.util.Set; -import java.util.stream.Collectors; -import java.util.stream.IntStream; - -import static org.apache.flink.state.table.SavepointConnectorOptions.FIELDS; -import static org.apache.flink.state.table.SavepointConnectorOptions.KEY_CLASS; -import static org.apache.flink.state.table.SavepointConnectorOptions.KEY_CLASS_PLACEHOLDER; -import static org.apache.flink.state.table.SavepointConnectorOptions.KEY_TYPE_FACTORY; -import static org.apache.flink.state.table.SavepointConnectorOptions.KEY_TYPE_INFO_FACTORY_PLACEHOLDER; +import java.util.function.Supplier; + import static org.apache.flink.state.table.SavepointConnectorOptions.OPERATOR_UID; import static org.apache.flink.state.table.SavepointConnectorOptions.OPERATOR_UID_HASH; import static org.apache.flink.state.table.SavepointConnectorOptions.STATE_BACKEND_TYPE; -import static org.apache.flink.state.table.SavepointConnectorOptions.STATE_NAME; import static org.apache.flink.state.table.SavepointConnectorOptions.STATE_NAME_PLACEHOLDER; import static org.apache.flink.state.table.SavepointConnectorOptions.STATE_PATH; -import static org.apache.flink.state.table.SavepointConnectorOptions.STATE_TYPE; -import static org.apache.flink.state.table.SavepointConnectorOptions.STATE_TYPE_PLACEHOLDER; -import static org.apache.flink.state.table.SavepointConnectorOptions.VALUE_CLASS; -import static org.apache.flink.state.table.SavepointConnectorOptions.VALUE_CLASS_PLACEHOLDER; -import static org.apache.flink.state.table.SavepointConnectorOptions.VALUE_TYPE_FACTORY; -import static org.apache.flink.state.table.SavepointConnectorOptions.VALUE_TYPE_INFO_FACTORY_PLACEHOLDER; +import static org.apache.flink.state.table.SavepointConnectorOptions.STATE_READER_MODE; import static org.apache.flink.state.table.SavepointConnectorOptionsUtil.getOperatorIdentifier; import static org.apache.flink.table.factories.FactoryUtil.CONNECTOR; /** Dynamic source factory for {@link SavepointDynamicTableSource}. */ public class SavepointDynamicTableSourceFactory implements DynamicTableSourceFactory { - private static final Logger LOG = - LoggerFactory.getLogger(SavepointDynamicTableSourceFactory.class); - @Override public DynamicTableSource createDynamicTableSource(Context context) { Configuration options = new Configuration(); @@ -89,207 +55,152 @@ public DynamicTableSource createDynamicTableSource(Context context) { final String statePath = options.get(STATE_PATH); final OperatorIdentifier operatorIdentifier = getOperatorIdentifier(options); - final Map preloadedStateMetadata = - preloadStateMetadata(statePath, operatorIdentifier); - - // Create resolver with preloaded metadata - SavepointTypeInfoResolver typeResolver = - new SavepointTypeInfoResolver(preloadedStateMetadata, serializerConfig); - - final Tuple2 keyValueProjections = - createKeyValueProjections(context.getCatalogTable()); + SavepointConnectorOptions.StateReaderMode readerMode = options.get(STATE_READER_MODE); + switch (readerMode) { + case KEYED: + return createKeyedDynamicTableSource( + context, + options, + serializerConfig, + stateBackendType, + statePath, + operatorIdentifier); + case KEYED_FLAT: + return createFlattenedDynamicTableSource( + context, + options, + serializerConfig, + stateBackendType, + statePath, + operatorIdentifier); + default: + throw new IllegalArgumentException("Unsupported state reader mode: " + readerMode); + } + } - LogicalType logicalType = context.getPhysicalRowDataType().getLogicalType(); - Preconditions.checkArgument(logicalType.is(LogicalTypeRoot.ROW), "Row data type expected."); - RowType rowType = (RowType) logicalType; + /** + * Creates a {@link SavepointDynamicTableSource} for the general keyed-state table (selected via + * {@link SavepointConnectorOptions#STATE_READER_MODE} being set to {@link + * SavepointConnectorOptions.StateReaderMode#KEYED}, the default). + */ + private DynamicTableSource createKeyedDynamicTableSource( + Context context, + Configuration options, + SerializerConfig serializerConfig, + String stateBackendType, + String statePath, + OperatorIdentifier operatorIdentifier) { Set> requiredOptions = new HashSet<>(requiredOptions()); Set> optionalOptions = new HashSet<>(optionalOptions()); - RowType.RowField keyRowField = rowType.getFields().get(keyValueProjections.f0); - ConfigOption keyFormatOption = - optionOf(keyRowField.getName(), VALUE_CLASS).stringType().noDefaultValue(); - optionalOptions.add(keyFormatOption); - - ConfigOption keyTypeInfoFactoryOption = - optionOf(keyRowField.getName(), VALUE_TYPE_FACTORY).stringType().noDefaultValue(); - optionalOptions.add(keyTypeInfoFactoryOption); - - TypeInformation keyTypeInfo = - typeResolver.resolveKeyType( - options, keyFormatOption, keyTypeInfoFactoryOption, keyRowField); - - final Tuple2> keyValueConfigProjections = - Tuple2.of( - keyValueProjections.f0, - Arrays.stream(keyValueProjections.f1) - .mapToObj( - columnIndex -> - createStateColumnConfiguration( - columnIndex, - rowType, - options, - optionalOptions, - typeResolver)) - .collect(Collectors.toList())); - FactoryUtil.validateFactoryOptions(requiredOptions, optionalOptions, options); + // Validate schema and register per-field options eagerly (no class loading) so that + // option validation passes at planning time. + int keyColumnIndex = + StateTableMapping.validateAndExtractKeyColumn( + context.getCatalogTable(), optionalOptions); - Set consumedOptionKeys = new HashSet<>(); - consumedOptionKeys.add(CONNECTOR.key()); - requiredOptions.stream().map(ConfigOption::key).forEach(consumedOptionKeys::add); - optionalOptions.stream().map(ConfigOption::key).forEach(consumedOptionKeys::add); - FactoryUtil.validateUnconsumedKeys( - factoryIdentifier(), options.keySet(), consumedOptionKeys); + validateOptions(options, requiredOptions, optionalOptions); + + // Defer I/O and class loading to scan time by creating the StateTableMapping lazily. + Supplier mappingSupplier = + () -> + StateTableMapping.from( + context.getCatalogTable(), + options, + statePath, + operatorIdentifier, + serializerConfig); - return new SavepointDynamicTableSource( + RowType rowType = (RowType) context.getPhysicalRowDataType().getLogicalType(); + + return new SavepointDynamicTableSource<>( stateBackendType, statePath, operatorIdentifier, - keyTypeInfo, - keyValueConfigProjections, - rowType); + keyColumnIndex, + mappingSupplier, + rowType, + "Savepoint Table Source", + SavepointDataStreamScanProvider::new); } - private StateValueColumnConfiguration createStateColumnConfiguration( - int columnIndex, - RowType rowType, + /** + * Creates a {@link FlattenedSavepointDynamicTableSource} for a table exposing a single + * flattened LIST/MAP state (selected via {@link SavepointConnectorOptions#STATE_READER_MODE} + * being set to {@link SavepointConnectorOptions.StateReaderMode#KEYED_FLAT}). The state name is + * resolved from {@link SavepointConnectorOptions#FLATTENED_STATE_NAME}. + */ + private DynamicTableSource createFlattenedDynamicTableSource( + Context context, Configuration options, - Set> optionalOptions, - SavepointTypeInfoResolver typeResolver) { - - RowType.RowField valueRowField = rowType.getFields().get(columnIndex); - - ConfigOption stateNameOption = - optionOf(valueRowField.getName(), STATE_NAME).stringType().noDefaultValue(); - optionalOptions.add(stateNameOption); - - ConfigOption stateTypeOption = - optionOf(valueRowField.getName(), STATE_TYPE) - .enumType(SavepointConnectorOptions.StateType.class) - .noDefaultValue(); - optionalOptions.add(stateTypeOption); - - ConfigOption mapKeyFormatOption = - optionOf(valueRowField.getName(), KEY_CLASS).stringType().noDefaultValue(); - optionalOptions.add(mapKeyFormatOption); - - ConfigOption mapKeyTypeInfoFactoryOption = - optionOf(valueRowField.getName(), KEY_TYPE_FACTORY).stringType().noDefaultValue(); - optionalOptions.add(mapKeyTypeInfoFactoryOption); - - ConfigOption valueFormatOption = - optionOf(valueRowField.getName(), VALUE_CLASS).stringType().noDefaultValue(); - optionalOptions.add(valueFormatOption); - - ConfigOption valueTypeInfoFactoryOption = - optionOf(valueRowField.getName(), VALUE_TYPE_FACTORY).stringType().noDefaultValue(); - optionalOptions.add(valueTypeInfoFactoryOption); - - LogicalType valueLogicalType = valueRowField.getType(); + SerializerConfig serializerConfig, + String stateBackendType, + String statePath, + OperatorIdentifier operatorIdentifier) { SavepointConnectorOptions.StateType stateType = - options.getOptional(stateTypeOption) - .orElseGet(() -> inferStateType(valueLogicalType)); - - TypeSerializer mapKeyTypeSerializer = - typeResolver.resolveSerializer( - options, - mapKeyFormatOption, - mapKeyTypeInfoFactoryOption, - valueRowField, - stateType.equals(SavepointConnectorOptions.StateType.MAP), - SavepointTypeInfoResolver.InferenceContext.MAP_KEY); - - TypeSerializer valueTypeSerializer = - typeResolver.resolveSerializer( - options, - valueFormatOption, - valueTypeInfoFactoryOption, - valueRowField, - true, - SavepointTypeInfoResolver.InferenceContext.VALUE); - - return new StateValueColumnConfiguration( - columnIndex, - options.getOptional(stateNameOption).orElse(valueRowField.getName()), - stateType, - mapKeyTypeSerializer, - valueTypeSerializer); - } + FlattenedStateTableMapping.validateFlattenedSchema(context.getCatalogTable()); - private static ConfigOptions.OptionBuilder optionOf(String rowField, String optionName) { - return ConfigOptions.key(String.format("%s.%s.%s", FIELDS, rowField, optionName)); - } + RowType rowType = (RowType) context.getPhysicalRowDataType().getLogicalType(); - private Tuple2 createKeyValueProjections(ResolvedCatalogTable catalogTable) { - ResolvedSchema schema = catalogTable.getResolvedSchema(); - if (schema.getPrimaryKey().isEmpty()) { - throw new ValidationException("Could not find the primary key in the table schema."); - } + String stateName = validateAndGetFlattenedStateName(options); - List keyFields = schema.getPrimaryKey().get().getColumns(); - if (keyFields.size() != 1) { - throw new ValidationException( - "Only a single primary key must be defined in the table schema."); - } - - DataType physicalDataType = schema.toPhysicalRowDataType(); - int keyProjection = createKeyFormatProjection(physicalDataType, keyFields.get(0)); - int[] valueProjection = createValueFormatProjection(physicalDataType, keyProjection); - - return Tuple2.of(keyProjection, valueProjection); - } + // Defer I/O to scan time by creating the mapping lazily. + Supplier mappingSupplier = + () -> + FlattenedStateTableMapping.from( + context.getCatalogTable(), + stateName, + statePath, + operatorIdentifier, + serializerConfig, + stateType); - private int createKeyFormatProjection(DataType physicalDataType, String keyField) { - final LogicalType physicalType = physicalDataType.getLogicalType(); - Preconditions.checkArgument( - physicalType.is(LogicalTypeRoot.ROW), "Row data type expected."); - final List physicalFields = LogicalTypeChecks.getFieldNames(physicalType); - return physicalFields.indexOf(keyField); - } - - private int[] createValueFormatProjection(DataType physicalDataType, int keyProjection) { - final LogicalType physicalType = physicalDataType.getLogicalType(); - Preconditions.checkArgument( - physicalType.is(LogicalTypeRoot.ROW), "Row data type expected."); - final int physicalFieldCount = LogicalTypeChecks.getFieldCount(physicalType); - final IntStream physicalFields = IntStream.range(0, physicalFieldCount); - - return physicalFields.filter(pos -> keyProjection != pos).toArray(); + return new FlattenedSavepointDynamicTableSource<>( + stateBackendType, + statePath, + operatorIdentifier, + FlattenedStateTableMapping.STATE_KEY_COLUMN_INDEX, + mappingSupplier, + rowType, + "Flattened Savepoint Table Source", + FlattenedSavepointDataStreamScanProvider::new); } - private SavepointConnectorOptions.StateType inferStateType(LogicalType logicalType) { - switch (logicalType.getTypeRoot()) { - case ARRAY: - return SavepointConnectorOptions.StateType.LIST; + /** + * Validates {@code options} against the required/optional option sets extended with {@link + * SavepointConnectorOptions#FLATTENED_STATE_NAME}, and returns the resolved state name — shared + * by every table kind whose columns represent a single named state's flattened value fields (or + * a single scalar value column) rather than encoding the state's name via the column layout + * itself. + */ + private String validateAndGetFlattenedStateName(Configuration options) { + Set> requiredOptions = new HashSet<>(requiredOptions()); + requiredOptions.add(SavepointConnectorOptions.FLATTENED_STATE_NAME); + Set> optionalOptions = new HashSet<>(optionalOptions()); - case MAP: - return SavepointConnectorOptions.StateType.MAP; + validateOptions(options, requiredOptions, optionalOptions); - default: - return SavepointConnectorOptions.StateType.VALUE; - } + return options.get(SavepointConnectorOptions.FLATTENED_STATE_NAME); } /** - * Preloads all state metadata for an operator in a single I/O operation. - * - * @param savepointPath Path to the savepoint - * @param operatorIdentifier Operator UID or hash - * @return Map from state name to StateMetaInfoSnapshot + * Validates {@code options} against the given required/optional option sets and ensures no + * unrecognized keys remain (shared by both the general and flattened table source paths). */ - private Map preloadStateMetadata( - String savepointPath, OperatorIdentifier operatorIdentifier) { - try { - return SavepointLoader.loadOperatorStateMetadata(savepointPath, operatorIdentifier); - } catch (Exception e) { - throw new RuntimeException( - String.format( - "Failed to load state metadata from savepoint '%s' for operator '%s'. " - + "Ensure the savepoint path is valid and the operator exists in the savepoint. ", - savepointPath, operatorIdentifier), - e); - } + private void validateOptions( + Configuration options, + Set> requiredOptions, + Set> optionalOptions) { + FactoryUtil.validateFactoryOptions(requiredOptions, optionalOptions, options); + + Set consumedOptionKeys = new HashSet<>(); + consumedOptionKeys.add(CONNECTOR.key()); + requiredOptions.stream().map(ConfigOption::key).forEach(consumedOptionKeys::add); + optionalOptions.stream().map(ConfigOption::key).forEach(consumedOptionKeys::add); + FactoryUtil.validateUnconsumedKeys( + factoryIdentifier(), options.keySet(), consumedOptionKeys); } @Override @@ -316,11 +227,15 @@ public Set> optionalOptions() { // Multiple values can be read so registering placeholders options.add(STATE_NAME_PLACEHOLDER); - options.add(STATE_TYPE_PLACEHOLDER); - options.add(KEY_CLASS_PLACEHOLDER); - options.add(KEY_TYPE_INFO_FACTORY_PLACEHOLDER); - options.add(VALUE_CLASS_PLACEHOLDER); - options.add(VALUE_TYPE_INFO_FACTORY_PLACEHOLDER); + + // Selects between the general and flattened keyed-state table schemas; set automatically + // by StateCatalog. + options.add(STATE_READER_MODE); + + // Required only for STATE_READER_MODE == KEYED_FLAT/WINDOWED_FLAT (enforced in + // validateAndGetFlattenedStateName); listed here as optional so that generic option + // introspection (docs, Table API tooling) can discover it regardless of mode. + options.add(SavepointConnectorOptions.FLATTENED_STATE_NAME); return options; } diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointFallbackSchemaLoader.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointFallbackSchemaLoader.java new file mode 100644 index 00000000000000..62c9a566dd8827 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointFallbackSchemaLoader.java @@ -0,0 +1,115 @@ +/* + * 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.annotation.Internal; +import org.apache.flink.api.common.typeutils.CustomRestoreSerializerFactory; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; +import org.apache.flink.runtime.checkpoint.OperatorState; +import org.apache.flink.runtime.checkpoint.metadata.CheckpointMetadata; +import org.apache.flink.state.api.OperatorIdentifier; +import org.apache.flink.state.api.input.deserializer.MissingClassSerializerFactory; +import org.apache.flink.state.api.runtime.SavepointLoader; +import org.apache.flink.state.api.schema.StateSchemaExtractor; +import org.apache.flink.state.api.schema.StateSchemaInfo; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +/** + * Shared helpers for loading fallback {@link StateSchemaInfo} (original serializer snapshots) from + * savepoint metadata, and for restoring a {@link TypeSerializer} from a snapshot in a way that + * tolerates missing POJO/Avro classes. + */ +@Internal +final class SavepointFallbackSchemaLoader { + + private SavepointFallbackSchemaLoader() {} + + /** + * Restores {@code snapshot} into a {@link TypeSerializer}, tolerating snapshots (e.g. {@code + * PojoSerializerSnapshot}, {@code AvroSerializerSnapshot}) whose class is missing from the + * classpath — including ones nested arbitrarily deep inside a composite snapshot (e.g. a {@code + * ListSerializerSnapshot} wrapping a POJO element type), since {@code + * CompositeTypeSerializerSnapshot#restoreSerializer()} eagerly restores all of its nested + * serializers. + * + *

The {@link CustomRestoreSerializerFactory} registered here is consulted by the snapshot's + * own {@code restoreSerializer()} whenever it encounters a missing class, at any nesting depth, + * and builds a schema-only deserializer instead of throwing (see {@link + * MissingClassSerializerFactory}). Snapshots whose class is present restore normally, + * unaffected by the factory. The factory is cleared once this single, synchronous restore + * completes so it cannot leak into unrelated work later scheduled on the same thread (e.g. a + * reused planner thread). + */ + static TypeSerializer buildFallbackSerializer(TypeSerializerSnapshot snapshot) { + CustomRestoreSerializerFactory.set(MissingClassSerializerFactory::create); + try { + return snapshot.restoreSerializer(); + } finally { + CustomRestoreSerializerFactory.remove(); + } + } + + static StateSchemaInfo getSchema(String name, Map fallbackSchemas) { + StateSchemaInfo schema = fallbackSchemas.get(name); + if (schema == null) { + throw new IllegalStateException( + "No schema found for state '" + name + "' in savepoint."); + } + return schema; + } + + /** + * Reads the savepoint header and returns a map of state name → {@link StateSchemaInfo}, but + * only when {@code anyNullTypeInfo} is {@code true}; otherwise returns an empty map to avoid + * the overhead of reading the savepoint header. + */ + static Map loadFallbackSchemas( + String statePath, OperatorIdentifier operatorIdentifier, boolean anyNullTypeInfo) { + if (!anyNullTypeInfo) { + return Collections.emptyMap(); + } + + try { + CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(statePath); + OperatorState operatorState = + metadata.getOperatorStates().stream() + .filter( + op -> + op.getOperatorID() + .equals(operatorIdentifier.getOperatorId())) + .findFirst() + .orElse(null); + if (operatorState == null) { + return Collections.emptyMap(); + } + + Map result = new HashMap<>(); + for (StateSchemaInfo info : StateSchemaExtractor.extractSchema(operatorState)) { + result.put(info.stateName, info); + } + return result; + } catch (Exception e) { + throw new RuntimeException("Failed to load fallback schemas from savepoint", e); + } + } +} 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 index f12cdf9a304e96..d225accb804dee 100644 --- 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 @@ -19,6 +19,7 @@ package org.apache.flink.state.table; import org.apache.flink.state.api.filter.SavepointKeyFilter; +import org.apache.flink.table.connector.source.abilities.SupportsFilterPushDown; import org.apache.flink.table.expressions.CallExpression; import org.apache.flink.table.expressions.FieldReferenceExpression; import org.apache.flink.table.expressions.ResolvedExpression; @@ -26,6 +27,8 @@ import org.apache.flink.table.functions.BuiltInFunctionDefinitions; import org.apache.flink.table.functions.FunctionDefinition; import org.apache.flink.table.types.DataType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.table.types.utils.TypeConversions; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -38,6 +41,7 @@ import java.util.Map; import java.util.Set; import java.util.function.BiFunction; +import java.util.function.Consumer; /** * Converts {@link ResolvedExpression} key filter predicates into {@link SavepointKeyFilter} @@ -62,13 +66,13 @@ class SavepointFilterTranslator { BuiltInFunctionDefinitions.BETWEEN, SavepointFilterTranslator::fromBetween, BuiltInFunctionDefinitions.GREATER_THAN, - SavepointFilterTranslator::fromGreaterThan, + (t, call) -> t.fromComparison(call, Comparison.GT), BuiltInFunctionDefinitions.GREATER_THAN_OR_EQUAL, - SavepointFilterTranslator::fromGreaterThanOrEqual, + (t, call) -> t.fromComparison(call, Comparison.GTE), BuiltInFunctionDefinitions.LESS_THAN, - SavepointFilterTranslator::fromLessThan, + (t, call) -> t.fromComparison(call, Comparison.LT), BuiltInFunctionDefinitions.LESS_THAN_OR_EQUAL, - SavepointFilterTranslator::fromLessThanOrEqual); + (t, call) -> t.fromComparison(call, Comparison.LTE)); private final int keyColumnIndex; private final DataType keyColumnType; @@ -97,6 +101,27 @@ Result apply(List filters) { return new Result(accepted, remaining, keyFilter); } + /** + * Shared {@code SupportsFilterPushDown.applyFilters} implementation for {@link + * SavepointDynamicTableSource} and {@link FlattenedSavepointDynamicTableSource}: translates + * {@code filters} against the key column at {@code keyColumnIndex}, reports the extracted key + * filter to {@code keyFilterSetter}, and returns the accepted/remaining split. + */ + static SupportsFilterPushDown.Result applyKeyColumnFilters( + int keyColumnIndex, + RowType rowType, + List filters, + Consumer keyFilterSetter) { + Result result = + new SavepointFilterTranslator( + keyColumnIndex, + TypeConversions.fromLogicalToDataType( + rowType.getTypeAt(keyColumnIndex))) + .apply(filters); + keyFilterSetter.accept(result.keyFilter()); + return SupportsFilterPushDown.Result.of(result.accepted(), result.remaining()); + } + @Nullable private SavepointKeyFilter extractFilter(ResolvedExpression expr) { final BiFunction extractor = @@ -205,26 +230,6 @@ private SavepointKeyFilter fromBetween(CallExpression call) { (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)) { diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointTypeInformationFactory.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointStateMapping.java similarity index 64% rename from flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointTypeInformationFactory.java rename to flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointStateMapping.java index bed6b09a9fc418..b95c5deee98ff8 100644 --- a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointTypeInformationFactory.java +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointStateMapping.java @@ -18,12 +18,20 @@ package org.apache.flink.state.table; -import org.apache.flink.annotation.Experimental; +import org.apache.flink.annotation.Internal; import org.apache.flink.api.common.typeinfo.TypeInformation; -/** {@link TypeInformation} factory for decoding savepoint value data. */ -@Experimental -public interface SavepointTypeInformationFactory { - /** Returns {@link TypeInformation} for data deserialization. */ - TypeInformation getTypeInformation(); +import javax.annotation.Nullable; + +/** + * Common surface shared by every keyed state table mapping, allowing {@link + * AbstractSavepointDataStreamScanProvider} and {@link AbstractSavepointDynamicTableSource} to + * operate on any of them generically. + */ +@Internal +interface SavepointStateMapping { + + /** The resolved key {@link TypeInformation} for the keyed state(s) backing this mapping. */ + @Nullable + TypeInformation getKeyTypeInfo(); } diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointTypeInfoResolver.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointTypeInfoResolver.java index c2571ce1988d31..55a965e65d22b2 100644 --- a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointTypeInfoResolver.java +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointTypeInfoResolver.java @@ -20,142 +20,189 @@ import org.apache.flink.annotation.Internal; import org.apache.flink.api.common.serialization.SerializerConfig; +import org.apache.flink.api.common.state.StateDescriptor; import org.apache.flink.api.common.typeinfo.TypeInformation; -import org.apache.flink.api.common.typeinfo.utils.TypeUtils; import org.apache.flink.api.common.typeutils.TypeSerializer; import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; +import org.apache.flink.api.common.typeutils.base.ListSerializer; import org.apache.flink.api.common.typeutils.base.MapSerializer; -import org.apache.flink.configuration.ConfigOption; -import org.apache.flink.configuration.Configuration; import org.apache.flink.runtime.state.metainfo.StateMetaInfoSnapshot; +import org.apache.flink.table.runtime.typeutils.ExternalTypeInfo; import org.apache.flink.table.types.logical.ArrayType; import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.LogicalTypeRoot; import org.apache.flink.table.types.logical.MapType; import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.table.types.utils.TypeConversions; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import javax.annotation.Nullable; + import java.math.BigDecimal; import java.util.Map; import java.util.Optional; -import static org.apache.flink.state.table.SavepointConnectorOptions.FIELDS; -import static org.apache.flink.state.table.SavepointConnectorOptions.VALUE_CLASS; - /** Resolver for TypeInformation from savepoint metadata and configuration. */ @Internal class SavepointTypeInfoResolver { private static final Logger LOG = LoggerFactory.getLogger(SavepointTypeInfoResolver.class); - /** Context for type inference to determine what aspect of the type we need. */ + /** Determines which serializer of a state's metadata entry is being resolved. */ enum InferenceContext { - /** Inferring the key type of keyed state (always primitive). */ + /** The key type of keyed state, or a {@code BroadcastState}'s (unwrapped) map key. */ KEY, - /** Inferring the key type of a MAP state. */ + /** The key type of a keyed MAP state, wrapped inside its {@code MapSerializer}. */ MAP_KEY, - /** Inferring the value type (behavior depends on logical type). */ - VALUE + /** The value type, unwrapping a keyed LIST/MAP state's List/MapSerializer. */ + VALUE, + /** + * The value type of a non-keyed (operator) state — {@code ListState}/{@code + * UnionState}/{@code BroadcastState} — whose value serializer is stored flat/unwrapped, so + * unlike {@link #VALUE} no ARRAY/MAP unwrapping is applied. + */ + FLAT_VALUE } private final Map preloadedStateMetadata; private final SerializerConfig serializerConfig; + @Nullable private final TypeSerializerSnapshot preloadedKeySnapshot; public SavepointTypeInfoResolver( Map preloadedStateMetadata, - SerializerConfig serializerConfig) { + SerializerConfig serializerConfig, + @Nullable TypeSerializerSnapshot preloadedKeySnapshot) { this.preloadedStateMetadata = preloadedStateMetadata; this.serializerConfig = serializerConfig; + this.preloadedKeySnapshot = preloadedKeySnapshot; } /** - * Resolves TypeInformation for keyed state keys (primitive types only). + * Resolves the {@link TypeInformation} for a keyed state key: primitive types directly from the + * {@link LogicalType}, complex ones (POJO, Avro) from the backend key serializer snapshot + * preloaded from the savepoint metadata (every state in the same keyed-state backend shares + * it). * - *

This is a simplified version of type resolution specifically for key types, which are - * always primitive and don't require complex metadata inference. - * - * @param options Configuration containing table options - * @param classOption Config option for explicit class specification - * @param typeInfoFactoryOption Config option for type factory specification - * @param rowField The row field containing name and LogicalType - * @return The resolved TypeInformation for the key - * @throws IllegalArgumentException If both class and factory options are specified - * @throws RuntimeException If type instantiation fails + * @throws IllegalArgumentException if the type cannot be inferred */ - public TypeInformation resolveKeyType( - Configuration options, - ConfigOption classOption, - ConfigOption typeInfoFactoryOption, - RowType.RowField rowField) { - try { - // Priority 1: Explicit configuration (backward compatibility) - TypeInformation explicitTypeInfo = - getExplicitTypeInfo(options, classOption, typeInfoFactoryOption); - if (explicitTypeInfo != null) { - return explicitTypeInfo; - } + public TypeInformation resolveKeyType(RowType.RowField rowField) { + LogicalType logicalType = rowField.getType(); + + Class primitiveClass = primitiveClass(logicalType); + if (primitiveClass != null) { + return TypeInformation.of(primitiveClass); + } - // Priority 2: Simple primitive type inference from LogicalType - LogicalType logicalType = rowField.getType(); - String columnName = rowField.getName(); - return TypeInformation.of(getPrimitiveClass(logicalType, columnName)); - } catch (ReflectiveOperationException e) { - throw new RuntimeException(e); + // A ROW-typed key (POJO or Avro) has no primitive class; restore its serializer from the + // savepoint metadata instead. + if (logicalType.is(LogicalTypeRoot.ROW) && preloadedKeySnapshot != null) { + return ExternalTypeInfo.of( + TypeConversions.fromLogicalToDataType(logicalType), + SavepointFallbackSchemaLoader.buildFallbackSerializer(preloadedKeySnapshot)); } + + throw new IllegalArgumentException( + "Cannot resolve key TypeInformation for column '" + + rowField.getName() + + "' with type " + + logicalType.getTypeRoot() + + "."); } /** - * Resolves TypeSerializer for a table field using a three-tier priority system with direct - * serializer extraction for metadata inference. - * - *

Three-Tier Priority System (Serializer-First)

- * - *
    - *
  1. Priority 1: Explicit Configuration (Highest priority)
    - * Uses user-specified class name or type factory from table options, then converts to - * serializer. - *
  2. Priority 2: Metadata Inference
    - * Directly extracts serializers from preloaded savepoint metadata (NO TypeInformation - * conversion). - *
  3. Priority 3: LogicalType Fallback (Lowest priority)
    - * Infers TypeInformation from table schema's LogicalType, then converts to serializer. - *
+ * Resolves the {@link TypeSerializer} for a MAP state's key, or {@code null} if {@code isMap} + * is {@code false} (the field's state type has no map key to resolve). + */ + @Nullable + public TypeSerializer resolveMapKeySerializer(RowType.RowField rowField, boolean isMap) { + return isMap ? resolveSerializer(rowField, InferenceContext.MAP_KEY) : null; + } + + /** + * Resolves the {@link TypeSerializer} for a state's key stored directly under {@code + * CommonSerializerKeys#KEY_SERIALIZER} (e.g. an operator {@code BroadcastState}'s map key, + * which — unlike a keyed {@code MapState}'s key — is not wrapped inside a {@code + * MapSerializer}). + */ + public TypeSerializer resolveKeySerializer(RowType.RowField rowField) { + return resolveSerializer(rowField, InferenceContext.KEY); + } + + /** Resolves the {@link TypeSerializer} for a state's value. */ + public TypeSerializer resolveValueSerializer(RowType.RowField rowField) { + return resolveSerializer(rowField, InferenceContext.VALUE); + } + + /** + * Resolves the {@link TypeSerializer} for a non-keyed (operator) state's value with no + * ARRAY/MAP unwrapping, unlike {@link #resolveValueSerializer} which unwraps a keyed LIST/MAP + * state's wrapping ListSerializer/MapSerializer. + */ + public TypeSerializer resolveFlatValueSerializer(RowType.RowField rowField) { + return resolveSerializer(rowField, InferenceContext.FLAT_VALUE); + } + + /** + * Resolves the precise {@link StateDescriptor.Type} (e.g. {@code REDUCING}/{@code AGGREGATING}) + * a state was originally registered under, read from the {@code KEYED_STATE_TYPE} option in the + * preloaded savepoint metadata. Returns {@code UNKNOWN} when the state (or the option) is not + * present, in which case callers fall back to treating the state as plain VALUE/LIST/MAP. + */ + public StateDescriptor.Type resolveStateKind(String stateName) { + StateMetaInfoSnapshot stateMetaInfo = preloadedStateMetadata.get(stateName); + if (stateMetaInfo == null) { + return StateDescriptor.Type.UNKNOWN; + } + String kind = + stateMetaInfo.getOption(StateMetaInfoSnapshot.CommonOptionsKeys.KEYED_STATE_TYPE); + if (kind == null) { + return StateDescriptor.Type.UNKNOWN; + } + try { + return StateDescriptor.Type.valueOf(kind); + } catch (IllegalArgumentException e) { + return StateDescriptor.Type.UNKNOWN; + } + } + + /** + * Resolves the {@link TypeSerializer} used for the namespace under which {@code stateName} is + * registered (e.g. a window serializer), read directly from the preloaded savepoint metadata. * - *

This approach eliminates TypeInformation extraction complexity for metadata inference, - * making it work with ANY serializer type (Avro, custom types, etc.). + *

Unlike {@link #resolveValueSerializer}, there is no LogicalType-based fallback: the + * namespace serializer is only ever needed for states already classified as namespaced, so the + * metadata is known to contain it. + */ + public TypeSerializer resolveNamespaceSerializer(String stateName) { + StateMetaInfoSnapshot stateMetaInfo = preloadedStateMetadata.get(stateName); + if (stateMetaInfo == null) { + throw new IllegalArgumentException( + "State '" + stateName + "' not found in preloaded metadata."); + } + TypeSerializerSnapshot namespaceSnapshot = + stateMetaInfo.getTypeSerializerSnapshot( + StateMetaInfoSnapshot.CommonSerializerKeys.NAMESPACE_SERIALIZER); + if (namespaceSnapshot == null) { + throw new IllegalArgumentException( + "State '" + stateName + "' has no namespace serializer in metadata."); + } + return SavepointFallbackSchemaLoader.buildFallbackSerializer(namespaceSnapshot); + } + + /** + * Resolves the serializer for a table field, preferring the serializer the state was actually + * written with (extracted from the preloaded savepoint metadata) and falling back to inference + * from the table schema's {@link LogicalType}. * - * @param options Configuration containing table options - * @param classOption Config option for explicit class specification - * @param typeInfoFactoryOption Config option for type factory specification - * @param rowField The table field containing name and LogicalType - * @param inferStateType Whether to enable automatic type inference. If false, returns null when - * no explicit configuration is provided. - * @param context The inference context determining what type aspect to extract. - * @return The resolved TypeSerializer, or null if inferStateType is false and no explicit - * configuration is provided. - * @throws IllegalArgumentException If both class and factory options are specified - * @throws RuntimeException If serializer creation fails + * @return the resolved serializer, or {@code null} for complex types (ROW/POJO) that cannot be + * inferred, which callers resolve from the savepoint header instead */ - public TypeSerializer resolveSerializer( - Configuration options, - ConfigOption classOption, - ConfigOption typeInfoFactoryOption, - RowType.RowField rowField, - boolean inferStateType, - InferenceContext context) { + @Nullable + private TypeSerializer resolveSerializer( + RowType.RowField rowField, InferenceContext context) { try { - // Priority 1: Explicit configuration (backward compatibility) - TypeInformation explicitTypeInfo = - getExplicitTypeInfo(options, classOption, typeInfoFactoryOption); - if (explicitTypeInfo != null) { - return explicitTypeInfo.createSerializer(serializerConfig); - } - if (!inferStateType) { - return null; - } - - // Priority 2: Direct serializer extraction from metadata Optional> metadataSerializer = getSerializerFromMetadata(rowField, context); if (metadataSerializer.isPresent()) { @@ -167,9 +214,10 @@ public TypeSerializer resolveSerializer( return metadataSerializer.get(); } - // Priority 3: Fallback to LogicalType-based inference TypeInformation fallbackTypeInfo = inferTypeFromLogicalType(rowField, context); - return fallbackTypeInfo.createSerializer(serializerConfig); + return fallbackTypeInfo == null + ? null + : fallbackTypeInfo.createSerializer(serializerConfig); } catch (Exception e) { throw new RuntimeException( "Failed to resolve serializer for field " + rowField.getName(), e); @@ -177,98 +225,29 @@ public TypeSerializer resolveSerializer( } /** - * Extracts explicit TypeInformation from user configuration (Priority 1). - * - * @param options Configuration containing table options - * @param classOption Config option for explicit class specification - * @param typeInfoFactoryOption Config option for type factory specification - * @return The explicit TypeInformation if specified, null otherwise - * @throws IllegalArgumentException If both class and factory options are specified - * @throws ReflectiveOperationException If type instantiation fails - */ - private TypeInformation getExplicitTypeInfo( - Configuration options, - ConfigOption classOption, - ConfigOption typeInfoFactoryOption) - throws ReflectiveOperationException { - - Optional clazz = options.getOptional(classOption); - Optional typeInfoFactory = options.getOptional(typeInfoFactoryOption); - - if (clazz.isPresent() && typeInfoFactory.isPresent()) { - throw new IllegalArgumentException( - "Either " - + classOption.key() - + " or " - + typeInfoFactoryOption.key() - + " can be specified, not both."); - } - - if (clazz.isPresent()) { - return TypeInformation.of(Class.forName(clazz.get())); - } else if (typeInfoFactory.isPresent()) { - SavepointTypeInformationFactory savepointTypeInformationFactory = - (SavepointTypeInformationFactory) - TypeUtils.getInstance(typeInfoFactory.get(), new Object[0]); - return savepointTypeInformationFactory.getTypeInformation(); - } - - return null; - } - - /** - * Directly extracts TypeSerializer from preloaded metadata (Priority 2). - * - *

This method performs NO I/O and NO TypeInformation conversion. It directly extracts the - * serializer that was used to write the state data. - * - * @param rowField The row field to extract serializer for - * @param context The inference context determining what serializer to extract - * @return The serializer if found in metadata, empty otherwise + * Extracts the serializer the state was written with from the preloaded metadata. Performs no + * I/O and no {@link TypeInformation} conversion, so it works with any serializer type (Avro, + * custom types, etc.). */ private Optional> getSerializerFromMetadata( RowType.RowField rowField, InferenceContext context) { + String stateName = rowField.getName(); try { - // Get state name for this field (defaults to field name) - String stateName = rowField.getName(); - - // Look up from preloaded metadata (NO I/O) StateMetaInfoSnapshot stateMetaInfo = preloadedStateMetadata.get(stateName); - if (stateMetaInfo == null) { LOG.debug("State '{}' not found in preloaded metadata", stateName); return Optional.empty(); } - // Extract appropriate serializer based on context - TypeSerializerSnapshot serializerSnapshot = null; - switch (context) { - case KEY: - serializerSnapshot = - stateMetaInfo.getTypeSerializerSnapshot( - StateMetaInfoSnapshot.CommonSerializerKeys.KEY_SERIALIZER); - break; - case MAP_KEY: - // For MAP_KEY, we need the key serializer from the value serializer - // (which is MapSerializer) - TypeSerializerSnapshot valueSnapshot = - stateMetaInfo.getTypeSerializerSnapshot( - StateMetaInfoSnapshot.CommonSerializerKeys.VALUE_SERIALIZER); - if (valueSnapshot != null) { - TypeSerializer valueSerializer = valueSnapshot.restoreSerializer(); - if (valueSerializer instanceof MapSerializer) { - serializerSnapshot = - ((MapSerializer) valueSerializer) - .getKeySerializer() - .snapshotConfiguration(); - } - } - break; - case VALUE: - serializerSnapshot = - stateMetaInfo.getTypeSerializerSnapshot( - StateMetaInfoSnapshot.CommonSerializerKeys.VALUE_SERIALIZER); - break; + TypeSerializerSnapshot serializerSnapshot; + if (context == InferenceContext.KEY) { + serializerSnapshot = + stateMetaInfo.getTypeSerializerSnapshot( + StateMetaInfoSnapshot.CommonSerializerKeys.KEY_SERIALIZER); + } else { + serializerSnapshot = + stateMetaInfo.getTypeSerializerSnapshot( + StateMetaInfoSnapshot.CommonSerializerKeys.VALUE_SERIALIZER); } if (serializerSnapshot == null) { @@ -279,44 +258,41 @@ private Optional> getSerializerFromMetadata( return Optional.empty(); } - // Restore serializer from snapshot - TypeSerializer serializer = serializerSnapshot.restoreSerializer(); + // Restore via the POJO-friendly path, so a missing POJO class does not fail the + // restore. + TypeSerializer serializer = + SavepointFallbackSchemaLoader.buildFallbackSerializer(serializerSnapshot); - // For VALUE context with complex types, extract the appropriate sub-serializer - if (context == InferenceContext.VALUE) { - return extractValueSerializerForLogicalType(serializer, rowField.getType()); + switch (context) { + case MAP_KEY: + // A keyed MAP state's key serializer lives inside its MapSerializer. + return serializer instanceof MapSerializer + ? Optional.of(((MapSerializer) serializer).getKeySerializer()) + : Optional.empty(); + case VALUE: + return unwrapValueSerializer(serializer, rowField.getType()); + default: + return Optional.of(serializer); } - - return Optional.of(serializer); - } catch (Exception e) { LOG.warn( "Failed to extract serializer from metadata for field '{}': {}", - rowField.getName(), + stateName, e.getMessage()); return Optional.empty(); } } /** - * Extracts the appropriate value serializer based on LogicalType for VALUE context. - * - * @param fullSerializer The complete serializer from metadata - * @param logicalType The LogicalType from the table schema - * @return The appropriate value serializer + * Unwraps the element/value serializer of a keyed LIST/MAP state (whose metadata serializer is + * a List/MapSerializer); other logical types use the serializer as-is. */ - private Optional> extractValueSerializerForLogicalType( + private Optional> unwrapValueSerializer( TypeSerializer fullSerializer, LogicalType logicalType) { - switch (logicalType.getTypeRoot()) { case ARRAY: - // ARRAY logical type → LIST state → extract element serializer - if (fullSerializer - instanceof org.apache.flink.api.common.typeutils.base.ListSerializer) { - org.apache.flink.api.common.typeutils.base.ListSerializer listSerializer = - (org.apache.flink.api.common.typeutils.base.ListSerializer) - fullSerializer; - return Optional.of(listSerializer.getElementSerializer()); + if (fullSerializer instanceof ListSerializer) { + return Optional.of(((ListSerializer) fullSerializer).getElementSerializer()); } LOG.debug( "Expected ListSerializer for ARRAY logical type but got: {}", @@ -324,7 +300,6 @@ private Optional> extractValueSerializerForLogicalType( return Optional.empty(); case MAP: - // MAP logical type → MAP state → extract value serializer if (fullSerializer instanceof MapSerializer) { return Optional.of(((MapSerializer) fullSerializer).getValueSerializer()); } @@ -334,161 +309,88 @@ private Optional> extractValueSerializerForLogicalType( return Optional.empty(); default: - // Primitive logical type → VALUE state → use serializer as-is return Optional.of(fullSerializer); } } /** - * Fallback inference using LogicalType when metadata extraction fails. - * - * @param rowField The row field to infer type for - * @param context The inference context - * @return The inferred TypeInformation + * Infers the {@link TypeInformation} from the table schema's {@link LogicalType} when the state + * is absent from the savepoint metadata. Returns {@code null} for complex value types, which + * the caller resolves from the savepoint header instead. */ + @Nullable private TypeInformation inferTypeFromLogicalType( RowType.RowField rowField, InferenceContext context) { - LogicalType logicalType = rowField.getType(); - String columnName = rowField.getName(); - try { - switch (context) { - case KEY: - // Keys are always primitive - return TypeInformation.of(getPrimitiveClass(logicalType, columnName)); + switch (context) { + case KEY: + case FLAT_VALUE: + // Keys are always primitive here (complex ROW keys are handled by + // resolveKeyType), and non-keyed state values are never ARRAY/MAP-wrapped. + Class flatClass = primitiveClass(logicalType); + if (flatClass == null) { + throw new UnsupportedOperationException( + "Cannot infer a " + context + " type from logical type " + logicalType); + } + return TypeInformation.of(flatClass); - case MAP_KEY: - // Extract key type from MAP logical type - if (logicalType instanceof MapType) { - LogicalType keyType = ((MapType) logicalType).getKeyType(); - return TypeInformation.of(getPrimitiveClass(keyType, columnName)); - } + case MAP_KEY: + if (!(logicalType instanceof MapType)) { throw new UnsupportedOperationException( "MAP_KEY context requires MAP logical type, but got: " + logicalType); + } + return TypeInformation.of(primitiveClass(((MapType) logicalType).getKeyType())); - case VALUE: - return inferValueTypeFromLogicalType(logicalType, columnName); - - default: - throw new UnsupportedOperationException("Unknown context: " + context); - } - } catch (ClassNotFoundException e) { - throw new RuntimeException("Failed to infer type for context " + context, e); - } - } - - /** - * Infers value type from LogicalType for VALUE context fallback. - * - * @param logicalType The LogicalType - * @param columnName The column name for error messages - * @return The inferred TypeInformation - */ - private TypeInformation inferValueTypeFromLogicalType( - LogicalType logicalType, String columnName) throws ClassNotFoundException { - - switch (logicalType.getTypeRoot()) { - case ARRAY: - // ARRAY logical type → LIST state → return element type - ArrayType arrayType = (ArrayType) logicalType; - return TypeInformation.of( - getPrimitiveClass(arrayType.getElementType(), columnName)); - - case MAP: - // MAP logical type → MAP state → return value type - MapType mapType = (MapType) logicalType; - return TypeInformation.of(getPrimitiveClass(mapType.getValueType(), columnName)); + case VALUE: + // primitiveClass() already unwraps ARRAY element and MAP value types. + Class valueClass = primitiveClass(logicalType); + return valueClass == null ? null : TypeInformation.of(valueClass); default: - // Primitive logical type → VALUE state → return primitive type - return TypeInformation.of(getPrimitiveClass(logicalType, columnName)); + throw new UnsupportedOperationException("Unknown context: " + context); } } /** - * Maps LogicalType to primitive Java class. - * - * @param logicalType The LogicalType to map - * @param columnName The column name for error messages - * @return The corresponding Java class + * Maps a {@link LogicalType} to its Java class, unwrapping ARRAY element and MAP value types, + * or {@code null} for complex/unknown types (e.g. ROW/POJO) whose class cannot be inferred from + * the schema alone. */ - private Class getPrimitiveClass(LogicalType logicalType, String columnName) - throws ClassNotFoundException { - String className = inferTypeInfoClassFromLogicalType(columnName, logicalType); - return Class.forName(className); - } - - private String inferTypeInfoClassFromLogicalType(String columnName, LogicalType logicalType) { + @Nullable + private static Class primitiveClass(LogicalType logicalType) { switch (logicalType.getTypeRoot()) { case CHAR: case VARCHAR: - return String.class.getName(); - + return String.class; case BOOLEAN: - return Boolean.class.getName(); - + return Boolean.class; case BINARY: case VARBINARY: - return byte[].class.getName(); - + return byte[].class; case DECIMAL: - return BigDecimal.class.getName(); - + return BigDecimal.class; case TINYINT: - return Byte.class.getName(); - + return Byte.class; case SMALLINT: - return Short.class.getName(); - + return Short.class; case INTEGER: - return Integer.class.getName(); - - case BIGINT: - return Long.class.getName(); - - case FLOAT: - return Float.class.getName(); - - case DOUBLE: - return Double.class.getName(); - case DATE: - return Integer.class.getName(); - + return Integer.class; + case BIGINT: case INTERVAL_YEAR_MONTH: case INTERVAL_DAY_TIME: - return Long.class.getName(); - + return Long.class; + case FLOAT: + return Float.class; + case DOUBLE: + return Double.class; case ARRAY: - return inferTypeInfoClassFromLogicalType( - columnName, ((ArrayType) logicalType).getElementType()); - + return primitiveClass(((ArrayType) logicalType).getElementType()); case MAP: - return inferTypeInfoClassFromLogicalType( - columnName, ((MapType) logicalType).getValueType()); - - case NULL: - return null; - - case ROW: - case MULTISET: - case TIME_WITHOUT_TIME_ZONE: - case TIMESTAMP_WITHOUT_TIME_ZONE: - case TIMESTAMP_WITH_TIME_ZONE: - case TIMESTAMP_WITH_LOCAL_TIME_ZONE: - case DISTINCT_TYPE: - case STRUCTURED_TYPE: - case RAW: - case SYMBOL: - case UNRESOLVED: - case DESCRIPTOR: + return primitiveClass(((MapType) logicalType).getValueType()); default: - throw new UnsupportedOperationException( - String.format( - "Unable to infer state format for SQL type: %s in column: %s. " - + "Please override the type with the following config parameter: %s.%s.%s", - logicalType, columnName, FIELDS, columnName, VALUE_CLASS)); + return null; } } } diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SingleColumnStateMapping.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SingleColumnStateMapping.java new file mode 100644 index 00000000000000..e2cbc94366e022 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SingleColumnStateMapping.java @@ -0,0 +1,47 @@ +/* + * 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.annotation.Internal; +import org.apache.flink.api.common.state.StateDescriptor; +import org.apache.flink.api.common.typeutils.TypeSerializer; + +import javax.annotation.Nullable; + +/** + * Mixed into mapping classes that back a table with exactly one flattened LIST/MAP state, described + * by a single fixed descriptor rather than a list of value columns. Implemented by {@link + * FlattenedStateTableMapping} and {@link WindowFlattenedStateTableMapping}, allowing {@link + * AbstractSingleColumnScanProvider} to build their state descriptor generically. + */ +@Internal +@SuppressWarnings("rawtypes") +interface SingleColumnStateMapping extends SavepointStateMapping { + + String getStateName(); + + SavepointConnectorOptions.StateType getStateType(); + + @Nullable + TypeSerializer getMapKeyTypeSerializer(); + + TypeSerializer getValueTypeSerializer(); + + void setStateDescriptor(StateDescriptor stateDescriptor); +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/StateTableMapping.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/StateTableMapping.java new file mode 100644 index 00000000000000..02db3097c949a0 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/StateTableMapping.java @@ -0,0 +1,197 @@ +/* + * 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.annotation.Internal; +import org.apache.flink.api.common.serialization.SerializerConfig; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.configuration.ConfigOption; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.state.api.OperatorIdentifier; +import org.apache.flink.table.api.ValidationException; +import org.apache.flink.table.catalog.ResolvedCatalogTable; +import org.apache.flink.table.catalog.ResolvedSchema; +import org.apache.flink.table.types.DataType; +import org.apache.flink.table.types.logical.RowType; + +import javax.annotation.Nullable; + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.List; +import java.util.Set; + +import static org.apache.flink.state.table.SavepointConnectorOptions.STATE_NAME; + +/** + * Maps the key column and state value columns to their positions in the output row. + * + *

After projection push-down, column indices reflect positions in the projected output row + * rather than the original table schema. The {@link #allValueColumns} list always contains the full + * original set of value columns (needed for key enumeration), while {@link #valueColumns} contains + * only the projected subset that is written to the output row. + */ +@Internal +public class StateTableMapping implements Serializable, MultiColumnStateMapping { + + private static final long serialVersionUID = 1L; + + private final int keyColumnIndex; + private final List valueColumns; + + /** + * Full original value columns, preserved across projections for state-descriptor registration. + * Key enumeration in {@code KeyedStateReaderOperator} requires at least one registered state; + * this list guarantees that even when all value columns are projected out. + */ + private final List allValueColumns; + + /** Resolved key type info; {@code null} when the mapping was not created via {@link #from}. */ + @Nullable private final TypeInformation keyTypeInfo; + + public StateTableMapping(int keyColumnIndex, List valueColumns) { + this(keyColumnIndex, null, valueColumns, valueColumns); + } + + private StateTableMapping( + int keyColumnIndex, + @Nullable TypeInformation keyTypeInfo, + List valueColumns, + List allValueColumns) { + this.keyColumnIndex = keyColumnIndex; + this.keyTypeInfo = keyTypeInfo; + this.valueColumns = valueColumns; + this.allValueColumns = allValueColumns; + } + + public int getKeyColumnIndex() { + return keyColumnIndex; + } + + /** Columns to write to the output row (may be a projected subset). */ + public List getValueColumns() { + return valueColumns; + } + + @Override + public List getAllValueColumns() { + return allValueColumns; + } + + @Override + @Nullable + public TypeInformation getKeyTypeInfo() { + return keyTypeInfo; + } + + /** + * Creates a new {@link StateTableMapping} with column indices remapped to the projected output. + * Only flat projections ({@code projectedFields[i].length == 1}) are supported; a key that was + * projected out (e.g. after constant folding with filter push-down) maps to {@code -1}, meaning + * it is not written to the output row. + */ + @Override + public StateTableMapping project(int[][] projectedFields) { + return new StateTableMapping( + TableMappingSupport.remapColumnIndex(projectedFields, this.keyColumnIndex), + keyTypeInfo, + TableMappingSupport.remapValueColumns(projectedFields, this.valueColumns), + allValueColumns); + } + + // ------------------------------------------------------------------------- + // Factory + // ------------------------------------------------------------------------- + + /** + * Validates the table schema, registers all per-field connector {@link ConfigOption}s into + * {@code optionalOptions} for option validation, and returns the index of the primary-key + * column in the physical row data type. + * + *

This is a purely structural operation: it analyses the table schema but performs no class + * loading or type resolution. Call this eagerly so that option validation passes, then wrap + * {@link #from} in a {@code Supplier} to defer class loading to scan time. + */ + public static int validateAndExtractKeyColumn( + ResolvedCatalogTable catalogTable, Set> optionalOptions) { + + ResolvedSchema schema = catalogTable.getResolvedSchema(); + if (schema.getPrimaryKey().isEmpty()) { + throw new ValidationException("Could not find the primary key in the table schema."); + } + + List keyFields = schema.getPrimaryKey().get().getColumns(); + if (keyFields.size() != 1) { + throw new ValidationException( + "Only a single primary key must be defined in the table schema."); + } + + DataType physicalDataType = schema.toPhysicalRowDataType(); + int keyIdx = TableMappingSupport.columnIndex(physicalDataType, keyFields.get(0)); + RowType rowType = (RowType) physicalDataType.getLogicalType(); + + for (int colIdx : TableMappingSupport.valueColumnIndices(physicalDataType, keyIdx)) { + RowType.RowField valueRowField = rowType.getFields().get(colIdx); + optionalOptions.add( + TableMappingSupport.fieldOption(valueRowField.getName(), STATE_NAME) + .stringType() + .noDefaultValue()); + } + + return keyIdx; + } + + /** + * Builds a complete {@link StateTableMapping} from a {@link ResolvedCatalogTable}, loading + * operator state metadata from the savepoint and resolving serializers and key type from it. + * + *

Assumes {@link #validateAndExtractKeyColumn} has already been called for this table. This + * performs I/O (savepoint metadata loading); callers should invoke it lazily, deferred to scan + * time, to keep planning free of savepoint access. + */ + public static StateTableMapping from( + ResolvedCatalogTable catalogTable, + Configuration options, + String statePath, + OperatorIdentifier operatorIdentifier, + SerializerConfig serializerConfig) { + + SavepointTypeInfoResolver typeResolver = + TableMappingSupport.createTypeResolver( + statePath, operatorIdentifier, serializerConfig); + + DataType physicalDataType = catalogTable.getResolvedSchema().toPhysicalRowDataType(); + RowType rowType = (RowType) physicalDataType.getLogicalType(); + List keyFields = + catalogTable.getResolvedSchema().getPrimaryKey().get().getColumns(); + int keyIdx = TableMappingSupport.columnIndex(physicalDataType, keyFields.get(0)); + + TypeInformation keyTypeInfo = + typeResolver.resolveKeyType(rowType.getFields().get(keyIdx)); + + List valueColumns = new ArrayList<>(); + for (int colIdx : TableMappingSupport.valueColumnIndices(physicalDataType, keyIdx)) { + valueColumns.add( + TableMappingSupport.createValueColumnConfig( + colIdx, rowType, options, typeResolver)); + } + + return new StateTableMapping(keyIdx, keyTypeInfo, valueColumns, valueColumns); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/StateValueColumnConfiguration.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/StateValueColumnConfiguration.java index 865077717fc7cb..cd36806d015838 100644 --- a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/StateValueColumnConfiguration.java +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/StateValueColumnConfiguration.java @@ -31,6 +31,7 @@ public class StateValueColumnConfiguration implements Serializable { private final int columnIndex; private final String stateName; private final SavepointConnectorOptions.StateType stateType; + private final StateDescriptor.Type actualStateKind; @Nullable private final TypeSerializer mapKeyTypeSerializer; @Nullable private final TypeSerializer valueTypeSerializer; @Nullable private StateDescriptor stateDescriptor; @@ -39,11 +40,13 @@ public StateValueColumnConfiguration( int columnIndex, final String stateName, final SavepointConnectorOptions.StateType stateType, + final StateDescriptor.Type actualStateKind, @Nullable final TypeSerializer mapKeyTypeSerializer, final TypeSerializer valueTypeSerializer) { this.columnIndex = columnIndex; this.stateName = stateName; this.stateType = stateType; + this.actualStateKind = actualStateKind; this.mapKeyTypeSerializer = mapKeyTypeSerializer; this.valueTypeSerializer = valueTypeSerializer; } @@ -60,6 +63,16 @@ public SavepointConnectorOptions.StateType getStateType() { return stateType; } + /** + * The precise {@link StateDescriptor.Type} the state was originally registered under (e.g. + * {@code REDUCING}/{@code AGGREGATING} for a {@code .reduce()}/{@code .aggregate()} window + * function's window-contents state), as opposed to {@link #getStateType()} which only + * distinguishes the coarse VALUE/LIST/MAP shape used for the SQL schema. + */ + public StateDescriptor.Type getActualStateKind() { + return actualStateKind; + } + @Nullable public TypeSerializer getMapKeyTypeSerializer() { return mapKeyTypeSerializer; @@ -77,4 +90,17 @@ public void setStateDescriptor(StateDescriptor stateDescriptor) { public StateDescriptor getStateDescriptor() { return stateDescriptor; } + + public StateValueColumnConfiguration withColumnIndex(int newColumnIndex) { + StateValueColumnConfiguration copy = + new StateValueColumnConfiguration( + newColumnIndex, + stateName, + stateType, + actualStateKind, + mapKeyTypeSerializer, + valueTypeSerializer); + copy.stateDescriptor = this.stateDescriptor; + return copy; + } } diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/StateValueConverter.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/StateValueConverter.java new file mode 100644 index 00000000000000..40b9c0d8b7e8e3 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/StateValueConverter.java @@ -0,0 +1,281 @@ +/* + * 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.annotation.Internal; +import org.apache.flink.api.common.state.AggregatingState; +import org.apache.flink.api.common.state.ListState; +import org.apache.flink.api.common.state.ReducingState; +import org.apache.flink.api.common.state.State; +import org.apache.flink.api.common.state.StateDescriptor; +import org.apache.flink.api.common.state.ValueState; +import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.state.api.input.deserializer.InternalTypeConverter; +import org.apache.flink.state.api.schema.AvroStateUtils; +import org.apache.flink.streaming.api.windowing.windows.TimeWindow; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.TimestampData; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.types.RowKind; +import org.apache.flink.util.Collector; + +import java.lang.reflect.Field; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Supplier; + +/** + * Converts Java objects returned by Flink state (POJOs, Avro records, primitives) into their + * internal {@link RowData} representation, driven by a target {@link LogicalType}. Shared by all + * state reader functions and non-keyed row mappers. + * + *

flink-avro is an optional dependency of this module (see the module {@code pom.xml}), so this + * class must never mention an Avro type directly: doing so - even in an {@code instanceof} branch - + * would make the JVM try to resolve that type the moment this method is reached for *any* + * row-shaped value, throwing {@code NoClassDefFoundError} for callers who never use Avro at all. + * Avro values are instead recognized and read via {@link AvroStateUtils}, which is only ever loaded + * once it has already confirmed, by interface name, that the value is genuinely Avro-backed. + */ +@Internal +@SuppressWarnings({"rawtypes", "unchecked"}) +class StateValueConverter implements java.io.Serializable { + + private static final long serialVersionUID = 1L; + + // Field/Method are not serializable; the reader function containing this converter is shipped + // to task managers via Java serialization, so these are rebuilt lazily on first use. The + // Optional value distinguishes a cached "not found" from "not yet looked up", so a missing + // field/getter is only ever resolved via reflection once per class instead of on every row. + private transient Map, Optional> classFieldCache; + private transient Map, Optional> classMethodCache; + + private Map, Optional> classFieldCache() { + if (classFieldCache == null) { + classFieldCache = new ConcurrentHashMap<>(); + } + return classFieldCache; + } + + private Map, Optional> classMethodCache() { + if (classMethodCache == null) { + classMethodCache = new ConcurrentHashMap<>(); + } + return classMethodCache; + } + + /** + * Reads the value from a VALUE-shaped state, dispatching to {@code ValueState.value()}, {@code + * ReducingState.get()} or {@code AggregatingState.get()} depending on {@code actualStateKind} + * (see {@link AbstractSavepointDataStreamScanProvider#buildStateDescriptor}). + */ + static Object readValueLikeState(State state, StateDescriptor.Type actualStateKind) + throws Exception { + switch (actualStateKind) { + case REDUCING: + return ((ReducingState) state).get(); + case AGGREGATING: + return ((AggregatingState) state).get(); + default: + return ((ValueState) state).value(); + } + } + + /** + * Reads a LIST state's elements for the {@code (state_key, ..., list_value)} (non-flattened) + * table, normalizing the state backend's {@code null} — returned by both {@code HeapListState} + * and {@code RocksDBListState} when the current key/namespace has no entries — to an empty + * {@link List}. This mirrors {@code MapState.entries()}, which already returns an empty (never + * {@code null}) {@link Iterable} in that same case, so LIST- and MAP-shaped state consistently + * always have a value (possibly empty) as documented on {@code StateTableUtils}, rather than + * the LIST column surfacing SQL {@code NULL} for a key that other states in the same row do + * have data for. + */ + static Iterable readListLikeState(ListState state) throws Exception { + Iterable values = state.get(); + return values == null ? Collections.emptyList() : values; + } + + /** + * Iterates a flattened LIST state's elements, emitting one row per element via {@code out}. + * Each row is created via {@code rowTemplate} (which supplies the leading columns already + * populated — e.g. the key, and for namespaced states, the window), and this method fills in + * the list index and value at {@code subKeyColumnIndex}/{@code valueColumnIndex}. + */ + void writeListRows( + Iterable values, + Supplier rowTemplate, + int subKeyColumnIndex, + int valueColumnIndex, + LogicalType indexLogicalType, + LogicalType valueLogicalType, + Collector out) { + if (values == null) { + return; + } + long index = 0; + for (Object value : values) { + GenericRowData row = rowTemplate.get(); + row.setField(subKeyColumnIndex, getValue(indexLogicalType, index)); + row.setField(valueColumnIndex, getValue(valueLogicalType, value)); + out.collect(row); + index++; + } + } + + /** + * Iterates a flattened MAP state's entries, emitting one row per entry via {@code out}. Mirrors + * {@link #writeListRows} for MAP-shaped state. + */ + void writeMapRows( + Iterable> entries, + Supplier rowTemplate, + int mapKeyColumnIndex, + int valueColumnIndex, + LogicalType mapKeyLogicalType, + LogicalType valueLogicalType, + Collector out) { + if (entries == null) { + return; + } + for (Map.Entry entry : entries) { + GenericRowData row = rowTemplate.get(); + row.setField(mapKeyColumnIndex, getValue(mapKeyLogicalType, entry.getKey())); + row.setField(valueColumnIndex, getValue(valueLogicalType, entry.getValue())); + out.collect(row); + } + } + + Object getValue(LogicalType logicalType, Object object) { + if (object == null) { + return null; + } + switch (logicalType.getTypeRoot()) { + case ROW: + if (object instanceof TimeWindow) { + TimeWindow window = (TimeWindow) object; + GenericRowData result = new GenericRowData(RowKind.INSERT, 2); + result.setField(0, TimestampData.fromEpochMillis(window.getStart())); + result.setField(1, TimestampData.fromEpochMillis(window.getEnd())); + return result; + } + if (object instanceof GenericRowData) { + // Copy by position into a fresh INSERT row rather than returning the source + // as-is: a coincidentally equal arity does not imply matching field ordering + // or semantics (e.g. after schema evolution or a column reorder). + GenericRowData sourceRow = (GenericRowData) object; + int targetFieldCount = ((RowType) logicalType).getFieldCount(); + GenericRowData result = new GenericRowData(RowKind.INSERT, targetFieldCount); + for (int i = 0; i < targetFieldCount; i++) { + result.setField(i, i < sourceRow.getArity() ? sourceRow.getField(i) : null); + } + return result; + } + return convertToRow(object, logicalType); + default: + return InternalTypeConverter.toInternal(object, logicalType); + } + } + + private GenericRowData convertToRow(Object object, LogicalType logicalType) { + RowType rowType = (RowType) logicalType; + GenericRowData result = new GenericRowData(RowKind.INSERT, rowType.getFieldCount()); + List fields = rowType.getFields(); + for (int i = 0; i < rowType.getFieldCount(); i++) { + RowType.RowField subRowField = fields.get(i); + result.setField( + i, getValue(subRowField.getType(), getObjectField(object, subRowField))); + } + return result; + } + + private Object getObjectField(Object object, RowType.RowField rowField) { + String rowFieldName = rowField.getName(); + Class objectClass = object.getClass(); + + // Avro GenericRecord: use the typed API directly. AvroStateUtils.isGenericRecord() caches + // its own answer per class, so this is cheap on every call after the first for a class. + if (AvroStateUtils.isGenericRecord(objectClass)) { + return AvroStateUtils.getGenericRecordField(object, rowFieldName); + } + + Optional field = + classFieldCache() + .computeIfAbsent( + Tuple2.of(objectClass, rowFieldName), + key -> lookupField(objectClass, rowFieldName)); + if (field.isPresent()) { + try { + return field.get().get(object); + } catch (IllegalAccessException e) { + throw new UnsupportedOperationException( + "Cannot access field by either public member or getter function: " + + rowFieldName); + } + } + + Method getter = getGetter(objectClass, rowFieldName); + if (getter == null) { + throw new UnsupportedOperationException( + "Cannot access field by either public member or getter function: " + + rowFieldName); + } + try { + return getter.invoke(object); + } catch (IllegalAccessException | InvocationTargetException e) { + throw new RuntimeException(e); + } + } + + private static Optional lookupField(Class objectClass, String fieldName) { + try { + return Optional.of(objectClass.getField(fieldName)); + } catch (NoSuchFieldException e) { + return Optional.empty(); + } + } + + private Method getGetter(Class objectClass, String rowFieldName) { + String capitalized = rowFieldName.substring(0, 1).toUpperCase() + rowFieldName.substring(1); + Optional getter = lookupMethod(objectClass, "get" + capitalized); + if (getter.isPresent()) { + return getter.get(); + } + return lookupMethod(objectClass, "is" + capitalized).orElse(null); + } + + private Optional lookupMethod(Class objectClass, String methodName) { + return classMethodCache() + .computeIfAbsent( + Tuple2.of(objectClass, methodName), + key -> { + try { + return Optional.of(objectClass.getMethod(methodName)); + } catch (NoSuchMethodException e) { + return Optional.empty(); + } + }); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/TableMappingSupport.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/TableMappingSupport.java new file mode 100644 index 00000000000000..52272e698ea894 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/TableMappingSupport.java @@ -0,0 +1,349 @@ +/* + * 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.api.common.serialization.SerializerConfig; +import org.apache.flink.api.common.state.StateDescriptor; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.configuration.ConfigOption; +import org.apache.flink.configuration.ConfigOptions; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.state.api.OperatorIdentifier; +import org.apache.flink.state.api.runtime.SavepointLoader; +import org.apache.flink.state.api.runtime.SavepointLoader.OperatorStateMetadata; +import org.apache.flink.table.api.ValidationException; +import org.apache.flink.table.types.DataType; +import org.apache.flink.table.types.logical.ArrayType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.LogicalTypeRoot; +import org.apache.flink.table.types.logical.MapType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.table.types.logical.utils.LogicalTypeChecks; +import org.apache.flink.util.Preconditions; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.List; +import java.util.stream.IntStream; + +import static org.apache.flink.state.table.SavepointConnectorOptions.FIELDS; +import static org.apache.flink.state.table.SavepointConnectorOptions.STATE_NAME; + +/** + * Shared static helpers for the table mapping classes. These helpers are pure schema/metadata + * plumbing with no dependency on any single mapping class's field layout, so they are factored out + * here rather than duplicated (or hung off one mapping class for the others to reach into). + */ +final class TableMappingSupport { + + private TableMappingSupport() {} + + /** Returns a {@link ConfigOptions.OptionBuilder} for a per-field connector option. */ + static ConfigOptions.OptionBuilder fieldOption(String fieldName, String suffix) { + return ConfigOptions.key(String.format("%s.%s.%s", FIELDS, fieldName, suffix)); + } + + /** Returns the index of {@code fieldName} in the physical row data type, or {@code -1}. */ + static int columnIndex(DataType physicalDataType, String fieldName) { + final LogicalType physicalType = physicalDataType.getLogicalType(); + Preconditions.checkArgument( + physicalType.is(LogicalTypeRoot.ROW), "Row data type expected."); + return LogicalTypeChecks.getFieldNames(physicalType).indexOf(fieldName); + } + + /** Returns the indices of all columns except {@code excludedIndices}. */ + static int[] valueColumnIndices(DataType physicalDataType, int... excludedIndices) { + final LogicalType physicalType = physicalDataType.getLogicalType(); + Preconditions.checkArgument( + physicalType.is(LogicalTypeRoot.ROW), "Row data type expected."); + final int fieldCount = LogicalTypeChecks.getFieldCount(physicalType); + return IntStream.range(0, fieldCount) + .filter(pos -> IntStream.of(excludedIndices).noneMatch(excluded -> excluded == pos)) + .toArray(); + } + + /** + * Remaps a list of value columns' indices under {@code projectedFields}, dropping any column + * that was projected away. + */ + static List remapValueColumns( + int[][] projectedFields, List valueColumns) { + List newValueColumns = new ArrayList<>(); + for (StateValueColumnConfiguration col : valueColumns) { + int newColumnIndex = remapColumnIndex(projectedFields, col.getColumnIndex()); + if (newColumnIndex >= 0) { + newValueColumns.add(col.withColumnIndex(newColumnIndex)); + } + } + return newValueColumns; + } + + /** + * Returns the output-row index that {@code sourceIndex} is remapped to under {@code + * projectedFields} (see {@link + * org.apache.flink.table.connector.source.abilities.SupportsProjectionPushDown}), or {@code -1} + * if {@code sourceIndex} was projected away. + * + *

Shared by the mappings' {@code project(...)} and the table sources' {@code + * applyProjection(...)}, so both sides compute "where does this column end up after projection" + * the same way. + */ + static int remapColumnIndex(int[][] projectedFields, int sourceIndex) { + for (int outputIdx = 0; outputIdx < projectedFields.length; outputIdx++) { + Preconditions.checkArgument( + projectedFields[outputIdx].length == 1, + "Only flat (non-nested) projections are supported."); + if (projectedFields[outputIdx][0] == sourceIndex) { + return outputIdx; + } + } + return -1; + } + + /** Infers the state type from its SQL logical type. */ + static SavepointConnectorOptions.StateType inferStateType(LogicalType logicalType) { + switch (logicalType.getTypeRoot()) { + case ARRAY: + return SavepointConnectorOptions.StateType.LIST; + case MAP: + return SavepointConnectorOptions.StateType.MAP; + default: + return SavepointConnectorOptions.StateType.VALUE; + } + } + + /** + * Preloads keyed operator metadata (state serializer snapshots + backend key serializer + * snapshot) and builds the {@link SavepointTypeInfoResolver} derived from it, in a single I/O + * operation. Shared by all keyed mapping classes' {@code from(...)} factories. + */ + static SavepointTypeInfoResolver createTypeResolver( + String statePath, + OperatorIdentifier operatorIdentifier, + SerializerConfig serializerConfig) { + OperatorStateMetadata operatorMetadata; + try { + operatorMetadata = SavepointLoader.loadOperatorMetadata(statePath, operatorIdentifier); + } catch (Exception e) { + throw metadataLoadFailure(statePath, operatorIdentifier, e); + } + return new SavepointTypeInfoResolver( + operatorMetadata.stateSnapshots, + serializerConfig, + operatorMetadata.keySerializerSnapshot); + } + + private static RuntimeException metadataLoadFailure( + String statePath, OperatorIdentifier operatorIdentifier, Exception cause) { + return new RuntimeException( + String.format( + "Failed to load state metadata from savepoint '%s' for operator '%s'. " + + "Ensure the savepoint path is valid and the operator exists in the savepoint. ", + statePath, operatorIdentifier), + cause); + } + + @SuppressWarnings("rawtypes") + static StateValueColumnConfiguration createValueColumnConfig( + int columnIndex, + RowType rowType, + Configuration options, + SavepointTypeInfoResolver typeResolver) { + + RowType.RowField valueRowField = rowType.getFields().get(columnIndex); + + ConfigOption stateNameOption = + fieldOption(valueRowField.getName(), STATE_NAME).stringType().noDefaultValue(); + String stateName = options.getOptional(stateNameOption).orElse(valueRowField.getName()); + + StateDescriptor.Type actualStateKind = typeResolver.resolveStateKind(stateName); + SavepointConnectorOptions.StateType stateType = + stateTypeFromActualKind(actualStateKind, valueRowField.getType()); + + TypeSerializer mapKeyTypeSerializer = + typeResolver.resolveMapKeySerializer( + valueRowField, stateType == SavepointConnectorOptions.StateType.MAP); + // A VALUE-shaped state (including REDUCING/AGGREGATING) whose value happens to be a + // List/Map (e.g. a collect-list accumulator) must resolve the serializer for the whole + // value, not the element/entry serializer resolveValueSerializer would unwrap it to for a + // genuine keyed LIST/MAP state. + TypeSerializer valueTypeSerializer = + stateType == SavepointConnectorOptions.StateType.VALUE + ? typeResolver.resolveFlatValueSerializer(valueRowField) + : typeResolver.resolveValueSerializer(valueRowField); + + return new StateValueColumnConfiguration( + columnIndex, + stateName, + stateType, + actualStateKind, + mapKeyTypeSerializer, + valueTypeSerializer); + } + + /** + * Determines the coarse VALUE/LIST/MAP shape used for the SQL schema, preferring the actual + * {@link StateDescriptor.Type} the state was registered under (resolved from the savepoint + * metadata) over inferring it from the column's SQL type. Without this, a {@code + * ValueState}/{@code ReducingState}/{@code AggregatingState} whose value happens to be a + * List/Map (e.g. a collect-list accumulator) would be misclassified as a keyed LIST/MAP state + * purely because its value's SQL type is ARRAY/MAP, causing {@code KeyedStateReader} to open it + * with the wrong state-getter ({@code getListState}/{@code getMapState} instead of {@code + * getState}/{@code getReducingState}/{@code getAggregatingState}) against state that is + * physically stored in an incompatible format. + * + *

Falls back to {@link #inferStateType} when the state is absent from the preloaded metadata + * (i.e. {@code actualStateKind == UNKNOWN}). + */ + private static SavepointConnectorOptions.StateType stateTypeFromActualKind( + StateDescriptor.Type actualStateKind, LogicalType logicalType) { + switch (actualStateKind) { + case LIST: + return SavepointConnectorOptions.StateType.LIST; + case MAP: + return SavepointConnectorOptions.StateType.MAP; + case VALUE: + case REDUCING: + case AGGREGATING: + case FOLDING: + return SavepointConnectorOptions.StateType.VALUE; + default: + return inferStateType(logicalType); + } + } + + /** + * Infers the flattened state type (LIST or MAP) from the sub-key column name and validates that + * the value column is named consistently with it. + * + * @param tableKindLabel e.g. "Flattened keyed state tables", used in validation error messages + * @param subKeyColumnOrdinal ordinal word for the sub-key column's position, e.g. "second" + * @param valueColumnOrdinal ordinal word for the value column's position, e.g. "third" + */ + static SavepointConnectorOptions.StateType inferFlattenedStateTypeAndValidateValueColumn( + String tableKindLabel, + String subKeyColumnOrdinal, + String valueColumnOrdinal, + String subKeyColumnName, + String valueColumnName) { + SavepointConnectorOptions.StateType stateType; + String expectedValueColumnName; + switch (subKeyColumnName) { + case "list_index": + stateType = SavepointConnectorOptions.StateType.LIST; + expectedValueColumnName = "list_value"; + break; + case "map_key": + stateType = SavepointConnectorOptions.StateType.MAP; + expectedValueColumnName = "map_value"; + break; + default: + throw new ValidationException( + tableKindLabel + + " must name their " + + subKeyColumnOrdinal + + " column either 'list_index' (LIST state) or 'map_key' (MAP " + + "state), but found '" + + subKeyColumnName + + "'."); + } + + if (!expectedValueColumnName.equals(valueColumnName)) { + throw new ValidationException( + tableKindLabel + + " must name their " + + valueColumnOrdinal + + " column '" + + expectedValueColumnName + + "', but found '" + + valueColumnName + + "'."); + } + + return stateType; + } + + /** Resolved key type and value-related serializers for a flattened (LIST/MAP) state mapping. */ + static final class FlattenedSerializers { + final TypeInformation keyTypeInfo; + @Nullable final TypeSerializer mapKeyTypeSerializer; + final TypeSerializer valueTypeSerializer; + + FlattenedSerializers( + TypeInformation keyTypeInfo, + @Nullable TypeSerializer mapKeyTypeSerializer, + TypeSerializer valueTypeSerializer) { + this.keyTypeInfo = keyTypeInfo; + this.mapKeyTypeSerializer = mapKeyTypeSerializer; + this.valueTypeSerializer = valueTypeSerializer; + } + } + + /** + * Resolves the key type and value-related serializers shared by {@link + * FlattenedStateTableMapping} and {@link WindowFlattenedStateTableMapping}'s {@code from(...)} + * factories. + */ + static FlattenedSerializers resolveFlattenedSerializers( + RowType rowType, + SavepointTypeInfoResolver typeResolver, + String stateName, + SavepointConnectorOptions.StateType stateType, + int stateKeyColumnIndex, + int subKeyColumnIndex, + int valueColumnIndex) { + TypeInformation keyTypeInfo = + typeResolver.resolveKeyType(rowType.getFields().get(stateKeyColumnIndex)); + + RowType.RowField compositeValueField = + buildCompositeValueField( + rowType, stateName, stateType, subKeyColumnIndex, valueColumnIndex); + + TypeSerializer mapKeyTypeSerializer = + typeResolver.resolveMapKeySerializer( + compositeValueField, stateType == SavepointConnectorOptions.StateType.MAP); + TypeSerializer valueTypeSerializer = + typeResolver.resolveValueSerializer(compositeValueField); + + return new FlattenedSerializers(keyTypeInfo, mapKeyTypeSerializer, valueTypeSerializer); + } + + /** + * Builds a synthetic {@link RowType.RowField} for a flattened LIST/MAP state's composite + * (ArrayType/MapType) value, keyed by {@code stateName} so metadata lookup succeeds, mirroring + * how the general (non-flattened) path resolves value columns. + */ + private static RowType.RowField buildCompositeValueField( + RowType rowType, + String stateName, + SavepointConnectorOptions.StateType stateType, + int subKeyColumnIndex, + int valueColumnIndex) { + LogicalType valueLogicalType = rowType.getFields().get(valueColumnIndex).getType(); + LogicalType compositeLogicalType = + stateType == SavepointConnectorOptions.StateType.LIST + ? new ArrayType(valueLogicalType) + : new MapType( + rowType.getFields().get(subKeyColumnIndex).getType(), + valueLogicalType); + return new RowType.RowField(stateName, compositeLogicalType); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/resources/META-INF/services/org.apache.flink.table.factories.Factory b/flink-libraries/flink-state-processing-api/src/main/resources/META-INF/services/org.apache.flink.table.factories.Factory index c5e2715f26dd17..c8bac30c71d822 100644 --- a/flink-libraries/flink-state-processing-api/src/main/resources/META-INF/services/org.apache.flink.table.factories.Factory +++ b/flink-libraries/flink-state-processing-api/src/main/resources/META-INF/services/org.apache.flink.table.factories.Factory @@ -15,3 +15,4 @@ org.apache.flink.state.table.module.StateModuleFactory org.apache.flink.state.table.SavepointDynamicTableSourceFactory +org.apache.flink.state.catalog.StateCatalogFactory diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/EmbeddedRocksDBKeyedStateReadingITCase.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/EmbeddedRocksDBKeyedStateReadingITCase.java new file mode 100644 index 00000000000000..3bfe8752f27268 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/EmbeddedRocksDBKeyedStateReadingITCase.java @@ -0,0 +1,31 @@ +/* + * 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; + +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.StateBackendOptions; + +/** Runs {@link KeyedStateReadingITCase} against the embedded RocksDB state backend. */ +public class EmbeddedRocksDBKeyedStateReadingITCase extends KeyedStateReadingITCase { + + @Override + protected Configuration getConfiguration() { + return new Configuration().set(StateBackendOptions.STATE_BACKEND, "rocksdb"); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/HashMapKeyedStateReadingITCase.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/HashMapKeyedStateReadingITCase.java new file mode 100644 index 00000000000000..bd9294ed2101d3 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/HashMapKeyedStateReadingITCase.java @@ -0,0 +1,31 @@ +/* + * 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; + +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.StateBackendOptions; + +/** Runs {@link KeyedStateReadingITCase} against the heap ({@code hashmap}) state backend. */ +public class HashMapKeyedStateReadingITCase extends KeyedStateReadingITCase { + + @Override + protected Configuration getConfiguration() { + return new Configuration().set(StateBackendOptions.STATE_BACKEND, "hashmap"); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/KeyedStateReadingITCase.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/KeyedStateReadingITCase.java new file mode 100644 index 00000000000000..4678ade760aae0 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/KeyedStateReadingITCase.java @@ -0,0 +1,404 @@ +/* + * 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; + +import org.apache.flink.api.common.functions.OpenContext; +import org.apache.flink.api.common.state.ListState; +import org.apache.flink.api.common.state.ListStateDescriptor; +import org.apache.flink.api.common.state.MapState; +import org.apache.flink.api.common.state.MapStateDescriptor; +import org.apache.flink.api.common.state.ValueState; +import org.apache.flink.api.common.state.ValueStateDescriptor; +import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.api.java.typeutils.runtime.PojoSerializerSnapshot; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.runtime.checkpoint.OperatorState; +import org.apache.flink.runtime.checkpoint.metadata.CheckpointMetadata; +import org.apache.flink.state.api.input.deserializer.PojoToRowDataDeserializer; +import org.apache.flink.state.api.runtime.SavepointLoader; +import org.apache.flink.state.api.schema.KeyedStateSchemaInfo; +import org.apache.flink.state.api.schema.StateSchemaExtractor; +import org.apache.flink.state.api.schema.StateSchemaInfo; +import org.apache.flink.state.api.utils.SavepointTestBase; +import org.apache.flink.state.catalog.StateCatalog; +import org.apache.flink.state.table.module.StateModule; +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.table.api.EnvironmentSettings; +import org.apache.flink.table.api.Schema; +import org.apache.flink.table.api.Table; +import org.apache.flink.table.api.TableEnvironment; +import org.apache.flink.table.api.TableResult; +import org.apache.flink.table.api.bridge.java.StreamTableEnvironment; +import org.apache.flink.table.types.logical.LogicalTypeRoot; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.types.Row; +import org.apache.flink.util.CloseableIterator; +import org.apache.flink.util.Collector; + +import org.junit.jupiter.api.Test; + +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Objects; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Integration tests that write real keyed state through a MiniCluster job, take a savepoint at + * runtime, and read it back — verified against both the heap ({@code hashmap}) and RocksDB state + * backends (see {@code HashMapKeyedStateReadingITCase} / {@code + * EmbeddedRocksDBKeyedStateReadingITCase}) so that schema extraction and reads are checked against + * both keyed-state-handle formats. + * + *

Unlike {@code StateCatalogGeneratedSavepointITCase}, which reads savepoints checked in as test + * resources (necessarily HashMap-only, since RocksDB fixtures can't be generated locally), every + * savepoint here is produced at test run time, so it runs on whichever backend the subclass + * configures. + */ +public abstract class KeyedStateReadingITCase extends SavepointTestBase { + + protected abstract Configuration getConfiguration(); + + /** Deliberately simple so no Kryo or special serializers are needed. */ + public static class PersonPojo { + public String name; + public int age; + public long score; + + public PersonPojo() {} + + public PersonPojo(String name, int age, long score) { + this.name = name; + this.age = age; + this.score = score; + } + + @Override + public boolean equals(Object o) { + if (!(o instanceof PersonPojo)) { + return false; + } + PersonPojo other = (PersonPojo) o; + return Objects.equals(name, other.name) && age == other.age && score == other.score; + } + } + + // ------------------------------------------------------------------------- + // Schema extraction: RowData-typed internal SQL operator state + // ------------------------------------------------------------------------- + + private static final ValueStateDescriptor PERSON_STATE_DESC = + new ValueStateDescriptor<>("person", PersonPojo.class); + + @Test + public void testGroupAggAccStateSchemaExtraction() throws Exception { + StreamExecutionEnvironment env = + StreamExecutionEnvironment.getExecutionEnvironment(getConfiguration()); + env.setParallelism(1); + + StreamTableEnvironment tEnv = StreamTableEnvironment.create(env); + + Tuple2[] data = + new Tuple2[] { + Tuple2.of("a", 1L), Tuple2.of("a", 2L), Tuple2.of("b", 3L), + }; + DataStream> source = env.addSource(createSource(data)); + tEnv.createTemporaryView( + "t", + source, + Schema.newBuilder().column("f0", "STRING").column("f1", "BIGINT").build()); + + Table result = + tEnv.sqlQuery( + "SELECT f0 AS `key`, COUNT(*) AS cnt, SUM(f1) AS total FROM t GROUP BY f0"); + + DataStream resultStream = tEnv.toChangelogStream(result); + resultStream.sinkTo(new DiscardingSink<>()); + + String savepointPath = takeSavepoint(env); + CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(savepointPath); + + // SQL-planned operators don't carry a user-assigned uid; find the aggregation operator by + // the state it registers. + OperatorIdentifier aggOpId = null; + for (OperatorIdentifier candidate : StateTableUtils.getOperatorIdentifiers(metadata)) { + List stateNames = StateTableUtils.getKeyedStates(metadata, candidate); + if (stateNames.contains("accState")) { + aggOpId = candidate; + break; + } + } + assertNotNull(aggOpId, "Could not find operator with 'accState'"); + + KeyedStateSchemaInfo schemaInfo = StateTableUtils.getKeyedStateSchema(metadata, aggOpId); + KeyedStateSchemaInfo.StateEntryInfo accEntry = schemaInfo.stateSchemas.get("accState"); + assertNotNull(accEntry, "'accState' not found in extracted schema"); + + assertEquals(LogicalTypeRoot.ROW, accEntry.logicalType.getTypeRoot()); + RowType rowType = (RowType) accEntry.logicalType; + + // The accumulator row holds one field per aggregate call: COUNT(*) and SUM(f1). + assertEquals(2, rowType.getFieldCount()); + assertEquals(LogicalTypeRoot.BIGINT, rowType.getFields().get(0).getType().getTypeRoot()); + assertEquals(LogicalTypeRoot.BIGINT, rowType.getFields().get(1).getType().getTypeRoot()); + } + + // ------------------------------------------------------------------------- + // Schema extraction: POJO value state + // ------------------------------------------------------------------------- + + private static final String POJO_UID = "pojo-state-operator"; + + @Test + public void testSchemaExtractionFromPojoState() throws Exception { + StreamExecutionEnvironment env = + StreamExecutionEnvironment.getExecutionEnvironment(getConfiguration()); + env.setParallelism(1); + + PersonPojo[] data = { + new PersonPojo("Alice", 30, 100L), + new PersonPojo("Bob", 25, 200L), + new PersonPojo("Carol", 35, 300L) + }; + env.addSource(createSource(data)) + .returns(PersonPojo.class) + .keyBy(p -> p.name) + .process(new PersonStateWriter()) + .uid(POJO_UID) + .sinkTo(new DiscardingSink<>()); + + String savepointPath = takeSavepoint(env); + CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(savepointPath); + + OperatorIdentifier opId = OperatorIdentifier.forUid(POJO_UID); + + List stateNames = StateTableUtils.getKeyedStates(metadata, opId); + assertTrue(stateNames.contains("person"), "Expected 'person' state"); + + KeyedStateSchemaInfo schemaInfo = StateTableUtils.getKeyedStateSchema(metadata, opId); + KeyedStateSchemaInfo.StateEntryInfo personEntry = schemaInfo.stateSchemas.get("person"); + assertNotNull(personEntry, "'person' state not found in schema"); + + // PersonPojo has 3 fields → the logicalType should be a RowType with 3 fields + assertEquals(LogicalTypeRoot.ROW, personEntry.logicalType.getTypeRoot()); + RowType rowType = (RowType) personEntry.logicalType; + assertEquals(3, rowType.getFieldCount()); + assertHasField(rowType, "name", LogicalTypeRoot.VARCHAR); + assertHasField(rowType, "age", LogicalTypeRoot.INTEGER); + assertHasField(rowType, "score", LogicalTypeRoot.BIGINT); + + // The same snapshot must also be usable to build a deserializer directly (lower-level API). + StateSchemaInfo personRaw = + StateSchemaExtractor.extractSchema(findOperatorState(metadata, opId)).stream() + .filter(s -> "person".equals(s.stateName)) + .findFirst() + .orElse(null); + assertNotNull(personRaw); + assertNotNull( + PojoToRowDataDeserializer.create( + (PojoSerializerSnapshot) personRaw.valueSnapshot)); + } + + // ------------------------------------------------------------------------- + // End-to-end read through StateCatalog + SQL: primitive, POJO, list and map state + // ------------------------------------------------------------------------- + // + // This is the backend-parameterized equivalent of + // StateCatalogGeneratedSavepointITCase.SchemaDiscoveryWithoutSourceClasses — same state + // shapes, but written and savepointed at test run time instead of read from a checked-in + // HashMap-only fixture, so it also runs on RocksDB. + + private static final String MIXED_STATE_UID = "mixed-state-operator"; + private static final ValueStateDescriptor COUNT_STATE_DESC = + new ValueStateDescriptor<>("count", Long.class); + private static final ListStateDescriptor ITEMS_STATE_DESC = + new ListStateDescriptor<>("items", Long.class); + private static final MapStateDescriptor COUNTS_STATE_DESC = + new MapStateDescriptor<>("counts", Long.class, Long.class); + + @Test + public void testReadPrimitivePojoListAndMapStateThroughCatalog() throws Exception { + StreamExecutionEnvironment env = + StreamExecutionEnvironment.getExecutionEnvironment(getConfiguration()); + env.setParallelism(1); + + Long[] keys = {1L, 2L, 3L}; + env.addSource(createSource(keys)) + .returns(Long.class) + .keyBy(k -> k) + .process(new MixedStateWriter()) + .uid(MIXED_STATE_UID) + .sinkTo(new DiscardingSink<>()); + + String savepointPath = takeSavepoint(env); + // takeSavepoint() returns a "file:" URI string, not a plain filesystem path. + Path catalogRoot = Paths.get(java.net.URI.create(savepointPath)).getParent(); + + StateCatalog catalog = + new StateCatalog( + "state", + Collections.singletonMap("test", catalogRoot.toAbsolutePath().toString())); + catalog.open(); + try { + List dbs = catalog.listDatabases(); + assertEquals(1, dbs.size()); + String dbName = dbs.get(0); + + TableEnvironment tableEnv = TableEnvironment.create(EnvironmentSettings.inBatchMode()); + tableEnv.loadModule("state", StateModule.INSTANCE); + tableEnv.registerCatalog("state", catalog); + tableEnv.useCatalog("state"); + tableEnv.useDatabase(dbName); + + String mainTable = + "`" + + StateCatalog.OPERATOR_UID_PREFIX + + MIXED_STATE_UID + + StateCatalog.OPERATOR_TABLE_SUFFIX + + "`"; + List rows = collectWithSql(tableEnv, "SELECT * FROM " + mainTable); + assertEquals(3, rows.size()); + for (Row row : rows) { + Long key = (Long) row.getField("state_key"); + assertNotNull(key); + assertEquals(key, row.getField("count")); + + Row person = (Row) row.getField("person"); + assertNotNull(person); + assertEquals("name-" + key, person.getField("name")); + assertEquals(key * 100, person.getField("score")); + } + + String listFlatTable = + "`" + + StateCatalog.OPERATOR_UID_PREFIX + + MIXED_STATE_UID + + "_items" + + StateCatalog.FLAT_STATE_TABLE_SUFFIX + + "`"; + List listRows = + collectWithSql( + tableEnv, "SELECT * FROM " + listFlatTable + " WHERE state_key = 2"); + assertEquals(1, listRows.size()); + assertEquals(2L, listRows.get(0).getField("state_key")); + assertEquals(20L, listRows.get(0).getField("list_value")); + + String mapFlatTable = + "`" + + StateCatalog.OPERATOR_UID_PREFIX + + MIXED_STATE_UID + + "_counts" + + StateCatalog.FLAT_STATE_TABLE_SUFFIX + + "`"; + List mapRows = + collectWithSql( + tableEnv, "SELECT * FROM " + mapFlatTable + " WHERE state_key = 2"); + assertEquals(1, mapRows.size()); + assertEquals(2L, mapRows.get(0).getField("state_key")); + assertEquals(2L, mapRows.get(0).getField("map_key")); + assertEquals(2L, mapRows.get(0).getField("map_value")); + } finally { + catalog.close(); + } + } + + // ------------------------------------------------------------------------- + // Helpers + // ------------------------------------------------------------------------- + + private static OperatorState findOperatorState( + CheckpointMetadata metadata, OperatorIdentifier opId) { + for (OperatorState op : metadata.getOperatorStates()) { + if (op.getOperatorID().equals(opId.getOperatorId())) { + return op; + } + } + throw new IllegalArgumentException("Operator not found: " + opId); + } + + private static void assertHasField(RowType row, String name, LogicalTypeRoot expectedRoot) { + RowType.RowField field = + row.getFields().stream() + .filter(f -> f.getName().equals(name)) + .findFirst() + .orElse(null); + assertNotNull(field, "Field '" + name + "' not found in row type"); + assertEquals( + expectedRoot, field.getType().getTypeRoot(), "Wrong type for field '" + name + "'"); + } + + private static List collectWithSql(TableEnvironment tEnv, String sql) throws Exception { + List rows = new ArrayList<>(); + TableResult result = tEnv.executeSql(sql); + try (CloseableIterator it = result.collect()) { + it.forEachRemaining(rows::add); + } + return rows; + } + + // ------------------------------------------------------------------------- + // Operators + // ------------------------------------------------------------------------- + + private static class PersonStateWriter extends KeyedProcessFunction { + private transient ValueState state; + + @Override + public void open(OpenContext ctx) throws Exception { + state = getRuntimeContext().getState(PERSON_STATE_DESC); + } + + @Override + public void processElement(PersonPojo value, Context ctx, Collector out) + throws Exception { + state.update(value); + } + } + + private static class MixedStateWriter extends KeyedProcessFunction { + private transient ValueState countState; + private transient ValueState personState; + private transient ListState itemsState; + private transient MapState countsState; + + @Override + public void open(OpenContext ctx) throws Exception { + countState = getRuntimeContext().getState(COUNT_STATE_DESC); + personState = getRuntimeContext().getState(PERSON_STATE_DESC); + itemsState = getRuntimeContext().getListState(ITEMS_STATE_DESC); + countsState = getRuntimeContext().getMapState(COUNTS_STATE_DESC); + } + + @Override + public void processElement(Long key, Context ctx, Collector out) throws Exception { + countState.update(key); + personState.update(new PersonPojo("name-" + key, key.intValue(), key * 100)); + itemsState.add(key * 10); + countsState.put(key, key); + } + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/StateTableUtilsTest.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/StateTableUtilsTest.java new file mode 100644 index 00000000000000..34a6e16bda1a8c --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/StateTableUtilsTest.java @@ -0,0 +1,129 @@ +/* + * 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; + +import org.apache.flink.runtime.checkpoint.OperatorState; +import org.apache.flink.runtime.checkpoint.metadata.CheckpointMetadata; +import org.apache.flink.runtime.jobgraph.OperatorID; +import org.apache.flink.runtime.state.StateBackendLoader; +import org.apache.flink.state.api.runtime.SavepointLoader; + +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Optional; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** Unit tests for {@link StateTableUtils} that do not require a running Flink cluster. */ +public class StateTableUtilsTest { + + // ------------------------------------------------------------------------- + // getOperatorIdentifiers — filters operators without keyed state + // ------------------------------------------------------------------------- + + @Test + public void testGetOperatorIdentifiersFiltersEmptyOperators() { + OperatorState opState1 = new OperatorState(null, null, new OperatorID(1L, 2L), 1, 128); + OperatorState opState2 = new OperatorState(null, null, new OperatorID(3L, 4L), 2, 128); + + // No subtasks means no keyed state, regardless of how many such operators are present. + List> cases = + Arrays.asList( + Collections.singletonList(opState1), + Collections.emptyList(), + Arrays.asList(opState1, opState2)); + + for (List operators : cases) { + CheckpointMetadata metadata = + new CheckpointMetadata(1L, operators, Collections.emptyList()); + List ids = StateTableUtils.getOperatorIdentifiers(metadata); + + assertNotNull(ids); + assertTrue( + ids.isEmpty(), + "Operators without keyed state should be filtered out, input: " + operators); + } + } + + // ------------------------------------------------------------------------- + // detectStateBackendType — reads real checkpoint metadata produced by other tests + // ------------------------------------------------------------------------- + + /** + * Checkpoint (native format) taken with the HashMap state backend, committed as a fixture for + * {@code StatefulJobSnapshotMigrationITCase} in flink-tests. Native-format keyed state handles + * retain their backend-specific type, so every keyed operator here must resolve to {@link + * StateBackendLoader#HASHMAP_STATE_BACKEND_NAME}. + */ + private static final String HASHMAP_CHECKPOINT_DIR = + "../../flink-tests/src/test/resources/" + + "new-stateful-udf-migration-itcase-flink1.20-hashmap-checkpoint"; + + /** + * Checkpoint (native format) taken with the RocksDB state backend, committed as a fixture for + * {@code StatefulJobSnapshotMigrationITCase} in flink-tests. Must resolve to {@link + * StateBackendLoader#ROCKSDB_STATE_BACKEND_NAME}. + */ + private static final String ROCKSDB_CHECKPOINT_DIR = + "../../flink-tests/src/test/resources/" + + "new-stateful-udf-migration-itcase-flink2.1-rocksdb-checkpoint"; + + @Test + public void testDetectStateBackendTypeFromHashMapCheckpoint() throws IOException { + assertAllKeyedOperatorsDetectAs( + HASHMAP_CHECKPOINT_DIR, StateBackendLoader.HASHMAP_STATE_BACKEND_NAME); + } + + @Test + public void testDetectStateBackendTypeFromRocksDBCheckpoint() throws IOException { + assertAllKeyedOperatorsDetectAs( + ROCKSDB_CHECKPOINT_DIR, StateBackendLoader.ROCKSDB_STATE_BACKEND_NAME); + } + + /** + * Loads the checkpoint metadata at {@code checkpointDir} and asserts that every operator + * carrying keyed state resolves to exactly {@code expectedType}, and that at least one operator + * did so (i.e. the fixture actually exercises the detection logic). + */ + private static void assertAllKeyedOperatorsDetectAs(String checkpointDir, String expectedType) + throws IOException { + CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(checkpointDir); + assertFalse(metadata.getOperatorStates().isEmpty()); + + int operatorsWithKeyedState = 0; + for (OperatorState opState : metadata.getOperatorStates()) { + Optional detected = StateTableUtils.detectStateBackendType(opState); + if (detected.isEmpty()) { + continue; + } + operatorsWithKeyedState++; + assertEquals(expectedType, detected.get(), "operator " + opState.getOperatorID()); + } + assertTrue( + operatorsWithKeyedState > 0, + "Expected at least one operator with keyed state in " + checkpointDir); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/MultiStateKeyIteratorTest.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/MultiStateKeyIteratorTest.java index c6f011721d2386..fbe0bfb80b1eaa 100644 --- a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/MultiStateKeyIteratorTest.java +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/MultiStateKeyIteratorTest.java @@ -173,7 +173,7 @@ void testIteratorPullsKeyFromAllDescriptors() throws Exception { List keys = new ArrayList<>(); while (iterator.hasNext()) { - keys.add(iterator.next()); + keys.add(iterator.next().f0); } assertThat(keys).containsExactly(1, 2); @@ -202,7 +202,7 @@ void testIteratorSkipsEmptyDescriptors() throws Exception { List keys = new ArrayList<>(); while (iterator.hasNext()) { - keys.add(iterator.next()); + keys.add(iterator.next().f0); } assertThat(keys).containsExactly(1, 2); @@ -260,9 +260,16 @@ public CountingKeysKeyedStateBackend( @Override public Stream getKeys(List states, N namespace) { + return getKeysAndKeyGroups(states, namespace).map(t -> t.f0); + } + + @Override + public Stream> getKeysAndKeyGroups( + List states, N namespace) { return IntStream.range(0, this.numberOfKeysGenerated) .boxed() - .peek(i -> numberOfKeysEnumerated++); + .peek(i -> numberOfKeysEnumerated++) + .map(i -> Tuple2.of(i, 0)); } @Override diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/deserializer/InternalTypeConverterTest.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/deserializer/InternalTypeConverterTest.java new file mode 100644 index 00000000000000..6efe86b3f86079 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/deserializer/InternalTypeConverterTest.java @@ -0,0 +1,301 @@ +/* + * 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.deserializer; + +import org.apache.flink.table.data.DecimalData; +import org.apache.flink.table.data.GenericArrayData; +import org.apache.flink.table.data.GenericMapData; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.StringData; +import org.apache.flink.table.data.TimestampData; +import org.apache.flink.table.types.logical.ArrayType; +import org.apache.flink.table.types.logical.BigIntType; +import org.apache.flink.table.types.logical.BooleanType; +import org.apache.flink.table.types.logical.DateType; +import org.apache.flink.table.types.logical.DayTimeIntervalType; +import org.apache.flink.table.types.logical.DecimalType; +import org.apache.flink.table.types.logical.DoubleType; +import org.apache.flink.table.types.logical.FloatType; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.LocalZonedTimestampType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.MapType; +import org.apache.flink.table.types.logical.MultisetType; +import org.apache.flink.table.types.logical.NullType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.table.types.logical.SmallIntType; +import org.apache.flink.table.types.logical.TimeType; +import org.apache.flink.table.types.logical.TimestampType; +import org.apache.flink.table.types.logical.TinyIntType; +import org.apache.flink.table.types.logical.VarBinaryType; +import org.apache.flink.table.types.logical.VarCharType; +import org.apache.flink.table.types.logical.YearMonthIntervalType; +import org.apache.flink.table.types.logical.ZonedTimestampType; +import org.apache.flink.types.Row; +import org.apache.flink.types.RowKind; + +import org.junit.jupiter.api.Test; + +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.sql.Timestamp; +import java.time.Instant; +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.util.Arrays; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; + +/** Unit tests for {@link InternalTypeConverter}. */ +public class InternalTypeConverterTest { + + @Test + public void testNullReturnsNull() { + assertNull(InternalTypeConverter.toInternal(null, new IntType())); + assertNull(InternalTypeConverter.toInternal(null, new VarCharType())); + assertNull(InternalTypeConverter.toInternal("anything", new NullType())); + } + + @Test + public void testVarChar() { + // String → StringData + assertEquals( + StringData.fromString("hello"), + InternalTypeConverter.toInternal("hello", new VarCharType())); + // StringData → pass-through + StringData sd = StringData.fromString("world"); + assertSame(sd, InternalTypeConverter.toInternal(sd, new VarCharType())); + // Other type → toString() + assertEquals( + StringData.fromString("42"), + InternalTypeConverter.toInternal(42, new VarCharType())); + } + + @Test + public void testPrimitivePassThroughs() { + // All of these are returned unchanged. + assertSame(Boolean.TRUE, InternalTypeConverter.toInternal(true, new BooleanType())); + Byte b = (byte) 7; + assertSame(b, InternalTypeConverter.toInternal(b, new TinyIntType())); + Short s = (short) 100; + assertSame(s, InternalTypeConverter.toInternal(s, new SmallIntType())); + Integer i = 42; + assertSame(i, InternalTypeConverter.toInternal(i, new IntType())); + Long l = 123L; + assertSame(l, InternalTypeConverter.toInternal(l, new BigIntType())); + Float f = 1.5f; + assertSame(f, InternalTypeConverter.toInternal(f, new FloatType())); + Double d = 3.14; + assertSame(d, InternalTypeConverter.toInternal(d, new DoubleType())); + Integer timeMillis = 3_600_000; + assertSame(timeMillis, InternalTypeConverter.toInternal(timeMillis, new TimeType())); + Long months = 13L; + assertSame( + months, + InternalTypeConverter.toInternal( + months, + new YearMonthIntervalType( + YearMonthIntervalType.YearMonthResolution.YEAR_TO_MONTH))); + Long dayMillis = 86_400_000L; + assertSame( + dayMillis, + InternalTypeConverter.toInternal( + dayMillis, + new DayTimeIntervalType(DayTimeIntervalType.DayTimeResolution.DAY))); + } + + @Test + public void testDecimal() { + DecimalType type = new DecimalType(10, 2); + // BigDecimal → DecimalData + BigDecimal bd = new BigDecimal("12.34"); + assertEquals( + DecimalData.fromBigDecimal(bd, 10, 2), InternalTypeConverter.toInternal(bd, type)); + // byte[] → DecimalData (unscaled bytes) + byte[] unscaledBytes = BigDecimal.valueOf(1234).unscaledValue().toByteArray(); + assertEquals( + DecimalData.fromUnscaledBytes(unscaledBytes, 10, 2), + InternalTypeConverter.toInternal(unscaledBytes, type)); + // ByteBuffer → DecimalData (unscaled bytes) + assertEquals( + DecimalData.fromUnscaledBytes(unscaledBytes, 10, 2), + InternalTypeConverter.toInternal(ByteBuffer.wrap(unscaledBytes), type)); + // DecimalData → pass-through + DecimalData dd = DecimalData.fromBigDecimal(new BigDecimal("9.99"), 10, 2); + assertSame(dd, InternalTypeConverter.toInternal(dd, type)); + } + + @Test + public void testDate() { + // Integer (epoch day) → pass-through + Integer epochDay = 19_000; + assertSame(epochDay, InternalTypeConverter.toInternal(epochDay, new DateType())); + // LocalDate → epoch day int + LocalDate ld = LocalDate.of(2022, 6, 15); + assertEquals((int) ld.toEpochDay(), InternalTypeConverter.toInternal(ld, new DateType())); + // java.sql.Date → epoch day int + java.sql.Date sqlDate = java.sql.Date.valueOf("2022-06-15"); + assertEquals( + (int) sqlDate.toLocalDate().toEpochDay(), + InternalTypeConverter.toInternal(sqlDate, new DateType())); + } + + @Test + public void testTimestamp() { + Timestamp ts = Timestamp.valueOf("2023-01-15 10:30:00"); + Instant instant = Instant.parse("2023-01-15T10:30:00Z"); + LocalDateTime ldt = LocalDateTime.of(2023, 1, 15, 10, 30, 0); + + // All three timestamp type roots accept the same source types. + for (LogicalType tsType : + new LogicalType[] { + new TimestampType(), new ZonedTimestampType(), new LocalZonedTimestampType() + }) { + assertEquals( + TimestampData.fromTimestamp(ts), InternalTypeConverter.toInternal(ts, tsType)); + assertEquals( + TimestampData.fromInstant(instant), + InternalTypeConverter.toInternal(instant, tsType)); + assertEquals( + TimestampData.fromLocalDateTime(ldt), + InternalTypeConverter.toInternal(ldt, tsType)); + } + // TimestampData → pass-through + TimestampData td = TimestampData.fromEpochMillis(1000L); + assertSame(td, InternalTypeConverter.toInternal(td, new TimestampType())); + } + + @Test + public void testBinary() { + byte[] bytes = {1, 2, 3}; + // byte[] → pass-through + assertSame(bytes, InternalTypeConverter.toInternal(bytes, new VarBinaryType())); + // ByteBuffer → extracted byte[] + assertArrayEquals( + bytes, + (byte[]) + InternalTypeConverter.toInternal( + ByteBuffer.wrap(bytes), new VarBinaryType())); + } + + @Test + public void testRow() { + RowType rowType = RowType.of(new VarCharType(), new IntType()); + // Flink Row → GenericRowData with recursive field conversion + Row row = Row.ofKind(RowKind.INSERT, "Alice", 30); + GenericRowData result = (GenericRowData) InternalTypeConverter.toInternal(row, rowType); + assertEquals(StringData.fromString("Alice"), result.getString(0)); + assertEquals(30, result.getInt(1)); + // GenericRowData → pass-through + GenericRowData grd = GenericRowData.of(StringData.fromString("x"), 1); + assertSame(grd, InternalTypeConverter.toInternal(grd, rowType)); + } + + @Test + public void testArray() { + ArrayType intArrayType = new ArrayType(new IntType()); + ArrayType strArrayType = new ArrayType(new VarCharType()); + + // List → GenericArrayData + GenericArrayData fromList = + (GenericArrayData) + InternalTypeConverter.toInternal(Arrays.asList(1, 2, 3), intArrayType); + assertEquals(3, fromList.size()); + assertEquals(1, fromList.getInt(0)); + assertEquals(3, fromList.getInt(2)); + + // Object[] → GenericArrayData with recursive element conversion + GenericArrayData fromObjectArray = + (GenericArrayData) + InternalTypeConverter.toInternal(new Object[] {"a", "b"}, strArrayType); + assertEquals(StringData.fromString("a"), fromObjectArray.getString(0)); + assertEquals(StringData.fromString("b"), fromObjectArray.getString(1)); + + // Iterable → GenericArrayData (ListState returns Iterable) + GenericArrayData fromIterable = + (GenericArrayData) + InternalTypeConverter.toInternal( + Arrays.asList(10L, 20L), new ArrayType(new BigIntType())); + assertEquals(10L, fromIterable.getLong(0)); + assertEquals(20L, fromIterable.getLong(1)); + + // GenericArrayData → pass-through + GenericArrayData gad = new GenericArrayData(new Object[] {1, 2}); + assertSame(gad, InternalTypeConverter.toInternal(gad, intArrayType)); + } + + @Test + public void testMap() { + MapType type = new MapType(new VarCharType(), new IntType()); + + // Map → GenericMapData with recursive key/value conversion + Map map = new LinkedHashMap<>(); + map.put("a", 1); + map.put("b", 2); + GenericMapData fromMap = (GenericMapData) InternalTypeConverter.toInternal(map, type); + assertEquals(2, fromMap.size()); + assertEquals(1, fromMap.get(StringData.fromString("a"))); + assertEquals(2, fromMap.get(StringData.fromString("b"))); + + // Iterable → GenericMapData (MapState.entries() returns this) + GenericMapData fromEntries = + (GenericMapData) InternalTypeConverter.toInternal(map.entrySet(), type); + assertEquals(1, fromEntries.get(StringData.fromString("a"))); + assertEquals(2, fromEntries.get(StringData.fromString("b"))); + + // GenericMapData → pass-through + Map inner = new HashMap<>(); + inner.put(StringData.fromString("k"), 99); + GenericMapData gmd = new GenericMapData(inner); + assertSame(gmd, InternalTypeConverter.toInternal(gmd, type)); + } + + @Test + public void testMultiset() { + // MultisetType is not a MapType: it only carries an element type, and is represented + // internally as Map (element -> multiplicity). + MultisetType type = new MultisetType(new VarCharType()); + + Map map = new LinkedHashMap<>(); + map.put("a", 3); + map.put("b", 1); + GenericMapData fromMap = (GenericMapData) InternalTypeConverter.toInternal(map, type); + assertEquals(2, fromMap.size()); + assertEquals(3, fromMap.get(StringData.fromString("a"))); + assertEquals(1, fromMap.get(StringData.fromString("b"))); + + // Iterable → GenericMapData + GenericMapData fromEntries = + (GenericMapData) InternalTypeConverter.toInternal(map.entrySet(), type); + assertEquals(3, fromEntries.get(StringData.fromString("a"))); + assertEquals(1, fromEntries.get(StringData.fromString("b"))); + + // GenericMapData → pass-through + Map inner = new HashMap<>(); + inner.put(StringData.fromString("k"), 5); + GenericMapData gmd = new GenericMapData(inner); + assertSame(gmd, InternalTypeConverter.toInternal(gmd, type)); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/deserializer/PojoToRowDataDeserializerTest.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/deserializer/PojoToRowDataDeserializerTest.java new file mode 100644 index 00000000000000..37905ce0220e1f --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/deserializer/PojoToRowDataDeserializerTest.java @@ -0,0 +1,239 @@ +/* + * 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.deserializer; + +import org.apache.flink.api.common.serialization.SerializerConfigImpl; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.api.java.typeutils.TypeExtractor; +import org.apache.flink.api.java.typeutils.runtime.PojoSerializer; +import org.apache.flink.api.java.typeutils.runtime.PojoSerializerSnapshot; +import org.apache.flink.core.memory.DataInputDeserializer; +import org.apache.flink.core.memory.DataOutputSerializer; +import org.apache.flink.state.api.schema.SerializerSnapshotToLogicalTypeConverter; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.StringData; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.RowType; + +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Unit tests for {@link PojoToRowDataDeserializer}. + * + *

Each test serializes a POJO using the real {@link PojoSerializer}, then deserializes with + * {@link PojoToRowDataDeserializer} — no POJO class needed on the deserialization side. + */ +public class PojoToRowDataDeserializerTest { + + // ------------------------------------------------------------------------- + // POJO classes + // ------------------------------------------------------------------------- + + public static class FlatPojo { + public String name; + public int age; + public long score; + public boolean active; + + public FlatPojo() {} + + public FlatPojo(String name, int age, long score, boolean active) { + this.name = name; + this.age = age; + this.score = score; + this.active = active; + } + } + + public static class PojoWithNullableField { + public String tag; // may be null + public int value; + + public PojoWithNullableField() {} + + public PojoWithNullableField(String tag, int value) { + this.tag = tag; + this.value = value; + } + } + + public static class NestedPojo { + public String label; + public FlatPojo inner; + + public NestedPojo() {} + + public NestedPojo(String label, FlatPojo inner) { + this.label = label; + this.inner = inner; + } + } + + // ------------------------------------------------------------------------- + // Tests + // ------------------------------------------------------------------------- + + @Test + public void testDeserializeFlatPojo() throws IOException { + FlatPojo original = new FlatPojo("Alice", 30, 12345L, true); + + PojoToRowDataDeserializer deserializer = buildDeserializer(FlatPojo.class); + GenericRowData row = (GenericRowData) roundtrip(original, FlatPojo.class, deserializer); + + assertNotNull(row); + assertEquals(4, row.getArity()); + assertEquals( + StringData.fromString("Alice"), + row.getString(indexOfField(FlatPojo.class, "name"))); + assertEquals(30, row.getInt(indexOfField(FlatPojo.class, "age"))); + assertEquals(12345L, row.getLong(indexOfField(FlatPojo.class, "score"))); + assertTrue(row.getBoolean(indexOfField(FlatPojo.class, "active"))); + } + + @Test + public void testDeserializeWithNullField() throws IOException { + PojoWithNullableField original = new PojoWithNullableField(null, 42); + PojoToRowDataDeserializer deserializer = buildDeserializer(PojoWithNullableField.class); + GenericRowData row = + (GenericRowData) roundtrip(original, PojoWithNullableField.class, deserializer); + + assertNotNull(row); + assertTrue(row.isNullAt(indexOfField(PojoWithNullableField.class, "tag"))); + assertEquals(42, row.getInt(indexOfField(PojoWithNullableField.class, "value"))); + } + + @Test + public void testDeserializeNullValue() throws IOException { + TypeSerializer pojoSer = buildPojoSerializer(FlatPojo.class); + DataOutputSerializer out = new DataOutputSerializer(64); + pojoSer.serialize(null, out); + + DataInputDeserializer in = new DataInputDeserializer(out.getSharedBuffer()); + PojoToRowDataDeserializer deserializer = buildDeserializer(FlatPojo.class); + RowData result = deserializer.deserialize(in); + assertNull(result); + } + + @Test + public void testDeserializeNestedPojo() throws IOException { + NestedPojo original = new NestedPojo("outer", new FlatPojo("Bob", 25, 999L, false)); + PojoToRowDataDeserializer deserializer = buildDeserializer(NestedPojo.class); + GenericRowData row = (GenericRowData) roundtrip(original, NestedPojo.class, deserializer); + + assertNotNull(row); + int labelIdx = indexOfField(NestedPojo.class, "label"); + int innerIdx = indexOfField(NestedPojo.class, "inner"); + assertEquals(StringData.fromString("outer"), row.getString(labelIdx)); + + // Nested POJO should be a GenericRowData + RowData innerRow = row.getRow(innerIdx, 4); + assertNotNull(innerRow); + } + + @Test + public void testUnregisteredSubclassThrowsIoException() throws IOException { + // Write a value normally using the serializer, then inject fake IS_SUBCLASS bytes. + DataOutputSerializer out = new DataOutputSerializer(64); + out.writeByte(PojoSerializer.IS_SUBCLASS); + out.writeUTF("com.example.UnknownSubclass"); + + DataInputDeserializer in = new DataInputDeserializer(out.getSharedBuffer()); + PojoToRowDataDeserializer deserializer = buildDeserializer(FlatPojo.class); + + IOException e = assertThrows(IOException.class, () -> deserializer.deserialize(in)); + assertTrue(e.getMessage().contains("UnknownSubclass")); + } + + @Test + public void testTaggedSubclassWithUnresolvableDeserializerThrowsIoException() + throws IOException { + // A registered subclass whose own serializer snapshot could not be turned into a + // PojoToRowDataDeserializer (e.g. it is not a POJO, or its snapshot was unreadable) is + // represented by a `null` entry in registeredSubclassDeserializers (see + // PojoSerializerSnapshot#getRegisteredSubclassSnapshotsOrdered). Deserializing a tagged + // subclass that resolves to such an entry must fail with a clear IOException rather than + // an NPE. + List registeredSubclassDeserializers = new ArrayList<>(); + registeredSubclassDeserializers.add(null); + PojoToRowDataDeserializer deserializer = + new PojoToRowDataDeserializer( + new TypeSerializer[0], + new LogicalType[0], + new String[0], + registeredSubclassDeserializers); + + DataOutputSerializer out = new DataOutputSerializer(8); + out.writeByte(PojoSerializer.IS_TAGGED_SUBCLASS); + out.writeByte(0); + DataInputDeserializer in = new DataInputDeserializer(out.getSharedBuffer()); + + IOException e = assertThrows(IOException.class, () -> deserializer.deserialize(in)); + assertTrue(e.getMessage().contains("tag 0")); + } + + // ------------------------------------------------------------------------- + // Helpers + // ------------------------------------------------------------------------- + + @SuppressWarnings("unchecked") + private static PojoSerializer buildPojoSerializer(Class clazz) { + return (PojoSerializer) + TypeExtractor.createTypeInfo(clazz).createSerializer(new SerializerConfigImpl()); + } + + @SuppressWarnings("unchecked") + private static PojoToRowDataDeserializer buildDeserializer(Class clazz) { + PojoSerializer ser = buildPojoSerializer(clazz); + PojoSerializerSnapshot snapshot = + (PojoSerializerSnapshot) ser.snapshotConfiguration(); + return PojoToRowDataDeserializer.create(snapshot); + } + + private static RowData roundtrip(T value, Class clazz, PojoToRowDataDeserializer deser) + throws IOException { + PojoSerializer ser = buildPojoSerializer(clazz); + DataOutputSerializer out = new DataOutputSerializer(256); + ser.serialize(value, out); + + DataInputDeserializer in = new DataInputDeserializer(out.getSharedBuffer()); + return deser.deserialize(in); + } + + /** Returns the field index as it appears in the PojoSerializer's field ordering. */ + private static int indexOfField(Class clazz, String fieldName) { + RowType rowType = + (RowType) + SerializerSnapshotToLogicalTypeConverter.convert( + buildPojoSerializer(clazz).snapshotConfiguration()); + int idx = rowType.getFieldNames().indexOf(fieldName); + assertTrue(idx >= 0, "Field '" + fieldName + "' not found"); + return idx; + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/SerializerSnapshotToLogicalTypeConverterTest.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/SerializerSnapshotToLogicalTypeConverterTest.java new file mode 100644 index 00000000000000..3773b67dd67575 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/SerializerSnapshotToLogicalTypeConverterTest.java @@ -0,0 +1,425 @@ +/* + * 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.schema; + +import org.apache.flink.api.common.serialization.SerializerConfigImpl; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.api.common.typeutils.base.BooleanSerializer; +import org.apache.flink.api.common.typeutils.base.DoubleSerializer; +import org.apache.flink.api.common.typeutils.base.FloatSerializer; +import org.apache.flink.api.common.typeutils.base.IntSerializer; +import org.apache.flink.api.common.typeutils.base.ListSerializer; +import org.apache.flink.api.common.typeutils.base.LongSerializer; +import org.apache.flink.api.common.typeutils.base.MapSerializer; +import org.apache.flink.api.common.typeutils.base.StringSerializer; +import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.api.java.typeutils.TypeExtractor; +import org.apache.flink.api.java.typeutils.runtime.NullableSerializer; +import org.apache.flink.api.java.typeutils.runtime.TupleSerializer; +import org.apache.flink.formats.avro.typeutils.AvroTypeInfo; +import org.apache.flink.formats.avro.typeutils.GenericRecordAvroTypeInfo; +import org.apache.flink.table.runtime.typeutils.RowDataSerializer; +import org.apache.flink.table.types.logical.ArrayType; +import org.apache.flink.table.types.logical.BigIntType; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.LogicalTypeRoot; +import org.apache.flink.table.types.logical.MapType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.table.types.logical.VarCharType; + +import com.example.state.writer.job.schema.avro.AvroRecord; +import org.apache.avro.Schema; +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.List; +import java.util.function.Supplier; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Unit tests for {@link SerializerSnapshotToLogicalTypeConverter}. */ +class SerializerSnapshotToLogicalTypeConverterTest { + + // ------------------------------------------------------------------------- + // POJO classes + // ------------------------------------------------------------------------- + + /** Simple POJO used to exercise field-by-field POJO schema extraction. */ + public static class SimplePojo { + public String name; + public int age; + public long score; + public boolean active; + } + + /** POJO with a nested POJO field, used to exercise recursive schema extraction. */ + public static class NestedPojo { + public String label; + public SimplePojo inner; + } + + // ------------------------------------------------------------------------- + // Primitive / scalar snapshots + // ------------------------------------------------------------------------- + + @Test + void testPrimitives() { + List cases = + Arrays.asList( + new PrimitiveCase( + IntSerializer.IntSerializerSnapshot::new, LogicalTypeRoot.INTEGER), + new PrimitiveCase( + LongSerializer.LongSerializerSnapshot::new, LogicalTypeRoot.BIGINT), + new PrimitiveCase( + FloatSerializer.FloatSerializerSnapshot::new, + LogicalTypeRoot.FLOAT), + new PrimitiveCase( + DoubleSerializer.DoubleSerializerSnapshot::new, + LogicalTypeRoot.DOUBLE), + new PrimitiveCase( + BooleanSerializer.BooleanSerializerSnapshot::new, + LogicalTypeRoot.BOOLEAN), + new PrimitiveCase( + StringSerializer.StringSerializerSnapshot::new, + LogicalTypeRoot.VARCHAR)); + + for (PrimitiveCase c : cases) { + var snapshot = c.snapshotSupplier.get(); + LogicalType t = convert(snapshot); + assertThat(t.getTypeRoot()) + .as("wrong type root for %s", snapshot.getClass().getSimpleName()) + .isEqualTo(c.expectedRoot); + } + + // Numeric primitives are non-nullable at the wire level; spot-check one representative. + assertThat(convert(new IntSerializer.IntSerializerSnapshot()).isNullable()).isFalse(); + } + + private static final class PrimitiveCase { + final Supplier> + snapshotSupplier; + final LogicalTypeRoot expectedRoot; + + PrimitiveCase( + Supplier> + snapshotSupplier, + LogicalTypeRoot expectedRoot) { + this.snapshotSupplier = snapshotSupplier; + this.expectedRoot = expectedRoot; + } + } + + // ------------------------------------------------------------------------- + // Composite types + // ------------------------------------------------------------------------- + + @Test + void testListOfString() { + ListSerializer ser = new ListSerializer<>(StringSerializer.INSTANCE); + LogicalType t = convert(ser.snapshotConfiguration()); + assertThat(t.getTypeRoot()).isEqualTo(LogicalTypeRoot.ARRAY); + ArrayType at = (ArrayType) t; + assertThat(at.getElementType().getTypeRoot()).isEqualTo(LogicalTypeRoot.VARCHAR); + } + + @Test + void testMapStringToLong() { + MapSerializer ser = + new MapSerializer<>(StringSerializer.INSTANCE, LongSerializer.INSTANCE); + LogicalType t = convert(ser.snapshotConfiguration()); + assertThat(t.getTypeRoot()).isEqualTo(LogicalTypeRoot.MAP); + MapType mt = (MapType) t; + assertThat(mt.getKeyType().getTypeRoot()).isEqualTo(LogicalTypeRoot.VARCHAR); + assertThat(mt.getValueType().getTypeRoot()).isEqualTo(LogicalTypeRoot.BIGINT); + } + + @Test + void testNullableWrapsNestedTypeAsNullable() { + // LongSerializer alone always maps to a non-nullable BIGINT (see testPrimitives); wrapping + // it in NullableSerializer must flip only the nullability, not the underlying type. + TypeSerializer wrapped = NullableSerializer.wrap(LongSerializer.INSTANCE, true); + LogicalType t = convert(wrapped.snapshotConfiguration()); + assertThat(t.getTypeRoot()).isEqualTo(LogicalTypeRoot.BIGINT); + assertThat(t.isNullable()).isTrue(); + } + + @Test + @SuppressWarnings({"unchecked", "rawtypes"}) + void testTuple() { + TupleSerializer> ser = + new TupleSerializer<>( + (Class) Tuple2.class, + new TypeSerializer[] { + IntSerializer.INSTANCE, StringSerializer.INSTANCE + }); + LogicalType t = convert(ser.snapshotConfiguration()); + assertThat(t.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType rt = (RowType) t; + assertThat(rt.getFieldCount()).isEqualTo(2); + // Tuples have no field names in the serializer snapshot, so fields fall back to + // positional names, same as the RowData-without-field-names case. + assertField(rt, "f0", LogicalTypeRoot.INTEGER); + assertField(rt, "f1", LogicalTypeRoot.VARCHAR); + } + + // ------------------------------------------------------------------------- + // POJO types + // ------------------------------------------------------------------------- + + @Test + void testSimplePojo() { + LogicalType t = convertPojoType(SimplePojo.class); + assertThat(t.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType rt = (RowType) t; + assertThat(rt.getFieldCount()).isEqualTo(4); + + assertField(rt, "name", LogicalTypeRoot.VARCHAR); + assertField(rt, "age", LogicalTypeRoot.INTEGER); + assertField(rt, "score", LogicalTypeRoot.BIGINT); + assertField(rt, "active", LogicalTypeRoot.BOOLEAN); + } + + @Test + void testNestedPojo() { + LogicalType t = convertPojoType(NestedPojo.class); + assertThat(t.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType rt = (RowType) t; + assertThat(rt.getFieldCount()).isEqualTo(2); + assertField(rt, "label", LogicalTypeRoot.VARCHAR); + + // The nested 'inner' field should map to ROW + RowType.RowField innerField = findField(rt, "inner"); + assertThat(innerField).isNotNull(); + assertThat(innerField.getType().getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType innerRow = (RowType) innerField.getType(); + assertField(innerRow, "name", LogicalTypeRoot.VARCHAR); + assertField(innerRow, "age", LogicalTypeRoot.INTEGER); + } + + @Test + void testPojoFieldNamesPreservedWithoutClass() { + // Even without the POJO class on the classpath, field names should be available. + var snapshot = buildPojoSnapshot(SimplePojo.class); + // Field name extraction happens via the snapshot, no class needed + LogicalType t = SerializerSnapshotToLogicalTypeConverter.convert(snapshot); + RowType rt = (RowType) t; + List names = rt.getFieldNames(); + assertThat(names).contains("name", "age", "score", "active"); + } + + // ------------------------------------------------------------------------- + // Avro types + // ------------------------------------------------------------------------- + + @Test + void testAvroSpecificRecord() { + // AvroRecord has one field: longData (long) + LogicalType t = convertAvroSpecificType(AvroRecord.class); + assertThat(t.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType rt = (RowType) t; + assertThat(rt.getFieldCount()).isEqualTo(1); + assertField(rt, "longData", LogicalTypeRoot.BIGINT); + } + + @Test + void testAvroGenericRecord() { + // convertAvro() reads only the embedded writer Schema (see + // SerializerSnapshotToLogicalTypeConverter#convertAvro), so this also covers the + // "specific record class missing at read time" fallback: AvroSerializerSnapshot degrades + // to GenericRecord.class in that case, but the schema-derived RowType is identical either + // way. The actual missing-class read path is covered end-to-end by + // StateCatalogGeneratedSavepointITCase#testReadAvroKeyedStateFromSchemaDiscovery. + Schema schema = AvroRecord.getClassSchema(); + LogicalType t = convertAvroGenericType(schema); + assertThat(t.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType rt = (RowType) t; + assertThat(rt.getFieldCount()).isEqualTo(1); + assertField(rt, "longData", LogicalTypeRoot.BIGINT); + } + + // ------------------------------------------------------------------------- + // RowData types + // ------------------------------------------------------------------------- + + @Test + void testRowDataWithFieldNames() { + RowType rowType = + RowType.of( + new LogicalType[] {new IntType(), VarCharType.STRING_TYPE}, + new String[] {"id", "name"}); + RowDataSerializer serializer = new RowDataSerializer(rowType); + + LogicalType t = convert(serializer.snapshotConfiguration()); + assertThat(t.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType rt = (RowType) t; + assertThat(rt.getFieldCount()).isEqualTo(2); + assertField(rt, "id", LogicalTypeRoot.INTEGER); + assertField(rt, "name", LogicalTypeRoot.VARCHAR); + } + + @Test + void testRowDataWithoutFieldNamesFallsBackToPositional() { + // Many production call sites (e.g. window operators) build a RowDataSerializer from a + // bare LogicalType[], so no field names are available. The converter must still produce + // a usable RowType, falling back to positional names like the Tuple case. + RowDataSerializer serializer = new RowDataSerializer(new IntType(), new BigIntType()); + + LogicalType t = convert(serializer.snapshotConfiguration()); + assertThat(t.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType rt = (RowType) t; + assertThat(rt.getFieldCount()).isEqualTo(2); + assertField(rt, "f0", LogicalTypeRoot.INTEGER); + assertField(rt, "f1", LogicalTypeRoot.BIGINT); + } + + @Test + void testNestedRowData() { + RowType innerType = + RowType.of( + new LogicalType[] {new IntType(), VarCharType.STRING_TYPE}, + new String[] {"innerId", "innerName"}); + RowType outerType = + RowType.of( + new LogicalType[] {VarCharType.STRING_TYPE, innerType}, + new String[] {"label", "inner"}); + RowDataSerializer serializer = new RowDataSerializer(outerType); + + LogicalType t = convert(serializer.snapshotConfiguration()); + assertThat(t.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType rt = (RowType) t; + assertThat(rt.getFieldCount()).isEqualTo(2); + assertField(rt, "label", LogicalTypeRoot.VARCHAR); + + RowType.RowField innerField = findField(rt, "inner"); + assertThat(innerField).isNotNull(); + assertThat(innerField.getType().getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType innerRow = (RowType) innerField.getType(); + assertField(innerRow, "innerId", LogicalTypeRoot.INTEGER); + assertField(innerRow, "innerName", LogicalTypeRoot.VARCHAR); + } + + // ------------------------------------------------------------------------- + // Window namespace types + // ------------------------------------------------------------------------- + + @Test + void testTimeWindow() { + LogicalType t = + convert( + new org.apache.flink.streaming.api.windowing.windows.TimeWindow.Serializer() + .snapshotConfiguration()); + assertThat(t.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + assertThat(t.isNullable()).isFalse(); + RowType rt = (RowType) t; + assertThat(rt.getFieldCount()).isEqualTo(2); + assertField(rt, "window_start", LogicalTypeRoot.TIMESTAMP_WITHOUT_TIME_ZONE); + assertField(rt, "window_end", LogicalTypeRoot.TIMESTAMP_WITHOUT_TIME_ZONE); + } + + @Test + void testGlobalWindow() { + LogicalType t = + convert( + new org.apache.flink.streaming.api.windowing.windows.GlobalWindow + .Serializer() + .snapshotConfiguration()); + assertThat(t.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + assertThat(t.isNullable()).isFalse(); + RowType rt = (RowType) t; + assertThat(rt.getFieldCount()).isEqualTo(0); + } + + // ------------------------------------------------------------------------- + // Null / unknown snapshot + // ------------------------------------------------------------------------- + + @Test + void testNullSnapshot() { + LogicalType t = SerializerSnapshotToLogicalTypeConverter.convert(null); + assertThat(t.getTypeRoot()).isEqualTo(LogicalTypeRoot.VARBINARY); + } + + @Test + void testVoidNamespaceSnapshotUnsupported() { + // VoidNamespace is filtered out by StateTableUtils before ever reaching the converter + // (plain per-key state has no namespace to convert); confirm it stays unsupported here. + assertThatThrownBy( + () -> + convert( + new org.apache.flink.runtime.state.VoidNamespaceSerializer() + .snapshotConfiguration())) + .isInstanceOf(UnsupportedOperationException.class); + } + + // ------------------------------------------------------------------------- + // Helpers + // ------------------------------------------------------------------------- + + private static LogicalType convert( + org.apache.flink.api.common.typeutils.TypeSerializerSnapshot snapshot) { + return SerializerSnapshotToLogicalTypeConverter.convert(snapshot); + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + private static LogicalType convertPojoType(Class pojoClass) { + var snapshot = buildPojoSnapshot(pojoClass); + return SerializerSnapshotToLogicalTypeConverter.convert(snapshot); + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + private static org.apache.flink.api.common.typeutils.TypeSerializerSnapshot + buildPojoSnapshot(Class clazz) { + return TypeExtractor.createTypeInfo(clazz) + .createSerializer(new SerializerConfigImpl()) + .snapshotConfiguration(); + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + private static + LogicalType convertAvroSpecificType(Class avroClass) { + return SerializerSnapshotToLogicalTypeConverter.convert( + new AvroTypeInfo<>(avroClass) + .createSerializer(new SerializerConfigImpl()) + .snapshotConfiguration()); + } + + private static LogicalType convertAvroGenericType(Schema schema) { + return SerializerSnapshotToLogicalTypeConverter.convert( + new GenericRecordAvroTypeInfo(schema) + .createSerializer(new SerializerConfigImpl()) + .snapshotConfiguration()); + } + + private static void assertField(RowType row, String name, LogicalTypeRoot expectedRoot) { + RowType.RowField field = findField(row, name); + assertThat(field).as("Field '%s' not found in row type", name).isNotNull(); + assertThat(field.getType().getTypeRoot()) + .as("Wrong type for field '%s'", name) + .isEqualTo(expectedRoot); + } + + private static RowType.RowField findField(RowType row, String name) { + return row.getFields().stream() + .filter(f -> f.getName().equals(name)) + .findFirst() + .orElse(null); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/StateSchemaExtractorTest.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/StateSchemaExtractorTest.java new file mode 100644 index 00000000000000..7849dcdd6721dc --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/StateSchemaExtractorTest.java @@ -0,0 +1,204 @@ +/* + * 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.schema; + +import org.apache.flink.api.common.state.StateDescriptor; +import org.apache.flink.api.common.typeutils.base.IntSerializer; +import org.apache.flink.api.common.typeutils.base.LongSerializer; +import org.apache.flink.api.common.typeutils.base.StringSerializer; +import org.apache.flink.core.memory.DataInputDeserializer; +import org.apache.flink.core.memory.DataOutputSerializer; +import org.apache.flink.runtime.state.KeyedBackendSerializationProxy; +import org.apache.flink.runtime.state.VoidNamespaceSerializer; +import org.apache.flink.runtime.state.metainfo.StateMetaInfoSnapshot; + +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Unit tests for {@link StateSchemaExtractor}. */ +class StateSchemaExtractorTest { + + // ------------------------------------------------------------------------- + // Tests + // ------------------------------------------------------------------------- + + @Test + void testExtractSingleValueState() throws IOException { + StateMetaInfoSnapshot stateSnap = + buildValueStateSnapshot("my-state", StateDescriptor.Type.VALUE); + + KeyedBackendSerializationProxy proxy = + new KeyedBackendSerializationProxy<>( + IntSerializer.INSTANCE, Collections.singletonList(stateSnap), false); + + List result = roundTrip(proxy); + + assertThat(result).hasSize(1); + StateSchemaInfo info = result.get(0); + assertThat(info.stateName).isEqualTo("my-state"); + assertThat(info.stateKind).isEqualTo(StateDescriptor.Type.VALUE); + assertThat(info.valueSnapshot).isInstanceOf(IntSerializer.IntSerializerSnapshot.class); + assertThat(info.keySnapshot).isInstanceOf(IntSerializer.IntSerializerSnapshot.class); + assertThat(info.mapKeySnapshot).isNull(); + } + + @Test + void testExtractMultipleStates() throws IOException { + StateMetaInfoSnapshot valueState = + buildValueStateSnapshot("int-state", StateDescriptor.Type.VALUE); + StateMetaInfoSnapshot stringState = + buildValueStateSnapshotWithStringValue("str-state", StateDescriptor.Type.VALUE); + + KeyedBackendSerializationProxy proxy = + new KeyedBackendSerializationProxy<>( + IntSerializer.INSTANCE, Arrays.asList(valueState, stringState), false); + + List result = roundTrip(proxy); + + assertThat(result).hasSize(2); + + assertThat(result.get(0).stateName).isEqualTo("int-state"); + assertThat(result.get(0).valueSnapshot) + .isInstanceOf(IntSerializer.IntSerializerSnapshot.class); + + assertThat(result.get(1).stateName).isEqualTo("str-state"); + assertThat(result.get(1).valueSnapshot) + .isInstanceOf(StringSerializer.StringSerializerSnapshot.class); + } + + @Test + void testExtractMapState() throws IOException { + Map options = new HashMap<>(); + options.put( + StateMetaInfoSnapshot.CommonOptionsKeys.KEYED_STATE_TYPE.toString(), + StateDescriptor.Type.MAP.toString()); + + Map> + serializerSnapshots = new LinkedHashMap<>(); + serializerSnapshots.put( + StateMetaInfoSnapshot.CommonSerializerKeys.NAMESPACE_SERIALIZER.toString(), + new VoidNamespaceSerializer.VoidNamespaceSerializerSnapshot()); + serializerSnapshots.put( + StateMetaInfoSnapshot.CommonSerializerKeys.USER_KEY_SERIALIZER.toString(), + new StringSerializer.StringSerializerSnapshot()); + serializerSnapshots.put( + StateMetaInfoSnapshot.CommonSerializerKeys.VALUE_SERIALIZER.toString(), + new LongSerializer.LongSerializerSnapshot()); + + StateMetaInfoSnapshot mapSnap = + new StateMetaInfoSnapshot( + "map-state", + StateMetaInfoSnapshot.BackendStateType.KEY_VALUE, + options, + serializerSnapshots); + + KeyedBackendSerializationProxy proxy = + new KeyedBackendSerializationProxy<>( + IntSerializer.INSTANCE, Collections.singletonList(mapSnap), false); + + List result = roundTrip(proxy); + + assertThat(result).hasSize(1); + StateSchemaInfo info = result.get(0); + assertThat(info.stateName).isEqualTo("map-state"); + assertThat(info.stateKind).isEqualTo(StateDescriptor.Type.MAP); + assertThat(info.mapKeySnapshot) + .isInstanceOf(StringSerializer.StringSerializerSnapshot.class); + assertThat(info.valueSnapshot).isInstanceOf(LongSerializer.LongSerializerSnapshot.class); + } + + @Test + void testEmptyStates() throws IOException { + KeyedBackendSerializationProxy proxy = + new KeyedBackendSerializationProxy<>( + IntSerializer.INSTANCE, Collections.emptyList(), false); + + List result = roundTrip(proxy); + + assertThat(result).isNotNull().isEmpty(); + } + + // ------------------------------------------------------------------------- + // Helpers + // ------------------------------------------------------------------------- + + private static StateMetaInfoSnapshot buildValueStateSnapshot( + String name, StateDescriptor.Type type) { + Map options = new HashMap<>(); + options.put( + StateMetaInfoSnapshot.CommonOptionsKeys.KEYED_STATE_TYPE.toString(), + type.toString()); + + Map> + serializerSnapshots = new LinkedHashMap<>(); + serializerSnapshots.put( + StateMetaInfoSnapshot.CommonSerializerKeys.NAMESPACE_SERIALIZER.toString(), + new VoidNamespaceSerializer.VoidNamespaceSerializerSnapshot()); + serializerSnapshots.put( + StateMetaInfoSnapshot.CommonSerializerKeys.VALUE_SERIALIZER.toString(), + new IntSerializer.IntSerializerSnapshot()); + + return new StateMetaInfoSnapshot( + name, + StateMetaInfoSnapshot.BackendStateType.KEY_VALUE, + options, + serializerSnapshots); + } + + private static StateMetaInfoSnapshot buildValueStateSnapshotWithStringValue( + String name, StateDescriptor.Type type) { + Map options = new HashMap<>(); + options.put( + StateMetaInfoSnapshot.CommonOptionsKeys.KEYED_STATE_TYPE.toString(), + type.toString()); + + Map> + serializerSnapshots = new LinkedHashMap<>(); + serializerSnapshots.put( + StateMetaInfoSnapshot.CommonSerializerKeys.NAMESPACE_SERIALIZER.toString(), + new VoidNamespaceSerializer.VoidNamespaceSerializerSnapshot()); + serializerSnapshots.put( + StateMetaInfoSnapshot.CommonSerializerKeys.VALUE_SERIALIZER.toString(), + new StringSerializer.StringSerializerSnapshot()); + + return new StateMetaInfoSnapshot( + name, + StateMetaInfoSnapshot.BackendStateType.KEY_VALUE, + options, + serializerSnapshots); + } + + private static List roundTrip(KeyedBackendSerializationProxy proxy) + throws IOException { + DataOutputSerializer out = new DataOutputSerializer(256); + proxy.write(out); + + DataInputDeserializer in = new DataInputDeserializer(out.getSharedBuffer()); + return StateSchemaExtractor.extractSchema(in); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/SnapshotDiscoveryTest.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/SnapshotDiscoveryTest.java new file mode 100644 index 00000000000000..fc0e3bd15892de --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/SnapshotDiscoveryTest.java @@ -0,0 +1,248 @@ +/* + * 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.catalog; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.attribute.FileTime; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Unit tests for {@link SnapshotDiscovery}. */ +class SnapshotDiscoveryTest { + + // Every metadata file created by createMetadataFile() gets this exact modification time, so + // the creationTs segment in derived database names is deterministic across all tests. + private static final Instant FIXED_TS = Instant.parse("2024-03-15T10:30:45Z"); + private static final String TS = "2024-03-15T10:30:45Z"; + + private static final String MISSING_DIR = "/nonexistent-snapshot-discovery-test-path"; + + @TempDir Path tempDir; + + private final List started = new ArrayList<>(); + + private SnapshotDiscovery discovery; + + @BeforeEach + void setUp() { + discovery = start(Collections.singletonMap("app", tempDir.toString()), true); + } + + @AfterEach + void tearDown() { + started.forEach(SnapshotDiscovery::stop); + } + + // ------------------------------------------------------------------------- + // Construction-time validation + // ------------------------------------------------------------------------- + + @Test + void testConstructionValidation() { + assertThatThrownBy(() -> new SnapshotDiscovery(Collections.emptyMap(), 2, true)) + .as("no directory configured") + .isInstanceOf(IllegalArgumentException.class); + + assertThatThrownBy(() -> new SnapshotDiscovery(dirs("/state/app", "/state/app"), 2, true)) + .as("same directory under two labels") + .isInstanceOf(IllegalArgumentException.class); + + assertThatThrownBy(() -> new SnapshotDiscovery(dirs("/state", "/state/app"), 2, true)) + .as("one directory nested inside the other") + .isInstanceOf(IllegalArgumentException.class); + } + + // ------------------------------------------------------------------------- + // find() + // ------------------------------------------------------------------------- + + @Test + void testFindRejectsInvalidDatabaseNames() { + assertThat(discovery.find(null)).isEmpty(); + assertThat(discovery.find("")).isEmpty(); + assertThat(discovery.find(" ")).isEmpty(); + assertThat(discovery.find("/savepoint-abc")).isEmpty(); + assertThat(discovery.find("unknown/" + TS + "/savepoint-abc")).isEmpty(); + + // With db-name.include-ts enabled (the default), a name with no '/' has no room for the + // mandatory creationTs segment, so it can never match. + assertThat(discovery.find("app")).isEmpty(); + assertThat(discovery.find("savepoint-abc")).isEmpty(); + } + + @Test + void testFindResolvesRelativePathVerbatim() throws IOException { + createMetadataFile(tempDir.resolve("savepoint-abc")); + createMetadataFile(tempDir.resolve("jobId").resolve("chk-3")); + createMetadataFile(tempDir.resolve("a").resolve("b").resolve("c")); + + assertThat(discovery.find("app/" + TS + "/savepoint-abc")) + .hasValue(tempDir.resolve("savepoint-abc").toString()); + assertThat(discovery.find("app/" + TS + "/jobId/chk-3")) + .hasValue(tempDir.resolve("jobId").resolve("chk-3").toString()); + assertThat(discovery.find("app/" + TS + "/a/b/c")) + .hasValue(tempDir.resolve("a").resolve("b").resolve("c").toString()); + + assertThat(discovery.find("app/" + TS + "/savepoint-nonexistent")).isEmpty(); + } + + @Test + void testFindSnapshotWithTrailingSlash() throws IOException { + createMetadataFile(tempDir.resolve("savepoint-abc")); + + // trailing slash after the creationTs → relativePath is empty, same as a ts-only path + assertThat(discovery.find("app/" + TS + "/")).isEmpty(); + } + + @Test + void testFindTsOnlyPathMatchesSnapshotDirectlyUnderConfiguredDir() throws IOException { + createMetadataFile(tempDir); + + assertThat(discovery.find("app/" + TS)).hasValue(tempDir.toString()); + } + + @Test + void testFindWithTsDisabled() throws IOException { + SnapshotDiscovery noTs = start(Collections.singletonMap("app", tempDir.toString()), false); + createMetadataFile(tempDir.resolve("savepoint-abc")); + + assertThat(noTs.find("app/savepoint-abc")) + .hasValue(tempDir.resolve("savepoint-abc").toString()); + // With db-name.include-ts disabled, everything after the label is taken verbatim as the + // relative path — no segment is skipped as a timestamp. + assertThat(noTs.find("app/extra-segment/savepoint-abc")).isEmpty(); + } + + // ------------------------------------------------------------------------- + // list() + // ------------------------------------------------------------------------- + + @Test + void testListReflectsFilesystemChangesWithoutCaching() throws IOException { + assertThat(discovery.list()).isEmpty(); + + createMetadataFile(tempDir.resolve("savepoint-new")); + assertThat(discovery.list()).containsExactly("app/" + TS + "/savepoint-new"); + + Files.delete(tempDir.resolve("savepoint-new").resolve("_metadata")); + assertThat(discovery.list()).isEmpty(); + } + + @Test + void testListMultipleSnapshots() throws IOException { + createMetadataFile(tempDir.resolve("savepoint-a")); + createMetadataFile(tempDir.resolve("savepoint-b")); + createMetadataFile(tempDir.resolve("jobId").resolve("chk-1")); + + assertThat(discovery.list()) + .containsExactlyInAnyOrder( + "app/" + TS + "/savepoint-a", + "app/" + TS + "/savepoint-b", + "app/" + TS + "/jobId/chk-1"); + } + + @Test + void testListNonMetadataFilesIgnored() throws IOException { + Files.createDirectories(tempDir.resolve("savepoint-a")); + Files.createFile(tempDir.resolve("savepoint-a").resolve("other.file")); + + assertThat(discovery.list()).isEmpty(); + } + + @Test + void testListReturnsSnapshotsFromHealthyDirectoryWhenOtherFails() throws IOException { + createMetadataFile(tempDir.resolve("savepoint-ok")); + + Map labelToDir = new LinkedHashMap<>(); + labelToDir.put("good", tempDir.toString()); + labelToDir.put("bad", MISSING_DIR); + + assertThat(start(labelToDir, true).list()).containsExactly("good/" + TS + "/savepoint-ok"); + } + + @Test + void testListThrowsWhenAllDirectoriesFail() { + SnapshotDiscovery allBad = start(Collections.singletonMap("bad", MISSING_DIR), true); + + assertThatThrownBy(allBad::list) + .isInstanceOf(IOException.class) + .hasMessageContaining("All configured directories failed") + .cause() + .isInstanceOf(IOException.class) + .hasMessageContaining("All directory listings failed"); + } + + @Test + void testListWithTsDisabled() throws IOException { + SnapshotDiscovery noTs = start(Collections.singletonMap("app", tempDir.toString()), false); + createMetadataFile(tempDir.resolve("savepoint-a")); + createMetadataFile(tempDir.resolve("jobId").resolve("chk-1")); + + assertThat(noTs.list()).containsExactlyInAnyOrder("app/savepoint-a", "app/jobId/chk-1"); + } + + @Test + void testDbNameCreationTsMatchesExpectedFormat() throws IOException { + // Uses the real (unset) modification time to verify the formatter itself, rather than the + // fixed FIXED_TS used elsewhere in this file. + Files.createDirectories(tempDir.resolve("savepoint-live")); + Files.createFile(tempDir.resolve("savepoint-live").resolve("_metadata")); + + assertThat(discovery.list().get(0)) + .matches("app/\\d{4}-\\d{2}-\\d{2}T\\d{2}:\\d{2}:\\d{2}Z/savepoint-live"); + } + + // ------------------------------------------------------------------------- + // Helpers + // ------------------------------------------------------------------------- + + private SnapshotDiscovery start(Map labelToDir, boolean dbNameIncludeTs) { + SnapshotDiscovery snapshotDiscovery = new SnapshotDiscovery(labelToDir, 2, dbNameIncludeTs); + snapshotDiscovery.start(); + started.add(snapshotDiscovery); + return snapshotDiscovery; + } + + private static Map dirs(String firstDir, String secondDir) { + Map labelToDir = new LinkedHashMap<>(); + labelToDir.put("a", firstDir); + labelToDir.put("b", secondDir); + return labelToDir; + } + + private static void createMetadataFile(Path snapshotDir) throws IOException { + Files.createDirectories(snapshotDir); + Path file = Files.createFile(snapshotDir.resolve("_metadata")); + Files.setLastModifiedTime(file, FileTime.from(FIXED_TS)); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/StateCatalogDiscoveryITCase.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/StateCatalogDiscoveryITCase.java new file mode 100644 index 00000000000000..a392bea2806374 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/StateCatalogDiscoveryITCase.java @@ -0,0 +1,210 @@ +/* + * 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.catalog; + +import org.apache.flink.runtime.checkpoint.Checkpoints; +import org.apache.flink.runtime.checkpoint.OperatorState; +import org.apache.flink.runtime.checkpoint.metadata.CheckpointMetadata; +import org.apache.flink.runtime.jobgraph.OperatorID; +import org.apache.flink.state.table.module.StateModule; +import org.apache.flink.table.api.EnvironmentSettings; +import org.apache.flink.table.api.TableEnvironment; +import org.apache.flink.table.api.TableResult; +import org.apache.flink.table.catalog.CatalogView; +import org.apache.flink.table.catalog.ObjectPath; +import org.apache.flink.types.Row; +import org.apache.flink.util.CloseableIterator; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.OutputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Integration tests for basic {@link StateCatalog} functionality driven through {@code CREATE + * CATALOG} DDL and SQL: multi-label discovery and the {@code metadata} view. Checkpoint metadata is + * written directly via {@link Checkpoints#storeCheckpointMetadataWithoutExclusiveDir} — no + * minicluster or real state backend is involved, so these tests are backend-agnostic by + * construction. + * + *

For reads of real (generated) keyed-state savepoints, see {@code + * StateCatalogGeneratedSavepointITCase} (HashMap-only, checked-in fixtures) and {@code + * KeyedStateReadingITCase} (parameterized across backends, savepoints taken at runtime). + */ +class StateCatalogDiscoveryITCase { + + @Test + void testMetadataQueryReturnsOperators(@TempDir Path tempDir) throws Exception { + OperatorID opId1 = new OperatorID(1, 2); + OperatorState op1 = new OperatorState("source", "source-uid", opId1, 2, 128); + OperatorID opId2 = new OperatorID(3, 4); + OperatorState op2 = new OperatorState("sink", null, opId2, 1, 128); + + Path savepointDir = Files.createDirectories(tempDir.resolve("savepoint-test")); + writeMetadata(savepointDir, 42L, Arrays.asList(op1, op2)); + + TableEnvironment tableEnv = newTableEnv(); + createCatalog(tableEnv, "state", directoryOption("app", tempDir)); + tableEnv.executeSql("USE CATALOG state"); + + StateCatalog catalog = getCatalog(tableEnv, "state"); + String dbName = catalog.listDatabases().get(0); + tableEnv.executeSql("USE `" + dbName + "`"); + + List rows = collectWithSql(tableEnv, "SELECT * FROM metadata"); + + assertThat(rows).hasSize(2); + rows.forEach(row -> assertThat(row.getField("checkpoint-id")).isEqualTo(42L)); + assertThat(rows.stream().map(r -> r.getField("operator-name")).collect(Collectors.toList())) + .containsExactlyInAnyOrder("source", "sink"); + + catalog.close(); + } + + @Test + void testMultipleLabelsDiscovered(@TempDir Path tempDir) throws Exception { + Path checkpointsDir = Files.createDirectories(tempDir.resolve("checkpoints")); + Path savepointsDir = Files.createDirectories(tempDir.resolve("savepoints")); + touchMetadata(checkpointsDir.resolve("savepoint-a")); + touchMetadata(savepointsDir.resolve("savepoint-b")); + + String directoryOptions = + directoryOption("ckpts", checkpointsDir) + + ", " + + directoryOption("svpts", savepointsDir); + + TableEnvironment tableEnv = newTableEnv(); + createCatalog(tableEnv, "with_ts", directoryOptions); + createCatalog( + tableEnv, "without_ts", directoryOptions + ", 'db-name.include-ts' = 'false'"); + + StateCatalog withTs = getCatalog(tableEnv, "with_ts"); + assertThat(withTs.listDatabases()) + .hasSize(2) + .allSatisfy( + dbName -> + assertThat(dbName) + .matches( + "(ckpts|svpts)/\\d{4}-\\d{2}-\\d{2}T\\d{2}:\\d{2}:\\d{2}Z/savepoint-[ab]")); + withTs.close(); + + StateCatalog withoutTs = getCatalog(tableEnv, "without_ts"); + assertThat(withoutTs.listDatabases()) + .containsExactlyInAnyOrder("ckpts/savepoint-a", "svpts/savepoint-b"); + withoutTs.close(); + } + + @Test + void testCatalogOperations(@TempDir Path tempDir) throws Exception { + Path savepointDir = Files.createDirectories(tempDir.resolve("savepoint-abc")); + writeMetadata(savepointDir, 7L, Collections.emptyList()); + + TableEnvironment tableEnv = newTableEnv(); + createCatalog(tableEnv, "state", directoryOption("app", tempDir)); + + StateCatalog catalog = getCatalog(tableEnv, "state"); + List dbs = catalog.listDatabases(); + assertThat(dbs).hasSize(1); + String dbName = dbs.get(0); + assertThat(dbName).matches("app/\\d{4}-\\d{2}-\\d{2}T\\d{2}:\\d{2}:\\d{2}Z/savepoint-abc"); + + assertThat(catalog.databaseExists(dbName)).isTrue(); + assertThat(catalog.databaseExists("app/nonexistent")).isFalse(); + + // listTables includes views (the "metadata" view), per the Catalog contract. + assertThat(catalog.listTables(dbName)).containsExactly(StateCatalog.METADATA_TABLE); + assertThat(catalog.listViews(dbName)).containsExactly(StateCatalog.METADATA_TABLE); + + assertThat(catalog.tableExists(new ObjectPath(dbName, StateCatalog.METADATA_TABLE))) + .isTrue(); + assertThat(catalog.tableExists(new ObjectPath(dbName, "other"))).isFalse(); + + CatalogView view = + (CatalogView) catalog.getTable(new ObjectPath(dbName, StateCatalog.METADATA_TABLE)); + assertThat(view.getOriginalQuery()) + .contains("savepoint_metadata") + .contains(savepointDir.toAbsolutePath().toString()); + + // Verify querying the metadata view via SQL works and returns expected rows + tableEnv.executeSql("USE CATALOG state"); + tableEnv.executeSql("USE `" + dbName + "`"); + assertThat(collectWithSql(tableEnv, "SELECT * FROM metadata")).isEmpty(); + + catalog.close(); + } + + // ------------------------------------------------------------------------- + // Helpers + // ------------------------------------------------------------------------- + + private static TableEnvironment newTableEnv() { + TableEnvironment env = TableEnvironment.create(EnvironmentSettings.inBatchMode()); + env.loadModule("state", StateModule.INSTANCE); + return env; + } + + private static void createCatalog( + TableEnvironment tableEnv, String catalogName, String withOptions) { + tableEnv.executeSql( + String.format( + "CREATE CATALOG %s WITH ('type' = '%s', %s)", + catalogName, StateCatalogFactory.IDENTIFIER, withOptions)); + } + + private static String directoryOption(String label, Path dir) { + return String.format( + "'directory.%s' = '%s'", label, dir.toAbsolutePath().toString().replace("'", "''")); + } + + private static StateCatalog getCatalog(TableEnvironment tableEnv, String name) { + return (StateCatalog) tableEnv.getCatalog(name).get(); + } + + private static void touchMetadata(Path snapshotDir) throws Exception { + Files.createDirectories(snapshotDir); + Files.createFile(snapshotDir.resolve("_metadata")); + } + + private static void writeMetadata( + Path snapshotDir, long checkpointId, List operators) throws Exception { + CheckpointMetadata metadata = + new CheckpointMetadata(checkpointId, operators, Collections.emptyList()); + try (OutputStream out = Files.newOutputStream(snapshotDir.resolve("_metadata"))) { + Checkpoints.storeCheckpointMetadataWithoutExclusiveDir(metadata, out); + } + } + + private static List collectWithSql(TableEnvironment tEnv, String sql) throws Exception { + List rows = new ArrayList<>(); + TableResult result = tEnv.executeSql(sql); + try (CloseableIterator it = result.collect()) { + it.forEachRemaining(rows::add); + } + return rows; + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/StateCatalogGeneratedSavepointITCase.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/StateCatalogGeneratedSavepointITCase.java new file mode 100644 index 00000000000000..c56fd71365a3be --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/StateCatalogGeneratedSavepointITCase.java @@ -0,0 +1,691 @@ +/* + * 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.catalog; + +import org.apache.flink.api.common.RuntimeExecutionMode; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.runtime.checkpoint.metadata.CheckpointMetadata; +import org.apache.flink.state.api.OperatorIdentifier; +import org.apache.flink.state.api.StateTableUtils; +import org.apache.flink.state.api.runtime.SavepointLoader; +import org.apache.flink.state.api.schema.KeyedStateSchemaInfo; +import org.apache.flink.state.table.module.StateModule; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.table.api.EnvironmentSettings; +import org.apache.flink.table.api.Schema; +import org.apache.flink.table.api.TableEnvironment; +import org.apache.flink.table.api.TableResult; +import org.apache.flink.table.api.bridge.java.StreamTableEnvironment; +import org.apache.flink.table.api.internal.TableEnvironmentImpl; +import org.apache.flink.table.catalog.CatalogManager; +import org.apache.flink.table.catalog.CatalogTable; +import org.apache.flink.table.catalog.ObjectPath; +import org.apache.flink.table.catalog.UnresolvedIdentifier; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.LogicalTypeRoot; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.types.Row; +import org.apache.flink.util.CloseableIterator; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import java.util.stream.LongStream; + +import static org.apache.flink.configuration.ExecutionOptions.RUNTIME_MODE; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Integration tests for {@link StateCatalog} and {@link StateTableUtils} against real keyed-state + * savepoints that are checked in as test resources, grouped by fixture/scenario in {@code @Nested} + * classes. + * + *

These savepoints were produced once with the {@code hashmap} state backend by the (disabled, + * manually-run) generator programs under {@code src/test/resources/generator} and are checked in + * under {@code src/test/resources/}. They cannot be regenerated for the RocksDB state backend + * without running those generators locally against a RocksDB-configured job, so every test in this + * class is inherently HashMap-only — see {@code KeyedStateReadingITCase} for the + * RocksDB-parameterized equivalent exercised against savepoints taken at runtime instead of + * checked-in fixtures. + */ +class StateCatalogGeneratedSavepointITCase { + + /** + * Schema discovery and reads for a savepoint whose POJO/Avro classes aren't on the classpath. + */ + @Nested + class SchemaDiscoveryWithoutSourceClasses { + + private static final String STATE_PATH = "src/test/resources/table-state-missing-class"; + private static final String OPERATOR_UID = "missing-class-operator"; + private static final String AVRO_STATE_PATH = "src/test/resources/table-state-missing-avro"; + private static final String AVRO_OPERATOR_UID = "missing-avro-operator"; + private final String[] avroStateNames = {"KeyedAvroSpecificValue", "KeyedAvroGenericValue"}; + private static final int NUM_KEYS = 10; + + @Test + @SuppressWarnings("unchecked") + void testReadKeyedStateFromSchemaDiscovery() throws Exception { + List result = readViaTemporaryTable(STATE_PATH, OPERATOR_UID, "state_table"); + + assertThat(result).hasSize(NUM_KEYS); + assertThat(stateKeys(result)).containsExactlyInAnyOrderElementsOf(longRange(NUM_KEYS)); + + Set primitiveValues = + result.stream() + .map(r -> (Long) r.getField("KeyedPrimitiveValue")) + .collect(Collectors.toSet()); + assertThat(primitiveValues).containsExactly(1L); + + Set pojoValues = + result.stream() + .map(r -> (Row) r.getField("KeyedPojoValue")) + .collect(Collectors.toSet()); + assertThat(pojoValues).hasSize(1); + Row pojoRow = pojoValues.iterator().next(); + assertThat(pojoRow.getField("privateLong")).isEqualTo(1L); + assertThat(pojoRow.getField("publicLong")).isEqualTo(1L); + + // Each key holds the single-element list [state_key] and the single map entry + // {state_key: state_key}. + for (Row row : result) { + Long key = (Long) row.getField("state_key"); + assertThat((Long[]) row.getField("KeyedPrimitiveValueList")).containsExactly(key); + assertThat((Map) row.getField("KeyedPrimitiveValueMap")) + .containsExactly(Map.entry(key, key)); + } + } + + @Test + void testFlattenedKeyedStateTables() throws Exception { + StateCatalog catalog = openCatalogOn(STATE_PATH); + try { + String dbName = catalog.listDatabases().get(0); + String listTable = flatKeyedTable(OPERATOR_UID, "KeyedPrimitiveValueList"); + String mapTable = flatKeyedTable(OPERATOR_UID, "KeyedPrimitiveValueMap"); + + // (a) the flattened tables exist and expose a composite primary key + assertThat(catalog.listTables(dbName)).contains(listTable, mapTable); + assertThat(catalog.tableExists(new ObjectPath(dbName, listTable))).isTrue(); + assertThat(catalog.tableExists(new ObjectPath(dbName, mapTable))).isTrue(); + assertThat(catalog.tableExists(new ObjectPath(dbName, listTable + "-nonexistent"))) + .isFalse(); + + Schema listSchema = schemaOf(catalog, dbName, listTable); + assertThat(columnNames(listSchema)) + .containsExactly("state_key", "list_index", "list_value"); + assertThat(listSchema.getPrimaryKey()).isPresent(); + assertThat(listSchema.getPrimaryKey().get().getColumnNames()) + .containsExactly("state_key", "list_index"); + + Schema mapSchema = schemaOf(catalog, dbName, mapTable); + assertThat(columnNames(mapSchema)) + .containsExactly("state_key", "map_key", "map_value"); + assertThat(mapSchema.getPrimaryKey()).isPresent(); + assertThat(mapSchema.getPrimaryKey().get().getColumnNames()) + .containsExactly("state_key", "map_key"); + + // (b) the flattened tables can be read correctly and return the expected data + TableEnvironment tableEnv = newCatalogTableEnv(catalog, dbName); + + // KeyedPrimitiveValueList holds a single-element list [state_key] per key. + List listRows = collectWithSql(tableEnv, "SELECT * FROM `" + listTable + "`"); + assertThat(listRows).hasSize(NUM_KEYS); + assertThat(stateKeys(listRows)) + .containsExactlyInAnyOrderElementsOf(longRange(NUM_KEYS)); + for (Row row : listRows) { + assertThat(row.getField("list_index")).isEqualTo(0L); + assertThat(row.getField("list_value")).isEqualTo(row.getField("state_key")); + } + + // KeyedPrimitiveValueMap holds a single entry {state_key: state_key} per key. + List mapRows = collectWithSql(tableEnv, "SELECT * FROM `" + mapTable + "`"); + assertThat(mapRows).hasSize(NUM_KEYS); + assertThat(stateKeys(mapRows)) + .containsExactlyInAnyOrderElementsOf(longRange(NUM_KEYS)); + for (Row row : mapRows) { + assertThat(row.getField("map_key")).isEqualTo(row.getField("state_key")); + assertThat(row.getField("map_value")).isEqualTo(row.getField("state_key")); + } + + // state_key filter push-down (SupportsFilterPushDown) prunes to a single key even + // though state_key is only part of the composite (state_key, list_index/map_key) + // primary key in the flattened schema. + for (String table : new String[] {listTable, mapTable}) { + List filtered = + collectWithSql( + tableEnv, "SELECT * FROM `" + table + "` WHERE state_key = 3"); + assertThat(filtered).hasSize(1); + assertThat(filtered.get(0).getField("state_key")).isEqualTo(3L); + } + } finally { + catalog.close(); + } + } + + @Test + void testSchemaExtractionWithoutPojoClass() throws Exception { + CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(STATE_PATH); + KeyedStateSchemaInfo schemaInfo = + StateTableUtils.getKeyedStateSchema( + metadata, OperatorIdentifier.forUid(OPERATOR_UID)); + + assertThat(schemaInfo.keyType.getTypeRoot()).isEqualTo(LogicalTypeRoot.BIGINT); + assertThat(stateTypeRoot(schemaInfo, "KeyedPrimitiveValue")) + .isEqualTo(LogicalTypeRoot.BIGINT); + assertThat(stateTypeRoot(schemaInfo, "KeyedPrimitiveValueList")) + .isEqualTo(LogicalTypeRoot.ARRAY); + assertThat(stateTypeRoot(schemaInfo, "KeyedPrimitiveValueMap")) + .isEqualTo(LogicalTypeRoot.MAP); + + assertThat(stateTypeRoot(schemaInfo, "KeyedPojoValue")).isEqualTo(LogicalTypeRoot.ROW); + RowType pojoRowType = + (RowType) schemaInfo.stateSchemas.get("KeyedPojoValue").logicalType; + assertThat(pojoRowType.getFieldNames()).contains("privateLong", "publicLong"); + assertThat(fieldTypeRoot(pojoRowType, "privateLong")).isEqualTo(LogicalTypeRoot.BIGINT); + assertThat(fieldTypeRoot(pojoRowType, "publicLong")).isEqualTo(LogicalTypeRoot.BIGINT); + } + + @Test + void testReadAvroKeyedStateFromSchemaDiscovery() throws Exception { + CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(AVRO_STATE_PATH); + KeyedStateSchemaInfo schemaInfo = + StateTableUtils.getKeyedStateSchema( + metadata, OperatorIdentifier.forUid(AVRO_OPERATOR_UID)); + + // Both the specific-record and the generic-record state degrade to ROW(longData). + for (String stateName : avroStateNames) { + assertThat(stateTypeRoot(schemaInfo, stateName)).isEqualTo(LogicalTypeRoot.ROW); + RowType rowType = (RowType) schemaInfo.stateSchemas.get(stateName).logicalType; + assertThat(rowType.getFieldNames()).containsExactly("longData"); + } + + List result = + readViaTemporaryTable(AVRO_STATE_PATH, AVRO_OPERATOR_UID, "avro_state_table"); + + assertThat(result).hasSize(NUM_KEYS); + assertThat(stateKeys(result)).containsExactlyInAnyOrderElementsOf(longRange(NUM_KEYS)); + for (String stateName : avroStateNames) { + Set values = + result.stream() + .map(r -> (Row) r.getField(stateName)) + .collect(Collectors.toSet()); + assertThat(values).hasSize(1); + assertThat(values.iterator().next().getField("longData")).isEqualTo(1L); + } + } + } + + /** Schema and data for four operator types (primitive, POJO, Avro-specific, Avro-generic). */ + @Nested + class MultiOperatorTypeCatalog { + + private static final String RESOURCES_DIR = "src/test/resources/keyed-state-catalog"; + private static final String UID_PRIMITIVE = "primitive-state-op"; + private static final String UID_POJO = "pojo-state-op"; + private static final String UID_AVRO_SPECIFIC = "avro-specific-state-op"; + private static final String UID_AVRO_GENERIC = "avro-generic-state-op"; + + private StateCatalog catalog; + private TableEnvironment tableEnv; + private String dbName; + + @BeforeEach + void openCatalog() throws Exception { + catalog = openCatalogOn(RESOURCES_DIR); + dbName = catalog.listDatabases().get(0); + tableEnv = newCatalogTableEnv(catalog, dbName); + } + + @AfterEach + void closeCatalog() { + catalog.close(); + } + + @Test + void testKeyedStateCatalog() throws Exception { + List tables = catalog.listTables(dbName); + assertThat(tables) + .contains( + keyedTable(UID_PRIMITIVE), + keyedTable(UID_POJO), + keyedTable(UID_AVRO_SPECIFIC), + keyedTable(UID_AVRO_GENERIC)); + + // metadata is a view, but listTables includes views too, per the Catalog contract. + assertThat(catalog.listViews(dbName)).containsExactly(StateCatalog.METADATA_TABLE); + assertThat(tables).contains(StateCatalog.METADATA_TABLE); + + assertThat(catalog.tableExists(new ObjectPath(dbName, keyedTable(UID_PRIMITIVE)))) + .isTrue(); + assertThat(catalog.tableExists(new ObjectPath(dbName, StateCatalog.METADATA_TABLE))) + .isTrue(); + assertThat(catalog.tableExists(new ObjectPath(dbName, "nonexistent"))).isFalse(); + + assertThat(columnNames(schemaOf(catalog, dbName, keyedTable(UID_PRIMITIVE)))) + .contains("state_key", "count"); + assertThat(columnNames(schemaOf(catalog, dbName, keyedTable(UID_POJO)))) + .contains("state_key", "profile"); + assertThat(columnNames(schemaOf(catalog, dbName, keyedTable(UID_AVRO_SPECIFIC)))) + .contains("state_key", "avro_specific"); + assertThat(columnNames(schemaOf(catalog, dbName, keyedTable(UID_AVRO_GENERIC)))) + .contains("state_key", "avro_generic"); + + // Primitive state: 5 distinct int keys + List primRows = collectAll(tableEnv, UID_PRIMITIVE); + assertThat(primRows).hasSize(5); + assertThat(stateKeys(primRows)).containsExactlyInAnyOrder(1, 2, 3, 4, 5); + + // The remaining operators all hold a nested ROW value under string keys. + assertNestedRowState(UID_POJO, "profile", "name", "score"); + assertNestedRowState(UID_AVRO_SPECIFIC, "avro_specific", "name", "value"); + assertNestedRowState(UID_AVRO_GENERIC, "avro_generic", "name", "value"); + } + + private void assertNestedRowState(String operatorUid, String column, String... nestedFields) + throws Exception { + List rows = collectAll(tableEnv, operatorUid); + assertThat(rows).hasSize(5); + assertThat(stateKeys(rows)).containsExactlyInAnyOrder("1", "2", "3", "4", "5"); + for (Row row : rows) { + Row nested = (Row) row.getField(column); + assertThat(nested).isNotNull(); + for (String nestedField : nestedFields) { + assertThat(nested.getField(nestedField)).isNotNull(); + } + } + } + + @Test + void testProjectionColumnReorder() throws Exception { + // Reorder: value column first, key column second + List primRows = + collectWithSql( + tableEnv, + "SELECT `count`, state_key FROM `" + + keyedTable(UID_PRIMITIVE) + + "` ORDER BY state_key"); + + assertThat(primRows).hasSize(5); + for (Row row : primRows) { + assertThat(row.getArity()).isEqualTo(2); + assertThat(row.getField(0)).isEqualTo(1); + assertThat(row.getField("count")).isEqualTo(1); + assertThat(row.getField("state_key")).isIn(1, 2, 3, 4, 5); + } + + // POJO operator: reorder profile (ROW) before state_key + List pojoRows = + collectWithSql( + tableEnv, + "SELECT profile, state_key FROM `" + + keyedTable(UID_POJO) + + "` ORDER BY state_key"); + + assertThat(pojoRows).hasSize(5); + for (Row row : pojoRows) { + assertThat(row.getArity()).isEqualTo(2); + Row profile = (Row) row.getField(0); + assertThat(profile).isNotNull(); + assertThat(profile.getField("name")).isNotNull(); + assertThat(profile.getField("score")).isNotNull(); + assertThat(row.getField("state_key")).isNotNull(); + } + } + + @Test + void testProjectionSubsets() throws Exception { + String primTable = "`" + keyedTable(UID_PRIMITIVE) + "`"; + + List keyOnlyRows = + collectWithSql( + tableEnv, "SELECT state_key FROM " + primTable + " ORDER BY state_key"); + assertThat(keyOnlyRows).hasSize(5); + for (Row row : keyOnlyRows) { + assertThat(row.getArity()).isEqualTo(1); + } + assertThat( + keyOnlyRows.stream() + .map(r -> r.getField("state_key")) + .collect(Collectors.toList())) + .containsExactly(1, 2, 3, 4, 5); + + List valueOnlyRows = collectWithSql(tableEnv, "SELECT `count` FROM " + primTable); + assertThat(valueOnlyRows).hasSize(5); + for (Row row : valueOnlyRows) { + assertThat(row.getArity()).isEqualTo(1); + assertThat(row.getField("count")).isEqualTo(1); + } + } + } + + /** POJO-key and Avro-specific-key savepoints — off-classpath key types. */ + @Nested + class OffClasspathKeyTypes { + + private static final String POJO_AVRO_KEY_DIR = + "src/test/resources/keyed-state-pojo-avro-key"; + private static final String UID_POJO_KEY = "pojo-key-state-op"; + private static final String UID_AVRO_SPECIFIC_KEY = "avro-specific-key-state-op"; + + private StateCatalog catalog; + private TableEnvironment tableEnv; + private String dbName; + private Path savepointPath; + + @BeforeEach + void openCatalog() throws Exception { + savepointPath = findSavepointDir(POJO_AVRO_KEY_DIR); + catalog = openCatalogOn(POJO_AVRO_KEY_DIR); + dbName = catalog.listDatabases().get(0); + tableEnv = newCatalogTableEnv(catalog, dbName); + } + + @AfterEach + void closeCatalog() { + catalog.close(); + } + + @Test + void testPojoAndAvroKeySchemaTypes() throws Exception { + CheckpointMetadata metadata = + SavepointLoader.loadSavepointMetadata(savepointPath.toString()); + + // POJO key (PersonKey{int id, String name}) → ROW(id INT, name VARCHAR) + KeyedStateSchemaInfo pojoSchema = + StateTableUtils.getKeyedStateSchema( + metadata, OperatorIdentifier.forUid(UID_POJO_KEY)); + assertThat(pojoSchema.keyType.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType pojoKeyType = (RowType) pojoSchema.keyType; + assertThat(pojoKeyType.getFieldNames()).containsExactlyInAnyOrder("id", "name"); + assertThat(fieldTypeRoot(pojoKeyType, "id")).isEqualTo(LogicalTypeRoot.INTEGER); + assertThat(fieldTypeRoot(pojoKeyType, "name")).isEqualTo(LogicalTypeRoot.VARCHAR); + + // Avro specific key (StateTestRecord{String name, long value}) → ROW(name VARCHAR, + // value BIGINT) + KeyedStateSchemaInfo avroSchema = + StateTableUtils.getKeyedStateSchema( + metadata, OperatorIdentifier.forUid(UID_AVRO_SPECIFIC_KEY)); + assertThat(avroSchema.keyType.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType avroKeyType = (RowType) avroSchema.keyType; + assertThat(avroKeyType.getFieldNames()).containsExactlyInAnyOrder("name", "value"); + } + + @Test + void testPojoAndAvroKeyedStateTables() throws Exception { + assertThat(catalog.listTables(dbName)) + .contains(keyedTable(UID_POJO_KEY), keyedTable(UID_AVRO_SPECIFIC_KEY)); + + for (String uid : new String[] {UID_POJO_KEY, UID_AVRO_SPECIFIC_KEY}) { + assertThat(columnNames(schemaOf(catalog, dbName, keyedTable(uid)))) + .contains("state_key", "count"); + } + + // POJO key: 5 rows; PersonKey{id, name} off-classpath → deserialized as ROW + List pojoKeyRows = collectAll(tableEnv, UID_POJO_KEY); + assertThat(pojoKeyRows).hasSize(5); + for (Row row : pojoKeyRows) { + Row key = (Row) row.getField("state_key"); + assertThat(key).isNotNull(); + assertThat(key.getField("id")).isIn(1, 2, 3, 4, 5); + assertThat(key.getField("name")).asString().startsWith("name-"); + assertThat(row.getField("count")).isEqualTo(1); + } + + // Avro-specific key: 5 rows; StateTestRecord{name, value} off-classpath → ROW via + // GenericRecord fallback + List avroKeyRows = collectAll(tableEnv, UID_AVRO_SPECIFIC_KEY); + assertThat(avroKeyRows).hasSize(5); + for (Row row : avroKeyRows) { + Row key = (Row) row.getField("state_key"); + assertThat(key).isNotNull(); + assertThat(key.getField("name")).asString().startsWith("key-"); + assertThat(key.getField("value")).isIn(1L, 2L, 3L, 4L, 5L); + assertThat(row.getField("count")).isEqualTo(1); + } + } + } + + /** TupleX key and TupleX value with mixed basic + POJO types. */ + @Nested + class TupleKeyAndValue { + + private static final String TUPLE_KEY_DIR = "src/test/resources/keyed-state-tuple-key"; + private static final String UID_TUPLE_KEY = "tuple-key-state-op"; + private static final String UID_TUPLE_POJO_VALUE = "tuple-pojo-value-state-op"; + + private StateCatalog catalog; + private TableEnvironment tableEnv; + private String dbName; + private Path savepointPath; + + @BeforeEach + void openCatalog() throws Exception { + savepointPath = findSavepointDir(TUPLE_KEY_DIR); + catalog = openCatalogOn(TUPLE_KEY_DIR); + dbName = catalog.listDatabases().get(0); + tableEnv = newCatalogTableEnv(catalog, dbName); + } + + @AfterEach + void closeCatalog() { + catalog.close(); + } + + @Test + void testTupleKeySchemaTypes() throws Exception { + CheckpointMetadata metadata = + SavepointLoader.loadSavepointMetadata(savepointPath.toString()); + + // Tuple2 key → ROW(f0 INT NOT NULL, f1 VARCHAR) + KeyedStateSchemaInfo tupleKeySchema = + StateTableUtils.getKeyedStateSchema( + metadata, OperatorIdentifier.forUid(UID_TUPLE_KEY)); + assertThat(tupleKeySchema.keyType.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType tupleKeyType = (RowType) tupleKeySchema.keyType; + assertThat(tupleKeyType.getFieldNames()).containsExactly("f0", "f1"); + assertThat(fieldTypeRoot(tupleKeyType, "f0")).isEqualTo(LogicalTypeRoot.INTEGER); + assertThat(fieldTypeRoot(tupleKeyType, "f1")).isEqualTo(LogicalTypeRoot.VARCHAR); + + // Integer key, Tuple2 value → value column is ROW(f0 BIGINT, f1 + // ROW(name VARCHAR, score BIGINT)) + KeyedStateSchemaInfo tuplePojoValueSchema = + StateTableUtils.getKeyedStateSchema( + metadata, OperatorIdentifier.forUid(UID_TUPLE_POJO_VALUE)); + assertThat(tuplePojoValueSchema.keyType.getTypeRoot()) + .isEqualTo(LogicalTypeRoot.INTEGER); + assertThat(tuplePojoValueSchema.stateSchemas).containsKey("tuple_pojo"); + LogicalType tupleValueType = + tuplePojoValueSchema.stateSchemas.get("tuple_pojo").logicalType; + assertThat(tupleValueType.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType tupleValueRowType = (RowType) tupleValueType; + assertThat(tupleValueRowType.getFieldNames()).containsExactly("f0", "f1"); + assertThat(fieldTypeRoot(tupleValueRowType, "f0")).isEqualTo(LogicalTypeRoot.BIGINT); + assertThat(fieldTypeRoot(tupleValueRowType, "f1")).isEqualTo(LogicalTypeRoot.ROW); + RowType innerPojoType = + (RowType) tupleValueRowType.getTypeAt(tupleValueRowType.getFieldIndex("f1")); + assertThat(innerPojoType.getFieldNames()).containsExactlyInAnyOrder("name", "score"); + } + + @Test + void testTupleKeyedStateTables() throws Exception { + assertThat(catalog.listTables(dbName)) + .contains(keyedTable(UID_TUPLE_KEY), keyedTable(UID_TUPLE_POJO_VALUE)); + + // Tuple2 key: 5 rows; key fields f0=int, f1=string + List tupleKeyRows = collectAll(tableEnv, UID_TUPLE_KEY); + assertThat(tupleKeyRows).hasSize(5); + for (Row row : tupleKeyRows) { + Row key = (Row) row.getField("state_key"); + assertThat(key).isNotNull(); + assertThat(key.getField("f0")).isIn(1, 2, 3, 4, 5); + assertThat(key.getField("f1")).asString().isEqualTo("k-" + key.getField("f0")); + assertThat(row.getField("count")).isEqualTo(1); + } + + // Tuple2 value: 5 rows; value row has f0=long, f1=row(name,score) + List tuplePojoValueRows = collectAll(tableEnv, UID_TUPLE_POJO_VALUE); + assertThat(tuplePojoValueRows).hasSize(5); + for (Row row : tuplePojoValueRows) { + Integer key = (Integer) row.getField("state_key"); + assertThat(key).isIn(1, 2, 3, 4, 5); + Row tupleValue = (Row) row.getField("tuple_pojo"); + assertThat(tupleValue).isNotNull(); + assertThat(tupleValue.getField("f0")).isEqualTo((long) key * 10); + Row pojoField = (Row) tupleValue.getField("f1"); + assertThat(pojoField).isNotNull(); + assertThat(pojoField.getField("name")).isEqualTo("name-" + key); + assertThat(pojoField.getField("score")).isEqualTo((long) key * 100); + } + } + } + + // ------------------------------------------------------------------------- + // Shared helpers + // ------------------------------------------------------------------------- + + private static StateCatalog openCatalogOn(String resourceDir) throws Exception { + String catalogRoot = Paths.get(resourceDir).toAbsolutePath().toString(); + StateCatalog catalog = + new StateCatalog("state", Collections.singletonMap("test", catalogRoot)); + catalog.open(); + return catalog; + } + + private static TableEnvironment newCatalogTableEnv(StateCatalog catalog, String dbName) { + TableEnvironment tableEnv = TableEnvironment.create(EnvironmentSettings.inBatchMode()); + tableEnv.loadModule("state", StateModule.INSTANCE); + tableEnv.registerCatalog("state", catalog); + tableEnv.useCatalog("state"); + tableEnv.useDatabase(dbName); + return tableEnv; + } + + /** Finds the single {@code savepoint-*} directory nested directly under {@code parentDir}. */ + private static Path findSavepointDir(String parentDir) throws IOException { + try (var stream = Files.list(Paths.get(parentDir))) { + return stream.filter( + p -> + Files.isDirectory(p) + && p.getFileName().toString().startsWith("savepoint-")) + .findFirst() + .orElseThrow(() -> new IOException("No savepoint found in " + parentDir)); + } + } + + private static String keyedTable(String operatorUid) { + return StateCatalog.OPERATOR_UID_PREFIX + operatorUid + StateCatalog.OPERATOR_TABLE_SUFFIX; + } + + private static String flatKeyedTable(String operatorUid, String stateName) { + return StateCatalog.OPERATOR_UID_PREFIX + + operatorUid + + "_" + + stateName + + StateCatalog.FLAT_STATE_TABLE_SUFFIX; + } + + private static Schema schemaOf(StateCatalog catalog, String dbName, String tableName) + throws Exception { + return ((CatalogTable) catalog.getTable(new ObjectPath(dbName, tableName))) + .getUnresolvedSchema(); + } + + private static List columnNames(Schema schema) { + return schema.getColumns().stream() + .map(Schema.UnresolvedColumn::getName) + .collect(Collectors.toList()); + } + + private static LogicalTypeRoot stateTypeRoot( + KeyedStateSchemaInfo schemaInfo, String stateName) { + KeyedStateSchemaInfo.StateEntryInfo entry = schemaInfo.stateSchemas.get(stateName); + assertThat(entry).as("state '%s'", stateName).isNotNull(); + return entry.logicalType.getTypeRoot(); + } + + private static LogicalTypeRoot fieldTypeRoot(RowType rowType, String fieldName) { + return rowType.getTypeAt(rowType.getFieldIndex(fieldName)).getTypeRoot(); + } + + private static Set stateKeys(List rows) { + return rows.stream().map(r -> r.getField("state_key")).collect(Collectors.toSet()); + } + + private static List longRange(int endExclusive) { + return LongStream.range(0, endExclusive).boxed().collect(Collectors.toList()); + } + + /** + * Registers the discovered keyed-state table of {@code operatorUid} as a temporary table and + * reads it in batch mode via {@code StreamTableEnvironment}, bypassing {@link StateCatalog}. + */ + private static List readViaTemporaryTable( + String statePath, String operatorUid, String tableName) throws Exception { + Configuration config = new Configuration(); + config.set(RUNTIME_MODE, RuntimeExecutionMode.BATCH); + StreamTableEnvironment tEnv = + StreamTableEnvironment.create( + StreamExecutionEnvironment.getExecutionEnvironment(config)); + + CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(statePath); + OperatorIdentifier opId = OperatorIdentifier.forUid(operatorUid); + CatalogTable catalogTable = + StateTableUtils.getStateCatalogTable( + metadata, + StateTableUtils.getKeyedStateSchema(metadata, opId), + statePath, + opId); + + CatalogManager catalogManager = ((TableEnvironmentImpl) tEnv).getCatalogManager(); + catalogManager.createTemporaryTable( + catalogTable, + catalogManager.qualifyIdentifier(UnresolvedIdentifier.of(tableName)), + false); + + return tEnv.toDataStream(tEnv.sqlQuery("SELECT * FROM " + tableName)) + .executeAndCollect(100); + } + + private static List collectAll(TableEnvironment tEnv, String operatorUid) + throws Exception { + return collectWithSql(tEnv, "SELECT * FROM `" + keyedTable(operatorUid) + "`"); + } + + private static List collectWithSql(TableEnvironment tEnv, String sql) throws Exception { + List rows = new ArrayList<>(); + TableResult result = tEnv.executeSql(sql); + try (CloseableIterator it = result.collect()) { + it.forEachRemaining(rows::add); + } + return rows; + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/StateCatalogTest.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/StateCatalogTest.java new file mode 100644 index 00000000000000..44e8ab8264fe61 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/StateCatalogTest.java @@ -0,0 +1,143 @@ +/* + * 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.catalog; + +import org.apache.flink.table.catalog.CatalogDatabaseImpl; +import org.apache.flink.table.catalog.CatalogView; +import org.apache.flink.table.catalog.ObjectPath; +import org.apache.flink.table.catalog.exceptions.DatabaseNotExistException; +import org.apache.flink.table.catalog.exceptions.TableNotExistException; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.attribute.FileTime; +import java.time.Instant; +import java.util.Collections; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Unit and functional tests for {@link StateCatalog} that are specific to the catalog layer + * (CatalogView/CatalogTable semantics, unsupported write operations). Directory scanning, db-name + * derivation, and dynamic re-discovery are {@link StateCatalog}'s delegation to {@link + * SnapshotDiscovery} and are covered exhaustively by {@link SnapshotDiscoveryTest} instead of being + * re-verified here. + */ +class StateCatalogTest { + + // Every metadata file created by createMetadataFile() gets this exact modification time, so + // the creationTs segment in derived database names is deterministic across all tests. + private static final Instant FIXED_TS = Instant.parse("2024-03-15T10:30:45Z"); + private static final String TS = "2024-03-15T10:30:45Z"; + private static final String METADATA_TABLE = StateCatalog.METADATA_TABLE; + + @TempDir Path tempDir; + + @Test + void testCatalogOperations() throws Exception { + createMetadataFile(tempDir.resolve("savepoint-abc")); + StateCatalog catalog = openCatalog("app1", tempDir); + String dbName = "app1/" + TS + "/savepoint-abc"; + + // databaseExists + assertThat(catalog.databaseExists(dbName)).isTrue(); + assertThat(catalog.databaseExists("app1/" + TS + "/savepoint-nonexistent")).isFalse(); + assertThat(catalog.databaseExists("unknown/" + TS + "/savepoint-abc")).isFalse(); + + // tableExists + assertThat(catalog.tableExists(new ObjectPath(dbName, METADATA_TABLE))).isTrue(); + assertThat(catalog.tableExists(new ObjectPath(dbName, "nonexistent"))).isFalse(); + assertThat( + catalog.tableExists( + new ObjectPath("app1/" + TS + "/nonexistent", METADATA_TABLE))) + .isFalse(); + + // getTable returns CatalogView with correct query + Path savepointDir = tempDir.resolve("savepoint-abc"); + CatalogView view = (CatalogView) catalog.getTable(new ObjectPath(dbName, METADATA_TABLE)); + assertThat(view.getOriginalQuery()) + .contains("savepoint_metadata") + .contains(savepointDir.toAbsolutePath().toString()); + + // getDatabase throws for unknown + assertThatThrownBy(() -> catalog.getDatabase("app1/" + TS + "/savepoint-nonexistent")) + .isInstanceOf(DatabaseNotExistException.class); + + // getTable throws for unknown snapshot + assertThatThrownBy( + () -> + catalog.getTable( + new ObjectPath( + "app1/" + TS + "/savepoint-nonexistent", + METADATA_TABLE))) + .isInstanceOf(TableNotExistException.class); + + catalog.close(); + } + + @Test + void testListFunctionsAlwaysReturnsEmpty() throws Exception { + StateCatalog catalog = openCatalog("app1", tempDir); + assertThat(catalog.listFunctions("nonexistent")).isEmpty(); + catalog.close(); + } + + @Test + void testWriteOperationsThrow() throws Exception { + StateCatalog catalog = openCatalog("app1", tempDir); + + assertThatThrownBy( + () -> + catalog.createDatabase( + "db", + new CatalogDatabaseImpl(Collections.emptyMap(), ""), + false)) + .isInstanceOf(UnsupportedOperationException.class); + assertThatThrownBy(() -> catalog.dropDatabase("db", true, false)) + .isInstanceOf(UnsupportedOperationException.class); + assertThatThrownBy(() -> catalog.createTable(new ObjectPath("db", "t"), null, false)) + .isInstanceOf(UnsupportedOperationException.class); + + catalog.close(); + } + + // ------------------------------------------------------------------------- + // Helpers + // ------------------------------------------------------------------------- + + private static void createMetadataFile(Path snapshotDir) throws IOException { + Files.createDirectories(snapshotDir); + Path file = Files.createFile(snapshotDir.resolve("_metadata")); + Files.setLastModifiedTime(file, FileTime.from(FIXED_TS)); + } + + private static StateCatalog openCatalog(String label, Path directory) throws Exception { + StateCatalog catalog = + new StateCatalog( + "state", + Collections.singletonMap(label, directory.toAbsolutePath().toString())); + catalog.open(); + return catalog; + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/TuplePojoField.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/TuplePojoField.java new file mode 100644 index 00000000000000..c963a306e3fce9 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/TuplePojoField.java @@ -0,0 +1,59 @@ +/* + * 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.catalog; + +import java.util.Objects; + +/** + * Simple POJO used as a nested field inside a Tuple value state in {@link + * StateCatalogGeneratedSavepointITCase.TupleKeyAndValue}. + * + *

Must remain in the normal test compilation scope (not in resources/generator/) so that it is + * on the classpath during test runs and the TupleSerializer can deserialize it. + */ +public class TuplePojoField { + public String name; + public long score; + + public TuplePojoField() {} + + public TuplePojoField(String name, long score) { + this.name = name; + this.score = score; + } + + @Override + public boolean equals(Object o) { + if (!(o instanceof TuplePojoField)) { + return false; + } + TuplePojoField other = (TuplePojoField) o; + return Objects.equals(name, other.name) && score == other.score; + } + + @Override + public int hashCode() { + return Objects.hash(name, score); + } + + @Override + public String toString() { + return "TuplePojoField{name='" + name + "', score=" + score + "}"; + } +} 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 e2a8ca574d0ce1..3b587f7ce6e73b 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 @@ -37,6 +37,7 @@ import static org.apache.flink.configuration.ExecutionOptions.RUNTIME_MODE; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatCode; /** Unit tests for the savepoint SQL reader. */ class SavepointDynamicTableSourceTest { @@ -410,6 +411,94 @@ void testUnsupportedFilterIsNotPushedDownButReturnsCorrectResult() throws Except assertThat(keys).containsExactly(0L, 2L, 4L, 6L, 8L); } + // ------------------------------------------------------------------------- + // Projection push-down tests + // ------------------------------------------------------------------------- + + @Test + void testProjectionPushDownSelectKeyAndOneColumn() throws Exception { + StreamTableEnvironment tEnv = createBatchTableEnv(); + tEnv.executeSql(STATE_TABLE_DDL); + + // Projection only: all 10 rows, 2 columns. + String sql = "SELECT k, KeyedPrimitiveValue FROM state_table ORDER BY k"; + List result = tEnv.toDataStream(tEnv.sqlQuery(sql)).executeAndCollect(100); + + assertThat(result).hasSize(10); + for (Row row : result) { + assertThat(row.getArity()).isEqualTo(2); + assertThat(row.getField("KeyedPrimitiveValue")).isEqualTo(1L); + } + List keys = + result.stream().map(r -> (Long) r.getField("k")).collect(Collectors.toList()); + assertThat(keys) + .containsExactlyElementsOf( + LongStream.range(0L, 10L).boxed().collect(Collectors.toList())); + + // Projection combined with filter push-down: applyProjection updates keyColumnIndex to its + // position in the projected row, but keyFilter holds only the key value so it remains + // valid regardless of how the key column moves in the output. + String filteredSql = "SELECT k, KeyedPrimitiveValue FROM state_table WHERE k = 5"; + assertThat(hasPushedDownFilter(tEnv, filteredSql)).isTrue(); + List filteredResult = + tEnv.toDataStream(tEnv.sqlQuery(filteredSql)).executeAndCollect(100); + assertThat(filteredResult).hasSize(1); + Row row = filteredResult.get(0); + assertThat(row.getArity()).isEqualTo(2); + assertThat(row.getField("k")).isEqualTo(5L); + assertThat(row.getField("KeyedPrimitiveValue")).isEqualTo(1L); + } + + @Test + @SuppressWarnings("unchecked") + void testProjectionPushDownAllColumns() throws Exception { + StreamTableEnvironment tEnv = createBatchTableEnv(); + tEnv.executeSql(STATE_TABLE_DDL); + + List result = + tEnv.toDataStream(tEnv.sqlQuery("SELECT * FROM state_table")) + .executeAndCollect(100); + + assertThat(result).hasSize(10); + for (Row row : result) { + assertThat(row.getArity()).isEqualTo(5); + assertThat(row.getField("KeyedPrimitiveValue")).isEqualTo(1L); + } + } + + // ------------------------------------------------------------------------- + // Lazy type resolution tests + // ------------------------------------------------------------------------- + + @Test + void testPlanningSucceedsWithNonexistentSavepointPath() { + // Planning must never touch the savepoint (metadata I/O or class loading); both are + // deferred to the scan runtime provider. A nonexistent path/operator would fail + // immediately if either were resolved eagerly during planning. + StreamTableEnvironment tEnv = createBatchTableEnv(); + + String ddl = + "CREATE TABLE state_table_missing (\n" + + " k bigint,\n" + + " KeyedPrimitiveValue bigint,\n" + + " PRIMARY KEY (k) NOT ENFORCED\n" + + ")\n" + + "with (\n" + + " 'connector' = 'savepoint',\n" + + " 'state.path' = 'src/test/resources/does-not-exist',\n" + + " 'operator.uid' = 'nonexistent-operator-uid'\n" + + ")"; + + tEnv.executeSql(ddl); + tEnv.executeSql("CREATE TABLE sink (k BIGINT) WITH ('connector' = 'blackhole')"); + + assertThatCode( + () -> + tEnv.compilePlanSql( + "INSERT INTO sink SELECT k FROM state_table_missing")) + .doesNotThrowAnyException(); + } + // ------------------------------------------------------------------------- // Helpers // ------------------------------------------------------------------------- @@ -423,7 +512,7 @@ private static StreamTableEnvironment createBatchTableEnv() { private static final Pattern PUSHED_DOWN_FILTER = Pattern.compile( - "TableSourceScan\\(table=\\[\\[default_catalog, default_database, state_table, filter=\\[.+?]]]"); + "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/SavepointTypeInfoResolverTest.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/table/SavepointTypeInfoResolverTest.java new file mode 100644 index 00000000000000..1d028a669a1d9a --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/table/SavepointTypeInfoResolverTest.java @@ -0,0 +1,104 @@ +/* + * 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.api.common.serialization.SerializerConfigImpl; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; +import org.apache.flink.api.common.typeutils.base.LongSerializer; +import org.apache.flink.api.common.typeutils.base.StringSerializer; +import org.apache.flink.runtime.state.metainfo.StateMetaInfoSnapshot; + +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Unit tests for {@link SavepointTypeInfoResolver#resolveNamespaceSerializer}. */ +class SavepointTypeInfoResolverTest { + + private static final String STATE_NAME = "window-state"; + + @Test + void resolveNamespaceSerializerReturnsRestoredSerializer() { + Map> serializerSnapshots = new HashMap<>(); + serializerSnapshots.put( + StateMetaInfoSnapshot.CommonSerializerKeys.NAMESPACE_SERIALIZER.toString(), + LongSerializer.INSTANCE.snapshotConfiguration()); + + StateMetaInfoSnapshot metaInfoSnapshot = + new StateMetaInfoSnapshot( + STATE_NAME, + StateMetaInfoSnapshot.BackendStateType.KEY_VALUE, + Collections.emptyMap(), + serializerSnapshots); + + SavepointTypeInfoResolver resolver = + new SavepointTypeInfoResolver( + Collections.singletonMap(STATE_NAME, metaInfoSnapshot), + new SerializerConfigImpl(), + null); + + TypeSerializer namespaceSerializer = resolver.resolveNamespaceSerializer(STATE_NAME); + + assertThat(namespaceSerializer).isInstanceOf(LongSerializer.class); + } + + @Test + void resolveNamespaceSerializerThrowsWhenStateNotFound() { + SavepointTypeInfoResolver resolver = + new SavepointTypeInfoResolver( + Collections.emptyMap(), new SerializerConfigImpl(), null); + + assertThatThrownBy(() -> resolver.resolveNamespaceSerializer(STATE_NAME)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining(STATE_NAME) + .hasMessageContaining("not found in preloaded metadata"); + } + + @Test + void resolveNamespaceSerializerThrowsWhenNamespaceSerializerMissing() { + Map> serializerSnapshots = new HashMap<>(); + serializerSnapshots.put( + StateMetaInfoSnapshot.CommonSerializerKeys.VALUE_SERIALIZER.toString(), + StringSerializer.INSTANCE.snapshotConfiguration()); + + StateMetaInfoSnapshot metaInfoSnapshot = + new StateMetaInfoSnapshot( + STATE_NAME, + StateMetaInfoSnapshot.BackendStateType.KEY_VALUE, + Collections.emptyMap(), + serializerSnapshots); + + SavepointTypeInfoResolver resolver = + new SavepointTypeInfoResolver( + Collections.singletonMap(STATE_NAME, metaInfoSnapshot), + new SerializerConfigImpl(), + null); + + assertThatThrownBy(() -> resolver.resolveNamespaceSerializer(STATE_NAME)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining(STATE_NAME) + .hasMessageContaining("no namespace serializer in metadata"); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/table/SavepointTypeInformationFactoryTest.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/table/SavepointTypeInformationFactoryTest.java deleted file mode 100644 index b8f456a9bc0ca0..00000000000000 --- a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/table/SavepointTypeInformationFactoryTest.java +++ /dev/null @@ -1,129 +0,0 @@ -/* - * 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.api.common.RuntimeExecutionMode; -import org.apache.flink.api.common.typeinfo.TypeInformation; -import org.apache.flink.configuration.Configuration; -import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; -import org.apache.flink.table.api.Table; -import org.apache.flink.table.api.bridge.java.StreamTableEnvironment; -import org.apache.flink.types.Row; - -import org.junit.jupiter.api.Test; - -import java.util.List; -import java.util.Set; -import java.util.stream.Collectors; - -import static org.apache.flink.configuration.ExecutionOptions.RUNTIME_MODE; -import static org.assertj.core.api.Assertions.assertThat; - -/** Unit tests for the SavepointTypeInformationFactory. */ -class SavepointTypeInformationFactoryTest { - - public static class TestLongTypeInformationFactory implements SavepointTypeInformationFactory { - private static volatile boolean wasCalled = false; - - public static boolean wasFactoryCalled() { - return wasCalled; - } - - public static void resetCallTracker() { - wasCalled = false; - } - - @Override - public TypeInformation getTypeInformation() { - wasCalled = true; - return TypeInformation.of(Long.class); - } - } - - private static class TestStringTypeInformationFactory - implements SavepointTypeInformationFactory { - @Override - public TypeInformation getTypeInformation() { - return TypeInformation.of(String.class); - } - } - - @Test - void testSavepointTypeInformationFactoryEndToEnd() throws Exception { - TestLongTypeInformationFactory.resetCallTracker(); - - Configuration config = new Configuration(); - config.set(RUNTIME_MODE, RuntimeExecutionMode.BATCH); - StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(config); - StreamTableEnvironment tEnv = StreamTableEnvironment.create(env); - - final String sql = - "CREATE TABLE state_table (\n" - + " k bigint,\n" - + " KeyedPrimitiveValue 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" - + " 'fields.KeyedPrimitiveValue.value-type-factory' = '" - + TestLongTypeInformationFactory.class.getName() - + "'\n" - + ")"; - - tEnv.executeSql(sql); - Table table = tEnv.sqlQuery("SELECT k, KeyedPrimitiveValue FROM state_table"); - List result = tEnv.toDataStream(table).executeAndCollect(100); - - assertThat(TestLongTypeInformationFactory.wasFactoryCalled()) - .as( - "Factory getTypeInformation() method must be called - this proves factory is used instead of metadata inference") - .isTrue(); - - assertThat(result).hasSize(10); - - Set keys = - result.stream().map(r -> (Long) r.getField("k")).collect(Collectors.toSet()); - assertThat(keys).hasSize(10); - assertThat(keys).containsExactlyInAnyOrder(0L, 1L, 2L, 3L, 4L, 5L, 6L, 7L, 8L, 9L); - - Set primitiveValues = - result.stream() - .map(r -> (Long) r.getField("KeyedPrimitiveValue")) - .collect(Collectors.toSet()); - assertThat(primitiveValues).containsExactly(1L); - } - - @Test - void testBasicFactoryFunctionality() { - TestLongTypeInformationFactory.resetCallTracker(); - - TestLongTypeInformationFactory longFactory = new TestLongTypeInformationFactory(); - TypeInformation longTypeInfo = longFactory.getTypeInformation(); - - assertThat(longTypeInfo).isEqualTo(TypeInformation.of(Long.class)); - assertThat(TestLongTypeInformationFactory.wasFactoryCalled()).isTrue(); - - TestStringTypeInformationFactory stringFactory = new TestStringTypeInformationFactory(); - TypeInformation stringTypeInfo = stringFactory.getTypeInformation(); - - assertThat(stringTypeInfo).isEqualTo(TypeInformation.of(String.class)); - } -} diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/table/TypeConversionDriftGuardTest.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/table/TypeConversionDriftGuardTest.java new file mode 100644 index 00000000000000..fcb0f415f2afbb --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/table/TypeConversionDriftGuardTest.java @@ -0,0 +1,273 @@ +/* + * 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.api.common.serialization.SerializerConfigImpl; +import org.apache.flink.api.common.typeutils.base.ListSerializer; +import org.apache.flink.api.common.typeutils.base.LongSerializer; +import org.apache.flink.api.common.typeutils.base.MapSerializer; +import org.apache.flink.api.common.typeutils.base.StringSerializer; +import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.api.java.typeutils.TypeExtractor; +import org.apache.flink.formats.avro.typeutils.AvroTypeInfo; +import org.apache.flink.state.api.schema.SerializerSnapshotToLogicalTypeConverter; +import org.apache.flink.streaming.api.windowing.windows.GlobalWindow; +import org.apache.flink.streaming.api.windowing.windows.TimeWindow; +import org.apache.flink.table.data.GenericArrayData; +import org.apache.flink.table.data.GenericMapData; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.StringData; +import org.apache.flink.table.data.TimestampData; +import org.apache.flink.table.runtime.typeutils.RowDataSerializer; +import org.apache.flink.table.types.logical.ArrayType; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.LogicalTypeRoot; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.table.types.logical.VarCharType; + +import com.example.state.writer.job.schema.avro.AvroRecord; +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Drift guard between {@link SerializerSnapshotToLogicalTypeConverter} (schema-time: {@code + * TypeSerializerSnapshot} -> {@link LogicalType}) and {@link StateValueConverter}/{@link + * org.apache.flink.state.api.input.deserializer.InternalTypeConverter} (runtime: raw deserialized + * value -> internal representation). + * + *

For every type shape the schema-time converter knows how to describe, this asserts that a + * representative runtime value of that shape round-trips through the runtime converters without + * throwing and lands on the internal representation the shape implies. This is a permanent + * regression guard, not a one-time check: whenever a new case is added to {@link + * SerializerSnapshotToLogicalTypeConverter#convert}, a matching case must be added here too, or + * this test stops actually covering the new shape. + */ +class TypeConversionDriftGuardTest { + + private final StateValueConverter converter = new StateValueConverter(); + + @Test + void testInt() { + LogicalType type = new IntType(false); + Object result = converter.getValue(type, 42); + assertThat(result).isEqualTo(42); + } + + @Test + void testLong() { + LogicalType type = + SerializerSnapshotToLogicalTypeConverter.convert( + new LongSerializer.LongSerializerSnapshot()); + Object result = converter.getValue(type, 42L); + assertThat(result).isEqualTo(42L); + } + + @Test + void testString() { + LogicalType type = + SerializerSnapshotToLogicalTypeConverter.convert( + new StringSerializer.StringSerializerSnapshot()); + Object result = converter.getValue(type, "hello"); + assertThat(result).isInstanceOf(StringData.class); + assertThat(result.toString()).isEqualTo("hello"); + } + + @Test + void testList() { + ListSerializer ser = new ListSerializer<>(StringSerializer.INSTANCE); + LogicalType type = + SerializerSnapshotToLogicalTypeConverter.convert(ser.snapshotConfiguration()); + assertThat(type.getTypeRoot()).isEqualTo(LogicalTypeRoot.ARRAY); + + List value = Arrays.asList("a", "b", "c"); + Object result = converter.getValue(type, value); + assertThat(result).isInstanceOf(GenericArrayData.class); + assertThat(((GenericArrayData) result).size()).isEqualTo(3); + } + + @Test + void testMap() { + MapSerializer ser = + new MapSerializer<>(StringSerializer.INSTANCE, LongSerializer.INSTANCE); + LogicalType type = + SerializerSnapshotToLogicalTypeConverter.convert(ser.snapshotConfiguration()); + assertThat(type.getTypeRoot()).isEqualTo(LogicalTypeRoot.MAP); + + Map value = Collections.singletonMap("k", 7L); + Object result = converter.getValue(type, value); + assertThat(result).isInstanceOf(GenericMapData.class); + assertThat(((GenericMapData) result).size()).isEqualTo(1); + } + + @Test + void testNullableInnerType() { + ListSerializer ser = new ListSerializer<>(StringSerializer.INSTANCE); + LogicalType innerType = + SerializerSnapshotToLogicalTypeConverter.convert(ser.snapshotConfiguration()) + .copy(true); + assertThat(converter.getValue(innerType, null)).isNull(); + } + + @Test + void testTimeWindow() { + LogicalType type = + SerializerSnapshotToLogicalTypeConverter.convert( + new TimeWindow.Serializer().snapshotConfiguration()); + assertThat(type.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + + TimeWindow window = new TimeWindow(100L, 200L); + Object result = converter.getValue(type, window); + assertThat(result).isInstanceOf(GenericRowData.class); + GenericRowData row = (GenericRowData) result; + assertThat(row.getArity()).isEqualTo(2); + assertThat(row.getField(0)).isEqualTo(TimestampData.fromEpochMillis(100L)); + assertThat(row.getField(1)).isEqualTo(TimestampData.fromEpochMillis(200L)); + } + + @Test + void testGlobalWindow() { + LogicalType type = + SerializerSnapshotToLogicalTypeConverter.convert( + new GlobalWindow.Serializer().snapshotConfiguration()); + assertThat(type.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + assertThat(((RowType) type).getFieldCount()).isEqualTo(0); + + Object result = converter.getValue(type, GlobalWindow.get()); + assertThat(result).isInstanceOf(GenericRowData.class); + assertThat(((GenericRowData) result).getArity()).isEqualTo(0); + } + + @Test + void testTuple() { + @SuppressWarnings({"unchecked", "rawtypes"}) + org.apache.flink.api.common.typeutils.TypeSerializerSnapshot snapshot = + TypeExtractor.getForObject(Tuple2.of(1, "x")) + .createSerializer(new SerializerConfigImpl()) + .snapshotConfiguration(); + LogicalType type = SerializerSnapshotToLogicalTypeConverter.convert(snapshot); + assertThat(type.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + + Object result = converter.getValue(type, Tuple2.of(1, "x")); + assertThat(result).isInstanceOf(GenericRowData.class); + GenericRowData row = (GenericRowData) result; + assertThat(row.getArity()).isEqualTo(2); + assertThat(row.getField(0)).isEqualTo(1); + assertThat(row.getField(1).toString()).isEqualTo("x"); + } + + /** POJO with the same shape used by {@code SerializerSnapshotToLogicalTypeConverterTest}. */ + public static class SamplePojo { + public String name; + public int age; + + public SamplePojo() {} + + public SamplePojo(String name, int age) { + this.name = name; + this.age = age; + } + } + + @Test + void testPojo() { + @SuppressWarnings({"unchecked", "rawtypes"}) + org.apache.flink.api.common.typeutils.TypeSerializerSnapshot snapshot = + TypeExtractor.createTypeInfo(SamplePojo.class) + .createSerializer(new SerializerConfigImpl()) + .snapshotConfiguration(); + LogicalType type = SerializerSnapshotToLogicalTypeConverter.convert(snapshot); + assertThat(type.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + + Object result = converter.getValue(type, new SamplePojo("bob", 30)); + assertThat(result).isInstanceOf(GenericRowData.class); + RowType rowType = (RowType) type; + GenericRowData row = (GenericRowData) result; + int nameIdx = rowType.getFieldIndex("name"); + int ageIdx = rowType.getFieldIndex("age"); + assertThat(row.getField(nameIdx).toString()).isEqualTo("bob"); + assertThat(row.getField(ageIdx)).isEqualTo(30); + } + + @Test + void testAvroSpecificRecord() { + LogicalType type = + SerializerSnapshotToLogicalTypeConverter.convert( + new AvroTypeInfo<>(AvroRecord.class) + .createSerializer(new SerializerConfigImpl()) + .snapshotConfiguration()); + assertThat(type.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + + AvroRecord avroRecord = AvroRecord.newBuilder().setLongData(99L).build(); + Object result = converter.getValue(type, avroRecord); + assertThat(result).isInstanceOf(GenericRowData.class); + RowType rowType = (RowType) type; + int idx = rowType.getFieldIndex("longData"); + assertThat(((GenericRowData) result).getField(idx)).isEqualTo(99L); + } + + @Test + void testRowDataPassThrough() { + RowType rowType = + RowType.of( + new LogicalType[] {new IntType(), VarCharType.STRING_TYPE}, + new String[] {"id", "name"}); + RowDataSerializer serializer = new RowDataSerializer(rowType); + LogicalType type = + SerializerSnapshotToLogicalTypeConverter.convert( + serializer.snapshotConfiguration()); + assertThat(type.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + + GenericRowData sourceRow = new GenericRowData(2); + sourceRow.setField(0, 1); + sourceRow.setField(1, StringData.fromString("abc")); + + Object result = converter.getValue(type, sourceRow); + assertThat(result).isInstanceOf(GenericRowData.class); + GenericRowData resultRow = (GenericRowData) result; + assertThat(resultRow.getField(0)).isEqualTo(1); + assertThat(resultRow.getField(1).toString()).isEqualTo("abc"); + } + + @Test + void testArrayOfArray() { + ListSerializer> serializer = + new ListSerializer<>(new ListSerializer<>(StringSerializer.INSTANCE)); + LogicalType type = + SerializerSnapshotToLogicalTypeConverter.convert( + serializer.snapshotConfiguration()); + assertThat(type.getTypeRoot()).isEqualTo(LogicalTypeRoot.ARRAY); + assertThat(((ArrayType) type).getElementType().getTypeRoot()) + .isEqualTo(LogicalTypeRoot.ARRAY); + + List> value = + Arrays.asList(Arrays.asList("a", "b"), Collections.singletonList("c")); + Object result = converter.getValue(type, value); + assertThat(result).isInstanceOf(GenericArrayData.class); + GenericArrayData outerArray = (GenericArrayData) result; + assertThat(outerArray.size()).isEqualTo(2); + assertThat(outerArray.getArray(0)).isInstanceOf(GenericArrayData.class); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/resources/generator/KeyedStateCatalogSavepointGenerator.java b/flink-libraries/flink-state-processing-api/src/test/resources/generator/KeyedStateCatalogSavepointGenerator.java new file mode 100644 index 00000000000000..2dda9bb6f9a796 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/resources/generator/KeyedStateCatalogSavepointGenerator.java @@ -0,0 +1,356 @@ +/* + * 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.catalog; + +import org.apache.flink.api.common.JobStatus; +import org.apache.flink.api.common.functions.OpenContext; +import org.apache.flink.api.common.state.ValueState; +import org.apache.flink.api.common.state.ValueStateDescriptor; +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.StateBackendOptions; +import org.apache.flink.core.execution.JobClient; +import org.apache.flink.core.execution.SavepointFormatType; +import org.apache.flink.formats.avro.typeutils.AvroTypeInfo; +import org.apache.flink.formats.avro.typeutils.GenericRecordAvroTypeInfo; +import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration; +import org.apache.flink.state.catalog.avro.StateTestRecord; +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.api.functions.source.legacy.RichSourceFunction; +import org.apache.flink.test.util.MiniClusterWithClientResource; +import org.apache.flink.util.Collector; + +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.file.FileVisitResult; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.nio.file.SimpleFileVisitor; +import java.nio.file.attribute.BasicFileAttributes; +import java.util.Objects; +import java.util.concurrent.TimeUnit; + +/** + * Generates the pre-built savepoint used by {@code + * StateCatalogGeneratedSavepointITCase.MultiOperatorTypeCatalog}. + * + *

This file lives in {@code src/test/resources/generator/} so it is NOT compiled as part of the + * normal build. The generated savepoint lives in {@code src/test/resources/keyed-state-catalog/} + * and is committed to the repository. + * + *

To regenerate the savepoint (e.g. after a Flink serializer format change): + * + *

    + *
  1. Copy this file to {@code src/test/java/org/apache/flink/state/catalog/}. + *
  2. Copy {@code src/test/resources/generator/StateTestRecord.avsc} to {@code + * src/test/resources/avro/StateTestRecord.avsc} so the Avro class is generated (the {@code + * avro-maven-plugin} and {@code flink-avro} dependency are already configured in pom.xml, no + * changes needed there). + *
  3. Remove the {@code @Disabled} annotation and run: {@code ./mvnw test -pl + * flink-libraries/flink-state-processing-api + * -Dtest=KeyedStateCatalogSavepointGenerator#generateSavepoint} + *
  4. Verify the new savepoint under {@code src/test/resources/keyed-state-catalog/}. + *
  5. Remove the copied {@code .java} file from the source tree and the copied {@code .avsc} + * file from {@code src/test/resources/avro/}. + *
+ */ +@Disabled("Run manually to regenerate the pre-built savepoint in test resources") +class KeyedStateCatalogSavepointGenerator { + + private static final String RESOURCES_DIR = "src/test/resources/keyed-state-catalog"; + + // Must match StateCatalogGeneratedSavepointITCase.MultiOperatorTypeCatalog's UID_* constants. + private static final String UID_PRIMITIVE = "primitive-state-op"; + private static final String UID_POJO = "pojo-state-op"; + private static final String UID_AVRO_SPECIFIC = "avro-specific-state-op"; + private static final String UID_AVRO_GENERIC = "avro-generic-state-op"; + + @Test + void generateSavepoint() throws Exception { + Path outputDir = Paths.get(RESOURCES_DIR).toAbsolutePath(); + Files.createDirectories(outputDir); + deleteExistingSavepoints(outputDir); + + var cluster = + new MiniClusterWithClientResource( + new MiniClusterResourceConfiguration.Builder() + .setNumberSlotsPerTaskManager(4) + .build()); + cluster.before(); + try { + Configuration cfg = new Configuration(); + cfg.set(StateBackendOptions.STATE_BACKEND, "hashmap"); + var env = StreamExecutionEnvironment.getExecutionEnvironment(cfg); + env.setParallelism(2); + + env.addSource(new BoundedWaitingSource(new int[] {1, 2, 3, 4, 5})) + .returns(Types.INT) + .keyBy(v -> v) + .process(new IntCountOperator()) + .uid(UID_PRIMITIVE) + .name(UID_PRIMITIVE) + .keyBy(v -> String.valueOf(v)) + .process(new ProfileOperator()) + .uid(UID_POJO) + .name(UID_POJO) + .keyBy(v -> String.valueOf(v)) + .process(new AvroSpecificOperator()) + .uid(UID_AVRO_SPECIFIC) + .name(UID_AVRO_SPECIFIC) + .keyBy(v -> String.valueOf(v)) + .process(new AvroGenericOperator()) + .uid(UID_AVRO_GENERIC) + .name(UID_AVRO_GENERIC) + .sinkTo(new DiscardingSink<>()); + + String savepointPath = takeSavepoint(env, outputDir.toString()); + System.out.println("Savepoint written to: " + savepointPath); + } finally { + cluster.after(); + } + } + + private static String takeSavepoint(StreamExecutionEnvironment env, String savepointDir) + throws Exception { + JobClient jobClient = env.executeAsync(); + try { + while (jobClient.getJobStatus().get() != JobStatus.RUNNING) { + Thread.sleep(100); + } + Exception lastEx = null; + for (int attempt = 0; attempt < 30; attempt++) { + try { + return jobClient + .triggerSavepoint(savepointDir, SavepointFormatType.CANONICAL) + .get(2, TimeUnit.MINUTES); + } catch (Exception e) { + lastEx = e; + Thread.sleep(200); + } + } + throw new RuntimeException("Could not trigger savepoint after 30 attempts", lastEx); + } finally { + try { + jobClient.cancel().get(10, TimeUnit.SECONDS); + } catch (Exception ignored) { + } + } + } + + /** Deletes any existing {@code savepoint-*} directories inside the given directory. */ + private static void deleteExistingSavepoints(Path dir) throws IOException { + if (!Files.isDirectory(dir)) { + return; + } + try (var stream = Files.list(dir)) { + stream.filter(p -> p.getFileName().toString().startsWith("savepoint-")) + .filter(Files::isDirectory) + .forEach( + p -> { + try { + deleteDirectory(p); + } catch (IOException e) { + throw new RuntimeException("Failed to delete " + p, e); + } + }); + } + } + + private static void deleteDirectory(Path dir) throws IOException { + Files.walkFileTree( + dir, + new SimpleFileVisitor<>() { + @Override + public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) + throws IOException { + Files.delete(file); + return FileVisitResult.CONTINUE; + } + + @Override + public FileVisitResult postVisitDirectory(Path d, IOException exc) + throws IOException { + Files.delete(d); + return FileVisitResult.CONTINUE; + } + }); + } + + // ------------------------------------------------------------------------- + // Source + // ------------------------------------------------------------------------- + + private static class BoundedWaitingSource extends RichSourceFunction { + + private final int[] elements; + private volatile boolean running = true; + + BoundedWaitingSource(int[] elements) { + this.elements = elements; + } + + @Override + public void run(SourceContext ctx) throws Exception { + for (int e : elements) { + ctx.collect(e); + } + while (running) { + Thread.sleep(50); + } + } + + @Override + public void cancel() { + running = false; + } + } + + // ------------------------------------------------------------------------- + // Operators + // ------------------------------------------------------------------------- + + private static class IntCountOperator extends KeyedProcessFunction { + + private transient ValueState count; + + @Override + public void open(OpenContext ctx) throws Exception { + count = + getRuntimeContext() + .getState(new ValueStateDescriptor<>("count", Integer.class)); + } + + @Override + public void processElement(Integer value, Context ctx, Collector out) + throws Exception { + Integer c = count.value(); + count.update(c == null ? 1 : c + 1); + out.collect(value); + } + } + + private static class ProfileOperator extends KeyedProcessFunction { + + private transient ValueState profile; + + @Override + public void open(OpenContext ctx) throws Exception { + profile = + getRuntimeContext() + .getState(new ValueStateDescriptor<>("profile", PersonProfile.class)); + } + + @Override + public void processElement(Integer value, Context ctx, Collector out) + throws Exception { + profile.update(new PersonProfile("name-" + value, value * 10L)); + out.collect(value); + } + } + + private static class AvroSpecificOperator + extends KeyedProcessFunction { + + private transient ValueState avroSpecific; + + @Override + public void open(OpenContext ctx) throws Exception { + avroSpecific = + getRuntimeContext() + .getState( + new ValueStateDescriptor<>( + "avro_specific", + new AvroTypeInfo<>(StateTestRecord.class))); + } + + @Override + public void processElement(Integer value, Context ctx, Collector out) + throws Exception { + var record = new StateTestRecord(); + record.setName("avro-specific-" + value); + record.setValue((long) value); + avroSpecific.update(record); + out.collect(value); + } + } + + private static class AvroGenericOperator + extends KeyedProcessFunction { + + private transient ValueState avroGeneric; + + @Override + public void open(OpenContext ctx) throws Exception { + avroGeneric = + getRuntimeContext() + .getState( + new ValueStateDescriptor<>( + "avro_generic", + new GenericRecordAvroTypeInfo( + StateTestRecord.getClassSchema()))); + } + + @Override + public void processElement(Integer value, Context ctx, Collector out) + throws Exception { + var record = new GenericData.Record(StateTestRecord.getClassSchema()); + record.put("name", "avro-generic-" + value); + record.put("value", (long) value); + avroGeneric.update(record); + out.collect(value); + } + } + + // ------------------------------------------------------------------------- + // POJO state type + // ------------------------------------------------------------------------- + + public static class PersonProfile { + public String name; + public long score; + + public PersonProfile() {} + + public PersonProfile(String name, long score) { + this.name = name; + this.score = score; + } + + @Override + public boolean equals(Object o) { + if (!(o instanceof PersonProfile)) { + return false; + } + var other = (PersonProfile) o; + return Objects.equals(name, other.name) && score == other.score; + } + + @Override + public int hashCode() { + return Objects.hash(name, score); + } + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/resources/generator/KeyedStatePojoAvroKeySavepointGenerator.java b/flink-libraries/flink-state-processing-api/src/test/resources/generator/KeyedStatePojoAvroKeySavepointGenerator.java new file mode 100644 index 00000000000000..4bff1559a83c2f --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/resources/generator/KeyedStatePojoAvroKeySavepointGenerator.java @@ -0,0 +1,321 @@ +/* + * 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.catalog; + +import org.apache.flink.api.common.JobStatus; +import org.apache.flink.api.common.functions.OpenContext; +import org.apache.flink.api.common.state.ValueState; +import org.apache.flink.api.common.state.ValueStateDescriptor; +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.StateBackendOptions; +import org.apache.flink.core.execution.JobClient; +import org.apache.flink.core.execution.SavepointFormatType; +import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration; +import org.apache.flink.state.catalog.avro.StateTestRecord; +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.api.functions.source.legacy.RichSourceFunction; +import org.apache.flink.test.util.MiniClusterWithClientResource; +import org.apache.flink.util.Collector; + +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.file.FileVisitResult; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.nio.file.SimpleFileVisitor; +import java.nio.file.attribute.BasicFileAttributes; +import java.util.Objects; +import java.util.concurrent.TimeUnit; + +/** + * Generates the pre-built savepoint used by {@code + * StateCatalogGeneratedSavepointITCase.OffClasspathKeyTypes} for POJO-key and Avro-specific-key + * scenarios. + * + *

This file lives in {@code src/test/java/} only during savepoint generation. After generation, + * move it to {@code src/test/resources/generator/} so it is NOT compiled as part of the normal + * build. The generated savepoint is committed to {@code + * src/test/resources/keyed-state-pojo-avro-key/}. + * + *

To regenerate the savepoint (e.g. after a Flink serializer format change): + * + *

    + *
  1. Copy this file to {@code src/test/java/org/apache/flink/state/catalog/} if not already + * there. + *
  2. Copy {@code src/test/resources/generator/StateTestRecord.avsc} to {@code + * src/test/resources/avro/StateTestRecord.avsc} so the Avro class is generated. + *
  3. Remove the {@code @Disabled} annotation and run: {@code ./mvnw test -pl + * flink-libraries/flink-state-processing-api + * -Dtest=KeyedStatePojoAvroKeySavepointGenerator#generateSavepoint} + *
  4. Verify the savepoint under {@code src/test/resources/keyed-state-pojo-avro-key/}. + *
  5. Remove the copied {@code StateTestRecord.avsc} from {@code src/test/resources/avro/} and + * move this {@code .java} file back to {@code resources/generator/}. + *
+ */ +@Disabled("Run manually to regenerate the pre-built savepoint in test resources") +class KeyedStatePojoAvroKeySavepointGenerator { + + /** Operator UIDs referenced by {@code StateCatalogGeneratedSavepointITCase.OffClasspathKeyTypes}. */ + static final String UID_POJO_KEY = "pojo-key-state-op"; + + static final String UID_AVRO_SPECIFIC_KEY = "avro-specific-key-state-op"; + + private static final String RESOURCES_DIR = "src/test/resources/keyed-state-pojo-avro-key"; + + @Test + void generateSavepoint() throws Exception { + Path outputDir = Paths.get(RESOURCES_DIR).toAbsolutePath(); + Files.createDirectories(outputDir); + deleteExistingSavepoints(outputDir); + + var cluster = + new MiniClusterWithClientResource( + new MiniClusterResourceConfiguration.Builder() + .setNumberSlotsPerTaskManager(4) + .build()); + cluster.before(); + try { + Configuration cfg = new Configuration(); + cfg.set(StateBackendOptions.STATE_BACKEND, "hashmap"); + var env = StreamExecutionEnvironment.getExecutionEnvironment(cfg); + env.setParallelism(2); + + // Pipeline 1: POJO key (PersonKey{id,name}). + // PersonKey is defined only in this file → NOT on the classpath during normal test + // runs, which exercises the off-classpath POJO-key reading path. + env.addSource(new BoundedWaitingSource(new int[] {1, 2, 3, 4, 5})) + .returns(Types.INT) + .keyBy(v -> new PersonKey(v, "name-" + v)) + .process(new PojoKeyCountOperator()) + .uid(UID_POJO_KEY) + .name(UID_POJO_KEY) + .sinkTo(new DiscardingSink<>()); + + // Pipeline 2: Avro-specific key (StateTestRecord). + // StateTestRecord is generated from StateTestRecord.avsc only when the .avsc file is + // present in src/test/resources/avro/. After generation that file is removed, so the + // class stays off the classpath during tests, exercising the Avro fallback path. + env.addSource(new BoundedWaitingSource(new int[] {1, 2, 3, 4, 5})) + .returns(Types.INT) + .keyBy( + (KeySelector) + v -> { + StateTestRecord r = new StateTestRecord(); + r.setName("key-" + v); + r.setValue((long) v); + return r; + }) + .process(new AvroSpecificKeyCountOperator()) + .uid(UID_AVRO_SPECIFIC_KEY) + .name(UID_AVRO_SPECIFIC_KEY) + .sinkTo(new DiscardingSink<>()); + + String savepointPath = takeSavepoint(env, outputDir.toString()); + System.out.println("Savepoint written to: " + savepointPath); + } finally { + cluster.after(); + } + } + + private static String takeSavepoint(StreamExecutionEnvironment env, String savepointDir) + throws Exception { + JobClient jobClient = env.executeAsync(); + try { + while (jobClient.getJobStatus().get() != JobStatus.RUNNING) { + Thread.sleep(100); + } + Exception lastEx = null; + for (int attempt = 0; attempt < 30; attempt++) { + try { + return jobClient + .triggerSavepoint(savepointDir, SavepointFormatType.CANONICAL) + .get(2, TimeUnit.MINUTES); + } catch (Exception e) { + lastEx = e; + Thread.sleep(200); + } + } + throw new RuntimeException("Could not trigger savepoint after 30 attempts", lastEx); + } finally { + try { + jobClient.cancel().get(10, TimeUnit.SECONDS); + } catch (Exception ignored) { + } + } + } + + private static void deleteExistingSavepoints(Path dir) throws IOException { + if (!Files.isDirectory(dir)) { + return; + } + try (var stream = Files.list(dir)) { + stream.filter(p -> p.getFileName().toString().startsWith("savepoint-")) + .filter(Files::isDirectory) + .forEach( + p -> { + try { + deleteDirectory(p); + } catch (IOException e) { + throw new RuntimeException("Failed to delete " + p, e); + } + }); + } + } + + private static void deleteDirectory(Path dir) throws IOException { + Files.walkFileTree( + dir, + new SimpleFileVisitor<>() { + @Override + public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) + throws IOException { + Files.delete(file); + return FileVisitResult.CONTINUE; + } + + @Override + public FileVisitResult postVisitDirectory(Path d, IOException exc) + throws IOException { + Files.delete(d); + return FileVisitResult.CONTINUE; + } + }); + } + + // ------------------------------------------------------------------------- + // Source + // ------------------------------------------------------------------------- + + private static class BoundedWaitingSource extends RichSourceFunction { + + private final int[] elements; + private volatile boolean running = true; + + BoundedWaitingSource(int[] elements) { + this.elements = elements; + } + + @Override + public void run(SourceContext ctx) throws Exception { + for (int e : elements) { + ctx.collect(e); + } + while (running) { + Thread.sleep(50); + } + } + + @Override + public void cancel() { + running = false; + } + } + + // ------------------------------------------------------------------------- + // Operators + // ------------------------------------------------------------------------- + + /** Counts elements per PersonKey key. Simple Integer value state avoids Avro complications. */ + private static class PojoKeyCountOperator + extends KeyedProcessFunction { + + private transient ValueState count; + + @Override + public void open(OpenContext ctx) throws Exception { + count = + getRuntimeContext() + .getState(new ValueStateDescriptor<>("count", Integer.class)); + } + + @Override + public void processElement(Integer value, Context ctx, Collector out) + throws Exception { + Integer c = count.value(); + count.update(c == null ? 1 : c + 1); + out.collect(value); + } + } + + /** Counts elements per StateTestRecord (Avro-specific) key. */ + private static class AvroSpecificKeyCountOperator + extends KeyedProcessFunction { + + private transient ValueState count; + + @Override + public void open(OpenContext ctx) throws Exception { + count = + getRuntimeContext() + .getState(new ValueStateDescriptor<>("count", Integer.class)); + } + + @Override + public void processElement(Integer value, Context ctx, Collector out) + throws Exception { + Integer c = count.value(); + count.update(c == null ? 1 : c + 1); + out.collect(value); + } + } + + // ------------------------------------------------------------------------- + // POJO key type — defined here so it is NOT on the classpath during normal + // test runs (this file must be excluded from the normal build). + // ------------------------------------------------------------------------- + + /** Key POJO for the {@value #UID_POJO_KEY} operator. */ + public static class PersonKey { + public int id; + public String name; + + public PersonKey() {} + + public PersonKey(int id, String name) { + this.id = id; + this.name = name; + } + + @Override + public boolean equals(Object o) { + if (!(o instanceof PersonKey)) { + return false; + } + PersonKey other = (PersonKey) o; + return id == other.id && Objects.equals(name, other.name); + } + + @Override + public int hashCode() { + return Objects.hash(id, name); + } + + @Override + public String toString() { + return "PersonKey{id=" + id + ", name='" + name + "'}"; + } + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/resources/generator/KeyedStateTupleKeySavepointGenerator.java b/flink-libraries/flink-state-processing-api/src/test/resources/generator/KeyedStateTupleKeySavepointGenerator.java new file mode 100644 index 00000000000000..4147924c36ae95 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/resources/generator/KeyedStateTupleKeySavepointGenerator.java @@ -0,0 +1,277 @@ +/* + * 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.catalog; + +import org.apache.flink.api.common.JobStatus; +import org.apache.flink.api.common.functions.OpenContext; +import org.apache.flink.api.common.state.ValueState; +import org.apache.flink.api.common.state.ValueStateDescriptor; +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.StateBackendOptions; +import org.apache.flink.core.execution.JobClient; +import org.apache.flink.core.execution.SavepointFormatType; +import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration; +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.api.functions.source.legacy.RichSourceFunction; +import org.apache.flink.test.util.MiniClusterWithClientResource; +import org.apache.flink.util.Collector; + +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.file.FileVisitResult; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.nio.file.SimpleFileVisitor; +import java.nio.file.attribute.BasicFileAttributes; +import java.util.concurrent.TimeUnit; + +/** + * Generates the pre-built savepoint used by {@code StateCatalogGeneratedSavepointITCase.TupleKeyAndValue} + * for Tuple-key and Tuple-value scenarios (including a Tuple value with a mixed basic + POJO field). + * + *

This file lives in {@code src/test/resources/generator/} so it is NOT compiled as part of the + * normal build. The generated savepoint is committed to {@code + * src/test/resources/keyed-state-tuple-key/}. + * + *

To regenerate the savepoint (e.g. after a Flink serializer format change): + * + *

    + *
  1. Copy this file to {@code src/test/java/org/apache/flink/state/catalog/} if not already + * there. + *
  2. Remove the {@code @Disabled} annotation and run: {@code ./mvnw test -pl + * flink-libraries/flink-state-processing-api + * -Dtest=KeyedStateTupleKeySavepointGenerator#generateSavepoint} + *
  3. Verify the savepoint under {@code src/test/resources/keyed-state-tuple-key/}. + *
  4. Move this {@code .java} file back to {@code resources/generator/}. + *
+ */ +@Disabled("Run manually to regenerate the pre-built savepoint in test resources") +class KeyedStateTupleKeySavepointGenerator { + + /** Operator UIDs referenced by {@code StateCatalogGeneratedSavepointITCase.TupleKeyAndValue}. */ + static final String UID_TUPLE_KEY = "tuple-key-state-op"; + + static final String UID_TUPLE_POJO_VALUE = "tuple-pojo-value-state-op"; + + private static final String RESOURCES_DIR = "src/test/resources/keyed-state-tuple-key"; + + @Test + void generateSavepoint() throws Exception { + Path outputDir = Paths.get(RESOURCES_DIR).toAbsolutePath(); + Files.createDirectories(outputDir); + deleteExistingSavepoints(outputDir); + + var cluster = + new MiniClusterWithClientResource( + new MiniClusterResourceConfiguration.Builder() + .setNumberSlotsPerTaskManager(4) + .build()); + cluster.before(); + try { + Configuration cfg = new Configuration(); + cfg.set(StateBackendOptions.STATE_BACKEND, "hashmap"); + var env = StreamExecutionEnvironment.getExecutionEnvironment(cfg); + env.setParallelism(2); + + // Pipeline 1: Tuple2 key, Integer value count. + // Tests schema discovery and reading of a Tuple key with basic element types. + env.addSource(new BoundedWaitingSource(new int[] {1, 2, 3, 4, 5})) + .returns(Types.INT) + .keyBy( + (KeySelector>) + v -> Tuple2.of(v, "k-" + v), + Types.TUPLE(Types.INT, Types.STRING)) + .process(new TupleKeyCountOperator()) + .uid(UID_TUPLE_KEY) + .name(UID_TUPLE_KEY) + .sinkTo(new DiscardingSink<>()); + + // Pipeline 2: Integer key, Tuple2 value state. + // Tests schema discovery and reading of a Tuple value containing a basic type (Long) + // and a POJO type (TuplePojoField), verifying mixed Tuple element type support. + env.addSource(new BoundedWaitingSource(new int[] {1, 2, 3, 4, 5})) + .returns(Types.INT) + .keyBy(v -> v) + .process(new TuplePojoValueOperator()) + .uid(UID_TUPLE_POJO_VALUE) + .name(UID_TUPLE_POJO_VALUE) + .sinkTo(new DiscardingSink<>()); + + String savepointPath = takeSavepoint(env, outputDir.toString()); + System.out.println("Savepoint written to: " + savepointPath); + } finally { + cluster.after(); + } + } + + private static String takeSavepoint(StreamExecutionEnvironment env, String savepointDir) + throws Exception { + JobClient jobClient = env.executeAsync(); + try { + while (jobClient.getJobStatus().get() != JobStatus.RUNNING) { + Thread.sleep(100); + } + Exception lastEx = null; + for (int attempt = 0; attempt < 30; attempt++) { + try { + return jobClient + .triggerSavepoint(savepointDir, SavepointFormatType.CANONICAL) + .get(2, TimeUnit.MINUTES); + } catch (Exception e) { + lastEx = e; + Thread.sleep(200); + } + } + throw new RuntimeException("Could not trigger savepoint after 30 attempts", lastEx); + } finally { + try { + jobClient.cancel().get(10, TimeUnit.SECONDS); + } catch (Exception ignored) { + } + } + } + + private static void deleteExistingSavepoints(Path dir) throws IOException { + if (!Files.isDirectory(dir)) { + return; + } + try (var stream = Files.list(dir)) { + stream.filter(p -> p.getFileName().toString().startsWith("savepoint-")) + .filter(Files::isDirectory) + .forEach( + p -> { + try { + deleteDirectory(p); + } catch (IOException e) { + throw new RuntimeException("Failed to delete " + p, e); + } + }); + } + } + + private static void deleteDirectory(Path dir) throws IOException { + Files.walkFileTree( + dir, + new SimpleFileVisitor<>() { + @Override + public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) + throws IOException { + Files.delete(file); + return FileVisitResult.CONTINUE; + } + + @Override + public FileVisitResult postVisitDirectory(Path d, IOException exc) + throws IOException { + Files.delete(d); + return FileVisitResult.CONTINUE; + } + }); + } + + // ------------------------------------------------------------------------- + // Source + // ------------------------------------------------------------------------- + + private static class BoundedWaitingSource extends RichSourceFunction { + + private final int[] elements; + private volatile boolean running = true; + + BoundedWaitingSource(int[] elements) { + this.elements = elements; + } + + @Override + public void run(SourceContext ctx) throws Exception { + for (int e : elements) { + ctx.collect(e); + } + while (running) { + Thread.sleep(50); + } + } + + @Override + public void cancel() { + running = false; + } + } + + // ------------------------------------------------------------------------- + // Operators + // ------------------------------------------------------------------------- + + /** Counts elements per {@code Tuple2} key. */ + private static class TupleKeyCountOperator + extends KeyedProcessFunction, Integer, Integer> { + + private transient ValueState count; + + @Override + public void open(OpenContext ctx) throws Exception { + count = + getRuntimeContext() + .getState(new ValueStateDescriptor<>("count", Integer.class)); + } + + @Override + public void processElement(Integer value, Context ctx, Collector out) + throws Exception { + Integer c = count.value(); + count.update(c == null ? 1 : c + 1); + out.collect(value); + } + } + + /** Stores a {@code Tuple2} value per Integer key. */ + private static class TuplePojoValueOperator + extends KeyedProcessFunction { + + private transient ValueState> tupleState; + + @Override + public void open(OpenContext ctx) throws Exception { + tupleState = + getRuntimeContext() + .getState( + new ValueStateDescriptor<>( + "tuple_pojo", + Types.TUPLE( + Types.LONG, Types.POJO(TuplePojoField.class)))); + } + + @Override + public void processElement(Integer value, Context ctx, Collector out) + throws Exception { + tupleState.update( + Tuple2.of( + (long) value * 10, new TuplePojoField("name-" + value, value * 100L))); + out.collect(value); + } + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/resources/generator/StateTestRecord.avsc b/flink-libraries/flink-state-processing-api/src/test/resources/generator/StateTestRecord.avsc new file mode 100644 index 00000000000000..681692f5ca6486 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/resources/generator/StateTestRecord.avsc @@ -0,0 +1,32 @@ +/* + * 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. + */ + +{ + "namespace": "org.apache.flink.state.catalog.avro", + "type": "record", + "name": "StateTestRecord", + "fields": [ + { + "name": "name", + "type": "string" + }, + { + "name": "value", + "type": "long" + } + ] +} diff --git a/flink-libraries/flink-state-processing-api/src/test/resources/keyed-state-catalog/savepoint-01d134-a82d2259b86b/_metadata b/flink-libraries/flink-state-processing-api/src/test/resources/keyed-state-catalog/savepoint-01d134-a82d2259b86b/_metadata new file mode 100644 index 00000000000000..77125045df3fe9 Binary files /dev/null and b/flink-libraries/flink-state-processing-api/src/test/resources/keyed-state-catalog/savepoint-01d134-a82d2259b86b/_metadata differ diff --git a/flink-libraries/flink-state-processing-api/src/test/resources/keyed-state-pojo-avro-key/savepoint-07a18b-da0f2ab0e5e0/_metadata b/flink-libraries/flink-state-processing-api/src/test/resources/keyed-state-pojo-avro-key/savepoint-07a18b-da0f2ab0e5e0/_metadata new file mode 100644 index 00000000000000..cea2d23d2398cc Binary files /dev/null and b/flink-libraries/flink-state-processing-api/src/test/resources/keyed-state-pojo-avro-key/savepoint-07a18b-da0f2ab0e5e0/_metadata differ diff --git a/flink-libraries/flink-state-processing-api/src/test/resources/keyed-state-tuple-key/savepoint-515de8-42f928682f3b/_metadata b/flink-libraries/flink-state-processing-api/src/test/resources/keyed-state-tuple-key/savepoint-515de8-42f928682f3b/_metadata new file mode 100644 index 00000000000000..186a78268ed05d Binary files /dev/null and b/flink-libraries/flink-state-processing-api/src/test/resources/keyed-state-tuple-key/savepoint-515de8-42f928682f3b/_metadata differ diff --git a/flink-libraries/flink-state-processing-api/src/test/resources/table-state-missing-avro/_metadata b/flink-libraries/flink-state-processing-api/src/test/resources/table-state-missing-avro/_metadata new file mode 100644 index 00000000000000..1e780952d66d64 Binary files /dev/null and b/flink-libraries/flink-state-processing-api/src/test/resources/table-state-missing-avro/_metadata differ diff --git a/flink-libraries/flink-state-processing-api/src/test/resources/table-state-missing-class/_metadata b/flink-libraries/flink-state-processing-api/src/test/resources/table-state-missing-class/_metadata new file mode 100644 index 00000000000000..a8295168fda81a Binary files /dev/null and b/flink-libraries/flink-state-processing-api/src/test/resources/table-state-missing-class/_metadata differ diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java index d773fdb4a0af1c..e11ca590ab353a 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java @@ -91,6 +91,22 @@ void applyToAllKeys( */ Stream getKeys(List states, N namespace); + /** + * @return A stream of all keys for the multiple states and a given namespace, paired with the + * key-group each key is actually stored under. Modifications to the states during iterating + * over its keys are not supported. + *

Unlike {@link #getKeys(List, Object)}, the key-group in the returned pair reflects how + * the key was physically partitioned when it was written, rather than being recomputed from + * {@code key.hashCode()}. Use this method instead of {@link #getKeys(List, Object)} + * whenever the returned key's {@code hashCode()} is not guaranteed to match the hash of the + * key that originally wrote the data; callers that need to restore the reading context for + * such a key (e.g. via {@link #setCurrentKeyAndKeyGroup}) must use the returned key-group + * rather than recomputing it from the key. + * @param states State variables for which existing keys will be returned. + * @param namespace Namespace for which existing keys will be returned. + */ + Stream> getKeysAndKeyGroups(List states, N namespace); + /** * @return A stream of all keys for the given state and namespace. Modifications to the state * during iterating over it keys are not supported. Implementations go not make any ordering diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java index ddc3405f18c80e..aec63ee7d0effc 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java @@ -395,6 +395,54 @@ public Stream getKeys(List states, N namespace) { return keyStreams.stream().reduce(Stream.empty(), Stream::concat); } + @SuppressWarnings("unchecked") + @Override + public Stream> getKeysAndKeyGroups(List states, N namespace) { + final List> tables = + states.stream() + .filter(registeredKVStates::containsKey) + .map(s -> (StateTable) registeredKVStates.get(s)) + .collect(Collectors.toList()); + final List>> keyStreams = new ArrayList<>(); + for (int i = 0; i < tables.size(); i++) { + int finalI = i; + Stream> keyStream = + tables.get(i) + .getEntriesWithKeyGroup() + .filter( + entryAndKeyGroup -> { + StateEntry entry = entryAndKeyGroup.f0; + if (!entry.getNamespace().equals(namespace)) { + return false; + } + // Deduplicate keys across tables. The entry is looked up in + // its own key-group instead of one recomputed from the + // key's hashCode(), since a substituted key (e.g. one + // deserialized without its original class) may not + // reproduce it faithfully. + int keyGroup = entryAndKeyGroup.f1; + for (int j = 0; j < finalI; ++j) { + if (tables.get(j) + .getMapForKeyGroup(keyGroup) + .get( + entry.getKey(), + entry.getNamespace()) + != null) { + return false; + } + } + return true; + }) + .map( + entryAndKeyGroup -> + Tuple2.of( + entryAndKeyGroup.f0.getKey(), + entryAndKeyGroup.f1)); + keyStreams.add(keyStream); + } + return keyStreams.stream().reduce(Stream.empty(), Stream::concat); + } + @SuppressWarnings("unchecked") @Override public Stream> getKeysAndNamespaces(String state) { diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/StateTable.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/StateTable.java index d79706ac872ae3..4fc0489cb85d53 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/StateTable.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/StateTable.java @@ -43,6 +43,7 @@ import java.util.Objects; import java.util.Spliterators; import java.util.stream.Collectors; +import java.util.stream.IntStream; import java.util.stream.Stream; import java.util.stream.StreamSupport; @@ -275,6 +276,24 @@ public Stream> getKeysAndNamespaces() { .map(entry -> Tuple2.of(entry.getKey(), entry.getNamespace())); } + /** + * Returns all entries paired with the key-group they are physically stored under, i.e. the + * {@link #keyGroupedStateMaps} slot holding them, rather than a key-group recomputed from the + * key's {@code hashCode()}. + */ + public Stream, Integer>> getEntriesWithKeyGroup() { + final int offset = getKeyGroupOffset(); + return IntStream.range(0, keyGroupedStateMaps.length) + .boxed() + .flatMap( + i -> + StreamSupport.stream( + Spliterators.spliteratorUnknownSize( + keyGroupedStateMaps[i].iterator(), 0), + false) + .map(entry -> Tuple2.of(entry, i + offset))); + } + public StateIncrementalVisitor getStateIncrementalVisitor( int recommendedMaxNumberOfReturnedRecords) { return new StateEntryIterator(recommendedMaxNumberOfReturnedRecords); diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/InternalTimeServiceManagerImpl.java b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/InternalTimeServiceManagerImpl.java index 2b7f31c060afcc..aad9dcc8f6d159 100644 --- a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/InternalTimeServiceManagerImpl.java +++ b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/InternalTimeServiceManagerImpl.java @@ -66,7 +66,7 @@ public class InternalTimeServiceManagerImpl implements InternalTimeServiceMan protected static final Logger LOG = LoggerFactory.getLogger(InternalTimeServiceManagerImpl.class); - @VisibleForTesting static final String TIMER_STATE_PREFIX = "_timer_state"; + public static final String TIMER_STATE_PREFIX = "_timer_state"; @VisibleForTesting static final String PROCESSING_TIMER_PREFIX = TIMER_STATE_PREFIX + "/processing_"; diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionKeyedStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionKeyedStateBackend.java index 1d7d18553ca5e2..457dd413b9fe61 100644 --- a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionKeyedStateBackend.java +++ b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionKeyedStateBackend.java @@ -165,6 +165,15 @@ public Stream getKeys(List states, N namespace) { return Stream.empty(); } + @Override + public Stream> getKeysAndKeyGroups(List states, N namespace) { + LOG.debug("Returning an empty stream in BATCH execution mode in getKeysAndKeyGroups()."); + // We return an empty Stream here. This is correct because the BATCH broadcast operators + // process the broadcast side first, meaning we know that the keyed side will always be + // empty when this is called + return Stream.empty(); + } + @Override public Stream> getKeysAndNamespaces(String state) { LOG.debug("Returning an empty stream in BATCH execution mode in getKeysAndNamespaces()."); diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperator.java b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperator.java index 4bdc82c5e5bc92..f674073c3b6700 100644 --- a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperator.java +++ b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperator.java @@ -105,6 +105,9 @@ public class WindowOperator private static final long serialVersionUID = 1L; + /** Name of the internal list state that stores merging-window bookkeeping. */ + public static final String MERGING_WINDOW_SET_STATE_NAME = "merging-window-set"; + // ------------------------------------------------------------------------ // Configuration values and user functions // ------------------------------------------------------------------------ @@ -269,7 +272,7 @@ public long getCurrentProcessingTime() { typedTuple, new TypeSerializer[] {windowSerializer, windowSerializer}); final ListStateDescriptor> mergingSetsStateDescriptor = - new ListStateDescriptor<>("merging-window-set", tupleSerializer); + new ListStateDescriptor<>(MERGING_WINDOW_SET_STATE_NAME, tupleSerializer); // get the state that stores the merging sets mergingSetsState = diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java index 98e467961fd572..3a243c0a99acc6 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java @@ -508,6 +508,50 @@ void testGetKeys() throws Exception { } } + @TestTemplate + void testGetKeysAndKeyGroups() throws Exception { + final int elementsNum = 1000; + final int numberOfKeyGroups = 10; + String fieldName = "get-keys-and-key-groups-test"; + CheckpointableKeyedStateBackend backend = + createKeyedBackend(IntSerializer.INSTANCE); + try { + final String ns = "ns1"; + ValueState keyedState = + backend.getPartitionedState( + ns, + StringSerializer.INSTANCE, + new ValueStateDescriptor<>(fieldName, IntSerializer.INSTANCE)); + + for (int key = 0; key < elementsNum; key++) { + backend.setCurrentKey(key); + keyedState.update(key * 2); + } + + try (Stream> stream = + backend.getKeysAndKeyGroups(Collections.singletonList(fieldName), ns)) { + final Set seenKeys = new HashSet<>(); + stream.forEach( + entry -> { + assertThat(seenKeys.add(entry.f0)) + .withFailMessage("Duplicate key") + .isTrue(); + assertThat(entry.f1) + .withFailMessage("Unexpected key-group") + .isEqualTo( + KeyGroupRangeAssignment.assignToKeyGroup( + entry.f0, numberOfKeyGroups)); + }); + assertThat(seenKeys.size()) + .withFailMessage("Unexpected keys count") + .isEqualTo(elementsNum); + } + } finally { + IOUtils.closeQuietly(backend); + backend.dispose(); + } + } + @TestTemplate void testGetKeysAndNamespaces() throws Exception { final int elementsNum = 1000; diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestUtils.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestUtils.java index 8fdaf601bf4b0d..549ba20a77c0e6 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestUtils.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestUtils.java @@ -282,6 +282,12 @@ public Stream getKeys(List states, N namespace) { return delegatedKeyedStateBackend.getKeys(states, namespace); } + @Override + public Stream> getKeysAndKeyGroups( + List states, N namespace) { + return delegatedKeyedStateBackend.getKeysAndKeyGroups(states, namespace); + } + @Override public Stream> getKeysAndNamespaces(String state) { return delegatedKeyedStateBackend.getKeysAndNamespaces(state); diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/ttl/mock/MockKeyedStateBackend.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/ttl/mock/MockKeyedStateBackend.java index 18cfa86c1ca65c..e71b55445fee9a 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/ttl/mock/MockKeyedStateBackend.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/ttl/mock/MockKeyedStateBackend.java @@ -33,6 +33,7 @@ import org.apache.flink.runtime.state.InternalKeyContext; import org.apache.flink.runtime.state.KeyExtractorFunction; import org.apache.flink.runtime.state.KeyGroupRange; +import org.apache.flink.runtime.state.KeyGroupRangeAssignment; import org.apache.flink.runtime.state.KeyGroupedInternalPriorityQueue; import org.apache.flink.runtime.state.Keyed; import org.apache.flink.runtime.state.KeyedStateHandle; @@ -222,6 +223,17 @@ public Stream getKeys(List states, N namespace) { .stream(); } + @Override + public Stream> getKeysAndKeyGroups(List states, N namespace) { + return getKeys(states, namespace) + .map( + key -> + Tuple2.of( + key, + KeyGroupRangeAssignment.assignToKeyGroup( + key, getNumberOfKeyGroups()))); + } + @Override @SuppressWarnings("unchecked") public Stream> getKeysAndNamespaces(String state) { diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/TestStateBackend.java b/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/TestStateBackend.java index c8ade06c8889b4..1b484dfeb980b4 100644 --- a/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/TestStateBackend.java +++ b/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/TestStateBackend.java @@ -171,6 +171,12 @@ public Stream getKeys(List states, N namespace) { throw new UnsupportedOperationException(); } + @Override + public Stream> getKeysAndKeyGroups( + List states, N namespace) { + throw new UnsupportedOperationException(); + } + @Override public Stream> getKeysAndNamespaces(String state) { throw new UnsupportedOperationException(); diff --git a/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/ChangelogKeyedStateBackend.java b/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/ChangelogKeyedStateBackend.java index 4182023900b37a..621aefb9f8b6e8 100644 --- a/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/ChangelogKeyedStateBackend.java +++ b/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/ChangelogKeyedStateBackend.java @@ -319,6 +319,11 @@ public Stream getKeys(List states, N namespace) { return keyedStateBackend.getKeys(states, namespace); } + @Override + public Stream> getKeysAndKeyGroups(List states, N namespace) { + return keyedStateBackend.getKeysAndKeyGroups(states, namespace); + } + @Override public Stream> getKeysAndNamespaces(String state) { return keyedStateBackend.getKeysAndNamespaces(state); diff --git a/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/restore/ChangelogMigrationRestoreTarget.java b/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/restore/ChangelogMigrationRestoreTarget.java index 8814e61a81a28c..ba02e93d2210b1 100644 --- a/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/restore/ChangelogMigrationRestoreTarget.java +++ b/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/restore/ChangelogMigrationRestoreTarget.java @@ -238,6 +238,12 @@ public Stream> getKeysAndNamespaces(String state) { return keyedStateBackend.getKeysAndNamespaces(state); } + @Override + public Stream> getKeysAndKeyGroups( + List states, N namespace) { + return keyedStateBackend.getKeysAndKeyGroups(states, namespace); + } + @Nonnull @Override public IS createOrUpdateInternalState( diff --git a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/sync/ForStMultiStateKeysIterator.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/sync/ForStMultiStateKeysIterator.java new file mode 100644 index 00000000000000..1fd0579ce90cf0 --- /dev/null +++ b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/sync/ForStMultiStateKeysIterator.java @@ -0,0 +1,211 @@ +/* + * 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.forst.sync; + +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.core.memory.DataInputDeserializer; +import org.apache.flink.runtime.state.CompositeKeySerializationUtils; +import org.apache.flink.util.FlinkRuntimeException; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Iterator; +import java.util.List; +import java.util.NoSuchElementException; +import java.util.Objects; + +/** + * Adapter class to bridge between {@link ForStIteratorWrapper} and {@link Iterator} to iterate over + * the keys of multiple states in namespace order. This class is not thread safe. + * + * @param the type of the iterated objects, which are keys in ForSt. + */ +public class ForStMultiStateKeysIterator implements AutoCloseable, Iterator { + + private final List iterators; + private final List states; + private final TypeSerializer keySerializer; + private final List ambiguousKeyPossibles; + private final int keyGroupPrefixBytes; + private final byte[] namespaceBytes; + private final DataInputDeserializer byteArrayDataInputView; + + private final byte[][] iteratorKeys; + private final int[] iteratorKeysToRemove; + private K previousKey; + private K nextKey; + private byte[] nextKeyBytes; + + public ForStMultiStateKeysIterator( + List iterators, + List states, + @Nonnull TypeSerializer keySerializer, + int keyGroupPrefixBytes, + List ambiguousKeyPossibles, + @Nonnull byte[] namespaceBytes) { + this.iterators = iterators; + this.states = states; + this.keySerializer = keySerializer; + this.ambiguousKeyPossibles = ambiguousKeyPossibles; + this.keyGroupPrefixBytes = keyGroupPrefixBytes; + this.namespaceBytes = namespaceBytes; + this.byteArrayDataInputView = new DataInputDeserializer(); + this.iteratorKeys = new byte[iterators.size()][]; + Arrays.fill(iteratorKeys, null); + this.iteratorKeysToRemove = new int[iterators.size()]; + Arrays.fill(iteratorKeysToRemove, -1); + this.previousKey = null; + this.nextKey = null; + } + + @Override + public boolean hasNext() { + try { + while (nextKey == null && hasDataToProcess()) { + pullKeysFromIterators(); + K smallestIteratorKey = calculateSmallestKeyFromLocalData(); + if (smallestIteratorKey != null) { + previousKey = smallestIteratorKey; + nextKey = smallestIteratorKey; + } + } + } catch (Exception e) { + throw new FlinkRuntimeException( + "Failed to access states [" + String.join(",", states) + "]", e); + } + return nextKey != null; + } + + private boolean hasDataToProcess() { + boolean result = iterators.stream().anyMatch(ForStIteratorWrapper::isValid); + if (!result) { + for (int i = 0; i < iterators.size(); ++i) { + if (iteratorKeys[i] != null) { + result = true; + break; + } + } + } + return result; + } + + private void pullKeysFromIterators() { + for (int i = 0; i < iterators.size(); ++i) { + ForStIteratorWrapper iterator = iterators.get(i); + if (iteratorKeys[i] == null && iterator.isValid()) { + iteratorKeys[i] = iterator.key(); + iterator.next(); + } + } + } + + @Nullable + private K calculateSmallestKeyFromLocalData() throws IOException { + int smallestIteratorKeyIndex = -1; + byte[] smallestIteratorKey = null; + int iteratorKeysToRemoveIndex = 0; + for (int i = 0; i < iteratorKeys.length; ++i) { + byte[] iteratorKey = iteratorKeys[i]; + if (iteratorKey != null) { + boolean update = smallestIteratorKey == null; + if (!update) { + int cmp = Arrays.compare(iteratorKey, smallestIteratorKey); + if (cmp < 0) { + update = true; + } else if (cmp == 0) { + iteratorKeysToRemove[iteratorKeysToRemoveIndex++] = i; + } + } + + if (update) { + smallestIteratorKeyIndex = i; + smallestIteratorKey = iteratorKey; + Arrays.fill(iteratorKeysToRemove, -1); + iteratorKeysToRemoveIndex = 0; + iteratorKeysToRemove[iteratorKeysToRemoveIndex++] = i; + } + } + } + + if (smallestIteratorKey != null) { + for (int i = 0; i < iteratorKeysToRemoveIndex; ++i) { + iteratorKeys[iteratorKeysToRemove[i]] = null; + } + byteArrayDataInputView.setBuffer( + smallestIteratorKey, + keyGroupPrefixBytes, + smallestIteratorKey.length - keyGroupPrefixBytes); + final K smallestIteratorKeyValue = + CompositeKeySerializationUtils.readKey( + keySerializer, + byteArrayDataInputView, + ambiguousKeyPossibles.get(smallestIteratorKeyIndex)); + if (isMatchingNameSpace( + smallestIteratorKey, + byteArrayDataInputView.getPosition(), + namespaceBytes) + && !Objects.equals(previousKey, smallestIteratorKeyValue)) { + nextKeyBytes = smallestIteratorKey; + return smallestIteratorKeyValue; + } + } + + return null; + } + + private static boolean isMatchingNameSpace( + @Nonnull byte[] key, int namespaceBytesStartPos, @Nonnull byte[] namespaceBytes) { + if (key.length >= namespaceBytesStartPos + namespaceBytes.length) { + for (int i = 0; i < namespaceBytes.length; ++i) { + if (key[namespaceBytesStartPos + i] != namespaceBytes[i]) { + return false; + } + } + return true; + } + return false; + } + + @Override + public K next() { + if (!hasNext()) { + throw new NoSuchElementException( + "Failed to access states [" + String.join(",", states) + "]"); + } + + K tmpKey = nextKey; + nextKey = null; + return tmpKey; + } + + /** Returns the key-group of the key most recently returned by {@link #next()}. */ + public int getKeyGroup() { + return CompositeKeySerializationUtils.extractKeyGroup(keyGroupPrefixBytes, nextKeyBytes); + } + + @Override + public void close() { + for (ForStIteratorWrapper iterator : iterators) { + iterator.close(); + } + } +} diff --git a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/sync/ForStSyncKeyedStateBackend.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/sync/ForStSyncKeyedStateBackend.java index 4f410ca4287ae9..97ab1dfbc1bccc 100644 --- a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/sync/ForStSyncKeyedStateBackend.java +++ b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/sync/ForStSyncKeyedStateBackend.java @@ -88,7 +88,9 @@ import java.io.Closeable; import java.io.IOException; import java.util.ArrayList; +import java.util.Arrays; import java.util.HashMap; +import java.util.Iterator; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; @@ -383,8 +385,98 @@ public Stream getKeys(String state, N namespace) { @Override public Stream getKeys(List states, N namespace) { - // TODO - throw new UnsupportedOperationException(); + final ForStMultiStateKeysIterator iteratorWrapper = + openMultiStateKeysIterator(states, namespace); + Stream targetStream = + StreamSupport.stream( + Spliterators.spliteratorUnknownSize(iteratorWrapper, Spliterator.ORDERED), + false); + return targetStream.onClose(iteratorWrapper::close); + } + + @Override + public Stream> getKeysAndKeyGroups(List states, N namespace) { + final ForStMultiStateKeysIterator iteratorWrapper = + openMultiStateKeysIterator(states, namespace); + Iterator> keyAndKeyGroupIterator = + new Iterator>() { + @Override + public boolean hasNext() { + return iteratorWrapper.hasNext(); + } + + @Override + public Tuple2 next() { + K key = iteratorWrapper.next(); + return Tuple2.of(key, iteratorWrapper.getKeyGroup()); + } + }; + Stream> targetStream = + StreamSupport.stream( + Spliterators.spliteratorUnknownSize( + keyAndKeyGroupIterator, Spliterator.ORDERED), + false); + return targetStream.onClose(iteratorWrapper::close); + } + + @SuppressWarnings("unchecked") + private ForStMultiStateKeysIterator openMultiStateKeysIterator( + List states, N namespace) { + final List ambiguousKeyPossibles = new ArrayList<>(); + final List iterators = new ArrayList<>(); + byte[] namespaceBytes = null; + + for (String state : states) { + ForStOperationUtils.ForStKvStateInfo columnInfo = kvStateInformation.get(state); + if (columnInfo == null + || !(columnInfo.metaInfo instanceof RegisteredKeyValueStateBackendMetaInfo)) { + continue; + } + + RegisteredKeyValueStateBackendMetaInfo registeredKeyValueStateBackendMetaInfo = + (RegisteredKeyValueStateBackendMetaInfo) columnInfo.metaInfo; + + final TypeSerializer namespaceSerializer = + registeredKeyValueStateBackendMetaInfo.getNamespaceSerializer(); + final DataOutputSerializer namespaceOutputView = new DataOutputSerializer(8); + boolean ambiguousKeyPossible = + CompositeKeySerializationUtils.isAmbiguousKeyPossible( + getKeySerializer(), namespaceSerializer); + ambiguousKeyPossibles.add(ambiguousKeyPossible); + try { + CompositeKeySerializationUtils.writeNameSpace( + namespace, namespaceSerializer, namespaceOutputView, ambiguousKeyPossible); + final byte[] stateNamespaceBytes = namespaceOutputView.getCopyOfBuffer(); + if (namespaceBytes == null) { + namespaceBytes = stateNamespaceBytes; + } else { + if (!Arrays.equals(namespaceBytes, stateNamespaceBytes)) { + throw new FlinkRuntimeException( + "Key namespaces are different for states [" + + String.join(",", states) + + "]"); + } + } + } catch (IOException ex) { + throw new FlinkRuntimeException( + "Failed to get keys from ForSt sync state backend.", ex); + } + + ForStIteratorWrapper iterator = + ForStOperationUtils.getForStIterator( + db, columnInfo.columnFamilyHandle, readOptions); + iterator.seekToFirst(); + iterators.add(iterator); + } + + Preconditions.checkNotNull(namespaceBytes, "Namespace must exist"); + return new ForStMultiStateKeysIterator<>( + iterators, + states, + getKeySerializer(), + keyGroupPrefixBytes, + ambiguousKeyPossibles, + namespaceBytes); } @Override diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/RocksDBKeyedStateBackend.java b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/RocksDBKeyedStateBackend.java index b77232cb067b3c..731425713732ab 100644 --- a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/RocksDBKeyedStateBackend.java +++ b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/RocksDBKeyedStateBackend.java @@ -95,6 +95,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.HashMap; +import java.util.Iterator; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; @@ -406,9 +407,45 @@ public Stream getKeys(String state, N namespace) { return targetStream.onClose(iteratorWrapper::close); } - @SuppressWarnings("unchecked") @Override public Stream getKeys(List states, N namespace) { + final RocksMultiStateKeysIterator iteratorWrapper = + openMultiStateKeysIterator(states, namespace); + Stream targetStream = + StreamSupport.stream( + Spliterators.spliteratorUnknownSize(iteratorWrapper, Spliterator.ORDERED), + false); + return targetStream.onClose(iteratorWrapper::close); + } + + @Override + public Stream> getKeysAndKeyGroups(List states, N namespace) { + final RocksMultiStateKeysIterator iteratorWrapper = + openMultiStateKeysIterator(states, namespace); + Iterator> keyAndKeyGroupIterator = + new Iterator>() { + @Override + public boolean hasNext() { + return iteratorWrapper.hasNext(); + } + + @Override + public Tuple2 next() { + K key = iteratorWrapper.next(); + return Tuple2.of(key, iteratorWrapper.getKeyGroup()); + } + }; + Stream> targetStream = + StreamSupport.stream( + Spliterators.spliteratorUnknownSize( + keyAndKeyGroupIterator, Spliterator.ORDERED), + false); + return targetStream.onClose(iteratorWrapper::close); + } + + @SuppressWarnings("unchecked") + private RocksMultiStateKeysIterator openMultiStateKeysIterator( + List states, N namespace) { final List ambiguousKeyPossibles = new ArrayList<>(); final List iterators = new ArrayList<>(); byte[] namespaceBytes = null; @@ -457,20 +494,13 @@ public Stream getKeys(List states, N namespace) { } Preconditions.checkNotNull(namespaceBytes, "Namespace must exist"); - final RocksMultiStateKeysIterator iteratorWrapper = - new RocksMultiStateKeysIterator<>( - iterators, - states, - getKeySerializer(), - keyGroupPrefixBytes, - ambiguousKeyPossibles, - namespaceBytes); - - Stream targetStream = - StreamSupport.stream( - Spliterators.spliteratorUnknownSize(iteratorWrapper, Spliterator.ORDERED), - false); - return targetStream.onClose(iteratorWrapper::close); + return new RocksMultiStateKeysIterator<>( + iterators, + states, + getKeySerializer(), + keyGroupPrefixBytes, + ambiguousKeyPossibles, + namespaceBytes); } @Override diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/iterator/RocksMultiStateKeysIterator.java b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/iterator/RocksMultiStateKeysIterator.java index fd6ec728f77473..ec0de632dd511f 100644 --- a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/iterator/RocksMultiStateKeysIterator.java +++ b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/iterator/RocksMultiStateKeysIterator.java @@ -56,6 +56,7 @@ public class RocksMultiStateKeysIterator implements AutoCloseable, Iterator iterators, @@ -166,6 +167,7 @@ private K calculateSmallestKeyFromLocalData() throws IOException { byteArrayDataInputView.getPosition(), namespaceBytes) && !Objects.equals(previousKey, smallestIteratorKeyValue)) { + nextKeyBytes = smallestIteratorKey; return smallestIteratorKeyValue; } } @@ -185,6 +187,16 @@ public K next() { return tmpKey; } + /** + * Returns the key-group of the key most recently returned by {@link #next()}, decoded from that + * key's RocksDB byte prefix rather than recomputed from {@code hashCode()}. Only valid + * immediately after a call to {@link #next()}; the decoding happens here so that callers which + * never need the key-group pay nothing for it. + */ + public int getKeyGroup() { + return CompositeKeySerializationUtils.extractKeyGroup(keyGroupPrefixBytes, nextKeyBytes); + } + @Override public void close() { for (RocksIteratorWrapper iterator : iterators) { diff --git a/flink-table/flink-table-type-utils/src/main/java/org/apache/flink/table/runtime/typeutils/ExternalTypeInfo.java b/flink-table/flink-table-type-utils/src/main/java/org/apache/flink/table/runtime/typeutils/ExternalTypeInfo.java index e3b98f27773699..49ebf02126d00d 100644 --- a/flink-table/flink-table-type-utils/src/main/java/org/apache/flink/table/runtime/typeutils/ExternalTypeInfo.java +++ b/flink-table/flink-table-type-utils/src/main/java/org/apache/flink/table/runtime/typeutils/ExternalTypeInfo.java @@ -80,6 +80,16 @@ public static ExternalTypeInfo of(DataType dataType, boolean isInternalIn return new ExternalTypeInfo<>(dataType, serializer); } + /** + * Creates type information for the given {@link DataType} with an explicitly provided + * serializer. Use this when the serializer cannot be derived from the data type alone, e.g. + * when it was restored from a savepoint serializer snapshot. + */ + @SuppressWarnings("unchecked") + public static ExternalTypeInfo of(DataType dataType, TypeSerializer typeSerializer) { + return new ExternalTypeInfo<>(dataType, (TypeSerializer) typeSerializer); + } + @SuppressWarnings("unchecked") private static TypeSerializer createExternalTypeSerializer( DataType dataType, boolean isInternalInput) { diff --git a/flink-table/flink-table-type-utils/src/main/java/org/apache/flink/table/runtime/typeutils/RowDataSerializer.java b/flink-table/flink-table-type-utils/src/main/java/org/apache/flink/table/runtime/typeutils/RowDataSerializer.java index 718fbc20be9209..6e0944975fa718 100644 --- a/flink-table/flink-table-type-utils/src/main/java/org/apache/flink/table/runtime/typeutils/RowDataSerializer.java +++ b/flink-table/flink-table-type-utils/src/main/java/org/apache/flink/table/runtime/typeutils/RowDataSerializer.java @@ -407,5 +407,22 @@ public TypeSerializerSchemaCompatibility resolveSchemaCompatibility( return intermediateResult.getFinalResult(); } + + /** Returns the logical types stored in this snapshot. */ + @Internal + public LogicalType[] getTypes() { + return types; + } + + /** + * Returns the field names stored in this snapshot, in the same order as {@link + * #getTypes()}, or {@code null} if the originating serializer was built from a bare {@link + * LogicalType} array or the snapshot predates field name tracking. + */ + @Internal + @Nullable + public String[] getFieldNames() { + return fieldNames; + } } }