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 d32601ce400f91266899e7dbbae6e267c49a1100 Author: Gyula Fora <[email protected]> AuthorDate: Mon Jul 20 08:48:03 2026 +0200 [FLINK-40180][state-processor-api] Non-keyed/operator state related utilities Adds the schema-extraction and catalog-table-building utilities needed to expose an operator's non-keyed state (ListState/UnionState/ BroadcastState): NonKeyedStateSchemaInfo and OperatorStateSchemaInfo, plus hasNonInternalOperatorState, getNonKeyedStateSchema, getOperatorStateCatalogTable, and getBroadcastStateCatalogTable in StateTableUtils, mirroring the keyed/window state utilities added previously. Also extends getOperatorIdentifiers to include operators that register only non-keyed state. These are not yet wired into the catalog or table factory. --- .../apache/flink/state/api/StateTableUtils.java | 311 ++++++++++++++------- .../state/api/schema/NonKeyedStateSchemaInfo.java | 68 +++++ .../state/api/schema/OperatorStateSchemaInfo.java | 66 +++++ .../state/api/schema/StateSchemaExtractor.java | 107 +++++-- 4 files changed, 421 insertions(+), 131 deletions(-) diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/StateTableUtils.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/StateTableUtils.java index 731c60c0fac..9bf12c20fbf 100644 --- a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/StateTableUtils.java +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/StateTableUtils.java @@ -32,6 +32,8 @@ import org.apache.flink.runtime.state.StateBackendLoader; import org.apache.flink.runtime.state.VoidNamespaceSerializer; import org.apache.flink.runtime.state.changelog.ChangelogStateBackendHandle; import org.apache.flink.state.api.schema.KeyedStateSchemaInfo; +import org.apache.flink.state.api.schema.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.ArrayType; import org.apache.flink.table.types.logical.BigIntType; import org.apache.flink.table.types.logical.LogicalType; import org.apache.flink.table.types.logical.MapType; +import org.apache.flink.table.types.logical.RowType; import org.apache.flink.table.types.logical.VarBinaryType; import org.apache.flink.table.types.utils.LogicalTypeDataTypeConverter; @@ -67,8 +70,8 @@ import java.util.Set; import java.util.stream.Collectors; /** - * High-level utility for inspecting and reading keyed state from a checkpoint / savepoint without - * requiring user POJO classes on the classpath. + * 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 { @@ -78,15 +81,11 @@ public final class StateTableUtils { private StateTableUtils() {} /** - * Returns the {@link OperatorIdentifier}s of all operators present in the given checkpoint - * metadata that have at least one non-internal keyed state. - * - * @param metadata the checkpoint metadata to inspect - * @return list of operator identifiers; never null, may be empty + * Returns identifiers of all operators with at least one non-internal keyed or non-keyed state. */ 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,15 +112,21 @@ public final class StateTableUtils { } } - /** - * 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 - */ + 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. */ public static List<String> getKeyedStates( CheckpointMetadata metadata, OperatorIdentifier operatorId) throws IOException { OperatorState opState = findOperatorState(metadata, operatorId); @@ -133,16 +138,10 @@ public final class StateTableUtils { } /** - * Returns the {@link KeyedStateSchemaInfo} for the plain per-key (void-namespace) states of the - * given operator — the ones exposed by the {@code _keyed}/{@code _keyed_flat} tables. - * - * <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. - * - * @param metadata the checkpoint metadata to inspect - * @param operatorId identifies the operator - * @return schema information covering the key type and all registered state entries - * @throws IOException if the state header cannot be read + * Returns schema info for the plain per-key (void-namespace) states of the given operator — the + * ones exposed by the {@code _keyed}/{@code _keyed_flat} tables. Extraction is lenient: field + * names/types come from the serializer snapshot, so the user POJO class need not be on the + * classpath. */ public static KeyedStateSchemaInfo getKeyedStateSchema( CheckpointMetadata metadata, OperatorIdentifier operatorId) throws IOException { @@ -204,11 +203,7 @@ public final class StateTableUtils { } /** - * Logs that a single state's schema could not be extracted and will therefore be excluded from - * the table schema, shared by {@link #buildKeyedStateSchemaInfo}. - * - * @param label a prefix inserted before "state" in the log message (e.g. {@code "non-keyed "} - * or {@code ""}), distinguishing which caller excluded the state + * Logs that a state's schema extraction failed and it will be excluded from the table schema. */ private static void logSchemaExtractionFailure( String label, @@ -226,21 +221,10 @@ public final class StateTableUtils { } /** - * Builds a {@link CatalogTable} representing all keyed states of an operator. - * - * <p>The resulting table has one column named {@code "state_key"} for the key and one column - * per keyed state. The connector options are pre-populated so the table can be registered - * directly in a {@link org.apache.flink.table.catalog.CatalogManager}. - * - * <p>When the state backend that produced the operator's keyed state can be unambiguously - * determined from the checkpoint metadata, {@link SavepointConnectorOptions#STATE_BACKEND_TYPE} - * is pre-populated as well, so callers don't need to specify it themselves. - * - * @param metadata the checkpoint metadata the operator belongs to - * @param schemaInfo the schema information returned by {@link #getKeyedStateSchema} - * @param statePath the path to the savepoint / checkpoint - * @param operatorIdentifier identifies the operator whose state to read - * @return a {@link CatalogTable} ready for registration + * Builds a {@link CatalogTable} for all keyed states of an operator: one {@code "state_key"} + * column plus one column per keyed state. Pre-populates {@link + * SavepointConnectorOptions#STATE_BACKEND_TYPE} when it can be unambiguously determined from + * the checkpoint metadata. */ public static CatalogTable getStateCatalogTable( CheckpointMetadata metadata, @@ -252,25 +236,13 @@ public final class StateTableUtils { /** * Builds a {@link CatalogTable} exposing a single keyed LIST or MAP state flattened into one - * row per list element / map entry, rather than one row per key. - * - * <p>The resulting table has 3 columns, with a composite primary key on {@code state_key} and - * the sub-key column (the {@code state_key} value repeats across rows belonging to the same - * key, but the pair uniquely identifies a row). The third column has a fixed name — not the - * state's own name, to avoid collisions with other (reserved) column names: + * row per list element / map entry, rather than one row per key. Value column names are fixed + * rather than the state's own name to avoid collisions with reserved columns: * * <ul> - * <li>LIST: {@code (state_key, list_index, list_value)}, primary key {@code (state_key, - * list_index)} - * <li>MAP: {@code (state_key, map_key, map_value)}, primary key {@code (state_key, map_key)} + * <li>LIST: {@code (state_key, list_index, list_value)}, PK {@code (state_key, list_index)} + * <li>MAP: {@code (state_key, map_key, map_value)}, PK {@code (state_key, map_key)} * </ul> - * - * @param metadata the checkpoint metadata the operator belongs to - * @param schemaInfo the schema information returned by {@link #getKeyedStateSchema} - * @param stateName the name of the LIST or MAP state to flatten - * @param statePath the path to the savepoint / checkpoint - * @param operatorIdentifier identifies the operator whose state to read - * @return a {@link CatalogTable} ready for registration */ public static CatalogTable getFlattenedStateCatalogTable( CheckpointMetadata metadata, @@ -466,16 +438,161 @@ public final class StateTableUtils { return subKeyColumnName; } + /** + * Returns schema info for the non-keyed (operator) states of the given operator — {@code + * ListState}, {@code UnionState}, {@code BroadcastState} — exposed by the {@code _list}/{@code + * _union}/{@code _broadcast} tables. Extraction is lenient (no user POJO class required on the + * classpath); states whose schema can't be determined are excluded with a logged error. + */ + 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}, with + * one row per list element: a ROW-typed value contributes one column per field, a scalar value + * gets a single column named after the state. No ordering column and no primary key, since no + * column is guaranteed unique across rows. Unlike the keyed table builders, this does not set + * {@link SavepointConnectorOptions#STATE_BACKEND_TYPE} since non-keyed state isn't part of a + * state backend. + */ + 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}, one row per map + * entry: {@code (map_key NOT NULL, map_value)} with PK {@code map_key}. Column names are fixed + * rather than the state's own name to avoid collisions with reserved columns. Like {@link + * #getOperatorStateCatalogTable}, this does not set {@link + * SavepointConnectorOptions#STATE_BACKEND_TYPE}. + */ + 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 // ------------------------------------------------------------------------- /** - * Resolves the SQL column {@link org.apache.flink.table.types.DataType} for a single state's - * value column, forcing it nullable for VALUE-shaped state: unlike LIST/MAP (which always have - * a value, possibly empty), a {@code ValueState}/{@code ReducingState}/{@code AggregatingState} - * can legitimately hold no value (e.g. never written, or cleared by a trigger such as {@code - * CountTrigger}), in which case a read returns {@code null}. + * 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 value column's {@link org.apache.flink.table.types.DataType}, forcing it + * nullable for VALUE-shaped state: unlike LIST/MAP, a {@code ValueState} can legitimately hold + * no value (never written, or cleared by a trigger), so a read may return {@code null}. */ private static DataType stateValueColumnDataType( KeyedStateSchemaInfo.StateEntryInfo entryInfo) { @@ -497,13 +614,9 @@ public final class StateTableUtils { } /** - * Returns the base connector options ({@link FactoryUtil#CONNECTOR}, {@link - * SavepointConnectorOptions#STATE_PATH}, {@link SavepointConnectorOptions#STATE_READER_MODE}, - * and the operator identifier option) shared by every savepoint-backed {@link CatalogTable}. - * - * <p>{@code readerMode} is always set explicitly rather than relying on {@link - * SavepointConnectorOptions#STATE_READER_MODE}'s default value, so that every catalog table's - * options unambiguously reflect its schema. + * Returns the base connector options shared by every savepoint-backed {@link CatalogTable}. + * {@code readerMode} is always set explicitly rather than relying on the option's default, so + * every table's options unambiguously reflect its schema. */ private static Map<String, String> buildBaseConnectorOptions( String statePath, @@ -525,10 +638,9 @@ public final class StateTableUtils { } /** - * Adds {@link SavepointConnectorOptions#STATE_BACKEND_TYPE} to {@code options} when it can be - * unambiguously determined from the checkpoint metadata. Only meaningful for keyed state - * tables: non-keyed (list/union/broadcast) state isn't stored in a state backend, so callers - * for those table kinds must not call this. + * Adds {@link SavepointConnectorOptions#STATE_BACKEND_TYPE} when it can be unambiguously + * determined. Only meaningful for keyed state tables — non-keyed state isn't stored in a state + * backend. */ private static void withStateBackendType( Map<String, String> options, @@ -543,16 +655,12 @@ public final class StateTableUtils { } /** - * Attempts to determine the state backend (shortcut name, see {@link - * StateBackendLoader#HASHMAP_STATE_BACKEND_NAME} / {@link - * StateBackendLoader#ROCKSDB_STATE_BACKEND_NAME}) that produced the operator's keyed state, by - * inspecting the concrete {@link KeyedStateHandle} subtype found in the checkpoint metadata: - * heap/HashMap backends produce {@link KeyGroupsStateHandle}, RocksDB/ForSt backends produce - * {@link IncrementalKeyedStateHandle}. - * - * <p>Canonical-format savepoints rewrite keyed state into the backend-agnostic {@link - * KeyGroupsSavepointStateHandle}, in which case the originating backend can no longer be - * determined from the handle alone; an empty result is returned rather than guessing. + * Attempts to determine the state backend that produced the operator's keyed state, from the + * concrete {@link KeyedStateHandle} subtype: heap/HashMap backends produce {@link + * KeyGroupsStateHandle}, RocksDB/ForSt produce {@link IncrementalKeyedStateHandle}. + * Canonical-format savepoints rewrite state into the backend-agnostic {@link + * KeyGroupsSavepointStateHandle}, in which case the backend can't be told and an empty result + * is returned. */ static Optional<String> detectStateBackendType(OperatorState opState) { Set<String> detectedTypes = new HashSet<>(); @@ -584,8 +692,7 @@ public final class StateTableUtils { } else if (handle instanceof IncrementalKeyedStateHandle) { detectedTypes.add(StateBackendLoader.ROCKSDB_STATE_BACKEND_NAME); } else if (handle instanceof KeyGroupsSavepointStateHandle) { - // Canonical-format savepoints rewrite keyed state into a backend-agnostic - // format; the originating backend can no longer be told apart from the handle. + // backend-agnostic format; origin can't be told apart from the handle } else if (handle instanceof KeyGroupsStateHandle) { detectedTypes.add(StateBackendLoader.HASHMAP_STATE_BACKEND_NAME); } else { @@ -601,11 +708,9 @@ public final class StateTableUtils { } /** - * Returns {@code true} if a state is plain per-key state (registered with {@code - * VoidNamespace}) and {@code false} if it is scoped by some other namespace (e.g. a window). - * - * <p>A missing namespace snapshot (e.g. from an older savepoint format) is treated as void, - * matching pre-existing behavior. + * Returns {@code true} for plain per-key state (registered with {@code VoidNamespace}), {@code + * false} if scoped by another namespace (e.g. a window). A missing snapshot (older savepoint + * format) is treated as void. */ private static boolean isVoidNamespace(TypeSerializerSnapshot<?> namespaceSnapshot) { return namespaceSnapshot == null @@ -633,16 +738,10 @@ public final class StateTableUtils { } /** - * Partitions the user-registered states of a single operator into plain per-key - * (void-namespace) states and namespaced states, resolving the namespaced states' shared {@link - * LogicalType} along the way. - * - * <p>An operator may register states under more than one distinct namespace <em>type</em> only - * via hand-rolled state access (no built-in windowing API does this); when that happens, the - * first namespace type whose schema can be determined is kept, and every other group is - * excluded with a logged warning. - * - * @param operatorLabel a human-readable operator identifier, used only for log messages + * Partitions user-registered states into plain per-key (void-namespace) and namespaced states, + * resolving the namespaced states' shared {@link LogicalType}. An operator can register more + * than one namespace type only via hand-rolled state access; when that happens, the first + * resolvable namespace type is kept and the rest excluded with a warning. */ private static ClassifiedStates classifyStates( String operatorLabel, List<StateSchemaInfo> schemas) { diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/NonKeyedStateSchemaInfo.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/NonKeyedStateSchemaInfo.java new file mode 100644 index 00000000000..8f33bb4ecb0 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/NonKeyedStateSchemaInfo.java @@ -0,0 +1,68 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.state.api.schema; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.state.table.SavepointConnectorOptions.StateReaderMode; +import org.apache.flink.table.types.logical.LogicalType; + +import javax.annotation.Nullable; + +import java.util.LinkedHashMap; + +/** + * Schema information for all non-keyed (operator) states of a single operator — {@code ListState}, + * {@code UnionState}, and {@code BroadcastState} — extracted from a savepoint without requiring the + * user POJO classes on the classpath. + * + * <p>Unlike {@link KeyedStateSchemaInfo}, there is no shared key type across entries: operator + * state has no per-key partitioning, and only {@link StateReaderMode#BROADCAST} entries carry a + * (per-entry) map key type. + */ +@Internal +public final class NonKeyedStateSchemaInfo { + + /** State names to entry info, in savepoint registration order. */ + public final LinkedHashMap<String, StateEntryInfo> stateSchemas; + + public NonKeyedStateSchemaInfo(LinkedHashMap<String, StateEntryInfo> stateSchemas) { + this.stateSchemas = stateSchemas; + } + + /** Schema information for one non-keyed state entry. */ + public static final class StateEntryInfo { + + public final StateReaderMode kind; + + /** Logical type of the state's value (list element type, or broadcast map's value type). */ + public final LogicalType valueLogicalType; + + /** Broadcast map's key type; non-null only for {@link StateReaderMode#BROADCAST}. */ + @Nullable public final LogicalType mapKeyLogicalType; + + public StateEntryInfo( + StateReaderMode kind, + LogicalType valueLogicalType, + @Nullable LogicalType mapKeyLogicalType) { + this.kind = kind; + this.valueLogicalType = valueLogicalType; + this.mapKeyLogicalType = mapKeyLogicalType; + } + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/OperatorStateSchemaInfo.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/OperatorStateSchemaInfo.java new file mode 100644 index 00000000000..31004151863 --- /dev/null +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/OperatorStateSchemaInfo.java @@ -0,0 +1,66 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.state.api.schema; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; +import org.apache.flink.state.table.SavepointConnectorOptions.StateReaderMode; + +import javax.annotation.Nullable; + +/** + * Carries the schema information extracted from a single non-keyed (operator) state entry — {@code + * ListState}, {@code UnionState}, or {@code BroadcastState} — in a savepoint, without requiring the + * user POJO classes on the classpath. + * + * <p>Unlike {@link StateSchemaInfo}, there is no per-key serializer: non-keyed state has no key + * concept. {@link #keySnapshot} is only present for {@link StateReaderMode#BROADCAST}, where it + * represents the broadcast map's key. + */ +@Internal +public final class OperatorStateSchemaInfo { + + /** Name of the state as registered by the operator. */ + public final String stateName; + + public final StateReaderMode kind; + + /** + * Serializer snapshot for the state's value type: the map's value type for {@link + * StateReaderMode#BROADCAST}, the list element type for {@link StateReaderMode#LIST}/{@link + * StateReaderMode#UNION}. + */ + public final TypeSerializerSnapshot<?> valueSnapshot; + + /** + * Broadcast map's key serializer snapshot; non-null only for {@link StateReaderMode#BROADCAST}. + */ + @Nullable public final TypeSerializerSnapshot<?> keySnapshot; + + public OperatorStateSchemaInfo( + String stateName, + StateReaderMode kind, + TypeSerializerSnapshot<?> valueSnapshot, + @Nullable TypeSerializerSnapshot<?> keySnapshot) { + this.stateName = stateName; + this.kind = kind; + this.valueSnapshot = valueSnapshot; + this.keySnapshot = keySnapshot; + } +} diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/StateSchemaExtractor.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/StateSchemaExtractor.java index c58c7bf7737..6774c0e5d84 100644 --- a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/StateSchemaExtractor.java +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/schema/StateSchemaExtractor.java @@ -29,23 +29,26 @@ import org.apache.flink.runtime.checkpoint.OperatorSubtaskState; import org.apache.flink.runtime.state.IncrementalKeyedStateHandle; import org.apache.flink.runtime.state.KeyedBackendSerializationProxy; import org.apache.flink.runtime.state.KeyedStateHandle; +import org.apache.flink.runtime.state.OperatorBackendSerializationProxy; +import org.apache.flink.runtime.state.OperatorStateHandle; import org.apache.flink.runtime.state.StreamStateHandle; import org.apache.flink.runtime.state.metainfo.StateMetaInfoSnapshot; import org.apache.flink.runtime.state.metainfo.StateMetaInfoSnapshot.CommonOptionsKeys; import org.apache.flink.runtime.state.metainfo.StateMetaInfoSnapshot.CommonSerializerKeys; import org.apache.flink.state.api.input.deserializer.MissingClassSerializerFactory; +import org.apache.flink.state.table.SavepointConnectorOptions.StateReaderMode; import java.io.IOException; +import java.io.InputStream; import java.util.ArrayList; import java.util.Collections; import java.util.List; /** - * Utility for extracting {@link StateSchemaInfo} from a savepoint without instantiating the full - * state backend or requiring user POJO classes on the classpath. - * - * <p>It reads the {@link KeyedBackendSerializationProxy} header that every heap/RocksDB keyed state - * file starts with. + * Utility for extracting {@link StateSchemaInfo} / {@link OperatorStateSchemaInfo} from a savepoint + * without instantiating the full state backend or requiring user POJO classes on the classpath, by + * reading the {@link KeyedBackendSerializationProxy} / {@link OperatorBackendSerializationProxy} + * header that every keyed / non-keyed state file starts with. */ @Internal public final class StateSchemaExtractor { @@ -53,23 +56,14 @@ public final class StateSchemaExtractor { private StateSchemaExtractor() {} /** - * Reads state schema information from the first available keyed state handle in the given - * operator state. - * - * <p>Returns an empty list rather than throwing when no keyed state handle is found: an - * operator may register only non-keyed (list/union/broadcast) state, in which case it has no - * keyed state to describe. + * Reads state schema info from the first available keyed state handle. Returns an empty list + * (rather than throwing) when the operator registers only non-keyed state. * - * <p>The metadata header lives in different places depending on the state backend: heap ({@code - * HashMapStateBackend}) savepoints hand back a {@code KeyGroupsStateHandle}, which is itself a - * {@link StreamStateHandle} starting with the header; RocksDB (incremental or full native) - * snapshots hand back an {@link IncrementalKeyedStateHandle}, whose own data stream starts with - * the SST payload instead, so the header must be read from {@link + * <p>The metadata header's location depends on the backend: heap ({@code HashMapStateBackend}) + * hands back a {@code KeyGroupsStateHandle} that is itself a {@link StreamStateHandle} starting + * with the header, while RocksDB hands back an {@link IncrementalKeyedStateHandle} whose stream + * starts with the SST payload instead, so the header must come from {@link * IncrementalKeyedStateHandle#getMetaDataStateHandle()}. - * - * @param operatorState the operator state from a loaded savepoint / checkpoint metadata - * @return list of schema info, one entry per registered state; never null, may be empty - * @throws IOException if the state header cannot be read */ public static List<StateSchemaInfo> extractSchema(OperatorState operatorState) throws IOException { @@ -84,7 +78,7 @@ public final class StateSchemaExtractor { metadataHandle = (StreamStateHandle) handle; } if (metadataHandle != null) { - try (java.io.InputStream stream = metadataHandle.openInputStream()) { + try (InputStream stream = metadataHandle.openInputStream()) { return extractSchema(new DataInputViewStreamWrapper(stream)); } } @@ -93,10 +87,7 @@ public final class StateSchemaExtractor { return Collections.emptyList(); } - /** - * Package-private overload that accepts a {@link DataInputView} directly. Allows unit tests to - * inject pre-built byte arrays without a real filesystem. - */ + /** Package-private overload for tests to inject pre-built bytes without a real filesystem. */ static List<StateSchemaInfo> extractSchema(DataInputView in) throws IOException { ClassLoader classLoader = Thread.currentThread().getContextClassLoader(); KeyedBackendSerializationProxy<?> proxy = new KeyedBackendSerializationProxy<>(classLoader); @@ -138,4 +129,70 @@ public final class StateSchemaExtractor { return result; } + + /** + * Reads non-keyed (operator) state schema info — {@code ListState}, {@code UnionState}, {@code + * BroadcastState} — from the first available operator state handle. Returns an empty list + * (rather than throwing) when none is found. Unlike keyed state, an {@link OperatorStateHandle} + * is itself a {@link StreamStateHandle} starting with the metadata header, for every backend. + */ + public static List<OperatorStateSchemaInfo> extractOperatorSchema(OperatorState operatorState) + throws IOException { + + for (OperatorSubtaskState subtask : operatorState.getSubtaskStates().values()) { + for (OperatorStateHandle handle : subtask.getManagedOperatorState()) { + try (InputStream stream = handle.openInputStream()) { + return extractOperatorSchema(new DataInputViewStreamWrapper(stream)); + } + } + } + return Collections.emptyList(); + } + + /** Package-private overload for tests to inject pre-built bytes without a real filesystem. */ + static List<OperatorStateSchemaInfo> extractOperatorSchema(DataInputView in) + throws IOException { + ClassLoader classLoader = Thread.currentThread().getContextClassLoader(); + OperatorBackendSerializationProxy proxy = + new OperatorBackendSerializationProxy(classLoader); + CustomRestoreSerializerFactory.set(MissingClassSerializerFactory::create); + proxy.read(in); + + List<OperatorStateSchemaInfo> result = new ArrayList<>(); + + for (StateMetaInfoSnapshot meta : proxy.getOperatorStateMetaInfoSnapshots()) { + TypeSerializerSnapshot<?> valueSnapshot = + meta.getTypeSerializerSnapshot(CommonSerializerKeys.VALUE_SERIALIZER); + if (valueSnapshot == null) { + continue; + } + + // Only union-distributed list state is redistributed as a whole on rescale; every + // other distribution mode behaves like a plain (split) ListState here. + String distributionMode = + meta.getOption(CommonOptionsKeys.OPERATOR_STATE_DISTRIBUTION_MODE); + StateReaderMode kind = + OperatorStateHandle.Mode.UNION.name().equals(distributionMode) + ? StateReaderMode.UNION + : StateReaderMode.LIST; + + result.add(new OperatorStateSchemaInfo(meta.getName(), kind, valueSnapshot, null)); + } + + for (StateMetaInfoSnapshot meta : proxy.getBroadcastStateMetaInfoSnapshots()) { + TypeSerializerSnapshot<?> valueSnapshot = + meta.getTypeSerializerSnapshot(CommonSerializerKeys.VALUE_SERIALIZER); + TypeSerializerSnapshot<?> keySnapshot = + meta.getTypeSerializerSnapshot(CommonSerializerKeys.KEY_SERIALIZER); + if (valueSnapshot == null || keySnapshot == null) { + continue; + } + + result.add( + new OperatorStateSchemaInfo( + meta.getName(), StateReaderMode.BROADCAST, valueSnapshot, keySnapshot)); + } + + return result; + } }
