CSV, GZIP와 ZIP을 표로 변환하기

파일별 Spark 작업으로 CSV 텍스트를 읽고 스키마를 추론하여 표 데이터셋으로 쓸 수 있습니다. 아래 예제는 UTF-8 CSV, .gz로 압축한 CSV와 CSV 항목들을 담은 .zip을 받습니다.

ZIP 안의 파일들은 같은 열 구조의 한 데이터셋에 속해야 합니다. 이 구현은 물리적인 줄마다 CSV 레코드를 전달하므로 따옴표 안에 줄바꿈이 들어 있는 multiline CSV에는 사용하지 않습니다. 기본 예시는 헤더 없는 CSV이며 헤더가 있는 입력을 처리하려면 reader의 header 설정과 모든 파일의 헤더 구성을 일치시켜야 합니다.

파일별 스트리밍 읽기

다음은 myproject/files/CompressedCsvToDataset.java입니다. 원천 파일을 통째로 메모리에 올리지 않고 iterator로 줄을 반환합니다. 파일명 끝의 .gz, .zip에 따라 압축 reader를 선택하고 그 밖의 입력은 일반 CSV로 읽습니다.

package myproject.files;

import com.google.common.collect.AbstractIterator;
import com.palantir.spark.binarystream.data.PortableFile;
import com.palantir.transforms.lang.java.api.Compute;
import com.palantir.transforms.lang.java.api.FoundryInput;
import com.palantir.transforms.lang.java.api.FoundryOutput;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.UncheckedIOException;
import java.nio.charset.StandardCharsets;
import java.util.zip.GZIPInputStream;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;
import org.apache.spark.TaskContext;
import org.apache.spark.api.java.function.FlatMapFunction;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.Row;
import org.apache.spark.util.TaskCompletionListener;

public final class CompressedCsvToDataset {
    @Compute
    public void compute(FoundryInput input, FoundryOutput output) {
        Dataset<PortableFile> files = input.asFiles().getFileSystem().filesAsDataset();
        Dataset<String> lines = files.flatMap(
                (FlatMapFunction<PortableFile, String>) file ->
                        file.convertToIterator(stream -> new FileLines(
                                stream, file.getLogicalPath().getFileName().toString())),
                Encoders.STRING());

        Dataset<Row> rows = files.sparkSession().read()
                .option("inferSchema", "true")
                .option("header", "false")
                .csv(lines);
        output.getDataFrameWriter(rows).write();
    }

    private static final class FileLines extends AbstractIterator<String>
            implements AutoCloseable {
        private final InputStream source;
        private final ZipInputStream zip;
        private BufferedReader reader;
        private boolean closed;

        FileLines(InputStream stream, String filename) throws IOException {
            if (filename.endsWith(".zip")) {
                zip = new ZipInputStream(stream);
                source = zip;
            } else {
                zip = null;
                source = filename.endsWith(".gz") ? new GZIPInputStream(stream) : stream;
                reader = newReader(source);
            }
            TaskContext task = TaskContext.get();
            if (task != null) {
                task.addTaskCompletionListener((TaskCompletionListener) ignored -> close());
            }
        }

        @Override
        protected String computeNext() {
            if (closed) {
                return endOfData();
            }
            try {
                while (true) {
                    if (reader == null) {
                        ZipEntry entry;
                        do {
                            entry = zip.getNextEntry();
                        } while (entry != null && entry.isDirectory());
                        if (entry == null) {
                            close();
                            return endOfData();
                        }
                        reader = newReader(zip);
                    }

                    String line = reader.readLine();
                    if (line != null) {
                        return line;
                    }
                    if (zip == null) {
                        close();
                        return endOfData();
                    }
                    zip.closeEntry();
                    reader = null;
                }
            } catch (IOException failure) {
                try {
                    close();
                } catch (UncheckedIOException closeFailure) {
                    failure.addSuppressed(closeFailure);
                }
                throw new UncheckedIOException(failure);
            }
        }

        private static BufferedReader newReader(InputStream stream) {
            return new BufferedReader(new InputStreamReader(stream, StandardCharsets.UTF_8));
        }

        @Override
        public void close() {
            if (!closed) {
                closed = true;
                try {
                    source.close();
                } catch (IOException failure) {
                    throw new UncheckedIOException(failure);
                }
            }
        }
    }
}

ZIP 디렉터리 항목과 빈 파일은 다음 항목으로 넘어갑니다. 파일 끝에 도달하거나 오류가 나면 stream을 닫고, Spark task가 끝날 때도 남은 자원을 닫습니다. TaskContext가 없는 실행에서 iterator를 반환하기 전에 stream을 미리 닫지 않습니다.

CSV reader는 inferSchema=true로 결과의 자료형을 추론합니다. 모든 파일의 값과 열 구조가 일관되어야 하며, 표의 결과는 getDataFrameWriter(...).write()로 스냅샷 저장합니다.

수동 등록

다음은 이 변환을 등록하는 전체 진입점 예입니다. 기존 Java 하위 프로젝트의 PipelineDefiner가 있으면 구현을 추가로 만들지 말고 해당 register를 기존 define()에 넣습니다.

package myproject;

import com.palantir.transforms.lang.java.api.LowLevelTransform;
import com.palantir.transforms.lang.java.api.Pipeline;
import com.palantir.transforms.lang.java.api.PipelineDefiner;
import myproject.files.CompressedCsvToDataset;

public final class CsvPipelineDefiner implements PipelineDefiner {
    @Override
    public void define(Pipeline pipeline) {
        pipeline.register(LowLevelTransform.builder()
                .computeFunctionInstance(new CompressedCsvToDataset())
                .putParameterToInputAlias("input", "/Example Project/Inputs/compressed_csv")
                .putParameterToOutputAlias("output", "/Example Project/Outputs/parsed_csv")
                .build());
    }
}

가상 경로를 실제 입력 및 출력으로 바꾸고, 이 구현을 사용하는 서비스 파일에는 myproject.CsvPipelineDefiner 한 줄을 지정합니다. 파일을 그대로 해제하거나 ZIP 하나로 묶는 경로는 파일 처리와 구분합니다.