Java Code Examples for org.apache.flink.runtime.state.FunctionInitializationContext#getOperatorStateStore()

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {

	readyToFail = false;

	if (context.isRestored()) {
		isRunning = false;
	} else {
		isRunning = true;

		OperatorStateStore operatorStateStore = context.getOperatorStateStore();
		for (int i = 0; i < numListStates; ++i) {

			ListStateDescriptor<String> listStateDescriptor =
				new ListStateDescriptor<>("test-list-state-" + i, String.class);

			ListState<String> unionListState =
				operatorStateStore.getUnionListState(listStateDescriptor);

			for (int j = 0; j < numPartitionsPerListState; ++j) {
				unionListState.add(String.valueOf(j));
			}
		}
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {

	readyToFail = false;

	if (context.isRestored()) {
		isRunning = false;
	} else {
		isRunning = true;

		OperatorStateStore operatorStateStore = context.getOperatorStateStore();
		for (int i = 0; i < numListStates; ++i) {

			ListStateDescriptor<String> listStateDescriptor =
				new ListStateDescriptor<>("test-list-state-" + i, String.class);

			ListState<String> unionListState =
				operatorStateStore.getUnionListState(listStateDescriptor);

			for (int j = 0; j < numPartitionsPerListState; ++j) {
				unionListState.add(String.valueOf(j));
			}
		}
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
    OperatorStateStore stateStore = context.getOperatorStateStore();

    unionOffsetStates = stateStore.getUnionListState(
            new ListStateDescriptor<>(
                    OFFSETS_STATE_NAME,
                    TypeInformation.of(new TypeHint<Tuple2<String, MessageId>>() {
                    })));

    if (context.isRestored()) {
        restoredState = new TreeMap<>();
        unionOffsetStates.get().forEach(e -> restoredState.put(e.f0, e.f1));
        log.info("Source subtask {} restored state {}",
                taskIndex,
                StringUtils.join(restoredState.entrySet()));
    } else {
        log.info("Source subtask {} has no restore state", taskIndex);
    }
}
 

@Override
public final void initializeState(FunctionInitializationContext context) throws Exception {

	OperatorStateStore stateStore = context.getOperatorStateStore();

	this.unionOffsetStates = stateStore.getUnionListState(new ListStateDescriptor<>(OFFSETS_STATE_NAME,
		createStateSerializer(getRuntimeContext().getExecutionConfig())));

	if (context.isRestored()) {
		restoredState = new TreeMap<>(new KafkaTopicPartition.Comparator());

		// populate actual holder for restored state
		for (Tuple2<KafkaTopicPartition, Long> kafkaOffset : unionOffsetStates.get()) {
			restoredState.put(kafkaOffset.f0, kafkaOffset.f1);
		}

		LOG.info("Consumer subtask {} restored state: {}.", getRuntimeContext().getIndexOfThisSubtask(), restoredState);
	} else {
		LOG.info("Consumer subtask {} has no restore state.", getRuntimeContext().getIndexOfThisSubtask());
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
  if (checkpointCoder == null) {
    // no checkpoint coder available in this source
    return;
  }

  OperatorStateStore stateStore = context.getOperatorStateStore();
  @SuppressWarnings("unchecked")
  CoderTypeInformation<KV<? extends UnboundedSource<OutputT, CheckpointMarkT>, CheckpointMarkT>>
      typeInformation = (CoderTypeInformation) new CoderTypeInformation<>(checkpointCoder);
  stateForCheckpoint =
      stateStore.getOperatorState(
          new ListStateDescriptor<>(
              DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME,
              typeInformation.createSerializer(new ExecutionConfig())));

  if (context.isRestored()) {
    isRestored = true;
    LOG.info("Restoring state in the UnboundedSourceWrapper.");
  } else {
    LOG.info("No restore state for UnboundedSourceWrapper.");
  }
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {

	readyToFail = false;

	if (context.isRestored()) {
		isRunning = false;
	} else {
		isRunning = true;

		OperatorStateStore operatorStateStore = context.getOperatorStateStore();
		for (int i = 0; i < numListStates; ++i) {

			ListStateDescriptor<String> listStateDescriptor =
				new ListStateDescriptor<>("test-list-state-" + i, String.class);

			ListState<String> unionListState =
				operatorStateStore.getUnionListState(listStateDescriptor);

			for (int j = 0; j < numPartitionsPerListState; ++j) {
				unionListState.add(String.valueOf(j));
			}
		}
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	this.helper = new StreamingFileSinkHelper<>(
			bucketsBuilder.createBuckets(getRuntimeContext().getIndexOfThisSubtask()),
			context.isRestored(),
			context.getOperatorStateStore(),
			((StreamingRuntimeContext) getRuntimeContext()).getProcessingTimeService(),
			bucketCheckInterval);
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	Preconditions.checkArgument(this.restoredBucketStates == null,
		"The " + getClass().getSimpleName() + " has already been initialized.");

	try {
		initFileSystem();
	} catch (IOException e) {
		LOG.error("Error while creating FileSystem when initializing the state of the RollingSink.", e);
		throw new RuntimeException("Error while creating FileSystem when initializing the state of the RollingSink.", e);
	}

	if (this.refTruncate == null) {
		this.refTruncate = reflectTruncate(fs);
	}

	OperatorStateStore stateStore = context.getOperatorStateStore();
	restoredBucketStates = stateStore.getSerializableListState("rolling-states");

	int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
	if (context.isRestored()) {
		LOG.info("Restoring state for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);

		for (BucketState bucketState : restoredBucketStates.get()) {
			handleRestoredBucketState(bucketState);
		}

		if (LOG.isDebugEnabled()) {
			LOG.debug("{} (taskIdx= {}) restored {}", getClass().getSimpleName(), subtaskIndex, bucketState);
		}
	} else {
		LOG.info("No state to restore for the {} (taskIdx= {}).", getClass().getSimpleName(), subtaskIndex);
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	Preconditions.checkArgument(this.restoredBucketStates == null, "The operator has already been initialized.");

	try {
		initFileSystem();
	} catch (IOException e) {
		LOG.error("Error while creating FileSystem when initializing the state of the BucketingSink.", e);
		throw new RuntimeException("Error while creating FileSystem when initializing the state of the BucketingSink.", e);
	}

	if (this.refTruncate == null) {
		this.refTruncate = reflectTruncate(fs);
	}

	// We are using JavaSerializer from the flink-runtime module here. This is very naughty and
	// we shouldn't be doing it because ideally nothing in the API modules/connector depends
	// directly on flink-runtime. We are doing it here because we need to maintain backwards
	// compatibility with old state and because we will have to rework/remove this code soon.
	OperatorStateStore stateStore = context.getOperatorStateStore();
	this.restoredBucketStates = stateStore.getListState(new ListStateDescriptor<>("bucket-states", new JavaSerializer<>()));

	int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
	if (context.isRestored()) {
		LOG.info("Restoring state for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);

		for (State<T> recoveredState : restoredBucketStates.get()) {
			handleRestoredBucketState(recoveredState);
			if (LOG.isDebugEnabled()) {
				LOG.debug("{} idx {} restored {}", getClass().getSimpleName(), subtaskIndex, recoveredState);
			}
		}
	} else {
		LOG.info("No state to restore for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);
	}
}
 

@Override
    public void initializeState(FunctionInitializationContext context)
            throws Exception
    {
        OperatorStateStore stateStore = context.getOperatorStateStore();
//        ListStateDescriptor<Tuple2<String, Long>> descriptor = new ListStateDescriptor<>(
//                "sink_partition_state",
//                TypeInformation.of(new TypeHint<Tuple2<String, Long>>() {}));
//        this.unionState = stateStore.getUnionListState(descriptor);
    }
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	final int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
	this.buckets = bucketsBuilder.createBuckets(subtaskIndex);

	final OperatorStateStore stateStore = context.getOperatorStateStore();
	bucketStates = stateStore.getListState(BUCKET_STATE_DESC);
	maxPartCountersState = stateStore.getUnionListState(MAX_PART_COUNTER_STATE_DESC);

	if (context.isRestored()) {
		buckets.initializeState(bucketStates, maxPartCountersState);
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	Preconditions.checkArgument(this.restoredBucketStates == null, "The operator has already been initialized.");

	try {
		initFileSystem();
	} catch (IOException e) {
		LOG.error("Error while creating FileSystem when initializing the state of the BucketingSink.", e);
		throw new RuntimeException("Error while creating FileSystem when initializing the state of the BucketingSink.", e);
	}

	if (this.refTruncate == null) {
		this.refTruncate = reflectTruncate(fs);
	}

	OperatorStateStore stateStore = context.getOperatorStateStore();
	restoredBucketStates = stateStore.getSerializableListState("bucket-states");

	int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
	if (context.isRestored()) {
		LOG.info("Restoring state for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);

		for (State<T> recoveredState : restoredBucketStates.get()) {
			handleRestoredBucketState(recoveredState);
			if (LOG.isDebugEnabled()) {
				LOG.debug("{} idx {} restored {}", getClass().getSimpleName(), subtaskIndex, recoveredState);
			}
		}
	} else {
		LOG.info("No state to restore for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	final int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
	this.buckets = bucketsBuilder.createBuckets(subtaskIndex);

	final OperatorStateStore stateStore = context.getOperatorStateStore();
	bucketStates = stateStore.getListState(BUCKET_STATE_DESC);
	maxPartCountersState = stateStore.getUnionListState(MAX_PART_COUNTER_STATE_DESC);

	if (context.isRestored()) {
		buckets.initializeState(bucketStates, maxPartCountersState);
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	Preconditions.checkArgument(this.restoredBucketStates == null, "The operator has already been initialized.");

	try {
		initFileSystem();
	} catch (IOException e) {
		LOG.error("Error while creating FileSystem when initializing the state of the BucketingSink.", e);
		throw new RuntimeException("Error while creating FileSystem when initializing the state of the BucketingSink.", e);
	}

	if (this.refTruncate == null) {
		this.refTruncate = reflectTruncate(fs);
	}

	OperatorStateStore stateStore = context.getOperatorStateStore();
	restoredBucketStates = stateStore.getSerializableListState("bucket-states");

	int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
	if (context.isRestored()) {
		LOG.info("Restoring state for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);

		for (State<T> recoveredState : restoredBucketStates.get()) {
			handleRestoredBucketState(recoveredState);
			if (LOG.isDebugEnabled()) {
				LOG.debug("{} idx {} restored {}", getClass().getSimpleName(), subtaskIndex, recoveredState);
			}
		}
	} else {
		LOG.info("No state to restore for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);
	}
}
 

/**
 * The actual paths in this depend on the binary checkpoint so it you update this the paths
 * here have to be updated as well.
 */
@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	OperatorStateStore stateStore = context.getOperatorStateStore();

	ListState<State<T>> restoredBucketStates = stateStore.getSerializableListState("bucket-states");

	if (context.isRestored()) {

		for (State<T> states : restoredBucketStates.get()) {
			for (String bucketPath : states.bucketStates.keySet()) {
				BucketState state = states.getBucketState(new Path(bucketPath));
				String current = state.currentFile;
				long validLength = state.currentFileValidLength;

				Assert.assertEquals(expectedBucketFilesPrefix + "4", current);
				Assert.assertEquals(6, validLength);

				List<String> pendingFiles = state.pendingFiles;
				assertTrue(pendingFiles.isEmpty());

				final Map<Long, List<String>> pendingFilesPerCheckpoint = state.pendingFilesPerCheckpoint;
				Assert.assertEquals(1, pendingFilesPerCheckpoint.size());

				for (Map.Entry<Long, List<String>> entry: pendingFilesPerCheckpoint.entrySet()) {
					long checkpoint = entry.getKey();
					List<String> files = entry.getValue();

					Assert.assertEquals(0L, checkpoint);
					Assert.assertEquals(4, files.size());

					for (int i = 0; i < 4; i++) {
						Assert.assertEquals(
								expectedBucketFilesPrefix + i,
								files.get(i));
					}
				}
			}
		}
	}

	initializeCalled = true;
	super.initializeState(context);
}
 

/**
 * The actual paths in this depend on the binary checkpoint so it you update this the paths
 * here have to be updated as well.
 */
@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	OperatorStateStore stateStore = context.getOperatorStateStore();

	// We are using JavaSerializer from the flink-runtime module here. This is very naughty and
	// we shouldn't be doing it because ideally nothing in the API modules/connector depends
	// directly on flink-runtime. We are doing it here because we need to maintain backwards
	// compatibility with old state and because we will have to rework/remove this code soon.
	ListState<State<T>> restoredBucketStates = stateStore.getListState(new ListStateDescriptor<>("bucket-states", new JavaSerializer<>()));

	if (context.isRestored()) {

		for (State<T> states : restoredBucketStates.get()) {
			for (String bucketPath : states.bucketStates.keySet()) {
				BucketState state = states.getBucketState(new Path(bucketPath));
				String current = state.currentFile;
				long validLength = state.currentFileValidLength;

				Assert.assertEquals(expectedBucketFilesPrefix + "4", current);
				Assert.assertEquals(6, validLength);

				List<String> pendingFiles = state.pendingFiles;
				assertTrue(pendingFiles.isEmpty());

				final Map<Long, List<String>> pendingFilesPerCheckpoint = state.pendingFilesPerCheckpoint;
				Assert.assertEquals(1, pendingFilesPerCheckpoint.size());

				for (Map.Entry<Long, List<String>> entry: pendingFilesPerCheckpoint.entrySet()) {
					long checkpoint = entry.getKey();
					List<String> files = entry.getValue();

					Assert.assertEquals(0L, checkpoint);
					Assert.assertEquals(4, files.size());

					for (int i = 0; i < 4; i++) {
						Assert.assertEquals(
								expectedBucketFilesPrefix + i,
								files.get(i));
					}
				}
			}
		}
	}

	initializeCalled = true;
	super.initializeState(context);
}
 

/**
 * The actual paths in this depend on the binary checkpoint so it you update this the paths
 * here have to be updated as well.
 */
@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	OperatorStateStore stateStore = context.getOperatorStateStore();

	ListState<State<T>> restoredBucketStates = stateStore.getSerializableListState("bucket-states");

	if (context.isRestored()) {

		for (State<T> states : restoredBucketStates.get()) {
			for (String bucketPath : states.bucketStates.keySet()) {
				BucketState state = states.getBucketState(new Path(bucketPath));
				String current = state.currentFile;
				long validLength = state.currentFileValidLength;

				Assert.assertEquals(expectedBucketFilesPrefix + "4", current);
				Assert.assertEquals(6, validLength);

				List<String> pendingFiles = state.pendingFiles;
				assertTrue(pendingFiles.isEmpty());

				final Map<Long, List<String>> pendingFilesPerCheckpoint = state.pendingFilesPerCheckpoint;
				Assert.assertEquals(1, pendingFilesPerCheckpoint.size());

				for (Map.Entry<Long, List<String>> entry: pendingFilesPerCheckpoint.entrySet()) {
					long checkpoint = entry.getKey();
					List<String> files = entry.getValue();

					Assert.assertEquals(0L, checkpoint);
					Assert.assertEquals(4, files.size());

					for (int i = 0; i < 4; i++) {
						Assert.assertEquals(
								expectedBucketFilesPrefix + i,
								files.get(i));
					}
				}
			}
		}
	}

	initializeCalled = true;
	super.initializeState(context);
}