This is an automated email from the ASF dual-hosted git repository. gyfora pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/flink.git
commit d28f941743b598398fd9388a9cebc379049f338d Author: Gyula Fora <[email protected]> AuthorDate: Sun Jul 19 21:49:24 2026 +0200 [FLINK-40178][state-processor-api] Keyed state reading integration tests Adds end-to-end integration tests for the general and flattened keyed-state tables introduced in the previous two commits, exercising real (generated) savepoint fixtures end-to-end via both direct StateTableUtils/CatalogTable construction and StateCatalog + SQL: - testReadKeyedStateFromSchemaDiscovery / testReadAvroKeyedStateFromSchemaDiscovery / testSchemaExtractionWithoutPojoClass: schema discovery and data reading for a savepoint whose POJO/Avro classes are unavailable on the classpath. - testKeyedStateCatalog / testPojoAndAvroKeyedStateTables / testTupleKeyedStateTables: full StateCatalog-backed catalog/table discovery and SQL reads across primitive, POJO, Avro-specific, Avro-generic, and Tuple key/value shapes. - testFlattenedKeyedStateTables: flattened LIST/MAP state table schema and SQL reads, including state_key filter push-down. - testProjectionColumnReorder / testProjectionSubsetKeyOnly / testProjectionSubsetValueOnly: projection push-down correctness for the general keyed-state table. - testPojoAndAvroKeySchemaTypes / testTupleKeySchemaTypes: key-type schema inference for POJO, Avro-specific, and Tuple key types. The savepoint fixtures are produced by the generator programs under src/test/resources/generator (not run as part of the build; their output is checked in), plus the pre-existing missing-class/missing-avro fixtures from the schema-discovery bootstrapping commit. --- .../flink/state/api/KeyedStateReadingITCase.java | 187 +++++- .../StateCatalogGeneratedSavepointITCase.java | 691 +++++++++++++++++++++ .../apache/flink/state/catalog/TuplePojoField.java | 59 ++ .../KeyedStateCatalogSavepointGenerator.java | 356 +++++++++++ .../KeyedStatePojoAvroKeySavepointGenerator.java | 321 ++++++++++ .../KeyedStateTupleKeySavepointGenerator.java | 277 +++++++++ .../test/resources/generator/StateTestRecord.avsc | 32 + .../savepoint-01d134-a82d2259b86b/_metadata | Bin 0 -> 17739 bytes .../savepoint-07a18b-da0f2ab0e5e0/_metadata | Bin 0 -> 12750 bytes .../savepoint-515de8-42f928682f3b/_metadata | Bin 0 -> 11181 bytes .../resources/table-state-missing-avro/_metadata | Bin 0 -> 4391 bytes .../resources/table-state-missing-class/_metadata | Bin 0 -> 4453 bytes 12 files changed, 1890 insertions(+), 33 deletions(-) diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/KeyedStateReadingITCase.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/KeyedStateReadingITCase.java index 5936b30bb5d..4678ade760a 100644 --- a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/KeyedStateReadingITCase.java +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/KeyedStateReadingITCase.java @@ -19,6 +19,10 @@ package org.apache.flink.state.api; import org.apache.flink.api.common.functions.OpenContext; +import org.apache.flink.api.common.state.ListState; +import org.apache.flink.api.common.state.ListStateDescriptor; +import org.apache.flink.api.common.state.MapState; +import org.apache.flink.api.common.state.MapStateDescriptor; import org.apache.flink.api.common.state.ValueState; import org.apache.flink.api.common.state.ValueStateDescriptor; import org.apache.flink.api.java.tuple.Tuple2; @@ -32,20 +36,30 @@ import org.apache.flink.state.api.schema.KeyedStateSchemaInfo; import org.apache.flink.state.api.schema.StateSchemaExtractor; import org.apache.flink.state.api.schema.StateSchemaInfo; import org.apache.flink.state.api.utils.SavepointTestBase; +import org.apache.flink.state.catalog.StateCatalog; +import org.apache.flink.state.table.module.StateModule; import org.apache.flink.streaming.api.datastream.DataStream; import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; import org.apache.flink.streaming.api.functions.KeyedProcessFunction; import org.apache.flink.streaming.api.functions.sink.v2.DiscardingSink; +import org.apache.flink.table.api.EnvironmentSettings; import org.apache.flink.table.api.Schema; import org.apache.flink.table.api.Table; +import org.apache.flink.table.api.TableEnvironment; +import org.apache.flink.table.api.TableResult; import org.apache.flink.table.api.bridge.java.StreamTableEnvironment; import org.apache.flink.table.types.logical.LogicalTypeRoot; import org.apache.flink.table.types.logical.RowType; import org.apache.flink.types.Row; +import org.apache.flink.util.CloseableIterator; import org.apache.flink.util.Collector; import org.junit.jupiter.api.Test; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.Objects; @@ -59,6 +73,11 @@ import static org.junit.jupiter.api.Assertions.assertTrue; * backends (see {@code HashMapKeyedStateReadingITCase} / {@code * EmbeddedRocksDBKeyedStateReadingITCase}) so that schema extraction and reads are checked against * both keyed-state-handle formats. + * + * <p>Unlike {@code StateCatalogGeneratedSavepointITCase}, which reads savepoints checked in as test + * resources (necessarily HashMap-only, since RocksDB fixtures can't be generated locally), every + * savepoint here is produced at test run time, so it runs on whichever backend the subclass + * configures. */ public abstract class KeyedStateReadingITCase extends SavepointTestBase { @@ -177,11 +196,9 @@ public abstract class KeyedStateReadingITCase extends SavepointTestBase { OperatorIdentifier opId = OperatorIdentifier.forUid(POJO_UID); - // Discover states via StateTableUtils List<String> 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"); @@ -193,48 +210,120 @@ public abstract class KeyedStateReadingITCase extends SavepointTestBase { assertHasField(rowType, "name", LogicalTypeRoot.VARCHAR); assertHasField(rowType, "age", LogicalTypeRoot.INTEGER); assertHasField(rowType, "score", LogicalTypeRoot.BIGINT); + + // The same snapshot must also be usable to build a deserializer directly (lower-level API). + StateSchemaInfo personRaw = + StateSchemaExtractor.extractSchema(findOperatorState(metadata, opId)).stream() + .filter(s -> "person".equals(s.stateName)) + .findFirst() + .orElse(null); + assertNotNull(personRaw); + assertNotNull( + PojoToRowDataDeserializer.create( + (PojoSerializerSnapshot<?>) personRaw.valueSnapshot)); } + // ------------------------------------------------------------------------- + // End-to-end read through StateCatalog + SQL: primitive, POJO, list and map state + // ------------------------------------------------------------------------- + // + // This is the backend-parameterized equivalent of + // StateCatalogGeneratedSavepointITCase.SchemaDiscoveryWithoutSourceClasses — same state + // shapes, but written and savepointed at test run time instead of read from a checked-in + // HashMap-only fixture, so it also runs on RocksDB. + + private static final String MIXED_STATE_UID = "mixed-state-operator"; + private static final ValueStateDescriptor<Long> COUNT_STATE_DESC = + new ValueStateDescriptor<>("count", Long.class); + private static final ListStateDescriptor<Long> ITEMS_STATE_DESC = + new ListStateDescriptor<>("items", Long.class); + private static final MapStateDescriptor<Long, Long> COUNTS_STATE_DESC = + new MapStateDescriptor<>("counts", Long.class, Long.class); + @Test - public void testDeserializerBuiltFromPojoSnapshot() throws Exception { + public void testReadPrimitivePojoListAndMapStateThroughCatalog() 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(POJO_UID) + Long[] keys = {1L, 2L, 3L}; + env.addSource(createSource(keys)) + .returns(Long.class) + .keyBy(k -> k) + .process(new MixedStateWriter()) + .uid(MIXED_STATE_UID) .sinkTo(new DiscardingSink<>()); String savepointPath = takeSavepoint(env); - CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(savepointPath); - - OperatorIdentifier opId = OperatorIdentifier.forUid(POJO_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<StateSchemaInfo> rawSchemas = - StateSchemaExtractor.extractSchema(findOperatorState(metadata, opId)); - StateSchemaInfo personRaw = - rawSchemas.stream() - .filter(s -> "person".equals(s.stateName)) - .findFirst() - .orElse(null); - assertNotNull(personRaw); + // takeSavepoint() returns a "file:" URI string, not a plain filesystem path. + Path catalogRoot = Paths.get(java.net.URI.create(savepointPath)).getParent(); + + StateCatalog catalog = + new StateCatalog( + "state", + Collections.singletonMap("test", catalogRoot.toAbsolutePath().toString())); + catalog.open(); + try { + List<String> dbs = catalog.listDatabases(); + assertEquals(1, dbs.size()); + String dbName = dbs.get(0); + + TableEnvironment tableEnv = TableEnvironment.create(EnvironmentSettings.inBatchMode()); + tableEnv.loadModule("state", StateModule.INSTANCE); + tableEnv.registerCatalog("state", catalog); + tableEnv.useCatalog("state"); + tableEnv.useDatabase(dbName); + + String mainTable = + "`" + + StateCatalog.OPERATOR_UID_PREFIX + + MIXED_STATE_UID + + StateCatalog.OPERATOR_TABLE_SUFFIX + + "`"; + List<Row> rows = collectWithSql(tableEnv, "SELECT * FROM " + mainTable); + assertEquals(3, rows.size()); + for (Row row : rows) { + Long key = (Long) row.getField("state_key"); + assertNotNull(key); + assertEquals(key, row.getField("count")); + + Row person = (Row) row.getField("person"); + assertNotNull(person); + assertEquals("name-" + key, person.getField("name")); + assertEquals(key * 100, person.getField("score")); + } - var deser = - PojoToRowDataDeserializer.create( - (PojoSerializerSnapshot<?>) personRaw.valueSnapshot); - assertNotNull(deser); - assertTrue( - deser instanceof PojoToRowDataDeserializer, - "Expected PojoToRowDataDeserializer, got: " + deser.getClass().getSimpleName()); + String listFlatTable = + "`" + + StateCatalog.OPERATOR_UID_PREFIX + + MIXED_STATE_UID + + "_items" + + StateCatalog.FLAT_STATE_TABLE_SUFFIX + + "`"; + List<Row> listRows = + collectWithSql( + tableEnv, "SELECT * FROM " + listFlatTable + " WHERE state_key = 2"); + assertEquals(1, listRows.size()); + assertEquals(2L, listRows.get(0).getField("state_key")); + assertEquals(20L, listRows.get(0).getField("list_value")); + + String mapFlatTable = + "`" + + StateCatalog.OPERATOR_UID_PREFIX + + MIXED_STATE_UID + + "_counts" + + StateCatalog.FLAT_STATE_TABLE_SUFFIX + + "`"; + List<Row> mapRows = + collectWithSql( + tableEnv, "SELECT * FROM " + mapFlatTable + " WHERE state_key = 2"); + assertEquals(1, mapRows.size()); + assertEquals(2L, mapRows.get(0).getField("state_key")); + assertEquals(2L, mapRows.get(0).getField("map_key")); + assertEquals(2L, mapRows.get(0).getField("map_value")); + } finally { + catalog.close(); + } } // ------------------------------------------------------------------------- @@ -262,6 +351,15 @@ public abstract class KeyedStateReadingITCase extends SavepointTestBase { expectedRoot, field.getType().getTypeRoot(), "Wrong type for field '" + name + "'"); } + private static List<Row> collectWithSql(TableEnvironment tEnv, String sql) throws Exception { + List<Row> rows = new ArrayList<>(); + TableResult result = tEnv.executeSql(sql); + try (CloseableIterator<Row> it = result.collect()) { + it.forEachRemaining(rows::add); + } + return rows; + } + // ------------------------------------------------------------------------- // Operators // ------------------------------------------------------------------------- @@ -280,4 +378,27 @@ public abstract class KeyedStateReadingITCase extends SavepointTestBase { state.update(value); } } + + private static class MixedStateWriter extends KeyedProcessFunction<Long, Long, Void> { + private transient ValueState<Long> countState; + private transient ValueState<PersonPojo> personState; + private transient ListState<Long> itemsState; + private transient MapState<Long, Long> countsState; + + @Override + public void open(OpenContext ctx) throws Exception { + countState = getRuntimeContext().getState(COUNT_STATE_DESC); + personState = getRuntimeContext().getState(PERSON_STATE_DESC); + itemsState = getRuntimeContext().getListState(ITEMS_STATE_DESC); + countsState = getRuntimeContext().getMapState(COUNTS_STATE_DESC); + } + + @Override + public void processElement(Long key, Context ctx, Collector<Void> out) throws Exception { + countState.update(key); + personState.update(new PersonPojo("name-" + key, key.intValue(), key * 100)); + itemsState.add(key * 10); + countsState.put(key, key); + } + } } diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/StateCatalogGeneratedSavepointITCase.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/StateCatalogGeneratedSavepointITCase.java new file mode 100644 index 00000000000..c56fd71365a --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/StateCatalogGeneratedSavepointITCase.java @@ -0,0 +1,691 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.state.catalog; + +import org.apache.flink.api.common.RuntimeExecutionMode; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.runtime.checkpoint.metadata.CheckpointMetadata; +import org.apache.flink.state.api.OperatorIdentifier; +import org.apache.flink.state.api.StateTableUtils; +import org.apache.flink.state.api.runtime.SavepointLoader; +import org.apache.flink.state.api.schema.KeyedStateSchemaInfo; +import org.apache.flink.state.table.module.StateModule; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.table.api.EnvironmentSettings; +import org.apache.flink.table.api.Schema; +import org.apache.flink.table.api.TableEnvironment; +import org.apache.flink.table.api.TableResult; +import org.apache.flink.table.api.bridge.java.StreamTableEnvironment; +import org.apache.flink.table.api.internal.TableEnvironmentImpl; +import org.apache.flink.table.catalog.CatalogManager; +import org.apache.flink.table.catalog.CatalogTable; +import org.apache.flink.table.catalog.ObjectPath; +import org.apache.flink.table.catalog.UnresolvedIdentifier; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.LogicalTypeRoot; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.types.Row; +import org.apache.flink.util.CloseableIterator; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import java.util.stream.LongStream; + +import static org.apache.flink.configuration.ExecutionOptions.RUNTIME_MODE; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Integration tests for {@link StateCatalog} and {@link StateTableUtils} against real keyed-state + * savepoints that are checked in as test resources, grouped by fixture/scenario in {@code @Nested} + * classes. + * + * <p>These savepoints were produced once with the {@code hashmap} state backend by the (disabled, + * manually-run) generator programs under {@code src/test/resources/generator} and are checked in + * under {@code src/test/resources/}. They cannot be regenerated for the RocksDB state backend + * without running those generators locally against a RocksDB-configured job, so every test in this + * class is inherently HashMap-only — see {@code KeyedStateReadingITCase} for the + * RocksDB-parameterized equivalent exercised against savepoints taken at runtime instead of + * checked-in fixtures. + */ +class StateCatalogGeneratedSavepointITCase { + + /** + * Schema discovery and reads for a savepoint whose POJO/Avro classes aren't on the classpath. + */ + @Nested + class SchemaDiscoveryWithoutSourceClasses { + + private static final String STATE_PATH = "src/test/resources/table-state-missing-class"; + private static final String OPERATOR_UID = "missing-class-operator"; + private static final String AVRO_STATE_PATH = "src/test/resources/table-state-missing-avro"; + private static final String AVRO_OPERATOR_UID = "missing-avro-operator"; + private final String[] avroStateNames = {"KeyedAvroSpecificValue", "KeyedAvroGenericValue"}; + private static final int NUM_KEYS = 10; + + @Test + @SuppressWarnings("unchecked") + void testReadKeyedStateFromSchemaDiscovery() throws Exception { + List<Row> result = readViaTemporaryTable(STATE_PATH, OPERATOR_UID, "state_table"); + + assertThat(result).hasSize(NUM_KEYS); + assertThat(stateKeys(result)).containsExactlyInAnyOrderElementsOf(longRange(NUM_KEYS)); + + Set<Long> primitiveValues = + result.stream() + .map(r -> (Long) r.getField("KeyedPrimitiveValue")) + .collect(Collectors.toSet()); + assertThat(primitiveValues).containsExactly(1L); + + Set<Row> pojoValues = + result.stream() + .map(r -> (Row) r.getField("KeyedPojoValue")) + .collect(Collectors.toSet()); + assertThat(pojoValues).hasSize(1); + Row pojoRow = pojoValues.iterator().next(); + assertThat(pojoRow.getField("privateLong")).isEqualTo(1L); + assertThat(pojoRow.getField("publicLong")).isEqualTo(1L); + + // Each key holds the single-element list [state_key] and the single map entry + // {state_key: state_key}. + for (Row row : result) { + Long key = (Long) row.getField("state_key"); + assertThat((Long[]) row.getField("KeyedPrimitiveValueList")).containsExactly(key); + assertThat((Map<Long, Long>) row.getField("KeyedPrimitiveValueMap")) + .containsExactly(Map.entry(key, key)); + } + } + + @Test + void testFlattenedKeyedStateTables() throws Exception { + StateCatalog catalog = openCatalogOn(STATE_PATH); + try { + String dbName = catalog.listDatabases().get(0); + String listTable = flatKeyedTable(OPERATOR_UID, "KeyedPrimitiveValueList"); + String mapTable = flatKeyedTable(OPERATOR_UID, "KeyedPrimitiveValueMap"); + + // (a) the flattened tables exist and expose a composite primary key + assertThat(catalog.listTables(dbName)).contains(listTable, mapTable); + assertThat(catalog.tableExists(new ObjectPath(dbName, listTable))).isTrue(); + assertThat(catalog.tableExists(new ObjectPath(dbName, mapTable))).isTrue(); + assertThat(catalog.tableExists(new ObjectPath(dbName, listTable + "-nonexistent"))) + .isFalse(); + + Schema listSchema = schemaOf(catalog, dbName, listTable); + assertThat(columnNames(listSchema)) + .containsExactly("state_key", "list_index", "list_value"); + assertThat(listSchema.getPrimaryKey()).isPresent(); + assertThat(listSchema.getPrimaryKey().get().getColumnNames()) + .containsExactly("state_key", "list_index"); + + Schema mapSchema = schemaOf(catalog, dbName, mapTable); + assertThat(columnNames(mapSchema)) + .containsExactly("state_key", "map_key", "map_value"); + assertThat(mapSchema.getPrimaryKey()).isPresent(); + assertThat(mapSchema.getPrimaryKey().get().getColumnNames()) + .containsExactly("state_key", "map_key"); + + // (b) the flattened tables can be read correctly and return the expected data + TableEnvironment tableEnv = newCatalogTableEnv(catalog, dbName); + + // KeyedPrimitiveValueList holds a single-element list [state_key] per key. + List<Row> listRows = collectWithSql(tableEnv, "SELECT * FROM `" + listTable + "`"); + assertThat(listRows).hasSize(NUM_KEYS); + assertThat(stateKeys(listRows)) + .containsExactlyInAnyOrderElementsOf(longRange(NUM_KEYS)); + for (Row row : listRows) { + assertThat(row.getField("list_index")).isEqualTo(0L); + assertThat(row.getField("list_value")).isEqualTo(row.getField("state_key")); + } + + // KeyedPrimitiveValueMap holds a single entry {state_key: state_key} per key. + List<Row> mapRows = collectWithSql(tableEnv, "SELECT * FROM `" + mapTable + "`"); + assertThat(mapRows).hasSize(NUM_KEYS); + assertThat(stateKeys(mapRows)) + .containsExactlyInAnyOrderElementsOf(longRange(NUM_KEYS)); + for (Row row : mapRows) { + assertThat(row.getField("map_key")).isEqualTo(row.getField("state_key")); + assertThat(row.getField("map_value")).isEqualTo(row.getField("state_key")); + } + + // state_key filter push-down (SupportsFilterPushDown) prunes to a single key even + // though state_key is only part of the composite (state_key, list_index/map_key) + // primary key in the flattened schema. + for (String table : new String[] {listTable, mapTable}) { + List<Row> filtered = + collectWithSql( + tableEnv, "SELECT * FROM `" + table + "` WHERE state_key = 3"); + assertThat(filtered).hasSize(1); + assertThat(filtered.get(0).getField("state_key")).isEqualTo(3L); + } + } finally { + catalog.close(); + } + } + + @Test + void testSchemaExtractionWithoutPojoClass() throws Exception { + CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(STATE_PATH); + KeyedStateSchemaInfo schemaInfo = + StateTableUtils.getKeyedStateSchema( + metadata, OperatorIdentifier.forUid(OPERATOR_UID)); + + assertThat(schemaInfo.keyType.getTypeRoot()).isEqualTo(LogicalTypeRoot.BIGINT); + assertThat(stateTypeRoot(schemaInfo, "KeyedPrimitiveValue")) + .isEqualTo(LogicalTypeRoot.BIGINT); + assertThat(stateTypeRoot(schemaInfo, "KeyedPrimitiveValueList")) + .isEqualTo(LogicalTypeRoot.ARRAY); + assertThat(stateTypeRoot(schemaInfo, "KeyedPrimitiveValueMap")) + .isEqualTo(LogicalTypeRoot.MAP); + + assertThat(stateTypeRoot(schemaInfo, "KeyedPojoValue")).isEqualTo(LogicalTypeRoot.ROW); + RowType pojoRowType = + (RowType) schemaInfo.stateSchemas.get("KeyedPojoValue").logicalType; + assertThat(pojoRowType.getFieldNames()).contains("privateLong", "publicLong"); + assertThat(fieldTypeRoot(pojoRowType, "privateLong")).isEqualTo(LogicalTypeRoot.BIGINT); + assertThat(fieldTypeRoot(pojoRowType, "publicLong")).isEqualTo(LogicalTypeRoot.BIGINT); + } + + @Test + void testReadAvroKeyedStateFromSchemaDiscovery() throws Exception { + CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(AVRO_STATE_PATH); + KeyedStateSchemaInfo schemaInfo = + StateTableUtils.getKeyedStateSchema( + metadata, OperatorIdentifier.forUid(AVRO_OPERATOR_UID)); + + // Both the specific-record and the generic-record state degrade to ROW(longData). + for (String stateName : avroStateNames) { + assertThat(stateTypeRoot(schemaInfo, stateName)).isEqualTo(LogicalTypeRoot.ROW); + RowType rowType = (RowType) schemaInfo.stateSchemas.get(stateName).logicalType; + assertThat(rowType.getFieldNames()).containsExactly("longData"); + } + + List<Row> result = + readViaTemporaryTable(AVRO_STATE_PATH, AVRO_OPERATOR_UID, "avro_state_table"); + + assertThat(result).hasSize(NUM_KEYS); + assertThat(stateKeys(result)).containsExactlyInAnyOrderElementsOf(longRange(NUM_KEYS)); + for (String stateName : avroStateNames) { + Set<Row> values = + result.stream() + .map(r -> (Row) r.getField(stateName)) + .collect(Collectors.toSet()); + assertThat(values).hasSize(1); + assertThat(values.iterator().next().getField("longData")).isEqualTo(1L); + } + } + } + + /** Schema and data for four operator types (primitive, POJO, Avro-specific, Avro-generic). */ + @Nested + class MultiOperatorTypeCatalog { + + private static final String RESOURCES_DIR = "src/test/resources/keyed-state-catalog"; + private static final String UID_PRIMITIVE = "primitive-state-op"; + private static final String UID_POJO = "pojo-state-op"; + private static final String UID_AVRO_SPECIFIC = "avro-specific-state-op"; + private static final String UID_AVRO_GENERIC = "avro-generic-state-op"; + + private StateCatalog catalog; + private TableEnvironment tableEnv; + private String dbName; + + @BeforeEach + void openCatalog() throws Exception { + catalog = openCatalogOn(RESOURCES_DIR); + dbName = catalog.listDatabases().get(0); + tableEnv = newCatalogTableEnv(catalog, dbName); + } + + @AfterEach + void closeCatalog() { + catalog.close(); + } + + @Test + void testKeyedStateCatalog() throws Exception { + List<String> tables = catalog.listTables(dbName); + assertThat(tables) + .contains( + keyedTable(UID_PRIMITIVE), + keyedTable(UID_POJO), + keyedTable(UID_AVRO_SPECIFIC), + keyedTable(UID_AVRO_GENERIC)); + + // metadata is a view, but listTables includes views too, per the Catalog contract. + assertThat(catalog.listViews(dbName)).containsExactly(StateCatalog.METADATA_TABLE); + assertThat(tables).contains(StateCatalog.METADATA_TABLE); + + assertThat(catalog.tableExists(new ObjectPath(dbName, keyedTable(UID_PRIMITIVE)))) + .isTrue(); + assertThat(catalog.tableExists(new ObjectPath(dbName, StateCatalog.METADATA_TABLE))) + .isTrue(); + assertThat(catalog.tableExists(new ObjectPath(dbName, "nonexistent"))).isFalse(); + + assertThat(columnNames(schemaOf(catalog, dbName, keyedTable(UID_PRIMITIVE)))) + .contains("state_key", "count"); + assertThat(columnNames(schemaOf(catalog, dbName, keyedTable(UID_POJO)))) + .contains("state_key", "profile"); + assertThat(columnNames(schemaOf(catalog, dbName, keyedTable(UID_AVRO_SPECIFIC)))) + .contains("state_key", "avro_specific"); + assertThat(columnNames(schemaOf(catalog, dbName, keyedTable(UID_AVRO_GENERIC)))) + .contains("state_key", "avro_generic"); + + // Primitive state: 5 distinct int keys + List<Row> primRows = collectAll(tableEnv, UID_PRIMITIVE); + assertThat(primRows).hasSize(5); + assertThat(stateKeys(primRows)).containsExactlyInAnyOrder(1, 2, 3, 4, 5); + + // The remaining operators all hold a nested ROW value under string keys. + assertNestedRowState(UID_POJO, "profile", "name", "score"); + assertNestedRowState(UID_AVRO_SPECIFIC, "avro_specific", "name", "value"); + assertNestedRowState(UID_AVRO_GENERIC, "avro_generic", "name", "value"); + } + + private void assertNestedRowState(String operatorUid, String column, String... nestedFields) + throws Exception { + List<Row> rows = collectAll(tableEnv, operatorUid); + assertThat(rows).hasSize(5); + assertThat(stateKeys(rows)).containsExactlyInAnyOrder("1", "2", "3", "4", "5"); + for (Row row : rows) { + Row nested = (Row) row.getField(column); + assertThat(nested).isNotNull(); + for (String nestedField : nestedFields) { + assertThat(nested.getField(nestedField)).isNotNull(); + } + } + } + + @Test + void testProjectionColumnReorder() throws Exception { + // Reorder: value column first, key column second + List<Row> primRows = + collectWithSql( + tableEnv, + "SELECT `count`, state_key FROM `" + + keyedTable(UID_PRIMITIVE) + + "` ORDER BY state_key"); + + assertThat(primRows).hasSize(5); + for (Row row : primRows) { + assertThat(row.getArity()).isEqualTo(2); + assertThat(row.getField(0)).isEqualTo(1); + assertThat(row.getField("count")).isEqualTo(1); + assertThat(row.getField("state_key")).isIn(1, 2, 3, 4, 5); + } + + // POJO operator: reorder profile (ROW) before state_key + List<Row> pojoRows = + collectWithSql( + tableEnv, + "SELECT profile, state_key FROM `" + + keyedTable(UID_POJO) + + "` ORDER BY state_key"); + + assertThat(pojoRows).hasSize(5); + for (Row row : pojoRows) { + assertThat(row.getArity()).isEqualTo(2); + Row profile = (Row) row.getField(0); + assertThat(profile).isNotNull(); + assertThat(profile.getField("name")).isNotNull(); + assertThat(profile.getField("score")).isNotNull(); + assertThat(row.getField("state_key")).isNotNull(); + } + } + + @Test + void testProjectionSubsets() throws Exception { + String primTable = "`" + keyedTable(UID_PRIMITIVE) + "`"; + + List<Row> keyOnlyRows = + collectWithSql( + tableEnv, "SELECT state_key FROM " + primTable + " ORDER BY state_key"); + assertThat(keyOnlyRows).hasSize(5); + for (Row row : keyOnlyRows) { + assertThat(row.getArity()).isEqualTo(1); + } + assertThat( + keyOnlyRows.stream() + .map(r -> r.getField("state_key")) + .collect(Collectors.toList())) + .containsExactly(1, 2, 3, 4, 5); + + List<Row> valueOnlyRows = collectWithSql(tableEnv, "SELECT `count` FROM " + primTable); + assertThat(valueOnlyRows).hasSize(5); + for (Row row : valueOnlyRows) { + assertThat(row.getArity()).isEqualTo(1); + assertThat(row.getField("count")).isEqualTo(1); + } + } + } + + /** POJO-key and Avro-specific-key savepoints — off-classpath key types. */ + @Nested + class OffClasspathKeyTypes { + + private static final String POJO_AVRO_KEY_DIR = + "src/test/resources/keyed-state-pojo-avro-key"; + private static final String UID_POJO_KEY = "pojo-key-state-op"; + private static final String UID_AVRO_SPECIFIC_KEY = "avro-specific-key-state-op"; + + private StateCatalog catalog; + private TableEnvironment tableEnv; + private String dbName; + private Path savepointPath; + + @BeforeEach + void openCatalog() throws Exception { + savepointPath = findSavepointDir(POJO_AVRO_KEY_DIR); + catalog = openCatalogOn(POJO_AVRO_KEY_DIR); + dbName = catalog.listDatabases().get(0); + tableEnv = newCatalogTableEnv(catalog, dbName); + } + + @AfterEach + void closeCatalog() { + catalog.close(); + } + + @Test + void testPojoAndAvroKeySchemaTypes() throws Exception { + CheckpointMetadata metadata = + SavepointLoader.loadSavepointMetadata(savepointPath.toString()); + + // POJO key (PersonKey{int id, String name}) → ROW(id INT, name VARCHAR) + KeyedStateSchemaInfo pojoSchema = + StateTableUtils.getKeyedStateSchema( + metadata, OperatorIdentifier.forUid(UID_POJO_KEY)); + assertThat(pojoSchema.keyType.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType pojoKeyType = (RowType) pojoSchema.keyType; + assertThat(pojoKeyType.getFieldNames()).containsExactlyInAnyOrder("id", "name"); + assertThat(fieldTypeRoot(pojoKeyType, "id")).isEqualTo(LogicalTypeRoot.INTEGER); + assertThat(fieldTypeRoot(pojoKeyType, "name")).isEqualTo(LogicalTypeRoot.VARCHAR); + + // Avro specific key (StateTestRecord{String name, long value}) → ROW(name VARCHAR, + // value BIGINT) + KeyedStateSchemaInfo avroSchema = + StateTableUtils.getKeyedStateSchema( + metadata, OperatorIdentifier.forUid(UID_AVRO_SPECIFIC_KEY)); + assertThat(avroSchema.keyType.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType avroKeyType = (RowType) avroSchema.keyType; + assertThat(avroKeyType.getFieldNames()).containsExactlyInAnyOrder("name", "value"); + } + + @Test + void testPojoAndAvroKeyedStateTables() throws Exception { + assertThat(catalog.listTables(dbName)) + .contains(keyedTable(UID_POJO_KEY), keyedTable(UID_AVRO_SPECIFIC_KEY)); + + for (String uid : new String[] {UID_POJO_KEY, UID_AVRO_SPECIFIC_KEY}) { + assertThat(columnNames(schemaOf(catalog, dbName, keyedTable(uid)))) + .contains("state_key", "count"); + } + + // POJO key: 5 rows; PersonKey{id, name} off-classpath → deserialized as ROW + List<Row> 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<Row> avroKeyRows = collectAll(tableEnv, UID_AVRO_SPECIFIC_KEY); + assertThat(avroKeyRows).hasSize(5); + for (Row row : avroKeyRows) { + Row key = (Row) row.getField("state_key"); + assertThat(key).isNotNull(); + assertThat(key.getField("name")).asString().startsWith("key-"); + assertThat(key.getField("value")).isIn(1L, 2L, 3L, 4L, 5L); + assertThat(row.getField("count")).isEqualTo(1); + } + } + } + + /** TupleX key and TupleX value with mixed basic + POJO types. */ + @Nested + class TupleKeyAndValue { + + private static final String TUPLE_KEY_DIR = "src/test/resources/keyed-state-tuple-key"; + private static final String UID_TUPLE_KEY = "tuple-key-state-op"; + private static final String UID_TUPLE_POJO_VALUE = "tuple-pojo-value-state-op"; + + private StateCatalog catalog; + private TableEnvironment tableEnv; + private String dbName; + private Path savepointPath; + + @BeforeEach + void openCatalog() throws Exception { + savepointPath = findSavepointDir(TUPLE_KEY_DIR); + catalog = openCatalogOn(TUPLE_KEY_DIR); + dbName = catalog.listDatabases().get(0); + tableEnv = newCatalogTableEnv(catalog, dbName); + } + + @AfterEach + void closeCatalog() { + catalog.close(); + } + + @Test + void testTupleKeySchemaTypes() throws Exception { + CheckpointMetadata metadata = + SavepointLoader.loadSavepointMetadata(savepointPath.toString()); + + // Tuple2<Integer, String> key → ROW(f0 INT NOT NULL, f1 VARCHAR) + KeyedStateSchemaInfo tupleKeySchema = + StateTableUtils.getKeyedStateSchema( + metadata, OperatorIdentifier.forUid(UID_TUPLE_KEY)); + assertThat(tupleKeySchema.keyType.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType tupleKeyType = (RowType) tupleKeySchema.keyType; + assertThat(tupleKeyType.getFieldNames()).containsExactly("f0", "f1"); + assertThat(fieldTypeRoot(tupleKeyType, "f0")).isEqualTo(LogicalTypeRoot.INTEGER); + assertThat(fieldTypeRoot(tupleKeyType, "f1")).isEqualTo(LogicalTypeRoot.VARCHAR); + + // Integer key, Tuple2<Long, TuplePojoField> value → value column is ROW(f0 BIGINT, f1 + // ROW(name VARCHAR, score BIGINT)) + KeyedStateSchemaInfo tuplePojoValueSchema = + StateTableUtils.getKeyedStateSchema( + metadata, OperatorIdentifier.forUid(UID_TUPLE_POJO_VALUE)); + assertThat(tuplePojoValueSchema.keyType.getTypeRoot()) + .isEqualTo(LogicalTypeRoot.INTEGER); + assertThat(tuplePojoValueSchema.stateSchemas).containsKey("tuple_pojo"); + LogicalType tupleValueType = + tuplePojoValueSchema.stateSchemas.get("tuple_pojo").logicalType; + assertThat(tupleValueType.getTypeRoot()).isEqualTo(LogicalTypeRoot.ROW); + RowType tupleValueRowType = (RowType) tupleValueType; + assertThat(tupleValueRowType.getFieldNames()).containsExactly("f0", "f1"); + assertThat(fieldTypeRoot(tupleValueRowType, "f0")).isEqualTo(LogicalTypeRoot.BIGINT); + assertThat(fieldTypeRoot(tupleValueRowType, "f1")).isEqualTo(LogicalTypeRoot.ROW); + RowType innerPojoType = + (RowType) tupleValueRowType.getTypeAt(tupleValueRowType.getFieldIndex("f1")); + assertThat(innerPojoType.getFieldNames()).containsExactlyInAnyOrder("name", "score"); + } + + @Test + void testTupleKeyedStateTables() throws Exception { + assertThat(catalog.listTables(dbName)) + .contains(keyedTable(UID_TUPLE_KEY), keyedTable(UID_TUPLE_POJO_VALUE)); + + // Tuple2<Integer, String> key: 5 rows; key fields f0=int, f1=string + List<Row> 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<Long, TuplePojoField> value: 5 rows; value row has f0=long, f1=row(name,score) + List<Row> tuplePojoValueRows = collectAll(tableEnv, UID_TUPLE_POJO_VALUE); + assertThat(tuplePojoValueRows).hasSize(5); + for (Row row : tuplePojoValueRows) { + Integer key = (Integer) row.getField("state_key"); + assertThat(key).isIn(1, 2, 3, 4, 5); + Row tupleValue = (Row) row.getField("tuple_pojo"); + assertThat(tupleValue).isNotNull(); + assertThat(tupleValue.getField("f0")).isEqualTo((long) key * 10); + Row pojoField = (Row) tupleValue.getField("f1"); + assertThat(pojoField).isNotNull(); + assertThat(pojoField.getField("name")).isEqualTo("name-" + key); + assertThat(pojoField.getField("score")).isEqualTo((long) key * 100); + } + } + } + + // ------------------------------------------------------------------------- + // Shared helpers + // ------------------------------------------------------------------------- + + private static StateCatalog openCatalogOn(String resourceDir) throws Exception { + String catalogRoot = Paths.get(resourceDir).toAbsolutePath().toString(); + StateCatalog catalog = + new StateCatalog("state", Collections.singletonMap("test", catalogRoot)); + catalog.open(); + return catalog; + } + + private static TableEnvironment newCatalogTableEnv(StateCatalog catalog, String dbName) { + TableEnvironment tableEnv = TableEnvironment.create(EnvironmentSettings.inBatchMode()); + tableEnv.loadModule("state", StateModule.INSTANCE); + tableEnv.registerCatalog("state", catalog); + tableEnv.useCatalog("state"); + tableEnv.useDatabase(dbName); + return tableEnv; + } + + /** Finds the single {@code savepoint-*} directory nested directly under {@code parentDir}. */ + private static Path findSavepointDir(String parentDir) throws IOException { + try (var stream = Files.list(Paths.get(parentDir))) { + return stream.filter( + p -> + Files.isDirectory(p) + && p.getFileName().toString().startsWith("savepoint-")) + .findFirst() + .orElseThrow(() -> new IOException("No savepoint found in " + parentDir)); + } + } + + private static String keyedTable(String operatorUid) { + return StateCatalog.OPERATOR_UID_PREFIX + operatorUid + StateCatalog.OPERATOR_TABLE_SUFFIX; + } + + private static String flatKeyedTable(String operatorUid, String stateName) { + return StateCatalog.OPERATOR_UID_PREFIX + + operatorUid + + "_" + + stateName + + StateCatalog.FLAT_STATE_TABLE_SUFFIX; + } + + private static Schema schemaOf(StateCatalog catalog, String dbName, String tableName) + throws Exception { + return ((CatalogTable) catalog.getTable(new ObjectPath(dbName, tableName))) + .getUnresolvedSchema(); + } + + private static List<String> columnNames(Schema schema) { + return schema.getColumns().stream() + .map(Schema.UnresolvedColumn::getName) + .collect(Collectors.toList()); + } + + private static LogicalTypeRoot stateTypeRoot( + KeyedStateSchemaInfo schemaInfo, String stateName) { + KeyedStateSchemaInfo.StateEntryInfo entry = schemaInfo.stateSchemas.get(stateName); + assertThat(entry).as("state '%s'", stateName).isNotNull(); + return entry.logicalType.getTypeRoot(); + } + + private static LogicalTypeRoot fieldTypeRoot(RowType rowType, String fieldName) { + return rowType.getTypeAt(rowType.getFieldIndex(fieldName)).getTypeRoot(); + } + + private static Set<Object> stateKeys(List<Row> rows) { + return rows.stream().map(r -> r.getField("state_key")).collect(Collectors.toSet()); + } + + private static List<Long> longRange(int endExclusive) { + return LongStream.range(0, endExclusive).boxed().collect(Collectors.toList()); + } + + /** + * Registers the discovered keyed-state table of {@code operatorUid} as a temporary table and + * reads it in batch mode via {@code StreamTableEnvironment}, bypassing {@link StateCatalog}. + */ + private static List<Row> readViaTemporaryTable( + String statePath, String operatorUid, String tableName) throws Exception { + Configuration config = new Configuration(); + config.set(RUNTIME_MODE, RuntimeExecutionMode.BATCH); + StreamTableEnvironment tEnv = + StreamTableEnvironment.create( + StreamExecutionEnvironment.getExecutionEnvironment(config)); + + CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(statePath); + OperatorIdentifier opId = OperatorIdentifier.forUid(operatorUid); + CatalogTable catalogTable = + StateTableUtils.getStateCatalogTable( + metadata, + StateTableUtils.getKeyedStateSchema(metadata, opId), + statePath, + opId); + + CatalogManager catalogManager = ((TableEnvironmentImpl) tEnv).getCatalogManager(); + catalogManager.createTemporaryTable( + catalogTable, + catalogManager.qualifyIdentifier(UnresolvedIdentifier.of(tableName)), + false); + + return tEnv.toDataStream(tEnv.sqlQuery("SELECT * FROM " + tableName)) + .executeAndCollect(100); + } + + private static List<Row> collectAll(TableEnvironment tEnv, String operatorUid) + throws Exception { + return collectWithSql(tEnv, "SELECT * FROM `" + keyedTable(operatorUid) + "`"); + } + + private static List<Row> collectWithSql(TableEnvironment tEnv, String sql) throws Exception { + List<Row> rows = new ArrayList<>(); + TableResult result = tEnv.executeSql(sql); + try (CloseableIterator<Row> it = result.collect()) { + it.forEachRemaining(rows::add); + } + return rows; + } +} diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/TuplePojoField.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/TuplePojoField.java new file mode 100644 index 00000000000..c963a306e3f --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/catalog/TuplePojoField.java @@ -0,0 +1,59 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.state.catalog; + +import java.util.Objects; + +/** + * Simple POJO used as a nested field inside a Tuple value state in {@link + * StateCatalogGeneratedSavepointITCase.TupleKeyAndValue}. + * + * <p>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/resources/generator/KeyedStateCatalogSavepointGenerator.java b/flink-libraries/flink-state-processing-api/src/test/resources/generator/KeyedStateCatalogSavepointGenerator.java new file mode 100644 index 00000000000..2dda9bb6f9a --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/resources/generator/KeyedStateCatalogSavepointGenerator.java @@ -0,0 +1,356 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.state.catalog; + +import org.apache.flink.api.common.JobStatus; +import org.apache.flink.api.common.functions.OpenContext; +import org.apache.flink.api.common.state.ValueState; +import org.apache.flink.api.common.state.ValueStateDescriptor; +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.StateBackendOptions; +import org.apache.flink.core.execution.JobClient; +import org.apache.flink.core.execution.SavepointFormatType; +import org.apache.flink.formats.avro.typeutils.AvroTypeInfo; +import org.apache.flink.formats.avro.typeutils.GenericRecordAvroTypeInfo; +import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration; +import org.apache.flink.state.catalog.avro.StateTestRecord; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.streaming.api.functions.KeyedProcessFunction; +import org.apache.flink.streaming.api.functions.sink.v2.DiscardingSink; +import org.apache.flink.streaming.api.functions.source.legacy.RichSourceFunction; +import org.apache.flink.test.util.MiniClusterWithClientResource; +import org.apache.flink.util.Collector; + +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.file.FileVisitResult; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.nio.file.SimpleFileVisitor; +import java.nio.file.attribute.BasicFileAttributes; +import java.util.Objects; +import java.util.concurrent.TimeUnit; + +/** + * Generates the pre-built savepoint used by {@code + * StateCatalogGeneratedSavepointITCase.MultiOperatorTypeCatalog}. + * + * <p>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. + * + * <p>To regenerate the savepoint (e.g. after a Flink serializer format change): + * + * <ol> + * <li>Copy this file to {@code src/test/java/org/apache/flink/state/catalog/}. + * <li>Copy {@code src/test/resources/generator/StateTestRecord.avsc} to {@code + * src/test/resources/avro/StateTestRecord.avsc} so the Avro class is generated (the {@code + * avro-maven-plugin} and {@code flink-avro} dependency are already configured in pom.xml, no + * changes needed there). + * <li>Remove the {@code @Disabled} annotation and run: {@code ./mvnw test -pl + * flink-libraries/flink-state-processing-api + * -Dtest=KeyedStateCatalogSavepointGenerator#generateSavepoint} + * <li>Verify the new savepoint under {@code src/test/resources/keyed-state-catalog/}. + * <li>Remove the copied {@code .java} file from the source tree and the copied {@code .avsc} + * file from {@code src/test/resources/avro/}. + * </ol> + */ +@Disabled("Run manually to regenerate the pre-built savepoint in test resources") +class KeyedStateCatalogSavepointGenerator { + + private static final String RESOURCES_DIR = "src/test/resources/keyed-state-catalog"; + + // Must match StateCatalogGeneratedSavepointITCase.MultiOperatorTypeCatalog's UID_* constants. + private static final String UID_PRIMITIVE = "primitive-state-op"; + private static final String UID_POJO = "pojo-state-op"; + private static final String UID_AVRO_SPECIFIC = "avro-specific-state-op"; + private static final String UID_AVRO_GENERIC = "avro-generic-state-op"; + + @Test + void generateSavepoint() throws Exception { + Path outputDir = Paths.get(RESOURCES_DIR).toAbsolutePath(); + Files.createDirectories(outputDir); + deleteExistingSavepoints(outputDir); + + var cluster = + new MiniClusterWithClientResource( + new MiniClusterResourceConfiguration.Builder() + .setNumberSlotsPerTaskManager(4) + .build()); + cluster.before(); + try { + Configuration cfg = new Configuration(); + cfg.set(StateBackendOptions.STATE_BACKEND, "hashmap"); + var env = StreamExecutionEnvironment.getExecutionEnvironment(cfg); + env.setParallelism(2); + + env.addSource(new BoundedWaitingSource(new int[] {1, 2, 3, 4, 5})) + .returns(Types.INT) + .keyBy(v -> v) + .process(new IntCountOperator()) + .uid(UID_PRIMITIVE) + .name(UID_PRIMITIVE) + .keyBy(v -> String.valueOf(v)) + .process(new ProfileOperator()) + .uid(UID_POJO) + .name(UID_POJO) + .keyBy(v -> String.valueOf(v)) + .process(new AvroSpecificOperator()) + .uid(UID_AVRO_SPECIFIC) + .name(UID_AVRO_SPECIFIC) + .keyBy(v -> String.valueOf(v)) + .process(new AvroGenericOperator()) + .uid(UID_AVRO_GENERIC) + .name(UID_AVRO_GENERIC) + .sinkTo(new DiscardingSink<>()); + + String savepointPath = takeSavepoint(env, outputDir.toString()); + System.out.println("Savepoint written to: " + savepointPath); + } finally { + cluster.after(); + } + } + + private static String takeSavepoint(StreamExecutionEnvironment env, String savepointDir) + throws Exception { + JobClient jobClient = env.executeAsync(); + try { + while (jobClient.getJobStatus().get() != JobStatus.RUNNING) { + Thread.sleep(100); + } + Exception lastEx = null; + for (int attempt = 0; attempt < 30; attempt++) { + try { + return jobClient + .triggerSavepoint(savepointDir, SavepointFormatType.CANONICAL) + .get(2, TimeUnit.MINUTES); + } catch (Exception e) { + lastEx = e; + Thread.sleep(200); + } + } + throw new RuntimeException("Could not trigger savepoint after 30 attempts", lastEx); + } finally { + try { + jobClient.cancel().get(10, TimeUnit.SECONDS); + } catch (Exception ignored) { + } + } + } + + /** Deletes any existing {@code savepoint-*} directories inside the given directory. */ + private static void deleteExistingSavepoints(Path dir) throws IOException { + if (!Files.isDirectory(dir)) { + return; + } + try (var stream = Files.list(dir)) { + stream.filter(p -> p.getFileName().toString().startsWith("savepoint-")) + .filter(Files::isDirectory) + .forEach( + p -> { + try { + deleteDirectory(p); + } catch (IOException e) { + throw new RuntimeException("Failed to delete " + p, e); + } + }); + } + } + + private static void deleteDirectory(Path dir) throws IOException { + Files.walkFileTree( + dir, + new SimpleFileVisitor<>() { + @Override + public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) + throws IOException { + Files.delete(file); + return FileVisitResult.CONTINUE; + } + + @Override + public FileVisitResult postVisitDirectory(Path d, IOException exc) + throws IOException { + Files.delete(d); + return FileVisitResult.CONTINUE; + } + }); + } + + // ------------------------------------------------------------------------- + // Source + // ------------------------------------------------------------------------- + + private static class BoundedWaitingSource extends RichSourceFunction<Integer> { + + private final int[] elements; + private volatile boolean running = true; + + BoundedWaitingSource(int[] elements) { + this.elements = elements; + } + + @Override + public void run(SourceContext<Integer> 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<Integer, Integer, Integer> { + + private transient ValueState<Integer> count; + + @Override + public void open(OpenContext ctx) throws Exception { + count = + getRuntimeContext() + .getState(new ValueStateDescriptor<>("count", Integer.class)); + } + + @Override + public void processElement(Integer value, Context ctx, Collector<Integer> out) + throws Exception { + Integer c = count.value(); + count.update(c == null ? 1 : c + 1); + out.collect(value); + } + } + + private static class ProfileOperator extends KeyedProcessFunction<String, Integer, Integer> { + + private transient ValueState<PersonProfile> profile; + + @Override + public void open(OpenContext ctx) throws Exception { + profile = + getRuntimeContext() + .getState(new ValueStateDescriptor<>("profile", PersonProfile.class)); + } + + @Override + public void processElement(Integer value, Context ctx, Collector<Integer> out) + throws Exception { + profile.update(new PersonProfile("name-" + value, value * 10L)); + out.collect(value); + } + } + + private static class AvroSpecificOperator + extends KeyedProcessFunction<String, Integer, Integer> { + + private transient ValueState<StateTestRecord> avroSpecific; + + @Override + public void open(OpenContext ctx) throws Exception { + avroSpecific = + getRuntimeContext() + .getState( + new ValueStateDescriptor<>( + "avro_specific", + new AvroTypeInfo<>(StateTestRecord.class))); + } + + @Override + public void processElement(Integer value, Context ctx, Collector<Integer> 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<String, Integer, Integer> { + + private transient ValueState<GenericRecord> avroGeneric; + + @Override + public void open(OpenContext ctx) throws Exception { + avroGeneric = + getRuntimeContext() + .getState( + new ValueStateDescriptor<>( + "avro_generic", + new GenericRecordAvroTypeInfo( + StateTestRecord.getClassSchema()))); + } + + @Override + public void processElement(Integer value, Context ctx, Collector<Integer> 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 00000000000..4bff1559a83 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/resources/generator/KeyedStatePojoAvroKeySavepointGenerator.java @@ -0,0 +1,321 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.state.catalog; + +import org.apache.flink.api.common.JobStatus; +import org.apache.flink.api.common.functions.OpenContext; +import org.apache.flink.api.common.state.ValueState; +import org.apache.flink.api.common.state.ValueStateDescriptor; +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.StateBackendOptions; +import org.apache.flink.core.execution.JobClient; +import org.apache.flink.core.execution.SavepointFormatType; +import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration; +import org.apache.flink.state.catalog.avro.StateTestRecord; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.streaming.api.functions.KeyedProcessFunction; +import org.apache.flink.streaming.api.functions.sink.v2.DiscardingSink; +import org.apache.flink.streaming.api.functions.source.legacy.RichSourceFunction; +import org.apache.flink.test.util.MiniClusterWithClientResource; +import org.apache.flink.util.Collector; + +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.file.FileVisitResult; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.nio.file.SimpleFileVisitor; +import java.nio.file.attribute.BasicFileAttributes; +import java.util.Objects; +import java.util.concurrent.TimeUnit; + +/** + * Generates the pre-built savepoint used by {@code + * StateCatalogGeneratedSavepointITCase.OffClasspathKeyTypes} for POJO-key and Avro-specific-key + * scenarios. + * + * <p>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/}. + * + * <p>To regenerate the savepoint (e.g. after a Flink serializer format change): + * + * <ol> + * <li>Copy this file to {@code src/test/java/org/apache/flink/state/catalog/} if not already + * there. + * <li>Copy {@code src/test/resources/generator/StateTestRecord.avsc} to {@code + * src/test/resources/avro/StateTestRecord.avsc} so the Avro class is generated. + * <li>Remove the {@code @Disabled} annotation and run: {@code ./mvnw test -pl + * flink-libraries/flink-state-processing-api + * -Dtest=KeyedStatePojoAvroKeySavepointGenerator#generateSavepoint} + * <li>Verify the savepoint under {@code src/test/resources/keyed-state-pojo-avro-key/}. + * <li>Remove the copied {@code StateTestRecord.avsc} from {@code src/test/resources/avro/} and + * move this {@code .java} file back to {@code resources/generator/}. + * </ol> + */ +@Disabled("Run manually to regenerate the pre-built savepoint in test resources") +class KeyedStatePojoAvroKeySavepointGenerator { + + /** Operator UIDs referenced by {@code StateCatalogGeneratedSavepointITCase.OffClasspathKeyTypes}. */ + static final String UID_POJO_KEY = "pojo-key-state-op"; + + static final String UID_AVRO_SPECIFIC_KEY = "avro-specific-key-state-op"; + + private static final String RESOURCES_DIR = "src/test/resources/keyed-state-pojo-avro-key"; + + @Test + void generateSavepoint() throws Exception { + Path outputDir = Paths.get(RESOURCES_DIR).toAbsolutePath(); + Files.createDirectories(outputDir); + deleteExistingSavepoints(outputDir); + + var cluster = + new MiniClusterWithClientResource( + new MiniClusterResourceConfiguration.Builder() + .setNumberSlotsPerTaskManager(4) + .build()); + cluster.before(); + try { + Configuration cfg = new Configuration(); + cfg.set(StateBackendOptions.STATE_BACKEND, "hashmap"); + var env = StreamExecutionEnvironment.getExecutionEnvironment(cfg); + env.setParallelism(2); + + // Pipeline 1: POJO key (PersonKey{id,name}). + // PersonKey is defined only in this file → NOT on the classpath during normal test + // runs, which exercises the off-classpath POJO-key reading path. + env.addSource(new BoundedWaitingSource(new int[] {1, 2, 3, 4, 5})) + .returns(Types.INT) + .keyBy(v -> new PersonKey(v, "name-" + v)) + .process(new PojoKeyCountOperator()) + .uid(UID_POJO_KEY) + .name(UID_POJO_KEY) + .sinkTo(new DiscardingSink<>()); + + // Pipeline 2: Avro-specific key (StateTestRecord). + // StateTestRecord is generated from StateTestRecord.avsc only when the .avsc file is + // present in src/test/resources/avro/. After generation that file is removed, so the + // class stays off the classpath during tests, exercising the Avro fallback path. + env.addSource(new BoundedWaitingSource(new int[] {1, 2, 3, 4, 5})) + .returns(Types.INT) + .keyBy( + (KeySelector<Integer, StateTestRecord>) + 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<Integer> { + + private final int[] elements; + private volatile boolean running = true; + + BoundedWaitingSource(int[] elements) { + this.elements = elements; + } + + @Override + public void run(SourceContext<Integer> 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<PersonKey, Integer, Integer> { + + private transient ValueState<Integer> count; + + @Override + public void open(OpenContext ctx) throws Exception { + count = + getRuntimeContext() + .getState(new ValueStateDescriptor<>("count", Integer.class)); + } + + @Override + public void processElement(Integer value, Context ctx, Collector<Integer> 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<StateTestRecord, Integer, Integer> { + + private transient ValueState<Integer> count; + + @Override + public void open(OpenContext ctx) throws Exception { + count = + getRuntimeContext() + .getState(new ValueStateDescriptor<>("count", Integer.class)); + } + + @Override + public void processElement(Integer value, Context ctx, Collector<Integer> 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 00000000000..4147924c36a --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/test/resources/generator/KeyedStateTupleKeySavepointGenerator.java @@ -0,0 +1,277 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.state.catalog; + +import org.apache.flink.api.common.JobStatus; +import org.apache.flink.api.common.functions.OpenContext; +import org.apache.flink.api.common.state.ValueState; +import org.apache.flink.api.common.state.ValueStateDescriptor; +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.StateBackendOptions; +import org.apache.flink.core.execution.JobClient; +import org.apache.flink.core.execution.SavepointFormatType; +import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.streaming.api.functions.KeyedProcessFunction; +import org.apache.flink.streaming.api.functions.sink.v2.DiscardingSink; +import org.apache.flink.streaming.api.functions.source.legacy.RichSourceFunction; +import org.apache.flink.test.util.MiniClusterWithClientResource; +import org.apache.flink.util.Collector; + +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.file.FileVisitResult; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.nio.file.SimpleFileVisitor; +import java.nio.file.attribute.BasicFileAttributes; +import java.util.concurrent.TimeUnit; + +/** + * Generates the pre-built savepoint used by {@code StateCatalogGeneratedSavepointITCase.TupleKeyAndValue} + * for Tuple-key and Tuple-value scenarios (including a Tuple value with a mixed basic + POJO field). + * + * <p>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/}. + * + * <p>To regenerate the savepoint (e.g. after a Flink serializer format change): + * + * <ol> + * <li>Copy this file to {@code src/test/java/org/apache/flink/state/catalog/} if not already + * there. + * <li>Remove the {@code @Disabled} annotation and run: {@code ./mvnw test -pl + * flink-libraries/flink-state-processing-api + * -Dtest=KeyedStateTupleKeySavepointGenerator#generateSavepoint} + * <li>Verify the savepoint under {@code src/test/resources/keyed-state-tuple-key/}. + * <li>Move this {@code .java} file back to {@code resources/generator/}. + * </ol> + */ +@Disabled("Run manually to regenerate the pre-built savepoint in test resources") +class KeyedStateTupleKeySavepointGenerator { + + /** Operator UIDs referenced by {@code StateCatalogGeneratedSavepointITCase.TupleKeyAndValue}. */ + static final String UID_TUPLE_KEY = "tuple-key-state-op"; + + static final String UID_TUPLE_POJO_VALUE = "tuple-pojo-value-state-op"; + + private static final String RESOURCES_DIR = "src/test/resources/keyed-state-tuple-key"; + + @Test + void generateSavepoint() throws Exception { + Path outputDir = Paths.get(RESOURCES_DIR).toAbsolutePath(); + Files.createDirectories(outputDir); + deleteExistingSavepoints(outputDir); + + var cluster = + new MiniClusterWithClientResource( + new MiniClusterResourceConfiguration.Builder() + .setNumberSlotsPerTaskManager(4) + .build()); + cluster.before(); + try { + Configuration cfg = new Configuration(); + cfg.set(StateBackendOptions.STATE_BACKEND, "hashmap"); + var env = StreamExecutionEnvironment.getExecutionEnvironment(cfg); + env.setParallelism(2); + + // Pipeline 1: Tuple2<Integer, String> 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<Integer, Tuple2<Integer, String>>) + 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<Long, TuplePojoField> 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<Integer> { + + private final int[] elements; + private volatile boolean running = true; + + BoundedWaitingSource(int[] elements) { + this.elements = elements; + } + + @Override + public void run(SourceContext<Integer> 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<Integer, String>} key. */ + private static class TupleKeyCountOperator + extends KeyedProcessFunction<Tuple2<Integer, String>, Integer, Integer> { + + private transient ValueState<Integer> count; + + @Override + public void open(OpenContext ctx) throws Exception { + count = + getRuntimeContext() + .getState(new ValueStateDescriptor<>("count", Integer.class)); + } + + @Override + public void processElement(Integer value, Context ctx, Collector<Integer> out) + throws Exception { + Integer c = count.value(); + count.update(c == null ? 1 : c + 1); + out.collect(value); + } + } + + /** Stores a {@code Tuple2<Long, TuplePojoField>} value per Integer key. */ + private static class TuplePojoValueOperator + extends KeyedProcessFunction<Integer, Integer, Integer> { + + private transient ValueState<Tuple2<Long, TuplePojoField>> tupleState; + + @Override + public void open(OpenContext ctx) throws Exception { + tupleState = + getRuntimeContext() + .getState( + new ValueStateDescriptor<>( + "tuple_pojo", + Types.TUPLE( + Types.LONG, Types.POJO(TuplePojoField.class)))); + } + + @Override + public void processElement(Integer value, Context ctx, Collector<Integer> 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 00000000000..681692f5ca6 --- /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 00000000000..77125045df3 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 00000000000..cea2d23d239 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 00000000000..186a78268ed 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 00000000000..1e780952d66 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 00000000000..a8295168fda Binary files /dev/null and b/flink-libraries/flink-state-processing-api/src/test/resources/table-state-missing-class/_metadata differ
