com.marklogic.spark.reader.customcode.CustomCodeBatch 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.reader.customcode;

import org.apache.spark.sql.connector.read.Batch;
import org.apache.spark.sql.connector.read.InputPartition;
import org.apache.spark.sql.connector.read.PartitionReaderFactory;

import java.util.List;

class CustomCodeBatch implements Batch {

    private CustomCodeContext customCodeContext;
    private InputPartition[] inputPartitions;

    CustomCodeBatch(CustomCodeContext customCodeContext, List partitions) {
        this.customCodeContext = customCodeContext;
        if (partitions != null && partitions.size() > 1) {
            inputPartitions = new InputPartition[partitions.size()];
            for (int i = 0; i < partitions.size(); i++) {
                inputPartitions[i] = new CustomCodePartition(partitions.get(i));
            }
        } else {
            inputPartitions = new InputPartition[]{new CustomCodePartition()};
        }
    }

    @Override
    public InputPartition[] planInputPartitions() {
        return inputPartitions;
    }

    @Override
    public PartitionReaderFactory createReaderFactory() {
        return new CustomCodePartitionReaderFactory(customCodeContext);
    }

}