@Override
    public CompletableFuture<Void> write(
            TieredStoragePartitionId partitionId, List<SubpartitionBufferContext> buffersToWrite) {
        List<CompletableFuture<Void>> completableFutures = new ArrayList<>();
        buffersToWrite.forEach(
                subpartitionBuffers -> {
                    int subpartitionId = subpartitionBuffers.getSubpartitionId();
                    List<SegmentBufferContext> segmentBufferContexts =
                            subpartitionBuffers.getSegmentBufferContexts();
                    segmentBufferContexts.forEach(
                            segmentBufferContext -> {
                                CompletableFuture<Void> flushSuccessNotifier =
                                        new CompletableFuture<>();
                                ioExecutor.execute(
                                        () ->
                                                flushOrFinishSegment(
                                                        partitionId,
                                                        subpartitionId,
                                                        segmentBufferContext,
                                                        flushSuccessNotifier));
                                completableFutures.add(flushSuccessNotifier);
                            });
                });
        return FutureUtils.waitForAll(completableFutures);
    }
@Test
    void testWrite() throws IOException {
        TieredStoragePartitionId partitionId =
                TieredStorageIdMappingUtils.convertId(new ResultPartitionID());
        int numSubpartitions = 5;
        int numSegments = 10;
        int numBuffersPerSegment = 10;
        int bufferSizeBytes = 3;

        Path tieredStorageDir = Path.fromLocalFile(tempFolder.toFile());

        SegmentPartitionFileWriter partitionFileWriter =
                new SegmentPartitionFileWriter(tieredStorageDir.getPath(), numSubpartitions);

        // Prepare the buffers to be written
        List<PartitionFileWriter.SubpartitionBufferContext> subpartitionBuffers =
                generateBuffersToWrite(
                        numSubpartitions, numSegments, numBuffersPerSegment, bufferSizeBytes);

        // Write the file
        partitionFileWriter.write(partitionId, subpartitionBuffers);
        partitionFileWriter.release();

        // Check the written files
        checkWrittenSegmentFiles(
                partitionId,
                numSubpartitions,
                numSegments,
                numBuffersPerSegment,
                bufferSizeBytes,
                tieredStorageDir);
    }