[FLINK-40180] Non-keyed/operator state support in State Catalog

flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/StateTableUtils.java

@@ -32,6 +32,8 @@
 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.NonKeyedStateSchemaInfo;
+import org.apache.flink.state.api.schema.OperatorStateSchemaInfo;
 import org.apache.flink.state.api.schema.SerializerSnapshotToLogicalTypeConverter;
 import org.apache.flink.state.api.schema.StateSchemaExtractor;
 import org.apache.flink.state.api.schema.StateSchemaInfo;
@@ -46,6 +48,7 @@
 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.RowType;
 import org.apache.flink.table.types.logical.VarBinaryType;
 import org.apache.flink.table.types.utils.LogicalTypeDataTypeConverter;
 
@@ -67,8 +70,8 @@
 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.
+ * High-level utility for inspecting and reading keyed and non-keyed state from a checkpoint /
+ * savepoint without requiring user POJO classes on the classpath.
  */
 @Internal
 public final class StateTableUtils {
@@ -79,14 +82,14 @@ 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.
+     * metadata that have at least one non-internal keyed or non-keyed state.
      *
      * @param metadata the checkpoint metadata to inspect
      * @return list of operator identifiers; never null, may be empty
      */
     public static List<OperatorIdentifier> getOperatorIdentifiers(CheckpointMetadata metadata) {
         return metadata.getOperatorStates().stream()
-                .filter(StateTableUtils::hasNonInternalKeyedState)
+                .filter(op -> hasNonInternalKeyedState(op) || hasNonInternalOperatorState(op))
                 .map(
                         op ->
                                 op.getOperatorUid()
@@ -113,12 +116,25 @@ private static boolean hasNonInternalKeyedState(OperatorState op) {
         }
     }
 
+    private static boolean hasNonInternalOperatorState(OperatorState op) {
+        try {
+            List<OperatorStateSchemaInfo> schemas = StateSchemaExtractor.extractOperatorSchema(op);
+            return schemas.stream().anyMatch(info -> !isInternalState(info.stateName));
+        } catch (Exception e) {
+            LOG.error(
+                    "Could not extract non-keyed 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
      */
@@ -466,10 +482,187 @@ private static String addFlattenedValueColumns(
         return subKeyColumnName;
     }
 
+    /**
+     * Returns the {@link NonKeyedStateSchemaInfo} for the non-keyed (operator) states of the given
+     * operator — {@code ListState}, {@code UnionState} and {@code BroadcastState} — the ones
+     * exposed by the {@code _list}/{@code _union}/{@code _broadcast} tables.
+     *
+     * <p>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. States whose schema
+     * cannot be determined are excluded with a logged error.
+     *
+     * @param metadata the checkpoint metadata to inspect
+     * @param operatorId identifies the operator
+     * @return schema information covering all registered non-keyed state entries
+     * @throws IOException if the state header cannot be read
+     */
+    public static NonKeyedStateSchemaInfo getNonKeyedStateSchema(
+            CheckpointMetadata metadata, OperatorIdentifier operatorId) throws IOException {
+        OperatorState opState = findOperatorState(metadata, operatorId);
+        List<OperatorStateSchemaInfo> schemas = StateSchemaExtractor.extractOperatorSchema(opState);
+
+        LinkedHashMap<String, NonKeyedStateSchemaInfo.StateEntryInfo> stateSchemas =
+                new LinkedHashMap<>();
+        for (OperatorStateSchemaInfo info : schemas) {
+            if (isInternalState(info.stateName)) {
+                continue;
+            }
+            try {
+                LogicalType valueLogicalType =
+                        SerializerSnapshotToLogicalTypeConverter.convert(info.valueSnapshot);
+                LogicalType mapKeyLogicalType =
+                        info.keySnapshot == null
+                                ? null
+                                : SerializerSnapshotToLogicalTypeConverter.convert(
+                                        info.keySnapshot);
+                stateSchemas.put(
+                        info.stateName,
+                        new NonKeyedStateSchemaInfo.StateEntryInfo(
+                                info.kind, valueLogicalType, mapKeyLogicalType));
+            } catch (Exception e) {
+                logSchemaExtractionFailure("non-keyed ", info.stateName, info.valueSnapshot, e);
+            }
+        }
+
+        return new NonKeyedStateSchemaInfo(stateSchemas);
+    }
+
+    /**
+     * Builds a {@link CatalogTable} exposing a single {@code ListState} or {@code UnionState} of an
+     * operator, with one row per list element.
+     *
+     * <p>The state's value is flattened directly into the table's columns instead of being wrapped
+     * in a single value column: a structured (ROW-typed) value contributes one column per field,
+     * while a scalar value gets a single column named after the state. There is no synthetic
+     * ordering column and no primary key, since no column is guaranteed unique across rows.
+     *
+     * <p>Unlike the keyed table builders, this does not add {@link
+     * SavepointConnectorOptions#STATE_BACKEND_TYPE}: non-keyed state is part of the operator's own
+     * snapshot rather than of a state backend, so the hint does not apply here.
+     *
+     * @param schemaInfo the schema information returned by {@link #getNonKeyedStateSchema}
+     * @param stateName the name of the LIST or UNION state to expose
+     * @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 getOperatorStateCatalogTable(
+            NonKeyedStateSchemaInfo schemaInfo,
+            String stateName,
+            String statePath,
+            OperatorIdentifier operatorIdentifier) {
+
+        NonKeyedStateSchemaInfo.StateEntryInfo entryInfo =
+                findNonKeyedStateEntry(schemaInfo, stateName, operatorIdentifier);
+        if (entryInfo.kind != SavepointConnectorOptions.StateReaderMode.LIST
+                && entryInfo.kind != SavepointConnectorOptions.StateReaderMode.UNION) {
+            throw new IllegalArgumentException(
+                    "Operator state tables are only supported for LIST and UNION states, but '"
+                            + stateName
+                            + "' is "
+                            + entryInfo.kind
+                            + ".");
+        }
+
+        Schema.Builder schemaBuilder = Schema.newBuilder();
+        if (entryInfo.valueLogicalType instanceof RowType) {
+            for (RowType.RowField field : ((RowType) entryInfo.valueLogicalType).getFields()) {
+                schemaBuilder.column(
+                        field.getName(), LogicalTypeDataTypeConverter.toDataType(field.getType()));
+            }
+        } else {
+            schemaBuilder.column(
+                    stateName, LogicalTypeDataTypeConverter.toDataType(entryInfo.valueLogicalType));
+        }
+
+        Map<String, String> options =
+                buildBaseConnectorOptions(statePath, operatorIdentifier, entryInfo.kind);
+        options.put(SavepointConnectorOptions.FLATTENED_STATE_NAME.key(), stateName);
+
+        return CatalogTable.newBuilder().schema(schemaBuilder.build()).options(options).build();
+    }
+
+    /**
+     * Builds a {@link CatalogTable} exposing a single {@code BroadcastState} of an operator, with
+     * one row per broadcast map entry: {@code (map_key <key-type> NOT NULL, map_value)}, with a
+     * primary key on {@code map_key}.
+     *
+     * <p>The value column has a fixed name rather than the state's own name, to avoid collisions
+     * with other (reserved) column names.
+     *
+     * <p>Unlike the keyed table builders, this does not add {@link
+     * SavepointConnectorOptions#STATE_BACKEND_TYPE}: non-keyed state is part of the operator's own
+     * snapshot rather than of a state backend, so the hint does not apply here.
+     *
+     * @param schemaInfo the schema information returned by {@link #getNonKeyedStateSchema}
+     * @param stateName the name of the BROADCAST state to expose
+     * @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 getBroadcastStateCatalogTable(
+            NonKeyedStateSchemaInfo schemaInfo,
+            String stateName,
+            String statePath,
+            OperatorIdentifier operatorIdentifier) {
+
+        NonKeyedStateSchemaInfo.StateEntryInfo entryInfo =
+                findNonKeyedStateEntry(schemaInfo, stateName, operatorIdentifier);
+        if (entryInfo.kind != SavepointConnectorOptions.StateReaderMode.BROADCAST) {
+            throw new IllegalArgumentException(
+                    "Broadcast state tables are only supported for BROADCAST states, but '"
+                            + stateName
+                            + "' is "
+                            + entryInfo.kind
+                            + ".");
+        }
+
+        Schema schema =
+                Schema.newBuilder()
+                        .column(
+                                "map_key",
+                                LogicalTypeDataTypeConverter.toDataType(entryInfo.mapKeyLogicalType)
+                                        .notNull())
+                        .column(
+                                "map_value",
+                                LogicalTypeDataTypeConverter.toDataType(entryInfo.valueLogicalType))
+                        .primaryKeyNamed("PK_map_key", "map_key")
+                        .build();
+
+        Map<String, String> options =
+                buildBaseConnectorOptions(
+                        statePath,
+                        operatorIdentifier,
+                        SavepointConnectorOptions.StateReaderMode.BROADCAST);
+        options.put(SavepointConnectorOptions.FLATTENED_STATE_NAME.key(), stateName);
+
+        return CatalogTable.newBuilder().schema(schema).options(options).build();
+    }
+
     // -------------------------------------------------------------------------
     // Private helpers
     // -------------------------------------------------------------------------
 
+    /**
+     * Looks up a single non-keyed state entry, failing rather than returning {@code null} when the
+     * operator has no such state.
+     */
+    private static NonKeyedStateSchemaInfo.StateEntryInfo findNonKeyedStateEntry(
+            NonKeyedStateSchemaInfo schemaInfo,
+            String stateName,
+            OperatorIdentifier operatorIdentifier) {
+        NonKeyedStateSchemaInfo.StateEntryInfo entryInfo = schemaInfo.stateSchemas.get(stateName);
+        if (entryInfo == null) {
+            throw new IllegalArgumentException(
+                    "State '"
+                            + stateName
+                            + "' not found for operator '"
+                            + operatorIdentifier
+                            + "'.");
+        }
+        return entryInfo;
+    }
+
     /**
      * 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

flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/OperatorStateInputFormat.java

@@ -24,6 +24,7 @@
 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.common.typeutils.CustomRestoreSerializerFactory;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.core.fs.CloseableRegistry;
 import org.apache.flink.core.io.InputSplitAssigner;
@@ -36,6 +37,7 @@
 import org.apache.flink.runtime.state.OperatorStateBackend;
 import org.apache.flink.runtime.state.OperatorStateHandle;
 import org.apache.flink.runtime.state.StateBackend;
+import org.apache.flink.state.api.input.deserializer.MissingClassSerializerFactory;
 import org.apache.flink.state.api.input.splits.OperatorStateInputSplit;
 import org.apache.flink.streaming.api.operators.StreamOperatorStateContext;
 import org.apache.flink.util.CollectionUtil;
@@ -176,6 +178,14 @@ public void open(OperatorStateInputSplit split) throws IOException {
         ExecutionConfig executionConfig =
                 KeyedStateInputFormat.deserialize(
                         serializedExecutionConfig, runtimeContext.getUserCodeClassLoader());
+
+        // Deserialize any POJO/Avro state whose class is missing from the classpath into
+        // RowData/GenericRecord instead of failing the restore. Must be registered before the
+        // backend is restored, since that eagerly restores the previous serializer of every
+        // registered state. Registering it unconditionally is safe: it is only ever consulted once
+        // a snapshot has determined that the class it needs is genuinely missing.
+        CustomRestoreSerializerFactory.set(MissingClassSerializerFactory::create);
+
         final StreamOperatorStateContext context =
                 new StreamOperatorContextBuilder(
                                 runtimeContext,