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/java/typeutils/runtime/PojoDeserializerCompatibilitySnapshot.java b/flink-core/src/main/java/org/apache/flink/api/java/typeutils/runtime/PojoDeserializerCompatibilitySnapshot.java new file mode 100644 index 00000000000000..0dace3f1537fd1 --- /dev/null +++ b/flink-core/src/main/java/org/apache/flink/api/java/typeutils/runtime/PojoDeserializerCompatibilitySnapshot.java @@ -0,0 +1,93 @@ +/* + * 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.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.core.memory.DataInputView; +import org.apache.flink.core.memory.DataOutputView; + +import java.io.IOException; + +/** + * A {@link TypeSerializerSnapshot} for deserializers that can read POJO binary data without the + * user POJO class being on the classpath. + * + *

This snapshot declares itself {@link TypeSerializerSchemaCompatibility#compatibleAsIs() + * compatible as-is} with any stored {@link PojoSerializerSnapshot}. It is used by the {@code + * PojoToRowDataDeserializer} in the state processing API (which converts them to {@code + * GenericRowData}). + * + *

This snapshot only ever exists in-memory, wrapping the live deserializer it was created from: + * composite schema compatibility checks (e.g. {@code + * CompositeTypeSerializerSnapshot#resolveOuterSchemaCompatibility}) restore the "new" side of a + * composite serializer (such as a timer serializer nesting the key serializer) to compare it + * against the old one, even when nested-level compatibility already short-circuited to {@link + * TypeSerializerSchemaCompatibility#compatibleAsIs()}. {@link #restoreSerializer()} therefore + * returns the wrapped instance instead of throwing, since it never actually needs to be + * reconstructed from a persisted snapshot. + */ +@Internal +public final class PojoDeserializerCompatibilitySnapshot implements TypeSerializerSnapshot { + + private final TypeSerializer restoredSerializer; + + public PojoDeserializerCompatibilitySnapshot() { + this.restoredSerializer = 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) { + return restoredSerializer; + } + throw new UnsupportedOperationException( + "PojoDeserializerCompatibilitySnapshot cannot reconstruct the deserializer on its own. " + + "Use PojoSerializerSnapshot.restoreSerializer() or " + + "PojoToRowDataDeserializer.create()."); + } + + @Override + public void writeSnapshot(DataOutputView out) throws IOException {} + + @Override + public void readSnapshot(int readVersion, DataInputView in, ClassLoader userCodeClassLoader) + throws IOException {} +} 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..bb78d725f8aca2 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 @@ -54,6 +54,42 @@ public class PojoSerializerSnapshot implements TypeSerializerSnapshot { */ private static final int VERSION = 2; + /** + * Optional thread-local factory for building a custom serializer when the POJO class is not on + * the classpath. + * + *

Used by the State Processing API to inject {@code PojoToRowDataDeserializer} so that the + * heap state backend can restore POJO state without the user class (values are converted to + * {@code GenericRowData} during restore instead of raising {@link ClassNotFoundException}). + * + *

Set via {@link #setPojoRestoreSerializerFactory} before the state backend restore and + * cleared via {@link #clearPojoRestoreSerializerFactory} afterward. + */ + private static final ThreadLocal< + java.util.function.Function, TypeSerializer>> + RESTORE_FACTORY = new ThreadLocal<>(); + + /** + * Registers a factory that provides a fallback {@link TypeSerializer} when the POJO class is + * absent from the classpath. + * + *

The factory receives the snapshot and should return a serializer that can read the POJO + * binary format. + * + *

Must be paired with {@link #clearPojoRestoreSerializerFactory()} in a try-finally block. + */ + @Internal + public static void setPojoRestoreSerializerFactory( + java.util.function.Function, TypeSerializer> factory) { + RESTORE_FACTORY.set(factory); + } + + /** Clears the factory registered via {@link #setPojoRestoreSerializerFactory}. */ + @Internal + public static void clearPojoRestoreSerializerFactory() { + RESTORE_FACTORY.remove(); + } + /** Contains the actual content for the serializer snapshot. */ private PojoSerializerSnapshotData snapshotData; @@ -143,6 +179,15 @@ public void readSnapshot(int readVersion, DataInputView in, ClassLoader userCode @Override @SuppressWarnings("unchecked") public TypeSerializer restoreSerializer() { + if (snapshotData.getPojoClass() == null) { + java.util.function.Function, TypeSerializer> factory = + RESTORE_FACTORY.get(); + if (factory != null) { + return (TypeSerializer) factory.apply(this); + } + throw new RuntimeException(new ClassNotFoundException(snapshotData.getPojoClassName())); + } + final int numFields = snapshotData.getFieldSerializerSnapshots().size(); final ArrayList restoredFields = new ArrayList<>(numFields); @@ -257,6 +302,73 @@ public TypeSerializerSchemaCompatibility resolveSchemaCompatibility( return TypeSerializerSchemaCompatibility.compatibleAsIs(); } + // --------------------------------------------------------------------------------------------- + // Schema extraction support + // --------------------------------------------------------------------------------------------- + + /** + * Returns the raw snapshot data. Exposed for schema extraction utilities that need to inspect + * field names and field serializer snapshots without the user POJO class being on the + * classpath. + * + * @return the snapshot data; field names are always present, field serializer snapshots may be + * null for individual fields if their snapshot could not be read + */ + @Internal + public PojoSerializerSnapshotData getSnapshotData() { + return snapshotData; + } + + /** + * Returns {@code true} if the POJO class can be loaded from the current classloader. + * + *

When {@code false}, {@link #restoreSerializer()} returns a (or a factory-supplied + * deserializer) instead of the normal {@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 TypeSerializerSnapshot) pairs. Field names are + * always present; snapshot values may be {@code null} when the field snapshot could not be + * read. + */ + @Internal + public java.util.List>> + getFieldSnapshotEntries() { + java.util.List>> + result = new java.util.ArrayList<>(); + snapshotData + .getFieldSerializerSnapshots() + .forEach( + (fieldName, field, fieldSnapshot) -> + result.add( + new java.util.AbstractMap.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 java.util.List> getRegisteredSubclassSnapshotsOrdered() { + return new java.util.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..12a3b58807f1f5 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 @@ -24,7 +24,6 @@ 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 +31,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 +114,7 @@ static PojoSerializerSnapshotData createFrom( return new PojoSerializerSnapshotData<>( pojoClass, + pojoClass.getName(), fieldSerializerSnapshots, optionalMapOf(registeredSubclassSerializerSnapshots, Class::getName), optionalMapOf(nonRegisteredSubclassSerializerSnapshots, Class::getName)); @@ -153,12 +155,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 +170,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 +192,7 @@ private PojoSerializerSnapshotData( // --------------------------------------------------------------------------------------------- void writeSnapshotData(DataOutputView out) throws IOException { - out.writeUTF(pojoClass.getName()); + out.writeUTF(pojoClassName); writeOptionalMap( out, fieldSerializerSnapshots, @@ -206,7 +212,17 @@ 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 = null; + try { + @SuppressWarnings("unchecked") + Class resolved = (Class) Class.forName(pojoClassName, false, userCodeClassLoader); + pojoClass = resolved; + } catch (ClassNotFoundException e) { + LOG.debug( + "POJO class '{}' not found on classpath; schema can still be read from field snapshots.", + pojoClassName); + } LinkedOptionalMap> fieldSerializerSnapshots = readOptionalMap( @@ -226,6 +242,7 @@ private static PojoSerializerSnapshotData readSnapshotData( return new PojoSerializerSnapshotData<>( pojoClass, + pojoClassName, fieldSerializerSnapshots, registeredSubclassSerializerSnapshots, nonRegisteredSubclassSerializerSnapshots); @@ -235,10 +252,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/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..f2789979f8e7db --- /dev/null +++ b/flink-core/src/test/java/org/apache/flink/api/java/typeutils/runtime/PojoSerializerSnapshotLenientReadTest.java @@ -0,0 +1,188 @@ +/* + * 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.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.core.memory.DataInputDeserializer; +import org.apache.flink.core.memory.DataOutputSerializer; + +import org.junit.Assert; +import org.junit.Test; + +import java.io.IOException; +import java.util.AbstractMap; +import java.util.List; + +/** + * Tests for the lenient POJO class-loading fix in {@link PojoSerializerSnapshotData}. + * + *

Verifies that a {@link PojoSerializerSnapshot} can be read back even when the POJO class is + * not on the classpath — field names and field serializer snapshots must remain accessible. + */ +public class PojoSerializerSnapshotLenientReadTest { + + // ------------------------------------------------------------------------- + // Simple POJO used for writing snapshots (available during write, absent during read) + // ------------------------------------------------------------------------- + + public static class SomePojo { + public String name; + public int age; + public long score; + } + + // ------------------------------------------------------------------------- + // Tests + // ------------------------------------------------------------------------- + + @Test + public void testReadWithClassPresent() throws IOException { + PojoSerializerSnapshot written = buildSnapshot(SomePojo.class); + PojoSerializerSnapshot read = + roundtripSnapshot(written, getClass().getClassLoader()); + + // Class should be found normally. + Assert.assertNotNull(read.getSnapshotData().getPojoClass()); + Assert.assertTrue( + "Class name should contain SomePojo", + read.getSnapshotData().getPojoClassName().contains("SomePojo")); + + List>> entries = + read.getFieldSnapshotEntries(); + Assert.assertEquals(3, entries.size()); + assertFieldEntry(entries, "name", StringSerializer.StringSerializerSnapshot.class); + assertFieldEntry(entries, "age", IntSerializer.IntSerializerSnapshot.class); + assertFieldEntry(entries, "score", LongSerializer.LongSerializerSnapshot.class); + } + + @Test + public void testReadWithClassAbsent() throws IOException { + PojoSerializerSnapshot written = buildSnapshot(SomePojo.class); + + // Use a classloader that cannot see SomePojo. + ClassLoader noPojoClassLoader = + new ClassLoader(getClass().getClassLoader()) { + @Override + public Class loadClass(String name) throws ClassNotFoundException { + if (name.equals(SomePojo.class.getName())) { + throw new ClassNotFoundException(name); + } + return super.loadClass(name); + } + }; + + PojoSerializerSnapshot read = roundtripSnapshot(written, noPojoClassLoader); + + // Class should be null but className should be preserved. + Assert.assertNull(read.getSnapshotData().getPojoClass()); + Assert.assertTrue(read.getPojoClassName().endsWith("SomePojo")); + + // Field names and snapshots must still be readable. + List>> entries = + read.getFieldSnapshotEntries(); + Assert.assertEquals(3, entries.size()); + assertFieldEntry(entries, "name", StringSerializer.StringSerializerSnapshot.class); + assertFieldEntry(entries, "age", IntSerializer.IntSerializerSnapshot.class); + assertFieldEntry(entries, "score", LongSerializer.LongSerializerSnapshot.class); + } + + @Test + public void testGetPojoClassNameAlwaysSet() throws IOException { + PojoSerializerSnapshot written = buildSnapshot(SomePojo.class); + PojoSerializerSnapshot read = roundtripSnapshot(written, getClass().getClassLoader()); + + // getPojoClassName must always return the class name, regardless of loading. + Assert.assertEquals(SomePojo.class.getName(), read.getPojoClassName()); + } + + @Test + public void testFieldNamePreservedWhenFieldClassAbsent() throws IOException { + // Write a snapshot for a POJO that has a field referencing another POJO class. + // When read back without the field's declaring class, the field name must still + // appear in the entries (key name is always written before the framed value). + PojoSerializerSnapshot written = buildSnapshot(SomePojo.class); + + ClassLoader noFieldLoader = + new ClassLoader(getClass().getClassLoader()) { + @Override + protected Class loadClass(String name, boolean resolve) + throws ClassNotFoundException { + if (name.contains("SomePojo")) { + throw new ClassNotFoundException(name); + } + return super.loadClass(name, resolve); + } + }; + + PojoSerializerSnapshot read = roundtripSnapshot(written, noFieldLoader); + + List>> entries = + read.getFieldSnapshotEntries(); + Assert.assertEquals(3, entries.size()); + // Names must be present even when class is absent (order may vary). + Assert.assertTrue(entries.stream().anyMatch(e -> "name".equals(e.getKey()))); + Assert.assertTrue(entries.stream().anyMatch(e -> "age".equals(e.getKey()))); + Assert.assertTrue(entries.stream().anyMatch(e -> "score".equals(e.getKey()))); + } + + // ------------------------------------------------------------------------- + // Helpers + // ------------------------------------------------------------------------- + + @SuppressWarnings({"unchecked", "rawtypes"}) + private static PojoSerializerSnapshot buildSnapshot(Class clazz) { + return (PojoSerializerSnapshot) + org.apache.flink.api.java.typeutils.TypeExtractor.createTypeInfo(clazz) + .createSerializer(new SerializerConfigImpl()) + .snapshotConfiguration(); + } + + @SuppressWarnings("unchecked") + private static PojoSerializerSnapshot roundtripSnapshot( + PojoSerializerSnapshot snapshot, ClassLoader cl) throws IOException { + DataOutputSerializer out = new DataOutputSerializer(256); + TypeSerializerSnapshot.writeVersionedSnapshot(out, snapshot); + + DataInputDeserializer in = new DataInputDeserializer(out.getSharedBuffer()); + return (PojoSerializerSnapshot) TypeSerializerSnapshot.readVersionedSnapshot(in, cl); + } + + private static void assertFieldEntry( + List>> entries, + String expectedName, + Class expectedSnapshotType) { + for (AbstractMap.SimpleEntry> entry : entries) { + if (expectedName.equals(entry.getKey())) { + Assert.assertNotNull( + "Snapshot for field '" + expectedName + "' should not be null", + entry.getValue()); + Assert.assertSame( + "Wrong snapshot type for field '" + expectedName + "'", + expectedSnapshotType, + entry.getValue().getClass()); + return; + } + } + Assert.fail("Field '" + expectedName + "' not found in snapshot entries"); + } +} 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..c4c8b622d1bab5 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 @@ -32,6 +32,8 @@ import org.apache.avro.reflect.ReflectData; import org.apache.avro.specific.SpecificData; import org.apache.avro.specific.SpecificRecord; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import javax.annotation.Nonnull; @@ -49,6 +51,7 @@ * @param The data type that the originating serializer of this configuration serializes. */ public class AvroSerializerSnapshot implements TypeSerializerSnapshot { + private static final Logger LOG = LoggerFactory.getLogger(AvroSerializerSnapshot.class); private Class runtimeType; private Schema schema; private Schema runtimeSchema; @@ -114,7 +117,8 @@ 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.runtimeType = + findClassOrFallbackToGeneric(userCodeClassLoader, previousRuntimeTypeName); this.schema = parseAvroSchema(previousSchemaDefinition); this.runtimeSchema = tryExtractAvroSchema(userCodeClassLoader, runtimeType); } @@ -123,7 +127,8 @@ private void readV3(DataInputView in, ClassLoader userCodeClassLoader) throws IO final String previousRuntimeTypeName = readString(in); final String previousSchemaDefinition = readString(in); - this.runtimeType = findClassOrThrow(userCodeClassLoader, previousRuntimeTypeName); + this.runtimeType = + findClassOrFallbackToGeneric(userCodeClassLoader, previousRuntimeTypeName); this.schema = parseAvroSchema(previousSchemaDefinition); this.runtimeSchema = tryExtractAvroSchema(userCodeClassLoader, runtimeType); } @@ -161,6 +166,11 @@ public TypeSerializer restoreSerializer() { // Helpers // ------------------------------------------------------------------------------------------------------------ + /** Returns the Avro writer schema stored in this snapshot. */ + public Schema getSchema() { + return schema; + } + /** * Resolves writer/reader schema compatibly. * @@ -252,6 +262,9 @@ private static Class findClassOrFallbackToGeneric( Class runtimeTarget = Class.forName(className, false, userCodeClassLoader); return (Class) runtimeTarget; } catch (ClassNotFoundException e) { + LOG.debug( + "Avro runtime type '{}' not found on classpath; falling back to GenericRecord.", + className); return (Class) GenericRecord.class; } } diff --git a/flink-libraries/flink-state-processing-api/pom.xml b/flink-libraries/flink-state-processing-api/pom.xml index eb027221a0c4f1..b9fff64d9a73bf 100644 --- a/flink-libraries/flink-state-processing-api/pom.xml +++ b/flink-libraries/flink-state-processing-api/pom.xml @@ -89,7 +89,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..0d2e269686bcaf --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/StateTableUtils.java @@ -0,0 +1,635 @@ +/* + * 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.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.warn( + "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.warn( + "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); + 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: + * + *

+ * + * @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) { + // Flink's built-in timer states use this prefix. + // "merging-window-set" is WindowOperator's internal session-window merging bookkeeping + // state, not a user-registered state. + return stateName.startsWith("_timer_state/") || stateName.equals("merging-window-set"); + } + + /** + * 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 (Exception 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/ExecutionConfigs.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/ExecutionConfigs.java new file mode 100644 index 00000000000000..679a8bf1ae8637 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/ExecutionConfigs.java @@ -0,0 +1,45 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.state.api.input; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.api.common.ExecutionConfig; +import org.apache.flink.util.SerializedValue; + +import java.io.IOException; + +/** + * Shared helper for deserializing the {@link ExecutionConfig} that {@link KeyedStateInputFormat} + * and {@link OperatorStateInputFormat} ship to task managers as a {@link SerializedValue}. + */ +@Internal +final class ExecutionConfigs { + + private ExecutionConfigs() {} + + 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/KeyedStateInputFormat.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/KeyedStateInputFormat.java index 00bb5c262e1931..3967812f064e1d 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.java.tuple.Tuple3; +import org.apache.flink.api.java.typeutils.runtime.PojoSerializerSnapshot; 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.PojoToRowDataDeserializer; 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,11 @@ 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 = + ExecutionConfigs.deserialize( + serializedExecutionConfig, runtimeContext.getUserCodeClassLoader()); + + StreamOperatorContextBuilder builder = new StreamOperatorContextBuilder( runtimeContext, configuration, @@ -229,19 +227,35 @@ 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())); + + // Inject PojoToRowDataDeserializer as the restore serializer for any PojoSerializerSnapshot + // whose POJO class is absent from the current classpath. PojoSerializerSnapshot only + // consults the factory when getPojoClass() == null, so setting it unconditionally is safe: + // snapshots whose class is present restore normally and never invoke the factory. + // The factory runs on this task thread, which is the same thread that drives all + // subsequent state backend access, so a ThreadLocal is safe here. The scope must stay + // active beyond open(): the heap backend resolves state schema compatibility eagerly + // during build(), but the RocksDB backend defers it to the first actual state access, + // which happens lazily while processing elements in nextRecord(). The factory is only + // cleared in close(), once no further state access can occur on this thread. + PojoSerializerSnapshot.setPojoRestoreSerializerFactory( + pojoSnapshot -> PojoToRowDataDeserializer.create(pojoSnapshot)); - AbstractKeyedStateBackend keyedStateBackend = - (AbstractKeyedStateBackend) context.keyedStateBackend(); + try { + final StreamOperatorStateContext context = builder.build(LOG); - final DefaultKeyedStateStore keyedStateStore = - new DefaultKeyedStateStore(keyedStateBackend, runtimeContext::createSerializer); - SavepointRuntimeContext ctx = new SavepointRuntimeContext(runtimeContext, keyedStateStore); + AbstractKeyedStateBackend keyedStateBackend = + (AbstractKeyedStateBackend) context.keyedStateBackend(); + + final DefaultKeyedStateStore keyedStateStore = + new DefaultKeyedStateStore(keyedStateBackend, runtimeContext::createSerializer); + SavepointRuntimeContext ctx = + new SavepointRuntimeContext(runtimeContext, keyedStateStore); + + InternalTimeServiceManager timeServiceManager = + (InternalTimeServiceManager) context.internalTimerServiceManager(); - InternalTimeServiceManager timeServiceManager = - (InternalTimeServiceManager) context.internalTimerServiceManager(); - try { operator.setup( runtimeContext::createSerializer, keyedStateBackend, timeServiceManager, ctx); operator.open(); @@ -254,6 +268,7 @@ public void open(KeyGroupRangeInputSplit split) throws IOException { keysAndNamespaces = operator.getKeysAndNamespaces(ctx); } } catch (Exception e) { + PojoSerializerSnapshot.clearPojoRestoreSerializerFactory(); throw new IOException("Failed to restore timer state", e); } } @@ -266,6 +281,8 @@ public void close() throws IOException { IOUtils.closeQuietly(registry); } catch (Exception e) { throw new IOException("Failed to close state backend", e); + } finally { + PojoSerializerSnapshot.clearPojoRestoreSerializerFactory(); } } @@ -280,8 +297,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); 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..23650914d30a15 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; @@ -39,10 +40,10 @@ * @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 +53,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 +73,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..6597e4f38cb1b0 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 = + ExecutionConfigs.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/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..8d4c942e7ce3b6 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/deserializer/InternalTypeConverter.java @@ -0,0 +1,288 @@ +/* + * 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.LogicalType; +import org.apache.flink.table.types.logical.MapType; +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.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 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 TIME_WITHOUT_TIME_ZONE: + 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 List) { + return listToArrayData((List) value, at.getElementType()); + } + if (value instanceof Object[]) { + return objectArrayToArrayData((Object[]) value, at.getElementType()); + } + if (value instanceof Iterable) { + return iterableToArrayData((Iterable) value, at.getElementType()); + } + return value; + + case MAP: + case MULTISET: + if (value instanceof GenericMapData) { + return value; + } + if (value instanceof Map) { + MapType mt = (MapType) type; + return mapToMapData((Map) value, mt.getKeyType(), mt.getValueType()); + } + if (value instanceof Iterable) { + MapType mt = (MapType) type; + return mapEntryIterableToMapData( + (Iterable) value, mt.getKeyType(), mt.getValueType()); + } + return value; + + default: + return value; + } + } + + 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 listToArrayData(List list, LogicalType elementType) { + Object[] arr = new Object[list.size()]; + for (int i = 0; i < list.size(); i++) { + arr[i] = toInternal(list.get(i), elementType); + } + return new GenericArrayData(arr); + } + + 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) { + java.util.LinkedHashMap converted = + new java.util.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) { + java.util.LinkedHashMap result = new java.util.LinkedHashMap<>(); + boolean typeChecked = false; + for (Object e : iterable) { + // instanceof is slow; checking only the first element is sufficient for state entries. + if (!typeChecked && !(e instanceof Map.Entry)) { + throw new UnsupportedOperationException( + "Map conversion supports only Iterable but received: " + + iterable.getClass().getName()); + } + typeChecked = true; + Map.Entry entry = (Map.Entry) e; + result.put( + toInternal(entry.getKey(), keyType), toInternal(entry.getValue(), valueType)); + } + return new GenericMapData(result); + } +} 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..653c4543c6acb6 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/deserializer/PojoToRowDataDeserializer.java @@ -0,0 +1,309 @@ +/* + * 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.PojoDeserializerCompatibilitySnapshot; +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.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); + + // Mirrors constants in PojoSerializer + static final int IS_NULL = 0x01; + static final int NO_SUBCLASS = 0x02; + static final int IS_SUBCLASS = 0x04; + static final int IS_TAGGED_SUBCLASS = 0x08; + + 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); + } + + int numFields = fieldDeserializerList.size(); + return new PojoToRowDataDeserializer( + numFields, + fieldDeserializerList.toArray(new TypeSerializer[0]), + fieldTypeList.toArray(new LogicalType[0]), + fieldNameList.toArray(new String[0]), + subDeserializers); + } + + PojoToRowDataDeserializer( + int numFields, + TypeSerializer[] fieldDeserializers, + LogicalType[] fieldTypes, + String[] fieldNames, + List registeredSubclassDeserializers) { + this.numFields = numFields; + 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 & IS_NULL) != 0) { + return null; + } + + if ((flags & NO_SUBCLASS) != 0) { + return readFields(source); + } + + if ((flags & IS_TAGGED_SUBCLASS) != 0) { + int tag = source.readByte() & 0xFF; + if (tag < registeredSubclassDeserializers.size()) { + return registeredSubclassDeserializers.get(tag).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 & 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; + } + Object raw; + try { + raw = fieldDeserializers[i].deserialize(source); + } catch (Exception e) { + // Stream position may be corrupted — null remaining fields and return partial row. + LOG.warn( + "Failed to deserialize field '{}' (index {}): {}. " + + "Setting this and remaining fields to null.", + fieldNames[i], + i, + e.getMessage()); + row.setField(i, null); + for (int j = i + 1; j < numFields; j++) { + row.setField(j, null); + } + return row; + } + try { + row.setField(i, InternalTypeConverter.toInternal(raw, fieldTypes[i])); + } catch (Exception e) { + // Bytes already consumed correctly; only this field's conversion failed. + 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; + } + + @Override + public int hashCode() { + return numFields; + } + + @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..ffcd492f4dec51 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,14 +19,10 @@ 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; @@ -36,11 +32,8 @@ 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 +47,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) { @@ -67,7 +60,13 @@ public void open() throws Exception { InternalTimerService timerService = getInternalTimerService(USER_TIMERS_NAME); - context = new Context<>(getKeyedStateBackend(), timerService); + TimerRegistration timerRegistration = + registerTimers( + getKeyedStateBackend(), + timerService, + USER_TIMERS_NAME, + namespace -> namespace.equals(VoidNamespace.INSTANCE)); + context = new Context(timerRegistration); } @Override @@ -77,7 +76,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 +85,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 +119,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..8f6ff92cb07336 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 passed to {@link #setCurrentKeyAndKeyGroup} meaning the key-group is unknown (e.g. + * for namespace types whose backend lookup API does not expose the physically stored + * key-group), in which case the key-group is derived from {@code key.hashCode()} instead. + */ + 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,81 @@ 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}, restricted to the timers whose namespace matches {@code + * namespaceFilter}, and returns a {@link TimerRegistration} exposing them. + * + *

Callers pass a filter rather than always snapshotting every timer because namespaced + * readers (e.g. window state) and non-namespaced readers (e.g. plain keyed state, which uses + * {@link org.apache.flink.runtime.state.VoidNamespace}) have different notions of which timers + * belong to the state being read. + */ + protected final TimerRegistration registerTimers( + KeyedStateBackend keyedStateBackend, + 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 implements Serializable { + + private static final long serialVersionUID = 1L; + + 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 +227,20 @@ public final void setCurrentKey(Object key) { keyedStateBackend.setCurrentKey((KEY) key); } + /** + * Restores the reading context for the given key, using the given key-group if known. + * + *

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..2ab5a9950eb0f3 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,13 @@ private WindowReaderOperator( public void open() throws Exception { super.open(); - ctx = new Context(getKeyedStateBackend(), getInternalTimerService(WINDOW_TIMER_NAME)); + TimerRegistration timerRegistration = + registerTimers( + getKeyedStateBackend(), + getInternalTimerService(WINDOW_TIMER_NAME), + WINDOW_TIMER_NAME, + namespace -> true); + ctx = new Context(getKeyedStateBackend(), timerRegistration); } @Override @@ -176,57 +176,33 @@ 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 { + // The backend's getKeysAndNamespaces(String) API does not expose the physically stored + // key-group, so we fall back to hashCode-based derivation via the -1 sentinel. + 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; - - ListState procTimers; + final TimerRegistration timerRegistration; private Context( - KeyedStateBackend keyedStateBackend, InternalTimerService timerService) - throws Exception { + KeyedStateBackend keyedStateBackend, TimerRegistration timerRegistration) { 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 +233,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..a8a8868491a75a 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,14 @@ package org.apache.flink.state.api.runtime; import org.apache.flink.annotation.Internal; +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; @@ -33,6 +35,8 @@ import org.apache.flink.runtime.state.metainfo.StateMetaInfoSnapshot; import org.apache.flink.state.api.OperatorIdentifier; +import javax.annotation.Nullable; + import java.io.DataInputStream; import java.io.IOException; import java.util.Map; @@ -44,6 +48,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 +105,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 +148,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()); 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..466654f881d65c --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/KeyedStateSchemaInfo.java @@ -0,0 +1,92 @@ +/* + * 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} the same way {@link + * #logicalType} is — never a fixed {@code TimeWindow}-shaped type. + * + *

Named "window" rather than "namespace" because this is the user-facing, + * post-conversion form; see {@link StateSchemaInfo}'s class-level note 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..b2f4da5fbe7b2e --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/SerializerSnapshotToLogicalTypeConverter.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.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.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.formats.avro.typeutils.AvroSchemaConverter; +import org.apache.flink.formats.avro.typeutils.AvroSerializerSnapshot; +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.Collections; +import java.util.List; + +/** + * 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, …}. + */ +@Internal +public final class SerializerSnapshotToLogicalTypeConverter { + + private SerializerSnapshotToLogicalTypeConverter() {} + + /** + * Converts the given snapshot to a {@link LogicalType}. + * + * @param snapshot the serializer snapshot to convert; must not 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) { + return convert(snapshot, 0); + } + + @SuppressWarnings("unchecked") + private static LogicalType convert(TypeSerializerSnapshot snapshot, int depth) { + if (snapshot == null) { + return new VarBinaryType(true, VarBinaryType.MAX_LENGTH); + } + + if (snapshot instanceof PojoSerializerSnapshot) { + return convertPojo((PojoSerializerSnapshot) snapshot, depth); + } + + // --- Avro types --- + // Class name check avoids a hard load of AvroSerializerSnapshot when flink-avro is absent; + // if flink-avro IS present the direct cast below is safe. + if (snapshot instanceof AvroSerializerSnapshot) { + return convertAvro((AvroSerializerSnapshot) snapshot); + } + + // --- Primitive / scalar types --- + if (snapshot instanceof IntSerializer.IntSerializerSnapshot) { + return new IntType(false); + } + if (snapshot instanceof LongSerializer.LongSerializerSnapshot) { + return new BigIntType(false); + } + if (snapshot instanceof FloatSerializer.FloatSerializerSnapshot) { + return new FloatType(false); + } + if (snapshot instanceof DoubleSerializer.DoubleSerializerSnapshot) { + return new DoubleType(false); + } + if (snapshot instanceof BooleanSerializer.BooleanSerializerSnapshot) { + return new BooleanType(false); + } + if (snapshot instanceof ByteSerializer.ByteSerializerSnapshot) { + return new TinyIntType(false); + } + if (snapshot instanceof ShortSerializer.ShortSerializerSnapshot) { + return new SmallIntType(false); + } + if (snapshot instanceof CharSerializer.CharSerializerSnapshot) { + return new CharType(false, 1); + } + if (snapshot instanceof StringSerializer.StringSerializerSnapshot) { + return new VarCharType(true, VarCharType.MAX_LENGTH); + } + + // --- Window namespace types --- + // Neither case is namespace-specific: both fire identically if these types were ever used + // as an ordinary VALUE, not just as a namespace. + if (snapshot instanceof TimeWindow.Serializer.TimeWindowSerializerSnapshot) { + return new RowType( + false, + List.of( + new RowType.RowField("window_start", new TimestampType(false, 3)), + new RowType.RowField("window_end", new TimestampType(false, 3)))); + } + if (snapshot instanceof GlobalWindow.Serializer.GlobalWindowSerializerSnapshot) { + return new RowType(false, Collections.emptyList()); + } + + // --- Composite types --- + if (snapshot instanceof ListSerializerSnapshot) { + TypeSerializerSnapshot elementSnapshot = getNestedSnapshot(snapshot, 0); + LogicalType elementType = convert(elementSnapshot, depth + 1); + return new ArrayType(true, elementType); + } + if (snapshot instanceof MapSerializerSnapshot) { + TypeSerializerSnapshot keySnapshot = getNestedSnapshot(snapshot, 0); + TypeSerializerSnapshot valueSnapshot = getNestedSnapshot(snapshot, 1); + return new MapType( + true, convert(keySnapshot, depth + 1), convert(valueSnapshot, depth + 1)); + } + if (snapshot instanceof NullableSerializerSnapshot) { + TypeSerializerSnapshot innerSnapshot = getNestedSnapshot(snapshot, 0); + return convert(innerSnapshot, depth + 1).copy(true); + } + if (snapshot instanceof TupleSerializerSnapshot) { + return convertTuple((TupleSerializerSnapshot) snapshot, depth); + } + if (snapshot instanceof RowDataSerializerSnapshot) { + return convertRowData((RowDataSerializerSnapshot) snapshot); + } + + // --- Fallback --- + 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 convertTuple(TupleSerializerSnapshot snapshot, int depth) { + TypeSerializerSnapshot[] nestedSnapshots = snapshot.getNestedSerializerSnapshots(); + int arity = nestedSnapshots == null ? 0 : nestedSnapshots.length; + List fields = new ArrayList<>(arity); + for (int i = 0; i < arity; i++) { + LogicalType fieldType = convert(nestedSnapshots[i], depth + 1); + fields.add(new RowType.RowField("f" + i, fieldType)); + } + return new RowType(true, fields); + } + + /** + * Converts a {@link RowDataSerializerSnapshot} directly from its {@link + * RowDataSerializerSnapshot#getTypes()}, without inspecting the nested field serializer + * snapshots. + * + *

Unlike POJO/Avro/Tuple snapshots, a {@link LogicalType} is already a complete, + * self-describing schema on its own (including nested {@link RowType}s with their own field + * names), and Flink table internal types (RowData, StringData, DecimalData, …) are always on + * the classpath. So there is nothing to gain from walking the nested serializer snapshot tree + * here, 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 convertAvro(AvroSerializerSnapshot snapshot) { + // getSchema() returns the writer schema embedded in the snapshot — always present + // regardless of whether the specific record class is on the classpath. + return AvroSchemaConverter.convertToDataType(snapshot.getSchema().toString()) + .getLogicalType(); + } + + private static LogicalType convertPojo(PojoSerializerSnapshot snapshot, int depth) { + List>> fieldEntries = + snapshot.getFieldSnapshotEntries(); + List fields = new ArrayList<>(fieldEntries.size()); + + for (AbstractMap.SimpleEntry> entry : fieldEntries) { + String fieldName = entry.getKey(); + TypeSerializerSnapshot fieldSnapshot = entry.getValue(); + LogicalType fieldType = + fieldSnapshot != null + ? convert(fieldSnapshot, depth + 1) + : new VarBinaryType(true, VarBinaryType.MAX_LENGTH); + fields.add(new RowType.RowField(fieldName, fieldType)); + } + + return new RowType(true, fields); + } + + private static TypeSerializerSnapshot getNestedSnapshot( + TypeSerializerSnapshot snapshot, int index) { + if (snapshot instanceof CompositeTypeSerializerSnapshot) { + TypeSerializerSnapshot[] nested = + ((CompositeTypeSerializerSnapshot) snapshot) + .getNestedSerializerSnapshots(); + if (nested != null && index < nested.length) { + return nested[index]; + } + } + return null; + } +} 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..4e1e6e480f3e40 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/StateSchemaExtractor.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.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 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 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); + 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..e70155fcab7c0a --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/StateSchemaInfo.java @@ -0,0 +1,102 @@ +/* + * 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, pre-conversion concept as Flink's runtime state backend knows it (e.g. {@link + * #namespaceSnapshot} here); "window" names the same concept once resolved to a user-facing + * table/column ({@link KeyedStateSchemaInfo.StateEntryInfo#windowLogicalType}). In practice a + * non-{@code VoidNamespace} namespace is always a window, but the two terms are kept distinct by + * layer rather than treated as synonyms picked arbitrarily. + */ +@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; + } + + @Override + public String toString() { + return "StateSchemaInfo{" + + "stateName='" + + stateName + + '\'' + + ", stateKind=" + + stateKind + + ", keySnapshot=" + + keySnapshot.getClass().getSimpleName() + + ", valueSnapshot=" + + valueSnapshot.getClass().getSimpleName() + + (mapKeySnapshot != null + ? ", mapKeySnapshot=" + mapKeySnapshot.getClass().getSimpleName() + : "") + + '}'; + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/EmptyCatalogDatabase.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/EmptyCatalogDatabase.java new file mode 100644 index 00000000000000..f81ad309dfd4a9 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/EmptyCatalogDatabase.java @@ -0,0 +1,67 @@ +/* + * 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.table.catalog.CatalogDatabase; + +import java.util.Collections; +import java.util.Map; +import java.util.Optional; + +/** + * A {@link CatalogDatabase} with no properties, used for every database in {@link StateCatalog}. + */ +@Internal +class EmptyCatalogDatabase implements CatalogDatabase { + + static final EmptyCatalogDatabase INSTANCE = new EmptyCatalogDatabase(); + + private EmptyCatalogDatabase() {} + + @Override + public Map getProperties() { + return Collections.emptyMap(); + } + + @Override + public String getComment() { + return ""; + } + + @Override + public CatalogDatabase copy() { + return this; + } + + @Override + public CatalogDatabase copy(Map properties) { + return this; + } + + @Override + public Optional getDescription() { + return Optional.empty(); + } + + @Override + public Optional getDetailedDescription() { + return Optional.empty(); + } +} 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..75fe589678c8d8 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/SnapshotDiscovery.java @@ -0,0 +1,368 @@ +/* + * 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}. + * + *

Two operations are provided: + * + *

    + *
  • {@link #list()} — full BFS scan of all configured directories; directories at the same + * depth are listed concurrently up to {@code listingParallelism} at a time with automatic + * backpressure via {@link ThreadPoolExecutor.CallerRunsPolicy}. + *
  • {@link #find(String)} — single {@code getFileStatus} check for a specific database name; + * splits on the first (and, if {@code dbNameIncludeTs} is enabled, second) {@code /} to + * recover the label and relative path, then checks exactly one file. + *
+ */ +@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) { + this( + labelToDirPath, + listingParallelism, + StateCatalogOptions.DB_NAME_INCLUDE_TS.defaultValue()); + } + + SnapshotDiscovery( + Map labelToDirPath, int listingParallelism, boolean dbNameIncludeTs) { + validateLabels(labelToDirPath); + this.labelToDir = convert(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. Returns one database name per discovered {@code + * _metadata} file. Used only by {@code listDatabases()}. + * + * @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; + } + + /** + * Single {@code getFileStatus} check for a specific database name. Splits on the first (and, if + * {@code dbNameIncludeTs} is enabled, second) {@code /} to recover the label and the verbatim + * relative path, then checks exactly one file. Returns the snapshot directory path if found. + */ + 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 database name with no '/' has no room for the mandatory creationTs segment that + // buildDatabaseName() always appends when dbNameIncludeTs is enabled. + 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 database name with no '/' means the configured directory itself is the snapshot + // (mirroring buildDatabaseName(), which returns the bare label in that case). + 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 + // ------------------------------------------------------------------------- + + // Package-private so SnapshotDiscoveryTest can call it directly. + List findMetadataFiles(Path directory) throws IOException { + List metadataFiles = new ArrayList<>(); + for (FileStatus status : findMetadataFileStatuses(directory)) { + metadataFiles.add(status.getPath()); + } + return metadataFiles; + } + + 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 + // ------------------------------------------------------------------------- + + private static void validateLabels(Map labelToDirPath) { + if (labelToDirPath.isEmpty()) { + throw new IllegalArgumentException( + "At least one directory must be configured via 'directory.{label}' options."); + } + Set seenDirs = new LinkedHashSet<>(); + for (Map.Entry entry : labelToDirPath.entrySet()) { + String normalizedDir = new Path(entry.getValue()).toUri().getPath(); + if (!seenDirs.add(normalizedDir)) { + throw new IllegalArgumentException( + String.format( + "Directory '%s' is assigned to more than one label.", + entry.getValue())); + } + } + List paths = new ArrayList<>(seenDirs); + for (int i = 0; i < paths.size(); i++) { + for (int j = 0; j < paths.size(); j++) { + if (i == j) { + continue; + } + String a = paths.get(i); + String b = paths.get(j); + if (b.startsWith(a.endsWith("/") ? a : a + "/")) { + throw new IllegalArgumentException( + String.format( + "Directory '%s' is an ancestor of '%s'. Providing both would " + + "discover the same snapshots under multiple labels.", + a, b)); + } + } + } + } + + private static Map convert(Map labelToDirPath) { + Map result = new LinkedHashMap<>(); + for (Map.Entry entry : labelToDirPath.entrySet()) { + result.put(entry.getKey(), new Path(entry.getValue())); + } + 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..c2ca6c4905045b --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/StateCatalog.java @@ -0,0 +1,675 @@ +/* + * 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.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.List; +import java.util.Map; +import java.util.Optional; + +/** + * 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 = EmptyCatalogDatabase.INSTANCE; + + 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<>(); + try { + CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(snapshotPath.get()); + for (OperatorIdentifier opId : StateTableUtils.getOperatorIdentifiers(metadata)) { + for (ResolvedTable candidate : candidateTablesForOperator(metadata, opId)) { + tables.add( + tableName( + candidate.operatorIdentifier, + candidate.kind, + candidate.stateName)); + } + } + } 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) { + 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) { + 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..3ed2ec45a8990d --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/StateCatalogFactory.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.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.HashSet; +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() { + Set> options = new HashSet<>(); + options.add(StateCatalogOptions.LISTING_PARALLELISM); + options.add(StateCatalogOptions.DB_NAME_INCLUDE_TS); + return options; + } + + @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); + int listingParallelism = configuration.get(StateCatalogOptions.LISTING_PARALLELISM); + boolean dbNameIncludeTs = configuration.get(StateCatalogOptions.DB_NAME_INCLUDE_TS); + + return new StateCatalog( + context.getName(), labelsToDirs, listingParallelism, dbNameIncludeTs); + } +} 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..4a385beba18491 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/catalog/StateCatalogOptions.java @@ -0,0 +1,75 @@ +/* + * 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}). + * + *

Example: + * + * 

{@code
+     * CREATE CATALOG state WITH (
+     *     'type'                 = 'state',
+     *     'directory.my-app'     = '/checkpoints/app1',
+     *     'directory.staging'    = '/savepoints/staging'
+     * );
+     * }
+ */ + 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."); + + /** + * Whether derived database names include the snapshot's creation timestamp as a segment, i.e. + * {@code label/creationTs/relativePath} instead of {@code label/relativePath}. The timestamp is + * the modification time of the snapshot's {@code _metadata} file, formatted with {@code + * yyyy-MM-dd'T'HH:mm:ssX} (e.g. {@code 2026-07-22T10:30:45Z}). + */ + 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..420be21620bc4a --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/AbstractMultiColumnScanProvider.java @@ -0,0 +1,75 @@ +/* + * 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); + + 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..a616ab933914ae --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/AbstractSavepointDataStreamScanProvider.java @@ -0,0 +1,311 @@ +/* + * 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.List; +import java.util.Map; +import java.util.function.Supplier; + +/** + * Shared scan-time logic for {@link SavepointDataStreamScanProvider} and {@link + * FlattenedSavepointDataStreamScanProvider}: opening the {@link SavepointReader} against the + * configured state backend and driving {@link SavepointReader#readKeyedState}. Subclasses supply + * the mapping-specific state-descriptor setup and {@link KeyedStateReaderFunction}. + */ +@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 + @SuppressWarnings({"rawtypes", "unchecked"}) + 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 org.apache.flink.api.common.state.StateDescriptor}(s) needed for this scan + * and registers them onto {@code mapping} (e.g. via {@code setStateDescriptor}), 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. The default implementation + * reads plain (VoidNamespace) keyed state via {@link #createReaderFunction}; namespaced + * (window) scan providers override this directly instead of implementing {@link + * #createReaderFunction}. + */ + @SuppressWarnings({"rawtypes", "unchecked"}) + protected DataStream readState(SavepointReader savepointReader, M mapping) + throws Exception { + TypeInformation outTypeInfo = InternalTypeInfo.of(rowType); + return savepointReader.readKeyedState( + operatorIdentifier, + createReaderFunction(mapping), + (TypeInformation) mapping.getKeyTypeInfo(), + outTypeInfo, + keyFilter); + } + + /** + * Creates the reader function that turns the keyed state(s) described by {@code mapping} into + * output rows. Not used by scan providers that override {@link #readState} directly. + */ + protected KeyedStateReaderFunction createReaderFunction(M mapping) { + throw new UnsupportedOperationException( + "This provider does not support createReaderFunction; it overrides readState directly."); + } + + /** + * Loads the configured (or overridden) {@link StateBackend} and opens a {@link SavepointReader} + * against it. + * + *

Package-private (rather than {@code private}) so that the non-keyed (operator state) scan + * providers, which cannot extend this keyed-specific abstraction, can reuse it directly. + */ + @SuppressWarnings("rawtypes") + 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} for + * columns whose {@link TypeSerializer} could not be resolved from the preloaded savepoint + * metadata. + * + *

Returns an empty map if all columns already have a resolved {@link TypeSerializer}, + * avoiding the overhead of reading the savepoint header. + */ + protected final Map loadFallbackSchemas( + List columns) { + boolean anyNullTypeInfo = + columns.stream() + .anyMatch( + c -> + c.getValueTypeSerializer() == null + || (c.getStateType() + == SavepointConnectorOptions + .StateType.MAP + && c.getMapKeyTypeSerializer() == null)); + return SavepointFallbackSchemaLoader.loadFallbackSchemas( + statePath, operatorIdentifier, anyNullTypeInfo); + } + + /** + * Builds the {@link StateDescriptor} for a single VALUE/LIST/MAP state. + * + *

When a {@link TypeSerializer} was already resolved from the preloaded savepoint metadata, + * the serializer-accepting descriptor constructors are used directly. Otherwise the original + * serializer is restored from {@code fallbackSchemas} (read from the savepoint header). + * + *

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 the originally-registered one. 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 (never + * write) this state. + */ + @SuppressWarnings({"rawtypes", "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 resolvedValueSerializer = + valueTypeSerializer != null + ? valueTypeSerializer + : SavepointFallbackSchemaLoader.buildFallbackSerializer( + SavepointFallbackSchemaLoader.getSchema( + name, fallbackSchemas) + .valueSnapshot); + switch (actualStateKind) { + case REDUCING: + return new ReducingStateDescriptor<>( + name, new NoOpReduceFunction(), resolvedValueSerializer); + case AGGREGATING: + return new AggregatingStateDescriptor<>( + name, new IdentityAggregateFunction(), resolvedValueSerializer); + default: + return new ValueStateDescriptor<>(name, resolvedValueSerializer); + } + + case LIST: + if (valueTypeSerializer != null) { + return new ListStateDescriptor<>(name, valueTypeSerializer); + } + return new ListStateDescriptor<>( + name, + SavepointFallbackSchemaLoader.buildFallbackSerializer( + SavepointFallbackSchemaLoader.getSchema(name, fallbackSchemas) + .valueSnapshot)); + + 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); + } + } + + /** + * Never-invoked {@link ReduceFunction} used to build a {@link ReducingStateDescriptor} for + * read-only access: callers only call {@code ReducingState.get()}, which never merges, so the + * reduce function itself is never called. + */ + 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). {@code createAccumulator}/{@code add}/{@code merge} are never + * invoked since callers only ever read (never write) this state. + */ + 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..becd4ef75a23b0 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/AbstractSavepointDynamicTableSource.java @@ -0,0 +1,109 @@ +/* + * 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. + */ +@Internal +abstract class AbstractSavepointDynamicTableSource + implements ScanTableSource, SupportsFilterPushDown { + + @Nullable protected final String stateBackendType; + protected final String statePath; + protected final OperatorIdentifier operatorIdentifier; + protected Supplier mappingSupplier; + protected RowType rowType; + @Nullable protected SavepointKeyFilter keyFilter; + + protected AbstractSavepointDynamicTableSource( + @Nullable final String stateBackendType, + final String statePath, + final OperatorIdentifier operatorIdentifier, + final Supplier mappingSupplier, + final RowType rowType) { + this.stateBackendType = stateBackendType; + this.statePath = statePath; + this.operatorIdentifier = operatorIdentifier; + this.mappingSupplier = mappingSupplier; + this.rowType = rowType; + } + + @Override + public Result applyFilters(List filters) { + return SavepointFilterTranslator.applyKeyColumnFilters( + getKeyColumnIndex(), rowType, filters, kf -> this.keyFilter = kf); + } + + /** Index of the (single) key column in {@link #rowType}, used for filter push-down. */ + protected abstract int getKeyColumnIndex(); + + @Override + public final DynamicTableSource copy() { + AbstractSavepointDynamicTableSource copy = newInstance(); + copy.keyFilter = this.keyFilter; + return copy; + } + + /** + * Creates a fresh instance carrying the same immutable 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(); + } + + /** + * Builds the {@link DataStreamScanProvider} for a given set of scan-time constructor arguments. + * Injected by {@link SavepointDynamicTableSourceFactory} as a constructor reference to the + * appropriate concrete {@code *DataStreamScanProvider} class (e.g. {@code + * SavepointDataStreamScanProvider::new}), letting a single table-source class serve every + * keyed/namespaced state table kind without a subclass per kind. + */ + interface ScanProviderFactory { + DataStreamScanProvider create( + @Nullable String stateBackendType, + String statePath, + OperatorIdentifier operatorIdentifier, + Supplier mappingSupplier, + RowType rowType, + @Nullable SavepointKeyFilter 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..1674bf2c00f8c9 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/AbstractSingleColumnScanProvider.java @@ -0,0 +1,74 @@ +/* + * 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) { + boolean anyNullTypeInfo = + mapping.getValueTypeSerializer() == null + || (mapping.getStateType() == SavepointConnectorOptions.StateType.MAP + && mapping.getMapKeyTypeSerializer() == null); + Map fallbackSchemas = + SavepointFallbackSchemaLoader.loadFallbackSchemas( + statePath, operatorIdentifier, anyNullTypeInfo); + + 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..d8346e767438ad --- /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 = StateValueConverter.INSTANCE; + + 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..afaa96d7967b34 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/FlattenedSavepointDataStreamScanProvider.java @@ -0,0 +1,53 @@ +/* + * 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.state.api.functions.KeyedStateReaderFunction; +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}). + */ +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 KeyedStateReaderFunction createReaderFunction( + FlattenedStateTableMapping mapping) { + return 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..2304f30e505f59 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/FlattenedSavepointDynamicTableSource.java @@ -0,0 +1,99 @@ +/* + * 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). The two kinds differ only in the fixed + * schema validated upstream and in which {@link AbstractSavepointDataStreamScanProvider} subclass + * performs the actual scan, which {@link SavepointDynamicTableSourceFactory} supplies via {@code + * scanProviderFactory}. + * + *

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 { + + private final int keyColumnIndex; + private final String summaryString; + private final ScanProviderFactory scanProviderFactory; + + public FlattenedSavepointDynamicTableSource( + @Nullable final String stateBackendType, + final String statePath, + final OperatorIdentifier operatorIdentifier, + int keyColumnIndex, + final Supplier mappingSupplier, + final RowType rowType, + String summaryString, + ScanProviderFactory scanProviderFactory) { + super(stateBackendType, statePath, operatorIdentifier, mappingSupplier, rowType); + this.keyColumnIndex = keyColumnIndex; + this.summaryString = summaryString; + this.scanProviderFactory = scanProviderFactory; + } + + @Override + protected int getKeyColumnIndex() { + return keyColumnIndex; + } + + @Override + protected AbstractSavepointDynamicTableSource newInstance() { + return new FlattenedSavepointDynamicTableSource<>( + stateBackendType, + statePath, + operatorIdentifier, + keyColumnIndex, + mappingSupplier, + rowType, + summaryString, + scanProviderFactory); + } + + @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/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..3867c3c7fd9850 --- /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 = + KeyedTableMappingSupport.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 = + KeyedTableMappingSupport.createTypeResolver( + statePath, operatorIdentifier, serializerConfig); + + DataType physicalDataType = catalogTable.getResolvedSchema().toPhysicalRowDataType(); + RowType rowType = (RowType) physicalDataType.getLogicalType(); + + KeyedTableMappingSupport.FlattenedSerializers serializers = + KeyedTableMappingSupport.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..34598c0f55fcd9 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 = StateValueConverter.INSTANCE; + + /** + * 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,40 @@ 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( - valueLogicalType, - ((ListState) states.get(columnConfig.getColumnIndex())).get())); + converter.getValue(valueLogicalType, ((ListState) state).get())); 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 +169,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/KeyedTableMappingSupport.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/KeyedTableMappingSupport.java new file mode 100644 index 00000000000000..2c68167f860edd --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/KeyedTableMappingSupport.java @@ -0,0 +1,318 @@ +/* + * 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.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 four keyed table mapping classes ({@link StateTableMapping}, {@link + * WindowStateTableMapping}, {@link FlattenedStateTableMapping}, {@link + * WindowFlattenedStateTableMapping}). 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 KeyedTableMappingSupport { + + private KeyedTableMappingSupport() {} + + /** 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 keyField} in the physical row data type. */ + static int keyColumnIndex(DataType physicalDataType, String keyField) { + final LogicalType physicalType = physicalDataType.getLogicalType(); + Preconditions.checkArgument( + physicalType.is(LogicalTypeRoot.ROW), "Row data type expected."); + return LogicalTypeChecks.getFieldNames(physicalType).indexOf(keyField); + } + + /** 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. Shared by {@code StateTableMapping#project} and {@code + * WindowStateTableMapping#project}. + */ + 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 {@code StateTableMapping}/{@code WindowStateTableMapping}'s {@code project(...)} + * (remapping their tracked key/window/value column indices) and the corresponding table + * sources' {@code applyProjection(...)} (eagerly tracking the key/window column index for + * filter push-down before the mapping itself is resolved), 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 operator metadata (state serializer snapshots + backend key serializer snapshot) for + * an operator in a single I/O operation. + */ + static OperatorStateMetadata preloadOperatorMetadata( + String savepointPath, OperatorIdentifier operatorIdentifier) { + try { + return SavepointLoader.loadOperatorMetadata(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); + } + } + + /** + * Preloads operator metadata and builds the {@link SavepointTypeInfoResolver} derived from it, + * in one step. Shared by all four mapping classes' {@code from(...)} factories. + */ + static SavepointTypeInfoResolver createTypeResolver( + String statePath, + OperatorIdentifier operatorIdentifier, + SerializerConfig serializerConfig) { + OperatorStateMetadata operatorMetadata = + preloadOperatorMetadata(statePath, operatorIdentifier); + return new SavepointTypeInfoResolver( + operatorMetadata.stateSnapshots, + serializerConfig, + operatorMetadata.keySerializerSnapshot); + } + + /** + * 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, shared by {@code FlattenedStateTableMapping} + * and {@code WindowFlattenedStateTableMapping}'s {@code validateFlattenedSchema}. + * + * @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 (composite value field, map-key + * serializer, value serializer) shared by {@code FlattenedStateTableMapping} and {@code + * WindowFlattenedStateTableMapping}'s {@code from(...)} factories. + */ + static FlattenedSerializers resolveFlattenedSerializers( + RowType rowType, + SavepointTypeInfoResolver typeResolver, + String stateName, + SavepointConnectorOptions.StateType stateType, + int stateKeyColumnIndex, + int subKeyColumnIndex, + int valueColumnIndex) { + RowType.RowField keyRowField = rowType.getFields().get(stateKeyColumnIndex); + TypeInformation keyTypeInfo = typeResolver.resolveKeyType(keyRowField); + + 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); + } + + @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()); + + SavepointConnectorOptions.StateType stateType = inferStateType(valueRowField.getType()); + + TypeSerializer mapKeyTypeSerializer = + typeResolver.resolveMapKeySerializer( + valueRowField, stateType.equals(SavepointConnectorOptions.StateType.MAP)); + + TypeSerializer valueTypeSerializer = typeResolver.resolveValueSerializer(valueRowField); + + return new StateValueColumnConfiguration( + columnIndex, + stateName, + stateType, + typeResolver.resolveStateKind(stateName), + 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. + */ + 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/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..1d83f15a88e5da --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/MultiColumnStateMapping.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 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 {@code getAllValueColumns()}. Implemented by {@link + * StateTableMapping} and {@link WindowStateTableMapping}, allowing {@link + * AbstractMultiColumnScanProvider} to build their state descriptors generically. + */ +@Internal +interface MultiColumnStateMapping extends SavepointStateMapping { + + /** + * Full original value columns required for state-descriptor registration / key(-and-namespace) + * enumeration. 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 across {@link + * StateTableMapping} and {@link WindowStateTableMapping}. + */ + 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/SavepointDataStreamScanProvider.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointDataStreamScanProvider.java index 35a9e04a7af8c5..c6023886f4d42d 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,47 @@ 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.state.api.functions.KeyedStateReaderFunction; 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. */ -@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 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); - } +/** + * Savepoint data stream scan provider. + * + *

State descriptors are built in two ways depending on whether {@link StateTableMapping} was + * able to resolve a {@link TypeSerializer} for the column from the savepoint metadata: + * + *

    + *
  • Resolved: the resolved serializer is passed directly to the serializer-accepting + * descriptor constructors. + *
  • Unresolved: when no serializer could be resolved (e.g. the state is missing from the + * preloaded metadata), the savepoint header is read to obtain the original type serializer + * instead. + *
+ */ +public class SavepointDataStreamScanProvider + extends AbstractMultiColumnScanProvider { 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 KeyedStateReaderFunction createReaderFunction( + StateTableMapping mapping) { + return 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..f75c876413afee 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,93 +18,95 @@ 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}). The two kinds differ only in which {@link + * AbstractSavepointDataStreamScanProvider} subclass performs the actual scan, which {@link + * SavepointDynamicTableSourceFactory} supplies via {@code scanProviderFactory}. + */ +public class SavepointDynamicTableSource + extends AbstractSavepointDynamicTableSource implements SupportsProjectionPushDown { + + // keyColumnIndex is kept eager so filter push-down can reference it during planning + // without resolving the lazy mapping. + private int keyColumnIndex; + + private final String summaryString; + private final ScanProviderFactory scanProviderFactory; 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; + int keyColumnIndex, + final Supplier mappingSupplier, + RowType rowType, + String summaryString, + ScanProviderFactory scanProviderFactory) { + super(stateBackendType, statePath, operatorIdentifier, mappingSupplier, rowType); + this.keyColumnIndex = keyColumnIndex; + this.summaryString = summaryString; + this.scanProviderFactory = scanProviderFactory; + } + + @Override + protected int getKeyColumnIndex() { + return keyColumnIndex; } @Override - public DynamicTableSource copy() { - SavepointDynamicTableSource copy = - new SavepointDynamicTableSource( - stateBackendType, - statePath, - operatorIdentifier, - keyTypeInfo, - keyValueProjections, - rowType); - copy.keyFilter = this.keyFilter; - return copy; + protected AbstractSavepointDynamicTableSource newInstance() { + return new SavepointDynamicTableSource<>( + stateBackendType, + statePath, + operatorIdentifier, + keyColumnIndex, + mappingSupplier, + rowType, + summaryString, + scanProviderFactory); } @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()); + public boolean supportsNestedProjection() { + return false; } @Override - public String asSummaryString() { - return "Savepoint Table Source"; + @SuppressWarnings("unchecked") + public void applyProjection(int[][] projectedFields, DataType producedDataType) { + this.rowType = (RowType) producedDataType.getLogicalType(); + // Track the projected key column index eagerly for filter push-down. + this.keyColumnIndex = + KeyedTableMappingSupport.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); } @Override - public ChangelogMode getChangelogMode() { - return ChangelogMode.insertOnly(); + public String asSummaryString() { + return summaryString; } @Override public ScanRuntimeProvider getScanRuntimeProvider(ScanContext scanContext) { - return new SavepointDataStreamScanProvider( + return scanProviderFactory.create( stateBackendType, statePath, operatorIdentifier, - keyTypeInfo, - keyValueProjections, + mappingSupplier, rowType, keyFilter); } 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..07cbf7e80d6199 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)); + FlattenedStateTableMapping.validateFlattenedSchema(context.getCatalogTable()); - 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); - } - - private static ConfigOptions.OptionBuilder optionOf(String rowField, String optionName) { - return ConfigOptions.key(String.format("%s.%s.%s", FIELDS, rowField, optionName)); - } - - 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."); - } - - 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."); - } + RowType rowType = (RowType) context.getPhysicalRowDataType().getLogicalType(); - DataType physicalDataType = schema.toPhysicalRowDataType(); - int keyProjection = createKeyFormatProjection(physicalDataType, keyFields.get(0)); - int[] valueProjection = createValueFormatProjection(physicalDataType, keyProjection); + String stateName = validateAndGetFlattenedStateName(options); - return Tuple2.of(keyProjection, valueProjection); - } - - 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); + // Defer I/O to scan time by creating the mapping lazily. + Supplier mappingSupplier = + () -> + FlattenedStateTableMapping.from( + context.getCatalogTable(), + stateName, + statePath, + operatorIdentifier, + serializerConfig, + stateType); - 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,10 @@ 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); 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..0877276c320ebe --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/SavepointFallbackSchemaLoader.java @@ -0,0 +1,117 @@ +/* + * 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.TypeSerializer; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; +import org.apache.flink.api.java.typeutils.runtime.PojoSerializerSnapshot; +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.PojoToRowDataDeserializer; +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 classes. + * + *

Used by {@link SavepointDataStreamScanProvider} and {@link + * FlattenedSavepointDataStreamScanProvider} when no serializer could be resolved from the preloaded + * savepoint metadata and the original serializer must be restored from the savepoint header. + */ +@Internal +final class SavepointFallbackSchemaLoader { + + private SavepointFallbackSchemaLoader() {} + + /** + * Restores {@code snapshot} into a {@link TypeSerializer}, tolerating {@link + * PojoSerializerSnapshot}s whose POJO class is missing from the classpath — including ones + * nested arbitrarily deep inside a composite snapshot (e.g. a {@code ListSerializerSnapshot} or + * {@code MapSerializerSnapshot} wrapping a POJO element/value type), since {@code + * CompositeTypeSerializerSnapshot#restoreSerializer()} eagerly restores all of its nested + * serializers. + * + *

For the duration of the {@link TypeSerializerSnapshot#restoreSerializer()} call, registers + * a factory (via {@link PojoSerializerSnapshot#setPojoRestoreSerializerFactory}) that {@link + * PojoSerializerSnapshot#restoreSerializer()} consults whenever it encounters a missing POJO + * class, at any nesting depth, and which builds a {@link PojoToRowDataDeserializer} instead of + * throwing. Snapshots whose class is present restore normally, unaffected by the factory. + */ + static TypeSerializer buildFallbackSerializer(TypeSerializerSnapshot snapshot) { + PojoSerializerSnapshot.setPojoRestoreSerializerFactory(PojoToRowDataDeserializer::create); + try { + return snapshot.restoreSerializer(); + } finally { + PojoSerializerSnapshot.clearPojoRestoreSerializerFactory(); + } + } + + 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..2530c5dd0b2735 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} @@ -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 = 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 63% 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..9e88a4e079e996 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 {@link StateTableMapping} and {@link FlattenedStateTableMapping}, + * allowing {@link AbstractSavepointDataStreamScanProvider} and {@link + * AbstractSavepointDynamicTableSource} to operate on either mapping 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..49b1e88ed1dadf 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,29 +20,31 @@ 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.MapSerializer; -import org.apache.flink.configuration.ConfigOption; -import org.apache.flink.configuration.Configuration; +import org.apache.flink.api.java.typeutils.runtime.PojoSerializerSnapshot; import org.apache.flink.runtime.state.metainfo.StateMetaInfoSnapshot; +import org.apache.flink.state.api.input.deserializer.PojoToRowDataDeserializer; +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 { @@ -56,106 +58,209 @@ enum InferenceContext { /** Inferring the key type of a MAP state. */ MAP_KEY, /** Inferring the value type (behavior depends on logical type). */ - VALUE + VALUE, + /** + * Inferring 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 even if the logical type happens + * to be an ARRAY or MAP. + */ + 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 TypeInformation for keyed state keys. * - *

This is a simplified version of type resolution specifically for key types, which are - * always primitive and don't require complex metadata inference. + *

Handles primitive types directly from the {@link LogicalType}, and complex types (POJO, + * Avro) by inspecting the key serializer snapshot stored in the savepoint metadata (every state + * in the same keyed-state backend shares the same key serializer snapshot; it is stored at the + * {@code KeyedBackendSerializationProxy} level, not per individual state, and preloaded at + * construction time via {@link + * org.apache.flink.state.api.runtime.SavepointLoader#loadOperatorMetadata}). * - * @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 IllegalArgumentException If the type cannot be inferred * @throws RuntimeException If type instantiation fails */ - public TypeInformation resolveKeyType( - Configuration options, - ConfigOption classOption, - ConfigOption typeInfoFactoryOption, - RowType.RowField rowField) { + public TypeInformation resolveKeyType(RowType.RowField rowField) { try { - // Priority 1: Explicit configuration (backward compatibility) - TypeInformation explicitTypeInfo = - getExplicitTypeInfo(options, classOption, typeInfoFactoryOption); - if (explicitTypeInfo != null) { - return explicitTypeInfo; - } - - // Priority 2: Simple primitive type inference from LogicalType LogicalType logicalType = rowField.getType(); String columnName = rowField.getName(); - return TypeInformation.of(getPrimitiveClass(logicalType, columnName)); + + // Priority 1: Simple primitive type inference from LogicalType + Class primitiveClass = getPrimitiveClass(logicalType, columnName); + if (primitiveClass != null) { + return TypeInformation.of(primitiveClass); + } + + // Priority 2: ROW type key (POJO or Avro) — resolve from savepoint metadata. + // getPrimitiveClass() returns null for ROW types; we must not call + // TypeInformation.of(null) as that throws NullPointerException. + if (logicalType.getTypeRoot() == LogicalTypeRoot.ROW && preloadedKeySnapshot != null) { + return resolveRowKeyTypeInfo(preloadedKeySnapshot, logicalType); + } + + throw new IllegalArgumentException( + "Cannot resolve key TypeInformation for column '" + + columnName + + "' with type " + + logicalType.getTypeRoot() + + "."); } catch (ReflectiveOperationException e) { throw new RuntimeException(e); } } /** - * Resolves TypeSerializer for a table field using a three-tier priority system with direct + * Resolves a {@link TypeInformation} for a ROW-type key column from the given serializer + * snapshot. + * + *

The POJO restore-serializer factory is activated for the duration of the call so that + * {@link TypeSerializerSnapshot#restoreSerializer()} on a POJO snapshot whose class is absent + * returns a {@link PojoToRowDataDeserializer} instead of throwing. For all other snapshot types + * (Avro, etc.) the factory is simply ignored. + */ + private TypeInformation resolveRowKeyTypeInfo( + TypeSerializerSnapshot keySnapshot, LogicalType keyLogicalType) { + PojoSerializerSnapshot.setPojoRestoreSerializerFactory( + pojoSnapshot -> PojoToRowDataDeserializer.create(pojoSnapshot)); + try { + return ExternalTypeInfo.of( + TypeConversions.fromLogicalToDataType(keyLogicalType), + keySnapshot.restoreSerializer()); + } finally { + PojoSerializerSnapshot.clearPojoRestoreSerializerFactory(); + } + } + + /** + * 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). + */ + public TypeSerializer resolveMapKeySerializer(RowType.RowField rowField, boolean isMap) { + return resolveSerializer(rowField, isMap, InferenceContext.MAP_KEY); + } + + /** + * 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, true, InferenceContext.KEY); + } + + /** Resolves the {@link TypeSerializer} for a state's value. */ + public TypeSerializer resolveValueSerializer(RowType.RowField rowField) { + return resolveSerializer(rowField, true, InferenceContext.VALUE); + } + + /** + * Resolves the {@link TypeSerializer} for a non-keyed (operator) state's value — {@code + * ListState}/{@code UnionState}/{@code BroadcastState} — stored directly under {@code + * CommonSerializerKeys#VALUE_SERIALIZER} 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, true, InferenceContext.FLAT_VALUE); + } + + /** + * Resolves the precise {@link StateDescriptor.Type} (e.g. {@code REDUCING}/{@code AGGREGATING}) + * a state was originally registered under, read directly from the {@code KEYED_STATE_TYPE} + * option stored in the preloaded savepoint metadata. Returns {@code UNKNOWN} when the state (or + * the option) is not present, in which case callers should 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 kindStr = + stateMetaInfo.getOption(StateMetaInfoSnapshot.CommonOptionsKeys.KEYED_STATE_TYPE); + try { + return StateDescriptor.Type.valueOf(kindStr); + } catch (IllegalArgumentException | NullPointerException 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. + * + *

Unlike {@link #resolveValueSerializer}, there is no LogicalType-based fallback: the + * namespace serializer is only ever needed for states that {@code StateTableUtils} has already + * classified as namespaced (i.e. metadata is known to contain it). Restoration goes through + * {@link SavepointFallbackSchemaLoader#buildFallbackSerializer} so that a namespace type that + * is (or contains) a POJO whose class is missing from the classpath falls back to a {@code + * PojoToRowDataDeserializer} instead of throwing. + */ + 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 TypeSerializer for a table field using a two-tier priority system with direct * serializer extraction for metadata inference. * - *

Three-Tier Priority System (Serializer-First)

+ *

Two-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
    + *
  3. Priority 1: Metadata Inference (Highest priority)
    * Directly extracts serializers from preloaded savepoint metadata (NO TypeInformation * conversion). - *
  4. Priority 3: LogicalType Fallback (Lowest priority)
    + *
  5. Priority 2: LogicalType Fallback (Lowest priority)
    * Infers TypeInformation from table schema's LogicalType, then converts to serializer. *
* *

This approach eliminates TypeInformation extraction complexity for metadata inference, * making it work with ANY serializer type (Avro, custom types, etc.). * - * @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 inferStateType Whether to enable automatic type inference. If false, returns null. * @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 + * @return The resolved TypeSerializer, or null if inferStateType is false or the type cannot be + * inferred (complex types are left to the caller's snapshot-based fallback). * @throws RuntimeException If serializer creation fails */ - public TypeSerializer resolveSerializer( - Configuration options, - ConfigOption classOption, - ConfigOption typeInfoFactoryOption, - RowType.RowField rowField, - boolean inferStateType, - InferenceContext context) { + private TypeSerializer resolveSerializer( + RowType.RowField rowField, boolean inferStateType, 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 + // Priority 1: Direct serializer extraction from metadata Optional> metadataSerializer = getSerializerFromMetadata(rowField, context); if (metadataSerializer.isPresent()) { @@ -167,8 +272,12 @@ public TypeSerializer resolveSerializer( return metadataSerializer.get(); } - // Priority 3: Fallback to LogicalType-based inference + // Priority 2: Fallback to LogicalType-based inference TypeInformation fallbackTypeInfo = inferTypeFromLogicalType(rowField, context); + if (fallbackTypeInfo == null) { + // Complex types (ROW/POJO) cannot be inferred; caller falls back to snapshot. + return null; + } return fallbackTypeInfo.createSerializer(serializerConfig); } catch (Exception e) { throw new RuntimeException( @@ -177,47 +286,7 @@ 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). + * Directly extracts TypeSerializer from preloaded metadata (Priority 1). * *

This method performs NO I/O and NO TypeInformation conversion. It directly extracts the * serializer that was used to write the state data. @@ -255,7 +324,9 @@ private Optional> getSerializerFromMetadata( stateMetaInfo.getTypeSerializerSnapshot( StateMetaInfoSnapshot.CommonSerializerKeys.VALUE_SERIALIZER); if (valueSnapshot != null) { - TypeSerializer valueSerializer = valueSnapshot.restoreSerializer(); + TypeSerializer valueSerializer = + SavepointFallbackSchemaLoader.buildFallbackSerializer( + valueSnapshot); if (valueSerializer instanceof MapSerializer) { serializerSnapshot = ((MapSerializer) valueSerializer) @@ -265,6 +336,7 @@ private Optional> getSerializerFromMetadata( } break; case VALUE: + case FLAT_VALUE: serializerSnapshot = stateMetaInfo.getTypeSerializerSnapshot( StateMetaInfoSnapshot.CommonSerializerKeys.VALUE_SERIALIZER); @@ -279,8 +351,9 @@ private Optional> getSerializerFromMetadata( return Optional.empty(); } - // Restore serializer from snapshot - TypeSerializer serializer = serializerSnapshot.restoreSerializer(); + // Restore serializer from snapshot, using POJO-friendly path when class may be absent + TypeSerializer serializer = + SavepointFallbackSchemaLoader.buildFallbackSerializer(serializerSnapshot); // For VALUE context with complex types, extract the appropriate sub-serializer if (context == InferenceContext.VALUE) { @@ -355,7 +428,7 @@ private TypeInformation inferTypeFromLogicalType( try { switch (context) { case KEY: - // Keys are always primitive + // Primitive keys only — complex (ROW) keys are handled by resolveKeyType() return TypeInformation.of(getPrimitiveClass(logicalType, columnName)); case MAP_KEY: @@ -370,6 +443,11 @@ private TypeInformation inferTypeFromLogicalType( case VALUE: return inferValueTypeFromLogicalType(logicalType, columnName); + case FLAT_VALUE: + // Non-keyed state values are never ARRAY/MAP-wrapped composites; infer + // directly from the (already flat) logical type. + return TypeInformation.of(getPrimitiveClass(logicalType, columnName)); + default: throw new UnsupportedOperationException("Unknown context: " + context); } @@ -392,17 +470,28 @@ private TypeInformation inferValueTypeFromLogicalType( case ARRAY: // ARRAY logical type → LIST state → return element type ArrayType arrayType = (ArrayType) logicalType; - return TypeInformation.of( - getPrimitiveClass(arrayType.getElementType(), columnName)); + Class elementClass = getPrimitiveClass(arrayType.getElementType(), columnName); + if (elementClass == null) { + return null; + } + return TypeInformation.of(elementClass); case MAP: // MAP logical type → MAP state → return value type MapType mapType = (MapType) logicalType; - return TypeInformation.of(getPrimitiveClass(mapType.getValueType(), columnName)); + Class mapValueClass = getPrimitiveClass(mapType.getValueType(), columnName); + if (mapValueClass == null) { + return null; + } + return TypeInformation.of(mapValueClass); default: // Primitive logical type → VALUE state → return primitive type - return TypeInformation.of(getPrimitiveClass(logicalType, columnName)); + Class primitiveClass = getPrimitiveClass(logicalType, columnName); + if (primitiveClass == null) { + return null; + } + return TypeInformation.of(primitiveClass); } } @@ -416,6 +505,9 @@ private TypeInformation inferValueTypeFromLogicalType( private Class getPrimitiveClass(LogicalType logicalType, String columnName) throws ClassNotFoundException { String className = inferTypeInfoClassFromLogicalType(columnName, logicalType); + if (className == null) { + return null; + } return Class.forName(className); } @@ -484,11 +576,9 @@ private String inferTypeInfoClassFromLogicalType(String columnName, LogicalType case UNRESOLVED: case DESCRIPTOR: 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)); + // Cannot infer a primitive class for complex/unknown types (e.g. ROW/POJO). + // Caller is expected to fall back to the snapshot serializer. + 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..bb3827019537fe --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/StateTableMapping.java @@ -0,0 +1,229 @@ +/* + * 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.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.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. Set by {@link #from} and carried through projections; {@code null} + * when the mapping was not created via the factory path. + */ + @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; + } + + /** + * Full original value columns required for state-descriptor registration / key enumeration. + * Never empty when the source table has at least one state column. + */ + @Override + public List getAllValueColumns() { + return allValueColumns; + } + + /** + * The resolved key {@link TypeInformation}, available when this mapping was created via {@link + * #from}. Returns {@code null} for mappings built through other constructors. + */ + @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. + * + * @param projectedFields output-to-source column index mapping from {@link + * org.apache.flink.table.connector.source.abilities.SupportsProjectionPushDown} + * @return remapped mapping for the projected schema + */ + @Override + public StateTableMapping project(int[][] projectedFields) { + // -1 if the key was projected out (e.g. after constant-folding with filter pushdown); + // -1 signals the key should not be written to the output row. + int newKeyColumnIndex = + KeyedTableMappingSupport.remapColumnIndex(projectedFields, this.keyColumnIndex); + List newValueColumns = + KeyedTableMappingSupport.remapValueColumns(projectedFields, this.valueColumns); + + return new StateTableMapping( + newKeyColumnIndex, keyTypeInfo, newValueColumns, allValueColumns); + } + + // ------------------------------------------------------------------------- + // Factory + // ------------------------------------------------------------------------- + + /** + * Validates the table schema, registers all per-field connector {@link ConfigOption}s into + * {@code optionalOptions} for option validation, and returns the key column index. + * + *

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. + * + * @return the index of the primary-key column in the physical row data type + */ + 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(); + Preconditions.checkArgument( + physicalDataType.getLogicalType().is(LogicalTypeRoot.ROW), + "Row data type expected."); + RowType rowType = (RowType) physicalDataType.getLogicalType(); + int keyIdx = KeyedTableMappingSupport.keyColumnIndex(physicalDataType, keyFields.get(0)); + + for (int colIdx : KeyedTableMappingSupport.valueColumnIndices(physicalDataType, keyIdx)) { + RowType.RowField valueRowField = rowType.getFields().get(colIdx); + optionalOptions.add( + KeyedTableMappingSupport.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 so + * that the schema is valid and per-field options have been registered. + * + *

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 options connector options for per-field overrides (e.g. 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 + */ + public static StateTableMapping from( + ResolvedCatalogTable catalogTable, + Configuration options, + String statePath, + OperatorIdentifier operatorIdentifier, + SerializerConfig serializerConfig) { + + SavepointTypeInfoResolver typeResolver = + KeyedTableMappingSupport.createTypeResolver( + statePath, operatorIdentifier, serializerConfig); + + DataType physicalDataType = catalogTable.getResolvedSchema().toPhysicalRowDataType(); + RowType rowType = (RowType) physicalDataType.getLogicalType(); + List keyFields = + catalogTable.getResolvedSchema().getPrimaryKey().get().getColumns(); + int keyIdx = KeyedTableMappingSupport.keyColumnIndex(physicalDataType, keyFields.get(0)); + + RowType.RowField keyRowField = rowType.getFields().get(keyIdx); + TypeInformation keyTypeInfo = typeResolver.resolveKeyType(keyRowField); + + List valueColumns = new ArrayList<>(); + for (int colIdx : KeyedTableMappingSupport.valueColumnIndices(physicalDataType, keyIdx)) { + valueColumns.add( + KeyedTableMappingSupport.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..99ea01078a3832 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; @@ -73,8 +86,20 @@ public void setStateDescriptor(StateDescriptor stateDescriptor) { this.stateDescriptor = stateDescriptor; } - @Nullable 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..c1310bdf0dea83 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/table/StateValueConverter.java @@ -0,0 +1,261 @@ +/* + * 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.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.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 org.apache.flink.shaded.guava33.com.google.common.cache.Cache; +import org.apache.flink.shaded.guava33.com.google.common.cache.CacheBuilder; + +import org.apache.avro.generic.GenericRecord; + +import java.lang.reflect.Field; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ExecutionException; +import java.util.function.Supplier; + +/** + * Converts Java objects returned by Flink state (POJOs, Avro records, primitives) into their + * internal {@link org.apache.flink.table.data.RowData} representation, driven by a target {@link + * LogicalType}. + * + *

Shared between {@link KeyedStateReader} (which converts whole rows of state values) and {@link + * FlattenedKeyedStateReader} (which converts a single list/map state's elements). + */ +@Internal +@SuppressWarnings({"rawtypes", "unchecked"}) +class StateValueConverter implements java.io.Serializable { + + private static final long serialVersionUID = 1L; + private static final long CACHE_MAX_SIZE = 64000L; + + /** + * Shared instance used by every state reader/scan provider: this class carries no state that + * differs meaningfully between instances (the reflection caches below are a pure optimization, + * rebuilt lazily and independently wherever they are actually needed), so a single instance can + * be reused instead of each caller instantiating its own. + */ + static final StateValueConverter INSTANCE = new StateValueConverter(); + + // Guava caches 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. + private transient Cache, Field> classFieldCache; + private transient Cache, Method> classMethodCache; + + private Cache, Field> classFieldCache() { + if (classFieldCache == null) { + classFieldCache = CacheBuilder.newBuilder().maximumSize(CACHE_MAX_SIZE).build(); + } + return classFieldCache; + } + + private Cache, Method> classMethodCache() { + if (classMethodCache == null) { + classMethodCache = CacheBuilder.newBuilder().maximumSize(CACHE_MAX_SIZE).build(); + } + 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}). Shared by {@link + * KeyedStateReader} and {@link WindowKeyedStateReader}. + */ + 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(); + } + } + + /** + * 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}. Shared by + * {@link FlattenedKeyedStateReader}/{@link WindowFlattenedKeyedStateReader}, whose only + * difference here is how many leading columns precede the sub-key/value pair. + */ + 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) { + GenericRowData sourceRow = (GenericRowData) object; + int targetFieldCount = ((RowType) logicalType).getFieldCount(); + // Always create a fresh INSERT row and copy by position. + // Returning the source as-is is wrong when the target field count + // coincidentally equals the source arity but the field ordering or + // semantics differ (e.g. schema evolution, column reorder). + 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 Object getObjectField(Object object, RowType.RowField rowField) { + String rowFieldName = rowField.getName(); + + // Avro GenericRecord: use the typed API directly. + if (object instanceof GenericRecord) { + return ((GenericRecord) object).get(rowFieldName); + } + + 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 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/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/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..fb3283be06af04 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/StateTableUtilsTest.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; + +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.Assert; +import org.junit.Test; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Optional; + +/** 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 with no subtasks has no keyed state → should be excluded + OperatorID opId = new OperatorID(1L, 2L); + OperatorState opState = new OperatorState(null, null, opId, 1, 128); + CheckpointMetadata metadata = + new CheckpointMetadata( + 1L, Collections.singletonList(opState), Collections.emptyList()); + + List ids = StateTableUtils.getOperatorIdentifiers(metadata); + + Assert.assertNotNull(ids); + Assert.assertTrue("Empty operator (no keyed state) should be filtered out", ids.isEmpty()); + } + + @Test + public void testGetOperatorIdentifiersEmpty() { + CheckpointMetadata metadata = + new CheckpointMetadata(3L, Collections.emptyList(), Collections.emptyList()); + + List ids = StateTableUtils.getOperatorIdentifiers(metadata); + + Assert.assertNotNull(ids); + Assert.assertTrue(ids.isEmpty()); + } + + @Test + public void testGetOperatorIdentifiersMultipleEmptyOpsAllFiltered() { + OperatorID opId1 = new OperatorID(1L, 2L); + OperatorID opId2 = new OperatorID(3L, 4L); + OperatorState opState1 = new OperatorState(null, null, opId1, 1, 128); + OperatorState opState2 = new OperatorState(null, null, opId2, 2, 128); + CheckpointMetadata metadata = + new CheckpointMetadata( + 2L, Arrays.asList(opState1, opState2), Collections.emptyList()); + + List ids = StateTableUtils.getOperatorIdentifiers(metadata); + + Assert.assertNotNull(ids); + Assert.assertTrue( + "All empty operators (no keyed state) should be filtered out", ids.isEmpty()); + } + + // ------------------------------------------------------------------------- + // 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); + Assert.assertFalse(metadata.getOperatorStates().isEmpty()); + + int operatorsWithKeyedState = 0; + for (OperatorState opState : metadata.getOperatorStates()) { + Optional detected = StateTableUtils.detectStateBackendType(opState); + if (detected.isEmpty()) { + continue; + } + operatorsWithKeyedState++; + Assert.assertEquals( + "operator " + opState.getOperatorID(), expectedType, detected.get()); + } + Assert.assertTrue( + "Expected at least one operator with keyed state in " + checkpointDir, + operatorsWithKeyedState > 0); + } +} 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..b1727c9d830df9 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/deserializer/InternalTypeConverterTest.java @@ -0,0 +1,272 @@ +/* + * 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.MapType; +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.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.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.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 (org.apache.flink.table.types.logical.LogicalType tsType : + new org.apache.flink.table.types.logical.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)); + } +} 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..efb13ee1e047d5 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/deserializer/PojoToRowDataDeserializerTest.java @@ -0,0 +1,213 @@ +/* + * 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.Assert; +import org.junit.Test; + +import java.io.IOException; +import java.util.List; + +/** + * 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); + + Assert.assertNotNull(row); + Assert.assertEquals(4, row.getArity()); + Assert.assertEquals( + StringData.fromString("Alice"), + row.getString(indexOfField(FlatPojo.class, "name"))); + Assert.assertEquals(30, row.getInt(indexOfField(FlatPojo.class, "age"))); + Assert.assertEquals(12345L, row.getLong(indexOfField(FlatPojo.class, "score"))); + Assert.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); + + Assert.assertNotNull(row); + Assert.assertTrue(row.isNullAt(indexOfField(PojoWithNullableField.class, "tag"))); + Assert.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); + Assert.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); + + Assert.assertNotNull(row); + int labelIdx = indexOfField(NestedPojo.class, "label"); + int innerIdx = indexOfField(NestedPojo.class, "inner"); + Assert.assertEquals(StringData.fromString("outer"), row.getString(labelIdx)); + + // Nested POJO should be a GenericRowData + RowData innerRow = row.getRow(innerIdx, 4); + Assert.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(PojoToRowDataDeserializer.IS_SUBCLASS); + out.writeUTF("com.example.UnknownSubclass"); + + DataInputDeserializer in = new DataInputDeserializer(out.getSharedBuffer()); + PojoToRowDataDeserializer deserializer = buildDeserializer(FlatPojo.class); + + try { + deserializer.deserialize(in); + Assert.fail("Expected IOException"); + } catch (IOException e) { + Assert.assertTrue(e.getMessage().contains("UnknownSubclass")); + } + } + + // ------------------------------------------------------------------------- + // 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); + } + + @SuppressWarnings("unchecked") + 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) throws IOException { + PojoToRowDataDeserializer deser = buildDeserializer(clazz); + LogicalType lt = + SerializerSnapshotToLogicalTypeConverter.convert( + buildPojoSerializer(clazz).snapshotConfiguration()); + RowType rt = (RowType) lt; + List names = rt.getFieldNames(); + int idx = names.indexOf(fieldName); + Assert.assertTrue("Field '" + fieldName + "' not found", idx >= 0); + return idx; + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/EmbeddedRocksDBPojoStateSchemaExtractionITCase.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/EmbeddedRocksDBPojoStateSchemaExtractionITCase.java new file mode 100644 index 00000000000000..d7b5853cb17e12 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/EmbeddedRocksDBPojoStateSchemaExtractionITCase.java @@ -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. + */ + +package org.apache.flink.state.api.schema; + +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.StateBackendOptions; + +/** Runs {@link PojoStateSchemaExtractionITCase} against the embedded RocksDB state backend. */ +public class EmbeddedRocksDBPojoStateSchemaExtractionITCase + extends PojoStateSchemaExtractionITCase { + + @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/schema/EmbeddedRocksDBStateCatalogSqlITCase.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/EmbeddedRocksDBStateCatalogSqlITCase.java new file mode 100644 index 00000000000000..26a9dedc4d3e21 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/EmbeddedRocksDBStateCatalogSqlITCase.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.schema; + +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.StateBackendOptions; + +/** Runs {@link StateCatalogSqlITCase} against the embedded RocksDB state backend. */ +public class EmbeddedRocksDBStateCatalogSqlITCase extends StateCatalogSqlITCase { + + @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/schema/HashMapPojoStateSchemaExtractionITCase.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/HashMapPojoStateSchemaExtractionITCase.java new file mode 100644 index 00000000000000..10ab066105c230 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/HashMapPojoStateSchemaExtractionITCase.java @@ -0,0 +1,33 @@ +/* + * 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.configuration.Configuration; +import org.apache.flink.configuration.StateBackendOptions; + +/** + * Runs {@link PojoStateSchemaExtractionITCase} against the heap ({@code hashmap}) state backend. + */ +public class HashMapPojoStateSchemaExtractionITCase extends PojoStateSchemaExtractionITCase { + + @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/schema/HashMapStateCatalogSqlITCase.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/HashMapStateCatalogSqlITCase.java new file mode 100644 index 00000000000000..2257f44bc8f377 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/HashMapStateCatalogSqlITCase.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.schema; + +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.StateBackendOptions; + +/** Runs {@link StateCatalogSqlITCase} against the heap ({@code hashmap}) state backend. */ +public class HashMapStateCatalogSqlITCase extends StateCatalogSqlITCase { + + @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/schema/PojoStateSchemaExtractionITCase.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/PojoStateSchemaExtractionITCase.java new file mode 100644 index 00000000000000..9b8aca3b865e65 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/PojoStateSchemaExtractionITCase.java @@ -0,0 +1,238 @@ +/* + * 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.functions.OpenContext; +import org.apache.flink.api.common.state.ValueState; +import org.apache.flink.api.common.state.ValueStateDescriptor; +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.OperatorIdentifier; +import org.apache.flink.state.api.StateTableUtils; +import org.apache.flink.state.api.input.deserializer.PojoToRowDataDeserializer; +import org.apache.flink.state.api.runtime.SavepointLoader; +import org.apache.flink.state.api.utils.SavepointTestBase; +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.types.logical.LogicalTypeRoot; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.util.Collector; + +import org.junit.jupiter.api.Test; + +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 test that verifies the full pipeline: + * + *

    + *
  1. Write POJO keyed state to a savepoint. + *
  2. Load the savepoint metadata without the POJO class on the classpath (simulated with a + * custom class loader). + *
  3. Extract state schema ({@link StateSchemaExtractor}). + *
  4. Convert to {@link RowType} ({@link SerializerSnapshotToLogicalTypeConverter}). + *
  5. Build a {@link PojoToRowDataDeserializer} ({@link PojoToRowDataDeserializer#create}). + *
  6. Verify field names and types are correct. + *
+ * + *

Subclassed per state backend (see {@code HashMapPojoStateSchemaExtractionITCase} and {@code + * EmbeddedRocksDBPojoStateSchemaExtractionITCase}) so that {@link StateSchemaExtractor} is verified + * against both the heap and RocksDB keyed-state-handle formats. + */ +public abstract class PojoStateSchemaExtractionITCase extends SavepointTestBase { + + protected abstract Configuration getConfiguration(); + + private static final String UID = "pojo-state-operator"; + + // ------------------------------------------------------------------------- + // POJO class used for state + // ------------------------------------------------------------------------- + + /** 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; + } + } + + private static final ValueStateDescriptor STATE_DESC = + new ValueStateDescriptor<>("person", PersonPojo.class); + + // ------------------------------------------------------------------------- + // Tests + // ------------------------------------------------------------------------- + + @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(UID) + .sinkTo(new DiscardingSink<>()); + + String savepointPath = takeSavepoint(env); + CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(savepointPath); + + OperatorIdentifier opId = OperatorIdentifier.forUid(UID); + + // Discover states via StateTableUtils + List stateNames = StateTableUtils.getKeyedStates(metadata, opId); + assertTrue(stateNames.contains("person"), "Expected 'person' state"); + + // Extract schema via StateTableUtils — all states in one call + 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); + + // POJO class name should be preserved in the schema info + assertNotNull(personEntry); + // The logicalType is a RowType derived from the field serializer snapshot + assertEquals(LogicalTypeRoot.ROW, personEntry.logicalType.getTypeRoot()); + } + + @Test + public void testDeserializerBuiltFromPojoSnapshot() throws Exception { + StreamExecutionEnvironment env = + StreamExecutionEnvironment.getExecutionEnvironment(getConfiguration()); + env.setParallelism(1); + + PersonPojo[] data = {new PersonPojo("Alice", 30, 100L)}; + env.addSource(createSource(data)) + .returns(PersonPojo.class) + .keyBy(p -> p.name) + .process(new PersonStateWriter()) + .uid(UID) + .sinkTo(new DiscardingSink<>()); + + String savepointPath = takeSavepoint(env); + CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(savepointPath); + + OperatorIdentifier opId = OperatorIdentifier.forUid(UID); + + KeyedStateSchemaInfo schemaInfo = StateTableUtils.getKeyedStateSchema(metadata, opId); + KeyedStateSchemaInfo.StateEntryInfo personEntry = schemaInfo.stateSchemas.get("person"); + assertNotNull(personEntry); + + // Building the PojoToRowDataDeserializer directly from the snapshot (lower-level API) + List rawSchemas = + StateSchemaExtractor.extractSchema(findOperatorState(metadata, opId)); + StateSchemaInfo personRaw = + rawSchemas.stream() + .filter(s -> "person".equals(s.stateName)) + .findFirst() + .orElse(null); + assertNotNull(personRaw); + + var deser = + PojoToRowDataDeserializer.create( + (PojoSerializerSnapshot) personRaw.valueSnapshot); + assertNotNull(deser); + assertTrue( + deser instanceof PojoToRowDataDeserializer, + "Expected PojoToRowDataDeserializer, got: " + deser.getClass().getSimpleName()); + } + + // ------------------------------------------------------------------------- + // 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 + "'"); + } + + // ------------------------------------------------------------------------- + // Operator + // ------------------------------------------------------------------------- + + private static class PersonStateWriter extends KeyedProcessFunction { + private transient ValueState state; + + @Override + public void open(OpenContext ctx) throws Exception { + state = getRuntimeContext().getState(STATE_DESC); + } + + @Override + public void processElement(PersonPojo value, Context ctx, Collector out) + throws Exception { + state.update(value); + } + } +} 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..8e8d24ed3669e1 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/SerializerSnapshotToLogicalTypeConverterTest.java @@ -0,0 +1,386 @@ +/* + * 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.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.typeutils.TypeExtractor; +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.Assert; +import org.junit.Test; + +import java.util.List; + +/** Unit tests for {@link SerializerSnapshotToLogicalTypeConverter}. */ +public class SerializerSnapshotToLogicalTypeConverterTest { + + // ------------------------------------------------------------------------- + // POJO classes + // ------------------------------------------------------------------------- + + public static class SimplePojo { + public String name; + public int age; + public long score; + public boolean active; + } + + public static class NestedPojo { + public String label; + public SimplePojo inner; + } + + // ------------------------------------------------------------------------- + // Primitive / scalar snapshots + // ------------------------------------------------------------------------- + + @Test + public void testInt() { + LogicalType t = convert(new IntSerializer.IntSerializerSnapshot()); + Assert.assertEquals(LogicalTypeRoot.INTEGER, t.getTypeRoot()); + Assert.assertFalse(t.isNullable()); + } + + @Test + public void testLong() { + LogicalType t = convert(new LongSerializer.LongSerializerSnapshot()); + Assert.assertEquals(LogicalTypeRoot.BIGINT, t.getTypeRoot()); + Assert.assertFalse(t.isNullable()); + } + + @Test + public void testFloat() { + LogicalType t = convert(new FloatSerializer.FloatSerializerSnapshot()); + Assert.assertEquals(LogicalTypeRoot.FLOAT, t.getTypeRoot()); + Assert.assertFalse(t.isNullable()); + } + + @Test + public void testDouble() { + LogicalType t = convert(new DoubleSerializer.DoubleSerializerSnapshot()); + Assert.assertEquals(LogicalTypeRoot.DOUBLE, t.getTypeRoot()); + Assert.assertFalse(t.isNullable()); + } + + @Test + public void testBoolean() { + LogicalType t = convert(new BooleanSerializer.BooleanSerializerSnapshot()); + Assert.assertEquals(LogicalTypeRoot.BOOLEAN, t.getTypeRoot()); + } + + @Test + public void testString() { + LogicalType t = convert(new StringSerializer.StringSerializerSnapshot()); + Assert.assertEquals(LogicalTypeRoot.VARCHAR, t.getTypeRoot()); + } + + // ------------------------------------------------------------------------- + // Composite types + // ------------------------------------------------------------------------- + + @Test + public void testListOfString() { + ListSerializer ser = new ListSerializer<>(StringSerializer.INSTANCE); + LogicalType t = convert(ser.snapshotConfiguration()); + Assert.assertEquals(LogicalTypeRoot.ARRAY, t.getTypeRoot()); + ArrayType at = (ArrayType) t; + Assert.assertEquals(LogicalTypeRoot.VARCHAR, at.getElementType().getTypeRoot()); + } + + @Test + public void testMapStringToLong() { + MapSerializer ser = + new MapSerializer<>(StringSerializer.INSTANCE, LongSerializer.INSTANCE); + LogicalType t = convert(ser.snapshotConfiguration()); + Assert.assertEquals(LogicalTypeRoot.MAP, t.getTypeRoot()); + MapType mt = (MapType) t; + Assert.assertEquals(LogicalTypeRoot.VARCHAR, mt.getKeyType().getTypeRoot()); + Assert.assertEquals(LogicalTypeRoot.BIGINT, mt.getValueType().getTypeRoot()); + } + + // ------------------------------------------------------------------------- + // POJO types + // ------------------------------------------------------------------------- + + @Test + public void testSimplePojo() { + LogicalType t = convertPojoType(SimplePojo.class); + Assert.assertEquals(LogicalTypeRoot.ROW, t.getTypeRoot()); + RowType rt = (RowType) t; + Assert.assertEquals(4, rt.getFieldCount()); + + assertField(rt, "name", LogicalTypeRoot.VARCHAR); + assertField(rt, "age", LogicalTypeRoot.INTEGER); + assertField(rt, "score", LogicalTypeRoot.BIGINT); + assertField(rt, "active", LogicalTypeRoot.BOOLEAN); + } + + @Test + public void testNestedPojo() { + LogicalType t = convertPojoType(NestedPojo.class); + Assert.assertEquals(LogicalTypeRoot.ROW, t.getTypeRoot()); + RowType rt = (RowType) t; + Assert.assertEquals(2, rt.getFieldCount()); + assertField(rt, "label", LogicalTypeRoot.VARCHAR); + + // The nested 'inner' field should map to ROW + RowType.RowField innerField = findField(rt, "inner"); + Assert.assertNotNull(innerField); + Assert.assertEquals(LogicalTypeRoot.ROW, innerField.getType().getTypeRoot()); + RowType innerRow = (RowType) innerField.getType(); + assertField(innerRow, "name", LogicalTypeRoot.VARCHAR); + assertField(innerRow, "age", LogicalTypeRoot.INTEGER); + } + + @Test + public 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(); + Assert.assertTrue(names.contains("name")); + Assert.assertTrue(names.contains("age")); + Assert.assertTrue(names.contains("score")); + Assert.assertTrue(names.contains("active")); + } + + // ------------------------------------------------------------------------- + // Avro types + // ------------------------------------------------------------------------- + + @Test + public void testAvroSpecificRecord() { + // AvroRecord has one field: longData (long) + LogicalType t = convertAvroSpecificType(AvroRecord.class); + Assert.assertEquals(LogicalTypeRoot.ROW, t.getTypeRoot()); + RowType rt = (RowType) t; + Assert.assertEquals(1, rt.getFieldCount()); + assertField(rt, "longData", LogicalTypeRoot.BIGINT); + } + + @Test + public void testAvroGenericRecord() { + Schema schema = AvroRecord.getClassSchema(); + LogicalType t = convertAvroGenericType(schema); + Assert.assertEquals(LogicalTypeRoot.ROW, t.getTypeRoot()); + RowType rt = (RowType) t; + Assert.assertEquals(1, rt.getFieldCount()); + assertField(rt, "longData", LogicalTypeRoot.BIGINT); + } + + @Test + public void testAvroSnapshotWithMissingClass() { + // When a specific record class is absent at read time, AvroSerializerSnapshot falls back to + // GenericRecord.class. The schema is still embedded in the snapshot, so we can infer a + // RowType from it. Simulate this by building the snapshot via GenericRecordAvroTypeInfo + // (the same logical path) and verify the correct RowType is returned. + Schema schema = AvroRecord.getClassSchema(); + LogicalType t = convertAvroGenericType(schema); + Assert.assertEquals(LogicalTypeRoot.ROW, t.getTypeRoot()); + RowType rt = (RowType) t; + Assert.assertEquals(1, rt.getFieldCount()); + assertField(rt, "longData", LogicalTypeRoot.BIGINT); + } + + // ------------------------------------------------------------------------- + // RowData types + // ------------------------------------------------------------------------- + + @Test + public 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()); + Assert.assertEquals(LogicalTypeRoot.ROW, t.getTypeRoot()); + RowType rt = (RowType) t; + Assert.assertEquals(2, rt.getFieldCount()); + assertField(rt, "id", LogicalTypeRoot.INTEGER); + assertField(rt, "name", LogicalTypeRoot.VARCHAR); + } + + @Test + public 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()); + Assert.assertEquals(LogicalTypeRoot.ROW, t.getTypeRoot()); + RowType rt = (RowType) t; + Assert.assertEquals(2, rt.getFieldCount()); + assertField(rt, "f0", LogicalTypeRoot.INTEGER); + assertField(rt, "f1", LogicalTypeRoot.BIGINT); + } + + @Test + public 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()); + Assert.assertEquals(LogicalTypeRoot.ROW, t.getTypeRoot()); + RowType rt = (RowType) t; + Assert.assertEquals(2, rt.getFieldCount()); + assertField(rt, "label", LogicalTypeRoot.VARCHAR); + + RowType.RowField innerField = findField(rt, "inner"); + Assert.assertNotNull(innerField); + Assert.assertEquals(LogicalTypeRoot.ROW, innerField.getType().getTypeRoot()); + RowType innerRow = (RowType) innerField.getType(); + assertField(innerRow, "innerId", LogicalTypeRoot.INTEGER); + assertField(innerRow, "innerName", LogicalTypeRoot.VARCHAR); + } + + // ------------------------------------------------------------------------- + // Window namespace types + // ------------------------------------------------------------------------- + + @Test + public void testTimeWindow() { + LogicalType t = + convert( + new org.apache.flink.streaming.api.windowing.windows.TimeWindow.Serializer() + .snapshotConfiguration()); + Assert.assertEquals(LogicalTypeRoot.ROW, t.getTypeRoot()); + Assert.assertFalse(t.isNullable()); + RowType rt = (RowType) t; + Assert.assertEquals(2, rt.getFieldCount()); + assertField(rt, "window_start", LogicalTypeRoot.TIMESTAMP_WITHOUT_TIME_ZONE); + assertField(rt, "window_end", LogicalTypeRoot.TIMESTAMP_WITHOUT_TIME_ZONE); + } + + @Test + public void testGlobalWindow() { + LogicalType t = + convert( + new org.apache.flink.streaming.api.windowing.windows.GlobalWindow + .Serializer() + .snapshotConfiguration()); + Assert.assertEquals(LogicalTypeRoot.ROW, t.getTypeRoot()); + Assert.assertFalse(t.isNullable()); + RowType rt = (RowType) t; + Assert.assertEquals(0, rt.getFieldCount()); + } + + // ------------------------------------------------------------------------- + // Null / unknown snapshot + // ------------------------------------------------------------------------- + + @Test + public void testNullSnapshot() { + LogicalType t = SerializerSnapshotToLogicalTypeConverter.convert(null); + Assert.assertEquals(LogicalTypeRoot.VARBINARY, t.getTypeRoot()); + } + + @Test(expected = UnsupportedOperationException.class) + public 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. + convert( + new org.apache.flink.runtime.state.VoidNamespaceSerializer() + .snapshotConfiguration()); + } + + // ------------------------------------------------------------------------- + // 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); + Assert.assertNotNull("Field '" + name + "' not found in row type", field); + Assert.assertEquals( + "Wrong type for field '" + name + "'", expectedRoot, field.getType().getTypeRoot()); + } + + 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/StateCatalogSqlITCase.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/StateCatalogSqlITCase.java new file mode 100644 index 00000000000000..06c9cc56525d47 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/StateCatalogSqlITCase.java @@ -0,0 +1,110 @@ +/* + * 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.java.tuple.Tuple2; +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.utils.SavepointTestBase; +import org.apache.flink.streaming.api.datastream.DataStream; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.streaming.api.functions.sink.v2.DiscardingSink; +import org.apache.flink.table.api.Schema; +import org.apache.flink.table.api.Table; +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.junit.jupiter.api.Test; + +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; + +/** + * Integration test verifying that {@link StateTableUtils} can extract the schema of RowData-typed + * internal SQL operator state, using the {@code accState} value state of a one-phase (non + * mini-batch) SQL {@code GROUP BY} aggregation as a representative example. + * + *

Subclassed per state backend (see {@code HashMapStateCatalogSqlITCase} and {@code + * EmbeddedRocksDBStateCatalogSqlITCase}) so keyed-state schema extraction is verified against both + * the heap and RocksDB keyed-state-handle formats. + */ +public abstract class StateCatalogSqlITCase extends SavepointTestBase { + + protected abstract Configuration getConfiguration(); + + @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()); + } +} 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..4767a12e019a47 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/schema/StateSchemaExtractorTest.java @@ -0,0 +1,206 @@ +/* + * 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.Assert; +import org.junit.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; + +/** Unit tests for {@link StateSchemaExtractor}. */ +public class StateSchemaExtractorTest { + + // ------------------------------------------------------------------------- + // Tests + // ------------------------------------------------------------------------- + + @Test + public 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); + + Assert.assertEquals(1, result.size()); + StateSchemaInfo info = result.get(0); + Assert.assertEquals("my-state", info.stateName); + Assert.assertEquals(StateDescriptor.Type.VALUE, info.stateKind); + Assert.assertNotNull(info.valueSnapshot); + Assert.assertTrue(info.valueSnapshot instanceof IntSerializer.IntSerializerSnapshot); + Assert.assertNotNull(info.keySnapshot); + Assert.assertTrue(info.keySnapshot instanceof IntSerializer.IntSerializerSnapshot); + Assert.assertNull(info.mapKeySnapshot); + } + + @Test + public 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); + + Assert.assertEquals(2, result.size()); + + Assert.assertEquals("int-state", result.get(0).stateName); + Assert.assertTrue( + result.get(0).valueSnapshot instanceof IntSerializer.IntSerializerSnapshot); + + Assert.assertEquals("str-state", result.get(1).stateName); + Assert.assertTrue( + result.get(1).valueSnapshot instanceof StringSerializer.StringSerializerSnapshot); + } + + @Test + public 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); + + Assert.assertEquals(1, result.size()); + StateSchemaInfo info = result.get(0); + Assert.assertEquals("map-state", info.stateName); + Assert.assertEquals(StateDescriptor.Type.MAP, info.stateKind); + Assert.assertNotNull(info.mapKeySnapshot); + Assert.assertTrue(info.mapKeySnapshot instanceof StringSerializer.StringSerializerSnapshot); + Assert.assertTrue(info.valueSnapshot instanceof LongSerializer.LongSerializerSnapshot); + } + + @Test + public void testEmptyStates() throws IOException { + KeyedBackendSerializationProxy proxy = + new KeyedBackendSerializationProxy<>( + IntSerializer.INSTANCE, Collections.emptyList(), false); + + List result = roundTrip(proxy); + + Assert.assertNotNull(result); + Assert.assertTrue(result.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..20e9157703d89a --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/SnapshotDiscoveryTest.java @@ -0,0 +1,352 @@ +/* + * 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.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"; + + @TempDir Path tempDir; + + private SnapshotDiscovery discovery; + + @BeforeEach + void setUp() { + discovery = new SnapshotDiscovery(Collections.singletonMap("app", tempDir.toString()), 2); + discovery.start(); + } + + @AfterEach + void tearDown() { + discovery.stop(); + } + + // ------------------------------------------------------------------------- + // Construction-time validation + // ------------------------------------------------------------------------- + + @Test + void testEmptyMapThrows() { + assertThatThrownBy(() -> new SnapshotDiscovery(Collections.emptyMap(), 2)) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + void testDuplicateDirThrows() { + Map m = new LinkedHashMap<>(); + m.put("a", "/state/app"); + m.put("b", "/state/app"); + assertThatThrownBy(() -> new SnapshotDiscovery(m, 2)) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + void testIntersectingDirsThrow() { + Map m = new LinkedHashMap<>(); + m.put("a", "/state"); + m.put("b", "/state/app"); + assertThatThrownBy(() -> new SnapshotDiscovery(m, 2)) + .isInstanceOf(IllegalArgumentException.class); + } + + // ------------------------------------------------------------------------- + // find() — defensive validation + // ------------------------------------------------------------------------- + + @Test + void testFindNullReturnsEmpty() { + assertThat(discovery.find(null)).isEmpty(); + } + + @Test + void testFindEmptyStringReturnsEmpty() { + assertThat(discovery.find("")).isEmpty(); + } + + @Test + void testFindWhitespaceOnlyReturnsEmpty() { + assertThat(discovery.find(" ")).isEmpty(); + } + + @Test + void testFindNoSlashReturnsEmpty() { + // 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 testFindSlashAtStartReturnsEmpty() { + assertThat(discovery.find("/savepoint-abc")).isEmpty(); + } + + @Test + void testFindUnknownLabelReturnsEmpty() { + assertThat(discovery.find("unknown/" + TS + "/savepoint-abc")).isEmpty(); + } + + // ------------------------------------------------------------------------- + // find() — path reconstruction + // ------------------------------------------------------------------------- + + @Test + void testFindSavepointSingleLevel() throws IOException { + createMetadataFile(tempDir.resolve("savepoint-abc")); + + assertThat(discovery.find("app/" + TS + "/savepoint-abc")) + .hasValue(tempDir.resolve("savepoint-abc").toString()); + } + + @Test + void testFindCheckpointTwoLevels() throws IOException { + String jobId = "00000000000000000000000000000001"; + createMetadataFile(tempDir.resolve(jobId).resolve("chk-3")); + + assertThat(discovery.find("app/" + TS + "/" + jobId + "/chk-3")) + .hasValue(tempDir.resolve(jobId).resolve("chk-3").toString()); + } + + @Test + void testFindDeeplyNested() throws IOException { + createMetadataFile(tempDir.resolve("a").resolve("b").resolve("c")); + + assertThat(discovery.find("app/" + TS + "/a/b/c")) + .hasValue(tempDir.resolve("a").resolve("b").resolve("c").toString()); + } + + @Test + void testFindNonexistentSnapshotReturnsEmpty() { + 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()); + } + + // ------------------------------------------------------------------------- + // find() — db-name.include-ts disabled + // ------------------------------------------------------------------------- + + @Test + void testFindWithTsDisabled() throws IOException { + SnapshotDiscovery noTs = + new SnapshotDiscovery( + Collections.singletonMap("app", tempDir.toString()), 2, false); + noTs.start(); + try { + 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(); + } finally { + noTs.stop(); + } + } + + // ------------------------------------------------------------------------- + // list() + // ------------------------------------------------------------------------- + + @Test + void testListEmptyDirectory() throws IOException { + assertThat(discovery.list()).isEmpty(); + } + + @Test + void testListSingleSnapshot() throws IOException { + createMetadataFile(tempDir.resolve("savepoint-abc")); + + assertThat(discovery.list()).containsExactly("app/" + TS + "/savepoint-abc"); + } + + @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 testListIoErrorDoesNotAffectOtherLabels() throws IOException { + Map m = new LinkedHashMap<>(); + m.put("good", tempDir.toString()); + m.put("bad", "/nonexistent-snapshot-discovery-test-path"); + SnapshotDiscovery multi = new SnapshotDiscovery(m, 2); + multi.start(); + try { + assertThat(multi.list()).isEmpty(); // good dir has no metadata files + } finally { + multi.stop(); + } + } + + @Test + void testListReturnsSnapshotsFromHealthyDirectoryWhenOtherFails() throws IOException { + createMetadataFile(tempDir.resolve("savepoint-ok")); + + Map m = new LinkedHashMap<>(); + m.put("good", tempDir.toString()); + m.put("bad", "/nonexistent-snapshot-discovery-test-path"); + SnapshotDiscovery multi = new SnapshotDiscovery(m, 2); + multi.start(); + try { + List snapshots = multi.list(); + assertThat(snapshots).containsExactly("good/" + TS + "/savepoint-ok"); + } finally { + multi.stop(); + } + } + + @Test + void testListThrowsWhenAllDirectoriesFail() { + SnapshotDiscovery allBad = + new SnapshotDiscovery( + Collections.singletonMap( + "bad", "/nonexistent-snapshot-discovery-test-path"), + 2); + allBad.start(); + try { + assertThatThrownBy(allBad::list) + .isInstanceOf(IOException.class) + .hasMessageContaining("All configured directories failed") + .hasCauseInstanceOf(IOException.class); + } finally { + allBad.stop(); + } + } + + // ------------------------------------------------------------------------- + // list() — db-name.include-ts disabled + // ------------------------------------------------------------------------- + + @Test + void testListWithTsDisabled() throws IOException { + SnapshotDiscovery noTs = + new SnapshotDiscovery( + Collections.singletonMap("app", tempDir.toString()), 2, false); + noTs.start(); + try { + createMetadataFile(tempDir.resolve("savepoint-a")); + createMetadataFile(tempDir.resolve("jobId").resolve("chk-1")); + + assertThat(noTs.list()).containsExactlyInAnyOrder("app/savepoint-a", "app/jobId/chk-1"); + } finally { + noTs.stop(); + } + } + + // ------------------------------------------------------------------------- + // creationTs — end-to-end format + // ------------------------------------------------------------------------- + + @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")); + + String dbName = discovery.list().get(0); + assertThat(dbName).matches("app/\\d{4}-\\d{2}-\\d{2}T\\d{2}:\\d{2}:\\d{2}Z/savepoint-live"); + } + + // ------------------------------------------------------------------------- + // findMetadataFiles() — all-fail error propagation + // ------------------------------------------------------------------------- + + @Test + void testFindMetadataFilesThrowsWhenAllListingsFail() { + org.apache.flink.core.fs.Path nonExistent = + new org.apache.flink.core.fs.Path(tempDir.resolve("does-not-exist").toString()); + + assertThatThrownBy(() -> discovery.findMetadataFiles(nonExistent)) + .isInstanceOf(IOException.class) + .hasMessageContaining("All directory listings failed") + .hasCauseInstanceOf(IOException.class); + } + + @Test + void testFindMetadataFilesSucceedsForEmptyDirectory() throws IOException { + org.apache.flink.core.fs.Path emptyDir = + new org.apache.flink.core.fs.Path(tempDir.toString()); + + assertThat(discovery.findMetadataFiles(emptyDir)).isEmpty(); + } + + // ------------------------------------------------------------------------- + // 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)); + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/StateCatalogITCase.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/StateCatalogITCase.java new file mode 100644 index 00000000000000..94a3dff9fc7532 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/StateCatalogITCase.java @@ -0,0 +1,1156 @@ +/* + * 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.api.java.tuple.Tuple2; +import org.apache.flink.configuration.Configuration; +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.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.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.api.internal.TableEnvironmentImpl; +import org.apache.flink.table.catalog.CatalogManager; +import org.apache.flink.table.catalog.CatalogTable; +import org.apache.flink.table.catalog.CatalogView; +import org.apache.flink.table.catalog.ObjectIdentifier; +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.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.IOException; +import java.io.OutputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.ArrayList; +import java.util.Arrays; +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}. Each test registers the catalog via {@code CREATE + * CATALOG} DDL and exercises a specific scenario against real {@code _metadata} files. + * + *

No minicluster is required — checkpoint metadata is written directly via {@link + * Checkpoints#storeCheckpointMetadata}. + */ +class StateCatalogITCase { + + // ------------------------------------------------------------------------- + // Schema-discovery savepoints (no POJO class on classpath) + // ------------------------------------------------------------------------- + + 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 int NUM_KEYS = 10; + + 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 static final int AVRO_NUM_KEYS = 10; + + // ------------------------------------------------------------------------- + // Keyed-state catalog savepoint (multiple operator types) + // ------------------------------------------------------------------------- + + 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 static final String RESOURCES_DIR = "src/test/resources/keyed-state-catalog"; + + // ------------------------------------------------------------------------- + // SQL query — reads real operator data + // ------------------------------------------------------------------------- + + @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(); + registerCatalog(tableEnv, "state", "app", tempDir); + tableEnv.executeSql("USE CATALOG state"); + + StateCatalog catalog = getCatalog(tableEnv, "state"); + String dbName = catalog.listDatabases().get(0); + tableEnv.executeSql("USE `" + dbName + "`"); + + TableResult result = tableEnv.executeSql("SELECT * FROM metadata"); + List rows = new ArrayList<>(); + try (CloseableIterator it = result.collect()) { + it.forEachRemaining(rows::add); + } + + 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(); + } + + // ------------------------------------------------------------------------- + // Discovery — multiple labels + // ------------------------------------------------------------------------- + + @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")); + + TableEnvironment tableEnv = newTableEnv(); + tableEnv.executeSql( + String.format( + "CREATE CATALOG state WITH (" + + "'type' = '%s', " + + "'directory.ckpts' = '%s', " + + "'directory.svpts' = '%s')", + StateCatalogFactory.IDENTIFIER, + checkpointsDir.toAbsolutePath(), + savepointsDir.toAbsolutePath())); + + StateCatalog catalog = getCatalog(tableEnv, "state"); + assertThat(catalog.listDatabases()) + .allSatisfy( + dbName -> + assertThat(dbName) + .matches( + "(ckpts|svpts)/\\d{4}-\\d{2}-\\d{2}T\\d{2}:\\d{2}:\\d{2}Z/savepoint-[ab]")); + catalog.close(); + } + + @Test + void testMultipleLabelsDiscoveredWithTsDisabled(@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")); + + TableEnvironment tableEnv = newTableEnv(); + tableEnv.executeSql( + String.format( + "CREATE CATALOG state WITH (" + + "'type' = '%s', " + + "'directory.ckpts' = '%s', " + + "'directory.svpts' = '%s', " + + "'db-name.include-ts' = 'false')", + StateCatalogFactory.IDENTIFIER, + checkpointsDir.toAbsolutePath(), + savepointsDir.toAbsolutePath())); + + StateCatalog catalog = getCatalog(tableEnv, "state"); + assertThat(catalog.listDatabases()) + .containsExactlyInAnyOrder("ckpts/savepoint-a", "svpts/savepoint-b"); + catalog.close(); + } + + // ------------------------------------------------------------------------- + // Catalog operations — databaseExists, tableExists, listViews, getTable + // ------------------------------------------------------------------------- + + @Test + void testCatalogOperations(@TempDir Path tempDir) throws Exception { + Path savepointDir = Files.createDirectories(tempDir.resolve("savepoint-abc")); + writeMetadata(savepointDir, 7L, Collections.emptyList()); + + TableEnvironment tableEnv = newTableEnv(); + registerCatalog(tableEnv, "state", "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(); + + assertThat(catalog.listTables(dbName)).isEmpty(); + 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 + "`"); + TableResult result = tableEnv.executeSql("SELECT * FROM metadata"); + List rows = new ArrayList<>(); + try (CloseableIterator it = result.collect()) { + it.forEachRemaining(rows::add); + } + assertThat(rows).hasSize(0); + + catalog.close(); + } + + // ------------------------------------------------------------------------- + // Schema-discovery: read keyed state without POJO class + // ------------------------------------------------------------------------- + + @Test + @SuppressWarnings("unchecked") + void testReadKeyedStateFromSchemaDiscovery() throws Exception { + Configuration config = new Configuration(); + config.set(RUNTIME_MODE, RuntimeExecutionMode.BATCH); + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(config); + StreamTableEnvironment tEnv = StreamTableEnvironment.create(env); + + CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(STATE_PATH); + KeyedStateSchemaInfo schemaInfo = + StateTableUtils.getKeyedStateSchema( + metadata, OperatorIdentifier.forUid(OPERATOR_UID)); + CatalogTable catalogTable = + StateTableUtils.getStateCatalogTable( + metadata, schemaInfo, STATE_PATH, OperatorIdentifier.forUid(OPERATOR_UID)); + + CatalogManager catalogManager = ((TableEnvironmentImpl) tEnv).getCatalogManager(); + ObjectIdentifier tableId = + catalogManager.qualifyIdentifier(UnresolvedIdentifier.of("state_table")); + catalogManager.createTemporaryTable(catalogTable, tableId, false); + + Table table = tEnv.sqlQuery("SELECT * FROM state_table"); + List result = tEnv.toDataStream(table).executeAndCollect(100); + + assertThat(result.size()).isEqualTo(NUM_KEYS); + + List keys = + result.stream() + .map(r -> (Long) r.getField("state_key")) + .collect(Collectors.toList()); + assertThat(keys) + .containsExactlyInAnyOrderElementsOf( + LongStream.range(0, NUM_KEYS).boxed().collect(Collectors.toList())); + + 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.size()).isEqualTo(1); + Row pojoRow = pojoValues.iterator().next(); + assertThat(pojoRow.getField("privateLong")).isEqualTo(1L); + assertThat(pojoRow.getField("publicLong")).isEqualTo(1L); + + Set> listValues = + result.stream() + .map( + r -> + Tuple2.of( + (Long) r.getField("state_key"), + (Long[]) r.getField("KeyedPrimitiveValueList"))) + .flatMap(l -> Set.of(l).stream()) + .collect(Collectors.toSet()); + assertThat(listValues.size()).isEqualTo(NUM_KEYS); + for (Tuple2 t : listValues) { + assertThat(t.f0).isEqualTo(t.f1[0]); + } + + Set>> mapValues = + result.stream() + .map( + r -> + Tuple2.of( + (Long) r.getField("state_key"), + (Map) + r.getField("KeyedPrimitiveValueMap"))) + .flatMap(l -> Set.of(l).stream()) + .collect(Collectors.toSet()); + assertThat(mapValues.size()).isEqualTo(NUM_KEYS); + for (Tuple2> t : mapValues) { + assertThat(t.f1.size()).isEqualTo(1); + assertThat(t.f0).isEqualTo(t.f1.get(t.f0)); + } + } + + // ------------------------------------------------------------------------- + // Flattened keyed LIST/MAP state tables — schema-discovery fixture + // ------------------------------------------------------------------------- + + @Test + void testFlattenedKeyedStateTables() throws Exception { + String catalogRoot = Paths.get(STATE_PATH).toAbsolutePath().toString(); + StateCatalog catalog = + new StateCatalog("state", Collections.singletonMap("test", catalogRoot)); + catalog.open(); + + try { + List dbs = catalog.listDatabases(); + assertThat(dbs).hasSize(1); + String dbName = dbs.get(0); + + String listTable = + StateCatalog.OPERATOR_UID_PREFIX + + OPERATOR_UID + + "_KeyedPrimitiveValueList" + + StateCatalog.FLAT_STATE_TABLE_SUFFIX; + String mapTable = + StateCatalog.OPERATOR_UID_PREFIX + + OPERATOR_UID + + "_KeyedPrimitiveValueMap" + + StateCatalog.FLAT_STATE_TABLE_SUFFIX; + + // (a) the flattened tables exist + 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(); + + var listSchema = + ((CatalogTable) catalog.getTable(new ObjectPath(dbName, listTable))) + .getUnresolvedSchema(); + assertThat( + listSchema.getColumns().stream() + .map(c -> c.getName()) + .collect(Collectors.toList())) + .containsExactly("state_key", "list_index", "list_value"); + assertThat(listSchema.getPrimaryKey()).isPresent(); + assertThat(listSchema.getPrimaryKey().get().getColumnNames()) + .containsExactly("state_key", "list_index"); + + var mapSchema = + ((CatalogTable) catalog.getTable(new ObjectPath(dbName, mapTable))) + .getUnresolvedSchema(); + assertThat( + mapSchema.getColumns().stream() + .map(c -> c.getName()) + .collect(Collectors.toList())) + .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 = TableEnvironment.create(EnvironmentSettings.inBatchMode()); + tableEnv.loadModule("state", StateModule.INSTANCE); + tableEnv.registerCatalog("state", catalog); + tableEnv.useCatalog("state"); + tableEnv.useDatabase(dbName); + + // KeyedPrimitiveValueList holds a single-element list [state_key] per key. + List listRows = collectWithSql(tableEnv, "SELECT * FROM `" + listTable + "`"); + assertThat(listRows).hasSize(NUM_KEYS); + for (Row row : listRows) { + Long key = (Long) row.getField("state_key"); + assertThat(key).isNotNull(); + assertThat(row.getField("list_index")).isEqualTo(0L); + assertThat(row.getField("list_value")).isEqualTo(key); + } + assertThat( + listRows.stream() + .map(r -> r.getField("state_key")) + .collect(Collectors.toSet())) + .containsExactlyInAnyOrderElementsOf( + LongStream.range(0, NUM_KEYS).boxed().collect(Collectors.toList())); + + // KeyedPrimitiveValueMap holds a single entry {state_key: state_key} per key. + List mapRows = collectWithSql(tableEnv, "SELECT * FROM `" + mapTable + "`"); + assertThat(mapRows).hasSize(NUM_KEYS); + for (Row row : mapRows) { + Long key = (Long) row.getField("state_key"); + assertThat(key).isNotNull(); + assertThat(row.getField("map_key")).isEqualTo(key); + assertThat(row.getField("map_value")).isEqualTo(key); + } + assertThat( + mapRows.stream() + .map(r -> r.getField("state_key")) + .collect(Collectors.toSet())) + .containsExactlyInAnyOrderElementsOf( + LongStream.range(0, NUM_KEYS).boxed().collect(Collectors.toList())); + + // 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. + List filteredListRows = + collectWithSql( + tableEnv, "SELECT * FROM `" + listTable + "` WHERE state_key = 3"); + assertThat(filteredListRows).hasSize(1); + assertThat(filteredListRows.get(0).getField("state_key")).isEqualTo(3L); + + List filteredMapRows = + collectWithSql( + tableEnv, "SELECT * FROM `" + mapTable + "` WHERE state_key = 3"); + assertThat(filteredMapRows).hasSize(1); + assertThat(filteredMapRows.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); + + KeyedStateSchemaInfo.StateEntryInfo pojoEntry = + schemaInfo.stateSchemas.get("KeyedPojoValue"); + assertThat(pojoEntry).isNotNull(); + assertThat(pojoEntry.logicalType.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType pojoRowType = (RowType) pojoEntry.logicalType; + assertThat(pojoRowType.getFieldNames()).contains("privateLong", "publicLong"); + assertThat(pojoRowType.getTypeAt(pojoRowType.getFieldIndex("privateLong")).getTypeRoot()) + .isEqualTo(LogicalTypeRoot.BIGINT); + assertThat(pojoRowType.getTypeAt(pojoRowType.getFieldIndex("publicLong")).getTypeRoot()) + .isEqualTo(LogicalTypeRoot.BIGINT); + + assertThat(schemaInfo.stateSchemas.get("KeyedPrimitiveValue").logicalType.getTypeRoot()) + .isEqualTo(LogicalTypeRoot.BIGINT); + assertThat(schemaInfo.stateSchemas.get("KeyedPrimitiveValueList").logicalType.getTypeRoot()) + .isEqualTo(LogicalTypeRoot.ARRAY); + assertThat(schemaInfo.stateSchemas.get("KeyedPrimitiveValueMap").logicalType.getTypeRoot()) + .isEqualTo(LogicalTypeRoot.MAP); + } + + @Test + @SuppressWarnings("unchecked") + void testReadAvroKeyedStateFromSchemaDiscovery() throws Exception { + Configuration config = new Configuration(); + config.set(RUNTIME_MODE, RuntimeExecutionMode.BATCH); + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(config); + StreamTableEnvironment tEnv = StreamTableEnvironment.create(env); + + CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(AVRO_STATE_PATH); + KeyedStateSchemaInfo schemaInfo = + StateTableUtils.getKeyedStateSchema( + metadata, OperatorIdentifier.forUid(AVRO_OPERATOR_UID)); + + KeyedStateSchemaInfo.StateEntryInfo specificEntry = + schemaInfo.stateSchemas.get("KeyedAvroSpecificValue"); + assertThat(specificEntry).isNotNull(); + assertThat(specificEntry.logicalType.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType specificRowType = (RowType) specificEntry.logicalType; + assertThat(specificRowType.getFieldNames()).containsExactly("longData"); + + KeyedStateSchemaInfo.StateEntryInfo genericEntry = + schemaInfo.stateSchemas.get("KeyedAvroGenericValue"); + assertThat(genericEntry).isNotNull(); + assertThat(genericEntry.logicalType.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType genericRowType = (RowType) genericEntry.logicalType; + assertThat(genericRowType.getFieldNames()).containsExactly("longData"); + + CatalogTable catalogTable = + StateTableUtils.getStateCatalogTable( + metadata, + schemaInfo, + AVRO_STATE_PATH, + OperatorIdentifier.forUid(AVRO_OPERATOR_UID)); + + CatalogManager catalogManager = ((TableEnvironmentImpl) tEnv).getCatalogManager(); + ObjectIdentifier tableId = + catalogManager.qualifyIdentifier(UnresolvedIdentifier.of("avro_state_table")); + catalogManager.createTemporaryTable(catalogTable, tableId, false); + + Table table = tEnv.sqlQuery("SELECT * FROM avro_state_table"); + List result = tEnv.toDataStream(table).executeAndCollect(100); + + assertThat(result.size()).isEqualTo(AVRO_NUM_KEYS); + + List keys = + result.stream() + .map(r -> (Long) r.getField("state_key")) + .collect(Collectors.toList()); + assertThat(keys) + .containsExactlyInAnyOrderElementsOf( + LongStream.range(0, AVRO_NUM_KEYS).boxed().collect(Collectors.toList())); + + Set specificValues = + result.stream() + .map(r -> (Row) r.getField("KeyedAvroSpecificValue")) + .collect(Collectors.toSet()); + assertThat(specificValues).hasSize(1); + assertThat(specificValues.iterator().next().getField("longData")).isEqualTo(1L); + + Set genericValues = + result.stream() + .map(r -> (Row) r.getField("KeyedAvroGenericValue")) + .collect(Collectors.toSet()); + assertThat(genericValues).hasSize(1); + assertThat(genericValues.iterator().next().getField("longData")).isEqualTo(1L); + } + + // ------------------------------------------------------------------------- + // POJO-key / Avro-key savepoint — off-classpath key types + // ------------------------------------------------------------------------- + + 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"; + + // ------------------------------------------------------------------------- + // Keyed-state catalog: schema and data for all four operator types + // ------------------------------------------------------------------------- + + @Test + void testKeyedStateCatalog() throws Exception { + String catalogRoot = Paths.get(RESOURCES_DIR).toAbsolutePath().toString(); + StateCatalog catalog = + new StateCatalog("state", Collections.singletonMap("test", catalogRoot)); + catalog.open(); + + try { + List dbs = catalog.listDatabases(); + assertThat(dbs).hasSize(1); + String dbName = dbs.get(0); + + List tables = catalog.listTables(dbName); + assertThat(tables) + .contains( + StateCatalog.OPERATOR_UID_PREFIX + + UID_PRIMITIVE + + StateCatalog.OPERATOR_TABLE_SUFFIX, + StateCatalog.OPERATOR_UID_PREFIX + + UID_POJO + + StateCatalog.OPERATOR_TABLE_SUFFIX, + StateCatalog.OPERATOR_UID_PREFIX + + UID_AVRO_SPECIFIC + + StateCatalog.OPERATOR_TABLE_SUFFIX, + StateCatalog.OPERATOR_UID_PREFIX + + UID_AVRO_GENERIC + + StateCatalog.OPERATOR_TABLE_SUFFIX); + + // metadata is a view, not a table + assertThat(catalog.listViews(dbName)).containsExactly(StateCatalog.METADATA_TABLE); + assertThat(tables).doesNotContain(StateCatalog.METADATA_TABLE); + + assertThat( + catalog.tableExists( + new ObjectPath( + dbName, + StateCatalog.OPERATOR_UID_PREFIX + + UID_PRIMITIVE + + StateCatalog.OPERATOR_TABLE_SUFFIX))) + .isTrue(); + assertThat(catalog.tableExists(new ObjectPath(dbName, StateCatalog.METADATA_TABLE))) + .isTrue(); + assertThat(catalog.tableExists(new ObjectPath(dbName, "nonexistent"))).isFalse(); + + // primitive: state_key INT, count INT + var primSchema = + ((CatalogTable) + catalog.getTable( + new ObjectPath( + dbName, + StateCatalog.OPERATOR_UID_PREFIX + + UID_PRIMITIVE + + StateCatalog.OPERATOR_TABLE_SUFFIX))) + .getUnresolvedSchema(); + assertThat( + primSchema.getColumns().stream() + .map(c -> c.getName()) + .collect(Collectors.toList())) + .contains("state_key", "count"); + + // pojo: state_key STRING, profile ROW + var pojoSchema = + ((CatalogTable) + catalog.getTable( + new ObjectPath( + dbName, + StateCatalog.OPERATOR_UID_PREFIX + + UID_POJO + + StateCatalog.OPERATOR_TABLE_SUFFIX))) + .getUnresolvedSchema(); + assertThat( + pojoSchema.getColumns().stream() + .map(c -> c.getName()) + .collect(Collectors.toList())) + .contains("state_key", "profile"); + + // avro specific + var avroSpecificSchema = + ((CatalogTable) + catalog.getTable( + new ObjectPath( + dbName, + StateCatalog.OPERATOR_UID_PREFIX + + UID_AVRO_SPECIFIC + + StateCatalog.OPERATOR_TABLE_SUFFIX))) + .getUnresolvedSchema(); + assertThat( + avroSpecificSchema.getColumns().stream() + .map(c -> c.getName()) + .collect(Collectors.toList())) + .contains("state_key", "avro_specific"); + + // avro generic + var avroGenericSchema = + ((CatalogTable) + catalog.getTable( + new ObjectPath( + dbName, + StateCatalog.OPERATOR_UID_PREFIX + + UID_AVRO_GENERIC + + StateCatalog.OPERATOR_TABLE_SUFFIX))) + .getUnresolvedSchema(); + assertThat( + avroGenericSchema.getColumns().stream() + .map(c -> c.getName()) + .collect(Collectors.toList())) + .contains("state_key", "avro_generic"); + + // Data checks via SQL + TableEnvironment tableEnv = TableEnvironment.create(EnvironmentSettings.inBatchMode()); + tableEnv.loadModule("state", StateModule.INSTANCE); + tableEnv.registerCatalog("state", catalog); + tableEnv.useCatalog("state"); + tableEnv.useDatabase(dbName); + + // Primitive state: 5 distinct int keys + List primRows = collectAll(tableEnv, UID_PRIMITIVE); + assertThat(primRows).hasSize(5); + assertThat( + primRows.stream() + .map(r -> r.getField("state_key")) + .collect(Collectors.toSet())) + .containsExactlyInAnyOrder(1, 2, 3, 4, 5); + + // POJO state: 5 rows with string keys + List pojoRows = collectAll(tableEnv, UID_POJO); + assertThat(pojoRows).hasSize(5); + assertThat( + pojoRows.stream() + .map(r -> r.getField("state_key")) + .collect(Collectors.toSet())) + .containsExactlyInAnyOrder("1", "2", "3", "4", "5"); + for (Row row : pojoRows) { + var profile = (Row) row.getField("profile"); + assertThat(profile).isNotNull(); + assertThat(profile.getField("name")).isNotNull(); + assertThat(profile.getField("score")).isNotNull(); + } + + // Avro specific: 5 rows + List avroSpecificRows = collectAll(tableEnv, UID_AVRO_SPECIFIC); + assertThat(avroSpecificRows).hasSize(5); + assertThat( + avroSpecificRows.stream() + .map(r -> r.getField("state_key")) + .collect(Collectors.toSet())) + .containsExactlyInAnyOrder("1", "2", "3", "4", "5"); + for (Row row : avroSpecificRows) { + var avroField = (Row) row.getField("avro_specific"); + assertThat(avroField).isNotNull(); + assertThat(avroField.getField("name")).isNotNull(); + assertThat(avroField.getField("value")).isNotNull(); + } + + // Avro generic: 5 rows + List avroGenericRows = collectAll(tableEnv, UID_AVRO_GENERIC); + assertThat(avroGenericRows).hasSize(5); + assertThat( + avroGenericRows.stream() + .map(r -> r.getField("state_key")) + .collect(Collectors.toSet())) + .containsExactlyInAnyOrder("1", "2", "3", "4", "5"); + for (Row row : avroGenericRows) { + var avroField = (Row) row.getField("avro_generic"); + assertThat(avroField).isNotNull(); + assertThat(avroField.getField("name")).isNotNull(); + assertThat(avroField.getField("value")).isNotNull(); + } + } finally { + catalog.close(); + } + } + + // ------------------------------------------------------------------------- + // Projection: column reorder and column subset + // ------------------------------------------------------------------------- + + @Test + void testProjectionColumnReorder() throws Exception { + TableEnvironment tableEnv = TableEnvironment.create(EnvironmentSettings.inBatchMode()); + tableEnv.loadModule("state", StateModule.INSTANCE); + String catalogRoot = Paths.get(RESOURCES_DIR).toAbsolutePath().toString(); + StateCatalog catalog = + new StateCatalog("state", Collections.singletonMap("test", catalogRoot)); + catalog.open(); + tableEnv.registerCatalog("state", catalog); + tableEnv.useCatalog("state"); + tableEnv.useDatabase(catalog.listDatabases().get(0)); + + try { + // Reorder: value column first, key column second + String primTable = + "`" + + StateCatalog.OPERATOR_UID_PREFIX + + UID_PRIMITIVE + + StateCatalog.OPERATOR_TABLE_SUFFIX + + "`"; + List primRows = + collectWithSql( + tableEnv, + "SELECT `count`, state_key FROM " + primTable + " ORDER BY state_key"); + + assertThat(primRows).hasSize(5); + for (Row row : primRows) { + assertThat(row.getArity()).isEqualTo(2); + // count is the first (index 0) field in the reordered projection + assertThat(row.getField(0)).isEqualTo(1); + assertThat(row.getField("count")).isEqualTo(1); + // state_key is the second (index 1) field + assertThat(row.getField("state_key")).isIn(1, 2, 3, 4, 5); + } + + // POJO operator: reorder profile (ROW) before state_key + String pojoTable = + "`" + + StateCatalog.OPERATOR_UID_PREFIX + + UID_POJO + + StateCatalog.OPERATOR_TABLE_SUFFIX + + "`"; + List pojoRows = + collectWithSql( + tableEnv, + "SELECT profile, state_key FROM " + pojoTable + " ORDER BY state_key"); + + assertThat(pojoRows).hasSize(5); + for (Row row : pojoRows) { + assertThat(row.getArity()).isEqualTo(2); + // profile is the first field in the reordered projection + Row profile = (Row) row.getField(0); + assertThat(profile).isNotNull(); + assertThat(profile.getField("name")).isNotNull(); + assertThat(profile.getField("score")).isNotNull(); + // state_key is the second field + assertThat(row.getField("state_key")).isNotNull(); + } + } finally { + catalog.close(); + } + } + + @Test + void testProjectionSubsetKeyOnly() throws Exception { + TableEnvironment tableEnv = TableEnvironment.create(EnvironmentSettings.inBatchMode()); + tableEnv.loadModule("state", StateModule.INSTANCE); + String catalogRoot = Paths.get(RESOURCES_DIR).toAbsolutePath().toString(); + StateCatalog catalog = + new StateCatalog("state", Collections.singletonMap("test", catalogRoot)); + catalog.open(); + tableEnv.registerCatalog("state", catalog); + tableEnv.useCatalog("state"); + tableEnv.useDatabase(catalog.listDatabases().get(0)); + + try { + String primTable = + "`" + + StateCatalog.OPERATOR_UID_PREFIX + + UID_PRIMITIVE + + StateCatalog.OPERATOR_TABLE_SUFFIX + + "`"; + List rows = + collectWithSql( + tableEnv, "SELECT state_key FROM " + primTable + " ORDER BY state_key"); + + assertThat(rows).hasSize(5); + for (Row row : rows) { + assertThat(row.getArity()).isEqualTo(1); + } + assertThat(rows.stream().map(r -> r.getField("state_key")).collect(Collectors.toList())) + .containsExactly(1, 2, 3, 4, 5); + } finally { + catalog.close(); + } + } + + @Test + void testProjectionSubsetValueOnly() throws Exception { + TableEnvironment tableEnv = TableEnvironment.create(EnvironmentSettings.inBatchMode()); + tableEnv.loadModule("state", StateModule.INSTANCE); + String catalogRoot = Paths.get(RESOURCES_DIR).toAbsolutePath().toString(); + StateCatalog catalog = + new StateCatalog("state", Collections.singletonMap("test", catalogRoot)); + catalog.open(); + tableEnv.registerCatalog("state", catalog); + tableEnv.useCatalog("state"); + tableEnv.useDatabase(catalog.listDatabases().get(0)); + + try { + String primTable = + "`" + + StateCatalog.OPERATOR_UID_PREFIX + + UID_PRIMITIVE + + StateCatalog.OPERATOR_TABLE_SUFFIX + + "`"; + List rows = collectWithSql(tableEnv, "SELECT `count` FROM " + primTable); + + assertThat(rows).hasSize(5); + for (Row row : rows) { + assertThat(row.getArity()).isEqualTo(1); + assertThat(row.getField("count")).isEqualTo(1); + } + } finally { + catalog.close(); + } + } + + // ------------------------------------------------------------------------- + // POJO key and Avro-specific key: schema discovery + // ------------------------------------------------------------------------- + + @Test + void testPojoAndAvroKeySchemaTypes() throws Exception { + Path savepointPath = + Files.list(Paths.get(POJO_AVRO_KEY_DIR)) + .filter( + p -> + Files.isDirectory(p) + && p.getFileName() + .toString() + .startsWith("savepoint-")) + .findFirst() + .orElseThrow( + () -> + new IOException( + "No savepoint found in " + POJO_AVRO_KEY_DIR)); + + 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(pojoKeyType.getTypeAt(pojoKeyType.getFieldIndex("id")).getTypeRoot()) + .isEqualTo(LogicalTypeRoot.INTEGER); + assertThat(pojoKeyType.getTypeAt(pojoKeyType.getFieldIndex("name")).getTypeRoot()) + .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"); + } + + // ------------------------------------------------------------------------- + // POJO key and Avro-specific key: end-to-end data reading via StateCatalog + // ------------------------------------------------------------------------- + + @Test + void testPojoAndAvroKeyedStateTables() throws Exception { + String catalogRoot = Paths.get(POJO_AVRO_KEY_DIR).toAbsolutePath().toString(); + StateCatalog catalog = + new StateCatalog("state", Collections.singletonMap("test", catalogRoot)); + catalog.open(); + + try { + List dbs = catalog.listDatabases(); + assertThat(dbs).hasSize(1); + String dbName = dbs.get(0); + + List tables = catalog.listTables(dbName); + assertThat(tables) + .contains( + StateCatalog.OPERATOR_UID_PREFIX + + UID_POJO_KEY + + StateCatalog.OPERATOR_TABLE_SUFFIX, + StateCatalog.OPERATOR_UID_PREFIX + + UID_AVRO_SPECIFIC_KEY + + StateCatalog.OPERATOR_TABLE_SUFFIX); + + // Both tables should have state_key and count columns + for (String uid : new String[] {UID_POJO_KEY, UID_AVRO_SPECIFIC_KEY}) { + var schema = + ((CatalogTable) + catalog.getTable( + new ObjectPath( + dbName, + StateCatalog.OPERATOR_UID_PREFIX + + uid + + StateCatalog + .OPERATOR_TABLE_SUFFIX))) + .getUnresolvedSchema(); + assertThat( + schema.getColumns().stream() + .map(c -> c.getName()) + .collect(Collectors.toList())) + .contains("state_key", "count"); + } + + TableEnvironment tableEnv = TableEnvironment.create(EnvironmentSettings.inBatchMode()); + tableEnv.loadModule("state", StateModule.INSTANCE); + tableEnv.registerCatalog("state", catalog); + tableEnv.useCatalog("state"); + tableEnv.useDatabase(dbName); + + // 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); + } + } finally { + catalog.close(); + } + } + + // ------------------------------------------------------------------------- + // TupleX key and TupleX value with mixed basic + POJO types + // ------------------------------------------------------------------------- + + 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"; + + @Test + void testTupleKeySchemaTypes() throws Exception { + Path savepointPath = + Files.list(Paths.get(TUPLE_KEY_DIR)) + .filter( + p -> + Files.isDirectory(p) + && p.getFileName() + .toString() + .startsWith("savepoint-")) + .findFirst() + .orElseThrow( + () -> new IOException("No savepoint found in " + TUPLE_KEY_DIR)); + + 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(tupleKeyType.getTypeAt(tupleKeyType.getFieldIndex("f0")).getTypeRoot()) + .isEqualTo(LogicalTypeRoot.INTEGER); + assertThat(tupleKeyType.getTypeAt(tupleKeyType.getFieldIndex("f1")).getTypeRoot()) + .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(tupleValueRowType.getTypeAt(tupleValueRowType.getFieldIndex("f0")).getTypeRoot()) + .isEqualTo(LogicalTypeRoot.BIGINT); + assertThat(tupleValueRowType.getTypeAt(tupleValueRowType.getFieldIndex("f1")).getTypeRoot()) + .isEqualTo(LogicalTypeRoot.ROW); + RowType innerPojoType = + (RowType) tupleValueRowType.getTypeAt(tupleValueRowType.getFieldIndex("f1")); + assertThat(innerPojoType.getFieldNames()).containsExactlyInAnyOrder("name", "score"); + } + + @Test + void testTupleKeyedStateTables() throws Exception { + String catalogRoot = Paths.get(TUPLE_KEY_DIR).toAbsolutePath().toString(); + StateCatalog catalog = + new StateCatalog("state", Collections.singletonMap("test", catalogRoot)); + catalog.open(); + + try { + List dbs = catalog.listDatabases(); + assertThat(dbs).hasSize(1); + String dbName = dbs.get(0); + + assertThat(catalog.listTables(dbName)) + .contains( + StateCatalog.OPERATOR_UID_PREFIX + + UID_TUPLE_KEY + + StateCatalog.OPERATOR_TABLE_SUFFIX, + StateCatalog.OPERATOR_UID_PREFIX + + UID_TUPLE_POJO_VALUE + + StateCatalog.OPERATOR_TABLE_SUFFIX); + + TableEnvironment tableEnv = TableEnvironment.create(EnvironmentSettings.inBatchMode()); + tableEnv.loadModule("state", StateModule.INSTANCE); + tableEnv.registerCatalog("state", catalog); + tableEnv.useCatalog("state"); + tableEnv.useDatabase(dbName); + + // 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); + } + } finally { + catalog.close(); + } + } + + // ------------------------------------------------------------------------- + // Helpers + // ------------------------------------------------------------------------- + + private static TableEnvironment newTableEnv() { + TableEnvironment env = TableEnvironment.create(EnvironmentSettings.inBatchMode()); + env.loadModule("state", StateModule.INSTANCE); + return env; + } + + private static void registerCatalog( + TableEnvironment tableEnv, String catalogName, String label, Path dir) { + tableEnv.executeSql( + String.format( + "CREATE CATALOG %s WITH ('type' = '%s', 'directory.%s' = '%s')", + catalogName, + StateCatalogFactory.IDENTIFIER, + label, + dir.toAbsolutePath().toString().replace("'", "''"))); + } + + private static StateCatalog getCatalog(TableEnvironment tableEnv, String name) { + return (StateCatalog) tableEnv.getCatalog(name).get(); + } + + private static Path touchMetadata(Path snapshotDir) throws IOException { + Files.createDirectories(snapshotDir); + return 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.storeCheckpointMetadata(metadata, out); + } + } + + private static List collectAll(TableEnvironment tEnv, String operatorUid) + throws Exception { + String tableName = + "`" + + StateCatalog.OPERATOR_UID_PREFIX + + operatorUid + + StateCatalog.OPERATOR_TABLE_SUFFIX + + "`"; + TableResult result = tEnv.executeSql("SELECT * FROM " + tableName); + List rows = new ArrayList<>(); + try (CloseableIterator it = result.collect()) { + it.forEachRemaining(rows::add); + } + return rows; + } + + private static List collectWithSql(TableEnvironment tEnv, String sql) throws Exception { + List rows = new ArrayList<>(); + try (CloseableIterator it = tEnv.executeSql(sql).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..581787149ec057 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/StateCatalogTest.java @@ -0,0 +1,266 @@ +/* + * 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.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 java.util.LinkedHashMap; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Unit and functional tests for {@link StateCatalog}. */ +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"; + + @TempDir Path tempDir; + + @Test + void testConstructionValidation() { + assertThatThrownBy(() -> new StateCatalog("state", Collections.emptyMap())) + .isInstanceOf(IllegalArgumentException.class); + + assertThatThrownBy( + () -> + new StateCatalog( + "state", mapOf("a", "/state/app", "b", "/state/app"))) + .isInstanceOf(IllegalArgumentException.class); + + assertThatThrownBy(() -> new StateCatalog("state", mapOf("a", "/state", "b", "/state/app"))) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + void testDatabaseNaming() throws Exception { + // savepoint-style directory + createMetadataFile(tempDir.resolve("app1").resolve("savepoint-abc123")); + StateCatalog catalog = openCatalog("app1", tempDir.resolve("app1")); + assertThat(catalog.listDatabases()).containsExactly("app1/" + TS + "/savepoint-abc123"); + catalog.close(); + + // checkpoint-style directory (label/creationTs/jobId/chk-N) + String jobId = "00000000000000000000000000000001"; + createMetadataFile(tempDir.resolve("app2").resolve(jobId).resolve("chk-3")); + catalog = openCatalog("app2", tempDir.resolve("app2")); + assertThat(catalog.listDatabases()).containsExactly("app2/" + TS + "/" + jobId + "/chk-3"); + catalog.close(); + + // arbitrary nesting + createMetadataFile(tempDir.resolve("app3").resolve("savepoints").resolve("savepoint-xyz")); + createMetadataFile( + tempDir.resolve("app3") + .resolve("checkpoints") + .resolve("a1b2c3d4") + .resolve("chk-1")); + catalog = openCatalog("staging", tempDir); + assertThat(catalog.listDatabases()) + .contains( + "staging/" + TS + "/app3/savepoints/savepoint-xyz", + "staging/" + TS + "/app3/checkpoints/a1b2c3d4/chk-1"); + catalog.close(); + + // multiple labels + createMetadataFile(tempDir.resolve("checkpoints").resolve("savepoint-a")); + createMetadataFile(tempDir.resolve("savepoints").resolve("savepoint-b")); + Map labels = new LinkedHashMap<>(); + labels.put("ckpts", tempDir.resolve("checkpoints").toAbsolutePath().toString()); + labels.put("svpts", tempDir.resolve("savepoints").toAbsolutePath().toString()); + catalog = openCatalog(labels); + assertThat(catalog.listDatabases()) + .containsExactlyInAnyOrder( + "ckpts/" + TS + "/savepoint-a", "svpts/" + TS + "/savepoint-b"); + catalog.close(); + } + + @Test + void testDatabaseNamingWithTsDisabled() throws Exception { + createMetadataFile(tempDir.resolve("savepoint-abc123")); + StateCatalog catalog = + new StateCatalog( + "state", + Collections.singletonMap("app1", tempDir.toAbsolutePath().toString()), + 2, + false); + catalog.open(); + try { + assertThat(catalog.listDatabases()).containsExactly("app1/savepoint-abc123"); + } finally { + catalog.close(); + } + } + + @Test + void testDirectoryScanning() throws Exception { + // empty directory + StateCatalog catalog = openCatalog("app1", tempDir); + assertThat(catalog.listDatabases()).isEmpty(); + catalog.close(); + + // non-metadata files are ignored + Files.createDirectories(tempDir.resolve("savepoint-a")); + Files.createFile(tempDir.resolve("savepoint-a").resolve("other.file")); + catalog = openCatalog("app1", tempDir); + assertThat(catalog.listDatabases()).isEmpty(); + catalog.close(); + + // I/O error in one directory does not affect others + createMetadataFile(tempDir.resolve("savepoint-b")); + Map labels = new LinkedHashMap<>(); + labels.put("good", tempDir.toAbsolutePath().toString()); + labels.put("bad", "/nonexistent-state-catalog-test-path"); + catalog = openCatalog(labels); + assertThat(catalog.listDatabases()).containsExactly("good/" + TS + "/savepoint-b"); + catalog.close(); + } + + @Test + void testListDatabasesReturnsEmptyWhenAllDirectoriesFail() throws Exception { + // e.g. the Flink job hasn't written any checkpoints/savepoints under this path yet + StateCatalog catalog = openCatalog("app1", tempDir.resolve("not-created-yet")); + assertThat(catalog.listDatabases()).isEmpty(); + catalog.close(); + } + + @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 testDynamicDiscovery() throws Exception { + StateCatalog catalog = openCatalog("app1", tempDir); + assertThat(catalog.listDatabases()).isEmpty(); + + createMetadataFile(tempDir.resolve("savepoint-new")); + assertThat(catalog.listDatabases()).containsExactly("app1/" + TS + "/savepoint-new"); + + Path metadataFile = tempDir.resolve("savepoint-new").resolve("_metadata"); + Files.delete(metadataFile); + assertThat(catalog.listDatabases()).isEmpty(); + + 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 = + new StateCatalog("state", Collections.singletonMap("app1", tempDir.toString())); + catalog.open(); + + assertThatThrownBy(() -> catalog.createDatabase("db", EmptyCatalogDatabase.INSTANCE, 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 final String METADATA_TABLE = StateCatalog.METADATA_TABLE; + + private static Path createMetadataFile(Path snapshotDir) throws IOException { + Files.createDirectories(snapshotDir); + Path file = Files.createFile(snapshotDir.resolve("_metadata")); + Files.setLastModifiedTime(file, FileTime.from(FIXED_TS)); + return file; + } + + private StateCatalog openCatalog(String label, Path directory) throws Exception { + return openCatalog(Collections.singletonMap(label, directory.toAbsolutePath().toString())); + } + + private StateCatalog openCatalog(Map labelsToDirs) throws Exception { + StateCatalog catalog = new StateCatalog("state", labelsToDirs, 2); + catalog.open(); + return catalog; + } + + private static Map mapOf(String k1, String v1, String k2, String v2) { + Map m = new LinkedHashMap<>(); + m.put(k1, v1); + m.put(k2, v2); + return m; + } +} 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..73796d75d3c6b5 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/TuplePojoField.java @@ -0,0 +1,58 @@ +/* + * 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 StateCatalogITCase}. + * + *

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..c06646dc7fbd53 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/table/TypeConversionDriftGuardTest.java @@ -0,0 +1,272 @@ +/* + * 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.Assert; +import org.junit.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +/** + * 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. + */ +public class TypeConversionDriftGuardTest { + + private final StateValueConverter converter = new StateValueConverter(); + + @Test + public void testInt() { + LogicalType type = new IntType(false); + Object result = converter.getValue(type, 42); + Assert.assertEquals(42, result); + } + + @Test + public void testLong() { + LogicalType type = + SerializerSnapshotToLogicalTypeConverter.convert( + new LongSerializer.LongSerializerSnapshot()); + Object result = converter.getValue(type, 42L); + Assert.assertEquals(42L, result); + } + + @Test + public void testString() { + LogicalType type = + SerializerSnapshotToLogicalTypeConverter.convert( + new StringSerializer.StringSerializerSnapshot()); + Object result = converter.getValue(type, "hello"); + Assert.assertTrue(result instanceof StringData); + Assert.assertEquals("hello", result.toString()); + } + + @Test + public void testList() { + ListSerializer ser = new ListSerializer<>(StringSerializer.INSTANCE); + LogicalType type = + SerializerSnapshotToLogicalTypeConverter.convert(ser.snapshotConfiguration()); + Assert.assertEquals(LogicalTypeRoot.ARRAY, type.getTypeRoot()); + + List value = Arrays.asList("a", "b", "c"); + Object result = converter.getValue(type, value); + Assert.assertTrue(result instanceof GenericArrayData); + Assert.assertEquals(3, ((GenericArrayData) result).size()); + } + + @Test + public void testMap() { + MapSerializer ser = + new MapSerializer<>(StringSerializer.INSTANCE, LongSerializer.INSTANCE); + LogicalType type = + SerializerSnapshotToLogicalTypeConverter.convert(ser.snapshotConfiguration()); + Assert.assertEquals(LogicalTypeRoot.MAP, type.getTypeRoot()); + + Map value = Collections.singletonMap("k", 7L); + Object result = converter.getValue(type, value); + Assert.assertTrue(result instanceof GenericMapData); + Assert.assertEquals(1, ((GenericMapData) result).size()); + } + + @Test + public void testNullableInnerType() { + ListSerializer ser = new ListSerializer<>(StringSerializer.INSTANCE); + LogicalType innerType = + SerializerSnapshotToLogicalTypeConverter.convert(ser.snapshotConfiguration()) + .copy(true); + Assert.assertNull(converter.getValue(innerType, null)); + } + + @Test + public void testTimeWindow() { + LogicalType type = + SerializerSnapshotToLogicalTypeConverter.convert( + new TimeWindow.Serializer().snapshotConfiguration()); + Assert.assertEquals(LogicalTypeRoot.ROW, type.getTypeRoot()); + + TimeWindow window = new TimeWindow(100L, 200L); + Object result = converter.getValue(type, window); + Assert.assertTrue(result instanceof GenericRowData); + GenericRowData row = (GenericRowData) result; + Assert.assertEquals(2, row.getArity()); + Assert.assertEquals(TimestampData.fromEpochMillis(100L), row.getField(0)); + Assert.assertEquals(TimestampData.fromEpochMillis(200L), row.getField(1)); + } + + @Test + public void testGlobalWindow() { + LogicalType type = + SerializerSnapshotToLogicalTypeConverter.convert( + new GlobalWindow.Serializer().snapshotConfiguration()); + Assert.assertEquals(LogicalTypeRoot.ROW, type.getTypeRoot()); + Assert.assertEquals(0, ((RowType) type).getFieldCount()); + + Object result = converter.getValue(type, GlobalWindow.get()); + Assert.assertTrue(result instanceof GenericRowData); + Assert.assertEquals(0, ((GenericRowData) result).getArity()); + } + + @Test + public 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); + Assert.assertEquals(LogicalTypeRoot.ROW, type.getTypeRoot()); + + Object result = converter.getValue(type, Tuple2.of(1, "x")); + Assert.assertTrue(result instanceof GenericRowData); + GenericRowData row = (GenericRowData) result; + Assert.assertEquals(2, row.getArity()); + Assert.assertEquals(1, row.getField(0)); + Assert.assertEquals("x", row.getField(1).toString()); + } + + /** 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 + public 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); + Assert.assertEquals(LogicalTypeRoot.ROW, type.getTypeRoot()); + + Object result = converter.getValue(type, new SamplePojo("bob", 30)); + Assert.assertTrue(result instanceof GenericRowData); + RowType rowType = (RowType) type; + GenericRowData row = (GenericRowData) result; + int nameIdx = rowType.getFieldIndex("name"); + int ageIdx = rowType.getFieldIndex("age"); + Assert.assertEquals("bob", row.getField(nameIdx).toString()); + Assert.assertEquals(30, row.getField(ageIdx)); + } + + @Test + public void testAvroSpecificRecord() { + LogicalType type = + SerializerSnapshotToLogicalTypeConverter.convert( + new AvroTypeInfo<>(AvroRecord.class) + .createSerializer(new SerializerConfigImpl()) + .snapshotConfiguration()); + Assert.assertEquals(LogicalTypeRoot.ROW, type.getTypeRoot()); + + AvroRecord avroRecord = AvroRecord.newBuilder().setLongData(99L).build(); + Object result = converter.getValue(type, avroRecord); + Assert.assertTrue(result instanceof GenericRowData); + RowType rowType = (RowType) type; + int idx = rowType.getFieldIndex("longData"); + Assert.assertEquals(99L, ((GenericRowData) result).getField(idx)); + } + + @Test + public 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()); + Assert.assertEquals(LogicalTypeRoot.ROW, type.getTypeRoot()); + + GenericRowData sourceRow = new GenericRowData(2); + sourceRow.setField(0, 1); + sourceRow.setField(1, StringData.fromString("abc")); + + Object result = converter.getValue(type, sourceRow); + Assert.assertTrue(result instanceof GenericRowData); + GenericRowData resultRow = (GenericRowData) result; + Assert.assertEquals(1, resultRow.getField(0)); + Assert.assertEquals("abc", resultRow.getField(1).toString()); + } + + @Test + public void testArrayOfArray() { + ListSerializer> serializer = + new ListSerializer<>(new ListSerializer<>(StringSerializer.INSTANCE)); + LogicalType type = + SerializerSnapshotToLogicalTypeConverter.convert( + serializer.snapshotConfiguration()); + Assert.assertEquals(LogicalTypeRoot.ARRAY, type.getTypeRoot()); + Assert.assertEquals( + LogicalTypeRoot.ARRAY, ((ArrayType) type).getElementType().getTypeRoot()); + + List> value = + Arrays.asList(Arrays.asList("a", "b"), Collections.singletonList("c")); + Object result = converter.getValue(type, value); + Assert.assertTrue(result instanceof GenericArrayData); + GenericArrayData outerArray = (GenericArrayData) result; + Assert.assertEquals(2, outerArray.size()); + Assert.assertTrue(outerArray.getArray(0) instanceof GenericArrayData); + } +} 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..61ca78ad31b194 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/resources/generator/KeyedStateCatalogSavepointGenerator.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.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.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 {@link KeyedStateCatalogITCase}. + * + *

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. Restore the {@code StateTestRecord.avsc} Avro schema to {@code + * src/test/resources/avro/StateTestRecord.avsc}. + *
  3. Re-add the {@code avro-maven-plugin} and {@code flink-avro} test dependency in pom.xml. + *
  4. Remove the {@code @Disabled} annotation and run: {@code ./mvnw test -pl apple-flink-state + * -Dtest=KeyedStateCatalogSavepointGenerator#generateSavepoint} + *
  5. Verify the new savepoint under {@code src/test/resources/keyed-state-catalog/}. + *
  6. Remove the copied {@code .java} file from the source tree and revert pom.xml changes. + *
+ */ +@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"; + + @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(KeyedStateCatalogITCase.UID_PRIMITIVE) + .name(KeyedStateCatalogITCase.UID_PRIMITIVE) + .keyBy(v -> String.valueOf(v)) + .process(new ProfileOperator()) + .uid(KeyedStateCatalogITCase.UID_POJO) + .name(KeyedStateCatalogITCase.UID_POJO) + .keyBy(v -> String.valueOf(v)) + .process(new AvroSpecificOperator()) + .uid(KeyedStateCatalogITCase.UID_AVRO_SPECIFIC) + .name(KeyedStateCatalogITCase.UID_AVRO_SPECIFIC) + .keyBy(v -> String.valueOf(v)) + .process(new AvroGenericOperator()) + .uid(KeyedStateCatalogITCase.UID_AVRO_GENERIC) + .name(KeyedStateCatalogITCase.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, org.apache.flink.util.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, org.apache.flink.util.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, org.apache.flink.util.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, org.apache.flink.util.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..5b53b515ccbcb0 --- /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.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 {@link StateCatalogITCase} 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 {@link StateCatalogITCase}. */ + 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, org.apache.flink.util.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, org.apache.flink.util.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..e9a9f275b91cbf --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/resources/generator/KeyedStateTupleKeySavepointGenerator.java @@ -0,0 +1,278 @@ +/* + * 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.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 {@link StateCatalogITCase} 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 {@link StateCatalogITCase}. */ + 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, org.apache.flink.util.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, org.apache.flink.util.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..70b5b3661e8eea 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()}. Callers that need to restore the reading context for a key + * obtained this way (e.g. via {@link #setCurrentKeyAndKeyGroup}) must use this key-group + * instead of recomputing it, since a key substituted during deserialization (e.g. read back + * without its original class on the classpath) may not reproduce the original {@code + * hashCode()}. + * @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..81881a5a954e46 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; + } + // This ensures key deduplication across all table entry + // keys. Look up the entry in its own key-group rather than + // recomputing the key-group 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..7d20c65ba132b3 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 (as determined + * by which {@link #keyGroupedStateMaps} slot holds them), rather than one 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/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/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/ForStSyncKeyedStateBackend.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/sync/ForStSyncKeyedStateBackend.java index 4f410ca4287ae9..7d45cf0fefc968 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 @@ -387,6 +387,12 @@ public Stream getKeys(List states, N namespace) { throw new UnsupportedOperationException(); } + @Override + public Stream> getKeysAndKeyGroups(List states, N namespace) { + // TODO + throw new UnsupportedOperationException(); + } + @Override public Stream> getKeysAndNamespaces(String state) { ForStOperationUtils.ForStKvStateInfo columnInfo = kvStateInformation.get(state); 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..45301ba26fc7e1 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,17 @@ public K next() { return tmpKey; } + /** + * Returns the key-group of the key most recently returned by {@link #next()}, as physically + * decoded from the RocksDB key's byte prefix rather than recomputed from {@code hashCode()}. + * The decoding is only performed when this method is actually invoked, so callers that never + * need the key-group (e.g. the plain key iteration codepath) pay no extra cost. Only valid + * immediately after a call to {@link #next()}. + */ + 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..be19f6cbaa1458 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} paired with an explicitly provided + * serializer. Use this when the serializer cannot be derived from the data type alone, e.g. + * when it was restored directly 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..e5ead49ed3b37e 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,23 @@ 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 field names are not available (the originating + * serializer was built from a bare {@link LogicalType} array, or the snapshot was persisted + * before field names were tracked). + */ + @Internal + @Nullable + public String[] getFieldNames() { + return fieldNames; + } } }