com.marklogic.spark.writer.WriteBatcherDataWriterFactory Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of marklogic-spark-connector Show documentation

Spark 3 connector for MarkLogic

There is a newer version: 2.4.2

Show newest version

/*
 * Copyright © 2024 MarkLogic Corporation. All Rights Reserved.
 */
package com.marklogic.spark.writer;

import org.apache.spark.sql.catalyst.InternalRow;
import org.apache.spark.sql.connector.write.DataWriter;
import org.apache.spark.sql.connector.write.DataWriterFactory;
import org.apache.spark.sql.connector.write.streaming.StreamingDataWriterFactory;
import org.apache.spark.util.SerializableConfiguration;

class WriteBatcherDataWriterFactory implements DataWriterFactory, StreamingDataWriterFactory {

    private final WriteContext writeContext;
    private final SerializableConfiguration hadoopConfiguration;

    WriteBatcherDataWriterFactory(WriteContext writeContext, SerializableConfiguration hadoopConfiguration) {
        this.writeContext = writeContext;
        this.hadoopConfiguration = hadoopConfiguration;
    }

    @Override
    public DataWriter createWriter(int partitionId, long taskId) {
        return new WriteBatcherDataWriter(writeContext, hadoopConfiguration, partitionId);
    }

    @Override
    public DataWriter createWriter(int partitionId, long taskId, long epochId) {
        return createWriter(partitionId, taskId);
    }
}