WriteData implements VoidFunction<Iterator<Element>> { @Override public void call(final Iterator<Element> elements) throws Exception { final TaskContext taskContext = TaskContext.get(); if (null == taskContext) { throw new OperationException("Method call(Iterator<Element>) should only be called from within a Spark job"); } call(elements, taskContext.partitionId(), taskContext.taskAttemptId()); } WriteData(final Function<String, String> groupToDirectory,
                     final Schema schema,
                     final CompressionCodecName compressionCodecName); @Override void call(final Iterator<Element> elements);  }
@Test public void testTwoWritesToSamePartitionDoesntThrowException(@TempDir java.nio.file.Path tempDir) throws Exception { final Schema schema = new Schema.Builder() .type("int", new TypeDefinition.Builder() .clazz(Integer.class) .serialiser(new IntegerParquetSerialiser()) .build()) .type("string", new TypeDefinition.Builder() .clazz(String.class) .serialiser(new StringParquetSerialiser()) .build()) .entity("entity", new SchemaEntityDefinition.Builder() .vertex("string") .property("property1", "int") .aggregate(false) .build()) .edge("edge", new SchemaEdgeDefinition.Builder() .source("string") .destination("string") .property("property2", "int") .aggregate(false) .build()) .vertexSerialiser(new StringParquetSerialiser()) .build(); final Function<String, String> groupToDirectory = group -> tempDir.toAbsolutePath().toString() + "/" + group; final List<Element> elements = new ArrayList<>(); elements.add(new Entity.Builder() .group("entity") .vertex("A") .property("property1", 1) .build()); elements.add(new Edge.Builder() .group("edge") .source("B") .dest("C") .property("property2", 100) .build()); final WriteData writeData = new WriteData(groupToDirectory, schema, CompressionCodecName.GZIP); final FileSystem fileSystem = FileSystem.get(new Configuration()); final ExecutorService executorService = Executors.newFixedThreadPool(3); final List<Callable<Void>> tasks = new ArrayList<>(); LongStream.range(1000L, 1003L) .forEach(l -> { tasks.add(() -> { writeData.call(elements.iterator(), 1, l); return null; }); }); executorService.invokeAll(tasks); assertTrue(fileSystem.exists(new Path(groupToDirectory.apply("entity") + "/" + "input-1.parquet"))); assertTrue(fileSystem.exists(new Path(groupToDirectory.apply("edge") + "/" + "input-1.parquet"))); }